From 4e09328c3bb42dcdcfe1f3de916803ad0ecc016d Mon Sep 17 00:00:00 2001
From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com>
Date: Sat, 18 Jul 2026 20:53:33 +0800
Subject: [PATCH 001/255] fix(tokenizer): check for tokenizer.model after
saving it, not before (#7194)
* fix(tokenizer): check for tokenizer.model after saving it, not before
`fix_sentencepiece_tokenizer` creates its temporary directory, then returns
early unless that directory already contains a tokenizer.model:
if not os.path.exists(temporary_location):
os.makedirs(temporary_location) # fresh, empty
if not os.path.isfile(f"{temporary_location}/tokenizer.model"):
return new_tokenizer # always true
old_tokenizer.save_pretrained(temporary_location) # writes that file
The file only appears on the line after the check, so the guard is always
true and the body never runs. Nothing else writes that path either --
`convert_to_fast_tokenizer` saves into a per-name subdirectory, not
`{temporary_location}/tokenizer.model`.
Both call sites are in `get_chat_template` and are commented "Must fix the
sentence piece tokenizer since there's no tokenizer.model file!" -- the
guard defeats the exact intent the caller states. The effect is silent: the
caller still gets a working `new_tokenizer`, but the sentencepiece piece
rename is skipped, so the mapped token (e.g. the eos token remapped to
`<|im_end|>`) is missing from tokenizer.model and GGUF/llama.cpp exports
carry the old piece.
`check_if_sentencepiece_model` in save.py does the same probe in the right
order -- makedirs, save_pretrained, then isfile. Match it.
Tests are added under tests/saving/ next to the existing sentencepiece
coverage, and to the two Bucket-A lists in consolidated-tests-ci.yml, since
Repo tests (CPU) --ignores tests/saving and these need protobuf.
Co-Authored-By: Claude Opus 4.8 (1M context)
* Clear stale tokenizer.model before the sentencepiece guard
The guard now runs after old_tokenizer.save_pretrained, but the default
temporary_location is a fixed reusable directory. A fast-only tokenizer writes
no tokenizer.model, so a stale file from an earlier sentencepiece call could
pass the guard and patch the wrong model (e.g. mixing models in one process,
like a long-running server). Remove any existing tokenizer.model first, and add
a regression test.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Empty the reusable sentencepiece scratch directory each call
The final AutoTokenizer.from_pretrained reloads the whole temporary_location, so
removing only a stale tokenizer.model still let other artifacts from a previous
tokenizer (added_tokens.json, chat template, etc.) leak into the reload when the
default reusable directory is used across models in one process. Recreate the
directory instead, and add a regression test for the leaked-artifact case.
* Clear only top-level scratch files, keep subdirectories
Recreating the whole reusable directory deleted the {name} subtree that
convert_to_fast_tokenizer stores a converted tokenizer's source vocab in, so
old_tokenizer.save_pretrained could not copy tokenizer.model and the guard
returned the tokenizer unpatched for those legacy converted tokenizers. Remove
only stale top-level files (all the final reload reads) and leave subdirectories
intact. Add a regression test for the converted-source subdirectory.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Keep the current tokenizer's own source vocab when clearing
On a repeated get_chat_template(map_eos_token=True) call, the returned tokenizer's
vocab_file points back at the top-level tokenizer.model, and the cleanup deleted
that source before old_tokenizer.save_pretrained could re-emit it, so the guard
returned the tokenizer unpatched. Skip removing the old tokenizer's own source
vocab while still clearing stale files from a different tokenizer, and add a
regression test.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Use a per-call temporary directory for the sentencepiece fix
The scratch directory defaulted to a single shared path, so concurrent or repeated
get_chat_template(map_eos_token=True) calls could delete or overwrite each other's
tokenizer.model between save and reload (tripping the piece assertion or reloading
the wrong model), and stale files from an earlier tokenizer could leak into the
reload. Work in a unique per-call subdirectory instead: this isolates every call
without deleting anything the caller owns, and replaces the earlier per-file cleanup.
Tests updated to read the patched model from the reloaded directory and to cover
isolation and source-vocab preservation.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Pass only the applied token mappings into the sentencepiece fix
get_chat_template mirrors token remaps into tokenizer.model via
fix_sentencepiece_tokenizer, but two caller paths passed a mapping that did not
match what they wrote to the fast tokenizer JSON, so once the sentencepiece patch
runs the model and JSON disagree:
- the mapped-token path skipped entries whose target already existed but still
passed the full mapping, renaming a piece the JSON never changed;
- the EOS-swap path swapped both tokens in the JSON but passed only one direction,
leaving two stop_word pieces and no old EOS piece.
Pass the applied mapping (and both swap directions) instead. Add regression tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten sentencepiece guard comments
* Add SPDX license identifier to sentencepiece guard test
* Reclaim the per-call sentencepiece scratch directory
The per-call tempfile.mkdtemp fixed the shared-directory race but never cleaned
up, so a long-running process leaked one scratch dir per call. The dir cannot be
deleted eagerly for sentencepiece tokenizers because the returned tokenizer's
vocab_file points into it (a later save_pretrained copies the patched
tokenizer.model from there). Reclaim it correctly instead: remove the dir right
away on the fast-only path (the returned tokenizer never references it), and
attach a weakref.finalize so the sentencepiece dir is removed once its tokenizer
is garbage collected. Add regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten the scratch-dir reclaim comment
---------
Co-authored-by: Claude Opus 4.8 (1M context)
Co-authored-by: danielhanchen
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen
---
.github/workflows/consolidated-tests-ci.yml | 2 +
.../test_fix_sentencepiece_tokenizer_guard.py | 307 ++++++++++++++++++
unsloth/chat_templates.py | 12 +-
unsloth/tokenizer_utils.py | 18 +-
4 files changed, 333 insertions(+), 6 deletions(-)
create mode 100644 tests/saving/test_fix_sentencepiece_tokenizer_guard.py
diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml
index fa84471d36..d1bea819eb 100644
--- a/.github/workflows/consolidated-tests-ci.yml
+++ b/.github/workflows/consolidated-tests-ci.yml
@@ -268,6 +268,7 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
+ tests/saving/test_fix_sentencepiece_tokenizer_guard.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
@@ -358,6 +359,7 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
+ tests/saving/test_fix_sentencepiece_tokenizer_guard.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
diff --git a/tests/saving/test_fix_sentencepiece_tokenizer_guard.py b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py
new file mode 100644
index 0000000000..1ee523d57b
--- /dev/null
+++ b/tests/saving/test_fix_sentencepiece_tokenizer_guard.py
@@ -0,0 +1,307 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+import gc
+import os
+
+os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python")
+
+import transformers
+from transformers.utils import sentencepiece_model_pb2
+
+from unsloth.tokenizer_utils import fix_sentencepiece_tokenizer
+
+
+NORMAL, CONTROL = 1, 3
+
+
+def _spm_bytes(pieces):
+ m = sentencepiece_model_pb2.ModelProto()
+ for piece, score, typ in pieces:
+ p = m.pieces.add()
+ p.piece = piece
+ p.score = score
+ p.type = typ
+ return m.SerializeToString()
+
+
+def _read_pieces(path):
+ m = sentencepiece_model_pb2.ModelProto()
+ with open(path, "rb") as f:
+ m.ParseFromString(f.read())
+ return [p.piece for p in m.pieces]
+
+
+class _FakeTokenizer:
+ """Minimal stand-in for a sentencepiece-backed slow tokenizer.
+
+ ``save_pretrained`` writes a tokenizer.model, which is what the real slow
+ tokenizers do and what fix_sentencepiece_tokenizer reads back.
+ """
+
+ def __init__(
+ self,
+ name,
+ spm_bytes = None,
+ vocab = None,
+ ):
+ self.name = name
+ self.eos_token = ""
+ self.pad_token = ""
+ self._spm_bytes = spm_bytes
+ self._vocab = vocab or {}
+ self.saved_to = []
+
+ def save_pretrained(self, location):
+ self.saved_to.append(location)
+ os.makedirs(location, exist_ok = True)
+ if self._spm_bytes is not None:
+ with open(os.path.join(location, "tokenizer.model"), "wb") as f:
+ f.write(self._spm_bytes)
+
+ def __call__(
+ self,
+ texts,
+ add_special_tokens = False,
+ ):
+ class _Encoded:
+ pass
+
+ encoded = _Encoded()
+ encoded.input_ids = [[self._vocab[text]] for text in texts]
+ return encoded
+
+
+def _tokenizers():
+ pieces = [("", 0.0, CONTROL), ("a", -1.0, NORMAL), (" ", 0.0, CONTROL)]
+ old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2})
+ new = _FakeTokenizer("new")
+ return old, new
+
+
+class _ReloadedTokenizer:
+ """Weakref-able stand-in for the tokenizer AutoTokenizer.from_pretrained returns."""
+
+ def __init__(self, location):
+ self.location = location
+
+
+def _stub_auto_tokenizer(monkeypatch):
+ """fix_sentencepiece_tokenizer reloads the patched directory through
+ AutoTokenizer at the end; that needs a full tokenizer on disk, which is
+ out of scope here. Record the reload location and hand back a sentinel.
+ """
+ loaded = []
+
+ class _StubAutoTokenizer:
+ @staticmethod
+ def from_pretrained(location, **kwargs):
+ loaded.append(location)
+ return _ReloadedTokenizer(location)
+
+ monkeypatch.setattr(transformers, "AutoTokenizer", _StubAutoTokenizer)
+ return loaded
+
+
+def test_old_tokenizer_is_saved_so_its_model_can_be_read(tmp_path, monkeypatch):
+ """The guard must not skip the body on a fresh temporary directory.
+
+ fix_sentencepiece_tokenizer creates its scratch directory itself and then
+ checks for a tokenizer.model inside it, but that file only appears once
+ old_tokenizer.save_pretrained() has run.
+ """
+ _stub_auto_tokenizer(monkeypatch)
+ old, new = _tokenizers()
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+
+ fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+
+ assert old.saved_to, "old tokenizer was never saved: the body did not run"
+
+
+def test_token_mapping_is_applied_to_the_sentencepiece_model(tmp_path, monkeypatch):
+ loaded = _stub_auto_tokenizer(monkeypatch)
+ old, new = _tokenizers()
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+
+ # Hold the returned tokenizer so its scratch dir survives until we read it.
+ tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+
+ assert "<|im_end|>" in _read_pieces(f"{loaded[-1]}/tokenizer.model")
+ assert tok is not None
+
+
+def test_tokenizer_without_a_sentencepiece_model_is_returned_untouched(tmp_path, monkeypatch):
+ """A fast-only tokenizer writes no tokenizer.model, so the guard still
+ short-circuits and the caller gets new_tokenizer back unchanged. Its scratch
+ dir is unreferenced and reclaimed immediately.
+ """
+ _stub_auto_tokenizer(monkeypatch)
+ old = _FakeTokenizer("old", spm_bytes = None)
+ new = _FakeTokenizer("new")
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+
+ result = fix_sentencepiece_tokenizer(
+ old, new, {"": "<|im_end|>"}, temporary_location = location
+ )
+
+ assert result is new
+ assert not any(
+ name.startswith("tokenizer_") for name in os.listdir(location)
+ ), "the fast-only scratch dir was not reclaimed"
+
+
+def test_each_call_uses_a_fresh_isolated_subdirectory(tmp_path, monkeypatch):
+ """Each call must work in its own unique subdirectory, so concurrent or
+ repeated calls never share scratch files, stale artifacts never leak into
+ the reload, and nothing the caller left in the scratch location is deleted.
+ """
+ loaded = _stub_auto_tokenizer(monkeypatch)
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+ os.makedirs(location, exist_ok = True)
+
+ # A pre-existing artifact in the shared scratch location.
+ marker = os.path.join(location, "leftover.json")
+ with open(marker, "w") as f:
+ f.write("{}")
+
+ old1, new1 = _tokenizers()
+ old2, new2 = _tokenizers()
+ # Hold both returned tokenizers so their scratch dirs stay alive.
+ tok1 = fix_sentencepiece_tokenizer(
+ old1, new1, {"": "<|im_end|>"}, temporary_location = location
+ )
+ tok2 = fix_sentencepiece_tokenizer(
+ old2, new2, {"": "<|im_end|>"}, temporary_location = location
+ )
+
+ work1, work2 = loaded[0], loaded[1]
+ assert work1 != work2, "two calls reused the same directory"
+ assert os.path.dirname(work1) == location and os.path.dirname(work2) == location
+ assert os.path.isdir(work1) and os.path.isdir(work2)
+ # Nothing the caller left behind is deleted, and it never leaks into a work dir.
+ assert os.path.isfile(marker), "a pre-existing scratch file was deleted"
+ assert not os.path.isfile(os.path.join(work1, "leftover.json"))
+ assert not os.path.isfile(os.path.join(work2, "leftover.json"))
+ assert tok1 is not None and tok2 is not None
+
+
+def test_sentencepiece_scratch_dir_is_reclaimed_once_the_tokenizer_is_gone(tmp_path, monkeypatch):
+ """The scratch dir must live as long as the returned tokenizer (its vocab_file
+ points there), then be reclaimed when the tokenizer is garbage collected.
+ """
+ loaded = _stub_auto_tokenizer(monkeypatch)
+ old, new = _tokenizers()
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+
+ tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+ work = loaded[-1]
+ assert os.path.isdir(work), "scratch dir vanished while the tokenizer was alive"
+
+ del tok
+ gc.collect()
+ assert not os.path.isdir(work), "scratch dir was not reclaimed after the tokenizer was freed"
+
+
+class _CopyFromSubdirTokenizer:
+ """A slow tokenizer whose sentencepiece source lives elsewhere (like the
+ tokenizers convert_to_fast_tokenizer produces under {location}/{name}).
+ save_pretrained copies that source into the destination, as HF slow
+ tokenizers copy their vocab_file.
+ """
+
+ def __init__(self, source_model_path):
+ self.eos_token = ""
+ self.pad_token = ""
+ self._source_model_path = source_model_path
+
+ def save_pretrained(self, location):
+ os.makedirs(location, exist_ok = True)
+ if os.path.isfile(self._source_model_path):
+ with open(self._source_model_path, "rb") as src:
+ data = src.read()
+ with open(os.path.join(location, "tokenizer.model"), "wb") as dst:
+ dst.write(data)
+
+ def __call__(
+ self,
+ texts,
+ add_special_tokens = False,
+ ):
+ class _Encoded:
+ pass
+
+ encoded = _Encoded()
+ encoded.input_ids = [[2] for _ in texts]
+ return encoded
+
+
+def test_source_vocab_outside_the_work_directory_is_not_disturbed(tmp_path, monkeypatch):
+ """A tokenizer whose sentencepiece source lives elsewhere (e.g. the subtree
+ convert_to_fast_tokenizer created) is copied into the fresh work directory
+ and patched there; the original source is left untouched.
+ """
+ loaded = _stub_auto_tokenizer(monkeypatch)
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+ subdir = os.path.join(location, "some_model")
+ os.makedirs(subdir, exist_ok = True)
+
+ pieces = [("", 0.0, CONTROL), ("a", -1.0, NORMAL), (" ", 0.0, CONTROL)]
+ source_model = os.path.join(subdir, "tokenizer.model")
+ with open(source_model, "wb") as f:
+ f.write(_spm_bytes(pieces))
+
+ old = _CopyFromSubdirTokenizer(source_model)
+ new = _FakeTokenizer("new")
+ tok = fix_sentencepiece_tokenizer(old, new, {"": "<|im_end|>"}, temporary_location = location)
+
+ assert _read_pieces(source_model) == [
+ "",
+ "a",
+ " ",
+ ], "the original source vocab was modified"
+ assert "<|im_end|>" in _read_pieces(f"{loaded[-1]}/tokenizer.model")
+ assert tok is not None
+
+
+def test_swap_mapping_swaps_both_pieces_without_duplicating(tmp_path, monkeypatch):
+ """When the caller swaps eos and stop_word in the fast JSON it must pass both
+ directions here; a one-way mapping would leave two stop_word pieces and no eos.
+ """
+ loaded = _stub_auto_tokenizer(monkeypatch)
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+
+ pieces = [("", 0.0, CONTROL), ("<|im_end|>", -1.0, NORMAL), (" ", 0.0, CONTROL)]
+ old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"": 2, "<|im_end|>": 1})
+ new = _FakeTokenizer("new")
+
+ tok = fix_sentencepiece_tokenizer(
+ old, new, {"": "<|im_end|>", "<|im_end|>": ""}, temporary_location = location
+ )
+
+ result = _read_pieces(f"{loaded[-1]}/tokenizer.model")
+ assert result.count("<|im_end|>") == 1 and result.count("") == 1, result
+ assert tok is not None
+
+
+def test_only_applied_mappings_are_patched(tmp_path, monkeypatch):
+ """When the caller skips a mapping whose target already exists, it must not
+ pass that mapping here, or the skipped source token gets renamed anyway and
+ duplicates the existing target in the model.
+ """
+ loaded = _stub_auto_tokenizer(monkeypatch)
+ location = str(tmp_path / "_unsloth_sentencepiece_temp")
+
+ pieces = [
+ ("", 0.0, CONTROL),
+ ("aa", -1.0, NORMAL),
+ ("bb", -1.0, NORMAL),
+ ("X", -1.0, NORMAL),
+ ]
+ old = _FakeTokenizer("old", spm_bytes = _spm_bytes(pieces), vocab = {"aa": 1, "bb": 2})
+ new = _FakeTokenizer("new")
+
+ # Caller skipped aa->X (X already exists) and applied bb->Y, so only bb->Y is passed.
+ tok = fix_sentencepiece_tokenizer(old, new, {"bb": "Y"}, temporary_location = location)
+
+ result = _read_pieces(f"{loaded[-1]}/tokenizer.model")
+ assert result.count("X") == 1 and "Y" in result and "aa" in result, result
+ assert tok is not None
diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index 2d3674fb04..bd612db1d5 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -1929,6 +1929,9 @@ def get_chat_template(
string_vocab = tokenizer._tokenizer.to_str()
skipped = 0
+ # Only mirror applied mappings into the spm model; a skipped one would
+ # rename a piece the JSON never changed and desync the two.
+ applied_mapping = {}
for old_token, new_token in token_mapping.items():
old_count = string_vocab.count(f'"{old_token}"')
new_count = string_vocab.count(f'"{new_token}"')
@@ -1939,6 +1942,7 @@ def get_chat_template(
raise RuntimeError(f"{old_token} was not part of the tokenizer!")
else:
string_vocab = string_vocab.replace(f'"{old_token}"', f'"{new_token}"')
+ applied_mapping[old_token] = new_token
pass
pass
@@ -1973,7 +1977,7 @@ def get_chat_template(
# Must fix the sentence piece tokenizer since there's no tokenizer.model file!
from .tokenizer_utils import fix_sentencepiece_tokenizer
- tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,)
+ tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, applied_mapping,)
else:
pass
@@ -1997,8 +2001,11 @@ def get_chat_template(
string_vocab = string_vocab.replace(old_eos_token, temporary_stop_token)
string_vocab = string_vocab.replace(stop_word, old_eos_token)
string_vocab = string_vocab.replace(temporary_stop_token, stop_word)
+ # JSON swapped both, so swap both here too; a one-way map leaves two stop_word pieces.
+ sentencepiece_mapping = { old_eos_token : stop_word, stop_word : old_eos_token, }
else:
string_vocab = string_vocab.replace(old_eos_token, stop_word)
+ sentencepiece_mapping = { old_eos_token : stop_word, }
pass
new_tokenizer = tokenizer._tokenizer.from_str(string_vocab)
@@ -2017,9 +2024,8 @@ def get_chat_template(
)
# Must fix the sentence piece tokenizer since there's no tokenizer.model file!
- token_mapping = { old_eos_token : stop_word, }
from .tokenizer_utils import fix_sentencepiece_tokenizer
- tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,)
+ tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, sentencepiece_mapping,)
pass
else:
diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py
index 3a91ef188d..d6e8247ec4 100644
--- a/unsloth/tokenizer_utils.py
+++ b/unsloth/tokenizer_utils.py
@@ -17,6 +17,9 @@ from transformers.convert_slow_tokenizer import convert_slow_tokenizer
from transformers import PreTrainedTokenizerFast
import re
import os
+import shutil
+import tempfile
+import weakref
from transformers.models.llama.modeling_llama import logger
from peft import PeftModelForCausalLM
import torch
@@ -370,13 +373,19 @@ def fix_sentencepiece_tokenizer(
if not os.path.exists(temporary_location):
os.makedirs(temporary_location)
- # Check if tokenizer.model exists
- if not os.path.isfile(f"{temporary_location}/tokenizer.model"):
- return new_tokenizer
+ # Fresh per-call subdir so concurrent/repeated calls can't clobber each other's
+ # tokenizer.model or leak stale files, without deleting anything the caller owns.
+ temporary_location = tempfile.mkdtemp(prefix = "tokenizer_", dir = temporary_location)
# First save the old tokenizer
old_tokenizer.save_pretrained(temporary_location)
+ # Only sentencepiece tokenizers write tokenizer.model, so check after the save.
+ if not os.path.isfile(f"{temporary_location}/tokenizer.model"):
+ # new_tokenizer was built in memory and never references this dir, so drop it.
+ shutil.rmtree(temporary_location, ignore_errors = True)
+ return new_tokenizer
+
tokenizer_file = sentencepiece_model_pb2.ModelProto()
tokenizer_file.ParseFromString(open(f"{temporary_location}/tokenizer.model", "rb").read())
@@ -414,6 +423,9 @@ def fix_sentencepiece_tokenizer(
eos_token = new_tokenizer.eos_token,
pad_token = new_tokenizer.pad_token,
)
+ # vocab_file points here, so the dir must outlive the tokenizer (a later
+ # save_pretrained copies the patched tokenizer.model from it); reclaim it on GC.
+ weakref.finalize(tokenizer, shutil.rmtree, temporary_location, ignore_errors = True)
return tokenizer
From 9db639f708f52906ecc7ada0939063d6e5b9b166 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Sat, 18 Jul 2026 05:54:00 -0700
Subject: [PATCH 002/255] Stabilize Studio regression tests (#7192)
* Stabilize Studio regression tests
Rebuild on current main. Restore the set-membership sidebar account-block matcher
(#6647, which fixed the same order-sensitive regex, was reverted on main, so the
guard is failing on main again) and keep the watchdog replacement-race fix, whose
blocked-watchdog stub now waits without a timeout so a superseded watchdog stays
alive until cleanup regardless of scheduler load.
* Tighten the blocked-watchdog stub comment
---------
Co-authored-by: Daniel Han
---
.../tests/test_training_stop_watchdog.py | 22 ++++++++++++----
.../test_studio_text_descender_clipping.py | 25 +++++++++++--------
2 files changed, 32 insertions(+), 15 deletions(-)
diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py
index 457dfc8ea2..0cd702bce2 100644
--- a/studio/backend/tests/test_training_stop_watchdog.py
+++ b/studio/backend/tests/test_training_stop_watchdog.py
@@ -258,15 +258,27 @@ def test_watchdog_no_op_when_worker_superseded(monkeypatch):
def test_new_run_gets_its_own_watchdog(monkeypatch):
# A stale watchdog sleeping on an old proc must not stop a new run's stop from
# creating its own watcher.
- monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0)
- monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0)
b = TrainingBackend()
- _record_force_terminate(monkeypatch, b)
+ started = []
+ release = threading.Event()
+
+ def _blocked_watchdog(
+ target_proc,
+ cancel,
+ watched_job_id = None,
+ ):
+ started.append(target_proc)
+ # No timeout: the finally always releases this, so a superseded watchdog stays
+ # alive through the assertions regardless of load; as a daemon it can't hang exit.
+ release.wait()
+
+ monkeypatch.setattr(b, "_stop_watchdog_loop", _blocked_watchdog)
old_proc = _FakeProc(alive = True)
b._proc = old_proc
b._start_stop_watchdog(cancel = False)
first_wd = b._stop_watchdog
+ assert _wait_until(lambda: started == [old_proc])
# New run: fresh worker replaces the handle; its stop must get a new watcher
# even though the old (superseded) watchdog is still alive.
@@ -276,12 +288,12 @@ def test_new_run_gets_its_own_watchdog(monkeypatch):
second_wd = b._stop_watchdog
try:
+ assert _wait_until(lambda: started == [old_proc, new_proc])
assert first_wd.is_alive()
assert second_wd is not first_wd, "a new run must get its own watchdog"
assert b._stop_watchdog_proc is new_proc
finally:
- old_proc._alive = False
- new_proc._alive = False
+ release.set()
first_wd.join(timeout = 5)
second_wd.join(timeout = 5)
diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py
index 7cad6cacfe..73d1244ee3 100644
--- a/tests/studio/test_studio_text_descender_clipping.py
+++ b/tests/studio/test_studio_text_descender_clipping.py
@@ -34,17 +34,22 @@ def test_model_selector_trigger_label_uses_leading_tight():
def test_sidebar_account_block_uses_leading_tight():
src = _read(APP_SIDEBAR)
- # Match the account-block parent div regardless of its gap utility; this
- # guard is about the leading-* class, not the spacing.
- pattern = re.compile(
- r'',
- )
- matches = pattern.findall(src)
+ class_names = re.findall(r'
Date: Sat, 18 Jul 2026 20:54:50 +0800
Subject: [PATCH 003/255] fix(dataprep): skip .jsonl lines that are valid JSON
but not objects (#7195)
* fix(dataprep): skip .jsonl lines that are valid JSON but not objects
`_read_file_by_format` json.loads each line and hands the result to
`_extract_text_from_json`, which assumes a dict:
for field in self._TEXT_FIELDS:
if field in data and isinstance(data[field], str):
A JSON line does not have to be an object -- `"context"`, `["text"]` and
`42` are all valid JSON. For those, `field in data` stops being a key
lookup and becomes a substring/membership test, so `data[field]` raises:
"context" -> "text" in "context" is True (substring!)
-> TypeError: string indices must be integers
["text", "foo"] -> TypeError: list indices must be integers
42 -> TypeError: argument of type 'int' is not iterable
The TypeError escapes past `except json.JSONDecodeError: continue`, so the
whole load dies on one odd line.
That except clause is also the tell: a *malformed* line is already skipped
gracefully. A *well-formed* line that happens not to be an object should be
too -- it carries no text either way. This makes the two agree.
Reachable from `unsloth-cli.py:253` (`--dataset foo.jsonl` auto-detect) and
`RawTextDataLoader` is exported from `unsloth/__init__.py`.
Co-Authored-By: Claude Opus 4.8 (1M context)
* Slim the non-object jsonl regression test and shorten the guard comment
---------
Co-authored-by: Claude Opus 4.8 (1M context)
Co-authored-by: Daniel Han
---
tests/test_raw_text.py | 18 ++++++++++++++++++
unsloth/dataprep/raw_text.py | 4 ++++
2 files changed, 22 insertions(+)
diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py
index ba16e0cfc4..18549adfe8 100644
--- a/tests/test_raw_text.py
+++ b/tests/test_raw_text.py
@@ -295,7 +295,25 @@ def test_smart_chunk_text_single_chunk_no_eos_returns_plain_list():
return True
+def test_load_from_file_skips_non_object_json_lines():
+ """Non-object .jsonl lines (valid JSON, not dicts) are skipped, not fatal."""
+ # "context" contains "text", ["text"] holds it, 42 isn't iterable -- each
+ # would reach data[field] and raise TypeError without the isinstance guard.
+ with tempfile.NamedTemporaryFile("w", suffix = ".jsonl", delete = False) as f:
+ f.write('"context"\n["text", "x"]\n42\n{"text": "keep this"}\n')
+ path = f.name
+ try:
+ text = RawTextDataLoader(None)._read_file_by_format(path, "json_lines")
+ assert text == "keep this", text
+ finally:
+ os.unlink(path)
+
+ print("test_load_from_file_skips_non_object_json_lines passed")
+ return True
+
+
if __name__ == "__main__":
success = test_raw_text_loader()
success = test_smart_chunk_text_single_chunk_no_eos_returns_plain_list() and success
+ success = test_load_from_file_skips_non_object_json_lines() and success
sys.exit(0 if success else 1)
diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py
index 128d966ecd..8623285a25 100644
--- a/unsloth/dataprep/raw_text.py
+++ b/unsloth/dataprep/raw_text.py
@@ -236,6 +236,10 @@ class RawTextDataLoader:
def _extract_text_from_json(self, data):
"""Extract text from JSON object using common field names."""
+ # Skip non-object lines (str/list/number): `field in data` would be a
+ # substring/membership test, not a key lookup, and `data[field]` raises.
+ if not isinstance(data, dict):
+ return ""
for field in self._TEXT_FIELDS:
if field in data and isinstance(data[field], str):
return data[field]
From 9073f07705488601ab4ff59e3828fe698083894e Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Sat, 18 Jul 2026 21:08:55 -0700
Subject: [PATCH 004/255] fix(studio): equal padding in the dataset source
segmented control (#7230)
* fix(studio): equal padding in dataset source segmented control
* fix(studio): scope dataset source pill layoutId per component instance
---
.../studio/sections/dataset-section.tsx | 43 ++++++++++++-------
1 file changed, 28 insertions(+), 15 deletions(-)
diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx
index e7bab1f47f..6aa9329609 100644
--- a/studio/frontend/src/features/studio/sections/dataset-section.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx
@@ -70,11 +70,13 @@ import {
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useNavigate } from "@tanstack/react-router";
+import { motion, useReducedMotion } from "motion/react";
import {
type ChangeEvent,
type DragEvent,
useCallback,
useEffect,
+ useId,
useMemo,
useRef,
useState,
@@ -153,6 +155,9 @@ function normalizeSliceInput(value: string): string | null {
export function DatasetSection() {
const t = useT();
const navigate = useNavigate();
+ const reducedMotion = useReducedMotion();
+ // Scopes the pill layoutId so multiple instances never share one.
+ const sourcePillLayoutId = useId();
const {
dataset,
datasetSource,
@@ -686,6 +691,9 @@ export function DatasetSection() {
{(() => {
// Hub-style sliding-pill segmented control, matching the Hub tabs
// via the shared .hub-tab-toggle / .hub-tab-toggle-pill classes.
+ // flex-auto buttons share leftover space equally so padding stays
+ // equal for all labels; the pill sits inside the active button so
+ // it always matches its bounds.
const sourceTabs: {
value: "huggingface" | "upload" | "s3";
label: string;
@@ -696,24 +704,12 @@ export function DatasetSection() {
? []
: [{ value: "s3" as const, label: "Amazon S3" }]),
];
- const activeIndex = Math.max(
- 0,
- sourceTabs.findIndex((item) => item.value === datasetSource),
- );
return (
-
{sourceTabs.map((item) => (
- {item.label}
+ {datasetSource === item.value && (
+
+ )}
+ {item.label}
))}
From d8aa0df66e728cacd677773adb9c4a9ce66cd13a Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Sun, 19 Jul 2026 11:09:49 +0530
Subject: [PATCH 005/255] Studio: keep stale canvas from surviving into a new
chat (#7229)
---
.../src/features/chat/artifacts/artifact-card.tsx | 7 +++++--
studio/frontend/src/features/chat/chat-page.tsx | 10 ++++------
2 files changed, 9 insertions(+), 8 deletions(-)
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
index ee8c26abf1..0345dc6e2a 100644
--- a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
+++ b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
@@ -8,7 +8,7 @@ import { cn } from "@/lib/utils";
import { useAuiState } from "@assistant-ui/react";
import { LayoutTwoColumnIcon as Layout2ColumnIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
-import { useLayoutEffect, useMemo } from "react";
+import { useLayoutEffect, useMemo, useRef } from "react";
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
import type { ArtifactViewMode } from "./html-frame";
import {
@@ -83,15 +83,18 @@ export function ArtifactCard({
],
);
const surface = artifactThreadId ? "panel" : "overlay";
+ // Once per mount, so a view-change cleanup can't re-trigger a stale open.
+ const autoOpenAttemptedRef = useRef(false);
useLayoutEffect(() => {
if (selectedArtifactId === artifact.id) {
updateArtifact(artifact);
}
- if (!autoOpen) {
+ if (!autoOpen || autoOpenAttemptedRef.current) {
return;
}
+ autoOpenAttemptedRef.current = true;
if (hasAutoOpenedArtifact(artifact.id)) {
return;
}
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 380ce0e0ab..ec0ad977bf 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -247,13 +247,13 @@ const SingleContent = memo(function SingleContent({
useState(false);
const [isArtifactSurfaceVisible, setIsArtifactSurfaceVisible] =
useState(false);
+ // Without a URL threadId the artifact must belong to the active thread.
const showArtifactPanel = Boolean(
artifact &&
artifactSurface === "panel" &&
(threadId
? !artifact.threadId || artifact.threadId === threadId
- : Boolean(newThreadNonce) ||
- Boolean(artifact.threadId && artifact.threadId === activeThreadId)),
+ : Boolean(artifact.threadId && artifact.threadId === activeThreadId)),
);
const artifactLayoutActive = showArtifactPanel || isArtifactPanelLayoutActive;
@@ -1771,10 +1771,8 @@ export function ChatPage({
useEffect(() => {
if (view.mode !== "single") return;
- if (view.threadId || view.newThreadNonce || !selectedArtifact) return;
- // view excludes __LOCALID_ threads (they fall through to mode:"single"
- // with no threadId/nonce). Don't close a canvas whose thread is the
- // active local thread.
+ if (view.threadId || !selectedArtifact) return;
+ // Close any canvas that doesn't belong to the active thread.
if (
selectedArtifact.threadId &&
selectedArtifact.threadId === activeThreadId
From 95fa3fbe30198917bdcce4000881530b9adab079 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Sat, 18 Jul 2026 22:47:00 -0700
Subject: [PATCH 006/255] Allow API key for Ollama connections (#7173)
The Connections form hid the API key field for the Ollama preset, which
blocked Ollama cloud (it requires a key). Show the optional field for
Ollama; the backend already sends Authorization: Bearer when a key is
set and omits the header when empty, so local keyless servers are
unaffected.
Fixes #7163
---
studio/backend/core/inference/providers.py | 5 +++--
studio/frontend/src/features/chat/chat-providers-dialog.tsx | 4 ++--
studio/frontend/src/features/chat/external-providers.ts | 5 +++--
3 files changed, 8 insertions(+), 6 deletions(-)
diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py
index 5b72373c03..d3bffc2f3d 100644
--- a/studio/backend/core/inference/providers.py
+++ b/studio/backend/core/inference/providers.py
@@ -276,8 +276,9 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": (
- "Local Ollama server. OpenAI-compatible /v1/chat/completions; "
- "no API key. Surfaced via CUSTOM_PROVIDER_PRESETS in the frontend."
+ "Ollama server (local or cloud). OpenAI-compatible "
+ "/v1/chat/completions; API key optional (required by Ollama "
+ "cloud). Surfaced via CUSTOM_PROVIDER_PRESETS in the frontend."
),
"hidden": True,
},
diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx
index 95e5cfbd79..e39955e576 100644
--- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx
+++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx
@@ -244,8 +244,8 @@ export function ChatProvidersSettings({
(s) => s.setConnectionsEnabled,
);
const isCustomProvider = isCustomProviderType(providerType);
- // Local presets (Ollama, llama.cpp) never use API keys — hide the field.
- // vLLM may optionally use a bearer token on secured deployments.
+ // llama.cpp hides the key field. Ollama and vLLM show an optional key:
+ // Ollama cloud and secured vLLM need one; local servers leave it empty.
const showApiKeyField = !customPresetSkipsApiKeyField(providerType);
const showReasoningToggle = supportsProviderReasoningToggle(providerType);
diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts
index bc718abbba..eb9e4656b0 100644
--- a/studio/frontend/src/features/chat/external-providers.ts
+++ b/studio/frontend/src/features/chat/external-providers.ts
@@ -184,11 +184,12 @@ export function supportsRemoteModelCatalog(
);
}
-/** Presets that skip the API-key field (local servers with no auth by default). */
+/** Presets that hide the API-key field. Ollama is not skipped: Ollama cloud
+ * requires a key; local servers leave the optional field empty. */
export function customPresetSkipsApiKeyField(
providerType: string | null | undefined,
): boolean {
- return providerType === "ollama" || providerType === "llama_cpp";
+ return providerType === "llama_cpp";
}
/** Catalog load plus optional manual model IDs. */
From c2cf2b4a1e023f3e9de80b8889971f4d04ded9c7 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Sat, 18 Jul 2026 23:15:13 -0700
Subject: [PATCH 007/255] Studio: keep the permission pill label when composer
pills collapse (#7231)
With 5 or more pills active the composer collapses every pill to an
icon, which hid the Bypass permissions label behind a small glyph.
Exempt the permission pill via data-keep-label so it always shows its
label, with the collapsed icons lining up to its right. Since the pill
is never icon-only now, drop the compact-mode fallthrough in the glyph
off switch so it works while the other pills are collapsed.
---
.../src/features/chat/permission-mode-select.tsx | 13 ++++---------
studio/frontend/src/index.css | 4 +++-
2 files changed, 7 insertions(+), 10 deletions(-)
diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx
index 4277c1bfcf..a9cb8ce5d1 100644
--- a/studio/frontend/src/features/chat/permission-mode-select.tsx
+++ b/studio/frontend/src/features/chat/permission-mode-select.tsx
@@ -278,28 +278,23 @@ export function PermissionModeComposerPill({
data-pill-label={active.label}
data-active={fullAccess ? "true" : "false"}
data-variant={fullAccess ? "danger" : undefined}
+ data-keep-label="true"
aria-label="Permission level for tool calls"
title={`${active.label}: ${active.description}`}
>
{/* The icon doubles as an off switch (mirrors the MCP pill): hover
swaps it to an X; clicking it turns bypass permissions Off (no
- prompts, sandbox on) without opening the menu. In compact
- icon-only mode the glyph is the whole button, so clicks fall
- through and open the menu instead. */}
+ prompts, sandbox on) without opening the menu. data-keep-label
+ exempts this pill from compact icon-only mode, so the off switch
+ stays clickable even while the other pills are collapsed. */}
{
- if (e.currentTarget.closest('[data-pill-compact="true"]')) {
- return;
- }
e.stopPropagation();
}}
onClick={(e) => {
- if (e.currentTarget.closest('[data-pill-compact="true"]')) {
- return;
- }
e.stopPropagation();
setPermissionMode("off");
}}
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 333ed70f6a..6d4c21eec8 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -1467,7 +1467,9 @@ html[data-chat-font] .aui-root {
}
/* With more than 4 tools on, drop pill labels to icons only to cut clutter.
- Compare keeps its label via data-keep-label. */
+ Compare and the bypass-permissions pill keep their labels via
+ data-keep-label; the permission pill sits before the collapsed icons, so
+ they line up to its right. */
[data-pill-compact="true"]
.composer-pill-btn:not([data-keep-label])
> span:not(.composer-pill-glyph) {
From 4e4af72b9ce8be38155da4f84b81542cd706981b Mon Sep 17 00:00:00 2001
From: Long Yixing
Date: Sun, 19 Jul 2026 15:27:55 +0800
Subject: [PATCH 008/255] fix(studio): honor MLX adapter state in compare mode
(#7196)
* fix(studio): add MLX adapter state control
* fix(studio): honor MLX adapter comparison state
* fix(studio): keep enabled MLX adapters permissive
* Studio: preserve public error message on MLX compare-mode adapter failures
generate_with_adapter_control raised a plain RuntimeError, which the compare
route handled with the generic handler that drops the operational message.
Raise GenStreamErrorRaised(public=chunk.public) instead and catch it in the
streaming and non-streaming consumers, matching the safetensors tool loop, so
errors like 'model is being unloaded' surface their real message.
* Studio: re-emit VLM think prefill inside the adapter context
The compare-mode merge dropped _generate_vlm's upfront yield of the prefilled
block. Restore it as the first snapshot inside the lock+adapter context
(matching _generate_text) so the UI renders the thinking block during prefill
and a cancel/error before the first token does not drop it. Adds a regression
test asserting the prefill is emitted first, after entering the adapter context.
---------
Co-authored-by: danielhanchen
---
.../backend/core/inference/mlx_inference.py | 93 +++++-
studio/backend/core/inference/orchestrator.py | 17 +-
studio/backend/core/inference/worker.py | 31 +-
studio/backend/routes/inference.py | 14 +
.../tests/test_mlx_inference_backend.py | 267 ++++++++++++++++--
.../tests/test_orchestrator_unload_cancel.py | 64 +++++
6 files changed, 443 insertions(+), 43 deletions(-)
diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py
index e7a90b4307..e78c93b6f3 100644
--- a/studio/backend/core/inference/mlx_inference.py
+++ b/studio/backend/core/inference/mlx_inference.py
@@ -8,6 +8,7 @@ instead of torch/transformers for model loading and generation.
import json
import os
import threading
+from contextlib import contextmanager
from typing import Optional, Generator
from core.inference.message_content import content_to_text
from core.inference.runtime_context import runtime_context_length
@@ -20,6 +21,63 @@ from loggers import get_logger
logger = get_logger(__name__)
+def _mlx_adapter_modules(model):
+ """Return bypassable adapter entries and unsupported wrapper paths."""
+ adapters = []
+ unsupported = []
+ for path, module in model.named_modules():
+ if not path or not (hasattr(module, "lora_a") and hasattr(module, "lora_b")):
+ continue
+ base = getattr(module, "linear", None)
+ if base is None:
+ base = getattr(module, "embedding", None)
+ if base is None:
+ unsupported.append(path)
+ else:
+ adapters.append((path, module, base))
+ return adapters, unsupported
+
+
+@contextmanager
+def _temporary_mlx_adapter_state(model, use_adapter):
+ """Select base or adapter modules for one request, then restore the tree."""
+ if use_adapter is None:
+ yield
+ return
+ if isinstance(use_adapter, str):
+ raise NotImplementedError(
+ "Unsloth MLX: named adapter selection is not supported; use True for "
+ "the loaded adapter or False for the base model."
+ )
+ if use_adapter is not True and use_adapter is not False:
+ raise TypeError("Unsloth MLX: use_adapter must be None, True, False, or a string.")
+
+ adapters, unsupported = _mlx_adapter_modules(model)
+ if use_adapter is True:
+ if not adapters and not unsupported:
+ logger.warning("MLX adapter requested, but the active model has no adapter layers")
+ yield
+ return
+ if unsupported:
+ raise RuntimeError(
+ "Unsloth MLX: cannot disable adapter layers without their base modules: "
+ + ", ".join(unsupported[:5])
+ )
+ if not adapters:
+ yield
+ return
+
+ from mlx.utils import tree_unflatten
+
+ base_modules = tree_unflatten([(path, base) for path, _, base in adapters])
+ adapter_modules = tree_unflatten([(path, wrapper) for path, wrapper, _ in adapters])
+ try:
+ model.update_modules(base_modules)
+ yield
+ finally:
+ model.update_modules(adapter_modules)
+
+
def _mlx_vlm_model_config(model):
"""Return the loaded MLX model config and its type, preferring whichever of
config / _config actually carries a model_type."""
@@ -508,6 +566,7 @@ class MLXInferenceBackend:
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
+ _adapter_state = None,
) -> Generator[str, None, None]:
if self._model is None:
raise RuntimeError("No model loaded")
@@ -552,6 +611,7 @@ class MLXInferenceBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
+ _adapter_state = _adapter_state,
)
else:
stream = self._generate_text(
@@ -568,6 +628,7 @@ class MLXInferenceBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
+ _adapter_state = _adapter_state,
)
yield from stream
@@ -587,6 +648,7 @@ class MLXInferenceBackend:
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
+ _adapter_state = None,
):
from mlx_lm import stream_generate
from mlx_lm.sample_utils import make_sampler, make_logits_processors
@@ -635,10 +697,6 @@ class MLXInferenceBackend:
think_prefix = detect_think_prefill(
prompt, getattr(self._tokenizer, "all_special_tokens", None)
)
- # Emit it before the first token so the block renders during prefill.
- if think_prefix:
- yield think_prefix
-
sampler = make_sampler(
temp = temperature,
top_p = top_p,
@@ -680,9 +738,12 @@ class MLXInferenceBackend:
type(self._model).__name__,
type(self._tokenizer).__name__,
)
- with self._generation_lock:
+ with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state):
final_response = None
try:
+ # Enter request-scoped model state before yielding any response.
+ if think_prefix:
+ yield think_prefix
gen_kwargs = dict(
prompt = prompt,
max_tokens = max_new_tokens,
@@ -749,6 +810,7 @@ class MLXInferenceBackend:
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
+ _adapter_state = None,
):
from mlx_vlm import stream_generate as vlm_stream
@@ -852,9 +914,6 @@ class MLXInferenceBackend:
# Re-emit an open prefill from the prompt (see _generate_text).
cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None))
- # Emit it before the first token so the block renders during prefill.
- if cumulative:
- yield cumulative
logger.info(
"VLM generating: prompt_len=%d, has_image=%s",
len(prompt),
@@ -891,9 +950,18 @@ class MLXInferenceBackend:
def _stream_vlm_snapshots():
nonlocal cumulative
- with self._generation_lock:
+ # Hold the generation lock AND the request-scoped adapter state for the
+ # whole stream so Base-vs-LoRA compare mode honors use_adapter and the
+ # wrapper tree is restored on completion, cancellation, or close.
+ with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state):
final_response = None
try:
+ # Emit any prefilled block before the first token so the
+ # UI renders it during prefill, matching _generate_text. Done
+ # inside the adapter context so an unsupported request raises
+ # before any output escapes.
+ if cumulative:
+ yield cumulative
for response in vlm_stream(
self._model,
self._processor,
@@ -927,8 +995,11 @@ class MLXInferenceBackend:
cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
- # MLX LoRA adapter toggling not yet supported; generate normally
- yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs)
+ yield from self.generate_chat_response(
+ cancel_event = cancel_event,
+ _adapter_state = use_adapter,
+ **gen_kwargs,
+ )
def reset_generation_state(self):
import mlx.core as mx
diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index 3afda74411..eaa474d9b8 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/studio/backend/core/inference/orchestrator.py
@@ -1502,14 +1502,27 @@ class InferenceOrchestrator:
Uses the dispatcher path (no _gen_lock) so compare-mode requests
don't block each other; the subprocess serializes them via its
- sequential command loop.
+ sequential command loop. Backend failures raise instead of becoming
+ assistant text.
"""
- yield from self._generate_dispatched(
+ stream = self._generate_dispatched(
use_adapter = use_adapter,
cancel_event = cancel_event,
stats_holder = stats_holder,
**gen_kwargs,
)
+ try:
+ for chunk in stream:
+ if isinstance(chunk, GenStreamError):
+ # Preserve the public/operational flag so the route can surface
+ # the real message (e.g. "model is being unloaded") instead of a
+ # generic error. Mirrors the safetensors tool loop's _single_turn.
+ raise GenStreamErrorRaised(str(chunk), public = chunk.public)
+ yield chunk
+ finally:
+ close = getattr(stream, "close", None)
+ if callable(close):
+ close()
def _generate_inner(
self,
diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index e4628dcea8..9f301ba37e 100644
--- a/studio/backend/core/inference/worker.py
+++ b/studio/backend/core/inference/worker.py
@@ -513,20 +513,25 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
logger.info("Starting text generation for request_id=%s", request_id)
- for cumulative_text in generator:
- # cancel_event is an mp.Event — checked instantly, no queue polling.
- if cancel_event.is_set():
- logger.info("Generation cancelled for request %s", request_id)
- break
+ try:
+ for cumulative_text in generator:
+ # cancel_event is an mp.Event — checked instantly, no queue polling.
+ if cancel_event.is_set():
+ logger.info("Generation cancelled for request %s", request_id)
+ break
- _send_response(
- resp_queue,
- {
- "type": "token",
- "request_id": request_id,
- "text": cumulative_text,
- },
- )
+ _send_response(
+ resp_queue,
+ {
+ "type": "token",
+ "request_id": request_id,
+ "text": cumulative_text,
+ },
+ )
+ finally:
+ close = getattr(generator, "close", None)
+ if callable(close):
+ close()
_send_response(
resp_queue,
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 52ea1f86a3..9299e26d56 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -9443,6 +9443,13 @@ async def openai_chat_completions(
backend.reset_generation_state()
api_monitor.finish(monitor_id, "cancelled")
raise
+ except GenStreamErrorRaised as exc:
+ # Adapter-controlled (compare-mode) backend failure. Honor the
+ # public flag so operational errors surface their real message.
+ backend.reset_generation_state()
+ _msg = _friendly_gen_stream_error(exc)
+ api_monitor.fail(monitor_id, _msg)
+ yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}})
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI streaming: {e}", exc_info = True)
@@ -9591,6 +9598,13 @@ async def openai_chat_completions(
except HTTPException:
raise
+ except GenStreamErrorRaised as exc:
+ # Adapter-controlled (compare-mode) backend failure. Honor the public
+ # flag so operational errors surface their real message.
+ backend.reset_generation_state()
+ _msg = _friendly_gen_stream_error(exc)
+ api_monitor.fail(monitor_id, _msg)
+ raise HTTPException(status_code = 500, detail = _msg)
except Exception as e:
backend.reset_generation_state()
logger.error(f"Error during OpenAI completion: {e}", exc_info = True)
diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py
index fa50cd84d6..7b8aefb722 100644
--- a/studio/backend/tests/test_mlx_inference_backend.py
+++ b/studio/backend/tests/test_mlx_inference_backend.py
@@ -2,6 +2,7 @@
import sys
import types
+from contextlib import contextmanager
from types import SimpleNamespace
import pytest
@@ -40,12 +41,16 @@ class _DummyModel:
def _install_fake_mlx(monkeypatch):
mlx_pkg = types.ModuleType("mlx")
mlx_core = types.ModuleType("mlx.core")
+ mlx_utils = types.ModuleType("mlx.utils")
mlx_core.metal = _DummyMetal()
mlx_core.set_wired_limit = _DummyMX.set_wired_limit
mlx_core.device_info = _DummyMX.device_info
+ mlx_utils.tree_unflatten = dict
mlx_pkg.core = mlx_core
+ mlx_pkg.utils = mlx_utils
monkeypatch.setitem(sys.modules, "mlx", mlx_pkg)
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core)
+ monkeypatch.setitem(sys.modules, "mlx.utils", mlx_utils)
def _install_fake_fast_mlx(monkeypatch, calls):
@@ -68,6 +73,99 @@ def _install_fake_fast_mlx(monkeypatch, calls):
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader)
+class _AdapterTree:
+ def __init__(self, modules):
+ self.modules = dict(modules)
+
+ def named_modules(self):
+ return list(self.modules.items())
+
+ def update_modules(self, modules):
+ self.modules.update(modules)
+
+
+def test_temporary_mlx_adapter_state_bypasses_and_restores_wrappers(monkeypatch):
+ _install_fake_mlx(monkeypatch)
+ from core.inference.mlx_inference import _temporary_mlx_adapter_state
+
+ base = object()
+ wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), linear = base, m = object())
+ model = _AdapterTree({"model.layers.0.proj": wrapper})
+
+ with pytest.raises(RuntimeError, match = "generation failed"):
+ with _temporary_mlx_adapter_state(model, False):
+ assert model.modules["model.layers.0.proj"] is base
+ raise RuntimeError("generation failed")
+ assert model.modules["model.layers.0.proj"] is wrapper
+
+
+def test_temporary_mlx_adapter_state_validates_requests():
+ from core.inference.mlx_inference import _temporary_mlx_adapter_state
+
+ wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), embedding = object())
+ model = _AdapterTree({"embed_tokens": wrapper})
+ with _temporary_mlx_adapter_state(model, True):
+ assert model.modules["embed_tokens"] is wrapper
+ with pytest.raises(NotImplementedError, match = "named adapter"):
+ with _temporary_mlx_adapter_state(model, "other"):
+ pass
+
+ base_model = _AdapterTree({"proj": object()})
+ with _temporary_mlx_adapter_state(base_model, None):
+ pass
+ with _temporary_mlx_adapter_state(base_model, True):
+ pass
+
+ unsupported = _AdapterTree({"proj": SimpleNamespace(lora_a = object(), lora_b = object())})
+ with _temporary_mlx_adapter_state(unsupported, True):
+ pass
+ with pytest.raises(RuntimeError, match = "without their base modules"):
+ with _temporary_mlx_adapter_state(unsupported, False):
+ pass
+
+
+def test_temporary_mlx_adapter_state_uses_real_mlx_module_tree():
+ nn = pytest.importorskip("mlx.nn")
+ pytest.importorskip("mlx_lm")
+ from mlx_lm.models.switch_layers import SwitchLinear
+ from mlx_lm.tuner.dora import DoRALinear
+ from mlx_lm.tuner.lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear
+
+ from core.inference.mlx_inference import _temporary_mlx_adapter_state
+
+ class _Layer(nn.Module):
+ def __init__(self):
+ super().__init__()
+ quantized = nn.QuantizedLinear.from_linear(nn.Linear(32, 32), group_size = 32, bits = 4)
+ self.quantized_proj = LoRALinear.from_base(quantized)
+ self.dora_proj = DoRALinear.from_base(nn.Linear(4, 4))
+
+ class _Model(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.layers = [_Layer()]
+ self.embed_tokens = LoRAEmbedding.from_base(nn.Embedding(16, 4))
+ self.experts = LoRASwitchLinear.from_base(SwitchLinear(4, 4, 2))
+
+ model = _Model()
+ wrappers = {
+ path: module
+ for path, module in model.named_modules()
+ if hasattr(module, "lora_a") and hasattr(module, "lora_b")
+ }
+ bases = {
+ path: getattr(module, "linear", getattr(module, "embedding", None))
+ for path, module in wrappers.items()
+ }
+
+ with _temporary_mlx_adapter_state(model, False):
+ live = dict(model.named_modules())
+ assert all(live[path] is base for path, base in bases.items())
+
+ restored = dict(model.named_modules())
+ assert all(restored[path] is wrapper for path, wrapper in wrappers.items())
+
+
def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):
_install_fake_mlx(monkeypatch)
calls = []
@@ -333,10 +431,87 @@ def test_mlx_generate_chat_response_accepts_template_kwargs():
), f"{name!r} must default to None so existing callers stay valid"
+def test_mlx_vlm_reemits_think_prefill_inside_adapter_context(monkeypatch):
+ """A prefilled block must be re-emitted as the first VLM snapshot,
+ inside the adapter context (so unsupported requests still raise first), so
+ the UI renders the thinking block during prefill and a pre-first-token
+ cancel does not drop it. Mirrors _generate_text."""
+ from core.inference import mlx_inference
+
+ MLXInferenceBackend = mlx_inference.MLXInferenceBackend
+
+ order = []
+
+ @contextmanager
+ def _adapter_state(_model, state):
+ assert backend._generation_lock.locked()
+ order.append("adapter_enter")
+ try:
+ yield
+ finally:
+ order.append("adapter_exit")
+
+ monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state)
+ monkeypatch.setattr(
+ "core.inference.chat_template_helpers.detect_think_prefill",
+ lambda *_a, **_k: "\n",
+ )
+
+ prompt_utils = SimpleNamespace(
+ MODEL_CONFIG = {"deepseek_vl_v2": object()},
+ apply_chat_template = lambda *_a, **_k: " model-aware",
+ )
+ mlx_vlm = types.ModuleType("mlx_vlm")
+ mlx_vlm.prompt_utils = prompt_utils
+
+ def _vlm_stream(*_a, **_k):
+ # The prefill must have been emitted before any generated token.
+ assert order[-1] == "adapter_enter"
+ yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)
+
+ mlx_vlm.stream_generate = _vlm_stream
+ monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm)
+ monkeypatch.setattr(
+ "core.inference.chat_template_helpers.apply_chat_template_for_generation",
+ lambda _t, _m, **_k: " model-aware",
+ )
+
+ backend = MLXInferenceBackend()
+ backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"})
+ backend._processor = SimpleNamespace(tokenizer = SimpleNamespace())
+ args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None)
+
+ gen = backend._generate_vlm(*args, _adapter_state = False)
+ # First snapshot is the prefill alone, emitted after entering the adapter context.
+ assert next(gen) == "\n"
+ assert order == ["adapter_enter"]
+ # Subsequent snapshots are cumulative (prefill + generated text).
+ assert next(gen) == "\nok"
+ gen.close()
+ assert order == ["adapter_enter", "adapter_exit"]
+
+
def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch):
- from core.inference.mlx_inference import MLXInferenceBackend
+ from core.inference import mlx_inference
+
+ MLXInferenceBackend = mlx_inference.MLXInferenceBackend
calls = {"generic": [], "model": [], "stream": []}
+ adapter_events = []
+ adapter_active = {"value": False}
+
+ @contextmanager
+ def _adapter_state(_model, state):
+ assert backend._generation_lock.locked()
+ adapter_events.append(("enter", state))
+ adapter_active["value"] = True
+ try:
+ yield
+ finally:
+ adapter_active["value"] = False
+ adapter_events.append(("exit", state))
+
+ monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state)
state = {"generic": "serialized", "model": " model-aware"}
prompt_utils = SimpleNamespace(
MODEL_CONFIG = {"deepseek_vl_v2": object()},
@@ -346,10 +521,13 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch):
)
mlx_vlm = types.ModuleType("mlx_vlm")
mlx_vlm.prompt_utils = prompt_utils
- mlx_vlm.stream_generate = lambda *_args, **kwargs: (
- calls["stream"].append((_args, kwargs))
- or iter([SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)])
- )
+
+ def _vlm_stream(*args, **kwargs):
+ assert adapter_active["value"]
+ calls["stream"].append((args, kwargs))
+ yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)
+
+ mlx_vlm.stream_generate = _vlm_stream
monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm)
def generic(_target, _messages, **kwargs):
@@ -369,7 +547,11 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch):
backend._processor = SimpleNamespace(tokenizer = SimpleNamespace())
args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None)
tools = [{"function": {"name": "search"}}]
- assert list(backend._generate_vlm(*args)) == ["ok"]
+ generator = backend._generate_vlm(*args, _adapter_state = False)
+ assert next(generator) == "ok"
+ assert adapter_active["value"] and backend._generation_lock.locked()
+ generator.close()
+ assert adapter_events == [("enter", False), ("exit", False)]
assert calls["model"][0]["num_images"] == 1
assert calls["stream"][0][0][2] == " model-aware"
with pytest.raises(RuntimeError, match = "dropping requested tools"):
@@ -449,7 +631,10 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
"""Mac text path must route through apply_chat_template_for_generation so
reasoning / tool kwargs reach the tokenizer."""
_install_fake_mlx(monkeypatch)
- from core.inference.mlx_inference import MLXInferenceBackend
+ from core.inference import mlx_inference
+
+ MLXInferenceBackend = mlx_inference.MLXInferenceBackend
+ real_adapter_state = mlx_inference._temporary_mlx_adapter_state
# The text path renders once with tools, then the native-template fallback makes a second no-
# tools probe call (tools=None) to detect whether the template dropped the schema.
@@ -474,11 +659,31 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
mlx_lm_sample.make_sampler = lambda **_kw: object()
mlx_lm_sample.make_logits_processors = lambda **_kw: None
+ adapter_events = []
+ adapter_active = {"value": False}
+ stream_state = {"fail": False}
+
+ @contextmanager
+ def _adapter_state(_model, state):
+ assert backend._generation_lock.locked()
+ adapter_events.append(("enter", state))
+ adapter_active["value"] = True
+ try:
+ yield
+ finally:
+ adapter_active["value"] = False
+ adapter_events.append(("exit", state))
+
+ monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state)
+
class _Resp:
def __init__(self, tok):
self.token = tok
def _stream_generate(_model, _tokenizer, **_kw):
+ assert adapter_active["value"]
+ if stream_state["fail"]:
+ raise RuntimeError("generation failed")
yield _Resp(1)
mlx_lm_pkg.stream_generate = _stream_generate
@@ -500,17 +705,45 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
backend._tokenizer = _Tok()
backend._is_vlm = False
- out = list(
- backend.generate_chat_response(
- messages = [{"role": "user", "content": "ping"}],
- tools = [{"function": {"name": "web_search"}}],
- enable_thinking = True,
- reasoning_effort = "medium",
- preserve_thinking = True,
- max_new_tokens = 1,
- )
+ generator = backend.generate_with_adapter_control(
+ use_adapter = False,
+ messages = [{"role": "user", "content": "ping"}],
+ tools = [{"function": {"name": "web_search"}}],
+ enable_thinking = True,
+ reasoning_effort = "medium",
+ preserve_thinking = True,
+ max_new_tokens = 1,
)
- assert out == ["hi"]
+ assert next(generator) == "hi"
+ assert adapter_active["value"] and backend._generation_lock.locked()
+ generator.close()
+ assert adapter_events == [("enter", False), ("exit", False)]
+ stream_state["fail"] = True
+ with pytest.raises(RuntimeError, match = "generation failed"):
+ list(
+ backend.generate_with_adapter_control(
+ use_adapter = False,
+ messages = [{"role": "user", "content": "ping"}],
+ max_new_tokens = 1,
+ )
+ )
+ assert adapter_events[-2:] == [("enter", False), ("exit", False)]
+ assert not backend._generation_lock.locked()
+
+ monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", real_adapter_state)
+ monkeypatch.setattr(
+ "core.inference.chat_template_helpers.detect_think_prefill",
+ lambda *_args, **_kwargs: "",
+ )
+ stream_state["fail"] = False
+ named = backend.generate_with_adapter_control(
+ use_adapter = "named",
+ messages = [{"role": "user", "content": "ping"}],
+ max_new_tokens = 1,
+ )
+ with pytest.raises(NotImplementedError, match = "named adapter"):
+ next(named)
+ assert not adapter_active["value"] and not backend._generation_lock.locked()
# The toggled kwargs must reach the chat-template helper on the real render
# (one of the calls carries the tools; the fallback probe passes tools=None).
tool_renders = [
diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py
index fb80b6d061..3a36500aee 100644
--- a/studio/backend/tests/test_orchestrator_unload_cancel.py
+++ b/studio/backend/tests/test_orchestrator_unload_cancel.py
@@ -34,6 +34,70 @@ def _bare_orchestrator():
return o
+def test_adapter_control_raises_stream_errors(monkeypatch):
+ o = _bare_orchestrator()
+ monkeypatch.setattr(
+ o,
+ "_generate_dispatched",
+ lambda **_kwargs: iter([orch_mod.GenStreamError("Error: adapter failed")]),
+ )
+
+ with pytest.raises(RuntimeError, match = "adapter failed"):
+ list(o.generate_with_adapter_control(use_adapter = False))
+
+ closed = []
+
+ def _stream(**_kwargs):
+ try:
+ yield "token"
+ yield "late token"
+ finally:
+ closed.append(True)
+
+ monkeypatch.setattr(o, "_generate_dispatched", _stream)
+ generator = o.generate_with_adapter_control(use_adapter = False)
+ assert next(generator) == "token"
+ generator.close()
+ assert closed == [True]
+
+
+def test_worker_closes_cancelled_generator_before_gen_done():
+ from core.inference.worker import _handle_generate
+
+ events = []
+
+ class _Backend:
+ last_generation_stats = None
+
+ def generate_with_adapter_control(self, **_kwargs):
+ try:
+ yield "token"
+ yield "late token"
+ finally:
+ events.append("closed")
+
+ class _Responses:
+ def __init__(self):
+ self.items = []
+
+ def put(self, item):
+ if item["type"] == "gen_done":
+ assert events == ["closed"]
+ self.items.append(item)
+
+ responses = _Responses()
+ cancel = threading.Event()
+ cancel.set()
+ _handle_generate(
+ _Backend(),
+ {"request_id": "r1", "messages": [], "use_adapter": False},
+ responses,
+ cancel,
+ )
+
+ assert [item["type"] for item in responses.items] == ["gen_done"]
+
+
def test_unload_cancels_inflight_generation_then_unloads(monkeypatch):
o = _bare_orchestrator()
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
From 030524ae8edd056fc37fc03190362d6fc59a0392 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Sun, 19 Jul 2026 00:34:13 -0700
Subject: [PATCH 009/255] security: refresh the fastapi C2-loop baseline entry
for the current release (#7223)
The pip scan-packages studio shard is red on main and on every open PR:
the baselined fastapi finding (the benign SSE keepalive `while True:`
loop in fastapi/routing.py, reviewed and suppressed long ago) records
its evidence at L586 with the span digest of the fastapi release current
at baseline time. The latest fastapi shifts that loop to L587 and its
span digest with it, so the evidence hash no longer matches and the
scanner reports the finding as new, failing the shard with one
unsuppressed CRITICAL.
Re-reviewed the flagged code in the current release before refreshing:
L587 is the same keepalive loop inside the streaming response machinery,
not a beacon. Only the one entry's evidence and evidence_hash change.
Verified with the scanner itself: `scan_packages.py fastapi
--no-baseline` reproduces the exact CI evidence string, and with the
updated baseline the same scan exits 0 with the finding suppressed as
1 CRITICAL baselined.
---
scripts/scan_packages_baseline.json | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json
index 929cc37bda..1f7bc8dcc0 100644
--- a/scripts/scan_packages_baseline.json
+++ b/scripts/scan_packages_baseline.json
@@ -95,8 +95,8 @@
"file": "fastapi/routing.py",
"check": "C2 polling/beaconing loop detected",
"severity": "CRITICAL",
- "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45",
- "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5"
+ "evidence": "L587: while True: sha256:06c2c7f15d73bf192e5e3272c5ff5fcaeff7f6774fef5f4eca6ef473ae50e2b3",
+ "evidence_hash": "57acd497f404c203e4450d0580ad85aa8a33406e8d64ad06fbac6cf47d97b24d"
},
{
"package": "fastmcp-slim",
From e9ef2ac60f35d8e980460763e09027e477064a38 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Sun, 19 Jul 2026 13:04:56 +0530
Subject: [PATCH 010/255] Studio: enforce 60s minimum on idle auto-unload TTL
(0 stays off) (#7185)
* Studio: enforce 60s minimum on idle auto-unload TTL (0 stays off)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Drop decorative section separator from idle TTL floor tests
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
Co-authored-by: danielhanchen
---
.../backend/tests/test_openai_auto_switch.py | 54 ++++++++++++++++++-
.../utils/openai_auto_switch_settings.py | 49 ++++++++++++++---
.../components/model-auto-switch-section.tsx | 10 +++-
studio/frontend/src/i18n/locales/ar.ts | 4 +-
studio/frontend/src/i18n/locales/de.ts | 4 +-
studio/frontend/src/i18n/locales/en.ts | 4 +-
studio/frontend/src/i18n/locales/es.ts | 4 +-
studio/frontend/src/i18n/locales/fr.ts | 4 +-
studio/frontend/src/i18n/locales/hi.ts | 4 +-
studio/frontend/src/i18n/locales/ja.ts | 4 +-
studio/frontend/src/i18n/locales/ko.ts | 4 +-
studio/frontend/src/i18n/locales/pt-br.ts | 4 +-
studio/frontend/src/i18n/locales/ru.ts | 4 +-
studio/frontend/src/i18n/locales/zh-CN.ts | 4 +-
14 files changed, 125 insertions(+), 32 deletions(-)
diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py
index d02a2a4f7e..8742b84ae7 100644
--- a/studio/backend/tests/test_openai_auto_switch.py
+++ b/studio/backend/tests/test_openai_auto_switch.py
@@ -1609,11 +1609,11 @@ def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch):
def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch):
# An explicit stored value wins over the env default and remains gated on the
# auto-switch toggle.
- store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 30}
+ store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 90}
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600")
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True)
- assert settings.get_auto_unload_idle_seconds() == 30 # stored wins, not env
+ assert settings.get_auto_unload_idle_seconds() == 90 # stored wins, not env
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False)
assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off
@@ -3094,3 +3094,53 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp
monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct"
)
assert "Model auto-switch" in non_gguf_loaded
+
+
+def test_setter_rejects_idle_below_floor(monkeypatch):
+ import storage.studio_db as db
+
+ writes = []
+ monkeypatch.setattr(db, "upsert_app_settings", lambda m: writes.append(dict(m)))
+ settings._cache.clear()
+
+ with pytest.raises(ValueError, match = "at least 60"):
+ settings.set_openai_auto_switch(True, 30)
+ assert writes == [] # rejected before any persist
+ # 0 (off) and >= 60 pass through unchanged.
+ assert settings.set_openai_auto_switch(True, 0)[1] == 0
+ assert settings.set_openai_auto_switch(True, 60)[1] == 60
+ assert settings.set_openai_auto_switch(True, 3600)[1] == 3600
+
+
+def test_put_route_rejects_idle_below_floor():
+ import routes.settings as settings_route
+ from fastapi import HTTPException
+
+ payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 30)
+ with pytest.raises(HTTPException) as excinfo:
+ settings_route.update_openai_auto_switch(payload, "tester")
+ assert excinfo.value.status_code == 400
+
+
+def test_stored_legacy_idle_below_floor_is_clamped(monkeypatch):
+ # Values persisted before the floor existed are raised to it on read, for
+ # both the effective TTL and the value the settings UI displays.
+ store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 5}
+ monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
+ monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True)
+ assert settings.get_auto_unload_idle_seconds() == 60
+ assert settings.get_stored_auto_unload_idle_seconds() == 60
+ store[settings.AUTO_UNLOAD_IDLE_SETTING_KEY] = 90
+ assert settings.get_auto_unload_idle_seconds() == 90
+
+
+def test_env_idle_below_floor_is_clamped(monkeypatch):
+ monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d)
+ monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "5")
+ assert settings.get_auto_unload_idle_seconds() == 60
+ monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "0")
+ assert settings.get_auto_unload_idle_seconds() == 0
+ monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600")
+ assert settings.get_auto_unload_idle_seconds() == 600
+ monkeypatch.delenv(settings.MODEL_IDLE_TTL_ENV_VAR)
+ assert settings.get_auto_unload_idle_seconds() == 0
diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py
index 1689395f40..462435e5d5 100644
--- a/studio/backend/utils/openai_auto_switch_settings.py
+++ b/studio/backend/utils/openai_auto_switch_settings.py
@@ -8,7 +8,9 @@ Two settings, both off by default so existing API behavior is unchanged:
names a downloaded local GGUF different from the loaded one transparently
loads it before serving (llama-swap-style). Unknown names pass through.
- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is
- unloaded after this many idle seconds to free VRAM.
+ unloaded after this many idle seconds to free VRAM. Enabled values have a
+ 60s floor (0 stays "off"): a tiny TTL tears the model down between turns of
+ an active chat, forcing a full weight reload + prompt re-prefill per turn.
The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env
var. Unlike the stored setting (which stays gated on auto-switch), the env value
@@ -33,6 +35,7 @@ MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
+MIN_AUTO_UNLOAD_IDLE_SECONDS = 60
_CACHE_TTL_S = 2.0
_cache_lock = threading.Lock()
@@ -58,6 +61,10 @@ def _coerce_int(value: Any) -> int | None:
return None
+def _apply_idle_floor(seconds: int) -> int:
+ return 0 if seconds <= 0 else max(MIN_AUTO_UNLOAD_IDLE_SECONDS, seconds)
+
+
def _cached_setting(key: str, default: Any) -> Any:
"""Read an app setting, memoized for _CACHE_TTL_S to spare the hot path."""
now = time.monotonic()
@@ -91,12 +98,34 @@ def _stored_idle_seconds() -> Optional[int]:
return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None))
+_env_floor_warned = False
+
+
def _env_idle_seconds() -> Optional[int]:
- """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid."""
+ """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid.
+
+ Floored to MIN_AUTO_UNLOAD_IDLE_SECONDS here (with a one-time warning) since
+ headless/container deploys have no UI to surface a validation error."""
raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR)
if raw is None or not raw.strip():
return None
- return _coerce_int(raw)
+ parsed = _coerce_int(raw)
+ if parsed is None:
+ return None
+ floored = _apply_idle_floor(parsed)
+ if floored != parsed:
+ global _env_floor_warned
+ if not _env_floor_warned:
+ _env_floor_warned = True
+ from loggers import get_logger
+ get_logger(__name__).warning(
+ "%s=%s is below the %ss minimum; using %ss",
+ MODEL_IDLE_TTL_ENV_VAR,
+ parsed,
+ MIN_AUTO_UNLOAD_IDLE_SECONDS,
+ floored,
+ )
+ return floored
def get_stored_auto_unload_idle_seconds() -> int:
@@ -108,7 +137,9 @@ def get_stored_auto_unload_idle_seconds() -> int:
"""
stored = _stored_idle_seconds()
if stored is not None:
- return stored
+ # Floor legacy values persisted before the minimum existed, so the UI
+ # displays the effective TTL and round-trips it cleanly.
+ return _apply_idle_floor(stored)
env = _env_idle_seconds()
return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS
@@ -118,8 +149,9 @@ def get_auto_unload_idle_seconds() -> int:
stored = _stored_idle_seconds()
if stored is not None:
# An explicit UI/API value stays gated on auto-switch: off reports 0 so the
- # off state is identical to pre-feature.
- return stored if get_openai_auto_switch_enabled() else 0
+ # off state is identical to pre-feature. Floored to cover values persisted
+ # before the minimum existed.
+ return _apply_idle_floor(stored) if get_openai_auto_switch_enabled() else 0
# No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that
# enables idle-unload even with auto-switch off (headless/container deploys).
env = _env_idle_seconds()
@@ -136,6 +168,11 @@ def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]:
parsed_idle = _coerce_int(idle_seconds)
if parsed_idle is None:
raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
+ if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS:
+ raise ValueError(
+ f"Auto-unload idle seconds must be 0 (off) or at least "
+ f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}."
+ )
from storage.studio_db import upsert_app_settings
upsert_app_settings(
diff --git a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
index 5bebefa84c..32b3e53a2c 100644
--- a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
+++ b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
@@ -14,6 +14,9 @@ import {
import { SettingsRow } from "./settings-row";
import { SettingsSection } from "./settings-section";
+// Mirrors MIN_AUTO_UNLOAD_IDLE_SECONDS in the backend settings store.
+const MIN_IDLE_SECONDS = 60;
+
export function ModelAutoSwitchSection() {
const t = useT();
const [settings, setSettings] = useState(
@@ -45,13 +48,16 @@ export function ModelAutoSwitchSection() {
};
}, [t]);
- // Parse the idle-seconds draft to a non-negative integer; empty/invalid -> null.
+ // Parse the idle-seconds draft: 0 (off) or >= MIN_IDLE_SECONDS; else null.
const parseIdleSeconds = (): number | null => {
if (!draftIdleSeconds.trim()) {
return null;
}
const parsed = Number(draftIdleSeconds);
- return Number.isInteger(parsed) && parsed >= 0 ? parsed : null;
+ if (!Number.isInteger(parsed)) {
+ return null;
+ }
+ return parsed === 0 || parsed >= MIN_IDLE_SECONDS ? parsed : null;
};
const persist = async (
diff --git a/studio/frontend/src/i18n/locales/ar.ts b/studio/frontend/src/i18n/locales/ar.ts
index 28f404a384..744a2002c9 100644
--- a/studio/frontend/src/i18n/locales/ar.ts
+++ b/studio/frontend/src/i18n/locales/ar.ts
@@ -155,14 +155,14 @@ export const ar = {
"عندما يسمّي طلب متوافق مع OpenAI ملف GGUF مُنزّلاً مختلفًا، يتم تحميله قبل الخدمة. مُعطّل افتراضيًا؛ الأسماء غير المعروفة تُبقي على النموذج المُحمَّل.",
idleUnload: "الإلغاء التلقائي عند الخمول",
idleUnloadDescription:
- "إلغاء تحميل النموذج بعد هذا العدد من ثواني الخمول لتحرير الـ VRAM؛ الطلب التالي يعيد تحميله. القيمة 0 تُبقيه محمَّلاً.",
+ "إلغاء تحميل النموذج بعد هذا العدد من ثواني الخمول لتحرير الـ VRAM؛ الطلب التالي يعيد تحميله. القيمة 0 تُبقيه محمَّلاً. الحد الأدنى 60 ثانية.",
idleNeedsEnable:
"فعّل تبديل النموذج حسب الطلب حتى يعاد تحميل النموذج غير المحمَّل عند الاستخدام التالي.",
idleActiveViaEnv:
"الإلغاء التلقائي عند الخمول مُفعَّل عبر متغير البيئة UNSLOTH_MODEL_IDLE_TTL.",
loadError: "فشل تحميل إعدادات التبديل التلقائي للنموذج.",
saveError: "فشل حفظ إعدادات التبديل التلقائي للنموذج.",
- idleError: "أدخل عددًا صحيحًا من الثواني (0 أو أكثر).",
+ idleError: "أدخل 0 لإبقاء النموذج محمَّلاً، أو 60 ثانية على الأقل.",
},
previewSharing: {
sectionTitle: "مشاركة المعاينة",
diff --git a/studio/frontend/src/i18n/locales/de.ts b/studio/frontend/src/i18n/locales/de.ts
index e38cdbfa0e..7d94e7656e 100644
--- a/studio/frontend/src/i18n/locales/de.ts
+++ b/studio/frontend/src/i18n/locales/de.ts
@@ -158,7 +158,7 @@ export const de = {
"Wenn eine OpenAI-kompatible Anfrage ein anderes heruntergeladenes GGUF nennt, wird dieses vor der Auslieferung geladen. Standardmäßig aus; unbekannte Namen liefern weiterhin das geladene Modell aus.",
idleUnload: "Automatisches Entladen bei Inaktivität",
idleUnloadDescription:
- "Entlädt das Modell nach dieser Anzahl inaktiver Sekunden, um VRAM freizugeben; die nächste Anfrage lädt es erneut. 0 hält es geladen.",
+ "Entlädt das Modell nach dieser Anzahl inaktiver Sekunden, um VRAM freizugeben; die nächste Anfrage lädt es erneut. 0 hält es geladen. Minimum 60 Sekunden.",
idleNeedsEnable:
"Aktivieren Sie \"Modell je Anfrage wechseln\", damit ein entladenes Modell bei der nächsten Nutzung erneut geladen wird.",
idleActiveViaEnv:
@@ -167,7 +167,7 @@ export const de = {
"Einstellungen für automatischen Modellwechsel konnten nicht geladen werden.",
saveError:
"Einstellungen für automatischen Modellwechsel konnten nicht gespeichert werden.",
- idleError: "Geben Sie eine ganze Anzahl an Sekunden ein (0 oder mehr).",
+ idleError: "Geben Sie 0 ein, um das Modell geladen zu halten, oder mindestens 60 Sekunden.",
},
previewSharing: {
sectionTitle: "Vorschau-Freigabe",
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index fe3a6f8542..de8ac17c29 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -224,14 +224,14 @@ export const en = {
"When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.",
idleUnload: "Idle auto-unload",
idleUnloadDescription:
- "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded.",
+ "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded. Minimum 60 seconds.",
idleNeedsEnable:
"Turn on Switch model by request so an unloaded model reloads on next use.",
idleActiveViaEnv:
"Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.",
loadError: "Failed to load model auto-switch settings.",
saveError: "Failed to save model auto-switch settings.",
- idleError: "Enter a whole number of seconds (0 or more).",
+ idleError: "Enter 0 to keep the model loaded, or at least 60 seconds.",
},
previewSharing: {
sectionTitle: "Preview sharing",
diff --git a/studio/frontend/src/i18n/locales/es.ts b/studio/frontend/src/i18n/locales/es.ts
index e5f9650bef..988c109a3f 100644
--- a/studio/frontend/src/i18n/locales/es.ts
+++ b/studio/frontend/src/i18n/locales/es.ts
@@ -157,7 +157,7 @@ export const es = {
"Cuando una solicitud compatible con OpenAI nombra un GGUF descargado distinto, se carga antes de responder. Desactivado por defecto; los nombres desconocidos siguen usando el modelo cargado.",
idleUnload: "Descarga automática por inactividad",
idleUnloadDescription:
- "Descarga el modelo tras este número de segundos inactivo para liberar VRAM; la siguiente solicitud lo recarga. 0 lo mantiene cargado.",
+ "Descarga el modelo tras este número de segundos inactivo para liberar VRAM; la siguiente solicitud lo recarga. 0 lo mantiene cargado. Mínimo 60 segundos.",
idleNeedsEnable:
"Activa Cambiar de modelo según la solicitud para que un modelo descargado se recargue en el próximo uso.",
idleActiveViaEnv:
@@ -166,7 +166,7 @@ export const es = {
"No se pudo cargar la configuración de cambio automático de modelo.",
saveError:
"No se pudo guardar la configuración de cambio automático de modelo.",
- idleError: "Introduce un número entero de segundos (0 o más).",
+ idleError: "Introduce 0 para mantener el modelo cargado, o al menos 60 segundos.",
},
previewSharing: {
sectionTitle: "Compartir vista previa",
diff --git a/studio/frontend/src/i18n/locales/fr.ts b/studio/frontend/src/i18n/locales/fr.ts
index 190284175d..e1f2a0c5ec 100644
--- a/studio/frontend/src/i18n/locales/fr.ts
+++ b/studio/frontend/src/i18n/locales/fr.ts
@@ -157,7 +157,7 @@ export const fr = {
"Lorsqu'une requête compatible OpenAI nomme un autre GGUF téléchargé, le charger avant de répondre. Désactivé par défaut ; les noms inconnus continuent de servir le modèle chargé.",
idleUnload: "Déchargement automatique en cas d'inactivité",
idleUnloadDescription:
- "Décharger le modèle après ce nombre de secondes d'inactivité pour libérer la VRAM ; la requête suivante le recharge. 0 le maintient chargé.",
+ "Décharger le modèle après ce nombre de secondes d'inactivité pour libérer la VRAM ; la requête suivante le recharge. 0 le maintient chargé. Minimum 60 secondes.",
idleNeedsEnable:
"Activez Changer de modèle par requête pour qu'un modèle déchargé se recharge à la prochaine utilisation.",
idleActiveViaEnv:
@@ -166,7 +166,7 @@ export const fr = {
"Échec du chargement des paramètres de changement automatique de modèle.",
saveError:
"Échec de l'enregistrement des paramètres de changement automatique de modèle.",
- idleError: "Saisissez un nombre entier de secondes (0 ou plus).",
+ idleError: "Saisissez 0 pour garder le modèle chargé, ou au moins 60 secondes.",
},
previewSharing: {
sectionTitle: "Partage de l'aperçu",
diff --git a/studio/frontend/src/i18n/locales/hi.ts b/studio/frontend/src/i18n/locales/hi.ts
index 97f55251c5..77b6265e7b 100644
--- a/studio/frontend/src/i18n/locales/hi.ts
+++ b/studio/frontend/src/i18n/locales/hi.ts
@@ -154,14 +154,14 @@ export const hi = {
"जब कोई OpenAI-संगत अनुरोध किसी अन्य डाउनलोड किए गए GGUF का नाम लेता है, तो सर्व करने से पहले उसे लोड करें। डिफ़ॉल्ट रूप से बंद; अज्ञात नाम लोड किए गए मॉडल को सर्व करते रहते हैं।",
idleUnload: "निष्क्रिय ऑटो-अनलोड",
idleUnloadDescription:
- "VRAM मुक्त करने के लिए इतने निष्क्रिय सेकंड के बाद मॉडल को अनलोड करें; अगला अनुरोध इसे फिर से लोड करता है। 0 इसे लोड रखता है।",
+ "VRAM मुक्त करने के लिए इतने निष्क्रिय सेकंड के बाद मॉडल को अनलोड करें; अगला अनुरोध इसे फिर से लोड करता है। 0 इसे लोड रखता है। न्यूनतम 60 सेकंड।",
idleNeedsEnable:
"अनुरोध के अनुसार मॉडल बदलें चालू करें ताकि अनलोड किया गया मॉडल अगले उपयोग पर फिर से लोड हो।",
idleActiveViaEnv:
"निष्क्रिय ऑटो-अनलोड UNSLOTH_MODEL_IDLE_TTL एनवायरनमेंट वेरिएबल के माध्यम से सक्रिय है।",
loadError: "मॉडल ऑटो-स्विच सेटिंग्स लोड करने में विफल।",
saveError: "मॉडल ऑटो-स्विच सेटिंग्स सहेजने में विफल।",
- idleError: "सेकंड की पूरी संख्या दर्ज करें (0 या अधिक)।",
+ idleError: "मॉडल को लोड रखने के लिए 0 दर्ज करें, या कम से कम 60 सेकंड।",
},
previewSharing: {
sectionTitle: "पूर्वावलोकन साझाकरण",
diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts
index 23752b6c26..a261994f03 100644
--- a/studio/frontend/src/i18n/locales/ja.ts
+++ b/studio/frontend/src/i18n/locales/ja.ts
@@ -149,12 +149,12 @@ export const ja = {
enable: "リクエストごとにモデルを切り替え",
enableDescription: "OpenAI互換のリクエストが別のダウンロード済み GGUF を指定した場合、応答する前にそのモデルを読み込みます。デフォルトはオフです。不明な名前の場合は、読み込み済みのモデルで応答を続けます。",
idleUnload: "アイドル時の自動アンロード",
- idleUnloadDescription: "指定した秒数だけアイドル状態が続くとモデルをアンロードして VRAM を解放します。次のリクエストで再読み込みされます。0 にすると読み込んだままにします。",
+ idleUnloadDescription: "指定した秒数だけアイドル状態が続くとモデルをアンロードして VRAM を解放します。次のリクエストで再読み込みされます。0 にすると読み込んだままにします。最小 60 秒。",
idleNeedsEnable: "アンロードされたモデルが次回使用時に再読み込みされるように、「リクエストごとにモデルを切り替え」をオンにしてください。",
idleActiveViaEnv: "アイドル時の自動アンロードは UNSLOTH_MODEL_IDLE_TTL 環境変数によって有効になっています。",
loadError: "モデル自動切り替え設定の読み込みに失敗しました。",
saveError: "モデル自動切り替え設定の保存に失敗しました。",
- idleError: "秒数を整数(0 以上)で入力してください。",
+ idleError: "モデルを読み込んだままにするには 0 を、それ以外は 60 秒以上を入力してください。",
},
previewSharing: {
sectionTitle: "プレビュー共有",
diff --git a/studio/frontend/src/i18n/locales/ko.ts b/studio/frontend/src/i18n/locales/ko.ts
index a11a94faf2..7c5691925e 100644
--- a/studio/frontend/src/i18n/locales/ko.ts
+++ b/studio/frontend/src/i18n/locales/ko.ts
@@ -153,14 +153,14 @@ export const ko = {
"OpenAI 호환 요청이 다운로드된 다른 GGUF를 지정하면, 응답하기 전에 해당 모델을 불러옵니다. 기본값은 꺼짐이며, 알 수 없는 이름은 불러온 모델을 계속 제공합니다.",
idleUnload: "유휴 시 자동 해제",
idleUnloadDescription:
- "지정한 유휴 시간(초)이 지나면 모델을 해제하여 VRAM을 확보합니다. 다음 요청 시 다시 불러옵니다. 0으로 설정하면 계속 로드된 상태로 유지됩니다.",
+ "지정한 유휴 시간(초)이 지나면 모델을 해제하여 VRAM을 확보합니다. 다음 요청 시 다시 불러옵니다. 0으로 설정하면 계속 로드된 상태로 유지됩니다. 최소 60초입니다.",
idleNeedsEnable:
"해제된 모델이 다음 사용 시 다시 로드되도록 하려면 요청에 따라 모델 전환을 켜세요.",
idleActiveViaEnv:
"유휴 시 자동 해제가 UNSLOTH_MODEL_IDLE_TTL 환경 변수를 통해 활성화되어 있습니다.",
loadError: "모델 자동 전환 설정을 불러오지 못했습니다.",
saveError: "모델 자동 전환 설정을 저장하지 못했습니다.",
- idleError: "정수(초)를 입력하세요(0 이상).",
+ idleError: "모델을 로드 상태로 유지하려면 0을, 그렇지 않으면 60초 이상을 입력하세요.",
},
previewSharing: {
sectionTitle: "미리보기 공유",
diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts
index 07832ef1d0..e6d2347c10 100644
--- a/studio/frontend/src/i18n/locales/pt-br.ts
+++ b/studio/frontend/src/i18n/locales/pt-br.ts
@@ -157,14 +157,14 @@ export const ptBR = {
"Quando uma requisição compatível com OpenAI nomear um GGUF baixado diferente, carrega-o antes de responder. Desativado por padrão; nomes desconhecidos continuam usando o modelo carregado.",
idleUnload: "Descarregamento automático por inatividade",
idleUnloadDescription:
- "Descarrega o modelo após esta quantidade de segundos de inatividade para liberar VRAM; a próxima requisição o recarrega. 0 mantém o modelo carregado.",
+ "Descarrega o modelo após esta quantidade de segundos de inatividade para liberar VRAM; a próxima requisição o recarrega. 0 mantém o modelo carregado. Mínimo de 60 segundos.",
idleNeedsEnable:
"Ative Trocar de modelo por requisição para que um modelo descarregado seja recarregado no próximo uso.",
idleActiveViaEnv:
"O descarregamento automático por inatividade está ativo por meio da variável de ambiente UNSLOTH_MODEL_IDLE_TTL.",
loadError: "Falha ao carregar as configurações de troca automática de modelo.",
saveError: "Falha ao salvar as configurações de troca automática de modelo.",
- idleError: "Insira um número inteiro de segundos (0 ou mais).",
+ idleError: "Insira 0 para manter o modelo carregado, ou pelo menos 60 segundos.",
},
previewSharing: {
sectionTitle: "Compartilhamento de pré-visualização",
diff --git a/studio/frontend/src/i18n/locales/ru.ts b/studio/frontend/src/i18n/locales/ru.ts
index 81d20cc2ea..c7464a3b44 100644
--- a/studio/frontend/src/i18n/locales/ru.ts
+++ b/studio/frontend/src/i18n/locales/ru.ts
@@ -154,14 +154,14 @@ export const ru = {
"Когда OpenAI-совместимый запрос указывает другую загруженную GGUF, загружать её перед обслуживанием. По умолчанию выключено; неизвестные имена продолжают обслуживать загруженную модель.",
idleUnload: "Автовыгрузка при простое",
idleUnloadDescription:
- "Выгружать модель после указанного числа секунд простоя, чтобы освободить VRAM; следующий запрос загрузит её снова. 0 оставляет модель загруженной.",
+ "Выгружать модель после указанного числа секунд простоя, чтобы освободить VRAM; следующий запрос загрузит её снова. 0 оставляет модель загруженной. Минимум 60 секунд.",
idleNeedsEnable:
"Включите «Переключать модель по запросу», чтобы выгруженная модель загружалась при следующем использовании.",
idleActiveViaEnv:
"Автовыгрузка при простое активна через переменную окружения UNSLOTH_MODEL_IDLE_TTL.",
loadError: "Не удалось загрузить настройки автопереключения модели.",
saveError: "Не удалось сохранить настройки автопереключения модели.",
- idleError: "Введите целое число секунд (0 или больше).",
+ idleError: "Введите 0, чтобы модель оставалась загруженной, или не менее 60 секунд.",
},
previewSharing: {
sectionTitle: "Публикация предпросмотра",
diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts
index 4c51755244..ff218adad2 100644
--- a/studio/frontend/src/i18n/locales/zh-CN.ts
+++ b/studio/frontend/src/i18n/locales/zh-CN.ts
@@ -152,14 +152,14 @@ export const zhCN = {
"当兼容 OpenAI 的请求指定了另一个已下载的 GGUF 时,先加载它再提供服务。默认关闭;未知名称将继续使用已加载的模型。",
idleUnload: "空闲自动卸载",
idleUnloadDescription:
- "空闲达到该秒数后卸载模型以释放 VRAM;下次请求会重新加载。设为 0 则保持加载。",
+ "空闲达到该秒数后卸载模型以释放 VRAM;下次请求会重新加载。设为 0 则保持加载。最小 60 秒。",
idleNeedsEnable:
"开启“按请求切换模型”,以便已卸载的模型在下次使用时重新加载。",
idleActiveViaEnv:
"空闲自动卸载已通过 UNSLOTH_MODEL_IDLE_TTL 环境变量启用。",
loadError: "加载模型自动切换设置失败。",
saveError: "保存模型自动切换设置失败。",
- idleError: "请输入整数秒数(0 或以上)。",
+ idleError: "输入 0 保持模型加载,或输入至少 60 秒。",
},
previewSharing: {
sectionTitle: "预览分享",
From 6d8c18cd1a630a231cb282799be13f8bb0b5ff3b Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Sun, 19 Jul 2026 00:47:04 -0700
Subject: [PATCH 011/255] Replace standalone Studio wording with Unsloth
(#7221)
* Replace standalone Studio wording with Unsloth
Replace the single word Studio with Unsloth wherever it is used as
shorthand for Unsloth Studio in docs, CLI output, UI strings, i18n
locales, workflow display names, comments and docstrings.
Kept unchanged: the full name Unsloth Studio, third party product
names (LM Studio, Visual Studio, Mac Studio), feature names
(Recipe Studio, Fine-tuning Studio and its translations), and all
identifiers such as env vars, commands, paths and filenames.
* Address review feedback on the Studio wording rename
Use "an" before Unsloth where the rename left the article as "a".
Restore the split brand where Unsloth and Studio render as two halves
of the full product name: the onboarding sidebar subtitle and the
IPv6 localhost warning. Scope two messages to the full name Unsloth
Studio where plain Unsloth was misleading: the AMD README bullet and
the CLI studio setup error.
---
.gitattributes | 2 +-
.github/scripts/agent-guides-drive.sh | 4 +-
.github/scripts/assert-llama-loads.sh | 2 +-
.github/scripts/assert-prompt-cache.sh | 2 +-
.github/scripts/hf-download-with-retry.sh | 4 +-
.github/workflows/lint-ci.yml | 4 +-
.github/workflows/local-agent-guides-ci.yml | 16 +--
.github/workflows/mlx-ci.yml | 12 +-
.github/workflows/release-desktop.yml | 8 +-
.github/workflows/security-audit.yml | 30 ++---
.github/workflows/studio-api-smoke.yml | 14 +--
.github/workflows/studio-backend-ci.yml | 2 +-
.../workflows/studio-export-capability-ci.yml | 2 +-
.github/workflows/studio-frontend-ci.yml | 4 +-
.github/workflows/studio-inference-smoke.yml | 56 ++++-----
.../workflows/studio-load-orchestrator-ci.yml | 4 +-
.github/workflows/studio-mac-api-smoke.yml | 10 +-
.../workflows/studio-mac-inference-smoke.yml | 58 +++++-----
.../workflows/studio-mac-install-matrix.yml | 4 +-
.github/workflows/studio-mac-ui-smoke.yml | 18 +--
.github/workflows/studio-mac-update-smoke.yml | 16 +--
.github/workflows/studio-tauri-smoke.yml | 2 +-
.github/workflows/studio-ui-smoke.yml | 34 +++---
.github/workflows/studio-update-smoke.yml | 12 +-
.../workflows/studio-windows-api-smoke.yml | 16 +--
.../studio-windows-inference-smoke.yml | 90 +++++++--------
.github/workflows/studio-windows-ui-smoke.yml | 32 +++---
.../workflows/studio-windows-update-smoke.yml | 26 ++---
.github/workflows/wheel-smoke.yml | 10 +-
README.md | 16 +--
build.sh | 8 +-
install.ps1 | 36 +++---
install.sh | 30 ++---
scripts/install_rocm_wsl_strixhalo.sh | 2 +-
scripts/lockfile_supply_chain_audit.py | 4 +-
scripts/scan_npm_packages.py | 2 +-
scripts/stamp_studio_release.py | 20 ++--
scripts/uninstall.ps1 | 14 +--
scripts/uninstall.sh | 14 +--
studio/MCP.md | 12 +-
studio/Unsloth_Studio_Colab.ipynb | 2 +-
.../assets/chat_templates/gemma-4-edge.jinja | 2 +-
.../assets/chat_templates/gemma-4.jinja | 2 +-
studio/backend/auth/authentication.py | 2 +-
studio/backend/auth/bootstrap_timeout.py | 10 +-
studio/backend/auth/storage.py | 6 +-
studio/backend/auth/terminal_prompt.py | 6 +-
studio/backend/cloudflare_tunnel.py | 8 +-
studio/backend/colab.py | 12 +-
studio/backend/core/data_recipe/jobs/parse.py | 6 +-
.../data_recipe/local_callable_validators.py | 2 +-
studio/backend/core/data_recipe/service.py | 4 +-
studio/backend/core/inference/__init__.py | 2 +-
.../core/inference/anthropic_compat.py | 2 +-
.../core/inference/chat_template_helpers.py | 2 +-
.../backend/core/inference/chat_templates.py | 4 +-
.../core/inference/external_provider.py | 14 +--
studio/backend/core/inference/llama_cpp.py | 80 ++++++-------
.../backend/core/inference/llama_keepwarm.py | 2 +-
.../core/inference/llama_server_args.py | 38 +++----
studio/backend/core/inference/llama_stats.py | 2 +-
.../core/inference/local_model_resolver.py | 4 +-
studio/backend/core/inference/mcp_client.py | 2 +-
.../core/inference/passthrough_healing.py | 4 +-
studio/backend/core/inference/pricing.py | 8 +-
.../core/inference/safetensors_agentic.py | 2 +-
.../inference/sandbox_site/sitecustomize.py | 2 +-
.../core/inference/tool_loop_controller.py | 2 +-
studio/backend/core/inference/tools.py | 14 +--
studio/backend/core/rag/captioner.py | 2 +-
studio/backend/core/rag/embed_llama_server.py | 2 +-
studio/backend/core/rag/embeddings.py | 2 +-
studio/backend/core/training/resume.py | 2 +-
studio/backend/core/training/trainer.py | 2 +-
studio/backend/core/training/training.py | 6 +-
studio/backend/core/training/worker.py | 12 +-
.../hub/services/download_lifecycle.py | 2 +-
.../hub/services/models/folder_browser.py | 2 +-
studio/backend/hub/services/models/ollama.py | 2 +-
studio/backend/hub/utils/state_dir.py | 2 +-
studio/backend/main.py | 14 +--
studio/backend/mcp_server.py | 26 ++---
studio/backend/models/inference.py | 16 +--
studio/backend/models/training.py | 2 +-
.../data-designer-github-repo-seed/README.md | 4 +-
.../__init__.py | 2 +-
.../data_designer_github_repo_seed/scraper.py | 2 +-
.../backend/requirements/extras-no-deps.txt | 2 +-
.../backend/requirements/no-torch-runtime.txt | 2 +-
.../requirements/single-env/constraints.txt | 2 +-
studio/backend/requirements/studio.txt | 4 +-
studio/backend/routes/auth.py | 4 +-
studio/backend/routes/data_recipe/jobs.py | 2 +-
studio/backend/routes/datasets.py | 2 +-
studio/backend/routes/inference.py | 64 +++++------
studio/backend/routes/mcp_servers.py | 2 +-
studio/backend/routes/models.py | 24 ++--
studio/backend/routes/training.py | 6 +-
studio/backend/run.py | 62 +++++-----
studio/backend/startup_banner.py | 4 +-
studio/backend/tests/conftest.py | 4 +-
.../tests/test_amd_apu_unified_memory.py | 2 +-
.../tests/test_anthropic_compaction.py | 2 +-
.../tests/test_anthropic_fast_mode_edge.py | 6 +-
.../backend/tests/test_anthropic_messages.py | 10 +-
studio/backend/tests/test_compute_buffer.py | 2 +-
studio/backend/tests/test_cpu_threads.py | 4 +-
.../backend/tests/test_frontend_resolution.py | 2 +-
studio/backend/tests/test_gemini_provider.py | 4 +-
.../test_gemma4_chat_template_override.py | 2 +-
studio/backend/tests/test_hf_xet_fallback.py | 18 +--
studio/backend/tests/test_identity.py | 2 +-
.../test_index_bootstrap_origin_extra.py | 2 +-
.../tests/test_llama_cpp_context_fit.py | 2 +-
.../tests/test_llama_cpp_mmproj_fallback.py | 4 +-
.../tests/test_llama_cpp_mtp_detection.py | 4 +-
.../tests/test_llama_cpp_no_context_shift.py | 2 +-
.../tests/test_llama_cpp_props_readback.py | 4 +-
.../backend/tests/test_llama_cpp_tool_loop.py | 2 +-
.../tests/test_llama_cpp_wait_for_health.py | 2 +-
.../test_llama_cpp_wait_for_vram_settle.py | 10 +-
.../test_llama_cpp_windows_nvidia_path.py | 2 +-
.../backend/tests/test_llama_server_args.py | 16 +--
.../tests/test_local_llama_cpp_link.py | 4 +-
studio/backend/tests/test_mcp_servers.py | 2 +-
.../tests/test_mcp_stdio_improvements.py | 2 +-
.../tests/test_mlx_inference_backend.py | 2 +-
studio/backend/tests/test_mlx_repair.py | 6 +-
studio/backend/tests/test_mtp_vram_budget.py | 24 ++--
.../backend/tests/test_multimodal_document.py | 2 +-
.../tests/test_nudge_tool_calls_wiring.py | 6 +-
.../tests/test_offline_gguf_cache_fallback.py | 2 +-
.../tests/test_offline_inference_parent.py | 2 +-
.../backend/tests/test_openai_auto_switch.py | 2 +-
.../backend/tests/test_openai_compaction.py | 2 +-
.../tests/test_openai_image_generation.py | 2 +-
.../tests/test_openai_tool_passthrough.py | 32 +++---
.../tests/test_password_prompt_backstop.py | 2 +-
studio/backend/tests/test_permission_mode.py | 4 +-
studio/backend/tests/test_providers_api.py | 4 +-
.../tests/test_rag_embed_llama_server.py | 2 +-
.../test_recommended_folders_permission.py | 2 +-
.../tests/test_responses_tool_passthrough.py | 2 +-
studio/backend/tests/test_rocm_oom_guard.py | 2 +-
.../tests/test_safetensors_tool_loop.py | 8 +-
.../backend/tests/test_secure_tunnel_gate.py | 2 +-
.../backend/tests/test_server_disk_logging.py | 2 +-
studio/backend/tests/test_slot_offload_fit.py | 2 +-
studio/backend/tests/test_studio_api.py | 6 +-
studio/backend/tests/test_tensor_parallel.py | 2 +-
.../backend/tests/test_tool_confirm_stream.py | 4 +-
.../tests/test_tool_message_empty_content.py | 2 +-
.../tests/test_tp_vision_regression.py | 2 +-
.../backend/tests/test_trained_model_scan.py | 2 +-
.../tests/test_training_nan_loss_handling.py | 2 +-
.../backend/tests/test_transformers_latest.py | 2 +-
studio/backend/utils/_studio_release_build.py | 2 +-
studio/backend/utils/api_errors.py | 6 +-
studio/backend/utils/client_ip.py | 4 +-
studio/backend/utils/cpu_threads.py | 2 +-
studio/backend/utils/datasets/cache_safe.py | 4 +-
.../backend/utils/hardware/VRAM_ESTIMATION.md | 2 +-
studio/backend/utils/hardware/amd.py | 2 +-
studio/backend/utils/hardware/hardware.py | 6 +-
.../backend/utils/helper_precache_settings.py | 4 +-
studio/backend/utils/hf_xet_fallback.py | 16 +--
studio/backend/utils/host_policy.py | 2 +-
studio/backend/utils/llama_cpp_update.py | 10 +-
studio/backend/utils/mlx_repair.py | 26 ++---
studio/backend/utils/models/checkpoints.py | 2 +-
studio/backend/utils/models/model_config.py | 4 +-
studio/backend/utils/paths/storage_roots.py | 6 +-
studio/backend/utils/preview_rate_limit.py | 2 +-
studio/backend/utils/process_lifetime.py | 4 +-
studio/backend/utils/studio_version.py | 6 +-
studio/backend/utils/training_runs.py | 2 +-
studio/backend/utils/transformers_latest.py | 10 +-
studio/backend/utils/transformers_version.py | 2 +-
studio/backend/utils/upload_limits.py | 2 +-
studio/frontend/.npmrc | 2 +-
studio/frontend/src/app/provider.tsx | 2 +-
.../assistant-ui/model-selector/pickers.tsx | 6 +-
.../src/components/assistant-ui/thread.tsx | 2 +-
.../frontend/src/components/ui/confetti.tsx | 2 +-
.../features/auth/components/auth-form.tsx | 2 +-
.../features/chat/artifacts/html-frame.tsx | 2 +-
.../chat/hooks/use-chat-sidebar-items.ts | 2 +-
.../src/features/chat/lib/friendly-names.ts | 2 +-
.../features/chat/provider-capabilities.ts | 2 +-
.../chat/stores/chat-runtime-store.ts | 4 +-
.../chat/utils/chat-history-storage.ts | 2 +-
.../src/features/hub/download-manager/api.ts | 2 +-
.../components/steps/model-selection-step.tsx | 2 +-
.../settings/components/usage-examples.tsx | 8 +-
.../studio/recent-trainings-section.tsx | 2 +-
.../training/stores/training-config-store.ts | 2 +-
.../transformers-upgrade-dialog.tsx | 2 +-
.../frontend/src/hooks/use-tauri-backend.ts | 4 +-
studio/frontend/src/i18n/README.md | 2 +-
studio/frontend/src/i18n/locales/ar.ts | 4 +-
studio/frontend/src/i18n/locales/de.ts | 4 +-
studio/frontend/src/i18n/locales/es.ts | 4 +-
studio/frontend/src/i18n/locales/fr.ts | 4 +-
studio/frontend/src/i18n/locales/hi.ts | 4 +-
studio/frontend/src/i18n/locales/ja.ts | 2 +-
studio/frontend/src/i18n/locales/ko.ts | 4 +-
studio/frontend/src/i18n/locales/ru.ts | 4 +-
studio/frontend/src/i18n/locales/zh-CN.ts | 2 +-
studio/frontend/src/index.css | 2 +-
studio/frontend/src/lib/tauri-diagnostics.ts | 2 +-
studio/install_llama_prebuilt.py | 8 +-
studio/install_node_prebuilt.py | 2 +-
studio/install_python_stack.py | 6 +-
studio/setup.ps1 | 46 ++++----
studio/setup.sh | 36 +++---
studio/src-tauri/src/commands.rs | 12 +-
studio/src-tauri/src/desktop_auth.rs | 6 +-
studio/src-tauri/src/main.rs | 2 +-
studio/src-tauri/src/native_path_policy.rs | 2 +-
tests/python/test_e2e_no_torch_sandbox.py | 8 +-
tests/python/test_studio_import_no_torch.py | 2 +-
.../test_prewarm_base_model_hub_cache.py | 2 +-
tests/studio/_playwright_robust.py | 4 +-
.../smoke_test_parallel_studio_home.py | 2 +-
.../install/test_launch_studio_launcher.py | 2 +-
.../install/test_managed_node_runtime.py | 2 +-
tests/studio/install/test_pr5940_followups.py | 6 +-
tests/studio/install/test_rocm_support.py | 4 +-
tests/studio/install/test_selection_logic.py | 12 +-
tests/studio/playwright_chat_ime_i18n.py | 2 +-
tests/studio/playwright_chat_ui.py | 10 +-
tests/studio/playwright_extra_ui.py | 10 +-
tests/studio/run_real_mlx_smoke.py | 2 +-
tests/studio/studio_api_smoke.py | 6 +-
tests/studio/test_auth_form_input_count.py | 2 +-
tests/studio/test_chat_title_generation.py | 2 +-
tests/studio/test_cli_studio_stop_windows.py | 2 +-
tests/studio/test_hardware_dispatch_matrix.py | 20 ++--
tests/studio/test_is_mlx_dispatch_gate.py | 4 +-
tests/studio/test_llama_cpp_wall_clock_cap.py | 2 +-
.../test_locale_root_direction_contract.py | 2 +-
tests/studio/test_node_decision.ps1 | 2 +-
.../test_studio_gguf_export_script_pin.py | 2 +-
.../test_studio_text_descender_clipping.py | 2 +-
tests/test_studio_install_workspace_guard.py | 8 +-
tests/test_studio_root_resilience.py | 2 +-
tests/test_studio_shutdown_thread_wait.py | 2 +-
unsloth/chat_templates.py | 2 +-
unsloth/import_fixes.py | 2 +-
unsloth/models/loader_utils.py | 2 +-
unsloth/save.py | 4 +-
unsloth/tokenizer_utils.py | 2 +-
unsloth_cli/__init__.py | 2 +-
unsloth_cli/_inference.py | 28 ++---
unsloth_cli/commands/chat.py | 6 +-
unsloth_cli/commands/inference.py | 4 +-
unsloth_cli/commands/start.py | 86 +++++++-------
unsloth_cli/commands/studio.py | 106 +++++++++---------
unsloth_cli/tests/test_inference_chat.py | 2 +-
unsloth_cli/tests/test_start.py | 38 +++----
.../tests/test_studio_cloudflare_flag.py | 2 +-
.../tests/test_studio_password_prompt.py | 4 +-
unsloth_cli/tests/test_studio_secure_flag.py | 2 +-
unsloth_cli/tests/test_studio_verbose_flag.py | 4 +-
264 files changed, 1095 insertions(+), 1095 deletions(-)
diff --git a/.gitattributes b/.gitattributes
index 5f04b5e9d1..0025f2a697 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -6,7 +6,7 @@
# them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -").
*.sh text eol=lf
-# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather
+# Normalize Unsloth frontend sources to LF. Scoped to the frontend tree (rather
# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files
# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts)
# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF.
diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh
index f4189a159e..2457f08407 100755
--- a/.github/scripts/agent-guides-drive.sh
+++ b/.github/scripts/agent-guides-drive.sh
@@ -166,8 +166,8 @@ parse_connect() {
echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw"
CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)"
# The launch command is the last non-export, non-status line. start.py
- # prints "Studio · model " and "Updated ..." status lines first.
- CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \
+ # prints "Unsloth · model " and "Updated ..." status lines first.
+ CONNECT_CMD="$(grep -vE '^(export |unset |Unsloth |Updated |Disabled |Warning|Loading)' "$raw" \
| grep -E '[^[:space:]]' | tail -1)"
[ -n "$CONNECT_CMD" ] || guide_fail "could not parse a launch command from connect --no-launch output"
redact "$raw"
diff --git a/.github/scripts/assert-llama-loads.sh b/.github/scripts/assert-llama-loads.sh
index c2ffe27469..62ef80d364 100755
--- a/.github/scripts/assert-llama-loads.sh
+++ b/.github/scripts/assert-llama-loads.sh
@@ -2,7 +2,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
-# Assert Studio installed a llama.cpp that loads and runs on THIS macOS. Tests
+# Assert Unsloth installed a llama.cpp that loads and runs on THIS macOS. Tests
# the contract that matters (binaries load and their minimum-OS is <= this host)
# instead of the old "did install.sh fall back to a source build?" grep, since a
# source build with a correct deployment target is a valid outcome.
diff --git a/.github/scripts/assert-prompt-cache.sh b/.github/scripts/assert-prompt-cache.sh
index f5b6b075eb..8c28569f77 100755
--- a/.github/scripts/assert-prompt-cache.sh
+++ b/.github/scripts/assert-prompt-cache.sh
@@ -31,7 +31,7 @@
# (llama_cpp.py:337-340). So default: ~/.unsloth/studio/logs/llama-server/.
#
# is the INTERNAL llama-server port (self._find_free_port(),
-# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Studio port. So we must
+# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Unsloth port. So we must
# NOT filter the log glob by STUDIO_PORT (the brief's `port-`
# glob would never match). We pick the newest llama-*.log instead.
#
diff --git a/.github/scripts/hf-download-with-retry.sh b/.github/scripts/hf-download-with-retry.sh
index 013a459f46..6dec93356a 100755
--- a/.github/scripts/hf-download-with-retry.sh
+++ b/.github/scripts/hf-download-with-retry.sh
@@ -3,7 +3,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
#
# Download a single file from a Hugging Face repo with a stall-retry
-# watchdog. Used by the Studio CI workflows so a hung hf-xet transfer
+# watchdog. Used by the Unsloth CI workflows so a hung hf-xet transfer
# kills + retries instead of silently consuming the job's timeout.
#
# Usage: hf-download-with-retry.sh REPO FILE LOCAL_DIR
@@ -35,7 +35,7 @@ REPO="${1:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}"
FILE="${2:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}"
# LOCAL_DIR is optional. If empty, hf falls back to HF_HUB_CACHE
# (~/.cache/huggingface/hub) which is the desired path for callers
-# that populate HF_HOME for a downstream Studio model load.
+# that populate HF_HOME for a downstream Unsloth model load.
LOCAL_DIR="${3:-}"
# Stall threshold per attempt, in seconds. Override with
diff --git a/.github/workflows/lint-ci.yml b/.github/workflows/lint-ci.yml
index bd859a6e9e..e1f0afd299 100644
--- a/.github/workflows/lint-ci.yml
+++ b/.github/workflows/lint-ci.yml
@@ -13,10 +13,10 @@
# committed YAML / JSON config.
#
# TypeScript and Rust are NOT duplicated here on purpose:
-# - Studio Frontend CI runs `npm run typecheck` (= `tsc --noEmit`)
+# - Unsloth Frontend CI runs `npm run typecheck` (= `tsc --noEmit`)
# and `npm run build` (vite/swc) on every studio/frontend/**
# change, which is a full TS AST + type check.
-# - Studio Tauri CI runs `tauri build --debug --no-bundle` on
+# - Unsloth Tauri CI runs `tauri build --debug --no-bundle` on
# every studio/src-tauri/** or studio/frontend/** change, which
# compiles the Rust crate (= cargo check + cargo build).
# Each is a stricter check than a parse-only step would be, so a
diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml
index 25796bd5cf..c48328e90f 100644
--- a/.github/workflows/local-agent-guides-ci.yml
+++ b/.github/workflows/local-agent-guides-ci.yml
@@ -154,7 +154,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
@@ -256,7 +256,7 @@ jobs:
done
fi
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
# Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make
@@ -359,7 +359,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
@@ -448,7 +448,7 @@ jobs:
done
fi
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
# Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make
@@ -543,7 +543,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
@@ -620,7 +620,7 @@ jobs:
done
fi
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then
@@ -706,7 +706,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
@@ -764,7 +764,7 @@ jobs:
done
fi
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
# Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make
diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml
index a2f716a93c..aadf0b54e6 100644
--- a/.github/workflows/mlx-ci.yml
+++ b/.github/workflows/mlx-ci.yml
@@ -130,7 +130,7 @@ jobs:
# MLX support landed after the most recent unsloth-zoo PyPI
# release; the wheel still raises NotImplementedError on
# Apple Silicon when device_type.get_device_type() runs
- # unguarded. Studio's own install.sh overlays unsloth-zoo
+ # unguarded. Unsloth's own install.sh overlays unsloth-zoo
# from git main for the same reason. Pulling deps lets pip
# resolve the platform-conditional MLX-only wheels (mlx,
# mlx-lm, mlx-vlm gated on darwin+arm64 in unsloth-zoo's
@@ -317,13 +317,13 @@ jobs:
echo
done
- # Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the
+ # Validates the macOS prebuilt path Unsloth'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)
+ - name: Unsloth prebuilt llama.cpp install + GGUF download (Mac M1)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -344,12 +344,12 @@ jobs:
# 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)
+ - name: Unsloth 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);
+ # Unsloth 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"
@@ -400,4 +400,4 @@ jobs:
tail -40 /tmp/llama-server.log
exit 1
fi
- echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works"
+ echo "OK: Unsloth prebuilt llama.cpp on Mac M1 + GGUF /completion works"
diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml
index 4daafae35d..081eda4e32 100644
--- a/.github/workflows/release-desktop.yml
+++ b/.github/workflows/release-desktop.yml
@@ -4,7 +4,7 @@ on:
workflow_dispatch:
inputs:
studio_version:
- description: 'Studio version tag to release (for example, v0.1.39-beta)'
+ description: 'Unsloth version tag to release (for example, v0.1.39-beta)'
type: string
required: true
pypi_version:
@@ -69,7 +69,7 @@ jobs:
if not studio_version:
sys.exit('studio_version is required, for example v0.1.39-beta')
if re.fullmatch(r'v?20\d{2}\.\d+\.\d+(?:[-+][0-9A-Za-z.-]+)?', studio_version):
- sys.exit(f'studio_version must be a Studio SemVer tag, not a date-style backend version: {studio_version}')
+ sys.exit(f'studio_version must be an Unsloth SemVer tag, not a date-style backend version: {studio_version}')
semver_tag = re.compile(
r'^v(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)'
@@ -146,7 +146,7 @@ jobs:
print(f'pypi_version={pypi_version}', file=output)
PY
- - name: Verify PyPI package and Studio stamp
+ - name: Verify PyPI package and Unsloth stamp
shell: bash
env:
STUDIO_VERSION: ${{ steps.prepare.outputs.studio_version }}
@@ -211,7 +211,7 @@ jobs:
fi
python3 scripts/stamp_studio_release.py --verify-dist "$RUNNER_TEMP/pypi-unsloth-dist" --expected "$STUDIO_VERSION"
else
- echo "scripts/stamp_studio_release.py not found; release-desktop requires #5308 to verify the PyPI Studio stamp." >&2
+ echo "scripts/stamp_studio_release.py not found; release-desktop requires #5308 to verify the PyPI Unsloth stamp." >&2
exit 1
fi
diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml
index 1275d12216..27eafbedea 100644
--- a/.github/workflows/security-audit.yml
+++ b/.github/workflows/security-audit.yml
@@ -36,8 +36,8 @@
# - unsloth `huggingfacenotorch` extras (the canonical install path
# for fine-tuning users; pulls transformers / peft / accelerate /
# trl / datasets / diffusers / sentence-transformers / etc.)
-# - all six Studio backend requirements files
-# - Studio frontend (npm) and Tauri shell (cargo)
+# - all six Unsloth backend requirements files
+# - Unsloth frontend (npm) and Tauri shell (cargo)
# Each Python step builds a filtered dep list from pyproject.toml +
# requirements/*.txt before auditing. We do NOT install any of these
# -- pip-audit resolves through PyPI metadata, scan_packages.py
@@ -218,7 +218,7 @@ jobs:
# on the runner). A comment line is left in place so the
# skipped specs are obvious in the artifact.
# The `huggingface` extra is `huggingfacenotorch` plus torch /
- # torchvision / triton, deliberately skipped: Studio backend
+ # torchvision / triton, deliberately skipped: Unsloth backend
# already pins a torch and the +cu* / +cpu local-version tags
# trip up the PyPI resolver in `-r` mode.
run: |
@@ -253,7 +253,7 @@ jobs:
# `-r requirements.txt` resolves the requirements through pip's
# dependency resolver against PyPI metadata and audits the
# resolved tree without ever executing setup.py / install
- # hooks. Way faster than installing the full Studio runtime
+ # hooks. Way faster than installing the full Unsloth runtime
# and -- critically -- safer: an attacker who has compromised
# a transitive dep cannot run code in this job.
#
@@ -326,9 +326,9 @@ jobs:
} >> "$GITHUB_STEP_SUMMARY"
# ─────────────────────────────────────────────────────────────
- # npm: Studio frontend
+ # npm: Unsloth frontend
# ─────────────────────────────────────────────────────────────
- - name: npm audit (Studio frontend)
+ - name: npm audit (Unsloth frontend)
# `npm audit` resolves the lockfile through the npmjs.com
# advisory DB. `--audit-level=high` filters the noise floor
# to only HIGH and CRITICAL. We do NOT pass --omit=dev: a
@@ -342,7 +342,7 @@ jobs:
# Always also write the full JSON for grep-ability.
npm audit --json > ../../logs-npm-audit.json || true
{
- echo "## npm audit (Studio frontend)"
+ echo "## npm audit (Unsloth frontend)"
echo
echo '```'
tail -200 ../../logs-npm-audit.txt
@@ -350,9 +350,9 @@ jobs:
} >> "$GITHUB_STEP_SUMMARY"
# ─────────────────────────────────────────────────────────────
- # cargo: Studio Tauri shell
+ # cargo: Unsloth Tauri shell
# ─────────────────────────────────────────────────────────────
- - name: cargo audit (Studio Tauri)
+ - name: cargo audit (Unsloth Tauri)
# `--deny warnings` would make the job fail on any advisory.
# Keep non-blocking initially; drop continue-on-error after
# the baseline closes.
@@ -362,7 +362,7 @@ jobs:
set +e
cargo audit | tee ../../logs-cargo-audit.txt
{
- echo "## cargo audit (Studio Tauri)"
+ echo "## cargo audit (Unsloth Tauri)"
echo
echo '```'
tail -200 ../../logs-cargo-audit.txt
@@ -559,7 +559,7 @@ jobs:
# ─────────────────────────────────────────────────────────────
# CycloneDX SBOM. Lets downstream consumers audit what's
- # actually shipped in unsloth wheels and the Studio backend
+ # actually shipped in unsloth wheels and the Unsloth backend
# runtime. Generates one JSON file per requirements input plus
# a combined SBOM keyed off pyproject.toml; uploads as a build
# artifact (and a future step can attest it via SLSA).
@@ -740,7 +740,7 @@ jobs:
# `--with-deps` makes the scan transitive: every package the
# declared set resolves to gets fetched and pattern-scanned, not
# just the top-level pins. Resolving the full transitive closure
- # of the unsloth + Studio dep tree downloads several hundred
+ # of the unsloth + Unsloth dep tree downloads several hundred
# archives, hence the longer timeout.
#
# Sharded across runners for wall-clock parallelism. Each shard
@@ -749,7 +749,7 @@ jobs:
# composition tries to balance load:
# - hf-stack: pyproject extras + no-torch-runtime
# (~150 archives, transformers/peft/accelerate/...)
- # - studio: FastAPI/Studio backend + overrides + extras-no-deps
+ # - studio: FastAPI/Unsloth backend + overrides + extras-no-deps
# (~150 archives, smaller scientific stack)
# - extras: the heavy openai-whisper / scikit-learn / librosa
# stack (~250 archives, dominant cost)
@@ -964,7 +964,7 @@ jobs:
# documented at scripts/scan_npm_packages.py top-of-file. The
# script is stdlib-only so adding it does not increase the
# transitive supply-chain surface.
- name: npm scan-packages (Studio frontend tarballs)
+ name: npm scan-packages (Unsloth frontend tarballs)
runs-on: ubuntu-latest
timeout-minutes: 30
needs: []
@@ -1173,7 +1173,7 @@ jobs:
with:
python-version: '3.12'
- - name: Install Studio frontend deps (--ignore-scripts)
+ - name: Install Unsloth frontend deps (--ignore-scripts)
# `npm audit signatures` requires node_modules to be populated.
# `--ignore-scripts` is mandatory: this is exactly the lever the
# new-install-script gate below protects against, and we must
diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml
index 15efee382e..cdf1f6bf12 100644
--- a/.github/workflows/studio-api-smoke.yml
+++ b/.github/workflows/studio-api-smoke.yml
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-# Studio API & Auth Tests -- HTTP-level integration tests for the
+# Unsloth API & Auth Tests -- HTTP-level integration tests for the
# FastAPI surface. No Playwright, no model UI; tests/studio/test_studio_api_smoke.py
# runs ~30 s and asserts:
# - CORS hardening (no wildcard + credentials, no bootstrap leak)
@@ -15,7 +15,7 @@
# Reuses the GGUF cache key from studio-ui-smoke.yml so the model
# download is one cache-hit on the second job.
-name: Studio API CI
+name: Unsloth API CI
on:
pull_request:
@@ -40,7 +40,7 @@ permissions:
jobs:
api-smoke:
- name: Studio API & Auth Tests
+ name: Unsloth API & Auth Tests
runs-on: ubuntu-latest
timeout-minutes: 12
env:
@@ -98,7 +98,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -111,7 +111,7 @@ jobs:
- name: Install pyjwt for the JWT-expiry forge test
run: pip install 'pyjwt>=2.6'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -144,7 +144,7 @@ jobs:
echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV"
echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV"
- - name: Run Studio API & Auth tests
+ - name: Run Unsloth API & Auth tests
# The script is named WITHOUT a `test_` prefix so it isn't
# auto-collected by pytest in Backend CI's `tests/` walk
# (which doesn't set BASE_URL and would crash at import).
@@ -153,7 +153,7 @@ jobs:
STUDIO_AUTH_DIR: /home/runner/.unsloth/studio/auth
run: python tests/studio/studio_api_smoke.py
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml
index 3022127a2b..b8f587b63e 100644
--- a/.github/workflows/studio-backend-ci.yml
+++ b/.github/workflows/studio-backend-ci.yml
@@ -64,7 +64,7 @@ jobs:
- name: Install backend test dependencies (CPU only)
run: |
python -m pip install --upgrade pip
- # Studio's declared backend deps:
+ # Unsloth's declared backend deps:
pip install -r studio/backend/requirements/studio.txt
# Extras that studio.txt does not list but the import chain needs
# (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography
diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml
index 1ee6489209..83df3ed476 100644
--- a/.github/workflows/studio-export-capability-ci.yml
+++ b/.github/workflows/studio-export-capability-ci.yml
@@ -9,7 +9,7 @@
# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block
# torch/unsloth, so the job installs only a CPU PyTorch plus import deps.
-name: Studio export capability
+name: Unsloth export capability
on:
pull_request:
diff --git a/.github/workflows/studio-frontend-ci.yml b/.github/workflows/studio-frontend-ci.yml
index b42086f191..3a9e373915 100644
--- a/.github/workflows/studio-frontend-ci.yml
+++ b/.github/workflows/studio-frontend-ci.yml
@@ -136,7 +136,7 @@ jobs:
- name: Build
run: npm run build
- - name: Built bundle must not contain Studio's unstable_Provider call site
+ - name: Built bundle must not contain Unsloth's unstable_Provider call site
run: |
set -e
JS=$(ls dist/assets/index-*.js | head -1)
@@ -144,7 +144,7 @@ jobs:
echo "main bundle: $JS"
echo "unstable_Provider: hits=$HITS (assistant-ui internals contribute up to 3)"
if [ "$HITS" -gt 3 ]; then
- echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Studio bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead."
+ echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Unsloth bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead."
exit 1
fi
diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml
index 58ef2558f3..c2d52eac22 100644
--- a/.github/workflows/studio-inference-smoke.yml
+++ b/.github/workflows/studio-inference-smoke.yml
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-# Three end-to-end smoke jobs that boot a freshly-installed Studio and
+# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and
# exercise the surfaces real users hit through the OpenAI / Anthropic
# SDKs and curl. Each job picks the smallest model that exercises the
# behaviour under test, primes HF_HOME via actions/cache, and shares
@@ -27,7 +27,7 @@
# All three jobs run in parallel. Total wall time is dominated by job 3
# on a cold cache; warm cache cuts that to ~3 min.
-name: Studio GGUF CI
+name: Unsloth GGUF CI
on:
pull_request:
@@ -112,7 +112,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -125,7 +125,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -142,7 +142,7 @@ jobs:
fi
sleep 1
done
- echo "Studio did not become healthy in 180s"
+ echo "Unsloth did not become healthy in 180s"
tail -200 logs/studio.log
exit 1
@@ -229,11 +229,11 @@ jobs:
return replies
def run_anthropic():
- # Two SDK quirks vs. Studio:
+ # Two SDK quirks vs. Unsloth:
# 1. base_url must NOT include /v1 -- the SDK appends
# /v1/messages itself; otherwise the request hits
# /v1/v1/messages and 405s.
- # 2. The SDK sends `x-api-key` by default, but Studio's
+ # 2. The SDK sends `x-api-key` by default, but Unsloth's
# auth layer is HTTPBearer-only. Override via
# default_headers so Authorization: Bearer ... is
# sent instead.
@@ -276,7 +276,7 @@ jobs:
print(
f"[{label}] WARN non-determinism at temperature=0.0 across "
f"{len(determinism_failures)} of {len(first)} turn(s); "
- f"small-quant model drift, not a Studio regression. "
+ f"small-quant model drift, not an Unsloth regression. "
f"Details: " + " | ".join(determinism_failures)
)
# Sanity: turn-2 reply should mention the earlier question, and
@@ -290,7 +290,7 @@ jobs:
print(f"[{label}] {status_word} -- 4 turns, history grounded ('paris' present)")
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@@ -323,7 +323,7 @@ jobs:
# store xet chunks + blobs + snapshots = ~4 GiB compressed --
# 4-5x file-size inflation, dominated by xet chunks. Use main's
# `--local-dir gguf-cache` pattern to cache the flat .gguf only.
- # Studio's /api/inference/load accepts either a HF repo (which
+ # Unsloth's /api/inference/load accepts either a HF repo (which
# uses HF_HOME) or an absolute file path; passing the absolute
# path keeps the test off HF_HOME entirely so the cache size
# tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images
@@ -380,7 +380,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -390,7 +390,7 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- - name: Reset auth + boot Studio (API-only, default tool policy)
+ - name: Reset auth + boot Unsloth (API-only, default tool policy)
# We deliberately use the API-only mode rather than
# `unsloth studio run` because the latter calls
# `set_tool_policy(...)` with a resolved bool: on loopback the
@@ -503,7 +503,7 @@ jobs:
that the tool path executed.
A shared CI runner can stall the stream transport (the
- connection opening, or a mid-stream read) even when Studio
+ connection opening, or a mid-stream read) even when Unsloth
is healthy, so retry a stall once with a fresh request
capped at 300s. A stall means the stream did NOT complete,
so partial events are normally NOT returned (an early
@@ -575,11 +575,11 @@ jobs:
def _tool_invoked(events):
"""Structural check: True iff some SSE payload is a real
- tool envelope (Studio tool_start/tool_end, Anthropic
+ tool envelope (Unsloth tool_start/tool_end, Anthropic
tool_use/tool_result, OpenAI non-empty delta.tool_calls /
message.tool_calls / finish_reason='tool_calls' /
role:'tool' / function_call). tool_status is NOT
- evidence: Studio emits empty tool_status events on
+ evidence: Unsloth emits empty tool_status events on
iteration boundaries even when no tool ran.
"""
for raw in events:
@@ -698,7 +698,7 @@ jobs:
attempt has structural invocation evidence. WARN (not
FAIL) if invoked but no attempt produces the expected
literal in tool_end.result -- small-quant Qwen3.5-2B can
- emit OpenAI tool_calls deltas without Studio's GGUF
+ emit OpenAI tool_calls deltas without Unsloth's GGUF
agentic loop intercepting them, and that GGUF-vs-OpenAI
format mismatch is out of scope for #5642.
"""
@@ -811,7 +811,7 @@ jobs:
# because (a) the search may legitimately return no results,
# and (b) DuckDuckGo upstream blocks GHA IP ranges often
# enough that requiring a tool_call marker would create
- # red-herring failures from infra rather than from Studio.
+ # red-herring failures from infra rather than from Unsloth.
try:
# Best-effort and bounded: a single 180s attempt keeps a stall
# from eating the job's timeout-minutes (it already WARNs, so a
@@ -834,7 +834,7 @@ jobs:
print(f"[tools] WARN web_search probe failed (non-blocking): {exc}")
# ── 5. Thinking on / off ─────────────────────────────────────
- # Studio strips think blocks from message.content for tools-mode
+ # Unsloth strips think blocks from message.content for tools-mode
# responses, so we toggle plain chat (no enable_tools) and look
# at the surfaced reasoning_content / message.thinking field.
def thinking_call(enable):
@@ -848,7 +848,7 @@ jobs:
})
assert status == 200
msg = data["choices"][0]["message"]
- # Studio surfaces thinking via reasoning_content (OpenAI
+ # Unsloth surfaces thinking via reasoning_content (OpenAI
# extension). Fall back to inline markers for
# robustness across template versions.
raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "")
@@ -868,7 +868,7 @@ jobs:
print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)")
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@@ -960,7 +960,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -973,7 +973,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
# See Job 2's comment: API-only mode keeps tool_policy=None so
# response_format requests aren't routed through the agentic
# tool loop.
@@ -1076,13 +1076,13 @@ jobs:
# llama.cpp's HTTP server supports OpenAI-compatible JSON
# mode: `response_format: {"type": "json_object"}` constrains
# the model to emit syntactically-valid JSON. We use raw HTTP
- # rather than the OpenAI SDK so that the field shape Studio
+ # rather than the OpenAI SDK so that the field shape Unsloth
# forwards to llama-server is unambiguous (the SDK rewrites
# response_format depending on which variant it recognises).
# We deliberately do NOT pass a strict JSON schema -- on
# small Gemma-4 quants the GBNF-from-schema path occasionally
# produces empty output, and JSON mode is the surface we care
- # about exposing through Studio.
+ # about exposing through Unsloth.
status, data = post("/v1/chat/completions", {
"model": "default",
"messages": [
@@ -1112,7 +1112,7 @@ jobs:
print(f"[json] PASS json_object -> {parsed}")
# ── 2. OpenAI image_url (data URI base64) ───────────────────
- # 64x64 solid-red PNG. stb_image (used by Studio's image
+ # 64x64 solid-red PNG. stb_image (used by Unsloth's image
# normaliser at routes/inference.py:3410) rejects 4x4 or
# smaller PNGs as truncated, so we go up to 64x64 -- still
# tiny in token cost. The assertion is loose: any non-empty
@@ -1148,9 +1148,9 @@ jobs:
print("[image/openai] PASS image_url accepted, non-empty response")
# ── 3. Anthropic source/base64 image ────────────────────────
- # Two SDK quirks vs. Studio: base_url must NOT include /v1
+ # Two SDK quirks vs. Unsloth: base_url must NOT include /v1
# (the SDK appends it itself; otherwise /v1/v1/messages -> 405),
- # and Studio's auth is HTTPBearer-only so the SDK's default
+ # and Unsloth's auth is HTTPBearer-only so the SDK's default
# x-api-key header is ignored -- send Authorization: Bearer
# via default_headers.
anthropic = Anthropic(
@@ -1184,7 +1184,7 @@ jobs:
print("[image/anthropic] PASS source/base64 accepted, non-empty response")
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-load-orchestrator-ci.yml b/.github/workflows/studio-load-orchestrator-ci.yml
index 93d1a7742d..8710efc2bd 100644
--- a/.github/workflows/studio-load-orchestrator-ci.yml
+++ b/.github/workflows/studio-load-orchestrator-ci.yml
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
-# Event-loop regression test for the Studio model-load orchestrator.
+# Event-loop regression test for the Unsloth model-load orchestrator.
# Pins down issue #5642 (Win10 UI freeze on model load): the /load
# route calls LlamaCppBackend.detect_audio_type synchronously, blocking
# the FastAPI event loop on a chain of sync httpx.Client.post() probes.
@@ -14,7 +14,7 @@
# danielhanchen/unsloth-staging-2 (Ubuntu / macOS / Windows all
# green at PR time).
-name: Studio load-orchestrator CI
+name: Unsloth load-orchestrator CI
on:
pull_request:
diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml
index 617ce189dc..1968885a1d 100644
--- a/.github/workflows/studio-mac-api-smoke.yml
+++ b/.github/workflows/studio-mac-api-smoke.yml
@@ -33,7 +33,7 @@ permissions:
jobs:
api-smoke:
- name: Studio API & Auth Tests
+ name: Unsloth API & Auth Tests
runs-on: macos-14
timeout-minutes: 25
env:
@@ -83,7 +83,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -99,7 +99,7 @@ jobs:
- name: Install pyjwt for the JWT-expiry forge test
run: pip install 'pyjwt>=2.6'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -129,13 +129,13 @@ jobs:
echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV"
echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV"
- - name: Run Studio API & Auth tests
+ - name: Run Unsloth API & Auth tests
env:
BASE_URL: http://127.0.0.1:18895
STUDIO_AUTH_DIR: /Users/runner/.unsloth/studio/auth
run: python tests/studio/studio_api_smoke.py
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml
index 946681706a..ce15eed5c8 100644
--- a/.github/workflows/studio-mac-inference-smoke.yml
+++ b/.github/workflows/studio-mac-inference-smoke.yml
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-# Three end-to-end smoke jobs that boot a freshly-installed Studio and
+# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and
# exercise the surfaces real users hit through the OpenAI / Anthropic
# SDKs and curl. Each job picks the smallest model that exercises the
# behaviour under test, primes a model cache via actions/cache, and
@@ -108,7 +108,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -124,7 +124,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -141,7 +141,7 @@ jobs:
fi
sleep 1
done
- echo "Studio did not become healthy in 180s"
+ echo "Unsloth did not become healthy in 180s"
tail -200 logs/studio.log
exit 1
@@ -228,11 +228,11 @@ jobs:
return replies
def run_anthropic():
- # Two SDK quirks vs. Studio:
+ # Two SDK quirks vs. Unsloth:
# 1. base_url must NOT include /v1 -- the SDK appends
# /v1/messages itself; otherwise the request hits
# /v1/v1/messages and 405s.
- # 2. The SDK sends `x-api-key` by default, but Studio's
+ # 2. The SDK sends `x-api-key` by default, but Unsloth's
# auth layer is HTTPBearer-only. Override via
# default_headers so Authorization: Bearer ... is
# sent instead.
@@ -283,7 +283,7 @@ jobs:
print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded")
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@@ -363,7 +363,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -376,7 +376,7 @@ jobs:
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- - name: Reset auth + boot Studio (API-only, default tool policy)
+ - name: Reset auth + boot Unsloth (API-only, default tool policy)
# We deliberately use the API-only mode rather than
# `unsloth studio run` because the latter calls
# `set_tool_policy(...)` with a resolved bool: on loopback the
@@ -478,7 +478,7 @@ jobs:
call with enable_tools=true must use this helper.
A shared CI runner can stall the stream transport (the
- connection opening, or a mid-stream read) even when Studio
+ connection opening, or a mid-stream read) even when Unsloth
is healthy, so harden the read three ways: retry a stall
once with a fresh request capped at 300s; return any text
already streamed before a stall (a stall on the trailing
@@ -574,11 +574,11 @@ jobs:
assert status == 200, f"tool call status {status}: {data}"
choice = data["choices"][0]
tool_calls = (choice.get("message") or {}).get("tool_calls") or []
- # Studio's contract: when tool_choice='required', llama.cpp's
+ # Unsloth's contract: when tool_choice='required', llama.cpp's
# grammar should force a tool_calls payload. On Mac that
# contract is sometimes broken by the underlying quant; the
# PASS path is "tool_calls present + correct schema", the
- # WARN path documents Studio still returned 200 with a
+ # WARN path documents Unsloth still returned 200 with a
# well-formed choices[] envelope.
if tool_calls:
tc = tool_calls[0]
@@ -660,7 +660,7 @@ jobs:
print(f"[tools] WARN web_search probe failed (non-blocking): {exc}")
# ── 4. Thinking on / off ─────────────────────────────────────
- # Studio strips think blocks from message.content for tools-mode
+ # Unsloth strips think blocks from message.content for tools-mode
# responses, so we toggle plain chat (no enable_tools) and look
# at the surfaced reasoning_content / message.thinking field.
def thinking_call(enable):
@@ -678,7 +678,7 @@ jobs:
}, timeout = 180)
assert status == 200
msg = data["choices"][0]["message"]
- # Studio surfaces thinking via reasoning_content (OpenAI
+ # Unsloth surfaces thinking via reasoning_content (OpenAI
# extension). Fall back to inline markers for
# robustness across template versions.
raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "")
@@ -704,7 +704,7 @@ jobs:
print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)")
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@@ -810,7 +810,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -826,7 +826,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
# See Job 2's comment: API-only mode keeps tool_policy=None so
# response_format requests aren't routed through the agentic
# tool loop.
@@ -929,13 +929,13 @@ jobs:
# llama.cpp's HTTP server supports OpenAI-compatible JSON
# mode: `response_format: {"type": "json_object"}` constrains
# the model to emit syntactically-valid JSON. We use raw HTTP
- # rather than the OpenAI SDK so that the field shape Studio
+ # rather than the OpenAI SDK so that the field shape Unsloth
# forwards to llama-server is unambiguous (the SDK rewrites
# response_format depending on which variant it recognises).
# We deliberately do NOT pass a strict JSON schema -- on
# small Gemma-4 quants the GBNF-from-schema path occasionally
# produces empty output, and JSON mode is the surface we care
- # about exposing through Studio.
+ # about exposing through Unsloth.
status, data = post("/v1/chat/completions", {
"model": "default",
"messages": [
@@ -1007,7 +1007,7 @@ jobs:
)
# ── 2. OpenAI image_url (data URI base64) ───────────────────
- # 64x64 solid-red PNG. stb_image (used by Studio's image
+ # 64x64 solid-red PNG. stb_image (used by Unsloth's image
# normaliser at routes/inference.py:3410) rejects 4x4 or
# smaller PNGs as truncated, so we go up to 64x64 -- still
# tiny in token cost. The assertion is loose: any non-empty
@@ -1023,11 +1023,11 @@ jobs:
# The Mac prebuilt llama.cpp server has a known crash when
# processing image inputs alongside the gemma-4-E2B mmproj
# (server disconnects mid-completion). This is upstream
- # llama.cpp behaviour, not Studio. Wrap both SDK calls in
+ # llama.cpp behaviour, not Unsloth. Wrap both SDK calls in
# try/except so an upstream crash registers as a WARN rather
- # than failing the whole job. Studio's contract (OpenAI/
+ # than failing the whole job. Unsloth's contract (OpenAI/
# Anthropic image fields are accepted and forwarded) is
- # validated by the request body Studio constructs, not by
+ # validated by the request body Unsloth constructs, not by
# whether llama.cpp can decode it on Mac Metal.
client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY)
try:
@@ -1053,14 +1053,14 @@ jobs:
except Exception as exc:
print(
f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: "
- f"{exc}. Likely upstream llama.cpp Mac+vision crash, NOT a Studio "
- f"regression. Studio successfully forwarded the request."
+ f"{exc}. Likely upstream llama.cpp Mac+vision crash, NOT an Unsloth "
+ f"regression. Unsloth successfully forwarded the request."
)
# ── 3. Anthropic source/base64 image ────────────────────────
- # Two SDK quirks vs. Studio: base_url must NOT include /v1
+ # Two SDK quirks vs. Unsloth: base_url must NOT include /v1
# (the SDK appends it itself; otherwise /v1/v1/messages -> 405),
- # and Studio's auth is HTTPBearer-only so the SDK's default
+ # and Unsloth's auth is HTTPBearer-only so the SDK's default
# x-api-key header is ignored -- send Authorization: Bearer
# via default_headers.
anthropic = Anthropic(
@@ -1099,11 +1099,11 @@ jobs:
print(
f"[image/anthropic] WARN anthropic image SDK call raised: "
f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp Mac+vision "
- f"crash, NOT a Studio regression."
+ f"crash, NOT an Unsloth regression."
)
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml
index 362305cdd4..e990f752d4 100644
--- a/.github/workflows/studio-mac-install-matrix.yml
+++ b/.github/workflows/studio-mac-install-matrix.yml
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-# Proves Studio's llama.cpp install loads on every supported macOS. The heavy
+# Proves Unsloth's llama.cpp install loads on every supported macOS. The heavy
# app smokes stay single-OS; this matrix covers the OS-version dimension cheaply
# (install.sh + binary-load assert). Regression guard for the macOS-version
# selection in studio/install_llama_prebuilt.py.
@@ -60,7 +60,7 @@ jobs:
with:
python-version: '3.12'
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml
index 20ca247b9f..378e8ee5a6 100644
--- a/.github/workflows/studio-mac-ui-smoke.yml
+++ b/.github/workflows/studio-mac-ui-smoke.yml
@@ -83,7 +83,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -143,7 +143,7 @@ jobs:
print(f"pipeTransport.js: patched JSON.parse calls in {path}")
PY
- - name: Reset auth + boot Studio
+ - name: Reset auth + boot Unsloth
run: |
unsloth studio reset-password
mkdir -p logs
@@ -188,7 +188,7 @@ jobs:
# 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
+ # guard redirects mid-navigation. The retry FULLY resets Unsloth
# (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
@@ -209,7 +209,7 @@ jobs:
|| 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..."
+ echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..."
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
unsloth studio reset-password
@@ -238,13 +238,13 @@ jobs:
exit "$rc"
done
- - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders)
+ - name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders)
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- - name: Reset auth + boot Studio for extra UI tests (port 18897)
+ - name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -271,7 +271,7 @@ jobs:
echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV"
echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV"
- - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright
+ - name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright
env:
BASE_URL: http://127.0.0.1:18897
STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }}
@@ -300,7 +300,7 @@ jobs:
|| 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..."
+ echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..."
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
unsloth studio reset-password
@@ -327,7 +327,7 @@ jobs:
exit "$rc"
done
- - name: Stop second Studio
+ - name: Stop second Unsloth
if: always()
run: |
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml
index d104306c7e..fe9880f3ca 100644
--- a/.github/workflows/studio-mac-update-smoke.yml
+++ b/.github/workflows/studio-mac-update-smoke.yml
@@ -4,15 +4,15 @@
# Mac counterpart to studio-update-smoke.yml. Verifies that on a real
# Apple Silicon (macos-14, M1) runner:
#
-# 1. install.sh --local --no-torch installs Studio AND auto-fetches
+# 1. install.sh --local --no-torch installs Unsloth AND auto-fetches
# the prebuilt llama.cpp Mac binary (llama-bNNNN-bin-macos-arm64
# from ggml-org/llama.cpp). Hitting the source-build fallback is
-# treated as an Unsloth bug -- Studio must always pick the
+# treated as an Unsloth bug -- Unsloth must always pick the
# prebuilt on Mac.
# 2. unsloth studio update --local is idempotent. Two consecutive
# runs both report "prebuilt up to date and validated", no
# source-build fallback.
-# 3. The installed Studio still boots and /api/health returns
+# 3. The installed Unsloth still boots and /api/health returns
# healthy after the update path.
name: Mac Studio Update CI
@@ -42,7 +42,7 @@ permissions:
jobs:
update-idempotency:
- name: Studio Updating Tests
+ name: Unsloth Updating Tests
runs-on: macos-14
timeout-minutes: 30
steps:
@@ -59,7 +59,7 @@ jobs:
python-version: '3.12'
cache: 'pip'
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -106,7 +106,7 @@ jobs:
grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log
echo "second update was clean"
- - name: Boot Studio briefly to confirm the install is still usable
+ - name: Boot Unsloth briefly to confirm the install is still usable
run: |
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \
@@ -123,13 +123,13 @@ jobs:
sleep 1
done
if [ -z "$HEALTHY" ]; then
- echo "Studio failed to come up after \`update\`"
+ echo "Unsloth failed to come up after \`update\`"
tail -200 logs/studio.log
kill "$PID" 2>/dev/null || true
exit 1
fi
kill "$PID" 2>/dev/null || true
- echo "post-update Studio /api/health OK"
+ echo "post-update Unsloth /api/health OK"
- name: Uninstall and verify clean
# Round-trip through scripts/uninstall.sh on real macOS. As a side
diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml
index 018857de68..8e26b9fd0c 100644
--- a/.github/workflows/studio-tauri-smoke.yml
+++ b/.github/workflows/studio-tauri-smoke.yml
@@ -12,7 +12,7 @@
# stay in release-desktop.yml (manual `workflow_dispatch`) because they need
# code-signing secrets and ~30 min of runner time each.
-name: Studio Tauri CI
+name: Unsloth Tauri CI
on:
pull_request:
diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml
index 297a585430..b6d6d7d6e2 100644
--- a/.github/workflows/studio-ui-smoke.yml
+++ b/.github/workflows/studio-ui-smoke.yml
@@ -1,8 +1,8 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-# End-to-end Studio chat UI smoke via Playwright + Chromium against a
-# headless Linux runner. Boots Studio with the smallest GGUF
+# End-to-end Unsloth chat UI smoke via Playwright + Chromium against a
+# headless Linux runner. Boots Unsloth with the smallest GGUF
# (gemma-3-270m-it UD-Q4_K_XL, ~254 MiB), drives the actual frontend
# bundle, and asserts the full bootstrap-password / change-password /
# send-message / persist-on-reload journey works end to end.
@@ -14,7 +14,7 @@
# frontend-only CI happily pass while the actual user-visible UI is
# broken (cf. the 2026.5.1 chat-history release).
-name: Studio UI CI
+name: Unsloth UI CI
on:
pull_request:
@@ -97,7 +97,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@@ -115,7 +115,7 @@ jobs:
# warm runner.
python -m playwright install --with-deps chromium
- - name: Reset auth + boot Studio
+ - name: Reset auth + boot Unsloth
run: |
unsloth studio reset-password
mkdir -p logs
@@ -147,7 +147,7 @@ jobs:
# NEW + NEW2 are generated freshly per CI run via secrets.token_urlsafe
# rather than hardcoded. If a workflow gets compromised, the
# attacker can't replay a known-good rotated password against
- # any future / parallel Studio install -- the rotated value
+ # any future / parallel Unsloth install -- the rotated value
# only ever exists for the lifetime of this single job, masked
# in the log via ::add-mask::.
run: |
@@ -165,18 +165,18 @@ jobs:
env:
BASE_URL: http://127.0.0.1:18892
# The test file lives in the repo so it can be run locally
- # against a freshly-installed Studio (BASE_URL=...; STUDIO_OLD_PW=
+ # against a freshly-installed Unsloth (BASE_URL=...; STUDIO_OLD_PW=
# $(cat ~/.unsloth/studio/auth/.bootstrap_password); python ...).
PW_ART_DIR: logs/playwright
# Strict mode: in CI a missing button / nav / dialog must
# FAIL the test. Locally the test still runs against partial
- # Studio installs without STUDIO_UI_STRICT.
+ # Unsloth installs without STUDIO_UI_STRICT.
STUDIO_UI_STRICT: '1'
run: |
mkdir -p logs/playwright
python tests/studio/playwright_chat_ui.py
- - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders)
+ - name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders)
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@@ -184,10 +184,10 @@ jobs:
# The chat UI test ends by clicking the Shutdown menuitem, which
# leaves the server dead. The extra UI test (Compare / Recipes /
- # Export / Studio / Settings) needs a fresh Studio, so we boot a
+ # Export / Unsloth / Settings) needs a fresh Unsloth, so we boot a
# second one on a different port. Boot is fast (~3-5s on the
# warm install we already did) so this adds little wall time.
- - name: Reset auth + boot Studio for extra UI tests (port 18894)
+ - name: Reset auth + boot Unsloth for extra UI tests (port 18894)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -214,7 +214,7 @@ jobs:
echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV"
echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV"
- - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright
+ - name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright
env:
BASE_URL: http://127.0.0.1:18894
STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }}
@@ -227,16 +227,16 @@ jobs:
mkdir -p logs/playwright_extra
python tests/studio/playwright_extra_ui.py
- - name: Stop second Studio
+ - name: Stop second Unsloth
if: always()
run: |
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
# IME + multilingual paste regression (issue #5318 / PR #5327).
- # Third Studio on its own port so a hang here cannot poison the
+ # Third Unsloth on its own port so a hang here cannot poison the
# earlier UI tests. No GGUF -- the bug surface is the composer.
- - name: Reset auth + boot Studio for IME / i18n tests (port 18896)
+ - name: Reset auth + boot Unsloth for IME / i18n tests (port 18896)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -256,7 +256,7 @@ jobs:
- name: Pass bootstrap pw for IME / i18n test
# IME smoke does the change-password against the bootstrap that
- # Studio's frontend injects into the page, so it only needs the
+ # Unsloth's frontend injects into the page, so it only needs the
# NEW password.
run: |
NEW="CIIme-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
@@ -273,7 +273,7 @@ jobs:
mkdir -p logs/playwright_ime
python tests/studio/playwright_chat_ime_i18n.py
- - name: Stop third Studio
+ - name: Stop third Unsloth
if: always()
run: |
kill "${STUDIO_IME_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml
index 08a79afacd..625c2c7811 100644
--- a/.github/workflows/studio-update-smoke.yml
+++ b/.github/workflows/studio-update-smoke.yml
@@ -9,7 +9,7 @@
# This catches regressions in setup.sh's update path that the existing
# GGUF / wheel jobs would miss because they only invoke install.sh once.
-name: Studio Update CI
+name: Unsloth Update CI
on:
pull_request:
@@ -36,7 +36,7 @@ permissions:
jobs:
update-idempotency:
- name: Studio Updating Tests
+ name: Unsloth Updating Tests
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
@@ -63,7 +63,7 @@ jobs:
# post-step then fatal-errors with "Cache folder path is
# retrieved for pip but doesn't exist on disk".
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
# Pass the workflow token so the llama.cpp prebuilt installer's
# GitHub-API call to list releases isn't rate-limited (60/hr
# unauthenticated). Without this, three consecutive install +
@@ -122,7 +122,7 @@ jobs:
grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log
echo "second update was clean"
- - name: Boot Studio briefly to confirm the install is still usable
+ - name: Boot Unsloth briefly to confirm the install is still usable
# If `update --local` accidentally broke the venv or wiped the
# llama-server binary, the server would fail to start here.
run: |
@@ -138,13 +138,13 @@ jobs:
sleep 1
done
if ! jq -e '.status == "healthy"' /tmp/health.json 2>/dev/null; then
- echo "Studio failed to come up after `update`"
+ echo "Unsloth failed to come up after `update`"
tail -200 logs/studio.log
kill "$PID" 2>/dev/null || true
exit 1
fi
kill "$PID" 2>/dev/null || true
- echo "post-update Studio /api/health OK"
+ echo "post-update Unsloth /api/health OK"
- name: Uninstall and verify clean
# Round-trip the installer through scripts/uninstall.sh: confirms the
diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml
index e9abd2d669..6dbcceebbd 100644
--- a/.github/workflows/studio-windows-api-smoke.yml
+++ b/.github/workflows/studio-windows-api-smoke.yml
@@ -9,7 +9,7 @@
# (Section 6) is Linux-only and short-circuits on non-POSIX; the rest
# is platform-portable.
-name: Windows Studio API CI
+name: Windows Unsloth API CI
on:
pull_request:
@@ -34,7 +34,7 @@ permissions:
jobs:
api-smoke:
- name: Studio API & Auth Tests
+ name: Unsloth API & Auth Tests
runs-on: windows-latest
timeout-minutes: 30
defaults:
@@ -105,7 +105,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
- # rebuild" and Studio boots with an empty dist directory.
+ # rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@@ -121,7 +121,7 @@ jobs:
}
}
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -161,7 +161,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- - name: Add Studio shim to GITHUB_PATH
+ - name: Add Unsloth shim to GITHUB_PATH
# install.ps1's User-PATH update doesn't propagate to a
# running Git Bash session; export the shim dir so the
# next `unsloth ...` invocation finds it.
@@ -177,7 +177,7 @@ jobs:
- name: Install pyjwt for the JWT-expiry forge test
run: python -m pip install 'pyjwt>=2.6'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -207,7 +207,7 @@ jobs:
echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV"
echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV"
- - name: Run Studio API & Auth tests
+ - name: Run Unsloth API & Auth tests
# Do NOT pin STUDIO_AUTH_DIR here. The Mac/Linux mirrors
# hardcode runner-specific paths (/Users/runner/...,
# /home/runner/...), but on Windows the path is
@@ -219,7 +219,7 @@ jobs:
BASE_URL: http://127.0.0.1:18895
run: python tests/studio/studio_api_smoke.py
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml
index 63a7e9dc8f..3ebe442f52 100644
--- a/.github/workflows/studio-windows-inference-smoke.yml
+++ b/.github/workflows/studio-windows-inference-smoke.yml
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-# Three end-to-end smoke jobs that boot a freshly-installed Studio and
+# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and
# exercise the surfaces real users hit through the OpenAI / Anthropic
# SDKs and curl, on the FREE windows-latest runner. Each job picks the
# smallest model that exercises the behaviour under test, primes
@@ -16,7 +16,7 @@
# Qwen3-VL-2B-Instruct UD-IQ2_XXS + mmproj-F16 (~1.4 GiB total).
# Within the 14 GB windows-latest SSD budget.
-name: Windows Studio GGUF CI
+name: Windows Unsloth GGUF CI
on:
pull_request:
@@ -57,7 +57,7 @@ jobs:
STUDIO_PORT: '18888'
HF_HOME: ${{ github.workspace }}/hf-cache
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
- # download / Studio CLI print "✓" checkmarks and crash
+ # download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@@ -160,7 +160,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
- # rebuild" and Studio boots with an empty dist directory.
+ # rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@@ -176,7 +176,7 @@ jobs:
}
}
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -214,7 +214,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- - name: Add Studio shim to GITHUB_PATH
+ - name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@@ -227,7 +227,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: python -m pip install 'openai>=1.50' 'anthropic>=0.40'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -244,7 +244,7 @@ jobs:
fi
sleep 1
done
- echo "Studio did not become healthy in 180s"
+ echo "Unsloth did not become healthy in 180s"
tail -200 logs/studio.log
exit 1
@@ -281,7 +281,7 @@ jobs:
# Retry the load step a few times so a transient TCP RST during
# llama-server warm-up (Windows runner image churn,
# windows-latest -> windows-2025-vs2026 rollout) doesn't fail
- # the whole job. The Studio backend's _wait_for_health now
+ # the whole job. The Unsloth backend's _wait_for_health now
# catches httpx.ReadError too; this retry layer covers the
# cases the backend can't recover from on its own.
LOAD_OK=0
@@ -382,15 +382,15 @@ jobs:
print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded")
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
- # test run. The runner reclaims the Studio child process at
+ # test run. The runner reclaims the Unsloth child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
- run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
+ run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
if: always()
@@ -398,10 +398,10 @@ jobs:
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
- # Copy llama-server's own stdout/stderr (teed by Studio under
+ # Copy llama-server's own stdout/stderr (teed by Unsloth under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
# upload-artifact can pick it up. Crucial for diagnosing a
- # subprocess crash where Studio's traceback only shows the
+ # subprocess crash where Unsloth's traceback only shows the
# symptom (httpx ReadError) but not the cause.
run: |
mkdir -p logs/llama-server
@@ -439,14 +439,14 @@ jobs:
# (211 s on first run; subsequent runs hit the cache, but the
# one-time cost recurs every time the cache key bumps). Use
# main's `--local-dir gguf-cache` pattern: cache the flat .gguf
- # only, pass an absolute path to Studio's /api/inference/load.
+ # only, pass an absolute path to Unsloth's /api/inference/load.
# The OpenAI/Anth and JSON+images jobs still cover the
# gguf_variant resolution path.
GGUF_REPO: unsloth/Qwen3.5-2B-GGUF
GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf
STUDIO_PORT: '18898'
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
- # download / Studio CLI print "✓" checkmarks and crash
+ # download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@@ -507,7 +507,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
- # rebuild" and Studio boots with an empty dist directory.
+ # rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@@ -523,7 +523,7 @@ jobs:
}
}
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -561,7 +561,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- - name: Add Studio shim to GITHUB_PATH
+ - name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@@ -571,7 +571,7 @@ jobs:
fi
cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH"
- - name: Reset auth + boot Studio (API-only, default tool policy)
+ - name: Reset auth + boot Unsloth (API-only, default tool policy)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -607,7 +607,7 @@ jobs:
# raw string, but we cannot embed `\a` etc. in JSON without
# JSON-string-escaping every backslash. Replace `\` with `/`
# via bash parameter expansion -- pathlib.Path on Windows
- # accepts forward slashes natively, so Studio's loader sees
+ # accepts forward slashes natively, so Unsloth's loader sees
# a normal path.
GGUF_PATH="${GITHUB_WORKSPACE//\\//}/gguf-cache/${GGUF_FILE}"
ls -lh "$GGUF_PATH"
@@ -680,7 +680,7 @@ jobs:
def post_sse(path, body, *, timeout = 600, retries = 1, soft = False):
# The server-side agentic loop always answers over SSE. A
# shared CI runner can stall the stream transport (the
- # connection opening, or a mid-stream read) even when Studio
+ # connection opening, or a mid-stream read) even when Unsloth
# is healthy, so harden the read three ways:
# * retry a transport stall once with a fresh request,
# capped at 300s (a healthy server answers a retry
@@ -882,15 +882,15 @@ jobs:
print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)")
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
- # test run. The runner reclaims the Studio child process at
+ # test run. The runner reclaims the Unsloth child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
- run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
+ run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
if: always()
@@ -898,10 +898,10 @@ jobs:
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
- # Copy llama-server's own stdout/stderr (teed by Studio under
+ # Copy llama-server's own stdout/stderr (teed by Unsloth under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
# upload-artifact can pick it up. Crucial for diagnosing a
- # subprocess crash where Studio's traceback only shows the
+ # subprocess crash where Unsloth's traceback only shows the
# symptom (httpx ReadError) but not the cause.
run: |
mkdir -p logs/llama-server
@@ -939,7 +939,7 @@ jobs:
STUDIO_PORT: '18899'
HF_HOME: ${{ github.workspace }}/hf-cache
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
- # download / Studio CLI print "✓" checkmarks and crash
+ # download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@@ -1005,7 +1005,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
- # rebuild" and Studio boots with an empty dist directory.
+ # rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@@ -1021,7 +1021,7 @@ jobs:
}
}
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -1059,7 +1059,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- - name: Add Studio shim to GITHUB_PATH
+ - name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@@ -1072,7 +1072,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: python -m pip install 'openai>=1.50' 'anthropic>=0.40'
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -1262,7 +1262,7 @@ jobs:
except Exception as exc:
print(
f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: "
- f"{exc}. Studio successfully forwarded the request; failure here is "
+ f"{exc}. Unsloth successfully forwarded the request; failure here is "
f"upstream llama.cpp vision behaviour."
)
@@ -1303,19 +1303,19 @@ jobs:
print(
f"[image/anthropic] WARN anthropic image SDK call raised: "
f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp vision "
- f"behaviour, NOT a Studio regression."
+ f"behaviour, NOT an Unsloth regression."
)
PY
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
- # test run. The runner reclaims the Studio child process at
+ # test run. The runner reclaims the Unsloth child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
- run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
+ run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
if: always()
@@ -1323,10 +1323,10 @@ jobs:
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
- # Copy llama-server's own stdout/stderr (teed by Studio under
+ # Copy llama-server's own stdout/stderr (teed by Unsloth under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
# upload-artifact can pick it up. Crucial for diagnosing a
- # subprocess crash where Studio's traceback only shows the
+ # subprocess crash where Unsloth's traceback only shows the
# symptom (httpx ReadError) but not the cause.
run: |
mkdir -p logs/llama-server
@@ -1348,7 +1348,7 @@ jobs:
# ── folded from studio-windows-no-vs-smoke.yml: install + run with no Visual Studio ──
no-vs-cpu:
- name: Studio install + inference without Visual Studio
+ name: Unsloth install + inference without Visual Studio
runs-on: windows-latest
timeout-minutes: 35
defaults:
@@ -1502,7 +1502,7 @@ jobs:
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
+ - name: Install Unsloth (--local, --no-torch) with no build tools present
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -1538,13 +1538,13 @@ jobs:
echo "Prebuilt installed with no build tools:"
cat "$INFO"
- - name: Add Studio shim to GITHUB_PATH
+ - name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
[ -f "$SHIM_DIR/unsloth.exe" ] || { echo "::error::unsloth.exe shim not found"; ls -la ~/.unsloth/studio/ || true; exit 1; }
cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH"
- - name: Reset auth + boot Studio (API-only)
+ - name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -1613,10 +1613,10 @@ jobs:
}
Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue
- - name: Stop Studio
+ - name: Stop Unsloth
if: always()
shell: cmd
- run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
+ run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
if: always()
diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml
index 405309916a..12d7475b53 100644
--- a/.github/workflows/studio-windows-ui-smoke.yml
+++ b/.github/workflows/studio-windows-ui-smoke.yml
@@ -4,11 +4,11 @@
# Windows counterpart to studio-ui-smoke.yml / studio-mac-ui-smoke.yml.
# Same Playwright + Chromium end-to-end chat UI flow + extra UI flow,
# but on the FREE windows-latest runner so we catch Windows-specific
-# regressions in the install path (install.ps1), the Studio CLI's
+# regressions in the install path (install.ps1), the Unsloth CLI's
# Windows process-management branches, and the llama.cpp prebuilt's
# Windows HTTP layer.
-name: Windows Studio UI CI
+name: Windows Unsloth UI CI
on:
pull_request:
@@ -49,7 +49,7 @@ jobs:
GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf
STUDIO_PORT: '18896'
HF_HOME: ${{ github.workspace }}/hf-cache
- # Force UTF-8 for stdio so Python tools (hf download, Studio
+ # Force UTF-8 for stdio so Python tools (hf download, Unsloth
# CLI, etc.) can print Unicode characters like the success
# checkmark "✓". Windows defaults to cp1252 / charmap and
# any tool that prints "OK ✓" hits a UnicodeEncodeError.
@@ -121,7 +121,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
- # rebuild" and Studio boots with an empty dist directory.
+ # rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@@ -148,7 +148,7 @@ jobs:
Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode
Write-Host "seeded legacy launch-studio.vbs at $appDir"
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
# install.ps1 is the supported Windows installer. install.sh
# has no Windows branch (apt-get / brew calls). The PS1
# script's `Install-UnslothStudio @args` line at the bottom
@@ -205,7 +205,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- - name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut)
+ - name: Assert Unsloth launcher chain (no VBS, hidden PowerShell shortcut)
# The shortcut launch path is otherwise untested here (the steps below
# boot `unsloth studio` directly). Guard against re-introducing the VBS
# that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk
@@ -234,7 +234,7 @@ jobs:
}
Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)"
- - name: Launch Studio via the shortcut and assert health
+ - name: Launch Unsloth via the shortcut and assert health
# Run the exact command the .lnk stores (hidden PowerShell over
# launch-studio.ps1) and confirm it brings the backend up. This is the
# only step that proves the shortcut launch is not silently broken.
@@ -265,10 +265,10 @@ jobs:
$owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess
if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null }
} catch {}
- if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" }
- Write-Host "Studio healthy on port $foundPort (launched via the shortcut)"
+ if (-not $foundPort) { throw "Unsloth did not become healthy when launched via the shortcut" }
+ Write-Host "Unsloth healthy on port $foundPort (launched via the shortcut)"
- - name: Add Studio shim to GITHUB_PATH
+ - name: Add Unsloth shim to GITHUB_PATH
# install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe
# and adds that dir to the User PATH via the Windows registry.
# Registry-level PATH updates don't propagate to a running
@@ -284,7 +284,7 @@ jobs:
fi
# GITHUB_PATH wants Windows-style paths; convert via cygpath.
cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH"
- echo "Added Studio shim dir to PATH: $(cygpath -w "$SHIM_DIR")"
+ echo "Added Unsloth shim dir to PATH: $(cygpath -w "$SHIM_DIR")"
- name: Install Playwright + Chromium
# No --with-deps on Windows: that flag installs Linux apt
@@ -294,7 +294,7 @@ jobs:
python -m pip install 'playwright>=1.45'
python -m playwright install chromium
- - name: Reset auth + boot Studio
+ - name: Reset auth + boot Unsloth
run: |
unsloth studio reset-password
mkdir -p logs
@@ -339,13 +339,13 @@ jobs:
mkdir -p logs/playwright
python tests/studio/playwright_chat_ui.py
- - name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders)
+ - name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders)
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- - name: Reset auth + boot Studio for extra UI tests (port 18897)
+ - name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
mkdir -p logs
@@ -372,7 +372,7 @@ jobs:
echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV"
echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV"
- - name: Drive Compare/Recipes/Export/Studio/Settings with Playwright
+ - name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright
env:
BASE_URL: http://127.0.0.1:18897
STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }}
@@ -386,7 +386,7 @@ jobs:
mkdir -p logs/playwright_extra
python tests/studio/playwright_extra_ui.py
- - name: Stop second Studio
+ - name: Stop second Unsloth
if: always()
run: |
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml
index 5b92f1a3e0..42d74d47d2 100644
--- a/.github/workflows/studio-windows-update-smoke.yml
+++ b/.github/workflows/studio-windows-update-smoke.yml
@@ -5,19 +5,19 @@
# studio-mac-update-smoke.yml. Verifies that on the FREE
# windows-latest runner:
#
-# 1. install.ps1 --local --no-torch installs Studio AND auto-fetches
+# 1. install.ps1 --local --no-torch installs Unsloth AND auto-fetches
# the prebuilt llama.cpp Windows binary (app--windows-x64-cpu
# from unslothai/llama.cpp). Hitting the source-build fallback is
-# treated as an Unsloth bug -- Studio must always pick the
+# treated as an Unsloth bug -- Unsloth must always pick the
# prebuilt on Windows.
# 2. unsloth studio update --local is idempotent. Two consecutive
# runs both report "prebuilt up to date and validated", no
# source-build fallback. The CLI's _find_setup_script picks
# setup.ps1 on Windows automatically.
-# 3. The installed Studio still boots and /api/health returns
+# 3. The installed Unsloth still boots and /api/health returns
# healthy after the update path.
-name: Windows Studio Update CI
+name: Windows Unsloth Update CI
on:
pull_request:
@@ -45,7 +45,7 @@ permissions:
jobs:
update-idempotency:
- name: Studio Updating Tests
+ name: Unsloth Updating Tests
runs-on: windows-latest
timeout-minutes: 30
defaults:
@@ -53,7 +53,7 @@ jobs:
shell: bash
env:
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
- # download / Studio CLI print "✓" checkmarks and crash
+ # download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@@ -90,7 +90,7 @@ jobs:
# reuses the existing Node with no download.
#
# (2) Defender. windows-latest's real-time scan opens / hashes
- # every file Studio writes during install (Vite output =
+ # every file Unsloth writes during install (Vite output =
# thousands of small chunks, uv pip = wheel-extraction =
# thousands of small files). The latency dominates the
# 200 s frontend build and the 90 s deps install. Adding
@@ -109,7 +109,7 @@ jobs:
# setup.ps1 line 1281-1296's mtime-based "is the frontend
# stale?" check into "up to date, skip rebuild", because the
# newly-created dist's mtime is younger than every source
- # file. Studio then boots with an empty dist and 500s on
+ # file. Unsloth then boots with an empty dist and 500s on
# GET / with FileNotFoundError: dist\index.html. See run
# 25546676715 / job 74984469728.
# Add-MpPreference accepts paths that do not yet exist; the
@@ -129,7 +129,7 @@ jobs:
}
}
- - name: Install Studio (--local, --no-torch)
+ - name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -168,7 +168,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- - name: Add Studio shim to GITHUB_PATH
+ - name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@@ -212,7 +212,7 @@ jobs:
grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log
echo "second update was clean"
- - name: Boot Studio briefly to confirm the install is still usable
+ - name: Boot Unsloth briefly to confirm the install is still usable
run: |
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \
@@ -239,13 +239,13 @@ jobs:
sleep 1
done
if [ -z "$HEALTHY" ]; then
- echo "Studio failed to come up after \`update\`"
+ echo "Unsloth failed to come up after \`update\`"
tail -200 logs/studio.log
kill "$PID" 2>/dev/null || true
exit 1
fi
kill "$PID" 2>/dev/null || true
- echo "post-update Studio /api/health OK"
+ echo "post-update Unsloth /api/health OK"
- name: Uninstall and verify clean
# Round-trip through scripts/uninstall.ps1 against the default
diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml
index 3de3c33ca2..cdad617027 100644
--- a/.github/workflows/wheel-smoke.yml
+++ b/.github/workflows/wheel-smoke.yml
@@ -3,7 +3,7 @@
# Builds the PyPI wheel from the PR branch, then verifies the built wheel
# actually contains what we expect to ship and does NOT contain the broken
-# Studio bundle that 2026.5.1 published. This is the single workflow that
+# Unsloth bundle that 2026.5.1 published. This is the single workflow that
# would have blocked the 2026.5.1 release before twine upload.
#
# Verified locally end-to-end against this branch:
@@ -12,7 +12,7 @@
# lockfile shipped, frontend dist shipped,
# no node_modules in wheel, no bun.lock in wheel,
# main bundle has unstable_Provider hits=1 (assistant-ui internals only).
-# - Studio backend imports cleanly from the installed wheel with the
+# - Unsloth backend imports cleanly from the installed wheel with the
# lightweight dep set below.
name: Wheel CI
@@ -101,7 +101,7 @@ jobs:
hits = data.count("unstable_Provider:")
print(f"main bundle: {js[0]}")
print(f"unstable_Provider hits: {hits} (>=4 indicates 2026.5.1 regression)")
- checks["bundle has no Studio unstable_Provider call site"] = (hits < 4)
+ checks["bundle has no Unsloth unstable_Provider call site"] = (hits < 4)
print()
for k, v in checks.items():
@@ -109,7 +109,7 @@ jobs:
sys.exit(0 if all(checks.values()) else 1)
PY
- - name: Studio backend import smoke
+ - name: Unsloth backend import smoke
# Imports `studio.backend.main:app` from the freshly-installed wheel in
# a clean venv. This catches the class of bug that 2026.5.1 shipped with:
# frontend dist missing, package-lock.json missing, or the wheel's Python
@@ -125,7 +125,7 @@ jobs:
/tmp/v/bin/pip install --no-deps dist/unsloth-*.whl
# Run from /tmp so Python imports the installed package, not the source tree.
cd /tmp
- /tmp/v/bin/python -c "from studio.backend.main import app; print('Studio backend OK:', app.title)"
+ /tmp/v/bin/python -c "from studio.backend.main import app; print('Unsloth backend OK:', app.title)"
- name: Upload wheel on failure
if: failure()
diff --git a/README.md b/README.md
index ef45b91430..085c7718e5 100644
--- a/README.md
+++ b/README.md
@@ -65,7 +65,7 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
* **CPU:** Supported for Chat and Data Recipes currently
* **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more
* **macOS:** Training, MLX and GGUF inference are ALL supported.
-* **AMD:** Chat + Data works. Train with [Unsloth Core](#unsloth-core-code-based). Studio support is out soon.
+* **AMD:** Chat + Data works. Train with [Unsloth Core](#unsloth-core-code-based). Unsloth Studio support is out soon.
* **Multi-GPU:** Available now, with a major upgrade on the way
#### macOS, Linux, WSL:
@@ -86,7 +86,7 @@ unsloth studio -p 8888
```
For LAN or cloud access, add `-H 0.0.0.0` (raw port only; add `--cloudflare` for a public URL). By default, Unsloth is accessible only locally.
-To reach Studio over HTTPS, use `unsloth studio --secure`. Studio stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Studio reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below).
+To reach Unsloth over HTTPS, use `unsloth studio --secure`. Unsloth stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Unsloth reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below).
#### Docker
Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run:
@@ -208,7 +208,7 @@ unsloth studio -p 8888
#### Remote access: `--secure` (HTTPS tunnel) vs raw port
By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach it from another device, pick one of:
-- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Studio stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed.
+- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Unsloth stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed.
```bash
unsloth studio --secure -p 8888
```
@@ -218,7 +218,7 @@ unsloth studio -H 0.0.0.0 -p 8888
```
The Cloudflare tunnel is **off by default**: `-H 0.0.0.0` exposes the raw port only, not a public internet URL. Pair the wildcard bind with `--cloudflare` (`unsloth studio -H 0.0.0.0 --cloudflare`) to also publish a public `https://*.trycloudflare.com` link, or prefer `--secure` (above), which keeps the raw port private. `--cloudflare` has no effect on a loopback bind.
-The first time Studio is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Studio shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI.
+The first time Unsloth is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Unsloth shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI.
For headless setups that cannot answer that prompt, set the initial admin password non-interactively with `--password` (only takes effect when no password is set yet; if one already exists it is a hard error, so rotate later with `unsloth studio reset-password`):
@@ -230,7 +230,7 @@ printf '%s\n' 'your-strong-password' | unsloth studio --secure --password - #
A literal `--password VALUE` is visible in the process list and shell history, so prefer the `UNSLOTH_STUDIO_PASSWORD` env var or `--password -` (stdin) for automation. This applies to any launch (public or a headless `-H 0.0.0.0` bind), and the password is set in the parent before the server binds, so it never reaches a re-executed child process.
-Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Studio.
+Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Unsloth.
#### Advanced launch options
Installer options can be passed as environment variables. On macOS, Linux and WSL place the variable after the pipe so the shell passes it to `sh`; on Windows set it with `$env:` before piping to `iex`.
@@ -243,7 +243,7 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh
$env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex
```
-Skip the post-install prompt that starts Studio (useful for automated installs):
+Skip the post-install prompt that starts Unsloth (useful for automated installs):
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh
```
@@ -279,9 +279,9 @@ UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh -
```powershell
$env:UNSLOTH_NPM_REGISTRY='https://artifactory.example.com/api/npm/npm/'; .\install.ps1 --local
```
-It is threaded as `--registry` into the Studio frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force.
+It is threaded as `--registry` into the Unsloth frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force.
-Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`.
+Cap Unsloth's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`.
#### Uninstall
The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows):
diff --git a/build.sh b/build.sh
index dc272f0de1..2a836e19d9 100644
--- a/build.sh
+++ b/build.sh
@@ -4,9 +4,9 @@
set -euo pipefail
-# PyPI/Studio release publishing must use `./build.sh publish` (or an
-# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Studio
-# artifacts include the display-only Studio release version.
+# PyPI/Unsloth release publishing must use `./build.sh publish` (or an
+# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Unsloth
+# artifacts include the display-only Unsloth release version.
# 1. Build frontend (Vite outputs to dist/)
cd studio/frontend
@@ -87,7 +87,7 @@ cd ../..
# 2. Clean old artifacts
rm -rf build dist *.egg-info
-# 3. Stamp display-only Studio release metadata for packaged builds.
+# 3. Stamp display-only Unsloth release metadata for packaged builds.
_STUDIO_BUILD_INFO="studio/backend/utils/_studio_release_build.py"
_STUDIO_BUILD_INFO_BACKUP="$(mktemp)"
cp "$_STUDIO_BUILD_INFO" "$_STUDIO_BUILD_INFO_BACKUP"
diff --git a/install.ps1 b/install.ps1
index c25d7e7b7f..df49414620 100644
--- a/install.ps1
+++ b/install.ps1
@@ -176,7 +176,7 @@ function Install-UnslothStudio {
$envOverride = $env:STUDIO_HOME.Trim()
}
- # Custom Studio roots are not supported with --tauri (desktop app still
+ # Custom Unsloth roots are not supported with --tauri (desktop app still
# resolves %USERPROFILE%\.unsloth\studio). Pass through if override == legacy.
if ($TauriMode -and $envOverride) {
$_tauriOverride = $envOverride
@@ -756,7 +756,7 @@ function Find-FreeLaunchPort {
return `$null
}
-# If Studio is already healthy on any expected port, just open it and exit.
+# If Unsloth is already healthy on any expected port, just open it and exit.
`$existingPort = Find-HealthyStudioPort
if (`$existingPort) {
Start-Process "http://localhost:`$existingPort"
@@ -772,7 +772,7 @@ try {
`$haveMutex = `$true
}
if (-not `$haveMutex) {
- # Another launcher is already running; wait for it to bring Studio up
+ # Another launcher is already running; wait for it to bring Unsloth up
`$deadline = (Get-Date).AddSeconds(`$timeoutSec)
while ((Get-Date) -lt `$deadline) {
`$port = Find-HealthyStudioPort
@@ -1438,7 +1438,7 @@ exit 0
if (Test-Path -LiteralPath $VenvPython) {
# why: matching guard to the .venv branch below -- in env-mode
# $StudioHome is a user-chosen workspace, so refuse to nuke an
- # existing $StudioHome\unsloth_studio that lacks Studio sentinels.
+ # existing $StudioHome\unsloth_studio that lacks Unsloth sentinels.
# -PathType Leaf rejects a directory at the sentinel path. Accept the
# in-VENV ownership marker so partial-install retries are not blocked.
if (
@@ -1449,7 +1449,7 @@ exit 0
) {
Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red
Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow
- throw "Refusing to delete non-Studio venv at $VenvDir"
+ throw "Refusing to delete non-Unsloth venv at $VenvDir"
}
# New layout already exists -- replace only after preserving rollback copy.
substep "preserving existing environment for rollback..."
@@ -1468,7 +1468,7 @@ exit 0
# workspace root (e.g. user's existing project Python venv).
$OldVenv = Join-Path $StudioHome ".venv"
$OldPy = Join-Path $OldVenv "Scripts\python.exe"
- substep "found legacy Studio environment, validating..."
+ substep "found legacy Unsloth environment, validating..."
$prevEAP2 = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
@@ -1498,7 +1498,7 @@ exit 0
# Skip in env-mode so we don't relocate the default-install venv into
# the workspace root.
$CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio"
- substep "found CWD-relative Studio environment, migrating to $VenvDir..."
+ substep "found CWD-relative Unsloth environment, migrating to $VenvDir..."
Move-Item -LiteralPath $CwdVenv -Destination $VenvDir -Force
substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio"
$_Migrated = $true
@@ -1517,7 +1517,7 @@ exit 0
substep "$VenvDir"
}
- # Mark the freshly-created venv as Studio-owned so a partial install can be
+ # Mark the freshly-created venv as Unsloth-owned so a partial install can be
# repaired by re-running install.ps1; the env-mode deletion guard above
# accepts this marker as the primary sentinel.
if (Test-Path -LiteralPath $VenvDir -PathType Container) {
@@ -1526,7 +1526,7 @@ exit 0
# ── Helper: run amd-smi without triggering a UAC elevation prompt ──
# amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing
- # DiskPart UAC prompt mid-install (Studio backend amd.py hits the same).
+ # DiskPart UAC prompt mid-install (Unsloth backend amd.py hits the same).
# __COMPAT_LAYER=RunAsInvoker forces it (and helpers it spawns) to run
# un-elevated; on failure the WMI name -> gfx fallback still resolves the arch.
function Invoke-AmdSmiNoElevate {
@@ -1653,7 +1653,7 @@ exit 0
function Test-HipinfoIsVenvInternal {
param([AllowNull()][string]$HipinfoPath)
if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false }
- # Also derive the venv from the setup python + default Studio home, so
+ # Also derive the venv from the setup python + default Unsloth home, so
# the venv hipInfo is caught when VenvDir/VIRTUAL_ENV are unset.
$venvRoots = @()
if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV }
@@ -1663,7 +1663,7 @@ exit 0
try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {}
}
if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") }
- # A custom Studio home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the
+ # A custom Unsloth home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the
# venv off the default path; seed it too or its hipInfo escapes the filter.
$studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null }
if ($studioHomeEnv) {
@@ -1942,7 +1942,7 @@ exit 0
substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow"
substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} elseif ($ROCmGfxArch) {
- # Known arch: Studio setup installs AMD's bundled-runtime ROCm PyTorch wheels
+ # Known arch: Unsloth setup installs AMD's bundled-runtime ROCm PyTorch wheels
# (repo.amd.com), which ship their own runtime -- HIP SDK optional.
step "gpu" "AMD ROCm ($ROCmGfxArch)" "Cyan"
substep "Detected: $ROCmGpuLabel" "Cyan"
@@ -2219,8 +2219,8 @@ exit 0
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
# Transient AMD-index failure: fall back to a CPU base so the install
- # still completes; Studio setup retries ROCm afterwards.
- substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Studio setup retries ROCm." "Yellow"
+ # still completes; Unsloth setup retries ROCm afterwards.
+ substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow"
# --force-reinstall: a failed ROCm install can leave an unpinned ROCm
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
@@ -2422,7 +2422,7 @@ exit 0
Write-TauriLog "ERROR" "unsloth CLI was not installed correctly"
Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red
Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow
- Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow
+ Write-Host " This usually means an older unsloth version was installed that does not include the Unsloth CLI." -ForegroundColor Yellow
Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow
return (Exit-InstallFailure "unsloth CLI was not installed correctly")
}
@@ -2533,7 +2533,7 @@ exit 0
Write-Host " Move or remove it manually, then re-run the installer." -ForegroundColor Yellow
throw "Cannot create unsloth launcher: $ShimExe is a directory."
}
- # try/catch: if unsloth.exe is locked (Studio running), keep the old shim.
+ # try/catch: if unsloth.exe is locked (Unsloth running), keep the old shim.
$shimUpdated = $false
try {
if (Test-Path -LiteralPath $ShimExe) { Remove-Item -LiteralPath $ShimExe -Force -ErrorAction Stop }
@@ -2551,7 +2551,7 @@ exit 0
if (Test-Path -LiteralPath $ShimExe) {
Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow
Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow
- Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow
+ Write-Host " Close Unsloth and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow
Write-Host " Continuing with the existing launcher." -ForegroundColor Yellow
} else {
Write-Host "[WARN] Could not create unsloth launcher at $ShimExe" -ForegroundColor Yellow
@@ -2616,7 +2616,7 @@ exit 0
# Diagnostic only; never block install on a probe failure.
}
- # In interactive terminals, ask the user before starting Studio unless the
+ # In interactive terminals, ask the user before starting Unsloth unless the
# caller explicitly disabled the post-install prompt.
# In non-interactive environments (CI, Docker) just print instructions.
$IsInteractive = (-not $SkipAutostart) -and [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
diff --git a/install.sh b/install.sh
index 4a3c4471fa..5972379d26 100755
--- a/install.sh
+++ b/install.sh
@@ -97,7 +97,7 @@ if [ "$_VERBOSE" = true ]; then
export UNSLOTH_VERBOSE=1
fi
-# Custom Studio roots are not supported with --tauri (desktop app still
+# Custom Unsloth roots are not supported with --tauri (desktop app still
# resolves ~/.unsloth/studio). Pass through if the override == legacy default.
if [ "$TAURI_MODE" = true ]; then
_tauri_override_var=""
@@ -663,7 +663,7 @@ POLL_INTERVAL_SEC=0.25
LOG_FILE="$DATA_DIR/studio.log"
# why: in env-override mode multiple installs share an OS user; namespace the
# lock and remember our own healthy port so we never attach to an unrelated
-# Studio listening on the global 8888..8908 range.
+# Unsloth listening on the global 8888..8908 range.
LOCK_DIR="${XDG_RUNTIME_DIR:-/tmp}/unsloth-studio-launcher-$(id -u).lock"
PORT_FILE=""
# why: gate on the install-time mode (baked above) instead of the runtime env
@@ -734,7 +734,7 @@ _candidate_ports() {
_find_healthy_port() {
if [ -n "$PORT_FILE" ] && [ -f "$PORT_FILE" ]; then
# why: env-mode installs only attach to a port we previously launched
- # ourselves; never to a sibling Studio that happens to be healthy.
+ # ourselves; never to a sibling Unsloth that happens to be healthy.
_p=$(cat "$PORT_FILE" 2>/dev/null || true)
case "$_p" in
''|*[!0-9]*) ;;
@@ -901,7 +901,7 @@ _acquire_lock() {
# Lock dir exists -- check if owner is still alive
_old_pid=$(cat "$LOCK_DIR/pid" 2>/dev/null || true)
if [ -n "$_old_pid" ] && kill -0 "$_old_pid" 2>/dev/null; then
- # Another launcher is running; wait for it to bring Studio up
+ # Another launcher is running; wait for it to bring Unsloth up
_deadline=$(($(date +%s) + TIMEOUT_SEC))
while [ "$(date +%s)" -lt "$_deadline" ]; do
_port=$(_find_healthy_port) && {
@@ -1371,7 +1371,7 @@ WSLPS1_EOF
# shortcut wasn't created; tell the user how to launch / re-enable it.
if [ "$_css_created" -ne 1 ]; then
substep "Couldn't create the Windows shortcut (WSL interop may be disabled)." "$C_WARN"
- substep " Launch Studio from Windows: wsl -d \"$_css_distro\" -- bash -lc 'unsloth studio'" "$C_WARN"
+ substep " Launch Unsloth from Windows: wsl -d \"$_css_distro\" -- bash -lc 'unsloth studio'" "$C_WARN"
substep " (re-enable shortcuts: turn WSL interop back on, e.g. run 'wsl --shutdown' then reopen WSL.)" "$C_WARN"
fi
fi
@@ -1439,7 +1439,7 @@ if [ "$MAC_INTEL" = true ]; then
echo ""
echo " NOTE: Intel Mac (x86_64) detected."
echo " PyTorch is unavailable for this platform (dropped Jan 2024)."
- echo " Studio will install in GGUF-only mode."
+ echo " Unsloth will install in GGUF-only mode."
echo " Chat, inference via GGUF, and data recipes will work."
echo " Training requires Apple Silicon or Linux with GPU."
echo ""
@@ -1671,7 +1671,7 @@ _maybe_reroute_strixhalo_to_2404() {
_maybe_reroute_strixhalo_to_2404 || true
# ── Check system dependencies ──
-# cmake/git are only needed to *build* llama.cpp from source. Studio downloads a
+# cmake/git are only needed to *build* llama.cpp from source. Unsloth downloads a
# prebuilt by default, and setup.sh self-skips the source build when they're
# absent -- so macOS doesn't block on cmake (requiring it would force a manual
# Homebrew install). Linux keeps requiring them; its package manager has them.
@@ -1825,7 +1825,7 @@ _MIGRATED=false
if [ -x "$VENV_DIR/bin/python" ]; then
# why: matching guard to the .venv branch below -- in env-mode
# $STUDIO_HOME is a user-chosen workspace, so refuse to nuke an
- # existing $STUDIO_HOME/unsloth_studio that lacks Studio sentinels.
+ # existing $STUDIO_HOME/unsloth_studio that lacks Unsloth sentinels.
# Accept the in-VENV ownership marker so partial-install retries are
# not blocked. Sentinels must be regular files: -f follows symlinks
# to files (the legitimate ln -s shim shape) but rejects directories
@@ -1846,7 +1846,7 @@ elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/pytho
# Skip in env-mode so we don't rm -rf an unrelated .venv at the
# workspace root (e.g. user's existing project Python venv).
# In no-torch mode, a missing torch package is expected; validate Python only.
- substep "found legacy Studio environment, validating..."
+ substep "found legacy Unsloth environment, validating..."
_legacy_ok=false
if [ "$SKIP_TORCH" = true ]; then
if "$STUDIO_HOME/.venv/bin/python" -c "import sys; print(sys.executable)" >/dev/null 2>&1; then
@@ -1903,7 +1903,7 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then
fi
fi
-# Mark the freshly-created venv as Studio-owned so a partial install can be
+# Mark the freshly-created venv as Unsloth-owned so a partial install can be
# repaired by re-running install.sh; the env-mode deletion guard above accepts
# this marker as the primary sentinel.
if [ -x "$VENV_DIR/bin/python" ]; then
@@ -2335,7 +2335,7 @@ _pick_radeon_wheel() {
# the installer -- always returns 0. Runs the idempotent helper (ROCm 7.2 +
# librocdxg), then sources the env it persisted so detection finds the GPU.
# Export the ROCm-on-WSL env into this process and persist it to /etc/profile.d
-# so non-login Studio/llama launches inherit it. Idempotent (writes only when
+# so non-login Unsloth/llama launches inherit it. Idempotent (writes only when
# the drop-in is missing); no-op without librocdxg, so never fires off WSL.
# /etc/profile.d is root-owned -- sudo-tee when not root, else ROCm vanishes
# after this shell on a non-root reinstall. Best-effort either way.
@@ -2380,7 +2380,7 @@ _maybe_bootstrap_rocm_wsl() {
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then
# rocminfo may work only via the transient env _ensure_rocm_probe_env
# just set, which dies with the installer. Persist the drop-in so login
- # shells (Studio, llama.cpp) inherit it -- else a reinstall over an
+ # shells (Unsloth, llama.cpp) inherit it -- else a reinstall over an
# existing /opt/rocm (uninstall keeps ROCm but drops it) loses the GPU.
_persist_rocm_wsl_dropin
return 0
@@ -2402,7 +2402,7 @@ _maybe_bootstrap_rocm_wsl() {
# shellcheck disable=SC1091
. /etc/profile.d/unsloth-rocm-wsl.sh || true
else
- # librocdxg present but the env drop-in is gone (e.g. a Studio
+ # librocdxg present but the env drop-in is gone (e.g. an Unsloth
# uninstall removed it while keeping shared ROCm). Restore the env.
_persist_rocm_wsl_dropin
fi
@@ -3033,7 +3033,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
fi
# ── Run studio setup ──
-tauri_log "STEP" "Running Studio setup"
+tauri_log "STEP" "Running Unsloth setup"
# When --local, use the repo's own setup.sh directly.
# Otherwise, find it inside the installed package.
SETUP_SH=""
@@ -3227,7 +3227,7 @@ printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!"
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
echo ""
-# In interactive terminals, ask the user before starting Studio unless the
+# In interactive terminals, ask the user before starting Unsloth unless the
# caller explicitly disabled the post-install prompt.
# In non-interactive environments (Docker, CI, cloud-init) just print instructions.
if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then
diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh
index aa560fc432..697aae933f 100644
--- a/scripts/install_rocm_wsl_strixhalo.sh
+++ b/scripts/install_rocm_wsl_strixhalo.sh
@@ -219,7 +219,7 @@ fi
echo "${ROCM_DIR}/lib" | $SUDO tee /etc/ld.so.conf.d/rocm.conf >/dev/null
$SUDO ldconfig
-# ── Step 4: persist environment (system-wide so Studio's worker inherits it) ──
+# ── Step 4: persist environment (system-wide so Unsloth's worker inherits it) ──
say "Persisting ROCm-on-WSL environment"
_envfile="/etc/profile.d/unsloth-rocm-wsl.sh"
$SUDO tee "$_envfile" >/dev/null < str:
return f'''# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-"""Build-stamped Studio release metadata."""
+"""Build-stamped Unsloth release metadata."""
STUDIO_RELEASE_VERSION = {literal}
'''
@@ -168,7 +168,7 @@ def stamp(require_release: bool) -> int:
version, source = resolve_version()
if version is not None and not is_valid_version(version):
print(
- f"Invalid Studio release version from {source}: {version!r}",
+ f"Invalid Unsloth release version from {source}: {version!r}",
file = sys.stderr,
)
return 2
@@ -196,9 +196,9 @@ def stamp(require_release: bool) -> int:
if version is None:
if require_release:
print(
- "No Studio release version available. Set "
+ "No Unsloth release version available. Set "
"UNSLOTH_STUDIO_RELEASE_VERSION, build from a GitHub tag, "
- "or run from an exact local Studio release tag.",
+ "or run from an exact local Unsloth release tag.",
file = sys.stderr,
)
return 2
@@ -207,7 +207,7 @@ def stamp(require_release: bool) -> int:
return 0
_atomic_write_text(BUILD_INFO_PATH, build_info_source(version), encoding = "utf-8")
- print(f"Stamping Studio release version {version} from {source}", file = sys.stderr)
+ print(f"Stamping Unsloth release version {version} from {source}", file = sys.stderr)
print(version)
return 0
@@ -233,7 +233,7 @@ def _read_sdist_member(path: Path) -> str | None:
def verify_dist(expected: str, dist_dir: Path) -> int:
if not is_valid_version(expected):
- print(f"Invalid expected Studio release version: {expected!r}", file = sys.stderr)
+ print(f"Invalid expected Unsloth release version: {expected!r}", file = sys.stderr)
return 2
artifacts = list(dist_dir.glob("*.whl")) + list(dist_dir.glob("*.tar.gz"))
@@ -251,14 +251,14 @@ def verify_dist(expected: str, dist_dir: Path) -> int:
if content is None:
failures.append(f"{artifact.name}: missing {BUILD_INFO_SUFFIX}")
elif expected_line not in content:
- failures.append(f"{artifact.name}: Studio release version mismatch")
+ failures.append(f"{artifact.name}: Unsloth release version mismatch")
if failures:
for failure in failures:
print(failure, file = sys.stderr)
return 2
- print(f"Verified Studio release version {expected} in {len(artifacts)} artifact(s)")
+ print(f"Verified Unsloth release version {expected} in {len(artifacts)} artifact(s)")
return 0
diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1
index 88defb9ea0..9b6e6ebb86 100644
--- a/scripts/uninstall.ps1
+++ b/scripts/uninstall.ps1
@@ -83,7 +83,7 @@ function Uninstall-UnslothStudio {
}
}
- # A path is a Studio-owned root iff one of install.ps1's sentinels exists:
+ # A path is an Unsloth-owned root iff one of install.ps1's sentinels exists:
# \share\studio.conf, \unsloth_studio\.unsloth-studio-owned,
# or \bin\unsloth.exe.
function _IsStudioRoot {
@@ -164,7 +164,7 @@ function Uninstall-UnslothStudio {
return $p
}
- # Discover non-default Studio roots from env vars + studio.conf files.
+ # Discover non-default Unsloth roots from env vars + studio.conf files.
# Mirrors install.ps1's precedence: UNSLOTH_STUDIO_HOME wins, STUDIO_HOME
# is ignored when both are set, so uninstalling install A doesn't also
# delete install B if the user has a stale STUDIO_HOME pointing at B.
@@ -207,7 +207,7 @@ function Uninstall-UnslothStudio {
# Return $true iff the PID's image path lives under one of $KnownRoots.
# Prevents killing an unrelated process that happens to listen on a stale
- # Studio port.
+ # Unsloth port.
function _PidUnderKnownRoot {
param([int]$Pid_, [string[]]$KnownRoots)
if (-not $KnownRoots -or $KnownRoots.Count -eq 0) { return $false }
@@ -223,8 +223,8 @@ function Uninstall-UnslothStudio {
return $false
}
- # Stop a Studio backend whose port is recorded in \studio.port.
- # Only kills if the listening PID's exe path is under a known Studio root.
+ # Stop an Unsloth backend whose port is recorded in \studio.port.
+ # Only kills if the listening PID's exe path is under a known Unsloth root.
function _StopByPortFile {
param([string]$PortFile, [string[]]$KnownRoots)
if (-not (Test-Path -LiteralPath $PortFile -PathType Leaf)) { return }
@@ -372,7 +372,7 @@ function Uninstall-UnslothStudio {
continue
}
if (-not (_IsStudioRoot $r)) {
- _Substep "refusing to remove non-Studio path: $r" "Yellow"
+ _Substep "refusing to remove non-Unsloth path: $r" "Yellow"
continue
}
_RemovePath $r
@@ -436,7 +436,7 @@ function Uninstall-UnslothStudio {
$entries = $rawPath -split ';'
$kept = New-Object System.Collections.ArrayList
$removedAny = $false
- # Only remove PATH entries that live inside a Studio root we
+ # Only remove PATH entries that live inside an Unsloth root we
# actually own (default or env-mode). A literal substring
# match on `unsloth_studio` would clobber unrelated user
# virtualenvs that happen to share the name.
diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh
index 31e851fcbb..957d2b7af2 100755
--- a/scripts/uninstall.sh
+++ b/scripts/uninstall.sh
@@ -12,7 +12,7 @@
set -e
-# Stop a Studio server via its PID file (written by install.sh's _spawn_terminal).
+# Stop an Unsloth server via its PID file (written by install.sh's _spawn_terminal).
_kill_pid_file() {
_pid_file="$1"
[ -f "$_pid_file" ] || return 0
@@ -47,7 +47,7 @@ _pkill_studio() {
command -v pkill >/dev/null 2>&1 || return 0
# Scope fallback patterns to the install roots we are removing so a
- # different Studio install (different UNSLOTH_STUDIO_HOME) is not touched.
+ # different Unsloth install (different UNSLOTH_STUDIO_HOME) is not touched.
_kill_roots="$HOME/.unsloth/studio"
_roots_from_conf=$(_custom_studio_roots 2>/dev/null || true)
[ -n "$_roots_from_conf" ] && _kill_roots="$_kill_roots
@@ -89,7 +89,7 @@ _remove_path() {
fi
}
-# Accept as Studio root only if Studio sentinels exist (matches install.sh's
+# Accept as Unsloth root only if Unsloth sentinels exist (matches install.sh's
# env-mode ownership guard at install.sh:1358-1361). A bare unsloth_studio/
# directory is NOT enough -- require the install-time owner marker so a user
# directory that happens to contain a folder named "unsloth_studio" is safe.
@@ -175,8 +175,8 @@ _custom_studio_roots() {
_from_conf "$HOME/.local/share/unsloth/studio.conf"
}
-# Remove $HOME/.local/bin/unsloth only if it's a Studio-managed symlink.
-# Studio's install.sh writes this as a symlink into the studio venv
+# Remove $HOME/.local/bin/unsloth only if it's an Unsloth-managed symlink.
+# Unsloth's install.sh writes this as a symlink into the studio venv
# (install.sh: `ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path"`). A
# pip-installed `unsloth` CLI is a regular file — leave it alone to avoid
# wiping an unrelated install.
@@ -206,7 +206,7 @@ _custom_studio_roots | while IFS= read -r _custom_root; do
continue
fi
if ! _is_studio_root "$_custom_root"; then
- echo " refusing to remove non-Studio path: $_custom_root" >&2
+ echo " refusing to remove non-Unsloth path: $_custom_root" >&2
continue
fi
_remove_path "$_custom_root"
@@ -234,7 +234,7 @@ _remove_path "$HOME/.unsloth/rocm-smoketest"
# Drop ~/.unsloth only if now empty (rmdir refuses non-empty, so user content is kept).
rmdir "$HOME/.unsloth" 2>/dev/null || true
_remove_path "$HOME/.local/share/unsloth"
-# CLI shim: only the symlink Studio created, never a pip-installed file.
+# CLI shim: only the symlink Unsloth created, never a pip-installed file.
_remove_cli_shim
echo "Removing desktop shortcut and launcher lock..."
diff --git a/studio/MCP.md b/studio/MCP.md
index 91b39fcc69..127a85a116 100644
--- a/studio/MCP.md
+++ b/studio/MCP.md
@@ -1,10 +1,10 @@
# Unsloth Studio MCP server
-Studio can expose a local MCP server so an MCP client can inspect models and
+Unsloth can expose a local MCP server so an MCP client can inspect models and
GPU state, validate recipes, start or stop training, inspect recipe output, and
export a loaded model.
-The server is disabled by default. Enable it for a local Studio process with:
+The server is disabled by default. Enable it for a local Unsloth process with:
```bash
UNSLOTH_STUDIO_ENABLE_MCP=1 \
@@ -12,8 +12,8 @@ UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \
unsloth studio
```
-The endpoint is `http://127.0.0.1:8888/mcp/` when Studio uses its default port
-(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Studio
+The endpoint is `http://127.0.0.1:8888/mcp/` when Unsloth uses its default port
+(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Unsloth
port when it is configured differently.
The high-impact tools are:
@@ -23,9 +23,9 @@ The high-impact tools are:
- `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset`
- `load_checkpoint` and `export_gguf`
-`start_training` accepts the same fields as the Studio `TrainingStartRequest`.
+`start_training` accepts the same fields as the Unsloth `TrainingStartRequest`.
The request is validated by the existing Pydantic model before a subprocess is
-started. Export paths use the existing Studio validation as well.
+started. Export paths use the existing Unsloth validation as well.
The endpoint always requires `UNSLOTH_STUDIO_MCP_TOKEN` and checks an exact
Bearer token for both HTTP and WebSocket connections. Keep it on localhost
diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb
index 619395bd6d..44282b2255 100644
--- a/studio/Unsloth_Studio_Colab.ipynb
+++ b/studio/Unsloth_Studio_Colab.ipynb
@@ -33,7 +33,7 @@
"\n",
"We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n",
"\n",
- "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
+ "[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
]
},
{
diff --git a/studio/backend/assets/chat_templates/gemma-4-edge.jinja b/studio/backend/assets/chat_templates/gemma-4-edge.jinja
index 0266127233..74fa73ddd3 100644
--- a/studio/backend/assets/chat_templates/gemma-4-edge.jinja
+++ b/studio/backend/assets/chat_templates/gemma-4-edge.jinja
@@ -3,7 +3,7 @@
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
- Studio-local changes vs PR #118:
+ Unsloth-local changes vs PR #118:
1. preserve_thinking defaults to false (see SETUP block below).
2. The empty "<|channel>thought\n" block on enable_thinking=false is
NOT emitted. Google ships a distinct template for E2B/E4B (google/gemma-4-E2B-it,
diff --git a/studio/backend/assets/chat_templates/gemma-4.jinja b/studio/backend/assets/chat_templates/gemma-4.jinja
index 65ab39df57..cc5f98065f 100644
--- a/studio/backend/assets/chat_templates/gemma-4.jinja
+++ b/studio/backend/assets/chat_templates/gemma-4.jinja
@@ -3,7 +3,7 @@
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
- Studio-local change: preserve_thinking defaults to false (see SETUP block below).
+ Unsloth-local change: preserve_thinking defaults to false (see SETUP block below).
Applied to unsloth/gemma-4-*-GGUF models so the embedded GGUF template does not
need re-downloading. Keep in sync with upstream if PR #118 changes.
-#}
diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py
index b13cd1c851..dfb8fc513e 100644
--- a/studio/backend/auth/authentication.py
+++ b/studio/backend/auth/authentication.py
@@ -148,7 +148,7 @@ async def authenticated_via_api_key(
) -> bool:
"""True when the caller used an sk-unsloth API key, not a UI session JWT.
- Lets routes treat programmatic API callers differently from the Studio UI
+ Lets routes treat programmatic API callers differently from the Unsloth UI
(e.g. refuse a teardown the UI would allow).
"""
return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX))
diff --git a/studio/backend/auth/bootstrap_timeout.py b/studio/backend/auth/bootstrap_timeout.py
index 728433dc54..97a8086f04 100644
--- a/studio/backend/auth/bootstrap_timeout.py
+++ b/studio/backend/auth/bootstrap_timeout.py
@@ -1,13 +1,13 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged.
+"""Auto-shutdown for an exposed first-run Unsloth whose admin password is unchanged.
On a fresh install the seeded bootstrap admin password stays a valid login
credential until first login changes it. When the web UI is put on the network
(``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within
-a deadline, tear Studio down so a fresh, unconfigured instance does not stay
-publicly reachable indefinitely. If the password was changed, Studio keeps
+a deadline, tear Unsloth down so a fresh, unconfigured instance does not stay
+publicly reachable indefinitely. If the password was changed, Unsloth keeps
running.
Scope: web UI launches only (never ``--api-only``, which authenticates by API
@@ -98,7 +98,7 @@ def enforce_bootstrap_password_deadline(
) -> bool:
"""Deadline handler: shut down iff the seeded admin password is still unchanged.
- Returns True if it shut Studio down, False if it left it running (the
+ Returns True if it shut Unsloth down, False if it left it running (the
password was changed in time).
"""
try:
@@ -106,7 +106,7 @@ def enforce_bootstrap_password_deadline(
except Exception:
return False
if not still_default:
- return False # password changed in time -> leave Studio running
+ return False # password changed in time -> leave Unsloth running
message = (
"\nUnsloth Studio was exposed on the network but its default admin "
diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py
index 9bb3ab5735..39fa691304 100644
--- a/studio/backend/auth/storage.py
+++ b/studio/backend/auth/storage.py
@@ -146,7 +146,7 @@ def get_connection() -> sqlite3.Connection:
pass
conn.row_factory = sqlite3.Row
# WAL lets token reads run concurrently with refresh-token writes;
- # busy_timeout bounds lock waits. Matches the other Studio SQLite stores.
+ # busy_timeout bounds lock waits. Matches the other Unsloth SQLite stores.
# Set busy_timeout first: switching journal_mode needs a lock, so if a
# refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY;
# with busy_timeout already in effect it waits instead of failing and leaving
@@ -305,8 +305,8 @@ def get_or_create_identity_secret() -> bytes:
def compute_identity_proof(nonce: bytes, host: str, port: int) -> str:
"""HMAC-SHA256 proof that the caller holds this install's identity secret,
bound to the loopback address and port the connection landed on. A proof
- relayed from a Studio on a different address/port (a squatter proxying to the
- real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was
+ relayed from an Unsloth on a different address/port (a squatter proxying to the
+ real one, e.g. localhost resolving to ::1 while Unsloth is on 127.0.0.1) was
computed for that other endpoint and won't match the one the client dialed."""
try:
host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms
diff --git a/studio/backend/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py
index 8491019ae9..e855f4078b 100644
--- a/studio/backend/auth/terminal_prompt.py
+++ b/studio/backend/auth/terminal_prompt.py
@@ -2,14 +2,14 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Interactive terminal prompt that forces a bootstrap password change before
-Studio is exposed on a public Cloudflare URL (``--secure`` / ``--cloudflare``).
+Unsloth is exposed on a public Cloudflare URL (``--secure`` / ``--cloudflare``).
Masked input echoes one ``*`` per keystroke (unlike ``getpass``). Works on
Windows (``msvcrt``) and Linux/macOS (``termios``). All output goes to stderr so
redirected stdout never swallows the prompt.
Mirrored for the CLI at ``unsloth_cli/commands/_password_prompt.py`` (the CLI
-cannot import the Studio backend package); keep the two in sync.
+cannot import the Unsloth backend package); keep the two in sync.
"""
from __future__ import annotations
@@ -252,7 +252,7 @@ def prompt_for_password_change(
out.flush()
return True
except (KeyboardInterrupt, EOFError):
- out.write("Password change aborted; not exposing Studio.\n")
+ out.write("Password change aborted; not exposing Unsloth.\n")
out.flush()
return False
diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py
index ef7bacba67..b1ddc74c32 100644
--- a/studio/backend/cloudflare_tunnel.py
+++ b/studio/backend/cloudflare_tunnel.py
@@ -1,13 +1,13 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-"""Free Cloudflare quick tunnel for Studio's 0.0.0.0 launches.
+"""Free Cloudflare quick tunnel for Unsloth's 0.0.0.0 launches.
The raw http://: is often unreachable (https-vs-http, blocked ports,
closed security groups); a cloudflared quick tunnel gives a free
https://*.trycloudflare.com URL that works anywhere, with no account or domain.
-Best-effort throughout: any failure collapses to "no URL" and Studio keeps
+Best-effort throughout: any failure collapses to "no URL" and Unsloth keeps
running. Stdlib only (back-end imports are lazy) so it is safe to import early.
"""
@@ -95,7 +95,7 @@ def _cache_path() -> Optional[Path]:
def find_cloudflared() -> Optional[str]:
- """Locate an existing cloudflared: PATH first, then the Studio bin cache."""
+ """Locate an existing cloudflared: PATH first, then the Unsloth bin cache."""
on_path = shutil.which("cloudflared")
if on_path:
return on_path
@@ -309,7 +309,7 @@ class CloudflareTunnel:
pass
-# Single serving process per Studio launch, so one module-level tunnel handle is
+# Single serving process per Unsloth launch, so one module-level tunnel handle is
# enough; the lock guards the start/stop/shutdown races.
_active_tunnel: Optional[CloudflareTunnel] = None
_active_lock = threading.Lock()
diff --git a/studio/backend/colab.py b/studio/backend/colab.py
index e04543b3aa..1762469bcf 100644
--- a/studio/backend/colab.py
+++ b/studio/backend/colab.py
@@ -129,7 +129,7 @@ def start_cloudflare_tunnel(port: int) -> "str | None":
logger.warning(
"Cloudflare link not started: the admin account still has its temporary "
"bootstrap password, which is exposed to anyone who can load the page. "
- "Open Studio in this tab, log in and change the admin password, then re-run "
+ "Open Unsloth in this tab, log in and change the admin password, then re-run "
"start(cloudflare=True) to get the shareable link."
)
return None
@@ -203,7 +203,7 @@ def _shareable_link_html(cloudflare_url: str) -> str:
display: flex; align-items: center; gap: 12px;">
- Shareable Studio Link is Ready!
+ Shareable Unsloth Link is Ready!
ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
- message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
+ message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."),
),
)
@@ -147,7 +147,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
status = "rate_limited",
retry_after_sec = seconds,
message = (
- "Waiting for GitHub secondary rate limit. Studio will resume automatically."
+ "Waiting for GitHub secondary rate limit. Unsloth will resume automatically."
),
),
)
@@ -161,7 +161,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
- message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
+ message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."),
),
)
diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py
index ebb1d39dfb..ffc81669ae 100644
--- a/studio/backend/core/data_recipe/local_callable_validators.py
+++ b/studio/backend/core/data_recipe/local_callable_validators.py
@@ -238,7 +238,7 @@ def _run_oxc_batch(
if not node_executable:
return _fallback_results(
len(code_values),
- "Node.js not found (install Node >= 20.19, or re-run Studio setup to provision it).",
+ "Node.js not found (install Node >= 20.19, or re-run Unsloth setup to provision it).",
)
try:
tmp_dir = ensure_dir(oxc_validator_tmp_root())
diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py
index 4647dc098d..9770e88b7f 100644
--- a/studio/backend/core/data_recipe/service.py
+++ b/studio/backend/core/data_recipe/service.py
@@ -280,8 +280,8 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None =
from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports]
if artifact_path is None:
- # DataDesigner defaults to cwd/artifacts; packaged Studio can run with
- # cwd=/, so keep default callers on Studio's writable recipe artifact root.
+ # DataDesigner defaults to cwd/artifacts; packaged Unsloth can run with
+ # cwd=/, so keep default callers on Unsloth's writable recipe artifact root.
artifact_path = str(recipe_datasets_root())
recipe = _strip_frontend_model_config_metadata(recipe)
diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py
index ad78157418..1491dfa749 100644
--- a/studio/backend/core/inference/__init__.py
+++ b/studio/backend/core/inference/__init__.py
@@ -11,7 +11,7 @@ subprocess and can be imported directly from .inference when needed.
Public names are resolved lazily (PEP 562): importing this package -- or a
dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull
the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML
-backend and its Studio dependencies). Those load only when a public name is
+backend and its Unsloth dependencies). Those load only when a public name is
actually accessed, so standalone helpers stay unit-testable without the full
inference stack.
"""
diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py
index 3c7a4cb182..34445cc58e 100644
--- a/studio/backend/core/inference/anthropic_compat.py
+++ b/studio/backend/core/inference/anthropic_compat.py
@@ -539,7 +539,7 @@ class AnthropicPassthroughEmitter:
Only calls naming a tool in ``allowed_tools`` (the client's declared
tools) are promoted; everything else streams as text exactly as before.
- Never enabled for Studio's own tool loop.
+ Never enabled for Unsloth's own tool loop.
"""
from core.inference.passthrough_healing import StreamToolCallHealer
diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py
index 5113eebb36..528c059fbc 100644
--- a/studio/backend/core/inference/chat_template_helpers.py
+++ b/studio/backend/core/inference/chat_template_helpers.py
@@ -150,7 +150,7 @@ def _split_partial_marker(text: str, marker: str) -> tuple[str, str]:
class ReasoningChannelNormalizer:
"""Incrementally convert one native reasoning channel to ````.
- The parser follows mlx-vlm's streaming boundary behavior but emits Studio's
+ The parser follows mlx-vlm's streaming boundary behavior but emits Unsloth's
established canonical text contract. Only the configured opening and
closing markers are consumed; tool-call and other control markers remain
available to downstream parsers.
diff --git a/studio/backend/core/inference/chat_templates.py b/studio/backend/core/inference/chat_templates.py
index 58f63ff61b..04c0db6aae 100644
--- a/studio/backend/core/inference/chat_templates.py
+++ b/studio/backend/core/inference/chat_templates.py
@@ -4,13 +4,13 @@
"""Bundled chat-template selection for GGUF inference.
Some shipped GGUF quants embed an older chat template. Rather than re-cutting and
-asking users to re-download every quant, Studio can override the embedded template
+asking users to re-download every quant, Unsloth can override the embedded template
at llama-server launch time with a bundled, up-to-date Jinja template for known
model families. The override is wired through the existing ``chat_template_override``
-> ``--chat-template-file`` path in ``LlamaCppBackend.load_model``.
Currently this covers ``unsloth/gemma-4-*-GGUF``, which gains the upstream PR #118
-``preserve_thinking`` flag (defaulted OFF here) so the Studio "Preserve thinking"
+``preserve_thinking`` flag (defaulted OFF here) so the Unsloth "Preserve thinking"
toggle appears while staying disabled by default.
"""
diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py
index 20312e067c..2debf946e9 100644
--- a/studio/backend/core/inference/external_provider.py
+++ b/studio/backend/core/inference/external_provider.py
@@ -473,7 +473,7 @@ def _apply_mistral_reasoning_controls(
# handles every provider without storing credentials.
def _create_shared_http_client() -> httpx.AsyncClient:
# Unsupported env proxy schemes (socks:// etc) raise at construction and
- # would crash Studio startup (#6090); retry ignoring env proxies instead.
+ # would crash Unsloth startup (#6090); retry ignoring env proxies instead.
try:
return httpx.AsyncClient()
except (ImportError, ValueError) as exc:
@@ -858,7 +858,7 @@ class ExternalProviderClient:
if not self._is_openai_compatible():
# Gemini speaks its own native REST shape (contents/parts);
# `_stream_gemini` translates request/response into the OpenAI
- # Chat Completions chunk format the rest of Studio expects.
+ # Chat Completions chunk format the rest of Unsloth expects.
# API ref: https://ai.google.dev/gemini-api/docs
if self.provider_type == "gemini":
async for line in self._stream_gemini(
@@ -1706,7 +1706,7 @@ class ExternalProviderClient:
# Translate OpenAI multimodal parts -> Anthropic native shapes.
# - `image_url` -> `{type:"image", source:...}`
# - `input_document` -> `{type:"document", source:...}`
- # (Studio extension; mirrors Anthropic's document block,
+ # (Unsloth extension; mirrors Anthropic's document block,
# which supports PDFs as base64 or URL per
# https://platform.claude.com/docs/en/build-with-claude/vision)
anthropic_parts: list[dict[str, Any]] = []
@@ -1749,7 +1749,7 @@ class ExternalProviderClient:
}
)
elif part.get("type") == "input_document":
- # Studio's normalised PDF/doc type (file_data data-URI or
+ # Unsloth's normalised PDF/doc type (file_data data-URI or
# file_url) -> Anthropic's native `document` block.
url = part.get("file_url") or ""
data_uri = part.get("file_data") or ""
@@ -4704,7 +4704,7 @@ class ExternalProviderClient:
{"type": "image_generation_call", "id": call_id}
)
elif part_type == "input_document":
- # Map Studio's `input_document` onto Responses' `input_file`.
+ # Map Unsloth's `input_document` onto Responses' `input_file`.
# https://developers.openai.com/api/docs/guides/images-vision
file_url = part.get("file_url")
file_data = part.get("file_data")
@@ -6010,7 +6010,7 @@ class ExternalProviderClient:
if not models and self.provider_type == "ollama":
models = await self._list_ollama_native_models()
# Gemini's native /v1beta/models uses a different shape; repackage
- # into the OpenAI-compatible one Studio expects.
+ # into the OpenAI-compatible one Unsloth expects.
if not models and self.provider_type == "gemini":
models = self._parse_gemini_models(data)
return models
@@ -6213,7 +6213,7 @@ def _friendly_provider_error_text(
*,
model: str | None = None,
) -> str:
- """Rewrite common provider errors into actionable Studio copy."""
+ """Rewrite common provider errors into actionable Unsloth copy."""
if status_code == 404 and model:
lowered = raw_message.lower()
if "not found" in lowered or "not_found" in lowered:
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index fec5214cf9..c9ab7eb83b 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -116,7 +116,7 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = (
# llama-server can serve HTTP 200 while running a model entirely on CPU when a
# GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so
-# Studio can warn. Priority: explicit "offloaded N/M layers to GPU" counts
+# Unsloth can warn. Priority: explicit "offloaded N/M layers to GPU" counts
# (authoritative), then GPU "model buffer size" lines (host-pinned _Host
# excluded), then the "device_info:" device table (disconfirm only).
_GPU_OFFLOAD_MARKERS = (
@@ -1363,7 +1363,7 @@ def _kv_bytes_per_elem(cache_type: Optional[str]) -> float:
def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]:
"""Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it
- exceeds the f16 default, else None. Studio emits --cache-type only for the
+ exceeds the f16 default, else None. Unsloth emits --cache-type only for the
param/extras path, so a heavier env (f32) would otherwise reach the child
unbudgeted; quantized env types stay over-reserved by f16 (-> None)."""
e = os.environ if env is None else env
@@ -1682,7 +1682,7 @@ def _build_ngram_mod_flags(
return []
-# Canonical Speculative Decoding modes exposed by the Studio chat UI.
+# Canonical Speculative Decoding modes exposed by the Unsloth chat UI.
# Dropdown renders five (auto, mtp, ngram, mtp+ngram, off); the load API
# also accepts legacy values the original Switch and external callers emit
# (default, draft-mtp, ngram-mod, ngram-simple).
@@ -1731,7 +1731,7 @@ def _backfill_usage_from_timings(usage, timings):
"""Synthesize ``usage`` from llama-server's ``timings`` when the
OpenAI-style usage block is missing or reports zero tokens.
- The Studio chat UI computes generation t/s from
+ The Unsloth chat UI computes generation t/s from
``meta.usage.completion_tokens / totalStreamTime``. llama-server always
populates ``timings.predicted_n`` (true decoded count) and
``timings.prompt_n``, but the final SSE chunk's ``usage`` can be absent
@@ -1804,7 +1804,7 @@ def _llama_lib_dir(binary: str) -> Path:
def _is_external_link(path: Path) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
- the user's own llama.cpp checkout, which Studio does not own."""
+ the user's own llama.cpp checkout, which Unsloth does not own."""
try:
if os.path.islink(path):
return True
@@ -1960,7 +1960,7 @@ class LlamaCppBackend:
# observes it (direct proxy endpoints, or nothing in flight).
self._mtp_watchdog_thread: Optional[threading.Thread] = None
self._mtp_watchdog_stop = threading.Event()
- # True when the launch actually runs MTP+tensor (Studio- or user/env-driven);
+ # True when the launch actually runs MTP+tensor (Unsloth- or user/env-driven);
# gates the probe, watchdog, and recovery so pass-through MTP is covered.
self._mtp_runtime_fallback_active = False
self._stdout_lines: list[str] = []
@@ -2353,7 +2353,7 @@ class LlamaCppBackend:
@staticmethod
def _resolved_studio_root_and_is_legacy() -> "tuple[Optional[Path], bool]":
- """Resolve the Studio install root and classify it as the legacy
+ """Resolve the Unsloth install root and classify it as the legacy
~/.unsloth/studio root vs. a custom (env/venv-inferred) root.
Returns (resolved_root, is_legacy). On any import/resolution failure the
@@ -3241,7 +3241,7 @@ class LlamaCppBackend:
return
prev = curr
- # Free-VRAM fraction at which Studio pins the GPU directly instead of
+ # Free-VRAM fraction at which Unsloth pins the GPU directly instead of
# deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in
# the fit, so this only guards fragmentation + multi-GPU per-device CUDA context
# (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106).
@@ -3800,7 +3800,7 @@ class LlamaCppBackend:
return total if total > 0 else None
return draft_kv + weights + target_ctx_copy
- _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it
+ _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Unsloth does not override it
_COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate
# Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682).
_CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB)
@@ -3940,7 +3940,7 @@ class LlamaCppBackend:
n_ubatch: Optional[int] = None,
) -> tuple[Optional[list[int]], bool, int]:
"""Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits,
- so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers
+ so Unsloth keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers
to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the
slot-independent footprint (weights + soft overhead + MTP + context-linear compute,
minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer
@@ -4416,7 +4416,7 @@ class LlamaCppBackend:
]
# Otherwise hand off to the resolver (cache / bootstrap / transformers / HF). Diffusion models
- # skip it: they do not use Studio's SWA pattern and the resolver can raise for them.
+ # skip it: they do not use Unsloth's SWA pattern and the resolver can raise for them.
if (
self._sliding_window_pattern is None
and self._sliding_window
@@ -4536,7 +4536,7 @@ class LlamaCppBackend:
) -> bool:
"""Launch the OpenAI-compat diffusion shim (which drives the on-device
visual decoder) and wait for health. Presents the same /v1 + /health
- interface as llama-server, so the rest of Studio is unchanged.
+ interface as llama-server, so the rest of Unsloth is unchanged.
"""
assets = self._find_diffusion_assets()
if assets is None:
@@ -4608,7 +4608,7 @@ class LlamaCppBackend:
logger.debug(f"Could not open diffusion runner log file: {e}")
# The shim (and its visual server) die with this backend process, so a
- # Studio crash/restart never orphans a GPU process.
+ # Unsloth crash/restart never orphans a GPU process.
self._process = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
@@ -5242,7 +5242,7 @@ class LlamaCppBackend:
return (
f"'{arch}' is a diffusion (image-generation) GGUF, which "
"llama-server cannot run as a chat/completion model. Use "
- "Studio's Images page to generate with local diffusion "
+ "Unsloth's Images page to generate with local diffusion "
"GGUFs such as FLUX and Qwen-Image."
)
if is_ollama:
@@ -6103,7 +6103,7 @@ class LlamaCppBackend:
and not bool(mtp_draft_path)
)
# LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras
- # nor Studio emit a spec flag (mode "off", no user --spec-type),
+ # nor Unsloth emit a spec flag (mode "off", no user --spec-type),
# since _build_speculative_flags emits one for every other mode.
# Consult the env for the reserve only then, else a stale MTP env
# would over-reserve.
@@ -6112,7 +6112,7 @@ class LlamaCppBackend:
if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off")
else {}
)
- # Extras can run MTP even when Studio suppresses its own emission.
+ # Extras can run MTP even when Unsloth suppresses its own emission.
_user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env)
# A non-MTP model-based draft mode (draft-simple/draft-eagle3) in
# extras also loads a separate draft model that needs reserving;
@@ -6178,7 +6178,7 @@ class LlamaCppBackend:
_mtp_eff_n_max = 2 if gpus else 3
# Separate-drafter weights live on GPU (an embedded head is
# already in model_size). Size the drafter the launch loads, by
- # precedence: extras --model-draft (last-wins), else Studio's
+ # precedence: extras --model-draft (last-wins), else Unsloth's
# emitted mtp_draft_path, else the env drafter. Sizing the wrong
# one would under-reserve and OOM.
_cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {})
@@ -7097,12 +7097,12 @@ class LlamaCppBackend:
# Vulkan pins via --device (a cmd arg, unlike the env-based
# CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's
- # last-wins parsing lets a user --device override Studio's pick.
+ # last-wins parsing lets a user --device override Unsloth's pick.
if is_vulkan_backend and gpu_indices is not None:
cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices)
# User pass-through args go last so llama.cpp's last-wins parsing
- # lets the user override Studio's auto-set flags. Already
+ # lets the user override Unsloth's auto-set flags. Already
# validated by the route via validate_extra_args().
if extra_args:
cmd.extend(str(a) for a in extra_args)
@@ -7118,9 +7118,9 @@ class LlamaCppBackend:
if "--threads" not in cmd:
env.pop("LLAMA_ARG_THREADS", None)
- # Reconcile the inherited LLAMA_ARG_* env with Studio's final
+ # Reconcile the inherited LLAMA_ARG_* env with Unsloth's final
# decision: stripping CLI extras on a tensor->layer downgrade
- # can't remove env vars, so the child could run a mode/KV Studio
+ # can't remove env vars, so the child could run a mode/KV Unsloth
# didn't budget.
if not tensor_parallel:
# Layer split: clear a non-layer inherited split mode (and any
@@ -7130,7 +7130,7 @@ class LlamaCppBackend:
env.pop("LLAMA_ARG_SPLIT_MODE", None)
env.pop("LLAMA_ARG_TENSOR_SPLIT", None)
else:
- # Studio owns the tensor split: it emits --tensor-split when it
+ # Unsloth owns the tensor split: it emits --tensor-split when it
# picks an uneven one (CLI wins) and nothing when an even split
# is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even
# case can't be overridden by a stale env (the layer branch above
@@ -7201,7 +7201,7 @@ class LlamaCppBackend:
# 'on') even when -ngl is explicit. That step has aborted on
# some ROCm hosts (ggml-cuda.cu ROCm error during worst-case
# estimation, e.g. MTP + mmproj models on gfx1151). When
- # Studio's own VRAM math already placed the model
+ # Unsloth's own VRAM math already placed the model
# (use_fit=False), the step is redundant second-guessing --
# retry once with --fit off before declaring the load failed.
# Never retry when fit was requested (use_fit) or the caller
@@ -7284,7 +7284,7 @@ class LlamaCppBackend:
and _startup_crashed
and not _split_axis_crash
):
- # We forced --fit off because Studio's (conservative) VRAM
+ # We forced --fit off because Unsloth's (conservative) VRAM
# math placed the model fully on GPU. A startup crash here
# means that estimate was optimistic, so fall back to --fit
# on and let llama.cpp offload rather than fail the load.
@@ -7296,7 +7296,7 @@ class LlamaCppBackend:
self._process.returncode,
self._llama_log_path,
)
- # Flip Studio's own --fit off (added first, before any
+ # Flip Unsloth's own --fit off (added first, before any
# user extra args) to on; a user's later --fit still wins
# by last-arg. Defensive: if absent, the default is already
# --fit on, so leave it.
@@ -7313,7 +7313,7 @@ class LlamaCppBackend:
):
logger.warning(
"llama-server crashed during startup (exit code %s) "
- "with the default memory-fit step enabled; Studio "
+ "with the default memory-fit step enabled; Unsloth "
"already verified the model fits, retrying once "
"with --fit off. Crash log: %s",
self._process.returncode,
@@ -7393,7 +7393,7 @@ class LlamaCppBackend:
cmd = _fa_cmd
healthy = _spawn_and_wait(_fa_cmd, label = "-noflash")
- # MTP from Studio's spec flags or the user's (extra_args
+ # MTP from Unsloth's spec flags or the user's (extra_args
# --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child
# only when neither emits a spec flag, so consult it only then.
_launch_spec_env: Mapping[str, str] = (
@@ -7587,11 +7587,11 @@ class LlamaCppBackend:
if self._gpu_offload_active is False:
logger.warning(
"llama-server appears to have loaded the model entirely "
- "on CPU even though Studio detected at least one GPU. "
+ "on CPU even though Unsloth detected at least one GPU. "
"This usually means the prebuilt binary's GPU backend "
"failed to load -- on Windows, cudart64_X.dll / "
"cublas64_X.dll could not be resolved. Reinstall the "
- "Studio llama.cpp prebuilt or install a matching CUDA "
+ "Unsloth llama.cpp prebuilt or install a matching CUDA "
"toolkit (issue unslothai/unsloth#5106).",
)
@@ -7888,7 +7888,7 @@ class LlamaCppBackend:
logger.info(
"Auto: MLA embedded-MTP model detected; llama.cpp's MLA/DSA "
"MTP path is slower than no speculation, so using ngram-mod "
- "instead. Override via the Studio Speculative Decoding "
+ "instead. Override via the Unsloth Speculative Decoding "
"dropdown or UNSLOTH_MLA_MTP_ENABLED=1."
)
_emit_ngram_mod()
@@ -7916,7 +7916,7 @@ class LlamaCppBackend:
f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
"is below the 3B speedup threshold; using ngram-mod "
"only (zero-VRAM, no draft head). Override via "
- "--spec-type or the Studio Speculative Decoding "
+ "--spec-type or the Unsloth Speculative Decoding "
"dropdown."
)
_emit_ngram_mod()
@@ -8294,7 +8294,7 @@ class LlamaCppBackend:
def _pid_parent_is_alive(pid: int) -> bool:
"""True if the recorded server's parent is still running, i.e. the server is
NOT orphaned. Lets the cross-session reap kill only a true orphan (parent
- gone) and never a live server owned by a running Studio, regardless of which
+ gone) and never a live server owned by a running Unsloth, regardless of which
process performs the sweep. Biased toward "alive" on uncertainty so a live
server is never mistakenly reaped."""
try:
@@ -8334,9 +8334,9 @@ class LlamaCppBackend:
@classmethod
def _reap_recorded_pid(cls) -> int:
"""Kill the exact llama-server PID recorded at spawn, but only when it is a
- genuine orphan -- its parent (the Studio that spawned it) is gone. This is
+ genuine orphan -- its parent (the Unsloth that spawned it) is gone. This is
the cross-session backstop the parent-death reaper (Job Object /
- PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Studio
+ PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Unsloth
(macOS, a best-effort failure, or a pre-existing orphan). Path-independent,
so it also catches an orphan the install-root match would miss.
@@ -8393,7 +8393,7 @@ class LlamaCppBackend:
"""Kill orphaned llama-server processes started by studio.
Only kills processes whose resolved binary lives under a known
- Studio install dir (or matches an exact env-var override), to avoid
+ Unsloth install dir (or matches an exact env-var override), to avoid
terminating unrelated llama-server instances. Mirrors every location
_find_llama_server_binary() can return, so orphans from any
supported install path are cleaned up.
@@ -8413,7 +8413,7 @@ class LlamaCppBackend:
try:
# -- Build the ownership allowlist --------------------------------
# exact_binaries -- env var overrides (exact path match).
- # install_roots -- Studio-owned dir trees (binary must be under one).
+ # install_roots -- Unsloth-owned dir trees (binary must be under one).
install_roots: list[Path] = []
# Env-mode custom root (mirrors _find_llama_server_binary).
@@ -8423,7 +8423,7 @@ class LlamaCppBackend:
install_roots.append(_resolved_sr / "llama.cpp")
# Primary install dir (default mode only). Env-mode skips this so a
- # custom-root Studio can't kill a default-install Studio's server.
+ # custom-root Unsloth can't kill a default-install Unsloth's server.
if not _is_custom_root:
install_roots.append(Path.home() / ".unsloth" / "llama.cpp")
@@ -8497,7 +8497,7 @@ class LlamaCppBackend:
if not is_ours:
continue
- # A live parent means a running Studio (or the user's
+ # A live parent means a running Unsloth (or the user's
# shell) still owns it -- not an orphan.
if LlamaCppBackend._pid_parent_is_alive(proc.info["pid"]):
continue
@@ -8577,7 +8577,7 @@ class LlamaCppBackend:
def _fit_off_retry_eligible(cmd: "list[str]", use_fit: bool) -> bool:
"""Whether a llama-server startup crash may be retried with --fit off.
- Only when Studio's own VRAM math placed the model (use_fit=False)
+ Only when Unsloth's own VRAM math placed the model (use_fit=False)
and nothing on the command line set the fit mode explicitly
(-fit / --fit, space- or equals-form). --fit-ctx / --fit-target /
-fitc / -fitt tune the fit step but do not select the mode, so
@@ -8821,7 +8821,7 @@ class LlamaCppBackend:
return None
def _reconcile_effective_ctx_with_server(self) -> None:
- """Adopt the server's real ``n_ctx`` when it is below Studio's value.
+ """Adopt the server's real ``n_ctx`` when it is below Unsloth's value.
Keeps ``context_length`` (load response, status route, passthrough
``max_tokens`` ceiling) honest; clients sized to the requested value
diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py
index 4ce663c3ce..86a8c8a404 100644
--- a/studio/backend/core/inference/llama_keepwarm.py
+++ b/studio/backend/core/inference/llama_keepwarm.py
@@ -59,7 +59,7 @@ _INFERENCE_SUFFIXES = (
"/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages
"/embeddings",
"/responses",
- "/generate/stream", # Studio's own streaming route on the same llama-server
+ "/generate/stream", # Unsloth's own streaming route on the same llama-server
"/audio/generate", # direct GGUF TTS; can outlive the idle TTL
)
diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py
index f400d2ae40..70d0dc774d 100644
--- a/studio/backend/core/inference/llama_server_args.py
+++ b/studio/backend/core/inference/llama_server_args.py
@@ -3,10 +3,10 @@
"""Boundary validator for user-supplied llama-server pass-through args.
-Reject only flags Studio manages (model identity, auth, network, parallel
+Reject only flags Unsloth manages (model identity, auth, network, parallel
slots). Everything else (sampling, ``-c``, ``-ngl``, ``--flash-attn``,
``--cache-type-*``, ``--spec-*``, ``--jinja``, ...) is appended after
-Studio's auto-set flags so llama.cpp's last-wins parser lets the user override.
+Unsloth's auto-set flags so llama.cpp's last-wins parser lets the user override.
Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
"""
@@ -22,12 +22,12 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Parallel slots: owned by typer --parallel; a pass-through would desync
# app.state.llama_parallel_slots from llama-server.
frozenset({"-np", "--parallel", "--n-parallel"}),
- # Model identity: Studio resolves it from LoadRequest; a second -m would
- # load a different model than Studio thinks it loaded.
+ # Model identity: Unsloth resolves it from LoadRequest; a second -m would
+ # load a different model than Unsloth thinks it loaded.
frozenset({"-m", "--model"}),
- # Public model id: Studio sets a sanitized --alias so the OpenAI API never
+ # Public model id: Unsloth 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
+ # Unsloth'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"}),
@@ -39,14 +39,14 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"-hft", "--hf-token"}),
frozenset({"-mm", "--mmproj"}),
frozenset({"-mmu", "--mmproj-url"}),
- # Networking: Studio binds + proxies; retargeting orphans the proxy.
+ # Networking: Unsloth binds + proxies; retargeting orphans the proxy.
frozenset({"--host"}),
frozenset({"--port"}),
frozenset({"--path"}),
frozenset({"--api-prefix"}),
frozenset({"--reuse-port"}),
- # Auth / TLS: Studio terminates auth; upstream --api-key / TLS shadows
- # Studio's key and breaks the proxy hop.
+ # Auth / TLS: Unsloth terminates auth; upstream --api-key / TLS shadows
+ # Unsloth's key and breaks the proxy hop.
frozenset({"--api-key"}),
frozenset({"--api-key-file"}),
frozenset({"--ssl-key-file"}),
@@ -64,11 +64,11 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"--models-max"}),
frozenset({"--models-autoload", "--no-models-autoload"}),
# Server-mode flips: --embedding / --rerank restrict llama-server to
- # those endpoints, breaking Studio's /v1/chat/completions hop.
+ # those endpoints, breaking Unsloth's /v1/chat/completions hop.
frozenset({"--embedding", "--embeddings"}),
frozenset({"--rerank", "--reranking"}),
# llama-server's own built-in tools flag would silently stack on top of
- # Studio's --enable-tools / --disable-tools policy resolver.
+ # Unsloth's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
)
@@ -120,7 +120,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
def is_managed_flag(flag: str) -> bool:
- """True if ``flag`` is Studio-managed. Normalises via ``_flag_name`` so
+ """True if ``flag`` is Unsloth-managed. Normalises via ``_flag_name`` so
`-np8` / `--parallel=8` classify like the canonical tokens."""
normalised = _flag_name(flag)
return normalised is not None and normalised in _DENYLIST
@@ -142,7 +142,7 @@ _SPEC_FLAGS: frozenset[str] = frozenset(
"--draft-min",
"--draft-max",
# MTP path (llama.cpp #22673). The drafter selectors (local --model-draft
- # and HF --spec-draft-hf aliases) are Studio-managed since the separate-
+ # and HF --spec-draft-hf aliases) are Unsloth-managed since the separate-
# drafter support (Gemma 4): an inherited copy must not last-wins-override
# the auto-detected drafter. Explicit extras for the current load are never
# stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld,
@@ -179,9 +179,9 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset(
# (--split-mode tensor). Pass-through stays allowed so users keep the
# row/none/layer modes the toggle doesn't expose, but it's stripped on
# inherit and reconciled into the round-tripped tensor_parallel state.
-# --tensor-split is coupled to the split mode and is stripped with it: Studio
+# --tensor-split is coupled to the split mode and is stripped with it: Unsloth
# owns the tensor-mode split ratios, so an inherited/stale --tensor-split must
-# not last-wins-override Studio's computed asymmetric split.
+# not last-wins-override Unsloth's computed asymmetric split.
_SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"})
_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"})
_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
@@ -197,7 +197,7 @@ _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja
def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
"""Return the last user-supplied ``-c`` / ``--ctx-size`` value.
- Mirrors llama.cpp's last-wins parsing for the one numeric knob Studio's
+ Mirrors llama.cpp's last-wins parsing for the one numeric knob Unsloth's
load-time fit logic needs.
"""
if not args:
@@ -286,7 +286,7 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
Mirrors parse_ctx_override but for cache type. Recognises both -ctk
(key) and -ctv (value). When both flags appear, returns the last-wins
value, treating key and value cache flags as the same setting because
- Studio's KV estimate has a single cache_type_kv knob.
+ Unsloth's KV estimate has a single cache_type_kv knob.
"""
return _last_flag_value(args, _CACHE_FLAGS)
@@ -341,7 +341,7 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral
def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool:
- """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio
+ """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Unsloth
emits --split-mode only on its tensor branch, so a tensor env on the layer
path would run the child tensor-parallel unbudgeted; this flips the budget
to tensor. Only tensor is heavier, so other modes are ignored."""
@@ -425,7 +425,7 @@ def strip_shadowing_flags(
strip_template: bool = True,
strip_split_mode: bool = True,
) -> list[str]:
- """Strip flags that shadow first-class Studio settings.
+ """Strip flags that shadow first-class Unsloth settings.
Used when inheriting a previous load's ``llama_extra_args`` so an
inherited `-c 4096` can't override the current `max_seq_length`
diff --git a/studio/backend/core/inference/llama_stats.py b/studio/backend/core/inference/llama_stats.py
index 6047aedbc0..ab0d287e8c 100644
--- a/studio/backend/core/inference/llama_stats.py
+++ b/studio/backend/core/inference/llama_stats.py
@@ -5,7 +5,7 @@
engine-stats log line (generation/prompt throughput, requests in flight).
llama-server already computes these (it needs `--metrics`); this lifts them
-into Studio's structured log so the terminal shows serving health, not just
+into Unsloth's structured log so the terminal shows serving health, not just
per-request access lines. Emitted only while there is activity.
"""
diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py
index 002cafe2c8..86ad8b9fd8 100644
--- a/studio/backend/core/inference/local_model_resolver.py
+++ b/studio/backend/core/inference/local_model_resolver.py
@@ -130,7 +130,7 @@ def info_has_local_gguf(info) -> bool:
def _build_index() -> dict[str, _LocalGgufEntry]:
"""Map normalized id/model_id/display_name -> local GGUF entry.
- Scans the same roots Studio's model picker lists (./models, the active plus
+ Scans the same roots Unsloth's model picker lists (./models, the active plus
legacy/default HF caches, LM Studio dirs, and user scan folders) so a named
local model is never missed and silently served as the loaded one. Ollama's
scanner is skipped: it creates symlinks as a side effect and this runs on the
@@ -199,7 +199,7 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
raw_id = getattr(info, "id", None)
if not raw_id:
continue
- # Skip what Studio hides from its pickers (validation probe, RAG embed
+ # Skip what Unsloth hides from its pickers (validation probe, RAG embed
# weights): not chat models, so never an auto-switch target.
if _is_hidden_model(raw_id, getattr(info, "path", None)):
continue
diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py
index 6b5ce02216..0256df944e 100644
--- a/studio/backend/core/inference/mcp_client.py
+++ b/studio/backend/core/inference/mcp_client.py
@@ -906,7 +906,7 @@ def _call_stdio_tool(
def _remaining() -> Optional[float]:
return None if deadline is None else max(0.0, deadline - time.monotonic())
- # Callers without a Studio session id must retain the former one-shot
+ # Callers without an Unsloth session id must retain the former one-shot
# behavior: no browser/cookie/tool state can leak into another request.
# Use an ephemeral key (and close it below) rather than the shared empty
# scope that the persistent-session cache used previously.
diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py
index ed7c7ecfcf..e6da0a22b0 100644
--- a/studio/backend/core/inference/passthrough_healing.py
+++ b/studio/backend/core/inference/passthrough_healing.py
@@ -5,7 +5,7 @@
With server-side tools disabled (``unsloth run --disable-tools``, every
``unsloth start`` coding agent), requests carrying the client's own ``tools``
-bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small
+bypass Unsloth's tool loop and are relayed to/from llama-server verbatim. Small
GGUF models often emit their tool calls as TEXT (``{...} ``,
Gemma ``<|tool_call>...``, ```` XML) instead of structured
``tool_calls`` -- on the passthrough that text reaches the agent as prose and
@@ -18,7 +18,7 @@ promotes calls whose function name exactly matches a declared tool. Promotion
removes EXACTLY the promoted calls' markup spans (the parser reports them):
undeclared calls, unparseable blocks, and suppressed alternate formats keep
every byte and relay as text, so healing can never silently delete model
-output. Responses without a tool signal, requests without tools, and Studio's
+output. Responses without a tool signal, requests without tools, and Unsloth's
own enable-tools loop are untouched. Per-request opt-out:
``auto_heal_tool_calls: false``. Process kill-switch:
``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``.
diff --git a/studio/backend/core/inference/pricing.py b/studio/backend/core/inference/pricing.py
index 3b611d3596..30fec47723 100644
--- a/studio/backend/core/inference/pricing.py
+++ b/studio/backend/core/inference/pricing.py
@@ -122,12 +122,12 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
"priced": bool(prices),
}
- # Accept raw (input_tokens/output_tokens) and Studio chat-style
+ # Accept raw (input_tokens/output_tokens) and Unsloth chat-style
# (prompt_tokens/completion_tokens) envelopes. Cache buckets differ:
# raw Anthropic: input_tokens EXCLUDES cache buckets
# raw OpenAI: input_tokens INCLUDES cache_read
- # Studio Anthropic: prompt_tokens INCLUDES cache_creation + cache_read
- # Studio OpenAI: prompt_tokens == raw input_tokens
+ # Unsloth Anthropic: prompt_tokens INCLUDES cache_creation + cache_read
+ # Unsloth OpenAI: prompt_tokens == raw input_tokens
# Clamp >=0 so corrupted payloads can't produce a negative bill.
cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0))
cache_read_native_present = (
@@ -160,7 +160,7 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
output_tokens = max(0, int(usage.get("completion_tokens") or 0))
if provider == "openai":
# Cached tokens land on input_tokens_details (raw Responses) or
- # prompt_tokens_details (Studio chat-style).
+ # prompt_tokens_details (Unsloth chat-style).
for key in ("input_tokens_details", "prompt_tokens_details"):
details = usage.get(key) or {}
if isinstance(details, dict):
diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py
index 9110315815..40731de57b 100644
--- a/studio/backend/core/inference/safetensors_agentic.py
+++ b/studio/backend/core/inference/safetensors_agentic.py
@@ -995,7 +995,7 @@ def run_safetensors_tool_loop(
if not safety_tc:
# Re-prompt once on plan-without-action, before any tool runs
# (GGUF loop parity). The retry is gated on nudge_tool_calls so
- # Studio callers (which send True) always nudge, while API callers
+ # Unsloth callers (which send True) always nudge, while API callers
# who omit the flag keep today's no-reprompt behavior (opt-in).
intent_text = _reprompt_intent_text(
content_accum,
diff --git a/studio/backend/core/inference/sandbox_site/sitecustomize.py b/studio/backend/core/inference/sandbox_site/sitecustomize.py
index d655e8e35a..244fa95145 100644
--- a/studio/backend/core/inference/sandbox_site/sitecustomize.py
+++ b/studio/backend/core/inference/sandbox_site/sitecustomize.py
@@ -4,7 +4,7 @@
"""Sandbox-side compatibility shim for ChatGPT code-interpreter paths.
Models habitually write to /mnt/data (or /mnt/outputs, /home/sandbox,
-/workspace), none of which exist in the Studio sandbox. This module sits on the
+/workspace), none of which exist in the Unsloth sandbox. This module sits on the
sandbox subprocess PYTHONPATH (see ``tools._build_safe_env``), so it loads at
interpreter startup in every sandboxed ``python`` run and any Python the
``terminal`` tool launches.
diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py
index f7ed450d11..61643b5795 100644
--- a/studio/backend/core/inference/tool_loop_controller.py
+++ b/studio/backend/core/inference/tool_loop_controller.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
-"""Shared controller state for Studio local agentic tool loops.
+"""Shared controller state for Unsloth local agentic tool loops.
This module is intentionally dependency-light: it owns only per-response
ledger state and value objects used by the GGUF and safetensors loops.
diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index 5fd57e1b2c..bc9ffe85c2 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -2502,7 +2502,7 @@ def _build_safe_env(workdir: str) -> dict[str, str]:
shim directory.
"""
# Start from the running interpreter's dir so 'python'/'pip' resolve to the
- # same environment the Studio server runs in.
+ # same environment the Unsloth server runs in.
exe_dir = os.path.dirname(sys.executable)
path_entries = [exe_dir] if exe_dir else []
@@ -2792,7 +2792,7 @@ def _bypass_preexec():
"""Minimal pre-exec for bypass exec: os.setsid() only.
Required, not a restriction: _kill_process_tree does killpg(getpgid(child)),
- so without a new session a timeout/cancel would kill the Studio server too.
+ so without a new session a timeout/cancel would kill the Unsloth server too.
"""
try:
os.setsid()
@@ -2800,13 +2800,13 @@ def _bypass_preexec():
pass
-# Hardening the Studio parent is done once (PR_SET_DUMPABLE is process-global
+# Hardening the Unsloth parent is done once (PR_SET_DUMPABLE is process-global
# and sticky); guarded so repeated bypass calls do not re-issue the prctl.
_parent_proc_hardened = False
def _harden_parent_against_proc_env_leak() -> bool:
- """Make the Studio process's /proc//environ unreadable to its children.
+ """Make the Unsloth process's /proc//environ unreadable to its children.
Stripping the child env is not enough on Linux: a bypassed same-UID child
can read /proc//environ to recover the parent's unfiltered
@@ -5482,7 +5482,7 @@ def _truncate(text: str, limit: int = _MAX_OUTPUT_CHARS) -> str:
# ChatGPT code-interpreter path conventions models write out of habit; none
-# exist in the Studio sandbox, so a failure on one earns the retry hint.
+# exist in the Unsloth sandbox, so a failure on one earns the retry hint.
_MISSING_PATH_PREFIXES = (
"/mnt/data",
"/mnt/outputs",
@@ -5688,7 +5688,7 @@ def _python_exec(
# Close the /proc//environ secret-recovery path first; if it
# cannot be applied, fail closed rather than leak the parent environ.
return (
- "Execution error: could not harden the Studio process against "
+ "Execution error: could not harden the Unsloth process against "
"/proc environment reads; refusing bypass execution."
)
@@ -5833,7 +5833,7 @@ def _bash_exec(
# Close the /proc//environ secret-recovery path first; if it
# cannot be applied, fail closed rather than leak the parent environ.
return (
- "Execution error: could not harden the Studio process against "
+ "Execution error: could not harden the Unsloth process against "
"/proc environment reads; refusing bypass execution."
)
diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py
index 6d1512a770..8398506f21 100644
--- a/studio/backend/core/rag/captioner.py
+++ b/studio/backend/core/rag/captioner.py
@@ -6,7 +6,7 @@
Both turn pixels into indexable text and are a no-op (never raise) without a loaded
vision model. They reuse the chat model's vision endpoint, so it must be served with
``--ubatch-size`` >= one image's tokens (some encoders, e.g. Gemma, attend
-non-causally and abort otherwise); Studio's vision chat already requires this."""
+non-causally and abort otherwise); Unsloth's vision chat already requires this."""
from __future__ import annotations
diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py
index 46a282c939..b141e59422 100644
--- a/studio/backend/core/rag/embed_llama_server.py
+++ b/studio/backend/core/rag/embed_llama_server.py
@@ -10,7 +10,7 @@ Opt-in (``RAG_EMBED_BACKEND=llama-server``). Runs a dedicated
Device is ``auto`` (GPU when present, else CPU, falling back to CPU if a GPU start
fails); ``RAG_EMBED_DEVICE`` forces it. We call only llama_cpp's *static* helpers
(no torch), copying the instance-coupled bits locally, since constructing a
-``LlamaCppBackend`` runs an ``__init__`` reaper that kills any Studio llama-server
+``LlamaCppBackend`` runs an ``__init__`` reaper that kills any Unsloth llama-server
-- so each request re-spawns ours if it died (self-heal).
"""
diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py
index b0ecedd593..15be7f1249 100644
--- a/studio/backend/core/rag/embeddings.py
+++ b/studio/backend/core/rag/embeddings.py
@@ -39,7 +39,7 @@ _model = None
_name: str | None = None
-# Studio device -> torch device string. Apple has no torch device -> CPU.
+# Unsloth device -> torch device string. Apple has no torch device -> CPU.
_TORCH_DEVICE = {DeviceType.CUDA: "cuda", DeviceType.XPU: "xpu"}
diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py
index 2a4a198610..bbd9a895ab 100644
--- a/studio/backend/core/training/resume.py
+++ b/studio/backend/core/training/resume.py
@@ -53,7 +53,7 @@ def get_resume_checkpoint_path(path_value: str) -> Optional[str]:
def normalize_resume_output_dir(path_value: str) -> str:
path = resolve_output_dir(path_value)
if not _is_under_outputs(path):
- raise ValueError("Resume checkpoint must be inside Studio outputs.")
+ raise ValueError("Resume checkpoint must be inside Unsloth outputs.")
return str(path)
diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 883a535a89..26720865f4 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -797,7 +797,7 @@ class UnslothTrainer:
)
logger.info("Loaded text model")
- raise_if_offloaded(self.model, device_map, "Studio training")
+ raise_if_offloaded(self.model, device_map, "Unsloth training")
if self.should_stop:
return False
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 38f6b92f6d..b407ba39a5 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -140,7 +140,7 @@ def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool:
def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]:
- """Build the normalized worker config shared by Studio and the CLI adapter."""
+ """Build the normalized worker config shared by Unsloth and the CLI adapter."""
config = {
"model_name": values["model_name"],
"project_name": values.get("project_name"),
@@ -307,7 +307,7 @@ PLOT_HEIGHT = 3.5
@dataclass
class TrainingProgress:
- """Shared training progress payload for Studio and backend-aware trainers."""
+ """Shared training progress payload for Unsloth and backend-aware trainers."""
epoch: float = 0
step: int = 0
@@ -328,7 +328,7 @@ class TrainingProgress:
class _MLXTrainerAdapter:
- """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path."""
+ """Adapts the legacy UnslothTrainer API to the shared Unsloth MLX worker path."""
def __init__(self):
self.model = None
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index c52adbe8fa..111f4fdd0f 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -1100,7 +1100,7 @@ _MLX_VLM_RESIZED_IMAGE_LAYOUT_CACHE = {}
def _mlx_vlm_resized_image_layout(processor = None) -> str | None:
- """Return the numpy image layout expected after Studio-side VLM resizing."""
+ """Return the numpy image layout expected after Unsloth-side VLM resizing."""
image_processor = getattr(processor, "image_processor", None)
if image_processor is None:
return None
@@ -1257,7 +1257,7 @@ _MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"}
# Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used
-# only when mlx (Apple Silicon) is not importable so Studio config validation
+# only when mlx (Apple Silicon) is not importable so Unsloth config validation
# still works on non-MLX hosts. The zoo function stays the source of truth.
_MLX_STUDIO_ADAMW_ALIASES = frozenset(
(
@@ -1309,7 +1309,7 @@ def _normalize_mlx_studio_scheduler(value):
def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]:
- """Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer."""
+ """Resolve CLI paths and Unsloth local dataset uploads without importing the GPU trainer."""
from utils.paths import resolve_dataset_path
all_files: list[str] = []
@@ -1912,7 +1912,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
if "max_grad_leaf_norm" in _supported_fields:
mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm
if "append_eos" in _supported_fields:
- # Studio SFT formatting owns rendered examples; raw/CPT text still
+ # Unsloth SFT formatting owns rendered examples; raw/CPT text still
# needs MLX to append EOS like the CUDA raw-text path.
mlx_config_kwargs["append_eos"] = bool(raw_text_mode)
@@ -2121,7 +2121,7 @@ def run_mlx_training_process(
config: dict,
transformers_activated: bool = False,
) -> None:
- """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter."""
+ """MLX worker entrypoint shared by Unsloth subprocesses and the CLI adapter."""
model_name = config["model_name"]
backend_path = str(Path(__file__).resolve().parent.parent.parent)
@@ -2780,7 +2780,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
# Unified Windows APUs: the WDDM budget is user-raisable, but
# nothing on the box says so -- users see "48 GB VRAM" on a
- # 96 GB machine and assume a Studio bug. Say where the limit
+ # 96 GB machine and assume an Unsloth bug. Say where the limit
# comes from and how to raise it.
if _is_unified and sys.platform == "win32":
try:
diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py
index e5d48872c1..23f8c7c911 100644
--- a/studio/backend/hub/services/download_lifecycle.py
+++ b/studio/backend/hub/services/download_lifecycle.py
@@ -76,7 +76,7 @@ def spawn_worker(
env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
env["HF_HUB_DISABLE_TELEMETRY"] = "1"
env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1"
- # No token in Studio settings: fall back to the backend's own HF_TOKEN so
+ # No token in Unsloth settings: fall back to the backend's own HF_TOKEN so
# private repos stay downloadable (needed while inkling repos are private).
if not hf_token:
hf_token = os.environ.get("HF_TOKEN") or None
diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py
index d56b62c318..7d9c3ac665 100644
--- a/studio/backend/hub/services/models/folder_browser.py
+++ b/studio/backend/hub/services/models/folder_browser.py
@@ -165,7 +165,7 @@ def _looks_like_model_dir(directory: Path) -> bool:
def _build_browse_allowlist(
media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None
) -> list[Path]:
- """Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary.
+ """Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Unsloth outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary.
*media_roots* / *drive_roots* let the caller pass already-probed
removable-media and Windows drive roots so they aren't scanned again (a
diff --git a/studio/backend/hub/services/models/ollama.py b/studio/backend/hub/services/models/ollama.py
index 96a4114620..2ccdbb44f1 100644
--- a/studio/backend/hub/services/models/ollama.py
+++ b/studio/backend/hub/services/models/ollama.py
@@ -85,7 +85,7 @@ def _contained_link_path(link_dir: Path, link_name: str) -> Optional[Path]:
def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
- """Writable directory for Ollama ``.gguf`` symlinks. Prefers ``/.studio_links/`` next to the blobs; falls back to Studio's cache (read-only system installs), then the temp dir (sandboxed installs)."""
+ """Writable directory for Ollama ``.gguf`` symlinks. Prefers ``/.studio_links/`` next to the blobs; falls back to Unsloth's cache (read-only system installs), then the temp dir (sandboxed installs)."""
def _ensure_writable_dir(path: Path) -> Optional[Path]:
try:
diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py
index 183e934724..898c03c87d 100644
--- a/studio/backend/hub/utils/state_dir.py
+++ b/studio/backend/hub/utils/state_dir.py
@@ -3,7 +3,7 @@
"""Filesystem layout for Hub download state.
-State directory sits beside HF's cache (under Studio's own cache root)
+State directory sits beside HF's cache (under Unsloth's own cache root)
so it survives ``huggingface-cli delete-cache`` and any other HF-side
cache lifecycle. Two subdirectories:
diff --git a/studio/backend/main.py b/studio/backend/main.py
index 6e16dc00ca..4797764ce7 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -19,7 +19,7 @@ os.environ["PYTHONWARNINGS"] = "ignore"
# Pin GPU index ordering to PCI bus id before any torch import creates a CUDA
# context. Without this, torch/CUDA default to FASTEST_FIRST while nvidia-smi
-# (and Studio's VRAM probes) use PCI-bus order, so a GPU index chosen from
+# (and Unsloth's VRAM probes) use PCI-bus order, so a GPU index chosen from
# nvidia-smi data can resolve to a different physical card via
# CUDA_VISIBLE_DEVICES. setdefault so an explicit user override wins. See
# utils/hardware/hardware.py for the full rationale; set here too so the entry
@@ -93,7 +93,7 @@ if sys.platform == "win32":
# ── Windows AMD ROCm: make hipInfo.exe resolvable for subprocess probes ──
# bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via PATH at import
# time; the AMD torch wheel ships it in the venv Scripts dir, which is on
- # PATH only when the venv is activated -- Studio launches python directly.
+ # PATH only when the venv is activated -- Unsloth launches python directly.
# Without this, every bitsandbytes import logs a scary (but harmless)
# "Could not detect ROCm GPU architecture: [WinError 2]" ERROR + WARNING.
# Gated on the file existing: only AMD ROCm wheels ship hipInfo.exe, so
@@ -252,7 +252,7 @@ def _read_studio_install_id() -> str:
Returns "" when absent or not a 64-char lowercase-hex token; then
/api/health emits "" and the launcher accepts any healthy backend.
- Carries no install-path info (matters when Studio runs -H 0.0.0.0)."""
+ Carries no install-path info (matters when Unsloth runs -H 0.0.0.0)."""
try:
token = (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
except (OSError, ValueError):
@@ -573,7 +573,7 @@ async def lifespan(app: FastAPI):
print("DEFAULT ADMIN ACCOUNT CREATED")
print(f" username: {storage.DEFAULT_ADMIN_USERNAME}")
print(f" password saved to: {bootstrap_path}")
- print(" Open the Studio UI to sign in and change it.")
+ print(" Open the Unsloth UI to sign in and change it.")
print("=" * 60 + "\n")
else:
app.state.bootstrap_password = (
@@ -613,7 +613,7 @@ app = FastAPI(
)
# The MCP surface is opt-in because it can start GPU jobs and write model
-# artifacts. Mount it only when explicitly enabled by the Studio process.
+# artifacts. Mount it only when explicitly enabled by the Unsloth process.
if os.environ.get("UNSLOTH_STUDIO_ENABLE_MCP") == "1":
from fastmcp.utilities.lifespan import combine_lifespans
@@ -973,7 +973,7 @@ app.include_router(training_router, prefix = "/api/train", tags = ["training"])
app.include_router(models_router, prefix = "/api/models", tags = ["models"])
app.include_router(chat_history_router, prefix = "/api/chat", tags = ["chat"])
app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
-# Studio-only inference endpoints (cancel, etc.) are NOT exposed on the /v1
+# Unsloth-only inference endpoints (cancel, etc.) are NOT exposed on the /v1
# OpenAI-compat prefix below.
app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"])
@@ -1080,7 +1080,7 @@ def studio_install_source(_current_subject: str = Depends(get_current_subject)):
@app.get("/api/studio/update-status")
def studio_update_status(_current_subject: str = Depends(get_current_subject)):
- """Return source-aware manual update status for browser-served Studio."""
+ """Return source-aware manual update status for browser-served Unsloth."""
return get_studio_update_status(UNSLOTH_VERSION)
diff --git a/studio/backend/mcp_server.py b/studio/backend/mcp_server.py
index f837f46425..e93490411d 100644
--- a/studio/backend/mcp_server.py
+++ b/studio/backend/mcp_server.py
@@ -3,7 +3,7 @@
"""Curated MCP tools for driving an Unsloth Studio instance.
-The MCP surface deliberately wraps the existing Studio services instead of
+The MCP surface deliberately wraps the existing Unsloth services instead of
duplicating training or export logic. It is opt-in because several tools can
start GPU work or write model artifacts.
"""
@@ -17,14 +17,14 @@ from fastmcp import FastMCP
class BearerTokenMiddleware:
- """Require an exact bearer token when Studio MCP is exposed remotely."""
+ """Require an exact bearer token when Unsloth MCP is exposed remotely."""
def __init__(self, app: Any, token: str) -> None:
if not token or not token.strip():
- raise ValueError("Studio MCP bearer token must be a non-empty value")
+ raise ValueError("Unsloth MCP bearer token must be a non-empty value")
if not token.isascii():
# A non-ASCII token cannot be sent in an HTTP header; reject it here.
- raise ValueError("Studio MCP bearer token must contain ASCII characters only")
+ raise ValueError("Unsloth MCP bearer token must contain ASCII characters only")
self.app = app
# Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII
# input, which would surface as a 500 instead of a clean 401.
@@ -76,18 +76,18 @@ def _dump(value: Any) -> Any:
def _clamp(value: int, low: int, high: int) -> int:
"""Clamp an MCP-supplied integer into an inclusive range.
- MCP tools call the Studio route functions directly, which skips FastAPI's
+ MCP tools call the Unsloth route functions directly, which skips FastAPI's
Query(ge=, le=) validation, so we re-apply the same bounds here.
"""
return max(low, min(value, high))
def create_studio_mcp() -> FastMCP:
- """Create the Studio MCP server and register the high-value tools."""
+ """Create the Unsloth MCP server and register the high-value tools."""
mcp = FastMCP(
"Unsloth Studio",
instructions = (
- "Use read tools to inspect the local Studio state before starting GPU work. "
+ "Use read tools to inspect the local Unsloth state before starting GPU work. "
"Training and export tools can consume substantial VRAM and write files. "
"Never expose tokens or local paths from tool results unless the user asks."
),
@@ -116,7 +116,7 @@ def create_studio_mcp() -> FastMCP:
@mcp.tool
async def list_local_models(models_dir: str = "./models") -> dict[str, Any]:
- """List local and cached models available to Studio."""
+ """List local and cached models available to Unsloth."""
from routes.models import list_local_models as list_models
return _dump(await list_models(models_dir = models_dir, current_subject = "mcp"))
@@ -128,9 +128,9 @@ def create_studio_mcp() -> FastMCP:
@mcp.tool
async def start_training(config: dict[str, Any]) -> dict[str, Any]:
- """Start a validated Studio training job from a TrainingStartRequest-shaped object.
+ """Start a validated Unsloth training job from a TrainingStartRequest-shaped object.
- The config is validated by the same Pydantic model used by the Studio UI.
+ The config is validated by the same Pydantic model used by the Unsloth UI.
Call get_training_status first and do not start work while another job runs.
"""
from models import TrainingStartRequest
@@ -138,7 +138,7 @@ def create_studio_mcp() -> FastMCP:
request = TrainingStartRequest.model_validate(config)
# Pass via_api_key explicitly (a direct call leaves it a Depends object).
- # MCP drives Studio like the UI session, so it coexists and frees VRAM.
+ # MCP drives Unsloth like the UI session, so it coexists and frees VRAM.
return _dump(await start(request, current_subject = "mcp", via_api_key = False))
@mcp.tool
@@ -159,7 +159,7 @@ def create_studio_mcp() -> FastMCP:
@mcp.tool
def validate_recipe(recipe: dict[str, Any]) -> dict[str, Any]:
- """Validate a Data Recipe with the same validator used by Studio."""
+ """Validate a Data Recipe with the same validator used by Unsloth."""
from models.data_recipe import RecipePayload
from routes.data_recipe.validate import validate
@@ -225,7 +225,7 @@ def create_studio_mcp() -> FastMCP:
imatrix: bool = False,
imatrix_path: str | None = None,
) -> dict[str, Any]:
- """Export the loaded model to GGUF using Studio's existing path validation.
+ """Export the loaded model to GGUF using Unsloth's existing path validation.
quantization_method may be a single method or a list to produce several
GGUFs from one load. Pass hf_token when push_to_hub is set (the backend
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 3ae974448e..f3ae0f70df 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -105,7 +105,7 @@ class LoadRequest(BaseModel):
description = (
"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, "
+ "Unsloth-managed flags (model identity, port, context length, GPU placement, "
"auth, UI/server mode) are rejected. Ignored for non-GGUF models."
),
)
@@ -151,13 +151,13 @@ class TransformersUpgradeInfo(BaseModel):
)
supported_in_pypi: bool = Field(
False,
- description = "True if the latest PyPI release ships this model_type; Studio can "
+ description = "True if the latest PyPI release ships this model_type; Unsloth can "
"install it into a persistent sidecar after user consent.",
)
supported_in_main: bool = Field(
False,
description = "True if transformers GitHub main ships this model_type (dev-only; "
- "not installable through Studio yet).",
+ "not installable through Unsloth yet).",
)
@@ -533,7 +533,7 @@ class ImageContentPart(BaseModel):
class InputDocumentContentPart(BaseModel):
"""Document (PDF / file) content part in a multimodal message.
- Studio-normalised shape (file_data or file_url, plus optional filename/media_type).
+ Unsloth-normalised shape (file_data or file_url, plus optional filename/media_type).
Mapped onto Anthropic ``document`` / OpenAI ``input_file`` for vision providers;
dropped for non-vision providers.
"""
@@ -689,7 +689,7 @@ class ThinkingConfig(BaseModel):
"""Anthropic-compatible thinking/reasoning configuration.
Use type='disabled' to turn off thinking, or type='enabled' to turn it on.
Only type is read; extra fields (e.g. budget_tokens) are ignored, since
- Studio sets provider thinking budgets itself.
+ Unsloth sets provider thinking budgets itself.
"""
type: Literal["disabled", "enabled"] = "disabled"
@@ -748,7 +748,7 @@ class ChatCompletionRequest(BaseModel):
None,
description = (
"OpenAI function-tool definitions. When provided without `enable_tools=true`, "
- "Studio forwards the tools to the backend so the model returns structured "
+ "Unsloth forwards the tools to the backend so the model returns structured "
"tool_calls for the client to execute (standard OpenAI function calling)."
),
)
@@ -1160,7 +1160,7 @@ class ChatCompletionRequest(BaseModel):
and (self.enable_tools is True or bool(self.mcp_enabled))
):
# "Ask" gates every call, so a direct API caller that omits the legacy
- # confirm flag must still hit the confirmation gate for Studio's own
+ # confirm flag must still hit the confirmation gate for Unsloth's own
# tool loop. An explicit confirm_tool_calls=False wins over the mode
# (mirrors _permission_mode_confirm and the Anthropic pre-switch guard),
# so only self-enable when the flag is unset. Only self-enable when that
@@ -1168,7 +1168,7 @@ class ChatCompletionRequest(BaseModel):
# (enable_tools / mcp_enabled) -- the router enters the loop on those
# signals, not on enabled_tools alone (which merely filters which tools
# run). A plain client-tool passthrough (client-supplied `tools` that
- # Studio does not execute) must route verbatim, and external-provider
+ # Unsloth does not execute) must route verbatim, and external-provider
# routing rejects confirm_tool_calls with tools, so skip the fold there.
#
# "auto" is deliberately NOT folded: it only prompts for a call the
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index ff815a2fa9..0b50f63b95 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -446,7 +446,7 @@ class TrainingStartRequest(BaseModel):
random_seed: int = Field(
3407,
description = (
- "Random seed; matches the Studio backend / MLX worker default "
+ "Random seed; matches the Unsloth backend / MLX worker default "
"and unsloth's historical recommended value."
),
)
diff --git a/studio/backend/plugins/data-designer-github-repo-seed/README.md b/studio/backend/plugins/data-designer-github-repo-seed/README.md
index 346d94b305..44519496f5 100644
--- a/studio/backend/plugins/data-designer-github-repo-seed/README.md
+++ b/studio/backend/plugins/data-designer-github-repo-seed/README.md
@@ -4,7 +4,7 @@ A Data Designer seed-reader plugin for **Unsloth Studio** that scrapes real
GitHub data (issues, pull requests, commits) from one or more repositories
and hands it to the recipe pipeline as a seed dataset.
-Designed to ship with Studio as a default seed source so any user with a
+Designed to ship with Unsloth as a default seed source so any user with a
GitHub token can build training datasets straight from live repos.
## What it does
@@ -64,7 +64,7 @@ sleeps until reset when the budget drops below a safety threshold.
## Install
-Shipped as a default Studio plugin. For development:
+Shipped as a default Unsloth plugin. For development:
```bash
pip install -e .
diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py
index 62ecb2e280..d4d46da370 100644
--- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py
+++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py
@@ -3,4 +3,4 @@
# Intentionally empty. Data-designer loads submodules lazily via qualified names
# in plugin.py, so importing this package must not touch data_designer.engine.*
-# during Studio bootstrap (circular import).
+# during Unsloth bootstrap (circular import).
diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py
index 637193e8b3..1af8133cc5 100644
--- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py
+++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.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
-"""Multi-repo GitHub scraper for the Studio seed plugin.
+"""Multi-repo GitHub scraper for the Unsloth seed plugin.
Drives the GraphQL scraper in `scraper_impl/` per repo, capped via trial_limits
to stop at `limit` items per resource. Then reads the per-resource JSONL shards
diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt
index 5830a47789..3361af50dd 100644
--- a/studio/backend/requirements/extras-no-deps.txt
+++ b/studio/backend/requirements/extras-no-deps.txt
@@ -5,7 +5,7 @@ julius
torchcodec==0.10.0
snac
-# peft 0.19.0 causes export subprocess shutdown issues in Studio;
+# peft 0.19.0 causes export subprocess shutdown issues in Unsloth;
# installing with --no-deps to avoid pulling in torch>=0.11.0
peft==0.18.1
diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt
index de321f80ed..378fb33a60 100644
--- a/studio/backend/requirements/no-torch-runtime.txt
+++ b/studio/backend/requirements/no-torch-runtime.txt
@@ -70,7 +70,7 @@ cut_cross_entropy
pillow
# RAG store + document parsing, mirroring studio.txt. Pinned here because
-# this file installs --no-deps; without them Studio runs with RAG disabled.
+# this file installs --no-deps; without them Unsloth runs with RAG disabled.
sqlite-vec==0.1.9
pymupdf==1.27.2.3
# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the
diff --git a/studio/backend/requirements/single-env/constraints.txt b/studio/backend/requirements/single-env/constraints.txt
index 0ed2bf8b26..0a5619924a 100644
--- a/studio/backend/requirements/single-env/constraints.txt
+++ b/studio/backend/requirements/single-env/constraints.txt
@@ -4,7 +4,7 @@ transformers==4.57.6
trl==0.23.1
huggingface-hub==0.36.2
-# Studio stack
+# Unsloth stack
datasets==4.3.0
pyarrow==23.0.1
diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt
index 6f4a5c3292..0c7503a5ca 100644
--- a/studio/backend/requirements/studio.txt
+++ b/studio/backend/requirements/studio.txt
@@ -1,4 +1,4 @@
-# Studio UI backend dependencies
+# Unsloth UI backend dependencies
typer
fastapi
uvicorn
@@ -9,7 +9,7 @@ pandas
nest_asyncio
datasets==4.3.0
pyjwt
-# gradio>=4.0.0 # 148 MB - Studio uses React + FastAPI, not Gradio
+# gradio>=4.0.0 # 148 MB - Unsloth uses React + FastAPI, not Gradio
huggingface-hub==0.36.2
structlog>=24.1.0
diceware
diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py
index c61c1a16e4..d779c8784e 100644
--- a/studio/backend/routes/auth.py
+++ b/studio/backend/routes/auth.py
@@ -338,11 +338,11 @@ def _clear_login_bucket(key: tuple[str, str]) -> None:
# so FastAPI runs it in the threadpool rather than blocking the event loop.
@router.get("/identity")
def identity(nonce: str, request: Request) -> dict:
- """Challenge-response proof this is the real local Studio: caller sends a nonce,
+ """Challenge-response proof this is the real local Unsloth: caller sends a nonce,
gets HMAC(install identity secret, nonce, connection address + port).
Unauthenticated and side-effect free; a process that can't read the same-user
secret can't forge a proof, and binding to the address/port the connection
- landed on stops a squatter relaying a proof from the real Studio elsewhere."""
+ landed on stops a squatter relaying a proof from the real Unsloth elsewhere."""
try:
raw = base64.urlsafe_b64decode(nonce)
except Exception:
diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py
index 59714380da..e870e8855e 100644
--- a/studio/backend/routes/data_recipe/jobs.py
+++ b/studio/backend/routes/data_recipe/jobs.py
@@ -37,7 +37,7 @@ def _resolve_local_v1_endpoint(request: Request) -> str:
Resolution order:
1. ``app.state.server_port`` (run.py, post-bind) - survives proxies/tunnels.
- 2. ``request.scope["server"]`` - when Studio starts outside ``run_server``.
+ 2. ``request.scope["server"]`` - when Unsloth starts outside ``run_server``.
3. parsed ``request.base_url`` - last resort for test fixtures.
"""
port: Any = getattr(request.app.state, "server_port", None)
diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py
index 46319ca2ba..5456080f34 100644
--- a/studio/backend/routes/datasets.py
+++ b/studio/backend/routes/datasets.py
@@ -485,7 +485,7 @@ async def upload_dataset(
# Stream to disk in chunks to avoid holding the whole file in memory. The
# route-level cap gives a clear training-dataset error and avoids leaving
- # oversized partial files in the Studio uploads directory.
+ # oversized partial files in the Unsloth uploads directory.
upload_limit_bytes = get_upload_limit_bytes()
total_bytes = 0
upload_complete = False
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 9299e26d56..3d527bf317 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -92,7 +92,7 @@ def _mlx_distributed_launch_detected() -> bool:
def _install_httpcore_asyncgen_silencer() -> None:
"""Silence benign httpx/httpcore asyncgen GC noise on Python 3.13.
- When Studio proxies a llama-server stream via httpx, the innermost
+ When Unsloth proxies a llama-server stream via httpx, the innermost
``HTTP11ConnectionByteStream.__aiter__`` async generator is finalised by
the asyncgen GC hook on a task different from the one that opened it. Its
``aclose`` calls ``anyio.Lock.acquire`` → ``cancel_shielded_checkpoint``,
@@ -229,14 +229,14 @@ def _friendly_upstream_error(text: str) -> str:
parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as
a hard 400 on every tool-bearing turn. It is a llama-server limitation with some
model/quant + tool-schema combinations, and recent llama.cpp builds handle the common
- coding-agent tools, so point the user at updating Studio rather than the raw body.
+ coding-agent tools, so point the user at updating Unsloth rather than the raw body.
"""
lowered = text.lower()
if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered:
return (
"The model couldn't compile a tool-calling grammar for this request. This is a "
"llama-server limitation with some model/quant and tool-schema combinations. "
- "Update Studio (it installs the latest llama.cpp, which handles the common "
+ "Update Unsloth (it installs the latest llama.cpp, which handles the common "
"coding-agent tools) or try a different GGUF model."
)
return f"llama-server error: {text}"
@@ -731,7 +731,7 @@ def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]:
Some llama-server builds can emit the logical final chunk (``finish_reason``)
and optional usage chunk, then keep the HTTP stream open without sending the
- OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close
+ OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Unsloth close
the client stream promptly while preserving an optional trailing usage chunk.
"""
if not raw_line.startswith("data:"):
@@ -1786,7 +1786,7 @@ import numpy as np
from datetime import date as _date
router = APIRouter()
-# Studio-only router (not mounted on /v1 OpenAI-compat).
+# Unsloth-only router (not mounted on /v1 OpenAI-compat).
studio_router = APIRouter()
@@ -2108,9 +2108,9 @@ def _effective_enable_tools(payload) -> Optional[bool]:
def _explicit_studio_tool_loop_requested(payload) -> bool:
- """True when the request itself asks Studio to execute local tools.
+ """True when the request itself asks Unsloth to execute local tools.
- Process-wide CLI policy can default Studio's tool loop on for ordinary chat,
+ Process-wide CLI policy can default Unsloth's tool loop on for ordinary chat,
but it must not steal OpenAI-compatible client tools or response_format
requests from the llama-server passthrough path. A policy of ``False``
(--disable-tools) vetoes even an explicit ``enable_tools: true`` ask.
@@ -2122,7 +2122,7 @@ def _explicit_studio_tool_loop_requested(payload) -> bool:
def _permission_mode_confirm(payload) -> bool:
- """Effective confirm-gate intent for Studio's own local tool loop.
+ """Effective confirm-gate intent for Unsloth's own local tool loop.
Honors the documented default that an unset permission_mode behaves as
"ask". An explicit confirm_tool_calls (True or False) wins; explicit
@@ -2144,7 +2144,7 @@ def _permission_mode_confirm(payload) -> bool:
def _confirm_gate_needs_stream(payload) -> bool:
- """Whether Studio's local tool-loop confirm gate still requires stream=true.
+ """Whether Unsloth's local tool-loop confirm gate still requires stream=true.
The gate can only prompt while streaming, so a non-streaming request that will
prompt must 400 up front. auto ("Approve for me") only prompts for a call the
@@ -3143,7 +3143,7 @@ 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
+ A bare tensor_parallel field is NOT a drop: the Unsloth 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
@@ -4197,8 +4197,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
raise HTTPException(
status_code = 400,
detail = (
- "Studio does not support distributed MLX inference under "
- "mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio "
+ "Unsloth does not support distributed MLX inference under "
+ "mlx.launch. Use `mlx.launch ... unsloth chat` or run Unsloth "
"without the distributed launcher."
),
)
@@ -4288,7 +4288,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# omits chat_template_override, so strip the inherited
# --chat-template-file in that case too -- otherwise the stale
# extra arg (appended last) shadows the bundled template while
- # Studio reports the bundled template's capabilities.
+ # Unsloth reports the bundled template's capabilities.
fields_set = getattr(request, "model_fields_set", set())
stripped = strip_shadowing_flags(
llama_backend.extra_args,
@@ -4727,7 +4727,7 @@ def _requires_trust_remote_code_for_model(
model_identifier: str, hf_token: Optional[str] = None
) -> bool:
"""Whether loading this model would execute custom repo code, so the consent
- dialog must run first. True if the Studio YAML default enables
+ dialog must run first. True if the Unsloth YAML default enables
``trust_remote_code`` OR the raw config declares an ``auto_map`` (Hub/local,
config.json or tokenizer_config.json). Reads raw JSON only; never imports
model code."""
@@ -5360,7 +5360,7 @@ async def confirm_tool_call(
@studio_router.get("/monitor")
async def get_api_monitor(current_subject: str = Depends(get_current_subject)):
- """Return recent OpenAI-compatible API activity for Studio."""
+ """Return recent OpenAI-compatible API activity for Unsloth."""
active_model = _monitor_active_model()
active_requests = api_monitor.active_count(subject = current_subject)
if active_requests:
@@ -5548,7 +5548,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
_display_model_id = os.path.basename(_model_id)
_inference_cfg = load_inference_config(_model_id) if _model_id else None
_audio_type = getattr(llama_backend, "_audio_type", None)
- # Don't surface Studio's auto-applied bundled family template (e.g. the
+ # Don't surface Unsloth's auto-applied bundled family template (e.g. the
# gemma-4 override) as a user-authored override: the frontend adopts
# status.chat_template_override as editable state and would otherwise
# re-send it as an explicit override for a later, unrelated model. Only
@@ -6173,7 +6173,7 @@ def _build_external_messages(
metadata; strip it for providers that can't parse the unknown key.
2. Marked server-side builtin cards (`_server_tool: true` on a
canonical builtin name, or a Gemini `native_part` payload) are
- Studio-internal tool cards from a prior native Gemini turn;
+ Unsloth-internal tool cards from a prior native Gemini turn;
forwarding them to OpenAI / Anthropic / custom OAI-compat gateways
sends an orphan `tool_calls` entry (no matching tool declaration,
often no matching `role="tool"` reply) that can be rejected. We
@@ -6859,7 +6859,7 @@ async def openai_chat_completions(
# is invalid and must not evict the resident model first.
#
# Enter the local-loop arm exactly when the passthrough router below would
- # run Studio's own tool loop. That gate is `_tools_on or _mcp_allowed`
+ # run Unsloth's own tool loop. That gate is `_tools_on or _mcp_allowed`
# (see the use_tools block): _effective_enable_tools (which lets a
# process-wide --enable-tools policy force the loop on) plus mcp_enabled
# honoring --disable-tools, and tool_choice="none" disabling it unless the
@@ -6884,7 +6884,7 @@ async def openai_chat_completions(
or bool(payload.openai_code_exec_container_id)
or bool(payload.anthropic_code_exec_container_id)
# A JSON-schema response_format is guided-decoding structured output the
- # router forwards to the llama-server passthrough, not Studio's tool
+ # router forwards to the llama-server passthrough, not Unsloth's tool
# loop, so a --enable-tools policy must not 400 it as a local-confirm
# request under ask/auto.
or bool(_extract_response_format(payload))
@@ -6961,7 +6961,7 @@ async def openai_chat_completions(
using_gguf = llama_backend.is_loaded
# OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, which
- # the SDK spreads into the request body at the top level. Studio's
+ # the SDK spreads into the request body at the top level. Unsloth's
# ChatCompletionRequest has ``extra="allow"`` so pydantic stashes them in
# ``model_extra``, but downstream generators consume the typed
# ``payload.enable_thinking``. Lift ``enable_thinking`` from the extra-body
@@ -7215,7 +7215,7 @@ async def openai_chat_completions(
# ── Standard OpenAI function-calling pass-through (GGUF only) ────
# When a client (opencode / Claude Code via OpenAI compat / Cursor /
- # Continue / ...) sends standard OpenAI `tools` without Studio's
+ # Continue / ...) sends standard OpenAI `tools` without Unsloth's
# `enable_tools` shorthand, forward the request to llama-server
# verbatim so structured `tool_calls` flow back to the client. This
# branch runs BEFORE `_extract_content_parts` because that helper is
@@ -7238,7 +7238,7 @@ async def openai_chat_completions(
_has_tool_catalog = bool(payload.tools and len(payload.tools) > 0)
_has_active_tool_catalog = _has_tool_catalog and payload.tool_choice != "none"
_has_client_tool_contract = _has_active_tool_catalog or _has_tool_messages
- # The Studio tool loop needs a tool-capable backend, so a request that asks
+ # The Unsloth tool loop needs a tool-capable backend, so a request that asks
# for it on a backend that can't run it (DiffusionGemma forces supports_tools
# off) must not steal client tools from the passthrough (#6851).
_studio_tool_loop_requested = (
@@ -7434,7 +7434,7 @@ async def openai_chat_completions(
use_tools = False
if use_tools:
- # permission_mode ask/auto require the confirm gate for Studio's own
+ # permission_mode ask/auto require the confirm gate for Unsloth's own
# tool loop. The request validator self-enables confirm only for
# request-level tool signals (enable_tools/enabled_tools/mcp_enabled);
# when a CLI policy (--enable-tools) forces the loop on without those,
@@ -8711,7 +8711,7 @@ async def openai_chat_completions(
_sf_model_info = backend.models.get(backend.active_model_name, {})
_sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template")
# Named templates may expose native reasoning only in their ``tool_use``
- # branch. Use a truthy placeholder for Studio-managed tools, whose concrete
+ # branch. Use a truthy placeholder for Unsloth-managed tools, whose concrete
# schemas are selected below, and the request schemas for client passthrough.
_sf_server_tool_intent = bool(
_effective_enable_tools(payload) or _explicit_studio_tool_loop_requested(payload)
@@ -8790,7 +8790,7 @@ async def openai_chat_completions(
_sf_use_tools = False
if _sf_use_tools:
- # permission_mode ask/auto require the confirm gate for Studio's own tool
+ # permission_mode ask/auto require the confirm gate for Unsloth's own tool
# loop; when a CLI policy (--enable-tools) forces the loop on without a
# request-level tool signal, derive confirm here so the mode still gates
# the call (matching the GGUF path). off/full never prompt.
@@ -12050,7 +12050,7 @@ def _anthropic_requested_studio_tools(tools: Optional[list]) -> set[str]:
def _select_anthropic_server_tools(
all_tools: list[dict], requested_studio_tools: set[str], enabled_tools: Optional[list[str]]
) -> list[dict]:
- """Select Studio tools requested through Anthropic tools and extensions."""
+ """Select Unsloth tools requested through Anthropic tools and extensions."""
if not requested_studio_tools and enabled_tools is None:
return all_tools
@@ -12289,7 +12289,7 @@ async def anthropic_messages(
),
)
- # Reject an unsupported confirm-gated permission mode for Studio's own
+ # Reject an unsupported confirm-gated permission mode for Unsloth's own
# ("server") Anthropic tools before the switch, mirroring the malformed- and
# mixed-tool checks above. ask always wants a per-call pause this passthrough
# cannot offer, so it 400s whenever server tools are selected. auto only needs
@@ -12778,11 +12778,11 @@ async def _anthropic_tool_stream(
ends_on_tool_use = True
elif etype == "tool_end":
tool_blocks_emitted += 1
- # A tool_end means Studio executed the tool server-side, so
+ # A tool_end means Unsloth executed the tool server-side, so
# the response no longer ends on a pending client action.
# Without this, a server tool that produces no trailing text
# would be mislabeled stop_reason "tool_use", telling the
- # client to run a tool Studio already ran.
+ # client to run a tool Unsloth already ran.
ends_on_tool_use = False
elif etype == "content" and event.get("text"):
ends_on_tool_use = False
@@ -13698,7 +13698,7 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
structured ``tool_calls``. Content-parts images already in the list are
left untouched.
- When a client uses Studio's legacy ``image_base64`` top-level field, the
+ When a client uses Unsloth's legacy ``image_base64`` top-level field, the
image is re-encoded to PNG (llama-server's stb_image has limited format
support) and spliced into the last user message as an OpenAI ``image_url``
content part so vision + function-calling requests work transparently.
@@ -13832,7 +13832,7 @@ def _build_openai_passthrough_body(
) -> dict:
"""Assemble the llama-server request body from a ChatCompletionRequest.
- Only known OpenAI / llama-server fields are forwarded, so Studio-specific
+ Only known OpenAI / llama-server fields are forwarded, so Unsloth-specific
extensions (``enable_tools``, ``enabled_tools``, ``session_id``, ...) never
leak to the backend.
"""
@@ -14082,7 +14082,7 @@ async def _openai_passthrough_stream_admitted(
admission_lease: LlamaAdmissionLease,
tracker,
):
- """Streaming client-side pass-through after Studio granted an upstream slot.
+ """Streaming client-side pass-through after Unsloth granted an upstream slot.
Forwards the client's OpenAI function-calling request to llama-server and
relays the SSE stream back with minimal normalization (reasoning-only
diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py
index 71f0fd2874..dc018d163a 100644
--- a/studio/backend/routes/mcp_servers.py
+++ b/studio/backend/routes/mcp_servers.py
@@ -82,7 +82,7 @@ def _validate_url(url: str) -> str:
if _looks_like_command(trimmed):
detail = (
"Local commands aren't enabled on this server. To allow them, "
- "set UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 and restart Studio, or use "
+ "set UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 and restart Unsloth, or use "
"an http:// or https:// URL instead."
)
else:
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index b8526c75e7..742ecde3ba 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -544,7 +544,7 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
"""Return a writable directory for Ollama ``.gguf`` symlinks.
Prefers ``/.studio_links/`` so links sit next to their
- blobs; falls back to a per-ollama-dir namespace under Studio's cache
+ blobs; falls back to a per-ollama-dir namespace under Unsloth's cache
when the models dir is read-only (common for system installs).
"""
from utils.paths.storage_roots import cache_root
@@ -555,7 +555,7 @@ def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]:
return primary
except OSError as e:
logger.debug(
- "Ollama dir %s not writable for .studio_links (%s); falling back to Studio cache",
+ "Ollama dir %s not writable for .studio_links (%s); falling back to Unsloth cache",
ollama_dir,
e,
)
@@ -594,7 +594,7 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
model, keyed by a short hash of the manifest path, so
``detect_mmproj_file`` only sees that model's projector). Links are
symlinks when possible, else hardlinks; the link dir is
- ``.studio_links/`` when writable, else Studio's cache.
+ ``.studio_links/`` when writable, else Unsloth's cache.
"""
manifests_root = ollama_dir / "manifests"
if not manifests_root.is_dir():
@@ -1194,7 +1194,7 @@ def _build_browse_allowlist(
"""Return the root directories the folder browser may walk.
The same list seeds the sidebar suggestion chips, so chip targets are
- always reachable. Roots: HOME, resolved HF cache dirs, Studio's
+ always reachable. Roots: HOME, resolved HF cache dirs, Unsloth's
outputs/exports/studio root, registered scan folders, and well-known
local-LLM dirs (LM Studio, Ollama, ``~/models``); each added only if
it resolves to a real directory.
@@ -1486,7 +1486,7 @@ def browse_folders(
"Directory to list. If omitted, defaults to the current user's "
"home directory. Tilde (`~`) and relative paths are expanded. "
"Must resolve inside the allowlist of browseable roots (HOME, "
- "HF cache, Studio dirs, registered scan folders, well-known "
+ "HF cache, Unsloth dirs, registered scan folders, well-known "
"model dirs)."
),
),
@@ -2251,15 +2251,15 @@ async def delete_finetuned_model(
gguf_variant: Optional[str] = Body(None),
current_subject: str = Depends(get_current_subject),
):
- """Delete a Studio-trained or exported model from disk.
+ """Delete an Unsloth-trained or exported model from disk.
- Only paths under Studio's outputs/exports roots are accepted.
+ Only paths under Unsloth's outputs/exports roots are accepted.
Exported GGUF entries can delete one quant variant at a time.
"""
if source not in {"training", "exported"}:
raise HTTPException(
status_code = 400,
- detail = "Only trained or exported Studio models can be deleted",
+ detail = "Only trained or exported Unsloth models can be deleted",
)
if not model_path or not model_path.strip():
@@ -2291,14 +2291,14 @@ async def delete_finetuned_model(
if not _is_path_under_lexically(delete_path, allowed_root):
raise HTTPException(
status_code = 400,
- detail = "Model path is outside Studio storage",
+ detail = "Model path is outside Unsloth storage",
)
if export_type == "gguf" and gguf_variant:
target_path = delete_path.resolve()
if not _is_path_under(target_path, allowed_root):
raise HTTPException(
status_code = 400,
- detail = "Model path is outside Studio storage",
+ detail = "Model path is outside Unsloth storage",
)
else:
target_path = delete_path
@@ -2311,7 +2311,7 @@ async def delete_finetuned_model(
if should_check_resolved_path and not _is_path_under(target_path, allowed_root):
raise HTTPException(
status_code = 400,
- detail = "Model path is outside Studio storage",
+ detail = "Model path is outside Unsloth storage",
)
if target_path == allowed_root:
raise HTTPException(
@@ -3456,7 +3456,7 @@ _EXPORT_SIZE_CACHE: dict[str, tuple[int, int, str]] = {}
def _is_sizable_local_path(model: str) -> bool:
- """True only for local paths under a Studio data root.
+ """True only for local paths under an Unsloth data root.
Containment is decided lexically (no filesystem access) before the path is
touched, then the path is symlink-resolved and re-checked so a symlink
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index d53e8f2bbc..53b1c4d991 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -127,9 +127,9 @@ async def start_training(
try:
logger.info(f"Starting training job with model: {request.model_name}")
- # When Studio is driven as an inference API (API-key auth), refuse to start
+ # When Unsloth is driven as an inference API (API-key auth), refuse to start
# training while a request is in flight: training frees VRAM by unloading
- # the chat model, which would kill the stream. The Studio UI (session auth)
+ # the chat model, which would kill the stream. The Unsloth UI (session auth)
# still starts training and coexists/frees VRAM as before. (A mixed UI+API
# session is not yet special-cased.)
if via_api_key is True:
@@ -139,7 +139,7 @@ async def start_training(
status_code = 409,
detail = (
"Cannot start training over the API while an inference request is in "
- "progress. Wait for it to finish, or start training from the Studio UI."
+ "progress. Wait for it to finish, or start training from the Unsloth UI."
),
)
diff --git a/studio/backend/run.py b/studio/backend/run.py
index 4f105b53b1..398943cc2c 100644
--- a/studio/backend/run.py
+++ b/studio/backend/run.py
@@ -232,7 +232,7 @@ def _working_local_url(port: int) -> "str | None":
def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None":
"""Return the IPv4 loopback URL when localhost won't reach 127.0.0.1.
- Local Studio binds to 127.0.0.1. Where localhost resolves to IPv6 only (::1),
+ Local Unsloth binds to 127.0.0.1. Where localhost resolves to IPv6 only (::1),
http://localhost: fails (or hits a different process on ::1) even though
http://127.0.0.1: works. Return the IPv4 URL for the caller to surface.
"""
@@ -243,7 +243,7 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None":
ipv4_url = f"http://127.0.0.1:{port}"
- # Only warn once Studio is confirmed answering on IPv4 loopback.
+ # Only warn once Unsloth is confirmed answering on IPv4 loopback.
if _working_local_url(port) != ipv4_url:
return None
@@ -265,7 +265,7 @@ def _localhost_ipv6_mismatch_url(bind_host: str, port: int) -> "str | None":
if host == "::1":
has_ipv6_loopback = True
- # A connection to ::1 is NOT evidence Studio is reachable there: Studio binds
+ # A connection to ::1 is NOT evidence Unsloth is reachable there: Unsloth binds
# 127.0.0.1 only, so anything on ::1 is a different process. Dual-stack
# localhost is fine (browsers fall back to 127.0.0.1), so only the IPv6-only
# case strands the user.
@@ -287,7 +287,7 @@ def _stdout_color_ok() -> bool:
def _print_localhost_ipv6_mismatch_warning(local_url: str, port: int) -> None:
- """Warn that localhost points at ::1 while Studio is bound to 127.0.0.1."""
+ """Warn that localhost points at ::1 while Unsloth is bound to 127.0.0.1."""
use_color = _stdout_color_ok()
warn_c = "\033[38;5;215;1m" if use_color else ""
reset = "\033[0m" if use_color else ""
@@ -303,7 +303,7 @@ def _print_localhost_ipv6_mismatch_warning(local_url: str, port: int) -> None:
def _verify_global_reachability(display_host: str, port: int) -> None:
"""Probe check-host.net to confirm display_host:port is reachable from the
public internet. Synchronous so output lands between the banner URLs and the
- stop hint. Bounded at ~15s; failures swallowed (verifier failing != Studio
+ stop hint. Bounded at ~15s; failures swallowed (verifier failing != Unsloth
failing). Only meaningful for a wildcard bind."""
global _public_reachable
# Reset to "unknown" each run; set True/False only when the probe decides.
@@ -563,15 +563,15 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
" Cloudflare tunnel: ON. This Cloudflare URL is PUBLIC, and the "
"raw port is also publicly reachable. --no-cloudflare disables "
f"only the Cloudflare URL; bind {loopback_host} or close firewall "
- "access to keep Studio private.",
+ "access to keep Unsloth private.",
warn,
)
else:
_emit(
" Cloudflare tunnel: ON. This is a PUBLIC internet URL: anyone "
- "who has it can reach this Studio. Relaunch with --no-cloudflare "
+ "who has it can reach this Unsloth. Relaunch with --no-cloudflare "
f"to disable the Cloudflare URL; bind {loopback_host} or close "
- "firewall access to keep Studio private.",
+ "firewall access to keep Unsloth private.",
warn,
)
return
@@ -580,12 +580,12 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
_emit(
" Cloudflare tunnel: requested but failed to start. The raw port is "
"still reachable from the public internet (see the reachability check "
- "above): anyone who can reach it can access this Studio.",
+ "above): anyone who can reach it can access this Unsloth.",
warn,
)
elif _public_reachable is False:
_emit(
- " Cloudflare tunnel: requested but failed to start. Studio is reachable "
+ " Cloudflare tunnel: requested but failed to start. Unsloth is reachable "
"on your local network only (no public link).",
warn,
)
@@ -593,7 +593,7 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
_emit(
" Cloudflare tunnel: requested but failed to start. There is no "
"Cloudflare public link. Raw port reachability was not verified; "
- f"bind {loopback_host} or close firewall access to keep Studio private.",
+ f"bind {loopback_host} or close firewall access to keep Unsloth private.",
warn,
)
elif _cloudflare_flag:
@@ -601,19 +601,19 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
_emit(
" Cloudflare tunnel: OFF for this mode. The raw port is still "
"reachable from the public internet (see the reachability check above): "
- "anyone who can reach it can access this Studio.",
+ "anyone who can reach it can access this Unsloth.",
warn,
)
elif _public_reachable is False:
_emit(
- " Cloudflare tunnel: OFF for this mode. Studio is reachable on your "
+ " Cloudflare tunnel: OFF for this mode. Unsloth is reachable on your "
"local network only (no public link)."
)
else:
_emit(
" Cloudflare tunnel: OFF for this mode. There is no Cloudflare public "
"link. Raw port reachability was not verified; "
- f"bind {loopback_host} or close firewall access to keep Studio private.",
+ f"bind {loopback_host} or close firewall access to keep Unsloth private.",
warn,
)
elif _cloudflare_flag is False or _cloudflare_flag is None:
@@ -624,12 +624,12 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
f" Cloudflare tunnel: OFF ({_reason}). The raw port is still "
"reachable from the public internet (see the reachability check above): "
"pass --cloudflare to also expose a public Cloudflare HTTPS link, or "
- f"bind {loopback_host} to keep Studio private.",
+ f"bind {loopback_host} to keep Unsloth private.",
warn,
)
elif _public_reachable is False:
_emit(
- f" Cloudflare tunnel: OFF ({_reason}). Studio is reachable on your "
+ f" Cloudflare tunnel: OFF ({_reason}). Unsloth is reachable on your "
"local network only. Pass --cloudflare to expose a public "
"Cloudflare HTTPS link."
)
@@ -638,7 +638,7 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
f" Cloudflare tunnel: OFF ({_reason}). There is no Cloudflare "
"public link. Raw port reachability was not verified; pass --cloudflare "
"to expose a public Cloudflare HTTPS link, or "
- f"bind {loopback_host} or close firewall access to keep Studio private.",
+ f"bind {loopback_host} or close firewall access to keep Unsloth private.",
warn,
)
@@ -674,7 +674,7 @@ def _is_port_free(host: str, port: int) -> bool:
For a ``0.0.0.0`` wildcard host, also check whether anything is listening on
``127.0.0.1`` (and ``::1`` when IPv6 exists): an SSH tunnel may hold loopback
- while the wildcard bind succeeds, making Studio unreachable via ``localhost``.
+ while the wildcard bind succeeds, making Unsloth unreachable via ``localhost``.
"""
import socket
@@ -1087,7 +1087,7 @@ def _terminal_password_gate(
) -> Tuple[bool, bool]:
"""Force a terminal password change before the public tunnel goes up.
- When the tunnel is about to publish Studio and the seeded admin password was
+ When the tunnel is about to publish Unsloth and the seeded admin password was
never changed, ask for a new one (masked, confirmed) before any public URL
exists. The CLI normally does this before re-exec'ing the backend; this is
the backstop for direct `python run.py` launches and older-CLI installs.
@@ -1147,7 +1147,7 @@ def _terminal_password_gate(
)
if not deadline_arms:
print(
- "Refusing to publish Studio on a public Cloudflare URL: the "
+ "Refusing to publish Unsloth on a public Cloudflare URL: the "
"default admin password was never changed, no terminal is "
"attached to change it here, and the bootstrap shutdown "
"deadline does not apply to this launch (api-only, or "
@@ -1163,11 +1163,11 @@ def _terminal_password_gate(
# terminal-attached run / reset-password instead of reading it from disk.
print(
" WARNING: the default admin password is still active while "
- "Studio is about to be published on a public Cloudflare URL, and "
+ "Unsloth is about to be published on a public Cloudflare URL, and "
"no terminal is attached to change it here. The public page will "
"NOT auto-fill the bootstrap credential. Set a new password by "
"running `unsloth studio` locally with a terminal attached, or "
- "`unsloth studio reset-password`. Studio shuts down after the "
+ "`unsloth studio reset-password`. Unsloth shuts down after the "
"bootstrap deadline (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT, default 1h) "
"unless the password is changed.",
file = sys.stderr,
@@ -1222,7 +1222,7 @@ def _apply_supplied_password(password_value: "Optional[str]") -> None:
_auth_storage.ensure_default_admin()
if not _auth_storage.requires_password_change(_admin):
print(
- "Error: a Studio admin password is already set; --password only sets "
+ "Error: an Unsloth admin password is already set; --password only sets "
"the initial password. Run `unsloth studio reset-password` first.",
file = sys.stderr,
flush = True,
@@ -1337,7 +1337,7 @@ def run_server(
pass
# Persist a session log + native-crash stacks BEFORE importing main, so
- # even import-time failures leave evidence on disk. Field report: Studio
+ # even import-time failures leave evidence on disk. Field report: Unsloth
# "terminates without a warning" -- a native crash in the GPU runtime
# kills the process with no Python traceback, and a desktop-shortcut
# console closes before anything can be read. Console-only logging made
@@ -1406,7 +1406,7 @@ def run_server(
ensure_studio_directories()
logger.info(
- "Ensured Studio directories in %.1fms",
+ "Ensured Unsloth directories in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
@@ -1455,7 +1455,7 @@ def run_server(
installer_bin = home / "unsloth_studio" / "bin" / "unsloth"
tried_lines = "\n".join(f" - {p}" for p in attempted) or " (none)"
raise SystemExit(
- "[ERROR] Studio frontend build not found.\n"
+ "[ERROR] Unsloth frontend build not found.\n"
f"Tried:\n{tried_lines}\n"
"\n"
"Likely cause: another 'unsloth' on PATH is shadowing the "
@@ -1557,7 +1557,7 @@ def run_server(
)
if not _pw_proceed:
print(
- "Not starting Studio; set a new admin password first, or launch "
+ "Not starting Unsloth; set a new admin password first, or launch "
"without --secure/--cloudflare.",
file = sys.stderr,
flush = True,
@@ -1695,7 +1695,7 @@ def run_server(
logger = logger,
)
logger.info(
- "Studio will shut down in %ds unless the default admin password is changed.",
+ "Unsloth will shut down in %ds unless the default admin password is changed.",
_bootstrap_timeout,
)
except Exception as e: # best-effort: never block startup on the timeout
@@ -1753,11 +1753,11 @@ def _build_arg_parser():
"--cloudflare",
action = argparse.BooleanOptionalAction,
default = None,
- help = "Expose Studio on a PUBLIC internet URL via a free Cloudflare HTTPS "
+ help = "Expose Unsloth on a PUBLIC internet URL via a free Cloudflare HTTPS "
"tunnel, for non-api-only wildcard binds (0.0.0.0 or ::). Off by default; "
"pass --cloudflare to enable it (--secure implies it), --no-cloudflare to "
"force it off. It does not change a raw wildcard bind. If the admin "
- "password was never changed, Studio asks for a new one in the terminal "
+ "password was never changed, Unsloth asks for a new one in the terminal "
"before publishing the URL.",
)
parser.add_argument(
@@ -1767,7 +1767,7 @@ def _build_arg_parser():
help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed "
"if the tunnel can't start. Without it, --no-secure also serves the raw "
"0.0.0.0 port, which is reachable from anywhere on the network. If the "
- "admin password was never changed, Studio asks for a new one in the "
+ "admin password was never changed, Unsloth asks for a new one in the "
"terminal before publishing the URL.",
)
# Back-compat: accept --not-secure as a hidden alias for --no-secure.
diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py
index ea951a4325..9ec7a4f91c 100644
--- a/studio/backend/startup_banner.py
+++ b/studio/backend/startup_banner.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
-"""Terminal banner for Studio startup.
+"""Terminal banner for Unsloth startup.
Stdlib only -- safe to import without the rest of the backend.
"""
@@ -172,7 +172,7 @@ def print_studio_access_banner(
secondary,
),
style(
- " Only on trusted networks -- anyone who reaches this machine can use Studio.",
+ " Only on trusted networks -- anyone who reaches this machine can use Unsloth.",
secondary,
),
]
diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py
index b0b9ee309c..c2216104a3 100644
--- a/studio/backend/tests/conftest.py
+++ b/studio/backend/tests/conftest.py
@@ -101,13 +101,13 @@ def studio_server(request):
@pytest.fixture
def base_url(studio_server):
- """Base URL for the e2e Studio server (from ``studio_server``)."""
+ """Base URL for the e2e Unsloth server (from ``studio_server``)."""
return studio_server[0]
@pytest.fixture
def api_key(studio_server):
- """API key for the e2e Studio server (from ``studio_server``)."""
+ """API key for the e2e Unsloth server (from ``studio_server``)."""
return studio_server[1]
diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py
index 4df9e85b30..9fd8260bf2 100644
--- a/studio/backend/tests/test_amd_apu_unified_memory.py
+++ b/studio/backend/tests/test_amd_apu_unified_memory.py
@@ -91,7 +91,7 @@ class TestApuRamShortfall:
"""On a unified-memory APU the weights load into system RAM, so a model
larger than available RAM (the field case: a 64.6 GB GGUF on a WSL VM capped
well below the ROCm-reported APU budget) must be refused before spawning,
- not left to OOM-kill the Studio process."""
+ not left to OOM-kill the Unsloth process."""
def test_field_case_wsl_cap_refuses(self):
# 64.6 GB weights, ~46 GB available (WSL VM): refuse with guidance.
diff --git a/studio/backend/tests/test_anthropic_compaction.py b/studio/backend/tests/test_anthropic_compaction.py
index 1528eebe8b..acc0acc2e0 100644
--- a/studio/backend/tests/test_anthropic_compaction.py
+++ b/studio/backend/tests/test_anthropic_compaction.py
@@ -4,7 +4,7 @@
"""Unit tests for Anthropic server-side context compaction wiring.
Compaction is a beta (header ``compact-2026-01-12``) gated to Opus 4.6/4.7,
-Sonnet 4.6, and Mythos preview. When enabled, Studio attaches
+Sonnet 4.6, and Mythos preview. When enabled, Unsloth attaches
``context_management.edits[{type:"compact_20260112", trigger:{type:"input_tokens",
value:N}}]``; the 50k-token minimum is clamped up so the request doesn't 400.
diff --git a/studio/backend/tests/test_anthropic_fast_mode_edge.py b/studio/backend/tests/test_anthropic_fast_mode_edge.py
index dd69d77590..03f5d1c0eb 100644
--- a/studio/backend/tests/test_anthropic_fast_mode_edge.py
+++ b/studio/backend/tests/test_anthropic_fast_mode_edge.py
@@ -330,7 +330,7 @@ def test_refusal_chunk_is_proper_openai_delta_shape(monkeypatch):
def test_refusal_tool_event_chunk_shape(monkeypatch):
- """Drop signal rides a Studio `_toolEvent` envelope (delta={},
+ """Drop signal rides an Unsloth `_toolEvent` envelope (delta={},
finish_reason=null); the frontend latches on
`_toolEvent.type == "anthropic_refusal"`."""
_, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7")
@@ -409,7 +409,7 @@ def _fast_speed_sse(model: str = "claude-opus-4-7", speed: str = "fast") -> byte
def test_usage_speed_propagates_to_final_usage_chunk_fast(monkeypatch):
- """``usage.speed == "fast"`` from upstream must reach the Studio usage chunk."""
+ """``usage.speed == "fast"`` from upstream must reach the Unsloth usage chunk."""
_, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "fast"))
usage_lines = [l for l in lines if l.startswith("data: ") and '"usage"' in l]
assert usage_lines, lines
@@ -428,7 +428,7 @@ def test_usage_speed_propagates_to_final_usage_chunk_standard(monkeypatch):
def test_usage_speed_absent_when_anthropic_does_not_report(monkeypatch):
- """Studio must not invent ``usage.speed`` when upstream omits it."""
+ """Unsloth must not invent ``usage.speed`` when upstream omits it."""
_, lines = _capture(monkeypatch)
parsed = [
json.loads(l[len("data: ") :]) for l in lines if l.startswith("data: ") and '"usage"' in l
diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py
index 3b0ea37372..9ccc3f44dd 100644
--- a/studio/backend/tests/test_anthropic_messages.py
+++ b/studio/backend/tests/test_anthropic_messages.py
@@ -1418,7 +1418,7 @@ class TestNormalizeAnthropicOpenAIImages:
# =====================================================================
-# Studio-tool alias detection (/v1/messages tool routing)
+# Unsloth-tool alias detection (/v1/messages tool routing)
# =====================================================================
@@ -1436,7 +1436,7 @@ class TestAnthropicRequestedStudioTools:
def test_client_tool_named_python_is_not_misclassified(self):
# input_schema is the client-tool discriminator; its presence must
- # prevent the name from being treated as a Studio alias.
+ # prevent the name from being treated as an Unsloth alias.
tools = [
{
"name": "python",
@@ -1747,9 +1747,9 @@ class TestAnthropicMessagesToolRouting:
assert "name" in exc.value.detail
def test_alias_named_client_tool_without_schema_rejected_with_400(self, monkeypatch):
- # Regression: a typo'd client tool whose name collides with a Studio
+ # Regression: a typo'd client tool whose name collides with an Unsloth
# alias (e.g. a custom "python" tool missing input_schema) must
- # surface a 400, not silently switch into Studio's built-in python
+ # surface a 400, not silently switch into Unsloth's built-in python
# execution.
_mock_backend(monkeypatch)
payload = _basic_payload(tools = [{"name": "python"}])
@@ -1770,7 +1770,7 @@ class TestAnthropicMessagesToolRouting:
def test_disable_tools_policy_overrides_server_tool_alias(self, monkeypatch):
# CLI `unsloth run --disable-tools` sets policy=False. A request with
- # a Studio server-tool alias must NOT enter the agentic loop then.
+ # an Unsloth server-tool alias must NOT enter the agentic loop then.
backend = _mock_backend(monkeypatch)
set_tool_policy(False)
payload = _basic_payload(
diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py
index 8408f8203d..3e95acc98d 100644
--- a/studio/backend/tests/test_compute_buffer.py
+++ b/studio/backend/tests/test_compute_buffer.py
@@ -152,7 +152,7 @@ class TestFallback:
class TestParallel1Default:
- """At Studio's default --parallel 1 the buffer is negligible in pipeline."""
+ """At Unsloth's default --parallel 1 the buffer is negligible in pipeline."""
def test_default_n_parallel(self):
est = _backend()._estimate_compute_buffer_bytes() / MIB
diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py
index 2930c9f081..9d8795b6c0 100644
--- a/studio/backend/tests/test_cpu_threads.py
+++ b/studio/backend/tests/test_cpu_threads.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
-"""Tests for Studio's early CPU thread-pool configuration."""
+"""Tests for Unsloth's early CPU thread-pool configuration."""
import ast
import os
@@ -30,7 +30,7 @@ def test_cpu_thread_cap_seeds_native_pool_limits():
}
-# Explicit per-library values win over the Studio knob via setdefault.
+# Explicit per-library values win over the Unsloth knob via setdefault.
def test_cpu_thread_cap_preserves_runtime_specific_override():
env = {"UNSLOTH_CPU_THREADS": "4", "OMP_NUM_THREADS": "2"}
diff --git a/studio/backend/tests/test_frontend_resolution.py b/studio/backend/tests/test_frontend_resolution.py
index c3e0524a30..7ac2717aae 100644
--- a/studio/backend/tests/test_frontend_resolution.py
+++ b/studio/backend/tests/test_frontend_resolution.py
@@ -218,7 +218,7 @@ def test_systemexit_message_contains_actionable_fixes(tmp_path, monkeypatch):
installer_bin = home / "unsloth_studio" / "bin" / "unsloth"
tried_lines = "\n".join(f" - {p}" for p in attempted)
message = (
- "[ERROR] Studio frontend build not found.\n"
+ "[ERROR] Unsloth frontend build not found.\n"
f"Tried:\n{tried_lines}\n"
"\n"
"Likely cause: another 'unsloth' on PATH is shadowing the "
diff --git a/studio/backend/tests/test_gemini_provider.py b/studio/backend/tests/test_gemini_provider.py
index 85ceb04d27..c6ffa798d0 100644
--- a/studio/backend/tests/test_gemini_provider.py
+++ b/studio/backend/tests/test_gemini_provider.py
@@ -768,7 +768,7 @@ def test_cached_content_pass_through(monkeypatch):
def test_boolean_caching_does_not_set_cached_content(monkeypatch):
- """Studio's existing True/False signals shouldn't fabricate a cache id."""
+ """Unsloth's existing True/False signals shouldn't fabricate a cache id."""
captured = _capture_body(monkeypatch, enable_prompt_caching = True)
assert "cachedContent" not in captured["body"]
@@ -2613,7 +2613,7 @@ def test_gemini_native_skips_orphan_function_response_for_native_part_replay(mon
def test_gemini_native_part_falls_back_to_args_google(monkeypatch):
"""Round 27: a direct OpenAI-compat API caller (or imported third-party
- thread) cannot use Studio's non-standard `tool_calls[].extra_content`
+ thread) cannot use Unsloth's non-standard `tool_calls[].extra_content`
field, so the native_part payload round-trips through `function.arguments`
as `{"google": {"native_part": {...}}}`. The synthetic-builtin detector
recognizes that location, but the replay branch was only reading from
diff --git a/studio/backend/tests/test_gemma4_chat_template_override.py b/studio/backend/tests/test_gemma4_chat_template_override.py
index f726741aa5..9fb24a4cf6 100644
--- a/studio/backend/tests/test_gemma4_chat_template_override.py
+++ b/studio/backend/tests/test_gemma4_chat_template_override.py
@@ -3,7 +3,7 @@
"""Auto-override of the chat template for ``unsloth/gemma-4-*-GGUF``.
-Studio ships a bundled ``gemma-4.jinja`` (PR #118 based, ``preserve_thinking``
+Unsloth ships a bundled ``gemma-4.jinja`` (PR #118 based, ``preserve_thinking``
defaulted off) and applies it to gemma-4 GGUF loads via the existing
``chat_template_override`` -> ``--chat-template-file`` path, so users do not need
to re-download quants. Pins the family matcher, the resolver precedence, the
diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py
index 2fff744b64..48aff29659 100644
--- a/studio/backend/tests/test_hf_xet_fallback.py
+++ b/studio/backend/tests/test_hf_xet_fallback.py
@@ -1,10 +1,10 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback.
+"""Tests for the Unsloth shim over the shared unsloth_zoo Xet -> HTTP fallback.
The transport-policy matrix is tested once in unsloth_zoo; here we assert only the
-Studio seam: re-exporting the shared API and injecting the marker-aware
+Unsloth seam: re-exporting the shared API and injecting the marker-aware
prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess.
"""
@@ -69,7 +69,7 @@ def test_child_should_disable_xet_truth_table():
def test_shim_injects_studio_prepare_on_http_retry(monkeypatch):
- """A Xet stall retries over HTTP and the shim runs Studio's marker-aware
+ """A Xet stall retries over HTTP and the shim runs Unsloth's marker-aware
``prepare_cache_for_transport(..., 'http')`` before the retry."""
_requires_shared()
for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"):
@@ -107,11 +107,11 @@ def test_shim_injects_studio_prepare_on_http_retry(monkeypatch):
out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
assert out == "/cache/model.gguf"
assert seen_disable_xet == [False, True] # Xet first, then HTTP
- assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep"
+ assert prepared == [("model", DL_REPO, "http")], "shim must run Unsloth's marker-aware prep"
def test_shim_snapshot_injects_studio_prepare(monkeypatch):
- """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper."""
+ """The snapshot wrapper forwards Unsloth's marker-aware prep, like the file wrapper."""
captured = {}
def fake_snapshot(repo_id, **kwargs):
@@ -127,7 +127,7 @@ def test_shim_snapshot_injects_studio_prepare(monkeypatch):
def test_degrades_gracefully_without_shared_helper(monkeypatch):
- """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio
+ """On an older unsloth_zoo lacking the shared helper, the shim still imports (Unsloth
boots) and exposes stub API doing plain HF downloads with the watchdog disabled."""
import importlib
@@ -206,7 +206,7 @@ def test_degrades_gracefully_without_shared_helper(monkeypatch):
def test_degrades_when_unsloth_zoo_entirely_absent():
"""When unsloth_zoo is absent entirely, the import raises
ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still
- degrades and does not re-raise, breaking every Studio import that pulls it in."""
+ degrades and does not re-raise, breaking every Unsloth import that pulls it in."""
import importlib
class _BlockZoo:
@@ -248,7 +248,7 @@ def test_degrades_when_unsloth_zoo_entirely_absent():
def test_degrades_when_shared_helper_import_raises_importerror():
"""unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only
- Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too."""
+ Unsloth), raising ImportError not ModuleNotFoundError. The shim must degrade for that too."""
import importlib
class _BlockWithImportError:
@@ -329,7 +329,7 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch):
# with it set); accessing DownloadStallError drives it via __getattr__.
stall_error = degraded.DownloadStallError
assert seen_env == [None, "1"], seen_env
- # Both attempts raised -> Studio still boots in degraded mode.
+ # Both attempts raised -> Unsloth still boots in degraded mode.
assert issubclass(stall_error, RuntimeError)
# The env override must not leak past the load.
assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None
diff --git a/studio/backend/tests/test_identity.py b/studio/backend/tests/test_identity.py
index 1e84ddef35..712348f7ca 100644
--- a/studio/backend/tests/test_identity.py
+++ b/studio/backend/tests/test_identity.py
@@ -3,7 +3,7 @@
"""Tests for the server identity handshake (`GET /api/auth/identity`).
-The endpoint lets a client confirm an endpoint is really this Studio install
+The endpoint lets a client confirm an endpoint is really this Unsloth install
before sending it a credential: the client sends a random nonce and checks the
returned HMAC against one computed from the install identity secret. A process
that cannot read this same-user secret cannot forge a matching proof.
diff --git a/studio/backend/tests/test_index_bootstrap_origin_extra.py b/studio/backend/tests/test_index_bootstrap_origin_extra.py
index feda88c14c..e1c52a653e 100644
--- a/studio/backend/tests/test_index_bootstrap_origin_extra.py
+++ b/studio/backend/tests/test_index_bootstrap_origin_extra.py
@@ -26,7 +26,7 @@ def _build_request(
def test_is_same_origin_request_ipv6_loopback_same_origin():
- """Studio supports ``-H ::1`` binds; netloc is ``[::1]:8902``. Bare
+ """Unsloth supports ``-H ::1`` binds; netloc is ``[::1]:8902``. Bare
``partition(":")`` mis-parses the bracketed form and would refuse the
bootstrap on legitimate same-origin navigation.
"""
diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py
index d3a10df8ca..2a4f6d19d2 100644
--- a/studio/backend/tests/test_llama_cpp_context_fit.py
+++ b/studio/backend/tests/test_llama_cpp_context_fit.py
@@ -567,7 +567,7 @@ class TestClassifyGpuOffload:
assert inst._classify_gpu_offload(False, []) is None
def test_user_did_not_intend_gpu_returns_none(self):
- # Studio called start_llama_server without expecting GPU; don't warn.
+ # Unsloth called start_llama_server without expecting GPU; don't warn.
inst = self._backend(
[
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py
index 04d4aac9e1..049058e511 100644
--- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py
+++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py
@@ -222,7 +222,7 @@ class TestFlashAttnOff:
assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"]
def test_flips_every_occurrence_last_wins(self):
- # extra_args can re-enable FA after Studio's flag; llama.cpp is last-wins,
+ # extra_args can re-enable FA after Unsloth's flag; llama.cpp is last-wins,
# so one leftover 'on' would re-crash the retry. Every enable must flip.
cmd = ["llama-server", "--flash-attn", "on", "--mmproj", "/p", "--flash-attn", "on"]
out = _flash_off(cmd)
@@ -234,7 +234,7 @@ class TestFlashAttnOff:
assert _flash_off(["llama-server", "--flash-attn=off"]) is None
def test_none_when_user_off_wins_last(self):
- # User appended 'off' after Studio's 'on'; effective (last-wins) is off,
+ # User appended 'off' after Unsloth's 'on'; effective (last-wins) is off,
# so there is nothing to retry.
assert _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn", "off"]) is None
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index 3f9d2a8f50..8fe04c0e39 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -1014,7 +1014,7 @@ def test_already_in_target_state_2b_falls_back_to_ngram_below_threshold(monkeypa
)
-# usage backfill from timings (Studio UI t/s widget fix).
+# usage backfill from timings (Unsloth UI t/s widget fix).
def test_backfill_usage_from_timings_fills_when_completion_tokens_zero():
@@ -1606,7 +1606,7 @@ def test_reload_forced_mtp_bounces_auto_mla():
)
-# ── Full named-repo resolver matrix (the shipping Studio families) ─────
+# ── Full named-repo resolver matrix (the shipping Unsloth families) ─────
#
# Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and
# gemma-4 (regular + QAT) GGUF repo, including the giant MoEs that stay
diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py
index 10b1dc7ff6..f320d29a02 100644
--- a/studio/backend/tests/test_llama_cpp_no_context_shift.py
+++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py
@@ -5,7 +5,7 @@
With llama-server's default context-shift behavior, the UI cannot tell the user
the KV cache was rotated -- earlier turns silently vanish from the conversation.
-The Studio backend always passes ``--no-context-shift`` so the server returns a
+The Unsloth backend always passes ``--no-context-shift`` so the server returns a
clean error instead, and the chat adapter can point the user at the
``Context Length`` input in the settings panel.
diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py
index 316956325f..488645ee5a 100644
--- a/studio/backend/tests/test_llama_cpp_props_readback.py
+++ b/studio/backend/tests/test_llama_cpp_props_readback.py
@@ -4,7 +4,7 @@
"""Tests for the post-launch /props context readback.
llama-server's memory-fit step or --parallel slot split can allocate less
-context than the requested -c while Studio keeps advertising the requested
+context than the requested -c while Unsloth keeps advertising the requested
value; clients sized to it then die on exceed_context_size_error 400s.
``_reconcile_effective_ctx_with_server`` must adopt the server's real
``default_generation_settings.n_ctx`` whenever it is smaller.
@@ -223,7 +223,7 @@ _CAPS_NONE = {"supports_kv_unified": False, "supports_fit_ctx": False}
def test_kv_unified_added_for_multi_slot():
"""Explicit --parallel N disables llama-server's auto-slots kv-unified
- default, splitting -c into per-slot windows of -c/N; Studio must restore
+ default, splitting -c into per-slot windows of -c/N; Unsloth must restore
the shared pool so one request can use the full advertised context."""
flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL)
assert "--kv-unified" in flags
diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py
index bd2c008589..e99e227d40 100644
--- a/studio/backend/tests/test_llama_cpp_tool_loop.py
+++ b/studio/backend/tests/test_llama_cpp_tool_loop.py
@@ -122,7 +122,7 @@ def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list
def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
"""llama-server may emit content first and then native delta.tool_calls.
- Studio must not drop that tool call after it has streamed the preface.
+ Unsloth must not drop that tool call after it has streamed the preface.
"""
tool_call_id = "call_render_late"
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py
index 82c5b4931a..423c3dd009 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_health.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py
@@ -224,7 +224,7 @@ class TestRetryLogFilenameUnique:
class TestFitOffRetryEligible:
"""Gate for the one-shot --fit off startup-crash retry.
- Retry only when Studio's own VRAM math placed the model and nothing
+ Retry only when Unsloth's own VRAM math placed the model and nothing
on the command line chose the fit mode explicitly."""
def test_eligible_for_plain_ngl_launch(self):
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
index d0213f6079..b28df7ec3f 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
@@ -346,7 +346,7 @@ def test_helper_is_static_method_callable_off_class():
def test_kill_orphaned_servers_returns_count():
"""The reaper reports how many owned orphans it killed, so __init__ can
- arm the settle wait. Only Studio-owned llama-server procs count."""
+ arm the settle wait. Only Unsloth-owned llama-server procs count."""
import os
mypid = os.getpid()
@@ -376,7 +376,7 @@ def test_kill_orphaned_servers_returns_count():
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
):
n = LlamaCppBackend._kill_orphaned_servers()
- assert n == 1, "only the Studio-owned orphan should be counted"
+ assert n == 1, "only the Unsloth-owned orphan should be counted"
assert killed == [mypid + 1]
# No owned orphans -> zero, so __init__ leaves the cold-start sentinel.
@@ -392,8 +392,8 @@ def test_kill_orphaned_servers_returns_count():
def test_kill_orphaned_servers_spares_live_parent():
- """A Studio-owned llama-server whose parent is still running is not an
- orphan (a live Studio or the user's shell owns it) and must never be
+ """An Unsloth-owned llama-server whose parent is still running is not an
+ orphan (a live Unsloth or the user's shell owns it) and must never be
killed; only the true orphan (parent gone) is reaped."""
import os
@@ -548,7 +548,7 @@ def test_record_then_reap_round_trip_identity_matches(tmp_path):
def test_reap_recorded_pid_spares_live_server(tmp_path):
- """A recorded server whose parent is still alive (the running Studio) is NEVER
+ """A recorded server whose parent is still alive (the running Unsloth) is NEVER
reaped, and its pidfile is kept. This is the finding-3 guard: a helper backend
constructed in-process must not kill the active chat server. Uses the REAL
_pid_parent_is_alive (the child's parent is this live test process)."""
diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
index 957de4bad6..489d9eb8d1 100644
--- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
+++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
@@ -3,7 +3,7 @@
"""Tests for the Windows pip-nvidia DLL dir resolver.
-Studio installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13,
+Unsloth installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13,
nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find those
DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH block.
See unslothai/unsloth#5106.
diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py
index deeb228026..ba52afad1c 100644
--- a/studio/backend/tests/test_llama_server_args.py
+++ b/studio/backend/tests/test_llama_server_args.py
@@ -75,7 +75,7 @@ validate_extra_args = _lsa.validate_extra_args
# Reasoning controls
["--reasoning-format", "deepseek"],
["-rea", "auto"],
- # Soft-managed: user flags last-wins over Studio's auto-set version.
+ # Soft-managed: user flags last-wins over Unsloth's auto-set version.
# --parallel / -np / --n-parallel are hard-denied (KV-cache + slot
# count would desync); use `unsloth studio run --parallel N` instead.
["-c", "131072"],
@@ -150,7 +150,7 @@ def test_non_flag_token_passes_through():
"--mmproj",
"-mmu",
"--mmproj-url",
- # Networking (Studio binds + proxies)
+ # Networking (Unsloth binds + proxies)
"--host",
"--port",
"--path",
@@ -176,12 +176,12 @@ def test_non_flag_token_passes_through():
"--models-autoload",
"--no-models-autoload",
# Server-mode flips: --embedding / --rerank restrict llama-server to
- # those endpoints and break Studio's chat hop.
+ # those endpoints and break Unsloth's chat hop.
"--embedding",
"--embeddings",
"--rerank",
"--reranking",
- # llama-server's own --tools clashes with Studio's tool policy.
+ # llama-server's own --tools clashes with Unsloth's tool policy.
"--tools",
],
)
@@ -194,7 +194,7 @@ def test_denylist_rejects_all_aliases(denied):
"args,offending",
[
# Pass-through --parallel would last-wins-override the real slot
- # count while Studio's KV-cache fit + llama_parallel_slots stay at
+ # count while Unsloth's KV-cache fit + llama_parallel_slots stay at
# the typer value -- plan vs. process disagree.
(["--parallel", "8"], "--parallel"),
(["--parallel=8"], "--parallel"),
@@ -656,7 +656,7 @@ def test_extra_args_disable_mmproj_last_wins():
def test_strip_shadowing_flags_drops_model_draft_with_spec():
- # --model-draft (and aliases) are Studio-managed since the separate
+ # --model-draft (and aliases) are Unsloth-managed since the separate
# MTP drafter support: an inherited copy must not last-wins-override
# the auto-detected drafter.
out = strip_shadowing_flags(
@@ -681,7 +681,7 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec():
)
def test_strip_shadowing_flags_drops_hf_drafter_selectors_with_spec(selector):
# HF drafter selectors must reset on inherit like local --model-draft, or a
- # stale inherited HF drafter last-wins over Studio's re-derived spec choice.
+ # stale inherited HF drafter last-wins over Unsloth's re-derived spec choice.
out = strip_shadowing_flags(
selector + ["--top-k", "20"],
strip_context = False,
@@ -769,7 +769,7 @@ def test_strip_split_mode_only_preserves_none_and_empty():
def test_strip_shadowing_flags_drops_tensor_split_with_split_mode():
# --tensor-split is coupled to the split mode: stripped together so a stale
- # ratio can't override Studio's computed tensor split. Other flags survive.
+ # ratio can't override Unsloth's computed tensor split. Other flags survive.
out = strip_shadowing_flags(
["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"],
strip_context = False,
diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py
index c78c029d91..6b44f61972 100644
--- a/studio/backend/tests/test_local_llama_cpp_link.py
+++ b/studio/backend/tests/test_local_llama_cpp_link.py
@@ -4,7 +4,7 @@
"""Behavioral tests for the --with-llama-cpp-dir 'unmanaged local link' contract.
When the canonical llama.cpp dir is a symlink (POSIX) / junction (Windows) to a
-user's own checkout, Studio must treat it as externally managed:
+user's own checkout, Unsloth must treat it as externally managed:
- the in-app updater must not offer or apply a prebuilt over the link
- orphan cleanup must not kill a llama-server the user launched from that tree
@@ -67,7 +67,7 @@ def test_active_install_is_local_link(tmp_path: Path) -> None:
binary = str(link / _server_subpath())
assert u._active_install_is_local_link(binary) is True
- # A plain (non-link) llama.cpp dir is Studio-managed, not a local link.
+ # A plain (non-link) llama.cpp dir is Unsloth-managed, not a local link.
plain = tmp_path / "plain" / "llama.cpp"
plain.mkdir(parents = True)
assert u._active_install_is_local_link(str(plain / _server_subpath())) is False
diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py
index 6432ffb8e1..c5c37f098f 100644
--- a/studio/backend/tests/test_mcp_servers.py
+++ b/studio/backend/tests/test_mcp_servers.py
@@ -577,7 +577,7 @@ def test_clear_oauth_tokens_swallows_constructor_errors(tmp_path, monkeypatch):
def test_tool_xml_parser_handles_hyphenated_function_names():
"""Hyphenated tool names like `mcp__srv__list-issues` must parse, else the
- model can call the tool but Studio can't dispatch."""
+ model can call the tool but Unsloth can't dispatch."""
from core.inference.tool_call_parser import parse_tool_calls_from_text
calls = parse_tool_calls_from_text(
diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py
index b0bfd45135..745c2cc447 100644
--- a/studio/backend/tests/test_mcp_stdio_improvements.py
+++ b/studio/backend/tests/test_mcp_stdio_improvements.py
@@ -188,7 +188,7 @@ def test_validate_url_allows_url_in_argument(monkeypatch):
# ── P6: Data Recipe stdio path obeys the same host gate ─────────────
-# build_mcp_providers needs the Studio-only data_designer plugin; skip if absent.
+# build_mcp_providers needs the Unsloth-only data_designer plugin; skip if absent.
_STDIO_RECIPE = {
"mcp_providers": [
diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py
index 7b8aefb722..fafaea0043 100644
--- a/studio/backend/tests/test_mlx_inference_backend.py
+++ b/studio/backend/tests/test_mlx_inference_backend.py
@@ -236,7 +236,7 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri
_install_fake_fast_mlx(monkeypatch, calls)
def _native_vlm_load(*_args, **_kwargs):
- raise AssertionError("Studio MLX VLM inference must use FastMLXModel")
+ raise AssertionError("Unsloth MLX VLM inference must use FastMLXModel")
mlx_vlm = types.ModuleType("mlx_vlm")
mlx_vlm.load = _native_vlm_load
diff --git a/studio/backend/tests/test_mlx_repair.py b/studio/backend/tests/test_mlx_repair.py
index 365cc46410..47a695ccbd 100644
--- a/studio/backend/tests/test_mlx_repair.py
+++ b/studio/backend/tests/test_mlx_repair.py
@@ -103,7 +103,7 @@ def test_repair_install_pins_transformers_and_cleans_up(monkeypatch):
assert mr.attempt_mlx_repair() is True
cmd = captured["cmd"]
# transformers is pinned via a constraint file so the mlx install cannot
- # upgrade it underneath Studio, and the temp constraint file is cleaned up.
+ # upgrade it underneath Unsloth, and the temp constraint file is cleaned up.
assert "--constraint" in cmd
assert "--upgrade" in cmd
reinstall_pairs = set(zip(cmd, cmd[1:]))
@@ -123,7 +123,7 @@ def test_install_requires_prebuilt_wheels(monkeypatch):
# A source distribution's PEP 517 build backend runs arbitrary code at install
# time, before the post-install stack check. The unattended self-heal must
# require pre-built wheels so a malicious resolver-selected sdist cannot execute
- # during ordinary Studio startup. mlx/mlx-metal ship wheels only and
+ # during ordinary Unsloth startup. mlx/mlx-metal ship wheels only and
# mlx-lm/mlx-vlm publish py3-none-any wheels, so a healthy self-heal still works.
pytest.importorskip("transformers")
captured = {}
@@ -143,7 +143,7 @@ def test_install_requires_prebuilt_wheels(monkeypatch):
def test_install_env_drops_secrets_and_source_redirects(monkeypatch):
- # The unattended self-heal must not hand resolver/build code the full Studio
+ # The unattended self-heal must not hand resolver/build code the full Unsloth
# environment: secrets and package-source redirects are dropped, while the
# variables uv genuinely needs are forwarded.
monkeypatch.setenv("HF_TOKEN", "secret-hf")
diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py
index 0efbbf596d..694d60cfc6 100644
--- a/studio/backend/tests/test_mtp_vram_budget.py
+++ b/studio/backend/tests/test_mtp_vram_budget.py
@@ -502,7 +502,7 @@ class TestExtraArgsMtpDetection:
assert _extra_args_mtp_draft_path([], env = dict(os.environ)) == "/large.gguf"
def test_load_model_gates_env_spec_type_on_off_mode(self):
- # LLAMA_ARG_SPEC_TYPE only reaches the child when Studio emits no spec
+ # LLAMA_ARG_SPEC_TYPE only reaches the child when Unsloth emits no spec
# flag (UI mode "off", no user --spec-type); otherwise the emitted
# --spec-type/--spec-default overrides the env, so the reserve must not
# consult it or a stale MTP env over-reserves (Finding F3). Whitespace-
@@ -530,8 +530,8 @@ class TestExtraArgsMtpDetection:
def test_load_model_drafter_budget_precedence(self):
# The budget sizes the drafter the launch actually loads: CLI extras win,
- # then Studio's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL),
- # then the env drafter -- not the env before Studio's (reviewer.py R3).
+ # then Unsloth's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL),
+ # then the env drafter -- not the env before Unsloth's (reviewer.py R3).
compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact
assert "_env_draft_for_budget=_extra_args_mtp_draft_path([],env=os.environ)" in compact
@@ -732,7 +732,7 @@ class TestExtraArgsMtpDetection:
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None
def test_env_main_cache_type_for_budget(self):
- # The child inherits LLAMA_ARG_CACHE_TYPE_K/_V, but Studio emits no
+ # The child inherits LLAMA_ARG_CACHE_TYPE_K/_V, but Unsloth emits no
# --cache-type when neither param nor extras set it -> a heavier env
# main KV (f32) must be adopted so the reserve matches the child.
assert _env_main_cache_type_for_budget(env = {}) is None
@@ -765,7 +765,7 @@ class TestExtraArgsMtpDetection:
assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact
def test_env_split_mode_is_tensor(self):
- # The child inherits LLAMA_ARG_SPLIT_MODE, but Studio emits --split-mode
+ # The child inherits LLAMA_ARG_SPLIT_MODE, but Unsloth emits --split-mode
# only on its tensor branch -> a tensor env must flip the budget so the
# heavier per-device compute buffer is reserved (not layer overhead).
assert _env_split_mode_is_tensor(env = {}) is False
@@ -918,7 +918,7 @@ class TestExtraArgsMtpDetection:
# Cluster A: when the final decision is layer split, an inherited
# non-layer LLAMA_ARG_SPLIT_MODE (and paired LLAMA_ARG_TENSOR_SPLIT) must
# be popped from the child env so the child cannot run tensor/row/none
- # against Studio's layer budget. Whitespace-stripped for formatter.
+ # against Unsloth's layer budget. Whitespace-stripped for formatter.
compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
assert 'env.get("LLAMA_ARG_SPLIT_MODE")' in compact
assert '_inherited_sm!="layer"' in compact
@@ -936,10 +936,10 @@ class TestExtraArgsMtpDetection:
assert "env.pop(_ct_var,None)" in compact
def test_load_model_clears_tensor_split_env_in_tensor_mode(self):
- # review run3 #2: Studio owns the tensor split. When it emits no
+ # review run3 #2: Unsloth owns the tensor split. When it emits no
# --tensor-split (even split), a stale inherited LLAMA_ARG_TENSOR_SPLIT must
# be cleared in the TENSOR branch too (not just the layer downgrade), or the
- # child runs a split Studio didn't budget. The else (tensor) branch pops it.
+ # child runs a split Unsloth didn't budget. The else (tensor) branch pops it.
src = inspect.getsource(LlamaCppBackend.load_model)
compact = "".join(src.split())
# appears in both the layer branch and the tensor branch.
@@ -1005,14 +1005,14 @@ def test_qwen36_class_regression_picks_lower_ctx_with_mtp():
def test_mtp_draft_budget_prefers_user_extras_drafter():
# A user --model-draft in extras is appended last and wins at launch, so the
- # VRAM budget must size it first; then Studio's emitted mtp_draft_path (which
+ # VRAM budget must size it first; then Unsloth's emitted mtp_draft_path (which
# overrides LLAMA_ARG_SPEC_DRAFT_MODEL), then the env drafter (load_model is too
# entangled to drive end-to-end; assert the precedence at the source level).
# Whitespace-stripped so the check survives any formatter line-wrapping.
compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
- # CLI extras sized first (env={} so the env doesn't pre-empt Studio's drafter).
+ # CLI extras sized first (env={} so the env doesn't pre-empt Unsloth's drafter).
assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact
- # Order: CLI extras, then Studio's mtp_draft_path, then the env drafter.
+ # Order: CLI extras, then Unsloth's mtp_draft_path, then the env drafter.
assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact
- # The env must not be consulted before Studio's resolved drafter.
+ # The env must not be consulted before Unsloth's resolved drafter.
assert "_extra_args_mtp_draft_path(extra_args)ormtp_draft_path" not in compact
diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py
index 5cd7c876cc..b347c4aef8 100644
--- a/studio/backend/tests/test_multimodal_document.py
+++ b/studio/backend/tests/test_multimodal_document.py
@@ -3,7 +3,7 @@
"""Tests for PDF / document attachment translation on external providers.
-Studio adds a normalised `input_document` content part on
+Unsloth adds a normalised `input_document` content part on
ChatCompletionRequest so the frontend needn't know the per-provider
attachment shape:
diff --git a/studio/backend/tests/test_nudge_tool_calls_wiring.py b/studio/backend/tests/test_nudge_tool_calls_wiring.py
index e03fd0c7d7..82a6543aeb 100644
--- a/studio/backend/tests/test_nudge_tool_calls_wiring.py
+++ b/studio/backend/tests/test_nudge_tool_calls_wiring.py
@@ -3,7 +3,7 @@
"""Wiring guard for the plan-without-action ``nudge_tool_calls`` policy.
-Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths
+Decided policy: the re-prompt is ALWAYS ON for the Unsloth inference paths
(safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat +
Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off).
@@ -16,7 +16,7 @@ Mechanism (verified here without loading a model):
opt-in), while the GGUF loop keeps its pre-existing default-on behaviour
(``None`` keeps nudging) so an omitted flag never disables GGUF;
* the API request models default the flag to ``None`` (opt-in / off);
- * the Studio-facing routes forward the request's flag, and the Studio frontend
+ * the Unsloth-facing routes forward the request's flag, and the Unsloth frontend
sends ``nudge_tool_calls: true`` -- exercised behaviourally in
``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``.
"""
@@ -87,7 +87,7 @@ def test_api_request_models_default_the_flag_off():
def test_studio_routes_forward_the_request_flag():
- # The Studio chat frontend posts to /v1/chat/completions and /v1/messages
+ # The Unsloth chat frontend posts to /v1/chat/completions and /v1/messages
# with nudge_tool_calls=true; the route handlers forward the request value
# (external API clients that omit it fall back to the opt-in default).
from routes import inference as routes_inference
diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py
index 0b9a1e704f..295549c443 100644
--- a/studio/backend/tests/test_offline_gguf_cache_fallback.py
+++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py
@@ -897,7 +897,7 @@ class TestHfOfflineIfDnsDead:
assert "HF_HUB_OFFLINE" not in os.environ
def test_user_set_hf_hub_offline_is_preserved(self, dns, clean_offline_env, monkeypatch):
- # User explicitly set offline before launching Studio.
+ # User explicitly set offline before launching Unsloth.
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
dns.fail()
with _hf_offline_if_dns_dead() as did_set:
diff --git a/studio/backend/tests/test_offline_inference_parent.py b/studio/backend/tests/test_offline_inference_parent.py
index 71331220d6..bd0014ea64 100644
--- a/studio/backend/tests/test_offline_inference_parent.py
+++ b/studio/backend/tests/test_offline_inference_parent.py
@@ -139,7 +139,7 @@ class TestLoraDetectOffline:
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
- # Studio catches Exception broadly; pin that the call still happens
+ # Unsloth catches Exception broadly; pin that the call still happens
# (so cached LoRAs aren't missed) and returns fast via the mock.
class _OfflineModeIsEnabled(Exception):
pass
diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py
index 8742b84ae7..90fbd19297 100644
--- a/studio/backend/tests/test_openai_auto_switch.py
+++ b/studio/backend/tests/test_openai_auto_switch.py
@@ -689,7 +689,7 @@ def test_v1_models_retrieve_is_case_insensitive(monkeypatch):
def test_index_excludes_hidden_models(tmp_path, monkeypatch):
# The llama.cpp validation probe and RAG embedding weights are hidden from
- # Studio's pickers; they must never become auto-switch targets.
+ # Unsloth's pickers; they must never become auto-switch targets.
from types import SimpleNamespace
import routes.models as models_route
diff --git a/studio/backend/tests/test_openai_compaction.py b/studio/backend/tests/test_openai_compaction.py
index c7de0a9aed..6fad2c5eaf 100644
--- a/studio/backend/tests/test_openai_compaction.py
+++ b/studio/backend/tests/test_openai_compaction.py
@@ -86,7 +86,7 @@ def test_cloud_openai_sets_compaction_block(monkeypatch):
def test_cloud_openai_below_default_threshold_passes_through(monkeypatch):
- # Studio doesn't clamp the OpenAI side -- the API accepts whatever the
+ # Unsloth doesn't clamp the OpenAI side -- the API accepts whatever the
# caller sends, so a small probe like 60k still goes through.
captured = _capture(
monkeypatch,
diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py
index ace57588d3..c2eef0381f 100644
--- a/studio/backend/tests/test_openai_image_generation.py
+++ b/studio/backend/tests/test_openai_image_generation.py
@@ -4,7 +4,7 @@
"""Unit tests for OpenAI Responses API image_generation tool wiring.
The tool is a server-side Responses-API tool (``{type: "image_generation"}``);
-the result comes back as an ``image_generation_call`` output item, which Studio
+the result comes back as an ``image_generation_call`` output item, which Unsloth
translates into ``_toolEvent`` chunks so the chat adapter renders it inline.
Tests pin: the tool is added to the body only on a cloud OpenAI base when asked
for, the done event produces the expected chunks, and non-cloud bases drop it.
diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py
index 8725b28ac8..161c8743c4 100644
--- a/studio/backend/tests/test_openai_tool_passthrough.py
+++ b/studio/backend/tests/test_openai_tool_passthrough.py
@@ -119,7 +119,7 @@ class TestFriendlyUpstreamError:
raw = '{"error":{"code":400,"message":"Failed to initialize samplers: failed to parse grammar","type":"invalid_request_error"}}'
msg = _friendly_upstream_error(raw)
assert "failed to parse grammar" not in msg # raw body is not surfaced verbatim
- assert "tool-calling grammar" in msg and "Update Studio" in msg
+ assert "tool-calling grammar" in msg and "Update Unsloth" in msg
def test_failed_to_initialize_samplers_alone_matches(self):
assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers")
@@ -262,7 +262,7 @@ class TestChatMessageToolRoles:
def test_tool_empty_content_accepted(self):
# Empty tool output (mkdir, git add, ...) is routine in agentic loops;
- # OpenAI and llama-server both accept it, so Studio must not 400.
+ # OpenAI and llama-server both accept it, so Unsloth must not 400.
msg = ChatMessage(role = "tool", tool_call_id = "call_1", content = "")
assert msg.content == ""
@@ -400,7 +400,7 @@ class TestChatCompletionRequestToolFields:
assert req.session_id == "abc"
def test_stream_defaults_false_matching_openai_spec(self):
- # OpenAI defaults `stream` to false. Studio used to default true,
+ # OpenAI defaults `stream` to false. Unsloth used to default true,
# breaking naive curl/.NET clients (#5047) that omit it. Pin the fix.
req = self._make()
assert req.stream is False
@@ -664,7 +664,7 @@ class TestChatCompletionRequestToolFields:
raise AssertionError("client tools must use passthrough")
def generate_chat_completion_with_tools(self, **_kwargs):
- raise AssertionError("Studio tool loop must stay disabled")
+ raise AssertionError("Unsloth tool loop must stay disabled")
async def fake_passthrough(llama_backend, payload, model_name, **kwargs):
captured["body"] = inference_route._build_openai_passthrough_body(
@@ -707,11 +707,11 @@ class TestChatCompletionRequestToolFields:
assert monitor.active_count() == 0
def test_permission_mode_does_not_reject_client_tool_passthrough(self, monkeypatch):
- # A non-streaming client-tool passthrough (client tools, no Studio tool
+ # A non-streaming client-tool passthrough (client tools, no Unsloth tool
# loop) that also carries permission_mode "ask"/"auto" must reach the
# provider passthrough, not the confirm-without-stream guard: the
# validator leaves confirm_tool_calls unset for passthrough, and a bare
- # permission_mode only gates Studio's own local tool loop. An explicit
+ # permission_mode only gates Unsloth's own local tool loop. An explicit
# confirm_tool_calls=True still forces the local-confirm rejection.
# The pre-switch guard only runs when an automatic load may run, so force
# that predicate on to exercise it against a resident passthrough backend.
@@ -732,7 +732,7 @@ class TestChatCompletionRequestToolFields:
raise AssertionError("client tools must use passthrough")
def generate_chat_completion_with_tools(self, **_kwargs):
- raise AssertionError("Studio tool loop must stay disabled")
+ raise AssertionError("Unsloth tool loop must stay disabled")
async def fake_passthrough(llama_backend, payload, model_name, **kwargs):
inference_route.api_monitor.finish(kwargs.get("monitor_id"))
@@ -757,7 +757,7 @@ class TestChatCompletionRequestToolFields:
return self._v1_client(monkeypatch, _GGUFBackend())
# A process --enable-tools policy must not turn a client-tool passthrough
- # into a Studio local loop, so a policy of None or True both keep the
+ # into an Unsloth local loop, so a policy of None or True both keep the
# passthrough (the guard mirrors _explicit_studio_tool_loop_requested).
for policy in (None, True):
for mode in ("ask", "auto"):
@@ -810,7 +810,7 @@ class TestChatCompletionRequestToolFields:
assert "requires stream=true" in resp.json()["error"]["message"]
def test_permission_mode_policy_forced_local_loop_rejected_before_switch(self, monkeypatch):
- # A process --enable-tools policy forces Studio's own tool loop on even
+ # A process --enable-tools policy forces Unsloth's own tool loop on even
# when the request omits enable_tools and carries no client tools. A
# non-streaming ask/auto request is then confirm-gated with no stream to
# prompt on, so it must 400 at the pre-switch guard -- before
@@ -863,7 +863,7 @@ class TestChatCompletionRequestToolFields:
def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch):
# DiffusionGemma forces supports_tools off while passthrough stays
# available (#6851): enable_tools=True must not steal client tools
- # from the passthrough into a Studio tool loop that cannot run.
+ # from the passthrough into an Unsloth tool loop that cannot run.
import routes.inference as inference_route
captured = {}
@@ -883,7 +883,7 @@ class TestChatCompletionRequestToolFields:
raise AssertionError("client tools must use passthrough")
def generate_chat_completion_with_tools(self, **_kwargs):
- raise AssertionError("Studio tool loop cannot run on a non-tool backend")
+ raise AssertionError("Unsloth tool loop cannot run on a non-tool backend")
async def fake_passthrough(llama_backend, payload, model_name, **kwargs):
captured["body"] = inference_route._build_openai_passthrough_body(
@@ -2581,7 +2581,7 @@ class TestGgufVisionToolRouting:
raise AssertionError("plain GGUF path should not be used")
def _tools(**_kwargs):
- raise AssertionError("Studio tool loop should not steal response_format")
+ raise AssertionError("Unsloth tool loop should not steal response_format")
backend = SimpleNamespace(
is_loaded = True,
@@ -2654,7 +2654,7 @@ class TestGgufVisionToolRouting:
raise AssertionError("plain GGUF path should not be used")
def _tools(**_kwargs):
- raise AssertionError("Studio tool loop should not replace client tools")
+ raise AssertionError("Unsloth tool loop should not replace client tools")
backend = SimpleNamespace(
is_loaded = True,
@@ -2726,7 +2726,7 @@ class TestGgufVisionToolRouting:
yield "plain response"
def _tools(**_kwargs):
- raise AssertionError("tool_choice='none' must not start Studio's tool loop")
+ raise AssertionError("tool_choice='none' must not start Unsloth's tool loop")
backend = SimpleNamespace(
is_loaded = True,
@@ -2780,7 +2780,7 @@ class TestGgufVisionToolRouting:
raise AssertionError("plain GGUF path should not be used")
def _tools(**_kwargs):
- raise AssertionError("enabled_tools alone must not start Studio's tool loop")
+ raise AssertionError("enabled_tools alone must not start Unsloth's tool loop")
backend = SimpleNamespace(
is_loaded = True,
@@ -2844,7 +2844,7 @@ class TestGgufVisionToolRouting:
raise AssertionError("plain GGUF path should not be used")
def _tools(**_kwargs):
- raise AssertionError("enabled_tools alone must not start Studio's tool loop")
+ raise AssertionError("enabled_tools alone must not start Unsloth's tool loop")
backend = SimpleNamespace(
is_loaded = True,
diff --git a/studio/backend/tests/test_password_prompt_backstop.py b/studio/backend/tests/test_password_prompt_backstop.py
index 597eac1625..3c2c1956f9 100644
--- a/studio/backend/tests/test_password_prompt_backstop.py
+++ b/studio/backend/tests/test_password_prompt_backstop.py
@@ -3,7 +3,7 @@
"""Pre-tunnel terminal password gate: never publish a public Cloudflare URL
while the seeded default admin password is active. Imports run.py directly,
-so run under the Studio venv."""
+so run under the Unsloth venv."""
from __future__ import annotations
diff --git a/studio/backend/tests/test_permission_mode.py b/studio/backend/tests/test_permission_mode.py
index 3b7197fc49..4fc64a6291 100644
--- a/studio/backend/tests/test_permission_mode.py
+++ b/studio/backend/tests/test_permission_mode.py
@@ -1438,7 +1438,7 @@ def test_unknown_permission_mode_normalizes_to_ask_on_request_models():
def test_ask_auto_self_enable_confirm_on_chat_request():
# "Ask" gates every call, so a direct /chat/completions caller that requests
- # ask but omits the legacy confirm flag self-enables it when Studio's own tool
+ # ask but omits the legacy confirm flag self-enables it when Unsloth's own tool
# loop is requested. Only the router's loop-entry signals count (enable_tools /
# mcp_enabled); enabled_tools alone never starts the loop.
for loop in ({"enable_tools": True}, {"mcp_enabled": True}):
@@ -1481,7 +1481,7 @@ def test_ask_auto_self_enable_confirm_on_chat_request():
confirm_tool_calls = False,
)
assert req.confirm_tool_calls is False
- # A plain client-tool passthrough (client-supplied tools that Studio does not
+ # A plain client-tool passthrough (client-supplied tools that Unsloth does not
# execute) must NOT self-enable confirm, or the route rejects the passthrough.
req = ChatCompletionRequest(
messages = [{"role": "user", "content": "hi"}],
diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py
index 5e24ed752d..7cac3a9e99 100644
--- a/studio/backend/tests/test_providers_api.py
+++ b/studio/backend/tests/test_providers_api.py
@@ -38,11 +38,11 @@ BASE_URL = os.getenv("STUDIO_TEST_URL", "http://localhost:8000")
USERNAME = os.getenv("STUDIO_TEST_USER", "unsloth")
PASSWORD = os.getenv("STUDIO_TEST_PASSWORD", "")
-# Skip the whole module when no live Studio server / bootstrap password is
+# Skip the whole module when no live Unsloth server / bootstrap password is
# available (e.g. on CI) so pytest discovery does not error out.
pytestmark = pytest.mark.skipif(
not PASSWORD,
- reason = "Integration test requires a running Studio server; set STUDIO_TEST_PASSWORD to enable.",
+ reason = "Integration test requires a running Unsloth server; set STUDIO_TEST_PASSWORD to enable.",
)
# provider_type → (env var name, model for inference test)
diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py
index 0e1f74cefe..3a332ee19b 100644
--- a/studio/backend/tests/test_rag_embed_llama_server.py
+++ b/studio/backend/tests/test_rag_embed_llama_server.py
@@ -149,7 +149,7 @@ def test_build_env_gpu_inherits_devices(monkeypatch):
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1")
b = LlamaServerBackend()
env = b._build_env("/bin/llama-server", use_gpu = True)
- assert env.get("CUDA_VISIBLE_DEVICES") == "0,1" # inherit Studio's selection
+ assert env.get("CUDA_VISIBLE_DEVICES") == "0,1" # inherit Unsloth's selection
def test_use_gpu_explicit_modes(monkeypatch):
diff --git a/studio/backend/tests/test_recommended_folders_permission.py b/studio/backend/tests/test_recommended_folders_permission.py
index 33a457755e..b65695ad93 100644
--- a/studio/backend/tests/test_recommended_folders_permission.py
+++ b/studio/backend/tests/test_recommended_folders_permission.py
@@ -112,7 +112,7 @@ def test_path_under_unreadable_parent_returns_false_not_raises(tmp_path):
)
def test_demonstrates_the_underlying_stdlib_regression(tmp_path):
"""Documents *why* _safe_is_dir exists: the old bare pattern raises on
- the interpreters Studio ships on (3.12+)."""
+ the interpreters Unsloth ships on (3.12+)."""
parent = tmp_path / "ollama"
parent.mkdir()
os.chmod(parent, 0o000)
diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py
index 46dd0d42e4..69715649b7 100644
--- a/studio/backend/tests/test_responses_tool_passthrough.py
+++ b/studio/backend/tests/test_responses_tool_passthrough.py
@@ -120,7 +120,7 @@ class TestResponsesRequestTools:
def test_builtin_tool_type_passes_validation(self):
"""Non-function built-in tools (web_search, file_search, mcp, ...)
must not raise at validation so SDKs that default to them don't
- fail on Studio; they're filtered out during translation."""
+ fail on Unsloth; they're filtered out during translation."""
req = ResponsesRequest(
input = "hi",
tools = [{"type": "web_search_preview"}],
diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py
index 6e70c7cde4..699d0b74f5 100644
--- a/studio/backend/tests/test_rocm_oom_guard.py
+++ b/studio/backend/tests/test_rocm_oom_guard.py
@@ -36,7 +36,7 @@ class TestIsIntegratedSignal:
"""hipDeviceProp_t.integrated wins when truthy; 0/absent never downgrades.
Same universal gate PR #5988's UMA safetensors fast-load uses -- keeps
- Studio's two unified-memory consumers on one signal."""
+ Unsloth's two unified-memory consumers on one signal."""
def test_integrated_upgrades_unknown_apu(self) -> None:
# gfx1103 Phoenix iGPU: outside the hardcoded arch set, but the
diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py
index 915f82ac8e..31c728afca 100644
--- a/studio/backend/tests/test_safetensors_tool_loop.py
+++ b/studio/backend/tests/test_safetensors_tool_loop.py
@@ -2230,8 +2230,8 @@ def _reprompt_loop(*, auto_heal_tool_calls):
tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}],
execute_tool = exec_fn,
auto_heal_tool_calls = auto_heal_tool_calls,
- # Studio always nudges (always-on for the Studio inference paths); the
- # API opts in per request. Model the Studio caller here.
+ # Unsloth always nudges (always-on for the Unsloth inference paths); the
+ # API opts in per request. Model the Unsloth caller here.
nudge_tool_calls = True,
max_tool_iterations = 3,
)
@@ -3203,7 +3203,7 @@ class TestLoopBehaviour:
class TestLoopRePrompt:
- """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``."""
+ """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Unsloth always nudges, so these drive the loop with ``nudge_tool_calls=True``."""
def test_reasoning_intent_does_not_reprompt_a_visible_answer(self):
generations = 0
@@ -4258,7 +4258,7 @@ class TestPlanWithoutActionReprompt:
def test_omitted_nudge_flag_is_not_reprompted(self):
# The retry is new on this loop: API callers who do not send the flag
- # must keep today's behavior. Studio opts in explicitly.
+ # must keep today's behavior. Unsloth opts in explicitly.
loop, exec_fn = _make_loop(
turns = [
["I'll search the web for that."],
diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py
index 2c13e13bbb..a8c0c2305f 100644
--- a/studio/backend/tests/test_secure_tunnel_gate.py
+++ b/studio/backend/tests/test_secure_tunnel_gate.py
@@ -2,7 +2,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Cloudflare tunnel start gate, incl. --secure on loopback. Imports run.py
-directly, so run under the Studio venv."""
+directly, so run under the Unsloth venv."""
from __future__ import annotations
diff --git a/studio/backend/tests/test_server_disk_logging.py b/studio/backend/tests/test_server_disk_logging.py
index 05d03d869c..ce733c2aaa 100644
--- a/studio/backend/tests/test_server_disk_logging.py
+++ b/studio/backend/tests/test_server_disk_logging.py
@@ -3,7 +3,7 @@
"""Tests for the server session log + native-crash capture in run.py.
-Field regression: Studio "terminates without a warning" -- a native crash in
+Field regression: Unsloth "terminates without a warning" -- a native crash in
the GPU runtime kills the process with no Python traceback, and a desktop-
shortcut console closes before anything can be read. The server must tee its
console output to disk and aim faulthandler at the same file so even hard
diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py
index ac606e4627..d354c7e113 100644
--- a/studio/backend/tests/test_slot_offload_fit.py
+++ b/studio/backend/tests/test_slot_offload_fit.py
@@ -3,7 +3,7 @@
"""Tests for the offload-avoidance serving-slot reduction (`_slots_that_fit_on_gpu`).
-When a pinned context does not fit at the requested `--parallel` slot count, Studio would
+When a pinned context does not fit at the requested `--parallel` slot count, Unsloth would
flip to `--fit on` and llama-server offloads layers to host RAM, collapsing decode ~3x
(oobabooga #6718). Instead the loader retries the on-GPU fit at fewer slots and keeps the
largest count that stays fully on GPU (`-ngl -1`). These tests drive the real helper with
diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py
index 928b636e3e..087c00b648 100644
--- a/studio/backend/tests/test_studio_api.py
+++ b/studio/backend/tests/test_studio_api.py
@@ -11,7 +11,7 @@ the CLI's ``--help`` output:
1. curl -- basic chat completions (non-streaming)
2. curl -- streaming chat completions
3. Python OpenAI SDK -- streaming completions
- 4. curl -- Studio server-side tools (enable_tools=true)
+ 4. curl -- Unsloth server-side tools (enable_tools=true)
5. curl -- Standard OpenAI function calling (non-streaming)
6. curl -- Standard OpenAI function calling (streaming)
7. curl -- Standard OpenAI function calling (multi-turn tool loop)
@@ -31,7 +31,7 @@ Usage:
python tests/test_studio_api.py
python tests/test_studio_api.py --model unsloth/... --gguf-variant ...
- # Pytest mode, external server — start a Studio server yourself,
+ # Pytest mode, external server — start an Unsloth server yourself,
# then point pytest at it. Fastest iteration loop.
unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL &
export UNSLOTH_E2E_BASE_URL=http://127.0.0.1:8080
@@ -341,7 +341,7 @@ def _final_finish_reason(chunks: list[dict]) -> str | None:
def test_openai_tools_nonstream(base_url: str, api_key: str):
"""Standard OpenAI function calling, non-streaming, tool_choice='required'.
- Regression: before the fix, Studio stripped `tools` and the model
+ Regression: before the fix, Unsloth stripped `tools` and the model
returned plain text with finish_reason='stop'. After the fix,
llama-server's response is forwarded verbatim so the client sees
finish_reason='tool_calls' with a structured tool_calls array and
diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py
index 0d71b89d87..06f72d3b9f 100644
--- a/studio/backend/tests/test_tensor_parallel.py
+++ b/studio/backend/tests/test_tensor_parallel.py
@@ -420,7 +420,7 @@ def test_runtime_recovery_fires_for_user_env_mtp(monkeypatch):
# MTP driven by user extra_args / LLAMA_ARG_SPEC_TYPE leaves _speculative_type
# unset, but the launch flag still gates recovery on (pass-through MTP).
b = _recovery_backend()
- b._speculative_type = None # Studio stepped back; user/env owns the spec
+ b._speculative_type = None # Unsloth stepped back; user/env owns the spec
done = threading.Event()
captured = {}
diff --git a/studio/backend/tests/test_tool_confirm_stream.py b/studio/backend/tests/test_tool_confirm_stream.py
index b8e0472e12..0813f6b68d 100644
--- a/studio/backend/tests/test_tool_confirm_stream.py
+++ b/studio/backend/tests/test_tool_confirm_stream.py
@@ -3,12 +3,12 @@
"""End-to-end handshake test for the tool-confirmation gate, no model.
-The real Studio stream wrappers in ``routes/inference.py`` drive the
+The real Unsloth stream wrappers in ``routes/inference.py`` drive the
synchronous agentic generator with ``await asyncio.to_thread(next, gen,
...)`` so the blocking ``threading.Event`` wait runs off the event loop.
This test rebuilds that exact pattern around the real
``state.tool_approvals`` functions, served by a real uvicorn process on
-loopback (the same server Studio uses), and proves the load-bearing
+loopback (the same server Unsloth uses), and proves the load-bearing
property:
* ``tool_start`` reaches the client before the gate blocks, and
diff --git a/studio/backend/tests/test_tool_message_empty_content.py b/studio/backend/tests/test_tool_message_empty_content.py
index d63b16ce80..636a35f5a9 100644
--- a/studio/backend/tests/test_tool_message_empty_content.py
+++ b/studio/backend/tests/test_tool_message_empty_content.py
@@ -4,7 +4,7 @@
"""Empty ``role="tool"`` content must be accepted on the OpenAI-compat surface.
Agentic clients send ``content: ""`` when a command produced no output;
-OpenAI and llama-server both accept it. Studio used to 400, which standard
+OpenAI and llama-server both accept it. Unsloth used to 400, which standard
clients treat as non-retryable and kill the session. The validator must
normalize empty/missing tool content to ``""`` instead of raising.
"""
diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py
index 09af876da6..fb0989b306 100644
--- a/studio/backend/tests/test_tp_vision_regression.py
+++ b/studio/backend/tests/test_tp_vision_regression.py
@@ -625,7 +625,7 @@ def _fallback_loaded_backend(layer_preserves_tensor_intent: bool) -> LlamaCppBac
def test_tensor_off_echo_preserves_multi_gpu_fallback():
- """The Studio UI always sends tensor_parallel and echoes the /load response's
+ """The Unsloth 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)."""
diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py
index 7bf572e214..5d74bb7d28 100644
--- a/studio/backend/tests/test_trained_model_scan.py
+++ b/studio/backend/tests/test_trained_model_scan.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
-"""Tests for Studio trained-model discovery used by Chat."""
+"""Tests for Unsloth trained-model discovery used by Chat."""
import json
from pathlib import Path
diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py
index a2dc78bee2..5a477a084d 100644
--- a/studio/backend/tests/test_training_nan_loss_handling.py
+++ b/studio/backend/tests/test_training_nan_loss_handling.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-"""Pin Studio's behavior when a training event reports non-finite (NaN/Inf) loss.
+"""Pin Unsloth's behavior when a training event reports non-finite (NaN/Inf) loss.
The training event handler used to filter NaN/Inf to None silently while
leaving the previous finite loss in progress.loss — so the API kept reporting
diff --git a/studio/backend/tests/test_transformers_latest.py b/studio/backend/tests/test_transformers_latest.py
index 20616dccba..af48d674cc 100644
--- a/studio/backend/tests/test_transformers_latest.py
+++ b/studio/backend/tests/test_transformers_latest.py
@@ -1036,7 +1036,7 @@ def test_upgrade_check_mixed_pypi_main_reports_dev_only(monkeypatch):
def test_install_endpoint_not_mounted_on_v1():
- """The consented pip-install endpoint is a Studio admin action; it must live
+ """The consented pip-install endpoint is an Unsloth admin action; it must live
on studio_router (kept off the OpenAI-compatible /v1 mount), not router."""
from routes import inference as ri
diff --git a/studio/backend/utils/_studio_release_build.py b/studio/backend/utils/_studio_release_build.py
index 267197a202..07ede36912 100644
--- a/studio/backend/utils/_studio_release_build.py
+++ b/studio/backend/utils/_studio_release_build.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
-"""Build-stamped Studio release metadata.
+"""Build-stamped Unsloth release metadata.
Release builds may rewrite this module in the build workspace before creating
Python artifacts. Keep the committed value neutral so source checkouts do not
diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py
index cae8daf287..a3686c3a26 100644
--- a/studio/backend/utils/api_errors.py
+++ b/studio/backend/utils/api_errors.py
@@ -20,7 +20,7 @@ client-error responses on the ``/v1/*`` surface:
CRITICAL: the exception handlers installed by :func:`install_api_error_handlers`
are global, but they ONLY transform responses for paths that start with ``/v1/``.
For every other path (``/api/...``, frontend routes) they reproduce FastAPI's
-default behavior byte-for-byte, because the Studio frontend depends on the
+default behavior byte-for-byte, because the Unsloth frontend depends on the
``{"detail": ...}`` shape for ``/api/*``.
Public contract (other modules depend on these):
@@ -107,7 +107,7 @@ def anthropic_error_body(
Returns ``{"type": "error", "request_id": None, "error": {"type", "message"}}``.
``request_id`` is a required (nullable) field on the spec's ErrorResponse;
- Studio has no request-id system, so it is null. ``err_type`` defaults to
+ Unsloth has no request-id system, so it is null. ``err_type`` defaults to
:data:`ANTHROPIC_TYPE_BY_STATUS` for ``status`` (``"api_error"`` fallback).
"""
return {
@@ -192,7 +192,7 @@ def install_api_error_handlers(app) -> None:
Both handlers are global but only transform responses for OpenAI/Anthropic-
compatible surfaces (see :func:`wants_api_error_envelope`: the ``/v1/*`` mount
and the preview ``/p/.../v1/*`` mount). Every other path reproduces FastAPI's
- default ``{"detail": ...}`` behavior exactly so the Studio frontend keeps working.
+ default ``{"detail": ...}`` behavior exactly so the Unsloth frontend keeps working.
"""
@app.exception_handler(RequestValidationError)
diff --git a/studio/backend/utils/client_ip.py b/studio/backend/utils/client_ip.py
index 94acbf1809..cc48a096d2 100644
--- a/studio/backend/utils/client_ip.py
+++ b/studio/backend/utils/client_ip.py
@@ -4,12 +4,12 @@
"""Resolve the caller's IP for rate limiting.
Trust model, in order:
- 1. If the operator opts in via ``UNSLOTH_STUDIO_TRUST_FORWARDED`` (Studio behind
+ 1. If the operator opts in via ``UNSLOTH_STUDIO_TRUST_FORWARDED`` (Unsloth behind
their own reverse proxy), honor the *rightmost* ``X-Forwarded-For`` hop -- the
one the trusted proxy appended. The leftmost entry is client-controlled and
spoofable, so this assumes a proxy that appends (or overwrites) the header;
only enable the env var behind such a proxy.
- 2. If the socket peer is loopback, honor ``CF-Connecting-IP``. Studio's managed
+ 2. If the socket peer is loopback, honor ``CF-Connecting-IP``. Unsloth's managed
Cloudflare tunnel terminates at 127.0.0.1, so every tunneled visitor would
otherwise collapse onto the same socket peer (the local cloudflared process)
and share one rate-limit bucket. ``CF-Connecting-IP`` is set by Cloudflare's
diff --git a/studio/backend/utils/cpu_threads.py b/studio/backend/utils/cpu_threads.py
index 4ed0021054..91d577408d 100644
--- a/studio/backend/utils/cpu_threads.py
+++ b/studio/backend/utils/cpu_threads.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
-"""Early CPU thread-pool configuration for Studio processes."""
+"""Early CPU thread-pool configuration for Unsloth processes."""
import os
from typing import MutableMapping, Optional
diff --git a/studio/backend/utils/datasets/cache_safe.py b/studio/backend/utils/datasets/cache_safe.py
index e629210f33..d2dc7b737a 100644
--- a/studio/backend/utils/datasets/cache_safe.py
+++ b/studio/backend/utils/datasets/cache_safe.py
@@ -7,7 +7,7 @@ A shared HF datasets cache can contain subtrees owned by another user (for
example populated by an earlier root-run job). datasets then raises
"[Errno 13] Permission denied: ..._builder.lock" while locking the cached
builder, killing the training run even though the dataset itself is fine.
-Retry such loads in a Studio-owned cache so the run proceeds; the worst case
+Retry such loads in an Unsloth-owned cache so the run proceeds; the worst case
is one rebuild of the dataset in the fallback location.
"""
@@ -26,7 +26,7 @@ def studio_datasets_cache() -> str:
def load_dataset_cache_safe(*args, **kwargs):
- """datasets.load_dataset, retried in a Studio-owned cache on EACCES."""
+ """datasets.load_dataset, retried in an Unsloth-owned cache on EACCES."""
from datasets import load_dataset
try:
return load_dataset(*args, **kwargs)
diff --git a/studio/backend/utils/hardware/VRAM_ESTIMATION.md b/studio/backend/utils/hardware/VRAM_ESTIMATION.md
index a6b4de29d2..68ca1d5ffd 100644
--- a/studio/backend/utils/hardware/VRAM_ESTIMATION.md
+++ b/studio/backend/utils/hardware/VRAM_ESTIMATION.md
@@ -106,7 +106,7 @@ Non_flash_attention = B * num_attention_heads * S^2 * 2 * 12.0 * effective_layer
Activations = max(Per_layer_with_gc, Non_flash_attention)
```
-Studio resolves the attention implementation with Unsloth's
+Unsloth resolves the attention implementation with Unsloth's
`resolve_attention_implementation` helper and uses that result directly. The
estimator does not duplicate model-family attention policy.
diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py
index f5b64c45d0..91a06c9a2a 100644
--- a/studio/backend/utils/hardware/amd.py
+++ b/studio/backend/utils/hardware/amd.py
@@ -125,7 +125,7 @@ def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optiona
# amd-smi does not exist on Windows (neither Adrenalin nor the HIP SDK
# ship a CLI) and can be absent on minimal Linux installs. Disable the
# poller in one step instead of burning the 3-strike circuit breaker
- # on guaranteed FileNotFoundError spawns. Studio's VRAM display falls
+ # on guaranteed FileNotFoundError spawns. Unsloth's VRAM display falls
# back to torch mem_get_info.
if not _amd_smi_disabled:
logger.info(
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 117ad7b780..adc9a54aab 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -37,7 +37,7 @@ logger = get_logger(__name__)
# ── GPU index ordering ──────────────────────────────────────────────────────
# CUDA defaults to CUDA_DEVICE_ORDER=FASTEST_FIRST, numbering GPUs by compute
-# performance. nvidia-smi -- and every free-VRAM probe in Studio -- numbers GPUs
+# performance. nvidia-smi -- and every free-VRAM probe in Unsloth -- numbers GPUs
# by PCI bus id instead. On a mixed-GPU host (e.g. an RTX 5090 alongside an RTX
# PRO 6000) the two orderings disagree, so an index picked from nvidia-smi data
# ("the emptiest card is GPU 1") gets written into CUDA_VISIBLE_DEVICES and then
@@ -49,7 +49,7 @@ logger = get_logger(__name__)
# and spawn workers copy os.environ. setdefault so an explicit user override wins.
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
-# Studio workers can import MLX without importing unsloth first, so mirror the
+# Unsloth workers can import MLX without importing unsloth first, so mirror the
# package bootstrap here. Keep an explicit user value authoritative.
if platform.system() == "Darwin" and platform.machine() == "arm64":
os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1")
@@ -117,7 +117,7 @@ def _has_mlx() -> bool:
def _has_usable_mlx_stack() -> bool:
- """True only when the FULL Studio MLX training/export stack is usable
+ """True only when the FULL Unsloth MLX training/export stack is usable
(mlx + mlx-lm + mlx-vlm at the minimum versions unsloth-zoo requires), not
just a bare ``import mlx.core``. A backtracked/old mlx-vlm still imports but
breaks VLM Train/Export, so the training gate must match the self-heal's own
diff --git a/studio/backend/utils/helper_precache_settings.py b/studio/backend/utils/helper_precache_settings.py
index db19a2d028..e7d3c0e6dd 100644
--- a/studio/backend/utils/helper_precache_settings.py
+++ b/studio/backend/utils/helper_precache_settings.py
@@ -32,7 +32,7 @@ def helper_model_disabled_by_env() -> bool:
def get_helper_precache_enabled() -> bool:
"""Read the persisted startup pre-cache preference.
- Missing or unreadable settings default to False so Studio startup never
+ Missing or unreadable settings default to False so Unsloth startup never
performs optional network work unless the user explicitly opted in.
"""
try:
@@ -45,7 +45,7 @@ def get_helper_precache_enabled() -> bool:
def set_helper_precache_enabled(value: Any) -> bool:
- """Persist whether Studio should pre-cache the Helper LLM at startup."""
+ """Persist whether Unsloth should pre-cache the Helper LLM at startup."""
parsed = _coerce_bool(value)
if parsed is None:
raise ValueError("Helper LLM startup pre-cache must be true or false.")
diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py
index 9bc4a60fad..2628b99a2d 100644
--- a/studio/backend/utils/hf_xet_fallback.py
+++ b/studio/backend/utils/hf_xet_fallback.py
@@ -1,9 +1,9 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback.
+"""Unsloth shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback.
-Re-exports the shared API and injects Studio's marker-aware cache purge
+Re-exports the shared API and injects Unsloth's marker-aware cache purge
(``prepare_cache_for_transport``) so the download manager keeps its ``.transport``
marker semantics on the HTTP retry.
@@ -68,7 +68,7 @@ def _load_shared() -> bool:
_shared_available = True
_shared_import_error = None
return True
- except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF
+ except Exception as exc2: # noqa: BLE001 - degrade so Unsloth still boots with plain HF
_shared_import_error = exc2
_shared_available = False
import logging as _logging
@@ -263,7 +263,7 @@ __all__ = [
def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None:
- """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport``
+ """Unsloth's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport``
accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged,
not fatal to the retry."""
try:
@@ -273,7 +273,7 @@ def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None:
try:
from loggers import get_logger
get_logger(__name__).debug(
- "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc
+ "Unsloth prepare_cache_for_transport failed for %s: %s", repo_id, exc
)
except ModuleNotFoundError as logger_exc:
if logger_exc.name != "loggers":
@@ -294,8 +294,8 @@ def hf_hub_download_with_xet_fallback(
on_status: Optional[Callable[[str], None]] = None,
force_download: bool = False,
) -> str:
- """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep.
- ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path)."""
+ """Single-file download via the shared fallback with Unsloth's marker-aware HTTP-retry prep.
+ ``force_download`` re-fetches a newer blob over a cached one (Unsloth's model-update path)."""
return _shared_hf_hub_download_with_xet_fallback(
repo_id,
filename,
@@ -313,6 +313,6 @@ def hf_hub_download_with_xet_fallback(
def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str:
- """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep."""
+ """Whole-repo download via the shared fallback with Unsloth's marker-aware HTTP-retry prep."""
kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http)
return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs)
diff --git a/studio/backend/utils/host_policy.py b/studio/backend/utils/host_policy.py
index f506eadc03..55565bb338 100644
--- a/studio/backend/utils/host_policy.py
+++ b/studio/backend/utils/host_policy.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
-"""Bind-host trust policy for the Studio backend.
+"""Bind-host trust policy for the Unsloth backend.
Stdlib only -- safe to import without the rest of the backend.
diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py
index f6d3635301..31dbda63ea 100644
--- a/studio/backend/utils/llama_cpp_update.py
+++ b/studio/backend/utils/llama_cpp_update.py
@@ -82,7 +82,7 @@ def _utcnow() -> str:
def _find_binary() -> Optional[str]:
"""Locate the active llama-server binary via the inference backend's own
- resolver, so update targets exactly what Studio runs. Lazy import keeps the
+ resolver, so update targets exactly what Unsloth runs. Lazy import keeps the
heavy inference module off this module's import path."""
try:
from core.inference.llama_cpp import LlamaCppBackend
@@ -109,7 +109,7 @@ def _installer_script() -> Optional[Path]:
"""Locate install_llama_prebuilt.py. Honours UNSLOTH_LLAMA_INSTALLER, then
searches up from this file for both ``/install_llama_prebuilt.py`` and
``/studio/install_llama_prebuilt.py`` so it works in the dev tree and
- in an installed Studio layout."""
+ in an installed Unsloth layout."""
env = os.environ.get("UNSLOTH_LLAMA_INSTALLER")
if env and Path(env).is_file():
return Path(env)
@@ -227,7 +227,7 @@ def _is_under(path: Path, root: Path) -> bool:
def _llama_install_root(binary: Optional[str]) -> Optional[Path]:
- """The Studio-managed llama.cpp root the active binary lives under, or None
+ """The Unsloth-managed llama.cpp root the active binary lives under, or None
when the binary is unmanaged. Installing anywhere the active binary is not
would not replace what _find_llama_server_binary runs (which prefers a pinned
LLAMA_SERVER_PATH, then UNSLOTH_LLAMA_CPP_PATH, then a llama.cpp tree), so we
@@ -327,7 +327,7 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
def _is_external_link(path: Optional[Path]) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
- the user's own llama.cpp checkout, so Studio must never auto-update it."""
+ the user's own llama.cpp checkout, so Unsloth must never auto-update it."""
if path is None:
return False
try:
@@ -635,7 +635,7 @@ def start_update() -> dict:
"reason": "local_link",
"message": (
"llama.cpp is a local directory linked with --with-llama-cpp-dir; "
- "Studio won't replace it. Update your own llama.cpp checkout instead."
+ "Unsloth won't replace it. Update your own llama.cpp checkout instead."
),
"job": get_update_status()["job"],
}
diff --git a/studio/backend/utils/mlx_repair.py b/studio/backend/utils/mlx_repair.py
index 7e1c9864c9..4ea1ec62f5 100644
--- a/studio/backend/utils/mlx_repair.py
+++ b/studio/backend/utils/mlx_repair.py
@@ -3,7 +3,7 @@
"""Best-effort MLX self-heal for Apple Silicon.
-On macOS, Studio enables Train/Export only when the MLX training/export stack is
+On macOS, Unsloth enables Train/Export only when the MLX training/export stack is
usable (see utils.hardware.hardware.detect_hardware -> CHAT_ONLY). MLX is pulled
only transitively via unsloth-zoo, and a resolver backtrack (mlx-vlm ->
transformers>=5 vs the single-env transformers pin) can silently drop it, leaving
@@ -13,7 +13,7 @@ a background thread, then re-detects so the gate re-opens without a manual
The install mirrors the main Apple Silicon installer (install_python_stack.py):
it points UV_OVERRIDE at overrides-darwin-arm64.txt so the resolver keeps the
-Studio transformers pin AND installs a current mlx-vlm, and it requires the same
+Unsloth transformers pin AND installs a current mlx-vlm, and it requires the same
minimum versions unsloth-zoo declares so a backtracked old mlx-vlm (which still
imports but breaks VLM Train/Export) is never accepted as healthy.
@@ -69,11 +69,11 @@ _MLX_REINSTALL_ARGS = tuple(
# reject anything. mlx/mlx-metal ship wheels only (no sdist on PyPI) and
# mlx-lm/mlx-vlm publish py3-none-any wheels, so requiring wheels does not break a
# healthy self-heal; if a wheel is genuinely unavailable the install fails and
-# Studio stays chat-only (the existing safe fallback) until `unsloth studio update`.
+# Unsloth stays chat-only (the existing safe fallback) until `unsloth studio update`.
_ONLY_BINARY_ARG = "--only-binary=:all:"
# Allowlist of environment variables forwarded to the install subprocess. The
# self-heal runs without confirmation on the default startup path, so it must not
-# hand resolver/build code the full Studio environment. Everything outside this
+# hand resolver/build code the full Unsloth environment. Everything outside this
# set is dropped, which excludes three dangerous classes by construction:
# * secrets (HF_TOKEN, AWS_*, WANDB_API_KEY, ...) that a malicious wheel/sdist
# build hook would otherwise read straight out of os.environ;
@@ -207,13 +207,13 @@ def _mlx_install_env() -> dict[str, str]:
The self-heal runs without confirmation on the default startup path, so it
forwards only the variables uv genuinely needs (see _MLX_ENV_ALLOWLIST) instead
- of the full Studio environment: secrets and package-source redirects in
+ of the full Unsloth environment: secrets and package-source redirects in
os.environ are dropped so a malicious resolver-selected artifact cannot read
- Studio secrets or be steered to a hostile index.
+ Unsloth secrets or be steered to a hostile index.
Mirror the main installer (install_python_stack.py) by pointing UV_OVERRIDE at
overrides-darwin-arm64.txt, which relaxes mlx-vlm/mlx-lm's transformers>=5
- requirement to >=4.57.6. Without it, uv keeps the Studio transformers pin only
+ requirement to >=4.57.6. Without it, uv keeps the Unsloth transformers pin only
by silently backtracking mlx-vlm to an old, unsupported version (uv honours
UV_OVERRIDE; plain pip ignores it, so the transformers constraint below is the
pip-path safety net). We set UV_OVERRIDE ourselves, so a poisoned one in the
@@ -234,17 +234,17 @@ def _mlx_install_env() -> dict[str, str]:
def _transformers_constraint_args() -> tuple[list[str], str | None]:
"""Pin transformers to the running version for the mlx install.
- The install must never upgrade transformers underneath a running Studio
+ The install must never upgrade transformers underneath a running Unsloth
(the single-env install pins transformers==4.57.6). With UV_OVERRIDE set this
is belt-and-suspenders; on the plain-pip path (no UV_OVERRIDE support) it is
the actual guard -- the resolver either finds an mlx build compatible with the
- pin or fails, leaving us chat-only rather than breaking Studio. Returns
+ pin or fails, leaving us chat-only rather than breaking Unsloth. Returns
(pip args, temp file path to clean up).
Read the version from installed metadata rather than `import transformers`:
transformers can have valid metadata yet fail to import (e.g. an incompatible
huggingface_hub), and in that case we still want to pin it so the mlx install
- cannot quietly upgrade it out from under Studio."""
+ cannot quietly upgrade it out from under Unsloth."""
from importlib.metadata import PackageNotFoundError, version as _dist_version
try:
@@ -263,10 +263,10 @@ def attempt_mlx_repair(*, timeout: int = _REPAIR_TIMEOUT_S) -> bool:
"""Install a usable mlx/mlx-lm/mlx-vlm stack by name into the running venv.
Best-effort; returns True iff the resulting stack meets unsloth-zoo's minimums
(so a backtracked old mlx-vlm is rejected, not accepted). transformers is held
- at its pinned version so the install can never upgrade it underneath Studio."""
+ at its pinned version so the install can never upgrade it underneath Unsloth."""
# Prepare the constraint inside the try: this runs on a daemon thread, so an
# exception here (e.g. tempfile.mkstemp failing on a full disk or bad TMPDIR)
- # must leave Studio chat-only, not crash the background self-heal thread.
+ # must leave Unsloth chat-only, not crash the background self-heal thread.
constraint_path = None
try:
constraint_args, constraint_path = _transformers_constraint_args()
@@ -279,7 +279,7 @@ def attempt_mlx_repair(*, timeout: int = _REPAIR_TIMEOUT_S) -> bool:
)
if cmd is None:
logger.warning(
- "MLX self-heal requires uv so Studio can apply dependency overrides; "
+ "MLX self-heal requires uv so Unsloth can apply dependency overrides; "
"staying chat-only. Run `unsloth studio update` to restore uv."
)
return False
diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py
index b6b080b1c4..f2125ad034 100644
--- a/studio/backend/utils/models/checkpoints.py
+++ b/studio/backend/utils/models/checkpoints.py
@@ -37,7 +37,7 @@ def _checkpoint_sort_key(checkpoint_path: Path) -> tuple[int, int, str]:
def _infer_base_model_from_history(checkpoint_dir: Path) -> Optional[str]:
- """Best-effort base-model lookup using persisted Studio run metadata."""
+ """Best-effort base-model lookup using persisted Unsloth run metadata."""
checkpoint_name = checkpoint_dir.name
resolved_checkpoint_dir = str(checkpoint_dir.resolve())
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 284bbb5745..dadf103cea 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -2083,7 +2083,7 @@ def _has_model_weight_files(model_dir: Path) -> bool:
def _detect_training_output_type(model_dir: Path) -> Optional[str]:
- """Classify a Studio training output as LoRA or full finetune."""
+ """Classify an Unsloth training output as LoRA or full finetune."""
adapter_config = model_dir / "adapter_config.json"
adapter_model = model_dir / "adapter_model.safetensors"
if adapter_config.exists() or adapter_model.exists():
@@ -2105,7 +2105,7 @@ def _looks_like_lora_adapter(model_dir: Path) -> bool:
def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str, str]]:
- """Scan outputs folder for trained Studio models.
+ """Scan outputs folder for trained Unsloth models.
Returns:
List of (display_name, model_path, model_type), where model_type is
diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py
index 759681da3f..1faa2b1281 100644
--- a/studio/backend/utils/paths/storage_roots.py
+++ b/studio/backend/utils/paths/storage_roots.py
@@ -36,7 +36,7 @@ def _infer_studio_home_from_venv() -> Path | None:
def studio_root() -> Path:
- """Studio install root.
+ """Unsloth install root.
Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then sys.prefix
inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins if
@@ -62,7 +62,7 @@ def cache_root() -> Path:
def studio_bin_root() -> Path:
- """Dir for Studio-managed executables (the `unsloth` shim, downloaded tools like cloudflared)."""
+ """Dir for Unsloth-managed executables (the `unsloth` shim, downloaded tools like cloudflared)."""
return studio_root() / "bin"
@@ -443,7 +443,7 @@ def resolve_export_write_dir(path_value: str | None = None) -> Path:
Unlike :func:`resolve_export_dir`, this function passes absolute
paths through as-is so users can target a different drive when
- their Studio install lives on a constrained system volume
+ their Unsloth install lives on a constrained system volume
(see :gh-issue:`6082`). Used only by the export write path.
"""
if not path_value or not str(path_value).strip():
diff --git a/studio/backend/utils/preview_rate_limit.py b/studio/backend/utils/preview_rate_limit.py
index dd38cfd5e7..c59a1bf5b3 100644
--- a/studio/backend/utils/preview_rate_limit.py
+++ b/studio/backend/utils/preview_rate_limit.py
@@ -5,7 +5,7 @@
A signed link stops ref guessing, but anyone with a link can still drive GPU
generation. This bounds sustained abuse from a single source. In-process and
-single-worker only (like the login limiter in ``routes/auth.py``); Studio runs as
+single-worker only (like the login limiter in ``routes/auth.py``); Unsloth runs as
one uvicorn process, so a shared store isn't needed.
"""
diff --git a/studio/backend/utils/process_lifetime.py b/studio/backend/utils/process_lifetime.py
index 3ffd54cc26..c63227ae86 100644
--- a/studio/backend/utils/process_lifetime.py
+++ b/studio/backend/utils/process_lifetime.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
-"""Bind Studio child processes to the parent's lifetime so none survive an
+"""Bind Unsloth child processes to the parent's lifetime so none survive an
abnormal parent exit (terminal-window close, Task Manager "End Task", SIGKILL,
crash) -- the cooperative shutdown path only runs on graceful exits.
@@ -139,7 +139,7 @@ def _install_windows_job() -> None:
kernel32.CloseHandle(job)
return
# AssignProcessToJobObject(parent) makes children inherit the job. May
- # fail if Studio already runs inside an incompatible host job (pre-Win8);
+ # fail if Unsloth already runs inside an incompatible host job (pre-Win8);
# degrade to the cooperative path rather than blocking startup.
if not kernel32.AssignProcessToJobObject(job, kernel32.GetCurrentProcess()):
kernel32.CloseHandle(job)
diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py
index 9c18070fbb..82ade74bba 100644
--- a/studio/backend/utils/studio_version.py
+++ b/studio/backend/utils/studio_version.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
-"""Network-free Studio release version resolution for display-only UI."""
+"""Network-free Unsloth release version resolution for display-only UI."""
from __future__ import annotations
@@ -20,7 +20,7 @@ _MAX_VERSION_LENGTH = 64
def is_valid_studio_release_version(value: object) -> bool:
- """Return True for Studio release tags such as ``v0.1.39-beta``."""
+ """Return True for Unsloth release tags such as ``v0.1.39-beta``."""
if not isinstance(value, str):
return False
version = value.strip()
@@ -102,7 +102,7 @@ def _git_branch(repo_root: Path) -> str | None:
def get_studio_version(repo_root: Path | None = None) -> str:
- """Return the installed Studio release tag for display, or ``dev``.
+ """Return the installed Unsloth release tag for display, or ``dev``.
Intentionally separate from the PyPI ``unsloth`` package version used by
update checks. Never performs network requests.
diff --git a/studio/backend/utils/training_runs.py b/studio/backend/utils/training_runs.py
index dc2535e570..dcdfa1395d 100644
--- a/studio/backend/utils/training_runs.py
+++ b/studio/backend/utils/training_runs.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
-"""Helpers for naming and describing Studio training runs."""
+"""Helpers for naming and describing Unsloth training runs."""
from __future__ import annotations
diff --git a/studio/backend/utils/transformers_latest.py b/studio/backend/utils/transformers_latest.py
index 40c8f729a5..9f1d11be5b 100644
--- a/studio/backend/utils/transformers_latest.py
+++ b/studio/backend/utils/transformers_latest.py
@@ -5,7 +5,7 @@
When a model's ``model_type`` is absent from every installed transformers overlay
(base 4.57.x plus the .venv_t5_530/550/510 sidecars and, if provisioned, .venv_t5_latest),
-Studio cannot load it today. This module answers, without authentication, code execution,
+Unsloth cannot load it today. This module answers, without authentication, code execution,
or trust_remote_code:
1. Does the LATEST transformers release on PyPI ship this ``model_type``?
@@ -387,7 +387,7 @@ def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dic
_SHADOWABLE_DEPS = frozenset({"tokenizers", "safetensors"})
# Provided by the sidecar recipe; checked against its pin, not the base env.
_SIDECAR_PROVIDED = {"huggingface-hub": "1.8.0", "hf-xet": "1.4.2"}
-# CLI-only; never imported at runtime in Studio's workers.
+# CLI-only; never imported at runtime in Unsloth's workers.
_IGNORED_DEPS = frozenset({"typer"})
@@ -538,7 +538,7 @@ def _install_latest_transformers_locked(version: str, before_swap = None) -> dic
return {
"success": False,
"version": version,
- "message": "Cannot install: Studio is in offline mode.",
+ "message": "Cannot install: Unsloth is in offline mode.",
}
# Re-verify against a LIVE snapshot (a release may land inside the cache TTL);
# fall back to the cached one on fetch failure.
@@ -573,13 +573,13 @@ def _install_latest_transformers_locked(version: str, before_swap = None) -> dic
"version": version,
"message": "Cannot install transformers "
f"{version}: this environment does not satisfy {', '.join(blockers)}. "
- "A Studio update is required first.",
+ "An Unsloth update is required first.",
}
if not ensure_latest_transformers_venv(version, extra_packages, before_swap = before_swap):
return {
"success": False,
"version": version,
- "message": f"Installing transformers {version} failed; see the Studio logs.",
+ "message": f"Installing transformers {version} failed; see the Unsloth logs.",
}
_invalidate_capability_caches()
return {
diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py
index 9f9f8aa3de..1fbcc9f46f 100644
--- a/studio/backend/utils/transformers_version.py
+++ b/studio/backend/utils/transformers_version.py
@@ -2151,7 +2151,7 @@ def end_sidecar_swap() -> None:
def sidecar_swap_in_progress() -> bool:
"""True while a .venv_t5_latest install or repair holds the reservation,
- in this process or any other Studio process (lock file)."""
+ in this process or any other Unsloth process (lock file)."""
return sidecar_swap_kind() is not None
diff --git a/studio/backend/utils/upload_limits.py b/studio/backend/utils/upload_limits.py
index c21ea69af7..fff0ac423f 100644
--- a/studio/backend/utils/upload_limits.py
+++ b/studio/backend/utils/upload_limits.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
-"""Shared Studio upload/request size limits."""
+"""Shared Unsloth upload/request size limits."""
from __future__ import annotations
diff --git a/studio/frontend/.npmrc b/studio/frontend/.npmrc
index 19783b5ff4..414379da6e 100644
--- a/studio/frontend/.npmrc
+++ b/studio/frontend/.npmrc
@@ -1,4 +1,4 @@
-# Studio frontend npm configuration.
+# Unsloth frontend npm configuration.
#
# Mini Shai-Hulud / Axios-style supply chain defense.
# Requires npm >=11.10.0. Refuses tarballs published less than 7 days ago,
diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx
index c35706e50a..a7e9469cfc 100644
--- a/studio/frontend/src/app/provider.tsx
+++ b/studio/frontend/src/app/provider.tsx
@@ -412,7 +412,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
{desktopBooting ? (
-
Preparing Studio
+
Preparing Unsloth
The local backend is ready. Signing in to your desktop session
before loading chats.
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 8e06181585..ff4bd4bebd 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -1679,7 +1679,7 @@ export function HubModelPicker({
const deviceType = usePlatformStore((s) => s.deviceType);
const isMac = deviceType === "mac";
- // Drop models Studio can't run for chat (diffusion / image / video / etc.)
+ // Drop models Unsloth can't run for chat (diffusion / image / video / etc.)
// using the Hub's classifier on the tags the listing already carries.
const isChatSupported = useCallback(
(r: HfModelResult) =>
@@ -1730,7 +1730,7 @@ export function HubModelPicker({
let rows = recommendedSearch.results
.filter((r) => !isHiddenModelId(r.id))
.filter((r) => !isMobileVariant(r.id));
- // Drop models Studio can't run for chat (diffusion / image / video / etc.).
+ // Drop models Unsloth can't run for chat (diffusion / image / video / etc.).
rows = rows.filter(isChatSupported);
// With no explicit format, show the device-recommended formats (GGUF, plus
// MLX on Mac). When the user picks a format, honor it instead so Safetensors
@@ -1868,7 +1868,7 @@ export function HubModelPicker({
// eslint-disable-next-line react-hooks/exhaustive-deps
[lmStudioModels, downloadedSort, formatFilter, loadTimes, localQuery],
);
- // Local ./models entries. Chat-only Studio runs GGUF (any host) and MLX (Mac
+ // Local ./models entries. Chat-only Unsloth runs GGUF (any host) and MLX (Mac
// only), so raw checkpoints there are hidden (mirrors the cached non-GGUF
// rule). An MLX build a Mac user dropped in ./models stays selectable.
const sortedLocalDir = useMemo(
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 62b8af6e3a..235dcb3c3d 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -1970,7 +1970,7 @@ function isNativeComposing(event: Event) {
}
// Fallback timeout for stuck IME composition. With Chrome on Windows against
-// a WSL-hosted Studio (issue #5546), `compositionend` never fires after the
+// a WSL-hosted Unsloth (issue #5546), `compositionend` never fires after the
// candidate commits, so `composingRef` stays true and Send stays disabled.
// Every compositionupdate / non-composing input resets the timer; only a true
// gap-after-commit lets it fire. 2500ms is above a normal candidate-window
diff --git a/studio/frontend/src/components/ui/confetti.tsx b/studio/frontend/src/components/ui/confetti.tsx
index 892bffdb18..35f5913240 100644
--- a/studio/frontend/src/components/ui/confetti.tsx
+++ b/studio/frontend/src/components/ui/confetti.tsx
@@ -34,7 +34,7 @@ export type ConfettiRef = Api | null;
const ConfettiContext = createContext
({} as Api);
-// Studio CSP blocks canvas-confetti's default blob: worker, so force
+// Unsloth CSP blocks canvas-confetti's default blob: worker, so force
// useWorker: false. Module-scoped so the prop default keeps stable
// identity across renders (`canvasRef` depends on `globalOptions`).
const DEFAULT_GLOBAL_OPTIONS: ConfettiGlobalOptions = {
diff --git a/studio/frontend/src/features/auth/components/auth-form.tsx b/studio/frontend/src/features/auth/components/auth-form.tsx
index 119471da10..73db10d41b 100644
--- a/studio/frontend/src/features/auth/components/auth-form.tsx
+++ b/studio/frontend/src/features/auth/components/auth-form.tsx
@@ -298,7 +298,7 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
// reset-password"), which the installer puts on PATH on every platform.
// Do NOT rewrite it to a relative Windows path like
// ".\unsloth_studio\Scripts\unsloth.exe ..." -- that only resolves inside
- // the Studio home dir and fails with CommandNotFoundException elsewhere.
+ // the Unsloth home dir and fails with CommandNotFoundException elsewhere.
// Show the backend message as-is.
const msg = err instanceof Error ? err.message : "Auth failed.";
setError(msg);
diff --git a/studio/frontend/src/features/chat/artifacts/html-frame.tsx b/studio/frontend/src/features/chat/artifacts/html-frame.tsx
index b26f2f6685..36e3ed8a5b 100644
--- a/studio/frontend/src/features/chat/artifacts/html-frame.tsx
+++ b/studio/frontend/src/features/chat/artifacts/html-frame.tsx
@@ -28,7 +28,7 @@ export function buildArtifactSrcDoc(code: string): string {
}
// Preview iframes intentionally omit allow-downloads: generated canvases can
-// offer their own UI, but downloads must go through Studio's explicit
+// offer their own UI, but downloads must go through Unsloth's explicit
// copy/download controls outside the no-same-origin sandbox.
export function ArtifactHtmlFrame({
code,
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
index 0a0df1139b..bfb3eeb14c 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
@@ -45,7 +45,7 @@ export function groupThreads(
for (const t of threads) {
// Coerce archived to a boolean before comparing. Legacy threads (from the
- // older browser-only Studio, or any record predating the archived field)
+ // older browser-only Unsloth, or any record predating the archived field)
// can have archived === undefined or null; a raw `!== archived` comparison
// would drop those from BOTH the Recents (archived=false) and Archived
// (archived=true) lists, hiding existing chats. Treat missing as false.
diff --git a/studio/frontend/src/features/chat/lib/friendly-names.ts b/studio/frontend/src/features/chat/lib/friendly-names.ts
index 79744b3181..bc9c77b12d 100644
--- a/studio/frontend/src/features/chat/lib/friendly-names.ts
+++ b/studio/frontend/src/features/chat/lib/friendly-names.ts
@@ -5,7 +5,7 @@
* Friendly default names for auto-created OpenAI shell containers, used by the
* chat-adapter's lazy-create path (Code pill on, no thread container, non-default
* TTL). Goal: a memorable label like "otter" instead of "chat-abc12345"; users
- * can still rename via the Studio alias map.
+ * can still rename via the Unsloth alias map.
*
* The list is curated to be unambiguous, non-offensive nouns from natural
* categories (animals, plants, geography, materials, weather), avoid
diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts
index 79c9a3205c..ec251cdada 100644
--- a/studio/frontend/src/features/chat/provider-capabilities.ts
+++ b/studio/frontend/src/features/chat/provider-capabilities.ts
@@ -409,7 +409,7 @@ function isGeminiImageModel(modelId: string): boolean {
* Whether the saved Gemini connection points at a custom OpenAI-compat gateway
* (any non-Google host). The backend `_is_openai_compatible` routes these
* through `/chat/completions` instead of the native translator, so native Gemini
- * tool envelopes never reach them. Hide the matching Studio pills here so the
+ * tool envelopes never reach them. Hide the matching Unsloth pills here so the
* request, builder, and UI agree.
*/
export function isGeminiCustomOpenAICompatBase(
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index af99458349..192ce1ec69 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -671,7 +671,7 @@ type ChatRuntimeStore = {
// Describe figures/charts at ingest time (vision model required).
ragCaptionFigures: boolean;
/**
- * When on, local Studio tool calls pause for an explicit allow/deny in the
+ * When on, local Unsloth tool calls pause for an explicit allow/deny in the
* chat before they run.
*/
confirmToolCalls: boolean;
@@ -1824,7 +1824,7 @@ export const useChatRuntimeStore = create((set, get) => ({
setContextUsage: (contextUsage) => set({ contextUsage }),
}));
-// Mirror token edits made through the shared store (e.g. Studio's field).
+// Mirror token edits made through the shared store (e.g. Unsloth's field).
const unsubscribeHfTokenMirror = mirrorHfTokenInto(useChatRuntimeStore);
if (import.meta.hot) {
import.meta.hot.dispose(unsubscribeHfTokenMirror);
diff --git a/studio/frontend/src/features/chat/utils/chat-history-storage.ts b/studio/frontend/src/features/chat/utils/chat-history-storage.ts
index 00df3657b1..2ed17c26a0 100644
--- a/studio/frontend/src/features/chat/utils/chat-history-storage.ts
+++ b/studio/frontend/src/features/chat/utils/chat-history-storage.ts
@@ -324,7 +324,7 @@ async function importLegacyChatsIfNeeded(): Promise {
if (legacyChatImportPromise) return legacyChatImportPromise;
legacyChatImportPromise = (async () => {
- // Fast-path: no Dexie DB -- new user, never had browser-only Studio.
+ // Fast-path: no Dexie DB -- new user, never had browser-only Unsloth.
if (await dexieDbAbsent()) {
markLegacyChatImportDone();
return;
diff --git a/studio/frontend/src/features/hub/download-manager/api.ts b/studio/frontend/src/features/hub/download-manager/api.ts
index 3373a47b04..2c55b90edd 100644
--- a/studio/frontend/src/features/hub/download-manager/api.ts
+++ b/studio/frontend/src/features/hub/download-manager/api.ts
@@ -12,7 +12,7 @@ function parseErrorText(status: number, body: unknown): string {
const detail = (body as { detail?: unknown }).detail;
const formatted = formatFastApiDetail(detail);
if (status === 405) {
- return `${formatted || "Method Not Allowed"} - the Studio backend did not accept this API method. Restart Studio so the frontend and backend are on the same build.`;
+ return `${formatted || "Method Not Allowed"} - the Unsloth backend did not accept this API method. Restart Unsloth so the frontend and backend are on the same build.`;
}
if (formatted) return formatted;
const message = (body as { message?: unknown }).message;
diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
index a5ce25dd32..e60786decb 100644
--- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
@@ -113,7 +113,7 @@ export function ModelSelectionStep() {
return applyPriorityOrdering(ids);
}, [hfResults]);
- // Match Studio: only show exception signals (OOM/TIGHT) in training flows.
+ // Match Unsloth: only show exception signals (OOM/TIGHT) in training flows.
const vramMap = useMemo(() => {
const fitMap = buildModelVramMap(
hfResults,
diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx
index 86125585cc..4ccc43d16c 100644
--- a/studio/frontend/src/features/settings/components/usage-examples.tsx
+++ b/studio/frontend/src/features/settings/components/usage-examples.tsx
@@ -229,12 +229,12 @@ curl.exe ${base}/v1/chat/completions \`
}
// A second OpenAI call naming a different downloaded GGUF: with auto-switch on,
-// Studio loads it before serving, so the model field selects the served model.
+// Unsloth loads it before serving, so the model field selects the served model.
function pythonSwitchDemo(): string {
return `
# "Switch model by request" is on: replace the model below with another GGUF you
-# have downloaded and Studio loads it before serving. Unknown names keep serving
+# have downloaded and Unsloth loads it before serving. Unknown names keep serving
# the current model.
response = client.chat.completions.create(
model=${j(SWITCH_MODEL)},
@@ -349,7 +349,7 @@ function javascriptSwitchDemo(): string {
return `
// "Switch model by request" is on: replace the model below with another GGUF you
-// have downloaded and Studio loads it before serving. Unknown names keep serving
+// have downloaded and Unsloth loads it before serving. Unknown names keep serving
// the current model.
const switchResponse = await client.chat.completions.create({
model: ${j(SWITCH_MODEL)},
@@ -512,7 +512,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
if (!localAgentDetection) {
setDetectedAgents([]);
// A previously auto-picked agent was only ever verified against the
- // Studio backend's PATH, which is meaningless now that this panel no
+ // Unsloth backend's PATH, which is meaningless now that this panel no
// longer targets a loopback base -- don't leave it selected, but
// never touch a choice the user made by hand.
if (!agentPickedByUserRef.current) {
diff --git a/studio/frontend/src/features/studio/recent-trainings-section.tsx b/studio/frontend/src/features/studio/recent-trainings-section.tsx
index 43808a7ca9..ba65d7f736 100644
--- a/studio/frontend/src/features/studio/recent-trainings-section.tsx
+++ b/studio/frontend/src/features/studio/recent-trainings-section.tsx
@@ -8,7 +8,7 @@ import { HistoryCardGrid } from "./history-card-grid";
/**
* Recent training runs surfaced on Data Recipes and Export. Selecting a run
- * stores its id and navigates to Studio, which auto-opens its History tab.
+ * stores its id and navigates to Unsloth, which auto-opens its History tab.
* Renders nothing once we know there are no runs.
*/
export function RecentTrainingsSection() {
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 a927f83fd8..d204e24f47 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -1004,7 +1004,7 @@ export const useTrainingConfigStore = create()(
}
if (version < 12) {
// hfToken moved to the shared hf-token-store; seed it once so an
- // existing Studio-only token isn't lost.
+ // existing Unsloth-only token isn't lost.
const legacyToken = typeof s.hfToken === "string" ? s.hfToken.trim() : "";
if (legacyToken && !getHfToken()) {
useHfTokenStore.getState().setToken(legacyToken);
diff --git a/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx b/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx
index 98c590e5e4..3e39a8ff1c 100644
--- a/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx
+++ b/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx
@@ -82,7 +82,7 @@ export function TransformersUpgradeDialog() {
<>
Even the latest transformers release on PyPI does not
support it yet: the architecture is only available on the
- transformers development branch (main), and Studio does
+ transformers development branch (main), and Unsloth does
not install development builds. Support arrives with the
next transformers release on PyPI.
>
diff --git a/studio/frontend/src/hooks/use-tauri-backend.ts b/studio/frontend/src/hooks/use-tauri-backend.ts
index db996768a6..53122864e7 100644
--- a/studio/frontend/src/hooks/use-tauri-backend.ts
+++ b/studio/frontend/src/hooks/use-tauri-backend.ts
@@ -71,8 +71,8 @@ function externalConflictMessage(preflight: DesktopPreflightResult) {
}
return preflight.port
- ? `A Unsloth server for this install is already running from a terminal on port ${preflight.port}. Stop that server, or run \`unsloth studio update\` from that terminal before using the desktop app.`
- : "A Unsloth server for this install is already running from a terminal. Stop that server, or run `unsloth studio update` from that terminal before using the desktop app.";
+ ? `An Unsloth server for this install is already running from a terminal on port ${preflight.port}. Stop that server, or run \`unsloth studio update\` from that terminal before using the desktop app.`
+ : "An Unsloth server for this install is already running from a terminal. Stop that server, or run `unsloth studio update` from that terminal before using the desktop app.";
}
async function waitForManagedServerPort(
diff --git a/studio/frontend/src/i18n/README.md b/studio/frontend/src/i18n/README.md
index ba8cce5d8d..7594e4723f 100644
--- a/studio/frontend/src/i18n/README.md
+++ b/studio/frontend/src/i18n/README.md
@@ -8,5 +8,5 @@
- Preserve interpolation variables exactly, for example `{count}`, `{model}`, and `{provider}`.
- Keep product and technical names unchanged unless there is an established localized name, for example `Unsloth Studio`, `LoRA`, `GGUF`, and `Hugging Face`.
- Keep translation changes small and reviewable. Prefer separate commits for runtime changes, UI migration, and locale text.
-- When adding user-facing Studio UI text, add the English message key first and add non-English overrides only when the translation is clear.
+- When adding user-facing Unsloth UI text, add the English message key first and add non-English overrides only when the translation is clear.
- Run `npx tsx src/i18n/check-parity.ts` before committing to ensure there are no shape mismatches or placeholder discrepancies in the non-English overlays.
diff --git a/studio/frontend/src/i18n/locales/ar.ts b/studio/frontend/src/i18n/locales/ar.ts
index 744a2002c9..4bd4328cce 100644
--- a/studio/frontend/src/i18n/locales/ar.ts
+++ b/studio/frontend/src/i18n/locales/ar.ts
@@ -113,7 +113,7 @@ export const ar = {
showToken: "إظهار التوكن",
tokenSaved: "تم حفظ التوكن",
password: "كلمة المرور",
- passwordDescription: "تغيير كلمة المرور لحساب Studio هذا.",
+ passwordDescription: "تغيير كلمة المرور لحساب Unsloth هذا.",
passwordDialog: {
trigger: "تغيير كلمة المرور",
title: "تغيير كلمة المرور",
@@ -284,7 +284,7 @@ export const ar = {
},
resources: {
title: "النظام",
- description: "مراقبة أجهزة خادم Studio هذا وتخزينه.",
+ description: "مراقبة أجهزة خادم Unsloth هذا وتخزينه.",
liveUpdates: "التحديثات المباشرة",
floatingWindow: "نافذة عائمة",
disableOverlay: "تعطيل التراكب",
diff --git a/studio/frontend/src/i18n/locales/de.ts b/studio/frontend/src/i18n/locales/de.ts
index 7d94e7656e..c28d07790f 100644
--- a/studio/frontend/src/i18n/locales/de.ts
+++ b/studio/frontend/src/i18n/locales/de.ts
@@ -115,7 +115,7 @@ export const de = {
tokenSaved: "Token gespeichert",
password: "Passwort",
passwordDescription:
- "Ändern Sie das Passwort für dieses Studio-Konto.",
+ "Ändern Sie das Passwort für dieses Unsloth-Konto.",
passwordDialog: {
trigger: "Passwort ändern",
title: "Passwort ändern",
@@ -295,7 +295,7 @@ export const de = {
resources: {
title: "System",
description:
- "Überwachen Sie Hardware und Speicher dieses Studio-Servers.",
+ "Überwachen Sie Hardware und Speicher dieses Unsloth-Servers.",
liveUpdates: "Live-Updates",
floatingWindow: "Schwebendes Fenster",
disableOverlay: "Overlay deaktivieren",
diff --git a/studio/frontend/src/i18n/locales/es.ts b/studio/frontend/src/i18n/locales/es.ts
index 988c109a3f..b7dfee10b8 100644
--- a/studio/frontend/src/i18n/locales/es.ts
+++ b/studio/frontend/src/i18n/locales/es.ts
@@ -115,7 +115,7 @@ export const es = {
tokenSaved: "Token guardado",
password: "Contraseña",
passwordDescription:
- "Cambia la contraseña de esta cuenta de Studio.",
+ "Cambia la contraseña de esta cuenta de Unsloth.",
passwordDialog: {
trigger: "Cambiar contraseña",
title: "Cambiar contraseña",
@@ -294,7 +294,7 @@ export const es = {
resources: {
title: "Sistema",
description:
- "Monitorea el hardware y el almacenamiento de este servidor de Studio.",
+ "Monitorea el hardware y el almacenamiento de este servidor de Unsloth.",
liveUpdates: "Actualizaciones en vivo",
floatingWindow: "Ventana flotante",
disableOverlay: "Desactivar superposición",
diff --git a/studio/frontend/src/i18n/locales/fr.ts b/studio/frontend/src/i18n/locales/fr.ts
index e1f2a0c5ec..6105cd8ccf 100644
--- a/studio/frontend/src/i18n/locales/fr.ts
+++ b/studio/frontend/src/i18n/locales/fr.ts
@@ -115,7 +115,7 @@ export const fr = {
tokenSaved: "Token enregistré",
password: "Mot de passe",
passwordDescription:
- "Changez le mot de passe de ce compte Studio.",
+ "Changez le mot de passe de ce compte Unsloth.",
passwordDialog: {
trigger: "Changer le mot de passe",
title: "Changer le mot de passe",
@@ -291,7 +291,7 @@ export const fr = {
resources: {
title: "Système",
description:
- "Surveillez le matériel et le stockage de ce serveur Studio.",
+ "Surveillez le matériel et le stockage de ce serveur Unsloth.",
liveUpdates: "Mises à jour en direct",
floatingWindow: "Fenêtre flottante",
disableOverlay: "Désactiver la superposition",
diff --git a/studio/frontend/src/i18n/locales/hi.ts b/studio/frontend/src/i18n/locales/hi.ts
index 77b6265e7b..732ae1d7fa 100644
--- a/studio/frontend/src/i18n/locales/hi.ts
+++ b/studio/frontend/src/i18n/locales/hi.ts
@@ -113,7 +113,7 @@ export const hi = {
showToken: "token दिखाएं",
tokenSaved: "Token सहेजा गया",
password: "पासवर्ड",
- passwordDescription: "इस Studio खाते के लिए पासवर्ड बदलें।",
+ passwordDescription: "इस Unsloth खाते के लिए पासवर्ड बदलें।",
passwordDialog: {
trigger: "पासवर्ड बदलें",
title: "पासवर्ड बदलें",
@@ -283,7 +283,7 @@ export const hi = {
},
resources: {
title: "सिस्टम",
- description: "इस Studio सर्वर के हार्डवेयर और स्टोरेज की निगरानी करें।",
+ description: "इस Unsloth सर्वर के हार्डवेयर और स्टोरेज की निगरानी करें।",
liveUpdates: "लाइव अपडेट",
floatingWindow: "फ्लोटिंग विंडो",
disableOverlay: "ओवरले अक्षम करें",
diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts
index a261994f03..de5e93c672 100644
--- a/studio/frontend/src/i18n/locales/ja.ts
+++ b/studio/frontend/src/i18n/locales/ja.ts
@@ -360,7 +360,7 @@ export const ja = {
},
resources: {
title: "システム",
- description: "この Studio サーバーのハードウェアとストレージを監視します。",
+ description: "この Unsloth サーバーのハードウェアとストレージを監視します。",
liveUpdates: "リアルタイム更新",
floatingWindow: "フローティングウィンドウ",
disableOverlay: "オーバーレイを無効化",
diff --git a/studio/frontend/src/i18n/locales/ko.ts b/studio/frontend/src/i18n/locales/ko.ts
index 7c5691925e..6ff9cdbb5e 100644
--- a/studio/frontend/src/i18n/locales/ko.ts
+++ b/studio/frontend/src/i18n/locales/ko.ts
@@ -113,7 +113,7 @@ export const ko = {
showToken: "토큰 표시",
tokenSaved: "토큰이 저장되었습니다",
password: "비밀번호",
- passwordDescription: "이 Studio 계정의 비밀번호를 변경합니다.",
+ passwordDescription: "이 Unsloth 계정의 비밀번호를 변경합니다.",
passwordDialog: {
trigger: "비밀번호 변경",
title: "비밀번호 변경",
@@ -282,7 +282,7 @@ export const ko = {
},
resources: {
title: "시스템",
- description: "이 Studio 서버의 하드웨어와 저장소를 모니터링합니다.",
+ description: "이 Unsloth 서버의 하드웨어와 저장소를 모니터링합니다.",
liveUpdates: "실시간 업데이트",
floatingWindow: "플로팅 창",
disableOverlay: "오버레이 비활성화",
diff --git a/studio/frontend/src/i18n/locales/ru.ts b/studio/frontend/src/i18n/locales/ru.ts
index c7464a3b44..60e939bb0d 100644
--- a/studio/frontend/src/i18n/locales/ru.ts
+++ b/studio/frontend/src/i18n/locales/ru.ts
@@ -113,7 +113,7 @@ export const ru = {
showToken: "Показать токен",
tokenSaved: "Токен сохранён",
password: "Пароль",
- passwordDescription: "Изменить пароль для этого аккаунта Studio.",
+ passwordDescription: "Изменить пароль для этого аккаунта Unsloth.",
passwordDialog: {
trigger: "Изменить пароль",
title: "Изменить пароль",
@@ -283,7 +283,7 @@ export const ru = {
},
resources: {
title: "Система",
- description: "Мониторинг оборудования и хранилища этого сервера Studio.",
+ description: "Мониторинг оборудования и хранилища этого сервера Unsloth.",
liveUpdates: "Обновления в реальном времени",
floatingWindow: "Плавающее окно",
disableOverlay: "Отключить оверлей",
diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts
index ff218adad2..37c73086cf 100644
--- a/studio/frontend/src/i18n/locales/zh-CN.ts
+++ b/studio/frontend/src/i18n/locales/zh-CN.ts
@@ -375,7 +375,7 @@ export const zhCN = {
},
resources: {
title: "系统",
- description: "监控此 Studio 服务器的硬件和存储。",
+ description: "监控此 Unsloth 服务器的硬件和存储。",
liveUpdates: "实时更新",
floatingWindow: "悬浮窗口",
disableOverlay: "禁用悬浮层",
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 6d4c21eec8..a43e3d974b 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -2599,7 +2599,7 @@ html[data-chat-font] .aui-root {
* the documented WCAG outcome (motion is "minimised, not removed").
*
* .animate-spin and generated image loading dots are the exceptions: loading
- * indicators across Studio (tool execution loaders, sonner toasts, Tauri
+ * indicators across Unsloth (tool execution loaders, sonner toasts, Tauri
* startup / update screens, the primitive, and image generation
* cards). Freezing them removes the only visual signal that work is in flight,
* so they keep animating.
diff --git a/studio/frontend/src/lib/tauri-diagnostics.ts b/studio/frontend/src/lib/tauri-diagnostics.ts
index 5c07931a86..2b687478fe 100644
--- a/studio/frontend/src/lib/tauri-diagnostics.ts
+++ b/studio/frontend/src/lib/tauri-diagnostics.ts
@@ -62,7 +62,7 @@ export function redactDiagnosticsText(text: string): string {
"$1$2",
);
- // Redact Studio paths before broader home-directory paths.
+ // Redact Unsloth paths before broader home-directory paths.
redacted = redacted.replace(
/(?:\/Users|\/home)\/[^\s/]+\/\.unsloth\/studio/gi,
"",
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index cca8886777..9bbd0cb8be 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -2897,7 +2897,7 @@ def _pick_rocm_gfx_target(out: str) -> str | None:
break
if _vis_raw is not None:
_vis = _vis_raw.strip()
- # Empty or "-1" means "no AMD GPU visible" (matches the rest of Studio).
+ # Empty or "-1" means "no AMD GPU visible" (matches the rest of Unsloth).
if _vis == "" or _vis == "-1":
return None
_first = _vis.split(",")[0].strip()
@@ -4437,7 +4437,7 @@ def ensure_diffusion_visual_server(
approved_checksums: ApprovedReleaseChecksums,
) -> None:
"""Best-effort placement of the DiffusionGemma visual-server binary next to
- llama-server in the install tree, so Studio can serve DiffusionGemma GGUFs
+ llama-server in the install tree, so Unsloth can serve DiffusionGemma GGUFs
without any DG_* env. This is an Unsloth artifact (not a ggml-org one), so it
is optional: if it is already present we just make it executable, otherwise we
try the published release and quietly skip on absence. A source build
@@ -4719,7 +4719,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
# repackage the SO/DLL set (e.g. ggml-org/llama.cpp#23462 split the
# per-binary entry code into paired ``lib-impl.so`` shared
# libraries between b9279 and b9283) without us re-enumerating
- # every new file. Studio invokes llama-server, llama-quantize, and the
+ # every new file. Unsloth invokes llama-server, llama-quantize, and the
# DiffusionGemma visual-server (when the bundle ships it, for native
# DiffusionGemma serving); other CLIs upstream ships (llama-cli,
# llama-bench, ...) are skipped.
@@ -6320,7 +6320,7 @@ def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) ->
from asset names."""
attempts: list[AssetChoice] = []
if host.has_usable_nvidia:
- # Prefer the cudart major Studio loads at runtime (torch's bundled
+ # Prefer the cudart major Unsloth loads at runtime (torch's bundled
# libcudart), not the newest detected on disk. Without this a stray
# cuda13 runtime outranks the torch cuda12 the binary links against.
torch_preference = detect_torch_cuda_runtime_preference(host)
diff --git a/studio/install_node_prebuilt.py b/studio/install_node_prebuilt.py
index 4038216eef..fb40634e95 100644
--- a/studio/install_node_prebuilt.py
+++ b/studio/install_node_prebuilt.py
@@ -6,7 +6,7 @@
Downloads an official Node.js archive from nodejs.org into an isolated
``/node`` and never touches the system Node/npm. Pinning Node 24+
-LTS clears the Studio frontend build floor (Vite 8: Node ^20.19 || >=22.12,
+LTS clears the Unsloth frontend build floor (Vite 8: Node ^20.19 || >=22.12,
npm >= 11) with the npm it bundles.
Archives are verified against sha256 digests pinned in ``node_prebuilt_pins.json``
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 19c492deaa..95c9356d4a 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -1008,7 +1008,7 @@ def _install_bnb_windows_rocm() -> bool:
# `hipinfo.exe` at import time to detect the GPU arch and logs a scary
# (harmless) ERROR + WARNING on every import when it is missing. The venv
# Scripts dir is on PATH only when the venv is activated, which neither
- # Studio nor the installer's child processes ever do.
+ # Unsloth nor the installer's child processes ever do.
_scripts_dir = os.path.dirname(sys.executable)
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")) and not shutil.which(
"hipinfo.exe"
@@ -2318,7 +2318,7 @@ def install_python_stack() -> int:
_progress("dependency overrides (skipped, no torch)")
elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm():
# No working Windows ROCm torchao build: it imports an absent c10d backend
- # and crashes transformers.quantizers. Studio stubs it at runtime, so
+ # and crashes transformers.quantizers. Unsloth stubs it at runtime, so
# installing it only ships a package that crashes on import -- skip it.
_progress("dependency overrides (skipped, Windows ROCm)")
_safe_print(" Windows ROCm -- skipping torchao (no working build; stubbed at runtime)")
@@ -2371,7 +2371,7 @@ def install_python_stack() -> int:
# "https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/main/unsloth/save.py",
# )
- # 8. Studio dependencies
+ # 8. Unsloth dependencies
_progress("studio deps")
pip_install(
"Installing studio dependencies",
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index dab1e1e73f..98e801cd3c 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -322,7 +322,7 @@ function Write-CudaDriverToolkitMismatch {
$driverMajor = $DriverMaxCuda.Split('.')[0]
substep "CUDA Toolkit $ToolkitVersion is a major-version mismatch: toolkit major $toolkitMajor exceeds driver CUDA major $driverMajor ($DriverMaxCuda)." $Color
substep "Update the NVIDIA GPU driver to run CUDA Toolkit $ToolkitVersion, or install a CUDA $driverMajor.x toolkit." $Color
- substep "Or let Studio use the prebuilt CUDA bundle; it does not need the local toolkit." $Color
+ substep "Or let Unsloth use the prebuilt CUDA bundle; it does not need the local toolkit." $Color
}
# Detect CUDA Compute Capability via nvidia-smi.
@@ -937,13 +937,13 @@ function Show-NpmRegistryHint {
Write-Host ""
step "frontend" "registry.npmjs.org looks blocked (corporate firewall/proxy?)" "Yellow"
if ($mirror) {
- substep "Studio pins the public npm registry; your mirror is being ignored."
+ substep "Unsloth pins the public npm registry; your mirror is being ignored."
substep "Detected a registry in your npm config:"
substep " $mirror"
- substep "Re-run pointing Studio at it:"
+ substep "Re-run pointing Unsloth at it:"
substep " `$env:UNSLOTH_NPM_REGISTRY='$mirror'; .\install.ps1 --local"
} else {
- substep "If you use a private mirror/proxy, point Studio at it and re-run:"
+ substep "If you use a private mirror/proxy, point Unsloth at it and re-run:"
substep " `$env:UNSLOTH_NPM_REGISTRY='https://your-mirror.example/api/npm/'; .\install.ps1 --local"
}
substep "(min-release-age and save-exact stay enforced.)"
@@ -1070,7 +1070,7 @@ if (-not $HasNvidiaSmi) {
}
# ── Helper: run amd-smi without triggering a UAC elevation prompt ──
# amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing
-# DiskPart UAC prompt mid-install (Studio backend amd.py hits the same). RunAsInvoker
+# DiskPart UAC prompt mid-install (Unsloth backend amd.py hits the same). RunAsInvoker
# forces it (and helpers it spawns) to run un-elevated; on failure the WMI name ->
# gfx fallback still resolves the arch.
function Invoke-AmdSmiNoElevate {
@@ -1132,7 +1132,7 @@ if (-not $HasNvidiaSmi) {
if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false }
# VenvDir/VIRTUAL_ENV can be unset this early (the update flow probes before
# VenvDir is set), so also derive the venv from the setup python + default
- # Studio home, else the venv hipInfo isn't caught.
+ # Unsloth home, else the venv hipInfo isn't caught.
$venvRoots = @()
if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV }
$vd = Get-Variable -Name VenvDir -ValueOnly -ErrorAction SilentlyContinue
@@ -1141,7 +1141,7 @@ if (-not $HasNvidiaSmi) {
try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {}
}
if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") }
- # A custom Studio home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the
+ # A custom Unsloth home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the
# venv off the default path; seed it too or its hipInfo escapes the filter.
$studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null }
if ($studioHomeEnv) {
@@ -1860,7 +1860,7 @@ $SysNpmVersion = ""
$NodeSource = $null
if (-not $IsPipInstall) {
- # Put Node beside the Studio root. OXC can still need npm when the
+ # Put Node beside the Unsloth root. OXC can still need npm when the
# frontend build is skipped.
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $NodeOverride = $env:UNSLOTH_STUDIO_HOME.Trim() }
elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $NodeOverride = $env:STUDIO_HOME.Trim() }
@@ -2137,17 +2137,17 @@ if ($NeedNodeForSetup) {
step "frontend" "skipped (no suitable Node; system left untouched)" "Yellow"
}
$NeedFrontendBuild = $false
- substep "found Node='$SysNodeVersion' npm='$SysNpmVersion'; Studio needs Node >=20.19/22.12/23 and npm >= 11" "Yellow"
+ substep "found Node='$SysNodeVersion' npm='$SysNpmVersion'; Unsloth needs Node >=20.19/22.12/23 and npm >= 11" "Yellow"
substep "install a suitable Node + npm, or unset UNSLOTH_SKIP_NODE_INSTALL to let Unsloth manage an isolated Node" "Yellow"
} elseif ($NodeSource -eq "bundled") {
New-Item -ItemType Directory -Force -Path $NodeParent -ErrorAction SilentlyContinue | Out-Null
- # Minimal ownership guard for a custom-home dir (the full Studio-owned
+ # Minimal ownership guard for a custom-home dir (the full Unsloth-owned
# helpers are defined later); never os.replace over a user-owned dir.
if ($NodeOverride -and (Test-Path -LiteralPath $NodeDir -PathType Container)) {
$nodeOwnedMarker = Join-Path $NodeDir ".unsloth-studio-owned"
$nodeMeta = Join-Path $NodeDir "UNSLOTH_NODE_PREBUILT_INFO.json"
if (-not (Test-Path -LiteralPath $nodeOwnedMarker) -and -not (Test-Path -LiteralPath $nodeMeta)) {
- Write-Host "[ERROR] $NodeDir already exists and is not a Studio-owned Node install." -ForegroundColor Red
+ Write-Host "[ERROR] $NodeDir already exists and is not an Unsloth-owned Node install." -ForegroundColor Red
Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow
exit 1
}
@@ -2160,7 +2160,7 @@ if ($NeedNodeForSetup) {
$nodeExit = $LASTEXITCODE
if ($nodeExit -eq 3) {
Write-Host $nodeOut -ForegroundColor DarkGray
- step "node" "install blocked by another active Studio install" "Red"
+ step "node" "install blocked by another active Unsloth install" "Red"
exit 3
} elseif ($nodeExit -ne 0) {
Write-Host $nodeOut -ForegroundColor DarkGray
@@ -2486,7 +2486,7 @@ if (Test-Path -LiteralPath $LegacyStudioHome -PathType Container) {
$LegacyStudioHome = (Resolve-Path -LiteralPath $LegacyStudioHome).Path
}
$StudioHomeIsCustom = ($_studioHomeCanon -ne $LegacyStudioHome)
-# Directory-local evidence that Studio created $Path, used to adopt a custom-home
+# Directory-local evidence that Unsloth created $Path, used to adopt a custom-home
# llama.cpp predating the .unsloth-studio-owned marker (see setup.sh). Only the
# prebuilt UNSLOTH_PREBUILT_INFO.json counts; source builds are indistinguishable
# from a user clone on Windows and stay under the strict guard.
@@ -2506,7 +2506,7 @@ function Assert-StudioOwnedOrAbsent {
Mark-StudioOwned $Path
return
}
- Write-Host "[ERROR] $Path already exists and is not marked as a Studio-owned $Label." -ForegroundColor Red
+ Write-Host "[ERROR] $Path already exists and is not marked as an Unsloth-owned $Label." -ForegroundColor Red
Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow
exit 1
}
@@ -2576,7 +2576,7 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
$reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" }
if ($InstallerManagedSetup) {
substep "Stale venv detected ($reason)." "Yellow"
- Write-Host " [ERROR] The existing Studio environment needs repair." -ForegroundColor Red
+ Write-Host " [ERROR] The existing Unsloth environment needs repair." -ForegroundColor Red
Write-Host " Re-run install.ps1 so it can replace the environment safely with rollback." -ForegroundColor Yellow
exit 1
}
@@ -2598,7 +2598,7 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
Remove-Item -LiteralPath $VenvDir -Recurse -Force -ErrorAction Stop
} catch {
Write-Host " [ERROR] Could not remove stale venv: $($_.Exception.Message)" -ForegroundColor Red
- Write-Host " Close any running Studio/Python processes and re-run setup." -ForegroundColor Red
+ Write-Host " Close any running Unsloth/Python processes and re-run setup." -ForegroundColor Red
exit 1
}
}
@@ -2764,7 +2764,7 @@ if ($script:UnslothVerbose) {
# The CUDA tag is chosen based on the driver's max supported CUDA version.
# Triton/inductor filenames are long and can hit Windows MAX_PATH (260). With long
-# paths on, cache under Studio home; else use a short drive-root dir for headroom.
+# paths on, cache under Unsloth home; else use a short drive-root dir for headroom.
if ($LongPathsEnabled) {
$TorchCacheDir = Join-Path $StudioHome "TORCHINDUCTOR_CACHE_DIR"
} else {
@@ -2793,7 +2793,7 @@ $ROCmIndexUrl = $null
# Install AMD ROCm PyTorch wheels when ROCm is confirmed OR a gfx arch is known
# (name-inferred on Adrenalin-only hosts). The per-arch wheels bundle the runtime
# (rocm-sdk-libraries-), so torch.cuda.is_available() is True without a HIP
-# SDK -- which flips Studio out of chat-only (CHAT_ONLY) and enables Train/Export.
+# SDK -- which flips Unsloth out of chat-only (CHAT_ONLY) and enables Train/Export.
# Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed
# ROCm install still falls back to CPU below, so this is safe.
if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
@@ -3231,7 +3231,7 @@ if ($LocalLlamaCppSrc) {
# Reusing a local dir disables both the prebuilt download and the source
# build, so a runnable llama-server.exe must already be present. Accept any
# layout LlamaCppBackend._layout_candidates() resolves (root-level, build\bin,
- # or build\bin\Release) so the flag never rejects a tree Studio could run.
+ # or build\bin\Release) so the flag never rejects a tree Unsloth could run.
$LocalLlamaServerFound = $false
foreach ($_cand in @(
(Join-Path $ResolvedLocal "llama-server.exe"),
@@ -3253,7 +3253,7 @@ if ($LocalLlamaCppSrc) {
}
} else {
# Fail clearly rather than junction an unbuilt or wrong-platform checkout
- # and leave Studio with no usable binary.
+ # and leave Unsloth with no usable binary.
if (-not $LocalLlamaServerFound) {
step "llama.cpp" "no llama-server.exe under $ResolvedLocal (looked for .\llama-server.exe, .\build\bin and .\build\bin\Release) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "Red"
exit 1
@@ -3281,7 +3281,7 @@ if ($LocalLlamaCppSrc) {
# prebuilt path's active-process handling and stop with a clear message.
if (Test-Path -LiteralPath $LlamaCppDir) {
step "llama.cpp" "install blocked by active llama.cpp process" "Yellow"
- substep "Close Studio or other llama.cpp users and retry" "Yellow"
+ substep "Close Unsloth or other llama.cpp users and retry" "Yellow"
exit 3
}
}
@@ -3414,7 +3414,7 @@ if ($LocalLlamaCppLinked) {
if (Test-Path -LiteralPath $LlamaCppDir) {
substep "Existing install was restored" "Yellow"
}
- substep "Close Studio or other llama.cpp users and retry" "Yellow"
+ substep "Close Unsloth or other llama.cpp users and retry" "Yellow"
exit 3
} else {
step "llama.cpp" "prebuilt install failed (continuing)" "Yellow"
@@ -4003,7 +4003,7 @@ if ($LocalLlamaCppLinked) {
}
# -- Step E: Build the DiffusionGemma visual server (optional, best-effort) --
- # An example target present on llama.cpp PR #24423; lets Studio serve
+ # An example target present on llama.cpp PR #24423; lets Unsloth serve
# DiffusionGemma GGUFs without DG_VISUAL_BIN. No-op when not configured.
if ($BuildOk) {
$null = cmake --build $BuildDir --config Release --target llama-diffusion-gemma-visual-server -j $NumCpu 2>&1 | Out-String
diff --git a/studio/setup.sh b/studio/setup.sh
index 3d67db3da7..8d47eecfda 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -127,13 +127,13 @@ _suggest_npm_registry() {
printf '\n' >&2
step "frontend" "registry.npmjs.org looks blocked (corporate firewall/proxy?)" "$C_WARN" >&2
if [ -n "$_mirror" ]; then
- substep "Studio pins the public npm registry; your mirror is being ignored." >&2
+ substep "Unsloth pins the public npm registry; your mirror is being ignored." >&2
substep "Detected a registry in your npm config:" >&2
substep " $_mirror" >&2
- substep "Re-run pointing Studio at it:" >&2
+ substep "Re-run pointing Unsloth at it:" >&2
substep " UNSLOTH_NPM_REGISTRY=$_mirror ./install.sh --local" >&2
else
- substep "If you use a private mirror/proxy, point Studio at it and re-run:" >&2
+ substep "If you use a private mirror/proxy, point Unsloth at it and re-run:" >&2
substep " UNSLOTH_NPM_REGISTRY=https://your-mirror.example/api/npm/ ./install.sh --local" >&2
fi
substep "(min-release-age and save-exact stay enforced.)" >&2
@@ -396,7 +396,7 @@ _print_cuda_driver_toolkit_mismatch() {
local _driver_major=${_driver_version%%.*}
substep "CUDA Toolkit $_toolkit_version is a major-version mismatch: toolkit major $_toolkit_major exceeds driver CUDA major $_driver_major ($_driver_version)." "$C_WARN"
substep "Update the NVIDIA GPU driver to run CUDA Toolkit $_toolkit_version, or install a CUDA $_driver_major.x toolkit." "$C_WARN"
- substep "Or let Studio use the prebuilt CUDA bundle; it does not need the local toolkit." "$C_WARN"
+ substep "Or let Unsloth use the prebuilt CUDA bundle; it does not need the local toolkit." "$C_WARN"
}
print_llama_error_log() {
@@ -531,7 +531,7 @@ _STUDIO_HOME_IS_CUSTOM=false
if [ "$_studio_home_canon" != "$_LEGACY_STUDIO_HOME" ]; then
_STUDIO_HOME_IS_CUSTOM=true
fi
-# Directory-local evidence Studio created "$1": only prebuilt-installer metadata
+# Directory-local evidence Unsloth created "$1": only prebuilt-installer metadata
# counts (UNSLOTH_PREBUILT_INFO.json for llama.cpp, UNSLOTH_NODE_PREBUILT_INFO.json
# for Node), both written only by our installers. Mirrors the setup.ps1 Node guard.
# A markerless source build stays strict since this runs right before an rm -rf.
@@ -549,7 +549,7 @@ _assert_studio_owned_or_absent() {
: > "$_aso_dir/$_STUDIO_OWNED_MARKER" 2>/dev/null || true
return 0
fi
- echo "ERROR: $_aso_dir already exists and is not marked as a Studio-owned $_aso_label." >&2
+ echo "ERROR: $_aso_dir already exists and is not marked as an Unsloth-owned $_aso_label." >&2
echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2
exit 1
fi
@@ -585,7 +585,7 @@ if [ "$_NEED_FRONTEND_BUILD" = false ] && [ ! -d "$_OXC_DIR" ]; then
else
# ── Node (isolated; never touches the system Node/npm) ──
-# Studio's frontend (Vite 8) needs Node ^20.19 || >=22.12 || >=23 and npm >= 11.
+# Unsloth's frontend (Vite 8) needs Node ^20.19 || >=22.12 || >=23 and npm >= 11.
# Three sources:
# system -- system Node + npm already satisfy both; used read-only.
# bundled -- install a pinned isolated Node under $UNSLOTH_HOME/node, build-only.
@@ -666,9 +666,9 @@ elif [ "$NODE_SOURCE" = bundled ]; then
fi
set -e
if [ "$_NODE_STATUS" -eq 3 ]; then
- step "node" "install blocked by another active Studio install" "$C_ERR"
+ step "node" "install blocked by another active Unsloth install" "$C_ERR"
sed 's/^/ | /' "$_NODE_LOG" >&2; rm -f "$_NODE_LOG"
- substep "close other Studio installs and retry"
+ substep "close other Unsloth installs and retry"
exit 3
elif [ "$_NODE_STATUS" -ne 0 ]; then
step "node" "isolated Node install failed" "$C_ERR"
@@ -692,7 +692,7 @@ elif [ "$NODE_SOURCE" = bundled ]; then
else
_FRONTEND_SKIP=true
step "frontend" "skipped (no suitable Node; system left untouched)" "$C_WARN"
- substep "found Node='${_SYS_NODE_VER:-none}' npm='${_SYS_NPM_VER:-none}'; Studio needs Node >=20.19/22.12/23 and npm >= 11"
+ substep "found Node='${_SYS_NODE_VER:-none}' npm='${_SYS_NPM_VER:-none}'; Unsloth needs Node >=20.19/22.12/23 and npm >= 11"
substep "install a suitable Node + npm, or unset UNSLOTH_SKIP_NODE_INSTALL to let Unsloth manage an isolated Node"
fi
verbose_substep "node source: $NODE_SOURCE (sys node=${_SYS_NODE_VER:-none} npm=${_SYS_NPM_VER:-none}) dir=$NODE_DIR"
@@ -873,11 +873,11 @@ _remove_agent_instruction_files \
_COLAB_NO_VENV=false
if [ ! -x "$VENV_DIR/bin/python" ]; then
if [ "$IS_COLAB" = true ]; then
- # On Colab there is no Studio venv -- install backend deps into system Python.
+ # On Colab there is no Unsloth venv -- install backend deps into system Python.
# Strip all version constraints so pip keeps Colab's pre-installed
# packages (huggingface-hub, datasets, transformers) and only pulls
# in genuinely missing ones (structlog, fastapi, etc.).
- substep "Colab detected, installing Studio backend dependencies..."
+ substep "Colab detected, installing Unsloth backend dependencies..."
_COLAB_REQS_TMP="$(mktemp)"
sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/studio.txt" \
| grep -v '^#' | grep -v '^$' > "$_COLAB_REQS_TMP"
@@ -1254,7 +1254,7 @@ _link_local_llama_quantize_shim() {
}
# Accept any layout LlamaCppBackend._layout_candidates() resolves so the flag
-# never rejects a tree Studio could actually run: a root-level llama-server (a
+# never rejects a tree Unsloth could actually run: a root-level llama-server (a
# `make` build or a flat-extracted release) or the CMake build/bin/llama-server.
_has_local_llama_server() {
[ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ]
@@ -1298,13 +1298,13 @@ if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then
# Reusing disables BOTH the prebuilt download and the source build, so the
# linked tree must already contain a runnable llama-server in one of the
# layouts the backend resolves (root-level or build/bin/). Fail clearly
- # rather than link an unbuilt or wrong-platform checkout and leave Studio
+ # rather than link an unbuilt or wrong-platform checkout and leave Unsloth
# with no usable binary.
if ! _has_local_llama_server "$_RESOLVED_LOCAL"; then
step "llama.cpp" "no llama-server under $_RESOLVED_LOCAL (looked for ./llama-server and ./build/bin/llama-server) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "$C_ERR"
exit 1
fi
- # A stale link from a previous --with-llama-cpp-dir run isn't Studio-owned
+ # A stale link from a previous --with-llama-cpp-dir run isn't Unsloth-owned
# content; drop it before the ownership check so re-runs stay idempotent
# for a custom UNSLOTH_STUDIO_HOME (the assert would otherwise follow the
# link into the user's dir and reject it as unowned).
@@ -1389,7 +1389,7 @@ else
if [ -d "$LLAMA_CPP_DIR" ]; then
substep "existing install was restored"
fi
- substep "close Studio or other llama.cpp users and retry"
+ substep "close Unsloth or other llama.cpp users and retry"
exit 3
else
step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
@@ -1886,7 +1886,7 @@ else
ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize"
fi
# DiffusionGemma visual server, if it was built (PR #24423): link next to
- # llama-server so Studio serves DiffusionGemma GGUFs without DG_VISUAL_BIN.
+ # llama-server so Unsloth serves DiffusionGemma GGUFs without DG_VISUAL_BIN.
if [ -f "$LLAMA_CPP_DIR/build/bin/llama-diffusion-gemma-visual-server" ]; then
ln -sf build/bin/llama-diffusion-gemma-visual-server "$LLAMA_CPP_DIR/llama-diffusion-gemma-visual-server"
fi
@@ -1983,7 +1983,7 @@ echo ""
# When called from install.sh (SKIP_STUDIO_BASE=1), exit non-zero so the
# installer can report the GGUF failure after finishing PATH/shortcut setup.
# When called directly via 'unsloth studio update', keep the install
-# successful -- the footer above already reports the limitation and Studio
+# successful -- the footer above already reports the limitation and Unsloth
# is still usable for non-GGUF workflows.
if [ "$_LLAMA_CPP_DEGRADED" = true ] && [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then
exit 1
diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs
index 6bc2116786..5d41cb9217 100644
--- a/studio/src-tauri/src/commands.rs
+++ b/studio/src-tauri/src/commands.rs
@@ -25,12 +25,12 @@ fn should_emit_repair_failed(msg: &str) -> bool {
fn external_conflict_message(conflict: &crate::preflight::ExternalBackendConflict) -> String {
if conflict.reason == "desktop_owned_backend_active" {
return format!(
- "A desktop-owned Studio server for this install is already running on port {}. Quit the other desktop app instance, then try again.",
+ "A desktop-owned Unsloth server for this install is already running on port {}. Quit the other desktop app instance, then try again.",
conflict.port
);
}
format!(
- "A Studio server for this install is already running from a terminal on port {}. Stop that server, or run `unsloth studio update` from that terminal before using desktop repair/update.",
+ "An Unsloth server for this install is already running from a terminal on port {}. Stop that server, or run `unsloth studio update` from that terminal before using desktop repair/update.",
conflict.port
)
}
@@ -475,7 +475,7 @@ pub async fn start_backend_update(
.map_err(|e| format!("Update task panicked: {e}"))?
}
-/// Repair a stale managed Studio install.
+/// Repair a stale managed Unsloth install.
#[tauri::command]
pub async fn start_managed_repair(
app: AppHandle,
@@ -522,7 +522,7 @@ pub async fn start_managed_repair(
let repair_group_id = install::take_pending_repair_group_for_resume(&install_state)
.unwrap_or_else(|| diagnostics::begin_repair_group(&diagnostics_state));
- let _ = app.emit("repair-progress", "Updating existing Studio install...");
+ let _ = app.emit("repair-progress", "Updating existing Unsloth install...");
let update_app = app.clone();
let update_state = update_state.inner().clone();
let update_diagnostics = diagnostics_state.clone();
@@ -549,7 +549,7 @@ pub async fn start_managed_repair(
warn!("Managed repair update finished, but preflight is still not ready; falling back to installer");
let _ = app.emit(
"repair-progress",
- "Update finished, but Studio is still not ready. Running bundled installer...",
+ "Update finished, but Unsloth is still not ready. Running bundled installer...",
);
}
Err(msg) => {
@@ -627,7 +627,7 @@ pub async fn start_managed_repair(
return Ok(());
}
- let msg = "Repair finished, but Studio install is still not desktop-ready.".to_string();
+ let msg = "Repair finished, but Unsloth install is still not desktop-ready.".to_string();
error!("{}", msg);
diagnostics::finish_repair_group(
&diagnostics_state,
diff --git a/studio/src-tauri/src/desktop_auth.rs b/studio/src-tauri/src/desktop_auth.rs
index 33605f6c65..db65b8796f 100644
--- a/studio/src-tauri/src/desktop_auth.rs
+++ b/studio/src-tauri/src/desktop_auth.rs
@@ -201,7 +201,7 @@ async fn exchange_desktop_secret(
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Err(AuthError::StaleResponder(
- "Running Studio backend is too old for this desktop app. Update that backend and restart."
+ "Running Unsloth backend is too old for this desktop app. Update that backend and restart."
.to_string(),
));
}
@@ -364,7 +364,7 @@ async fn desktop_auth_inner(
}
Err(
- "Desktop auth failed. Update or repair the managed Studio install, then restart Studio."
+ "Desktop auth failed. Update or repair the managed Unsloth install, then restart Unsloth."
.to_string(),
)
}
@@ -465,7 +465,7 @@ mod tests {
.message();
assert_eq!(
error,
- "Running Studio backend is too old for this desktop app. Update that backend and restart."
+ "Running Unsloth backend is too old for this desktop app. Update that backend and restart."
);
}
}
diff --git a/studio/src-tauri/src/main.rs b/studio/src-tauri/src/main.rs
index 4ed12051ed..4724317601 100644
--- a/studio/src-tauri/src/main.rs
+++ b/studio/src-tauri/src/main.rs
@@ -107,7 +107,7 @@ fn cleanup_child_processes(app: &tauri::AppHandle) {
}
fn setup_tray(app: &tauri::App) -> Result<(), Box> {
- let open = MenuItemBuilder::with_id("open", "Open Studio").build(app)?;
+ let open = MenuItemBuilder::with_id("open", "Open Unsloth").build(app)?;
let toggle = MenuItemBuilder::with_id("toggle", "Start/Stop Server").build(app)?;
let quit = MenuItemBuilder::with_id("quit", "Quit").build(app)?;
let menu = MenuBuilder::new(app)
diff --git a/studio/src-tauri/src/native_path_policy.rs b/studio/src-tauri/src/native_path_policy.rs
index b2ebb34621..b82e516e7a 100644
--- a/studio/src-tauri/src/native_path_policy.rs
+++ b/studio/src-tauri/src/native_path_policy.rs
@@ -191,7 +191,7 @@ fn reject_sensitive_artifact(path: &Path) -> Result<(), String> {
"\\pid",
] {
if lowered.contains(needle) {
- return Err("Sensitive Studio state cannot be registered as an artifact.".to_string());
+ return Err("Sensitive Unsloth state cannot be registered as an artifact.".to_string());
}
}
if let Some(ext) = path.extension().and_then(|ext| ext.to_str()) {
diff --git a/tests/python/test_e2e_no_torch_sandbox.py b/tests/python/test_e2e_no_torch_sandbox.py
index f412de2063..bb61af462d 100644
--- a/tests/python/test_e2e_no_torch_sandbox.py
+++ b/tests/python/test_e2e_no_torch_sandbox.py
@@ -30,7 +30,7 @@ VLM_PROCESSING = DATASETS_DIR / "vlm_processing.py"
ITERABLE = DATASETS_DIR / "iterable.py"
HARDWARE_PY = HARDWARE_DIR / "hardware.py"
-# Studio venv for server tests
+# Unsloth venv for server tests
STUDIO_VENV = Path.home() / ".unsloth" / "studio" / "unsloth_studio"
sys.path.insert(0, str(STUDIO_DIR))
@@ -957,20 +957,20 @@ server = pytest.mark.server
@server
class TestLiveServerStartup:
- """Live server startup against the existing Studio venv with torch made unimportable (pytest -m server)."""
+ """Live server startup against the existing Unsloth venv with torch made unimportable (pytest -m server)."""
@pytest.fixture(autouse = True)
def _check_studio_venv(self):
py = _studio_venv_python()
if py is None:
- pytest.skip("Studio venv not found at ~/.unsloth/studio/unsloth_studio")
+ pytest.skip("Unsloth venv not found at ~/.unsloth/studio/unsloth_studio")
@pytest.fixture(scope = "class")
def server_process(self):
"""Start the studio backend server without torch, yield (proc, port), then stop."""
py = _studio_venv_python()
if py is None:
- pytest.skip("Studio venv not found")
+ pytest.skip("Unsloth venv not found")
port = _server_port()
backend_dir = BACKEND_DIR
diff --git a/tests/python/test_studio_import_no_torch.py b/tests/python/test_studio_import_no_torch.py
index c4efbc8cea..f551519de9 100644
--- a/tests/python/test_studio_import_no_torch.py
+++ b/tests/python/test_studio_import_no_torch.py
@@ -1,4 +1,4 @@
-"""Sandbox tests: Studio dataset modules load/run in isolated no-torch venvs."""
+"""Sandbox tests: Unsloth dataset modules load/run in isolated no-torch venvs."""
from __future__ import annotations
diff --git a/tests/saving/test_prewarm_base_model_hub_cache.py b/tests/saving/test_prewarm_base_model_hub_cache.py
index cbb52863ba..4269f9d61c 100644
--- a/tests/saving/test_prewarm_base_model_hub_cache.py
+++ b/tests/saving/test_prewarm_base_model_hub_cache.py
@@ -4,7 +4,7 @@
"""Regression tests for #6890: repeated base-model downloads across checkpoint exports.
merge_and_overwrite_lora downloads missing 16-bit shards with hf_hub_download(local_dir),
-which never populates the persistent HF hub cache; a temporary merge directory (Studio
+which never populates the persistent HF hub cache; a temporary merge directory (Unsloth
GGUF exports delete it) means every checkpoint export re-downloads the full base model.
_prewarm_base_model_hub_cache snapshot-downloads the base into the hub cache first so
the zoo's cache-copy fast path is hit on later exports.
diff --git a/tests/studio/_playwright_robust.py b/tests/studio/_playwright_robust.py
index a4590066d4..b2f6df751e 100644
--- a/tests/studio/_playwright_robust.py
+++ b/tests/studio/_playwright_robust.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
-"""Shared CI-runner workarounds for the Studio Playwright tests (Chromium flags,
+"""Shared CI-runner workarounds for the Unsloth Playwright tests (Chromium flags,
view-transition killer, page recovery, post-action response wait). Imported
directly by the standalone scripts; does NOT depend on pytest.
"""
@@ -130,7 +130,7 @@ def wait_for_health(
timeout = 3.0,
)
last_status, last_body = status, body
- # Accept any 200 -- different Studio builds report status differently.
+ # Accept any 200 -- different Unsloth builds report status differently.
if status == 200:
if info is not None:
info(f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}")
diff --git a/tests/studio/install/smoke_test_parallel_studio_home.py b/tests/studio/install/smoke_test_parallel_studio_home.py
index 9d840355e1..55b6df5190 100644
--- a/tests/studio/install/smoke_test_parallel_studio_home.py
+++ b/tests/studio/install/smoke_test_parallel_studio_home.py
@@ -80,7 +80,7 @@ def _launch_backend(
env = os.environ.copy()
env["HOME"] = str(fake_home)
# Pin UNSLOTH_STUDIO_HOME and clear the alias so the child can't inherit a
- # Studio root from the caller's shell and resolve to the wrong install.
+ # Unsloth root from the caller's shell and resolve to the wrong install.
env["UNSLOTH_STUDIO_HOME"] = str(studio_home)
env.pop("STUDIO_HOME", None)
# Popen dups stdout into the child, so closing the parent's handle here is safe.
diff --git a/tests/studio/install/test_launch_studio_launcher.py b/tests/studio/install/test_launch_studio_launcher.py
index a7396aaf5d..1c12024f6d 100644
--- a/tests/studio/install/test_launch_studio_launcher.py
+++ b/tests/studio/install/test_launch_studio_launcher.py
@@ -1,4 +1,4 @@
-"""Guard install.ps1's Studio launcher against the AV-heuristic shape (Kaspersky
+"""Guard install.ps1's Unsloth launcher against the AV-heuristic shape (Kaspersky
HEUR:Trojan.VBS.Agent.gen): a WScript .vbs spawning a hidden ExecutionPolicy-Bypass PowerShell.
The shortcut must stay windowless via powershell.exe -WindowStyle Hidden over launch-studio.ps1,
never a .vbs/WScript.Shell.Run wrapper, and any pre-existing .vbs must be deleted on upgrade."""
diff --git a/tests/studio/install/test_managed_node_runtime.py b/tests/studio/install/test_managed_node_runtime.py
index 0dbc6788ca..17c7e3e60f 100644
--- a/tests/studio/install/test_managed_node_runtime.py
+++ b/tests/studio/install/test_managed_node_runtime.py
@@ -3,7 +3,7 @@
"""Tests for the runtime managed-Node resolver (studio/backend/utils/node_runtime.py).
-The Studio frontend installer may provision an isolated Node under
+The Unsloth frontend installer may provision an isolated Node under
``/node`` that is never added to the user's PATH. The backend OXC
validator must still find a usable Node at runtime: a version-adequate system
Node, else the managed isolated one. These tests pin that resolution and the
diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py
index ac6a96167a..d6dc8b2f8e 100644
--- a/tests/studio/install/test_pr5940_followups.py
+++ b/tests/studio/install/test_pr5940_followups.py
@@ -420,7 +420,7 @@ def test_ps_installers_gate_amd_smi_on_windows():
assert (
"UNSLOTH_SETUP_PYTHON" in text
), f"{ps.name} venv-internal check must seed the venv root from UNSLOTH_SETUP_PYTHON"
- # A custom Studio home moves the venv off the default path; it must be
+ # A custom Unsloth home moves the venv off the default path; it must be
# seeded too or its hipInfo escapes the filter and reopens the gate.
assert (
"UNSLOTH_STUDIO_HOME" in text
@@ -429,7 +429,7 @@ def test_ps_installers_gate_amd_smi_on_windows():
@pytest.mark.parametrize("ps", [_INSTALL_PS1, _SETUP_PS1], ids = ["install.ps1", "setup.ps1"])
def test_ps_venv_probe_expands_tilde_for_custom_studio_home(ps):
- # The probe seeds the venv root from a custom Studio home; a ~\studio form
+ # The probe seeds the venv root from a custom Unsloth home; a ~\studio form
# must expand to USERPROFILE like the canonical resolver, else GetFullPath
# keeps the literal ~ (cwd-relative) and the hipInfo escapes the filter.
text = ps.read_text(encoding = "utf-8")
@@ -439,7 +439,7 @@ def test_ps_venv_probe_expands_tilde_for_custom_studio_home(ps):
block = text[i:j]
assert "USERPROFILE" in block and ".Substring(1)" in block, (
f"{ps.name}: the venv-internal probe must expand a leading ~ in the custom "
- "Studio home before seeding the venv root (mirroring the canonical resolver)"
+ "Unsloth home before seeding the venv root (mirroring the canonical resolver)"
)
# The ~ expansion must be guarded on a non-empty USERPROFILE; otherwise
# Join-Path $env:USERPROFILE throws on a service/SYSTEM account with no profile,
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index 5cabf41f57..e7ac0ec82d 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -360,7 +360,7 @@ class TestRuntimePatterns:
install_kind = "windows-hip",
)
patterns = runtime_patterns_for_choice(choice)
- # Narrowed from "*.exe" to the two binaries Studio actually invokes.
+ # Narrowed from "*.exe" to the two binaries Unsloth actually invokes.
assert "llama-server.exe" in patterns
assert "llama-quantize.exe" in patterns
assert "*.dll" in patterns
@@ -378,7 +378,7 @@ class TestRuntimePatterns:
assert "lib*.dylib" in patterns
def test_diffusion_visual_server_kept(self):
- # The DiffusionGemma visual-server must survive the prune so Studio can
+ # The DiffusionGemma visual-server must survive the prune so Unsloth can
# serve DiffusionGemma GGUFs natively.
for kind, name in (
("linux-cuda", "llama-diffusion-gemma-visual-server"),
diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py
index 7575dea581..e372a3bbeb 100644
--- a/tests/studio/install/test_selection_logic.py
+++ b/tests/studio/install/test_selection_logic.py
@@ -270,13 +270,13 @@ def mock_windows_runtime(monkeypatch, lines):
# ===========================================================================
-# Studio run.py localhost warning
+# Unsloth run.py localhost warning
# ===========================================================================
class TestStudioLocalhostIpv6Warning:
def _prepare_loopback(self, run_module, monkeypatch):
- # Studio confirmed answering on the IPv4 loopback.
+ # Unsloth confirmed answering on the IPv4 loopback.
monkeypatch.setattr(
run_module,
"_working_local_url",
@@ -327,7 +327,7 @@ class TestStudioLocalhostIpv6Warning:
assert "http://localhost:8888" in captured.out
def test_ipv6_listener_does_not_suppress_warning(self, monkeypatch):
- # A process on ::1 is NOT Studio (binds 127.0.0.1 only), so the warning must
+ # A process on ::1 is NOT Unsloth (binds 127.0.0.1 only), so the warning must
# still fire -- that is exactly when http://localhost opens the wrong service.
run_module = load_studio_run_module(monkeypatch)
self._prepare_loopback(run_module, monkeypatch)
@@ -356,7 +356,7 @@ class TestStudioLocalhostIpv6Warning:
assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", port) is None
def test_ipv4_not_answering_suppresses_warning(self, monkeypatch):
- # Studio not confirmed on 127.0.0.1 -> no warning.
+ # Unsloth not confirmed on 127.0.0.1 -> no warning.
run_module = load_studio_run_module(monkeypatch)
monkeypatch.setattr(run_module, "_working_local_url", lambda port: None)
self._set_getaddrinfo(monkeypatch, [self._ipv6()])
@@ -3668,7 +3668,7 @@ class TestCpuFallback:
# ===========================================================================
-@pytest.mark.skipif(sys.platform == "win32", reason = "bash-only Studio installer tests")
+@pytest.mark.skipif(sys.platform == "win32", reason = "bash-only Unsloth installer tests")
class TestCudaDriverToolkitMismatchMessage:
_SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
_SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
@@ -3873,7 +3873,7 @@ class TestCudaDriverToolkitMismatchMessage:
"or install a CUDA $driverMajor.x toolkit." in source
)
assert (
- "Or let Studio use the prebuilt CUDA bundle; it does not need the local toolkit."
+ "Or let Unsloth use the prebuilt CUDA bundle; it does not need the local toolkit."
) in source
assert (
"Write-CudaDriverToolkitMismatch -ToolkitVersion $IncompatibleToolkit "
diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py
index 9c01e95fd4..5bebf0a9e8 100644
--- a/tests/studio/playwright_chat_ime_i18n.py
+++ b/tests/studio/playwright_chat_ime_i18n.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
-"""Studio chat composer IME + multilingual regression smoke.
+"""Unsloth chat composer IME + multilingual regression smoke.
Covers: stuck IME composition (#5318 / PR #5327), multilingual paste round-trip,
stuck compositionend (#5546), and Mac input-method switch recovery (keydown/blur).
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index 065ba7a745..35b18756ff 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.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
-"""Comprehensive Studio chat UI test, run locally + in CI."""
+"""Comprehensive Unsloth chat UI test, run locally + in CI."""
import json
import os
@@ -156,7 +156,7 @@ with sync_playwright() as p:
# pointer events and break Playwright's actionability check.
reduced_motion = "reduce",
)
- # Hard-disable CSS view-transitions: Studio's theme toggle + sidebar
+ # Hard-disable CSS view-transitions: Unsloth's theme toggle + sidebar
# collapse run startViewTransition() which can leave intercepting
# pointer events for a beat after each route swap. See _playwright_robust.py.
install_view_transition_killer(ctx)
@@ -477,7 +477,7 @@ with sync_playwright() as p:
fail(f"/api/inference/load returned {load_resp['status']}: {load_resp.get('body')!r}")
info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}")
- # Studio caches model state in zustand; reload so the composer picks
+ # Unsloth caches model state in zustand; reload so the composer picks
# up the loaded model.
page.reload()
composer = page.locator('textarea[aria-label="Message input"]')
@@ -493,7 +493,7 @@ with sync_playwright() as p:
# (app-sidebar.tsx) -- as stable as anything in the codebase.
picker_btn = page.locator('[data-tour="chat-model-selector"]').first
if picker_btn.count() == 0:
- # Fall back to text-based locators for older Studio builds.
+ # Fall back to text-based locators for older Unsloth builds.
picker_btn = page.locator(
'button:has-text("gemma-3-270m"), '
'button:has-text("Gemma 3"), '
@@ -893,7 +893,7 @@ with sync_playwright() as p:
if len(observed) < 3:
soft_fail(f"theme toggle ran only {len(observed)} cycle(s), expected 3")
# Don't strict-fail on both polarities: the runner's
- # prefers-color-scheme + Studio's "system" default can collapse
+ # prefers-color-scheme + Unsloth's "system" default can collapse
# to one polarity even when .dark toggles correctly. The 3-cycle
# completion above is the real invariant.
if light_seen and dark_seen:
diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py
index 209a8a06f1..dde6c5d635 100644
--- a/tests/studio/playwright_extra_ui.py
+++ b/tests/studio/playwright_extra_ui.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
-"""Studio extra-UI Playwright test: Compare tab, Recipes editor, /export, /studio, Settings tabs."""
+"""Unsloth extra-UI Playwright test: Compare tab, Recipes editor, /export, /studio, Settings tabs."""
import json
import os
@@ -90,11 +90,11 @@ with sync_playwright() as p:
)
install_view_transition_killer(ctx)
page = ctx.new_page()
- # 60s default for slow macos-14 --single-process Chromium (second Studio boot of the job).
+ # 60s default for slow macos-14 --single-process Chromium (second Unsloth boot of the job).
page.set_default_timeout(60_000)
page_errors = []
- # Filter out known-benign React errors (timing artefacts on slow CI runners, not Studio bugs);
+ # Filter out known-benign React errors (timing artefacts on slow CI runners, not Unsloth bugs);
# shared base list lives in _playwright_robust.BENIGN_PAGE_ERROR_PATTERNS.
def _on_pageerror(e):
msg = str(e)
@@ -451,9 +451,9 @@ with sync_playwright() as p:
)
# ─────────────────────────────────────────────────────
- # 4. Studio training route.
+ # 4. Unsloth training route.
# ─────────────────────────────────────────────────────
- step(f"Studio route ({'chat-only redirect' if chat_only else 'tabs + sections'})")
+ step(f"Unsloth route ({'chat-only redirect' if chat_only else 'tabs + sections'})")
page.goto(f"{BASE}/studio")
page.wait_for_timeout(1500)
shoot("08-studio")
diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py
index 275fe7ac57..63bc0dbba9 100644
--- a/tests/studio/run_real_mlx_smoke.py
+++ b/tests/studio/run_real_mlx_smoke.py
@@ -106,7 +106,7 @@ def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, flo
import mlx.nn as nn
from mlx.utils import tree_flatten
- # Match Studio's text dataset path: no EOS appended behind the user's back.
+ # Match Unsloth's text dataset path: no EOS appended behind the user's back.
ids = list(tokenizer.encode(text))
if len(ids) < 2:
raise RuntimeError(f"text too short to compute loss: {len(ids)} tokens")
diff --git a/tests/studio/studio_api_smoke.py b/tests/studio/studio_api_smoke.py
index 845a9ed021..d30bd11dca 100644
--- a/tests/studio/studio_api_smoke.py
+++ b/tests/studio/studio_api_smoke.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
-"""End-to-end Studio API & Auth HTTP integration tests against an externally-booted Studio."""
+"""End-to-end Unsloth API & Auth HTTP integration tests against an externally-booted Unsloth."""
import json
import os
@@ -571,7 +571,7 @@ EXPECTED_AUTH_ENDPOINTS = [
for method, path in EXPECTED_AUTH_ENDPOINTS:
if (method, path) in PUBLIC:
continue
- # Don't actually shut Studio down: an unauthenticated call must 401/403 before the trigger fires.
+ # Don't actually shut Unsloth down: an unauthenticated call must 401/403 before the trigger fires.
if path == "/api/shutdown":
code, _ = http(method, path)
if code in (401, 403):
@@ -610,6 +610,6 @@ if _failed:
sys.exit(1)
_emit(
"",
- "PASS all Studio API & Auth assertions"
+ "PASS all Unsloth API & Auth assertions"
+ (f" ({len(_warned)} audit findings logged)" if _warned else ""),
)
diff --git a/tests/studio/test_auth_form_input_count.py b/tests/studio/test_auth_form_input_count.py
index 4d5d72d20b..75e6cfd1fb 100644
--- a/tests/studio/test_auth_form_input_count.py
+++ b/tests/studio/test_auth_form_input_count.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-"""Fast source and runtime contracts for Studio's frontend authentication flows.
+"""Fast source and runtime contracts for Unsloth's frontend authentication flows.
PR #5490 added a third "Current password" input, regressing first-boot UX to
three inputs; PR #5545 restores two by rendering it only when BOOTSTRAP is absent.
diff --git a/tests/studio/test_chat_title_generation.py b/tests/studio/test_chat_title_generation.py
index b568a51400..6a47cfbce4 100644
--- a/tests/studio/test_chat_title_generation.py
+++ b/tests/studio/test_chat_title_generation.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
-"""Regression checks for Studio chat title generation context."""
+"""Regression checks for Unsloth chat title generation context."""
from __future__ import annotations
diff --git a/tests/studio/test_cli_studio_stop_windows.py b/tests/studio/test_cli_studio_stop_windows.py
index 778679c73b..2267d7feda 100644
--- a/tests/studio/test_cli_studio_stop_windows.py
+++ b/tests/studio/test_cli_studio_stop_windows.py
@@ -5,7 +5,7 @@
`stop` once used `os.kill(pid, 0)`, which raises WinError 87 on Windows before
reaching taskkill; the fix adds cross-platform `_pid_alive` (tasklist on Windows,
-signal-0 elsewhere). AST + mock-only; no real processes, no Studio deps imported.
+signal-0 elsewhere). AST + mock-only; no real processes, no Unsloth deps imported.
"""
import ast
diff --git a/tests/studio/test_hardware_dispatch_matrix.py b/tests/studio/test_hardware_dispatch_matrix.py
index bccddac967..19619bc318 100644
--- a/tests/studio/test_hardware_dispatch_matrix.py
+++ b/tests/studio/test_hardware_dispatch_matrix.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: AGPL-3.0-only
-"""Studio hardware dispatch matrix: spoofs platform/torch/mlx per PROFILES to exercise CUDA/ROCm/XPU/MLX/CPU paths without real hardware."""
+"""Unsloth hardware dispatch matrix: spoofs platform/torch/mlx per PROFILES to exercise CUDA/ROCm/XPU/MLX/CPU paths without real hardware."""
from __future__ import annotations
@@ -31,9 +31,9 @@ class HardwareProfile:
mps_available: bool # torch.backends.mps.is_available() value
expect_is_mlx: bool # unsloth._IS_MLX
- expect_device_type: str # Studio DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU")
- expect_is_rocm: bool # Studio IS_ROCM
- expect_apple_silicon: bool # Studio is_apple_silicon()
+ expect_device_type: str # Unsloth DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU")
+ expect_is_rocm: bool # Unsloth IS_ROCM
+ expect_apple_silicon: bool # Unsloth is_apple_silicon()
extra_notes: str = ""
@@ -66,7 +66,7 @@ PROFILES = [
expect_is_rocm = True,
expect_apple_silicon = False,
extra_notes = "PyTorch ROCm reuses torch.cuda.* over HIP; "
- "Studio still uses DeviceType.CUDA but flips IS_ROCM=True.",
+ "Unsloth still uses DeviceType.CUDA but flips IS_ROCM=True.",
),
HardwareProfile(
name = "intel_xpu",
@@ -154,7 +154,7 @@ def spoof_hardware(monkeypatch):
import platform
import torch
- # platform spoof (used by both the unsloth gate and Studio's helpers)
+ # platform spoof (used by both the unsloth gate and Unsloth's helpers)
monkeypatch.setattr(platform, "system", lambda: profile.system)
monkeypatch.setattr(platform, "machine", lambda: profile.machine)
@@ -227,7 +227,7 @@ def spoof_hardware(monkeypatch):
monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)
- # Studio's _has_mlx() does `import mlx.core`, not find_spec; block it
+ # Unsloth's _has_mlx() does `import mlx.core`, not find_spec; block it
# with a meta_path finder that raises ImportError for mlx.*.
class _BlockMLXFinder:
def find_spec(
@@ -266,7 +266,7 @@ def _evaluate_unsloth_is_mlx_gate() -> bool:
def _import_studio_hardware_module():
- """Lazy-load Studio's hardware module under the bare-imports layout."""
+ """Lazy-load Unsloth's hardware module under the bare-imports layout."""
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
# Fresh import so detect_hardware re-runs under the current spoofs.
@@ -290,7 +290,7 @@ def test_unsloth_is_mlx_gate_matches_profile(profile, spoof_hardware):
@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_studio_detect_hardware_matches_profile(profile, spoof_hardware):
- """Studio's detect_hardware() routes to the right DeviceType per profile."""
+ """Unsloth's detect_hardware() routes to the right DeviceType per profile."""
spoof_hardware(profile)
hw = _import_studio_hardware_module()
detected = hw.detect_hardware()
@@ -306,7 +306,7 @@ def test_studio_detect_hardware_matches_profile(profile, spoof_hardware):
@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_studio_is_apple_silicon_matches_profile(profile, spoof_hardware):
- """Studio's is_apple_silicon() helper agrees with platform spoof."""
+ """Unsloth's is_apple_silicon() helper agrees with platform spoof."""
spoof_hardware(profile)
hw = _import_studio_hardware_module()
assert hw.is_apple_silicon() is profile.expect_apple_silicon, (
diff --git a/tests/studio/test_is_mlx_dispatch_gate.py b/tests/studio/test_is_mlx_dispatch_gate.py
index f31f6d1655..0e5de1b789 100644
--- a/tests/studio/test_is_mlx_dispatch_gate.py
+++ b/tests/studio/test_is_mlx_dispatch_gate.py
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
-"""Regression tests for the CUDA-vs-MLX dispatch gates Studio relies on.
+"""Regression tests for the CUDA-vs-MLX dispatch gates Unsloth relies on.
Two gates: (1) ``unsloth._IS_MLX`` (import-time, delegates to the zoo MLX
runtime gate behind a local precheck barrier); (2)
@@ -147,7 +147,7 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
def _import_studio_hardware():
- """Lazy import of the Studio hardware module (studio/backend on sys.path)."""
+ """Lazy import of the Unsloth hardware module (studio/backend on sys.path)."""
studio_backend = REPO_ROOT / "studio" / "backend"
if str(studio_backend) not in sys.path:
sys.path.insert(0, str(studio_backend))
diff --git a/tests/studio/test_llama_cpp_wall_clock_cap.py b/tests/studio/test_llama_cpp_wall_clock_cap.py
index f173cfbc55..b7b6917092 100644
--- a/tests/studio/test_llama_cpp_wall_clock_cap.py
+++ b/tests/studio/test_llama_cpp_wall_clock_cap.py
@@ -1,4 +1,4 @@
-"""Timeout policy checks for Studio's local llama-server path."""
+"""Timeout policy checks for Unsloth's local llama-server path."""
from __future__ import annotations
diff --git a/tests/studio/test_locale_root_direction_contract.py b/tests/studio/test_locale_root_direction_contract.py
index 1baabdbbca..20d10aa805 100644
--- a/tests/studio/test_locale_root_direction_contract.py
+++ b/tests/studio/test_locale_root_direction_contract.py
@@ -1,4 +1,4 @@
-"""Regression guard for locale changes affecting the entire Studio layout."""
+"""Regression guard for locale changes affecting the entire Unsloth layout."""
from pathlib import Path
diff --git a/tests/studio/test_node_decision.ps1 b/tests/studio/test_node_decision.ps1
index bd5d5c8677..44f3ef0e1e 100644
--- a/tests/studio/test_node_decision.ps1
+++ b/tests/studio/test_node_decision.ps1
@@ -60,7 +60,7 @@ $globalBunOffset = $source.IndexOf('npm install -g bun')
Check "NodeSource initialized before SKIP_STUDIO_FRONTEND branch" (
$nodeSourceOffset -ge 0 -and $skipFrontendBranchOffset -ge 0 -and $nodeSourceOffset -lt $skipFrontendBranchOffset
)
-Check "custom Studio home validated before Node parent creation" (
+Check "custom Unsloth home validated before Node parent creation" (
$customHomeErrorOffset -ge 0 -and $nodeParentMkdirOffset -ge 0 -and $customHomeErrorOffset -lt $nodeParentMkdirOffset
)
Check "bundled Node pins npm prefix and clears NODE_PATH" (
diff --git a/tests/studio/test_studio_gguf_export_script_pin.py b/tests/studio/test_studio_gguf_export_script_pin.py
index 1f7e7adaa4..defd0d49d4 100644
--- a/tests/studio/test_studio_gguf_export_script_pin.py
+++ b/tests/studio/test_studio_gguf_export_script_pin.py
@@ -1,4 +1,4 @@
-"""Studio GGUF export pins convert_hf_to_gguf.py via UNSLOTH_LLAMA_CPP_SCRIPTS_DIR, with a once-per-process warning fallback when unsloth_zoo lacks the local-script resolver."""
+"""Unsloth GGUF export pins convert_hf_to_gguf.py via UNSLOTH_LLAMA_CPP_SCRIPTS_DIR, with a once-per-process warning fallback when unsloth_zoo lacks the local-script resolver."""
from __future__ import annotations
diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py
index 73d1244ee3..359ca4873a 100644
--- a/tests/studio/test_studio_text_descender_clipping.py
+++ b/tests/studio/test_studio_text_descender_clipping.py
@@ -1,4 +1,4 @@
-"""Regression guard: Studio text spans must not pair `leading-none` with
+"""Regression guard: Unsloth text spans must not pair `leading-none` with
`truncate`, which clips glyph descenders (g, p, q, y, j) in visible labels.
"""
diff --git a/tests/test_studio_install_workspace_guard.py b/tests/test_studio_install_workspace_guard.py
index 89836cfb0d..ea5b34672c 100644
--- a/tests/test_studio_install_workspace_guard.py
+++ b/tests/test_studio_install_workspace_guard.py
@@ -1,4 +1,4 @@
-"""install.sh/install.ps1 must refuse to rm -rf an existing Studio venv in env-mode without a sentinel."""
+"""install.sh/install.ps1 must refuse to rm -rf an existing Unsloth venv in env-mode without a sentinel."""
from __future__ import annotations
@@ -127,7 +127,7 @@ def test_install_ps1_has_matching_env_mode_guard():
), "install.ps1 must gate Remove-Item $VenvDir on env-mode"
assert "share\\studio.conf" in block, "install.ps1 guard must check share\\studio.conf sentinel"
assert "bin\\unsloth.exe" in block, "install.ps1 guard must check bin\\unsloth.exe sentinel"
- assert "Refusing to delete non-Studio venv" in block
+ assert "Refusing to delete non-Unsloth venv" in block
def test_setup_ps1_has_writability_probe():
@@ -160,7 +160,7 @@ def test_env_mode_blocks_when_bin_unsloth_is_a_directory(tmp_path):
capture_output = True,
)
assert res.returncode != 0, (
- "directory at bin/unsloth must NOT satisfy the Studio sentinel; "
+ "directory at bin/unsloth must NOT satisfy the Unsloth sentinel; "
f"stdout={res.stdout!r} stderr={res.stderr!r}"
)
assert (venv / "important.txt").is_file(), "unrelated workspace data must survive"
@@ -205,7 +205,7 @@ def test_install_ps1_sentinel_uses_pathtype_leaf():
def test_setup_ps1_stale_venv_has_env_mode_guard():
- """setup.ps1 stale-venv branch must gate Remove-Item $VenvDir on a custom-root Studio sentinel."""
+ """setup.ps1 stale-venv branch must gate Remove-Item $VenvDir on a custom-root Unsloth sentinel."""
src = SETUP_PS1.read_text()
idx = src.index("Stale venv detected")
block = src[idx : idx + 1500]
diff --git a/tests/test_studio_root_resilience.py b/tests/test_studio_root_resilience.py
index 0dfb826376..779ff2f3f1 100644
--- a/tests/test_studio_root_resilience.py
+++ b/tests/test_studio_root_resilience.py
@@ -1,4 +1,4 @@
-"""Studio install-root inference must not crash under hostile filesystem conditions (PermissionError/OSError swallowed; custom root kept when resolve() fails)."""
+"""Unsloth install-root inference must not crash under hostile filesystem conditions (PermissionError/OSError swallowed; custom root kept when resolve() fails)."""
from __future__ import annotations
diff --git a/tests/test_studio_shutdown_thread_wait.py b/tests/test_studio_shutdown_thread_wait.py
index 4ec2afc0f0..8299116d9a 100644
--- a/tests/test_studio_shutdown_thread_wait.py
+++ b/tests/test_studio_shutdown_thread_wait.py
@@ -130,4 +130,4 @@ def test_cli_entrypoints_wait_before_returning_to_shell():
assert (
_calls_shutdown_wait_getattr(tree) >= 3
- ), "Studio CLI terminal paths must wait for the backend thread after requesting shutdown"
+ ), "Unsloth CLI terminal paths must wait for the backend thread after requesting shutdown"
diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index bd612db1d5..f47c78ba80 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -2047,7 +2047,7 @@ def get_chat_template(
.replace("'assistant'", "'" + mapping["assistant"] + "'")
if use_zoo_tokenizer_patch:
- # Studio MLX avoids the model-utils tokenizer wrapper because that
+ # Unsloth MLX avoids the model-utils tokenizer wrapper because that
# import path pulls in Torch/GPU-specific modules before MLX training.
from unsloth_zoo.tokenizer_utils import patch_tokenizer
else:
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index 09de248c7b..5d54815705 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -3008,7 +3008,7 @@ def maybe_set_windows_rocm_bnb_version():
No-op unless ALL of: Windows, a real HIP torch build (env hints like
HIP_PATH do not count), a ROCm DLL installed, and no explicit user value.
- Linux is untouched. Values seeded by Studio's venv sitecustomize.py
+ Linux is untouched. Values seeded by Unsloth's venv sitecustomize.py
(marked ``UNSLOTH_BNB_ROCM_VERSION_SOURCE=sitecustomize``) are
redetectable defaults, not overrides; ``UNSLOTH_SKIP_BNB_ROCM_VERSION=1``
opts out and drops a seeded default. Returns the value set, else None.
diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py
index 1e9cc641e9..7661b0d714 100644
--- a/unsloth/models/loader_utils.py
+++ b/unsloth/models/loader_utils.py
@@ -821,7 +821,7 @@ def _exclude_rope_inv_freq_from_ddp(model):
# =============================================================================
# Offline loading - single source of truth (shared by vision.py, loader.py and
-# the Studio exporter). Decide offline ONCE at the load boundary and force it
+# the Unsloth exporter). Decide offline ONCE at the load boundary and force it
# ONCE around the whole load, so every nested HF call inherits it.
# =============================================================================
diff --git a/unsloth/save.py b/unsloth/save.py
index 7d5774aa97..0e2650b174 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -228,7 +228,7 @@ def _loaded_via_remote_code(obj):
Transformers loads auto_map code into the ``transformers_modules`` package, so a
``transformers_modules`` class proves the original load actually ran that remote code
- (which the caller's / Studio's consent gate scans at load time). Export paths derive their
+ (which the caller's / Unsloth's consent gate scans at load time). Export paths derive their
reload trust_remote_code from this - the already approved load decision - instead of from a
checkpoint's static ``auto_map``: a model that loads with built-in classes must not have its
unvetted remote code run when it is re-read during quantization export. Walks PEFT / wrapper
@@ -3858,7 +3858,7 @@ def _prewarm_base_model_hub_cache(
from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download
# Resolve the cache from the live env like the merge, not huggingface_hub's frozen
- # constants: a runtime cache redirect (read-only default, Studio) would else miss (#6890).
+ # constants: a runtime cache redirect (read-only default, Unsloth) would else miss (#6890).
try:
from unsloth_zoo.hf_cache import _active_caches
_hub_cache = _active_caches()[1]
diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py
index d6e8247ec4..c7f61288d5 100644
--- a/unsloth/tokenizer_utils.py
+++ b/unsloth/tokenizer_utils.py
@@ -1476,7 +1476,7 @@ def get_tokenizer_info(tokenizer) -> dict:
"""Return a concise diagnostic summary of a tokenizer instance.
Collects key properties into a JSON-safe dict for logging, debugging, or the
- Studio UI. Missing attributes fall back to ``None`` rather than raising.
+ Unsloth UI. Missing attributes fall back to ``None`` rather than raising.
Example output::
diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py
index b3831f5314..121b26f03f 100644
--- a/unsloth_cli/__init__.py
+++ b/unsloth_cli/__init__.py
@@ -81,7 +81,7 @@ app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")
app.add_typer(
start_app,
name = "start",
- help = "Start a coding agent (Claude, Codex, OpenClaw, OpenCode, Hermes, Pi) against Studio.",
+ help = "Start a coding agent (Claude, Codex, OpenClaw, OpenCode, Hermes, Pi) against Unsloth.",
)
# Backwards-compatible hidden alias: `unsloth connect` routes to `unsloth start`.
app.add_typer(
diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py
index c3b188710e..551bef4787 100644
--- a/unsloth_cli/_inference.py
+++ b/unsloth_cli/_inference.py
@@ -18,7 +18,7 @@ _THINK_OPEN = ""
_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*? ", re.DOTALL)
_STREAMED_ERROR_PREFIX = "Error: "
-# Cloudflare (in front of remote Studio proxies like RunPod) 403s the default
+# Cloudflare (in front of remote Unsloth proxies like RunPod) 403s the default
# "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request.
_USER_AGENT = "unsloth-cli"
_MPI_ENV_PAIRS = (
@@ -36,7 +36,7 @@ _no_redirect_opener = None
def urlopen_no_redirect(request, timeout):
"""urlopen that errors on any redirect: following a 3xx would send a bearer
token (or accept an identity proof) to a base we never vetted, letting a port
- squatter relay a real Studio's response."""
+ squatter relay a real Unsloth's response."""
global _no_redirect_opener
if _no_redirect_opener is None:
import urllib.error
@@ -540,7 +540,7 @@ def find_studio_server(timeout: float = 3.0) -> Optional[str]:
def is_loopback_url(base: str) -> bool:
"""True only when *base* resolves to loopback. find_studio_server() trusts a
base after only a health probe, so credentials are auto-sent only to loopback
- (a local Studio or an SSH tunnel on 127.0.0.1), the targets the auto flows mean."""
+ (a local Unsloth or an SSH tunnel on 127.0.0.1), the targets the auto flows mean."""
from urllib.parse import urlparse
host = (urlparse(base).hostname or "").lower()
@@ -554,7 +554,7 @@ def is_loopback_url(base: str) -> bool:
def verify_studio_identity(base: str, timeout: float = 3.0) -> bool:
- """Confirm `base` is really this machine's Studio before sending a secret.
+ """Confirm `base` is really this machine's Unsloth before sending a secret.
Send a random nonce to /api/auth/identity and check the returned HMAC against
the one computed from the local same-user secret; an endpoint without that
@@ -578,7 +578,7 @@ def verify_studio_identity(base: str, timeout: float = 3.0) -> bool:
port = parsed.port or (443 if parsed.scheme == "https" else 80)
# Resolve to one concrete address and talk to *that* address, then bind the
# proof to (address, port). A name like localhost can resolve to a squatter on
- # ::1 while the real Studio is on 127.0.0.1; connecting to the resolved IP and
+ # ::1 while the real Unsloth is on 127.0.0.1; connecting to the resolved IP and
# binding to it means a proof relayed from a different address/port won't match.
try:
ip = socket.getaddrinfo(host, port, type = socket.SOCK_STREAM)[0][4][0]
@@ -592,7 +592,7 @@ def verify_studio_identity(base: str, timeout: float = 3.0) -> bool:
headers = {"User-Agent": _USER_AGENT, "Host": parsed.netloc},
)
try:
- # No redirects: a 302 could relay a real Studio's proof (see urlopen_no_redirect).
+ # No redirects: a 302 could relay a real Unsloth's proof (see urlopen_no_redirect).
# Cap the read: the server is still unverified, so don't trust its length.
with urlopen_no_redirect(request, timeout = timeout) as response:
proof = json.loads(response.read(65536).decode() or "{}").get("proof")
@@ -623,7 +623,7 @@ def _studio_token() -> Optional[str]:
class HttpChatBackend:
- """Chat against a running Studio server over its OpenAI-compatible API.
+ """Chat against a running Unsloth server over its OpenAI-compatible API.
close() leaves the model loaded on purpose — the next session (or the
UI) starts instantly.
@@ -666,7 +666,7 @@ class HttpChatBackend:
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
) -> None:
- typer.echo(f"Loading {model} on the Studio server", err = True)
+ typer.echo(f"Loading {model} on the Unsloth server", err = True)
payload = {
"model_path": model,
"hf_token": hf_token,
@@ -769,7 +769,7 @@ def connect_studio_server(
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
):
- """Backend on a running Studio server, or None (caller loads locally)."""
+ """Backend on a running Unsloth server, or None (caller loads locally)."""
base_url = find_studio_server()
if not base_url:
return None
@@ -782,20 +782,20 @@ def connect_studio_server(
if not explicit:
return None
typer.echo(
- f"Can't attach to the Studio server at {base_url}: {reason} Run Studio "
+ f"Can't attach to the Unsloth server at {base_url}: {reason} Run Unsloth "
"on this machine, or unset UNSLOTH_STUDIO_URL to load the model locally.",
err = True,
)
raise typer.Exit(code = 1)
# Only hand the self-issued JWT (signed with the local secret) to loopback: a
- # remote URL is unverified and a real remote Studio would reject it anyway.
+ # remote URL is unverified and a real remote Unsloth would reject it anyway.
if not is_loopback_url(base_url):
return _refuse(
- "it isn't a local Studio, so a self-issued token can't "
+ "it isn't a local Unsloth, so a self-issued token can't "
"authenticate to it and must not be sent to it."
)
- # Confirm the loopback responder is really our Studio (not a port squatter).
+ # Confirm the loopback responder is really our Unsloth (not a port squatter).
if not verify_studio_identity(base_url):
return _refuse(
"its identity couldn't be verified (it may be running as a "
@@ -803,7 +803,7 @@ def connect_studio_server(
)
token = _studio_token()
if not token:
- return _refuse("couldn't self-issue a Studio token (is Studio set up here?).")
+ return _refuse("couldn't self-issue an Unsloth token (is Unsloth set up here?).")
backend = HttpChatBackend(base_url, token)
backend.ensure_loaded(
model,
diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py
index bc4a72f36c..bba5fab08e 100644
--- a/unsloth_cli/commands/chat.py
+++ b/unsloth_cli/commands/chat.py
@@ -206,7 +206,7 @@ def chat(
no_server: bool = typer.Option(
False,
"--no-server",
- help = "Load the model in-process even if a Studio server is running.",
+ help = "Load the model in-process even if an Unsloth server is running.",
),
):
"""Start an interactive chat with a model (loads once, stays warm)."""
@@ -262,14 +262,14 @@ def chat(
llama_extra_args = llama_extra_args,
)
- # Prefer a running Studio server: instant starts, model shared with the UI.
+ # Prefer a running Unsloth server: instant starts, model shared with the UI.
chat_backend = (
None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts)
)
server_mode = chat_backend is not None
if server_mode and should_print:
console.print(
- "(Studio server connected — model stays warm after /exit)",
+ "(Unsloth server connected — model stays warm after /exit)",
style = "bright_black",
)
else:
diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py
index 84a126163e..524d8fd015 100644
--- a/unsloth_cli/commands/inference.py
+++ b/unsloth_cli/commands/inference.py
@@ -67,7 +67,7 @@ def inference(
no_server: bool = typer.Option(
False,
"--no-server",
- help = "Load the model in-process even if a Studio server is running.",
+ help = "Load the model in-process even if an Unsloth server is running.",
),
):
"""Run a single inference using the specified model."""
@@ -85,7 +85,7 @@ def inference(
)
raise typer.Exit(code = 1)
- # A running Studio server keeps the model warm between runs. Under
+ # A running Unsloth server keeps the model warm between runs. Under
# mlx.launch, every rank must enter the local MLX path instead of rank 0
# alone talking to a server.
load_opts = dict(
diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py
index fa39fdf761..9257a7fbcb 100644
--- a/unsloth_cli/commands/start.py
+++ b/unsloth_cli/commands/start.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
-"""`unsloth start` — launch a coding agent against a running Studio server."""
+"""`unsloth start` — launch a coding agent against a running Unsloth server."""
import atexit
import contextlib
@@ -35,7 +35,7 @@ from unsloth_cli._inference import (
)
start_app = typer.Typer(
- help = "Start a coding agent against a running Studio server.",
+ help = "Start a coding agent against a running Unsloth server.",
no_args_is_help = True,
context_settings = {"help_option_names": ["-h", "--help"]},
)
@@ -75,14 +75,14 @@ _CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN")
# Shared by every agent command; only the config/env/command differ.
_MODEL_OPTION = typer.Option(
- None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio."
+ None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Unsloth."
)
_KEY_OPTION = typer.Option(
None,
"--api-key",
envvar = "UNSLOTH_API_KEY",
help = (
- "Studio API key. For a local Studio it is minted automatically and "
+ "Unsloth API key. For a local Unsloth it is minted automatically and "
"remembered per server. For a remote server, pass one with --api-key "
"(or UNSLOTH_API_KEY); it is remembered for next time."
),
@@ -96,7 +96,7 @@ _SERVE_OPTION = typer.Option(
True,
"--serve/--no-serve",
help = (
- "If no Studio server is running, auto-start one for --model and stop it when the "
+ "If no Unsloth server is running, auto-start one for --model and stop it when the "
"agent exits. --no-serve keeps the old behavior of erroring out."
),
)
@@ -347,7 +347,7 @@ def _shutdown_auto_served() -> None:
global _auto_served_server
server, _auto_served_server = _auto_served_server, None
if server is not None and server.poll() is None:
- typer.echo("Stopping the auto-started Studio server…")
+ typer.echo("Stopping the auto-started Unsloth server…")
_shutdown_server(server)
@@ -381,7 +381,7 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess
log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log"
typer.echo(
- f"No Studio server at {base}. Starting one for {model} (loading the model can take a while)…"
+ f"No Unsloth server at {base}. Starting one for {model} (loading the model can take a while)…"
)
typer.echo(f"Server log: {log_path}")
# 0600: the `unsloth run` banner in this log carries the minted sk-unsloth- key, and
@@ -408,16 +408,16 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess
if server.poll() is not None:
tail = _log_tail(log_path)
_shutdown_auto_served()
- _fail(f"The Studio server stopped before it was ready. Last log lines:\n{tail}")
+ _fail(f"The Unsloth server stopped before it was ready. Last log lines:\n{tail}")
# `unsloth run` prints the minted key only after the server is up AND the model is
# loaded, so it is the fully-ready signal (same contract serve-unsloth-run.sh uses).
if _studio_healthy(base) and "sk-unsloth-" in _log_tail(log_path, lines = 400):
- typer.echo(f"Studio server ready at {base}.")
+ typer.echo(f"Unsloth server ready at {base}.")
return server
time.sleep(2.0)
_shutdown_auto_served()
_fail(
- f"The Studio server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}."
+ f"The Unsloth server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}."
)
@@ -463,7 +463,7 @@ def _require_studio(
return expected, _start_studio_server(expected, model, load or LoadOptions())
model_hint = "" if model else " Pass --model to have it start one for you, or"
_fail(
- f"No running Studio server found at {expected}.{model_hint} start one with "
+ f"No running Unsloth server found at {expected}.{model_hint} start one with "
"`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server."
)
@@ -567,12 +567,12 @@ def _key_accepted(base: str, key: str) -> bool:
if exc.code in (401, 403):
return False
_fail(
- f"Studio server error while checking an API key ({exc.code}). "
+ f"Unsloth server error while checking an API key ({exc.code}). "
"The server may be starting up or unhealthy; try again shortly."
)
except (urllib.error.URLError, TimeoutError) as exc:
_fail(
- "Couldn't reach the Studio server while checking an API key: "
+ "Couldn't reach the Unsloth server while checking an API key: "
f"{getattr(exc, 'reason', None) or exc}"
)
@@ -592,10 +592,10 @@ def _agent_api_key(
# UNSLOTH_API_KEY meant for some other server must not fail the
# launch: the loopback mint path below is guaranteed to work.
# (An explicit key that the fresh server accepts, e.g. one persisted
- # in this Studio home's auth db, is still honored above.)
+ # in this Unsloth home's auth db, is still honored above.)
# Replay a key the user saved for *this exact* server first (scoped per base,
- # so it only goes back there -- including a remote/SSH-tunnelled Studio whose
+ # so it only goes back there -- including a remote/SSH-tunnelled Unsloth whose
# secret the local handshake can't match). Skip ones the server rejects.
for key in _cached_keys(cache, base, "saved"):
if _key_accepted(base, key):
@@ -608,15 +608,15 @@ def _agent_api_key(
if not is_loopback_url(base):
_fail(
f"No saved API key for {base} and automatic minting only runs against "
- "a local Studio. Create an API key in Studio → Settings → API and "
+ "a local Unsloth. Create an API key in Unsloth → Settings → API and "
"pass it with --api-key (it is remembered per server), or set "
"UNSLOTH_API_KEY."
)
if not verify_studio_identity(base):
_fail(
- f"Couldn't verify that {base} is your Studio (it may be running as a "
+ f"Couldn't verify that {base} is your Unsloth (it may be running as a "
"different OS user, or another process took the port). Create an API "
- "key in Studio → Settings → API and pass it with --api-key, or set "
+ "key in Unsloth → Settings → API and pass it with --api-key, or set "
"UNSLOTH_API_KEY."
)
@@ -630,8 +630,8 @@ def _agent_api_key(
token = _studio_token()
if token is None:
_fail(
- "Couldn't authenticate with the Studio server automatically. Create "
- "an API key in Studio → Settings → API and pass it with --api-key, "
+ "Couldn't authenticate with the Unsloth server automatically. Create "
+ "an API key in Unsloth → Settings → API and pass it with --api-key, "
"or set UNSLOTH_API_KEY."
)
key = _http_json(
@@ -664,7 +664,7 @@ def _is_hub_model_id(value: object) -> bool:
return False
# A hub id is exactly "namespace/name" over a restricted charset. Anything with
# extra path segments (e.g. a server-side relative path such as
- # models/Llama/Foo.gguf on a remote Studio) is not a hub id and must not be
+ # models/Llama/Foo.gguf on a remote Unsloth) is not a hub id and must not be
# casefold-matched against a differently cased path on a case-sensitive
# filesystem. This is host independent, unlike the existence probe below which
# cannot see a path that only exists on the server.
@@ -690,8 +690,8 @@ def _model_id_matches(
if actual == requested:
return True
# Case-insensitive matching is only safe when the local existence probe in
- # _is_hub_model_id is authoritative, i.e. against a loopback Studio on this host.
- # Against a remote Studio a two-segment string is indistinguishable from a
+ # _is_hub_model_id is authoritative, i.e. against a loopback Unsloth on this host.
+ # Against a remote Unsloth a two-segment string is indistinguishable from a
# server-side relative path (e.g. Models/Foo vs models/foo), so casefolding it
# could attach to the wrong model on a case-sensitive server; defer to an exact
# match there and let the load endpoint resolve the requested path.
@@ -709,7 +709,7 @@ def _resolve_model(
load: LoadOptions = LoadOptions(),
) -> dict:
models = _loaded_models(base, key)
- # Only casefold-match ids against a loopback Studio, where _is_hub_model_id's
+ # Only casefold-match ids against a loopback Unsloth, where _is_hub_model_id's
# local existence probe can actually reject a server-side path; see the note there.
allow_casefold = is_loopback_url(base)
# /v1/models reports the model id but not the active GGUF variant or runtime load
@@ -741,7 +741,7 @@ def _resolve_model(
typer.echo(
f"Ensuring {requested} is loaded with the requested settings…"
if load_has_overrides
- else f"Loading {requested} on the Studio server (this can take a while)…"
+ else f"Loading {requested} on the Unsloth server (this can take a while)…"
)
# Mirror `unsloth run`'s load knobs; keep the default payload as just
# model_path so a bare `--model` load is unchanged.
@@ -762,7 +762,7 @@ def _resolve_model(
timeout = 3600,
error = "Model load failed",
)
- # Studio registers the model under a canonical id (resolved identifier,
+ # Unsloth registers the model under a canonical id (resolved identifier,
# casing) that /v1/models echoes but which may differ from the path we
# passed; match on the id the load reports so we don't silently fall
# through to models[0] and connect to a different loaded model.
@@ -783,22 +783,22 @@ def _resolve_model(
if match is not None:
return match
if requested:
- # We asked Studio to load it and it didn't surface in /v1/models; don't
+ # We asked Unsloth to load it and it didn't surface in /v1/models; don't
# silently hand back an unrelated loaded model.
_fail(
- f"Studio didn't report '{requested}' as loaded. Double-check the model "
+ f"Unsloth didn't report '{requested}' as loaded. Double-check the model "
"id, or load it from the model dropdown in the UI."
)
if not models:
_fail(
- "No model is loaded in Studio. Load one from the model dropdown in "
+ "No model is loaded in Unsloth. Load one from the model dropdown in "
"the UI, or pass --model to load it from here."
)
return models[0]
def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None:
- # Codex always streams, and Studio only streams /v1/responses from llama-server.
+ # Codex always streams, and Unsloth only streams /v1/responses from llama-server.
try:
status = _http_json("GET", f"{base}/api/inference/status", key)
except urllib.error.HTTPError as exc:
@@ -901,7 +901,7 @@ def _codex_supports_model_catalog() -> bool:
def _codex_model_catalog(model: dict) -> dict:
- """Return conservative metadata for a Studio model unknown to Codex's built-in catalog."""
+ """Return conservative metadata for an Unsloth model unknown to Codex's built-in catalog."""
model_id = model["id"]
window = model.get("context_length") or model.get("max_context_length")
entry = {
@@ -1202,7 +1202,7 @@ def _connect(
# `--model org/name:QUANT` is shorthand for `--model org/name --gguf-variant QUANT`.
# Split it before we match/serve so the attach path resolves against the already-loaded
# `org/name` (listed without the suffix) instead of reloading a `:`-suffixed repo id --
- # which Studio rejects and which would evict a model another session is using.
+ # which Unsloth rejects and which would evict a model another session is using.
if model:
repo, variant = _split_repo_variant(model)
if variant:
@@ -1240,7 +1240,7 @@ def _run(
# --no-launch recipes stay intact.
if launch and clear_screen:
click.clear()
- typer.echo(f"Studio {base} · model {entry['id']}")
+ typer.echo(f"Unsloth {base} · model {entry['id']}")
wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else ()
if not launch:
_print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge)
@@ -1306,7 +1306,7 @@ def write_openclaw_config(
)
return
before = json.dumps(config, sort_keys = True)
- # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path).
+ # Unsloth is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path).
provider_model = {"id": model["id"], "name": model["id"]}
window = model.get("context_length") or model.get("max_context_length")
if window:
@@ -1570,14 +1570,14 @@ def write_pi_config(base: str, key: str, model: dict, path: Path) -> None:
return
before = json.dumps(config, sort_keys = True)
# Pi reads custom providers from ~/.pi/agent/models.json (HOME-relocated for the
- # session). Studio is a generic OpenAI-compatible /v1 endpoint, and the key lives
+ # session). Unsloth is a generic OpenAI-compatible /v1 endpoint, and the key lives
# in the config rather than the env (matching openclaw/opencode).
provider_model = {"id": model["id"]}
window = model.get("context_length") or model.get("max_context_length")
if window:
window = int(window)
# An unspecified model defaults to contextWindow 128000 / maxTokens 16384,
- # far larger than a small Studio context, so Pi compacts too late and overflows
+ # far larger than a small Unsloth context, so Pi compacts too late and overflows
# the server. Pin the real window and a sane output cap (mirrors OpenCode).
provider_model["contextWindow"] = window
provider_model["maxTokens"] = min(window // 4, 8192)
@@ -1606,7 +1606,7 @@ def claude(
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
- """Point Claude Code at the running Studio server and start it."""
+ """Point Claude Code at the running Unsloth server and start it."""
base, key, entry = _connect(
api_key,
model,
@@ -1690,7 +1690,7 @@ def codex(
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
- """Point OpenAI Codex at the running Studio server and start it."""
+ """Point OpenAI Codex at the running Unsloth server and start it."""
base, key, entry = _connect(
api_key,
model,
@@ -1734,7 +1734,7 @@ def openclaw(
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
- """Point OpenClaw at the running Studio server and start it."""
+ """Point OpenClaw at the running Unsloth server and start it."""
base, key, entry = _connect(
api_key,
model,
@@ -1791,7 +1791,7 @@ def opencode(
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
- """Point OpenCode at the running Studio server and start it."""
+ """Point OpenCode at the running Unsloth server and start it."""
base, key, entry = _connect(
api_key,
model,
@@ -1835,7 +1835,7 @@ def opencode(
# setting them in the highest-priority inline overlay neutralizes any user allowlist
# or denylist for the launch. It is session-only: it lives in OPENCODE_CONFIG_CONTENT
# for this invocation and never touches the user's config files, so their normal
- # `opencode` is unchanged; only this session is limited to the Studio provider.
+ # `opencode` is unchanged; only this session is limited to the Unsloth provider.
# small_model is opencode's separate model for lightweight tasks; pin it to the
# session model too, or a user/project small_model on another (now filtered)
# provider would resolve a not-found error mid-session. The session serves one
@@ -1869,7 +1869,7 @@ def hermes(
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
- """Point Hermes (Nous Research) at the running Studio server and start it."""
+ """Point Hermes (Nous Research) at the running Unsloth server and start it."""
native_args = [*_yolo_command_flags("hermes", yolo), *ctx.args]
command = ["hermes", *_hermes_resume_oneshot_args(native_args)]
base, key, entry = _connect(
@@ -1902,7 +1902,7 @@ def pi(
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
- """Point Pi (coding agent) at the running Studio server and start it."""
+ """Point Pi (coding agent) at the running Unsloth server and start it."""
base, key, entry = _connect(
api_key,
model,
diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py
index 09355bd454..f2f41fc583 100644
--- a/unsloth_cli/commands/studio.py
+++ b/unsloth_cli/commands/studio.py
@@ -40,7 +40,7 @@ def _enable_verbose_access_logs() -> None:
# UNSLOTH_STUDIO_HOME wins when both env vars are set.
def _looks_like_installer_managed_studio_home(candidate: Path) -> bool:
"""Sentinel check (studio.conf or bin shim) so a dev venv named
- unsloth_studio is not misidentified as a custom Studio root.
+ unsloth_studio is not misidentified as a custom Unsloth root.
"""
shim_name = "unsloth.exe" if platform.system() == "Windows" else "unsloth"
return (candidate / "share" / "studio.conf").is_file() or (
@@ -212,7 +212,7 @@ def _find_run_py() -> Optional[Path]:
run_py = _PACKAGE_ROOT / "studio" / "backend" / "run.py"
if run_py.is_file():
return run_py
- # 2. Studio venv's site-packages (Linux + Windows layouts)
+ # 2. Unsloth venv's site-packages (Linux + Windows layouts)
for pattern in (
"lib/python*/site-packages/studio/backend/run.py",
"Lib/site-packages/studio/backend/run.py",
@@ -273,7 +273,7 @@ def _find_setup_script() -> Optional[Path]:
s = _PACKAGE_ROOT / "studio" / name
if s.is_file():
return s
- # 2. Studio venv's site-packages
+ # 2. Unsloth venv's site-packages
for pattern in (
f"lib/python*/site-packages/studio/{name}",
f"Lib/site-packages/studio/{name}",
@@ -641,7 +641,7 @@ def _create_desktop_secret_in_cli() -> str:
def _should_prompt_password_change(
*, cloudflare: Optional[bool], host: str, secure: bool, api_only: bool
) -> bool:
- """Whether this launch will expose Studio through the Cloudflare tunnel.
+ """Whether this launch will expose Unsloth through the Cloudflare tunnel.
CLI mirror of run.py's _cloudflare_tunnel_should_start, minus the Colab
case (Colab launches never come through this CLI path). --secure implies
@@ -747,7 +747,7 @@ def _apply_supplied_password_before_launch(supplied_password: "str | None") -> N
conn = _connect_auth_db()
except (OSError, sqlite3.Error) as exc:
typer.echo(
- f"Error: --password could not open the Studio auth database ({exc}); not starting.",
+ f"Error: --password could not open the Unsloth auth database ({exc}); not starting.",
err = True,
)
raise typer.Exit(1)
@@ -767,7 +767,7 @@ def _apply_supplied_password_before_launch(supplied_password: "str | None") -> N
raise typer.Exit(1)
if not row[2]:
typer.echo(
- "Error: a Studio admin password is already set; --password only sets "
+ "Error: an Unsloth admin password is already set; --password only sets "
"the initial password. Run `unsloth studio reset-password` first "
"(or change it in the UI).",
err = True,
@@ -790,7 +790,7 @@ def _apply_supplied_password_before_launch(supplied_password: "str | None") -> N
# Any DB failure fails closed (typer.Exit is not caught here, so the
# deliberate Exit(1) branches above propagate unchanged).
typer.echo(
- f"Error: --password could not update the Studio auth database ({exc}); not starting.",
+ f"Error: --password could not update the Unsloth auth database ({exc}); not starting.",
err = True,
)
raise typer.Exit(1)
@@ -813,9 +813,9 @@ def _strip_seeded_bootstrap_password_or_exit(*, context: str) -> None:
bootstrap_file.unlink(missing_ok = True)
except OSError as exc:
typer.echo(
- "Error: refusing to publish Studio on a public Cloudflare URL: "
+ "Error: refusing to publish Unsloth on a public Cloudflare URL: "
f"could not remove the seeded bootstrap password file ({exc}), so an "
- f"older Studio child could still serve the default credential ({context}). "
+ f"older Unsloth child could still serve the default credential ({context}). "
"Delete it manually or change the admin password (run `unsloth studio` "
"locally with a terminal attached, or `unsloth studio reset-password`), "
"then retry.",
@@ -851,7 +851,7 @@ def _require_servable_frontend_or_exit(
return frontend
typer.echo(
"Error: --frontend points at a directory with no index.html, so a "
- "public Studio launch would have no login page to change the seeded "
+ "public Unsloth launch would have no login page to change the seeded "
"admin password. Point --frontend at a built dist, rebuild it (re-run "
"install.sh), or use --api-only.",
err = True,
@@ -862,7 +862,7 @@ def _require_servable_frontend_or_exit(
if resolved is not None:
return resolved
typer.echo(
- "Error: the Studio frontend is not built, so a public launch would have "
+ "Error: the Unsloth frontend is not built, so a public launch would have "
"no login page to change the seeded admin password. Build it (re-run "
"install.sh), pass --frontend PATH to a built dist, or use --api-only.",
err = True,
@@ -892,8 +892,8 @@ def _validate_inproc_backend_before_strip(
_load_run_module()
except Exception as exc:
typer.echo(
- f"Error: the Studio backend could not be loaded ({exc}); refusing to "
- "expose Studio publicly before it is confirmed runnable. Re-run: "
+ f"Error: the Unsloth backend could not be loaded ({exc}); refusing to "
+ "expose Unsloth publicly before it is confirmed runnable. Re-run: "
"unsloth studio setup",
err = True,
)
@@ -902,7 +902,7 @@ def _validate_inproc_backend_before_strip(
def _tunnel_binary_confirmed_unavailable() -> bool:
"""True only if cloudflared is provably unavailable (found nowhere on PATH or
- in the Studio cache AND the download failed), so the tunnel cannot start.
+ in the Unsloth cache AND the download failed), so the tunnel cannot start.
Used on the --secure path (loopback bind, so the tunnel is the ONLY public
exposure) to skip stripping the seeded recovery password before a public URL
@@ -921,7 +921,7 @@ def _tunnel_binary_confirmed_unavailable() -> bool:
if not tunnel_py.is_file():
return False
# ensure_cloudflared() lazily imports utils.paths.storage_roots to resolve the
- # Studio bin cache. The outer CLI hasn't added studio/backend to sys.path yet,
+ # Unsloth bin cache. The outer CLI hasn't added studio/backend to sys.path yet,
# so that import would fail and return None (a false "unavailable" that wrongly
# refuses --secure). Add the backend dir so the cache path resolves as in the child.
added_backend_path = False
@@ -946,7 +946,7 @@ def _tunnel_binary_confirmed_unavailable() -> bool:
def _child_self_suppresses(*, in_studio_venv: bool, child_run_py: Optional[Path]) -> bool:
- """True when the child that will serve Studio is provably THIS install's
+ """True when the child that will serve Unsloth is provably THIS install's
backend, whose pre-bind gate sets app.state.suppress_bootstrap_injection and
so never serves the seeded credential publicly -- even with .bootstrap_password
on disk. The parent-side strip is then unnecessary and can be skipped to avoid
@@ -1002,8 +1002,8 @@ def _enforce_password_change_before_exposure(
# Refuse rather than risk a child serving the default login; a transient
# lock clears on retry.
typer.echo(
- "Error: refusing to publish Studio on a public Cloudflare URL: could "
- f"not open the Studio auth database ({exc}) to confirm the admin "
+ "Error: refusing to publish Unsloth on a public Cloudflare URL: could "
+ f"not open the Unsloth auth database ({exc}) to confirm the admin "
"password was changed. Retry (a transient database lock clears), or "
"change the password first (run `unsloth studio` locally with a "
"terminal attached, or `unsloth studio reset-password`).",
@@ -1028,8 +1028,8 @@ def _enforce_password_change_before_exposure(
except OSError:
pass
typer.echo(
- "Error: refusing to publish Studio on a public Cloudflare URL: could "
- f"not initialize the admin account ({exc}), so a re-exec'd Studio "
+ "Error: refusing to publish Unsloth on a public Cloudflare URL: could "
+ f"not initialize the admin account ({exc}), so a re-exec'd Unsloth "
"child could regenerate and serve a default credential. Retry (a "
"transient database lock clears), or change the password first (run "
"`unsloth studio` locally with a terminal attached, or `unsloth "
@@ -1053,7 +1053,7 @@ def _enforce_password_change_before_exposure(
# regenerate; we just couldn't read must_change back. Strip the seeded
# file so nothing serves it, failing closed if the strip itself fails.
typer.echo(
- f"Warning: could not read the Studio admin state back ({exc}); "
+ f"Warning: could not read the Unsloth admin state back ({exc}); "
"removing the seeded bootstrap password before public exposure.",
err = True,
)
@@ -1066,7 +1066,7 @@ def _enforce_password_change_before_exposure(
# the launch: it never arms for api-only, and TIMEOUT=0 disables it.
if api_only or not _bootstrap_deadline_active():
typer.echo(
- "Error: refusing to publish Studio on a public Cloudflare "
+ "Error: refusing to publish Unsloth on a public Cloudflare "
"URL: the default admin password was never changed, no "
"terminal is attached to change it here, and the bootstrap "
"shutdown deadline does not apply to this launch (api-only, "
@@ -1085,12 +1085,12 @@ def _enforce_password_change_before_exposure(
# fails). Keep the file for LOCAL recovery; must_change stays set
# and the deadline arms.
typer.echo(
- "Warning: Studio is being exposed publicly while the admin "
+ "Warning: Unsloth is being exposed publicly while the admin "
"account still uses its auto-generated bootstrap password. The "
"login page forces a change and the credential is never served "
"on the public page. Set a new password by running `unsloth "
"studio` locally with a terminal attached, or `unsloth studio "
- "reset-password`; Studio shuts down after ~1h if the password "
+ "reset-password`; Unsloth shuts down after ~1h if the password "
"stays unchanged (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT).",
err = True,
)
@@ -1104,7 +1104,7 @@ def _enforce_password_change_before_exposure(
# uncertainty.)
if secure and _tunnel_binary_confirmed_unavailable():
typer.echo(
- "Error: refusing to expose Studio: the Cloudflare tunnel binary "
+ "Error: refusing to expose Unsloth: the Cloudflare tunnel binary "
"(cloudflared) is unavailable and could not be downloaded, so no "
"public URL can start. The seeded bootstrap password is preserved "
"for recovery; fix connectivity and retry, or change the password "
@@ -1121,11 +1121,11 @@ def _enforce_password_change_before_exposure(
# forces a change and the timer still arms; only the on-disk copy goes.
_strip_seeded_bootstrap_password_or_exit(context = "no terminal to change it")
typer.echo(
- "Warning: Studio is being exposed publicly while the admin account "
+ "Warning: Unsloth is being exposed publicly while the admin account "
"still uses its auto-generated bootstrap password. The seeded password "
"file has been removed so it is not served on the public page. Set a new "
"password by running `unsloth studio` locally with a terminal attached, "
- "or `unsloth studio reset-password`; Studio shuts down after ~1h if the "
+ "or `unsloth studio reset-password`; Unsloth shuts down after ~1h if the "
"password stays unchanged (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT).",
err = True,
)
@@ -1146,7 +1146,7 @@ def _enforce_password_change_before_exposure(
new_password = _password_prompt.prompt_new_password(_is_current_password)
except (KeyboardInterrupt, EOFError):
typer.echo(
- "\nError: password change aborted; refusing to expose Studio "
+ "\nError: password change aborted; refusing to expose Unsloth "
"with the default admin password. Re-run and set a password, "
"or launch without --secure/--cloudflare.",
err = True,
@@ -1247,7 +1247,7 @@ def studio_default(
cloudflare: Optional[bool] = typer.Option(
None,
"--cloudflare/--no-cloudflare",
- help = "Expose Studio on a PUBLIC internet URL via a free Cloudflare HTTPS "
+ help = "Expose Unsloth on a PUBLIC internet URL via a free Cloudflare HTTPS "
"tunnel, for non-api-only wildcard binds (0.0.0.0 or ::). Off by default; "
"pass --cloudflare to enable it (--secure implies it). --no-cloudflare forces "
"it off but does not change a raw wildcard bind.",
@@ -1404,7 +1404,7 @@ def studio_default(
studio_python = _studio_venv_python()
run_py = _find_run_py()
if not (studio_python and run_py):
- typer.echo("Studio not set up. Run install.sh first.")
+ typer.echo("Unsloth Studio not set up. Run install.sh first.")
raise typer.Exit(1)
# A public UI launch must have a servable login page BEFORE the gate can
# strip the seeded .bootstrap_password, or the child has no way to change
@@ -1510,7 +1510,7 @@ def studio_default(
rc = proc.wait()
if rc != 0:
typer.echo(
- f"\nError: Studio server exited unexpectedly (code {rc}).",
+ f"\nError: Unsloth server exited unexpectedly (code {rc}).",
err = True,
)
typer.echo(
@@ -1522,7 +1522,7 @@ def studio_default(
else:
os.execvp(str(studio_python), args)
else:
- typer.echo("Studio not set up. Run install.sh first.")
+ typer.echo("Unsloth Studio not set up. Run install.sh first.")
raise typer.Exit(1)
run_mod = _load_run_module()
@@ -1733,7 +1733,7 @@ def run(
cloudflare: Optional[bool] = typer.Option(
None,
"--cloudflare/--no-cloudflare",
- help = "Expose Studio on a PUBLIC internet URL via a free Cloudflare HTTPS "
+ help = "Expose Unsloth on a PUBLIC internet URL via a free Cloudflare HTTPS "
"tunnel, for non-api-only wildcard binds (0.0.0.0 or ::). Off by default; "
"pass --cloudflare to enable it (--secure implies it). --no-cloudflare forces "
"it off but does not change a raw wildcard bind.",
@@ -1769,16 +1769,16 @@ def run(
"process list and shell history. Rotate later with `unsloth studio reset-password`.",
),
):
- """Start Studio, load a model, print an API key -- one-liner server.
+ """Start Unsloth, load a model, print an API key -- one-liner server.
- Unknown flags pass through to llama-server (GGUF only). Studio
+ Unknown flags pass through to llama-server (GGUF only). Unsloth
rejects managed flags with HTTP 400: model identity, network
(--host/--port/--path/--api-prefix/--reuse-port), auth/TLS
(--api-key/--ssl-*), single-model UI (--ui/--models-*/--webui),
and parallel slots (use --parallel above). Full denylist in
studio/backend/core/inference/llama_server_args.py. Other knobs
(-c, -ngl, --jinja, --flash-attn, -t, ...) pass through and
- last-wins-override Studio's auto-set value.
+ last-wins-override Unsloth's auto-set value.
Example:
unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL
@@ -1792,7 +1792,7 @@ def run(
# Set before any re-exec so the in-venv server inherits it via the env.
# `run --verbose` used to pass through to llama-server (its own -v); keep
- # that by forwarding --log-verbose so we add Studio logs without dropping it.
+ # that by forwarding --log-verbose so we add Unsloth logs without dropping it.
if verbose:
_enable_verbose_access_logs()
if not any(a in ("--verbose", "-v", "--log-verbose") for a in extra_llama_args):
@@ -1878,14 +1878,14 @@ def run(
if not in_studio_venv:
studio_python = _studio_venv_python()
if not studio_python:
- typer.echo("Studio not set up. Run install.sh first.")
+ typer.echo("Unsloth Studio not set up. Run install.sh first.")
raise typer.Exit(1)
# Re-exec via the studio venv's `unsloth` console-script.
studio_bin = studio_python.parent / "unsloth"
if not studio_bin.is_file():
- typer.echo("Studio venv missing 'unsloth' entry point. Re-run: unsloth studio setup")
+ typer.echo("Unsloth venv missing 'unsloth' entry point. Re-run: unsloth studio setup")
raise typer.Exit(1)
- # `run` serves the same Studio UI (unless --api-only); a public launch must
+ # `run` serves the same Unsloth UI (unless --api-only); a public launch must
# have a servable login page BEFORE the gate strips the seeded password, or
# the child has no way to change it. Validate here and forward the resolved
# dist so a shadowed child that can't self-resolve one still serves it.
@@ -2216,7 +2216,7 @@ def stop():
import signal as _signal
if not _PID_FILE.is_file():
- typer.echo("No running Studio server found (no PID file).")
+ typer.echo("No running Unsloth server found (no PID file).")
raise typer.Exit(0)
pid_text = _PID_FILE.read_text().strip()
@@ -2229,7 +2229,7 @@ def stop():
# Check if still alive (os.kill(pid, 0) is invalid on Windows -- see _pid_alive).
if not _pid_alive(pid):
- typer.echo(f"Studio server (PID {pid}) is not running. Cleaning up stale PID file.")
+ typer.echo(f"Unsloth server (PID {pid}) is not running. Cleaning up stale PID file.")
_PID_FILE.unlink(missing_ok = True)
raise typer.Exit(0)
@@ -2239,13 +2239,13 @@ def stop():
subprocess.run(["taskkill", "/PID", str(pid), "/F"], check = True)
else:
os.kill(pid, _signal.SIGTERM)
- typer.echo(f"Sent shutdown signal to Studio server (PID {pid}).")
+ typer.echo(f"Sent shutdown signal to Unsloth server (PID {pid}).")
except ProcessLookupError:
- typer.echo(f"Studio server (PID {pid}) already exited.")
+ typer.echo(f"Unsloth server (PID {pid}) already exited.")
_PID_FILE.unlink(missing_ok = True)
raise typer.Exit(0)
except Exception as e:
- typer.echo(f"Failed to stop Studio server (PID {pid}): {e}", err = True)
+ typer.echo(f"Failed to stop Unsloth server (PID {pid}): {e}", err = True)
raise typer.Exit(1)
# Wait briefly for the process to exit and clean up.
@@ -2253,10 +2253,10 @@ def stop():
time.sleep(0.5)
if not _pid_alive(pid):
_PID_FILE.unlink(missing_ok = True)
- typer.echo("Studio server stopped.")
+ typer.echo("Unsloth server stopped.")
raise typer.Exit(0)
- typer.echo("Studio server is shutting down (may take a few seconds).")
+ typer.echo("Unsloth server is shutting down (may take a few seconds).")
# ── unsloth studio setup / update ─────────────────────────────────────
@@ -2471,7 +2471,7 @@ def setup(
help = "Full pip/build output during setup for troubleshooting.",
),
):
- """Run Studio setup (called by install.ps1 / install.sh)."""
+ """Run Unsloth setup (called by install.ps1 / install.sh)."""
_run_setup_script(verbose = verbose)
@@ -2632,10 +2632,10 @@ def provision_desktop_auth():
@studio_app.command("reset-password")
def reset_password():
- """Reset the Studio admin password.
+ """Reset the Unsloth admin password.
Deletes the auth database so that a fresh admin account with a new
- random password is created on the next server start. The Studio
+ random password is created on the next server start. The Unsloth
server must be restarted after running this command.
"""
auth_dir = STUDIO_HOME / "auth"
@@ -2647,7 +2647,7 @@ def reset_password():
had_db = db_file.exists()
# Delete auth.db FIRST and prove it is gone before touching the seeded
- # credential files. If it cannot be removed (a running Studio or Windows
+ # credential files. If it cannot be removed (a running Unsloth or Windows
# holds it open, or a read-only auth dir), abort with the credential files
# untouched: deleting them while an un-resettable DB (must_change_password=1)
# survives would lock a forgotten-password reset out of any recovery
@@ -2657,7 +2657,7 @@ def reset_password():
except OSError as exc:
typer.echo(
f"Error: could not delete the auth database ({exc}). Stop any running "
- "Studio and retry; no credential files were changed.",
+ "Unsloth and retry; no credential files were changed.",
err = True,
)
raise typer.Exit(1)
@@ -2679,7 +2679,7 @@ def reset_password():
except OSError as exc:
typer.echo(
f"Error: could not remove or clear {path.name} ({exc}); delete "
- "it manually before restarting Studio or the old password may "
+ "it manually before restarting Unsloth or the old password may "
"be reused.",
err = True,
)
diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py
index 56633408fb..ae6f8dcfd4 100644
--- a/unsloth_cli/tests/test_inference_chat.py
+++ b/unsloth_cli/tests/test_inference_chat.py
@@ -372,7 +372,7 @@ def test_find_studio_server_none_when_not_running(monkeypatch):
def test_find_studio_server_prefers_ipv4_loopback_for_localhost(monkeypatch):
- # localhost resolving ::1-first must not hide a Studio bound to 127.0.0.1:
+ # localhost resolving ::1-first must not hide an Unsloth bound to 127.0.0.1:
# discovery tries each loopback address and returns the one that answers.
import socket
import urllib.request
diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py
index 7e76465144..2405ba0480 100644
--- a/unsloth_cli/tests/test_start.py
+++ b/unsloth_cli/tests/test_start.py
@@ -648,7 +648,7 @@ def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch
def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch):
- # Against a remote Studio the local existence probe cannot see server-side paths,
+ # Against a remote Unsloth the local existence probe cannot see server-side paths,
# so a case-variant loaded id must NOT attach without a load: it could be a distinct
# server-side path on a case-sensitive host. The load endpoint resolves the request.
calls = []
@@ -738,7 +738,7 @@ def test_no_launch_output_is_parseable(fake_studio):
result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"])
assert result.exit_code == 0, result.output
lines = [ln for ln in result.output.splitlines() if ln.strip()]
- skip = ("export ", "unset ", "Studio ", "Updated ", "Disabled ", "Warning", "Loading")
+ skip = ("export ", "unset ", "Unsloth ", "Updated ", "Disabled ", "Warning", "Loading")
body = [ln for ln in lines if not ln.startswith(skip)]
assert "codex --oss --profile unsloth_api" in body[-1]
assert any(ln.startswith("export CODEX_HOME=") for ln in lines)
@@ -767,7 +767,7 @@ def test_no_launch_last_line_is_self_contained(fake_studio, tmp_path):
def test_no_launch_claude_last_line_blanks_conflicting_auth(fake_studio):
# The unset vars must be neutralized inline too, or a partial copy would send the
- # user's own ANTHROPIC_API_KEY to the Studio base.
+ # user's own ANTHROPIC_API_KEY to the Unsloth base.
result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"])
assert result.exit_code == 0, result.output
last = [ln for ln in result.output.splitlines() if ln.strip()][-1]
@@ -814,7 +814,7 @@ def test_https_loopback_never_auto_serves(fake_studio, monkeypatch):
)
result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"])
assert result.exit_code == 1
- assert "No running Studio server" in result.output
+ assert "No running Unsloth server" in result.output
assert started["called"] is False
@@ -967,7 +967,7 @@ def test_connect_model_flag_forwards_load_options(fake_studio):
def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch):
- # Studio registers a loaded model under a canonical id (resolved identifier
+ # Unsloth registers a loaded model under a canonical id (resolved identifier
# / casing) that can differ from the path we passed. The agent must connect
# to that model, not silently fall through to the first loaded one.
requested = "Unsloth/Qwen3.5-35B-A3B"
@@ -1024,7 +1024,7 @@ def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio):
def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio):
# `--model repo:QUANT` splits into a VALID load payload (bare repo + gguf_variant),
- # never the `:`-suffixed repo id Studio rejects. The variant knob defers to
+ # never the `:`-suffixed repo id Unsloth rejects. The variant knob defers to
# /api/inference/load, whose already-loaded dedup answers without reloading when the
# active variant+settings match -- so a second session running the same command
# attaches without evicting the first, while a genuinely different quant reloads.
@@ -1116,7 +1116,7 @@ def test_connect_no_model_loaded_errors(fake_studio, monkeypatch):
def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch):
- # Studio never surfaces the requested model; fail loudly rather than
+ # Unsloth never surfaces the requested model; fail loudly rather than
# silently connecting to whatever else happens to be loaded.
inner = start._http_json
@@ -1187,7 +1187,7 @@ def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch):
def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch):
- # A key saved for a remote (non-loopback) Studio is replayed on keyless runs;
+ # A key saved for a remote (non-loopback) Unsloth is replayed on keyless runs;
# auto-minting stays blocked for non-loopback.
remote = "http://studio.example:8888"
monkeypatch.setattr(start, "find_studio_server", lambda: remote)
@@ -1201,7 +1201,7 @@ def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatc
def test_connect_studio_server_errors_on_explicit_remote(monkeypatch):
- # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an
+ # A user who pointed UNSLOTH_STUDIO_URL at a remote Unsloth should get an
# error, not a silent local model load (which they did not ask for).
import typer
@@ -1242,7 +1242,7 @@ def test_connect_unverified_loopback_without_cached_key_refuses_to_mint(
def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch):
- # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match)
+ # A "saved" key (e.g. for an SSH-tunnelled Unsloth the handshake can't match)
# replays on keyless runs without the handshake, scoped to its own base.
cache = tmp_path / "agent_api_key.json"
cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}))
@@ -1358,7 +1358,7 @@ def _serve_redirect(target):
def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch):
- # A squatter could 302 /api/auth/identity to the real Studio and relay its
+ # A squatter could 302 /api/auth/identity to the real Unsloth and relay its
# proof; redirects must be refused so the squatter's base isn't accepted.
import unsloth_cli._inference as inference
@@ -1384,7 +1384,7 @@ def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch):
def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch):
- # A squatter that proxies the nonce to the real Studio on another port gets a
+ # A squatter that proxies the nonce to the real Unsloth on another port gets a
# proof bound to *that* port; the client expects one bound to the port it
# connected to, so the relayed proof is rejected.
import unsloth_cli._inference as inference
@@ -1435,7 +1435,7 @@ def test_connect_no_studio_errors(fake_studio, monkeypatch):
monkeypatch.setattr(start, "find_studio_server", lambda: None)
result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"])
assert result.exit_code == 1
- assert "No running Studio server" in result.output
+ assert "No running Unsloth server" in result.output
@pytest.fixture(autouse = True)
@@ -1564,7 +1564,7 @@ def test_no_serve_preserves_error(fake_studio, monkeypatch):
start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-serve"]
)
assert result.exit_code == 1
- assert "No running Studio server" in result.output
+ assert "No running Unsloth server" in result.output
assert started["called"] is False
@@ -1578,7 +1578,7 @@ def test_no_launch_never_serves(fake_studio, monkeypatch):
start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-launch"]
)
assert result.exit_code == 1
- assert "No running Studio server" in result.output
+ assert "No running Unsloth server" in result.output
assert started["called"] is False
@@ -1871,7 +1871,7 @@ def _opencode_inline_config(output: str) -> dict:
def test_opencode_inline_scopes_session_to_studio_provider(fake_studio):
# opencode filters even config-defined providers through enabled/disabled_providers,
# and a model pin does not bypass that gate. The inline overlay (session-only, highest
- # layer, arrays replace) allowlists our provider and clears the denylist so the Studio
+ # layer, arrays replace) allowlists our provider and clears the denylist so the Unsloth
# model always loads regardless of the user's config, without reading or editing it.
result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"])
assert result.exit_code == 0, result.output
@@ -2657,7 +2657,7 @@ def test_agent_api_key_auto_started_rejected_env_key_falls_back(fake_studio, tmp
def test_agent_api_key_auto_started_accepted_key_is_honored(fake_studio, tmp_path):
- # An explicit key the fresh server accepts (e.g. persisted in this Studio
+ # An explicit key the fresh server accepts (e.g. persisted in this Unsloth
# home's auth db across restarts) keeps working exactly as before.
key = start._agent_api_key(BASE, "sk-unsloth-deadbeefdeadbeef", auto_started = True)
assert key == "sk-unsloth-deadbeefdeadbeef"
@@ -2937,11 +2937,11 @@ def test_hermes_resume_oneshot_rejects_usage_file(monkeypatch, usage_arg):
def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch):
# The persistence flag is --persist, NOT --resume, so an agent's own
# `--resume ` (e.g. `unsloth start claude --resume `) still flows
- # through to the agent verbatim and is not swallowed as a Studio option.
+ # through to the agent verbatim and is not swallowed as an Unsloth option.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
monkeypatch.setattr(start, "_claude_flags", lambda: [])
captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"])
assert captured["command"][-2:] == ["--resume", "some-session-guid"]
- # Studio never auto-appends its own resume token when the user drives resume.
+ # Unsloth never auto-appends its own resume token when the user drives resume.
assert captured["command"].count("--resume") == 1
assert "--continue" not in captured["command"]
diff --git a/unsloth_cli/tests/test_studio_cloudflare_flag.py b/unsloth_cli/tests/test_studio_cloudflare_flag.py
index fb57d7aaf4..7287737f75 100644
--- a/unsloth_cli/tests/test_studio_cloudflare_flag.py
+++ b/unsloth_cli/tests/test_studio_cloudflare_flag.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
-"""Tests for the `--cloudflare/--no-cloudflare` Studio flag.
+"""Tests for the `--cloudflare/--no-cloudflare` Unsloth flag.
Pins the typer Option (tri-state, default off / None) on both `unsloth studio`
and `unsloth studio run`, and that the chosen polarity reaches the re-exec'd
diff --git a/unsloth_cli/tests/test_studio_password_prompt.py b/unsloth_cli/tests/test_studio_password_prompt.py
index 7bbdfe3703..6e9a2c1d52 100644
--- a/unsloth_cli/tests/test_studio_password_prompt.py
+++ b/unsloth_cli/tests/test_studio_password_prompt.py
@@ -885,7 +885,7 @@ def test_run_non_tty_deletes_bootstrap_password_file(monkeypatch, tmp_path):
def test_run_missing_frontend_exits_before_stripping_bootstrap(monkeypatch, tmp_path):
# Regression (item B / reviewer finding 4): `unsloth studio run` serves the
- # same Studio UI and strips the seeded password on a headless public launch,
+ # same Unsloth UI and strips the seeded password on a headless public launch,
# so a missing frontend dist must abort BEFORE the strip -- the same lockout
# guard as `unsloth studio`, not just `studio run`'s model-load residual.
import typer as _typer
@@ -1104,7 +1104,7 @@ def test_cli_update_password_truncates_locked_bootstrap_after_change(monkeypatch
def test_reset_password_fails_closed_when_db_cannot_be_deleted(monkeypatch, tmp_path):
- # If auth.db cannot be removed (running Studio / Windows lock, read-only dir),
+ # If auth.db cannot be removed (running Unsloth / Windows lock, read-only dir),
# reset must abort BEFORE touching the credential files -- deleting them while
# an un-resettable must_change_password=1 DB survives would lock a
# forgotten-password reset out with no recovery credential.
diff --git a/unsloth_cli/tests/test_studio_secure_flag.py b/unsloth_cli/tests/test_studio_secure_flag.py
index 5e5895309c..2a67aad95a 100644
--- a/unsloth_cli/tests/test_studio_secure_flag.py
+++ b/unsloth_cli/tests/test_studio_secure_flag.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
-"""Tests for the `--secure/--no-secure` Studio flag: option registration,
+"""Tests for the `--secure/--no-secure` Unsloth flag: option registration,
re-exec/run_server forwarding, the forced 127.0.0.1 bind, and rejection
alongside --no-cloudflare or before a subcommand. Modeled on
test_studio_cloudflare_flag.py."""
diff --git a/unsloth_cli/tests/test_studio_verbose_flag.py b/unsloth_cli/tests/test_studio_verbose_flag.py
index 4af32fd4a2..20468b5f02 100644
--- a/unsloth_cli/tests/test_studio_verbose_flag.py
+++ b/unsloth_cli/tests/test_studio_verbose_flag.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
-"""Tests for the `--verbose/-v` Studio flag: option registration on both the
+"""Tests for the `--verbose/-v` Unsloth flag: option registration on both the
plain callback and the `run` subcommand, re-exec forwarding, the access-log
env override, and rejection before a subcommand. Modeled on
test_studio_secure_flag.py."""
@@ -123,7 +123,7 @@ def test_run_without_verbose_leaves_env_unset(monkeypatch):
def test_run_verbose_preserves_llama_server_verbosity(monkeypatch):
- # Studio consumes --verbose but still forwards llama-server's own verbosity.
+ # Unsloth consumes --verbose but still forwards llama-server's own verbosity.
monkeypatch.delenv(_DEDUP, raising = False)
monkeypatch.delenv(_POLL, raising = False)
captured = _invoke_run(monkeypatch, _BASE + ["--verbose"])
From 74d1a284ebe2fcb7ee0123e0a47b9f4bac8a7690 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Sun, 19 Jul 2026 03:20:56 -0700
Subject: [PATCH 012/255] Studio: hide the RAG embedder and llama.cpp probe
from the hub cached inventory (#7018)
* Studio: hide infra models from the hub cached inventory
The hub inventory scans behind /api/hub/cached-gguf and /api/hub/cached-models
returned the llama.cpp install validation probe (ggml-org/models) and the RAG
embedder (unsloth/bge-small-en-v1.5[-GGUF]) as on-device models. Share the
hidden-model check from routes/models.py via utils/models/hidden_models.py and
apply it in both scans. A GGUF infra repo stays visible when the user
explicitly downloaded a variant through the Hub, since variant manifests only
exist for user-initiated downloads.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: make On Device trust the hub inventory, match repo ids exactly, lighten the hidden-model import
Follow-up on the hub cached-inventory hidden-model change, addressing the review.
On Device now trusts the Hub inventory API for cached rows. The backend already
hides the RAG embedder and the llama.cpp probe and re-includes a GGUF infra repo
once the user downloads a variant through the Hub, but the frontend was
re-hiding it by repo id, so the user-downloaded variant never appeared in the On
Device list or the count. isVisibleInventoryRow now short-circuits cached rows
(kind === "cache") to visible and keeps client-side needle hiding only for local
filesystem rows and Discover.
is_hidden_model matches Hub repo ids exactly (case-insensitive) against the probe
plus the effective embedder and its GGUF companion, instead of substring
matching the configured-embedder basename. A custom embedder with a generic
basename like org/model no longer hides unrelated cached repos such as
user/model-chat or org/model-instruct. The probe filename and local-path
embedders keep exact matching.
The helper moves to utils/hidden_models.py and is imported at module scope in the
hub cache scanner, so it no longer pulls in utils/models/__init__ (the eager
model-config/checkpoint stack) and a broken import fails at startup instead of
being swallowed per-repo and silently emptying the inventory. routes.models
keeps the _is_hidden_model and _safe_resolve aliases and drops the unused
_HF_REPO_ID_RE re-export that was failing source lint.
Tests: exact repo-id matching with a custom embedder, the cached-models scan
keeping an unrelated repo, and a clean-interpreter check that the helper imports
without the model-config stack.
* Studio: match the llama.cpp probe filename on both path separators
The hidden-model check compared the probe's on-disk filename with
Path(value).name, which on a POSIX interpreter does not split a Windows-style
path ("...\stories260K.gguf") and would let the probe through. Split on both
separators so the probe is matched regardless of which OS produced the path,
matching the tolerance of the previous substring check. Adds a Windows-path
assertion to the probe test.
* Studio: harden hidden infra model handling
* Fix hidden cache row confirmation
* Fix hidden local rows and confirmed hint merges
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Handle snapshot-configured hidden models
* Hide basename-only default embedders
* Fix dynamic embedder inventory filtering
* Studio: hide the configured RAG embedder from Discover and feed rows
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen
Co-authored-by: Daniel Han <23090290+danielhanchen@users.noreply.github.com>
---
.../core/inference/local_model_resolver.py | 6 +-
studio/backend/core/rag/config.py | 25 +-
.../hub/services/models/cache_inventory.py | 40 ++-
.../hub/services/models/local_inventory.py | 17 +-
.../backend/hub/tests/test_model_services.py | 338 ++++++++++++++++++
studio/backend/routes/models.py | 61 +---
studio/backend/routes/settings.py | 5 +
.../backend/tests/test_cached_gguf_routes.py | 139 +++++++
.../test_embedding_model_security_gate.py | 17 +
.../tests/test_embedding_model_settings.py | 7 +
.../backend/tests/test_openai_auto_switch.py | 27 +-
studio/backend/utils/hidden_models.py | 142 ++++++++
.../hub/hooks/use-hidden-embedding-models.ts | 46 +++
studio/frontend/src/features/hub/hub-page.tsx | 53 ++-
studio/frontend/src/features/hub/index.ts | 1 +
.../features/hub/inventory/inventory-hints.ts | 11 +-
.../src/features/hub/inventory/types.ts | 1 +
.../hub/inventory/use-hub-inventory.ts | 1 +
.../src/features/hub/inventory/view-models.ts | 18 +-
.../src/features/hub/lib/hidden-models.ts | 26 +-
.../features/settings/api/embedding-model.ts | 17 +-
.../frontend/src/features/settings/index.ts | 1 +
22 files changed, 899 insertions(+), 100 deletions(-)
create mode 100644 studio/backend/utils/hidden_models.py
create mode 100644 studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts
diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py
index 86ad8b9fd8..64ab38ec75 100644
--- a/studio/backend/core/inference/local_model_resolver.py
+++ b/studio/backend/core/inference/local_model_resolver.py
@@ -201,7 +201,11 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
continue
# Skip what Unsloth hides from its pickers (validation probe, RAG embed
# weights): not chat models, so never an auto-switch target.
- if _is_hidden_model(raw_id, getattr(info, "path", None)):
+ if _is_hidden_model(
+ raw_id,
+ getattr(info, "model_id", None),
+ getattr(info, "path", None),
+ ):
continue
# Advertise a client-facing alias, not an absolute filesystem path.
loader_id = _advertised_loader_id(info)
diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py
index 2de32a68e4..f54d795731 100644
--- a/studio/backend/core/rag/config.py
+++ b/studio/backend/core/rag/config.py
@@ -87,6 +87,22 @@ def _names_gguf(model: str) -> bool:
return "gguf" in re.split(r"[^a-z0-9]+", model.lower())
+def gguf_repo_for_embedding_model(model: str) -> str:
+ """GGUF repo for ``model``, honoring an explicit companion override."""
+ if "RAG_EMBED_GGUF_REPO" in os.environ:
+ return EMBED_GGUF_REPO
+ if model == DEFAULT_EMBEDDING_MODEL:
+ return EMBED_GGUF_REPO
+ if _names_gguf(model):
+ return model
+ return f"{model}-GGUF"
+
+
+def default_gguf_repo() -> str:
+ """GGUF companion for the env/default embedding model."""
+ return gguf_repo_for_embedding_model(EMBEDDING_MODEL)
+
+
def effective_gguf_repo() -> str:
"""GGUF repo for the llama-server backend, tracking the effective model.
@@ -95,14 +111,7 @@ def effective_gguf_repo() -> str:
``-GGUF`` companion repo (the unsloth convention the default pair follows),
or is used as-is when it already names a GGUF repo.
"""
- if "RAG_EMBED_GGUF_REPO" in os.environ:
- return EMBED_GGUF_REPO
- model = effective_embedding_model()
- if model == DEFAULT_EMBEDDING_MODEL:
- return EMBED_GGUF_REPO
- if _names_gguf(model):
- return model
- return f"{model}-GGUF"
+ return gguf_repo_for_embedding_model(effective_embedding_model())
# llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this
diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py
index 1f38af9381..54a25482f2 100644
--- a/studio/backend/hub/services/models/cache_inventory.py
+++ b/studio/backend/hub/services/models/cache_inventory.py
@@ -37,6 +37,13 @@ from hub.services.models.common import (
_runtime_for_format,
)
+# Imported at module scope (not inside the per-repo scan loop) so a broken
+# import surfaces at startup instead of silently emptying the inventory: the
+# scan loop swallows per-repo exceptions and would drop every repo. Lives under
+# ``utils`` (not ``utils.models``) to avoid the eager model-config/checkpoint
+# imports in ``utils/models/__init__.py``.
+from utils.hidden_models import is_hidden_model
+
logger = get_logger(__name__)
_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = (
@@ -243,6 +250,13 @@ def invalidate_hf_cache_scans() -> None:
hf_cache_scan.invalidate_hf_cache_scans()
+def _is_hidden_infra_repo(*values: str | None) -> bool:
+ """True for infra-only repos (the RAG embedder and the llama.cpp install
+ validation probe) that are cached as a side effect of Studio itself and are
+ not usable chat models."""
+ return is_hidden_model(*values)
+
+
def _scan_cached_gguf() -> list[dict]:
"""Synchronous HF-cache disk walk for GGUF repos; runs in a worker thread."""
cache_scans = all_hf_cache_scans()
@@ -254,13 +268,24 @@ def _scan_cached_gguf() -> list[dict]:
if str(repo_info.repo_type) != "model":
continue
repo_id = repo_info.repo_id
+ repo_path = Path(repo_info.repo_path)
+ snapshot_path = _cached_model_snapshot_path(repo_path)
total_size = _repo_gguf_size_bytes(repo_info)
has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id)
+ is_hidden_infra = _is_hidden_infra_repo(
+ repo_id,
+ str(repo_path),
+ str(snapshot_path) if snapshot_path is not None else None,
+ )
+ # Hide infra repos unless the user downloaded a variant via
+ # the Hub; variant state only exists for user downloads.
+ if is_hidden_infra and not has_variant_state:
+ continue
if total_size == 0 and not has_variant_state:
continue
partial = hf_cache_scan.is_gguf_repo_partial(
repo_id,
- Path(repo_info.repo_path),
+ repo_path,
)
if total_size == 0 and not partial:
continue
@@ -283,6 +308,9 @@ def _scan_cached_gguf() -> list[dict]:
requires_variant = True,
)
)
+ # Visible infra variants remain management-only.
+ if is_hidden_infra:
+ row["capabilities"]["can_chat"] = False
if _prefer_cache_row(row, existing):
seen_lower[key] = row
except Exception as e:
@@ -475,6 +503,15 @@ def _scan_cached_models() -> list[dict]:
if str(repo_info.repo_type) != "model":
continue
repo_id = repo_info.repo_id
+ repo_path = Path(repo_info.repo_path)
+ snapshot_path = _cached_model_snapshot_path(repo_path)
+ # The non-GGUF embedder has no variant downloads; always hide.
+ if _is_hidden_infra_repo(
+ repo_id,
+ str(repo_path),
+ str(snapshot_path) if snapshot_path is not None else None,
+ ):
+ continue
has_main_gguf = _repo_has_gguf_files(repo_info)
payload = _repo_non_gguf_model_payload(repo_info)
if payload.size_bytes == 0:
@@ -486,7 +523,6 @@ def _scan_cached_models() -> list[dict]:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
- repo_path = Path(repo_info.repo_path)
snapshot_partial = hf_cache_scan.is_snapshot_partial(
"model",
repo_id,
diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py
index a3782efead..b34532fa35 100644
--- a/studio/backend/hub/services/models/local_inventory.py
+++ b/studio/backend/hub/services/models/local_inventory.py
@@ -36,6 +36,7 @@ from hub.utils.paths import (
)
from hub.services.models import common as model_common
from hub.services.models.ollama import scan_ollama_dir
+from utils.hidden_models import is_hidden_model
logger = get_logger(__name__)
_MAX_MODELS_PER_CUSTOM_FOLDER = 200
@@ -623,6 +624,20 @@ def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelI
)
+def _filter_hidden_models(local_models: List[LocalModelInfo]) -> list[LocalModelInfo]:
+ """Remove infrastructure-only models from the shared local inventory."""
+ visible: list[LocalModelInfo] = []
+ for model in local_models:
+ resolved_cache_path = (
+ hf_cache_scan.resolve_hf_cache_realpath(Path(model.path))
+ if model.source == "hf_cache"
+ else None
+ )
+ if not is_hidden_model(model.id, model.model_id, model.path, resolved_cache_path):
+ visible.append(model)
+ return visible
+
+
async def list_local_models_response(models_dir: str = "./models") -> LocalModelListResponse:
"""List local model candidates from every supported on-device source."""
hf_cache_dir = _resolve_hf_cache_dir()
@@ -653,7 +668,7 @@ async def list_local_models_response(models_dir: str = "./models") -> LocalModel
ollama_dirs,
)
local_models += await _collect_models_from_custom_folders()
- models = _dedupe_local_models(local_models)
+ models = _dedupe_local_models(_filter_hidden_models(local_models))
return LocalModelListResponse(
models_dir = str(models_root),
diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py
index 2c33e09b2b..693d945ee1 100644
--- a/studio/backend/hub/tests/test_model_services.py
+++ b/studio/backend/hub/tests/test_model_services.py
@@ -439,6 +439,287 @@ def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypa
assert row["capabilities"]["requires_variant"] is True
+def test_cached_gguf_scan_hides_infra_repos_without_user_downloads(monkeypatch, tmp_path):
+ probe = _repo(
+ "ggml-org/models",
+ [_file("tinyllamas/stories260K.gguf", 1_200_000)],
+ tmp_path / "probe",
+ )
+ embedder = _repo(
+ "unsloth/bge-small-en-v1.5-GGUF",
+ [_file("bge-small-en-v1.5-f16.gguf", 60_000_000)],
+ tmp_path / "embedder",
+ )
+ chat = _repo("Org/Chat-GGUF", [_file("Q4_K_M.gguf", 100)], tmp_path / "chat")
+ monkeypatch.setattr(
+ cache_inventory,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [probe, embedder, chat])],
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_gguf_repo_partial",
+ lambda _repo_id, _path: False,
+ )
+
+ result = {"cached": cache_inventory._scan_cached_gguf()}
+
+ assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat-GGUF"]
+
+
+def test_cached_gguf_scan_keeps_infra_repo_with_user_downloaded_variant(monkeypatch, tmp_path):
+ monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
+ embedder = _repo(
+ "unsloth/bge-small-en-v1.5-GGUF",
+ [
+ _file("bge-small-en-v1.5-f16.gguf", 60_000_000),
+ _file("bge-small-en-v1.5-Q8_0.gguf", 35_000_000),
+ ],
+ tmp_path / "embedder",
+ )
+ # Variant manifests only exist for user Hub downloads, not auto-downloads.
+ assert download_manifest.write_manifest(
+ "model",
+ "unsloth/bge-small-en-v1.5-GGUF",
+ "Q8_0",
+ [download_manifest.ExpectedFile(path = "bge-small-en-v1.5-Q8_0.gguf", size = 35_000_000)],
+ "http",
+ )
+ monkeypatch.setattr(
+ cache_inventory,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [embedder])],
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_gguf_repo_partial",
+ lambda _repo_id, _path: False,
+ )
+
+ result = {"cached": cache_inventory._scan_cached_gguf()}
+
+ assert [row["repo_id"] for row in result["cached"]] == ["unsloth/bge-small-en-v1.5-GGUF"]
+ assert result["cached"][0]["capabilities"]["can_chat"] is False
+
+
+def test_cached_models_scan_hides_non_gguf_embedder(monkeypatch, tmp_path):
+ embedder_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5"
+ embedder_path.mkdir(parents = True)
+ embedder = _repo(
+ "unsloth/bge-small-en-v1.5",
+ [_file("config.json", 12), _file("model.safetensors", 130_000_000)],
+ embedder_path,
+ )
+ chat_path = tmp_path / "hub" / "models--Org--Chat"
+ chat_path.mkdir(parents = True)
+ chat = _repo(
+ "Org/Chat",
+ [_file("config.json", 12), _file("model.safetensors", 100)],
+ chat_path,
+ )
+ monkeypatch.setattr(
+ cache_inventory,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [embedder, chat])],
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_snapshot_partial",
+ lambda _kind, _repo_id, _path: False,
+ )
+
+ result = {"cached": cache_inventory._scan_cached_models()}
+
+ assert [row["repo_id"] for row in result["cached"]] == ["Org/Chat"]
+
+
+def test_cached_scans_hide_embedders_configured_by_cache_path(monkeypatch, tmp_path):
+ from core.rag import config as rag_config
+
+ gguf_path = tmp_path / "hub" / "models--Org--PathEmbedder-GGUF"
+ gguf_path.mkdir(parents = True)
+ gguf = _repo(
+ "Org/PathEmbedder-GGUF",
+ [_file("model-F16.gguf", 60_000_000)],
+ gguf_path,
+ )
+ model_path = tmp_path / "hub" / "models--Org--PathEmbedder"
+ model_path.mkdir(parents = True)
+ model = _repo(
+ "Org/PathEmbedder",
+ [_file("config.json", 12), _file("model.safetensors", 130_000_000)],
+ model_path,
+ )
+ monkeypatch.setattr(
+ rag_config,
+ "effective_embedding_model",
+ lambda: str(model_path),
+ )
+ monkeypatch.setattr(
+ rag_config,
+ "effective_gguf_repo",
+ lambda: str(gguf_path),
+ )
+ monkeypatch.setattr(
+ cache_inventory,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [gguf, model])],
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_gguf_repo_partial",
+ lambda _repo_id, _path: False,
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_snapshot_partial",
+ lambda _kind, _repo_id, _path: False,
+ )
+
+ assert cache_inventory._scan_cached_gguf() == []
+ assert cache_inventory._scan_cached_models() == []
+
+
+def test_cached_scans_hide_embedders_configured_by_snapshot_path(monkeypatch, tmp_path):
+ from core.rag import config as rag_config
+
+ gguf_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder-GGUF"
+ gguf_snapshot = gguf_path / "snapshots" / "gguf-revision"
+ gguf_snapshot.mkdir(parents = True)
+ gguf = _repo(
+ "Org/SnapshotEmbedder-GGUF",
+ [_file("model-F16.gguf", 60_000_000)],
+ gguf_path,
+ )
+ model_path = tmp_path / "hub" / "models--Org--SnapshotEmbedder"
+ model_snapshot = model_path / "snapshots" / "model-revision"
+ model_snapshot.mkdir(parents = True)
+ model = _repo(
+ "Org/SnapshotEmbedder",
+ [_file("config.json", 12), _file("model.safetensors", 130_000_000)],
+ model_path,
+ )
+ monkeypatch.setattr(
+ rag_config,
+ "effective_embedding_model",
+ lambda: str(model_snapshot),
+ )
+ monkeypatch.setattr(
+ rag_config,
+ "effective_gguf_repo",
+ lambda: str(gguf_snapshot),
+ )
+ monkeypatch.setattr(
+ cache_inventory,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [gguf, model])],
+ )
+
+ def _resolve_snapshot(repo_path):
+ return str(
+ {
+ gguf_path: gguf_snapshot,
+ model_path: model_snapshot,
+ }.get(Path(repo_path), Path(repo_path))
+ )
+
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "resolve_hf_cache_realpath",
+ _resolve_snapshot,
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_gguf_repo_partial",
+ lambda _repo_id, _path: False,
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_snapshot_partial",
+ lambda _kind, _repo_id, _path: False,
+ )
+
+ assert cache_inventory._scan_cached_gguf() == []
+ assert cache_inventory._scan_cached_models() == []
+
+
+def test_cached_models_scan_keeps_unrelated_repo_with_custom_generic_embedder(
+ monkeypatch, tmp_path
+):
+ # A custom embedder with a generic basename ("org/model") must be hidden by
+ # EXACT repo-id match only. An unrelated cached chat model whose id merely
+ # contains "model" (e.g. "user/model-chat") must stay on device: substring
+ # basename matching used to drop real chat models from the inventory.
+ from core.rag import config as rag_config
+
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
+
+ def _model_repo(repo_id: str):
+ path = tmp_path / "hub" / f"models--{repo_id.replace('/', '--')}"
+ path.mkdir(parents = True)
+ return _repo(
+ repo_id,
+ [_file("config.json", 12), _file("model.safetensors", 100)],
+ path,
+ )
+
+ embedder = _model_repo("org/model")
+ chat = _model_repo("user/model-chat")
+ monkeypatch.setattr(
+ cache_inventory,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [embedder, chat])],
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_snapshot_partial",
+ lambda _kind, _repo_id, _path: False,
+ )
+
+ result = {"cached": cache_inventory._scan_cached_models()}
+
+ assert [row["repo_id"] for row in result["cached"]] == ["user/model-chat"]
+
+
+def test_cached_scans_hide_stale_default_embedder_after_custom_setting(monkeypatch, tmp_path):
+ from core.rag import config as rag_config
+
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF")
+
+ gguf = _repo(
+ "unsloth/bge-small-en-v1.5-GGUF",
+ [_file("bge-small-en-v1.5-f16.gguf", 60_000_000)],
+ tmp_path / "default-gguf",
+ )
+ weights_path = tmp_path / "hub" / "models--unsloth--bge-small-en-v1.5"
+ weights_path.mkdir(parents = True)
+ weights = _repo(
+ "unsloth/bge-small-en-v1.5",
+ [_file("config.json", 12), _file("model.safetensors", 130_000_000)],
+ weights_path,
+ )
+ monkeypatch.setattr(
+ cache_inventory,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [gguf, weights])],
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_gguf_repo_partial",
+ lambda _repo_id, _path: False,
+ )
+ monkeypatch.setattr(
+ cache_inventory.hf_cache_scan,
+ "is_snapshot_partial",
+ lambda _kind, _repo_id, _path: False,
+ )
+
+ assert cache_inventory._scan_cached_gguf() == []
+ assert cache_inventory._scan_cached_models() == []
+
+
def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj():
requirements = gguf_variants._build_gguf_variant_requirements(
[
@@ -1610,6 +1891,63 @@ def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_
assert rows[0].capabilities.requires_variant is True
+def test_local_inventory_filters_custom_embedder_hf_cache_row(monkeypatch, tmp_path):
+ from core.rag import config as rag_config
+
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/embedder")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF")
+
+ def _row(repo_id: str):
+ repo_path = tmp_path / f"models--{repo_id.replace('/', '--')}"
+ return model_common._local_model_info(
+ scan_path = repo_path,
+ load_path = repo_path,
+ source = "hf_cache",
+ model_format = "safetensors",
+ model_id = repo_id,
+ )
+
+ rows = local_inventory._filter_hidden_models([_row("org/embedder"), _row("org/chat-model")])
+
+ assert [row.model_id for row in rows] == ["org/chat-model"]
+
+
+def test_local_inventory_filters_embedder_configured_by_snapshot_path(monkeypatch, tmp_path):
+ from core.rag import config as rag_config
+
+ embedder_path = tmp_path / "hub" / "models--org--embedder"
+ embedder_snapshot = embedder_path / "snapshots" / "revision"
+ embedder_snapshot.mkdir(parents = True)
+ chat_path = tmp_path / "hub" / "models--org--chat-model"
+ chat_path.mkdir(parents = True)
+ monkeypatch.setattr(
+ rag_config,
+ "effective_embedding_model",
+ lambda: str(embedder_snapshot),
+ )
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF")
+ monkeypatch.setattr(
+ local_inventory.hf_cache_scan,
+ "resolve_hf_cache_realpath",
+ lambda path: str(embedder_snapshot) if Path(path) == embedder_path else str(path),
+ )
+
+ def _row(repo_id: str, repo_path: Path):
+ return model_common._local_model_info(
+ scan_path = repo_path,
+ load_path = repo_path,
+ source = "hf_cache",
+ model_format = "safetensors",
+ model_id = repo_id,
+ )
+
+ rows = local_inventory._filter_hidden_models(
+ [_row("org/embedder", embedder_path), _row("org/chat-model", chat_path)]
+ )
+
+ assert [row.model_id for row in rows] == ["org/chat-model"]
+
+
def test_model_download_job_helpers_preserve_idle_shape():
key = downloads._download_job_key("Org/Model", None)
status = downloads._job_status(key)
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 742ecde3ba..bb321695cd 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -59,59 +59,12 @@ def _safe_is_dir(path) -> bool:
return False
-# Hub repo id shape ("owner/name", no leading separator); anything else is
-# treated as a local filesystem path.
-_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$")
-
-
-def _is_hidden_model(*values: str | None) -> bool:
- """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or
- EMBED_GGUF_REPO basename) or the llama.cpp install validation probe
- (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF).
- None are usable chat models; the probe can be cached as a side effect of
- installing the prebuilt llama-server and otherwise sorts smallest, so it
- would be auto-selected. A local-path embedder is matched by exact resolved
- path only: a generic basename like "model" must not substring-hide
- unrelated chat models."""
- from core.rag import config as rag_config
-
- needles = [
- # The validation probe's repo (matches the cached repo id) and its exact
- # filename (matches the on-disk path). The filename carries the .gguf so
- # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``.
- "ggml-org/models",
- "stories260k.gguf",
- ]
- exact_paths: list[str] = []
- for model in (
- rag_config.effective_embedding_model(),
- rag_config.effective_gguf_repo(),
- ):
- if _HF_REPO_ID_RE.match(model):
- needles.append(model.split("/")[-1].lower())
- else:
- resolved = _safe_resolve(Path(model).expanduser())
- if resolved:
- exact_paths.append(resolved.lower())
- for v in values:
- if not v:
- continue
- low = v.lower()
- if any(n in low for n in needles):
- return True
- if exact_paths:
- resolved = _safe_resolve(Path(v).expanduser())
- if resolved and resolved.lower() in exact_paths:
- return True
- return False
-
-
-def _safe_resolve(path: Path) -> Optional[str]:
- """resolve() to a string, or None when the path is inaccessible."""
- try:
- return str(path.resolve())
- except OSError:
- return None
+# Shared with the hub inventory scans; keep the private aliases so existing
+# importers (core.inference.local_model_resolver, tests) stay valid.
+from utils.hidden_models import (
+ _safe_resolve,
+ is_hidden_model as _is_hidden_model,
+)
backend_path = Path(__file__).parent.parent.parent
@@ -853,7 +806,7 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
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)]
+ return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)]
@router.get("/local", response_model = LocalModelListResponse)
diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py
index 1ddfc0eacb..ab0fd2fd99 100644
--- a/studio/backend/routes/settings.py
+++ b/studio/backend/routes/settings.py
@@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
from auth.authentication import get_current_subject
from auth.storage import rotate_preview_link_secret
+from core.rag.config import default_gguf_repo, effective_gguf_repo
from loggers import get_logger
from utils.utils import safe_error_detail, log_and_http_error
from utils.personalization_settings import (
@@ -263,14 +264,18 @@ class EmbeddingModelPayload(BaseModel):
class EmbeddingModelResponse(BaseModel):
embedding_model: str
+ embedding_gguf_repo: str
default_embedding_model: str
+ default_embedding_gguf_repo: str
is_custom: bool
def _embedding_model_response() -> EmbeddingModelResponse:
return EmbeddingModelResponse(
embedding_model = get_rag_embedding_model(),
+ embedding_gguf_repo = effective_gguf_repo(),
default_embedding_model = default_embedding_model(),
+ default_embedding_gguf_repo = default_gguf_repo(),
is_custom = get_stored_embedding_model() is not None,
)
diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py
index d4a7cae208..b3e6255d55 100644
--- a/studio/backend/tests/test_cached_gguf_routes.py
+++ b/studio/backend/tests/test_cached_gguf_routes.py
@@ -120,12 +120,151 @@ def test_is_hidden_model_hides_validation_probe_everywhere():
assert models_route._is_hidden_model(
None, "/hf/models--ggml-org--models/snapshots/abc/tinyllamas/stories260K.gguf"
)
+ # A Windows-style snapshot path must match too, even on a POSIX interpreter
+ # (the filename check splits on both separators).
+ assert models_route._is_hidden_model(
+ r"C:\Users\u\.cache\huggingface\hub\models--ggml-org--models\snapshots\abc\tinyllamas\stories260K.gguf"
+ )
assert not models_route._is_hidden_model("unsloth/gemma-3-270m-it-GGUF")
# The exact-filename needle must not hide a real repo that merely
# references stories260K in its name.
assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF")
+def test_is_hidden_model_matches_repo_ids_exactly(monkeypatch):
+ """A custom embedder with a generic basename is hidden by EXACT repo-id
+ match only, so unrelated cached repos that merely contain the basename stay
+ visible. Regression: substring basename matching hid real chat models like
+ ``user/model-chat`` from the On Device inventory."""
+ from core.rag import config as rag_config
+
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
+
+ # The exact embedder repo and its GGUF companion are hidden.
+ assert models_route._is_hidden_model("org/model")
+ assert models_route._is_hidden_model("org/model-GGUF")
+ # Unrelated repos that merely contain "model" must NOT be hidden.
+ assert not models_route._is_hidden_model("user/model-chat")
+ assert not models_route._is_hidden_model("org/model-instruct")
+ assert not models_route._is_hidden_model("acme/remodelled-chat")
+ # The validation probe stays hidden regardless of embedder config.
+ assert models_route._is_hidden_model("ggml-org/models")
+
+
+def test_is_hidden_model_matches_repo_derived_local_paths(monkeypatch):
+ """Match exact repo-derived cache and LM Studio paths."""
+ from core.rag import config as rag_config
+
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
+
+ assert models_route._is_hidden_model(
+ "/cache/models--org--model/snapshots/abc/model.safetensors"
+ )
+ assert models_route._is_hidden_model(
+ r"C:\Users\u\.cache\huggingface\hub\models--org--model-GGUF\snapshots\abc"
+ )
+ assert models_route._is_hidden_model("/lm-studio/org/model-GGUF/model-Q8_0.gguf")
+ assert not models_route._is_hidden_model("/lm-studio/user/model-chat/model-Q8_0.gguf")
+ assert not models_route._is_hidden_model("/cache/models--org--model-instruct")
+
+
+def test_is_hidden_model_prefers_existing_relative_path(monkeypatch, tmp_path):
+ """Prefer an existing relative path over repo-id syntax."""
+ from core.rag import config as rag_config
+
+ embedder = tmp_path / "models" / "embedder"
+ embedder.mkdir(parents = True)
+ monkeypatch.chdir(tmp_path)
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/embedder-GGUF")
+
+ assert models_route._is_hidden_model(str(embedder))
+
+
+def test_is_hidden_model_keeps_stale_default_embedder_hidden(monkeypatch):
+ """Keep default embedders hidden after a settings change."""
+ from core.rag import config as rag_config
+
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF")
+
+ assert models_route._is_hidden_model("unsloth/bge-small-en-v1.5")
+ assert models_route._is_hidden_model("unsloth/bge-small-en-v1.5-GGUF")
+ assert models_route._is_hidden_model("/models/bge-small-en-v1.5")
+ assert models_route._is_hidden_model("/models/bge-small-en-v1.5-F16.gguf")
+ assert models_route._is_hidden_model(r"C:\models\bge-small-en-v1.5-Q8_0.gguf")
+ # Repo IDs still use exact matching, and similar local basenames must have
+ # a real separator after the static default name.
+ assert not models_route._is_hidden_model("user/bge-small-en-v1.5-chat")
+ assert not models_route._is_hidden_model("/models/bge-small-en-v1.50")
+
+
+def test_is_hidden_model_keeps_env_default_hidden_after_override(monkeypatch):
+ """A persisted override must not expose the deployment's env default."""
+ from core.rag import config as rag_config
+
+ monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False)
+ monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", "org/env-default")
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/custom")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/custom-GGUF")
+
+ assert models_route._is_hidden_model("org/env-default")
+ assert models_route._is_hidden_model("org/env-default-GGUF")
+ assert models_route._is_hidden_model("org/custom")
+ assert models_route._is_hidden_model("org/custom-GGUF")
+ assert not models_route._is_hidden_model("org/env-default-chat")
+
+
+def test_hidden_models_importable_without_heavy_model_stack():
+ """The hub cache scanner imports ``is_hidden_model`` at module scope, so it
+ must not drag in ``utils/models/__init__`` (the model-config + checkpoint
+ stack). Verify in a clean interpreter that importing the helper touches
+ neither ``utils.models`` nor those heavy submodules, and still classifies
+ the probe."""
+ import os
+ import subprocess
+ import textwrap
+
+ backend = Path(__file__).resolve().parents[1]
+ code = textwrap.dedent(
+ """
+ import sys
+
+ class _Blocker:
+ _blocked = (
+ "utils.models",
+ "utils.models.model_config",
+ "utils.models.checkpoints",
+ )
+
+ def find_spec(self, name, path=None, target=None):
+ if name in self._blocked:
+ raise ImportError("blocked heavy import: " + name)
+ return None
+
+ sys.meta_path.insert(0, _Blocker())
+ from utils.hidden_models import is_hidden_model
+
+ loaded = sorted(m for m in sys.modules if m.startswith("utils.models"))
+ assert not loaded, loaded
+ assert is_hidden_model("ggml-org/models") is True
+ assert is_hidden_model("unsloth/gemma-3-270m-it-GGUF") is False
+ print("HIDDEN_MODELS_IMPORT_OK")
+ """
+ )
+ env = dict(os.environ, PYTHONPATH = str(backend))
+ proc = subprocess.run(
+ [sys.executable, "-c", code],
+ capture_output = True,
+ text = True,
+ env = env,
+ )
+ assert proc.returncode == 0, proc.stderr
+ assert "HIDDEN_MODELS_IMPORT_OK" in proc.stdout
+
+
def test_list_cached_gguf_hides_llama_validation_probe(monkeypatch, tmp_path):
"""The ggml-org/models / stories260K install validation probe can land in
the HF cache as a side effect of installing the prebuilt llama-server.
diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py
index 940b35d7ba..b3fa98b604 100644
--- a/studio/backend/tests/test_embedding_model_security_gate.py
+++ b/studio/backend/tests/test_embedding_model_security_gate.py
@@ -52,6 +52,16 @@ def client(monkeypatch):
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
+ monkeypatch.setattr(
+ settings,
+ "effective_gguf_repo",
+ lambda: f"{saved.get('model', 'unsloth/default-embed')}-GGUF",
+ )
+ monkeypatch.setattr(
+ settings,
+ "default_gguf_repo",
+ lambda: "unsloth/default-embed-GGUF",
+ )
app = FastAPI()
app.include_router(settings.router)
@@ -257,6 +267,13 @@ def test_clean_repo_saves_under_force(client, monkeypatch):
r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True})
assert r.status_code == 200
assert saved.get("model") == "acme/clean-embed"
+ assert r.json() == {
+ "embedding_model": "acme/clean-embed",
+ "embedding_gguf_repo": "acme/clean-embed-GGUF",
+ "default_embedding_model": "unsloth/default-embed",
+ "default_embedding_gguf_repo": "unsloth/default-embed-GGUF",
+ "is_custom": True,
+ }
def test_load_sink_refuses_flagged_model(monkeypatch):
diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py
index 3be4af0e32..bcf3ded71c 100644
--- a/studio/backend/tests/test_embedding_model_settings.py
+++ b/studio/backend/tests/test_embedding_model_settings.py
@@ -53,3 +53,10 @@ def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeyp
assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL
assert ems.get_stored_embedding_model() is None
+
+
+def test_env_default_derives_its_gguf_companion(monkeypatch):
+ monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False)
+ monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", "org/env-default-embedder")
+
+ assert rag_config.default_gguf_repo() == "org/env-default-embedder-GGUF"
diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py
index 90fbd19297..c4c0ce15c9 100644
--- a/studio/backend/tests/test_openai_auto_switch.py
+++ b/studio/backend/tests/test_openai_auto_switch.py
@@ -697,6 +697,10 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch):
normal.write_bytes(b"x" * 32)
probe = tmp_path / "stories260K.gguf" # llama.cpp install-validation probe
probe.write_bytes(b"x" * 32)
+ embedder = tmp_path / "embedding-Q8_0.gguf"
+ embedder.write_bytes(b"x" * 32)
+ local_default_embedder = tmp_path / "bge-small-en-v1.5-F16.gguf"
+ local_default_embedder.write_bytes(b"x" * 32)
def _info(mid, path):
return SimpleNamespace(id = mid, path = str(path), model_id = mid, display_name = mid)
@@ -704,7 +708,22 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch):
monkeypatch.setattr(
models_route,
"_scan_models_dir",
- lambda *a, **k: [_info("org/Normal-GGUF", normal), _info("ggml-org/models", probe)],
+ lambda *a, **k: [
+ _info("org/Normal-GGUF", normal),
+ _info("ggml-org/models", probe),
+ SimpleNamespace(
+ id = str(embedder),
+ path = str(embedder),
+ model_id = "unsloth/bge-small-en-v1.5-GGUF",
+ display_name = "embedding-Q8_0",
+ ),
+ SimpleNamespace(
+ id = str(local_default_embedder),
+ path = str(local_default_embedder),
+ model_id = None,
+ display_name = local_default_embedder.name,
+ ),
+ ],
)
monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: [])
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path)
@@ -713,6 +732,8 @@ def test_index_excludes_hidden_models(tmp_path, monkeypatch):
index = resolver._index()
assert "org/normal-gguf" in index # keys are normalized to lowercase
assert "ggml-org/models" not in index
+ assert "unsloth/bge-small-en-v1.5-gguf" not in index
+ assert str(local_default_embedder).lower() not in index
# And the hidden probe cannot be auto-switched to by name.
resolver._scan = (0.0, {})
assert resolver.resolve_local_gguf("ggml-org/models") is None
@@ -1729,6 +1750,8 @@ def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch):
# host path in /v1/models, yet the model stays resolvable by that path too.
from types import SimpleNamespace
import routes.models as models_route
+ from storage import studio_db
+ import utils.paths as paths
gguf = tmp_path / "model-Q4_K_M.gguf"
gguf.write_bytes(b"x" * 32)
@@ -1742,6 +1765,8 @@ def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch):
monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: [])
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path)
monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False)
+ monkeypatch.setattr(paths, "lmstudio_model_dirs", lambda: [])
+ monkeypatch.setattr(studio_db, "list_scan_folders", lambda: [])
resolver._scan = (0.0, {})
# The advertised id is the alias, never the absolute path.
diff --git a/studio/backend/utils/hidden_models.py b/studio/backend/utils/hidden_models.py
new file mode 100644
index 0000000000..20d0bb966e
--- /dev/null
+++ b/studio/backend/utils/hidden_models.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Infra-only model detection shared by the model routes and the hub
+inventory. Lives directly under ``utils`` (not ``utils.models``) so the hub
+cache scanner can import it without pulling in ``utils/models/__init__.py``,
+which eagerly loads the model-config/checkpoint stack, and without importing
+``routes.models`` (import-time side effects, would cycle)."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Optional
+
+# Hub repo id shape ("owner/name", no leading separator); anything else is
+# treated as a local filesystem path.
+_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$")
+
+# The llama.cpp install-validation probe repo. Always hidden.
+_PROBE_REPO_ID = "ggml-org/models"
+# The probe's on-disk filename. Carries the ".gguf" so it stays specific and
+# does not hide unrelated repos like ``user/stories260K-finetune-GGUF``.
+_PROBE_FILENAME = "stories260k.gguf"
+# Keep previously cached defaults hidden after settings changes.
+_DEFAULT_EMBEDDING_REPO_IDS = {
+ "unsloth/bge-small-en-v1.5",
+ "unsloth/bge-small-en-v1.5-GGUF",
+}
+# Local copies do not always retain the repo id. Keep a narrow basename
+# fallback for Studio's static default embedder only; configured custom repos
+# remain exact-match-only.
+_DEFAULT_EMBEDDING_PATH_BASENAMES = {"bge-small-en-v1.5"}
+
+
+def _safe_resolve(path: Path) -> Optional[str]:
+ """resolve() to a string, or None when the path is inaccessible."""
+ try:
+ return str(path.resolve())
+ except OSError:
+ return None
+
+
+def _existing_resolved_path(value: str) -> Optional[str]:
+ """Resolve an existing local path."""
+ path = Path(value).expanduser()
+ try:
+ if not path.exists():
+ return None
+ except OSError:
+ return None
+ return _safe_resolve(path)
+
+
+def _path_contains_repo_id(value: str, repo_ids: set[str]) -> bool:
+ """Match exact repo-derived path segments."""
+ parts = [part for part in value.lower().replace("\\", "/").split("/") if part]
+ for repo_id in repo_ids:
+ owner, name = repo_id.split("/", 1)
+ if f"models--{owner}--{name}" in parts:
+ return True
+ if any(
+ parts[index] == owner and parts[index + 1] == name for index in range(len(parts) - 1)
+ ):
+ return True
+ return False
+
+
+def _path_basename_is_default_embedder(value: str) -> bool:
+ """Match a default embedder folder or a suffixed local weight filename."""
+ normalized = value.lower().replace("\\", "/").rstrip("/")
+ basename = normalized.rsplit("/", 1)[-1]
+ return any(
+ basename == needle
+ or any(basename.startswith(f"{needle}{separator}") for separator in ("-", "_", "."))
+ for needle in _DEFAULT_EMBEDDING_PATH_BASENAMES
+ )
+
+
+def is_hidden_model(*values: str | None) -> bool:
+ """True if any id/path is the RAG embedding model (the effective embedder
+ or its GGUF companion repo) or the llama.cpp install validation probe
+ (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF).
+ None are usable chat models; the probe can be cached as a side effect of
+ installing the prebuilt llama-server and otherwise sorts smallest, so it
+ would be auto-selected.
+
+ Hub repo ids are matched EXACTLY (case-insensitive full "owner/name"), so a
+ custom embedder with a generic basename like "org/model" cannot substring
+ hide unrelated cached repos such as "user/model-chat" or "org/model-GGUF".
+ Existing paths take precedence over the identical ``owner/name`` repo
+ shape. Cache and LM Studio paths use exact repo-derived segments. Local
+ copies of the static default embedder also use a boundary-aware basename
+ fallback; configured custom repos never do."""
+ from core.rag import config as rag_config
+
+ hidden_repo_ids = {
+ _PROBE_REPO_ID.lower(),
+ *(repo_id.lower() for repo_id in _DEFAULT_EMBEDDING_REPO_IDS),
+ }
+ exact_paths: list[str] = []
+ for model in {
+ rag_config.EMBEDDING_MODEL,
+ rag_config.default_gguf_repo(),
+ rag_config.effective_embedding_model(),
+ rag_config.effective_gguf_repo(),
+ }:
+ existing_path = _existing_resolved_path(model)
+ if existing_path:
+ exact_paths.append(existing_path.lower())
+ elif _HF_REPO_ID_RE.match(model):
+ hidden_repo_ids.add(model.lower())
+ else:
+ resolved = _safe_resolve(Path(model).expanduser())
+ if resolved:
+ exact_paths.append(resolved.lower())
+ for v in values:
+ if not v:
+ continue
+ low = v.lower()
+ if _HF_REPO_ID_RE.match(v):
+ # A repo id ("owner/name"): match the hidden set exactly. It is
+ # never a filesystem path, so skip the path/filename checks.
+ if low in hidden_repo_ids:
+ return True
+ continue
+ # Anything else is treated as a filesystem path (the cached snapshot
+ # path, or a local model id). Match the probe by its exact filename and
+ # any configured local-path embedder by exact resolved path. Split on
+ # both separators so a Windows-style path ("...\\stories260K.gguf") is
+ # matched even when this runs on a POSIX interpreter (and vice versa).
+ if low.replace("\\", "/").rsplit("/", 1)[-1] == _PROBE_FILENAME:
+ return True
+ if _path_basename_is_default_embedder(v):
+ return True
+ if _path_contains_repo_id(v, hidden_repo_ids):
+ return True
+ if exact_paths:
+ resolved = _safe_resolve(Path(v).expanduser())
+ if resolved and resolved.lower() in exact_paths:
+ return True
+ return False
diff --git a/studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts b/studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts
new file mode 100644
index 0000000000..f78679310f
--- /dev/null
+++ b/studio/frontend/src/features/hub/hooks/use-hidden-embedding-models.ts
@@ -0,0 +1,46 @@
+// 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 { loadEmbeddingModelSettings } from "@/features/settings";
+import { useEffect, useState } from "react";
+import { useInventoryVersion } from "../stores/inventory-events";
+
+/** Backend-resolved embedding repos that optimistic inventory rows must hide. */
+export function useHiddenEmbeddingModelIds(
+ enabled: boolean,
+): ReadonlySet {
+ const inventoryVersion = useInventoryVersion();
+ const [hiddenIds, setHiddenIds] = useState>(
+ () => new Set(),
+ );
+
+ // biome-ignore lint/correctness/useExhaustiveDependencies: inventory invalidation must reload backend-resolved embedder ids
+ useEffect(() => {
+ if (!enabled) {
+ return;
+ }
+ let cancelled = false;
+ loadEmbeddingModelSettings()
+ .then((settings) => {
+ if (cancelled) {
+ return;
+ }
+ setHiddenIds(
+ new Set(
+ [
+ settings.embeddingModel,
+ settings.embeddingGgufRepo,
+ settings.defaultEmbeddingModel,
+ settings.defaultEmbeddingGgufRepo,
+ ].map((value) => value.trim().toLowerCase()),
+ ),
+ );
+ })
+ .catch(() => undefined);
+ return () => {
+ cancelled = true;
+ };
+ }, [enabled, inventoryVersion]);
+
+ return hiddenIds;
+}
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx
index 630daa48ad..d57f9636fe 100644
--- a/studio/frontend/src/features/hub/hub-page.tsx
+++ b/studio/frontend/src/features/hub/hub-page.tsx
@@ -63,6 +63,7 @@ import { useDiscoverSearch } from "./hooks/use-discover-search";
import { useFeedWriteBack } from "./hooks/use-feed-write-back";
import { useHubFeed } from "./hooks/use-hub-feed";
import { useHubModelVram } from "./hooks/use-hub-model-vram";
+import { useHiddenEmbeddingModelIds } from "./hooks/use-hidden-embedding-models";
import { useModelsSelection } from "./hooks/use-models-selection";
import {
CHANNEL_TO_SECTION,
@@ -73,7 +74,10 @@ import {
SECTION_TO_CHANNEL,
findChannel,
} from "./lib/channels";
-import { isHiddenModelId } from "./lib/hidden-models";
+import {
+ isConfiguredHiddenModelId,
+ isHiddenModelId,
+} from "./lib/hidden-models";
import { inventoryRowMatches, tokenizeQuery } from "./lib/inventory-search";
import { resolveOwnerProviderLogo } from "./lib/provider-logos";
import {
@@ -386,6 +390,7 @@ export function ModelsPage() {
useState("all");
const isDiscoverTab = tab === "discover";
const isDatasetMode = resourceType === "datasets";
+ const hiddenEmbeddingModelIds = useHiddenEmbeddingModelIds(!isDatasetMode);
const urlSection = hubSearch.section ?? null;
const isModelDiscover = isDiscoverTab && !isDatasetMode;
const sectionChannelId: ChannelId | null = urlSection
@@ -700,6 +705,7 @@ export function ModelsPage() {
return discoverRows.filter(
(row) =>
!isHiddenModelId(row.id) &&
+ !isConfiguredHiddenModelId(hiddenEmbeddingModelIds, row.id) &&
// The default feed only shows models with a provider logo.
(!isFeedMode ||
resolveOwnerProviderLogo(row.owner, row.repo) !== null) &&
@@ -714,6 +720,7 @@ export function ModelsPage() {
);
}, [
discoverRows,
+ hiddenEmbeddingModelIds,
isDatasetMode,
isFeedMode,
effectiveDiscoverFormat,
@@ -739,7 +746,11 @@ export function ModelsPage() {
effectiveCachedRows,
effectiveLocalRows,
)
- .filter((row) => !isHiddenModelId(row.id))
+ .filter(
+ (row) =>
+ !isHiddenModelId(row.id) &&
+ !isConfiguredHiddenModelId(hiddenEmbeddingModelIds, row.id),
+ )
.filter((row) => matchesFormat(row.result.isGguf, "gguf"))
// Same fit filter as the main Discover list, so the feed carousel
// honors the toggle too.
@@ -751,6 +762,7 @@ export function ModelsPage() {
),
[
hubFeed.trending.results,
+ hiddenEmbeddingModelIds,
modelDiscoveryInventorySignature,
fitOnDeviceOnly,
gpu,
@@ -778,22 +790,29 @@ export function ModelsPage() {
() => (isDiscoverTab ? [] : tokenizeQuery(deferredDebouncedQuery)),
[isDiscoverTab, deferredDebouncedQuery],
);
- // Hide infra models (e.g. the RAG embedder bge-small-en-v1.5) from the On
- // Device list like Discover, but reveal a row when a query matches it so the
- // user can confirm it is already downloaded.
+ // Server cache rows already apply variant-aware infra hiding. Optimistic
+ // rows are not server-confirmed, so apply the client filter first.
const isVisibleInventoryRow = useCallback(
- (row: CachedInventoryRow | LocalInventoryRow) =>
- // Local rows can have a null repoId and an id that is a hash rather than
- // the file path/name, so also check path/title (the backend's
- // _is_hidden_model checks the on-disk path for the same reason).
- !isHiddenModelId(
- row.id,
- row.repoId,
- row.kind !== "cache" ? row.path : undefined,
- row.kind !== "cache" ? row.title : undefined,
- ) ||
- (inventoryTokens.length > 0 && inventoryRowMatches(row, inventoryTokens)),
- [inventoryTokens],
+ (row: CachedInventoryRow | LocalInventoryRow) => {
+ if (row.kind === "cache") {
+ return (
+ !row.optimistic ||
+ (!isHiddenModelId(row.id, row.repoId, row.cachePath) &&
+ !isConfiguredHiddenModelId(
+ hiddenEmbeddingModelIds,
+ row.id,
+ row.repoId,
+ row.cachePath,
+ ))
+ );
+ }
+ // Local rows may lack a repo id, so also check path and title.
+ return (
+ !isHiddenModelId(row.id, row.repoId, row.path, row.title) ||
+ (inventoryTokens.length > 0 && inventoryRowMatches(row, inventoryTokens))
+ );
+ },
+ [hiddenEmbeddingModelIds, inventoryTokens],
);
// Format filter is a deliberate scope narrowing, so hard-filter it out. The
// text query instead drives dim-not-filter on On Device (see ModelsCatalog) so
diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts
index 3515f6ca76..5d4151e87d 100644
--- a/studio/frontend/src/features/hub/index.ts
+++ b/studio/frontend/src/features/hub/index.ts
@@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { cancelStagedModelDownload } from "./download-manager";
+export { bumpInventoryVersion } from "./stores/inventory-events";
export {
getHfToken,
mirrorHfTokenInto,
diff --git a/studio/frontend/src/features/hub/inventory/inventory-hints.ts b/studio/frontend/src/features/hub/inventory/inventory-hints.ts
index af9f254ab3..5e202e3150 100644
--- a/studio/frontend/src/features/hub/inventory/inventory-hints.ts
+++ b/studio/frontend/src/features/hub/inventory/inventory-hints.ts
@@ -12,6 +12,7 @@ export type InventoryHintRow = {
repo_id: string;
size_bytes: number;
partial?: boolean;
+ optimistic?: boolean;
};
export type InventoryHintReconciliation = {
@@ -41,6 +42,7 @@ function optimisticRow(hint: InventoryHint): InventoryHintRow {
repo_id: hint.repoId,
size_bytes: hint.bytes ?? 0,
partial: false,
+ optimistic: true,
};
}
@@ -101,9 +103,14 @@ function mergeInventoryHint(
if (idx === -1) {
return [...rows, seed];
}
+ const serverRow = rows[idx];
const merged = {
- ...rows[idx],
- ...seed,
+ ...serverRow,
+ // A completed hint may arrive before a partial server scan catches up. In
+ // that case keep the synthetic row non-runnable. A complete server row is
+ // already authoritative even when its runnable-weight size is smaller than
+ // the hint's full-snapshot byte count, so do not mark that merge optimistic.
+ ...(serverRow.partial ? seed : { optimistic: false }),
size_bytes: Math.max(rowSizeBytes(rows[idx]), rowSizeBytes(seed)),
};
return [...rows.slice(0, idx), merged, ...rows.slice(idx + 1)];
diff --git a/studio/frontend/src/features/hub/inventory/types.ts b/studio/frontend/src/features/hub/inventory/types.ts
index c86ffb1d86..6f65a56037 100644
--- a/studio/frontend/src/features/hub/inventory/types.ts
+++ b/studio/frontend/src/features/hub/inventory/types.ts
@@ -54,6 +54,7 @@ export interface CachedInventoryRow {
libraryName?: string | null;
quantMethod?: string | null;
liveDownload?: boolean;
+ optimistic?: boolean;
}
export interface LocalInventoryRow {
diff --git a/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts b/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts
index a7dc7ae3f3..fea7b3d331 100644
--- a/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts
+++ b/studio/frontend/src/features/hub/inventory/use-hub-inventory.ts
@@ -204,6 +204,7 @@ function liveDownloadInventoryRows(
size_bytes: job.displayBytes,
partial: true,
partial_transport: null,
+ optimistic: true,
},
modelFormat,
),
diff --git a/studio/frontend/src/features/hub/inventory/view-models.ts b/studio/frontend/src/features/hub/inventory/view-models.ts
index 63d70418be..334050fab4 100644
--- a/studio/frontend/src/features/hub/inventory/view-models.ts
+++ b/studio/frontend/src/features/hub/inventory/view-models.ts
@@ -176,6 +176,7 @@ export function buildCachedInventoryRow(
runtime?: string | null;
format_variant?: string | null;
capabilities?: BackendModelCapabilities | null;
+ optimistic?: boolean;
},
fallbackFormat: ModelInventoryFormat,
): CachedInventoryRow {
@@ -185,6 +186,15 @@ export function buildCachedInventoryRow(
const inferredFromEndpoint =
rawModelFormat === "unknown" && modelFormat !== "unknown";
const requiresVariant = modelFormat === "gguf";
+ const capabilities = normalizeCapabilities(
+ inferredFromEndpoint ? null : row.capabilities,
+ modelFormat,
+ row.partial ?? false,
+ requiresVariant,
+ );
+ if (row.optimistic) {
+ capabilities.canChat = false;
+ }
return {
kind: "cache",
id:
@@ -202,12 +212,7 @@ export function buildCachedInventoryRow(
modelFormat,
),
formatVariant: row.format_variant ?? null,
- capabilities: normalizeCapabilities(
- inferredFromEndpoint ? null : row.capabilities,
- modelFormat,
- row.partial ?? false,
- requiresVariant,
- ),
+ capabilities,
bytes: row.size_bytes,
cachePath: row.cache_path ?? null,
partial: row.partial ?? false,
@@ -216,6 +221,7 @@ export function buildCachedInventoryRow(
tags: row.tags,
libraryName: row.library_name ?? null,
quantMethod: row.quant_method ?? null,
+ optimistic: row.optimistic,
};
}
diff --git a/studio/frontend/src/features/hub/lib/hidden-models.ts b/studio/frontend/src/features/hub/lib/hidden-models.ts
index 2dbe257947..634a061e0c 100644
--- a/studio/frontend/src/features/hub/lib/hidden-models.ts
+++ b/studio/frontend/src/features/hub/lib/hidden-models.ts
@@ -1,11 +1,13 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-// Infra models hidden from every browse/preview list (Hub discover and the chat
-// model selector). Mirrors the backend `_is_hidden_model`: the RAG embedding
-// model and the llama.cpp validation probe are not usable chat models. Per-repo
-// file/download views are NOT filtered, so a reinstall still shows the model as
-// already downloaded.
+// Infra models hidden from browse/preview lists (Hub Discover, the chat model
+// selector, and local on-device rows). Mirrors the backend
+// `utils.hidden_models`: the RAG embedding model and the llama.cpp validation
+// probe are not usable chat models. Server-confirmed cache rows are trusted
+// because the backend applies variant-aware filtering. Optimistic cache rows
+// still use these needles until the server confirms them. Per-repo views are
+// not filtered, so reinstall flows still show downloaded files.
const HIDDEN_NEEDLES = [
"bge-small-en-v1.5", // RAG embedder: unsloth/bge-small-en-v1.5[-GGUF]
"ggml-org/models", // llama.cpp validation probe repo
@@ -17,8 +19,20 @@ export function isHiddenModelId(
...values: (string | null | undefined)[]
): boolean {
return values.some((v) => {
- if (!v) return false;
+ if (!v) {
+ return false;
+ }
const lower = v.toLowerCase();
return HIDDEN_NEEDLES.some((needle) => lower.includes(needle));
});
}
+
+/** Exact-match configured infra repos without hiding similarly named models. */
+export function isConfiguredHiddenModelId(
+ configuredIds: ReadonlySet,
+ ...values: (string | null | undefined)[]
+): boolean {
+ return values.some(
+ (value) => value != null && configuredIds.has(value.trim().toLowerCase()),
+ );
+}
diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts
index 9a61142f73..cc21559f38 100644
--- a/studio/frontend/src/features/settings/api/embedding-model.ts
+++ b/studio/frontend/src/features/settings/api/embedding-model.ts
@@ -2,11 +2,14 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
+import { bumpInventoryVersion } from "@/features/hub";
import { readFastApiError } from "@/lib/format-fastapi-error";
export type EmbeddingModelSettings = {
embeddingModel: string;
+ embeddingGgufRepo: string;
defaultEmbeddingModel: string;
+ defaultEmbeddingGgufRepo: string;
isCustom: boolean;
};
@@ -14,8 +17,12 @@ type ApiEmbeddingModelSettings = {
// biome-ignore lint/style/useNamingConvention: API schema
embedding_model: string;
// biome-ignore lint/style/useNamingConvention: API schema
+ embedding_gguf_repo: string;
+ // biome-ignore lint/style/useNamingConvention: API schema
default_embedding_model: string;
// biome-ignore lint/style/useNamingConvention: API schema
+ default_embedding_gguf_repo: string;
+ // biome-ignore lint/style/useNamingConvention: API schema
is_custom: boolean;
};
@@ -30,7 +37,9 @@ export class EmbeddingModelBlockedError extends Error {}
function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings {
return {
embeddingModel: settings.embedding_model,
+ embeddingGgufRepo: settings.embedding_gguf_repo,
defaultEmbeddingModel: settings.default_embedding_model,
+ defaultEmbeddingGgufRepo: settings.default_embedding_gguf_repo,
isCustom: settings.is_custom,
};
}
@@ -75,7 +84,9 @@ export async function updateEmbeddingModelSettings(
await readFastApiError(res, "Failed to save embedding model"),
);
}
- return fromApi(await res.json());
+ const settings = fromApi(await res.json());
+ bumpInventoryVersion();
+ return settings;
}
export async function resetEmbeddingModelSettings(): Promise {
@@ -87,5 +98,7 @@ export async function resetEmbeddingModelSettings(): Promise
Date: Sun, 19 Jul 2026 18:37:23 +0800
Subject: [PATCH 013/255] fix(registry): don't register deepseek models at
import time (#7227)
* fix(registry): don't register deepseek models at import time
`_deepseek.py` called `register_deepseek_models(include_original_model=True)`
at module scope, so merely importing `unsloth.registry` registered models
(and reached the hub via `list_models`) as a side effect. None of the other
five families (`_gemma`/`_llama`/`_mistral`/`_phi`/`_qwen`) do this; they only
register when `register_models()` asks them to.
Two consequences:
- Importing the registry populated MODEL_REGISTRY on its own (32 entries,
including 10 `deepseek-ai` original models that no other family leaks) and
did network I/O at import time.
- Because the import-time call set the `_IS_DEEPSEEK_*_REGISTERED` guards with
`include_original_model=True`, the later `register_models()` call (which uses
the default `include_original_model=False`) early-returned, so the
original-model set won permanently.
Remove the stray module-level call. The `if __name__ == "__main__"` block below
still registers with `include_original_model=True` for standalone use, so the
generator script is unaffected.
Co-Authored-By: Claude Opus 4.8 (1M context)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* test(registry): make import-side-effect test pass on CPU-only runners
The new test spawned a fresh `python -c "import unsloth.registry"` that did
not inherit tests/conftest.py's GPU-free harness, so on no-accelerator CI
runners the child raised NotImplementedError from unsloth_zoo.device_type
before printing REGISTRY_SIZE. With check=True this surfaced only as an
opaque CalledProcessError, turning the "Repo tests (CPU)" job red even
though the registry fix is correct.
Import this directory's conftest inside the child first so it applies the
same device_type stubs and torch.cuda probe patches. Also use check=False
and include the child stdout/stderr in the assertion message so a future
import regression is legible instead of an opaque non-zero exit.
* test(registry): assert register_models() leaks no upstream originals
Adds a fresh-interpreter test that register_models() registers only
unsloth-org models (deepseek still present via the normal path) and never
leaks the upstream deepseek-ai originals that the import-time guard poisoning
used to leak (129 -> 139). Factors the conftest-harness subprocess runner
into a shared helper reused by both registry import tests.
---------
Co-authored-by: Claude Opus 4.8 (1M context)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
tests/test_model_registry.py | 85 +++++++++++++++++++++++++++++++++++
unsloth/registry/_deepseek.py | 2 -
2 files changed, 85 insertions(+), 2 deletions(-)
diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py
index 283b099107..76f462c5cf 100644
--- a/tests/test_model_registry.py
+++ b/tests/test_model_registry.py
@@ -1,5 +1,8 @@
"""Register each model set and check the registered ids exist on the HF Hub."""
+import os
+import subprocess
+import sys
from dataclasses import dataclass
import pytest
@@ -77,3 +80,85 @@ def test_quant_type():
assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models)
quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH]
assert all(quant_tag in m.model_path for m in dynamic_quant_models)
+
+
+def _run_registry_child(body: str) -> subprocess.CompletedProcess:
+ """Run ``body`` in a fresh interpreter that first imports this directory's
+ ``conftest`` so it inherits the same GPU-free harness the pytest session
+ uses (device_type stubs plus torch.cuda probe patches). Without it,
+ ``import unsloth.registry`` raises ``NotImplementedError`` from
+ ``unsloth_zoo.device_type`` on no-accelerator CI runners, so the child
+ would exit non-zero and the test would fail even though the registry code
+ is correct. A fresh process also keeps each check independent of any
+ ``register_models()`` calls other tests make on the shared registry.
+ """
+ tests_dir = os.path.dirname(os.path.abspath(__file__))
+ prelude = (
+ f"import sys; sys.path.insert(0, {tests_dir!r})\n"
+ "try:\n"
+ " import conftest # noqa: F401 GPU-free harness on no-accelerator runners\n"
+ "except Exception:\n"
+ " pass\n"
+ )
+ return subprocess.run(
+ [sys.executable, "-c", prelude + body],
+ capture_output = True,
+ text = True,
+ check = False,
+ )
+
+
+def test_importing_registry_does_not_register_models():
+ """Importing the registry must not populate MODEL_REGISTRY on its own.
+
+ ``_deepseek`` used to call ``register_deepseek_models(...)`` at module
+ scope, so merely importing ``unsloth.registry`` registered models as an
+ import side effect, unlike every other family which only registers on
+ demand.
+ """
+ result = _run_registry_child(
+ "import unsloth.registry\n"
+ "from unsloth.registry.registry import MODEL_REGISTRY\n"
+ "print('REGISTRY_SIZE', len(MODEL_REGISTRY))"
+ )
+ assert result.returncode == 0, (
+ f"registry import subprocess exited {result.returncode}\n"
+ f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+ )
+ size_lines = [line for line in result.stdout.splitlines() if line.startswith("REGISTRY_SIZE")]
+ assert size_lines == ["REGISTRY_SIZE 0"], result.stdout + result.stderr
+
+
+def test_register_models_registers_no_upstream_originals():
+ """``register_models()`` must register each family's ``unsloth``-org models
+ and must NOT leak upstream vendor "original" models.
+
+ Before the fix, ``_deepseek``'s import-time
+ ``register_deepseek_models(include_original_model = True)`` set the
+ ``_IS_DEEPSEEK_*_REGISTERED`` guards, so the later default
+ ``register_models()`` early-returned for deepseek and its 10 ``deepseek-ai``
+ originals leaked permanently (129 -> 139). This asserts the whole registry
+ is ``unsloth``-org after ``register_models()`` while deepseek is still
+ registered via the normal path. Runs in a fresh interpreter so it is
+ independent of other tests' registry mutations.
+ """
+ result = _run_registry_child(
+ "import unsloth.registry\n"
+ "from unsloth.registry import register_models\n"
+ "from unsloth.registry.registry import MODEL_REGISTRY\n"
+ "register_models()\n"
+ "orgs = sorted({m.org for m in MODEL_REGISTRY.values()})\n"
+ "deepseek = [k for k in MODEL_REGISTRY if 'deepseek' in k.lower()]\n"
+ "print('ORGS', orgs)\n"
+ "print('NUM_DEEPSEEK', len(deepseek))"
+ )
+ assert result.returncode == 0, (
+ f"register_models subprocess exited {result.returncode}\n"
+ f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+ )
+ out = result.stdout
+ # Every registered model is unsloth-org: no upstream "original" leaked.
+ assert "ORGS ['unsloth']" in out, out + result.stderr
+ # Deepseek is still registered via the normal path, just without originals.
+ deepseek_lines = [line for line in out.splitlines() if line.startswith("NUM_DEEPSEEK")]
+ assert deepseek_lines and int(deepseek_lines[0].split()[1]) > 0, out + result.stderr
diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py
index e29190f0f2..618453fb82 100644
--- a/unsloth/registry/_deepseek.py
+++ b/unsloth/registry/_deepseek.py
@@ -171,8 +171,6 @@ def _list_deepseek_r1_distill_models():
return distill_models
-register_deepseek_models(include_original_model = True)
-
if __name__ == "__main__":
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
From e8db1cecff48cfda2a3376f5d85b0a10a1170416 Mon Sep 17 00:00:00 2001
From: Hakan Baysal
Date: Sun, 19 Jul 2026 13:45:29 +0300
Subject: [PATCH 014/255] studio: show the active run's saved config in the
Training Progress popover (#7217)
* studio: show the active run's saved config in the Training Progress popover
The Training Config popover on the live Training Progress page read the
editable form store (useTrainingConfigStore), so it showed stale/static values
whenever the form changed after the run started; only the History view read the
run's saved config snapshot, which is why re-opening the same run from Recents
showed the correct values (#6853).
Wire the live view to the same authoritative source History already uses:
- Extract History's field mapping into sections/run-config-override.ts
(mapRunConfigToOverride) so both views share one mapper over
GET /api/train/runs/{id} config.
- LiveTrainingView fetches the run record as soon as the job id is known and
passes the mapped override to ProgressSection; the fetched config is keyed
by job id, and until it loads (or if the fetch fails) the form store remains
the fallback. The run record is created at job start, so it is available
while the run is live.
- ProgressSection prefers configOverride whenever one is present instead of
only when isHistorical, so the live override takes effect.
Adds a source-level regression test pinning the wiring and the mapper's
backend config keys.
Fixes #6853
* studio: retry the run-config fetch after the first step, carry the saved method
Two review fixes on the live Training Config popover source:
1. The backend creates the run row only on the first progress event, so the
fetch issued as soon as the job id appeared commonly 404'd during
model/dataset preparation and never retried -- leaving the popover on the
form store for the whole run. The effect is now also keyed on
firstStepReceived (and skips once resolved for the job), so it re-fetches
exactly when the row is guaranteed to exist.
2. The popover's method label and LoRA-row visibility came from
viewData.trainingMethod, still read from the editable form store; changing
the form (e.g. LoRA -> Full) after starting a run relabeled it and hid its
saved LoRA rows. The run-config mapper now derives trainingMethod from the
snapshot's training_type/load_in_4bit (via parseBackendTrainingMethod, now
exported from the feature index) and the live view prefers it.
* studio: fetch the run config on a terminal phase too, not just the first step
The live config-popover fetch was keyed on firstStepReceived, which the runtime
store sets only when step > 0. A run that fails or completes during preparation
(before step 1) creates and finalizes its row from the terminal error/complete
event, but neither the job id nor firstStepReceived changed, so the fetch never
ran and the popover stayed on the editable form store -- showing the wrong
config/method if the form was edited afterward (Configure re-enables on failure).
Gate the fetch on a runRowReady signal = firstStepReceived OR a terminal phase
(completed/error/stopped), the states in which the backend guarantees the row
exists. This also stops the earlier fetch-then-404 churn during preparation and
lets the effect depend only on values it reads (no lint suppression needed).
* studio: retry the run-config lookup and accept a hydrated step as row-ready
Two ways the popover could stay stuck on the editable form store for a whole
run:
- The backend publishes the progress event that reveals the run before
create_run commits, so the first lookup can lose that race and 404. The catch
changed neither runRowReady nor fetchedRunConfig, leaving every effect
dependency identical, so no further attempt was ever made for that job. The
failure path now schedules an explicit retry, bounded and keyed by job id, so
a genuinely absent row falls back to the form store instead of polling.
- A run recovered through status/metrics polling (SSE unavailable or blocked)
has currentStep restored by applyStatus/applyMetrics but never
firstStepReceived, and the phase stays training, so the row was treated as
not ready even at step > 0. currentStep > 0 is now a readiness signal of its
own.
* studio: fetch the saved run config as soon as the job id exists
start_training() inserts the run row before the pump can consume any event --
deliberately, so the run appears in history during model loading -- and /status
exposes the job id throughout the pre-step phases. Gating the lookup on a first
step or a terminal phase therefore held the popover on the editable form store
for the whole configuring/loading/downloading window, which on a long model or
dataset load is minutes, and indefinitely for a run adopted from another client.
The job id is now the entire readiness condition; the existing bounded retry
still covers the instant before the insert commits.
* Fix Training Config popover fallback for history runs without a saved config; tighten popover comments
---------
Co-authored-by: danielhanchen
---
.../test_training_config_popover_source.py | 109 ++++++++++++++++++
.../studio/historical-training-view.tsx | 21 +---
.../features/studio/live-training-view.tsx | 91 ++++++++++++++-
.../studio/sections/progress-section.tsx | 40 +++----
.../studio/sections/run-config-override.ts | 54 +++++++++
.../frontend/src/features/training/index.ts | 1 +
6 files changed, 270 insertions(+), 46 deletions(-)
create mode 100644 studio/backend/tests/test_training_config_popover_source.py
create mode 100644 studio/frontend/src/features/studio/sections/run-config-override.ts
diff --git a/studio/backend/tests/test_training_config_popover_source.py b/studio/backend/tests/test_training_config_popover_source.py
new file mode 100644
index 0000000000..4263b012eb
--- /dev/null
+++ b/studio/backend/tests/test_training_config_popover_source.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Source-level regression guards for the Training Config popover data source
+(#6853).
+
+The live Training Progress popover used to read the editable form store
+(useTrainingConfigStore) while a run was active, so it showed stale/static
+values whenever the user touched the form after starting the run; only the
+History view read the run's saved config snapshot. These guards pin the fixed
+wiring: both views feed ProgressSection a config override mapped from
+GET /api/train/runs/{id}, and ProgressSection prefers that override whenever
+one is present -- not only for historical views.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+_STUDIO_FRONTEND = Path(__file__).resolve().parents[2] / "frontend" / "src" / "features" / "studio"
+
+
+def _read(rel: str) -> str:
+ return (_STUDIO_FRONTEND / rel).read_text(encoding = "utf-8")
+
+
+def test_progress_section_prefers_override_over_form_store():
+ src = _read("sections/progress-section.tsx")
+ # Fields key on the override's presence, not isHistorical: a live view passing
+ # an override wins over the store; without one, live keeps the store while
+ # History shows blanks rather than unrelated live form values.
+ assert "const cfg = configOverride ?? (isHistorical ? undefined : config)" in src
+ assert "const cfgEpochs = cfg?.epochs" in src
+ assert "isHistorical ? configOverride?.epochs" not in src
+
+
+def test_live_view_fetches_the_active_run_config():
+ src = _read("live-training-view.tsx")
+ # Live view resolves the run's saved config snapshot by job id...
+ assert "getTrainingRun(" in src
+ assert "mapRunConfigToOverride(" in src
+ # ...and hands it to the popover.
+ assert "configOverride={runConfigOverride}" in src
+
+
+def test_live_view_fetches_as_soon_as_the_job_id_exists():
+ # start_training() inserts the run row BEFORE the pump consumes any event, so
+ # the saved config is available during configuring/loading/downloading. The
+ # job id is therefore the whole readiness condition: gating on a first step
+ # or a terminal phase would show the wrong config for the entire pre-step
+ # window of a long load, or for a run adopted from another client.
+ src = _read("live-training-view.tsx")
+ assert "if (!runtime.jobId) {" in src
+ assert "[runtime.jobId, fetchedRunConfig, fetchAttempt]" in src
+ # No step/phase readiness gate may creep back in.
+ assert "runRowReady" not in src
+
+
+def test_live_view_retries_the_transient_row_miss():
+ # start_training() creates the row before the pump, but a lookup racing that
+ # commit can still 404. Nothing else in the effect deps changes on failure, so
+ # the retry must be explicit and bounded, else a genuinely absent row would
+ # poll forever instead of falling back to the form store.
+ src = _read("live-training-view.tsx")
+ assert "RUN_CONFIG_FETCH_RETRIES" in src
+ assert "RUN_CONFIG_FETCH_RETRY_MS" in src
+ assert "setFetchAttempt(" in src
+ assert "attempts >= RUN_CONFIG_FETCH_RETRIES" in src
+ # The budget is keyed by job so a new run always starts fresh.
+ assert "fetchAttempt?.jobId === jobId ? fetchAttempt.count : 0" in src
+ # The pending retry must be cancelled with the effect.
+ assert "clearTimeout(retryTimer)" in src
+
+
+def test_live_view_prefers_saved_training_method():
+ # The method label / LoRA-row visibility must come from the run snapshot,
+ # not the editable form (which may have changed since the run started).
+ src = _read("live-training-view.tsx")
+ assert "runConfigOverride?.trainingMethod ?? config.trainingMethod" in src
+
+
+def test_history_view_uses_the_shared_mapper():
+ src = _read("historical-training-view.tsx")
+ # Shared mapper, not a re-inlined field-by-field copy that could drift.
+ assert "mapRunConfigToOverride(detail.config)" in src
+ assert "num_epochs" not in src
+
+
+def test_shared_mapper_matches_backend_config_keys():
+ src = _read("sections/run-config-override.ts")
+ # The mapper reads the run config JSON the backend snapshots at job start;
+ # keep the key set pinned so a silent rename breaks loudly here.
+ for key in (
+ "training_type",
+ "load_in_4bit",
+ "num_epochs",
+ "batch_size",
+ "learning_rate",
+ "max_steps",
+ "max_seq_length",
+ "warmup_steps",
+ "optim",
+ "lora_r",
+ "lora_alpha",
+ "lora_dropout",
+ "use_rslora",
+ "use_loftq",
+ ):
+ assert key in src, f"run-config mapper lost backend key {key}"
diff --git a/studio/frontend/src/features/studio/historical-training-view.tsx b/studio/frontend/src/features/studio/historical-training-view.tsx
index 2f80fc29ca..b6ec06b06a 100644
--- a/studio/frontend/src/features/studio/historical-training-view.tsx
+++ b/studio/frontend/src/features/studio/historical-training-view.tsx
@@ -8,6 +8,7 @@ import { parseBackendTrainingMethod } from "@/features/training/lib/training-met
import { type ReactElement, useEffect, useState } from "react";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
+import { mapRunConfigToOverride } from "./sections/run-config-override";
import { translate, useT } from "@/i18n";
type StudioT = ReturnType;
@@ -147,25 +148,7 @@ export function HistoricalTrainingView({
}
const viewData = mapToViewData(detail, t);
- const configOverride = detail.config
- ? {
- epochs: detail.config.num_epochs as number | undefined,
- batchSize: detail.config.batch_size as number | undefined,
- learningRate: detail.config.learning_rate as string | undefined,
- maxSteps: detail.config.max_steps as number | undefined,
- contextLength: detail.config.max_seq_length as number | undefined,
- warmupSteps: detail.config.warmup_steps as number | undefined,
- optimizerType: detail.config.optim as string | undefined,
- loraRank: detail.config.lora_r as number | undefined,
- loraAlpha: detail.config.lora_alpha as number | undefined,
- loraDropout: detail.config.lora_dropout as number | undefined,
- loraVariant: detail.config.use_rslora
- ? "rslora"
- : detail.config.use_loftq
- ? "loftq"
- : "lora",
- }
- : undefined;
+ const configOverride = mapRunConfigToOverride(detail.config);
return (
diff --git a/studio/frontend/src/features/studio/live-training-view.tsx b/studio/frontend/src/features/studio/live-training-view.tsx
index cce39adbf4..0aecc7030e 100644
--- a/studio/frontend/src/features/studio/live-training-view.tsx
+++ b/studio/frontend/src/features/studio/live-training-view.tsx
@@ -1,18 +1,42 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import { cn } from "@/lib/utils";
import {
+ getTrainingRun,
useTrainingConfigStore,
useTrainingRuntimeStore,
} from "@/features/training";
import type { TrainingViewData } from "@/features/training";
+import { cn } from "@/lib/utils";
import type { ReactElement } from "react";
+import { useEffect, useState } from "react";
import { useShallow } from "zustand/react/shallow";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
+import {
+ type RunConfigOverride,
+ mapRunConfigToOverride,
+} from "./sections/run-config-override";
import { TrainingStartOverlay } from "./training-start-overlay";
+/** Retry budget for the run-config lookup. The row is inserted at
+ * start_training(), but a lookup issued in the same instant can still miss it;
+ * a few short retries cover that without polling a genuinely absent row. */
+const RUN_CONFIG_FETCH_RETRIES = 5;
+const RUN_CONFIG_FETCH_RETRY_MS = 1000;
+
+/** The fetched run config only applies while it belongs to the active job;
+ * a stale record from a previous run falls back to the form store. */
+function activeRunOverride(
+ fetched: { jobId: string; override: RunConfigOverride | undefined } | null,
+ jobId: string | null,
+): RunConfigOverride | undefined {
+ if (fetched === null || fetched.jobId !== jobId) {
+ return undefined;
+ }
+ return fetched.override;
+}
+
export function LiveTrainingView(): ReactElement {
const runtime = useTrainingRuntimeStore(
useShallow((state) => ({
@@ -52,6 +76,59 @@ export function LiveTrainingView(): ReactElement {
})),
);
+ // Show the ACTIVE run's saved config, not the editable form store the user may
+ // have changed since starting (#6853). start_training() commits the run row
+ // before the pump, so the job id alone gates the fetch; the bounded retry below
+ // covers the narrow uncommitted window, and until it loads ProgressSection falls
+ // back to the form store. The result is keyed by job id and filtered at render.
+ const [fetchedRunConfig, setFetchedRunConfig] = useState<{
+ jobId: string;
+ override: RunConfigOverride | undefined;
+ } | null>(null);
+ // Retry budget for the transient 404 below, keyed by job so a new run always
+ // starts with a fresh budget.
+ const [fetchAttempt, setFetchAttempt] = useState<{
+ jobId: string;
+ count: number;
+ } | null>(null);
+ useEffect(() => {
+ if (!runtime.jobId) {
+ return;
+ }
+ const jobId = runtime.jobId;
+ if (fetchedRunConfig !== null && fetchedRunConfig.jobId === jobId) {
+ return; // already resolved for this job
+ }
+ const attempts = fetchAttempt?.jobId === jobId ? fetchAttempt.count : 0;
+ const controller = new AbortController();
+ let retryTimer: ReturnType
| undefined;
+ getTrainingRun(jobId, controller.signal)
+ .then((detail) => {
+ setFetchedRunConfig({
+ jobId,
+ override: mapRunConfigToOverride(detail.config),
+ });
+ })
+ .catch(() => {
+ // A lookup racing the row commit can miss transiently; nothing else in
+ // the deps changes on failure, so retry explicitly. Bounded so a genuinely
+ // absent row falls back to the form store instead of polling forever.
+ if (controller.signal.aborted || attempts >= RUN_CONFIG_FETCH_RETRIES) {
+ return;
+ }
+ retryTimer = setTimeout(() => {
+ setFetchAttempt({ jobId, count: attempts + 1 });
+ }, RUN_CONFIG_FETCH_RETRY_MS);
+ });
+ return () => {
+ controller.abort();
+ if (retryTimer !== undefined) {
+ clearTimeout(retryTimer);
+ }
+ };
+ }, [runtime.jobId, fetchedRunConfig, fetchAttempt]);
+ const runConfigOverride = activeRunOverride(fetchedRunConfig, runtime.jobId);
+
const activeProjectName =
runtime.startProjectName !== null
? runtime.startProjectName.trim() || null
@@ -76,7 +153,11 @@ export function LiveTrainingView(): ReactElement {
isTrainingRunning: runtime.isTrainingRunning,
modelName: runtime.startModelName ?? config.selectedModel ?? "",
projectName: activeProjectName,
- trainingMethod: config.trainingMethod ?? "",
+ // Prefer the saved run's method: the form may have been edited (e.g. LoRA
+ // -> Full) after the run started, which would relabel the run and hide its
+ // saved LoRA rows in the popover.
+ trainingMethod:
+ runConfigOverride?.trainingMethod ?? config.trainingMethod ?? "",
lossHistory: runtime.lossHistory,
lrHistory: runtime.lrHistory,
gradNormHistory: runtime.gradNormHistory,
@@ -105,7 +186,11 @@ export function LiveTrainingView(): ReactElement {
)}
>
o.value === cfgOptimizerType)?.label ??
diff --git a/studio/frontend/src/features/studio/sections/run-config-override.ts b/studio/frontend/src/features/studio/sections/run-config-override.ts
new file mode 100644
index 0000000000..a1272bfeb0
--- /dev/null
+++ b/studio/frontend/src/features/studio/sections/run-config-override.ts
@@ -0,0 +1,54 @@
+// 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 { parseBackendTrainingMethod } from "@/features/training";
+
+/** Shape of the Training Config popover's data when it is driven by a saved
+ * run snapshot instead of the editable form store. */
+export interface RunConfigOverride {
+ trainingMethod?: string;
+ epochs?: number;
+ batchSize?: number;
+ learningRate?: string;
+ maxSteps?: number;
+ contextLength?: number;
+ warmupSteps?: number;
+ optimizerType?: string;
+ loraRank?: number;
+ loraAlpha?: number;
+ loraDropout?: number;
+ loraVariant?: string;
+}
+
+/** Map a saved run's config (GET /api/train/runs/{id} `detail.config`) into the
+ * Training Config popover's override shape. Shared by the History view and the
+ * live Current Run view so both read the same authoritative run snapshot
+ * instead of the editable form store (#6853). */
+export function mapRunConfigToOverride(
+ config: Record | null | undefined,
+): RunConfigOverride | undefined {
+ if (!config) {
+ return undefined;
+ }
+ return {
+ trainingMethod: parseBackendTrainingMethod(
+ config.training_type,
+ config.load_in_4bit,
+ ),
+ epochs: config.num_epochs as number | undefined,
+ batchSize: config.batch_size as number | undefined,
+ learningRate: config.learning_rate as string | undefined,
+ maxSteps: config.max_steps as number | undefined,
+ contextLength: config.max_seq_length as number | undefined,
+ warmupSteps: config.warmup_steps as number | undefined,
+ optimizerType: config.optim as string | undefined,
+ loraRank: config.lora_r as number | undefined,
+ loraAlpha: config.lora_alpha as number | undefined,
+ loraDropout: config.lora_dropout as number | undefined,
+ loraVariant: config.use_rslora
+ ? "rslora"
+ : config.use_loftq
+ ? "loftq"
+ : "lora",
+ };
+}
diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts
index 553dcc2af5..a0d249ff1b 100644
--- a/studio/frontend/src/features/training/index.ts
+++ b/studio/frontend/src/features/training/index.ts
@@ -12,6 +12,7 @@ export {
getTrainingRunDisplayTitle,
getTrainingRunModelSubtitle,
} from "./lib/run-display";
+export { parseBackendTrainingMethod } from "./lib/training-methods";
export { useTrainingHistorySidebarItems } from "./hooks/use-training-history-sidebar";
export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
export { useTrainingCompletionWatch } from "./hooks/use-training-completion-watch";
From ecd97a935a2c71f918b93653e36b5320fd8aa872 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Sun, 19 Jul 2026 04:54:17 -0700
Subject: [PATCH 015/255] test(version-compat): keep GRPO fake-run logits
finite on CPU (#7247)
* test(version-compat): keep GRPO fake-run logits finite on CPU
The GRPO fake-run test samples completions from a tiny untrained model on
CPU. Such a model can emit non-finite logits, so torch.multinomial inside
generate() intermittently raises "probability tensor contains either inf,
nan or element < 0" -- a nondeterministic sampling failure, not a regression
(the Trainer already fixes the seed, but CPU reduction order is not
bit-reproducible). Add a forward hook that sanitizes the LM head logits to a
finite bounded range before sampling, so the fake run reliably exercises the
whole train loop; the test checks the loop runs, not the numerics.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* test(version-compat): drop redundant nan_to_num bounds (clamp handles them)
* test(version-compat): scope GRPO finite-logits guard to the GRPO test
Only test_grpo_trains_on_cpu autoregressively samples completions, so it is
the only canary that can hit the non-finite-logits torch.multinomial crash.
Move the _guard_finite_logits hook out of the shared _load_plain() and into
test_grpo_trains_on_cpu so the SFT and DPO canaries keep asserting against the
model's true, unclamped logits.
---------
Co-authored-by: Daniel Han
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
.../version_compat/test_trl_fake_train_cpu.py | 36 +++++++++++++++++++
1 file changed, 36 insertions(+)
diff --git a/tests/version_compat/test_trl_fake_train_cpu.py b/tests/version_compat/test_trl_fake_train_cpu.py
index 4dae696282..2f7b469428 100644
--- a/tests/version_compat/test_trl_fake_train_cpu.py
+++ b/tests/version_compat/test_trl_fake_train_cpu.py
@@ -158,6 +158,37 @@ except Exception:
_MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM"
+def _guard_finite_logits(model):
+ """Keep the LM head logits finite so GRPO sampling can't crash.
+
+ ``test_grpo_trains_on_cpu`` samples completions from a tiny, *untrained*
+ random model on CPU. Driven autoregressively -- and nudged by the fake
+ reward's optimizer step between the two train steps -- such a model can emit
+ non-finite logits, so ``torch.multinomial`` inside ``generate()``
+ intermittently raises "probability tensor contains either `inf`, `nan` or
+ element < 0". That is a well-known nondeterministic sampling failure, not an
+ Unsloth/TRL regression: the Trainer already fixes the seed, but CPU reduction
+ order is not bit-reproducible, so the blow-up still surfaces every so often.
+
+ Sanitize the logits to a finite, bounded range (out of place, so autograd
+ stays valid) before they reach the sampler. This test asserts the train loop
+ runs end to end, not the (deliberately meaningless) numerics, so bounding the
+ logits changes nothing it checks while making the run reliable.
+ """
+
+ def _finite_logits_hook(_module, _inputs, output):
+ logits = getattr(output, "logits", None)
+ if logits is None:
+ return output
+ # nan_to_num maps nan -> 0 and the infinities to large finite values;
+ # clamp then bounds everything to [-30, 30].
+ output.logits = torch.nan_to_num(logits).clamp(-30.0, 30.0)
+ return output
+
+ model.register_forward_hook(_finite_logits_hook)
+ return model
+
+
def _load_plain():
"""Tiny plain HF model + tokenizer on CPU. Skips (not fails) if the model
cannot be fetched -- that is a network/hub issue, not an unsloth regression."""
@@ -233,6 +264,11 @@ def test_grpo_trains_on_cpu(tmp_path):
assert GRPOTrainer.__name__ == "UnslothGRPOTrainer", "GRPO patch did not apply"
model, tok = _load_plain()
+ # GRPO is the only canary that autoregressively samples completions, so it is
+ # the only one that can hit the non-finite-logits multinomial crash. Install
+ # the guard here (not in _load_plain) so the SFT/DPO canaries keep asserting
+ # against the model's true, unclamped outputs.
+ _guard_finite_logits(model)
ds = Dataset.from_list([{"prompt": "hi there"}] * 4)
cfg = GRPOConfig(
output_dir = str(tmp_path / "ci_grpo"),
From 5f1f30ec82d097d92d1093d1d557427dcbf079e6 Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Sun, 19 Jul 2026 09:46:22 -0300
Subject: [PATCH 016/255] Studio: GPU memory configuration for GGUF models
(#6414)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* Studio: GPU memory dropdown — llama.cpp --fit on and manual gpu-layers/cpu-moe
* Studio: simplify GPU memory changes (reuse ParamSlider, GPU_LAYERS_ALL, loadedGpuMemoryFields helper)
* Studio: GPU picker — choose which GPUs a GGUF model loads on (gpu_ids)
* Studio: simplify GPU picker (share /api/system fetch, validate gpu_ids)
* Studio: GPU picker review fixes (gate relative indices, no cross-model leak, validate, types)
* Studio: group GPU controls under a collapsible GPU section
* Studio: GPU feature review fixes (fix fit-ctx test, behavior-test the floor, comment accuracy)
* Studio: make GPU a top-level settings section (not nested under Model)
* Studio: flatten GPU controls into the Model section, group by GPU/context/generation
* Studio: move GPU Memory to the bottom of Model with its dependent controls beneath it
* Studio: move GPU Memory below Tensor Parallelism and GPUs below GPU Memory
* Studio: tighten GPU Memory and GPU Layers tooltip copy
* Studio: fix fit-mode context slider track-click, restore GPU Memory tooltip, shorten fit dropdown label
* Studio: GPU Memory tooltip one mode per line, briefer
* Studio: note HIP_VISIBLE_DEVICES (ROCm) in the GPUs picker tooltip
* Studio: narrow the GPU Memory dropdown to fit the shortened label
* Studio: use 'llama.cpp --fit' in the GPU Memory tooltip for consistency
* Studio: allow Tensor Parallelism in Manual GPU mode
* Studio: graduated MoE-on-CPU offload (--n-cpu-moe) replacing the all-or-nothing toggle
* Studio: size the MoE-offload slider for staged (deferred-load) models
* Studio: share one GGUF header walk for the context-length and MoE-count readers
* Studio: size the GPU Layers slider for staged models (one staged-header read)
* Studio: move Tensor Parallelism below the GPUs picker
* Studio: GPU split (--tensor-split) per-GPU model share in Manual mode
* Studio: tolerate whitespace in GPU split input, move it below GPU Layers
* Studio: rename the GPU split control to "Split ratio"
* Studio: Split ratio sends explicit even input; fix blank=free-VRAM (not even) copy
* Studio: tighten llama.cpp --fit VRAM margin with --fit-target 512
* Studio: GPU memory review fixes (rollback re-baseline, single-GPU TP gate, accurate copy)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: move Split ratio below MoE Layers on CPU
* Studio: address PR review (fix GPU-info hydration race, share fit context-length across load paths)
* Studio: address codex review (manual single-GPU TP guard, GPU-aware spec defaults in fit/manual, GGUF-only context/preference)
* Studio: address codex review round 2 (gpu_present seed, single-GPU tensor-split guard, staged manual-knob reset, strip inherited offload flags)
* Studio: address codex review round 3 (strip inherited --n-cpu-moe, CPU-fallback warning in Manual mode)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: address codex review round 4 (preserve pinned fit context across a later Apply)
* Studio: address codex review round 5 (honor GPU picker for diffusion GGUFs, clear fit pin on cross-model switch)
* Studio: preserve the pending GPU Memory mode when staging a model
* Studio: pin diffusion GPU device order and reset GPU-memory state for diffusion loads
* Studio: address codex review round 6 (fit-Auto rollback context, preserve manual non-tensor split modes, persist GPU mode on load not select)
* Studio: persist the applied GPU Memory mode, not the requested one (skip diffusion loads)
* Studio: replace Manual-mode split-ratio field with per-GPU layer sliders
* Studio: clarify per-GPU layer split hint for tensor-parallel mode
* Studio: address codex review round 7 (allow GGUF gpu_ids past the legacy guard, replay GPU-memory fields on respawn)
* Studio: address codex review round 8 (size the validate preflight like the load in fit mode, across both load paths)
* Studio: skip the training-OOM guard for llama.cpp --fit GGUF loads (they spill to RAM)
* Studio: drop the now-redundant compare-path validate sizing (the --fit guard skip makes it moot)
* Studio: address codex review round 9 (keep the training guard for fit loads, forward gpu_ids to validate, strip inherited manual tensor-split)
* Studio: address codex review round 10 (gate GPU-memory adoption on is_gguf, record manual knobs only in Manual mode)
* Studio: handle diffusion GGUFs symmetrically in the GPU Memory controls (preserve the standing mode preference, hide the inapplicable mode/TP controls)
* Studio: remember the GPU Memory settings per model
* Studio: consolidate --fit mode and Manual mode into a single Manual mode
* Studio: preserve the per-GPU layer split across GPU Layers changes
* Studio: trim overly long GPU Memory comments
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address GPU memory config review comments
* trim redundant GPU memory tests
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Reconcile manual-mode TP drops with the #6659 drop-site invariants
* Preserve quantized KV in manual --fit, charge GGUF companions in full, reconcile GPU pick on load
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Clear stale GPU baseline on non-GGUF loads so it can't read as dirty
* Fix no-context-shift test for the conditional -c flag
* Credit manual GPU-layer offload for cached HF GGUFs
* Reset per-model load knobs on GGUF quant switch
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Strip inherited tensor-split when manual ratio is cleared
* Match auto-load validation to safetensors placement
* Reset editable manual knobs after Auto GGUF loads
* Record a single device for diffusion GPU picks
* Reset per-model GPU knobs before applying saved settings
* Address review comments
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Guard manual tensor splits and keep remembered context on auto-load
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Snapshot compare knobs, seed splits from free VRAM, flag zero-offload loads
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Exempt CPU-only loads from the guard floor and harden compare and reseed paths
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Reach full offload from the layers slider and charge extras drafters in the guard
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Warm the GPU device cache before pick reconciles and disable staged GPU controls
* Align the training guard with inherited extras, spec mode, and compare targets
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Hide GPUs from companion-less zero-offload loads
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Size diffusion picks per device, own manual offload flags, reject XPU picks
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Drop tensor flags at zero layers and exempt CPU-pinned drafters
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Allowlist the zero-layer tensor parallel drop site
* Keep validate and load guards on the same extras and refresh stale baselines
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Drop mismatched manual tensor splits before launch
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gate XPU picks on the real backend field and harden split and hydration paths
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Weight full GPUs as zero, clamp split shares, and refine the zero-layer mask gate
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Carry fit context across mode changes and align drafter and picker gates
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Catch variant switches, uncached diffusion repos, and text-only mmproj skips
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Check companions on the first device and size native and remote zero-layer loads
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Replace the training guard's precise VRAM modeling with a conservative bound
* Baseline context pins on non-GGUF hydration and reprobe list-seeded staged GGUFs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Size manual splits by their largest share and preserve resolved context from Default
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Default-deny unsized required companions and price KV at the effective cache dtype
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Reserve MTP draft KV and MLA target-copy in the training guard
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Size tensor-parallel loads per device and show GPU controls for native GGUFs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Reserve MTP overhead for uncached remote GGUFs and the mmproj runtime factor
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Drop the training-coexistence VRAM estimation this PR added
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gate remembered load settings to GGUF picks
* Lock the remaining load-time controls during a staged load
* Clear the stale native-path token on compare loads
* Drop a stale guard reference from the zero-offload masking comment
* Seed GPU baselines from the rollback response and drop never-emitted offload flags
* Match validate's training guard to load and keep the native reload token
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim verbose GPU-memory comments
* Thread the variants header walk off the event loop, honor device pins on zero-offload, and hold staged GPU edits
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Honor manual placement and classify pinned zero-offload loads
* Close diffusion admission and status hydration gaps
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Check the actual diffusion GPU during training
* Align staged baselines and manual reload dedupe
* Fix GGUF placement and rollback state
* Harden manual GGUF placement boundaries
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Remove unused resolve_tensor_parallel import in llama_cpp.py
The name is used only in llama_server_args.py, routes/inference.py, and tests,
not in llama_cpp.py; the unused hoisted import trips the import-hoist verifier
in the source-lint CI job.
* Fix diffusion GPU dedup and training guard for non-numeric device tokens
The diffusion runner drives only its single lowest device and the backend
records that one device (self._gpu_ids = [sorted(gpu_ids)[0]]), but the reload
dedupe compared it against the full requested list, so a multi-GPU pick that
resolves to the same device forced a needless reload. Normalize the request the
same way for a loaded diffusion model in both _already_in_target_state and the
route _request_matches_loaded_settings.
The chat-during-training coexistence guard called int() on the single-device
token and hard-rejected when it could not parse. A non-numeric token (a CUDA
UUID / MIG handle) now sizes against the whole visible pool like the GGUF guard
instead of falsely blocking the load, and an empty token (a CPU-only runner such
as a CPU diffusion GGUF) is allowed outright since it uses no GPU VRAM.
* Tighten comments added by the GPU memory config changes
* Harden GGUF placement from independent review: VRAM sizing, diffusion TP reset, tensor_split validation
- Training coexistence guard: a single-device runner pinned through an
unresolvable UUID/MIG token was sized against the aggregate visible-VRAM pool,
so a load could pass on capacity it cannot use and then OOM active training.
Size against the worst-case visible device (min free) instead, keeping the
guard's documented default-deny contract. The empty-token (CPU-only runner)
allow path is unchanged.
- Diffusion startup: _start_diffusion_server now resets self._tensor_parallel to
False alongside the other placement resets. A prior tensor-parallel chat load
(process killed but not fully unload-reset) otherwise left /status misreporting
tensor parallelism and made an identical diffusion re-Apply reload against the
stale state.
- tensor_split: reject negative / non-finite / all-zero splits up front. They
were dropped at launch but still compared raw in the reload dedupe, so an
identical Apply reloaded indefinitely.
- Tests: the shared httpx stub was incomplete and, installed via setdefault
before real httpx loaded, broke a combined pytest run (collection errors on
httpx.Response). Import the real installed httpx instead.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen
Co-authored-by: danielhanchen
---
studio/backend/core/inference/llama_cpp.py | 715 +++++++++++++-
.../core/inference/llama_server_args.py | 27 +-
studio/backend/main.py | 14 +
studio/backend/models/inference.py | 144 ++-
studio/backend/routes/inference.py | 518 ++++++++---
studio/backend/routes/models.py | 6 +-
studio/backend/routes/training_vram.py | 42 +-
.../tests/test_chat_load_during_training.py | 333 ++++++-
studio/backend/tests/test_gguf_metadata.py | 73 ++
studio/backend/tests/test_gpu_memory_mode.py | 879 ++++++++++++++++++
studio/backend/tests/test_gpu_selection.py | 18 +-
.../tests/test_llama_cpp_no_context_shift.py | 12 +-
.../tests/test_llama_cpp_props_readback.py | 31 +-
.../backend/tests/test_llama_server_args.py | 45 +
studio/backend/tests/test_tensor_parallel.py | 5 +-
.../tests/test_tp_vision_regression.py | 17 +-
studio/backend/utils/models/gguf_metadata.py | 109 ++-
.../remembered-load-settings.ts | 24 +-
.../src/features/chat/api/chat-adapter.ts | 128 ++-
.../src/features/chat/api/chat-api.ts | 38 +-
.../frontend/src/features/chat/chat-page.tsx | 13 +-
.../src/features/chat/chat-settings-sheet.tsx | 450 ++++++++-
.../chat/hooks/use-chat-model-runtime.ts | 217 ++++-
.../hooks/use-staged-model-preparation.ts | 46 +-
.../lib/apply-inference-status-to-store.ts | 117 ++-
.../features/chat/presets/preset-policy.ts | 31 +
.../src/features/chat/shared-composer.tsx | 109 ++-
.../chat/stores/chat-runtime-store.ts | 369 +++++++-
.../frontend/src/features/chat/types/api.ts | 38 +
studio/frontend/src/hooks/use-gpu-info.ts | 186 ++--
studio/frontend/src/hooks/use-system.ts | 3 +
31 files changed, 4356 insertions(+), 401 deletions(-)
create mode 100644 studio/backend/tests/test_gpu_memory_mode.py
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index c9ab7eb83b..d7c7eed518 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -11,6 +11,7 @@ import atexit
import contextlib
import functools
import json
+import math
import os
import re
import struct
@@ -23,11 +24,22 @@ import sys
import threading
import time
from pathlib import Path
-from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union
+from typing import (
+ Callable,
+ Collection,
+ Generator,
+ Iterable,
+ List,
+ Literal,
+ Mapping,
+ Optional,
+ Union,
+)
import httpx
from core.inference.llama_server_args import (
+ _LAYER_OFFLOAD_FLAGS,
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
extra_args_disable_mmproj,
@@ -234,8 +246,7 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]":
return out
-# Plan-without-action re-prompt state (intent signal, caps, message) now lives
-# in tool_call_parser, imported above under its old aliases.
+# Plan-without-action re-prompt state now lives in tool_call_parser (imported above).
# Default max_tokens to the effective context when known. The floor is high
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
@@ -1431,7 +1442,10 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"})
-_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"})
+# Layer-offload override detection. Single-sourced from llama_server_args, which
+# also strips these (plus the MoE flags) from inherited extras; sharing the layer
+# set keeps detection and stripping from drifting.
+_GPU_OFFLOAD_OVERRIDE_FLAGS = _LAYER_OFFLOAD_FLAGS
_THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"})
@@ -1895,6 +1909,17 @@ class LlamaCppBackend:
self._cache_type_kv: Optional[str] = None
# Whether --split-mode tensor was applied on the active load.
self._tensor_parallel: bool = False
+ # GPU memory strategy applied on the active load ("auto"/"manual").
+ self._gpu_memory_mode: str = "auto"
+ # Manual-mode load options (echoed back so the UI round-trips them).
+ self._gpu_layers: int = -1
+ # MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none.
+ self._n_cpu_moe: int = 0
+ # Relative model share per GPU (--tensor-split), in GPU order; None =
+ # default (llama.cpp splits by free VRAM).
+ self._tensor_split: Optional[List[float]] = None
+ # User-picked physical GPU indices (None = automatic selection).
+ self._gpu_ids: Optional[List[int]] = None
# 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
@@ -1909,6 +1934,11 @@ class LlamaCppBackend:
self._spec_draft_n_max: Optional[int] = None
# KV-cache estimation fields (populated by _read_gguf_metadata)
self._n_layers: Optional[int] = None
+ # MoE metadata (populated by _read_gguf_metadata): expert count (>0 =
+ # MoE) and leading dense-layer count (offsets --n-cpu-moe, which counts
+ # from layer 0). See the n_moe_layers property.
+ self._n_experts: Optional[int] = None
+ self._leading_dense_block_count: Optional[int] = None
self._n_kv_heads: Optional[int] = None
self._n_kv_heads_by_layer: Optional[list[int]] = None
self._n_heads: Optional[int] = None
@@ -2329,6 +2359,79 @@ class LlamaCppBackend:
"""Whether --split-mode tensor is active on the loaded server."""
return self._tensor_parallel
+ @property
+ def gpu_memory_mode(self) -> str:
+ """Active GPU memory strategy: 'auto' or 'manual' (gpu_layers < 0 = Auto/--fit, >= 0 = pinned)."""
+ return self._gpu_memory_mode
+
+ @property
+ def gpu_layers(self) -> int:
+ """Requested --gpu-layers for manual mode (-1 when not manual)."""
+ return self._gpu_layers
+
+ @property
+ def n_cpu_moe(self) -> int:
+ """MoE expert layers manual mode kept on CPU (--n-cpu-moe); 0 = none."""
+ return self._n_cpu_moe
+
+ @property
+ def tensor_split(self) -> Optional[List[float]]:
+ """Manual-mode relative model share per GPU (--tensor-split); None =
+ default (split by free VRAM)."""
+ return self._tensor_split
+
+ @property
+ def gpu_ids(self) -> Optional[List[int]]:
+ """User-picked physical GPU indices, or None for automatic selection."""
+ return self._gpu_ids
+
+ @property
+ def n_layers(self) -> Optional[int]:
+ """Model layer count (GGUF block_count), or None if unknown."""
+ return self._n_layers
+
+ @property
+ def n_moe_layers(self) -> int:
+ """Number of MoE expert layers (the --n-cpu-moe ceiling), 0 if not MoE.
+
+ block_count minus the leading dense layers (which carry no experts):
+ --n-cpu-moe counts from layer 0, so those dense layers are no-ops.
+ """
+ if not self._n_experts or not self._n_layers:
+ return 0
+ return max(0, self._n_layers - (self._leading_dense_block_count or 0))
+
+ @staticmethod
+ def _resolve_cpu_moe_flag(
+ n_cpu_moe: int, n_moe_layers: int, leading_dense: int
+ ) -> Optional[int]:
+ """The --n-cpu-moe value (absolute first-N layers), or None to omit it.
+
+ Clamps the requested count to the model's MoE layers, then offsets past
+ the leading dense layers (--n-cpu-moe counts from layer 0). Returns None
+ for nothing-to-offload (0 requested) or a non-MoE model.
+ """
+ if n_cpu_moe <= 0 or n_moe_layers <= 0:
+ return None
+ return leading_dense + min(n_cpu_moe, n_moe_layers)
+
+ @staticmethod
+ def _sanitize_tensor_split(tensor_split: Optional[List[float]]) -> List[float]:
+ """Per-GPU shares with negative and non-finite entries clamped to 0.
+
+ A direct caller's negative entry would launch a placement different
+ from the ratio the UI showed, and inf would pass a plain ``> 0`` total
+ gate and emit ``--tensor-split inf,...``. Returns [] for input that
+ can't be read as floats (the length gate at the call site then drops
+ the split).
+ """
+ try:
+ return [
+ x if math.isfinite(x) and x > 0.0 else 0.0 for x in (float(v) for v in tensor_split)
+ ]
+ except (TypeError, ValueError, OverflowError):
+ return []
+
@property
def layer_preserves_tensor_intent(self) -> bool:
"""True when a downgraded tensor request kept this layer load multi-GPU."""
@@ -2530,6 +2633,7 @@ class LlamaCppBackend:
"spec_draft_n_max_flag": None,
"supports_kv_unified": False,
"supports_fit_ctx": False,
+ "supports_fit_target": False,
"supports_cache_ram": False,
"supports_ctx_checkpoints": False,
"supports_no_cache_prompt": False,
@@ -2549,6 +2653,7 @@ class LlamaCppBackend:
spec_draft_n_max_flag: Optional[str] = None
supports_kv_unified = False
supports_fit_ctx = False
+ supports_fit_target = False
supports_cache_ram = False
supports_ctx_checkpoints = False
supports_no_cache_prompt = False
@@ -2646,6 +2751,7 @@ class LlamaCppBackend:
supports_kv_unified = _is_real("--kv-unified")
supports_fit_ctx = _is_real("--fit-ctx")
+ supports_fit_target = _is_real("--fit-target")
supports_cache_ram = _is_real("--cache-ram")
supports_ctx_checkpoints = _is_real("--ctx-checkpoints")
supports_no_cache_prompt = _is_real("--no-cache-prompt")
@@ -2662,6 +2768,7 @@ class LlamaCppBackend:
"spec_draft_n_max_flag": spec_draft_n_max_flag,
"supports_kv_unified": supports_kv_unified,
"supports_fit_ctx": supports_fit_ctx,
+ "supports_fit_target": supports_fit_target,
"supports_cache_ram": supports_cache_ram,
"supports_ctx_checkpoints": supports_ctx_checkpoints,
"supports_no_cache_prompt": supports_no_cache_prompt,
@@ -2746,6 +2853,57 @@ class LlamaCppBackend:
except ValueError:
return None
+ @staticmethod
+ def _emit_child_gpu_visibility(env: dict, pinned: str) -> None:
+ """Write the child's GPU visibility mask (CUDA, plus the HIP mirror on
+ ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child
+ seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP
+ mask at different layers, so the same indices apply twice -- ROCR reduces
+ and re-indexes from 0, then a non-zero HIP pin points out of range, HIP
+ enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone
+ narrows correctly; clear any inherited ROCR mask so it can't double up."""
+ env["CUDA_VISIBLE_DEVICES"] = pinned
+ try:
+ import torch as _torch
+ if getattr(_torch.version, "hip", None) is not None:
+ env["HIP_VISIBLE_DEVICES"] = pinned
+ env.pop("ROCR_VISIBLE_DEVICES", None)
+ except Exception as e:
+ logger.debug("Failed to set ROCm visibility env vars for child: %s", e)
+
+ @staticmethod
+ def _pin_visible_gpu_order_for_split(env: dict) -> None:
+ """Pin the child's GPU enumeration to the picker's order for a manual
+ ``--tensor-split`` across the whole visible set. CUDA's default
+ FASTEST_FIRST enumeration applies the shares to the wrong cards on
+ heterogeneous hosts (#5025), and CUDA_DEVICE_ORDER only fixes the
+ numbering base: an inherited numeric visibility mask ALSO defines
+ enumeration order, so a reordered parent mask (CUDA_VISIBLE_DEVICES=3,1)
+ would still hand the shares to the wrong cards. The UI built the split
+ positionally over get_backend_visible_gpu_info's device list (ascending
+ physical via nvidia-smi, inherited mask order on the torch fallback), so
+ re-emit the same set in that report order -- not an assumed ascending
+ sort. The visible set itself never changes. No mask, an empty mask, or a
+ UUID/MIG mask (which resolves to None) is left alone -- the multi-GPU
+ controls are hidden for the latter."""
+ env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+ inherited = LlamaCppBackend._resolve_visible_physical_ids()
+ if not inherited:
+ return
+ order = None
+ try:
+ from utils.hardware import get_backend_visible_gpu_info
+ info = get_backend_visible_gpu_info()
+ if info.get("available") and info.get("index_kind") == "physical":
+ reported = [d["index"] for d in info.get("devices", [])]
+ if sorted(reported) == sorted(inherited):
+ order = reported
+ except Exception as e:
+ logger.debug("Could not read reported GPU order for split pin: %s", e)
+ if order is None:
+ order = sorted(inherited)
+ LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order))
+
@staticmethod
def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool:
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
@@ -3262,6 +3420,20 @@ class LlamaCppBackend:
# aborts a --split-mode tensor load, so it's dropped for the tensor attempt.
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
+ # Main-model placement settings that Manual mode owns. They must not leak
+ # from Studio's parent environment into llama-server and silently override
+ # the command assembled from the current request. Draft-model placement is
+ # intentionally separate and remains available to speculative decoding.
+ _MANUAL_PLACEMENT_ENV_VARS = (
+ "LLAMA_ARG_CPU_MOE",
+ "LLAMA_ARG_N_CPU_MOE",
+ "LLAMA_ARG_N_GPU_LAYERS",
+ "LLAMA_ARG_TENSOR_SPLIT",
+ "LLAMA_ARG_FIT",
+ "LLAMA_ARG_FIT_TARGET",
+ "LLAMA_ARG_FIT_CTX",
+ )
+
# (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.
@@ -3426,6 +3598,12 @@ class LlamaCppBackend:
return env
+ @classmethod
+ def _clear_manual_placement_env(cls, env: dict[str, str]) -> None:
+ """Remove inherited main-model placement owned by Manual mode."""
+ for name in cls._MANUAL_PLACEMENT_ENV_VARS:
+ env.pop(name, None)
+
@staticmethod
def _select_gpus(
model_size_bytes: int,
@@ -4246,6 +4424,8 @@ class LlamaCppBackend:
self._supports_preserve_thinking = False
self._supports_tools = False
self._n_layers = None
+ self._n_experts = None
+ self._leading_dense_block_count = None
self._n_kv_heads = None
self._n_kv_heads_by_layer = None
self._n_heads = None
@@ -4335,6 +4515,8 @@ class LlamaCppBackend:
arch_keys = {
f"{arch}.context_length": "context_length",
f"{arch}.block_count": "n_layers",
+ f"{arch}.expert_count": "n_experts",
+ f"{arch}.leading_dense_block_count": "leading_dense_block_count",
f"{arch}.attention.head_count_kv": "n_kv_heads",
f"{arch}.attention.head_count": "n_heads",
f"{arch}.embedding_length": "embedding_length",
@@ -4523,6 +4705,28 @@ class LlamaCppBackend:
return None
+ @staticmethod
+ def _diffusion_gpu_arg(gpu_ids: Optional[List[int]], *, cpu_only: bool = False) -> str:
+ """Device token passed to the diffusion visual-server child.
+
+ The visual engine replaces its child's CUDA visibility mask with this
+ token, so an unpinned load must carry forward the first token from the
+ parent's mask rather than turning a parent-relative ordinal into a new
+ physical selection.
+ """
+ if gpu_ids:
+ return str(sorted(gpu_ids)[0])
+ if cpu_only:
+ return ""
+ if "DG_GPU" in os.environ:
+ return os.environ["DG_GPU"]
+ parent_mask = os.environ.get("CUDA_VISIBLE_DEVICES")
+ if parent_mask:
+ first = next((token.strip() for token in parent_mask.split(",") if token.strip()), "")
+ if first and first != "-1":
+ return first
+ return "0"
+
def _start_diffusion_server(
self,
*,
@@ -4533,6 +4737,7 @@ class LlamaCppBackend:
model_identifier: str,
n_ctx: int,
extra_args: Optional[List[str]],
+ gpu_ids: Optional[List[int]] = None,
) -> bool:
"""Launch the OpenAI-compat diffusion shim (which drives the on-device
visual decoder) and wait for health. Presents the same /v1 + /health
@@ -4558,7 +4763,11 @@ class LlamaCppBackend:
# CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child
# CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default.
cpu_only = self._effective_gpu_count() == 0
- gpu = "" if cpu_only else os.environ.get("DG_GPU", "0")
+ # Honor the GPU picker first: the diffusion runner takes a single device,
+ # so use the lowest selected GPU (matches the sorted set recorded below, so
+ # the device used == the echoed gpu_ids[0]). With no pick, fall back to the
+ # CPU-only mask, else DG_GPU / 0.
+ gpu = self._diffusion_gpu_arg(gpu_ids, cpu_only = cpu_only)
cmd = list(shim_cmd) + [
"--gguf",
@@ -4586,6 +4795,11 @@ class LlamaCppBackend:
env.setdefault("UNSLOTH_ALLOW_CPU", "1")
env["DG_VISUAL_BIN"] = visual_bin
env["DG_GPU"] = gpu
+ if gpu_ids:
+ # The visual server remasks via CUDA_VISIBLE_DEVICES=; pin PCI
+ # order (as the llama-server path does) so the picked physical id maps
+ # to the GPU the picker showed, not CUDA's default fastest-first order.
+ env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH.
# (The zoo-package shim is an installed module and needs no PYTHONPATH change.)
if extra_pythonpath:
@@ -4631,6 +4845,23 @@ class LlamaCppBackend:
self._model_identifier = model_identifier
self._cache_type_kv = None
self._gpu_offload_active = True
+ # Diffusion doesn't use the llama.cpp GPU-memory knobs; reset them to
+ # defaults (the picked device is still recorded below) so /load, /status
+ # and reload dedup don't report a previous GGUF's manual settings.
+ self._gpu_memory_mode = "auto"
+ self._gpu_layers = -1
+ self._n_cpu_moe = 0
+ self._tensor_split = None
+ # Diffusion is never tensor-parallel; clear any state left by a prior TP
+ # chat load (load_model phase 1 only kills the process, it doesn't run
+ # the unload reset) so /status doesn't misreport TP and an identical
+ # re-Apply doesn't reload against stale tensor-parallel state.
+ self._tensor_parallel = False
+ # Record only the single device the runner actually uses (the lowest
+ # selected GPU, chosen above) -- not the whole pick. The diffusion runner
+ # is single-device, so echoing a multi-GPU list would misreport placement
+ # in /status and let a re-Apply dedup against GPUs the runner never used.
+ self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None
if hf_variant:
self._hf_variant = hf_variant
elif gguf_path:
@@ -5721,6 +5952,11 @@ class LlamaCppBackend:
speculative_type: Optional[str] = None,
spec_draft_n_max: Optional[int] = None,
tensor_parallel: bool = False,
+ gpu_memory_mode: Literal["auto", "manual"] = "auto",
+ gpu_layers: int = -1,
+ n_cpu_moe: int = 0,
+ tensor_split: Optional[List[float]] = None,
+ gpu_ids: Optional[List[int]] = None,
n_threads: Optional[int] = None,
n_gpu_layers: Optional[int] = None, # caller compat, unused
n_parallel: int = 1,
@@ -5753,6 +5989,14 @@ class LlamaCppBackend:
"speculative_type": speculative_type,
"spec_draft_n_max": spec_draft_n_max,
"tensor_parallel": tensor_parallel,
+ # GPU-memory placement: replayed on respawn so a server SIGKILL'd by
+ # GPU/RAM pressure reloads onto the same devices with the same
+ # offload, not the auto defaults.
+ "gpu_memory_mode": gpu_memory_mode,
+ "gpu_layers": gpu_layers,
+ "n_cpu_moe": n_cpu_moe,
+ "tensor_split": list(tensor_split) if tensor_split is not None else None,
+ "gpu_ids": list(gpu_ids) if gpu_ids is not None else None,
"n_threads": n_threads,
"n_gpu_layers": n_gpu_layers,
"n_parallel": n_parallel,
@@ -5779,6 +6023,11 @@ class LlamaCppBackend:
speculative_type = speculative_type,
spec_draft_n_max = spec_draft_n_max,
tensor_parallel = tensor_parallel,
+ gpu_memory_mode = gpu_memory_mode,
+ gpu_layers = gpu_layers,
+ n_cpu_moe = n_cpu_moe,
+ tensor_split = tensor_split,
+ gpu_ids = gpu_ids,
chat_template_override = chat_template_override,
extra_args = extra_args,
is_vision = is_vision,
@@ -5899,6 +6148,7 @@ class LlamaCppBackend:
model_identifier = model_identifier,
n_ctx = n_ctx,
extra_args = extra_args,
+ gpu_ids = gpu_ids,
)
if not binary:
@@ -5960,6 +6210,59 @@ class LlamaCppBackend:
# use the same helper so a healthy env-driven tensor server matches.
split_mode_override = parse_split_mode_override(extra_args)
tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel)
+ # gpu_layers=0 leaves nothing to split, yet --split-mode tensor or
+ # a per-GPU ratio still launches tensor mode -- and under the
+ # CPU-only mask below (no visible devices) that aborts the server
+ # instead of loading on CPU. Drop both here (nothing to split).
+ if gpu_memory_mode == "manual" and gpu_layers == 0:
+ if tensor_parallel or tensor_split:
+ logger.info(
+ "Manual gpu_layers=0: dropping tensor split/parallel "
+ "flags (nothing to split on the GPU)"
+ )
+ tensor_parallel = False
+ tensor_split = None
+ # Record the requested strategy for /status and the load
+ # response. 'manual' has no fallback, so the request value is the
+ # value actually applied.
+ self._gpu_memory_mode = gpu_memory_mode
+ # The layer/MoE/split knobs apply only with an explicit offload
+ # (manual + gpu_layers >= 0); else record defaults so /status and
+ # /load don't report knobs the server never applied.
+ if gpu_memory_mode == "manual" and gpu_layers >= 0:
+ self._gpu_layers = gpu_layers
+ self._n_cpu_moe = n_cpu_moe
+ self._tensor_split = tensor_split
+ else:
+ self._gpu_layers = -1
+ self._n_cpu_moe = 0
+ self._tensor_split = None
+ self._gpu_ids = sorted(gpu_ids) if gpu_ids else None
+ # Manual offload skips the TP planner but still emits --split-mode
+ # tensor at launch; drop it when fewer than 2 GPUs are in use --
+ # tensor split is a no-op there and aborts on some architectures.
+ # Done before the cache-drop below so a quantized KV survives.
+ if (
+ tensor_parallel
+ and gpu_memory_mode == "manual"
+ and gpu_layers >= 0
+ and self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2
+ ):
+ logger.info(
+ "Tensor parallelism requested in manual mode but fewer "
+ "than 2 GPUs are in use; ignoring (needs >= 2)."
+ )
+ tensor_parallel = False
+ # Drop TP for manual + Auto layers before the cache-drop below (like
+ # the <2-GPU guard above), so a requested quantized KV survives into
+ # the --fit load rather than being stripped for a tensor attempt.
+ if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:
+ logger.info(
+ "Manual mode with Auto layers hands memory management to "
+ "llama.cpp --fit, which is incompatible with tensor "
+ "parallelism; ignoring the tensor split."
+ )
+ tensor_parallel = False
# Tensor mode aborts on a quantized KV cache, so drop it for the
# tensor attempt (and strip any inherited/explicit --cache-type
# that would re-impose it when appended last). Layer split does
@@ -6040,6 +6343,12 @@ class LlamaCppBackend:
"Vision-capable GGUF loaded without a usable mmproj; "
"image input will be disabled for this session"
)
+ # Seed before the try: the except (GPU-selection failure ->
+ # --fit on) falls through to the launch which reads this, and the
+ # probe that assigns it may throw first. Captured before manual
+ # empty `gpus` so the speculative defaults stay GPU-aware and the
+ # CPU-fallback check still knows GPUs were present.
+ _detected_gpus: list[tuple[int, int]] = []
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).
@@ -6057,6 +6366,18 @@ class LlamaCppBackend:
_gpu_mem = self._get_gpu_memory(binary)
gpus = [(idx, free) for idx, free, _t in _gpu_mem]
total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
+ # GPU picker: restrict every mode to the chosen devices, so
+ # auto selection only considers them and manual mask to
+ # them (the env block below pins CUDA/HIP_VISIBLE_DEVICES).
+ if gpu_ids:
+ _picked = set(gpu_ids)
+ gpus = [g for g in gpus if g[0] in _picked]
+
+ # GPUs the model will run on -- captured before manual
+ # empty `gpus` to bypass the planner. bool() drives the
+ # GPU-aware speculative defaults; the list feeds the
+ # CPU-fallback check.
+ _detected_gpus = list(gpus)
def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION):
# Per-GPU usable budget for ranking: free - (1-frac)*total.
@@ -6088,6 +6409,44 @@ class LlamaCppBackend:
# GPU/VRAM-fit logic below may shrink it on limited HW.
max_available_ctx = self._context_length or effective_ctx
+ # Manual + Auto layers (the Manual default): hand memory
+ # management to llama.cpp's --fit. Emptying the probed GPU set
+ # no-ops the selection/TP planning below, leaving gpu_indices
+ # None (an explicit gpu_ids pick still pins below) and use_fit
+ # True. An explicit context is honored (--fit optimizes around
+ # it); 0 lets --fit size it.
+ if gpu_memory_mode == "manual" and gpu_layers < 0:
+ # Tensor parallelism was already dropped above (before the
+ # cache-drop), so a quantized KV survives into this --fit load.
+ gpus = []
+ effective_ctx = requested_ctx if requested_ctx > 0 else 0
+ original_ctx = effective_ctx
+ # --fit aborts under --split-mode tensor; a raw extras
+ # --split-mode/--tensor-split (appended last) would
+ # otherwise reach llama-server. Strip it like the TP
+ # downgrade does.
+ extra_args = strip_split_mode_only(extra_args)
+ elif gpu_memory_mode == "manual":
+ # Manual offload (--gpu-layers + --fit off): no automatic
+ # device masking (a gpu_ids pick still pins below) or
+ # context cap -- the user owns both. tensor_parallel is
+ # honored but skips the memory-based planner (gpus = []);
+ # the toggle just emits --split-mode tensor (split by free
+ # VRAM, or by the Split ratio if set).
+ gpus = []
+ effective_ctx = (
+ requested_ctx if requested_ctx > 0 else (self._context_length or 0)
+ )
+ original_ctx = effective_ctx
+ # Strip the user --split-mode when the toggle owns the split
+ # (TP engaged -> Studio emits --split-mode tensor) or when the
+ # user asked for tensor (which aborts on a single GPU even if
+ # the manual <2-GPU guard downgraded TP). Otherwise keep their
+ # non-tensor mode (row/none/layer) -- the toggle can't express
+ # those.
+ if tensor_parallel or split_mode_override == "tensor":
+ extra_args = strip_split_mode_only(extra_args)
+
# Will MTP engage? If so, auto-fit reserves draft-model VRAM.
# Mirrors _build_speculative_flags: forced mtp/mtp+ngram always
# engage; auto only on an MTP model >= 3B; ngram/off never. A
@@ -6175,7 +6534,10 @@ class LlamaCppBackend:
_extra_n_max = _extra_args_spec_draft_n_max(extra_args)
_mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max
if _mtp_eff_n_max is None:
- _mtp_eff_n_max = 2 if gpus else 3
+ # _detected_gpus (not gpus) so manual -- which empty
+ # gpus to bypass the planner -- keep the GPU draft depth the
+ # launch flags also use, instead of the CPU default.
+ _mtp_eff_n_max = 2 if _detected_gpus else 3
# Separate-drafter weights live on GPU (an embedded head is
# already in model_size). Size the drafter the launch loads, by
# precedence: extras --model-draft (last-wins), else Unsloth's
@@ -6313,7 +6675,8 @@ class LlamaCppBackend:
# honor it, cap only if it fits no combination. Auto (native):
# prefer fewer GPUs with reduced context (multi-GPU is slower).
gpu_indices, use_fit = None, True
- # Per-GPU weight proportions for tensor mode (None = even).
+ # Per-GPU weight proportions for tensor mode (None lets
+ # llama.cpp split by free VRAM).
tp_tensor_split: Optional[list[int]] = None
explicit_ctx = requested_ctx > 0
# Flat MTP reserve fraction: used only as the fallback when the
@@ -6388,7 +6751,12 @@ class LlamaCppBackend:
# GPUs below that reserve from the set up front (gpu_indices
# becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them).
tp_gpus = gpus
- if tensor_parallel:
+ # Manual mode owns the layer count and context, so it skips
+ # the memory-based planner; its toggle still emits
+ # --split-mode tensor below (split by free VRAM, or by the
+ # Split ratio if set). auto plans here.
+ plan_tp = tensor_parallel and gpu_memory_mode != "manual"
+ if plan_tp:
# Deterministic per-device compute buffer (replicated on
# every device in tensor mode); flat fallback when dims
# are unavailable. _plan_tensor_parallel uses the same.
@@ -6407,7 +6775,7 @@ class LlamaCppBackend:
# free yet have no budget left.
tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib]
- if tensor_parallel and len(tp_gpus) < 2:
+ if plan_tp and len(tp_gpus) < 2:
# Tensor parallelism needs >= 2 usable GPUs. On a single
# GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only
# or probe failed) it must not reach llama-server; and a
@@ -6823,6 +7191,12 @@ class LlamaCppBackend:
tp_tensor_split = None
effective_ctx = requested_ctx # fall back to original
+ # GPU picker: when no narrower subset was chosen (manual, or
+ # a failed/file-size selection), pin the whole picked set so the
+ # model can't spill onto an unpicked GPU.
+ if gpu_ids and gpu_indices is None:
+ gpu_indices = sorted(gpu_ids)
+
# Unified-memory APUs load weights into system RAM (under WSL the VM
# cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
# oversize load the OS would otherwise kill mid-flight. Base model
@@ -6859,8 +7233,6 @@ class LlamaCppBackend:
model_path,
"--port",
str(self._port),
- "-c",
- str(effective_ctx) if effective_ctx > 0 else "0",
"--parallel",
str(n_parallel),
"--flash-attn",
@@ -6868,6 +7240,17 @@ class LlamaCppBackend:
# Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length".
"--no-context-shift",
]
+ # A positive context is always passed (in auto-fit, --fit then
+ # optimizes the gpu-layer offload around it). When auto-fit has
+ # no explicit context, omit -c so --fit sizes it to fit VRAM:
+ # "-c 0" would instead pin the FULL native context (llama.cpp's
+ # -c handler sets fit_params_min_ctx = UINT32_MAX on value 0,
+ # disabling --fit's reduction). See gpu_memory_mode.
+ auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0
+ if effective_ctx > 0:
+ cmd.extend(["-c", str(effective_ctx)])
+ elif not auto_fit:
+ cmd.extend(["-c", "0"])
# 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
@@ -6879,7 +7262,63 @@ class LlamaCppBackend:
cmd.extend(["--alias", _alias])
fully_gpu_offloaded = False
- if use_fit:
+ # Set when a positional --tensor-split is emitted, so the env block
+ # can pin CUDA to PCI order even without a GPU subset (see below).
+ manual_tensor_split_emitted = False
+ if gpu_memory_mode == "manual" and gpu_layers >= 0:
+ # Pin the user's layer count and disable auto-fit. --fit off
+ # also means _ctx_integrity_flags must not add --fit-ctx.
+ use_fit = False
+ cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])
+ # Keep the first n_cpu_moe MoE layers' experts on CPU.
+ moe_flag = self._resolve_cpu_moe_flag(
+ n_cpu_moe,
+ self.n_moe_layers,
+ self._leading_dense_block_count or 0,
+ )
+ if moe_flag is not None:
+ cmd.extend(["--n-cpu-moe", str(moe_flag)])
+ elif n_cpu_moe:
+ # Requested on a dense model: nothing was emitted, so
+ # don't report a count llama-server never received.
+ self._n_cpu_moe = 0
+ # Distribute the model across GPUs by the user's per-GPU shares
+ # (--tensor-split). Works with layer split and tensor
+ # parallelism; --fit off means no fit/tensor abort. Only emit
+ # when >1 GPU is in use AND the list length matches that count:
+ # the field is hidden (not cleared) when the picker narrows to
+ # one, and a direct caller can send a stale ratio for a different
+ # GPU set. Studio drops any mismatch to the free-VRAM default
+ # (llama.cpp would silently zero-pad a short list, or abort past
+ # its 16-device cap).
+ _split_gpus = self._effective_gpu_count(gpu_indices)
+ if tensor_split and _split_gpus > 1:
+ # An all-zero/non-positive sanitized split assigns nothing
+ # anywhere, so fall through to the free-VRAM default in
+ # that case.
+ _sanitized_split = self._sanitize_tensor_split(tensor_split)
+ _split_total = sum(_sanitized_split)
+ if len(_sanitized_split) == _split_gpus and _split_total > 0:
+ cmd.extend(
+ ["--tensor-split", ",".join(f"{x:g}" for x in _sanitized_split)]
+ )
+ self._tensor_split = _sanitized_split
+ manual_tensor_split_emitted = True
+ else:
+ logger.warning(
+ "Dropping manual --tensor-split (%d entries for "
+ "%d GPUs, sanitized total %s); llama.cpp's "
+ "free-VRAM split applies instead",
+ len(tensor_split),
+ _split_gpus,
+ _split_total,
+ )
+ self._tensor_split = None
+ elif tensor_split:
+ # Single effective GPU: the split is never emitted, so
+ # don't report it as active via /status and /load.
+ self._tensor_split = None
+ elif use_fit:
cmd.extend(["--fit", "on"])
elif gpu_indices is not None:
# Fits on selected GPU(s) -- force all layers on GPU. --fit off is
@@ -6897,6 +7336,7 @@ class LlamaCppBackend:
self._ctx_integrity_flags(
n_parallel,
use_fit,
+ auto_fit,
requested_ctx,
effective_ctx,
server_caps,
@@ -6960,9 +7400,11 @@ class LlamaCppBackend:
self._cache_type_kv = None
# Tensor parallelism: split the model across GPUs by tensor
- # rather than by layer. Multi-GPU only -- a no-op on a single
- # GPU. Default (layer split) is left implicit by omitting the
- # flag. See llama.cpp --split-mode.
+ # rather than by layer. The UI only offers it on multi-GPU; a
+ # direct single-GPU caller is redundant (supported archs no-op,
+ # unsupported ones abort and the /load path retries layer split).
+ # Default (layer split) is left implicit by omitting the flag.
+ # See llama.cpp --split-mode.
if tensor_parallel:
cmd.extend(["--split-mode", "tensor"])
if tp_tensor_split and len(tp_tensor_split) > 1:
@@ -6994,7 +7436,7 @@ class LlamaCppBackend:
extra_args = extra_args,
model_identifier = model_identifier,
model_path = model_path,
- gpus = bool(gpus),
+ gpus = bool(_detected_gpus),
binary = binary,
mtp_draft_path = launch_mtp_draft_path,
)
@@ -7112,6 +7554,8 @@ class LlamaCppBackend:
# Library paths so llama-server finds its shared libs and CUDA DLLs.
env = self._llama_server_env_for_binary(binary)
+ if gpu_memory_mode == "manual":
+ self._clear_manual_placement_env(env)
# Omitting --threads relies on llama.cpp's physical-core default, so
# drop an inherited LLAMA_ARG_THREADS that would otherwise feed the
# arg handler and silently force hardware_concurrency(). #5692
@@ -7170,28 +7614,39 @@ class LlamaCppBackend:
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so
# set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device
# (above), not here.
- if gpu_indices is not None and not is_vulkan_backend:
- pinned = ",".join(str(i) for i in gpu_indices)
- env["CUDA_VISIBLE_DEVICES"] = pinned
- try:
- import torch as _torch
- if getattr(_torch.version, "hip", None) is not None:
- env["HIP_VISIBLE_DEVICES"] = pinned
- # Do NOT also set ROCR_VISIBLE_DEVICES to the same
- # value. ROCR_VISIBLE_DEVICES filters at the HSA/ROCr
- # layer and HIP_VISIBLE_DEVICES at the HIP layer, so
- # setting both with the same physical indices applies
- # the mask twice: ROCR reduces the visible set and
- # re-indexes it from 0, then HIP indexes into the
- # already-reduced set. A single non-zero pin (e.g.
- # "1") then points out of range at the HIP layer, HIP
- # enumerates 0 devices, and llama.cpp falls back to
- # CPU ("ggml_cuda_init: no ROCm-capable device is
- # detected"). The HIP mask alone narrows correctly;
- # clear any inherited ROCR mask so it can't double up.
- env.pop("ROCR_VISIBLE_DEVICES", None)
- except Exception as e:
- logger.debug("Failed to set ROCm visibility env vars for child: %s", e)
+ # A deliberate zero-offload load with no GPU companions runs
+ # entirely on CPU, yet a visible CUDA device still costs the child
+ # ~0.5 GB (context + compute scratch) that the CPU-only
+ # classification below reports as free. Hide the GPUs so the load
+ # is exactly what it claims: zero VRAM (verified: GPU stays at idle
+ # baseline and generation runs). Companion loads keep the normal
+ # masking, and a user device pin (in extras or an inherited
+ # LLAMA_ARG_DEVICE) keeps control of its own devices -- the child
+ # aborts on a pin it can't see. The draft-device forms count too:
+ # llama-server parses them even with no drafter loaded.
+ _cpu_only_zero_offload = (
+ gpu_memory_mode == "manual"
+ and gpu_layers == 0
+ and not is_vulkan_backend
+ and not self._zero_offload_keeps_gpu_visible(cmd, env)
+ )
+ if _cpu_only_zero_offload:
+ self._emit_child_gpu_visibility(env, "-1")
+ elif gpu_indices is not None and not is_vulkan_backend:
+ # When the user picked GPUs by index, align CUDA's ordering
+ # with the PCI-bus order the picker enumerated (nvidia-smi),
+ # so "GPU 1" in the UI is GPU 1 to llama.cpp -- not CUDA's
+ # default FASTEST_FIRST order (#5025).
+ if gpu_ids:
+ env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+ self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices))
+ elif manual_tensor_split_emitted and not is_vulkan_backend:
+ # A manual per-GPU ratio across ALL GPUs (no explicit pick, so
+ # no CUDA_VISIBLE_DEVICES mask above): the UI built the
+ # --tensor-split list in ascending physical/PCI index order,
+ # so pin the child's enumeration to that order too. The whole
+ # visible set stays in use; only its ordering is fixed.
+ self._pin_visible_gpu_order_for_split(env)
# Captured before any text-only fallback strips it from cmd.
launched_with_mmproj = "--mmproj" in cmd
@@ -7350,7 +7805,6 @@ class LlamaCppBackend:
self._effective_context_length = (
effective_ctx if effective_ctx > 0 else self._context_length
)
- self._reconcile_effective_ctx_with_server()
self._max_context_length = (
max_available_ctx if max_available_ctx > 0 else self._effective_context_length
)
@@ -7535,6 +7989,10 @@ class LlamaCppBackend:
"session; run 'unsloth studio update' to enable vision."
)
cmd = self._strip_mmproj_args(_last_spawn_cmd)
+ # This retry bypasses _spawn_and_wait, so refresh the
+ # launched-argv snapshot itself -- the zero-offload
+ # classification below must not see the stripped --mmproj.
+ _last_spawn_cmd = list(cmd)
self._is_vision = False
self._mmproj_has_audio = False
self._start_llama_process(cmd, env)
@@ -7566,6 +8024,13 @@ class LlamaCppBackend:
self._healthy = True
self._commit_effective_parallel_slots(n_parallel)
+ # Server is up: adopt the real per-request context it allocated
+ # -- the length --fit chose, or a --parallel slot split -- so the
+ # reported context_length matches reality. (Querying /props
+ # before the spawn above always failed; the seeded value was the
+ # requested/native length.)
+ self._reconcile_effective_ctx_with_server()
+
# Commit caller intent only after _healthy=True so a failed start
# can't poison the next inheritance check. None keeps prior, []
# clears, list sets. Source records hf_variant for the route's
@@ -7580,11 +8045,24 @@ class LlamaCppBackend:
self._mtp_runtime_fallback_active = _mtp_active_for_launched_server
self._start_mtp_crash_watchdog()
- # Catch silent CPU fallback when GPU was intended (#5106).
- self._gpu_offload_active = self._classify_gpu_offload(
- gpu_indices is not None or use_fit, gpus or []
- )
- if self._gpu_offload_active is False:
+ # Catch silent CPU fallback when GPU was intended (#5106). Manual
+ # offload (no picker) leaves gpu_indices None and use_fit False, so
+ # include its GPU-layer intent; use the preserved probe since
+ # auto-layers/manual empty `gpus`. A deliberate zero-offload load
+ # classifies by its launched argv instead: the main model is
+ # CPU-only by construction and must read False (not None), or
+ # training needlessly unloads a server holding no VRAM.
+ _deliberate_cpu_only = gpu_memory_mode == "manual" and gpu_layers == 0
+ if _deliberate_cpu_only:
+ self._gpu_offload_active = self._zero_offload_gpu_flag(
+ _last_spawn_cmd, _detected_gpus, env
+ )
+ else:
+ self._gpu_offload_active = self._classify_gpu_offload(
+ gpu_indices is not None or use_fit or gpu_memory_mode == "manual",
+ _detected_gpus,
+ )
+ if self._gpu_offload_active is False and not _deliberate_cpu_only:
logger.warning(
"llama-server appears to have loaded the model entirely "
"on CPU even though Unsloth detected at least one GPU. "
@@ -7947,6 +8425,11 @@ class LlamaCppBackend:
gguf_path: Optional[str] = None,
spec_draft_n_max: Optional[int] = None,
tensor_parallel: bool = False,
+ gpu_memory_mode: Literal["auto", "manual"] = "auto",
+ gpu_layers: int = -1,
+ n_cpu_moe: int = 0,
+ tensor_split: Optional[List[float]] = None,
+ gpu_ids: Optional[List[int]] = None,
mtp_draft_path: Optional[str] = None,
preserve_multi_gpu_on_layer: bool = False,
) -> bool:
@@ -8003,6 +8486,38 @@ class LlamaCppBackend:
):
return False
+ # The diffusion runner is mode-agnostic (always "auto", ignores the
+ # layer/MoE/split knobs), so a standing manual preference in the
+ # request must not force a needless reload -- only the GPU pick matters.
+ if not self._is_diffusion:
+ # A GPU-memory-mode flip (Unsloth / manual) must always reload.
+ if self._gpu_memory_mode != gpu_memory_mode:
+ return False
+ # Manual: a layer-count change always reloads (covers Auto(-1) <-> a
+ # pinned count); MoE/split only matter with an explicit offload.
+ if gpu_memory_mode == "manual" and (
+ self._gpu_layers != gpu_layers
+ or (
+ gpu_layers >= 0
+ and (
+ self._n_cpu_moe != n_cpu_moe
+ or (self._tensor_split or None) != (tensor_split or None)
+ )
+ )
+ ):
+ return False
+ # A changed GPU pick must reload (compare order-insensitively; None/[]
+ # both mean automatic). The diffusion runner collapses a multi-GPU pick
+ # to its single lowest device, so self._gpu_ids holds just that device;
+ # normalize the request the same way, or a multi-GPU pick that resolves
+ # to the same device needlessly reloads.
+ if self._is_diffusion:
+ requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None
+ else:
+ requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None
+ if (self._gpu_ids or None) != requested_gpu_pick:
+ return False
+
# Compare on the canonical requested mode. With --spec-type in
# extra_args the backend stores None; mirror that here.
if _extra_args_set_spec_type(extra_args):
@@ -8071,6 +8586,78 @@ class LlamaCppBackend:
return None
return classify_gpu_offload_lines(self._stdout_lines)
+ @staticmethod
+ def _cmd_has_gpu_companion(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool:
+ """True when the argv/env carries a GPU companion: any --mmproj form, or
+ a drafter (Studio's --model-draft, the extras aliases, or the
+ LLAMA_ARG_SPEC_DRAFT_* env) -- these offload to the GPU regardless of
+ the main ``--gpu-layers``. A drafter explicitly forced to CPU
+ (--spec-draft-ngl 0 / --spec-draft-device cpu) doesn't count."""
+ if any(str(a).startswith("--mmproj") for a in cmd):
+ return True
+ if _extra_args_mtp_draft_path(cmd, env) is None:
+ return False
+ return not _extra_args_draft_offloaded_to_cpu(cmd, env)
+
+ @staticmethod
+ def _zero_offload_keeps_gpu_visible(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool:
+ """Whether a zero-layer launch still has a reason to use visible GPUs.
+
+ Keep this shared by child masking and post-launch residency bookkeeping:
+ a device pin, surviving tensor mode, mmproj, or GPU drafter prevents the
+ launch from being a confirmed zero-VRAM server.
+ """
+ return (
+ LlamaCppBackend._cmd_has_gpu_device_pin(cmd, env)
+ or _effective_tensor_parallel(cmd, False, env)
+ or LlamaCppBackend._cmd_has_gpu_companion(cmd, env)
+ )
+
+ @staticmethod
+ def _cmd_has_gpu_device_pin(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool:
+ """True when the effective main or draft ``--device`` pin names a GPU."""
+ main_flags = {"--device", "-dev"}
+ draft_flags = {"--spec-draft-device", "-devd", "--device-draft"}
+ last_main: Optional[str] = None
+ last_draft: Optional[str] = None
+ args = [str(arg) for arg in cmd]
+ for index, raw in enumerate(args):
+ flag, equals, inline = raw.partition("=")
+ if flag not in main_flags and flag not in draft_flags:
+ continue
+ value = inline if equals else (args[index + 1] if index + 1 < len(args) else "")
+ if flag in main_flags:
+ last_main = value
+ else:
+ last_draft = value
+ if last_main is None:
+ last_main = (env or {}).get("LLAMA_ARG_DEVICE")
+
+ def _names_gpu(value: Optional[str]) -> bool:
+ if value is None:
+ return False
+ devices = [item.strip().lower() for item in value.split(",") if item.strip()]
+ return not devices or any(item not in ("cpu", "none") for item in devices)
+
+ return _names_gpu(last_main) or _names_gpu(last_draft)
+
+ @staticmethod
+ def _zero_offload_gpu_flag(
+ spawn_cmd: list,
+ detected_gpus: list,
+ env: Optional[Mapping[str, str]] = None,
+ ) -> Optional[bool]:
+ """GPU-residency flag for a deliberate manual zero-offload load. The
+ main model is CPU-only by construction, but device pins, tensor mode,
+ mmproj, and GPU drafters can still make the server hold VRAM. The counted
+ offload classifier cannot see those allocations. This uses the same
+ predicate as the launch-time zero-VRAM mask; None means no GPU signal."""
+ if not detected_gpus:
+ return None
+ if LlamaCppBackend._is_vulkan_backend():
+ return True
+ return LlamaCppBackend._zero_offload_keeps_gpu_visible(spawn_cmd, env)
+
def load_cancelled(self) -> bool:
"""True if a load was cancelled (e.g. via unload/_cancel_event) and not
yet consumed by the next load_model. Lets the tensor->layer fallback
@@ -8114,11 +8701,18 @@ class LlamaCppBackend:
self._supports_tools = False
self._cache_type_kv = None
self._tensor_parallel = False
+ self._gpu_memory_mode = "auto"
+ self._gpu_layers = -1
+ self._n_cpu_moe = 0
+ self._tensor_split = None
+ self._gpu_ids = None
self._layer_preserves_tensor_intent = False
self._speculative_type = None
self._requested_spec_mode = None
self._spec_draft_n_max = None
self._n_layers = None
+ self._n_experts = None
+ self._leading_dense_block_count = None
self._n_kv_heads = None
self._n_kv_heads_by_layer = None
self._n_heads = None
@@ -8181,6 +8775,10 @@ class LlamaCppBackend:
# Clear healthy so a /load during the replacement's warm-up can't
# short-circuit against the previous server's health (#5401).
self._healthy = False
+ # Reset to unknown so the training guard treats the next (still
+ # loading) server as VRAM-resident rather than reading the killed
+ # server's stale zero-offload flag until the health probe reclassifies.
+ self._gpu_offload_active = None
# Drives _wait_for_vram_settle in the next load_model; set in finally
# so both in-process and frontend Apply paths record the kill.
self._last_kill_monotonic = time.monotonic()
@@ -8785,7 +9383,12 @@ class LlamaCppBackend:
@staticmethod
def _ctx_integrity_flags(
- n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict
+ n_parallel: int,
+ use_fit: bool,
+ auto_fit: bool,
+ requested_ctx: int,
+ effective_ctx: int,
+ caps: dict,
) -> list[str]:
"""Flags that keep the per-request window equal to the advertised ctx.
@@ -8793,14 +9396,28 @@ class LlamaCppBackend:
``--kv-unified`` default, silently splitting ``-c`` into per-slot
windows of ``-c / N``; restore the shared pool so one request can use
the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step
- at an explicitly requested ctx (default floor is 4096) so it offloads
- or fails instead of silently shrinking the window.
+ at an explicitly requested ctx so it offloads or fails instead of
+ silently shrinking the window. The 8192 auto-floor and the tighter
+ ``--fit-target`` margin apply only under Manual + Auto (``auto_fit``),
+ which omits ``-c``: on the legacy auto path ``-c 0`` already pins the
+ native window and ``--fit-ctx 8192`` would override it down to 8192.
"""
flags: list[str] = []
if n_parallel > 1 and caps.get("supports_kv_unified"):
flags.append("--kv-unified")
- if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"):
- flags.extend(["--fit-ctx", str(effective_ctx)])
+ if use_fit and caps.get("supports_fit_ctx"):
+ if requested_ctx > 0 and effective_ctx > 0:
+ # Floor the fit step at the explicitly requested ctx.
+ flags.extend(["--fit-ctx", str(effective_ctx)])
+ elif auto_fit:
+ # Manual + Auto omits -c, so floor at 8192 so --fit doesn't
+ # shrink the window below a usable size.
+ flags.extend(["--fit-ctx", "8192"])
+ if use_fit and auto_fit and caps.get("supports_fit_target"):
+ # llama.cpp's --fit leaves 1 GiB free per device by default;
+ # tighten that to 512 MiB so it packs more of the model onto
+ # the GPU before spilling to system RAM.
+ flags.extend(["--fit-target", "512"])
return flags
def _query_server_n_ctx(self) -> Optional[int]:
diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py
index 70d0dc774d..e72e10e071 100644
--- a/studio/backend/core/inference/llama_server_args.py
+++ b/studio/backend/core/inference/llama_server_args.py
@@ -186,12 +186,25 @@ _SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"})
_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"})
_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
+# GPU-offload flags. Stripped only when the GPU Memory mode owns offload
+# (manual emits --fit / --gpu-layers / --n-cpu-moe); in auto, a user's
+# inherited -ngl is respected (the offload_overridden path), so this group is
+# opt-in, not default. Layer flags are shared with llama_cpp's override
+# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them).
+_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset(
+ {"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}
+)
+_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"})
+_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS
+
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS
)
# Shadowing flags that take no value -- strip the flag only, not the next token.
-_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"})
+_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
+ {"--spec-default", "--jinja", "--no-jinja", "-cmoe", "--cpu-moe"}
+)
def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
@@ -424,6 +437,8 @@ def strip_shadowing_flags(
strip_spec: bool = True,
strip_template: bool = True,
strip_split_mode: bool = True,
+ strip_tensor_split: bool = False,
+ strip_offload: bool = False,
) -> list[str]:
"""Strip flags that shadow first-class Unsloth settings.
@@ -432,6 +447,12 @@ def strip_shadowing_flags(
(same for cache / spec / template / split-mode). Each ``strip_*``
toggle controls one group; the route only strips groups whose
first-class field the caller actually supplied.
+
+ ``strip_split_mode`` removes both ``--split-mode`` and the coupled
+ ``--tensor-split`` (the Tensor Parallelism toggle owns the whole split).
+ ``strip_tensor_split`` removes ``--tensor-split`` *alone*, so manual mode can
+ replace an inherited per-GPU ratio while leaving the user's ``--split-mode``
+ row/none/layer choice intact.
"""
shadowing: set[str] = set()
if strip_context:
@@ -444,6 +465,10 @@ def strip_shadowing_flags(
shadowing |= _TEMPLATE_FLAGS
if strip_split_mode:
shadowing |= _SPLIT_SHADOWING_FLAGS
+ if strip_tensor_split:
+ shadowing |= _TENSOR_SPLIT_FLAGS
+ if strip_offload:
+ shadowing |= _OFFLOAD_SHADOWING_FLAGS
tokens = [str(a) for a in (args or [])]
out: list[str] = []
diff --git a/studio/backend/main.py b/studio/backend/main.py
index 4797764ce7..81d4c16e52 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -1156,9 +1156,23 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]:
enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct")
enriched_devices.append(enriched_dev)
+ # Whether GGUF loads accept an explicit gpu_ids pick: /load and
+ # /validate 400 picks on XPU hosts (no visibility mask speaks torch-xpu
+ # ordinals) and on Vulkan-only builds (--device pins ggml's own
+ # ordinals), so the picker must not offer them.
+ try:
+ from core.inference.llama_cpp import LlamaCppBackend
+ from utils.hardware import DeviceType, get_device
+ gpu_ids_supported = (
+ get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend()
+ )
+ except Exception as e:
+ logger.debug(f"Could not resolve gpu_ids support: {e}")
+ gpu_ids_supported = True
gpu_info = {
"available": visibility_info.get("available", False),
"devices": enriched_devices,
+ "gguf_gpu_ids_supported": gpu_ids_supported,
}
_system_gpu_cache = (time.monotonic(), gpu_info)
return gpu_info
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index f3ae0f70df..d51d35189b 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -64,7 +64,7 @@ class LoadRequest(BaseModel):
)
gpu_ids: Optional[List[int]] = Field(
None,
- description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
+ description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.",
)
speculative_type: Optional[str] = Field(
None,
@@ -100,6 +100,66 @@ class LoadRequest(BaseModel):
"No effect on a single GPU. Ignored for non-GGUF models."
),
)
+ gpu_memory_mode: Literal["auto", "manual"] = Field(
+ "auto",
+ description = (
+ "GPU memory strategy for GGUF models. 'auto' (default): Unsloth "
+ "selects GPUs and caps context to fit VRAM. 'manual': you own the "
+ "offload. Leave gpu_layers at -1 (Auto) to hand memory management to "
+ "llama.cpp's --fit (no device masking, no context auto-reduce, no "
+ "gpu-layer/tensor-split planning); set gpu_layers >= 0 to pin layers "
+ "and n_cpu_moe yourself (--fit off), with tensor_parallel still "
+ "applying (split by free VRAM unless tensor_split is set, no planner). "
+ "Ignored for non-GGUF."
+ ),
+ )
+ gpu_layers: int = Field(
+ -1,
+ ge = -1,
+ description = (
+ "Manual mode only: number of layers to offload to the GPU "
+ "(--gpu-layers, with --fit off). A value >= the model's layer count "
+ "offloads all of them. -1 = Auto: hand layer + context sizing to "
+ "llama.cpp's --fit. Ignored unless gpu_memory_mode is 'manual'."
+ ),
+ )
+ n_cpu_moe: int = Field(
+ 0,
+ ge = 0,
+ description = (
+ "Manual mode only: keep the first N MoE expert layers on the CPU "
+ "(--n-cpu-moe) to save VRAM on MoE models. 0 = none, N = number of "
+ "MoE layers offloaded (the backend offsets past any leading dense "
+ "layers). Ignored unless gpu_memory_mode is 'manual' with gpu_layers >= 0."
+ ),
+ )
+ tensor_split: Optional[List[float]] = Field(
+ None,
+ description = (
+ "Manual mode only: relative share of the model per GPU (--tensor-split), "
+ "in the order of the GPUs in use, e.g. [2, 1] for 2:1. Omit it to let "
+ "llama.cpp use its default, which splits by free VRAM. Any list given is "
+ "passed through as-is, so send [1, 1] to force an even split. Ignored "
+ "unless gpu_memory_mode is 'manual' with gpu_layers >= 0."
+ ),
+ )
+
+ @field_validator("tensor_split")
+ @classmethod
+ def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Optional[List[float]]:
+ # A negative / non-finite / all-zero split is silently dropped at launch
+ # (stored as None) yet still compared raw in the reload dedupe, so an
+ # identical Apply reloads forever. Reject it up front; [] = no split.
+ if not value:
+ return value
+ import math
+
+ if any((not math.isfinite(v)) or v < 0 for v in value):
+ raise ValueError("tensor_split entries must be finite and non-negative")
+ if sum(value) <= 0:
+ raise ValueError("tensor_split must have a positive total")
+ return value
+
llama_extra_args: Optional[List[str]] = Field(
None,
description = (
@@ -133,6 +193,14 @@ class ValidateModelRequest(BaseModel):
max_seq_length: int = Field(0, ge = 0, le = 1048576)
load_in_4bit: bool = Field(True)
gpu_ids: Optional[List[int]] = Field(None)
+ gpu_memory_mode: Literal["auto", "manual"] = Field(
+ "auto",
+ description = (
+ "GGUF GPU-memory strategy intended for the follow-up load. Manual "
+ "placement bypasses the training coexistence estimate: Auto layers "
+ "delegate fitting to llama.cpp, while explicit layers are user-owned."
+ ),
+ )
include_context_length: bool = Field(
False,
description = "Also read the native context length from the local GGUF header. "
@@ -188,6 +256,16 @@ class ValidateModelResponse(BaseModel):
description = "Native training context length, read from the GGUF header when the file "
"is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
)
+ layer_count: Optional[int] = Field(
+ None,
+ description = "Total layer count (GGUF block_count), the manual gpu-layers ceiling, read "
+ "from the header alongside context_length; None when not read.",
+ )
+ moe_layer_count: Optional[int] = Field(
+ None,
+ description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF "
+ "header alongside context_length; 0 for dense models, None when not read.",
+ )
# Additive fields; the consuming consent dialog ships in a follow-up frontend PR.
requires_transformers_upgrade: bool = Field(
False,
@@ -333,6 +411,34 @@ class LoadResponse(BaseModel):
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
+ gpu_memory_mode: Literal["auto", "manual"] = Field(
+ "auto",
+ description = "Active GPU memory strategy ('auto' or 'manual').",
+ )
+ gpu_layers: int = Field(
+ -1,
+ description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).",
+ )
+ n_cpu_moe: int = Field(
+ 0,
+ description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.",
+ )
+ tensor_split: Optional[List[float]] = Field(
+ None,
+ description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).",
+ )
+ n_layers: Optional[int] = Field(
+ None,
+ description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.",
+ )
+ n_moe_layers: int = Field(
+ 0,
+ description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.",
+ )
+ gpu_ids: Optional[List[int]] = Field(
+ None,
+ description = "Physical GPU indices the model is pinned to, or None for automatic selection.",
+ )
class UnloadResponse(BaseModel):
@@ -461,6 +567,42 @@ class InferenceStatusResponse(BaseModel):
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
+ gpu_memory_mode: Literal["auto", "manual"] = Field(
+ "auto",
+ description = "Active GPU memory strategy ('auto' or 'manual').",
+ )
+ gpu_layers: int = Field(
+ -1,
+ description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).",
+ )
+ n_cpu_moe: int = Field(
+ 0,
+ description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.",
+ )
+ tensor_split: Optional[List[float]] = Field(
+ None,
+ description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).",
+ )
+ requested_context_length: Optional[int] = Field(
+ None,
+ description = (
+ "The n_ctx the active GGUF load was invoked with (0 = Auto). Lets the "
+ "UI re-seed a Manual + Auto-layers context pin on hydration, where "
+ "context_length only exposes the resolved value. None for non-GGUF."
+ ),
+ )
+ n_layers: Optional[int] = Field(
+ None,
+ description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.",
+ )
+ n_moe_layers: int = Field(
+ 0,
+ description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.",
+ )
+ gpu_ids: Optional[List[int]] = Field(
+ None,
+ description = "Physical GPU indices the model is pinned to, or None for automatic selection.",
+ )
llama_cpp_supports_mtp: bool = Field(
True,
description = (
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 3d527bf317..136e4f7645 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -13,7 +13,7 @@ from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request, status
from fastapi.responses import StreamingResponse, JSONResponse, Response
from starlette.requests import ClientDisconnect
-from typing import Any, Callable, List, Optional, Union
+from typing import Any, Callable, List, Literal, Optional, Union
import json
import httpx
from loggers import get_logger
@@ -3115,13 +3115,16 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]:
def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool:
- """Whether an inherited --split-mode should be stripped on reload.
+ """Whether an inherited --split-mode (and its coupled --tensor-split) should
+ be stripped on reload.
The binary Tensor Parallelism toggle can't carry --split-mode's row/none/
layer modes, so only strip when the toggle overrides it: tensor being turned
on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes
- survive. Shared by the inheritance strip and the already-loaded stale check
- so they agree on what reload would do.
+ survive. A manual per-GPU ratio is handled by _should_strip_tensor_split,
+ which strips only --tensor-split so the inherited mode is kept. Shared by the
+ inheritance strip and the already-loaded stale check so they agree on what
+ reload would do.
"""
fields_set = getattr(request, "model_fields_set", set())
return "tensor_parallel" in fields_set and (
@@ -3129,6 +3132,25 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[
)
+def _should_strip_tensor_split(request: LoadRequest) -> bool:
+ """Whether an inherited --tensor-split alone should be stripped on reload.
+
+ Manual explicit offload (gpu_layers >= 0) owns the per-GPU split: with a ratio
+ it emits its own --tensor-split (an inherited one, appended last, would
+ override it), and with the ratio cleared it wants llama.cpp's default
+ free-VRAM split. Either way an inherited --tensor-split must go, else the
+ cleared case silently keeps the stale ratio while status reports None.
+ Unlike _should_strip_split_mode this leaves --split-mode untouched, so a
+ user's row/none/layer mode survives a Studio split-ratio edit. When the
+ Tensor Parallelism toggle IS overriding the mode, _should_strip_split_mode
+ (called alongside this at every site) strips --split-mode anyway.
+ """
+ return (
+ getattr(request, "gpu_memory_mode", "auto") == "manual"
+ and getattr(request, "gpu_layers", -1) >= 0
+ )
+
+
def _carry_preserved_tensor_intent(
*, preserved: bool, same_model: bool, explicit_drop: bool
) -> bool:
@@ -3187,12 +3209,44 @@ def _request_matches_loaded_settings(
else strip_shadowing_flags(
backend_extra,
strip_split_mode = _should_strip_split_mode(request, backend_extra),
+ strip_tensor_split = _should_strip_tensor_split(request),
+ strip_offload = request.gpu_memory_mode == "manual",
)
)
if not _tensor_parallel_matches_loaded(
effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
):
return False
+ # The diffusion runner is mode-agnostic (it always reports "auto" and ignores
+ # the layer/MoE/split knobs), so a standing manual preference in the request
+ # must not force a needless reload -- only the GPU pick matters.
+ if not llama_backend.is_diffusion:
+ if request.gpu_memory_mode != llama_backend.gpu_memory_mode:
+ return False
+ # Manual: a layer-count change always reloads; MoE/split only matter with
+ # an explicit offload (gpu_layers >= 0), so a leftover value under Auto
+ # must not force one. Mirrors LlamaCppBackend._already_in_target_state.
+ if request.gpu_memory_mode == "manual" and (
+ request.gpu_layers != llama_backend.gpu_layers
+ or (
+ request.gpu_layers >= 0
+ and (
+ request.n_cpu_moe != llama_backend.n_cpu_moe
+ or (request.tensor_split or None) != (llama_backend.tensor_split or None)
+ )
+ )
+ ):
+ return False
+ # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU
+ # request to its single lowest device (it drives one device only), so the
+ # backend records just that device; compare the request the same way, or a
+ # multi-GPU pick that resolves to the same device needlessly reloads.
+ if llama_backend.is_diffusion:
+ _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None
+ else:
+ _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None
+ if _req_gpu_ids != llama_backend.gpu_ids:
+ 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
@@ -3235,14 +3289,17 @@ def _request_matches_loaded_settings(
# contain any shadow flag, so the reload path strips them rather than
# leaving a stale override in effect. (backend_extra computed above.)
if request.llama_extra_args is None:
- # Mirror the reload's conditional split-mode strip, so a preserved
- # non-tensor mode (row/none/layer) isn't seen as stale and doesn't
- # trigger a needless reload of a healthy server.
+ # Mirror the reload's conditional strips, so a preserved non-tensor mode
+ # (row/none/layer) isn't seen as stale and doesn't trigger a needless
+ # reload of a healthy server, while an inherited offload/ratio flag that
+ # the reload *would* strip is correctly seen as stale.
if (
backend_extra
and strip_shadowing_flags(
backend_extra,
strip_split_mode = _should_strip_split_mode(request, backend_extra),
+ strip_tensor_split = _should_strip_tensor_split(request),
+ strip_offload = request.gpu_memory_mode == "manual",
)
!= backend_extra
):
@@ -3861,6 +3918,46 @@ def _estimate_gguf_required_gb(
return None
+def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]:
+ """Classify a GGUF as diffusion, normal, or unknown before it is loaded.
+
+ ``None`` is important here: a remote GGUF whose header is not cached can
+ still be routed to the single-GPU diffusion runner after download. Treating
+ that case as normal would let Manual mode skip the training guard even
+ though the runner ignores Manual's llama-server placement controls.
+ """
+ identity = " ".join(
+ str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file")
+ ).lower()
+ if "diffusion" in identity:
+ return True
+
+ try:
+ main = getattr(config, "gguf_file", None)
+ if not (main and Path(main).is_file()):
+ repo = getattr(config, "gguf_hf_repo", None)
+ variant = getattr(config, "gguf_variant", None)
+ if repo and variant:
+ from hub.utils.gguf import resolve_local_gguf_path
+ main = resolve_local_gguf_path(repo, variant)
+ if not main or not Path(main).is_file():
+ return None
+
+ probe = LlamaCppBackend()
+ probe._read_gguf_metadata(str(main))
+ if probe.is_diffusion:
+ return True
+ # A successfully decoded architecture proves that this is a normal
+ # llama-server GGUF. No architecture means the lightweight probe could
+ # not establish the routing decision, so preserve the unknown state.
+ if getattr(probe, "_architecture", None):
+ return False
+ return None
+ except Exception as e:
+ logger.debug("Could not identify diffusion GGUF for training guard: %s", e)
+ return None
+
+
def _guard_chat_load_against_training(
config: ModelConfig,
*,
@@ -3871,11 +3968,19 @@ def _guard_chat_load_against_training(
requested_gpu_ids: Optional[List[int]],
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
+ gpu_memory_mode: Literal["auto", "manual"] = "auto",
) -> None:
- """Refuse loading a local chat model that would OOM an active training run.
+ """Protect active training from automatically placed chat-model loads.
+
No-op when training is inactive or unknown. `load_in_4bit` must be the
- effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the
- model would not fit alongside training."""
+ effective quantization (see _effective_load_in_4bit). Manual chat-GGUF
+ placement is an explicit override: Auto layers delegate fitting to
+ llama.cpp's ``--fit`` and pinned layers are owned by the user, so neither is
+ estimated here. Diffusion is still guarded because its mode-agnostic runner
+ ignores those controls and uses one GPU. An unclassified GGUF is guarded as
+ potentially diffusion until its local header proves otherwise. Other loads
+ raise HTTP 409 when they would not fit beside training.
+ """
from core.training import get_training_backend
from routes.training_vram import can_load_chat_during_training
@@ -3887,6 +3992,19 @@ def _guard_chat_load_against_training(
return
is_gguf = bool(getattr(config, "is_gguf", False))
+ diffusion_kind = _classify_diffusion_gguf(config) if is_gguf else False
+ if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False:
+ return
+
+ diffusion_gpu = None
+ if is_gguf and diffusion_kind is not False:
+ # Use the same token selection as the runner: an explicit pick wins,
+ # followed by DG_GPU, the first parent-visible token, then GPU 0.
+ diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg(
+ requested_gpu_ids,
+ cpu_only = LlamaCppBackend._effective_gpu_count() == 0,
+ )
+
required_override_gb = (
_estimate_gguf_required_gb(
config,
@@ -3907,6 +4025,7 @@ def _guard_chat_load_against_training(
requested_gpu_ids = requested_gpu_ids,
is_gguf = is_gguf,
required_override_gb = required_override_gb,
+ single_device_gpu = diffusion_gpu,
)
if ok:
return
@@ -3934,6 +4053,98 @@ def _guard_chat_load_against_training(
raise HTTPException(status_code = 409, detail = detail)
+def _resolve_inherited_extra_args(
+ request,
+ config: ModelConfig,
+ model_identifier: str,
+ extra_llama_args: Optional[list[str]],
+ effective_chat_template_override: Optional[str] = None,
+) -> Optional[list[str]]:
+ """Effective pass-through extras for a GGUF request that omitted the field:
+ the previous same-model load's extras, shadow-stripped, so a settings-Apply
+ reload (which does not round-trip the extras field) keeps them (#5401)."""
+ if getattr(request, "llama_extra_args", None) is not None:
+ return extra_llama_args
+ if not getattr(config, "is_gguf", False):
+ return extra_llama_args
+ llama_backend = get_llama_cpp_backend()
+ if not llama_backend.extra_args:
+ return extra_llama_args
+ # Inherit the previous load's extras (the chat-settings Apply path doesn't
+ # round-trip them; an explicit [] still clears). Gated on (model_identifier,
+ # hf_variant) to refuse cross-model pickup, and shadowing flags are
+ # stripped so an inherited override can't win the last-wins CLI
+ # parse against a freshly-supplied first-class field.
+ source = llama_backend.extra_args_source
+ # Compare against the resolved variant, not the request field: callers
+ # commonly omit gguf_variant for local ``.gguf`` paths and HF auto-pick
+ # flows. ``config.gguf_variant`` is the variant load_model was actually
+ # invoked with, so both sides of the comparison key off the same string.
+ resolved_variant = (config.gguf_variant or "").lower()
+ request_variant = (request.gguf_variant or "").lower()
+ stored_variant = (source[1] or "").lower() if source else ""
+ same_model = bool(source and source[0] and source[0].lower() == model_identifier.lower())
+ if request.gguf_variant:
+ variant_mismatch = request_variant != stored_variant
+ else:
+ variant_mismatch = bool(stored_variant and resolved_variant != stored_variant)
+ same_source = same_model and not variant_mismatch
+ if not same_source:
+ logger.info(
+ "Not inheriting llama_extra_args: stored args came from %s, loading %s",
+ source,
+ (model_identifier, resolved_variant),
+ )
+ # Cross-model: clear explicitly so the backend doesn't
+ # inherit via "no opinion" semantics.
+ extra_llama_args = []
+ else:
+ # Strip only the groups whose first-class field was set by the caller, so
+ # an inherited --chat-template-file survives an Apply that omits
+ # chat_template_override. A bundled family template (e.g. gemma-4) counts as
+ # a first-class template even when the request omits chat_template_override,
+ # so strip the inherited --chat-template-file then too -- else the stale arg
+ # (appended last) shadows the bundled template while Studio reports its caps.
+ fields_set = getattr(request, "model_fields_set", set())
+ stripped = strip_shadowing_flags(
+ llama_backend.extra_args,
+ strip_context = "max_seq_length" in fields_set,
+ strip_cache = "cache_type_kv" in fields_set,
+ strip_spec = ("speculative_type" in fields_set or "spec_draft_n_max" in fields_set),
+ strip_template = (
+ "chat_template_override" in fields_set
+ or effective_chat_template_override is not None
+ ),
+ strip_split_mode = _should_strip_split_mode(request, llama_backend.extra_args),
+ # manual + per-GPU ratio emits its own --tensor-split; drop
+ # an inherited one (appended last would override it) while
+ # keeping the user's --split-mode row/none/layer choice.
+ strip_tensor_split = _should_strip_tensor_split(request),
+ # manual emits its own --fit/--gpu-layers, so an inherited offload flag
+ # must not last-wins-override it. auto leaves a user's inherited -ngl
+ # alone. getattr: a validate request reuses this resolver, no offload fields.
+ strip_offload = getattr(request, "gpu_memory_mode", "auto") == "manual",
+ )
+ try:
+ extra_llama_args = validate_extra_args(stripped)
+ except ValueError:
+ # Shouldn't happen on already-validated args; degrade to
+ # no-extras rather than 400 if managed flags changed.
+ logger.warning(
+ "Stored llama_extra_args failed revalidation; loading without them: %s",
+ stripped,
+ )
+ extra_llama_args = []
+ else:
+ if extra_llama_args:
+ logger.info(
+ "Inheriting llama_extra_args from previous "
+ "load (same model, shadow-stripped): %s",
+ extra_llama_args,
+ )
+ return extra_llama_args
+
+
def _model_json_response(model, status_code: int = 200) -> Response:
"""Serialize a pydantic response once via pydantic-core.
@@ -4040,6 +4251,35 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
None if request.llama_extra_args is None else extra_llama_args
)
+ # Manual mode owns the offload flags: strip them from EXPLICIT extras
+ # too (the inherited path already does), or a last-wins --gpu-layers /
+ # --fit in extras re-enables GPU offload on a load status reports as
+ # CPU-only. Manual + per-GPU ratio owns --tensor-split the same way.
+ if request.gpu_memory_mode == "manual" and extra_llama_args:
+ _stripped_explicit = strip_shadowing_flags(
+ extra_llama_args,
+ strip_context = False,
+ strip_cache = False,
+ strip_spec = False,
+ strip_template = False,
+ strip_split_mode = False,
+ strip_tensor_split = _should_strip_tensor_split(request),
+ strip_offload = True,
+ )
+ if _stripped_explicit != extra_llama_args:
+ logger.info(
+ "Manual GPU memory owns the offload flags; stripping them "
+ "from explicit llama_extra_args: %s -> %s",
+ extra_llama_args,
+ _stripped_explicit,
+ )
+ extra_llama_args = _stripped_explicit
+
+ # Keep every downstream consumer on the normalized explicit list. In
+ # particular, the already-loaded comparator must not compare the raw
+ # request's managed offload flags against the stripped launch state.
+ request = request.model_copy(update = {"llama_extra_args": extra_llama_args})
+
model_identifier, model_log_label, native_grant_backed = (
_resolve_model_identifier_for_request(request, operation = "load-model")
)
@@ -4121,6 +4361,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
+ gpu_memory_mode = llama_backend.gpu_memory_mode,
+ gpu_layers = llama_backend.gpu_layers,
+ n_cpu_moe = llama_backend.n_cpu_moe,
+ tensor_split = llama_backend.tensor_split,
+ n_layers = llama_backend.n_layers,
+ n_moe_layers = llama_backend.n_moe_layers,
+ gpu_ids = llama_backend.gpu_ids,
)
else:
if (
@@ -4187,12 +4434,41 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Normalize gpu_ids: empty list means auto-selection, same as None
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
- # Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409.
+ # GGUF supports gpu_ids: validate the pick up front (before the training
+ # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects
+ # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts
+ # are rejected outright: the picker's indices are torch-xpu ordinals neither
+ # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin
+ # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device.
if config.is_gguf and effective_gpu_ids is not None:
- raise HTTPException(
- status_code = 400,
- detail = "gpu_ids is not supported for GGUF models yet.",
- )
+ from utils.hardware import DeviceType, get_device
+ from utils.hardware.hardware import resolve_requested_gpu_ids
+
+ if get_device() == DeviceType.XPU:
+ raise HTTPException(
+ status_code = 400,
+ detail = (
+ "GPU selection (gpu_ids) is not supported on Intel XPU. "
+ "Omit gpu_ids to use all devices."
+ ),
+ )
+ # Same reasoning for a Vulkan-only build: --device pins ggml's own
+ # Vulkan ordinals, so a physical pick can land on the wrong card on
+ # masked or non-contiguous hosts.
+ if LlamaCppBackend._is_vulkan_backend():
+ raise HTTPException(
+ status_code = 400,
+ detail = (
+ "GPU selection (gpu_ids) is not supported with a Vulkan "
+ "llama.cpp build: physical GPU ids have no defined "
+ "mapping to Vulkan device ordinals. Omit gpu_ids to use "
+ "all devices."
+ ),
+ )
+ try:
+ resolve_requested_gpu_ids(effective_gpu_ids)
+ except ValueError as exc:
+ raise HTTPException(status_code = 400, detail = str(exc)) from exc
if not config.is_gguf and _mlx_distributed_launch_detected():
raise HTTPException(
status_code = 400,
@@ -4222,8 +4498,20 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
"architectures)"
)
- # Refuse a load that would OOM active training, before the unload step below
- # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work.
+ # Inherit the previous same-model load's pass-through extras when this
+ # request omits the field (a settings-Apply reload doesn't round-trip
+ # them); shadow-stripped so an inherited flag can't override a
+ # first-class field the caller did set (#5401).
+ extra_llama_args = _resolve_inherited_extra_args(
+ request,
+ config,
+ model_identifier,
+ extra_llama_args,
+ effective_chat_template_override,
+ )
+
+ # Apply the training coexistence policy before the unload step below
+ # frees the resident model. Off-loop: the default-mode guard does sync work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
@@ -4234,6 +4522,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
requested_gpu_ids = effective_gpu_ids,
llama_extra_args = extra_llama_args,
n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1),
+ gpu_memory_mode = request.gpu_memory_mode,
)
# ── GGUF path: load via llama-server ──────────────────────
@@ -4245,84 +4534,6 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
from core.inference.llama_cpp import gguf_load_in_flight
gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo))
- # Inherit llama_extra_args from the previous load when the request
- # omits the field (the chat-settings Apply path doesn't round-trip
- # them; explicit [] still clears). Gated on (model_identifier,
- # hf_variant) to refuse cross-model pickup, and shadowing flags are
- # stripped so an inherited override can't win the last-wins CLI
- # parse against a freshly-supplied first-class field.
- if request.llama_extra_args is None and llama_backend.extra_args:
- source = llama_backend.extra_args_source
- # Compare against the resolved variant, not the request
- # field: callers commonly omit gguf_variant for local
- # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_
- # variant`` is the variant load_model was actually
- # invoked with (see the HF / local branches below), so
- # both sides of the comparison key off the same string.
- resolved_variant = (config.gguf_variant or "").lower()
- request_variant = (request.gguf_variant or "").lower()
- stored_variant = (source[1] or "").lower() if source else ""
- same_model = bool(
- source and source[0] and source[0].lower() == model_identifier.lower()
- )
- if request.gguf_variant:
- variant_mismatch = request_variant != stored_variant
- else:
- variant_mismatch = bool(stored_variant and resolved_variant != stored_variant)
- same_source = same_model and not variant_mismatch
- if not same_source:
- logger.info(
- "Not inheriting llama_extra_args: stored args came from %s, loading %s",
- source,
- (model_identifier, resolved_variant),
- )
- # Cross-model: clear explicitly so the backend doesn't
- # inherit via "no opinion" semantics.
- extra_llama_args = []
- else:
- # Strip only the groups whose first-class field was set by
- # the caller, so an inherited --chat-template-file survives
- # an Apply that omits chat_template_override. A bundled family
- # template (e.g. the gemma-4 override) is an effective
- # first-class template setting even when the raw request
- # omits chat_template_override, so strip the inherited
- # --chat-template-file in that case too -- otherwise the stale
- # extra arg (appended last) shadows the bundled template while
- # Unsloth reports the bundled template's capabilities.
- fields_set = getattr(request, "model_fields_set", set())
- stripped = strip_shadowing_flags(
- llama_backend.extra_args,
- strip_context = "max_seq_length" in fields_set,
- strip_cache = "cache_type_kv" in fields_set,
- strip_spec = (
- "speculative_type" in fields_set or "spec_draft_n_max" in fields_set
- ),
- strip_template = (
- "chat_template_override" in fields_set
- or effective_chat_template_override is not None
- ),
- strip_split_mode = _should_strip_split_mode(
- request, llama_backend.extra_args
- ),
- )
- try:
- extra_llama_args = validate_extra_args(stripped)
- except ValueError:
- # Shouldn't happen on already-validated args; degrade to
- # no-extras rather than 400 if managed flags changed.
- logger.warning(
- "Stored llama_extra_args failed revalidation; loading without them: %s",
- stripped,
- )
- extra_llama_args = []
- else:
- if extra_llama_args:
- logger.info(
- "Inheriting llama_extra_args from previous "
- "load (same model, shadow-stripped): %s",
- extra_llama_args,
- )
-
# Block cache writes that would race the download manager. This runs
# after pass-through argument inheritance so a carried --no-mmproj
# changes the companion requirement exactly as it does for the load.
@@ -4370,6 +4581,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
cache_type_kv = request.cache_type_kv,
speculative_type = request.speculative_type,
spec_draft_n_max = request.spec_draft_n_max,
+ gpu_memory_mode = request.gpu_memory_mode,
+ gpu_layers = request.gpu_layers,
+ n_cpu_moe = request.n_cpu_moe,
+ tensor_split = request.tensor_split,
+ gpu_ids = effective_gpu_ids,
n_parallel = _n_parallel,
)
if config.gguf_hf_repo:
@@ -4537,6 +4753,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
+ gpu_memory_mode = llama_backend.gpu_memory_mode,
+ gpu_layers = llama_backend.gpu_layers,
+ n_cpu_moe = llama_backend.n_cpu_moe,
+ tensor_split = llama_backend.tensor_split,
+ n_layers = llama_backend.n_layers,
+ n_moe_layers = llama_backend.n_moe_layers,
+ gpu_ids = llama_backend.gpu_ids,
)
# ── Standard path: load via Unsloth/transformers ──────────
@@ -4795,7 +5018,9 @@ def _requires_security_review_for_model(
@router.post("/validate", response_model = ValidateModelResponse)
async def validate_model(
- request: ValidateModelRequest, current_subject: str = Depends(get_current_subject)
+ request: ValidateModelRequest,
+ fastapi_request: Request = None,
+ current_subject: str = Depends(get_current_subject),
):
"""
Lightweight validation endpoint for model identifiers.
@@ -4823,15 +5048,39 @@ async def validate_model(
detail = f"Invalid model identifier: {model_log_label}",
)
- # Refuse early (before the frontend unloads to load this) if it can't fit
- # alongside training, using the same settings /load uses so they agree.
+ # Apply the same training coexistence policy as /load before the frontend
+ # unloads the current model.
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
- # Mirror /load: reject GGUF + gpu_ids before the guard so both return 400.
+ # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is
+ # a clean 400) before the guard sizes the model against training VRAM.
+ # XPU-host picks are rejected like /load (no defined mapping from the
+ # picker's torch-xpu ordinals to the launcher's device spaces).
if config.is_gguf and effective_gpu_ids is not None:
- raise HTTPException(
- status_code = 400,
- detail = "gpu_ids is not supported for GGUF models yet.",
- )
+ from utils.hardware import DeviceType, get_device
+ from utils.hardware.hardware import resolve_requested_gpu_ids
+
+ if get_device() == DeviceType.XPU:
+ raise HTTPException(
+ status_code = 400,
+ detail = (
+ "GPU selection (gpu_ids) is not supported on Intel XPU. "
+ "Omit gpu_ids to use all devices."
+ ),
+ )
+ if LlamaCppBackend._is_vulkan_backend():
+ raise HTTPException(
+ status_code = 400,
+ detail = (
+ "GPU selection (gpu_ids) is not supported with a Vulkan "
+ "llama.cpp build: physical GPU ids have no defined "
+ "mapping to Vulkan device ordinals. Omit gpu_ids to use "
+ "all devices."
+ ),
+ )
+ try:
+ resolve_requested_gpu_ids(effective_gpu_ids)
+ except ValueError as exc:
+ raise HTTPException(status_code = 400, detail = str(exc)) from exc
effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
# Both checks cover the [adapter, base] set (matching the scan route and workers):
@@ -4895,16 +5144,32 @@ async def validate_model(
latest_tier_active_for, config.identifier, request.hf_token
):
effective_load_in_4bit = False
- # Off-loop: guard does sync nvidia-smi / HF work.
- await asyncio.to_thread(
- _guard_chat_load_against_training,
- config,
- model_identifier = model_identifier,
- hf_token = request.hf_token,
- load_in_4bit = effective_load_in_4bit,
- max_seq_length = request.max_seq_length,
- requested_gpu_ids = effective_gpu_ids,
- )
+ # A metadata-only probe just reads the GGUF header and allocates no VRAM,
+ # so it must not be refused by the training guard. Real loads validate
+ # without include_context_length and /load applies the guard again.
+ if not request.include_context_length:
+ # Match /load's inherited llama.cpp extras and parallel slot count so
+ # validation cannot pass a smaller estimate than the subsequent load.
+ effective_extra_args = _resolve_inherited_extra_args(
+ request, config, model_identifier, None
+ )
+ # Off-loop: guard does sync nvidia-smi / HF work.
+ await asyncio.to_thread(
+ _guard_chat_load_against_training,
+ config,
+ model_identifier = model_identifier,
+ hf_token = request.hf_token,
+ load_in_4bit = effective_load_in_4bit,
+ max_seq_length = request.max_seq_length,
+ requested_gpu_ids = effective_gpu_ids,
+ llama_extra_args = effective_extra_args,
+ n_parallel = (
+ getattr(fastapi_request.app.state, "llama_parallel_slots", 1)
+ if fastapi_request is not None
+ else 1
+ ),
+ gpu_memory_mode = request.gpu_memory_mode,
+ )
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
@@ -4918,10 +5183,15 @@ async def validate_model(
# Native context length, read from the local GGUF header when present.
# Lets the staged ("Load on selection" off) flow populate the context
# slider before the GPU load; None until the file is downloaded.
+ # Staged header dims (one read): native context, total layer count, and
+ # MoE expert-layer count -- let the staged flow size the context, GPU-
+ # layers and manual --n-cpu-moe sliders before the load.
context_length: Optional[int] = None
+ layer_count: Optional[int] = None
+ moe_layer_count: Optional[int] = None
if request.include_context_length and is_gguf:
from hub.utils.gguf import resolve_local_gguf_path
- from utils.models.gguf_metadata import read_gguf_context_length
+ from utils.models.gguf_metadata import read_gguf_staged_dims
# Best-effort: a header-read failure must never fail validation of an
# otherwise-valid model (the outer except turns it into a 400).
@@ -4937,9 +5207,15 @@ async def validate_model(
model_identifier, request.gguf_variant
)
if local_gguf:
- context_length = read_gguf_context_length(local_gguf)
+ # Header walk reads tokenizer arrays for dense models (tens of
+ # ms); keep it off the event loop.
+ dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf)
+ if dims:
+ context_length = dims["context_length"]
+ layer_count = dims["layer_count"]
+ moe_layer_count = dims["moe_layer_count"]
except Exception as e:
- logger.debug("Context-length probe failed for %s: %s", model_log_label, e)
+ logger.debug("Header probe failed for %s: %s", model_log_label, e)
return ValidateModelResponse(
valid = True,
@@ -4954,6 +5230,8 @@ async def validate_model(
requires_trust_remote_code = requires_trust_remote_code,
requires_security_review = requires_security_review,
context_length = context_length,
+ layer_count = layer_count,
+ moe_layer_count = moe_layer_count,
requires_transformers_upgrade = transformers_upgrade is not None,
transformers_upgrade = transformers_upgrade,
)
@@ -5593,6 +5871,14 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
+ gpu_memory_mode = llama_backend.gpu_memory_mode,
+ gpu_layers = llama_backend.gpu_layers,
+ n_cpu_moe = llama_backend.n_cpu_moe,
+ tensor_split = llama_backend.tensor_split,
+ requested_context_length = llama_backend.requested_n_ctx,
+ n_layers = llama_backend.n_layers,
+ n_moe_layers = llama_backend.n_moe_layers,
+ gpu_ids = llama_backend.gpu_ids,
llama_cpp_supports_mtp = _supports_mtp,
spec_fallback_reason = llama_backend.spec_fallback_reason,
llama_cpp_prebuilt_stale = _stale,
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index bb321695cd..0806c2f513 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -2731,7 +2731,11 @@ async def get_gguf_variants(
],
has_vision = response.has_vision,
default_variant = response.default_variant,
- context_length = _read_native_context_length(repo_id, is_local = local),
+ # The header walk reads tokenizer arrays on dense models (tens of
+ # ms per uncached file); keep it off the event loop.
+ context_length = await asyncio.to_thread(
+ _read_native_context_length, repo_id, is_local = local
+ ),
)
except HTTPException:
raise
diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py
index fb361d3359..fd96fe2175 100644
--- a/studio/backend/routes/training_vram.py
+++ b/studio/backend/routes/training_vram.py
@@ -197,15 +197,18 @@ def can_load_chat_during_training(
requested_gpu_ids: Optional[List[int]],
is_gguf: bool = False,
required_override_gb: Optional[float] = None,
+ single_device_gpu: Optional[str] = None,
) -> Tuple[bool, Dict[str, Any]]:
"""Decide if a NEW chat model can load without OOMing active training (inverse
of can_keep_chat_during_training: training is already resident, so size the
chat model against the free VRAM that remains). Sizes/places it the same way
the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an
even-share per-GPU floor for device_map="balanced"; GGUF sizes from
- required_override_gb over the visible pool. `load_in_4bit` must be effective
- (LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any
- CUDA case it can't size, so a load never OOMs training."""
+ required_override_gb over the visible pool. ``single_device_gpu`` is the
+ exact physical device token selected by a single-device runner.
+ `load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA
+ allows the load; default-deny on any CUDA case it can't size, so a load never
+ OOMs training."""
try:
from utils.hardware import (
DeviceType,
@@ -251,26 +254,49 @@ def can_load_chat_during_training(
}
# Explicit GPUs, or GGUF: size directly and check live free VRAM.
+ if single_device_gpu is not None:
+ mode = "single_device"
+ elif is_gguf:
+ mode = "gguf"
+ else:
+ mode = "explicit"
required_gb = required_override_gb
if required_gb is None:
required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs)
if required_gb is None:
- mode = "explicit" if requested_gpu_ids else "gguf"
return False, {"mode": mode, "reason": "estimate_unavailable"}
free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
- if requested_gpu_ids:
+ if single_device_gpu is not None:
+ token = str(single_device_gpu).strip()
+ if not token:
+ # Empty token = a CPU-only single-device runner (e.g. a CPU
+ # diffusion GGUF): it uses no GPU VRAM, so it never threatens
+ # active training and can always load.
+ return True, {"mode": "single_device", "reason": "cpu_only"}
+ try:
+ selected_gpu = int(token)
+ if selected_gpu < 0:
+ raise ValueError
+ except (TypeError, ValueError):
+ # A non-numeric device token (e.g. a CUDA UUID / MIG handle)
+ # can't be mapped to a free-VRAM index, but the runner still
+ # drives ONE device. Size against the worst-case visible device
+ # (min free), never the aggregate pool, so a single-device load
+ # is never OK'd on capacity it can't use and OOMs training.
+ free_vals = [min(free_by_index.values())] if free_by_index else []
+ else:
+ free_vals = [free_by_index.get(selected_gpu, 0.0)]
+ elif requested_gpu_ids:
# Invalid ids -> load_model 400s first, so don't block; missing id = 0.
try:
resolved = resolve_requested_gpu_ids(requested_gpu_ids)
except ValueError:
- return True, {"mode": "explicit", "reason": "invalid_gpu_ids"}
+ return True, {"mode": mode, "reason": "invalid_gpu_ids"}
free_vals = [free_by_index.get(i, 0.0) for i in resolved]
- mode = "explicit"
else:
# GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate.
free_vals = list(free_by_index.values())
- mode = "gguf"
if not free_vals:
return False, {"mode": mode, "reason": "no_visible_gpus"}
diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py
index 63dba8579c..7daa4224aa 100644
--- a/studio/backend/tests/test_chat_load_during_training.py
+++ b/studio/backend/tests/test_chat_load_during_training.py
@@ -168,11 +168,14 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
devices,
required_override = None,
estimate = None,
+ single_device_gpu = None,
+ gpu_ids = None,
):
with (
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})),
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
+ patch("utils.hardware.resolve_requested_gpu_ids", return_value = gpu_ids),
patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
):
ok, info = tv.can_load_chat_during_training(
@@ -180,9 +183,10 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
- requested_gpu_ids = None,
+ requested_gpu_ids = gpu_ids,
is_gguf = True,
required_override_gb = required_override,
+ single_device_gpu = single_device_gpu,
)
return ok, info, auto_mock
@@ -198,6 +202,88 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0)
self.assertTrue(ok)
+ def test_no_per_gpu_floor_for_gguf_with_explicit_gpu_ids(self):
+ # gpu_ids narrows llama.cpp's candidate pool but does not turn its
+ # self-placement into HF device_map="balanced". The uneven selected
+ # pair therefore keeps the aggregate GGUF check without an even-share
+ # floor on the nearly-full card.
+ ok, info, _ = self._run(
+ devices = _devices((0, 80, 35), (1, 80, 70), (2, 80, 0)),
+ required_override = 20.0,
+ gpu_ids = [0, 1],
+ )
+ self.assertTrue(ok)
+ self.assertEqual(info["mode"], "gguf")
+
+ def test_single_device_uses_selected_gpu(self):
+ # The model needs 27 GB with headroom. GPU 0 has 45 GB free, while an
+ # unrelated training-heavy GPU 1 has only 10 GB free.
+ ok, info, _ = self._run(
+ devices = _devices((0, 80, 35), (1, 80, 70)),
+ required_override = 20.0,
+ single_device_gpu = "0",
+ )
+ self.assertTrue(ok)
+ self.assertEqual(info["usable_gb"], 45.0)
+
+ blocked, blocked_info, _ = self._run(
+ devices = _devices((0, 80, 35), (1, 80, 70)),
+ required_override = 20.0,
+ single_device_gpu = "1",
+ )
+ self.assertFalse(blocked)
+ self.assertEqual(blocked_info["usable_gb"], 10.0)
+
+ def test_single_device_unresolved_token_sizes_against_worst_device(self):
+ # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a
+ # free-VRAM index. The runner still drives ONE device, so size against the
+ # worst-case visible device (min free), not the aggregate pool: one GPU
+ # with 80 GB free vs a 20 GB model -> allow.
+ ok, info, _ = self._run(
+ devices = _devices((0, 80, 0)),
+ required_override = 20.0,
+ single_device_gpu = "GPU-uuid",
+ )
+ self.assertTrue(ok)
+ self.assertEqual(info["mode"], "single_device")
+ self.assertNotIn("reason", info)
+
+ def test_single_device_unresolved_token_refuses_when_worst_device_full(self):
+ # Same UUID fallback, worst-case device nearly full (2 GB for a 20 GB
+ # model) -> refuse (default-deny), not on an unresolved-token technicality.
+ ok, info, _ = self._run(
+ devices = _devices((0, 80, 78)),
+ required_override = 20.0,
+ single_device_gpu = "GPU-uuid",
+ )
+ self.assertFalse(ok)
+ self.assertNotEqual(info.get("reason"), "unresolved_gpu_id")
+
+ def test_single_device_unresolved_token_uses_min_free_not_aggregate(self):
+ # The single-device runner uses ONE device but we can't tell which from a
+ # UUID token. Sizing against the aggregate pool would let a 20 GB model
+ # "fit" 160 GB of pooled free VRAM while landing on a 2 GB card and OOMing
+ # training. Min-free (2 GB) is the safe worst case -> refuse.
+ ok, info, _ = self._run(
+ devices = _devices((0, 80, 78), (1, 80, 0), (2, 80, 0)),
+ required_override = 20.0,
+ single_device_gpu = "GPU-uuid",
+ )
+ self.assertFalse(ok)
+ self.assertEqual(info["mode"], "single_device")
+
+ def test_single_device_cpu_token_allows(self):
+ # An empty device token = a CPU-only single-device runner (CPU diffusion
+ # GGUF): it uses no GPU VRAM, so it never threatens training -> allow
+ # regardless of how full the GPUs are.
+ ok, info, _ = self._run(
+ devices = _devices((0, 80, 78)),
+ required_override = 20.0,
+ single_device_gpu = "",
+ )
+ self.assertTrue(ok)
+ self.assertEqual(info["reason"], "cpu_only")
+
def test_estimate_unavailable_refuses(self):
# No override and the estimator can't size it -> default-deny.
ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None)
@@ -309,6 +395,8 @@ class TestChatLoadGuardRoute(unittest.TestCase):
captured = None,
training_active,
decision,
+ gpu_memory_mode = "auto",
+ requested_gpu_ids = None,
):
config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None)
with _stub_guard_deps(
@@ -320,7 +408,8 @@ class TestChatLoadGuardRoute(unittest.TestCase):
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
- requested_gpu_ids = None,
+ requested_gpu_ids = requested_gpu_ids,
+ gpu_memory_mode = gpu_memory_mode,
)
def test_noop_when_training_inactive(self):
@@ -332,6 +421,141 @@ class TestChatLoadGuardRoute(unittest.TestCase):
def test_allows_when_fits(self):
self._guard(training_active = True, decision = (True, {"mode": "auto"}))
+ def test_diffusion_detection_uses_name_before_download(self):
+ config = SimpleNamespace(
+ identifier = "unsloth/DiffusionGemma-GGUF",
+ gguf_hf_repo = "unsloth/DiffusionGemma-GGUF",
+ gguf_file = None,
+ )
+ self.assertTrue(self.route._classify_diffusion_gguf(config))
+
+ def test_uncached_gguf_classification_remains_unknown(self):
+ config = SimpleNamespace(
+ identifier = "owner/renamed-model",
+ gguf_hf_repo = "owner/renamed-model",
+ gguf_variant = "Q4_K_M",
+ gguf_file = None,
+ )
+ self.assertIsNone(self.route._classify_diffusion_gguf(config))
+
+ def test_diffusion_detection_reuses_loader_metadata_probe(self):
+ import tempfile
+
+ seen = []
+
+ class _Probe:
+ is_diffusion = False
+ _architecture = None
+
+ def _read_gguf_metadata(self, path):
+ seen.append(path)
+ self.is_diffusion = True
+
+ with tempfile.TemporaryDirectory() as d:
+ model = Path(d) / "renamed.gguf"
+ model.write_bytes(b"GGUF")
+ config = SimpleNamespace(identifier = "local", gguf_file = str(model))
+ with patch.object(self.route, "LlamaCppBackend", _Probe):
+ self.assertTrue(self.route._classify_diffusion_gguf(config))
+ self.assertEqual(seen, [str(model)])
+
+ def test_local_chat_gguf_classification_is_definitive(self):
+ import tempfile
+ class _Probe:
+ is_diffusion = False
+ _architecture = "llama"
+
+ def _read_gguf_metadata(self, _path):
+ pass
+
+ with tempfile.TemporaryDirectory() as d:
+ model = Path(d) / "renamed.gguf"
+ model.write_bytes(b"GGUF")
+ config = SimpleNamespace(identifier = "local", gguf_file = str(model))
+ with patch.object(self.route, "LlamaCppBackend", _Probe):
+ self.assertFalse(self.route._classify_diffusion_gguf(config))
+
+ def test_manual_known_normal_gguf_bypasses_training_estimate(self):
+ captured = []
+ config = SimpleNamespace(is_gguf = True)
+ with patch.object(self.route, "_classify_diffusion_gguf", return_value = False):
+ self._guard(
+ config = config,
+ captured = captured,
+ training_active = True,
+ decision = (False, {"reason": "must not run"}),
+ gpu_memory_mode = "manual",
+ )
+ self.assertEqual(captured, [])
+
+ def test_manual_unknown_gguf_keeps_single_device_training_guard(self):
+ captured = []
+ config = SimpleNamespace(is_gguf = True)
+ with (
+ patch.object(self.route, "_classify_diffusion_gguf", return_value = None),
+ patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5),
+ patch.object(
+ self.route.LlamaCppBackend,
+ "_diffusion_gpu_arg",
+ return_value = "2",
+ ),
+ ):
+ self._guard(
+ config = config,
+ captured = captured,
+ training_active = True,
+ decision = (True, {"mode": "single_device"}),
+ gpu_memory_mode = "manual",
+ )
+ self.assertEqual(len(captured), 1)
+ self.assertEqual(captured[0]["single_device_gpu"], "2")
+
+ def test_manual_diffusion_uses_single_device_guard(self):
+ captured = []
+ config = SimpleNamespace(is_gguf = True)
+ with (
+ patch.object(self.route, "_classify_diffusion_gguf", return_value = True),
+ patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5),
+ ):
+ self._guard(
+ config = config,
+ captured = captured,
+ training_active = True,
+ decision = (True, {"mode": "gguf"}),
+ gpu_memory_mode = "manual",
+ requested_gpu_ids = [3, 1],
+ )
+ self.assertEqual(len(captured), 1)
+ self.assertEqual(captured[0]["single_device_gpu"], "1")
+ self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1])
+
+ def test_unpinned_diffusion_uses_runner_default_gpu(self):
+ captured = []
+ config = SimpleNamespace(is_gguf = True)
+ with (
+ patch.object(self.route, "_classify_diffusion_gguf", return_value = True),
+ patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5),
+ patch.object(
+ self.route.LlamaCppBackend,
+ "_effective_gpu_count",
+ return_value = 2,
+ ),
+ patch.object(
+ self.route.LlamaCppBackend,
+ "_diffusion_gpu_arg",
+ return_value = "3",
+ ) as gpu_arg,
+ ):
+ self._guard(
+ config = config,
+ captured = captured,
+ training_active = True,
+ decision = (True, {"mode": "single_device"}),
+ gpu_memory_mode = "manual",
+ )
+ gpu_arg.assert_called_once_with(None, cpu_only = False)
+ self.assertEqual(captured[0]["single_device_gpu"], "3")
+
def test_refuses_with_headroom_number(self):
info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"}
with self.assertRaises(HTTPException) as exc:
@@ -467,36 +691,115 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
self.assertEqual(captured[0]["load_in_4bit"], False)
self.assertEqual(captured[0]["max_seq_length"], 4096)
- def test_rejects_gguf_with_gpu_ids_before_guard(self):
- # /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard.
+ def test_validate_forwards_manual_gpu_memory_mode_to_guard(self):
from models.inference import ValidateModelRequest
- request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0])
+ request = ValidateModelRequest(
+ model_path = "unsloth/model-GGUF",
+ gguf_variant = "Q4_K_M",
+ gpu_memory_mode = "manual",
+ )
cfg = SimpleNamespace(
- identifier = "x.gguf",
- display_name = "x",
+ identifier = "unsloth/model-GGUF",
+ display_name = "model-GGUF",
is_gguf = True,
is_lora = False,
is_vision = False,
path = None,
base_model = None,
)
- captured = []
+ captured = {}
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
- return_value = ("x.gguf", "x.gguf", False),
+ return_value = ("unsloth/model-GGUF", "unsloth/model-GGUF", False),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
- _stub_guard_deps(training_active = True, decision = (True, {}), captured = captured),
+ patch.object(
+ self.route,
+ "_guard_chat_load_against_training",
+ lambda config, **kw: captured.update(kw),
+ ),
):
- with self.assertRaises(HTTPException) as exc:
- asyncio.run(self.route.validate_model(request, current_subject = "u"))
- self.assertEqual(exc.exception.status_code, 400)
- self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail)
- self.assertEqual(captured, []) # guard never reached
+ asyncio.run(self.route.validate_model(request, current_subject = "u"))
+ self.assertEqual(captured.get("gpu_memory_mode"), "manual")
+
+ def test_validate_forwards_inherited_extras_and_parallel_to_guard(self):
+ # Regression: /load resolves inherited same-model extras and passes the
+ # real slot count to the guard; validate must do the same, else it sizes
+ # a smaller estimate (no inherited -c/--model-draft, n_parallel=1) and
+ # /load then 409s after the frontend has already unloaded.
+ from models.inference import ValidateModelRequest
+
+ request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096)
+ cfg = SimpleNamespace(
+ identifier = "unsloth/Qwen3-1.7B",
+ display_name = "Qwen3-1.7B",
+ is_gguf = False,
+ is_lora = False,
+ is_vision = False,
+ path = None,
+ base_model = None,
+ )
+ captured = {}
+ with (
+ patch.object(
+ self.route,
+ "_resolve_model_identifier_for_request",
+ return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
+ ),
+ patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
+ patch.object(self.route, "load_inference_config", return_value = {}),
+ patch.object(self.route, "_resolve_inherited_extra_args", return_value = ["-c", "32768"]),
+ patch.object(
+ self.route,
+ "_guard_chat_load_against_training",
+ lambda config, **kw: captured.update(kw),
+ ),
+ ):
+ asyncio.run(self.route.validate_model(request, current_subject = "u"))
+ self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"])
+ self.assertIn("n_parallel", captured)
+
+ def test_metadata_probe_skips_training_guard(self):
+ # A header-only probe (include_context_length) allocates no VRAM, so the
+ # training guard must not run -- else the staging GPU-layers / MoE sliders
+ # it feeds are hidden exactly when a during-training user needs them.
+ from models.inference import ValidateModelRequest
+
+ request = ValidateModelRequest(
+ model_path = "unsloth/Qwen3-1.7B",
+ max_seq_length = 4096,
+ include_context_length = True,
+ )
+ cfg = SimpleNamespace(
+ identifier = "unsloth/Qwen3-1.7B",
+ display_name = "Qwen3-1.7B",
+ is_gguf = False,
+ is_lora = False,
+ is_vision = False,
+ path = None,
+ base_model = None,
+ )
+ guard_called = []
+ with (
+ patch.object(
+ self.route,
+ "_resolve_model_identifier_for_request",
+ return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
+ ),
+ patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
+ patch.object(self.route, "load_inference_config", return_value = {}),
+ patch.object(
+ self.route,
+ "_guard_chat_load_against_training",
+ lambda *a, **kw: guard_called.append(True),
+ ),
+ ):
+ asyncio.run(self.route.validate_model(request, current_subject = "u"))
+ self.assertEqual(guard_called, [])
# ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ──────
diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py
index a5be07f8e3..ec0330ce05 100644
--- a/studio/backend/tests/test_gguf_metadata.py
+++ b/studio/backend/tests/test_gguf_metadata.py
@@ -15,6 +15,7 @@ from utils.models.gguf_metadata import (
pairing_score,
read_gguf_context_length,
read_gguf_general_metadata,
+ read_gguf_staged_dims,
read_mmproj_audio_capability,
)
@@ -153,6 +154,78 @@ def test_context_length_ignores_foreign_arch_key(tmp_path: Path):
assert read_gguf_context_length(str(p)) is None
+# --- read_gguf_staged_dims (one pass: context + layer + moe counts) ----
+
+
+def test_staged_dims_none_for_missing_or_non_gguf(tmp_path: Path):
+ assert read_gguf_staged_dims(str(tmp_path / "nope.gguf")) is None
+ p = tmp_path / "garbage.gguf"
+ p.write_bytes(b"not a gguf at all")
+ assert read_gguf_staged_dims(str(p)) is None
+
+
+def test_staged_dims_moe_with_leading_dense(tmp_path: Path):
+ # GLM-4.7-Flash shape: context + total layers + MoE layers in one read.
+ p = _write_synthetic_gguf(
+ tmp_path / "glm.gguf",
+ {"general.architecture": "deepseek2"},
+ extra_uint32 = {
+ "deepseek2.context_length": 202752,
+ "deepseek2.block_count": 47,
+ "deepseek2.expert_count": 64,
+ "deepseek2.leading_dense_block_count": 1,
+ },
+ )
+ assert read_gguf_staged_dims(str(p)) == {
+ "context_length": 202752,
+ "layer_count": 47,
+ "moe_layer_count": 46,
+ }
+
+
+def test_staged_dims_dense_model(tmp_path: Path):
+ # Dense: layer_count present, moe_layer_count 0 (slider hidden).
+ p = _write_synthetic_gguf(
+ tmp_path / "dense.gguf",
+ {"general.architecture": "qwen3"},
+ extra_uint32 = {"qwen3.context_length": 40960, "qwen3.block_count": 36},
+ )
+ assert read_gguf_staged_dims(str(p)) == {
+ "context_length": 40960,
+ "layer_count": 36,
+ "moe_layer_count": 0,
+ }
+
+
+def test_staged_dims_all_moe_no_leading_dense(tmp_path: Path):
+ # Experts present, no leading_dense key -> every block is a MoE layer.
+ p = _write_synthetic_gguf(
+ tmp_path / "moe.gguf",
+ {"general.architecture": "qwen35moe"},
+ extra_uint32 = {"qwen35moe.block_count": 40, "qwen35moe.expert_count": 256},
+ )
+ assert read_gguf_staged_dims(str(p)) == {
+ "context_length": None,
+ "layer_count": 40,
+ "moe_layer_count": 40,
+ }
+
+
+def test_staged_dims_uint64_block_count(tmp_path: Path):
+ # block_count stored as uint64 (vtype 10) still parses; moe == block_count.
+ p = _write_synthetic_gguf(
+ tmp_path / "moe64.gguf",
+ {"general.architecture": "gpt-oss"},
+ extra_uint32 = {"gpt-oss.expert_count": 32},
+ extra_uint64 = {"gpt-oss.block_count": 24},
+ )
+ assert read_gguf_staged_dims(str(p)) == {
+ "context_length": None,
+ "layer_count": 24,
+ "moe_layer_count": 24,
+ }
+
+
def test_context_length_read_from_uint64(tmp_path: Path):
# Some models store context_length as a uint64 (vtype 10).
p = _write_synthetic_gguf(
diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py
new file mode 100644
index 0000000000..b17274197f
--- /dev/null
+++ b/studio/backend/tests/test_gpu_memory_mode.py
@@ -0,0 +1,879 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Backend contract for the GPU Memory mode dropdown.
+
+The dropdown threads a single ``gpu_memory_mode`` ("auto" | "manual") from the
+chat UI through the load request. "manual" lets the user own the offload: with
+``gpu_layers < 0`` (Auto, the default) it hands all memory management to
+llama.cpp's ``--fit on`` (no CUDA/HIP device masking, no context auto-reduce, no
+gpu-layer or tensor-split planning); with ``gpu_layers >= 0`` it pins the layers
+and MoE offload itself (``--fit off``). These tests pin:
+
+ * the pydantic request/response/status contract (snake_case key, default
+ "auto", unknown values rejected),
+ * the backend ``gpu_memory_mode`` property and its reset on unload,
+ * the ``_already_in_target_state`` reload-detection branch, and
+ * that the manual + Auto-layers branch in ``load_model`` empties the probed
+ GPU set and drops tensor parallelism so the selection below no-ops, while
+ the explicit-offload branch emits ``--gpu-layers`` / ``--fit off``.
+"""
+
+from __future__ import annotations
+
+import inspect
+import sys
+import types as _types
+from pathlib import Path
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+ sys.path.insert(0, _BACKEND_DIR)
+
+# Same external-dep stubs as the other llama_cpp unit tests so importing
+# the backend doesn't drag in structlog / httpx / loggers.
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+
+_structlog_stub = _types.ModuleType("structlog")
+_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
+sys.modules.setdefault("structlog", _structlog_stub)
+
+# httpx is a real, installed backend dependency: import it so the genuine module
+# is in sys.modules. A hand-rolled stub here is inevitably incomplete and, since
+# setdefault installs it before real httpx loads, would poison a combined pytest
+# run -- routes/inference references httpx.Response (and other attrs) at def time.
+import httpx # noqa: F401
+
+from core.inference import llama_cpp as llama_cpp_module
+from core.inference.llama_cpp import LlamaCppBackend
+from models.inference import (
+ InferenceStatusResponse,
+ LoadRequest,
+ LoadResponse,
+)
+
+
+# ── Pydantic contract (snake_case key, default "auto") ───────────────
+
+
+def test_load_request_defaults_gpu_memory_mode_auto():
+ assert LoadRequest(model_path = "owner/repo").gpu_memory_mode == "auto"
+
+
+def test_load_request_round_trips_json_key():
+ req = LoadRequest.model_validate({"model_path": "owner/repo", "gpu_memory_mode": "manual"})
+ assert req.gpu_memory_mode == "manual"
+ assert req.model_dump()["gpu_memory_mode"] == "manual"
+
+
+def test_load_request_rejects_unknown_mode():
+ with pytest.raises(ValueError):
+ LoadRequest(model_path = "owner/repo", gpu_memory_mode = "bogus")
+
+
+@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
+def test_response_models_emit_gpu_memory_mode(model_cls):
+ if model_cls is LoadResponse:
+ default = model_cls(
+ status = "loaded",
+ model = "owner/repo",
+ display_name = "repo",
+ inference = {},
+ )
+ manual = model_cls(
+ status = "loaded",
+ model = "owner/repo",
+ display_name = "repo",
+ inference = {},
+ gpu_memory_mode = "manual",
+ )
+ else:
+ default = model_cls()
+ manual = model_cls(gpu_memory_mode = "manual")
+ assert default.model_dump()["gpu_memory_mode"] == "auto"
+ assert manual.model_dump()["gpu_memory_mode"] == "manual"
+
+
+# ── Backend property + reset ─────────────────────────────────────────
+
+
+class _FakeProcess:
+ """Stand-in for subprocess.Popen so _kill_process is a no-op."""
+
+ def terminate(self):
+ pass
+
+ def wait(self, timeout = None):
+ return 0
+
+ def kill(self):
+ pass
+
+ def poll(self):
+ return 0
+
+
+def test_gpu_memory_mode_property_defaults_auto():
+ assert LlamaCppBackend().gpu_memory_mode == "auto"
+
+
+def test_gpu_memory_mode_property_reflects_field():
+ backend = LlamaCppBackend()
+ backend._gpu_memory_mode = "manual"
+ assert backend.gpu_memory_mode == "manual"
+
+
+def test_unload_resets_gpu_memory_mode():
+ backend = LlamaCppBackend()
+ backend._process = _FakeProcess()
+ backend._gpu_memory_mode = "manual"
+ backend.unload_model()
+ assert backend.gpu_memory_mode == "auto"
+
+
+# ── _already_in_target_state reload-detection branch ─────────────────
+
+
+def _loaded_backend(gpu_memory_mode: str) -> LlamaCppBackend:
+ backend = LlamaCppBackend()
+ backend._process = _FakeProcess() # is_loaded only checks "is not None"
+ backend._healthy = True
+ backend._model_identifier = "owner/repo"
+ backend._hf_variant = "Q4_K_M"
+ backend._requested_n_ctx = 8192
+ backend._cache_type_kv = None
+ backend._requested_spec_mode = "auto"
+ backend._chat_template_override = None
+ backend._is_vision = False
+ backend._extra_args = None
+ backend._gguf_path = None
+ backend._gpu_memory_mode = gpu_memory_mode
+ return backend
+
+
+def _target_state(backend: LlamaCppBackend, gpu_memory_mode: str) -> bool:
+ return backend._already_in_target_state(
+ gguf_path = None,
+ model_identifier = "owner/repo",
+ hf_variant = "Q4_K_M",
+ n_ctx = 8192,
+ cache_type_kv = None,
+ speculative_type = "auto",
+ chat_template_override = None,
+ extra_args = None,
+ is_vision = False,
+ gpu_memory_mode = gpu_memory_mode,
+ )
+
+
+@pytest.mark.parametrize("mode", ["auto", "manual"])
+def test_already_in_target_state_matches_same_mode(mode):
+ assert _target_state(_loaded_backend(mode), mode) is True
+
+
+@pytest.mark.parametrize("loaded,requested", [("auto", "manual"), ("manual", "auto")])
+def test_already_in_target_state_reloads_on_mode_change(loaded, requested):
+ # Flipping the dropdown either direction must force a reload so the command
+ # is rebuilt with/without the Unsloth GPU masking.
+ assert _target_state(_loaded_backend(loaded), requested) is False
+
+
+def test_already_in_target_state_ignores_mode_for_diffusion():
+ # The diffusion runner is mode-agnostic (always "auto"), so a standing manual
+ # preference must not force a needless reload.
+ backend = _loaded_backend("auto")
+ backend._is_diffusion = True
+ assert _target_state(backend, "manual") is True
+
+
+# ── load_model: manual + Auto layers bypasses Unsloth GPU management ──
+
+
+def _load_model_source() -> str:
+ return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)
+
+
+def test_auto_layers_branch_empties_gpus_and_drops_tensor_parallel():
+ # Emptying the probed set makes the selection / TP planning below no-op, so
+ # gpu_indices stays None and use_fit True (--fit on).
+ src = _load_model_source()
+ gate = src.find('if gpu_memory_mode == "manual" and gpu_layers < 0:')
+ assert gate != -1, "load_model must branch on manual + Auto layers (gpu_layers < 0)"
+ block = src[gate : gate + 1400]
+ assert "gpus = []" in block, "Auto-layers branch must empty the probed GPU set"
+ # --fit aborts under --split-mode tensor, so a raw-extras split-mode is stripped.
+ assert "strip_split_mode_only(extra_args)" in block
+ assert "requested_ctx if requested_ctx > 0 else 0" in block
+ # The branch sits before GPU selection assigns gpu_indices; --fit on is its emission.
+ assert gate < src.find("gpu_indices, use_fit = None, True")
+ assert 'cmd.extend(["--fit", "on"])' in src
+ # TP drops for this path, but at a guard BEFORE the quantized-KV cache-drop, so
+ # a requested quantized cache survives into the --fit load.
+ tp_drop = src.find('if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:')
+ assert tp_drop != -1, "manual + Auto layers must drop tensor_parallel"
+ assert "tensor_parallel = False" in src[tp_drop : tp_drop + 400]
+ cache_drop = src.find("Tensor parallelism requires a non-quantized KV cache")
+ assert cache_drop != -1
+ assert (
+ tp_drop < cache_drop
+ ), "TP must drop before the cache-drop so a quantized KV survives --fit"
+
+
+def test_auto_layers_never_sends_ctx_size_zero():
+ # Sending "-c 0" sets fit_params_min_ctx = UINT32_MAX in llama.cpp, pinning
+ # the full native context and disabling --fit's reduction. So the base cmd
+ # must never carry -c, "-c 0" is emitted only outside the Auto-layers (--fit)
+ # case, and a positive context is passed through (which --fit optimizes
+ # layers around).
+ src = _load_model_source()
+ base_start = src.find("cmd = [")
+ base_end = src.find("\n ]", base_start)
+ base_block = src[base_start:base_end]
+ assert '"-c"' not in base_block, "-c must be conditional, not in the base cmd list"
+ assert 'cmd.extend(["-c", str(effective_ctx)])' in src, "positive ctx must pass -c"
+ assert 'auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0' in src
+ zero = src.find('cmd.extend(["-c", "0"])')
+ assert zero != -1, '"-c 0" emission must exist outside the Auto-layers case'
+ guard = src.rfind("elif not auto_fit:", 0, zero)
+ assert guard != -1 and zero - guard < 120, '"-c 0" must sit under the not-auto_fit guard'
+
+
+def test_manual_mode_clears_inherited_main_model_placement_env():
+ env = {name: "inherited" for name in LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS}
+ env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] = "7"
+ env["UNRELATED"] = "kept"
+
+ LlamaCppBackend._clear_manual_placement_env(env)
+
+ assert not (set(env) & set(LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS))
+ assert env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] == "7"
+ assert env["UNRELATED"] == "kept"
+
+
+def test_load_model_sanitizes_manual_env_after_building_child_env():
+ src = _load_model_source()
+ env_build = src.find("env = self._llama_server_env_for_binary(binary)")
+ env_clear = src.find("self._clear_manual_placement_env(env)", env_build)
+ launch = src.find("subprocess.Popen", env_build)
+ assert env_build != -1
+ assert env_build < env_clear < launch
+
+
+# ── Manual offload (--gpu-layers + --fit off + --n-cpu-moe) ───────────
+
+
+def test_load_request_accepts_manual():
+ req = LoadRequest(
+ model_path = "owner/repo",
+ gpu_memory_mode = "manual",
+ gpu_layers = 20,
+ n_cpu_moe = 8,
+ tensor_split = [2, 1],
+ )
+ assert req.gpu_memory_mode == "manual"
+ assert req.gpu_layers == 20
+ assert req.n_cpu_moe == 8
+ assert req.tensor_split == [2, 1]
+
+
+def test_load_request_manual_defaults():
+ req = LoadRequest(model_path = "owner/repo")
+ assert req.gpu_layers == -1
+ assert req.n_cpu_moe == 0
+ assert req.tensor_split is None
+
+
+@pytest.mark.parametrize("bad", [[0, 0], [-1, 2], [float("inf"), 1], [float("nan"), 1]])
+def test_load_request_rejects_degenerate_tensor_split(bad):
+ # A negative/non-finite/all-zero split is dropped at launch but compared raw
+ # in the reload dedupe, so it would reload forever -- reject it up front.
+ with pytest.raises(ValueError):
+ LoadRequest(model_path = "owner/repo", tensor_split = bad)
+
+
+@pytest.mark.parametrize("good", [[2, 1], [1, 1], [], None])
+def test_load_request_accepts_valid_tensor_split(good):
+ assert LoadRequest(model_path = "owner/repo", tensor_split = good).tensor_split == good
+
+
+def test_route_normalizes_explicit_extras_before_reload_dedupe():
+ route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
+ load_impl = route_src[route_src.index("async def _load_model_impl") :]
+ strip = load_impl.index("_stripped_explicit = strip_shadowing_flags")
+ normalize = load_impl.index(
+ 'request = request.model_copy(update = {"llama_extra_args": extra_llama_args})'
+ )
+ dedupe = load_impl.index("and _request_matches_loaded_settings(")
+ assert strip < normalize < dedupe
+
+
+@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
+def test_response_models_emit_manual_fields(model_cls):
+ if model_cls is LoadResponse:
+ obj = model_cls(
+ status = "loaded",
+ model = "owner/repo",
+ display_name = "repo",
+ inference = {},
+ gpu_memory_mode = "manual",
+ gpu_layers = 20,
+ n_cpu_moe = 8,
+ tensor_split = [2, 1],
+ n_layers = 32,
+ n_moe_layers = 32,
+ )
+ else:
+ obj = model_cls(
+ gpu_memory_mode = "manual",
+ gpu_layers = 20,
+ n_cpu_moe = 8,
+ tensor_split = [2, 1],
+ n_layers = 32,
+ n_moe_layers = 32,
+ )
+ dumped = obj.model_dump()
+ assert dumped["gpu_memory_mode"] == "manual"
+ assert dumped["gpu_layers"] == 20
+ assert dumped["n_cpu_moe"] == 8
+ assert dumped["tensor_split"] == [2, 1]
+ assert dumped["n_layers"] == 32
+ assert dumped["n_moe_layers"] == 32
+
+
+def test_manual_properties_default_and_reflect_and_reset():
+ backend = LlamaCppBackend()
+ assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0
+ assert backend.tensor_split is None
+ backend._gpu_layers = 20
+ backend._n_cpu_moe = 8
+ backend._tensor_split = [2, 1]
+ assert backend.gpu_layers == 20 and backend.n_cpu_moe == 8
+ assert backend.tensor_split == [2, 1]
+ backend._process = _FakeProcess()
+ backend.unload_model()
+ assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0
+ assert backend.tensor_split is None
+
+
+def test_n_moe_layers_property():
+ # 0 for a dense model (hides the slider); block_count for all-MoE;
+ # block_count - leading_dense otherwise (GLM-4.7-Flash: 47 - 1 -> 46).
+ b = LlamaCppBackend()
+ b._n_layers = 36
+ b._n_experts = None
+ assert b.n_moe_layers == 0
+ b._n_experts = 128
+ b._leading_dense_block_count = None
+ assert b.n_moe_layers == 36
+ b._n_layers = 47
+ b._leading_dense_block_count = 1
+ assert b.n_moe_layers == 46
+
+
+def _target_state_manual(
+ backend,
+ *,
+ gpu_layers,
+ n_cpu_moe,
+ tensor_split = None,
+):
+ return backend._already_in_target_state(
+ gguf_path = None,
+ model_identifier = "owner/repo",
+ hf_variant = "Q4_K_M",
+ n_ctx = 8192,
+ cache_type_kv = None,
+ speculative_type = "auto",
+ chat_template_override = None,
+ extra_args = None,
+ is_vision = False,
+ gpu_memory_mode = "manual",
+ gpu_layers = gpu_layers,
+ n_cpu_moe = n_cpu_moe,
+ tensor_split = tensor_split,
+ )
+
+
+def test_manual_reloads_on_gpu_layers_or_n_cpu_moe_or_split_change():
+ backend = _loaded_backend("manual")
+ backend._gpu_layers = 20
+ backend._n_cpu_moe = 0
+ backend._tensor_split = None
+ # Same knobs -> no reload.
+ assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is True
+ # Changed layer count -> reload.
+ assert _target_state_manual(backend, gpu_layers = 16, n_cpu_moe = 0) is False
+ # Changed MoE offload -> reload.
+ assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 8) is False
+ # Added a GPU split -> reload.
+ assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is False
+ # Same GPU split -> no reload.
+ backend._tensor_split = [2, 1]
+ assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is True
+
+
+def test_auto_layers_reload_tracks_only_gpu_layers():
+ # Under Auto (gpu_layers < 0) the MoE/split knobs don't apply, so a leftover
+ # request value must not reload -- only a gpu_layers change (Auto -> pinned) does.
+ backend = _loaded_backend("manual")
+ backend._gpu_layers = -1
+ backend._n_cpu_moe = 0
+ backend._tensor_split = None
+ # Same Auto, leftover MoE/split in the request -> still no reload.
+ assert _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) is True
+ # Auto -> explicit offload reloads.
+ assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is False
+
+
+def test_manual_offload_emits_gpu_layers_fit_off_and_n_cpu_moe():
+ src = _load_model_source()
+ gate = src.find('elif gpu_memory_mode == "manual":')
+ assert gate != -1, "load_model must have an explicit-offload manual branch"
+ block = src[gate : gate + 700]
+ # Empties the probed set (skips the planner) but keeps the user's TP choice
+ # (only the Auto-layers branch above drops TP).
+ assert "gpus = []" in block
+ assert "tensor_parallel = False" not in block
+ # The cmd emits the layer count with fit disabled, gated on gpu_layers >= 0.
+ assert 'if gpu_memory_mode == "manual" and gpu_layers >= 0:' in src
+ assert 'cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])' in src
+ # MoE offload uses --n-cpu-moe via _resolve_cpu_moe_flag (tested behaviorally below).
+ assert "_resolve_cpu_moe_flag(" in src
+ assert 'cmd.extend(["--n-cpu-moe", str(moe_flag)])' in src
+ # A count requested on a dense model is never emitted, so it must also be
+ # dropped from the recorded state -- else /status and /load report a count
+ # llama-server never received (same rule as the tensor-split drop below).
+ moe_emit = src.find('cmd.extend(["--n-cpu-moe", str(moe_flag)])')
+ assert "elif n_cpu_moe:" in src[moe_emit : moe_emit + 300]
+ assert "self._n_cpu_moe = 0" in src[moe_emit : moe_emit + 300]
+ # The offload path forces use_fit False so --fit-ctx is never added under --fit off.
+ emit = src.find('cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])')
+ assert "use_fit = False" in src[src.rfind("\n", 0, emit) - 200 : emit + 80]
+
+
+def test_status_reports_requested_context_length():
+ # The hydration path re-seeds a Manual+Auto context pin from the REQUESTED
+ # n_ctx (0 = Auto); context_length only exposes the resolved value.
+ assert "requested_context_length" in InferenceStatusResponse.model_fields
+ s = InferenceStatusResponse(requested_context_length = 8192)
+ assert s.model_dump()["requested_context_length"] == 8192
+ assert InferenceStatusResponse().model_dump()["requested_context_length"] is None
+ # The /status route must actually wire it from the backend (a declared-but-
+ # never-populated field would leave hydration silently reverting the pin).
+ from pathlib import Path as _P
+
+ route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
+ assert "requested_context_length = llama_backend.requested_n_ctx" in route_src
+
+
+def test_manual_offload_emits_tensor_split():
+ # The offload path emits --tensor-split from the per-GPU shares, only when
+ # provided, with >1 GPU in use, AND matching that count (a stale ratio on a
+ # narrowed picker or a mismatched direct-API list must not emit -- llama-
+ # server aborts on a split/GPU-count mismatch).
+ src = _load_model_source()
+ assert "if tensor_split and _split_gpus > 1:" in src
+ # Emit only on a length match AND a positive sanitized total: a mismatched
+ # or all-zero split aborts llama-server / assigns nothing, so it's dropped.
+ # The emitted list is the sanitized one (clamping tested behaviorally below).
+ assert "_sanitized_split = self._sanitize_tensor_split(tensor_split)" in src
+ assert "if len(_sanitized_split) == _split_gpus and _split_total > 0:" in src
+ assert '"--tensor-split"' in src
+ # Joined as a comma list (e.g. "2,1") within the explicit-offload cmd branch.
+ gate = src.find('if gpu_memory_mode == "manual" and gpu_layers >= 0:')
+ nxt = src.find("elif use_fit:", gate)
+ assert '","' in src[gate:nxt] and "tensor_split" in src[gate:nxt]
+ # A split with a single effective GPU is never emitted, so it must also be
+ # dropped from the recorded state -- else /status and /load report a ratio
+ # llama-server never received and the dedupe baseline preserves it.
+ assert "elif tensor_split:" in src[gate:nxt]
+ drop = src.find("elif tensor_split:", gate, nxt)
+ assert "self._tensor_split = None" in src[drop : drop + 250]
+
+
+def test_sanitize_tensor_split_clamps_negative_and_non_finite():
+ # Negative entries would launch a placement different from the ratio the
+ # UI showed; inf passes a plain > 0 total gate and would emit
+ # "--tensor-split inf,..." (llama.cpp normalizes shares by the running
+ # total, so an inf poisons the shares from that entry on). Both clamp to 0.
+ sanitize = LlamaCppBackend._sanitize_tensor_split
+ assert sanitize([2, 1]) == [2.0, 1.0]
+ assert sanitize([-1, 2]) == [0.0, 2.0]
+ assert sanitize([float("inf"), 1]) == [0.0, 1.0]
+ assert sanitize([float("nan"), 1]) == [0.0, 1.0]
+ # All-zero survives sanitization; the call site's total gate drops it.
+ assert sanitize([0, 0]) == [0.0, 0.0]
+ # Unreadable input -> []; the call site's length gate drops it.
+ assert sanitize(["x", 1]) == []
+ assert sanitize([10**400, 1]) == []
+
+
+def test_zero_offload_mask_honors_device_pin_spellings():
+ # A user device pin must keep the GPUs visible: llama-server aborts on a
+ # pin it can't see ('error: invalid device'). The pin can arrive as
+ # --device or its -dev alias, as the draft forms (parsed even with no
+ # drafter loaded), or as an inherited LLAMA_ARG_DEVICE env var.
+ load_src = _load_model_source()
+ assert "self._zero_offload_keeps_gpu_visible(cmd, env)" in load_src
+ block = inspect.getsource(LlamaCppBackend._cmd_has_gpu_device_pin)
+ for flag in (
+ '"--device"',
+ '"-dev"',
+ '"--spec-draft-device"',
+ '"-devd"',
+ '"--device-draft"',
+ ):
+ assert flag in block
+ assert '"LLAMA_ARG_DEVICE"' in block
+
+
+def test_resolve_cpu_moe_flag():
+ # Clamp the requested MoE-layer count to the model's MoE layers, then offset
+ # past leading dense layers (--n-cpu-moe counts from layer 0).
+ R = LlamaCppBackend._resolve_cpu_moe_flag
+ assert R(0, 40, 0) is None # nothing requested
+ assert R(8, 0, 0) is None # dense model (no MoE layers)
+ assert R(8, 40, 0) == 8 # all-MoE: direct
+ assert R(100, 40, 0) == 40 # clamp to the MoE layer count
+ # GLM-4.7-Flash (deepseek2): block_count 47, leading_dense 1, n_moe 46.
+ assert R(5, 46, 1) == 6 # offset past the 1 dense layer
+ assert R(46, 46, 1) == 47 # all MoE on CPU == block_count
+
+
+def test_manual_allows_tensor_parallel_via_split_mode():
+ # Manual offload keeps the user's TP choice but skips the memory-based planner
+ # (plan_tp excludes manual, so its empty gpu set can't downgrade TP). The
+ # --split-mode tensor emission gates on tensor_parallel alone, so manual
+ # reaches it -- with tp_tensor_split None it's an even split (no
+ # --tensor-split). --fit off means no fit/tensor abort.
+ src = _load_model_source()
+ assert 'plan_tp = tensor_parallel and gpu_memory_mode != "manual"' in src
+ assert "if plan_tp:" in src
+ assert "if plan_tp and len(tp_gpus) < 2:" in src
+ sm = src.find('cmd.extend(["--split-mode", "tensor"])')
+ assert sm != -1, "TP must emit --split-mode tensor"
+ guard = src.rfind("if tensor_parallel:", 0, sm)
+ assert guard != -1 and sm - guard < 200, "split-mode gates on tensor_parallel"
+ # The tensor-split is only emitted for a planned (non-even) split, which
+ # manual never produces, so manual stays an even split.
+ assert "if tp_tensor_split and len(tp_tensor_split) > 1:" in src
+
+
+def test_fit_sets_target_margin():
+ # Manual + Auto (auto_fit) tightens the per-device VRAM margin to 512 MiB.
+ caps = {"supports_fit_target": True}
+ flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 0, caps)
+ assert flags[flags.index("--fit-target") + 1] == "512"
+ # Not emitted on the legacy auto path (fit on but not auto_fit): -c 0 pins
+ # native there, so the tighter margin must not ride along.
+ assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, True, False, 0, 0, caps)
+ # Not emitted when fit is off.
+ assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, False, False, 0, 0, caps)
+ # Not emitted when the binary lacks support.
+ assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(
+ 1, True, True, 0, 0, {"supports_fit_target": False}
+ )
+
+
+# ── GPU picker (gpu_ids -> CUDA_VISIBLE_DEVICES) ─────────────────────
+
+
+def test_load_request_accepts_gpu_ids():
+ req = LoadRequest(model_path = "owner/repo", gpu_ids = [1, 0])
+ assert req.gpu_ids == [1, 0]
+ assert LoadRequest(model_path = "owner/repo").gpu_ids is None
+
+
+@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
+def test_response_models_emit_gpu_ids(model_cls):
+ if model_cls is LoadResponse:
+ obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1])
+ else:
+ obj = model_cls(gpu_ids = [1])
+ assert obj.model_dump()["gpu_ids"] == [1]
+
+
+def test_gpu_ids_property_default_and_reset():
+ backend = LlamaCppBackend()
+ assert backend.gpu_ids is None
+ backend._gpu_ids = [0, 1]
+ assert backend.gpu_ids == [0, 1]
+ backend._process = _FakeProcess()
+ backend.unload_model()
+ assert backend.gpu_ids is None
+
+
+def _target_state_gpu_ids(backend, gpu_ids):
+ return backend._already_in_target_state(
+ gguf_path = None,
+ model_identifier = "owner/repo",
+ hf_variant = "Q4_K_M",
+ n_ctx = 8192,
+ cache_type_kv = None,
+ speculative_type = "auto",
+ chat_template_override = None,
+ extra_args = None,
+ is_vision = False,
+ gpu_ids = gpu_ids,
+ )
+
+
+def test_gpu_ids_reload_detection_is_order_insensitive():
+ backend = _loaded_backend("auto")
+ backend._gpu_ids = [0, 1]
+ # Same set, different order -> no reload.
+ assert _target_state_gpu_ids(backend, [1, 0]) is True
+ # Different set -> reload.
+ assert _target_state_gpu_ids(backend, [0]) is False
+ # Dropping the pick (auto) -> reload.
+ assert _target_state_gpu_ids(backend, None) is False
+
+
+def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device():
+ # The diffusion runner drives only its single lowest device, so the backend
+ # records [lowest]. A later multi-GPU request that still resolves to that
+ # same lowest device must dedupe (no needless reload); a request whose lowest
+ # device moves, or that drops the pick, must reload.
+ backend = _loaded_backend("auto")
+ backend._is_diffusion = True
+ backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick
+ assert _target_state_gpu_ids(backend, [3, 1]) is True
+ assert _target_state_gpu_ids(backend, [1]) is True
+ # Lowest device changes (2, not 1) -> reload.
+ assert _target_state_gpu_ids(backend, [3, 2]) is False
+ # Dropping the pick (auto) -> reload.
+ assert _target_state_gpu_ids(backend, None) is False
+
+
+def test_start_diffusion_server_resets_tensor_parallel():
+ # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model
+ # phase 1 only kills the process, it skips the unload reset). Diffusion is never
+ # TP, so startup must clear it -- else /status misreports TP and an identical
+ # diffusion re-Apply reloads against stale tensor-parallel state.
+ src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server)
+ assert "self._tensor_parallel = False" in src
+
+
+def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids():
+ # The route-level reload dedupe mirrors the backend: for a loaded diffusion
+ # model it compares the request against the single recorded device, not the
+ # full requested list, or a same-device multi-GPU pick reloads needlessly.
+ route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
+ match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :]
+ guard = match_impl.index("if llama_backend.is_diffusion:")
+ collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None")
+ compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:")
+ assert guard < collapse < compare
+
+
+# ── Manual tensor split: child enumeration pinned to the picker's order ──────
+
+
+def _patch_split_pin_env(monkeypatch, *, inherited, reported):
+ """Point the pin helper at a fake inherited mask and picker report.
+ ``reported`` None = enumeration unavailable (falls back to ascending)."""
+ import utils.hardware as hw
+
+ monkeypatch.setattr(
+ LlamaCppBackend, "_resolve_visible_physical_ids", staticmethod(lambda: inherited)
+ )
+ info = (
+ {"available": False}
+ if reported is None
+ else {
+ "available": True,
+ "index_kind": "physical",
+ "devices": [{"index": i} for i in reported],
+ }
+ )
+ monkeypatch.setattr(hw, "get_backend_visible_gpu_info", lambda: info)
+
+
+def test_split_pin_reorders_inherited_numeric_mask(monkeypatch):
+ # Parent CUDA_VISIBLE_DEVICES=3,1 makes the child enumerate dev0=phys3, but
+ # nvidia-smi reported the picker's list ascending -- the mask must be
+ # re-emitted in that order or the per-GPU shares land on the wrong cards.
+ _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
+ env = {"CUDA_VISIBLE_DEVICES": "3,1"}
+ LlamaCppBackend._pin_visible_gpu_order_for_split(env)
+ assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID"
+ assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
+
+
+def test_split_pin_keeps_mask_order_when_picker_reported_it(monkeypatch):
+ # Torch-fallback enumeration (no nvidia-smi) reports devices in inherited
+ # mask order, so the picker's split list follows the mask -- the pin must
+ # keep that order, not re-sort it into a mismatch.
+ _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [3, 1])
+ env = {"CUDA_VISIBLE_DEVICES": "3,1"}
+ LlamaCppBackend._pin_visible_gpu_order_for_split(env)
+ assert env["CUDA_VISIBLE_DEVICES"] == "3,1"
+
+
+def test_split_pin_falls_back_to_ascending_without_report(monkeypatch):
+ # Enumeration unavailable: ascending physical is the best guess (it matches
+ # the dominant nvidia-smi report order).
+ _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = None)
+ env = {"CUDA_VISIBLE_DEVICES": "3,1"}
+ LlamaCppBackend._pin_visible_gpu_order_for_split(env)
+ assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
+
+
+def test_split_pin_without_mask_only_sets_pci_order(monkeypatch):
+ # No inherited mask (or a UUID/MIG one resolving to None): enumeration order
+ # is fully fixed by CUDA_DEVICE_ORDER, so no mask is written.
+ _patch_split_pin_env(monkeypatch, inherited = None, reported = None)
+ env = {}
+ LlamaCppBackend._pin_visible_gpu_order_for_split(env)
+ assert env == {"CUDA_DEVICE_ORDER": "PCI_BUS_ID"}
+
+
+def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch):
+ # ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR
+ # mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP
+ # would index into the already-reduced set).
+ _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = "6.0")
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"}
+ LlamaCppBackend._pin_visible_gpu_order_for_split(env)
+ assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
+ assert env["HIP_VISIBLE_DEVICES"] == "1,3"
+ assert "ROCR_VISIBLE_DEVICES" not in env
+
+
+# ── Diffusion single-device selection ───────────────────────────────────────
+
+
+def test_diffusion_gpu_arg_uses_lowest_explicit_physical_id(monkeypatch):
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1")
+ monkeypatch.setenv("DG_GPU", "7")
+ assert LlamaCppBackend._diffusion_gpu_arg([3, 1]) == "1"
+
+
+def test_diffusion_gpu_arg_preserves_parent_mask_order(monkeypatch):
+ monkeypatch.delenv("DG_GPU", raising = False)
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1")
+ assert LlamaCppBackend._diffusion_gpu_arg(None) == "3"
+
+
+def test_diffusion_gpu_arg_honors_override_and_cpu_mask(monkeypatch):
+ monkeypatch.setenv("DG_GPU", "GPU-abc")
+ assert LlamaCppBackend._diffusion_gpu_arg(None) == "GPU-abc"
+ assert LlamaCppBackend._diffusion_gpu_arg(None, cpu_only = True) == ""
+
+
+# ── Deliberate zero-offload (manual gpu_layers=0): training-skip flag ─────────
+
+
+def test_zero_offload_flag_false_without_companions():
+ # CPU-only by construction: False lets training skip unloading a server that
+ # holds no VRAM.
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--fit", "off"]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is False
+
+
+@pytest.mark.parametrize(
+ "companion",
+ ["--mmproj", "--model-draft", "-md", "--spec-draft-model", "-hfd"],
+)
+def test_zero_offload_flag_true_with_companion(companion):
+ # mmproj / a drafter offload to GPU regardless of --gpu-layers, so the
+ # server still holds VRAM and training must unload it. Drafter detection
+ # reuses the extras parser, so pass-through aliases count too.
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", companion, "x.gguf"]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
+
+
+def test_zero_offload_flag_true_with_inline_companion_forms():
+ cmd = ["llama-server", "-m", "model.gguf", "--spec-draft-model=x.gguf"]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
+ cmd = ["llama-server", "-m", "model.gguf", "--mmproj=proj.gguf"]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
+
+
+def test_zero_offload_flag_true_with_env_drafter():
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
+ env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "x.gguf"}
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True
+
+
+@pytest.mark.parametrize(
+ "device_args",
+ [
+ ["--device", "CUDA0"],
+ ["--device=CUDA0"],
+ ["-dev", "CUDA0"],
+ ["--spec-draft-device", "CUDA0"],
+ ["--device-draft=CUDA0"],
+ ],
+)
+def test_zero_offload_flag_true_with_device_pin(device_args):
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
+
+
+def test_zero_offload_flag_true_with_env_device_pin():
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
+ env = {"LLAMA_ARG_DEVICE": "CUDA0"}
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True
+
+
+@pytest.mark.parametrize(
+ ("device_args", "env"),
+ [
+ (["--device", "cpu"], {}),
+ (["--device=none"], {}),
+ (["--spec-draft-device", "cpu"], {}),
+ ([], {"LLAMA_ARG_DEVICE": "none"}),
+ (["--device", "CUDA0", "--device", "cpu"], {}),
+ ],
+)
+def test_zero_offload_flag_false_with_cpu_device_pin(device_args, env):
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is False
+
+
+def test_zero_offload_flag_true_with_surviving_tensor_mode():
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--split-mode", "tensor"]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
+
+
+def test_zero_offload_flag_true_for_unmasked_vulkan(monkeypatch):
+ monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True))
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
+
+
+def test_zero_offload_flag_none_without_gpus():
+ cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
+ assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [], {}) is None
+
+
+def test_cmd_has_gpu_companion_detection():
+ # The env mask for CPU-only zero-offload loads keys off this scan: any
+ # --mmproj form or a drafter (flag aliases / env) keeps the GPUs visible.
+ has = LlamaCppBackend._cmd_has_gpu_companion
+ assert has(["llama-server", "-m", "m.gguf"], {}) is False
+ assert has(["llama-server", "--mmproj", "p.gguf"], {}) is True
+ assert has(["llama-server", "--mmproj=p.gguf"], {}) is True
+ assert has(["llama-server", "-md", "d.gguf"], {}) is True
+ assert has(["llama-server"], {"LLAMA_ARG_SPEC_DRAFT_MODEL": "d.gguf"}) is True
+
+
+def test_cmd_companion_ignores_cpu_forced_drafter():
+ # A CPU-pinned drafter holds no VRAM: the zero-offload mask may hide the GPUs
+ # and training may leave the server alone.
+ has = LlamaCppBackend._cmd_has_gpu_companion
+ cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0"]
+ assert has(cmd, {}) is False
+ cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-device", "cpu"]
+ assert has(cmd, {}) is False
+ # mmproj still counts even alongside a CPU drafter.
+ cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0", "--mmproj", "p.gguf"]
+ assert has(cmd, {}) is True
diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py
index 69ad560788..d4f2fbe993 100644
--- a/studio/backend/tests/test_gpu_selection.py
+++ b/studio/backend/tests/test_gpu_selection.py
@@ -853,7 +853,13 @@ class TestRouteErrors(unittest.TestCase):
self.assertIn("only supported on CUDA devices", str(exc_info.exception))
- def test_inference_route_rejects_gpu_ids_for_gguf(self):
+ def test_inference_route_validates_gpu_ids_for_gguf(self):
+ # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still
+ # validated: a rejected pick surfaces as a clean 400, not the old
+ # "not supported for GGUF" rejection. Patch the validator so the test
+ # is deterministic regardless of the host's (or a prior test's) GPU env.
+ import utils.hardware.hardware as hardware_mod
+
inference_route = _load_route_module(
"inference_route_module_for_gguf_gpu_ids_test",
"routes/inference.py",
@@ -887,6 +893,11 @@ class TestRouteErrors(unittest.TestCase):
),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
+ patch.object(
+ hardware_mod,
+ "resolve_requested_gpu_ids",
+ side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"),
+ ),
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
@@ -901,8 +912,11 @@ class TestRouteErrors(unittest.TestCase):
)
)
+ # The validator's ValueError becomes a clean 400 (not the removed
+ # "not supported for GGUF" rejection).
self.assertEqual(exc_info.exception.status_code, 400)
- self.assertIn("GGUF", exc_info.exception.detail)
+ self.assertIn("gpu_ids", exc_info.exception.detail.lower())
+ self.assertNotIn("not supported", exc_info.exception.detail.lower())
def test_training_route_returns_400_for_invalid_gpu_ids(self):
training_route = _load_route_module(
diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py
index f320d29a02..662c918305 100644
--- a/studio/backend/tests/test_llama_cpp_no_context_shift.py
+++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py
@@ -118,9 +118,17 @@ def test_flag_sits_inside_the_base_cmd_list():
"conditional branch -- otherwise some code paths would still "
"run with silent context shift enabled."
)
- # Pin that it sits next to -c / --ctx so the grouping makes sense.
- assert '"-c"' in block
assert '"--flash-attn"' in block
+ # -c is emitted in the conditional right after the base list, not inside
+ # it: auto-fit (--fit on with no pinned context) must omit -c entirely,
+ # because "-c 0" pins the full native context and disables --fit's
+ # VRAM-based sizing. Pin that it still sits next to the base block so the
+ # context grouping stays intact.
+ after = rest[end_rel : end_rel + 1000]
+ assert '"-c"' in after, (
+ "-c must still be emitted in the conditional immediately after the "
+ "base cmd list (omitted only in auto-fit, where --fit sizes context)."
+ )
def _iter_lines_with_offset(text: str):
diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py
index 488645ee5a..fe1e67edad 100644
--- a/studio/backend/tests/test_llama_cpp_props_readback.py
+++ b/studio/backend/tests/test_llama_cpp_props_readback.py
@@ -225,31 +225,46 @@ def test_kv_unified_added_for_multi_slot():
"""Explicit --parallel N disables llama-server's auto-slots kv-unified
default, splitting -c into per-slot windows of -c/N; Unsloth must restore
the shared pool so one request can use the full advertised context."""
- flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL)
+ flags = LlamaCppBackend._ctx_integrity_flags(4, False, False, 98304, 98304, _CAPS_ALL)
assert "--kv-unified" in flags
def test_kv_unified_skipped_for_single_slot_or_old_build():
assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags(
- 1, False, 98304, 98304, _CAPS_ALL
+ 1, False, False, 98304, 98304, _CAPS_ALL
)
assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags(
- 4, False, 98304, 98304, _CAPS_NONE
+ 4, False, False, 98304, 98304, _CAPS_NONE
)
def test_fit_ctx_floors_explicit_request_under_fit():
- flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL)
+ # An explicit requested ctx floors --fit-ctx at that value on any --fit
+ # path, including legacy auto (auto_fit False).
+ flags = LlamaCppBackend._ctx_integrity_flags(1, True, False, 98304, 98304, _CAPS_ALL)
assert flags[flags.index("--fit-ctx") + 1] == "98304"
-def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support():
+def test_fit_ctx_skipped_without_fit_or_support():
+ # No --fit on -> no --fit-ctx.
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
- 1, False, 98304, 98304, _CAPS_ALL
+ 1, False, False, 98304, 98304, _CAPS_ALL
)
- assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL)
+ # --fit on but the binary doesn't support --fit-ctx.
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
- 1, True, 98304, 98304, _CAPS_NONE
+ 1, True, True, 98304, 98304, _CAPS_NONE
+ )
+
+
+def test_fit_ctx_floors_auto_request_at_8192_only_under_auto_fit():
+ # Manual + Auto (auto_fit) floors the auto window at 8192 so --fit can't
+ # shrink it to a tiny size.
+ flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 262144, _CAPS_ALL)
+ assert flags[flags.index("--fit-ctx") + 1] == "8192"
+ # Legacy auto (fit on but not auto_fit) emits -c 0 to pin native, so the
+ # 8192 floor must NOT ride along and override that pin.
+ assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
+ 1, True, False, 0, 262144, _CAPS_ALL
)
diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py
index ba52afad1c..c6d16363f8 100644
--- a/studio/backend/tests/test_llama_server_args.py
+++ b/studio/backend/tests/test_llama_server_args.py
@@ -747,6 +747,34 @@ def test_strip_shadowing_flags_defaults_strip_split_mode_too():
assert strip_shadowing_flags(["--split-mode", "tensor"]) == []
+def test_strip_offload_is_opt_in_and_covers_moe():
+ base = dict(
+ strip_context = False,
+ strip_cache = False,
+ strip_spec = False,
+ strip_template = False,
+ strip_split_mode = False,
+ )
+ # Default: offload (incl. MoE) flags are NOT stripped.
+ assert strip_shadowing_flags(["--n-cpu-moe", "8", "--top-k", "20"], **base) == [
+ "--n-cpu-moe",
+ "8",
+ "--top-k",
+ "20",
+ ]
+ # Opt-in strips layer AND MoE offload flags (value-aware), keeps the rest.
+ assert strip_shadowing_flags(
+ ["--n-cpu-moe", "8", "--gpu-layers", "33", "--fit", "off", "--top-k", "20"],
+ **base,
+ strip_offload = True,
+ ) == ["--top-k", "20"]
+ # Boolean --cpu-moe drops the flag only, not the following value.
+ assert strip_shadowing_flags(["--cpu-moe", "--seed", "-1"], **base, strip_offload = True) == [
+ "--seed",
+ "-1",
+ ]
+
+
@pytest.mark.parametrize(
"args",
[
@@ -796,6 +824,23 @@ def test_strip_split_mode_only_drops_tensor_split_too():
assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == []
+def test_strip_tensor_split_alone_preserves_split_mode():
+ # Manual mode emits its own --tensor-split, so an inherited ratio is dropped
+ # -- but the user's --split-mode row/none/layer choice (which the manual
+ # ratio toggle can't express) must survive. strip_tensor_split removes only
+ # the ratio, unlike strip_split_mode which removes the whole group.
+ out = strip_shadowing_flags(
+ ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"],
+ strip_context = False,
+ strip_cache = False,
+ strip_spec = False,
+ strip_template = False,
+ strip_split_mode = False,
+ strip_tensor_split = True,
+ )
+ assert out == ["--split-mode", "row", "--top-k", "20"]
+
+
def test_strip_shadowing_flags_keeps_model_draft_without_spec():
out = strip_shadowing_flags(
["--model-draft", "/custom/mtp.gguf"],
diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py
index 06f72d3b9f..00c7aeac69 100644
--- a/studio/backend/tests/test_tensor_parallel.py
+++ b/studio/backend/tests/test_tensor_parallel.py
@@ -262,9 +262,12 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode():
src = _load_model_source()
assert '"--tensor-split"' in src
gate = src.find("if tensor_parallel:")
- ts = src.find('"--tensor-split"')
+ # Find the TP block's emission (after the gate); manual mode emits its own
+ # --tensor-split earlier in the source from the user's per-GPU shares.
+ ts = src.find('"--tensor-split"', gate)
nxt_else = src.find("self._tensor_parallel = False")
assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`"
+ assert "tp_tensor_split" in src[gate:nxt_else]
def test_mtp_decode_probe_wired_under_tensor_parallel():
diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py
index fb0989b306..d1372ca415 100644
--- a/studio/backend/tests/test_tp_vision_regression.py
+++ b/studio/backend/tests/test_tp_vision_regression.py
@@ -126,10 +126,21 @@ _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: tensor needs >= 2 GPUs clearing the compute-buffer reserve. Gated
+ # on plan_tp (not raw tensor_parallel) so manual mode skips this planner (#6414).
+ "plan_tp and len(tp_gpus) < 2",
# Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split.
"_tp_weight_budget_mib <= _tp_required_mib",
+ # Manual mode, Auto layers: --fit owns memory and is incompatible with a
+ # tensor split, so TP is dropped (surfaced via logger.info) before the
+ # cache-drop, so a quantized KV survives into the --fit load (#6414).
+ "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers < 0)",
+ # Manual mode, explicit layers: a tensor split still needs >= 2 GPUs in use.
+ "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers >= 0) and (self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2)",
+ # Manual mode, zero layers: nothing to split on the GPU, and a tensor-mode
+ # launch under the CPU-only GPU mask (no visible devices) aborts the server
+ # instead of the intended CPU-only load (#6414).
+ "gpu_memory_mode == 'manual' and gpu_layers == 0",
}
@@ -364,7 +375,7 @@ def test_compute_buffer_downgrade_preserves_multi_gpu_intent():
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")
+ gate = src.find("plan_tp 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)
diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py
index c24ec28e1d..50b3cd3513 100644
--- a/studio/backend/utils/models/gguf_metadata.py
+++ b/studio/backend/utils/models/gguf_metadata.py
@@ -50,9 +50,11 @@ _CACHE_MAX_ENTRIES = 4096
# keyed by (file cache key, wanted key). None = key absent / file unreadable.
_BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {}
-# Native training context length (``{arch}.context_length``). None = absent /
-# unreadable. Lets the UI show the real context ceiling before a model loads.
-_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {}
+# GGUF header dims for the staged/deferred-load UI: context_length, layer_count
+# (block_count), and moe_layer_count (block_count minus leading dense layers; 0
+# if not MoE). One cached pass fills all three so the staged sheet can size every
+# slider before the model loads. None = unreadable / not a GGUF.
+_DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {}
def _cache_key(path: str) -> Optional[_CacheKey]:
@@ -142,32 +144,45 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]:
return out
-def read_gguf_context_length(path: str) -> Optional[int]:
- """Return the GGUF's native training context length (``{arch}.context_length``),
- or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size).
- Lets the UI populate the context slider before the model is loaded."""
+def read_gguf_staged_dims(path: str) -> Optional[Dict[str, Optional[int]]]:
+ """GGUF header dims for the staged-load UI in one cached pass:
+ ``{"context_length", "layer_count", "moe_layer_count"}``. Each may be None
+ when absent (moe_layer_count is 0 for a dense model). Returns ``None`` if not
+ a GGUF / unreadable. Cached by (path, mtime, size). Lets the staged sheet size
+ the context, GPU-layers and MoE sliders before the model loads."""
key = _cache_key(path)
if key is None:
return None
with _CACHE_LOCK:
- if key in _CONTEXT_CACHE:
- return _CONTEXT_CACHE[key]
- result = _parse_gguf_context_length(path)
+ if key in _DIMS_CACHE:
+ return _DIMS_CACHE[key]
+ result = _parse_gguf_staged_dims(path)
with _CACHE_LOCK:
- while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES:
+ while len(_DIMS_CACHE) >= _CACHE_MAX_ENTRIES:
try:
- _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE)))
+ _DIMS_CACHE.pop(next(iter(_DIMS_CACHE)))
except StopIteration:
break
- _CONTEXT_CACHE[key] = result
+ _DIMS_CACHE[key] = result
return result
-def _parse_gguf_context_length(path: str) -> Optional[int]:
- # The context key is architecture-namespaced (``llama.context_length`` etc.),
- # so we learn the key only after reading ``general.architecture``. GGUF writes
- # general.* before arch.* keys, matching the loader's own parser.
- ctx_key: Optional[str] = None
+def read_gguf_context_length(path: str) -> Optional[int]:
+ """Native training context length (``{arch}.context_length``), or ``None``.
+ Thin accessor over read_gguf_staged_dims."""
+ dims = read_gguf_staged_dims(path)
+ return dims["context_length"] if dims else None
+
+
+def _parse_gguf_arch_uints(path: str, wanted_suffixes: frozenset[str]) -> Optional[Dict[str, int]]:
+ """Walk a GGUF header once and return the requested architecture-namespaced
+ uint (vtype 4/10) keys, e.g. ``{"block_count": 32}``. Keys are
+ ``{arch}.``; the arch is learned from ``general.architecture`` (GGUF
+ writes general.* before arch.* keys, matching the loader's own parser).
+ Returns ``None`` if not a GGUF / unreadable, else a dict (possibly empty or
+ partial when some keys are absent)."""
+ arch: Optional[str] = None
+ found: Dict[str, int] = {}
try:
with open(path, "rb") as f:
head = f.read(24)
@@ -204,28 +219,68 @@ def _parse_gguf_context_length(path: str) -> Optional[int]:
sbytes = f.read(slen)
if len(sbytes) < slen:
break
- ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length"
- elif ctx_key is not None and key == ctx_key and vtype in (4, 10):
+ arch = sbytes.decode("utf-8", "replace")
+ elif (
+ arch is not None
+ and vtype in (4, 10)
+ and key.startswith(f"{arch}.")
+ and key[len(arch) + 1 :] in wanted_suffixes
+ ):
width = 4 if vtype == 4 else 8
n_bytes = f.read(width)
if len(n_bytes) < width:
break
- value = struct.unpack(" 0 else None
+ found[key[len(arch) + 1 :]] = struct.unpack(
+ " Optional[Dict[str, Optional[int]]]:
+ vals = _parse_gguf_arch_uints(
+ path,
+ frozenset(
+ {
+ "context_length",
+ "block_count",
+ "expert_count",
+ "leading_dense_block_count",
+ }
+ ),
+ )
+ if vals is None:
+ return None
+ ctx = vals.get("context_length")
+ block = vals.get("block_count")
+ # A real context/layer count is positive; treat 0/garbage as absent so the
+ # UI never builds a slider with max < min.
+ context_length = ctx if ctx and ctx > 0 else None
+ layer_count = block if block and block > 0 else None
+ # MoE layer count = block_count - leading dense layers, only when experts
+ # exist; else 0 (dense -> slider hidden). Mirrors n_moe_layers in
+ # core/inference/llama_cpp.py.
+ if not vals.get("expert_count") or not block:
+ moe_layer_count: Optional[int] = 0
+ else:
+ moe_layer_count = max(0, block - (vals.get("leading_dense_block_count") or 0))
+ return {
+ "context_length": context_length,
+ "layer_count": layer_count,
+ "moe_layer_count": moe_layer_count,
+ }
# Strings (8) and arrays (9) are handled inline.
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
index ec75b17f20..08492ab480 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
+++ b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
@@ -2,7 +2,9 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Per-model pre-load inference settings, persisted in localStorage so the load
-// dialog can offer "Remember settings for ".
+// dialog can offer "Remember settings for ". GGUF picks only: every
+// field is a llama.cpp load knob, so all save/restore call sites gate on
+// GGUF-ness (a non-GGUF blob would only snapshot leftover standing values).
const KEY = "unsloth_load_settings";
@@ -12,14 +14,22 @@ export interface RememberedLoadSettings {
speculativeType: string | null;
specDraftNMax: number | null;
tensorParallel: boolean;
+ // GPU Memory controls. Optional so an older blob (which lacked them) still
+ // parses, leaving the live knobs untouched on apply. The mode is kept with the
+ // manual knobs (gpuLayers/nCpuMoe are ignored outside Manual mode). A null
+ // selectedGpuIds is meaningful (all GPUs), so it's distinguished from absent.
+ // The per-GPU split ratio is deliberately NOT remembered: it's positionally
+ // bound to the exact GPU set/order and unvalidated, so it would mismatch.
+ gpuMemoryMode?: "auto" | "manual";
+ gpuLayers?: number;
+ nCpuMoe?: number;
+ selectedGpuIds?: number[] | null;
}
-// Storage key for a pick's remembered settings. The remembered knobs are
-// VRAM-budget driven (context override, KV-cache dtype, tensor-parallel), so the
-// right values differ per quant. An HF repo collapses all its GGUF variants into
-// one `id`, so fold the variant in to scope settings per quant. Local .gguf
-// paths key by their file path (already file-specific); native drag-drop files
-// key by display label, so same-named files in different folders share an entry.
+// Storage key for a pick's remembered settings, scoped per quant (the VRAM-budget
+// knobs differ per quant). An HF repo collapses its GGUF variants into one `id`,
+// so fold the variant in. Local .gguf paths are already file-specific; native
+// drag-drop files key by display label, so same-named files share an entry.
export function rememberedLoadSettingsKey(selection: {
id: string;
ggufVariant?: string | null;
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 0bf46e7343..7083f02288 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -45,12 +45,18 @@ import {
import {
type PendingImageEditReference,
type RagAutoInject,
+ GPU_LAYERS_AUTO,
+ loadedGpuMemoryFieldsUnlessStaged,
+ reconcilePersistedGpuIds,
resolveLoadedSpeculativeSettings,
resolveSpeculativeSettingsForLoad,
+ persistGpuMemoryModeOnLoad,
resolveToolsEnabledOnLoad,
saveSpeculativeType,
useChatRuntimeStore,
} from "../stores/chat-runtime-store";
+import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "../presets/preset-policy";
+import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info";
import { useExternalProvidersStore } from "../stores/external-providers-store";
import {
shouldPreserveFullOutput,
@@ -1489,6 +1495,13 @@ async function autoLoadSmallestModel(): Promise<{
max_seq_length: number;
is_lora: boolean;
gguf_variant?: string | null;
+ // GGUF-only: scopes the training guard to the same placement policy /load
+ // will use. Manual mode must match because it makes placement user-owned.
+ // The layer/MoE/split/KV/spec knobs are deliberately not sent: Auto mode's
+ // guard sizes conservatively, while Manual mode bypasses that estimate.
+ // The safetensors fallback omits both fields and uses HF auto-placement.
+ gpu_ids?: number[];
+ gpu_memory_mode?: "auto" | "manual";
}): Promise {
const validation = await validateModel({
...payload,
@@ -1520,12 +1533,18 @@ async function autoLoadSmallestModel(): Promise<{
return false;
}
const currentStore = useChatRuntimeStore.getState();
- const remembered = loadRememberedLoadSettings(
- rememberedLoadSettingsKey({
- id: candidate.id,
- ggufVariant: candidate.ggufVariant,
- }),
- );
+ // Blobs are saved for GGUF picks only (the sheet gates on it), so don't
+ // let a legacy non-GGUF blob feed a stale context/spec choice into a
+ // safetensors auto-load.
+ const remembered =
+ candidate.kind === "gguf"
+ ? loadRememberedLoadSettings(
+ rememberedLoadSettingsKey({
+ id: candidate.id,
+ ggufVariant: candidate.ggufVariant,
+ }),
+ )
+ : null;
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
modelId: candidate.id,
ggufVariant: candidate.ggufVariant,
@@ -1537,6 +1556,38 @@ async function autoLoadSmallestModel(): Promise<{
maxSeqLength: candidate.maxSeqLength,
presetSource: currentStore.activePresetSource,
});
+ // The GPU knobs are per-model, so read them from the same remembered
+ // settings that fed effectiveMaxSeqLength -- on a background auto-load the
+ // live store holds session defaults, not the saved Manual mode / layer pin /
+ // GPU pick. Absent fields fall back like applyRememberedLoadSettings: the
+ // mode to the store (a persisted standing preference), the per-model knobs to
+ // their defaults. The saved GPU pick is reconciled against the GPUs present
+ // now, like the interactive restore.
+ const effectiveGpuMemoryMode =
+ remembered?.gpuMemoryMode ?? currentStore.gpuMemoryMode;
+ const effectiveGpuLayers = remembered?.gpuLayers ?? GPU_LAYERS_AUTO;
+ const effectiveNCpuMoe = remembered?.nCpuMoe ?? 0;
+ if (remembered?.selectedGpuIds != null) {
+ // Warm the device cache first: on a cold cache the reconcile passes the
+ // saved pick through unvalidated, and a stale cross-host pick then fails
+ // the load with the picker hidden.
+ await ensureGpuDeviceCache();
+ }
+ const effectiveGpuIds =
+ remembered?.selectedGpuIds !== undefined
+ ? reconcilePersistedGpuIds(remembered.selectedGpuIds)
+ : null;
+ // Under Manual GPU memory + Auto layers, llama.cpp's --fit owns context
+ // sizing, so send 0 (or the pinned length). GGUF-only; a no-op otherwise.
+ // The context pin is per-model too, so it comes from remembered settings,
+ // not the live store.
+ const fitMaxSeqLength = resolveFitMaxSeqLength(
+ candidate.kind === "gguf",
+ effectiveGpuMemoryMode,
+ effectiveGpuLayers,
+ remembered?.contextLength ?? null,
+ effectiveMaxSeqLength,
+ );
const effectiveSpeculativeType =
remembered?.speculativeType ?? specSettings.speculativeType;
const effectiveSpecDraftNMax =
@@ -1544,9 +1595,16 @@ async function autoLoadSmallestModel(): Promise<{
if (
!(await canAutoLoad({
model_path: candidate.id,
- max_seq_length: effectiveMaxSeqLength,
+ max_seq_length: fitMaxSeqLength,
is_lora: false,
gguf_variant: candidate.ggufVariant,
+ // The same remembered-derived GPU pick the load below sends.
+ ...(candidate.kind === "gguf"
+ ? {
+ gpu_ids: effectiveGpuIds ?? undefined,
+ gpu_memory_mode: effectiveGpuMemoryMode,
+ }
+ : {}),
}))
) {
skippedAutoLoadCandidates.add(
@@ -1558,7 +1616,7 @@ async function autoLoadSmallestModel(): Promise<{
const loadResp = await loadModel({
model_path: candidate.id,
hf_token: hfToken,
- max_seq_length: effectiveMaxSeqLength,
+ max_seq_length: fitMaxSeqLength,
load_in_4bit: true,
is_lora: false,
gguf_variant: candidate.ggufVariant,
@@ -1567,8 +1625,22 @@ async function autoLoadSmallestModel(): Promise<{
speculative_type: effectiveSpeculativeType,
spec_draft_n_max: effectiveSpecDraftNMax,
tensor_parallel: remembered?.tensorParallel ?? false,
+ // GGUF-only: the safetensors fallback loads via HF auto-placement (no
+ // explicit pins). The split ratio is deliberately never remembered
+ // (positionally bound to an exact GPU set), so auto-load leaves llama.cpp's
+ // free-VRAM default in charge rather than sending a stale store value.
+ ...(candidate.kind === "gguf"
+ ? {
+ gpu_memory_mode: effectiveGpuMemoryMode,
+ gpu_layers: effectiveGpuLayers,
+ n_cpu_moe: effectiveNCpuMoe,
+ gpu_ids: effectiveGpuIds ?? undefined,
+ }
+ : {}),
});
saveSpeculativeType(effectiveSpeculativeType);
+ // Self-gates on is_gguf (skips diffusion), so persists only for a real GGUF load.
+ persistGpuMemoryModeOnLoad(loadResp, effectiveGpuMemoryMode);
useChatRuntimeStore
.getState()
.setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined);
@@ -1597,6 +1669,15 @@ async function autoLoadSmallestModel(): Promise<{
store.setModels([...store.models, autoModel]);
}
if (candidate.kind === "gguf") {
+ // Keep an explicit Manual+Auto context pin the load just applied (so a
+ // later Apply doesn't silently revert it to auto-fit sizing), mirroring
+ // the interactive path's keepCustomCtx; other cases baseline on
+ // ggufContextLength.
+ const keepCustomCtx = resolveManualAutoCtxPin(
+ effectiveGpuMemoryMode,
+ effectiveGpuLayers,
+ remembered?.contextLength ?? null,
+ );
useChatRuntimeStore.setState({
ggufContextLength: loadResp.context_length ?? 131072,
ggufMaxContextLength:
@@ -1613,6 +1694,10 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
+ ...loadedGpuMemoryFieldsUnlessStaged(loadResp, {
+ customContextLength: keepCustomCtx,
+ }),
+ loadedCustomContextLength: keepCustomCtx,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
@@ -1633,6 +1718,9 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
+ // Non-GGUF response: clears any stale GPU baseline a prior manual-GPU
+ // GGUF load left, matching the interactive/status sibling load paths.
+ ...loadedGpuMemoryFieldsUnlessStaged(loadResp),
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
@@ -1820,12 +1908,17 @@ async function autoLoadSmallestModel(): Promise<{
duration: 30000,
});
try {
+ const rt = useChatRuntimeStore.getState();
if (
!(await canAutoLoad({
model_path: "unsloth/Qwen3.5-4B-MTP-GGUF",
max_seq_length: 0,
is_lora: false,
gguf_variant: "UD-Q4_K_XL",
+ // The same live-store GPU pick the load below sends (a fresh default
+ // model has no remembered settings to prefer).
+ gpu_ids: rt.selectedGpuIds ?? undefined,
+ gpu_memory_mode: rt.gpuMemoryMode,
}))
) {
toast.dismiss(toastId);
@@ -1835,6 +1928,9 @@ async function autoLoadSmallestModel(): Promise<{
const loadResp = await loadModel({
model_path: "unsloth/Qwen3.5-4B-MTP-GGUF",
hf_token: hfToken,
+ // Model default under both modes: Auto layers + no pin means
+ // resolveFitMaxSeqLength returns 0 for every mode (the canAutoLoad
+ // preflight above sends the same).
max_seq_length: 0,
load_in_4bit: true,
is_lora: false,
@@ -1842,8 +1938,20 @@ async function autoLoadSmallestModel(): Promise<{
trust_remote_code: trustRemoteCode,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,
+ // GPU Memory mode is a standing preference, so honor it on auto-load.
+ // The layer/MoE/split knobs and the context pin are per-model: the live
+ // store may hold edits drafted for a staged pick, and a fresh default
+ // model has no remembered settings, so those stay at their defaults like
+ // the cached-candidate path. The GPU pick deliberately differs (it's the
+ // picker's current on-screen selection, which the canAutoLoad preflight
+ // above already committed to).
+ gpu_memory_mode: rt.gpuMemoryMode,
+ gpu_layers: GPU_LAYERS_AUTO,
+ n_cpu_moe: 0,
+ gpu_ids: rt.selectedGpuIds ?? undefined,
});
saveSpeculativeType(specSettings.speculativeType);
+ persistGpuMemoryModeOnLoad(loadResp, rt.gpuMemoryMode);
useChatRuntimeStore
.getState()
.setCheckpoint("unsloth/Qwen3.5-4B-MTP-GGUF", "UD-Q4_K_XL");
@@ -1880,6 +1988,10 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
+ ...loadedGpuMemoryFieldsUnlessStaged(loadResp),
+ // Drives the GPU Memory controls' diffusion gate; set alongside the
+ // GPU fields on every load path so the gate can't read stale.
+ loadedIsDiffusion: loadResp.is_diffusion ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedIsMultimodal: isMultimodalResponse(loadResp),
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index ebf9461172..0f6af38033 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -127,28 +127,38 @@ export async function validateModel(
native_path_lease: payload.nativePathLease ?? null,
hf_token: payload.hf_token,
gguf_variant: payload.gguf_variant ?? null,
- // Send the intended load settings so validate's VRAM check matches the
- // follow-up /load and doesn't unload for a load /load would then reject.
+ // Intended load settings so validate's preflight matches the follow-up
+ // /load. Default placement is sized against the selected GPUs.
max_seq_length: payload.max_seq_length,
load_in_4bit: payload.load_in_4bit,
+ gpu_ids: payload.gpu_ids,
+ // Manual placement is an explicit override: Auto layers use llama.cpp
+ // --fit, while a pinned layer count is owned by the user. Tell validate
+ // so it applies the same training-guard policy as /load.
+ gpu_memory_mode: payload.gpu_memory_mode,
}),
});
return parseJsonOrThrow(response);
}
/**
- * Read a GGUF's native context length from its local header (no GPU load, no
- * download). Returns null when the file isn't downloaded yet, the model isn't a
- * GGUF, or it's gated. For a native (drag-drop / picked) file, pass
- * `nativePathToken` so the backend reads the granted local path. Used by the
- * deferred-load staging flow to fill the context slider before the single load.
+ * Read a GGUF's header dims (native context length, total layer count, MoE
+ * expert-layer count) from its local file (no GPU load, no download). All are
+ * null when the file isn't downloaded yet, the model isn't a GGUF, or it's
+ * gated. For a native (drag-drop / picked) file, pass `nativePathToken` so the
+ * backend reads the granted local path. Used by the deferred-load staging flow
+ * to size the context, GPU-layers and MoE sliders before the single load.
*/
-export async function fetchGgufContextLength(payload: {
+export async function fetchGgufStagedMetadata(payload: {
model_path: string;
gguf_variant?: string | null;
hf_token?: string | null;
nativePathToken?: string | null;
-}): Promise {
+}): Promise<{
+ contextLength: number | null;
+ layerCount: number | null;
+ moeLayerCount: number | null;
+}> {
let nativePathLease: string | null = null;
if (payload.nativePathToken) {
try {
@@ -156,8 +166,8 @@ export async function fetchGgufContextLength(payload: {
await consumeNativePathToken(payload.nativePathToken, "validate-model")
).nativePathLease;
} catch {
- // Lease expired / revoked: degrade to no context (the load can re-mint).
- return null;
+ // Lease expired / revoked: degrade to no metadata (the load can re-mint).
+ return { contextLength: null, layerCount: null, moeLayerCount: null };
}
}
const response = await authFetch("/api/inference/validate", {
@@ -172,7 +182,11 @@ export async function fetchGgufContextLength(payload: {
}),
});
const res = await parseJsonOrThrow(response);
- return res.context_length ?? null;
+ return {
+ contextLength: res.context_length ?? null,
+ layerCount: res.layer_count ?? null,
+ moeLayerCount: res.moe_layer_count ?? null,
+ };
}
export async function unloadModel(payload: UnloadModelRequest): Promise {
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index ec0ad977bf..217eaf8b6d 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -1445,9 +1445,11 @@ export function ChatPage({
// were already seeded on stage, so keepSpeculative only when a config was
// saved -- otherwise the standing speculative preference should win.
autoLoadStagedRef.current = (pending) => {
- const remembered = loadRememberedLoadSettings(
- rememberedLoadSettingsKey(pending),
- );
+ // Blobs are saved for GGUF picks only (the sheet gates on it), so don't
+ // let a legacy non-GGUF blob claim a seeded config here.
+ const remembered = hasGgufSource(pending)
+ ? loadRememberedLoadSettings(rememberedLoadSettingsKey(pending))
+ : null;
void selectModel({
...pending,
isDownloaded: true,
@@ -2813,6 +2815,11 @@ export function ChatPage({
selectModel({
id: state.params.checkpoint,
ggufVariant: state.activeGgufVariant ?? undefined,
+ // A native (drag-drop / picked) GGUF's checkpoint is only a display
+ // label, so the reload needs its path token to re-mint a lease --
+ // else applying the now-exposed GPU/context controls can't resolve
+ // the file. Null for non-native loads, which reload by id as before.
+ nativePathToken: state.activeNativePathToken ?? undefined,
forceReload: true,
isDownloaded: true,
loadingDescription: "Reloading with updated chat template.",
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index cedd298ecf..b368a811fa 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -55,6 +55,7 @@ import { Switch } from "@/components/ui/switch";
import { Textarea } from "@/components/ui/textarea";
import { InfoHint } from "@/components/ui/info-hint";
import { Tooltip, TooltipContent } from "@/components/ui/tooltip";
+import { useGpuDevices } from "@/hooks/use-gpu-info";
import { useIsMobile } from "@/hooks/use-mobile";
import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check";
import { cn } from "@/lib/utils";
@@ -99,8 +100,11 @@ import {
providerSupportsFastMode,
} from "./provider-capabilities";
import {
+ GPU_LAYERS_AUTO,
+ distributeByWeight,
isPendingGguf,
pendingSelectionMatches,
+ rebalanceSplit,
useChatRuntimeStore,
} from "./stores/chat-runtime-store";
import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section";
@@ -250,6 +254,7 @@ function ParamSlider({
displayValue,
info,
valueSize,
+ disabled,
}: {
label: string;
value: number;
@@ -260,6 +265,7 @@ function ParamSlider({
displayValue?: string;
info?: ReactNode;
valueSize?: number;
+ disabled?: boolean;
}) {
return (
@@ -279,6 +285,7 @@ function ParamSlider({
displayValue={displayValue}
ariaLabel={label}
size={valueSize ?? 4}
+ disabled={disabled}
/>
onChange(snapToStep(v, step, min, max))}
className="panel-slider"
+ disabled={disabled}
/>
);
@@ -540,8 +548,17 @@ export function ChatSettingsPanel({
const base = slash >= 0 ? id.slice(slash + 1) : id;
return base || id;
})();
+ const activeNativePathToken = useChatRuntimeStore(
+ (s) => s.activeNativePathToken,
+ );
+ const loadedGgufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
+ // A GGUF loaded from a native path / direct .gguf has no HF variant, so key
+ // off the same signal the status hydration uses -- variant OR native token OR
+ // a GGUF context -- else the GPU Memory controls hide for a loaded local GGUF.
const isLoadedGguf =
- useChatRuntimeStore((s) => s.activeGgufVariant) != null;
+ useChatRuntimeStore((s) => s.activeGgufVariant) != null ||
+ activeNativePathToken != null ||
+ loadedGgufContextLength != null;
// While a pick is staged the sheet configures *that* model, so its GGUF-ness
// (not the currently loaded model's) decides whether the GGUF-only controls
// show. Otherwise a staged non-GGUF Hub repo would inherit the loaded GGUF's
@@ -607,6 +624,25 @@ export function ChatSettingsPanel({
const loadedTensorParallel = useChatRuntimeStore(
(s) => s.loadedTensorParallel,
);
+ const gpuMemoryMode = useChatRuntimeStore((s) => s.gpuMemoryMode);
+ const setGpuMemoryMode = useChatRuntimeStore((s) => s.setGpuMemoryMode);
+ const loadedGpuMemoryMode = useChatRuntimeStore((s) => s.loadedGpuMemoryMode);
+ const loadedIsDiffusion = useChatRuntimeStore((s) => s.loadedIsDiffusion);
+ const gpuLayers = useChatRuntimeStore((s) => s.gpuLayers);
+ const setGpuLayers = useChatRuntimeStore((s) => s.setGpuLayers);
+ const loadedGpuLayers = useChatRuntimeStore((s) => s.loadedGpuLayers);
+ const nCpuMoe = useChatRuntimeStore((s) => s.nCpuMoe);
+ const setNCpuMoe = useChatRuntimeStore((s) => s.setNCpuMoe);
+ const loadedNCpuMoe = useChatRuntimeStore((s) => s.loadedNCpuMoe);
+ const splitRatio = useChatRuntimeStore((s) => s.splitRatio);
+ const setSplitRatio = useChatRuntimeStore((s) => s.setSplitRatio);
+ const loadedSplitRatio = useChatRuntimeStore((s) => s.loadedSplitRatio);
+ const ggufLayerCount = useChatRuntimeStore((s) => s.ggufLayerCount);
+ const moeLayerCount = useChatRuntimeStore((s) => s.moeLayerCount);
+ const selectedGpuIds = useChatRuntimeStore((s) => s.selectedGpuIds);
+ const setSelectedGpuIds = useChatRuntimeStore((s) => s.setSelectedGpuIds);
+ const loadedGpuIds = useChatRuntimeStore((s) => s.loadedGpuIds);
+ const gpuDevices = useGpuDevices();
const chatTemplateOverride = useChatRuntimeStore(
(s) => s.chatTemplateOverride,
);
@@ -614,6 +650,9 @@ export function ChatSettingsPanel({
(s) => s.loadedChatTemplateOverride,
);
const customContextLength = useChatRuntimeStore((s) => s.customContextLength);
+ const loadedCustomContextLength = useChatRuntimeStore(
+ (s) => s.loadedCustomContextLength,
+ );
const setCustomContextLength = useChatRuntimeStore(
(s) => s.setCustomContextLength,
);
@@ -641,10 +680,14 @@ export function ChatSettingsPanel({
: null;
useEffect(() => {
if (!pendingKey) return;
- const saved = loadRememberedLoadSettings(pendingKey);
+ // GGUF-only, like the stageOrLoad / Hub restore paths: every remembered
+ // field is a llama.cpp knob, so a non-GGUF pick has nothing to restore --
+ // and applying its blob would clobber the standing gpuMemoryMode with a
+ // stale snapshot (the save on Load below is gated the same way).
+ const saved = pendingIsGguf ? loadRememberedLoadSettings(pendingKey) : null;
setRemember(saved != null);
if (saved) applyRememberedLoadSettings(saved);
- }, [pendingKey, applyRememberedLoadSettings]);
+ }, [pendingKey, pendingIsGguf, applyRememberedLoadSettings]);
// While staging, the sheet reflects the STAGED model, so its header context
// takes precedence over the loaded model's (which may differ or be larger).
const baseContext = pendingIsGguf ? stagedContextLength : ggufContextLength;
@@ -661,15 +704,132 @@ export function ChatSettingsPanel({
const ctxDisplayValue = customContextLength ?? baseContext ?? "";
const ctxMaxValue = baseNativeContext ?? baseContext ?? null;
const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
- const ctxDirty = customContextLength !== null;
+ const ctxDirty = customContextLength !== loadedCustomContextLength;
const specDirty = speculativeType !== loadedSpeculativeType;
const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax;
const tpDirty = tensorParallel !== (loadedTensorParallel ?? false);
+ // A loaded diffusion GGUF runs mode-agnostic (pins all layers on one GPU,
+ // ignores --fit/--gpu-layers), so the GPU Memory mode + manual controls don't
+ // apply -- hide them and don't let the preserved standing mode read as dirty.
+ // The GPU picker still applies (diffusion pins the chosen device). A staged pick
+ // keeps the controls (a pending pick's diffusion-ness isn't known until load).
+ const gpuModeApplies =
+ isGguf && (pendingSelection != null || !loadedIsDiffusion);
+ const gpuDirty =
+ gpuModeApplies && gpuMemoryMode !== (loadedGpuMemoryMode ?? "auto");
+ const isManual = gpuModeApplies && gpuMemoryMode === "manual";
+ // Manual with the GPU Layers slider at "Auto" (leftmost): --fit owns the whole
+ // layout, so the offload knobs (MoE, split, TP) don't apply.
+ const autoLayers = isManual && gpuLayers < 0;
+ // GPUs actually in use: the picked subset, or all visible when none picked.
+ const gpusInUse = selectedGpuIds ?? gpuDevices.map((d) => d.index);
+ // TP is off with fewer than 2 GPUs in use (single GPU, or the picker narrowed
+ // to one): tensor split is a no-op there and aborts on some archs. Mirrors the
+ // multi-GPU gate on the GPU picker / Split ratio. (Under Auto layers the whole
+ // TP control is hidden -- llama.cpp's --fit aborts under --split-mode tensor.)
+ const tpDisabled = gpusInUse.length <= 1;
+ // Manual gpu-layers ceiling = model layer count + 1 (else a safe fallback):
+ // llama.cpp counts the output layer as one more offloadable layer past the
+ // repeating blocks ("offloaded 33/33" needs -ngl 33 on a 32-block model), so
+ // the slider max must reach it or full offload is unreachable. While staging,
+ // use the staged model's layer count (read from its header).
+ const stagedLayerCount = pendingSelection?.layerCount ?? null;
+ const modelLayerCount = pendingIsGguf ? stagedLayerCount : ggufLayerCount;
+ const gpuLayersMax = modelLayerCount != null ? modelLayerCount + 1 : 256;
+ // MoE-offload slider: shown only for MoE models, capped at their MoE-layer
+ // count. While staging, use the staged model's count (read from its header);
+ // otherwise the loaded model's.
+ const stagedMoeLayerCount = pendingSelection?.moeLayerCount ?? null;
+ const moeLayersMax = pendingIsGguf
+ ? (stagedMoeLayerCount ?? 0)
+ : (moeLayerCount ?? 0);
+ const showMoeSlider = isManual && !autoLayers && moeLayersMax > 0;
+ // gpuLayers always counts; MoE only with an explicit layer count (see above).
+ const manualDirty =
+ isManual &&
+ (gpuLayers !== loadedGpuLayers ||
+ (!autoLayers && nCpuMoe !== (loadedNCpuMoe ?? 0)));
+ // GPU picker: only meaningful on multi-GPU, and only when the reported
+ // indices are physical (relative ordinals from a parent CUDA_VISIBLE_DEVICES
+ // mask can't be mapped back to pin a device). null = use all (auto).
+ const showGpuPicker =
+ isGguf &&
+ gpuDevices.length > 1 &&
+ gpuDevices.every((d) => d.physicalIndex);
+ const isGpuChecked = (index: number) =>
+ selectedGpuIds === null || selectedGpuIds.includes(index);
+ const toggleGpu = (index: number) => {
+ const all = gpuDevices.map((d) => d.index);
+ const current = selectedGpuIds ?? all;
+ const next = current.includes(index)
+ ? current.filter((i) => i !== index)
+ : [...current, index].sort((a, b) => a - b);
+ if (next.length === 0) return; // keep at least one GPU selected
+ setSelectedGpuIds(next.length === all.length ? null : next);
+ // The per-GPU split is positional, so any change to the set of GPUs in use
+ // invalidates it: drop it (the sliders fall back to the VRAM-weighted
+ // default). TP needs 2+ GPUs, so disable it when only one remains.
+ setSplitRatio(null);
+ if (next.length <= 1) {
+ setTensorParallel(false);
+ }
+ };
+ const gpuIdsKey = (ids: number[] | null) => (ids === null ? "auto" : ids.join(","));
+ const gpuIdsDirty = gpuIdsKey(selectedGpuIds) !== gpuIdsKey(loadedGpuIds);
+ // Per-GPU layer split (--tensor-split): manual + 2+ GPUs in use. One slider
+ // per GPU, each a layer count; together they sum to the GPU Layers total.
+ const showSplitRatio =
+ isManual && !autoLayers && showGpuPicker && gpusInUse.length > 1;
+ // The total the per-GPU counts sum to (the GPU Layers slider value); 0 under
+ // Auto, where the split is hidden. The devices behind the GPUs in use, for
+ // labels + the VRAM-weighted default.
+ const splitTotal = Math.max(0, Math.min(gpuLayers, gpuLayersMax));
+ const gpusInUseDevices = gpusInUse.map(
+ (i) => gpuDevices.find((d) => d.index === i) ?? null,
+ );
+ // Displayed per-GPU counts. splitRatio is a stable reference balance (only a
+ // slider edit changes it), rescaled to the current total; deriving rather than
+ // mutating it on GPU Layers changes keeps the balance intact when the total
+ // passes through low values or Auto. No saved split: free-VRAM-weighted default
+ // (llama.cpp's unset default splits by free VRAM, so the first edit starts from
+ // the default's placement, not a total-VRAM ratio that can land layers on a
+ // busy GPU). A genuine 0 (a full GPU) is a real weight, not missing data: the
+ // probe's no-data case degrades to the total server-side, and an all-zero list
+ // falls back to an even split in distributeByWeight. Not yet sent.
+ const splitCounts =
+ splitRatio && splitRatio.length === gpusInUse.length
+ ? distributeByWeight(splitTotal, splitRatio)
+ : distributeByWeight(
+ splitTotal,
+ gpusInUseDevices.map((d) => d?.memoryFreeGb ?? d?.memoryTotalGb ?? 1),
+ );
+ const setSplitCount = (k: number, v: number) =>
+ setSplitRatio(rebalanceSplit(splitTotal, splitCounts, k, v));
+ const splitRatioDirty =
+ isManual &&
+ !autoLayers &&
+ JSON.stringify(splitRatio ?? null) !== JSON.stringify(loadedSplitRatio ?? null);
+ // Auto-fit context (Manual + Auto layers): <= 0 means "Auto" (--fit sizes it);
+ // a positive value pins it. Surface the length --fit chose once it's loaded.
+ const fitCtxAuto = autoLayers && (customContextLength ?? 0) <= 0;
+ const loadedAutoLayers =
+ loadedGpuMemoryMode === "manual" && (loadedGpuLayers ?? GPU_LAYERS_AUTO) < 0;
+ const fitResolvedCtx =
+ fitCtxAuto && loadedAutoLayers ? ggufContextLength : null;
// A saved chat-template override is a reload-time setting too, so surface
// Apply for a template-only edit (otherwise it could never be applied).
const templateDirty = chatTemplateOverride !== loadedChatTemplateOverride;
const modelSettingsDirty =
- kvDirty || ctxDirty || specDirty || specDraftDirty || tpDirty || templateDirty;
+ kvDirty ||
+ ctxDirty ||
+ specDirty ||
+ specDraftDirty ||
+ tpDirty ||
+ gpuDirty ||
+ manualDirty ||
+ gpuIdsDirty ||
+ splitRatioDirty ||
+ templateDirty;
const [presetNameInput, setPresetNameInput] = useState(activePreset);
const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false);
const [systemPromptDraft, setSystemPromptDraft] = useState("");
@@ -980,7 +1140,64 @@ export function ChatSettingsPanel({
)}
{isGguf && (
<>
- {showContextControl && (
+ {showContextControl && (autoLayers ? (
+
+
+
+
+ Context Length
+
+
+ Auto: llama.cpp's --fit sizes the context to fit VRAM.
+ Set a length to pin it instead -- --fit then optimizes
+ GPU layer offload around it. The length --fit chose
+ shows here after loading.
+
+
+
{
+ setCustomContextLength(v > 0 ? v : null);
+ }}
+ ariaLabel="Context Length"
+ size={8}
+ disabled={modelControlsDisabled}
+ />
+
+
{
+ // Far-left snaps to Auto; otherwise to the nearest 1024.
+ if (v < 512) {
+ setCustomContextLength(null);
+ } else {
+ setCustomContextLength(Math.round(v / 1024) * 1024);
+ }
+ }}
+ className="panel-slider"
+ disabled={modelControlsDisabled}
+ />
+ {fitResolvedCtx != null && (
+
+ llama.cpp loaded {fitResolvedCtx.toLocaleString()} tokens.
+
+ )}
+
+ ) : (
@@ -1036,7 +1253,7 @@ export function ChatSettingsPanel({
)}
- )}
+ ))}
@@ -1191,6 +1408,163 @@ export function ChatSettingsPanel({
)}
>
)}
+ {gpuModeApplies && (
+
+
+
+ GPU Memory
+
+
+
+
+ Default: Unsloth
+ fits the model and context to your GPUs.
+
+
+ Manual: set GPU
+ Layers yourself. Leave it on Auto to let llama.cpp size
+ the context and offload overflow (including MoE experts)
+ to RAM.
+
+
+
+
+
+ {
+ setGpuMemoryMode(v as "auto" | "manual");
+ }}
+ // An in-flight staged load already snapshotted its
+ // settings, so edits here could not apply -- disable like
+ // the sibling context/KV/spec controls.
+ disabled={modelControlsDisabled}
+ >
+
+
+
+
+ Default
+ Manual
+
+
+
+
+ )}
+ {isManual && (
+ <>
+
+ Layers to keep on the GPU (--gpu-layers); the rest run
+ on CPU. Auto lets llama.cpp size the split (and the
+ context) to fit VRAM. At the maximum, the whole model
+ is on the GPU.
+ >
+ }
+ />
+ {showMoeSlider && (
+
+ Keep the experts of this many MoE layers on the CPU
+ (--n-cpu-moe) to save VRAM. 0 = all experts on the
+ GPU; at the maximum, all are on the CPU.
+ >
+ }
+ />
+ )}
+ {showSplitRatio && (
+
+
+
+ Layers per GPU
+
+
+ Splits GPU Layers across GPUs (--tensor-split).
+ Without Tensor Parallelism each value is the layer
+ count on that GPU; with it, every GPU holds a slice
+ of each layer, so the values are only a ratio.
+
+
+ {gpusInUseDevices.map((d, k) => (
+
setSplitCount(k, v)}
+ valueSize={6}
+ disabled={modelControlsDisabled}
+ />
+ ))}
+
+ )}
+ >
+ )}
+ {showGpuPicker && (
+
+
+
+ GPUs
+
+
+ Which GPUs this model may use. Unchecked GPUs are hidden
+ from llama.cpp (CUDA_VISIBLE_DEVICES, or
+ HIP_VISIBLE_DEVICES on ROCm). Leave all checked to use
+ every GPU.
+
+
+
+ {gpuDevices.map((d) => (
+
+
+ GPU {d.index}: {d.name}
+ {d.memoryTotalGb
+ ? ` · ${Math.round(d.memoryTotalGb)} GB`
+ : ""}
+
+ toggleGpu(d.index)}
+ data-test-id={`gpu-pick-${d.index}`}
+ disabled={modelControlsDisabled}
+ />
+
+ ))}
+
+
+ )}
+ {gpuModeApplies && !autoLayers && (
@@ -1206,10 +1580,11 @@ export function ChatSettingsPanel({
className="panel-switch shrink-0"
checked={tensorParallel}
onCheckedChange={setTensorParallel}
- disabled={modelControlsDisabled}
+ disabled={tpDisabled || modelControlsDisabled}
data-test-id="tensor-parallel-switch"
/>
+ )}
>
)}
{/* No persistent "enable custom code" toggle: it is consented per model
@@ -1228,14 +1603,21 @@ export function ChatSettingsPanel({
{Math.round((stagedDownloadFraction ?? 0) * 100)}%
)}
-
- setRemember(v === true)}
- />
- Remember settings next time
-
+ {/* GGUF picks only: a non-GGUF pick shows none of the load
+ knobs the blob captures, so there is nothing to remember. */}
+ {pendingIsGguf && (
+
+ setRemember(v === true)}
+ // The save/clear already ran in the Load click handler, so
+ // a mid-load toggle could not apply -- lock it like the knobs.
+ disabled={modelControlsDisabled}
+ />
+ Remember settings next time
+
+ )}
{stagedLoading ? (
// Mid-load: nothing to load or abandon until it settles, so disable.
) : null}
-
+ {/* The template override is a load-time knob too (applied on the next
+ reload) and the in-flight load already snapshotted it, so lock its
+ editors like the sibling controls -- a mid-load save would be
+ silently clobbered by the load response despite its toast. */}
+
)}
@@ -2086,7 +2477,7 @@ function BypassPermissionsToggle() {
);
}
-function ChatTemplateFields() {
+function ChatTemplateFields({ disabled = false }: { disabled?: boolean }) {
const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate);
const override = useChatRuntimeStore((s) => s.chatTemplateOverride);
const setOverride = useChatRuntimeStore((s) => s.setChatTemplateOverride);
@@ -2120,7 +2511,8 @@ function ChatTemplateFields() {
Chat Template
@@ -2131,7 +2523,8 @@ function ChatTemplateFields() {
setOverride(null)}
- className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white"
+ disabled={disabled}
+ className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white disabled:pointer-events-none disabled:opacity-50"
aria-label="Revert chat template"
>
Cancel
-
+ {/* Also locked mid-load: an autoLoad can start with this dialog
+ already open, and a save then would be silently clobbered. */}
+
Save
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
index 10e0904e4f..3003b52230 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
@@ -29,9 +29,14 @@ import {
} from "../api/chat-api";
import { formatEta, formatRate } from "../utils/format-transfer";
import {
+ GPU_LAYERS_AUTO,
isLocalModelPath,
+ loadedGpuMemoryFields,
+ loadedGpuMemoryFieldsUnlessStaged,
pendingSelectionMatches,
+ persistGpuMemoryModeOnLoad,
readPersistedSpeculativeType,
+ reconcilePersistedGpuIds,
resolveToolsEnabledOnLoad,
saveSpeculativeType,
useChatRuntimeStore,
@@ -46,9 +51,12 @@ import {
} from "../lib/apply-inference-status-to-store";
import {
mergeBackendRecommendedInference,
+ resolveFitMaxSeqLength,
resolveLoadMaxSeqLength,
+ resolveManualAutoCtxPin,
} from "../presets/preset-policy";
import { recordLastLocalModelLoad } from "../utils/last-local-model-load";
+import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info";
import {
isMultimodalResponse,
} from "../types/api";
@@ -291,9 +299,12 @@ async function syncInferenceStatusToStore(options?: {
if (statusRes.active_model && !isExternalSelectionActive) {
const checkpointId = resolveInferenceCheckpointId(statusRes);
if (checkpointId) {
+ const previousGgufVariant =
+ useChatRuntimeStore.getState().activeGgufVariant;
setCheckpoint(checkpointId, statusRes.gguf_variant);
applyActiveModelStatusToStore(statusRes, {
previousCheckpoint: selectedCheckpoint,
+ previousGgufVariant,
});
// setModels(listRes...) above used catalog data, which omits audio
// capability. Re-apply live status so attach gates survive a refresh.
@@ -511,7 +522,11 @@ export function useChatModelRuntime() {
typeof selection === "string" ? false : selection.isDownloaded ?? false;
const model = models.find((entry) => entry.id === modelId);
const lora = loras.find((entry) => entry.id === modelId);
- const isGguf = explicitIsGguf ?? model?.isGguf ?? false;
+ // A native path-token selection is a local GGUF by construction (the
+ // native model intents only grant .gguf files), but its id is a display
+ // label that need not end in ".gguf" -- without this, Manual + Auto
+ // layers would pin the UI context instead of letting --fit size it.
+ const isGguf = explicitIsGguf ?? model?.isGguf ?? nativePathToken != null;
const loraIsAdapter = lora?.exportType === "lora";
const isLora =
explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false;
@@ -578,18 +593,27 @@ export function useChatModelRuntime() {
let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false;
let approvedRemoteCodeFingerprint: string | null = null;
const maxSeqLength = stateBeforeUnload.params.maxSeqLength;
+ const previousActiveNativePathToken =
+ stateBeforeUnload.activeNativePathToken;
const previousIsGguf =
previousModel?.isGguf === true
|| previousVariant != null
+ || previousActiveNativePathToken != null
|| (previousCheckpoint?.toLowerCase().endsWith(".gguf") ?? false);
- const rollbackMaxSeqLength = previousIsGguf
- ? (stateBeforeUnload.ggufContextLength ?? 0)
- : maxSeqLength;
+ // Respect the rolled-back model's auto-layers mode: a Manual+Auto model
+ // with an unpinned (auto) context must reload with 0 (so --fit
+ // re-auto-sizes), not the positive context it happened to pick (which
+ // the backend would treat as a pin).
+ const rollbackMaxSeqLength = resolveFitMaxSeqLength(
+ previousIsGguf,
+ stateBeforeUnload.loadedGpuMemoryMode ?? "auto",
+ stateBeforeUnload.loadedGpuLayers ?? GPU_LAYERS_AUTO,
+ stateBeforeUnload.loadedCustomContextLength,
+ previousIsGguf ? (stateBeforeUnload.ggufContextLength ?? 0) : maxSeqLength,
+ );
const hfToken = stateBeforeUnload.hfToken || null;
const previousModelRequiresTrustRemoteCode =
stateBeforeUnload.modelRequiresTrustRemoteCode;
- const previousActiveNativePathToken =
- stateBeforeUnload.activeNativePathToken;
// Snapshot the load settings at click time, before the awaits below
// (validation, the trust dialog, unload). For a staged Load these knobs
// stay editable and a sheet-close revert (abandonStagedModel) can fire
@@ -598,11 +622,29 @@ export function useChatModelRuntime() {
// updates this snapshot in lock-step so non-staged loads are unchanged.
const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride;
const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype;
- const loadCustomContextLength = stateBeforeUnload.customContextLength;
+ // gpuMemoryMode is a standing preference (kept across a model switch);
+ // the rest are per-model knobs the reset below clears, so they are
+ // re-baselined there in lock-step with the store.
+ let loadCustomContextLength = stateBeforeUnload.customContextLength;
const loadGgufContextLength = stateBeforeUnload.ggufContextLength;
const loadTensorParallel = stateBeforeUnload.tensorParallel;
const loadActivePresetSource = stateBeforeUnload.activePresetSource;
const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant;
+ const loadGpuMemoryMode = stateBeforeUnload.gpuMemoryMode;
+ let loadGpuLayers = stateBeforeUnload.gpuLayers;
+ let loadNCpuMoe = stateBeforeUnload.nCpuMoe;
+ let loadSplitRatio = stateBeforeUnload.splitRatio;
+ // Reconcile the persisted pick against the GPUs present now, so a stale
+ // cross-host / now-hidden pick is dropped before /load rather than
+ // rejected there. Warm the device cache first: load-on-selection can
+ // run before any GPU hook mounted, and a cold cache would pass the
+ // pick through unvalidated. validateGpuIds derives from this too.
+ if (stateBeforeUnload.selectedGpuIds != null) {
+ await ensureGpuDeviceCache();
+ }
+ let loadSelectedGpuIds = reconcilePersistedGpuIds(
+ stateBeforeUnload.selectedGpuIds,
+ );
let loadSpeculativeType = stateBeforeUnload.speculativeType;
let loadSpecDraftNMax = stateBeforeUnload.specDraftNMax;
try {
@@ -615,16 +657,47 @@ export function useChatModelRuntime() {
// context can exceed maxSeqLength, so sizing on raw maxSeqLength could
// pass, unload, then have /load refuse it. Uses the click-time
// snapshot (same values loadModel uses below), so the two agree.
- const validateMaxSeqLength = resolveLoadMaxSeqLength({
- modelId,
- ggufVariant,
- customContextLength: loadCustomContextLength,
- ggufContextLength: loadGgufContextLength,
- currentCheckpoint,
- activeGgufVariant: loadActiveGgufVariant,
- maxSeqLength,
- presetSource: loadActivePresetSource,
- });
+ // Mirror what /load does on a cross-model switch: the reset below
+ // clears the per-model Auto-layers context pin + GPU pick, and
+ // Manual+Auto sizes context through resolveFitMaxSeqLength.
+ // gpuMemoryMode is a standing preference, kept across the switch.
+ // A same-repo quant switch (same checkpoint, different gguf_variant)
+ // is a different model for per-model knobs: the pinned context,
+ // gpuLayers, GPU pick, and MoE offload are scoped per variant, so
+ // treat a variant change like a model switch and re-baseline them.
+ const switchingModelOrVariant =
+ currentCheckpoint !== modelId ||
+ (loadActiveGgufVariant ?? null) !== (ggufVariant ?? null);
+ const resetsPerModelSettings = Boolean(
+ currentCheckpoint && switchingModelOrVariant && !keepSpeculative,
+ );
+ const validateCustomContextLength = resetsPerModelSettings
+ ? null
+ : loadCustomContextLength;
+ const validateGpuIds = resetsPerModelSettings
+ ? null
+ : loadSelectedGpuIds;
+ // The reset below re-baselines gpuLayers to Auto; mirror it here.
+ const validateGpuLayers = resetsPerModelSettings
+ ? GPU_LAYERS_AUTO
+ : loadGpuLayers;
+ const validateMaxSeqLength = resolveFitMaxSeqLength(
+ isGguf,
+ loadGpuMemoryMode,
+ validateGpuLayers,
+ validateCustomContextLength,
+ resolveLoadMaxSeqLength({
+ modelId,
+ ggufVariant,
+ isGguf,
+ customContextLength: validateCustomContextLength,
+ ggufContextLength: loadGgufContextLength,
+ currentCheckpoint,
+ activeGgufVariant: loadActiveGgufVariant,
+ maxSeqLength,
+ presetSource: loadActivePresetSource,
+ }),
+ );
const validation = await validateModel({
model_path: modelId,
nativePathLease: validateNativePathLease,
@@ -633,6 +706,8 @@ export function useChatModelRuntime() {
load_in_4bit: true,
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
+ gpu_ids: validateGpuIds ?? undefined,
+ ...(isGguf ? { gpu_memory_mode: loadGpuMemoryMode } : {}),
});
// Upgrade consent runs before the security dialogs; Accept installs and the load continues.
if (validation.requires_transformers_upgrade) {
@@ -697,18 +772,52 @@ export function useChatModelRuntime() {
// keepSpeculative skips this for a staged Load: the user picked the
// mode for this model on the sidebar, so honor it (the backend still
// falls back at runtime if the model has no MTP head).
- if (currentCheckpoint && currentCheckpoint !== modelId && !keepSpeculative) {
+ if (resetsPerModelSettings) {
const persistedSpeculativeType = readPersistedSpeculativeType();
useChatRuntimeStore.setState({
speculativeType: persistedSpeculativeType,
loadedSpeculativeType: persistedSpeculativeType,
specDraftNMax: null,
loadedSpecDraftNMax: null,
+ // Per-model GPU knobs must not follow onto a different model
+ // (gpuMemoryMode is a standing preference and is kept).
+ selectedGpuIds: null,
+ gpuLayers: GPU_LAYERS_AUTO,
+ nCpuMoe: 0,
+ splitRatio: null,
+ // A Manual+Auto context pin is per-model; clear it so a different
+ // model loads at Auto/native, not the previous model's pin.
+ customContextLength: null,
});
loadSpeculativeType = persistedSpeculativeType;
loadSpecDraftNMax = null;
+ // Keep the click-time snapshot in lock-step with the store reset so
+ // the load below sizes against the cleared per-model knobs, not the
+ // previous model's (gpuMemoryMode is standing, so left as captured).
+ loadCustomContextLength = null;
+ loadSelectedGpuIds = null;
+ loadGpuLayers = GPU_LAYERS_AUTO;
+ loadNCpuMoe = 0;
+ loadSplitRatio = null;
}
+ // Pinning layers on the SAME model keeps the currently resolved
+ // context: with no explicit pin, a manual+pinned reload would send 0,
+ // which the backend's --fit off branch treats as the NATIVE context --
+ // far larger than the sheet shows when the load was fit-sized (Default
+ // or Manual + Auto layers may auto-reduce context to fit VRAM), a
+ // likely OOM. ggufContextLength is that resolved value; a model already
+ // at native reloads unchanged, so this is safe for any prior mode.
+ if (
+ isGguf &&
+ !switchingModelOrVariant &&
+ loadGpuMemoryMode === "manual" &&
+ loadGpuLayers >= 0 &&
+ loadCustomContextLength == null &&
+ (loadGgufContextLength ?? 0) > 0
+ ) {
+ loadCustomContextLength = loadGgufContextLength;
+ }
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
modelId,
ggufVariant,
@@ -720,13 +829,20 @@ export function useChatModelRuntime() {
maxSeqLength,
presetSource: loadActivePresetSource,
});
+ const loadMaxSeqLength = resolveFitMaxSeqLength(
+ isGguf,
+ loadGpuMemoryMode,
+ loadGpuLayers,
+ loadCustomContextLength,
+ effectiveMaxSeqLength,
+ );
const effectiveChatTemplateOverride =
loadChatTemplateOverride?.trim() ? loadChatTemplateOverride : null;
const loadResponse = await loadModel({
model_path: modelId,
nativePathLease: loadNativePathLease,
hf_token: hfToken,
- max_seq_length: effectiveMaxSeqLength,
+ max_seq_length: loadMaxSeqLength,
load_in_4bit: true,
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
@@ -737,6 +853,11 @@ export function useChatModelRuntime() {
speculative_type: loadSpeculativeType,
spec_draft_n_max: loadSpecDraftNMax,
tensor_parallel: loadTensorParallel,
+ gpu_memory_mode: loadGpuMemoryMode,
+ gpu_layers: loadGpuLayers,
+ n_cpu_moe: loadNCpuMoe,
+ tensor_split: loadSplitRatio ?? undefined,
+ gpu_ids: loadSelectedGpuIds ?? undefined,
});
// If cancelled while loading, don't update UI to show
@@ -747,6 +868,9 @@ export function useChatModelRuntime() {
// preference now (the requested intent, not the resolved echo;
// saveSpeculativeType keeps only the universal auto/ngram/off).
saveSpeculativeType(loadSpeculativeType);
+ // Persist the GPU Memory mode only on a successful load (not on
+ // dropdown change), so an abandoned selection doesn't stick.
+ persistGpuMemoryModeOnLoad(loadResponse, loadGpuMemoryMode);
const currentParams = useChatRuntimeStore.getState().params;
setParams(
@@ -782,9 +906,13 @@ export function useChatModelRuntime() {
const reportedNativeCtx = loadResponse.is_gguf
? (loadResponse.native_context_length ?? null)
: null;
- // A successful reload has applied settings, so clear pending custom
- // context state and display the backend-reported effective context.
- const keepCustomCtx = null;
+ // Keep an explicit Manual+Auto context pin (so a later Apply doesn't
+ // revert it to Auto); other cases baseline on ggufContextLength.
+ const keepCustomCtx = resolveManualAutoCtxPin(
+ loadGpuMemoryMode,
+ loadGpuLayers,
+ loadCustomContextLength,
+ );
const reasoningAlwaysOn = loadResponse.reasoning_always_on ?? false;
const reasoningStyle = loadResponse.reasoning_style ?? "enable_thinking";
const supportsReasoning = loadResponse.supports_reasoning ?? false;
@@ -837,11 +965,13 @@ export function useChatModelRuntime() {
loadedKvCacheDtype: loadedKv,
tensorParallel: loadedTp,
loadedTensorParallel: loadedTp,
+ ...loadedGpuMemoryFields(loadResponse),
speculativeType: loadedSpec,
loadedSpeculativeType: loadedSpec,
specDraftNMax: loadResponse.spec_draft_n_max ?? null,
loadedSpecDraftNMax: loadResponse.spec_draft_n_max ?? null,
customContextLength: keepCustomCtx,
+ loadedCustomContextLength: keepCustomCtx,
defaultChatTemplate: loadResponse.chat_template ?? null,
chatTemplateOverride: effectiveChatTemplateOverride,
loadedChatTemplateOverride: effectiveChatTemplateOverride,
@@ -938,7 +1068,7 @@ export function useChatModelRuntime() {
}
}
try {
- await loadModel({
+ const rollbackResponse = await loadModel({
model_path: previousCheckpoint,
nativePathLease: rollbackNativePathLease,
hf_token: hfToken,
@@ -951,14 +1081,51 @@ export function useChatModelRuntime() {
// Resend the previous model's pinned approval so restoring it is not re-blocked.
approved_remote_code_fingerprint:
approvedRemoteCodeFingerprints.get(previousCheckpoint) ?? null,
+ chat_template_override:
+ stateBeforeUnload.loadedChatTemplateOverride,
+ cache_type_kv: stateBeforeUnload.loadedKvCacheDtype,
+ speculative_type:
+ stateBeforeUnload.loadedSpeculativeType,
+ spec_draft_n_max:
+ stateBeforeUnload.loadedSpecDraftNMax,
// Restore the previous model in the split mode it was running,
// not the default layer split.
tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false,
+ gpu_memory_mode: stateBeforeUnload.loadedGpuMemoryMode ?? "auto",
+ gpu_layers: stateBeforeUnload.loadedGpuLayers ?? -1,
+ n_cpu_moe: stateBeforeUnload.loadedNCpuMoe ?? 0,
+ tensor_split: stateBeforeUnload.loadedSplitRatio ?? undefined,
+ gpu_ids: stateBeforeUnload.loadedGpuIds ?? undefined,
});
+ const rollbackSpeculativeType = normalizeSpeculativeType(
+ rollbackResponse.speculative_type,
+ );
useChatRuntimeStore.setState({
activeNativePathToken: previousActiveNativePathToken ?? null,
- loadedSpeculativeType: null,
- loadedSpecDraftNMax: null,
+ loadedSpeculativeType: rollbackSpeculativeType,
+ loadedSpecDraftNMax:
+ rollbackResponse.spec_draft_n_max ?? null,
+ loadedKvCacheDtype: rollbackResponse.cache_type_kv ?? null,
+ loadedChatTemplateOverride:
+ stateBeforeUnload.loadedChatTemplateOverride,
+ // Re-baseline the GPU knobs from the rolled-back load's own
+ // response (the shared seeding every load path uses): the
+ // refresh() below can't do it, since the status reseed is
+ // gated off while modelLoading is still true. A failed staged
+ // Load stays staged for retry, so the staged hold applies.
+ ...loadedGpuMemoryFieldsUnlessStaged(rollbackResponse, {
+ tensorParallel: rollbackResponse.tensor_parallel ?? false,
+ loadedTensorParallel:
+ rollbackResponse.tensor_parallel ?? false,
+ // refresh() is held while modelLoading remains true, so
+ // restore the rolled-back model's context pin directly.
+ customContextLength:
+ stateBeforeUnload.loadedCustomContextLength,
+ }),
+ loadedTensorParallel:
+ rollbackResponse.tensor_parallel ?? false,
+ loadedCustomContextLength:
+ stateBeforeUnload.loadedCustomContextLength,
});
await refresh();
} catch {
diff --git a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
index d8076c720b..a3e7a2d264 100644
--- a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
+++ b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
@@ -7,7 +7,7 @@ import { useRepoDownload } from "@/features/hub/download-manager/use-repo-downlo
import type { DownloadJob } from "@/features/hub/download-manager/use-repo-download";
import { useLatestRef } from "@/features/hub/hooks/use-latest-ref";
-import { fetchGgufContextLength } from "../api/chat-api";
+import { fetchGgufStagedMetadata } from "../api/chat-api";
import {
isPendingGguf,
pendingSelectionMatches,
@@ -46,8 +46,16 @@ export function useStagedModelPreparation(opts?: {
const pendingDownloaded = useChatRuntimeStore(
(s) => s.pendingSelection?.isDownloaded ?? false,
);
- const pendingHasContext = useChatRuntimeStore(
- (s) => s.pendingSelection?.contextLength != null,
+ // "Already probed" must key off layerCount / moeLayerCount, which only the
+ // full header probe fills (it sets all three together, so either is a
+ // reliable marker). contextLength alone can be list-seeded from
+ // /gguf-variants, which returns no layer/MoE counts -- treating it as
+ // complete would skip the probe and leave the GPU Layers slider at its 256
+ // fallback and the MoE slider hidden until the model loads.
+ const pendingHasMetadata = useChatRuntimeStore(
+ (s) =>
+ s.pendingSelection?.layerCount != null ||
+ s.pendingSelection?.moeLayerCount != null,
);
const setPendingSelection = useChatRuntimeStore((s) => s.setPendingSelection);
const onAutoLoadRef = useLatestRef(opts?.onAutoLoad);
@@ -69,25 +77,31 @@ export function useStagedModelPreparation(opts?: {
if (!current?.id || !isPendingGguf(current)) return;
const { id, ggufVariant, nativePathToken } = current;
try {
- const contextLength = await fetchGgufContextLength({
- model_path: id,
- gguf_variant: ggufVariant,
- hf_token: useChatRuntimeStore.getState().hfToken || null,
- nativePathToken,
- });
+ const { contextLength, layerCount, moeLayerCount } =
+ await fetchGgufStagedMetadata({
+ model_path: id,
+ gguf_variant: ggufVariant,
+ hf_token: useChatRuntimeStore.getState().hfToken || null,
+ nativePathToken,
+ });
// Apply only if the same model is still staged (the user may have switched
// picks or loaded/cancelled while the request was in flight).
const latest = useChatRuntimeStore.getState().pendingSelection;
if (
latest &&
- contextLength != null &&
- pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken })
+ pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken }) &&
+ (contextLength != null || layerCount != null || moeLayerCount != null)
) {
- setPendingSelection({ ...latest, contextLength });
+ setPendingSelection({
+ ...latest,
+ contextLength,
+ layerCount,
+ moeLayerCount,
+ });
}
} catch {
- // Leave contextLength null: the context slider stays hidden and the user
- // can still load (context fills in from the load response afterwards).
+ // Leave metadata null: the context/MoE sliders stay hidden and the user
+ // can still load (they fill in from the load response afterwards).
}
}, [setPendingSelection]);
@@ -125,7 +139,7 @@ export function useStagedModelPreparation(opts?: {
if (
!pendingId ||
(!pendingIsGguf && !pendingIsHubRepo) ||
- pendingHasContext
+ pendingHasMetadata
) {
return;
}
@@ -146,7 +160,7 @@ export function useStagedModelPreparation(opts?: {
pendingIsGguf,
pendingIsHubRepo,
pendingDownloaded,
- pendingHasContext,
+ pendingHasMetadata,
startDownloadRef,
fetchMetadataRef,
]);
diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
index a4b5f848e2..69bb38bbbe 100644
--- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
+++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
@@ -2,13 +2,17 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getInferenceStatus } from "../api/chat-api";
-import { mergeBackendRecommendedInference } from "../presets/preset-policy";
+import {
+ mergeBackendRecommendedInference,
+ resolveManualAutoCtxPin,
+} from "../presets/preset-policy";
import { clampReasoningEffortToLevels } from "../provider-capabilities";
import {
CHAT_REASONING_ENABLED_KEY,
type ReasoningEffort,
type ReasoningStyle,
loadOptionalBool,
+ loadedGpuMemoryFields,
resolveToolsEnabledOnLoad,
useChatRuntimeStore,
} from "../stores/chat-runtime-store";
@@ -20,6 +24,10 @@ import type { ChatModelSummary } from "../types/runtime";
type LocalReasoningEffort = Extract
;
+function sameArray(a: T[] | null, b: T[] | null): boolean {
+ return JSON.stringify(a) === JSON.stringify(b);
+}
+
// Canonicalises backend / persisted speculative mode values onto the UI modes.
export function normalizeSpeculativeType(
v: string | null | undefined,
@@ -119,6 +127,10 @@ function ensureActiveModelInStoreList(
export type ApplyInferenceStatusOptions = {
previousCheckpoint?: string;
+ /** activeGgufVariant BEFORE the caller's setCheckpoint synced it to the
+ * status -- without it a variant-only switch underneath the tab reads as
+ * steady state and the hydration reseed keeps the old quant's baselines. */
+ previousGgufVariant?: string | null;
};
/** Mirror refresh() hydration so adopted CLI models get reasoning/tools flags. */
@@ -144,9 +156,13 @@ export function applyActiveModelStatusToStore(
);
}
+ const previousGgufVariant =
+ options.previousGgufVariant !== undefined
+ ? options.previousGgufVariant
+ : store.activeGgufVariant;
const hydratingExistingModel =
previousCheckpoint !== checkpointId ||
- store.activeGgufVariant !== (status.gguf_variant ?? null);
+ previousGgufVariant !== (status.gguf_variant ?? null);
const supportsReasoning = status.supports_reasoning ?? false;
const reasoningAlwaysOn = status.reasoning_always_on ?? false;
const reasoningStyle = status.reasoning_style ?? "enable_thinking";
@@ -185,6 +201,66 @@ export function applyActiveModelStatusToStore(
// While a load is in flight, performLoad owns the load params. Seeding them
// from a stale poll here would clobber the values the load dialog just set.
const seedLoadParams = !prevState.modelLoading;
+ // A Manual + Auto-layers load sent its positive context pin as max_seq_length,
+ // and status only exposes the RESOLVED context; re-seed the pin from the
+ // requested value (parity with the load paths' keepCustomCtx). Baselines
+ // unconditionally: anything but an applicable pin is null, so a previous
+ // model's pin can't survive a model change underneath and reload at the old length.
+ const gpuPin = status.is_gguf
+ ? resolveManualAutoCtxPin(
+ status.gpu_memory_mode ?? "auto",
+ status.gpu_layers ?? -1,
+ status.requested_context_length ?? null,
+ )
+ : null;
+ const incomingGpuMode = status.is_gguf
+ ? (status.gpu_memory_mode ?? "auto")
+ : null;
+ const incomingGpuLayers =
+ incomingGpuMode === "manual" ? (status.gpu_layers ?? null) : null;
+ const incomingNCpuMoe =
+ incomingGpuMode === "manual" ? (status.n_cpu_moe ?? null) : null;
+ const incomingSplit =
+ incomingGpuMode === "manual" ? (status.tensor_split ?? null) : null;
+ const incomingGpuIds = status.is_gguf ? (status.gpu_ids ?? null) : null;
+ const gpuStatusChanged =
+ prevState.loadedGpuMemoryMode !== incomingGpuMode ||
+ prevState.loadedGpuLayers !== incomingGpuLayers ||
+ prevState.loadedNCpuMoe !== incomingNCpuMoe ||
+ !sameArray(prevState.loadedSplitRatio, incomingSplit) ||
+ !sameArray(prevState.loadedGpuIds, incomingGpuIds) ||
+ prevState.loadedCustomContextLength !== gpuPin;
+ const gpuMemoryEditsPending =
+ (prevState.loadedGpuMemoryMode !== null &&
+ prevState.gpuMemoryMode !== prevState.loadedGpuMemoryMode) ||
+ (prevState.loadedGpuMemoryMode === "manual" &&
+ (prevState.gpuLayers !== prevState.loadedGpuLayers ||
+ prevState.nCpuMoe !== prevState.loadedNCpuMoe ||
+ !sameArray(prevState.splitRatio, prevState.loadedSplitRatio))) ||
+ prevState.customContextLength !== prevState.loadedCustomContextLength;
+ const gpuIdsEditPending = !sameArray(
+ prevState.selectedGpuIds,
+ prevState.loadedGpuIds,
+ );
+ const incomingGpuFields = loadedGpuMemoryFields(status);
+ // A same-model reload from another client advances every loaded baseline.
+ // Preserve each editable group only when this tab has an unapplied change.
+ const preserveSameModelEdits = gpuStatusChanged && !hydratingExistingModel;
+ const gpuStatusFields = {
+ ...incomingGpuFields,
+ customContextLength: gpuPin,
+ loadedCustomContextLength: gpuPin,
+ ...(preserveSameModelEdits &&
+ gpuMemoryEditsPending && {
+ gpuMemoryMode: prevState.gpuMemoryMode,
+ gpuLayers: prevState.gpuLayers,
+ nCpuMoe: prevState.nCpuMoe,
+ splitRatio: prevState.splitRatio,
+ customContextLength: prevState.customContextLength,
+ }),
+ ...(preserveSameModelEdits &&
+ gpuIdsEditPending && { selectedGpuIds: prevState.selectedGpuIds }),
+ };
useChatRuntimeStore.setState({
supportsReasoning,
@@ -215,30 +291,51 @@ export function applyActiveModelStatusToStore(
loadedIsMultimodal: isMultimodalResponse(status),
loadedIsDiffusion: status.is_diffusion ?? false,
specFallbackReason: status.spec_fallback_reason ?? null,
+ // The spec / KV seeds share the GPU-fields reseed mechanism below: a
+ // non-GGUF status leaves their loaded baselines null, so the "unseeded"
+ // guard re-fires every refresh -- hold them too while a staged pick's
+ // settings are being edited, or the refresh resets the staged edit.
+ // hydratingExistingModel reopens every load-param seed: when the active
+ // model changed underneath this tab (auto-switch, another client), the
+ // old model's baselines are stale and must adopt the new status.
...(seedLoadParams &&
- prevState.loadedSpeculativeType === null && {
+ prevState.pendingSelection == null &&
+ (prevState.loadedSpeculativeType === null || hydratingExistingModel) && {
speculativeType: currentSpecType,
loadedSpeculativeType: currentSpecType,
}),
...(seedLoadParams &&
+ prevState.pendingSelection == null &&
status.spec_draft_n_max !== undefined &&
- prevState.loadedSpecDraftNMax === null &&
- prevState.specDraftNMax === null && {
+ (hydratingExistingModel ||
+ (prevState.loadedSpecDraftNMax === null &&
+ prevState.specDraftNMax === null)) && {
specDraftNMax: status.spec_draft_n_max ?? null,
loadedSpecDraftNMax: status.spec_draft_n_max ?? null,
}),
...(seedLoadParams &&
+ prevState.pendingSelection == null &&
status.cache_type_kv !== undefined &&
- prevState.loadedKvCacheDtype === null && {
+ (prevState.loadedKvCacheDtype === null || hydratingExistingModel) && {
kvCacheDtype: status.cache_type_kv,
loadedKvCacheDtype: status.cache_type_kv,
}),
...(seedLoadParams &&
+ prevState.pendingSelection == null &&
status.tensor_parallel !== undefined &&
- prevState.loadedTensorParallel === null && {
+ (prevState.loadedTensorParallel === null || hydratingExistingModel) && {
tensorParallel: status.tensor_parallel,
loadedTensorParallel: status.tensor_parallel,
}),
+ // Re-seed on first hydration, model/variant changes, or a same-model backend
+ // placement change. gpuStatusFields preserves dirty local edits in the last
+ // case while advancing their loaded baselines.
+ ...(seedLoadParams &&
+ prevState.pendingSelection == null &&
+ (prevState.loadedGpuMemoryMode === null ||
+ hydratingExistingModel ||
+ gpuStatusChanged) &&
+ gpuStatusFields),
...(status.chat_template_override !== undefined &&
prevState.loadedChatTemplateOverride === null &&
prevState.chatTemplateOverride === null && {
@@ -298,7 +395,11 @@ export async function tryAdoptServerActiveModel(): Promise {
if (previousCheckpoint) {
return true;
}
+ const previousGgufVariant = useChatRuntimeStore.getState().activeGgufVariant;
store.setCheckpoint(checkpointId, status.gguf_variant);
- applyActiveModelStatusToStore(status, { previousCheckpoint });
+ applyActiveModelStatusToStore(status, {
+ previousCheckpoint,
+ previousGgufVariant,
+ });
return true;
}
diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts
index f96ee91f1b..23d79a35e1 100644
--- a/studio/frontend/src/features/chat/presets/preset-policy.ts
+++ b/studio/frontend/src/features/chat/presets/preset-policy.ts
@@ -339,3 +339,34 @@ export function resolveLoadMaxSeqLength({
}
return maxSeqLength;
}
+
+/**
+ * Adjust a resolved max-seq-length for the GPU Memory mode. Under Manual + Auto
+ * layers (GGUF, gpuLayers < 0) llama.cpp's --fit owns context sizing, so send 0
+ * (the backend omits -c) unless the user pinned a length; every other case keeps
+ * the resolved fallback. Shared by every GGUF load path so they can't drift.
+ */
+export function resolveFitMaxSeqLength(
+ isGguf: boolean | null | undefined,
+ gpuMemoryMode: "auto" | "manual",
+ gpuLayers: number,
+ customContextLength: number | null,
+ fallback: number,
+): number {
+ if (!isGguf || gpuMemoryMode !== "manual" || gpuLayers >= 0) return fallback;
+ return customContextLength && customContextLength > 0 ? customContextLength : 0;
+}
+
+// A Manual + Auto-layers load sends its positive context pin as max_seq_length;
+// keep it across a status reseed/Apply so the model isn't reverted to auto-fit
+// sizing. Anything else (Auto mode, pinned layers, no pin) baselines to null.
+// The caller keeps its own isGguf/targetIsGguf guard inline.
+export function resolveManualAutoCtxPin(
+ gpuMemoryMode: "auto" | "manual",
+ gpuLayers: number,
+ customContextLength: number | null,
+): number | null {
+ return gpuMemoryMode === "manual" && gpuLayers < 0 && (customContextLength ?? 0) > 0
+ ? customContextLength
+ : null;
+}
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index a0813fe27b..31e50ee60c 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -84,6 +84,8 @@ import {
useTransformersUpgradeDialogStore,
} from "@/features/transformers-upgrade";
import { loadModel, validateModel } from "./api/chat-api";
+import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "./presets/preset-policy";
+import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info";
import {
parseExternalModelId,
providerTypeSupportsVision,
@@ -95,8 +97,11 @@ import {
usePlusMenuPrefsStore,
} from "./stores/plus-menu-prefs-store";
import {
+ loadedGpuMemoryFieldsUnlessStaged,
type ReasoningEffort,
+ reconcilePersistedGpuIds,
resolveLoadedSpeculativeSettings,
+ persistGpuMemoryModeOnLoad,
resolveSpeculativeSettingsForLoad,
saveSpeculativeType,
useChatRuntimeStore,
@@ -1037,10 +1042,32 @@ export function SharedComposer({
return parts[parts.length - 1] || id;
}
+ // Warm the device cache before the snapshot below reconciles the GPU
+ // pick: on a cold cache the reconcile passes a stale pick through.
+ if (store.selectedGpuIds != null) {
+ await ensureGpuDeviceCache();
+ }
+ // The GPU/offload knobs both compare loads must use, snapshotted at Send.
+ // ensureModelLoaded runs sequentially and the first load's response echo
+ // (loadedGpuMemoryFields) rewrites the live store -- a non-GGUF or Auto
+ // first model resets gpuLayers/nCpuMoe/split/pick to defaults -- so
+ // reading the store per load would hand model 2 the first model's echoed
+ // defaults instead of the settings the user pressed Send with.
+ const compareLoadKnobs = {
+ gpuMemoryMode: store.gpuMemoryMode,
+ gpuLayers: store.gpuLayers,
+ nCpuMoe: store.nCpuMoe,
+ splitRatio: store.splitRatio,
+ // Reconcile the pick against the GPUs present now, like the model-switch
+ // path: an early remember-restore can hold a stale cross-host pick that
+ // /load would reject (the device cache is populated by send time).
+ selectedGpuIds: reconcilePersistedGpuIds(store.selectedGpuIds),
+ tensorParallel: store.tensorParallel,
+ customContextLength: store.customContextLength,
+ };
// Set when an accepted transformers install unloaded the active model
// server-side; a later failure must then clear the stale checkpoint.
let upgradeUnloadedActive = false;
-
// Helper: load a model and update store checkpoint
async function ensureModelLoaded(
sel: CompareModelSelection,
@@ -1057,15 +1084,35 @@ export function SharedComposer({
if (isAlreadyActive) {
return "ready";
}
+ const targetIsGguf =
+ sel.id.toLowerCase().endsWith(".gguf") || sel.ggufVariant != null;
+ // Size validation exactly as the load below, so the training-guard
+ // preflight checks the footprint that actually loads (under Manual + Auto
+ // layers the load sends 0 / the pinned context, not raw maxSeqLength).
+ const compareMaxSeqLength = resolveFitMaxSeqLength(
+ targetIsGguf,
+ compareLoadKnobs.gpuMemoryMode,
+ compareLoadKnobs.gpuLayers,
+ compareLoadKnobs.customContextLength,
+ maxSeqLength,
+ );
const validation = await validateModel({
model_path: sel.id,
hf_token: currentStore.hfToken || null,
- max_seq_length: maxSeqLength,
+ max_seq_length: compareMaxSeqLength,
load_in_4bit: true,
is_lora: sel.isLora,
gguf_variant: sel.ggufVariant ?? null,
trust_remote_code: loadTrustRemoteCode,
chat_template_override: effectiveChatTemplateOverride,
+ // Scope the validate to the picked GPUs. GGUF-only, like the load
+ // below: a non-GGUF target must not inherit a hidden GGUF GPU pick.
+ ...(targetIsGguf
+ ? {
+ gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined,
+ gpu_memory_mode: compareLoadKnobs.gpuMemoryMode,
+ }
+ : {}),
});
// Upgrade dialog first (mirrors the primary load path).
if (validation.requires_transformers_upgrade) {
@@ -1114,7 +1161,7 @@ export function SharedComposer({
const resp = await loadModel({
model_path: sel.id,
hf_token: useChatRuntimeStore.getState().hfToken || null,
- max_seq_length: maxSeqLength,
+ max_seq_length: compareMaxSeqLength,
load_in_4bit: true,
is_lora: sel.isLora,
gguf_variant: sel.ggufVariant ?? null,
@@ -1123,10 +1170,25 @@ export function SharedComposer({
chat_template_override: effectiveChatTemplateOverride,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,
- // Honor the Tensor Parallelism toggle on compare loads too.
- tensor_parallel: currentStore.tensorParallel,
+ // Honor the Tensor Parallelism + GPU Memory choices on compare loads.
+ // GGUF-only, like the auto-load path: the picker is a GGUF control,
+ // so a non-GGUF target loads via HF auto-placement instead of being
+ // pinned to a leftover GGUF pick it can't even show.
+ tensor_parallel: compareLoadKnobs.tensorParallel,
+ ...(targetIsGguf
+ ? {
+ gpu_memory_mode: compareLoadKnobs.gpuMemoryMode,
+ gpu_layers: compareLoadKnobs.gpuLayers,
+ n_cpu_moe: compareLoadKnobs.nCpuMoe,
+ tensor_split: compareLoadKnobs.splitRatio ?? undefined,
+ gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined,
+ }
+ : {}),
});
saveSpeculativeType(specSettings.speculativeType);
+ // Persist the GPU Memory mode on a non-diffusion GGUF compare-load too,
+ // so an applied manual choice survives a restart.
+ persistGpuMemoryModeOnLoad(resp, compareLoadKnobs.gpuMemoryMode);
upgradeUnloadedActive = false;
const store = useChatRuntimeStore.getState();
store.setCheckpoint(
@@ -1136,6 +1198,17 @@ export function SharedComposer({
store.setModelRequiresTrustRemoteCode(
resp.requires_trust_remote_code ?? false,
);
+ // Keep an explicit Manual+Auto context pin the load just applied (so a
+ // later Apply/Reset doesn't silently revert the model to auto-fit
+ // sizing), mirroring the interactive path's keepCustomCtx. Non-GGUF
+ // compare loads don't send the pin, so their baseline clears.
+ const keepCustomCtx = targetIsGguf
+ ? resolveManualAutoCtxPin(
+ compareLoadKnobs.gpuMemoryMode,
+ compareLoadKnobs.gpuLayers,
+ compareLoadKnobs.customContextLength,
+ )
+ : null;
useChatRuntimeStore.setState({
supportsReasoning: resp.supports_reasoning ?? false,
reasoningAlwaysOn: resp.reasoning_always_on ?? false,
@@ -1144,6 +1217,32 @@ export function SharedComposer({
supportsTools: resp.supports_tools ?? false,
tensorParallel: resp.tensor_parallel ?? false,
loadedTensorParallel: resp.tensor_parallel ?? false,
+ customContextLength: keepCustomCtx,
+ loadedCustomContextLength: keepCustomCtx,
+ // Seed the loaded GGUF context (interactive/auto-load parity): the
+ // settings sheet keys the GGUF GPU controls off it for a direct .gguf
+ // with no variant, and a later Apply reads it as the resolved context.
+ ...(targetIsGguf
+ ? {
+ ggufContextLength: resp.context_length ?? 131072,
+ ggufMaxContextLength:
+ resp.max_context_length ?? resp.context_length ?? 131072,
+ ggufNativeContextLength: resp.native_context_length ?? null,
+ }
+ : { ggufContextLength: null }),
+ // Compare loads resolve by id (HF repo / local path), never through a
+ // native-path lease, so a token left by a previously loaded native
+ // GGUF is stale here -- isLoadedGguf keys off it, and a stale token
+ // would dress a non-GGUF compare load in GGUF controls. Mirror the
+ // interactive path, which writes it on every load success.
+ activeNativePathToken: null,
+ // Held under an open staged pick: setCheckpoint preserves a stage on
+ // the empty->active transition, so a compare load can complete with
+ // staged GPU edits still on screen.
+ ...loadedGpuMemoryFieldsUnlessStaged(resp),
+ // Drives the GPU Memory controls' diffusion gate; set alongside the
+ // GPU fields on every load path so the gate can't read stale.
+ loadedIsDiffusion: resp.is_diffusion ?? false,
loadedIsMultimodal: isMultimodalResponse(resp),
...resolveLoadedSpeculativeSettings(resp),
});
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index 192ce1ec69..5786947118 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -7,6 +7,10 @@ import {
mirrorHfTokenInto,
useHfTokenStore,
} from "@/features/hub";
+import {
+ cachedPinnableGpuIndices,
+ ensureGpuDeviceCache,
+} from "@/hooks/use-gpu-info";
import { toast } from "@/lib/toast";
import { create } from "zustand";
import { isExternalModelId, parseExternalModelId } from "../external-providers";
@@ -74,6 +78,7 @@ export const CHAT_RAG_AUTOINJECT_MIN_SCORE_KEY =
export const CHAT_RAG_OCR_KEY = "unsloth_chat_rag_ocr_scanned";
export const CHAT_RAG_CAPTION_KEY = "unsloth_chat_rag_caption_figures";
export const CHAT_SPECULATIVE_TYPE_KEY = "unsloth_chat_speculative_type";
+export const CHAT_GPU_MEMORY_MODE_KEY = "unsloth_chat_gpu_memory_mode";
// Persist only the model-agnostic intents (auto/ngram/off). MTP modes
// (mtp/mtp+ngram) and spec_draft_n_max stay session-only: a persisted MTP
@@ -497,6 +502,213 @@ export function saveSpeculativeType(value: string | null): void {
}
}
+// GPU Memory strategy is a standing preference (like speculative type), not a
+// per-model setting: a "manual" choice persists across model switches and reloads.
+export function readPersistedGpuMemoryMode(): "auto" | "manual" {
+ return loadString(CHAT_GPU_MEMORY_MODE_KEY, "auto") === "manual" ? "manual" : "auto";
+}
+
+export function saveGpuMemoryMode(value: "auto" | "manual"): void {
+ saveString(CHAT_GPU_MEMORY_MODE_KEY, value);
+}
+
+/** Persist the GPU Memory mode after a load, but only for a non-diffusion GGUF:
+ * non-GGUF has no such mode, and diffusion runs mode-agnostic (reports "auto"),
+ * so neither must clobber the standing manual preference. */
+export function persistGpuMemoryModeOnLoad(
+ resp: { is_gguf?: boolean; is_diffusion?: boolean },
+ mode: "auto" | "manual",
+): void {
+ if (resp.is_gguf && !resp.is_diffusion) saveGpuMemoryMode(mode);
+}
+
+// Manual-mode gpu_layers sentinel: -1 = Auto (hand layer + context sizing to
+// llama.cpp's --fit). The Manual default; "all on GPU" is the slider's max.
+export const GPU_LAYERS_AUTO = -1;
+
+// Round real-valued shares to integers summing exactly to `total`, giving the
+// leftover units to the largest fractional parts (largest-remainder method).
+function largestRemainder(shares: number[], total: number): number[] {
+ const out = shares.map((x) => Math.floor(x));
+ let rem = total - out.reduce((a, b) => a + b, 0);
+ const byFrac = shares
+ .map((x, i) => ({ i, frac: x - Math.floor(x) }))
+ .sort((a, b) => b.frac - a.frac);
+ for (let k = 0; rem > 0 && k < byFrac.length; k++, rem--) out[byFrac[k].i] += 1;
+ return out;
+}
+
+// Spread `total` layers across GPUs in proportion to `weights` (e.g. per-GPU
+// VRAM), as integers summing exactly to `total`; even split for all-zero/empty
+// weights. Default per-GPU layer split before the user edits it (mirrors
+// llama.cpp's free-VRAM default).
+export function distributeByWeight(total: number, weights: number[]): number[] {
+ if (weights.length === 0) return [];
+ const t = Math.max(0, Math.floor(total));
+ const sum = weights.reduce((a, b) => a + b, 0);
+ const w = sum > 0 ? weights : weights.map(() => 1);
+ const wSum = w.reduce((a, b) => a + b, 0);
+ return largestRemainder(
+ w.map((x) => (t * x) / wSum),
+ t,
+ );
+}
+
+// Set GPU `index` to `value` and rebalance the rest so per-GPU counts still sum
+// to `total`; others absorb the remainder in proportion to their counts (evenly
+// if all zero). The --tensor-split editor: counts are sent verbatim, and
+// llama.cpp gives each GPU exactly its count when gpu_layers == sum(counts).
+export function rebalanceSplit(
+ total: number,
+ counts: number[],
+ index: number,
+ value: number,
+): number[] {
+ const v = Math.max(0, Math.min(value, total));
+ const out = counts.slice();
+ const otherIdx = counts.map((_, i) => i).filter((i) => i !== index);
+ // No other GPU to absorb the remainder: this one holds everything.
+ if (otherIdx.length === 0) {
+ out[index] = total;
+ return out;
+ }
+ out[index] = v;
+ const dist = distributeByWeight(
+ total - v,
+ otherIdx.map((i) => counts[i]),
+ );
+ otherIdx.forEach((i, k) => (out[i] = dist[k]));
+ return out;
+}
+
+// Validate a persisted gpu_ids pick against the GPUs present right now, before
+// restoring it from remembered settings. Returns null (= automatic) when the
+// pick is stale (none of the saved ids exist, or the host can't pin a multi-GPU
+// set), so a saved [1] on a now-1-GPU host doesn't get sent and rejected with no
+// way to clear it. A null pick (= automatic) passes through unchanged, and an
+// unpopulated device cache leaves the pick alone (the backend still guards).
+export function reconcilePersistedGpuIds(
+ ids: number[] | null,
+): number[] | null {
+ if (ids == null) return ids;
+ const pinnable = cachedPinnableGpuIndices();
+ if (pinnable === null) return ids; // cache not ready: can't validate, keep it
+ const kept = ids.filter((i) => pinnable.includes(i));
+ return kept.length > 0 ? kept : null;
+}
+
+// Store fields derived from a load/status response's GPU-memory settings.
+// Shared by every load path so the manual-knob round-trip can't drift.
+export function loadedGpuMemoryFields(resp: {
+ is_gguf?: boolean;
+ is_diffusion?: boolean;
+ gpu_memory_mode?: "auto" | "manual";
+ gpu_layers?: number;
+ n_cpu_moe?: number;
+ tensor_split?: number[] | null;
+ n_layers?: number | null;
+ n_moe_layers?: number;
+ gpu_ids?: number[] | null;
+}) {
+ // GPU-memory state is meaningful only for a GGUF chat load. A non-GGUF response
+ // still carries gpu_memory_mode (its default "auto" is serialized), so gate on
+ // the authoritative is_gguf flag, not the field's presence -- otherwise loading
+ // a transformers model would reset the standing manual preference.
+ if (!resp.is_gguf) {
+ // Clear the GPU pick / offload baseline a prior GGUF load may have left, so it
+ // reflects the non-GGUF model (no pin) -- else a stale loadedGpuIds reads as
+ // dirty (gpuIdsDirty is ungated) and Reset restores it while the picker is
+ // hidden. gpuMemoryMode (the standing preference) is kept, but its loaded
+ // baseline clears to null so Reset preserves the preference, not a stale mode.
+ return {
+ selectedGpuIds: null,
+ loadedGpuIds: null,
+ loadedGpuMemoryMode: null,
+ gpuLayers: GPU_LAYERS_AUTO,
+ loadedGpuLayers: null,
+ nCpuMoe: 0,
+ loadedNCpuMoe: null,
+ splitRatio: null,
+ loadedSplitRatio: null,
+ ggufLayerCount: null,
+ moeLayerCount: null,
+ };
+ }
+ const mode = resp.gpu_memory_mode ?? "auto";
+ const gpuIds = resp.gpu_ids ?? null;
+ // Layer/MoE/split knobs apply (and are reported) only in manual mode; in auto
+ // the server ignores them, so don't seed the loaded baseline or the editable
+ // knobs with values it never applied. In manual, the server reports gpu_layers
+ // = -1 under Auto, which round-trips the slider back to its Auto position.
+ const manualKnobs =
+ mode === "manual"
+ ? {
+ loadedGpuLayers: resp.gpu_layers ?? null,
+ loadedNCpuMoe: resp.n_cpu_moe ?? null,
+ loadedSplitRatio: resp.tensor_split ?? null,
+ gpuLayers: resp.gpu_layers ?? GPU_LAYERS_AUTO,
+ nCpuMoe: resp.n_cpu_moe ?? 0,
+ splitRatio: resp.tensor_split ?? null,
+ }
+ : {
+ loadedGpuLayers: null,
+ loadedNCpuMoe: null,
+ loadedSplitRatio: null,
+ // Auto ignores these, so reset the editable knobs too (not just the
+ // loaded baseline) -- else a later switch back to Manual would snapshot
+ // and send a previous model's stale gpuLayers/nCpuMoe/split that this
+ // load never applied. Mirrors the non-GGUF branch above.
+ gpuLayers: GPU_LAYERS_AUTO,
+ nCpuMoe: 0,
+ splitRatio: null,
+ };
+ return {
+ // A diffusion GGUF runs mode-agnostic (pins all layers on one GPU, reports
+ // "auto"), so adopt everything a chat GGUF does EXCEPT the live standing
+ // preference -- the next chat load must still honor the user's manual choice.
+ // The loaded baseline is still "auto", but the UI hides mode controls for a
+ // loaded diffusion model so it can't read as dirty against the preference.
+ ...(resp.is_diffusion ? {} : { gpuMemoryMode: mode }),
+ loadedGpuMemoryMode: mode,
+ ggufLayerCount: resp.n_layers ?? null,
+ // MoE expert-layer count: the n_cpu_moe slider max, and 0 hides the slider.
+ moeLayerCount: resp.n_moe_layers ?? null,
+ // The picker reflects what loaded (the request sent the user's pick).
+ selectedGpuIds: gpuIds,
+ loadedGpuIds: gpuIds,
+ ...manualKnobs,
+ };
+}
+
+/** loadedGpuMemoryFields (plus any seedExtras), unless a staged pick is open.
+ *
+ * With a staged pick open (the load fired mid-staging), preserve its editable
+ * GPU knobs and seedExtras, but still advance every loaded baseline. Otherwise
+ * cancelling the stage restores its edits onto the newly loaded model. The
+ * status reseed cannot repair that while pendingSelection holds it off.
+ */
+export function loadedGpuMemoryFieldsUnlessStaged(
+ resp: Parameters[0],
+ seedExtras?: T,
+) {
+ const fields = loadedGpuMemoryFields(resp);
+ if (useChatRuntimeStore.getState().pendingSelection != null) {
+ return {
+ loadedGpuMemoryMode: fields.loadedGpuMemoryMode,
+ loadedGpuLayers: fields.loadedGpuLayers,
+ loadedNCpuMoe: fields.loadedNCpuMoe,
+ loadedSplitRatio: fields.loadedSplitRatio,
+ loadedGpuIds: fields.loadedGpuIds,
+ // These are metadata ceilings for the model that actually loaded, not
+ // editable values from the open stage. Advance them with the baselines
+ // so abandoning the stage cannot expose the previous model's limits.
+ ggufLayerCount: fields.ggufLayerCount,
+ moeLayerCount: fields.moeLayerCount,
+ };
+ }
+ return { ...fields, ...seedExtras };
+}
+
/** A local model staged for a deferred load (see `pendingSelection`). Shape is
* a subset of the load hook's `SelectedModelInput`, structurally assignable. */
export type PendingModelSelection = {
@@ -515,6 +727,13 @@ export type PendingModelSelection = {
* Scoped here (not the shared `ggufContextLength`) so a staged model's
* metadata never pollutes the currently-loaded model's context display. */
contextLength?: number | null;
+ /** Total layer count (GGUF block_count); the manual gpu-layers ceiling is
+ * this + 1 (llama.cpp counts the output layer as offloadable too);
+ * scoped here like contextLength. */
+ layerCount?: number | null;
+ /** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling);
+ * 0 for dense models, scoped here like contextLength. */
+ moeLayerCount?: number | null;
/** "Load on selection" on + un-cached GGUF: download via the manager (global
* indicator) without opening the sheet, then load once the download finishes. */
autoLoad?: boolean;
@@ -743,6 +962,32 @@ type ChatRuntimeStore = {
tensorParallel: boolean;
/** Backend-reported tensor-parallel state; null until first hydrated. */
loadedTensorParallel: boolean | null;
+ /** GPU memory strategy for GGUF loads. "auto" = Unsloth picks GPUs and context
+ * to fit; "manual" = you own the offload (gpuLayers < 0 = Auto/--fit, >= 0
+ * pins layers + nCpuMoe). */
+ gpuMemoryMode: "auto" | "manual";
+ /** Backend-reported gpu memory mode; null until first hydrated. */
+ loadedGpuMemoryMode: "auto" | "manual" | null;
+ /** Manual mode: layers to offload to GPU. -1 = Auto (--fit); >= model layer
+ * count = all. */
+ gpuLayers: number;
+ loadedGpuLayers: number | null;
+ /** Manual mode: MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. */
+ nCpuMoe: number;
+ loadedNCpuMoe: number | null;
+ /** Manual mode: per-GPU layer counts (--tensor-split), in GPU-in-use order;
+ * null = unset (llama.cpp splits by free VRAM). */
+ splitRatio: number[] | null;
+ /** Backend-reported per-GPU split ratio (--tensor-split); null = unset. */
+ loadedSplitRatio: number[] | null;
+ /** Model layer count (GGUF block_count); the manual gpu-layers ceiling is
+ * this + 1 (the output layer is offloadable too). */
+ ggufLayerCount: number | null;
+ /** MoE expert-layer count: the nCpuMoe slider max; 0/null hides the slider. */
+ moeLayerCount: number | null;
+ /** Picked physical GPU indices (null = use all / automatic). */
+ selectedGpuIds: number[] | null;
+ loadedGpuIds: number[] | null;
/** Persisted: when false, picking a local model stages it as
* `pendingSelection` (and opens settings) instead of loading immediately,
* so load settings can be set before the single load. */
@@ -766,6 +1011,9 @@ type ChatRuntimeStore = {
* per step, cleared when the run ends, never persisted into the transcript. */
activeDiffusionCanvas: DiffusionCanvasFrame | null;
customContextLength: number | null;
+ /** The pinned context the loaded model used (null = Auto), so dirty-tracking
+ * and a later fit Apply can tell an explicit pin apart from Auto. */
+ loadedCustomContextLength: number | null;
defaultChatTemplate: string | null;
chatTemplateOverride: string | null;
loadedChatTemplateOverride: string | null;
@@ -884,6 +1132,11 @@ type ChatRuntimeStore = {
* which skip the sheet but must still honor a saved config. */
applyRememberedLoadSettings: (settings: RememberedLoadSettings) => void;
setTensorParallel: (value: boolean) => void;
+ setGpuMemoryMode: (mode: "auto" | "manual") => void;
+ setGpuLayers: (value: number) => void;
+ setNCpuMoe: (value: number) => void;
+ setSplitRatio: (value: number[] | null) => void;
+ setSelectedGpuIds: (ids: number[] | null) => void;
setLoadOnSelection: (value: boolean) => void;
setExpandQuantizations: (value: boolean) => void;
setShowAllQuantizations: (value: boolean) => void;
@@ -1101,11 +1354,12 @@ function setScalarSettingVersion(
/** The "revert to the loaded model" baseline for the editable load knobs.
* Shared by resetModelSettingsToLoaded (full revert) and stageModel (which
- * overrides speculative to start a fresh pick from the standing default). */
+ * overrides speculative and the per-model GPU knobs to start a fresh pick). */
function loadedBaselineSettings(s: ChatRuntimeStore) {
const hasLoadedModel = Boolean(s.params.checkpoint);
return {
- customContextLength: null,
+ // Revert to the loaded model's pin (null = Auto), not a blanket Auto.
+ customContextLength: s.loadedCustomContextLength,
kvCacheDtype: s.loadedKvCacheDtype,
tensorParallel: s.loadedTensorParallel ?? false,
speculativeType: hasLoadedModel
@@ -1113,6 +1367,20 @@ function loadedBaselineSettings(s: ChatRuntimeStore) {
: readPersistedSpeculativeType(),
specDraftNMax: hasLoadedModel ? s.loadedSpecDraftNMax : null,
chatTemplateOverride: s.loadedChatTemplateOverride,
+ // GPU memory mode is a standing preference; revert to the loaded model's
+ // mode (or the persisted default when nothing is loaded). Manual knobs and
+ // the GPU pick are per-model and revert to their loaded baseline. A loaded
+ // model with no applicable mode -- diffusion ("auto" baseline) or non-GGUF
+ // (null baseline) -- keeps the live preference so Reset can't drop it.
+ gpuMemoryMode: !hasLoadedModel
+ ? readPersistedGpuMemoryMode()
+ : s.loadedIsDiffusion
+ ? s.gpuMemoryMode
+ : (s.loadedGpuMemoryMode ?? s.gpuMemoryMode),
+ gpuLayers: s.loadedGpuLayers ?? GPU_LAYERS_AUTO,
+ nCpuMoe: s.loadedNCpuMoe ?? 0,
+ splitRatio: s.loadedSplitRatio ?? null,
+ selectedGpuIds: s.loadedGpuIds,
};
}
@@ -1213,6 +1481,18 @@ export const useChatRuntimeStore = create((set, get) => ({
loadedSpecDraftNMax: null,
tensorParallel: false,
loadedTensorParallel: null,
+ gpuMemoryMode: readPersistedGpuMemoryMode(),
+ loadedGpuMemoryMode: null,
+ gpuLayers: GPU_LAYERS_AUTO,
+ loadedGpuLayers: null,
+ nCpuMoe: 0,
+ loadedNCpuMoe: null,
+ splitRatio: null,
+ loadedSplitRatio: null,
+ ggufLayerCount: null,
+ moeLayerCount: null,
+ selectedGpuIds: null,
+ loadedGpuIds: null,
loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true),
expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false),
showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true),
@@ -1221,6 +1501,7 @@ export const useChatRuntimeStore = create((set, get) => ({
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
+ loadedCustomContextLength: null,
defaultChatTemplate: null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
@@ -1455,9 +1736,23 @@ export const useChatRuntimeStore = create((set, get) => ({
loadedSpecDraftNMax: null,
tensorParallel: false,
loadedTensorParallel: null,
+ // Standing preference: survives unload, unlike the per-model knobs above.
+ gpuMemoryMode: readPersistedGpuMemoryMode(),
+ loadedGpuMemoryMode: null,
+ gpuLayers: GPU_LAYERS_AUTO,
+ loadedGpuLayers: null,
+ nCpuMoe: 0,
+ loadedNCpuMoe: null,
+ splitRatio: null,
+ loadedSplitRatio: null,
+ ggufLayerCount: null,
+ moeLayerCount: null,
+ selectedGpuIds: null,
+ loadedGpuIds: null,
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
+ loadedCustomContextLength: null,
defaultChatTemplate: null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
@@ -1753,17 +2048,67 @@ export const useChatRuntimeStore = create((set, get) => ({
setSpeculativeType: (speculativeType) => set({ speculativeType }),
setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }),
setTensorParallel: (tensorParallel) => set({ tensorParallel }),
+ // Standing preference, but persisted only on a successful load (see
+ // use-chat-model-runtime), not on selection -- so an unapplied pick the user
+ // resets/abandons doesn't stick to the next session.
+ setGpuMemoryMode: (gpuMemoryMode) => set({ gpuMemoryMode }),
+ setGpuLayers: (gpuLayers) => set({ gpuLayers }),
+ setNCpuMoe: (nCpuMoe) => set({ nCpuMoe }),
+ setSplitRatio: (splitRatio) => set({ splitRatio }),
+ setSelectedGpuIds: (selectedGpuIds) => set({ selectedGpuIds }),
resetModelSettingsToLoaded: () => set((s) => loadedBaselineSettings(s)),
- applyRememberedLoadSettings: (settings) =>
+ applyRememberedLoadSettings: (settings) => {
+ const gpuCacheWasCold = cachedPinnableGpuIndices() === null;
+ const restoredGpuIds =
+ settings.selectedGpuIds !== undefined
+ ? reconcilePersistedGpuIds(settings.selectedGpuIds)
+ : undefined;
// Coalesce every field: a blob persisted by an older/newer build can omit
// keys, and a raw spread would push `undefined` into fields typed non-null.
+ // The GPU knobs are spread only when present, but first reset the per-model
+ // ones to defaults: this path (load-on-selection) starts from the loaded
+ // model's baseline and skips the model-switch reset, so a blob omitting
+ // gpuLayers/nCpuMoe/selectedGpuIds (older build) or splitRatio (never
+ // remembered) must not inherit the previous model's placement. gpuMemoryMode
+ // (standing preference) is NOT reset, only applied when the blob carries it;
+ // selectedGpuIds keeps a meaningful null (all GPUs), so it keys off undefined.
set({
+ gpuLayers: GPU_LAYERS_AUTO,
+ nCpuMoe: 0,
+ splitRatio: null,
+ selectedGpuIds: null,
customContextLength: settings.contextLength ?? null,
kvCacheDtype: settings.kvCacheDtype ?? null,
speculativeType: settings.speculativeType ?? "auto",
specDraftNMax: settings.specDraftNMax ?? null,
tensorParallel: settings.tensorParallel ?? false,
- }),
+ ...(settings.gpuMemoryMode != null && {
+ gpuMemoryMode: settings.gpuMemoryMode,
+ }),
+ ...(settings.gpuLayers != null && { gpuLayers: settings.gpuLayers }),
+ ...(settings.nCpuMoe != null && { nCpuMoe: settings.nCpuMoe }),
+ ...(restoredGpuIds !== undefined && {
+ // Reconcile against the GPUs present now (see reconcilePersistedGpuIds):
+ // a saved [1] on a 1-GPU host (or under relative/UUID visibility) would
+ // hide the picker yet still send gpu_ids, which the backend rejects.
+ selectedGpuIds: restoredGpuIds,
+ }),
+ });
+ // A cold cache makes the synchronous restore provisional. Reconcile again
+ // when the shared fetch completes, but only if this exact restored array is
+ // still current so a user edit, stage change, or load cannot be overwritten.
+ if (gpuCacheWasCold && restoredGpuIds != null) {
+ void ensureGpuDeviceCache().then(() => {
+ set((state) => {
+ if (state.selectedGpuIds !== restoredGpuIds) return state;
+ const reconciled = reconcilePersistedGpuIds(restoredGpuIds);
+ return reconciled === restoredGpuIds
+ ? state
+ : { selectedGpuIds: reconciled };
+ });
+ });
+ }
+ },
setLoadOnSelection: (loadOnSelection) => {
saveBool(CHAT_LOAD_ON_SELECTION_KEY, loadOnSelection);
set({ loadOnSelection });
@@ -1798,6 +2143,22 @@ export const useChatRuntimeStore = create((set, get) => ({
// Load's keepSpeculative) a forced MTP mode onto a model that may lack it.
speculativeType: readPersistedSpeculativeType(),
specDraftNMax: null,
+ // Keep the on-screen GPU Memory selection (loadedBaselineSettings would
+ // otherwise revert it to the loaded model's mode, dropping a Manual choice
+ // just made). Use the live store value, not the persisted one, which can
+ // lag a mode hydrated from an out-of-band load.
+ gpuMemoryMode: s.gpuMemoryMode,
+ // Per-model GPU knobs start from defaults too so a fresh pick doesn't
+ // inherit the loaded model's layer/MoE/split/GPU choices, matching the
+ // immediate-switch reset.
+ gpuLayers: GPU_LAYERS_AUTO,
+ nCpuMoe: 0,
+ splitRatio: null,
+ selectedGpuIds: null,
+ // Fresh pick starts at Auto context (loadedBaselineSettings would
+ // otherwise restore the current model's pin). Leaves the baseline
+ // intact, like the GPU knobs, so abandoning restores the loaded pin.
+ customContextLength: null,
};
});
},
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index d72c406fdd..c24ddde5f5 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -65,6 +65,18 @@ export interface LoadModelRequest {
* of by layer for GGUF models. Multi-GPU only; no effect on a single GPU.
*/
tensor_parallel?: boolean | null;
+ /** GPU memory strategy for GGUF models. "auto" (default): Unsloth selects GPUs
+ * and caps context to fit VRAM. "manual": you own the offload -- gpu_layers
+ * -1 (Auto) hands sizing to llama.cpp's --fit, >= 0 pins layers/n_cpu_moe. */
+ gpu_memory_mode?: "auto" | "manual";
+ /** Manual mode: layers to offload to GPU (--gpu-layers, --fit off); -1 = Auto (--fit). */
+ gpu_layers?: number;
+ /** Manual mode: MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. */
+ n_cpu_moe?: number;
+ /** Manual mode: relative model share per GPU (--tensor-split), in GPU order. */
+ tensor_split?: number[] | null;
+ /** Picked physical GPU indices (omit/empty = automatic). */
+ gpu_ids?: number[];
}
export interface ValidateModelResponse {
@@ -80,6 +92,13 @@ export interface ValidateModelResponse {
requires_security_review?: boolean;
/** Native context length from the local GGUF header; null until downloaded. */
context_length?: number | null;
+ /** Total layer count (GGUF block_count); the manual gpu-layers ceiling is
+ * this + 1 (llama.cpp counts the output layer as offloadable too); null
+ * until downloaded. */
+ layer_count?: number | null;
+ /** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling);
+ * 0 for dense models, null until downloaded. */
+ moe_layer_count?: number | null;
/** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */
requires_transformers_upgrade?: boolean;
/** Set only when requires_transformers_upgrade. */
@@ -159,6 +178,14 @@ export interface LoadModelResponse {
spec_draft_n_max?: number | null;
/** Whether tensor-parallel split (--split-mode tensor) is active. */
tensor_parallel?: boolean;
+ gpu_memory_mode?: "auto" | "manual";
+ gpu_layers?: number;
+ n_cpu_moe?: number;
+ tensor_split?: number[] | null;
+ n_layers?: number | null;
+ /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */
+ n_moe_layers?: number;
+ gpu_ids?: number[] | null;
}
export interface UnloadModelRequest {
@@ -203,6 +230,17 @@ export interface InferenceStatusResponse {
spec_draft_n_max?: number | null;
/** Whether tensor-parallel split (--split-mode tensor) is active. */
tensor_parallel?: boolean;
+ gpu_memory_mode?: "auto" | "manual";
+ gpu_layers?: number;
+ n_cpu_moe?: number;
+ tensor_split?: number[] | null;
+ /** n_ctx the active GGUF load was invoked with (0 = Auto); re-seeds a
+ * Manual + Auto-layers context pin on hydration. Null for non-GGUF. */
+ requested_context_length?: number | null;
+ gpu_ids?: number[] | null;
+ n_layers?: number | null;
+ /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */
+ n_moe_layers?: number;
/**
* Why MTP was disabled on the loaded model despite being requested.
* "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable
diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts
index 1e313acdf3..db2cc021be 100644
--- a/studio/frontend/src/hooks/use-gpu-info.ts
+++ b/studio/frontend/src/hooks/use-gpu-info.ts
@@ -15,6 +15,19 @@ export interface GpuInfo {
systemRamTotalGb: number
}
+export interface SystemGpuDevice {
+ index: number;
+ name: string;
+ memoryTotalGb: number;
+ /** Free VRAM at fetch time. Degrades to the total when the utilization
+ * probe had no usage data; 0 only when the total is unknown too. */
+ memoryFreeGb: number;
+ /** "physical" = `index` is a stable physical/PCI id safe to pin via gpu_ids;
+ * "relative" = an ordinal into a parent CUDA_VISIBLE_DEVICES mask, which the
+ * backend can't map back, so the picker must not offer it. */
+ physicalIndex: boolean;
+}
+
const DEFAULT_GPU: GpuInfo = {
available: false,
name: "Unknown",
@@ -25,70 +38,135 @@ const DEFAULT_GPU: GpuInfo = {
systemRamTotalGb: 0
};
-// Module-level cache so multiple components share one fetch.
-let cachedGpu: GpuInfo | null = null;
-let fetchPromise: Promise | null = null;
+// One module-level cache so every GPU hook shares a single /api/system fetch.
+let cachedSystem: SystemInfoResponse | null = null;
+let systemPromise: Promise | null = null;
-async function fetchGpuOnce(): Promise {
- if (cachedGpu) return cachedGpu;
- if (fetchPromise) return fetchPromise;
-
- fetchPromise = (async () => {
+async function fetchSystemOnce(): Promise {
+ if (cachedSystem) return cachedSystem;
+ if (systemPromise) return systemPromise;
+ systemPromise = (async () => {
try {
const res = await authFetch("/api/system");
if (!res.ok) throw new Error(`HTTP ${res.status}`);
-
- const data = await res.json() as SystemInfoResponse;
- const gpuData = data?.gpu;
-
- // CPU/RAM exist even on hosts without a GPU, so populate them on every path.
- // No discrete GPU (e.g. Mac): still surface system RAM so memory math
- // (unified memory) has a budget to work with.
- const base = {
- cpuCore: data?.cpu?.physical_count ?? 0,
- cpuThread: data?.cpu?.logical_count ?? 0,
- systemRamAvailableGb: data?.memory?.available_gb ?? 0,
- systemRamTotalGb: data?.memory?.total_gb ?? 0,
- };
-
- const devices = gpuData?.devices ?? [];
- const info: GpuInfo =
- gpuData?.available && devices.length
- ? {
- ...base,
- available: true,
- name: devices[0]?.name ?? "Unknown",
- memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0),
- }
- : { ...DEFAULT_GPU, ...base };
- cachedGpu = info;
- return info;
+ cachedSystem = (await res.json()) as SystemInfoResponse;
+ return cachedSystem;
} catch {
- // Reset promise so subsequent calls retry (e.g. backend wasn't ready)
- fetchPromise = null;
- return DEFAULT_GPU;
+ systemPromise = null; // reset so a later call retries (backend not ready)
+ return null;
}
})();
+ return systemPromise;
+}
- return fetchPromise;
+function toGpuInfo(data: SystemInfoResponse | null): GpuInfo {
+ // CPU/RAM exist even on GPU-less hosts (e.g. Mac), so populate them on every
+ // path: unified-memory math still needs a RAM budget to work with.
+ const base = {
+ cpuCore: data?.cpu?.physical_count ?? 0,
+ cpuThread: data?.cpu?.logical_count ?? 0,
+ systemRamAvailableGb: data?.memory?.available_gb ?? 0,
+ systemRamTotalGb: data?.memory?.total_gb ?? 0,
+ };
+ const gpuData = data?.gpu;
+ const devices = gpuData?.devices ?? [];
+ if (!gpuData?.available || !devices.length) {
+ return { ...DEFAULT_GPU, ...base };
+ }
+ return {
+ ...base,
+ available: true,
+ name: devices[0]?.name ?? "Unknown",
+ memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0),
+ };
+}
+
+function toGpuDevices(data: SystemInfoResponse | null): SystemGpuDevice[] {
+ // Unpinnable configurations must hide every pick surface: XPU indices are
+ // torch-xpu ordinals no applicator speaks, and Vulkan-only builds pin ggml's
+ // own ordinals -- /load and /validate 400 picks on both, so the backend
+ // reports gpu.gguf_gpu_ids_supported and every gate keyed on physicalIndex
+ // (picker, persisted-pick reconcile) follows it. The device flavor lives on
+ // the TOP-LEVEL device_backend field; absent support info defaults to
+ // pinnable (older backend).
+ const pinnableBackend =
+ data?.device_backend !== "xpu" &&
+ data?.gpu?.gguf_gpu_ids_supported !== false;
+ return (data?.gpu?.devices ?? [])
+ .filter((d) => typeof d.index === "number")
+ .map((d) => ({
+ index: d.index as number,
+ name: d.name ?? `GPU ${d.index}`,
+ memoryTotalGb: d.memory_total_gb ?? 0,
+ memoryFreeGb: d.vram_free_gb ?? 0,
+ physicalIndex: pinnableBackend && d.index_kind === "physical",
+ }));
+}
+
+/** Aggregate GPU info from /api/system; shares one module-level fetch across all GPU hooks. */
+export function useGpuInfo(): GpuInfo {
+ const [gpu, setGpu] = useState(
+ cachedSystem ? toGpuInfo(cachedSystem) : DEFAULT_GPU,
+ );
+ useEffect(() => {
+ // No early return on cachedSystem: a consumer mounting as the cache fills
+ // (between render and effect) would otherwise stay stuck at the default.
+ let cancelled = false;
+ fetchSystemOnce().then((d) => {
+ if (!cancelled) setGpu(toGpuInfo(d));
+ });
+ return () => {
+ cancelled = true;
+ };
+ }, []);
+ return gpu;
+}
+
+/** All backend-visible GPUs (index, name, total VRAM); shares the same fetch. */
+export function useGpuDevices(): SystemGpuDevice[] {
+ const [devices, setDevices] = useState(
+ cachedSystem ? toGpuDevices(cachedSystem) : [],
+ );
+ useEffect(() => {
+ // No early return on cachedSystem: a consumer mounting as the cache fills
+ // (between render and effect) would otherwise stay stuck at the default.
+ let cancelled = false;
+ fetchSystemOnce().then((d) => {
+ if (!cancelled) setDevices(toGpuDevices(d));
+ });
+ return () => {
+ cancelled = true;
+ };
+ }, []);
+ return devices;
}
/**
- * Fetch GPU info from /api/system. Cached at module level, so only one request
- * is made no matter how many components call this hook.
+ * Await the shared /api/system fetch so cachedPinnableGpuIndices (and the
+ * store's reconcilePersistedGpuIds) can validate a persisted pick before a
+ * load path sends it -- on a cold cache the reconcile passes ids through
+ * unvalidated, and a stale cross-host pick then fails /load with the picker
+ * hidden. Resolves immediately once the module cache is warm; a failed fetch
+ * keeps the cache cold, preserving the "can't validate, backend guards"
+ * degradation.
*/
-export function useGpuInfo(): GpuInfo {
- const [gpu, setGpu] = useState(cachedGpu ?? DEFAULT_GPU);
+export async function ensureGpuDeviceCache(): Promise {
+ await fetchSystemOnce();
+}
- useEffect(() => {
- if (cachedGpu) return;
-
- let cancelled = false;
- fetchGpuOnce().then((info) => {
- if (!cancelled) setGpu(info);
- });
- return () => { cancelled = true; };
- }, []);
-
- return gpu;
-}
\ No newline at end of file
+/**
+ * Pinnable physical GPU indices from the already-fetched /api/system cache, for
+ * non-React code (the store) that needs to validate a persisted `gpu_ids` pick
+ * without triggering a fetch. Returns:
+ * - `null` when the cache isn't populated yet (caller can't validate, so keep
+ * the pick and let the backend guard reject a truly bad one);
+ * - `[]` when the host has no pinnable multi-GPU set (single GPU, or relative/
+ * UUID-masked indices) -- the picker is hidden, so any saved pick is stale;
+ * - the physical indices otherwise.
+ */
+export function cachedPinnableGpuIndices(): number[] | null {
+ if (!cachedSystem) return null;
+ const physical = toGpuDevices(cachedSystem).filter((d) => d.physicalIndex);
+ // Mirrors the sheet's showGpuPicker gate: only a 2+ physical-GPU host can pin.
+ return physical.length > 1 ? physical.map((d) => d.index) : [];
+}
diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts
index a135cce86e..8cfe2bace4 100644
--- a/studio/frontend/src/hooks/use-system.ts
+++ b/studio/frontend/src/hooks/use-system.ts
@@ -40,6 +40,9 @@ export interface SystemInfoResponse {
gpu: {
available: boolean;
backend?: string;
+ /** Whether GGUF loads accept an explicit gpu_ids pick (false on XPU hosts
+ * and Vulkan-only builds, where /load and /validate 400 picks). */
+ gguf_gpu_ids_supported?: boolean;
backend_cuda_visible_devices?: string | null;
parent_visible_gpu_ids?: number[];
index_kind?: string;
From 03590f696e97401361d59d61e1b9b367238ea229 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Sun, 19 Jul 2026 06:08:54 -0700
Subject: [PATCH 017/255] Give opencode real timeout headroom in Local Agent
Guides CI (#7235)
* Raise the opencode invoke timeout in Local Agent Guides CI
The connection (opencode) cell flakes with a 600s timeout reported as guide drift, but it is not a hang: in a passing run the same opencode run finishes in ~482s (08:12:31 to 08:20:33), right against the shared AGENT_INVOKE_TIMEOUT of 600s, so about one run in six drifts past the cap.
opencode is the slow outlier. The print-mode agents (claude -p, codex exec) run one turn against a minimal injected system prompt, while opencode run runs its own full turn with opencode's large system prompt plus a separate small_model call to name the session (start.py pins small_model to the same 4B the server hosts). On a CPU-served gemma-4-E4B that is about 8 minutes, leaving no margin under 600s.
Double opencode's per-invoke timeout in agent-guides-drive.sh and keep the tight 600s cap for the fast agents, so a genuine headless-TTY hang still fails quickly. 1200s stays well under the 40-minute job budget.
* Normalize the agent invoke timeout before doubling it for opencode
Strip an optional trailing 's' from AGENT_INVOKE_TIMEOUT so the opencode
arithmetic, and the "${TIMEOUT}s" timeout message, stay valid if a
timeout(1)-style suffix is ever configured.
* Only double the opencode timeout for a bare-integer seconds value
Guard the arithmetic so a GNU timeout(1) duration suffix (s/m/h/d, including
floats like 0.5s) is passed through unchanged instead of breaking the
expansion; timeout(1) parses those directly. Bare seconds still double.
---------
Co-authored-by: danielhanchen
---
.github/scripts/agent-guides-drive.sh | 17 +++++++++++++++++
1 file changed, 17 insertions(+)
diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh
index 2457f08407..b63ac94b93 100755
--- a/.github/scripts/agent-guides-drive.sh
+++ b/.github/scripts/agent-guides-drive.sh
@@ -36,6 +36,23 @@ AGENT="${2:?usage: agent-guides-drive.sh }"
# Determinism (seed/temp) is applied at the server level by
# serve-unsloth-run.sh --extra; agents inherit it through the API.
TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}"
+# opencode is the slow outlier. Unlike the print-mode agents (claude -p, codex
+# exec) it runs a full turn AND a separate small_model call to name the session,
+# so one connection reply takes ~8 min on a CPU-served 4B -- right at the shared
+# 600s cap, so the cell flaked when a run drifted past a ~480s success. Give it
+# headroom (still well under the 40-min job budget); the fast agents keep the
+# tight cap that still catches a real headless-TTY hang.
+case "$AGENT" in
+ opencode)
+ # Double it, but only for a bare-integer seconds value. A GNU timeout(1)
+ # duration suffix (s/m/h/d, including floats like 0.5s) is left unchanged so
+ # the arithmetic never sees a non-number; timeout(1) parses it directly.
+ case "$TIMEOUT" in
+ *[!0-9]*) ;;
+ *) TIMEOUT=$(( TIMEOUT * 2 )) ;;
+ esac
+ ;;
+esac
# Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner
# IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless
From a9be36830eb4ec731dcd10008d2f6dc3bb102d40 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Sun, 19 Jul 2026 06:19:29 -0700
Subject: [PATCH 018/255] Installer: allow torch 2.11.x on the CUDA install
path (fresh install + studio) (#6959)
* Studio: allow torch 2.11.x on the CUDA install path
The CUDA torch repair path (_ensure_cuda_torch) installs torch/torchvision/
torchaudio from an exclusive --index-url, so _CUDA_TORCH_PKG_SPEC decides
exactly which torch the Studio venv gets. It was capped at torch<2.11.0, so on
a cu128/cu130 host the venv resolved torch 2.10.x even though the CUDA indexes
now publish torch 2.11.0. That left the Studio venv a torch minor behind the
torch 2.11.0 Docker base image, so the CUDA dedup step would relink base libs
under a mismatched torch.
Raise the upper bound to <2.12.0 (torchvision <0.27.0, torchaudio <2.12.0) so
the CUDA install path lands on torch 2.11.x, matching the rocm7.2 spec and the
base image. The torchao selector already maps torch 2.11 -> torchao 0.17.0, and
_ensure_flash_attn degrades gracefully when no prebuilt wheel matches (Blackwell
skips it outright; non-Blackwell prints a warning and continues), so no other
pin needs to move.
Add test_cuda_torch_spec.py to lock the bound (torch 2.11.x in, 2.12.x out) and
assert the CUDA and rocm7.2 upper bounds stay in lockstep.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* test: use zip(strict=True) so a spec length mismatch fails loudly
* install.sh: widen the CUDA torch ceiling to <2.12.0 so a fresh install matches the base
Raising _CUDA_TORCH_PKG_SPEC alone was not enough: that spec only feeds
_ensure_cuda_torch(), the ROCm-poisoning repair path that early-returns on a
normal NVIDIA host. A fresh CUDA install (including the studio Docker build,
which runs `bash install.sh --local`) takes its torch from install.sh's
TORCH_CONSTRAINT, which was still capped at torch>=2.4,<2.11.0, so cu12x/cu13x
resolved torch 2.10.x and the venv landed a minor behind the torch 2.11.0 base
image.
Extend the existing `case "$TORCH_INDEX_URL"` block (which already relaxes
rocm7.2) with a `*/cu[0-9]*` branch that widens the ceiling to <2.12.0, keeping
the >=2.4 floor so an older CUDA index (e.g. cu118) that tops out below 2.11
still resolves. The CPU wheel and older ROCm tags stay on <2.11.0 (the glob
does not match /cpu). torchvision/torchaudio are bare on this install line and
resolve their compatible companions via wheel metadata, matching the rocm7.2
pattern.
Add behavioral tests (Python + shell) exercising the case block: cu118/124/126/
128/130 widen to <2.12.0, rocm7.2 stays 2.11.x, and /cpu plus older ROCm keep
the default <2.11.0.
* install.sh: key the CUDA torch widening off the index leaf, not the full URL
The `*/cu[0-9]*` glob matched a `cu` segment anywhere in TORCH_INDEX_URL,
so a custom UNSLOTH_PYTORCH_MIRROR whose base path contains e.g. cu128 but whose
final leaf is cpu or an older ROCm tag would still widen TORCH_CONSTRAINT to
<2.12.0, contradicting the block's own comment and letting a CPU / older-ROCm
mirror resolve torch 2.11.x. Match on _torch_index_leaf (the final path segment
the backend classification just above already computes) so only a real cu*/
rocm7.2 leaf is affected; cpu and older ROCm keep the default <2.11.0. Update
the Python + shell tests to mirror the leaf-anchored case and add regression
cases for a mirror base that contains cu128 but resolves to a cpu / rocm7.1 leaf.
* install: freeze the torch trio during the with-deps unsloth installs
Released unsloth wheels can pin an older torch than Step 1 installed
(unsloth 2026.7.2 declares torch<2.11.0), so the with-deps resolve from
PyPI silently downgrades the pinned +cuXXX torch trio to PyPI's default
wheel. The flavor guard cannot catch every such swap: PyPI's torch 2.10
default is itself cu128-flavored, so the cuXXX tag comparison still
matches while the version silently drops. Freeze the just-installed trio
with uv --overrides (overrides replace dependency requirements during
resolution), keeping torch 2.11.0+cuXXX in place while unsloth's other
dependencies resolve normally. Verified on the cu128 path: without the
override torch drops 2.11.0+cu128 -> 2.10.0; with it the trio survives
and unsloth 2026.7.2 + unsloth-zoo install cleanly.
* install: fold UV_OVERRIDE env files into the torch-trio overrides file
The CLI --overrides flag is the command-line form of UV_OVERRIDE, so
passing it replaced any overrides file already exported for the process;
macOS arm64 exports UV_OVERRIDE=overrides-darwin-arm64.txt for the same
generic install path and would have lost those pins. Concatenate any
UV_OVERRIDE files into the temp trio file so both keep applying.
* install: extend the torch-trio overrides guard to migrated installs
Four follow-ups to the Step-2 --overrides guard, all empirically verified:
1. The migrated-environment with-deps unsloth install resolved
unsloth>=2026.7.2 (which pins torch<2.11.0) without the overrides file,
so a migrated CUDA venv on torch 2.11 was silently downgraded -- the
exact bug this branch fixes on the fresh path. The overrides build is
now a function (_build_unsloth_torch_overrides, reading the trio
installed at call time) invoked by both with-deps paths; the migrated
no-torch path installs --no-deps and stays unguarded.
2. The overrides temp file is now cleaned by the EXIT trap (same pattern
as _UV_OVERRIDE_TMPDIR, pre-initialized empty so an inherited value can
never reach the trap's rm); previously any Step-2 failure leaked it.
3. Folding UV_OVERRIDE files used cat, which joins the last requirement of
a file lacking a trailing newline onto the next file's first requirement
(reproduced: idna==3.10certifi==2025.1.31 makes uv fail parsing).
4. Inherited torch/torchvision/torchaudio override lines are now filtered
out when folding: uv intersects duplicate overrides rather than
last-wins (verified on uv 0.10.12: direct conflict is unsatisfiable,
transitive conflict silently backtracks), so a conflicting inherited
trio pin would break the resolve the generated exact pins protect.
Both 3 and 4 are handled by a single newline-terminating awk filter
that preserves non-trio overrides (torchmetrics, torchao, ...).
test_unsloth_torch_override.sh extended: migrated-path coverage, trap
assertion, and a functional fold test (14 checks).
* installer: tighten comments
* install: keep the existing torch release when re-running the installer
Re-running `curl -fsSL https://unsloth.ai/install.sh | sh` over an existing
install rebuilds the venv for clean state, which silently moved users to the
newest torch in range (2.10 -> 2.11 once the constraint widened). A torch the
user already validated must survive an unsloth update.
Before the old venv is moved aside for rollback, its torch version is probed
(last stdout line only, so sitecustomize noise cannot corrupt it). After the
index leaf is chosen, _previous_torch_pin turns that version into a
torch==X.Y.Z pin, but only when it cannot do harm:
- cu*/cpu leaves only; rocm leaves keep their floors (rocm7.2 must land 2.11
for the Strix _grouped_mm fix) and the Radeon wheel-matching path is
untouched.
- The wheel's flavor tag must match the freshly chosen leaf, so a flavor
change (cpu -> cuda, cu126 -> cu130) still installs the correct new build.
- The base must look like a release, so probe noise never becomes a pin.
- UNSLOTH_TORCH_UPGRADE=1 opts out and restores the old always-newest
behavior; the substep line advertises it.
The supported range is kept in _PREV_FALLBACK_CONSTRAINT: if the exact
release is not resolvable from the chosen index (custom mirrors prune old
wheels), the install warns and falls back to the newest supported release
instead of failing the whole run. The later flavor-mismatch repair reuses
TORCH_CONSTRAINT, so a mid-install clobber is repaired back to the kept
release rather than the newest one.
Verified end to end: a venv seeded with torch 2.10.0+cu130 re-run through the
full installer finishes with torch 2.10.0+cu130 (previously 2.11.0+cu130).
Tests: tests/sh/test_previous_torch_pin.sh covers keep/flavor-change/rocm/
noise/opt-out plus wiring (probe ordering before venv replacement, fallback
present, SKIP_TORCH gate).
* install: constrain kept torch pins to the supported window
Review caught that _previous_torch_pin pinned the previous venv's torch on
flavor match alone, so a release outside the installer's active range (a
2.3.x manual install below the >=2.4 floor, or a 2.12.x manual upgrade above
the ceiling) replaced the bounds computed just above it and a rerun kept a
torch the installer otherwise deliberately excludes.
New _torch_release_in_window checks the probed base against the active
TORCH_CONSTRAINT ("torch>=A.B[,
---
install.sh | 168 +++++++++++++++++-
studio/backend/tests/test_cuda_torch_spec.py | 73 ++++++++
.../test_tokenizers_and_torch_constraint.py | 80 +++++++++
tests/sh/test_previous_torch_pin.sh | 101 +++++++++++
tests/sh/test_torch_constraint.sh | 14 ++
tests/sh/test_unsloth_torch_override.sh | 131 ++++++++++++++
6 files changed, 558 insertions(+), 9 deletions(-)
create mode 100644 studio/backend/tests/test_cuda_torch_spec.py
create mode 100644 tests/sh/test_previous_torch_pin.sh
create mode 100644 tests/sh/test_unsloth_torch_override.sh
diff --git a/install.sh b/install.sh
index 5972379d26..6076721540 100755
--- a/install.sh
+++ b/install.sh
@@ -472,11 +472,13 @@ _on_install_exit() {
_restore_studio_venv_replacement
fi
[ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
+ [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true
exit "$_status"
}
-# Empty so an inherited value can never reach the trap's rm; only a temp dir
-# this script creates below (Apple Silicon, spaced path) is ever removed.
+# Empty so an inherited value never reaches the trap's rm; only temp paths this
+# script creates below (spaced-path dir, torch-trio overrides) are removed.
_UV_OVERRIDE_TMPDIR=""
+_UNSLOTH_TORCH_OVERRIDES=""
trap _on_install_exit EXIT
# ── Helper: download a URL to a file (supports curl and wget) ──
@@ -1821,6 +1823,8 @@ tauri_log "STEP" "Creating virtual environment"
mkdir -p "$STUDIO_HOME"
_MIGRATED=false
+# Empty so an inherited value can never masquerade as a probed torch version.
+_PREV_TORCH_VER=""
if [ -x "$VENV_DIR/bin/python" ]; then
# why: matching guard to the .venv branch below -- in env-mode
@@ -1838,6 +1842,12 @@ if [ -x "$VENV_DIR/bin/python" ]; then
echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2
exit 1
fi
+ # Record the existing venv's torch BEFORE the replacement moves it aside: a re-run
+ # rebuilds the venv for clean state, but must keep the torch release the user
+ # already has (see _previous_torch_pin below). Last line only: sitecustomize or
+ # import-hook noise on stdout must not corrupt the version.
+ _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \
+ "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true)
# New layout already exists — replace only after preserving rollback copy.
substep "preserving existing environment for rollback..."
_start_studio_venv_replacement "$VENV_DIR"
@@ -2187,6 +2197,68 @@ _torch_flavor_tag() {
esac
}
+# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2
+# ("torch>=A.B[.C],="*",<"*) ;;
+ *) echo "no"; return ;;
+ esac
+ _trw_floor="${_trw_con#torch>=}"; _trw_floor="${_trw_floor%%,*}"
+ _trw_ceil="${_trw_con##*,<}"
+ _v_maj="${1%%.*}"; _v_rest="${1#*.}"; _v_min="${_v_rest%%.*}"
+ _f_maj="${_trw_floor%%.*}"; _f_rest="${_trw_floor#*.}"; _f_min="${_f_rest%%.*}"
+ _c_maj="${_trw_ceil%%.*}"; _c_rest="${_trw_ceil#*.}"; _c_min="${_c_rest%%.*}"
+ for _trw_n in "$_v_maj" "$_v_min" "$_f_maj" "$_f_min" "$_c_maj" "$_c_min"; do
+ case "$_trw_n" in ''|*[!0-9]*) echo "no"; return ;; esac
+ done
+ if [ "$_v_maj" -gt "$_f_maj" ] || { [ "$_v_maj" -eq "$_f_maj" ] && [ "$_v_min" -ge "$_f_min" ]; }; then
+ if [ "$_v_maj" -lt "$_c_maj" ] || { [ "$_v_maj" -eq "$_c_maj" ] && [ "$_v_min" -lt "$_c_min" ]; }; then
+ echo "yes"
+ return
+ fi
+ fi
+ echo "no"
+}
+
+# Whether a re-run should keep the previous venv's torch: echo "torch==X.Y.Z" when the
+# probed previous version ($1) has a flavor tag matching the freshly chosen cu*/cpu index
+# leaf ($2) AND sits inside the active constraint window ($3), else "". Re-running
+# `curl | sh` rebuilds the venv for clean state, but a healthy torch the user already
+# validated must not be silently moved to a newer release (2.10 -> 2.11); a flavor
+# change (cpu <-> cuda, cu126 -> cu130) still installs the correct new build, rocm
+# leaves keep their floors (rocm7.2 must land 2.11 for the Strix _grouped_mm fix), and
+# a release outside the window (2.3.x manual install, 2.12.x manual upgrade) is never
+# kept: the installer's own bounds win. Opt out with UNSLOTH_TORCH_UPGRADE=1 to get
+# the newest release.
+_previous_torch_pin() {
+ _ptp_ver="$1"
+ _ptp_leaf="$2"
+ _ptp_con="$3"
+ [ -n "$_ptp_ver" ] || { echo ""; return; }
+ [ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; }
+ case "$_ptp_leaf" in
+ cu[0-9]*|cpu) ;;
+ *) echo ""; return ;;
+ esac
+ _ptp_base="${_ptp_ver%%+*}"
+ # The base must look like a release (probe noise / garbage must never become a pin).
+ case "$_ptp_base" in
+ [0-9]*.[0-9]*) ;;
+ *) echo ""; return ;;
+ esac
+ [ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; }
+ if [ "$(_torch_flavor_tag "$_ptp_ver")" = "$_ptp_leaf" ]; then
+ echo "torch==$_ptp_base"
+ else
+ echo ""
+ fi
+}
+
# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* ->
# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops.
_expected_torch_flavor_tag() {
@@ -2478,12 +2550,32 @@ case "$_torch_index_leaf" in
*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
esac
-# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it.
-# All other ROCm tags and CUDA stay within <2.11.0.
-case "$TORCH_INDEX_URL" in
- */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
+# rocm7.2 and the CUDA cu12x/cu13x indexes now ship torch 2.11.x, so widen the
+# ceiling to <2.12.0 (matches the base image and _CUDA_TORCH_PKG_SPEC in
+# studio/install_python_stack.py). Keep the >=2.4 floor so an older CUDA index
+# (e.g. cu118) still resolves. Match on _torch_index_leaf, not the full URL, so
+# a mirror whose base path contains cu*/rocm7.2 but resolves to a cpu/older-rocm
+# leaf keeps the default <2.11.0.
+case "$_torch_index_leaf" in
+ rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
+ cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
esac
+# Re-run over an existing install: keep the previous venv's torch release instead of
+# resolving the newest in range. The range stays in _PREV_FALLBACK_CONSTRAINT so the
+# install can fall back when the exact release is not on the chosen index (custom
+# mirrors may prune old wheels). Skipped for --no-torch (no previous probe runs).
+_PREV_TORCH_PIN=""
+_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT"
+if [ "$SKIP_TORCH" = false ]; then
+ _prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$_torch_index_leaf" "$TORCH_CONSTRAINT")
+ if [ -n "$_prev_pin" ]; then
+ _PREV_TORCH_PIN="$_prev_pin"
+ TORCH_CONSTRAINT="$_prev_pin"
+ substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)"
+ fi
+fi
+
# Auto-detect GPU for AMD ROCm based
# get_torch_index_url must have chosen */rocm*
# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
@@ -2705,6 +2797,43 @@ esac
# ── Install unsloth directly into the venv (no activation needed) ──
tauri_log "STEP" "Installing PyTorch"
_VENV_PY="$VENV_DIR/bin/python"
+
+# A released unsloth wheel can pin an older torch (unsloth 2026.7.2 declares
+# torch<2.11.0); a with-deps PyPI resolve then downgrades the whole trio,
+# swapping the pinned +cuXXX/+rocm build for PyPI's default. The flavor guard
+# below misses this (PyPI's torch 2.10 default is itself cu128-flavored), so
+# freeze the trio via uv --overrides (overrides replace dependency requirements
+# during resolution) while unsloth's other deps resolve normally. Sets
+# _UNSLOTH_TORCH_OVERRIDES from the trio in the venv; every with-deps unsloth
+# install (migrated and fresh) must call this before resolving and rm it after.
+_build_unsloth_torch_overrides() {
+ _UNSLOTH_TORCH_OVERRIDES=""
+ [ "$SKIP_TORCH" = false ] || return 0
+ _torch_trio_pins=$("$_VENV_PY" -c "
+from importlib.metadata import version, PackageNotFoundError
+for _p in ('torch', 'torchvision', 'torchaudio'):
+ try:
+ print(_p + '==' + version(_p))
+ except PackageNotFoundError:
+ pass
+" 2>/dev/null) || _torch_trio_pins=""
+ case "$_torch_trio_pins" in
+ torch==*)
+ _UNSLOTH_TORCH_OVERRIDES=$(mktemp)
+ printf '%s\n' "$_torch_trio_pins" > "$_UNSLOTH_TORCH_OVERRIDES"
+ # The CLI --overrides flag replaces any UV_OVERRIDE env file (same
+ # uv setting; macOS arm64 exports one here), so fold its pins in.
+ # awk, not cat: it drops inherited torch-trio lines (uv intersects
+ # duplicate overrides, so a conflicting pin would make resolution
+ # unsatisfiable) and newline-terminates the last line so an
+ # unterminated file cannot join two requirements into one.
+ for _ov_file in ${UV_OVERRIDE:-}; do
+ [ -f "$_ov_file" ] && awk '!/^[[:space:]]*torch(vision|audio)?([[:space:]<>=!~;@[]|$)/' "$_ov_file" >> "$_UNSLOTH_TORCH_OVERRIDES"
+ done
+ ;;
+ esac
+}
+
if [ "$_MIGRATED" = true ]; then
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
@@ -2729,9 +2858,13 @@ if [ "$_MIGRATED" = true ]; then
else
# Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
+ _build_unsloth_torch_overrides
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
+ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-}
+ [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES"
+ _UNSLOTH_TORCH_OVERRIDES=""
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@@ -2913,8 +3046,20 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
- run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL"
+ if [ -n "$_PREV_TORCH_PIN" ]; then
+ # Kept previous release: fall back to the supported range if the exact
+ # release is not resolvable from the chosen index (pruned mirror).
+ if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
+ --default-index "$TORCH_INDEX_URL"; then
+ substep "[WARN] $_PREV_TORCH_PIN is not installable from $TORCH_INDEX_URL -- installing the newest supported release instead" "$C_WARN"
+ TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT"
+ run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
+ --default-index "$TORCH_INDEX_URL"
+ fi
+ else
+ run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
+ --default-index "$TORCH_INDEX_URL"
+ fi
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
@@ -2927,9 +3072,10 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
;;
esac
fi
- # Fresh: Step 2 - install unsloth, preserving pre-installed torch
+ # Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
+ _build_unsloth_torch_overrides
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
@@ -2953,6 +3099,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
+ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
@@ -2962,8 +3109,11 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
else
run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
+ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-}
fi
+ [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES"
+ _UNSLOTH_TORCH_OVERRIDES=""
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
if [ "$SKIP_TORCH" = false ]; then
diff --git a/studio/backend/tests/test_cuda_torch_spec.py b/studio/backend/tests/test_cuda_torch_spec.py
new file mode 100644
index 0000000000..928cef787e
--- /dev/null
+++ b/studio/backend/tests/test_cuda_torch_spec.py
@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Tests for _CUDA_TORCH_PKG_SPEC in install_python_stack.py.
+
+The CUDA repair path installs the torch trio from an exclusive --index-url (no
+PyPI fallback), so these pinned ranges decide which torch the venv gets. The
+upper bound is locked to the 2.11.x family to match the base image and rocm7.2
+spec and to keep the companions off a torch-2.12 wheel that would ABI-mismatch.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+from packaging.requirements import Requirement
+
+# install_python_stack.py lives at repo_root/studio/install_python_stack.py
+_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py"
+
+
+def _load_module(monkeypatch):
+ """(Re-)import and return install_python_stack (mirrors test_torchao_select)."""
+ sys.modules.pop("install_python_stack", None)
+ monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent))
+ import install_python_stack
+
+ return install_python_stack
+
+
+def _spec_of(pkg_spec: str):
+ """Parse 'torch>=2.4,<2.12.0' into a packaging SpecifierSet."""
+ return Requirement(pkg_spec).specifier
+
+
+@pytest.mark.parametrize(
+ "index, allowed, rejected",
+ [
+ # torch: 2.11.x allowed (matches base image); 2.12.x excluded.
+ (0, ["2.11.0", "2.11.2", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0", "1.13.1"]),
+ # torchvision: 0.26.x (torch 2.11 companion) allowed; 0.27.x (torch 2.12) out.
+ (1, ["0.26.0", "0.26.1", "0.19.0"], ["0.27.0", "0.18.0"]),
+ # torchaudio: same 2.11.x window as torch.
+ (2, ["2.11.0", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0"]),
+ ],
+)
+def test_cuda_spec_bounds(monkeypatch, index, allowed, rejected):
+ mod = _load_module(monkeypatch)
+ spec = _spec_of(mod._CUDA_TORCH_PKG_SPEC[index])
+ for v in allowed:
+ assert spec.contains(v, prereleases = True), f"{v} should satisfy {spec}"
+ for v in rejected:
+ assert not spec.contains(v, prereleases = True), f"{v} should not satisfy {spec}"
+
+
+def test_cuda_spec_matches_rocm72_upper_bound(monkeypatch):
+ """CUDA and rocm7.2 target the same torch 2.11.x family, so their upper
+ bounds must stay in lockstep (bump both together at 2.12.x)."""
+ mod = _load_module(monkeypatch)
+ rocm72 = mod._ROCM_TORCH_PKG_SPECS["rocm7.2"]
+
+ def _upper(pkg_spec: str) -> str:
+ for clause in _spec_of(pkg_spec):
+ if clause.operator == "<":
+ return clause.version
+ raise AssertionError(f"no upper bound in {pkg_spec!r}")
+
+ for cuda_pkg, rocm_pkg in zip(mod._CUDA_TORCH_PKG_SPEC, rocm72, strict = True):
+ assert _upper(cuda_pkg) == _upper(
+ rocm_pkg
+ ), f"CUDA {cuda_pkg!r} upper bound must match rocm7.2 {rocm_pkg!r}"
diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py
index 4322f0c7d6..c58808689b 100644
--- a/tests/python/test_tokenizers_and_torch_constraint.py
+++ b/tests/python/test_tokenizers_and_torch_constraint.py
@@ -69,6 +69,21 @@ class TestStructuralTorchConstraint:
def test_tightened_assignment_exists(self):
assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh
+ def test_cuda_constraint_widened_to_2_12(self):
+ """A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x
+ land torch 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC);
+ without it cu128/cu130 resolves torch 2.10.x."""
+ assert 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' in self._sh
+
+ def test_cuda_case_widens_via_index_leaf(self):
+ """The cu* branch of the _torch_index_leaf case sets the widened
+ constraint (parallel to rocm7.2), anchored on the leaf."""
+ m = re.search(
+ r'cu\[0-9\]\*\)\s*TORCH_CONSTRAINT="torch>=2\.4,<2\.12\.0"',
+ self._sh,
+ )
+ assert m is not None, "CUDA (cu*) TORCH_CONSTRAINT widening case not found"
+
def test_variable_used_in_pip_install(self):
"""$TORCH_CONSTRAINT must appear in a uv pip install line."""
assert '"$TORCH_CONSTRAINT"' in self._sh
@@ -384,6 +399,71 @@ class TestTorchConstraintShell:
logged = log_file.read_text()
assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}"
+ # Mirrors the _torch_index_leaf case in install.sh: rocm7.2 -> 2.11.x floor,
+ # CUDA -> widened <2.12.0 ceiling, else (CPU/older ROCm) -> default. Anchored
+ # on the final path segment, so a mirror base path containing cu*/rocm7.2 but
+ # ending in a cpu/older-rocm leaf keeps the default.
+ _INDEX_SNIPPET = textwrap.dedent(r"""
+ #!/bin/bash
+ set -e
+ TORCH_INDEX_URL="{index_url}"
+ TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+ _torch_index_leaf="${TORCH_INDEX_URL%/}"
+ _torch_index_leaf="${_torch_index_leaf##*/}"
+ case "$_torch_index_leaf" in
+ rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
+ cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
+ esac
+ echo "$TORCH_CONSTRAINT"
+ """).strip()
+
+ def _resolve_index(self, tmp_path: pathlib.Path, index_url: str) -> str:
+ script_file = tmp_path / "index_snippet.sh"
+ script_file.write_text(self._INDEX_SNIPPET.replace("{index_url}", index_url))
+ script_file.chmod(0o755)
+ result = subprocess.run(
+ ["bash", str(script_file)],
+ capture_output = True,
+ text = True,
+ timeout = 10,
+ )
+ assert result.returncode == 0, f"Script failed: {result.stderr}"
+ return result.stdout.strip()
+
+ @pytest.mark.parametrize("leaf", ["cu118", "cu124", "cu126", "cu128", "cu130"])
+ def test_cuda_index_widens_to_2_12(self, tmp_path, leaf):
+ url = f"https://download.pytorch.org/whl/{leaf}"
+ assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
+
+ def test_rocm72_index_uses_211_floor(self, tmp_path):
+ url = "https://download.pytorch.org/whl/rocm7.2"
+ assert self._resolve_index(tmp_path, url) == "torch>=2.11.0,<2.12.0"
+
+ def test_cpu_index_keeps_default(self, tmp_path):
+ # /cpu must NOT match the */cu[0-9]* branch.
+ url = "https://download.pytorch.org/whl/cpu"
+ assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
+
+ def test_older_rocm_index_keeps_default(self, tmp_path):
+ url = "https://download.pytorch.org/whl/rocm7.1"
+ assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
+
+ def test_cuda_index_custom_mirror_widens(self, tmp_path):
+ url = "https://internal.example.com/pytorch/cu128"
+ assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
+
+ @pytest.mark.parametrize(
+ "url",
+ [
+ "https://internal.example.com/pytorch/cu128/cpu",
+ "https://internal.example.com/cu128/whl/rocm7.1",
+ ],
+ )
+ def test_cuda_in_mirror_path_but_noncuda_leaf_keeps_default(self, tmp_path, url):
+ # A cu128 in the mirror base path must not widen when the leaf is cpu /
+ # older ROCm: the case anchors on _torch_index_leaf, not the whole URL.
+ assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
+
# Group 3 -- E2E tokenizers fix (requires network, ~2-5 min)
@pytest.mark.e2e
diff --git a/tests/sh/test_previous_torch_pin.sh b/tests/sh/test_previous_torch_pin.sh
new file mode 100644
index 0000000000..253ede8a27
--- /dev/null
+++ b/tests/sh/test_previous_torch_pin.sh
@@ -0,0 +1,101 @@
+#!/bin/bash
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+# Unit tests for install.sh's _previous_torch_pin, which keeps the previous
+# venv's torch release on a re-run (curl | sh over an existing install) instead
+# of silently moving the user to a newer release. Helpers are extracted from
+# install.sh and sourced.
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+INSTALL_SH="$SCRIPT_DIR/../../install.sh"
+PASS=0
+FAIL=0
+
+# Extract _previous_torch_pin and its dependencies _torch_flavor_tag and
+# _torch_release_in_window.
+_FUNC_FILE=$(mktemp)
+{
+ sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
+ echo ""
+ sed -n '/^_torch_release_in_window()/,/^}/p' "$INSTALL_SH"
+ echo ""
+ sed -n '/^_previous_torch_pin()/,/^}/p' "$INSTALL_SH"
+} > "$_FUNC_FILE"
+# shellcheck disable=SC1090
+. "$_FUNC_FILE"
+rm -f "$_FUNC_FILE"
+
+assert_eq() {
+ _label="$1"; _expected="$2"; _actual="$3"
+ if [ "$_actual" = "$_expected" ]; then
+ echo " PASS: $_label"; PASS=$((PASS + 1))
+ else
+ echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1))
+ fi
+}
+
+unset UNSLOTH_TORCH_UPGRADE
+
+echo "=== _previous_torch_pin: matching flavor keeps the release ==="
+assert_eq "cu126 wheel on cu126 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "cu130 wheel on cu130 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
+assert_eq "cpu wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cpu' 'cpu' 'torch>=2.4,<2.12.0')"
+assert_eq "untagged wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0' 'cpu' 'torch>=2.4,<2.12.0')"
+assert_eq "local suffix stripped" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+cu128' 'cu128' 'torch>=2.4,<2.12.0')"
+
+echo "=== _previous_torch_pin: flavor change installs the new build ==="
+assert_eq "cu126 wheel on cu130 leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu130' 'torch>=2.4,<2.12.0')"
+assert_eq "cpu wheel on cu126 leaf" "" "$(_previous_torch_pin '2.10.0+cpu' 'cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "cu126 wheel on cpu leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cpu' 'torch>=2.4,<2.12.0')"
+
+echo "=== _previous_torch_pin: rocm and unknown leaves never pin ==="
+assert_eq "rocm7.2 leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'rocm7.2' 'torch>=2.4,<2.12.0')"
+assert_eq "gfx leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'gfx120X-all' 'torch>=2.4,<2.12.0')"
+assert_eq "unknown mirror leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'simple' 'torch>=2.4,<2.12.0')"
+
+echo "=== _previous_torch_pin: probe noise never becomes a pin ==="
+assert_eq "empty version" "" "$(_previous_torch_pin '' 'cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "garbage version" "" "$(_previous_torch_pin 'not-a-version' 'cpu' 'torch>=2.4,<2.12.0')"
+assert_eq "traceback fragment" "" "$(_previous_torch_pin "ModuleNotFoundError: No module named 'torch'" 'cpu' 'torch>=2.4,<2.12.0')"
+
+echo "=== _previous_torch_pin: out-of-window releases never pin ==="
+assert_eq "2.3.x below the cu floor" "" "$(_previous_torch_pin '2.3.1+cu118' 'cu118' 'torch>=2.4,<2.12.0')"
+assert_eq "2.12.x above the cu ceiling" "" "$(_previous_torch_pin '2.12.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
+assert_eq "floor boundary 2.4.0 kept" "torch==2.4.0" "$(_previous_torch_pin '2.4.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "ceiling-adjacent 2.11.x kept" "torch==2.11.1" "$(_previous_torch_pin '2.11.1+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
+assert_eq "cpu window excludes 2.11.x" "" "$(_previous_torch_pin '2.11.0+cpu' 'cpu' 'torch>=2.4,<2.11.0')"
+assert_eq "mac floor excludes 2.5.x" "" "$(_previous_torch_pin '2.5.1' 'cpu' 'torch>=2.6,<2.11.0')"
+assert_eq "malformed window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch')"
+assert_eq "empty window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' '')"
+
+echo "=== _torch_release_in_window ==="
+assert_eq "in window" "yes" "$(_torch_release_in_window '2.10.0' 'torch>=2.4,<2.12.0')"
+assert_eq "at floor" "yes" "$(_torch_release_in_window '2.4.0' 'torch>=2.4,<2.12.0')"
+assert_eq "below floor" "no" "$(_torch_release_in_window '2.3.1' 'torch>=2.4,<2.12.0')"
+assert_eq "at ceiling" "no" "$(_torch_release_in_window '2.12.0' 'torch>=2.4,<2.12.0')"
+assert_eq "next major" "no" "$(_torch_release_in_window '3.0.0' 'torch>=2.4,<2.12.0')"
+assert_eq "patch-level floor" "yes" "$(_torch_release_in_window '2.11.5' 'torch>=2.11.0,<2.12.0')"
+assert_eq "no ceiling -> no" "no" "$(_torch_release_in_window '2.10.0' 'torch>=2.4')"
+assert_eq "garbage minor -> no" "no" "$(_torch_release_in_window '2.x' 'torch>=2.4,<2.12.0')"
+
+echo "=== _previous_torch_pin: UNSLOTH_TORCH_UPGRADE=1 opts out ==="
+assert_eq "upgrade env set" "" "$(UNSLOTH_TORCH_UPGRADE=1 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "upgrade env 0" "torch==2.10.0" "$(UNSLOTH_TORCH_UPGRADE=0 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
+
+echo "=== install.sh wiring ==="
+# The probe must run against the OLD venv, before it is moved aside for rollback.
+_probe_line=$(grep -n '_PREV_TORCH_VER=\$(' "$INSTALL_SH" | head -1 | cut -d: -f1)
+_move_line=$(grep -n '_start_studio_venv_replacement "\$VENV_DIR"' "$INSTALL_SH" | head -1 | cut -d: -f1)
+assert_eq "probe exists" "yes" "$([ -n "$_probe_line" ] && echo yes)"
+assert_eq "probe before venv replacement" "yes" "$([ -n "$_probe_line" ] && [ -n "$_move_line" ] && [ "$_probe_line" -lt "$_move_line" ] && echo yes)"
+# A kept release that vanished from the index must fall back to the supported range.
+assert_eq "resolve-failure fallback wired" "yes" "$(grep -q 'TORCH_CONSTRAINT="\$_PREV_FALLBACK_CONSTRAINT"' "$INSTALL_SH" && echo yes)"
+assert_eq "pin gated on SKIP_TORCH" "yes" "$(grep -q 'if \[ "\$SKIP_TORCH" = false \]; then' "$INSTALL_SH" && echo yes)"
+
+echo ""
+if [ "$FAIL" -gt 0 ]; then
+ echo "$FAIL check(s) FAILED"
+ exit 1
+fi
+echo "All $PASS checks passed"
diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh
index 293a709360..d60dfc9f90 100644
--- a/tests/sh/test_torch_constraint.sh
+++ b/tests/sh/test_torch_constraint.sh
@@ -108,6 +108,20 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var"
_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded"
+# A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x land torch
+# 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC).
+_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true)
+assert_eq "CUDA TORCH_CONSTRAINT widened to <2.12.0" "1" "$_cuda_widen"
+
+# Widening keys off the final leaf (_torch_index_leaf), not the full URL, so a
+# mirror base path with cu*/rocm7.2 but a cpu/older-rocm leaf is not mis-widened.
+_cuda_case=$(grep -c 'cu\[0-9\]\*)' "$INSTALL_SH" || true)
+_has_cuda_case=$([ "$_cuda_case" -ge 1 ] && echo "yes" || echo "no")
+assert_eq "cu* index case adjusts TORCH_CONSTRAINT" "yes" "$_has_cuda_case"
+_leaf_case=$(grep -c 'case "\$_torch_index_leaf" in' "$INSTALL_SH" || true)
+_has_leaf_constraint=$([ "$_leaf_case" -ge 2 ] && echo "yes" || echo "no")
+assert_eq "constraint case anchors on _torch_index_leaf" "yes" "$_has_leaf_constraint"
+
echo ""
echo "=== Structural: tokenizers in no-torch-runtime.txt ==="
diff --git a/tests/sh/test_unsloth_torch_override.sh b/tests/sh/test_unsloth_torch_override.sh
new file mode 100644
index 0000000000..7e8e3f5b5b
--- /dev/null
+++ b/tests/sh/test_unsloth_torch_override.sh
@@ -0,0 +1,131 @@
+#!/bin/bash
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+# Tests for the torch-trio --overrides guard on the Step-2 unsloth installs in
+# install.sh. A released unsloth wheel can pin an older torch (2026.7.2 declares
+# torch<2.11.0); without the overrides file a with-deps PyPI resolve downgrades
+# the trio Step 1 installed, and the flavor guard misses it (PyPI's torch 2.10
+# default is itself cu128-flavored). Same assertion pattern as test_torch_constraint.sh.
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+INSTALL_SH="$SCRIPT_DIR/../../install.sh"
+PASS=0
+FAIL=0
+
+assert_true() {
+ _label="$1"; _ok="$2"
+ if [ "$_ok" = "0" ]; then
+ echo " PASS: $_label"
+ PASS=$((PASS + 1))
+ else
+ echo " FAIL: $_label"
+ FAIL=$((FAIL + 1))
+ fi
+}
+
+echo "=== test_unsloth_torch_override ==="
+
+# 1. Every with-deps unsloth install carries the overrides expansion (local,
+# generic, migrated); the --no-deps no-torch paths need no guard.
+_local_block=$(grep -A2 '"install unsloth (local)"' "$INSTALL_SH")
+printf '%s' "$_local_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"'
+assert_true "local (with-deps) unsloth install passes --overrides" "$?"
+
+_generic_block=$(grep -A2 '"install unsloth" uv pip install' "$INSTALL_SH")
+printf '%s' "$_generic_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"'
+assert_true "generic (with-deps) unsloth install passes --overrides" "$?"
+
+_migrated_block=$(grep -A3 '"install unsloth (migrated)"' "$INSTALL_SH")
+printf '%s' "$_migrated_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"'
+assert_true "migrated (with-deps) unsloth install passes --overrides" "$?"
+
+_no_torch_block=$(grep -A2 '"install unsloth (no-torch)"' "$INSTALL_SH")
+if printf '%s' "$_no_torch_block" | grep -q -- '--overrides'; then _rc=1; else _rc=0; fi
+assert_true "no-torch (--no-deps) unsloth install has no overrides" "$_rc"
+
+_migrated_nt_block=$(grep -A2 '"install unsloth (migrated no-torch)"' "$INSTALL_SH")
+if printf '%s' "$_migrated_nt_block" | grep -q -- '--overrides'; then _rc=1; else _rc=0; fi
+assert_true "migrated no-torch (--no-deps) unsloth install has no overrides" "$_rc"
+
+# 2. The overrides file is only built when SKIP_TORCH=false.
+grep -B2 '_torch_trio_pins=\$(' "$INSTALL_SH" | grep -q 'SKIP_TORCH" = false'
+assert_true "overrides file build is gated on SKIP_TORCH=false" "$?"
+
+# 3. The pin-collection snippet emits exact ==pins for the installed trio (run
+# the embedded python against this test's interpreter).
+_snippet=$(sed -n '/_torch_trio_pins=\$("\$_VENV_PY" -c "/,/^" 2>\/dev\/null)/p' "$INSTALL_SH" \
+ | sed '1s/.*-c "//' | sed '$d')
+_out=$(python3 -c "$_snippet" 2>&1) || true
+# torch may or may not be importable on the test host; the snippet must not
+# crash and every line it does emit must be an exact pkg==version pin.
+if [ -n "$_out" ]; then
+ printf '%s\n' "$_out" | grep -vqE '^(torch|torchvision|torchaudio)==.+$' && _rc=1 || _rc=0
+else
+ _rc=0
+fi
+assert_true "pin snippet emits only exact trio ==pins (or nothing)" "$_rc"
+
+# 4. The temp overrides file is cleaned up after Step 2.
+grep -q 'rm -f "\$_UNSLOTH_TORCH_OVERRIDES"' "$INSTALL_SH"
+assert_true "overrides temp file is removed after the unsloth installs" "$?"
+
+# 5. Any UV_OVERRIDE env file is folded in (the CLI --overrides flag would
+# otherwise replace it, dropping e.g. the macOS arm64 darwin overrides).
+grep -q 'for _ov_file in \${UV_OVERRIDE:-}' "$INSTALL_SH"
+assert_true "UV_OVERRIDE env files are merged into the overrides file" "$?"
+
+# 6. The EXIT trap also removes the overrides file, so a failed Step 2 (set -e
+# fires before the normal-path rm) cannot leak it.
+sed -n '/_on_install_exit() {/,/^}/p' "$INSTALL_SH" \
+ | grep -q 'rm -f "\$_UNSLOTH_TORCH_OVERRIDES"'
+assert_true "EXIT trap removes the overrides temp file on failure" "$?"
+
+# 7. The UV_OVERRIDE fold filters inherited files instead of cat-ing them (run
+# the extracted awk program on sample files): (a) inherited torch-trio lines
+# are dropped so the generated exact pins win (uv intersects duplicates);
+# (b) every line is newline-terminated so an unterminated file cannot join
+# two requirements into one.
+_awk_prog=$(sed -n "s/.*awk '\(.*\)' \"\$_ov_file\".*/\1/p" "$INSTALL_SH")
+[ -n "$_awk_prog" ]
+assert_true "UV_OVERRIDE fold uses the trio-filtering awk program" "$?"
+
+_ov_dir=$(mktemp -d)
+printf '%s' 'transformers>=4.57.6' > "$_ov_dir/ov1.txt" # no trailing newline
+cat > "$_ov_dir/ov2.txt" <<'EOF'
+# comment survives
+torch<2.11.0
+torchvision==0.25.0
+torchaudio!=2.11.0
+torchmetrics==1.0
+anyio<4.14.0
+EOF
+_merged="$_ov_dir/merged.txt"
+printf '%s\n' 'torch==2.11.0+cu128' > "$_merged"
+for _f in "$_ov_dir/ov1.txt" "$_ov_dir/ov2.txt"; do
+ awk "$_awk_prog" "$_f" >> "$_merged"
+done
+
+grep -qx 'transformers>=4.57.6' "$_merged"
+assert_true "no-trailing-newline override stays a separate requirement line" "$?"
+
+if grep -qx 'torchmetrics==1.0' "$_merged" && grep -qx 'anyio<4.14.0' "$_merged"; then
+ _rc=0
+else
+ _rc=1
+fi
+assert_true "unrelated inherited overrides are preserved" "$_rc"
+
+if grep -qE '^(torch|torchvision|torchaudio)([[:space:]<>=!~;@[]|$)' "$_merged" \
+ && [ "$(grep -cE '^(torch|torchvision|torchaudio)([[:space:]<>=!~;@[]|$)' "$_merged")" != "1" ]; then
+ _rc=1
+else
+ _rc=0
+fi
+grep -qx 'torch==2.11.0+cu128' "$_merged" || _rc=1
+assert_true "inherited torch-trio lines are dropped; generated pin wins" "$_rc"
+rm -rf "$_ov_dir"
+
+echo ""
+echo "Results: $PASS passed, $FAIL failed"
+[ "$FAIL" -eq 0 ] || exit 1
From b307823b1daf7013632340bceac5b2f70dbc04a8 Mon Sep 17 00:00:00 2001
From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com>
Date: Sun, 19 Jul 2026 21:33:48 +0800
Subject: [PATCH 019/255] fix(chat_templates): bind loop_messages when
default_system_message is None (#7199)
* fix(chat_templates): bind loop_messages when default_system_message is None
construct_chat_template(default_system_message=None) built a system part that
binds loop_messages only inside the `{% if messages[0]['role'] == 'system' %}`
arm. The `Fix missing loop_messages` step right below then found no
unconditional `{% set loop_messages = messages %}`, concluded loop_messages was
missing, and rewrote `{% for message in loop_messages %}` back to
`{% for message in messages %}` -- undoing the `messages[1:]` skip.
A caller-supplied system message therefore reached the loop and tripped
raise_exception:
Only user and assistant roles are supported!
Add the `{% else %}` arm so loop_messages is always bound, mirroring the
default_system_message is not None branch minus the default text. That also
stops the rewrite from firing, since the unconditional binding is now present.
Renders before / after, same template, same inputs:
default_system_message input before after
None system msg raise_exception 'Be terse.\n### User: Hi\n'
None no system '### User: Hi\n' unchanged
'You are helpful.' system msg 'Be terse.\n### User: Hi\n' unchanged
'You are helpful.' no system 'You are helpful.\n...' unchanged
The rewrite still fires for templates with no {SYSTEM} part, which is what it
was there for -- verified unchanged.
Co-Authored-By: Claude Opus 4.8
* Scope loop_messages binding to {SYSTEM} templates for PR #7199
The None branch now only adds the else arm when system_part contains
{SYSTEM}, so a static prefix with no {SYSTEM} placeholder keeps raising on a
caller system message instead of silently dropping it. Strengthen the tests:
assert the default does not leak when a caller system message is present, and
add a regression test for the static prefix case.
---------
Co-authored-by: Claude Opus 4.8
Co-authored-by: danielhanchen
---
...test_construct_chat_template_validation.py | 90 +++++++++++++++++++
unsloth/chat_templates.py | 6 ++
2 files changed, 96 insertions(+)
diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py
index 66d3d80920..53d281d435 100644
--- a/tests/python/test_construct_chat_template_validation.py
+++ b/tests/python/test_construct_chat_template_validation.py
@@ -104,3 +104,93 @@ def test_chat_template_does_not_leak_sentinel_when_section_starts_with_it(chat_t
)
assert "{INPUT}" not in jinja_template
assert "{OUTPUT}" not in jinja_template
+
+
+_SYSTEM_CHAT_TEMPLATE = (
+ "{SYSTEM}\n"
+ "### User: {INPUT}\n### Assistant: {OUTPUT}"
+ "### User: {INPUT}\n### Assistant: {OUTPUT}"
+)
+
+
+def _render(jinja_template, messages):
+ from jinja2.sandbox import ImmutableSandboxedEnvironment
+
+ env = ImmutableSandboxedEnvironment()
+ env.globals["raise_exception"] = lambda message: (_ for _ in ()).throw(RuntimeError(message))
+ return env.from_string(jinja_template).render(
+ messages = messages,
+ bos_token = "",
+ eos_token = " ",
+ add_generation_prompt = False,
+ )
+
+
+@pytest.mark.parametrize("default_system_message", [None, "You are helpful."])
+def test_system_message_is_consumed_by_the_system_part(default_system_message):
+ """A caller-supplied system message must be rendered by the system part and
+ skipped by the message loop, whatever `default_system_message` is.
+
+ With `default_system_message = None` the generated template used to bind
+ `loop_messages` only inside the `{% if %}` arm. The `Fix missing
+ loop_messages` step then saw no unconditional binding, rewrote the loop back
+ to `messages`, and the system message reached the loop and tripped
+ `raise_exception`.
+ """
+ _, jinja_template, _, _ = construct_chat_template(
+ tokenizer = _SuccessFakeTokenizer(),
+ chat_template = _SYSTEM_CHAT_TEMPLATE,
+ default_system_message = default_system_message,
+ extra_eos_tokens = [""],
+ )
+ rendered = _render(
+ jinja_template,
+ [
+ {"role": "system", "content": "Be terse."},
+ {"role": "user", "content": "Hi"},
+ ],
+ )
+ assert rendered.count("Be terse.") == 1
+ assert rendered.count("Hi") == 1
+ # A caller system message overrides the default; the default must not leak in.
+ if default_system_message is not None:
+ assert default_system_message not in rendered
+
+
+def test_absent_system_message_still_renders_without_default():
+ """`default_system_message = None` with no system message in the input must
+ keep working -- the `{% else %}` arm has to bind `loop_messages = messages`."""
+ _, jinja_template, _, _ = construct_chat_template(
+ tokenizer = _SuccessFakeTokenizer(),
+ chat_template = _SYSTEM_CHAT_TEMPLATE,
+ default_system_message = None,
+ extra_eos_tokens = [""],
+ )
+ rendered = _render(jinja_template, [{"role": "user", "content": "Hi"}])
+ assert "Hi" in rendered
+
+
+_NO_SYSTEM_CHAT_TEMPLATE = (
+ "PREAMBLE\n"
+ "### User: {INPUT}\n### Assistant: {OUTPUT}"
+ "### User: {INPUT}\n### Assistant: {OUTPUT}"
+)
+
+
+def test_static_prefix_without_system_still_rejects_system_message():
+ """A template with a static prefix but no {SYSTEM} placeholder cannot render a
+ caller system message, so it must still raise rather than silently drop it."""
+ _, jinja_template, _, _ = construct_chat_template(
+ tokenizer = _SuccessFakeTokenizer(),
+ chat_template = _NO_SYSTEM_CHAT_TEMPLATE,
+ default_system_message = None,
+ extra_eos_tokens = [""],
+ )
+ with pytest.raises(RuntimeError, match = "Only user and assistant roles are supported!"):
+ _render(
+ jinja_template,
+ [
+ {"role": "system", "content": "Be terse."},
+ {"role": "user", "content": "Hi"},
+ ],
+ )
diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index f47c78ba80..b857c34bcb 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -2652,6 +2652,12 @@ extra_eos_tokens = None,
"{{ '" + full_system + "' }}"\
"{% set loop_messages = messages %}"\
"{% endif %}"
+ elif "{SYSTEM}" in system_part:
+ # Only bind loop_messages when the template can render a caller system
+ # message. A static prefix with no {SYSTEM} must still raise, not drop it.
+ partial_system += "{% else %}"\
+ "{% set loop_messages = messages %}"\
+ "{% endif %}"
else:
partial_system += "{% endif %}"
From b3c0259cffdccb91362e7a16dc856632319f7304 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Sun, 19 Jul 2026 07:55:06 -0700
Subject: [PATCH 020/255] Installer: preserve the previous torch release across
every flavor and vendor on re-runs (#7250)
* install: preserve the previous torch release across every flavor and vendor
A re-run of curl | sh over an existing install was supposed to keep the
user's validated torch release, but the pin required the old build's
local flavor tag to match the freshly chosen index leaf. That gate was
wrong in practice: a PyPI-sourced torch reports a BARE version (on Linux
the PyPI wheel IS a CUDA build), which classified as cpu and never
matched a cu leaf, so a healthy 2.10 on a cu130 host was silently moved
to 2.11 (reproduced end to end); the same happened for any flavor drift
such as cu128 to cu130 after a driver upgrade, and AMD ROCm leaves were
excluded from preservation entirely.
The rule is now release-based and flavor-agnostic: the probed previous
release is pinned whenever it sits inside the final constraint window,
and the pin installs from the freshly chosen index, so the flavor always
follows the machine (NVIDIA cu*, AMD rocm/gfx, Intel/CPU, mac) while the
release follows the user. The pin is evaluated AFTER every index and
constraint decision including the Strix reroute, so raised floors
(rocm7.2 / Strix gfx need torch 2.11 for the _grouped_mm fix) correctly
reject an older release and win. UNSLOTH_TORCH_UPGRADE=1 still opts out,
out-of-window releases are never kept, and probe noise never becomes a
pin.
The kept-release install with its range fallback (for indexes that do
not carry the exact release) is factored into
_install_torch_default_index and used by every --default-index torch
path: the default NVIDIA/CPU/mac path and all three ROCm-index
fallbacks, which previously bypassed the fallback. The Radeon-repo
direct-wheel path keeps its curated per-arch wheel set (those wheels are
already exact-pinned per rocm release).
Platform coverage: install.sh serves Linux, WSL (including the WoA
fallback), and macOS for all vendors; native Windows install.ps1 still
caps at <2.11.0 everywhere, so the silent 2.10-to-2.11 move cannot occur
there (2.11 alignment is a separate follow-up).
Verified: 35-check unit suite rewritten to the new spec (any-flavor
keep, floor rejection, noise, window edges, opt-out, wiring including
pin-after-reroute and helper coverage); end-to-end matrix against
sandboxed UNSLOTH_STUDIO_HOME installs on a cu130 host covering PyPI
bare, cu128 drift, cu130 same-flavor, out-of-window 2.3, the upgrade
opt-out, the hidden-GPU cpu leaf, and a fresh-install control.
* install: honor the kept torch release on the Radeon direct-wheel path
The Radeon repo path installs an explicit wheel trio selected by
_pick_radeon_wheel, bypassing --default-index, so the kept-release pin
only took effect when the listing failed and the install fell back to
the ROCm index. On a re-run over an in-window Radeon install the trio
search started at the newest common minor and silently moved the user
forward (2.9 to 2.10 whenever the repo offered both).
The trio search now starts at the kept release's minor when
_PREV_TORCH_PIN is set and the listing still offers a torch wheel for
that minor. Radeon wheels are patch-curated per rocm release, so the
minor is the unit of preservation there; the raised rocm7.2 / Strix
floors still win because the pin is window-checked against the final
constraint before this point, and gaps keep the existing downward
search / ROCm-index fallback.
Verified with a simulated listing carrying both a 2.9 and a 2.10 trio:
no pin selects the 2.10 trio, a kept 2.9 release selects the matched
2.9 / 0.24 / 2.9 trio, and an unavailable minor degrades to the newest
trio. Added a structural wiring check to test_previous_torch_pin.sh
(now 36 checks).
* install: tighten comments in the torch preservation paths
* install: exact kept release on the Radeon path, pin fallback in ROCm repairs
The minor-level clamp on the Radeon direct-wheel path still allowed
patch drift (a kept 2.10.0 could become 2.10.1 when the listing carried
both) and the downward gap search could settle below the kept minor,
both breaking the exact preservation guarantee the other vendor paths
honor. The kept release now gets an exact-first trio attempt before the
newest-trio search: pick the kept patch (else the newest patch of the
kept minor, for listings that pruned the exact patch) together with the
paired torchvision/torchaudio wheels for that minor. Any gap warns and
falls back to the unchanged newest-trio search, mirroring
_install_torch_default_index, so a rerun installs either the kept
release or the same set a fresh install would choose, never something
in between.
The two ROCm torch repair sites (torch overwritten by dependency
resolution, on the migrated and fresh paths) installed TORCH_CONSTRAINT
directly, so a pinned release missing from the generic ROCm index would
abort the rerun instead of falling back. Both now route through
_install_torch_default_index, which passes extra uv args through
(--force-reinstall) and clears the pin once the fallback fires so later
paths stay consistent.
Verified against synthetic listings: both patches listed keeps exactly
2.10.0; a kept minor missing vision/audio warns and yields the newest
complete trio rather than a silent undercut; a pruned patch stays on
the kept minor; no pin keeps the existing newest-trio behavior. Unit
suite now 39 checks, all passing.
* install: never pin nightly/dev/source torch builds on a rerun
A survey of published torch version strings (PyPI bare, +cpu, +cu116
through +cu132, +rocmX.Y and +rocmX.Y.Z, +xpu, nightly .devYYYYMMDD,
source a0+git, rc tags) showed one gap: nightly, dev, rc, and source
builds passed the loose release-shape check, producing a pin such as
torch==2.11.0.dev20250704 that no stable index carries. The range
fallback rescued the install, but it printed "keeping it" and then
burned a doomed resolve first. The base must now be a plain numeric
X.Y[.Z] release, so those builds skip the pin and go straight to the
newest supported release.
Added unit checks for +xpu and three-component +rocm7.2.1 tags (both
already preserved correctly) and for nightly, a0 source, and rc builds
(never pinned). Suite now 44 checks, all passing.
* install: pair kept-release companions, protect the flavor repair, note substitutions
Three fixes from a 12-way review pass over the preservation work:
The kept-release install left torchvision and torchaudio unconstrained
next to the exact torch pin. torchvision exact-pins its torch in wheel
metadata so it always paired correctly, but torchaudio no longer does:
a kept torch 2.9.0 on cu130 resolved torchaudio 2.11.0 (verified with
uv dry-runs). The helper now pairs both companions to the kept minor
(torchvision 0.minor+15, torchaudio 2.minor); if the index lacks the
paired set the existing range fallback fires. Verified resolving
correctly on cu130, cu126, and rocm6.4.
The wrong-flavor repair at the end of the install was the one remaining
default-index torch install outside the helper. It runs under set -e,
so a retained pin absent from the repair index (reachable when the
Radeon direct-wheel path installed the kept release and dependency
resolution later overwrote it) aborted the installer at the last step
instead of falling back. It now routes through the helper with its
reinstall flags passed through.
The Radeon kept-release path installed a same-series build silently
when the listing had pruned the exact patch; it now prints what it is
substituting.
Unit suite extended with wiring checks for all three (46 checks, all
passing).
---
install.sh | 180 +++++++++++++++++-----------
tests/sh/test_previous_torch_pin.sh | 106 ++++++++++------
2 files changed, 181 insertions(+), 105 deletions(-)
diff --git a/install.sh b/install.sh
index 6076721540..7918a2bd23 100755
--- a/install.sh
+++ b/install.sh
@@ -2225,37 +2225,67 @@ _torch_release_in_window() {
echo "no"
}
-# Whether a re-run should keep the previous venv's torch: echo "torch==X.Y.Z" when the
-# probed previous version ($1) has a flavor tag matching the freshly chosen cu*/cpu index
-# leaf ($2) AND sits inside the active constraint window ($3), else "". Re-running
-# `curl | sh` rebuilds the venv for clean state, but a healthy torch the user already
-# validated must not be silently moved to a newer release (2.10 -> 2.11); a flavor
-# change (cpu <-> cuda, cu126 -> cu130) still installs the correct new build, rocm
-# leaves keep their floors (rocm7.2 must land 2.11 for the Strix _grouped_mm fix), and
-# a release outside the window (2.3.x manual install, 2.12.x manual upgrade) is never
-# kept: the installer's own bounds win. Opt out with UNSLOTH_TORCH_UPGRADE=1 to get
-# the newest release.
+# Keep the previous venv's torch on a re-run: echo "torch==X.Y.Z" when the probed
+# version ($1) is inside the active constraint window ($2), else "". The RELEASE is kept
+# regardless of flavor tag; the pin installs from the freshly chosen index, so flavor
+# follows the machine (cpu <-> cuda, cu126 -> cu130, PyPI bare -> +cu130) while the
+# release follows the user. Gating on flavor was wrong: a PyPI torch reports a BARE
+# version (on Linux the PyPI wheel IS CUDA), misclassified "cpu", so a healthy 2.10 on a
+# cu130 host was moved to 2.11. Per-leaf floors still win (rocm7.2 / gfx >=2.11 for the
+# Strix _grouped_mm fix, out-of-window manual installs) and are never pinned; the caller's
+# _PREV_FALLBACK_CONSTRAINT installs the newest supported release when the index lacks the
+# exact one. Opt out with UNSLOTH_TORCH_UPGRADE=1.
_previous_torch_pin() {
_ptp_ver="$1"
- _ptp_leaf="$2"
- _ptp_con="$3"
+ _ptp_con="$2"
[ -n "$_ptp_ver" ] || { echo ""; return; }
[ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; }
- case "$_ptp_leaf" in
- cu[0-9]*|cpu) ;;
- *) echo ""; return ;;
- esac
_ptp_base="${_ptp_ver%%+*}"
- # The base must look like a release (probe noise / garbage must never become a pin).
+ # Base must be a plain numeric release (X.Y[.Z]); probe noise and
+ # nightly/dev/source builds (2.11.0.dev20250704, 2.9.0a0) must never
+ # become a pin -- no stable index carries them, so pinning would only
+ # print "keeping it" and then burn a doomed resolve before falling back.
case "$_ptp_base" in
+ *[!0-9.]* | *..* | .* | *.) echo ""; return ;;
[0-9]*.[0-9]*) ;;
*) echo ""; return ;;
esac
[ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; }
- if [ "$(_torch_flavor_tag "$_ptp_ver")" = "$_ptp_leaf" ]; then
- echo "torch==$_ptp_base"
+ echo "torch==$_ptp_base"
+}
+
+# Install torch from TORCH_INDEX_URL honoring a kept-release pin: with _PREV_TORCH_PIN
+# set, TORCH_CONSTRAINT is the exact previous release; fall back to the supported range
+# if the index lacks it (pruned mirror) rather than failing. Used by every --default-index
+# path (NVIDIA cu*, AMD rocm/gfx fallbacks, cpu/mac, ROCm repairs) so preservation is
+# uniform. Extra args (e.g. --force-reinstall) are passed through to uv.
+_install_torch_default_index() {
+ if [ -n "$_PREV_TORCH_PIN" ]; then
+ # Pair the companions with the kept torch minor: torchaudio no longer
+ # exact-pins torch in its metadata, so leaving it unconstrained resolves
+ # a newer mismatched build (a kept torch 2.9.0 pulled torchaudio 2.11.0).
+ _itdi_base="${_PREV_TORCH_PIN#torch==}"
+ _itdi_minor="${_itdi_base#*.}"
+ _itdi_minor="${_itdi_minor%%.*}"
+ _itdi_tv="torchvision"
+ _itdi_ta="torchaudio"
+ case "$_itdi_base" in
+ 2.*)
+ _itdi_tv="torchvision==0.$((_itdi_minor + 15)).*"
+ _itdi_ta="torchaudio==2.${_itdi_minor}.*"
+ ;;
+ esac
+ if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$_itdi_tv" "$_itdi_ta" \
+ --default-index "$TORCH_INDEX_URL" "$@"; then
+ substep "[WARN] $_PREV_TORCH_PIN is not installable from $TORCH_INDEX_URL -- installing the newest supported release instead" "$C_WARN"
+ TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT"
+ _PREV_TORCH_PIN=""
+ run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
+ --default-index "$TORCH_INDEX_URL" "$@"
+ fi
else
- echo ""
+ run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
+ --default-index "$TORCH_INDEX_URL" "$@"
fi
}
@@ -2561,21 +2591,6 @@ case "$_torch_index_leaf" in
cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
esac
-# Re-run over an existing install: keep the previous venv's torch release instead of
-# resolving the newest in range. The range stays in _PREV_FALLBACK_CONSTRAINT so the
-# install can fall back when the exact release is not on the chosen index (custom
-# mirrors may prune old wheels). Skipped for --no-torch (no previous probe runs).
-_PREV_TORCH_PIN=""
-_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT"
-if [ "$SKIP_TORCH" = false ]; then
- _prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$_torch_index_leaf" "$TORCH_CONSTRAINT")
- if [ -n "$_prev_pin" ]; then
- _PREV_TORCH_PIN="$_prev_pin"
- TORCH_CONSTRAINT="$_prev_pin"
- substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)"
- fi
-fi
-
# Auto-detect GPU for AMD ROCm based
# get_torch_index_url must have chosen */rocm*
# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
@@ -2660,6 +2675,23 @@ case "$TORCH_INDEX_URL" in
fi
;;
esac
+# Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh
+# index above supplies the right flavor for this machine. Evaluated HERE, after every
+# index/constraint decision including the Strix reroute, so the window checked is the
+# final one and a raised floor (rocm7.2 / Strix gfx) rejects an older release.
+# _PREV_FALLBACK_CONSTRAINT keeps the range so the install can fall back when the exact
+# release is not on the chosen index (mirrors may prune old wheels). Skipped for --no-torch.
+_PREV_TORCH_PIN=""
+_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT"
+if [ "$SKIP_TORCH" = false ]; then
+ _prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$TORCH_CONSTRAINT")
+ if [ -n "$_prev_pin" ]; then
+ _PREV_TORCH_PIN="$_prev_pin"
+ TORCH_CONSTRAINT="$_prev_pin"
+ substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)"
+ fi
+fi
+
_TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL")
if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then
_TAURI_TORCH_INDEX_FAMILY="radeon"
@@ -2885,10 +2917,7 @@ if [ "$_MIGRATED" = true ]; then
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
- run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
- "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL" \
- --force-reinstall
+ _install_torch_default_index --force-reinstall
fi
;;
esac
@@ -2953,7 +2982,42 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_ta_ver=$(_extract_version "$_ta_whl" "torchaudio")
_radeon_versions_match=false
- if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
+ # Kept release (_PREV_TORCH_PIN) wins here too: pick its exact
+ # patch (else the newest patch of its minor) plus the paired
+ # vision/audio wheels. Any gap falls back to the newest-trio
+ # search below, mirroring _install_torch_default_index, so a
+ # rerun never drifts to another release nor below the kept one.
+ if [ -n "$_PREV_TORCH_PIN" ]; then
+ _prev_kept_base="${_PREV_TORCH_PIN#torch==}"
+ _prev_kept_minor="${_prev_kept_base#*.}"
+ _prev_kept_minor="${_prev_kept_minor%%.*}"
+ case "$_prev_kept_minor" in
+ ''|*[!0-9]*) ;;
+ *)
+ _kept_torch=$(_pick_radeon_wheel "torch" "${_prev_kept_base}" 2>/dev/null) || _kept_torch=""
+ [ -z "$_kept_torch" ] && { _kept_torch=$(_pick_radeon_wheel "torch" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_torch=""; }
+ _kept_tv=$(_pick_radeon_wheel "torchvision" "0.$((_prev_kept_minor + 15))." 2>/dev/null) || _kept_tv=""
+ _kept_ta=$(_pick_radeon_wheel "torchaudio" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_ta=""
+ if [ -n "$_kept_torch" ] && [ -n "$_kept_tv" ] && [ -n "$_kept_ta" ]; then
+ _torch_whl=$_kept_torch
+ _tv_whl=$_kept_tv
+ _ta_whl=$_kept_ta
+ _tri_whl=""
+ _radeon_versions_match=true
+ # Say so when the listing pruned the exact patch
+ # and a same-series build is installed instead.
+ case "$(printf '%s' "${_kept_torch##*/}" | sed 's/%2[Bb]/+/g')" in
+ "torch-${_prev_kept_base}"[+-]*) ;;
+ *) substep "kept release ${_prev_kept_base} is not in the Radeon listing -- installing the closest 2.${_prev_kept_minor} series build instead" ;;
+ esac
+ else
+ substep "[WARN] Radeon repo lacks a complete wheel set for kept $_PREV_TORCH_PIN -- installing the newest compatible set instead" "$C_WARN"
+ fi
+ ;;
+ esac
+ fi
+ if [ "$_radeon_versions_match" != true ] && \
+ [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
_torch_minor=${_torch_ver#*.}
_ta_minor=${_ta_ver#*.}
_tv_minor=${_tv_ver#*.}
@@ -3011,9 +3075,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
- run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
- "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL"
+ _install_torch_default_index
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
# Pass explicit wheel URLs so the matched trio is
@@ -3034,32 +3096,15 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
else
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
- run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
- "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL"
+ _install_torch_default_index
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
- run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
- "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL"
+ _install_torch_default_index
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
- if [ -n "$_PREV_TORCH_PIN" ]; then
- # Kept previous release: fall back to the supported range if the exact
- # release is not resolvable from the chosen index (pruned mirror).
- if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL"; then
- substep "[WARN] $_PREV_TORCH_PIN is not installable from $TORCH_INDEX_URL -- installing the newest supported release instead" "$C_WARN"
- TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT"
- run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL"
- fi
- else
- run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL"
- fi
+ _install_torch_default_index
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
@@ -3122,10 +3167,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
- run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
- "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL" \
- --force-reinstall
+ _install_torch_default_index --force-reinstall
fi
;;
esac
@@ -3164,9 +3206,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
- run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
- "$TORCH_CONSTRAINT" torchvision torchaudio \
- --default-index "$TORCH_INDEX_URL" \
+ _install_torch_default_index \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
diff --git a/tests/sh/test_previous_torch_pin.sh b/tests/sh/test_previous_torch_pin.sh
index 253ede8a27..1bc0d1f27f 100644
--- a/tests/sh/test_previous_torch_pin.sh
+++ b/tests/sh/test_previous_torch_pin.sh
@@ -2,9 +2,12 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Unit tests for install.sh's _previous_torch_pin, which keeps the previous
-# venv's torch release on a re-run (curl | sh over an existing install) instead
-# of silently moving the user to a newer release. Helpers are extracted from
-# install.sh and sourced.
+# venv's torch RELEASE on a re-run instead of moving the user to a newer one.
+# The release is kept regardless of the old build's flavor tag (PyPI bare,
+# +cuXXX, +rocm, +cpu): the pin installs from the freshly chosen index, so the
+# flavor follows the machine while the release follows the user. Per-leaf
+# windows still win (rocm7.2 / Strix floors, out-of-window manual installs).
+# Helpers are extracted from install.sh and sourced.
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
@@ -12,12 +15,9 @@ INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
-# Extract _previous_torch_pin and its dependencies _torch_flavor_tag and
-# _torch_release_in_window.
+# Extract _previous_torch_pin and its dependency _torch_release_in_window.
_FUNC_FILE=$(mktemp)
{
- sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
- echo ""
sed -n '/^_torch_release_in_window()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_previous_torch_pin()/,/^}/p' "$INSTALL_SH"
@@ -37,37 +37,43 @@ assert_eq() {
unset UNSLOTH_TORCH_UPGRADE
-echo "=== _previous_torch_pin: matching flavor keeps the release ==="
-assert_eq "cu126 wheel on cu126 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
-assert_eq "cu130 wheel on cu130 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
-assert_eq "cpu wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cpu' 'cpu' 'torch>=2.4,<2.12.0')"
-assert_eq "untagged wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0' 'cpu' 'torch>=2.4,<2.12.0')"
-assert_eq "local suffix stripped" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+cu128' 'cu128' 'torch>=2.4,<2.12.0')"
+echo "=== _previous_torch_pin: in-window releases are kept, any flavor ==="
+assert_eq "cu126 wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "cu130 wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu130' 'torch>=2.4,<2.12.0')"
+assert_eq "cpu wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cpu' 'torch>=2.4,<2.12.0')"
+assert_eq "PyPI bare version (CUDA build on Linux)" "torch==2.10.0" "$(_previous_torch_pin '2.10.0' 'torch>=2.4,<2.12.0')"
+assert_eq "rocm wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+rocm6.4' 'torch>=2.4,<2.11.0')"
+assert_eq "rocm three-component tag" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+rocm7.2.1' 'torch>=2.4,<2.12.0')"
+assert_eq "Intel xpu wheel" "torch==2.9.0" "$(_previous_torch_pin '2.9.0+xpu' 'torch>=2.4,<2.12.0')"
+assert_eq "local suffix stripped" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+cu128' 'torch>=2.4,<2.12.0')"
-echo "=== _previous_torch_pin: flavor change installs the new build ==="
-assert_eq "cu126 wheel on cu130 leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu130' 'torch>=2.4,<2.12.0')"
-assert_eq "cpu wheel on cu126 leaf" "" "$(_previous_torch_pin '2.10.0+cpu' 'cu126' 'torch>=2.4,<2.12.0')"
-assert_eq "cu126 wheel on cpu leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cpu' 'torch>=2.4,<2.12.0')"
-
-echo "=== _previous_torch_pin: rocm and unknown leaves never pin ==="
-assert_eq "rocm7.2 leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'rocm7.2' 'torch>=2.4,<2.12.0')"
-assert_eq "gfx leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'gfx120X-all' 'torch>=2.4,<2.12.0')"
-assert_eq "unknown mirror leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'simple' 'torch>=2.4,<2.12.0')"
+echo "=== _previous_torch_pin: raised floors reject older releases ==="
+# rocm7.2 / Strix gfx leaves raise TORCH_CONSTRAINT to >=2.11.0 BEFORE the pin
+# is evaluated, so an old 2.10 is out of window there and the floor wins.
+assert_eq "old 2.10 vs rocm7.2 floor" "" "$(_previous_torch_pin '2.10.0+rocm7.1' 'torch>=2.11.0,<2.12.0')"
+assert_eq "2.11 passes the rocm7.2 floor" "torch==2.11.0" "$(_previous_torch_pin '2.11.0+rocm7.2' 'torch>=2.11.0,<2.12.0')"
echo "=== _previous_torch_pin: probe noise never becomes a pin ==="
-assert_eq "empty version" "" "$(_previous_torch_pin '' 'cu126' 'torch>=2.4,<2.12.0')"
-assert_eq "garbage version" "" "$(_previous_torch_pin 'not-a-version' 'cpu' 'torch>=2.4,<2.12.0')"
-assert_eq "traceback fragment" "" "$(_previous_torch_pin "ModuleNotFoundError: No module named 'torch'" 'cpu' 'torch>=2.4,<2.12.0')"
+assert_eq "empty version" "" "$(_previous_torch_pin '' 'torch>=2.4,<2.12.0')"
+assert_eq "garbage version" "" "$(_previous_torch_pin 'not-a-version' 'torch>=2.4,<2.12.0')"
+assert_eq "traceback fragment" "" "$(_previous_torch_pin "ModuleNotFoundError: No module named 'torch'" 'torch>=2.4,<2.12.0')"
+
+echo "=== _previous_torch_pin: nightly / dev / source builds never pin ==="
+# No stable index carries these, so pinning would print "keeping it" and then
+# burn a doomed resolve before the range fallback rescues the install.
+assert_eq "nightly dev build" "" "$(_previous_torch_pin '2.11.0.dev20250704+cu128' 'torch>=2.4,<2.12.0')"
+assert_eq "source build a0 tag" "" "$(_previous_torch_pin '2.9.0a0+gitabc1234' 'torch>=2.4,<2.12.0')"
+assert_eq "release candidate" "" "$(_previous_torch_pin '2.11.0rc1+cu130' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: out-of-window releases never pin ==="
-assert_eq "2.3.x below the cu floor" "" "$(_previous_torch_pin '2.3.1+cu118' 'cu118' 'torch>=2.4,<2.12.0')"
-assert_eq "2.12.x above the cu ceiling" "" "$(_previous_torch_pin '2.12.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
-assert_eq "floor boundary 2.4.0 kept" "torch==2.4.0" "$(_previous_torch_pin '2.4.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
-assert_eq "ceiling-adjacent 2.11.x kept" "torch==2.11.1" "$(_previous_torch_pin '2.11.1+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
-assert_eq "cpu window excludes 2.11.x" "" "$(_previous_torch_pin '2.11.0+cpu' 'cpu' 'torch>=2.4,<2.11.0')"
-assert_eq "mac floor excludes 2.5.x" "" "$(_previous_torch_pin '2.5.1' 'cpu' 'torch>=2.6,<2.11.0')"
-assert_eq "malformed window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch')"
-assert_eq "empty window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' '')"
+assert_eq "2.3.x below the cu floor" "" "$(_previous_torch_pin '2.3.1+cu118' 'torch>=2.4,<2.12.0')"
+assert_eq "2.12.x above the cu ceiling" "" "$(_previous_torch_pin '2.12.0+cu130' 'torch>=2.4,<2.12.0')"
+assert_eq "floor boundary 2.4.0 kept" "torch==2.4.0" "$(_previous_torch_pin '2.4.0+cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "ceiling-adjacent 2.11.x kept" "torch==2.11.1" "$(_previous_torch_pin '2.11.1+cu130' 'torch>=2.4,<2.12.0')"
+assert_eq "cpu window excludes 2.11.x" "" "$(_previous_torch_pin '2.11.0+cpu' 'torch>=2.4,<2.11.0')"
+assert_eq "mac floor excludes 2.5.x" "" "$(_previous_torch_pin '2.5.1' 'torch>=2.6,<2.11.0')"
+assert_eq "malformed window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'torch')"
+assert_eq "empty window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' '')"
echo "=== _torch_release_in_window ==="
assert_eq "in window" "yes" "$(_torch_release_in_window '2.10.0' 'torch>=2.4,<2.12.0')"
@@ -80,8 +86,8 @@ assert_eq "no ceiling -> no" "no" "$(_torch_release_in_window '2.10.0' 'tor
assert_eq "garbage minor -> no" "no" "$(_torch_release_in_window '2.x' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: UNSLOTH_TORCH_UPGRADE=1 opts out ==="
-assert_eq "upgrade env set" "" "$(UNSLOTH_TORCH_UPGRADE=1 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
-assert_eq "upgrade env 0" "torch==2.10.0" "$(UNSLOTH_TORCH_UPGRADE=0 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "upgrade env set" "" "$(UNSLOTH_TORCH_UPGRADE=1 _previous_torch_pin '2.10.0+cu126' 'torch>=2.4,<2.12.0')"
+assert_eq "upgrade env 0" "torch==2.10.0" "$(UNSLOTH_TORCH_UPGRADE=0 _previous_torch_pin '2.10.0+cu126' 'torch>=2.4,<2.12.0')"
echo "=== install.sh wiring ==="
# The probe must run against the OLD venv, before it is moved aside for rollback.
@@ -89,9 +95,39 @@ _probe_line=$(grep -n '_PREV_TORCH_VER=\$(' "$INSTALL_SH" | head -1 | cut -d: -f
_move_line=$(grep -n '_start_studio_venv_replacement "\$VENV_DIR"' "$INSTALL_SH" | head -1 | cut -d: -f1)
assert_eq "probe exists" "yes" "$([ -n "$_probe_line" ] && echo yes)"
assert_eq "probe before venv replacement" "yes" "$([ -n "$_probe_line" ] && [ -n "$_move_line" ] && [ "$_probe_line" -lt "$_move_line" ] && echo yes)"
+# The pin must be evaluated AFTER the last index/constraint decision (the Strix
+# reroute raises the floor), so a raised floor rejects an older kept release.
+_pin_line=$(grep -n '_prev_pin=\$(_previous_torch_pin' "$INSTALL_SH" | head -1 | cut -d: -f1)
+_strix_line=$(grep -n 'Strix Halo / Strix Point: force rocm7.2 wheels' "$INSTALL_SH" | head -1 | cut -d: -f1)
+assert_eq "pin evaluated after the Strix reroute" "yes" "$([ -n "$_pin_line" ] && [ -n "$_strix_line" ] && [ "$_pin_line" -gt "$_strix_line" ] && echo yes)"
# A kept release that vanished from the index must fall back to the supported range.
assert_eq "resolve-failure fallback wired" "yes" "$(grep -q 'TORCH_CONSTRAINT="\$_PREV_FALLBACK_CONSTRAINT"' "$INSTALL_SH" && echo yes)"
assert_eq "pin gated on SKIP_TORCH" "yes" "$(grep -q 'if \[ "\$SKIP_TORCH" = false \]; then' "$INSTALL_SH" && echo yes)"
+# Every --default-index torch install path must go through the kept-release
+# helper (definition + default path + three ROCm-index fallbacks + two ROCm
+# repairs + the flavor repair), so a pinned release missing from the index
+# never aborts a rerun.
+_helper_uses=$(grep -c '_install_torch_default_index' "$INSTALL_SH")
+assert_eq "kept-release helper used by all default-index paths" "yes" "$([ "$_helper_uses" -ge 8 ] && echo yes)"
+_repair_uses=$(grep -c '_install_torch_default_index --force-reinstall' "$INSTALL_SH")
+assert_eq "ROCm repairs routed through the kept-release helper" "yes" "$([ "$_repair_uses" -ge 2 ] && echo yes)"
+# The wrong-flavor repair must use the helper too (it runs under set -e, so a
+# direct uv call with an unresolvable pin would abort the whole installer).
+assert_eq "flavor repair routed through the kept-release helper" "yes" "$(grep -q '_install_torch_default_index \\' "$INSTALL_SH" && grep -q -- '--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio' "$INSTALL_SH" && echo yes)"
+# The kept-release install must pair the companions with the kept minor:
+# torchaudio no longer exact-pins torch, so unconstrained it resolves a newer
+# mismatched build (verified: torch==2.9.0 pulled torchaudio 2.11.0 on cu130).
+assert_eq "kept-release install pairs torchvision/torchaudio to the kept minor" "yes" "$(grep -q 'torchaudio==2.\${_itdi_minor}.\*' "$INSTALL_SH" && grep -q 'torchvision==0.\$((_itdi_minor + 15)).\*' "$INSTALL_SH" && echo yes)"
+# The Radeon direct-wheel path must also honor the pin: an exact-first kept-trio
+# attempt (exact patch, else the kept minor's newest patch, with paired
+# vision/audio) runs BEFORE the newest-trio search, and the newest-trio search
+# only runs when that attempt did not produce a match, so a kept release can
+# neither drift to another patch/minor nor be undercut by the gap search.
+_radeon_kept_line=$(grep -n '_kept_torch=\$(_pick_radeon_wheel "torch" *"\${_prev_kept_base}"' "$INSTALL_SH" | head -1 | cut -d: -f1)
+_radeon_loop_line=$(grep -n 'Loop downwards to find the first complete matching trio' "$INSTALL_SH" | head -1 | cut -d: -f1)
+assert_eq "Radeon kept-trio attempt before the newest-trio search" "yes" "$([ -n "$_radeon_kept_line" ] && [ -n "$_radeon_loop_line" ] && [ "$_radeon_kept_line" -lt "$_radeon_loop_line" ] && echo yes)"
+assert_eq "Radeon newest-trio search gated on no kept match" "yes" "$(grep -q 'if \[ "\$_radeon_versions_match" != true \] &&' "$INSTALL_SH" && echo yes)"
+assert_eq "Radeon kept-trio gap falls back with a warning" "yes" "$(grep -q 'lacks a complete wheel set for kept' "$INSTALL_SH" && echo yes)"
echo ""
if [ "$FAIL" -gt 0 ]; then
From 17fd6c8ec6c3788c1da2f9e86452fc340e451444 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Sun, 19 Jul 2026 17:19:15 -0700
Subject: [PATCH 021/255] studio: fix stale GGUF load-marker ordering test
after inheritance relocation (#7252)
#6414 moved the llama_extra_args inheritance out of the GGUF branch in
_load_model_impl into _guard_chat_load_against_training, which runs before the
branch, so 'if request.llama_extra_args is None' is no longer inside the
gguf_branch slice that test_load_marker_precedes_hub_guard_and_unload checks.
The assertion failed on that now-missing landmark even though the guarantee it
protects (the gguf_load_in_flight marker is entered before the hub-download
guard and the unload) is intact. Drop the relocated landmark from the ordering
so the test matches the current structure.
Co-authored-by: danielhanchen
---
studio/backend/tests/test_gguf_load_cache_reuse.py | 6 +++++-
1 file changed, 5 insertions(+), 1 deletion(-)
diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py
index 15d91cd324..62596fcc8a 100644
--- a/studio/backend/tests/test_gguf_load_cache_reuse.py
+++ b/studio/backend/tests/test_gguf_load_cache_reuse.py
@@ -728,9 +728,13 @@ class TestLoadHubDownloadExclusion:
source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text()
gguf_branch = source[source.index("if config.is_gguf:") :]
+ # The gguf_load_in_flight marker must be entered before the hub-download
+ # guard and the unload so a concurrent load can't race the download
+ # manager. The llama_extra_args inheritance that used to sit between the
+ # marker and the guard now runs in _guard_chat_load_against_training, ahead
+ # of the GGUF branch, so it is no longer a landmark inside this slice.
assert (
gguf_branch.index("enter_context(gguf_load_in_flight")
- < gguf_branch.index("if request.llama_extra_args is None")
< gguf_branch.index("_hub_download_blocks_gguf_load")
< gguf_branch.index("unsloth_backend.unload_model")
)
From 8fab1c5310e6d4117a939f30c8d1546ffca023bd Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 20 Jul 2026 05:03:50 +0100
Subject: [PATCH 022/255] Route OpenCode yolo aliases to native auto mode
(#7187)
Route --yolo to OpenCode native --auto for the default TUI and run; keep the config permission fallback for no-auto subcommands (including hidden console/generate) and for --mini, which ignores --auto.
---
unsloth_cli/commands/start.py | 115 ++++++++++++++++++--
unsloth_cli/tests/test_start.py | 185 +++++++++++++++++++++++++++++++-
2 files changed, 285 insertions(+), 15 deletions(-)
diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py
index 9257a7fbcb..8128447a02 100644
--- a/unsloth_cli/commands/start.py
+++ b/unsloth_cli/commands/start.py
@@ -149,8 +149,8 @@ _PERSIST_OPTION = typer.Option(
),
)
-# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no
-# such flag (config only) and are handled in their config writers, so they are absent.
+# Per-agent CLI flag for "run tools without prompting". OpenCode (native --auto is
+# command-scoped, handled below) and OpenClaw (config-only) are absent from this prefix map.
_YOLO_COMMAND_FLAGS = {
"claude": ["--dangerously-skip-permissions"],
"codex": ["--dangerously-bypass-approvals-and-sandbox"],
@@ -166,6 +166,84 @@ def _yolo_command_flags(agent: str, yolo: bool) -> list:
return _YOLO_COMMAND_FLAGS.get(agent, []) if yolo else []
+# Subcommands that reject --auto (OpenCode exposes it only on the default TUI and `run`),
+# so `opencode serve --auto` is never emitted. Includes console/generate, hidden from
+# `opencode --help` but still registered. Unknown first positionals are TUI paths -> --auto.
+_OPENCODE_NON_AUTO_SUBCOMMANDS = frozenset(
+ "completion acp mcp attach debug providers auth agent upgrade uninstall serve web "
+ "models stats export import github pr session plugin plug db console generate".split()
+)
+_OPENCODE_GLOBAL_BOOLEAN_OPTIONS = frozenset(
+ "-h --help -v --version --print-logs --pure --mdns".split()
+)
+_OPENCODE_GLOBAL_VALUE_OPTIONS = frozenset(
+ "--log-level --port --hostname --mdns-domain --cors".split()
+)
+_OPENCODE_NATIVE_AUTO_MIN_VERSION = (1, 17, 12)
+
+
+def _opencode_supports_native_auto() -> bool:
+ executable = shutil.which("opencode")
+ if executable is None:
+ # No local binary: a --no-launch recipe may run elsewhere, and _run installs the
+ # current release on launch -- either way assume native --auto is available.
+ return True
+ try:
+ output = subprocess.check_output(
+ [executable, "--version"],
+ text = True,
+ timeout = 10,
+ stderr = subprocess.DEVNULL,
+ )
+ except Exception:
+ return False
+ match = re.search(r"(\d+)\.(\d+)\.(\d+)", output)
+ return bool(match) and tuple(int(part) for part in match.groups()) >= (
+ _OPENCODE_NATIVE_AUTO_MIN_VERSION
+ )
+
+
+def _opencode_subcommand(args: list[str]) -> Optional[str]:
+ """Return an explicit OpenCode subcommand after supported global options."""
+ index = 0
+ while index < len(args):
+ arg = args[index]
+ if arg == "--":
+ return None
+ if arg in _OPENCODE_GLOBAL_BOOLEAN_OPTIONS:
+ index += 1
+ continue
+ if arg in _OPENCODE_GLOBAL_VALUE_OPTIONS:
+ index += 2
+ continue
+ if any(arg.startswith(f"{option}=") for option in _OPENCODE_GLOBAL_VALUE_OPTIONS):
+ index += 1
+ continue
+ # A non-global option (e.g. --session) is a TUI flag; stop before its value is
+ # mistaken for a subcommand.
+ if arg.startswith("-"):
+ return None
+ return arg
+ return None
+
+
+def _opencode_native_auto_args(args: list[str], yolo: bool) -> tuple[list[str], bool]:
+ """Add OpenCode's native --auto when the selected command supports it."""
+ routed = list(args)
+ if not yolo:
+ return routed, False
+ if _opencode_subcommand(routed) in _OPENCODE_NON_AUTO_SUBCOMMANDS:
+ return routed, False
+ separator = routed.index("--") if "--" in routed else len(routed)
+ # --mini's runMini TUI forces auto=false and never forwards --auto, so appending it is
+ # useless; fall back to the config permission block so --yolo still auto-approves.
+ if any(arg == "--mini" or arg.startswith("--mini=") for arg in routed[:separator]):
+ return routed, False
+ if "--auto" not in routed[:separator]:
+ routed.insert(separator, "--auto")
+ return routed, True
+
+
def _hermes_install_hint() -> str:
return _HERMES_WINDOWS_INSTALL_HINT if os.name == "nt" else _HERMES_POSIX_INSTALL_HINT
@@ -1465,10 +1543,10 @@ def write_opencode_config(
compaction["reserved"] = max(1, window // 10)
tools = ("edit", "bash", "webfetch")
if yolo:
- # OpenCode has no --yolo flag; auto-approve is the config `permission` block
- # (singular). Allow the prompting tools and paths outside the launch directory so
- # tool calls don't block on the TUI. This rides inline (OPENCODE_CONFIG_CONTENT) so
- # --yolo works even over a project config.
+ # Fallback for commands without native --auto and for the append-safe bare
+ # --no-launch command (subcommand unknown yet). Rides inline (OPENCODE_CONFIG_CONTENT)
+ # so it wins over a project config. TUI and `run` launches use --auto and call here
+ # with yolo=False, letting OpenCode preserve explicit deny rules.
session_permission = {t: "allow" for t in tools}
session_permission["external_directory"] = {"*": "allow"}
config["permission"] = dict(session_permission)
@@ -1807,11 +1885,20 @@ def opencode(
# --no-launch, where the printed command is consumed by drivers that append a
# subcommand such as `run `; a leading --model would land before that
# subcommand and break it. Those paths rely on the inline pin instead.
+ native_auto = False
+ route_native_auto = yolo and _opencode_supports_native_auto()
if ctx.args:
- command = ["opencode", *ctx.args]
+ opencode_args, native_auto = _opencode_native_auto_args(list(ctx.args), route_native_auto)
+ command = ["opencode", *opencode_args]
elif launch:
- command = ["opencode", "--model", opencode_model]
+ opencode_args, native_auto = _opencode_native_auto_args(
+ ["--model", opencode_model],
+ route_native_auto,
+ )
+ command = ["opencode", *opencode_args]
else:
+ # Append-safe base: `opencode --auto run ...` parses as the TUI with a project
+ # "run", not the run subcommand. Command unknown here, so keep the config fallback.
command = ["opencode"]
# opencode keeps sessions in ~/.local/share/opencode (never relocated), so resume
# already survives exit; reopen the last one by passing `opencode --continue` through.
@@ -1820,12 +1907,18 @@ def opencode(
# OPENCODE_CONFIG is an overlay (loaded between the user's global and project
# configs), so this adds the Unsloth provider/model for the session without
# changing the user's default model. Key lives in the config, not the env.
- session_permission = write_opencode_config(base, key, entry, config_path, yolo = yolo)
+ session_permission = write_opencode_config(
+ base,
+ key,
+ entry,
+ config_path,
+ yolo = yolo and not native_auto,
+ )
# A project's own opencode.json outranks OPENCODE_CONFIG, so the session model pin
# would silently lose to a repo config. Carry it in OPENCODE_CONFIG_CONTENT, which
# outranks project config; the API key stays in the private file, never the env.
- # Only --yolo carries a permission here (its allow must win over a project config);
- # a non-yolo session returns no permission, so the project's own rules are honored.
+ # Only the config fallback carries a permission. Native --auto omits it (auto-approve
+ # asks, keep explicit denies); a non-yolo session omits it too, honoring project rules.
# opencode filters every provider (a config-defined custom one included) through
# its enabled_providers allowlist and disabled_providers denylist, and a model pin
# does not bypass that gate -- a filtered provider resolves to ModelNotFoundError.
diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py
index 2405ba0480..5b2806be12 100644
--- a/unsloth_cli/tests/test_start.py
+++ b/unsloth_cli/tests/test_start.py
@@ -2190,8 +2190,18 @@ def test_yolo_aliases_are_interchangeable(fake_studio, alias):
assert "--dangerously-bypass-approvals-and-sandbox" in codex.output
assert "--dangerously-skip-permissions" not in codex.output
+ opencode = CliRunner().invoke(
+ start.start_app,
+ ["opencode", alias, "--no-launch", "run", "hello"],
+ )
+ assert opencode.exit_code == 0, opencode.output
+ assert _launch_command(opencode.output) == ["opencode", "run", "hello", "--auto"]
+ assert "permission" not in _opencode_inline_config(opencode.output)
-def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path):
+
+def test_yolo_opencode_bare_no_launch_uses_permission_fallback(fake_studio, tmp_path):
+ # A bare --no-launch recipe stays append-safe (callers add a subcommand later);
+ # `opencode --auto run ...` would select the TUI, not `run`, so keep the config fallback.
result = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"])
assert result.exit_code == 0, result.output
config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text())
@@ -2203,6 +2213,172 @@ def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path):
}
+def test_yolo_opencode_run_uses_native_auto(fake_studio):
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--yolo", "--no-launch", "run", "hello"],
+ )
+ assert result.exit_code == 0, result.output
+ command = _launch_command(result.output)
+ assert command == ["opencode", "run", "hello", "--auto"]
+ assert "permission" not in _opencode_inline_config(result.output)
+
+
+def test_yolo_opencode_tui_resume_uses_native_auto(fake_studio):
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--yolo", "--no-launch", "--session", "sid"],
+ )
+ assert result.exit_code == 0, result.output
+ command = _launch_command(result.output)
+ assert command == ["opencode", "--session", "sid", "--auto"]
+ assert "permission" not in _opencode_inline_config(result.output)
+
+
+def test_no_yolo_opencode_run_omits_native_auto(fake_studio):
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--no-launch", "run", "hello"],
+ )
+ assert result.exit_code == 0, result.output
+ assert _launch_command(result.output) == ["opencode", "run", "hello"]
+ assert "permission" not in _opencode_inline_config(result.output)
+
+
+def test_yolo_opencode_bare_launch_uses_native_auto(fake_studio, monkeypatch):
+ monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
+ monkeypatch.setattr(start, "_opencode_supports_native_auto", lambda: True)
+ captured = _capture_launch(monkeypatch, ["opencode", "--yolo"])
+ assert captured["command"][1:] == [
+ "--model",
+ f"{start._OPENCODE_PROVIDER}/{MODEL['id']}",
+ "--auto",
+ ]
+ assert "permission" not in json.loads(captured["env"]["OPENCODE_CONFIG_CONTENT"])
+
+
+def test_yolo_opencode_native_auto_clears_prior_config_fallback(fake_studio, tmp_path):
+ fallback = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--yolo", "--no-launch"],
+ )
+ assert fallback.exit_code == 0, fallback.output
+
+ native = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--yolo", "--no-launch", "run", "hello"],
+ )
+ assert native.exit_code == 0, native.output
+ assert _launch_command(native.output) == ["opencode", "run", "hello", "--auto"]
+ assert "permission" not in _opencode_inline_config(native.output)
+ config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text())
+ assert config["permission"] == {
+ "edit": "ask",
+ "bash": "ask",
+ "webfetch": "ask",
+ "external_directory": {"*": "ask"},
+ }
+
+
+@pytest.mark.parametrize(
+ ("version", "expected"),
+ [
+ ("1.17.11", False),
+ ("1.17.12", True),
+ ("opencode 1.18.2", True),
+ ("development build", False),
+ ],
+)
+def test_opencode_native_auto_version_gate(monkeypatch, version, expected):
+ monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
+ monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: version)
+ assert start._opencode_supports_native_auto() is expected
+
+
+def test_opencode_native_auto_assumes_current_without_local_binary(monkeypatch):
+ monkeypatch.setattr(start.shutil, "which", lambda _: None)
+ assert start._opencode_supports_native_auto() is True
+
+
+def test_yolo_opencode_old_version_uses_config_fallback(fake_studio, monkeypatch):
+ monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
+ monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: "1.17.11")
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--yolo", "--no-launch", "run", "hello"],
+ )
+ assert result.exit_code == 0, result.output
+ assert _launch_command(result.output) == ["opencode", "run", "hello"]
+ assert _opencode_inline_config(result.output)["permission"] == {
+ "edit": "allow",
+ "bash": "allow",
+ "webfetch": "allow",
+ "external_directory": {"*": "allow"},
+ }
+
+
+@pytest.mark.parametrize(
+ ("args", "expected", "native"),
+ [
+ ([], ["--auto"], True),
+ (["run", "hello"], ["run", "hello", "--auto"], True),
+ (
+ ["run", "hello", "--", "--literal"],
+ ["run", "hello", "--auto", "--", "--literal"],
+ True,
+ ),
+ (["--print-logs", "run", "hello"], ["--print-logs", "run", "hello", "--auto"], True),
+ (["--session", "serve"], ["--session", "serve", "--auto"], True),
+ (["serve"], ["serve"], False),
+ (["--print-logs", "serve"], ["--print-logs", "serve"], False),
+ (["run", "--auto", "hello"], ["run", "--auto", "hello"], True),
+ # Hidden commands that reject --auto fall back like the visible utility ones.
+ (["generate"], ["generate"], False),
+ (["console", "login"], ["console", "login"], False),
+ # --mini ignores --auto (runMini forces auto=false), so use the config fallback.
+ (["--mini"], ["--mini"], False),
+ (["--session", "sid", "--mini"], ["--session", "sid", "--mini"], False),
+ ],
+)
+def test_opencode_native_auto_args(args, expected, native):
+ assert start._opencode_native_auto_args(args, True) == (expected, native)
+ assert start._opencode_native_auto_args(args, False) == (args, False)
+
+
+def test_yolo_opencode_non_agent_subcommand_uses_config_fallback(fake_studio):
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--yolo", "--no-launch", "serve"],
+ )
+ assert result.exit_code == 0, result.output
+ command = _launch_command(result.output)
+ assert command == ["opencode", "serve"]
+ assert _opencode_inline_config(result.output)["permission"] == {
+ "edit": "allow",
+ "bash": "allow",
+ "webfetch": "allow",
+ "external_directory": {"*": "allow"},
+ }
+
+
+@pytest.mark.parametrize("passthrough", (["generate"], ["console", "login"], ["--mini"]))
+def test_yolo_opencode_no_auto_command_uses_config_fallback(fake_studio, passthrough):
+ # generate/console are hidden and reject --auto, --mini ignores it: none get --auto,
+ # all keep the config permission fallback.
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--yolo", "--no-launch", *passthrough],
+ )
+ assert result.exit_code == 0, result.output
+ assert _launch_command(result.output) == ["opencode", *passthrough]
+ assert _opencode_inline_config(result.output)["permission"] == {
+ "edit": "allow",
+ "bash": "allow",
+ "webfetch": "allow",
+ "external_directory": {"*": "allow"},
+ }
+
+
def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path):
result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"])
assert result.exit_code == 0, result.output
@@ -2549,15 +2725,16 @@ def test_openclaw_non_yolo_preserves_full_mode(tmp_path):
def test_yolo_command_flags_unmapped_agent_is_empty():
- # Config-based agents (and any typo) must yield no flag, not a KeyError.
+ # Placement-aware/config-based agents (and any typo) must yield no prefix flag.
assert start._yolo_command_flags("opencode", True) == []
assert start._yolo_command_flags("openclaw", True) == []
assert start._yolo_command_flags("claude", True) == ["--dangerously-skip-permissions"]
assert start._yolo_command_flags("claude", False) == []
-def test_yolo_config_agents_add_no_command_flag(fake_studio):
- # opencode/openclaw auto-approve is config-only; nothing should leak onto argv.
+def test_yolo_config_fallbacks_add_no_legacy_command_flag(fake_studio):
+ # OpenClaw is config-only; OpenCode's append-safe bare recipe uses its config fallback.
+ # Neither should leak a legacy yolo/dangerous alias onto argv.
for agent in ("opencode", "openclaw"):
result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"])
assert result.exit_code == 0, result.output
From e0132b6d6c414cece2bced7eaf164eeebe088dd1 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 20 Jul 2026 05:04:28 +0100
Subject: [PATCH 023/255] Pin the Hermes remote installer and harden consent
(#7179)
Pin the fetched Hermes install.sh/install.ps1 and the checkout they perform to an immutable upstream commit, and distinguish pinned from unpinned sources in the consent warning.
---
unsloth_cli/commands/start.py | 54 +++++++++++++++++++++++++++------
unsloth_cli/tests/test_start.py | 38 ++++++++++++++++++-----
2 files changed, 75 insertions(+), 17 deletions(-)
diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py
index 8128447a02..6da31229f8 100644
--- a/unsloth_cli/commands/start.py
+++ b/unsloth_cli/commands/start.py
@@ -49,13 +49,22 @@ _HERMES_PROVIDER = "unsloth"
# the wizard's global API-key/model prompts would block the launch and point the
# user at a different (global) provider than the one Unsloth just configured.
# Both installers expose a skip flag: `-SkipSetup` (PowerShell) and
-# `--skip-setup` (POSIX; passed to the piped script via `bash -s --`).
+# `--skip-setup` (POSIX; passed to the piped script via `bash -s --`). Pin both
+# the fetched script and the repository checkout it performs to the same full
+# commit so a later change to either upstream branch cannot silently replace
+# code that Unsloth executes with the user's privileges.
+_HERMES_INSTALL_COMMIT = "f1af945f6c576eccb126fa955edc9be258b33020"
+_HERMES_INSTALL_BASE = (
+ "https://raw.githubusercontent.com/NousResearch/hermes-agent/"
+ f"{_HERMES_INSTALL_COMMIT}/scripts"
+)
_HERMES_WINDOWS_INSTALL_HINT = (
- "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup"
+ f"& ([scriptblock]::Create((irm {_HERMES_INSTALL_BASE}/install.ps1)))"
+ f" -SkipSetup -Commit {_HERMES_INSTALL_COMMIT}"
)
_HERMES_POSIX_INSTALL_HINT = (
- "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent"
- "/main/scripts/install.sh | bash -s -- --skip-setup"
+ f"curl -fsSL {_HERMES_INSTALL_BASE}/install.sh | bash -s --"
+ f" --skip-setup --commit {_HERMES_INSTALL_COMMIT}"
)
# Hermes refuses to initialize when the model window is under 64,000 tokens; its
# error message points at the model.context_length / auxiliary.compression
@@ -1199,6 +1208,16 @@ def _install_source(install_hint: str) -> Optional[str]:
return match.group(0) if match else None
+def _pinned_raw_github_commit(source: str) -> Optional[str]:
+ """Return the immutable full commit in a raw GitHub URL, if present."""
+ match = re.match(
+ r"^https://raw\.githubusercontent\.com/[^/]+/[^/]+/([0-9a-f]{40})/",
+ source,
+ flags = re.IGNORECASE,
+ )
+ return match.group(1).lower() if match else None
+
+
def _install_agent(name: str, install_hint: str) -> Optional[str]:
# Missing agent under --launch: offer to run its documented install command, then
# re-resolve it on PATH. Consent-based (we never auto-run a remote install script
@@ -1212,12 +1231,27 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]:
# and nothing checks a signature or hash on the fetched content. Naming the source
# turns a blind "yes" into informed consent.
source = _install_source(install_hint)
- warning = (
- f"This will download and RUN a script from {source} with your privileges"
- if source
- else f"This will RUN `{install_hint}` with your privileges"
- )
- typer.secho(f"{warning}; there is no signature or hash check.", fg = "yellow", err = True)
+ if source:
+ pinned_commit = _pinned_raw_github_commit(source)
+ if pinned_commit:
+ warning = (
+ "Security warning: This will download and execute a third-party script "
+ f"from {source} with your privileges. Unsloth pins this content to "
+ f"immutable upstream commit {pinned_commit}, but does not independently "
+ "verify or sandbox it. Continue only if you trust this source and commit."
+ )
+ else:
+ warning = (
+ "Security warning: This will download and execute an unverified third-party "
+ f"script from {source} with your privileges. Unsloth does not pin or verify "
+ "the downloaded content. Continue only if you trust this source."
+ )
+ else:
+ warning = (
+ f"This will RUN `{install_hint}` with your privileges; "
+ "there is no signature or hash check."
+ )
+ typer.secho(warning, fg = "yellow", err = True)
if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False):
return None
# Run each hint through the shell it is written for: PowerShell (irm | iex, or npm)
diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py
index 5b2806be12..98bd9f8157 100644
--- a/unsloth_cli/tests/test_start.py
+++ b/unsloth_cli/tests/test_start.py
@@ -7,6 +7,7 @@ from __future__ import annotations
import json
import os
+import re
import shlex
import sys
import urllib.error
@@ -128,7 +129,7 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch):
assert ran == [["powershell", "-NoProfile", "-Command", install_hint]]
-def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys):
+def test_install_agent_warns_remote_installer_is_unverified_third_party(monkeypatch, capsys):
# Before the confirm, a remote installer must name the URL it fetches so the
# user consents to a specific source rather than blindly accepting.
monkeypatch.setattr(start.os, "name", "nt")
@@ -137,9 +138,23 @@ def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys):
hint = "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup"
assert start._install_agent("hermes", hint) is None
err = capsys.readouterr().err
+ assert "Security warning" in err
+ assert "unverified third-party script" in err
assert "https://hermes-agent.nousresearch.com/install.ps1" in err
- assert "download and RUN" in err
- assert "signature or hash" in err
+ assert "Unsloth does not pin or verify the downloaded content" in err
+ assert "Continue only if you trust this source" in err
+
+
+def test_install_agent_reports_immutable_remote_installer_pin(monkeypatch, capsys):
+ monkeypatch.setattr(start.os, "name", "posix")
+ monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
+ monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False)
+ assert start._install_agent("hermes", start._HERMES_POSIX_INSTALL_HINT) is None
+ err = capsys.readouterr().err
+ assert start._HERMES_INSTALL_COMMIT in err
+ assert "immutable upstream commit" in err
+ assert "does not independently verify or sandbox it" in err
+ assert "does not pin or verify" not in err
def test_install_agent_warns_for_package_installer(monkeypatch, capsys):
@@ -160,8 +175,8 @@ def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch):
# Scriptblock form so `-SkipSetup` reaches the installer and the interactive
# setup wizard is skipped during the unattended `unsloth start hermes` run.
assert start._hermes_install_hint() == (
- "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1)))"
- " -SkipSetup"
+ f"& ([scriptblock]::Create((irm {start._HERMES_INSTALL_BASE}/install.ps1)))"
+ f" -SkipSetup -Commit {start._HERMES_INSTALL_COMMIT}"
)
@@ -170,11 +185,20 @@ def test_hermes_install_hint_is_bash_on_posix(monkeypatch):
# `bash -s -- --skip-setup` forwards the skip flag to the piped installer.
assert start._hermes_install_hint() == (
- "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent"
- "/main/scripts/install.sh | bash -s -- --skip-setup"
+ f"curl -fsSL {start._HERMES_INSTALL_BASE}/install.sh | bash -s --"
+ f" --skip-setup --commit {start._HERMES_INSTALL_COMMIT}"
)
+def test_hermes_install_hints_pin_script_and_checkout_to_full_commit():
+ commit = start._HERMES_INSTALL_COMMIT
+ assert re.fullmatch(r"[0-9a-f]{40}", commit)
+ for hint in (start._HERMES_WINDOWS_INSTALL_HINT, start._HERMES_POSIX_INSTALL_HINT):
+ assert hint.count(commit) == 2
+ assert "/main/" not in hint
+ assert "hermes-agent.nousresearch.com" not in hint
+
+
def test_refresh_windows_path_noop_off_windows(monkeypatch):
monkeypatch.setattr(start.os, "name", "posix")
before = os.environ.get("PATH", "")
From 39497e6516bdc7d7edc2b09493b222c1dc2ce49c Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 20 Jul 2026 05:05:22 +0100
Subject: [PATCH 024/255] Translate PWD for WSL-launched Windows agents (#7111)
Bridge PWD through WSLENV /p when launching a Windows npm shim from WSL so project-root discovery uses the live cwd. The no-launch recipe adds PWD/p without freezing PWD; the concrete cwd override applies only on direct launch.
---
unsloth_cli/commands/start.py | 16 ++++++++++++++--
unsloth_cli/tests/test_start.py | 16 ++++++++++++++--
2 files changed, 28 insertions(+), 4 deletions(-)
diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py
index 6da31229f8..707c2b3c90 100644
--- a/unsloth_cli/commands/start.py
+++ b/unsloth_cli/commands/start.py
@@ -1274,6 +1274,16 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]:
return executable
+def _wsl_shim_env(command: list, env: dict, unset_env: tuple) -> tuple[dict, tuple]:
+ wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else ()
+ if not wsl_env_bridge:
+ return env, wsl_env_bridge
+ # Bridge PWD via WSLENV (PWD/p) so the Windows shim finds its project root from the
+ # live cwd, not a stale inherited Linux PWD. Don't freeze env["PWD"]: a --no-launch
+ # recipe must translate the live PWD when run, not when generated; _launch overrides it.
+ return env, (*wsl_env_bridge, "PWD/p")
+
+
def _launch(
command: list,
env: dict,
@@ -1283,9 +1293,11 @@ def _launch(
executable = shutil.which(command[0]) or _install_agent(command[0], install_hint)
if executable is None:
_fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}")
- wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else ()
+ env, wsl_env_bridge = _wsl_shim_env(command, env, unset_env)
child_env = dict(os.environ)
if wsl_env_bridge:
+ # Override stale inherited PWD with the real cwd so the shim resolves the project root.
+ env = {**env, "PWD": os.getcwd()}
child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge)
for name in unset_env:
child_env[name] = ""
@@ -1353,8 +1365,8 @@ def _run(
if launch and clear_screen:
click.clear()
typer.echo(f"Unsloth {base} · model {entry['id']}")
- wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else ()
if not launch:
+ env, wsl_env_bridge = _wsl_shim_env(command, env, unset_env)
_print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge)
return
try:
diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py
index 98bd9f8157..227918f63e 100644
--- a/unsloth_cli/tests/test_start.py
+++ b/unsloth_cli/tests/test_start.py
@@ -476,8 +476,10 @@ def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypa
reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); "
"os.name is 'posix' under WSL, so this path can't run on a native Windows runner.",
)
-def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch):
+def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch, tmp_path):
captured = {}
+ monkeypatch.chdir(tmp_path)
+ monkeypatch.setenv("PWD", "/stale/outer/repo")
monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale")
monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale")
@@ -505,6 +507,8 @@ def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypat
assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface"
assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE
assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"]
+ assert captured["env"]["PWD"] == str(tmp_path)
+ assert "PWD/p" in captured["env"]["WSLENV"].split(":")
for name in (
"ANTHROPIC_AUTH_TOKEN",
"ANTHROPIC_BASE_URL",
@@ -520,7 +524,11 @@ def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypat
reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); "
"os.name is 'posix' under WSL, so this path can't run on a native Windows runner.",
)
-def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch):
+def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(
+ fake_studio, monkeypatch, tmp_path
+):
+ monkeypatch.chdir(tmp_path)
+ monkeypatch.setenv("PWD", "/stale/outer/repo")
monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
monkeypatch.setattr(
start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude"
@@ -532,6 +540,10 @@ def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studi
assert "export ANTHROPIC_API_KEY=" in result.output
assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output
assert "export WSLENV=" in result.output
+ # PWD must NOT be frozen into the recipe (no `export PWD=`): WSLENV PWD/p translates the
+ # shell's live PWD at run time, so a recipe reused from another dir resolves the project root.
+ assert "export PWD=" not in result.output
+ assert "PWD/p" in result.output
assert "ANTHROPIC_AUTH_TOKEN" in result.output
assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output
From 95d9970233ff3f248c9a5f89a987084e088623e9 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Mon, 20 Jul 2026 12:42:42 +0530
Subject: [PATCH 025/255] persist llama.cpp KV cache across idle auto-unload
(slot save/restore) (#7204)
* Studio: persist llama.cpp KV cache across idle auto-unload (slot save/restore)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: address KV persistence review feedback
* Studio: guard KV restore on launch config
* Studio: fix KV resume purge race, fingerprint requested ctx, purge on disable
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: re-check idle/keep-KV settings after slot save, ns file identity
* Studio: shard-aware KV guard, honor user --no-cache-prompt, early save cap
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: honor LLAMA_ARG_CACHE_PROMPT env in slot-save guard
* Studio: derive prompt-cache state from final argv for slot saves
* Studio: stat LoRA/control-vector sidecars in KV restore fingerprint
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: parse csv and FNAME:SCALE sidecar syntax in KV fingerprint
* Studio: address codex review on idle-unload KV resume
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: harden slot-save cleanup, cap accounting, stale-KV guard, save timeout
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: treat unavailable KV estimate as full-cap for slot-save disk check
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
studio/backend/core/inference/llama_cpp.py | 293 +++++++++
.../backend/core/inference/llama_keepwarm.py | 138 +++-
.../core/inference/llama_server_args.py | 2 +
studio/backend/main.py | 3 +-
studio/backend/routes/inference.py | 2 +-
studio/backend/routes/settings.py | 21 +-
.../tests/test_llama_cpp_mtp_detection.py | 19 +
.../tests/test_llama_cpp_slot_resume.py | 494 +++++++++++++++
.../backend/tests/test_llama_server_args.py | 12 +
.../backend/tests/test_openai_auto_switch.py | 588 +++++++++++++++++-
.../utils/openai_auto_switch_settings.py | 61 +-
studio/backend/utils/paths/storage_roots.py | 5 +
.../settings/api/openai-auto-switch.ts | 19 +-
.../components/model-auto-switch-section.tsx | 26 +-
studio/frontend/src/i18n/locales/en.ts | 3 +
15 files changed, 1644 insertions(+), 42 deletions(-)
create mode 100644 studio/backend/tests/test_llama_cpp_slot_resume.py
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index d7c7eed518..d4cab81bdf 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -23,6 +23,7 @@ import subprocess
import sys
import threading
import time
+import uuid
from pathlib import Path
from typing import (
Callable,
@@ -453,6 +454,23 @@ def _hf_offline_if_dns_dead():
os.environ.pop("TRANSFORMERS_OFFLINE", None)
+try:
+ _SLOT_SAVE_MAX_BYTES = int(os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30))
+except ValueError:
+ _SLOT_SAVE_MAX_BYTES = 10 << 30
+
+# The idle loop holds the lifecycle gate across a slot save, so a newly arriving
+# request waits on the in-flight save's HTTP call. Bound it (was 120s) so a slow
+# or stuck save can't stall the next request for minutes; best-effort save just
+# falls back to a plain unload. Override with UNSLOTH_SLOT_SAVE_TIMEOUT (seconds).
+try:
+ _SLOT_SAVE_HTTP_TIMEOUT = float(os.environ.get("UNSLOTH_SLOT_SAVE_TIMEOUT") or 30.0)
+except ValueError:
+ _SLOT_SAVE_HTTP_TIMEOUT = 30.0
+if _SLOT_SAVE_HTTP_TIMEOUT <= 0:
+ _SLOT_SAVE_HTTP_TIMEOUT = 30.0
+
+
def _swa_cache_path() -> Path:
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
base = Path(home) if home else Path.home() / ".unsloth" / "studio"
@@ -2000,6 +2018,12 @@ class LlamaCppBackend:
self._llama_log_path: Optional[Path] = None
self._cancel_event = threading.Event()
self._api_key: Optional[str] = None
+ self._slot_save_dir: Optional[str] = None
+ self._slot_save_binary: Optional[tuple[str, int]] = None
+ # (gguf_identity, launch_fingerprint) snapshotted at load, so a later slot
+ # save can tell whether the model files were swapped on disk since load.
+ self._slot_loaded_identity: Optional[tuple] = None
+ self._prompt_cache_disabled: bool = False
# True once a probe has completed; cleared on transient failure.
self._is_audio: bool = False
self._audio_type: Optional[str] = None
@@ -2638,6 +2662,7 @@ class LlamaCppBackend:
"supports_ctx_checkpoints": False,
"supports_no_cache_prompt": False,
"supports_metrics": False,
+ "supports_slot_save": False,
}
try:
mtime = int(Path(bin_path).stat().st_mtime)
@@ -2658,6 +2683,7 @@ class LlamaCppBackend:
supports_ctx_checkpoints = False
supports_no_cache_prompt = False
supports_metrics = False
+ supports_slot_save = False
try:
probe_env = cls._llama_server_env_for_binary(bin_path)
result = subprocess.run(
@@ -2756,6 +2782,7 @@ class LlamaCppBackend:
supports_ctx_checkpoints = _is_real("--ctx-checkpoints")
supports_no_cache_prompt = _is_real("--no-cache-prompt")
supports_metrics = _is_real("--metrics")
+ supports_slot_save = _is_real("--slot-save-path")
except (OSError, subprocess.SubprocessError) as exc:
logger.debug(f"llama-server --help probe failed: {exc}")
@@ -2773,6 +2800,7 @@ class LlamaCppBackend:
"supports_ctx_checkpoints": supports_ctx_checkpoints,
"supports_no_cache_prompt": supports_no_cache_prompt,
"supports_metrics": supports_metrics,
+ "supports_slot_save": supports_slot_save,
}
cls._capability_cache[cache_key] = info
return info
@@ -7332,6 +7360,26 @@ class LlamaCppBackend:
# when the binary advertises it (older/custom binaries may not).
if server_caps.get("supports_metrics"):
cmd.append("--metrics")
+ self._slot_save_dir = None
+ self._slot_save_binary = None
+ self._prompt_cache_disabled = False
+ if server_caps.get("supports_slot_save"):
+ try:
+ from utils.paths.storage_roots import ( # noqa: WPS433
+ llama_slot_cache_root,
+ )
+
+ slot_dir = llama_slot_cache_root()
+ slot_dir.mkdir(parents = True, exist_ok = True)
+ # Saved KV encodes chat content; keep it from other local users.
+ with contextlib.suppress(OSError):
+ os.chmod(slot_dir, 0o700)
+ cmd.extend(["--slot-save-path", str(slot_dir)])
+ self._slot_save_dir = str(slot_dir)
+ self._slot_save_binary = (binary, Path(binary).stat().st_mtime_ns)
+ except OSError:
+ self._slot_save_dir = None
+ self._slot_save_binary = None
cmd.extend(
self._ctx_integrity_flags(
n_parallel,
@@ -7529,6 +7577,7 @@ class LlamaCppBackend:
unsupported_cache_flags.append("--ctx-checkpoints")
if server_caps.get("supports_no_cache_prompt"):
cmd.append("--no-cache-prompt")
+ self._prompt_cache_disabled = True
else:
unsupported_cache_flags.append("--no-cache-prompt")
if unsupported_cache_flags:
@@ -8105,6 +8154,15 @@ class LlamaCppBackend:
if not self._healthy:
return False
+ # Snapshot the files the server actually loaded. If a GGUF shard or a
+ # LoRA/control-vector sidecar is swapped on disk afterwards while the
+ # old weights stay mapped, save_slots_for_resume() compares against
+ # this and refuses to persist KV that a reload could misapply.
+ if self._slot_save_dir:
+ self._slot_loaded_identity = (
+ self._gguf_file_identity(self._gguf_path),
+ self._slot_launch_fingerprint(),
+ )
return True
def _build_speculative_flags(
@@ -8690,6 +8748,10 @@ class LlamaCppBackend:
self._effective_context_length = None
self._max_context_length = None
self._reset_effective_parallel_slots()
+ self._slot_save_dir = None
+ self._slot_save_binary = None
+ self._slot_loaded_identity = None
+ self._prompt_cache_disabled = False
self._chat_template = None
self._chat_template_override = None
self._supports_reasoning = False
@@ -9216,6 +9278,237 @@ class LlamaCppBackend:
return False
return True
+ def _slot_launch_fingerprint(self) -> tuple:
+ # KV validity keys on extra args, stat'd sidecar weights, effective ctx.
+ sidecars = []
+ for path in self._sidecar_weight_files():
+ try:
+ st = os.stat(path)
+ sidecars.append((path, st.st_size, st.st_mtime_ns))
+ except OSError:
+ sidecars.append((path, None, None))
+ return (
+ tuple(self._extra_args or ()),
+ tuple(sidecars),
+ self._requested_n_ctx,
+ self._effective_context_length,
+ getattr(self, "_cache_type_kv", None),
+ self.effective_parallel_slots,
+ )
+
+ def _gguf_file_identity(self, path) -> Optional[tuple]:
+ # (size, mtime_ns) per shard: a split GGUF keys KV validity on every sibling.
+ p = Path(path)
+ paths = [p]
+ m = _SHARD_FULL_RE.match(p.name)
+ if m:
+ prefix, _first, total = m.groups()
+ paths = [
+ p.with_name(f"{prefix}-{i:05d}-of-{total}{p.suffix}")
+ for i in range(1, int(total) + 1)
+ ]
+ try:
+ return tuple((sp.stat().st_size, sp.stat().st_mtime_ns) for sp in paths)
+ except OSError:
+ return None
+
+ _SIDECAR_WEIGHT_FLAGS = (
+ "--lora",
+ "--lora-scaled",
+ "--control-vector",
+ "--control-vector-scaled",
+ )
+
+ def _sidecar_weight_files(self) -> list[str]:
+ # llama.cpp: comma-separated paths, FNAME:SCALE on -scaled (older builds: FNAME SCALE).
+ args = [str(a).strip() for a in (self._extra_args or ())]
+ files: list[str] = []
+ for i, arg in enumerate(args):
+ flag, sep, inline = arg.partition("=")
+ if flag not in self._SIDECAR_WEIGHT_FLAGS:
+ continue
+ operand = inline if sep else (args[i + 1] if i + 1 < len(args) else "")
+ if not operand:
+ continue
+ candidates = [operand]
+ pieces = [p for p in operand.split(",") if p]
+ if len(pieces) > 1:
+ candidates.extend(pieces)
+ if flag.endswith("-scaled"):
+ for item in list(candidates):
+ # ":" tail is a scale; rpartition spares drive letters.
+ head, colon, tail = item.rpartition(":")
+ if not (colon and head):
+ continue
+ try:
+ float(tail)
+ except ValueError:
+ continue
+ candidates.append(head)
+ for cand in candidates:
+ if cand not in files:
+ files.append(cand)
+ return files
+
+ def _prompt_cache_off(self) -> bool:
+ # Caching off makes restores useless; last prompt-cache flag wins, env only when unset.
+ last = None
+ for arg in self._extra_args or ():
+ flag = arg.strip().split("=", 1)[0]
+ if flag in ("--cache-prompt", "--no-cache-prompt"):
+ last = flag
+ if last is not None:
+ return last == "--no-cache-prompt"
+ if self._prompt_cache_disabled:
+ return True
+ if os.environ.get("LLAMA_ARG_NO_CACHE_PROMPT") is not None:
+ return True
+ env = (os.environ.get("LLAMA_ARG_CACHE_PROMPT") or "").strip().lower()
+ return env in {"off", "disabled", "false", "0"}
+
+ def save_slots_for_resume(
+ self, should_abort: Optional[Callable[[], bool]] = None
+ ) -> Optional[dict]:
+ if (
+ not self.is_loaded
+ or not self._slot_save_dir
+ or not self._gguf_path
+ or self._prompt_cache_off()
+ ):
+ return None
+ save_dir = Path(self._slot_save_dir)
+ gguf_stat = self._gguf_file_identity(self._gguf_path)
+ if gguf_stat is None:
+ return None
+ launch = self._slot_launch_fingerprint()
+ # If the GGUF or a sidecar was swapped on disk while the original weights
+ # stayed mapped, the live KV belongs to the old weights but a reload would
+ # load the new file. Persisting it would let restore misapply stale KV.
+ if self._slot_loaded_identity is not None and self._slot_loaded_identity != (
+ gguf_stat,
+ launch,
+ ):
+ logger.debug("Skipping slot save: model files changed on disk since load")
+ return None
+ try:
+ estimate = self._estimate_kv_cache_bytes(
+ self._effective_context_length or self._context_length or 0,
+ self._cache_type_kv,
+ n_parallel = self.effective_parallel_slots,
+ )
+ # Skip before writing anything when the estimate alone blows the cap,
+ # rather than fully writing a slot and discarding it afterwards.
+ if estimate > _SLOT_SAVE_MAX_BYTES:
+ logger.debug(
+ "Skipping slot save: estimated %d bytes exceeds cap %d",
+ estimate,
+ _SLOT_SAVE_MAX_BYTES,
+ )
+ return None
+ # A 0 estimate means metadata was insufficient, not a zero-byte cache:
+ # a slot can still be many GiB, so demand room for the whole cap before
+ # trusting the post-write check.
+ required = (estimate if estimate > 0 else _SLOT_SAVE_MAX_BYTES) + (1 << 30)
+ if shutil.disk_usage(save_dir).free < required:
+ logger.debug("Skipping slot save: insufficient free disk")
+ return None
+ except Exception:
+ pass
+ token = uuid.uuid4().hex[:8]
+ entries: list[dict] = []
+ total_bytes = 0
+ for slot in range(self.effective_parallel_slots):
+ # A request pending mid-save waits on the gate; stop wasting its time.
+ if should_abort is not None and should_abort():
+ break
+ filename = f"resume-{token}-slot{slot}.bin"
+ path = save_dir / filename
+ try:
+ resp = httpx.post(
+ f"{self.base_url}/slots/{slot}",
+ params = {"action": "save"},
+ json = {"filename": filename},
+ headers = self._auth_headers,
+ timeout = _SLOT_SAVE_HTTP_TIMEOUT,
+ trust_env = False,
+ )
+ except Exception as e:
+ logger.debug(f"slot {slot} save failed: {e}")
+ with contextlib.suppress(OSError):
+ path.unlink()
+ break
+ if resp.status_code != 200:
+ logger.debug(f"slot {slot} save returned HTTP {resp.status_code}")
+ with contextlib.suppress(OSError):
+ path.unlink()
+ continue
+ try:
+ body = resp.json()
+ if not isinstance(body, dict):
+ raise ValueError("slot save response was not a JSON object")
+ n_saved = int(body.get("n_saved") or 0)
+ except Exception as e:
+ # A 200 that still wrote a file but returns a malformed body must
+ # clean up like the transport/HTTP error paths above, or the file
+ # (which holds chat KV) is orphaned until the next startup sweep.
+ logger.debug(f"slot {slot} save returned an invalid response: {e}")
+ with contextlib.suppress(OSError):
+ path.unlink()
+ continue
+ if n_saved <= 0:
+ with contextlib.suppress(OSError):
+ path.unlink()
+ continue
+ # Account by the bytes actually on disk, not the server-reported
+ # count, so the cap holds even if a custom binary under-reports.
+ try:
+ n_written = path.stat().st_size
+ except OSError:
+ n_written = 0
+ total_bytes += n_written
+ entries.append({"id": slot, "filename": filename, "n_saved": n_saved})
+ if total_bytes > _SLOT_SAVE_MAX_BYTES:
+ break # already over the cap; the discard below cleans up
+ if not entries:
+ return None
+ if total_bytes > _SLOT_SAVE_MAX_BYTES:
+ logger.debug(
+ "Discarding slot save: %d bytes exceeds cap %d",
+ total_bytes,
+ _SLOT_SAVE_MAX_BYTES,
+ )
+ for entry in entries:
+ with contextlib.suppress(OSError):
+ (save_dir / entry["filename"]).unlink()
+ return None
+ return {
+ "dir": self._slot_save_dir,
+ "binary": self._slot_save_binary,
+ "gguf": str(self._gguf_path),
+ "gguf_stat": gguf_stat,
+ "launch": launch,
+ "slots": entries,
+ }
+
+ def restore_slots_for_resume(self, manifest: dict) -> None:
+ if not self.is_loaded or not self._slot_save_dir:
+ return
+ for entry in manifest.get("slots") or []:
+ try:
+ resp = httpx.post(
+ f"{self.base_url}/slots/{int(entry['id'])}",
+ params = {"action": "restore"},
+ json = {"filename": str(entry["filename"])},
+ headers = self._auth_headers,
+ timeout = _SLOT_SAVE_HTTP_TIMEOUT,
+ trust_env = False,
+ )
+ except Exception as e:
+ logger.debug(f"slot restore failed: {e}")
+ break
+ if resp.status_code != 200:
+ logger.debug(f"slot {entry.get('id')} restore returned HTTP {resp.status_code}")
+
def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool:
"""Schedule one background reload without MTP after a mid-generation death.
diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py
index 86a8c8a404..3380ebf5f5 100644
--- a/studio/backend/core/inference/llama_keepwarm.py
+++ b/studio/backend/core/inference/llama_keepwarm.py
@@ -15,6 +15,7 @@ import asyncio
import contextlib
import threading
import time
+from pathlib import Path
from loggers import get_logger
@@ -30,6 +31,8 @@ _last_active = time.monotonic()
# otherwise 503 against an empty backend can reload it (set on unload, cleared on
# reload). Storing the quant means the reload restores the exact freed variant.
_last_unloaded_model = None
+# Slot KV manifest saved by the idle unload; whoever pops it owns deleting its files.
+_kv_resume = None
# Guards inflight bumps against the idle-check-then-unload race, and blocks new
# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is
# shared across every event loop in the process, so a per-loop gate would let a
@@ -161,11 +164,17 @@ def inference_lifecycle_gate():
return _unload_gate()
-def note_model_loaded() -> None:
- """Record a successful GGUF load: stamp activity and drop any reload stash so
- a manual load clears it synchronously, not only on the next idle poll."""
+def note_model_loaded(backend = None) -> None:
+ """Stamp activity and synchronously drop any reload stash."""
_note_activity()
+ resume = take_kv_resume()
_set_last_unloaded(None)
+ if resume is None:
+ return
+ if backend is not None:
+ restore_kv_resume(backend, resume)
+ else:
+ _delete_resume_files(resume)
def note_model_unloaded() -> None:
@@ -182,9 +191,81 @@ def get_last_unloaded_model():
def _set_last_unloaded(value) -> None:
- global _last_unloaded_model
+ global _last_unloaded_model, _kv_resume
+ stale = None
with _lock:
_last_unloaded_model = value
+ if value is None and _kv_resume is not None:
+ stale, _kv_resume = _kv_resume, None
+ if stale:
+ _delete_resume_files(stale)
+
+
+def _delete_resume_files(manifest) -> None:
+ try:
+ base = Path(manifest.get("dir") or "")
+ for entry in manifest.get("slots") or []:
+ with contextlib.suppress(OSError):
+ (base / str(entry.get("filename"))).unlink()
+ except Exception:
+ pass
+
+
+def _set_kv_resume(value) -> None:
+ global _kv_resume
+ stale = None
+ with _lock:
+ if _kv_resume is not None and _kv_resume is not value:
+ stale = _kv_resume
+ _kv_resume = value
+ if stale:
+ _delete_resume_files(stale)
+
+
+def take_kv_resume():
+ global _kv_resume
+ with _lock:
+ manifest, _kv_resume = _kv_resume, None
+ return manifest
+
+
+def purge_kv_resume() -> None:
+ resume = take_kv_resume()
+ if resume:
+ _delete_resume_files(resume)
+
+
+def restore_kv_resume(backend, manifest) -> None:
+ try:
+ gguf = manifest.get("gguf")
+ binary = manifest.get("binary")
+ current = getattr(backend, "_gguf_path", None)
+ same_gguf = bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve()
+ if same_gguf:
+ # Same path is not enough: shards may have been rewritten meanwhile.
+ identity = getattr(backend, "_gguf_file_identity", None)
+ same_gguf = callable(identity) and identity(current) == manifest.get("gguf_stat")
+ if same_gguf:
+ # Nor the same file: launch overrides can invalidate KV numerics.
+ fingerprint = getattr(backend, "_slot_launch_fingerprint", None)
+ same_gguf = callable(fingerprint) and manifest.get("launch") == fingerprint()
+ if same_gguf and binary and binary == getattr(backend, "_slot_save_binary", None):
+ logger.info("Restoring saved slot KV onto the reloaded model")
+ backend.restore_slots_for_resume(manifest)
+ except Exception as exc:
+ logger.debug("slot restore after reload failed: %s", exc)
+ finally:
+ _delete_resume_files(manifest)
+
+
+def sweep_slot_save_dir() -> None:
+ try:
+ from utils.paths.storage_roots import llama_slot_cache_root
+ for path in llama_slot_cache_root().glob("resume-*.bin"):
+ with contextlib.suppress(OSError):
+ path.unlink()
+ except Exception:
+ pass
class LlamaKeepWarmMiddleware:
@@ -266,7 +347,10 @@ def _loaded_identity(backend):
async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
"""Unload the loaded GGUF once idle past the configured TTL. Inert when off."""
- from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds
+ from utils.openai_auto_switch_settings import (
+ get_auto_unload_idle_seconds,
+ get_auto_unload_keep_kv,
+ )
seen_model = None
while True:
@@ -281,17 +365,47 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
# Track by (id, variant): a (re)loaded model -- including the same repo
# at a different quant -- counts as activity so it survives one TTL
# before its first request (loads bypass the activity middleware).
- current = _loaded_identity(backend)
- if current != seen_model:
- seen_model = current
- if current is not None:
- _note_activity()
- _set_last_unloaded(None) # a model is loaded; drop stale stash
async with _unload_gate():
+ # Purging the stash mid-reload would race the restore.
+ current = _loaded_identity(backend)
+ if current != seen_model:
+ seen_model = current
+ if current is not None:
+ _note_activity()
+ _set_last_unloaded(None) # a model is loaded; drop stale stash
if backend.is_loaded and _is_idle(ttl):
freed = _loaded_identity(backend)
- await asyncio.to_thread(backend.unload_model)
+ manifest = None
+ if get_auto_unload_keep_kv():
+ try:
+ manifest = await asyncio.to_thread(
+ backend.save_slots_for_resume,
+ lambda: not _is_idle(ttl),
+ )
+ except Exception as exc:
+ logger.debug("slot save before idle unload failed: %s", exc)
+ # Re-read settings: the save can outlive a settings change.
+ ttl = get_auto_unload_idle_seconds()
+ if ttl <= 0 or not _is_idle(ttl):
+ if manifest:
+ _delete_resume_files(manifest)
+ continue
+ if manifest and not get_auto_unload_keep_kv():
+ _delete_resume_files(manifest)
+ manifest = None
+ try:
+ await asyncio.to_thread(backend.unload_model)
+ except Exception:
+ # Failed unload means nothing will stash the manifest.
+ if manifest:
+ _delete_resume_files(manifest)
+ raise
_set_last_unloaded(freed) # let an alias request reload it
+ if manifest and freed:
+ _set_kv_resume({"identity": freed, **manifest})
+ logger.info("Idle auto-unload: saved slot KV for restore on reload")
+ elif manifest:
+ _delete_resume_files(manifest)
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)
seen_model = None
except Exception as exc:
diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py
index e72e10e071..7b42d2f40d 100644
--- a/studio/backend/core/inference/llama_server_args.py
+++ b/studio/backend/core/inference/llama_server_args.py
@@ -70,6 +70,8 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# llama-server's own built-in tools flag would silently stack on top of
# Unsloth's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
+ # Slot-state dir: Studio owns it for KV persistence across idle unload.
+ frozenset({"--slot-save-path"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
diff --git a/studio/backend/main.py b/studio/backend/main.py
index 81d4c16e52..a1ff4d60da 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -547,8 +547,9 @@ async def lifespan(app: FastAPI):
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
# Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set).
- from core.inference.llama_keepwarm import idle_unload_loop
+ from core.inference.llama_keepwarm import idle_unload_loop, sweep_slot_save_dir
+ sweep_slot_save_dir()
app.state.idle_unload_task = asyncio.create_task(idle_unload_loop())
# Initialize RSA key pair for API key encryption (external providers).
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 136e4f7645..9c08ea4b79 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -4710,7 +4710,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Clear any idle-unload reload stash now, not only on the next poll.
from core.inference.llama_keepwarm import note_model_loaded
- note_model_loaded()
+ await asyncio.to_thread(note_model_loaded, llama_backend)
# A plain load advertises its own identifier; auto-switch overwrites
# this with the repo id right after _load_model_impl returns.
llama_backend._openai_advertised_id = None
diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py
index ab0fd2fd99..17e64df918 100644
--- a/studio/backend/routes/settings.py
+++ b/studio/backend/routes/settings.py
@@ -36,9 +36,10 @@ from utils.helper_precache_settings import (
)
from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents
from utils.openai_auto_switch_settings import (
- DEFAULT_AUTO_UNLOAD_IDLE_SECONDS,
+ DEFAULT_AUTO_UNLOAD_KEEP_KV,
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
get_auto_unload_idle_seconds,
+ get_auto_unload_keep_kv,
get_model_overrides,
get_openai_auto_switch_enabled,
get_stored_auto_unload_idle_seconds,
@@ -90,7 +91,9 @@ class HelperPrecacheResponse(BaseModel):
class OpenAIAutoSwitchPayload(BaseModel):
enabled: bool
- auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0)
+ # None leaves the stored value untouched (partial updates can't clobber it).
+ auto_unload_idle_seconds: Optional[int] = Field(default = None, ge = 0)
+ auto_unload_keep_kv: Optional[bool] = None
class OpenAIAutoSwitchResponse(BaseModel):
@@ -101,6 +104,7 @@ class OpenAIAutoSwitchResponse(BaseModel):
# UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled
# is false, so the UI can show idle-unload as active instead of "needs enable".
idle_unload_active: bool = False
+ auto_unload_keep_kv: bool = DEFAULT_AUTO_UNLOAD_KEEP_KV
class ModelOverridePayload(BaseModel):
@@ -198,6 +202,7 @@ def get_openai_auto_switch(
enabled = get_openai_auto_switch_enabled(),
auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(),
idle_unload_active = get_auto_unload_idle_seconds() > 0,
+ auto_unload_keep_kv = get_auto_unload_keep_kv(),
)
@@ -206,8 +211,8 @@ def update_openai_auto_switch(
payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject)
) -> OpenAIAutoSwitchResponse:
try:
- enabled, idle_seconds = set_openai_auto_switch(
- payload.enabled, payload.auto_unload_idle_seconds
+ enabled, idle_seconds, keep_kv = set_openai_auto_switch(
+ payload.enabled, payload.auto_unload_idle_seconds, payload.auto_unload_keep_kv
)
except ValueError as exc:
raise log_and_http_error(
@@ -217,10 +222,16 @@ def update_openai_auto_switch(
event = "settings.update_openai_auto_switch_failed",
log = logger,
) from exc
+ idle_unload_active = get_auto_unload_idle_seconds() > 0
+ if not keep_kv or not idle_unload_active:
+ # Keep-KV off or idle unload disabled: drop already-saved chat context too.
+ from core.inference.llama_keepwarm import purge_kv_resume
+ purge_kv_resume()
return OpenAIAutoSwitchResponse(
enabled = enabled,
auto_unload_idle_seconds = idle_seconds,
- idle_unload_active = get_auto_unload_idle_seconds() > 0,
+ idle_unload_active = idle_unload_active,
+ auto_unload_keep_kv = keep_kv,
)
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index 8fe04c0e39..68b706ebf9 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -741,6 +741,25 @@ def test_probe_reports_windows_cache_flags_absent_for_older_binary(tmp_path):
assert caps["supports_no_cache_prompt"] is False
+@_NEEDS_BASH
+def test_probe_detects_slot_save_path(tmp_path):
+ fake = _make_fake_llama_server(
+ tmp_path / "llama-server",
+ "--slot-save-path PATH path to save slot kv cache\n--threads N\n",
+ )
+ _clear_caps_cache()
+ caps = LlamaCppBackend.probe_server_capabilities(str(fake))
+ assert caps["supports_slot_save"] is True
+
+
+@_NEEDS_BASH
+def test_probe_reports_slot_save_absent_for_older_binary(tmp_path):
+ fake = _make_fake_llama_server(tmp_path / "llama-server", "--threads N\n")
+ _clear_caps_cache()
+ caps = LlamaCppBackend.probe_server_capabilities(str(fake))
+ assert caps["supports_slot_save"] is False
+
+
def test_build_ngram_mod_flags_new():
flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"})
assert flags == [
diff --git a/studio/backend/tests/test_llama_cpp_slot_resume.py b/studio/backend/tests/test_llama_cpp_slot_resume.py
new file mode 100644
index 0000000000..8b20c952c4
--- /dev/null
+++ b/studio/backend/tests/test_llama_cpp_slot_resume.py
@@ -0,0 +1,494 @@
+# 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 os
+from types import SimpleNamespace
+
+import core.inference.llama_cpp as llama_cpp
+from core.inference.llama_cpp import LlamaCppBackend
+
+
+def _resume_backend(tmp_path, n_slots = 1):
+ backend = LlamaCppBackend()
+ backend._healthy = True
+ # No-op lifecycle methods so the atexit cleanup can kill the fake quietly.
+ backend._process = SimpleNamespace(
+ poll = lambda: None,
+ terminate = lambda: None,
+ wait = lambda *a, **k: 0,
+ kill = lambda: None,
+ pid = 0,
+ )
+ backend._port = 8081
+ backend._slot_save_dir = str(tmp_path)
+ backend._slot_save_binary = ("/bin/llama-server", 1)
+ (tmp_path / "model.gguf").write_bytes(b"gguf")
+ backend._gguf_path = str(tmp_path / "model.gguf")
+ backend._effective_parallel_slots = n_slots
+ backend._estimate_kv_cache_bytes = lambda *a, **k: 0
+ return backend
+
+
+def _fake_disk(monkeypatch, free = 1 << 40):
+ monkeypatch.setattr(llama_cpp.shutil, "disk_usage", lambda _p: SimpleNamespace(free = free))
+
+
+class _Resp:
+ def __init__(
+ self,
+ status_code = 200,
+ body = None,
+ ):
+ self.status_code = status_code
+ self._body = body or {}
+
+ def json(self):
+ return self._body
+
+
+def test_save_returns_none_when_slot_save_disabled(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ backend._slot_save_dir = None
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
+
+
+def test_save_skipped_when_prompt_cache_disabled(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ backend._prompt_cache_disabled = True
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
+
+
+def test_save_skipped_when_insufficient_free_disk(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40
+ _fake_disk(monkeypatch, free = 1 << 20)
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
+
+
+def test_save_collects_manifest_across_slots(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path, n_slots = 2)
+ _fake_disk(monkeypatch)
+ calls = []
+
+ def fake_post(url, **kwargs):
+ calls.append((url, kwargs["params"], kwargs["json"]))
+ return _Resp(200, {"n_saved": 40, "n_written": 100})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ manifest = backend.save_slots_for_resume()
+ assert manifest is not None
+ assert manifest["dir"] == str(tmp_path)
+ assert manifest["binary"] == ("/bin/llama-server", 1)
+ assert manifest["gguf"] == str(tmp_path / "model.gguf")
+ st = os.stat(manifest["gguf"])
+ assert manifest["gguf_stat"] == ((st.st_size, st.st_mtime_ns),)
+ assert manifest["launch"] == backend._slot_launch_fingerprint()
+ assert [e["id"] for e in manifest["slots"]] == [0, 1]
+ assert all(e["n_saved"] == 40 for e in manifest["slots"])
+ assert [c[1] for c in calls] == [{"action": "save"}] * 2
+ assert "/slots/0" in calls[0][0] and "/slots/1" in calls[1][0]
+
+
+def test_save_unlinks_empty_slot_and_returns_none(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ _fake_disk(monkeypatch)
+
+ def fake_post(url, **kwargs):
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"")
+ return _Resp(200, {"n_saved": 0, "n_written": 0})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None
+ assert list(tmp_path.glob("resume-*.bin")) == [] # empty-slot file removed
+
+
+def test_save_cap_breach_discards_all_files(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path, n_slots = 2)
+ _fake_disk(monkeypatch)
+ monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
+
+ def fake_post(url, **kwargs):
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100)
+ return _Resp(200, {"n_saved": 40, "n_written": 100})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None # 200 bytes > 150 cap
+ assert list(tmp_path.glob("resume-*.bin")) == []
+
+
+def test_save_transport_error_aborts_remaining_slots(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path, n_slots = 3)
+ _fake_disk(monkeypatch)
+ calls = []
+
+ def fake_post(url, **kwargs):
+ calls.append(url)
+ raise OSError("connection refused")
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None
+ assert len(calls) == 1 # no retries against a dead server
+
+
+def test_save_transport_error_unlinks_partial_file(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ _fake_disk(monkeypatch)
+
+ def fake_post(url, **kwargs):
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"partial")
+ raise OSError("timed out")
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None
+ assert list(tmp_path.glob("resume-*.bin")) == []
+
+
+def test_fingerprint_tracks_lora_sidecar_rewrite(tmp_path):
+ backend = _resume_backend(tmp_path)
+ adapter = tmp_path / "adapter.gguf"
+ adapter.write_bytes(b"v1")
+ backend._extra_args = ["--lora", str(adapter)]
+
+ before = backend._slot_launch_fingerprint()
+ adapter.write_bytes(b"v2-different") # re-exported adapter, same path
+ assert backend._slot_launch_fingerprint() != before
+
+ backend._extra_args = [f"--lora={adapter}"]
+ assert backend._sidecar_weight_files() == [str(adapter)]
+ backend._extra_args = ["--lora-scaled", str(adapter), "0.5"]
+ assert backend._sidecar_weight_files() == [str(adapter)]
+ backend._extra_args = ["--control-vector", str(adapter), "--threads", "4"]
+ assert backend._sidecar_weight_files() == [str(adapter)]
+
+
+def test_sidecar_files_parse_csv_and_colon_scale(tmp_path):
+ backend = _resume_backend(tmp_path)
+ a, b = tmp_path / "a.gguf", tmp_path / "b.gguf"
+
+ backend._extra_args = ["--lora", f"{a},{b}"]
+ files = backend._sidecar_weight_files()
+ assert str(a) in files and str(b) in files
+
+ backend._extra_args = ["--lora-scaled", f"{a}:0.5"]
+ assert str(a) in backend._sidecar_weight_files()
+
+ backend._extra_args = ["--control-vector-scaled", f"{a}:1.0,{b}:2.0"]
+ files = backend._sidecar_weight_files()
+ assert str(a) in files and str(b) in files
+
+ # Windows drive letter must not be mistaken for a scale separator.
+ backend._extra_args = ["--lora-scaled", "C:\\adapters\\a.gguf:0.75"]
+ assert "C:\\adapters\\a.gguf" in backend._sidecar_weight_files()
+ backend._extra_args = ["--lora", "C:\\adapters\\a.gguf"]
+ assert backend._sidecar_weight_files() == ["C:\\adapters\\a.gguf"]
+
+
+def test_fingerprint_tracks_colon_scaled_adapter_rewrite(tmp_path):
+ backend = _resume_backend(tmp_path)
+ adapter = tmp_path / "adapter.gguf"
+ adapter.write_bytes(b"v1")
+ backend._extra_args = ["--lora-scaled", f"{adapter}:0.5"]
+
+ before = backend._slot_launch_fingerprint()
+ adapter.write_bytes(b"v2-different") # re-exported adapter, same path
+ assert backend._slot_launch_fingerprint() != before
+
+
+def test_fingerprint_tracks_effective_context_length(tmp_path):
+ backend = _resume_backend(tmp_path)
+ backend._effective_context_length = 8192
+
+ before = backend._slot_launch_fingerprint()
+ backend._effective_context_length = 4096 # auto-fit landed smaller on reload
+ assert backend._slot_launch_fingerprint() != before
+
+
+def test_gguf_file_identity_covers_split_shards(tmp_path):
+ backend = _resume_backend(tmp_path)
+ first = tmp_path / "m-00001-of-00002.gguf"
+ second = tmp_path / "m-00002-of-00002.gguf"
+ first.write_bytes(b"a")
+ second.write_bytes(b"bb")
+
+ before = backend._gguf_file_identity(str(first))
+ st1, st2 = os.stat(first), os.stat(second)
+ assert before == ((st1.st_size, st1.st_mtime_ns), (st2.st_size, st2.st_mtime_ns))
+
+ second.write_bytes(b"rewritten") # sibling changes, primary untouched
+ after = backend._gguf_file_identity(str(first))
+ assert after is not None and after != before
+ assert after[0] == before[0] # primary shard unchanged
+
+ second.unlink()
+ assert backend._gguf_file_identity(str(first)) is None # missing shard
+
+
+def test_save_skipped_when_user_disabled_prompt_cache(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ backend._extra_args = ["--no-cache-prompt"]
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
+
+
+def test_save_skipped_when_env_disables_prompt_cache(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0")
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
+ monkeypatch.delenv("LLAMA_ARG_CACHE_PROMPT")
+ monkeypatch.setenv("LLAMA_ARG_NO_CACHE_PROMPT", "1") # legacy negative form
+ assert backend.save_slots_for_resume() is None
+
+
+def test_explicit_cache_prompt_flag_overrides_env(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0")
+ backend._extra_args = ["--cache-prompt"] # CLI wins over env in llama.cpp
+ _fake_disk(monkeypatch)
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is not None
+
+
+def test_user_cache_prompt_overrides_studio_no_cache_flag(monkeypatch, tmp_path):
+ # User extras follow Studio's flags, so an explicit --cache-prompt wins.
+ backend = _resume_backend(tmp_path)
+ backend._prompt_cache_disabled = True
+ backend._extra_args = ["--cache-prompt"]
+ _fake_disk(monkeypatch)
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is not None
+ # Last flag wins when both appear in extras.
+ backend._extra_args = ["--cache-prompt", "--no-cache-prompt"]
+ assert backend.save_slots_for_resume() is None
+
+
+def test_save_stops_writing_once_cap_exceeded(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path, n_slots = 3)
+ _fake_disk(monkeypatch)
+ monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
+ calls = []
+
+ def fake_post(url, **kwargs):
+ calls.append(url)
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100)
+ return _Resp(200, {"n_saved": 1, "n_written": 100})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None
+ assert len(calls) == 2 # cap blown after slot 1; slot 2 never attempted
+ assert list(tmp_path.glob("resume-*.bin")) == []
+
+
+def test_save_aborts_between_slots_when_no_longer_idle(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path, n_slots = 3)
+ _fake_disk(monkeypatch)
+ calls = []
+
+ def fake_post(url, **kwargs):
+ calls.append(url)
+ return _Resp(200, {"n_saved": 5, "n_written": 10})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ aborts = iter([False, True, True])
+ manifest = backend.save_slots_for_resume(should_abort = lambda: next(aborts))
+ assert len(calls) == 1 # slots 1 and 2 skipped
+ assert manifest is not None
+ assert [e["id"] for e in manifest["slots"]] == [0]
+
+
+def test_save_non_200_slot_is_skipped_but_others_kept(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path, n_slots = 2)
+ _fake_disk(monkeypatch)
+
+ def fake_post(url, **kwargs):
+ if "/slots/0" in url:
+ return _Resp(500)
+ return _Resp(200, {"n_saved": 5, "n_written": 10})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ manifest = backend.save_slots_for_resume()
+ assert manifest is not None
+ assert [e["id"] for e in manifest["slots"]] == [1]
+
+
+def test_restore_posts_each_slot_and_tolerates_failures(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ calls = []
+
+ def fake_post(url, **kwargs):
+ calls.append((url, kwargs["params"], kwargs["json"]))
+ return _Resp(500 if "/slots/0" in url else 200, {"n_restored": 5})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ backend.restore_slots_for_resume(
+ {
+ "slots": [
+ {"id": 0, "filename": "resume-a-slot0.bin", "n_saved": 5},
+ {"id": 1, "filename": "resume-a-slot1.bin", "n_saved": 5},
+ ]
+ }
+ )
+ assert [c[1] for c in calls] == [{"action": "restore"}] * 2
+ assert calls[0][2] == {"filename": "resume-a-slot0.bin"}
+
+
+def test_restore_transport_error_stops_early(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ calls = []
+
+ def fake_post(url, **kwargs):
+ calls.append(url)
+ raise OSError("connection refused")
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ backend.restore_slots_for_resume(
+ {"slots": [{"id": 0, "filename": "a.bin"}, {"id": 1, "filename": "b.bin"}]}
+ )
+ assert len(calls) == 1
+
+
+def test_save_deletes_orphan_on_malformed_response(monkeypatch, tmp_path):
+ # A 200 that writes a file but returns a non-numeric counter must be cleaned
+ # up like any other save failure, not left orphaned holding chat KV.
+ backend = _resume_backend(tmp_path)
+ _fake_disk(monkeypatch)
+
+ def fake_post(url, **kwargs):
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv")
+ return _Resp(200, {"n_saved": "not-an-int"})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None
+ assert list(tmp_path.glob("resume-*.bin")) == []
+
+
+def test_save_deletes_orphan_on_non_dict_response(monkeypatch, tmp_path):
+ backend = _resume_backend(tmp_path)
+ _fake_disk(monkeypatch)
+
+ def fake_post(url, **kwargs):
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv")
+ return _Resp(200, ["unexpected", "list"])
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None
+ assert list(tmp_path.glob("resume-*.bin")) == []
+
+
+def test_save_cap_uses_actual_file_size_not_reported_bytes(monkeypatch, tmp_path):
+ # A binary under-reporting n_written must not slip past the disk cap: the
+ # cap is enforced against the bytes actually on disk.
+ backend = _resume_backend(tmp_path)
+ _fake_disk(monkeypatch)
+ monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
+
+ def fake_post(url, **kwargs):
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 200)
+ return _Resp(200, {"n_saved": 5, "n_written": 1}) # under-reported
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ assert backend.save_slots_for_resume() is None # 200 real bytes > 150 cap
+ assert list(tmp_path.glob("resume-*.bin")) == []
+
+
+def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path):
+ # An estimate over the cap skips before writing any slot at all.
+ backend = _resume_backend(tmp_path)
+ backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40
+ monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 1 << 20)
+ _fake_disk(monkeypatch)
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
+
+
+def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path):
+ # The GGUF/sidecars were swapped on disk after the server loaded them, so the
+ # live KV belongs to the old weights: refuse to persist it (no POST at all).
+ backend = _resume_backend(tmp_path)
+ backend._slot_loaded_identity = ((("stale", 0),), ()) # != current identity
+ _fake_disk(monkeypatch)
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
+
+
+def test_save_proceeds_when_load_identity_matches(monkeypatch, tmp_path):
+ # Matching load-time snapshot: the save runs normally.
+ backend = _resume_backend(tmp_path)
+ backend._slot_loaded_identity = (
+ backend._gguf_file_identity(backend._gguf_path),
+ backend._slot_launch_fingerprint(),
+ )
+ _fake_disk(monkeypatch)
+
+ def fake_post(url, **kwargs):
+ (tmp_path / kwargs["json"]["filename"]).write_bytes(b"kv")
+ return _Resp(200, {"n_saved": 5, "n_written": 2})
+
+ monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
+ manifest = backend.save_slots_for_resume()
+ assert manifest is not None
+ assert [e["id"] for e in manifest["slots"]] == [0]
+
+
+def test_save_skipped_when_estimate_unavailable_and_low_disk(monkeypatch, tmp_path):
+ # A 0 estimate means metadata was insufficient, not a zero-byte cache: the save
+ # must demand room for the whole cap, not just 1 GiB, on a low-disk host.
+ backend = _resume_backend(tmp_path)
+ backend._estimate_kv_cache_bytes = lambda *a, **k: 0 # metadata unavailable
+ monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 8 << 30) # 8 GiB cap
+ _fake_disk(monkeypatch, free = 2 << 30) # 2 GiB free < 8 + 1 GiB required
+ monkeypatch.setattr(
+ llama_cpp.httpx,
+ "post",
+ lambda *a, **k: (_ for _ in ()).throw(AssertionError),
+ raising = False,
+ )
+ assert backend.save_slots_for_resume() is None
diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py
index c6d16363f8..fa4ba71791 100644
--- a/studio/backend/tests/test_llama_server_args.py
+++ b/studio/backend/tests/test_llama_server_args.py
@@ -183,6 +183,8 @@ def test_non_flag_token_passes_through():
"--reranking",
# llama-server's own --tools clashes with Unsloth's tool policy.
"--tools",
+ # Slot-state dir: Studio owns it for KV persistence across idle unload.
+ "--slot-save-path",
],
)
def test_denylist_rejects_all_aliases(denied):
@@ -224,6 +226,16 @@ def test_denylist_rejects_equals_form():
validate_extra_args(["--port=9000"])
+def test_slot_save_path_is_managed_in_all_forms():
+ for args in (["--slot-save-path", "/tmp/x"], ["--slot-save-path=/tmp/x"], ["--slot-save-path"]):
+ with pytest.raises(ValueError, match = "--slot-save-path"):
+ validate_extra_args(args)
+ assert is_managed_flag("--slot-save-path") is True
+ assert is_managed_flag("--slot-save-path=/tmp/x") is True
+ # --slots (read-only diagnostics endpoint) stays a user choice.
+ assert is_managed_flag("--slots") is False
+
+
@pytest.mark.parametrize(
"padded",
[" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"],
diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py
index c4c0ce15c9..1ee9ef36d3 100644
--- a/studio/backend/tests/test_openai_auto_switch.py
+++ b/studio/backend/tests/test_openai_auto_switch.py
@@ -8,6 +8,7 @@ tests/test_gguf_completion_usage.py.
"""
import asyncio
+import os
import pytest
@@ -18,6 +19,10 @@ from utils import openai_auto_switch_settings as settings
class _FakeBackend:
+ effective_parallel_slots = 1
+ _slot_save_binary = None
+ _gguf_path = None
+
def __init__(
self,
loaded_id = None,
@@ -29,6 +34,22 @@ class _FakeBackend:
self.hf_variant = hf_variant
self._openai_advertised_id = advertised_id
+ def save_slots_for_resume(self, should_abort = None):
+ return None
+
+ def restore_slots_for_resume(self, manifest):
+ return None
+
+ def _slot_launch_fingerprint(self):
+ return ((), None, None, 1)
+
+ def _gguf_file_identity(self, path):
+ try:
+ st = os.stat(path)
+ except OSError:
+ return None
+ return ((st.st_size, st.st_mtime_ns),)
+
class _LoadRecorder:
"""Stand-in for the load route: records calls and simulates a load."""
@@ -53,10 +74,15 @@ class _LoadRecorder:
from fastapi import HTTPException
raise HTTPException(status_code = 503, detail = "load failed")
self.backend.model_identifier = request.model_path
+ self.backend.hf_variant = getattr(request, "gguf_variant", None)
+ self.backend._gguf_path = request.model_path
self.backend.is_loaded = True
# Mirror _load_model_impl: a load advertises its own id until the
# auto-switch caller overwrites it with the repo id.
self.backend._openai_advertised_id = None
+ from core.inference import llama_keepwarm as kw
+
+ kw.note_model_loaded(self.backend)
return None
@@ -446,6 +472,75 @@ def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch):
assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M"
+def test_idle_loop_deletes_saved_kv_when_unload_fails(monkeypatch, tmp_path):
+ import time
+ from core.inference import llama_keepwarm as kw
+
+ monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
+ monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
+ kw._inflight = 0
+ kw._pending = 0
+ kw._last_active = time.monotonic() - 3600
+ kw._last_unloaded_model = None
+ kw._kv_resume = None
+
+ saved = tmp_path / "resume-abc-slot0.bin"
+ backend = _FakeBackend("unsloth/Idle-GGUF")
+ manifests = []
+
+ def _save(should_abort = None):
+ if manifests:
+ return None
+ saved.write_bytes(b"kv")
+ manifest = {"dir": str(tmp_path), "slots": [{"id": 0, "filename": saved.name}]}
+ manifests.append(manifest)
+ return manifest
+
+ def _unload():
+ raise RuntimeError("cuda teardown failed")
+
+ backend.save_slots_for_resume = _save
+ backend.unload_model = _unload
+ monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
+
+ async def _drive():
+ task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01))
+ for _ in range(200):
+ await asyncio.sleep(0.01)
+ if manifests and not saved.exists():
+ break
+ task.cancel()
+ try:
+ await task
+ except asyncio.CancelledError:
+ pass
+
+ asyncio.run(_drive())
+ assert manifests and not saved.exists()
+ assert kw._kv_resume is None
+
+
+def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path):
+ # PUT leaves keep-KV on but makes idle unload inactive: saved KV must go too.
+ import routes.settings as settings_route
+ from core.inference import llama_keepwarm as kw
+
+ saved = tmp_path / "resume-abc-slot0.bin"
+ saved.write_bytes(b"kv")
+ kw._kv_resume = {
+ "identity": ("m", None, "m"),
+ "dir": str(tmp_path),
+ "slots": [{"id": 0, "filename": saved.name}],
+ }
+ monkeypatch.setattr(settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True))
+ monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0)
+
+ payload = settings_route.OpenAIAutoSwitchPayload(enabled = False)
+ resp = settings_route.update_openai_auto_switch(payload, "tester")
+ assert resp.idle_unload_active is False and resp.auto_unload_keep_kv is True
+ assert kw._kv_resume is None and not saved.exists()
+
+
def test_audio_generate_is_tracked_as_inference_path():
# Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so
# the keep-warm middleware must count it as in-flight inference.
@@ -2912,8 +3007,10 @@ def test_non_gguf_load_clears_reload_stash():
# A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF
# branch, so it never lingers until the idle poll (or forever, idle-unload off).
import inspect
+
src = inspect.getsource(inference_route._load_model_impl)
- assert src.count("note_model_loaded()") >= 2
+ assert src.count("note_model_loaded()") >= 1 # non-GGUF branch
+ assert "to_thread(note_model_loaded, llama_backend)" in src # GGUF branch
def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch):
@@ -3121,6 +3218,495 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp
assert "Model auto-switch" in non_gguf_loaded
+# ── idle-unload KV persistence (slot save/restore) ──────────────────
+
+
+def _seed_kv_manifest(
+ tmp_path,
+ identity = ("unsloth/A-GGUF", "Q4_K_M", "unsloth/A-GGUF"),
+ gguf = None,
+):
+ if gguf is None:
+ gguf_file = tmp_path / "model.gguf"
+ gguf_file.write_bytes(b"gguf")
+ gguf = str(gguf_file)
+ st = os.stat(gguf)
+ state_file = tmp_path / "resume-abc-slot0.bin"
+ state_file.write_bytes(b"kv")
+ return state_file, {
+ "identity": identity,
+ "dir": str(tmp_path),
+ "binary": ("/bin/llama-server", 111),
+ "gguf": gguf,
+ "gguf_stat": ((st.st_size, st.st_mtime_ns),),
+ "launch": ((), None, None, 1),
+ "slots": [{"id": 0, "filename": state_file.name, "n_saved": 42}],
+ }
+
+
+def _drive_idle_loop(
+ kw,
+ poll_seconds = 0.02,
+ run_for = 0.2,
+):
+ async def _drive():
+ task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = poll_seconds))
+ await asyncio.sleep(run_for)
+ task.cancel()
+ try:
+ await task
+ except asyncio.CancelledError:
+ pass
+
+ asyncio.run(_drive())
+
+
+def test_idle_unload_saves_slots_before_unload_and_stashes_manifest(monkeypatch, tmp_path):
+ import time
+ from core.inference import llama_keepwarm as kw
+
+ monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
+ monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
+ kw._inflight = 0
+ kw._pending = 0
+ kw._last_active = time.monotonic() - 3600
+ kw._last_unloaded_model = None
+ kw._kv_resume = None
+
+ events = []
+ backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
+ manifest = {
+ "dir": str(tmp_path),
+ "binary": ("bin", 1),
+ "slots": [{"id": 0, "filename": "f.bin", "n_saved": 42}],
+ }
+
+ def _save(should_abort = None):
+ events.append("save")
+ return manifest
+
+ def _unload():
+ events.append("unload")
+ backend.is_loaded = False
+
+ backend.save_slots_for_resume = _save
+ backend.unload_model = _unload
+ monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
+
+ _drive_idle_loop(kw)
+ # KV must be saved while the server is still alive, then exactly one unload.
+ assert events == ["save", "unload"]
+ assert kw.get_last_unloaded_model()[:2] == ("unsloth/Idle-GGUF", "Q4_K_M")
+ resume = kw.take_kv_resume()
+ assert resume is not None
+ assert resume["identity"][:2] == ("unsloth/Idle-GGUF", "Q4_K_M")
+ assert resume["slots"][0]["filename"] == "f.bin"
+
+
+def test_idle_save_failure_still_unloads_plain(monkeypatch):
+ import time
+ from core.inference import llama_keepwarm as kw
+
+ monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
+ monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
+ kw._inflight = 0
+ kw._pending = 0
+ kw._last_active = time.monotonic() - 3600
+ kw._last_unloaded_model = None
+ kw._kv_resume = None
+
+ unloads = []
+ backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
+
+ def _save(should_abort = None):
+ raise RuntimeError("slot save exploded")
+
+ def _unload():
+ unloads.append(1)
+ backend.is_loaded = False
+
+ backend.save_slots_for_resume = _save
+ backend.unload_model = _unload
+ monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
+
+ _drive_idle_loop(kw)
+ assert unloads == [1] # the save failure must not skip the unload
+ assert kw.get_last_unloaded_model() is not None
+ assert kw.take_kv_resume() is None
+
+
+def test_keep_kv_setting_off_skips_save(monkeypatch):
+ import time
+ from core.inference import llama_keepwarm as kw
+
+ monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
+ monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: False)
+ kw._inflight = 0
+ kw._pending = 0
+ kw._last_active = time.monotonic() - 3600
+ kw._last_unloaded_model = None
+ kw._kv_resume = None
+
+ saves, unloads = [], []
+ backend = _FakeBackend("unsloth/Idle-GGUF")
+
+ def _unload():
+ unloads.append(1)
+ backend.is_loaded = False
+
+ backend.save_slots_for_resume = lambda *a, **k: saves.append(1)
+ backend.unload_model = _unload
+ monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
+
+ _drive_idle_loop(kw)
+ assert saves == []
+ assert unloads == [1]
+ assert kw.take_kv_resume() is None
+
+
+def test_keep_kv_disabled_mid_save_discards_manifest(monkeypatch, tmp_path):
+ import time
+ from core.inference import llama_keepwarm as kw
+
+ keep = {"on": True}
+ monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
+ monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: keep["on"])
+ kw._inflight = 0
+ kw._pending = 0
+ kw._last_active = time.monotonic() - 3600
+ kw._last_unloaded_model = None
+ kw._kv_resume = None
+
+ unloads = []
+ backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
+ state_file = tmp_path / "resume-mid-slot0.bin"
+ state_file.write_bytes(b"kv")
+ manifest = {
+ "dir": str(tmp_path),
+ "binary": ("bin", 1),
+ "slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}],
+ }
+
+ def _save(should_abort = None):
+ keep["on"] = False # user flips the toggle while the save runs
+ return manifest
+
+ def _unload():
+ unloads.append(1)
+ backend.is_loaded = False
+
+ backend.save_slots_for_resume = _save
+ backend.unload_model = _unload
+ monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
+
+ _drive_idle_loop(kw)
+ assert unloads == [1] # still unloads; only the stash is dropped
+ assert kw.take_kv_resume() is None
+ assert not state_file.exists()
+
+
+def test_idle_ttl_disabled_mid_save_skips_unload(monkeypatch, tmp_path):
+ import time
+ from core.inference import llama_keepwarm as kw
+
+ ttl = {"v": 0.005}
+ monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: ttl["v"])
+ monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
+ kw._inflight = 0
+ kw._pending = 0
+ kw._last_active = time.monotonic() - 3600
+ kw._last_unloaded_model = None
+ kw._kv_resume = None
+
+ unloads = []
+ backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
+ state_file = tmp_path / "resume-mid-slot0.bin"
+ state_file.write_bytes(b"kv")
+ manifest = {
+ "dir": str(tmp_path),
+ "binary": ("bin", 1),
+ "slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}],
+ }
+
+ def _save(should_abort = None):
+ ttl["v"] = 0 # user turns idle unload off while the save runs
+ return manifest
+
+ backend.save_slots_for_resume = _save
+ backend.unload_model = lambda: unloads.append(1)
+ monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
+
+ _drive_idle_loop(kw)
+ assert unloads == [] # the unload was cancelled by the setting change
+ assert kw.take_kv_resume() is None
+ assert not state_file.exists()
+
+
+def test_alias_reload_restores_slots_and_deletes_files(monkeypatch, tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ backend = _FakeBackend(None) # idle-unload emptied the backend
+ backend._slot_save_binary = ("/bin/llama-server", 111)
+ restored = []
+ backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
+
+ rec = _LoadRecorder(backend)
+ _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec)
+ monkeypatch.setattr(kw, "_inflight", 0)
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ monkeypatch.setattr(kw, "_last_unloaded_model", (manifest["gguf"], "Q4_K_M"))
+ monkeypatch.setattr(kw, "_kv_resume", manifest)
+
+ _run_hook("gpt-4o-mini")
+ assert len(rec.calls) == 1
+ assert len(restored) == 1 # same model + binary: restore ran
+ assert not state_file.exists() # state file deleted after the restore
+ assert kw._kv_resume is None
+
+
+def test_no_restore_when_different_model_loads(monkeypatch, tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ backend = _FakeBackend(None)
+ backend._slot_save_binary = ("/bin/llama-server", 111)
+ restored = []
+ backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
+ rec = _LoadRecorder(backend)
+ _wire(
+ monkeypatch,
+ enabled = True,
+ resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"),
+ backend = backend,
+ recorder = rec,
+ )
+ monkeypatch.setattr(kw, "_inflight", 0)
+ state_file, manifest = _seed_kv_manifest(tmp_path) # manifest is for model A
+ monkeypatch.setattr(kw, "_kv_resume", manifest)
+
+ _run_hook("unsloth/B-GGUF")
+ assert len(rec.calls) == 1
+ assert restored == [] # different model: never restored
+ assert not state_file.exists() # but the stale files are gone
+ assert kw._kv_resume is None
+
+
+def test_restore_skipped_when_binary_changed(monkeypatch, tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
+ backend._gguf_path = manifest["gguf"]
+ backend._slot_save_binary = ("/bin/llama-server", 222) # newer mtime
+ restored = []
+ backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
+
+ kw.restore_kv_resume(backend, manifest)
+ assert restored == []
+ assert not state_file.exists()
+
+
+def test_restore_skipped_when_launch_config_changed(tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
+ backend._gguf_path = manifest["gguf"]
+ backend._slot_save_binary = ("/bin/llama-server", 111)
+ backend._slot_launch_fingerprint = lambda: (("--rope-freq-scale", "0.5"), None, None, 1)
+ restored = []
+ backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
+
+ kw.restore_kv_resume(backend, manifest)
+ assert restored == []
+ assert not state_file.exists()
+
+
+def test_restore_skipped_when_gguf_rewritten_in_place(tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ with open(manifest["gguf"], "wb") as fh:
+ fh.write(b"different weights") # same path, new content
+ backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
+ backend._gguf_path = manifest["gguf"]
+ backend._slot_save_binary = ("/bin/llama-server", 111)
+ restored = []
+ backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
+
+ kw.restore_kv_resume(backend, manifest)
+ assert restored == []
+ assert not state_file.exists()
+
+
+def test_note_model_unloaded_purges_manifest_and_files(tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M"))
+ kw._set_kv_resume(manifest)
+ kw.note_model_unloaded()
+ assert kw.get_last_unloaded_model() is None
+ assert kw.take_kv_resume() is None
+ assert not state_file.exists()
+
+
+def test_note_model_loaded_purges_manifest_and_files(tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M"))
+ kw._set_kv_resume(manifest)
+ kw.note_model_loaded()
+ assert kw.get_last_unloaded_model() is None
+ assert kw.take_kv_resume() is None
+ assert not state_file.exists()
+
+
+def test_new_idle_save_purges_previous_manifest_files(tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ old_file, old_manifest = _seed_kv_manifest(tmp_path)
+ kw._set_kv_resume(old_manifest)
+ new_file = tmp_path / "resume-def-slot0.bin"
+ new_file.write_bytes(b"kv2")
+ kw._set_kv_resume(
+ {
+ "identity": ("unsloth/B-GGUF", None, "unsloth/B-GGUF"),
+ "dir": str(tmp_path),
+ "binary": ("/bin/llama-server", 111),
+ "slots": [{"id": 0, "filename": new_file.name, "n_saved": 7}],
+ }
+ )
+ assert not old_file.exists() # replaced manifest's files purged
+ assert new_file.exists()
+ assert kw.take_kv_resume()["slots"][0]["filename"] == new_file.name
+
+
+def test_sweep_slot_save_dir_removes_only_resume_files(monkeypatch, tmp_path):
+ from core.inference import llama_keepwarm as kw
+ from utils.paths import storage_roots
+
+ monkeypatch.setattr(storage_roots, "llama_slot_cache_root", lambda: tmp_path)
+ stale = tmp_path / "resume-old-slot0.bin"
+ stale.write_bytes(b"kv")
+ other = tmp_path / "unrelated.txt"
+ other.write_text("keep")
+ kw.sweep_slot_save_dir()
+ assert not stale.exists()
+ assert other.exists()
+
+
+def test_keep_kv_setting_roundtrip_and_default(monkeypatch):
+ import storage.studio_db as db
+
+ store = {}
+ monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
+ monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
+
+ assert settings.get_auto_unload_keep_kv() is True # default when never stored
+ assert settings.set_openai_auto_switch(True, 60, False)[2] is False
+ assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False
+ assert settings.get_auto_unload_keep_kv() is False
+ # None leaves the stored value untouched (older clients can't reset it).
+ assert settings.set_openai_auto_switch(True, 60, None)[2] is False
+ assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False
+ with pytest.raises(ValueError, match = "true or false"):
+ settings.set_openai_auto_switch(True, 60, "garbage")
+
+
+def test_stale_stash_cleanup_waits_for_lifecycle_gate(monkeypatch, tmp_path):
+ # The loop's stale-stash purge must wait on the gate a mid-reload holds.
+ import time
+ from core.inference import llama_keepwarm as kw
+
+ monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 3600)
+ kw._inflight = 0
+ kw._pending = 0
+ kw._last_active = time.monotonic()
+ backend = _FakeBackend("unsloth/New-GGUF")
+ monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ kw._kv_resume = manifest
+ kw._last_unloaded_model = ("unsloth/A-GGUF", "Q4_K_M")
+
+ assert kw._lifecycle_lock.acquire(blocking = False) # simulate in-flight reload
+ try:
+ _drive_idle_loop(kw)
+ assert kw._kv_resume is manifest # purge deferred while the gate is held
+ assert state_file.exists()
+ finally:
+ kw._lifecycle_lock.release()
+ _drive_idle_loop(kw)
+ assert kw._kv_resume is None # gate freed: genuinely stale stash purged
+ assert not state_file.exists()
+
+
+def test_put_route_disabling_keep_kv_purges_saved_state(monkeypatch, tmp_path):
+ import routes.settings as settings_route
+ import storage.studio_db as db
+ from core.inference import llama_keepwarm as kw
+
+ store = {}
+ monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
+ monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
+ state_file, manifest = _seed_kv_manifest(tmp_path)
+ monkeypatch.setattr(kw, "_kv_resume", manifest)
+
+ payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_keep_kv = False)
+ resp = settings_route.update_openai_auto_switch(payload, "tester")
+ assert resp.auto_unload_keep_kv is False
+ assert kw._kv_resume is None
+ assert not state_file.exists()
+
+
+def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch):
+ # A keep-KV-only update must not materialize the env TTL as a stored value.
+ import routes.settings as settings_route
+ import storage.studio_db as db
+
+ store = {}
+ monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
+ monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
+ monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600")
+
+ assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None
+ enabled, idle, keep_kv = settings.set_openai_auto_switch(False, None, False)
+ assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched
+ assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active
+ assert (enabled, idle, keep_kv) == (False, 600, False)
+
+
+def test_load_impl_notes_loaded_with_backend_off_loop():
+ import inspect
+ src = inspect.getsource(inference_route._load_model_impl)
+ assert "to_thread(note_model_loaded, llama_backend)" in src
+
+
+def test_restore_matches_gguf_realpath_across_naming(tmp_path):
+ from core.inference import llama_keepwarm as kw
+
+ blob = tmp_path / "blob.gguf"
+ blob.write_bytes(b"gguf")
+ link = tmp_path / "snapshot.gguf"
+ try:
+ link.symlink_to(blob)
+ except OSError:
+ pytest.skip("symlinks unsupported on this host")
+
+ backend = _FakeBackend("/hf/snapshots/d7f5", hf_variant = None)
+ backend._gguf_path = str(link) # reload resolved the symlink spelling
+ backend._slot_save_binary = ("/bin/llama-server", 111)
+ restored = []
+ backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
+ state_file, manifest = _seed_kv_manifest(
+ tmp_path, identity = ("unsloth/A-GGUF", None, "unsloth/A-GGUF"), gguf = str(blob)
+ )
+
+ kw.restore_kv_resume(backend, manifest)
+ assert len(restored) == 1 # names differ, file identical: restore ran
+ assert not state_file.exists()
+
+
def test_setter_rejects_idle_below_floor(monkeypatch):
import storage.studio_db as db
diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py
index 462435e5d5..7007440f4c 100644
--- a/studio/backend/utils/openai_auto_switch_settings.py
+++ b/studio/backend/utils/openai_auto_switch_settings.py
@@ -30,11 +30,13 @@ from typing import Any, Optional
OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model"
AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds"
+AUTO_UNLOAD_KEEP_KV_SETTING_KEY = "openai_api_auto_unload_keep_kv"
MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides"
MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
+DEFAULT_AUTO_UNLOAD_KEEP_KV = True
MIN_AUTO_UNLOAD_IDLE_SECONDS = 60
_CACHE_TTL_S = 2.0
@@ -158,29 +160,54 @@ def get_auto_unload_idle_seconds() -> int:
return env if env is not None else 0
-def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]:
- """Set both auto-switch flags in one transaction so a settings PUT can't leave
- one key updated and the other stale. Both values are coerced before any write,
- so an invalid value raises without persisting either."""
+def get_auto_unload_keep_kv() -> bool:
+ """Whether the idle unload persists slot KV to disk for restore on reload."""
+ parsed = _coerce_bool(_cached_setting(AUTO_UNLOAD_KEEP_KV_SETTING_KEY, None))
+ return parsed if parsed is not None else DEFAULT_AUTO_UNLOAD_KEEP_KV
+
+
+def set_openai_auto_switch(
+ enabled: Any,
+ idle_seconds: Any,
+ keep_kv: Any = None,
+) -> tuple[bool, int, bool]:
+ """One-transaction write; ``None`` leaves a stored value untouched."""
parsed_enabled = _coerce_bool(enabled)
if parsed_enabled is None:
raise ValueError("OpenAI auto-switch must be true or false.")
- parsed_idle = _coerce_int(idle_seconds)
- if parsed_idle is None:
- raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
- if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS:
- raise ValueError(
- f"Auto-unload idle seconds must be 0 (off) or at least "
- f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}."
- )
+ parsed_idle = None
+ if idle_seconds is not None:
+ parsed_idle = _coerce_int(idle_seconds)
+ if parsed_idle is None:
+ raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
+ if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS:
+ raise ValueError(
+ f"Auto-unload idle seconds must be 0 (off) or at least "
+ f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}."
+ )
+ parsed_keep_kv = None
+ if keep_kv is not None:
+ parsed_keep_kv = _coerce_bool(keep_kv)
+ if parsed_keep_kv is None:
+ raise ValueError("Keep KV on idle unload must be true or false.")
from storage.studio_db import upsert_app_settings
- upsert_app_settings(
- {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle}
- )
+ updates: dict[str, Any] = {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled}
+ if parsed_idle is not None:
+ updates[AUTO_UNLOAD_IDLE_SETTING_KEY] = parsed_idle
+ if parsed_keep_kv is not None:
+ updates[AUTO_UNLOAD_KEEP_KV_SETTING_KEY] = parsed_keep_kv
+ upsert_app_settings(updates)
_invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY)
- _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
- return parsed_enabled, parsed_idle
+ if parsed_idle is not None:
+ _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
+ if parsed_keep_kv is not None:
+ _invalidate(AUTO_UNLOAD_KEEP_KV_SETTING_KEY)
+ return (
+ parsed_enabled,
+ parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(),
+ parsed_keep_kv if parsed_keep_kv is not None else get_auto_unload_keep_kv(),
+ )
def get_model_overrides() -> dict[str, dict]:
diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py
index 1faa2b1281..35b8c57e9b 100644
--- a/studio/backend/utils/paths/storage_roots.py
+++ b/studio/backend/utils/paths/storage_roots.py
@@ -61,6 +61,11 @@ def cache_root() -> Path:
return studio_root() / "cache"
+def llama_slot_cache_root() -> Path:
+ """Dir llama-server saves/restores slot KV state in across idle unloads."""
+ return cache_root() / "llama-slots"
+
+
def studio_bin_root() -> Path:
"""Dir for Unsloth-managed executables (the `unsloth` shim, downloaded tools like cloudflared)."""
return studio_root() / "bin"
diff --git a/studio/frontend/src/features/settings/api/openai-auto-switch.ts b/studio/frontend/src/features/settings/api/openai-auto-switch.ts
index 80ffc084d0..47bad56eab 100644
--- a/studio/frontend/src/features/settings/api/openai-auto-switch.ts
+++ b/studio/frontend/src/features/settings/api/openai-auto-switch.ts
@@ -11,6 +11,8 @@ export type OpenAIAutoSwitchSettings = {
// True when the idle-unload loop will actually unload (e.g. enabled via the
// UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off).
idleUnloadActive: boolean;
+ // Persist the KV cache to disk on idle unload and restore it on reload.
+ autoUnloadKeepKv: boolean;
};
type ApiOpenAIAutoSwitchSettings = {
@@ -21,6 +23,8 @@ type ApiOpenAIAutoSwitchSettings = {
default_enabled: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
idle_unload_active?: boolean;
+ // biome-ignore lint/style/useNamingConvention: API schema
+ auto_unload_keep_kv?: boolean;
};
let cachedSettings: OpenAIAutoSwitchSettings | null = null;
@@ -34,6 +38,7 @@ function fromApi(
autoUnloadIdleSeconds: settings.auto_unload_idle_seconds,
defaultEnabled: settings.default_enabled,
idleUnloadActive: settings.idle_unload_active ?? false,
+ autoUnloadKeepKv: settings.auto_unload_keep_kv ?? true,
};
}
@@ -66,15 +71,23 @@ export async function loadOpenAIAutoSwitchSettings() {
export async function updateOpenAIAutoSwitchSettings(
enabled: boolean,
- autoUnloadIdleSeconds: number,
+ autoUnloadIdleSeconds?: number,
+ autoUnloadKeepKv?: boolean,
): Promise {
const res = await authFetch("/api/settings/openai-auto-switch", {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
enabled,
- // biome-ignore lint/style/useNamingConvention: API schema
- auto_unload_idle_seconds: autoUnloadIdleSeconds,
+ // Omitted fields keep their stored value.
+ ...(autoUnloadIdleSeconds === undefined
+ ? {}
+ : // biome-ignore lint/style/useNamingConvention: API schema
+ { auto_unload_idle_seconds: autoUnloadIdleSeconds }),
+ ...(autoUnloadKeepKv === undefined
+ ? {}
+ : // biome-ignore lint/style/useNamingConvention: API schema
+ { auto_unload_keep_kv: autoUnloadKeepKv }),
}),
});
if (!res.ok) {
diff --git a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
index 32b3e53a2c..aa6857cff5 100644
--- a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
+++ b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
@@ -62,13 +62,18 @@ export function ModelAutoSwitchSection() {
const persist = async (
enabled: boolean,
- idleSeconds: number,
+ idleSeconds: number | undefined,
syncDraft = true,
+ keepKv?: boolean,
) => {
setIsSaving(true);
setError(null);
try {
- const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds);
+ const saved = await updateOpenAIAutoSwitchSettings(
+ enabled,
+ idleSeconds,
+ keepKv,
+ );
setSettings(saved);
if (syncDraft) {
setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds));
@@ -107,6 +112,11 @@ export function ModelAutoSwitchSection() {
void persist(true, idleSeconds);
};
+ const handleKeepKvToggle = (keepKv: boolean) => {
+ if (!settings) return;
+ void persist(settings.enabled, undefined, false, keepKv);
+ };
+
return (
+ {settings?.idleUnloadActive ? (
+
+
+
+ ) : null}
);
}
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index de8ac17c29..cbddc9f0c2 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -232,6 +232,9 @@ export const en = {
loadError: "Failed to load model auto-switch settings.",
saveError: "Failed to save model auto-switch settings.",
idleError: "Enter 0 to keep the model loaded, or at least 60 seconds.",
+ keepKv: "Keep chat context across idle unload",
+ keepKvDescription:
+ "Save the model's KV cache to disk before an idle unload and restore it on reload, so resumed chats skip re-reading their history. Chat context is written to disk (up to 10 GB) until it is restored or cleaned up.",
},
previewSharing: {
sectionTitle: "Preview sharing",
From 9e334d552c77de8cc4b52ff1d2b13a93891d1366 Mon Sep 17 00:00:00 2001
From: alkinun
Date: Mon, 20 Jul 2026 10:23:37 +0300
Subject: [PATCH 026/255] Fix text-only VLM CPT packing truncation (#7211)
* Fix text-only VLM CPT packing truncation
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Handle streaming vision datasets in packing
* Harden multimodal packing detection
* Preserve safe packing boundaries
* Scope stream packing checks to VLMs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Narrow VLM packing detection
* Align packing mode and eval safety
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add qwen3_5/qwen3_next to PADDING_FREE_BLOCKLIST to avoid packed-sequence contamination
* Detect hybrid linear-attention models structurally instead of by name for packing guard
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Install wrapped-packing setup at the signature, not the Zoo license comment
The _unsloth_wrapped_packing / _inspect setup block was injected by matching the
exact 'All Unsloth Zoo code licensed under LGPLv3' comment line in the sourced
sft_prepare_dataset. The unsloth_zoo dependency is only lower-bounded, so a newer
Zoo that moves or drops that header made the setup a silent no-op while the
truncation and pack_dataset rewrites still emitted references to those names,
raising NameError on every SFT dataset preparation.
Anchor the setup on the function signature instead (a structural location that
always exists) and fail loudly if it cannot be found, so the helper variables are
always defined before they are referenced across Zoo versions.
Adds a regression test that patches in a Zoo source without the license header.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Etherl <61019402+Etherll@users.noreply.github.com>
Co-authored-by: danielhanchen
---
studio/backend/core/training/trainer.py | 12 +-
tests/utils/test_packing.py | 371 +++++++++++++++++++++++-
unsloth/models/rl_replacements.py | 88 ++++--
unsloth/trainer.py | 177 ++++++++++-
4 files changed, 610 insertions(+), 38 deletions(-)
diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 26720865f4..8e419849cb 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -3425,15 +3425,19 @@ class UnslothTrainer:
logger.info(
f"CPT: using UnslothTrainer with embedding_learning_rate={embedding_lr}\n"
)
+ cpt_args = _UnslothTrainingArguments(
+ embedding_learning_rate = embedding_lr,
+ **config_args,
+ )
+ if config_args.get("packing", False):
+ cpt_args.packing_strategy = "wrapped"
+ logger.info("CPT packing strategy: wrapped\n")
trainer_kwargs = {
"model": self.model,
"tokenizer": sft_tokenizer,
"train_dataset": dataset["dataset"],
"data_collator": data_collator,
- "args": _UnslothTrainingArguments(
- embedding_learning_rate = embedding_lr,
- **config_args,
- ),
+ "args": cpt_args,
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset
diff --git a/tests/utils/test_packing.py b/tests/utils/test_packing.py
index a8557d8533..98c29d9f0f 100644
--- a/tests/utils/test_packing.py
+++ b/tests/utils/test_packing.py
@@ -14,6 +14,7 @@
# along with this program. If not, see .
from unsloth import FastLanguageModel
+import unsloth.trainer as trainer_module
from unsloth.utils import attention_dispatch as attention_dispatch_utils
from unsloth.utils.packing import (
configure_padding_free,
@@ -29,7 +30,7 @@ from unittest.mock import patch
import pytest
import torch
-from datasets import Dataset
+from datasets import Dataset, IterableDataset
from trl import SFTConfig, SFTTrainer
from trl.trainer.sft_trainer import DataCollatorForLanguageModeling
@@ -160,6 +161,374 @@ def test_configure_padding_free():
assert config.remove_unused_columns is False
+def _patch_fake_sft_trainer():
+ class FakeSFTTrainer:
+ def __init__(self, *args, **kwargs):
+ self.model = args[0] if len(args) >= 1 else kwargs["model"]
+ self.args = args[1] if len(args) >= 2 else kwargs["args"]
+ self.data_collator = args[2] if len(args) >= 3 else kwargs.get("data_collator")
+
+ trainer_module._patch_sft_trainer_auto_packing(SimpleNamespace(SFTTrainer = FakeSFTTrainer))
+ return FakeSFTTrainer
+
+
+def _vlm_model():
+ return SimpleNamespace(
+ config = SimpleNamespace(
+ architectures = ["Gemma4ForConditionalGeneration"],
+ model_type = "gemma4",
+ vision_config = SimpleNamespace(),
+ ),
+ max_seq_length = 16,
+ )
+
+
+def _text_model():
+ return SimpleNamespace(
+ config = SimpleNamespace(
+ architectures = ["LlamaForCausalLM"],
+ model_type = "llama",
+ ),
+ max_seq_length = 16,
+ )
+
+
+class _CharacterTokenizer:
+ bos_token = None
+ eos_token = None
+ chat_template = None
+
+ def __call__(self, texts, **kwargs):
+ is_batched = isinstance(texts, list)
+ if not is_batched:
+ texts = [texts]
+ input_ids = [[ord(char) for char in text] for text in texts]
+ if kwargs.get("truncation") and kwargs.get("max_length") is not None:
+ input_ids = [ids[: kwargs["max_length"]] for ids in input_ids]
+ return {"input_ids": input_ids if is_batched else input_ids[0]}
+
+
+def test_vlm_text_dataset_allows_explicit_packing():
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+
+ trainer = fake_trainer(
+ model = _vlm_model(),
+ args = config,
+ processing_class = object(),
+ train_dataset = Dataset.from_dict({"text": ["text-only CPT sample"]}),
+ )
+
+ assert config.packing is True
+ assert config.padding_free is True
+ assert trainer.model._unsloth_allow_packed_overlength is True
+
+
+def test_vlm_without_processing_class_still_disables_packing():
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+
+ fake_trainer(
+ _vlm_model(),
+ config,
+ None,
+ Dataset.from_dict({"text": ["text-only sample"]}),
+ )
+
+ assert config.packing is False
+ assert config.padding_free is False
+
+
+@pytest.mark.parametrize(
+ ("model_type", "architecture"),
+ (
+ ("t5", "T5ForConditionalGeneration"),
+ ("bart", "BartForConditionalGeneration"),
+ ("whisper", "WhisperForConditionalGeneration"),
+ ("csm", "CsmForConditionalGeneration"),
+ ),
+)
+def test_nonvision_conditional_generation_keeps_packing(model_type, architecture):
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ model = SimpleNamespace(
+ config = SimpleNamespace(model_type = model_type, architectures = [architecture]),
+ max_seq_length = 16,
+ )
+
+ trainer = fake_trainer(
+ model,
+ config,
+ None,
+ Dataset.from_dict({"text": ["text-only sample"]}),
+ )
+
+ assert config.packing is True
+ assert config.padding_free is True
+ assert trainer.model._unsloth_allow_packed_overlength is True
+
+
+def test_vlm_vision_dataset_still_disables_packing():
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+
+ fake_trainer(
+ _vlm_model(),
+ config,
+ None,
+ Dataset.from_dict({"images": [None], "text": ["multimodal sample"]}),
+ None,
+ object(),
+ )
+
+ assert config.packing is False
+ assert config.padding_free is False
+
+
+@pytest.mark.parametrize(
+ "vision_column",
+ ("pixel_values", "pixel_attention_mask", "image_grid_thw"),
+)
+def test_vlm_preprocessed_vision_dataset_disables_packing(vision_column):
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+
+ fake_trainer(
+ model = _vlm_model(),
+ args = config,
+ processing_class = object(),
+ train_dataset = Dataset.from_dict({"input_ids": [[1]], vision_column: [None]}),
+ )
+
+ assert config.packing is False
+ assert config.padding_free is False
+
+
+@pytest.mark.parametrize("dict_eval", (False, True))
+def test_vlm_vision_eval_dataset_disables_packing(dict_eval):
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ eval_dataset = Dataset.from_dict({"input_ids": [[1]], "pixel_values": [None]})
+ if dict_eval:
+ eval_dataset = {"vision": eval_dataset}
+
+ fake_trainer(
+ model = _vlm_model(),
+ args = config,
+ processing_class = object(),
+ train_dataset = Dataset.from_dict({"text": ["text-only training sample"]}),
+ eval_dataset = eval_dataset,
+ )
+
+ assert config.packing is False
+ assert config.padding_free is False
+
+
+def test_vlm_streaming_vision_dataset_without_metadata_disables_packing():
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ dataset = IterableDataset.from_generator(
+ lambda: iter([{"images": [None], "text": "multimodal sample"}])
+ )
+ assert dataset.column_names is None
+
+ fake_trainer(
+ model = _vlm_model(),
+ args = config,
+ processing_class = object(),
+ train_dataset = dataset,
+ )
+
+ assert config.packing is False
+ assert config.padding_free is False
+ assert next(iter(dataset))["text"] == "multimodal sample"
+
+
+@pytest.mark.parametrize("data_collator", (None, object()))
+def test_stateful_stream_is_not_consumed_during_detection(data_collator):
+ class StatefulDataset:
+ def __init__(self):
+ self.rows = iter([{"text": "first"}, {"text": "second"}])
+
+ def __iter__(self):
+ return (row for row in self.rows)
+
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ dataset = StatefulDataset()
+
+ fake_trainer(
+ model = _vlm_model(),
+ args = config,
+ processing_class = object(),
+ data_collator = data_collator,
+ train_dataset = dataset,
+ )
+
+ assert config.packing is False
+ assert config.padding_free is False
+ assert next(iter(dataset))["text"] == "first"
+
+
+def test_text_model_stream_without_metadata_keeps_packing():
+ class StatefulDataset:
+ def __init__(self):
+ self.rows = iter([{"text": "first"}, {"text": "second"}])
+
+ def __iter__(self):
+ return (row for row in self.rows)
+
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ dataset = StatefulDataset()
+
+ trainer = fake_trainer(
+ model = _text_model(),
+ args = config,
+ processing_class = object(),
+ train_dataset = dataset,
+ )
+
+ assert config.packing is True
+ assert config.padding_free is True
+ assert trainer.model._unsloth_allow_packed_overlength is True
+ assert next(iter(dataset))["text"] == "first"
+
+
+def test_bfd_packing_truncates_before_packing(monkeypatch):
+ args = SimpleNamespace(
+ dataset_num_proc = 1,
+ dataset_text_field = "text",
+ max_length = 4,
+ packing_strategy = "bfd",
+ )
+ trainer = SimpleNamespace(model = None)
+ dataset = Dataset.from_dict({"prompt": ["abc"], "completion": ["defghij"]})
+ prepare_globals = SFTTrainer._prepare_dataset.__globals__
+
+ def passthrough_pack_dataset(dataset, seq_length, strategy, map_kwargs):
+ return dataset
+
+ monkeypatch.setitem(prepare_globals, "pack_dataset", passthrough_pack_dataset)
+ packed = SFTTrainer._prepare_dataset(
+ trainer,
+ dataset,
+ _CharacterTokenizer(),
+ args,
+ True,
+ None,
+ "train",
+ )
+
+ assert len(packed["input_ids"][0]) == args.max_length
+
+
+def test_wrapped_strategy_without_packing_still_truncates():
+ args = SimpleNamespace(
+ dataset_num_proc = 1,
+ dataset_text_field = "text",
+ max_length = 4,
+ packing_strategy = "wrapped",
+ )
+ trainer = SimpleNamespace(model = None)
+ dataset = Dataset.from_dict({"text": ["abcdefghi"]})
+
+ prepared = SFTTrainer._prepare_dataset(
+ trainer,
+ dataset,
+ _CharacterTokenizer(),
+ args,
+ False,
+ None,
+ "train",
+ )
+
+ assert len(prepared["input_ids"][0]) == args.max_length
+
+
+@pytest.mark.parametrize("legacy_api", (False, True))
+def test_wrapped_packing_preserves_overlength_tokens(monkeypatch, legacy_api):
+ args_kwargs = {
+ "dataset_num_proc": 1,
+ "dataset_text_field": "text",
+ "max_length": 4,
+ }
+ if not legacy_api:
+ args_kwargs["packing_strategy"] = "wrapped"
+ args = SimpleNamespace(**args_kwargs)
+ trainer = SimpleNamespace(model = None)
+ dataset = Dataset.from_dict({"text": ["abcdefghi"]})
+ prepare_globals = SFTTrainer._prepare_dataset.__globals__
+ pack_dataset = prepare_globals["pack_dataset"]
+
+ def legacy_pack_dataset(
+ dataset,
+ seq_length,
+ map_kwargs = None,
+ ):
+ return pack_dataset(dataset, seq_length, "wrapped", map_kwargs)
+
+ if legacy_api:
+ monkeypatch.setitem(prepare_globals, "pack_dataset", legacy_pack_dataset)
+
+ packed = SFTTrainer._prepare_dataset(
+ trainer,
+ dataset,
+ _CharacterTokenizer(),
+ args,
+ True,
+ None,
+ "train",
+ )
+
+ packed_ids = packed["input_ids"]
+ assert sum(len(input_ids) for input_ids in packed_ids) == 9
+ assert all(len(input_ids) <= args.max_length for input_ids in packed_ids)
+
+
+# Named to match the unsloth_zoo helper: sft_trainer_prepare_dataset sources it by
+# name and renames "def sft_prepare_dataset" -> "def _prepare_dataset". This fixture
+# deliberately omits the "All Unsloth Zoo code licensed under LGPLv3" header to emulate
+# a newer, compatible Zoo whose header moved (the dependency is only lower-bounded).
+def sft_prepare_dataset(
+ self, dataset, processing_class, args, packing, formatting_func, dataset_text_field
+):
+ do_truncation = True
+ # Mirror the Zoo call so the "truncation = do_truncation," injection anchor
+ # survives formatting (a bare tuple assignment gets rewritten to a paren form).
+ dataset = processing_class(
+ dataset,
+ truncation = do_truncation,
+ )
+ return dataset
+
+
+def test_wrapped_packing_setup_survives_missing_zoo_header(monkeypatch):
+ # Regression: the wrapped-packing setup used to anchor on the Zoo license comment,
+ # so a header change made it a no-op while the truncation reference still landed,
+ # NameError-ing every SFT dataset preparation. It must now install via the
+ # signature and always precede the reference.
+ import ast
+ import textwrap
+ import unsloth.models.rl_replacements as rlr
+
+ monkeypatch.setitem(rlr.RL_REPLACEMENTS, "sft_prepare_dataset", sft_prepare_dataset)
+
+ source = (
+ "def _prepare_dataset(self, dataset, processing_class, args, packing, "
+ "formatting_func, dataset_text_field):\n return dataset\n"
+ )
+ patched = rlr.sft_trainer_prepare_dataset("_prepare_dataset", source)
+
+ assert "_unsloth_wrapped_packing = packing" in patched
+ assert "import inspect as _inspect" in patched
+ assert "not _unsloth_wrapped_packing" in patched
+ assert patched.index("_unsloth_wrapped_packing = packing") < patched.index(
+ "truncation = do_truncation and not _unsloth_wrapped_packing"
+ )
+ ast.parse(textwrap.dedent(patched))
+
+
class _DummyChild(torch.nn.Module):
def __init__(self):
super().__init__()
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index ffb845b04f..b0709f7376 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -276,17 +276,13 @@ def dpo_trainer_vision_signature_columns(function_name, function):
_extra_columns = "".join(f' "{_k}",\n' for _k in _DPO_VISION_KEYS)
new_function = function.replace(
' "image_sizes",\n "token_type_ids",\n',
- f' "image_sizes",\n'
- f"{_extra_columns}"
- f' "token_type_ids",\n',
+ f' "image_sizes",\n{_extra_columns} "token_type_ids",\n',
)
if new_function != function:
return new_function
return function.replace(
' "image_sizes",\n "ref_chosen_logps",\n',
- f' "image_sizes",\n'
- f"{_extra_columns}"
- f' "ref_chosen_logps",\n',
+ f' "image_sizes",\n{_extra_columns} "ref_chosen_logps",\n',
)
@@ -458,6 +454,60 @@ def sft_trainer_prepare_dataset(function_name, function):
if matched:
# Use fast version!
function = inspect.getsource(fast_sft_prepare_dataset)
+ # why: install the wrapped-packing setup (and the `_inspect` import the
+ # truncation / pack_dataset rewrites below depend on) at the function
+ # signature, a structural anchor that always exists, rather than the
+ # unsloth_zoo license-comment line. That header is only lower-bounded, so a
+ # newer Zoo may move or drop it; anchoring there let the setup silently
+ # no-op while the references still landed, NameError-ing every SFT dataset
+ # preparation. Fail loudly if even the signature cannot be located.
+ _wrapped_packing_setup = (
+ " import inspect as _inspect\n"
+ " try:\n"
+ ' _unsloth_pack_has_strategy = "strategy" in _inspect.signature(pack_dataset).parameters\n'
+ " except Exception:\n"
+ " _unsloth_pack_has_strategy = True\n"
+ " _unsloth_wrapped_packing = packing and (\n"
+ ' getattr(args, "packing_strategy", None) == "wrapped"\n'
+ " or not _unsloth_pack_has_strategy\n"
+ " )\n"
+ )
+ function, _n_setup = re.subn(
+ r"(def sft_prepare_dataset\s*\(.*?\)\s*(?:->[^:\n]*)?:[ \t]*\n)",
+ lambda match: match.group(1) + _wrapped_packing_setup,
+ function,
+ count = 1,
+ flags = re.DOTALL,
+ )
+ if _n_setup != 1:
+ raise RuntimeError(
+ "Unsloth: failed to install wrapped-packing support into "
+ "sft_prepare_dataset (signature not found); please file a bug report."
+ )
+ function = function.replace(
+ "truncation = do_truncation,",
+ "truncation = do_truncation and not _unsloth_wrapped_packing,",
+ )
+ function = function.replace(
+ "if do_truncation and max_seq_length > 0:",
+ "if do_truncation and not _unsloth_wrapped_packing and max_seq_length > 0:",
+ )
+ function = function.replace(
+ """dataset = pack_dataset(
+ dataset.select_columns(used_column_names),
+ max_seq_length,
+ getattr(args, "packing_strategy", "bfd"),
+ map_kwargs,
+ )""",
+ """_pack_kwargs = {"map_kwargs": map_kwargs}
+ if "strategy" in _inspect.signature(pack_dataset).parameters:
+ _pack_kwargs["strategy"] = getattr(args, "packing_strategy", "bfd")
+ dataset = pack_dataset(
+ dataset.select_columns(used_column_names),
+ max_seq_length,
+ **_pack_kwargs,
+ )""",
+ )
function = function.split("\n")
function = "\n".join(" " * 4 + x for x in function)
function = function.replace("def sft_prepare_dataset", "def _prepare_dataset")
@@ -2120,19 +2170,21 @@ def grpo_trainer_compute_loss(function_name, function):
logits_to_keep,
batch_size = None,
compute_entropy = False,
- compute_efficient = False: self._get_per_token_logps(
- model, input_ids, attention_mask, logits_to_keep, compute_efficient
+ compute_efficient = False: (
+ self._get_per_token_logps(
+ model, input_ids, attention_mask, logits_to_keep, compute_efficient
+ )
+ if hasattr(self, "_get_per_token_logps")
+ else self._get_per_token_logps_and_entropies(
+ model,
+ input_ids,
+ attention_mask,
+ logits_to_keep,
+ batch_size,
+ compute_entropy,
+ compute_efficient,
+ )[0]
)
- if hasattr(self, "_get_per_token_logps")
- else self._get_per_token_logps_and_entropies(
- model,
- input_ids,
- attention_mask,
- logits_to_keep,
- batch_size,
- compute_entropy,
- compute_efficient,
- )[0]
) # logps
per_token_logps = get_logps_func(
diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index 83cb1758f0..61d41aad21 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -100,6 +100,10 @@ PADDING_FREE_BLOCKLIST = {
"gemma2", # - gemma2: Uses slow_attention_softcapping which has torch.compile issues
"gpt_oss", # - gpt_oss: Uses Flex Attention which doesn't handle padding_free correctly
}
+# Hybrid linear-attention / state-space models (Qwen3.5, Qwen3-Next, ...) carry a
+# recurrent gated-delta state plus a causal conv1d. Sample packing / padding-free
+# flattens the batch, so those ops leak state across sequence boundaries. Detected
+# structurally by _is_hybrid_linear_attention_model rather than by model name.
def _should_pack(config) -> bool:
@@ -137,6 +141,132 @@ def _should_skip_auto_packing_error(exc: Exception) -> bool:
return any(msg in message for msg in _AUTO_PACK_SKIP_MESSAGES)
+_VISION_DATASET_KEYS = frozenset(
+ {
+ "image",
+ "images",
+ "image_grid_thw",
+ "image_position_ids",
+ "image_sizes",
+ "mm_token_type_ids",
+ "pixel_attention_mask",
+ "pixel_position_ids",
+ "pixel_values",
+ "pixel_values_videos",
+ "video",
+ "videos",
+ "video_grid_thw",
+ }
+)
+
+
+def _is_vlm_config(config, model_types = ()) -> bool:
+ if any(
+ hasattr(config, attr)
+ for attr in ("vision_config", "img_processor", "image_token_index", "projector_config")
+ ):
+ return True
+
+ architectures = getattr(config, "architectures", None) or ()
+ try:
+ from transformers.models.auto import modeling_auto
+
+ mappings = (
+ getattr(modeling_auto, "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES", {}) or {},
+ getattr(modeling_auto, "MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES", {}) or {},
+ )
+ registry_types = set().union(*(mapping.keys() for mapping in mappings))
+ registry_classes = set().union(*(mapping.values() for mapping in mappings))
+ config_types = set(model_types or ())
+ model_type = getattr(config, "model_type", None)
+ if model_type is not None:
+ config_types.add(model_type)
+ if not config_types.isdisjoint(registry_types) or any(
+ architecture in registry_classes for architecture in architectures
+ ):
+ return True
+ except Exception:
+ pass
+ return any(
+ isinstance(architecture, str) and architecture.endswith("ForVisionText2Text")
+ for architecture in architectures
+ )
+
+
+def _is_vision_dataset(dataset, *, unknown_is_vision = False) -> bool:
+ if dataset is None:
+ return False
+ column_names = getattr(dataset, "column_names", None)
+ if column_names is not None:
+ return not _VISION_DATASET_KEYS.isdisjoint(column_names)
+ # Unknown-schema streams cannot be safely probed without potentially dropping a sample.
+ return unknown_is_vision
+
+
+def _is_vision_eval_dataset(dataset, *, unknown_is_vision = False) -> bool:
+ if isinstance(dataset, dict):
+ return any(
+ _is_vision_dataset(split, unknown_is_vision = unknown_is_vision)
+ for split in dataset.values()
+ )
+ return _is_vision_dataset(dataset, unknown_is_vision = unknown_is_vision)
+
+
+_HYBRID_CONFIG_MARKERS = (
+ "linear_conv_kernel_dim",
+ "linear_key_head_dim",
+ "linear_value_head_dim",
+ "full_attention_interval",
+)
+
+
+def _is_hybrid_linear_attention_model(model) -> bool:
+ """Detect models mixing linear-attention / state-space mixers (gated-delta,
+ Mamba-style) with a causal conv1d, e.g. Qwen3.5 / Qwen3-Next. Packing and
+ padding-free flatten the batch, and those recurrent + conv ops leak state
+ across sequence boundaries, so they must not be packed. Uses composite
+ structural evidence rather than a model-name match."""
+ if model is None:
+ return False
+
+ # Config-level: explicit hybrid layer schedule or linear-attn markers.
+ for config in (
+ getattr(model, "config", None),
+ getattr(getattr(model, "config", None), "text_config", None),
+ ):
+ if config is None:
+ continue
+ layer_types = getattr(config, "layer_types", None)
+ if isinstance(layer_types, (list, tuple)) and any(
+ isinstance(t, str) and "linear_attention" in t for t in layer_types
+ ):
+ return True
+ if any(hasattr(config, marker) for marker in _HYBRID_CONFIG_MARKERS):
+ return True
+
+ # Module-level: a mixer carrying a recurrent gated-delta op plus a conv1d.
+ named_modules = getattr(model, "named_modules", None)
+ if named_modules is None:
+ return False
+ seen = set()
+ for _, module in named_modules():
+ if id(module) in seen:
+ continue
+ seen.add(id(module))
+ cls = type(module).__name__
+ if not (
+ cls.endswith("GatedDeltaNet") or "LinearAttention" in cls or cls.endswith("Mamba2Mixer")
+ ):
+ continue
+ has_recurrent = any(
+ hasattr(module, attr)
+ for attr in ("chunk_gated_delta_rule", "recurrent_gated_delta_rule", "A_log")
+ )
+ if has_recurrent and hasattr(module, "conv1d"):
+ return True
+ return False
+
+
# Unsloth gradient accumulation fix:
from transformers import __version__ as transformers_version, ProcessorMixin
@@ -498,30 +628,43 @@ def _patch_sft_trainer_auto_packing(trl_module):
else:
config_arg = kwargs.get("args")
- model = kwargs.get("model")
- is_unsupported_model = False
+ model = args[0] if len(args) >= 1 else kwargs.get("model")
is_vlm = False
+ is_unsupported_model = False
+ is_hybrid = False
if model is not None:
model_config = getattr(model, "config", None)
if model_config is not None:
model_types = get_transformers_model_type(model_config)
is_unsupported_model = any(x in PADDING_FREE_BLOCKLIST for x in model_types)
+ is_vlm = _is_vlm_config(model_config, model_types)
+ is_hybrid = _is_hybrid_linear_attention_model(model)
- architectures = getattr(model_config, "architectures", None)
- if architectures is None:
- architectures = []
- is_vlm = any(x.endswith("ForConditionalGeneration") for x in architectures)
- is_vlm = is_vlm or hasattr(model_config, "vision_config")
-
- processing_class = kwargs.get("processing_class") or kwargs.get("tokenizer")
- data_collator = kwargs.get("data_collator")
+ processing_class = (
+ args[5] if len(args) >= 6 else kwargs.get("processing_class") or kwargs.get("tokenizer")
+ )
+ data_collator = args[2] if len(args) >= 3 else kwargs.get("data_collator")
+ train_dataset = args[3] if len(args) >= 4 else kwargs.get("train_dataset")
+ eval_dataset = args[4] if len(args) >= 5 else kwargs.get("eval_dataset")
+ is_processor = isinstance(processing_class, ProcessorMixin)
+ is_auto_processor_vlm = is_vlm and processing_class is None
+ is_vision_dataset = (
+ data_collator is None
+ and not is_processor
+ and (
+ _is_vision_dataset(train_dataset, unknown_is_vision = is_vlm)
+ or _is_vision_eval_dataset(eval_dataset, unknown_is_vision = is_vlm)
+ )
+ )
# Disable padding-free for VLMs / custom collators / blocklisted models
blocked = (
(data_collator is not None)
- or isinstance(processing_class, ProcessorMixin)
- or is_vlm
+ or is_processor
+ or is_auto_processor_vlm
+ or is_vision_dataset
or is_unsupported_model
+ or is_hybrid
or (
os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
) # Disable padding free on forced logits
@@ -535,10 +678,14 @@ def _patch_sft_trainer_auto_packing(trl_module):
if blocked and requested_pack:
reason = "custom data collator"
- if data_collator is None and isinstance(processing_class, ProcessorMixin):
+ if data_collator is None and is_processor:
reason = "processor-based model"
- elif is_vlm:
- reason = "vision-language model"
+ elif is_auto_processor_vlm:
+ reason = "vision-language model with auto processor"
+ elif is_vision_dataset:
+ reason = "vision dataset"
+ elif is_hybrid:
+ reason = "hybrid linear-attention model"
elif is_unsupported_model:
reason = f"unsupported model type(s): {', '.join(model_types)}"
message = f"Unsloth: Sample packing skipped ({reason} detected)."
From cf912cbd881190f411147b8c93294408efe5f90c Mon Sep 17 00:00:00 2001
From: Naitik Pal
Date: Mon, 20 Jul 2026 13:03:56 +0530
Subject: [PATCH 027/255] feat(studio): add UNSLOTH_LLAMA_CPP_BACKEND env var
to force CPU fallback #7213 (#7228)
* test(studio): add e2e test for cpu-fallback overriding vulkan
* feat(studio): add UNSLOTH_LLAMA_CPP_BACKEND env var
* feat(studio): add UNSLOTH_LLAMA_CPP_BACKEND env var
* Preserve UNSLOTH_LLAMA_CPP_BACKEND=cpu across llama.cpp updates for PR #7228
The in-app updater rebuilt the installer command without --cpu-fallback and
only re-asserted Vulkan, so accepting a llama.cpp update after forcing CPU on
an Intel iGPU host re-ran host detection and routed back to the crashing Vulkan
bundle (#7213). Record install_kind in the prebuilt marker and re-assert
--cpu-fallback on update when the installed bundle is CPU.
Also make setup.sh's UNSLOTH_LLAMA_CPP_BACKEND check case-insensitive to match
setup.ps1, and add tests for the updater CPU preservation and the setup.sh flag
plumbing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim and validate UNSLOTH_LLAMA_CPP_BACKEND, warn on unknown values for PR #7228
Trim surrounding whitespace and lowercase the value in both setup.sh and
setup.ps1, so values like ' cpu ' or 'CPU' still force the CPU-only prebuilt.
An unrecognized value (e.g. 'gpu') now prints a warning instead of silently
falling back to auto. Extend test_setup_llama_cpp_backend.py to cover both
scripts, including trimmed, empty and unknown values.
* Preserve arm64 CPU installs on update and honor CPU override in Windows prune for PR #7228
The update-path CPU preservation only matched install_kind ending in -cpu, so
arm64 CPU bundles (linux-arm64, windows-arm64) were re-routed to a GPU or source
build on update. Match the full set of CPU-only kinds instead.
Persisting install_kind also activated the previously inert Windows
mismatch-prune in setup.ps1: on a GPU host with UNSLOTH_LLAMA_CPP_BACKEND=cpu it
saw the windows-cpu marker as mismatched and deleted it every rerun. Normalize
the override once and make CPU expected so a deliberate CPU install is kept.
Extend the tests to cover both.
* Document legacy llama.cpp markers keep heal-to-GPU on update for PR #7228
Legacy prebuilt markers written before install_kind was persisted intentionally
do not force --cpu-fallback on update: the in-app updater lets them re-resolve
(heal to a GPU bundle) per the existing behavior from #6097, and only markers
that explicitly record a CPU install_kind are pinned to CPU. Add a comment and a
regression case documenting the boundary.
* Tighten llama.cpp CPU-fallback comments for PR #7228
* Fix Windows install-prune to keep valid Intel/fallback bundles for PR #7228
Persisting install_kind activated the setup.ps1 mismatch-prune, whose
expectedKinds was incomplete: the non-NVIDIA/non-AMD branch omitted
windows-vulkan (the Intel auto-route) and the GPU branches omitted the
windows-cpu/windows-arm64 fallback the installer uses when a GPU prebuilt is
missing. That made every setup rerun delete and re-download a valid Intel Vulkan
(or CPU-fallback) install. List all kinds the installer can produce per host so
only a bundle the host cannot run is pruned. Cover the full matrix in tests.
* Persist force_cpu marker flag so only forced CPU installs re-assert on update for PR #7228
* Add --force-cpu for deliberate CPU installs and warn on macOS for PR #7228
* Record force_cpu when reusing a matching CPU bundle for PR #7228
* Accept force_cpu keyword in installer test validator fakes for PR #7228
---------
Co-authored-by: danielhanchen
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
.../tests/test_install_resolve_prebuilt.py | 95 +++++++++++
studio/backend/tests/test_llama_cpp_update.py | 58 ++++++-
.../tests/test_setup_llama_cpp_backend.py | 154 ++++++++++++++++++
studio/backend/utils/llama_cpp_update.py | 12 +-
studio/install_llama_prebuilt.py | 69 +++++++-
studio/setup.ps1 | 13 ++
studio/setup.sh | 20 +++
.../test_install_llama_prebuilt_logic.py | 4 +
8 files changed, 410 insertions(+), 15 deletions(-)
create mode 100644 studio/backend/tests/test_setup_llama_cpp_backend.py
diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py
index e97ca47717..3ebad861ad 100644
--- a/studio/backend/tests/test_install_resolve_prebuilt.py
+++ b/studio/backend/tests/test_install_resolve_prebuilt.py
@@ -445,6 +445,101 @@ def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
assert routed is host
+@pytest.mark.parametrize("cpu_flag", ["--cpu-fallback", "--force-cpu"])
+def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan(monkeypatch, capsys, cpu_flag):
+ """Either CPU flag via CLI must suppress Vulkan even on an Intel GPU host: both
+ drop GPU detection (--force-cpu additionally persists, on the install path)."""
+ monkeypatch.setattr(
+ ilp,
+ "detect_host",
+ lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True),
+ )
+ seen = {}
+
+ def _resolver(tag, host, repo, published_release_tag):
+ seen["host"] = host
+ seen["repo"] = repo
+ raise ilp.PrebuiltFallback("no asset")
+
+ monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver)
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "install_llama_prebuilt.py",
+ "--resolve-prebuilt",
+ "latest",
+ cpu_flag,
+ "--output-format",
+ "json",
+ ],
+ )
+ assert ilp.main() == ilp.EXIT_SUCCESS
+ # The CPU flag must suppress Intel GPU, route to fork (not upstream Vulkan)
+ assert seen["host"].has_intel_gpu is False
+ assert seen["repo"] == FORK
+
+
+@pytest.mark.parametrize(
+ "flags, expect_force, expect_persist",
+ [
+ ([], False, False),
+ # Automatic/transient last resort (arm64 GPU-build recovery): drops GPU but
+ # does NOT persist, so a later update heals to a GPU bundle (#6097).
+ (["--cpu-fallback"], True, False),
+ # Deliberate CPU-only (UNSLOTH_LLAMA_CPP_BACKEND=cpu): drops GPU AND persists so
+ # the updater re-asserts it and never revives the Intel iGPU crash (#7213).
+ (["--force-cpu"], True, True),
+ (["--cpu-fallback", "--force-cpu"], True, True),
+ ],
+)
+def test_cli_cpu_flags_thread_force_and_persist(
+ monkeypatch, tmp_path, flags, expect_force, expect_persist
+):
+ captured = {}
+ monkeypatch.setattr(ilp, "install_prebuilt", lambda **kw: captured.update(kw))
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ ["install_llama_prebuilt.py", "--install-dir", str(tmp_path / "llama.cpp"), *flags],
+ )
+ assert ilp.main() == ilp.EXIT_SUCCESS
+ assert captured["force_cpu"] is expect_force
+ assert captured["persist_force_cpu"] is expect_persist
+
+
+@pytest.mark.parametrize(
+ "existing, requested, expected",
+ [
+ # A deliberate --force-cpu on top of a naturally-installed CPU bundle (same
+ # asset, install skipped) must still flip the marker to true (#7213).
+ (False, True, True),
+ (None, True, True),
+ # No spurious writes when already in sync, and a released force syncs down.
+ (True, True, True),
+ (False, False, False),
+ (True, False, False),
+ ],
+)
+def test_sync_marker_force_cpu(tmp_path, existing, requested, expected):
+ marker = {"tag": "b9585", "asset": "llama-b9585-bin-ubuntu-x64.tar.gz"}
+ if existing is not None:
+ marker["force_cpu"] = existing
+ marker_path = tmp_path / "UNSLOTH_PREBUILT_INFO.json"
+ marker_path.write_text(json.dumps(marker))
+ ilp.sync_marker_force_cpu(tmp_path, requested)
+ written = json.loads(marker_path.read_text())
+ assert written["force_cpu"] is expected
+ # Unrelated fields are preserved.
+ assert written["asset"] == "llama-b9585-bin-ubuntu-x64.tar.gz"
+
+
+def test_sync_marker_force_cpu_missing_marker_is_noop(tmp_path):
+ # No marker (or unreadable) must not crash the reuse path.
+ ilp.sync_marker_force_cpu(tmp_path, True)
+ assert not (tmp_path / "UNSLOTH_PREBUILT_INFO.json").exists()
+
+
def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
# A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
# physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or
diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py
index 83ea07a066..f12384231f 100644
--- a/studio/backend/tests/test_llama_cpp_update.py
+++ b/studio/backend/tests/test_llama_cpp_update.py
@@ -83,6 +83,7 @@ def _write_install(
repo: str = "unslothai/llama.cpp",
asset: str | None = None,
release_tag: str | None = None,
+ force_cpu: bool | None = None,
) -> str:
"""Create a fake prebuilt install and return the llama-server path."""
bin_dir = dir_ / "build" / "bin"
@@ -99,6 +100,8 @@ def _write_install(
}
if asset is not None:
marker["asset"] = asset
+ if force_cpu is not None:
+ marker["force_cpu"] = force_cpu
(dir_ / MARKER).write_text(json.dumps(marker))
return str(binary)
@@ -493,6 +496,47 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
+@pytest.mark.parametrize(
+ "force_cpu, expect_flag",
+ [
+ # A deliberate CPU install (marker force_cpu=True) re-asserts --force-cpu on
+ # update so detect_host on a GPU host cannot re-route and revive the crash
+ # (#7213); --force-cpu also re-persists the flag for the next update.
+ (True, True),
+ # A transient fallback (or a legacy marker without the flag) stays free to
+ # heal to a GPU bundle (#6097).
+ (False, False),
+ (None, False),
+ ],
+)
+def test_start_update_cpu_fallback_preserved_by_flag(monkeypatch, tmp_path, force_cpu, expect_flag):
+ asset = "llama-b9493-bin-ubuntu-x64.tar.gz"
+ install_dir = tmp_path / "llama.cpp"
+ binary = _write_install(install_dir, "b9493", asset = asset, force_cpu = force_cpu)
+ monkeypatch.setattr(upd, "_find_binary", lambda: binary)
+ monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
+ monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
+
+ captured: dict = {}
+
+ def _on_start(cmd):
+ captured["cmd"] = cmd
+ _write_install(install_dir, "b9518", asset = asset, force_cpu = force_cpu)
+
+ _patch_installer_popen(monkeypatch, lines = ["installed\n"], on_start = _on_start)
+
+ assert upd.start_update()["started"] is True
+ deadline = time.time() + 10
+ while time.time() < deadline:
+ job = upd.get_update_status()["job"]
+ if job["state"] in ("success", "error"):
+ break
+ time.sleep(0.05)
+ assert job["state"] == "success", job
+ assert ("--force-cpu" in captured["cmd"]) is expect_flag
+ assert "--cpu-fallback" not in captured["cmd"]
+
+
def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")
@@ -676,7 +720,7 @@ def test_install_cmd_rocm_marker_forwards_gfx(monkeypatch, tmp_path):
assert "--rocm-gfx" in cmd
assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x"
assert "--has-rocm" not in cmd
- assert "--cpu-fallback" not in cmd
+ assert "--force-cpu" not in cmd
assert "--simple-policy" not in cmd
assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd
@@ -690,17 +734,17 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path):
def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path):
- # Legacy CPU installs recorded a ggml-org marker (new installs use the fork).
- # Re-running into the same install-dir/repo reproduces the same CPU bundle;
- # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's
- # arm64 rescue and must not appear here.
+ # Legacy CPU installs recorded a ggml-org marker (new installs use the fork) with
+ # no force_cpu field. Re-running into the same install-dir/repo reproduces the same
+ # CPU bundle; --force-cpu (the persisted-CPU re-assert) must not appear for a marker
+ # that never recorded a deliberate CPU choice, so it can still heal to GPU (#6097).
cmd = _capture_install_cmd(
monkeypatch,
tmp_path,
repo = "ggml-org/llama.cpp",
asset = "llama-b9334-bin-ubuntu-x64.tar.gz",
)
- assert "--cpu-fallback" not in cmd
+ assert "--force-cpu" not in cmd
assert "--rocm-gfx" not in cmd
assert "--has-rocm" not in cmd
assert "--simple-policy" not in cmd
@@ -714,7 +758,7 @@ def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tm
assert "--simple-policy" not in cmd
assert "--rocm-gfx" not in cmd
assert "--has-rocm" not in cmd
- assert "--cpu-fallback" not in cmd
+ assert "--force-cpu" not in cmd
def test_install_cmd_pins_offered_release_tag(monkeypatch, tmp_path):
diff --git a/studio/backend/tests/test_setup_llama_cpp_backend.py b/studio/backend/tests/test_setup_llama_cpp_backend.py
new file mode 100644
index 0000000000..36928c680c
--- /dev/null
+++ b/studio/backend/tests/test_setup_llama_cpp_backend.py
@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""setup.sh and setup.ps1 must map UNSLOTH_LLAMA_CPP_BACKEND=cpu to
+install_llama_prebuilt.py's --force-cpu so users can force the CPU-only prebuilt
+on GPU hosts (#7213). The match is case-insensitive and whitespace-trimmed, an
+unrecognized value warns instead of silently falling back, and macOS warns (no
+CPU-only bundle). Runs the real block extracted from each script so the tests
+track the shipped logic.
+"""
+
+import os
+import re
+import shutil
+import subprocess
+from pathlib import Path
+
+import pytest
+
+_STUDIO = Path(__file__).resolve().parents[2]
+_SETUP_SH = _STUDIO / "setup.sh"
+_SETUP_PS1 = _STUDIO / "setup.ps1"
+_SKIP_NO_BASH = pytest.mark.skipif(shutil.which("bash") is None, reason = "bash unavailable")
+_SKIP_NO_PWSH = pytest.mark.skipif(shutil.which("pwsh") is None, reason = "pwsh unavailable")
+
+
+def _backend_block() -> str:
+ text = _SETUP_SH.read_text(encoding = "utf-8")
+ m = re.search(r"_llama_backend=.*?esac", text, re.DOTALL)
+ assert m, "UNSLOTH_LLAMA_CPP_BACKEND block not found in setup.sh"
+ return m.group(0)
+
+
+def _run(value: str | None, system: str = "Linux") -> tuple[list[str], str]:
+ # Pass the value through env (not the script text) so whitespace survives, and
+ # stub the setup.sh logging helpers the unknown-value branch calls. system sets
+ # _HOST_SYSTEM so the macOS (Darwin) no-op branch can be exercised.
+ env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"}
+ if value is not None:
+ env["UNSLOTH_LLAMA_CPP_BACKEND"] = value
+ harness = (
+ f'_PREBUILT_CMD=()\nC_WARN=""\n_HOST_SYSTEM="{system}"\n'
+ 'step() { printf "STEP: %s\\n" "$*" >&2; }\n'
+ f"{_backend_block()}\n"
+ 'printf "%s\\n" "${_PREBUILT_CMD[@]}"'
+ )
+ out = subprocess.run(
+ ["bash", "-c", harness], capture_output = True, text = True, env = env, check = True
+ )
+ return out.stdout.split(), out.stderr
+
+
+@_SKIP_NO_BASH
+@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"])
+def test_backend_cpu_appends_flag(value):
+ # A deliberate CPU choice persists, so it uses --force-cpu (not the transient
+ # --cpu-fallback the arm64 GPU-build recovery uses).
+ args, stderr = _run(value)
+ assert "--force-cpu" in args
+ assert "--cpu-fallback" not in args
+ assert "Ignoring" not in stderr
+
+
+@_SKIP_NO_BASH
+@pytest.mark.parametrize("value", ["cpu", "CPU", " cpu "])
+def test_backend_cpu_macos_warns_no_flag(value):
+ # macOS has no CPU-only bundle (the universal build already runs on CPU), so the
+ # override warns instead of writing a misleading forced-CPU marker.
+ args, stderr = _run(value, system = "Darwin")
+ assert "--force-cpu" not in args
+ assert "--cpu-fallback" not in args
+ assert "macOS" in stderr
+
+
+@_SKIP_NO_BASH
+@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "])
+def test_backend_auto_no_flag_no_warn(value):
+ args, stderr = _run(value)
+ assert "--force-cpu" not in args
+ assert "Ignoring" not in stderr
+
+
+@_SKIP_NO_BASH
+@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"])
+def test_backend_unknown_warns_and_no_flag(value):
+ args, stderr = _run(value)
+ assert "--force-cpu" not in args
+ assert "Ignoring" in stderr
+
+
+@_SKIP_NO_BASH
+def test_arm64_recovery_uses_transient_cpu_fallback():
+ # The arm64 Linux GPU-build recovery must stay transient (--cpu-fallback), never
+ # the persisted --force-cpu, so a later update can still heal to a GPU bundle (#6097).
+ text = _SETUP_SH.read_text(encoding = "utf-8")
+ m = re.search(r"_ARM64_CPU_CMD=\((.*?)\)", text, re.DOTALL)
+ assert m, "arm64 CPU recovery command not found in setup.sh"
+ block = m.group(1)
+ assert "--cpu-fallback" in block
+ assert "--force-cpu" not in block
+
+
+def _ps1_search(pattern: str, flags = 0) -> str:
+ m = re.search(pattern, _SETUP_PS1.read_text(encoding = "utf-8"), flags)
+ assert m, f"setup.ps1 block not found: {pattern}"
+ return m.group(0)
+
+
+def _run_ps1(value: str | None) -> str:
+ # The override is normalized (assign + warn) at the top of the prebuilt block and
+ # applied to $prebuiltArgs lower down; compose both real snippets.
+ normalize = _ps1_search(
+ r'\$llamaBackend = "\$\(\$env:UNSLOTH_LLAMA_CPP_BACKEND\)".*?Write-Host.*?\n\s*\}',
+ re.DOTALL,
+ )
+ apply_flag = _ps1_search(
+ r'if \(\$llamaBackend -eq "cpu"\) \{\s*\$prebuiltArgs \+= "--force-cpu"\s*\}'
+ )
+ env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"}
+ if value is not None:
+ env["UNSLOTH_LLAMA_CPP_BACKEND"] = value
+ harness = f'$prebuiltArgs = @()\n{normalize}\n{apply_flag}\n"ARGS:" + ($prebuiltArgs -join ",")'
+ out = subprocess.run(
+ ["pwsh", "-NoProfile", "-Command", harness],
+ capture_output = True,
+ text = True,
+ env = env,
+ check = True,
+ )
+ return out.stdout
+
+
+@_SKIP_NO_PWSH
+@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"])
+def test_ps1_backend_cpu_appends_flag(value):
+ out = _run_ps1(value)
+ assert "--force-cpu" in out
+ assert "Ignoring" not in out
+
+
+@_SKIP_NO_PWSH
+@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "])
+def test_ps1_backend_auto_no_flag_no_warn(value):
+ out = _run_ps1(value)
+ assert "--force-cpu" not in out
+ assert "Ignoring" not in out
+
+
+@_SKIP_NO_PWSH
+@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"])
+def test_ps1_backend_unknown_warns_and_no_flag(value):
+ out = _run_ps1(value)
+ assert "--force-cpu" not in out
+ assert "Ignoring" in out
diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py
index 31dbda63ea..67733bde35 100644
--- a/studio/backend/utils/llama_cpp_update.py
+++ b/studio/backend/utils/llama_cpp_update.py
@@ -479,6 +479,7 @@ def _run_update(
asset: Optional[str],
script: Path,
pin_release_tag: Optional[str] = None,
+ force_cpu: bool = False,
) -> None:
"""Worker: put the backend into a maintenance state, run the installer for
the latest prebuilt, then refresh caches so the next load uses the new build.
@@ -522,6 +523,12 @@ def _run_update(
if pin_release_tag:
cmd.extend(["--published-release-tag", pin_release_tag])
cmd.extend(_rocm_install_args(asset))
+ # Re-assert a deliberate CPU install (--force-cpu) so detect_host on a GPU host
+ # does not re-route to a GPU/Vulkan bundle and revive the crash (#7213). --force-cpu
+ # (not --cpu-fallback) also re-persists force_cpu, keeping the choice across future
+ # updates. A natural fallback (or a legacy marker without the flag) heals to GPU (#6097).
+ if force_cpu:
+ cmd.append("--force-cpu")
logger.info("llama update: installing", cmd = " ".join(cmd))
# Stream progress lines into job["progress"].
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
@@ -671,6 +678,7 @@ def start_update() -> dict:
repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO
from_tag = marker.get("tag") or marker.get("release_tag")
asset = marker.get("asset")
+ force_cpu = bool(marker.get("force_cpu"))
# Install exactly the release the banner offered: the installer's own
# "latest" is commit-date ordered and can lag the published_at pick
# above, reinstalling the current build in a loop (the #6219 class).
@@ -705,6 +713,8 @@ def start_update() -> dict:
repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO
from_tag = None
asset = (res or {}).get("asset")
+ # Source builds carry no forced-CPU marker, so nothing to preserve here.
+ force_cpu = False
# No pin: source-build detection resolves via --resolve-prebuilt latest,
# the same resolver the unpinned apply uses, so the two already agree.
pin_release_tag = None
@@ -735,7 +745,7 @@ def start_update() -> dict:
thread = threading.Thread(
target = _run_update,
- args = (install_dir, repo, asset, script, pin_release_tag),
+ args = (install_dir, repo, asset, script, pin_release_tag, force_cpu),
name = "llama-cpp-update",
daemon = True,
)
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index 9bbd0cb8be..b8182a534b 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -6450,6 +6450,7 @@ def write_prebuilt_metadata(
choice: AssetChoice,
approved_checksums: ApprovedReleaseChecksums,
prebuilt_fallback_used: bool,
+ force_cpu: bool = False,
) -> None:
source_asset_name, source_sha256 = selected_source_archive_metadata(
approved_checksums,
@@ -6474,6 +6475,10 @@ def write_prebuilt_metadata(
"release_tag": release_tag,
"published_repo": approved_checksums.repo,
"asset": choice.name,
+ # True only for a deliberate CPU choice (--force-cpu). The updater re-asserts it
+ # so a forced CPU install is not re-routed to a GPU bundle (#7213). An automatic
+ # --cpu-fallback (e.g. arm64 GPU-build recovery) stays False so it can heal to GPU.
+ "force_cpu": force_cpu,
"asset_sha256": choice.expected_sha256,
"source": choice.source_label,
# Binary-side repo/tag for non-fork sources (e.g. the ggml-org upstream
@@ -6501,6 +6506,24 @@ def write_prebuilt_metadata(
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(metadata, indent = 2) + "\n")
+def sync_marker_force_cpu(install_dir: Path, persist_force_cpu: bool) -> None:
+ """Sync only the force_cpu flag of an existing marker when the resolved bundle is
+ unchanged, so the install is skipped without a full metadata rewrite. A deliberate
+ --force-cpu on top of a naturally installed CPU bundle (same asset) must still be
+ recorded, else the updater will not re-assert it and can re-route the install to a
+ GPU/Vulkan bundle that revives the crash (#7213)."""
+ marker_path = install_dir / "UNSLOTH_PREBUILT_INFO.json"
+ try:
+ marker = json.loads(marker_path.read_text())
+ except (OSError, ValueError):
+ return
+ if not isinstance(marker, dict) or bool(marker.get("force_cpu")) == persist_force_cpu:
+ return
+ marker["force_cpu"] = persist_force_cpu
+ marker_path.write_text(json.dumps(marker, indent = 2) + "\n")
+ log(f"existing install reused; recorded force_cpu={persist_force_cpu} from this run")
+
+
def expected_install_fingerprint(
*,
llama_tag: str,
@@ -6746,6 +6769,7 @@ def validate_prebuilt_choice(
approved_checksums: ApprovedReleaseChecksums,
prebuilt_fallback_used: bool,
quantized_path: Path,
+ force_cpu: bool = False,
) -> tuple[Path, Path]:
source_repo, source_ref, source_archive, exact_source = preferred_source_archive(
approved_checksums, llama_tag
@@ -6786,6 +6810,7 @@ def validate_prebuilt_choice(
choice = choice,
approved_checksums = approved_checksums,
prebuilt_fallback_used = prebuilt_fallback_used,
+ force_cpu = force_cpu,
)
# Hashless external prebuilts are not in the approved-sha256
# manifest and rely on the functional smoke test as their only integrity gate,
@@ -6828,6 +6853,7 @@ def validate_prebuilt_attempts(
approved_checksums: ApprovedReleaseChecksums,
initial_fallback_used: bool = False,
existing_install_dir: Path | None = None,
+ force_cpu: bool = False,
) -> tuple[AssetChoice, Path, bool]:
attempt_list = list(attempts)
if not attempt_list:
@@ -6880,6 +6906,7 @@ def validate_prebuilt_attempts(
approved_checksums = approved_checksums,
prebuilt_fallback_used = tried_fallback,
quantized_path = quantized_path,
+ force_cpu = force_cpu,
)
except Exception as exc:
remove_tree(staging_dir)
@@ -6939,8 +6966,8 @@ def _route_to_vulkan_prebuilt(
"""Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes
- from UPSTREAM_REPO. Two triggers route here, both suppressed under
- --cpu-fallback (the explicit "give me CPU" last resort wins):
+ from UPSTREAM_REPO. Two triggers route here, both suppressed when a CPU flag
+ (--cpu-fallback or --force-cpu, folded into force_cpu) wins:
* UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend;
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose
of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
@@ -7020,8 +7047,11 @@ def install_prebuilt(
override_has_rocm: bool = False,
override_rocm_gfx: str | None = None,
force_cpu: bool = False,
+ persist_force_cpu: bool = False,
instruction_cleanup_root: Path | None = None,
) -> None:
+ # force_cpu drops GPU detection (mechanism, both --cpu-fallback and --force-cpu);
+ # persist_force_cpu records the deliberate choice so the updater re-asserts it.
host = detect_host()
host = _apply_host_overrides(
host,
@@ -7072,6 +7102,9 @@ def install_prebuilt(
"existing llama.cpp install already matches selected release "
f"{current.release_tag} upstream_tag={current.llama_tag}; skipping download and install"
)
+ # Reused bundle is unchanged, but a fresh --force-cpu still must be
+ # recorded so the updater re-asserts it (#7213).
+ sync_marker_force_cpu(install_dir, persist_force_cpu)
return
with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp:
work_dir = Path(tmp)
@@ -7092,6 +7125,7 @@ def install_prebuilt(
"existing llama.cpp install already matches fallback release "
f"{plan.release_tag} upstream_tag={plan.llama_tag}; skipping reinstall"
)
+ sync_marker_force_cpu(install_dir, persist_force_cpu)
return
log(
"selected "
@@ -7112,6 +7146,8 @@ def install_prebuilt(
initial_fallback_used = release_index > 0,
# Skip is gated per-attempt inside, so pass the dir always.
existing_install_dir = install_dir,
+ # Persist only the deliberate choice, not a transient fallback.
+ force_cpu = persist_force_cpu,
)
except ExistingInstallSatisfied:
return
@@ -7209,8 +7245,21 @@ def parse_args() -> argparse.Namespace:
default = False,
help = (
"Select the CPU prebuilt for this OS/arch even when a GPU is present. "
- "setup.sh uses this as a last resort for arm64 Linux GPU hosts whose "
- "source build failed (no arm64 CUDA prebuilt exists anywhere)."
+ "Automatic/transient: setup.sh uses this as a last resort for arm64 Linux "
+ "GPU hosts whose source build failed. Does NOT persist, so a later update "
+ "heals back to a GPU bundle once one is available (#6097). Use --force-cpu "
+ "for a deliberate CPU-only choice that survives updates."
+ ),
+ )
+ parser.add_argument(
+ "--force-cpu",
+ action = "store_true",
+ default = False,
+ help = (
+ "Deliberate CPU-only install (UNSLOTH_LLAMA_CPP_BACKEND=cpu). Drops GPU "
+ "detection like --cpu-fallback but also records force_cpu in the marker, so "
+ "the in-app updater re-asserts CPU and never re-routes to a GPU/Vulkan "
+ "bundle that would revive the Intel iGPU crash (#7213)."
),
)
resolve_group = parser.add_mutually_exclusive_group()
@@ -7333,16 +7382,19 @@ def main() -> int:
# Host-aware "is a prebuilt available" probe, no download. Every host now
# plans against the fork (args.published_repo defaults to it); an explicit
# --published-repo overrides. PrebuiltFallback == source build.
+ # Both flags drop GPU detection; --force-cpu additionally persists (install
+ # path only). The probe only needs the mechanism, so OR them.
+ _cpu_mechanism = args.cpu_fallback or args.force_cpu
host = _apply_host_overrides(
detect_host(),
override_has_rocm = args.has_rocm,
override_rocm_gfx = args.rocm_gfx,
- force_cpu = args.cpu_fallback,
+ force_cpu = _cpu_mechanism,
)
# Same Vulkan routing the install path applies, so the probe's answer
# matches what would install (an Intel/forced-Vulkan host -> upstream).
host, repo, release_tag = _route_to_vulkan_prebuilt(
- host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback
+ host, args.published_repo, args.published_release_tag or "", force_cpu = _cpu_mechanism
)
try:
_requested, plans = resolve_simple_install_release_plans(
@@ -7380,7 +7432,10 @@ def main() -> int:
published_release_tag = args.published_release_tag or "",
override_has_rocm = args.has_rocm,
override_rocm_gfx = args.rocm_gfx,
- force_cpu = args.cpu_fallback,
+ # Both drop GPU detection; only --force-cpu (deliberate) is recorded so the
+ # updater re-asserts it. --cpu-fallback stays transient and heals to GPU.
+ force_cpu = args.cpu_fallback or args.force_cpu,
+ persist_force_cpu = args.force_cpu,
instruction_cleanup_root = install_arg.absolute(),
)
return EXIT_SUCCESS
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 98e801cd3c..f7d33a1142 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -32,6 +32,10 @@ $PackageDir = Split-Path -Parent $ScriptDir
# (no matching GitHub release), forces a source build, and causes HTTP 422
# errors. Only use "master" temporarily when the latest release is missing
# support for a new model architecture.
+#
+# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default) or "cpu". When "cpu", forces
+# the CPU-only prebuilt bundle on GPU hosts. Fixes Intel iGPU Vulkan
+# crashes (#7213).
$DefaultLlamaPrForce = ""
$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"
$DefaultLlamaTag = "latest"
@@ -3367,6 +3371,15 @@ if ($LocalLlamaCppLinked) {
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
$prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG)
}
+ # UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, whitespace-trimmed) forces the
+ # CPU-only prebuilt via --force-cpu (persisted so updates keep it). Fixes Intel
+ # iGPU Vulkan crash (#7213).
+ $llamaBackend = "$($env:UNSLOTH_LLAMA_CPP_BACKEND)".Trim().ToLowerInvariant()
+ if ($llamaBackend -eq "cpu") {
+ $prebuiltArgs += "--force-cpu"
+ } elseif ($llamaBackend -and $llamaBackend -ne "auto") {
+ Write-Host "[WARN] Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$($env:UNSLOTH_LLAMA_CPP_BACKEND)' (expected 'auto' or 'cpu')" -ForegroundColor Yellow
+ }
$prevEAPPrebuilt = $ErrorActionPreference
$ErrorActionPreference = "Continue"
$previousNativeErrorPreference = $null
diff --git a/studio/setup.sh b/studio/setup.sh
index 8d47eecfda..df7178c662 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -36,6 +36,10 @@ fi
# forces a source build, and causes HTTP 422 errors.
# Only use "master" temporarily when the latest release
# is missing support for a new model architecture.
+#
+# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default) or "cpu". When "cpu", forces
+# the CPU-only prebuilt bundle on GPU hosts.
+# Fixes Intel iGPU Vulkan crashes (#7213).
# ──────────────────────────────────────────────────────────────────────────
_DEFAULT_LLAMA_PR_FORCE=""
_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"
@@ -1359,6 +1363,22 @@ else
# present so it can still attempt a prebuilt. Mirrors setup.ps1 behaviour.
_PREBUILT_CMD+=(--has-rocm)
fi
+ # UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, trimmed) forces the CPU-only
+ # prebuilt via --force-cpu, bypassing Vulkan/CUDA/ROCm. Fixes Intel iGPU crash (#7213).
+ # No effect on macOS: the universal bundle already runs on CPU (Metal is a runtime
+ # -ngl choice), so warn instead of writing a misleading forced-CPU marker.
+ _llama_backend="$(printf '%s' "${UNSLOTH_LLAMA_CPP_BACKEND:-auto}" | awk '{$1=$1; print tolower($0)}')"
+ case "$_llama_backend" in
+ cpu)
+ if [ "$_HOST_SYSTEM" = "Darwin" ]; then
+ step "llama.cpp" "UNSLOTH_LLAMA_CPP_BACKEND=cpu has no effect on macOS (universal build; use -ngl 0 at runtime for CPU-only)" "$C_WARN" >&2
+ else
+ _PREBUILT_CMD+=(--force-cpu)
+ fi
+ ;;
+ ""|auto) ;;
+ *) step "llama.cpp" "Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$UNSLOTH_LLAMA_CPP_BACKEND' (expected 'auto' or 'cpu')" "$C_WARN" >&2 ;;
+ esac
_PREBUILT_LOG="$(mktemp)"
set +e
if _is_verbose; then
diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py
index e995e5033e..9a094ddc0e 100644
--- a/tests/studio/install/test_install_llama_prebuilt_logic.py
+++ b/tests/studio/install/test_install_llama_prebuilt_logic.py
@@ -1348,6 +1348,7 @@ def test_install_prebuilt_falls_back_to_older_release_plan(
approved_checksums,
initial_fallback_used = False,
existing_install_dir = None,
+ force_cpu = False,
):
call_log.append((llama_tag, initial_fallback_used))
if llama_tag == "b9002":
@@ -2551,6 +2552,7 @@ def test_install_prebuilt_skips_when_older_release_fallback_matches_existing_ins
approved_checksums,
initial_fallback_used = False,
existing_install_dir = None,
+ force_cpu = False,
):
call_log.append(llama_tag)
raise PrebuiltFallback("validation failed for latest release")
@@ -2698,6 +2700,7 @@ def test_install_prebuilt_skips_same_release_fallback_attempt_when_installed(
approved_checksums,
prebuilt_fallback_used,
quantized_path,
+ force_cpu = False,
):
attempted_names.append(choice.name)
if choice.name == first_choice.name:
@@ -2824,6 +2827,7 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p
approved_checksums,
initial_fallback_used = False,
existing_install_dir = None,
+ force_cpu = False,
):
attempted.append((llama_tag, release_tag, attempts[0].source_label))
if llama_tag == "b9002":
From 07272b9278eaa2813c30f3b12f712276ff97fa01 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 00:57:02 -0700
Subject: [PATCH 028/255] Experimental: correct varlen sample packing for
hybrid linear-attention models (#7249)
* Fix text-only VLM CPT packing truncation
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Handle streaming vision datasets in packing
* Harden multimodal packing detection
* Preserve safe packing boundaries
* Scope stream packing checks to VLMs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Narrow VLM packing detection
* Align packing mode and eval safety
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add qwen3_5/qwen3_next to PADDING_FREE_BLOCKLIST to avoid packed-sequence contamination
* Detect hybrid linear-attention models structurally instead of by name for packing guard
* Add experimental varlen packing for hybrid linear-attention models
Feed seq_idx to the causal conv and cu_seqlens to the gated-delta scan so
sample packing / padding-free reset state at sequence boundaries for hybrid
linear-attention models (Qwen3.5, Qwen3-Next). Gated behind
UNSLOTH_EXPERIMENTAL_HYBRID_PACKING and fail-closed: when the flag is off or
the accelerated kernels (causal_conv1d + fla) are unavailable, the guard keeps
these models on the padded path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden hybrid linear-attention varlen packing shim
Make patch_hybrid_linear_attention_varlen robust across transformers 4.57.6
through 5.x and TRL 0.22.2 through 1.x, following the import_fixes.py style:
- Read UNSLOTH_EXPERIMENTAL_HYBRID_PACKING at call time so the flag takes effect
when set after importing unsloth.
- Idempotent: repeat calls on a patched model return True without re-validating
the wrappers or double-wrapping; signatures are checked on captured originals.
- Prefer the authoritative packed_seq_lengths (via get_packed_info_from_kwargs)
over position_ids resets, handling pad_to_multiple_of trailing tokens.
- Suppress injection for cached forwards (use_cache / past_key_values) so
generation and eval are left on the untouched decode path.
- Validate every gated-delta module before mutating any (transactional).
- Bind position_ids / use_cache from both positional and keyword args.
- Verify dispatch at runtime (Unsloth wraps each module forward, so the mixer
source is not statically inspectable) and warn once if the shim is never hit.
- Emit one deduped diagnostic on each fail-closed path.
Add CPU unit tests covering the hybrid guard detection, the boundary builders,
and the shim (fail-closed, active, idempotent, cached no-op, runtime handshake).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Abort hybrid packing when the varlen shim is not fully dispatched
The runtime handshake used a single per-module hit flag written by both the conv
and scan wrappers, so a partial dispatch (only one kernel routed through
self.) passed the any() check and trained on contaminated data, and a
missing dispatch only logged a warning. Track conv and scan dispatch separately,
require both on every gated-delta module on the first packed forward, and raise
before loss/backward when either is missing (the batch is already flattened, so
there is no padded recovery at that point). Also skip an empty packed_seq_lengths
before it reaches max(), and document the position_ids fallback's left-pad
assumption.
Add tests for no-dispatch and partial (conv-only / scan-only) abort, the
packed_seq_lengths preference over a competing position_ids, MRoPE 3D position
ids, and the pad_to_multiple_of trailing-segment path through the metadata builder.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Import the hybrid packing patch from its submodule to satisfy the import-hoist lint
* Fail closed for hybrid packing on encoder-decoder, chunked-loss, and string-name models
The varlen shim only helps decoder-only hybrid models that run their mixer
through self. on a live nn.Module forward. Three cases slipped past
the guard:
- Encoder-decoder configs (is_encoder_decoder) reached the packing path even
though flattening a cross-attention batch is unsound. Block them explicitly.
- TRL's chunked_nll loss (the 1.x default) calls the backbone directly and
bypasses model.forward, so the per-instance forward wrapper that refreshes
the varlen stash never runs. Detect that path and keep the model padded.
- A string model_name reaches the trainer before the module exists, so the
instance shim has nothing to patch. Resolve the config up front and keep
string hybrids on the padded path.
Adds encoder-decoder / decoder-only / chunked-loss / string-model tests.
* Harden the SFT source-injection replacements and forward auth args for string models
The wrapped-packing injection rewrote the sourced unsloth_zoo sft_prepare_dataset
with str.replace anchored on the exact 'All Unsloth Zoo code licensed under
LGPLv3' comment. str.replace never raises on a missing anchor, so a supported
newer unsloth_zoo (the dependency is only lower-bounded) that moved that header
would silently drop the setup while the truncation and pack_dataset edits still
referenced _unsloth_wrapped_packing / _inspect, raising NameError on every SFT
dataset preparation.
- Install the setup at the sft_prepare_dataset signature via re.subn (a structural
anchor that always exists) and raise if even that is missing.
- Route the remaining edits through a _require_replace helper that fails loudly on a
missing required anchor (or warns once for an optional one), formalizing the
verify-then-replace idiom the DPO patchers in this file already use.
- Reuse the guarded _unsloth_pack_has_strategy at the pack_dataset call instead of
re-calling inspect.signature(pack_dataset) unguarded, so a non-introspectable
pack_dataset cannot crash there after the setup already handled it.
- _resolve_string_model_config now forwards token / use_auth_token / cache_dir /
code_revision, so a private hybrid resolves its config instead of falling through
as non-hybrid and enabling packing without the varlen shim.
Adds regression tests for the drift-resistant injection, the helper, and the
string-model auth forwarding.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Honor top-level SFTConfig.trust_remote_code when resolving a string model
TRL merges the top-level args.trust_remote_code into the load via
model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) before
create_model_from_path, so a remote-code hybrid is commonly set with
SFTConfig(trust_remote_code=True) rather than inside model_init_kwargs. The config
probe only read model_init_kwargs, so AutoConfig could fail for such a model, leave
model_config None, and let the guard treat it as non-hybrid, enabling packing
without the varlen shim. Mirror TRL's setdefault (model_init_kwargs wins).
* Tighten hybrid-packing comments for concision
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: alkinun
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Etherl <61019402+Etherll@users.noreply.github.com>
---
tests/utils/test_packing.py | 579 +++++++++++++++++++++++++++---
unsloth/models/rl_replacements.py | 90 +++--
unsloth/trainer.py | 93 ++++-
unsloth/utils/packing.py | 302 ++++++++++++++++
4 files changed, 984 insertions(+), 80 deletions(-)
diff --git a/tests/utils/test_packing.py b/tests/utils/test_packing.py
index 98c29d9f0f..1b8bb65058 100644
--- a/tests/utils/test_packing.py
+++ b/tests/utils/test_packing.py
@@ -15,6 +15,7 @@
from unsloth import FastLanguageModel
import unsloth.trainer as trainer_module
+import unsloth.utils.packing as packing_module
from unsloth.utils import attention_dispatch as attention_dispatch_utils
from unsloth.utils.packing import (
configure_padding_free,
@@ -22,6 +23,7 @@ from unsloth.utils.packing import (
enable_padding_free_metadata,
enable_sample_packing,
mask_packed_sequence_boundaries,
+ patch_hybrid_linear_attention_varlen,
)
from contextlib import ExitStack
@@ -161,6 +163,327 @@ def test_configure_padding_free():
assert config.remove_unused_columns is False
+# --- Hybrid linear-attention guard + varlen shim (PR #7211 / #7249) ---------------
+
+
+def _hybrid_config_model():
+ # Qwen3.5 / Qwen3-Next style: explicit linear_attention layer schedule.
+ return SimpleNamespace(
+ config = SimpleNamespace(layer_types = ["linear_attention", "full_attention"])
+ )
+
+
+def _gemma3_model():
+ # Has layer_types but no linear_attention -> must NOT be flagged as hybrid.
+ return SimpleNamespace(
+ config = SimpleNamespace(
+ model_type = "gemma3", layer_types = ["sliding_attention", "full_attention"]
+ ),
+ )
+
+
+def _dense_qwen3_model():
+ return SimpleNamespace(
+ config = SimpleNamespace(model_type = "qwen3", architectures = ["Qwen3ForCausalLM"])
+ )
+
+
+class _FakeGatedDeltaNet(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv1d = torch.nn.Conv1d(4, 4, 3, groups = 4)
+ self.A_log = torch.nn.Parameter(torch.zeros(4))
+
+ def forward(self, hidden_states, **kwargs): # dispatch through self.
+ return self.chunk_gated_delta_rule(self.causal_conv1d_fn(hidden_states))
+
+
+class _FakeHybridModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.config = SimpleNamespace() # no markers -> forces module-level detection
+ self.linear_attn = _FakeGatedDeltaNet()
+
+
+def test_is_hybrid_linear_attention_detects_and_excludes():
+ is_hybrid = trainer_module._is_hybrid_linear_attention_model
+ assert is_hybrid(_hybrid_config_model()) is True
+ assert is_hybrid(_FakeHybridModel()) is True # module-structural evidence
+ assert is_hybrid(_text_model()) is False # Llama
+ assert is_hybrid(_gemma3_model()) is False # layer_types without linear_attention
+ assert is_hybrid(_dense_qwen3_model()) is False # dense Qwen3
+ assert is_hybrid(None) is False
+
+
+def test_varlen_from_position_ids():
+ cu, seq_idx = packing_module._varlen_from_position_ids(torch.tensor([[0, 1, 0, 0, 1, 2]]))
+ assert cu.tolist() == [0, 2, 3, 6]
+ assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2]]
+ assert (
+ packing_module._varlen_from_position_ids(torch.tensor([[0, 1, 2, 3]])) is None
+ ) # single sequence
+ assert packing_module._varlen_from_position_ids(torch.tensor([[1, 2, 3]])) is None # first != 0
+ assert (
+ packing_module._varlen_from_position_ids(torch.tensor([[0, 1], [0, 1]])) is None
+ ) # normal 2-row batch
+ assert packing_module._varlen_from_position_ids(None) is None
+
+
+def test_seq_idx_from_cu_seqlens_handles_trailing_pad():
+ cu = torch.tensor([0, 2, 5], dtype = torch.int32)
+ boundaries, seq_idx = packing_module._seq_idx_from_cu_seqlens(cu, total = 8) # pad_to_multiple_of
+ assert boundaries.tolist() == [0, 2, 5, 8]
+ assert seq_idx.tolist() == [[0, 0, 1, 1, 1, 2, 2, 2]]
+ boundaries2, _ = packing_module._seq_idx_from_cu_seqlens(cu, total = 5) # exact fit
+ assert boundaries2.tolist() == [0, 2, 5]
+ assert (
+ packing_module._seq_idx_from_cu_seqlens(torch.tensor([1, 2], dtype = torch.int32), total = 2)
+ is None
+ )
+ assert packing_module._seq_idx_from_cu_seqlens(cu, total = 3) is None # boundaries exceed total
+
+
+def test_hybrid_varlen_metadata_prefers_packed_seq_lengths():
+ # A competing position_ids would segment [0, 3, 6]; packed_seq_lengths must win.
+ kwargs = {
+ "input_ids": torch.zeros(1, 6, dtype = torch.long),
+ "packed_seq_lengths": torch.tensor([2, 1, 3], dtype = torch.int32),
+ "position_ids": torch.tensor([[0, 1, 2, 0, 1, 2]]),
+ }
+ cu, seq_idx = packing_module._hybrid_varlen_metadata(kwargs)
+ assert cu.tolist() == [0, 2, 3, 6]
+ assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2]]
+
+
+def test_hybrid_varlen_metadata_suppressed_when_cached():
+ base = {
+ "input_ids": torch.zeros(1, 6, dtype = torch.long),
+ "packed_seq_lengths": torch.tensor([2, 1, 3], dtype = torch.int32),
+ }
+ assert packing_module._hybrid_varlen_metadata({**base, "use_cache": True}) is None
+ assert packing_module._hybrid_varlen_metadata({**base, "past_key_values": object()}) is None
+
+
+def test_hybrid_varlen_metadata_none_for_plain_batch():
+ kwargs = {
+ "input_ids": torch.zeros(1, 4, dtype = torch.long),
+ "position_ids": torch.tensor([[0, 1, 2, 3]]),
+ }
+ assert packing_module._hybrid_varlen_metadata(kwargs) is None
+
+
+def _make_fake_kernels():
+ def causal_conv1d_fn(
+ x,
+ weight = None,
+ bias = None,
+ activation = None,
+ seq_idx = None,
+ ):
+ causal_conv1d_fn.calls.append(seq_idx)
+ return x
+
+ causal_conv1d_fn.calls = []
+
+ def chunk_gated_delta_rule(
+ q,
+ k = None,
+ v = None,
+ cu_seqlens = None,
+ **kw,
+ ):
+ chunk_gated_delta_rule.calls.append(cu_seqlens)
+ return q
+
+ chunk_gated_delta_rule.calls = []
+ return causal_conv1d_fn, chunk_gated_delta_rule
+
+
+class _ShimGatedDeltaNet(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv1d = torch.nn.Conv1d(4, 4, 3, groups = 4)
+ self.causal_conv1d_fn, self.chunk_gated_delta_rule = _make_fake_kernels()
+
+ def forward(self, hidden_states, **kwargs):
+ return self.chunk_gated_delta_rule(self.causal_conv1d_fn(hidden_states))
+
+
+class _ShimHybridModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.config = SimpleNamespace(layer_types = ["linear_attention", "full_attention"])
+ self.linear_attn = _ShimGatedDeltaNet()
+
+ def forward(
+ self,
+ input_ids = None,
+ position_ids = None,
+ packed_seq_lengths = None,
+ use_cache = None,
+ **kwargs,
+ ):
+ return self.linear_attn(input_ids.float())
+
+
+def test_patch_hybrid_varlen_flag_off(monkeypatch):
+ monkeypatch.delenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", raising = False)
+ model = _ShimHybridModel()
+ assert patch_hybrid_linear_attention_varlen(model) is False
+ assert not getattr(model, "_unsloth_varlen_forward_wrapped", False)
+
+
+def test_patch_hybrid_varlen_active_and_idempotent(monkeypatch):
+ monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
+ model = _ShimHybridModel()
+ conv_orig, scan_orig = (
+ model.linear_attn.causal_conv1d_fn,
+ model.linear_attn.chunk_gated_delta_rule,
+ )
+
+ assert patch_hybrid_linear_attention_varlen(model) is True
+ assert model._unsloth_varlen_forward_wrapped is True
+ assert model.linear_attn._unsloth_varlen_wrapped is True
+ assert patch_hybrid_linear_attention_varlen(model) is True # idempotent, no double-wrap
+
+ conv_orig.calls.clear()
+ scan_orig.calls.clear()
+ packing_module._HYBRID_WARNED.clear()
+ ids = torch.zeros(1, 6, dtype = torch.long)
+ model(
+ input_ids = ids,
+ packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
+ use_cache = False,
+ )
+ assert conv_orig.calls[-1] is not None # seq_idx injected
+ assert scan_orig.calls[-1].tolist() == [0, 2, 3, 6] # cu_seqlens injected
+ assert not packing_module._HYBRID_WARNED # handshake passed, no rejection
+
+ conv_orig.calls.clear()
+ scan_orig.calls.clear()
+ model(
+ input_ids = ids, packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32), use_cache = True
+ )
+ assert conv_orig.calls[-1] is None # cached forward -> no injection
+ assert scan_orig.calls[-1] is None
+
+
+def test_patch_hybrid_varlen_torch_fallback_fail_closed(monkeypatch):
+ monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
+ model = _ShimHybridModel()
+
+ def torch_chunk_gated_delta_rule(
+ q,
+ cu_seqlens = None,
+ **kw,
+ ):
+ return q
+
+ model.linear_attn.chunk_gated_delta_rule = torch_chunk_gated_delta_rule
+ assert patch_hybrid_linear_attention_varlen(model) is False
+ assert not getattr(model, "_unsloth_varlen_forward_wrapped", False)
+
+
+def test_patch_hybrid_varlen_bad_signature_fail_closed(monkeypatch):
+ monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
+ model = _ShimHybridModel()
+
+ def scan_no_cu(q, **kw): # missing cu_seqlens
+ return q
+
+ model.linear_attn.chunk_gated_delta_rule = scan_no_cu
+ assert patch_hybrid_linear_attention_varlen(model) is False
+
+
+def _hybrid_model_with_gdn(gdn_forward):
+ # Build a fake hybrid model whose gated-delta mixer forward is `gdn_forward`.
+ class _GatedDeltaNet(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv1d = torch.nn.Conv1d(4, 4, 3, groups = 4)
+ self.causal_conv1d_fn, self.chunk_gated_delta_rule = _make_fake_kernels()
+
+ forward = gdn_forward
+
+ class _Model(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.config = SimpleNamespace(layer_types = ["linear_attention", "full_attention"])
+ self.linear_attn = _GatedDeltaNet()
+
+ def forward(
+ self,
+ input_ids = None,
+ packed_seq_lengths = None,
+ use_cache = None,
+ **kwargs,
+ ):
+ return self.linear_attn(input_ids.float())
+
+ return _Model()
+
+
+def test_patch_hybrid_varlen_no_dispatch_aborts(monkeypatch):
+ # Dispatch is verified at runtime, not statically. A mixer that never calls
+ # self. installs the shim, but the first packed forward aborts (both
+ # boundary kernels are load-bearing).
+ monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
+ model = _hybrid_model_with_gdn(lambda self, hidden_states, **kw: hidden_states)
+ assert patch_hybrid_linear_attention_varlen(model) is True # kernels valid -> installs
+ with pytest.raises(RuntimeError, match = "both invoked"):
+ model(
+ input_ids = torch.zeros(1, 6),
+ packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
+ use_cache = False,
+ )
+
+
+def test_patch_hybrid_varlen_partial_dispatch_aborts(monkeypatch):
+ # Only the conv fires; the scan would leak state. Both must be invoked, so abort.
+ monkeypatch.setenv("UNSLOTH_EXPERIMENTAL_HYBRID_PACKING", "1")
+ conv_only = _hybrid_model_with_gdn(
+ lambda self, hidden_states, **kw: self.causal_conv1d_fn(hidden_states)
+ )
+ assert patch_hybrid_linear_attention_varlen(conv_only) is True
+ with pytest.raises(RuntimeError, match = "both invoked"):
+ conv_only(
+ input_ids = torch.zeros(1, 6),
+ packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
+ use_cache = False,
+ )
+
+ scan_only = _hybrid_model_with_gdn(
+ lambda self, hidden_states, **kw: self.chunk_gated_delta_rule(hidden_states)
+ )
+ assert patch_hybrid_linear_attention_varlen(scan_only) is True
+ with pytest.raises(RuntimeError, match = "both invoked"):
+ scan_only(
+ input_ids = torch.zeros(1, 6),
+ packed_seq_lengths = torch.tensor([2, 1, 3], dtype = torch.int32),
+ use_cache = False,
+ )
+
+
+def test_varlen_from_position_ids_mrope_3d():
+ pos = (
+ torch.tensor([[0, 1, 0, 0, 1, 2]]).unsqueeze(0).expand(3, 1, 6).clone()
+ ) # [3,1,T] text plane
+ cu, seq_idx = packing_module._varlen_from_position_ids(pos)
+ assert cu.tolist() == [0, 2, 3, 6]
+ assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2]]
+
+
+def test_hybrid_varlen_metadata_trailing_pad():
+ # packed_seq_lengths sum to 6 but the flattened input is 8 (pad_to_multiple_of).
+ kwargs = {
+ "input_ids": torch.zeros(1, 8, dtype = torch.long),
+ "packed_seq_lengths": torch.tensor([2, 1, 3], dtype = torch.int32),
+ }
+ cu, seq_idx = packing_module._hybrid_varlen_metadata(kwargs)
+ assert cu.tolist() == [0, 2, 3, 6, 8]
+ assert seq_idx.tolist() == [[0, 0, 1, 2, 2, 2, 3, 3]]
+
+
def _patch_fake_sft_trainer():
class FakeSFTTrainer:
def __init__(self, *args, **kwargs):
@@ -245,29 +568,101 @@ def test_vlm_without_processing_class_still_disables_packing():
("t5", "T5ForConditionalGeneration"),
("bart", "BartForConditionalGeneration"),
("whisper", "WhisperForConditionalGeneration"),
- ("csm", "CsmForConditionalGeneration"),
),
)
-def test_nonvision_conditional_generation_keeps_packing(model_type, architecture):
+def test_encoder_decoder_disables_packing(model_type, architecture):
+ # Text-only encoder-decoder models are not VLMs, but their bidirectional encoder
+ # attends across concatenated samples once padding-free drops attention_mask.
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
model = SimpleNamespace(
- config = SimpleNamespace(model_type = model_type, architectures = [architecture]),
+ config = SimpleNamespace(
+ model_type = model_type,
+ architectures = [architecture],
+ is_encoder_decoder = True,
+ ),
max_seq_length = 16,
)
- trainer = fake_trainer(
- model,
- config,
- None,
- Dataset.from_dict({"text": ["text-only sample"]}),
+ trainer = fake_trainer(model, config, None, Dataset.from_dict({"text": ["text-only sample"]}))
+
+ assert config.packing is False
+ assert config.padding_free is False
+
+
+def test_decoder_only_conditional_generation_keeps_packing():
+ # CSM is decoder-only despite the ForConditionalGeneration name -> packing stays on.
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ model = SimpleNamespace(
+ config = SimpleNamespace(
+ model_type = "csm",
+ architectures = ["CsmForConditionalGeneration"],
+ is_encoder_decoder = False,
+ ),
+ max_seq_length = 16,
)
+ trainer = fake_trainer(model, config, None, Dataset.from_dict({"text": ["text-only sample"]}))
+
assert config.packing is True
assert config.padding_free is True
assert trainer.model._unsloth_allow_packed_overlength is True
+def _hybrid_trainer_model():
+ return SimpleNamespace(
+ config = SimpleNamespace(
+ model_type = "qwen3_next",
+ architectures = ["Qwen3NextForCausalLM"],
+ layer_types = ["linear_attention", "full_attention"],
+ ),
+ max_seq_length = 16,
+ )
+
+
+def test_hybrid_varlen_active_enables_packing(monkeypatch):
+ # Baseline: shim active + no forward bypass -> hybrid packing is allowed.
+ monkeypatch.setattr(trainer_module, "_chunked_loss_bypasses_forward", lambda config: False)
+ monkeypatch.setattr(trainer_module, "patch_hybrid_linear_attention_varlen", lambda model: True)
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ fake_trainer(_hybrid_trainer_model(), config, None, Dataset.from_dict({"text": ["x"]}))
+ assert config.packing is True
+ assert config.padding_free is True
+
+
+def test_hybrid_chunked_loss_stays_on_padded_path(monkeypatch):
+ # TRL's chunked-loss forward bypass leaves the varlen shim off -> block packing.
+ monkeypatch.setattr(trainer_module, "_chunked_loss_bypasses_forward", lambda config: True)
+ monkeypatch.setattr(trainer_module, "patch_hybrid_linear_attention_varlen", lambda model: True)
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ fake_trainer(_hybrid_trainer_model(), config, None, Dataset.from_dict({"text": ["x"]}))
+ assert config.packing is False
+ assert config.padding_free is False
+
+
+def test_string_hybrid_model_disables_packing(monkeypatch):
+ # A string model= is materialized after init; a hybrid string is blocked because the
+ # shim cannot patch a not-yet-built model.
+ monkeypatch.setattr(
+ trainer_module,
+ "_resolve_string_model_config",
+ lambda name, cfg: SimpleNamespace(
+ model_type = "qwen3_next",
+ architectures = ["Qwen3NextForCausalLM"],
+ layer_types = ["linear_attention", "full_attention"],
+ ),
+ )
+ monkeypatch.setattr(trainer_module, "patch_hybrid_linear_attention_varlen", lambda model: True)
+ fake_trainer = _patch_fake_sft_trainer()
+ config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
+ fake_trainer("Qwen/Qwen3-Next-80B-A3B", config, None, Dataset.from_dict({"text": ["x"]}))
+ assert config.packing is False
+ assert config.padding_free is False
+
+
def test_vlm_vision_dataset_still_disables_packing():
fake_trainer = _patch_fake_sft_trainer()
config = SimpleNamespace(packing = True, padding_free = None, remove_unused_columns = True)
@@ -486,49 +881,6 @@ def test_wrapped_packing_preserves_overlength_tokens(monkeypatch, legacy_api):
assert all(len(input_ids) <= args.max_length for input_ids in packed_ids)
-# Named to match the unsloth_zoo helper: sft_trainer_prepare_dataset sources it by
-# name and renames "def sft_prepare_dataset" -> "def _prepare_dataset". This fixture
-# deliberately omits the "All Unsloth Zoo code licensed under LGPLv3" header to emulate
-# a newer, compatible Zoo whose header moved (the dependency is only lower-bounded).
-def sft_prepare_dataset(
- self, dataset, processing_class, args, packing, formatting_func, dataset_text_field
-):
- do_truncation = True
- # Mirror the Zoo call so the "truncation = do_truncation," injection anchor
- # survives formatting (a bare tuple assignment gets rewritten to a paren form).
- dataset = processing_class(
- dataset,
- truncation = do_truncation,
- )
- return dataset
-
-
-def test_wrapped_packing_setup_survives_missing_zoo_header(monkeypatch):
- # Regression: the wrapped-packing setup used to anchor on the Zoo license comment,
- # so a header change made it a no-op while the truncation reference still landed,
- # NameError-ing every SFT dataset preparation. It must now install via the
- # signature and always precede the reference.
- import ast
- import textwrap
- import unsloth.models.rl_replacements as rlr
-
- monkeypatch.setitem(rlr.RL_REPLACEMENTS, "sft_prepare_dataset", sft_prepare_dataset)
-
- source = (
- "def _prepare_dataset(self, dataset, processing_class, args, packing, "
- "formatting_func, dataset_text_field):\n return dataset\n"
- )
- patched = rlr.sft_trainer_prepare_dataset("_prepare_dataset", source)
-
- assert "_unsloth_wrapped_packing = packing" in patched
- assert "import inspect as _inspect" in patched
- assert "not _unsloth_wrapped_packing" in patched
- assert patched.index("_unsloth_wrapped_packing = packing") < patched.index(
- "truncation = do_truncation and not _unsloth_wrapped_packing"
- )
- ast.parse(textwrap.dedent(patched))
-
-
class _DummyChild(torch.nn.Module):
def __init__(self):
super().__init__()
@@ -759,3 +1111,128 @@ def test_packing_sdpa(tmp_path):
if hasattr(trainer, "accelerator"):
trainer.accelerator.free_memory()
+
+
+# --- wrapped-packing source-injection robustness (reviewer.py / fork findings) --------
+
+
+# fmt: off
+# Named to match the unsloth_zoo helper (sourced by name, "def sft_prepare_dataset" ->
+# "def _prepare_dataset"). Deliberately OMITS the "licensed under LGPLv3" header to
+# emulate a newer Zoo whose header moved (dependency is only lower-bounded). Source only.
+def sft_prepare_dataset(
+ self, dataset, processing_class, args, packing, formatting_func, dataset_text_field
+):
+ do_truncation = True
+ max_seq_length = 4
+ used_column_names = ["text"]
+ map_kwargs = {}
+ dataset = processing_class(dataset, truncation = do_truncation,)
+ if do_truncation and max_seq_length > 0:
+ pass
+ if packing:
+ dataset = pack_dataset(
+ dataset.select_columns(used_column_names),
+ max_seq_length,
+ getattr(args, "packing_strategy", "bfd"),
+ map_kwargs,
+ )
+ return dataset
+# fmt: on
+
+
+def test_wrapped_packing_injection_is_drift_resistant(monkeypatch):
+ # Regression: the setup used to anchor on the Zoo license comment, so a header
+ # change silently no-op'd it while the truncation/pack edits still referenced its
+ # variables -> NameError on every SFT prep. It must now install via the signature
+ # before those references, and the pack edit must reuse the guarded
+ # _unsloth_pack_has_strategy instead of re-calling _inspect.signature(pack_dataset).
+ import ast
+ import textwrap
+ import unsloth.models.rl_replacements as rlr
+
+ monkeypatch.setitem(rlr.RL_REPLACEMENTS, "sft_prepare_dataset", sft_prepare_dataset)
+
+ source = (
+ "def _prepare_dataset(self, dataset, processing_class, args, packing, "
+ "formatting_func, dataset_text_field):\n return dataset\n"
+ )
+ patched = rlr.sft_trainer_prepare_dataset("_prepare_dataset", source)
+
+ # setup installed despite the missing header, and before it is referenced
+ assert "_unsloth_wrapped_packing = packing" in patched
+ assert "import inspect as _inspect" in patched
+ assert patched.index("_unsloth_wrapped_packing = packing") < patched.index(
+ "truncation = do_truncation and not _unsloth_wrapped_packing"
+ )
+ # the pack edit reuses the guarded flag (signature inspected exactly once, in setup)
+ assert "if _unsloth_pack_has_strategy:" in patched
+ assert patched.count("_inspect.signature(pack_dataset)") == 1
+ ast.parse(textwrap.dedent(patched))
+
+
+def test_require_replace_raises_on_missing_anchor():
+ from unsloth.models.rl_replacements import _require_replace
+
+ assert _require_replace("abc", "b", "B") == "aBc"
+ with pytest.raises(RuntimeError):
+ _require_replace("abc", "z", "Z", where = "unit test")
+ # an optional edit warns once and returns the source unchanged (no dangling ref)
+ assert _require_replace("abc", "z", "Z", required = False, where = "optional") == "abc"
+
+
+def test_resolve_string_model_config_forwards_token(monkeypatch):
+ import transformers
+
+ captured = {}
+
+ class _FakeAutoConfig:
+ @staticmethod
+ def from_pretrained(name, **kwargs):
+ captured.update(kwargs)
+ return SimpleNamespace(is_encoder_decoder = False)
+
+ monkeypatch.setattr(transformers, "AutoConfig", _FakeAutoConfig)
+
+ config_arg = SimpleNamespace(
+ model_init_kwargs = {
+ "token": "hf_secret",
+ "trust_remote_code": True,
+ "cache_dir": "/tmp/cache",
+ "torch_dtype": "bfloat16", # not a config arg -> must NOT be forwarded
+ }
+ )
+ result = trainer_module._resolve_string_model_config("org/private-hybrid", config_arg)
+
+ assert result is not None
+ assert captured.get("token") == "hf_secret"
+ assert captured.get("trust_remote_code") is True
+ assert captured.get("cache_dir") == "/tmp/cache"
+ assert "torch_dtype" not in captured
+
+
+def test_resolve_string_model_config_merges_top_level_trust_remote_code(monkeypatch):
+ import transformers
+
+ captured = {}
+
+ class _FakeAutoConfig:
+ @staticmethod
+ def from_pretrained(name, **kwargs):
+ captured.update(kwargs)
+ return SimpleNamespace(is_encoder_decoder = False)
+
+ monkeypatch.setattr(transformers, "AutoConfig", _FakeAutoConfig)
+
+ # SFTConfig(trust_remote_code=True) with no model_init_kwargs entry is honored
+ config_arg = SimpleNamespace(model_init_kwargs = {}, trust_remote_code = True)
+ trainer_module._resolve_string_model_config("org/remote-hybrid", config_arg)
+ assert captured.get("trust_remote_code") is True
+
+ # model_init_kwargs wins over the top-level flag (mirrors TRL's setdefault)
+ captured.clear()
+ config_arg = SimpleNamespace(
+ model_init_kwargs = {"trust_remote_code": False}, trust_remote_code = True
+ )
+ trainer_module._resolve_string_model_config("org/remote-hybrid", config_arg)
+ assert captured.get("trust_remote_code") is False
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index b0709f7376..4ef3af6add 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -437,6 +437,52 @@ RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_compute_loss_liger)
RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_data_collator_vision_keys)
+_WRAPPED_PACKING_SETUP = (
+ " import inspect as _inspect\n"
+ " try:\n"
+ ' _unsloth_pack_has_strategy = "strategy" in _inspect.signature(pack_dataset).parameters\n'
+ " except Exception:\n"
+ " _unsloth_pack_has_strategy = True\n"
+ " _unsloth_wrapped_packing = packing and (\n"
+ ' getattr(args, "packing_strategy", None) == "wrapped"\n'
+ " or not _unsloth_pack_has_strategy\n"
+ " )\n"
+)
+
+_WARNED_MISSING_ANCHORS = set()
+
+
+def _require_replace(
+ function,
+ old,
+ new,
+ *,
+ count = 1,
+ required = True,
+ where = "",
+):
+ """str.replace that never silently no-ops a load-bearing source edit.
+
+ Plain str.replace returns the source unchanged when the anchor is absent, so a
+ drifted anchor in a newer TRL / unsloth_zoo would skip the edit while later edits
+ still reference helper variables it should have introduced (NameError at runtime).
+ Fail loudly for a required edit, warn once and skip for an optional one, so a
+ drifted source can never corrupt the patched function silently.
+ """
+ if old not in function:
+ detail = f" ({where})" if where else ""
+ if required:
+ raise RuntimeError(
+ f"Unsloth: source anchor not found{detail}; the patched function is out "
+ "of sync with this TRL / unsloth_zoo version. Please file a bug report."
+ )
+ if where not in _WARNED_MISSING_ANCHORS:
+ _WARNED_MISSING_ANCHORS.add(where)
+ logger.warning(f"Unsloth: skipped an optional source edit{detail} (anchor not found).")
+ return function
+ return function.replace(old, new, count)
+
+
# Fix tokenizer double BOS
def sft_trainer_prepare_dataset(function_name, function):
if function_name != "_prepare_non_packed_dataloader" and function_name != "_prepare_dataset":
@@ -454,27 +500,14 @@ def sft_trainer_prepare_dataset(function_name, function):
if matched:
# Use fast version!
function = inspect.getsource(fast_sft_prepare_dataset)
- # why: install the wrapped-packing setup (and the `_inspect` import the
- # truncation / pack_dataset rewrites below depend on) at the function
- # signature, a structural anchor that always exists, rather than the
- # unsloth_zoo license-comment line. That header is only lower-bounded, so a
- # newer Zoo may move or drop it; anchoring there let the setup silently
- # no-op while the references still landed, NameError-ing every SFT dataset
- # preparation. Fail loudly if even the signature cannot be located.
- _wrapped_packing_setup = (
- " import inspect as _inspect\n"
- " try:\n"
- ' _unsloth_pack_has_strategy = "strategy" in _inspect.signature(pack_dataset).parameters\n'
- " except Exception:\n"
- " _unsloth_pack_has_strategy = True\n"
- " _unsloth_wrapped_packing = packing and (\n"
- ' getattr(args, "packing_strategy", None) == "wrapped"\n'
- " or not _unsloth_pack_has_strategy\n"
- " )\n"
- )
+ # why: anchor the wrapped-packing setup on the function signature -- a
+ # structural anchor that always exists -- not the unsloth_zoo license comment,
+ # which is only lower-bounded and a newer Zoo may move or drop. Anchoring there
+ # let the setup silently no-op while edits below referenced its variables,
+ # NameError-ing every SFT dataset prep. Fail loudly if the signature is missing.
function, _n_setup = re.subn(
r"(def sft_prepare_dataset\s*\(.*?\)\s*(?:->[^:\n]*)?:[ \t]*\n)",
- lambda match: match.group(1) + _wrapped_packing_setup,
+ lambda match: match.group(1) + _WRAPPED_PACKING_SETUP,
function,
count = 1,
flags = re.DOTALL,
@@ -484,15 +517,25 @@ def sft_trainer_prepare_dataset(function_name, function):
"Unsloth: failed to install wrapped-packing support into "
"sft_prepare_dataset (signature not found); please file a bug report."
)
- function = function.replace(
+ # why: route each edit through _require_replace so a drifted anchor fails
+ # loudly instead of leaving a dangling reference to the setup variables.
+ function = _require_replace(
+ function,
"truncation = do_truncation,",
"truncation = do_truncation and not _unsloth_wrapped_packing,",
+ where = "sft_prepare_dataset truncation flag",
)
- function = function.replace(
+ function = _require_replace(
+ function,
"if do_truncation and max_seq_length > 0:",
"if do_truncation and not _unsloth_wrapped_packing and max_seq_length > 0:",
+ where = "sft_prepare_dataset truncation guard",
)
- function = function.replace(
+ # why: reuse the guarded _unsloth_pack_has_strategy from the setup instead of
+ # re-calling _inspect.signature(pack_dataset) here -- the setup wraps that call
+ # in try/except, so a non-introspectable pack_dataset must not crash here.
+ function = _require_replace(
+ function,
"""dataset = pack_dataset(
dataset.select_columns(used_column_names),
max_seq_length,
@@ -500,13 +543,14 @@ def sft_trainer_prepare_dataset(function_name, function):
map_kwargs,
)""",
"""_pack_kwargs = {"map_kwargs": map_kwargs}
- if "strategy" in _inspect.signature(pack_dataset).parameters:
+ if _unsloth_pack_has_strategy:
_pack_kwargs["strategy"] = getattr(args, "packing_strategy", "bfd")
dataset = pack_dataset(
dataset.select_columns(used_column_names),
max_seq_length,
**_pack_kwargs,
)""",
+ where = "sft_prepare_dataset pack_dataset call",
)
function = function.split("\n")
function = "\n".join(" " * 4 + x for x in function)
diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index 61d41aad21..1c30192301 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -17,6 +17,7 @@ import os
import psutil
import warnings
from dataclasses import dataclass, field
+from types import SimpleNamespace
from typing import Optional, List
from functools import wraps
@@ -32,6 +33,7 @@ from unsloth.utils import (
enable_padding_free_metadata,
enable_sample_packing,
)
+from unsloth.utils.packing import patch_hybrid_linear_attention_varlen
from unsloth_zoo.training_utils import (
unsloth_train as _unsloth_train,
)
@@ -101,9 +103,9 @@ PADDING_FREE_BLOCKLIST = {
"gpt_oss", # - gpt_oss: Uses Flex Attention which doesn't handle padding_free correctly
}
# Hybrid linear-attention / state-space models (Qwen3.5, Qwen3-Next, ...) carry a
-# recurrent gated-delta state plus a causal conv1d. Sample packing / padding-free
-# flattens the batch, so those ops leak state across sequence boundaries. Detected
-# structurally by _is_hybrid_linear_attention_model rather than by model name.
+# recurrent gated-delta state plus a causal conv1d that leak across sequence
+# boundaries once packing flattens the batch. Detected structurally by
+# _is_hybrid_linear_attention_model, not by model name.
def _should_pack(config) -> bool:
@@ -267,6 +269,57 @@ def _is_hybrid_linear_attention_model(model) -> bool:
return False
+def _resolve_string_model_config(model_name, config_arg):
+ """TRL materializes a string ``model=`` inside ``__init__``; resolve its config
+ up front so the packing guards run before the dataset is packed. Best-effort:
+ returns None if the config cannot be loaded."""
+ try:
+ from transformers import AutoConfig
+
+ init_kwargs = getattr(config_arg, "model_init_kwargs", None) or {}
+ # why: forward auth + cache args too. Dropping token/use_auth_token made a
+ # private hybrid fail to load (resolve as None) -> treated as non-hybrid ->
+ # packing enabled without the shim even though TRL later loads it with the token.
+ forward = {
+ key: init_kwargs[key]
+ for key in (
+ "trust_remote_code",
+ "revision",
+ "subfolder",
+ "token",
+ "use_auth_token",
+ "cache_dir",
+ "code_revision",
+ )
+ if key in init_kwargs
+ }
+ # why: TRL merges top-level args.trust_remote_code into the load via setdefault
+ # before create_model_from_path, so honor it here (model_init_kwargs wins), else
+ # a remote-code hybrid with SFTConfig(trust_remote_code=True) resolves as None
+ # and skips the guard.
+ top_level_trust_remote_code = getattr(config_arg, "trust_remote_code", None)
+ if top_level_trust_remote_code is not None:
+ forward.setdefault("trust_remote_code", top_level_trust_remote_code)
+ return AutoConfig.from_pretrained(model_name, **forward)
+ except Exception:
+ return None
+
+
+def _chunked_loss_bypasses_forward(config) -> bool:
+ """TRL's default ``loss_type="chunked_nll"`` patches the model forward and calls
+ the backbone directly, so a forward wrapper never runs. Detect it so hybrid
+ packing stays on the padded path instead of silently skipping the varlen shim."""
+ try:
+ import trl.trainer.sft_trainer as _sft_trainer
+ except Exception:
+ return False
+ if not hasattr(_sft_trainer, "_patch_chunked_ce_lm_head"):
+ return False # TRL has no chunked-CE path -> forward is not bypassed
+ if getattr(config, "use_liger_kernel", False):
+ return False # liger forces loss_type="nll" -> normal forward
+ return getattr(config, "loss_type", None) in (None, "chunked_nll")
+
+
# Unsloth gradient accumulation fix:
from transformers import __version__ as transformers_version, ProcessorMixin
@@ -632,13 +685,38 @@ def _patch_sft_trainer_auto_packing(trl_module):
is_vlm = False
is_unsupported_model = False
is_hybrid = False
+ is_encoder_decoder = False
+ hybrid_varlen_active = False
if model is not None:
model_config = getattr(model, "config", None)
+ if model_config is None and isinstance(model, str):
+ # TRL builds a string model inside __init__; resolve its config now.
+ model_config = _resolve_string_model_config(model, config_arg)
if model_config is not None:
model_types = get_transformers_model_type(model_config)
is_unsupported_model = any(x in PADDING_FREE_BLOCKLIST for x in model_types)
is_vlm = _is_vlm_config(model_config, model_types)
- is_hybrid = _is_hybrid_linear_attention_model(model)
+ is_encoder_decoder = bool(getattr(model_config, "is_encoder_decoder", False))
+ hybrid_target = (
+ SimpleNamespace(config = model_config)
+ if isinstance(model, str) and model_config is not None
+ else model
+ )
+ is_hybrid = _is_hybrid_linear_attention_model(hybrid_target)
+ # Hybrid models corrupt packed batches unless the gated-delta conv + scan
+ # reset at sequence boundaries. Enable the experimental varlen shim (flag +
+ # kernels) so packing stays correct, else keep them blocked. A string model
+ # (patched only after init) and TRL's chunked-loss forward bypass both leave
+ # the shim off, so hybrid packing falls back to the padded path.
+ if (
+ is_hybrid
+ and not isinstance(model, str)
+ and not _chunked_loss_bypasses_forward(config_arg)
+ ):
+ try:
+ hybrid_varlen_active = patch_hybrid_linear_attention_varlen(model)
+ except Exception:
+ hybrid_varlen_active = False
processing_class = (
args[5] if len(args) >= 6 else kwargs.get("processing_class") or kwargs.get("tokenizer")
@@ -664,7 +742,8 @@ def _patch_sft_trainer_auto_packing(trl_module):
or is_auto_processor_vlm
or is_vision_dataset
or is_unsupported_model
- or is_hybrid
+ or is_encoder_decoder
+ or (is_hybrid and not hybrid_varlen_active)
or (
os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
) # Disable padding free on forced logits
@@ -684,7 +763,9 @@ def _patch_sft_trainer_auto_packing(trl_module):
reason = "vision-language model with auto processor"
elif is_vision_dataset:
reason = "vision dataset"
- elif is_hybrid:
+ elif is_encoder_decoder:
+ reason = "encoder-decoder model"
+ elif is_hybrid and not hybrid_varlen_active:
reason = "hybrid linear-attention model"
elif is_unsupported_model:
reason = f"unsupported model type(s): {', '.join(model_types)}"
diff --git a/unsloth/utils/packing.py b/unsloth/utils/packing.py
index dd0a1bfb62..f8d539fb93 100644
--- a/unsloth/utils/packing.py
+++ b/unsloth/utils/packing.py
@@ -17,8 +17,11 @@
from __future__ import annotations
+import inspect
import logging
+import os
from collections import OrderedDict
+from functools import wraps
from typing import Any, Iterable, Optional, Sequence, Tuple
import torch
@@ -218,6 +221,305 @@ def enable_padding_free_metadata(model, trainer):
collator._unsloth_padding_free_lengths_wrapped = True
+# --- Experimental: correct packing / padding-free for hybrid linear-attention ---
+# Qwen3.5 / Qwen3-Next mix a gated-delta recurrence with a causal conv1d. Packing
+# flattens the batch, and both ops leak state across sequence boundaries unless we
+# pass seq_idx (conv) and cu_seqlens (scan). Only the accelerated kernels accept
+# these, so we fail closed on the pure-torch fallbacks. Gated behind an env flag.
+#
+# Overrides only the per-module prefill kernels (causal_conv1d_fn /
+# chunk_gated_delta_rule), leaving decode untouched so generation is unaffected.
+# Recompute-safe under gradient checkpointing; never fires for cached forwards.
+# Feature-detect (never version-detect), fail closed, idempotent, one deduped
+# diagnostic when it declines to activate.
+_HYBRID_PACKING_ENV_VAR = "UNSLOTH_EXPERIMENTAL_HYBRID_PACKING"
+_HYBRID_LOGGER = logging.getLogger("unsloth.hybrid_packing")
+_HYBRID_WARNED: set = set()
+
+
+def _hybrid_packing_enabled() -> bool:
+ # Read at call time so setting the flag after `import unsloth` still takes effect.
+ return os.environ.get(_HYBRID_PACKING_ENV_VAR, "0").strip().lower() in {
+ "1",
+ "true",
+ "yes",
+ "on",
+ }
+
+
+def _hybrid_reject(reason: str) -> bool:
+ # One deduped diagnostic explaining why hybrid packing stayed on the padded path.
+ if reason not in _HYBRID_WARNED:
+ _HYBRID_WARNED.add(reason)
+ _HYBRID_LOGGER.warning(
+ "Unsloth: hybrid linear-attention packing disabled (padded path): %s.",
+ reason,
+ )
+ return False
+
+
+def _iter_gated_delta_modules(model):
+ modules, seen = [], set()
+ for module in model.modules():
+ if id(module) in seen:
+ continue
+ seen.add(id(module))
+ if type(module).__name__.endswith("GatedDeltaNet") and hasattr(module, "conv1d"):
+ modules.append(module)
+ return modules
+
+
+def _hybrid_varlen_kernels_available(gated_delta_modules) -> Optional[str]:
+ """None if every module can use the accelerated varlen path, else a short
+ reason string. All modules are validated before any are mutated; signatures
+ are read off the captured originals when already wrapped.
+
+ Dispatch (the mixer actually calling self.causal_conv1d_fn /
+ self.chunk_gated_delta_rule) is verified at RUNTIME by the forward-wrapper
+ handshake, not statically: Unsloth's compile-disable shim hides it from
+ inspect.getsource, and every supported transformers release dispatches
+ through the instance attribute."""
+ if not gated_delta_modules:
+ return "no gated-delta modules found"
+ for module in gated_delta_modules:
+ conv = getattr(module, "_unsloth_varlen_orig_conv", None) or getattr(
+ module,
+ "causal_conv1d_fn",
+ None,
+ )
+ scan = getattr(module, "_unsloth_varlen_orig_scan", None) or getattr(
+ module,
+ "chunk_gated_delta_rule",
+ None,
+ )
+ if conv is None or scan is None:
+ return "accelerated kernels missing (install causal_conv1d and fla)"
+ if getattr(scan, "__name__", "").startswith("torch_") or getattr(
+ conv,
+ "__name__",
+ "",
+ ).startswith("torch_"):
+ return "pure-torch kernel fallback in use"
+ try:
+ if "seq_idx" not in inspect.signature(conv).parameters:
+ return "conv kernel does not accept seq_idx"
+ if "cu_seqlens" not in inspect.signature(scan).parameters:
+ return "scan kernel does not accept cu_seqlens"
+ except (TypeError, ValueError):
+ return "kernel signature not introspectable"
+ return None
+
+
+def _varlen_from_position_ids(position_ids):
+ """(cu_seqlens int32[n+1], seq_idx int32[1,T]) for a flattened padding-free
+ batch, else None. Padding-free position_ids reset to 0 at each sequence start;
+ accepts only a validated single-row pack (normal batch or single sequence ->
+ None). Fallback used only when packed_seq_lengths is absent: it assumes
+ right-packed reset position_ids and would mis-segment a left-padded row, which
+ is why packed_seq_lengths is always preferred."""
+ if position_ids is None:
+ return None
+ pos = position_ids
+ if pos.dim() == 3: # MRoPE [n_planes, 1, T] -> text plane is index 0
+ pos = pos[0]
+ if pos.dim() != 2 or pos.shape[0] != 1:
+ return None
+ row = pos[0]
+ total = row.shape[0]
+ starts = (row == 0).nonzero(as_tuple = False).flatten()
+ if starts.numel() <= 1 or int(starts[0].item()) != 0:
+ return None
+ cu_seqlens = torch.cat(
+ [
+ starts.to(torch.int32),
+ torch.tensor([total], dtype = torch.int32, device = row.device),
+ ]
+ )
+ return _seq_idx_from_cu_seqlens(cu_seqlens, total)
+
+
+def _seq_idx_from_cu_seqlens(cu_seqlens, total):
+ """(cu_seqlens int32[n+1], seq_idx int32[1,total]) partitioning [0, total),
+ else None. Appends a trailing segment for pad_to_multiple_of zero tokens so the
+ boundaries always cover the full flattened length the kernels see."""
+ if cu_seqlens is None or cu_seqlens.numel() < 2 or int(cu_seqlens[0].item()) != 0:
+ return None
+ boundaries = cu_seqlens.to(torch.int32)
+ last = int(boundaries[-1].item())
+ if last > total:
+ return None
+ if last < total: # trailing pad tokens -> one final segment
+ boundaries = torch.cat(
+ [
+ boundaries,
+ torch.tensor([total], dtype = torch.int32, device = boundaries.device),
+ ]
+ )
+ lengths = boundaries[1:] - boundaries[:-1]
+ if not bool((lengths > 0).all()):
+ return None
+ seq_idx = torch.repeat_interleave(
+ torch.arange(lengths.numel(), dtype = torch.int32, device = boundaries.device),
+ lengths.to(torch.int64),
+ ).unsqueeze(0)
+ return boundaries, seq_idx
+
+
+def _hybrid_varlen_metadata(kwargs):
+ """Boundary metadata (cu_seqlens, seq_idx) for one flattened packed forward,
+ else None. Prefers the authoritative packed_seq_lengths, falls back to
+ reset-style position_ids. Returns None for cached forwards and non-packed
+ batches so decode / eval / normal batches are a strict no-op."""
+ if kwargs.get("use_cache"):
+ return None
+ if kwargs.get("past_key_values") is not None or kwargs.get("cache_params") is not None:
+ return None
+ total, device = None, None
+ for key in ("input_ids", "inputs_embeds", "position_ids"):
+ tensor = kwargs.get(key)
+ if tensor is not None and hasattr(tensor, "shape"):
+ total = tensor.shape[1] if key == "inputs_embeds" else tensor.shape[-1]
+ device = tensor.device
+ break
+ if total is None:
+ return None
+ psl = kwargs.get("packed_seq_lengths")
+ if psl is not None and getattr(psl, "numel", lambda: 1)() > 0: # skip empty (no max())
+ info = get_packed_info_from_kwargs(kwargs, device)
+ if info is not None:
+ _, cu_seqlens, _ = info
+ built = _seq_idx_from_cu_seqlens(cu_seqlens, total)
+ if built is not None:
+ return built
+ return _varlen_from_position_ids(kwargs.get("position_ids"))
+
+
+def patch_hybrid_linear_attention_varlen(model) -> bool:
+ """Feed seq_idx / cu_seqlens to the gated-delta conv + scan so packing and
+ padding-free reset state at sequence boundaries. Gated by
+ UNSLOTH_EXPERIMENTAL_HYBRID_PACKING and fail-closed. Returns True when the
+ varlen path is active, so the caller may allow packing for the model.
+ Idempotent: repeat calls on an already-patched model return True."""
+ if not _hybrid_packing_enabled():
+ return False
+ gated_delta_modules = _iter_gated_delta_modules(model)
+
+ # Idempotency: an already fully-patched model stays active without re-validation.
+ if (
+ getattr(model, "_unsloth_varlen_forward_wrapped", False)
+ and gated_delta_modules
+ and all(getattr(m, "_unsloth_varlen_wrapped", False) for m in gated_delta_modules)
+ ):
+ return True
+
+ reason = _hybrid_varlen_kernels_available(gated_delta_modules)
+ if reason is not None:
+ return _hybrid_reject(reason)
+
+ # Transactional: every module validated above, now wrap each and stash originals.
+ for module in gated_delta_modules:
+ if getattr(module, "_unsloth_varlen_wrapped", False):
+ continue
+ conv_orig, scan_orig = module.causal_conv1d_fn, module.chunk_gated_delta_rule
+ module._unsloth_varlen_orig_conv = conv_orig
+ module._unsloth_varlen_orig_scan = scan_orig
+
+ @wraps(conv_orig)
+ def conv_fn(
+ *args,
+ _orig = conv_orig,
+ _module = module,
+ **kwargs,
+ ):
+ varlen = getattr(_module, "_unsloth_varlen", None)
+ if varlen is not None:
+ _module._unsloth_varlen_conv_hit = True # runtime dispatch handshake
+ if kwargs.get("seq_idx") is None:
+ kwargs["seq_idx"] = varlen[1]
+ return _orig(*args, **kwargs)
+
+ @wraps(scan_orig)
+ def scan_fn(
+ *args,
+ _orig = scan_orig,
+ _module = module,
+ **kwargs,
+ ):
+ varlen = getattr(_module, "_unsloth_varlen", None)
+ if varlen is not None:
+ _module._unsloth_varlen_scan_hit = True
+ if kwargs.get("cu_seqlens") is None:
+ kwargs["cu_seqlens"] = varlen[0]
+ return _orig(*args, **kwargs)
+
+ module.causal_conv1d_fn = conv_fn
+ module.chunk_gated_delta_rule = scan_fn
+ module._unsloth_varlen = None
+ module._unsloth_varlen_wrapped = True
+
+ # Refresh the boundary stash on the outermost forward (once per step, outside
+ # gradient-checkpoint recompute, so it stays valid for recomputed inner
+ # forwards). Read from both positional and keyword args via the bound signature.
+ if not getattr(model, "_unsloth_varlen_forward_wrapped", False):
+ forward_orig = model.forward
+ try:
+ forward_sig = inspect.signature(forward_orig)
+ except (TypeError, ValueError):
+ forward_sig = None
+
+ @wraps(forward_orig)
+ def forward_with_varlen(*args, **kwargs):
+ try:
+ bound = dict(kwargs)
+ if forward_sig is not None and args:
+ bound.update(forward_sig.bind_partial(*args).arguments)
+ varlen = _hybrid_varlen_metadata(bound)
+ except Exception:
+ varlen = None
+ first_pack = varlen is not None and not getattr(
+ model,
+ "_unsloth_varlen_handshake_done",
+ False,
+ )
+ for module in gated_delta_modules:
+ module._unsloth_varlen = varlen
+ if first_pack:
+ module._unsloth_varlen_conv_hit = False
+ module._unsloth_varlen_scan_hit = False
+ out = forward_orig(*args, **kwargs)
+ # Runtime dispatch handshake: on the first packed forward, confirm BOTH
+ # boundary kernels ran for EVERY module. seq_idx (conv) and cu_seqlens
+ # (scan) are both load-bearing, so a partial/absent dispatch (a future
+ # version no longer routing through self.) leaves cross-sequence
+ # contamination. The batch is already flattened with no padded recovery,
+ # so abort before loss/backward rather than train on corrupted data.
+ if first_pack:
+ model._unsloth_varlen_handshake_done = True
+ missing = [
+ type(m).__name__
+ for m in gated_delta_modules
+ if not (
+ getattr(m, "_unsloth_varlen_conv_hit", False)
+ and getattr(m, "_unsloth_varlen_scan_hit", False)
+ )
+ ]
+ if missing:
+ for m in gated_delta_modules:
+ m._unsloth_varlen = None
+ _hybrid_reject("varlen conv/scan not both dispatched (dispatch changed?)")
+ raise RuntimeError(
+ "Unsloth: experimental hybrid packing cannot continue because the "
+ "varlen conv/scan wrappers were not both invoked for "
+ f"{sorted(set(missing))}. Unset UNSLOTH_EXPERIMENTAL_HYBRID_PACKING "
+ "to train these models on the padded path."
+ )
+ return out
+
+ model.forward = forward_with_varlen
+ model._unsloth_varlen_forward_wrapped = True
+ return True
+
+
def get_packed_info_from_kwargs(
kwargs: dict, device: torch.device
) -> Optional[Tuple[torch.Tensor, torch.Tensor, int]]:
From 3ab8dce97a95923b2b4e6741e9df9cba1a9baaca Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 00:58:52 -0700
Subject: [PATCH 029/255] install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL
override CUDA wheel detection (#6692)
* install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL override CUDA wheel detection
get_torch_index_url (and the studio-update mirror _detect_cuda_torch_index_url)
chose the torch wheel family solely by probing the host GPU, with no override.
In a headless / container / CI build the host driver is visible via the
/proc/driver/nvidia/gpus fallback but nvidia-smi cannot report a CUDA version,
so the function fell back to its cu126 default and installed the wrong wheels
(e.g. a cu128 image got cu126 torch).
Add an explicit override checked before any probing, in both the shell installer
and the Python studio-update path:
- UNSLOTH_TORCH_INDEX_URL full index URL, used verbatim (wins)
- UNSLOTH_TORCH_INDEX_FAMILY family (cpu, cu128, rocm6.4, ...) appended to the
mirror base (UNSLOTH_PYTORCH_MIRROR still honoured)
This matches how the published GPU images select CUDA -- vLLM and SGLang take the
CUDA version from an explicit build ARG rather than detecting it, and the Unsloth
Docker base image already pins the cu128 index directly. Desktop installs are
unchanged: with no override set, detection runs exactly as before.
Adds test_get_torch_index_url.sh cases for the override (family, full URL,
precedence, mirror base, trailing-slash strip, empty-ignored).
* install: make the torch-index override authoritative across ROCm paths
Address review feedback on the override added in this PR so a pinned index is
honoured everywhere, not just in get_torch_index_url:
- Skip the WSL ROCm bootstrap (root privilege + large downloads, probes
/dev/dxg) when UNSLOTH_TORCH_INDEX_URL / _FAMILY is set; it previously ran
before the override was consulted.
- Skip the Radeon/Strix rerouting (which re-probes the GPU and overwrites the
resolved URL with repo.radeon.com / repo.amd.com) when the index is pinned, so
an explicit ROCm override (e.g. UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) is kept.
- install_python_stack.py: derive _TORCH_BACKEND from the override when
UNSLOTH_TORCH_BACKEND is unset (standalone studio update), so _ensure_rocm_torch
/ _ensure_cuda_torch repair to the requested family instead of re-detecting.
- Strip ALL leading/trailing slashes in the shell override to match the Python
side (avoids 404s on strict pip proxies).
Adds test cases for double-slash and leading/trailing-slash overrides.
* install: honor pinned torch index in CUDA/ROCm repair paths
Follow-up to the override work in this PR: the get_torch_index_url / install.sh
reroute already respect a pinned UNSLOTH_TORCH_INDEX_URL / _FAMILY, but the
Python repair helpers in install_python_stack.py still re-probed the GPU and
could overwrite the pinned family. Make the pin authoritative there too:
- _ensure_cuda_torch: an explicit cu* pin commits to CUDA wheels, so repair a
ROCm-poisoned venv even when no NVIDIA GPU is visible here (headless /
container / CI cross-install), instead of bailing on the GPU-presence gate.
- _ensure_rocm_torch: skip the AMD per-gfx (Strix) reroute when a ROCm index is
pinned, and in the generic reinstall path install from the pinned URL verbatim
rather than re-detecting the host ROCm version. gfx*/rocm7.2 indexes serve
torch 2.11+, so select the 2.11 package specs for a gfx leaf.
- install.sh: raise the torch constraint to 2.11 for */gfx* indexes too, matching
rocm7.2, so a pinned full-URL/family override that returns early keeps a valid
constraint.
Add _explicit_torch_index_url / _explicit_rocm_torch_index_url helpers and tests
covering the no-GPU CUDA pin repair and the explicit gfx index honored verbatim.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: honor torch-index override on the Windows installers too
The pinned-index work landed for install.sh and install_python_stack.py, but the
Windows installers still picked the wheel index from GPU probing. Extend the same
UNSLOTH_TORCH_INDEX_URL / _FAMILY contract so a pinned index wins on every platform:
- install.ps1: Get-TorchIndexUrl returns the pinned URL/family before nvidia-smi
probing; the AMD ROCm reroute is skipped when the index is pinned, so an explicit
cpu/cu* pin on an AMD host is not overwritten.
- studio/setup.ps1: add shared Get-PinnedTorchIndexUrl / Get-TorchIndexLeaf helpers;
the stale-venv check, the install selection and the AMD reroute all honor the pin,
and the CPU/CUDA install pulls from the resolved index URL.
- tests: parity test that all four installers read both override vars and the two
Windows installers gate the AMD reroute on the pinned flag.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: complete pinned-index handling for ROCm/Windows edge cases
Follow-ups to the override work flagged in review:
- install.ps1: a pinned gfx*/rocm>=7.2 index previously skipped the AMD reroute
that sets the torch>=2.11 floor, so the generic install used torch>=2.4,<2.11
and could resolve the known-bad _grouped_mm wheel. Route a pinned ROCm index
through the ROCm install path with the 2.11 floor + companions, and guard the
companion-spec lookup so a skipped reroute block cannot null-deref.
- studio/setup.ps1: the stale-venv check compared the installed flavor (cuXXX/cpu,
with +rocm misread as cpu) against the raw pinned leaf (gfx1151 / rocm6.4), so a
correct pinned ROCm venv was always marked stale. Classify +rocm wheels as the
generic 'rocm' flavor and normalize a pinned rocm*/gfx* leaf to 'rocm' before
comparing (cu* stays specific so cu126-vs-cu128 still rebuilds).
- install_python_stack.py: _ensure_cuda_torch now also reinstalls from a pinned
CUDA index when the venv carries a CPU wheel (headless CPU-venv-to-CUDA
cross-install via 'studio update'), not only when it finds a ROCm build.
- tests: parity assertions already cover all four installers honoring the override.
* install: finish pinned ROCm/CUDA edge cases on Windows + repair path
Follow-ups to the previous round:
- studio/setup.ps1: a pinned gfx*/rocm>=7.2 index now routes through the ROCm
install path with the 2.11 floor + companions (it previously fell through to the
CUDA branch with bare torch/torchvision/torchaudio against the ROCm index). The
CPU/CUDA fallback index is forced to the CPU wheel index when a ROCm index is
active, so a failed pinned-ROCm install does not retry the ROCm mirror.
- studio/setup.ps1: the stale-venv check no longer treats an unrecognized pinned
URL leaf (e.g. a PEP 503 mirror ending in /simple) as a torch flavor tag, which
was marking a correct venv stale; cu*/cpu/rocm/gfx leaves are still compared.
- install.ps1: the post-failure CPU fallback uses an explicit CPU index instead of
, which for a pinned ROCm index was the ROCm mirror itself (so the
'fallback' just retried the failing index and aborted the installer).
- install_python_stack.py: _ensure_cuda_torch now also reinstalls when the venv's
CUDA family differs from a pinned one (installed cu126 vs pinned cu128), not only
CPU->CUDA; the probe reports the installed cuXXX tag for the comparison.
* install: keep the ROCm to CPU fallback install inside the retry-helper window
The pinned-ROCm CPU fallback computes an explicit CPU index, but the comment
explaining why it cannot reuse $TorchIndexUrl pushed the actual
Invoke-InstallCommandRetry / --force-reinstall call more than 600 chars past the
"ROCm PyTorch install failed" message, so test_pr5940_followups's window check
no longer saw the retry helper. Move the CPU-index computation and its comment
above the failure substep so the retrying force-reinstall stays adjacent to the
message. No behavior change: same explicit CPU index, same retry, same
--force-reinstall.
* install: address #6692 review round 5 (ROCm/CPU pin edge cases)
setup.ps1:
- Stale-venv check: treat an AMD/ROCm host (HasROCm or a resolved gfx arch) with
no explicit pin as expecting "rocm", not "cpu", so a healthy +rocm venv is not
flagged stale (which made installer-managed setup exit and direct update rebuild).
- Pinned-ROCm install failure now routes into the force-reinstall CPU branch:
CuTag stays the rocm/gfx leaf on failure, so the condition also checks
ROCmCpuFallback; otherwise the CUDA branch installed from the CPU index without
--force-reinstall and kept the partial ROCm torch.
- Explicit ROCm pin compare no longer collapses gfx*/rocm* to a generic "rocm":
it compares the +rocmX.Y version (and the torch 2.11 line for gfx pins) so
changing the pinned family (e.g. rocm6.4 -> gfx1151) rebuilds and applies it.
install_python_stack.py:
- _ensure_rocm_torch: an explicit ROCm wheel-index pin now bypasses the
NVIDIA-present / no-AMD-GPU / unreadable-ROCm gates (headless/container/CI
cross-install), mirroring the explicit-CUDA-pin bypass in _ensure_cuda_torch.
- Add _ensure_cpu_torch: an explicit CPU pin (FAMILY=cpu or /cpu URL) now has a
repair path that reinstalls CPU torch over an existing CUDA/ROCm build on a
standalone update (which skips install.sh's flavor enforcement).
install.sh:
- Pin torchvision/torchaudio companions alongside torch for the rocm7.2 / per-gfx
index and the Strix reroute (those AMD indexes publish companions independently
and a bare name can resolve a torch-2.12-built wheel, an ABI mismatch).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* torch-index override: classify CUDA pin by leaf; trim blank shell overrides
_ensure_cuda_torch only overrode the NVIDIA-presence gate for *any* pinned index,
so a non-CUDA mirror URL (or a ROCm/CPU pin) on a non-NVIDIA host with ROCm torch
could force a CUDA reinstall over a working ROCm venv. Add
_explicit_cuda_torch_index_url() (leaf cu*), matching the ROCm/CPU helpers, and
gate on it instead.
install.sh::get_torch_index_url treated a whitespace-only UNSLOTH_TORCH_INDEX_URL
/ _FAMILY as authoritative (yielding an invalid index), unlike the Python .strip()
and PowerShell IsNullOrWhiteSpace paths; trim leading/trailing whitespace first.
* install: honor pinned torch index over CVD/GPU gates and fix leaf-based ROCm classification
- install_python_stack.py: an explicit cu* pin now clears the CUDA_VISIBLE_DEVICES
empty/-1 hide gate as well as the NVIDIA-presence gate, so
CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update repairs to CUDA wheels
(parity with install.sh's get_torch_index_url override, which skips all GPU
probing). Unpinned CVD=-1 still skips.
- install_python_stack.py: _ensure_cpu_torch installs the bounded _CPU_TORCH_PKG_SPEC
instead of a bare torch/torchvision/torchaudio trio; the /cpu index now also
serves torch 2.11+, which is outside the supported <2.11 range.
- install.sh: the torch>=2.11 constraint case matches the index leaf (rocm7.2|gfx*)
instead of the whole URL, so a mirror base path containing a gfx/rocm7.2 segment
with a cu*/cpu family is not false-matched onto the 2.11 line.
- setup.ps1: the stale-venv check expects rocm torch only for arches the install
path maps to a repo.amd.com wheel index; an unmapped/unreadable arch installs
CPU, so a correct CPU venv is no longer marked stale.
- Tests for each of the above.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: tighten pinned torch-index override edge cases
- install.sh: trim whitespace-only UNSLOTH_TORCH_INDEX_URL/_FAMILY before the
_torch_index_pinned guard, matching get_torch_index_url, so a blank override no
longer skips the WSL bootstrap and Radeon/Strix reroutes while detection still
picks the normal index.
- install.sh / install.ps1 / setup.ps1 / install_python_stack.py: force the torch
2.11 floor only for the gfx families with the <2.11 _grouped_mm bug (gfx120X-all,
gfx1151, gfx1150). A pinned override to gfx110X-all/gfx90a/gfx908 stays on the
default range, matching the automatic AMD path.
- install_python_stack.py _ensure_cuda_torch: treat an untagged CUDA build under a
CUDA pin as a family mismatch (reinstall), and match cuXXX pins narrowly (cu +
digits) so a custom/current mirror leaf no longer forces CUDA over a CPU/ROCm venv.
- install_python_stack.py _ensure_rocm_torch: reinstall when an explicit ROCm pin
names a different ROCm family than the already-installed ROCm torch (the ROCm
analogue of the CUDA cuXXX mismatch repair).
Adds tests for each case.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: fix second-order edge cases in pinned torch-index ROCm/CUDA handling
Parse the ROCm torch probe positionally so an empty HIP marker is kept:
CPU/CUDA torch no longer reads as HIP, so the ROCm reinstall is not skipped.
Emit one "|" line (like the CUDA probe) for a robust parse.
Limit the gfx torch 2.11 expectation to the install allowlist
(gfx120X-all/gfx1151/gfx1150). A pinned gfx110X-all/gfx90a/gfx908 index stays
on the default <2.11 specs, so a correct 2.10+rocm wheel is no longer judged a
mismatch and force-reinstalled every update.
Distinguish an AMD per-arch wheel (three-part +rocmA.B.C) from a generic
pytorch.org wheel (two-part +rocmA.B): a gfx per-arch pin over a generic 2.11
wheel now reinstalls the per-arch wheel, while an already-installed per-arch
wheel is not re-flagged (no reinstall loop).
Mirror all of the above in setup.ps1 via new Test-RocmGfx211Leaf /
Test-CudaFamilyLeaf / Get-RocmPinStaleTags helpers, reused by both the
install-spec path and the stale-venv check so they cannot diverge again.
Require a digit after "cu" (^cu[0-9]) in setup.ps1, install.ps1 and install.sh
so a mirror leaf like /custom or /current is not branded CUDA and does not
rebuild the venv every run.
Add tests: CPU/CUDA probe -> has_hip_torch False; gfx110X-all pin + 2.10 wheel
not stale; gfx1151 pin + generic 2.11 wheel stale; gfx1151 pin + per-arch wheel
not stale; /custom and /current not CUDA; plus cross-language allowlist and
cu-digit parity guards, and a PowerShell unit test for the new setup.ps1 helpers.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix ROCm/gfx pin case normalization, ROCm-tag requirement, and CUDA-leaf classification
Normalize torch-index leaves to lowercase before the gfx*/rocm*/cu* allowlist
matches so the canonical gfx120X-all (capital X) gets the torch 2.11 floor in
install.sh (leaf, flavor and repairable helpers). Require an installed +rocm
local tag before a rocmX.Y or non-2.11 gfx pin is judged satisfied in
setup.ps1 Get-RocmPinStaleTags and the Python _rocm_pin_family_mismatch, so an
untagged CPU/CUDA wheel never leaves the pin unapplied. Classify a leaf as CUDA
only via ^cu[0-9]: the Python _TORCH_BACKEND derivation now uses
_is_cuda_family_leaf, and install.sh brands cuda only on cu[0-9]* (unset on an
unknown /current /custom mirror leaf) so the stack probes the GPU instead of
skipping ROCm repair. Add bash, Python and PowerShell tests for capital
gfx120X-all floor, current/custom not-cuda, and untagged-wheel ROCm pins.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: converge torch-index pin detection via a per-venv marker
Introduce a torch-index MARKER that records the exact wheel --index-url used
after each successful torch install, so `unsloth studio update` / repair makes
the "did the pinned index change?" decision by an EXACT string compare rather
than inferring it from the wheel +rocm/+cu version tag. The tag cannot encode
the AMD per-arch gfx family (two 2.11 gfx indexes both install +rocm7.13.0), so
the tag heuristic missed a gfx1151 -> gfx120X-all switch and a custom-URL swap.
Marker path is per-venv (.unsloth-torch-index), one line = the resolved index
URL, written atomically (temp + rename). Path, format and normalization are
shared across all four installers (install.sh, install_python_stack.py,
setup.ps1, install.ps1).
- Reapply gfx pins on a per-arch target change: the marker's exact compare
reinstalls when the pinned index differs, even when both wheels share a tag.
- Honor custom ROCm URL pins during repair: an explicit index whose leaf is not
rocm/gfx/cu/cpu (e.g. simple, current) now reinstalls torch VERBATIM from the
pin when it differs from the marker ("URL wins verbatim").
- Align the KNOWN-2.11 rocm/gfx set to exactly rocm7.2 plus the gfx allowlist
gfx120x-all/gfx1151/gfx1150 in every language; stop treating an unknown newer
rocm (rocm7.3, which does not exist) as the 2.11 line speculatively.
Backward compatible: with no marker (old venvs, torch installed out-of-band) the
existing +rocm/version-tag heuristics still decide, and a matching marker never
reinstall-loops. A cu128 CUDA pin stays a CUDA pin; custom and current leaves are
not CUDA. Adds marker tests (py/sh/ps) plus cross-installer parity checks.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: keep the torch-index marker additive to flavor validation
Three narrow fixes in the marker-based stale-venv detection:
- setup.ps1: a matching marker no longer overwrites the detected installed
flavor. The marker compare is now an additional rebuild trigger, so a stale
wheel (torch swapped to a +cpu build while the marker still records a cuXXX
pin) is still caught by the flavor check instead of being masked as up to date.
- setup.ps1: a supported AMD arch carrying CPU torch is no longer marked stale
and wiped. The downstream AMD Windows ROCm override upgrades CPU torch to ROCm
in place, so wiping first would delete the venv and abort with "Virtual
environment not found". Only a genuinely wrong CUDA wheel still rebuilds.
- install.sh: the Radeon --find-links path records its repo.radeon.com base in
the marker instead of the generic pytorch.org ROCm fallback index, so a later
pin to that generic family correctly reinstalls rather than comparing equal.
Mirrors install.ps1/setup.ps1, which already record the real AMD index.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: honor custom pins and repair pinned venvs in place
Four follow-ups to the torch-index marker work:
- install_python_stack.py: _ensure_cuda_torch/_ensure_rocm_torch now bail when an
explicit custom-index pin names no known torch family, so a verbatim URL override
(a private/simple mirror) is not clobbered by auto-detected CUDA/ROCm wheels
before _ensure_verbatim_torch_index applies it.
- install_python_stack.py: the ROCm marker is additive, not a substitute -- a
matching marker still runs the family/version check so a wheel swapped after the
marker was written is caught. Mirrors setup.ps1.
- setup.ps1: a stale venv under an explicit pin, whose torch still imports, is
repaired in place (force-reinstall torch from the pin in the dependency pass)
instead of wiped. The wipe path only delegates to install.ps1, so on a direct
update it stranded the user at "Virtual environment not found" instead of
applying the new pin. A broken venv or unpinned drift still wipes/delegates.
- install.ps1: when a pinned ROCm install fails over to a CPU base, the marker now
records the CPU index actually used instead of the ROCm pin, so the next managed
setup does not see CPU torch under a ROCm pin and abort as stale.
* setup.ps1: keep the ROCm CPU-fallback force line the pr5940 test guards
5c93ffd4 folded the pin-change force-reinstall into the ROCm CPU-fallback
condition on one line, so the exact literal that test_pr5940_followups.py checks
(if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }) no longer appeared
and the test failed. Split the two conditions into separate if lines: the ROCm
fallback line is restored verbatim and the pin-change force is its own line. Both
still set $cpuForce to the array, so @splat passes one arg.
* install: honor exact CUDA/custom index URL pins in the torch-index marker
Address three Codex review findings on the torch-index marker mechanism:
- install.sh: after the ROCm CPU repair reinstalls torch from the generic
$TORCH_INDEX_URL, record that as the marker source. A Radeon --find-links
install set _TORCH_MARKER_INDEX_URL to its repo.radeon.com base earlier, so
leaving it made the marker misreport Radeon wheels and a later Radeon pin would
compare equal and skip a needed reinstall.
- install_python_stack.py: _ensure_cuda_torch now consults the exact-URL marker
(_marker_pin_mismatch) when the installed +cuXXX tag matches the pinned leaf,
so a same-leaf CUDA mirror change (official cu128 to an internal cu128 mirror)
is reinstalled and re-recorded instead of skipped.
- _normalize_index_url / _normalize_family_leaf (install.sh, setup.ps1,
install_python_stack.py): lowercase only KNOWN wheel-family leaves (rocm/gfx/
cpu/cuXXX) so gfx120X-all still matches gfx120x-all, while a custom
(unknown-family) leaf keeps its case so a verbatim URL pin like /Current does
not compare equal to /current. Tests updated to assert the refined behavior.
* install: fix 3 torch-index marker edge cases (CPU mirror pin, Radeon leaf, migrated venv)
Addresses three review findings on the torch-index override path:
1. CPU index URL change on an already-CPU venv. _ensure_cpu_torch returned
early whenever torch was already a CPU build, so a standalone update that
moved the pin (official /cpu -> a private UNSLOTH_PYTORCH_MIRROR /cpu, same
+cpu tag) never reinstalled. It now consults the exact-URL marker and
reinstalls only when _marker_pin_mismatch reports a different index,
mirroring the CUDA/ROCm same-family handling. A matching marker (or none)
still leaves CPU torch untouched, so there is no reinstall loop.
2. Radeon find-links directory misclassified as a pip ROCm family. A
repo.radeon.com/.../rocm-rel-7.2.1 leaf starts with "rocm" but is a
find-links listing, not a pip --index-url. The old startswith(("rocm",
"gfx")) test routed it into a --index-url reinstall that fails against
find-links. New _is_pip_rocm_family_leaf gates on ^rocm\d / gfx (matching
install.sh's rocm[0-9]* and setup.ps1's ^(rocm[0-9]|gfx)), so a Radeon URL
routes to the verbatim/marker path instead.
3. Migrated venv rewriting its marker to a pin it did not install. install.sh
and install.ps1 write the marker unconditionally, so a migration that
preserves existing torch recorded the newly requested pin and a later
update then found a matching marker and skipped the reinstall the pin
needs (e.g. a per-arch gfx1151 -> gfx120X-all switch, identical +rocm tag).
Both now track _TORCH_INSTALLED_THIS_RUN and write the marker only when
torch was actually installed or repaired this run.
Also add Get-NormalizedFamilyLeaf to the setup.ps1 helper-extraction list in
test_torch_index_marker.ps1 (it was added to setup.ps1 and the shell test in an
earlier round but missed here) and add two unit tests covering findings 1 and 2.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: keep pinned torch repairs on the pinned index
Two fixes for explicit index pins (UNSLOTH_TORCH_INDEX_FAMILY / _URL):
1. install_python_stack.py's repair paths ran uv without clearing the
inherited uv index env vars. uv resolves the default index (--index-url
or --default-index) at the LOWEST priority, so a UV_INDEX or
UV_EXTRA_INDEX_URL mirror in the environment won for any package it
served: a cu128-pinned repair could install torch from the mirror and
then record the cu128 marker it never used. Verified empirically: with
UV_EXTRA_INDEX_URL=.../cu126 exported, uv pip install torch
--index-url .../cu128 resolves torch 2.13.0+cu126. Strip the four uv
index env vars for pinned-index commands only, mirroring the gate
install.sh, install.ps1 and setup.ps1 already have; non-pinned installs
keep the user's mirror.
2. install.ps1 routed any pinned leaf matching rocm* through the ROCm
--default-index path, so a custom find-links leaf like rocm-rel-7.2.1
was treated as a PEP 503 ROCm index and could silently fall back to CPU
torch on resolution failure. Require a digit after rocm, matching
install.sh's rocm[0-9]* and install_python_stack.py's ^rocm\d.
Adds parity + unit tests for both (11 new tests).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: keep pinned repairs off UV_TORCH_BACKEND and narrow setup.ps1's rocm pin match
Round 2 of the pinned-index hardening:
1. _build_uv_cmd converted UV_TORCH_BACKEND into --torch-backend before the
new env isolation could act, and uv's torch backend redirects torch
resolution to its own per-backend index even when --index-url is given
(verified: a cu128-pinned dry run with UV_TORCH_BACKEND=cpu resolves
torch 2.13.0+cpu). Pinned-index commands now never receive the flag and
UV_TORCH_BACKEND joins the stripped env vars, so uv cannot re-read it.
2. setup.ps1's pinned reroute had the same bare rocm* glob install.ps1 had:
a custom find-links leaf like rocm-rel-7.2.1 was routed through the ROCm
--index-url path instead of the verbatim unknown-pin path. Now requires
a digit after rocm, matching install.ps1, install.sh and
_is_pip_rocm_family_leaf.
3. The marker test's case-normalization checks used -eq, which is
case-insensitive in PowerShell, making them vacuous, and the unknown-leaf
expectation was written lowercased while the implementation deliberately
preserves custom-leaf case. Tightened to -ceq with the case-preserving
expected value.
Adds unit + parity tests for 1 and 2 (5 new tests).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: extend the pinned-index guards to every remaining surface
Round 3 of the pinned-index hardening, closing the same holes on the
surfaces the earlier rounds missed:
1. install.sh's pinned-install env scrub now clears UV_TORCH_BACKEND (uv's
torch backend redirects torch resolution to its own per-backend index
even against --default-index), and both PowerShell wrappers clear it in
their pinned-install scrubs, matching install_python_stack.py.
2. setup.ps1's marker stale check still classified any rocm* leaf as a
PyTorch ROCm family while the install selection is digit-gated, so a
custom rocm-current / rocm-rel-7.2.1 pin stale-compared as
not-rocm vs rocm and force-reinstalled on every studio update. The
stale check now uses the same ^rocm\d gate.
3. install_python_stack.py's pinned-command scrub also strips
PIP_EXTRA_INDEX_URL for the pip fallback: pip adds the env extra index
in addition to --index-url, so an inherited mirror could satisfy torch
off the pin while the marker recorded the pinned URL. PIP_INDEX_URL
needs no strip since the explicit --index-url flag overrides it.
Parity + unit tests extended (4 new tests).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: scrub find-links and carry the pinned scrub through pip fallbacks
Round 4 of the pinned-index hardening:
1. UV_FIND_LINKS joins every pinned-install scrub (install.sh, install.ps1,
setup.ps1, install_python_stack.py): uv's --find-links locations can
satisfy torch off the pinned index the same way an extra index does.
2. setup.ps1's Fast-Install restored the scrubbed vars in its finally
BEFORE the pip fallback ran, and never touched the pip env vars at all,
so a failed uv attempt fell back to python -m pip with an inherited
PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS able to win over the pinned
--index-url. The scrub now wraps the whole function (uv attempt + pip
fallback) and includes the pip vars; restore happens after both.
3. install_python_stack.py's scrub also strips PIP_FIND_LINKS for its own
pip fallback, completing the PIP_EXTRA_INDEX_URL fix from round 3.
Parity tests extended (2 new tests).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: digit-gate rocm leaves in marker normalization and ROCm side effects
Round 5 of the pinned-index hardening (three custom-rocm-leaf edge cases):
1. _normalize_family_leaf lowercased every leaf starting with rocm, so a
custom mirror leaf like rocm-Current compared equal to its lowercase form
and a case-only pin change was skipped. URL paths can be case-sensitive.
The rocm prefix is now digit-gated (rocm[0-9]*, matching
_is_pip_rocm_family_leaf) in install.sh, setup.ps1 and
install_python_stack.py, so only true family leaves (rocm7.2) are
lowercased; a custom rocm-* leaf keeps its case.
2. setup.ps1 Test-MarkerPinMismatch compared normalized URLs with -ne, which
is case-insensitive in PowerShell, so a case-only marker change (Simple
vs simple) was treated as matching and the reinstall skipped. Now -cne.
3. install.sh gated the AMD bitsandbytes install and the "repair ROCm torch"
--default-index reinstall on a bare whole-URL rocm glob, so a custom
CPU/CUDA/private index whose leaf merely starts with rocm (rocm-current)
was force-repaired from the wrong ROCm-only path whenever torch.version.hip
was empty. Both now gate on _torch_index_is_rocm_family, computed once from
the digit-gated leaf (rocm[0-9]*/gfx*).
Tests: 4 new parity assertions plus 2 case-sensitivity marker checks.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: apply an explicit custom torch-index pin on the first update
Round 6: an explicitly-set custom (unknown-family) UNSLOTH_TORCH_INDEX_URL
was silently ignored on the first `studio update` of a venv that predates
the marker feature, on both platforms, because the no-marker case was
treated as "do nothing" and the version-tag heuristics cannot judge an
unknown leaf.
1. install_python_stack.py _ensure_verbatim_torch_index now reinstalls
verbatim when the marker is ABSENT (None), not only when it differs, and
short-circuits only when the marker already records this exact pin. It
then writes the marker, so every later update is a no-op. A user who did
not set the override gets pin=None and is untouched, so an out-of-band
torch install is never clobbered.
2. setup.ps1: for an unknown-family pin on a marker-less venv the stale-venv
check now sets PinChangedForceReinstall so the torch block reinstalls in
place from the pin. It deliberately does NOT set shouldRebuild, which
would wipe the venv and strand a direct `studio update`.
3. setup.sh (the Linux `studio update` entry point) skipped
install_python_stack.py entirely when unsloth was already current, so the
marker-driven reinstall (both the verbatim custom pin and the cu/rocm
flavor and family-change repair, e.g. gfx1151 to gfx120X-all) never ran.
It now forces the dependency pass when a torch-index pin env var is set;
the pass is idempotent and no-ops when the marker already matches. This
mirrors setup.ps1's stale-venv pre-check.
Tests: 3 new parity assertions.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* test: expect first-update reinstall for a no-marker custom index pin
Follow-up to d671d8fb2: _ensure_verbatim_torch_index now applies an
explicit unknown-family URL pin verbatim on the first update when the
marker is absent (instead of no-op), so the old
test_verbatim_custom_url_no_marker_is_noop assertion was stale. Rewritten
as test_verbatim_custom_url_no_marker_reinstalls_once: asserts the one
verbatim reinstall from the pinned URL, that the marker is written, and
that a second call with the pin still set is idempotent (no reinstall
loop).
* install: gate the pinned update pass on the marker and record a pin baseline
Round 8, two follow-ups to the round-6 first-update pin fix:
1. setup.sh forced the full dependency pass on EVERY `studio update` while a
torch-index pin stayed exported, even after the marker already recorded the
same pin, turning quick updates into the expensive pass every time. It now
probes install_python_stack.py --torch-pin-needs-apply (which reuses the
exact marker normalization) and forces the pass only when the pin is not yet
applied (marker absent or different); an already-applied persistent pin keeps
the fast path. A probe error fails safe toward running the pass. setup.ps1
gets the same probe in its fast path for parity.
2. A known-family full-URL pin on a venv predating the marker (e.g. an installed
cu128 build and UNSLOTH_TORCH_INDEX_URL pointing at a same-family mirror) left
the marker absent forever: the _ensure_* helpers deliberately do not force a
multi-GB reinstall of identical-family wheels on an old venv, so nothing
recorded the pin and every update re-entered the pass. _record_torch_index_pin_baseline
now records the resolved pin as a baseline after the ensure sequence when the
family already matches and no marker exists, so the pin is tracked (a later
genuine change is detected and applied) and the update loop is broken, without
the redundant reinstall.
Tests: 3 new baseline unit tests, 4 new parity assertions, and the CLI probe.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* setup.sh: keep the pin probe's exit 1 from killing the update under set -e
The --torch-pin-needs-apply probe deliberately exits 1 for the common
steady-state answer (pin already recorded, keep the fast path), but it ran
as a bare command under set -euo pipefail, so the whole studio update
aborted before the exit code was even captured. Absorb the status with
|| _PIN_NEEDS_APPLY=$? and pre-seed 0 so all three outcomes route as
documented: 0 runs the pass, 1 keeps the fast path, anything else fails
safe into the pass. Parity test asserts the guard.
* install: strip pin credentials, disable uv config discovery, bound verbatim installs
Four verified fix groups from a 12-reviewer audit of the torch-index
override feature, each reproduced before fixing:
1. Credential persistence: all four marker writers stored the raw pin URL,
so an authenticated pin (https://user:token@mirror/simple) persisted its
credentials in .unsloth-torch-index (mode 0644 under a default POSIX
umask) and install_python_stack.py printed pin URLs verbatim in repair
messages. Userinfo is now stripped before persisting and in every
log/substep that interpolates a pin, via lockstep helpers
(_strip_index_url_credentials in install.sh / install_python_stack.py,
Remove-IndexUrlCredentials in install.ps1 / setup.ps1). The three
normalizers strip too, so an OLD marker that already carries credentials
still compares equal to the same pin: no reinstall loop on upgrade.
Query strings deliberately stay in the marker; two indexes distinguished
only by query must not compare equal.
2. uv configuration discovery beat the explicit pin: with a discovered
uv.toml declaring torch-backend = "cpu" or a [[index]] entry, uv 0.10.12
resolves torch 2.13.0+cpu against an explicit --index-url/.../cu126 pin;
UV_NO_CONFIG=1 restores +cu126 (reproduced both ways). The pinned-install
scrub in all four installers now sets UV_NO_CONFIG=1 and drops
UV_CONFIG_FILE.
3. The verbatim custom-index update path installed a bare, unconstrained
torch trio while fresh installs from the same unknown-leaf pin apply the
supported range; _ensure_verbatim_torch_index now installs the bounded
trio spec, closing the fresh-vs-update asymmetry.
4. Query-bearing pins (.../cu128?token=x) classified by raw leaf split and
force-reinstalled on every update (the installed cu128 never equals
cu128?token=x). Query/fragment are now stripped before leaf
classification in all four implementations; the marker comparison keeps
the query per (1).
Rejected after verification (no change): the pin-baseline record cannot
produce a wrong later decision (every pin change still mismatches and
reinstalls from the new pin); the venv temp-file symlink scenarios require
an attacker who already owns the environment; pathological inputs like
" / cu128 / " have no realistic caller and fail loudly.
Parity, stack, rocm-support, marker (sh + ps1), pin-stale, index-url and
flavor suites all pass (455 python + full shell/ps1 batteries).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: harden custom-pin repair against clobber, broken torch, and pip config
Four follow-ups to the pinned-index audit fixes:
1. setup.ps1 routed an unknown-leaf custom pin through the CUDA branch with
a bare torch trio while install.ps1 (fresh) and the Python verbatim path
bound the supported range; the pinned unknown-leaf route now applies the
same torch>=2.4,<2.11.0 bound. Known cu* leaves and unpinned runs are
unchanged.
2. The final torch safety pass could not repair a clobbered unknown-family
pin: intermediate dependency steps can pull torch from PyPI (the pass
exists for exactly that reason), but the verbatim helper short-circuited
on marker==pin and no flavor tag exists to probe. The helper now keeps a
per-run snapshot of the installed trio (taken after a verbatim reinstall
or on the first matching-marker pass) and reinstalls from the pin when
the final pass sees the trio drifted. Probe failure skips the
comparison; a reinstall refreshes the snapshot, so no loop.
3. _record_torch_index_pin_baseline could freeze a known-family pin as
applied on a venv whose torch is missing or broken (every family helper
returns without reinstalling when its probe fails), making
--torch-pin-needs-apply report done forever. The baseline now probes the
installed flavor and records only on a match: a cuXXX pin requires the
matching +cuXXX tag, cpu requires a cpu build, rocm/gfx requires hip;
probe failure records nothing.
4. The pinned pip fallback stripped PIP_* env vars but user/site pip config
files still applied (a configured global.extra-index-url can satisfy
torch off the pin). PIP_CONFIG_FILE is now pointed at the null device
for pinned commands (pip loads no config files then), in
_install_env_for_cmd and setup.ps1's Fast-Install pinned scrub.
install.sh / install.ps1 have no pip fallback (uv-only), verified.
Tests: 7 new rocm_support tests (snapshot reset fixture), 1 stack test,
2 parity tests. Full battery green (464 python, sh and ps1 suites).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: complete the pin-repair coverage across the fast path and platforms
Three cross-platform follow-ups to the round-2 pin-repair fixes:
1. The --torch-pin-needs-apply probe only compared marker==pin, so a torch
trio clobbered to the wrong family (a cpu wheel replacing cu128 via a
later pip install) with a still-matching marker reported "already
applied" and the _ensure_{cuda,rocm,cpu} repair never ran on the Linux
fast path. The probe is now a testable _torch_pin_needs_apply() that also
checks the installed flavor against a known-family pin (via a shared
_torch_flavor_matches_pin() helper, so the baseline and the probe cannot
drift). An unknown-family pin has no flavor to validate and a failed
probe cannot prove drift, so both keep the fast path.
2. macOS ARM (real CPU/MPS torch, not NO_TORCH) never applied an unknown-
family custom pin on update: both the verbatim path and the baseline
returned on IS_MACOS while fresh install.sh honors the pin, so the marker
was never written and setup.sh forced the dependency pass on every update
forever. The guards are now IS_MAC_INTEL (Intel mac is already NO_TORCH),
and the final pass applies the pin on macOS ARM.
3. The round-2 final verbatim repair sat in the step-13 sequence guarded
not IS_WINDOWS, so on Windows a dependency step that clobbered torch after
the pin was applied was masked by the matching marker (setup.ps1 does not
re-validate the main venv's torch after calling this script -- verified).
Step 13 now runs the verbatim snapshot-drift repair on Windows and macOS
ARM too; the Linux-oriented cuda/rocm/cpu family helpers stay Linux-only.
Tests: 13 new rocm_support cases (flavor drift, macOS ARM, Windows repair),
parity updates. Full battery green (475 python, sh and ps1 suites).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: strip query tokens from the marker and tighten the pin-drift probe
Four follow-ups to the round-3 pin-repair fixes:
1. The credential stripper feeding the torch-index marker and the logged repair
messages dropped only user:pass@ userinfo, so a private feed that carries its
auth token in the query string (.../simple?token=SECRET) persisted the token
in the world-readable marker (mode 0644 under a default umask) and printed it
in substep output. All four strippers (install.sh, install.ps1,
studio/setup.ps1, install_python_stack.py) now drop the query and fragment
before building the sanitized URL. A query is not part of a PEP 503 index's
identity, so this also stops a rotated token from spuriously mismatching the
marker and forcing a needless reinstall.
2. The --torch-pin-needs-apply fast-path probe accepted an untagged CUDA build
(no +cuXXX local tag) under a specific cuXXX pin, but _ensure_cuda_torch
reinstalls exactly that build to enforce the pin. The probe was more lenient
than the repair, so the repair pass was skipped on the fast path.
_torch_flavor_matches_pin now reports a mismatch for an untagged build under a
cuXXX pin, forcing the pass.
3. The probe's ROCm branch accepted any HIP build for a rocm/gfx pin, while
_ensure_rocm_torch decides a reinstall with the per-arch
_rocm_pin_family_mismatch predicate (a generic +rocm7.2 wheel under a per-arch
gfx pin, or a wrong ROCm version, is a mismatch). The probe now reuses that
predicate, so it is as strict as the repair. This needs the installed torch
version, so _probe_torch_flavor now returns (marker, cutag, version) and
_torch_flavor_matches_pin takes the pin URL (extracting the leaf internally).
4. On Windows a known-family cu*/cpu pin is applied to the main venv by setup.ps1
before install_python_stack.py runs; a later dependency step can clobber it,
and the GPU-aware _ensure_{cuda,cpu}_torch self-skip on Windows while the
verbatim helper handles only unknown-family pins, so nothing repaired the
clobber (setup.ps1 does not re-validate the main venv's torch afterward,
verified). New _ensure_pinned_known_family_torch reinstalls a drifted cu*/cpu
pin in the step-13 Windows/macOS-ARM branch; rocm/gfx per-arch specs stay owned
by setup.ps1, unknown-family by the verbatim helper.
A speculative ROCm 2.11 floor was also raised but is unreachable: the rocm7.2
index publishes no 2.x wheel below 2.11.0, and an unknown newer rocm is not
floored speculatively.
Tests: query/fragment strip cases in the sh + ps1 marker suites and the Python
strip/marker tests; the tri-state helper and the probe/baseline harnesses moved
to the (marker, cutag, version) flavor with matching versions; new probe cases
(untagged CUDA, generic-rocm-under-gfx) and 8 _ensure_pinned_known_family_torch
tests; a four-way query-strip parity assertion. Full battery green (1150 python,
sh 26/26 marker, ps1 marker/flavor/pin-stale).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: reinstall markerless gfx pins and cap custom-index updates at torch 2.11
Two follow-ups from the pin-marker audit:
1. A markerless venv with a gfx per-arch 2.11 pin trusted the wheel version
tag, which is byte-identical (+rocm7.13.0) across gfx120X-all / gfx1151 /
gfx1150. A pre-marker install holding one gfx arch's wheel that is now
pinned to a DIFFERENT gfx index was therefore never switched:
_rocm_pin_family_mismatch returns no-mismatch for any three-part +rocm
2.11 wheel, and _ensure_rocm_torch's absent-marker branch fell through to
that heuristic. _ensure_rocm_torch now forces a one-time reinstall when the
marker is absent AND the pin leaf is a 2.11 gfx per-arch index; the reinstall
writes the marker, so the next update compares exactly and does not loop
(the correctly-pinned no-reinstall guarantee then comes from the exact marker
compare, not the ambiguous tag). Non-gfx-2.11 pins (rocmX.Y, non-2.11 gfx)
stay on the tag heuristic -- their tags are distinguishable.
2. The verbatim custom-index update path used _CUDA_TORCH_PKG_SPEC (torch
<2.12.0) while a FRESH install of the same unknown leaf caps torch at
<2.11.0 (install.sh's default TORCH_CONSTRAINT, and setup.ps1's custom-pin
branch), so a private /simple mirror publishing torch 2.11 could upgrade a
`studio update` to a state the fresh installer never produces. Added
_CUSTOM_INDEX_TORCH_PKG_SPEC (torch>=2.4,<2.11.0), used only by the verbatim
path; companions stay pinned for the same exclusive --index-url ABI reason
as _CUDA_TORCH_PKG_SPEC (a bare name could pull a torch-2.12-built
torchvision). _CUDA_TORCH_PKG_SPEC is unchanged (known-family cu/cpu repair
correctly tracks install.sh's widened cu ceiling).
Tests: 2 new markerless-gfx cases (one-time reinstall + marker write + no-loop
second run, and the rocmX.Y absent-marker no-op), the pre-existing markerless
gfx no-reinstall test flipped to assert the one-time reinstall (it had encoded
the old tag-trusting behavior), and the custom-index bound assertions. 488
passed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: a matching marker must not mask a broken, clobbered, or misclassified torch
Four round-6 follow-ups, all closing cases where a matching torch-index
marker wrongly vouched for a torch that is not actually the pinned one:
1. _is_cuda_family_leaf matched cu+digits by PREFIX (^cu[0-9]), so a custom
mirror leaf like cu128-private classified as CUDA family; the flavor check
then compared the installed cu128 tag to the whole leaf cu128-private and
forced a reinstall on EVERY update (never converging). The cu family is
now matched EXACTLY (re.fullmatch cu[0-9]+), so a cu-suffixed custom leaf
routes through the verbatim/unknown path with a stable marker. Mirrored in
install.sh (_normalize_family_leaf: strip cu, require an all-digit
remainder) and setup.ps1 / install.ps1 (^cu[0-9]+$).
2. _torch_pin_needs_apply returned False on a failed torch probe (missing or
unimportable) under a matching marker, so setup.sh kept the fast path and
a broken torch was never repaired. A failed probe now forces the pass: the
marker cannot vouch for a torch that does not import, forcing is idempotent,
and once torch imports again the probe succeeds and the forcing stops
(self-resolving). Reverses the round-4 conservative choice for this case.
3. _ensure_verbatim_torch_index snapshotted the installed trio on the first
pass with a matching marker and treated an unimportable torch (snapshot
None) as "no drift, skip", so a torch clobbered to a broken state before
the run was masked. A None snapshot now reapplies the pin. A torch
clobbered to a WORKING-but-wrong build under an unknown-family pin remains
undetectable from metadata (no flavor tag; reinstalling every update would
be the loop this avoids) and is documented as a known limitation.
4. The step-13 Windows final repair reran only the verbatim (unknown-family)
and known-family cu*/cpu paths, so a clobbered explicit rocm/gfx pin (the
wheel setup.ps1 installed from AMD's per-arch index) was left in place. The
branch now also runs _ensure_rocm_torch on Windows for an explicit rocm/gfx
pin; it has a Windows path and no-ops when torch already links HIP, so it
only reinstalls a genuinely clobbered ROCm venv (loop-safe).
Tests: the round-4 failed-probe-trusts-marker test flipped to force the pass;
new cases for the cu-suffix no-loop, the broken-torch verbatim reinstall, and
the Windows rocm final-repair structure; item-2 exact-cu parity assertions.
490 passed. sh/ps1 marker + flavor + pin-stale suites all green.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: repair Windows ROCm pins from the pinned URL and honor NO_TORCH
Four round-7 review items, two of them regressions in the round-6 work:
1. _torch_pin_needs_apply ignored UNSLOTH_NO_TORCH. With a torch-index env
var set and no marker, the failed-probe branch forced the dependency pass
on every `studio update`, and the pass (which also honors NO_TORCH) never
installs torch or writes a marker, so nothing could ever stop the forcing.
It now returns False immediately under NO_TORCH: the pin only matters once
torch is actually installed.
2. The step-13 Windows final repair (round-6) restored a clobbered explicit
rocm/gfx pin by calling _ensure_rocm_torch, whose Windows path reinstalls
from the arch AUTO-DETECTED via hipinfo, not from the pin. A user pinning a
different gfx family or a private mirror was restored from the wrong source
(and the wrong marker written), and a headless box was skipped entirely
(the arch probe returns nothing). The repair now goes through
_ensure_pinned_known_family_torch, which reinstalls from the PINNED url with
the same per-arch floor setup.ps1 uses (2.11-line gfx leaves) or a bare trio
(older arches, rocmN mirrors). It is gated on IS_WINDOWS since macOS ARM has
no ROCm, and the existing flavor check keeps it loop-safe (a matching HIP
wheel is left alone).
3. _ensure_verbatim_torch_index's broken-torch check (round-6) used
"_installed_trio_snapshot() is None", but that helper reports a REMOVED torch
as "torch==absent" (a non-None tuple) and a broken import as the stale
on-disk version, so a missing or unimportable torch under a matching marker
was read as "no drift" and skipped. The matching-marker path now confirms
torch health with an import probe (_probe_torch_flavor): a torch that does
not import reapplies the pin, while a healthy torch keeps the snapshot-based
intra-run drift detection.
4. A unit test for _ensure_cpu_torch did not pin NO_TORCH False like its
siblings, so a suite run with UNSLOTH_NO_TORCH=1 in the environment made the
guard return early and the reinstall assertions fail spuriously.
Tests: the round-6 broken-torch verbatim test re-encodes the non-None
"torch==absent" snapshot case (the exact state the old "is None" check missed);
new Windows-ROCm pinned-repair cases (reinstall from the pin, per-arch floor vs
bare spec, matching-wheel no-op, off-Windows no-op); a NO_TORCH fast-path probe
case; the parity test now asserts the Windows final branch does not auto-detect
the ROCm index and that the helper reinstalls from the explicit pin. 494 passed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: floor the rocm7.2 index in the Windows pin repair; isolate marker tests
Three round-8 review items, two of them downstream of the round-7 changes:
1. _ensure_pinned_known_family_torch gave a rocm index leaf a bare
torch/torchvision/torchaudio trio while flooring only gfx* leaves, so a
Windows venv clobbered under an explicit rocm7.2 pin could reinstall an
unbounded or ABI-mismatched trio from that exclusive --index-url. It now
mirrors the spec the initial ROCm paths pin: the rocm7.2 floor for 2.11-line
gfx leaves and rocm leaves that serve torch 2.11, the <2.11 default for
older rocm versions, and a bare trio only for older gfx per-arch leaves
(which publish no floor), matching _ROCM_TORCH_PKG_SPECS / _ensure_rocm_torch.
2. test_verbatim_custom_url_no_marker_reinstalls_once called
_ensure_verbatim_torch_index twice; the second call now hits the
matching-marker health probe, and with pip_install mocked torch never becomes
importable, so in a no-torch environment _probe_torch_flavor returned None and
forced another reinstall, failing the idempotence assertion. The test now pins
a healthy flavor so the idempotence check is about the marker, not ambient
torch.
3. The TestEnsureRocmTorchMarker fixture patched os.environ per test but not
_TORCH_BACKEND, which install_python_stack.py computes once at import from
UNSLOTH_TORCH_BACKEND. A runner starting with a cuda/cpu backend made
_ensure_rocm_torch early-return and skip the mocked repair these tests
exercise. The fixture now neutralizes _TORCH_BACKEND so the marker tests are
independent of the caller's installer-pin environment.
Tests: the Windows floor-spec test now asserts a rocm7.2 mirror pin uses the
rocm7.2 floor (not bare), plus a new rocm7.1 case that must fall back to the
<2.11 default; the marker suite passes under a hostile
UNSLOTH_TORCH_BACKEND=cuda / UNSLOTH_TORCH_INDEX_URL env. 495 passed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: apply same-flavor pin repoints, keep ROCm fallback nonfatal, bound custom companions
Four round-9 review items, two of them regressions in the round-7 pin helper:
1. _ensure_pinned_known_family_torch returned as satisfied whenever the installed
flavor matched the pin, so a same-flavor SOURCE change (one /cpu or /cu128
mirror to another, or a gfx1151 -> gfx120x-all per-arch switch, both carrying
the same wheel tag) was never applied, while _torch_pin_needs_apply kept forcing
the pass on the marker mismatch forever. It now also reinstalls when the marker
records a DIFFERENT index of the same flavor, rewriting the marker so the next
update matches (no loop), exactly as the Linux _ensure_{cuda,cpu}_torch helpers
do. An absent marker on an already-matching venv is still left to the baseline
recorder (no forced reinstall of a correct pre-marker venv).
2. That helper reinstalled a Windows ROCm pin with the FATAL pip_install, so when
setup.ps1 had taken its CPU fallback (the pinned AMD index unavailable), the
final repair re-hit the same missing index and aborted the whole install. The
ROCm reinstall is now nonfatal (pip_install_try): on failure it leaves the CPU
base in place and writes no ROCm marker, so the install completes -- matching
_ensure_rocm_torch's Windows path. cu*/cpu pins stay fatal (authoritative source).
3. install.sh left torchvision/torchaudio bare for a pinned custom/unknown-leaf
index (a private /simple mirror), unlike the Python update path's
_CUSTOM_INDEX_TORCH_PKG_SPEC, so a mirror also exposing newer companion wheels
could resolve a torch-2.12-built torchvision against the capped <2.11 torch. It
now bounds the companions (torchvision>=0.19,<0.26.0 / torchaudio>=2.4,<2.11.0)
for a custom leaf, gated on an empty _expected_torch_flavor_tag so known families
keep their curated bare/floored companions.
4. install.sh's _expected_torch_flavor_tag matched cu[0-9]* by prefix, so a custom
leaf like cu128-private classified as the cu128 family and force-reinstalled a
correct +cu128 wheel on every run. It now requires exact cu+digits (routing the
suffixed leaf to the custom path), matching the Python re.fullmatch(cu[0-9]+) and
PowerShell, and feeding item 3's custom-leaf detection.
Tests: new cases for the same-flavor marker-change reinstall, the nonfatal ROCm
fallback (no marker on failure), the rocm7.2/older-rocm floor selection now split
across the nonfatal path, cu-suffixed custom leaves in test_torch_flavor.sh, and the
custom-leaf companion bounds in test_torch_constraint.sh. 497 python + 143 shell
assertions pass; the marker suite still passes under a hostile
UNSLOTH_TORCH_BACKEND=cuda env.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: bound custom-pin companions on the Windows setup path; isolate pin-probe tests
Two round-10 review items:
1. setup.ps1's custom/unknown-leaf pin branch capped only torch ($cudaTorchSpec)
and still asked the exclusive index for bare torchvision/torchaudio, so a
private mirror that also serves newer companion wheels could install a
torch<2.11 wheel alongside a torchvision>=0.26 / torchaudio>=2.11 built for a
newer torch ABI, after which the marker records the pin as applied. It now
bounds the whole trio (torch>=2.4,<2.11.0 / torchvision>=0.19,<0.26.0 /
torchaudio>=2.4,<2.11.0) for a pinned non-cu-family leaf, matching install.sh,
install.ps1's fresh pinned install, and install_python_stack.py's
_CUSTOM_INDEX_TORCH_PKG_SPEC. This completes the companion-bounds fix across all
three installers; known cu* leaves keep bare specs (the family index bounds them).
2. The _torch_pin_needs_apply probe tests did not pin NO_TORCH False, so a test
process launched with UNSLOTH_NO_TORCH=1 short-circuited the probe (the round-7
guard) and returned False for cases that expect the pass to run. The _needs_apply
helper now patches NO_TORCH (default False) around the call, and the dedicated
no-torch case passes no_torch=True explicitly.
Tests: the cross-platform parity test now asserts setup.ps1 bounds the full trio
(not just torch) for a custom leaf; the pin-probe suite passes under a hostile
UNSLOTH_NO_TORCH=1 environment. setup.ps1 parses clean; 497 python + shell suites
green.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: bound custom rocm-* pins, redact diag tokens, snapshot custom pins before base update
Three round-11 review items, all reproduced before fixing:
1. install.sh's custom-index companion bounds gated on _expected_torch_flavor_tag
returning empty, but that helper returned "rocm" for ANY rocm* leaf, so a custom
mirror whose leaf starts with rocm but is not a pip family (a private rocm-current
mirror, a Radeon find-links rocm-rel-7.2.1) escaped the bounds and installed bare
torchvision/torchaudio. It now digit-gates rocm to rocm[0-9]* (matching the Python
_is_pip_rocm_family_leaf ^rocm\d), so those custom leaves return "" and the <2.11
companion caps apply; real rocm7.2 / gfx per-arch indexes still classify as rocm.
2. _tauri_torch_index_family classified by the raw last path segment, so a pinned URL
carrying auth in the query (.../rocm7.2?token=SECRET) had the token echoed verbatim
into the emitted [TAURI:DIAG] line. It now strips query/fragment before classifying
(mirroring the marker/log credential stripping), so no token reaches the diagnostic
output; as a side effect .../cu128?token=x now classifies as cu128 instead of auto.
3. On studio update, the core package step (a newer unsloth can require a torch the
custom pin does not satisfy, pulling a default PyPI trio) runs BEFORE the step-2b
verbatim check, which then recorded the already-clobbered trio as the baseline for a
matching marker and left the pin unapplied. A new _capture_verbatim_baseline() records
the pre-clobber trio before the core step, so the verbatim pass detects the drift and
reapplies the pin. Captures only for a matching custom pin with importable torch; a
mismatched/absent marker or broken torch is left to _ensure_verbatim_torch_index.
Tests: _expected_torch_flavor_tag rocm-current / rocm-rel cases; _tauri_torch_index_family
token/fragment redaction with a no-leak regression guard; _capture_verbatim_baseline
record/skip cases plus an end-to-end clobber-detection scenario; a structural guard that
the capture runs before the core step. 501 python + shell suites pass; install.sh bash -n
clean, shellcheck unchanged from base.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: match rocm family leaves exactly, enforce the rocm7.2 torch line, repair a broken pinned torch
A pinned index is a pip ROCm --index-url family only when its leaf is an exact
rocm / rocm. (rocm7.2) or a gfx* per-arch leaf. The prior
^rocm[0-9] prefix match also caught suffixed private-mirror leaves (rocm7.2-private,
rocm7-current), routing them through the ROCm/companion-family path instead of the
verbatim pin: the companion bounds were skipped and, on a pre-marker venv with a
compatible +rocm wheel, the pin was never applied. Match the family exactly through one
shared helper at every site:
- install_python_stack.py: _is_pip_rocm_family_leaf (re.fullmatch), plus the two other
loose gates it feeds (_normalize_family_leaf, _torch_flavor_matches_pin).
- install.sh: a new _is_pip_rocm_family_leaf routes _expected_torch_flavor_tag,
_torch_index_repairable, _normalize_family_leaf and the ROCm side-effect gate.
- setup.ps1: a new Test-PipRocmFamilyLeaf routes Get-NormalizedFamilyLeaf and both
pinned reroutes; install.ps1 anchors its reroute regex.
_rocm_pin_family_mismatch (and its setup.ps1 mirror Get-RocmPinStaleTags) compared only
the ROCm version, so a +rocm7.2 wheel whose torch release drifted off the 2.11 line
(2.12/2.13 from an out-of-band upgrade or a custom rocm7.2 mirror) satisfied the family
check while violating _ROCM_TORCH_PKG_SPECS['rocm7.2'] (torch>=2.11,<2.12). Flag it stale
so the repair reinstalls to floor; >=2.11 alone is not enough, so the release is compared
exactly against the 2.11 line for a KNOWN-2.11 rocm pin.
_ensure_pinned_known_family_torch returned on a failed import probe, but
_torch_pin_needs_apply forces the dependency pass on that same failed probe: a broken
torch under a known-family pin was left in place and the pass was forced on every update.
Treat an unimportable torch as drift and reinstall the pinned trio (the spec and marker
derive from the pinned leaf, not the absent flavor); once it lands the probe succeeds and
the fast path returns.
Tests: exact-match cases across test_torch_flavor.sh, test_rocm_support.py,
test_cross_platform_parity.py and the two .ps1 helper suites; the rocm7.2 release-line
and broken-probe-reinstall cases; extraction lists updated for the new helpers.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: anchor the PS pinned-ROCm floor gate and bound install.ps1 custom-pin companions
Round 12 made every family CLASSIFIER exact, but the Windows install-flow floor gate reads
$_pinRocm211 directly from the raw pinned leaf with an unanchored -match '^rocm(\d+)\.(\d+)'
BEFORE any exact classification runs. A suffixed custom leaf (rocm7.2-private) matches that
rocm7.2 prefix, so it takes the 2.11-floor branch and is force-routed through the ROCm
install path before the exact-match elseif can send it to the verbatim install. Anchor the
match ($) in both install.ps1 and setup.ps1 so only an exact rocmX.Y leaf is floored; a
suffixed or newer-suffix leaf falls through to the verbatim path. The Python floor
selection is already exact (dict lookups gated on _is_pip_rocm_family_leaf), so only the two
PS scripts needed this.
install.ps1's custom (non-cu-family) pinned-torch install bounded torch>=2.4,<2.11.0 but
left torchvision/torchaudio bare, so a private mirror serving newer companions could pull a
wheel built for a newer torch ABI while the marker records the pin as applied. Bound both
companions (torchvision>=0.19,<0.26.0 / torchaudio>=2.4,<2.11.0) when the leaf is not a
cu family index (a cu index bounds its own resolution), matching setup.ps1's
Test-CudaFamilyLeaf gate and _CUSTOM_INDEX_TORCH_PKG_SPEC.
Tests: parity guards for the anchored floor gate in both PS scripts and for install.ps1's
bounded custom-pin companions.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: tighten comments in the torch-index-override paths
Collapse the verbose comment and docstring blocks added across the installer
scripts and their tests to fewer, clearer lines without changing behaviour.
Remove a duplicated CUDA-spec comment block. Comments/docstrings only; no code
changes (AST-verified).
* install: repair a broken pinned torch on Linux, strip trailing slash in tauri family, count the final step
_ensure_cuda_torch / _ensure_cpu_torch returned on a failed import probe (torch present but
unimportable). With an explicit CUDA/CPU pin, _torch_pin_needs_apply forces the dependency
pass on that same failed probe, and the base package update does not force-reinstall an
already-installed torch distribution, so the broken torch was left in place and the pass
reran every update without repairing it. Treat a failed probe under a pin as drift and
reinstall from the pinned index (the reinstall rewrites the marker and the next probe
imports, so no loop). This is the Linux counterpart of the known-family repair fix.
_tauri_torch_index_family stripped the query/fragment before classifying but not a trailing
slash, so a token-authenticated pin like .../cu128/?token=x collapsed to .../cu128/ and fell
through the exact-suffix */cu128 and */cpu arms to "auto". Strip a trailing slash too,
mirroring _torch_index_url_leaf.
The Windows / macOS-ARM final torch-repair step (_ensure_pinned_known_family_torch) runs a
progress step that base_total never counted (the final-step increment was gated to Linux),
so _STEP ran one past _TOTAL on those platforms. Add the missing increment.
Tests: broken-probe reinstall for the CUDA (family and URL pins) and CPU paths; trailing
slash / slash+token cases for _tauri_torch_index_family; a full-flow progress-count guard
asserting _STEP == _TOTAL on Windows and Linux.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: tighten comments in the torch-index-override paths
* install: harden the torch-index pin across all four installers
Redact index-URL credentials from captured install logs before they print on
failure. uv/pip failure text embeds the failing --index-url verbatim, so a
user:token@ or ?token= secret could leak into the console. Add a shared
redaction pass (_redact_install_output / Redact-InstallOutput) wired into the
error-output dump in install.sh, install.ps1, setup.ps1 and
install_python_stack.py. Verbose mode still streams live uncaptured output, so
it is intentionally left unredacted (developer opt-in).
Trim trailing slashes on the PATH only for a verbatim UNSLOTH_TORCH_INDEX_URL
override, preserving a ?query/#fragment token. A whole-URL rstrip corrupted a
base64 token ending in "/", and a single-slash strip left .../cu128//
classifying as an empty leaf. Add _trim_index_path_slashes /
Trim-IndexPathSlashes and route the override through it; strip ALL trailing
slashes in the backend-branding leaf classifier so a double slash still yields
the real leaf.
Reject a trailing-dot ROCm leaf (rocm7.) in the bash family validator so it
matches Python re.fullmatch(rocm\d+(?:\.\d+)?) and the PowerShell regex: both the
major and the minor must be non-empty digits, so rocm7. is a custom verbatim pin,
not a pip ROCm family.
Scrub PIP_NO_INDEX and PIP_INDEX_URL for a pinned install in the two installers
that have a plain-pip fallback (install_python_stack.py, setup.ps1):
PIP_NO_INDEX=1 makes the fallback ignore every index including the pinned
--index-url, and PIP_INDEX_URL replaces it. install.sh and install.ps1 install
via uv --default-index (which ignores pip config/env), so they are unaffected.
Add unit tests (bash, Python, PowerShell) and cross-platform parity tests
covering credential redaction, path-only slash trimming, the rocm7. validator,
the double-slash leaf, and the PIP_NO_INDEX/PIP_INDEX_URL scrub.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: redact captured torch-install output and warn on a failed pinned ROCm repair
Close a redaction gap the earlier pass missed: setup.ps1's direct
`Fast-Install ... | Out-String` branches (ROCm from $ROCmIndexUrl, CPU/CUDA from
$TorchInstallIndexUrl, plus the Triton and T5 sub-venv installs) printed the
captured $output verbatim on failure, bypassing Redact-InstallOutput. A private
index carrying userinfo or a ?token= in the pin could leak into Windows Studio
setup logs. Route every `Write-Host $output` through Redact-InstallOutput.
Warn on a failed pinned Windows ROCm reinstall in
_ensure_pinned_known_family_torch: the branch printed "reinstalling from it" then
called pip_install_try, but had no else, so a failure continued silently and left
the user believing the pin was applied while the old CPU/wrong torch survived.
Mirror the auto-ROCm Windows path and warn, telling the user to retry.
* install: redact captured output on the pip fallback and optional-install failure paths
The uv install path already redacted its captured output, but pip_install's pip
fallback runs through run(), which printed result.stdout verbatim on failure, and
_print_optional_install_failure did the same. A pinned --index-url carrying
userinfo or a ?token= could still leak there when uv is unavailable or the pip
fallback also fails. Route both through _redact_install_output. The verbose
pip_install_try path stays raw (developer opt-in), matching the other installers.
* install: split the survive-updates marker subsystem into a follow-up
The torch-index override PR grew a persisted per-venv marker plus repair
machinery (stale-pin detection, verbatim re-apply, update-time reinstall
triggers) that roughly doubled it. That subsystem is orthogonal to the core
feature and is being reworked in a follow-up (versioned/hashed marker,
full-URL pin baseline), so it moves there wholesale instead of shipping
twice.
What this PR still does: UNSLOTH_TORCH_INDEX_URL / UNSLOTH_TORCH_INDEX_FAMILY
pick the torch wheel index at install time in all four installers, with the
exact rocm/gfx/cpu/cu leaf classification, the torch 2.11 floor for the
per-arch AMD indexes, bounded companions for custom leaves, credential
redaction of captured installer output, path-only slash trimming, and the
uv/pip index env scrubs. Flavor-based repair keeps honoring the pin: a wrong
family under an explicit pin still reinstalls from the pinned URL, and
setup.ps1 repairs a pinned stale venv in place instead of wiping it.
What moves to the follow-up: the .unsloth-torch-index marker file and its
writers/readers/normalizers, exact-URL pin-change detection on update
(same-tag gfx switches, custom-mirror repoints), the verbatim trio snapshot
and clobber re-apply, the pin-baseline recorder, and the
--torch-pin-needs-apply fast-path probe in setup.sh / setup.ps1. Their tests
(the marker sh/ps1 suites, the stale-pin suite, and the marker classes in the
rocm/cuda/parity suites) move with them; the removed code is preserved on a
local archive branch to seed that PR.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: re-apply a ROCm pin over an existing HIP wheel via the version tag
The subsystem split left an explicit ROCm/gfx pin unenforced on `studio
update` whenever the venv already imported ANY ROCm torch: the pinned
reinstall lived inside the `elif not has_hip_torch` branch, so a rocm6.4 to
rocm7.2 switch, a gfx1151 pin over a generic +rocm7.2 wheel, or a broken
2.12+rocm7.2 drift never re-applied the pin.
Restore the markerless half of that detection: _rocm_pin_family_mismatch
compares the pinned leaf against the installed wheel tag (exact rocmX.Y
compare, the 2.11 gfx per-arch allowlist, the untagged-wheel rule), the HIP
probe emits "|" again so the installed tag is available,
and _ensure_rocm_torch reinstalls from the pinned URL when the tag mismatches
even though HIP torch is present. setup.ps1 mirrors it: the stale-venv check
routes a pinned rocm/gfx leaf through Get-RocmPinStaleTags instead of
collapsing it to a generic "rocm" flavor, and the existing pinned in-place
repair (no wipe) applies the change.
What still waits for the follow-up marker PR, by design: pin changes the
wheel tag cannot see -- a per-arch switch between two 2.11 gfx indexes
(identical +rocm7.13.0 tag), a custom-mirror URL repoint under the same
family leaf, and unknown-family verbatim pins. Those need the persisted
index record.
Tests restored with the code: the _rocm_pin_family_mismatch table, the five
update-path cases (older-rocm reinstall, gfx-over-pre-2.11 reinstall,
matching-pin no-reinstall, non-2.11 gfx no-reinstall, gfx-over-generic-2.11
reinstall), the "|" probe-format guards, and the AST-extracted
Get-RocmPinStaleTags suite for setup.ps1.
* install: compare major-only rocm pins, redact URL fragments, bound pinned CPU trio
Three review fixes on the restored pin-repair path.
The family classifier accepts a major-only rocm leaf (rocm7), but the
mismatch comparators only parsed rocmX.Y, so a rocm7 pin fell through to the
2.11-line fallback and INVERTED both verdicts: an installed +rocm6.4 wheel
compared as satisfied (pin never re-applied) while a matching +rocm7.2 wheel
compared as stale (reinstall loop). Major-only pins now compare on the major
alone in _rocm_pin_family_mismatch and Get-RocmPinStaleTags: rocm6.x under a
rocm7 pin is a mismatch, any rocm7.x satisfies it, an untagged wheel never
does, and a bare +rocm tag with an unreadable version is accepted (matching
the existing lenient unreadable fallback).
The output redactors scrubbed userinfo and ?query= values but not #fragments,
so a pin like https://mirror/whl/cu128#token=secret leaked the secret in
captured uv/pip failure text -- inconsistent with the URL handling itself,
which already treats fragments as sensitive. All four redactors gain a
URL-anchored fragment rule (anchored so a bare "# comment" line in tool
output is never touched).
setup.ps1's CPU branch installed a bare torch/torchvision/torchaudio trio;
fine for the unpinned host default, but a PINNED cpu index routes through the
same branch and the /cpu index serves newer torch, so a fresh pinned CPU
install could land an unsupported trio that _ensure_cpu_torch then keeps
(it accepts any CPU build). Under a pin the branch now installs the bounded
trio mirroring _CPU_TORCH_PKG_SPEC (torch>=2.4,<2.12.0 and matching
companions); the unpinned path is unchanged.
Tests: major-only rows in the Python mismatch table and the AST-extracted
setup.ps1 suite; fragment + query-plus-fragment + bare-hash-comment cases in
all four redactor suites; a parity check that the pinned CPU trio bounds
exist, are gated on the pin, and mirror the Python repair spec.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install: tighten comments in the torch index override paths
* tests: track the moved pass-through inheritance in the gguf order check
Main moved the llama_extra_args pass-through inheritance out of the
GGUF branch into _resolve_inherited_extra_args, which runs before it,
so the source-order assertion's "if request.llama_extra_args is None"
anchor no longer exists inside the branch and the check failed after
the main merge. The test now asserts the same property in the current
shape: inheritance before the GGUF branch (a carried --no-mmproj still
shapes the hub guard's companion requirement), and marker, hub guard,
unload in order within the branch. Full file passes (32 tests).
* tests: anchor the inheritance order check on the call, not the definition
source.index("_resolve_inherited_extra_args(") matched the function
definition, which always precedes the endpoint, so the ordering
assertion was vacuously true. Anchoring on "= _resolve_inherited_
extra_args(" pins the first call site inside the load endpoint (line
4505), which is the statement whose position relative to the GGUF
branch the test is meant to guard. 32 tests pass.
* tests: align the gguf order test with main
Main fixed the stale ordering assertion in PR 7252; adopting its
version verbatim removes this file from the branch diff entirely and
avoids a conflict on the next main merge. 32 tests pass.
* install: bound the companion constraints to torch's window everywhere
A full platform x vendor validation matrix over this branch surfaced a
real trio mismatch on the cpu/mac paths: torch is capped <2.11 (installs
2.10.0+cpu) but the bare torchaudio companion resolves 2.11.0+cpu,
because torchaudio 2.11 dropped its exact torch pin. Reproduced in a
sandboxed end to end cpu install. torchvision still exact-pins torch and
self-corrected.
The default companion constraints are now bounded to torch's window
(<0.26 / <2.11) and widen together with the cu* torch window (<0.27 /
<2.12), so every leaf resolves a paired trio. Verified with uv dry-runs
on the cpu, cu130, and rocm6.4 leaves (2.10.0/0.25.0/2.10.0,
2.11.0/0.26.0/2.11.0, 2.9.1/0.24.1/2.9.1) and a rerun of the sandboxed
cpu install, which now lands torch 2.10.0+cpu with torchaudio
2.10.0+cpu.
The Strix WSL reroute now also forwards UNSLOTH_TORCH_INDEX_URL and
UNSLOTH_TORCH_INDEX_FAMILY into the rerouted 24.04 distro; dropping
them silently reverted the child install to auto-detection, defeating
the pin this branch introduces.
test_torch_constraint.sh updated: the bounded companions must appear at
the defaults and the custom-leaf block, no bare companion may remain,
and the cu* widen must carry the companions with it.
* install: harden the override path against reroute drift and credential leaks
Review sweep focused on default-path idempotency found no defects on the
unset path; these fixes cover the override path and failure reporting.
install.sh:
- The early WSL Strix Halo distro reroute now honors an explicit index
pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY): the pin is used in the current
distro instead of probing the GPU and re-entering another distribution,
matching the contract of the later Radeon and Strix guards. Whitespace
only values do not gate, in parity with get_torch_index_url.
- Verbose mode now streams installer output through the credential
redactor; it previously bypassed the redaction the quiet path applies.
The exit code survives the pipe via an rc file since the script runs
under plain sh with no pipefail.
- The kept-release fallback warning now strips credentials from the
index URL before printing it.
install.ps1:
- Bounded torchvision and torchaudio next to every capped torch install
(custom pin, ROCm CPU fallback, CUDA flavor repair). torchaudio 2.11
dropped its exact torch pin from the wheel metadata, so a bare
companion beside torch<2.11 can resolve a mismatched 2.11.0 build,
cu family indexes included. Mirrors the install.sh companion bounds.
studio/install_python_stack.py:
- The verbose failure path now redacts index URLs in pip and uv output
before printing, matching every other output site in the file.
All sh, ps1 and python installer test suites pass (the host-defaults
suite has a known pre-existing failure unrelated to this change).
* install: redact verbose Windows installer output and repair the parity tests
Follow-ups to the override-hardening commit, from review:
- install.ps1 Invoke-InstallCommand and setup.ps1 Invoke-SetupCommand now
pipe verbose output through Redact-InstallOutput per record, and the
three verbose Fast-Install torch call sites (ROCm, CPU, CUDA) do the
same: uv and pip echo the pinned index URL, credentials included, in
their errors, and verbose mode previously bypassed the redaction the
quiet paths apply. ForEach-Object and Out-Host leave $LASTEXITCODE
untouched, verified with a native command exiting 7 behind the pipe.
- test_cross_platform_parity.py: the install.ps1 companion-bounds
assertion now matches the implemented behavior (bounds on every index,
no cu-family exemption, since torchaudio 2.11 dropped its exact torch
pin) instead of requiring the removed $_pinCuLeaf gate.
- test_rocm_support.py: the WSL reroute guard test slices the whole
function body to its closing brace instead of a fixed 1200-character
window, which the new pin-gate preamble had outgrown.
428 tests pass across the parity, install stack and rocm support suites;
the sh and ps1 installer suites pass unchanged.
* install: tighten comments in the torch-index and ROCm/CUDA repair paths
* install: digit-gate the gfx family leaf and honor ROCm pins in the Windows repair
Two review follow-ups on the override path:
- The pip ROCm family predicate accepted ANY gfx-prefixed leaf, so a
custom verbatim pin like /gfx-private classified as a ROCm family and
enabled the ROCm-only side effects (AMD bitsandbytes, ROCm torch
repair) on a mirror that may serve CPU/CUDA wheels. gfx now requires a
following digit (gfx90a, gfx1151, gfx120X-all), consistently in
install.sh, install_python_stack.py, install.ps1 (family gate and
expected-flavor classifier) and setup.ps1, matching the strictness the
rocm side already had (rocm7.2-private stays verbatim). The broader
backend BRANDING globs are unchanged on purpose: radeon repo leaves
(rocm-rel-X.Y) must still brand the rocm backend without being
force-repaired as a family.
- The Windows branch of the ROCm torch repair always installed from the
public per-arch index, ignoring an explicit ROCm-family pin: after a
pinned setup.ps1 install failed to a CPU base, the repair retried
repo.amd.com instead of the pinned index. The branch now resolves
_explicit_rocm_torch_index_url() first, uses it as the install index
when set, and mirrors the Linux pin contract by skipping the NVIDIA
and gfx-detection gates a pin is documented to override.
Source-assertion tests updated to the tightened predicate and the new
repair label. 1165 tests pass across the parity, install stack and
studio install suites; the sh and ps1 suites pass; both PowerShell
installers parse clean.
* Remove scratch archives accidentally committed with the comment pass
The temp/ archive copies of installer and test files were working
scratch, not PR content, and inflated the diff by about nine thousand
lines.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
install.ps1 | 153 ++-
install.sh | 343 +++++--
studio/install_python_stack.py | 872 +++++++++++++-----
studio/setup.ps1 | 402 +++++++-
tests/python/test_cross_platform_parity.py | 606 ++++++++++++
tests/python/test_install_python_stack.py | 134 +++
tests/run_all.sh | 1 +
tests/sh/test_get_torch_index_url.sh | 57 ++
tests/sh/test_redact_install_output.sh | 89 ++
tests/sh/test_torch_constraint.sh | 74 ++
tests/sh/test_torch_flavor.sh | 89 +-
tests/studio/install/test_cuda_repair.py | 189 +++-
.../install/test_gpu_detection_followups.py | 131 ++-
tests/studio/install/test_pr5940_followups.py | 2 +-
tests/studio/install/test_rocm_support.py | 419 ++++++++-
tests/studio/test_setup_pin_stale.ps1 | 114 +++
tests/studio/test_torch_flavor.ps1 | 15 +-
.../studio/test_torch_index_pin_hardening.ps1 | 78 ++
18 files changed, 3382 insertions(+), 386 deletions(-)
create mode 100755 tests/sh/test_redact_install_output.sh
create mode 100644 tests/studio/test_setup_pin_stale.ps1
create mode 100644 tests/studio/test_torch_index_pin_hardening.ps1
diff --git a/install.ps1 b/install.ps1
index df49414620..6e059ee0dd 100644
--- a/install.ps1
+++ b/install.ps1
@@ -53,7 +53,8 @@ function Install-UnslothStudio {
param([string]$TorchIndexUrl)
if ($SkipTorch) { return "none" }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" }
- $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
+ # Drop query/fragment first so a token-authenticated pin classifies by family.
+ $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf }
if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf }
return "auto"
@@ -62,7 +63,8 @@ function Install-UnslothStudio {
function Get-TauriGpuBranch {
param([string]$TorchIndexFamily)
if ($SkipTorch) { return "no_torch" }
- if ($TorchIndexFamily -like "cu*") { return "cuda" }
+ # Require a digit after "cu" so /current or /custom isn't branded CUDA (parity ^cu[0-9]).
+ if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" }
if ($TorchIndexFamily -like "rocm*") { return "rocm" }
if ($TorchIndexFamily -eq "cpu") { return "cpu" }
return "unknown"
@@ -467,22 +469,35 @@ function Install-UnslothStudio {
}
}
+ # Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
+ # output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
+ # Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
+ function Redact-InstallOutput {
+ param([string]$Text)
+ if (-not $Text) { return $Text }
+ $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@'
+ $Text = $Text -replace '([?&][^=\s&`]+)=[^\s`]+', '$1='
+ # A #token=... fragment is as sensitive as a query; URL-anchored.
+ return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#'
+ }
+
# Run native commands quietly by default to match install.sh behavior.
# Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
function Invoke-InstallCommand {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
- # Installer-pinned index installs (torch) must beat an inherited uv mirror
- # (#6898): when the command pins an index, clear every uv index env var so
- # it wins, then restore in finally. Other installs keep the user's mirror.
+ # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
+ # for --default-index, clear the uv index env vars (restore in finally) and set
+ # UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10).
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
- foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
+ foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'UV_CONFIG_FILE', 'UV_NO_CONFIG') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
+ $env:UV_NO_CONFIG = '1'
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
@@ -493,17 +508,23 @@ function Install-UnslothStudio {
# Merge stderr into stdout so progress/warning output stays visible
# without flipping $? on successful native commands (PS 5.1 treats
# stderr records as errors that set $? = $false even on exit code 0).
- & $Command 2>&1 | Out-Host
+ # Redact per record: uv echoes index URLs (credentials and all) in
+ # its errors, and verbose mode must not bypass the quiet path's
+ # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched.
+ & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
} else {
$output = & $Command 2>&1 | Out-String
if ($LASTEXITCODE -ne 0) {
- Write-Host $output -ForegroundColor Red
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
}
}
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
- if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
+ if ($savedUvIndex) {
+ Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue
+ foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } }
+ }
}
}
@@ -1960,10 +1981,31 @@ exit 0
# On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint.
if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint }
+ # Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
+ # TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
+ function Trim-IndexPathSlashes {
+ param([string]$Url)
+ $value = $Url.Trim()
+ $idx = $value.IndexOfAny([char[]]@('?', '#'))
+ if ($idx -lt 0) {
+ return $value.TrimEnd('/')
+ }
+ return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx)
+ }
+
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
# Mirrors Get-PytorchCudaTag in setup.ps1.
function Get-TorchIndexUrl {
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
+ # Explicit pin -- skip ALL GPU probing (headless / CI / cross-install).
+ # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended
+ # to the mirror base. Matches install.sh / install_python_stack.py.
+ if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
+ return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL)
+ }
+ if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
+ return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
+ }
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
@@ -1984,6 +2026,25 @@ exit 0
return "$baseUrl/cu126"
}
+ # Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with
+ # _strip_index_url_credentials (install.sh / py / setup.ps1).
+ function Remove-IndexUrlCredentials {
+ param([string]$Url)
+ $sep = $Url.IndexOf('://')
+ if ($sep -lt 0) { return $Url }
+ $scheme = $Url.Substring(0, $sep)
+ $rest = $Url.Substring($sep + 3)
+ # Drop query / fragment (may hold auth tokens).
+ $q = $rest.IndexOfAny([char[]]('?', '#'))
+ if ($q -ge 0) { $rest = $rest.Substring(0, $q) }
+ $slash = $rest.IndexOf('/')
+ $authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest }
+ $at = $authority.LastIndexOf('@')
+ $host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority }
+ if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" }
+ return "${scheme}://${host_}"
+ }
+
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
# torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu,
# matching setup.ps1's stale-venv parse.
@@ -2002,11 +2063,13 @@ exit 0
param([string]$TorchIndexUrl, [string]$ROCmIndexUrl)
if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null }
- $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
+ # Drop query/fragment first so .../cu128?token=x classifies as cu128 (else it reinstalls every run).
+ $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if ($leaf -match '^cu\d+$') { return $leaf }
if ($leaf -eq 'cpu') { return 'cpu' }
if ($leaf -match '^rocm') { return 'rocm' }
- if ($leaf -match '^gfx') { return 'rocm' }
+ # gfx must be followed by a digit (an architecture leaf); gfx-private is custom.
+ if ($leaf -match '^gfx[0-9]') { return 'rocm' }
return $null
}
@@ -2041,6 +2104,10 @@ exit 0
} catch { return $null }
}
+ # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it
+ # (e.g. a deliberate cpu pin on an AMD host).
+ $TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or `
+ (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY))
$TorchIndexUrl = Get-TorchIndexUrl
# ── GPU arch → newest compatible Windows ROCm wheel release ──
@@ -2052,7 +2119,9 @@ exit 0
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
$ROCmIndexUrl = $null
$ROCmTorchFloor = $null
- if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
+ $PinnedRocmVisionSpec = $null
+ $PinnedRocmAudioSpec = $null
+ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
@@ -2102,6 +2171,32 @@ exit 0
}
}
+ # A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below
+ # would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2
+ # indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path.
+ if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) {
+ $_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower()
+ $_pinRocm211 = $false
+ # Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to verbatim.
+ if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') {
+ # Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version.
+ $_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2)
+ }
+ # Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare.
+ $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf
+ if ($_pinGfx211 -or $_pinRocm211) {
+ $ROCmIndexUrl = $TorchIndexUrl
+ $ROCmTorchFloor = "torch>=2.11.0,<2.12.0"
+ $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0"
+ $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
+ substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan"
+ } elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') {
+ # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with
+ # bare specs. Only EXACT rocm/gfx* are families; a suffixed leaf is verbatim.
+ $ROCmIndexUrl = $TorchIndexUrl
+ }
+ }
+
if ($ROCmIndexUrl) {
$TorchIndexFamily = "rocm"
} else {
@@ -2164,8 +2259,8 @@ exit 0
}
if ($_Migrated) {
- # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
- # in the new venv location, while preserving existing torch/CUDA
+ # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
+ # existing torch/CUDA unless the flavor repair below re-lands it.
Write-TauriLog "STEP" "Installing unsloth"
substep "upgrading unsloth in migrated environment..."
if ($SkipTorch) {
@@ -2210,22 +2305,24 @@ exit 0
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
} elseif ($ROCmIndexUrl) {
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
- substep "installing PyTorch from $ROCmIndexUrl..."
+ substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..."
$torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin the companions to match $torchSpec; bare names can resolve an
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
- $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
- $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
+ $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
+ $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
- # Transient AMD-index failure: fall back to a CPU base so the install
- # still completes; Unsloth setup retries ROCm afterwards.
+ # Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries
+ # ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS
+ # the ROCm mirror, so reusing it would just retry it.
+ $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow"
# --force-reinstall: a failed ROCm install can leave an unpinned ROCm
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
# the companions -- a mismatched venv the flavor-repair block won't fix.
- $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
+ $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@@ -2238,8 +2335,14 @@ exit 0
}
} else {
Write-TauriLog "STEP" "Installing PyTorch"
- substep "installing PyTorch ($TorchIndexUrl)..."
- $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
+ substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..."
+ # Bound the companions to the capped torch on EVERY index, cu
+ # families included: torchaudio 2.11 dropped its exact torch pin from
+ # the wheel metadata, so a bare companion next to torch<2.11 can
+ # resolve a mismatched 2.11.0 build. Mirrors install.sh.
+ $_pinVisionSpec = "torchvision>=0.19,<0.26.0"
+ $_pinAudioSpec = "torchaudio>=2.4,<2.11.0"
+ $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@@ -2335,8 +2438,8 @@ exit 0
$rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin companions like the fresh ROCm path (bare names can pull an
# ABI-incompatible torchvision/torchaudio from the per-arch index).
- $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
- $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
+ $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
+ $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
@@ -2347,7 +2450,7 @@ exit 0
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
- $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
+ $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
diff --git a/install.sh b/install.sh
index 7918a2bd23..c02552628f 100755
--- a/install.sh
+++ b/install.sh
@@ -159,18 +159,58 @@ run_maybe_quiet() {
fi
}
+# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
+# strip corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
+_trim_index_path_slashes() {
+ _tips_v="$1"
+ case "$_tips_v" in
+ *[?#]*)
+ _tips_head="${_tips_v%%[?#]*}"
+ _tips_tail="${_tips_v#"$_tips_head"}"
+ ;;
+ *)
+ _tips_head="$_tips_v"
+ _tips_tail=""
+ ;;
+ esac
+ while [ -n "$_tips_head" ] && [ "${_tips_head%/}" != "$_tips_head" ]; do
+ _tips_head="${_tips_head%/}"
+ done
+ printf '%s%s' "$_tips_head" "$_tips_tail"
+}
+
+# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
+# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
+# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
+_redact_install_output() {
+ sed -E \
+ -e 's#(https?://)[^/@[:space:]`]+@#\1@#g' \
+ -e 's#([?&][^=[:space:]&`]+)=[^[:space:]`]+#\1=#g' \
+ -e 's|(https?://[^[:space:]`#]+)#[^[:space:]`]+|\1#|g' \
+ "$@"
+}
+
run_install_cmd() {
_label="$1"
shift
- # Installer-pinned index installs (torch) must beat an inherited uv mirror
- # (#6898): when we pass --default-index, neutralize every uv index env var so
- # the pinned index wins. Other installs keep the user's mirror.
+ # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
+ # for --default-index, neutralize the uv index/backend/config vars (UV_TORCH_BACKEND
+ # redirects torch; UV_NO_CONFIG=1 + dropping UV_CONFIG_FILE stops a uv.toml/pyproject
+ # index outranking the CLI pin, uv 0.10).
case " $* " in
- *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
+ *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS -u UV_CONFIG_FILE UV_NO_CONFIG=1 "$@" ;;
esac
if _is_verbose; then
- "$@" && return 0
- _rc=$?
+ # Stream through the redactor: uv echoes index URLs (credentials and
+ # all) in its errors, and verbose mode previously bypassed the
+ # redaction the quiet path applies. The rc file preserves the
+ # command's exit code across the pipe without relying on pipefail
+ # (this script runs under plain sh).
+ _rcf=$(mktemp)
+ { "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output
+ _rc=$(cat "$_rcf" 2>/dev/null || echo 1)
+ rm -f "$_rcf"
+ [ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
return "$_rc"
fi
@@ -178,7 +218,7 @@ run_install_cmd() {
"$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; }
_rc=$?
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
- cat "$_log" >&2
+ _redact_install_output "$_log" >&2
rm -f "$_log"
return $_rc
}
@@ -257,7 +297,7 @@ _install_bnb_rocm() {
fi
_bnb_rc=$?
if _is_verbose; then
- cat "$_bnb_log" >&2
+ _redact_install_output "$_bnb_log" >&2
fi
rm -f "$_bnb_log"
step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2
@@ -310,6 +350,11 @@ _tauri_torch_index_family() {
return
fi
_diag_url="${1:-}"
+ # Strip query/fragment AND a trailing slash before classifying (like _torch_index_url_leaf):
+ # a token isn't echoed into [TAURI:DIAG], and .../cu128/?token=x still classifies as cu128.
+ _diag_url="${_diag_url%%\?*}"
+ _diag_url="${_diag_url%%#*}"
+ _diag_url="${_diag_url%/}"
case "$_diag_url" in
*/cu118) echo "cu118" ;;
*/cu124) echo "cu124" ;;
@@ -343,7 +388,8 @@ _tauri_gpu_branch() {
return
fi
case "$_diag_family" in
- cu*) echo "cuda" ;;
+ # Require a digit after cu so /current or /custom isn't branded CUDA (parity ^cu[0-9]).
+ cu[0-9]*) echo "cuda" ;;
rocm*)
if [ "$_diag_radeon" = true ]; then
echo "rocm_radeon"
@@ -1575,6 +1621,12 @@ _has_usable_nvidia_gpu() {
# the STUDIO_HOME mkdir/venv so the origin distro is untouched.
_maybe_reroute_strixhalo_to_2404() {
[ "${OS:-}" = "wsl" ] || return 0
+ # An explicit index pin skips every GPU-driven reroute (same contract as
+ # the later Radeon/Strix guard): the pin is honored in THIS distro rather
+ # than probing the GPU and switching distributions. Whitespace-only
+ # overrides do not gate (parity with get_torch_index_url).
+ _rr_pin=$(printf '%s' "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" | tr -d '[:space:]')
+ [ -n "$_rr_pin" ] && return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
@@ -1636,6 +1688,10 @@ _maybe_reroute_strixhalo_to_2404() {
# Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the
# GPU instead of falling back to the desktop-app prompt path.
[ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1"
+ # Forward a pinned torch index into the rerouted distro; dropping it would
+ # silently revert the child install to auto-detection.
+ [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_URL=$(_rr_q "$UNSLOTH_TORCH_INDEX_URL")"
+ [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_FAMILY=$(_rr_q "$UNSLOTH_TORCH_INDEX_FAMILY")"
[ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1"
_rr_args=""
[ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")"
@@ -2001,6 +2057,15 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t
TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
fi
fi
+# Companion (torchvision/torchaudio) constraints, bounded to torch's window.
+# torchaudio 2.11 dropped its exact torch pin, so a bare companion next to a
+# <2.11-capped torch resolves torchaudio 2.11 (verified: cpu leaf installed
+# torch 2.10.0+cpu with torchaudio 2.11.0+cpu). torchvision still exact-pins
+# torch and self-corrects, but is bounded for symmetry. Widened alongside the
+# cu* torch window below; the torch-2.11 AMD paths (rocm7.2 / per-gfx / Strix)
+# pin their own trio.
+TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
+TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
# ── Resolve repo root (for --local installs) ──
_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
@@ -2069,6 +2134,24 @@ _has_amd_rocm_gpu() {
get_torch_index_url() {
_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
_base="${_base%/}"
+ # Explicit override -- skip ALL GPU probing (headless / container / CI / cross-install).
+ # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf (cpu, cu128, ...)
+ # appended to the mirror base. Trim whitespace so a whitespace-only value is unset.
+ _url="${UNSLOTH_TORCH_INDEX_URL:-}"
+ _url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}"
+ if [ -n "$_url" ]; then
+ # Trim trailing PATH slashes (a multi-slash path 404s on strict pip proxies) while
+ # preserving a ?query/#fragment token (a whole-URL strip would eat a "/"-ending token).
+ _url=$(_trim_index_path_slashes "$_url")
+ echo "$_url"; return
+ fi
+ _family="${UNSLOTH_TORCH_INDEX_FAMILY:-}"
+ _family="${_family#"${_family%%[![:space:]]*}"}"; _family="${_family%"${_family##*[![:space:]]}"}"
+ if [ -n "$_family" ]; then
+ while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done
+ while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done
+ echo "$_base/$_family"; return
+ fi
# macOS: always CPU (no CUDA support)
case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
# Try nvidia-smi -- require the binary to actually list a usable GPU.
@@ -2197,6 +2280,45 @@ _torch_flavor_tag() {
esac
}
+# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped first
+# so a token-authenticated pin (.../cu128?token=x) classifies as cu128 (else it reinstalls
+# every update). Classification only. Shared with the py / ps1 leaf extractors.
+_torch_index_url_leaf() {
+ _tl_u="${1%%\?*}"
+ _tl_u="${_tl_u%%#*}"
+ # Strip ALL trailing slashes, not one: .../rocm7.2// must yield rocm7.2, not an empty leaf.
+ while [ -n "$_tl_u" ] && [ "${_tl_u%/}" != "$_tl_u" ]; do
+ _tl_u="${_tl_u%/}"
+ done
+ printf '%s' "${_tl_u##*/}" | tr '[:upper:]' '[:lower:]'
+}
+
+# True (exit 0) when a lowercased leaf is an EXACT pip ROCm family: rocm[.]
+# or a gfx ARCHITECTURE leaf (gfx followed by a digit: gfx90a, gfx1151, gfx120x-all). A leaf
+# that merely starts with rocm/gfx (rocm7.2-private, gfx-private) is a custom verbatim pin.
+# Matches the py / ps1 sides.
+_is_pip_rocm_family_leaf() {
+ case "$1" in
+ gfx[0-9]*) return 0 ;;
+ rocm[0-9]*)
+ # Exact rocm[.]: both major and minor must be non-empty all-digits
+ # (rocm7., rocm7.2.1, rocm7.2-private are all custom pins, not a family).
+ _rocm_rest="${1#rocm}"
+ case "$_rocm_rest" in
+ *.*.*) return 1 ;;
+ *.*)
+ _rocm_minor="${_rocm_rest#*.}"
+ case "${_rocm_rest%%.*}" in "" | *[!0-9]*) return 1 ;; esac
+ case "$_rocm_minor" in "" | *[!0-9]*) return 1 ;; esac
+ ;;
+ *[!0-9]*) return 1 ;;
+ esac
+ return 0
+ ;;
+ *) return 1 ;;
+ esac
+}
+
# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2
# ("torch>=A.B[.C],
# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops.
_expected_torch_flavor_tag() {
- _u="${1%/}"
- _leaf="${_u##*/}"
+ _leaf=$(_torch_index_url_leaf "$1")
case "$_leaf" in
- cu[0-9]*) echo "$_leaf" ;;
- cpu) echo "cpu" ;;
- rocm*|gfx*) echo "rocm" ;;
- *) echo "" ;;
+ cu[0-9]*)
+ # Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom),
+ # else a correct +cu128 wheel is force-reinstalled every run.
+ case "${_leaf#cu}" in
+ *[!0-9]*) echo "" ;;
+ *) echo "$_leaf" ;;
+ esac
+ ;;
+ cpu) echo "cpu" ;;
+ # Exact rocm/gfx families only; a custom rocm*-suffixed leaf -> "" (custom).
+ *)
+ if _is_pip_rocm_family_leaf "$_leaf"; then echo "rocm"; else echo ""; fi
+ ;;
esac
}
@@ -2308,14 +2438,42 @@ _expected_torch_flavor_tag() {
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
- _u="${1%/}"
- _leaf="${_u##*/}"
+ _leaf=$(_torch_index_url_leaf "$1")
case "$_leaf" in
- cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;;
- *) echo "no" ;;
+ cu[0-9]*) echo "yes" ;;
+ # Only EXACT rocm/gfx families resolve via --default-index; a suffixed leaf is verbatim.
+ *)
+ if _is_pip_rocm_family_leaf "$_leaf"; then echo "yes"; else echo "no"; fi
+ ;;
esac
}
+# Remove credentials from a wheel index URL ($1) so an authenticated pin never leaks:
+# drops userinfo AND query/fragment; scheme/host/path stay exact. Shared with py / ps1.
+_strip_index_url_credentials() {
+ _sic_url="$1"
+ case "$_sic_url" in
+ *://*) ;;
+ *) printf '%s' "$_sic_url"; return ;;
+ esac
+ _sic_scheme="${_sic_url%%://*}"
+ _sic_rest="${_sic_url#*://}"
+ # Drop query / fragment (may hold auth tokens).
+ _sic_rest="${_sic_rest%%\?*}"
+ _sic_rest="${_sic_rest%%#*}"
+ _sic_auth="${_sic_rest%%/*}"
+ # Drop user:pass@ userinfo if present.
+ case "$_sic_auth" in
+ *@*) _sic_host="${_sic_auth##*@}" ;;
+ *) _sic_host="$_sic_auth" ;;
+ esac
+ if [ "$_sic_auth" = "$_sic_rest" ]; then
+ printf '%s://%s' "$_sic_scheme" "$_sic_host"
+ else
+ printf '%s://%s/%s' "$_sic_scheme" "$_sic_host" "${_sic_rest#*/}"
+ fi
+}
+
get_radeon_wheel_url() {
# Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing
# contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/,
@@ -2561,7 +2719,19 @@ _maybe_bootstrap_rocm_wsl() {
[ -n "$_rw_tmp" ] && rm -f "$_rw_tmp"
return 0
}
-_maybe_bootstrap_rocm_wsl || true
+# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), honour it
+# everywhere: skip the WSL ROCm bootstrap and the Radeon/Strix reroute below (which would
+# re-probe the GPU and overwrite the pin). Trim whitespace first (parity with
+# get_torch_index_url): a whitespace-only override is unset there, so must not flip this true.
+_torch_index_pinned=false
+_ti_url_trim="${UNSLOTH_TORCH_INDEX_URL:-}"
+_ti_url_trim="${_ti_url_trim#"${_ti_url_trim%%[![:space:]]*}"}"; _ti_url_trim="${_ti_url_trim%"${_ti_url_trim##*[![:space:]]}"}"
+_ti_family_trim="${UNSLOTH_TORCH_INDEX_FAMILY:-}"
+_ti_family_trim="${_ti_family_trim#"${_ti_family_trim%%[![:space:]]*}"}"; _ti_family_trim="${_ti_family_trim%"${_ti_family_trim##*[![:space:]]}"}"
+if [ -n "$_ti_url_trim" ] || [ -n "$_ti_family_trim" ]; then
+ _torch_index_pinned=true
+fi
+[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true
TORCH_INDEX_URL=$(get_torch_index_url)
@@ -2572,29 +2742,74 @@ TORCH_INDEX_URL=$(get_torch_index_url)
# whose base path happens to contain "rocm" or "gfx" must not mislabel a
# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix
# overrides in gfxNNNN/, so the trailing slash is stripped first).
-_torch_index_leaf="${TORCH_INDEX_URL%/}"
+# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD
+# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror
+# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so
+# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in
+# lockstep with the shared _torch_index_url_leaf extractor).
+_torch_index_leaf="${TORCH_INDEX_URL%%\?*}"
+_torch_index_leaf="${_torch_index_leaf%%#*}"
+# Strip ALL trailing slashes, not one: .../cu128// must yield cu128, not an empty leaf.
+while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_index_leaf" ]; do
+ _torch_index_leaf="${_torch_index_leaf%/}"
+done
_torch_index_leaf="${_torch_index_leaf##*/}"
+_torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]')
case "$_torch_index_leaf" in
rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;;
cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;;
- *) export UNSLOTH_TORCH_BACKEND="cuda" ;;
+ cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
+ # Unknown leaf (odd mirror, /current): unset so a stale inherited value can't leak and
+ # the stack probes the GPU.
+ *) unset UNSLOTH_TORCH_BACKEND ;;
esac
-# rocm7.2 and the CUDA cu12x/cu13x indexes now ship torch 2.11.x, so widen the
-# ceiling to <2.12.0 (matches the base image and _CUDA_TORCH_PKG_SPEC in
-# studio/install_python_stack.py). Keep the >=2.4 floor so an older CUDA index
-# (e.g. cu118) still resolves. Match on _torch_index_leaf, not the full URL, so
-# a mirror whose base path contains cu*/rocm7.2 but resolves to a cpu/older-rocm
-# leaf keeps the default <2.11.0.
+# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm* / gfx*), gating the
+# ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf
+# merely STARTING with "rocm" isn't force-repaired from the wrong path.
+if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then
+ _torch_index_is_rocm_family=true
+else
+ _torch_index_is_rocm_family=false
+fi
+
+# rocm7.2 and the per-gfx indexes with the _grouped_mm <2.11 bug (gfx120X-all, gfx1151,
+# gfx1150) ship torch 2.11.0 -- raise the floor (also covers a pinned override that skipped
+# the Strix reroute). Pin the companions too: the per-gfx index publishes them independently
+# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a
+# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced.
case "$_torch_index_leaf" in
- rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
- cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
+ rocm7.2|gfx120x-all|gfx1151|gfx1150)
+ TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
+ TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
+ ;;
+ # CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the ceiling to <2.12.0 (matches
+ # _CUDA_TORCH_PKG_SPEC) and widen the companions with it so the trio stays paired.
+ cu[0-9]*)
+ TORCH_CONSTRAINT="torch>=2.4,<2.12.0"
+ TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"
+ ;;
esac
+# A pinned custom/unknown-leaf index (/simple, /current, /cu128-private) has no curated
+# companion set, so bound torchvision/torchaudio to the same <2.11 range the Python path pins
+# (else a mirror with newer companions resolves a 2.12 ABI-mismatched wheel). Known families
+# keep their curated companions above (_expected_torch_flavor_tag returns "" only for custom).
+if [ "$_torch_index_pinned" = true ] && \
+ [ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then
+ TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
+fi
+
# Auto-detect GPU for AMD ROCm based
# get_torch_index_url must have chosen */rocm*
# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
+# Skipped when the index is pinned: an explicit override must not be rerouted to the
+# Radeon/Strix repos by GPU probing.
_amd_gpu_radeon=false
+if [ "$_torch_index_pinned" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \
@@ -2671,10 +2886,14 @@ case "$TORCH_INDEX_URL" in
done
TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/"
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
+ # Pin companions to 2.11 (per-gfx index publishes them independently).
+ TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
_amd_gpu_radeon=false
fi
;;
esac
+fi # _torch_index_pinned guard (Radeon + Strix reroute)
# Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh
# index above supplies the right flavor for this machine. Evaluated HERE, after every
# index/constraint decision including the Strix reroute, so the window checked is the
@@ -2821,7 +3040,7 @@ case "$TORCH_INDEX_URL" in
if [ "$_amd_gpu_radeon" = true ]; then
substep "wheels: repo.radeon.com (Radeon)"
else
- substep "wheels: $TORCH_INDEX_URL"
+ substep "wheels: $(_strip_index_url_credentials "$TORCH_INDEX_URL")"
fi
;;
esac
@@ -2867,8 +3086,8 @@ for _p in ('torch', 'torchvision', 'torchaudio'):
}
if [ "$_MIGRATED" = true ]; then
- # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
- # in the new venv location, while preserving existing torch/CUDA
+ # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
+ # existing torch/CUDA unless the ROCm repair below fires.
substep "upgrading unsloth in migrated environment..."
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps (current
@@ -2909,18 +3128,14 @@ if [ "$_MIGRATED" = true ]; then
# AMD ROCm: install bitsandbytes even in migrated environments so
# existing ROCm installs gain the AMD bitsandbytes build without a
# fresh reinstall.
- if [ "$SKIP_TORCH" = false ]; then
- case "$TORCH_INDEX_URL" in
- */rocm*|*/gfx*)
- _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
- # Repair ROCm torch if overwritten during migrated install
- _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
- if [ -z "$_has_hip" ]; then
- substep "repairing ROCm torch (overwritten by dependency resolution)..."
- _install_torch_default_index --force-reinstall
- fi
- ;;
- esac
+ if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
+ _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
+ # Repair ROCm torch if overwritten during migrated install
+ _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
+ if [ -z "$_has_hip" ]; then
+ substep "repairing ROCm torch (overwritten by dependency resolution)..."
+ _install_torch_default_index --force-reinstall
+ fi
fi
elif [ -n "$TORCH_INDEX_URL" ]; then
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
@@ -3074,7 +3289,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
- substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
+ substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN"
_install_torch_default_index
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
@@ -3095,7 +3310,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
fi
else
- substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
+ substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN"
_install_torch_default_index
fi
else
@@ -3103,19 +3318,15 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_install_torch_default_index
fi
else
- substep "installing PyTorch ($TORCH_INDEX_URL)..."
+ substep "installing PyTorch ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))..."
_install_torch_default_index
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
# host stays in GGUF-only mode rather than pulling in bitsandbytes,
# which is only useful once torch is present for training.
- if [ "$SKIP_TORCH" = false ]; then
- case "$TORCH_INDEX_URL" in
- */rocm*|*/gfx*)
- _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
- ;;
- esac
+ if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
+ _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
fi
# Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed
tauri_log "STEP" "Installing Unsloth"
@@ -3161,16 +3372,12 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_UNSLOTH_TORCH_OVERRIDES=""
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
- if [ "$SKIP_TORCH" = false ]; then
- case "$TORCH_INDEX_URL" in
- */rocm*|*/gfx*)
- _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
- if [ -z "$_has_hip" ]; then
- substep "repairing ROCm torch (overwritten by dependency resolution)..."
- _install_torch_default_index --force-reinstall
- fi
- ;;
- esac
+ if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
+ _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
+ if [ -z "$_has_hip" ]; then
+ substep "repairing ROCm torch (overwritten by dependency resolution)..."
+ _install_torch_default_index --force-reinstall
+ fi
fi
else
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
@@ -3217,7 +3424,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
- substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
+ substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $(_strip_index_url_credentials "$TORCH_INDEX_URL") --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 95c9356d4a..9921b83543 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -44,11 +44,10 @@ IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64"
IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64"
IS_LINUX = sys.platform.startswith("linux")
-# DiskPart-prompt suppression: amd-smi auto-elevates on Windows, popping a
-# UAC/DiskPart prompt mid-install. This installer only spawns probes and pip/uv
-# (none need elevation), so set __COMPAT_LAYER=RunAsInvoker process-wide -- every
-# amd-smi subprocess then runs un-elevated, no per-call guard needed. setup.ps1
-# keeps per-call guards since it ALSO spawns winget installers that need elevation.
+# amd-smi auto-elevates on Windows (UAC/DiskPart prompt mid-install). This installer
+# only spawns probes and pip/uv (no elevation), so set __COMPAT_LAYER=RunAsInvoker
+# process-wide; amd-smi then runs un-elevated. setup.ps1 keeps per-call guards (it
+# also spawns winget installers that need elevation).
if IS_WINDOWS:
os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker")
# torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64,
@@ -74,6 +73,14 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
(6, 0): "rocm6.0",
}
+# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug).
+# Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare.
+_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"})
+
+# pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (rocm7.2 only today); don't
+# floor an unknown newer rocm speculatively. Match install.sh / setup.ps1 / install.ps1.
+_ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)})
+
# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x).
_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"rocm7.2": (
@@ -81,18 +88,16 @@ _ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
),
- # Default for rocm7.1 and earlier: torch 2.x below 2.11
+ # rocm7.1 and earlier: torch 2.x below 2.11
"_default": (
"torch>=2.4,<2.11.0",
"torchvision>=0.19,<0.26.0",
"torchaudio>=2.4,<2.11.0",
),
}
-# Windows AMD per-arch companion pins for the repo.amd.com index, mirroring the
-# install.ps1 / setup.ps1 floor maps (gfx120X and Strix Halo/Point use the rocm7.2
-# torch 2.11 trio). Pinning the companions keeps AMD's per-arch index -- which
-# publishes each independently -- from resolving an ABI-mismatched one. Unlisted
-# arches have no published floor, so stay bare. Bump with the PS maps at 2.12.x.
+# Windows AMD per-arch companion pins for the repo.amd.com index (mirrors the install.ps1 /
+# setup.ps1 floor maps): pinning stops the per-arch index (each published independently) from
+# resolving an ABI-mismatched companion. Unlisted arches have no floor, so stay bare.
_WINDOWS_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = {
"gfx1201": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
"gfx1200": _ROCM_TORCH_PKG_SPECS["rocm7.2"],
@@ -103,19 +108,79 @@ _PYTORCH_WHL_BASE = (
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
).rstrip("/")
-# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed: its
-# torchao 0.17 cpp kernels load cleanly (0.16 crashes on cu130), and the flash-attn
-# / causal-conv1d / mamba torch2.10 wheels load and pass their upstream suites on
-# 2.11 (see wheel_utils._PREBUILT_WHEEL_TORCH_MM). torchvision/torchaudio are pinned
-# (not bare) because the install uses an exclusive --index-url (no PyPI fallback), so
-# a bare name could resolve one built against a different torch major (e.g. 0.27 for
-# torch 2.12) and fail at runtime with an ABI mismatch.
+
+def _strip_index_url_credentials(url: str) -> str:
+ """Strip userinfo (user:password@) AND query/fragment from a wheel index URL.
+
+ An authenticated pin must not leak credentials in printed output; query/fragment
+ may hold tokens and aren't part of the PEP 503 index identity. Host/path stay
+ exact. MUST match install.sh / setup.ps1 / install.ps1.
+ """
+ scheme, sep, rest = url.partition("://")
+ if not sep:
+ return url
+ rest = rest.split("?", 1)[0].split("#", 1)[0] # drop query / fragment
+ authority, slash, tail = rest.partition("/")
+ host = authority.rpartition("@")[2] # drop user:pass@ userinfo
+ return f"{scheme}://{host}{slash}{tail}"
+
+
+_URL_USERINFO_RE = re.compile(r"(https?://)[^/@\s`]+@")
+_URL_QUERY_VALUE_RE = re.compile(r"([?&][^=\s&`]+)=[^\s`]+")
+# URL-anchored so a bare "#..." (a shell comment in tool output) is never touched.
+_URL_FRAGMENT_RE = re.compile(r"(https?://[^\s`#]+)#[^\s`]+")
+
+
+def _redact_install_output(output: "bytes | str") -> str:
+ """Redact index-URL credentials (userinfo + query values + fragments) from captured
+ installer output before printing. uv/pip failure text embeds the failing --index-url
+ verbatim, which can carry a user:token@, ?token= or #token= secret. MUST match
+ install.sh / setup.ps1 / install.ps1's output sanitizers."""
+ text = output.decode(errors = "replace") if isinstance(output, bytes) else output
+ text = _URL_USERINFO_RE.sub(r"\1@", text)
+ text = _URL_QUERY_VALUE_RE.sub(r"\1=", text)
+ return _URL_FRAGMENT_RE.sub(r"\1#", text)
+
+
+def _trim_index_path_slashes(url: str) -> str:
+ """Trim trailing slashes from the URL PATH only, preserving ?query / #fragment. A
+ whole-URL rstrip("/") corrupts a token that ends in "/" (e.g. base64 ...abc/) and a
+ single-slash strip leaves .../cu128// classifying as an empty leaf. MUST match
+ install.sh / setup.ps1 / install.ps1."""
+ value = url.strip()
+ match = re.fullmatch(r"([^?#]*)([?#].*)?", value)
+ if match is None:
+ return value.rstrip("/")
+ return match.group(1).rstrip("/") + (match.group(2) or "")
+
+
+def _torch_index_leaf(url: str) -> str:
+ """Final URL path segment, lowercased, query/fragment removed first.
+
+ So a token-authenticated pin (.../cu128?token=x) classifies as cu128 (a raw leaf
+ keeps the query, never equals the +cu128 tag, and force-reinstalls every update).
+ CLASSIFICATION only; the install keeps the full URL. MUST match install.sh /
+ setup.ps1 / install.ps1.
+ """
+ path = url.split("?", 1)[0].split("#", 1)[0]
+ return path.rstrip("/").rsplit("/", 1)[-1].lower()
+
+
+# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed (torchao
+# 0.17 cpp loads cleanly, and the flash-attn/causal-conv1d/mamba wheels pass on 2.11).
+# torchvision/torchaudio are pinned (not bare) so the exclusive --index-url can't
+# resolve one built against a different torch major -> ABI mismatch.
_CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
"torch>=2.4,<2.12.0",
"torchvision>=0.19,<0.27.0",
"torchaudio>=2.4,<2.12.0",
)
+# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as CUDA: the
+# /cpu index also serves newer torch, so a bare trio could resolve out of range or ABI-
+# mismatched.
+_CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC
+
# torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch
# mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails
# to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a
@@ -408,9 +473,8 @@ def _detect_rocm_version() -> tuple[int, int] | None:
try:
with open(path) as fh:
parts = fh.read().strip().split("-")[0].split(".")
- # Explicit length guard so we don't rely on the broad except
- # below to swallow IndexError when the version file has a
- # single component (e.g. "6\n" on a partial install).
+ # Explicit length guard: don't rely on the broad except below to
+ # swallow IndexError on a single-component version (e.g. "6\n").
if len(parts) >= 2:
return int(parts[0]), int(parts[1])
except Exception:
@@ -455,11 +519,10 @@ def _detect_rocm_version() -> tuple[int, int] | None:
except Exception:
pass
- # Distro package-manager fallbacks. Package-managed ROCm installs can
- # expose GPUs via rocminfo/amd-smi but lack /opt/rocm/.info/version and
- # hipconfig, so probe dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE)
- # for the rocm-core version. Matches install.sh::get_torch_index_url so
- # `unsloth studio update` behaves like a fresh `curl | sh` install.
+ # Distro package-manager fallbacks: package-managed ROCm can expose GPUs via
+ # rocminfo/amd-smi but lack /opt/rocm/.info/version and hipconfig, so probe
+ # dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE) for the rocm-core version.
+ # Matches install.sh::get_torch_index_url so `studio update` == fresh install.
for cmd in (
["dpkg-query", "-W", "-f=${Version}\n", "rocm-core"],
["rpm", "-q", "--qf", "%{VERSION}\n", "rocm-core"],
@@ -561,11 +624,10 @@ def _detect_windows_gfx_arch() -> str | None:
stderr = subprocess.DEVNULL,
timeout = 10,
)
- # Accept partial output even when hipinfo crashes (e.g. exit code
- # 0xC0000005 / STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): if
- # gcnArchName is present in stdout the device was enumerated before
- # the crash, so the arch is trustworthy. Ignoring it causes a
- # silent CPU PyTorch fallback (issue #6043).
+ # Accept partial output even when hipinfo crashes (e.g. 0xC0000005 /
+ # STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): a gcnArchName in stdout
+ # means the device was enumerated pre-crash, so the arch is trustworthy.
+ # Ignoring it causes a silent CPU PyTorch fallback (issue #6043).
text = result.stdout.decode(errors = "replace")
# findall gets every gcnArchName line so multi-GPU hosts are
# enumerable and HIP_VISIBLE_DEVICES selects correctly.
@@ -706,9 +768,8 @@ def _detect_bnb_rocm_dll_ver() -> str | None:
m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll))
if m:
all_vers.append(m.group(1))
- # Pick the highest numeric suffix so e.g. "713" wins over "72" when both
- # variants are present. Glob order is not guaranteed, so always sort
- # rather than stopping at the first match.
+ # Highest numeric suffix wins (e.g. "713" over "72"); glob order is not
+ # guaranteed, so sort rather than take the first match.
return max(all_vers, key = lambda v: int(v)) if all_vers else None
@@ -825,17 +886,14 @@ def _has_rocm_gpu() -> bool:
if result.returncode == 0 and result.stdout.strip():
if check_fn(result.stdout):
return True
- # sysfs KFD topology fallback (Linux only) -- matches install.sh's
- # runtime-only detection. On minimal package-managed installs (no
- # rocminfo / no amd-smi tools), the kernel exposes AMD GPUs via
- # /sys/class/kfd so `studio update` can still detect and repair.
+ # sysfs KFD topology fallback (Linux only) -- matches install.sh's runtime-only
+ # detection. On minimal package-managed installs (no rocminfo / amd-smi), the
+ # kernel exposes AMD GPUs via /sys/class/kfd so `studio update` can still repair.
#
- # Guard: reject any KFD node whose properties file reports a non-AMD
- # vendor. With the NVIDIA open kernel module (driver 560+), NVIDIA GPUs
- # can register KFD topology nodes with a non-zero gpu_id; those nodes
- # have vendor_id 4318 (0x10DE) rather than the AMD value 4098 (0x1002).
- # Without this check the fallback returns True on NVIDIA-only systems,
- # causing _ensure_rocm_torch to install ROCm wheels on NVIDIA hardware.
+ # Guard: reject any KFD node whose properties file reports a non-AMD vendor. The
+ # NVIDIA open kernel module (driver 560+) registers KFD nodes with a non-zero
+ # gpu_id and vendor_id 4318 (0x10DE), not the AMD 4098 (0x1002); without this
+ # check the fallback returns True on NVIDIA-only hosts, installing ROCm wheels.
if sys.platform != "win32":
try:
kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
@@ -849,12 +907,10 @@ def _has_rocm_gpu() -> bool:
continue
if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node
continue
- # Require AMD vendor_id 4098 (0x1002) in the properties file.
- # KFD properties files exist on every kernel that exposes
- # /sys/class/kfd, so absence of the file means we cannot
- # confirm AMD ownership -- skip the node rather than risk a
- # false positive (e.g. NVIDIA open driver KFD nodes that
- # lack a properties file on some kernel versions).
+ # Require AMD vendor_id 4098 (0x1002). KFD properties files exist
+ # on every kernel exposing /sys/class/kfd, so a missing file means
+ # AMD ownership is unconfirmed -- skip the node rather than risk a
+ # false positive (e.g. NVIDIA open-driver KFD nodes lacking it).
props_path = os.path.join(kfd_nodes, entry, "properties")
try:
with open(props_path) as fh:
@@ -981,13 +1037,10 @@ def _install_bnb_windows_rocm() -> bool:
)
if not _ok:
return False
- # After install: detect the actual ROCm DLL suffix shipped in the wheel and
- # set BNB_ROCM_VERSION so bitsandbytes loads the correct DLL regardless of
- # what torch.version.hip reports. The wheel may ship an older suffix (e.g.
- # "72") while torch reports a newer HIP version (e.g. 7.13); the env var
- # override ensures bitsandbytes does not fail looking for a non-existent DLL.
- # The worker subprocess inherits this env var automatically.
- # Fall back to "72" if detection fails (e.g. install was a no-op / dry-run).
+ # Detect the actual ROCm DLL suffix in the wheel and set BNB_ROCM_VERSION so bnb
+ # loads the right DLL regardless of torch.version.hip (the wheel may ship "72"
+ # while torch reports 7.13). The worker subprocess inherits it; fall back to "72"
+ # if detection fails (e.g. a no-op / dry-run install).
_env_ver = os.environ.get("BNB_ROCM_VERSION")
_env_is_persisted_default = (
os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE
@@ -1002,13 +1055,11 @@ def _install_bnb_windows_rocm() -> bool:
_persist_detected_version = True
if _persist_detected_version:
_persist_bnb_rocm_version(_ver)
- # Make hipInfo.exe (shipped into the venv Scripts dir by the AMD torch
- # wheel) resolvable via PATH for this process and every child python the
- # installer spawns (import checks, precompile): bitsandbytes runs
- # `hipinfo.exe` at import time to detect the GPU arch and logs a scary
- # (harmless) ERROR + WARNING on every import when it is missing. The venv
- # Scripts dir is on PATH only when the venv is activated, which neither
- # Unsloth nor the installer's child processes ever do.
+ # Make hipInfo.exe (shipped into venv Scripts by the AMD torch wheel) resolvable
+ # via PATH for this process and every child python (import checks, precompile):
+ # bitsandbytes runs hipinfo.exe at import to detect the GPU arch and logs a scary
+ # (harmless) ERROR + WARNING when it is missing. Scripts is on PATH only for an
+ # activated venv, which neither Unsloth nor the installer's children ever do.
_scripts_dir = os.path.dirname(sys.executable)
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")) and not shutil.which(
"hipinfo.exe"
@@ -1020,13 +1071,18 @@ def _install_bnb_windows_rocm() -> bool:
def _detect_cuda_torch_index_url() -> str:
"""Return the pytorch.org CUDA wheel index URL for the host's NVIDIA driver.
- Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update`
- repairs to the same wheel family a fresh `curl | sh` install would pick.
- Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the
- legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings.
- Defaults to cu126 when nvidia-smi is missing or the version is unreadable
- (e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback).
+ Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update` repairs
+ to the same wheel family a fresh install would pick. Honours the explicit
+ overrides first (UNSLOTH_TORCH_INDEX_URL / _FAMILY) so a headless / CI install
+ never lets the host GPU decide. Otherwise probes nvidia-smi (parsing both "CUDA
+ Version:" and "CUDA UMD Version:"), defaulting to cu126 when unreadable.
"""
+ _override_url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
+ if _override_url:
+ return _trim_index_path_slashes(_override_url)
+ _override_family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
+ if _override_family:
+ return f"{_PYTORCH_WHL_BASE}/{_override_family.strip('/')}"
exe = shutil.which("nvidia-smi")
if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
exe = "/usr/bin/nvidia-smi"
@@ -1061,6 +1117,157 @@ def _detect_cuda_torch_index_url() -> str:
return f"{_PYTORCH_WHL_BASE}/{tag}"
+def _explicit_torch_index_url() -> "str | None":
+ """The wheel index URL pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY, else None.
+
+ Lets the CUDA/ROCm repair helpers honour the exact pinned family/URL instead
+ of re-probing the GPU. Mirrors install.sh::get_torch_index_url's override.
+ """
+ url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
+ if url:
+ return _trim_index_path_slashes(url)
+ family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
+ if family:
+ return f"{_PYTORCH_WHL_BASE}/{family.strip('/')}"
+ return None
+
+
+def _is_pip_rocm_family_leaf(leaf: str) -> bool:
+ """True when a lowercased leaf names a pip --index-url ROCm family: an EXACT
+ rocm[.] leaf or a gfx leaf. A suffixed leaf (rocm-rel-7.2.1,
+ rocm7.2-private) starts with "rocm" but is a custom pin the verbatim path owns, so
+ match EXACTLY. Mirrors install.sh / setup.ps1.
+ """
+ # gfx must be followed by a digit (gfx90a, gfx1151, gfx120X-all): a gfx-prefixed
+ # custom leaf (gfx-private) is a verbatim pin, like rocm7.2-private.
+ return bool(re.fullmatch(r"rocm\d+(?:\.\d+)?", leaf)) or bool(re.match(r"gfx\d", leaf))
+
+
+def _explicit_rocm_torch_index_url() -> "str | None":
+ """The pinned wheel index URL when it names a pip ROCm family (rocm/gfx*), else None."""
+ url = _explicit_torch_index_url()
+ if url is None:
+ return None
+ return url if _is_pip_rocm_family_leaf(_torch_index_leaf(url)) else None
+
+
+def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool:
+ """True when an explicit ROCm pin names a different ROCm family than the installed
+ ROCm torch, so the pin needs a reinstall. Mirrors setup.ps1's stale-venv comparison;
+ same three pin-leaf cases as _ensure_rocm_torch. A same-family pin is NOT a mismatch.
+ """
+ leaf = _torch_index_leaf(pin_url)
+ # Pinned ROCm version. The family classifier accepts a major-only rocm leaf too,
+ # so parse the minor as optional; a major-only pin compares on the major alone.
+ _pin_rocm = re.match(r"^rocm(\d+)(?:\.(\d+))?", leaf)
+ _pin_major = int(_pin_rocm.group(1)) if _pin_rocm else None
+ _pin_ver = (
+ (int(_pin_rocm.group(1)), int(_pin_rocm.group(2)))
+ if _pin_rocm and _pin_rocm.group(2) is not None
+ else None
+ )
+ # Installed +rocmX.Y version; a THREE-part +rocmA.B.C tag is the AMD per-arch
+ # (repo.amd.com/gfx*) signature vs a two-part pytorch.org wheel.
+ _inst_rocm = re.search(r"\+rocm(\d+)\.(\d+)", installed_ver)
+ _inst_ver = (int(_inst_rocm.group(1)), int(_inst_rocm.group(2))) if _inst_rocm else None
+ _inst_is_perarch = re.search(r"\+rocm\d+\.\d+\.\d+", installed_ver) is not None
+ # A ROCm build MUST carry a +rocm tag; an untagged wheel never satisfies a ROCm pin.
+ _inst_has_rocm = re.search(r"\+rocm", installed_ver) is not None
+ # Installed torch RELEASE (before "+") is 2.11+.
+ _inst_rel = re.match(r"^(\d+)\.(\d+)", installed_ver)
+ _inst_is_211 = (
+ (int(_inst_rel.group(1)), int(_inst_rel.group(2))) >= (2, 11) if _inst_rel else False
+ )
+
+ if leaf.startswith("gfx"):
+ # 2.11-allowlist arches expect the AMD per-arch wheel (three-part +rocmA.B.C,
+ # torch 2.11+); a generic or pre-2.11 build is a mismatch.
+ if leaf in _ROCM_GFX_TORCH211_LEAVES:
+ return not (_inst_is_211 and _inst_is_perarch)
+ # Non-2.11 gfx leaf (<2.11 specs): mismatch on an untagged wheel or torch 2.11+.
+ return (not _inst_has_rocm) or _inst_is_211
+
+ # Major-only rocm pin (rocm7): compare majors only -- a +rocm6.4 wheel under a rocm7
+ # pin is a mismatch, any +rocm7.x wheel satisfies it (there is no pinned minor to
+ # compare, and the 2.11-line fallback below would invert both verdicts).
+ if _pin_major is not None and _pin_ver is None:
+ if _inst_ver is not None:
+ return _inst_ver[0] != _pin_major
+ # Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable
+ # version is accepted (matches the lenient unreadable fallback below).
+ return not _inst_has_rocm
+
+ # rocmX.Y pin. Only KNOWN-2.11 rocm is the 2.11 line (no speculative floor).
+ _pin_is_211 = _pin_ver in _ROCM_KNOWN_TORCH211_VERSIONS if _pin_ver is not None else False
+ if _pin_ver is not None and _inst_ver is not None:
+ # Both readable: exact (major, minor) compare (rocm7.2 pin over +rocm7.13.x ->
+ # mismatch, reinstall the pinned wheel).
+ if _pin_ver != _inst_ver:
+ return True
+ # Same family: a KNOWN-2.11 pin whose release drifted off 2.11 (2.12+rocm7.2)
+ # violates the spec -> reinstall to floor (exact compare, not >=2.11).
+ if _pin_is_211 and _inst_rel is not None:
+ if (int(_inst_rel.group(1)), int(_inst_rel.group(2))) != (2, 11):
+ return True
+ return False
+ # rocm pin, unreadable installed version: compare on the 2.11 line, but an untagged
+ # wheel never satisfies a rocmX.Y pin -> mismatch.
+ if not _inst_has_rocm:
+ return True
+ return _pin_is_211 != _inst_is_211
+
+
+def _explicit_cpu_torch_index_url() -> "str | None":
+ """The pinned wheel index URL when it names the CPU family (leaf == cpu), else None.
+
+ An explicit CPU pin (UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu)
+ is authoritative -- see _ensure_cpu_torch.
+ """
+ url = _explicit_torch_index_url()
+ if url is None:
+ return None
+ return url if _torch_index_leaf(url) == "cpu" else None
+
+
+def _is_cuda_family_leaf(leaf: str) -> bool:
+ """True only for a real CUDA wheel-family leaf: "cu" + digits (cu118, cu128, ...).
+
+ A bare startswith("cu") would match "custom"/"current". The match is EXACT so
+ "cu128-private" is NOT a family leaf and routes to the verbatim path instead.
+ """
+ return re.fullmatch(r"cu[0-9]+", leaf) is not None
+
+
+def _explicit_cuda_torch_index_url() -> "str | None":
+ """The pinned wheel index URL when it names a CUDA family (leaf cuXXX), else None.
+
+ Mirrors _explicit_rocm/cpu_torch_index_url so _ensure_cuda_torch only treats a
+ *CUDA* pin as authority to override the NVIDIA-presence gate (an arbitrary mirror
+ or a ROCm/CPU pin must not force a CUDA reinstall on a non-NVIDIA host).
+ """
+ url = _explicit_torch_index_url()
+ if url is None:
+ return None
+ return url if _is_cuda_family_leaf(_torch_index_leaf(url)) else None
+
+
+def _explicit_unknown_family_torch_index_url() -> "str | None":
+ """The pinned index URL when its leaf names NO known torch family, else None.
+
+ Known = rocm* / gfx* / cpu / cuXXX. Anything else (a private mirror /simple,
+ /current) is UNKNOWN: version-tag heuristics can't judge it, so the family
+ repair helpers must leave it alone (the install applied it verbatim).
+ Matches install.sh / setup.ps1 / install.ps1.
+ """
+ url = _explicit_torch_index_url()
+ if url is None:
+ return None
+ leaf = _torch_index_leaf(url)
+ if _is_pip_rocm_family_leaf(leaf) or leaf == "cpu" or _is_cuda_family_leaf(leaf):
+ return None
+ return url
+
+
def _ensure_cuda_torch() -> None:
"""Repair a venv whose torch is a ROCm build on an NVIDIA host.
@@ -1073,44 +1280,47 @@ def _ensure_cuda_torch() -> None:
Only repairs when torch actually links against HIP/ROCm. Healthy CUDA
torch and deliberate CPU-only torch are left untouched.
"""
- # Respect an explicit backend choice from install.sh: only "" (standalone
- # `studio update`) or "cuda" should ever force CUDA wheels. "rocm"/"cpu"
- # (or any unrecognised value) are deliberate and must not be overridden.
+ # Respect install.sh's backend: only "" (standalone update) or "cuda" force CUDA
+ # wheels; "rocm"/"cpu"/unrecognised are deliberate.
if _TORCH_BACKEND not in ("", "cuda"):
return
- # No CUDA torch on macOS; Windows venv/torch lifecycle is owned by
- # install.ps1 (and the KFD poisoning bug is Linux-only), so skip both.
+ # An explicit unknown-family pin was applied VERBATIM at install time; leave it alone.
+ if _explicit_unknown_family_torch_index_url() is not None:
+ return
+ # No CUDA torch on macOS; Windows torch is owned by install.ps1 (KFD bug is Linux-only).
if IS_MACOS or IS_WINDOWS or NO_TORCH:
return
# Never undo a deliberate ROCm install (setup.ps1 sets this marker).
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
return
- # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for
- # example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card);
- # never force CUDA wheels over that choice.
+ # An explicit CUDA pin (headless / CI cross-install) commits to CUDA wheels and skips ALL
+ # GPU probing, so it clears both the CUDA_VISIBLE_DEVICES hide gate and the NVIDIA gate below.
+ _cuda_pinned = _explicit_cuda_torch_index_url() is not None
+ # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU; never force CUDA
+ # wheels over that unless a CUDA index is pinned.
_cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
- if _cvd is not None and _cvd.strip() in ("", "-1"):
+ if not _cuda_pinned and _cvd is not None and _cvd.strip() in ("", "-1"):
return
- # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu()
- # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent.
- if not _has_usable_nvidia_gpu():
+ # Only NVIDIA hosts carry CUDA torch (the CUDA pin overrides this gate too).
+ if not _cuda_pinned and not _has_usable_nvidia_gpu():
return
- # Classify the installed torch: "hip" (ROCm build -- the poisoning
- # signature), "cuda" (healthy), or "cpu" (deliberate CPU wheel). A
- # non-zero exit means torch is missing or un-importable; the base install
- # step handles that, so leave it alone.
+ # Classify the installed torch: "hip" (ROCm poisoning signature), "cuda" (healthy),
+ # or "cpu". A non-zero exit means torch is missing/un-importable: without a pin the
+ # base install owns it, but a pinned CUDA index reinstalls it below.
try:
probe = subprocess.run(
[
sys.executable,
"-c",
(
- "import torch; "
+ "import torch, re; "
"hip = getattr(torch.version, 'hip', '') or ''; "
"cuda = getattr(torch.version, 'cuda', '') or ''; "
"ver = getattr(torch, '__version__', '').lower(); "
- "print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))"
+ "m = re.search(r'\\+(cu\\d+)', ver); "
+ "marker = 'hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'); "
+ "print(marker + '|' + (m.group(1) if m else ''))"
),
],
stdout = subprocess.PIPE,
@@ -1120,22 +1330,60 @@ def _ensure_cuda_torch() -> None:
except (OSError, subprocess.TimeoutExpired):
return
if probe.returncode != 0:
+ # torch present but can't import. Without a pin the base install owns it; but an
+ # explicit CUDA pin forces this pass (failed probe) and the base update won't
+ # reinstall an already-installed torch, so reinstall from the pin (self-resolving).
+ if not _cuda_pinned:
+ return
+ index_url = _detect_cuda_torch_index_url()
+ _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
+ print(
+ f" torch cannot import but an explicit CUDA index is pinned -- reinstalling "
+ f"CUDA torch from {_strip_index_url_credentials(index_url)}"
+ )
+ pip_install(
+ "CUDA torch repair",
+ "--force-reinstall",
+ "--no-cache-dir",
+ _torch_pkg,
+ _vision_pkg,
+ _audio_pkg,
+ "--index-url",
+ index_url,
+ constrain = False,
+ )
return
- # Take the last non-empty stdout line: stray output from sitecustomize or
- # an import hook must not mask the marker (fail-closed either way).
+ # Last non-empty line: stray sitecustomize/import-hook output must not mask the marker.
_marker_lines = [
line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()
]
- if not _marker_lines or _marker_lines[-1] != "hip":
- return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
+ if not _marker_lines:
+ return
+ _marker, _, _installed_cu = _marker_lines[-1].partition("|")
+ # Reinstall CUDA torch on a ROCm build on an NVIDIA host (poisoning signature), or when a
+ # CUDA index is pinned but the venv has the wrong family (CPU or a different cuXXX). A
+ # healthy match, or a CPU wheel with no CUDA pin, is left alone.
+ _pin = _explicit_torch_index_url()
+ _pin_leaf = _torch_index_leaf(_pin) if _pin else ""
+ _pinned_cuda = _is_cuda_family_leaf(_pin_leaf)
+ if _marker == "hip":
+ _why = "torch is a ROCm build on an NVIDIA host"
+ elif _marker == "cpu" and _pinned_cuda:
+ _why = "torch is a CPU build but an explicit CUDA index is pinned"
+ elif _marker == "cuda" and _pinned_cuda and _installed_cu != _pin_leaf:
+ # Installed cuXXX differs from the pin. An untagged build (empty) counts too:
+ # the family can't be confirmed, so reinstall to enforce it (idempotent).
+ _installed_desc = _installed_cu if _installed_cu else "an untagged CUDA build"
+ _why = f"torch is {_installed_desc} but the pinned CUDA index is {_pin_leaf}"
+ else:
+ return # healthy CUDA torch matching the pin, or a deliberate CPU wheel
index_url = _detect_cuda_torch_index_url()
_torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
print(
- f" torch is a ROCm build on an NVIDIA host -- reinstalling "
- f"CUDA torch from {index_url}\n"
- f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch "
- f"on a mixed AMD+NVIDIA host)"
+ f" {_why} -- reinstalling CUDA torch from {_strip_index_url_credentials(index_url)}\n"
+ f" (set UNSLOTH_TORCH_BACKEND=rocm or cpu to keep a deliberate "
+ f"non-CUDA torch)"
)
pip_install(
"CUDA torch repair",
@@ -1150,6 +1398,90 @@ def _ensure_cuda_torch() -> None:
)
+def _ensure_cpu_torch() -> None:
+ """Reinstall CPU torch when an explicit CPU pin is set but the venv has a GPU build.
+
+ Counterpart to _ensure_cuda/rocm_torch for the explicit-CPU case (those treat a CPU
+ backend as a skip, so a standalone `studio update` would ignore the authoritative CPU
+ pin). Only fires for an EXPLICIT pin.
+ """
+ if NO_TORCH:
+ return
+ pin = _explicit_cpu_torch_index_url()
+ if pin is None:
+ return
+
+ # Classify the installed torch family. A non-zero exit means torch is missing or
+ # un-importable: the explicit CPU pin reinstalls it below.
+ try:
+ probe = subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ (
+ "import torch, re; "
+ "hip = getattr(torch.version, 'hip', '') or ''; "
+ "cuda = getattr(torch.version, 'cuda', '') or ''; "
+ "ver = getattr(torch, '__version__', '').lower(); "
+ "gpu = bool(hip) or 'rocm' in ver or bool(cuda) or bool(re.search(r'\\+cu\\d+', ver)); "
+ "print('gpu' if gpu else 'cpu')"
+ ),
+ ],
+ stdout = subprocess.PIPE,
+ stderr = subprocess.DEVNULL,
+ timeout = 90,
+ )
+ except (OSError, subprocess.TimeoutExpired):
+ return
+ if probe.returncode != 0:
+ # torch present but can't import. The explicit CPU pin forces this pass (failed
+ # probe) and the base update won't reinstall an already-installed torch, so
+ # reinstall from the pin (self-resolving, no loop).
+ _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC
+ print(
+ f" torch cannot import but an explicit CPU index is pinned -- reinstalling "
+ f"CPU torch from {_strip_index_url_credentials(pin)}"
+ )
+ pip_install(
+ "CPU torch repair",
+ "--force-reinstall",
+ "--no-cache-dir",
+ _torch_pkg,
+ _vision_pkg,
+ _audio_pkg,
+ "--index-url",
+ pin,
+ constrain = False,
+ )
+ return
+ _lines = [
+ line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()
+ ]
+ if not _lines:
+ return # unreadable -- the base install step handles a missing torch
+ if _lines[-1] != "gpu":
+ return # already a CPU build
+
+ print(
+ " torch is a GPU build but an explicit CPU index is pinned -- reinstalling "
+ f"CPU torch from {_strip_index_url_credentials(pin)}"
+ )
+ # Pin the supported torch<2.11 family (the /cpu index now serves 2.11+, so a bare
+ # trio could resolve out of range or ABI-mismatched).
+ _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC
+ pip_install(
+ "CPU torch repair",
+ "--force-reinstall",
+ "--no-cache-dir",
+ _torch_pkg,
+ _vision_pkg,
+ _audio_pkg,
+ "--index-url",
+ pin,
+ constrain = False,
+ )
+
+
def _ensure_rocm_torch() -> None:
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
@@ -1160,16 +1492,15 @@ def _ensure_rocm_torch() -> None:
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
- # install.sh sets UNSLOTH_TORCH_BACKEND to the resolved wheel family
- # ("cuda", "rocm", "cpu"). Skip ROCm operations entirely when install.sh
- # already selected a non-ROCm backend -- this is the authoritative signal
- # and avoids re-running GPU detection in a subprocess that may see a
- # different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.).
+ # install.sh's resolved backend is authoritative: skip ROCm when it already chose a
+ # non-ROCm family (avoids re-detecting in a subprocess that may see a different env).
if _TORCH_BACKEND in ("cuda", "cpu"):
return
- # setup.ps1 sets this after installing AMD wheels; skip the probe only when
- # torch is actually importable as ROCm. If the venv was wiped between runs,
- # the stale env-var would suppress a needed reinstall.
+ # An explicit unknown-family pin was applied VERBATIM at install time; leave it alone.
+ if _explicit_unknown_family_torch_index_url() is not None:
+ return
+ # setup.ps1 sets this after installing AMD wheels; skip only when torch is actually
+ # importable as ROCm (a wiped venv leaves a stale env-var that must not suppress it).
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_torch_ok = False
try:
@@ -1193,9 +1524,8 @@ def _ensure_rocm_torch() -> None:
pass
if _torch_ok:
_rocm_windows_torch_installed = True
- # setup.ps1 already installed ROCm torch, but we still need the AMD
- # Windows BNB wheel here -- the PyPI bitsandbytes wheel ships only
- # CUDA DLLs and fails to load on ROCm.
+ # ROCm torch is already installed, but the AMD Windows BNB wheel is still
+ # needed (the PyPI bitsandbytes ships only CUDA DLLs, fails on ROCm).
_install_bnb_windows_rocm()
return
# torch was wiped between runs; fall through to the full install path
@@ -1203,10 +1533,15 @@ def _ensure_rocm_torch() -> None:
return
if IS_WINDOWS:
- if _has_usable_nvidia_gpu():
+ # An explicit ROCm-family pin commits to ROCm wheels regardless of the visible
+ # GPU and overrides the public per-arch index (mirrors the Linux pin handling
+ # below): after a pinned setup.ps1 install fails to CPU, this repair must retry
+ # the PINNED index, not repo.amd.com.
+ _win_rocm_pin = _explicit_rocm_torch_index_url()
+ if _win_rocm_pin is None and _has_usable_nvidia_gpu():
return
gfx_arch = _detect_windows_gfx_arch()
- if not gfx_arch:
+ if not gfx_arch and _win_rocm_pin is None:
return # no AMD GPU visible via hipinfo
# Probe whether torch already links against HIP.
_torch_already_rocm = False
@@ -1231,23 +1566,24 @@ def _ensure_rocm_torch() -> None:
except (OSError, subprocess.TimeoutExpired):
pass
if not _torch_already_rocm:
- index_url = _windows_rocm_index_url(gfx_arch)
+ index_url = _win_rocm_pin or _windows_rocm_index_url(gfx_arch)
if index_url is None:
print(f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping")
return
- print(f" {gfx_arch} (Windows) -- installing torch from {index_url}")
- # Pin companions for the arches install.ps1/setup.ps1 pin (gfx120X /
- # Strix) so the per-arch index resolves an ABI-consistent trio; other
- # arches stay bare (no published floor), matching the PowerShell side.
+ print(
+ f" {gfx_arch or 'pinned ROCm index'} (Windows) -- installing torch from "
+ f"{_strip_index_url_credentials(index_url)}"
+ )
+ # Pin companions for the arches install.ps1/setup.ps1 pin (gfx120X / Strix)
+ # so the per-arch index resolves an ABI-consistent trio; other arches stay bare.
_torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get(
gfx_arch, ("torch", "torchvision", "torchaudio")
)
- # Nonfatal: a transient AMD-index failure must not abort the whole
- # install once the PowerShell side has fallen back to CPU torch.
- # --force-reinstall resolves before uninstalling, so a failed index
- # leaves the existing build intact; keep it and let the user retry.
+ # Nonfatal: a transient AMD-index failure must not abort the install.
+ # --force-reinstall resolves before uninstalling, so a failed index keeps the
+ # existing build intact; let the user retry.
if not pip_install_try(
- f"ROCm torch (Windows, {gfx_arch})",
+ f"ROCm torch (Windows, {gfx_arch or 'pinned'})",
"--force-reinstall",
"--index-url",
index_url,
@@ -1257,7 +1593,7 @@ def _ensure_rocm_torch() -> None:
constrain = False,
):
print(
- f" Warning: AMD Windows ROCm torch install failed for {gfx_arch}; "
+ f" Warning: AMD Windows ROCm torch install failed for {gfx_arch or 'the pinned index'}; "
"keeping the existing torch build. Re-run 'unsloth studio update' "
"later to retry ROCm."
)
@@ -1280,26 +1616,30 @@ def _ensure_rocm_torch() -> None:
# ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ──
if platform.machine().lower() not in {"x86_64", "amd64"}:
return
- # NVIDIA takes precedence on mixed hosts -- but only if a GPU is usable
- if _has_usable_nvidia_gpu():
- return
- # Use _has_rocm_gpu() (rocminfo / amd-smi GPU data rows) as the
- # authoritative "is this an AMD ROCm host?" signal. The old gate required
- # /opt/rocm or hipcc to exist, which breaks runtime-only ROCm installs
- # (minimal package-managed installs, Radeon software) that ship
- # amd-smi/rocminfo without /opt/rocm or hipcc, leaving `unsloth studio
- # update` unable to repair a CPU-only venv on those systems.
- if not _has_rocm_gpu():
- return # no AMD GPU visible
+ # An explicit ROCm pin commits to ROCm wheels regardless of the visible GPU (headless / CI).
+ # Mirror _ensure_cuda_torch: skip the NVIDIA/no-AMD/unreadable gates.
+ _rocm_pin = _explicit_rocm_torch_index_url()
+ if _rocm_pin is None:
+ # NVIDIA takes precedence on mixed hosts (only if a GPU is usable).
+ if _has_usable_nvidia_gpu():
+ return
+ # _has_rocm_gpu() (rocminfo / amd-smi rows) is the authoritative AMD-host signal;
+ # the old /opt/rocm-or-hipcc gate broke runtime-only ROCm installs.
+ if not _has_rocm_gpu():
+ return # no AMD GPU visible
ver = _detect_rocm_version()
if ver is None:
- print(" ROCm detected but version unreadable -- skipping torch reinstall")
- return
+ if _rocm_pin is None:
+ print(" ROCm detected but version unreadable -- skipping torch reinstall")
+ return
+ # Explicit pin: the pinned leaf drives the install, so an unreadable host version
+ # is fine (sentinel keeps ver comparisons defined).
+ ver = (0, 0)
- # Probe whether torch already links against HIP (ROCm already working).
- # Do NOT skip for CUDA-only builds: they are unusable on AMD-only hosts
- # (the NVIDIA check above already handled mixed AMD+NVIDIA setups).
+ # Probe whether torch links against HIP, capturing the installed ROCm tag for pin-mismatch
+ # detection. Emit ONE "|" line: marker (HIP version, "rocm" sentinel,
+ # or empty for CPU/CUDA) before "|", wheel version after.
try:
probe = subprocess.run(
[
@@ -1309,10 +1649,10 @@ def _ensure_rocm_torch() -> None:
"import torch; "
"hip=getattr(torch.version,'hip','') or ''; "
"ver=getattr(torch,'__version__','').lower(); "
- # Print the HIP version when present (back-compat), else a
- # "rocm" sentinel when only torch.__version__ flags ROCm
- # (AMD SDK / Radeon wheels). Empty string = CPU/CUDA.
- "print(hip if hip else ('rocm' if 'rocm' in ver else ''))"
+ # HIP version if present, else a "rocm" sentinel when only the
+ # version string flags ROCm; empty marker = CPU/CUDA torch.
+ "marker=hip if hip else ('rocm' if 'rocm' in ver else ''); "
+ "print(marker + '|' + ver)"
),
],
stdout = subprocess.PIPE,
@@ -1321,29 +1661,42 @@ def _ensure_rocm_torch() -> None:
)
except (OSError, subprocess.TimeoutExpired):
probe = None
- has_hip_torch = (
- probe is not None and probe.returncode == 0 and probe.stdout.decode().strip() != ""
+ # Last non-empty line, split on the FIRST "|" so the empty HIP field is preserved.
+ _marker_lines = (
+ [ln.strip() for ln in probe.stdout.decode(errors = "replace").splitlines() if ln.strip()]
+ if (probe is not None and probe.returncode == 0)
+ else []
+ )
+ _hip_marker, _sep, _installed_torch_ver = (
+ _marker_lines[-1].partition("|") if _marker_lines else ("", "", "")
+ )
+ # A "|"-delimited line is required; without it treat HIP as absent -> reinstall.
+ has_hip_torch = bool(_sep) and _hip_marker != ""
+
+ # An explicit ROCm pin whose family differs from the installed torch must reinstall, else a
+ # rocm7.2/gfx* pin over an older +rocm6.4/7.1 build never applies. Version-tag heuristic
+ # only: a same-tag per-arch switch (gfx1151 -> gfx120X-all, both +rocm7.13.0) isn't detectable.
+ _rocm_pin_mismatch = (
+ _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver)
+ if (has_hip_torch and _rocm_pin is not None)
+ else False
)
- rocm_torch_ready = has_hip_torch
+ rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
- # Strix Halo / Strix Point (gfx1151 / gfx1150) segfault under ROCm 7.1
- # in torch._grouped_mm. AMD's per-gfx repo ships torch 2.11.0+rocm7.13.0
- # with the real fix, so route those hosts there instead of the generic
- # pytorch.org rocm7.1 wheel. Mirrors install.sh's Strix override.
- # On mixed hosts (Strix iGPU + non-Strix dGPU), route to the AMD per-gfx
- # index only when HIP's runtime GPU is the Strix one -- else the dGPU gets
- # an incompatible wheel. Use HIP_VISIBLE_DEVICES for the runtime target.
+ # Strix Halo / Point (gfx1151 / gfx1150) segfault under ROCm 7.1 in torch._grouped_mm;
+ # AMD's per-gfx repo ships 2.11.0+rocm7.13.0 with the fix, so route those hosts there
+ # (mirrors install.sh). On mixed hosts, reroute only when HIP's runtime GPU is the Strix one.
_strix_override_url: "str | None" = None
_strix_override_pkgs: "tuple[str, str, str] | None" = None
- if ver < (7, 2):
+ # An explicit ROCm pin is authoritative: never auto-reroute it.
+ if ver < (7, 2) and _explicit_rocm_torch_index_url() is None:
gfx_codes = _detect_amd_gfx_codes()
_strix_gfx = {"gfx1151", "gfx1150"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
if _detected_strix:
- # Pick the runtime-visible GPU: use the HIP_VISIBLE_DEVICES index
- # into gfx_codes, else default to the first GPU. Skip the override
- # unless the resolved GPU is Strix.
+ # Runtime-visible GPU (HIP_VISIBLE_DEVICES index into gfx_codes, else first);
+ # skip the override unless it's Strix.
_runtime_gfx = gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None
if _runtime_gfx in _strix_gfx:
_selected_gfx = _runtime_gfx
@@ -1353,12 +1706,8 @@ def _ensure_rocm_torch() -> None:
_strix_override_url = f"{_amd_mirror}/{_selected_gfx}/"
_strix_override_pkgs = (
"torch>=2.11.0,<2.12.0",
- # Pin torchvision/torchaudio to the 2.11.x-compatible range.
- # The install uses --index-url (exclusive, no PyPI fallback),
- # so bare unversioned names risk resolving an AMD-index build
- # targeting a different torch major (e.g. 0.27 built against
- # torch 2.12), which fails at runtime with an ABI/version
- # mismatch. Matches _ROCM_TORCH_CONSTRAINT["rocm7.2"].
+ # Pin companions to the 2.11.x range: the exclusive --index-url could
+ # otherwise resolve a build for a different torch major (ABI mismatch).
"torchvision>=0.26.0,<0.27.0",
"torchaudio>=2.11.0,<2.12.0",
)
@@ -1378,14 +1727,15 @@ def _ensure_rocm_torch() -> None:
f" skipping AMD per-gfx index override.\n"
)
- # Strix override on ROCm 7.1 must fire even when has_hip_torch is True --
- # an existing torch with `torch.version.hip == "7.1"` is exactly the broken
- # combo the override repairs, so skipping it leaves users on the known
- # _grouped_mm segfault.
+ # The Strix override must fire even when has_hip_torch is True: an existing
+ # torch.version.hip == "7.1" is exactly the broken combo it repairs.
if _strix_override_url is not None and _strix_override_pkgs is not None:
index_url = _strix_override_url
_torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
- print(f" Strix ROCm 7.1 override -- installing torch from {index_url}")
+ print(
+ f" Strix ROCm 7.1 override -- installing torch from "
+ f"{_strip_index_url_credentials(index_url)}"
+ )
pip_install(
"ROCm torch (Strix arch-specific)",
"--force-reinstall",
@@ -1398,24 +1748,38 @@ def _ensure_rocm_torch() -> None:
constrain = False,
)
rocm_torch_ready = True
- elif not has_hip_torch:
- # Select best matching wheel tag (newest ROCm version <= installed)
- tag = next(
- (
- t
- for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
- if ver >= (maj, mn)
- ),
- None,
- )
- if tag is None:
- print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall")
+ elif not has_hip_torch or _rocm_pin_mismatch:
+ # Reinstall when torch is not ROCm yet, OR a ROCm build's family differs from a pin.
+ # Honour a ROCm pin verbatim; else pick the newest wheel tag <= host.
+ _override_idx = _explicit_rocm_torch_index_url()
+ if _override_idx is not None:
+ index_url = _override_idx
+ tag = _torch_index_leaf(index_url)
else:
- index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
- print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
- _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
- tag, _ROCM_TORCH_PKG_SPECS["_default"]
+ tag = next(
+ (
+ t
+ for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
+ if ver >= (maj, mn)
+ ),
+ None,
)
+ if tag is None:
+ print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping torch reinstall")
+ else:
+ if _override_idx is None:
+ index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
+ print(f" ROCm torch -- installing from {_strip_index_url_credentials(index_url)}")
+ # Only the _grouped_mm-bug gfx arches need the 2.11 spec; other gfx indexes ship
+ # <2.11 and stay on the default range (matches install.ps1 / setup.ps1).
+ if tag in _ROCM_GFX_TORCH211_LEAVES:
+ _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"]
+ elif tag.startswith("gfx"):
+ _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["_default"]
+ else:
+ _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
+ tag, _ROCM_TORCH_PKG_SPECS["_default"]
+ )
pip_install(
f"ROCm torch ({tag})",
"--force-reinstall",
@@ -1504,11 +1868,26 @@ def _infer_no_torch() -> bool:
NO_TORCH = _infer_no_torch()
-# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() so
-# that this script knows which torch variant was selected without re-running
-# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown
-# (standalone `unsloth studio update` runs, where we re-detect normally).
+# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() ("cuda", "rocm",
+# "cpu"; empty = standalone `studio update`, where we re-detect).
_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
+# Standalone update with an explicit pin: derive the backend from the override (classify on
+# the final URL/family segment, mirroring install.sh) instead of re-probing the GPU.
+if not _TORCH_BACKEND:
+ _idx_override = (
+ os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
+ or os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
+ )
+ _idx_leaf = _torch_index_leaf(_idx_override)
+ if _idx_leaf.startswith(("rocm", "gfx")):
+ _TORCH_BACKEND = "rocm"
+ elif _idx_leaf == "cpu":
+ _TORCH_BACKEND = "cpu"
+ elif _is_cuda_family_leaf(_idx_leaf):
+ # Require a digit after "cu" so /current or /custom is NOT branded CUDA (a wrong backend
+ # makes _ensure_rocm_torch return early on AMD hosts). An unknown leaf keeps "" so the
+ # helpers probe the GPU.
+ _TORCH_BACKEND = "cuda"
def _torch_step_label(suffix: str) -> str:
@@ -1724,12 +2103,15 @@ def run(
cmd,
stdout = subprocess.PIPE if quiet else None,
stderr = subprocess.STDOUT if quiet else None,
+ env = _install_env_for_cmd(cmd),
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
_step("error", f"{label} failed (exit code {result.returncode})", _red)
if result.stdout:
- print(result.stdout.decode(errors = "replace"))
+ # Redact before printing: the failing pip command may carry a pinned --index-url
+ # with userinfo/?token= creds, so raw pip error text would leak them.
+ print(_redact_install_output(result.stdout))
sys.exit(result.returncode)
return result
@@ -1737,15 +2119,13 @@ def run(
# Packages to skip on Windows (require special build steps)
WINDOWS_SKIP_PACKAGES = {"triton_kernels"}
-# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode).
-# These either *are* torch extensions or have unconditional
-# ``Requires-Dist: torch``, so installing them would pull torch back in.
-# ``librosa`` is here too despite not requiring torch: upstream ``llvmlite``
-# dropped its macOS x86_64 wheel between 0.42.0 and 0.46.0+ (see
-# https://pypi.org/project/llvmlite/0.47.0/#files -- only
-# macosx_arm64 / manylinux / win_amd64 remain), so on Intel Mac the
-# librosa -> numba -> llvmlite chain triggers a from-source build that fails
-# in CI and on hosts without LLVM 14/15 headers. Tracked in unslothai/unsloth#5046.
+# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode). These
+# either *are* torch extensions or have unconditional ``Requires-Dist: torch``, so
+# installing them pulls torch back in. ``librosa`` is here despite not requiring
+# torch: upstream ``llvmlite`` dropped its macOS x86_64 wheel (0.46.0+ ships only
+# macosx_arm64 / manylinux / win_amd64), so on Intel Mac the librosa -> numba ->
+# llvmlite chain triggers a from-source build that fails without LLVM 14/15 headers.
+# Tracked in unslothai/unsloth#5046.
NO_TORCH_SKIP_PACKAGES = {
"torch-stoi",
"timm",
@@ -1767,7 +2147,8 @@ def _build_flash_attn_wheel_url(env: dict[str, str]) -> str | None:
def _print_optional_install_failure(label: str, result: subprocess.CompletedProcess[str]) -> None:
_step("warning", f"{label} failed (exit code {result.returncode})", _cyan)
if result.stdout:
- print(result.stdout.strip())
+ # Redact any pinned --index-url credentials before printing captured output.
+ print(_redact_install_output(result.stdout).strip())
def _flash_attn_install_disabled() -> bool:
@@ -1913,15 +2294,60 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
# Colab and similar).
cmd.extend(["--python", sys.executable])
cmd.extend(_translate_pip_args_for_uv(args))
- # Torch is pre-installed by install.sh/setup.ps1. Do not add
- # --torch-backend by default -- it can cause solver dead-ends on CPU-only
- # machines. Callers that need it can set UV_TORCH_BACKEND.
+ # Torch is pre-installed, so don't add --torch-backend by default (solver dead-ends on
+ # CPU-only machines); callers can set UV_TORCH_BACKEND. Never add it to a pinned-index
+ # command: uv's torch backend redirects torch to its own per-backend index, defeating the pin.
_tb = os.environ.get("UV_TORCH_BACKEND", "")
- if _tb:
+ if _tb and not _is_pinned_index_cmd(cmd):
cmd.append(f"--torch-backend={_tb}")
return cmd
+# uv resolves --index-url / --default-index at LOWEST priority, so an inherited UV_INDEX /
+# UV_EXTRA_INDEX_URL mirror wins and a pinned torch repair silently ignores the pin.
+# Neutralise these for pinned installs (as install.sh #6898 / install.ps1 / setup.ps1 do).
+# UV_TORCH_BACKEND redirects torch; PIP_* matter for the pip FALLBACK; UV_CONFIG_FILE is
+# stripped + UV_NO_CONFIG=1 (a discovered uv.toml outranks the CLI pin, uv 0.10).
+_UV_INDEX_ENV_VARS = (
+ "UV_CONFIG_FILE",
+ "UV_DEFAULT_INDEX",
+ "UV_INDEX_URL",
+ "UV_INDEX",
+ "UV_EXTRA_INDEX_URL",
+ "UV_TORCH_BACKEND",
+ "UV_FIND_LINKS",
+ "PIP_EXTRA_INDEX_URL",
+ "PIP_FIND_LINKS",
+ # PIP_NO_INDEX=1 makes the pip fallback ignore ALL indexes (defeating --index-url);
+ # PIP_INDEX_URL is dropped too so a stale mirror env can't outrank the pin.
+ "PIP_NO_INDEX",
+ "PIP_INDEX_URL",
+)
+
+
+def _is_pinned_index_cmd(cmd: "list[str] | tuple[str, ...]") -> bool:
+ """True when the command pins an index via --index-url / --default-index."""
+ return any(arg in ("--index-url", "--default-index") for arg in cmd)
+
+
+def _install_env_for_cmd(cmd: "list[str]") -> "dict[str, str] | None":
+ """Return an env with the uv index vars stripped for a pinned-index install.
+
+ None (inherit env) when the command does NOT pin an index, so ordinary installs honour
+ the user's mirror. For pinned commands, the uv index/backend vars are removed,
+ UV_NO_CONFIG=1 set (a discovered uv.toml outranks the CLI pin), and PIP_CONFIG_FILE
+ pointed at os.devnull for the pip fallback. Mirrors install.sh's gate (#6898).
+ """
+ if not _is_pinned_index_cmd(cmd):
+ return None
+ env = os.environ.copy()
+ for name in _UV_INDEX_ENV_VARS:
+ env.pop(name, None)
+ env["UV_NO_CONFIG"] = "1"
+ env["PIP_CONFIG_FILE"] = os.devnull
+ return env
+
+
def pip_install_try(
label: str,
*args: str,
@@ -1948,11 +2374,13 @@ def pip_install_try(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
+ env = _install_env_for_cmd(cmd),
)
if result.returncode == 0:
return True
if VERBOSE and result.stdout:
- print(result.stdout.decode(errors = "replace"))
+ # pip/uv echo index URLs (credentials included) in failure output.
+ print(_redact_install_output(result.stdout))
return False
@@ -2000,13 +2428,14 @@ def pip_install(
uv_cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
+ env = _install_env_for_cmd(uv_cmd),
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode == 0:
return
print(_red(f" uv failed, falling back to pip..."))
if result.stdout:
- print(result.stdout.decode(errors = "replace"))
+ print(_redact_install_output(result.stdout))
pip_cmd = _build_pip_cmd(args) + constraint_args_pip + req_args_pip
run(f"{label} (pip)" if USE_UV else label, pip_cmd)
@@ -2054,10 +2483,9 @@ def install_python_stack() -> int:
global USE_UV, _STEP, _TOTAL
_STEP = 0
- # install.sh (which already installed unsloth) sets SKIP_STUDIO_BASE=1 to
- # avoid reinstalling base packages. "unsloth studio update" does NOT set it,
- # so base packages (unsloth + unsloth-zoo) are reinstalled to pick up new
- # versions.
+ # install.sh sets SKIP_STUDIO_BASE=1 to avoid reinstalling base packages;
+ # `studio update` does NOT, so unsloth + unsloth-zoo are reinstalled to pick
+ # up new versions.
skip_base = os.environ.get("SKIP_STUDIO_BASE", "0") == "1"
# --package installs a different package name (for testing).
package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth")
@@ -2067,9 +2495,9 @@ def install_python_stack() -> int:
if IS_MACOS:
base_total -= 1 # triton step is skipped on macOS
if not IS_MACOS and not NO_TORCH:
- base_total += 1 # ROCm torch check (line 1526) -- all non-macOS platforms
+ base_total += 1 # ROCm torch check (step 2b), non-macOS
if not IS_WINDOWS:
- base_total += 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only
+ base_total += 2 # flash-attn + torch final repair (step 13), Linux
_TOTAL = (base_total - 1) if skip_base else base_total
# 1. Try uv for faster installs (before pip upgrade -- uv venvs don't
@@ -2134,9 +2562,8 @@ def install_python_stack() -> int:
if skip_base:
pass
elif NO_TORCH:
- # No-torch update path: install unsloth + unsloth-zoo with --no-deps
- # (PyPI metadata still declares torch as a hard dep), then runtime deps
- # with --no-deps (avoids transitive torch).
+ # No-torch update path: install unsloth + unsloth-zoo, then runtime deps,
+ # both with --no-deps (PyPI metadata declares torch a hard dep; avoid it).
_progress("base packages (no torch)")
pip_install(
f"Updating {package_name} + unsloth-zoo (no-torch mode)",
@@ -2149,10 +2576,9 @@ def install_python_stack() -> int:
package_name,
"unsloth-zoo",
)
- # Resolve pydantic WITH deps so pip pins pydantic-core to the exact
- # version pydantic's metadata declares. Under --no-deps pip picks the
- # latest of each and trips pydantic's _ensure_pydantic_core_version
- # check. Transitive deps are torch-free.
+ # Resolve pydantic WITH deps so pip pins pydantic-core to the exact version
+ # its metadata declares (under --no-deps pip picks the latest of each and
+ # trips pydantic's _ensure_pydantic_core_version check). Deps are torch-free.
pip_install(
"Installing pydantic (with deps for compatible core)",
"--no-cache-dir",
@@ -2244,6 +2670,7 @@ def install_python_stack() -> int:
_progress(_torch_step_label("check"))
_ensure_cuda_torch()
_ensure_rocm_torch()
+ _ensure_cpu_torch()
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
# version or unknown ROCm version).
@@ -2309,11 +2736,10 @@ def install_python_stack() -> int:
req = REQ_ROOT / "extras-no-deps.txt",
)
- # 4. Overrides (torchao) -- force-reinstall. The torchao version is chosen to
- # match the torch installed in the venv so its C++ extensions load (see
- # _select_torchao_spec). Skip when torch is unavailable (e.g. Intel Mac
- # GGUF-only mode): torchao requires torch. Also skipped on Windows ROCm
- # (no working build; see below).
+ # 4. Overrides (torchao) -- force-reinstall to a version matching the venv's
+ # torch so its C++ extensions load (see _select_torchao_spec). Skipped when
+ # torch is unavailable (Intel Mac GGUF-only) and on Windows ROCm (no working
+ # build; see below).
if NO_TORCH:
_progress("dependency overrides (skipped, no torch)")
elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm():
@@ -2430,14 +2856,12 @@ def install_python_stack() -> int:
[sys.executable, str(SINGLE_ENV / "patch_metadata.py")],
)
- # 13. AMD ROCm: final torch repair. Several steps above can pull in CUDA
- # torch from PyPI (base packages, extras, overrides, studio deps, etc.).
- # Running the repair last ensures ROCm torch is in place at runtime,
- # whichever intermediate step clobbered it.
+ # 13. Final torch repair. Steps above can pull CUDA torch from PyPI, so repair last.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
_progress(_torch_step_label("final"))
_ensure_cuda_torch()
_ensure_rocm_torch()
+ _ensure_cpu_torch()
# 14. Final check (silent; third-party conflicts are expected)
subprocess.run(
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index f7d33a1142..f523b9ff14 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -431,6 +431,167 @@ function Get-PytorchCudaTag {
return "cu126"
}
+# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
+# TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
+function Trim-IndexPathSlashes {
+ param([string]$Url)
+ $value = $Url.Trim()
+ $idx = $value.IndexOfAny([char[]]@('?', '#'))
+ if ($idx -lt 0) {
+ return $value.TrimEnd('/')
+ }
+ return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx)
+}
+
+# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the stale-venv check
+# and install selection so a pinned index wins over GPU probing (parity with the other
+# installers). URL is verbatim; _FAMILY is the leaf joined to the mirror base.
+function Get-PinnedTorchIndexUrl {
+ if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
+ return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL)
+ }
+ if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
+ $base = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
+ return "$base/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
+ }
+ return $null
+}
+
+# Last path segment of a wheel index URL, query/fragment dropped first so a token-authenticated
+# pin (.../cu128?token=x) classifies as cu128 (else it reinstalls every update). Classification
+# only. Shared with the py / install.sh leaf extractors.
+function Get-TorchIndexLeaf {
+ param([string]$Url)
+ if ([string]::IsNullOrWhiteSpace($Url)) { return $null }
+ $path = ($Url -split '[?#]', 2)[0]
+ if ([string]::IsNullOrWhiteSpace($path)) { return $null }
+ return ($path.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
+}
+
+# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
+# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
+# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
+function Redact-InstallOutput {
+ param([string]$Text)
+ if (-not $Text) { return $Text }
+ $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@'
+ $Text = $Text -replace '([?&][^=\s&`]+)=[^\s`]+', '$1='
+ # A #token=... fragment is as sensitive as a query; URL-anchored.
+ return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#'
+}
+
+# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). MUST match
+# the install-spec path below and the other installers; other leaves ship <2.11 and stay default.
+function Test-RocmGfx211Leaf {
+ param([string]$Leaf)
+ return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf
+}
+
+# rocmX.Y versions KNOWN to ship torch 2.11: rocm7.2 only today. Do NOT floor an unknown newer
+# rocm speculatively. MUST match _ROCM_KNOWN_TORCH211_VERSIONS and the rocm7.2 leaf elsewhere.
+function Test-RocmKnown211Version {
+ param([int]$Major, [int]$Minor)
+ return ($Major -eq 7 -and $Minor -eq 2)
+}
+
+# True only for a real CUDA family leaf: "cu" + digits (cu118, cu128, ...). A bare -like 'cu*'
+# would match "custom"/"current" and rebuild the venv every run. Mirrors _is_cuda_family_leaf.
+function Test-CudaFamilyLeaf {
+ param([string]$Leaf)
+ if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false }
+ # EXACT cu+digits: cu128-private routes through the unknown-leaf path instead.
+ return $Leaf -match '^cu[0-9]+$'
+}
+
+# True only for a real pip ROCm family leaf: EXACT rocm[.] or a gfx leaf. A leaf
+# that merely STARTS with rocm (rocm-rel-7.2.1, rocm7.2-private) is a custom pin the verbatim
+# path owns, so anchor the match. Mirrors _is_pip_rocm_family_leaf / install.sh.
+function Test-PipRocmFamilyLeaf {
+ param([string]$Leaf)
+ if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false }
+ # gfx must be followed by a digit (an architecture leaf); gfx-private is custom.
+ return ($Leaf -match '^gfx[0-9]') -or ($Leaf -match '^rocm[0-9]+(\.[0-9]+)?$')
+}
+
+# Stale-venv ROCm comparison for a pinned gfx*/rocm* index. Returns @{ Expected; Installed } so
+# the caller rebuilds when they differ. Mirrors _rocm_pin_family_mismatch (same rocmX.Y / gfx
+# cases). An untagged (no +rocm) wheel never satisfies a ROCm pin -> stale.
+function Get-RocmPinStaleTags {
+ param([string]$PinLeaf, [string]$TorchVersion)
+ $_pinRocm = [regex]::Match($PinLeaf, '^rocm(\d+)\.(\d+)')
+ $_pinVer = if ($_pinRocm.Success) { "$($_pinRocm.Groups[1].Value).$($_pinRocm.Groups[2].Value)" } else { $null }
+ # The family classifier accepts a major-only rocm leaf too (rocm7).
+ $_pinMajorOnly = [regex]::Match($PinLeaf, '^rocm(\d+)$')
+ # Installed rocm version and whether the wheel is a per-arch (three-part) build.
+ $_instRocm = [regex]::Match($TorchVersion, '\+rocm(\d+)\.(\d+)')
+ $_instVer = if ($_instRocm.Success) { "$($_instRocm.Groups[1].Value).$($_instRocm.Groups[2].Value)" } else { $null }
+ $_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+')
+ # A ROCm build MUST carry a +rocm tag; an untagged wheel can't satisfy any ROCm pin.
+ $_instHasRocm = [regex]::IsMatch($TorchVersion, '\+rocm')
+ $_instRel = [regex]::Match($TorchVersion, '^(\d+)\.(\d+)')
+ $_instIs211 = $false
+ if ($_instRel.Success) {
+ $_instIs211 = ([int]$_instRel.Groups[1].Value -gt 2) -or ([int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -ge 11)
+ }
+
+ if ($PinLeaf -like 'gfx*') {
+ if (Test-RocmGfx211Leaf $PinLeaf) {
+ # Expect the AMD per-arch (three-part) 2.11 wheel: satisfied only when BOTH
+ # a 2.11 release AND a three-part rocm tag are installed.
+ $installed = if ($_instIs211 -and $_instPerArch) { "rocm-perarch(torch>=2.11)" } else { "rocm-generic-or-old" }
+ return @{ Expected = "rocm-perarch(torch>=2.11)"; Installed = $installed }
+ }
+ # Non-2.11 gfx leaf (<2.11 spec): stale on an untagged wheel or a 2.11+ build.
+ $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
+ return @{
+ Expected = "rocm(torch<2.11)"
+ Installed = $installed
+ }
+ }
+
+ # Major-only rocm pin (rocm7): compare majors only -- a +rocm6.4 wheel under a rocm7
+ # pin is stale, any +rocm7.x wheel satisfies it (no pinned minor to compare, and the
+ # 2.11-line fallback below would invert both verdicts). Mirrors _rocm_pin_family_mismatch.
+ if ($_pinMajorOnly.Success) {
+ $_pinMaj = [int]$_pinMajorOnly.Groups[1].Value
+ if ($_instVer) {
+ $_instMaj = [int]$_instRocm.Groups[1].Value
+ $expected = if ($_instMaj -eq $_pinMaj) { "rocm$_instVer" } else { "rocm$_pinMaj.x" }
+ return @{ Expected = $expected; Installed = "rocm$_instVer" }
+ }
+ # Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable
+ # version is accepted (matches the lenient unreadable fallback below).
+ $installed = if ($_instHasRocm) { "rocm" } else { "not-rocm" }
+ return @{ Expected = "rocm"; Installed = $installed }
+ }
+
+ # rocmX.Y pin.
+ if ($_pinVer -and $_instVer) {
+ # Both readable: exact compare. When they match AND the pin is KNOWN-2.11, the
+ # installed release must also be 2.11 (a +rocm7.2 wheel drifted to 2.12 shares the
+ # tag but violates the spec), so fold the release into the tag. Mirrors _rocm_pin_family_mismatch.
+ $_pinKnown211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value)
+ $_instOn211 = $_instRel.Success -and [int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -eq 11
+ if ($_pinKnown211 -and -not $_instOn211) {
+ return @{ Expected = "rocm$_pinVer(torch2.11)"; Installed = "rocm$_instVer(torch-off-2.11)" }
+ }
+ return @{ Expected = "rocm$_pinVer"; Installed = "rocm$_instVer" }
+ }
+ $_pinNeeds211 = $false
+ if ($_pinRocm.Success) {
+ # Only KNOWN-2.11 rocm (rocm7.2) is on the 2.11 line (no speculative floor).
+ # Matches _ROCM_KNOWN_TORCH211_VERSIONS.
+ $_pinNeeds211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value)
+ }
+ # Fallback (installed rocm version unreadable): compare on the 2.11 line; an untagged
+ # wheel never satisfies a rocmX.Y pin -> stale.
+ $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
+ return @{
+ Expected = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
+ Installed = $installed
+ }
+}
+
# VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major
# (18->v180, 17->v170), defaulting to v170 when unparseable.
function Get-VcBuildCustomizationsDir {
@@ -813,11 +974,14 @@ function Invoke-SetupCommand {
# Merge stderr into stdout so progress/warning output stays visible
# without flipping $? on successful native commands (PS 5.1 treats
# stderr records as errors that set $? = $false even on exit code 0).
- & $Command 2>&1 | Out-Host
+ # Redact per record: uv/pip echo index URLs (credentials and all) in
+ # their errors, and verbose mode must not bypass the quiet path's
+ # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched.
+ & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
} else {
$output = & $Command 2>&1 | Out-String
if ($LASTEXITCODE -ne 0) {
- Write-Host $output -ForegroundColor Red
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
}
}
return [int]$LASTEXITCODE
@@ -2535,6 +2699,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
$VenvPyExe = Join-Path $VenvDir "Scripts\python.exe"
$installedTorchTag = $null
$shouldRebuild = $false
+ # Set when a stale venv under a pin is repaired in place (force-reinstall) not wiped.
+ $script:PinChangedForceReinstall = $false
if (Test-Path -LiteralPath $VenvPyExe) {
try {
@@ -2551,10 +2717,14 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) {
if ($torchVer -match '\+(cu\d+)') {
$installedTorchTag = $Matches[1]
+ } elseif ($torchVer -match '\+rocm') {
+ # Any +rocm / gfx wheel -> generic "rocm" flavor (the exact version is
+ # repaired later by install_python_stack.py; here we only need the flavor).
+ $installedTorchTag = "rocm"
} elseif ($torchVer -match '\+cpu') {
$installedTorchTag = "cpu"
} else {
- # Untagged wheel (plain "2.x.y" from PyPI) -- treat as cpu
+ # Untagged wheel (plain "2.x.y" from PyPI) -> cpu.
$installedTorchTag = "cpu"
}
} else {
@@ -2570,12 +2740,71 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
}
if (-not $shouldRebuild) {
- $expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
- if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
+ $_pinnedIdx = Get-PinnedTorchIndexUrl
+ $_expectedKnown = $true
+ if ($_pinnedIdx) {
+ $_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx
+ # Digit-gated like the install selection: a custom rocm-* leaf (rocm-current /
+ # rocm-rel-7.2.1) is NOT a ROCm family and must not be stale-compared.
+ if (Test-PipRocmFamilyLeaf $_pinLeaf) {
+ # Don't collapse a pinned ROCm/gfx leaf to a generic "rocm" (would mask a family
+ # change, rocm6.4 -> gfx1151). Get-RocmPinStaleTags uses the SAME 2.11 allowlist
+ # as the install path, so a gfx110X-all/gfx90a/gfx908 pin on a <2.11 wheel is NOT stale.
+ $_rocmTags = Get-RocmPinStaleTags -PinLeaf $_pinLeaf -TorchVersion $torchVer
+ $expectedTorchTag = $_rocmTags.Expected
+ $installedTorchTag = $_rocmTags.Installed
+ } elseif ((Test-CudaFamilyLeaf $_pinLeaf) -or $_pinLeaf -eq 'cpu') {
+ # cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds;
+ # /custom and /current fall through to the unknown-index branch below.
+ $expectedTorchTag = $_pinLeaf
+ } else {
+ # Custom index whose leaf is not a torch flavor (a /simple mirror): the
+ # flavor can't be inferred, so never treat the venv as stale over it.
+ $_expectedKnown = $false
+ $expectedTorchTag = $installedTorchTag
+ }
+ } elseif ($HasNvidiaSmi) {
+ $expectedTorchTag = Get-PytorchCudaTag
+ } elseif ($HasROCm -or $script:ROCmGfxArch) {
+ # AMD/ROCm host with no explicit pin: an existing +rocm wheel is correct (gfx arch
+ # counts even when $HasROCm is false). But only the arches the install path maps to a
+ # repo.amd.com index get ROCm torch; an unmapped arch installs CPU, so expect "cpu"
+ # for those or a correct CPU venv rebuilds every update.
+ $_rocmWheelArches = @(
+ "gfx1201", "gfx1200", # RDNA 4
+ "gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point)
+ "gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3
+ "gfx90a", "gfx908" # MI200 / MI100
+ )
+ if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) {
+ # A correct +rocm wheel is not stale. A CPU wheel on a supported AMD arch is
+ # NOT wiped either (the AMD Windows ROCm override below upgrades it in place);
+ # expect "cpu" for that case. A wrong CUDA wheel still rebuilds.
+ if ($installedTorchTag -eq "cpu") {
+ $expectedTorchTag = "cpu"
+ } else {
+ $expectedTorchTag = "rocm"
+ }
+ } else {
+ $expectedTorchTag = "cpu"
+ }
+ } else {
+ $expectedTorchTag = "cpu"
+ }
+ if ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
$shouldRebuild = $true
}
}
+ # A stale venv under a pin whose torch still imports is repaired IN PLACE (the dependency
+ # pass force-reinstalls from the pin). The rebuild path wipes the venv and would strand a
+ # direct `studio update`; only a broken venv or an unpinned drift wipes.
+ if ($shouldRebuild -and $_pinnedIdx -and $installedTorchTag) {
+ substep "Torch-index pin changed ($installedTorchTag) -- reinstalling torch from the pin in place." "Cyan"
+ $script:PinChangedForceReinstall = $true
+ $shouldRebuild = $false
+ }
+
if ($shouldRebuild) {
$reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" }
if ($InstallerManagedSetup) {
@@ -2653,23 +2882,41 @@ if (Get-Command uv -ErrorAction SilentlyContinue) {
# Helper: install a package, preferring uv with pip fallback
function Fast-Install {
param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
- if ($UseUv) {
- $VenvPy = (Get-Command python).Source
- # An explicit --index-url must win. Inherited uv index env vars otherwise
- # override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop
- # them only for index-pinned installs; mirrors still apply elsewhere.
- $saved = @{}
- if (@($Args_) -contains '--index-url') {
- foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
- $saved[$n] = [Environment]::GetEnvironmentVariable($n)
- Remove-Item "Env:$n" -ErrorAction SilentlyContinue
- }
+ # An explicit --index-url must win: inherited uv index vars otherwise pull CPU torch over
+ # the CUDA/ROCm build (#6898), so drop them for pinned installs (scrub covers the whole
+ # function since the pip fallback honours PIP_* too). UV_TORCH_BACKEND / UV_FIND_LINKS also
+ # reroute; UV_NO_CONFIG=1 (+ dropping UV_CONFIG_FILE) stops a uv.toml index outranking the
+ # pin (uv 0.10); PIP_NO_INDEX / PIP_INDEX_URL would defeat the pinned --index-url in pip.
+ $saved = @{}
+ $pinned = @($Args_) -contains '--index-url'
+ if ($pinned) {
+ foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL',
+ 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'PIP_EXTRA_INDEX_URL', 'PIP_FIND_LINKS',
+ 'PIP_NO_INDEX', 'PIP_INDEX_URL',
+ 'UV_CONFIG_FILE', 'UV_NO_CONFIG', 'PIP_CONFIG_FILE') {
+ $saved[$n] = [Environment]::GetEnvironmentVariable($n)
+ Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
- try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 }
- finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } }
- if ($LASTEXITCODE -eq 0) { return }
+ $env:UV_NO_CONFIG = '1'
+ # A `pip config` global.extra-index-url still adds indexes to the pip FALLBACK;
+ # PIP_CONFIG_FILE = 'nul' (Windows devnull) loads NO config (uv ignores pip config).
+ $env:PIP_CONFIG_FILE = 'nul'
+ }
+ try {
+ if ($UseUv) {
+ $VenvPy = (Get-Command python).Source
+ $result = & uv pip install --python $VenvPy @Args_ 2>&1
+ if ($LASTEXITCODE -eq 0) { return }
+ }
+ & python -m pip install @Args_ 2>&1
+ }
+ finally {
+ if ($pinned) {
+ Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue
+ Remove-Item "Env:PIP_CONFIG_FILE" -ErrorAction SilentlyContinue
+ }
+ foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } }
}
- & python -m pip install @Args_ 2>&1
}
# ── Check if Python deps need updating ──
@@ -2752,6 +2999,10 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1)
# pip install unsloth 2>&1 | Out-Null
# }
+# A torch-index pin change repairs in place: force the dependency pass so the torch install
+# below force-reinstalls from the new pin (else the fast path keeps the old wheel).
+if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false }
+
if (-not $SkipPythonDeps) {
if ($script:UnslothVerbose) {
@@ -2779,7 +3030,13 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir
[Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User')
substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)"
-if ($HasNvidiaSmi) {
+# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD reroute below;
+# matches install.sh / install.ps1 / install_python_stack.py.
+$PinnedTorchIndexUrl = Get-PinnedTorchIndexUrl
+$TorchIndexPinned = [bool]$PinnedTorchIndexUrl
+if ($PinnedTorchIndexUrl) {
+ $CuTag = Get-TorchIndexLeaf $PinnedTorchIndexUrl
+} elseif ($HasNvidiaSmi) {
$CuTag = Get-PytorchCudaTag
} else {
$CuTag = "cpu"
@@ -2800,7 +3057,7 @@ $ROCmIndexUrl = $null
# SDK -- which flips Unsloth out of chat-only (CHAT_ONLY) and enables Train/Export.
# Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed
# ROCm install still falls back to CPU below, so this is safe.
-if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
+if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
@@ -2850,8 +3107,45 @@ if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
}
}
+# A pinned gfx*/rocm index skips the auto-reroute above; route it through the ROCm install path
+# with the same floor/companions the unpinned AMD path uses (mirrors install.ps1), else the CUDA
+# branch installs bare torch and resolves a known-bad wheel for gfx115x/gfx120x/rocm>=7.2.
+if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) {
+ $_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl
+ $_pinRocm211 = $false
+ # Anchor the match ($) so a suffixed custom leaf (rocm7.2-private) falls through to the
+ # verbatim install instead of being floored by its rocm7.2 prefix.
+ if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') {
+ # Only KNOWN-2.11 rocm (rocm7.2) gets the floor (no speculative floor). Matches
+ # Test-RocmKnown211Version / _ROCM_KNOWN_TORCH211_VERSIONS.
+ $_pinRocm211 = Test-RocmKnown211Version -Major ([int]$Matches[1]) -Minor ([int]$Matches[2])
+ }
+ # Only the 2.11 gfx arches need the floor; others publish <2.11 and stay bare. Reuse
+ # Test-RocmGfx211Leaf so this allowlist and the stale-venv check never diverge.
+ $_pinGfx211 = Test-RocmGfx211Leaf $_pinLeaf
+ if ($_pinGfx211 -or $_pinRocm211) {
+ $ROCmIndexUrl = $PinnedTorchIndexUrl
+ $ROCmTorchSpec = "torch>=2.11.0,<2.12.0"
+ $ROCmVisionSpec = "torchvision>=0.26.0,<0.27.0"
+ $ROCmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
+ substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "Cyan"
+ } elseif (Test-PipRocmFamilyLeaf $_pinLeaf) {
+ # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with
+ # bare specs. Only EXACT rocm and gfx* are --index-url families; a suffixed
+ # leaf stays on the verbatim path. Mirrors install.ps1 / _is_pip_rocm_family_leaf.
+ $ROCmIndexUrl = $PinnedTorchIndexUrl
+ $ROCmTorchSpec = "torch"
+ $ROCmVisionSpec = "torchvision"
+ $ROCmAudioSpec = "torchaudio"
+ }
+}
+
$PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
+# A full URL pin is used verbatim; a family pin already set $CuTag. A pinned ROCm install
+# goes through $ROCmIndexUrl; on failure the fallback uses the CPU index, not the ROCm pin.
+$TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
+
$ROCmCpuFallback = $false
if ($ROCmIndexUrl) {
substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..."
@@ -2859,7 +3153,7 @@ if ($ROCmIndexUrl) {
substep " enforcing $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec (known _grouped_mm bug in older wheels)" "Cyan"
}
if ($script:UnslothVerbose) {
- Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl
+ Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
@@ -2868,7 +3162,7 @@ if ($ROCmIndexUrl) {
}
if ($torchInstallExit -ne 0) {
Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow
- Write-Host $output -ForegroundColor Yellow
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow
$ROCmIndexUrl = $null
$ROCmCpuFallback = $true
} else {
@@ -2878,42 +3172,70 @@ if ($ROCmIndexUrl) {
}
}
-if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
+if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) {
substep "installing PyTorch (CPU-only)..."
- # After an AMD ROCm fallback, force-reinstall so a partially-installed ROCm torch
- # (which still satisfies the CPU torch>= range) is replaced by the CPU build. Skip
- # the forced reinstall on a genuine CPU-only host so the common path stays fast.
- # Build the array directly: an if-expression collapses @("x") to a scalar string,
- # which @splat would then enumerate char-by-char into broken single-letter args.
+ # After an AMD ROCm fallback, force-reinstall so a partial ROCm torch (which satisfies the
+ # CPU torch>= range) is replaced by the CPU build; skip on a genuine CPU host to stay fast.
+ # $ROCmCpuFallback matters when a PINNED ROCm index failed ($CuTag is still the rocm leaf).
+ # Build the array directly: an if-expression collapses @("x") to a scalar @splat would
+ # enumerate char-by-char.
$cpuForce = @()
if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }
+ # --force-reinstall on a pin change: a stale +cu / +rocm wheel still satisfies the CPU
+ # torch>= range, so uv would keep it and only swap companions.
+ if ($script:PinChangedForceReinstall) { $cpuForce = @("--force-reinstall") }
+ # A PINNED cpu index installs the bounded trio (parity with _CPU_TORCH_PKG_SPEC): the /cpu
+ # index serves newer torch, and _ensure_cpu_torch keeps any CPU build, so a bare trio could
+ # land an unsupported version. Unpinned CPU hosts keep the bare trio (pre-pin behavior).
+ $cpuTorchSpec = "torch"; $cpuVisionSpec = "torchvision"; $cpuAudioSpec = "torchaudio"
+ if ($TorchIndexPinned) {
+ $cpuTorchSpec = "torch>=2.4,<2.12.0"
+ $cpuVisionSpec = "torchvision>=0.19,<0.27.0"
+ $cpuAudioSpec = "torchaudio>=2.4,<2.12.0"
+ }
if ($script:UnslothVerbose) {
- Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu"
+ Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
- $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String
+ $output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String
$torchInstallExit = $LASTEXITCODE
}
if ($torchInstallExit -ne 0) {
Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red
- Write-Host $output -ForegroundColor Red
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
exit 1
}
} elseif (-not $ROCmIndexUrl) {
substep "installing PyTorch with CUDA support ($CuTag)..."
substep "(This download is ~2.8 GB -- may take a few minutes)"
+ # --force-reinstall on a pin change: an installed cuXXX wheel satisfies the bare torch
+ # requirement (PEP 440 ignores the +cuXXX tag), so without it a changed CUDA pin (cu126
+ # -> cu128) never applies.
+ $cudaForce = @()
+ if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") }
+ # An unknown-leaf custom pin (/simple, /current) routes here with $CuTag as that leaf. Bound
+ # the trio like the fresh custom-pin paths so a mirror can't pull an ABI-newer companion
+ # against the capped torch. Known cu* leaves keep bare specs.
+ $cudaTorchSpec = "torch"
+ $cudaVisionSpec = "torchvision"
+ $cudaAudioSpec = "torchaudio"
+ if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) {
+ $cudaTorchSpec = "torch>=2.4,<2.11.0"
+ $cudaVisionSpec = "torchvision>=0.19,<0.26.0"
+ $cudaAudioSpec = "torchaudio>=2.4,<2.11.0"
+ }
if ($script:UnslothVerbose) {
- Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag"
+ Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
- $output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String
+ $output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String
$torchInstallExit = $LASTEXITCODE
}
if ($torchInstallExit -ne 0) {
Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red
- Write-Host $output -ForegroundColor Red
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
exit 1
}
@@ -2929,7 +3251,7 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
}
if ($tritonInstallExit -ne 0) {
substep "Triton install failed -- torch.compile may not work" "Yellow"
- Write-Host $output -ForegroundColor Yellow
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow
} else {
substep "Triton for Windows installed (enables torch.compile)"
}
@@ -3026,7 +3348,7 @@ foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.8.0", "hf_xet==1.4
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_530/" -ForegroundColor Red
- Write-Host $output -ForegroundColor Red
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
@@ -3061,7 +3383,7 @@ foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_550/" -ForegroundColor Red
- Write-Host $output -ForegroundColor Red
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
@@ -3096,7 +3418,7 @@ foreach ($pkg in @("transformers==5.10.2", "huggingface_hub==1.8.0", "hf_xet==1.
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_510/" -ForegroundColor Red
- Write-Host $output -ForegroundColor Red
+ Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py
index 666ea7ce10..b3a9b99c55 100644
--- a/tests/python/test_cross_platform_parity.py
+++ b/tests/python/test_cross_platform_parity.py
@@ -10,6 +10,8 @@ import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
INSTALL_SH = REPO_ROOT / "install.sh"
INSTALL_PS1 = REPO_ROOT / "install.ps1"
+SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1"
+STACK_PY = REPO_ROOT / "studio" / "install_python_stack.py"
class TestNoTorchBackendAutoInInstallSh:
@@ -180,3 +182,607 @@ class TestUvBytecodeCompileTimeout:
assert (
'$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text
), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
+
+
+class TestTorchIndexOverrideParity:
+ """Every installer must honor UNSLOTH_TORCH_INDEX_URL / _FAMILY so a pinned wheel
+ index wins over GPU probing on all platforms (no asymmetric, per-OS coverage)."""
+
+ @pytest.mark.parametrize(
+ "path",
+ [INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY],
+ ids = ["install.sh", "install.ps1", "setup.ps1", "install_python_stack.py"],
+ )
+ def test_installer_reads_override_env(self, path):
+ text = path.read_text(encoding = "utf-8")
+ for var in ("UNSLOTH_TORCH_INDEX_URL", "UNSLOTH_TORCH_INDEX_FAMILY"):
+ assert var in text, f"{path.name} does not honor {var}"
+
+ @pytest.mark.parametrize(
+ "path",
+ [INSTALL_PS1, SETUP_PS1],
+ ids = ["install.ps1", "setup.ps1"],
+ )
+ def test_amd_reroute_guarded_when_pinned(self, path):
+ # The AMD ROCm reroute must be skipped when the index is explicitly pinned,
+ # so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
+ text = path.read_text(encoding = "utf-8")
+ assert (
+ "TorchIndexPinned" in text
+ ), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
+
+ def test_cuda_pin_overrides_cvd_hide_gate(self):
+ # A pinned cu* index skips ALL host-GPU probing, so the CUDA repair must clear the
+ # CUDA_VISIBLE_DEVICES hide gate too (else the GPU-less CI case bails).
+ text = STACK_PY.read_text(encoding = "utf-8")
+ m = re.search(r"def _ensure_cuda_torch\(\).*?(?=\ndef )", text, re.DOTALL)
+ assert m, "could not locate _ensure_cuda_torch"
+ body = m.group(0)
+ assert "_cuda_pinned" in body, (
+ "_ensure_cuda_torch should compute a CUDA-pin flag so the pin can "
+ "override the CVD hide gate"
+ )
+ assert re.search(
+ r"if not _cuda_pinned and _cvd is not None", body
+ ), "the CVD hide gate must be bypassed when a CUDA index is pinned"
+
+ def test_cpu_repair_pins_supported_torch_range(self):
+ # The explicit-CPU repair must use the bounded CPU/CUDA spec, not a bare trio (the
+ # /cpu index serves torch 2.11+, so a bare install could resolve out of range).
+ text = STACK_PY.read_text(encoding = "utf-8")
+ m = re.search(r"def _ensure_cpu_torch\(\).*?(?=\ndef )", text, re.DOTALL)
+ assert m, "could not locate _ensure_cpu_torch"
+ body = m.group(0)
+ assert "_CPU_TORCH_PKG_SPEC" in body, (
+ "_ensure_cpu_torch should install the bounded _CPU_TORCH_PKG_SPEC, "
+ "not a bare torch/torchvision/torchaudio trio"
+ )
+
+ def test_setup_ps1_stale_check_gates_rocm_on_supported_arch(self):
+ # The stale check must expect ROCm torch only for arches the install path maps to a
+ # repo.amd.com index; expecting "rocm" for an unmapped arch marks a good CPU venv stale.
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ assert "_rocmWheelArches" in text, (
+ "setup.ps1 stale check should restrict the ROCm expected-tag to the "
+ "supported gfx wheel arches"
+ )
+
+
+class TestGfx211AllowlistParity:
+ """The gfx per-arch 2.11-floor leaves (gfx120X-all / gfx1151 / gfx1150) must be the
+ SAME set in every installer and its stale/mismatch check. When they diverged, a
+ pinned gfx110X-all / gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update."""
+
+ EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150"}
+
+ def test_install_sh_allowlist(self):
+ text = INSTALL_SH.read_text(encoding = "utf-8").lower()
+ # install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx120x-all|gfx1151|gfx1150).
+ m = re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150", text)
+ assert m, "install.sh gfx-2.11 allowlist case not found / changed"
+
+ def test_install_ps1_allowlist(self):
+ text = INSTALL_PS1.read_text(encoding = "utf-8").lower()
+ m = re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text)
+ assert m, "install.ps1 $_pinGfx211 allowlist not found / changed"
+
+ def test_setup_ps1_defines_single_allowlist_helper(self):
+ # setup.ps1 must define the allowlist once (Test-RocmGfx211Leaf) and reuse it, so
+ # the stale check and install spec can't disagree.
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ assert (
+ "function Test-RocmGfx211Leaf" in text
+ ), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper"
+ assert re.search(
+ r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()
+ ), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
+ assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, (
+ "setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not "
+ "re-hardcode the allowlist (they must not diverge)"
+ )
+
+ def test_stack_py_allowlist(self):
+ text = STACK_PY.read_text(encoding = "utf-8").lower()
+ assert (
+ '"gfx120x-all", "gfx1151", "gfx1150"' in text
+ ), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
+
+
+class TestCudaLeafDigitParity:
+ """A wheel-family leaf is CUDA only when it is "cu" + digits (cu118/cu128/...).
+ A bare cu* glob wrongly catches mirror leaves like /custom or /current; when
+ that happened the venv was marked stale and rebuilt on every run. Every
+ installer must require a digit after "cu" in its family/CUDA classification."""
+
+ def test_stack_py_requires_cu_digit(self):
+ text = STACK_PY.read_text(encoding = "utf-8")
+ # EXACT cu+digits: a custom leaf like cu128-private must route to the
+ # verbatim/unknown path, not be compared against the installed +cu128 tag.
+ assert re.search(
+ r'r"cu\[0-9\]\+"', text
+ ), "install_python_stack.py _is_cuda_family_leaf must fullmatch cu[0-9]+"
+
+ def test_setup_ps1_requires_cu_digit(self):
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ # EXACT cu+digits: cu128-private must not classify as CUDA (it would become
+ # the expected tag and rebuild the venv on every update).
+ assert re.search(
+ r"'\^cu\[0-9\]\+\$'", text
+ ), "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9]+$, not a cu* prefix"
+ # The stale-venv branch must go through the digit-guarded helper.
+ assert (
+ "Test-CudaFamilyLeaf $_pinLeaf" in text
+ ), "setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf"
+
+ def test_install_ps1_requires_cu_digit_in_gpu_branch(self):
+ text = INSTALL_PS1.read_text(encoding = "utf-8")
+ assert re.search(
+ r"'\^cu\[0-9\]'", text
+ ), "install.ps1 Get-TauriGpuBranch must require a digit after cu"
+
+ def test_install_sh_requires_cu_digit_in_gpu_branch(self):
+ text = INSTALL_SH.read_text(encoding = "utf-8")
+ # The _tauri_gpu_branch cuda case must be cu[0-9]*, not a bare cu*.
+ assert re.search(
+ r"cu\[0-9\]\*\)\s*echo \"cuda\"", text
+ ), "install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*"
+
+ def test_install_sh_backend_export_requires_cu_digit(self):
+ text = INSTALL_SH.read_text(encoding = "utf-8")
+ # Brand CUDA only on cu[0-9]*; a bare catch-all *) -> cuda would mis-brand
+ # /current, /custom pins and skip ROCm repair on AMD hosts.
+ assert re.search(
+ r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text
+ ), "install.sh backend export must brand cuda only on cu[0-9]*"
+ # An unknown leaf must NOT commit a cuda backend (it unsets instead).
+ assert re.search(
+ r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text
+ ), "install.sh backend export must unset (not force cuda) on an unknown leaf"
+
+ def test_install_sh_lowercases_backend_leaf(self):
+ text = INSTALL_SH.read_text(encoding = "utf-8")
+ # The leaf feeding both the backend case and the 2.11 floor case must be
+ # lowercased so the canonical gfx120X-all (capital X) matches.
+ assert re.search(
+ r"_torch_index_leaf=\$\(printf '%s' \"\$_torch_index_leaf\" \| tr '\[:upper:\]' '\[:lower:\]'\)",
+ text,
+ ), "install.sh must lowercase _torch_index_leaf before the gfx/rocm/cu case matches"
+
+
+class TestKnown211SetParity:
+ """The KNOWN-2.11 rocm/gfx set must be identical across all four installers:
+ exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}.
+ rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively."""
+
+ def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self):
+ text = INSTALL_SH.read_text(encoding = "utf-8")
+ # The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves.
+ assert re.search(
+ r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text
+ ), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
+ # No speculative rocm7.3 anywhere.
+ assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"
+
+ def test_python_known_211_versions_is_only_rocm72(self):
+ text = STACK_PY.read_text(encoding = "utf-8")
+ assert "_ROCM_KNOWN_TORCH211_VERSIONS" in text
+ # The frozenset literal is exactly {(7, 2)}.
+ m = re.search(r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text)
+ assert m is not None, "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS"
+ assert "(7, 2)" in m.group(1)
+ assert "7, 3" not in m.group(1) and "7, 1" not in m.group(1)
+
+ def test_setup_ps1_known_211_helper_is_only_rocm72(self):
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ assert "Test-RocmKnown211Version" in text
+ # The predicate is Major -eq 7 -and Minor -eq 2 (only rocm7.2).
+ assert re.search(
+ r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2", text
+ ), "setup.ps1 Test-RocmKnown211Version must accept only rocm7.2"
+
+ def test_install_ps1_pin_floor_is_only_rocm72(self):
+ text = INSTALL_PS1.read_text(encoding = "utf-8")
+ # The pinned-ROCm install-spec floor must be Major -eq 7 -and Minor -eq 2,
+ # not the speculative >= 2 that would floor a non-existent rocm7.3.
+ assert re.search(
+ r"\$_pinRocm211 = \(\[int\]\$Matches\[1\] -eq 7 -and \[int\]\$Matches\[2\] -eq 2\)",
+ text,
+ ), "install.ps1 pinned-ROCm floor must be rocm7.2 only (no speculative >= 2)"
+
+ def test_ps1_pin_floor_gate_is_anchored(self):
+ """The floor-selection gate that reads $_pinRocm211 from the raw leaf must anchor
+ the rocm match ($), or a suffixed custom leaf (rocm7.2-private) matches the rocm7.2
+ prefix, takes the 2.11-floor branch, and is force-routed through the ROCm path
+ before the exact-match elseif can send it to the verbatim install (Codex P2)."""
+ for path, label in ((INSTALL_PS1, "install.ps1"), (SETUP_PS1, "setup.ps1")):
+ text = path.read_text(encoding = "utf-8")
+ assert "-match '^rocm(\\d+)\\.(\\d+)$'" in text, (
+ f"{label} floor gate must anchor the rocm match (^rocm(\\d+)\\.(\\d+)$) so a "
+ "suffixed custom leaf is not floored/routed as rocm7.2"
+ )
+ assert (
+ "-match '^rocm(\\d+)\\.(\\d+)'\n" not in text
+ ), f"{label} floor gate must not use the unanchored ^rocm(\\d+)\\.(\\d+) prefix"
+
+ def test_install_ps1_bounds_unknown_leaf_pinned_torch(self):
+ """install.ps1's pinned-torch install must bound BOTH companions on EVERY
+ index, cu families included: torchaudio 2.11 dropped its exact torch
+ pin from the wheel metadata, so a bare companion beside torch<2.11 can
+ resolve a mismatched 2.11.0 build (Codex P2, then unconditional per the
+ torchaudio 2.11 unpinning)."""
+ text = INSTALL_PS1.read_text(encoding = "utf-8")
+ assert (
+ '$_pinVisionSpec = "torchvision>=0.19,<0.26.0"' in text
+ ), "install.ps1 custom-pin install must bound torchvision (>=0.19,<0.26.0)"
+ assert (
+ '$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"' in text
+ ), "install.ps1 custom-pin install must bound torchaudio (>=2.4,<2.11.0)"
+ # No cu-family exemption: the bounds apply unconditionally.
+ assert (
+ "$_pinCuLeaf" not in text
+ ), "install.ps1 must bound companions on every index (no cu-family exemption)"
+ # The bounded companions must actually be passed to the install command.
+ assert re.search(
+ r'"torch>=2\.4,<2\.11\.0" \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl',
+ text,
+ ), "install.ps1 custom-pin install must pass the bounded companion specs to uv"
+
+ def test_gfx_allowlist_matches_across_installers(self):
+ # The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each.
+ gfx = ("gfx120x-all", "gfx1151", "gfx1150")
+ for path, label in (
+ (INSTALL_SH, "install.sh"),
+ (INSTALL_PS1, "install.ps1"),
+ (SETUP_PS1, "setup.ps1"),
+ (STACK_PY, "install_python_stack.py"),
+ ):
+ low = path.read_text(encoding = "utf-8").lower()
+ for g in gfx:
+ assert g in low, f"{label} missing gfx 2.11 allowlist member {g}"
+
+
+class TestPinnedRocmLeafDigitParity:
+ """A pinned index is a pip ROCm --default-index family only when its leaf is an
+ EXACT rocm+digits (rocm7 / rocm7.2) or gfx*. A ^rocm[0-9] PREFIX (or a bare rocm*
+ glob) wrongly catches a custom mirror / find-links leaf (rocm-current /
+ rocm-rel-7.2.1) AND a suffixed private-mirror leaf (rocm7.2-private / rocm7-current),
+ routing it through the ROCm install path (which silently falls back to CPU on
+ failure) or skipping the custom-index companion bounds, instead of the verbatim
+ --default-index install. All installers must match the family EXACTLY: Python and
+ install.sh via a shared _is_pip_rocm_family_leaf, setup.ps1 via Test-PipRocmFamilyLeaf,
+ install.ps1 via an anchored ^rocm[0-9]+(\\.[0-9]+)?$ reroute."""
+
+ def test_install_ps1_pinned_reroute_requires_rocm_digit(self):
+ text = INSTALL_PS1.read_text(encoding = "utf-8")
+ # The pinned gfx*/rocm reroute must match rocm EXACTLY (anchored), so a suffixed
+ # rocm7.2-private / rocm-current falls through to the verbatim --default-index path.
+ assert "-match '^rocm[0-9]+(\\.[0-9]+)?$'" in text, (
+ "install.ps1 pinned-index reroute must anchor the rocm match "
+ "(^rocm[0-9]+(\\.[0-9]+)?$), not a bare -like 'rocm*' or an unanchored ^rocm\\d"
+ )
+ # Neither the broad glob nor the unanchored prefix may drive that reroute.
+ assert (
+ "-like 'rocm*'" not in text
+ ), "install.ps1 must not route a pinned index on a bare -like 'rocm*' glob"
+ assert (
+ "-match '^rocm\\d'" not in text
+ ), "install.ps1 must not route a pinned index on an unanchored -match '^rocm\\d'"
+
+ def test_setup_ps1_pinned_reroute_requires_rocm_digit(self):
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ # setup.ps1 routes every family decision through Test-PipRocmFamilyLeaf, which
+ # anchors the rocm match so a suffixed custom leaf stays on the verbatim path.
+ assert (
+ "function Test-PipRocmFamilyLeaf" in text
+ ), "setup.ps1 must define Test-PipRocmFamilyLeaf (the exact rocm/gfx family gate)"
+ assert "'^rocm[0-9]+(\\.[0-9]+)?$'" in text, (
+ "setup.ps1 Test-PipRocmFamilyLeaf must anchor the rocm match "
+ "(^rocm[0-9]+(\\.[0-9]+)?$) so rocm7.2-private / rocm-current stay verbatim"
+ )
+ pinned_block = text[text.find("$_pinGfx211 = Test-RocmGfx211Leaf") :][:2000]
+ assert (
+ "-like 'rocm*'" not in pinned_block
+ ), "setup.ps1 pinned reroute must not route on a bare -like 'rocm*' glob"
+
+ def test_install_sh_repairable_requires_rocm_digit(self):
+ text = INSTALL_SH.read_text(encoding = "utf-8")
+ # _torch_index_repairable routes rocm/gfx through the exact-match helper.
+ assert (
+ "_is_pip_rocm_family_leaf" in text
+ ), "install.sh must define/use _is_pip_rocm_family_leaf for the exact rocm gate"
+ # gfx needs a following digit: gfx-private / gfxfoo are custom verbatim pins.
+ assert re.search(
+ r'case "\$1" in\n\s*gfx\[0-9\]\*\) return 0', text
+ ), "install.sh _is_pip_rocm_family_leaf must treat only gfx* as a family"
+ assert not re.search(
+ r'case "\$1" in\n\s*gfx\*\) return 0', text
+ ), "install.sh _is_pip_rocm_family_leaf must not family-match a bare gfx* glob"
+
+ def test_stack_py_pip_rocm_family_requires_digit(self):
+ text = STACK_PY.read_text(encoding = "utf-8")
+ assert re.search(
+ r'fullmatch\(r"rocm\\d\+\(\?:\\\.\\d\+\)\?", leaf\)', text
+ ), "install_python_stack.py _is_pip_rocm_family_leaf must fullmatch rocm\\d+(?:\\.\\d+)?"
+ # The unanchored prefix must be gone from the family/flavor gates.
+ assert (
+ 're.match(r"^rocm\\d"' not in text
+ ), "install_python_stack.py must not gate a family on an unanchored re.match(^rocm\\d)"
+
+ def test_install_sh_rocm_side_effects_digit_gated(self):
+ """The AMD bitsandbytes + 'repair ROCm torch' side effects must fire only on
+ an EXACT ROCm family (rocm7.2/gfx*), not a bare */rocm* whole-URL glob nor a
+ ^rocm[0-9] prefix that catches a custom CPU/CUDA index like /rocm-current or a
+ suffixed /rocm7.2-private and force-repairs it from the wrong --default-index."""
+ text = INSTALL_SH.read_text(encoding = "utf-8")
+ assert (
+ 'if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then\n _torch_index_is_rocm_family=true'
+ in text
+ ), "install.sh must set _torch_index_is_rocm_family from the exact-match helper"
+ assert (
+ '[ "$_torch_index_is_rocm_family" = true ]' in text
+ ), "install.sh ROCm bnb/repair hooks must gate on _torch_index_is_rocm_family"
+ assert (
+ "*/rocm*|*/gfx*)\n _install_bnb_rocm" not in text
+ ), "install.sh must not gate _install_bnb_rocm on a bare */rocm* whole-URL glob"
+
+
+class TestPinnedIndexClearsUvEnvParity:
+ """Every installer must neutralise the uv index env vars for a pinned torch
+ install (#6898). uv treats the default index (--index-url / --default-index) as
+ lowest priority, so an inherited UV_INDEX / UV_EXTRA_INDEX_URL mirror would win
+ under uv's first-index strategy and pull torch from the wrong index -- after
+ which the pinned wheel index is silently never used."""
+
+ UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL")
+
+ def test_install_sh_clears_uv_index_vars(self):
+ text = INSTALL_SH.read_text(encoding = "utf-8")
+ assert (
+ "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in text
+ ), "install.sh run_install_cmd must clear the uv index vars for --default-index installs"
+
+ def test_install_ps1_clears_uv_index_vars(self):
+ text = INSTALL_PS1.read_text(encoding = "utf-8")
+ for var in self.UV_VARS:
+ assert var in text, f"install.ps1 must clear {var} for pinned installs"
+
+ def test_setup_ps1_clears_uv_index_vars(self):
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ for var in self.UV_VARS:
+ assert var in text, f"setup.ps1 must clear {var} for pinned installs"
+
+ def test_stack_py_clears_uv_index_vars(self):
+ text = STACK_PY.read_text(encoding = "utf-8")
+ assert "_install_env_for_cmd" in text, (
+ "install_python_stack.py must scrub inherited uv index vars for pinned "
+ "installs via _install_env_for_cmd (parity with install.sh #6898)"
+ )
+ for var in self.UV_VARS:
+ assert var in text, f"install_python_stack.py must clear {var} for pinned installs"
+
+ def test_all_installers_clear_uv_torch_backend(self):
+ """uv's torch backend redirects torch resolution to its own per-backend
+ index even against an explicit pin, so every installer's pinned-install
+ scrub must clear UV_TORCH_BACKEND too."""
+ sh = INSTALL_SH.read_text(encoding = "utf-8")
+ assert "-u UV_TORCH_BACKEND" in sh, "install.sh pinned scrub must clear UV_TORCH_BACKEND"
+ for path in (INSTALL_PS1, SETUP_PS1):
+ text = path.read_text(encoding = "utf-8")
+ assert (
+ "'UV_TORCH_BACKEND'" in text
+ ), f"{path.name} pinned scrub must clear UV_TORCH_BACKEND"
+ stack = STACK_PY.read_text(encoding = "utf-8")
+ assert (
+ '"UV_TORCH_BACKEND",' in stack
+ ), "install_python_stack.py strip tuple must include UV_TORCH_BACKEND"
+
+ def test_stack_py_strips_pip_extra_index_for_pip_fallback(self):
+ """The pip fallback honours PIP_EXTRA_INDEX_URL (pip adds it IN ADDITION
+ to --index-url), so the pinned-command scrub must strip it."""
+ stack = STACK_PY.read_text(encoding = "utf-8")
+ assert (
+ '"PIP_EXTRA_INDEX_URL",' in stack
+ ), "install_python_stack.py strip tuple must include PIP_EXTRA_INDEX_URL"
+
+ def test_all_installers_scrub_find_links(self):
+ """uv's --find-links (env UV_FIND_LINKS) adds candidate locations that can
+ satisfy torch off a pinned index; every pinned-install scrub must clear it."""
+ sh = INSTALL_SH.read_text(encoding = "utf-8")
+ assert "-u UV_FIND_LINKS" in sh
+ for path in (INSTALL_PS1, SETUP_PS1):
+ assert "'UV_FIND_LINKS'" in path.read_text(encoding = "utf-8"), path.name
+ stack = STACK_PY.read_text(encoding = "utf-8")
+ assert '"UV_FIND_LINKS",' in stack and '"PIP_FIND_LINKS",' in stack
+
+ def test_setup_ps1_scrub_covers_pip_fallback(self):
+ """setup.ps1's Fast-Install must keep the scrub active through the pip
+ fallback (pip honours PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS in addition to
+ --index-url); restoring the vars before the fallback reopens the hole."""
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ fi = text[text.find("function Fast-Install") :][:2500]
+ assert "'PIP_EXTRA_INDEX_URL'" in fi and "'PIP_FIND_LINKS'" in fi
+ # the pip fallback must sit INSIDE the try whose finally restores the vars
+ assert fi.find("python -m pip install") < fi.find(
+ "finally"
+ ), "pip fallback must run before the scrub is restored"
+
+ def test_all_installers_disable_uv_config_for_pinned_installs(self):
+ """A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin
+ (verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default
+ [[index]] both resolve torch+cpu against an explicit --index-url /
+ --default-index cu126 pin; UV_NO_CONFIG=1 restores the pin). Every
+ installer's pinned scrub must set UV_NO_CONFIG=1 and drop UV_CONFIG_FILE."""
+ sh = INSTALL_SH.read_text(encoding = "utf-8")
+ assert "-u UV_CONFIG_FILE UV_NO_CONFIG=1" in sh, (
+ "install.sh run_install_cmd must set UV_NO_CONFIG=1 and drop "
+ "UV_CONFIG_FILE for --default-index installs"
+ )
+ for path in (INSTALL_PS1, SETUP_PS1):
+ text = path.read_text(encoding = "utf-8")
+ assert "'UV_CONFIG_FILE'" in text, f"{path.name} must drop UV_CONFIG_FILE"
+ assert (
+ "$env:UV_NO_CONFIG = '1'" in text
+ ), f"{path.name} must set UV_NO_CONFIG=1 for pinned installs"
+ stack = STACK_PY.read_text(encoding = "utf-8")
+ assert (
+ '"UV_CONFIG_FILE",' in stack
+ ), "install_python_stack.py strip tuple must include UV_CONFIG_FILE"
+ assert (
+ 'env["UV_NO_CONFIG"] = "1"' in stack
+ ), "_install_env_for_cmd must set UV_NO_CONFIG=1 for pinned installs"
+
+ def test_pip_fallbacks_disable_pip_config_files(self):
+ """The pip FALLBACK (uv missing/failed) honours user/site pip config files
+ even with the PIP_* env vars stripped: `pip config set
+ global.extra-index-url` still adds indexes to a pinned install. pip loads
+ NO configuration files when PIP_CONFIG_FILE is the platform devnull, so
+ the two installers that HAVE a pip fallback (install_python_stack.py and
+ setup.ps1's Fast-Install) must set it in their pinned scrub. install.sh
+ and install.ps1 are uv-only (no python -m pip fallback) and need no
+ equivalent."""
+ stack = STACK_PY.read_text(encoding = "utf-8")
+ assert 'env["PIP_CONFIG_FILE"] = os.devnull' in stack, (
+ "_install_env_for_cmd must point PIP_CONFIG_FILE at os.devnull for "
+ "pinned installs (pip fallback isolation)"
+ )
+ setup = SETUP_PS1.read_text(encoding = "utf-8")
+ assert "$env:PIP_CONFIG_FILE = 'nul'" in setup, (
+ "setup.ps1 Fast-Install pinned scrub must point PIP_CONFIG_FILE at nul "
+ "(Windows devnull) so the pip fallback ignores user/site pip config"
+ )
+ assert (
+ "'PIP_CONFIG_FILE'" in setup
+ ), "setup.ps1 must save/restore PIP_CONFIG_FILE around the pinned scrub"
+
+ def test_setup_ps1_bounds_unknown_leaf_pinned_torch(self):
+ """A first-time/changed unknown-leaf custom pin routes through setup.ps1's
+ CUDA branch; install.ps1's fresh pinned install, install.sh, and the Python
+ verbatim path bound the WHOLE trio, so the Windows update path must too -- a
+ private mirror serving newer torch OR newer companions must not lift the venv
+ above the supported range under the pin."""
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ # The custom-leaf branch bounds torch AND both companions (parity with the
+ # other installers' custom-pin trio bounds), gated on a non-cu-family leaf.
+ for spec in (
+ '$cudaTorchSpec = "torch>=2.4,<2.11.0"',
+ '$cudaVisionSpec = "torchvision>=0.19,<0.26.0"',
+ '$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"',
+ ):
+ assert spec in text, f"setup.ps1 must bound the custom-leaf trio: {spec}"
+ assert (
+ "if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) {" in text
+ ), "the custom-leaf trio bounds must be gated on a pinned non-cu-family leaf"
+ assert (
+ "Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec" in text
+ ), "setup.ps1's CUDA branch must install via the bounded spec variables"
+
+ def test_setup_ps1_bounds_pinned_cpu_torch(self):
+ """setup.ps1's CPU branch must bound the trio under an explicit pin (parity with
+ _CPU_TORCH_PKG_SPEC): the /cpu index serves newer torch, and _ensure_cpu_torch
+ keeps any CPU build, so a bare pinned trio could land an unsupported version.
+ An unpinned CPU host keeps the bare trio (pre-pin behavior unchanged)."""
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ for spec in (
+ '$cpuTorchSpec = "torch>=2.4,<2.12.0"',
+ '$cpuVisionSpec = "torchvision>=0.19,<0.27.0"',
+ '$cpuAudioSpec = "torchaudio>=2.4,<2.12.0"',
+ ):
+ assert spec in text, f"setup.ps1 must bound the pinned CPU trio: {spec}"
+ assert (
+ "if ($TorchIndexPinned) {" in text
+ ), "the CPU trio bounds must be gated on an explicit pin"
+ assert (
+ "Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce" in text
+ ), "setup.ps1's CPU branch must install via the spec variables"
+ # The ceilings mirror the Python repair spec exactly.
+ stack = STACK_PY.read_text(encoding = "utf-8")
+ spec_block = re.search(r"_CUDA_TORCH_PKG_SPEC[^(]*\(\s*(.*?)\)", stack, re.DOTALL)
+ assert spec_block and '"torch>=2.4,<2.12.0"' in spec_block.group(1), (
+ "_CPU_TORCH_PKG_SPEC (via _CUDA_TORCH_PKG_SPEC) must keep the torch<2.12 "
+ "ceiling the setup.ps1 pinned CPU branch mirrors"
+ )
+
+ def test_setup_ps1_stale_check_requires_rocm_digit(self):
+ """The stale-venv check must use the same EXACT rocm/gfx gate as the install
+ selection (Test-PipRocmFamilyLeaf), or a custom rocm-* / suffixed rocm7.2-private
+ leaf is stale-compared as a family and force-reinstalls on every studio update."""
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ anchor = text.find("$_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx")
+ assert anchor >= 0, "setup.ps1 stale check must classify the pinned leaf"
+ stale = text[anchor:][:2500]
+ assert (
+ "Test-PipRocmFamilyLeaf" in stale
+ ), "setup.ps1 stale check must gate rocm leaves via the exact Test-PipRocmFamilyLeaf"
+ assert (
+ stale.count("-like 'rocm*'") == 0
+ ), "setup.ps1 stale check must not use a bare -like 'rocm*' glob"
+ assert (
+ "-match '^rocm\\d'" not in stale
+ ), "setup.ps1 stale check must not use an unanchored -match '^rocm\\d'"
+
+
+class TestIndexPathSlashTrimParity:
+ """Every installer must trim trailing PATH slashes only on the verbatim
+ UNSLOTH_TORCH_INDEX_URL override, preserving a ?query/#fragment token: a whole-URL
+ strip corrupts a base64 token ending in "/", a single strip leaves a double-slash leaf
+ empty. The helper must be DEFINED and WIRED into the override return in all four."""
+
+ def test_helper_defined_in_all_installers(self):
+ assert "def _trim_index_path_slashes(" in STACK_PY.read_text(encoding = "utf-8")
+ assert "_trim_index_path_slashes()" in INSTALL_SH.read_text(encoding = "utf-8")
+ assert "function Trim-IndexPathSlashes" in INSTALL_PS1.read_text(encoding = "utf-8")
+ assert "function Trim-IndexPathSlashes" in SETUP_PS1.read_text(encoding = "utf-8")
+
+ def test_helper_wired_into_override_in_all_installers(self):
+ assert "_trim_index_path_slashes(url)" in STACK_PY.read_text(encoding = "utf-8")
+ assert '_url=$(_trim_index_path_slashes "$_url")' in INSTALL_SH.read_text(encoding = "utf-8")
+ assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in INSTALL_PS1.read_text(
+ encoding = "utf-8"
+ )
+ assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in SETUP_PS1.read_text(
+ encoding = "utf-8"
+ )
+
+
+class TestInstallOutputRedactionParity:
+ """uv/pip failure text embeds the failing --index-url verbatim, so a captured install
+ log dumped on error can leak a user:token@ or ?token= secret. Every installer must
+ DEFINE a redaction helper and WIRE it into the captured-output print path."""
+
+ def test_helper_defined_in_all_installers(self):
+ assert "def _redact_install_output(" in STACK_PY.read_text(encoding = "utf-8")
+ assert "_redact_install_output()" in INSTALL_SH.read_text(encoding = "utf-8")
+ assert "function Redact-InstallOutput" in INSTALL_PS1.read_text(encoding = "utf-8")
+ assert "function Redact-InstallOutput" in SETUP_PS1.read_text(encoding = "utf-8")
+
+ def test_helper_wired_into_failure_print(self):
+ # install.sh dumps the captured log through the redactor on failure.
+ assert '_redact_install_output "$_log"' in INSTALL_SH.read_text(encoding = "utf-8")
+ # Both ps1 installers redact the captured $output before Write-Host on non-zero exit.
+ assert (
+ "Write-Host (Redact-InstallOutput $output) -ForegroundColor Red"
+ in INSTALL_PS1.read_text(encoding = "utf-8")
+ )
+ assert (
+ "Write-Host (Redact-InstallOutput $output) -ForegroundColor Red"
+ in SETUP_PS1.read_text(encoding = "utf-8")
+ )
+ # Python redacts the captured stdout before printing.
+ assert "_redact_install_output(" in STACK_PY.read_text(encoding = "utf-8")
+
+
+class TestPipNoIndexScrubParity:
+ """The plain-pip fallback honours PIP_*: PIP_NO_INDEX=1 makes it ignore ALL indexes
+ (defeating the pinned --index-url) and PIP_INDEX_URL replaces the pin. The two installers
+ that HAVE a plain-pip fallback (Python + setup.ps1) must scrub both for a pinned install.
+ install.sh / install.ps1 are uv-only (--default-index), which ignores pip config/env."""
+
+ def test_python_scrubs_pip_no_index_and_pip_index_url(self):
+ text = STACK_PY.read_text(encoding = "utf-8")
+ assert '"PIP_NO_INDEX"' in text
+ assert '"PIP_INDEX_URL"' in text
+
+ def test_setup_ps1_scrubs_pip_no_index_and_pip_index_url(self):
+ text = SETUP_PS1.read_text(encoding = "utf-8")
+ assert "'PIP_NO_INDEX'" in text
+ assert "'PIP_INDEX_URL'" in text
diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py
index 9015ff8c9d..3a12e53f95 100644
--- a/tests/python/test_install_python_stack.py
+++ b/tests/python/test_install_python_stack.py
@@ -54,6 +54,24 @@ class TestBuildUvCmdTorchBackend:
a.startswith("--torch-backend") for a in cmd
), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}"
+ def test_uv_torch_backend_skipped_for_pinned_index(self):
+ """A pinned-index command must NOT get --torch-backend: uv's torch backend
+ redirects torch resolution to its own per-backend index even when
+ --index-url is given (verified: cu128 pin + backend cpu installs
+ torch+cpu), defeating the pin."""
+ for pin_flag in ("--index-url", "--default-index"):
+ with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
+ cmd = self._call(("torch", pin_flag, "https://download.pytorch.org/whl/cu128"))
+ assert not any(
+ a.startswith("--torch-backend") for a in cmd
+ ), f"{pin_flag} command must not carry --torch-backend, got: {cmd}"
+
+ def test_uv_torch_backend_kept_for_unpinned(self):
+ """Non-pinned commands still honour UV_TORCH_BACKEND."""
+ with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
+ cmd = self._call(("somepackage",))
+ assert "--torch-backend=cpu" in cmd
+
class TestUvSafePath:
"""_uv_safe_path hands uv a space-free `-c`/`-r` path (issue #6503)."""
@@ -148,3 +166,119 @@ class TestUvSafePathHardening:
assert " " not in value
assert Path(value).read_text() == "transformers>=4.57.6\n"
+
+
+class TestPinnedIndexClearsUvEnv:
+ """A pinned torch install (--index-url / --default-index) must neutralise an
+ inherited UV_INDEX / UV_EXTRA_INDEX_URL so the pinned wheel index wins.
+
+ uv treats the default index (--index-url / --default-index) as LOWEST priority,
+ so an inherited UV_INDEX / UV_EXTRA_INDEX_URL (a corporate/CPU mirror) would be
+ searched first and, under uv's default first-index strategy, resolve torch from
+ the wrong mirror -- after which the marker records a wheel index that was never
+ used. install.sh (#6898), install.ps1 and setup.ps1 already clear these for
+ pinned installs; install_python_stack must match (parity across all installers).
+ """
+
+ UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL")
+
+ def test_pinned_index_url_strips_uv_index_vars(self):
+ cmd = [
+ "uv",
+ "pip",
+ "install",
+ "--force-reinstall",
+ "torch",
+ "torchvision",
+ "torchaudio",
+ "--index-url",
+ "https://download.pytorch.org/whl/cu128",
+ ]
+ with mock.patch.dict(
+ os.environ,
+ {
+ "UV_INDEX": "https://mirror.corp/simple",
+ "UV_EXTRA_INDEX_URL": "https://mirror.corp/extra",
+ "UV_INDEX_URL": "https://mirror.corp/root",
+ "UV_DEFAULT_INDEX": "https://mirror.corp/default",
+ },
+ ):
+ env = ips._install_env_for_cmd(cmd)
+ assert env is not None, "a --index-url install must run with a scrubbed env"
+ for var in self.UV_VARS:
+ assert var not in env, f"{var} must be cleared for a pinned-index install"
+
+ def test_pinned_default_index_strips_uv_index_vars(self):
+ # --default-index must be gated too (matches install.sh / install.ps1).
+ cmd = ["uv", "pip", "install", "torch", "--default-index", "https://x/cu126"]
+ with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}):
+ env = ips._install_env_for_cmd(cmd)
+ assert env is not None
+ assert "UV_INDEX" not in env
+
+ def test_non_pinned_install_keeps_user_mirror(self):
+ # A plain install (no --index-url) must NOT scrub the env, so a user's mirror
+ # still applies to base packages.
+ cmd = ["uv", "pip", "install", "unsloth", "unsloth-zoo"]
+ with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}):
+ env = ips._install_env_for_cmd(cmd)
+ assert env is None, "non-pinned installs must inherit the caller env unchanged"
+
+ def test_scrubbed_env_preserves_other_vars(self):
+ cmd = ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
+ with mock.patch.dict(
+ os.environ,
+ {"UV_INDEX": "https://mirror.corp/simple", "PATH_SENTINEL_XYZ": "keepme"},
+ ):
+ env = ips._install_env_for_cmd(cmd)
+ assert env is not None
+ assert env.get("PATH_SENTINEL_XYZ") == "keepme", "only uv index vars are removed"
+
+ def test_pinned_cmd_strips_pip_extra_index_url(self):
+ """PIP_EXTRA_INDEX_URL is stripped for pinned commands so the pip
+ fallback cannot satisfy torch from an inherited extra index."""
+ with mock.patch.dict(os.environ, {"PIP_EXTRA_INDEX_URL": "https://mirror/simple"}):
+ env = ips._install_env_for_cmd(
+ ["pip", "install", "torch", "--index-url", "https://x/cu128"]
+ )
+ assert env is not None and "PIP_EXTRA_INDEX_URL" not in env
+
+ def test_pinned_cmd_strips_uv_torch_backend(self):
+ """UV_TORCH_BACKEND is stripped for pinned commands so uv cannot read it
+ from the environment and reroute torch off the pinned index."""
+ with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
+ env = ips._install_env_for_cmd(
+ ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
+ )
+ assert env is not None and "UV_TORCH_BACKEND" not in env
+
+ def test_pinned_cmd_disables_uv_config_discovery(self):
+ """A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin too
+ (verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default
+ [[index]] both resolve torch+cpu against an explicit --index-url /
+ --default-index cu126 pin). Pinned commands must run with UV_NO_CONFIG=1
+ and without an inherited UV_CONFIG_FILE."""
+ with mock.patch.dict(os.environ, {"UV_CONFIG_FILE": "/etc/uv/uv.toml"}):
+ env = ips._install_env_for_cmd(
+ ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
+ )
+ assert env is not None
+ assert env.get("UV_NO_CONFIG") == "1"
+ assert "UV_CONFIG_FILE" not in env
+
+ def test_pinned_cmd_disables_pip_config_files(self):
+ """The pip FALLBACK honours user/site pip config files (pip config set
+ global.extra-index-url) even with the PIP_* env vars stripped; pip loads
+ NO configuration files when PIP_CONFIG_FILE is os.devnull. Harmless for
+ uv, decisive for the fallback."""
+ env = ips._install_env_for_cmd(
+ ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
+ )
+ assert env is not None
+ assert env.get("PIP_CONFIG_FILE") == os.devnull
+
+ def test_non_pinned_cmd_keeps_uv_config_discovery(self):
+ """Non-pinned installs inherit the caller env unchanged, so a user's uv
+ configuration still applies to base packages."""
+ env = ips._install_env_for_cmd(["uv", "pip", "install", "unsloth"])
+ assert env is None
diff --git a/tests/run_all.sh b/tests/run_all.sh
index d03f4c4d4f..a31103a85b 100755
--- a/tests/run_all.sh
+++ b/tests/run_all.sh
@@ -15,6 +15,7 @@ sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh"
sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh"
sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh"
sh "$TESTS_DIR/sh/test_torch_flavor.sh"
+sh "$TESTS_DIR/sh/test_redact_install_output.sh"
sh "$TESTS_DIR/sh/test_install_uv_override_space.sh"
echo ""
diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh
index 6656142625..23902097ef 100755
--- a/tests/sh/test_get_torch_index_url.sh
+++ b/tests/sh/test_get_torch_index_url.sh
@@ -23,6 +23,8 @@ _FAKE_SMI_DIR=$(mktemp -d)
echo ""
sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
+ sed -n '/^_trim_index_path_slashes()/,/^}/p' "$INSTALL_SH"
+ echo ""
sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH"
} | sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \
> "$_FUNC_FILE"
@@ -379,6 +381,61 @@ _result=$(run_func "$_dir" " -1 ")
assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
+# --- explicit overrides (headless / container / CI; no GPU probing) ----------
+# 39) UNSLOTH_TORCH_INDEX_FAMILY pins the family with no GPU present (not the cpu fallback).
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
+assert_eq "family override (no GPU) -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+
+# 40) Family override beats real detection: an nvidia-smi 12.6 host still gets cu128
+# (the Docker-build case -- builder sees the host driver but publishes a cu128 image).
+_dir=$(make_mock_smi "12.6")
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "$_dir")
+assert_eq "family override beats detected 12.6 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+rm -rf "$_dir"
+
+# 41) UNSLOTH_TORCH_INDEX_URL is used verbatim and wins over detection.
+_dir=$(make_mock_smi "12.6")
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu999" run_func "$_dir")
+assert_eq "url override beats detection -> verbatim" "https://mirror.example.com/whl/cu999" "$_result"
+rm -rf "$_dir"
+
+# 42) Family override is appended to UNSLOTH_PYTORCH_MIRROR (mirror still honoured).
+_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
+assert_eq "mirror + family override -> mirror/cu128" "https://mirror.example.com/whl/cu128" "$_result"
+
+# 43) Trailing slash in UNSLOTH_TORCH_INDEX_URL is stripped.
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128/" run_func "none")
+assert_eq "url override trailing slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
+
+# 44) URL override takes precedence over family override.
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu130" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
+assert_eq "url override beats family override -> url" "https://mirror.example.com/whl/cu130" "$_result"
+
+# 45) An empty override is ignored (falls through to normal detection).
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="" UNSLOTH_TORCH_INDEX_URL="" run_func "none")
+assert_eq "empty overrides ignored -> detected cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+
+# 46) ALL trailing slashes are stripped from a URL override (not just one).
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128///" run_func "none")
+assert_eq "url override double slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
+
+# 47) Leading and trailing slashes stripped from a family override.
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="//cu128//" run_func "none")
+assert_eq "family override slashes stripped" "https://download.pytorch.org/whl/cu128" "$_result"
+
+# 48) A ?query token that ends in "/" is PRESERVED: only PATH slashes are trimmed, so a
+# base64 token ending in "/" is not corrupted (path-only trim, not whole-URL rstrip).
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128?token=ab12cd/" run_func "none")
+assert_eq "url override preserves query token slash" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result"
+
+# 49) Double PATH slash before a query is collapsed while the query survives intact.
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128//?token=ab12cd/" run_func "none")
+assert_eq "url override path slash trimmed, query kept" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result"
+
+# 50) A #fragment ending in "/" is likewise preserved.
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128#anchor/" run_func "none")
+assert_eq "url override preserves fragment slash" "https://mirror.example.com/whl/cu128#anchor/" "$_result"
+
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"
diff --git a/tests/sh/test_redact_install_output.sh b/tests/sh/test_redact_install_output.sh
new file mode 100755
index 0000000000..0f10122aea
--- /dev/null
+++ b/tests/sh/test_redact_install_output.sh
@@ -0,0 +1,89 @@
+#!/bin/bash
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+# Unit tests for install.sh's _redact_install_output helper. uv/pip failure text embeds the
+# failing --index-url verbatim, so a captured install log dumped on error can leak a
+# user:token@ or ?token= secret. The helper redacts both before printing. Mirrors
+# _redact_install_output (install_python_stack.py) / Redact-InstallOutput (install.ps1 /
+# setup.ps1).
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+INSTALL_SH="$SCRIPT_DIR/../../install.sh"
+PASS=0
+FAIL=0
+
+_FUNC_FILE=$(mktemp)
+sed -n '/^_redact_install_output()/,/^}/p' "$INSTALL_SH" > "$_FUNC_FILE"
+# shellcheck disable=SC1090
+. "$_FUNC_FILE"
+rm -f "$_FUNC_FILE"
+
+assert_eq() {
+ _label="$1"; _expected="$2"; _actual="$3"
+ if [ "$_actual" = "$_expected" ]; then
+ echo " PASS: $_label"; PASS=$((PASS + 1))
+ else
+ echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1))
+ fi
+}
+
+# Redact from a file (the actual call site passes a captured-log tempfile).
+redact_str() {
+ _rs_tmp=$(mktemp)
+ printf '%s\n' "$1" > "$_rs_tmp"
+ _rs_out=$(_redact_install_output "$_rs_tmp")
+ rm -f "$_rs_tmp"
+ printf '%s' "$_rs_out"
+}
+
+echo "=== _redact_install_output ==="
+assert_eq "userinfo user:token@ redacted" \
+ "ERROR: failed https://@download.pytorch.org/whl/cu128" \
+ "$(redact_str 'ERROR: failed https://alice:s3cr3t@download.pytorch.org/whl/cu128')"
+
+assert_eq "bare-token@ userinfo redacted" \
+ "fetch https://@host/whl/cu128 failed" \
+ "$(redact_str 'fetch https://ghp_deadbeef@host/whl/cu128 failed')"
+
+assert_eq "single ?token= query redacted" \
+ "url https://host/whl/cu128?token= unreachable" \
+ "$(redact_str 'url https://host/whl/cu128?token=abcd1234 unreachable')"
+
+assert_eq "multiple query values redacted" \
+ "https://host/whl/cu128?token=&channel=" \
+ "$(redact_str 'https://host/whl/cu128?token=abcd1234&channel=beta')"
+
+assert_eq "http (not https) userinfo redacted" \
+ "http://@host/simple" \
+ "$(redact_str 'http://u:p@host/simple')"
+
+assert_eq "fragment token redacted" \
+ "ERROR: could not fetch https://mirror.local/whl/cu128# (403)" \
+ "$(redact_str 'ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)')"
+
+assert_eq "query and fragment both redacted" \
+ "https://host/whl/cu128?token=# done" \
+ "$(redact_str 'https://host/whl/cu128?token=abc#sig=xyz done')"
+
+# Non-secret text is untouched (no false positives on ordinary log lines).
+assert_eq "plain line untouched" \
+ "Resolved 42 packages in 1.2s" \
+ "$(redact_str 'Resolved 42 packages in 1.2s')"
+assert_eq "plain url without creds untouched" \
+ "downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl" \
+ "$(redact_str 'downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl')"
+assert_eq "bare hash comment untouched" \
+ "# retrying with --no-cache-dir" \
+ "$(redact_str '# retrying with --no-cache-dir')"
+
+# Regression guard: no secret substring survives.
+_leak=$(redact_str 'https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET')
+case "$_leak" in
+ *s3cr3t*|*SUPERSECRET*|*ALSOSECRET*) assert_eq "no secret leak" "clean" "leaked:$_leak" ;;
+ *) assert_eq "no secret leak" "clean" "clean" ;;
+esac
+
+echo ""
+echo "Results: $PASS passed, $FAIL failed"
+[ "$FAIL" -eq 0 ]
diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh
index d60dfc9f90..bfafbd161b 100644
--- a/tests/sh/test_torch_constraint.sh
+++ b/tests/sh/test_torch_constraint.sh
@@ -108,6 +108,25 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var"
_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded"
+# Companions must be bounded to torch's window everywhere: the <2.11 bound appears
+# twice (default assignments + the pinned custom-leaf block), never bare. torchaudio
+# 2.11 dropped its exact torch pin, so a bare companion next to a <2.11-capped torch
+# resolves a mismatched 2.11 build.
+_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"' "$INSTALL_SH" || true)
+assert_eq "torchvision bounded (<0.26) at default + custom-leaf" "2" "$_count"
+_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"' "$INSTALL_SH" || true)
+assert_eq "torchaudio bounded (<2.11) at default + custom-leaf" "2" "$_count"
+_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision"$' "$INSTALL_SH" || true)
+assert_eq "no bare torchvision companion remains" "0" "$_count"
+_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio"$' "$INSTALL_SH" || true)
+assert_eq "no bare torchaudio companion remains" "0" "$_count"
+# The cu* widen must carry the companions with it (torch <2.12 with torchaudio <2.11
+# would cap a mismatched pair the other way).
+assert_eq "cu widen pairs torchaudio (<2.12)" "1" "$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"' "$INSTALL_SH" || true)"
+_gated=$(grep -c '_expected_torch_flavor_tag "$TORCH_INDEX_URL"' "$INSTALL_SH" || true)
+_has_gate=$([ "$_gated" -ge 1 ] && echo "yes" || echo "no")
+assert_eq "custom-companion bound gated on empty flavor tag" "yes" "$_has_gate"
+
# A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x land torch
# 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC).
_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true)
@@ -285,6 +304,61 @@ bash -c "
_uv_got2=$(cat "$_UV_LOG2" 2>/dev/null || echo "")
assert_contains "mock uv arm64+py312 receives torch>=2.4" "$_uv_got2" "torch>=2.4,<2.11.0"
+# ======================================================================
+# ROCm 2.11 floor: leaf is lowercased before the gfx*/rocm* allowlist match
+# ======================================================================
+echo ""
+echo "=== ROCm 2.11 floor case (leaf normalization) ==="
+
+# Structural: install.sh lowercases _torch_index_leaf before the floor case, so the
+# canonical gfx120X-all (capital X) matches gfx120x-all.
+_has_lc=$(grep -c '_torch_index_leaf=$(printf .* | tr .\[:upper:\]. .\[:lower:\].)' "$INSTALL_SH" || true)
+_has_lc_ok=$([ "$_has_lc" -ge 1 ] && echo "yes" || echo "no")
+assert_eq "install.sh lowercases _torch_index_leaf" "yes" "$_has_lc_ok"
+
+# Runtime: replicate install.sh's normalization + floor case and assert both gfx120X-all
+# and gfx120x-all get the floor, while non-2.11 leaves keep the default.
+run_floor_case() {
+ _url="$1"
+ bash -c '
+ TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+ TORCHVISION_CONSTRAINT="torchvision"
+ TORCHAUDIO_CONSTRAINT="torchaudio"
+ _torch_index_leaf="${1%/}"
+ _torch_index_leaf="${_torch_index_leaf##*/}"
+ _torch_index_leaf=$(printf "%s" "$_torch_index_leaf" | tr "[:upper:]" "[:lower:]")
+ case "$_torch_index_leaf" in
+ rocm7.2|gfx120x-all|gfx1151|gfx1150)
+ TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
+ TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
+ ;;
+ esac
+ echo "$TORCH_CONSTRAINT"
+ ' _ "$_url"
+}
+
+assert_eq "gfx120X-all (capital) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
+ "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all')"
+assert_eq "gfx120X-all trailing slash -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
+ "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all/')"
+assert_eq "gfx120x-all (lowercase) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
+ "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120x-all')"
+assert_eq "gfx1151 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
+ "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1151')"
+assert_eq "gfx1150 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
+ "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1150')"
+assert_eq "rocm7.2 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
+ "$(run_floor_case 'https://download.pytorch.org/whl/rocm7.2')"
+assert_eq "gfx110X-all -> default (no floor)" "torch>=2.4,<2.11.0" \
+ "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx110X-all')"
+assert_eq "rocm6.4 -> default (no floor)" "torch>=2.4,<2.11.0" \
+ "$(run_floor_case 'https://download.pytorch.org/whl/rocm6.4')"
+assert_eq "cu128 -> default (no floor)" "torch>=2.4,<2.11.0" \
+ "$(run_floor_case 'https://download.pytorch.org/whl/cu128')"
+assert_eq "cpu -> default (no floor)" "torch>=2.4,<2.11.0" \
+ "$(run_floor_case 'https://download.pytorch.org/whl/cpu')"
+
# ======================================================================
# Summary
# ======================================================================
diff --git a/tests/sh/test_torch_flavor.sh b/tests/sh/test_torch_flavor.sh
index ead2c4164f..55da2c0f07 100755
--- a/tests/sh/test_torch_flavor.sh
+++ b/tests/sh/test_torch_flavor.sh
@@ -11,14 +11,21 @@ INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
-# Extract the three helper functions from install.sh and source them.
+# Extract the helper functions from install.sh and source them
+# (_torch_index_url_leaf is the shared leaf extractor the classifiers call).
_FUNC_FILE=$(mktemp)
{
sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
+ sed -n '/^_torch_index_url_leaf()/,/^}/p' "$INSTALL_SH"
+ echo ""
+ sed -n '/^_is_pip_rocm_family_leaf()/,/^}/p' "$INSTALL_SH"
+ echo ""
sed -n '/^_expected_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_torch_index_repairable()/,/^}/p' "$INSTALL_SH"
+ echo ""
+ sed -n '/^_tauri_torch_index_family()/,/^}/p' "$INSTALL_SH"
} > "$_FUNC_FILE"
# shellcheck disable=SC1090
. "$_FUNC_FILE"
@@ -56,6 +63,26 @@ assert_eq "amd gfx index" "rocm" "$(_expected_torch_flavor_tag 'https://re
assert_eq "mirror cu130 leaf" "cu130" "$(_expected_torch_flavor_tag 'https://my.mirror/pytorch/whl/cu130')"
assert_eq "unrecognized leaf" "" "$(_expected_torch_flavor_tag 'https://my.mirror/whl/simple')"
assert_eq "empty url" "" "$(_expected_torch_flavor_tag '')"
+# Query/fragment dropped before classification: .../cu128?token=x classifies as cu128,
+# not an opaque leaf that reinstalls every run.
+assert_eq "query-bearing cu128" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128?token=x')"
+assert_eq "fragment-bearing cpu" "cpu" "$(_expected_torch_flavor_tag 'https://m/whl/cpu#frag')"
+# A cu-suffixed CUSTOM leaf (cu128-private, cu128x) is NOT the cu128 family (exact
+# cu+digits only). Mirrors Python re.fullmatch(cu[0-9]+) / PowerShell.
+assert_eq "cu-suffix custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128-private')"
+assert_eq "cu-alnum custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128x')"
+assert_eq "bare cu digits stays" "cu126" "$(_expected_torch_flavor_tag 'https://m/whl/cu126')"
+# A custom leaf merely STARTING with rocm (rocm-current, rocm-rel-7.2.1) is NOT a pip
+# rocm family -> "" (custom); real families (rocm7.2) and gfx indexes stay "rocm".
+assert_eq "custom rocm-current" "" "$(_expected_torch_flavor_tag 'https://mirror/whl/rocm-current')"
+assert_eq "radeon rocm-rel leaf" "" "$(_expected_torch_flavor_tag 'https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1')"
+assert_eq "real rocm7.2 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7.2')"
+# A rocm-SUFFIX private mirror (rocm7.2-private, rocm7-current) is a custom pin ->
+# "" (custom); match the family exactly, not the prefix.
+assert_eq "suffixed rocm7.2-private" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2-private')"
+assert_eq "suffixed rocm7-current" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7-current')"
+assert_eq "two-dot rocm7.2.1" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2.1')"
+assert_eq "bare rocm7 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7')"
echo "=== _torch_index_repairable ==="
assert_eq "cu130 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cu130')"
@@ -64,6 +91,66 @@ assert_eq "gfx repairable" "yes" "$(_torch_index_repairable 'https://repo.
assert_eq "gfx1151 repairable" "yes" "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx1151/')"
assert_eq "cpu NOT repairable" "no" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cpu')"
assert_eq "unknown NOT repair" "no" "$(_torch_index_repairable 'https://my.mirror/whl/simple')"
+# A suffixed rocm leaf is a verbatim pin, not a --default-index repairable family.
+assert_eq "rocm-private NOT repair" "no" "$(_torch_index_repairable 'https://co.internal/whl/rocm7.2-private')"
+
+echo "=== _is_pip_rocm_family_leaf ==="
+assert_family() {
+ _label="$1"; _expected="$2"; _leaf="$3"
+ if _is_pip_rocm_family_leaf "$_leaf"; then _actual="yes"; else _actual="no"; fi
+ assert_eq "$_label" "$_expected" "$_actual"
+}
+assert_family "rocm7.2 family" "yes" "rocm7.2"
+assert_family "rocm6.4 family" "yes" "rocm6.4"
+assert_family "bare rocm7 family" "yes" "rocm7"
+assert_family "gfx120x-all family" "yes" "gfx120x-all"
+assert_family "gfx1151 family" "yes" "gfx1151"
+assert_family "rocm7.2-private custom" "no" "rocm7.2-private"
+assert_family "rocm7-current custom" "no" "rocm7-current"
+assert_family "rocm-current custom" "no" "rocm-current"
+assert_family "rocm-rel-7.2.1 custom" "no" "rocm-rel-7.2.1"
+assert_family "rocm7.2.1 custom" "no" "rocm7.2.1"
+# A trailing dot (rocm7.) or leading/double dot is NOT a family: both major and minor must
+# be non-empty all-digits, matching Python re.fullmatch(rocm\d+(?:\.\d+)?). Bash previously
+# accepted rocm7. via a bare %/-style trim while Python rejected it (validator asymmetry).
+assert_family "rocm7. trailing-dot custom" "no" "rocm7."
+assert_family "rocm.7 leading-dot custom" "no" "rocm.7"
+assert_family "rocm7..2 double-dot custom" "no" "rocm7..2"
+assert_family "cpu not rocm" "no" "cpu"
+assert_family "cu128 not rocm" "no" "cu128"
+assert_family "simple not rocm" "no" "simple"
+
+echo "=== _torch_index_url_leaf (ALL trailing slashes stripped -> non-empty leaf) ==="
+# A double (or triple) trailing slash must yield the real leaf, not an empty string that
+# fails every classifier arm. Python .rstrip("/") drops them all; bash must match (a bare
+# %/ left .../cu128// classifying as "").
+assert_eq "double slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//')"
+assert_eq "triple slash rocm7.2 leaf" "rocm7.2" "$(_torch_index_url_leaf 'https://m/whl/rocm7.2///')"
+assert_eq "double slash + token leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//?token=x')"
+assert_eq "single slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128/')"
+# The classifier that consumes the leaf must therefore still tag a double-slash index.
+assert_eq "double-slash cu128 tag" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128//')"
+assert_eq "double-slash rocm7.2 tag" "rocm" "$(_expected_torch_flavor_tag 'https://m/whl/rocm7.2//')"
+
+echo "=== _tauri_torch_index_family (credential redaction) ==="
+# A token/fragment must be stripped BEFORE classification so it never reaches the
+# [TAURI:DIAG] line (the family is the last path segment, which else carries the query).
+SKIP_TORCH=false
+assert_eq "token stripped from rocm" "rocm7.2" "$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET')"
+assert_eq "token-bearing cu classifies" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128?token=x')"
+assert_eq "fragment stripped cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu#frag')"
+assert_eq "plain rocm7.2 unchanged" "rocm7.2" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/rocm7.2')"
+# A trailing slash must be stripped too, or the */cu128 and */cpu arms miss .../cu128/
+# and it falls through to "auto".
+assert_eq "trailing slash cu128" "cu128" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/cu128/')"
+assert_eq "slash + token cu128" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128/?token=x')"
+assert_eq "trailing slash cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu/')"
+# Regression guard: no secret token substring may survive in any classification.
+_leak=$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET')
+case "$_leak" in
+ *SECRET*|*token*) assert_eq "no token leak in family" "clean" "leaked:$_leak" ;;
+ *) assert_eq "no token leak in family" "clean" "clean" ;;
+esac
echo ""
echo "Results: $PASS passed, $FAIL failed"
diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py
index cea4383268..c6d2b95316 100644
--- a/tests/studio/install/test_cuda_repair.py
+++ b/tests/studio/install/test_cuda_repair.py
@@ -64,15 +64,23 @@ def _run_cuda_repair(
rocm_marker = False,
smi_path = "/usr/bin/nvidia-smi",
cvd = None,
+ index_family = None,
+ index_url = None,
):
"""Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock.
- cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it."""
+ cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it.
+ index_family sets UNSLOTH_TORCH_INDEX_FAMILY (the explicit wheel-index pin).
+ index_url sets UNSLOTH_TORCH_INDEX_URL (the full-URL pin form)."""
env = {}
if rocm_marker:
env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
if cvd is not None:
env["CUDA_VISIBLE_DEVICES"] = cvd
+ if index_family is not None:
+ env["UNSLOTH_TORCH_INDEX_FAMILY"] = index_family
+ if index_url is not None:
+ env["UNSLOTH_TORCH_INDEX_URL"] = index_url
def _which(name, *a, **k):
if name == "nvidia-smi":
@@ -99,6 +107,10 @@ def _run_cuda_repair(
stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None)
if cvd is None:
stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ if index_family is None:
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ if index_url is None:
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
_ensure_cuda_torch()
return mock_pip
@@ -123,11 +135,74 @@ class TestCudaRepairFires:
assert mock_pip.call_args.kwargs["constrain"] is False
def test_rocm_in_version_string_triggers_repair(self):
- # AMD SDK / Radeon wheels may encode rocm in __version__ without
- # torch.version.hip; the probe prints "hip" for both.
+ # AMD SDK / Radeon wheels may encode rocm in __version__ without torch.version.hip;
+ # the probe prints "hip" for both.
mock_pip = _run_cuda_repair(torch_state = "hip")
assert mock_pip.call_count == 1
+ def test_no_gpu_but_explicit_cuda_pin_repairs(self):
+ # Headless / CI cross-install: an explicit cu* pin commits to CUDA wheels with no
+ # NVIDIA GPU visible, so a ROCm-poisoned venv is still repaired to the pinned family.
+ mock_pip = _run_cuda_repair(
+ nvidia = False,
+ backend = "cuda",
+ index_family = "cu128",
+ torch_state = "hip",
+ )
+ assert mock_pip.call_count == 1
+ assert "cu128" in _index_url(mock_pip)
+
+ def test_cvd_hidden_but_explicit_cuda_pin_repairs(self):
+ # CVD=-1/"" hides the GPU, but an explicit cu* pin skips ALL host-GPU probing, so the
+ # CVD hide gate must not suppress the repair (GPU-less CI: CVD=-1, FAMILY=cu128).
+ for _cvd in ("-1", ""):
+ mock_pip = _run_cuda_repair(
+ nvidia = False,
+ backend = "cuda",
+ cvd = _cvd,
+ index_family = "cu128",
+ torch_state = "hip",
+ )
+ assert mock_pip.call_count == 1
+ assert "cu128" in _index_url(mock_pip)
+
+ def test_tagged_cuda_mismatch_repairs(self):
+ # A healthy CUDA torch whose +cuXXX differs from the pin is repaired.
+ mock_pip = _run_cuda_repair(
+ index_family = "cu128",
+ torch_state = "cuda|cu126",
+ cuda_version = "12.8",
+ )
+ assert mock_pip.call_count == 1
+ assert "cu128" in _index_url(mock_pip)
+
+ def test_untagged_cuda_build_under_pin_repairs(self):
+ # An untagged CUDA build (no +cuXXX tag -> empty installed cu) can't be confirmed
+ # to match the pin, so the pin is enforced with a reinstall.
+ mock_pip = _run_cuda_repair(
+ index_family = "cu128",
+ torch_state = "cuda", # marker cuda, empty installed cu
+ cuda_version = "12.8",
+ )
+ assert mock_pip.call_count == 1
+ assert "cu128" in _index_url(mock_pip)
+
+ def test_broken_probe_with_cuda_pin_repairs(self):
+ # torch present but unimportable under a CUDA pin: the base update won't repair a
+ # broken already-installed torch, so reinstall from the pin instead of stranding it.
+ mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1, index_family = "cu128")
+ assert mock_pip.call_count == 1
+ assert "cu128" in _index_url(mock_pip)
+
+ def test_broken_probe_with_cuda_url_pin_repairs(self):
+ mock_pip = _run_cuda_repair(
+ torch_state = "cpu",
+ torch_rc = 1,
+ index_url = "https://mirror.local/cu128",
+ )
+ assert mock_pip.call_count == 1
+ assert "https://mirror.local/cu128" in _index_url(mock_pip)
+
# No-op cases.
@@ -157,8 +232,9 @@ class TestCudaRepairSkips:
mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip")
mock_pip.assert_not_called()
- def test_torch_missing_skips(self):
- # Non-zero probe exit = torch missing / un-importable.
+ def test_torch_missing_no_pin_skips(self):
+ # Non-zero probe exit = torch missing/un-importable. With NO CUDA pin the base
+ # install owns it, so leave it alone (a pinned build reinstalls).
mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1)
mock_pip.assert_not_called()
@@ -191,6 +267,109 @@ class TestCudaRepairSkips:
mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip")
assert mock_pip.call_count == 1
+ def test_matching_tagged_cuda_pin_no_repair(self):
+ # Healthy CUDA torch whose +cuXXX already matches the pin: no reinstall.
+ mock_pip = _run_cuda_repair(
+ index_family = "cu128",
+ torch_state = "cuda|cu128",
+ cuda_version = "12.8",
+ )
+ mock_pip.assert_not_called()
+
+ def test_custom_mirror_leaf_not_treated_as_cuda_pin(self):
+ # A mirror leaf starting with "cu" but not cuXXX (.../custom, .../current) must
+ # NOT be treated as a CUDA pin, so it can't bypass the NVIDIA gate.
+ for _leaf in ("custom", "current"):
+ mock_pip = _run_cuda_repair(
+ nvidia = False,
+ backend = "cuda",
+ index_url = f"https://mymirror.example/{_leaf}",
+ torch_state = "hip",
+ )
+ mock_pip.assert_not_called()
+
+ def test_explicit_cuda_family_leaf_helper(self):
+ # _explicit_cuda_torch_index_url matches cuXXX narrowly, not any cu* leaf.
+ import contextlib
+
+ def _with(url):
+ with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ return stack_mod._explicit_cuda_torch_index_url()
+
+ assert _with("https://download.pytorch.org/whl/cu128") is not None
+ assert _with("https://download.pytorch.org/whl/cu126") is not None
+ assert _with("https://mymirror.example/custom") is None
+ assert _with("https://mymirror.example/current") is None
+ assert _with("https://download.pytorch.org/whl/cpu") is None
+ with contextlib.suppress(Exception):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
+
+
+class TestTorchBackendDerivationFromPin:
+ """The module-level _TORCH_BACKEND derivation (standalone `studio update`
+ with no install.sh-set UNSLOTH_TORCH_BACKEND) must classify the pinned index
+ leaf via _is_cuda_family_leaf (^cu[0-9]), NOT a bare startswith("cu"). A
+ full-override URL ending in /current or /custom must fall through to backend
+ "" (probe the GPU) so _ensure_rocm_torch() still repairs a wrong/CPU torch on
+ AMD hosts, instead of being wrongly branded "cuda" and returning early."""
+
+ @staticmethod
+ def _derive(env):
+ # Re-run the module's import-time derivation, using its own _is_cuda_family_leaf
+ # so this stays in lockstep.
+ idx_override = (
+ env.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
+ or env.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
+ )
+ backend = env.get("UNSLOTH_TORCH_BACKEND", "").lower()
+ if not backend:
+ leaf = idx_override.rstrip("/").rsplit("/", 1)[-1].lower()
+ if leaf.startswith(("rocm", "gfx")):
+ backend = "rocm"
+ elif leaf == "cpu":
+ backend = "cpu"
+ elif stack_mod._is_cuda_family_leaf(leaf):
+ backend = "cuda"
+ return backend
+
+ def test_cu128_pin_is_cuda(self):
+ assert (
+ self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://download.pytorch.org/whl/cu128"})
+ == "cuda"
+ )
+
+ def test_cu128_family_is_cuda(self):
+ assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cu128"}) == "cuda"
+
+ def test_current_leaf_not_cuda(self):
+ # ^cu[0-9] rejects /current -> backend stays "" (probe GPU), so an AMD host still
+ # repairs a CPU/wrong torch instead of short-circuiting.
+ assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/current"}) == ""
+
+ def test_custom_leaf_not_cuda(self):
+ assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/custom"}) == ""
+
+ def test_rocm_and_gfx_pins_are_rocm(self):
+ assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}) == "rocm"
+ assert (
+ self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"})
+ == "rocm"
+ )
+
+ def test_cpu_pin_is_cpu(self):
+ assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cpu"}) == "cpu"
+
+ def test_source_uses_helper_not_bare_startswith(self):
+ # Guard against a regression back to elif _idx_leaf.startswith("cu").
+ src = _STACK_PATH.read_text(encoding = "utf-8")
+ assert (
+ "elif _is_cuda_family_leaf(_idx_leaf):" in src
+ ), "_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf"
+ assert (
+ 'elif _idx_leaf.startswith("cu"):' not in src
+ ), "_TORCH_BACKEND derivation must not use a bare startswith('cu')"
+
# CUDA index ladder.
diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py
index f5ad9566d7..d2fd7ae8db 100644
--- a/tests/studio/install/test_gpu_detection_followups.py
+++ b/tests/studio/install/test_gpu_detection_followups.py
@@ -284,7 +284,9 @@ class TestBackendExportLeafClassification:
def test_export_block_uses_leaf(self, install_src):
anchor = install_src.find("_torch_index_leaf=")
assert anchor >= 0, "backend export must classify on the final path segment"
- window = install_src[anchor : anchor + 500]
+ # Window spans the leaf-normalization prelude (query/frag drop + all-slash trim loop)
+ # through the export case arms.
+ window = install_src[anchor : anchor + 900]
assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window
assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window
assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window
@@ -459,3 +461,130 @@ class TestHiddenCvdNotUsable:
cvd,
)
assert out == expected
+
+
+class TestRedactInstallOutput:
+ """_redact_install_output scrubs index-URL credentials from a captured install log
+ before it is printed on failure (uv/pip embeds the failing --index-url verbatim)."""
+
+ def test_userinfo_redacted(self):
+ out = stack_mod._redact_install_output(
+ "ERROR: failed https://alice:s3cr3t@download.pytorch.org/whl/cu128"
+ )
+ assert out == "ERROR: failed https://@download.pytorch.org/whl/cu128"
+
+ def test_bytes_input_decoded_and_redacted(self):
+ out = stack_mod._redact_install_output(b"fetch https://ghp_deadbeef@host/whl/cu128 failed")
+ assert out == "fetch https://@host/whl/cu128 failed"
+
+ def test_query_values_redacted(self):
+ out = stack_mod._redact_install_output(
+ "url https://host/whl/cu128?token=abcd1234&channel=beta unreachable"
+ )
+ assert out == "url https://host/whl/cu128?token=&channel= unreachable"
+
+ def test_fragment_redacted(self):
+ out = stack_mod._redact_install_output(
+ "ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)"
+ )
+ assert out == "ERROR: could not fetch https://mirror.local/whl/cu128# (403)"
+
+ def test_query_and_fragment_both_redacted(self):
+ out = stack_mod._redact_install_output("https://host/whl/cu128?token=abc#sig=xyz done")
+ assert out == "https://host/whl/cu128?token=# done"
+
+ def test_bare_hash_comment_untouched(self):
+ # The fragment redaction is URL-anchored: a shell comment in tool output survives.
+ assert (
+ stack_mod._redact_install_output("# retrying with --no-cache-dir")
+ == "# retrying with --no-cache-dir"
+ )
+
+ def test_plain_line_untouched(self):
+ assert (
+ stack_mod._redact_install_output("Resolved 42 packages in 1.2s")
+ == "Resolved 42 packages in 1.2s"
+ )
+
+ def test_no_secret_substring_survives(self):
+ out = stack_mod._redact_install_output(
+ "https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET"
+ )
+ assert "s3cr3t" not in out and "SUPERSECRET" not in out and "ALSOSECRET" not in out
+
+
+class TestTrimIndexPathSlashes:
+ """_trim_index_path_slashes strips trailing PATH slashes only; a ?query/#fragment token
+ ending in "/" must survive (a whole-URL rstrip would corrupt a base64 token)."""
+
+ def test_double_path_slash_collapsed(self):
+ assert stack_mod._trim_index_path_slashes("https://h/whl/cu128//") == "https://h/whl/cu128"
+
+ def test_query_token_slash_preserved(self):
+ assert (
+ stack_mod._trim_index_path_slashes("https://h/whl/cu128?token=ab12cd/")
+ == "https://h/whl/cu128?token=ab12cd/"
+ )
+
+ def test_path_slash_trimmed_query_kept(self):
+ assert (
+ stack_mod._trim_index_path_slashes("https://h/whl/cu128//?token=ab12cd/")
+ == "https://h/whl/cu128?token=ab12cd/"
+ )
+
+ def test_fragment_slash_preserved(self):
+ assert (
+ stack_mod._trim_index_path_slashes("https://h/whl/cu128#anchor/")
+ == "https://h/whl/cu128#anchor/"
+ )
+
+
+class TestRocmFamilyLeafParity:
+ """_is_pip_rocm_family_leaf must match re.fullmatch(rocm\\d+(?:\\.\\d+)?): a trailing dot
+ (rocm7.) is a CUSTOM pin, not a family (the historical bash/py validator asymmetry)."""
+
+ @pytest.mark.parametrize(
+ "leaf, expected",
+ [
+ ("rocm7", True),
+ ("rocm7.2", True),
+ ("gfx1151", True),
+ ("rocm7.", False),
+ ("rocm.7", False),
+ ("rocm7..2", False),
+ ("rocm7.2.1", False),
+ ("rocm7.2-private", False),
+ ("cpu", False),
+ ("cu128", False),
+ ],
+ )
+ def test_family_classification(self, leaf, expected):
+ assert stack_mod._is_pip_rocm_family_leaf(leaf) is expected
+
+
+class TestTorchIndexLeafAllSlashes:
+ """_torch_index_leaf drops query/fragment then strips ALL trailing slashes, so a
+ double-slash index still yields the real leaf (not an empty string)."""
+
+ @pytest.mark.parametrize(
+ "url, expected",
+ [
+ ("https://m/whl/cu128//", "cu128"),
+ ("https://m/whl/rocm7.2///", "rocm7.2"),
+ ("https://m/whl/cu128//?token=x", "cu128"),
+ ("https://m/whl/cu128/", "cu128"),
+ ],
+ )
+ def test_leaf_never_empty_on_double_slash(self, url, expected):
+ assert stack_mod._torch_index_leaf(url) == expected
+
+
+class TestUvIndexEnvVarsScrub:
+ """The pinned-install env scrub must drop PIP_NO_INDEX (which makes the pip fallback
+ ignore ALL indexes, defeating the pin) and PIP_INDEX_URL (replaces the pinned index)."""
+
+ def test_pip_no_index_scrubbed(self):
+ assert "PIP_NO_INDEX" in stack_mod._UV_INDEX_ENV_VARS
+
+ def test_pip_index_url_scrubbed(self):
+ assert "PIP_INDEX_URL" in stack_mod._UV_INDEX_ENV_VARS
diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py
index d6dc8b2f8e..dd2a7ec487 100644
--- a/tests/studio/install/test_pr5940_followups.py
+++ b/tests/studio/install/test_pr5940_followups.py
@@ -784,7 +784,7 @@ def test_install_python_stack_windows_rocm_repair_pins_and_is_nonfatal():
assert re.search(
r'"' + gfx + r'":\s*_ROCM_TORCH_PKG_SPECS\["rocm7\.2"\]', text
), f"{gfx} must pin to the rocm7.2 trio like install.ps1/setup.ps1"
- i = text.find('f"ROCm torch (Windows, {gfx_arch})"')
+ i = text.find("f\"ROCm torch (Windows, {gfx_arch or 'pinned'})\"")
assert i != -1, "Windows ROCm repair pip call not found"
# The nearest preceding call must be the nonfatal pip_install_try, not pip_install.
j = text.rfind("pip_install_try(", 0, i)
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index e7ac0ec82d..b94578369d 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -569,19 +569,27 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
- def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
- """If torch already has CUDA, should skip ROCm reinstall."""
+ def test_cuda_torch_on_amd_host_reinstalls(
+ self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """A CUDA-only torch build is unusable on an AMD-only host, so it must be
+ reinstalled to ROCm (has_hip_torch is driven by the empty HIP marker, not
+ by treating the CUDA version string as a HIP marker)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
- mock_probe.stdout = b"12.6\n" # CUDA version
+ # Single-line probe: empty HIP marker before "|" for a CUDA build.
+ mock_probe.stdout = b"|2.10.0+cu126\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
- mock_pip.assert_not_called()
+ assert mock_pip.call_count == 1
+ assert "rocm7.1" in str(mock_pip.call_args_list[0])
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@@ -591,12 +599,31 @@ class TestEnsureRocmTorch:
"""If torch already has HIP, should skip ROCm reinstall."""
mock_probe = MagicMock()
mock_probe.returncode = 0
- mock_probe.stdout = b"7.1.12345\n" # HIP version
+ mock_probe.stdout = b"7.1.12345|2.10.0+rocm7.1\n" # HIP marker + version
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
mock_pip.assert_not_called()
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
+ def test_cpu_torch_probe_line_not_read_as_hip(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
+ """A CPU build's probe line ("|2.10.0+cpu") must not read as HIP: the version
+ after the "|" separator is data, not a HIP marker, so has_hip_torch stays False
+ and the reinstall fires."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"|2.10.0+cpu\n"
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ with patch.object(stack_mod, "pip_install_try", return_value = True):
+ _ensure_rocm_torch()
+ assert mock_pip.call_count == 1
+ assert "rocm7.1" in str(mock_pip.call_args_list[0])
+
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@@ -680,6 +707,295 @@ class TestEnsureRocmTorch:
torch_call = mock_pip.call_args_list[0]
assert "rocm7.2" in str(torch_call)
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
+ def test_explicit_gfx_index_honored_and_skips_strix_reroute(
+ self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """An explicit gfx wheel-index pin is authoritative: install from it verbatim
+ with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4
+ would otherwise pick the rocm6.4 wheel / trigger the Strix re-route)."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"\n" # cpu torch -> reinstall
+ env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
+ with patch.dict(stack_mod.os.environ, env, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ # Would raise if the Strix block ran (it is skipped on an explicit pin).
+ with patch.object(
+ stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
+ ):
+ _ensure_rocm_torch()
+ assert mock_pip.call_count == 1
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "gfx1151" in torch_call
+ assert "torch>=2.11.0,<2.12.0" in torch_call
+
+ def test_rocm_pin_family_mismatch_helper(self):
+ """_rocm_pin_family_mismatch: exact rocm compare, else the 2.11 line."""
+ f = stack_mod._rocm_pin_family_mismatch
+ base = "https://download.pytorch.org/whl"
+ amd = "https://repo.amd.com/rocm/whl"
+ # Exact rocm version comparison.
+ assert f(f"{base}/rocm7.2", "2.11.0+rocm7.2") is False
+ assert f(f"{base}/rocm7.2", "2.10.0+rocm6.4") is True
+ assert f(f"{base}/rocm6.4", "2.10.0+rocm6.4") is False
+ # rocm7.2 is KNOWN-2.11. A +rocm7.2 wheel whose RELEASE drifted off 2.11 shares the
+ # tag but violates the spec -> mismatch (a plain version compare would accept it).
+ assert f(f"{base}/rocm7.2", "2.12.0+rocm7.2") is True
+ assert f(f"{base}/rocm7.2", "2.13.0+rocm7.2") is True
+ assert f(f"{base}/rocm7.2", "2.11.5+rocm7.2") is False # patch on 2.11 is in-spec
+ # An UNKNOWN newer rocm (not on the 2.11 allowlist) is not floored to 2.11, so a
+ # matching rocm version at any release line is NOT a mismatch on this branch.
+ assert f(f"{base}/rocm8.0", "2.12.0+rocm8.0") is False
+ # gfx pin (2.11 line) vs installed release line.
+ assert f(f"{amd}/gfx1151", "2.10.0+rocm6.4") is True
+ assert f(f"{amd}/gfx1151", "2.11.0+rocm7.13.0") is False
+ # rocm7.2 pin vs an untagged (no +rocm) wheel: a CPU/CUDA build never
+ # satisfies a ROCm pin, regardless of its release line -> always a mismatch.
+ assert f(f"{base}/rocm7.2", "2.10.0") is True
+ assert f(f"{base}/rocm7.2", "2.11.0") is True
+ assert f(f"{base}/rocm6.4", "2.10.0") is True
+ # A 2.11-allowlist gfx pin over a GENERIC (two-part +rocm7.2) 2.11 wheel mismatches:
+ # the user wants AMD's per-arch (three-part) wheel, not the generic one.
+ assert f(f"{amd}/gfx1151", "2.11.0+rocm7.2") is True
+ assert f(f"{amd}/gfx120X-all", "2.11.0+rocm7.2") is True
+ # ...but an already-installed per-arch (three-part) wheel is NOT re-flagged
+ # (no reinstall loop once the correct gfx wheel is present).
+ assert f(f"{amd}/gfx120X-all", "2.11.0+rocm7.13.0") is False
+ assert f(f"{amd}/gfx1150", "2.11.0+rocm7.13.0") is False
+ # A NON-2.11 gfx pin (gfx110X-all/gfx90a/gfx908) tracks the default <2.11 spec: a
+ # correct 2.10+rocm wheel is NOT a mismatch, a 2.11 build is.
+ assert f(f"{amd}/gfx110X-all", "2.10.0+rocm6.4") is False
+ assert f(f"{amd}/gfx90a", "2.10.0+rocm6.3") is False
+ assert f(f"{amd}/gfx908", "2.10.0+rocm7.0") is False
+ assert f(f"{amd}/gfx110X-all", "2.11.0+rocm7.2") is True
+ # A non-2.11 gfx pin over an untagged (no +rocm) wheel is a mismatch even
+ # when torch is already <2.11: a CPU/CUDA build never satisfies the ROCm pin.
+ assert f(f"{amd}/gfx110X-all", "2.10.0") is True
+ assert f(f"{amd}/gfx90a", "2.10.0") is True
+ # A major-only rocm pin (rocm7) compares on the major alone: rocm6.x mismatches,
+ # any rocm7.x satisfies it, an untagged wheel never does, a bare +rocm is lenient.
+ assert f(f"{base}/rocm7", "2.10.0+rocm6.4") is True
+ assert f(f"{base}/rocm7", "2.11.0+rocm7.2") is False
+ assert f(f"{base}/rocm7", "2.11.0+rocm7.13.0") is False
+ assert f(f"{base}/rocm7", "2.10.0") is True
+ assert f(f"{base}/rocm7", "2.10.0+rocm") is False
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
+ def test_rocm_pin_mismatch_over_installed_rocm_reinstalls(
+ self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """A rocm7.2 pin over an already-installed OLDER +rocm6.4 build must reinstall,
+ even though has_hip_torch is True (the ROCm analogue of the CUDA cuXXX mismatch)."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ # HIP marker present (has_hip_torch=True) + installed +rocm6.4 wheel.
+ mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
+ env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}
+ with patch.dict(stack_mod.os.environ, env, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "rocm7.2" in torch_call
+ assert "torch>=2.11.0,<2.12.0" in torch_call
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
+ def test_gfx_pin_over_installed_pre211_rocm_reinstalls(
+ self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """A gfx* pin (2.11 line) over an installed pre-2.11 +rocm6.4 build reinstalls."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
+ env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
+ with patch.dict(stack_mod.os.environ, env, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ with patch.object(
+ stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
+ ):
+ _ensure_rocm_torch()
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "gfx1151" in torch_call
+ assert "torch>=2.11.0,<2.12.0" in torch_call
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
+ def test_rocm_pin_matches_installed_no_torch_reinstall(
+ self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """A rocm7.2 pin over an already-matching +rocm7.2 build must NOT reinstall torch
+ (no false reinstall of a correct ROCm venv)."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n"
+ env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}
+ with patch.dict(stack_mod.os.environ, env, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ # No torch reinstall: any pip_install call must not target a torch index.
+ for _call in mock_pip.call_args_list:
+ _args = [str(a) for a in _call.args]
+ if "--index-url" in _args:
+ _url = _args[_args.index("--index-url") + 1]
+ assert "rocm7.2" not in _url or "torch" not in " ".join(
+ _args
+ ), "torch must not be reinstalled when the pin already matches"
+ # A torch reinstall would pass torch>=... as a positional; assert none did.
+ assert not any(
+ any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
+ )
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
+ def test_non211_gfx_pin_over_210_rocm_no_reinstall(
+ self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """A gfx110X-all pin (NOT in the 2.11 allowlist) over a correct 2.10+rocm
+ wheel must NOT be flagged stale -- the install path uses the default <2.11
+ specs for that arch, so re-flagging would reinstall-loop on every update."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
+ env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx110X-all"}
+ with patch.dict(stack_mod.os.environ, env, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ # has_hip_torch True + no mismatch -> torch must NOT be reinstalled.
+ assert not any(
+ any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
+ )
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
+ def test_gfx_pin_over_generic_rocm211_reinstalls(
+ self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """A gfx1151 pin over a GENERIC (two-part +rocm7.2) 2.11 wheel must reinstall
+ the AMD per-arch wheel -- even though both are torch 2.11, the generic wheel
+ is not the per-arch build the user pinned (Strix stays off the generic wheel)."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n"
+ env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
+ with patch.dict(stack_mod.os.environ, env, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ with patch.object(
+ stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
+ ):
+ _ensure_rocm_torch()
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "gfx1151" in torch_call
+ assert "torch>=2.11.0,<2.12.0" in torch_call
+
+ def test_radeon_url_not_classified_as_pip_rocm_family(self):
+ """A repo.radeon.com find-links dir (leaf rocm-rel-7.2.1) starts with "rocm" but is
+ NOT a pip --index-url ROCm family: it must route to the verbatim path, not a
+ --index-url reinstall that fails against a find-links listing."""
+ leaf_f = stack_mod._is_pip_rocm_family_leaf
+ # Real pip ROCm families (download.pytorch.org/whl/rocmX.Y, repo.amd.com gfx).
+ assert leaf_f("rocm7.2") is True
+ assert leaf_f("rocm6.4") is True
+ assert leaf_f("gfx120x-all") is True
+ assert leaf_f("gfx1151") is True
+ # A bare rocm (no minor) is still an exact family.
+ assert leaf_f("rocm7") is True
+ # A Radeon find-links dir leaf, a custom mirror, cpu and cuda are NOT pip rocm.
+ assert leaf_f("rocm-rel-7.2.1") is False
+ assert leaf_f("simple") is False
+ assert leaf_f("current") is False
+ assert leaf_f("cpu") is False
+ assert leaf_f("cu128") is False
+ # A rocm-SUFFIX private mirror shares the family prefix but is a custom pin
+ # the verbatim path owns: a ^rocm\d PREFIX match would wrongly treat it as a
+ # --index-url family. Match EXACTLY.
+ assert leaf_f("rocm7.2-private") is False
+ assert leaf_f("rocm7-current") is False
+ assert leaf_f("rocm7.2.1") is False # two-part local suffix -> custom, not rocm7.2
+
+ radeon = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1"
+ pip_rocm = "https://download.pytorch.org/whl/rocm7.2"
+ amd_gfx = "https://repo.amd.com/rocm/whl/gfx120X-all"
+
+ def _classify(url, fn):
+ with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ return fn()
+
+ rocm_fn = stack_mod._explicit_rocm_torch_index_url
+ unk_fn = stack_mod._explicit_unknown_family_torch_index_url
+ # Real pip rocm/gfx pins ARE a ROCm family (reinstallable via --index-url) and
+ # are NOT "unknown".
+ assert _classify(pip_rocm, rocm_fn) == pip_rocm
+ assert _classify(amd_gfx, rocm_fn) == amd_gfx
+ assert _classify(pip_rocm, unk_fn) is None
+ assert _classify(amd_gfx, unk_fn) is None
+ # The Radeon find-links URL is NOT a pip ROCm family (so _ensure_rocm_torch skips
+ # it) and IS unknown, so the family repair helpers leave it alone.
+ assert _classify(radeon, rocm_fn) is None
+ assert _classify(radeon, unk_fn) == radeon
+
+ # A rocm-suffix private mirror routes the same way: NOT a pip rocm family,
+ # IS an unknown-family (verbatim) pin.
+ suffixed = "https://co.internal/whl/rocm7.2-private"
+ assert _classify(suffixed, rocm_fn) is None
+ assert _classify(suffixed, unk_fn) == suffixed
+
+ @patch.object(stack_mod, "pip_install")
+ def test_ensure_cpu_torch_broken_probe_reinstalls(self, mock_pip):
+ """_ensure_cpu_torch: torch present but unimportable (probe exit != 0) under an
+ explicit CPU pin must reinstall from the pin, not return -- the base update does
+ not repair a broken installed torch, so returning would strand it (Codex P2)."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 1 # torch present but cannot import
+ mock_probe.stdout = b""
+ env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/cpu"}
+ with patch.dict(stack_mod.os.environ, env, clear = False):
+ stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+ with patch("subprocess.run", return_value = mock_probe):
+ with patch.object(stack_mod, "NO_TORCH", False):
+ stack_mod._ensure_cpu_torch()
+ assert mock_pip.call_count == 1
+ assert "https://mirror.local/cpu" in str(mock_pip.call_args)
+
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@@ -732,7 +1048,7 @@ class TestEnsureRocmTorch:
mock_pip.assert_not_called()
-# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard
+# TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692)
class TestHasRocmGpuKfdVendorGuard:
@@ -1716,9 +2032,8 @@ class TestDetectWindowsGfxArch:
assert result == "gfx1200"
def test_returns_arch_on_crash_with_gcnarchname_in_output(self):
- # Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after
- # printing gcnArchName. Accept the arch whenever gcnArchName is in
- # stdout, regardless of exit code (previously a CPU fallback).
+ # Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after printing
+ # gcnArchName. Accept the arch whenever gcnArchName is in stdout, any exit code.
mock_result = MagicMock()
mock_result.returncode = -1073741819 # 0xC0000005 STATUS_ACCESS_VIOLATION
mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n"
@@ -2360,6 +2675,49 @@ class TestWindowsRocmTorchaoGuard:
assert not any("torchao" in arg for arg in installed_specs)
+class TestProgressStepCountMatchesTotal:
+ """The progress bar must reach exactly _TOTAL: every _progress() step is counted in
+ base_total. Regression for a repair step added without incrementing base_total,
+ which pushed _STEP past _TOTAL (Codex P2)."""
+
+ def _run_stack(self, tmp_path, *, is_windows, is_macos, is_mac_arm):
+ unstructured_plugin = tmp_path / "unstructured"
+ github_plugin = tmp_path / "github"
+ unstructured_plugin.mkdir()
+ github_plugin.mkdir()
+ sub = MagicMock()
+ sub.returncode = 0
+ sub.stdout = ""
+ with (
+ patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}),
+ patch.object(stack_mod, "IS_WINDOWS", is_windows),
+ patch.object(stack_mod, "IS_MACOS", is_macos),
+ patch.object(stack_mod, "IS_MAC_ARM", is_mac_arm),
+ patch.object(stack_mod, "NO_TORCH", False),
+ patch.object(stack_mod, "_rocm_windows_torch_installed", False),
+ patch.object(stack_mod, "_bootstrap_uv", return_value = False),
+ patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = False),
+ patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True),
+ patch.object(stack_mod, "_repair_bad_anyio"),
+ patch.object(stack_mod, "_ensure_cuda_torch"),
+ patch.object(stack_mod, "_ensure_rocm_torch"),
+ patch.object(stack_mod, "_ensure_cpu_torch"),
+ patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin),
+ patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin),
+ patch.object(stack_mod.subprocess, "run", return_value = sub),
+ ):
+ assert stack_mod.install_python_stack() == 0
+ return stack_mod._STEP, stack_mod._TOTAL
+
+ def test_windows_progress_reaches_total(self, tmp_path):
+ step, total = self._run_stack(tmp_path, is_windows = True, is_macos = False, is_mac_arm = False)
+ assert step == total, f"Windows progress {step} != total {total} (final step uncounted)"
+
+ def test_linux_progress_reaches_total(self, tmp_path):
+ step, total = self._run_stack(tmp_path, is_windows = False, is_macos = False, is_mac_arm = False)
+ assert step == total, f"Linux progress {step} != total {total}"
+
+
# TEST: worker.py -- Windows ROCm patches (source-level checks)
@@ -2846,6 +3204,24 @@ class TestStrixRocm71Override:
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
assert "TORCH_CONSTRAINT" in source and "2.11" in source
+ def test_torch_constraint_211_matches_leaf_not_whole_url(self):
+ """The 2.11 constraint case must match the index LEAF, not the whole URL.
+
+ A custom UNSLOTH_PYTORCH_MIRROR whose base path contains a gfx/rocm7.2
+ segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family must
+ not be pushed to the torch 2.11 line -- same leaf-only reasoning the
+ UNSLOTH_TORCH_BACKEND classification uses.
+ """
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ # The 2.11 constraint block must switch on $_torch_index_leaf, not the full
+ # $TORCH_INDEX_URL (a */gfx* match false-positives on a mirror base path). Only the
+ # _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150) are pushed to 2.11;
+ # a bare gfx* would also floor gfx110X-all/gfx90a/gfx908, left bare on purpose.
+ assert 'case "$_torch_index_leaf" in\n rocm7.2|gfx120x-all|gfx1151|gfx1150)' in source, (
+ "the torch>=2.11 constraint must match the specific gfx leaves that need "
+ "it (rocm7.2|gfx120x-all|gfx1151|gfx1150), not a bare gfx* or the whole URL"
+ )
+
def test_amd_rocm_mirror_env_var_respected(self):
"""install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
@@ -2934,9 +3310,9 @@ class TestServerStartupRocmFixes:
assert '"BNB_ROCM_VERSION" not in os.environ' in source
# ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ────────────────
- # bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD
- # wheel ships it in venv Scripts (on PATH only for activated venvs), so
- # without the prepend bnb logs "[WinError 2]" when launched directly.
+ # bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD wheel ships it
+ # in venv Scripts (on PATH only for activated venvs), so without the prepend bnb logs
+ # "[WinError 2]" when launched directly.
def test_main_py_prepends_hipinfo_dir_to_path(self):
"""main.py must make hipInfo.exe resolvable before bnb imports."""
@@ -3178,11 +3554,10 @@ class TestRocmGfxForwarding:
assert '$HelperReleaseRepo = "unslothai/llama.cpp"' in source
assert "$HelperReleaseRepo = if (" not in source
- # The text pins above guard the literal. The tests below *execute* the real
- # routing line from setup.sh / setup.ps1 and assert the resolved release repo,
- # so a refactor that reintroduces a conditional (or a ggml-org branch) is still
- # caught. Inputs are varied -- CPU-only, inferred/forwarded gfx, usable NVIDIA --
- # to prove no host slips back onto ggml-org. No GPU, no tooling, no network.
+ # The text pins above guard the literal. The tests below execute the real routing line
+ # from setup.sh / setup.ps1 and assert the resolved release repo, so a refactor that
+ # reintroduces a conditional (or a ggml-org branch) is still caught. Inputs vary
+ # (CPU-only, inferred/forwarded gfx, usable NVIDIA) to prove no host hits ggml-org.
@staticmethod
def _resolve_setup_sh_repo(
@@ -3277,8 +3652,8 @@ class TestRocmGfxForwarding:
# TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output.
-# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt
-# for the selected GPU, not GPU 0.
+# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt for the
+# selected GPU, not GPU 0.
_pick_rocm_gfx_target = prebuilt_mod._pick_rocm_gfx_target
@@ -3454,7 +3829,11 @@ class TestWslRerouteNvidiaGuard:
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
start = source.find("_maybe_reroute_strixhalo_to_2404()")
assert start != -1
- body = source[start : start + 1200]
+ # Slice the WHOLE function body (to its closing brace at column 0), not a
+ # fixed-length window: preamble growth must not push the signals out of view.
+ end = source.find("\n}", start)
+ assert end != -1
+ body = source[start:end]
nv = body.find("_has_usable_nvidia_gpu")
wmi = body.find("_wsl_amd_gpu_name")
assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute"
diff --git a/tests/studio/test_setup_pin_stale.ps1 b/tests/studio/test_setup_pin_stale.ps1
new file mode 100644
index 0000000000..2c92ae317f
--- /dev/null
+++ b/tests/studio/test_setup_pin_stale.ps1
@@ -0,0 +1,114 @@
+#!/usr/bin/env pwsh
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+# Unit test for studio/setup.ps1's pinned-torch-index stale-venv helpers
+# (Test-RocmGfx211Leaf, Test-CudaFamilyLeaf, Get-RocmPinStaleTags). Pure helpers,
+# AST-extracted and run in-process. Mirrors the Python _rocm_pin_family_mismatch /
+# _is_cuda_family_leaf tests.
+# Run: pwsh -NoProfile -File tests/studio/test_setup_pin_stale.ps1
+
+$ErrorActionPreference = "Stop"
+$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1")
+$setupPath = (Resolve-Path $setupPath).Path
+
+# --- Parse setup.ps1 (also serves as a syntax gate) and extract the helpers ---
+$tokens = $null; $errors = $null
+$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
+if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
+
+foreach ($name in @("Test-RocmGfx211Leaf", "Test-RocmKnown211Version", "Test-CudaFamilyLeaf", "Get-RocmPinStaleTags")) {
+ $fn = $ast.FindAll({ param($n)
+ $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
+ }, $true)
+ if ($fn.Count -ne 1) { throw "expected exactly one $name in setup.ps1, found $($fn.Count)" }
+ # Pure helpers (no exit / external calls) -- safe to define in this scope.
+ Invoke-Expression $fn[0].Extent.Text
+}
+
+$failures = 0
+function Check($name, $cond) {
+ if ($cond) { Write-Host " PASS $name" }
+ else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ }
+}
+
+# A pinned gfx/rocm index is stale when Expected != Installed.
+function IsStale($leaf, $ver) {
+ $t = Get-RocmPinStaleTags -PinLeaf $leaf -TorchVersion $ver
+ return $t.Expected -ne $t.Installed
+}
+
+Write-Host "Test-RocmGfx211Leaf (the 2.11 gfx allowlist)"
+Check "gfx1151 -> true" (Test-RocmGfx211Leaf "gfx1151")
+Check "gfx1150 -> true" (Test-RocmGfx211Leaf "gfx1150")
+Check "gfx120x-all -> true" (Test-RocmGfx211Leaf "gfx120x-all")
+Check "gfx110x-all -> false" (-not (Test-RocmGfx211Leaf "gfx110x-all"))
+Check "gfx90a -> false" (-not (Test-RocmGfx211Leaf "gfx90a"))
+Check "gfx908 -> false" (-not (Test-RocmGfx211Leaf "gfx908"))
+
+Write-Host "Test-CudaFamilyLeaf (^cu[0-9])"
+Check "cu118 -> true" (Test-CudaFamilyLeaf "cu118")
+Check "cu128 -> true" (Test-CudaFamilyLeaf "cu128")
+Check "cu130 -> true" (Test-CudaFamilyLeaf "cu130")
+Check "custom -> false" (-not (Test-CudaFamilyLeaf "custom"))
+Check "current -> false" (-not (Test-CudaFamilyLeaf "current"))
+Check "cpu -> false" (-not (Test-CudaFamilyLeaf "cpu"))
+Check "empty -> false" (-not (Test-CudaFamilyLeaf ""))
+
+Write-Host "Get-RocmPinStaleTags (mirror of _rocm_pin_family_mismatch)"
+# Exact rocm version comparison.
+Check "rocm7.2 pin + 2.11.0+rocm7.2 -> not stale" (-not (IsStale "rocm7.2" "2.11.0+rocm7.2"))
+Check "rocm7.2 pin + 2.10.0+rocm6.4 -> stale" (IsStale "rocm7.2" "2.10.0+rocm6.4")
+Check "rocm6.4 pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "rocm6.4" "2.10.0+rocm6.4"))
+# rocm7.2 is a KNOWN-2.11 index. A +rocm7.2 wheel whose RELEASE drifted off 2.11 shares
+# the tag but violates the spec -> stale (mirror of _rocm_pin_family_mismatch).
+Check "rocm7.2 pin + 2.12.0+rocm7.2 -> stale" (IsStale "rocm7.2" "2.12.0+rocm7.2")
+Check "rocm7.2 pin + 2.13.0+rocm7.2 -> stale" (IsStale "rocm7.2" "2.13.0+rocm7.2")
+Check "rocm7.2 pin + 2.11.5+rocm7.2 -> not stale" (-not (IsStale "rocm7.2" "2.11.5+rocm7.2"))
+# An UNKNOWN newer rocm (off the 2.11 allowlist) isn't floored, so a matching version at
+# any release line is NOT stale on this exact-compare branch.
+Check "rocm8.0 pin + 2.12.0+rocm8.0 -> not stale" (-not (IsStale "rocm8.0" "2.12.0+rocm8.0"))
+# An untagged (no +rocm) wheel never satisfies a ROCm pin -> always stale.
+Check "rocm7.2 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm7.2" "2.10.0")
+Check "rocm7.2 pin + 2.11.0 (untagged) -> stale" (IsStale "rocm7.2" "2.11.0")
+Check "rocm6.4 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm6.4" "2.10.0")
+# 2.11-allowlist gfx pin: per-arch (three-part) wheel is satisfied, generic is stale.
+Check "gfx1151 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1151" "2.11.0+rocm7.13.0"))
+Check "gfx1150 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1150" "2.11.0+rocm7.13.0"))
+Check "gfx120x-all pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx120x-all" "2.11.0+rocm7.13.0"))
+Check "gfx1151 pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx1151" "2.11.0+rocm7.2")
+Check "gfx1151 pin + 2.10.0+rocm6.4 -> stale" (IsStale "gfx1151" "2.10.0+rocm6.4")
+# Non-2.11 gfx pin (gfx110X-all/gfx90a/gfx908): a valid <2.11 wheel is NOT stale.
+Check "gfx110x-all pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "gfx110x-all" "2.10.0+rocm6.4"))
+Check "gfx90a pin + 2.10.0+rocm6.3 -> not stale" (-not (IsStale "gfx90a" "2.10.0+rocm6.3"))
+Check "gfx908 pin + 2.10.0+rocm7.0 -> not stale" (-not (IsStale "gfx908" "2.10.0+rocm7.0"))
+Check "gfx110x-all pin + 2.11.0+rocm7.2 -> stale" (IsStale "gfx110x-all" "2.11.0+rocm7.2")
+# Non-2.11 gfx pin over an untagged wheel: never satisfies the pin -> stale, so the
+# explicit ROCm index is applied even when torch is already <2.11.
+Check "gfx110x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx110x-all" "2.10.0")
+Check "gfx90a pin + 2.10.0 (untagged) -> stale" (IsStale "gfx90a" "2.10.0")
+# Capital gfx120X-all is lowercased by Get-TorchIndexLeaf before this helper, so the
+# 2.11-allowlist branch fires and a generic/untagged wheel is stale.
+Check "gfx120x-all pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx120x-all" "2.11.0+rocm7.2")
+Check "gfx120x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx120x-all" "2.10.0")
+
+# Major-only rocm pin (rocm7): majors compared alone; mirrors _rocm_pin_family_mismatch.
+Check "rocm7 pin + 2.10.0+rocm6.4 -> stale" (IsStale "rocm7" "2.10.0+rocm6.4")
+Check "rocm7 pin + 2.11.0+rocm7.2 -> not stale" (-not (IsStale "rocm7" "2.11.0+rocm7.2"))
+Check "rocm7 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "rocm7" "2.11.0+rocm7.13.0"))
+Check "rocm7 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm7" "2.10.0")
+Check "rocm7 pin + 2.10.0+rocm (unreadable) -> not stale" (-not (IsStale "rocm7" "2.10.0+rocm"))
+
+Write-Host "Test-RocmKnown211Version + KNOWN-2.11 fallback (rocm7.2 only; no speculative rocm7.3)"
+Check "rocm7.2 -> known 2.11" (Test-RocmKnown211Version -Major 7 -Minor 2)
+Check "rocm7.1 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 1))
+Check "rocm7.3 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 3))
+Check "rocm8.0 -> not known" (-not (Test-RocmKnown211Version -Major 8 -Minor 0))
+# Unreadable-installed fallback: a rocm7.3 pin (unknown -> <2.11 line) over a <2.11 +rocm
+# wheel with an unreadable version is NOT stale; rocm7.2 (KNOWN-2.11) over the same wheel
+# IS stale (#2534 alignment).
+Check "rocm7.3 pin + 2.10.0+rocm (unreadable ver) -> not stale" (-not (IsStale "rocm7.3" "2.10.0+rocm"))
+Check "rocm7.2 pin + 2.10.0+rocm (unreadable ver) -> stale" (IsStale "rocm7.2" "2.10.0+rocm")
+
+Write-Host ""
+if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
+Write-Host "All checks passed" -ForegroundColor Green
diff --git a/tests/studio/test_torch_flavor.ps1 b/tests/studio/test_torch_flavor.ps1
index 50be4814b0..f2cc55c21c 100644
--- a/tests/studio/test_torch_flavor.ps1
+++ b/tests/studio/test_torch_flavor.ps1
@@ -15,7 +15,7 @@ $tokens = $null; $errors = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors)
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" }
-foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag")) {
+foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag", "Trim-IndexPathSlashes", "Redact-InstallOutput")) {
$fn = $ast.FindAll({ param($n)
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
}, $true)
@@ -49,6 +49,19 @@ Check "mirror cu130 leaf -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl
Check "unrecognized leaf -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/simple"))
Check "empty url -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl ""))
+Write-Host "Trim-IndexPathSlashes (install.ps1 parity: path-only, token-preserving)"
+Check "double path slash collapsed" ((Trim-IndexPathSlashes "https://h/whl/cu128//") -eq "https://h/whl/cu128")
+Check "single trailing slash trimmed" ((Trim-IndexPathSlashes "https://h/whl/cu128/") -eq "https://h/whl/cu128")
+Check "query token slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
+Check "path slash trimmed, query kept" ((Trim-IndexPathSlashes "https://h/whl/cu128//?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
+
+Write-Host "Redact-InstallOutput (install.ps1 parity: credential redaction)"
+Check "userinfo redacted" ((Redact-InstallOutput "ERROR https://alice:s3cr3t@download.pytorch.org/whl/cu128") -eq "ERROR https://@download.pytorch.org/whl/cu128")
+Check "query value redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abcd1234&channel=beta") -eq "https://host/whl/cu128?token=&channel=")
+Check "fragment token redacted" ((Redact-InstallOutput "ERROR https://mirror.local/whl/cu128#token=SECRET123 (403)") -eq "ERROR https://mirror.local/whl/cu128# (403)")
+Check "bare hash comment untouched" ((Redact-InstallOutput "# retrying with --no-cache-dir") -eq "# retrying with --no-cache-dir")
+Check "plain line untouched" ((Redact-InstallOutput "Resolved 42 packages in 1.2s") -eq "Resolved 42 packages in 1.2s")
+
Write-Host ""
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
Write-Host "All checks passed" -ForegroundColor Green
diff --git a/tests/studio/test_torch_index_pin_hardening.ps1 b/tests/studio/test_torch_index_pin_hardening.ps1
new file mode 100644
index 0000000000..b8edf3da25
--- /dev/null
+++ b/tests/studio/test_torch_index_pin_hardening.ps1
@@ -0,0 +1,78 @@
+#!/usr/bin/env pwsh
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+# Unit tests for setup.ps1's torch-index pin-hardening helpers: Trim-IndexPathSlashes
+# (path-only slash trim, token-preserving), Redact-InstallOutput (credential redaction of
+# captured install logs), Get-TorchIndexLeaf (ALL trailing slashes stripped) and
+# Test-PipRocmFamilyLeaf (rocm7. is a custom pin, not a family). Pure helpers, AST-extracted
+# and run in-process. Run: pwsh -NoProfile -File tests/studio/test_torch_index_pin_hardening.ps1
+
+$ErrorActionPreference = "Stop"
+$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1")
+$setupPath = (Resolve-Path $setupPath).Path
+$setupText = Get-Content -Raw $setupPath
+
+# --- Parse setup.ps1 (also a syntax gate) and extract the pure helpers ---
+$tokens = $null; $errors = $null
+$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
+if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
+
+foreach ($name in @("Trim-IndexPathSlashes", "Redact-InstallOutput", "Get-TorchIndexLeaf", "Test-PipRocmFamilyLeaf")) {
+ $fn = $ast.FindAll({ param($n)
+ $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
+ }, $true)
+ if ($fn.Count -ne 1) { throw "expected exactly one $name in setup.ps1, found $($fn.Count)" }
+ Invoke-Expression $fn[0].Extent.Text
+}
+
+$failures = 0
+function Check($name, $cond) {
+ if ($cond) { Write-Host " PASS $name" }
+ else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ }
+}
+
+Write-Host "Trim-IndexPathSlashes (path-only, token-preserving)"
+Check "double path slash collapsed" ((Trim-IndexPathSlashes "https://h/whl/cu128//") -eq "https://h/whl/cu128")
+Check "single trailing slash trimmed" ((Trim-IndexPathSlashes "https://h/whl/cu128/") -eq "https://h/whl/cu128")
+Check "no slash unchanged" ((Trim-IndexPathSlashes "https://h/whl/cu128") -eq "https://h/whl/cu128")
+Check "query token slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
+Check "path slash trimmed, query kept" ((Trim-IndexPathSlashes "https://h/whl/cu128//?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
+Check "fragment slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128#anchor/") -eq "https://h/whl/cu128#anchor/")
+
+Write-Host "Redact-InstallOutput (credential redaction)"
+Check "userinfo redacted" ((Redact-InstallOutput "ERROR https://alice:s3cr3t@download.pytorch.org/whl/cu128") -eq "ERROR https://@download.pytorch.org/whl/cu128")
+Check "bare-token@ redacted" ((Redact-InstallOutput "fetch https://ghp_deadbeef@host/whl/cu128 failed") -eq "fetch https://@host/whl/cu128 failed")
+Check "single query value redacted" ((Redact-InstallOutput "url https://host/whl/cu128?token=abcd1234 unreachable") -eq "url https://host/whl/cu128?token= unreachable")
+Check "multiple query values redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abcd1234&channel=beta") -eq "https://host/whl/cu128?token=&channel=")
+Check "fragment token redacted" ((Redact-InstallOutput "ERROR https://mirror.local/whl/cu128#token=SECRET123 (403)") -eq "ERROR https://mirror.local/whl/cu128# (403)")
+Check "query and fragment both redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abc#sig=xyz done") -eq "https://host/whl/cu128?token=# done")
+Check "bare hash comment untouched" ((Redact-InstallOutput "# retrying with --no-cache-dir") -eq "# retrying with --no-cache-dir")
+Check "plain line untouched" ((Redact-InstallOutput "Resolved 42 packages in 1.2s") -eq "Resolved 42 packages in 1.2s")
+$leak = Redact-InstallOutput "https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET"
+Check "no secret substring survives" (($leak -notmatch "s3cr3t") -and ($leak -notmatch "SUPERSECRET") -and ($leak -notmatch "ALSOSECRET"))
+
+Write-Host "Get-TorchIndexLeaf (ALL trailing slashes stripped)"
+Check "double slash cu128 -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128//") -eq "cu128")
+Check "triple slash rocm7.2 -> rocm7.2" ((Get-TorchIndexLeaf "https://m/whl/rocm7.2///") -eq "rocm7.2")
+Check "double slash + token -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128//?token=x") -eq "cu128")
+Check "single slash cu128 -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128/") -eq "cu128")
+
+Write-Host "Test-PipRocmFamilyLeaf (rocm7. is a custom pin, not a family)"
+Check "rocm7 family" (Test-PipRocmFamilyLeaf "rocm7")
+Check "rocm7.2 family" (Test-PipRocmFamilyLeaf "rocm7.2")
+Check "gfx1151 family" (Test-PipRocmFamilyLeaf "gfx1151")
+Check "rocm7. trailing-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7."))
+Check "rocm.7 leading-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm.7"))
+Check "rocm7.2.1 two-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7.2.1"))
+Check "rocm7.2-private NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7.2-private"))
+Check "cu128 NOT family" (-not (Test-PipRocmFamilyLeaf "cu128"))
+
+Write-Host "Fast-Install pinned-install env scrub (source assertion)"
+# The pip fallback honours PIP_*; PIP_NO_INDEX=1 would make it ignore the pinned --index-url
+# and PIP_INDEX_URL would replace it, so both must be scrubbed for a pinned install.
+Check "PIP_NO_INDEX scrubbed" ($setupText -match "'PIP_NO_INDEX'")
+Check "PIP_INDEX_URL scrubbed" ($setupText -match "'PIP_INDEX_URL'")
+
+Write-Host ""
+if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
+Write-Host "All checks passed" -ForegroundColor Green
From 65587c2be7b6167e369bc5d95155d315c2717ac8 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Mon, 20 Jul 2026 04:57:44 -0700
Subject: [PATCH 030/255] Studio: Data settings tab, uploaded files manager,
quant pinning, and chat image preview fix (#7029)
* Studio: Data settings tab, uploaded files manager, quant pinning, image preview fix
Settings
- New Data tab in the settings sidebar, under Connections. Chat data
management (archived chats, confirm before deleting, exports, import,
clear all) moved there from the Chat tab.
- New Archive all chats action with confirmation. Archives every chat in
Recents and Projects; compare pairs count as one chat.
- New Uploaded files manager listing RAG documents (chats, projects,
knowledge bases) and chat message attachments with location, size and
date. Files can be opened in a new tab or deleted. Deleting a chat
attachment keeps the message text.
Backend
- GET /api/rag/documents lists all uploaded RAG documents with file size
plus KB and project names.
- GET /api/chat/attachments lists chat message attachments; per
attachment file and delete endpoints included.
Model selector
- Downloaded GGUF quants can be pinned from the quant row (next to the
settings and delete actions). Pinned quants show at the top of On
Device under a Pinned heading as model name plus a grey quant chip and
load directly with one click. Non GGUF cached repos pin as a whole.
- Toned down the green of the downloaded label.
Fix
- Clicking an image attachment in chat now opens the preview overlay.
The tooltip trigger wrapper called preventDefault before composed
handlers ran, which made Radix DialogTrigger skip opening.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: image previews and file type chips in uploaded files list
Image attachments now show a small thumbnail (lazy loaded from the
stored bytes, object URL revoked on unmount) and every row shows a grey
uppercase type chip derived from the extension or content type. Non
image rows keep a file icon. Name cell floors its width and clips
overflow so narrow dialogs stay aligned.
* Harden attachment serving, add tests, and polish pinned rows and previews
- Strict base64 decoding for attachment files: corrupt payloads now return
422 instead of silently serving empty or garbled bytes; whitespace,
missing padding, the URL-safe alphabet, and RFC 2397 percent-encoded
data URLs are all handled
- New backend test suite covering attachment listing, size accounting,
malformed rows, deletion semantics, and every file-serving edge case
- Pinned quant rows show a Loaded tag when that exact quant is active,
and reveal unpin, settings, and delete actions on hover
- Uploaded files dialog is wider and chat locations link straight to the
thread the attachment belongs to
- Chat image preview is now a chrome-free lightbox: dimmed backdrop,
rounded image, corner close button, click outside to dismiss
- File opens go through a synchronous window.open so Safari and Firefox
popup blockers do not eat them
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Uploaded files: click a file to jump to its chat, square thumbs, new Data icon
- Clicking a file row (thumbnail or name) now goes straight to the chat it
belongs to; files without a chat open directly as before
- File thumbnails pin a small 7px radius: the theme scales rounded-md up
to a near circle at this size
- Settings Data tab now uses the database-setting icon
* Uploaded files is now a Data tab subpage instead of a popup
- Manage swaps the tab body for an inline Uploaded files page with a back
header, matching the rest of settings navigation
- Size column header and values are left aligned like the other columns
- Column widths tightened so the table fits the settings panel
* Lightbox polish and Data tab row order
- Image preview close button is transparent until hovered
- Preview image no longer rounds its corners
- Import chats now sits below Clear all chats in the Data tab
* Data tab: export chats as fine-tuning data and open them in Recipes
- New Fine-tuning section in Settings > Data converts every chat into a
JSONL dataset in the OpenAI messages format, one conversation per line
with string-only system/user/assistant turns
- The Train tab detects this file as chatml natively: no column mapping
and no standardization pass, and it works with train on completions
since every assistant turn sits behind the chat template response marker
- Consecutive same-role turns merge, trailing turns without an assistant
reply drop, and reasoning, tool calls, and images are excluded so chat
templates format the data cleanly
- Open in Recipes stages the JSONL as a local seed upload, creates a new
Data Recipe with the seed block preconfigured, and jumps to the editor
* Data tab: load chats straight into the Train tab, row moved to the top
- New Load in Train tab button uploads the fine-tuning JSONL through the
training dataset endpoint, selects it in the training config store, and
opens the Train tab with the dataset loaded and format-checked
- Use chats as training data now sits at the very top of the Data tab
- The Chats subheading is gone; chat rows flow directly under it
* Address review findings on the uploads manager and quant pins
- Deleting the last attachment stores '[]' instead of NULL: a NULL reads
back as a missing field and triggers the legacy IndexedDB backfill,
which resurrected the deleted attachment on the next chat load
- The attachment file endpoint now serves audio: adapter parts store
{data, format} raw base64 and compare chats store a bare base64 string;
media type comes from the attachment contentType or the format
- Compare-chat uploads live in message content parts, not attachments;
the uploads list now includes those blobs via synthetic content-part
ids that the same get and delete routes resolve
- Deleting a quant from the expanded repo row also unpins it so a pinned
row cannot try to load a file that no longer exists
- Thumbnails in the uploads list fetch their blob only once the row is
visible, so a long screenshot history does not download everything
- Nine new backend tests cover audio serving, content-part listing,
serving, deletion, and the empty-list delete behavior
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Data tab: single action dropdown with format choices for chat training data
- The three fine-tune buttons collapse into one dropdown plus a run
button; pick Load in Train tab, Open in Recipes, or Export JSONL,
then click the arrow to run it
- The dropdown's Format section adds ShareGPT and Alpaca alongside the
default OpenAI messages format, ticked like a checklist; all three
shapes are auto-detected by the Train tab's format check
- Alpaca is single-turn, so each user to assistant pair becomes its own
record with the system prompt and earlier turns carried in the input
column
- Shorter description on the training data row
- Uploaded files rows show the size under the file name instead of a
separate column, matching the tighter layout
* Polish the training data action control
- Run button is a true circle (icon-sm plus rounded-full) with a
heavier arrow stroke
- Dropdown trigger uses the shared standard chevron and a fixed width
so switching actions no longer resizes the control
* Shorten the training data row description
* Use the standard chevron for the run button and enlarge the ticks
- Run button uses the shared standard right chevron so it matches the
dropdown chevron instead of the hugeicons arrow
- Dropdown ticks bumped up a size for legibility
* Reword the training data row description
* Shorten Data Recipes to Recipes in the training data description
* List Export JSONL first and rename the default format to Chat Completions
* Handle legacy string content in fine-tune exports and gate Train on chat-only hosts
- messageToPlainText now accepts plain-string message content, the shape
legacy and imported histories store, so those conversations export
instead of being skipped as having no exchange
- The Load in Train tab action is disabled on chat-only hosts the same
way the sidebar gates Train; the default action falls back to Export
JSONL there so the run button never uploads a dataset that /studio
would immediately redirect away from
* Narrow the training data action dropdown slightly
* Drop the format picker from the training data dropdown
Chat Completions (OpenAI messages) is the only export format we ship, so
the ShareGPT and Alpaca options and the Format section are removed. The
export always uses the OpenAI messages shape.
* Address the second round of review findings
Security
- Chat attachment data URLs no longer echo their embedded media type:
anything that is not a plain raster image serves as octet-stream, so
imported text/html or SVG payloads cannot render under the app origin
- Uploaded .html/.htm RAG documents serve as text/plain for the same
reason; the preview sheet only uses the file URL for PDFs
Uploads manager
- Remote image URLs in imported chats are no longer listed as stored
uploads (nothing to serve, and delete would strip the chat reference);
the delete guard mirrors the same data:-only rule
- Deleting a content-part upload refetches the list since the remaining
parts re-index, keeping sibling row ids current
- Deleting a project document from the Data tab invalidates the project
sources cache like the sources panel does
- Data-tab deletions now patch the loaded thread's in-memory copy via a
small event, so a later repo sync cannot write the attachment back
Fine-tune export
- Branch siblings from retries stay out of the exported conversation;
only the selected chain converts (full exports still keep everything)
- Assistant turns before the first user turn drop, preserving leading
system prompts, so no unconditioned assistant targets are emitted
Four new backend tests cover the media type clamp and remote-URL rows;
two existing tests updated for the clamped types
* Fix uploaded file lifecycle and model state
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Make archived chats a Data settings subpage
* Studio: fix attachment route tests and pinned quant edge cases
- test_chat_attachments: drop asyncio.run around the synchronous
/attachments routes (list/get/delete are plain def, so asyncio.run
raised 'a coroutine was expected' and failed the Repo tests CI job).
- test_chat_attachments: align compare-chat content-part assertions with
the stable content-hash id scheme (content-part-sha256-...) instead of
the removed array-index ids; resolve ids from the listing.
- pickers: pass disabled={deleteDisabled} to the pinned-quant delete
action so a quant cannot be deleted mid model-load, matching the
expanded variant rows.
- pickers: build the pinned-quant existence set from the query-unfiltered
cached GGUF repos (format filter still applied) so a pinned quant stays
findable when the search term matches only its quant name.
* Fix Studio review regressions
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Guard fine-tune export content blocks
* Add Export button for archived chats
Adds an Export action to the Archived chats view in Settings > Data that
downloads only the archived chats as a JSON backup (their threads, messages
and projects). The button sits in the archived header row and appears only
when archived chats exist.
* Refactor archived export into pure, testable units
Split the archived-chats export into a dependency-free filter
(archived-chat-export.ts) and a shared JSON download helper
(download-json.ts). Skip the download when nothing is archived so a
stray call never drops an empty file. No behavior change to the button.
---------
Co-authored-by: shimmyshimmer
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: Unsloth
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
studio/backend/core/rag/store.py | 10 +
studio/backend/routes/chat_history.py | 134 ++-
studio/backend/routes/rag.py | 39 +-
studio/backend/storage/studio_db.py | 813 +++++++++++++++++-
studio/backend/tests/test_chat_attachments.py | 634 ++++++++++++++
.../frontend/src/components/app-sidebar.tsx | 4 +-
.../components/assistant-ui/attachment.tsx | 27 +-
.../assistant-ui/model-selector/pickers.tsx | 519 ++++++++++-
.../model-selector/pinned-models.ts | 72 ++
studio/frontend/src/components/ui/tooltip.tsx | 7 +-
.../src/features/chat/api/chat-api.ts | 74 +-
.../chat/hooks/use-chat-sidebar-items.ts | 34 +
studio/frontend/src/features/chat/index.ts | 19 +-
.../prompt-storage/prompt-storage-dialog.tsx | 223 ++++-
.../src/features/chat/runtime-provider.tsx | 167 ++++
.../chat/utils/archived-chat-export.ts | 62 ++
.../chat/utils/chat-attachment-events.ts | 123 +++
.../src/features/chat/utils/download-json.ts | 18 +
.../chat/utils/export-chat-history.ts | 35 +-
.../frontend/src/features/rag/api/rag-api.ts | 27 +-
.../rag/components/use-rag-documents.ts | 136 +--
studio/frontend/src/features/rag/index.ts | 7 +-
studio/frontend/src/features/rag/types/rag.ts | 7 +
.../components/archived-chats-dialog.tsx | 135 ++-
.../settings/components/finetune-recipe.ts | 98 +++
.../components/uploaded-files-dialog.tsx | 644 ++++++++++++++
.../src/features/settings/settings-dialog.tsx | 11 +
.../src/features/settings/settings-search.ts | 13 +-
.../settings/stores/settings-dialog-store.ts | 6 +-
.../src/features/settings/tabs/chat-tab.tsx | 329 +------
.../src/features/settings/tabs/data-tab.tsx | 727 ++++++++++++++++
studio/frontend/src/i18n/locales/en.ts | 50 +-
32 files changed, 4671 insertions(+), 533 deletions(-)
create mode 100644 studio/backend/tests/test_chat_attachments.py
create mode 100644 studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts
create mode 100644 studio/frontend/src/features/chat/utils/archived-chat-export.ts
create mode 100644 studio/frontend/src/features/chat/utils/chat-attachment-events.ts
create mode 100644 studio/frontend/src/features/chat/utils/download-json.ts
create mode 100644 studio/frontend/src/features/settings/components/finetune-recipe.ts
create mode 100644 studio/frontend/src/features/settings/components/uploaded-files-dialog.tsx
create mode 100644 studio/frontend/src/features/settings/tabs/data-tab.tsx
diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py
index f9128d1715..1165b6bb0e 100644
--- a/studio/backend/core/rag/store.py
+++ b/studio/backend/core/rag/store.py
@@ -158,6 +158,16 @@ def list_documents(conn: sqlite3.Connection, scope: str) -> list[dict]:
return [dict(r) for r in rows]
+def list_all_documents(conn: sqlite3.Connection) -> list[dict]:
+ """Every uploaded document across all scopes (KBs, threads, projects)."""
+ rows = conn.execute(
+ "SELECT id, scope, kb_id, thread_id, project_id, filename, sha256, status, error, "
+ "num_chunks, stored_path, created_at "
+ "FROM documents ORDER BY created_at DESC"
+ ).fetchall()
+ return [dict(r) for r in rows]
+
+
def get_document(conn: sqlite3.Connection, document_id: str) -> dict | None:
row = conn.execute("SELECT * FROM documents WHERE id=?", (document_id,)).fetchone()
return dict(row) if row else None
diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py
index 7a27a58a52..24b6dfb36d 100644
--- a/studio/backend/routes/chat_history.py
+++ b/studio/backend/routes/chat_history.py
@@ -5,7 +5,7 @@
Chat history API routes backed by studio.db.
"""
-from typing import Any, Literal, Optional
+from typing import Annotated, Any, Literal, Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, ConfigDict, Field, ValidationError
@@ -19,13 +19,16 @@ from storage.studio_db import (
clear_chat_history,
count_chat_threads,
count_forks_for_message,
+ delete_chat_attachment,
delete_chat_threads,
delete_chat_project,
ensure_chat_project_workspace,
fork_chat_thread,
+ get_chat_attachment,
get_chat_project,
get_chat_thread,
get_chat_message,
+ list_chat_attachments_page,
list_chat_projects,
list_chat_legacy_imports,
list_chat_settings,
@@ -279,6 +282,131 @@ async def delete_threads(
return {"status": "deleted"}
+@router.get("/attachments")
+def list_attachments(
+ limit: Annotated[int, Query(ge = 1, le = 100)] = 50,
+ offset: Annotated[int, Query(ge = 0)] = 0,
+ current_subject: str = Depends(get_current_subject),
+) -> dict:
+ """One bounded page of chat uploads for the settings Data tab."""
+ attachments, next_offset = list_chat_attachments_page(limit = limit, offset = offset)
+ return {"attachments": attachments, "nextOffset": next_offset}
+
+
+def _decode_attachment_base64(payload: str) -> bytes:
+ """Strict base64 decode of a stored payload.
+
+ Normalizes first: strips whitespace, fixes padding, accepts the URL-safe
+ alphabet. validate=False would silently drop bad characters and serve
+ corrupted bytes instead of failing, so raise 422 on anything else.
+ """
+ import base64
+
+ normalized = "".join(payload.split())
+ altchars = b"-_" if ("-" in normalized or "_" in normalized) else None
+ normalized += "=" * (-len(normalized) % 4)
+ try:
+ return base64.b64decode(normalized, altchars = altchars, validate = True)
+ except Exception as exc: # noqa: BLE001 - corrupt stored payload
+ raise HTTPException(status_code = 422, detail = "Attachment data is corrupt") from exc
+
+
+_AUDIO_FORMAT_MEDIA_TYPES = {
+ "mp3": "audio/mpeg",
+ "wav": "audio/wav",
+ "ogg": "audio/ogg",
+ "flac": "audio/flac",
+}
+
+
+def _safe_image_media_type(media_type: str) -> str:
+ """Clamp a data-URL media type to something inert to render.
+
+ Imported chats store image parts verbatim, so the embedded type can be
+ text/html or image/svg+xml; echoing those would execute markup with the
+ app origin when opened. Anything not a plain raster type downloads as
+ bytes instead.
+ """
+ lowered = media_type.strip().lower()
+ if lowered.startswith("image/") and lowered != "image/svg+xml":
+ return lowered
+ return "application/octet-stream"
+
+
+@router.get("/attachments/{message_id}/{attachment_id}/file")
+def get_attachment_file(
+ message_id: str,
+ attachment_id: str,
+ current_subject: str = Depends(get_current_subject),
+):
+ """Serve one attachment's stored content: image or audio bytes, or
+ extracted text."""
+ import urllib.parse
+
+ from fastapi.responses import Response
+
+ attachment = get_chat_attachment(message_id, attachment_id)
+ if attachment is None:
+ raise HTTPException(status_code = 404, detail = "Attachment not found")
+
+ attachment_content_type = attachment.get("contentType")
+ texts: list[str] = []
+ for part in attachment.get("content") or []:
+ if not isinstance(part, dict):
+ continue
+ image = part.get("image")
+ if isinstance(image, str) and image[:5].lower() == "data:":
+ header, _, payload = image.partition(",")
+ media_type = _safe_image_media_type(
+ header[5:].split(";", 1)[0] or "application/octet-stream"
+ )
+ if "base64" not in header.lower():
+ # RFC 2397 non-base64 form stores percent-encoded bytes.
+ data = urllib.parse.unquote_to_bytes(payload)
+ return Response(content = data, media_type = media_type)
+ data = _decode_attachment_base64(payload)
+ return Response(content = data, media_type = media_type)
+ # Audio parts: the attachment adapter stores {data, format} with raw
+ # base64; compare chats store a bare base64 string.
+ audio = part.get("audio")
+ if isinstance(audio, dict) or (isinstance(audio, str) and audio):
+ if isinstance(audio, dict):
+ payload = audio.get("data")
+ audio_format = audio.get("format")
+ else:
+ payload = audio.rsplit(",", 1)[-1]
+ audio_format = None
+ if isinstance(payload, str) and payload:
+ data = _decode_attachment_base64(payload)
+ media_type = (
+ attachment_content_type
+ if isinstance(attachment_content_type, str)
+ and attachment_content_type.startswith("audio/")
+ else _AUDIO_FORMAT_MEDIA_TYPES.get(
+ str(audio_format or "").lower(), "application/octet-stream"
+ )
+ )
+ return Response(content = data, media_type = media_type)
+ text = part.get("text")
+ if isinstance(text, str) and text:
+ texts.append(text)
+ if texts:
+ return Response(content = "\n".join(texts), media_type = "text/plain; charset=utf-8")
+ raise HTTPException(status_code = 404, detail = "Attachment has no stored content")
+
+
+@router.delete("/attachments/{message_id}/{attachment_id}")
+def delete_attachment(
+ message_id: str,
+ attachment_id: str,
+ current_subject: str = Depends(get_current_subject),
+) -> dict:
+ """Remove one attachment from its chat message."""
+ if not delete_chat_attachment(message_id, attachment_id):
+ raise HTTPException(status_code = 404, detail = "Attachment not found")
+ return {"ok": True}
+
+
@router.get("/projects", response_model = ChatProjectListResponse)
async def list_projects(
include_archived: bool = Query(False), current_subject: str = Depends(get_current_subject)
@@ -409,7 +537,7 @@ async def get_thread_message(
@router.put("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage)
-async def save_thread_message(
+def save_thread_message(
thread_id: str,
message_id: str,
payload: ChatMessage,
@@ -432,7 +560,7 @@ async def save_thread_message(
@router.put("/threads/{thread_id}/messages", response_model = ChatMessageListResponse)
-async def replace_thread_messages(
+def replace_thread_messages(
thread_id: str,
payload: ChatMessageSyncRequest,
current_subject: str = Depends(get_current_subject),
diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py
index e20fea74a3..392a4e0d02 100644
--- a/studio/backend/routes/rag.py
+++ b/studio/backend/routes/rag.py
@@ -318,6 +318,39 @@ def list_project_documents(project_id: str, subject: str = Depends(get_current_s
conn.close()
+@router.get("/documents")
+def list_all_uploaded_documents(subject: str = Depends(get_current_subject)) -> dict:
+ """Every uploaded file across chats, projects, and knowledge bases (settings
+ Data tab)."""
+ _require_rag()
+ conn = rag_db.get_connection()
+ try:
+ docs = store.list_all_documents(conn)
+ kb_names = {kb["id"]: kb["name"] for kb in store.list_kbs(conn)}
+ finally:
+ conn.close()
+
+ from storage.studio_db import list_chat_projects
+
+ project_names = {p["id"]: p["name"] for p in list_chat_projects(include_archived = True)}
+
+ out = []
+ for doc in docs:
+ view = _doc_view(doc)
+ stored_path = doc.get("stored_path")
+ size = None
+ if stored_path:
+ try:
+ size = os.path.getsize(stored_path)
+ except OSError:
+ size = None
+ view["sizeBytes"] = size
+ view["kbName"] = kb_names.get(doc.get("kb_id"))
+ view["projectName"] = project_names.get(doc.get("project_id"))
+ out.append(view)
+ return {"documents": out}
+
+
@router.delete("/documents/{document_id}")
def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict:
_require_rag()
@@ -424,8 +457,10 @@ _CONTENT_TYPES = {
".txt": "text/plain; charset=utf-8",
".md": "text/markdown; charset=utf-8",
".markdown": "text/markdown; charset=utf-8",
- ".html": "text/html; charset=utf-8",
- ".htm": "text/html; charset=utf-8",
+ # Served as plain text, never text/html: an uploaded HTML document rendered
+ # same-origin would execute its scripts with access to the app's storage.
+ ".html": "text/plain; charset=utf-8",
+ ".htm": "text/plain; charset=utf-8",
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
}
diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py
index 4e0c711b69..d889894d04 100644
--- a/studio/backend/storage/studio_db.py
+++ b/studio/backend/storage/studio_db.py
@@ -7,6 +7,7 @@ Like auth/storage.py (module-level functions, raw sqlite3, per-function
connections) plus WAL mode and PRAGMA foreign_keys = ON for CASCADE deletes.
"""
+import hashlib
import json
import logging
import os
@@ -100,6 +101,7 @@ _schema_lock = threading.Lock()
_schema_ready = False
_SQLITE_IN_CHUNK_SIZE = 900
_PROJECT_WORKSPACE_SUBDIRS = ("sandbox",)
+_CHAT_ATTACHMENT_INVENTORY_VERSION = 1
def _project_slug(name: str) -> str:
@@ -313,6 +315,141 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
)
"""
)
+ tombstone_schema = """
+ CREATE TABLE chat_attachment_tombstones (
+ thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE,
+ message_id TEXT NOT NULL,
+ attachment_id TEXT NOT NULL,
+ deleted_at INTEGER NOT NULL,
+ PRIMARY KEY(thread_id, message_id, attachment_id)
+ ) WITHOUT ROWID
+ """
+ tombstone_table = conn.execute(
+ """
+ SELECT 1 FROM sqlite_master
+ WHERE type = 'table' AND name = 'chat_attachment_tombstones'
+ """
+ ).fetchone()
+ if tombstone_table is None:
+ conn.execute(tombstone_schema)
+ else:
+ tombstone_columns = {
+ row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_tombstones)")
+ }
+ tombstone_fk_targets = {
+ row[2] for row in conn.execute("PRAGMA foreign_key_list(chat_attachment_tombstones)")
+ }
+ if "thread_id" not in tombstone_columns or "chat_threads" not in tombstone_fk_targets:
+ # The first implementation cascaded through chat_messages, which
+ # erased deletion knowledge during pruneMissing. Rebuild once,
+ # retaining every tombstone whose owning thread still exists.
+ conn.execute("SAVEPOINT migrate_chat_attachment_tombstones")
+ try:
+ conn.execute(
+ "ALTER TABLE chat_attachment_tombstones "
+ "RENAME TO chat_attachment_tombstones_legacy"
+ )
+ conn.execute(tombstone_schema)
+ if "thread_id" in tombstone_columns:
+ conn.execute(
+ """
+ INSERT OR IGNORE INTO chat_attachment_tombstones
+ (thread_id, message_id, attachment_id, deleted_at)
+ SELECT legacy.thread_id, legacy.message_id,
+ legacy.attachment_id, legacy.deleted_at
+ FROM chat_attachment_tombstones_legacy legacy
+ JOIN chat_threads thread ON thread.id = legacy.thread_id
+ """
+ )
+ else:
+ conn.execute(
+ """
+ INSERT OR IGNORE INTO chat_attachment_tombstones
+ (thread_id, message_id, attachment_id, deleted_at)
+ SELECT message.thread_id, legacy.message_id,
+ legacy.attachment_id, legacy.deleted_at
+ FROM chat_attachment_tombstones_legacy legacy
+ JOIN chat_messages message ON message.id = legacy.message_id
+ """
+ )
+ conn.execute("DROP TABLE chat_attachment_tombstones_legacy")
+ conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones")
+ except Exception:
+ conn.execute("ROLLBACK TO SAVEPOINT migrate_chat_attachment_tombstones")
+ conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones")
+ raise
+ conn.execute(
+ """
+ CREATE TABLE IF NOT EXISTS chat_attachment_inventory (
+ message_id TEXT NOT NULL REFERENCES chat_messages(id) ON DELETE CASCADE,
+ attachment_id TEXT NOT NULL,
+ name TEXT NOT NULL,
+ type TEXT,
+ content_type TEXT,
+ size_bytes INTEGER,
+ PRIMARY KEY(message_id, attachment_id)
+ ) WITHOUT ROWID
+ """
+ )
+ conn.execute(
+ """
+ CREATE TABLE IF NOT EXISTS chat_attachment_inventory_state (
+ singleton INTEGER NOT NULL PRIMARY KEY CHECK(singleton = 1),
+ inventory_version INTEGER NOT NULL DEFAULT 0,
+ dirty INTEGER NOT NULL DEFAULT 1,
+ backfilled_at INTEGER NOT NULL
+ )
+ """
+ )
+ inventory_state_columns = {
+ row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_inventory_state)")
+ }
+ if "inventory_version" not in inventory_state_columns:
+ conn.execute(
+ "ALTER TABLE chat_attachment_inventory_state "
+ "ADD COLUMN inventory_version INTEGER NOT NULL DEFAULT 0"
+ )
+ if "dirty" not in inventory_state_columns:
+ conn.execute(
+ "ALTER TABLE chat_attachment_inventory_state "
+ "ADD COLUMN dirty INTEGER NOT NULL DEFAULT 1"
+ )
+ conn.execute(
+ """
+ CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_insert
+ AFTER INSERT ON chat_messages
+ BEGIN
+ INSERT INTO chat_attachment_inventory_state
+ (singleton, inventory_version, dirty, backfilled_at)
+ VALUES (1, 0, 1, 0)
+ ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
+ END
+ """
+ )
+ conn.execute(
+ """
+ CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_update
+ AFTER UPDATE ON chat_messages
+ BEGIN
+ INSERT INTO chat_attachment_inventory_state
+ (singleton, inventory_version, dirty, backfilled_at)
+ VALUES (1, 0, 1, 0)
+ ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
+ END
+ """
+ )
+ conn.execute(
+ """
+ CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_delete
+ AFTER DELETE ON chat_messages
+ BEGIN
+ INSERT INTO chat_attachment_inventory_state
+ (singleton, inventory_version, dirty, backfilled_at)
+ VALUES (1, 0, 1, 0)
+ ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
+ END
+ """
+ )
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)"
)
@@ -391,6 +528,21 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_prompt_lists_created_at ON prompt_lists(created_at)"
)
+ inventory_state = conn.execute(
+ """
+ SELECT inventory_version, dirty
+ FROM chat_attachment_inventory_state
+ WHERE singleton = 1
+ """
+ ).fetchone()
+ if (
+ inventory_state is None
+ or inventory_state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION
+ or inventory_state["dirty"]
+ ):
+ _rebuild_chat_attachment_inventory(conn)
+ _mark_chat_attachment_inventory_clean(conn)
+ conn.commit()
def _prompt_entry_from_row(row: sqlite3.Row) -> dict:
@@ -1219,7 +1371,14 @@ def delete_chat_threads(ids: list[str]) -> None:
return
conn = get_connection()
try:
+ conn.execute("BEGIN IMMEDIATE")
+ _ensure_chat_attachment_inventory_current(conn)
+ conn.executemany(
+ "DELETE FROM chat_attachment_tombstones WHERE thread_id = ?",
+ [(id,) for id in ids],
+ )
conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids])
+ _mark_chat_attachment_inventory_clean(conn)
conn.commit()
finally:
conn.close()
@@ -1228,7 +1387,11 @@ def delete_chat_threads(ids: list[str]) -> None:
def clear_chat_history() -> None:
conn = get_connection()
try:
+ conn.execute("BEGIN IMMEDIATE")
+ _ensure_chat_attachment_inventory_current(conn)
+ conn.execute("DELETE FROM chat_attachment_tombstones")
conn.execute("DELETE FROM chat_threads")
+ _mark_chat_attachment_inventory_clean(conn)
conn.commit()
finally:
conn.close()
@@ -1354,6 +1517,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
+ _ensure_chat_attachment_inventory_current(conn)
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
if row is None:
conn.rollback()
@@ -1361,6 +1525,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]:
project = _chat_project_from_row(row)
conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,))
conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,))
+ _mark_chat_attachment_inventory_clean(conn)
conn.commit()
if delete_files:
_delete_project_workspace(project)
@@ -1483,15 +1648,285 @@ def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str)
)
+_CONTENT_PART_ID_PREFIX = "content-part-sha256-"
+_URI_SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*:")
+
+
+def _is_locally_stored_blob(value: str) -> bool:
+ """True for data URIs or bare base64, never external/blob URI references."""
+ candidate = value.lstrip()
+ if not candidate:
+ return False
+ if candidate[:5].lower() == "data:":
+ return True
+ if candidate.startswith(("//", "\\\\")):
+ return False
+ return _URI_SCHEME_RE.match(candidate) is None
+
+
+def _managed_content_part_payload(part: dict) -> Optional[tuple[str, Any]]:
+ """Return the locally stored blob payload used to identify a content part."""
+ image = part.get("image")
+ if isinstance(image, str) and image[:5].lower() == "data:":
+ return "image", image
+
+ audio = part.get("audio")
+ if isinstance(audio, str) and _is_locally_stored_blob(audio):
+ return "audio", audio
+ if isinstance(audio, dict):
+ data = audio.get("data")
+ if isinstance(data, str) and _is_locally_stored_blob(data):
+ return "audio", audio
+ return None
+
+
+def _content_part_id(part: dict) -> Optional[str]:
+ """Stable managed id derived from blob data, without mutating inference content."""
+ payload = _managed_content_part_payload(part)
+ if payload is None:
+ return None
+ canonical = json.dumps(
+ payload,
+ ensure_ascii = False,
+ separators = (",", ":"),
+ sort_keys = True,
+ ).encode("utf-8")
+ return f"{_CONTENT_PART_ID_PREFIX}{hashlib.sha256(canonical).hexdigest()}"
+
+
+def _chat_attachment_tombstones_for_messages(
+ conn: sqlite3.Connection, thread_id: str, message_ids: list[str]
+) -> dict[str, set[str]]:
+ tombstones = {message_id: set() for message_id in message_ids}
+ unique_ids = list(dict.fromkeys(message_ids))
+ for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE):
+ chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
+ placeholders = ",".join("?" for _ in chunk)
+ rows = conn.execute(
+ f"""
+ SELECT message_id, attachment_id
+ FROM chat_attachment_tombstones
+ WHERE thread_id = ? AND message_id IN ({placeholders})
+ """,
+ (thread_id, *chunk),
+ ).fetchall()
+ for row in rows:
+ tombstones[row["message_id"]].add(row["attachment_id"])
+ return tombstones
+
+
+def _reconcile_chat_message_uploads(message: dict, tombstones: set[str]) -> dict:
+ """Strip uploads previously deleted through the Data tab from a stale write."""
+ if not tombstones:
+ return message
+
+ reconciled = dict(message)
+ attachments = message.get("attachments")
+ if isinstance(attachments, list):
+ reconciled["attachments"] = [
+ attachment
+ for attachment in attachments
+ if not (isinstance(attachment, dict) and str(attachment.get("id") or "") in tombstones)
+ ]
+
+ content = message.get("content")
+ if isinstance(content, list):
+ reconciled["content"] = [
+ part
+ for part in content
+ if not (isinstance(part, dict) and (_content_part_id(part) or "") in tombstones)
+ ]
+ return reconciled
+
+
+def _chat_attachment_metadata_text(value, fallback: Optional[str] = None) -> Optional[str]:
+ """Keep untyped legacy/import metadata safe for SQLite binding."""
+ if value is None:
+ return fallback
+ if isinstance(value, str):
+ return value or fallback
+ if isinstance(value, (bool, int, float)):
+ return str(value)
+ # Objects and arrays are not useful display metadata and sqlite3 rejects
+ # binding them directly.
+ return fallback
+
+
+def _chat_attachment_inventory_entries(
+ attachments_json: Optional[str],
+ content_json: Optional[str],
+ tombstones: Optional[set[str]] = None,
+) -> list[dict]:
+ tombstones = tombstones or set()
+ attachments = _json_loads(attachments_json, None)
+ if not isinstance(attachments, list):
+ attachments = []
+ attachments = [
+ attachment
+ for attachment in attachments
+ if isinstance(attachment, dict) and attachment.get("id")
+ ]
+ attachments.extend(_content_part_attachments(content_json))
+
+ entries: list[dict] = []
+ seen: set[str] = set()
+ for attachment in attachments:
+ attachment_id = str(attachment["id"])
+ if attachment_id in seen or attachment_id in tombstones:
+ continue
+ seen.add(attachment_id)
+ entries.append(
+ {
+ "id": attachment_id,
+ "name": _chat_attachment_metadata_text(attachment.get("name"), "attachment"),
+ "type": _chat_attachment_metadata_text(attachment.get("type")),
+ "contentType": _chat_attachment_metadata_text(attachment.get("contentType")),
+ "sizeBytes": _chat_attachment_size_bytes(attachment),
+ }
+ )
+ return entries
+
+
+def _replace_chat_attachment_inventory(
+ conn: sqlite3.Connection,
+ message_id: str,
+ attachments_json: Optional[str],
+ content_json: Optional[str],
+ tombstones: Optional[set[str]] = None,
+) -> None:
+ conn.execute("DELETE FROM chat_attachment_inventory WHERE message_id = ?", (message_id,))
+ entries = _chat_attachment_inventory_entries(
+ attachments_json,
+ content_json,
+ tombstones,
+ )
+ conn.executemany(
+ """
+ INSERT INTO chat_attachment_inventory
+ (message_id, attachment_id, name, type, content_type, size_bytes)
+ VALUES (?, ?, ?, ?, ?, ?)
+ """,
+ [
+ (
+ message_id,
+ entry["id"],
+ entry["name"],
+ entry["type"],
+ entry["contentType"],
+ entry["sizeBytes"],
+ )
+ for entry in entries
+ ],
+ )
+
+
+def _mark_chat_attachment_inventory_clean(conn: sqlite3.Connection) -> None:
+ conn.execute(
+ """
+ INSERT INTO chat_attachment_inventory_state
+ (singleton, inventory_version, dirty, backfilled_at)
+ VALUES (1, ?, 0, ?)
+ ON CONFLICT(singleton) DO UPDATE SET
+ inventory_version = excluded.inventory_version,
+ dirty = 0,
+ backfilled_at = excluded.backfilled_at
+ """,
+ (
+ _CHAT_ATTACHMENT_INVENTORY_VERSION,
+ int(datetime.now(timezone.utc).timestamp() * 1000),
+ ),
+ )
+
+
+def _rebuild_chat_attachment_inventory(conn: sqlite3.Connection) -> None:
+ """Rebuild after schema upgrade or a write from an older Studio build."""
+ conn.execute("DELETE FROM chat_attachment_inventory")
+ tombstones: dict[tuple[str, str], set[str]] = {}
+ for row in conn.execute(
+ "SELECT thread_id, message_id, attachment_id FROM chat_attachment_tombstones"
+ ).fetchall():
+ tombstones.setdefault((row["thread_id"], row["message_id"]), set()).add(
+ row["attachment_id"]
+ )
+ rows = conn.execute(
+ "SELECT id, thread_id, attachments_json, content_json FROM chat_messages"
+ ).fetchall()
+ for row in rows:
+ _replace_chat_attachment_inventory(
+ conn,
+ row["id"],
+ row["attachments_json"],
+ row["content_json"],
+ tombstones.get((row["thread_id"], row["id"]), set()),
+ )
+
+
+def _ensure_chat_attachment_inventory_current(conn: sqlite3.Connection) -> None:
+ state = conn.execute(
+ """
+ SELECT inventory_version, dirty
+ FROM chat_attachment_inventory_state
+ WHERE singleton = 1
+ """
+ ).fetchone()
+ if (
+ state is not None
+ and state["inventory_version"] == _CHAT_ATTACHMENT_INVENTORY_VERSION
+ and not state["dirty"]
+ ):
+ return
+
+ owns_transaction = not conn.in_transaction
+ if owns_transaction:
+ conn.execute("BEGIN IMMEDIATE")
+ try:
+ state = conn.execute(
+ """
+ SELECT inventory_version, dirty
+ FROM chat_attachment_inventory_state
+ WHERE singleton = 1
+ """
+ ).fetchone()
+ if (
+ state is None
+ or state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION
+ or state["dirty"]
+ ):
+ _rebuild_chat_attachment_inventory(conn)
+ _mark_chat_attachment_inventory_clean(conn)
+ if owns_transaction:
+ conn.commit()
+ except Exception:
+ if owns_transaction:
+ conn.rollback()
+ raise
+
+
def upsert_chat_message(message: dict) -> dict:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
+ _ensure_chat_attachment_inventory_current(conn)
_raise_if_chat_message_thread_conflicts(
conn,
message["threadId"],
[message["id"]],
)
+ tombstones = _chat_attachment_tombstones_for_messages(
+ conn,
+ message["threadId"],
+ [message["id"]],
+ )
+ reconciled = _reconcile_chat_message_uploads(
+ message,
+ tombstones.get(message["id"], set()),
+ )
+ content_json = json.dumps(reconciled.get("content", []))
+ attachments_json = (
+ json.dumps(reconciled.get("attachments"))
+ if reconciled.get("attachments") is not None
+ else None
+ )
conn.execute(
"""
INSERT INTO chat_messages
@@ -1507,23 +1942,32 @@ def upsert_chat_message(message: dict) -> dict:
WHERE excluded.thread_id = chat_messages.thread_id
""",
(
- message["id"],
- message["threadId"],
- message.get("parentId"),
- message["role"],
- json.dumps(message.get("content", [])),
- json.dumps(message.get("attachments"))
- if message.get("attachments") is not None
+ reconciled["id"],
+ reconciled["threadId"],
+ reconciled.get("parentId"),
+ reconciled["role"],
+ content_json,
+ attachments_json,
+ json.dumps(reconciled.get("metadata"))
+ if reconciled.get("metadata") is not None
else None,
- json.dumps(message.get("metadata"))
- if message.get("metadata") is not None
- else None,
- int(message["createdAt"]),
+ int(reconciled["createdAt"]),
),
)
- _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"]))
+ _replace_chat_attachment_inventory(
+ conn,
+ reconciled["id"],
+ attachments_json,
+ content_json,
+ )
+ _bump_chat_thread_updated_at(
+ conn,
+ reconciled["threadId"],
+ int(reconciled["createdAt"]),
+ )
+ _mark_chat_attachment_inventory_clean(conn)
conn.commit()
- return message
+ return reconciled
except Exception:
conn.rollback()
raise
@@ -1539,13 +1983,28 @@ def sync_chat_messages(
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
+ _ensure_chat_attachment_inventory_current(conn)
_raise_if_chat_message_thread_conflicts(
conn,
thread_id,
[m["id"] for m in messages],
)
- if prune_missing:
- conn.execute("DELETE FROM chat_messages WHERE thread_id = ?", (thread_id,))
+ tombstones = _chat_attachment_tombstones_for_messages(
+ conn,
+ thread_id,
+ [m["id"] for m in messages],
+ )
+ reconciled_messages = [
+ _reconcile_chat_message_uploads(m, tombstones.get(m["id"], set())) for m in messages
+ ]
+ serialized_messages = [
+ (
+ m,
+ json.dumps(m.get("content", [])),
+ json.dumps(m.get("attachments")) if m.get("attachments") is not None else None,
+ )
+ for m in reconciled_messages
+ ]
conn.executemany(
"""
INSERT INTO chat_messages
@@ -1566,20 +2025,46 @@ def sync_chat_messages(
thread_id,
m.get("parentId"),
m["role"],
- json.dumps(m.get("content", [])),
- json.dumps(m.get("attachments")) if m.get("attachments") is not None else None,
+ content_json,
+ attachments_json,
json.dumps(m.get("metadata")) if m.get("metadata") is not None else None,
int(m["createdAt"]),
)
- for m in messages
+ for m, content_json, attachments_json in serialized_messages
],
)
- if prune_missing:
- _recompute_chat_thread_updated_at(conn, thread_id)
- elif messages:
- _bump_chat_thread_updated_at(
- conn, thread_id, max(int(m["createdAt"]) for m in messages)
+ for m, content_json, attachments_json in serialized_messages:
+ _replace_chat_attachment_inventory(
+ conn,
+ m["id"],
+ attachments_json,
+ content_json,
)
+ if prune_missing:
+ retained_ids = {m["id"] for m in reconciled_messages}
+ existing_ids = {
+ row["id"]
+ for row in conn.execute(
+ "SELECT id FROM chat_messages WHERE thread_id = ?",
+ (thread_id,),
+ ).fetchall()
+ }
+ missing_ids = sorted(existing_ids - retained_ids)
+ for start in range(0, len(missing_ids), _SQLITE_IN_CHUNK_SIZE):
+ chunk = missing_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
+ placeholders = ",".join("?" for _ in chunk)
+ conn.execute(
+ f"DELETE FROM chat_messages WHERE thread_id = ? AND id IN ({placeholders})",
+ (thread_id, *chunk),
+ )
+ _recompute_chat_thread_updated_at(conn, thread_id)
+ elif reconciled_messages:
+ _bump_chat_thread_updated_at(
+ conn,
+ thread_id,
+ max(int(m["createdAt"]) for m in reconciled_messages),
+ )
+ _mark_chat_attachment_inventory_clean(conn)
conn.commit()
return list_chat_messages(thread_id)
except ChatMessageConflictError:
@@ -1613,6 +2098,7 @@ def fork_chat_thread(
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
+ _ensure_chat_attachment_inventory_current(conn)
src = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,)
).fetchone()
@@ -1686,6 +2172,14 @@ def fork_chat_thread(
for row in ancestry
],
)
+ for row in ancestry:
+ _replace_chat_attachment_inventory(
+ conn,
+ id_map[row["id"]],
+ row["attachments_json"],
+ row["content_json"],
+ )
+ _mark_chat_attachment_inventory_clean(conn)
conn.commit()
thread_row = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,)
@@ -1744,6 +2238,279 @@ def get_chat_message(thread_id: str, message_id: str) -> Optional[dict]:
conn.close()
+def _blob_part_base64_len(part: dict) -> int:
+ """Base64 payload length of an image or audio content part, or 0."""
+ image = part.get("image")
+ if isinstance(image, str) and image[:5].lower() == "data:":
+ return len(image.rsplit(",", 1)[-1])
+ audio = part.get("audio")
+ if isinstance(audio, str) and _is_locally_stored_blob(audio):
+ return len(audio.rsplit(",", 1)[-1])
+ if isinstance(audio, dict):
+ data = audio.get("data")
+ if isinstance(data, str) and _is_locally_stored_blob(data):
+ return len(data)
+ return 0
+
+
+def _chat_attachment_size_bytes(attachment: dict) -> Optional[int]:
+ """Approximate stored size of one attachment's content parts.
+
+ Image and audio parts hold base64 payloads (decoded bytes ~= 3/4 of the
+ encoded length); text parts count their character length. None when there
+ is no sizable content (e.g. a stripped/legacy attachment).
+ """
+ total = 0
+ found = False
+ for part in attachment.get("content") or []:
+ if not isinstance(part, dict):
+ continue
+ blob_len = _blob_part_base64_len(part)
+ if blob_len > 0:
+ total += (blob_len * 3) // 4
+ found = True
+ continue
+ text = part.get("text")
+ if isinstance(text, str) and text:
+ total += len(text.encode("utf-8", errors = "ignore"))
+ found = True
+ return total if found else None
+
+
+def _content_part_attachments(content_json: Optional[str]) -> list[dict]:
+ """Managed local blobs stored in content_json, with stable payload ids.
+
+ Exact duplicate blobs intentionally share one inventory id. Deleting that
+ id removes every identical copy, avoiding ambiguous index-based addressing.
+ """
+ content = _json_loads(content_json, None)
+ if not isinstance(content, list):
+ return []
+ out: list[dict] = []
+ seen: set[str] = set()
+ for part in content:
+ if not isinstance(part, dict):
+ continue
+ attachment_id = _content_part_id(part)
+ payload = _managed_content_part_payload(part)
+ if attachment_id is None or payload is None or attachment_id in seen:
+ continue
+ seen.add(attachment_id)
+ kind, value = payload
+ content_type = None
+ if kind == "image" and isinstance(value, str):
+ content_type = value[5:].split(";", 1)[0].split(",", 1)[0] or None
+ out.append(
+ {
+ "id": attachment_id,
+ "type": kind,
+ "name": "Chat image" if kind == "image" else "Chat audio",
+ "contentType": content_type,
+ "content": [part],
+ }
+ )
+ return out
+
+
+def list_chat_attachments_page(
+ limit: int = 50, offset: int = 0
+) -> tuple[list[dict], Optional[int]]:
+ """One bounded page from the normalized attachment inventory."""
+ if not 1 <= limit <= 100:
+ raise ValueError("limit must be between 1 and 100")
+ if offset < 0:
+ raise ValueError("offset must be non-negative")
+
+ conn = get_connection()
+ try:
+ _ensure_chat_attachment_inventory_current(conn)
+ rows = conn.execute(
+ """
+ SELECT i.attachment_id, i.name, i.type, i.content_type,
+ i.size_bytes, m.id AS message_id, m.thread_id,
+ m.created_at, t.title AS thread_title, t.pair_id
+ FROM chat_attachment_inventory i
+ JOIN chat_messages m ON m.id = i.message_id
+ LEFT JOIN chat_threads t ON t.id = m.thread_id
+ ORDER BY m.created_at DESC, m.id ASC, i.attachment_id ASC
+ LIMIT ? OFFSET ?
+ """,
+ (limit + 1, offset),
+ ).fetchall()
+ finally:
+ conn.close()
+
+ has_more = len(rows) > limit
+ page_rows = rows[:limit]
+ attachments = [
+ {
+ "id": row["attachment_id"],
+ "messageId": row["message_id"],
+ "threadId": row["thread_id"],
+ "pairId": row["pair_id"],
+ "threadTitle": row["thread_title"],
+ "name": row["name"],
+ "type": row["type"],
+ "contentType": row["content_type"],
+ "sizeBytes": row["size_bytes"],
+ "createdAt": row["created_at"],
+ }
+ for row in page_rows
+ ]
+ return attachments, offset + limit if has_more else None
+
+
+def list_chat_attachments() -> list[dict]:
+ """Compatibility helper returning the full normalized inventory."""
+ attachments: list[dict] = []
+ offset = 0
+ while True:
+ page, next_offset = list_chat_attachments_page(limit = 100, offset = offset)
+ attachments.extend(page)
+ if next_offset is None:
+ return attachments
+ offset = next_offset
+
+
+def get_chat_attachment(message_id: str, attachment_id: str) -> Optional[dict]:
+ """One attachment record (full content) from a message, or None."""
+ conn = get_connection()
+ try:
+ row = conn.execute(
+ """
+ SELECT message.attachments_json, message.content_json,
+ EXISTS(
+ SELECT 1 FROM chat_attachment_tombstones tombstone
+ WHERE tombstone.thread_id = message.thread_id
+ AND tombstone.message_id = message.id
+ AND tombstone.attachment_id = ?
+ ) AS tombstoned
+ FROM chat_messages message
+ WHERE message.id = ?
+ """,
+ (attachment_id, message_id),
+ ).fetchone()
+ finally:
+ conn.close()
+ if row is None or row["tombstoned"]:
+ return None
+ attachments = _json_loads(row["attachments_json"], None)
+ if isinstance(attachments, list):
+ for attachment in attachments:
+ if isinstance(attachment, dict) and str(attachment.get("id") or "") == attachment_id:
+ return attachment
+ if attachment_id.startswith(_CONTENT_PART_ID_PREFIX):
+ for attachment in _content_part_attachments(row["content_json"]):
+ if attachment["id"] == attachment_id:
+ return attachment
+ return None
+
+
+def _record_chat_attachment_tombstone(
+ conn: sqlite3.Connection, thread_id: str, message_id: str, attachment_id: str
+) -> None:
+ conn.execute(
+ """
+ INSERT INTO chat_attachment_tombstones
+ (thread_id, message_id, attachment_id, deleted_at)
+ VALUES (?, ?, ?, ?)
+ ON CONFLICT(thread_id, message_id, attachment_id) DO UPDATE SET
+ deleted_at = excluded.deleted_at
+ """,
+ (
+ thread_id,
+ message_id,
+ attachment_id,
+ int(datetime.now(timezone.utc).timestamp() * 1000),
+ ),
+ )
+
+
+def delete_chat_attachment(message_id: str, attachment_id: str) -> bool:
+ """Remove one stored upload from a message.
+
+ The tombstone is retained while the thread exists, so pruning and later
+ recreating the same message id cannot restore the deleted upload. If an
+ ordinary attachment id collides with a content-blob id, both are deleted as
+ one managed item.
+ """
+ conn = get_connection()
+ try:
+ conn.execute("BEGIN IMMEDIATE")
+ _ensure_chat_attachment_inventory_current(conn)
+ row = conn.execute(
+ """
+ SELECT thread_id, attachments_json, content_json
+ FROM chat_messages WHERE id = ?
+ """,
+ (message_id,),
+ ).fetchone()
+ if row is None:
+ conn.rollback()
+ return False
+
+ attachments = _json_loads(row["attachments_json"], None)
+ updated_attachments_json = row["attachments_json"]
+ deleted_attachment = False
+ if isinstance(attachments, list):
+ remaining_attachments = [
+ attachment
+ for attachment in attachments
+ if not (
+ isinstance(attachment, dict)
+ and str(attachment.get("id") or "") == attachment_id
+ )
+ ]
+ deleted_attachment = len(remaining_attachments) != len(attachments)
+ if deleted_attachment:
+ updated_attachments_json = json.dumps(remaining_attachments)
+
+ content = _json_loads(row["content_json"], None)
+ updated_content_json = row["content_json"]
+ deleted_content = False
+ if attachment_id.startswith(_CONTENT_PART_ID_PREFIX) and isinstance(content, list):
+ remaining_content = [
+ part
+ for part in content
+ if not (isinstance(part, dict) and _content_part_id(part) == attachment_id)
+ ]
+ deleted_content = len(remaining_content) != len(content)
+ if deleted_content:
+ updated_content_json = json.dumps(remaining_content)
+
+ if not deleted_attachment and not deleted_content:
+ conn.rollback()
+ return False
+ conn.execute(
+ """
+ UPDATE chat_messages
+ SET attachments_json = ?, content_json = ?
+ WHERE id = ?
+ """,
+ (updated_attachments_json, updated_content_json, message_id),
+ )
+ _record_chat_attachment_tombstone(
+ conn,
+ row["thread_id"],
+ message_id,
+ attachment_id,
+ )
+ _replace_chat_attachment_inventory(
+ conn,
+ message_id,
+ updated_attachments_json,
+ updated_content_json,
+ )
+ _mark_chat_attachment_inventory_clean(conn)
+ conn.commit()
+ return True
+ except Exception:
+ conn.rollback()
+ raise
+ finally:
+ conn.close()
+
+
def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]:
if not thread_ids:
return []
diff --git a/studio/backend/tests/test_chat_attachments.py b/studio/backend/tests/test_chat_attachments.py
new file mode 100644
index 0000000000..459587ca9e
--- /dev/null
+++ b/studio/backend/tests/test_chat_attachments.py
@@ -0,0 +1,634 @@
+# 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 base64
+import json
+import os
+import sqlite3
+import sys
+
+import pytest
+from fastapi import HTTPException
+
+_backend = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, _backend)
+
+from routes import chat_history
+from storage import studio_db
+from utils.paths import studio_db_path
+
+PNG_BYTES = base64.b64decode(
+ "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
+)
+PNG_DATA_URL = "data:image/png;base64," + base64.b64encode(PNG_BYTES).decode("ascii")
+
+
+def _reset_studio_db(tmp_path, monkeypatch):
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setenv("UNSLOTH_STUDIO_PROJECTS_HOME", str(tmp_path / "Projects"))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+
+def _thread(
+ thread_id: str = "thread-1",
+ title: str = "Test Chat",
+ pair_id: str | None = None,
+) -> dict:
+ return {
+ "id": thread_id,
+ "title": title,
+ "modelType": "base",
+ "modelId": "test-model",
+ "pairId": pair_id,
+ "archived": False,
+ "createdAt": 1_700_000_000_000,
+ }
+
+
+def _message(
+ message_id: str,
+ created_at: int = 1_700_000_000_000,
+ attachments = None,
+ thread_id: str = "thread-1",
+) -> dict:
+ message = {
+ "id": message_id,
+ "threadId": thread_id,
+ "parentId": None,
+ "role": "user",
+ "content": [{"type": "text", "text": "hello"}],
+ "createdAt": created_at,
+ }
+ if attachments is not None:
+ message["attachments"] = attachments
+ return message
+
+
+def _image_attachment(attachment_id: str = "att-1", name: str = "photo.png") -> dict:
+ return {
+ "id": attachment_id,
+ "type": "image",
+ "name": name,
+ "contentType": "image/png",
+ "content": [{"type": "image", "image": PNG_DATA_URL}],
+ "status": {"type": "complete"},
+ }
+
+
+def _seed(
+ tmp_path,
+ monkeypatch,
+ attachments,
+ message_id: str = "msg-1",
+):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ studio_db.upsert_chat_message(_message(message_id, attachments = attachments))
+
+
+def _set_raw_attachments_json(message_id: str, raw: str) -> None:
+ conn = sqlite3.connect(studio_db_path())
+ try:
+ conn.execute(
+ "UPDATE chat_messages SET attachments_json = ? WHERE id = ?",
+ (raw, message_id),
+ )
+ conn.commit()
+ finally:
+ conn.close()
+
+
+def _raw_attachments_json(message_id: str):
+ conn = sqlite3.connect(studio_db_path())
+ try:
+ row = conn.execute(
+ "SELECT attachments_json FROM chat_messages WHERE id = ?",
+ (message_id,),
+ ).fetchone()
+ return row[0] if row is not None else None
+ finally:
+ conn.close()
+
+
+# ---------------------------------------------------------------------------
+# Storage: list_chat_attachments
+# ---------------------------------------------------------------------------
+
+
+def test_list_chat_attachments_empty_db(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ assert studio_db.list_chat_attachments() == []
+
+
+def test_list_chat_attachments_round_trip(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ records = studio_db.list_chat_attachments()
+ assert len(records) == 1
+ record = records[0]
+ assert record["id"] == "att-1"
+ assert record["messageId"] == "msg-1"
+ assert record["threadId"] == "thread-1"
+ assert record["threadTitle"] == "Test Chat"
+ assert record["name"] == "photo.png"
+ assert record["type"] == "image"
+ assert record["contentType"] == "image/png"
+ assert record["createdAt"] == 1_700_000_000_000
+ # Base64 length estimate is within padding error of the decoded size.
+ assert abs(record["sizeBytes"] - len(PNG_BYTES)) <= 2
+
+
+def test_list_chat_attachments_counts_text_utf8(tmp_path, monkeypatch):
+ text = "héllo wörld é世界"
+ attachment = {
+ "id": "att-txt",
+ "type": "document",
+ "name": "notes.txt",
+ "content": [{"type": "text", "text": text}],
+ }
+ _seed(tmp_path, monkeypatch, [attachment])
+ records = studio_db.list_chat_attachments()
+ assert records[0]["sizeBytes"] == len(text.encode("utf-8"))
+
+
+def test_list_chat_attachments_no_content_size_is_none(tmp_path, monkeypatch):
+ attachment = {"id": "att-empty", "name": "ghost.bin", "content": []}
+ _seed(tmp_path, monkeypatch, [attachment])
+ records = studio_db.list_chat_attachments()
+ assert records[0]["sizeBytes"] is None
+ assert records[0]["name"] == "ghost.bin"
+
+
+def test_list_chat_attachments_defaults_missing_name(tmp_path, monkeypatch):
+ attachment = {"id": "att-noname", "content": []}
+ _seed(tmp_path, monkeypatch, [attachment])
+ assert studio_db.list_chat_attachments()[0]["name"] == "attachment"
+
+
+def test_list_chat_attachments_sanitizes_structured_metadata(tmp_path, monkeypatch):
+ attachment = {
+ "id": "att-weird",
+ "name": {"nested": "name"},
+ "type": ["image"],
+ "contentType": {"mime": "image/png"},
+ "content": [],
+ }
+ _seed(tmp_path, monkeypatch, [attachment])
+ record = studio_db.list_chat_attachments()[0]
+ assert record["name"] == "attachment"
+ assert record["type"] is None
+ assert record["contentType"] is None
+
+
+def test_list_chat_attachments_skips_malformed_rows(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ for i, raw in enumerate(
+ [
+ "not json at all",
+ '{"id": "att-obj"}',
+ "null",
+ "[]",
+ '[{"noid": true}, "just a string", 42]',
+ '[{"id": ""}]',
+ ]
+ ):
+ message_id = f"msg-bad-{i}"
+ studio_db.upsert_chat_message(_message(message_id))
+ _set_raw_attachments_json(message_id, raw)
+ studio_db.upsert_chat_message(_message("msg-good", attachments = [_image_attachment("att-ok")]))
+ records = studio_db.list_chat_attachments()
+ assert [r["id"] for r in records] == ["att-ok"]
+
+
+def test_list_chat_attachments_orders_newest_first(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ studio_db.upsert_chat_message(
+ _message("msg-old", 1_700_000_000_000, [_image_attachment("att-old")])
+ )
+ studio_db.upsert_chat_message(
+ _message("msg-new", 1_700_000_100_000, [_image_attachment("att-new")])
+ )
+ assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-new", "att-old"]
+
+
+def test_list_chat_attachments_survives_missing_thread_row(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ studio_db.upsert_chat_message(_message("msg-1", attachments = [_image_attachment()]))
+ conn = sqlite3.connect(studio_db_path())
+ try:
+ conn.execute("DELETE FROM chat_threads WHERE id = 'thread-1'")
+ conn.commit()
+ finally:
+ conn.close()
+ records = studio_db.list_chat_attachments()
+ assert len(records) == 1
+ assert records[0]["threadTitle"] is None
+
+
+def test_list_chat_attachments_includes_compare_pair_id(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread(pair_id = "pair-1"))
+ studio_db.upsert_chat_message(_message("msg-compare", attachments = [_image_attachment()]))
+ record = studio_db.list_chat_attachments()[0]
+ assert record["threadId"] == "thread-1"
+ assert record["pairId"] == "pair-1"
+
+
+def test_list_chat_attachments_gone_after_thread_delete(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ studio_db.delete_chat_threads(["thread-1"])
+ assert studio_db.list_chat_attachments() == []
+
+
+# ---------------------------------------------------------------------------
+# Storage: get_chat_attachment / delete_chat_attachment
+# ---------------------------------------------------------------------------
+
+
+def test_get_chat_attachment_found_and_missing(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ attachment = studio_db.get_chat_attachment("msg-1", "att-1")
+ assert attachment is not None
+ assert attachment["content"][0]["image"] == PNG_DATA_URL
+ assert studio_db.get_chat_attachment("msg-1", "att-missing") is None
+ assert studio_db.get_chat_attachment("msg-missing", "att-1") is None
+
+
+def test_delete_chat_attachment_keeps_others(tmp_path, monkeypatch):
+ _seed(
+ tmp_path,
+ monkeypatch,
+ [_image_attachment("att-1"), _image_attachment("att-2", "other.png")],
+ )
+ assert studio_db.delete_chat_attachment("msg-1", "att-1") is True
+ assert studio_db.get_chat_attachment("msg-1", "att-1") is None
+ assert studio_db.get_chat_attachment("msg-1", "att-2") is not None
+ assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-2"]
+
+
+def test_delete_last_chat_attachment_stores_empty_list(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ assert studio_db.delete_chat_attachment("msg-1", "att-1") is True
+ # '[]' rather than NULL: a NULL attachments field reads back as missing
+ # and triggers the legacy IndexedDB backfill, resurrecting the deleted
+ # attachment on the next chat load.
+ assert _raw_attachments_json("msg-1") == "[]"
+ assert studio_db.list_chat_attachments() == []
+ # The message itself must survive with its content intact.
+ message = studio_db.get_chat_message("thread-1", "msg-1")
+ assert message is not None
+ assert message["content"] == [{"type": "text", "text": "hello"}]
+ assert message["attachments"] == []
+
+
+def test_delete_chat_attachment_missing_targets(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ assert studio_db.delete_chat_attachment("msg-missing", "att-1") is False
+ assert studio_db.delete_chat_attachment("msg-1", "att-missing") is False
+ _set_raw_attachments_json("msg-1", "not json")
+ assert studio_db.delete_chat_attachment("msg-1", "att-1") is False
+
+
+# ---------------------------------------------------------------------------
+# Routes: /attachments endpoints (real storage, direct calls)
+# ---------------------------------------------------------------------------
+
+
+def test_list_attachments_route(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ result = chat_history.list_attachments(current_subject = "unsloth")
+ assert [a["id"] for a in result["attachments"]] == ["att-1"]
+
+
+def test_attachment_file_serves_image_bytes(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert response.body == PNG_BYTES
+ assert response.media_type == "image/png"
+
+
+def test_attachment_file_tolerates_whitespace_in_base64(tmp_path, monkeypatch):
+ encoded = base64.b64encode(PNG_BYTES).decode("ascii")
+ wrapped = "\n".join(encoded[i : i + 8] for i in range(0, len(encoded), 8))
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + wrapped}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert response.body == PNG_BYTES
+
+
+def test_attachment_file_corrupt_base64_is_422(tmp_path, monkeypatch):
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "data:image/png;base64,%%%"}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ with pytest.raises(HTTPException) as excinfo:
+ chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert excinfo.value.status_code == 422
+
+
+def test_attachment_file_accepts_urlsafe_base64(tmp_path, monkeypatch):
+ data = bytes(range(251, 256)) * 3 # encodes to characters remapped by urlsafe
+ payload = base64.urlsafe_b64encode(data).decode("ascii")
+ assert "-" in payload or "_" in payload
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert response.body == data
+
+
+def test_attachment_file_accepts_missing_padding(tmp_path, monkeypatch):
+ payload = base64.b64encode(PNG_BYTES).decode("ascii").rstrip("=")
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert response.body == PNG_BYTES
+
+
+def test_attachment_file_serves_percent_encoded_data_url(tmp_path, monkeypatch):
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "data:text/plain,hello%20world"}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert response.body == b"hello world"
+ # Non-image data URL types are clamped so markup never renders same-origin.
+ assert response.media_type == "application/octet-stream"
+
+
+def test_attachment_file_serves_text_parts(tmp_path, monkeypatch):
+ attachment = {
+ "id": "att-txt",
+ "type": "document",
+ "name": "notes.txt",
+ "content": [
+ {"type": "text", "text": "first"},
+ {"type": "text", "text": "second"},
+ ],
+ }
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-txt", current_subject = "unsloth")
+ assert response.body.decode("utf-8") == "first\nsecond"
+ assert response.media_type.startswith("text/plain")
+
+
+def test_attachment_file_no_content_is_404(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [{"id": "att-empty", "name": "ghost", "content": []}])
+ with pytest.raises(HTTPException) as excinfo:
+ chat_history.get_attachment_file("msg-1", "att-empty", current_subject = "unsloth")
+ assert excinfo.value.status_code == 404
+
+
+def test_attachment_file_missing_message_is_404(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ with pytest.raises(HTTPException) as excinfo:
+ chat_history.get_attachment_file("nope", "att-1", current_subject = "unsloth")
+ assert excinfo.value.status_code == 404
+
+
+def test_attachment_file_non_data_url_image_is_404(tmp_path, monkeypatch):
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "https://example.com/a.png"}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ with pytest.raises(HTTPException) as excinfo:
+ chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert excinfo.value.status_code == 404
+
+
+def test_attachment_file_defaults_media_type(tmp_path, monkeypatch):
+ payload = base64.b64encode(b"raw-bytes").decode("ascii")
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "data:;base64," + payload}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert response.body == b"raw-bytes"
+ assert response.media_type == "application/octet-stream"
+
+
+def test_attachment_file_svg_media_type(tmp_path, monkeypatch):
+ svg = b" "
+ payload = base64.b64encode(svg).decode("ascii")
+ attachment = _image_attachment()
+ attachment["content"] = [{"type": "image", "image": "data:image/svg+xml;base64," + payload}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
+ assert response.body == svg
+ # SVG can carry scripts, so it downloads as bytes instead of rendering.
+ assert response.media_type == "application/octet-stream"
+
+
+def test_delete_attachment_route_then_404(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_image_attachment()])
+ result = chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth")
+ assert result == {"ok": True}
+ with pytest.raises(HTTPException) as excinfo:
+ chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth")
+ assert excinfo.value.status_code == 404
+
+
+# ---------------------------------------------------------------------------
+# Audio attachments (adapter {data, format} and compare-chat bare base64)
+# ---------------------------------------------------------------------------
+
+WAV_BYTES = b"RIFF$\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00"
+WAV_B64 = base64.b64encode(WAV_BYTES).decode("ascii")
+
+
+def _audio_attachment(attachment_id: str = "att-audio") -> dict:
+ return {
+ "id": attachment_id,
+ "type": "file",
+ "name": "clip.wav",
+ "contentType": "audio/wav",
+ "content": [{"type": "audio", "audio": {"data": WAV_B64, "format": "wav"}}],
+ "status": {"type": "complete"},
+ }
+
+
+def test_audio_attachment_lists_with_size(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_audio_attachment()])
+ records = studio_db.list_chat_attachments()
+ assert len(records) == 1
+ assert records[0]["id"] == "att-audio"
+ assert abs(records[0]["sizeBytes"] - len(WAV_BYTES)) <= 2
+
+
+def test_audio_attachment_file_serves_bytes(tmp_path, monkeypatch):
+ _seed(tmp_path, monkeypatch, [_audio_attachment()])
+ response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth")
+ assert response.body == WAV_BYTES
+ assert response.media_type == "audio/wav"
+
+
+def test_audio_attachment_media_type_from_format(tmp_path, monkeypatch):
+ attachment = _audio_attachment()
+ attachment["contentType"] = None
+ attachment["content"] = [{"type": "audio", "audio": {"data": WAV_B64, "format": "mp3"}}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth")
+ assert response.media_type == "audio/mpeg"
+
+
+def test_audio_attachment_corrupt_payload_is_422(tmp_path, monkeypatch):
+ attachment = _audio_attachment()
+ attachment["content"] = [{"type": "audio", "audio": {"data": "%%%", "format": "wav"}}]
+ _seed(tmp_path, monkeypatch, [attachment])
+ with pytest.raises(HTTPException) as excinfo:
+ chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth")
+ assert excinfo.value.status_code == 422
+
+
+# ---------------------------------------------------------------------------
+# Compare-chat uploads stored as message content parts
+# ---------------------------------------------------------------------------
+
+
+def _compare_message(message_id: str = "msg-cmp") -> dict:
+ return {
+ "id": message_id,
+ "threadId": "thread-1",
+ "parentId": None,
+ "role": "user",
+ "content": [
+ {"type": "image", "image": PNG_DATA_URL},
+ {"type": "audio", "audio": WAV_B64},
+ {"type": "text", "text": "compare these"},
+ ],
+ "createdAt": 1_700_000_000_000,
+ }
+
+
+def _seed_compare(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ studio_db.upsert_chat_message(_compare_message())
+
+
+_CONTENT_PART_PREFIX = "content-part-sha256-"
+
+
+def _content_part_id_for(message_id: str, kind: str) -> str:
+ """Resolve the stable content-hash id for a message's stored blob.
+
+ Content-part ids are SHA-256 hashes of the blob payload, not array
+ indices, so tests look them up from the listing instead of hardcoding an
+ index that would shift when an earlier part is deleted.
+ """
+ for record in studio_db.list_chat_attachments():
+ if record["messageId"] == message_id and record["type"] == kind:
+ return record["id"]
+ raise AssertionError(f"no {kind} content-part upload for {message_id}")
+
+
+def test_content_part_uploads_are_listed(tmp_path, monkeypatch):
+ _seed_compare(tmp_path, monkeypatch)
+ records = studio_db.list_chat_attachments()
+ # Ids are stable content hashes, not array indices.
+ assert all(r["id"].startswith(_CONTENT_PART_PREFIX) for r in records)
+ assert {r["type"] for r in records} == {"image", "audio"}
+ image = next(r for r in records if r["type"] == "image")
+ assert image["contentType"] == "image/png"
+ assert abs(image["sizeBytes"] - len(PNG_BYTES)) <= 2
+ audio = next(r for r in records if r["type"] == "audio")
+ assert audio["type"] == "audio"
+
+
+def test_content_part_file_serves_image_bytes(tmp_path, monkeypatch):
+ _seed_compare(tmp_path, monkeypatch)
+ image_id = _content_part_id_for("msg-cmp", "image")
+ response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth")
+ assert response.body == PNG_BYTES
+ assert response.media_type == "image/png"
+
+
+def test_content_part_delete_keeps_text(tmp_path, monkeypatch):
+ _seed_compare(tmp_path, monkeypatch)
+ image_id = _content_part_id_for("msg-cmp", "image")
+ assert studio_db.delete_chat_attachment("msg-cmp", image_id) is True
+ message = studio_db.get_chat_message("thread-1", "msg-cmp")
+ types = [p["type"] for p in message["content"]]
+ assert types == ["audio", "text"]
+ # The surviving audio blob keeps its own stable hash id after the delete.
+ remaining = studio_db.list_chat_attachments()
+ assert [r["type"] for r in remaining] == ["audio"]
+ assert remaining[0]["id"].startswith(_CONTENT_PART_PREFIX)
+ assert remaining[0]["id"] != image_id
+
+
+def test_content_part_delete_rejects_non_blob(tmp_path, monkeypatch):
+ _seed_compare(tmp_path, monkeypatch)
+ # The text part is not a stored upload, so it never gets an id: only the
+ # image and audio blobs are addressable.
+ assert len(studio_db.list_chat_attachments()) == 2
+ # A well-formed but unknown content-hash id, and malformed ids, all no-op.
+ assert studio_db.delete_chat_attachment("msg-cmp", _CONTENT_PART_PREFIX + "0" * 64) is False
+ assert studio_db.delete_chat_attachment("msg-cmp", "content-part-99") is False
+ assert studio_db.delete_chat_attachment("msg-cmp", "content-part-x") is False
+
+
+def test_text_only_messages_not_listed_as_uploads(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ # The word "image" inside text must not create phantom upload rows.
+ message = _message("msg-txt")
+ message["content"] = [{"type": "text", "text": 'discussing an "image" and "audio" here'}]
+ studio_db.upsert_chat_message(message)
+ assert studio_db.list_chat_attachments() == []
+
+
+def test_remote_image_urls_are_not_listed_as_uploads(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ message = _message("msg-remote")
+ message["content"] = [
+ {"type": "image", "image": "https://example.com/cat.png"},
+ {"type": "text", "text": "look at this"},
+ ]
+ studio_db.upsert_chat_message(message)
+ # No stored bytes: nothing to list, open, or delete.
+ assert studio_db.list_chat_attachments() == []
+ assert studio_db.get_chat_attachment("msg-remote", "content-part-0") is None
+ assert studio_db.delete_chat_attachment("msg-remote", "content-part-0") is False
+ stored = studio_db.get_chat_message("thread-1", "msg-remote")
+ assert [p["type"] for p in stored["content"]] == ["image", "text"]
+
+
+def test_html_data_url_serves_as_octet_stream(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ html_b64 = base64.b64encode(b"").decode()
+ message = _message("msg-html")
+ message["content"] = [
+ {"type": "image", "image": f"data:text/html;base64,{html_b64}"},
+ ]
+ studio_db.upsert_chat_message(message)
+ attachment_id = _content_part_id_for("msg-html", "image")
+ response = chat_history.get_attachment_file(
+ "msg-html", attachment_id, current_subject = "unsloth"
+ )
+ # Never echo a script-capable media type back under the app origin.
+ assert response.media_type == "application/octet-stream"
+ assert response.body == b""
+
+
+def test_svg_data_url_serves_as_octet_stream(tmp_path, monkeypatch):
+ _reset_studio_db(tmp_path, monkeypatch)
+ studio_db.upsert_chat_thread(_thread())
+ svg_b64 = base64.b64encode(b" ").decode()
+ message = _message("msg-svg")
+ message["content"] = [
+ {"type": "image", "image": f"data:image/svg+xml;base64,{svg_b64}"},
+ ]
+ studio_db.upsert_chat_message(message)
+ attachment_id = _content_part_id_for("msg-svg", "image")
+ response = chat_history.get_attachment_file("msg-svg", attachment_id, current_subject = "unsloth")
+ assert response.media_type == "application/octet-stream"
+
+
+def test_png_data_url_keeps_its_media_type(tmp_path, monkeypatch):
+ _seed_compare(tmp_path, monkeypatch)
+ image_id = _content_part_id_for("msg-cmp", "image")
+ response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth")
+ assert response.media_type == "image/png"
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index b6e218793a..b8601b00f6 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -969,10 +969,10 @@ export function AppSidebar() {
))}
- {/* Bulk export and import live in Settings -> Chat -> Data. */}
+ {/* Bulk export and import live in Settings -> Data. */}
- useSettingsDialogStore.getState().openDialog("chat")
+ useSettingsDialogStore.getState().openDialog("data")
}
>
Export all chats…
diff --git a/studio/frontend/src/components/assistant-ui/attachment.tsx b/studio/frontend/src/components/assistant-ui/attachment.tsx
index 98ebe5ab5f..b26840cc02 100644
--- a/studio/frontend/src/components/assistant-ui/attachment.tsx
+++ b/studio/frontend/src/components/assistant-ui/attachment.tsx
@@ -7,6 +7,7 @@
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import {
Dialog,
+ DialogClose,
DialogContent,
DialogTitle,
DialogTrigger,
@@ -27,12 +28,7 @@ import {
import { AudioWave01Icon, File02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { PlusIcon, XIcon } from "lucide-react";
-import {
- type FC,
- type PropsWithChildren,
- useEffect,
- useState,
-} from "react";
+import { type FC, type PropsWithChildren, useEffect, useState } from "react";
import { useShallow } from "zustand/shallow";
const useFileSrc = (file: File | undefined): string | undefined => {
@@ -83,7 +79,7 @@ const AttachmentPreview: FC = ({ src }) => {
src={src}
alt="Preview"
className={cn(
- "block h-auto max-h-[80vh] w-auto max-w-full object-contain",
+ "block h-auto max-h-[90dvh] w-auto max-w-[92vw] object-contain",
isLoaded
? "aui-attachment-preview-image-loaded"
: "aui-attachment-preview-image-loading invisible",
@@ -108,12 +104,23 @@ const AttachmentPreviewDialog: FC = ({ children }) => {
>
{children}
-
+ {/* Chrome-free lightbox: the image floats on the dimmed backdrop with
+ no dialog panel, and the close button sits in the screen corner. */}
+
Image Attachment Preview
-
-
+ {/* Clicking the backdrop (anywhere off the image) closes the preview. */}
+
+
+
+
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 ff4bd4bebd..23bf68c042 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -67,6 +67,8 @@ import {
Download01Icon,
Flag01Icon,
Folder02Icon,
+ PinIcon,
+ PinOffIcon,
RemoveCircleIcon,
Search01Icon,
ViewIcon,
@@ -99,6 +101,11 @@ import {
loadedAt,
useModelLoadTimes,
} from "./model-usage";
+import {
+ pinKey,
+ pinnedQuantEntries,
+ usePinnedModelsStore,
+} from "./pinned-models";
import {
type FormatFilter,
estimateQuantBytes,
@@ -384,10 +391,64 @@ function CapabilityIcons({ caps }: { caps: ModelCapabilities }) {
);
}
+function normalizeModelIdForPicker(modelId: string): string {
+ const trimmed = modelId.trim();
+ const slashPath = trimmed.replace(/\\/g, "/").replace(/\/+$/, "");
+ const caseInsensitive =
+ !/^(\/|\.{1,2}\/|~\/)/.test(slashPath) ||
+ /^[A-Za-z]:\//.test(slashPath) ||
+ slashPath.startsWith("//") ||
+ /^\/mnt\/[A-Za-z](?:\/|$)/.test(slashPath);
+ return caseInsensitive ? slashPath.toLowerCase() : slashPath;
+}
+
+function modelIdsMatchForPicker(
+ left: string | null | undefined,
+ right: string | null | undefined,
+): boolean {
+ return Boolean(
+ left &&
+ right &&
+ normalizeModelIdForPicker(left) === normalizeModelIdForPicker(right),
+ );
+}
+
+function normalizeGgufVariantForPicker(variant: string | null | undefined) {
+ return variant?.trim().toLowerCase() ?? "";
+}
+
+function ggufVariantsMatchForPicker(
+ left: string | null | undefined,
+ right: string | null | undefined,
+): boolean {
+ return (
+ normalizeGgufVariantForPicker(left) ===
+ normalizeGgufVariantForPicker(right)
+ );
+}
+
+function isRuntimeLoadedModel(
+ loadedModelId: string | undefined,
+ activeGgufVariant: string | null | undefined,
+ modelId: string,
+ variantPolicy: "none" | "required" | "ignore",
+): boolean {
+ if (!modelIdsMatchForPicker(loadedModelId, modelId)) return false;
+ if (variantPolicy === "ignore") return true;
+ const hasActiveGgufVariant = !ggufVariantsMatchForPicker(
+ activeGgufVariant,
+ null,
+ );
+ return variantPolicy === "required"
+ ? hasActiveGgufVariant
+ : !hasActiveGgufVariant;
+}
+
function ModelRow({
label,
meta,
selected,
+ loaded = false,
onClick,
vramStatus,
vramEst,
@@ -405,6 +466,8 @@ function ModelRow({
label: string;
meta?: string | null;
selected?: boolean;
+ /** Override badge state when authoritative runtime state is available. */
+ loaded?: boolean;
onClick: () => void;
vramStatus?: VramFitStatus | null;
vramEst?: number;
@@ -494,7 +557,7 @@ function ModelRow({
)}
- {selected && (
+ {loaded && (
)}
- {downloaded && !selected && (
+ {downloaded && !loaded && (
void;
}) {
+ const pinnedKeys = usePinnedModelsStore((s) => s.pinned);
+ const togglePinnedQuant = usePinnedModelsStore((s) => s.togglePinned);
const onUpdateVariant = variantActions?.onUpdate;
const updateVariantTitle = variantActions?.updateTitle ?? "Update cached model?";
const renderUpdateVariantDescription = variantActions?.renderUpdateDescription;
@@ -946,7 +1015,7 @@ function GgufVariantExpander({
{v.downloaded ? (
<>
-
+
downloaded
{v.update_available ? (
@@ -1000,6 +1069,43 @@ function GgufVariantExpander({
onUpdated={() => setRefreshKey((key) => key + 1)}
/>
)}
+ {v.downloaded && allowPin && (
+
+
+ togglePinnedQuant(repoId, v.quant)}
+ aria-label={
+ pinnedKeys.includes(pinKey(repoId, v.quant))
+ ? `Unpin ${repoId} ${v.quant}`
+ : `Pin ${repoId} ${v.quant}`
+ }
+ aria-pressed={pinnedKeys.includes(pinKey(repoId, v.quant))}
+ className={cn(
+ "shrink-0 rounded-md p-1 transition-colors hover:bg-black/5 hover:text-foreground dark:hover:bg-white/10",
+ pinnedKeys.includes(pinKey(repoId, v.quant))
+ ? "text-foreground/80"
+ : "text-muted-foreground/60",
+ )}
+ >
+
+
+
+
+ {pinnedKeys.includes(pinKey(repoId, v.quant))
+ ? "Unpin quant"
+ : "Pin quant to the top"}
+
+
+ )}
{v.downloaded && (
onDeleteVariant(v.quant)}
+ onConfirm={async () => {
+ await onDeleteVariant(v.quant);
+ // Drop the pin too: a pinned row for a deleted file
+ // would try to load something that no longer exists.
+ if (pinnedKeys.includes(pinKey(repoId, v.quant))) {
+ togglePinnedQuant(repoId, v.quant);
+ }
+ }}
/>
)}
@@ -1273,6 +1386,8 @@ export function HubModelPicker({
// Live model id from the runtime store (backend-mirrored active_model), not the dropdown
// highlight which can be a staged pick. Disables the update action for it.
const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint);
+ // Loaded GGUF quant of the active model; marks the matching pinned row.
+ const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
// Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date).
const loadTimes = useModelLoadTimes(value);
// Fade the list's top edge once scrolled, and its bottom edge while more
@@ -1396,6 +1511,7 @@ export function HubModelPicker({
[expandQuantizations],
);
+ const [pinnedCollapsed, setPinnedCollapsed] = useState(false);
const [downloadedCollapsed, setDownloadedCollapsed] = useState(false);
const [otherModelsCollapsed, setOtherModelsCollapsed] = useState(false);
const [customFoldersCollapsed, setCustomFoldersCollapsed] = useState(false);
@@ -1964,6 +2080,110 @@ export function HubModelPicker({
// logic must use this (not visibleCachedModels) or the picker can go blank.
const visibleCachedModelRows = chatOnly ? [] : visibleCachedModels;
+ // Pinned entries surface in their own section above the Unsloth heading.
+ // GGUF quants pin individually and their repo stays listed below; non-GGUF
+ // repos pin whole and leave the Unsloth / Other models groups.
+ const pinnedIds = usePinnedModelsStore((s) => s.pinned);
+ const togglePinned = usePinnedModelsStore((s) => s.togglePinned);
+ const pinnedSet = useMemo(() => new Set(pinnedIds), [pinnedIds]);
+
+ // Candidate pins whose repo still exists in the managed cache. Per-quant
+ // validation below is required because deleting one variant can leave a
+ // sibling quant (and therefore the repo row) cached.
+ const pinnedQuantCandidates = useMemo(() => {
+ // The existence check ignores the text query (but keeps the format filter)
+ // so a pinned quant stays findable by its quant name even when the repo id
+ // does not match the query; querying visibleCachedGguf here would drop the
+ // repo before the later `${repoId} ${quant}` predicate could surface it.
+ const cached = new Set(
+ sortedCachedGguf
+ .filter((c) => matchesFormatFilter(c.repo_id, true, formatFilter))
+ .map((c) => c.repo_id),
+ );
+ return pinnedQuantEntries(pinnedIds).filter((entry) =>
+ cached.has(entry.repoId),
+ );
+ }, [pinnedIds, sortedCachedGguf, formatFilter]);
+ const pinnedQuantValidationKey = useMemo(() => {
+ const cacheByRepo = new Map(
+ sortedCachedGguf.map((repo) => [repo.repo_id, repo]),
+ );
+ return pinnedQuantCandidates
+ .map((entry) => {
+ const cached = cacheByRepo.get(entry.repoId);
+ return `${pinKey(entry.repoId, entry.quant)}@${cached?.size_bytes ?? 0}:${cached?.last_modified ?? 0}`;
+ })
+ .join("\u0000");
+ }, [pinnedQuantCandidates, sortedCachedGguf]);
+ const [pinnedQuantValidation, setPinnedQuantValidation] = useState<{
+ key: string;
+ downloaded: ReadonlySet;
+ }>({ key: "", downloaded: new Set() });
+
+ useEffect(() => {
+ let cancelled = false;
+ const repoIds = Array.from(
+ new Set(pinnedQuantCandidates.map((entry) => entry.repoId)),
+ );
+ if (repoIds.length === 0) return;
+
+ void Promise.all(
+ repoIds.map(async (repoId) => {
+ try {
+ const response = await listGgufVariants(
+ repoId,
+ hfToken || undefined,
+ );
+ return normalizeGgufVariantsResponse(response).variants
+ .filter((variant) => variant.downloaded === true)
+ .map((variant) => pinKey(repoId, variant.quant));
+ } catch {
+ // If the backend cannot verify a quant, hiding the direct-load row
+ // is safer than claiming a missing file is downloaded.
+ return [];
+ }
+ }),
+ ).then((groups) => {
+ if (!cancelled) {
+ setPinnedQuantValidation({
+ key: pinnedQuantValidationKey,
+ downloaded: new Set(groups.flat()),
+ });
+ }
+ });
+
+ return () => {
+ cancelled = true;
+ };
+ }, [hfToken, pinnedQuantCandidates, pinnedQuantValidationKey]);
+ const downloadedPinnedQuantKeys = useMemo>(
+ () =>
+ pinnedQuantValidation.key === pinnedQuantValidationKey
+ ? pinnedQuantValidation.downloaded
+ : new Set(),
+ [pinnedQuantValidation, pinnedQuantValidationKey],
+ );
+
+ // Verified downloaded quants, in pin order and filtered by repo id or quant.
+ const pinnedQuants = useMemo(() => {
+ const q = normalizeForSearch(debouncedQuery.trim());
+ return pinnedQuantCandidates.filter(
+ (entry) =>
+ downloadedPinnedQuantKeys.has(pinKey(entry.repoId, entry.quant)) &&
+ (!q ||
+ normalizeForSearch(`${entry.repoId} ${entry.quant}`).includes(q)),
+ );
+ }, [
+ debouncedQuery,
+ downloadedPinnedQuantKeys,
+ pinnedQuantCandidates,
+ ]);
+
+ const pinnedCachedModelRows = useMemo(
+ () => visibleCachedModelRows.filter((c) => pinnedSet.has(pinKey(c.repo_id))),
+ [visibleCachedModelRows, pinnedSet],
+ );
+
// Split downloaded models so non-Unsloth repos get their own "Other models"
// section above Fine-tuned.
const unslothCachedGguf = useMemo(
@@ -1975,12 +2195,18 @@ export function HubModelPicker({
[visibleCachedGguf],
);
const unslothCachedModelRows = useMemo(
- () => visibleCachedModelRows.filter((c) => isUnslothRepoId(c.repo_id)),
- [visibleCachedModelRows],
+ () =>
+ visibleCachedModelRows.filter(
+ (c) => isUnslothRepoId(c.repo_id) && !pinnedSet.has(pinKey(c.repo_id)),
+ ),
+ [visibleCachedModelRows, pinnedSet],
);
const otherCachedModelRows = useMemo(
- () => visibleCachedModelRows.filter((c) => !isUnslothRepoId(c.repo_id)),
- [visibleCachedModelRows],
+ () =>
+ visibleCachedModelRows.filter(
+ (c) => !isUnslothRepoId(c.repo_id) && !pinnedSet.has(pinKey(c.repo_id)),
+ ),
+ [visibleCachedModelRows, pinnedSet],
);
// Param counts come straight off the unsloth listings the picker already
@@ -2076,6 +2302,25 @@ export function HubModelPicker({
const hubOptionKeys = useMemo(() => {
const keys: string[] = [];
+ // Pinned rows sit above the Unsloth heading on the On Device tab.
+ if (
+ section === "downloaded" &&
+ cachedReady &&
+ !pinnedCollapsed &&
+ (pinnedQuants.length > 0 || pinnedCachedModelRows.length > 0)
+ ) {
+ keys.push(
+ ...pinnedQuants.map((entry) =>
+ makeModelOptionKey("pinned-quant", pinKey(entry.repoId, entry.quant)),
+ ),
+ );
+ keys.push(
+ ...pinnedCachedModelRows.map((model) =>
+ makeModelOptionKey("downloaded-model", model.repo_id),
+ ),
+ );
+ }
+
// Downloaded (Unsloth) rows (query-filtered) on the On Device tab only.
if (
section === "downloaded" &&
@@ -2167,6 +2412,9 @@ export function HubModelPicker({
chatOnly,
sortedCustomFolderModels,
customFoldersCollapsed,
+ pinnedQuants,
+ pinnedCachedModelRows,
+ pinnedCollapsed,
downloadedCollapsed,
fineTunedRows,
fineTunedCollapsed,
@@ -2475,6 +2723,151 @@ export function HubModelPicker({
selected && "bg-[#ececec] dark:bg-[var(--sidebar-accent)]",
);
+ // Pin toggle at a row's right edge: hidden until the row is hovered (or the
+ // button is focused), always visible while pinned so pinned rows read as such.
+ // `small` matches the compact quant-row action sizing; it also skips the
+ // hide-until-hover classes since small pins render inside a hover-gated group.
+ const renderPinAction = (
+ repoId: string,
+ quant?: string,
+ opts?: { className?: string; small?: boolean },
+ ) => {
+ const pinned = pinnedSet.has(pinKey(repoId, quant));
+ const target = quant ? `${repoId} ${quant}` : repoId;
+ return (
+
+
+ {
+ e.stopPropagation();
+ togglePinned(repoId, quant);
+ }}
+ aria-label={pinned ? `Unpin ${target}` : `Pin ${target}`}
+ aria-pressed={pinned}
+ className={cn(
+ "shrink-0 rounded-md transition-colors hover:bg-black/5 dark:hover:bg-white/10",
+ opts?.small ? "p-1" : "p-1.5",
+ pinned
+ ? "text-foreground/80 hover:text-foreground"
+ : "text-muted-foreground/60 hover:text-foreground",
+ !pinned &&
+ !opts?.small &&
+ "opacity-0 focus-visible:opacity-100 group-hover:opacity-100 group-focus-within:opacity-100",
+ opts?.className,
+ )}
+ >
+
+
+
+
+ {pinned
+ ? quant
+ ? "Unpin quant"
+ : "Unpin model"
+ : quant
+ ? "Pin quant to the top"
+ : "Pin model to the top"}
+
+
+ );
+ };
+
+ // A pinned quant: repo name with the quant as a grey chip. One click loads
+ // that quant directly, no expansion needed.
+ const renderPinnedQuantRow = (entry: { repoId: string; quant: string }) => {
+ const optionKey = makeModelOptionKey(
+ "pinned-quant",
+ pinKey(entry.repoId, entry.quant),
+ );
+ const { owner, name } = splitRepoLabel(entry.repoId);
+ const isSelected = value === entry.repoId && activeGgufVariant === entry.quant;
+ const isLoaded =
+ modelIdsMatchForPicker(loadedModelId, entry.repoId) &&
+ !ggufVariantsMatchForPicker(activeGgufVariant, null) &&
+ ggufVariantsMatchForPicker(activeGgufVariant, entry.quant);
+ return (
+
+
+ onSelect(entry.repoId, {
+ source: "hub",
+ isLora: false,
+ ggufVariant: entry.quant,
+ isDownloaded: true,
+ })
+ }
+ className={cn(
+ "flex min-w-0 flex-1 items-center gap-2 rounded-full px-2 py-1.5 text-left text-sm transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring/45",
+ downloadedRowButtonClassName,
+ )}
+ title={`${entry.repoId} (${entry.quant})`}
+ >
+
+ {owner ? (
+
+ {owner}
+ /
+
+ ) : null}
+ {name}
+
+
+ {entry.quant}
+
+ {isLoaded && (
+
+ )}
+
+
+ {renderPinAction(entry.repoId, entry.quant, { small: true })}
+
+
+ This will remove{" "}
+
+ {entry.repoId} ({entry.quant})
+ {" "}
+ from disk. You can re-download it later.
+ >
+ }
+ successMessage={`Deleted ${entry.repoId} ${entry.quant}`}
+ buttonClassName="p-1"
+ iconClassName="size-3"
+ disabled={deleteDisabled}
+ onConfirm={async () => {
+ await deleteCachedModel(entry.repoId, entry.quant);
+ refreshCachedLists();
+ // The file is gone, so drop its pin too.
+ togglePinned(entry.repoId, entry.quant);
+ }}
+ />
+
+
+ );
+ };
+
// Shared row renderers so Downloaded (Unsloth) and Other models render alike.
const renderDownloadedGgufRow = (c: (typeof visibleCachedGguf)[number]) => {
const optionKey = makeModelOptionKey("downloaded-gguf", c.repo_id);
@@ -2489,6 +2882,12 @@ export function HubModelPicker({
meta="GGUF"
showVision={c.has_vision ?? visionByRepo[c.repo_id]}
selected={isSelected}
+ loaded={isRuntimeLoadedModel(
+ loadedModelId,
+ activeGgufVariant,
+ c.repo_id,
+ "required",
+ )}
optionProps={hubModelList.getOptionProps(optionKey, isSelected)}
onClick={() => toggleGgufExpanded(c.repo_id)}
onArrowDownIntoChildren={
@@ -2506,6 +2905,7 @@ export function HubModelPicker({
reportVision(c.repo_id, v)}
onSelect={onSelect}
hfToken={hfToken || undefined}
@@ -2523,6 +2923,7 @@ export function HubModelPicker({
await deleteCachedModel(c.repo_id, quant);
refreshCachedLists();
},
+ deleteDisabled,
}}
/>
)}
@@ -2547,6 +2948,12 @@ export function HubModelPicker({
c.size_bytes,
)}`}
selected={isSelected}
+ loaded={isRuntimeLoadedModel(
+ loadedModelId,
+ activeGgufVariant,
+ c.repo_id,
+ "none",
+ )}
optionProps={hubModelList.getOptionProps(
optionKey,
isSelected,
@@ -2562,6 +2969,7 @@ export function HubModelPicker({
className={downloadedRowButtonClassName}
/>
+ {renderPinAction(c.repo_id)}
deleteCachedModel(c.repo_id)}
+ disabled={deleteDisabled}
+ onConfirm={async () => {
+ await deleteCachedModel(c.repo_id);
+ if (pinnedSet.has(pinKey(c.repo_id))) {
+ togglePinned(c.repo_id);
+ }
+ }}
onDeleted={refreshCachedLists}
/>
@@ -2749,12 +3163,36 @@ export function HubModelPicker({
) : null}
+ {/* Pinned quants and models sit above the Unsloth heading so
+ favorites are always first. Filtered by the query like the
+ sections below. */}
+ {showDownloaded &&
+ (pinnedQuants.length > 0 ||
+ pinnedCachedModelRows.length > 0) ? (
+ <>
+
}
+ collapsed={pinnedCollapsed}
+ onToggle={() => setPinnedCollapsed((v) => !v)}
+ >
+ Pinned
+
+ {!pinnedCollapsed && pinnedQuants.map(renderPinnedQuantRow)}
+ {!pinnedCollapsed &&
+ pinnedCachedModelRows.map(renderDownloadedModelRow)}
+ >
+ ) : null}
+
{/* Downloaded (Unsloth) stays visible (filtered) while searching. */}
{showDownloaded &&
(unslothCachedGguf.length > 0 ||
unslothCachedModelRows.length > 0) ? (
<>
0 ||
+ pinnedCachedModelRows.length > 0
+ }
collapsed={downloadedCollapsed}
onToggle={() => setDownloadedCollapsed((v) => !v)}
action={
@@ -2896,6 +3334,8 @@ export function HubModelPicker({
)}
@@ -3497,6 +3974,12 @@ export function HubModelPicker({
: (vram?.detail ?? extractParamLabel(id))
}
selected={value === id}
+ loaded={isRuntimeLoadedModel(
+ loadedModelId,
+ activeGgufVariant,
+ id,
+ isKnownGgufRepo(id) ? "required" : "none",
+ )}
optionProps={hubModelList.getOptionProps(
optionKey,
value === id,
@@ -3546,6 +4029,7 @@ export function HubModelPicker({
await deleteCachedModel(id, quant);
refreshCachedLists();
},
+ deleteDisabled,
}}
/>
)}
@@ -3586,6 +4070,12 @@ export function HubModelPicker({
.join(" · ")
}
selected={value === id}
+ loaded={isRuntimeLoadedModel(
+ loadedModelId,
+ activeGgufVariant,
+ id,
+ isSearchGguf ? "required" : "none",
+ )}
optionProps={hubModelList.getOptionProps(
optionKey,
value === id,
@@ -3637,6 +4127,7 @@ export function HubModelPicker({
await deleteCachedModel(id, quant);
refreshCachedLists();
},
+ deleteDisabled,
}}
/>
)}
@@ -3687,6 +4178,8 @@ export function HubModelPicker({
function FineTunedRows({
adapters,
value,
+ loadedModelId,
+ activeGgufVariant,
onSelect,
onModelsChange,
deleteDisabled = false,
@@ -3697,6 +4190,8 @@ function FineTunedRows({
}: {
adapters: LoraModelOption[];
value?: string;
+ loadedModelId?: string;
+ activeGgufVariant?: string | null;
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
deleteDisabled?: boolean;
@@ -3753,6 +4248,12 @@ function FineTunedRows({
label={adapter.name}
meta={meta}
selected={value === adapter.id}
+ loaded={isRuntimeLoadedModel(
+ loadedModelId,
+ activeGgufVariant,
+ adapter.id,
+ isLocalGgufDir || isExportedGguf ? "required" : "none",
+ )}
optionProps={loraModelList.getOptionProps(
optionKey,
value === adapter.id,
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts b/studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts
new file mode 100644
index 0000000000..4835c4c0cf
--- /dev/null
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts
@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// Pinned models for the model selector's On Device list, persisted in
+// localStorage so pins survive reloads. GGUF quants pin individually
+// (repoId + quant); non-GGUF repos pin as a whole. Pinned entries surface
+// in a "Pinned" section above the Unsloth/Downloaded group.
+
+import { create } from "zustand";
+
+const KEY = "unsloth_pinned_models";
+
+// Entries are stored as strings: "repoId" pins a whole (non-GGUF) repo,
+// "repoId::quant" pins one GGUF quant. Neither part contains "::".
+export function pinKey(repoId: string, quant?: string): string {
+ return quant ? `${repoId}::${quant}` : repoId;
+}
+
+export interface PinnedQuantEntry {
+ repoId: string;
+ quant: string;
+}
+
+/** The pinned GGUF quants, in pin order. Plain repo pins are excluded. */
+export function pinnedQuantEntries(pinned: string[]): PinnedQuantEntry[] {
+ const out: PinnedQuantEntry[] = [];
+ for (const key of pinned) {
+ const sep = key.indexOf("::");
+ if (sep <= 0) continue;
+ const repoId = key.slice(0, sep);
+ const quant = key.slice(sep + 2);
+ if (repoId && quant) out.push({ repoId, quant });
+ }
+ return out;
+}
+
+function readPinned(): string[] {
+ try {
+ const raw = JSON.parse(localStorage.getItem(KEY) ?? "[]");
+ return Array.isArray(raw)
+ ? raw.filter((v): v is string => typeof v === "string")
+ : [];
+ } catch {
+ return [];
+ }
+}
+
+function writePinned(pinned: string[]): void {
+ try {
+ localStorage.setItem(KEY, JSON.stringify(pinned));
+ } catch {
+ // Ignore unavailable storage; pins stay session-only.
+ }
+}
+
+interface PinnedModelsState {
+ pinned: string[];
+ togglePinned: (repoId: string, quant?: string) => void;
+}
+
+export const usePinnedModelsStore = create((set) => ({
+ pinned: readPinned(),
+ togglePinned: (repoId, quant) =>
+ set((state) => {
+ const key = pinKey(repoId, quant);
+ const next = state.pinned.includes(key)
+ ? state.pinned.filter((id) => id !== key)
+ : [...state.pinned, key];
+ writePinned(next);
+ return { pinned: next };
+ }),
+}));
diff --git a/studio/frontend/src/components/ui/tooltip.tsx b/studio/frontend/src/components/ui/tooltip.tsx
index cbc1a09a8f..91745af5ec 100644
--- a/studio/frontend/src/components/ui/tooltip.tsx
+++ b/studio/frontend/src/components/ui/tooltip.tsx
@@ -67,9 +67,14 @@ function TooltipTrigger({
const handleClick = useCallback(
(e: React.MouseEvent) => {
+ // Run the composed handler first: when this trigger wraps another Radix
+ // trigger (e.g. DialogTrigger around an attachment tile), that trigger's
+ // action is skipped if the event is already default-prevented.
+ onClick?.(e);
+ // preventDefault keeps Radix Tooltip's internal close-on-click from
+ // undoing the tap-toggle below (its composed handler checks it).
e.preventDefault();
toggle?.();
- onClick?.(e);
},
[toggle, onClick],
);
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index 0f6af38033..3ba1ad6fe4 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -2,7 +2,11 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
+// These helpers are deliberately API-layer-only and are not part of their
+// features' React-facing public barrels.
+// eslint-disable-next-line no-restricted-imports
import { hubTokenHeader } from "@/features/hub/lib/hub-token-header";
+// eslint-disable-next-line no-restricted-imports
import { consumeNativePathToken } from "@/features/native-intents/api";
import { formatFastApiDetail } from "@/lib/format-fastapi-error";
import type {
@@ -437,6 +441,73 @@ export async function listChatThreads(
return Array.isArray(data.threads) ? data.threads : [];
}
+/** One chat message attachment, as listed for the settings uploaded-files view. */
+export interface ChatAttachmentRecord {
+ id: string;
+ messageId: string;
+ threadId: string;
+ pairId?: string | null;
+ threadTitle?: string | null;
+ name: string;
+ type?: string | null;
+ contentType?: string | null;
+ sizeBytes?: number | null;
+ createdAt?: number | null;
+}
+
+export interface ChatAttachmentPage {
+ attachments: ChatAttachmentRecord[];
+ nextOffset: number | null;
+}
+
+export async function listChatAttachments(
+ offset = 0,
+ limit = 50,
+): Promise {
+ const params = new URLSearchParams({
+ limit: String(limit),
+ offset: String(offset),
+ });
+ const response = await authFetch(`/api/chat/attachments?${params}`);
+ const data = await parseJsonOrThrow<{
+ attachments: ChatAttachmentRecord[];
+ nextOffset: number | null;
+ }>(response);
+ return {
+ attachments: Array.isArray(data.attachments) ? data.attachments : [],
+ nextOffset:
+ typeof data.nextOffset === "number" && Number.isFinite(data.nextOffset)
+ ? data.nextOffset
+ : null,
+ };
+}
+
+/** Stored attachment content (image bytes or extracted text) as a Blob. */
+export async function fetchChatAttachmentBlob(
+ messageId: string,
+ attachmentId: string,
+): Promise {
+ const response = await authFetch(
+ `/api/chat/attachments/${encodeURIComponent(messageId)}/${encodeURIComponent(attachmentId)}/file`,
+ );
+ if (!response.ok) {
+ const body = await response.json().catch(() => null);
+ throw new Error(parseErrorText(response.status, body));
+ }
+ return response.blob();
+}
+
+export async function deleteChatAttachment(
+ messageId: string,
+ attachmentId: string,
+): Promise {
+ const response = await authFetch(
+ `/api/chat/attachments/${encodeURIComponent(messageId)}/${encodeURIComponent(attachmentId)}`,
+ { method: "DELETE" },
+ );
+ await parseJsonOrThrow<{ ok: boolean }>(response);
+}
+
export async function getChatThread(
threadId: string,
): Promise {
@@ -960,7 +1031,8 @@ export async function* streamChatCompletions(
parsed.type === "reasoning_summary"
) {
yield {
- _reasoningDurationMs: (parsed as { duration_ms?: number }).duration_ms,
+ _reasoningDurationMs: (parsed as { duration_ms?: number })
+ .duration_ms,
} as unknown as OpenAIChatChunk;
separatorIndex = buffer.search(/\r?\n\r?\n/);
continue;
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
index bfb3eeb14c..a08bd5fa54 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
@@ -217,6 +217,40 @@ export async function archiveChatItem(
notifyChatHistoryUpdated();
}
+export async function archiveAllChatItems(
+ activeId?: string,
+ onSelect?: (view: { mode: "single"; newThreadNonce: string }) => void,
+): Promise {
+ const threads = await listStoredChatThreads({ includeArchived: true });
+ // Boolean() mirrors groupThreads: legacy records may have archived
+ // undefined/null, which must count as "not archived".
+ const toArchive = threads.filter((t) => !t.archived);
+ if (toArchive.length === 0) return 0;
+
+ for (const t of toArchive) cancelIfRunning(t.id);
+
+ await Promise.all(
+ toArchive.map((t) => updateStoredChatThread(t.id, { archived: true })),
+ );
+
+ // Reset only when this action archived the active single thread or compare
+ // pair. An already-archived chat opened from the archive is not in
+ // toArchive and must stay open.
+ const archivedActive =
+ activeId !== undefined &&
+ toArchive.some(
+ (thread) => thread.id === activeId || thread.pairId === activeId,
+ );
+ if (archivedActive) {
+ useChatRuntimeStore.getState().setActiveThreadId(null);
+ onSelect?.({ mode: "single", newThreadNonce: crypto.randomUUID() });
+ }
+
+ notifyChatHistoryUpdated();
+ // Report sidebar items, not raw threads: a compare pair reads as one chat.
+ return groupThreads(toArchive).length;
+}
+
export async function unarchiveChatItem(item: SidebarItem): Promise {
const threadIds: string[] =
item.type === "single"
diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts
index a8b5fc23ad..b0059b57b1 100644
--- a/studio/frontend/src/features/chat/index.ts
+++ b/studio/frontend/src/features/chat/index.ts
@@ -3,10 +3,15 @@
export { ChatPage, validateChatSearch, type ChatSearch } from "./chat-page";
export {
+ deleteChatAttachment,
+ fetchChatAttachmentBlob,
getInferenceStatus,
+ listChatAttachments,
listGgufVariants,
listLocalModels,
loadModel,
+ type ChatAttachmentPage,
+ type ChatAttachmentRecord,
type LocalModelInfo,
} from "./api/chat-api";
export type { GgufVariantDetail } from "./types/api";
@@ -17,6 +22,10 @@ export {
type Preset,
} from "./chat-settings-sheet";
export { useChatRuntimeStore } from "./stores/chat-runtime-store";
+export {
+ CHAT_RAG_CAPTION_KEY,
+ CHAT_RAG_OCR_KEY,
+} from "./stores/chat-runtime-store";
export {
preferFullToolOutput,
toolOutputKey,
@@ -46,12 +55,16 @@ export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
export type { ProjectRecord } from "./types";
export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
export { listStoredChatThreads } from "./utils/chat-history-storage";
+export { emitChatAttachmentDeleted } from "./utils/chat-attachment-events";
export { ArtifactCard } from "./artifacts/artifact-card";
export {
useChatArtifactsStore,
useSelectedChatArtifact,
} from "./artifacts/store";
-export { downloadChatExport } from "./utils/export-chat-history";
+export {
+ downloadChatExport,
+ downloadArchivedChatExport,
+} from "./utils/export-chat-history";
export {
clearNewChatDraft,
composerDraftKey,
@@ -60,10 +73,14 @@ export {
} from "./utils/composer-draft";
export {
EXPORT_FORMATS_LIST,
+ buildFineTuneJsonl,
bulkExportConversationsByScope,
+ exportFineTuneJsonl,
importConversationsFromFile,
+ type FineTuneFormat,
} from "./prompt-storage/prompt-storage-dialog";
export {
+ archiveAllChatItems,
archiveChatItem,
deleteChatItem,
renameChatItem,
diff --git a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
index 33fc6286c0..68f24a7b08 100644
--- a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
+++ b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
@@ -186,7 +186,16 @@ function contentBlocksToText(content: unknown): string {
// predate the user's next message); the parent chain is timestamp-independent.
type _Msg = { id: string; parentId?: string | null; createdAt?: number };
-function orderByParentChain(messages: T[]): T[] {
+function orderByParentChain(
+ messages: T[],
+ options: {
+ /** Append messages off the selected chain (abandoned branches) at the
+ * end. Full exports keep everything; fine-tune conversion must not,
+ * since alternate replies would merge into one conversation. */
+ includeSiblings?: boolean;
+ } = {},
+): T[] {
+ const { includeSiblings = true } = options;
const byId = new Map(messages.map((m) => [m.id, m]));
const childrenOf = new Map();
for (const m of messages) {
@@ -207,7 +216,9 @@ function orderByParentChain(messages: T[]): T[] {
byId.delete(next.id);
}
- for (const [, m] of byId) result.push(m);
+ if (includeSiblings) {
+ for (const [, m] of byId) result.push(m);
+ }
return result;
}
@@ -543,6 +554,214 @@ export async function exportProjectConversations(
);
}
+// ── Fine-tuning export ─────────────────────────────────────────────────────
+// One JSONL line per conversation: {"messages": [{"role", "content"}]} with
+// string-only content in system/user/assistant turns. Unsloth's training tab
+// detects this as ChatML natively (no column mapping, no standardization) and
+// it works with train-on-completions masking, which only trains on assistant
+// turns. Reasoning, tool calls, and images are dropped: clean SFT targets.
+
+export type FineTuneMessage = {
+ role: "system" | "user" | "assistant";
+ content: string;
+};
+
+const FINE_TUNE_ROLES = new Set(["system", "user", "assistant"]);
+
+/** Plain text of a message: text blocks plus text-type attachment parts. */
+function messageToPlainText(msg: {
+ content: unknown;
+ attachments?: unknown;
+}): string {
+ const parts: string[] = [];
+ const collect = (blocks: unknown) => {
+ // Legacy and imported histories can store content as a plain string.
+ if (typeof blocks === "string") {
+ if (blocks.trim()) parts.push(blocks);
+ return;
+ }
+ if (!Array.isArray(blocks)) return;
+ for (const b of blocks) {
+ if (!b || typeof b !== "object") {
+ continue;
+ }
+ const block = b as Record;
+ if (block.type === "text" && typeof block.text === "string" && block.text) {
+ parts.push(block.text);
+ }
+ }
+ };
+ collect(msg.content);
+ if (Array.isArray(msg.attachments)) {
+ for (const attachment of msg.attachments as Array<{ content?: unknown }>) {
+ collect(attachment?.content);
+ }
+ }
+ return parts.join("\n\n").trim();
+}
+
+/** Merge consecutive same-role turns so chat templates format cleanly. */
+function mergeSameRoleTurns(turns: FineTuneMessage[]): FineTuneMessage[] {
+ const merged: FineTuneMessage[] = [];
+ for (const turn of turns) {
+ const last = merged[merged.length - 1];
+ if (last && last.role === turn.role) {
+ last.content += `\n\n${turn.content}`;
+ } else {
+ merged.push({ ...turn });
+ }
+ }
+ return merged;
+}
+
+/** Conversation turns for fine-tuning, or null when the thread has no
+ * usable user + assistant exchange. Consecutive same-role turns merge,
+ * assistant turns before the first user turn drop (an assistant target
+ * with no prompt teaches nothing), and trailing non-assistant turns drop
+ * so chat templates format cleanly. */
+function messagesToFineTuneTurns(
+ messages: Array<{ role: unknown; content: unknown; attachments?: unknown }>,
+): FineTuneMessage[] | null {
+ const raw: FineTuneMessage[] = [];
+ for (const msg of messages) {
+ const role = msg.role as FineTuneMessage["role"];
+ if (!FINE_TUNE_ROLES.has(role)) continue;
+ const content = messageToPlainText(msg);
+ if (!content) continue;
+ raw.push({ role, content });
+ }
+ const firstUser = raw.findIndex((t) => t.role === "user");
+ if (firstUser === -1) return null;
+ const turns = mergeSameRoleTurns(
+ raw.filter((t, i) => i >= firstUser || t.role === "system"),
+ );
+ while (turns.length > 0 && turns[turns.length - 1].role !== "assistant") {
+ turns.pop();
+ }
+ const hasUser = turns.some((t) => t.role === "user");
+ const hasAssistant = turns.some((t) => t.role === "assistant");
+ return hasUser && hasAssistant ? turns : null;
+}
+
+export type FineTuneExportResult = {
+ lines: string[];
+ conversations: number;
+ skipped: number;
+};
+
+/** Dataset shapes the Train tab detects without column mapping. */
+export type FineTuneFormat = "openai" | "sharegpt" | "alpaca";
+
+const SHAREGPT_FROM: Record = {
+ system: "system",
+ user: "human",
+ assistant: "gpt",
+};
+
+/** JSONL lines for one conversation in the chosen format. Alpaca is
+ * single-turn, so each user to assistant pair becomes its own record with
+ * the system prompt and earlier exchange carried in the input field. */
+function turnsToFineTuneLines(
+ turns: FineTuneMessage[],
+ format: FineTuneFormat,
+): string[] {
+ if (format === "sharegpt") {
+ return [
+ JSON.stringify({
+ conversations: turns.map((t) => ({
+ from: SHAREGPT_FROM[t.role],
+ value: t.content,
+ })),
+ }),
+ ];
+ }
+ if (format === "alpaca") {
+ const lines: string[] = [];
+ const context: string[] = [];
+ let system = "";
+ let pendingUser: string | null = null;
+ for (const t of turns) {
+ if (t.role === "system") {
+ system = system ? `${system}\n\n${t.content}` : t.content;
+ continue;
+ }
+ if (t.role === "user") {
+ pendingUser = t.content;
+ continue;
+ }
+ if (pendingUser === null) continue;
+ const inputParts = [];
+ if (system) inputParts.push(system);
+ if (context.length > 0) inputParts.push(context.join("\n"));
+ lines.push(
+ JSON.stringify({
+ instruction: pendingUser,
+ input: inputParts.join("\n\n"),
+ output: t.content,
+ }),
+ );
+ context.push(`User: ${pendingUser}`, `Assistant: ${t.content}`);
+ pendingUser = null;
+ }
+ return lines;
+ }
+ return [JSON.stringify({ messages: turns })];
+}
+
+/** Every non-archived chat (Recents and Projects) as training-ready JSONL. */
+export async function buildFineTuneJsonl(
+ format: FineTuneFormat = "openai",
+): Promise {
+ const threads = await listStoredChatThreads({ includeArchived: false });
+ const ids = [...new Set(threads.map((t) => t.id))];
+ const lines: string[] = [];
+ let conversations = 0;
+ let skipped = 0;
+ for (const id of ids) {
+ const raw = await listStoredChatMessages(id);
+ const hasParentIds = raw.some(
+ (m) => (m as { parentId?: unknown }).parentId != null,
+ );
+ // Chain only: retries/regenerations leave sibling branches, and mixing
+ // alternate replies into one conversation corrupts the training targets.
+ const ordered = hasParentIds
+ ? (orderByParentChain(raw, { includeSiblings: false }) as typeof raw)
+ : raw;
+ const turns = messagesToFineTuneTurns(ordered);
+ const converted = turns ? turnsToFineTuneLines(turns, format) : [];
+ if (converted.length === 0) {
+ skipped += 1;
+ continue;
+ }
+ conversations += 1;
+ lines.push(...converted);
+ }
+ return { lines, conversations, skipped };
+}
+
+/** Download the fine-tuning JSONL; returns the conversation count. */
+export async function exportFineTuneJsonl(
+ format: FineTuneFormat = "openai",
+): Promise {
+ const { lines, conversations, skipped } = await buildFineTuneJsonl(format);
+ if (conversations === 0) {
+ toast.info("No chats with a user and assistant exchange to export.");
+ return 0;
+ }
+ const suffix = format === "openai" ? "" : `-${format}`;
+ downloadBlob(
+ lines.join("\n"),
+ `chat-finetune${suffix}-${exportTs()}.jsonl`,
+ "application/x-ndjson",
+ );
+ if (skipped > 0) {
+ toast.success(
+ `Exported ${conversations} conversation${conversations === 1 ? "" : "s"} (${skipped} without a full exchange skipped).`,
+ );
+ }
+ return conversations;
+}
+
// role:"tool" results are absorbed into the preceding assistant tool-call
// part's `result` field rather than becoming separate records.
function oaiMessagesToRecords(
diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx
index 019d0bfd8a..4a8740ac9b 100644
--- a/studio/frontend/src/features/chat/runtime-provider.tsx
+++ b/studio/frontend/src/features/chat/runtime-provider.tsx
@@ -56,6 +56,11 @@ import { AudioAttachmentAdapter } from "./audio-attachment-adapter";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import { ToolPaneScopeContext, toolPaneScope } from "./tool-output-scope";
import type { MessageRecord, ModelType, ThreadRecord } from "./types";
+import {
+ chatContentPartAttachmentIdFromSignature,
+ chatContentPartAttachmentSignature,
+ onChatAttachmentDeleted,
+} from "./utils/chat-attachment-events";
import {
deleteStoredChatThreads,
ensureStoredChatThread,
@@ -890,6 +895,168 @@ function useStudioRuntimeAdapters(
): StudioRuntimeAdapters {
const aui = useAui();
+ // Mirror Data-tab attachment deletions into the loaded thread. The in-memory
+ // repository otherwise keeps the attachment, and a later repo-to-storage sync
+ // (e.g. deleting a message in the thread) would write it back.
+ useEffect(() => {
+ let active = true;
+ let pendingDeletion = Promise.resolve();
+ const unsubscribe = onChatAttachmentDeleted((event) => {
+ pendingDeletion = pendingDeletion.then(async () => {
+ if (!active) return;
+ const { messageId, attachmentId } = event;
+ try {
+ const thread = aui.thread();
+ if (attachmentId.startsWith("content-part-sha256-")) {
+ for (let attempt = 0; attempt < 3 && active; attempt += 1) {
+ const exported = thread.export();
+ const target = exported.messages.find(
+ (item) => item.message.id === messageId,
+ );
+ if (!target || !Array.isArray(target.message.content)) return;
+ const content = target.message.content;
+
+ const signatures = content.map((part) =>
+ chatContentPartAttachmentSignature(part),
+ );
+ const ids = await Promise.all(
+ signatures.map((signature) =>
+ signature === null
+ ? null
+ : chatContentPartAttachmentIdFromSignature(signature),
+ ),
+ );
+ const targetAttachments = (
+ target.message as {
+ attachments?: readonly { id: string }[];
+ }
+ ).attachments;
+ const hasTargetAttachment =
+ Array.isArray(targetAttachments) &&
+ targetAttachments.some(
+ (attachment) => attachment.id === attachmentId,
+ );
+ if (
+ (!ids.includes(attachmentId) && !hasTargetAttachment) ||
+ !active
+ ) {
+ return;
+ }
+
+ // Preserve any messages added or streamed while WebCrypto ran.
+ // Retry if the target's managed content itself changed.
+ const latest = thread.export();
+ const latestTarget = latest.messages.find(
+ (item) => item.message.id === messageId,
+ );
+ const latestContent = latestTarget?.message.content;
+ if (!Array.isArray(latestContent)) return;
+ const latestSignatures = latestContent.map((part) =>
+ chatContentPartAttachmentSignature(part),
+ );
+ if (
+ signatures.length !== latestSignatures.length ||
+ signatures.some(
+ (signature, index) => signature !== latestSignatures[index],
+ )
+ ) {
+ continue;
+ }
+
+ const messages = latest.messages.map((item) => {
+ if (item.message.id !== messageId) return item;
+ const attachments = (
+ item.message as {
+ attachments?: readonly { id: string }[];
+ }
+ ).attachments;
+ return {
+ ...item,
+ message: {
+ ...item.message,
+ content: latestContent.filter(
+ (_, index) => ids[index] !== attachmentId,
+ ),
+ ...(Array.isArray(attachments)
+ ? {
+ attachments: attachments.filter(
+ (attachment) =>
+ attachment.id !== attachmentId,
+ ),
+ }
+ : {}),
+ } as typeof item.message,
+ };
+ });
+ if (active) thread.import({ ...latest, messages });
+ return;
+ }
+ return;
+ }
+
+ const exported = thread.export();
+ let changed = false;
+ const messages = exported.messages.map((item) => {
+ if (item.message.id !== messageId) return item;
+ const message = item.message;
+ const attachments = (
+ message as { attachments?: readonly { id: string }[] }
+ ).attachments;
+ if (
+ Array.isArray(attachments) &&
+ attachments.some(
+ (attachment) => attachment.id === attachmentId,
+ )
+ ) {
+ changed = true;
+ return {
+ ...item,
+ message: {
+ ...message,
+ attachments: attachments.filter(
+ (attachment) => attachment.id !== attachmentId,
+ ),
+ } as typeof message,
+ };
+ }
+ if (/^content-part-[0-9]+$/.test(attachmentId)) {
+ // Legacy synthetic id for a blob stored as a message content part.
+ const idx = Number(attachmentId.slice("content-part-".length));
+ const content = message.content;
+ if (
+ !Array.isArray(content) ||
+ !Number.isInteger(idx) ||
+ idx < 0 ||
+ idx >= content.length
+ ) {
+ return item;
+ }
+ const part = content[idx] as { type?: string };
+ if (part?.type !== "image" && part?.type !== "audio") return item;
+ changed = true;
+ return {
+ ...item,
+ message: {
+ ...message,
+ content: content.filter((_, i) => i !== idx),
+ } as typeof message,
+ };
+ }
+ return item;
+ });
+ if (changed && active) thread.import({ ...exported, messages });
+ } catch {
+ // No active thread mounted: storage already holds the truth.
+ }
+ });
+ return pendingDeletion;
+ });
+ return () => {
+ active = false;
+ unsubscribe();
+ };
+ }, [aui]);
+
const history = useMemo(
() => ({
async load() {
diff --git a/studio/frontend/src/features/chat/utils/archived-chat-export.ts b/studio/frontend/src/features/chat/utils/archived-chat-export.ts
new file mode 100644
index 0000000000..834dfdcf5f
--- /dev/null
+++ b/studio/frontend/src/features/chat/utils/archived-chat-export.ts
@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// Minimal views over the `unknown[]` export fields we filter on.
+type ExportThreadView = {
+ id?: string;
+ archived?: boolean;
+ projectId?: string | null;
+};
+type ExportMessageView = { threadId?: string };
+type ExportProjectView = { id?: string };
+
+// Full chat-export backup shape, kept structural so the pure filter below
+// stays decoupled from the storage layer that produces it.
+export interface ChatExportData {
+ exportedAt?: string;
+ version?: number;
+ threadCount: number;
+ projects?: unknown[];
+ threads: unknown[];
+ messages: unknown[];
+}
+
+// Restrict a full chat export to archived threads, their messages and the
+// projects those threads belong to. Pure: never mutates the input, and keeps
+// the original thread/message objects so the backup re-imports unchanged.
+export function filterArchivedChatExport(
+ full: T,
+): { data: T; archivedCount: number } {
+ const archivedThreads = (full.threads as ExportThreadView[]).filter(
+ (thread) => thread.archived === true,
+ );
+ const archivedThreadIds = new Set(
+ archivedThreads
+ .map((thread) => thread.id)
+ .filter((id): id is string => typeof id === "string"),
+ );
+ const messages = (full.messages as ExportMessageView[]).filter(
+ (message) =>
+ typeof message.threadId === "string" &&
+ archivedThreadIds.has(message.threadId),
+ );
+ const referencedProjectIds = new Set(
+ archivedThreads
+ .map((thread) => thread.projectId)
+ .filter((id): id is string => typeof id === "string"),
+ );
+ const projects = (full.projects as ExportProjectView[] | undefined)?.filter(
+ (project) =>
+ typeof project.id === "string" && referencedProjectIds.has(project.id),
+ );
+ return {
+ data: {
+ ...full,
+ threadCount: archivedThreads.length,
+ projects: projects ?? [],
+ threads: archivedThreads as unknown[],
+ messages: messages as unknown[],
+ },
+ archivedCount: archivedThreads.length,
+ };
+}
diff --git a/studio/frontend/src/features/chat/utils/chat-attachment-events.ts b/studio/frontend/src/features/chat/utils/chat-attachment-events.ts
new file mode 100644
index 0000000000..dbde157890
--- /dev/null
+++ b/studio/frontend/src/features/chat/utils/chat-attachment-events.ts
@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+/**
+ * Notifies loaded chat runtimes when the Data tab deletes a stored attachment.
+ * Without this, the active thread's in-memory repository still holds the
+ * attachment, and any later repo-to-storage sync (e.g. deleting a message in
+ * that thread) writes it back, undoing the deletion.
+ */
+
+import forge from "node-forge";
+
+export type ChatAttachmentDeletedEvent = {
+ messageId: string;
+ attachmentId: string;
+};
+
+const CONTENT_PART_ID_PREFIX = "content-part-sha256-";
+const URI_SCHEME_RE = /^[A-Za-z][A-Za-z0-9+.-]*:/;
+
+function isLocallyStoredBlob(value: string): boolean {
+ const candidate = value.trimStart();
+ if (!candidate) return false;
+ if (candidate.slice(0, 5).toLowerCase() === "data:") return true;
+ if (candidate.startsWith("//") || candidate.startsWith("\\\\")) {
+ return false;
+ }
+ return !URI_SCHEME_RE.test(candidate);
+}
+
+function stableJson(value: unknown): string {
+ if (Array.isArray(value)) {
+ return `[${value
+ .map((item) => (item === undefined ? "null" : stableJson(item)))
+ .join(",")}]`;
+ }
+ if (value && typeof value === "object") {
+ const record = value as Record;
+ return `{${Object.keys(record)
+ .filter((key) => record[key] !== undefined)
+ .sort()
+ .map((key) => `${JSON.stringify(key)}:${stableJson(record[key])}`)
+ .join(",")}}`;
+ }
+ return JSON.stringify(value) ?? "null";
+}
+
+/** Canonical payload used to detect whether an async hash still describes the
+ * current message content. */
+export function chatContentPartAttachmentSignature(
+ part: unknown,
+): string | null {
+ if (!part || typeof part !== "object") return null;
+ const record = part as Record;
+ let payload: ["image" | "audio", unknown] | null = null;
+ if (
+ typeof record.image === "string" &&
+ record.image.slice(0, 5).toLowerCase() === "data:"
+ ) {
+ payload = ["image", record.image];
+ } else if (
+ typeof record.audio === "string" &&
+ isLocallyStoredBlob(record.audio)
+ ) {
+ payload = ["audio", record.audio];
+ } else if (record.audio && typeof record.audio === "object") {
+ const data = (record.audio as Record).data;
+ if (typeof data === "string" && isLocallyStoredBlob(data)) {
+ payload = ["audio", record.audio];
+ }
+ }
+ if (!payload) return null;
+
+ return stableJson(payload);
+}
+
+/** Mirrors the backend's stable content-part identity without adding private
+ * metadata to the message payload sent to inference. */
+export async function chatContentPartAttachmentIdFromSignature(
+ signature: string,
+): Promise {
+ let hex: string | null = null;
+ const subtle = globalThis.crypto?.subtle;
+ if (subtle) {
+ try {
+ const digest = await subtle.digest(
+ "SHA-256",
+ new TextEncoder().encode(signature),
+ );
+ hex = Array.from(new Uint8Array(digest), (byte) =>
+ byte.toString(16).padStart(2, "0"),
+ ).join("");
+ } catch {
+ // Fall through to the pure-JS implementation below. Some embedded
+ // browsers expose crypto.subtle but reject it outside a secure context.
+ }
+ }
+ if (hex === null) {
+ const digest = forge.md.sha256.create();
+ digest.update(signature, "utf8");
+ hex = digest.digest().toHex();
+ }
+ return `${CONTENT_PART_ID_PREFIX}${hex}`;
+}
+
+type Listener = (event: ChatAttachmentDeletedEvent) => void | Promise;
+
+const listeners = new Set();
+
+export function onChatAttachmentDeleted(listener: Listener): () => void {
+ listeners.add(listener);
+ return () => {
+ listeners.delete(listener);
+ };
+}
+
+export function emitChatAttachmentDeleted(
+ event: ChatAttachmentDeletedEvent,
+): void {
+ for (const listener of [...listeners]) {
+ void listener(event);
+ }
+}
diff --git a/studio/frontend/src/features/chat/utils/download-json.ts b/studio/frontend/src/features/chat/utils/download-json.ts
new file mode 100644
index 0000000000..0c5ce00ff9
--- /dev/null
+++ b/studio/frontend/src/features/chat/utils/download-json.ts
@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// Save `data` as a pretty-printed JSON file via a temporary object URL. Uses
+// only the standard Blob/anchor download path so it works in every browser.
+export function triggerJsonDownload(data: unknown, filename: string): void {
+ const blob = new Blob([JSON.stringify(data, null, 2)], {
+ type: "application/json",
+ });
+ const url = URL.createObjectURL(blob);
+ const a = document.createElement("a");
+ a.href = url;
+ a.download = filename;
+ document.body.appendChild(a);
+ a.click();
+ a.remove();
+ URL.revokeObjectURL(url);
+}
diff --git a/studio/frontend/src/features/chat/utils/export-chat-history.ts b/studio/frontend/src/features/chat/utils/export-chat-history.ts
index 5faf4dc08a..b4bc64d053 100644
--- a/studio/frontend/src/features/chat/utils/export-chat-history.ts
+++ b/studio/frontend/src/features/chat/utils/export-chat-history.ts
@@ -1,21 +1,34 @@
// 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 { filterArchivedChatExport } from "./archived-chat-export";
import { buildStoredChatExport } from "./chat-history-storage";
+import { triggerJsonDownload } from "./download-json";
export const buildChatExport = buildStoredChatExport;
+function dateStamp(): string {
+ // Date only (no colons) so the filename is valid on every OS.
+ return new Date().toISOString().slice(0, 10);
+}
+
export async function downloadChatExport(): Promise {
const data = await buildChatExport();
- const blob = new Blob([JSON.stringify(data, null, 2)], {
- type: "application/json",
- });
- const url = URL.createObjectURL(blob);
- const a = document.createElement("a");
- a.href = url;
- a.download = `unsloth-chats-${new Date().toISOString().slice(0, 10)}.json`;
- document.body.appendChild(a);
- a.click();
- a.remove();
- URL.revokeObjectURL(url);
+ triggerJsonDownload(data, `unsloth-chats-${dateStamp()}.json`);
+}
+
+// Full backup restricted to archived chats. Returns the archived thread count.
+export async function buildArchivedChatExport() {
+ return filterArchivedChatExport(await buildChatExport());
+}
+
+// Download only the archived chats. Returns how many were exported; skips the
+// download entirely when there are none.
+export async function downloadArchivedChatExport(): Promise {
+ const { data, archivedCount } = await buildArchivedChatExport();
+ if (archivedCount === 0) {
+ return 0;
+ }
+ triggerJsonDownload(data, `unsloth-archived-chats-${dateStamp()}.json`);
+ return archivedCount;
}
diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts
index 20800230b5..c0bad1a9d5 100644
--- a/studio/frontend/src/features/rag/api/rag-api.ts
+++ b/studio/frontend/src/features/rag/api/rag-api.ts
@@ -10,6 +10,7 @@ import type {
KnowledgeBase,
PreviewTarget,
RagDocument,
+ UploadedDocument,
} from "../types/rag";
const RAG_BASE = "/api/rag";
@@ -194,10 +195,25 @@ export function invalidateProjectSources(projectId: string): void {
projectSourcesCache.delete(projectId);
}
-export function deleteDocument(documentId: string): Promise<{ ok: boolean }> {
- return ragRequest(`/documents/${encodeURIComponent(documentId)}`, {
- method: "DELETE",
- });
+export async function listAllDocuments(): Promise {
+ const data = await ragRequest<{ documents: UploadedDocument[] }>(
+ "/documents",
+ );
+ return data.documents ?? [];
+}
+
+export async function deleteDocument(
+ documentId: string,
+ projectId?: string | null,
+): Promise<{ ok: boolean }> {
+ const result = await ragRequest<{ ok: boolean }>(
+ `/documents/${encodeURIComponent(documentId)}`,
+ {
+ method: "DELETE",
+ },
+ );
+ if (projectId) invalidateProjectSources(projectId);
+ return result;
}
export function getJob(jobId: string): Promise {
@@ -237,7 +253,8 @@ export async function* streamJobEvents(
const dataLines: string[] = [];
for (const line of rawEvent.split(/\r?\n/)) {
- if (line.startsWith("data:")) dataLines.push(line.slice(5).trimStart());
+ if (line.startsWith("data:"))
+ dataLines.push(line.slice(5).trimStart());
}
if (dataLines.length > 0) {
const dataText = dataLines.join("\n");
diff --git a/studio/frontend/src/features/rag/components/use-rag-documents.ts b/studio/frontend/src/features/rag/components/use-rag-documents.ts
index 8e756b2782..8d6433d8c3 100644
--- a/studio/frontend/src/features/rag/components/use-rag-documents.ts
+++ b/studio/frontend/src/features/rag/components/use-rag-documents.ts
@@ -2,12 +2,11 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useCallback, useEffect, useRef, useState } from "react";
-import { useChatRuntimeStore } from "@/features/chat";
-
import {
CHAT_RAG_CAPTION_KEY,
CHAT_RAG_OCR_KEY,
-} from "@/features/chat/stores/chat-runtime-store";
+ useChatRuntimeStore,
+} from "@/features/chat";
import { toast } from "@/lib/toast";
import {
deleteDocument,
@@ -60,9 +59,7 @@ export function useRagDocuments(
if (ids.size === 0) return false;
const docs = documentsRef.current.filter((d) => ids.has(d.id));
if (docs.length === 0) return false; // sig tracked but doc gone -> allow re-upload
- return docs.some(
- (d) => d.status !== "completed" || (d.numChunks ?? 0) > 0,
- );
+ return docs.some((d) => d.status !== "completed" || (d.numChunks ?? 0) > 0);
}, []);
// True while upload() runs, so the scope-change effect can tell a real switch
// from lazy thread materialization mid-upload (which must not reset).
@@ -80,9 +77,7 @@ export function useRagDocuments(
const patchDoc = useCallback(
(documentId: string, patch: Partial) => {
setDocuments((rows) =>
- rows.map((row) =>
- row.id === documentId ? { ...row, ...patch } : row,
- ),
+ rows.map((row) => (row.id === documentId ? { ...row, ...patch } : row)),
);
},
[],
@@ -176,49 +171,61 @@ export function useRagDocuments(
[patchDoc],
);
- const refresh = useCallback(async (opts?: { quiet?: boolean }) => {
- if (!scope) return;
- if (!opts?.quiet) setLoading(true);
- try {
- // Merge server truth with local progress so a refresh mid-index keeps a
- // live "running %" chip. Failed docs hidden (toast warned at upload).
- const rows = (await lister()).filter((row) => row.status !== "failed");
- setDocuments((prev) => {
- const merged = rows.map((row) => {
- const tracked = prev.find((p) => p.id === row.id);
- return tracked && tracked.progress != null && row.status !== "completed"
- ? { ...row, progress: tracked.progress }
- : row;
+ const refresh = useCallback(
+ async (opts?: { quiet?: boolean }) => {
+ if (!scope) return;
+ if (!opts?.quiet) setLoading(true);
+ try {
+ // Merge server truth with local progress so a refresh mid-index keeps a
+ // live "running %" chip. Failed docs hidden (toast warned at upload).
+ const rows = (await lister()).filter((row) => row.status !== "failed");
+ setDocuments((prev) => {
+ const merged = rows.map((row) => {
+ const tracked = prev.find((p) => p.id === row.id);
+ return tracked &&
+ tracked.progress != null &&
+ row.status !== "completed"
+ ? { ...row, progress: tracked.progress }
+ : row;
+ });
+ // Keep optimistic chips (not yet listed) so a refresh racing an upload
+ // can't make them vanish.
+ const serverIds = new Set(rows.map((row) => row.id));
+ const pendingLocal = prev.filter(
+ (row) => row.id.startsWith("pending_") && !serverIds.has(row.id),
+ );
+ return [...merged, ...pendingLocal];
});
- // Keep optimistic chips (not yet listed) so a refresh racing an upload
- // can't make them vanish.
- const serverIds = new Set(rows.map((row) => row.id));
- const pendingLocal = prev.filter(
- (row) => row.id.startsWith("pending_") && !serverIds.has(row.id),
- );
- return [...merged, ...pendingLocal];
- });
- } catch (err) {
- toast.error("Failed to load documents", {
- description: err instanceof Error ? err.message : String(err),
- });
- } finally {
- if (!opts?.quiet) setLoading(false);
- }
- }, [scope, lister]);
+ } catch (err) {
+ toast.error("Failed to load documents", {
+ description: err instanceof Error ? err.message : String(err),
+ });
+ } finally {
+ if (!opts?.quiet) setLoading(false);
+ }
+ },
+ [scope, lister],
+ );
// A real switch (thread/KB swap) resets + reloads; first acquiring a scope just
// loads. Skip both during materialization mid-upload (scope null -> new thread
// while upload() runs) so we don't abort tracking or wipe optimistic chips.
useEffect(() => {
+ const jobs = trackedJobs.current;
const prev = prevScopeKeyRef.current;
prevScopeKeyRef.current = scopeKey;
if (prev !== null && prev !== scopeKey) {
- for (const controller of trackedJobs.current.values()) controller.abort();
- trackedJobs.current.clear();
+ for (const controller of jobs.values()) controller.abort();
+ jobs.clear();
sigByDocId.current.clear();
+ // Scope changes intentionally clear the old scope before fetching the new
+ // one. Keep this synchronous so React StrictMode's setup/cleanup replay
+ // cannot cancel the only refresh after prevScopeKeyRef has advanced.
setDocuments([]);
- if (scope) void refresh();
+ if (scope) {
+ // eslint-disable-next-line react-hooks/set-state-in-effect
+ void refresh();
+ }
} else if (prev === null && scope && !uploadInFlightRef.current) {
void refresh();
}
@@ -226,8 +233,8 @@ export function useRagDocuments(
// Preserve in-flight tracking when cleanup is the materialization flip,
// not a real switch/unmount.
if (uploadInFlightRef.current) return;
- for (const controller of trackedJobs.current.values()) controller.abort();
- trackedJobs.current.clear();
+ for (const controller of jobs.values()) controller.abort();
+ jobs.clear();
};
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [scopeKey]);
@@ -260,21 +267,41 @@ export function useRagDocuments(
// otherwise backend env defaults own the ingest policy.
const state = useChatRuntimeStore.getState();
const hasLocal = (key: string) =>
- typeof window !== "undefined" && window.localStorage.getItem(key) !== null;
- const ocr = hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined;
+ typeof window !== "undefined" &&
+ window.localStorage.getItem(key) !== null;
+ const ocr = hasLocal(CHAT_RAG_OCR_KEY)
+ ? state.ragOcrScanned
+ : undefined;
const caption = hasLocal(CHAT_RAG_CAPTION_KEY)
? state.ragCaptionFigures
: undefined;
const result =
activeScope.type === "kb"
- ? await uploadKnowledgeBaseDocument(activeScope.kbId, file, ocr, caption)
+ ? await uploadKnowledgeBaseDocument(
+ activeScope.kbId,
+ file,
+ ocr,
+ caption,
+ )
: activeScope.type === "project"
- ? await uploadProjectDocument(activeScope.projectId, file, ocr, caption)
- : await uploadThreadDocument(activeScope.threadId, file, ocr, caption);
+ ? await uploadProjectDocument(
+ activeScope.projectId,
+ file,
+ ocr,
+ caption,
+ )
+ : await uploadThreadDocument(
+ activeScope.threadId,
+ file,
+ ocr,
+ caption,
+ );
sigByDocId.current.set(result.documentId, fileSignature(file));
if (seenIds.has(result.documentId)) {
setDocuments((rows) => rows.filter((row) => row.id !== tempId));
- toast.info(`${result.filename || file.name} is already indexed - skipping`);
+ toast.info(
+ `${result.filename || file.name} is already indexed - skipping`,
+ );
return;
}
seenIds.add(result.documentId);
@@ -341,7 +368,9 @@ export function useRagDocuments(
]);
const resolved =
- overrideScope instanceof Promise ? await overrideScope : overrideScope;
+ overrideScope instanceof Promise
+ ? await overrideScope
+ : overrideScope;
const activeScope = resolved ?? scope;
if (!activeScope) {
// Materialization failed: drop the chips so they don't hang "pending".
@@ -377,7 +406,10 @@ export function useRagDocuments(
const prevSig = sigByDocId.current.get(documentId);
sigByDocId.current.delete(documentId);
try {
- await deleteDocument(documentId);
+ await deleteDocument(
+ documentId,
+ scope?.type === "project" ? scope.projectId : undefined,
+ );
} catch (err) {
setDocuments(prev);
if (prevSig !== undefined) sigByDocId.current.set(documentId, prevSig);
@@ -386,7 +418,7 @@ export function useRagDocuments(
});
}
},
- [documents],
+ [documents, scope],
);
return { documents, loading, uploading, refresh, upload, remove };
diff --git a/studio/frontend/src/features/rag/index.ts b/studio/frontend/src/features/rag/index.ts
index 9e35e345ee..b06c7e09cc 100644
--- a/studio/frontend/src/features/rag/index.ts
+++ b/studio/frontend/src/features/rag/index.ts
@@ -5,4 +5,9 @@ export { KnowledgeBaseComposerButton } from "./components/knowledge-base-compose
export { KnowledgeBaseDialog } from "./components/knowledge-base-dialog";
export { RetrievalSettingsSection } from "./components/retrieval-settings-section";
export { ThreadDocumentsBar } from "./components/thread-documents-bar";
-export type { KnowledgeBase, RagDocument } from "./types/rag";
+export {
+ deleteDocument,
+ getDocumentFileUrl,
+ listAllDocuments,
+} from "./api/rag-api";
+export type { KnowledgeBase, RagDocument, UploadedDocument } from "./types/rag";
diff --git a/studio/frontend/src/features/rag/types/rag.ts b/studio/frontend/src/features/rag/types/rag.ts
index 1277500ae6..9922c740af 100644
--- a/studio/frontend/src/features/rag/types/rag.ts
+++ b/studio/frontend/src/features/rag/types/rag.ts
@@ -24,6 +24,13 @@ export interface RagDocument {
createdAt?: string | null;
}
+/** RagDocument enriched for the global uploaded-files list (settings Data tab). */
+export interface UploadedDocument extends RagDocument {
+ sizeBytes?: number | null;
+ kbName?: string | null;
+ projectName?: string | null;
+}
+
export interface DocumentUploadResult {
documentId: string;
jobId: string;
diff --git a/studio/frontend/src/features/settings/components/archived-chats-dialog.tsx b/studio/frontend/src/features/settings/components/archived-chats-dialog.tsx
index 836b4d25d4..c4932ccc5b 100644
--- a/studio/frontend/src/features/settings/components/archived-chats-dialog.tsx
+++ b/studio/frontend/src/features/settings/components/archived-chats-dialog.tsx
@@ -12,18 +12,12 @@ import {
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import {
- Dialog,
- DialogContent,
- DialogHeader,
- DialogTitle,
-} from "@/components/ui/dialog";
-import {
+ type SidebarItem,
deleteChatItem,
unarchiveChatItem,
useChatPreferencesStore,
useChatRuntimeStore,
useChatSidebarItems,
- type SidebarItem,
} from "@/features/chat";
import { toast } from "@/lib/toast";
import { ArchiveRestoreIcon, Delete02Icon } from "@hugeicons/core-free-icons";
@@ -40,13 +34,7 @@ function formatCreatedAt(ms: number): string {
});
}
-export function ArchivedChatsDialog({
- open,
- onOpenChange,
-}: {
- open: boolean;
- onOpenChange: (open: boolean) => void;
-}) {
+export function ArchivedChatsView() {
const { archivedItems } = useChatSidebarItems({ requireMessages: false });
const navigate = useNavigate();
const closeSettings = useSettingsDialogStore((s) => s.closeDialog);
@@ -74,7 +62,6 @@ export function ArchivedChatsDialog({
search:
item.type === "single" ? { thread: item.id } : { compare: item.id },
});
- onOpenChange(false);
closeSettings();
}
@@ -114,72 +101,66 @@ export function ArchivedChatsDialog({
}
return (
-
-
-
- Archived chats
-
-
- {archivedItems.length === 0 ? (
-
- No archived chats.
-
- ) : (
-
-
- Name
- Date created
-
-
- {archivedItems.map((item) => (
-
+ {archivedItems.length === 0 ? (
+
+ No archived chats.
+
+ ) : (
+
+
+ Name
+ Date created
+
+
+ {archivedItems.map((item) => (
+
+ openChat(item)}
+ className="min-w-0 flex-1 truncate text-left text-primary hover:underline"
+ title={item.title}
>
+ {item.title}
+
+
+ {formatCreatedAt(item.createdAt)}
+
+
openChat(item)}
- className="min-w-0 flex-1 truncate text-left text-primary hover:underline"
- title={item.title}
+ onClick={() => void handleUnarchive(item)}
+ aria-label="Unarchive chat"
+ title="Unarchive"
+ className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
>
- {item.title}
+
-
- {formatCreatedAt(item.createdAt)}
-
-
- void handleUnarchive(item)}
- aria-label="Unarchive chat"
- title="Unarchive"
- className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
- >
-
-
- requestDelete(item)}
- aria-label="Delete chat"
- title="Delete"
- className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive"
- >
-
-
-
-
- ))}
-
- )}
-
+
requestDelete(item)}
+ aria-label="Delete chat"
+ title="Delete"
+ className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive"
+ >
+
+
+
+
+ ))}
+
+ )}
-
+
);
}
diff --git a/studio/frontend/src/features/settings/components/finetune-recipe.ts b/studio/frontend/src/features/settings/components/finetune-recipe.ts
new file mode 100644
index 0000000000..c98f422a35
--- /dev/null
+++ b/studio/frontend/src/features/settings/components/finetune-recipe.ts
@@ -0,0 +1,98 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// Settings Data tab glue: turn chat history into a fine-tuning JSONL, stage
+// it as a Data Recipe seed upload, and open a new recipe on that file.
+
+import { type FineTuneFormat, buildFineTuneJsonl } from "@/features/chat";
+import { saveRecipe } from "@/features/data-recipes/data/recipes-db";
+import { createEmptyRecipePayload } from "@/features/recipe-studio";
+import { inspectSeedUpload } from "@/features/recipe-studio/api";
+import { uploadTrainingDataset } from "@/features/training/api/datasets-api";
+import { useTrainingConfigStore } from "@/features/training/stores/training-config-store";
+import { toast } from "@/lib/toast";
+
+/** btoa cannot handle code points above latin-1, so encode UTF-8 bytes. */
+function base64FromString(value: string): string {
+ const bytes = new TextEncoder().encode(value);
+ let binary = "";
+ const CHUNK = 0x8000;
+ for (let i = 0; i < bytes.length; i += CHUNK) {
+ binary += String.fromCharCode(...bytes.subarray(i, i + CHUNK));
+ }
+ return btoa(binary);
+}
+
+/** Builds the JSONL, uploads it as a local recipe seed, and saves a new
+ * recipe whose seed block points at the file. Returns the recipe id, or
+ * null when there is nothing to export. */
+export async function createFineTuneRecipeFromChats(
+ format: FineTuneFormat = "openai",
+): Promise {
+ const { lines, conversations } = await buildFineTuneJsonl(format);
+ if (conversations === 0) {
+ toast.info("No chats with a user and assistant exchange to export.");
+ return null;
+ }
+
+ const dateLabel = new Date().toISOString().slice(0, 10);
+ const suffix = format === "openai" ? "" : `-${format}`;
+ const filename = `chat-finetune${suffix}-${dateLabel}.jsonl`;
+ const inspected = await inspectSeedUpload({
+ filename,
+ // biome-ignore lint/style/useNamingConvention: api schema
+ content_base64: base64FromString(lines.join("\n")),
+ // biome-ignore lint/style/useNamingConvention: api schema
+ preview_size: 10,
+ });
+
+ const payload = createEmptyRecipePayload();
+ payload.recipe.seed_config = {
+ source: {
+ // biome-ignore lint/style/useNamingConvention: api schema
+ seed_type: "local",
+ path: inspected.resolved_path,
+ },
+ // biome-ignore lint/style/useNamingConvention: api schema
+ sampling_strategy: "ordered",
+ // biome-ignore lint/style/useNamingConvention: api schema
+ selection_strategy: null,
+ };
+ payload.ui.nodes = [{ id: "seed", x: 0, y: 0, width: 400 }];
+ payload.ui.seed_source_type = "local";
+ payload.ui.seed_columns = inspected.columns;
+ payload.ui.seed_preview_rows = inspected.preview_rows ?? [];
+ payload.ui.local_file_name = filename;
+
+ const record = await saveRecipe({
+ name: `Chat fine-tuning ${dateLabel}`,
+ payload,
+ });
+ return record.id;
+}
+
+/** Builds the JSONL, uploads it as a training dataset, and selects it in the
+ * Train tab's config store so the Train page opens with it loaded. Returns
+ * false when there is nothing to export. */
+export async function loadFineTuneDatasetInTrainTab(
+ format: FineTuneFormat = "openai",
+): Promise {
+ const { lines, conversations } = await buildFineTuneJsonl(format);
+ if (conversations === 0) {
+ toast.info("No chats with a user and assistant exchange to export.");
+ return false;
+ }
+
+ const dateLabel = new Date().toISOString().slice(0, 10);
+ const suffix = format === "openai" ? "" : `-${format}`;
+ const file = new File(
+ [lines.join("\n")],
+ `chat-finetune${suffix}-${dateLabel}.jsonl`,
+ { type: "application/x-ndjson" },
+ );
+ const uploaded = await uploadTrainingDataset(file);
+ // Selecting also kicks off the dataset format check, so the Train tab
+ // shows the detected format as soon as it mounts.
+ useTrainingConfigStore.getState().selectLocalDataset(uploaded.stored_path);
+ return true;
+}
diff --git a/studio/frontend/src/features/settings/components/uploaded-files-dialog.tsx b/studio/frontend/src/features/settings/components/uploaded-files-dialog.tsx
new file mode 100644
index 0000000000..3368f07ff2
--- /dev/null
+++ b/studio/frontend/src/features/settings/components/uploaded-files-dialog.tsx
@@ -0,0 +1,644 @@
+// 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 {
+ AlertDialog,
+ AlertDialogAction,
+ AlertDialogCancel,
+ AlertDialogContent,
+ AlertDialogDescription,
+ AlertDialogFooter,
+ AlertDialogHeader,
+ AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { Spinner } from "@/components/ui/spinner";
+import {
+ type ChatAttachmentRecord,
+ deleteChatAttachment,
+ emitChatAttachmentDeleted,
+ fetchChatAttachmentBlob,
+ listChatAttachments,
+} from "@/features/chat";
+import {
+ deleteDocument,
+ getDocumentFileUrl,
+ listAllDocuments,
+ type UploadedDocument,
+} from "@/features/rag";
+import { toast } from "@/lib/toast";
+import {
+ ArrowUpRight01Icon,
+ Delete02Icon,
+ File02Icon,
+} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { useNavigate } from "@tanstack/react-router";
+import { type ReactNode, useEffect, useRef, useState } from "react";
+import { useSettingsDialogStore } from "../stores/settings-dialog-store";
+
+function formatUploadedAt(value: string | number | null | undefined): string {
+ if (value === null || value === undefined || value === "") return "-";
+ // Chat attachments carry ms epoch numbers; RAG documents carry SQLite
+ // ISO-ish strings (no timezone). Unparseable strings fall through raw.
+ const parsed = new Date(value);
+ if (Number.isNaN(parsed.getTime())) return String(value);
+ return parsed.toLocaleDateString(undefined, {
+ year: "numeric",
+ month: "long",
+ day: "numeric",
+ });
+}
+
+function formatSize(bytes: number | null | undefined): string {
+ if (bytes === null || bytes === undefined) return "-";
+ if (bytes < 1024) return `${bytes} B`;
+ const units = ["KB", "MB", "GB"];
+ let value = bytes;
+ let unit = "B";
+ for (const next of units) {
+ if (value < 1024) break;
+ value /= 1024;
+ unit = next;
+ }
+ return `${value >= 10 ? Math.round(value) : value.toFixed(1)} ${unit}`;
+}
+
+function ragLocationLabel(doc: UploadedDocument): string {
+ if (doc.kbId) return doc.kbName ? `KB · ${doc.kbName}` : "Knowledge base";
+ if (doc.projectId) {
+ return doc.projectName ? `Project · ${doc.projectName}` : "Project";
+ }
+ if (doc.threadId) return "Chat files (RAG)";
+ return "-";
+}
+
+/** Short uppercase file-type label from the filename extension, falling back
+ * to the content-type subtype (e.g. "image/webp" gives WEBP). */
+function fileTypeLabel(
+ name: string,
+ contentType?: string | null,
+): string | null {
+ const dot = name.lastIndexOf(".");
+ const ext = dot > 0 ? name.slice(dot + 1).trim() : "";
+ if (ext && ext.length <= 5) return ext.toUpperCase();
+ const subtype = contentType?.split("/")[1]?.split("+")[0]?.trim();
+ return subtype && subtype.length <= 10 ? subtype.toUpperCase() : null;
+}
+
+/** Lazy image thumbnail for a chat attachment; a file icon until it loads.
+ * The stored blob only downloads once the row scrolls into view, so a long
+ * history of screenshots does not fetch every image on open. */
+function ChatImageThumb({
+ messageId,
+ attachmentId,
+}: {
+ messageId: string;
+ attachmentId: string;
+}) {
+ const [src, setSrc] = useState(null);
+ const [visible, setVisible] = useState(false);
+ const holderRef = useRef(null);
+
+ useEffect(() => {
+ const el = holderRef.current;
+ if (!el) return;
+ if (typeof IntersectionObserver === "undefined") {
+ return;
+ }
+ const observer = new IntersectionObserver((entries) => {
+ if (entries.some((entry) => entry.isIntersecting)) {
+ setVisible(true);
+ observer.disconnect();
+ }
+ });
+ observer.observe(el);
+ return () => observer.disconnect();
+ }, []);
+
+ useEffect(() => {
+ if (!visible) return;
+ let cancelled = false;
+ let url: string | null = null;
+ fetchChatAttachmentBlob(messageId, attachmentId)
+ .then((blob) => {
+ if (cancelled) return;
+ url = URL.createObjectURL(blob);
+ setSrc(url);
+ })
+ .catch(() => {
+ // Keep the file icon on failure.
+ });
+ return () => {
+ cancelled = true;
+ if (url) URL.revokeObjectURL(url);
+ };
+ }, [visible, messageId, attachmentId]);
+
+ if (!src) {
+ return (
+
+
+
+ );
+ }
+ return ;
+}
+
+function FileIconThumb() {
+ return (
+
+ );
+}
+
+/** One display row: a RAG document or a chat message attachment. */
+interface UploadedFileRow {
+ key: string;
+ source: "rag" | "chat";
+ name: string;
+ location: string;
+ sizeBytes?: number | null;
+ createdAt?: string | number | null;
+ failed?: boolean;
+ /** Epoch ms for sorting; rows with unknown dates sort last. */
+ sortTime: number;
+ typeLabel: string | null;
+ /** Image rows render a thumbnail; others show a file icon. */
+ thumb: ReactNode;
+ /** Chat rows link back to their thread. */
+ threadId?: string | null;
+ /** Compare-chat rows navigate by pair id instead of opening one pane alone. */
+ pairId?: string | null;
+ open: () => Promise;
+ remove: () => Promise;
+ deleteDescription: string;
+}
+
+function toSortTime(value: string | number | null | undefined): number {
+ if (value === null || value === undefined || value === "") return 0;
+ const parsed = new Date(value).getTime();
+ return Number.isNaN(parsed) ? 0 : parsed;
+}
+
+// Safari and Firefox block window.open after an await (the user gesture is
+// gone), so open a blank tab synchronously and point it at the URL once
+// resolved. A blocked synchronous open is surfaced instead of silently losing
+// the file after the asynchronous URL lookup.
+async function openResolvedUrl(resolve: () => Promise): Promise {
+ const win = window.open("", "_blank");
+ if (!win) {
+ throw new Error(
+ "Your browser blocked the new tab. Allow popups and retry.",
+ );
+ }
+ win.opener = null;
+ let url: string;
+ try {
+ url = await resolve();
+ } catch (err) {
+ win.close();
+ throw err;
+ }
+ win.location.replace(url);
+}
+
+function ragRow(doc: UploadedDocument): UploadedFileRow {
+ return {
+ key: `rag-${doc.id}`,
+ source: "rag",
+ name: doc.filename,
+ location: ragLocationLabel(doc),
+ sizeBytes: doc.sizeBytes,
+ createdAt: doc.createdAt,
+ failed: doc.status === "failed",
+ sortTime: toSortTime(doc.createdAt),
+ typeLabel: fileTypeLabel(doc.filename),
+ // RAG uploads are documents (pdf, txt, md, docx, html), not images.
+ thumb: ,
+ open: () => openResolvedUrl(() => getDocumentFileUrl(doc.id)),
+ remove: async () => {
+ await deleteDocument(doc.id, doc.projectId);
+ },
+ deleteDescription:
+ "The file and its indexed content are removed. This cannot be undone.",
+ };
+}
+
+function chatAttachmentRow(att: ChatAttachmentRecord): UploadedFileRow {
+ const isImage =
+ att.type === "image" || Boolean(att.contentType?.startsWith("image/"));
+ return {
+ key: `chat-${att.messageId}-${att.id}`,
+ source: "chat",
+ name: att.name,
+ location: att.threadTitle ? `Chat · ${att.threadTitle}` : "Chat",
+ sizeBytes: att.sizeBytes,
+ createdAt: att.createdAt,
+ sortTime: toSortTime(att.createdAt),
+ typeLabel: fileTypeLabel(att.name, att.contentType),
+ threadId: att.threadId,
+ pairId: att.pairId,
+ thumb: isImage ? (
+
+ ) : (
+
+ ),
+ open: () =>
+ openResolvedUrl(async () => {
+ const blob = await fetchChatAttachmentBlob(att.messageId, att.id);
+ const url = URL.createObjectURL(blob);
+ // Give the new tab time to load the blob before revoking.
+ setTimeout(() => URL.revokeObjectURL(url), 60_000);
+ return url;
+ }),
+ remove: async () => {
+ await deleteChatAttachment(att.messageId, att.id);
+ // Patch any loaded runtime copy so a later repo sync cannot write the
+ // deleted attachment back to storage.
+ emitChatAttachmentDeleted({
+ messageId: att.messageId,
+ attachmentId: att.id,
+ });
+ },
+ deleteDescription:
+ "The attachment is removed from its chat message; the message text is kept. This cannot be undone.",
+ };
+}
+
+type SourceLoad = {
+ status: "loading" | "ready" | "error";
+ data: T;
+ error: string | null;
+};
+
+function errorMessage(error: unknown, fallback: string): string {
+ return error instanceof Error ? error.message : fallback;
+}
+
+/** Inline settings page listing uploaded files from each available source. */
+export function UploadedFilesView() {
+ const [ragFiles, setRagFiles] = useState>({
+ status: "loading",
+ data: [],
+ error: null,
+ });
+ const [chatFiles, setChatFiles] = useState<
+ SourceLoad
+ >({ status: "loading", data: [], error: null });
+ const [chatNextOffset, setChatNextOffset] = useState(null);
+ const [loadingMore, setLoadingMore] = useState(false);
+ const [confirmingDelete, setConfirmingDelete] =
+ useState(null);
+ const navigate = useNavigate();
+
+ const rows = [
+ ...ragFiles.data.map(ragRow),
+ ...chatFiles.data.map(chatAttachmentRow),
+ ].sort((a, b) => b.sortTime - a.sortTime);
+
+ // Jump to the chat thread the attachment lives in, closing the settings
+ // dialog so the thread is actually visible.
+ function goToChat(row: UploadedFileRow) {
+ if (!row.threadId) return;
+ useSettingsDialogStore.getState().closeDialog();
+ if (row.pairId) {
+ void navigate({ to: "/chat", search: { compare: row.pairId } });
+ } else {
+ void navigate({ to: "/chat", search: { thread: row.threadId } });
+ }
+ }
+
+ useEffect(() => {
+ let cancelled = false;
+ void listAllDocuments().then(
+ (data) => {
+ if (!cancelled) setRagFiles({ status: "ready", data, error: null });
+ },
+ (error: unknown) => {
+ if (!cancelled) {
+ setRagFiles({
+ status: "error",
+ data: [],
+ error: errorMessage(error, "Failed to load RAG documents"),
+ });
+ }
+ },
+ );
+ void listChatAttachments().then(
+ (page) => {
+ if (!cancelled) {
+ setChatFiles({
+ status: "ready",
+ data: page.attachments,
+ error: null,
+ });
+ setChatNextOffset(page.nextOffset);
+ }
+ },
+ (error: unknown) => {
+ if (!cancelled) {
+ setChatFiles({
+ status: "error",
+ data: [],
+ error: errorMessage(error, "Failed to load chat attachments"),
+ });
+ }
+ },
+ );
+ return () => {
+ cancelled = true;
+ };
+ }, []);
+
+ function retryRagFiles() {
+ setRagFiles((current) => ({ ...current, status: "loading", error: null }));
+ void listAllDocuments().then(
+ (data) => setRagFiles({ status: "ready", data, error: null }),
+ (error: unknown) =>
+ setRagFiles((current) => ({
+ ...current,
+ status: "error",
+ error: errorMessage(error, "Failed to load RAG documents"),
+ })),
+ );
+ }
+
+ async function loadChatPage(offset: number, append: boolean) {
+ setLoadingMore(true);
+ setChatFiles((current) => ({ ...current, status: "loading", error: null }));
+ try {
+ const page = await listChatAttachments(offset);
+ setChatFiles((current) => ({
+ status: "ready",
+ data: append
+ ? [
+ ...current.data,
+ ...page.attachments.filter(
+ (incoming) =>
+ !current.data.some(
+ (existing) =>
+ existing.id === incoming.id &&
+ existing.messageId === incoming.messageId,
+ ),
+ ),
+ ]
+ : page.attachments,
+ error: null,
+ }));
+ setChatNextOffset(page.nextOffset);
+ } catch (error) {
+ setChatFiles((current) => ({
+ ...current,
+ status: "error",
+ error: errorMessage(error, "Failed to load chat attachments"),
+ }));
+ } finally {
+ setLoadingMore(false);
+ }
+ }
+
+ function retryChatFiles() {
+ const append = chatFiles.data.length > 0 && chatNextOffset !== null;
+ void loadChatPage(append ? chatNextOffset : 0, append);
+ }
+
+ async function handleOpen(row: UploadedFileRow) {
+ try {
+ await row.open();
+ } catch (err) {
+ toast.error("Failed to open file", {
+ description: err instanceof Error ? err.message : undefined,
+ });
+ }
+ }
+
+ async function handleDelete(row: UploadedFileRow) {
+ // Offset pages and destructive mutations must not race: a deletion shifts
+ // the boundary used by an in-flight page request.
+ if (loadingMore) return;
+ try {
+ await row.remove();
+ if (row.source === "rag") {
+ setRagFiles((current) => ({
+ ...current,
+ data: current.data.filter((doc) => `rag-${doc.id}` !== row.key),
+ }));
+ } else {
+ setChatFiles((current) => ({
+ ...current,
+ data: current.data.filter(
+ (attachment) =>
+ `chat-${attachment.messageId}-${attachment.id}` !== row.key,
+ ),
+ }));
+ // Offset pagination is relative to the current server inventory. A
+ // deletion before the next page shifts every later row back by one.
+ setChatNextOffset((current) =>
+ current === null ? null : Math.max(0, current - 1),
+ );
+ }
+ toast.success("File deleted");
+ } catch (err) {
+ toast.error("Failed to delete file", {
+ description: err instanceof Error ? err.message : undefined,
+ });
+ }
+ }
+
+ return (
+
+ {ragFiles.status === "error" ? (
+
+ RAG documents unavailable: {ragFiles.error}
+
+ Retry
+
+
+ ) : null}
+ {chatFiles.status === "error" ? (
+
+ Chat attachments unavailable: {chatFiles.error}
+
+ Retry
+
+
+ ) : null}
+
+ {rows.length === 0 &&
+ (ragFiles.status === "loading" || chatFiles.status === "loading") ? (
+
+
+
+ ) : rows.length === 0 &&
+ ragFiles.status !== "error" &&
+ chatFiles.status !== "error" ? (
+
+ No uploaded files.
+
+ ) : rows.length > 0 ? (
+
+
+ Name
+ Location
+ Uploaded
+
+
+ {rows.map((row) => (
+
+ {/* Clicking the file jumps to its chat; files without one
+ open directly. The theme scales rounded-md up to a near
+ circle at this size, so the thumb pins a small radius. */}
+
+ row.threadId ? goToChat(row) : void handleOpen(row)
+ }
+ title={
+ row.threadId ? `Go to ${row.location}` : `Open ${row.name}`
+ }
+ className="group/name flex min-w-0 flex-1 basis-[calc(100%-5rem)] items-center gap-2.5 overflow-hidden text-left sm:basis-auto"
+ >
+
+ {row.thumb}
+
+
+
+ {/* Floor keeps the name visible when the chip and fixed
+ columns squeeze the cell at narrow widths. */}
+
+ {row.name}
+
+ {row.typeLabel ? (
+
+ {row.typeLabel}
+
+ ) : null}
+ {row.failed ? (
+
+ failed
+
+ ) : null}
+
+
+ {formatSize(row.sizeBytes)}
+
+
+
+ {row.threadId ? (
+ goToChat(row)}
+ title={`Go to ${row.location}`}
+ className="order-3 w-full truncate pl-10 text-left text-muted-foreground underline-offset-2 transition-colors hover:text-foreground hover:underline sm:order-none sm:w-36 sm:pl-0"
+ >
+ {row.location}
+
+ ) : (
+
+ {row.location}
+
+ )}
+
+ {formatUploadedAt(row.createdAt)}
+
+
+ void handleOpen(row)}
+ aria-label={`Open ${row.name}`}
+ title="Open"
+ className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
+ >
+
+
+ setConfirmingDelete(row)}
+ aria-label={`Delete ${row.name}`}
+ title="Delete"
+ className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive disabled:cursor-wait disabled:opacity-50"
+ >
+
+
+
+
+ ))}
+ {chatNextOffset !== null ? (
+
+ void loadChatPage(chatNextOffset, true)}
+ className="rounded-md border border-border px-3 py-1.5 text-sm font-medium hover:bg-muted disabled:cursor-wait disabled:opacity-60"
+ >
+ {loadingMore ? "Loading..." : "Load more chat attachments"}
+
+
+ ) : null}
+
+ ) : null}
+
+
{
+ if (!o) setConfirmingDelete(null);
+ }}
+ >
+
+
+ Delete file
+
+ Delete{" "}
+
+ "{confirmingDelete?.name}"
+
+ ? {confirmingDelete?.deleteDescription}
+
+
+
+ Cancel
+ {
+ const row = confirmingDelete;
+ setConfirmingDelete(null);
+ if (row) void handleDelete(row);
+ }}
+ >
+ Delete
+
+
+
+
+
+ );
+}
diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx
index 08588be8fe..d98d1b8ac0 100644
--- a/studio/frontend/src/features/settings/settings-dialog.tsx
+++ b/studio/frontend/src/features/settings/settings-dialog.tsx
@@ -15,6 +15,7 @@ import {
Cancel01Icon,
CloudIcon,
CpuIcon,
+ DatabaseSettingIcon,
Globe02Icon,
HelpCircleIcon,
Message01Icon,
@@ -43,6 +44,7 @@ import { ApiKeysTab } from "./tabs/api-keys-tab";
import { AppearanceTab } from "./tabs/appearance-tab";
import { ChatTab } from "./tabs/chat-tab";
import { ConnectionsTab } from "./tabs/connections-tab";
+import { DataTab } from "./tabs/data-tab";
import { GeneralTab } from "./tabs/general-tab";
import { ProfileTab } from "./tabs/profile-tab";
import { ResourcesTab } from "./tabs/resources-tab";
@@ -93,6 +95,12 @@ const TABS: TabDef[] = [
iconComponent: MicIcon,
badgeKey: "common.new",
},
+ {
+ id: "data",
+ labelKey: "settings.tabs.data",
+ icon: DatabaseSettingIcon,
+ badgeKey: "common.new",
+ },
{ id: "about", labelKey: "settings.tabs.about", icon: HelpCircleIcon },
];
@@ -112,6 +120,8 @@ function renderTab(tab: SettingsTab) {
return ;
case "connections":
return ;
+ case "data":
+ return ;
case "api-keys":
return ;
case "about":
@@ -210,6 +220,7 @@ export function SettingsDialog() {
chat: null,
voice: null,
connections: null,
+ data: null,
"api-keys": null,
about: null,
});
diff --git a/studio/frontend/src/features/settings/settings-search.ts b/studio/frontend/src/features/settings/settings-search.ts
index 786ee11d0a..c19b27c732 100644
--- a/studio/frontend/src/features/settings/settings-search.ts
+++ b/studio/frontend/src/features/settings/settings-search.ts
@@ -85,12 +85,19 @@ export const SETTINGS_SEARCH_INDEX: Record = {
"settings.chat.artifacts.title",
"settings.chat.artifacts.collapseHtmlBlocks",
"settings.chat.artifacts.allowNetworkAccess",
- "settings.chat.data",
+ "settings.chat.modelDisclaimer",
+ ],
+ // Chat data management moved to the Data tab; keep these rows findable there.
+ data: [
+ "settings.data.fineTuneExport",
+ "settings.data.archivedChats",
+ "settings.data.archiveAllChats",
+ "settings.data.confirmBeforeDeleting",
+ "settings.data.uploadedFiles",
+ "settings.chat.exportHistory",
"settings.chat.exportConversations",
"settings.chat.importChats",
"settings.chat.clearAllChats",
- "settings.chat.exportHistory",
- "settings.chat.modelDisclaimer",
],
"api-keys": [
"settings.apiKeys.title",
diff --git a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts
index b9e9c75b14..51908a5ad0 100644
--- a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts
+++ b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts
@@ -11,6 +11,7 @@ export type SettingsTab =
| "chat"
| "voice"
| "connections"
+ | "data"
| "api-keys"
| "about";
@@ -30,7 +31,7 @@ interface SettingsDialogState {
// explicitly via onCloseAutoFocus.
opener: HTMLElement | null;
// Set when something asks to jump straight to the archived chats list (the
- // archive toast). ChatTab consumes it to open the dialog, then clears it.
+ // archive toast). DataTab uses it as its initial subpage, then clears it.
archivedChatsRequested: boolean;
openDialog: (tab?: SettingsTab, options?: OpenDialogOptions) => void;
openArchivedChats: () => void;
@@ -66,6 +67,7 @@ function loadInitialTab(): SettingsTab {
"chat",
"voice",
"connections",
+ "data",
"api-keys",
"about",
];
@@ -90,7 +92,7 @@ export const useSettingsDialogStore = create((set) => ({
openArchivedChats: () =>
set({
open: true,
- activeTab: "chat",
+ activeTab: "data",
scrollTarget: null,
archivedChatsRequested: true,
opener: captureOpener(),
diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
index d5c9f10a4f..3e419af78d 100644
--- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
@@ -1,43 +1,16 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import { Button } from "@/components/ui/button";
-import {
- Dialog,
- DialogContent,
- DialogDescription,
- DialogFooter,
- DialogHeader,
- DialogTitle,
-} from "@/components/ui/dialog";
-import {
- DropdownMenu,
- DropdownMenuContent,
- DropdownMenuItem,
- DropdownMenuSeparator,
- DropdownMenuSub,
- DropdownMenuSubContent,
- DropdownMenuSubTrigger,
- DropdownMenuTrigger,
-} from "@/components/ui/dropdown-menu";
import { Switch } from "@/components/ui/switch";
import {
- EXPORT_FORMATS_LIST,
type PlusMenuItemId,
- bulkExportConversationsByScope,
- clearAllChats,
- countAllChats,
- downloadChatExport,
- importConversationsFromFile,
useChatPreferencesStore,
useChatRuntimeStore,
usePlusMenuPrefsStore,
} from "@/features/chat";
import { useT } from "@/i18n";
-import { toast } from "@/lib/toast";
import {
Bookmark02Icon,
- Delete02Icon,
Download01Icon,
FileDatabaseIcon,
Folder01Icon,
@@ -45,19 +18,16 @@ import {
PencilRulerIcon,
Settings02Icon,
ShieldBanIcon,
- Upload01Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Columns2Icon, PlusIcon } from "lucide-react";
-import { useEffect, useRef, useState } from "react";
+import { useEffect } from "react";
import type { ReactNode } from "react";
-import { ArchivedChatsDialog } from "../components/archived-chats-dialog";
import { SettingsRow } from "../components/settings-row";
import {
SettingsGroupDivider,
SettingsSection,
} from "../components/settings-section";
-import { useSettingsDialogStore } from "../stores/settings-dialog-store";
// Adjustable "+" menu items shown in settings, in display order. Icons mirror
// the ones used in the composer + menu itself.
@@ -155,24 +125,6 @@ export function ChatTab() {
const t = useT();
const plusPins = usePlusMenuPrefsStore((state) => state.pins);
const togglePlusPin = usePlusMenuPrefsStore((state) => state.togglePin);
- const [confirmOpen, setConfirmOpen] = useState(false);
- const [archivedOpen, setArchivedOpen] = useState(false);
- const [count, setCount] = useState(null);
- const archivedChatsRequested = useSettingsDialogStore(
- (s) => s.archivedChatsRequested,
- );
- const consumeArchivedChatsRequest = useSettingsDialogStore(
- (s) => s.consumeArchivedChatsRequest,
- );
-
- // Open the archived list when the archive toast asked to jump here.
- useEffect(() => {
- if (!archivedChatsRequested) return;
- setArchivedOpen(true);
- consumeArchivedChatsRequest();
- }, [archivedChatsRequested, consumeArchivedChatsRequest]);
- const [exporting, setExporting] = useState(false);
- const [clearing, setClearing] = useState(false);
const autoTitle = useChatRuntimeStore((state) => state.autoTitle);
const setAutoTitle = useChatRuntimeStore((state) => state.setAutoTitle);
const showCanvasMenuItem = useChatRuntimeStore(
@@ -212,12 +164,6 @@ export function ChatTab() {
const setShowAllQuantizations = useChatRuntimeStore(
(state) => state.setShowAllQuantizations,
);
- const confirmDeleteChats = useChatPreferencesStore(
- (state) => state.confirmDeleteChats,
- );
- const setConfirmDeleteChats = useChatPreferencesStore(
- (state) => state.setConfirmDeleteChats,
- );
const showModelDisclaimer = useChatPreferencesStore(
(state) => state.showModelDisclaimer,
);
@@ -232,95 +178,9 @@ export function ChatTab() {
);
useEffect(() => {
- void countAllChats().then(setCount);
void hydratePersistedSettings();
}, [hydratePersistedSettings]);
- const handleExport = async () => {
- setExporting(true);
- try {
- await downloadChatExport();
- } finally {
- setExporting(false);
- }
- };
-
- const importInputRef = useRef(null);
- const handleImport = async (file: File) => {
- try {
- const imported = await importConversationsFromFile(file, null);
- if (imported === 0) {
- toast.info(t("settings.chat.importNoConversations"));
- } else {
- toast.success(
- imported === 1
- ? t("settings.chat.importedOneChat")
- : t("settings.chat.importedChatCount", { count: imported }),
- );
- setCount(await countAllChats().catch(() => count));
- }
- } catch {
- toast.error(t("settings.chat.importFailed"));
- }
- };
-
- const handleClear = async () => {
- setClearing(true);
- try {
- const result = await clearAllChats();
- const clearedCount = result.deletedThreadIds.length;
- const hasFailedStore =
- result.backend === "failed" || result.legacy === "failed";
- if (!hasFailedStore && result.failedThreadIds.length === 0) {
- setCount(0);
- setConfirmOpen(false);
- toast.success(
- clearedCount === 0
- ? t("settings.chat.clearedAllChats")
- : clearedCount === 1
- ? t("settings.chat.clearedOneChat")
- : t("settings.chat.clearedChatCount", { count: clearedCount }),
- );
- return;
- }
-
- const fallbackRemaining =
- result.failedThreadIds.length > 0
- ? result.failedThreadIds.length
- : (count ?? 0);
- const remaining = await countAllChats().catch(() => fallbackRemaining);
- setCount(remaining);
- setConfirmOpen(false);
- toast.warning(t("settings.chat.someChatsCouldNotBeCleared"), {
- description:
- result.failedThreadIds.length > 0
- ? clearedCount === 1 && result.failedThreadIds.length === 1
- ? t("settings.chat.oneChatClearedRemainOne")
- : clearedCount === 1
- ? t("settings.chat.oneChatClearedRemain", {
- remainingCount: result.failedThreadIds.length,
- })
- : result.failedThreadIds.length === 1
- ? t("settings.chat.chatsClearedRemainOne", { clearedCount })
- : t("settings.chat.chatsClearedRemain", {
- clearedCount,
- remainingCount: result.failedThreadIds.length,
- })
- : remaining === 1
- ? t("settings.chat.storageClearFailedOne")
- : t("settings.chat.storageClearFailed", { count: remaining }),
- });
- } catch (error) {
- const remaining = await countAllChats().catch(() => count);
- setCount(remaining);
- toast.error(t("settings.chat.failedToClearChats"), {
- description: error instanceof Error ? error.message : undefined,
- });
- } finally {
- setClearing(false);
- }
- };
-
return (
@@ -347,7 +207,7 @@ export function ChatTab() {
Q4_K_M
-
+
downloaded
16 GB
@@ -481,191 +341,6 @@ export function ChatTab() {
/>
-
-
-
- setArchivedOpen(true)}
- >
- Manage
-
-
-
-
-
-
-
-
-
-
- {exporting
- ? t("settings.chat.exportingAction")
- : t("settings.chat.exportAction")}
-
-
-
-
-
-
-
-
- {t("settings.chat.exportConversationsAction")}
-
-
-
- {(
- [
- { scope: "recents", label: "exportScopeRecents" },
- { scope: "all", label: "exportScopeAll" },
- ] as const
- ).map(({ scope, label }) => (
-
-
-
- {t(`settings.chat.${label}`)}
-
-
- {EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
-
- void bulkExportConversationsByScope(scope, fmt, true)
- }
- >
- {fmtLabel} {t("settings.chat.exportCombinedSuffix")}
-
- ))}
-
- {EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
-
- void bulkExportConversationsByScope(scope, fmt, false)
- }
- >
- {fmtLabel} {t("settings.chat.exportPerChatSuffix")}
-
- ))}
-
-
- ))}
-
-
-
-
-
- importInputRef.current?.click()}
- >
-
- {t("settings.chat.importChatsAction")}
-
- {
- const file = e.target.files?.[0];
- e.target.value = "";
- if (file) void handleImport(file);
- }}
- />
-
-
-
- setConfirmOpen(true)}
- disabled={count === 0}
- className="text-destructive hover:text-destructive hover:border-destructive/60"
- >
-
- {t("settings.chat.clearChatsAction")}
-
-
-
-
-
-
-
-
-
-
- {count === 1
- ? t("settings.chat.clearOneChatTitle")
- : t("settings.chat.clearChatsTitle", { count: count ?? 0 })}
-
-
- {t("settings.chat.clearChatsConfirmDescription")}
-
-
-
- setConfirmOpen(false)}>
- {t("common.cancel")}
-
-
- {clearing
- ? t("settings.chat.clearingAction")
- : count === 1
- ? t("settings.chat.clearOneChatAction")
- : t("settings.chat.clearChatCountAction", {
- count: count ?? 0,
- })}
-
-
-
-
);
}
diff --git a/studio/frontend/src/features/settings/tabs/data-tab.tsx b/studio/frontend/src/features/settings/tabs/data-tab.tsx
new file mode 100644
index 0000000000..dc9e707ea0
--- /dev/null
+++ b/studio/frontend/src/features/settings/tabs/data-tab.tsx
@@ -0,0 +1,727 @@
+// 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 { Button } from "@/components/ui/button";
+import {
+ Dialog,
+ DialogContent,
+ DialogDescription,
+ DialogFooter,
+ DialogHeader,
+ DialogTitle,
+} from "@/components/ui/dialog";
+import {
+ DropdownMenu,
+ DropdownMenuContent,
+ DropdownMenuItem,
+ DropdownMenuSeparator,
+ DropdownMenuSub,
+ DropdownMenuSubContent,
+ DropdownMenuSubTrigger,
+ DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { Spinner } from "@/components/ui/spinner";
+import { Switch } from "@/components/ui/switch";
+import { usePlatformStore } from "@/config/env";
+import {
+ EXPORT_FORMATS_LIST,
+ type FineTuneFormat,
+ archiveAllChatItems,
+ bulkExportConversationsByScope,
+ clearAllChats,
+ countAllChats,
+ downloadArchivedChatExport,
+ downloadChatExport,
+ exportFineTuneJsonl,
+ importConversationsFromFile,
+ useChatPreferencesStore,
+ useChatRuntimeStore,
+ useChatSidebarItems,
+} from "@/features/chat";
+import { useT } from "@/i18n";
+import {
+ ChevronDownStandardIcon,
+ ChevronRightStandardIcon,
+} from "@/lib/chevron-icons";
+import { toast } from "@/lib/toast";
+import {
+ Archive02Icon,
+ ArrowLeft01Icon,
+ Delete02Icon,
+ Download01Icon,
+ Tick02Icon,
+ Upload01Icon,
+} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { useNavigate, useRouterState } from "@tanstack/react-router";
+import { useEffect, useRef, useState } from "react";
+import { ArchivedChatsView } from "../components/archived-chats-dialog";
+import {
+ createFineTuneRecipeFromChats,
+ loadFineTuneDatasetInTrainTab,
+} from "../components/finetune-recipe";
+import { SettingsRow } from "../components/settings-row";
+import { SettingsSection } from "../components/settings-section";
+import { UploadedFilesView } from "../components/uploaded-files-dialog";
+import { useSettingsDialogStore } from "../stores/settings-dialog-store";
+
+export function DataTab() {
+ const t = useT();
+ const navigate = useNavigate();
+ const archivedChatsRequested = useSettingsDialogStore(
+ (s) => s.archivedChatsRequested,
+ );
+ const consumeArchivedChatsRequest = useSettingsDialogStore(
+ (s) => s.consumeArchivedChatsRequest,
+ );
+ const [confirmOpen, setConfirmOpen] = useState(false);
+ const [archiveConfirmOpen, setArchiveConfirmOpen] = useState(false);
+ // Subpages swap the Data tab body instead of opening nested dialogs.
+ const [subpage, setSubpage] = useState<"main" | "archived" | "files">(
+ archivedChatsRequested ? "archived" : "main",
+ );
+ const [count, setCount] = useState(null);
+ const [exporting, setExporting] = useState(false);
+ const [archivedExporting, setArchivedExporting] = useState(false);
+ // Gates the archived subpage Export button.
+ const { archivedItems } = useChatSidebarItems({ requireMessages: false });
+ const [clearing, setClearing] = useState(false);
+ const [archiving, setArchiving] = useState(false);
+ const [fineTuneExporting, setFineTuneExporting] = useState(false);
+ const [openingRecipe, setOpeningRecipe] = useState(false);
+ const [loadingTraining, setLoadingTraining] = useState(false);
+ // Chat-only hosts redirect /studio back to /chat, so loading a dataset in
+ // the Train tab would upload it and then strand the user; gate the action
+ // the same way the sidebar gates Train.
+ const chatOnly = usePlatformStore((s) => s.isChatOnly());
+ const [fineTuneAction, setFineTuneAction] = useState<
+ "train" | "recipes" | "export"
+ >(chatOnly ? "export" : "train");
+ // Chat Completions (OpenAI messages) is the only export format we ship.
+ const fineTuneFormat: FineTuneFormat = "openai";
+
+ // The MLX self-heal can flip chat-only while the dialog is open.
+ useEffect(() => {
+ if (chatOnly) {
+ setFineTuneAction((a) => (a === "train" ? "export" : a));
+ }
+ }, [chatOnly]);
+ // Requests can arrive after Data is already mounted (for example from the
+ // archive-all toast), so always switch before consuming the flag.
+ useEffect(() => {
+ if (!archivedChatsRequested) return;
+ let cancelled = false;
+ queueMicrotask(() => {
+ if (cancelled) return;
+ setSubpage("archived");
+ consumeArchivedChatsRequest();
+ });
+ return () => {
+ cancelled = true;
+ };
+ }, [archivedChatsRequested, consumeArchivedChatsRequest]);
+
+ const confirmDeleteChats = useChatPreferencesStore(
+ (state) => state.confirmDeleteChats,
+ );
+ const setConfirmDeleteChats = useChatPreferencesStore(
+ (state) => state.setConfirmDeleteChats,
+ );
+
+ const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
+ // Open chat id from the route (single thread or compare pair), mirroring
+ // ArchivedChatsView: compare panes only live in the search params.
+ const openChatId = useRouterState({
+ select: (s) => {
+ if (!s.location.pathname.startsWith("/chat")) return undefined;
+ const search = s.location.search as Record;
+ return search.thread ?? search.compare ?? storeThreadId ?? undefined;
+ },
+ });
+
+ useEffect(() => {
+ void countAllChats().then(setCount);
+ }, []);
+
+ const handleExport = async () => {
+ setExporting(true);
+ try {
+ await downloadChatExport();
+ } finally {
+ setExporting(false);
+ }
+ };
+
+ const handleExportArchived = async () => {
+ setArchivedExporting(true);
+ try {
+ const exported = await downloadArchivedChatExport();
+ toast.success(
+ exported === 0
+ ? t("settings.data.noArchivedChatsToExport")
+ : exported === 1
+ ? t("settings.data.exportedOneArchivedChat")
+ : t("settings.data.exportedArchivedChatCount", { count: exported }),
+ );
+ } catch (error) {
+ toast.error(t("settings.data.failedToExportArchivedChats"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ } finally {
+ setArchivedExporting(false);
+ }
+ };
+
+ const importInputRef = useRef(null);
+ const handleImport = async (file: File) => {
+ try {
+ const imported = await importConversationsFromFile(file, null);
+ if (imported === 0) {
+ toast.info(t("settings.chat.importNoConversations"));
+ } else {
+ toast.success(
+ imported === 1
+ ? t("settings.chat.importedOneChat")
+ : t("settings.chat.importedChatCount", { count: imported }),
+ );
+ setCount(await countAllChats().catch(() => count));
+ }
+ } catch {
+ toast.error(t("settings.chat.importFailed"));
+ }
+ };
+
+ const handleArchiveAll = async () => {
+ setArchiving(true);
+ try {
+ const archived = await archiveAllChatItems(openChatId, (view) => {
+ navigate({ to: "/chat", search: { new: view.newThreadNonce } });
+ });
+ setArchiveConfirmOpen(false);
+ toast.success(
+ archived === 0
+ ? t("settings.data.noChatsToArchive")
+ : archived === 1
+ ? t("settings.data.archivedOneChat")
+ : t("settings.data.archivedChatCount", { count: archived }),
+ );
+ } catch (error) {
+ toast.error(t("settings.data.failedToArchiveChats"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ } finally {
+ setArchiving(false);
+ }
+ };
+
+ const handleFineTuneExport = async () => {
+ setFineTuneExporting(true);
+ try {
+ await exportFineTuneJsonl(fineTuneFormat);
+ } catch (error) {
+ toast.error(t("settings.data.fineTuneExportFailed"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ } finally {
+ setFineTuneExporting(false);
+ }
+ };
+
+ const handleOpenInRecipes = async () => {
+ setOpeningRecipe(true);
+ try {
+ const recipeId = await createFineTuneRecipeFromChats(fineTuneFormat);
+ if (!recipeId) return;
+ useSettingsDialogStore.getState().closeDialog();
+ void navigate({ to: "/data-recipes/$recipeId", params: { recipeId } });
+ } catch (error) {
+ toast.error(t("settings.data.fineTuneRecipeFailed"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ } finally {
+ setOpeningRecipe(false);
+ }
+ };
+
+ const handleUseInTraining = async () => {
+ setLoadingTraining(true);
+ try {
+ const loaded = await loadFineTuneDatasetInTrainTab(fineTuneFormat);
+ if (!loaded) return;
+ useSettingsDialogStore.getState().closeDialog();
+ void navigate({ to: "/studio" });
+ } catch (error) {
+ toast.error(t("settings.data.fineTuneTrainFailed"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ } finally {
+ setLoadingTraining(false);
+ }
+ };
+
+ const fineTuneActionLabels = {
+ train: t("settings.data.fineTuneTrainAction"),
+ recipes: t("settings.data.fineTuneOpenRecipesAction"),
+ export: t("settings.data.fineTuneExportAction"),
+ } as const;
+ const fineTuneBusy = loadingTraining || openingRecipe || fineTuneExporting;
+ const runFineTuneAction = () => {
+ if (fineTuneAction === "train") {
+ if (chatOnly) return;
+ void handleUseInTraining();
+ } else if (fineTuneAction === "recipes") void handleOpenInRecipes();
+ else void handleFineTuneExport();
+ };
+
+ const handleClear = async () => {
+ setClearing(true);
+ try {
+ const result = await clearAllChats();
+ const clearedCount = result.deletedThreadIds.length;
+ const hasFailedStore =
+ result.backend === "failed" || result.legacy === "failed";
+ if (!hasFailedStore && result.failedThreadIds.length === 0) {
+ setCount(0);
+ setConfirmOpen(false);
+ toast.success(
+ clearedCount === 0
+ ? t("settings.chat.clearedAllChats")
+ : clearedCount === 1
+ ? t("settings.chat.clearedOneChat")
+ : t("settings.chat.clearedChatCount", { count: clearedCount }),
+ );
+ return;
+ }
+
+ const fallbackRemaining =
+ result.failedThreadIds.length > 0
+ ? result.failedThreadIds.length
+ : (count ?? 0);
+ const remaining = await countAllChats().catch(() => fallbackRemaining);
+ setCount(remaining);
+ setConfirmOpen(false);
+ toast.warning(t("settings.chat.someChatsCouldNotBeCleared"), {
+ description:
+ result.failedThreadIds.length > 0
+ ? clearedCount === 1 && result.failedThreadIds.length === 1
+ ? t("settings.chat.oneChatClearedRemainOne")
+ : clearedCount === 1
+ ? t("settings.chat.oneChatClearedRemain", {
+ remainingCount: result.failedThreadIds.length,
+ })
+ : result.failedThreadIds.length === 1
+ ? t("settings.chat.chatsClearedRemainOne", { clearedCount })
+ : t("settings.chat.chatsClearedRemain", {
+ clearedCount,
+ remainingCount: result.failedThreadIds.length,
+ })
+ : remaining === 1
+ ? t("settings.chat.storageClearFailedOne")
+ : t("settings.chat.storageClearFailed", { count: remaining }),
+ });
+ } catch (error) {
+ const remaining = await countAllChats().catch(() => count);
+ setCount(remaining);
+ toast.error(t("settings.chat.failedToClearChats"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ } finally {
+ setClearing(false);
+ }
+ };
+
+ if (subpage === "archived") {
+ return (
+
+
+ setSubpage("main")}
+ aria-label={`Back to ${t("settings.data.title")}`}
+ className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
+ >
+
+
+
+ {t("settings.data.title")}
+
+
+
+
+
+ {t("settings.data.archivedChats")}
+
+
+ {t("settings.data.archivedChatsDescription")}
+
+
+ {archivedItems.length > 0 && (
+
+ {archivedExporting ? (
+
+ ) : (
+
+ )}
+ {archivedExporting
+ ? t("settings.data.exportingArchivedChats")
+ : t("settings.data.exportArchivedChats")}
+
+ )}
+
+
+
+ );
+ }
+
+ if (subpage === "files") {
+ return (
+
+
+ setSubpage("main")}
+ aria-label={`Back to ${t("settings.data.title")}`}
+ className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
+ >
+
+
+
+ {t("settings.data.title")}
+
+
+
+
+ {t("settings.data.uploadedFiles")}
+
+
+ {t("settings.data.uploadedFilesDescription")}
+
+
+
+
+ );
+ }
+
+ return (
+
+
+
+
+
+
+
+
+ {/* Fixed width so switching actions never resizes the row. */}
+
+
+ {fineTuneActionLabels[fineTuneAction]}
+
+
+
+
+
+ {(["export", "train", "recipes"] as const).map((action) => (
+ setFineTuneAction(action)}
+ >
+
+ {fineTuneActionLabels[action]}
+
+ {fineTuneAction === action ? (
+
+ ) : null}
+
+ ))}
+
+
+
+ {fineTuneBusy ? (
+
+ ) : (
+
+ )}
+
+
+
+
+
+ setSubpage("archived")}
+ >
+ {t("settings.data.manageAction")}
+
+
+
+
+ setArchiveConfirmOpen(true)}
+ >
+
+ {t("settings.data.archiveAllAction")}
+
+
+
+
+
+
+
+
+
+
+ {exporting
+ ? t("settings.chat.exportingAction")
+ : t("settings.chat.exportAction")}
+
+
+
+
+
+
+
+
+ {t("settings.chat.exportConversationsAction")}
+
+
+
+ {(
+ [
+ { scope: "recents", label: "exportScopeRecents" },
+ { scope: "all", label: "exportScopeAll" },
+ ] as const
+ ).map(({ scope, label }) => (
+
+
+
+ {t(`settings.chat.${label}`)}
+
+
+ {EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
+
+ void bulkExportConversationsByScope(scope, fmt, true)
+ }
+ >
+ {fmtLabel} {t("settings.chat.exportCombinedSuffix")}
+
+ ))}
+
+ {EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
+
+ void bulkExportConversationsByScope(scope, fmt, false)
+ }
+ >
+ {fmtLabel} {t("settings.chat.exportPerChatSuffix")}
+
+ ))}
+
+
+ ))}
+
+
+
+
+
+ setConfirmOpen(true)}
+ disabled={count === 0}
+ className="text-destructive hover:text-destructive hover:border-destructive/60"
+ >
+
+ {t("settings.chat.clearChatsAction")}
+
+
+
+
+ importInputRef.current?.click()}
+ >
+
+ {t("settings.chat.importChatsAction")}
+
+ {
+ const file = e.target.files?.[0];
+ e.target.value = "";
+ if (file) void handleImport(file);
+ }}
+ />
+
+
+
+
+
+ setSubpage("files")}
+ >
+ {t("settings.data.manageAction")}
+
+
+
+
+
+
+
+ {t("settings.data.archiveAllChatsTitle")}
+
+ {t("settings.data.archiveAllChatsConfirmDescription")}
+
+
+
+ setArchiveConfirmOpen(false)}
+ >
+ {t("common.cancel")}
+
+
+ {archiving
+ ? t("settings.data.archivingAction")
+ : t("settings.data.archiveAllAction")}
+
+
+
+
+
+
+
+
+
+ {count === 1
+ ? t("settings.chat.clearOneChatTitle")
+ : t("settings.chat.clearChatsTitle", { count: count ?? 0 })}
+
+
+ {t("settings.chat.clearChatsConfirmDescription")}
+
+
+
+ setConfirmOpen(false)}>
+ {t("common.cancel")}
+
+
+ {clearing
+ ? t("settings.chat.clearingAction")
+ : count === 1
+ ? t("settings.chat.clearOneChatAction")
+ : t("settings.chat.clearChatCountAction", {
+ count: count ?? 0,
+ })}
+
+
+
+
+
+ );
+}
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index cbddc9f0c2..5ee2805a33 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -99,6 +99,7 @@ export const en = {
chat: "Chat",
voice: "Voice",
connections: "Connections",
+ data: "Data",
apiKeys: "API",
about: "About",
},
@@ -511,7 +512,7 @@ export const en = {
},
chat: {
title: "Chat",
- description: "Manage chat history stored on this device.",
+ description: "Customize how chat behaves on this device.",
modelDisclaimer: "Show model disclaimer",
modelDisclaimerDescription:
'Show "LLMs can make mistakes" under the chat box.',
@@ -580,6 +581,53 @@ export const en = {
"A storage clear failed; {count} chats may remain. Please retry.",
failedToClearChats: "Failed to clear chats",
},
+ data: {
+ title: "Data",
+ description:
+ "Manage chat history and uploaded files stored on this device.",
+ archivedChats: "Archived chats",
+ archivedChatsDescription: "View and manage chats you have archived.",
+ manageAction: "Manage",
+ exportArchivedChats: "Export",
+ exportingArchivedChats: "Exporting...",
+ exportedOneArchivedChat: "Exported 1 archived chat",
+ exportedArchivedChatCount: "Exported {count} archived chats",
+ noArchivedChatsToExport: "No archived chats to export.",
+ failedToExportArchivedChats: "Failed to export archived chats",
+ archiveAllChats: "Archive all chats",
+ archiveAllChatsDescription:
+ "Move every chat in Recents and Projects to the archive.",
+ noChatsToArchive: "No chats to archive.",
+ archiveAllAction: "Archive all",
+ archivingAction: "Archiving...",
+ archiveAllChatsTitle: "Archive all chats?",
+ archiveAllChatsConfirmDescription:
+ "Moves every chat on this device to the archive. Archived chats stay available and can be unarchived at any time.",
+ archivedAllChats: "Archived all chats",
+ archivedOneChat: "Archived 1 chat",
+ archivedChatCount: "Archived {count} chats",
+ failedToArchiveChats: "Failed to archive chats",
+ confirmBeforeDeleting: "Confirm before deleting",
+ confirmBeforeDeletingDescription:
+ "Ask for confirmation before a chat is deleted. Turn off to delete instantly.",
+ filesSection: "Files",
+ uploadedFiles: "Uploaded files",
+ uploadedFilesDescription:
+ "View and manage files uploaded to chats, projects, and knowledge bases.",
+ fineTuneExport: "Use chats as training data",
+ fineTuneExportDescription:
+ "Create a fine-tuning JSONL dataset from your chats. Load it in Train, refine in Recipes, or export it.",
+ fineTuneExportAction: "Export JSONL",
+ fineTuneRunAction: "Run",
+ fineTuneExportingAction: "Exporting...",
+ fineTuneOpenRecipesAction: "Open in Recipes",
+ fineTuneOpeningRecipesAction: "Opening...",
+ fineTuneTrainAction: "Load in Train tab",
+ fineTuneTrainingAction: "Loading...",
+ fineTuneExportFailed: "Failed to export training data",
+ fineTuneRecipeFailed: "Failed to open chats in Recipes",
+ fineTuneTrainFailed: "Failed to load dataset in the Train tab",
+ },
connections: {
title: "Connections",
description: "Manage providers and external connections.",
From 66808ab25dcb7655dab05a0837a0ec0bb6cf080b Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 05:27:53 -0700
Subject: [PATCH 031/255] Studio: fix per-GPU VRAM reporting on Windows ROCm
(#7238)
* Studio: fix per-GPU VRAM reporting on Windows ROCm
On Windows ROCm without a HIP SDK, amd-smi is disabled and the System tab fell
back to torch mem_get_info, which reports free==total there (ROCm/ROCm#1909), so
used VRAM showed as 0. The perf-counter fallback also summed every adapter into a
single device with only GPU 0's total, hiding the second GPU.
Read per-adapter Dedicated Usage (LUID-instanced) for used and take each GPU's
total from torch properties, and treat the free==total case as unknown rather
than 0, so every GPU shows real usage. NVIDIA, Linux ROCm, Apple and CPU paths
are unchanged. Final validation needs a real Windows AMD box.
Fixes #7072
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: report unknown VRAM instead of fabricating or zeroing it
Two gaps in the Windows ROCm VRAM path. When more adapters are actively using
VRAM than are visible to the process (a GPU outside the visibility mask), the
per-adapter attribution paired usage by size and fabricated a per-GPU value;
report unknown for every device in that case rather than mis-assign. And the
System API turned an unknown (None) used value into 0 with ``or 0``, then
reported the full card as free, re-hiding the exact case this change surfaces;
keep None so the UI shows unknown.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Render unknown VRAM as Unknown instead of zero in the System tab
The backend reports null usage when it is unknown (e.g. the Windows ROCm
perf counter is unavailable or localized), but the System tab coerced
null to 0 and derived free from it, fabricating a 0-used/full-free total.
Preserve null and render the translated Unknown for per-device used, free
and utilization, and mark the aggregate VRAM tile unknown when any device
is unknown.
* Render unknown VRAM as Unknown in the floating monitor and the util tile
The floating VRAM monitor and the aggregate utilization ring both still
coerced a null usage to 0, showing a fabricated 0.00 GiB / full free / 0%
on the same Windows ROCm no-counter case the resources tab already
handles. Guard both on whether every device reports a finite usage and
render Unknown (value and percent) instead of a concrete 0.
* Attribute per-adapter VRAM usage only when capacity forces the mapping
On Windows/ROCm there is no shared key between LUID performance-counter
instances and torch ordinals, so usage was paired to devices purely by capacity
ranking. That pairing is only trustworthy when capacity forces it (a usage
larger than every smaller device can sit on one card). When a smaller-capacity
device could equally hold a strictly larger usage (for example an 8 GiB card
near full beside a lightly used 48 GiB card), the two values are swappable
without violating any capacity, so the ranking is a guess with no key to break
the tie. A wrong guess both mislabels the System tab and feeds
routes/training_vram.py a wrong per-index free value, driving a wrong
keep-resident decision.
Report unknown for every device when the assignment is ambiguous, keeping the
attribution only for the capacity-forced case. Returning None is the
conservative direction: training_vram treats a missing index as zero free, so it
never keeps a chat model into an OOM. Add regression tests for the
not-capacity-ordered, same-capacity, single-fits-both, and capacity-forced
cases.
* Report unknown VRAM usage when a hidden adapter survives the noise filter
When HIP_VISIBLE_DEVICES exposes a subset of the physical adapters, the LUID
usage counters cover cards outside the visibility mask too. The sub-64 MiB noise
filter could drop a genuinely-idle visible card's real usage while keeping a
hidden larger card's high usage, which was then clamped onto the smaller visible
device and reported as fully used (for example a hidden 48 GiB card at 40 GiB
shown as a visible 8 GiB card fully used, with its true 10 MiB usage filtered
out). That fabricated reading also feeds routes/training_vram.py a wrong
per-index free value.
Flag extra adapters on the raw counter count (before the noise filter, since an
idle visible card can itself fall below the floor) and, when a kept usage exceeds
its ranked visible capacity, report unknown rather than clamp a hidden card's
usage onto a visible device. The genuinely-idle-noise and capacity-forced
single-model cases are unchanged. Add a regression test for the hidden
high-use-adapter case in both counter orders.
* Report unknown when only a placeholder adapter counter survives the noise filter
When more raw counters than visible devices are present but every counter sits
below the 64 MiB noise floor (an idle real GPU alongside a Windows Basic Render
Driver placeholder), the non_trivial-or-raw fallback resurrected the raw
magnitude-sorted counters and could attribute the placeholder to a real GPU while
dropping a real card's reading. With a single visible device the swap-ambiguity
check cannot catch it (it needs at least two ranks), so the fabricated value
reached the System tab and automatic GPU selection.
Return unknown for every device in that case instead of falling back to raw
counters. With the earlier guards this completes the invariant: a concrete
per-GPU usage is emitted only when the assignment is capacity-forced, and every
ambiguous, extra-adapter, placeholder-fallback, or count-mismatch path reports
unknown. Add a regression test for the placeholder fallback in both counter
orders and the two-idle-GPU case.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: attribute Windows/ROCm VRAM only when capacity forces a clean bijection
With more raw adapter counters than visible devices, a survivor that merely
fits a visible card was pinned to it by magnitude ranking, fabricating a hidden
GPU's usage onto an idle visible card whose true reading was dropped by the
sub-threshold noise filter (two visible 48/8 GiB cards using 40 GiB / 10 MiB
beside a hidden 6 GiB adapter returned [40, 6]). Emit a concrete per-device
value only when the supra-threshold counters number exactly the visible devices
(every visible card has one real reading, the extras were sub-threshold
placeholders) AND the ranked usage strictly exceeds every smaller visible card's
capacity. When a visible card is idle (fewer supra-threshold counters than
devices) a survivor could be the hidden GPU's usage, so every device reports
unknown; more active counters than visible cards, the smallest card, and any
merely-fitting usage stay unknown too. The reporter's loaded-card display is
preserved (40 GiB / 0.5 GiB across 48/8 GiB -> [40, None]). Adds a regression
test for the reported case plus an exhaustive capacity-forced/bijection matrix.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: keep the unified-memory total when Windows-ROCm used is unknown
_apply_unified_memory_correction gated both the total and the used update on
torch_used_gb being known, so on a unified-memory APU (Strix Halo) where torch
reports used=None (the Windows-ROCm free==total sentinel) but an authoritative
full-GTT total, the device kept amd-smi's small dedicated carve-out and
underreported its capacity on the System tab. Adopt torch's larger total
independently of used; overwrite used only when torch's is known (otherwise keep
amd-smi's dedicated-usage figure) and recompute utilization against the
corrected total. Adds regression tests.
* Tighten comments in the ROCm/Windows VRAM reporting path
* Tighten comments further in the ROCm/Windows VRAM reporting path
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen
---
studio/backend/main.py | 8 +-
.../tests/test_rocm_windows_vram_7072.py | 361 ++++++++++++++++++
studio/backend/utils/hardware/hardware.py | 308 ++++++++++++---
.../src/components/floating-monitor.tsx | 26 +-
.../features/settings/tabs/resources-tab.tsx | 99 +++--
5 files changed, 716 insertions(+), 86 deletions(-)
create mode 100644 studio/backend/tests/test_rocm_windows_vram_7072.py
diff --git a/studio/backend/main.py b/studio/backend/main.py
index a1ff4d60da..0ffc4489a1 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -1149,11 +1149,15 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]:
util = util_devices.get(idx, {})
total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0
- used_vram = util.get("vram_used_gb") or 0
+ # Keep None (usage unknown, e.g. Windows ROCm perf counter) so the UI
+ # shows unknown, not a fabricated 0 used / full free.
+ used_vram = util.get("vram_used_gb")
enriched_dev = dict(dev)
enriched_dev["vram_used_gb"] = used_vram
- enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0
+ enriched_dev["vram_free_gb"] = (
+ round(total_vram - used_vram, 2) if total_vram and used_vram is not None else None
+ )
enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct")
enriched_devices.append(enriched_dev)
diff --git a/studio/backend/tests/test_rocm_windows_vram_7072.py b/studio/backend/tests/test_rocm_windows_vram_7072.py
new file mode 100644
index 0000000000..b4079831b7
--- /dev/null
+++ b/studio/backend/tests/test_rocm_windows_vram_7072.py
@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Regression tests for issue #7072 -- "VRAM Usage in System Tab is wrong".
+
+Reporter: dual AMD (Radeon PRO W7900 ~48GB + W7500 8GB), Windows 10, ROCm 7.13,
+torch 2.11.0+rocm7.13. On Windows without a HIP SDK, amd-smi is permanently
+disabled (avoids a UAC/DiskPart prompt) and hipMemGetInfo returns free==total
+(used 0). Two symptoms followed:
+
+ * System tab (/api/system -> get_visible_gpu_utilization) showed ~0 VRAM used
+ on every GPU (torch mem_get_info free==total quirk; ROCm/ROCm#1909).
+ * get_gpu_utilization()'s Windows fallback SUMMED "GPU Adapter Memory\\Dedicated
+ Usage" across all adapters into ONE fake device with only GPU 0's total, so
+ the second GPU never appeared.
+
+The fix reads the per-adapter (LUID-instanced) Dedicated Usage performance
+counter -- Task Manager's source -- for per-GPU used, takes per-GPU total from
+torch device properties, and guards the free==total mem_get_info quirk. CI has no
+AMD GPU/Windows, so torch, the performance counter, and platform are all mocked.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import types
+
+import pytest
+
+from utils.hardware import hardware as hw
+
+GB = 1024**3
+MiB = 1024**2
+
+
+# ----------------------------------------------------------------------------- #
+# Fakes
+# ----------------------------------------------------------------------------- #
+def _fake_torch(
+ devices,
+ *,
+ free_equals_total = False,
+ used_per_device = None,
+):
+ """Build a fake `torch` module. devices: list of (name, total_bytes)."""
+ dev = list(devices)
+
+ class _Props:
+ def __init__(self, name, total):
+ self.name = name
+ self.total_memory = total
+
+ def get_device_properties(i):
+ name, total = dev[i]
+ return _Props(name, total)
+
+ def mem_get_info(i):
+ _, total = dev[i]
+ if free_equals_total:
+ return (total, total)
+ used = used_per_device[i] if used_per_device is not None else 0
+ return (total - used, total)
+
+ t = types.ModuleType("torch")
+ t.__version__ = "2.11.0+rocm7.13"
+ t.version = types.SimpleNamespace(hip = "7.13", cuda = None)
+ t.cuda = types.SimpleNamespace(
+ is_available = lambda: len(dev) > 0,
+ device_count = lambda: len(dev),
+ current_device = lambda: 0,
+ get_device_properties = get_device_properties,
+ mem_get_info = mem_get_info,
+ memory_allocated = lambda i: 0,
+ memory_reserved = lambda i: 0,
+ )
+ return t
+
+
+def _adapter_output(adapters):
+ if not adapters:
+ return "__NONE__\n"
+ return "".join(f"{name}|{int(used)}\n" for name, used in adapters)
+
+
+def _subprocess_run(*, adapter_output = "__NONE__\n", util_output = "12.0\n"):
+ def fake_run(cmd, *a, **k):
+ joined = " ".join(cmd) if isinstance(cmd, list) else str(cmd)
+ if "GPU Adapter Memory" in joined and "InstanceName" in joined:
+ out = adapter_output
+ elif "engtype_3D" in joined or "GPU Engine" in joined:
+ out = util_output
+ else:
+ out = "-1\n"
+ return subprocess.CompletedProcess(args = cmd, returncode = 0, stdout = out, stderr = "")
+
+ return fake_run
+
+
+@pytest.fixture
+def win_rocm(monkeypatch):
+ """Configure the hardware module as a Windows ROCm host with 2 visible GPUs."""
+ monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
+ monkeypatch.setattr(hw, "IS_ROCM", True)
+ monkeypatch.setattr(hw.platform, "system", lambda: "Windows")
+ monkeypatch.setattr(hw.sys, "platform", "win32")
+ monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) # amd-smi disabled
+ # Visible set via HIP mask so we don't shell out to amd-smi for the count.
+ monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1")
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
+ monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
+ return monkeypatch
+
+
+REPORTER_ADAPTERS = [
+ ("luid_0x00000000_0x0000d1e2_phys_0", 40.0 * GB), # W7900, model loaded
+ ("luid_0x00000000_0x0000e34a_phys_0", 0.5 * GB), # W7500, idle
+ ("luid_0x00000000_0x0000f001_phys_0", 3 * MiB), # Basic Render Driver
+]
+DEVICES = [("AMD Radeon PRO W7900", 48 * GB), ("AMD Radeon PRO W7500", 8 * GB)]
+
+
+# ----------------------------------------------------------------------------- #
+# System tab (get_visible_gpu_utilization) -- the reporter's screenshot
+# ----------------------------------------------------------------------------- #
+def test_system_tab_shows_per_gpu_used(win_rocm, monkeypatch):
+ monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True))
+ monkeypatch.setattr(
+ hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS))
+ )
+
+ devices = hw.get_visible_gpu_utilization()["devices"]
+ by_idx = {d["index"]: d for d in devices}
+ assert len(devices) == 2
+ assert by_idx[0]["vram_total_gb"] == 48.0
+ assert by_idx[0]["vram_used_gb"] == pytest.approx(40.0, abs = 0.01) # not 0
+ assert by_idx[1]["vram_total_gb"] == 8.0 # own total
+ # The 3 MiB Basic Render Driver counter makes this a hidden-adapter case: only
+ # the 40 GiB is forced onto the 48 GiB card; the idle card reads Unknown.
+ assert by_idx[1]["vram_used_gb"] is None
+ assert by_idx[1]["vram_utilization_pct"] is None
+ assert all(
+ d["vram_used_gb"] <= d["vram_total_gb"] for d in devices if d["vram_used_gb"] is not None
+ )
+
+
+def test_gpu_utilization_does_not_collapse(win_rocm, monkeypatch):
+ monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True))
+ monkeypatch.setattr(
+ hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS))
+ )
+
+ result = hw.get_gpu_utilization()
+ devices = result["devices"]
+ assert sorted(d["index"] for d in devices) == [0, 1] # both GPUs, no collapse
+ assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0}
+ assert result["vram_total_gb"] == 48.0 # legacy primary mirror preserved
+
+
+def test_localized_counter_reports_unknown_not_zero(win_rocm, monkeypatch):
+ monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True))
+ monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n"))
+
+ devices = hw.get_visible_gpu_utilization()["devices"]
+ assert len(devices) == 2 # both still shown with correct totals
+ assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0}
+ assert all(d["vram_used_gb"] is None for d in devices) # unknown, not fake 0
+ assert all(d["vram_utilization_pct"] is None for d in devices)
+
+
+# ----------------------------------------------------------------------------- #
+# mem_get_info free==total guard scoping
+# ----------------------------------------------------------------------------- #
+def test_mem_get_info_guard_scopes_to_windows_rocm(monkeypatch):
+ torch_mod = _fake_torch(DEVICES, free_equals_total = True)
+ monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
+ monkeypatch.setitem(sys.modules, "torch", torch_mod)
+
+ # Windows ROCm -> used unknown (None), total kept.
+ monkeypatch.setattr(hw, "IS_ROCM", True)
+ monkeypatch.setattr(hw.sys, "platform", "win32")
+ win = hw._torch_get_per_device_info([0, 1])
+ assert [d["used_gb"] for d in win] == [None, None]
+ assert [d["total_gb"] for d in win] == [48.0, 8.0]
+
+ # Linux ROCm -> unchanged numeric used.
+ monkeypatch.setattr(hw.sys, "platform", "linux")
+ assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0]
+
+ # Windows NVIDIA -> guard must not fire.
+ monkeypatch.setattr(hw, "IS_ROCM", False)
+ monkeypatch.setattr(hw.sys, "platform", "win32")
+ assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0]
+
+
+# ----------------------------------------------------------------------------- #
+# Per-adapter attribution helpers (pure unit)
+# ----------------------------------------------------------------------------- #
+def test_match_adapter_pairs_and_clamps():
+ assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [48 * GB, 8 * GB]) == [
+ 40 * GB,
+ 0.5 * GB,
+ ]
+ assert hw._match_adapter_used_to_devices([100 * GB], [48 * GB]) == [48 * GB] # clamp
+ assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None]
+
+
+def test_match_adapter_reports_unknown_when_more_active_than_visible():
+ # More adapters actively using VRAM than are visible (a GPU outside the mask):
+ # attribution would fabricate a value, so report unknown for every device.
+ assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [8 * GB]) == [None]
+
+
+def test_match_adapter_reports_unknown_when_hidden_high_use_adapter_survives_filter():
+ # Idle 8 GiB card (10 MiB noise) beside a hidden 48 GiB card at 40 GiB: the
+ # 40 GiB can't fit the 8 GiB device, so clamping there would fabricate. Unknown.
+ assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB], [8 * GB]) == [None]
+ # Order of the counters must not matter.
+ assert hw._match_adapter_used_to_devices([10 * MiB, 40 * GB], [8 * GB]) == [None]
+
+
+def test_match_adapter_reports_unknown_for_placeholder_fallback():
+ # Every counter below the 64 MiB floor plus a placeholder: no LUID-to-ordinal
+ # mapping tells placeholder from idle GPU, so report unknown, not fabricate.
+ # Single visible 8 GiB card idle (10 MiB) beside a 50 MiB placeholder counter.
+ assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB], [8 * GB]) == [None]
+ # Order of the counters must not matter.
+ assert hw._match_adapter_used_to_devices([10 * MiB, 50 * MiB], [8 * GB]) == [None]
+ # Two idle visible GPUs plus a placeholder: all three counters below the floor.
+ assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [
+ None,
+ None,
+ ]
+
+
+def test_match_adapter_reports_unknown_when_usage_not_capacity_ordered():
+ # 8 GiB card at 7 GiB beside a 48 GiB card at 5 GiB: the bigger usage still fits
+ # the smaller card, so both pairings are feasible -> unknown.
+ assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [8 * GB, 48 * GB]) == [None, None]
+ # Device order must not matter (same physical situation, ordinals flipped).
+ assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [48 * GB, 8 * GB]) == [None, None]
+ # Same-capacity cards with unequal usage are equally unattributable.
+ assert hw._match_adapter_used_to_devices([12 * GB, 8 * GB], [24 * GB, 24 * GB]) == [None, None]
+ # A single usage that fits both cards can sit on either -> unknown.
+ assert hw._match_adapter_used_to_devices([5 * GB], [48 * GB, 8 * GB]) == [None, None]
+ # But a capacity-forced assignment (usage exceeds the smaller card) is kept:
+ # 40 GiB can only be the 48 GiB card, so it is not fabrication.
+ assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None]
+
+
+def test_match_adapter_reports_unknown_when_hidden_usage_fits_visible_card():
+ # A survivor that merely *fits* a visible card must not be pinned onto it. Two
+ # cards (48/8 GiB) at 40 GiB / 10 MiB beside a hidden 6 GiB adapter: the 6 GiB
+ # fits the idle 8 GiB card but isn't forced -> Unknown; only 40 GiB is forced.
+ assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB, 6 * GB], [48 * GB, 8 * GB]) == [
+ 40 * GB,
+ None,
+ ]
+ # Counter order must not matter.
+ assert hw._match_adapter_used_to_devices([6 * GB, 40 * GB, 10 * MiB], [48 * GB, 8 * GB]) == [
+ 40 * GB,
+ None,
+ ]
+ # A single visible card with a hidden adapter is never attributable: a fitting
+ # survivor could be the hidden GPU's while the visible card is idle.
+ assert hw._match_adapter_used_to_devices([6 * GB, 10 * MiB], [8 * GB]) == [None]
+
+
+def test_match_adapter_capacity_forced_matrix():
+ """Exhaustive hidden-adapter matrix for the capacity-forced rule.
+
+ A value is emitted only when the supra-threshold counters number exactly the
+ visible devices AND a device's ranked usage strictly exceeds every smaller
+ card's capacity. Otherwise (a visible card idle, a merely-fitting usage, or the
+ smallest card) every device reports unknown.
+ """
+ m = hw._match_adapter_used_to_devices
+ # -- exactly-n supra-threshold counters, capacity-forced survivors are kept - #
+ # Both visible cards have a real reading (the 3 MiB is a placeholder): 40 GiB
+ # forced onto the 48 GiB card, 0.5 GiB not forced -> None.
+ assert m([40 * GB, 0.5 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [40 * GB, None]
+ # Three visible cards all active (supra-threshold) + placeholder: 40 > 24 and
+ # 20 > 8, both forced; the 8 GiB card is not forced -> None.
+ assert m([40 * GB, 20 * GB, 5 * GB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [
+ 40 * GB,
+ 20 * GB,
+ None,
+ ]
+ # -- fewer supra-threshold counters than visible cards -> all unknown ------ #
+ # A visible card is idle, so even a "forced" 40 could be the hidden GPU's.
+ assert m([40 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None]
+ assert m([40 * GB, 10 * MiB, 10 * MiB], [48 * GB, 8 * GB]) == [None, None]
+ assert m([40 * GB, 20 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [
+ None,
+ None,
+ None,
+ ]
+ # Middle usage (6 GiB) fits both the 24 and 8 GiB cards, and only two cards are
+ # active for three visible -> not a bijection -> all unknown.
+ assert m([40 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [
+ None,
+ None,
+ None,
+ ]
+ # -- hidden larger than every visible card -> all unknown ----------------- #
+ assert m([40 * GB, 10 * MiB], [8 * GB]) == [None]
+ assert m([48 * GB, 3 * MiB, 3 * MiB], [24 * GB, 8 * GB]) == [None, None]
+ # -- more active adapters than visible cards -> all unknown --------------- #
+ assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None]
+ assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None]
+ # -- every counter below the noise floor (placeholder fallback) -> unknown - #
+ assert m([50 * MiB, 10 * MiB], [8 * GB]) == [None]
+ assert m([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [None, None]
+ # -- equal-capacity cards with a hidden adapter: nothing is forced -------- #
+ assert m([40 * GB, 40 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None]
+ assert m([40 * GB, 30 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None]
+
+
+def test_perf_counter_parser_and_sentinel(monkeypatch):
+ monkeypatch.setattr(hw.platform, "system", lambda: "Windows")
+ monkeypatch.setattr(
+ hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS))
+ )
+ parsed = hw._rocm_windows_perf_counter_vram_by_adapter()
+ assert parsed is not None and len(parsed) == 3
+ assert parsed[0][0].startswith("luid_")
+ monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n"))
+ assert hw._rocm_windows_perf_counter_vram_by_adapter() is None
+
+
+# ----------------------------------------------------------------------------- #
+# Unified-memory (Strix Halo APU) total reconciliation (Codex #7238)
+# ----------------------------------------------------------------------------- #
+def test_unified_memory_adopts_torch_total_even_when_used_unknown():
+ """Windows ROCm unified-memory APU: torch's used is None but its total (the full
+ GTT pool) is authoritative. The correction must still adopt the larger total;
+ used stays at amd-smi's figure when torch's is unknown."""
+ metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0}
+ hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": None, "index": 0})
+ assert metrics["vram_total_gb"] == 124.0 # full unified pool, not the 8 GB carve-out
+ assert metrics["vram_used_gb"] == 2.0 # amd-smi used preserved (torch's was None)
+ assert metrics["vram_utilization_pct"] == pytest.approx(round(2.0 / 124.0 * 100, 1))
+
+
+def test_unified_memory_overwrites_used_when_torch_used_known():
+ """When torch reports both a larger total and a known used, both are adopted
+ and utilization is recomputed against the corrected total (unchanged path)."""
+ metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0}
+ hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": 40.0, "index": 0})
+ assert metrics["vram_total_gb"] == 124.0
+ assert metrics["vram_used_gb"] == 40.0
+ assert metrics["vram_utilization_pct"] == pytest.approx(round(40.0 / 124.0 * 100, 1))
+
+
+def test_unified_memory_no_op_when_torch_total_not_larger():
+ """A discrete GPU where torch total does not exceed amd-smi's is left untouched."""
+ metrics = {"vram_total_gb": 48.0, "vram_used_gb": 10.0, "vram_utilization_pct": 20.8}
+ hw._apply_unified_memory_correction(metrics, {"total_gb": 48.0, "used_gb": None, "index": 0})
+ assert metrics["vram_total_gb"] == 48.0
+ assert metrics["vram_used_gb"] == 10.0
+ assert metrics["vram_utilization_pct"] == 20.8
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index adc9a54aab..9fef53e65e 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -538,21 +538,31 @@ def _torch_get_physical_gpu_count() -> Optional[int]:
def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]]:
- """Query torch for per-GPU name, total VRAM, and used VRAM."""
+ """Query torch for per-GPU name, total VRAM, and used VRAM.
+
+ ``used_gb`` is ``None`` on Windows ROCm when ``hipMemGetInfo`` reports
+ ``free == total`` (ROCm/ROCm#1909): that 0 means unknown, not empty.
+ """
mod, _ = _torch_get_device_module()
if mod is None:
return []
+ # free==total is a Windows-ROCm-only quirk.
+ _win_rocm = sys.platform == "win32" and IS_ROCM
devices = []
for ordinal, phys_idx in enumerate(device_indices):
try:
# torch ordinals are 0-based relative to CUDA_VISIBLE_DEVICES.
props = mod.get_device_properties(ordinal)
total_bytes = props.total_memory
+ used_bytes: Optional[int]
# Prefer mem_get_info (system-wide) so auto-select sees other consumers.
if hasattr(mod, "mem_get_info"):
free_bytes, total_bytes = mod.mem_get_info(ordinal)
used_bytes = total_bytes - free_bytes
+ # free==total is the broken-API sentinel, not an idle GPU.
+ if _win_rocm and free_bytes == total_bytes:
+ used_bytes = None
else:
used_bytes = mod.memory_allocated(ordinal)
devices.append(
@@ -561,7 +571,7 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]
"visible_ordinal": ordinal,
"name": props.name,
"total_gb": round(total_bytes / (1024**3), 2),
- "used_gb": round(used_bytes / (1024**3), 2),
+ "used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None,
}
)
except Exception as e:
@@ -724,20 +734,30 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
return None, None
-def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]:
- """Query system-wide dedicated GPU VRAM via Windows Performance Counters.
+# ── Windows AMD/ROCm per-adapter VRAM (issue #7072) ──────────────────────────
+# amd-smi is disabled and hipMemGetInfo reports free==total, so read used from the
+# per-LUID "GPU Adapter Memory" perf counters and take each total from torch, so
+# every GPU shows instead of one fake device with GPU 0's total.
+# Placeholder adapters (Basic Render Driver / idle iGPU) drop only when they would
+# outnumber the real torch devices.
+_ROCM_WIN_ADAPTER_MIN_BYTES = 64 * 1024 * 1024 # 64 MiB
- Same data source as Task Manager, so cross-process usage is accurate.
- Works for any GPU vendor without amd-smi or nvidia-smi.
- Returns (used_gb, total_gb) or (None, None) on failure.
+
+def _rocm_windows_perf_counter_vram_by_adapter() -> Optional[list[tuple[str, float]]]:
+ """Per-adapter dedicated VRAM usage on Windows via Performance Counters.
+
+ Returns ``[(instance_name, used_bytes)]`` (one per LUID-named adapter), or
+ ``None`` when the counter is unavailable/localized/empty so callers fall back.
"""
if platform.system() != "Windows":
- return None, None
+ return None
try:
+ # Emit "|" per sample, or a __NONE__ sentinel.
ps = (
"$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'"
" -ErrorAction SilentlyContinue).CounterSamples;"
- "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}"
+ "if($s){$s|ForEach-Object{'{0}|{1}' -f $_.InstanceName,[int64]$_.CookedValue}}"
+ "else{'__NONE__'}"
)
r = subprocess.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
@@ -746,16 +766,167 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa
timeout = 5,
)
if r.returncode != 0 or not r.stdout.strip():
- return None, None
- used_bytes = float(r.stdout.strip())
- if used_bytes < 0:
- return None, None
- import torch as _torch
-
- total_bytes = _torch.cuda.get_device_properties(0).total_memory
- return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
+ return None
+ adapters: list[tuple[str, float]] = []
+ for line in r.stdout.splitlines():
+ line = line.strip()
+ if not line or line == "__NONE__" or "|" not in line:
+ continue
+ instance, _, raw = line.rpartition("|")
+ try:
+ used = float(raw.strip())
+ except (ValueError, TypeError):
+ continue
+ if used < 0:
+ continue
+ adapters.append((instance.strip(), used))
+ return adapters or None
except Exception:
- return None, None
+ return None
+
+
+def _match_adapter_used_to_devices(
+ adapter_useds: list[float], device_totals: list[float]
+) -> list[Optional[float]]:
+ """Attribute per-adapter used bytes to torch devices by capacity ranking.
+
+ Windows shares no key between LUID counters and torch ordinals, so usages are
+ ranked against device totals and each is trusted only when capacity *forces* it
+ (it exceeds every smaller device); an ambiguous ranking reports unknown
+ (``None``) rather than fabricate a per-index free.
+
+ Extra counters mean a hidden/display adapter, and the noise filter may have
+ dropped a real reading, so values are emitted only when the supra-threshold
+ counters number EXACTLY the visible devices AND capacity forces the mapping;
+ otherwise every device is unknown. Best-effort but correct for the common
+ loaded-card case (#7072). Returns a list aligned to ``device_totals``.
+ """
+ n = len(device_totals)
+ if n == 0:
+ return []
+ useds = sorted(adapter_useds, reverse = True)
+ ranked_positions = sorted(range(n), key = lambda i: -device_totals[i])
+ ranked_totals = [device_totals[pos] for pos in ranked_positions]
+ assigned: list[Optional[float]]
+ # More counters than devices -> a hidden/display adapter (check before noise filter).
+ if len(useds) > n:
+ non_trivial = [u for u in useds if u >= _ROCM_WIN_ADAPTER_MIN_BYTES]
+ if len(non_trivial) != n:
+ # Not a clean bijection (a masked GPU is busy or a visible card idle):
+ # no counter maps to a specific card, so report unknown.
+ return [None] * n
+ # Exactly n supra-threshold counters: extras were placeholders, so a
+ # capacity-ranked bijection is plausible.
+ useds = non_trivial
+ ranked_useds = [useds[rank] for rank in range(n)]
+ # A usage above its ranked capacity is a hidden larger GPU; clamping onto the
+ # smaller card would fabricate a fully-used reading.
+ for rank in range(n):
+ if ranked_useds[rank] > ranked_totals[rank]:
+ return [None] * n
+ # Capacity forces the mapping only when the usage exceeds the next-smaller
+ # capacity; the smallest card and merely-fitting usages stay unknown.
+ # Keeps 40 GiB over 48/8 GiB -> [40, None].
+ assigned = [None] * n
+ for rank, pos in enumerate(ranked_positions):
+ if rank + 1 < n and ranked_useds[rank] > ranked_totals[rank + 1]:
+ assigned[pos] = min(ranked_useds[rank], device_totals[pos])
+ return assigned
+ # No hidden adapters: every counter is a visible card, so ranking is a permutation.
+ ranked_useds = [useds[rank] if rank < len(useds) else 0.0 for rank in range(n)]
+ # Ambiguous if a strictly larger usage also fits the next smaller card: the two
+ # could be swapped without breaking capacity, so ranking can't tell them apart.
+ for rank in range(n - 1):
+ upper, lower = ranked_useds[rank], ranked_useds[rank + 1]
+ if upper > lower and upper <= ranked_totals[rank + 1]:
+ return [None] * n
+ assigned = [None] * n
+ for rank, pos in enumerate(ranked_positions):
+ if rank < len(useds):
+ assigned[pos] = min(useds[rank], device_totals[pos])
+ return assigned
+
+
+def _rocm_windows_per_device_vram(device_indices: list[int]) -> list[Dict[str, Any]]:
+ """Per-GPU VRAM on Windows AMD/ROCm: total from torch properties (reliable),
+ used from the per-adapter Dedicated Usage counter.
+
+ Returns ``{index, visible_ordinal, name, used_gb, total_gb}`` per visible GPU
+ (``used_gb`` may be ``None`` when the counter is unavailable), or ``[]`` when
+ torch can't enumerate devices so callers fall through to the torch last resort.
+ """
+ if platform.system() != "Windows":
+ return []
+ mod, _ = _torch_get_device_module()
+ if mod is None:
+ return []
+ # Totals/names from torch properties (mem_get_info's free==total quirk zeroes used).
+ dev_meta: list[Dict[str, Any]] = []
+ for ordinal, phys_idx in enumerate(device_indices):
+ try:
+ props = mod.get_device_properties(ordinal)
+ dev_meta.append(
+ {
+ "index": phys_idx,
+ "visible_ordinal": ordinal,
+ "name": props.name,
+ "total_bytes": int(props.total_memory),
+ }
+ )
+ except Exception as e:
+ logger.debug("torch property probe failed for ordinal %d: %s", ordinal, e)
+ if not dev_meta:
+ return []
+
+ adapters = _rocm_windows_perf_counter_vram_by_adapter()
+ if adapters:
+ assigned = _match_adapter_used_to_devices(
+ [used for _, used in adapters],
+ [d["total_bytes"] for d in dev_meta],
+ )
+ else:
+ # Counter unavailable: show every GPU with a correct total, used unknown.
+ assigned = [None] * len(dev_meta)
+
+ devices: list[Dict[str, Any]] = []
+ for meta, used_bytes in zip(dev_meta, assigned):
+ total_gb = round(meta["total_bytes"] / (1024**3), 2)
+ used_gb = round(used_bytes / (1024**3), 2) if used_bytes is not None else None
+ devices.append(
+ {
+ "index": meta["index"],
+ "visible_ordinal": meta["visible_ordinal"],
+ "name": meta["name"],
+ "used_gb": used_gb,
+ "total_gb": total_gb,
+ }
+ )
+ return devices
+
+
+def _rocm_windows_device_payload_entry(
+ device: DeviceType, dev: Dict[str, Any], gpu_util_pct: Optional[float]
+) -> Dict[str, Any]:
+ """Build a ``get_gpu_utilization`` device entry from a per-device VRAM dict."""
+ total_gb = dev["total_gb"]
+ used_gb = dev["used_gb"]
+ return {
+ "available": True,
+ "backend": _backend_label(device),
+ "index": dev["index"],
+ "visible_ordinal": dev["visible_ordinal"],
+ "name": dev.get("name", "Unknown"),
+ "gpu_utilization_pct": gpu_util_pct,
+ "temperature_c": None,
+ "vram_used_gb": used_gb,
+ "vram_total_gb": total_gb,
+ "vram_utilization_pct": round((used_gb / total_gb) * 100, 1)
+ if total_gb and total_gb > 0 and used_gb is not None
+ else None,
+ "power_draw_w": None,
+ "power_limit_w": None,
+ "power_utilization_pct": None,
+ }
def _gpu_utilization_payload(
@@ -821,30 +992,24 @@ def get_gpu_utilization() -> Dict[str, Any]:
index_kind = result.get("index_kind"),
)
- # Fallback Windows ROCm
+ # Fallback Windows ROCm: per-adapter VRAM attribution (issue #7072), so
+ # every visible GPU is shown instead of a sum collapsed onto one device.
if IS_ROCM and platform.system() == "Windows":
- _win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
- if _win_used is not None and _win_total is not None:
- _win_util = _rocm_windows_perf_counter_gpu_util_pct()
+ _win_ids = _get_parent_visible_gpu_spec().get("numeric_ids")
+ if not _win_ids:
+ _win_ids = list(range(_torch_get_physical_gpu_count() or 0))
+ _win_devices = _rocm_windows_per_device_vram(_win_ids)
+ if _win_devices:
+ # A single visible GPU can own the aggregate 3D-engine utilization;
+ # across several GPUs the sum isn't per-device, so leave it unset.
+ _win_util = (
+ _rocm_windows_perf_counter_gpu_util_pct() if len(_win_devices) == 1 else None
+ )
return _gpu_utilization_payload(
device,
[
- {
- "available": True,
- "backend": _backend_label(device),
- "index": 0,
- "visible_ordinal": 0,
- "gpu_utilization_pct": _win_util,
- "temperature_c": None,
- "vram_used_gb": _win_used,
- "vram_total_gb": _win_total,
- "vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
- if _win_total > 0
- else None,
- "power_draw_w": None,
- "power_limit_w": None,
- "power_utilization_pct": None,
- }
+ _rocm_windows_device_payload_entry(device, _wd, _win_util)
+ for _wd in _win_devices
],
)
@@ -901,7 +1066,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
"vram_used_gb": _used,
"vram_total_gb": _total,
"vram_utilization_pct": round((_used / _total) * 100, 1)
- if _total > 0
+ if _total > 0 and _used is not None
else None,
"power_draw_w": None,
"power_limit_w": None,
@@ -995,19 +1160,27 @@ def _apply_unified_memory_correction(
endpoints stay in sync on AMD iGPUs with unified memory.
"""
torch_total_gb = torch_info["total_gb"]
+ torch_used_gb = torch_info.get("used_gb")
smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
+ # torch sees the full unified (GTT) pool; amd-smi only the dedicated carve-out.
+ # Adopt torch's larger total regardless of used: on Windows ROCm torch_used is
+ # None (free==total sentinel) but its total stays authoritative. Overwrite used
+ # only when torch's is known, then recompute utilization against whatever remains.
if torch_total_gb > smi_total_gb:
- torch_used_gb = torch_info["used_gb"]
device_metrics["vram_total_gb"] = torch_total_gb
- device_metrics["vram_used_gb"] = torch_used_gb
+ if torch_used_gb is not None:
+ device_metrics["vram_used_gb"] = torch_used_gb
+ _used_for_pct = device_metrics.get("vram_used_gb")
device_metrics["vram_utilization_pct"] = (
- round((torch_used_gb / torch_total_gb) * 100, 1) if torch_total_gb > 0 else None
+ round((_used_for_pct / torch_total_gb) * 100, 1)
+ if torch_total_gb > 0 and _used_for_pct is not None
+ else None
)
logger.debug(
- "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with "
- "torch mem_get_info total (%.2f GB) for device %s",
- smi_total_gb,
+ "ROCm unified memory: adopted torch mem_get_info total (%.2f GB) over "
+ "amd-smi (%.2f GB) for device %s",
torch_total_gb,
+ smi_total_gb,
torch_info.get("index"),
)
@@ -1067,6 +1240,49 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
_reconcile_rocm_unified_memory(result, numeric_ids)
return result
+ # Windows AMD/ROCm (issue #7072): the System tab's VRAM source. The torch
+ # fallback below would report used==0 (free==total), so read per-adapter
+ # Dedicated Usage instead; total from torch properties.
+ if IS_ROCM and platform.system() == "Windows":
+ win_numeric_ids = parent_visible_spec.get("numeric_ids")
+ if win_numeric_ids:
+ win_ids = win_numeric_ids
+ win_index_kind = "physical"
+ else:
+ win_ids = list(range(_torch_get_physical_gpu_count() or 0))
+ win_index_kind = "relative"
+ win_devices = _rocm_windows_per_device_vram(win_ids)
+ if win_devices:
+ devices = []
+ for wd in win_devices:
+ total = wd["total_gb"]
+ used = wd["used_gb"]
+ devices.append(
+ {
+ "index": wd["index"],
+ "index_kind": win_index_kind,
+ "visible_ordinal": wd["visible_ordinal"],
+ "name": wd.get("name"),
+ "gpu_utilization_pct": None,
+ "temperature_c": None,
+ "vram_used_gb": used,
+ "vram_total_gb": total,
+ "vram_utilization_pct": round((used / total) * 100, 1)
+ if total and total > 0 and used is not None
+ else None,
+ "power_draw_w": None,
+ "power_limit_w": None,
+ "power_utilization_pct": None,
+ }
+ )
+ return {
+ "available": True,
+ "backend": _backend_label(device),
+ "parent_visible_gpu_ids": win_numeric_ids or [],
+ "devices": devices,
+ "index_kind": win_index_kind,
+ }
+
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
if device in (DeviceType.CUDA, DeviceType.XPU):
parent_ids = get_parent_visible_gpu_ids()
@@ -1094,7 +1310,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
"vram_used_gb": used,
"vram_total_gb": total,
"vram_utilization_pct": round((used / total) * 100, 1)
- if total > 0
+ if total > 0 and used is not None
else None,
"power_draw_w": None,
"power_limit_w": None,
diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx
index 1272a577e9..e38e2e5882 100644
--- a/studio/frontend/src/components/floating-monitor.tsx
+++ b/studio/frontend/src/components/floating-monitor.tsx
@@ -70,13 +70,18 @@ export function FloatingMonitor() {
(sum, device) => sum + (device.memory_total_gb ?? 0),
0,
);
- const vramUsed = devices.reduce(
- (sum, device) => sum + (device.vram_used_gb ?? 0),
- 0,
- );
+ // null usage = unknown (e.g. Windows ROCm perf counter): treating it as 0
+ // fabricates a 0-used readout, so the aggregate is unknown if any device is.
+ const vramUsageKnown =
+ devices.length > 0 &&
+ devices.every((device) => Number.isFinite(device.vram_used_gb));
+ const vramUsed = vramUsageKnown
+ ? devices.reduce((sum, device) => sum + (device.vram_used_gb ?? 0), 0)
+ : 0;
const vramPercent = clampPercent(
- vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0,
+ vramUsageKnown && vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0,
);
+ const unknownLabel = t("settings.resources.environment.unknown");
const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0;
@@ -164,17 +169,20 @@ export function FloatingMonitor() {
- {Math.round(vramPercent)}%
+ {vramUsageKnown ? `${Math.round(vramPercent)}%` : "--"}
- {formatGiB(vramUsed)} / {formatGiB(vramTotal)}
+ {vramUsageKnown ? formatGiB(vramUsed) : unknownLabel} /{" "}
+ {formatGiB(vramTotal)}
diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx
index 9a359f0f38..eb7fa03cf9 100644
--- a/studio/frontend/src/features/settings/tabs/resources-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx
@@ -90,8 +90,11 @@ function MetricTile({
label: string;
value: string;
detail: string;
- percent: number;
+ // null = usage unknown (e.g. Windows ROCm perf counter): show a dash and
+ // empty bar rather than a fabricated 0%.
+ percent: number | null;
}) {
+ const percentKnown = isFiniteNumber(percent);
const safePercent = clampPercent(percent);
return (
@@ -102,10 +105,10 @@ function MetricTile({
- {formatPercent(safePercent)}
+ {percentKnown ? formatPercent(safePercent) : "--"}
@@ -117,7 +120,7 @@ function MetricTile({
sum + (device.memory_total_gb ?? 0),
0,
);
- const vramUsed = devices.reduce(
- (sum, device) => sum + (device.vram_used_gb ?? 0),
- 0,
- );
- const vramFree = devices.reduce(
- (sum, device) =>
- sum +
- (device.vram_free_gb ??
- Math.max(0, (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0))),
- 0,
- );
- const vramPercent = vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0;
+ // null usage = unknown (e.g. Windows ROCm perf counter): treating it as 0
+ // fabricates a 0-used total, so the aggregate is unknown if any device is.
+ const vramUsageKnown =
+ devices.length > 0 &&
+ devices.every((device) => isFiniteNumber(device.vram_used_gb));
+ const vramUsed = vramUsageKnown
+ ? devices.reduce((sum, device) => sum + (device.vram_used_gb ?? 0), 0)
+ : null;
+ const vramFree = vramUsageKnown
+ ? devices.reduce(
+ (sum, device) =>
+ sum +
+ (device.vram_free_gb ??
+ Math.max(
+ 0,
+ (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0),
+ )),
+ 0,
+ )
+ : null;
+ const vramPercent =
+ vramUsageKnown && isFiniteNumber(vramUsed) && vramTotal > 0
+ ? (vramUsed / vramTotal) * 100
+ : 0;
return {
devices,
@@ -218,6 +233,7 @@ export function ResourcesTab() {
vramUsed,
vramFree,
vramPercent,
+ vramUsageKnown,
};
}, [systemInfo]);
@@ -259,6 +275,7 @@ export function ResourcesTab() {
: modelsFolderLoaded
? t("settings.resources.environment.unknown")
: t("common.loading");
+ const unknownLabel = t("settings.resources.environment.unknown");
return (
@@ -327,17 +344,21 @@ export function ResourcesTab() {
label={t("settings.resources.liveMonitor.vram")}
value={
hasGpu
- ? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}`
+ ? metrics.vramUsageKnown
+ ? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}`
+ : `${unknownLabel} / ${formatGiB(metrics.vramTotal)}`
: t("settings.resources.liveMonitor.noGpu")
}
detail={
hasGpu
- ? t("settings.resources.liveMonitor.free", {
- value: formatGiB(metrics.vramFree),
- })
+ ? metrics.vramUsageKnown
+ ? t("settings.resources.liveMonitor.free", {
+ value: formatGiB(metrics.vramFree),
+ })
+ : unknownLabel
: backendLabel
}
- percent={metrics.vramPercent}
+ percent={metrics.vramUsageKnown ? metrics.vramPercent : null}
/>
@@ -346,13 +367,33 @@ export function ResourcesTab() {
{hasGpu ? (
metrics.devices.map((device, index) => {
const ordinal = deviceOrdinal(device);
- const total = device.memory_total_gb ?? 0;
- const used = device.vram_used_gb ?? 0;
- const free = device.vram_free_gb ?? Math.max(0, total - used);
+ // Preserve null (unknown, e.g. Windows ROCm perf counter); coercing
+ // to 0 would render a fabricated 0 used / full free.
+ const total = device.memory_total_gb ?? null;
+ const used = device.vram_used_gb ?? null;
+ const free =
+ device.vram_free_gb ??
+ (isFiniteNumber(total) && isFiniteNumber(used)
+ ? Math.max(0, total - used)
+ : null);
const percent =
device.vram_utilization_pct ??
- (total > 0 ? (used / total) * 100 : null);
+ (isFiniteNumber(total) && total > 0 && isFiniteNumber(used)
+ ? (used / total) * 100
+ : null);
const safePercent = clampPercent(percent);
+ const usedText = isFiniteNumber(used)
+ ? formatGiB(used)
+ : unknownLabel;
+ const freeText = isFiniteNumber(free)
+ ? formatGiB(free)
+ : unknownLabel;
+ const totalText = isFiniteNumber(total)
+ ? formatGiB(total)
+ : unknownLabel;
+ const percentText = isFiniteNumber(percent)
+ ? formatPercent(safePercent)
+ : unknownLabel;
return (
- {formatPercent(safePercent)}{" "}
+ {percentText}{" "}
{t("settings.resources.gpu.vramUtilization")}
@@ -382,17 +423,17 @@ export function ResourcesTab() {
{t("settings.resources.gpu.used", {
- value: formatGiB(used),
+ value: usedText,
})}
{t("settings.resources.gpu.free", {
- value: formatGiB(free),
+ value: freeText,
})}
{t("settings.resources.gpu.total", {
- value: formatGiB(total),
+ value: totalText,
})}
From 6c78147980c017f748a4bc2234782baff46eff70 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 05:28:31 -0700
Subject: [PATCH 032/255] Studio: do not show Run for embedding-only non-GGUF
models in the Model Hub (#7245)
* Studio: do not show Run for embedding-only non-GGUF models in the Model Hub
A downloaded embedding-only repo (sentence-transformers, feature-extraction)
reports canChat by safetensors format and classifies as supported, so the Model
Hub showed a Run button that dead-ends at load. Keep embedding-only non-GGUF
models out of the Run gate. GGUF is unaffected (llama.cpp resolves embed vs
generate at load time), and these models stay trainable.
* Gate embedding-only models on the pipeline tag, not just capabilities
An embedding repo whose name or tags also imply code, vision, audio or
reasoning (e.g. jina-embeddings-v2-base-code picks up code from its
-code suffix) slipped the embedding-only Run gate and dead-ended on a
chat load. Treat a feature-extraction or sentence-similarity pipeline
tag as authoritative for the gate; the change only ever widens it.
* studio: tighten comments in the hub embedding run guard
* Tighten comments in the hub embedding run guard
---------
Co-authored-by: danielhanchen
---
.../features/hub/catalog/model-inspector.tsx | 23 +++++++++++++++++++
1 file changed, 23 insertions(+)
diff --git a/studio/frontend/src/features/hub/catalog/model-inspector.tsx b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
index 63cb34678b..5a6ae1615a 100644
--- a/studio/frontend/src/features/hub/catalog/model-inspector.tsx
+++ b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
@@ -59,6 +59,13 @@ import { ModelReadme } from "./model-readme";
import { OwnerAvatar } from "./owner-avatar";
import { AccessChip, CapabilityPill } from "./shared";
+// HF pipeline_tag values authoritative for embedding-only repos; capability
+// labels (code/vision/audio) can leak onto them via name or tags.
+const EMBEDDING_PIPELINE_TAGS: ReadonlySet = new Set([
+ "feature-extraction",
+ "sentence-similarity",
+]);
+
function ViewRepositoryButton({
repoId,
isDataset,
@@ -531,11 +538,27 @@ export const ModelInspector = memo(function ModelInspector({
? formatCompact(model.totalParams)
: "N/A";
const unslothSupported = unslothSupport.status !== "unsupported";
+ // Embedding-only non-GGUF repos have no generative head, so keep them out of
+ // the Run gate. Prefer the pipeline tag, else the capability heuristic.
+ const isEmbeddingOnly =
+ !model.isGguf &&
+ model.capabilities.some((c) => c.key === "embedding") &&
+ (EMBEDDING_PIPELINE_TAGS.has(model.pipelineTag?.toLowerCase() ?? "") ||
+ !model.capabilities.some(
+ (c) =>
+ c.key === "conversational" ||
+ c.key === "tools" ||
+ c.key === "reasoning" ||
+ c.key === "code" ||
+ c.key === "vision" ||
+ c.key === "audio",
+ ));
// Chat-only hosts (no supported GPU / usable MLX) run inference only through
// llama.cpp, so only GGUF is loadable.
const canRunModel =
!isDataset &&
(model.runtimeCapabilities?.canChat ?? true) &&
+ !isEmbeddingOnly &&
(model.isGguf || (!chatOnly && unslothSupported));
const canTrainModel =
!isDataset &&
From bdf51525ea1013443de8a6cd03561b56ff2f934c Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 05:29:18 -0700
Subject: [PATCH 033/255] Studio: make Stop and stall deadlines interrupt a
wedged stream portably (#7236)
* Studio: make Stop and stall deadlines interrupt a wedged stream portably
The cancel watcher unblocks a stalled read by shutting the socket down from
another thread, which works on POSIX but not reliably on native Windows, where
Winsock does not dependably wake a recv() already in progress on another thread.
Wrap the httpcore network stream so the reader loops each read in short slices
and polls the cancel event itself. Stop and the stall deadlines now interrupt a
wedged mid-stream read without any cross-thread socket teardown, and a slow but
still-alive stream is never torn down. The POSIX shutdown path is preserved.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Honor the post-first-token stall timeout in the cancel-aware read
httpcore snapshots request.extensions timeout read once when the body
starts, so lowering it to the stall timeout after the first token never
reached the socket read and a one-token-then-silent server hung for the
full prefill window. Re-read the live extensions timeout per call and
bound each read by it, falling back to the httpcore-passed timeout when
absent so prefill and normal completion are unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: tighten comments in the llama.cpp stall timeout path
* Tighten comments in the stream stall cancel path
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen
---
studio/backend/core/inference/llama_cpp.py | 74 +++++++++++
.../tests/test_llama_cpp_stall_timeout.py | 125 ++++++++++++++++++
2 files changed, 199 insertions(+)
create mode 100644 studio/backend/tests/test_llama_cpp_stall_timeout.py
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index d4cab81bdf..dee507fae2 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -9907,6 +9907,75 @@ class LlamaCppBackend:
except Exception:
logger.debug("Could not close httpx client", exc_info = True)
+ @staticmethod
+ def _install_cancel_aware_read(
+ client: "httpx.Client",
+ cancel_event: threading.Event,
+ response: Optional["httpx.Response"] = None,
+ poll_s: float = 0.2,
+ ) -> None:
+ """Wrap the httpcore stream so the reader interrupts its own blocked recv() on cancel.
+
+ A cross-thread socket shutdown wakes a parked recv() on POSIX but not on
+ Windows (Winsock), so read in short slices and poll cancel_event between them
+ (plain or TLS); slice timeouts are swallowed so a slow-but-alive stream survives.
+ httpcore snapshots request.extensions["timeout"]["read"] once at body start, so
+ given ``response`` we re-read the live value per call to honor the post-first-token
+ stall timeout instead of the long prefill timeout."""
+ import httpcore
+
+ def _live_read_timeout() -> Optional[float]:
+ if response is None:
+ return None
+ try:
+ ext = response.request.extensions.get("timeout")
+ if isinstance(ext, dict):
+ value = ext.get("read")
+ if isinstance(value, (int, float)):
+ return float(value)
+ except Exception:
+ pass
+ return None
+
+ try:
+ pool = getattr(getattr(client, "_transport", None), "_pool", None)
+ for connection in list(getattr(pool, "_connections", []) or []):
+ inner = getattr(connection, "_connection", None)
+ stream = getattr(inner, "_network_stream", None)
+ if stream is None or getattr(stream, "_unsloth_cancel_wrapped", False):
+ continue
+ orig_read = stream.read
+
+ def read(
+ max_bytes,
+ timeout = None,
+ _orig = orig_read,
+ ):
+ live = _live_read_timeout()
+ effective = live if live is not None else timeout
+ deadline = None if effective is None else time.monotonic() + effective
+ while True:
+ if cancel_event.is_set():
+ raise httpcore.ReadError("stream cancelled by user")
+ if deadline is None:
+ step = poll_s
+ else:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ raise httpcore.ReadTimeout("read operation timed out")
+ step = min(poll_s, remaining)
+ try:
+ return _orig(max_bytes, timeout = step)
+ except httpcore.ReadTimeout:
+ if deadline is not None and time.monotonic() >= deadline:
+ raise
+ continue # slow but alive: keep reading
+
+ stream.read = read
+ stream._unsloth_cancel_wrapped = True
+ except Exception:
+ logger.debug("Could not install cancel-aware read", exc_info = True)
+
@staticmethod
@contextlib.contextmanager
def _stream_with_retry(
@@ -9964,6 +10033,11 @@ class LlamaCppBackend:
headers = headers,
) as response:
_response_ref[0] = response
+ if cancel_event is not None:
+ # Portable mid-stream cancel: the reader polls cancel itself, so
+ # Stop interrupts a stalled read where the watcher's Windows socket
+ # shutdown does not. Pass response to honor the live stall timeout.
+ LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response)
if cancel_event is not None and cancel_event.is_set():
raise _LlamaStreamCancelled
yield response
diff --git a/studio/backend/tests/test_llama_cpp_stall_timeout.py b/studio/backend/tests/test_llama_cpp_stall_timeout.py
new file mode 100644
index 0000000000..da36f75e8e
--- /dev/null
+++ b/studio/backend/tests/test_llama_cpp_stall_timeout.py
@@ -0,0 +1,125 @@
+# 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 test for the post-first-token stall timeout in the cancel-aware read.
+
+httpcore snapshots ``request.extensions["timeout"]["read"]`` once at body start, so
+when ``_iter_text_cancellable`` lowers it after the first token, a one-token-then-silent
+server hangs for the full prefill window. The fix re-reads the live extensions timeout
+per call; a fake clock and always-silent stream check the read gives up after the live
+stall timeout, not the stale prefill one.
+"""
+
+from __future__ import annotations
+
+import inspect
+import sys
+import threading
+import types as _types
+from pathlib import Path
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+ sys.path.insert(0, _BACKEND_DIR)
+
+# Mirror sibling tests' stubbing so the module imports without fastapi.
+_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"))
+
+import httpcore # noqa: E402
+
+from core.inference import llama_cpp as llama_cpp_mod # noqa: E402
+from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
+
+_PREFILL_TIMEOUT = 1200.0 # what httpcore snapshots from the prefill timeout
+_STALL_TIMEOUT = 120.0 # the post-first-token stall timeout the wrapper must honor
+
+
+class _Obj:
+ pass
+
+
+def _install(response, clock, silent_stream):
+ """Wire fake client/pool so _install_cancel_aware_read finds the stream; return the wrapped stream.read."""
+ inner = _Obj()
+ inner._network_stream = silent_stream
+ connection = _Obj()
+ connection._connection = inner
+ pool = _Obj()
+ pool._connections = [connection]
+ transport = _Obj()
+ transport._pool = pool
+ client = _Obj()
+ client._transport = transport
+
+ cancel_event = threading.Event() # never set: we test the stall path, not cancel
+ sig = inspect.signature(LlamaCppBackend._install_cancel_aware_read)
+ if "response" in sig.parameters:
+ # Fixed signature: wrapper reads the live extensions timeout.
+ LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response)
+ else:
+ # Pre-fix signature: no response, so the stall assertion fails (proves the bug).
+ LlamaCppBackend._install_cancel_aware_read(client, cancel_event)
+ return silent_stream.read
+
+
+def test_stall_timeout_honored_after_first_token(monkeypatch):
+ clock = {"t": 0.0}
+ monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"])
+
+ # One token then silence: every read times out, advancing fake time by its timeout.
+ def silent_read(max_bytes, timeout = None):
+ clock["t"] += timeout if timeout is not None else 0.0
+ raise httpcore.ReadTimeout("slice timed out on silence")
+
+ stream = _Obj()
+ stream.read = silent_read
+
+ # First token seen: the live read timeout is lowered to the stall timeout.
+ request = _Obj()
+ request.extensions = {"timeout": {"read": _STALL_TIMEOUT}}
+ response = _Obj()
+ response.request = request
+
+ wrapped_read = _install(response, clock, stream)
+
+ # httpcore still passes the stale prefill timeout it snapshotted at body start.
+ with pytest.raises(httpcore.ReadTimeout):
+ wrapped_read(65536, timeout = _PREFILL_TIMEOUT)
+
+ # Must give up ~stall timeout after the last token, not the prefill window.
+ assert clock["t"] <= _STALL_TIMEOUT * 1.5, (
+ f"stall timeout not honored: waited {clock['t']}s "
+ f"(expected ~{_STALL_TIMEOUT}s, not {_PREFILL_TIMEOUT}s)"
+ )
+ assert clock["t"] >= _STALL_TIMEOUT * 0.5
+
+
+def test_prefill_timeout_used_when_no_live_override(monkeypatch):
+ """Without a lowered live timeout, the wrapper honors the passed prefill timeout, so the normal first-token wait is unchanged."""
+ clock = {"t": 0.0}
+ monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"])
+
+ def silent_read(max_bytes, timeout = None):
+ clock["t"] += timeout if timeout is not None else 0.0
+ raise httpcore.ReadTimeout("slice timed out on silence")
+
+ stream = _Obj()
+ stream.read = silent_read
+
+ # No timeout extension: wrapper falls back to httpcore's passed timeout.
+ request = _Obj()
+ request.extensions = {}
+ response = _Obj()
+ response.request = request
+
+ wrapped_read = _install(response, clock, stream)
+
+ with pytest.raises(httpcore.ReadTimeout):
+ wrapped_read(65536, timeout = _PREFILL_TIMEOUT)
+
+ assert clock["t"] >= _PREFILL_TIMEOUT * 0.9
From 2916e8449900ed97b8462c2108aeaa8c7bb3ca40 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Mon, 20 Jul 2026 05:55:39 -0700
Subject: [PATCH 034/255] Studio: clarify tool permission controls (#7181)
---
.../scripts/run-studio-permission-browser.sh | 69 ++++++++
.github/workflows/studio-mac-ui-smoke.yml | 11 +-
.github/workflows/studio-ui-smoke.yml | 16 +-
.github/workflows/studio-windows-ui-smoke.yml | 7 +
.../src/components/assistant-ui/thread.tsx | 49 ++----
.../chat/bypass-permissions-menu-item.tsx | 20 +--
.../src/features/chat/chat-settings-sheet.tsx | 11 +-
.../features/chat/permission-mode-select.tsx | 57 ++-----
.../src/features/chat/shared-composer.tsx | 12 +-
.../src/features/settings/tabs/chat-tab.tsx | 2 +-
studio/frontend/src/i18n/locales/en.ts | 2 +-
studio/frontend/src/index.css | 4 +-
tests/studio/playwright_chat_ui.py | 153 +++++++++++++++++-
13 files changed, 296 insertions(+), 117 deletions(-)
create mode 100755 .github/scripts/run-studio-permission-browser.sh
diff --git a/.github/scripts/run-studio-permission-browser.sh b/.github/scripts/run-studio-permission-browser.sh
new file mode 100755
index 0000000000..2007789035
--- /dev/null
+++ b/.github/scripts/run-studio-permission-browser.sh
@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+set -euo pipefail
+
+port="${1:?usage: $0 PORT BROWSER [CHANNEL]}"
+browser="${2:?usage: $0 PORT BROWSER [CHANNEL]}"
+channel="${3:-}"
+slug="$browser${channel:+-$channel}"
+artifact_dir="logs/playwright-permissions-$slug"
+server_log="logs/studio-permissions-$slug.log"
+studio_home="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}"
+set --
+if [ -n "${STUDIO_PERMISSION_FRONTEND:-}" ]; then
+ set -- -f "$STUDIO_PERMISSION_FRONTEND"
+fi
+
+mkdir -p "$artifact_dir"
+unsloth studio reset-password
+UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$port" "$@" \
+ >"$server_log" 2>&1 &
+studio_pid=$!
+
+cleanup() {
+ kill "$studio_pid" 2>/dev/null || true
+ wait "$studio_pid" 2>/dev/null || true
+}
+trap cleanup EXIT
+
+healthy=0
+for _ in $(seq 1 180); do
+ if curl -fs "http://127.0.0.1:$port/api/health" >/dev/null; then
+ healthy=1
+ break
+ fi
+ if ! kill -0 "$studio_pid" 2>/dev/null; then
+ tail -100 "$server_log" || true
+ exit 1
+ fi
+ sleep 1
+done
+if [ "$healthy" -ne 1 ]; then
+ tail -100 "$server_log" || true
+ exit 1
+fi
+
+old_password=$(cat "$studio_home/auth/.bootstrap_password")
+new_password="CIPerm-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
+if [ "${GITHUB_ACTIONS:-}" = "true" ]; then
+ echo "::add-mask::$old_password"
+ echo "::add-mask::$new_password"
+fi
+
+export BASE_URL="http://127.0.0.1:$port"
+export STUDIO_OLD_PW="$old_password"
+export STUDIO_NEW_PW="$new_password"
+export STUDIO_UI_STRICT=1
+export STUDIO_UI_PERMISSION_ONLY=1
+export STUDIO_UI_WALL_TIMEOUT_S=240
+export STUDIO_PLAYWRIGHT_BROWSER="$browser"
+export PW_ART_DIR="$artifact_dir"
+if [ -n "$channel" ]; then
+ export STUDIO_PLAYWRIGHT_CHANNEL="$channel"
+else
+ unset STUDIO_PLAYWRIGHT_CHANNEL || true
+fi
+
+python tests/studio/playwright_chat_ui.py
diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml
index 378e8ee5a6..7375e9bcbf 100644
--- a/.github/workflows/studio-mac-ui-smoke.yml
+++ b/.github/workflows/studio-mac-ui-smoke.yml
@@ -19,6 +19,7 @@ on:
- 'install.sh'
- 'pyproject.toml'
- 'tests/studio/**'
+ - '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-mac-ui-smoke.yml'
push:
branches: [main, pip]
@@ -96,7 +97,7 @@ jobs:
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- - name: Install Playwright + Chromium
+ - name: Install Playwright browsers
# No --with-deps on Mac: that flag installs Linux apt packages.
# GitHub-hosted macos-14 ships the system frameworks Chromium
# needs already.
@@ -112,7 +113,7 @@ jobs:
# in-script retry recover from any residual flakes.
run: |
pip install 'playwright>=1.55,<1.58'
- python -m playwright install chromium
+ python -m playwright install chromium webkit
- name: Patch Playwright pipeTransport.js to tolerate malformed JSON
# In Playwright 1.55-1.58, pipeTransport.js does
@@ -244,6 +245,10 @@ jobs:
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
+ - name: Cross-browser permission controls
+ run: |
+ bash .github/scripts/run-studio-permission-browser.sh 18895 webkit
+
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
@@ -343,5 +348,7 @@ jobs:
logs/studio_extra.log
logs/install.log
logs/playwright
+ logs/playwright-permissions-*
logs/playwright_extra
+ logs/studio-permissions-*.log
retention-days: 7
diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml
index b6d6d7d6e2..30280c281e 100644
--- a/.github/workflows/studio-ui-smoke.yml
+++ b/.github/workflows/studio-ui-smoke.yml
@@ -27,6 +27,7 @@ on:
# The Playwright test files themselves -- a PR that ONLY edits
# the test must still trigger UI CI.
- 'tests/studio/**'
+ - '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-ui-smoke.yml'
push:
branches: [main, pip]
@@ -107,13 +108,10 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- - name: Install Playwright + Chromium
+ - name: Install Playwright browsers
run: |
pip install 'playwright>=1.45'
- # --with-deps installs the OS-level runtime libs Chromium
- # needs (libnss3, libxkbcommon, etc.). About 30 s on a
- # warm runner.
- python -m playwright install --with-deps chromium
+ python -m playwright install --with-deps chromium firefox webkit
- name: Reset auth + boot Unsloth
run: |
@@ -182,6 +180,12 @@ jobs:
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
+ - name: Cross-browser permission controls
+ run: |
+ bash .github/scripts/run-studio-permission-browser.sh 18893 firefox
+ bash .github/scripts/run-studio-permission-browser.sh 18893 webkit
+ bash .github/scripts/run-studio-permission-browser.sh 18893 chromium chrome
+
# The chat UI test ends by clicking the Shutdown menuitem, which
# leaves the server dead. The extra UI test (Compare / Recipes /
# Export / Unsloth / Settings) needs a fresh Unsloth, so we boot a
@@ -297,6 +301,8 @@ jobs:
logs/install.log
logs/server-logs/
logs/playwright
+ logs/playwright-permissions-*
logs/playwright_extra
logs/playwright_ime
+ logs/studio-permissions-*.log
retention-days: 7
diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml
index 12d7475b53..f401f7be44 100644
--- a/.github/workflows/studio-windows-ui-smoke.yml
+++ b/.github/workflows/studio-windows-ui-smoke.yml
@@ -19,6 +19,7 @@ on:
- 'install.ps1'
- 'pyproject.toml'
- 'tests/studio/**'
+ - '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-windows-ui-smoke.yml'
push:
branches: [main, pip]
@@ -345,6 +346,10 @@ jobs:
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
+ - name: Edge permission controls
+ run: |
+ bash .github/scripts/run-studio-permission-browser.sh 18895 chromium msedge
+
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
@@ -402,5 +407,7 @@ jobs:
logs/studio_extra.log
logs/install.log
logs/playwright
+ logs/playwright-permissions-*
logs/playwright_extra
+ logs/studio-permissions-*.log
retention-days: 7
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 235dcb3c3d..32fbd61e09 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -1434,13 +1434,10 @@ const Composer: FC<{
const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const ragEnabled = useChatRuntimeStore((s) => s.ragEnabled);
- const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
- // More than 4 pills: collapse to icons only. Search and Code always show; the
- // permission pill shows in every mode except "off" (it renders null there);
- // Images, RAG, Canvas and MCP are conditional.
+ // More than 4 pills: collapse to icons only. Search, Code, and permissions
+ // always show; Images, RAG, Canvas and MCP are conditional.
const pillsCompact =
- 2 +
- (permissionMode !== "off" ? 1 : 0) +
+ 3 +
(ragEnabled ? 1 : 0) +
(supportsBuiltinImageGeneration ? 1 : 0) +
(artifactsEnabled ? 1 : 0) +
@@ -1556,20 +1553,6 @@ const Composer: FC<{
const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300);
return () => clearTimeout(t);
}, [composerText, draftKey]);
- // Two-row layout shows once the input wraps or a tool is on. Tools can
- // pre-select before a model loads, so an active toggle expands it either way.
- // Keep the composer expanded whenever the permission pill is visible.
- const composerExpanded =
- isMultiline ||
- hasAttachments ||
- hasPendingAudio ||
- toolsEnabled ||
- codeToolsEnabled ||
- imageToolsEnabled ||
- ragEnabled ||
- artifactsEnabled ||
- mcpEnabledForChat ||
- permissionMode !== "off";
// react-textarea-autosize re-measures only on value change or window resize,
// not on the width swap from expanding, so it keeps the taller height and
// leaves a stray blank row. Nudge a resize whenever input width changes.
@@ -1856,27 +1839,25 @@ const Composer: FC<{
- {/* Permission-level pill: always visible, even while the pill row
- is collapsed; opens the permission level dropdown. */}
+ {/* Permission-level pill: always visible and opens the permission
+ level dropdown. */}
- {composerExpanded ? (
- <>
-
-
-
-
- {artifactsEnabled ?
: null}
- {mcpEnabledForChat ? (
-
- ) : null}
- >
+
+
+
+
+ {artifactsEnabled ?
: null}
+ {mcpEnabledForChat ? (
+
) : null}
More menu. Like the MCP
-// pill, it opens a submenu where the user picks the permission level (Ask for
-// approval / Approve for me / Full access). Picking Full access demands the
-// danger warning; the other levels apply immediately. The menu closes normally
-// on select (no preventDefault) -- the warning dialog lives outside the menu
-// (BypassPermissionsConfirmDialog, mounted once at the chat-page root and
-// driven by the store), so it survives the menu unmounting and the "+"/More
-// popovers don't stay frozen.
+// Tool permissions entry for the composer "+" menu.
export function BypassPermissionsMenuItem() {
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const setBypassConfirmOpen = useChatRuntimeStore(
@@ -44,7 +40,7 @@ export function BypassPermissionsMenuItem() {
}
>
- Bypass permissions
+ Tool permissions
Enable Full access?
- Full access (Bypass permissions) is dangerous since the AI model
- might delete, corrupt your machine, and or cause real world damage
- to you or the world - only accept if you are certain
+ {FULL_ACCESS_WARNING}
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index b368a811fa..ecb6707a04 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -2432,13 +2432,14 @@ function ConfirmToolCallsToggle() {
When on, every local Unsloth tool call pauses for your approval
before it runs (the "Ask for approval" level). When off, tool calls
- run without prompts inside the sandbox (the "Off" level).
+ run without prompts inside the sandbox (the "Run automatically"
+ level).
Provider-hosted tools are not gated here.
{permissionMode === "full" ? (
- Overridden by Full access (Bypass permissions)
+ Overridden by Full access
) : null}
@@ -2459,11 +2460,11 @@ function BypassPermissionsToggle() {
- Bypass permissions
+ Tool permissions
- How Unsloth approves tool calls before they run. Full access is
- dangerous: it disables confirmations and the code sandbox.
+ Choose how Unsloth approves tool calls before they run. Full access
+ disables confirmations and the code sandbox.
{/* Full width, styled like the panel selects/preset input. */}
diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx
index a9cb8ce5d1..e6c89cf54a 100644
--- a/studio/frontend/src/features/chat/permission-mode-select.tsx
+++ b/studio/frontend/src/features/chat/permission-mode-select.tsx
@@ -7,7 +7,6 @@ import {
CircleOff,
Hand,
ShieldCheck,
- XIcon,
} from "lucide-react";
import { useState } from "react";
@@ -39,9 +38,8 @@ import {
} from "./stores/chat-runtime-store";
/**
- * Permission levels for the Bypass permissions dropdowns (General settings,
- * chat settings sheet, composer "+" menu). Off sits last as the toggle that
- * turns the feature off entirely.
+ * Permission levels for tool calls. Full access stays last because it disables
+ * both approval prompts and the code sandbox.
*/
export const PERMISSION_MODE_OPTIONS: readonly {
value: PermissionMode;
@@ -61,6 +59,12 @@ export const PERMISSION_MODE_OPTIONS: readonly {
description: "Only ask for actions detected as potentially unsafe",
icon: ShieldCheck,
},
+ {
+ value: "off",
+ label: "Run automatically",
+ description: "Run tool calls without approval prompts inside the sandbox",
+ icon: CircleOff,
+ },
{
value: "full",
label: "Full access",
@@ -68,14 +72,11 @@ export const PERMISSION_MODE_OPTIONS: readonly {
"Unrestricted: no approval prompts and the code sandbox is disabled",
icon: CircleAlert,
},
- {
- value: "off",
- label: "Off",
- description: "Turn off bypass permissions",
- icon: CircleOff,
- },
] as const;
+export const FULL_ACCESS_WARNING =
+ "Full access lets tool calls run without approval prompts or the code sandbox. They can modify or delete files, run commands, and make network requests. Enable it only when you trust the current task.";
+
export function permissionModeOption(mode: PermissionMode) {
return (
PERMISSION_MODE_OPTIONS.find((option) => option.value === mode) ??
@@ -100,10 +101,10 @@ export function PermissionModeMenuItems({
{
- // Reselecting the active level toggles the feature off.
if (option.value === permissionMode) {
- setPermissionMode("off");
- } else if (option.value === "full") {
+ return;
+ }
+ if (option.value === "full") {
onRequestFullAccess();
} else {
setPermissionMode(option.value);
@@ -154,9 +155,7 @@ export function FullAccessConfirmDialog({
Enable Full access?
- Full access (Bypass permissions) is dangerous since the AI model
- might delete, corrupt your machine, and or cause real world damage
- to you or the world - only accept if you are certain
+ {FULL_ACCESS_WARNING}
@@ -260,15 +259,10 @@ export function PermissionModeComposerPill({
const setBypassConfirmOpen = useChatRuntimeStore(
(s) => s.setBypassConfirmOpen,
);
- const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode);
const active = permissionModeOption(permissionMode);
const ActiveIcon = active.icon;
const fullAccess = permissionMode === "full";
- // Off means the feature is off: no pill (re-enable via the "+" menu or
- // settings, like the pre-levels bypass badge).
- if (permissionMode === "off") return null;
-
return (
@@ -278,30 +272,11 @@ export function PermissionModeComposerPill({
data-pill-label={active.label}
data-active={fullAccess ? "true" : "false"}
data-variant={fullAccess ? "danger" : undefined}
- data-keep-label="true"
aria-label="Permission level for tool calls"
title={`${active.label}: ${active.description}`}
>
- {/* The icon doubles as an off switch (mirrors the MCP pill): hover
- swaps it to an X; clicking it turns bypass permissions Off (no
- prompts, sandbox on) without opening the menu. data-keep-label
- exempts this pill from compact icon-only mode, so the off switch
- stays clickable even while the other pills are collapsed. */}
- {
- e.stopPropagation();
- }}
- onClick={(e) => {
- e.stopPropagation();
- setPermissionMode("off");
- }}
- className="composer-pill-glyph cursor-pointer"
- >
+
-
{active.label}
s.artifactsEnabled);
const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled);
const showCanvasMenuItem = useChatRuntimeStore((s) => s.showCanvasMenuItem);
- const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const setMcpEnabledForChat = useChatRuntimeStore(
(s) => s.setMcpEnabledForChat,
@@ -790,15 +789,12 @@ export function SharedComposer({
// can still be pre-selected, matching Web search/Code/MCP.
const ragDisabled = modelLoaded && (isExternalModel || !supportsTools);
const showRagPill = !isExternalModel;
- // Above 4 pills, collapse to icons only to cut clutter. Compare, Search and
- // Code always show; the permission pill shows in every mode except "off"
- // (it renders null there); the rest are conditional.
- const permissionPillVisible = permissionMode !== "off";
+ // Above 4 pills, collapse to icons only. Compare, Search, Code, and
+ // permissions always show; the rest are conditional.
const pillsCompact =
- 3 +
- (permissionPillVisible ? 1 : 0) +
+ 4 +
(showImagePill ? 1 : 0) +
- (showRagPill && ragEnabled && !ragDisabled ? 1 : 0) +
+ (showRagPill && ragEnabled ? 1 : 0) +
(showWebFetchPill ? 1 : 0) +
(artifactsEnabled ? 1 : 0) +
(mcpEnabledForChat ? 1 : 0) >
diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
index 3e419af78d..f7f3bccad6 100644
--- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
@@ -110,7 +110,7 @@ const PLUS_MENU_SETTINGS: {
},
{
id: "bypassPermissions",
- label: "Bypass permissions",
+ label: "Tool permissions",
icon: (
:not(.composer-pill-x),
+ .composer-pill-btn[data-active="true"]:has(.composer-pill-x):hover .composer-pill-glyph > :not(.composer-pill-x),
.unsloth-thinking-pill[data-active="true"]:hover .composer-pill-glyph > :not(.composer-pill-x) {
@apply opacity-0;
}
- .composer-pill-btn[data-active="true"]:hover .composer-pill-x,
+ .composer-pill-btn[data-active="true"]:has(.composer-pill-x):hover .composer-pill-x,
.unsloth-thinking-pill[data-active="true"]:hover .composer-pill-x {
@apply opacity-100;
}
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index 35b18756ff..b00b45f97a 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.py
@@ -52,6 +52,14 @@ TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000"))
# Wall-clock cap for the whole script (healthy run is 5-9 min).
WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720"))
+# Run only bootstrap plus permission controls for fast cross-browser checks.
+PERMISSION_ONLY = os.environ.get("STUDIO_UI_PERMISSION_ONLY", "0") == "1"
+
+# Default stays Chromium for CI. Local runs can select firefox/webkit or a
+# Chromium channel such as chrome/msedge.
+PLAYWRIGHT_BROWSER = os.environ.get("STUDIO_PLAYWRIGHT_BROWSER", "chromium").lower()
+PLAYWRIGHT_CHANNEL = os.environ.get("STUDIO_PLAYWRIGHT_CHANNEL") or None
+
# Per-fetch budget; /api/inference/load is the slowest (cold-cache GGUF load).
FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000"))
LOAD_FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_LOAD_TIMEOUT_MS", "180000"))
@@ -116,6 +124,126 @@ def soft_fail(m):
info(f"WARN (strict-off): {m}")
+def exercise_permission_mode_controls(page, shoot):
+ """Exercise labels, migration, persistence, confirmation, and focus."""
+ step("permission levels: labels, persistence, confirmation, and focus")
+ pill = page.locator('button[aria-label="Permission level for tool calls"]:visible').first
+ expect(pill).to_be_visible()
+
+ def expect_mode(label):
+ expect(pill).to_have_attribute("data-pill-label", label)
+ expect(pill).to_contain_text(label)
+
+ def open_menu():
+ pill.click()
+ menu = page.get_by_role("menu").last
+ expect(menu).to_be_visible()
+ return menu
+
+ def choose(label):
+ menu = open_menu()
+ item = menu.get_by_role("menuitem").filter(has_text = label).first
+ expect(item).to_be_visible()
+ item.click()
+
+ # Fresh profiles default to Approve for me.
+ expect_mode("Approve for me")
+ menu = open_menu()
+ for label in (
+ "Ask for approval",
+ "Approve for me",
+ "Run automatically",
+ "Full access",
+ ):
+ expect(menu.get_by_role("menuitem").filter(has_text = label).first).to_be_visible()
+ if menu.get_by_text("Off", exact = True).count() != 0:
+ fail("legacy Off label is still visible")
+ if menu.locator('[role="menuitem"] button, [role="menuitem"] [role="button"]').count():
+ fail("permission menu contains nested interactive controls")
+ page.keyboard.press("Escape")
+ expect(pill).to_be_focused()
+
+ # The active row is a no-op and must not open the Full access dialog.
+ choose("Approve for me")
+ expect_mode("Approve for me")
+ expect(page.get_by_role("alertdialog")).to_have_count(0)
+
+ # Pointer and compact-layout coverage.
+ page.set_viewport_size({"width": 390, "height": 844})
+ expect(pill).to_be_visible()
+ box = pill.bounding_box()
+ if box is None or box["x"] < 0 or box["x"] + box["width"] > 390:
+ fail(f"permission pill is clipped in compact layout: {box!r}")
+ page.set_viewport_size({"width": 1280, "height": 900})
+
+ # Legacy setting migration: true -> ask, false -> off, absent -> auto.
+ migration_cases = (
+ ("true", "Ask for approval"),
+ ("false", "Run automatically"),
+ (None, "Approve for me"),
+ )
+ for legacy_value, expected_label in migration_cases:
+ page.evaluate(
+ """(legacyValue) => {
+ localStorage.removeItem("unsloth_chat_permission_mode");
+ if (legacyValue === null) {
+ localStorage.removeItem("unsloth_chat_confirm_tool_calls");
+ } else {
+ localStorage.setItem(
+ "unsloth_chat_confirm_tool_calls",
+ legacyValue,
+ );
+ }
+ }""",
+ legacy_value,
+ )
+ page.reload(wait_until = "domcontentloaded")
+ expect(pill).to_be_visible()
+ expect_mode(expected_label)
+
+ choose("Run automatically")
+ expect_mode("Run automatically")
+ expect(page.locator('button[data-pill-label="Search"]:visible').first).to_be_visible()
+ expect(page.locator('button[data-pill-label="Code"]:visible').first).to_be_visible()
+ stored = page.evaluate("() => localStorage.getItem('unsloth_chat_permission_mode')")
+ if stored != "off":
+ fail(f"Run automatically persisted {stored!r}, expected 'off'")
+
+ # Full access requires explicit consent and never overwrites persistence.
+ choose("Full access")
+ dialog = page.get_by_role("alertdialog")
+ expect(dialog).to_be_visible()
+ expect(dialog.get_by_role("heading", name = "Enable Full access?")).to_be_visible()
+ expect(dialog).to_contain_text("the code sandbox")
+ dialog.get_by_role("button", name = "Cancel").click()
+ expect(dialog).to_be_hidden()
+ expect_mode("Run automatically")
+
+ choose("Full access")
+ expect(dialog).to_be_visible()
+ dialog.get_by_role("button", name = "I understand").click()
+ expect_mode("Full access")
+ expect(pill).to_have_attribute("data-variant", "danger")
+ active_icon = pill.locator(".composer-pill-glyph > :first-child")
+ pill.hover()
+ page.wait_for_timeout(200)
+ icon_opacity = float(active_icon.evaluate("el => getComputedStyle(el).opacity"))
+ if icon_opacity < 0.5:
+ fail(f"Full access icon disappeared on hover (opacity={icon_opacity})")
+ stored = page.evaluate("() => localStorage.getItem('unsloth_chat_permission_mode')")
+ if stored != "off":
+ fail(f"Full access overwrote persisted mode with {stored!r}")
+
+ page.reload(wait_until = "domcontentloaded")
+ expect(pill).to_be_visible()
+ expect_mode("Run automatically")
+
+ # Leave the full chat smoke in the fresh-install default.
+ choose("Approve for me")
+ expect_mode("Approve for me")
+ shoot("04-permission-levels")
+
+
def login_via_api(pw):
req = urllib.request.Request(
f"{BASE}/api/auth/login",
@@ -145,11 +273,17 @@ with sync_playwright() as p:
# DB is still migrating; this 30s probe catches that gap before we
# sink 60s into a change-password timeout. Diagnostic only.
wait_for_health(BASE, timeout = 30.0, info = info)
- # Chromium launch args: see `tests/studio/_playwright_robust.py`.
- browser = p.chromium.launch(
- headless = True,
- args = chromium_launch_args(),
- )
+ if PLAYWRIGHT_BROWSER not in ("chromium", "firefox", "webkit"):
+ fail(f"unsupported STUDIO_PLAYWRIGHT_BROWSER={PLAYWRIGHT_BROWSER!r}")
+ browser_type = getattr(p, PLAYWRIGHT_BROWSER)
+ launch_kwargs = {"headless": True}
+ if PLAYWRIGHT_BROWSER == "chromium":
+ launch_kwargs["args"] = chromium_launch_args()
+ if PLAYWRIGHT_CHANNEL:
+ launch_kwargs["channel"] = PLAYWRIGHT_CHANNEL
+ elif PLAYWRIGHT_CHANNEL:
+ fail("STUDIO_PLAYWRIGHT_CHANNEL requires chromium")
+ browser = browser_type.launch(**launch_kwargs)
ctx = browser.new_context(
viewport = {"width": 1280, "height": 900},
# Reduce motion so view-transition animations don't intercept
@@ -364,6 +498,15 @@ with sync_playwright() as p:
raise last_err
shoot("03-chat-loaded")
+ exercise_permission_mode_controls(page, shoot)
+ if PERMISSION_ONLY:
+ info(
+ "permission-only run passed "
+ f"(browser={PLAYWRIGHT_BROWSER}, channel={PLAYWRIGHT_CHANNEL or 'bundled'})"
+ )
+ browser.close()
+ sys.exit(0)
+
# /api/models/list and /api/inference/load need a bearer; the
# frontend stores it under "unsloth_auth_token" (auth/session.ts).
token = robust_evaluate(
From 39a999c0566fc3721651866c8f724a515d616dc2 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 20 Jul 2026 13:57:14 +0100
Subject: [PATCH 035/255] Update README with latest features and Unsloth Start
(#7258)
* Document local agent connections
* Refresh README features and news
* Tighten README feature copy
* Restore selective README emphasis
* Add Unsloth Start quickstart
* Update
* Mention
* Update README.md
* Reduce
* Restore-inference-order
* Split-agent-API-features
---
README.md | 62 ++++++++++++++++++++++++++++++++++++++++++++-----------
1 file changed, 50 insertions(+), 12 deletions(-)
diff --git a/README.md b/README.md
index 085c7718e5..6aa8f4f4c3 100644
--- a/README.md
+++ b/README.md
@@ -11,6 +11,7 @@ Unsloth Studio lets you run and train models locally.
Features •
+ News •
Quickstart •
Notebooks •
Documentation
@@ -47,15 +48,44 @@ Unsloth Studio (Beta) lets you run and train text, [audio](https://unsloth.ai/do
* [Auto set inference settings](https://unsloth.ai/docs/new/studio/chat#auto-parameter-tuning) and customize chat templates.
* We work directly with teams behind [gpt-oss](https://docs.unsloth.ai/new/gpt-oss-how-to-run-and-fine-tune#unsloth-fixes-for-gpt-oss), [Qwen3](https://www.reddit.com/r/LocalLLaMA/comments/1kaodxu/qwen3_unsloth_dynamic_ggufs_128k_context_bug_fixes/), [Llama 4](https://github.com/ggml-org/llama.cpp/pull/12889), [Mistral](https://huggingface.co/mistralai/Mistral-Medium-3.5-128B/discussions/18), [Gemma 1-3](https://news.ycombinator.com/item?id=39671146), and [Phi-4](https://unsloth.ai/blog/phi4), where we’ve fixed bugs that improve model accuracy.
* Chat with images, audio, PDFs, code, DOCX and more. [Connect API providers](https://unsloth.ai/docs/integrations/connections) (OpenAI, Anthropic) or servers (vLLM, Ollama).
+* [**Compare any two models**](https://unsloth.ai/docs/new/studio/chat#model-arena) side by side with the same prompt.
+* **OpenAI/Anthropic-compatible APIs**: Serve local models through `/v1/chat/completions`, `/v1/responses` and `/v1/messages`.
+* **Connect local models to agents**: Use `unsloth start` with Claude Code, Codex, Hermes and more.
+* **Web/PDF search** can read PDF papers, manuals and other PDF results.
+* **GGUF hardware controls**: Choose GPUs/layers, offload MoE experts, use multi-GPU or Tensor Parallelism.
+* The opt-in **MCP control endpoint** lets AI clients manage models, training, recipes and exports.
### Training
-* Train and RL **500+ models** up to **2x faster** with up to **70% less VRAM**, with no accuracy loss.
-* Custom Triton and mathematical **kernels**. See some collabs we did with [PyTorch](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) and [Hugging Face](https://unsloth.ai/docs/new/faster-moe).
+* Train and RL **500+ models** up to **2x faster** with **70% less VRAM**; MoE up to **12x faster**.
+* Train and run RL on [AMD GPUs](https://unsloth.ai/docs/basics/amd) across Windows, WSL and Linux.
* **Data Recipes**: [Auto-create datasets](https://unsloth.ai/docs/new/studio/data-recipe) from **PDF, CSV, DOCX** etc. Edit data in a visual-node workflow.
-* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** (RL): The most efficient [RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide) library, using **80% less VRAM** for GRPO, [FP8](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) etc.
-* Supports full fine-tuning, RL, pretraining, 4-bit, 16-bit and, FP8 training.
+* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** uses **80% less VRAM** for GRPO, FP8 and vision RL, with 7x longer contexts.
+* [**Long-context training**](https://unsloth.ai/docs/new/3x-faster-training-packing): **3x faster**, 30% less VRAM and 500K+ context.
+* Supports LoRA/QLoRA, full fine-tuning, RL, pretraining, 4-bit, 16-bit and FP8.
+* Custom Triton and mathematical **kernels** built with PyTorch and Hugging Face.
* **Observability**: Monitor training live, track loss and GPU usage and customize graphs.
* [Multi-GPU](https://unsloth.ai/docs/basics/multi-gpu-training-with-unsloth) training is supported, with major improvements coming soon.
+## 🚀 Unsloth Start
+
+[Unsloth Start](https://unsloth.ai/docs/integrations/unsloth-start) connects [Claude Code](https://unsloth.ai/docs/basics/claude-code), [Codex](https://unsloth.ai/docs/basics/codex) and other agents to local models with one command.
+
+Start Unsloth, load a model, open your project folder, then run:
+
+```bash
+unsloth start claude
+```
+
+Replace `claude` with any supported agent:
+
+| Agent | Command |
+| --- | --- |
+| Claude Code | `unsloth start claude` |
+| OpenAI Codex | `unsloth start codex` |
+| Hermes Agent | `unsloth start hermes` |
+| OpenClaw | `unsloth start openclaw` |
+| OpenCode | `unsloth start opencode` |
+| Pi Coding Agent | `unsloth start pi` |
+
## 📥 Install
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.
@@ -65,7 +95,8 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
* **CPU:** Supported for Chat and Data Recipes currently
* **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more
* **macOS:** Training, MLX and GGUF inference are ALL supported.
-* **AMD:** Chat + Data works. Train with [Unsloth Core](#unsloth-core-code-based). Unsloth Studio support is out soon.
+* **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd).
+* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819).
* **Multi-GPU:** Available now, with a major upgrade on the way
#### macOS, Linux, WSL:
@@ -122,7 +153,7 @@ You can use the same Docker image as Unsloth Studio.
#### AMD, Intel:
For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth).
-To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
+To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/basics/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
## 📒 Free Notebooks
@@ -148,13 +179,20 @@ Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Ad
- See detailed documentation for Unsloth [here](https://unsloth.ai/docs)
## 🦥 Unsloth News
-- **Connections**: Connect any API provider (OpenAI, Anthropic) or server (vLLM, Ollama). [Guide](https://unsloth.ai/docs/integrations/connections)
-- **MTP**: Run Qwen3.6 MTP in Unsloth. MTP settings are autoset specific to your hardware. [Guide](https://unsloth.ai/docs/models/qwen3.6#mtp-guide)
-- **API inference endpoint**: Deploy and run local LLMs in Claude Code, Codex tools. [Guide](https://unsloth.ai/docs/basics/api)
-- **Qwen3.6**: Qwen3.6-35B-A3B can now be trained and run in Unsloth Studio. [Blog](https://unsloth.ai/docs/models/qwen3.6)
-- **Gemma 4**: Run and train Google’s new models directly in Unsloth. [Blog](https://unsloth.ai/docs/models/gemma-4)
+- **AMD training**: Train, run RL, chat and deploy on AMD GPUs across Windows, WSL and Linux. [Guide](https://unsloth.ai/docs/basics/amd)
+- **GGUF hardware controls**: Choose GPU/layer placement, offload MoE experts and use multi-GPU or Tensor Parallelism. [#6414](https://github.com/unslothai/unsloth/pull/6414)
+- **Local models for any agent**: Use `unsloth start` with Claude Code, Codex, Hermes, OpenCode, OpenClaw, Pi and more through Unsloth's OpenAI- and Anthropic-compatible APIs. [Guide](https://unsloth.ai/docs/basics/api)
+- **MCP control endpoint**: Let compatible clients manage models, training, recipes, checkpoints and exports. [#7191](https://github.com/unslothai/unsloth/pull/7191)
+- **Local inference reliability**: Resume long chats faster, recover stalled downloads and reuse existing GGUF files. [#7204](https://github.com/unslothai/unsloth/pull/7204) • [#6858](https://github.com/unslothai/unsloth/pull/6858) • [#7209](https://github.com/unslothai/unsloth/pull/7209)
+- **New models**: [Qwen-AgentWorld](https://huggingface.co/unsloth/Qwen-AgentWorld-35B-A3B-GGUF), [Ornith](https://huggingface.co/unsloth/models?search=ornith), [Kimi K2.7 Code](https://unsloth.ai/docs/models/kimi-k2.7-code) and [MiniMax M3](https://unsloth.ai/docs/models/minimax-m3)
+- **GLM-5.2**: Run Z.ai's 744B-parameter, 1M-context open model locally with Unsloth Dynamic GGUFs. [Guide](https://unsloth.ai/docs/models/glm-5.2)
+- **DeepSeek-V4**: Run DeepSeek-V4-Flash locally with corrected multi-turn and tool-calling behavior. [Guide](https://unsloth.ai/docs/models/deepseek-v4)
+- **DiffusionGemma**: Run and fine-tune Google's diffusion language model with 1.8x faster inference in Unsloth Studio. [Guide](https://unsloth.ai/docs/models/diffusiongemma)
+- **Qwen3.6**: Run and train Qwen3.6 with MTP for 1.4-2.2x faster inference and NVFP4 quants for supported GPUs. [Guide](https://unsloth.ai/docs/models/qwen3.6)
+- **Gemma 4**: Run and train Gemma 4 text, image and audio models with QAT, MTP, GGUF and MLX support. [Guide](https://unsloth.ai/docs/models/gemma-4)
+- **MCP servers**: Connect local models to files, apps, databases and external tools through Model Context Protocol. [Guide](https://unsloth.ai/docs/basics/mcp)
+- **Connections**: Mix local models with API providers (OpenAI, Anthropic) or servers (vLLM, Ollama) in the same interface. [Guide](https://unsloth.ai/docs/integrations/connections)
- **Introducing Unsloth Studio**: our new web UI for running and training LLMs. [Blog](https://unsloth.ai/docs/new/studio)
-- **Qwen3.5** - 0.8B, 2B, 4B, 9B, 27B, 35-A3B, 112B-A10B are now supported. [Guide + notebooks](https://unsloth.ai/docs/models/qwen3.5/fine-tune)
- Train **MoE LLMs 12x faster** with 35% less VRAM - DeepSeek, GLM, Qwen and gpt-oss. [Blog](https://unsloth.ai/docs/new/faster-moe)
- **Embedding models**: Unsloth now supports ~1.8-3.3x faster embedding fine-tuning. [Blog](https://unsloth.ai/docs/new/embedding-finetuning) • [Notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks#embedding-models)
- New **7x longer context RL** vs. all other setups, via our new batching algorithms. [Blog](https://unsloth.ai/docs/new/grpo-long-context)
From 7313be1466fa10406b1dacb672859047e0236956 Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Mon, 20 Jul 2026 09:57:49 -0300
Subject: [PATCH 036/255] Studio: lock the last enabled GPU switch instead of
silently ignoring it (#7259)
---
.../src/features/chat/chat-settings-sheet.tsx | 17 ++++++++++-------
1 file changed, 10 insertions(+), 7 deletions(-)
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index ecb6707a04..bd22cc4f55 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -723,11 +723,11 @@ export function ChatSettingsPanel({
const autoLayers = isManual && gpuLayers < 0;
// GPUs actually in use: the picked subset, or all visible when none picked.
const gpusInUse = selectedGpuIds ?? gpuDevices.map((d) => d.index);
- // TP is off with fewer than 2 GPUs in use (single GPU, or the picker narrowed
- // to one): tensor split is a no-op there and aborts on some archs. Mirrors the
- // multi-GPU gate on the GPU picker / Split ratio. (Under Auto layers the whole
- // TP control is hidden -- llama.cpp's --fit aborts under --split-mode tensor.)
- const tpDisabled = gpusInUse.length <= 1;
+ // The picker must keep one GPU selected.
+ const singleGpuInUse = gpusInUse.length <= 1;
+ // TP needs at least two GPUs because tensor split is a no-op on one and may
+ // abort. Auto layers hides TP because --fit aborts under --split-mode tensor.
+ const tpDisabled = singleGpuInUse;
// Manual gpu-layers ceiling = model layer count + 1 (else a safe fallback):
// llama.cpp counts the output layer as one more offloadable layer past the
// repeating blocks ("offloaded 33/33" needs -ngl 33 on a 32-block model), so
@@ -1537,7 +1537,7 @@ export function ChatSettingsPanel({
Which GPUs this model may use. Unchecked GPUs are hidden
from llama.cpp (CUDA_VISIBLE_DEVICES, or
HIP_VISIBLE_DEVICES on ROCm). Leave all checked to use
- every GPU.
+ every GPU. At least one GPU must stay selected.
@@ -1557,7 +1557,10 @@ export function ChatSettingsPanel({
checked={isGpuChecked(d.index)}
onCheckedChange={() => toggleGpu(d.index)}
data-test-id={`gpu-pick-${d.index}`}
- disabled={modelControlsDisabled}
+ disabled={
+ modelControlsDisabled ||
+ (isGpuChecked(d.index) && singleGpuInUse)
+ }
/>
))}
From c7b17c455bb545fdfd54a00716331918fc88840d Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 06:05:28 -0700
Subject: [PATCH 037/255] Fix `unsloth start` on Windows: agent install, PATH
resolution, and local model selection (#7257)
* unsloth start: fix Windows agent install/launch and local model selection
- claude: pin availableModels to the served model in the session --settings
overlay so a user's ~/.claude/settings.json allowlist no longer substitutes
the org default for the local Unsloth model. The allowlist covers --model,
ANTHROPIC_MODEL and the model setting, and an empty [] is ignored, so the pin
lists the model explicitly.
- installs: run the Windows installer under -ExecutionPolicy Bypass
(process-scoped, nothing persistent) so npm's npm.ps1 and irm|iex scripts run
under the default Restricted policy; on failure, hint at Set-ExecutionPolicy
-Scope CurrentUser -ExecutionPolicy RemoteSigned for a hand-run retry.
- PATH: resolve agents installed to ~/.local/bin (claude) and %APPDATA%\npm (npm
agents) in-process, so a fresh install launches without opening a new shell and
an already-installed agent is not re-prompted for install.
- load message: "Loading - please wait" while a model loads.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* unsloth start: resolve agent version against the launch PATH
The claude/codex/opencode version probes ran shutil.which while building the
command, before _launch augments PATH with the known install dirs. An agent
present only in ~/.local/bin or %APPDATA%\npm was therefore missed, assumed to
be a current build, and launched with flags an older build rejects (claude
aborts on the unknown flags). Route the three probes through a new
_which_with_install_dirs() so each resolves the same binary _launch will,
restoring PATH afterward so only _launch persists the augmentation.
Add regression tests for the three probes (POSIX and the Windows npm dir) and
make the Windows-branch tests run on POSIX hosts (pinning Path to the native
flavour so a simulated os.name does not make pathlib build WindowsPath).
* unsloth start: keep os.defpath when augmenting an unset PATH
_augment_path_with_install_dirs collapsed an unset PATH to just the install
dirs, dropping the os.defpath fallback (/bin:/usr/bin) that shutil.which and
exec*p* use when PATH is absent. A system-installed agent then looked missing
and the launched child lost its normal PATH. Seed os.defpath when PATH is
unset; an explicitly empty PATH is left as-is (search nothing), matching
shutil.which. Add regression tests for the augment helper and the version-probe
wrapper.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
unsloth_cli/commands/start.py | 116 ++++++++++---
unsloth_cli/tests/test_start.py | 277 ++++++++++++++++++++++++++++++--
2 files changed, 356 insertions(+), 37 deletions(-)
diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py
index 707c2b3c90..3d73df65be 100644
--- a/unsloth_cli/commands/start.py
+++ b/unsloth_cli/commands/start.py
@@ -192,7 +192,7 @@ _OPENCODE_NATIVE_AUTO_MIN_VERSION = (1, 17, 12)
def _opencode_supports_native_auto() -> bool:
- executable = shutil.which("opencode")
+ executable = _which_with_install_dirs("opencode")
if executable is None:
# No local binary: a --no-launch recipe may run elsewhere, and _run installs the
# current release on launch -- either way assume native --auto is available.
@@ -826,7 +826,7 @@ def _resolve_model(
)
if requested and match is None:
typer.echo(
- f"Ensuring {requested} is loaded with the requested settings…"
+ f"Loading {requested} - please wait…"
if load_has_overrides
else f"Loading {requested} on the Unsloth server (this can take a while)…"
)
@@ -902,17 +902,21 @@ def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None:
_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections"
-# Session overlay applied via `claude --settings`; suppresses the attribution header
-# for THIS run only (no ~/.claude write) so llama.cpp KV-cache reuse is preserved. It
-# reinforces the CLAUDE_CODE_ATTRIBUTION_HEADER env var on builds that read the setting
-# only from settings.json.
-_CLAUDE_SETTINGS_OVERLAY = '{"env":{"CLAUDE_CODE_ATTRIBUTION_HEADER":"0"}}'
+
+
+def _claude_settings_overlay(model_id: str) -> str:
+ # Session-only `claude --settings` overlay (command-line tier, no ~/.claude write):
+ # suppress the attribution header, and pin availableModels to the served model so a
+ # user allowlist can't reject it. The pin must be non-empty; [] is ignored.
+ return json.dumps(
+ {"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}, "availableModels": [model_id]}
+ )
def _claude_version() -> Optional[tuple]:
# None = no local `claude` (a --no-launch printout for another machine; assume a
# current build). An unparseable version is treated as too old for the new flags.
- executable = shutil.which("claude")
+ executable = _which_with_install_dirs("claude")
if executable is None:
return None
try:
@@ -929,16 +933,14 @@ def _claude_version() -> Optional[tuple]:
return (0,)
-def _claude_flags() -> list:
- # Both knobs preserve llama.cpp KV-cache reuse: --exclude-dynamic-system-prompt-sections
- # moves per-session context out of the system prompt, and --settings suppresses the
- # attribution header for this session only (no persistent ~/.claude write; the env var
- # sets it too). Claude Code < 2.1.98 aborts on unknown flags, so gate on the version;
- # no local binary means a printout for another machine, so assume a current build.
+def _claude_flags(model_id: str) -> list:
+ # KV-cache-preserving flags: move per-session context out of the system prompt and pass
+ # the session overlay. claude < 2.1.98 rejects unknown flags; no local binary means a
+ # printout for another machine, so assume a current build.
version = _claude_version()
if version is not None and version < (2, 1, 98):
return []
- return [_DYNAMIC_SECTIONS_FLAG, "--settings", _CLAUDE_SETTINGS_OVERLAY]
+ return [_DYNAMIC_SECTIONS_FLAG, "--settings", _claude_settings_overlay(model_id)]
def _merge_codex_config(existing: str, base: str) -> str:
@@ -971,7 +973,7 @@ _CODEX_MODEL_CATALOG_MIN_VERSION = (0, 110, 0)
def _codex_supports_model_catalog() -> bool:
- executable = shutil.which("codex")
+ executable = _which_with_install_dirs("codex")
if executable is None:
# A --no-launch recipe may be copied to another machine; assume a current Codex.
return True
@@ -1202,6 +1204,53 @@ def _refresh_windows_path() -> None:
os.environ["PATH"] = os.pathsep.join(entries)
+def _augment_path_with_install_dirs() -> None:
+ # Append known install dirs to PATH so a freshly installed agent resolves without a new
+ # shell: some installers write the binary but not PATH (claude drops ~/.local/bin and
+ # only prints a note; npm -g shims land in %APPDATA%\npm). Appended, so precedence holds.
+ try:
+ home = Path.home()
+ except (RuntimeError, OSError):
+ return
+ candidates = [home / ".local" / "bin"]
+ if os.name == "nt":
+ appdata = os.environ.get("APPDATA")
+ if appdata:
+ candidates.append(Path(appdata) / "npm")
+ current = os.environ.get("PATH")
+ if current is None:
+ # PATH unset: shutil.which() and exec*p* fall back to os.defpath (e.g. /bin:/usr/bin), so
+ # keep that default instead of collapsing to just the install dirs (which would hide a
+ # system-installed agent and strip the launched child's normal PATH). An explicitly empty
+ # PATH is left as-is: like shutil.which, it means "search nothing", not os.defpath.
+ current = os.defpath
+ seen = {os.path.normcase(entry) for entry in current.split(os.pathsep) if entry}
+ additions = [
+ str(directory)
+ for directory in candidates
+ if directory.is_dir() and os.path.normcase(str(directory)) not in seen
+ ]
+ if additions:
+ os.environ["PATH"] = os.pathsep.join([current, *additions] if current else additions)
+
+
+def _which_with_install_dirs(name: str) -> Optional[str]:
+ # shutil.which(name), but searching the known agent install dirs too, so a version probe
+ # resolves the same binary _launch() will (it augments PATH before it runs). Without this an
+ # agent present only in ~/.local/bin / %APPDATA%\npm is missed, wrongly assumed current, and
+ # launched with flags an older build rejects. PATH is restored afterward: only _launch()
+ # should persist the augmentation for the child process.
+ original = os.environ.get("PATH")
+ _augment_path_with_install_dirs()
+ try:
+ return shutil.which(name)
+ finally:
+ if original is None:
+ os.environ.pop("PATH", None)
+ else:
+ os.environ["PATH"] = original
+
+
def _install_source(install_hint: str) -> Optional[str]:
"""The first http(s) URL an install hint fetches, or None (e.g. an npm install)."""
match = re.search(r"https?://[^\s'\")]+", install_hint)
@@ -1254,17 +1303,35 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]:
typer.secho(warning, fg = "yellow", err = True)
if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False):
return None
- # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm)
- # on Windows, /bin/sh (curl | bash, or npm) everywhere else.
+ # Run each hint through its shell: PowerShell on Windows, /bin/sh elsewhere.
+ # -ExecutionPolicy Bypass is process-scoped (nothing persistent) so npm's npm.ps1 and
+ # irm | iex run under the Windows default Restricted policy instead of failing with a
+ # PSSecurityException.
if os.name == "nt":
- install_command = ["powershell", "-NoProfile", "-Command", install_hint]
+ install_command = [
+ "powershell",
+ "-NoProfile",
+ "-ExecutionPolicy",
+ "Bypass",
+ "-Command",
+ install_hint,
+ ]
else:
install_command = ["/bin/sh", "-c", install_hint]
if subprocess.run(install_command).returncode != 0:
- _fail(f"Install command failed. Run it yourself, then re-run: {install_hint}")
- # The installer just wrote PATH to the registry (Windows); pull it into this
- # process so the freshly installed agent resolves without a shell restart.
+ message = f"Install command failed. Run it yourself, then re-run: {install_hint}"
+ if os.name == "nt":
+ # A hand-run retry can still hit the policy; point at the one-time per-user fix.
+ message += (
+ "\nIf it fails because running scripts is disabled (PSSecurityException), "
+ "allow local scripts for your user, then retry:\n"
+ " Set-ExecutionPolicy -Scope CurrentUser -ExecutionPolicy RemoteSigned"
+ )
+ _fail(message)
+ # Resolve the freshly installed agent without a shell restart: pull registry PATH
+ # (Windows) plus well-known install dirs the installer may not have added to PATH.
_refresh_windows_path()
+ _augment_path_with_install_dirs()
executable = shutil.which(name)
if executable is None:
_fail(
@@ -1290,6 +1357,9 @@ def _launch(
install_hint: str,
unset_env: tuple = (),
) -> NoReturn:
+ # Resolve well-known install dirs (e.g. ~/.local/bin) first, so an already-installed
+ # agent not yet on PATH is found instead of prompting a needless reinstall.
+ _augment_path_with_install_dirs()
executable = shutil.which(command[0]) or _install_agent(command[0], install_hint)
if executable is None:
_fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}")
@@ -1780,7 +1850,7 @@ def claude(
"claude",
"--model",
model_id,
- *_claude_flags(),
+ *_claude_flags(model_id),
*_yolo_command_flags("claude", yolo),
*ctx.args,
]
diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py
index 227918f63e..1e03d390d1 100644
--- a/unsloth_cli/tests/test_start.py
+++ b/unsloth_cli/tests/test_start.py
@@ -61,36 +61,73 @@ def _fake_claude(monkeypatch, version_output: str) -> None:
)
+def _path_aware_which(binaries: dict):
+ # A shutil.which fake that resolves a name only when its directory is on PATH at call time.
+ # Lets a test prove a version probe augments PATH before resolving: an agent present only in
+ # an install dir (~/.local/bin, %APPDATA%\npm) must still be found and version-checked.
+ def _which(name):
+ directory = binaries.get(name)
+ if directory is None:
+ return None
+ entries = os.environ.get("PATH", "").split(os.pathsep)
+ # os.path.join (not Path()) so this works when a test has flipped os.name to "nt": under
+ # a simulated os.name, pathlib would build the non-native flavour and raise.
+ return os.path.join(str(directory), name) if str(directory) in entries else None
+
+ return _which
+
+
+def _simulate_windows(monkeypatch) -> None:
+ # Exercise the `os.name == "nt"` branch on any host. Flipping os.name alone makes pathlib
+ # pick the non-native flavour (WindowsPath on POSIX, PosixPath on Windows) when a Path is
+ # constructed, which raises; pin Path to the host-native class (captured before the flip)
+ # so the branch logic runs without that crash. Keeps these tests green on Linux/Mac/WSL too.
+ monkeypatch.setattr(start, "Path", type(Path()))
+ monkeypatch.setattr(start.os, "name", "nt")
+
+
def test_claude_flags_passed_to_supported_claude(monkeypatch):
_fake_claude(monkeypatch, "2.1.98 (Claude Code)\n")
- assert start._claude_flags() == [
+ assert start._claude_flags(MODEL["id"]) == [
"--exclude-dynamic-system-prompt-sections",
"--settings",
- start._CLAUDE_SETTINGS_OVERLAY,
+ start._claude_settings_overlay(MODEL["id"]),
]
def test_claude_flags_skipped_on_old_claude(monkeypatch):
_fake_claude(monkeypatch, "2.0.14 (Claude Code)\n")
- assert start._claude_flags() == []
+ assert start._claude_flags(MODEL["id"]) == []
def test_claude_flags_skipped_on_unparseable_version(monkeypatch):
_fake_claude(monkeypatch, "weird build string\n")
- assert start._claude_flags() == []
+ assert start._claude_flags(MODEL["id"]) == []
def test_claude_flags_detected_when_version_not_first_token(monkeypatch):
# The X.Y.Z is pulled from anywhere in the output, so a format change (version not
# the first token) doesn't silently drop the optimization flags.
_fake_claude(monkeypatch, "claude version 2.1.98\n")
- assert start._claude_flags() == [
+ assert start._claude_flags(MODEL["id"]) == [
"--exclude-dynamic-system-prompt-sections",
"--settings",
- start._CLAUDE_SETTINGS_OVERLAY,
+ start._claude_settings_overlay(MODEL["id"]),
]
+def test_claude_settings_overlay_pins_served_model():
+ # The session overlay must pin availableModels to the served model: a user's allowlist
+ # in ~/.claude/settings.json otherwise rejects the Unsloth --model ("restricted by your
+ # organization's settings"), and no env var can bypass it. The override must be a
+ # NON-EMPTY array to take effect (an empty [] is ignored and the user's list still
+ # applies), so it lists exactly this model, for this session only.
+ overlay = json.loads(start._claude_settings_overlay(MODEL["id"]))
+ assert overlay["availableModels"] == [MODEL["id"]]
+ # The attribution-header suppression is preserved alongside it.
+ assert overlay["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0"
+
+
def test_install_agent_prompts_then_installs(monkeypatch):
# TTY + yes: run the documented install command, then re-resolve the now-present binary.
monkeypatch.setattr(start.os, "name", "posix")
@@ -126,7 +163,52 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch):
executable = start._install_agent("hermes", install_hint)
assert executable == r"C:\Users\samle\bin\hermes.exe"
- assert ran == [["powershell", "-NoProfile", "-Command", install_hint]]
+ # -ExecutionPolicy Bypass (process-scoped) lets npm's npm.ps1 wrapper and irm|iex
+ # scripts run even when the machine policy is the Windows default Restricted.
+ assert ran == [
+ ["powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-Command", install_hint]
+ ]
+
+
+def test_install_agent_windows_failure_hints_execution_policy(monkeypatch, capsys):
+ # A failed install on Windows points the user at the per-user execution-policy fix:
+ # our subprocess bypasses the policy, but their own shell may still block npm.ps1
+ # (PSSecurityException) when they run the install by hand.
+ monkeypatch.setattr(start.os, "name", "nt")
+ monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
+ monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True)
+ monkeypatch.setattr(
+ start.subprocess,
+ "run",
+ lambda *a, **k: SimpleNamespace(returncode = 1),
+ )
+ monkeypatch.setattr(start.shutil, "which", lambda _: None)
+
+ with pytest.raises(start.typer.Exit):
+ start._install_agent("codex", "npm install -g @openai/codex")
+
+ err = capsys.readouterr().err
+ assert "Install command failed" in err
+ assert "Set-ExecutionPolicy -Scope CurrentUser -ExecutionPolicy RemoteSigned" in err
+
+
+def test_install_agent_posix_failure_omits_execution_policy_hint(monkeypatch, capsys):
+ # The execution-policy hint is Windows-only; a POSIX install failure must not mention it.
+ monkeypatch.setattr(start.os, "name", "posix")
+ monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
+ monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True)
+ monkeypatch.setattr(
+ start.subprocess,
+ "run",
+ lambda *a, **k: SimpleNamespace(returncode = 1),
+ )
+
+ with pytest.raises(start.typer.Exit):
+ start._install_agent("codex", "npm install -g @openai/codex")
+
+ err = capsys.readouterr().err
+ assert "Install command failed" in err
+ assert "Set-ExecutionPolicy" not in err
def test_install_agent_warns_remote_installer_is_unverified_third_party(monkeypatch, capsys):
@@ -254,6 +336,172 @@ def test_refresh_windows_path_merges_registry_hives(monkeypatch):
]
+def test_augment_path_adds_existing_local_bin(monkeypatch, tmp_path):
+ # Claude's installer drops its binary in ~/.local/bin but only *suggests* adding it to
+ # PATH, so Unsloth appends it in-process to resolve the freshly installed agent.
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata")) # skip the npm candidate
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ start._augment_path_with_install_dirs()
+ entries = os.environ["PATH"].split(os.pathsep)
+ assert str(local_bin) in entries
+ # Appended (lowest precedence), so it never shadows an existing PATH entry.
+ assert entries[-1] == str(local_bin)
+
+
+def test_augment_path_skips_missing_and_duplicate_dirs(monkeypatch, tmp_path):
+ # A non-existent ~/.local/bin is not added; an already-present one is not duplicated.
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path) # no .local/bin created yet
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata")) # skip the npm candidate
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ start._augment_path_with_install_dirs()
+ assert os.environ["PATH"] == str(tmp_path / "existing")
+
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setenv("PATH", os.pathsep.join([str(tmp_path / "existing"), str(local_bin)]))
+ start._augment_path_with_install_dirs()
+ assert os.environ["PATH"].split(os.pathsep).count(str(local_bin)) == 1
+
+
+def test_augment_path_adds_npm_global_bin_on_windows(monkeypatch, tmp_path):
+ # npm -g shims (codex/opencode/pi) land in %APPDATA%\npm on Windows; add it so a freshly
+ # installed npm agent resolves even when that dir isn't on PATH yet.
+ _simulate_windows(monkeypatch)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path) # no ~/.local/bin created
+ npm_dir = tmp_path / "Roaming" / "npm"
+ npm_dir.mkdir(parents = True)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "Roaming"))
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ start._augment_path_with_install_dirs()
+ assert str(npm_dir) in os.environ["PATH"].split(os.pathsep)
+
+
+def test_which_with_install_dirs_finds_agent_and_restores_path(monkeypatch, tmp_path):
+ # The probe helper resolves against the augmented PATH but must NOT persist it: only
+ # _launch() should mutate PATH for the child process. Here `claude` is only in ~/.local/bin.
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata")) # skip the npm candidate
+ original = str(tmp_path / "existing")
+ monkeypatch.setenv("PATH", original) # local_bin NOT on PATH yet
+ monkeypatch.setattr(start.shutil, "which", _path_aware_which({"claude": local_bin}))
+ assert start._which_with_install_dirs("claude") == str(local_bin / "claude")
+ assert os.environ["PATH"] == original # restored, no global pollution
+
+
+def test_claude_flags_probes_old_agent_only_in_install_dir(monkeypatch, tmp_path):
+ # Regression: the version probe must augment PATH before resolving, so an OLD claude present
+ # only in ~/.local/bin (not yet on PATH) is detected as old and the unsupported flags are
+ # dropped -- the same binary _launch() will run. Before the fix the probe saw no binary,
+ # assumed a current build, and emitted flags the old claude rejects.
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata"))
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ monkeypatch.setattr(start.shutil, "which", _path_aware_which({"claude": local_bin}))
+ monkeypatch.setattr(
+ start.subprocess, "run", lambda *a, **k: SimpleNamespace(stdout = "2.0.14 (Claude Code)\n")
+ )
+ assert start._claude_flags(MODEL["id"]) == []
+
+
+def test_claude_flags_detects_supported_agent_only_in_install_dir(monkeypatch, tmp_path):
+ # The counterpart: a SUPPORTED claude present only in ~/.local/bin is now resolved and gets
+ # the flags, instead of being missed and (coincidentally) also assumed current.
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata"))
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ monkeypatch.setattr(start.shutil, "which", _path_aware_which({"claude": local_bin}))
+ monkeypatch.setattr(
+ start.subprocess, "run", lambda *a, **k: SimpleNamespace(stdout = "2.1.98 (Claude Code)\n")
+ )
+ assert start._claude_flags(MODEL["id"]) == [
+ "--exclude-dynamic-system-prompt-sections",
+ "--settings",
+ start._claude_settings_overlay(MODEL["id"]),
+ ]
+
+
+def test_claude_flags_probes_npm_install_dir_on_windows(monkeypatch, tmp_path):
+ # npm -g shims land in %APPDATA%\npm on Windows; an old claude there (not on PATH) must still
+ # be version-checked so the unsupported flags are dropped.
+ _simulate_windows(monkeypatch)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path) # no ~/.local/bin created
+ npm_dir = tmp_path / "Roaming" / "npm"
+ npm_dir.mkdir(parents = True)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "Roaming"))
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ monkeypatch.setattr(start.shutil, "which", _path_aware_which({"claude": npm_dir}))
+ monkeypatch.setattr(
+ start.subprocess, "run", lambda *a, **k: SimpleNamespace(stdout = "2.0.14 (Claude Code)\n")
+ )
+ assert start._claude_flags(MODEL["id"]) == []
+
+
+def test_codex_catalog_probes_old_codex_only_in_install_dir(monkeypatch, tmp_path):
+ # Same ordering fix for codex: an old codex present only in an install dir is detected so the
+ # model-catalog config is omitted (the old binary can't consume it).
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata"))
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ monkeypatch.setattr(start.shutil, "which", _path_aware_which({"codex": local_bin}))
+ monkeypatch.setattr(start.subprocess, "check_output", lambda *a, **k: "codex-cli 0.109.0")
+ assert start._codex_supports_model_catalog() is False
+
+
+def test_opencode_native_auto_probes_old_opencode_only_in_install_dir(monkeypatch, tmp_path):
+ # Same ordering fix for opencode: an old opencode present only in an install dir is detected
+ # so native --auto is not assumed (the old binary rejects it).
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata"))
+ monkeypatch.setenv("PATH", str(tmp_path / "existing"))
+ monkeypatch.setattr(start.shutil, "which", _path_aware_which({"opencode": local_bin}))
+ monkeypatch.setattr(start.subprocess, "check_output", lambda *a, **k: "1.17.11")
+ assert start._opencode_supports_native_auto() is False
+
+
+def test_augment_path_preserves_defpath_when_path_unset(monkeypatch, tmp_path):
+ # PATH unset: shutil.which() and exec*p* fall back to os.defpath (e.g. /bin:/usr/bin), so the
+ # augmentation must keep those default dirs instead of collapsing to just the install dir
+ # (which would hide a system-installed agent and strip the launched child's normal PATH).
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata"))
+ monkeypatch.delenv("PATH", raising = False)
+ start._augment_path_with_install_dirs()
+ entries = os.environ["PATH"].split(os.pathsep)
+ for default_dir in os.defpath.split(os.pathsep):
+ if default_dir:
+ assert default_dir in entries
+ assert str(local_bin) in entries
+
+
+def test_which_with_install_dirs_keeps_defpath_when_path_unset(monkeypatch, tmp_path):
+ # With PATH unset, a system agent on os.defpath (e.g. /usr/bin) must still resolve; the
+ # install-dir augmentation must not drop the default search path. PATH is restored to unset.
+ local_bin = tmp_path / ".local" / "bin"
+ local_bin.mkdir(parents = True)
+ monkeypatch.setattr(start.Path, "home", lambda: tmp_path)
+ monkeypatch.setenv("APPDATA", str(tmp_path / "no-appdata"))
+ monkeypatch.delenv("PATH", raising = False)
+ sysdir = next(part for part in reversed(os.defpath.split(os.pathsep)) if part)
+ monkeypatch.setattr(start.shutil, "which", _path_aware_which({"claude": Path(sysdir)}))
+ assert start._which_with_install_dirs("claude") == os.path.join(sysdir, "claude")
+ assert "PATH" not in os.environ
+
+
def test_install_agent_declined_returns_none(monkeypatch):
# TTY + no: never runs anything; caller falls back to the print-hint failure.
monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
@@ -451,7 +699,7 @@ def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypa
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale")
monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale")
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
- monkeypatch.setattr(start, "_claude_flags", lambda: [])
+ monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
def run(command, env):
captured["command"] = command
@@ -486,7 +734,7 @@ def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypat
monkeypatch.setattr(
start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude"
)
- monkeypatch.setattr(start, "_claude_flags", lambda: [])
+ monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
def run(command, env):
captured["command"] = command
@@ -2175,8 +2423,9 @@ def test_powershell_quote_single_quotes_json():
# keeps the embedded double quotes literal (list2cmdline's backslashes would not).
assert start._powershell_quote("--settings") == "--settings"
assert start._powershell_quote("unsloth/gemma-4-26B") == "unsloth/gemma-4-26B"
- quoted = start._powershell_quote(start._CLAUDE_SETTINGS_OVERLAY)
- assert quoted == "'" + start._CLAUDE_SETTINGS_OVERLAY + "'"
+ overlay = start._claude_settings_overlay("unsloth/gemma-4-26B")
+ quoted = start._powershell_quote(overlay)
+ assert quoted == "'" + overlay + "'"
assert "\\" not in quoted # no cmd.exe backslash escaping
assert start._powershell_quote("a'b") == "'a''b'" # embedded quote doubled
@@ -2811,7 +3060,7 @@ def test_claude_launch_does_not_clear(fake_studio, monkeypatch):
calls = []
monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear"))
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
- monkeypatch.setattr(start, "_claude_flags", lambda: [])
+ monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0))
result = CliRunner().invoke(start.start_app, ["claude"])
assert result.exit_code == 0, result.output
@@ -2988,7 +3237,7 @@ def test_persist_bare_opencode_launch_has_no_resume_token(fake_studio, monkeypat
def test_persist_bare_claude_launch_has_no_resume_token(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
- monkeypatch.setattr(start, "_claude_flags", lambda: [])
+ monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
captured = _capture_launch(monkeypatch, ["claude", "--persist"])
assert "--continue" not in captured["command"]
assert captured["command"][1:] == ["--model", MODEL["id"]]
@@ -3152,7 +3401,7 @@ def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch):
# `--resume ` (e.g. `unsloth start claude --resume `) still flows
# through to the agent verbatim and is not swallowed as an Unsloth option.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
- monkeypatch.setattr(start, "_claude_flags", lambda: [])
+ monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"])
assert captured["command"][-2:] == ["--resume", "some-session-guid"]
# Unsloth never auto-appends its own resume token when the user drives resume.
From e092895e01bc90d21b6b0af3e54440978266868e Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 20 Jul 2026 14:17:45 +0100
Subject: [PATCH 038/255] Fix On Device model picker startup ordering (#6994)
* Fix on-device model picker startup ordering
* Fix on-device picker cached local remounts
* fix(studio): stabilize on-device picker readiness
* fix(studio): prevent stale picker refreshes
* fix(studio): retry incomplete picker scans
* fix(studio): preserve replacement picker readiness
* fix(studio): preserve slow local scans
---------
Co-authored-by: Long Yixing
---
.../assistant-ui/model-selector/pickers.tsx | 285 ++++++++++++++----
.../src/features/chat/api/chat-api.ts | 14 +-
2 files changed, 239 insertions(+), 60 deletions(-)
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index 23bf68c042..766139e2b4 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -1170,6 +1170,17 @@ let _lmStudioCache: LocalModelInfo[] = [];
let _localDirCache: LocalModelInfo[] = [];
let _customFolderCache: LocalModelInfo[] = [];
let _scanFoldersCache: ScanFolderInfo[] = [];
+let _onDeviceCachesReady = false;
+let _cachedGgufRequestVersion = 0;
+let _cachedModelsRequestVersion = 0;
+let _localModelsRequestVersion = 0;
+const _onDeviceCacheListeners = new Set<(settled?: boolean) => void>();
+
+const ON_DEVICE_CACHE_TIMEOUT_MS = 30_000;
+
+function notifyOnDeviceCachesChanged(settled = false): void {
+ for (const listener of _onDeviceCacheListeners) listener(settled);
+}
/** True when any on-device model (downloaded GGUF, cached repo, LM Studio, or
* custom-folder model) is known. Reads the module caches, which persist across
@@ -1569,8 +1580,7 @@ export function HubModelPicker({
useState(_cachedGgufCache);
const [cachedModels, setCachedModels] =
useState(_cachedModelsCache);
- const alreadyCached =
- _cachedGgufCache.length > 0 || _cachedModelsCache.length > 0;
+ const alreadyCached = _onDeviceCachesReady || hasDownloadedModels();
const [cachedReady, setCachedReady] = useState(alreadyCached);
const [updateConflictKey, setUpdateConflictKey] = useState(
null,
@@ -1605,6 +1615,25 @@ export function HubModelPicker({
const [customFolderModels, setCustomFolderModels] =
useState(_customFolderCache);
+ useEffect(() => {
+ const syncModuleCaches = (settled = false) => {
+ setCachedGguf(_cachedGgufCache);
+ setCachedModels(_cachedModelsCache);
+ setLmStudioModels(_lmStudioCache);
+ setLocalDirModels(_localDirCache);
+ setCustomFolderModels(_customFolderCache);
+ setCachedReady(
+ (ready) =>
+ ready || settled || _onDeviceCachesReady || hasDownloadedModels(),
+ );
+ };
+ _onDeviceCacheListeners.add(syncModuleCaches);
+ syncModuleCaches();
+ return () => {
+ _onDeviceCacheListeners.delete(syncModuleCaches);
+ };
+ }, []);
+
// Custom scan folders management
const [scanFolders, setScanFolders] =
useState(_scanFoldersCache);
@@ -1615,23 +1644,94 @@ export function HubModelPicker({
const [showFolderBrowser, setShowFolderBrowser] = useState(false);
const [recommendedFolders, setRecommendedFolders] = useState([]);
+ const applyLocalModels = useCallback(
+ (res: Awaited>) => {
+ const lm = sortLmStudio(
+ res.models.filter((m) => m.source === "lmstudio"),
+ );
+ _lmStudioCache = lm;
+ setLmStudioModels(lm);
+ const ld = res.models.filter((m) => m.source === "models_dir");
+ _localDirCache = ld;
+ setLocalDirModels(ld);
+ const cf = res.models.filter((m) => m.source === "custom");
+ _customFolderCache = cf;
+ setCustomFolderModels(cf);
+ notifyOnDeviceCachesChanged();
+ },
+ [],
+ );
+
+ const refreshColdOnDeviceCaches = useCallback(() => {
+ const ggufRequestVersion = ++_cachedGgufRequestVersion;
+ const modelsRequestVersion = ++_cachedModelsRequestVersion;
+ const localRequestVersion = ++_localModelsRequestVersion;
+ let ggufResult: Awaited> | undefined;
+ let modelsResult: Awaited> | undefined;
+ let localResult: Awaited> | undefined;
+ let released = false;
+
+ const ggufRequest = listCachedGguf().then(
+ (value) => { if (!released) ggufResult = value; },
+ () => {},
+ );
+ const modelsRequest = listCachedModels(hfToken || undefined).then(
+ (value) => { if (!released) modelsResult = value; },
+ () => {},
+ );
+ const localRequest = listLocalModels().then(
+ (value) => {
+ localResult = value;
+ if (released && localRequestVersion === _localModelsRequestVersion) {
+ if (ggufResult !== undefined && modelsResult !== undefined) _onDeviceCachesReady = true;
+ applyLocalModels(value);
+ }
+ },
+ () => {},
+ );
+ const isCurrent = () =>
+ ggufRequestVersion === _cachedGgufRequestVersion &&
+ modelsRequestVersion === _cachedModelsRequestVersion &&
+ localRequestVersion === _localModelsRequestVersion;
+ const publish = (invalidate = false) => {
+ if (!isCurrent()) return;
+ if (invalidate) {
+ released = true;
+ ++_cachedGgufRequestVersion;
+ ++_cachedModelsRequestVersion;
+ }
+ if (ggufResult !== undefined) {
+ _cachedGgufCache = ggufResult;
+ setCachedGguf(ggufResult);
+ }
+ if (modelsResult !== undefined) {
+ _cachedModelsCache = modelsResult;
+ setCachedModels(modelsResult);
+ }
+ if (localResult !== undefined) applyLocalModels(localResult);
+ if (ggufResult !== undefined && modelsResult !== undefined && localResult !== undefined) {
+ _onDeviceCachesReady = true;
+ }
+ notifyOnDeviceCachesChanged(true);
+ };
+ const timeout = window.setTimeout(() => publish(true), ON_DEVICE_CACHE_TIMEOUT_MS);
+ void Promise.all([ggufRequest, modelsRequest, localRequest]).then(() => {
+ window.clearTimeout(timeout);
+ publish();
+ });
+ }, [applyLocalModels, hfToken]);
+
const refreshLocalModelsList = useCallback(() => {
+ if (!_onDeviceCachesReady && !hasDownloadedModels()) return refreshColdOnDeviceCaches();
+ const requestVersion = ++_localModelsRequestVersion;
listLocalModels()
.then((res) => {
- const lm = sortLmStudio(
- res.models.filter((m) => m.source === "lmstudio"),
- );
- _lmStudioCache = lm;
- setLmStudioModels(lm);
- const ld = res.models.filter((m) => m.source === "models_dir");
- _localDirCache = ld;
- setLocalDirModels(ld);
- const cf = res.models.filter((m) => m.source === "custom");
- _customFolderCache = cf;
- setCustomFolderModels(cf);
+ if (requestVersion === _localModelsRequestVersion) {
+ applyLocalModels(res);
+ }
})
.catch(() => {});
- }, []);
+ }, [applyLocalModels, refreshColdOnDeviceCaches]);
const refreshScanFolders = useCallback(() => {
listScanFolders()
@@ -1710,20 +1810,27 @@ export function HubModelPicker({
);
const refreshCachedLists = useCallback(() => {
+ if (!_onDeviceCachesReady && !hasDownloadedModels()) return refreshColdOnDeviceCaches();
+ const ggufRequestVersion = ++_cachedGgufRequestVersion;
listCachedGguf()
.then((v) => {
+ if (ggufRequestVersion !== _cachedGgufRequestVersion) return;
_cachedGgufCache = v;
setCachedGguf(v);
+ notifyOnDeviceCachesChanged();
})
.catch(() => {});
+ const modelsRequestVersion = ++_cachedModelsRequestVersion;
listCachedModels(hfToken || undefined)
.then((v) => {
+ if (modelsRequestVersion !== _cachedModelsRequestVersion) return;
_cachedModelsCache = v;
setCachedModels(v);
+ notifyOnDeviceCachesChanged();
})
.catch(() => {});
refreshLocalModelsList();
- }, [hfToken, refreshLocalModelsList]);
+ }, [hfToken, refreshColdOnDeviceCaches, refreshLocalModelsList]);
// Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking
// call. The worker pulls only changed blobs, so the cached copy stays usable until done.
@@ -1751,36 +1858,100 @@ export function HubModelPicker({
);
useEffect(() => {
- // Always refresh LM Studio + custom folder models (not gated by alreadyCached).
- refreshLocalModelsList();
refreshScanFolders();
listRecommendedFolders()
.then(setRecommendedFolders)
.catch(() => {});
- // Always refetch cached GGUF/model lists. The module-level caches render
- // instantly with stale data (no spinner flash), but newly downloaded
- // repos need a fresh backend hit. cachedReady=alreadyCached initially,
- // so the background refresh is invisible when we already had data.
- let done = 0;
- const check = () => {
- if (++done >= 2) setCachedReady(true);
+ // Publish downloaded and local rows as one bounded snapshot. Existing data
+ // stays visible during background refreshes, and a failed source keeps its
+ // last successful cache instead of clearing or durably marking it ready.
+ const controller = new AbortController();
+ const timeout = window.setTimeout(
+ () => controller.abort(),
+ ON_DEVICE_CACHE_TIMEOUT_MS,
+ );
+ const aborted = new Promise((_, reject) => {
+ controller.signal.addEventListener(
+ "abort",
+ () => reject(controller.signal.reason),
+ { once: true },
+ );
+ });
+ const bounded = (request: Promise) =>
+ Promise.race([request, aborted]);
+ let cancelled = false;
+ const ggufRequestVersion = ++_cachedGgufRequestVersion;
+ const modelsRequestVersion = ++_cachedModelsRequestVersion;
+ const localRequestVersion = ++_localModelsRequestVersion;
+ const localRequest = listLocalModels();
+
+ void Promise.allSettled([
+ bounded(listCachedGguf(controller.signal)),
+ bounded(listCachedModels(hfToken || undefined, controller.signal)),
+ bounded(localRequest),
+ ]).then(([ggufResult, modelsResult, localResult]) => {
+ window.clearTimeout(timeout);
+ if (cancelled) return;
+
+ const ggufIsCurrent =
+ ggufRequestVersion === _cachedGgufRequestVersion;
+ const modelsAreCurrent =
+ modelsRequestVersion === _cachedModelsRequestVersion;
+ const localIsCurrent =
+ localRequestVersion === _localModelsRequestVersion;
+
+ if (ggufResult.status === "fulfilled" && ggufIsCurrent) {
+ _cachedGgufCache = ggufResult.value;
+ setCachedGguf(ggufResult.value);
+ notifyOnDeviceCachesChanged();
+ }
+ if (modelsResult.status === "fulfilled" && modelsAreCurrent) {
+ _cachedModelsCache = modelsResult.value;
+ setCachedModels(modelsResult.value);
+ notifyOnDeviceCachesChanged();
+ }
+ if (localResult.status === "fulfilled" && localIsCurrent) {
+ applyLocalModels(localResult.value);
+ }
+ if (localResult.status === "rejected" && controller.signal.aborted) {
+ void localRequest.then((value) => {
+ if (cancelled || localRequestVersion !== _localModelsRequestVersion) return;
+ if (ggufResult.status === "fulfilled" && modelsResult.status === "fulfilled") _onDeviceCachesReady = true;
+ applyLocalModels(value);
+ }).catch(() => {});
+ }
+ const snapshotIsCurrent =
+ ggufIsCurrent && modelsAreCurrent && localIsCurrent;
+ if (
+ ggufResult.status === "fulfilled" &&
+ modelsResult.status === "fulfilled" &&
+ localResult.status === "fulfilled" &&
+ snapshotIsCurrent
+ ) {
+ _onDeviceCachesReady = true;
+ }
+ notifyOnDeviceCachesChanged(snapshotIsCurrent);
+ });
+
+ return () => {
+ cancelled = true;
+ window.clearTimeout(timeout);
+ controller.abort();
+ queueMicrotask(() => {
+ if (
+ ggufRequestVersion === _cachedGgufRequestVersion &&
+ modelsRequestVersion === _cachedModelsRequestVersion &&
+ localRequestVersion === _localModelsRequestVersion
+ ) {
+ ++_cachedGgufRequestVersion;
+ ++_cachedModelsRequestVersion;
+ ++_localModelsRequestVersion;
+ notifyOnDeviceCachesChanged(true);
+ }
+ });
};
- listCachedGguf()
- .then((v) => {
- _cachedGgufCache = v;
- setCachedGguf(v);
- })
- .catch(() => {})
- .finally(check);
- listCachedModels(hfToken || undefined)
- .then((v) => {
- _cachedModelsCache = v;
- setCachedModels(v);
- })
- .catch(() => {})
- .finally(check);
- }, [hfToken, refreshLocalModelsList, refreshScanFolders]);
+ }, [applyLocalModels, hfToken, refreshScanFolders]);
// Hide downloaded models from the recommended list. Case-insensitive
// since the HF cache lowercases repo IDs.
@@ -2371,12 +2542,12 @@ export function HubModelPicker({
}
// Fine-tuned models sit below downloaded, above custom folders.
- if (section === "downloaded" && !fineTunedCollapsed) {
+ if (section === "downloaded" && cachedReady && !fineTunedCollapsed) {
keys.push(...fineTunedRows.map((m) => makeModelOptionKey("lora", m.id)));
}
// Custom folders sit right below the downloaded models on On Device.
- if (section === "downloaded" && !customFoldersCollapsed) {
+ if (section === "downloaded" && cachedReady && !customFoldersCollapsed) {
keys.push(
...sortedCustomFolderModels.map((model) =>
makeModelOptionKey("custom-folder", model.id),
@@ -2384,7 +2555,7 @@ export function HubModelPicker({
);
}
- if (section === "downloaded" && !lmStudioCollapsed) {
+ if (section === "downloaded" && cachedReady && !lmStudioCollapsed) {
keys.push(
...sortedLmStudio.map((model) =>
makeModelOptionKey("lm-studio", model.id),
@@ -2392,7 +2563,7 @@ export function HubModelPicker({
);
}
- if (section === "downloaded" && !localDirCollapsed) {
+ if (section === "downloaded" && cachedReady && !localDirCollapsed) {
keys.push(
...sortedLocalDir.map((model) =>
makeModelOptionKey("local-dir", model.id),
@@ -2585,6 +2756,7 @@ export function HubModelPicker({
const showDownloaded = section === "downloaded";
const showCustom = section === "downloaded";
const showRecommendedSection = !showHfSection && section === "recommended";
+ const onDeviceCacheLoading = showDownloaded && !cachedReady;
const downloadedEmpty =
visibleCachedGguf.length === 0 &&
visibleCachedModelRows.length === 0 &&
@@ -3135,11 +3307,8 @@ export function HubModelPicker({
)
) : (
<>
- {/* First-load spinner only when nothing cached is shown yet. */}
- {showDownloaded &&
- !cachedReady &&
- !showHfSection &&
- downloadedEmpty ? (
+ {/* First-load spinner while downloaded/local scans are resolving. */}
+ {onDeviceCacheLoading ? (
@@ -3185,6 +3354,7 @@ export function HubModelPicker({
{/* Downloaded (Unsloth) stays visible (filtered) while searching. */}
{showDownloaded &&
+ cachedReady &&
(unslothCachedGguf.length > 0 ||
unslothCachedModelRows.length > 0) ? (
<>
@@ -3278,7 +3448,7 @@ export function HubModelPicker({
{/* Other models: non-Unsloth downloads, grouped just above
Fine-tuned. Shown only when such models exist. */}
- {showDownloaded && hasOtherModels ? (
+ {showDownloaded && cachedReady && hasOtherModels ? (
) : null}
- {/* Fine-tuned models: a section above Custom Folders. Always shown on
- On Device so the train shortcut always has a target, with an empty
- state when none exist. */}
- {section === "downloaded" ? (
+ {/* Fine-tuned models: shown after the On Device scans resolve so
+ downloaded sections do not reorder during startup. */}
+ {section === "downloaded" && cachedReady ? (
<>
) : null}
- {showCustom ? (
+ {showCustom && cachedReady ? (
<>
) : null}
- {section === "downloaded" && sortedLmStudio.length > 0 ? (
+ {section === "downloaded" &&
+ cachedReady &&
+ sortedLmStudio.length > 0 ? (
<>
) : null}
- {section === "downloaded" && sortedLocalDir.length > 0 ? (
+ {section === "downloaded" &&
+ cachedReady &&
+ sortedLocalDir.length > 0 ? (
<>
{
- const response = await authFetch("/api/models/local");
+export async function listLocalModels(
+ signal?: AbortSignal,
+): Promise {
+ const response = await authFetch("/api/models/local", { signal });
return parseJsonOrThrow(response);
}
-export async function listCachedGguf(): Promise {
- const response = await authFetch("/api/models/cached-gguf");
+export async function listCachedGguf(
+ signal?: AbortSignal,
+): Promise {
+ const response = await authFetch("/api/models/cached-gguf", { signal });
const data = await parseJsonOrThrow<{ cached: CachedGgufRepo[] }>(response);
return data.cached;
}
@@ -349,9 +353,11 @@ export interface CachedModelRepo {
export async function listCachedModels(
hfToken?: string | null,
+ signal?: AbortSignal,
): Promise {
const response = await authFetch("/api/models/cached-models", {
headers: hubTokenHeader(hfToken),
+ signal,
});
const data = await parseJsonOrThrow<{ cached: CachedModelRepo[] }>(response);
return data.cached;
From 1b3bce0530a92a6e7aa934c4086049237b74dfe3 Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Mon, 20 Jul 2026 10:40:14 -0300
Subject: [PATCH 039/255] Studio: validate Hugging Face tokens before use
(#7261)
* Studio: validate Hugging Face tokens before use
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep token validation failures non-blocking
* Studio: harden Hugging Face token preflight
* Studio: make token validation effect lint-safe
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
studio/backend/hub/routes/__init__.py | 2 +
studio/backend/hub/routes/token.py | 44 ++++
studio/backend/main.py | 2 +
studio/backend/routes/inference.py | 10 +-
.../backend/tests/test_hf_token_validation.py | 165 ++++++++++++++
studio/backend/tests/test_utils.py | 16 +-
studio/backend/utils/hf_token_validation.py | 208 ++++++++++++++++++
studio/backend/utils/utils.py | 20 ++
studio/frontend/src/app/routes/__root.tsx | 2 +
.../src/features/chat/api/chat-api.ts | 8 +-
.../src/features/export/export-page.tsx | 13 +-
studio/frontend/src/features/hf-auth/api.ts | 44 ++++
.../src/features/hf-auth/confirm-token.ts | 63 ++++++
.../hf-auth/hf-token-warning-dialog.tsx | 66 ++++++
studio/frontend/src/features/hf-auth/index.ts | 10 +
studio/frontend/src/features/hf-auth/store.ts | 33 +++
.../features/settings/tabs/general-tab.tsx | 108 +++++----
.../src/features/training/api/train-api.ts | 5 +-
.../training/hooks/use-training-actions.ts | 17 +-
.../src/hooks/use-hf-token-validation.ts | 110 ++++++---
20 files changed, 869 insertions(+), 77 deletions(-)
create mode 100644 studio/backend/hub/routes/token.py
create mode 100644 studio/backend/tests/test_hf_token_validation.py
create mode 100644 studio/backend/utils/hf_token_validation.py
create mode 100644 studio/frontend/src/features/hf-auth/api.ts
create mode 100644 studio/frontend/src/features/hf-auth/confirm-token.ts
create mode 100644 studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx
create mode 100644 studio/frontend/src/features/hf-auth/index.ts
create mode 100644 studio/frontend/src/features/hf-auth/store.ts
diff --git a/studio/backend/hub/routes/__init__.py b/studio/backend/hub/routes/__init__.py
index e9579635b0..7c5cfb9b3c 100644
--- a/studio/backend/hub/routes/__init__.py
+++ b/studio/backend/hub/routes/__init__.py
@@ -5,8 +5,10 @@
from hub.routes.inventory import router as inventory_router
from hub.routes.datasets import router as datasets_router
+from hub.routes.token import router as token_router
__all__ = [
"inventory_router",
"datasets_router",
+ "token_router",
]
diff --git a/studio/backend/hub/routes/token.py b/studio/backend/hub/routes/token.py
new file mode 100644
index 0000000000..1b7ad733a2
--- /dev/null
+++ b/studio/backend/hub/routes/token.py
@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Hugging Face token validation endpoint."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Literal, Optional
+
+from fastapi import APIRouter, Depends, Request
+from pydantic import BaseModel
+
+from auth.authentication import get_current_subject
+from hub.dependencies import get_hf_token
+from utils.client_ip import client_ip
+from utils.hf_token_validation import validate_hf_token
+
+
+router = APIRouter()
+
+
+class HfTokenValidationResponse(BaseModel):
+ status: Literal["missing", "valid", "invalid", "rate_limited", "unavailable"]
+ retry_after_seconds: Optional[int] = None
+
+
+@router.post("/token/validate", response_model = HfTokenValidationResponse)
+async def validate_token(
+ request: Request,
+ hf_token: Optional[str] = Depends(get_hf_token),
+ current_subject: str = Depends(get_current_subject),
+):
+ if not hf_token:
+ return HfTokenValidationResponse(status = "missing")
+ result = await asyncio.to_thread(
+ validate_hf_token,
+ hf_token,
+ rate_key = f"{current_subject}:{client_ip(request)}",
+ )
+ return HfTokenValidationResponse(
+ status = result.status,
+ retry_after_seconds = result.retry_after_seconds,
+ )
diff --git a/studio/backend/main.py b/studio/backend/main.py
index 0ffc4489a1..f686e29bf5 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -312,6 +312,7 @@ from routes.preview import router as preview_router
from hub.routes import (
inventory_router as hub_inventory_router,
datasets_router as hub_datasets_router,
+ token_router as hub_token_router,
)
from hub.schemas.downloads import TransportCapabilities
from hub.utils.download_registry import (
@@ -993,6 +994,7 @@ app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])
app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"])
app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"])
+app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"])
# Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic
# error envelopes; non-/v1 paths keep FastAPI's default {"detail": ...} shape.
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 9c08ea4b79..afd942e9a5 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -1778,7 +1778,7 @@ from core.inference.providers import get_base_url
from core.inference.external_provider import ExternalProviderClient
from core.inference.chat_templates import resolve_effective_chat_template_override
from storage import providers_db
-from utils.utils import safe_error_detail, log_and_http_error
+from utils.utils import is_hf_authentication_error, safe_error_detail, log_and_http_error
import io
import base64
@@ -5244,6 +5244,14 @@ async def validate_model(
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
redacted_msg = redact_native_paths(str(e))
+ if is_hf_authentication_error(e):
+ raise HTTPException(
+ status_code = 400,
+ detail = (
+ "Hugging Face authentication failed. Check or clear the token "
+ "in Settings, and confirm access to this gated repository."
+ ),
+ )
if _is_unsupported_nvfp4_inference_error(redacted_msg):
logger.warning(
"NVFP4 inference is not supported yet while validating '%s'",
diff --git a/studio/backend/tests/test_hf_token_validation.py b/studio/backend/tests/test_hf_token_validation.py
new file mode 100644
index 0000000000..31b30fc37d
--- /dev/null
+++ b/studio/backend/tests/test_hf_token_validation.py
@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Focused coverage for cached, rate-limited HF token validation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import sys
+
+import httpx
+import pytest
+
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+ sys.path.insert(0, _BACKEND_DIR)
+
+import utils.hf_token_validation as validation
+
+
+@pytest.fixture(autouse = True)
+def _reset_validation_state():
+ validation.reset_hf_token_validation_state()
+ yield
+ validation.reset_hf_token_validation_state()
+
+
+def test_cached_token_does_not_spend_another_attempt(monkeypatch):
+ calls = []
+
+ def _check(token):
+ calls.append(token)
+ return validation.TokenValidationResult(status = "valid")
+
+ monkeypatch.setattr(validation, "_check_remote", _check)
+ first = validation.validate_hf_token("hf_valid", rate_key = "user:ip")
+ second = validation.validate_hf_token("hf_valid", rate_key = "user:ip")
+
+ assert first.status == second.status == "valid"
+ assert calls == ["hf_valid"]
+
+
+def test_three_uncached_attempts_per_hour(monkeypatch):
+ monkeypatch.setattr(
+ validation,
+ "_check_remote",
+ lambda _token: validation.TokenValidationResult(status = "invalid"),
+ )
+
+ for index in range(3):
+ result = validation.validate_hf_token(f"hf_bad_{index}", rate_key = "user:ip")
+ assert result.status == "invalid"
+
+ limited = validation.validate_hf_token("hf_bad_4", rate_key = "user:ip")
+ assert limited.status == "rate_limited"
+ assert limited.retry_after_seconds is not None
+ assert limited.retry_after_seconds > 0
+
+ other_user = validation.validate_hf_token("hf_other", rate_key = "other:ip")
+ assert other_user.status == "invalid"
+
+
+def test_window_rolls_forward(monkeypatch):
+ clock = {"now": 100.0}
+ monkeypatch.setattr(validation.time, "monotonic", lambda: clock["now"])
+ monkeypatch.setattr(validation, "_MAX_ATTEMPTS", 1)
+ monkeypatch.setattr(validation, "_WINDOW_SECONDS", 10.0)
+ monkeypatch.setattr(
+ validation,
+ "_check_remote",
+ lambda _token: validation.TokenValidationResult(status = "invalid"),
+ )
+
+ assert validation.validate_hf_token("hf_a", rate_key = "user:ip").status == "invalid"
+ assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "rate_limited"
+ clock["now"] += 11.0
+ assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "invalid"
+
+
+@pytest.mark.parametrize(
+ ("status_code", "expected"),
+ [(200, "valid"), (401, "invalid"), (429, "rate_limited"), (500, "unavailable")],
+)
+def test_remote_status_classification(monkeypatch, status_code, expected):
+ response = httpx.Response(
+ status_code,
+ request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"),
+ headers = {"Retry-After": "42"} if status_code == 429 else None,
+ )
+
+ class _Session:
+ def get(self, url, *, headers, timeout):
+ assert url == "https://huggingface.co/api/whoami-v2"
+ assert headers["authorization"] == "Bearer hf_test"
+ assert timeout == validation._REMOTE_TIMEOUT_SECONDS
+ return response
+
+ monkeypatch.setattr(validation, "get_session", lambda: _Session())
+ result = validation._check_remote("hf_test")
+ assert result.status == expected
+ if status_code == 429:
+ assert result.retry_after_seconds == 42
+
+
+def test_wrapped_http_401_is_invalid(monkeypatch):
+ response = httpx.Response(
+ 401,
+ request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"),
+ )
+
+ class _Session:
+ def get(self, _url, **_kwargs):
+ error = RuntimeError("Invalid user token.")
+ error.response = response
+ raise error
+
+ monkeypatch.setattr(validation, "get_session", lambda: _Session())
+ assert validation._check_remote("hf_test").status == "invalid"
+
+
+def test_remote_timeout_is_bounded_and_unavailable(monkeypatch):
+ class _Session:
+ def get(self, _url, *, headers, timeout):
+ assert headers["authorization"] == "Bearer hf_test"
+ assert timeout == validation._REMOTE_TIMEOUT_SECONDS
+ raise TimeoutError("timed out")
+
+ monkeypatch.setattr(validation, "get_session", lambda: _Session())
+ assert validation._check_remote("hf_test").status == "unavailable"
+
+
+def test_raw_token_is_not_retained(monkeypatch):
+ monkeypatch.setattr(
+ validation,
+ "_check_remote",
+ lambda _token: validation.TokenValidationResult(status = "valid"),
+ )
+ token = "hf_do_not_store_this_value"
+ validation.validate_hf_token(token, rate_key = "user:ip")
+
+ assert token not in repr(validation._cache)
+ assert token not in repr(validation._attempts)
+
+
+def test_unexpected_remote_exception_releases_singleflight(monkeypatch):
+ calls = 0
+ monkeypatch.setattr(validation, "_INFLIGHT_WAIT_SECONDS", 0.0)
+
+ def _check(_token):
+ nonlocal calls
+ calls += 1
+ if calls == 1:
+ raise RuntimeError("unexpected failure")
+ return validation.TokenValidationResult(status = "valid")
+
+ monkeypatch.setattr(validation, "_check_remote", _check)
+
+ with pytest.raises(RuntimeError, match = "unexpected failure"):
+ validation.validate_hf_token("hf_test", rate_key = "user:ip")
+
+ result = validation.validate_hf_token("hf_test", rate_key = "user:ip")
+ assert result.status == "valid"
+ assert calls == 2
+ assert validation._inflight == {}
diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py
index 64a3c62156..741f19c67a 100644
--- a/studio/backend/tests/test_utils.py
+++ b/studio/backend/tests/test_utils.py
@@ -38,7 +38,7 @@ from utils.hardware import (
DeviceType,
)
import utils.hardware.hardware as _hw_module
-from utils.utils import format_error_message
+from utils.utils import format_error_message, is_hf_authentication_error
# ========== Helpers ==========
@@ -439,6 +439,20 @@ class TestFormatErrorMessage:
msg = format_error_message(err, "any/model")
assert "invalid" in msg.lower()
+ def test_hf_authentication_error_follows_wrapped_401(self):
+ response = type("Response", (), {"status_code": 401})()
+ auth_error = Exception("request failed")
+ auth_error.response = response
+ wrapper = RuntimeError("model validation failed")
+ wrapper.__cause__ = auth_error
+ assert is_hf_authentication_error(wrapper) is True
+
+ def test_hf_authentication_error_does_not_treat_429_as_invalid(self):
+ response = type("Response", (), {"status_code": 429})()
+ rate_error = Exception("too many requests")
+ rate_error.response = response
+ assert is_hf_authentication_error(rate_error) is False
+
# --- OOM on CUDA ---
@needs_torch
diff --git a/studio/backend/utils/hf_token_validation.py b/studio/backend/utils/hf_token_validation.py
new file mode 100644
index 0000000000..7247c6e756
--- /dev/null
+++ b/studio/backend/utils/hf_token_validation.py
@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Cached, rate-limited Hugging Face token validation."""
+
+from __future__ import annotations
+
+import hashlib
+import threading
+import time
+from collections import deque
+from dataclasses import dataclass
+from typing import Literal
+
+from huggingface_hub import HfApi
+from huggingface_hub.utils import build_hf_headers, get_session
+
+
+TokenValidationStatus = Literal["valid", "invalid", "rate_limited", "unavailable"]
+
+
+@dataclass(frozen = True)
+class TokenValidationResult:
+ status: TokenValidationStatus
+ retry_after_seconds: int | None = None
+
+
+_WINDOW_SECONDS = 3600.0
+_MAX_ATTEMPTS = 3
+_CACHE_TTL_SECONDS = 3600.0
+_TEMPORARY_CACHE_TTL_SECONDS = 15.0
+_MAX_BUCKETS = 4096
+_MAX_CACHE_ENTRIES = 4096
+_INFLIGHT_WAIT_SECONDS = 30.0
+_REMOTE_TIMEOUT_SECONDS = 10.0
+
+_attempts: dict[str, deque[float]] = {}
+_cache: dict[str, tuple[float, TokenValidationResult]] = {}
+_inflight: dict[str, threading.Event] = {}
+_lock = threading.Lock()
+
+
+def _fingerprint(token: str) -> str:
+ return hashlib.sha256(token.encode("utf-8")).hexdigest()
+
+
+def _prune_attempts(bucket: deque[float], now: float) -> None:
+ while bucket and now - bucket[0] >= _WINDOW_SECONDS:
+ bucket.popleft()
+
+
+def _prune_locked(now: float) -> None:
+ for key in list(_attempts):
+ bucket = _attempts[key]
+ _prune_attempts(bucket, now)
+ if not bucket:
+ del _attempts[key]
+ for key, (expires_at, _result) in list(_cache.items()):
+ if expires_at <= now:
+ del _cache[key]
+
+
+def _cached_locked(fingerprint: str, now: float) -> TokenValidationResult | None:
+ cached = _cache.get(fingerprint)
+ if cached is None:
+ return None
+ expires_at, result = cached
+ if expires_at <= now:
+ del _cache[fingerprint]
+ return None
+ return result
+
+
+def _retry_after(bucket: deque[float], now: float) -> int:
+ return max(1, int(_WINDOW_SECONDS - (now - bucket[0])) + 1)
+
+
+def _reserve_attempt_locked(rate_key: str, now: float) -> TokenValidationResult | None:
+ bucket = _attempts.get(rate_key)
+ if bucket is None:
+ if len(_attempts) >= _MAX_BUCKETS:
+ _prune_locked(now)
+ if len(_attempts) >= _MAX_BUCKETS:
+ return TokenValidationResult(
+ status = "rate_limited",
+ retry_after_seconds = max(1, int(_WINDOW_SECONDS)),
+ )
+ bucket = _attempts[rate_key] = deque()
+ _prune_attempts(bucket, now)
+ if len(bucket) >= _MAX_ATTEMPTS:
+ return TokenValidationResult(
+ status = "rate_limited",
+ retry_after_seconds = _retry_after(bucket, now),
+ )
+ bucket.append(now)
+ return None
+
+
+def _http_status(response: object | None) -> int | None:
+ status = getattr(response, "status_code", None)
+ try:
+ return int(status) if status is not None else None
+ except (TypeError, ValueError):
+ return None
+
+
+def _remote_retry_after(response: object | None) -> int | None:
+ headers = getattr(response, "headers", None)
+ if not headers:
+ return None
+ raw = headers.get("Retry-After")
+ try:
+ return max(1, int(float(raw))) if raw is not None else None
+ except (TypeError, ValueError):
+ return None
+
+
+def _classify_response(response: object | None) -> TokenValidationResult:
+ status = _http_status(response)
+ if status is not None and 200 <= status < 300:
+ return TokenValidationResult(status = "valid")
+ if status == 401:
+ return TokenValidationResult(status = "invalid")
+ if status == 429:
+ return TokenValidationResult(
+ status = "rate_limited",
+ retry_after_seconds = _remote_retry_after(response),
+ )
+ return TokenValidationResult(status = "unavailable")
+
+
+def _check_remote(token: str) -> TokenValidationResult:
+ api = HfApi()
+ try:
+ # HfApi.whoami has no timeout parameter in the pinned Hub client.
+ # Use its session and headers against the same whoami endpoint.
+ response = get_session().get(
+ f"{api.endpoint}/api/whoami-v2",
+ headers = build_hf_headers(token = token),
+ timeout = _REMOTE_TIMEOUT_SECONDS,
+ )
+ except Exception as exc:
+ # huggingface-hub 0.36.x can wrap a 401 as requests.HTTPError.
+ return _classify_response(getattr(exc, "response", None))
+ return _classify_response(response)
+
+
+def validate_hf_token(token: str, *, rate_key: str) -> TokenValidationResult:
+ """Validate ``token`` without retaining it, sharing results across callers.
+
+ Cached checks do not consume the caller's three-per-hour network budget. A
+ single-flight event also prevents simultaneously mounted UI surfaces from
+ sending duplicate ``whoami`` requests for the same token.
+ """
+ normalized = token.strip()
+ if not normalized:
+ return TokenValidationResult(status = "invalid")
+ token_fingerprint = _fingerprint(normalized)
+ owner_event: threading.Event | None = None
+
+ try:
+ while True:
+ now = time.monotonic()
+ with _lock:
+ cached = _cached_locked(token_fingerprint, now)
+ if cached is not None:
+ return cached
+ waiting = _inflight.get(token_fingerprint)
+ if waiting is None:
+ limited = _reserve_attempt_locked(rate_key, now)
+ if limited is not None:
+ return limited
+ owner_event = threading.Event()
+ _inflight[token_fingerprint] = owner_event
+ break
+ if not waiting.wait(_INFLIGHT_WAIT_SECONDS):
+ return TokenValidationResult(status = "unavailable")
+
+ result = _check_remote(normalized)
+ now = time.monotonic()
+ ttl = (
+ _CACHE_TTL_SECONDS
+ if result.status in ("valid", "invalid")
+ else max(_TEMPORARY_CACHE_TTL_SECONDS, float(result.retry_after_seconds or 0))
+ )
+ with _lock:
+ if len(_cache) >= _MAX_CACHE_ENTRIES:
+ _prune_locked(now)
+ if len(_cache) < _MAX_CACHE_ENTRIES:
+ _cache[token_fingerprint] = (now + ttl, result)
+ return result
+ finally:
+ if owner_event is not None:
+ with _lock:
+ event = _inflight.get(token_fingerprint)
+ if event is owner_event:
+ _inflight.pop(token_fingerprint, None)
+ event.set()
+
+
+def reset_hf_token_validation_state() -> None:
+ """Clear process state for test isolation."""
+ with _lock:
+ for event in _inflight.values():
+ event.set()
+ _inflight.clear()
+ _attempts.clear()
+ _cache.clear()
diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py
index 3818253ac9..31f5f31bee 100644
--- a/studio/backend/utils/utils.py
+++ b/studio/backend/utils/utils.py
@@ -123,6 +123,26 @@ def without_hf_auth():
os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None)
+def is_hf_authentication_error(error: Exception) -> bool:
+ """Return whether an exception chain contains a definitive HF auth failure."""
+ seen: set[int] = set()
+ current: BaseException | None = error
+ while current is not None and id(current) not in seen:
+ seen.add(id(current))
+ response = getattr(current, "response", None)
+ status = getattr(response, "status_code", None)
+ try:
+ if status is not None and int(status) == 401:
+ return True
+ except (TypeError, ValueError):
+ pass
+ message = str(current).lower()
+ if "invalid user token" in message or "invalid hf token" in message:
+ return True
+ current = current.__cause__ or current.__context__
+ return False
+
+
def format_error_message(error: Exception, model_name: str) -> str:
"""
Format a user-friendly error message for common load issues.
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index ba56ce7525..e23892e020 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -16,6 +16,7 @@ import {
type ChatSearch,
} from "@/features/chat";
import { RemoteCodeConsentDialog } from "@/features/security";
+import { HfTokenWarningDialog } from "@/features/hf-auth";
import { TransformersUpgradeDialog } from "@/features/transformers-upgrade";
import { useTrainingUnloadGuard } from "@/features/training";
import { useExportRuntimeLifecycle } from "@/features/export";
@@ -230,6 +231,7 @@ function RootLayout() {
{!isAuthFlowRoute && }
+
{hideNavbar ? (
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index d9d0493a03..475bd0f801 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
+import { prepareHfTokenForUse } from "@/features/hf-auth";
// These helpers are deliberately API-layer-only and are not part of their
// features' React-facing public barrels.
// eslint-disable-next-line no-restricted-imports
@@ -108,11 +109,14 @@ export async function getApiMonitorEntry(id: string): Promise {
export async function loadModel(
payload: LoadModelRequest,
): Promise {
+ const preparedToken = await prepareHfTokenForUse(payload.hf_token);
+ if (!preparedToken.proceed) throw new Error("Model load cancelled.");
const response = await authFetch("/api/inference/load", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
...payload,
+ hf_token: preparedToken.token,
native_path_lease: payload.nativePathLease ?? null,
nativePathLease: undefined,
}),
@@ -123,13 +127,15 @@ export async function loadModel(
export async function validateModel(
payload: LoadModelRequest,
): Promise {
+ const preparedToken = await prepareHfTokenForUse(payload.hf_token);
+ if (!preparedToken.proceed) throw new Error("Model load cancelled.");
const response = await authFetch("/api/inference/validate", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
model_path: payload.model_path,
native_path_lease: payload.nativePathLease ?? null,
- hf_token: payload.hf_token,
+ hf_token: preparedToken.token,
gguf_variant: payload.gguf_variant ?? null,
// Intended load settings so validate's preflight matches the follow-up
// /load. Default placement is sized against the selected GPUs.
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx
index 4cd7958fb8..51e9f209dd 100644
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@@ -44,6 +44,7 @@ import {
import { usePlatformStore } from "@/config/env";
import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
+import { prepareHfTokenForUse } from "@/features/hf-auth";
import { GuidedTour, useGuidedTourController } from "@/features/tour";
import {
type LocalModelInfo,
@@ -724,11 +725,17 @@ export function ExportPage() {
const checkpointPath = selectedCp?.path ?? null;
const pushToHub = destination === "hub";
+ const preparedToken = await prepareHfTokenForUse(hfToken, {
+ allowAnonymous: !pushToHub,
+ });
+ if (!preparedToken.proceed) return;
+ const actionHfToken = preparedToken.token ?? "";
+
const repoId =
pushToHub && hfUsername && modelName
? `${hfUsername}/${modelName}`
: undefined;
- const token = pushToHub && hfToken ? hfToken : undefined;
+ const token = pushToHub && actionHfToken ? actionHfToken : undefined;
// The GGUF method with the LoRA target reuses the LoRA-adapter export path.
const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod;
const emitLoraGguf =
@@ -747,7 +754,7 @@ export function ExportPage() {
if (sourceMode !== "checkpoint") {
const remoteCodeOk = await confirmRemoteCodeIfNeeded({
modelName: source,
- hfToken: hfToken || null,
+ hfToken: actionHfToken || null,
// An HF source can need trust_remote_code via its YAML default with no
// auto_map to review; signal it so a YAML-only model does not export
// with it false.
@@ -767,7 +774,7 @@ export function ExportPage() {
modelSource,
trustRemoteCode,
approvedRemoteCodeFingerprint,
- loadToken: hfToken || null,
+ loadToken: actionHfToken || null,
exportMethod: effectiveMethod,
isAdapter: adapterExport,
quantLevels,
diff --git a/studio/frontend/src/features/hf-auth/api.ts b/studio/frontend/src/features/hf-auth/api.ts
new file mode 100644
index 0000000000..ab4b049566
--- /dev/null
+++ b/studio/frontend/src/features/hf-auth/api.ts
@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { authFetch } from "@/features/auth";
+// This header helper is API-layer-only and is not part of the feature's
+// React-facing public barrel.
+// eslint-disable-next-line no-restricted-imports
+import { hubTokenHeader } from "@/features/hub/lib/hub-token-header";
+
+export type HfTokenValidationStatus =
+ | "missing"
+ | "valid"
+ | "invalid"
+ | "rate_limited"
+ | "unavailable";
+
+export interface HfTokenValidationResult {
+ status: HfTokenValidationStatus;
+ retryAfterSeconds: number | null;
+}
+
+export async function validateHfToken(
+ token: string | null | undefined,
+): Promise {
+ const normalized = token?.trim() ?? "";
+ if (!normalized) {
+ return { status: "missing", retryAfterSeconds: null };
+ }
+ const response = await authFetch("/api/hub/token/validate", {
+ method: "POST",
+ headers: hubTokenHeader(normalized),
+ });
+ if (!response.ok) {
+ return { status: "unavailable", retryAfterSeconds: null };
+ }
+ const body = (await response.json()) as {
+ status?: HfTokenValidationStatus;
+ retry_after_seconds?: number | null;
+ };
+ return {
+ status: body.status ?? "unavailable",
+ retryAfterSeconds: body.retry_after_seconds ?? null,
+ };
+}
diff --git a/studio/frontend/src/features/hf-auth/confirm-token.ts b/studio/frontend/src/features/hf-auth/confirm-token.ts
new file mode 100644
index 0000000000..e1f6e6c705
--- /dev/null
+++ b/studio/frontend/src/features/hf-auth/confirm-token.ts
@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// These stores are used outside React and are not part of their features'
+// React-facing public barrels.
+// eslint-disable-next-line no-restricted-imports
+import { useHfTokenStore } from "@/features/hub/stores/hf-token-store";
+// eslint-disable-next-line no-restricted-imports
+import { useSettingsDialogStore } from "@/features/settings/stores/settings-dialog-store";
+import { validateHfToken } from "./api";
+import { useHfTokenWarningStore } from "./store";
+
+export interface PreparedHfToken {
+ proceed: boolean;
+ token: string | null;
+}
+
+interface PrepareHfTokenOptions {
+ allowAnonymous?: boolean;
+}
+
+// A caller can retain the pre-dialog payload while the shared store is cleared.
+// Remember that one-session choice so a follow-up /load does not prompt again
+// after its preceding /validate already continued anonymously.
+const anonymousForSession = new Set();
+
+export async function prepareHfTokenForUse(
+ token: string | null | undefined,
+ options: PrepareHfTokenOptions = {},
+): Promise {
+ const normalized = token?.trim() ?? "";
+ if (!normalized) return { proceed: true, token: null };
+ const allowAnonymous = options.allowAnonymous ?? true;
+ if (allowAnonymous && anonymousForSession.has(normalized)) {
+ return { proceed: true, token: null };
+ }
+
+ let validation;
+ try {
+ validation = await validateHfToken(normalized);
+ } catch {
+ // Validation is advisory. Let the real operation retain its own error.
+ return { proceed: true, token: normalized };
+ }
+ if (validation.status !== "invalid") {
+ // A connectivity failure or rate limit cannot prove that a token is bad.
+ // Let the real operation proceed and retain its repository-specific error.
+ return { proceed: true, token: normalized };
+ }
+
+ const decision = await useHfTokenWarningStore
+ .getState()
+ .requestDecision(allowAnonymous);
+ if (decision === "anonymous") {
+ anonymousForSession.add(normalized);
+ useHfTokenStore.getState().clearToken();
+ return { proceed: true, token: null };
+ }
+ if (decision === "replace") {
+ useSettingsDialogStore.getState().openDialog("general");
+ }
+ return { proceed: false, token: normalized };
+}
diff --git a/studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx b/studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx
new file mode 100644
index 0000000000..4370c39693
--- /dev/null
+++ b/studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx
@@ -0,0 +1,66 @@
+// 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 {
+ AlertDialog,
+ AlertDialogAction,
+ AlertDialogCancel,
+ AlertDialogContent,
+ AlertDialogDescription,
+ AlertDialogFooter,
+ AlertDialogHeader,
+ AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { Button } from "@/components/ui/button";
+import { AlertTriangle } from "lucide-react";
+import { useHfTokenWarningStore } from "./store";
+
+export function HfTokenWarningDialog() {
+ const open = useHfTokenWarningStore((state) => state.open);
+ const allowAnonymous = useHfTokenWarningStore(
+ (state) => state.allowAnonymous,
+ );
+ const resolve = useHfTokenWarningStore((state) => state.resolve);
+
+ return (
+ {
+ if (!next) resolve("cancel");
+ }}
+ >
+
+
+
+
+
+
Hugging Face token is invalid
+
+ {allowAnonymous
+ ? "Hugging Face rejected the saved token. Replace it to access private or gated repositories, or continue without it for public and fully downloaded models."
+ : "Hugging Face rejected the saved token. Replace it before uploading to the Hub."}
+
+
+
+
+
+ resolve("cancel")}>
+ Cancel
+
+
+ {allowAnonymous ? (
+
resolve("anonymous")}>
+ Continue without token
+
+ ) : null}
+
resolve("replace")}>
+ Replace token
+
+
+
+
+
+ );
+}
diff --git a/studio/frontend/src/features/hf-auth/index.ts b/studio/frontend/src/features/hf-auth/index.ts
new file mode 100644
index 0000000000..e8bcb48193
--- /dev/null
+++ b/studio/frontend/src/features/hf-auth/index.ts
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export { validateHfToken } from "./api";
+export type {
+ HfTokenValidationResult,
+ HfTokenValidationStatus,
+} from "./api";
+export { prepareHfTokenForUse } from "./confirm-token";
+export { HfTokenWarningDialog } from "./hf-token-warning-dialog";
diff --git a/studio/frontend/src/features/hf-auth/store.ts b/studio/frontend/src/features/hf-auth/store.ts
new file mode 100644
index 0000000000..faa2543a8a
--- /dev/null
+++ b/studio/frontend/src/features/hf-auth/store.ts
@@ -0,0 +1,33 @@
+// 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 { create } from "zustand";
+
+export type HfTokenWarningDecision = "anonymous" | "replace" | "cancel";
+type Resolver = (decision: HfTokenWarningDecision) => void;
+
+let pendingResolver: Resolver | null = null;
+
+interface HfTokenWarningStore {
+ open: boolean;
+ allowAnonymous: boolean;
+ requestDecision: (allowAnonymous: boolean) => Promise;
+ resolve: (decision: HfTokenWarningDecision) => void;
+}
+
+export const useHfTokenWarningStore = create((set) => ({
+ open: false,
+ allowAnonymous: true,
+ requestDecision: (allowAnonymous) =>
+ new Promise((resolve) => {
+ pendingResolver?.("cancel");
+ pendingResolver = resolve;
+ set({ open: true, allowAnonymous });
+ }),
+ resolve: (decision) => {
+ const resolver = pendingResolver;
+ pendingResolver = null;
+ set({ open: false, allowAnonymous: true });
+ resolver?.(decision);
+ },
+}));
diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx
index 73a33fd155..d5a85c01c6 100644
--- a/studio/frontend/src/features/settings/tabs/general-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx
@@ -21,6 +21,7 @@ import {
setShowLlamaUpdateBanner,
useShowLlamaUpdateBanner,
} from "@/hooks/use-llama-update-pref";
+import { useHfTokenValidation } from "@/hooks";
import { LOCALE_STORAGE_KEY, useT } from "@/i18n";
import { isTauri } from "@/lib/api-base";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
@@ -216,11 +217,18 @@ export function GeneralTab() {
if (trimmed !== hfToken) setHfToken(trimmed);
};
+ const clearHfToken = () => {
+ draftRef.current = "";
+ setDraftToken("");
+ setHfToken("");
+ };
+
// Show an "accepted" tick once a non-empty token has been committed to the
// store and the field still matches it (i.e. not mid-edit). Gives the user
// feedback that a pasted token was saved.
const tokenSaved =
draftToken.trim().length > 0 && draftToken.trim() === (hfToken ?? "");
+ const tokenValidation = useHfTokenValidation(hfToken ?? "");
useEffect(() => {
let cancelled = false;
@@ -498,46 +506,70 @@ export function GeneralTab() {
label={t("settings.general.huggingFaceToken")}
description={t("settings.general.huggingFaceTokenDescription")}
>
-
-
setDraftToken(e.target.value)}
- onBlur={commitToken}
- className={cn(
- "h-8 w-full font-mono text-xs",
- tokenSaved ? "pr-14" : "pr-8",
- )}
- />
- {tokenSaved ? (
- // Decorative: pointer-events-none lets clicks reach the input
- // underneath so the field still focuses anywhere.
-
+
+
+ setDraftToken(e.target.value)}
+ onBlur={commitToken}
+ className={cn(
+ "h-8 w-full font-mono text-xs",
+ tokenSaved ? "pr-14" : "pr-8",
+ )}
+ />
+ {tokenSaved ? (
+ // Decorative: pointer-events-none lets clicks reach the input
+ // underneath so the field still focuses anywhere.
+
+
+
+ ) : null}
+ setShowToken((s) => !s)}
+ className="absolute right-1.5 top-1/2 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground"
+ aria-label={
+ showToken
+ ? t("settings.general.hideToken")
+ : t("settings.general.showToken")
+ }
+ tabIndex={-1}
+ >
+ {showToken ? (
+
+ ) : (
+
+ )}
+
+
+
-
-
+ Clear
+
+
+ {tokenValidation.isChecking ? (
+ Checking token…
+ ) : tokenValidation.error ? (
+
+ {tokenValidation.error}
+
) : null}
- setShowToken((s) => !s)}
- className="absolute right-1.5 top-1/2 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground"
- aria-label={
- showToken
- ? t("settings.general.hideToken")
- : t("settings.general.showToken")
- }
- tabIndex={-1}
- >
- {showToken ? (
-
- ) : (
-
- )}
-
{/* The desktop app authenticates via desktop auto-auth with a generated
diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts
index 609e4f8591..c5f09c8bd0 100644
--- a/studio/frontend/src/features/training/api/train-api.ts
+++ b/studio/frontend/src/features/training/api/train-api.ts
@@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
+import { prepareHfTokenForUse } from "@/features/hf-auth";
import { readFastApiError } from "@/lib/format-fastapi-error";
import type {
TrainingStartRequest,
@@ -30,10 +31,12 @@ async function parseJson(response: Response): Promise {
export async function startTraining(
payload: TrainingStartRequest,
): Promise {
+ const preparedToken = await prepareHfTokenForUse(payload.hf_token);
+ if (!preparedToken.proceed) throw new Error("Training start cancelled.");
const response = await authFetch("/api/train/start", {
method: "POST",
headers: { "Content-Type": "application/json" },
- body: JSON.stringify(payload),
+ body: JSON.stringify({ ...payload, hf_token: preparedToken.token }),
});
return parseJson(response);
}
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 2f04656c23..0d32e066d5 100644
--- a/studio/frontend/src/features/training/hooks/use-training-actions.ts
+++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts
@@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { primeNativeNotificationPermission } from "@/lib/native-notifications";
+import { prepareHfTokenForUse } from "@/features/hf-auth";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import { useCallback } from "react";
import { toast } from "@/lib/toast";
@@ -43,7 +44,7 @@ export function useTrainingActions() {
const startError = useTrainingRuntimeStore((state) => state.startError);
const startTrainingRun = useCallback(async (): Promise => {
- const config = useTrainingConfigStore.getState();
+ let config = useTrainingConfigStore.getState();
const runtimeStore = useTrainingRuntimeStore.getState();
const dialogStore = useDatasetPreviewDialogStore.getState();
@@ -54,6 +55,13 @@ export function useTrainingActions() {
return false;
}
+ const preparedToken = await prepareHfTokenForUse(config.hfToken);
+ if (!preparedToken.proceed) return false;
+ if ((preparedToken.token ?? "") !== config.hfToken) {
+ config.setHfToken(preparedToken.token ?? "");
+ config = useTrainingConfigStore.getState();
+ }
+
primeNativeNotificationPermission().catch(() => undefined);
runtimeStore.setStartResources(
@@ -226,6 +234,13 @@ export function useTrainingActions() {
resume_from_checkpoint: outputDir,
} as TrainingStartRequest;
+ const preparedToken = await prepareHfTokenForUse(payload.hf_token);
+ if (!preparedToken.proceed) {
+ runtimeStore.setStarting(false);
+ return false;
+ }
+ payload.hf_token = preparedToken.token;
+
runtimeStore.setStartResources(
payload.model_name,
payload.hf_dataset,
diff --git a/studio/frontend/src/hooks/use-hf-token-validation.ts b/studio/frontend/src/hooks/use-hf-token-validation.ts
index 11f56a2151..e0ed02cca2 100644
--- a/studio/frontend/src/hooks/use-hf-token-validation.ts
+++ b/studio/frontend/src/hooks/use-hf-token-validation.ts
@@ -1,8 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import { whoAmI } from "@huggingface/hub";
-import { useCallback, useEffect, useRef, useState } from "react";
+import { validateHfToken } from "@/features/hf-auth";
+import { useEffect, useRef, useState } from "react";
import { useDebouncedValue } from "./use-debounced-value";
export interface HfTokenValidationState {
@@ -17,6 +17,20 @@ const INITIAL: HfTokenValidationState = {
isChecking: false,
};
+interface CompletedValidation extends HfTokenValidationState {
+ token: string;
+}
+
+const NO_COMPLETED_VALIDATION: CompletedValidation = {
+ ...INITIAL,
+ token: "",
+};
+
+// Current user access tokens contain 34 characters after the hf_ prefix.
+// Action-time validation still accepts legacy shapes without spending quota
+// on every intermediate value typed into a live form field.
+const COMPLETE_HF_TOKEN = /^hf_[A-Za-z0-9]{34}$/;
+
/**
* Validates the HF token via the whoami-v2 API, debounced to avoid excessive
* requests while typing. isValid is null until checked.
@@ -26,39 +40,73 @@ export function useHfTokenValidation(token: string): HfTokenValidationState {
token.trim().replace(/^["']+|["']+$/g, ""),
500,
);
- const [state, setState] = useState(INITIAL);
+ const [completed, setCompleted] = useState(
+ NO_COMPLETED_VALIDATION,
+ );
const versionRef = useRef(0);
-
- const runCheck = useCallback(async (t: string) => {
- if (!t) {
- setState({ isValid: null, error: null, isChecking: false });
- return;
- }
-
- const v = ++versionRef.current;
- setState((prev) => ({ ...prev, isChecking: true, error: null }));
-
- try {
- await whoAmI({ accessToken: t });
- if (versionRef.current !== v) return;
- setState({ isValid: true, error: null, isChecking: false });
- } catch {
- if (versionRef.current !== v) return;
- setState({
- isValid: false,
- error: "invalid or expired token",
- isChecking: false,
- });
- }
- }, []);
+ const shouldValidate = COMPLETE_HF_TOKEN.test(debouncedToken);
useEffect(() => {
- if (!debouncedToken) {
- setState(INITIAL);
+ if (!shouldValidate) {
+ versionRef.current += 1;
return;
}
- runCheck(debouncedToken);
- }, [debouncedToken, runCheck]);
+ const version = ++versionRef.current;
+ void validateHfToken(debouncedToken).then(
+ (result) => {
+ if (versionRef.current !== version) return;
+ if (result.status === "valid") {
+ setCompleted({
+ token: debouncedToken,
+ isValid: true,
+ error: null,
+ isChecking: false,
+ });
+ } else if (result.status === "invalid") {
+ setCompleted({
+ token: debouncedToken,
+ isValid: false,
+ error: "invalid or expired token",
+ isChecking: false,
+ });
+ } else if (result.status === "rate_limited") {
+ const wait = result.retryAfterSeconds
+ ? ` Try again in about ${Math.ceil(result.retryAfterSeconds / 60)} minute(s).`
+ : " Try again later.";
+ setCompleted({
+ token: debouncedToken,
+ isValid: null,
+ error: `Token verification is rate limited.${wait}`,
+ isChecking: false,
+ });
+ } else {
+ setCompleted({
+ token: debouncedToken,
+ isValid: null,
+ error: "Could not verify the token. Check your connection and try again.",
+ isChecking: false,
+ });
+ }
+ },
+ () => {
+ if (versionRef.current !== version) return;
+ setCompleted({
+ token: debouncedToken,
+ isValid: null,
+ error: "Could not verify the token. Check your connection and try again.",
+ isChecking: false,
+ });
+ },
+ );
+ }, [debouncedToken, shouldValidate]);
- return state;
+ if (!shouldValidate) return INITIAL;
+ if (completed.token !== debouncedToken) {
+ return { isValid: null, error: null, isChecking: true };
+ }
+ return {
+ isValid: completed.isValid,
+ error: completed.error,
+ isChecking: false,
+ };
}
From 796f8497e79843aec1afdd3f5190d5ec15177201 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 06:47:08 -0700
Subject: [PATCH 040/255] Studio (Windows): keep prompt caching on full GPU
offload (#7260)
* Studio (Windows): keep prompt caching on full GPU offload (#5692 follow-up)
The #5692 full-offload tuning also added --no-cache-prompt, which disables
in-VRAM prompt-prefix reuse. That is unrelated to the host-RAM KV checkpoints
#5692 fixed (--cache-ram 0 / --ctx-checkpoints 0): a fully offloaded model keeps
its KV cache in VRAM, so reusing a common prefix does not copy to system RAM and
does not cause the PCI-E overhead. --no-cache-prompt only forces every request to
re-prefill the whole prompt, which is small for short chats but severe for large
stable system prompts reused across calls (coding agents, long multi-turn chats).
Remove --no-cache-prompt; keep the checkpoint disables and the thread/OMP tuning.
_prompt_cache_disabled stays False (its default), so slot save/restore is intact.
Verified on a fully offloaded gemma GGUF: an identical repeated prompt reprefills
1 token instead of 2220.
* Guard against re-adding --no-cache-prompt to any llama-server command
Add a backend-wide test that AST-scans studio/backend and fails if
--no-cache-prompt is appended/extended/+= into a command. This locks in
the #7260 fix across every code path, not just load_model. Detecting the
flag or honouring a user-supplied one stays allowed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
studio/backend/core/inference/llama_cpp.py | 10 +---
.../tests/test_llama_cpp_mtp_detection.py | 55 ++++++++++++++++++-
2 files changed, 57 insertions(+), 8 deletions(-)
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index dee507fae2..2c7433f7a4 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -7563,8 +7563,9 @@ class LlamaCppBackend:
else:
self._api_key = None
- # Windows + full offload: disable KV checkpoints (WDDM/PCI-E
- # overhead). CPU/partial offload keeps prompt caching. #5692.
+ # Windows + full offload: drop the host-RAM KV checkpoints that cause
+ # WDDM/PCI-E overhead, but keep prompt caching (in-VRAM prefix reuse) so
+ # a repeated prompt is not re-prefilled on every request. #5692.
if sys.platform == "win32" and full_offload_tuning_active:
unsupported_cache_flags: list[str] = []
if server_caps.get("supports_cache_ram"):
@@ -7575,11 +7576,6 @@ class LlamaCppBackend:
cmd.extend(["--ctx-checkpoints", "0"])
else:
unsupported_cache_flags.append("--ctx-checkpoints")
- if server_caps.get("supports_no_cache_prompt"):
- cmd.append("--no-cache-prompt")
- self._prompt_cache_disabled = True
- else:
- unsupported_cache_flags.append("--no-cache-prompt")
if unsupported_cache_flags:
logger.info(
"Skipping unsupported Windows cache flags for llama-server: %s",
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index 68b706ebf9..1d15647967 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -9,6 +9,7 @@ the _already_in_target_state mirror that prevents needless reloads.
from __future__ import annotations
+import ast
import inspect
import os
import struct
@@ -345,10 +346,62 @@ def test_windows_full_offload_flags_use_current_llama_server_args():
stale_checkpoint_flag = "--checkpoint-" + "every-n-tokens"
assert '"--cache-ram"' in src
assert '"--ctx-checkpoints"' in src
- assert '"--no-cache-prompt"' in src
+ # Prompt caching stays on (in-VRAM prefix reuse); #5692 only needed the host-RAM
+ # checkpoints (--cache-ram / --ctx-checkpoints) disabled, not prompt reuse.
+ assert '"--no-cache-prompt"' not in src
assert stale_checkpoint_flag not in src
+# Backend-wide guard: Unsloth must never inject --no-cache-prompt into a llama-server
+# command. It disables in-VRAM prompt-prefix reuse, re-prefilling every repeated prompt
+# (#5692 only needed --cache-ram / --ctx-checkpoints off; #7260 dropped the stray flag).
+# Detecting it (_is_real) or honouring a user-supplied one (_prompt_cache_off) is fine.
+_NO_CACHE_PROMPT_FLAG = "--no-cache-prompt"
+_LIST_MUTATORS = frozenset({"append", "extend", "insert"})
+
+
+def _has_flag_literal(node: ast.AST) -> bool:
+ return any(
+ isinstance(n, ast.Constant) and n.value == _NO_CACHE_PROMPT_FLAG for n in ast.walk(node)
+ )
+
+
+def _no_cache_prompt_injections(source: str, filename: str) -> list[tuple[str, int]]:
+ """(file, lineno) for each spot adding --no-cache-prompt to a list."""
+ hits: list[tuple[str, int]] = []
+ for node in ast.walk(ast.parse(source, filename = filename)):
+ # cmd.append/extend/insert(... flag ...) or cmd += [... flag ...]
+ if (
+ isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Attribute)
+ and node.func.attr in _LIST_MUTATORS
+ and any(_has_flag_literal(a) for a in node.args)
+ ) or (
+ isinstance(node, ast.AugAssign)
+ and isinstance(node.op, ast.Add)
+ and _has_flag_literal(node.value)
+ ):
+ hits.append((filename, node.lineno))
+ return hits
+
+
+def test_unsloth_never_injects_no_cache_prompt_into_any_command():
+ root = Path(_BACKEND_DIR)
+ files = [p for p in root.rglob("*.py") if "tests" not in p.relative_to(root).parts]
+ violations: list[tuple[str, int]] = []
+ for path in files:
+ try:
+ violations += _no_cache_prompt_injections(path.read_text(encoding = "utf-8"), str(path))
+ except (OSError, UnicodeDecodeError, SyntaxError):
+ continue
+ assert files, "no backend source files were scanned"
+ assert violations == [], (
+ "Unsloth must never add --no-cache-prompt to a llama-server command "
+ "(it disables prompt-prefix reuse); detecting or honouring a user-supplied "
+ f"one is fine. Offending sites: {violations}"
+ )
+
+
def test_load_model_sets_threads_once():
src = inspect.getsource(LlamaCppBackend.load_model)
assert src.count('cmd.extend(["--threads", str(') == 1
From 691aebd567f68e6d348a96dc1576a360051c1f32 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 07:07:30 -0700
Subject: [PATCH 041/255] AMD
---
pyproject.toml | 6 +++---
unsloth/models/_utils.py | 2 +-
2 files changed, 4 insertions(+), 4 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index 7b8fdd100d..071258eb8f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -74,7 +74,7 @@ triton = [
]
huggingfacenotorch = [
- "unsloth_zoo>=2026.7.3",
+ "unsloth_zoo>=2026.7.4",
"wheel>=0.42.0",
"packaging",
"numpy",
@@ -95,7 +95,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
- "unsloth_zoo>=2026.7.3",
+ "unsloth_zoo>=2026.7.4",
"torchvision",
"unsloth[triton]",
]
@@ -580,7 +580,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
- "unsloth_zoo>=2026.7.3",
+ "unsloth_zoo>=2026.7.4",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 3d41b505df..57169fa3de 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-__version__ = "2026.7.3"
+__version__ = "2026.7.4"
__all__ = [
"SUPPORTS_BFLOAT16",
From 1c77b4d1496fe5d7f26bf624315acc682ef5cd6e Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Mon, 20 Jul 2026 07:17:31 -0700
Subject: [PATCH 042/255] Bump install.sh / install.ps1 pin to
unsloth>=2026.7.4 (#7263)
PyPI release unsloth 2026.7.4 is live; bump the pinned floor so fresh installs resolve to the new wheel.
---
install.ps1 | 10 +++++-----
install.sh | 10 +++++-----
2 files changed, 10 insertions(+), 10 deletions(-)
diff --git a/install.ps1 b/install.ps1
index 6e059ee0dd..a525d4df56 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2266,7 +2266,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@@ -2280,7 +2280,7 @@ exit 0
}
}
} else {
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -2354,7 +2354,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@@ -2366,7 +2366,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@@ -2394,7 +2394,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
- $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto }
+ $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
diff --git a/install.sh b/install.sh
index c02552628f..0acccdf049 100755
--- a/install.sh
+++ b/install.sh
@@ -3096,7 +3096,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
- "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
+ "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@@ -3113,7 +3113,7 @@ if [ "$_MIGRATED" = true ]; then
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
- "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-}
+ "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" ${_MLX_LM_EXCLUDE_ARG:-}
[ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES"
_UNSLOTH_TORCH_OVERRIDES=""
fi
@@ -3337,7 +3337,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
- "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
+ "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@@ -3356,7 +3356,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
- --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
+ --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@@ -3384,7 +3384,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
- run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto
+ run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
From c9d479f9e36103ae17a55b21cb561c10f1a474ba Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Mon, 20 Jul 2026 20:55:42 -0300
Subject: [PATCH 043/255] Studio: don't let a malformed HF token empty the
model picker's Recommended list (#7266)
---
.../src/components/assistant-ui/model-selector/pickers.tsx | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index 766139e2b4..84119cc992 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -47,7 +47,7 @@ import {
import { useOnlineStatus } from "@/features/hub/hooks/use-online-status";
import { isHiddenModelId } from "@/features/hub/lib/hidden-models";
import { classifyUnslothSupport } from "@/features/hub/lib/unsloth-support";
-import { useHfTokenStore } from "@/features/hub/stores/hf-token-store";
+import { hfApiToken, useHfTokenStore } from "@/features/hub/stores/hf-token-store";
import {
downloadManager,
jobKeyOf,
@@ -1411,7 +1411,8 @@ export function HubModelPicker({
// Shared Hub search stack (the same hooks the Hub page uses) so the picker
// and Hub run one implementation. Scoped to unsloth like the old listing.
const online = useOnlineStatus();
- const accessToken = hfToken || undefined;
+ // Sanitize to anonymous on a malformed token, matching the Hub page.
+ const accessToken = hfApiToken(hfToken);
// Recommended section: a live unsloth listing sorted by the dropdown. The
// same sort drives the search results so the dropdown works while searching.
const [recommendedSort, setRecommendedSort] =
From 3d379cdb81ea6b1688eee4812ff5a6e1e85c2c74 Mon Sep 17 00:00:00 2001
From: Long Yixing
Date: Tue, 21 Jul 2026 10:14:58 +0800
Subject: [PATCH 044/255] Fix local CLI streamed generation error handling
(#7135)
---
studio/backend/core/inference/orchestrator.py | 5 +-
unsloth_cli/_inference.py | 2 +-
unsloth_cli/commands/chat.py | 2 +-
unsloth_cli/commands/inference.py | 5 +-
unsloth_cli/tests/test_inference_chat.py | 98 +++++++++++++++++++
5 files changed, 103 insertions(+), 9 deletions(-)
diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index eaa474d9b8..75ef9c2399 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/studio/backend/core/inference/orchestrator.py
@@ -54,9 +54,8 @@ class GenStreamError(str):
"""A stream chunk carrying a real backend/generation error, not model text.
Subclasses str so existing display/logging consumers are unaffected, while
- callers that must abort a distributed run on error (raise_on_streamed_error)
- can distinguish a real error from model output whose visible text starts with
- "Error:" by checking isinstance(chunk, GenStreamError).
+ callers can distinguish a real error from model output whose visible text
+ starts with "Error:" by checking isinstance(chunk, GenStreamError).
"""
__slots__ = ("public",)
diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py
index 551bef4787..a2b0f9c04f 100644
--- a/unsloth_cli/_inference.py
+++ b/unsloth_cli/_inference.py
@@ -235,7 +235,7 @@ def collect_stream(stream, show_thinking: bool) -> str:
def raise_on_streamed_error(stream):
# Match real backend errors by type (GenStreamError), not the "Error:" text
# prefix, so a completion whose text opens with "Error:" is not misread as a
- # failure that aborts a distributed run.
+ # backend failure.
try:
ensure_studio_backend_path()
from core.inference.orchestrator import GenStreamError
diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py
index bba5fab08e..fdc5577700 100644
--- a/unsloth_cli/commands/chat.py
+++ b/unsloth_cli/commands/chat.py
@@ -331,7 +331,7 @@ def chat(
enable_thinking = show_thinking,
use_adapter = use_adapter,
)
- return raise_on_streamed_error(stream) if is_mlx_distributed else stream
+ return raise_on_streamed_error(stream)
if should_print:
console.print()
diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py
index 524d8fd015..7df05cef94 100644
--- a/unsloth_cli/commands/inference.py
+++ b/unsloth_cli/commands/inference.py
@@ -111,15 +111,12 @@ def inference(
repetition_penalty = repetition_penalty,
enable_thinking = think,
)
- if is_mlx_distributed:
- stream = raise_on_streamed_error(stream)
+ stream = raise_on_streamed_error(stream)
if rank == 0:
typer.echo("Assistant:")
try:
stream_to_stdout(stream, show_thinking = think)
except RuntimeError as exc:
- if not is_mlx_distributed:
- raise
typer.echo(f"Error: {exc}", err = True)
raise typer.Exit(code = 1)
else:
diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py
index ae6f8dcfd4..07eed01e97 100644
--- a/unsloth_cli/tests/test_inference_chat.py
+++ b/unsloth_cli/tests/test_inference_chat.py
@@ -873,6 +873,104 @@ def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch):
assert set(closed) == {"tuned", "base"}
+@pytest.mark.parametrize(
+ ("chunk_kind", "expected_exit"),
+ [
+ ("answer", 0),
+ ("model_text_error", 0),
+ ("real_error", 1),
+ ],
+)
+def test_inference_local_handles_stream(monkeypatch, chunk_kind, expected_exit):
+ from unsloth_cli.commands import inference as infermod
+ from unsloth_cli._inference import ensure_studio_backend_path
+
+ ensure_studio_backend_path()
+ from core.inference.orchestrator import GenStreamError
+
+ chunks = {
+ "answer": ["answer"],
+ "model_text_error": ["Error: printed by the model, not a backend failure"],
+ "real_error": [GenStreamError("Error: generation failed")],
+ }[chunk_kind]
+ closed = []
+
+ class _FakeBackend:
+ def stream(self, messages, **kwargs):
+ return iter(chunks)
+
+ def close(self):
+ closed.append(True)
+
+ monkeypatch.setattr(
+ infermod,
+ "connect_studio_server",
+ lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
+ )
+ monkeypatch.setattr(infermod, "load_chat_backend", lambda *a, **k: _FakeBackend())
+
+ result = CliRunner().invoke(
+ _inference_app(),
+ ["fake-model", "hello", "--no-server"],
+ )
+
+ assert result.exit_code == expected_exit, result.output
+ assert closed == [True]
+ if chunk_kind == "real_error":
+ assert result.stdout == "Assistant:\n"
+ assert result.stderr == "Error: generation failed\n"
+ else:
+ assert chunks[0] in result.output
+
+
+@pytest.mark.parametrize("chunk_kind", ["answer", "model_text_error", "real_error"])
+def test_chat_local_handles_stream(monkeypatch, chunk_kind):
+ from unsloth_cli._inference import ensure_studio_backend_path
+
+ ensure_studio_backend_path()
+ from core.inference.orchestrator import GenStreamError
+
+ first_chunk = {
+ "answer": "answer",
+ "model_text_error": "Error: printed by the model, not a backend failure",
+ "real_error": GenStreamError("Error: generation failed"),
+ }[chunk_kind]
+ calls, closed = [], []
+
+ class _FakeChatBackend:
+ def stream(self, messages, **kwargs):
+ calls.append([dict(message) for message in messages])
+ return iter([first_chunk if len(calls) == 1 else "second answer"])
+
+ def close(self):
+ closed.append(True)
+
+ monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
+ monkeypatch.setattr(chatmod, "connect_studio_server", lambda *a, **k: None)
+ monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
+ monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
+
+ result = CliRunner().invoke(
+ _chat_app(),
+ ["fake-model"],
+ input = "first\nsecond\n/exit\n",
+ )
+
+ assert result.exit_code == 0, result.output
+ assert closed == [True]
+ if chunk_kind == "real_error":
+ assert calls[1] == [{"role": "user", "content": "second"}]
+ assert "(error: generation failed)" in result.output
+ assert "Error: generation failed" not in result.output
+ else:
+ assert calls[1] == [
+ {"role": "user", "content": "first"},
+ {"role": "assistant", "content": first_chunk},
+ {"role": "user", "content": "second"},
+ ]
+ assert first_chunk in result.output
+
+
@pytest.mark.parametrize(
("chunk_kind", "expected_exit"),
[
From 27f3473c7eb4930c7aadce20945d3d6984029411 Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Tue, 21 Jul 2026 02:39:43 -0300
Subject: [PATCH 045/255] Studio: make tab navigation feel immediate (#7271)
* Studio: make repeated tab switches feel immediate
* Keep cached Studio navigation data fresh
* Make first Studio tab visits responsive
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Serve range requests uncompressed for immutable assets (PR #7271)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: test
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
studio/backend/main.py | 36 ++++++-
studio/backend/tests/test_middleware.py | 66 ++++++++++++
studio/frontend/src/app/auth-guards.ts | 61 ++++++++---
studio/frontend/src/app/routes/__root.tsx | 2 +-
.../frontend/src/app/routes/data-recipes.tsx | 10 +-
studio/frontend/src/app/routes/export.tsx | 10 +-
studio/frontend/src/app/routes/hub.tsx | 16 +--
studio/frontend/src/app/routes/projects.tsx | 10 +-
studio/frontend/src/app/routes/studio.tsx | 10 +-
.../frontend/src/components/app-sidebar.tsx | 46 +++++++-
.../src/features/chat/api/chat-api.ts | 14 ++-
.../features/chat/hooks/use-chat-projects.ts | 100 ++++++++++++++----
.../features/data-recipes/data/recipes-db.ts | 46 +++++++-
.../src/features/data-recipes/index.ts | 1 +
.../export/export-navigation-cache.ts | 61 +++++++++++
.../src/features/export/export-page.tsx | 50 +++++----
.../hub/hooks/use-hub-paginated-search.ts | 48 +++++++--
studio/frontend/src/features/hub/hub-page.tsx | 36 ++++++-
18 files changed, 516 insertions(+), 107 deletions(-)
create mode 100644 studio/frontend/src/features/export/export-navigation-cache.ts
diff --git a/studio/backend/main.py b/studio/backend/main.py
index f686e29bf5..48675b9539 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -289,6 +289,7 @@ from fastapi import Depends, FastAPI, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse, Response
+from starlette.middleware.gzip import GZipMiddleware
from pathlib import Path
from datetime import datetime
@@ -1509,6 +1510,34 @@ def _should_inject_bootstrap(request: Request) -> bool:
return _is_local_bootstrap_request(request)
+_IMMUTABLE_ASSET_CACHE_CONTROL = "public, max-age=31536000, immutable"
+
+
+class ImmutableStaticFiles(StaticFiles):
+ """Serve Vite's content-hashed assets without browser revalidation."""
+
+ def file_response(
+ self,
+ full_path,
+ stat_result,
+ scope,
+ status_code = 200,
+ ):
+ response = super().file_response(full_path, stat_result, scope, status_code)
+ response.headers["Cache-Control"] = _IMMUTABLE_ASSET_CACHE_CONTROL
+ return response
+
+
+class _AssetGZipMiddleware(GZipMiddleware):
+ """Serve range requests uncompressed; gzip + 206 mislabels Content-Range."""
+
+ async def __call__(self, scope, receive, send):
+ if scope["type"] == "http" and any(key == b"range" for key, _ in scope["headers"]):
+ await self.app(scope, receive, send)
+ return
+ await super().__call__(scope, receive, send)
+
+
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
@@ -1516,7 +1545,12 @@ def setup_frontend(app: FastAPI, build_path: Path):
assets_dir = build_path / "assets"
if assets_dir.exists():
- app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
+ assets_app = _AssetGZipMiddleware(
+ ImmutableStaticFiles(directory = assets_dir),
+ minimum_size = 1024,
+ compresslevel = 6,
+ )
+ app.mount("/assets", assets_app, name = "assets")
def _build_index_response(request: Request) -> Response:
content = (build_path / "index.html").read_bytes()
diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py
index 11aeee6d77..209c6cb90a 100644
--- a/studio/backend/tests/test_middleware.py
+++ b/studio/backend/tests/test_middleware.py
@@ -14,6 +14,7 @@ import pytest
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import Response
from fastapi.testclient import TestClient
+from starlette.middleware.gzip import GZipMiddleware
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
@@ -471,6 +472,71 @@ class TestSecurityHeadersMiddleware:
assert b"server" in names
+class TestFrontendAssets:
+ def test_hashed_assets_are_compressed_and_cached(self, tmp_path, main_module):
+ content = b"export const value = 'responsive';\n" * 200
+ (tmp_path / "page-abc123.js").write_bytes(content)
+ app = FastAPI()
+ assets_app = GZipMiddleware(
+ main_module.ImmutableStaticFiles(directory = tmp_path),
+ minimum_size = 1024,
+ compresslevel = 6,
+ )
+ app.mount("/assets", assets_app, name = "assets")
+
+ response = TestClient(app).get(
+ "/assets/page-abc123.js",
+ headers = {"Accept-Encoding": "gzip"},
+ )
+
+ assert response.status_code == 200
+ assert response.content == content
+ assert response.headers["content-encoding"] == "gzip"
+ assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
+ assert "accept-encoding" in response.headers["vary"].lower()
+
+ def test_asset_revalidation_keeps_immutable_cache_header(self, tmp_path, main_module):
+ (tmp_path / "page-abc123.js").write_text("export {};", encoding = "utf-8")
+ app = FastAPI()
+ app.mount(
+ "/assets",
+ main_module.ImmutableStaticFiles(directory = tmp_path),
+ name = "assets",
+ )
+ client = TestClient(app)
+ first = client.get("/assets/page-abc123.js")
+
+ response = client.get(
+ "/assets/page-abc123.js",
+ headers = {"If-None-Match": first.headers["etag"]},
+ )
+
+ assert response.status_code == 304
+ assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
+
+ def test_range_request_is_not_compressed(self, tmp_path, main_module):
+ content = b"export const value = 'responsive';\n" * 200
+ (tmp_path / "page-abc123.js").write_bytes(content)
+ app = FastAPI()
+ assets_app = main_module._AssetGZipMiddleware(
+ main_module.ImmutableStaticFiles(directory = tmp_path),
+ minimum_size = 1024,
+ compresslevel = 6,
+ )
+ app.mount("/assets", assets_app, name = "assets")
+
+ response = TestClient(app).get(
+ "/assets/page-abc123.js",
+ headers = {"Accept-Encoding": "gzip", "Range": "bytes=0-99"},
+ )
+
+ assert response.status_code == 206
+ assert response.headers.get("content-encoding") != "gzip"
+ assert response.headers["content-range"] == f"bytes 0-99/{len(content)}"
+ assert response.content == content[:100]
+ assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
+
+
# /api/health auth gate
diff --git a/studio/frontend/src/app/auth-guards.ts b/studio/frontend/src/app/auth-guards.ts
index 6849f380b8..a3523ac580 100644
--- a/studio/frontend/src/app/auth-guards.ts
+++ b/studio/frontend/src/app/auth-guards.ts
@@ -23,19 +23,47 @@ interface AuthStatus {
requires_password_change: boolean;
}
+const AUTH_STATUS_TTL_MS = 30_000;
+let authStatusCheckedAt = 0;
+let authStatusRequest: Promise | null = null;
+
+function hasFreshAuthStatus(): boolean {
+ return (
+ authStatusCheckedAt !== 0 &&
+ Date.now() - authStatusCheckedAt < AUTH_STATUS_TTL_MS
+ );
+}
+
async function fetchAuthStatus(): Promise {
- try {
- const res = await fetch(apiUrl("/api/auth/status"));
- if (!res.ok) return { initialized: true, requires_password_change: mustChangePassword() };
- const status = (await res.json()) as AuthStatus;
- // Server truth wins; keep localStorage in sync both ways.
- if (status.requires_password_change !== mustChangePassword()) {
- setMustChangePassword(status.requires_password_change);
+ if (authStatusRequest) return authStatusRequest;
+
+ const request = (async () => {
+ try {
+ const res = await fetch(apiUrl("/api/auth/status"));
+ if (!res.ok) {
+ return {
+ initialized: true,
+ requires_password_change: mustChangePassword(),
+ };
+ }
+ const status = (await res.json()) as AuthStatus;
+ authStatusCheckedAt = Date.now();
+ // Server truth wins; keep localStorage in sync both ways.
+ if (status.requires_password_change !== mustChangePassword()) {
+ setMustChangePassword(status.requires_password_change);
+ }
+ return status;
+ } catch {
+ return {
+ initialized: true,
+ requires_password_change: mustChangePassword(),
+ };
}
- return status;
- } catch {
- return { initialized: true, requires_password_change: mustChangePassword() };
- }
+ })().finally(() => {
+ authStatusRequest = null;
+ });
+ authStatusRequest = request;
+ return request;
}
function authRedirect(to: "/login" | "/change-password"): never {
@@ -49,12 +77,17 @@ export async function requireAuth(): Promise {
}
if (await hasActiveSession()) {
- const { requires_password_change } = await fetchAuthStatus();
- if (requires_password_change || mustChangePassword()) {
- authRedirect("/change-password");
+ // Reconcile periodically so local-only routes cannot outlive a server-side
+ // password-change requirement, while nearby route switches stay local.
+ if (mustChangePassword() || !hasFreshAuthStatus()) {
+ const { requires_password_change } = await fetchAuthStatus();
+ if (requires_password_change || mustChangePassword()) {
+ authRedirect("/change-password");
+ }
}
return;
}
+
const status = await fetchAuthStatus();
if (status.requires_password_change || mustChangePassword()) {
authRedirect("/change-password");
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index e23892e020..57e890dd5a 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -281,7 +281,7 @@ function RootLayout() {
initial={{ opacity: 0 }}
animate={{ opacity: 1 }}
exit={{ opacity: 0 }}
- transition={{ duration: 0.15 }}
+ transition={{ duration: 0.06 }}
className="flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-visible"
>
}>
diff --git a/studio/frontend/src/app/routes/data-recipes.tsx b/studio/frontend/src/app/routes/data-recipes.tsx
index c35e63da5f..22f87821af 100644
--- a/studio/frontend/src/app/routes/data-recipes.tsx
+++ b/studio/frontend/src/app/routes/data-recipes.tsx
@@ -1,15 +1,13 @@
// 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 { createRoute } from "@tanstack/react-router";
-import { lazy } from "react";
+import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
-const DataRecipesPage = lazy(() =>
- import("@/features/data-recipes").then((m) => ({
- default: m.DataRecipesPage,
- })),
+const DataRecipesPage = lazyRouteComponent(
+ () => import("@/features/data-recipes"),
+ "DataRecipesPage",
);
export const Route = createRoute({
diff --git a/studio/frontend/src/app/routes/export.tsx b/studio/frontend/src/app/routes/export.tsx
index 40118c6a92..5a7b586f19 100644
--- a/studio/frontend/src/app/routes/export.tsx
+++ b/studio/frontend/src/app/routes/export.tsx
@@ -1,15 +1,13 @@
// 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 { createRoute } from "@tanstack/react-router";
-import { lazy } from "react";
+import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
-const ExportPage = lazy(() =>
- import("@/features/export/export-page").then((m) => ({
- default: m.ExportPage,
- })),
+const ExportPage = lazyRouteComponent(
+ () => import("@/features/export/export-page"),
+ "ExportPage",
);
export type ExportSearch = {
diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx
index c623ef9848..2207490e44 100644
--- a/studio/frontend/src/app/routes/hub.tsx
+++ b/studio/frontend/src/app/routes/hub.tsx
@@ -1,15 +1,13 @@
// 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 { createRoute } from "@tanstack/react-router";
-import { lazy } from "react";
+import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
-const ModelsPage = lazy(() =>
- import("@/features/hub/hub-page").then((m) => ({
- default: m.ModelsPage,
- })),
+const ModelsPage = lazyRouteComponent(
+ () => import("@/features/hub/hub-page"),
+ "ModelsPage",
);
export interface ModelsSearch {
@@ -31,7 +29,11 @@ export const Route = createRoute({
const model = search.model;
if (typeof model === "string" && model.length > 0) next.model = model;
const section = search.section;
- if (section === "trending" || section === "latest" || section === "finetune") {
+ if (
+ section === "trending" ||
+ section === "latest" ||
+ section === "finetune"
+ ) {
next.section = section;
}
const kind = search.kind;
diff --git a/studio/frontend/src/app/routes/projects.tsx b/studio/frontend/src/app/routes/projects.tsx
index c63b1d5838..17f58ef631 100644
--- a/studio/frontend/src/app/routes/projects.tsx
+++ b/studio/frontend/src/app/routes/projects.tsx
@@ -1,15 +1,13 @@
// 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 { createRoute } from "@tanstack/react-router";
-import { lazy } from "react";
+import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
-const ProjectsPage = lazy(() =>
- import("@/features/chat/projects-page").then((m) => ({
- default: m.ProjectsPage,
- })),
+const ProjectsPage = lazyRouteComponent(
+ () => import("@/features/chat/projects-page"),
+ "ProjectsPage",
);
export const Route = createRoute({
diff --git a/studio/frontend/src/app/routes/studio.tsx b/studio/frontend/src/app/routes/studio.tsx
index ae7f445e94..798044bf64 100644
--- a/studio/frontend/src/app/routes/studio.tsx
+++ b/studio/frontend/src/app/routes/studio.tsx
@@ -1,15 +1,13 @@
// 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 { createRoute } from "@tanstack/react-router";
-import { lazy } from "react";
+import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
-const StudioPage = lazy(() =>
- import("@/features/studio/studio-page").then((m) => ({
- default: m.StudioPage,
- })),
+const StudioPage = lazyRouteComponent(
+ () => import("@/features/studio/studio-page"),
+ "StudioPage",
);
export const Route = createRoute({
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index b8601b00f6..8eab03133b 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -93,7 +93,12 @@ import {
import { Tooltip as TooltipPrimitive } from "radix-ui";
import { HugeiconsIcon } from "@hugeicons/react";
import { ChevronDown, Moon } from "lucide-react";
-import { Link, useNavigate, useRouterState } from "@tanstack/react-router";
+import {
+ Link,
+ useNavigate,
+ useRouter,
+ useRouterState,
+} from "@tanstack/react-router";
import {
archiveChatItem,
ChatSearchDialog,
@@ -256,6 +261,10 @@ function createNavigationNonce(): string {
return `${Date.now()}-${Math.random().toString(36).slice(2, 10)}`;
}
+function preloadSilently(request: Promise): void {
+ void request.catch(() => undefined);
+}
+
function NavItem({
icon,
label,
@@ -267,6 +276,7 @@ function NavItem({
className,
spinner,
tooltip,
+ onIntent,
}: {
icon: typeof ZapIcon;
label: string;
@@ -277,6 +287,7 @@ function NavItem({
dataTour?: string;
className?: string;
spinner?: boolean;
+ onIntent?: () => void;
// Overrides the hover tooltip (defaults to `label`). Used to explain why a
// disabled item (e.g. Train/Export on a chat-only host) is greyed out.
tooltip?: string;
@@ -288,6 +299,8 @@ function NavItem({
tooltip={tooltip ?? label}
disabled={disabled}
onClick={onClick}
+ onPointerEnter={disabled ? undefined : onIntent}
+ onFocus={disabled ? undefined : onIntent}
isActive={active}
data-tour={dataTour}
className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:mx-auto"
@@ -324,6 +337,7 @@ export function AppSidebar() {
});
const { togglePinned, isMobile, setOpenMobile } = useSidebar();
const navigate = useNavigate();
+ const router = useRouter();
// Web update detection: `webUpdate` is non-null only when the installed
// (PyPI) version is behind the latest release, so the card is hidden by
@@ -1218,6 +1232,9 @@ export function AppSidebar() {
navigate({ to: "/projects" });
closeMobileIfOpen();
}}
+ onIntent={() => {
+ preloadSilently(router.preloadRoute({ to: "/projects" }));
+ }}
className="group/projects-item relative"
>
{
+ preloadSilently(router.preloadRoute({ to: "/hub" }));
+ }}
/>
{/* Train has a labelled section when expanded; plain icon here only when collapsed. */}
{
+ preloadSilently(router.preloadRoute({ to: "/studio" }));
+ }}
className="hidden group-data-[collapsible=icon]:block"
/>
@@ -1293,6 +1316,9 @@ export function AppSidebar() {
navigate({ to: "/studio" });
closeMobileIfOpen();
}}
+ onIntent={() => {
+ preloadSilently(router.preloadRoute({ to: "/studio" }));
+ }}
/>
{
+ preloadSilently(
+ router.preloadRoute({ to: "/data-recipes" }),
+ );
+ preloadSilently(
+ import("@/features/data-recipes").then((module) =>
+ module.preloadRecipes(),
+ ),
+ );
+ }}
/>
{
+ preloadSilently(router.preloadRoute({ to: "/export" }));
+ preloadSilently(
+ import(
+ "@/features/export/export-navigation-cache"
+ ).then((module) => module.preloadExportData()),
+ );
+ }}
/>
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index 475bd0f801..631474c39a 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -33,6 +33,7 @@ import type {
} from "../types/api";
export const CHAT_HISTORY_UPDATED_EVENT = "unsloth-chat-history-updated";
+export const CHAT_PROJECTS_UPDATED_EVENT = "unsloth-chat-projects-updated";
/**
* Thrown when the chat SSE stream ends without a terminal signal (`[DONE]` or a
@@ -55,6 +56,13 @@ export function notifyChatHistoryUpdated(): void {
}
}
+function notifyChatProjectsUpdated(): void {
+ notifyChatHistoryUpdated();
+ if (typeof window !== "undefined") {
+ window.dispatchEvent(new Event(CHAT_PROJECTS_UPDATED_EVENT));
+ }
+}
+
function parseErrorText(status: number, body: unknown): string {
if (body && typeof body === "object") {
const detail = (body as { detail?: unknown }).detail;
@@ -644,7 +652,7 @@ export async function saveChatProject(
body: JSON.stringify(project),
});
const saved = await parseJsonOrThrow(response);
- notifyChatHistoryUpdated();
+ notifyChatProjectsUpdated();
return saved;
}
@@ -661,7 +669,7 @@ export async function updateChatProject(
},
);
const project = await parseJsonOrThrow(response);
- notifyChatHistoryUpdated();
+ notifyChatProjectsUpdated();
return project;
}
@@ -677,7 +685,7 @@ export async function deleteChatProject(
{ method: "DELETE" },
);
await parseJsonOrThrow(response);
- notifyChatHistoryUpdated();
+ notifyChatProjectsUpdated();
}
export async function listChatMessages(
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-projects.ts b/studio/frontend/src/features/chat/hooks/use-chat-projects.ts
index 3f0d46982d..86f540654d 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-projects.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-projects.ts
@@ -1,8 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import { useEffect, useState } from "react";
-import { CHAT_HISTORY_UPDATED_EVENT } from "../api/chat-api";
+import { useEffect, useState, useSyncExternalStore } from "react";
+import { CHAT_PROJECTS_UPDATED_EVENT } from "../api/chat-api";
import type { ProjectRecord } from "../types";
import {
createStoredChatProject,
@@ -15,32 +15,83 @@ import {
import type { SidebarItem } from "./use-chat-sidebar-items";
let cachedProjects: ProjectRecord[] = [];
+let projectsLoaded = false;
+let projectsRequest: Promise | null = null;
+let projectsRefreshPending = false;
+let lastProjectsUpdateEvent: Event | null = null;
+const projectSubscribers = new Set<() => void>();
+
+function subscribeToProjects(onStoreChange: () => void): () => void {
+ projectSubscribers.add(onStoreChange);
+ return () => projectSubscribers.delete(onStoreChange);
+}
+
+function getProjectsSnapshot(): ProjectRecord[] {
+ return cachedProjects;
+}
+
+function publishProjects(projects: ProjectRecord[]): void {
+ cachedProjects = projects;
+ projectsLoaded = true;
+ for (const onStoreChange of projectSubscribers) onStoreChange();
+}
+
+function loadProjects(
+ force = false,
+ followUpIfPending = false,
+): Promise {
+ if (projectsRequest) {
+ if (followUpIfPending) projectsRefreshPending = true;
+ return projectsRequest;
+ }
+ if (!force && projectsLoaded) {
+ return Promise.resolve(cachedProjects);
+ }
+
+ async function run(): Promise {
+ let nextProjects: ProjectRecord[] | null = null;
+ do {
+ projectsRefreshPending = false;
+ try {
+ const next = await listStoredChatProjects({ includeArchived: false });
+ nextProjects = Array.isArray(next) ? next : [];
+ } catch (error) {
+ if (!isExpectedBackgroundChatStorageError(error)) throw error;
+ nextProjects = null;
+ }
+ } while (projectsRefreshPending);
+ if (nextProjects !== null) publishProjects(nextProjects);
+ return cachedProjects;
+ }
+
+ const request = run().finally(() => {
+ projectsRequest = null;
+ });
+ projectsRequest = request;
+ return request;
+}
export function useChatProjects(): {
projects: ProjectRecord[];
isLoading: boolean;
hasLoaded: boolean;
} {
- // Stay null-safe even if the cache was poisoned by a bad response.
- const cached = Array.isArray(cachedProjects) ? cachedProjects : [];
- const [projects, setProjects] = useState(cached);
- const [isLoading, setIsLoading] = useState(cached.length === 0);
- const [hasLoaded, setHasLoaded] = useState(cached.length > 0);
+ const projects = useSyncExternalStore(
+ subscribeToProjects,
+ getProjectsSnapshot,
+ getProjectsSnapshot,
+ );
+ const [isLoading, setIsLoading] = useState(!projectsLoaded);
+ const [hasLoaded, setHasLoaded] = useState(projectsLoaded);
useEffect(() => {
let cancelled = false;
- async function load() {
- if (!cancelled) setIsLoading(true);
+ async function refresh(force = false, followUpIfPending = false) {
+ if (!force && projectsLoaded) return;
+ if (!cancelled && !projectsLoaded) setIsLoading(true);
try {
- const next = await listStoredChatProjects({ includeArchived: false });
- cachedProjects = Array.isArray(next) ? next : [];
- if (!cancelled) setProjects(cachedProjects);
- } catch (error) {
- if (isExpectedBackgroundChatStorageError(error)) {
- return;
- }
- if (!cancelled) throw error;
+ await loadProjects(force, followUpIfPending);
} finally {
if (!cancelled) {
setHasLoaded(true);
@@ -49,15 +100,18 @@ export function useChatProjects(): {
}
}
- const onHistoryUpdated = () => {
- void load();
+ const onProjectsUpdated = (event: Event) => {
+ const followUpIfPending = event !== lastProjectsUpdateEvent;
+ lastProjectsUpdateEvent = event;
+ void refresh(true, followUpIfPending);
};
-
- void load();
- window.addEventListener(CHAT_HISTORY_UPDATED_EVENT, onHistoryUpdated);
+ // Cached rows render immediately, then one shared request reconciles
+ // changes made by another browser tab or API client.
+ void refresh(projectsLoaded);
+ window.addEventListener(CHAT_PROJECTS_UPDATED_EVENT, onProjectsUpdated);
return () => {
cancelled = true;
- window.removeEventListener(CHAT_HISTORY_UPDATED_EVENT, onHistoryUpdated);
+ window.removeEventListener(CHAT_PROJECTS_UPDATED_EVENT, onProjectsUpdated);
};
}, []);
diff --git a/studio/frontend/src/features/data-recipes/data/recipes-db.ts b/studio/frontend/src/features/data-recipes/data/recipes-db.ts
index f89707e159..3ffe09c2d2 100644
--- a/studio/frontend/src/features/data-recipes/data/recipes-db.ts
+++ b/studio/frontend/src/features/data-recipes/data/recipes-db.ts
@@ -16,11 +16,40 @@ db.version(1).stores({
});
const recentRecipeCache = new Map();
+let cachedRecipeList: RecipeRecord[] = [];
+let recipeListReady = false;
+let recipeListRequest: Promise | null = null;
export function listRecipes(): Promise {
return db.recipes.orderBy("updatedAt").reverse().toArray();
}
+function cacheRecipeList(recipes: RecipeRecord[]): RecipeRecord[] {
+ for (const recipe of recipes) {
+ writeRecipeCache(recipe);
+ }
+ cachedRecipeList = recipes;
+ recipeListReady = true;
+ return recipes;
+}
+
+export function preloadRecipes(): Promise {
+ if (recipeListReady) {
+ return Promise.resolve(cachedRecipeList);
+ }
+ if (recipeListRequest) {
+ return recipeListRequest;
+ }
+
+ const request = listRecipes()
+ .then(cacheRecipeList)
+ .finally(() => {
+ recipeListRequest = null;
+ });
+ recipeListRequest = request;
+ return request;
+}
+
export function getRecipe(id: string): Promise {
return db.recipes.get(id);
}
@@ -55,12 +84,21 @@ export async function saveRecipe(
};
await db.recipes.put(record);
writeRecipeCache(record);
+ if (recipeListReady) {
+ cachedRecipeList = [
+ record,
+ ...cachedRecipeList.filter((recipe) => recipe.id !== record.id),
+ ].sort((a, b) => b.updatedAt - a.updatedAt);
+ }
return record;
}
export async function deleteRecipe(id: string): Promise {
await db.recipes.delete(id);
recentRecipeCache.delete(id);
+ if (recipeListReady) {
+ cachedRecipeList = cachedRecipeList.filter((recipe) => recipe.id !== id);
+ }
}
export function createRecipeDraft(): Promise {
@@ -87,15 +125,13 @@ export function useRecipes(): {
recipes: RecipeRecord[];
ready: boolean;
} {
- const [recipes, setRecipes] = useState([]);
- const [ready, setReady] = useState(false);
+ const [recipes, setRecipes] = useState(cachedRecipeList);
+ const [ready, setReady] = useState(recipeListReady);
useEffect(() => {
const sub = liveQuery(() => listRecipes()).subscribe({
next: (value) => {
- for (const recipe of value) {
- writeRecipeCache(recipe);
- }
+ cacheRecipeList(value);
setRecipes(value);
setReady(true);
},
diff --git a/studio/frontend/src/features/data-recipes/index.ts b/studio/frontend/src/features/data-recipes/index.ts
index d085096e0d..a75dd7075a 100644
--- a/studio/frontend/src/features/data-recipes/index.ts
+++ b/studio/frontend/src/features/data-recipes/index.ts
@@ -3,3 +3,4 @@
export { DataRecipesPage } from "./pages/data-recipes-page";
export { EditRecipePage } from "./pages/edit-recipe-page";
+export { preloadRecipes } from "./data/recipes-db";
diff --git a/studio/frontend/src/features/export/export-navigation-cache.ts b/studio/frontend/src/features/export/export-navigation-cache.ts
new file mode 100644
index 0000000000..f026f64525
--- /dev/null
+++ b/studio/frontend/src/features/export/export-navigation-cache.ts
@@ -0,0 +1,61 @@
+// 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 LocalModelInfo, listLocalModels } from "@/features/training";
+import { type ModelCheckpoints, fetchCheckpoints } from "./api/export-api";
+
+let cachedCheckpoints: ModelCheckpoints[] | null = null;
+let checkpointsRequest: Promise | null = null;
+let cachedLocalModels: LocalModelInfo[] | null = null;
+let localModelsRequest: Promise | null = null;
+
+export function getCachedCheckpoints(): ModelCheckpoints[] | null {
+ return cachedCheckpoints;
+}
+
+export function getCachedLocalModels(): LocalModelInfo[] | null {
+ return cachedLocalModels;
+}
+
+export function refreshCheckpoints(): Promise {
+ if (checkpointsRequest) {
+ return checkpointsRequest;
+ }
+ const request = fetchCheckpoints()
+ .then((data) => {
+ cachedCheckpoints = data.models;
+ return data.models;
+ })
+ .finally(() => {
+ checkpointsRequest = null;
+ });
+ checkpointsRequest = request;
+ return request;
+}
+
+export function refreshLocalModels(): Promise {
+ if (localModelsRequest) {
+ return localModelsRequest;
+ }
+ const request = listLocalModels()
+ .then((models) => {
+ cachedLocalModels = models;
+ return models;
+ })
+ .finally(() => {
+ localModelsRequest = null;
+ });
+ localModelsRequest = request;
+ return request;
+}
+
+export async function preloadExportData(): Promise {
+ const requests: Promise[] = [];
+ if (cachedCheckpoints === null) {
+ requests.push(refreshCheckpoints());
+ }
+ if (cachedLocalModels === null) {
+ requests.push(refreshLocalModels());
+ }
+ await Promise.allSettled(requests);
+}
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx
index 51e9f209dd..3a970713ac 100644
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@@ -48,7 +48,6 @@ import { prepareHfTokenForUse } from "@/features/hf-auth";
import { GuidedTour, useGuidedTourController } from "@/features/tour";
import {
type LocalModelInfo,
- listLocalModels,
useTrainingConfigStore,
} from "@/features/training";
import { useDebouncedValue, useHfTokenValidation } from "@/hooks";
@@ -67,7 +66,6 @@ import { useSearch } from "@tanstack/react-router";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { useShallow } from "zustand/react/shallow";
import type { ModelCheckpoints } from "./api/export-api";
-import { fetchCheckpoints } from "./api/export-api";
import { ExportRunPanel } from "./components/export-run-panel";
import { MethodPicker } from "./components/method-picker";
import { QuantPicker } from "./components/quant-picker";
@@ -83,6 +81,12 @@ import {
mergedFormatPayload,
} from "./constants";
import { useExportSizeEstimate } from "./hooks/use-export-size-estimate";
+import {
+ getCachedCheckpoints,
+ getCachedLocalModels,
+ refreshCheckpoints,
+ refreshLocalModels,
+} from "./export-navigation-cache";
import {
isExportPanelActive,
useExportRuntimeStore,
@@ -172,8 +176,12 @@ export function ExportPage() {
);
// ---- API-driven checkpoint state ----
- const [models, setModels] = useState([]);
- const [loadingCheckpoints, setLoadingCheckpoints] = useState(true);
+ const [models, setModels] = useState(
+ () => getCachedCheckpoints() ?? [],
+ );
+ const [loadingCheckpoints, setLoadingCheckpoints] = useState(
+ getCachedCheckpoints() === null,
+ );
const [checkpointError, setCheckpointError] = useState(null);
const [selectedModelIdx, setSelectedModelIdx] = useState(null);
@@ -185,8 +193,12 @@ export function ExportPage() {
null,
);
const [localModelInput, setLocalModelInput] = useState("");
- const [localModels, setLocalModels] = useState([]);
- const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(true);
+ const [localModels, setLocalModels] = useState(
+ () => getCachedLocalModels() ?? [],
+ );
+ const [isLoadingLocalModels, setIsLoadingLocalModels] = useState(
+ getCachedLocalModels() === null,
+ );
const [localModelsError, setLocalModelsError] = useState(null);
const debouncedModelQuery = useDebouncedValue(modelInput);
const debouncedHfToken = useDebouncedValue(hfToken, 500);
@@ -295,16 +307,15 @@ export function ExportPage() {
// ---- Fetch checkpoints on mount ----
useEffect(() => {
let cancelled = false;
- setLoadingCheckpoints(true);
- setCheckpointError(null);
- fetchCheckpoints()
- .then((data) => {
+ const hadCache = getCachedCheckpoints() !== null;
+ refreshCheckpoints()
+ .then((models) => {
if (!cancelled) {
- setModels(data.models);
+ setModels(models);
}
})
.catch((err) => {
- if (!cancelled) {
+ if (!cancelled && !hadCache) {
setCheckpointError(
err instanceof Error ? err.message : "Failed to load checkpoints",
);
@@ -343,14 +354,15 @@ export function ExportPage() {
// ---- Fetch local models for direct export ----
useEffect(() => {
- const controller = new AbortController();
- void listLocalModels(controller.signal)
+ let cancelled = false;
+ const hadCache = getCachedLocalModels() !== null;
+ void refreshLocalModels()
.then((models) => {
- if (controller.signal.aborted) return;
+ if (cancelled) return;
setLocalModels(models);
})
.catch((error) => {
- if (controller.signal.aborted) return;
+ if (cancelled || hadCache) return;
setLocalModelsError(
error instanceof Error
? error.message
@@ -358,10 +370,12 @@ export function ExportPage() {
);
})
.finally(() => {
- if (controller.signal.aborted) return;
+ if (cancelled) return;
setIsLoadingLocalModels(false);
});
- return () => controller.abort();
+ return () => {
+ cancelled = true;
+ };
}, []);
// ---- Derived state ----
diff --git a/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts b/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts
index 882c530782..d598065d17 100644
--- a/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts
+++ b/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts
@@ -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
-import { useCallback, useEffect, useRef, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
interface HfPaginatedState {
results: T[];
@@ -12,6 +12,10 @@ interface HfPaginatedState {
error: string | null;
}
+interface InternalPaginatedState extends HfPaginatedState {
+ queryKey: object | null;
+}
+
const INITIAL: HfPaginatedState = {
results: [],
scannedCount: 0,
@@ -74,10 +78,15 @@ export function useHubPaginatedSearch(
options?: { enabled?: boolean },
): HfPaginatedState & { fetchMore: () => boolean; retry: () => void } {
const enabled = options?.enabled ?? true;
- const [state, setState] = useState>(
- INITIAL as HfPaginatedState,
- );
const [retryNonce, setRetryNonce] = useState(0);
+ const queryKey = useMemo(
+ () => ({ createIter, mapItem, retryNonce }),
+ [createIter, mapItem, retryNonce],
+ );
+ const [state, setState] = useState>({
+ ...(INITIAL as HfPaginatedState),
+ queryKey: null,
+ });
const stateRef = useRef(state);
useEffect(() => {
stateRef.current = state;
@@ -187,6 +196,7 @@ export function useHubPaginatedSearch(
setState({
...(INITIAL as HfPaginatedState),
isLoading: true,
+ queryKey,
});
const iter = createIter(controller.signal);
@@ -203,6 +213,7 @@ export function useHubPaginatedSearch(
isLoadingMore: false,
hasMore: !done,
error: null,
+ queryKey,
});
})
.catch((err) => {
@@ -214,6 +225,7 @@ export function useHubPaginatedSearch(
isLoadingMore: false,
hasMore: false,
error: err instanceof Error ? err.message : "Search failed",
+ queryKey,
});
})
.finally(() => {
@@ -226,7 +238,14 @@ export function useHubPaginatedSearch(
return () => {
clearDeferredFetch();
};
- }, [createIter, mapItem, enabled, retryNonce, clearDeferredFetch]);
+ }, [
+ createIter,
+ mapItem,
+ enabled,
+ retryNonce,
+ queryKey,
+ clearDeferredFetch,
+ ]);
const retry = useCallback(() => {
setRetryNonce((n) => n + 1);
@@ -358,5 +377,22 @@ export function useHubPaginatedSearch(
};
}, [enabled, fetchMore]);
- return { ...state, fetchMore, retry };
+ const visibleState: InternalPaginatedState =
+ state.queryKey === queryKey
+ ? state
+ : {
+ ...(INITIAL as HfPaginatedState),
+ isLoading: enabled,
+ queryKey,
+ };
+ return {
+ results: visibleState.results,
+ scannedCount: visibleState.scannedCount,
+ isLoading: visibleState.isLoading,
+ isLoadingMore: visibleState.isLoadingMore,
+ hasMore: visibleState.hasMore,
+ error: visibleState.error,
+ fetchMore,
+ retry,
+ };
}
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx
index d57f9636fe..29e2f3f277 100644
--- a/studio/frontend/src/features/hub/hub-page.tsx
+++ b/studio/frontend/src/features/hub/hub-page.tsx
@@ -23,6 +23,10 @@ import {
hfApiToken,
useHfTokenStore,
} from "@/features/hub/stores/hf-token-store";
+import {
+ isChannelEntryFresh,
+ useHubFeedStore,
+} from "./stores/hub-feed-store";
import {
getInferenceStatus,
isExternalModelId,
@@ -80,6 +84,7 @@ import {
} from "./lib/hidden-models";
import { inventoryRowMatches, tokenizeQuery } from "./lib/inventory-search";
import { resolveOwnerProviderLogo } from "./lib/provider-logos";
+import { fingerprintToken } from "./lib/token-fingerprint";
import {
buildDiscoverRows,
detectResultFormat,
@@ -569,6 +574,10 @@ export function ModelsPage() {
const hfToken = useHfTokenStore((s) => s.token);
const debouncedHfToken = useDebouncedValue(hfToken, 500);
const apiHfToken = hfApiToken(debouncedHfToken);
+ const tokenFingerprint = useMemo(
+ () => fingerprintToken(apiHfToken),
+ [apiHfToken],
+ );
const deferredFormatFilter = useDeferredValue(formatFilter);
const deferredCapabilityFilter = useDeferredValue(capabilityFilter);
@@ -633,8 +642,22 @@ export function ModelsPage() {
online,
});
+ const cachedListEntry = useHubFeedStore((state) =>
+ liveListChannel ? state.channels[liveListChannel.id] : undefined,
+ );
+ const visibleResults =
+ results.length === 0 &&
+ liveListChannel &&
+ isChannelEntryFresh(
+ cachedListEntry,
+ liveListChannel.id,
+ tokenFingerprint,
+ )
+ ? (cachedListEntry?.results ?? results)
+ : results;
+
useFeedWriteBack({
- channelId: isChannelListMode ? activeChannelId : null,
+ channelId: liveListChannel?.id ?? null,
results,
isLoading,
accessToken: apiHfToken,
@@ -656,8 +679,13 @@ export function ModelsPage() {
[effectiveCachedRows, effectiveLocalRows],
);
const modelDiscoverRows = useMemo(
- () => buildDiscoverRows(results, effectiveCachedRows, effectiveLocalRows),
- [results, modelDiscoveryInventorySignature],
+ () =>
+ buildDiscoverRows(
+ visibleResults,
+ effectiveCachedRows,
+ effectiveLocalRows,
+ ),
+ [visibleResults, modelDiscoveryInventorySignature],
);
const datasetDiscoverRows = useMemo(() => {
@@ -784,7 +812,7 @@ export function ModelsPage() {
const selectionFilteredDiscoverRows = isFeedMode
? feedRows
: filteredDiscoverRows;
- const selectionResults = isFeedMode ? feedResults : results;
+ const selectionResults = isFeedMode ? feedResults : visibleResults;
const inventoryTokens = useMemo(
() => (isDiscoverTab ? [] : tokenizeQuery(deferredDebouncedQuery)),
From 27b6d553fe9dcd5a55d6e9e005825c7842e7c014 Mon Sep 17 00:00:00 2001
From: Eyera
Date: Tue, 21 Jul 2026 07:53:22 +0200
Subject: [PATCH 046/255] Feat/model picker per model config v2 (#7207)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* refactor(studio): move chat model picker into features/model-picker
Relocate model-selector + its support files from components/assistant-ui
into a self-contained features/model-picker feature (own barrel), mirroring
the modular Hub layout. Pure move + import repoint; no behaviour change.
* feat(model-picker): add per-model config persistence layer
Superset PerModelConfig (customContextLength, kvCacheDtype, speculativeType,
specDraftNMax, tensorParallel, chatTemplateOverride, trustRemoteCode) persisted
to localStorage (unsloth_model_configs) with schema versioning + LRU budget.
KV-dtype and speculative value sets match main's sidebar (no q4_0/ngram-simple).
Reuses features/hub/lib/model-identity for normalization; adds storage-key layer
and applyPerModelConfigToRuntime (sets tensorParallel, which the old PR omitted).
* feat(picker): modular backend for chat-template validate + default fetch
New studio/backend/picker package (schemas/service/routes) mounted at /api/picker:
- POST /api/picker/validate-chat-template (Jinja syntax validation, no false positives)
- GET /api/picker/chat-template/{model_name} (default template from tokenizer_config.json,
reusing get_cache_path/resolve_cached_repo_id_case; graceful null, no model-code exec)
Frontend api/templates.ts client + hooks/use-model-defaults lazy cache. No backend
changes to the existing inference load route (per-model load fields already supported).
* feat(model-picker): bind picker on-device list to shared hub inventory
Picker now sources cached + local models from useHubInventory (the Hub's shared
store) via a thin adapter, replacing its own /api/models/* fetchers + module
caches. Hub, download manager, and picker now share one source of truth, so
completed downloads reflect in the picker automatically. Partial/live-download
rows are filtered from the cached lists (unchanged rendering). Local naming/search
preserved via additive LocalInventoryRow modelId/displayName. Variant expander,
scan-folder management, recommended-fit, search, external providers untouched.
Known minor: cached 'Downloaded date' sort tiebreak degrades to alphabetical
(hub cached rows carry no mtime); default 'recent' (load-time) sort preserved.
* feat(model-picker): per-model config step inside the picker
Picking a (non-external) model now opens an in-picker config view built from
main's current load controls (context length, KV cache dtype, speculative
decoding, draft tokens, tensor parallel) plus a chat-template editor backed by
the picker validate/default endpoints. 'Remember for this model' persists the
config per model+variant; Run forwards the config to the existing load flow via
meta.config. External models bypass the step. Two-view orchestration lives in
model-selector (single interception point); pickers.tsx call sites untouched.
trustRemoteCode dropped from PerModelConfig to preserve main's per-load consent.
* feat(chat): apply/persist per-model config through the load flow
handleCheckpointChange threads meta.config into the selection; stageOrLoad and
the autoload/Hub-run paths now apply the picker config (explicit pick or saved
remembered config) via applyPerModelConfigToRuntime before staging/loading, with
keepSpeculative set so a remembered speculative mode survives the model switch.
Replaces the old remembered-load-settings seeding (resolveInitialConfig now the
single source). SelectedModelInput carries config.
* refactor(chat): remove per-model load config from the right sidebar
The load knobs (context, KV cache, speculative, draft tokens, tensor parallel)
and the chat-template editor now live only in the picker config step. The sheet's
Model section keeps the staged Load/Cancel flow (config is applied at pick time);
sampling params, system prompt, and RAG are unchanged. Deletes the superseded
remembered-load-settings module + the store's applyRememberedLoadSettings action,
removes the now-dead sheet state/imports, and points the settings reset at
unsloth_model_configs. Delete-cleanup deferred (stale config is LRU-capped).
* fix(model-picker): remove leftover sidebar-staging cogwheel + empty Model section
The downloaded-variant gear (ModelLoadSettingsAction) staged a model straight
into the right-sidebar Run-settings flow -- the old 'configure before load' path
now fully replaced by the in-picker config step. Removed the gear + its component.
Also gate the sheet's 'Model' section to staged picks only (pendingSelection):
after the load-knob strip its content is staged-only, so it was rendering an
empty section header whenever a model was merely loaded.
* chore(chat): remove dead per-model-config setters + modelControlsDisabled
After the load-config UI moved into the picker, the store's per-model setters
(setKvCacheDtype/setSpeculativeType/setSpecDraftNMax/setTensorParallel/
setCustomContextLength/setChatTemplateOverride) had zero callers
(applyPerModelConfigToRuntime writes via setState), and the sheet's
modelControlsDisabled was unreferenced. Verified dead across the whole tree.
* fix(chat): config-step Load actually loads (ignore Load-on-selection)
Root cause: with Settings > Chat > 'Load on selection' turned OFF, the config
step's load went down the deferred-staging path -- opening the right sidebar with
' is staged, not loaded yet / Choose Load model'. The in-picker config step
IS the deliberate load action, so its Load now loads immediately (or downloads +
auto-loads when not cached) regardless of the toggle. Renamed the button
'Run model' -> 'Load model' to match. Native/dropped picks still honor the toggle.
* refactor(chat,hub): retire 'Load on selection' — config step is the only load flow
The in-picker config step (and the Hub Run button) now fully supersede the old
stage-to-sidebar flow, so the Load-on-selection toggle is removed everywhere:
- chat stageOrLoad: every pick loads immediately, or downloads + auto-loads when
not cached (the previous default behaviour, now universal).
- hub Run: drops the stage branch; downloaded GGUFs load directly with their saved
per-model config (no collision with the chat config step — both end at selectModel).
- store: removed loadOnSelection field/setter/key/default; Settings>Chat toggle and
its settings-reset entry removed.
- staged sidebar section is now a download-progress view (auto-loads on completion).
No manual staging remains; stageModel is used only for background auto-load downloads.
* feat(model-picker): default chat template from GGUF + thread variant through config flow
Read the embedded tokenizer.chat_template from GGUF files (read_gguf_chat_template
in gguf_metadata) and use it as the per-model default. Plumb gguf_variant through
the picker service, /api/picker/chat-template route, frontend templates API, and
use-model-defaults so the right variant's template is fetched.
Also refine the picker config-page/model-selector wiring, drop the dead
ggufNativeContextLength runtime path, and add the per-model-config storage keys to
the settings prefs export.
* feat(model-picker): read safetensors chat template + hide editor where it has no effect
Resolve the default chat template for safetensors models: prefer the modern
chat_template.jinja, fall back to the tokenizer_config.json chat_template field,
then chat_template.json (multimodal processor), then the GGUF embedded template.
Applied to local dirs, the HF cache snapshot scan, and the HF remote fetch.
Hide the chat-template editor in the picker for safetensors models — the override
is only applied at load by the GGUF/llama.cpp backend, so editing it on safetensors
currently has no effect. GGUF keeps the editor. Nothing removed; the dialog stays
for when the safetensors apply path is wired up in a later branch.
* fix(model-picker): set legacy-migration flag only after the write succeeds
Set unsloth_model_configs_migrated only once writeMap confirms the migrated
map persisted, so a quota/storage failure no longer marks migration done and
silently drops the user's pre-existing remembered settings — the next load retries.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* MVP model picker fixes
* MVP picker config fix
* MVP safetensors config
* MVP max seq config
* MVP max seq fix
* Fix static max tokens cap ignoring model context
* Fix picker GGUF scan parity
* fix(studio): harden model picker config loading
Apply remembered per-model configs consistently from picker and Hub loads, keep default configs from overriding standing speculative settings, add config access for direct local GGUF files, and support saving or forgetting active model settings without a reload.
* Fix model picker config flow
* Fix model picker config loads
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Avoid recursive per-model config migration reads
* Apply the displayed context length when loading a GGUF
* Fix template validation, cached template lookup, and failed load rollback
- Validate chat templates with the loopcontrols extension so templates
that use break or continue tags pass the picker validator, matching the
inference renderer that already accepts them.
- Read the default chat template from the newest cache snapshot rather than
an arbitrary iterdir order, so an older cached revision no longer prefills
a stale template.
- Capture the runtime per-model config before a load and reapply it when the
load fails, so a failed switch leaves the active model context, KV cache,
template, and speculative settings as they were.
* Make chat template view only for safetensors models
Custom chat template overrides are applied at inference only for GGUF
models, which pass the template to llama-server. The safetensors backend
renders with the model built-in template and ignores the override, so
editing it would save a value that never loads. For safetensors the
config page now opens the template as a read-only preview with a note
that editing is not available yet. This can become editable once
inference support for custom safetensors templates lands in main.
* Fix model picker config edge cases
- Restore prior runtime config when a load no-ops for the active model
- Cap the picker validator request body via the protected prefixes
- Keep the GGUF context slider max above the loaded context
- Fetch subfolder chat templates for uncached Hub repos
- Show the compare side config when reopening the picker
* Keep saved GGUF context above the fallback ceiling
* Show the model config in the run settings sidebar
* Fix model config sidebar reset and context slider
- Stack the remember toggle and action buttons in the sidebar
- Reset the config to defaults instead of the loaded values
- Fetch the native context so the slider max is not the loaded value
* Fix model picker config and download regressions
- Run picker chat template routes off the event loop
- Depth and root guard local template directory scans
- Restore download manager flow for uncached hub picks
- Apply per model context length on reload
- Import model picker symbols from the feature barrel
* Fix model picker config and cached download sorting
- Restore load settings when a Hub run is rejected mid load
- Reuse one NumericValueInput instead of a duplicate copy
- Fix double decode of the model name in the template route
- Remove the unused reset-to-loaded settings action
- Fix cached model download sorting
* Fix model picker per-model config edge cases
Honor a saved or typed max seq length above the model's native context so
RoPE extended values are no longer clamped and silently overwritten. Allow
typing past native while the slider keeps native as a soft ceiling.
Guard the fetch success paths in use-model-defaults against an aborted
signal, and refetch when the HF token changes.
Hash the chat template content in the sidebar remount key instead of its
length. Enable reset for a GGUF whose native context is unknown, and floor
the context slider max so it can never fall below the min.
* Fix GGUF context auto-fit and gated model config token
Stop forcing a 32768 context when a GGUF native context is unknown so the backend auto-fits to VRAM again, while still honoring an explicit context edit.
Send the HF token as a query param so gated safetensors models resolve their max position embeddings.
Derive model default state during render to drop the set-state-in-effect calls.
* Fix native GGUF context ceiling and guard picker template reads
Restore the native context store field so the sidebar slider keeps the
full ceiling for drag and drop GGUFs. Limit local chat template reads to
the browse allowlist, skip malformed repo ids, and drop unused model
picker exports.
* Fix model picker lint boundaries
* Fix model picker review findings
Chat template editor never seeded its draft. Radix only calls onOpenChange
from internal events, so the seed in the nextOpen branch was dead and a model
with a saved override opened empty. Saving then cleared the override. Drop the
dead branch, treat draft as an untouched sentinel, and reset it on every close.
Uncached Hub picks could auto load a model after the user left the chat. Main
detached the staged pick on route exit and on chat context change. Carry the
context key on the pending pick and skip the load when it no longer matches.
Also clear configTarget when the picker closes, restore the onUpdated ref so
variant rows stop resubscribing on every parent render, skip the LRU write when
the entry is already most recent, import NumericValueInput relatively, and drop
the unused ModelUpdateAction barrel export.
* Preserve GGUF context on active reload
* Fix model picker per-model config regressions
- Stop reloading the already loaded model on re-pick
- Hide infra models from the chat picker
- Detect vision support on cached GGUF repos
- Honor saved maxSeqLength on auto load
- Restore default chat template for local GGUFs
- Warn on save failure and revert config on cancel
- Refetch picker inventory on open
- Persist read only per model config safely
* Fix stale model auto load
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix model picker numeric input sizing and constraints
Size value inputs to their content so long context lengths are not clipped,
restrict them to numeric characters, and stop the speculative decoding label
from truncating in the sidebar.
* Fix picker CI tests and harden chat template resolution for PR #6647
- tests: point the descender guard at the moved model-selector.tsx path
- tests: exclude the disabled Reload model button from the regenerate locator so .first targets the real Regenerate
- picker/service.py: reject symlinked template/gguf leaves that resolve outside the browse allowlist (HF cache reads unchanged)
- compare mode: resolve each pane's own remembered chat template instead of inheriting the other pane's from the store
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Protect future-schema per-model configs from deletion for PR #6647
savePerModelConfig already refuses to overwrite a stored config whose schema version is newer than this client understands, but deletePerModelConfig did not. Unchecking Remember on an older client therefore silently destroyed a newer client's saved config. Apply the same guard on delete and surface the blocked case through the existing saveFailed toast.
* Protect future-schema per-model configs from quota eviction for PR #6647
The save and delete guards already refuse to touch a stored config whose schema version is newer than this client understands, but the quota-eviction path did not, so a full store on an older client could still evict a newer client's config. Skip future-schema entries when evicting and fail the save if the budget cannot be met without them.
* Fix GGUF context persistence, compare context, and rollback settings for PR #6647
Persist a GGUF context override from the user's intent instead of collapsing it against the loaded context, which reintroduced the context-reset (f4838782cb reverted the native-baseline fix). model-config-page now collapses the saved value against native, and use-chat-model-runtime and chat-adapter retain the requested context on load so re-saving another setting keeps the override; a null request stays null so a VRAM auto-fit never becomes a stored override.
shared-composer: a compare pane with no explicit GGUF context now loads at native (0) like single-view, not the session maxSeqLength that silently shrank the shown context.
use-chat-model-runtime: restore the previous model's KV cache dtype and chat template on a failed-load rollback so it runs as it was, not with backend defaults.
* Preserve native path token when reloading the active model for PR #6647
handleReloadActiveModel rebuilt the selection without the store's activeNativePathToken, so reloading a file-picked GGUF after a settings change validated the display label as a repo/path and failed. Thread the active native token through the reload selection so native-loaded models reopen correctly.
* Make picker template validation resilient and accept HF generation tags for PR #6647
Import Jinja lazily inside validate_chat_template so a backend without the optional jinja2 package (GGUF-only installs) still starts instead of raising ModuleNotFoundError at import time. Register a no-op extension for the Transformers {% generation %} assistant-mask tag so pasting a valid HF chat template validates, matching the renderer, rather than being rejected as an unknown tag.
* Honor remembered compare config and parse processor chat_template.json for PR #6647
* Fix failed-load rollback context and processor template map fallback for PR #6647
* Restore speculative decoding config on failed-switch rollback
When a model switch fails after the previous model was unloaded, the
rollback reload restored tensor_parallel, KV cache dtype and the chat
template override, but omitted speculative_type and spec_draft_n_max and
cleared their loaded shadows to null. The previous model therefore came
back running at backend defaults (speculation off) while the UI still
showed it enabled, and the status resync confirmed the off state. Resend
the previous model's speculative settings in the rollback load and keep
the store's active and loaded speculative fields in sync with them.
* Reset max sequence length when a model has no saved config
applyPerModelConfigToRuntime reset every per-model field except
maxSeqLength, which it only wrote when the incoming config had one.
maxSeqLength is the sole field carried on store.params, so selecting a
model with no remembered config left the previous model's value in place
and later loaded the new model at that leaked length. Fall back to the
standing default so an unremembered model loads at its own default.
* Surface a message when a variant update cannot start
startManagedUpdate handled the conflict and error start outcomes but let
busy fall through as if the update began, so the confirm dialog closed
with no job created and the cached variant stayed stale. Show an info
message when the repo is busy with a sibling transfer so the click is
not silently dropped.
* Keep per-model speculative choices out of the global default
A staged load with a per-model or one-off config sets keepSpeculative,
which already skips reading the global speculative preference. The
matching save still ran unconditionally, so the model-specific choice was
written to the global unsloth_chat_speculative_type and a later model with
no saved config started from it instead of Auto. Skip saveSpeculativeType
when keepSpeculative so the per-model choice stays isolated.
* Seed non-active model settings from the app default max length
The Run settings page captured initialMaxSeqLength from the loaded
model's runtime params and fell back to it for a model with no saved
config. Opening settings for a different, unloaded model and clicking
Load then sent the active model's context (for example 64k) instead of
the 4096 default, risking validation failures or OOMs. Seed the default
for non-active models and keep the runtime value only for the active one.
* Prefer sidecar tokenizer chat template over the GGUF copy for variants
_chat_template_from_dir returned the embedded GGUF template first when a
variant was selected, reversing the tokenizer-first precedence of the
no-variant path. A model whose chat_template.jinja or tokenizer_config.json
supersedes a stale embedded template then got the wrong template on
variant selection. Keep tokenizer files first regardless of variant; the
variant only picks which GGUF is the fallback. Adds regression tests for
both the tokenizer-wins and gguf-fallback cases.
* Keep per-model speculative choices load-local in autoload and compare
The interactive load path treats a per-model speculative choice as
load-local and skips writing it to the global default. Autoload and
generalized compare still called saveSpeculativeType unconditionally, so a
remembered off or ngram setting leaked into unsloth_chat_speculative_type
and later models with no saved config inherited it. Persist the global
preference only when the value came from the global settings.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: record the compare pane's loaded context in runtime state so the active model's settings and any reload or save use it, not the previous context
* Studio: notify the user when a Hub autoload can't start because another download for the model is already running, instead of silently dropping it
* Studio: drop the merge's orphaned staged-model store helpers and unused alert imports
The main merge left isPendingGguf and pendingSelectionMatches referencing the
removed PendingModelSelection type, and the alert-dialog/alert imports unused
after the permission-mode dropdown replaced the bypass dialog, so tsc -b failed.
* Studio: cache a null default chat template so the viewer stops re-fetching it
A model with no sidecar or embedded template resolves to a terminal null, but
that result was never cached, so reopening the template viewer re-ran the
backend and Hugging Face lookup every time.
* Studio: detect direct-file GGUFs in run settings so Max Tokens uses their context
A GGUF loaded from a local file or custom folder has no variant label, so the
run-settings panel treated it as non-GGUF and clamped Max Tokens to the session
max_seq_length instead of the loaded GGUF context. Detect it via the reported
GGUF context and the .gguf checkpoint suffix, matching the chat page.
* Studio: prompt to re-select a local model file when its lease expired before reload
A file-picked GGUF is reachable only through a native path token that the
desktop host prunes after a TTL. Reloading reused that token blindly, so a
reload long after the initial load failed with an opaque error. Track the
token's expiry and, when it has passed, ask the user to re-select the file
instead of attempting a doomed reload.
* Fix descender-clipping test to tolerate sidebar layout utilities
The sidebar account-block div carries layout utilities (min-w-0, flex-1)
between 'flex' and 'flex-col', so the descender-clipping guard's regex,
which required 'flex' immediately followed by 'flex-col', no longer matched
and the test failed to locate the account-block div. Generalize the prefix
to allow intervening flex utilities while still capturing the leading-*
class before the collapsible visibility utility and asserting leading-tight,
so the guard against clipped glyph descenders is fully preserved.
* Harden picker chat-template resolution
Enforce the 64 KiB chat-template contract at the validate endpoint's request
model so a direct caller cannot submit a template far larger than the frontend
allows (MaxBodyMiddleware only bounds the whole request body, not this field);
oversized templates now return a clean 422.
Apply sidecar-over-GGUF template precedence globally across cached snapshots
instead of per snapshot. A repo with multiple cached revisions previously
returned the first snapshot's template, so a newer GGUF-only revision could
win over an older revision's maintained chat_template.jinja sidecar, which
contradicted the documented intent that sidecars supersede the embedded copy.
* Guard per-model config against future-schema and lossy migration
Two forward-compatibility gaps in the versioned per-model config store:
- The load/apply path returned and normalized a stored record without checking
its schema version, so a record written by a newer client was reinterpreted
under the current schema and applied to a live model load, even though save,
delete and eviction all refuse to touch future-schema records. Reject
future-schema records on load too.
- The one-time legacy migration enforced the storage budget without protecting
the entries it had just migrated and set the completion flag unconditionally.
When storage was already full of future-schema records (which are unevictable
by an older client), the migrated entries were the only evictable ones and
could be dropped while migration was still marked complete. Protect the
migrated keys during eviction and only mark migration complete when they
survive, so it retries once space frees up.
* Discard chat-template validation results after the dialog closes
Server-side template validation is async, but closing or cancelling the editor
did not abort it, so a late-arriving valid response still called onSave and
applied a template the user had already dismissed. Track a validation token
that is bumped on close and ignore any validation result whose token is stale.
* Record native lease expiry when loading a picked GGUF from the chip
The pending-native-model chip loaded via stageOrLoad directly, bypassing
loadNativeModelIntent, so activeNativePathExpiresAtMs was never recorded for a
chip-loaded file. A later reload then either skipped the lease-expiry guard
entirely (expiry left null) or compared against a previously loaded file's
stale expiry, so reload could reuse an already-pruned token or wrongly block a
still-valid one. Route the chip through loadNativeModelIntent, which builds the
same selection and records the expiry.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Prefer sidecar template for a directly selected local GGUF file
A direct .gguf file path read its embedded chat template without checking the
parent directory for a maintained sidecar (chat_template.jinja /
tokenizer_config.json), while directory and variant selections already prefer
the sidecar. That let the config editor preview or save a stale embedded
template for the same model depending on how it was selected. Check the parent
directory sidecars first, then fall back to the embedded copy, and cover both
paths with tests.
* Resolve cached chat template per revision, newest first
The earlier change searched every cached snapshot for a sidecar before
considering any snapshot's embedded GGUF template, which let an obsolete sidecar
from an older revision override the newest revision's template. Restore
per-snapshot resolution (newest first): a revision's sidecar still supersedes
its own embedded GGUF copy, but a newer revision is no longer overridden by an
older revision's sidecar.
* Preserve autoload transport conflicts and surface background busy downloads
- When a Hub autoload hits a transport conflict, keep pendingHubAutoLoad bound
instead of clearing it. Clearing it re-keyed the download surface and its
cleanup cancelled the conflict the toast tells the user to resolve, so the
Hub resume affordance was gone the moment it appeared. Return early on
conflict, mirroring the started branch, so resolving it from the Hub still
auto-loads on completion.
- The background-download branch handled started and conflict but silently
dropped a busy outcome, leaving the user with no feedback when a peer variant
of the same repo was already downloading. Surface the same busy toast the
autoload path uses.
* Fix context length, GGUF template, fetch state and lease expiry bugs
Keep explicit context length values instead of collapsing to null at
native. The collapse made the slider jump back at the native maximum
and made Reload load the previous context instead of the chosen one.
Prefer the first split when resolving a GGUF without a variant. Later
splits carry no chat template metadata, so picking the largest file
could return no template for a sharded model.
Clear stale fetch state when template and metadata lookups retry, so
a previous terminal error is not shown while a new fetch is running.
Record native path lease expiry together with the token when a load
commits. The expiry was written by only one load path and even when
the load did not start, so a reload could be blocked with an expired
file message for a still valid token.
* fix(model-picker): resolve review findings across config, inventory, and templates
- Apply remembered per-model config in the training-compare chat handoff so a
prior model's customContextLength no longer leaks into the next load
- Match GGUF variant labels with the inventory extractor too, so cached
no-quant-token files resolve their default chat template
- Show "Auto" instead of a fabricated 32768 when native context is unknown
- Reuse the identical staged auto-load object on same-pick so a re-pick during
download pre-flight no longer disarms auto-load via "busy"
- Union supports_vision when deduping cross-cache inventory rows
- Serve hidden-model needles from a new GET /api/hub/hidden-models endpoint and
merge them client-side, covering runtime-configured RAG embedders
- Clamp GET chat templates to MAX_CHAT_TEMPLATE_BYTES (route + jinja sidecar),
matching the validate endpoint's contract
- Lower-clamp stored customContextLength to shared CONTEXT_LENGTH_MIN
- Wipe unsloth_chat_load_on_selection in Settings "Reset all"
- Drop stale pendingHasContext comment describing deleted staging machinery
* Fix stale defaults cache, token in query string and rounded up context ceiling
Refresh cached chat template and max position data when a model update
completes. Send the HF token for model config requests in the dedicated
header instead of the URL. Snap the native sequence length ceiling down
to the nearest step so the slider cannot exceed the declared maximum.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix compare pane reverting active checkpoint on non-GGUF load
Re-read runtime params after setCheckpoint so the fresh checkpoint is
kept instead of being overwritten by the pre-setCheckpoint snapshot.
* Send the HF token via header for the vision and embedding checks
checkVisionModel and checkEmbeddingModel still passed the HuggingFace
token as a ?hf_token= query parameter, so it landed in server access
logs, proxy logs, and browser history. Move them to the
X-Unsloth-HF-Token header like getModelConfig already does, and accept
the header on the check-vision and check-embedding routes with the
existing query parameter kept as a fallback for older clients.
* Cap the chat template on the model load path
The load endpoint accepted an unbounded chat_template_override, so a
direct caller could hand llama.cpp an arbitrarily large Jinja template
even though the frontend, the validate endpoint, and the read paths all
enforce the 64 KiB limit. Reuse MAX_CHAT_TEMPLATE_BYTES in the
LoadRequest validator, rejecting oversized templates with a fast
character-count check before the exact UTF-8 byte check.
* Protect existing per-model configs during legacy migration
When the one-time legacy import pushes the store over budget, eviction
now protects the entries the user already has and drops only the
just-migrated legacy entries, so importing old load settings can never
discard a newer per-model config.
* Reset clears the context override instead of pinning the native value
Reset wrote the discovered native context into customContextLength for
GGUF models, but isDefaultConfig treats any non-null customContextLength
as an explicit pin, so Reset with Remember enabled persisted a fixed
context and future loads stopped using the native auto context. Reset
now restores the full default (customContextLength null); the native
value is still shown through the existing display fallback.
* Bound chat-template sidecar reads to a size limit
The chat_template.json, tokenizer_config.json, and Hub-downloaded sidecar
readers decoded and json-parsed the whole file before the extracted
template hit the 64 KiB response cap, so an oversized metadata file could
exhaust memory. Read them through a bounded reader (4 MiB envelope) that
returns None when the file is larger, matching the existing chat_template.jinja
size guard. Adds tests for oversized tokenizer_config.json and chat_template.json.
* Keep the native-path token and lease expiry in sync
Rollback after a failed reload restored the previous token but left the
failed load's expiry in the store, so a later reload could be falsely
blocked as expired (token A paired with load B's lease). Restore the
previous lease alongside the token, and clear the expiry wherever the
token is cleared on a non-GGUF transition, so the two never diverge.
* Clear the native file lease on compare-pane loads
* Studio: add regression tests for the model-picker per-model-config
Guard the specific regressions that reverted the predecessor change:
- backend pytest (studio/backend/tests/test_model_picker_regression.py):
infra-model hiding, HF token via header with query fallback, and the
chat-template byte caps.
- source contracts (tests/studio/test_model_picker_contracts.py): the token
stays out of the URL, the context ceiling is floored, the native lease is
cleared on compare-load and restored on rollback, the default caches key on
the inventory version, and the hidden needles stay present.
- Playwright E2E (tests/studio/playwright_model_config.py) wired into
studio-ui-smoke.yml on port 18898: Context Length persists across a reload,
Reset clears the stored override, and infra models are absent from the picker.
- optional GPU-gated inference smoke (tests/studio/test_gpu_inference_smoke.py)
that auto-skips on GPU-less CI and stays short on a GPU.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: model pinning, row menus, hub inference settings, and inventory filters
Pinning
- Add a pinned models store (localStorage) with repo and per-quant pins
- Pinned section in the model selector's On Device list and the hub inventory,
with newest pins first so Pin to top lands on top
- Deleting a repo drops its pins
Row menus
- Replace loose row icons with a shared 3-dots menu (pin, reveal in file
manager, copy identifier, copy path, delete) on picker rows, hub quant rows,
the hub run bar, and on-device inventory rows
- Menus only render for models actually on disk; platform-aware reveal labels
- Backend: cached-model-path and reveal-cached-model endpoints resolving
managed HF-cache repos only
Hub inference settings
- Gear in the GGUF run bar opens an Inference settings dialog reusing the chat
page's controls: model config (context length, KV cache, speculative
decoding, chat template), system prompt, reasoning, sampling, tools and
retrieval
Inventory
- Model-type filter (text, vision, embedding, STT, TTS, diffusion) beside the
sort pill, both with a sort icon, capped widths and truncation so the
On device heading never wraps
- Unsloth-owned repos without an upstream provider logo fall back to the
Unsloth mascot avatar
- Discover / On Device tabs widened; hub search bar narrowed to match
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: revert the Unsloth mascot avatar fallback
Unsloth-owned repos without an upstream provider match go back to the
colored-initial tile, and unslothai is no longer a relabeled owner.
* Studio: run-bar options on single models, and aligned type/capability filters
- Give single-model (non-GGUF) run bars the same 3-dots options menu and
settings gear as GGUF, at repo level
- Drop Pin to top from the run-bar menus; pinning stays in the On Device list
- Add an Image to text (diffusion) capability with detection, and surface it
in both the hub Discover capability filter and the On Device type filter
- Align the On Device type filter with the Discover capability options and
share the same detection so both dropdowns match
* Studio: apply hub inference config on reload, eject action, and run-bar polish
- Fix inference settings not applying: the hub dialog now writes the config to
the runtime before reload, matching the chat page (selectModel reads runtime
state, not the selection)
- Order the settings gear before the 3-dots menu in the run bars
- Replace the loaded-model run-bar action (New Chat) with Eject, wired through
the inspector to the hub's ejectModel
- Truncate the results heading so a long search query clips instead of
overlapping the header pills in split view
- Use a plain magnifying-glass icon for the no-results empty state
* Studio: fix GPU settings loss, load guards, pins, filters, and cached paths
Reloading a model from the chat sidebar or the hub gear dialog rebuilt the
per-model config without the GPU memory fields, so manual GPU layers, MoE
placement, and the GPU pick were reset on every reload and could be saved
over a remembered config. The active config now comes from a shared
useActiveModelConfig hook that carries the GPU fields for GGUF models, and
the sidebar remount signature tracks them through a shared gpuFieldsSignature
helper.
The in-flight load guard lived in a ref inside each useChatModelRuntime
instance, so the chat page, hub page, and gear dialog could not see each
other's loads. A load started from the gear dialog left the hub page free to
eject the model mid-reload or start a second concurrent load. The runtime
store now records the loading pick, selectModel checks it across instances,
and ejectModel refuses with a toast while any load is in flight.
The cached-model-path endpoint matched GGUF files by basename and excluded
only mmproj, so Copy path and Reveal could return an MTP drafter for a quant
and returned 404 for directory layouts like BF16/model-00001.gguf. Variant
files are now resolved from snapshot-relative paths with the same drafter,
mmproj, and big-endian exclusions as the load path, shared through a new
_main_variant_gguf_label helper.
Hub and picker fixes:
- rename the diffusion capability label from "Image to text" to
"Image generation", since it detects image generators
- validate pinned quants through the cached variant listing, keep the last
verified set while revalidating, and drop deleted quants immediately
- pass a measured scroll margin to the on-device virtual list so rows past
the overscan stay visible below the pinned block
- keep the delete menu for stopped partial safetensors downloads
- give the inventory type filter a reset in Clear filters, a truthful empty
state with a Show all types action, and hide it on the datasets view
- order picker pinned rows by pin recency, include pinned matches in the
empty-state check, and sync pins across browser tabs
- count only the visible rows in the On device list header
Tests: contract checks for each fix in test_model_picker_contracts.py and a
backend test for the variant label selection.
* Studio: reveal cached models in Windows Explorer under WSL
The reveal endpoint only branched on macOS, Windows, and generic Linux.
Under WSL the Linux branch spawned xdg-open, which is missing on a stock
distro without a Linux desktop, so the request failed with a 500 and the
UI showed a failed to open file manager error.
WSL is now detected with the existing helper and the path is converted
with wslpath before opening explorer.exe, selecting the file the same
way native Windows does. Directories open directly. When interop is
unavailable the old xdg-open fallback still runs. The macOS, native
Windows, and native Linux branches are unchanged, and the Tauri app is
covered since its hub reveal calls this same local endpoint.
Tests: platform guards for the WSL reveal, the interop fallback, and
the unchanged native Linux behavior in tests/studio/test_reveal_file_manager.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Adjust model picker row spacing and cogwheel hover consistency
* Studio: exact hidden model ids and newest revision cached paths
A custom RAG embedder repo was published to the frontend as a basename
substring needle, so a generic name like org/model could hide unrelated
models in the pickers. The hidden-models endpoint now sends full repo ids
that are matched exactly.
Copy path and Reveal picked a GGUF variant from an arbitrary cache
revision when the same file existed in more than one. The newest revision
now wins, matching the whole repo lookup.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix model picker GPU config, metadata, and cache selection
Load each compare model with its saved GPU memory mode, GPU layers, CPU MoE layers, and selected GPU IDs. Reconcile saved GPU IDs with the current hardware. Include the active native GGUF path token in metadata checks. Search all Hugging Face cache roots when resolving cached models and select the largest visible cache entry. Remove obsolete barrel exports and the staging-only GPU memory helper.
* Studio: hide hub inference settings gear for now
The cogwheel in the hub download cards is out of scope for this PR. The
dialog component stays in place and a TODO marks where the button
returns in a future PR.
* Refresh hidden model matchers
* Fix GGUF detection, compare context pin, and picker delete staleness
Treat any pick with a GGUF variant as GGUF in selectModel so the first
load after downloading an uncached quant validates and sizes with the
right GPU settings instead of unloading the current model on a wrong
preflight. Variant picks now also set isGguf on their selection meta.
Stop compare panes from inheriting the active model's context pin when
their own saved config says Auto. Null context in a remembered config
now means no pin, matching how the pane settings are shown.
Route picker deletes through the hub inventory client, which
invalidates the HF cache scan and the variants cache. The legacy
delete route left the scan cache warm, so deleted models reappeared
in the picker until the TTL expired. Removed the now unused legacy
delete client and updated the contract test to match.
* Studio: fix stale GGUF load-marker ordering test
The load-in-flight marker still precedes the hub-download guard and the
unload, but the llama_extra_args inheritance that used to sit between the
marker and the guard now runs ahead of the GGUF branch, so it is no
longer a landmark inside the sliced source. Drop it from the ordering
assertion and keep the marker -> guard -> unload invariant.
* Studio: fix per-model config edge cases in compare loads and saved defaults
- chat-settings-sheet: gate the MTP fallback note and context/VRAM warning on
the broader isGguf (variant, loaded gguf context, or .gguf suffix) instead of
isLoadedGguf, so direct-file and custom-folder GGUF loads still surface
those diagnostics.
- shared-composer: a compare pane's context now comes from its own config only
(a saved pin, else null for Auto/native). It no longer inherits the active
model's shared snapshot, which resolveFitMaxSeqLength treated as an explicit
pin and could load a pane at another model's context (VRAM/OOM), matching the
single-model load path.
- model-config-page: when an auto-fit GGUF is saved with fixed GPU layers
(Manual) and Remember, pin the displayed fitted context so a later fresh load
keeps the placement instead of sending native/0 and recreating the OOM.
- per-model-config: treat Auto GPU memory mode and Auto/default speculative type
as follow-global defaults; do not persist them as per-model overrides so later
global preference changes keep applying.
* Studio: gate vision capability on GGUF projectors and bound remote template downloads
- cache_inventory: only mark a cached repo vision-capable when it holds an actual
GGUF mmproj projector, not any file whose name merely contains "mmproj" (e.g.
mmproj_config.json), matching the runtime's GGUF-only projector detection.
- picker/service: pre-check the remote file size before downloading an uncached
repo's chat template / tokenizer config, so a maliciously large sidecar is
skipped instead of fetched and retained in full, mirroring the size gate the
local-file path already applies.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: add source-contract guards for the per-model-config edge-case fixes
Guard the four per-model-config fixes against silent regression in CI:
- local GGUF diagnostics gate on the broad isGguf, not the variant-only isLoadedGguf
- fixed-layer GGUF saves pin the displayed context
- Auto GPU mode and Auto/default speculative are not persisted as per-model overrides
- a compare pane's context comes from its own config, not the active model's snapshot
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: clear manual GPU knobs on Default and resolve local embedders before repo-id
- model-config-page: switching GPU Memory back to Default now clears the Manual-only
knobs (gpuLayers/nCpuMoe/selectedGpuIds); otherwise a remembered config kept stale
pins that a later load re-applied when the global GPU preference was Manual, despite
the page showing Default.
- routes/models hidden_model_matchers: resolve an existing local path before the repo-id
regex, mirroring is_hidden_model, so a local embedder shaped like "models/embedder" is
hidden by exact path instead of leaking as a chat model.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: add _is_mtp_drafter to the model_config stub in the export-paths test
routes/models.py imports _is_mtp_drafter from utils.models.model_config at module
load, but the lightweight stub in test_export_absolute_paths.py did not provide it,
so loading the module under the stub raised ImportError on Backend CI. Add the stub.
* Studio: read a picked GGUF's chat template through the native path lease
The picker chat-template GET has no native-path-lease plumbing, so a
desktop-picked (drag-drop) GGUF could not show its default chat template
in Run Settings until the model was loaded: the endpoint only receives
the display label, not the leased file path.
Read the embedded template through the existing lease-aware
/api/inference/validate probe instead. A new include_chat_template flag
resolves the granted canonical path and returns the GGUF's own embedded
template, never a sibling sidecar (the grant authorizes just that one
file); it skips the training guard like include_context_length and is
bounded by MAX_CHAT_TEMPLATE_BYTES. The frontend fetch mints a one-shot
validate-model lease when a native token is present and keeps the plain
GET path for HF and allowlisted local models.
Adds backend and source-contract regression tests.
* Studio: call worker.direct_wheel_url in the ROCm wheel-url test
The ROCm Mamba/SSM test referenced worker.py's private _direct_wheel_url,
but the worker imports the wheel helper under its public name
direct_wheel_url (utils.wheel_utils). When the worker module loads (its
imports resolve in CI), worker_mod._direct_wheel_url raised AttributeError;
the test only masked it by skipping when the worker could not be imported.
Call the name that actually exists so the assertion runs; it still returns
None for an empty cuda_major (ROCm).
* Studio: reset max sequence length to the app default, not the loaded value
For a non-GGUF active model, the per-model config seeds maxSeqLength from
the loaded runtime value so the panel opens showing the running context.
Reset set config.maxSeqLength to null, but the null fallback resolved back
to that captured runtime value, so the field kept showing the old custom
length and the config saved/reloaded it again. A remembered or active
max-length override therefore could not be cleared from Run settings.
Fall the null/default case back to the app default (clamped to the model's
native ceiling) instead of the active runtime snapshot, so Reset actually
clears the override. The initial view is unaffected: an active model's
config.maxSeqLength is already non-null, so it still shows the loaded value.
Adds a source-contract regression guard.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: persist default max length, refresh deleted quants, hide non-chat locals
Three follow-up fixes from review of the per-model-config picker:
- Max sequence length: the persisted per-model record now keeps config's
maxSeqLength (null after Reset) so isDefaultConfig can clear a remembered
override; the resolved app-default is substituted only into the load
request, never the saved record. Previously Reset saved the concrete
default and left the model pinned/remembered.
- GGUF variant expander: deleting a downloaded quant from a repo that still
has other cached quants now bumps the expander refresh key, so the removed
quant stops showing as downloaded and clickable (which would try to reload
the deleted file) until the repo is collapsed and reopened.
- Local picker rows: require capabilities.canChat before listing a local
models-folder / LM Studio row. A weightless folder (only config.json) is
classified non-chat, and toLocalModelInfo drops capabilities, so selecting
such a row would try to load a path the inventory already marked non-chat.
Adds source-contract regression guards for all three.
* Fix compare-pane and Reset context defaults in model picker
Two related per-model-config default regressions:
- A non-GGUF compare pane with no saved maxSeqLength fell back to the
active model's shared runtime snapshot, so comparing a saved 128K model
against an unconfigured pane loaded the latter at 128K and could OOM. It
now falls back to the shared app default (DEFAULT_MAX_SEQ_LENGTH), the
same fallback the single-model config path uses.
- contextAtDefault treated an explicit customContextLength equal to the
native ceiling as a default, which wedged the Reset button disabled for
a deliberate pin-to-native. It now counts as default only when there is
no override at all.
DEFAULT_MAX_SEQ_LENGTH becomes a single exported constant in
per-model-config.ts so the single-model config and the compare path share
one source of truth. Adds source-contract guards for both fixes.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip over-cap remote Jinja templates so the tokenizer template wins
The remote chat-template resolver bounded raw chat_template.jinja downloads
only by MAX_TEMPLATE_METADATA_BYTES (4 MiB), then returned the first
non-empty Jinja unconditionally. The picker route drops any template larger
than MAX_CHAT_TEMPLATE_BYTES (64 KiB), so an uncached repo whose
chat_template.jinja sits between 64 KiB and 4 MiB returned no template at
all, even when a valid smaller tokenizer_config.json template existed. The
local path already skips oversized .jinja files and falls through.
Gate the extracted Jinja on MAX_CHAT_TEMPLATE_BYTES and continue searching
when it exceeds the cap, matching _chat_template_from_jinja_file. The 4 MiB
download bound stays for JSON files that merely embed a small template. Adds
a regression test that a big Jinja plus a valid tokenizer config resolves to
the tokenizer template.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Guard legacy per-model-config migration idempotency
The v1->v2 localStorage migration (unsloth_load_settings ->
unsloth_model_configs) runs on every store read, so it must migrate exactly
once and never re-run, duplicate, or clobber a newer per-model config on a
reload or restart. That was covered only by a manual proof, so add durable
guards:
- Source-contract test pinning the three idempotency layers (the in-memory
legacyMigrationChecked guard, the persistent unsloth_model_configs_migrated
flag set in every terminal branch, and the non-overwriting Object.hasOwn
merge-skip) plus the readMap invocation. Reddens if any layer is dropped.
- Playwright model-config E2E: promote the legacy-migration step to a gating
check (soft_fail, which gates under the CI STUDIO_UI_STRICT=1) that the
migrated value is preserved and the flag is set, then reload again with a
fresh legacy seed present and assert the stored key set is unchanged, so a
second reload cannot re-migrate, duplicate, or clobber.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Note the migration E2E now gates idempotency under STUDIO_UI_STRICT
* Tighten model-picker per-model-config code comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: danielhanchen
Co-authored-by: shimmyshimmer
Co-authored-by: Unsloth
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
---
.github/workflows/studio-ui-smoke.yml | 50 +
studio/backend/hub/routes/inventory.py | 11 +
studio/backend/hub/schemas/inventory.py | 7 +
.../hub/services/models/cache_inventory.py | 84 +-
studio/backend/main.py | 3 +
studio/backend/models/inference.py | 24 +-
studio/backend/picker/__init__.py | 2 +
studio/backend/picker/routes/__init__.py | 6 +
studio/backend/picker/routes/templates.py | 45 +
studio/backend/picker/schemas.py | 32 +
studio/backend/picker/service.py | 426 +++
studio/backend/routes/inference.py | 44 +-
studio/backend/routes/models.py | 249 +-
.../tests/test_chat_load_during_training.py | 74 +
.../tests/test_export_absolute_paths.py | 1 +
.../tests/test_model_picker_regression.py | 232 ++
.../tests/test_model_update_robustness.py | 61 +
studio/backend/tests/test_picker_service.py | 266 ++
studio/backend/utils/models/gguf_metadata.py | 83 +-
studio/frontend/src/app/routes/__root.tsx | 7 -
.../frontend/src/components/app-sidebar.tsx | 22 -
.../remembered-load-settings.ts | 79 -
.../src/features/chat/api/chat-adapter.ts | 99 +-
.../src/features/chat/api/chat-api.ts | 25 +-
.../frontend/src/features/chat/chat-page.tsx | 513 +--
.../src/features/chat/chat-settings-sheet.tsx | 1348 +------
.../chat/hooks/use-chat-model-runtime.ts | 195 +-
.../hooks/use-staged-model-preparation.ts | 169 -
studio/frontend/src/features/chat/index.ts | 23 +-
.../lib/apply-inference-status-to-store.ts | 14 +-
.../src/features/chat/shared-composer.tsx | 198 +-
.../chat/stores/chat-runtime-store.ts | 323 +-
.../frontend/src/features/chat/types/api.ts | 3 +
.../export/components/export-run-panel.tsx | 71 +-
.../features/hub/catalog/catalog-states.tsx | 4 +
.../hub/catalog/dataset-download-section.tsx | 6 +-
.../features/hub/catalog/download-section.tsx | 6 +-
.../hub/catalog/gguf-download-card.tsx | 349 +-
.../features/hub/catalog/hub-option-menu.tsx | 17 +-
.../hub/catalog/local-dataset-card.tsx | 6 +-
.../hub/catalog/local-on-device-card.tsx | 119 +-
.../features/hub/catalog/model-inspector.tsx | 16 +-
.../hub/catalog/models-catalog-lists.tsx | 200 +-
.../hub/catalog/models-catalog-rows.tsx | 134 +-
.../features/hub/catalog/models-catalog.tsx | 4 +
.../src/features/hub/catalog/models-table.tsx | 59 +-
.../features/hub/catalog/models-toolbar.tsx | 188 +-
.../hub/catalog/on-device-folders-dialog.tsx | 33 +-
.../features/hub/catalog/path-info-button.tsx | 155 +-
.../hub/catalog/safetensors-download-card.tsx | 85 +-
.../hub/catalog/sampling-settings-dialog.tsx | 435 +++
.../src/features/hub/catalog/shared.tsx | 20 +-
.../download-manager-controller.ts | 23 -
.../features/hub/download-manager/index.ts | 1 -
studio/frontend/src/features/hub/hub-page.tsx | 242 +-
studio/frontend/src/features/hub/index.ts | 49 +-
.../src/features/hub/inventory/api.ts | 2 +
.../src/features/hub/inventory/types.ts | 3 +
.../hub/inventory/use-device-inventory.ts | 4 +
.../src/features/hub/inventory/view-models.ts | 9 +
.../src/features/hub/lib/hidden-models.ts | 69 +-
.../features/hub/lib/model-capabilities.ts | 28 +-
.../src/features/hub/lib/model-type-filter.ts | 52 +
.../src/features/hub/lib/view-models.ts | 7 +-
.../model-picker/api/model-metadata.ts | 20 +
.../features/model-picker/api/templates.ts | 87 +
.../chat-template-editor-dialog.tsx | 191 +
.../components/model-config-page.tsx | 991 +++++
.../components}/model-selector.tsx | 153 +-
.../model-selector/folder-browser.tsx | 90 +-
.../model-selector/model-capabilities.ts | 0
.../model-selector/model-delete-action.tsx | 9 +-
.../model-load-settings-action.tsx | 24 +-
.../model-selector/model-row-menu.tsx | 289 ++
.../model-selector/model-update-action.tsx | 25 +-
.../components}/model-selector/model-usage.ts | 3 +-
.../components}/model-selector/pickers.tsx | 3335 ++++++++---------
.../components}/model-selector/pill-tabs.tsx | 3 +-
.../model-selector/pinned-models.ts | 19 +-
.../model-selector/recommended-fit.ts | 10 +-
.../components}/model-selector/row-meta.ts | 0
.../components}/model-selector/source-tabs.ts | 0
.../components}/model-selector/types.ts | 14 +
.../components/numeric-value-input.tsx | 113 +
.../components/sidebar-model-config.tsx | 91 +
.../hooks/use-active-model-config.ts | 77 +
.../model-picker/hooks/use-model-defaults.ts | 191 +
.../src/features/model-picker/index.ts | 39 +
.../inventory/use-chat-picker-inventory.ts | 123 +
.../model-config/apply-per-model-config.ts | 127 +
.../model-config/model-identity.ts | 69 +
.../model-config/per-model-config.ts | 665 ++++
.../src/features/settings/tabs/chat-tab.tsx | 39 -
.../features/settings/tabs/general-tab.tsx | 4 +-
.../src/features/training/api/models-api.ts | 17 +-
.../frontend/src/features/training/index.ts | 4 +-
tests/studio/install/test_rocm_support.py | 4 +-
tests/studio/playwright_chat_ui.py | 14 +-
tests/studio/playwright_model_config.py | 740 ++++
.../test_cached_model_path_selection.py | 241 ++
tests/studio/test_gpu_inference_smoke.py | 65 +
tests/studio/test_model_picker_contracts.py | 407 ++
tests/studio/test_reveal_file_manager.py | 128 +
.../test_studio_text_descender_clipping.py | 9 +-
104 files changed, 10755 insertions(+), 4789 deletions(-)
create mode 100644 studio/backend/picker/__init__.py
create mode 100644 studio/backend/picker/routes/__init__.py
create mode 100644 studio/backend/picker/routes/templates.py
create mode 100644 studio/backend/picker/schemas.py
create mode 100644 studio/backend/picker/service.py
create mode 100644 studio/backend/tests/test_model_picker_regression.py
create mode 100644 studio/backend/tests/test_picker_service.py
delete mode 100644 studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
delete mode 100644 studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
create mode 100644 studio/frontend/src/features/hub/catalog/sampling-settings-dialog.tsx
create mode 100644 studio/frontend/src/features/hub/lib/model-type-filter.ts
create mode 100644 studio/frontend/src/features/model-picker/api/model-metadata.ts
create mode 100644 studio/frontend/src/features/model-picker/api/templates.ts
create mode 100644 studio/frontend/src/features/model-picker/components/chat-template-editor-dialog.tsx
create mode 100644 studio/frontend/src/features/model-picker/components/model-config-page.tsx
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector.tsx (83%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/folder-browser.tsx (84%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-capabilities.ts (100%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-delete-action.tsx (90%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-load-settings-action.tsx (63%)
create mode 100644 studio/frontend/src/features/model-picker/components/model-selector/model-row-menu.tsx
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-update-action.tsx (82%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-usage.ts (93%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/pickers.tsx (58%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/pill-tabs.tsx (98%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/pinned-models.ts (78%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/recommended-fit.ts (91%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/row-meta.ts (100%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/source-tabs.ts (100%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/types.ts (74%)
create mode 100644 studio/frontend/src/features/model-picker/components/numeric-value-input.tsx
create mode 100644 studio/frontend/src/features/model-picker/components/sidebar-model-config.tsx
create mode 100644 studio/frontend/src/features/model-picker/hooks/use-active-model-config.ts
create mode 100644 studio/frontend/src/features/model-picker/hooks/use-model-defaults.ts
create mode 100644 studio/frontend/src/features/model-picker/index.ts
create mode 100644 studio/frontend/src/features/model-picker/inventory/use-chat-picker-inventory.ts
create mode 100644 studio/frontend/src/features/model-picker/model-config/apply-per-model-config.ts
create mode 100644 studio/frontend/src/features/model-picker/model-config/model-identity.ts
create mode 100644 studio/frontend/src/features/model-picker/model-config/per-model-config.ts
create mode 100644 tests/studio/playwright_model_config.py
create mode 100644 tests/studio/test_cached_model_path_selection.py
create mode 100644 tests/studio/test_gpu_inference_smoke.py
create mode 100644 tests/studio/test_model_picker_contracts.py
create mode 100644 tests/studio/test_reveal_file_manager.py
diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml
index 30280c281e..0ad55ebd6d 100644
--- a/.github/workflows/studio-ui-smoke.yml
+++ b/.github/workflows/studio-ui-smoke.yml
@@ -237,6 +237,54 @@ jobs:
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
+ # Model-picker per-model-config regression (PR #7207 re-land of #6647).
+ # Fourth Unsloth on its own port; loads the tiny GGUF and drives the
+ # picker's run-settings surface: Context Length persists across a reload,
+ # Reset clears the stored override (never pins it), and the infra models
+ # (RAG embedder + llama.cpp probe) stay hidden from the picker.
+ - name: Reset auth + boot Unsloth for model-config tests (port 18898)
+ run: |
+ unsloth studio reset-password
+ mkdir -p logs
+ UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \
+ > logs/studio_modelcfg.log 2>&1 &
+ echo "STUDIO_MODELCFG_PID=$!" >> "$GITHUB_ENV"
+
+ - name: Wait for /api/health on 18898
+ run: |
+ for i in $(seq 1 180); do
+ if curl -fs "http://127.0.0.1:18898/api/health" > /tmp/health4.json; then
+ jq -e '.status == "healthy"' /tmp/health4.json && break
+ fi
+ sleep 1
+ done
+ jq -e '.status == "healthy"' /tmp/health4.json
+
+ - name: Pass bootstrap pw for model-config test
+ run: |
+ NEW="CIModelCfg-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
+ echo "::add-mask::$NEW"
+ echo "STUDIO_MODELCFG_NEW_PW=$NEW" >> "$GITHUB_ENV"
+
+ - name: Drive model-picker per-model-config with Playwright
+ env:
+ BASE_URL: http://127.0.0.1:18898
+ STUDIO_NEW_PW: ${{ env.STUDIO_MODELCFG_NEW_PW }}
+ PW_ART_DIR: logs/playwright_modelcfg
+ STUDIO_UI_STRICT: '1'
+ GGUF_REPO: ${{ env.GGUF_REPO }}
+ GGUF_VARIANT: ${{ env.GGUF_VARIANT }}
+ STUDIO_MODEL_HINT: gemma-3-270m
+ run: |
+ mkdir -p logs/playwright_modelcfg
+ python tests/studio/playwright_model_config.py
+
+ - name: Stop fourth Unsloth
+ if: always()
+ run: |
+ kill "${STUDIO_MODELCFG_PID}" 2>/dev/null || true
+ sleep 2
+
# IME + multilingual paste regression (issue #5318 / PR #5327).
# Third Unsloth on its own port so a hang here cannot poison the
# earlier UI tests. No GGUF -- the bug surface is the composer.
@@ -297,12 +345,14 @@ jobs:
path: |
logs/studio.log
logs/studio_extra.log
+ logs/studio_modelcfg.log
logs/studio_ime.log
logs/install.log
logs/server-logs/
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
+ logs/playwright_modelcfg
logs/playwright_ime
logs/studio-permissions-*.log
retention-days: 7
diff --git a/studio/backend/hub/routes/inventory.py b/studio/backend/hub/routes/inventory.py
index 4b6c179a2b..1ffadf0544 100644
--- a/studio/backend/hub/routes/inventory.py
+++ b/studio/backend/hub/routes/inventory.py
@@ -28,6 +28,7 @@ from hub.schemas.inventory import (
CachedModelsResponse,
DeleteCachedModelResponse,
GgufVariantsResponse,
+ HiddenModelsResponse,
LocalModelListResponse,
ModelsFolderResponse,
RecommendedFoldersResponse,
@@ -214,6 +215,16 @@ async def list_cached_models(
return await cache_inventory.list_cached_models_response(hf_token)
+@router.get("/hidden-models", response_model = HiddenModelsResponse)
+async def list_hidden_models(current_subject: str = Depends(get_current_subject)):
+ import asyncio
+
+ from routes.models import hidden_model_matchers
+
+ needles, exact_ids, exact_paths = await asyncio.to_thread(hidden_model_matchers)
+ return HiddenModelsResponse(needles = needles, exact_ids = exact_ids, exact_paths = exact_paths)
+
+
@router.delete(
"/delete-cached",
response_model = DeleteCachedModelResponse,
diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py
index ef95efe2f2..19d6da3e11 100644
--- a/studio/backend/hub/schemas/inventory.py
+++ b/studio/backend/hub/schemas/inventory.py
@@ -160,6 +160,7 @@ class CachedRepoBase(BaseModel):
repo_id: str
size_bytes: int = 0
cache_path: Optional[str] = None
+ last_modified: Optional[float] = None
partial: bool = False
partial_transport: Optional[str] = None
inventory_id: Optional[str] = None
@@ -189,6 +190,12 @@ class CachedModelsResponse(BaseModel):
cached: List[CachedModelRepo] = Field(default_factory = list)
+class HiddenModelsResponse(BaseModel):
+ needles: List[str] = Field(default_factory = list)
+ exact_ids: List[str] = Field(default_factory = list)
+ exact_paths: List[str] = Field(default_factory = list)
+
+
class AddScanFolderRequest(BaseModel):
"""Request body for adding a custom scan folder."""
diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py
index 54a25482f2..c1b864bb63 100644
--- a/studio/backend/hub/services/models/cache_inventory.py
+++ b/studio/backend/hub/services/models/cache_inventory.py
@@ -31,6 +31,7 @@ from hub.services.models.common import (
_is_checkpoint_weight_name,
_is_gguf_filename,
_is_main_gguf_filename,
+ _is_mmproj_filename,
_is_transformers_safetensors_weight_name,
_local_inventory_id,
_prefer_complete_larger,
@@ -132,6 +133,39 @@ def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
+def _blob_mtime(file_obj) -> float:
+ ts = getattr(file_obj, "blob_last_modified", None)
+ if isinstance(ts, (int, float)) and ts > 0:
+ return float(ts)
+ blob_path = getattr(file_obj, "blob_path", None)
+ if blob_path:
+ try:
+ return float(Path(blob_path).stat().st_mtime)
+ except OSError:
+ pass
+ return 0.0
+
+
+def _repo_gguf_last_modified(repo_info) -> float:
+ latest = 0.0
+ for revision in repo_info.revisions:
+ for f in revision.files:
+ if _is_main_gguf_filename(f.file_name):
+ latest = max(latest, _blob_mtime(f))
+ return latest
+
+
+def _repo_has_mmproj(repo_info) -> bool:
+ # An mmproj file only makes a repo vision-capable when it is an actual GGUF
+ # projector; a non-GGUF sidecar (e.g. mmproj_config.json) does not, and the
+ # runtime's projector detection is GGUF-only.
+ return any(
+ _is_gguf_filename(f.file_name) and _is_mmproj_filename(f.file_name)
+ for revision in repo_info.revisions
+ for f in revision.files
+ )
+
+
def _cached_repo_file_name(file_obj) -> str:
file_path = getattr(file_obj, "file_path", None)
if file_path:
@@ -291,6 +325,7 @@ def _scan_cached_gguf() -> list[dict]:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
+ last_modified = _repo_gguf_last_modified(repo_info)
row = {
"repo_id": repo_id,
"size_bytes": max(total_size, variant_state_size),
@@ -300,6 +335,9 @@ def _scan_cached_gguf() -> list[dict]:
# per-variant detail lives on GgufVariantDetail.
"partial_transport": None,
}
+ last_modified = max(last_modified, (existing or {}).get("last_modified", 0.0))
+ if last_modified > 0:
+ row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
@@ -308,11 +346,20 @@ def _scan_cached_gguf() -> list[dict]:
requires_variant = True,
)
)
+ if _repo_has_mmproj(repo_info):
+ row["capabilities"]["supports_vision"] = True
# Visible infra variants remain management-only.
if is_hidden_infra:
row["capabilities"]["can_chat"] = False
if _prefer_cache_row(row, existing):
+ if existing and existing["capabilities"].get("supports_vision"):
+ row["capabilities"]["supports_vision"] = True
seen_lower[key] = row
+ else:
+ if last_modified > existing.get("last_modified", 0.0):
+ existing["last_modified"] = last_modified
+ if row["capabilities"].get("supports_vision"):
+ existing["capabilities"]["supports_vision"] = True
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "")
logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}")
@@ -340,13 +387,14 @@ class _CachedNonGgufPayload(NamedTuple):
size_bytes: int
has_runnable_weights: bool
model_format: ModelFormat
+ last_modified: float
def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
- all_weight_blobs: dict[str, int] = {}
- adapter_blobs: dict[str, int] = {}
- safetensors_blobs: dict[str, int] = {}
- checkpoint_blobs: dict[str, int] = {}
+ all_weight_blobs: dict[str, tuple[int, float]] = {}
+ adapter_blobs: dict[str, tuple[int, float]] = {}
+ safetensors_blobs: dict[str, tuple[int, float]] = {}
+ checkpoint_blobs: dict[str, tuple[int, float]] = {}
has_config = False
has_adapter_config = False
has_adapter_weights = False
@@ -354,12 +402,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
has_transformers_safetensors = False
has_checkpoint = False
- def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None:
+ def _record_blob(
+ target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str
+ ) -> None:
blob_path = getattr(file_obj, "blob_path", None)
size = int(file_obj.size_on_disk or 0)
key = str(blob_path) if blob_path else f"{rev_id}:{file_name}"
- target[key] = size
- all_weight_blobs[key] = size
+ value = (size, _blob_mtime(file_obj))
+ target[key] = value
+ all_weight_blobs[key] = value
for revision in repo_info.revisions:
rev_id = getattr(revision, "commit_hash", None) or str(id(revision))
@@ -403,18 +454,19 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
or "unknown"
)
if model_format == "adapter":
- size_bytes = sum(adapter_blobs.values())
+ selected_blobs = adapter_blobs
elif model_format == "safetensors":
- size_bytes = sum(safetensors_blobs.values())
+ selected_blobs = safetensors_blobs
elif model_format == "checkpoint":
- size_bytes = sum(checkpoint_blobs.values())
+ selected_blobs = checkpoint_blobs
else:
- size_bytes = sum(all_weight_blobs.values())
+ selected_blobs = all_weight_blobs
return _CachedNonGgufPayload(
- size_bytes = size_bytes,
+ size_bytes = sum(size for size, _mtime in selected_blobs.values()),
has_runnable_weights = model_format != "unknown",
model_format = model_format,
+ last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0),
)
@@ -544,6 +596,12 @@ def _scan_cached_models() -> list[dict]:
),
**_cached_model_local_metadata(repo_path),
}
+ last_modified = max(
+ payload.last_modified,
+ (existing or {}).get("last_modified", 0.0),
+ )
+ if last_modified > 0:
+ row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
@@ -553,6 +611,8 @@ def _scan_cached_models() -> list[dict]:
)
if _prefer_cache_row(row, existing):
seen_lower[key] = row
+ elif last_modified > existing.get("last_modified", 0.0):
+ existing["last_modified"] = last_modified
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")
diff --git a/studio/backend/main.py b/studio/backend/main.py
index 48675b9539..3f244dc22e 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -315,6 +315,7 @@ from hub.routes import (
datasets_router as hub_datasets_router,
token_router as hub_token_router,
)
+from picker.routes import templates_router as picker_templates_router
from hub.schemas.downloads import TransportCapabilities
from hub.utils.download_registry import (
get_download_transport_capabilities,
@@ -764,6 +765,7 @@ _BODY_PROTECTED_PREFIXES = (
"/v1/completions",
"/p/",
"/api/inference",
+ "/api/picker",
"/api/data-recipe",
"/api/datasets",
"/api/hub",
@@ -995,6 +997,7 @@ app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])
app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"])
app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"])
+app.include_router(picker_templates_router, prefix = "/api/picker", tags = ["picker"])
app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"])
# Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index d51d35189b..580a74dddf 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -18,6 +18,8 @@ from pydantic import (
model_validator,
)
+from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
+
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
@@ -54,8 +56,16 @@ class LoadRequest(BaseModel):
@field_validator("chat_template_override")
@classmethod
def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]:
- if value is not None and value.strip() == "":
+ if value is None:
return None
+ # Char count is a lower bound on UTF-8 byte length: reject an oversized
+ # template before spending work encoding it.
+ if len(value) > MAX_CHAT_TEMPLATE_BYTES:
+ raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
+ if value.strip() == "":
+ return None
+ if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
+ raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
return value
cache_type_kv: Optional[str] = Field(
@@ -206,6 +216,13 @@ class ValidateModelRequest(BaseModel):
description = "Also read the native context length from the local GGUF header. "
"Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.",
)
+ include_chat_template: bool = Field(
+ False,
+ description = "Also read the embedded chat template from the local GGUF header, so a "
+ "native (picked / drag-drop) file's default template can be shown before it is loaded. "
+ "Opt-in and, like include_context_length, a metadata-only probe that skips the training "
+ "guard. Only the leased file's own embedded template is read, never sibling sidecars.",
+ )
class TransformersUpgradeInfo(BaseModel):
@@ -266,6 +283,11 @@ class ValidateModelResponse(BaseModel):
description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF "
"header alongside context_length; 0 for dense models, None when not read.",
)
+ chat_template: Optional[str] = Field(
+ None,
+ description = "Embedded GGUF chat template, read from the header when include_chat_template "
+ "is set (native lease-backed picks); None for non-GGUF, over-cap, or not-read templates.",
+ )
# Additive fields; the consuming consent dialog ships in a follow-up frontend PR.
requires_transformers_upgrade: bool = Field(
False,
diff --git a/studio/backend/picker/__init__.py b/studio/backend/picker/__init__.py
new file mode 100644
index 0000000000..32014236c6
--- /dev/null
+++ b/studio/backend/picker/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
diff --git a/studio/backend/picker/routes/__init__.py b/studio/backend/picker/routes/__init__.py
new file mode 100644
index 0000000000..c0e988c8bb
--- /dev/null
+++ b/studio/backend/picker/routes/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from .templates import router as templates_router
+
+__all__ = ["templates_router"]
diff --git a/studio/backend/picker/routes/templates.py b/studio/backend/picker/routes/templates.py
new file mode 100644
index 0000000000..02b8bf7184
--- /dev/null
+++ b/studio/backend/picker/routes/templates.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
+
+from __future__ import annotations
+
+import asyncio
+from typing import Optional
+
+from fastapi import APIRouter, Body, Depends, Query
+
+from auth.authentication import get_current_subject
+from hub.dependencies import get_hf_token
+
+from ..schemas import (
+ MAX_CHAT_TEMPLATE_BYTES,
+ ModelTemplateResponse,
+ ValidateChatTemplateRequest,
+ ValidateChatTemplateResponse,
+)
+from ..service import read_default_chat_template, validate_chat_template
+
+router = APIRouter()
+
+
+@router.post("/validate-chat-template", response_model = ValidateChatTemplateResponse)
+async def validate_chat_template_route(
+ body: ValidateChatTemplateRequest = Body(...),
+ current_subject: str = Depends(get_current_subject),
+) -> ValidateChatTemplateResponse:
+ return await asyncio.to_thread(validate_chat_template, body.template)
+
+
+@router.get("/chat-template/{model_name:path}", response_model = ModelTemplateResponse)
+async def get_default_chat_template_route(
+ model_name: str,
+ gguf_variant: Optional[str] = Query(None),
+ hf_token: Optional[str] = Depends(get_hf_token),
+ current_subject: str = Depends(get_current_subject),
+) -> ModelTemplateResponse:
+ template = await asyncio.to_thread(
+ read_default_chat_template, model_name, hf_token, gguf_variant
+ )
+ if template is not None and len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
+ template = None
+ return ModelTemplateResponse(model_name = model_name, chat_template = template)
diff --git a/studio/backend/picker/schemas.py b/studio/backend/picker/schemas.py
new file mode 100644
index 0000000000..b4f956188f
--- /dev/null
+++ b/studio/backend/picker/schemas.py
@@ -0,0 +1,32 @@
+# 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 pydantic import BaseModel, Field, field_validator
+
+# Mirror the frontend's 64 KiB chat-template contract (per-model-config.ts) at
+# the API boundary so a direct caller cannot make Jinja parse an oversized
+# template. MaxBodyMiddleware only caps the whole request body, not this field.
+MAX_CHAT_TEMPLATE_BYTES = 65_536
+
+
+class ValidateChatTemplateRequest(BaseModel):
+ template: str = Field(default = "")
+
+ @field_validator("template")
+ @classmethod
+ def _enforce_template_size(cls, value: str) -> str:
+ if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
+ raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
+ return value
+
+
+class ValidateChatTemplateResponse(BaseModel):
+ valid: bool
+ error: Optional[str] = None
+
+
+class ModelTemplateResponse(BaseModel):
+ model_name: str
+ chat_template: Optional[str] = None
diff --git a/studio/backend/picker/service.py b/studio/backend/picker/service.py
new file mode 100644
index 0000000000..13065b2920
--- /dev/null
+++ b/studio/backend/picker/service.py
@@ -0,0 +1,426 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+from pathlib import Path
+from typing import Optional
+
+from hub.services.models.folder_browser import (
+ _build_browse_allowlist,
+ _is_path_inside_allowlist,
+)
+from hub.utils.gguf import extract_quant_label, iter_hf_cache_snapshots
+from utils.models.gguf_metadata import read_gguf_chat_template
+from utils.models.model_config import (
+ _extract_quant_label,
+ _is_big_endian_gguf_path,
+ _is_mmproj,
+ _is_mtp_drafter,
+)
+from utils.paths.path_utils import (
+ is_local_path,
+ normalize_path,
+ resolve_cached_repo_id_case,
+)
+
+from .schemas import MAX_CHAT_TEMPLATE_BYTES, ValidateChatTemplateResponse
+
+logger = logging.getLogger(__name__)
+
+_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")
+
+
+def _is_valid_repo_id(repo_id: str) -> bool:
+ return bool(_VALID_REPO_ID.fullmatch(repo_id))
+
+
+_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json")
+_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja")
+_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json")
+
+# Cap sidecar reads so a malformed or hostile metadata file cannot exhaust memory
+# before its template is size-checked. The JSON envelope may exceed a bare template
+# (it carries other tokenizer metadata); the extracted template is still bounded by
+# MAX_CHAT_TEMPLATE_BYTES downstream.
+MAX_TEMPLATE_METADATA_BYTES = 4 * 1024 * 1024
+
+
+def _read_bounded_text(path: Path, limit: int) -> Optional[str]:
+ """Read at most `limit` bytes of UTF-8 text; None if larger or unreadable."""
+ try:
+ with path.open("rb") as f:
+ data = f.read(limit + 1)
+ except OSError:
+ return None
+ if len(data) > limit:
+ return None
+ try:
+ return data.decode("utf-8")
+ except UnicodeError:
+ return None
+
+
+def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool:
+ # Block symlinked children from escaping the validated directory (realpath-checked).
+ # None = trusted caller (HF cache / remote download).
+ return allow_roots is None or _is_path_inside_allowlist(path, allow_roots)
+
+
+def validate_chat_template(template: str) -> ValidateChatTemplateResponse:
+ text = (template or "").strip()
+ if not text:
+ return ValidateChatTemplateResponse(valid = True, error = None)
+ # Import Jinja lazily: optional at runtime (e.g. GGUF-only installs), so a
+ # missing dependency must not crash API startup.
+ try:
+ from jinja2 import TemplateError
+ from jinja2.ext import Extension
+ from jinja2.sandbox import ImmutableSandboxedEnvironment
+ except ImportError:
+ return ValidateChatTemplateResponse(valid = True, error = None)
+
+ class _GenerationTag(Extension):
+ # Accept Transformers' {% generation %} assistant-mask tag so a pasted HF
+ # chat template validates (we only parse it).
+ tags = {"generation"}
+
+ def parse(self, parser):
+ next(parser.stream)
+ return parser.parse_statements(["name:endgeneration"], drop_needle = True)
+
+ try:
+ env = ImmutableSandboxedEnvironment(
+ trim_blocks = True,
+ lstrip_blocks = True,
+ extensions = ["jinja2.ext.loopcontrols", _GenerationTag],
+ )
+ env.parse(text)
+ return ValidateChatTemplateResponse(valid = True, error = None)
+ except TemplateError as exc:
+ message = getattr(exc, "message", None) or str(exc)
+ lineno = getattr(exc, "lineno", None)
+ if lineno:
+ message = f"Line {lineno}: {message}"
+ return ValidateChatTemplateResponse(valid = False, error = message)
+ except Exception as exc:
+ return ValidateChatTemplateResponse(valid = False, error = str(exc))
+
+
+def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]:
+ if not isinstance(config, dict):
+ return None
+ raw = config.get("chat_template")
+ if isinstance(raw, str) and raw.strip():
+ return raw
+ if isinstance(raw, list):
+ fallback: Optional[str] = None
+ for entry in raw:
+ if not isinstance(entry, dict):
+ continue
+ template = entry.get("template")
+ if not isinstance(template, str):
+ continue
+ if entry.get("name") == "default":
+ return template
+ if fallback is None:
+ fallback = template
+ return fallback
+ return None
+
+
+def _chat_template_from_jinja_file(
+ dir_path: Path, allow_roots: Optional[list[Path]] = None
+) -> Optional[str]:
+ for rel in _JINJA_TEMPLATE_PATHS:
+ template_file = dir_path / rel
+ if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots):
+ continue
+ try:
+ if template_file.stat().st_size > MAX_CHAT_TEMPLATE_BYTES:
+ continue
+ template = template_file.read_text(encoding = "utf-8")
+ except Exception:
+ continue
+ if template.strip():
+ return template
+ return None
+
+
+def _chat_template_from_processor_payload(payload: object) -> Optional[str]:
+ # processor chat_template.json may be the template string itself or a
+ # {name: template} map, not only a tokenizer_config-shaped object.
+ if isinstance(payload, str):
+ return payload if payload.strip() else None
+ template = _chat_template_from_tokenizer_config(payload) # type: ignore[arg-type]
+ if template:
+ return template
+ if isinstance(payload, dict):
+ # Named-template map: prefer "default", else the first non-empty entry
+ # (mirrors the tokenizer-config list fallback).
+ default = payload.get("default")
+ if isinstance(default, str) and default.strip():
+ return default
+ for value in payload.values():
+ if isinstance(value, str) and value.strip():
+ return value
+ return None
+
+
+def _chat_template_from_processor_json(
+ dir_path: Path, allow_roots: Optional[list[Path]] = None
+) -> Optional[str]:
+ for rel in _PROCESSOR_TEMPLATE_PATHS:
+ config_file = dir_path / rel
+ if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
+ continue
+ raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
+ if raw is None:
+ continue
+ try:
+ payload = json.loads(raw)
+ except Exception:
+ continue
+ template = _chat_template_from_processor_payload(payload)
+ if template:
+ return template
+ return None
+
+
+def _chat_template_from_tokenizer_dir(
+ dir_path: Path, allow_roots: Optional[list[Path]] = None
+) -> Optional[str]:
+ jinja = _chat_template_from_jinja_file(dir_path, allow_roots)
+ if jinja:
+ return jinja
+ for rel in _TOKENIZER_CONFIG_PATHS:
+ config_file = dir_path / rel
+ if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
+ continue
+ raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
+ if raw is None:
+ continue
+ try:
+ config = json.loads(raw)
+ except Exception:
+ continue
+ template = _chat_template_from_tokenizer_config(config)
+ if template:
+ return template
+ return _chat_template_from_processor_json(dir_path, allow_roots)
+
+
+_GGUF_SCAN_MAX_DEPTH = 2
+
+
+def _iter_ggufs(dir_path: Path) -> list[Path]:
+ if dir_path == dir_path.parent:
+ return []
+ root = str(dir_path)
+ found: list[Path] = []
+ for current, dirs, files in os.walk(root, followlinks = False):
+ rel = os.path.relpath(current, root)
+ depth = 0 if rel == os.curdir else rel.count(os.sep) + 1
+ if depth >= _GGUF_SCAN_MAX_DEPTH:
+ dirs[:] = []
+ for name in files:
+ if not name.lower().endswith(".gguf") or _is_mmproj(name):
+ continue
+ path = Path(current) / name
+ try:
+ rel = path.relative_to(dir_path).as_posix()
+ except ValueError:
+ rel = name
+ quant = _extract_quant_label(rel)
+ if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
+ continue
+ found.append(path)
+ return found
+
+
+def _variant_matches(relative_path: str, needle: str) -> bool:
+ quant = _extract_quant_label(relative_path).lower()
+ if quant == needle:
+ return True
+ if extract_quant_label(relative_path).lower() == needle:
+ return True
+ prefix = f"{needle}-"
+ if not quant.startswith(prefix):
+ return False
+ suffix = quant[len(prefix) :]
+ if not suffix.endswith("bpw"):
+ return False
+ value = suffix[:-3]
+ return bool(value) and value.replace(".", "", 1).isdigit()
+
+
+_GGUF_SPLIT_INDEX_RE = re.compile(r"-(\d{3,})-of-\d{3,}$", re.IGNORECASE)
+
+
+def _is_nonfirst_gguf_split(path: Path) -> bool:
+ match = _GGUF_SPLIT_INDEX_RE.search(path.stem)
+ return match is not None and int(match.group(1)) != 1
+
+
+def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]:
+ try:
+ ggufs = sorted(_iter_ggufs(dir_path))
+ except OSError:
+ return None
+ if not ggufs:
+ return None
+ needle = (gguf_variant or "").strip().lower()
+ if needle:
+ for path in ggufs:
+ try:
+ relative = path.relative_to(dir_path).as_posix()
+ except ValueError:
+ relative = path.name
+ if _variant_matches(relative, needle):
+ return path
+ return None
+ candidates = [path for path in ggufs if not _is_nonfirst_gguf_split(path)] or ggufs
+ try:
+ return max(candidates, key = lambda path: path.stat().st_size)
+ except OSError:
+ return candidates[0]
+
+
+def _chat_template_from_dir(
+ dir_path: Path,
+ gguf_variant: Optional[str] = None,
+ allow_roots: Optional[list[Path]] = None,
+) -> Optional[str]:
+ def from_gguf() -> Optional[str]:
+ gguf = _find_gguf_in_dir(dir_path, gguf_variant)
+ if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots):
+ return None
+ return read_gguf_chat_template(str(gguf))
+
+ # Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are the
+ # author's maintained template and supersede the GGUF's possibly-stale embedded
+ # copy. The variant only picks the GGUF fallback, so tokenizer-first precedence
+ # holds whether or not a variant is given.
+ return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf()
+
+
+def read_default_chat_template(
+ model_name: str,
+ hf_token: Optional[str] = None,
+ gguf_variant: Optional[str] = None,
+) -> Optional[str]:
+ if not isinstance(model_name, str) or not model_name.strip():
+ return None
+ name = model_name.strip()
+
+ if is_local_path(name):
+ try:
+ target = Path(normalize_path(name)).expanduser()
+ allow_roots = _build_browse_allowlist()
+ if not _is_path_inside_allowlist(target, allow_roots):
+ logger.debug("Refused chat template read outside allowed folders: %s", name)
+ return None
+ if name.lower().endswith(".gguf"):
+ # Prefer a maintained sidecar next to the file over the GGUF's
+ # embedded copy (tokenizer-first precedence, as elsewhere).
+ sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots)
+ if sidecar:
+ return sidecar
+ return read_gguf_chat_template(str(target))
+ return _chat_template_from_dir(target, gguf_variant, allow_roots)
+ except Exception as exc:
+ logger.debug("Could not read local chat template for %s: %s", name, exc)
+ return None
+
+ if not _is_valid_repo_id(name):
+ return None
+
+ resolved = resolve_cached_repo_id_case(name)
+
+ try:
+ # Resolve within each cached revision, newest first. A revision's sidecar
+ # supersedes its own embedded GGUF copy, but must not override a newer
+ # revision, so precedence stays per-snapshot rather than global.
+ for snapshot in iter_hf_cache_snapshots(resolved):
+ template = _chat_template_from_dir(snapshot, gguf_variant)
+ if template:
+ return template
+ except Exception as exc:
+ logger.debug("Could not read cached chat template for %s: %s", resolved, exc)
+
+ try:
+ from huggingface_hub import HfApi, hf_hub_download
+
+ _api = HfApi()
+
+ def _remote_exceeds_cap(rel: str) -> bool:
+ # Best-effort: skip the download when the remote's advertised size
+ # exceeds the cap, so a maliciously large sidecar is never fetched.
+ try:
+ infos = _api.get_paths_info(resolved, [rel], repo_type = "model", token = hf_token)
+ except Exception:
+ return False
+ for info in infos:
+ size = getattr(info, "size", None)
+ if (
+ getattr(info, "path", None) == rel
+ and isinstance(size, int)
+ and size > MAX_TEMPLATE_METADATA_BYTES
+ ):
+ return True
+ return False
+
+ def _download_text(rel: str) -> Optional[str]:
+ if _remote_exceeds_cap(rel):
+ return None
+ try:
+ path = hf_hub_download(resolved, rel, token = hf_token)
+ return _read_bounded_text(Path(path), MAX_TEMPLATE_METADATA_BYTES)
+ except Exception:
+ return None
+
+ for rel in _JINJA_TEMPLATE_PATHS:
+ template = _download_text(rel)
+ if not template or not template.strip():
+ continue
+ # A raw Jinja sidecar is the whole template, so it must fit the route's
+ # response cap (the local path skips oversized .jinja too). Download stays
+ # bounded at MAX_TEMPLATE_METADATA_BYTES so a large JSON embedding a small
+ # template still extracts below, but an over-cap Jinja is dropped so the
+ # search falls through to the tokenizer/processor template.
+ if len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
+ continue
+ return template
+
+ for rel in _TOKENIZER_CONFIG_PATHS:
+ raw = _download_text(rel)
+ if not raw:
+ continue
+ try:
+ config = json.loads(raw)
+ except Exception:
+ continue
+ template = _chat_template_from_tokenizer_config(config)
+ if template:
+ return template
+
+ for rel in _PROCESSOR_TEMPLATE_PATHS:
+ raw = _download_text(rel)
+ if not raw:
+ continue
+ try:
+ payload = json.loads(raw)
+ except Exception:
+ continue
+ template = _chat_template_from_processor_payload(payload)
+ if template:
+ return template
+
+ return None
+ except Exception as exc:
+ logger.debug("Could not fetch chat template for %s: %s", resolved, exc)
+ return None
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index afd942e9a5..d3e588bb0b 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -5144,10 +5144,10 @@ async def validate_model(
latest_tier_active_for, config.identifier, request.hf_token
):
effective_load_in_4bit = False
- # A metadata-only probe just reads the GGUF header and allocates no VRAM,
- # so it must not be refused by the training guard. Real loads validate
- # without include_context_length and /load applies the guard again.
- if not request.include_context_length:
+ # A metadata-only probe reads the GGUF header and allocates no VRAM, so the
+ # training guard must not refuse it. Real loads omit include_context_length /
+ # include_chat_template, and /load applies the guard again.
+ if not (request.include_context_length or request.include_chat_template):
# Match /load's inherited llama.cpp extras and parallel slot count so
# validation cannot pass a smaller estimate than the subsequent load.
effective_extra_args = _resolve_inherited_extra_args(
@@ -5189,9 +5189,15 @@ async def validate_model(
context_length: Optional[int] = None
layer_count: Optional[int] = None
moe_layer_count: Optional[int] = None
- if request.include_context_length and is_gguf:
+ chat_template: Optional[str] = None
+ # Both header probes read the same local GGUF, so resolve it once.
+ if (request.include_context_length or request.include_chat_template) and is_gguf:
from hub.utils.gguf import resolve_local_gguf_path
- from utils.models.gguf_metadata import read_gguf_staged_dims
+ from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
+ from utils.models.gguf_metadata import (
+ read_gguf_chat_template,
+ read_gguf_staged_dims,
+ )
# Best-effort: a header-read failure must never fail validation of an
# otherwise-valid model (the outer except turns it into a 400).
@@ -5207,13 +5213,24 @@ async def validate_model(
model_identifier, request.gguf_variant
)
if local_gguf:
- # Header walk reads tokenizer arrays for dense models (tens of
- # ms); keep it off the event loop.
- dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf)
- if dims:
- context_length = dims["context_length"]
- layer_count = dims["layer_count"]
- moe_layer_count = dims["moe_layer_count"]
+ if request.include_context_length:
+ # Header walk reads tokenizer arrays (tens of ms); keep it
+ # off the event loop.
+ dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf)
+ if dims:
+ context_length = dims["context_length"]
+ layer_count = dims["layer_count"]
+ moe_layer_count = dims["moe_layer_count"]
+ if request.include_chat_template:
+ # Read only the leased GGUF's own embedded template (the copy
+ # llama.cpp loads), never a sibling sidecar: the native grant
+ # authorizes just this path, so neighbours would be scope escalation.
+ raw_template = await asyncio.to_thread(read_gguf_chat_template, local_gguf)
+ if (
+ raw_template is not None
+ and len(raw_template.encode("utf-8")) <= MAX_CHAT_TEMPLATE_BYTES
+ ):
+ chat_template = raw_template
except Exception as e:
logger.debug("Header probe failed for %s: %s", model_log_label, e)
@@ -5232,6 +5249,7 @@ async def validate_model(
context_length = context_length,
layer_count = layer_count,
moe_layer_count = moe_layer_count,
+ chat_template = chat_template,
requires_transformers_upgrade = transformers_upgrade is not None,
transformers_upgrade = transformers_upgrade,
)
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 0806c2f513..a5ce1a72f0 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -60,13 +60,52 @@ def _safe_is_dir(path) -> bool:
# Shared with the hub inventory scans; keep the private aliases so existing
-# importers (core.inference.local_model_resolver, tests) stay valid.
+# importers stay valid. ``_HF_REPO_ID_RE`` is the Hub repo id shape ("owner/name");
+# anything else is treated as a local filesystem path.
from utils.hidden_models import (
+ _HF_REPO_ID_RE,
+ _existing_resolved_path,
_safe_resolve,
is_hidden_model as _is_hidden_model,
)
+def hidden_model_matchers() -> tuple[list[str], list[str], list[str]]:
+ """Substring needles, exact repo ids, and exact resolved paths identifying
+ infra models (the RAG embedder and the llama.cpp install validation probe)
+ that pickers hide. Served by the ``/api/hub/hidden-models`` endpoint. A
+ configured HF-repo embedder is published as its exact lowercased repo id
+ (mirroring ``utils.hidden_models.is_hidden_model``) and a local-path
+ embedder as its exact resolved path only: a generic basename like "model"
+ must not substring-hide unrelated chat models."""
+ from core.rag import config as rag_config
+
+ needles = [
+ # The validation probe's repo and its exact filename. The filename carries
+ # .gguf so it won't hide unrelated repos like ``user/stories260K-finetune-GGUF``.
+ "ggml-org/models",
+ "stories260k.gguf",
+ ]
+ exact_ids: list[str] = []
+ exact_paths: list[str] = []
+ for model in (
+ rag_config.effective_embedding_model(),
+ rag_config.effective_gguf_repo(),
+ ):
+ # Resolve an existing local path before the repo-id regex: a local embedder
+ # shaped like "models/embedder" is an exact path, not a Hub repo id.
+ existing_path = _existing_resolved_path(model)
+ if existing_path:
+ exact_paths.append(existing_path.lower())
+ elif _HF_REPO_ID_RE.match(model):
+ exact_ids.append(model.lower())
+ else:
+ resolved = _safe_resolve(Path(model).expanduser())
+ if resolved:
+ exact_paths.append(resolved.lower())
+ return needles, exact_ids, exact_paths
+
+
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
@@ -91,6 +130,7 @@ try:
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
+ _is_mtp_drafter,
is_audio_input_type,
)
from core.inference import get_inference_backend
@@ -123,6 +163,7 @@ except ImportError:
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
+ _is_mtp_drafter,
is_audio_input_type,
)
from core.inference import get_inference_backend
@@ -803,7 +844,7 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
models = sorted(
deduped.values(),
- key = lambda item: (item.updated_at or 0),
+ key = lambda item: item.updated_at or 0,
reverse = True,
)
return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)]
@@ -1750,9 +1791,11 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op
async def get_model_config(
model_name: str,
hf_token: Optional[str] = Query(None),
+ header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""Get configuration for a specific model (wraps load_model_defaults)."""
+ hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
if not is_local_path(model_name):
resolved = resolve_cached_repo_id_case(model_name)
@@ -2471,6 +2514,7 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get
async def check_vision_model(
model_name: str,
hf_token: Optional[str] = Query(None),
+ header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""
@@ -2478,6 +2522,7 @@ async def check_vision_model(
This endpoint wraps the backend is_vision_model function.
"""
+ hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
logger.info(f"Checking if vision model: {model_name}")
# Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision).
@@ -2503,6 +2548,7 @@ async def check_vision_model(
async def check_embedding_model(
model_name: str,
hf_token: Optional[str] = Query(None),
+ header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""
@@ -2510,6 +2556,7 @@ async def check_embedding_model(
This endpoint wraps the backend is_embedding_model function.
"""
+ hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
logger.info(f"Checking if embedding model: {model_name}")
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
@@ -2573,12 +2620,6 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio
Q8_0 weights). Never raises.
"""
try:
- from utils.models.model_config import (
- _extract_quant_label,
- _is_big_endian_gguf_path,
- _is_mtp_drafter,
- )
-
if is_local:
roots = [Path(repo_id)]
else:
@@ -2595,25 +2636,19 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio
if snaps.is_dir():
roots.extend(s for s in snaps.iterdir() if s.is_dir())
- want = quant.lower().replace("-", "").replace("_", "")
+ want = _normalized_quant_label(quant)
best_total = 0
best_first: Optional[str] = None
for root in roots:
matches: list[tuple[str, Path]] = []
total = 0
for f in _iter_gguf_paths(root):
- if _is_mmproj_filename(f.name):
- continue
try:
rel = f.relative_to(root).as_posix()
except ValueError:
rel = f.name
- if _is_mtp_drafter(rel):
- continue
- q = _extract_quant_label(rel)
- if _is_big_endian_gguf_path(rel, q):
- continue
- if q.lower().replace("-", "").replace("_", "") != want:
+ q = _main_variant_gguf_label(rel)
+ if q is None or _normalized_quant_label(q) != want:
continue
try:
total += f.stat().st_size
@@ -3035,6 +3070,22 @@ def _is_main_gguf_filename(name: str) -> bool:
return _is_gguf_filename(name) and not _is_mmproj_filename(name)
+def _main_variant_gguf_label(rel_path: str) -> Optional[str]:
+ name = rel_path.rsplit("/", 1)[-1]
+ if not _is_main_gguf_filename(name):
+ return None
+ if _is_mtp_drafter(rel_path):
+ return None
+ label = _extract_quant_label(rel_path)
+ if _is_big_endian_gguf_path(rel_path, label):
+ return None
+ return label
+
+
+def _normalized_quant_label(label: str) -> str:
+ return label.lower().replace("-", "").replace("_", "")
+
+
def _repo_has_mmproj(repo_info) -> bool:
"""True if the repo ships a GGUF vision adapter (mmproj), so it can
take image inputs. Cheap: scans already-listed file names only."""
@@ -3362,6 +3413,170 @@ async def delete_cached_model(
)
+def _resolve_cached_model_path(repo_id: str, variant: Optional[str]) -> Path:
+ """Absolute path of a cached repo (newest snapshot dir) or, with *variant*,
+ that quant's main GGUF file (first split of a sharded quant). Paths come
+ from the HF cache scan only, so callers can't probe arbitrary paths."""
+ cache_scans = _all_hf_cache_scans()
+
+ matching_repos = []
+ for hf_cache in cache_scans:
+ for repo_info in hf_cache.repos:
+ if repo_info.repo_type != "model":
+ continue
+ if repo_info.repo_id.lower() == repo_id.lower():
+ matching_repos.append(repo_info)
+ if not matching_repos:
+ raise HTTPException(status_code = 404, detail = "Model not found in cache")
+
+ if variant:
+ want = _normalized_quant_label(variant)
+ candidate_revisions = sorted(
+ (rev for repo_info in matching_repos for rev in repo_info.revisions),
+ key = lambda rev: getattr(rev, "last_modified", 0) or 0,
+ reverse = True,
+ )
+ for rev in candidate_revisions:
+ snapshot = getattr(rev, "snapshot_path", None)
+ matches = []
+ for f in rev.files:
+ p = Path(f.file_path)
+ rel = f.file_name
+ if snapshot:
+ try:
+ rel = p.relative_to(snapshot).as_posix()
+ except ValueError:
+ pass
+ label = _main_variant_gguf_label(rel)
+ if label is None or _normalized_quant_label(label) != want:
+ continue
+ if p.exists() or p.is_symlink():
+ matches.append((rel, p))
+ if matches:
+ # Path-sorted so a sharded quant deterministically yields its first split.
+ return sorted(matches, key = lambda m: m[0].lower())[0][1]
+ raise HTTPException(
+ status_code = 404,
+ detail = f"Variant {variant} not found in cache for {repo_id}",
+ )
+
+ def repo_size(repo_info) -> int:
+ gguf_size = _repo_gguf_size_bytes(repo_info)
+ if gguf_size > 0:
+ return gguf_size
+ return sum(
+ (getattr(f, "size_on_disk", None) or 0)
+ for rev in repo_info.revisions
+ for f in rev.files
+ )
+
+ def repo_last_modified(repo_info) -> float:
+ return max(
+ (getattr(rev, "last_modified", 0) or 0 for rev in repo_info.revisions),
+ default = 0,
+ )
+
+ target_repo = max(
+ matching_repos,
+ key = lambda repo_info: (repo_size(repo_info), repo_last_modified(repo_info)),
+ )
+
+ # Whole repo: the newest revision's snapshot dir holds the visible files.
+ revisions = sorted(
+ (rev for rev in target_repo.revisions if getattr(rev, "snapshot_path", None)),
+ key = lambda rev: getattr(rev, "last_modified", 0) or 0,
+ reverse = True,
+ )
+ for rev in revisions:
+ p = Path(rev.snapshot_path)
+ if p.exists():
+ return p
+ p = Path(target_repo.repo_path)
+ if p.exists():
+ return p
+ raise HTTPException(status_code = 404, detail = "Cached model path not found")
+
+
+def _wsl_reveal_in_explorer(path: Path) -> bool:
+ import subprocess
+
+ from utils.paths.path_utils import _IS_WSL
+
+ if not _IS_WSL:
+ return False
+ try:
+ windows_path = subprocess.run(
+ ["wslpath", "-w", str(path)],
+ capture_output = True,
+ text = True,
+ check = True,
+ timeout = 10,
+ ).stdout.strip()
+ if not windows_path:
+ return False
+ argument = f"/select,{windows_path}" if path.is_file() else windows_path
+ subprocess.Popen(["explorer.exe", argument])
+ return True
+ except (OSError, subprocess.SubprocessError):
+ return False
+
+
+def _reveal_in_file_manager(path: Path) -> None:
+ """Open the OS file manager with *path* selected (best effort per platform)."""
+ import subprocess
+
+ target = str(path)
+ if sys.platform == "darwin":
+ cmd = ["open", "-R", target] if path.is_file() else ["open", target]
+ subprocess.Popen(cmd)
+ elif os.name == "nt":
+ if path.is_file():
+ subprocess.Popen(["explorer", f"/select,{target}"])
+ else:
+ os.startfile(target) # noqa: S606 - local user's own file manager
+ elif not _wsl_reveal_in_explorer(path):
+ # No cross-desktop "select file" standard on Linux; open the directory.
+ directory = target if path.is_dir() else str(path.parent)
+ subprocess.Popen(["xdg-open", directory])
+
+
+class CachedModelPathResponse(BaseModel):
+ path: str
+ is_dir: bool
+
+
+@router.get("/cached-model-path", response_model = CachedModelPathResponse)
+async def get_cached_model_path(
+ repo_id: str = Query(..., description = "HuggingFace repo ID"),
+ variant: str = Query("", description = "Quantization variant (empty for whole repo)"),
+ current_subject: str = Depends(get_current_subject),
+):
+ """Absolute on-disk path of a cached repo or one of its GGUF variants."""
+ if not _is_valid_repo_id(repo_id):
+ raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
+ path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant.strip() or None)
+ return {"path": str(path), "is_dir": path.is_dir()}
+
+
+@router.post("/reveal-cached-model")
+async def reveal_cached_model(
+ repo_id: str = Body(...),
+ variant: Optional[str] = Body(None),
+ current_subject: str = Depends(get_current_subject),
+):
+ """Reveal a cached repo (or one GGUF variant's file) in the OS file manager."""
+ if not _is_valid_repo_id(repo_id):
+ raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
+ variant = (variant or "").strip() or None
+ path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant)
+ try:
+ await asyncio.to_thread(_reveal_in_file_manager, path)
+ except Exception as e:
+ logger.error(f"Failed to reveal {path}: {e}")
+ raise HTTPException(status_code = 500, detail = "Failed to open file manager")
+ return {"status": "ok", "path": str(path)}
+
+
@router.get("/checkpoints", response_model = CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(
diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py
index 7daa4224aa..a5fd71b6a0 100644
--- a/studio/backend/tests/test_chat_load_during_training.py
+++ b/studio/backend/tests/test_chat_load_during_training.py
@@ -801,6 +801,80 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(guard_called, [])
+ def _validate_gguf_template(
+ self,
+ *,
+ template,
+ canonical_path = "/picked/model.gguf",
+ ):
+ # Drive validate_model for a native lease-backed GGUF template probe and
+ # capture what the embedded-template reader was called with.
+ from models.inference import ValidateModelRequest
+
+ request = ValidateModelRequest(
+ model_path = "model.gguf",
+ gguf_variant = "Q4_K_M",
+ native_path_lease = "signed-lease",
+ include_chat_template = True,
+ )
+ cfg = SimpleNamespace(
+ identifier = canonical_path,
+ display_name = "model.gguf",
+ is_gguf = True,
+ is_lora = False,
+ is_vision = False,
+ gguf_file = canonical_path,
+ path = None,
+ base_model = None,
+ )
+ import utils.models.gguf_metadata as gguf_meta
+
+ seen = {}
+
+ def _fake_read(path):
+ seen["path"] = path
+ return template
+
+ guard_called = []
+ with (
+ patch.object(
+ self.route,
+ "_resolve_model_identifier_for_request",
+ return_value = (canonical_path, "model.gguf", True),
+ ),
+ patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
+ patch.object(self.route, "load_inference_config", return_value = {}),
+ patch.object(gguf_meta, "read_gguf_chat_template", _fake_read),
+ patch.object(
+ self.route,
+ "_guard_chat_load_against_training",
+ lambda *a, **kw: guard_called.append(True),
+ ),
+ ):
+ resp = asyncio.run(self.route.validate_model(request, current_subject = "u"))
+ return resp, seen, guard_called
+
+ def test_include_chat_template_reads_leased_gguf_embedded_template(self):
+ # The picker chat-template GET has no lease plumbing, so a native picked
+ # GGUF surfaces its default template through this lease-aware probe: the
+ # embedded template is read from the granted canonical path and returned.
+ resp, seen, _ = self._validate_gguf_template(template = "{{ messages }}")
+ self.assertEqual(resp.chat_template, "{{ messages }}")
+ # Read strictly the leased file's own embedded template, never a sibling
+ # sidecar: the grant authorizes just this one path.
+ self.assertEqual(seen["path"], "/picked/model.gguf")
+
+ def test_include_chat_template_skips_training_guard(self):
+ # A template-only probe allocates no VRAM, so like include_context_length
+ # it must not be refused by the training guard.
+ _, _, guard_called = self._validate_gguf_template(template = "{{ messages }}")
+ self.assertEqual(guard_called, [])
+
+ def test_include_chat_template_over_cap_is_dropped(self):
+ from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
+ resp, _, _ = self._validate_gguf_template(template = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1))
+ self.assertIsNone(resp.chat_template)
+
# ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ──────
diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py
index 761ea08e3f..5097f9f53a 100644
--- a/studio/backend/tests/test_export_absolute_paths.py
+++ b/studio/backend/tests/test_export_absolute_paths.py
@@ -158,6 +158,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None
utils_model_config._extract_quant_label = lambda value: value
utils_model_config._is_big_endian_gguf_path = lambda *args, **kwargs: False
+ utils_model_config._is_mtp_drafter = lambda *args, **kwargs: False
utils_model_config.is_audio_input_type = lambda *args, **kwargs: None
monkeypatch.setitem(
sys.modules,
diff --git a/studio/backend/tests/test_model_picker_regression.py b/studio/backend/tests/test_model_picker_regression.py
new file mode 100644
index 0000000000..f38a4d0b8d
--- /dev/null
+++ b/studio/backend/tests/test_model_picker_regression.py
@@ -0,0 +1,232 @@
+# 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 the model-picker per-model-config feature (the set of
+bugs that got the predecessor PR reverted). Pure-function / validation checks
+only, so they run on CPU in the backend pytest job with no model download.
+
+Covers, at the backend layer:
+ - infra-model hiding: the RAG embedder (bge-small-en-v1.5) and the llama.cpp
+ install-validation probe (ggml-org/models / stories260K) stay hidden, while
+ normal chat repos are not hidden;
+ - the HF token is honored from the dedicated header with the query string as a
+ fallback, never the other way around;
+ - the chat-template byte caps reject oversized overrides (both the char-count
+ fast path and the UTF-8 byte path) and the sidecar reader is size-bounded.
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+# Keep this test runnable without the optional structlog dependency (mirrors
+# tests/test_cached_gguf_routes.py), since importing routes.models pulls it in.
+if "structlog" not in sys.modules:
+
+ class _DummyLogger:
+ def __getattr__(self, _name):
+ return lambda *args, **kwargs: None
+
+ sys.modules["structlog"] = types.SimpleNamespace(
+ BoundLogger = _DummyLogger,
+ get_logger = lambda *args, **kwargs: _DummyLogger(),
+ )
+
+import routes.models as models_route
+from core.rag import config as rag_config
+from hub.dependencies import get_hf_token
+from models.inference import LoadRequest
+from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
+from picker.service import _read_bounded_text
+from utils.hidden_models import is_hidden_model
+
+
+@pytest.fixture(autouse = True)
+def _pin_default_embedder(monkeypatch):
+ """Pin the effective embedder to Studio's static default so hiding is
+ deterministic and cannot depend on ambient RAG config / env."""
+ default = "unsloth/bge-small-en-v1.5"
+ monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", default, raising = False)
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: default)
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: default)
+ monkeypatch.setattr(rag_config, "default_gguf_repo", lambda: default)
+
+
+# --------------------------------------------------------------------------- #
+# Infra-model hiding (the "infra models resurfaced in the picker" regression) #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+ "value",
+ [
+ "ggml-org/models", # the probe repo id
+ "unsloth/bge-small-en-v1.5", # the RAG embedder repo
+ "unsloth/bge-small-en-v1.5-GGUF", # its GGUF companion
+ "/root/.cache/huggingface/hub/x/stories260K.gguf", # probe on disk
+ "/root/.cache/x/Stories260K.GGUF", # case-insensitive
+ r"C:\\models\\stories260K.gguf", # windows-style path
+ "/opt/models/bge-small-en-v1.5", # embedder basename folder
+ "/opt/models/bge-small-en-v1.5-Q8_0.gguf", # suffixed local weight
+ ],
+)
+def test_infra_models_are_hidden(value):
+ assert is_hidden_model(value) is True
+
+
+@pytest.mark.parametrize(
+ "value",
+ [
+ "unsloth/gemma-3-270m-it-GGUF", # a normal small chat GGUF
+ "unsloth/Qwen3-0.6B", # a normal non-GGUF chat model
+ "user/stories260K-finetune-GGUF", # repo id merely contains "stories260k"
+ "user/model-chat", # generic repo must not be hidden
+ "meta-llama/Llama-3.1-8B-Instruct",
+ ],
+)
+def test_normal_models_are_not_hidden(value):
+ assert is_hidden_model(value) is False
+
+
+def test_is_hidden_model_ignores_empty_values():
+ assert is_hidden_model(None) is False
+ assert is_hidden_model("") is False
+ assert is_hidden_model(None, "", "unsloth/gemma-3-270m-it-GGUF") is False
+
+
+def test_hidden_model_matchers_expose_probe_needles():
+ needles, exact_ids, _exact_paths = models_route.hidden_model_matchers()
+ lowered = [n.lower() for n in needles]
+ assert "ggml-org/models" in lowered
+ assert "stories260k.gguf" in lowered
+ # The configured embedder is exposed as an exact repo id, never as a
+ # basename needle that would substring-hide unrelated chat models.
+ assert "bge-small-en-v1.5" not in lowered
+ assert "unsloth/bge-small-en-v1.5" in exact_ids
+
+
+def test_hidden_model_matchers_custom_repo_publishes_exact_ids(monkeypatch):
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
+ needles, exact_ids, exact_paths = models_route.hidden_model_matchers()
+ assert needles == ["ggml-org/models", "stories260k.gguf"]
+ assert "org/model" in exact_ids
+ assert "org/model-gguf" in exact_ids
+ assert exact_paths == []
+
+
+def test_hidden_model_matchers_local_owner_name_path_is_exact_path(monkeypatch, tmp_path):
+ # A local embedder shaped like owner/name that exists on disk must be an
+ # exact resolved path, not a Hub repo id (mirroring is_hidden_model), so the
+ # local row stays hidden instead of showing as a chat model.
+ (tmp_path / "models" / "embedder").mkdir(parents = True)
+ monkeypatch.chdir(tmp_path)
+ monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder")
+ monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "ggml-org/models")
+ _needles, exact_ids, exact_paths = models_route.hidden_model_matchers()
+ resolved = str((tmp_path / "models" / "embedder").resolve()).lower()
+ assert resolved in exact_paths
+ assert "models/embedder" not in exact_ids
+
+
+# --------------------------------------------------------------------------- #
+# HF token via header, query string only as a fallback (the token-leak fix) #
+# --------------------------------------------------------------------------- #
+
+
+def test_get_hf_token_strips_and_returns():
+ assert get_hf_token(" hf_abc ") == "hf_abc"
+
+
+@pytest.mark.parametrize("value", [None, "", " ", "\n\t"])
+def test_get_hf_token_blank_is_none(value):
+ assert get_hf_token(value) is None
+
+
+@pytest.mark.parametrize(
+ "value,expected",
+ [(" hf_x ", "hf_x"), ("", None), (" ", None), (None, None), (1234, None)],
+)
+def test_normalize_hf_token(value, expected):
+ assert models_route._normalize_hf_token(value) == expected
+
+
+def test_header_token_wins_over_query():
+ header, query = "hf_header", "hf_query"
+ resolved = models_route._normalize_hf_token(header) or models_route._normalize_hf_token(query)
+ assert resolved == "hf_header"
+
+
+def test_query_token_is_fallback_when_header_absent():
+ resolved = models_route._normalize_hf_token(None) or models_route._normalize_hf_token(
+ "hf_query"
+ )
+ assert resolved == "hf_query"
+
+
+# --------------------------------------------------------------------------- #
+# Chat-template byte caps (the unbounded-template hardening) #
+# --------------------------------------------------------------------------- #
+
+
+def _load_request(**overrides):
+ data = {"model_path": "unsloth/test-model-GGUF", "gguf_variant": "Q4_K_M"}
+ data.update(overrides)
+ return LoadRequest.model_validate(data)
+
+
+def test_blank_chat_template_override_normalizes_to_none():
+ assert _load_request(chat_template_override = " \n\t").chat_template_override is None
+
+
+def test_nonblank_chat_template_override_preserved_verbatim():
+ template = " {{ messages }} "
+ assert _load_request(chat_template_override = template).chat_template_override == template
+
+
+def test_chat_template_at_byte_limit_is_accepted():
+ template = "a" * MAX_CHAT_TEMPLATE_BYTES # exactly the limit, 1 byte/char
+ assert (
+ len(_load_request(chat_template_override = template).chat_template_override)
+ == MAX_CHAT_TEMPLATE_BYTES
+ )
+
+
+def test_chat_template_over_char_limit_is_rejected():
+ with pytest.raises(Exception): # pydantic ValidationError wrapping ValueError
+ _load_request(chat_template_override = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1))
+
+
+def test_chat_template_over_byte_limit_is_rejected():
+ # Char count stays under the limit but UTF-8 bytes exceed it (3 bytes/char),
+ # so only the byte-count branch can catch this.
+ multibyte = "€" * (MAX_CHAT_TEMPLATE_BYTES // 2) # euro sign, 3 bytes each
+ assert len(multibyte) <= MAX_CHAT_TEMPLATE_BYTES
+ assert len(multibyte.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES
+ with pytest.raises(Exception):
+ _load_request(chat_template_override = multibyte)
+
+
+def test_read_bounded_text_reads_within_limit(tmp_path):
+ p = tmp_path / "t.json"
+ p.write_text("hello", encoding = "utf-8")
+ assert _read_bounded_text(p, 16) == "hello"
+
+
+def test_read_bounded_text_rejects_over_limit(tmp_path):
+ p = tmp_path / "big.json"
+ p.write_bytes(b"x" * 100)
+ assert _read_bounded_text(p, 50) is None
+
+
+def test_read_bounded_text_at_limit_is_read(tmp_path):
+ p = tmp_path / "exact.json"
+ p.write_bytes(b"x" * 50)
+ assert _read_bounded_text(p, 50) == "x" * 50
+
+
+def test_read_bounded_text_missing_file_is_none(tmp_path):
+ assert _read_bounded_text(tmp_path / "nope.json", 50) is None
diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py
index edf55812e2..7d98766616 100644
--- a/studio/backend/tests/test_model_update_robustness.py
+++ b/studio/backend/tests/test_model_update_robustness.py
@@ -314,6 +314,7 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
file_name = "model.safetensors",
size_on_disk = 100,
blob_path = str(repo_path / "blobs" / "modelsha"),
+ blob_last_modified = 3_000.0,
),
]
)
@@ -336,6 +337,51 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
assert rows[0]["repo_id"] == "Org/SafeTensorRepo"
assert rows[0]["model_format"] == "safetensors"
assert rows[0]["size_bytes"] == 100
+ assert rows[0]["last_modified"] == 3_000.0
+
+
+def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path):
+ repo_path = tmp_path / "models--Org--GgufRepo"
+ repo = SimpleNamespace(
+ repo_id = "Org/GgufRepo",
+ repo_type = "model",
+ repo_path = repo_path,
+ revisions = [
+ SimpleNamespace(
+ files = [
+ SimpleNamespace(
+ file_name = "model-Q4_K_M.gguf",
+ size_on_disk = 100,
+ blob_path = None,
+ blob_last_modified = 5_000.0,
+ ),
+ ]
+ )
+ ],
+ )
+ monkeypatch.setattr(
+ CI,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [repo])],
+ )
+ monkeypatch.setattr(
+ CI.hf_cache_scan,
+ "is_gguf_repo_partial",
+ lambda *args, **kwargs: False,
+ )
+ monkeypatch.setattr(
+ CI,
+ "_gguf_variant_state_summary",
+ lambda _repo_id: (False, 0),
+ )
+
+ rows = CI._scan_cached_gguf()
+
+ assert len(rows) == 1
+ assert rows[0]["repo_id"] == "Org/GgufRepo"
+ assert rows[0]["model_format"] == "gguf"
+ assert rows[0]["size_bytes"] == 100
+ assert rows[0]["last_modified"] == 5_000.0
# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ───
@@ -636,3 +682,18 @@ def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch
assert snap.exists() is True # the current file must survive
assert result["removed_snapshots"] == 0
assert result["deleted_blobs"] == 0
+
+
+def _mmproj_repo(*file_names: str):
+ return SimpleNamespace(
+ revisions = [SimpleNamespace(files = [SimpleNamespace(file_name = n) for n in file_names])]
+ )
+
+
+def test_repo_has_mmproj_requires_gguf_projector():
+ # A non-GGUF sidecar whose name merely contains "mmproj" must NOT mark the
+ # repo vision-capable; the runtime's projector detection is GGUF-only.
+ assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj_config.json")) is False
+ assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "README-mmproj.md")) is False
+ # A real GGUF projector still marks the repo vision-capable.
+ assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj-F16.gguf")) is True
diff --git a/studio/backend/tests/test_picker_service.py b/studio/backend/tests/test_picker_service.py
new file mode 100644
index 0000000000..be7ea18f03
--- /dev/null
+++ b/studio/backend/tests/test_picker_service.py
@@ -0,0 +1,266 @@
+# 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 types import SimpleNamespace
+
+from picker.service import (
+ MAX_TEMPLATE_METADATA_BYTES,
+ _chat_template_from_dir,
+ _chat_template_from_processor_json,
+ _chat_template_from_tokenizer_config,
+ _chat_template_from_tokenizer_dir,
+ _find_gguf_in_dir,
+ _iter_ggufs,
+ read_default_chat_template,
+ validate_chat_template,
+)
+
+
+def test_iter_ggufs_skips_gguf_companions(tmp_path):
+ mtp_dir = tmp_path / "MTP"
+ mtp_dir.mkdir()
+ main = tmp_path / "model-Q8_0.gguf"
+ main.write_bytes(b"")
+ (tmp_path / "mmproj-F16.gguf").write_bytes(b"")
+ (tmp_path / "mtp-model-Q8_0.gguf").write_bytes(b"")
+ (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
+ (tmp_path / "model-Q8_0-be.gguf").write_bytes(b"")
+
+ assert _iter_ggufs(tmp_path) == [main]
+
+
+def test_find_gguf_in_dir_matches_quant_label(tmp_path):
+ mtp_dir = tmp_path / "MTP"
+ mtp_dir.mkdir()
+ main = tmp_path / "model-Q8_0.gguf"
+ main.write_bytes(b"")
+ (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+
+ assert _find_gguf_in_dir(tmp_path, "Q8_0") == main
+ assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
+
+
+def test_find_gguf_in_dir_without_variant_prefers_largest_model(tmp_path):
+ smaller = tmp_path / "a-model-Q4_K_M.gguf"
+ larger = tmp_path / "z-model-Q8_0.gguf"
+ smaller.write_bytes(b"0")
+ larger.write_bytes(b"00")
+
+ assert _find_gguf_in_dir(tmp_path, None) == larger
+
+
+def test_find_gguf_in_dir_without_variant_prefers_first_split(tmp_path):
+ first = tmp_path / "model-Q4_K_M-00001-of-00003.gguf"
+ second = tmp_path / "model-Q4_K_M-00002-of-00003.gguf"
+ third = tmp_path / "model-Q4_K_M-00003-of-00003.gguf"
+ first.write_bytes(b"0")
+ second.write_bytes(b"000")
+ third.write_bytes(b"00")
+
+ assert _find_gguf_in_dir(tmp_path, None) == first
+
+ first.unlink()
+ assert _find_gguf_in_dir(tmp_path, None) == second
+
+
+def test_find_gguf_in_dir_matches_bpw_variant_base_label(tmp_path):
+ target = tmp_path / "model-IQ4_XS-3.53bpw.gguf"
+ target.write_bytes(b"")
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+
+ assert _find_gguf_in_dir(tmp_path, "IQ4_XS") == target
+ assert _find_gguf_in_dir(tmp_path, "IQ4_XS-3.53bpw") == target
+ assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
+
+
+def test_validate_chat_template_accepts_valid_and_empty():
+ assert validate_chat_template("{{ messages[0].content }}").valid is True
+ assert validate_chat_template("").valid is True
+ assert validate_chat_template(" ").valid is True
+
+
+def test_validate_chat_template_reports_syntax_error_with_line():
+ result = validate_chat_template("{% if %}{% endif %}")
+ assert result.valid is False
+ assert result.error is not None
+ assert result.error.startswith("Line ")
+
+
+def test_chat_template_from_tokenizer_config_reads_string():
+ assert _chat_template_from_tokenizer_config({"chat_template": "HELLO"}) == "HELLO"
+ assert _chat_template_from_tokenizer_config({"chat_template": " "}) is None
+ assert _chat_template_from_tokenizer_config({}) is None
+
+
+def test_chat_template_from_tokenizer_config_prefers_named_default():
+ config = {
+ "chat_template": [
+ {"name": "tool_use", "template": "TOOL"},
+ {"name": "default", "template": "DEFAULT"},
+ ]
+ }
+ assert _chat_template_from_tokenizer_config(config) == "DEFAULT"
+
+
+def test_chat_template_from_tokenizer_config_falls_back_to_first_entry():
+ config = {
+ "chat_template": [
+ {"name": "tool_use", "template": "TOOL"},
+ {"name": "other", "template": "OTHER"},
+ ]
+ }
+ assert _chat_template_from_tokenizer_config(config) == "TOOL"
+
+
+def test_chat_template_from_tokenizer_dir_prefers_jinja_file(tmp_path):
+ (tmp_path / "chat_template.jinja").write_text("FROM_JINJA", encoding = "utf-8")
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_JINJA"
+
+
+def test_chat_template_from_tokenizer_dir_reads_tokenizer_config(tmp_path):
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG"
+
+
+def test_chat_template_from_dir_without_variant_prefers_tokenizer(tmp_path):
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ assert _chat_template_from_dir(tmp_path) == "FROM_CONFIG"
+
+
+def test_chat_template_from_dir_with_variant_still_prefers_tokenizer(tmp_path, monkeypatch):
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # Selecting a variant must not flip precedence to the embedded GGUF template.
+ assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_CONFIG"
+
+
+def test_chat_template_from_dir_with_variant_falls_back_to_gguf(tmp_path, monkeypatch):
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # With no tokenizer sidecar, the embedded GGUF template is still the fallback.
+ assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_GGUF"
+
+
+def test_chat_template_from_dir_returns_none_when_absent(tmp_path):
+ assert _chat_template_from_dir(tmp_path) is None
+
+
+def test_read_default_chat_template_direct_gguf_prefers_sidecar(tmp_path, monkeypatch):
+ gguf = tmp_path / "model-Q4_K_M.gguf"
+ gguf.write_bytes(b"")
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # A directly selected .gguf must prefer a maintained sidecar over its embedded copy.
+ assert read_default_chat_template(str(gguf)) == "FROM_CONFIG"
+
+
+def test_read_default_chat_template_direct_gguf_falls_back_to_embedded(tmp_path, monkeypatch):
+ gguf = tmp_path / "model-Q4_K_M.gguf"
+ gguf.write_bytes(b"")
+ monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # With no sidecar next to the file, the embedded GGUF template is the fallback.
+ assert read_default_chat_template(str(gguf)) == "FROM_GGUF"
+
+
+def test_tokenizer_config_over_size_limit_is_skipped_not_parsed(tmp_path):
+ # An oversized tokenizer_config.json must be skipped before json.loads so a
+ # hostile sidecar cannot exhaust memory.
+ padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024)
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "HELLO", "_pad": padding}), encoding = "utf-8"
+ )
+ assert _chat_template_from_tokenizer_dir(tmp_path) is None
+
+
+def test_processor_json_over_size_limit_is_skipped_not_parsed(tmp_path):
+ padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024)
+ (tmp_path / "chat_template.json").write_text(
+ json.dumps({"default": "HELLO", "_pad": padding}), encoding = "utf-8"
+ )
+ assert _chat_template_from_processor_json(tmp_path) is None
+
+
+def test_tokenizer_config_at_size_limit_is_still_read(tmp_path):
+ # A normal-sized config is unaffected by the bound (regression guard).
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG"
+
+
+def test_remote_template_over_size_limit_is_skipped_before_download(monkeypatch):
+ # An uncached Hub repo whose template exceeds the cap must be skipped via the
+ # remote size pre-check, never downloaded.
+ import huggingface_hub
+
+ monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name)
+ monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: [])
+
+ def _fail_download(*args, **kwargs):
+ raise AssertionError("oversized remote template must not be downloaded")
+
+ def _fake_get_paths_info(self, repo_id, paths, **kwargs):
+ return [SimpleNamespace(path = p, size = MAX_TEMPLATE_METADATA_BYTES + 1) for p in paths]
+
+ monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fail_download)
+ monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info)
+
+ assert read_default_chat_template("org/oversized-model") is None
+
+
+def test_remote_oversized_jinja_falls_through_to_tokenizer_template(tmp_path, monkeypatch):
+ # A raw chat_template.jinja between the response cap (MAX_CHAT_TEMPLATE_BYTES)
+ # and the download bound (MAX_TEMPLATE_METADATA_BYTES) must not be returned: the
+ # route drops it, so the remote path must skip the oversized Jinja and fall
+ # through to the smaller tokenizer_config.json.
+ import huggingface_hub
+ from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
+
+ big_jinja = tmp_path / "chat_template.jinja"
+ big_jinja.write_text("{{ x }}" * (MAX_CHAT_TEMPLATE_BYTES // 4), encoding = "utf-8")
+ assert MAX_CHAT_TEMPLATE_BYTES < big_jinja.stat().st_size < MAX_TEMPLATE_METADATA_BYTES
+ tokenizer_config = tmp_path / "tokenizer_config.json"
+ tokenizer_config.write_text(json.dumps({"chat_template": "SMALL_TEMPLATE"}), encoding = "utf-8")
+ files = {
+ "chat_template.jinja": big_jinja,
+ "tokenizer_config.json": tokenizer_config,
+ }
+
+ monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name)
+ monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: [])
+
+ def _fake_download(repo_id, rel, **kwargs):
+ target = files.get(rel)
+ if target is None:
+ raise FileNotFoundError(rel)
+ return str(target)
+
+ def _fake_get_paths_info(self, repo_id, paths, **kwargs):
+ return [
+ SimpleNamespace(
+ path = p,
+ size = files[p].stat().st_size if p in files else 0,
+ )
+ for p in paths
+ ]
+
+ monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download)
+ monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info)
+
+ assert read_default_chat_template("org/big-jinja-model") == "SMALL_TEMPLATE"
diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py
index 50b3cd3513..749f2c9234 100644
--- a/studio/backend/utils/models/gguf_metadata.py
+++ b/studio/backend/utils/models/gguf_metadata.py
@@ -50,10 +50,14 @@ _CACHE_MAX_ENTRIES = 4096
# keyed by (file cache key, wanted key). None = key absent / file unreadable.
_BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {}
+_STRING_CACHE: Dict[Tuple[_CacheKey, str], Optional[str]] = {}
+
# GGUF header dims for the staged/deferred-load UI: context_length, layer_count
# (block_count), and moe_layer_count (block_count minus leading dense layers; 0
# if not MoE). One cached pass fills all three so the staged sheet can size every
-# slider before the model loads. None = unreadable / not a GGUF.
+# slider before the model loads. None = unreadable / not a GGUF. The native
+# training context length (``{arch}.context_length``) the UI shows before a model
+# loads is read from here via read_gguf_context_length.
_DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {}
@@ -408,6 +412,83 @@ def _read_gguf_bool(path: str, wanted_key: str) -> Optional[bool]:
return result
+def _parse_gguf_string(path: str, wanted_key: str) -> Optional[str]:
+ try:
+ with open(path, "rb") as f:
+ head = f.read(24)
+ if len(head) < 24:
+ return None
+ magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20:
+ break
+ kbytes = f.read(klen)
+ if len(kbytes) < klen:
+ break
+ key = kbytes.decode("utf-8", "replace")
+ vt_bytes = f.read(4)
+ if len(vt_bytes) < 4:
+ break
+ vtype = struct.unpack(" 1 << 22:
+ break
+ sbytes = f.read(slen)
+ if len(sbytes) < slen:
+ break
+ return sbytes.decode("utf-8", "replace")
+ if not _skip_gguf_value(f, vtype):
+ break
+ except (struct.error, UnicodeDecodeError):
+ break
+ except OSError as e:
+ logger.debug(f"_parse_gguf_string: cannot open {path}: {e}")
+ return None
+ except Exception as e:
+ logger.debug(f"_parse_gguf_string: parse failure on {path}: {e}")
+ return None
+ return None
+
+
+def _read_gguf_string(path: str, wanted_key: str) -> Optional[str]:
+ fkey = _cache_key(path)
+ if fkey is None:
+ return None
+ ckey = (fkey, wanted_key)
+ with _CACHE_LOCK:
+ if ckey in _STRING_CACHE:
+ return _STRING_CACHE[ckey]
+ result = _parse_gguf_string(path, wanted_key)
+ with _CACHE_LOCK:
+ while len(_STRING_CACHE) >= _CACHE_MAX_ENTRIES:
+ try:
+ _STRING_CACHE.pop(next(iter(_STRING_CACHE)))
+ except StopIteration:
+ break
+ _STRING_CACHE[ckey] = result
+ return result
+
+
+def read_gguf_chat_template(path: str) -> Optional[str]:
+ template = _read_gguf_string(path, "tokenizer.chat_template")
+ if isinstance(template, str) and template.strip():
+ return template
+ return None
+
+
def read_mmproj_audio_capability(path: str) -> Optional[bool]:
"""``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's
gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable.
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index 57e890dd5a..7137fd6f96 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -196,9 +196,6 @@ function RootLayout() {
chatRuntime.setActiveThreadId(null);
chatRuntime.setActiveProjectId(null);
chatRuntime.setIncognito(false);
- // Detach the staging UI but keep any in-flight download running, like Hub.
- if (chatRuntime.pendingSelection)
- chatRuntime.abandonStagedModel({ keepDownload: true });
void navigate({
to: "/chat",
search: { new: crypto.randomUUID() },
@@ -221,10 +218,6 @@ function RootLayout() {
chatRuntime.setActiveProjectId(null);
chatRuntime.setActiveThreadId(null);
chatRuntime.setIncognito(false);
- // Leaving chat must not kill an in-flight download: detach the staging UI
- // but keep the transfer running in the manager, like a Hub download.
- if (chatRuntime.pendingSelection)
- chatRuntime.abandonStagedModel({ keepDownload: true });
}, [isChatRoute]);
return (
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 8eab03133b..293971904f 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -1011,28 +1011,6 @@ export function AppSidebar() {
- {isPinned ? (
-
-
- {
- e.stopPropagation();
- unpinChat(item.id);
- }}
- aria-label="Unpin chat"
- className={cn(actionClass, "is-unpin-action")}
- >
-
-
-
-
-
-
- Unpin
-
-
- ) : null}
);
}
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
deleted file mode 100644
index 08492ab480..0000000000
--- a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
+++ /dev/null
@@ -1,79 +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
-
-// Per-model pre-load inference settings, persisted in localStorage so the load
-// dialog can offer "Remember settings for ". GGUF picks only: every
-// field is a llama.cpp load knob, so all save/restore call sites gate on
-// GGUF-ness (a non-GGUF blob would only snapshot leftover standing values).
-
-const KEY = "unsloth_load_settings";
-
-export interface RememberedLoadSettings {
- contextLength: number | null;
- kvCacheDtype: string | null;
- speculativeType: string | null;
- specDraftNMax: number | null;
- tensorParallel: boolean;
- // GPU Memory controls. Optional so an older blob (which lacked them) still
- // parses, leaving the live knobs untouched on apply. The mode is kept with the
- // manual knobs (gpuLayers/nCpuMoe are ignored outside Manual mode). A null
- // selectedGpuIds is meaningful (all GPUs), so it's distinguished from absent.
- // The per-GPU split ratio is deliberately NOT remembered: it's positionally
- // bound to the exact GPU set/order and unvalidated, so it would mismatch.
- gpuMemoryMode?: "auto" | "manual";
- gpuLayers?: number;
- nCpuMoe?: number;
- selectedGpuIds?: number[] | null;
-}
-
-// Storage key for a pick's remembered settings, scoped per quant (the VRAM-budget
-// knobs differ per quant). An HF repo collapses its GGUF variants into one `id`,
-// so fold the variant in. Local .gguf paths are already file-specific; native
-// drag-drop files key by display label, so same-named files share an entry.
-export function rememberedLoadSettingsKey(selection: {
- id: string;
- ggufVariant?: string | null;
-}): string {
- return selection.ggufVariant
- ? `${selection.id}::${selection.ggufVariant}`
- : selection.id;
-}
-
-function readAll(): Record {
- try {
- return JSON.parse(localStorage.getItem(KEY) ?? "{}");
- } catch {
- return {};
- }
-}
-
-function writeAll(all: Record) {
- try {
- localStorage.setItem(KEY, JSON.stringify(all));
- } catch {
- // Ignore quota / unavailable storage.
- }
-}
-
-export function loadRememberedLoadSettings(
- key: string,
-): RememberedLoadSettings | null {
- return readAll()[key] ?? null;
-}
-
-export function saveRememberedLoadSettings(
- key: string,
- settings: RememberedLoadSettings,
-) {
- const all = readAll();
- all[key] = settings;
- writeAll(all);
-}
-
-export function clearRememberedLoadSettings(key: string) {
- const all = readAll();
- if (key in all) {
- delete all[key];
- writeAll(all);
- }
-}
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 7083f02288..b0127b5e40 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -2,10 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getAuthToken } from "@/features/auth";
-import {
- loadRememberedLoadSettings,
- rememberedLoadSettingsKey,
-} from "@/components/assistant-ui/model-selector/remembered-load-settings";
+import { resolveInitialConfig } from "@/features/model-picker";
import { projectHasSources } from "@/features/rag/api/rag-api";
import { apiUrl } from "@/lib/api-base";
import { parseParamCountB } from "@/lib/model-size";
@@ -46,7 +43,7 @@ import {
type PendingImageEditReference,
type RagAutoInject,
GPU_LAYERS_AUTO,
- loadedGpuMemoryFieldsUnlessStaged,
+ loadedGpuMemoryFields,
reconcilePersistedGpuIds,
resolveLoadedSpeculativeSettings,
resolveSpeculativeSettingsForLoad,
@@ -1533,65 +1530,56 @@ async function autoLoadSmallestModel(): Promise<{
return false;
}
const currentStore = useChatRuntimeStore.getState();
- // Blobs are saved for GGUF picks only (the sheet gates on it), so don't
- // let a legacy non-GGUF blob feed a stale context/spec choice into a
- // safetensors auto-load.
- const remembered =
- candidate.kind === "gguf"
- ? loadRememberedLoadSettings(
- rememberedLoadSettingsKey({
- id: candidate.id,
- ggufVariant: candidate.ggufVariant,
- }),
- )
- : null;
+ const { config } = resolveInitialConfig(candidate.id, candidate.ggufVariant);
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
modelId: candidate.id,
ggufVariant: candidate.ggufVariant,
isGguf: candidate.kind === "gguf",
- customContextLength: remembered?.contextLength ?? null,
+ customContextLength: config.customContextLength,
ggufContextLength: null,
currentCheckpoint: currentStore.params.checkpoint,
activeGgufVariant: currentStore.activeGgufVariant,
- maxSeqLength: candidate.maxSeqLength,
+ maxSeqLength: config.maxSeqLength ?? candidate.maxSeqLength,
presetSource: currentStore.activePresetSource,
});
- // The GPU knobs are per-model, so read them from the same remembered
- // settings that fed effectiveMaxSeqLength -- on a background auto-load the
- // live store holds session defaults, not the saved Manual mode / layer pin /
- // GPU pick. Absent fields fall back like applyRememberedLoadSettings: the
- // mode to the store (a persisted standing preference), the per-model knobs to
- // their defaults. The saved GPU pick is reconciled against the GPUs present
- // now, like the interactive restore.
+ // The GPU knobs are per-model, so read them from the same per-model config
+ // that fed effectiveMaxSeqLength -- on a background auto-load the live store
+ // holds session defaults, not the saved Manual mode / layer pin / GPU pick.
+ // Absent fields fall back like the interactive restore: the mode to the store
+ // (a persisted standing preference), the per-model knobs to their defaults.
+ // The saved GPU pick is reconciled against the GPUs present now.
const effectiveGpuMemoryMode =
- remembered?.gpuMemoryMode ?? currentStore.gpuMemoryMode;
- const effectiveGpuLayers = remembered?.gpuLayers ?? GPU_LAYERS_AUTO;
- const effectiveNCpuMoe = remembered?.nCpuMoe ?? 0;
- if (remembered?.selectedGpuIds != null) {
+ config.gpuMemoryMode ?? currentStore.gpuMemoryMode;
+ const effectiveGpuLayers = config.gpuLayers ?? GPU_LAYERS_AUTO;
+ const effectiveNCpuMoe = config.nCpuMoe ?? 0;
+ if (config.selectedGpuIds != null) {
// Warm the device cache first: on a cold cache the reconcile passes the
// saved pick through unvalidated, and a stale cross-host pick then fails
// the load with the picker hidden.
await ensureGpuDeviceCache();
}
const effectiveGpuIds =
- remembered?.selectedGpuIds !== undefined
- ? reconcilePersistedGpuIds(remembered.selectedGpuIds)
+ config.selectedGpuIds !== undefined
+ ? reconcilePersistedGpuIds(config.selectedGpuIds)
: null;
// Under Manual GPU memory + Auto layers, llama.cpp's --fit owns context
// sizing, so send 0 (or the pinned length). GGUF-only; a no-op otherwise.
- // The context pin is per-model too, so it comes from remembered settings,
- // not the live store.
+ // The context pin is per-model too, so it comes from the saved config, not
+ // the live store.
const fitMaxSeqLength = resolveFitMaxSeqLength(
candidate.kind === "gguf",
effectiveGpuMemoryMode,
effectiveGpuLayers,
- remembered?.contextLength ?? null,
+ config.customContextLength ?? null,
effectiveMaxSeqLength,
);
const effectiveSpeculativeType =
- remembered?.speculativeType ?? specSettings.speculativeType;
+ config.speculativeType ?? specSettings.speculativeType;
const effectiveSpecDraftNMax =
- remembered?.specDraftNMax ?? specSettings.specDraftNMax;
+ config.specDraftNMax ?? specSettings.specDraftNMax;
+ const effectiveChatTemplateOverride = config.chatTemplateOverride?.trim()
+ ? config.chatTemplateOverride
+ : null;
if (
!(await canAutoLoad({
model_path: candidate.id,
@@ -1621,10 +1609,11 @@ async function autoLoadSmallestModel(): Promise<{
is_lora: false,
gguf_variant: candidate.ggufVariant,
trust_remote_code: trustRemoteCode,
- cache_type_kv: remembered?.kvCacheDtype ?? null,
+ chat_template_override: effectiveChatTemplateOverride,
+ cache_type_kv: config.kvCacheDtype,
speculative_type: effectiveSpeculativeType,
spec_draft_n_max: effectiveSpecDraftNMax,
- tensor_parallel: remembered?.tensorParallel ?? false,
+ tensor_parallel: config.tensorParallel,
// GGUF-only: the safetensors fallback loads via HF auto-placement (no
// explicit pins). The split ratio is deliberately never remembered
// (positionally bound to an exact GPU set), so auto-load leaves llama.cpp's
@@ -1638,7 +1627,12 @@ async function autoLoadSmallestModel(): Promise<{
}
: {}),
});
- saveSpeculativeType(effectiveSpeculativeType);
+ // Only persist the global preference when the value came from the global
+ // settings. A per-model config's choice must stay load-local, or autoloading
+ // a remembered model on startup would rewrite the global default.
+ if (config.speculativeType == null) {
+ saveSpeculativeType(effectiveSpeculativeType);
+ }
// Self-gates on is_gguf (skips diffusion), so persists only for a real GGUF load.
persistGpuMemoryModeOnLoad(loadResp, effectiveGpuMemoryMode);
useChatRuntimeStore
@@ -1650,6 +1644,9 @@ async function autoLoadSmallestModel(): Promise<{
);
store.setParams({
...store.params,
+ ...(candidate.kind === "gguf"
+ ? {}
+ : { maxSeqLength: effectiveMaxSeqLength }),
maxTokens:
candidate.kind === "gguf"
? loadResp.context_length ?? 131072
@@ -1676,7 +1673,7 @@ async function autoLoadSmallestModel(): Promise<{
const keepCustomCtx = resolveManualAutoCtxPin(
effectiveGpuMemoryMode,
effectiveGpuLayers,
- remembered?.contextLength ?? null,
+ config.customContextLength ?? null,
);
useChatRuntimeStore.setState({
ggufContextLength: loadResp.context_length ?? 131072,
@@ -1694,13 +1691,14 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
- ...loadedGpuMemoryFieldsUnlessStaged(loadResp, {
- customContextLength: keepCustomCtx,
- }),
+ ...loadedGpuMemoryFields(loadResp),
loadedCustomContextLength: keepCustomCtx,
defaultChatTemplate: loadResp.chat_template ?? null,
- chatTemplateOverride: null,
- loadedChatTemplateOverride: null,
+ chatTemplateOverride: effectiveChatTemplateOverride,
+ loadedChatTemplateOverride: effectiveChatTemplateOverride,
+ // Retain the saved requested context so re-saving the config keeps the
+ // override; null stays null (auto/VRAM-fit).
+ customContextLength: config.customContextLength,
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
...resolveLoadedSpeculativeSettings(loadResp),
@@ -1720,10 +1718,11 @@ async function autoLoadSmallestModel(): Promise<{
loadedTensorParallel: loadResp.tensor_parallel ?? false,
// Non-GGUF response: clears any stale GPU baseline a prior manual-GPU
// GGUF load left, matching the interactive/status sibling load paths.
- ...loadedGpuMemoryFieldsUnlessStaged(loadResp),
+ ...loadedGpuMemoryFields(loadResp),
defaultChatTemplate: loadResp.chat_template ?? null,
- chatTemplateOverride: null,
- loadedChatTemplateOverride: null,
+ chatTemplateOverride: effectiveChatTemplateOverride,
+ loadedChatTemplateOverride: effectiveChatTemplateOverride,
+ customContextLength: null,
...resolveLoadedSpeculativeSettings(loadResp),
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
@@ -1988,7 +1987,7 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
- ...loadedGpuMemoryFieldsUnlessStaged(loadResp),
+ ...loadedGpuMemoryFields(loadResp),
// Drives the GPU Memory controls' diffusion gate; set alongside the
// GPU fields on every load path so the gate can't read stale.
loadedIsDiffusion: loadResp.is_diffusion ?? false,
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index 631474c39a..de3e5e370c 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -377,14 +377,33 @@ export async function listCachedModels(
return data.cached;
}
-export async function deleteCachedModel(
+export interface CachedModelPath {
+ path: string;
+ is_dir: boolean;
+}
+
+/** Absolute on-disk path of a cached repo or one of its GGUF variants. */
+export async function getCachedModelPath(
+ repoId: string,
+ variant?: string,
+): Promise {
+ const params = new URLSearchParams({ repo_id: repoId });
+ if (variant) params.set("variant", variant);
+ const response = await authFetch(
+ `/api/models/cached-model-path?${params.toString()}`,
+ );
+ return parseJsonOrThrow(response);
+}
+
+/** Reveal a cached repo (or one GGUF variant's file) in the OS file manager. */
+export async function revealCachedModel(
repoId: string,
variant?: string,
): Promise {
const payload: Record = { repo_id: repoId };
if (variant) payload.variant = variant;
- const response = await authFetch("/api/models/delete-cached", {
- method: "DELETE",
+ const response = await authFetch("/api/models/reveal-cached-model", {
+ method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 217eaf8b6d..ef018445e0 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -2,16 +2,19 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
+ applyModelLoadConfigToRuntime,
+ currentRuntimePerModelConfig,
type DeletedModelRef,
type ExternalModelOption,
type LoraModelOption,
type ModelOption,
ModelSelector,
-} from "@/components/assistant-ui/model-selector";
-import {
- loadRememberedLoadSettings,
- rememberedLoadSettingsKey,
-} from "@/components/assistant-ui/model-selector/remembered-load-settings";
+ type ModelSelectorChangeMeta,
+ type PerModelConfig,
+ resolveInitialConfig,
+ SidebarModelConfig,
+ useActiveModelConfig,
+} from "@/features/model-picker";
import { ProjectComposer, Thread } from "@/components/assistant-ui/thread";
import { CopyableErrorChip } from "@/components/ui/copyable-error-chip";
import {
@@ -27,10 +30,10 @@ import {
} from "@/components/ui/resizable";
import { useSidebar } from "@/components/ui/sidebar";
import { Tooltip, TooltipContent } from "@/components/ui/tooltip";
-import { useLatestRef } from "@/features/hub/hooks/use-latest-ref";
import {
DOWNLOAD_KIND,
downloadManager,
+ useRepoDownload,
} from "@/features/hub/download-manager";
import {
type NativeIntent,
@@ -93,7 +96,6 @@ import {
renameChatItem,
useChatSidebarItems,
} from "./hooks/use-chat-sidebar-items";
-import { useStagedModelPreparation } from "./hooks/use-staged-model-preparation";
import {
clearTrainingCompareHandoff,
getTrainingCompareHandoff,
@@ -128,10 +130,8 @@ import {
hasGgufSource,
isDownloadableHubRepo,
loadOptionalBool,
- pendingSelectionMatches,
useChatRuntimeStore,
} from "./stores/chat-runtime-store";
-import type { PendingModelSelection } from "./stores/chat-runtime-store";
import { useChatPreferencesStore } from "./stores/chat-preferences-store";
import { useExternalProvidersStore } from "./stores/external-providers-store";
import { buildChatTourSteps } from "./tour";
@@ -385,6 +385,7 @@ type CompareModelSelection = {
id: string;
isLora: boolean;
ggufVariant?: string;
+ config?: PerModelConfig;
};
function modelMatchesDeleted(
@@ -645,6 +646,8 @@ function GeneralCompareHeader({
loraModels,
externalModels,
value,
+ selectedConfig,
+ selectedGgufVariant,
onValueChange,
onFoldersChange,
onModelsChange,
@@ -655,9 +658,11 @@ function GeneralCompareHeader({
loraModels: LoraModelOption[];
externalModels: ExternalModelOption[];
value: string;
+ selectedConfig?: PerModelConfig | null;
+ selectedGgufVariant?: string | null;
onValueChange: (
id: string,
- meta: { isLora: boolean; ggufVariant?: string },
+ meta: ModelSelectorChangeMeta,
) => void;
onFoldersChange?: () => void;
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
@@ -684,6 +689,8 @@ function GeneralCompareHeader({
loraModels={loraModels}
externalModels={externalModels}
value={value}
+ selectedConfig={selectedConfig}
+ selectedGgufVariant={selectedGgufVariant}
onValueChange={onValueChange}
onFoldersChange={onFoldersChange}
onModelsChange={onModelsChange}
@@ -811,11 +818,14 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
loraModels={loraModels}
externalModels={externalModels}
value={model1.id}
+ selectedConfig={model1.config}
+ selectedGgufVariant={model1.ggufVariant}
onValueChange={(id, meta) =>
setModel1({
id,
isLora: meta.isLora,
ggufVariant: meta.ggufVariant,
+ config: meta.config,
})
}
onFoldersChange={onFoldersChange}
@@ -838,11 +848,14 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
loraModels={loraModels}
externalModels={externalModels}
value={model2.id}
+ selectedConfig={model2.config}
+ selectedGgufVariant={model2.ggufVariant}
onValueChange={(id, meta) =>
setModel2({
id,
isLora: meta.isLora,
ggufVariant: meta.ggufVariant,
+ config: meta.config,
})
}
onFoldersChange={onFoldersChange}
@@ -1236,6 +1249,13 @@ export function validateChatSearch(search: Record): ChatSearch
};
}
+type PendingHubAutoLoad = {
+ selection: SelectedModelInput;
+ contextKey: string;
+ originCheckpoint: string;
+ originGgufVariant: string | null;
+};
+
// `search` comes from RootLayout (not useSearch) so ChatPage stays mounted off-route
// (keeping an in-flight generation alive), frozen to the last /chat search. `active`
// is false off-route: close body-portaled surfaces and stop route-specific listeners
@@ -1248,30 +1268,6 @@ export function ChatPage({
const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen);
const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen);
- // Deferred-load staging: downloads a staged GGUF (if needed) and reads its
- // header context so the sheet can show the context slider before the load.
- // autoLoad picks instead load the cached file as soon as the download ends;
- // selectModel is defined below, so the load runs through a ref.
- const autoLoadStagedRef = useRef<
- ((pending: PendingModelSelection) => void) | null
- >(null);
- const stagedDownload = useStagedModelPreparation({
- onAutoLoad: (pending) => autoLoadStagedRef.current?.(pending),
- });
- // Abandon a staged pick: the store action cancels its in-flight download and
- // reverts the edited knobs, so nothing lingers after the user walks away.
- const abandonStaged = useCallback(() => {
- useChatRuntimeStore.getState().abandonStagedModel();
- }, []);
- // Detach a staged pick on navigation without cancelling its download: the
- // transfer keeps running in the manager and lands in cache, like Hub.
- const detachStaged = useCallback(() => {
- useChatRuntimeStore.getState().abandonStagedModel({ keepDownload: true });
- }, []);
- // Tracks whether the chat page is still mounted, so a staged-load failure that
- // resolves after the user left chat doesn't resurrect the abandoned pick.
- const mountedRef = useRef(true);
- useEffect(() => () => void (mountedRef.current = false), []);
const incognito = useChatRuntimeStore((s) => s.incognito);
const setIncognito = useChatRuntimeStore((s) => s.setIncognito);
const incognitoLabel = incognito
@@ -1363,6 +1359,9 @@ export function ChatPage({
const ggufContextLength = useChatRuntimeStore(
(state) => state.ggufContextLength,
);
+ const ggufNativeContextLength = useChatRuntimeStore(
+ (state) => state.ggufNativeContextLength,
+ );
const contextUsage = useChatRuntimeStore((state) => state.contextUsage);
const modelsFromStore = useChatRuntimeStore((state) => state.models);
const lorasFromStore = useChatRuntimeStore((state) => state.loras);
@@ -1440,39 +1439,37 @@ export function ChatPage({
refreshRef.current = refresh;
selectModelRef.current = selectModel;
}, [refresh, selectModel]);
- // Load a cached autoLoad pick once its download finishes. The sheet was never
- // opened, so on a load failure just drop the orphaned staged knobs. The knobs
- // were already seeded on stage, so keepSpeculative only when a config was
- // saved -- otherwise the standing speculative preference should win.
- autoLoadStagedRef.current = (pending) => {
- // Blobs are saved for GGUF picks only (the sheet gates on it), so don't
- // let a legacy non-GGUF blob claim a seeded config here.
- const remembered = hasGgufSource(pending)
- ? loadRememberedLoadSettings(rememberedLoadSettingsKey(pending))
- : null;
- void selectModel({
- ...pending,
- isDownloaded: true,
- forceReload: true,
- keepSpeculative: remembered != null,
- throwOnError: true,
- }).catch(() => {
- const store = useChatRuntimeStore.getState();
- // selectModel only clears pendingSelection on success, so a failed
- // auto-load leaves our staged pick (and its edited load knobs) behind.
- // Abandon it when it is still the active stage; otherwise just revert the
- // settings if the stage was already cleared by something else.
- if (pendingSelectionMatches(store.pendingSelection, pending)) {
- store.abandonStagedModel();
- } else if (!store.pendingSelection) {
- store.resetModelSettingsToLoaded();
- }
- });
- };
+ const rememberedConfigFor = useCallback(
+ (selection: {
+ id: string;
+ ggufVariant?: string | null;
+ source?: string;
+ }) => {
+ if (selection.source === "external") return null;
+ const resolved = resolveInitialConfig(selection.id, selection.ggufVariant);
+ return resolved.remembered ? resolved.config : null;
+ },
+ [],
+ );
const isExternalModel = useMemo(
() => isExternalModelId(inferenceParams.checkpoint),
[inferenceParams.checkpoint],
);
+ const {
+ checkpoint: runtimeCheckpoint,
+ isGguf: runtimeModelIsGguf,
+ config: activeModelConfig,
+ } = useActiveModelConfig();
+ const activeModelIsGguf =
+ runtimeCheckpoint != null && !isExternalModel && runtimeModelIsGguf;
+ const activeModelIsLora = useMemo(() => {
+ const checkpoint = inferenceParams.checkpoint;
+ if (!checkpoint || isExternalModel) return false;
+ const model = modelsFromStore.find((entry) => entry.id === checkpoint);
+ if (model) return model.isLora;
+ const lora = lorasFromStore.find((entry) => entry.id === checkpoint);
+ return lora?.exportType === "lora";
+ }, [inferenceParams.checkpoint, isExternalModel, modelsFromStore, lorasFromStore]);
const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled);
const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle);
const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort);
@@ -1783,75 +1780,21 @@ export function ChatPage({
closeArtifactSurface();
}, [activeThreadId, closeArtifactSurface, selectedArtifact, view]);
- // Abandon a staged (not-yet-loaded) pick when the chat context actually
- // changes — switching threads, leaving single view, or starting a new chat /
- // project — so a stale Load button can't resurface in a different context.
- // New Chat keeps activeThreadId null and only bumps the `new` search nonce, so
- // the key includes the route identity, not just the thread. Mirrors the
- // incognito reset pattern. (Route exit is handled in __root.tsx, which runs
- // after this unmounts.) Clear only on a real change, never on mount: staging
- // from the Hub sets pendingSelection then navigates here, and clearing on
- // mount would wipe it. Comparing the previous context (rather than a first-run
- // flag) is also safe under StrictMode's double-invoke and component remounts.
- const chatContextKey = `${view.mode}|${activeThreadId ?? ""}|${search.new ?? ""}|${search.project ?? ""}`;
- const chatContextKeyRef = useLatestRef(chatContextKey);
- const prevChatContextRef = useRef(null);
- useEffect(() => {
- const prev = prevChatContextRef.current;
- prevChatContextRef.current = chatContextKey;
- if (prev === null || prev === chatContextKey) return;
- detachStaged();
- }, [chatContextKey, detachStaged]);
-
const hasActiveModel = Boolean(inferenceParams.checkpoint);
- // Load immediately, or — when "Load on selection" is off — stage the pick so
- // its load options can be set first. Shared by the main selector, native
- // drag-drop/picker, and the dropped-file chip (the Hub stages via the store).
+ const chatContextKey = `${view.mode}|${activeThreadId ?? ""}|${search.new ?? ""}|${search.project ?? ""}`;
+ const [pendingHubAutoLoad, setPendingHubAutoLoad] =
+ useState(null);
const stageOrLoad = useCallback(
async (selection: SelectedModelInput) => {
const store = useChatRuntimeStore.getState();
- // An un-cached HF repo (GGUF variant or a full non-GGUF snapshot) downloads
- // through the manager first (global indicator), then auto-loads. Everything
- // else -- cached picks, local/native files, LoRA, external -- loads now.
const wantManagerDownload =
isDownloadableHubRepo(selection) && !selection.isDownloaded;
- if (
- (!hasGgufSource(selection) && !wantManagerDownload) ||
- (store.loadOnSelection && selection.isDownloaded)
- ) {
- // Detach any staged pick first so its edited knobs (e.g. a custom
- // context length) don't leak into this immediate load -- resolveLoad
- // reads customContextLength before checking the target is GGUF. Detach
- // (not abandon) keeps its download running.
- detachStaged();
- // Load-on-selection skips the sheet, so seed the saved knobs here the
- // way the sheet's restore effect would; the switch would otherwise reset
- // the remembered speculative choice (keepSpeculative below prevents it).
- const remembered = hasGgufSource(selection)
- ? loadRememberedLoadSettings(rememberedLoadSettingsKey(selection))
- : null;
- if (remembered) store.applyRememberedLoadSettings(remembered);
- await selectModel(
- remembered ? { ...selection, keepSpeculative: true } : selection,
- );
- return;
- }
- // Loads can't queue behind each other, but a download is independent: if
- // the pick needs downloading, start it in the manager so it runs alongside
- // the load. Nothing to download (already on device) just waits.
if (store.modelLoading) {
- // Both an uncached non-GGUF snapshot (wantManagerDownload) and an
- // uncached remote GGUF quant download through the manager, so either can
- // run in the background while another model loads. wantManagerDownload
- // excludes GGUF by design, so the GGUF case is checked separately.
const wantBackgroundDownload =
wantManagerDownload ||
(selection.source === "hub" &&
hasGgufSource(selection) &&
!selection.isDownloaded);
- // The model currently loading already downloads as part of its own load
- // (the /load flow fetches before setting the checkpoint), so re-picking
- // it must not kick off a second transfer against the same cache.
const isLoadingThisPick =
!!loadingModel &&
normalizeModelRef(loadingModel.id) ===
@@ -1862,11 +1805,6 @@ export function ChatPage({
description: "It's downloading as part of the load in progress.",
});
} else if (wantBackgroundDownload) {
- // Only claim the download started once a job is actually created. A
- // transport conflict records state that is only resolvable from the
- // Hub download card, so point the user there instead of showing a
- // success toast for a transfer that never began; "busy" and "error"
- // already surface their own toasts.
const outcome = await downloadManager.requestStart({
kind: DOWNLOAD_KIND.MODEL,
repoId: selection.id,
@@ -1883,6 +1821,11 @@ export function ChatPage({
description:
"An earlier partial download used a different transport. Open the Hub tab to resume or restart it.",
});
+ } else if (outcome === "busy") {
+ toast.info("Download already in progress", {
+ description:
+ "Another download for this model is still running. Reselect it once that finishes to load it.",
+ });
}
} else {
toast.info("Another model is already loading", {
@@ -1891,23 +1834,128 @@ export function ChatPage({
}
return;
}
- // Detach the prior staged pick (keeping its download) before rebinding, so
- // a second pick downloads alongside the first instead of cancelling it.
- detachStaged();
- store.stageModel({
- id: selection.id,
- isLora: selection.isLora,
- ggufVariant: selection.ggufVariant,
- isDownloaded: selection.isDownloaded,
- expectedBytes: selection.expectedBytes,
- nativePathToken: selection.nativePathToken,
- isGguf: selection.isGguf,
- isHubRepo: wantManagerDownload || undefined,
- autoLoad: store.loadOnSelection,
+ const wantManagerStage =
+ wantManagerDownload ||
+ (selection.source === "hub" &&
+ hasGgufSource(selection) &&
+ !selection.isDownloaded);
+ if (wantManagerStage) {
+ setPendingHubAutoLoad((current) =>
+ current &&
+ current.selection.id === selection.id &&
+ (current.selection.ggufVariant ?? null) ===
+ (selection.ggufVariant ?? null) &&
+ current.contextKey === chatContextKey &&
+ current.originCheckpoint === store.params.checkpoint &&
+ current.originGgufVariant === store.activeGgufVariant
+ ? current
+ : {
+ selection,
+ contextKey: chatContextKey,
+ originCheckpoint: store.params.checkpoint,
+ originGgufVariant: store.activeGgufVariant,
+ },
+ );
+ return;
+ }
+ setPendingHubAutoLoad(null);
+ const previousConfig = currentRuntimePerModelConfig({
+ includeMaxSeqLength: true,
+ });
+ const hasAppliedConfig = applyModelLoadConfigToRuntime(
+ selection.config ?? rememberedConfigFor(selection),
+ );
+ await selectModel({
+ ...selection,
+ ...(hasAppliedConfig ? { keepSpeculative: true } : {}),
+ previousConfig,
});
},
- [detachStaged, selectModel, loadingModel],
+ [selectModel, loadingModel, rememberedConfigFor, chatContextKey],
);
+ useRepoDownload({
+ kind: DOWNLOAD_KIND.MODEL,
+ repoId: pendingHubAutoLoad?.selection.id ?? "__hub_autoload_idle__",
+ activeVariant: pendingHubAutoLoad?.selection.ggufVariant ?? null,
+ onComplete: (variant) => {
+ const pending = pendingHubAutoLoad;
+ if (
+ !pending ||
+ (pending.selection.ggufVariant ?? null) !== (variant ?? null)
+ ) {
+ return;
+ }
+ setPendingHubAutoLoad(null);
+ const store = useChatRuntimeStore.getState();
+ if (
+ !active ||
+ pending.contextKey !== chatContextKey ||
+ normalizeModelRef(pending.originCheckpoint) !==
+ normalizeModelRef(store.params.checkpoint) ||
+ pending.originGgufVariant !== store.activeGgufVariant
+ ) {
+ return;
+ }
+ void stageOrLoad({ ...pending.selection, isDownloaded: true });
+ },
+ onError: (variant) => {
+ if (
+ pendingHubAutoLoad &&
+ (pendingHubAutoLoad.selection.ggufVariant ?? null) === (variant ?? null)
+ ) {
+ setPendingHubAutoLoad(null);
+ }
+ },
+ onCancelled: (variant) => {
+ if (
+ pendingHubAutoLoad &&
+ (pendingHubAutoLoad.selection.ggufVariant ?? null) === (variant ?? null)
+ ) {
+ setPendingHubAutoLoad(null);
+ }
+ },
+ });
+ useEffect(() => {
+ const pending = pendingHubAutoLoad;
+ if (!pending) return;
+ let active = true;
+ void (async () => {
+ const outcome = await downloadManager.requestStart({
+ kind: DOWNLOAD_KIND.MODEL,
+ repoId: pending.selection.id,
+ variant: pending.selection.ggufVariant ?? null,
+ expectedBytes: pending.selection.expectedBytes ?? 0,
+ });
+ if (!active) return;
+ if (outcome === "started") {
+ toast.info("Downloading model", {
+ description: "It'll load automatically once the download finishes.",
+ });
+ return;
+ }
+ if (outcome === "conflict") {
+ // Keep pendingHubAutoLoad bound so this surface's cleanup does not wipe
+ // the conflict just recorded by requestStart (which the toast points the
+ // user to); resolving it from the Hub completes the download and this
+ // surface's onComplete auto-loads, mirroring the "started" branch.
+ toast.info("Resume this download from the Hub", {
+ description:
+ "An earlier partial download used a different transport. Open the Hub tab to resume or restart it.",
+ });
+ return;
+ }
+ if (outcome === "busy") {
+ toast.info("Download already in progress", {
+ description:
+ "Another download for this model is still running. Reselect it once that finishes to load it.",
+ });
+ }
+ setPendingHubAutoLoad((current) => (current === pending ? null : current));
+ })();
+ return () => {
+ active = false;
+ };
+ }, [pendingHubAutoLoad]);
const loadNativeModelIntent = useCallback(
async (intent: NativeIntent, loadingDescription: string) => {
const label =
@@ -1915,6 +1963,7 @@ export function ChatPage({
await stageOrLoad({
id: label,
nativePathToken: intent.path.token,
+ nativePathExpiresAtMs: intent.path.expiresAtMs ?? null,
isDownloaded: true,
loadingDescription,
forceReload: true,
@@ -1965,28 +2014,20 @@ export function ChatPage({
const handleCheckpointChange = useCallback(
(
value: string,
- meta?: {
- source?: string;
- isLora: boolean;
- ggufVariant?: string;
- isDownloaded?: boolean;
- expectedBytes?: number;
- isGguf?: boolean;
- },
+ meta?: ModelSelectorChangeMeta,
) => {
const store = useChatRuntimeStore.getState();
const currentCheckpoint = store.params.checkpoint;
const currentVariant = store.activeGgufVariant;
- if (
- !value ||
- (value === currentCheckpoint &&
- (meta?.ggufVariant ?? null) === (currentVariant ?? null))
- )
+ if (!value) return;
+ setPendingHubAutoLoad(null);
+ const isSameLoadedModel =
+ value === currentCheckpoint &&
+ (meta?.ggufVariant ?? null) === (currentVariant ?? null);
+ if (isSameLoadedModel && !meta?.forceReload) {
return;
+ }
if (meta?.source === "external" || isExternalModelId(value)) {
- // Switching to an external model abandons any staged local pick: cancel
- // its download too (setCheckpoint below only clears the pending + knobs).
- abandonStaged();
const selectedExternal = parseExternalModelId(value);
const selectedProvider = selectedExternal
? externalProvidersForChat.find(
@@ -2087,6 +2128,7 @@ export function ChatPage({
ggufMaxContextLength: null,
ggufNativeContextLength: null,
activeNativePathToken: null,
+ activeNativePathExpiresAtMs: null,
// Clear previous-model counters, else the relaxed external-provider
// render gate shows stale stats until the next completion.
contextUsage: null,
@@ -2158,19 +2200,18 @@ export function ChatPage({
source: meta?.source,
isLora: meta?.isLora,
ggufVariant: meta?.ggufVariant,
- isDownloaded: meta?.isDownloaded,
+ isDownloaded: meta?.isDownloaded || isSameLoadedModel,
expectedBytes: meta?.expectedBytes,
isGguf: meta?.isGguf,
+ config: meta?.config,
+ nativePathToken: meta?.nativePathToken,
+ nativePathExpiresAtMs: meta?.nativePathExpiresAtMs,
+ forceReload: isSameLoadedModel || undefined,
};
- // "Load on selection" off: stage the model and open settings so its
- // load knobs (tensor parallel, context length…) can be set, then it
- // loads once via the sheet's Load button. The currently loaded model
- // stays put until the user commits.
await stageOrLoad(selection);
})();
},
[
- abandonStaged,
activeThreadId,
externalProvidersForChat,
modelsFromStore,
@@ -2178,6 +2219,45 @@ export function ChatPage({
view,
],
);
+ const handleReloadActiveModel = useCallback(
+ (config: PerModelConfig) => {
+ const checkpoint = inferenceParams.checkpoint;
+ if (!checkpoint) return;
+ const runtime = useChatRuntimeStore.getState();
+ const nativeToken = runtime.activeNativePathToken;
+ const nativeExpiry = runtime.activeNativePathExpiresAtMs;
+ // A file-picked GGUF is reachable only via its native path token, which
+ // the desktop host prunes after a TTL. Reusing an expired token makes the
+ // reload fail with an opaque error, so prompt the user to re-select the
+ // file instead.
+ if (nativeToken && nativeExpiry != null && Date.now() >= nativeExpiry) {
+ toast.error("This local model file's access has expired.", {
+ description: "Re-select the model file to reload it.",
+ });
+ return;
+ }
+ handleCheckpointChange(checkpoint, {
+ source: "local",
+ isLora: activeModelIsLora,
+ ggufVariant: activeGgufVariant ?? undefined,
+ // Without the native token the reload validates the display label as a
+ // repo and fails.
+ nativePathToken: nativeToken ?? undefined,
+ nativePathExpiresAtMs: nativeExpiry,
+ isGguf: activeModelIsGguf,
+ isDownloaded: true,
+ config,
+ forceReload: true,
+ });
+ },
+ [
+ inferenceParams.checkpoint,
+ activeGgufVariant,
+ activeModelIsLora,
+ activeModelIsGguf,
+ handleCheckpointChange,
+ ],
+ );
const handleEject = useCallback(() => {
void (async () => {
if (await ejectModel()) {
@@ -2446,12 +2526,27 @@ export function ChatPage({
const state = useChatRuntimeStore.getState();
const targetLora = pickBestLoraForBase(state.loras, handoff.baseModel);
+ const selectWithConfig = async (
+ selection: Pick,
+ ) => {
+ const previousConfig = currentRuntimePerModelConfig({
+ includeMaxSeqLength: true,
+ });
+ const hasAppliedConfig = applyModelLoadConfigToRuntime(
+ rememberedConfigFor(selection),
+ );
+ await selectModelRef.current({
+ ...selection,
+ ...(hasAppliedConfig ? { keepSpeculative: true } : {}),
+ previousConfig,
+ });
+ };
if (targetLora) {
console.info("[chat-handoff] loading lora", {
id: targetLora.id,
baseModel: targetLora.baseModel,
});
- await selectModelRef.current({ id: targetLora.id, isLora: true });
+ await selectWithConfig({ id: targetLora.id, isLora: true });
if (canceled) return;
useChatRuntimeStore.getState().setActiveThreadId(null);
useChatRuntimeStore.getState().setContextUsage(null);
@@ -2468,10 +2563,7 @@ export function ChatPage({
console.info("[chat-handoff] no lora match, loading base", {
id: handoff.baseModel,
});
- await selectModelRef.current({
- id: handoff.baseModel,
- isLora: false,
- });
+ await selectWithConfig({ id: handoff.baseModel, isLora: false });
if (canceled) return;
} else {
console.warn("[chat-handoff] no lora/base match found", {
@@ -2491,7 +2583,7 @@ export function ChatPage({
return () => {
canceled = true;
};
- }, [active, navigate]);
+ }, [active, navigate, rememberedConfigFor]);
const tourSteps = useMemo(
() =>
@@ -2580,6 +2672,8 @@ export function ChatPage({
externalModels={externalModels}
value={inferenceParams.checkpoint}
activeGgufVariant={activeGgufVariant}
+ activeModelConfig={activeModelConfig}
+ activeGgufContextLength={ggufContextLength}
onValueChange={handleCheckpointChange}
onEject={handleEject}
onFoldersChange={refreshLocalModels}
@@ -2633,7 +2727,12 @@ export function ChatPage({
stageOrLoad(selection)}
+ onLoad={() =>
+ loadNativeModelIntent(
+ pendingNativeModelIntent,
+ "Loading selected local GGUF model.",
+ )
+ }
/>
) : null}
{loadingModel && loadToastDismissed ? (
@@ -2790,13 +2889,22 @@ export function ChatPage({
open={active && settingsOpen}
onOpenChange={(open) => {
setSettingsOpen(open);
- // Closing the sheet abandons a staged (not-yet-loaded) pick: cancel its
- // download and revert the staged knobs so nothing lingers as a dirty
- // edit (or a background download) on the loaded model.
- if (!open) abandonStaged();
}}
params={inferenceParams}
onParamsChange={setInferenceParams}
+ modelConfig={
+ view.mode !== "compare" && activeModelConfig && !modelLoading ? (
+
+ ) : null
+ }
isExternalModel={isExternalModel}
providerCapabilities={activeProviderCapabilities}
activeExternalProvider={activeExternalProvider}
@@ -2808,67 +2916,6 @@ export function ChatPage({
);
}}
externalProviderType={activeExternalProviderType}
- loadingModel={loadingModel}
- onReloadModel={() => {
- const state = useChatRuntimeStore.getState();
- if (state.params.checkpoint) {
- selectModel({
- id: state.params.checkpoint,
- ggufVariant: state.activeGgufVariant ?? undefined,
- // A native (drag-drop / picked) GGUF's checkpoint is only a display
- // label, so the reload needs its path token to re-mint a lease --
- // else applying the now-exposed GPU/context controls can't resolve
- // the file. Null for non-native loads, which reload by id as before.
- nativePathToken: state.activeNativePathToken ?? undefined,
- forceReload: true,
- isDownloaded: true,
- loadingDescription: "Reloading with updated chat template.",
- });
- }
- }}
- onLoadPendingModel={() => {
- const pending = useChatRuntimeStore.getState().pendingSelection;
- if (!pending) return;
- const keyAtLoad = chatContextKey;
- // forceReload: the staged model isn't loaded yet, so bypass the
- // same-checkpoint dedupe. keepSpeculative: honor the speculative mode
- // set on the sidebar.
- void selectModel({
- ...pending,
- forceReload: true,
- keepSpeculative: true,
- throwOnError: true,
- }).catch(() => {
- // Recoverable failure (expired token, gated repo, OOM…): the pick is
- // cleared only on success, so it normally stays staged with edited
- // knobs intact — nothing to restore.
- const store = useChatRuntimeStore.getState();
- // Still staged (this pick, or a newer one queued meanwhile): leave it.
- if (store.pendingSelection) return;
- // Cleared mid-load (sheet closed / switched chats). Re-stage only if
- // the staged-load is still wanted: same chat context, sheet still
- // open, page still mounted.
- const stillWanted =
- mountedRef.current &&
- store.settingsPanelOpen &&
- chatContextKeyRef.current === keyAtLoad;
- if (stillWanted) {
- store.setPendingSelection(pending);
- } else {
- // Abandoned (closed the sheet / switched chats / left chat): drop
- // the orphaned staged knob edits so they don't linger as dirty
- // settings over the loaded model.
- store.resetModelSettingsToLoaded();
- }
- });
- }}
- stagedDownloadFraction={stagedDownload.progress?.fraction ?? null}
- onCancelStagedDownload={() =>
- stagedDownload.cancelDownload(
- useChatRuntimeStore.getState().pendingSelection?.ggufVariant ??
- null,
- )
- }
/>
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index bd22cc4f55..d4f154882c 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -1,19 +1,7 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import {
- Alert,
- AlertDescription,
- AlertTitle,
-} from "@/components/ui/alert";
import { Button } from "@/components/ui/button";
-import { Checkbox } from "@/components/ui/checkbox";
-import {
- clearRememberedLoadSettings,
- loadRememberedLoadSettings,
- rememberedLoadSettingsKey,
- saveRememberedLoadSettings,
-} from "@/components/assistant-ui/model-selector/remembered-load-settings";
import {
Dialog,
DialogContent,
@@ -29,7 +17,7 @@ import {
DropdownMenuSeparator,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
-import { Input } from "@/components/ui/input";
+import { InfoHint } from "@/components/ui/info-hint";
import {
InputGroup,
InputGroupAddon,
@@ -50,27 +38,22 @@ import {
SheetTitle,
} from "@/components/ui/sheet";
import { Slider } from "@/components/ui/slider";
-import { Spinner } from "@/components/ui/spinner";
import { Switch } from "@/components/ui/switch";
import { Textarea } from "@/components/ui/textarea";
-import { InfoHint } from "@/components/ui/info-hint";
import { Tooltip, TooltipContent } from "@/components/ui/tooltip";
-import { useGpuDevices } from "@/hooks/use-gpu-info";
-import { useIsMobile } from "@/hooks/use-mobile";
+import { NumericValueInput, snapToStep } from "@/features/model-picker";
+import { RetrievalSettingsSection } from "@/features/rag";
import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check";
-import { cn } from "@/lib/utils";
-import {
- ArrowTurnBackwardIcon,
- Edit03Icon,
- LayoutAlignRightIcon,
-} from "@hugeicons/core-free-icons";
+import { useIsMobile } from "@/hooks/use-mobile";
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
+import { toast } from "@/lib/toast";
+import { cn } from "@/lib/utils";
+import { Edit03Icon, LayoutAlignRightIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Braces, ChevronDown, ExternalLink } from "lucide-react";
import { Tooltip as TooltipPrimitive } from "radix-ui";
import { Fragment, type ReactNode } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
-import { toast } from "@/lib/toast";
import { OpenAICodeExecSection } from "./components/openai-code-exec-section";
import { PermissionModeDropdown } from "./permission-mode-select";
import { resyncInferenceStatusAfterServerModelChange } from "./hooks/use-chat-model-runtime";
@@ -78,8 +61,8 @@ import {
type ExternalProviderConfig,
getExternalProviderApiKey,
parseExternalModelId,
- supportsProviderPromptCaching,
supportsProviderPromptCacheTtl,
+ supportsProviderPromptCaching,
} from "./external-providers";
import {
BUILTIN_PRESETS,
@@ -99,15 +82,7 @@ import {
providerSupportsBuiltinCodeExecution,
providerSupportsFastMode,
} from "./provider-capabilities";
-import {
- GPU_LAYERS_AUTO,
- distributeByWeight,
- isPendingGguf,
- pendingSelectionMatches,
- rebalanceSplit,
- useChatRuntimeStore,
-} from "./stores/chat-runtime-store";
-import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section";
+import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import type { InferenceParams } from "./types/runtime";
export { defaultInferenceParams, type Preset } from "./presets/preset-policy";
@@ -130,7 +105,7 @@ function getPromptVariablesError(raw: string): string | null {
return null;
}
} catch {
- return "Use valid JSON, for example { \"env\": \"staging\" }.";
+ return 'Use valid JSON, for example { "env": "staging" }.';
}
return "Variables must be a JSON object.";
}
@@ -139,112 +114,7 @@ function hasPromptVariableSyntax(prompt: string): boolean {
return PROMPT_VARIABLE_PATTERN.test(prompt);
}
-/**
- * Editable numeric value display, shared by every slider value and the Context
- * Length input. An that looks like text (shows `displayValue ?? value`,
- * so "Off"/"Max" labels render) until focus, when it swaps to the raw number,
- * selects it, and accepts free text. Commits on blur/Enter, reverts on Escape.
- * Clamping happens on commit so typing intermediate values isn't fought.
- */
-function snapToStep(
- value: number,
- step: number,
- min?: number,
- max?: number,
-): number {
- const lo = min ?? Number.NEGATIVE_INFINITY;
- const hi = max ?? Number.POSITIVE_INFINITY;
- const clamped = Math.min(Math.max(value, lo), hi);
- const stepStr = String(step);
- const decimals = stepStr.includes(".") ? stepStr.split(".")[1].length : 0;
- const base = Number.isFinite(lo) ? lo : 0;
- const snapped = base + Math.round((clamped - base) / step) * step;
- const reclamped = Math.min(Math.max(snapped, lo), hi);
- return Number(reclamped.toFixed(decimals));
-}
-
-function NumericValueInput({
- value,
- min,
- max,
- step,
- onChange,
- displayValue,
- className,
- ariaLabel,
- size: sizeAttr,
- disabled = false,
-}: {
- value: number;
- min?: number;
- max?: number;
- step: number;
- onChange: (v: number) => void;
- displayValue?: string;
- className?: string;
- ariaLabel?: string;
- size?: number;
- disabled?: boolean;
-}) {
- const [focused, setFocused] = useState(false);
- const [draft, setDraft] = useState("");
- const cancelBlurCommitRef = useRef(false);
-
- const commit = (raw: string) => {
- const parsed = Number.parseFloat(raw);
- if (!Number.isFinite(parsed)) {
- return;
- }
- const final = snapToStep(parsed, step, min, max);
- if (final !== value) {
- onChange(final);
- }
- };
-
- const displayed = focused ? draft : (displayValue ?? String(value));
-
- return (
- {
- cancelBlurCommitRef.current = false;
- setDraft(String(value));
- setFocused(true);
- // Defer select() so it runs after the value swap above.
- const target = e.currentTarget;
- requestAnimationFrame(() => target.select());
- }}
- onBlur={() => {
- if (cancelBlurCommitRef.current) {
- cancelBlurCommitRef.current = false;
- } else {
- commit(draft);
- }
- setFocused(false);
- }}
- onChange={(e) => setDraft(e.target.value)}
- onKeyDown={(e) => {
- if (e.key === "Enter") {
- e.currentTarget.blur();
- } else if (e.key === "Escape") {
- cancelBlurCommitRef.current = true;
- setDraft(String(value));
- e.currentTarget.blur();
- }
- }}
- className={cn("panel-number-input", className)}
- />
- );
-}
-
-function ParamSlider({
+export function ParamSlider({
label,
value,
min,
@@ -285,6 +155,7 @@ function ParamSlider({
displayValue={displayValue}
ariaLabel={label}
size={valueSize ?? 4}
+ className="panel-number-input"
disabled={disabled}
/>
@@ -385,8 +256,7 @@ function CollapsibleSection({
return (
{labelHref ? (
@@ -458,6 +328,7 @@ interface ChatSettingsPanelProps {
onOpenChange?: (open: boolean) => void;
params: InferenceParams;
onParamsChange: (params: InferenceParams) => void;
+ modelConfig?: ReactNode;
isExternalModel?: boolean;
/**
* Sampling-param capabilities for the active external provider, or `null` for
@@ -472,21 +343,6 @@ interface ChatSettingsPanelProps {
* Max Tokens floor in the slider.
*/
externalProviderType?: string | null;
- onReloadModel?: () => void;
- /** The in-flight load (id + GGUF variant + native path token), or null when
- * idle. Used to show a loading state for the staged pick only — not for an
- * unrelated load or a cancel's background unload. */
- loadingModel?: {
- id: string;
- ggufVariant?: string | null;
- nativePathToken?: string | null;
- } | null;
- /** Loads the staged `pendingSelection` (deferred "Load on selection" flow). */
- onLoadPendingModel?: () => void;
- /** Download progress (0–1) for a staged GGUF being fetched, or null when idle. */
- stagedDownloadFraction?: number | null;
- /** Cancels the in-flight staged download (paired with abandoning the stage). */
- onCancelStagedDownload?: () => void;
}
export function ChatSettingsPanel({
@@ -494,16 +350,12 @@ export function ChatSettingsPanel({
onOpenChange,
params,
onParamsChange,
+ modelConfig = null,
isExternalModel = false,
providerCapabilities = null,
activeExternalProvider = null,
onExternalProviderChange,
externalProviderType = null,
- onReloadModel,
- loadingModel = null,
- onLoadPendingModel,
- stagedDownloadFraction,
- onCancelStagedDownload,
}: ChatSettingsPanelProps) {
// Local models show every knob; providerCapabilities is only consulted when
// isExternalModel. Unknown providers fall back to the OpenAI-compat shape via
@@ -518,64 +370,23 @@ export function ChatSettingsPanel({
const showPresencePenalty =
!isExternalModel || Boolean(providerCapabilities?.presencePenalty);
const isMobile = useIsMobile();
- const pendingSelection = useChatRuntimeStore((s) => s.pendingSelection);
- // "Loading" only when the in-flight load IS this staged pick (full id + GGUF
- // variant + native token match), not an unrelated load or a cancel's
- // background unload. The variant matters: a different quant of the same repo
- // staged mid-load must not read as this one loading.
- const stagedLoading =
- loadingModel != null &&
- pendingSelectionMatches(pendingSelection, {
- id: loadingModel.id,
- ggufVariant: loadingModel.ggufVariant,
- nativePathToken: loadingModel.nativePathToken,
- });
- // Load settings are snapshotted at click time; lock them while loading.
- const modelControlsDisabled = stagedLoading;
- const abandonStagedModel = useChatRuntimeStore((s) => s.abandonStagedModel);
- const resetModelSettingsToLoaded = useChatRuntimeStore(
- (s) => s.resetModelSettingsToLoaded,
+ const isLoadedGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
+ const currentCheckpoint = params.checkpoint;
+ const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
+ // Direct-file / custom-folder GGUFs load without a variant label but still
+ // report a GGUF context, so detect them via the context and the checkpoint
+ // suffix too (mirrors the chat page's activeModelIsGguf). Otherwise Max Tokens
+ // would fall back to params.maxSeqLength instead of the loaded GGUF context.
+ const isGguf =
+ isLoadedGguf ||
+ ggufContextLength != null ||
+ (currentCheckpoint?.toLowerCase().endsWith(".gguf") ?? false);
+ const ggufMaxContextLength = useChatRuntimeStore(
+ (s) => s.ggufMaxContextLength,
);
- // A staged GGUF pick (deferred load) shows the GGUF load knobs so they can be
- // set before the single load.
- const pendingIsGguf = isPendingGguf(pendingSelection);
- // Short, human-readable name for the staged pick (HF ids carry an org prefix;
- // native picks are already a display label). Drives the "staged, not loaded"
- // callout so it's obvious the selection hasn't loaded yet.
- const stagedLabel = (() => {
- const id = pendingSelection?.id ?? "";
- const slash = id.lastIndexOf("/");
- const base = slash >= 0 ? id.slice(slash + 1) : id;
- return base || id;
- })();
- const activeNativePathToken = useChatRuntimeStore(
- (s) => s.activeNativePathToken,
- );
- const loadedGgufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
- // A GGUF loaded from a native path / direct .gguf has no HF variant, so key
- // off the same signal the status hydration uses -- variant OR native token OR
- // a GGUF context -- else the GPU Memory controls hide for a loaded local GGUF.
- const isLoadedGguf =
- useChatRuntimeStore((s) => s.activeGgufVariant) != null ||
- activeNativePathToken != null ||
- loadedGgufContextLength != null;
- // While a pick is staged the sheet configures *that* model, so its GGUF-ness
- // (not the currently loaded model's) decides whether the GGUF-only controls
- // show. Otherwise a staged non-GGUF Hub repo would inherit the loaded GGUF's
- // context/KV/speculative controls.
- const isGguf = pendingSelection != null ? pendingIsGguf : isLoadedGguf;
- // The Model section (and Load button) shows for any staged pick, even when the
- // currently active model is external.
- const hasModelContent =
- pendingSelection != null ||
- (!isExternalModel && (isGguf || Boolean(params.checkpoint)));
+ const customContextLength = useChatRuntimeStore((s) => s.customContextLength);
const speculativeType = useChatRuntimeStore((s) => s.speculativeType);
- const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType);
- const loadedSpeculativeType = useChatRuntimeStore(
- (s) => s.loadedSpeculativeType,
- );
const specFallbackReason = useChatRuntimeStore((s) => s.specFallbackReason);
- // Only binary fallback states are solved by a newer prebuilt.
const mtpUpdatable =
specFallbackReason === "binary_no_mtp" ||
specFallbackReason === "binary_outdated";
@@ -597,65 +408,27 @@ export function ChatSettingsPanel({
`llama.cpp updated to ${result.tag ?? "the latest build"}.${reloadHint}`,
);
} else {
- toast.error(`llama.cpp update failed: ${result.error ?? "unknown error"}`);
+ toast.error(
+ `llama.cpp update failed: ${result.error ?? "unknown error"}`,
+ );
}
}, [applyLlamaUpdate]);
- const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax);
- const setSpecDraftNMax = useChatRuntimeStore((s) => s.setSpecDraftNMax);
- const loadedSpecDraftNMax = useChatRuntimeStore(
- (s) => s.loadedSpecDraftNMax,
- );
- const currentCheckpoint = params.checkpoint;
- const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
- const ggufMaxContextLength = useChatRuntimeStore(
- (s) => s.ggufMaxContextLength,
- );
- const ggufNativeContextLength = useChatRuntimeStore(
- (s) => s.ggufNativeContextLength,
- );
- const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype);
- const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype);
- const applyRememberedLoadSettings = useChatRuntimeStore(
- (s) => s.applyRememberedLoadSettings,
- );
- const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype);
- const tensorParallel = useChatRuntimeStore((s) => s.tensorParallel);
- const setTensorParallel = useChatRuntimeStore((s) => s.setTensorParallel);
- const loadedTensorParallel = useChatRuntimeStore(
- (s) => s.loadedTensorParallel,
- );
- const gpuMemoryMode = useChatRuntimeStore((s) => s.gpuMemoryMode);
- const setGpuMemoryMode = useChatRuntimeStore((s) => s.setGpuMemoryMode);
- const loadedGpuMemoryMode = useChatRuntimeStore((s) => s.loadedGpuMemoryMode);
- const loadedIsDiffusion = useChatRuntimeStore((s) => s.loadedIsDiffusion);
- const gpuLayers = useChatRuntimeStore((s) => s.gpuLayers);
- const setGpuLayers = useChatRuntimeStore((s) => s.setGpuLayers);
- const loadedGpuLayers = useChatRuntimeStore((s) => s.loadedGpuLayers);
- const nCpuMoe = useChatRuntimeStore((s) => s.nCpuMoe);
- const setNCpuMoe = useChatRuntimeStore((s) => s.setNCpuMoe);
- const loadedNCpuMoe = useChatRuntimeStore((s) => s.loadedNCpuMoe);
- const splitRatio = useChatRuntimeStore((s) => s.splitRatio);
- const setSplitRatio = useChatRuntimeStore((s) => s.setSplitRatio);
- const loadedSplitRatio = useChatRuntimeStore((s) => s.loadedSplitRatio);
- const ggufLayerCount = useChatRuntimeStore((s) => s.ggufLayerCount);
- const moeLayerCount = useChatRuntimeStore((s) => s.moeLayerCount);
- const selectedGpuIds = useChatRuntimeStore((s) => s.selectedGpuIds);
- const setSelectedGpuIds = useChatRuntimeStore((s) => s.setSelectedGpuIds);
- const loadedGpuIds = useChatRuntimeStore((s) => s.loadedGpuIds);
- const gpuDevices = useGpuDevices();
- const chatTemplateOverride = useChatRuntimeStore(
- (s) => s.chatTemplateOverride,
- );
- const loadedChatTemplateOverride = useChatRuntimeStore(
- (s) => s.loadedChatTemplateOverride,
- );
- const customContextLength = useChatRuntimeStore((s) => s.customContextLength);
- const loadedCustomContextLength = useChatRuntimeStore(
- (s) => s.loadedCustomContextLength,
- );
- const setCustomContextLength = useChatRuntimeStore(
- (s) => s.setCustomContextLength,
- );
+ const loadedEffectiveContext = customContextLength ?? ggufContextLength;
+ const showSpecFallback =
+ !isExternalModel &&
+ isGguf &&
+ specFallbackReason != null &&
+ (speculativeType === "auto" ||
+ speculativeType === "mtp" ||
+ speculativeType === "mtp+ngram");
+ const showContextVramWarning =
+ !isExternalModel &&
+ isGguf &&
+ ggufMaxContextLength != null &&
+ loadedEffectiveContext != null &&
+ loadedEffectiveContext > ggufMaxContextLength;
+ const showLoadedDiagnostics = showSpecFallback || showContextVramWarning;
+ const hasModelContent = showLoadedDiagnostics;
const setActivePresetSource = useChatRuntimeStore(
(s) => s.setActivePresetSource,
);
@@ -666,170 +439,7 @@ export function ChatSettingsPanel({
const setActivePreset = useChatRuntimeStore((s) => s.setActivePreset);
const settingsHydrated = useChatRuntimeStore((s) => s.settingsHydrated);
- // A staged (not-yet-loaded) GGUF carries its own header context length on
- // pendingSelection, so the slider can use the staged model's real ceiling
- // without reading the loaded model's `ggufContextLength`.
- const stagedContextLength = pendingSelection?.contextLength ?? null;
- // "Remember settings next time" tick for a staged model. Seeds the store from
- // the saved per-model settings on stage, so the sheet opens with what was used
- // last time; the tick reflects whether a saved entry exists.
- const [remember, setRemember] = useState(false);
- // Keyed per quant: a different variant of the same repo has its own settings.
- const pendingKey = pendingSelection
- ? rememberedLoadSettingsKey(pendingSelection)
- : null;
- useEffect(() => {
- if (!pendingKey) return;
- // GGUF-only, like the stageOrLoad / Hub restore paths: every remembered
- // field is a llama.cpp knob, so a non-GGUF pick has nothing to restore --
- // and applying its blob would clobber the standing gpuMemoryMode with a
- // stale snapshot (the save on Load below is gated the same way).
- const saved = pendingIsGguf ? loadRememberedLoadSettings(pendingKey) : null;
- setRemember(saved != null);
- if (saved) applyRememberedLoadSettings(saved);
- }, [pendingKey, pendingIsGguf, applyRememberedLoadSettings]);
- // While staging, the sheet reflects the STAGED model, so its header context
- // takes precedence over the loaded model's (which may differ or be larger).
- const baseContext = pendingIsGguf ? stagedContextLength : ggufContextLength;
- const baseNativeContext = pendingIsGguf
- ? stagedContextLength
- : ggufNativeContextLength;
- // Context controls render once we actually have a ceiling: for a staged GGUF,
- // once its header metadata arrives (post-download); otherwise post-load.
- const showContextControl = pendingIsGguf
- ? stagedContextLength != null
- : isLoadedGguf;
- const stagedDownloading =
- stagedDownloadFraction != null && stagedDownloadFraction < 1;
- const ctxDisplayValue = customContextLength ?? baseContext ?? "";
- const ctxMaxValue = baseNativeContext ?? baseContext ?? null;
- const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
- const ctxDirty = customContextLength !== loadedCustomContextLength;
- const specDirty = speculativeType !== loadedSpeculativeType;
- const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax;
- const tpDirty = tensorParallel !== (loadedTensorParallel ?? false);
- // A loaded diffusion GGUF runs mode-agnostic (pins all layers on one GPU,
- // ignores --fit/--gpu-layers), so the GPU Memory mode + manual controls don't
- // apply -- hide them and don't let the preserved standing mode read as dirty.
- // The GPU picker still applies (diffusion pins the chosen device). A staged pick
- // keeps the controls (a pending pick's diffusion-ness isn't known until load).
- const gpuModeApplies =
- isGguf && (pendingSelection != null || !loadedIsDiffusion);
- const gpuDirty =
- gpuModeApplies && gpuMemoryMode !== (loadedGpuMemoryMode ?? "auto");
- const isManual = gpuModeApplies && gpuMemoryMode === "manual";
- // Manual with the GPU Layers slider at "Auto" (leftmost): --fit owns the whole
- // layout, so the offload knobs (MoE, split, TP) don't apply.
- const autoLayers = isManual && gpuLayers < 0;
- // GPUs actually in use: the picked subset, or all visible when none picked.
- const gpusInUse = selectedGpuIds ?? gpuDevices.map((d) => d.index);
- // The picker must keep one GPU selected.
- const singleGpuInUse = gpusInUse.length <= 1;
- // TP needs at least two GPUs because tensor split is a no-op on one and may
- // abort. Auto layers hides TP because --fit aborts under --split-mode tensor.
- const tpDisabled = singleGpuInUse;
- // Manual gpu-layers ceiling = model layer count + 1 (else a safe fallback):
- // llama.cpp counts the output layer as one more offloadable layer past the
- // repeating blocks ("offloaded 33/33" needs -ngl 33 on a 32-block model), so
- // the slider max must reach it or full offload is unreachable. While staging,
- // use the staged model's layer count (read from its header).
- const stagedLayerCount = pendingSelection?.layerCount ?? null;
- const modelLayerCount = pendingIsGguf ? stagedLayerCount : ggufLayerCount;
- const gpuLayersMax = modelLayerCount != null ? modelLayerCount + 1 : 256;
- // MoE-offload slider: shown only for MoE models, capped at their MoE-layer
- // count. While staging, use the staged model's count (read from its header);
- // otherwise the loaded model's.
- const stagedMoeLayerCount = pendingSelection?.moeLayerCount ?? null;
- const moeLayersMax = pendingIsGguf
- ? (stagedMoeLayerCount ?? 0)
- : (moeLayerCount ?? 0);
- const showMoeSlider = isManual && !autoLayers && moeLayersMax > 0;
- // gpuLayers always counts; MoE only with an explicit layer count (see above).
- const manualDirty =
- isManual &&
- (gpuLayers !== loadedGpuLayers ||
- (!autoLayers && nCpuMoe !== (loadedNCpuMoe ?? 0)));
- // GPU picker: only meaningful on multi-GPU, and only when the reported
- // indices are physical (relative ordinals from a parent CUDA_VISIBLE_DEVICES
- // mask can't be mapped back to pin a device). null = use all (auto).
- const showGpuPicker =
- isGguf &&
- gpuDevices.length > 1 &&
- gpuDevices.every((d) => d.physicalIndex);
- const isGpuChecked = (index: number) =>
- selectedGpuIds === null || selectedGpuIds.includes(index);
- const toggleGpu = (index: number) => {
- const all = gpuDevices.map((d) => d.index);
- const current = selectedGpuIds ?? all;
- const next = current.includes(index)
- ? current.filter((i) => i !== index)
- : [...current, index].sort((a, b) => a - b);
- if (next.length === 0) return; // keep at least one GPU selected
- setSelectedGpuIds(next.length === all.length ? null : next);
- // The per-GPU split is positional, so any change to the set of GPUs in use
- // invalidates it: drop it (the sliders fall back to the VRAM-weighted
- // default). TP needs 2+ GPUs, so disable it when only one remains.
- setSplitRatio(null);
- if (next.length <= 1) {
- setTensorParallel(false);
- }
- };
- const gpuIdsKey = (ids: number[] | null) => (ids === null ? "auto" : ids.join(","));
- const gpuIdsDirty = gpuIdsKey(selectedGpuIds) !== gpuIdsKey(loadedGpuIds);
- // Per-GPU layer split (--tensor-split): manual + 2+ GPUs in use. One slider
- // per GPU, each a layer count; together they sum to the GPU Layers total.
- const showSplitRatio =
- isManual && !autoLayers && showGpuPicker && gpusInUse.length > 1;
- // The total the per-GPU counts sum to (the GPU Layers slider value); 0 under
- // Auto, where the split is hidden. The devices behind the GPUs in use, for
- // labels + the VRAM-weighted default.
- const splitTotal = Math.max(0, Math.min(gpuLayers, gpuLayersMax));
- const gpusInUseDevices = gpusInUse.map(
- (i) => gpuDevices.find((d) => d.index === i) ?? null,
- );
- // Displayed per-GPU counts. splitRatio is a stable reference balance (only a
- // slider edit changes it), rescaled to the current total; deriving rather than
- // mutating it on GPU Layers changes keeps the balance intact when the total
- // passes through low values or Auto. No saved split: free-VRAM-weighted default
- // (llama.cpp's unset default splits by free VRAM, so the first edit starts from
- // the default's placement, not a total-VRAM ratio that can land layers on a
- // busy GPU). A genuine 0 (a full GPU) is a real weight, not missing data: the
- // probe's no-data case degrades to the total server-side, and an all-zero list
- // falls back to an even split in distributeByWeight. Not yet sent.
- const splitCounts =
- splitRatio && splitRatio.length === gpusInUse.length
- ? distributeByWeight(splitTotal, splitRatio)
- : distributeByWeight(
- splitTotal,
- gpusInUseDevices.map((d) => d?.memoryFreeGb ?? d?.memoryTotalGb ?? 1),
- );
- const setSplitCount = (k: number, v: number) =>
- setSplitRatio(rebalanceSplit(splitTotal, splitCounts, k, v));
- const splitRatioDirty =
- isManual &&
- !autoLayers &&
- JSON.stringify(splitRatio ?? null) !== JSON.stringify(loadedSplitRatio ?? null);
- // Auto-fit context (Manual + Auto layers): <= 0 means "Auto" (--fit sizes it);
- // a positive value pins it. Surface the length --fit chose once it's loaded.
- const fitCtxAuto = autoLayers && (customContextLength ?? 0) <= 0;
- const loadedAutoLayers =
- loadedGpuMemoryMode === "manual" && (loadedGpuLayers ?? GPU_LAYERS_AUTO) < 0;
- const fitResolvedCtx =
- fitCtxAuto && loadedAutoLayers ? ggufContextLength : null;
- // A saved chat-template override is a reload-time setting too, so surface
- // Apply for a template-only edit (otherwise it could never be applied).
- const templateDirty = chatTemplateOverride !== loadedChatTemplateOverride;
- const modelSettingsDirty =
- kvDirty ||
- ctxDirty ||
- specDirty ||
- specDraftDirty ||
- tpDirty ||
- gpuDirty ||
- manualDirty ||
- gpuIdsDirty ||
- splitRatioDirty ||
- templateDirty;
+ const baseContext = ggufContextLength;
const [presetNameInput, setPresetNameInput] = useState(activePreset);
const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false);
const [systemPromptDraft, setSystemPromptDraft] = useState("");
@@ -855,8 +465,7 @@ export function ChatSettingsPanel({
BUILTIN_PRESETS.find((preset) => preset.name === activePreset) ?? null,
[activePreset],
);
- const hasUnsavedPresetChanges = useMemo(
- () => {
+ const hasUnsavedPresetChanges = useMemo(() => {
if (activePresetDefinition == null) {
return false;
}
@@ -864,9 +473,7 @@ export function ChatSettingsPanel({
return activePresetSource === "modified";
}
return !isSamePresetConfig(activePresetDefinition.params, params);
- },
- [activePresetDefinition, activePresetSource, params],
- );
+ }, [activePresetDefinition, activePresetSource, params]);
const presetSaveState = useMemo(
() =>
getPresetSaveState({
@@ -895,6 +502,14 @@ export function ChatSettingsPanel({
const externalSelection = currentCheckpoint
? parseExternalModelId(currentCheckpoint)
: null;
+ const maxTokensMax = isExternalModel
+ ? getExternalMaxOutputTokens(
+ externalProviderType,
+ externalSelection?.modelId,
+ )
+ : isGguf && baseContext
+ ? baseContext
+ : Math.max(64, params.maxSeqLength);
const showOpenAICodeExecSection =
activeExternalProvider != null &&
providerSupportsBuiltinCodeExecution(
@@ -977,8 +592,7 @@ export function ChatSettingsPanel({
return;
}
const fallbackPreset =
- BUILTIN_PRESETS.find((preset) => preset.name === "Default") ??
- null;
+ BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? null;
const next = customPresets.filter((preset) => preset.name !== name);
setCustomPresets(next);
if (activePreset === name) {
@@ -1090,7 +704,7 @@ export function ChatSettingsPanel({
Run settings
-
+
onOpenChange?.(false)}
@@ -1121,611 +735,57 @@ export function ChatSettingsPanel({
className="run-settings-scroll relative min-h-0 flex-1 overflow-y-auto"
>
- {hasModelContent && (
-
-
- {pendingSelection && (
-
-
- {stagedLoading
- ? `Loading ${stagedLabel}…`
- : `${stagedLabel} is staged, not loaded yet`}
-
-
- {stagedLoading
- ? "Applying your settings."
- : "Set the options below, then choose Load model to load it."}
-
-
- )}
- {isGguf && (
- <>
- {showContextControl && (autoLayers ? (
-
-
-
-
- Context Length
-
-
- Auto: llama.cpp's --fit sizes the context to fit VRAM.
- Set a length to pin it instead -- --fit then optimizes
- GPU layer offload around it. The length --fit chose
- shows here after loading.
-
-
-
{
- setCustomContextLength(v > 0 ? v : null);
- }}
- ariaLabel="Context Length"
- size={8}
- disabled={modelControlsDisabled}
- />
-
-
{
- // Far-left snaps to Auto; otherwise to the nearest 1024.
- if (v < 512) {
- setCustomContextLength(null);
- } else {
- setCustomContextLength(Math.round(v / 1024) * 1024);
- }
- }}
- className="panel-slider"
- disabled={modelControlsDisabled}
- />
- {fitResolvedCtx != null && (
-
- llama.cpp loaded {fitResolvedCtx.toLocaleString()} tokens.
-
- )}
-
- ) : (
-
-
-
- Context Length
-
- {
- setCustomContextLength(
- v === (baseContext ?? 0) ? null : v,
- );
- }}
- ariaLabel="Context Length"
- size={8}
- disabled={modelControlsDisabled}
- />
-
-
{
- const snapped = Math.round(v);
- setCustomContextLength(
- snapped === (baseContext ?? 0) ? null : snapped,
- );
- }}
- className="panel-slider"
- disabled={modelControlsDisabled}
- />
- {ggufMaxContextLength != null &&
- typeof ctxDisplayValue === "number" &&
- ctxDisplayValue > ggufMaxContextLength && (
-
- Exceeds estimated VRAM capacity (
- {ggufMaxContextLength.toLocaleString()} tokens). The
- model may use system RAM.
-
- )}
-
- ))}
-
-
-
- KV Cache Dtype
-
-
- Lower KV cache precision to save VRAM at the cost of some
- quality. f16/bf16 are full precision; q8_0/q5_1/q4_1 are
- quantized.
-
-
-
- {
- setKvCacheDtype(v === "f16" ? null : v);
- }}
- >
-
-
-
-
- f16
- bf16
- q8_0
- q5_1
- q4_1
-
-
-
-
- {isGguf && (
- <>
-
-
-
- Speculative Decoding
-
-
- Faster generation with 0% accuracy hit. Auto picks
- MTP / ngram-mod based on the model and platform.
- Pick MTP, Ngram, or MTP+Ngram to force a specific
- strategy on both GPU and CPU.
-
-
-
- {
- setSpeculativeType(v);
- if (v !== "mtp" && v !== "mtp+ngram") {
- setSpecDraftNMax(null);
- }
- }}
- >
-
-
-
-
- Auto
- MTP
- Ngram
- MTP+Ngram
- Off
-
-
-
-
- {specFallbackReason &&
- (speculativeType === "auto" ||
- speculativeType === "mtp" ||
- speculativeType === "mtp+ngram") && (
-
-
- {specFallbackReason === "mla_mtp_disabled"
- ? "MTP is disabled by default for this model architecture because it currently runs slower than standard decoding. Select MTP above to force it."
- : specFallbackReason === "runtime_error"
- ? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding."
- : specFallbackReason === "drafter_not_found"
- ? "This model supports MTP, but its drafter file could not be downloaded, so MTP is off and it falls back to n-gram speculative decoding where the llama.cpp build supports it. Check your network connection or Hugging Face access, then reload the model to retry the drafter."
- : "MTP is not available in the installed llama.cpp build, so this model is running without it." +
- (llamaUpdateStatus?.update_available
+ {(hasModelContent || modelConfig) && (
+
+
+ {modelConfig}
+ {showSpecFallback && (
+
+
+ {specFallbackReason === "mla_mtp_disabled"
+ ? "MTP is disabled by default for this model architecture because it currently runs slower than standard decoding. Choose MTP in the model picker to force it."
+ : specFallbackReason === "runtime_error"
+ ? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding."
+ : specFallbackReason === "drafter_not_found"
+ ? "This model supports MTP, but its drafter file could not be downloaded, so MTP is off and it falls back to n-gram speculative decoding where the llama.cpp build supports it. Check your network connection or Hugging Face access, then reload the model to retry the drafter."
+ : `MTP is not available in the installed llama.cpp build, so this model is running without it.${
+ llamaUpdateStatus?.update_available
? " Update llama.cpp to enable it."
- : "")}
-
- {mtpUpdatable && llamaUpdateStatus?.update_available && (
-
- {llamaUpdating ? "Updating..." : "Update llama.cpp"}
-
- )}
-
- )}
- {(speculativeType === "mtp" ||
- speculativeType === "mtp+ngram") && (
-
-
-
- Draft Tokens
-
-
- Max MTP draft tokens per step
- (--spec-draft-n-max). Lower = less wasted
- draft decode; higher = bigger speedup when
- acceptance stays high. Default: 2 on GPU,
- 3 on CPU/Mac.
-
-
-
{
- const raw = e.target.value;
- if (raw === "") {
- setSpecDraftNMax(null);
- return;
- }
- const parsed = Number.parseInt(raw, 10);
- if (Number.isFinite(parsed)) {
- const clamped = Math.max(1, Math.min(16, parsed));
- setSpecDraftNMax(clamped);
- }
- }}
- data-test-id="spec-draft-n-max-input"
- aria-label="Speculative decoding draft tokens"
- className="h-7 w-[88px] rounded-full border-border bg-background hover:bg-accent/50 dark:border-transparent dark:bg-white/[0.05] dark:hover:bg-white/[0.1] pl-3 py-0 text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0"
- />
-
- )}
- >
- )}
- {gpuModeApplies && (
-
-
-
- GPU Memory
-
-
-
-
- Default: Unsloth
- fits the model and context to your GPUs.
-
-
- Manual: set GPU
- Layers yourself. Leave it on Auto to let llama.cpp size
- the context and offload overflow (including MoE experts)
- to RAM.
-
-
-
-
-
- {
- setGpuMemoryMode(v as "auto" | "manual");
- }}
- // An in-flight staged load already snapshotted its
- // settings, so edits here could not apply -- disable like
- // the sibling context/KV/spec controls.
- disabled={modelControlsDisabled}
- >
-
-
-
-
- Default
- Manual
-
-
-
-
- )}
- {isManual && (
- <>
-
- Layers to keep on the GPU (--gpu-layers); the rest run
- on CPU. Auto lets llama.cpp size the split (and the
- context) to fit VRAM. At the maximum, the whole model
- is on the GPU.
- >
- }
- />
- {showMoeSlider && (
-
- Keep the experts of this many MoE layers on the CPU
- (--n-cpu-moe) to save VRAM. 0 = all experts on the
- GPU; at the maximum, all are on the CPU.
- >
- }
- />
- )}
- {showSplitRatio && (
-
-
-
- Layers per GPU
-
-
- Splits GPU Layers across GPUs (--tensor-split).
- Without Tensor Parallelism each value is the layer
- count on that GPU; with it, every GPU holds a slice
- of each layer, so the values are only a ratio.
-
-
- {gpusInUseDevices.map((d, k) => (
-
setSplitCount(k, v)}
- valueSize={6}
- disabled={modelControlsDisabled}
- />
- ))}
-
- )}
- >
- )}
- {showGpuPicker && (
-
-
-
- GPUs
-
-
- Which GPUs this model may use. Unchecked GPUs are hidden
- from llama.cpp (CUDA_VISIBLE_DEVICES, or
- HIP_VISIBLE_DEVICES on ROCm). Leave all checked to use
- every GPU. At least one GPU must stay selected.
-
-
-
- {gpuDevices.map((d) => (
-
-
- GPU {d.index}: {d.name}
- {d.memoryTotalGb
- ? ` · ${Math.round(d.memoryTotalGb)} GB`
- : ""}
-
- toggleGpu(d.index)}
- data-test-id={`gpu-pick-${d.index}`}
- disabled={
- modelControlsDisabled ||
- (isGpuChecked(d.index) && singleGpuInUse)
- }
- />
-
- ))}
-
-
- )}
- {gpuModeApplies && !autoLayers && (
-
-
-
- Tensor Parallelism
-
-
- No effect on a single GPU. On multi-GPU setups, improves
- tokens/sec during generation when using dense models. MoE
- models don't benefit and can be much slower.
-
-
-
-
- )}
- >
- )}
- {/* No persistent "enable custom code" toggle: it is consented per model
- via the load-time review dialog. */}
- {/* Apply/Reset belongs to the model-reload settings above (context
- length, KV cache, speculative decoding). Render it here, before
- the Chat Template row, so it never reads as attached to Chat
- Template (which is edited via its own dialog). When a model is
- staged (deferred load), Load/Cancel takes its place: there's
- nothing loaded to "apply" against yet. */}
- {pendingSelection ? (
-
- {stagedDownloading && (
-
- Downloading…{" "}
- {Math.round((stagedDownloadFraction ?? 0) * 100)}%
+ : ""
+ }`}
- )}
- {/* GGUF picks only: a non-GGUF pick shows none of the load
- knobs the blob captures, so there is nothing to remember. */}
- {pendingIsGguf && (
-
- setRemember(v === true)}
- // The save/clear already ran in the Load click handler, so
- // a mid-load toggle could not apply -- lock it like the knobs.
- disabled={modelControlsDisabled}
- />
- Remember settings next time
-
- )}
- {stagedLoading ? (
- // Mid-load: nothing to load or abandon until it settles, so disable.
-
-
- Loading…
-
- ) : (
-
+ {mtpUpdatable && llamaUpdateStatus?.update_available && (
{
- // Persist (or clear) this model's load knobs before loading.
- // Context is stored as the override (null = auto), never the
- // resolved native value, so restoring can't force an OOM.
- // GGUF-only, like the restore effect: saving for a
- // non-GGUF pick would snapshot leftover standing values
- // its hidden controls never showed.
- const pid = pendingIsGguf ? pendingKey : null;
- if (pid) {
- if (remember) {
- saveRememberedLoadSettings(pid, {
- contextLength: customContextLength,
- kvCacheDtype,
- speculativeType,
- specDraftNMax,
- tensorParallel,
- gpuMemoryMode,
- gpuLayers,
- nCpuMoe,
- selectedGpuIds,
- });
- } else {
- clearRememberedLoadSettings(pid);
- }
- }
- onLoadPendingModel?.();
- }}
- // Disabled while a different model is mid-load: selectModel
- // refuses a concurrent load, so the click could only toast.
- disabled={stagedDownloading || loadingModel != null}
size="sm"
- className="h-9 w-full rounded-full text-[13px] font-medium tracking-nav bg-primary text-primary-foreground hover:bg-primary/90"
+ className="corner-squircle mt-2 h-7 text-[12px]"
+ onClick={handleMtpUpdate}
+ disabled={llamaUpdating}
+ data-test-id="mtp-update-button"
>
- {loadingModel != null ? "Another model loading…" : "Load model"}
+ {llamaUpdating ? "Updating..." : "Update llama.cpp"}
- {
- // Cancel abandons the stage; if a download is mid-flight,
- // stop it too rather than leaving it running headless.
- if (stagedDownloading) onCancelStagedDownload?.();
- abandonStagedModel();
- }}
- className="h-9 w-full rounded-full text-[13px] font-medium tracking-nav text-muted-foreground"
- >
- Cancel
-
-
- )}
-
- ) : modelSettingsDirty ? (
-
- onReloadModel?.()}
- size="sm"
- className="h-7 px-3 text-[12px] font-medium tracking-nav bg-primary/92 text-primary-foreground hover:bg-primary"
- >
- Apply
-
- resetModelSettingsToLoaded()}
- className="h-7 px-3 text-[12px] font-medium tracking-nav text-muted-foreground"
- >
- Reset
-
-
- ) : null}
- {/* The template override is a load-time knob too (applied on the next
- reload) and the in-flight load already snapshotted it, so lock its
- editors like the sibling controls -- a mid-load save would be
- silently clobbered by the load response despite its toast. */}
-
-
-
+ )}
+
+ )}
+ {showContextVramWarning && (
+
+ Context length exceeds the estimated VRAM capacity (
+ {ggufMaxContextLength?.toLocaleString()} tokens). The
+ model may use system RAM.
+
+ )}
+
+
)}
-
+
savePresetWithName(presetNameInput)}
disabled={!(settingsHydrated && presetSaveState.canSubmit)}
- variant={presetSaveState.isSaveReady ? "default" : "outline"}
+ variant={
+ presetSaveState.isSaveReady ? "default" : "outline"
+ }
size="sm"
className={cn(
"h-9 w-full rounded-full text-[13px] font-medium tracking-nav",
@@ -1850,7 +912,8 @@ export function ChatSettingsPanel({
Prompt caching
- Reuse compatible prompt prefixes for lower latency and cost.
+ Reuse compatible prompt prefixes for lower latency and
+ cost.
Anthropic exposes a 5 minute and a 1 hour ephemeral
- cache pool. The 1 hour pool costs 2x base input on
- write vs 1.25x for 5 minute, but reads stay 0.1x for
- both, so a single read landing more than 5 minutes
- after the write pays off the premium.
+ cache pool. The 1 hour pool costs 2x base input on write
+ vs 1.25x for 5 minute, but reads stay 0.1x for both, so
+ a single read landing more than 5 minutes after the
+ write pays off the premium.
-
+
-
{showTemperature ? (
@@ -2071,21 +1133,12 @@ export function ChatSettingsPanel({
max={2}
step={0.1}
onChange={set("presencePenalty")}
- displayValue={params.presencePenalty === 0 ? "Off" : undefined}
+ displayValue={
+ params.presencePenalty === 0 ? "Off" : undefined
+ }
info="Penalizes any token that has already appeared at least once, encouraging the model to introduce new topics. 0 = off."
/>
) : null}
- {!isExternalModel && !isGguf && (
-
- )}
= baseContext
? "Max"
- : undefined
+ : !isExternalModel &&
+ !isGguf &&
+ params.maxTokens >= maxTokensMax
+ ? "Max"
+ : undefined
}
info="Maximum number of tokens to generate per response. Generation stops at this limit or when the model emits an end-of-sequence token."
/>
- {!isExternalModel ? (
+ {isExternalModel ? null : (
@@ -2129,13 +1175,13 @@ export function ChatSettingsPanel({
- ) : null}
+ )}
- {!isExternalModel ? (
+ {isExternalModel ? null : (
- ) : null}
+ )}
@@ -2480,155 +1526,3 @@ function BypassPermissionsToggle() {
);
}
-
-function ChatTemplateFields({ disabled = false }: { disabled?: boolean }) {
- const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate);
- const override = useChatRuntimeStore((s) => s.chatTemplateOverride);
- const setOverride = useChatRuntimeStore((s) => s.setChatTemplateOverride);
- const [editorOpen, setEditorOpen] = useState(false);
- const [draft, setDraft] = useState("");
-
- if (!defaultTemplate) return null;
-
- const displayValue = override ?? defaultTemplate;
- const isModified = override !== null;
- const draftDirty = draft !== displayValue;
-
- const openEditor = () => {
- setDraft(displayValue);
- setEditorOpen(true);
- };
- const saveEditor = () => {
- const cleared = draft.trim().length === 0 || draft === defaultTemplate;
- setOverride(cleared ? null : draft);
- setEditorOpen(false);
- toast.success(
- cleared
- ? "Chat template reset to default. It applies on the next model reload."
- : "Chat template saved. It applies on the next model reload.",
- );
- };
-
- return (
- <>
-
-
- Chat Template
-
-
- {isModified && (
-
-
- setOverride(null)}
- disabled={disabled}
- className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white disabled:pointer-events-none disabled:opacity-50"
- aria-label="Revert chat template"
- >
-
-
-
-
- Revert changes
-
-
- )}
-
-
-
-
-
-
-
- Edit template
-
-
-
-
-
-
-
- Edit Chat Template
-
- Override the model's chat template. The change applies on the
- next model reload.
-
-
-
-
-
Template editor
-
- Jinja syntax. Save matching the default clears the override.
-
-
-
-
- setDraft(defaultTemplate)}
- disabled={draft === defaultTemplate}
- className="text-muted-foreground"
- >
- Reset
-
-
- setEditorOpen(false)}
- >
- Cancel
-
- {/* Also locked mid-load: an autoLoad can start with this dialog
- already open, and a save then would be silently clobbered. */}
-
- Save
-
-
-
-
-
- >
- );
-}
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
index 3003b52230..b72a7a95c2 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
@@ -32,14 +32,13 @@ import {
GPU_LAYERS_AUTO,
isLocalModelPath,
loadedGpuMemoryFields,
- loadedGpuMemoryFieldsUnlessStaged,
- pendingSelectionMatches,
persistGpuMemoryModeOnLoad,
readPersistedSpeculativeType,
reconcilePersistedGpuIds,
resolveToolsEnabledOnLoad,
saveSpeculativeType,
useChatRuntimeStore,
+ type LoadingModelPick,
type ReasoningEffort,
} from "../stores/chat-runtime-store";
import { clampReasoningEffortToLevels } from "../provider-capabilities";
@@ -61,7 +60,10 @@ import {
isMultimodalResponse,
} from "../types/api";
import { isExternalModelId } from "../external-providers";
-import { cancelStagedModelDownload } from "@/features/hub";
+import {
+ applyPerModelConfigToRuntime,
+ type PerModelConfig,
+} from "@/features/model-picker";
import type {
ChatLoraSummary,
ChatModelSummary,
@@ -81,14 +83,16 @@ export type SelectedModelInput = {
expectedBytes?: number;
forceReload?: boolean;
nativePathToken?: string;
+ nativePathExpiresAtMs?: number | null;
/** Direct local .gguf file (no HF variant / native token) — still a GGUF
* source, so the staging flow treats it as one. */
isGguf?: boolean;
throwOnError?: boolean;
/** Keep the current speculative-decoding choice across the model switch
- * instead of resetting it to the standing preference. Set by the deferred
- * ("Load on selection") Load, where the user picked it for this model. */
+ * instead of resetting it to the standing preference. */
keepSpeculative?: boolean;
+ config?: PerModelConfig;
+ previousConfig?: PerModelConfig;
};
// Approved fingerprints by checkpoint, so a rollback after a failed switch can resend
@@ -347,6 +351,18 @@ export async function resyncInferenceStatusAfterServerModelChange(): Promise state.params);
const models = useChatRuntimeStore((state) => state.models);
@@ -385,12 +401,16 @@ export function useChatModelRuntime() {
}, []);
const resetLoadingUi = useCallback(() => {
+ const inFlight = loadingModelRef.current;
setLoadingModel(null);
setLoadProgress(null);
loadingModelRef.current = null;
loadAbortRef.current = null;
loadToastIdRef.current = null;
setLoadToastDismissedState(false);
+ if (inFlight) {
+ useChatRuntimeStore.getState().clearLoadingModelPick(pickOf(inFlight));
+ }
if (!cancelUnloadPendingRef.current) {
useChatRuntimeStore.getState().setModelLoading(false);
}
@@ -424,6 +444,7 @@ export function useChatModelRuntime() {
loadAbortRef.current?.abort();
loadAbortRef.current = null;
loadingModelRef.current = null;
+ useChatRuntimeStore.getState().clearLoadingModelPick(pickOf(model));
const tid = loadToastIdRef.current;
loadToastIdRef.current = null;
setLoadingModel(null);
@@ -460,45 +481,36 @@ export function useChatModelRuntime() {
typeof selection === "string" ? false : selection.forceReload ?? false;
const nativePathToken =
typeof selection === "string" ? undefined : selection.nativePathToken;
+ const nativePathExpiresAtMs =
+ typeof selection === "string"
+ ? null
+ : selection.nativePathExpiresAtMs ?? null;
const explicitIsGguf =
typeof selection === "string" ? undefined : selection.isGguf;
const throwOnError =
typeof selection === "string" ? false : selection.throwOnError ?? false;
const keepSpeculative =
typeof selection === "string" ? false : selection.keepSpeculative ?? false;
- // Picking/loading any model abandons a staged (deferred) selection.
- // Before the early-returns below so even a no-op re-select clears the
- // stage.
- const staged = useChatRuntimeStore.getState().pendingSelection;
- if (staged) {
- // Loading a DIFFERENT model abandons this stage. Loading the staged pick
- // ITSELF keeps it so the sidebar can show its load settings (context, KV
- // cache, …) during the load. Cleared on success below; on failure it's
- // left staged so the user can retry (see onLoadPendingModel's catch).
- const loadingStagedPick = pendingSelectionMatches(staged, {
- id: modelId,
- ggufVariant,
- nativePathToken,
- });
- if (!loadingStagedPick) {
- cancelStagedModelDownload(staged);
- useChatRuntimeStore.getState().setPendingSelection(null);
- }
- }
const currentVariant = useChatRuntimeStore.getState().activeGgufVariant;
if (!forceReload && (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null)))) {
+ if (typeof selection !== "string" && selection.previousConfig) {
+ applyPerModelConfigToRuntime(selection.previousConfig);
+ }
return;
}
- // A load is already in flight. If it's this exact pick (id + GGUF variant +
- // native path token), ignore the duplicate click. If it's a DIFFERENT model
- // -- crucially including a different GGUF variant of the same repo, which the
- // old id+token-only guard wrongly treated as a duplicate and silently
- // no-op'd -- don't start a second concurrent load (the load path has no clean
- // supersession) and don't silently swallow the request: surface it so the
- // user knows to wait for, or cancel, the in-flight load. Centralized here so
- // every entry point is covered, not just the staged Load button.
- const inFlightLoad = loadingModelRef.current;
+ // A load is already in flight. If it's this exact pick (id + variant + token),
+ // ignore the duplicate click. If it's a DIFFERENT model (including a different
+ // GGUF variant of the same repo, which the old id+token guard wrongly treated
+ // as a duplicate), don't start a second concurrent load and don't swallow the
+ // request: surface it so the user waits or cancels. Centralized here so every
+ // entry point is covered, not just the staged Load button.
+ const inFlightLoad =
+ loadingModelRef.current ??
+ useChatRuntimeStore.getState().loadingModelPick;
if (inFlightLoad) {
+ if (typeof selection !== "string" && selection.previousConfig) {
+ applyPerModelConfigToRuntime(selection.previousConfig);
+ }
const loadingSamePick =
inFlightLoad.id === modelId &&
(inFlightLoad.ggufVariant ?? null) === (ggufVariant ?? null) &&
@@ -526,7 +538,11 @@ export function useChatModelRuntime() {
// native model intents only grant .gguf files), but its id is a display
// label that need not end in ".gguf" -- without this, Manual + Auto
// layers would pin the UI context instead of letting --fit size it.
- const isGguf = explicitIsGguf ?? model?.isGguf ?? nativePathToken != null;
+ const isGguf =
+ explicitIsGguf ??
+ (ggufVariant != null ||
+ nativePathToken != null ||
+ model?.isGguf === true);
const loraIsAdapter = lora?.exportType === "lora";
const isLora =
explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false;
@@ -575,6 +591,7 @@ export function useChatModelRuntime() {
};
setLoadingModel(loadInfo);
useChatRuntimeStore.getState().setModelLoading(true);
+ useChatRuntimeStore.getState().setLoadingModelPick(pickOf(loadInfo));
setLoadProgress(
isDownloaded || isCachedLora
? { percent: null, label: null, phase: "starting" }
@@ -600,26 +617,33 @@ export function useChatModelRuntime() {
|| previousVariant != null
|| previousActiveNativePathToken != null
|| (previousCheckpoint?.toLowerCase().endsWith(".gguf") ?? false);
+ // Roll back to the previous model's own context. previousConfig was
+ // snapshotted before this load pre-applied the next model's config, so
+ // params.maxSeqLength may already be the next model's; use it only when
+ // no snapshot exists.
+ const previousMaxSeqLength =
+ (typeof selection !== "string"
+ ? selection.previousConfig?.maxSeqLength
+ : null) ?? maxSeqLength;
// Respect the rolled-back model's auto-layers mode: a Manual+Auto model
- // with an unpinned (auto) context must reload with 0 (so --fit
- // re-auto-sizes), not the positive context it happened to pick (which
- // the backend would treat as a pin).
+ // with an unpinned context must reload with 0 (so --fit re-auto-sizes),
+ // not the positive context it picked (which the backend treats as a pin).
const rollbackMaxSeqLength = resolveFitMaxSeqLength(
previousIsGguf,
stateBeforeUnload.loadedGpuMemoryMode ?? "auto",
stateBeforeUnload.loadedGpuLayers ?? GPU_LAYERS_AUTO,
stateBeforeUnload.loadedCustomContextLength,
- previousIsGguf ? (stateBeforeUnload.ggufContextLength ?? 0) : maxSeqLength,
+ previousIsGguf
+ ? (stateBeforeUnload.ggufContextLength ?? 0)
+ : previousMaxSeqLength,
);
const hfToken = stateBeforeUnload.hfToken || null;
const previousModelRequiresTrustRemoteCode =
stateBeforeUnload.modelRequiresTrustRemoteCode;
+ const previousActiveNativePathExpiresAtMs =
+ stateBeforeUnload.activeNativePathExpiresAtMs;
// Snapshot the load settings at click time, before the awaits below
- // (validation, the trust dialog, unload). For a staged Load these knobs
- // stay editable and a sheet-close revert (abandonStagedModel) can fire
- // mid-load; reading them live just before loadModel would let the load
- // use post-click values. The model-switch speculative reset below
- // updates this snapshot in lock-step so non-staged loads are unchanged.
+ // (validation, the trust dialog, unload).
const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride;
const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype;
// gpuMemoryMode is a standing preference (kept across a model switch);
@@ -657,14 +681,11 @@ export function useChatModelRuntime() {
// context can exceed maxSeqLength, so sizing on raw maxSeqLength could
// pass, unload, then have /load refuse it. Uses the click-time
// snapshot (same values loadModel uses below), so the two agree.
- // Mirror what /load does on a cross-model switch: the reset below
- // clears the per-model Auto-layers context pin + GPU pick, and
- // Manual+Auto sizes context through resolveFitMaxSeqLength.
- // gpuMemoryMode is a standing preference, kept across the switch.
- // A same-repo quant switch (same checkpoint, different gguf_variant)
- // is a different model for per-model knobs: the pinned context,
- // gpuLayers, GPU pick, and MoE offload are scoped per variant, so
- // treat a variant change like a model switch and re-baseline them.
+ // Mirror /load on a cross-model switch: the reset below clears the
+ // per-model Auto-layers context pin + GPU pick; gpuMemoryMode is a
+ // standing preference kept across the switch. A same-repo quant switch
+ // (different gguf_variant) is a different model for per-model knobs
+ // (context/gpuLayers/pick/MoE are per variant), so re-baseline them too.
const switchingModelOrVariant =
currentCheckpoint !== modelId ||
(loadActiveGgufVariant ?? null) !== (ggufVariant ?? null);
@@ -867,7 +888,11 @@ export function useChatModelRuntime() {
// The load applied this spec mode, so persist the user's standing
// preference now (the requested intent, not the resolved echo;
// saveSpeculativeType keeps only the universal auto/ngram/off).
- saveSpeculativeType(loadSpeculativeType);
+ // Skip for a per-model config (keepSpeculative): that choice is
+ // model-specific and must not overwrite the global default.
+ if (!keepSpeculative) {
+ saveSpeculativeType(loadSpeculativeType);
+ }
// Persist the GPU Memory mode only on a successful load (not on
// dropdown change), so an abandoned selection doesn't stick.
persistGpuMemoryModeOnLoad(loadResponse, loadGpuMemoryMode);
@@ -907,7 +932,9 @@ export function useChatModelRuntime() {
? (loadResponse.native_context_length ?? null)
: null;
// Keep an explicit Manual+Auto context pin (so a later Apply doesn't
- // revert it to Auto); other cases baseline on ggufContextLength.
+ // revert it to Auto) and retain the user's requested context so
+ // re-open/re-save keeps the intended override, not the backend's
+ // auto-fit context; null stays null.
const keepCustomCtx = resolveManualAutoCtxPin(
loadGpuMemoryMode,
loadGpuLayers,
@@ -978,6 +1005,9 @@ export function useChatModelRuntime() {
loadedIsMultimodal: isMultimodalResponse(loadResponse),
loadedIsDiffusion: loadResponse.is_diffusion ?? false,
activeNativePathToken: nativePathToken ?? null,
+ activeNativePathExpiresAtMs: nativePathToken
+ ? nativePathExpiresAtMs
+ : null,
});
// Unlock attach menus for capabilities the catalog entry lacked.
syncModelCapabilities(modelId, loadResponse);
@@ -1031,25 +1061,6 @@ export function useChatModelRuntime() {
recordLastLocalModelLoad({ id: modelId, kind: "model" });
}
}
- // A successful load owns the shared (pick-unscoped) settings fields,
- // so any surviving stage is stale: the just-loaded pick itself, or a
- // pick queued for a different model mid-load whose knobs this load
- // overwrote. Drop it. Only a DIFFERENT pick's download needs
- // cancelling; the loaded pick's is already consumed, and cancelling
- // it inside its post-complete linger window would flicker its card.
- const staleStage = useChatRuntimeStore.getState().pendingSelection;
- if (staleStage) {
- if (
- !pendingSelectionMatches(staleStage, {
- id: modelId,
- ggufVariant,
- nativePathToken,
- })
- ) {
- cancelStagedModelDownload(staleStage);
- }
- useChatRuntimeStore.getState().setPendingSelection(null);
- }
} catch (error) {
// Skip rollback if user cancelled -- model is already being unloaded.
if (abortCtrl.signal.aborted) throw error;
@@ -1091,8 +1102,9 @@ export function useChatModelRuntime() {
// Restore the previous model in the split mode it was running,
// not the default layer split.
tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false,
+ // Restore the previous model's GPU Memory placement, not backend defaults.
gpu_memory_mode: stateBeforeUnload.loadedGpuMemoryMode ?? "auto",
- gpu_layers: stateBeforeUnload.loadedGpuLayers ?? -1,
+ gpu_layers: stateBeforeUnload.loadedGpuLayers ?? GPU_LAYERS_AUTO,
n_cpu_moe: stateBeforeUnload.loadedNCpuMoe ?? 0,
tensor_split: stateBeforeUnload.loadedSplitRatio ?? undefined,
gpu_ids: stateBeforeUnload.loadedGpuIds ?? undefined,
@@ -1102,28 +1114,27 @@ export function useChatModelRuntime() {
);
useChatRuntimeStore.setState({
activeNativePathToken: previousActiveNativePathToken ?? null,
+ // Restore the previous token's lease together with the token so a
+ // rollback never pairs restored token A with failed load B's expiry.
+ activeNativePathExpiresAtMs: previousActiveNativePathToken
+ ? (previousActiveNativePathExpiresAtMs ?? null)
+ : null,
+ // Restore the editable speculative knobs to the rolled-back
+ // model's; the loaded baselines below come from its reload echo.
+ speculativeType: stateBeforeUnload.loadedSpeculativeType ?? null,
+ specDraftNMax: stateBeforeUnload.loadedSpecDraftNMax ?? null,
loadedSpeculativeType: rollbackSpeculativeType,
loadedSpecDraftNMax:
rollbackResponse.spec_draft_n_max ?? null,
loadedKvCacheDtype: rollbackResponse.cache_type_kv ?? null,
loadedChatTemplateOverride:
stateBeforeUnload.loadedChatTemplateOverride,
- // Re-baseline the GPU knobs from the rolled-back load's own
- // response (the shared seeding every load path uses): the
- // refresh() below can't do it, since the status reseed is
- // gated off while modelLoading is still true. A failed staged
- // Load stays staged for retry, so the staged hold applies.
- ...loadedGpuMemoryFieldsUnlessStaged(rollbackResponse, {
- tensorParallel: rollbackResponse.tensor_parallel ?? false,
- loadedTensorParallel:
- rollbackResponse.tensor_parallel ?? false,
- // refresh() is held while modelLoading remains true, so
- // restore the rolled-back model's context pin directly.
- customContextLength:
- stateBeforeUnload.loadedCustomContextLength,
- }),
+ ...loadedGpuMemoryFields(rollbackResponse),
+ tensorParallel: rollbackResponse.tensor_parallel ?? false,
loadedTensorParallel:
rollbackResponse.tensor_parallel ?? false,
+ customContextLength:
+ stateBeforeUnload.loadedCustomContextLength,
loadedCustomContextLength:
stateBeforeUnload.loadedCustomContextLength,
});
@@ -1444,6 +1455,9 @@ export function useChatModelRuntime() {
resetLoadingUi();
}
} catch (error) {
+ if (typeof selection !== "string" && selection.previousConfig) {
+ applyPerModelConfigToRuntime(selection.previousConfig);
+ }
if (abortCtrl.signal.aborted) return; // User cancelled, nothing to report
resetLoadingUi();
const message =
@@ -1474,6 +1488,13 @@ export function useChatModelRuntime() {
if (!params.checkpoint) {
return false;
}
+ const runtime = useChatRuntimeStore.getState();
+ if (runtime.modelLoading || runtime.loadingModelPick) {
+ toast.info("A model is loading", {
+ description: "Wait for it to finish or cancel it first.",
+ });
+ return false;
+ }
setModelsError(null);
if (isExternalModelId(params.checkpoint)) {
clearCheckpoint();
diff --git a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
deleted file mode 100644
index a3e7a2d264..0000000000
--- a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
+++ /dev/null
@@ -1,169 +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
-
-import { useCallback, useEffect } from "react";
-
-import { useRepoDownload } from "@/features/hub/download-manager/use-repo-download";
-import type { DownloadJob } from "@/features/hub/download-manager/use-repo-download";
-import { useLatestRef } from "@/features/hub/hooks/use-latest-ref";
-
-import { fetchGgufStagedMetadata } from "../api/chat-api";
-import {
- isPendingGguf,
- pendingSelectionMatches,
- useChatRuntimeStore,
-} from "../stores/chat-runtime-store";
-import type { PendingModelSelection } from "../stores/chat-runtime-store";
-
-/**
- * Drives the deferred ("Load on selection" off) staging flow for a GGUF:
- * download the file if needed (HF repo) or read it in place (native drag-drop /
- * picked file), then read its header context length so the settings sheet can
- * show the real context slider before the single GPU load. The staged context
- * lands on `pendingSelection.contextLength` (scoped to the staged model, never
- * the loaded model's `ggufContextLength`). Returns the live download job so the
- * sheet can render progress / cancel. Mount once on the chat page.
- */
-export function useStagedModelPreparation(opts?: {
- /** Load the cached file once an autoLoad pick's download completes. */
- onAutoLoad?: (pending: PendingModelSelection) => void;
-}): DownloadJob {
- const pendingId = useChatRuntimeStore((s) => s.pendingSelection?.id ?? null);
- const pendingVariant = useChatRuntimeStore(
- (s) => s.pendingSelection?.ggufVariant ?? null,
- );
- const pendingNativeToken = useChatRuntimeStore(
- (s) => s.pendingSelection?.nativePathToken ?? null,
- );
- // Only GGUF picks (HF variant or native file) have a header worth reading.
- const pendingIsGguf = useChatRuntimeStore((s) =>
- isPendingGguf(s.pendingSelection),
- );
- // Non-GGUF HF repos download a full snapshot (variant null) but have no header.
- const pendingIsHubRepo = useChatRuntimeStore(
- (s) => s.pendingSelection?.isHubRepo ?? false,
- );
- const pendingDownloaded = useChatRuntimeStore(
- (s) => s.pendingSelection?.isDownloaded ?? false,
- );
- // "Already probed" must key off layerCount / moeLayerCount, which only the
- // full header probe fills (it sets all three together, so either is a
- // reliable marker). contextLength alone can be list-seeded from
- // /gguf-variants, which returns no layer/MoE counts -- treating it as
- // complete would skip the probe and leave the GPU Layers slider at its 256
- // fallback and the MoE slider hidden until the model loads.
- const pendingHasMetadata = useChatRuntimeStore(
- (s) =>
- s.pendingSelection?.layerCount != null ||
- s.pendingSelection?.moeLayerCount != null,
- );
- const setPendingSelection = useChatRuntimeStore((s) => s.setPendingSelection);
- const onAutoLoadRef = useLatestRef(opts?.onAutoLoad);
-
- // A failed or cancelled autoLoad download has no sheet to retry from, so drop
- // the staged pick rather than leave it waiting on a load that won't come.
- const handleAutoLoadAbort = useCallback((variant: string | null) => {
- const latest = useChatRuntimeStore.getState().pendingSelection;
- if (
- latest?.autoLoad &&
- (latest.ggufVariant ?? null) === (variant ?? null)
- ) {
- useChatRuntimeStore.getState().abandonStagedModel();
- }
- }, []);
-
- const fetchContextMetadata = useCallback(async () => {
- const current = useChatRuntimeStore.getState().pendingSelection;
- if (!current?.id || !isPendingGguf(current)) return;
- const { id, ggufVariant, nativePathToken } = current;
- try {
- const { contextLength, layerCount, moeLayerCount } =
- await fetchGgufStagedMetadata({
- model_path: id,
- gguf_variant: ggufVariant,
- hf_token: useChatRuntimeStore.getState().hfToken || null,
- nativePathToken,
- });
- // Apply only if the same model is still staged (the user may have switched
- // picks or loaded/cancelled while the request was in flight).
- const latest = useChatRuntimeStore.getState().pendingSelection;
- if (
- latest &&
- pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken }) &&
- (contextLength != null || layerCount != null || moeLayerCount != null)
- ) {
- setPendingSelection({
- ...latest,
- contextLength,
- layerCount,
- moeLayerCount,
- });
- }
- } catch {
- // Leave metadata null: the context/MoE sliders stay hidden and the user
- // can still load (they fill in from the load response afterwards).
- }
- }, [setPendingSelection]);
-
- const job = useRepoDownload({
- kind: "model",
- // useRepoDownload must be called unconditionally; an idle repo id keeps it
- // inert until something is staged.
- repoId: pendingId ?? "__staged_idle__",
- activeVariant: pendingVariant,
- onComplete: (variant) => {
- // autoLoad picks load the cached file now; staged picks read the header so
- // the sheet's context slider can show before a manual load.
- const latest = useChatRuntimeStore.getState().pendingSelection;
- if (
- latest?.autoLoad &&
- (latest.ggufVariant ?? null) === (variant ?? null)
- ) {
- onAutoLoadRef.current?.(latest);
- return;
- }
- void fetchContextMetadata();
- },
- onError: handleAutoLoadAbort,
- onCancelled: handleAutoLoadAbort,
- });
-
- // job.requestStartDownload's identity changes per render; hold it in a ref so
- // the staging effect re-runs only when the staged model itself changes.
- const startDownloadRef = useLatestRef(job.requestStartDownload);
- const fetchMetadataRef = useLatestRef(fetchContextMetadata);
-
- useEffect(() => {
- // GGUF picks (header worth reading) and uncached non-GGUF hub repos (full
- // snapshot, no header) both run here; everything else is loaded directly.
- if (
- !pendingId ||
- (!pendingIsGguf && !pendingIsHubRepo) ||
- pendingHasMetadata
- ) {
- return;
- }
- // Native files and already-downloaded HF files are local: read the header
- // now. Otherwise download first (a GGUF variant, or a null-variant snapshot
- // for a hub repo); onComplete then reads the header or auto-loads.
- if (pendingNativeToken || pendingDownloaded) {
- void fetchMetadataRef.current();
- } else {
- const expectedBytes =
- useChatRuntimeStore.getState().pendingSelection?.expectedBytes ?? 0;
- void startDownloadRef.current(pendingVariant, expectedBytes);
- }
- }, [
- pendingId,
- pendingVariant,
- pendingNativeToken,
- pendingIsGguf,
- pendingIsHubRepo,
- pendingDownloaded,
- pendingHasMetadata,
- startDownloadRef,
- fetchMetadataRef,
- ]);
-
- return job;
-}
diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts
index b0059b57b1..c5b24bbb55 100644
--- a/studio/frontend/src/features/chat/index.ts
+++ b/studio/frontend/src/features/chat/index.ts
@@ -3,20 +3,34 @@
export { ChatPage, validateChatSearch, type ChatSearch } from "./chat-page";
export {
+ addScanFolder,
+ browseFolders,
deleteChatAttachment,
+ deleteFineTunedModel,
fetchChatAttachmentBlob,
+ fetchGgufStagedMetadata,
+ getCachedModelPath,
getInferenceStatus,
listChatAttachments,
listGgufVariants,
listLocalModels,
+ listRecommendedFolders,
+ listScanFolders,
loadModel,
+ removeScanFolder,
+ revealCachedModel,
+ type BrowseFoldersResponse,
+ type CachedGgufRepo,
+ type CachedModelRepo,
type ChatAttachmentPage,
type ChatAttachmentRecord,
type LocalModelInfo,
+ type ScanFolderInfo,
} from "./api/chat-api";
export type { GgufVariantDetail } from "./types/api";
export {
ChatSettingsPanel,
+ ParamSlider,
defaultInferenceParams,
type InferenceParams,
type Preset,
@@ -25,6 +39,11 @@ export { useChatRuntimeStore } from "./stores/chat-runtime-store";
export {
CHAT_RAG_CAPTION_KEY,
CHAT_RAG_OCR_KEY,
+ normalizeSpeculativeType,
+ readPersistedSpeculativeType,
+ readPersistedGpuMemoryMode,
+ reconcilePersistedGpuIds,
+ GPU_LAYERS_AUTO,
} from "./stores/chat-runtime-store";
export {
preferFullToolOutput,
@@ -46,9 +65,11 @@ export {
} from "./hooks/use-chat-model-runtime";
export {
customProviderDisplayName,
+ isCustomProviderType,
isExternalModelId,
parseExternalModelId,
} from "./external-providers";
+export { ApiProviderLogo } from "./api-provider-logo";
export { useExternalProvidersStore } from "./stores/external-providers-store";
export { ChatSearchDialog } from "./components/chat-search-dialog";
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
@@ -62,8 +83,8 @@ export {
useSelectedChatArtifact,
} from "./artifacts/store";
export {
- downloadChatExport,
downloadArchivedChatExport,
+ downloadChatExport,
} from "./utils/export-chat-history";
export {
clearNewChatDraft,
diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
index 69bb38bbbe..547c3f374e 100644
--- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
+++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
@@ -280,12 +280,9 @@ export function applyActiveModelStatusToStore(
ggufContextLength: currentGgufContextLength,
ggufMaxContextLength,
ggufNativeContextLength,
- // A non-GGUF status must also drop a stale native-path token: without this the
- // isGguf OR (activeGgufVariant || activeNativePathToken || ggufContextLength)
- // stays true after switching from a native GGUF to a transformers model, so a
- // Codex-only detection would auto-select for a model its preflight rejects. A real
- // GGUF load reports is_gguf: true, so its token is preserved (the load path owns it).
- ...(status.is_gguf ? {} : { activeNativePathToken: null }),
+ ...(status.is_gguf
+ ? {}
+ : { activeNativePathToken: null, activeNativePathExpiresAtMs: null }),
modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false,
defaultChatTemplate: nextDefaultChatTemplate,
loadedIsMultimodal: isMultimodalResponse(status),
@@ -299,13 +296,11 @@ export function applyActiveModelStatusToStore(
// model changed underneath this tab (auto-switch, another client), the
// old model's baselines are stale and must adopt the new status.
...(seedLoadParams &&
- prevState.pendingSelection == null &&
(prevState.loadedSpeculativeType === null || hydratingExistingModel) && {
speculativeType: currentSpecType,
loadedSpeculativeType: currentSpecType,
}),
...(seedLoadParams &&
- prevState.pendingSelection == null &&
status.spec_draft_n_max !== undefined &&
(hydratingExistingModel ||
(prevState.loadedSpecDraftNMax === null &&
@@ -314,14 +309,12 @@ export function applyActiveModelStatusToStore(
loadedSpecDraftNMax: status.spec_draft_n_max ?? null,
}),
...(seedLoadParams &&
- prevState.pendingSelection == null &&
status.cache_type_kv !== undefined &&
(prevState.loadedKvCacheDtype === null || hydratingExistingModel) && {
kvCacheDtype: status.cache_type_kv,
loadedKvCacheDtype: status.cache_type_kv,
}),
...(seedLoadParams &&
- prevState.pendingSelection == null &&
status.tensor_parallel !== undefined &&
(prevState.loadedTensorParallel === null || hydratingExistingModel) && {
tensorParallel: status.tensor_parallel,
@@ -331,7 +324,6 @@ export function applyActiveModelStatusToStore(
// placement change. gpuStatusFields preserves dirty local edits in the last
// case while advancing their loaded baselines.
...(seedLoadParams &&
- prevState.pendingSelection == null &&
(prevState.loadedGpuMemoryMode === null ||
hydratingExistingModel ||
gpuStatusChanged) &&
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 461eef99b2..70dac70a23 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -79,6 +79,12 @@ import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge
import { NewProjectDialog } from "./components/new-project-dialog";
import { useChatProjects } from "./hooks/use-chat-projects";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
+import {
+ DEFAULT_MAX_SEQ_LENGTH,
+ normalizeMaxSeqLength,
+ resolveInitialConfig,
+ type PerModelConfig,
+} from "@/features/model-picker";
import {
confirmTransformersUpgradeIfNeeded,
useTransformersUpgradeDialogStore,
@@ -97,7 +103,7 @@ import {
usePlusMenuPrefsStore,
} from "./stores/plus-menu-prefs-store";
import {
- loadedGpuMemoryFieldsUnlessStaged,
+ loadedGpuMemoryFields,
type ReasoningEffort,
reconcilePersistedGpuIds,
resolveLoadedSpeculativeSettings,
@@ -495,8 +501,24 @@ type CompareModelSelection = {
id: string;
isLora: boolean;
ggufVariant?: string;
+ config?: PerModelConfig;
};
+function cleanCompareChatTemplate(
+ value: string | null | undefined,
+): string | null {
+ return value?.trim() ? value : null;
+}
+
+function resolveCompareSpecDraftNMax(
+ speculativeType: string | null,
+ value: number | null,
+): number | null {
+ return speculativeType === "mtp" || speculativeType === "mtp+ngram"
+ ? value
+ : null;
+}
+
// Tool icon plus an X overlay CSS reveals on hover when the pill is active.
function PillGlyph({ children }: { children: ReactNode }) {
return (
@@ -1023,15 +1045,12 @@ export function SharedComposer({
// Generalized compare: load each model before dispatching to its side
if (isGeneralizedCompare) {
const store = useChatRuntimeStore.getState();
- const maxSeqLength = store.params.maxSeqLength;
const trustRemoteCode = store.params.trustRemoteCode ?? false;
- const chatTemplateOverride = store.chatTemplateOverride;
- const effectiveChatTemplateOverride = chatTemplateOverride?.trim()
- ? chatTemplateOverride
- : null;
+ const fallbackTensorParallel = store.tensorParallel;
const specSettings = resolveSpeculativeSettingsForLoad({
usePersistedPreference: true,
});
+ let loadedFromConfig = false;
function modelDisplayName(id: string): string {
const parts = id.split("/");
@@ -1058,7 +1077,6 @@ export function SharedComposer({
// path: an early remember-restore can hold a stale cross-host pick that
// /load would reject (the device cache is populated by send time).
selectedGpuIds: reconcilePersistedGpuIds(store.selectedGpuIds),
- tensorParallel: store.tensorParallel,
customContextLength: store.customContextLength,
};
// Set when an accepted transformers install unloaded the active model
@@ -1069,15 +1087,68 @@ export function SharedComposer({
sel: CompareModelSelection,
): Promise {
const currentStore = useChatRuntimeStore.getState();
+ const config = sel.config ?? null;
+ // This pane's effective config: an explicit selection config, else the
+ // remembered store config for this model/quant (never the other pane's).
+ // No saved config resolves to all-null defaults, so settings below fall
+ // through to their session default.
+ const resolved = config
+ ? { config, remembered: true }
+ : resolveInitialConfig(sel.id, sel.ggufVariant ?? null);
+ const ownConfig = resolved.config;
+ const ownRemembered = resolved.remembered;
+ // Mirror single-view resolveLoadMaxSeqLength: a GGUF pane with no explicit
+ // context loads at native (0 -> n_ctx_train), not the session maxSeqLength,
+ // which would silently shrink the shown context.
+ const isGgufLoad =
+ (sel.ggufVariant ?? null) != null ||
+ sel.id.toLowerCase().endsWith(".gguf");
+ // A non-GGUF pane with no saved maxSeqLength falls back to the app default,
+ // not the active model's shared runtime snapshot: else comparing a saved
+ // 128K model against an unconfigured one loads the latter at 128K and OOMs.
+ const effectiveMaxSeqLength =
+ ownConfig.customContextLength ??
+ normalizeMaxSeqLength(ownConfig.maxSeqLength) ??
+ (isGgufLoad ? 0 : DEFAULT_MAX_SEQ_LENGTH);
+ const effectiveChatTemplateOverride = cleanCompareChatTemplate(
+ ownConfig.chatTemplateOverride,
+ );
+ const effectiveSpeculativeType =
+ ownConfig.speculativeType ?? specSettings.speculativeType;
+ const effectiveSpecDraftNMax = ownRemembered
+ ? resolveCompareSpecDraftNMax(
+ effectiveSpeculativeType,
+ ownConfig.specDraftNMax,
+ )
+ : specSettings.specDraftNMax;
+ const effectiveTensorParallel = ownRemembered
+ ? ownConfig.tensorParallel
+ : fallbackTensorParallel;
+ if (ownConfig.selectedGpuIds != null) {
+ await ensureGpuDeviceCache();
+ }
+ const effectiveGpuMemoryMode =
+ ownConfig.gpuMemoryMode ?? compareLoadKnobs.gpuMemoryMode;
+ const effectiveGpuLayers =
+ ownConfig.gpuLayers ?? compareLoadKnobs.gpuLayers;
+ const effectiveNCpuMoe =
+ ownConfig.nCpuMoe ?? compareLoadKnobs.nCpuMoe;
+ const effectiveSelectedGpuIds =
+ ownConfig.selectedGpuIds !== undefined
+ ? reconcilePersistedGpuIds(ownConfig.selectedGpuIds)
+ : compareLoadKnobs.selectedGpuIds;
+ // A pane's context comes from its own config only: a saved pin, or null
+ // (Auto/native). It must not inherit the active model's shared snapshot --
+ // resolveFitMaxSeqLength would treat that as a pin and load this pane at
+ // the other model's context (changing VRAM/results or OOMing).
+ const effectiveCustomContextLength = ownConfig.customContextLength;
let loadTrustRemoteCode = trustRemoteCode;
let approvedRemoteCodeFingerprint: string | null = null;
const isAlreadyActive =
currentStore.params.checkpoint === sel.id &&
(currentStore.activeGgufVariant ?? null) ===
(sel.ggufVariant ?? null);
- // Already loaded (gate passed at first load): skip a redundant reload that would
- // re-trigger the gate without the approval fingerprint and fail for HIGH custom code.
- if (isAlreadyActive) {
+ if (isAlreadyActive && !config && !loadedFromConfig) {
return "ready";
}
const targetIsGguf =
@@ -1087,10 +1158,13 @@ export function SharedComposer({
// layers the load sends 0 / the pinned context, not raw maxSeqLength).
const compareMaxSeqLength = resolveFitMaxSeqLength(
targetIsGguf,
- compareLoadKnobs.gpuMemoryMode,
- compareLoadKnobs.gpuLayers,
- compareLoadKnobs.customContextLength,
- maxSeqLength,
+ effectiveGpuMemoryMode,
+ effectiveGpuLayers,
+ // Prefer this pane's own saved context pin over the shared snapshot,
+ // falling back to its per-pane effective context (GGUF with no saved
+ // context loads at native, not the session maxSeqLength).
+ effectiveCustomContextLength,
+ effectiveMaxSeqLength,
);
const validation = await validateModel({
model_path: sel.id,
@@ -1105,8 +1179,8 @@ export function SharedComposer({
// below: a non-GGUF target must not inherit a hidden GGUF GPU pick.
...(targetIsGguf
? {
- gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined,
- gpu_memory_mode: compareLoadKnobs.gpuMemoryMode,
+ gpu_ids: effectiveSelectedGpuIds ?? undefined,
+ gpu_memory_mode: effectiveGpuMemoryMode,
}
: {}),
});
@@ -1164,27 +1238,28 @@ export function SharedComposer({
trust_remote_code: loadTrustRemoteCode,
approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
chat_template_override: effectiveChatTemplateOverride,
- speculative_type: specSettings.speculativeType,
- spec_draft_n_max: specSettings.specDraftNMax,
- // Honor the Tensor Parallelism + GPU Memory choices on compare loads.
- // GGUF-only, like the auto-load path: the picker is a GGUF control,
- // so a non-GGUF target loads via HF auto-placement instead of being
- // pinned to a leftover GGUF pick it can't even show.
- tensor_parallel: compareLoadKnobs.tensorParallel,
+ cache_type_kv: ownConfig.kvCacheDtype ?? null,
+ speculative_type: effectiveSpeculativeType,
+ spec_draft_n_max: effectiveSpecDraftNMax,
+ tensor_parallel: effectiveTensorParallel,
...(targetIsGguf
? {
- gpu_memory_mode: compareLoadKnobs.gpuMemoryMode,
- gpu_layers: compareLoadKnobs.gpuLayers,
- n_cpu_moe: compareLoadKnobs.nCpuMoe,
+ gpu_memory_mode: effectiveGpuMemoryMode,
+ gpu_layers: effectiveGpuLayers,
+ n_cpu_moe: effectiveNCpuMoe,
tensor_split: compareLoadKnobs.splitRatio ?? undefined,
- gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined,
+ gpu_ids: effectiveSelectedGpuIds ?? undefined,
}
: {}),
});
- saveSpeculativeType(specSettings.speculativeType);
+ // Keep a compare pane's per-model speculative choice load-local: persist
+ // the global preference only when it came from global settings.
+ if (ownConfig.speculativeType == null) {
+ saveSpeculativeType(effectiveSpeculativeType);
+ }
// Persist the GPU Memory mode on a non-diffusion GGUF compare-load too,
// so an applied manual choice survives a restart.
- persistGpuMemoryModeOnLoad(resp, compareLoadKnobs.gpuMemoryMode);
+ persistGpuMemoryModeOnLoad(resp, effectiveGpuMemoryMode);
upgradeUnloadedActive = false;
const store = useChatRuntimeStore.getState();
store.setCheckpoint(
@@ -1200,9 +1275,9 @@ export function SharedComposer({
// compare loads don't send the pin, so their baseline clears.
const keepCustomCtx = targetIsGguf
? resolveManualAutoCtxPin(
- compareLoadKnobs.gpuMemoryMode,
- compareLoadKnobs.gpuLayers,
- compareLoadKnobs.customContextLength,
+ effectiveGpuMemoryMode,
+ effectiveGpuLayers,
+ effectiveCustomContextLength,
)
: null;
useChatRuntimeStore.setState({
@@ -1211,37 +1286,52 @@ export function SharedComposer({
...reasoningCapsFromLoad(resp),
supportsPreserveThinking: resp.supports_preserve_thinking ?? false,
supportsTools: resp.supports_tools ?? false,
+ kvCacheDtype: resp.cache_type_kv ?? null,
+ loadedKvCacheDtype: resp.cache_type_kv ?? null,
tensorParallel: resp.tensor_parallel ?? false,
loadedTensorParallel: resp.tensor_parallel ?? false,
- customContextLength: keepCustomCtx,
+ defaultChatTemplate: resp.chat_template ?? null,
+ chatTemplateOverride: effectiveChatTemplateOverride,
+ loadedChatTemplateOverride: effectiveChatTemplateOverride,
+ // The context baseline this pane loaded with (see keepCustomCtx above),
+ // so a later Apply/Reset can't silently revert a Manual+Auto pin.
loadedCustomContextLength: keepCustomCtx,
- // Seed the loaded GGUF context (interactive/auto-load parity): the
- // settings sheet keys the GGUF GPU controls off it for a direct .gguf
- // with no variant, and a later Apply reads it as the resolved context.
- ...(targetIsGguf
- ? {
- ggufContextLength: resp.context_length ?? 131072,
- ggufMaxContextLength:
- resp.max_context_length ?? resp.context_length ?? 131072,
- ggufNativeContextLength: resp.native_context_length ?? null,
- }
- : { ggufContextLength: null }),
- // Compare loads resolve by id (HF repo / local path), never through a
- // native-path lease, so a token left by a previously loaded native
- // GGUF is stale here -- isLoadedGguf keys off it, and a stale token
- // would dress a non-GGUF compare load in GGUF controls. Mirror the
- // interactive path, which writes it on every load success.
- activeNativePathToken: null,
- // Held under an open staged pick: setCheckpoint preserves a stage on
- // the empty->active transition, so a compare load can complete with
- // staged GPU edits still on screen.
- ...loadedGpuMemoryFieldsUnlessStaged(resp),
+ // Adopt the load response's GPU-memory fields (mode/layers/MoE/split/pick
+ // plus loaded baselines) so the GPU controls round-trip. (gguf context,
+ // customContextLength and native-path token/expiry clear in the tail below.)
+ ...loadedGpuMemoryFields(resp),
// Drives the GPU Memory controls' diffusion gate; set alongside the
// GPU fields on every load path so the gate can't read stale.
loadedIsDiffusion: resp.is_diffusion ?? false,
loadedIsMultimodal: isMultimodalResponse(resp),
+ // Record the context this pane loaded with (like the single-model path)
+ // so when it becomes the active model, the UI and later reload/save use
+ // its context, not the previous/default one.
+ customContextLength: isGgufLoad
+ ? (ownConfig.customContextLength ?? keepCustomCtx)
+ : null,
+ ggufContextLength: resp.is_gguf ? (resp.context_length ?? null) : null,
+ ggufNativeContextLength: resp.is_gguf
+ ? (resp.native_context_length ?? null)
+ : null,
+ ggufMaxContextLength: resp.is_gguf
+ ? (resp.max_context_length ?? null)
+ : null,
+ // Compare selections load by repo/variant, never from the file picker,
+ // so they carry no native lease. Clear any prior picked file's
+ // token/expiry so the reload path never sends a stale lease.
+ activeNativePathToken: null,
+ activeNativePathExpiresAtMs: null,
...resolveLoadedSpeculativeSettings(resp),
});
+ if (!isGgufLoad) {
+ // Non-GGUF panes carry their context in params.maxSeqLength.
+ store.setParams({
+ ...useChatRuntimeStore.getState().params,
+ maxSeqLength: effectiveMaxSeqLength,
+ });
+ }
+ loadedFromConfig = config != null;
// Sync the models[] entry with the load response so attach/send gates
// read fresh capabilities. /api/models/list can lag a model's actual
// state (e.g. a GGUF whose mmproj arrived after the snapshot).
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index 5786947118..89bc21ee18 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -1,16 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import type { RememberedLoadSettings } from "@/components/assistant-ui/model-selector/remembered-load-settings";
-import {
- cancelStagedModelDownload,
- mirrorHfTokenInto,
- useHfTokenStore,
-} from "@/features/hub";
-import {
- cachedPinnableGpuIndices,
- ensureGpuDeviceCache,
-} from "@/hooks/use-gpu-info";
+import { mirrorHfTokenInto, useHfTokenStore } from "@/features/hub";
+import { cachedPinnableGpuIndices } from "@/hooks/use-gpu-info";
import { toast } from "@/lib/toast";
import { create } from "zustand";
import { isExternalModelId, parseExternalModelId } from "../external-providers";
@@ -46,7 +38,6 @@ export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY =
"unsloth_chat_allow_artifact_network_access";
export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled";
export const CHAT_CONFIRM_TOOL_CALLS_KEY = "unsloth_chat_confirm_tool_calls";
-export const CHAT_LOAD_ON_SELECTION_KEY = "unsloth_chat_load_on_selection";
export const CHAT_EXPAND_QUANTIZATIONS_KEY =
"unsloth_chat_expand_quantizations";
export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY =
@@ -229,6 +220,11 @@ export type PendingImageEditReference = {
openaiResponseId?: string;
openaiReasoningItem?: unknown;
};
+export type LoadingModelPick = {
+ id: string;
+ ggufVariant: string | null;
+ nativePathToken: string | null;
+};
export type ReasoningEffort =
| "none"
| "minimal"
@@ -680,70 +676,7 @@ export function loadedGpuMemoryFields(resp: {
};
}
-/** loadedGpuMemoryFields (plus any seedExtras), unless a staged pick is open.
- *
- * With a staged pick open (the load fired mid-staging), preserve its editable
- * GPU knobs and seedExtras, but still advance every loaded baseline. Otherwise
- * cancelling the stage restores its edits onto the newly loaded model. The
- * status reseed cannot repair that while pendingSelection holds it off.
- */
-export function loadedGpuMemoryFieldsUnlessStaged(
- resp: Parameters[0],
- seedExtras?: T,
-) {
- const fields = loadedGpuMemoryFields(resp);
- if (useChatRuntimeStore.getState().pendingSelection != null) {
- return {
- loadedGpuMemoryMode: fields.loadedGpuMemoryMode,
- loadedGpuLayers: fields.loadedGpuLayers,
- loadedNCpuMoe: fields.loadedNCpuMoe,
- loadedSplitRatio: fields.loadedSplitRatio,
- loadedGpuIds: fields.loadedGpuIds,
- // These are metadata ceilings for the model that actually loaded, not
- // editable values from the open stage. Advance them with the baselines
- // so abandoning the stage cannot expose the previous model's limits.
- ggufLayerCount: fields.ggufLayerCount,
- moeLayerCount: fields.moeLayerCount,
- };
- }
- return { ...fields, ...seedExtras };
-}
-
-/** A local model staged for a deferred load (see `pendingSelection`). Shape is
- * a subset of the load hook's `SelectedModelInput`, structurally assignable. */
-export type PendingModelSelection = {
- id: string;
- isLora?: boolean;
- ggufVariant?: string;
- isDownloaded?: boolean;
- expectedBytes?: number;
- /** Native (drag-drop / picked-from-disk) GGUF: the path token used to read
- * the header and to load. Absent for HF-repo models. */
- nativePathToken?: string;
- /** Direct local .gguf file (custom folder / LM Studio): a GGUF source even
- * though it carries neither an HF variant nor a native path token. */
- isGguf?: boolean;
- /** Native context length read from the GGUF header once the file is local.
- * Scoped here (not the shared `ggufContextLength`) so a staged model's
- * metadata never pollutes the currently-loaded model's context display. */
- contextLength?: number | null;
- /** Total layer count (GGUF block_count); the manual gpu-layers ceiling is
- * this + 1 (llama.cpp counts the output layer as offloadable too);
- * scoped here like contextLength. */
- layerCount?: number | null;
- /** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling);
- * 0 for dense models, scoped here like contextLength. */
- moeLayerCount?: number | null;
- /** "Load on selection" on + un-cached GGUF: download via the manager (global
- * indicator) without opening the sheet, then load once the download finishes. */
- autoLoad?: boolean;
- /** Uncached non-GGUF HF repo: download the full snapshot via the manager
- * (variant null) the same way GGUF picks download a variant. */
- isHubRepo?: boolean;
-};
-
-/** A pick is a GGUF (HF variant, native file, or a direct local .gguf) and so
- * has pre-load options worth staging. Works on a selection or a staged pick. */
+/** A pick is a GGUF: HF variant, native file, or a direct local .gguf. */
export function hasGgufSource(x: {
ggufVariant?: string;
nativePathToken?: string;
@@ -781,30 +714,6 @@ export function isDownloadableHubRepo(x: {
);
}
-export function isPendingGguf(pending: PendingModelSelection | null): boolean {
- return pending != null && hasGgufSource(pending);
-}
-
-/** Whether `pending` refers to the same model as `pick` (id + GGUF variant +
- * native path token, optionals null-normalized). Native ids are display labels
- * that can collide, so the token must match too — id alone can land on the
- * wrong file. */
-export function pendingSelectionMatches(
- pending: PendingModelSelection | null,
- pick: {
- id: string;
- ggufVariant?: string | null;
- nativePathToken?: string | null;
- },
-): boolean {
- return (
- pending != null &&
- pending.id === pick.id &&
- (pending.ggufVariant ?? null) === (pick.ggufVariant ?? null) &&
- (pending.nativePathToken ?? null) === (pick.nativePathToken ?? null)
- );
-}
-
type ChatRuntimeStore = {
settingsHydrated: boolean;
params: InferenceParams;
@@ -988,10 +897,6 @@ type ChatRuntimeStore = {
/** Picked physical GPU indices (null = use all / automatic). */
selectedGpuIds: number[] | null;
loadedGpuIds: number[] | null;
- /** Persisted: when false, picking a local model stages it as
- * `pendingSelection` (and opens settings) instead of loading immediately,
- * so load settings can be set before the single load. */
- loadOnSelection: boolean;
/** Persisted: expand every On Device GGUF repo's quantizations by default
* instead of waiting for a click. */
expandQuantizations: boolean;
@@ -1000,9 +905,6 @@ type ChatRuntimeStore = {
/** Persisted, shared by the chat model selector and the Hub page: list only
* models whose size fits this device's memory budget. */
fitOnDeviceOnly: boolean;
- /** A local model picked while `loadOnSelection` is off: staged, not loaded.
- * The settings sheet shows its load knobs and a Load button. */
- pendingSelection: PendingModelSelection | null;
loadedIsMultimodal: boolean;
/** Active model is a block-diffusion model (DiffusionGemma): drives the
* denoising-canvas artifact auto-render. */
@@ -1041,9 +943,16 @@ type ChatRuntimeStore = {
cacheWriteTokens?: number;
} | null;
modelLoading: boolean;
+ loadingModelPick: LoadingModelPick | null;
activeNativePathToken: string | null;
+ // Wall-clock expiry (ms) of the active native path token. The desktop host
+ // prunes file leases after a TTL, so a reload checks this to prompt
+ // re-selection instead of reusing a dead token.
+ activeNativePathExpiresAtMs: number | null;
hydratePersistedSettings: () => Promise;
setModelLoading: (loading: boolean) => void;
+ setLoadingModelPick: (pick: LoadingModelPick | null) => void;
+ clearLoadingModelPick: (expected: LoadingModelPick) => void;
setModelRequiresTrustRemoteCode: (required: boolean) => void;
setParams: (params: InferenceParams) => void;
setCustomPresets: (presets: Preset[]) => void;
@@ -1119,38 +1028,14 @@ type ChatRuntimeStore = {
setNudgeToolCalls: (enabled: boolean) => void;
setMaxToolCallsPerMessage: (value: number) => void;
setToolCallTimeout: (value: number) => void;
- setKvCacheDtype: (dtype: string | null) => void;
- setSpeculativeType: (type: string | null) => void;
- setSpecDraftNMax: (value: number | null) => void;
- /** Revert the editable load knobs to the loaded model's baseline (or defaults
- * when nothing is loaded). Used by the settings-sheet Reset button and to
- * start each deferred-staging session clean so one staged pick's settings
- * don't leak onto the next. */
- resetModelSettingsToLoaded: () => void;
- /** Seed the editable load knobs from a model's remembered settings. Shared by
- * the settings sheet's restore effect and the "Load on selection" paths,
- * which skip the sheet but must still honor a saved config. */
- applyRememberedLoadSettings: (settings: RememberedLoadSettings) => void;
- setTensorParallel: (value: boolean) => void;
setGpuMemoryMode: (mode: "auto" | "manual") => void;
setGpuLayers: (value: number) => void;
setNCpuMoe: (value: number) => void;
setSplitRatio: (value: number[] | null) => void;
setSelectedGpuIds: (ids: number[] | null) => void;
- setLoadOnSelection: (value: boolean) => void;
setExpandQuantizations: (value: boolean) => void;
setShowAllQuantizations: (value: boolean) => void;
setFitOnDeviceOnly: (value: boolean) => void;
- setPendingSelection: (selection: PendingModelSelection | null) => void;
- /** Stage a pick for a deferred load: revert knobs to the loaded baseline,
- * record the selection, and open the settings sheet. */
- stageModel: (selection: PendingModelSelection) => void;
- /** Abandon a staged pick without loading: revert knobs to the loaded baseline
- * and clear the pending selection. Cancels its in-flight download too, unless
- * `keepDownload` is set (navigation keeps the transfer running, like Hub). */
- abandonStagedModel: (opts?: { keepDownload?: boolean }) => void;
- setCustomContextLength: (v: number | null) => void;
- setChatTemplateOverride: (template: string | null) => void;
setPendingAudio: (base64: string, name: string) => void;
clearPendingAudio: () => void;
setPendingImageEditReference: (
@@ -1352,38 +1237,6 @@ function setScalarSettingVersion(
saveSettingsPatch({ [key]: value });
}
-/** The "revert to the loaded model" baseline for the editable load knobs.
- * Shared by resetModelSettingsToLoaded (full revert) and stageModel (which
- * overrides speculative and the per-model GPU knobs to start a fresh pick). */
-function loadedBaselineSettings(s: ChatRuntimeStore) {
- const hasLoadedModel = Boolean(s.params.checkpoint);
- return {
- // Revert to the loaded model's pin (null = Auto), not a blanket Auto.
- customContextLength: s.loadedCustomContextLength,
- kvCacheDtype: s.loadedKvCacheDtype,
- tensorParallel: s.loadedTensorParallel ?? false,
- speculativeType: hasLoadedModel
- ? s.loadedSpeculativeType
- : readPersistedSpeculativeType(),
- specDraftNMax: hasLoadedModel ? s.loadedSpecDraftNMax : null,
- chatTemplateOverride: s.loadedChatTemplateOverride,
- // GPU memory mode is a standing preference; revert to the loaded model's
- // mode (or the persisted default when nothing is loaded). Manual knobs and
- // the GPU pick are per-model and revert to their loaded baseline. A loaded
- // model with no applicable mode -- diffusion ("auto" baseline) or non-GGUF
- // (null baseline) -- keeps the live preference so Reset can't drop it.
- gpuMemoryMode: !hasLoadedModel
- ? readPersistedGpuMemoryMode()
- : s.loadedIsDiffusion
- ? s.gpuMemoryMode
- : (s.loadedGpuMemoryMode ?? s.gpuMemoryMode),
- gpuLayers: s.loadedGpuLayers ?? GPU_LAYERS_AUTO,
- nCpuMoe: s.loadedNCpuMoe ?? 0,
- splitRatio: s.loadedSplitRatio ?? null,
- selectedGpuIds: s.loadedGpuIds,
- };
-}
-
export const useChatRuntimeStore = create((set, get) => ({
settingsHydrated: false,
// Hydrate the last external checkpoint so the external picker survives a
@@ -1493,11 +1346,9 @@ export const useChatRuntimeStore = create((set, get) => ({
moeLayerCount: null,
selectedGpuIds: null,
loadedGpuIds: null,
- loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true),
expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false),
showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true),
fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false),
- pendingSelection: null,
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
@@ -1515,7 +1366,9 @@ export const useChatRuntimeStore = create((set, get) => ({
pendingImageEditReference: null,
contextUsage: null,
modelLoading: false,
+ loadingModelPick: null,
activeNativePathToken: null,
+ activeNativePathExpiresAtMs: null,
hydratePersistedSettings: async () => {
if (get().settingsHydrated) {
return;
@@ -1554,6 +1407,20 @@ export const useChatRuntimeStore = create((set, get) => ({
return settingsHydrationPromise;
},
setModelLoading: (loading) => set({ modelLoading: loading }),
+ setLoadingModelPick: (pick) => set({ loadingModelPick: pick }),
+ clearLoadingModelPick: (expected) =>
+ set((state) => {
+ const current = state.loadingModelPick;
+ if (
+ !current ||
+ current.id !== expected.id ||
+ current.ggufVariant !== expected.ggufVariant ||
+ current.nativePathToken !== expected.nativePathToken
+ ) {
+ return state;
+ }
+ return { loadingModelPick: null };
+ }),
setModelRequiresTrustRemoteCode: (modelRequiresTrustRemoteCode) =>
set({ modelRequiresTrustRemoteCode }),
setParams: (params) =>
@@ -1634,13 +1501,6 @@ export const useChatRuntimeStore = create((set, get) => ({
// Clear stale per-turn usage on model change; the relaxed external-provider
// render gate would otherwise show old counters until the next completion.
const checkpointChanged = state.params.checkpoint !== modelId;
- const pendingToClear =
- checkpointChanged && state.params.checkpoint
- ? state.pendingSelection
- : null;
- if (pendingToClear) {
- cancelStagedModelDownload(pendingToClear);
- }
// Clamp maxTokens to the new model's cap when switching into an external
// model so a value carried over from a local session doesn't exceed the
// slider's max.
@@ -1668,14 +1528,6 @@ export const useChatRuntimeStore = create((set, get) => ({
},
activeGgufVariant: ggufVariant ?? null,
...(checkpointChanged ? { contextUsage: null } : {}),
- // Switching away from a loaded model (e.g. picking an external provider)
- // abandons any staged pick, so its Load button and edited knobs don't
- // linger over the newly active model. Same revert as abandonStagedModel.
- // Guarded on a non-empty current checkpoint: an establishing set from a
- // background status sync (empty -> active) must not wipe a fresh stage.
- ...(pendingToClear
- ? { ...loadedBaselineSettings(state), pendingSelection: null }
- : {}),
};
}),
setActiveThreadId: (activeThreadId) =>
@@ -1689,7 +1541,6 @@ export const useChatRuntimeStore = create((set, get) => ({
// clear any stored external selection so the next refresh doesn't snap
// back to a model the user intentionally cleared.
saveLastExternalCheckpoint(null);
- cancelStagedModelDownload(get().pendingSelection);
return set((state) => ({
params: {
...state.params,
@@ -1697,7 +1548,7 @@ export const useChatRuntimeStore = create((set, get) => ({
},
activeGgufVariant: null,
activeNativePathToken: null,
- pendingSelection: null,
+ activeNativePathExpiresAtMs: null,
ggufContextLength: null,
ggufMaxContextLength: null,
ggufNativeContextLength: null,
@@ -2044,10 +1895,6 @@ export const useChatRuntimeStore = create((set, get) => ({
);
return { toolCallTimeout };
}),
- setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }),
- setSpeculativeType: (speculativeType) => set({ speculativeType }),
- setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }),
- setTensorParallel: (tensorParallel) => set({ tensorParallel }),
// Standing preference, but persisted only on a successful load (see
// use-chat-model-runtime), not on selection -- so an unapplied pick the user
// resets/abandons doesn't stick to the next session.
@@ -2056,63 +1903,6 @@ export const useChatRuntimeStore = create((set, get) => ({
setNCpuMoe: (nCpuMoe) => set({ nCpuMoe }),
setSplitRatio: (splitRatio) => set({ splitRatio }),
setSelectedGpuIds: (selectedGpuIds) => set({ selectedGpuIds }),
- resetModelSettingsToLoaded: () => set((s) => loadedBaselineSettings(s)),
- applyRememberedLoadSettings: (settings) => {
- const gpuCacheWasCold = cachedPinnableGpuIndices() === null;
- const restoredGpuIds =
- settings.selectedGpuIds !== undefined
- ? reconcilePersistedGpuIds(settings.selectedGpuIds)
- : undefined;
- // Coalesce every field: a blob persisted by an older/newer build can omit
- // keys, and a raw spread would push `undefined` into fields typed non-null.
- // The GPU knobs are spread only when present, but first reset the per-model
- // ones to defaults: this path (load-on-selection) starts from the loaded
- // model's baseline and skips the model-switch reset, so a blob omitting
- // gpuLayers/nCpuMoe/selectedGpuIds (older build) or splitRatio (never
- // remembered) must not inherit the previous model's placement. gpuMemoryMode
- // (standing preference) is NOT reset, only applied when the blob carries it;
- // selectedGpuIds keeps a meaningful null (all GPUs), so it keys off undefined.
- set({
- gpuLayers: GPU_LAYERS_AUTO,
- nCpuMoe: 0,
- splitRatio: null,
- selectedGpuIds: null,
- customContextLength: settings.contextLength ?? null,
- kvCacheDtype: settings.kvCacheDtype ?? null,
- speculativeType: settings.speculativeType ?? "auto",
- specDraftNMax: settings.specDraftNMax ?? null,
- tensorParallel: settings.tensorParallel ?? false,
- ...(settings.gpuMemoryMode != null && {
- gpuMemoryMode: settings.gpuMemoryMode,
- }),
- ...(settings.gpuLayers != null && { gpuLayers: settings.gpuLayers }),
- ...(settings.nCpuMoe != null && { nCpuMoe: settings.nCpuMoe }),
- ...(restoredGpuIds !== undefined && {
- // Reconcile against the GPUs present now (see reconcilePersistedGpuIds):
- // a saved [1] on a 1-GPU host (or under relative/UUID visibility) would
- // hide the picker yet still send gpu_ids, which the backend rejects.
- selectedGpuIds: restoredGpuIds,
- }),
- });
- // A cold cache makes the synchronous restore provisional. Reconcile again
- // when the shared fetch completes, but only if this exact restored array is
- // still current so a user edit, stage change, or load cannot be overwritten.
- if (gpuCacheWasCold && restoredGpuIds != null) {
- void ensureGpuDeviceCache().then(() => {
- set((state) => {
- if (state.selectedGpuIds !== restoredGpuIds) return state;
- const reconciled = reconcilePersistedGpuIds(restoredGpuIds);
- return reconciled === restoredGpuIds
- ? state
- : { selectedGpuIds: reconciled };
- });
- });
- }
- },
- setLoadOnSelection: (loadOnSelection) => {
- saveBool(CHAT_LOAD_ON_SELECTION_KEY, loadOnSelection);
- set({ loadOnSelection });
- },
setExpandQuantizations: (expandQuantizations) => {
saveBool(CHAT_EXPAND_QUANTIZATIONS_KEY, expandQuantizations);
set({ expandQuantizations });
@@ -2125,55 +1915,6 @@ export const useChatRuntimeStore = create((set, get) => ({
saveBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, fitOnDeviceOnly);
set({ fitOnDeviceOnly });
},
- setPendingSelection: (pendingSelection) => set({ pendingSelection }),
- stageModel: (selection) => {
- // Refuse staging mid-load: post-load cleanup would silently drop the queued
- // pick. stageOrLoad toasts first for callers that can.
- if (get().modelLoading) return;
- // Rebinding to a new pick keeps the prior pick's download running so the
- // user can queue multiple downloads at once (Hub-style).
- set((s) => {
- return {
- ...loadedBaselineSettings(s),
- pendingSelection: selection,
- // autoLoad downloads silently and loads on completion, so keep the sheet shut.
- settingsPanelOpen: !selection.autoLoad,
- // Speculative starts from the standing default, not the loaded model's
- // mode, so a fresh pick doesn't inherit (and then carry, via the staged
- // Load's keepSpeculative) a forced MTP mode onto a model that may lack it.
- speculativeType: readPersistedSpeculativeType(),
- specDraftNMax: null,
- // Keep the on-screen GPU Memory selection (loadedBaselineSettings would
- // otherwise revert it to the loaded model's mode, dropping a Manual choice
- // just made). Use the live store value, not the persisted one, which can
- // lag a mode hydrated from an out-of-band load.
- gpuMemoryMode: s.gpuMemoryMode,
- // Per-model GPU knobs start from defaults too so a fresh pick doesn't
- // inherit the loaded model's layer/MoE/split/GPU choices, matching the
- // immediate-switch reset.
- gpuLayers: GPU_LAYERS_AUTO,
- nCpuMoe: 0,
- splitRatio: null,
- selectedGpuIds: null,
- // Fresh pick starts at Auto context (loadedBaselineSettings would
- // otherwise restore the current model's pin). Leaves the baseline
- // intact, like the GPU knobs, so abandoning restores the loaded pin.
- customContextLength: null,
- };
- });
- },
- abandonStagedModel: (opts) => {
- const { pendingSelection } = get();
- if (!pendingSelection) return;
- // Cancel the staged pick's in-flight download (centralized for every abandon
- // path: sheet close, thread switch, route exit, new chat). `keepDownload`
- // opts out so navigation leaves the transfer running, like a Hub download.
- if (!opts?.keepDownload) cancelStagedModelDownload(pendingSelection);
- set((s) => ({ ...loadedBaselineSettings(s), pendingSelection: null }));
- },
- setCustomContextLength: (customContextLength) => set({ customContextLength }),
- setChatTemplateOverride: (chatTemplateOverride) =>
- set({ chatTemplateOverride }),
setPendingAudio: (base64, name) =>
set({ pendingAudioBase64: base64, pendingAudioName: name }),
clearPendingAudio: () =>
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index c24ddde5f5..c9c06834c1 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -99,6 +99,9 @@ export interface ValidateModelResponse {
/** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling);
* 0 for dense models, null until downloaded. */
moe_layer_count?: number | null;
+ /** Embedded GGUF chat template, returned when include_chat_template is set
+ * (native lease-backed picks); null for non-GGUF, over-cap, or not read. */
+ chat_template?: string | null;
/** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */
requires_transformers_upgrade?: boolean;
/** Set only when requires_transformers_upgrade. */
diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx
index 29ba5a703b..6c2794420b 100644
--- a/studio/frontend/src/features/export/components/export-run-panel.tsx
+++ b/studio/frontend/src/features/export/components/export-run-panel.tsx
@@ -2,7 +2,6 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
-import { FolderBrowser } from "@/components/assistant-ui/model-selector/folder-browser";
import { Input } from "@/components/ui/input";
import {
InputGroup,
@@ -17,6 +16,7 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
+import { FolderBrowser } from "@/features/model-picker";
import {
AlertCircleIcon,
ArrowRight01Icon,
@@ -28,17 +28,17 @@ import {
import { HugeiconsIcon } from "@hugeicons/react";
import { useEffect, useRef, useState } from "react";
import { useShallow } from "zustand/react/shallow";
+import type { ExportLogEntry } from "../api/export-api";
import {
EXPORT_METHODS,
type ExportMethod,
findMergedFormat,
} from "../constants";
-import type { ExportLogEntry } from "../api/export-api";
import { getExportLogLineClass } from "../lib/log-style";
import {
+ type ExportDestination,
selectExportProgressPercent,
useExportRuntimeStore,
- type ExportDestination,
} from "../stores/export-runtime-store";
function useElapsedSeconds(startedAt: number | null, running: boolean): number {
@@ -165,7 +165,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
const isExporting = run.isExporting;
const isTerminal =
- run.phase === "success" || run.phase === "error" || run.phase === "canceled";
+ run.phase === "success" ||
+ run.phase === "error" ||
+ run.phase === "canceled";
const showConfig = run.phase === "idle";
// Gate the log area on the active run's method (from the store) as well as the
// local form selection, so it stays visible after navigating away and back
@@ -197,7 +199,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
setFollowTail(nearBottom);
};
- const methodTitle = EXPORT_METHODS.find((m) => m.value === exportMethod)?.title;
+ const methodTitle = EXPORT_METHODS.find(
+ (m) => m.value === exportMethod,
+ )?.title;
const summary = run.summary;
const summaryBaseModel = summary?.baseModelName ?? baseModelName;
const summaryCheckpoint = summary?.checkpointLabel ?? checkpoint;
@@ -290,7 +294,10 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
onClick={() => setFolderBrowserOpen(true)}
aria-label="Browse save folder"
>
-
+
Browse
@@ -301,8 +308,8 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
<>Default: {defaultSaveDirectory}>
) : (
<>
- Paste an absolute path if the folder browser cannot reach the
- drive.
+ Paste an absolute path if the folder browser cannot reach
+ the drive.
>
)}
@@ -410,7 +417,10 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
: [];
const showLabels = items.length > 1;
return items.map((o, i) => (
-
+
{showLabels && o.label ? (
{o.label}
@@ -431,14 +441,22 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
{run.phase === "canceled" && (
-
- Export canceled. Training and inference were not affected.
+
+
+ Export canceled. Training and inference were not affected.
+
)}
{run.phase === "error" && run.error && (
-
+
{run.error}
)}
@@ -447,15 +465,21 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
Base Model
- {summaryBaseModel}
+
+ {summaryBaseModel}
+
{isAdapter ? "Checkpoint" : "Model"}
- {summaryCheckpoint}
+
+ {summaryCheckpoint}
+
Export Method
- {summaryMethodLabel}
+
+ {summaryMethodLabel}
+
{summaryMethod === "merged" && summaryFormats.length > 0 && (
@@ -484,7 +508,12 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
{summaryMethod === "gguf" && run.quantTotal > 1 && (
- Quant {Math.min(run.quantIndex + (isExporting ? 1 : 0), run.quantTotal)} of {run.quantTotal}
+ Quant{" "}
+ {Math.min(
+ run.quantIndex + (isExporting ? 1 : 0),
+ run.quantTotal,
+ )}{" "}
+ of {run.quantTotal}
)}
@@ -506,7 +535,10 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
}
/>
{run.stage && (
-
+
{run.stage}
)}
@@ -556,10 +588,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
) : (
{run.logLines.map((entry, idx) => (
-
+
{formatLogLine(entry)}
))}
diff --git a/studio/frontend/src/features/hub/catalog/catalog-states.tsx b/studio/frontend/src/features/hub/catalog/catalog-states.tsx
index 1d9eb63718..a36c4bb2da 100644
--- a/studio/frontend/src/features/hub/catalog/catalog-states.tsx
+++ b/studio/frontend/src/features/hub/catalog/catalog-states.tsx
@@ -10,6 +10,7 @@ import {
} from "@hugeicons/core-free-icons";
import type { IconSvgElement } from "@hugeicons/react";
import { HugeiconsIcon } from "@hugeicons/react";
+import type { ReactNode } from "react";
import { useLayoutEffect, useRef, useState } from "react";
export function NetworkErrorState({
@@ -199,10 +200,12 @@ export function EmptyState({
title,
body,
icon = CubeIcon,
+ action,
}: {
title: string;
body: string;
icon?: IconSvgElement;
+ action?: ReactNode;
}) {
return (
@@ -217,6 +220,7 @@ export function EmptyState({
{body}
+ {action}
);
}
diff --git a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
index d407517663..b821be4b0b 100644
--- a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
+++ b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
@@ -152,11 +152,7 @@ export function DatasetDownloadSection({
/>
)}
{isDownloaded && cachePath && (
-
+
)}
diff --git a/studio/frontend/src/features/hub/catalog/download-section.tsx b/studio/frontend/src/features/hub/catalog/download-section.tsx
index fb2921649b..b2dd4592a1 100644
--- a/studio/frontend/src/features/hub/catalog/download-section.tsx
+++ b/studio/frontend/src/features/hub/catalog/download-section.tsx
@@ -1,9 +1,9 @@
// 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 { ModelInventoryFormat } from "../inventory";
import { GgufDownloadCard } from "./gguf-download-card";
import { SafetensorsDownloadCard } from "./safetensors-download-card";
-import type { ModelInventoryFormat } from "../inventory";
export function DownloadSection({
repoId,
@@ -22,6 +22,7 @@ export function DownloadSection({
knownBytes,
onLoad,
onUseInChat,
+ onEject,
onTrain,
onChange,
}: {
@@ -41,6 +42,7 @@ export function DownloadSection({
knownBytes?: number | null;
onLoad: (opts: { ggufVariant?: string; expectedBytes?: number }) => void;
onUseInChat?: () => void;
+ onEject?: () => void;
onTrain?: () => void;
onChange?: () => void;
}) {
@@ -57,6 +59,7 @@ export function DownloadSection({
isPartial={isPartial}
onLoad={onLoad}
onUseInChat={onUseInChat}
+ onEject={onEject}
onChange={onChange}
/>
);
@@ -75,6 +78,7 @@ export function DownloadSection({
knownBytes={knownBytes}
onLoad={onLoad}
onUseInChat={onUseInChat}
+ onEject={onEject}
onTrain={onTrain}
onChange={onChange}
/>
diff --git a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
index f31212e4fd..0d6878b687 100644
--- a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
@@ -1,6 +1,13 @@
// 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 {
+ DropdownMenu,
+ DropdownMenuContent,
+ DropdownMenuItem,
+ DropdownMenuSeparator,
+ DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
import {
Popover,
PopoverContent,
@@ -12,49 +19,55 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
-import {
- downloadManager,
- useDownloadManagerStore,
- useRepoDownload,
-} from "../download-manager";
-import {
- type GgufVariantDetail,
- deleteCachedModel,
-} from "../inventory";
-import { formatBytes } from "../lib/format";
-import { type GgufFitClass, classifyGgufFit } from "../lib/gguf-fit";
-import { HUB_GGUF_RUN_ACTIONS_VISIBLE } from "../lib/hub-feature-flags";
-import {
- ggufVariantsMatch,
- normalizeGgufVariantIdentity,
-} from "../lib/model-identity";
+import { usePlatformStore } from "@/config/env";
+import { getCachedModelPath, revealCachedModel } from "@/features/chat";
+import { pinKey, usePinnedModelsStore } from "@/features/model-picker";
+import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
+import { copyToClipboard } from "@/lib/copy-to-clipboard";
+import { toast } from "@/lib/toast";
import { cn } from "@/lib/utils";
-import { useHfTokenStore } from "../stores/hf-token-store";
-import { useOnlineStatus } from "../hooks/use-online-status";
import {
ArrowReloadHorizontalIcon,
+ Copy01Icon,
Delete02Icon,
Download01Icon,
+ Folder01Icon,
InformationCircleIcon,
- PencilEdit02Icon,
+ MoreVerticalIcon,
+ PinIcon,
+ PinOffIcon,
PlayIcon,
+ RemoveCircleIcon,
} from "@hugeicons/core-free-icons";
-import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import {
+ type KeyboardEventHandler,
memo,
useCallback,
useEffect,
useMemo,
useState,
- type KeyboardEventHandler,
- type MouseEventHandler,
} from "react";
+import {
+ downloadManager,
+ useDownloadManagerStore,
+ useRepoDownload,
+} from "../download-manager";
+import { useOnlineStatus } from "../hooks/use-online-status";
+import { type GgufVariantDetail, deleteCachedModel } from "../inventory";
+import { formatBytes } from "../lib/format";
+import { type GgufFitClass, classifyGgufFit } from "../lib/gguf-fit";
import {
ggufVariantDisplayLabel,
ggufVariantDownloadSizeBytes,
sortDownloadableGgufVariants,
} from "../lib/gguf-variant-sort";
+import { HUB_GGUF_RUN_ACTIONS_VISIBLE } from "../lib/hub-feature-flags";
+import {
+ ggufVariantsMatch,
+ normalizeGgufVariantIdentity,
+} from "../lib/model-identity";
+import { useHfTokenStore } from "../stores/hf-token-store";
import { DotTag } from "./dot-tag";
import { DownloadCancelIndicator } from "./download-cancel-indicator";
import {
@@ -72,7 +85,6 @@ import {
GgufDownloadStatusCard,
GgufDownloadingFallbackCard,
} from "./gguf-status-cards";
-import { PathInfoButton } from "./path-info-button";
import { useDeleteConfirmAction } from "./use-delete-confirm-action";
import { useDownloadCardState } from "./use-download-card-state";
import { useGgufVariantFetchState } from "./use-gguf-variant-fetch-state";
@@ -204,7 +216,7 @@ function QuantBadge({
onOpenChange={tooltipMode === "lazy" ? setTooltipOpen : undefined}
>
@@ -245,7 +257,181 @@ function createGgufVariantMenuItems(
}));
}
+// Shared options menu: used on every variant row, the run bar, and the
+// single-model (non-GGUF) run bar. Omit `quant` for a repo-level model. The
+// identifier uses llama.cpp's repo:quant syntax so it pastes into `-hf`.
+export function QuantOptionsMenu({
+ repoId,
+ quant,
+ label,
+ downloaded,
+ canDelete,
+ onDelete,
+ showPin = true,
+ buttonClassName,
+ iconClassName,
+}: {
+ repoId: string;
+ quant?: string;
+ label: string;
+ downloaded: boolean;
+ canDelete: boolean;
+ onDelete: (quant?: string) => void;
+ // Hidden in the run bar; pinning belongs to the On Device list.
+ showPin?: boolean;
+ buttonClassName?: string;
+ iconClassName?: string;
+}) {
+ const pinnedKeys = usePinnedModelsStore((s) => s.pinned);
+ const togglePinned = usePinnedModelsStore((s) => s.togglePinned);
+ const pinned = pinnedKeys.includes(pinKey(repoId, quant));
+ const deviceType = usePlatformStore((s) => s.deviceType);
+ const revealLabel =
+ deviceType === "mac"
+ ? "Reveal in Finder"
+ : deviceType === "windows"
+ ? "Reveal in File Explorer"
+ : "Reveal in File Manager";
+ const handleCopyPath = useCallback(async () => {
+ try {
+ const { path } = await getCachedModelPath(repoId, quant);
+ if (await copyToClipboard(path)) {
+ toast.success("Copied path");
+ } else {
+ toast.error("Failed to copy");
+ }
+ } catch (err) {
+ toast.error(
+ err instanceof Error ? err.message : "Failed to resolve model path",
+ );
+ }
+ }, [repoId, quant]);
+ const handleCopyId = useCallback(async () => {
+ const id = quant ? `${repoId}:${quant}` : repoId;
+ if (await copyToClipboard(id)) {
+ toast.success("Copied identifier");
+ } else {
+ toast.error("Failed to copy");
+ }
+ }, [repoId, quant]);
+
+ return (
+
+
+ e.stopPropagation()}
+ aria-label={`More options for ${label}`}
+ className={cn(
+ "inline-flex size-6 shrink-0 cursor-pointer items-center justify-center rounded-full",
+ "text-muted-foreground transition-colors hover:bg-muted hover:text-foreground",
+ "data-[state=open]:bg-muted data-[state=open]:text-foreground",
+ buttonClassName,
+ )}
+ >
+
+
+
+
+ {showPin && downloaded && (
+ {
+ e.stopPropagation();
+ togglePinned(repoId, quant);
+ }}
+ >
+
+ {pinned ? "Unpin" : "Pin to top"}
+
+ )}
+ {downloaded && (
+ {
+ e.stopPropagation();
+ revealCachedModel(repoId, quant).catch((err) => {
+ toast.error(
+ err instanceof Error
+ ? err.message
+ : "Failed to open file manager",
+ );
+ });
+ }}
+ >
+
+ {revealLabel}
+
+ )}
+ {
+ e.stopPropagation();
+ void handleCopyId();
+ }}
+ >
+
+ Copy identifier
+
+ {downloaded && (
+ {
+ e.stopPropagation();
+ void handleCopyPath();
+ }}
+ >
+
+ Copy path
+
+ )}
+ {canDelete && (
+ <>
+
+ {
+ e.stopPropagation();
+ onDelete(quant);
+ }}
+ >
+
+ Delete
+
+ >
+ )}
+
+
+ );
+}
+
const GgufVariantMenuRow = memo(function GgufVariantMenuRow({
+ repoId,
item,
selected,
loaded,
@@ -254,6 +440,7 @@ const GgufVariantMenuRow = memo(function GgufVariantMenuRow({
onSelect,
onDelete,
}: {
+ repoId: string;
item: GgufVariantMenuItem;
selected: boolean;
loaded: boolean;
@@ -275,13 +462,6 @@ const GgufVariantMenuRow = memo(function GgufVariantMenuRow({
},
[selectVariant],
);
- const handleDelete = useCallback>(
- (e) => {
- e.stopPropagation();
- onDelete(item.quant);
- },
- [item.quant, onDelete],
- );
const canDelete = (item.downloaded || item.partial) && !loaded && !liveActive;
return (
@@ -317,7 +497,7 @@ const GgufVariantMenuRow = memo(function GgufVariantMenuRow({
)}
{!item.downloaded && item.partial && (
-
+
-
-
- {item.downloadSizeLabel}
-
- {canDelete && (
-
-
-
- )}
+
+ {item.downloadSizeLabel}
+ {/* Options only apply to files on disk; placeholder keeps the size
+ chips column-aligned across rows. */}
+ {item.downloaded || item.partial ? (
+ q && onDelete(q)}
+ />
+ ) : (
+
+ )}
);
@@ -374,7 +547,7 @@ export function GgufDownloadCard({
preferLocalCache = false,
isPartial = false,
onLoad,
- onUseInChat,
+ onEject,
onChange,
}: {
repoId: string;
@@ -387,7 +560,9 @@ export function GgufDownloadCard({
preferLocalCache?: boolean;
isPartial?: boolean;
onLoad: (opts: { ggufVariant?: string; expectedBytes?: number }) => void;
+ /** Accepted for API parity; the run bar ejects instead of opening chat. */
onUseInChat?: () => void;
+ onEject?: () => void;
onChange?: () => void;
}) {
const hfToken = useHfTokenStore((s) => s.token);
@@ -422,7 +597,9 @@ export function GgufDownloadCard({
() => createLiveGgufVariantStatesSelector(repoId),
[repoId],
);
- const liveVariantStates = useDownloadManagerStore(selectLiveGgufVariantStates);
+ const liveVariantStates = useDownloadManagerStore(
+ selectLiveGgufVariantStates,
+ );
const sortedVariants = useMemo(() => {
if (!rawSortedVariants) return null;
const withLive = applyLiveGgufVariantStates(
@@ -499,12 +676,7 @@ export function GgufDownloadCard({
if (expectedBytes > progress.expectedBytes) {
setExpectedBytes(expectedBytes, progress.variant);
}
- }, [
- variants,
- progress?.variant,
- progress?.expectedBytes,
- setExpectedBytes,
- ]);
+ }, [variants, progress?.variant, progress?.expectedBytes, setExpectedBytes]);
useEffect(() => {
setCompletedVariantKeys(new Set());
@@ -595,7 +767,8 @@ export function GgufDownloadCard({
if (!deleteTarget) return;
await deleteCachedModel(repoId, deleteTarget, hfToken || undefined);
},
- successMessage: () => `Deleted ${repoId} ${deleteTargetLabel ?? deleteTarget}`,
+ successMessage: () =>
+ `Deleted ${repoId} ${deleteTargetLabel ?? deleteTarget}`,
errorToast: (err) => ({
title: err instanceof Error ? err.message : "Failed to delete",
}),
@@ -654,7 +827,7 @@ export function GgufDownloadCard({
return (
);
@@ -666,7 +839,7 @@ export function GgufDownloadCard({
void refresh()}
@@ -729,7 +902,7 @@ export function GgufDownloadCard({
}
>
-
+
{
@@ -762,7 +935,7 @@ export function GgufDownloadCard({
)}
{selected && !selected.downloaded && selected.partial && (
-
+
- {selected?.downloaded && cachePath && (
-
- )}
+ {/* TODO: inference settings gear hidden for now, work on it in a future PR. */}
+ {/* Options only resolve managed HF-cache repos, so skip local paths;
+ they also only apply to quants actually on disk. */}
+ {selected &&
+ Boolean(selected.downloaded || selected.partial) &&
+ !/^([/\\~.]|[A-Za-z]:)/.test(repoId) && (
+ q && handleDeleteVariant(q)}
+ showPin={false}
+ buttonClassName="ml-0.5 size-7"
+ iconClassName="size-4"
+ />
+ )}
{!isGgufRunCta && }
@@ -859,7 +1048,7 @@ export function GgufDownloadCard({
return;
}
if (selectedIsActive) {
- onUseInChat?.();
+ onEject?.();
return;
}
if (!selected) return;
@@ -874,7 +1063,7 @@ export function GgufDownloadCard({
}}
aria-label={downloadAction.ariaLabel}
className={cn(
- isGgufRunCta ? "hub-run-action-btn w-28" : "hub-action-btn w-28",
+ isGgufRunCta ? "hub-run-action-btn w-24" : "hub-action-btn w-24",
isGgufRunCta && "ml-2",
ctaDisabled &&
!selectedIsActive &&
@@ -917,8 +1106,8 @@ export function GgufDownloadCard({
) : selectedIsActive ? (
<>
-
- New Chat
+
+ Eject
>
) : selected?.downloaded ? (
<>
diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx
index 38464d36e4..df30888d57 100644
--- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx
+++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx
@@ -6,9 +6,9 @@ import {
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
-import { cn } from "@/lib/utils";
-import { Tick02Icon } from "@/lib/tick-icon";
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
+import { Tick02Icon } from "@/lib/tick-icon";
+import { cn } from "@/lib/utils";
import { HugeiconsIcon } from "@hugeicons/react";
import {
type KeyboardEvent,
@@ -71,7 +71,9 @@ export function HubOptionMenu
({
? -1
: Math.min(activeIndex, options.length - 1);
const activeOptionId =
- resolvedActiveIndex >= 0 ? `${idBase}-option-${resolvedActiveIndex}` : undefined;
+ resolvedActiveIndex >= 0
+ ? `${idBase}-option-${resolvedActiveIndex}`
+ : undefined;
const activateIndex = useCallback((index: number) => {
setActiveIndex((current) => (current === index ? current : index));
@@ -161,7 +163,7 @@ export function HubOptionMenu({
return (
-
+
({
aria-label={ariaLabel}
title={title}
className={cn(
- "field-trigger hub-menu-trigger field-soft field-filter inline-flex h-9 shrink-0 cursor-pointer items-center justify-between gap-0.5 rounded-full pl-3 pr-2 text-[12.5px] transition-colors",
+ "field-trigger hub-menu-trigger field-soft field-filter inline-flex h-9 shrink-0 cursor-pointer items-center justify-between gap-2.5 rounded-full pl-3 pr-2.5 text-[12.5px] transition-colors",
"focus-visible:border-border focus-visible:ring-0 focus-visible:ring-offset-0",
className,
)}
@@ -183,7 +185,10 @@ export function HubOptionMenu({
}}
>
- {triggerContent ?? selected?.triggerLabel ?? selected?.label ?? value}
+ {triggerContent ??
+ selected?.triggerLabel ??
+ selected?.label ??
+ value}
{showChevron && (
{HUB_POST_DOWNLOAD_ACTIONS_VISIBLE && (
diff --git a/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx b/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx
index 3c020a7199..ba97cb0c53 100644
--- a/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx
@@ -1,12 +1,6 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import { TrainIcon } from "../components/train-icon";
-import {
- HUB_GGUF_RUN_ACTIONS_VISIBLE,
- HUB_NON_GGUF_RUN_ACTIONS_VISIBLE,
- HUB_POST_DOWNLOAD_ACTIONS_VISIBLE,
-} from "../lib/hub-feature-flags";
import {
Popover,
PopoverContent,
@@ -18,37 +12,44 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
+import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
+import { cn } from "@/lib/utils";
import {
- type BaseModelSource,
- type LocalModelInfo,
- type ModelInventoryFormat,
- deleteCachedModel,
-} from "../inventory";
+ Alert02Icon,
+ CubeIcon,
+ PlayIcon,
+ RemoveCircleIcon,
+ Share05Icon,
+} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { useCallback, useMemo, useState } from "react";
+import { TrainIcon } from "../components/train-icon";
import {
downloadManager,
jobKeyOf,
selectActiveJob,
useDownloadManagerStore,
} from "../download-manager";
-import { formatBytes } from "../lib/format";
-import { ggufVariantsMatch } from "../lib/model-identity";
-import { cn } from "@/lib/utils";
-import { confirmExternalLink } from "../stores/external-link-confirm";
-import { useHfTokenStore } from "../stores/hf-token-store";
+import { useOnlineStatus } from "../hooks/use-online-status";
import {
- Alert02Icon,
- CubeIcon,
- PencilEdit02Icon,
- PlayIcon,
- Share05Icon,
-} from "@hugeicons/core-free-icons";
-import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
-import { HugeiconsIcon } from "@hugeicons/react";
-import { useCallback, useMemo, useState } from "react";
+ type BaseModelSource,
+ type LocalModelInfo,
+ type ModelInventoryFormat,
+ deleteCachedModel,
+} from "../inventory";
+import { formatBytes } from "../lib/format";
import {
ggufVariantDisplayLabel,
sortLocalGgufVariants,
} from "../lib/gguf-variant-sort";
+import {
+ HUB_GGUF_RUN_ACTIONS_VISIBLE,
+ HUB_NON_GGUF_RUN_ACTIONS_VISIBLE,
+ HUB_POST_DOWNLOAD_ACTIONS_VISIBLE,
+} from "../lib/hub-feature-flags";
+import { ggufVariantsMatch } from "../lib/model-identity";
+import { confirmExternalLink } from "../stores/external-link-confirm";
+import { useHfTokenStore } from "../stores/hf-token-store";
import { DotTag } from "./dot-tag";
import {
CardDeleteButton,
@@ -60,7 +61,6 @@ import { PathInfoButton } from "./path-info-button";
import { TransportConflictDialog } from "./transport-conflict-dialog";
import { useCardDelete } from "./use-card-delete";
import { useGgufVariantFetchState } from "./use-gguf-variant-fetch-state";
-import { useOnlineStatus } from "../hooks/use-online-status";
type LocalLoadOptions = {
ggufVariant?: string;
@@ -91,7 +91,9 @@ interface LocalOnDeviceCardProps {
systemRamGb?: number;
unsupportedReason?: string | null;
onLoad: (opts?: LocalLoadOptions) => void;
+ /** Accepted for API parity; the run bar ejects instead of opening chat. */
onUseInChat: () => void;
+ onEject?: () => void;
onTrain?: () => void;
onChange?: () => void;
}
@@ -151,7 +153,7 @@ function BaseModelReference({
{canOpenHub && (
-
+
{
event.stopPropagation();
- if (confirmExternalLink(`https://huggingface.co/${baseModelHubId}`)) {
+ if (
+ confirmExternalLink(
+ `https://huggingface.co/${baseModelHubId}`,
+ )
+ ) {
event.preventDefault();
}
}}
@@ -205,7 +211,7 @@ export function LocalOnDeviceCard({
systemRamGb,
unsupportedReason,
onLoad,
- onUseInChat,
+ onEject,
onTrain,
onChange,
}: LocalOnDeviceCardProps) {
@@ -371,18 +377,20 @@ export function LocalOnDeviceCard({
const handleConfirmUpdate = () => {
if (!repoId || !updateTargetVariant) return;
setUpdateOpen(false);
- void downloadManager.requestStart({
- kind: "model",
- repoId,
- variant: updateTargetVariant,
- expectedBytes: updateExpectedBytes,
- }).then((outcome) => {
- if (outcome === "conflict") {
- setUpdateConflictKey(jobKeyOf("model", repoId, updateTargetVariant));
- }
- void currentVariantState.refresh();
- void remoteVariantState.refresh();
- });
+ void downloadManager
+ .requestStart({
+ kind: "model",
+ repoId,
+ variant: updateTargetVariant,
+ expectedBytes: updateExpectedBytes,
+ })
+ .then((outcome) => {
+ if (outcome === "conflict") {
+ setUpdateConflictKey(jobKeyOf("model", repoId, updateTargetVariant));
+ }
+ void currentVariantState.refresh();
+ void remoteVariantState.refresh();
+ });
};
const selectedVariantIsActive =
needsVariantSelection && selectedQuant
@@ -428,7 +436,8 @@ export function LocalOnDeviceCard({
can still keep it on disk, or delete it to free space.
- )}
+ )}
+
@@ -540,7 +549,7 @@ export function LocalOnDeviceCard({
{canUpdate && (
setUpdateOpen(true)}
/>
)}
@@ -550,11 +559,7 @@ export function LocalOnDeviceCard({
onClick={() => setDeleteOpen(true)}
/>
)}
-
+
{onTrain && HUB_POST_DOWNLOAD_ACTIONS_VISIBLE && (
@@ -585,7 +590,7 @@ export function LocalOnDeviceCard({
onClick={() => {
if (!canRun) return;
if (selectedVariantIsActive) {
- onUseInChat();
+ onEject?.();
return;
}
if (needsVariantSelection) {
@@ -615,24 +620,24 @@ export function LocalOnDeviceCard({
>
) : selectedVariantIsActive ? (
<>
-
- Chat
+
+ Eject
>
) : variantActionPending ? (
<>
Loading…
>
- ) : !canRun ? (
- <>
-
- No run
- >
- ) : (
+ ) : canRun ? (
<>
Run
>
+ ) : (
+ <>
+
+ No run
+ >
)}
diff --git a/studio/frontend/src/features/hub/catalog/model-inspector.tsx b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
index 5a6ae1615a..90663b5d76 100644
--- a/studio/frontend/src/features/hub/catalog/model-inspector.tsx
+++ b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
@@ -17,9 +17,9 @@ import {
formatRelativeShort,
formatShortDate,
} from "@/features/hub/lib/format";
-import { cn, formatCompact } from "@/lib/utils";
-import { confirmExternalLink } from "../stores/external-link-confirm";
import { useHfTokenStore } from "@/features/hub/stores/hf-token-store";
+import { Tick02Icon } from "@/lib/tick-icon";
+import { cn, formatCompact } from "@/lib/utils";
import {
Calendar03Icon,
CalendarAdd01Icon,
@@ -37,10 +37,10 @@ import {
RamMemoryIcon,
Share05Icon,
} from "@hugeicons/core-free-icons";
-import { Tick02Icon } from "@/lib/tick-icon";
import type { IconSvgElement } from "@hugeicons/react";
import { HugeiconsIcon } from "@hugeicons/react";
import { memo, useDeferredValue, useMemo } from "react";
+import { selectActiveJob, useDownloadManagerStore } from "../download-manager";
import { useCopyFeedback } from "../hooks/use-copy-feedback";
import { useDatasetSize } from "../hooks/use-dataset-size";
import {
@@ -49,8 +49,8 @@ import {
formatPipelineTag,
parseLanguageTags,
} from "../lib/view-models";
+import { confirmExternalLink } from "../stores/external-link-confirm";
import type { SelectedModelView } from "../types";
-import { selectActiveJob, useDownloadManagerStore } from "../download-manager";
import { DatasetDownloadSection } from "./dataset-download-section";
import { DownloadSection } from "./download-section";
import { LocalDatasetCard } from "./local-dataset-card";
@@ -399,6 +399,7 @@ export type ModelInspectorActions = {
expectedBytes?: number;
}) => void;
onUseInChat: () => void;
+ onEject?: () => void;
onTrain?: () => void;
onInventoryChange?: () => void;
onSearchHub?: (query: string) => void;
@@ -433,6 +434,7 @@ export const ModelInspector = memo(function ModelInspector({
onLoad,
onLoadLocal,
onUseInChat,
+ onEject,
onTrain,
onInventoryChange,
onSearchHub,
@@ -517,7 +519,9 @@ export const ModelInspector = memo(function ModelInspector({
? formatRelativeShort(model.updatedAt)
: formatLocalUpdated(model.localUpdatedAt);
const updatedLabel = updatedRaw === "Unknown update" ? "N/A" : updatedRaw;
- const createdLabel = model.createdAt ? formatShortDate(model.createdAt) : null;
+ const createdLabel = model.createdAt
+ ? formatShortDate(model.createdAt)
+ : null;
const libraryLabel = isDataset ? null : formatLibrary(model.libraryName);
const gatedAccess = model.gated !== false && model.gated !== undefined;
const downloadsTooltip =
@@ -696,6 +700,7 @@ export const ModelInspector = memo(function ModelInspector({
}
onLoad={onLoadLocal}
onUseInChat={onUseInChat}
+ onEject={onEject}
onTrain={
model.isDownloaded && canTrainModel ? onTrain : undefined
}
@@ -719,6 +724,7 @@ export const ModelInspector = memo(function ModelInspector({
knownBytes={model.cachedBytes}
onLoad={model.isLocal ? onLoadLocal : onLoad}
onUseInChat={onUseInChat}
+ onEject={onEject}
onTrain={
model.isDownloaded && canTrainModel ? onTrain : undefined
}
diff --git a/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx b/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx
index 355c38a91e..8cf5fc491e 100644
--- a/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx
@@ -2,13 +2,20 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Spinner } from "@/components/ui/spinner";
+import {
+ makePinRank,
+ pinKey,
+ usePinnedModelsStore,
+} from "@/features/model-picker";
import {
CubeIcon,
DownloadCircle02Icon,
- FolderSearchIcon,
+ PinIcon,
+ Search01Icon,
} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
import type { RefObject } from "react";
-import { useMemo } from "react";
+import { useLayoutEffect, useMemo, useState } from "react";
import {
inventoryRowMatches,
scoreInventoryRow,
@@ -208,7 +215,7 @@ export function DiscoverList({
) : (
void;
scrollElement: HTMLDivElement | null;
columns?: number;
activeCheckpoint: string | null;
@@ -274,11 +285,20 @@ export function DownloadedList({
sort: InventorySort;
onInventoryChange?: () => void;
}) {
+ // Pinned repos surface first regardless of the active sort; the chosen sort
+ // still orders rows within the pinned and unpinned groups.
+ const pinnedIds = usePinnedModelsStore((s) => s.pinned);
+ const pinnedSet = useMemo(() => new Set(pinnedIds), [pinnedIds]);
const inventoryItems = useMemo(() => {
const merged: InventoryItem[] = [
...cachedRows.map((row) => ({ variant: "cached" as const, row })),
...localRows.map((row) => ({ variant: "local" as const, row })),
];
+ // Pinned rows order by pin recency (newest pin first), not the active
+ // sort, so "Pin to top" puts the row exactly where the user expects.
+ const rank = makePinRank(pinnedIds);
+ const pinRank = (item: InventoryItem) =>
+ item.row.repoId ? rank(pinKey(item.row.repoId)) : Number.MAX_SAFE_INTEGER;
if (inventoryTokens.length > 0) {
return merged
.map((item, index) => ({
@@ -286,25 +306,85 @@ export function DownloadedList({
index,
score: scoreInventoryRow(item.row, inventoryTokens),
}))
- .sort((a, b) => b.score - a.score || a.index - b.index)
+ .sort(
+ (a, b) =>
+ pinRank(a.item) - pinRank(b.item) ||
+ b.score - a.score ||
+ a.index - b.index,
+ )
.map((entry) => entry.item);
}
if (sort === "recent") {
- return merged;
+ return merged
+ .map((item, index) => ({ item, index }))
+ .sort((a, b) => pinRank(a.item) - pinRank(b.item) || a.index - b.index)
+ .map((entry) => entry.item);
}
return merged
.map((item, index) => ({ item, index }))
- .sort((a, b) =>
- sort === "name"
- ? inventoryItemTitle(a.item).localeCompare(
- inventoryItemTitle(b.item),
- ) || a.index - b.index
- : inventoryItemSize(b.item) - inventoryItemSize(a.item) ||
- a.index - b.index,
+ .sort(
+ (a, b) =>
+ pinRank(a.item) - pinRank(b.item) ||
+ (sort === "name"
+ ? inventoryItemTitle(a.item).localeCompare(
+ inventoryItemTitle(b.item),
+ ) || a.index - b.index
+ : inventoryItemSize(b.item) - inventoryItemSize(a.item) ||
+ a.index - b.index),
)
.map((entry) => entry.item);
- }, [cachedRows, localRows, inventoryTokens, sort]);
+ }, [cachedRows, localRows, inventoryTokens, sort, pinnedIds]);
const hasInventoryRows = cachedRows.length > 0 || localRows.length > 0;
+ // Pinned repos get their own labelled section so it's clear why they lead
+ // the list; inventoryItems already sorts them first, so this is a prefix.
+ const pinnedCount = useMemo(
+ () =>
+ inventoryItems.filter(
+ (item) => item.row.repoId && pinnedSet.has(pinKey(item.row.repoId)),
+ ).length,
+ [inventoryItems, pinnedSet],
+ );
+ const pinnedItems = inventoryItems.slice(0, pinnedCount);
+ const unpinnedItems = inventoryItems.slice(pinnedCount);
+ const [virtualRowsWrapper, setVirtualRowsWrapper] =
+ useState(null);
+ const [scrollMargin, setScrollMargin] = useState(0);
+ useLayoutEffect(() => {
+ if (!virtualRowsWrapper || !scrollElement) return;
+ const measure = () => {
+ const margin = Math.max(
+ 0,
+ Math.round(
+ virtualRowsWrapper.getBoundingClientRect().top -
+ scrollElement.getBoundingClientRect().top +
+ scrollElement.scrollTop,
+ ),
+ );
+ setScrollMargin((current) => (current === margin ? current : margin));
+ };
+ measure();
+ const observer = new ResizeObserver(measure);
+ observer.observe(virtualRowsWrapper.parentElement ?? scrollElement);
+ return () => observer.disconnect();
+ }, [virtualRowsWrapper, scrollElement]);
+ const rowHeightPx = compact
+ ? RESULT_SPLIT_ROW_HEIGHT_PX
+ : RESULT_GRID_ROW_HEIGHT_PX;
+ const cellHeightPx = compact ? RESULT_SPLIT_HEIGHT_PX : RESULT_GRID_HEIGHT_PX;
+ const renderInventoryRow = (item: InventoryItem) => (
+
+ );
if (!downloadedReady && !hasInventoryRows) {
return (
@@ -325,9 +405,29 @@ export function DownloadedList({
}
if (cachedRows.length === 0 && localRows.length === 0) {
+ if (!query.trim() && typeFilterActive) {
+ return (
+
+ Show all types
+
+ )
+ }
+ />
+ );
+ }
return (
`${item.variant}-${item.row.id}`}
- renderRow={(item) => (
-
+ <>
+ {pinnedItems.length > 0 && (
+ <>
+
+
+ Pinned
+
+ {/* Pinned rows are few, so render them as a plain grid matching the
+ virtualized list's lane count and row spacing. */}
+
+ {pinnedItems.map((item) => (
+
+ {renderInventoryRow(item)}
+
+ ))}
+
+ {unpinnedItems.length > 0 && (
+
+ All {isDataset ? "datasets" : "models"}
+
+ )}
+ >
)}
- />
+
+ `${item.variant}-${item.row.id}`}
+ renderRow={renderInventoryRow}
+ />
+
+ >
);
}
diff --git a/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx b/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx
index 111d78f00b..6d1dc20414 100644
--- a/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx
@@ -1,7 +1,6 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import { ModelDeleteAction } from "@/components/assistant-ui/model-selector/model-delete-action";
import {
Tooltip,
TooltipContent,
@@ -9,18 +8,26 @@ import {
} from "@/components/ui/tooltip";
import {
type GgufVariantDetail,
- deleteCachedModel,
deleteCachedDataset,
+ deleteCachedModel,
formatLocalUpdated,
listGgufVariants,
useGgufVariantsCacheVersion,
-} from "@/features/hub/inventory";
-import { classifyUnslothSupport } from "@/features/hub/hooks/use-hub-model-search";
-import { formatBytes, formatRelativeShort } from "@/features/hub/lib/format";
-import { ggufVariantDisplayLabel } from "@/features/hub/lib/gguf-variant-sort";
-import { modelIdsMatch } from "@/features/hub/lib/model-identity";
+} from "../inventory";
+import {
+ classifyUnslothSupport,
+ formatBytes,
+ formatRelativeShort,
+ ggufVariantDisplayLabel,
+ useHfTokenStore,
+} from "@/features/hub";
+import { modelIdsMatch } from "../lib/model-identity";
+import {
+ ModelRowMenu,
+ pinKey,
+ usePinnedModelsStore,
+} from "@/features/model-picker";
import { cn, formatCompact } from "@/lib/utils";
-import { useHfTokenStore } from "@/features/hub/stores/hf-token-store";
import {
Download01Icon,
FavouriteIcon,
@@ -39,6 +46,7 @@ import {
useRef,
useState,
} from "react";
+import { paramLabelFromId } from "../lib/view-models";
import type {
CachedInventoryRow,
DiscoverRow,
@@ -46,7 +54,6 @@ import type {
} from "../types";
import { OwnerAvatar } from "./owner-avatar";
import { AccessGlyphs } from "./shared";
-import { paramLabelFromId } from "../lib/view-models";
const COARSE_POINTER =
typeof window !== "undefined" &&
@@ -142,15 +149,15 @@ function CachedSizeChipLive({
);
const rows: Array<{ label: string; size_bytes: number }> | null =
- !needsVariantFetch
- ? [{ label: repoId, size_bytes: totalBytes }]
- : currentVariantState.status === "loaded" &&
- currentVariantState.variants.length > 0
+ needsVariantFetch
+ ? currentVariantState.status === "loaded" &&
+ currentVariantState.variants.length > 0
? currentVariantState.variants.map((variant) => ({
label: ggufVariantDisplayLabel(variant),
size_bytes: variant.size_bytes,
}))
- : null;
+ : null
+ : [{ label: repoId, size_bytes: totalBytes }];
const variantMessage =
currentVariantState.status === "loading"
? "Loading downloaded variants..."
@@ -275,7 +282,9 @@ function CatalogRow({
)}
/>
-
+
{children}
@@ -653,7 +662,9 @@ export const InventoryRow = memo(function InventoryRow({
{/* Format already shows as the status dot, so the pill stays neutral. */}
{formatLabel &&
{formatLabel} }
- {paramLabel &&
{paramLabel} }
+ {paramLabel && (
+
{paramLabel}
+ )}
{quantLabel && (
{quantLabel}
@@ -697,9 +708,7 @@ export const InventoryRow = memo(function InventoryRow({
const compactMarkers =
partialRepoId || unsupported ? (
- {partialRepoId && (
-
- )}
+ {partialRepoId && }
{unsupported && (
)}
@@ -718,37 +727,72 @@ export const InventoryRow = memo(function InventoryRow({
);
+ const pinnedKeys = usePinnedModelsStore((s) => s.pinned);
+ const togglePinned = usePinnedModelsStore((s) => s.togglePinned);
+ const rowPinned =
+ cacheDeletableRepoId != null &&
+ pinnedKeys.includes(pinKey(cacheDeletableRepoId));
const deleteAction =
canDelete && cacheDeletableRepoId ? (
-
- This will remove{" "}
-
- {cacheDeletableRepoId}
- {" "}
- {isDataset
- ? "and its downloaded files"
- : row.isGguf
- ? "and all of its downloaded quantizations"
- : "and all of its downloaded files"}
- {row.kind === "cache" ? ` (${formatBytes(row.bytes)})` : ""} from
- disk. You can re-download it later.
- >
- }
- successMessage={`Deleted ${cacheDeletableRepoId}`}
+ {
- if (isDataset) {
- await deleteCachedDataset(cacheDeletableRepoId);
- } else {
- await deleteCachedModel(cacheDeletableRepoId);
- }
+ pin={
+ isDataset
+ ? undefined
+ : {
+ pinned: rowPinned,
+ pinLabel: "Pin to top",
+ unpinLabel: "Unpin",
+ onToggle: () => togglePinned(cacheDeletableRepoId),
+ }
+ }
+ cachePath={isDataset ? undefined : { repoId: cacheDeletableRepoId }}
+ del={{
+ title: isDataset ? "Delete cached dataset?" : "Delete cached model?",
+ description: (
+ <>
+ This will remove{" "}
+
+ {cacheDeletableRepoId}
+ {" "}
+ {isDataset
+ ? "and its downloaded files"
+ : row.isGguf
+ ? "and all of its downloaded quantizations"
+ : "and all of its downloaded files"}
+ {row.kind === "cache" ? ` (${formatBytes(row.bytes)})` : ""} from
+ disk. You can re-download it later.
+ >
+ ),
+ successMessage: `Deleted ${cacheDeletableRepoId}`,
+ onConfirm: async () => {
+ if (isDataset) {
+ await deleteCachedDataset(cacheDeletableRepoId);
+ } else {
+ await deleteCachedModel(cacheDeletableRepoId);
+ // Deleted repos can't stay pinned: drop the repo pin and any of
+ // its per-quant pins so stale rows don't linger up top.
+ const { pinned, togglePinned: toggle } =
+ usePinnedModelsStore.getState();
+ for (const key of pinned) {
+ if (
+ key === pinKey(cacheDeletableRepoId) ||
+ key.startsWith(`${cacheDeletableRepoId}::`)
+ ) {
+ toggle(
+ cacheDeletableRepoId,
+ key.includes("::")
+ ? key.slice(key.indexOf("::") + 2)
+ : undefined,
+ );
+ }
+ }
+ }
+ },
+ onDeleted: onChange,
}}
- onDeleted={onChange}
/>
) : null;
diff --git a/studio/frontend/src/features/hub/catalog/models-catalog.tsx b/studio/frontend/src/features/hub/catalog/models-catalog.tsx
index 344f5deb49..02b9b62fce 100644
--- a/studio/frontend/src/features/hub/catalog/models-catalog.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-catalog.tsx
@@ -54,6 +54,7 @@ export interface ModelsCatalogState {
hasMore: boolean;
manualFetchAvailable: boolean;
hasActiveFilters: boolean;
+ typeFilterActive: boolean;
}
export interface ModelsCatalogPagination {
@@ -117,6 +118,7 @@ export const ModelsCatalog = memo(function ModelsCatalog({
loadingIntentCount,
hasMore,
hasActiveFilters,
+ typeFilterActive,
} = state;
const { scrollRef, sentinelRef, isLoadingMore } = pagination;
const {
@@ -469,6 +471,8 @@ export const ModelsCatalog = memo(function ModelsCatalog({
downloadedReady={downloadedReady}
inventoryError={inventoryError}
query={query}
+ typeFilterActive={typeFilterActive}
+ onClearFilters={onClearFilters}
scrollElement={downloadedScrollEl}
activeCheckpoint={activeCheckpoint}
activeGgufVariant={activeGgufVariant}
diff --git a/studio/frontend/src/features/hub/catalog/models-table.tsx b/studio/frontend/src/features/hub/catalog/models-table.tsx
index f3887fee28..3f91685b47 100644
--- a/studio/frontend/src/features/hub/catalog/models-table.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-table.tsx
@@ -17,6 +17,10 @@ import {
formatRelativeLong,
formatRelativeShort,
} from "@/features/hub/lib/format";
+import {
+ MODEL_TYPE_FILTER_OPTIONS,
+ type ModelTypeFilter,
+} from "@/features/hub/lib/model-type-filter";
import {
formatModelParamLabel,
formatPipelineTag,
@@ -25,6 +29,7 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { cn, formatCompact } from "@/lib/utils";
import {
ArrowLeft01Icon,
+ ArrowUpDownIcon,
Copy01Icon,
Download01Icon,
FavouriteIcon,
@@ -124,6 +129,7 @@ export function InventorySortControl({
value: InventorySort;
onChange: (value: InventorySort) => void;
}) {
+ const selected = INVENTORY_SORTS.find((option) => option.value === value);
return (
value={value}
@@ -131,7 +137,46 @@ export function InventorySortControl({
onValueChange={onChange}
ariaLabel="Sort downloads"
align="end"
- className="h-8 text-[11.5px]"
+ title={selected?.label}
+ // Capped and shrinkable so a long label truncates instead of wrapping
+ // the "On device" heading beside these pills in the narrow split pane.
+ className="h-8 min-w-[72px] max-w-[124px] shrink text-[11.5px]"
+ triggerContent={
+
+
+ {selected?.label ?? value}
+
+ }
+ />
+ );
+}
+
+// Model-type filter pill (Text / Vision / Embedding / …) beside the sort pill.
+export function InventoryTypeFilterControl({
+ value,
+ onChange,
+}: {
+ value: ModelTypeFilter;
+ onChange: (value: ModelTypeFilter) => void;
+}) {
+ const selected = MODEL_TYPE_FILTER_OPTIONS.find(
+ (option) => option.value === value,
+ );
+ return (
+
+ value={value}
+ options={MODEL_TYPE_FILTER_OPTIONS}
+ onValueChange={onChange}
+ ariaLabel="Filter by model type"
+ align="end"
+ title={selected?.label}
+ // Capped and shrinkable so a long label ("Speech to text") truncates
+ // instead of wrapping the "On device" heading beside these pills.
+ className="h-8 min-w-[72px] max-w-[124px] shrink text-[11.5px]"
/>
);
}
@@ -183,7 +228,9 @@ export function HubListHeader({
)}
-
+ {/* truncate keeps the heading on one line and clips a long search
+ query with an ellipsis instead of overflowing the pills. */}
+
{title}
{subtitle && (
@@ -216,7 +263,9 @@ export function HubListHeader({
)}
{(actions || onViewChange) && (
-
+ // min-w-0 (not shrink-0) so shrinkable actions (the On-device filter
+ // pills) compress before the title is forced onto two lines.
+
{actions}
{onViewChange && (
{tab === "downloaded" && !isDataset && (
-
+
-
+
Only show models that fit
@@ -402,7 +402,7 @@ export const ModelsToolbar = memo(function ModelsToolbar({
)}
/>
-
+
-
+
{
- const withoutDuplicate = current.filter((row) => row.id !== folder.id);
+ const withoutDuplicate = current.filter(
+ (row) => row.id !== folder.id,
+ );
return [...withoutDuplicate, folder];
});
toast.success("Location added", {
@@ -184,9 +186,12 @@ export function OnDeviceFoldersDialog({
overlayClassName="bg-black/20 backdrop-blur-none"
>
- On-device locations
+
+ On-device locations
+
- Hugging Face model folders, GGUF files, and adapters are indexed here.
+ Hugging Face model folders, GGUF files, and adapters are indexed
+ here.
@@ -342,9 +347,7 @@ export function OnDeviceFoldersDialog({
-
+
{folder.path}
@@ -372,7 +375,10 @@ export function OnDeviceFoldersDialog({
/>
-
+
Open in file manager
@@ -397,7 +403,10 @@ export function OnDeviceFoldersDialog({
)}
-
+
Remove from list
diff --git a/studio/frontend/src/features/hub/catalog/path-info-button.tsx b/studio/frontend/src/features/hub/catalog/path-info-button.tsx
index 0d5ca6e00e..403bb0148e 100644
--- a/studio/frontend/src/features/hub/catalog/path-info-button.tsx
+++ b/studio/frontend/src/features/hub/catalog/path-info-button.tsx
@@ -1,38 +1,83 @@
// 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 {
- Dialog,
- DialogContent,
- DialogDescription,
- DialogHeader,
- DialogTitle,
-} from "@/components/ui/dialog";
import {
Tooltip,
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
+import { usePlatformStore } from "@/config/env";
+import { revealCachedModel } from "@/features/chat";
+import { toast } from "@/lib/toast";
import { cn } from "@/lib/utils";
-import { Copy01Icon, FolderSearchIcon } from "@hugeicons/core-free-icons";
+import { Copy01Icon, Folder01Icon } from "@hugeicons/core-free-icons";
import { Tick02Icon } from "@/lib/tick-icon";
import { HugeiconsIcon } from "@hugeicons/react";
import type { MouseEvent } from "react";
-import { useState } from "react";
import { useCopyFeedback } from "../hooks/use-copy-feedback";
+/** Reveal a cached repo (or one GGUF variant's file) in the OS file manager.
+ * Resolved server-side from the HF cache, so only managed repos qualify. */
+export function RevealPathButton({
+ repoId,
+ variant,
+ className,
+}: {
+ repoId: string;
+ variant?: string | null;
+ className?: string;
+}) {
+ const deviceType = usePlatformStore((s) => s.deviceType);
+ const revealLabel =
+ deviceType === "mac"
+ ? "Reveal in Finder"
+ : deviceType === "windows"
+ ? "Reveal in File Explorer"
+ : "Reveal in File Manager";
+
+ return (
+
+
+ {
+ e.stopPropagation();
+ revealCachedModel(repoId, variant ?? undefined).catch((err) => {
+ toast.error(
+ err instanceof Error
+ ? err.message
+ : "Failed to open file manager",
+ );
+ });
+ }}
+ className={cn(
+ "inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-muted hover:text-foreground focus-visible:opacity-100 group-hover/dl:opacity-100",
+ className,
+ )}
+ >
+
+
+
+
+ {revealLabel}
+
+
+ );
+}
+
+/** Copies the on-disk path straight to the clipboard, no dialog. */
export function PathInfoButton({
path,
- title = "On-device location",
- description = "Where this model lives on disk.",
className,
}: {
path: string;
- title?: string;
- description?: string;
className?: string;
}) {
- const [open, setOpen] = useState(false);
const { copied, copy } = useCopyFeedback();
const handleCopy = async (event: MouseEvent
) => {
@@ -42,67 +87,27 @@ export function PathInfoButton({
};
return (
- <>
-
-
- {
- e.stopPropagation();
- setOpen(true);
- }}
- className={cn(
- "inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-muted hover:text-foreground focus-visible:opacity-100 group-hover/dl:opacity-100",
- className,
- )}
- >
-
-
-
-
- Show path
-
-
-
- e.stopPropagation()}
+
+
+
-
- {title}
- {description}
-
-
-
- {path}
-
-
-
-
-
-
-
-
- {copied ? "Copied" : "Copy path"}
-
-
-
-
-
- >
+
+
+
+
+ {copied ? "Copied" : "Copy path"}
+
+
);
}
diff --git a/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx b/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx
index 424afa2e05..5cba9c229b 100644
--- a/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx
@@ -7,37 +7,36 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
-import { useRepoDownload } from "../download-manager";
-import { deleteCachedModel } from "../inventory";
import { cn } from "@/lib/utils";
import {
Alert02Icon,
- PencilEdit02Icon,
PlayIcon,
+ RemoveCircleIcon,
} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { useEffect, useState } from "react";
import { TrainIcon } from "../components/train-icon";
+import { useRepoDownload } from "../download-manager";
+import { useOnlineStatus } from "../hooks/use-online-status";
+import { deleteCachedModel } from "../inventory";
+import type { ModelInventoryFormat } from "../inventory";
+import { fetchModelSize } from "../lib/dataset-size";
+import { formatBytes } from "../lib/format";
import {
HUB_NON_GGUF_RUN_ACTIONS_VISIBLE,
HUB_POST_DOWNLOAD_ACTIONS_VISIBLE,
} from "../lib/hub-feature-flags";
-import { HugeiconsIcon } from "@hugeicons/react";
-import { useEffect, useState } from "react";
-import { useHfTokenStore } from "../stores/hf-token-store";
-import { fetchModelSize } from "../lib/dataset-size";
-import { formatBytes } from "../lib/format";
import { fingerprintToken } from "../lib/token-fingerprint";
-import { useOnlineStatus } from "../hooks/use-online-status";
+import { useHfTokenStore } from "../stores/hf-token-store";
+import { DotTag } from "./dot-tag";
import {
CardDivider,
- CardDeleteButton,
DeleteConfirmDialog,
DownloadActionButton,
DownloadCard,
} from "./download-card";
-import { DotTag } from "./dot-tag";
-import { PathInfoButton } from "./path-info-button";
+import { QuantOptionsMenu } from "./gguf-download-card";
import { useCardDelete } from "./use-card-delete";
-import type { ModelInventoryFormat } from "../inventory";
import { useDownloadCardState } from "./use-download-card-state";
function formatModelLabel(modelFormat?: ModelInventoryFormat | null): string {
@@ -62,10 +61,9 @@ export function SafetensorsDownloadCard({
canRun = true,
isActive,
isLoadingThisModel,
- cachePath,
knownBytes,
onLoad,
- onUseInChat,
+ onEject,
onTrain,
onChange,
}: {
@@ -77,10 +75,13 @@ export function SafetensorsDownloadCard({
canRun?: boolean;
isActive: boolean;
isLoadingThisModel: boolean;
+ /** Accepted for API parity; the options menu resolves the path itself. */
cachePath?: string | null;
knownBytes?: number | null;
onLoad: (opts: { ggufVariant?: string; expectedBytes?: number }) => void;
+ /** Accepted for API parity; the run bar ejects instead of opening chat. */
onUseInChat?: () => void;
+ onEject?: () => void;
onTrain?: () => void;
onChange?: () => void;
}) {
@@ -95,8 +96,8 @@ export function SafetensorsDownloadCard({
knownBytes && knownBytes > 0
? knownBytes
: modelSize.key === sizeKey
- ? modelSize.bytes
- : null;
+ ? modelSize.bytes
+ : null;
const [deleteRepoOpen, setDeleteRepoOpen] = useState(false);
const { deleting, runDelete } = useCardDelete({
action: () => deleteCachedModel(repoId, undefined, hfToken || undefined),
@@ -170,7 +171,8 @@ export function SafetensorsDownloadCard({
!isLoadingThisModel;
return (
-
+
-
+
@@ -227,17 +229,20 @@ export function SafetensorsDownloadCard({
)}
- {canDelete && (
-
setDeleteRepoOpen(true)}
- />
- )}
- {isDownloaded && cachePath && (
- setDeleteRepoOpen(true)}
+ showPin={false}
+ buttonClassName="ml-0.5 size-7"
+ iconClassName="size-4"
/>
)}
@@ -268,7 +273,7 @@ export function SafetensorsDownloadCard({
onClick={() => {
if (!canRun) return;
if (isActive) {
- onUseInChat?.();
+ onEject?.();
return;
}
onLoad({});
@@ -287,26 +292,26 @@ export function SafetensorsDownloadCard({
>
) : isActive ? (
<>
-
- Chat
+
+ Eject
>
- ) : !canRun ? (
- <>
-
- No run
- >
- ) : (
+ ) : canRun ? (
<>
Run
>
+ ) : (
+ <>
+
+ No run
+ >
)}
) : showUnavailableAction ? (
diff --git a/studio/frontend/src/features/hub/catalog/sampling-settings-dialog.tsx b/studio/frontend/src/features/hub/catalog/sampling-settings-dialog.tsx
new file mode 100644
index 0000000000..4981f7734b
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/sampling-settings-dialog.tsx
@@ -0,0 +1,435 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// Gear button for the GGUF run bar: model config, system prompt, reasoning,
+// sampling, tools and retrieval, using the same controls as the chat page's
+// Run settings. Edits write to the chat runtime store's persisted state.
+
+import {
+ Dialog,
+ DialogContent,
+ DialogDescription,
+ DialogHeader,
+ DialogTitle,
+} from "@/components/ui/dialog";
+import { InfoHint } from "@/components/ui/info-hint";
+import { Switch } from "@/components/ui/switch";
+import { Textarea } from "@/components/ui/textarea";
+import {
+ Tooltip,
+ TooltipContent,
+ TooltipTrigger,
+} from "@/components/ui/tooltip";
+import {
+ ParamSlider,
+ useChatModelRuntime,
+ useChatRuntimeStore,
+} from "@/features/chat";
+import {
+ type PerModelConfig,
+ SidebarModelConfig,
+ applyPerModelConfigToRuntime,
+ currentRuntimePerModelConfig,
+ useActiveModelConfig,
+} from "@/features/model-picker";
+import { RetrievalSettingsSection } from "@/features/rag";
+import { toast } from "@/lib/toast";
+import { cn } from "@/lib/utils";
+import { Settings02Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { type ReactNode, useCallback, useState } from "react";
+import { HubOptionMenu } from "./hub-option-menu";
+
+function SettingsSection({
+ label,
+ labelClassName,
+ children,
+}: {
+ label: string;
+ labelClassName?: string;
+ children: ReactNode;
+}) {
+ return (
+
+
+ {label}
+
+
{children}
+
+ );
+}
+
+function ToggleRow({
+ label,
+ info,
+ checked,
+ disabled,
+ onCheckedChange,
+}: {
+ label: string;
+ info?: string;
+ checked: boolean;
+ disabled?: boolean;
+ onCheckedChange: (checked: boolean) => void;
+}) {
+ return (
+
+
+ {label}
+ {info && {info} }
+
+
onCheckedChange(value === true)}
+ aria-label={label}
+ />
+
+ );
+}
+
+export function SamplingSettingsButton({ className }: { className?: string }) {
+ const [open, setOpen] = useState(false);
+ const params = useChatRuntimeStore((s) => s.params);
+ const setParams = useChatRuntimeStore((s) => s.setParams);
+
+ // Loaded model's per-model config (context length etc.), mirroring the chat
+ // page's Run settings Model section.
+ const { selectModel } = useChatModelRuntime();
+ const modelLoading = useChatRuntimeStore((s) => s.modelLoading);
+ const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
+ const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
+ const ggufNativeContextLength = useChatRuntimeStore(
+ (s) => s.ggufNativeContextLength,
+ );
+ const {
+ checkpoint,
+ isGguf: activeModelIsGguf,
+ config: activeModelConfig,
+ } = useActiveModelConfig();
+ const handleReloadActiveModel = useCallback(
+ (config: PerModelConfig) => {
+ const runtime = useChatRuntimeStore.getState();
+ const activeCheckpoint = runtime.params.checkpoint;
+ if (!activeCheckpoint) return;
+ const nativeToken = runtime.activeNativePathToken;
+ const nativeExpiry = runtime.activeNativePathExpiresAtMs;
+ // Mirrors the chat page: an expired native-path token can't reload.
+ if (nativeToken && nativeExpiry != null && Date.now() >= nativeExpiry) {
+ toast.error("This local model file's access has expired.", {
+ description: "Re-select the model file to reload it.",
+ });
+ return;
+ }
+ // selectModel reads config from the runtime store, not the selection, so
+ // apply it first (snapshotting the current one for rollback).
+ const previousConfig = currentRuntimePerModelConfig({
+ includeMaxSeqLength: true,
+ });
+ applyPerModelConfigToRuntime(config);
+ void selectModel({
+ id: activeCheckpoint,
+ source: "local",
+ ggufVariant: runtime.activeGgufVariant ?? undefined,
+ nativePathToken: nativeToken ?? undefined,
+ nativePathExpiresAtMs: nativeExpiry,
+ isGguf: activeModelIsGguf,
+ isDownloaded: true,
+ keepSpeculative: true,
+ previousConfig,
+ forceReload: true,
+ });
+ },
+ [selectModel, activeModelIsGguf],
+ );
+
+ // Reasoning + tools: same store bindings as the chat page.
+ const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled);
+ const reasoningAlwaysOn = useChatRuntimeStore((s) => s.reasoningAlwaysOn);
+ const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled);
+ const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort);
+ const reasoningEffortLevels = useChatRuntimeStore(
+ (s) => s.reasoningEffortLevels,
+ );
+ const setReasoningEffort = useChatRuntimeStore((s) => s.setReasoningEffort);
+ const supportsPreserveThinking = useChatRuntimeStore(
+ (s) => s.supportsPreserveThinking,
+ );
+ const preserveThinking = useChatRuntimeStore((s) => s.preserveThinking);
+ const setPreserveThinking = useChatRuntimeStore((s) => s.setPreserveThinking);
+ const maxToolCalls = useChatRuntimeStore((s) => s.maxToolCallsPerMessage);
+ const setMaxToolCalls = useChatRuntimeStore(
+ (s) => s.setMaxToolCallsPerMessage,
+ );
+ const toolCallTimeout = useChatRuntimeStore((s) => s.toolCallTimeout);
+ const setToolCallTimeout = useChatRuntimeStore((s) => s.setToolCallTimeout);
+ const autoHealToolCalls = useChatRuntimeStore((s) => s.autoHealToolCalls);
+ const setAutoHealToolCalls = useChatRuntimeStore(
+ (s) => s.setAutoHealToolCalls,
+ );
+ const nudgeToolCalls = useChatRuntimeStore((s) => s.nudgeToolCalls);
+ const setNudgeToolCalls = useChatRuntimeStore((s) => s.setNudgeToolCalls);
+ const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
+ const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls);
+
+ const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled);
+ const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled);
+
+ const set = (key: keyof typeof params) => (value: number) =>
+ setParams({ ...params, [key]: value });
+
+ // Slider 0-41; 41 maps to 9999 ("Max"), mirroring the chat page.
+ const toolCallsSliderValue =
+ maxToolCalls >= 9999 ? 41 : Math.min(maxToolCalls, 40);
+ // Slider 1-31; 31 maps to 9999 ("Max").
+ const timeoutSliderValue =
+ toolCallTimeout >= 9999 ? 31 : Math.min(Math.max(toolCallTimeout, 1), 30);
+
+ return (
+ <>
+
+
+ {
+ e.stopPropagation();
+ setOpen(true);
+ }}
+ className={cn(
+ "inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground",
+ className,
+ )}
+ >
+
+
+
+
+ Inference settings
+
+
+
+ e.stopPropagation()}
+ >
+
+ Inference settings
+
+ Applies to chats with local models.
+
+
+
+ {checkpoint && activeModelConfig && !modelLoading && (
+
+
+
+ )}
+
+
+
+
+
+
+ {reasoningEffortLevels.length > 0 && (
+
+
+ Reasoning effort
+
+ ({
+ value: level,
+ label: level.charAt(0).toUpperCase() + level.slice(1),
+ }))}
+ onValueChange={setReasoningEffort}
+ ariaLabel="Reasoning effort"
+ align="end"
+ className="h-8 text-[11.5px]"
+ />
+
+ )}
+ {supportsPreserveThinking && (
+
+ )}
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ setMaxToolCalls(v >= 41 ? 9999 : v)}
+ displayValue={
+ toolCallsSliderValue >= 41
+ ? "Max"
+ : toolCallsSliderValue === 0
+ ? "Off"
+ : undefined
+ }
+ info="Cap on tool/function calls the model may invoke within a single response. 0 disables tool use; Max removes the cap."
+ />
+ setToolCallTimeout(v >= 31 ? 9999 : v)}
+ displayValue={
+ timeoutSliderValue >= 31
+ ? "Max"
+ : timeoutSliderValue === 1
+ ? "1 minute"
+ : `${timeoutSliderValue} minutes`
+ }
+ valueSize={10}
+ info="Per-call wall-clock limit. Long-running tool executions are terminated when this elapses; the model continues with what completed."
+ />
+
+
+
+
+
+
+
+
+ >
+ );
+}
diff --git a/studio/frontend/src/features/hub/catalog/shared.tsx b/studio/frontend/src/features/hub/catalog/shared.tsx
index c435afcfe6..fbb3f6bb1d 100644
--- a/studio/frontend/src/features/hub/catalog/shared.tsx
+++ b/studio/frontend/src/features/hub/catalog/shared.tsx
@@ -1,12 +1,19 @@
// 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 {
+ Tooltip,
+ TooltipContent,
+ TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { cn } from "@/lib/utils";
import {
BrainIcon,
Chat01Icon,
CodeIcon,
GlobeIcon,
HeadphonesIcon,
+ ImageIcon,
LockIcon,
LockKeyIcon,
SparklesIcon,
@@ -15,12 +22,6 @@ import {
} from "@hugeicons/core-free-icons";
import type { IconSvgElement } from "@hugeicons/react";
import { HugeiconsIcon } from "@hugeicons/react";
-import {
- Tooltip,
- TooltipContent,
- TooltipTrigger,
-} from "@/components/ui/tooltip";
-import { cn } from "@/lib/utils";
import type { Capability, CapabilityKey } from "../lib/model-capabilities";
const CAPABILITY_ICON: Record = {
@@ -30,6 +31,7 @@ const CAPABILITY_ICON: Record = {
reasoning: BrainIcon,
code: CodeIcon,
embedding: SparklesIcon,
+ diffusion: ImageIcon,
multilingual: GlobeIcon,
conversational: Chat01Icon,
};
@@ -45,6 +47,8 @@ const CAPABILITY_TONE: Record = {
code: "bg-cyan-500/10 text-cyan-800 dark:bg-cyan-400/20 dark:text-cyan-300",
embedding:
"bg-emerald-500/10 text-emerald-700 dark:bg-emerald-400/20 dark:text-emerald-300",
+ diffusion:
+ "bg-pink-500/10 text-pink-700 dark:bg-pink-400/20 dark:text-pink-300",
multilingual:
"bg-sky-500/10 text-sky-700 dark:bg-sky-400/20 dark:text-sky-300",
conversational:
@@ -59,9 +63,7 @@ export function AccessChip({ label }: { label: string }) {
);
}
-function isGatedAccess(
- gated: false | "auto" | "manual" | undefined,
-): boolean {
+function isGatedAccess(gated: false | "auto" | "manual" | undefined): boolean {
return gated !== false && gated !== undefined;
}
diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts b/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts
index da653ddeb7..2cea1d8304 100644
--- a/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts
+++ b/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts
@@ -1,14 +1,10 @@
// 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 { DOWNLOAD_KIND } from "./constants";
import {
createDownloadManagerInitialState,
- jobKeyOf,
removeJob,
- selectActiveJob,
setState,
- useDownloadManagerStore,
} from "./download-manager-state";
import { resetDownloadApiAdapterState } from "./download-api-adapter";
import {
@@ -69,25 +65,6 @@ export const downloadManager: DownloadManagerController = {
dismiss: removeJob,
};
-/** Cancel the in-flight download for a staged model pick. No-op when nothing is
- * downloading (e.g. a native/local file that was never fetched). Lets non-React
- * callers (the chat store's abandon paths) stop a staged transfer without the
- * useRepoDownload hook. */
-export function cancelStagedModelDownload(
- pending: { id: string; ggufVariant?: string | null } | null,
-): void {
- if (!pending) return;
- const variant = pending.ggufVariant ?? null;
- const activeJob = selectActiveJob(
- useDownloadManagerStore.getState(),
- DOWNLOAD_KIND.MODEL,
- pending.id,
- variant,
- );
- void downloadManager.cancel(
- activeJob?.key ?? jobKeyOf(DOWNLOAD_KIND.MODEL, pending.id, variant),
- );
-}
if (import.meta.hot) {
import.meta.hot.dispose(() => {
diff --git a/studio/frontend/src/features/hub/download-manager/index.ts b/studio/frontend/src/features/hub/download-manager/index.ts
index dd88aaf3f0..60ef3851f8 100644
--- a/studio/frontend/src/features/hub/download-manager/index.ts
+++ b/studio/frontend/src/features/hub/download-manager/index.ts
@@ -20,7 +20,6 @@ export {
} from "./constants";
export {
__resetDownloadManagerForTests,
- cancelStagedModelDownload,
clearCompletedInventoryHint,
downloadManager,
hydrateDownloadManager,
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx
index 29e2f3f277..426f816dbe 100644
--- a/studio/frontend/src/features/hub/hub-page.tsx
+++ b/studio/frontend/src/features/hub/hub-page.tsx
@@ -1,28 +1,7 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import {
- loadRememberedLoadSettings,
- rememberedLoadSettingsKey,
-} from "@/components/assistant-ui/model-selector/remembered-load-settings";
-import { hfModelFitsDevice } from "@/components/assistant-ui/model-selector/recommended-fit";
-import { useHubInventory } from "@/features/hub/inventory";
-import { useDebouncedValue } from "@/hooks/use-debounced-value";
-import { useGpuInfo } from "@/hooks/use-gpu-info";
-import {
- type HfModelSearchChannel,
- type HfSortDirection,
- type HfSortKey,
-} from "@/features/hub/hooks/use-hub-model-search";
-import { useOnlineStatus } from "@/features/hub/hooks/use-online-status";
-import { useHubInfiniteScroll } from "@/features/hub/hooks/use-hub-infinite-scroll";
-import { ggufVariantsMatch, modelIdsMatch } from "@/features/hub/lib/model-identity";
-import { cn } from "@/lib/utils";
import { usePlatformStore } from "@/config/env";
-import {
- hfApiToken,
- useHfTokenStore,
-} from "@/features/hub/stores/hf-token-store";
import {
isChannelEntryFresh,
useHubFeedStore,
@@ -33,6 +12,25 @@ import {
useChatModelRuntime,
useChatRuntimeStore,
} from "@/features/chat";
+import { useHubInventory } from "./inventory";
+import type {
+ HfModelSearchChannel,
+ HfSortDirection,
+ HfSortKey,
+} from "./hooks/use-hub-model-search";
+import { useOnlineStatus } from "@/features/hub";
+import { useHubInfiniteScroll } from "@/features/hub";
+import { ggufVariantsMatch, modelIdsMatch } from "./lib/model-identity";
+import { hfApiToken, useHfTokenStore } from "./stores/hf-token-store";
+import {
+ applyModelLoadConfigToRuntime,
+ currentRuntimePerModelConfig,
+ hfModelFitsDevice,
+ resolveInitialConfig,
+} from "@/features/model-picker";
+import { useDebouncedValue } from "@/hooks/use-debounced-value";
+import { useGpuInfo } from "@/hooks/use-gpu-info";
+import { cn } from "@/lib/utils";
import { useNavigate, useSearch } from "@tanstack/react-router";
import {
useCallback,
@@ -42,17 +40,10 @@ import {
useRef,
useState,
} from "react";
+import { ExternalLinkConfirmDialog } from "./catalog/external-link-confirm-dialog";
import { HubDetailView } from "./catalog/hub-detail-view";
-import { HubTopBar } from "./catalog/hub-top-bar";
import { HubFeed } from "./catalog/hub-feed";
-import { OwnerScopeToggle } from "./catalog/owner-scope-toggle";
-import {
- type AllModelsView,
- HubListHeader,
- type InventorySort,
- InventorySortControl,
- ResultListHeader,
-} from "./catalog/models-table";
+import { HubTopBar } from "./catalog/hub-top-bar";
import {
ModelsCatalog,
type ModelsCatalogHandlers,
@@ -60,14 +51,22 @@ import {
type ModelsCatalogState,
} from "./catalog/models-catalog";
import { ModelsHeader } from "./catalog/models-header";
+import {
+ type AllModelsView,
+ HubListHeader,
+ type InventorySort,
+ InventorySortControl,
+ InventoryTypeFilterControl,
+ ResultListHeader,
+} from "./catalog/models-table";
import { ModelsToolbar } from "./catalog/models-toolbar";
-import { ExternalLinkConfirmDialog } from "./catalog/external-link-confirm-dialog";
import { OnDeviceFoldersDialog } from "./catalog/on-device-folders-dialog";
+import { OwnerScopeToggle } from "./catalog/owner-scope-toggle";
import { useDiscoverSearch } from "./hooks/use-discover-search";
import { useFeedWriteBack } from "./hooks/use-feed-write-back";
+import { useHiddenEmbeddingModelIds } from "./hooks/use-hidden-embedding-models";
import { useHubFeed } from "./hooks/use-hub-feed";
import { useHubModelVram } from "./hooks/use-hub-model-vram";
-import { useHiddenEmbeddingModelIds } from "./hooks/use-hidden-embedding-models";
import { useModelsSelection } from "./hooks/use-models-selection";
import {
CHANNEL_TO_SECTION,
@@ -83,6 +82,10 @@ import {
isHiddenModelId,
} from "./lib/hidden-models";
import { inventoryRowMatches, tokenizeQuery } from "./lib/inventory-search";
+import {
+ type ModelTypeFilter,
+ matchesModelType,
+} from "./lib/model-type-filter";
import { resolveOwnerProviderLogo } from "./lib/provider-logos";
import { fingerprintToken } from "./lib/token-fingerprint";
import {
@@ -456,6 +459,8 @@ export function ModelsPage() {
setInventorySortState(sort);
writeInventorySortPreference(sort);
}, []);
+ const [inventoryTypeFilter, setInventoryTypeFilter] =
+ useState("all");
const [foldersDialogOpen, setFoldersDialogOpen] = useState(false);
const [discoverFetchIntent, setDiscoverFetchIntent] = useState(0);
const [sortBrowseActive, setSortBrowseActive] = useState(false);
@@ -582,15 +587,15 @@ export function ModelsPage() {
const deferredCapabilityFilter = useDeferredValue(capabilityFilter);
const hasQuery = deferredDebouncedQuery.trim() !== "";
- const mode: DiscoverMode = !isModelDiscover
- ? "search"
- : hasQuery
+ const mode: DiscoverMode = isModelDiscover
+ ? hasQuery
? "search"
: urlSection != null
? "channel-list"
: sortBrowseActive
? "search"
- : "feed";
+ : "feed"
+ : "search";
const isFeedMode = mode === "feed";
const isChannelListMode = mode === "channel-list";
const isSortBrowseMode =
@@ -737,7 +742,10 @@ export function ModelsPage() {
// The default feed only shows models with a provider logo.
(!isFeedMode ||
resolveOwnerProviderLogo(row.owner, row.repo) !== null) &&
- matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) &&
+ matchesFormat(
+ detectResultFormat(row.result),
+ effectiveDiscoverFormat,
+ ) &&
matchesCapability(row.capabilities, deferredCapabilityFilter) &&
(!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)) &&
// Models already on disk stay visible regardless of device fit,
@@ -807,7 +815,10 @@ export function ModelsPage() {
}
return merged;
}, [isFeedMode, feedTrendingRows, filteredDiscoverRows]);
- const feedResults = useMemo(() => feedRows.map((row) => row.result), [feedRows]);
+ const feedResults = useMemo(
+ () => feedRows.map((row) => row.result),
+ [feedRows],
+ );
const selectionDiscoverRows = isFeedMode ? feedRows : discoverRows;
const selectionFilteredDiscoverRows = isFeedMode
? feedRows
@@ -837,7 +848,8 @@ export function ModelsPage() {
// Local rows may lack a repo id, so also check path and title.
return (
!isHiddenModelId(row.id, row.repoId, row.path, row.title) ||
- (inventoryTokens.length > 0 && inventoryRowMatches(row, inventoryTokens))
+ (inventoryTokens.length > 0 &&
+ inventoryRowMatches(row, inventoryTokens))
);
},
[hiddenEmbeddingModelIds, inventoryTokens],
@@ -855,6 +867,7 @@ export function ModelsPage() {
// id/title/path happens to contain an infra needle is not dropped.
isDatasetMode ||
(matchesFormat(row.modelFormat, deferredFormatFilter) &&
+ matchesModelType(row, inventoryTypeFilter) &&
isVisibleInventoryRow(row)),
),
inventoryTokens,
@@ -863,6 +876,7 @@ export function ModelsPage() {
effectiveCachedRows,
isDatasetMode,
deferredFormatFilter,
+ inventoryTypeFilter,
inventoryTokens,
isVisibleInventoryRow,
],
@@ -878,6 +892,7 @@ export function ModelsPage() {
// id/title/path happens to contain an infra needle is not dropped.
isDatasetMode ||
(matchesFormat(row.modelFormat, deferredFormatFilter) &&
+ matchesModelType(row, inventoryTypeFilter) &&
isVisibleInventoryRow(row)),
),
inventoryTokens,
@@ -886,6 +901,7 @@ export function ModelsPage() {
effectiveLocalRows,
isDatasetMode,
deferredFormatFilter,
+ inventoryTypeFilter,
inventoryTokens,
isVisibleInventoryRow,
],
@@ -918,6 +934,7 @@ export function ModelsPage() {
resourceType,
deferredFormatFilter,
deferredCapabilityFilter,
+ inventoryTypeFilter,
effectiveSort,
effectiveDirection,
activeChannelId,
@@ -928,6 +945,7 @@ export function ModelsPage() {
resourceType,
deferredFormatFilter,
deferredCapabilityFilter,
+ inventoryTypeFilter,
effectiveSort,
effectiveDirection,
activeChannelId,
@@ -946,6 +964,7 @@ export function ModelsPage() {
}
} else {
setDownloadedFormat("all");
+ setInventoryTypeFilter("all");
}
setCapabilityFilter("all");
}, [isDiscoverTab, urlSection, navigate]);
@@ -1155,50 +1174,22 @@ export function ModelsPage() {
(opts: ModelLoadOptions, isDownloaded: boolean) => {
if (!selectedModel) return;
const runId = selectedModel.resource.runId;
- // "Load on selection" off: stage GGUF picks instead of loading, so the
- // chat page's staging flow can read the header and show the load options.
- // Non-GGUF models have nothing to configure pre-load, so they load now.
- if (
- !useChatRuntimeStore.getState().loadOnSelection &&
- (opts.ggufVariant != null || selectedModel.isGguf)
- ) {
- useChatRuntimeStore.getState().stageModel({
- id: runId,
- ggufVariant: opts.ggufVariant,
- isGguf: selectedModel.isGguf,
- isDownloaded,
- expectedBytes: opts.expectedBytes,
- });
- openNewChat();
- return;
- }
- // Detach any leftover staged pick first so its edited knobs (e.g. a custom
- // context length) don't leak into this load -- mirrors the chat page's
- // detachStaged(); keepDownload keeps any staged download running.
- useChatRuntimeStore.getState().abandonStagedModel({ keepDownload: true });
- // Load-on-selection skips the chat sheet, so seed this GGUF pick's saved
- // load knobs here the way the sheet's restore effect would; otherwise the
- // remembered config is silently ignored on the Hub run path. keepSpeculative
- // then honors the restored speculative choice across the switch.
- const remembered =
- opts.ggufVariant != null || selectedModel.isGguf
- ? loadRememberedLoadSettings(
- rememberedLoadSettingsKey({
- id: runId,
- ggufVariant: opts.ggufVariant,
- }),
- )
- : null;
- if (remembered) {
- useChatRuntimeStore.getState().applyRememberedLoadSettings(remembered);
- }
+ const resolvedConfig = resolveInitialConfig(runId, opts.ggufVariant);
+ const rememberedConfig = resolvedConfig.remembered
+ ? resolvedConfig.config
+ : null;
+ const previousConfig = currentRuntimePerModelConfig({
+ includeMaxSeqLength: true,
+ });
+ const hasAppliedConfig = applyModelLoadConfigToRuntime(rememberedConfig);
void selectModel({
id: runId,
ggufVariant: opts.ggufVariant,
isDownloaded,
expectedBytes: opts.expectedBytes,
- keepSpeculative: remembered != null,
+ keepSpeculative: hasAppliedConfig,
throwOnError: true,
+ previousConfig,
})
.then(() => {
// Read fresh: the load is async, so the checkpoint may have changed.
@@ -1280,6 +1271,7 @@ export function ModelsPage() {
onLoad: handleLoad,
onLoadLocal: handleLoadLocal,
onUseInChat: openNewChat,
+ onEject: () => void ejectModel(),
onTrain: handleTrain,
onInventoryChange: refreshInventory,
onSearchHub: handleSearchHub,
@@ -1288,6 +1280,7 @@ export function ModelsPage() {
handleLoad,
handleLoadLocal,
openNewChat,
+ ejectModel,
handleTrain,
handleSearchHub,
refreshInventory,
@@ -1295,31 +1288,38 @@ export function ModelsPage() {
);
const catalogState = useMemo(
- () => ({
- tab,
- discoverRows: listRows,
- cachedRows: filteredCachedRows,
- localRows: filteredLocalRows,
- selectedId,
- isLoading,
- downloadedReady,
- inventoryError,
- inventoryWarning,
- query,
- activeCheckpoint,
- activeGgufVariant,
- searchError,
- online,
- isDataset: isDatasetMode,
- inventoryTokens,
- scannedCount,
- loadingIntentCount: discoverFetchIntent,
- hasMore,
- manualFetchAvailable: discoverManualFetchAvailable,
- hasActiveFilters:
- !isFeedMode &&
- (deferredFormatFilter !== "all" || deferredCapabilityFilter !== "all"),
- }),
+ () => {
+ const typeFilterActive =
+ !isDatasetMode && inventoryTypeFilter !== "all";
+ return {
+ tab,
+ discoverRows: listRows,
+ cachedRows: filteredCachedRows,
+ localRows: filteredLocalRows,
+ selectedId,
+ isLoading,
+ downloadedReady,
+ inventoryError,
+ inventoryWarning,
+ query,
+ activeCheckpoint,
+ activeGgufVariant,
+ searchError,
+ online,
+ isDataset: isDatasetMode,
+ inventoryTokens,
+ scannedCount,
+ loadingIntentCount: discoverFetchIntent,
+ hasMore,
+ manualFetchAvailable: discoverManualFetchAvailable,
+ hasActiveFilters:
+ !isFeedMode &&
+ (deferredFormatFilter !== "all" ||
+ deferredCapabilityFilter !== "all" ||
+ (tab === "downloaded" && typeFilterActive)),
+ typeFilterActive,
+ };
+ },
[
tab,
isFeedMode,
@@ -1344,6 +1344,7 @@ export function ModelsPage() {
discoverManualFetchAvailable,
deferredFormatFilter,
deferredCapabilityFilter,
+ inventoryTypeFilter,
],
);
@@ -1422,16 +1423,18 @@ export function ModelsPage() {
);
}
- const ownerToggle = !isDatasetMode ? (
+ const ownerToggle = isDatasetMode ? undefined : (
- ) : undefined;
+ );
// Compact pill so it stays beside the view-mode tabs even in the narrow
// split pane instead of dropping to its own row.
return (
{isChannelListMode ? (
{
- const sortControl = (
-
+ // Compact pills so they stay beside the view-mode tabs even in the narrow
+ // split pane instead of dropping to their own row.
+ const controls = (
+
+ {!isDatasetMode && (
+
+ )}
+
+
);
- // Compact pill so it stays beside the view-mode tabs even in the narrow
- // split pane instead of dropping to its own row.
return (
);
}, [
- visibleCachedCount,
- visibleLocalCount,
+ filteredCachedRows,
+ filteredLocalRows,
allModelsView,
setAllModelsView,
inventorySort,
setInventorySort,
+ inventoryTypeFilter,
+ isDatasetMode,
]);
const detailOpen = urlModel !== null;
diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts
index 5d4151e87d..b464ddf0cf 100644
--- a/studio/frontend/src/features/hub/index.ts
+++ b/studio/frontend/src/features/hub/index.ts
@@ -1,10 +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
-export { cancelStagedModelDownload } from "./download-manager";
+export {
+ downloadManager,
+ jobKeyOf,
+ subscribeJobListeners,
+ useDownloadManagerStore,
+} from "./download-manager";
+export {
+ useHubInventory,
+ type CachedInventoryRow,
+ type GgufVariantDetail,
+ type LocalInventoryRow,
+ type LocalSource,
+ type ScanFolderInfo,
+ addScanFolder,
+ deleteCachedModel,
+ invalidateGgufVariantsCache,
+ listGgufVariants,
+ listScanFolders,
+ removeScanFolder,
+} from "./inventory";
+export {
+ type HfModelResult,
+ type HfSortKey,
+ useHubModelSearch,
+} from "./hooks/use-hub-model-search";
+export { useOnlineStatus } from "./hooks/use-online-status";
+export { useHubInfiniteScroll } from "./hooks/use-hub-infinite-scroll";
export { bumpInventoryVersion } from "./stores/inventory-events";
export {
getHfToken,
+ hfApiToken,
mirrorHfTokenInto,
useHfTokenStore,
} from "./stores/hf-token-store";
+export { useInventoryVersion } from "./stores/inventory-events";
+export { looksLikeLocalPath } from "./lib/local-path";
+export { hubTokenHeader } from "./lib/hub-token-header";
+export {
+ ggufVariantsMatch,
+ normalizeGgufVariantIdentity,
+ normalizeModelIdentity,
+} from "./lib/model-identity";
+export { formatBytes, formatRelativeShort } from "./lib/format";
+export { ggufVariantDisplayLabel } from "./lib/gguf-variant-sort";
+export {
+ DeleteConfirmDialog,
+ UpdateConfirmDialog,
+} from "./catalog/download-card";
+export { HubOptionMenu, type HubOption } from "./catalog/hub-option-menu";
+export { DotTag } from "./catalog/dot-tag";
+export { TransportConflictDialog } from "./catalog/transport-conflict-dialog";
+export { TrainIcon } from "./components/train-icon";
+export { isHiddenModelId } from "./lib/hidden-models";
+export { classifyUnslothSupport } from "./lib/unsloth-support";
diff --git a/studio/frontend/src/features/hub/inventory/api.ts b/studio/frontend/src/features/hub/inventory/api.ts
index 0e01c5c0cb..8c9214c9e5 100644
--- a/studio/frontend/src/features/hub/inventory/api.ts
+++ b/studio/frontend/src/features/hub/inventory/api.ts
@@ -48,6 +48,7 @@ export interface CachedGgufRepo {
capabilities?: BackendModelCapabilities | null;
size_bytes: number;
cache_path?: string;
+ last_modified?: number | null;
partial?: boolean;
partial_transport?: string | null;
pipeline_tag?: string | null;
@@ -65,6 +66,7 @@ export interface CachedModelRepo {
capabilities?: BackendModelCapabilities | null;
size_bytes: number;
cache_path?: string;
+ last_modified?: number | null;
partial?: boolean;
partial_transport?: string | null;
pipeline_tag?: string | null;
diff --git a/studio/frontend/src/features/hub/inventory/types.ts b/studio/frontend/src/features/hub/inventory/types.ts
index 6f65a56037..b300be5fb0 100644
--- a/studio/frontend/src/features/hub/inventory/types.ts
+++ b/studio/frontend/src/features/hub/inventory/types.ts
@@ -47,6 +47,7 @@ export interface CachedInventoryRow {
capabilities: ModelInventoryCapabilities;
bytes: number;
cachePath?: string | null;
+ lastModified?: number | null;
partial?: boolean;
partialTransport?: string | null;
pipelineTag?: string | null;
@@ -66,6 +67,8 @@ export interface LocalInventoryRow {
title: string;
source: LocalSource;
sourceLabel: string;
+ modelId?: string | null;
+ displayName?: string;
path: string;
isGguf: boolean;
modelFormat: ModelInventoryFormat;
diff --git a/studio/frontend/src/features/hub/inventory/use-device-inventory.ts b/studio/frontend/src/features/hub/inventory/use-device-inventory.ts
index 00eaa31a22..5ff57dfff4 100644
--- a/studio/frontend/src/features/hub/inventory/use-device-inventory.ts
+++ b/studio/frontend/src/features/hub/inventory/use-device-inventory.ts
@@ -14,6 +14,7 @@ import {
listLocalModels,
} from "./api";
import { useDebouncedValue } from "@/hooks/use-debounced-value";
+import { ensureHiddenModelMatchers } from "../lib/hidden-models";
import { fingerprintToken } from "@/features/hub/lib/token-fingerprint";
import { useInventoryVersion } from "@/features/hub/stores/inventory-events";
import { useCallback, useEffect, useMemo } from "react";
@@ -146,12 +147,15 @@ async function runSourceFetch(
): Promise {
switch (source) {
case "cachedGguf":
+ await ensureHiddenModelMatchers();
return (await listCachedGguf(hfToken)) as DeviceInventoryRows[K];
case "cachedModels":
+ await ensureHiddenModelMatchers();
return (await listCachedModels(hfToken)) as DeviceInventoryRows[K];
case "cachedDatasets":
return (await listCachedDatasets()) as DeviceInventoryRows[K];
case "localModels":
+ await ensureHiddenModelMatchers();
return (await listLocalModels()).models as DeviceInventoryRows[K];
case "localDatasets":
return (await listLocalDatasets()).datasets as DeviceInventoryRows[K];
diff --git a/studio/frontend/src/features/hub/inventory/view-models.ts b/studio/frontend/src/features/hub/inventory/view-models.ts
index 334050fab4..5fd9cc2228 100644
--- a/studio/frontend/src/features/hub/inventory/view-models.ts
+++ b/studio/frontend/src/features/hub/inventory/view-models.ts
@@ -176,6 +176,7 @@ export function buildCachedInventoryRow(
runtime?: string | null;
format_variant?: string | null;
capabilities?: BackendModelCapabilities | null;
+ last_modified?: number | null;
optimistic?: boolean;
},
fallbackFormat: ModelInventoryFormat,
@@ -215,6 +216,12 @@ export function buildCachedInventoryRow(
capabilities,
bytes: row.size_bytes,
cachePath: row.cache_path ?? null,
+ lastModified:
+ typeof row.last_modified === "number" &&
+ Number.isFinite(row.last_modified) &&
+ row.last_modified > 0
+ ? row.last_modified
+ : null,
partial: row.partial ?? false,
partialTransport: row.partial_transport ?? null,
pipelineTag: row.pipeline_tag ?? null,
@@ -278,6 +285,8 @@ export function buildLocalInventoryRows(
title,
source: model.source,
sourceLabel: localSourceLabel(model.source),
+ modelId: model.model_id ?? null,
+ displayName: model.display_name,
path: model.path,
isGguf: modelFormat === "gguf",
modelFormat,
diff --git a/studio/frontend/src/features/hub/lib/hidden-models.ts b/studio/frontend/src/features/hub/lib/hidden-models.ts
index 634a061e0c..ce63832f99 100644
--- a/studio/frontend/src/features/hub/lib/hidden-models.ts
+++ b/studio/frontend/src/features/hub/lib/hidden-models.ts
@@ -1,19 +1,77 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+import { authFetch } from "@/features/auth";
+import { getInventoryVersion } from "../stores/inventory-events";
+
// Infra models hidden from browse/preview lists (Hub Discover, the chat model
// selector, and local on-device rows). Mirrors the backend
// `utils.hidden_models`: the RAG embedding model and the llama.cpp validation
// probe are not usable chat models. Server-confirmed cache rows are trusted
// because the backend applies variant-aware filtering. Optimistic cache rows
-// still use these needles until the server confirms them. Per-repo views are
-// not filtered, so reinstall flows still show downloaded files.
+// still use these needles until the server confirms them. The dynamic matchers
+// fetched from `/api/hub/hidden-models` add the user's configured embedder as
+// exact repo ids and exact resolved paths, never substring needles. Per-repo
+// views are not filtered, so reinstall flows still show downloaded files.
const HIDDEN_NEEDLES = [
"bge-small-en-v1.5", // RAG embedder: unsloth/bge-small-en-v1.5[-GGUF]
"ggml-org/models", // llama.cpp validation probe repo
"stories260k.gguf", // probe filename (carries .gguf so it stays specific)
];
+let dynamicNeedles: readonly string[] = [];
+let dynamicExactIds: readonly string[] = [];
+let dynamicExactPaths: readonly string[] = [];
+let matchersFetch: Promise | null = null;
+let matchersFetchVersion = -1;
+
+function toLowerStrings(value: unknown): string[] {
+ if (!Array.isArray(value)) {
+ return [];
+ }
+ return value
+ .filter((v): v is string => typeof v === "string" && v.length > 0)
+ .map((v) => v.toLowerCase());
+}
+
+export function ensureHiddenModelMatchers(): Promise {
+ const version = getInventoryVersion();
+ if (matchersFetch && matchersFetchVersion === version) {
+ return matchersFetch;
+ }
+ matchersFetchVersion = version;
+ matchersFetch = (async () => {
+ try {
+ const response = await authFetch("/api/hub/hidden-models");
+ if (!response.ok) {
+ throw new Error(`HTTP ${response.status}`);
+ }
+ const data = (await response.json()) as {
+ needles?: unknown;
+ exact_ids?: unknown;
+ exact_paths?: unknown;
+ };
+ if (
+ getInventoryVersion() !== version ||
+ matchersFetchVersion !== version
+ ) {
+ return;
+ }
+ dynamicNeedles = toLowerStrings(data.needles);
+ dynamicExactIds = toLowerStrings(data.exact_ids);
+ dynamicExactPaths = toLowerStrings(data.exact_paths);
+ } catch {
+ if (
+ getInventoryVersion() === version &&
+ matchersFetchVersion === version
+ ) {
+ matchersFetch = null;
+ }
+ }
+ })();
+ return matchersFetch;
+}
+
/** True if any id/path is a hidden infra model. */
export function isHiddenModelId(
...values: (string | null | undefined)[]
@@ -23,7 +81,12 @@ export function isHiddenModelId(
return false;
}
const lower = v.toLowerCase();
- return HIDDEN_NEEDLES.some((needle) => lower.includes(needle));
+ return (
+ HIDDEN_NEEDLES.some((needle) => lower.includes(needle)) ||
+ dynamicNeedles.some((needle) => lower.includes(needle)) ||
+ dynamicExactIds.includes(lower) ||
+ dynamicExactPaths.includes(lower)
+ );
});
}
diff --git a/studio/frontend/src/features/hub/lib/model-capabilities.ts b/studio/frontend/src/features/hub/lib/model-capabilities.ts
index 033a3b6f1f..1baa3fb8e0 100644
--- a/studio/frontend/src/features/hub/lib/model-capabilities.ts
+++ b/studio/frontend/src/features/hub/lib/model-capabilities.ts
@@ -10,6 +10,7 @@ export type CapabilityKey =
| "reasoning"
| "code"
| "embedding"
+ | "diffusion"
| "multilingual"
| "conversational";
@@ -62,6 +63,20 @@ const REASONING_TAGS = new Set([
"step-by-step",
]);
+// Image generation / diffusion (surfaced as "Image generation" in filters).
+const DIFFUSION_TAGS = new Set([
+ "diffusers",
+ "diffusion",
+ "stable-diffusion",
+ "latent-diffusion",
+ "flux",
+ "text-to-image",
+ "image-to-image",
+ "text-to-video",
+ "image-to-video",
+ "unconditional-image-generation",
+]);
+
const CODE_TAGS = new Set([
"code",
"code-generation",
@@ -186,10 +201,7 @@ export function detectCapabilities(
) {
out.push({ key: "code", label: "Code" });
}
- if (
- hasAny(CONVERSATIONAL_TAGS) ||
- CONVERSATIONAL_ID_RE.test(lowerId)
- ) {
+ if (hasAny(CONVERSATIONAL_TAGS) || CONVERSATIONAL_ID_RE.test(lowerId)) {
out.push({ key: "conversational", label: "Conversational" });
}
if (
@@ -200,6 +212,14 @@ export function detectCapabilities(
) {
out.push({ key: "embedding", label: "Embeddings" });
}
+ if (
+ hasAny(DIFFUSION_TAGS) ||
+ /stable[-_]?diffusion|\bsdxl\b|\bflux\b|qwen[-_]?image|hunyuan[-_]?(?:video|image)|wan2|latent[-_]?consistency|[-_]lcm\b|dreamshaper/.test(
+ lowerId,
+ )
+ ) {
+ out.push({ key: "diffusion", label: "Image generation" });
+ }
const languageCodes = new Set();
for (const tag of tags ?? []) {
const lower = tag.toLowerCase();
diff --git a/studio/frontend/src/features/hub/lib/model-type-filter.ts b/studio/frontend/src/features/hub/lib/model-type-filter.ts
new file mode 100644
index 0000000000..2a2e31437e
--- /dev/null
+++ b/studio/frontend/src/features/hub/lib/model-type-filter.ts
@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// Type filter for the On Device list. Mirrors the hub Discover capability
+// options and shares its detection, so both dropdowns behave the same.
+
+import type {
+ CachedInventoryRow,
+ LocalInventoryRow,
+} from "@/features/hub/inventory/types";
+import { type CapabilityKey, detectCapabilities } from "./model-capabilities";
+
+export type ModelTypeFilter =
+ | "all"
+ | "reasoning"
+ | "vision"
+ | "audio"
+ | "embedding"
+ | "diffusion";
+
+export const MODEL_TYPE_FILTER_OPTIONS: ReadonlyArray<{
+ value: ModelTypeFilter;
+ label: string;
+}> = [
+ { value: "all", label: "All types" },
+ { value: "reasoning", label: "Reasoning" },
+ { value: "vision", label: "Vision" },
+ { value: "audio", label: "Audio" },
+ { value: "embedding", label: "Embeddings" },
+ { value: "diffusion", label: "Image generation" },
+];
+
+function rowName(row: CachedInventoryRow | LocalInventoryRow): string {
+ return row.kind === "local"
+ ? `${row.id} ${row.repoId ?? ""} ${row.title} ${row.modelId ?? ""}`
+ : `${row.id} ${row.repoId}`;
+}
+
+export function matchesModelType(
+ row: CachedInventoryRow | LocalInventoryRow,
+ filter: ModelTypeFilter,
+): boolean {
+ if (filter === "all") return true;
+ // Honor the row's own vision flag before falling back to tag detection.
+ if (filter === "vision" && row.capabilities.supportsVision) return true;
+ const caps = detectCapabilities(
+ row.tags ?? undefined,
+ row.pipelineTag ?? undefined,
+ rowName(row),
+ );
+ return caps.some((cap: { key: CapabilityKey }) => cap.key === filter);
+}
diff --git a/studio/frontend/src/features/hub/lib/view-models.ts b/studio/frontend/src/features/hub/lib/view-models.ts
index 9ee6c5de5d..d4efdb6a65 100644
--- a/studio/frontend/src/features/hub/lib/view-models.ts
+++ b/studio/frontend/src/features/hub/lib/view-models.ts
@@ -7,18 +7,18 @@ import type {
CachedInventoryRow,
LocalInventoryRow,
} from "@/features/hub/inventory/types";
+import { ownerOf, repoOf } from "@/features/hub/lib/format";
import type {
CapabilityFilter,
DiscoverRow,
ModelFormatFilter,
} from "../types";
+import { estimateSizeFromDtypes, isGgufLike } from "./hf-model-meta";
import {
+ type CapabilityKey,
detectBaseModel,
detectCapabilities,
- type CapabilityKey,
} from "./model-capabilities";
-import { ownerOf, repoOf } from "@/features/hub/lib/format";
-import { estimateSizeFromDtypes, isGgufLike } from "./hf-model-meta";
export {
detectResultFormat,
isUnslothFinetunable,
@@ -39,6 +39,7 @@ export const CAPABILITY_FILTER_OPTIONS: ReadonlyArray<{
{ value: "vision", label: "Vision" },
{ value: "audio", label: "Audio" },
{ value: "embedding", label: "Embeddings" },
+ { value: "diffusion", label: "Image generation" },
];
export const FORMAT_FILTER_OPTIONS: ReadonlyArray<{
diff --git a/studio/frontend/src/features/model-picker/api/model-metadata.ts b/studio/frontend/src/features/model-picker/api/model-metadata.ts
new file mode 100644
index 0000000000..098ab51271
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/api/model-metadata.ts
@@ -0,0 +1,20 @@
+// 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 { getModelConfig } from "@/features/training";
+
+export async function fetchModelMaxPositionEmbeddings(
+ modelName: string,
+ hfToken?: string | null,
+ signal?: AbortSignal,
+): Promise {
+ const config = await getModelConfig(
+ modelName,
+ signal,
+ hfToken?.trim() || undefined,
+ );
+ const value = config.max_position_embeddings;
+ return typeof value === "number" && Number.isFinite(value) && value > 0
+ ? Math.floor(value)
+ : null;
+}
diff --git a/studio/frontend/src/features/model-picker/api/templates.ts b/studio/frontend/src/features/model-picker/api/templates.ts
new file mode 100644
index 0000000000..29f18cce29
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/api/templates.ts
@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { authFetch } from "@/features/auth";
+import { hubTokenHeader } from "@/features/hub";
+import { consumeNativePathToken } from "@/features/native-intents/api";
+import { readFastApiError } from "@/lib/format-fastapi-error";
+
+export interface ValidateChatTemplateResult {
+ valid: boolean;
+ error: string | null;
+}
+
+async function parseJsonOrThrow(response: Response): Promise {
+ if (!response.ok) {
+ throw new Error(await readFastApiError(response));
+ }
+ return response.json();
+}
+
+export async function validateChatTemplate(
+ template: string,
+ signal?: AbortSignal,
+): Promise {
+ const response = await authFetch("/api/picker/validate-chat-template", {
+ method: "POST",
+ headers: { "Content-Type": "application/json" },
+ body: JSON.stringify({ template }),
+ signal,
+ });
+ return parseJsonOrThrow(response);
+}
+
+export async function fetchDefaultChatTemplate(
+ modelName: string,
+ ggufVariant?: string | null,
+ hfToken?: string | null,
+ signal?: AbortSignal,
+ nativePathToken?: string | null,
+): Promise {
+ // A native (picked / drag-drop) GGUF lives at a path only its signed lease
+ // knows, and the picker chat-template GET has no lease plumbing, so redeem a
+ // one-shot validate-model lease and read the embedded template through the
+ // lease-aware /api/inference/validate probe instead (mirrors the staged
+ // header-dims fetch). Non-native models keep the plain GET path.
+ if (nativePathToken) {
+ let nativePathLease: string | null = null;
+ try {
+ nativePathLease = (
+ await consumeNativePathToken(nativePathToken, "validate-model")
+ ).nativePathLease;
+ } catch {
+ // Lease expired / revoked: no readable path, so no default template (the
+ // subsequent load re-mints its own lease).
+ return null;
+ }
+ const response = await authFetch("/api/inference/validate", {
+ method: "POST",
+ headers: { "Content-Type": "application/json" },
+ body: JSON.stringify({
+ model_path: modelName,
+ gguf_variant: ggufVariant ?? null,
+ hf_token: hfToken ?? null,
+ native_path_lease: nativePathLease,
+ include_chat_template: true,
+ }),
+ signal,
+ });
+ const data = await parseJsonOrThrow<{ chat_template?: string | null }>(
+ response,
+ );
+ return data.chat_template ?? null;
+ }
+
+ const query = ggufVariant
+ ? `?gguf_variant=${encodeURIComponent(ggufVariant)}`
+ : "";
+ const response = await authFetch(
+ `/api/picker/chat-template/${encodeURIComponent(modelName)}${query}`,
+ { headers: hubTokenHeader(hfToken), signal },
+ );
+ const data = await parseJsonOrThrow<{
+ model_name: string;
+ chat_template: string | null;
+ }>(response);
+ return data.chat_template ?? null;
+}
diff --git a/studio/frontend/src/features/model-picker/components/chat-template-editor-dialog.tsx b/studio/frontend/src/features/model-picker/components/chat-template-editor-dialog.tsx
new file mode 100644
index 0000000000..65e5a2026d
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/components/chat-template-editor-dialog.tsx
@@ -0,0 +1,191 @@
+// 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 { Button } from "@/components/ui/button";
+import {
+ Dialog,
+ DialogContent,
+ DialogDescription,
+ DialogFooter,
+ DialogHeader,
+ DialogTitle,
+} from "@/components/ui/dialog";
+import { Spinner } from "@/components/ui/spinner";
+import { Textarea } from "@/components/ui/textarea";
+import { useRef, useState } from "react";
+import { validateChatTemplate } from "../api/templates";
+import {
+ MAX_CHAT_TEMPLATE_BYTES,
+ chatTemplateByteLength,
+ isChatTemplateWithinLimit,
+} from "../model-config/per-model-config";
+
+interface ChatTemplateEditorDialogProps {
+ open: boolean;
+ onOpenChange: (open: boolean) => void;
+ value: string | null;
+ defaultTemplate: string | null;
+ defaultLoading: boolean;
+ onSave: (override: string | null) => void;
+ readOnly?: boolean;
+}
+
+export function ChatTemplateEditorDialog({
+ open,
+ onOpenChange,
+ value,
+ defaultTemplate,
+ defaultLoading,
+ onSave,
+ readOnly = false,
+}: ChatTemplateEditorDialogProps) {
+ const [draft, setDraft] = useState(null);
+ const [error, setError] = useState(null);
+ const [validating, setValidating] = useState(false);
+ // Bumped whenever the dialog closes so a validation still in flight cannot
+ // apply a template the user has already dismissed.
+ const validationToken = useRef(0);
+ const renderedDraft = draft ?? value ?? defaultTemplate ?? "";
+
+ const byteLength = chatTemplateByteLength(renderedDraft);
+ const overLimit = !isChatTemplateWithinLimit(renderedDraft);
+ const matchesDefault =
+ defaultTemplate != null && renderedDraft === defaultTemplate;
+
+ const handleClose = () => {
+ validationToken.current += 1;
+ setDraft(null);
+ setError(null);
+ setValidating(false);
+ onOpenChange(false);
+ };
+
+ const handleSave = async () => {
+ if (renderedDraft.trim().length === 0 || matchesDefault) {
+ onSave(null);
+ handleClose();
+ return;
+ }
+ if (overLimit) {
+ setError("Template exceeds the size limit.");
+ return;
+ }
+ setValidating(true);
+ const token = validationToken.current;
+ try {
+ const result = await validateChatTemplate(renderedDraft);
+ // Dialog was closed (or reopened) while validating; drop the result so a
+ // discarded template is never applied.
+ if (token !== validationToken.current) {
+ return;
+ }
+ if (!result.valid) {
+ setError(result.error ?? "Invalid Jinja template.");
+ return;
+ }
+ onSave(renderedDraft);
+ handleClose();
+ } catch {
+ if (token === validationToken.current) {
+ setError("Could not validate the template.");
+ }
+ } finally {
+ if (token === validationToken.current) {
+ setValidating(false);
+ }
+ }
+ };
+
+ return (
+ {
+ if (nextOpen) {
+ onOpenChange(true);
+ return;
+ }
+ handleClose();
+ }}
+ >
+
+
+
+ {readOnly ? "Chat Template" : "Edit Chat Template"}
+
+
+ {readOnly
+ ? "This is the model's chat template. Custom templates apply to GGUF models for now, so it is view only for safetensors models."
+ : "Override the model's chat template with custom Jinja. The change applies when the model loads. Saving an empty template or one that matches the default clears the override."}
+
+
+
+
+ );
+}
diff --git a/studio/frontend/src/features/model-picker/components/model-config-page.tsx b/studio/frontend/src/features/model-picker/components/model-config-page.tsx
new file mode 100644
index 0000000000..9c1bf093c1
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/components/model-config-page.tsx
@@ -0,0 +1,991 @@
+// 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 { Button } from "@/components/ui/button";
+import { Checkbox } from "@/components/ui/checkbox";
+import { InfoHint } from "@/components/ui/info-hint";
+import {
+ Select,
+ SelectContent,
+ SelectItem,
+ SelectTrigger,
+ SelectValue,
+} from "@/components/ui/select";
+import { Slider } from "@/components/ui/slider";
+import { Switch } from "@/components/ui/switch";
+import {
+ GPU_LAYERS_AUTO,
+ fetchGgufStagedMetadata,
+ readPersistedSpeculativeType,
+ useChatRuntimeStore,
+} from "@/features/chat";
+import { useGpuDevices } from "@/hooks/use-gpu-info";
+import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
+import { toast } from "@/lib/toast";
+import { ArrowLeft01Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { type ReactNode, useEffect, useId, useState } from "react";
+import {
+ useDefaultChatTemplate,
+ useModelMaxPositionEmbeddings,
+} from "../hooks/use-model-defaults";
+import { perModelConfigsEqual } from "../model-config/apply-per-model-config";
+import {
+ CONTEXT_LENGTH_MIN,
+ DEFAULT_MAX_SEQ_LENGTH,
+ DEFAULT_PER_MODEL_CONFIG,
+ KV_CACHE_DTYPES,
+ MAX_SEQ_LENGTH_MAX,
+ MAX_SEQ_LENGTH_MIN,
+ MAX_SEQ_LENGTH_STEP,
+ MTP_SPECULATIVE_TYPES,
+ type PerModelConfig,
+ SPECULATIVE_TYPES,
+ deletePerModelConfig,
+ floorMaxSeqLength,
+ isDefaultConfig,
+ normalizeMaxSeqLength,
+ resolveInitialConfig,
+ savePerModelConfig,
+} from "../model-config/per-model-config";
+import { ChatTemplateEditorDialog } from "./chat-template-editor-dialog";
+import type { ModelPickTarget } from "./model-selector/types";
+import { NumericValueInput } from "./numeric-value-input";
+
+const ROW_CLASS = "flex min-h-8 items-center justify-between gap-3";
+const LABEL_CLASS =
+ "min-w-0 truncate text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg";
+const LABEL_CLASS_WRAP =
+ "min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg";
+const CONTROL_SURFACE =
+ "rounded-full border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.1]";
+const SELECT_TRIGGER_CLASS = `grid h-8 min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 ${CONTROL_SURFACE} pl-3 pr-2 py-0 text-[13px]! font-medium text-nav-fg focus-visible:ring-0 focus-visible:border-transparent [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0`;
+const NUMBER_INPUT_CLASS = `h-8 w-[92px] ${CONTROL_SURFACE} pl-3 pr-2 py-0 text-right text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0`;
+
+const KV_CACHE_DTYPE_DEFAULT = "f16";
+const SPECULATIVE_TYPE_LABELS: Record<(typeof SPECULATIVE_TYPES)[number], string> =
+ {
+ auto: "Auto",
+ mtp: "MTP",
+ ngram: "Ngram",
+ "mtp+ngram": "MTP+Ngram",
+ off: "Off",
+ };
+
+function hasNonDefaultAdvanced(config: PerModelConfig): boolean {
+ return (
+ config.kvCacheDtype != null ||
+ (config.speculativeType ?? "auto") !== "auto" ||
+ config.specDraftNMax != null ||
+ config.tensorParallel ||
+ config.chatTemplateOverride != null ||
+ (config.gpuMemoryMode ?? "auto") !== "auto" ||
+ (config.gpuLayers != null && config.gpuLayers >= 0) ||
+ (config.nCpuMoe ?? 0) > 0 ||
+ config.selectedGpuIds != null
+ );
+}
+
+function ChatTemplateSetting({
+ config,
+ onEditTemplate,
+ readOnly = false,
+}: {
+ config: PerModelConfig;
+ onEditTemplate: () => void;
+ readOnly?: boolean;
+}) {
+ return (
+
+
+ Chat Template
+
+ {readOnly
+ ? "Preview the model's chat template. Custom overrides apply to GGUF models for now."
+ : "Override the model's chat template with custom Jinja. Applies when the model loads."}
+
+
+
+ {readOnly ? null : (
+
+ {config.chatTemplateOverride ? "Custom" : "Default"}
+
+ )}
+
+ {readOnly ? "View" : "Edit"}
+
+
+
+ );
+}
+
+function MaxSeqLengthSetting({
+ value,
+ max,
+ inputMax,
+ onChange,
+}: {
+ value: number;
+ max: number;
+ inputMax: number;
+ onChange: (value: number) => void;
+}) {
+ return (
+
+
+
+ Max Seq Length
+
+ Maximum context window size in tokens. Applies when the model loads.
+
+
+
+
+
onChange(next)}
+ className="panel-slider"
+ aria-label="Max Seq Length"
+ />
+
+ );
+}
+
+function clampMaxSeqLength(value: number, max: number): number {
+ const normalized = normalizeMaxSeqLength(value) ?? MAX_SEQ_LENGTH_MIN;
+ return Math.max(MAX_SEQ_LENGTH_MIN, Math.min(max, normalized));
+}
+
+function AdvancedGpuSlider({
+ label,
+ value,
+ min,
+ max,
+ onChange,
+ displayValue,
+ info,
+}: {
+ label: string;
+ value: number;
+ min: number;
+ max: number;
+ onChange: (value: number) => void;
+ displayValue?: string;
+ info?: ReactNode;
+}) {
+ return (
+
+
+
+ {label}
+ {info && {info} }
+
+
+
+
onChange(next)}
+ className="panel-slider"
+ aria-label={label}
+ />
+
+ );
+}
+
+// GPU Memory placement controls (mode / GPU Layers / MoE offload / GPU picker),
+// GGUF only. Slider ceilings come from the GGUF header dims, the picker from the
+// live device list. --tensor-split is not persisted per model, so not exposed here.
+function GpuMemorySettings({
+ config,
+ update,
+ layerCount,
+ moeLayerCount,
+}: {
+ config: PerModelConfig;
+ update: (patch: Partial) => void;
+ layerCount: number | null;
+ moeLayerCount: number | null;
+}) {
+ const gpuDevices = useGpuDevices();
+ const mode = config.gpuMemoryMode ?? "auto";
+ const isManual = mode === "manual";
+ const gpuLayers = config.gpuLayers ?? GPU_LAYERS_AUTO;
+ // Slider at Auto: llama.cpp --fit owns the layout, so MoE-offload doesn't apply.
+ const autoLayers = isManual && gpuLayers < 0;
+ // Ceiling = layer count + 1 (llama.cpp counts the output layer as offloadable),
+ // else a safe fallback.
+ const gpuLayersMax = layerCount != null ? layerCount + 1 : 256;
+ const nCpuMoe = config.nCpuMoe ?? 0;
+ const moeLayersMax = moeLayerCount ?? 0;
+ const showMoeSlider = isManual && !autoLayers && moeLayersMax > 0;
+ const selectedGpuIds = config.selectedGpuIds ?? null;
+ const singleGpuInUse =
+ (selectedGpuIds ?? gpuDevices.map((device) => device.index)).length <= 1;
+ // Multi-GPU only, and only with physical indices (relative ordinals from a
+ // CUDA_VISIBLE_DEVICES mask can't be mapped back to pin a device). null = all (auto).
+ const showGpuPicker =
+ gpuDevices.length > 1 && gpuDevices.every((d) => d.physicalIndex);
+ const isGpuChecked = (index: number) =>
+ selectedGpuIds === null || selectedGpuIds.includes(index);
+ const toggleGpu = (index: number) => {
+ const all = gpuDevices.map((d) => d.index);
+ const current = selectedGpuIds ?? all;
+ const next = current.includes(index)
+ ? current.filter((i) => i !== index)
+ : [...current, index].sort((a, b) => a - b);
+ if (next.length === 0) return; // keep at least one GPU selected
+ update({ selectedGpuIds: next.length === all.length ? null : next });
+ };
+ return (
+ <>
+
+
+
GPU Memory
+
+
+
+ Default: Unsloth fits the
+ model and context to your GPUs.
+
+
+ Manual: set GPU Layers
+ yourself. Leave it on Auto to let llama.cpp size the context and
+ offload overflow (including MoE experts) to RAM.
+
+
+
+
+
+ // Returning to Default must clear the Manual-only knobs, else a
+ // remembered config keeps stale gpuLayers/nCpuMoe/GPU pick that a
+ // later load re-applies while the page shows Default.
+ update(
+ v === "manual"
+ ? { gpuMemoryMode: "manual" }
+ : {
+ gpuMemoryMode: "auto",
+ gpuLayers: undefined,
+ nCpuMoe: undefined,
+ selectedGpuIds: undefined,
+ },
+ )
+ }
+ >
+
+
+
+
+ Default
+ Manual
+
+
+
+ {isManual && (
+ <>
+ update({ gpuLayers: v })}
+ displayValue={autoLayers ? "Auto" : undefined}
+ info={
+ <>
+ Layers to keep on the GPU (--gpu-layers); the rest run on CPU.
+ Auto lets llama.cpp size the split (and the context) to fit VRAM.
+ At the maximum, the whole model is on the GPU.
+ >
+ }
+ />
+ {showMoeSlider && (
+ update({ nCpuMoe: v })}
+ info={
+ <>
+ Keep the experts of this many MoE layers on the CPU
+ (--n-cpu-moe) to save VRAM. 0 = all experts on the GPU; at the
+ maximum, all are on the CPU.
+ >
+ }
+ />
+ )}
+ >
+ )}
+ {showGpuPicker && (
+
+
+ GPUs
+
+ Which GPUs this model may use. Unchecked GPUs are hidden from
+ llama.cpp (CUDA_VISIBLE_DEVICES, or HIP_VISIBLE_DEVICES on ROCm).
+ Leave all checked to use every GPU. At least one GPU must stay
+ selected.
+
+
+
+ {gpuDevices.map((d) => (
+
+
+ GPU {d.index}: {d.name}
+ {d.memoryTotalGb
+ ? ` · ${Math.round(d.memoryTotalGb)} GB`
+ : ""}
+
+ toggleGpu(d.index)}
+ disabled={isGpuChecked(d.index) && singleGpuInUse}
+ />
+
+ ))}
+
+
+ )}
+ >
+ );
+}
+
+function GgufAdvancedSettings({
+ config,
+ update,
+ isMtp,
+ speculativeFallback,
+ onEditTemplate,
+ layerCount,
+ moeLayerCount,
+}: {
+ config: PerModelConfig;
+ update: (patch: Partial) => void;
+ isMtp: boolean;
+ speculativeFallback: string;
+ onEditTemplate: () => void;
+ layerCount: number | null;
+ moeLayerCount: number | null;
+}) {
+ return (
+ <>
+
+
+ KV Cache Dtype
+
+ Lower KV cache precision to save VRAM at the cost of some quality.
+ f16/bf16 are full precision; q8_0/q5_1/q4_1 are quantized.
+
+
+
+ update({ kvCacheDtype: v === KV_CACHE_DTYPE_DEFAULT ? null : v })
+ }
+ >
+
+
+
+
+
+ {KV_CACHE_DTYPE_DEFAULT}
+
+ {KV_CACHE_DTYPES.map((dtype) => (
+
+ {dtype}
+
+ ))}
+
+
+
+
+
+
+ Speculative Decoding
+
+ Faster generation with no accuracy hit. Auto picks MTP / ngram based
+ on the model and platform. Pick a strategy to force it.
+
+
+
+ update({
+ speculativeType: v,
+ specDraftNMax:
+ v === "mtp" || v === "mtp+ngram" ? config.specDraftNMax : null,
+ })
+ }
+ >
+
+
+
+
+ {SPECULATIVE_TYPES.map((type) => (
+
+ {SPECULATIVE_TYPE_LABELS[type]}
+
+ ))}
+
+
+
+
+ {isMtp && (
+
+
+ Draft Tokens
+
+ Max MTP draft tokens per step. Leave blank for the platform
+ default (2 on GPU, 3 on CPU/Mac).
+
+
+
{
+ const raw = event.target.value;
+ if (raw === "") {
+ update({ specDraftNMax: null });
+ return;
+ }
+ const parsed = Number.parseInt(raw, 10);
+ if (Number.isFinite(parsed)) {
+ update({ specDraftNMax: Math.max(1, Math.min(16, parsed)) });
+ }
+ }}
+ aria-label="Speculative decoding draft tokens"
+ className={NUMBER_INPUT_CLASS}
+ />
+
+ )}
+
+
+
+ Tensor Parallelism
+
+ No effect on a single GPU. On multi-GPU setups, improves tokens/sec
+ for dense models. MoE models don't benefit.
+
+
+
update({ tensorParallel: checked })}
+ />
+
+
+
+
+
+ >
+ );
+}
+
+interface ModelConfigPageProps {
+ target: ModelPickTarget;
+ onBack?: () => void;
+ onRun: (config: PerModelConfig) => void;
+ loadedConfig?: PerModelConfig | null;
+ loadedContextLength?: number | null;
+ initialConfig?: PerModelConfig | null;
+ variant?: "page" | "sidebar";
+}
+
+export function ModelConfigPage({
+ target,
+ onBack,
+ onRun,
+ loadedConfig = null,
+ loadedContextLength = null,
+ initialConfig = null,
+ variant = "page",
+}: ModelConfigPageProps) {
+ const rememberId = useId();
+ const isActiveModel = loadedConfig != null;
+ const hfToken = useChatRuntimeStore((s) => s.hfToken);
+ const activeNativePathToken = useChatRuntimeStore(
+ (s) => s.activeNativePathToken,
+ );
+ const loadedDefaultChatTemplate = useChatRuntimeStore(
+ (s) => s.defaultChatTemplate,
+ );
+ const loadedMaxContextLength = useChatRuntimeStore(
+ (s) => s.ggufMaxContextLength,
+ );
+ const resolveInitial = () => {
+ const resolved = resolveInitialConfig(target.id, target.ggufVariant);
+ if (loadedConfig) {
+ return { config: loadedConfig, remembered: resolved.remembered };
+ }
+ if (initialConfig) {
+ return {
+ config: initialConfig,
+ remembered:
+ resolved.remembered &&
+ perModelConfigsEqual(initialConfig, resolved.config),
+ };
+ }
+ return resolved;
+ };
+ const [initial] = useState(resolveInitial);
+ const [config, setConfig] = useState(() => initial.config);
+ const [remember, setRemember] = useState(() => initial.remembered);
+ const [savedRemember, setSavedRemember] = useState(() => initial.remembered);
+ const [speculativeFallback] = useState(readPersistedSpeculativeType);
+ const [templateOpen, setTemplateOpen] = useState(false);
+ const [showAdvanced, setShowAdvanced] = useState(() =>
+ hasNonDefaultAdvanced(config),
+ );
+ const nativePathToken =
+ target.meta.nativePathToken ??
+ (isActiveModel ? activeNativePathToken : null);
+ const templateDefaults = useDefaultChatTemplate(
+ target.id,
+ target.ggufVariant,
+ templateOpen,
+ nativePathToken,
+ );
+ const modelMaxPosition = useModelMaxPositionEmbeddings(
+ target.id,
+ !target.isGguf,
+ );
+ const hasLoadedDefaultTemplate =
+ isActiveModel && loadedDefaultChatTemplate != null;
+ const resolvedDefaultTemplate = hasLoadedDefaultTemplate
+ ? loadedDefaultChatTemplate
+ : templateDefaults.template;
+ const resolvedDefaultLoading = hasLoadedDefaultTemplate
+ ? false
+ : templateDefaults.loading;
+
+ const update = (patch: Partial) =>
+ setConfig((current) => ({ ...current, ...patch }));
+
+ // Fetch GGUF header dims (context + layer/MoE counts) to size the GPU Memory
+ // sliders; the context also fills in below when target.meta lacks it.
+ const contextFetchKey = target.isGguf
+ ? `${target.id}\n${target.ggufVariant ?? ""}\n${hfToken || ""}\n${nativePathToken ?? ""}`
+ : null;
+ const [fetchedStagedDims, setFetchedStagedDims] = useState<{
+ key: string;
+ contextLength: number | null;
+ layerCount: number | null;
+ moeLayerCount: number | null;
+ } | null>(null);
+ useEffect(() => {
+ if (contextFetchKey == null) {
+ return;
+ }
+ let cancelled = false;
+ void fetchGgufStagedMetadata({
+ model_path: target.id,
+ gguf_variant: target.ggufVariant ?? null,
+ hf_token: hfToken || null,
+ nativePathToken,
+ })
+ .then((dims) => {
+ if (!cancelled) {
+ setFetchedStagedDims({ key: contextFetchKey, ...dims });
+ }
+ })
+ .catch(() => {
+ if (!cancelled) {
+ setFetchedStagedDims({
+ key: contextFetchKey,
+ contextLength: null,
+ layerCount: null,
+ moeLayerCount: null,
+ });
+ }
+ });
+ return () => {
+ cancelled = true;
+ };
+ }, [
+ contextFetchKey,
+ target.id,
+ target.ggufVariant,
+ hfToken,
+ nativePathToken,
+ ]);
+ const stagedDims =
+ fetchedStagedDims?.key === contextFetchKey ? fetchedStagedDims : null;
+
+ const isMtp =
+ config.speculativeType != null &&
+ MTP_SPECULATIVE_TYPES.has(config.speculativeType);
+ const nativeContextLength =
+ target.meta.contextLength ?? stagedDims?.contextLength ?? null;
+ const activeLoadedContext =
+ isActiveModel && target.isGguf ? loadedContextLength : null;
+ const minContext = CONTEXT_LENGTH_MIN;
+ const maxContext = Math.max(
+ minContext,
+ Math.max(
+ nativeContextLength ?? 0,
+ activeLoadedContext ?? 0,
+ config.customContextLength ?? 0,
+ ) || 32768,
+ );
+ const contextValue = Math.min(
+ Math.max(
+ config.customContextLength ??
+ activeLoadedContext ??
+ nativeContextLength ??
+ maxContext,
+ minContext,
+ ),
+ maxContext,
+ );
+ const setContextLength = (v: number) =>
+ update({ customContextLength: v });
+ const baseline = loadedConfig ?? DEFAULT_PER_MODEL_CONFIG;
+ const atBaseline = perModelConfigsEqual(config, baseline);
+ // An explicit customContextLength equal to the native ceiling is still an
+ // override (Reset stays enabled). "At default" means no override at all AND the
+ // shown context matches native (or no native context length is exposed).
+ const contextAtDefault =
+ !target.isGguf ||
+ (config.customContextLength == null &&
+ (nativeContextLength == null || contextValue === nativeContextLength));
+ const atDefault =
+ contextAtDefault &&
+ perModelConfigsEqual(
+ { ...config, customContextLength: null },
+ DEFAULT_PER_MODEL_CONFIG,
+ );
+ const nativeMaxSeqLength =
+ floorMaxSeqLength(modelMaxPosition.maxPositionEmbeddings) ??
+ MAX_SEQ_LENGTH_MAX;
+ // A non-GGUF active model seeds maxSeqLength from its loaded value. Once cleared
+ // (Reset sets null), fall back to the app default, not the loaded runtime value,
+ // else a remembered/active override can never be cleared.
+ const maxSeqLengthValue =
+ normalizeMaxSeqLength(config.maxSeqLength) ??
+ clampMaxSeqLength(DEFAULT_MAX_SEQ_LENGTH, nativeMaxSeqLength);
+ const maxSeqLengthMax = Math.max(nativeMaxSeqLength, maxSeqLengthValue);
+ // An auto-fit-below-native GGUF shows activeLoadedContext while
+ // customContextLength stays null. If the user fixes GPU Layers (Manual) and
+ // remembers, pin that shown context so a later fresh load keeps the fitted
+ // placement instead of sending native/0 for fixed layers and recreating the OOM.
+ const pinFixedLayerContext =
+ target.isGguf &&
+ config.gpuMemoryMode === "manual" &&
+ config.gpuLayers != null &&
+ config.gpuLayers >= 0 &&
+ config.customContextLength == null &&
+ activeLoadedContext != null;
+ // Persisted record: keep config as-is (non-GGUF keeps maxSeqLength null) so
+ // isDefaultConfig recognises it and clears a remembered override instead of
+ // pinning the app default.
+ const runtimeConfig = target.isGguf
+ ? pinFixedLayerContext
+ ? { ...config, customContextLength: activeLoadedContext }
+ : config
+ : config;
+ // Load request needs a concrete max length; substitute the fallback here only,
+ // never in the persisted runtimeConfig.
+ const loadConfig = target.isGguf
+ ? runtimeConfig
+ : { ...runtimeConfig, maxSeqLength: maxSeqLengthValue };
+ const rememberChanged = remember !== savedRemember;
+ const persistenceOnly = isActiveModel && atBaseline && rememberChanged;
+ const primaryActionLabel = persistenceOnly
+ ? remember
+ ? "Save settings"
+ : "Forget settings"
+ : isActiveModel
+ ? "Reload model"
+ : "Load model";
+
+ const handleRun = () => {
+ const defaultConfig = isDefaultConfig(runtimeConfig);
+ let saveFailed = false;
+ if (remember) {
+ saveFailed = !savePerModelConfig(
+ target.id,
+ target.ggufVariant,
+ runtimeConfig,
+ );
+ } else {
+ saveFailed = !deletePerModelConfig(target.id, target.ggufVariant);
+ }
+ if (persistenceOnly) {
+ if (saveFailed) {
+ toast.error("Couldn't save settings for this model.");
+ return;
+ }
+ const nextRemember = remember && !defaultConfig;
+ setSavedRemember(nextRemember);
+ setRemember(nextRemember);
+ toast.success(
+ nextRemember
+ ? "Settings saved."
+ : remember
+ ? "Default settings kept."
+ : "Settings forgotten.",
+ );
+ return;
+ }
+ if (saveFailed) {
+ toast.error("Couldn't save these settings, loading with them anyway.");
+ }
+ onRun(loadConfig);
+ };
+
+ return (
+
+ {variant === "page" && (
+
+ {onBack && (
+
+
+
+ )}
+
+
+ Run settings
+
+
+ {target.displayName}
+
+
+
+ )}
+
+
+ {target.isGguf && (
+ <>
+
+
+
+ Context Length
+
+ Tokens of context to allocate. Higher uses more VRAM.
+ {nativeContextLength != null
+ ? ` This model's native context is ${nativeContextLength.toLocaleString()} tokens.`
+ : ""}
+
+
+
+
+ {nativeContextLength != null ? (
+
setContextLength(v)}
+ className="panel-slider"
+ aria-label="Context Length"
+ />
+ ) : null}
+ {isActiveModel &&
+ loadedMaxContextLength != null &&
+ contextValue > loadedMaxContextLength && (
+
+ Exceeds estimated VRAM capacity (
+ {loadedMaxContextLength.toLocaleString()} tokens). The model
+ may use system RAM.
+
+ )}
+
+
+ {showAdvanced && (
+
setTemplateOpen(true)}
+ layerCount={stagedDims?.layerCount ?? null}
+ moeLayerCount={stagedDims?.moeLayerCount ?? null}
+ />
+ )}
+
+
+
+
+ Advanced settings
+
+
+ Extra options for how the model loads. Most setups don't need
+ these.
+
+
+
+
+ >
+ )}
+ {!target.isGguf && (
+ <>
+
+ update({
+ maxSeqLength: clampMaxSeqLength(value, MAX_SEQ_LENGTH_MAX),
+ })
+ }
+ />
+ setTemplateOpen(true)}
+ readOnly={true}
+ />
+ >
+ )}
+
+
+
+
+ setRemember(checked === true)}
+ />
+
+ Remember for this model
+
+
+
+ setConfig({ ...DEFAULT_PER_MODEL_CONFIG })}
+ >
+ Reset
+
+
+ {primaryActionLabel}
+
+
+
+
+
update({ chatTemplateOverride: override })}
+ />
+
+ );
+}
diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/features/model-picker/components/model-selector.tsx
similarity index 83%
rename from studio/frontend/src/components/assistant-ui/model-selector.tsx
rename to studio/frontend/src/features/model-picker/components/model-selector.tsx
index 6bfd1276ac..1cbce297dd 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector.tsx
+++ b/studio/frontend/src/features/model-picker/components/model-selector.tsx
@@ -3,6 +3,7 @@
"use client";
+import { Input } from "@/components/ui/input";
import {
Popover,
PopoverContent,
@@ -10,7 +11,7 @@ import {
} from "@/components/ui/popover";
import { TooltipProvider } from "@/components/ui/tooltip";
import { usePlatformStore } from "@/config/env";
-import { isCustomProviderType } from "@/features/chat/external-providers";
+import { isCustomProviderType } from "@/features/chat";
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
import { cn } from "@/lib/utils";
import {
@@ -33,7 +34,11 @@ import {
useRef,
useState,
} from "react";
-import { Input } from "../ui/input";
+import {
+ type PerModelConfig,
+ resolveInitialConfig,
+} from "../model-config/per-model-config";
+import { ModelConfigPage } from "./model-config-page";
import { HubModelPicker, hasDownloadedModels } from "./model-selector/pickers";
import { PillTabs } from "./model-selector/pill-tabs";
import {
@@ -45,6 +50,7 @@ import type {
ExternalModelOption,
LoraModelOption,
ModelOption,
+ ModelPickTarget,
ModelSelectorChangeMeta,
} from "./model-selector/types";
@@ -122,6 +128,10 @@ interface ModelSelectorProps {
value?: string;
defaultValue?: string;
activeGgufVariant?: string | null;
+ activeModelConfig?: PerModelConfig | null;
+ activeGgufContextLength?: number | null;
+ selectedConfig?: PerModelConfig | null;
+ selectedGgufVariant?: string | null;
onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
onEject?: () => void;
onFoldersChange?: () => void;
@@ -183,15 +193,12 @@ function ModelSelectorTrigger({
>
{isLoaded &&
(onEject ? (
- // Loaded status doubles as a mouse eject shortcut: green checkmark
- // at rest, red eject icon on pill hover, click to eject. A plain
- // span (no role/tabIndex) keeps it out of the trigger button's
- // content model, which forbids focusable descendants. Keyboard and
- // screen-reader users eject via the picker's "Eject model" button.
- // aria-hidden marks it decorative; stopPropagation stops the
- // popover from toggling. On touch (no hover) the eject icon and
- // tooltip never reveal, so pointer-events-none disables the
- // shortcut there and taps open the picker instead of ejecting.
+ // Loaded status doubles as a mouse eject shortcut (checkmark at rest,
+ // eject icon on hover). A plain span keeps it out of the trigger
+ // button's content model (no focusable descendants); keyboard/SR users
+ // eject via the "Eject model" button. aria-hidden marks it decorative;
+ // stopPropagation stops the popover toggling. On touch (no hover)
+ // pointer-events-none disables it so taps open the picker instead.
void;
onEject?: () => void;
onFoldersChange?: () => void;
@@ -337,8 +355,7 @@ function ModelSelectorContent({
const chatOnly = usePlatformStore((s) => s.isChatOnly());
const hasExternal = externalModels.length > 0;
// The Fine-tuned tab is for fine-tuned models only. Local models (LM Studio,
- // Ollama, custom folders) carry source "local" and live in the Hub tab's
- // Downloaded / Custom sections instead.
+ // Ollama, custom folders) carry source "local" and live in the Hub tab instead.
const fineTunedModels = useMemo(
() => loraModels.filter((model) => isFineTunedSource(model.source)),
[loraModels],
@@ -391,9 +408,12 @@ function ModelSelectorContent({
const effectiveHubSection: HubSection =
hubSection === "connected" && !hasExternal ? "recommended" : hubSection;
- // The picker below remounts on each open, but this tab state does not, so a
- // persisted selection that lands in lora/external after async load would
- // reopen on Hub. Re-derive the default tab on the open edge.
+ const [configTarget, setConfigTarget] = useState(
+ null,
+ );
+
+ // The picker remounts on each open but this tab state does not, so re-derive
+ // the default tab on the open edge (else a lora/external selection reopens on Hub).
const wasOpen = useRef(open);
useEffect(() => {
if (open && !wasOpen.current) {
@@ -402,6 +422,9 @@ function ModelSelectorContent({
// user has downloads, else their last section.
setHubSection(wantsConnectedDefault ? "connected" : defaultHubSection());
}
+ if (!open && wasOpen.current) {
+ setConfigTarget(null);
+ }
wasOpen.current = open;
}, [
open,
@@ -452,6 +475,29 @@ function ModelSelectorContent({
}
}
+ const visibleConfigTarget = open ? configTarget : null;
+ const openConfigPage = (id: string, meta: ModelSelectorChangeMeta) => {
+ const leaf = id.includes("/") ? id.slice(id.lastIndexOf("/") + 1) : id;
+ setConfigTarget({
+ id,
+ displayName: meta.ggufVariant ? `${leaf} · ${meta.ggufVariant}` : leaf,
+ ggufVariant: meta.ggufVariant ?? null,
+ isGguf: meta.isGguf ?? Boolean(meta.ggufVariant),
+ meta,
+ });
+ };
+ const handlePick = (id: string, meta: ModelSelectorChangeMeta) => {
+ if (meta.source === "external") {
+ onSelect(id, meta);
+ return;
+ }
+ const resolved = resolveInitialConfig(id, meta.ggufVariant);
+ onSelect(id, {
+ ...meta,
+ ...(resolved.remembered ? { config: resolved.config } : {}),
+ });
+ };
+
return (
@@ -477,6 +528,42 @@ function ModelSelectorContent({
skipDelayDuration={0}
disableHoverableContent={true}
>
+ {visibleConfigTarget ? (
+ setConfigTarget(null)}
+ onRun={(config) =>
+ onSelect(visibleConfigTarget.id, {
+ ...visibleConfigTarget.meta,
+ config,
+ forceReload: true,
+ })
+ }
+ loadedConfig={
+ value === visibleConfigTarget.id &&
+ (activeGgufVariant ?? null) ===
+ (visibleConfigTarget.ggufVariant ?? null)
+ ? (activeModelConfig ?? null)
+ : null
+ }
+ loadedContextLength={
+ value === visibleConfigTarget.id &&
+ (activeGgufVariant ?? null) ===
+ (visibleConfigTarget.ggufVariant ?? null)
+ ? (activeGgufContextLength ?? null)
+ : null
+ }
+ initialConfig={
+ value === visibleConfigTarget.id &&
+ (selectedGgufVariant ?? null) ===
+ (visibleConfigTarget.ggufVariant ?? null)
+ ? (selectedConfig ?? null)
+ : null
+ }
+ />
+ ) : (
+ <>
{tabs.length > 1 ? (
) : null}
- {/* Hub renders Eject inline as the last list row; other tabs keep the
- footer button. */}
{effectiveTab !== "hub" && hasSelection && onEject ? (
-
+
) : null}
+ >
+ )}
);
@@ -565,6 +653,10 @@ export function ModelSelector({
value,
defaultValue,
activeGgufVariant,
+ activeModelConfig,
+ activeGgufContextLength,
+ selectedConfig,
+ selectedGgufVariant,
onValueChange,
onEject,
onFoldersChange,
@@ -693,6 +785,11 @@ export function ModelSelector({
loraModels={loraModels}
externalModels={externalModels}
value={selected}
+ activeGgufVariant={activeGgufVariant}
+ activeModelConfig={activeModelConfig}
+ activeGgufContextLength={activeGgufContextLength}
+ selectedConfig={selectedConfig}
+ selectedGgufVariant={selectedGgufVariant}
onSelect={handleSelect}
onEject={onEject ? handleEject : undefined}
onFoldersChange={onFoldersChange}
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx b/studio/frontend/src/features/model-picker/components/model-selector/folder-browser.tsx
similarity index 84%
rename from studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx
rename to studio/frontend/src/features/model-picker/components/model-selector/folder-browser.tsx
index 16cc8a1956..6335721271 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx
+++ b/studio/frontend/src/features/model-picker/components/model-selector/folder-browser.tsx
@@ -14,10 +14,7 @@ import {
DialogTitle,
} from "@/components/ui/dialog";
import { Spinner } from "@/components/ui/spinner";
-import {
- type BrowseFoldersResponse,
- browseFolders,
-} from "@/features/chat/api/chat-api";
+import { type BrowseFoldersResponse, browseFolders } from "@/features/chat";
import { ChevronUpStandardIcon } from "@/lib/chevron-icons";
import { cn } from "@/lib/utils";
import { Folder02Icon } from "@hugeicons/core-free-icons";
@@ -35,9 +32,9 @@ export interface FolderBrowserProps {
function splitBreadcrumb(path: string): { label: string; value: string }[] {
if (!path) return [];
- // Detect path style BEFORE normalizing: on POSIX, `\` is a valid filename
- // char, so blindly rewriting `\` -> `/` mangles names like `my\backup` into
- // 404ing breadcrumbs. Only Windows-style paths (drive letter, or UNC) convert.
+ // Detect path style BEFORE normalizing: on POSIX `\` is a valid filename char,
+ // so rewriting `\` -> `/` would mangle names like `my\backup`. Only Windows
+ // paths (drive letter or UNC) convert.
const isWindowsDrive =
/^[A-Za-z]:[\\/]/.test(path) || /^[A-Za-z]:$/.test(path);
const isUnc = /^\\\\/.test(path);
@@ -58,9 +55,8 @@ function splitBreadcrumb(path: string): { label: string; value: string }[] {
return parts;
}
- // Windows drive path (C:, D:): first segment is the drive. Use `C:/` as the
- // crumb value so clicking the drive root navigates to the drive root, not the
- // drive-relative CWD (`C:` alone resolves to CWD-on-C, not `C:\`).
+ // Windows drive path: use `C:/` as the crumb value so clicking the drive root
+ // goes to the drive root, not the drive-relative CWD (`C:` alone is CWD-on-C).
if (/^[A-Za-z]:$/.test(segments[0])) {
const driveRoot = `${segments[0]}/`;
let cur = driveRoot;
@@ -90,47 +86,43 @@ export function FolderBrowser({
const [error, setError] = useState
(null);
const abortRef = useRef(null);
- const navigate = useCallback(
- (
- target: string | undefined,
- hidden: boolean,
- opts?: { fallbackOnError?: boolean },
- ) => {
- abortRef.current?.abort();
- const ctrl = new AbortController();
- abortRef.current = ctrl;
- setLoading(true);
- setError(null);
- // Forward the signal so cancelled navigation aborts the backend
- // enumeration, not just the response.
- browseFolders(target, hidden, ctrl.signal)
- .then((res) => {
- if (ctrl.signal.aborted) return;
- setData(res);
- setPath(res.current);
- })
- .catch((err) => {
- if (ctrl.signal.aborted) return;
- // Surface the error; if the first request (e.g. a bad initialPath)
- // fails, fall back to HOME so the modal stays navigable.
- const message = err instanceof Error ? err.message : String(err);
- setError(message);
- if (opts?.fallbackOnError && target !== undefined) {
- // Re-issue without a target -> backend defaults to HOME.
- // Don't recurse if HOME itself fails (allowlist always has HOME).
- queueMicrotask(() => navigate(undefined, hidden));
- }
- })
- .finally(() => {
- if (!ctrl.signal.aborted) setLoading(false);
- });
- },
- [],
- );
+ function navigate(
+ target: string | undefined,
+ hidden: boolean,
+ opts?: { fallbackOnError?: boolean },
+ ) {
+ abortRef.current?.abort();
+ const ctrl = new AbortController();
+ abortRef.current = ctrl;
+ setLoading(true);
+ setError(null);
+ // Forward the signal so cancelled navigation aborts the backend
+ // enumeration, not just the response.
+ browseFolders(target, hidden, ctrl.signal)
+ .then((res) => {
+ if (ctrl.signal.aborted) return;
+ setData(res);
+ setPath(res.current);
+ })
+ .catch((err) => {
+ if (ctrl.signal.aborted) return;
+ // Surface the error; if the first request (e.g. a bad initialPath)
+ // fails, fall back to HOME so the modal stays navigable.
+ const message = err instanceof Error ? err.message : String(err);
+ setError(message);
+ if (opts?.fallbackOnError && target !== undefined) {
+ // Re-issue without a target -> backend defaults to HOME.
+ // Don't recurse if HOME itself fails (allowlist always has HOME).
+ queueMicrotask(() => navigate(undefined, hidden));
+ }
+ })
+ .finally(() => {
+ if (!ctrl.signal.aborted) setLoading(false);
+ });
+ }
// Fetch only on closed -> open; later navigation is driven by `navigate()`,
// so `path` is deliberately kept out of the dependency list.
- // eslint-disable-next-line react-hooks/exhaustive-deps
useEffect(() => {
if (!open) return;
// fallbackOnError: recover into HOME if initialPath is bad, rather than
@@ -147,7 +139,7 @@ export function FolderBrowser({
const crumbs = useMemo(
() => (data?.current ? splitBreadcrumb(data.current) : []),
- [data?.current],
+ [data],
);
return (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-capabilities.ts b/studio/frontend/src/features/model-picker/components/model-selector/model-capabilities.ts
similarity index 100%
rename from studio/frontend/src/components/assistant-ui/model-selector/model-capabilities.ts
rename to studio/frontend/src/features/model-picker/components/model-selector/model-capabilities.ts
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx b/studio/frontend/src/features/model-picker/components/model-selector/model-delete-action.tsx
similarity index 90%
rename from studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx
rename to studio/frontend/src/features/model-picker/components/model-selector/model-delete-action.tsx
index 4de96d3648..09bb43abdc 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx
+++ b/studio/frontend/src/features/model-picker/components/model-selector/model-delete-action.tsx
@@ -1,12 +1,12 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-import { DeleteConfirmDialog } from "@/features/hub/catalog/download-card";
+import { DeleteConfirmDialog } from "@/features/hub";
+import { toast } from "@/lib/toast";
import { cn } from "@/lib/utils";
import { Delete02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
-import { useCallback, useState, type ReactNode } from "react";
-import { toast } from "@/lib/toast";
+import { type ReactNode, useCallback, useState } from "react";
interface ModelDeleteActionProps {
ariaLabel: string;
@@ -63,7 +63,8 @@ export function ModelDeleteAction({
disabled={disabled}
className={cn(
"shrink-0 rounded-md p-1.5 text-muted-foreground/60 transition-colors hover:bg-destructive/10 hover:text-destructive",
- disabled && "cursor-not-allowed opacity-40 hover:bg-transparent hover:text-muted-foreground/60",
+ disabled &&
+ "cursor-not-allowed opacity-40 hover:bg-transparent hover:text-muted-foreground/60",
buttonClassName,
)}
>
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-load-settings-action.tsx b/studio/frontend/src/features/model-picker/components/model-selector/model-load-settings-action.tsx
similarity index 63%
rename from studio/frontend/src/components/assistant-ui/model-selector/model-load-settings-action.tsx
rename to studio/frontend/src/features/model-picker/components/model-selector/model-load-settings-action.tsx
index 58510762d4..bbef42063f 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/model-load-settings-action.tsx
+++ b/studio/frontend/src/features/model-picker/components/model-selector/model-load-settings-action.tsx
@@ -6,24 +6,18 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
-import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { cn } from "@/lib/utils";
import { Settings02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
-/** Gear button on a downloaded quant row. Stages the model into the Run
- * settings sidebar (always, regardless of the Load-on-selection toggle) so the
- * user can set load options, then click Load model. */
export function ModelLoadSettingsAction({
ariaLabel,
- repoId,
- quant,
- maxContext,
+ onConfigure,
+ className,
}: {
ariaLabel: string;
- repoId: string;
- quant: string;
- maxContext?: number | null;
+ onConfigure: () => void;
+ className?: string;
}) {
return (
@@ -32,16 +26,12 @@ export function ModelLoadSettingsAction({
type="button"
onClick={(e) => {
e.stopPropagation();
- useChatRuntimeStore.getState().stageModel({
- id: repoId,
- ggufVariant: quant,
- isDownloaded: true,
- contextLength: maxContext ?? null,
- });
+ onConfigure();
}}
aria-label={ariaLabel}
className={cn(
- "shrink-0 rounded-md p-1 text-muted-foreground/60 transition-colors hover:bg-muted hover:text-foreground",
+ "shrink-0 rounded-md p-1 text-muted-foreground/60 transition-colors hover:bg-black/5 hover:text-foreground dark:hover:bg-white/10",
+ className,
)}
>
void;
+}
+
+interface ModelRowMenuUpdate {
+ title: string;
+ description: ReactNode;
+ /** Repo + variant the update targets (see ModelUpdateAction). */
+ repoId: string;
+ variant?: string | null;
+ disabled?: boolean;
+ onConfirm: () => Promise | void;
+ onUpdated?: () => void;
+}
+
+interface ModelRowMenuDelete {
+ title: string;
+ description: ReactNode;
+ successMessage: string;
+ disabled?: boolean;
+ onConfirm: () => Promise | void;
+ onDeleted?: () => void;
+}
+
+/** Managed-cache location for "Reveal in Finder" (resolved server-side). */
+interface ModelRowMenuCachePath {
+ repoId: string;
+ variant?: string;
+}
+
+export function ModelRowMenu({
+ ariaLabel,
+ buttonClassName,
+ iconClassName,
+ cachePath,
+ pin,
+ update,
+ del,
+}: {
+ ariaLabel: string;
+ buttonClassName?: string;
+ iconClassName?: string;
+ /** Enables "Reveal in Finder" for cached repos. */
+ cachePath?: ModelRowMenuCachePath;
+ pin?: ModelRowMenuPin;
+ update?: ModelRowMenuUpdate;
+ del?: ModelRowMenuDelete;
+}) {
+ const deviceType = usePlatformStore((s) => s.deviceType);
+ const revealLabel =
+ deviceType === "mac"
+ ? "Reveal in Finder"
+ : deviceType === "windows"
+ ? "Reveal in File Explorer"
+ : "Reveal in File Manager";
+ const [deleteOpen, setDeleteOpen] = useState(false);
+ const [deleting, setDeleting] = useState(false);
+ const [updateOpen, setUpdateOpen] = useState(false);
+
+ // Refresh the caller when this repo+variant's managed update completes
+ // (mirrors ModelUpdateAction).
+ const onUpdatedRef = useRef(update?.onUpdated);
+ useEffect(() => {
+ onUpdatedRef.current = update?.onUpdated;
+ }, [update?.onUpdated]);
+ const updateRepoId = update?.repoId;
+ const updateVariant = update?.variant ?? null;
+ useEffect(() => {
+ if (!updateRepoId) return;
+ return subscribeJobListeners("model", updateRepoId, {
+ onComplete: (completedVariant) => {
+ const matches = updateVariant
+ ? ggufVariantsMatch(completedVariant, updateVariant)
+ : !completedVariant;
+ if (matches) onUpdatedRef.current?.();
+ },
+ });
+ }, [updateRepoId, updateVariant]);
+
+ const onDeleteConfirm = del?.onConfirm;
+ const onDeleted = del?.onDeleted;
+ const deleteSuccessMessage = del?.successMessage;
+ const handleDeleteConfirm = useCallback(async () => {
+ if (!onDeleteConfirm) return;
+ setDeleting(true);
+ try {
+ await onDeleteConfirm();
+ if (deleteSuccessMessage) toast.success(deleteSuccessMessage);
+ onDeleted?.();
+ setDeleteOpen(false);
+ } catch (err) {
+ toast.error(
+ err instanceof Error ? err.message : "Failed to delete model",
+ );
+ } finally {
+ setDeleting(false);
+ }
+ }, [onDeleteConfirm, onDeleted, deleteSuccessMessage]);
+
+ const onUpdateConfirm = update?.onConfirm;
+ const handleUpdateConfirm = useCallback(() => {
+ // Start the re-download and close the dialog; the Downloads panel owns
+ // progress + cancel. Only a failure to START toasts.
+ void Promise.resolve()
+ .then(onUpdateConfirm)
+ .catch((err) => {
+ toast.error(
+ err instanceof Error ? err.message : "Failed to start update",
+ );
+ });
+ setUpdateOpen(false);
+ }, [onUpdateConfirm]);
+
+ const cachePathRepoId = cachePath?.repoId;
+ const cachePathVariant = cachePath?.variant;
+ const handleReveal = useCallback(() => {
+ if (!cachePathRepoId) return;
+ revealCachedModel(cachePathRepoId, cachePathVariant).catch((err) => {
+ toast.error(
+ err instanceof Error ? err.message : "Failed to open file manager",
+ );
+ });
+ }, [cachePathRepoId, cachePathVariant]);
+
+ if (!pin && !update && !del && !cachePath) return null;
+
+ return (
+ <>
+
+
+ e.stopPropagation()}
+ aria-label={ariaLabel}
+ className={cn(
+ "shrink-0 rounded-md p-1 text-muted-foreground/60 transition-colors hover:bg-black/5 hover:text-foreground dark:hover:bg-white/10",
+ buttonClassName,
+ )}
+ >
+
+
+
+
+ {pin && (
+ {
+ e.stopPropagation();
+ pin.onToggle();
+ }}
+ >
+
+ {pin.pinned ? pin.unpinLabel : pin.pinLabel}
+
+ )}
+ {cachePath && (
+ {
+ e.stopPropagation();
+ handleReveal();
+ }}
+ >
+
+ {revealLabel}
+
+ )}
+ {update && (
+ {
+ e.stopPropagation();
+ setUpdateOpen(true);
+ }}
+ >
+
+ Update
+
+ )}
+ {del && (
+ <>
+ {(cachePath || pin || update) && }
+ {
+ e.stopPropagation();
+ setDeleteOpen(true);
+ }}
+ >
+
+ Delete
+
+ >
+ )}
+
+
+
+ {del && (
+ {
+ if (!nextOpen && deleting) return;
+ setDeleteOpen(nextOpen);
+ }}
+ title={del.title}
+ description={del.description}
+ deleting={deleting}
+ onConfirm={() => void handleDeleteConfirm()}
+ />
+ )}
+
+ {update && (
+
+ )}
+ >
+ );
+}
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx b/studio/frontend/src/features/model-picker/components/model-selector/model-update-action.tsx
similarity index 82%
rename from studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx
rename to studio/frontend/src/features/model-picker/components/model-selector/model-update-action.tsx
index db7628777a..b13ed33d04 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx
+++ b/studio/frontend/src/features/model-picker/components/model-selector/model-update-action.tsx
@@ -1,12 +1,20 @@
// 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 { subscribeJobListeners } from "@/features/hub/download-manager";
-import { UpdateConfirmDialog } from "@/features/hub/catalog/download-card";
-import { ggufVariantsMatch } from "@/features/hub/lib/model-identity";
+import {
+ UpdateConfirmDialog,
+ ggufVariantsMatch,
+ subscribeJobListeners,
+} from "@/features/hub";
import { cn } from "@/lib/utils";
import { RefreshCw } from "lucide-react";
-import { useCallback, useEffect, useRef, useState, type ReactNode } from "react";
+import {
+ type ReactNode,
+ useCallback,
+ useEffect,
+ useRef,
+ useState,
+} from "react";
import { toast } from "sonner";
interface ModelUpdateActionProps {
@@ -42,10 +50,10 @@ export function ModelUpdateAction({
}: ModelUpdateActionProps) {
const [open, setOpen] = useState(false);
- // Refresh the caller when this repo+variant's download finishes so the "update available" cue
- // clears. A ref keeps the subscription stable across renders.
const onUpdatedRef = useRef(onUpdated);
- onUpdatedRef.current = onUpdated;
+ useEffect(() => {
+ onUpdatedRef.current = onUpdated;
+ }, [onUpdated]);
useEffect(() => {
return subscribeJobListeners("model", repoId, {
onComplete: (completedVariant) => {
@@ -83,7 +91,8 @@ export function ModelUpdateAction({
disabled={disabled}
className={cn(
"shrink-0 rounded-md p-1.5 text-muted-foreground/60 transition-colors hover:bg-amber-500/10 hover:text-amber-700 dark:hover:bg-amber-500/15 dark:hover:text-amber-300",
- disabled && "cursor-not-allowed opacity-40 hover:bg-transparent hover:text-muted-foreground/60",
+ disabled &&
+ "cursor-not-allowed opacity-40 hover:bg-transparent hover:text-muted-foreground/60",
buttonClassName,
)}
>
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-usage.ts b/studio/frontend/src/features/model-picker/components/model-selector/model-usage.ts
similarity index 93%
rename from studio/frontend/src/components/assistant-ui/model-selector/model-usage.ts
rename to studio/frontend/src/features/model-picker/components/model-selector/model-usage.ts
index dbcd4b9a1b..c6665e7658 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/model-usage.ts
+++ b/studio/frontend/src/features/model-picker/components/model-selector/model-usage.ts
@@ -40,7 +40,8 @@ export function loadedAt(times: ModelLoadTimes, id: string): number {
export function useModelLoadTimes(currentValue?: string): ModelLoadTimes {
const [times, setTimes] = useState(() => readLoadTimes());
useEffect(() => {
- if (currentValue) setTimes(recordModelLoaded(currentValue));
+ if (!currentValue) return;
+ queueMicrotask(() => setTimes(recordModelLoaded(currentValue)));
}, [currentValue]);
return times;
}
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx
similarity index 58%
rename from studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
rename to studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx
index 84119cc992..4df87b876a 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx
@@ -10,49 +10,48 @@ import {
TooltipTrigger,
} from "@/components/ui/tooltip";
import { usePlatformStore } from "@/config/env";
-import { ApiProviderLogo } from "@/features/chat/api-provider-logo";
+import { ApiProviderLogo } from "@/features/chat";
import {
type ScanFolderInfo,
addScanFolder,
- deleteCachedModel,
deleteFineTunedModel,
- listCachedGguf,
- listCachedModels,
listGgufVariants,
- listLocalModels,
listRecommendedFolders,
listScanFolders,
removeScanFolder,
-} from "@/features/chat/api/chat-api";
-import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
+} from "@/features/chat";
+import { useChatRuntimeStore } from "@/features/chat";
import type {
CachedGgufRepo,
CachedModelRepo,
+ GgufVariantDetail,
LocalModelInfo,
-} from "@/features/chat/api/chat-api";
-import type { GgufVariantDetail } from "@/features/chat/types/api";
-import { DotTag } from "@/features/hub/catalog/dot-tag";
+} from "@/features/chat";
import {
+ DotTag,
type HubOption,
HubOptionMenu,
-} from "@/features/hub/catalog/hub-option-menu";
-import { TransportConflictDialog } from "@/features/hub/catalog/transport-conflict-dialog";
-import { TrainIcon } from "@/features/hub/components/train-icon";
-import { useHubInfiniteScroll } from "@/features/hub/hooks/use-hub-infinite-scroll";
+ TrainIcon,
+ TransportConflictDialog,
+ deleteCachedModel,
+ listGgufVariants as listGgufVariantsCached,
+ useHubInfiniteScroll,
+} from "@/features/hub";
import {
type HfModelResult,
type HfSortKey,
useHubModelSearch,
-} from "@/features/hub/hooks/use-hub-model-search";
-import { useOnlineStatus } from "@/features/hub/hooks/use-online-status";
-import { isHiddenModelId } from "@/features/hub/lib/hidden-models";
-import { classifyUnslothSupport } from "@/features/hub/lib/unsloth-support";
-import { hfApiToken, useHfTokenStore } from "@/features/hub/stores/hf-token-store";
+} from "@/features/hub";
import {
+ classifyUnslothSupport,
downloadManager,
+ hfApiToken,
+ isHiddenModelId,
jobKeyOf,
useDownloadManagerStore,
-} from "@/features/hub/download-manager";
+ useHfTokenStore,
+ useOnlineStatus,
+} from "@/features/hub";
import { useDebouncedValue, useGpuInfo } from "@/hooks";
import { extractParamLabel } from "@/lib/model-size";
import { toast } from "@/lib/toast";
@@ -61,6 +60,7 @@ import type { VramFitStatus } from "@/lib/vram";
import { checkVramFit, estimateLoadingVram } from "@/lib/vram";
import {
Add01Icon,
+ ArrowUpDownIcon,
AudioWave01Icon,
Cancel01Icon,
DashboardCircleIcon,
@@ -68,7 +68,6 @@ import {
Flag01Icon,
Folder02Icon,
PinIcon,
- PinOffIcon,
RemoveCircleIcon,
Search01Icon,
ViewIcon,
@@ -87,6 +86,7 @@ import {
useRef,
useState,
} from "react";
+import { useChatPickerInventory } from "../../inventory/use-chat-picker-inventory";
import { FolderBrowser } from "./folder-browser";
import {
type ModelCapabilities,
@@ -94,14 +94,15 @@ import {
hasAnyCapability,
} from "./model-capabilities";
import { ModelDeleteAction } from "./model-delete-action";
-import { ModelUpdateAction } from "./model-update-action";
import { ModelLoadSettingsAction } from "./model-load-settings-action";
+import { ModelRowMenu } from "./model-row-menu";
import {
type ModelLoadTimes,
loadedAt,
useModelLoadTimes,
} from "./model-usage";
import {
+ makePinRank,
pinKey,
pinnedQuantEntries,
usePinnedModelsStore,
@@ -349,15 +350,12 @@ function ListLabel({
/** Format bytes to a human-readable size string. */
function formatBytes(bytes: number): string {
- // Guard non-positive / non-finite sizes (0, missing -> NaN, Infinity) so we
- // never render "NaN undefined" or a negative unit index.
+ // Guard non-positive / non-finite sizes so we never render "NaN undefined".
if (!Number.isFinite(bytes) || bytes <= 0) return "0 B";
- // Decimal (base-1000) units to match what Hugging Face reports for a repo's
- // file sizes -- e.g. 217 GB, not the 201.8 GiB a base-1024 divide would show.
- // (GPU-fit math below stays base-1024 since VRAM is binary.)
- // Divide iteratively rather than via Math.log, which has float error at exact
- // powers of 1000 (log(1e12)/log(1000) = 3.9999... would mislabel 1 TB as
- // "1000 GB"); the loop also can't run off the end of units.
+ // Decimal (base-1000) units to match Hugging Face's reported file sizes (GPU-fit
+ // math below stays base-1024 since VRAM is binary). Divide iteratively rather
+ // than via Math.log, which has float error at exact powers of 1000 (mislabeling
+ // 1 TB as "1000 GB") and could run off the end of units.
const units = ["B", "KB", "MB", "GB", "TB"];
let i = 0;
let value = bytes;
@@ -422,8 +420,7 @@ function ggufVariantsMatchForPicker(
right: string | null | undefined,
): boolean {
return (
- normalizeGgufVariantForPicker(left) ===
- normalizeGgufVariantForPicker(right)
+ normalizeGgufVariantForPicker(left) === normalizeGgufVariantForPicker(right)
);
}
@@ -674,12 +671,17 @@ function isValidGgufVariant(variant: unknown): variant is GgufVariantDetail {
);
}
-function normalizeGgufVariantsResponse(res: {
- variants?: unknown;
- default_variant?: unknown;
- has_vision?: unknown;
- context_length?: unknown;
-} | null | undefined): {
+function normalizeGgufVariantsResponse(
+ res:
+ | {
+ variants?: unknown;
+ default_variant?: unknown;
+ has_vision?: unknown;
+ context_length?: unknown;
+ }
+ | null
+ | undefined,
+): {
variants: GgufVariantDetail[];
defaultVariant: string | null;
hasVision: boolean;
@@ -722,6 +724,7 @@ function GgufVariantExpander({
parentOptionKey,
onNavigatePastStart,
onNavigatePastEnd,
+ onConfigure,
sourceOverride,
variantActions,
onDevice = false,
@@ -738,6 +741,7 @@ function GgufVariantExpander({
parentOptionKey?: string;
onNavigatePastStart?: () => void;
onNavigatePastEnd?: () => void;
+ onConfigure?: (id: string, meta: ModelSelectorChangeMeta) => void;
sourceOverride?: ModelSelectorChangeMeta["source"];
/** Update/delete actions for cached variant rows. Omitted by browse-only
* expanders (Recommended, etc.) that don't manage on-disk variants. */
@@ -765,13 +769,18 @@ function GgufVariantExpander({
const pinnedKeys = usePinnedModelsStore((s) => s.pinned);
const togglePinnedQuant = usePinnedModelsStore((s) => s.togglePinned);
const onUpdateVariant = variantActions?.onUpdate;
- const updateVariantTitle = variantActions?.updateTitle ?? "Update cached model?";
- const renderUpdateVariantDescription = variantActions?.renderUpdateDescription;
+ const updateVariantTitle =
+ variantActions?.updateTitle ?? "Update cached model?";
+ const renderUpdateVariantDescription =
+ variantActions?.renderUpdateDescription;
const updateDisabled = variantActions?.updateDisabled ?? false;
const onDeleteVariant = variantActions?.onDelete;
- const deleteVariantTitle = variantActions?.deleteTitle ?? "Delete cached model?";
- const renderDeleteVariantDescription = variantActions?.renderDeleteDescription;
- const getDeleteVariantSuccessMessage = variantActions?.getDeleteSuccessMessage;
+ const deleteVariantTitle =
+ variantActions?.deleteTitle ?? "Delete cached model?";
+ const renderDeleteVariantDescription =
+ variantActions?.renderDeleteDescription;
+ const getDeleteVariantSuccessMessage =
+ variantActions?.getDeleteSuccessMessage;
const deleteDisabled = variantActions?.deleteDisabled ?? false;
const [variants, setVariants] = useState(null);
const [defaultVariant, setDefaultVariant] = useState(null);
@@ -784,8 +793,11 @@ function GgufVariantExpander({
useEffect(() => {
let canceled = false;
- setLoading(true);
- setError(null);
+ queueMicrotask(() => {
+ if (canceled) return;
+ setLoading(true);
+ setError(null);
+ });
listGgufVariants(repoId, hfToken)
.then((res) => {
@@ -813,17 +825,12 @@ function GgufVariantExpander({
}, [repoId, refreshKey, hfToken]);
// Covers Unix absolute (/), Windows drive (C:\, D:/), UNC (\\server), relative (./, ../), tilde (~/)
- const isLocalPath = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test(
+ const isLocalPath = /^(\/|\.{1,2}[\\/]|~[\\/]|[A-Za-z]:[\\/]|\\\\)/.test(
repoId,
);
const handleVariantClick = useCallback(
(quant: string, downloaded?: boolean, sizeBytes?: number) => {
- // Only seed the staged context for picks whose weights are already on
- // disk. The staging effect short-circuits on a known contextLength
- // (pendingHasContext) before starting the download, so attaching it to an
- // undownloaded quant from a partially cached repo would skip the download
- // entirely (and, with Load on selection, never load).
const isAvailable = isLocalPath || downloaded === true;
onSelect(repoId, {
source: sourceOverride ?? (isLocalPath ? "local" : "hub"),
@@ -832,6 +839,7 @@ function GgufVariantExpander({
isDownloaded: isLocalPath ? true : downloaded,
expectedBytes: sizeBytes,
contextLength: isAvailable ? nativeContext : undefined,
+ isGguf: true,
});
},
[repoId, isLocalPath, onSelect, sourceOverride, nativeContext],
@@ -846,13 +854,12 @@ function GgufVariantExpander({
const getGgufFit = useCallback(
(sizeBytes: number): "fits" | "tight" | "oom" => {
- // No device budget at all (no GPU and no known system RAM): can't
- // classify, so don't scare the user with OOM badges.
+ // No device budget at all: can't classify, so don't show OOM badges.
if (totalBudgetGb <= 0) return "fits";
const gb = sizeBytes / 1024 ** 3;
if (gb <= 0 || gb <= gpuBudgetGb) return "fits";
- // No-GPU / unified-memory hosts (Mac) have only the RAM budget, so the
- // tier collapses to fit-or-oom against system RAM rather than GPU+offload.
+ // No-GPU / unified-memory hosts (Mac) have only the RAM budget, so the tier
+ // collapses to fit-or-oom against system RAM.
if (gpuBudgetGb <= 0) return gb <= totalBudgetGb ? "fits" : "oom";
if (gb <= totalBudgetGb) return "tight";
return "oom";
@@ -870,17 +877,13 @@ function GgufVariantExpander({
if (defaultV && getGgufFit(defaultV.size_bytes) !== "oom")
return defaultVariant;
// Largest non-OOM variant (best quality that fits)
- const fitting = variants.filter(
- (v) => getGgufFit(v.size_bytes) !== "oom",
- );
+ const fitting = variants.filter((v) => getGgufFit(v.size_bytes) !== "oom");
if (fitting.length > 0) {
fitting.sort((a, b) => b.size_bytes - a.size_bytes);
return fitting[0].quant;
}
// All OOM -- recommend smallest (most likely to partially run)
- const sorted = [...variants].sort(
- (a, b) => a.size_bytes - b.size_bytes,
- );
+ const sorted = [...variants].sort((a, b) => a.size_bytes - b.size_bytes);
return sorted[0]?.quant ?? defaultVariant;
}, [variants, defaultVariant, totalBudgetGb, getGgufFit]);
@@ -996,7 +999,7 @@ function GgufVariantExpander({
const keyBase = `${repoId}:${v.filename}`;
const variantOptionKey = makeModelOptionKey("gguf-variant", keyBase);
return (
-
+
@@ -1022,7 +1025,7 @@ function GgufVariantExpander({
update available
- ): null}
+ ) : null}
>
) : v.quant === effectiveRecommended ? (
@@ -1046,106 +1049,104 @@ function GgufVariantExpander({
- {v.downloaded && v.update_available && onUpdateVariant && (
-
- This will update{" "}
-
- {repoId} ({v.quant})
- {"."}
- >
- )
- }
- repoId={repoId}
- variant={v.quant}
- buttonClassName="p-1"
- iconClassName="size-3"
- disabled={updateDisabled}
- onConfirm={() => onUpdateVariant(v.quant, expectedBytes)}
- onUpdated={() => setRefreshKey((key) => key + 1)}
- />
- )}
- {v.downloaded && allowPin && (
-
-
- togglePinnedQuant(repoId, v.quant)}
- aria-label={
- pinnedKeys.includes(pinKey(repoId, v.quant))
- ? `Unpin ${repoId} ${v.quant}`
- : `Pin ${repoId} ${v.quant}`
- }
- aria-pressed={pinnedKeys.includes(pinKey(repoId, v.quant))}
- className={cn(
- "shrink-0 rounded-md p-1 transition-colors hover:bg-black/5 hover:text-foreground dark:hover:bg-white/10",
- pinnedKeys.includes(pinKey(repoId, v.quant))
- ? "text-foreground/80"
- : "text-muted-foreground/60",
- )}
- >
-
-
-
-
- {pinnedKeys.includes(pinKey(repoId, v.quant))
- ? "Unpin quant"
- : "Pin quant to the top"}
-
-
- )}
- {v.downloaded && (
+ {v.downloaded && onConfigure && (
+ onConfigure(repoId, {
+ source: sourceOverride ?? (isLocalPath ? "local" : "hub"),
+ isLora: false,
+ ggufVariant: v.quant,
+ isDownloaded: true,
+ expectedBytes,
+ contextLength: nativeContext,
+ isGguf: true,
+ })
+ }
/>
)}
- {v.downloaded && onDeleteVariant && (
-
- This will remove{" "}
-
- {repoId} ({v.quant})
- {" "}
- from disk. You can re-download it later.
- >
- )
- }
- successMessage={
- getDeleteVariantSuccessMessage?.(v.quant) ??
- `Deleted ${repoId} ${v.quant}`
- }
- buttonClassName="p-1"
- iconClassName="size-3"
- disabled={deleteDisabled}
- onConfirm={async () => {
- await onDeleteVariant(v.quant);
- // Drop the pin too: a pinned row for a deleted file
- // would try to load something that no longer exists.
- if (pinnedKeys.includes(pinKey(repoId, v.quant))) {
- togglePinnedQuant(repoId, v.quant);
+ {v.downloaded &&
+ (allowPin ||
+ (v.update_available && onUpdateVariant) ||
+ onDeleteVariant ||
+ !isLocalPath) && (
+
- )}
+ pin={
+ allowPin
+ ? {
+ pinned: pinnedKeys.includes(pinKey(repoId, v.quant)),
+ pinLabel: "Pin to top",
+ unpinLabel: "Unpin",
+ onToggle: () => togglePinnedQuant(repoId, v.quant),
+ }
+ : undefined
+ }
+ update={
+ v.update_available && onUpdateVariant
+ ? {
+ title: updateVariantTitle,
+ description: renderUpdateVariantDescription?.(
+ v.quant,
+ ) ?? (
+ <>
+ This will update{" "}
+
+ {repoId} ({v.quant})
+
+ {"."}
+ >
+ ),
+ repoId,
+ variant: v.quant,
+ disabled: updateDisabled,
+ onConfirm: () =>
+ onUpdateVariant(v.quant, expectedBytes),
+ onUpdated: () => setRefreshKey((key) => key + 1),
+ }
+ : undefined
+ }
+ del={
+ onDeleteVariant
+ ? {
+ title: deleteVariantTitle,
+ description: renderDeleteVariantDescription?.(
+ v.quant,
+ ) ?? (
+ <>
+ This will remove{" "}
+
+ {repoId} ({v.quant})
+ {" "}
+ from disk. You can re-download it later.
+ >
+ ),
+ successMessage:
+ getDeleteVariantSuccessMessage?.(v.quant) ??
+ `Deleted ${repoId} ${v.quant}`,
+ disabled: deleteDisabled,
+ onConfirm: async () => {
+ await onDeleteVariant(v.quant);
+ // Drop the pin too: a pinned row for a deleted file
+ // would try to load something that no longer exists.
+ if (pinnedKeys.includes(pinKey(repoId, v.quant))) {
+ togglePinnedQuant(repoId, v.quant);
+ }
+ // Re-fetch this expander's variants so the deleted
+ // quant stops showing as downloaded (and clickable to
+ // reload) while the repo still has other cached quants.
+ setRefreshKey((key) => key + 1);
+ },
+ }
+ : undefined
+ }
+ />
+ )}
);
})}
@@ -1170,17 +1171,6 @@ let _lmStudioCache: LocalModelInfo[] = [];
let _localDirCache: LocalModelInfo[] = [];
let _customFolderCache: LocalModelInfo[] = [];
let _scanFoldersCache: ScanFolderInfo[] = [];
-let _onDeviceCachesReady = false;
-let _cachedGgufRequestVersion = 0;
-let _cachedModelsRequestVersion = 0;
-let _localModelsRequestVersion = 0;
-const _onDeviceCacheListeners = new Set<(settled?: boolean) => void>();
-
-const ON_DEVICE_CACHE_TIMEOUT_MS = 30_000;
-
-function notifyOnDeviceCachesChanged(settled = false): void {
- for (const listener of _onDeviceCacheListeners) listener(settled);
-}
/** True when any on-device model (downloaded GGUF, cached repo, LM Studio, or
* custom-folder model) is known. Reads the module caches, which persist across
@@ -1333,6 +1323,19 @@ function localPathTooltip(name: string, path: string): ReactNode {
);
}
+function localModelMeta(isGguf = false): ModelSelectorChangeMeta {
+ return {
+ source: "local",
+ isLora: false,
+ isDownloaded: true,
+ ...(isGguf ? { isGguf: true } : {}),
+ };
+}
+
+function localDirectGgufMeta(): ModelSelectorChangeMeta {
+ return localModelMeta(true);
+}
+
/** Hugging Face address for an online/Hub row, or undefined when the repo id is
* missing so the row shows no (empty) address line on hover. */
function hubRepoUrl(id: string | null | undefined): string | undefined {
@@ -1343,9 +1346,7 @@ function hubRepoUrl(id: string | null | undefined): string | undefined {
/** Whether a local model is an MLX build (name hint). MLX runs on Mac only, so
* callers gate visibility on the host being a Mac. */
function localModelIsMlx(m: LocalModelInfo): boolean {
- return (
- isMlxId(m.id) || isMlxId(m.display_name) || isMlxId(m.model_id ?? "")
- );
+ return isMlxId(m.id) || isMlxId(m.display_name) || isMlxId(m.model_id ?? "");
}
/** Whether a local model matches the format toggle (GGUF detected by name/path). */
@@ -1369,6 +1370,7 @@ export function HubModelPicker({
onFoldersChange,
onBrowseHub,
onModelsChange,
+ onConfigure,
deleteDisabled = false,
section = "downloaded",
sectionToggle,
@@ -1385,12 +1387,12 @@ export function HubModelPicker({
/** Open the full Hub page to browse more models. */
onBrowseHub?: () => void;
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
+ onConfigure?: (id: string, meta: ModelSelectorChangeMeta) => void;
deleteDisabled?: boolean;
/** Section shown when not searching. Search spans all sections. */
section?: "downloaded" | "recommended" | "custom" | "connected";
/** Section toggle rendered under the search bar. */
sectionToggle?: ReactNode;
- /** Eject the loaded model. Rendered as the last list row when set. */
onEject?: () => void;
}) {
const gpu = useGpuInfo();
@@ -1434,9 +1436,8 @@ export function HubModelPicker({
pinUnslothFirst: true,
keepUnsupportedTags: true,
accessToken,
- // Only the Recommended section renders Hub results (On Device / Connected
- // use local data), so keep the Hub hooks idle on the other tabs to avoid
- // needless requests/spinner and to preserve offline-local behavior.
+ // Only Recommended renders Hub results, so keep the Hub hooks idle on other
+ // tabs to avoid needless requests and preserve offline-local behavior.
enabled: online && section === "recommended",
});
const recommendedSearch = useHubModelSearch("", {
@@ -1449,10 +1450,9 @@ export function HubModelPicker({
enabled: online && section === "recommended",
});
- // Lowercased repo ids confirmed GGUF by the store or HF search.
- // Absence means "no hint" -> hasGgufSuffix is the fallback (don't
- // conflate unknown with known-not-GGUF). Lowercased so store and HF
- // IDs differing only by casing match the same hint.
+ // Lowercased repo ids confirmed GGUF by the store or HF search. Absence means
+ // "no hint" -> hasGgufSuffix is the fallback (don't conflate unknown with
+ // known-not-GGUF). Lowercased so store and HF IDs match regardless of casing.
const modelGgufIds = useMemo(() => {
const ids = new Set
();
for (const model of models) {
@@ -1494,12 +1494,14 @@ export function HubModelPicker({
const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly);
// Repos the user clicked to collapse while expand-by-default is on. Kept in
// memory only, so it resets on reload (and when the setting is toggled).
- const [collapsedGguf, setCollapsedGguf] = useState>(
- () => new Set(),
- );
- useEffect(() => {
- setCollapsedGguf(new Set());
- }, [expandQuantizations]);
+ const [collapsedGgufState, setCollapsedGgufState] = useState<{
+ expandQuantizations: boolean;
+ value: Set;
+ }>(() => ({ expandQuantizations, value: new Set() }));
+ const collapsedGguf =
+ collapsedGgufState.expandQuantizations === expandQuantizations
+ ? collapsedGgufState.value
+ : new Set();
const isGgufExpanded = useCallback(
(id: string) =>
expandQuantizations ? !collapsedGguf.has(id) : expandedGguf === id,
@@ -1510,11 +1512,15 @@ export function HubModelPicker({
const toggleGgufExpanded = useCallback(
(id: string) => {
if (expandQuantizations) {
- setCollapsedGguf((prev) => {
- const next = new Set(prev);
+ setCollapsedGgufState((prev) => {
+ const current =
+ prev.expandQuantizations === expandQuantizations
+ ? prev.value
+ : new Set();
+ const next = new Set(current);
if (next.has(id)) next.delete(id);
else next.add(id);
- return next;
+ return { expandQuantizations, value: next };
});
} else {
setExpandedGguf((prev) => (prev === id ? null : id));
@@ -1575,14 +1581,37 @@ export function HubModelPicker({
});
}, []);
- // Cached (downloaded) repos -- module-level cache avoids flashing an
- // empty "Downloaded" section when the popover re-mounts.
- const [cachedGguf, setCachedGguf] =
- useState(_cachedGgufCache);
- const [cachedModels, setCachedModels] =
- useState(_cachedModelsCache);
- const alreadyCached = _onDeviceCachesReady || hasDownloadedModels();
- const [cachedReady, setCachedReady] = useState(alreadyCached);
+ const pickerInventory = useChatPickerInventory({ enabled: true });
+ const { cachedGguf, cachedModels, cachedReady, refreshInventory } =
+ pickerInventory;
+ const lmStudioModels = useMemo(
+ () =>
+ sortLmStudio(
+ pickerInventory.localModels.filter((m) => m.source === "lmstudio"),
+ ),
+ [pickerInventory.localModels],
+ );
+ const localDirModels = useMemo(
+ () => pickerInventory.localModels.filter((m) => m.source === "models_dir"),
+ [pickerInventory.localModels],
+ );
+ const customFolderModels = useMemo(
+ () => pickerInventory.localModels.filter((m) => m.source === "custom"),
+ [pickerInventory.localModels],
+ );
+ useEffect(() => {
+ _cachedGgufCache = cachedGguf;
+ _cachedModelsCache = cachedModels;
+ _lmStudioCache = lmStudioModels;
+ _localDirCache = localDirModels;
+ _customFolderCache = customFolderModels;
+ }, [
+ cachedGguf,
+ cachedModels,
+ lmStudioModels,
+ localDirModels,
+ customFolderModels,
+ ]);
const [updateConflictKey, setUpdateConflictKey] = useState(
null,
);
@@ -1606,35 +1635,6 @@ export function HubModelPicker({
setUpdateConflictKey(null);
}, [updateConflictKey]);
- // LM Studio local models -- module-level cache, same pattern as above.
- const [lmStudioModels, setLmStudioModels] =
- useState(_lmStudioCache);
- // Models found under the local models directory (./models), so they stay
- // selectable on the On Device tab after leaving the Fine-tuned tab.
- const [localDirModels, setLocalDirModels] =
- useState(_localDirCache);
- const [customFolderModels, setCustomFolderModels] =
- useState(_customFolderCache);
-
- useEffect(() => {
- const syncModuleCaches = (settled = false) => {
- setCachedGguf(_cachedGgufCache);
- setCachedModels(_cachedModelsCache);
- setLmStudioModels(_lmStudioCache);
- setLocalDirModels(_localDirCache);
- setCustomFolderModels(_customFolderCache);
- setCachedReady(
- (ready) =>
- ready || settled || _onDeviceCachesReady || hasDownloadedModels(),
- );
- };
- _onDeviceCacheListeners.add(syncModuleCaches);
- syncModuleCaches();
- return () => {
- _onDeviceCacheListeners.delete(syncModuleCaches);
- };
- }, []);
-
// Custom scan folders management
const [scanFolders, setScanFolders] =
useState(_scanFoldersCache);
@@ -1645,94 +1645,9 @@ export function HubModelPicker({
const [showFolderBrowser, setShowFolderBrowser] = useState(false);
const [recommendedFolders, setRecommendedFolders] = useState([]);
- const applyLocalModels = useCallback(
- (res: Awaited>) => {
- const lm = sortLmStudio(
- res.models.filter((m) => m.source === "lmstudio"),
- );
- _lmStudioCache = lm;
- setLmStudioModels(lm);
- const ld = res.models.filter((m) => m.source === "models_dir");
- _localDirCache = ld;
- setLocalDirModels(ld);
- const cf = res.models.filter((m) => m.source === "custom");
- _customFolderCache = cf;
- setCustomFolderModels(cf);
- notifyOnDeviceCachesChanged();
- },
- [],
- );
-
- const refreshColdOnDeviceCaches = useCallback(() => {
- const ggufRequestVersion = ++_cachedGgufRequestVersion;
- const modelsRequestVersion = ++_cachedModelsRequestVersion;
- const localRequestVersion = ++_localModelsRequestVersion;
- let ggufResult: Awaited> | undefined;
- let modelsResult: Awaited> | undefined;
- let localResult: Awaited> | undefined;
- let released = false;
-
- const ggufRequest = listCachedGguf().then(
- (value) => { if (!released) ggufResult = value; },
- () => {},
- );
- const modelsRequest = listCachedModels(hfToken || undefined).then(
- (value) => { if (!released) modelsResult = value; },
- () => {},
- );
- const localRequest = listLocalModels().then(
- (value) => {
- localResult = value;
- if (released && localRequestVersion === _localModelsRequestVersion) {
- if (ggufResult !== undefined && modelsResult !== undefined) _onDeviceCachesReady = true;
- applyLocalModels(value);
- }
- },
- () => {},
- );
- const isCurrent = () =>
- ggufRequestVersion === _cachedGgufRequestVersion &&
- modelsRequestVersion === _cachedModelsRequestVersion &&
- localRequestVersion === _localModelsRequestVersion;
- const publish = (invalidate = false) => {
- if (!isCurrent()) return;
- if (invalidate) {
- released = true;
- ++_cachedGgufRequestVersion;
- ++_cachedModelsRequestVersion;
- }
- if (ggufResult !== undefined) {
- _cachedGgufCache = ggufResult;
- setCachedGguf(ggufResult);
- }
- if (modelsResult !== undefined) {
- _cachedModelsCache = modelsResult;
- setCachedModels(modelsResult);
- }
- if (localResult !== undefined) applyLocalModels(localResult);
- if (ggufResult !== undefined && modelsResult !== undefined && localResult !== undefined) {
- _onDeviceCachesReady = true;
- }
- notifyOnDeviceCachesChanged(true);
- };
- const timeout = window.setTimeout(() => publish(true), ON_DEVICE_CACHE_TIMEOUT_MS);
- void Promise.all([ggufRequest, modelsRequest, localRequest]).then(() => {
- window.clearTimeout(timeout);
- publish();
- });
- }, [applyLocalModels, hfToken]);
-
const refreshLocalModelsList = useCallback(() => {
- if (!_onDeviceCachesReady && !hasDownloadedModels()) return refreshColdOnDeviceCaches();
- const requestVersion = ++_localModelsRequestVersion;
- listLocalModels()
- .then((res) => {
- if (requestVersion === _localModelsRequestVersion) {
- applyLocalModels(res);
- }
- })
- .catch(() => {});
- }, [applyLocalModels, refreshColdOnDeviceCaches]);
+ void pickerInventory.refreshInventory();
+ }, [pickerInventory.refreshInventory]);
const refreshScanFolders = useCallback(() => {
listScanFolders()
@@ -1752,9 +1667,8 @@ export function HubModelPicker({
if (!trimmed || folderLoading) return;
setFolderError(null);
setFolderLoading(true);
- // From the folder browser's one-click "Use this folder": the typed-
- // input panel is closed, so the inline folderError is invisible.
- // Surface failures (denylisted path, sandbox 403, etc.) via toast.
+ // From the folder browser's "Use this folder": the typed-input panel is
+ // closed, so surface failures (denylisted path, sandbox 403) via toast.
const fromBrowser = overridePath !== undefined;
try {
const created = await addScanFolder(trimmed);
@@ -1811,46 +1725,37 @@ export function HubModelPicker({
);
const refreshCachedLists = useCallback(() => {
- if (!_onDeviceCachesReady && !hasDownloadedModels()) return refreshColdOnDeviceCaches();
- const ggufRequestVersion = ++_cachedGgufRequestVersion;
- listCachedGguf()
- .then((v) => {
- if (ggufRequestVersion !== _cachedGgufRequestVersion) return;
- _cachedGgufCache = v;
- setCachedGguf(v);
- notifyOnDeviceCachesChanged();
- })
- .catch(() => {});
- const modelsRequestVersion = ++_cachedModelsRequestVersion;
- listCachedModels(hfToken || undefined)
- .then((v) => {
- if (modelsRequestVersion !== _cachedModelsRequestVersion) return;
- _cachedModelsCache = v;
- setCachedModels(v);
- notifyOnDeviceCachesChanged();
- })
- .catch(() => {});
- refreshLocalModelsList();
- }, [hfToken, refreshColdOnDeviceCaches, refreshLocalModelsList]);
+ void pickerInventory.refreshInventory();
+ }, [pickerInventory.refreshInventory]);
// Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking
// call. The worker pulls only changed blobs, so the cached copy stays usable until done.
- const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => {
- return downloadManager
- .requestStart({
- kind: "model",
- repoId,
- variant,
- expectedBytes,
- })
- .then((outcome) => {
- if (outcome === "conflict") {
- setUpdateConflictKey(jobKeyOf("model", repoId, variant));
- } else if (outcome === "error") {
- throw new Error("Failed to start update");
- }
- });
- }, []);
+ const startManagedUpdate = useCallback(
+ (repoId: string, variant: string, expectedBytes: number) => {
+ return downloadManager
+ .requestStart({
+ kind: "model",
+ repoId,
+ variant,
+ expectedBytes,
+ })
+ .then((outcome) => {
+ if (outcome === "conflict") {
+ setUpdateConflictKey(jobKeyOf("model", repoId, variant));
+ } else if (outcome === "busy") {
+ // A sibling variant/snapshot for this repo is already downloading,
+ // so this update did not start. Say so instead of closing the
+ // dialog as if it began and leaving the cached copy stale.
+ toast.info("A download for this model is already in progress", {
+ description: "Try updating again once it finishes.",
+ });
+ } else if (outcome === "error") {
+ throw new Error("Failed to start update");
+ }
+ });
+ },
+ [],
+ );
const updateGgufVariant = useCallback(
(repoId: string, quant: string, expectedBytes: number) =>
@@ -1863,96 +1768,11 @@ export function HubModelPicker({
listRecommendedFolders()
.then(setRecommendedFolders)
.catch(() => {});
+ }, [refreshScanFolders]);
- // Publish downloaded and local rows as one bounded snapshot. Existing data
- // stays visible during background refreshes, and a failed source keeps its
- // last successful cache instead of clearing or durably marking it ready.
- const controller = new AbortController();
- const timeout = window.setTimeout(
- () => controller.abort(),
- ON_DEVICE_CACHE_TIMEOUT_MS,
- );
- const aborted = new Promise((_, reject) => {
- controller.signal.addEventListener(
- "abort",
- () => reject(controller.signal.reason),
- { once: true },
- );
- });
- const bounded = (request: Promise) =>
- Promise.race([request, aborted]);
- let cancelled = false;
- const ggufRequestVersion = ++_cachedGgufRequestVersion;
- const modelsRequestVersion = ++_cachedModelsRequestVersion;
- const localRequestVersion = ++_localModelsRequestVersion;
- const localRequest = listLocalModels();
-
- void Promise.allSettled([
- bounded(listCachedGguf(controller.signal)),
- bounded(listCachedModels(hfToken || undefined, controller.signal)),
- bounded(localRequest),
- ]).then(([ggufResult, modelsResult, localResult]) => {
- window.clearTimeout(timeout);
- if (cancelled) return;
-
- const ggufIsCurrent =
- ggufRequestVersion === _cachedGgufRequestVersion;
- const modelsAreCurrent =
- modelsRequestVersion === _cachedModelsRequestVersion;
- const localIsCurrent =
- localRequestVersion === _localModelsRequestVersion;
-
- if (ggufResult.status === "fulfilled" && ggufIsCurrent) {
- _cachedGgufCache = ggufResult.value;
- setCachedGguf(ggufResult.value);
- notifyOnDeviceCachesChanged();
- }
- if (modelsResult.status === "fulfilled" && modelsAreCurrent) {
- _cachedModelsCache = modelsResult.value;
- setCachedModels(modelsResult.value);
- notifyOnDeviceCachesChanged();
- }
- if (localResult.status === "fulfilled" && localIsCurrent) {
- applyLocalModels(localResult.value);
- }
- if (localResult.status === "rejected" && controller.signal.aborted) {
- void localRequest.then((value) => {
- if (cancelled || localRequestVersion !== _localModelsRequestVersion) return;
- if (ggufResult.status === "fulfilled" && modelsResult.status === "fulfilled") _onDeviceCachesReady = true;
- applyLocalModels(value);
- }).catch(() => {});
- }
- const snapshotIsCurrent =
- ggufIsCurrent && modelsAreCurrent && localIsCurrent;
- if (
- ggufResult.status === "fulfilled" &&
- modelsResult.status === "fulfilled" &&
- localResult.status === "fulfilled" &&
- snapshotIsCurrent
- ) {
- _onDeviceCachesReady = true;
- }
- notifyOnDeviceCachesChanged(snapshotIsCurrent);
- });
-
- return () => {
- cancelled = true;
- window.clearTimeout(timeout);
- controller.abort();
- queueMicrotask(() => {
- if (
- ggufRequestVersion === _cachedGgufRequestVersion &&
- modelsRequestVersion === _cachedModelsRequestVersion &&
- localRequestVersion === _localModelsRequestVersion
- ) {
- ++_cachedGgufRequestVersion;
- ++_cachedModelsRequestVersion;
- ++_localModelsRequestVersion;
- notifyOnDeviceCachesChanged(true);
- }
- });
- };
- }, [applyLocalModels, hfToken, refreshScanFolders]);
+ useEffect(() => {
+ void refreshInventory();
+ }, [refreshInventory]);
// Hide downloaded models from the recommended list. Case-insensitive
// since the HF cache lowercases repo IDs.
@@ -1989,7 +1809,8 @@ export function HubModelPicker({
// Chat-only keeps runnable formats: GGUF anywhere, plus MLX/safetensors
// on Mac (matches the empty Recommended view so search stays consistent).
.filter(
- (id) => !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac),
+ (id) =>
+ !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac),
)
.filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id));
// Sort: GGUFs first, then hub models
@@ -2259,14 +2080,12 @@ export function HubModelPicker({
const togglePinned = usePinnedModelsStore((s) => s.togglePinned);
const pinnedSet = useMemo(() => new Set(pinnedIds), [pinnedIds]);
- // Candidate pins whose repo still exists in the managed cache. Per-quant
- // validation below is required because deleting one variant can leave a
- // sibling quant (and therefore the repo row) cached.
+ // Candidate pins whose repo still exists in the cache. Per-quant validation
+ // below is needed because deleting one variant can leave a sibling cached.
const pinnedQuantCandidates = useMemo(() => {
- // The existence check ignores the text query (but keeps the format filter)
- // so a pinned quant stays findable by its quant name even when the repo id
- // does not match the query; querying visibleCachedGguf here would drop the
- // repo before the later `${repoId} ${quant}` predicate could surface it.
+ // The existence check ignores the text query (keeps the format filter) so a
+ // pinned quant stays findable by quant name; querying visibleCachedGguf would
+ // drop the repo before the `${repoId} ${quant}` predicate could surface it.
const cached = new Set(
sortedCachedGguf
.filter((c) => matchesFormatFilter(c.repo_id, true, formatFilter))
@@ -2276,21 +2095,22 @@ export function HubModelPicker({
cached.has(entry.repoId),
);
}, [pinnedIds, sortedCachedGguf, formatFilter]);
- const pinnedQuantValidationKey = useMemo(() => {
- const cacheByRepo = new Map(
- sortedCachedGguf.map((repo) => [repo.repo_id, repo]),
- );
- return pinnedQuantCandidates
- .map((entry) => {
- const cached = cacheByRepo.get(entry.repoId);
- return `${pinKey(entry.repoId, entry.quant)}@${cached?.size_bytes ?? 0}:${cached?.last_modified ?? 0}`;
- })
- .join("\u0000");
- }, [pinnedQuantCandidates, sortedCachedGguf]);
const [pinnedQuantValidation, setPinnedQuantValidation] = useState<{
- key: string;
+ validated: boolean;
downloaded: ReadonlySet;
- }>({ key: "", downloaded: new Set() });
+ }>({ validated: false, downloaded: new Set() });
+ const prunePinnedQuantValidation = useCallback(
+ (repoId: string, quant: string) => {
+ const key = pinKey(repoId, quant);
+ setPinnedQuantValidation((prev) => {
+ if (!prev.downloaded.has(key)) return prev;
+ const downloaded = new Set(prev.downloaded);
+ downloaded.delete(key);
+ return { ...prev, downloaded };
+ });
+ },
+ [],
+ );
useEffect(() => {
let cancelled = false;
@@ -2302,12 +2122,13 @@ export function HubModelPicker({
void Promise.all(
repoIds.map(async (repoId) => {
try {
- const response = await listGgufVariants(
+ const response = await listGgufVariantsCached(
repoId,
hfToken || undefined,
+ { preferLocalCache: true },
);
- return normalizeGgufVariantsResponse(response).variants
- .filter((variant) => variant.downloaded === true)
+ return normalizeGgufVariantsResponse(response)
+ .variants.filter((variant) => variant.downloaded === true)
.map((variant) => pinKey(repoId, variant.quant));
} catch {
// If the backend cannot verify a quant, hiding the direct-load row
@@ -2318,7 +2139,7 @@ export function HubModelPicker({
).then((groups) => {
if (!cancelled) {
setPinnedQuantValidation({
- key: pinnedQuantValidationKey,
+ validated: true,
downloaded: new Set(groups.flat()),
});
}
@@ -2327,13 +2148,13 @@ export function HubModelPicker({
return () => {
cancelled = true;
};
- }, [hfToken, pinnedQuantCandidates, pinnedQuantValidationKey]);
+ }, [hfToken, pinnedQuantCandidates]);
const downloadedPinnedQuantKeys = useMemo>(
() =>
- pinnedQuantValidation.key === pinnedQuantValidationKey
+ pinnedQuantValidation.validated
? pinnedQuantValidation.downloaded
: new Set(),
- [pinnedQuantValidation, pinnedQuantValidationKey],
+ [pinnedQuantValidation],
);
// Verified downloaded quants, in pin order and filtered by repo id or quant.
@@ -2345,17 +2166,32 @@ export function HubModelPicker({
(!q ||
normalizeForSearch(`${entry.repoId} ${entry.quant}`).includes(q)),
);
- }, [
- debouncedQuery,
- downloadedPinnedQuantKeys,
- pinnedQuantCandidates,
- ]);
+ }, [debouncedQuery, downloadedPinnedQuantKeys, pinnedQuantCandidates]);
const pinnedCachedModelRows = useMemo(
- () => visibleCachedModelRows.filter((c) => pinnedSet.has(pinKey(c.repo_id))),
+ () =>
+ visibleCachedModelRows.filter((c) => pinnedSet.has(pinKey(c.repo_id))),
[visibleCachedModelRows, pinnedSet],
);
+ const pinnedRows = useMemo(() => {
+ const rank = makePinRank(pinnedIds);
+ const rows = [
+ ...pinnedQuants.map((entry) => ({
+ key: pinKey(entry.repoId, entry.quant),
+ entry,
+ model: null,
+ })),
+ ...pinnedCachedModelRows.map((model) => ({
+ key: pinKey(model.repo_id),
+ entry: null,
+ model,
+ })),
+ ];
+ rows.sort((a, b) => rank(a.key) - rank(b.key));
+ return rows;
+ }, [pinnedIds, pinnedQuants, pinnedCachedModelRows]);
+
// Split downloaded models so non-Unsloth repos get their own "Other models"
// section above Fine-tuned.
const unslothCachedGguf = useMemo(
@@ -2395,26 +2231,28 @@ export function HubModelPicker({
const filteredRecommendedIds = useMemo(() => {
if (!showHfSection) return [];
const q = normalizeForSearch(debouncedQuery.trim());
- return recommendedIds
- .filter((id) => normalizeForSearch(id).includes(q))
- .filter((id) =>
- matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter),
- )
- // Curated defaults obey the fit toggle like the live HF rows, else large
- // defaults resurface in search results with the filter on.
- .filter(
- (id) =>
- !fitOnDeviceOnly ||
- downloadedSet.has(id.toLowerCase()) ||
- hfModelFitsDevice(
- {
- id,
- totalParams: recommendedParamCountById.get(id),
- isGguf: isKnownGgufRepo(id),
- },
- gpu,
- ),
- );
+ return (
+ recommendedIds
+ .filter((id) => normalizeForSearch(id).includes(q))
+ .filter((id) =>
+ matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter),
+ )
+ // Curated defaults obey the fit toggle like the live HF rows, else large
+ // defaults resurface in search results with the filter on.
+ .filter(
+ (id) =>
+ !fitOnDeviceOnly ||
+ downloadedSet.has(id.toLowerCase()) ||
+ hfModelFitsDevice(
+ {
+ id,
+ totalParams: recommendedParamCountById.get(id),
+ isGguf: isKnownGgufRepo(id),
+ },
+ gpu,
+ ),
+ )
+ );
}, [
showHfSection,
debouncedQuery,
@@ -2435,27 +2273,30 @@ export function HubModelPicker({
const hfIds = useMemo(() => {
// Only the Unsloth tab searches the HF listing, and only Unsloth models.
if (!showHfSection || section !== "recommended") return [];
- return results
- .filter(isChatSupported)
- .filter(
- (r) =>
- !fitOnDeviceOnly ||
- downloadedSet.has(r.id.toLowerCase()) ||
- hfModelFitsDevice(r, gpu),
- )
- .map((result) => result.id)
- .filter((id) => !isHiddenModelId(id))
- .filter((id) => id.toLowerCase().startsWith("unsloth/"))
- .filter((id) => !recommendedSet.has(id))
- // Chat-only keeps runnable formats: GGUF anywhere, plus MLX/safetensors
- // on Mac (matches the empty Recommended view so search stays consistent).
- .filter(
- (id) => !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac),
- )
- .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id))
- .filter((id) =>
- matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter),
- );
+ return (
+ results
+ .filter(isChatSupported)
+ .filter(
+ (r) =>
+ !fitOnDeviceOnly ||
+ downloadedSet.has(r.id.toLowerCase()) ||
+ hfModelFitsDevice(r, gpu),
+ )
+ .map((result) => result.id)
+ .filter((id) => !isHiddenModelId(id))
+ .filter((id) => id.toLowerCase().startsWith("unsloth/"))
+ .filter((id) => !recommendedSet.has(id))
+ // Chat-only keeps runnable formats: GGUF anywhere, plus MLX/safetensors
+ // on Mac (matches the empty Recommended view so search stays consistent).
+ .filter(
+ (id) =>
+ !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac),
+ )
+ .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id))
+ .filter((id) =>
+ matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter),
+ )
+ );
}, [
recommendedSet,
results,
@@ -2479,16 +2320,13 @@ export function HubModelPicker({
section === "downloaded" &&
cachedReady &&
!pinnedCollapsed &&
- (pinnedQuants.length > 0 || pinnedCachedModelRows.length > 0)
+ pinnedRows.length > 0
) {
keys.push(
- ...pinnedQuants.map((entry) =>
- makeModelOptionKey("pinned-quant", pinKey(entry.repoId, entry.quant)),
- ),
- );
- keys.push(
- ...pinnedCachedModelRows.map((model) =>
- makeModelOptionKey("downloaded-model", model.repo_id),
+ ...pinnedRows.map((row) =>
+ row.entry
+ ? makeModelOptionKey("pinned-quant", row.key)
+ : makeModelOptionKey("downloaded-model", row.model.repo_id),
),
);
}
@@ -2543,12 +2381,12 @@ export function HubModelPicker({
}
// Fine-tuned models sit below downloaded, above custom folders.
- if (section === "downloaded" && cachedReady && !fineTunedCollapsed) {
+ if (section === "downloaded" && !fineTunedCollapsed) {
keys.push(...fineTunedRows.map((m) => makeModelOptionKey("lora", m.id)));
}
// Custom folders sit right below the downloaded models on On Device.
- if (section === "downloaded" && cachedReady && !customFoldersCollapsed) {
+ if (section === "downloaded" && !customFoldersCollapsed) {
keys.push(
...sortedCustomFolderModels.map((model) =>
makeModelOptionKey("custom-folder", model.id),
@@ -2556,7 +2394,7 @@ export function HubModelPicker({
);
}
- if (section === "downloaded" && cachedReady && !lmStudioCollapsed) {
+ if (section === "downloaded" && !lmStudioCollapsed) {
keys.push(
...sortedLmStudio.map((model) =>
makeModelOptionKey("lm-studio", model.id),
@@ -2564,7 +2402,7 @@ export function HubModelPicker({
);
}
- if (section === "downloaded" && cachedReady && !localDirCollapsed) {
+ if (section === "downloaded" && !localDirCollapsed) {
keys.push(
...sortedLocalDir.map((model) =>
makeModelOptionKey("local-dir", model.id),
@@ -2584,8 +2422,7 @@ export function HubModelPicker({
chatOnly,
sortedCustomFolderModels,
customFoldersCollapsed,
- pinnedQuants,
- pinnedCachedModelRows,
+ pinnedRows,
pinnedCollapsed,
downloadedCollapsed,
fineTunedRows,
@@ -2706,9 +2543,8 @@ export function HubModelPicker({
}, [scrollRef, updateListFades]);
// Sentinel + IntersectionObserver for recommended infinite scroll. Re-running
- // on each loaded page (results length) re-attaches the observer so a heavily
- // filtered list keeps paging until the viewport fills or the listing ends;
- // fetchMore is a no-op while a page is in flight. Callback ref tracks mount.
+ // per loaded page re-attaches the observer so a heavily filtered list keeps
+ // paging until the viewport fills; fetchMore is a no-op while a page is in flight.
const [recommendedSentinel, setRecommendedSentinel] =
useState(null);
const recommendedSentinelRef = useCallback((node: HTMLDivElement | null) => {
@@ -2757,8 +2593,8 @@ export function HubModelPicker({
const showDownloaded = section === "downloaded";
const showCustom = section === "downloaded";
const showRecommendedSection = !showHfSection && section === "recommended";
- const onDeviceCacheLoading = showDownloaded && !cachedReady;
const downloadedEmpty =
+ pinnedRows.length === 0 &&
visibleCachedGguf.length === 0 &&
visibleCachedModelRows.length === 0 &&
sortedLmStudio.length === 0 &&
@@ -2767,25 +2603,21 @@ export function HubModelPicker({
// non-empty Fine-tuned section.
fineTunedRows.length === 0;
- // Sort dropdown shown inline to the right of the section toggle. Options
- // depend on the tab and stay visible while searching so results can be
- // sorted. Fixed width matching the Search Hub button so it and the format
- // dropdown always line up; text-xs matches that button too. The trigger label
- // clips (no ellipsis) when long; the open menu expands to show it in full.
+ // Sort dropdown inline right of the section toggle; options depend on the tab
+ // and stay visible while searching. Fixed width matches the Search Hub button
+ // so it and the format dropdown line up. Trigger label clips; the menu shows full.
const sortTriggerClassName =
"w-[110px] shrink-0 justify-between pr-2.5 !border-0 text-xs [&>span]:!text-clip";
- // Tighter menu like the Projects activity Select: less left/top padding and
- // text-xs to match the trigger. Keep the option's right padding so the
- // selected-item checkmark never overlaps the label.
+ // Tighter menu (less padding, text-xs) matching the trigger. Keep the option's
+ // right padding so the selected-item checkmark never overlaps the label.
const sortMenuContentClassName =
"!p-1 !rounded-[14px] [&_[role=option]]:!pl-2 [&_[role=option]]:!py-1.5 [&_[role=option]]:!text-xs [&_[role=option]]:!rounded-[10px]";
- // Device-fit toggle lives inside the sort menu (shared with the Hub page).
- // The whole row is the click target (a button): a Checkbox renders as a
- // , and label-click forwarding to a button is unreliable, so the row
- // owns the toggle and the Checkbox is presentational (pointer-events-none).
+ // Device-fit toggle inside the sort menu (shared with the Hub page). The whole
+ // row is the button: a Checkbox renders as a and label-click forwarding
+ // to it is unreliable, so the row owns the toggle and the Checkbox is presentational.
const fitOnDeviceFooter = (
-
+
Only show models that fit
@@ -2808,6 +2640,19 @@ export function HubModelPicker({
);
+ // Sort icon + selected label inside the trigger pill.
+ const sortTriggerContent = (label: ReactNode) => (
+
+
+ {label}
+
+ );
+ // On Device / Custom rows are already on disk, so the device-fit filter
+ // only applies to the Unsloth listing.
const sectionSortDropdown =
section === "recommended" ? (
o.value === recommendedSort)
+ ?.label ?? recommendedSort,
+ )}
footer={fitOnDeviceFooter}
/>
) : section === "downloaded" ? (
@@ -2829,7 +2678,10 @@ export function HubModelPicker({
align="end"
className={sortTriggerClassName}
contentClassName={sortMenuContentClassName}
- footer={fitOnDeviceFooter}
+ triggerContent={sortTriggerContent(
+ LOCAL_SORT_OPTIONS.find((o) => o.value === downloadedSort)?.label ??
+ downloadedSort,
+ )}
/>
) : (
o.value === customSort)?.label ??
+ customSort,
+ )}
/>
);
@@ -2896,60 +2751,6 @@ export function HubModelPicker({
selected && "bg-[#ececec] dark:bg-[var(--sidebar-accent)]",
);
- // Pin toggle at a row's right edge: hidden until the row is hovered (or the
- // button is focused), always visible while pinned so pinned rows read as such.
- // `small` matches the compact quant-row action sizing; it also skips the
- // hide-until-hover classes since small pins render inside a hover-gated group.
- const renderPinAction = (
- repoId: string,
- quant?: string,
- opts?: { className?: string; small?: boolean },
- ) => {
- const pinned = pinnedSet.has(pinKey(repoId, quant));
- const target = quant ? `${repoId} ${quant}` : repoId;
- return (
-
-
- {
- e.stopPropagation();
- togglePinned(repoId, quant);
- }}
- aria-label={pinned ? `Unpin ${target}` : `Pin ${target}`}
- aria-pressed={pinned}
- className={cn(
- "shrink-0 rounded-md transition-colors hover:bg-black/5 dark:hover:bg-white/10",
- opts?.small ? "p-1" : "p-1.5",
- pinned
- ? "text-foreground/80 hover:text-foreground"
- : "text-muted-foreground/60 hover:text-foreground",
- !pinned &&
- !opts?.small &&
- "opacity-0 focus-visible:opacity-100 group-hover:opacity-100 group-focus-within:opacity-100",
- opts?.className,
- )}
- >
-
-
-
-
- {pinned
- ? quant
- ? "Unpin quant"
- : "Unpin model"
- : quant
- ? "Pin quant to the top"
- : "Pin model to the top"}
-
-
- );
- };
-
// A pinned quant: repo name with the quant as a grey chip. One click loads
// that quant directly, no expansion needed.
const renderPinnedQuantRow = (entry: { repoId: string; quant: string }) => {
@@ -2958,16 +2759,14 @@ export function HubModelPicker({
pinKey(entry.repoId, entry.quant),
);
const { owner, name } = splitRepoLabel(entry.repoId);
- const isSelected = value === entry.repoId && activeGgufVariant === entry.quant;
+ const isSelected =
+ value === entry.repoId && activeGgufVariant === entry.quant;
const isLoaded =
modelIdsMatchForPicker(loadedModelId, entry.repoId) &&
!ggufVariantsMatchForPicker(activeGgufVariant, null) &&
ggufVariantsMatchForPicker(activeGgufVariant, entry.quant);
return (
-
+
)}
-
- {renderPinAction(entry.repoId, entry.quant, { small: true })}
-
-
- This will remove{" "}
-
- {entry.repoId} ({entry.quant})
- {" "}
- from disk. You can re-download it later.
- >
- }
- successMessage={`Deleted ${entry.repoId} ${entry.quant}`}
- buttonClassName="p-1"
+
+ {onConfigure && (
+
+ onConfigure(entry.repoId, {
+ source: "hub",
+ isLora: false,
+ ggufVariant: entry.quant,
+ isDownloaded: true,
+ isGguf: true,
+ })
+ }
+ />
+ )}
+ {
- await deleteCachedModel(entry.repoId, entry.quant);
- refreshCachedLists();
- // The file is gone, so drop its pin too.
- togglePinned(entry.repoId, entry.quant);
+ cachePath={{ repoId: entry.repoId, variant: entry.quant }}
+ pin={{
+ pinned: true,
+ pinLabel: "Pin to top",
+ unpinLabel: "Unpin",
+ onToggle: () => togglePinned(entry.repoId, entry.quant),
+ }}
+ del={{
+ title: "Delete cached model?",
+ description: (
+ <>
+ This will remove{" "}
+
+ {entry.repoId} ({entry.quant})
+ {" "}
+ from disk. You can re-download it later.
+ >
+ ),
+ successMessage: `Deleted ${entry.repoId} ${entry.quant}`,
+ disabled: deleteDisabled,
+ onConfirm: async () => {
+ await deleteCachedModel(
+ entry.repoId,
+ entry.quant,
+ hfToken || undefined,
+ );
+ refreshCachedLists();
+ // The file is gone, so drop its pin too.
+ togglePinned(entry.repoId, entry.quant);
+ },
}}
/>
@@ -3081,6 +2900,7 @@ export function HubModelPicker({
allowPin={true}
onHasVision={(v) => reportVision(c.repo_id, v)}
onSelect={onSelect}
+ onConfigure={onConfigure}
hfToken={hfToken || undefined}
parentOptionKey={optionKey}
onNavigatePastStart={() => hubModelList.focusOption(optionKey)}
@@ -3090,13 +2910,12 @@ export function HubModelPicker({
variantActions={{
onUpdate: (quant, expectedBytes) =>
updateGgufVariant(c.repo_id, quant, expectedBytes),
- // Can't update the model that's live in memory under itself.
updateDisabled: loadedModelId === c.repo_id,
onDelete: async (quant) => {
- await deleteCachedModel(c.repo_id, quant);
+ await deleteCachedModel(c.repo_id, quant, hfToken || undefined);
+ prunePinnedQuantValidation(c.repo_id, quant);
refreshCachedLists();
},
- deleteDisabled,
}}
/>
)}
@@ -3109,10 +2928,7 @@ export function HubModelPicker({
const optionKey = makeModelOptionKey("downloaded-model", c.repo_id);
const isSelected = value === c.repo_id;
return (
-
+
onSelect(c.repo_id, {
source: "hub",
@@ -3142,27 +2955,48 @@ export function HubModelPicker({
className={downloadedRowButtonClassName}
/>
- {renderPinAction(c.repo_id)}
-
- This will remove{" "}
- {c.repo_id} {" "}
- from disk. You can re-download it later.
- >
- }
- successMessage={`Deleted ${c.repo_id}`}
- buttonClassName="mr-1"
- disabled={deleteDisabled}
- onConfirm={async () => {
- await deleteCachedModel(c.repo_id);
- if (pinnedSet.has(pinKey(c.repo_id))) {
- togglePinned(c.repo_id);
+ {onConfigure && (
+
+ onConfigure(c.repo_id, {
+ source: "hub",
+ isLora: false,
+ isDownloaded: true,
+ isGguf: false,
+ })
}
+ />
+ )}
+ togglePinned(c.repo_id),
+ }}
+ del={{
+ title: "Delete cached model?",
+ description: (
+ <>
+ This will remove{" "}
+ {c.repo_id} {" "}
+ from disk. You can re-download it later.
+ >
+ ),
+ successMessage: `Deleted ${c.repo_id}`,
+ disabled: deleteDisabled,
+ onConfirm: async () => {
+ await deleteCachedModel(c.repo_id, undefined, hfToken || undefined);
+ if (pinnedSet.has(pinKey(c.repo_id))) {
+ togglePinned(c.repo_id);
+ }
+ },
+ onDeleted: refreshCachedLists,
}}
- onDeleted={refreshCachedLists}
/>
);
@@ -3171,214 +3005,237 @@ export function HubModelPicker({
return (
<>
- {/* A small right inset shortens the search bar so Search Hub lands on the
+ {/* A small right inset shortens the search bar so Search Hub lands on the
last dropdown's right edge (none on the wider Connected box). */}
-
-
-
- setQuery(event.target.value)}
- placeholder={
- section === "downloaded"
- ? "Search local models"
- : "Search Unsloth models"
- }
- data-model-picker-search-input={true}
- className="field-soft h-9 border-0 pl-8 pr-8"
- />
- {isLoading && (
-
- )}
-
- {onBrowseHub ? (
-
-
-
-
- Search Hub
-
-
- Search all models
-
- ) : null}
-
-
- {/* Section tabs then the format and sort dropdowns, packed left with one
- uniform gap between every control. The box is sized so the last
- dropdown still lands on Search Hub's edge. Dropdowns hide on Connected. */}
-
- {sectionToggle}
- {showConnected ? null : (
-
-
- {sectionSortDropdown}
-
- )}
-
-
-
updateListFades(e.currentTarget)}
- className={cn(
- // List sits within the menu padding so left and right gaps match.
- // Height tracks the content up to the cap, so short lists do not
- // leave white space. scroll-py + symmetric px keep the focus ring off
- // the overflow clip edges during keyboard nav.
- "model-list-scroll max-h-[335px] overflow-y-auto scroll-py-1.5 px-0.5 mr-1",
- listScrolled && "is-scrolled",
- listMoreBelow && "is-bottom-faded",
- )}
- {...hubModelList.listboxProps}
- >
- {/* Clear space for the floating Eject pill when scrolled to the end, so
- its gap above the last row matches its gap below (applies to every
- section, including Recommended). */}
- {showConnected ? (
- connectedGroups.length === 0 ? (
-
- {externalModels.length === 0
- ? "No models from your connections. Set up in Settings then Connections."
- : "No models match your search."}
-
+
+
+ setQuery(event.target.value)}
+ placeholder={
+ section === "downloaded"
+ ? "Search local models"
+ : "Search Unsloth models"
+ }
+ data-model-picker-search-input={true}
+ className="field-soft h-9 border-0 pl-8 pr-8"
+ />
+ {isLoading && (
+
+ )}
+
+ {onBrowseHub ? (
+
+
+
+
+ Search Hub
+
+
+ Search all models
+
+ ) : null}
+
+
+ {/* Section tabs then the format and sort dropdowns, packed left with one
+ uniform gap between every control. The box is sized so the last
+ dropdown still lands on Search Hub's edge. Dropdowns hide on Connected. */}
+
+ {sectionToggle}
+ {showConnected ? null : (
+
+
+ {sectionSortDropdown}
+
+ )}
+
+
+
updateListFades(e.currentTarget)}
+ className={cn(
+ // List sits within the menu padding so gaps match; height tracks content
+ // up to the cap. scroll-py + symmetric px keep the focus ring off the
+ // overflow clip edges during keyboard nav.
+ "model-list-scroll max-h-[335px] overflow-y-auto scroll-py-1.5 px-0.5 mr-1",
+ listScrolled && "is-scrolled",
+ listMoreBelow && "is-bottom-faded",
+ )}
+ {...hubModelList.listboxProps}
+ >
+
+ {showConnected ? (
+ connectedGroups.length === 0 ? (
+
+ {externalModels.length === 0
+ ? "No models from your connections. Set up in Settings then Connections."
+ : "No models match your search."}
+
+ ) : (
+ connectedGroups.map((group) => (
+
+
+
+
+ {group.providerName}
+
+
+ {group.models.map((model) => (
+
+ onSelect(model.id, {
+ source: "external",
+ isLora: false,
+ })
+ }
+ className={cn(
+ "flex w-full items-center rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-[#ececec] dark:hover:bg-[var(--sidebar-accent)]",
+ value === model.id &&
+ "bg-[#ececec] dark:bg-[var(--sidebar-accent)]",
+ )}
+ >
+ {model.name}
+
+ ))}
+
+ ))
+ )
) : (
- connectedGroups.map((group) => (
-
-
-
-
- {group.providerName}
+ <>
+ {/* First-load spinner only when nothing cached is shown yet. */}
+ {showDownloaded &&
+ !cachedReady &&
+ !showHfSection &&
+ downloadedEmpty ? (
+
+
+
+ Loading models…
- {group.models.map((model) => (
-
- onSelect(model.id, {
- source: "external",
- isLora: false,
- })
- }
- className={cn(
- "flex w-full items-center rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-[#ececec] dark:hover:bg-[var(--sidebar-accent)]",
- value === model.id &&
- "bg-[#ececec] dark:bg-[var(--sidebar-accent)]",
- )}
- >
- {model.name}
-
- ))}
-
- ))
- )
- ) : (
- <>
- {/* First-load spinner while downloaded/local scans are resolving. */}
- {onDeviceCacheLoading ? (
-
-
-
- Loading models…
-
-
- ) : null}
+ ) : null}
- {/* Empty On Device: a search miss vs nothing downloaded yet. Hidden
+ {/* Empty On Device: a search miss vs nothing downloaded yet. Hidden
when custom folders below still have matches. */}
- {showDownloaded &&
- cachedReady &&
- downloadedEmpty &&
- sortedCustomFolderModels.length === 0 ? (
-
- {showHfSection
- ? "No matching models on device."
- : formatFilter === "all"
- ? "No downloaded models yet. Search above or pick Recommended."
- : `No downloaded ${FORMAT_FILTER_LABELS[formatFilter]} models yet.`}
-
- ) : null}
+ {showDownloaded &&
+ cachedReady &&
+ downloadedEmpty &&
+ sortedCustomFolderModels.length === 0 ? (
+
+ {showHfSection
+ ? "No matching models on device."
+ : formatFilter === "all"
+ ? "No downloaded models yet. Search above or pick Recommended."
+ : `No downloaded ${FORMAT_FILTER_LABELS[formatFilter]} models yet.`}
+
+ ) : null}
- {/* Pinned quants and models sit above the Unsloth heading so
+ {/* Pinned quants and models sit above the Unsloth heading so
favorites are always first. Filtered by the query like the
sections below. */}
- {showDownloaded &&
- (pinnedQuants.length > 0 ||
- pinnedCachedModelRows.length > 0) ? (
- <>
-
}
- collapsed={pinnedCollapsed}
- onToggle={() => setPinnedCollapsed((v) => !v)}
- >
- Pinned
-
- {!pinnedCollapsed && pinnedQuants.map(renderPinnedQuantRow)}
- {!pinnedCollapsed &&
- pinnedCachedModelRows.map(renderDownloadedModelRow)}
- >
- ) : null}
+ {showDownloaded && pinnedRows.length > 0 ? (
+ <>
+
+ }
+ collapsed={pinnedCollapsed}
+ onToggle={() => setPinnedCollapsed((v) => !v)}
+ >
+ Pinned
+
+ {!pinnedCollapsed &&
+ pinnedRows.map((row) =>
+ row.entry
+ ? renderPinnedQuantRow(row.entry)
+ : renderDownloadedModelRow(row.model),
+ )}
+ >
+ ) : null}
- {/* Downloaded (Unsloth) stays visible (filtered) while searching. */}
- {showDownloaded &&
- cachedReady &&
- (unslothCachedGguf.length > 0 ||
- unslothCachedModelRows.length > 0) ? (
- <>
-
0 ||
- pinnedCachedModelRows.length > 0
- }
- collapsed={downloadedCollapsed}
- onToggle={() => setDownloadedCollapsed((v) => !v)}
- action={
- <>
- {hasOtherModels ? (
+ {/* Downloaded (Unsloth) stays visible (filtered) while searching. */}
+ {showDownloaded &&
+ (unslothCachedGguf.length > 0 ||
+ unslothCachedModelRows.length > 0) ? (
+ <>
+ 0}
+ collapsed={downloadedCollapsed}
+ onToggle={() => setDownloadedCollapsed((v) => !v)}
+ action={
+ <>
+ {hasOtherModels ? (
+
+
+
+
+
+
+
+ Other non-Unsloth models
+
+
+ ) : null}
@@ -3387,845 +3244,809 @@ export function HubModelPicker({
side="bottom"
className="tooltip-compact"
>
- Other non-Unsloth models
+ Go to fine-tuned models
- ) : null}
-
-
-
+
+
+
+
+
+
-
-
-
-
- Go to fine-tuned models
-
-
-
-
-
-
-
-
-
- Go to custom folders
-
-
- >
- }
- >
- {/* When other providers (LM Studio/Ollama) also show here, name
- this group "Unsloth" so the two are easy to tell apart. */}
- {sortedLmStudio.length > 0 ? "Unsloth" : "Downloaded"}
-
- {!downloadedCollapsed &&
- unslothCachedGguf.map(renderDownloadedGgufRow)}
- {!downloadedCollapsed &&
- unslothCachedModelRows.map(renderDownloadedModelRow)}
- >
- ) : null}
-
- {/* Other models: non-Unsloth downloads, grouped just above
- Fine-tuned. Shown only when such models exist. */}
- {showDownloaded && cachedReady && hasOtherModels ? (
-
-
- }
- collapsed={otherModelsCollapsed}
- onToggle={() => setOtherModelsCollapsed((v) => !v)}
- >
- Other models
-
- {!otherModelsCollapsed &&
- otherCachedGguf.map(renderDownloadedGgufRow)}
- {!otherModelsCollapsed &&
- otherCachedModelRows.map(renderDownloadedModelRow)}
-
- ) : null}
-
- {/* Fine-tuned models: shown after the On Device scans resolve so
- downloaded sections do not reorder during startup. */}
- {section === "downloaded" && cachedReady ? (
- <>
-
-
-
- Fine-tuned
-
-
- setFineTunedCollapsed((v) => !v)}
- className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
- >
- {fineTunedCollapsed ? (
-
- ) : (
-
- )}
-
-
-
- {!fineTunedCollapsed && fineTunedRows.length > 0 && (
-
- )}
- >
- ) : null}
-
- {showCustom && cachedReady ? (
- <>
-
-
setShowFolderBrowser(true)}
- title="Browse folders on the server"
- className="flex items-center gap-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground transition-colors hover:text-foreground"
+ Go to custom folders
+
+
+ >
+ }
>
-
- Custom Folders
-
-
- {
- setShowFolderInput((open) => {
- if (open) {
- setFolderInput("");
- setFolderError(null);
- }
- return !open;
- });
- }}
- className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
- >
-
-
- setShowFolderBrowser(true)}
- className="shrink-0 rounded p-0.5 text-muted-foreground/60 transition-colors hover:text-foreground"
- >
-
-
-
-
- setCustomFoldersCollapsed((v) => !v)}
- className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
- >
- {customFoldersCollapsed ? (
-
- ) : (
-
- )}
-
-
-
+ {/* When other providers (LM Studio/Ollama) also show here, name
+ this group "Unsloth" so the two are easy to tell apart. */}
+ {sortedLmStudio.length > 0 ? "Unsloth" : "Downloaded"}
+
+ {!downloadedCollapsed &&
+ unslothCachedGguf.map(renderDownloadedGgufRow)}
+ {!downloadedCollapsed &&
+ unslothCachedModelRows.map(renderDownloadedModelRow)}
+ >
+ ) : null}
- {/* Folder paths */}
- {!customFoldersCollapsed &&
- scanFolders.map((f) => (
-
+
+ }
+ collapsed={otherModelsCollapsed}
+ onToggle={() => setOtherModelsCollapsed((v) => !v)}
+ >
+ Other models
+
+ {!otherModelsCollapsed &&
+ otherCachedGguf.map(renderDownloadedGgufRow)}
+ {!otherModelsCollapsed &&
+ otherCachedModelRows.map(renderDownloadedModelRow)}
+
+ ) : null}
+
+ {/* Fine-tuned models: a section above Custom Folders. Always shown on
+ On Device so the train shortcut always has a target, with an empty
+ state when none exist. */}
+ {section === "downloaded" ? (
+ <>
+
+
+
+ Fine-tuned
+
+
+ setFineTunedCollapsed((v) => !v)}
+ className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
+ >
+ {fineTunedCollapsed ? (
+
+ ) : (
+
+ )}
+
+
+
+ {!fineTunedCollapsed && fineTunedRows.length > 0 && (
+
+ )}
+ >
+ ) : null}
+
+ {showCustom ? (
+ <>
+
+
setShowFolderBrowser(true)}
+ title="Browse folders on the server"
+ className="flex items-center gap-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground transition-colors hover:text-foreground"
>
-
- {f.path}
-
+ Custom Folders
+
+
handleRemoveFolder(f.id)}
- aria-label={`Remove folder ${f.path}`}
- className="shrink-0 rounded p-1 text-foreground/70 transition-colors hover:bg-destructive/10 hover:text-destructive focus-visible:bg-destructive/10 focus-visible:text-destructive"
+ aria-label={
+ showFolderInput
+ ? "Cancel adding folder"
+ : "Add scan folder by path"
+ }
+ title={
+ showFolderInput ? "Cancel" : "Add by typing a path"
+ }
+ onClick={() => {
+ setShowFolderInput((open) => {
+ if (open) {
+ setFolderInput("");
+ setFolderError(null);
+ }
+ return !open;
+ });
+ }}
+ className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
>
-
- ))}
-
- {/* Recommended folders */}
- {!customFoldersCollapsed &&
- (() => {
- const registered = new Set(
- scanFolders.map((f) => f.path),
- );
- const unregistered = recommendedFolders.filter(
- (p) => !registered.has(p),
- );
- if (unregistered.length === 0) return null;
- return (
-
- {unregistered.map((p) => (
- void handleAddFolder(p)}
- disabled={folderLoading}
- title={`Add ${p}`}
- className="rounded-full border border-dashed border-border/50 px-2 py-0.5 font-mono text-[10px] text-muted-foreground/70 transition-colors hover:border-foreground/30 hover:bg-accent hover:text-foreground disabled:opacity-40"
- >
-
- +
- {" "}
- {p.length > 30 ? `...${p.slice(-27)}` : p}
-
- ))}
-
- );
- })()}
-
- {/* Add folder input */}
- {!customFoldersCollapsed && showFolderInput && (
-
-
-
- {
- setFolderInput(e.target.value);
- setFolderError(null);
- }}
- onKeyDown={(e) => {
- if (e.key === "Enter") {
- e.preventDefault();
- handleAddFolder();
- }
- if (e.key === "Escape") {
- e.preventDefault();
- e.stopPropagation();
- setShowFolderInput(false);
- setFolderInput("");
- setFolderError(null);
- }
- }}
- placeholder="/path/to/models"
- className="h-6 min-w-0 flex-1 rounded border border-border/50 bg-transparent px-1.5 font-mono text-[10px] text-foreground outline-none placeholder:text-muted-foreground/40 focus:border-foreground/20"
- disabled={folderLoading}
- autoFocus={true}
- />
setShowFolderBrowser(true)}
- disabled={folderLoading}
- aria-label="Browse for folder"
+ aria-label="Browse for a folder on the server"
title="Browse folders on the server"
- className="flex h-6 shrink-0 items-center justify-center rounded border border-border/50 px-1.5 text-muted-foreground transition-colors hover:bg-accent hover:text-foreground disabled:opacity-40"
+ onClick={() => setShowFolderBrowser(true)}
+ className="shrink-0 rounded p-0.5 text-muted-foreground/60 transition-colors hover:text-foreground"
>
+
+
{
- void handleAddFolder();
- }}
- disabled={folderLoading || !folderInput.trim()}
- className="h-6 shrink-0 rounded border border-border/50 px-1.5 text-[10px] text-muted-foreground transition-colors hover:bg-accent disabled:opacity-40"
+ aria-label={
+ customFoldersCollapsed
+ ? "Expand custom folders"
+ : "Collapse custom folders"
+ }
+ title={customFoldersCollapsed ? "Expand" : "Collapse"}
+ onClick={() => setCustomFoldersCollapsed((v) => !v)}
+ className="shrink-0 rounded p-1 text-muted-foreground/60 transition-colors hover:text-foreground"
>
- Add
+ {customFoldersCollapsed ? (
+
+ ) : (
+
+ )}
- {folderError && (
-
- {folderError}
-
- )}
- )}
-
{
- setFolderInput(picked);
- setFolderError(null);
- // Pass the path explicitly: `folderInput` state hasn't
- // flushed yet when "Use this folder" submits.
- void handleAddFolder(picked);
- }}
- />
-
- {/* Models from custom folders */}
- {!customFoldersCollapsed &&
- sortedCustomFolderModels.map((m) => {
- const isGgufFile = m.path.toLowerCase().endsWith(".gguf");
- // Honor the backend model_format hint (suffixless GGUF
- // folders) in addition to name/path so the row classifies
- // and loads through the same GGUF path as the filter.
- const isGguf = localModelIsGguf(m);
- // Single .gguf files (e.g. Ollama blobs) load directly;
- // GGUF repos/directories expand to pick a variant.
- const isDirectGguf = isGgufFile;
- const optionKey = makeModelOptionKey(
- "custom-folder",
- m.id,
- );
- return (
-
-
{
- if (isDirectGguf) {
- onSelect(m.id, {
- source: "local",
- isLora: false,
- isDownloaded: true,
- // Mark GGUF so "Load on selection = off" stages
- // through Run settings (matches LM Studio path).
- isGguf: true,
- });
- } else if (isGguf) {
- toggleGgufExpanded(m.id);
- } else {
- onSelect(m.id, {
- source: "local",
- isLora: false,
- isDownloaded: true,
- });
- }
- }}
- onArrowDownIntoChildren={
- isGguf && !isDirectGguf && isGgufExpanded(m.id)
- ? () => {
- const focused =
- focusFirstChildOption(optionKey);
- return focused;
- }
- : undefined
- }
- vramStatus={null}
+ {/* Folder paths */}
+ {!customFoldersCollapsed &&
+ scanFolders.map((f) => (
+
+
- {isGguf && !isDirectGguf && isGgufExpanded(m.id) && (
-
- hubModelList.focusOption(optionKey)
- }
- onNavigatePastEnd={() =>
- hubModelList.moveFocus(optionKey, "next")
- }
- gpuGb={
- gpu.available ? gpu.memoryTotalGb : undefined
- }
- systemRamGb={gpu.systemRamAvailableGb || undefined}
+
+ {f.path}
+
+ handleRemoveFolder(f.id)}
+ aria-label={`Remove folder ${f.path}`}
+ className="shrink-0 rounded p-1 text-foreground/70 transition-colors hover:bg-destructive/10 hover:text-destructive focus-visible:bg-destructive/10 focus-visible:text-destructive"
+ >
+
- )}
+
- );
- })}
- {!customFoldersCollapsed &&
- showHfSection &&
- sortedCustomFolderModels.length === 0 ? (
-
- No matching models in custom folders.
-
- ) : null}
- >
- ) : null}
+ ))}
- {section === "downloaded" &&
- cachedReady &&
- sortedLmStudio.length > 0 ? (
- <>
- setLmStudioCollapsed((v) => !v)}
- >
- LM Studio
-
- {!lmStudioCollapsed &&
- sortedLmStudio.map((m) => {
- const isGgufFile = m.path.toLowerCase().endsWith(".gguf");
- // LM Studio dirs are GGUF but rarely carry a -GGUF suffix;
- // use the shared helper (model_format hint) so the row,
- // filter, and load path agree.
- const isGguf = localModelIsGguf(m);
- const optionKey = makeModelOptionKey("lm-studio", m.id);
- return (
-
-
{
- if (isGgufFile) {
- onSelect(m.id, {
- source: "local",
- isLora: false,
- isDownloaded: true,
- isGguf: true,
- });
- } else if (isGguf) {
- toggleGgufExpanded(m.id);
- } else {
- onSelect(m.id, {
- source: "local",
- isLora: false,
- isDownloaded: true,
- });
+ {/* Recommended folders */}
+ {!customFoldersCollapsed &&
+ (() => {
+ const registered = new Set(
+ scanFolders.map((f) => f.path),
+ );
+ const unregistered = recommendedFolders.filter(
+ (p) => !registered.has(p),
+ );
+ if (unregistered.length === 0) return null;
+ return (
+
+ {unregistered.map((p) => (
+ void handleAddFolder(p)}
+ disabled={folderLoading}
+ title={`Add ${p}`}
+ className="rounded-full border border-dashed border-border/50 px-2 py-0.5 font-mono text-[10px] text-muted-foreground/70 transition-colors hover:border-foreground/30 hover:bg-accent hover:text-foreground disabled:opacity-40"
+ >
+
+ +
+ {" "}
+ {p.length > 30 ? `...${p.slice(-27)}` : p}
+
+ ))}
+
+ );
+ })()}
+
+ {/* Add folder input */}
+ {!customFoldersCollapsed && showFolderInput && (
+
+
+
+ {
+ setFolderInput(e.target.value);
+ setFolderError(null);
+ }}
+ onKeyDown={(e) => {
+ if (e.key === "Enter") {
+ e.preventDefault();
+ handleAddFolder();
+ }
+ if (e.key === "Escape") {
+ e.preventDefault();
+ e.stopPropagation();
+ setShowFolderInput(false);
+ setFolderInput("");
+ setFolderError(null);
}
}}
- onArrowDownIntoChildren={
- isGguf && !isGgufFile && isGgufExpanded(m.id)
- ? () => {
- const focused =
- focusFirstChildOption(optionKey);
- return focused;
- }
- : undefined
- }
- vramStatus={null}
+ placeholder="/path/to/models"
+ className="h-6 min-w-0 flex-1 rounded border border-border/50 bg-transparent px-1.5 font-mono text-[10px] text-foreground outline-none placeholder:text-muted-foreground/40 focus:border-foreground/20"
+ disabled={folderLoading}
+ autoFocus={true}
/>
- {isGguf && !isGgufFile && isGgufExpanded(m.id) && (
-
- hubModelList.focusOption(optionKey)
- }
- onNavigatePastEnd={() =>
- hubModelList.moveFocus(optionKey, "next")
- }
- gpuGb={
- gpu.available ? gpu.memoryTotalGb : undefined
- }
- systemRamGb={gpu.systemRamAvailableGb || undefined}
+ setShowFolderBrowser(true)}
+ disabled={folderLoading}
+ aria-label="Browse for folder"
+ title="Browse folders on the server"
+ className="flex h-6 shrink-0 items-center justify-center rounded border border-border/50 px-1.5 text-muted-foreground transition-colors hover:bg-accent hover:text-foreground disabled:opacity-40"
+ >
+
- )}
-
- );
- })}
- >
- ) : null}
-
- {section === "downloaded" &&
- cachedReady &&
- sortedLocalDir.length > 0 ? (
- <>
-
setLocalDirCollapsed((v) => !v)}
- >
- Local models
-
- {!localDirCollapsed &&
- sortedLocalDir.map((m) => {
- // A loose ./models/*.gguf file loads directly; a GGUF repo
- // directory expands to pick a variant. The backend's local
- // variant scanner returns nothing for a config-less loose
- // file, so expanding it would dead-end at "No GGUF variants".
- const isGgufFile = m.path.toLowerCase().endsWith(".gguf");
- const isGguf = localModelIsGguf(m);
- const optionKey = makeModelOptionKey("local-dir", m.id);
- return (
-
-
+ {
- if (isGgufFile) {
- onSelect(m.id, {
- source: "local",
- isLora: false,
- isDownloaded: true,
- isGguf: true,
- });
- } else if (isGguf) {
- toggleGgufExpanded(m.id);
- } else {
- onSelect(m.id, {
- source: "local",
- isLora: false,
- isDownloaded: true,
- });
- }
+ void handleAddFolder();
}}
- onArrowDownIntoChildren={
- isGguf && !isGgufFile && isGgufExpanded(m.id)
- ? () => focusFirstChildOption(optionKey)
- : undefined
- }
- vramStatus={null}
- />
- {isGguf && !isGgufFile && isGgufExpanded(m.id) && (
-
- hubModelList.focusOption(optionKey)
- }
- onNavigatePastEnd={() =>
- hubModelList.moveFocus(optionKey, "next")
- }
- gpuGb={
- gpu.available ? gpu.memoryTotalGb : undefined
- }
- systemRamGb={gpu.systemRamAvailableGb || undefined}
- />
- )}
+ disabled={folderLoading || !folderInput.trim()}
+ className="h-6 shrink-0 rounded border border-border/50 px-1.5 text-[10px] text-muted-foreground transition-colors hover:bg-accent disabled:opacity-40"
+ >
+ Add
+
- );
- })}
- >
- ) : null}
-
- {showRecommendedSection ? (
- <>
- {recommendedSearch.isLoading &&
- recommendedRows.length === 0 ? (
-
-
-
- Loading models…
-
-
- ) : recommendedRows.length === 0 ? (
-
- No models found.
-
- ) : (
- recommendedRows.map((r) => {
- const id = r.id;
- const info = recommendedMeta.get(id);
- const isG = isKnownGgufRepo(id);
- const optionKey = makeModelOptionKey("recommended", id);
- return (
-
- {
- if (isG) {
- setExpandedGguf((prev) =>
- prev === id ? null : id,
- );
- } else {
- handleModelClick(id);
- }
- }}
- vramStatus={info?.status ?? null}
- vramEst={info?.est}
- gpuGb={
- gpu.available ? gpu.memoryTotalGb : undefined
- }
- onArrowDownIntoChildren={
- expandedGguf === id
- ? () => focusFirstChildOption(optionKey)
- : undefined
- }
- />
- {expandedGguf === id && (
-
- hubModelList.focusOption(optionKey)
- }
- onNavigatePastEnd={() =>
- hubModelList.moveFocus(optionKey, "next")
- }
- gpuGb={
- gpu.available ? gpu.memoryTotalGb : undefined
- }
- systemRamGb={gpu.systemRamAvailableGb || undefined}
- variantActions={{
- onDelete: async (quant) => {
- await deleteCachedModel(id, quant);
- refreshCachedLists();
- },
- deleteDisabled,
- }}
- />
- )}
-
- );
- })
- )}
- {recommendedSearch.hasMore && (
- <>
-
-
-
-
- >
- )}
- >
- ) : null}
-
- {showHfSection &&
- section === "recommended" &&
- filteredRecommendedIds.length > 0 ? (
- <>
- {filteredRecommendedIds.map((id) => {
- const vram = recommendedVramMap.get(id);
- const optionKey = makeModelOptionKey(
- "search-recommended",
- id,
- );
- return (
-
-
{
- if (isKnownGgufRepo(id)) {
- setExpandedGguf((prev) =>
- prev === id ? null : id,
- );
- } else {
- handleModelClick(id);
- }
- }}
- vramStatus={
- isKnownGgufRepo(id) ? null : (vram?.status ?? null)
- }
- vramEst={isKnownGgufRepo(id) ? undefined : vram?.est}
- gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
- onArrowDownIntoChildren={
- expandedGguf === id
- ? () => {
- const focused =
- focusFirstChildOption(optionKey);
- return focused;
- }
- : undefined
- }
- />
- {expandedGguf === id && (
-
- hubModelList.focusOption(optionKey)
- }
- onNavigatePastEnd={() =>
- hubModelList.moveFocus(optionKey, "next")
- }
- gpuGb={
- gpu.available ? gpu.memoryTotalGb : undefined
- }
- systemRamGb={gpu.systemRamAvailableGb || undefined}
- variantActions={{
- onDelete: async (quant) => {
- await deleteCachedModel(id, quant);
- refreshCachedLists();
- },
- deleteDisabled,
- }}
- />
+ {folderError && (
+
+ {folderError}
+
)}
- );
- })}
- >
- ) : null}
+ )}
- {showHfSection && section === "recommended" ? (
- <>
- {hfIds.length === 0 && !isLoading ? (
- filteredRecommendedIds.length === 0 ? (
+
{
+ setFolderInput(picked);
+ setFolderError(null);
+ // Pass the path explicitly: `folderInput` state hasn't
+ // flushed yet when "Use this folder" submits.
+ void handleAddFolder(picked);
+ }}
+ />
+
+ {/* Models from custom folders */}
+ {!customFoldersCollapsed &&
+ sortedCustomFolderModels.map((m) => {
+ const isGgufFile = m.path
+ .toLowerCase()
+ .endsWith(".gguf");
+ // Honor the backend model_format hint (suffixless GGUF
+ // folders) in addition to name/path so the row classifies
+ // and loads through the same GGUF path as the filter.
+ const isGguf = localModelIsGguf(m);
+ // Single .gguf files (e.g. Ollama blobs) load directly;
+ // GGUF repos/directories expand to pick a variant.
+ const isDirectGguf = isGgufFile;
+ const optionKey = makeModelOptionKey(
+ "custom-folder",
+ m.id,
+ );
+ return (
+
+
+
+ {
+ if (isDirectGguf) {
+ onSelect(m.id, localDirectGgufMeta());
+ } else if (isGguf) {
+ toggleGgufExpanded(m.id);
+ } else {
+ onSelect(m.id, localModelMeta());
+ }
+ }}
+ onArrowDownIntoChildren={
+ isGguf &&
+ !isDirectGguf &&
+ isGgufExpanded(m.id)
+ ? () => {
+ const focused =
+ focusFirstChildOption(optionKey);
+ return focused;
+ }
+ : undefined
+ }
+ vramStatus={null}
+ />
+
+ {isDirectGguf && onConfigure && (
+
+ onConfigure(m.id, localDirectGgufMeta())
+ }
+ />
+ )}
+ {!isGguf && onConfigure && (
+
+ onConfigure(m.id, localModelMeta())
+ }
+ />
+ )}
+
+ {isGguf &&
+ !isDirectGguf &&
+ isGgufExpanded(m.id) && (
+
+ hubModelList.focusOption(optionKey)
+ }
+ onNavigatePastEnd={() =>
+ hubModelList.moveFocus(optionKey, "next")
+ }
+ gpuGb={
+ gpu.available
+ ? gpu.memoryTotalGb
+ : undefined
+ }
+ systemRamGb={
+ gpu.systemRamAvailableGb || undefined
+ }
+ />
+ )}
+
+ );
+ })}
+ {!customFoldersCollapsed &&
+ showHfSection &&
+ sortedCustomFolderModels.length === 0 ? (
- No matching Unsloth models.
+ No matching models in custom folders.
- ) : null
- ) : (
- hfIds.map((id) => {
- const vram = vramMap.get(id);
- const isSearchGguf = isKnownGgufRepo(id);
- const optionKey = makeModelOptionKey("search-hf", id);
+ ) : null}
+ >
+ ) : null}
+
+ {section === "downloaded" && sortedLmStudio.length > 0 ? (
+ <>
+ setLmStudioCollapsed((v) => !v)}
+ >
+ LM Studio
+
+ {!lmStudioCollapsed &&
+ sortedLmStudio.map((m) => {
+ const isGgufFile = m.path
+ .toLowerCase()
+ .endsWith(".gguf");
+ // LM Studio dirs are GGUF but rarely carry a -GGUF suffix;
+ // use the shared helper (model_format hint) so the row,
+ // filter, and load path agree.
+ const isGguf = localModelIsGguf(m);
+ const optionKey = makeModelOptionKey("lm-studio", m.id);
+ return (
+
+
+
+ {
+ if (isGgufFile) {
+ onSelect(m.id, localDirectGgufMeta());
+ } else if (isGguf) {
+ toggleGgufExpanded(m.id);
+ } else {
+ onSelect(m.id, localModelMeta());
+ }
+ }}
+ onArrowDownIntoChildren={
+ isGguf &&
+ !isGgufFile &&
+ isGgufExpanded(m.id)
+ ? () => {
+ const focused =
+ focusFirstChildOption(optionKey);
+ return focused;
+ }
+ : undefined
+ }
+ vramStatus={null}
+ />
+
+ {isGgufFile && onConfigure && (
+
+ onConfigure(m.id, localDirectGgufMeta())
+ }
+ />
+ )}
+ {!isGguf && onConfigure && (
+
+ onConfigure(m.id, localModelMeta())
+ }
+ />
+ )}
+
+ {isGguf && !isGgufFile && isGgufExpanded(m.id) && (
+
+ hubModelList.focusOption(optionKey)
+ }
+ onNavigatePastEnd={() =>
+ hubModelList.moveFocus(optionKey, "next")
+ }
+ gpuGb={
+ gpu.available ? gpu.memoryTotalGb : undefined
+ }
+ systemRamGb={
+ gpu.systemRamAvailableGb || undefined
+ }
+ />
+ )}
+
+ );
+ })}
+ >
+ ) : null}
+
+ {section === "downloaded" && sortedLocalDir.length > 0 ? (
+ <>
+ setLocalDirCollapsed((v) => !v)}
+ >
+ Local models
+
+ {!localDirCollapsed &&
+ sortedLocalDir.map((m) => {
+ // A loose ./models/*.gguf loads directly; a GGUF repo dir
+ // expands to pick a variant. The variant scanner returns
+ // nothing for a config-less loose file, so expanding it would
+ // dead-end at "No GGUF variants".
+ const isGgufFile = m.path
+ .toLowerCase()
+ .endsWith(".gguf");
+ const isGguf = localModelIsGguf(m);
+ const optionKey = makeModelOptionKey("local-dir", m.id);
+ return (
+
+
+
+ {
+ if (isGgufFile) {
+ onSelect(m.id, localDirectGgufMeta());
+ } else if (isGguf) {
+ toggleGgufExpanded(m.id);
+ } else {
+ onSelect(m.id, localModelMeta());
+ }
+ }}
+ onArrowDownIntoChildren={
+ isGguf &&
+ !isGgufFile &&
+ isGgufExpanded(m.id)
+ ? () => focusFirstChildOption(optionKey)
+ : undefined
+ }
+ vramStatus={null}
+ />
+
+ {isGgufFile && onConfigure && (
+
+ onConfigure(m.id, localDirectGgufMeta())
+ }
+ />
+ )}
+ {!isGguf && onConfigure && (
+
+ onConfigure(m.id, localModelMeta())
+ }
+ />
+ )}
+
+ {isGguf && !isGgufFile && isGgufExpanded(m.id) && (
+
+ hubModelList.focusOption(optionKey)
+ }
+ onNavigatePastEnd={() =>
+ hubModelList.moveFocus(optionKey, "next")
+ }
+ gpuGb={
+ gpu.available ? gpu.memoryTotalGb : undefined
+ }
+ systemRamGb={
+ gpu.systemRamAvailableGb || undefined
+ }
+ />
+ )}
+
+ );
+ })}
+ >
+ ) : null}
+
+ {showRecommendedSection ? (
+ <>
+ {recommendedSearch.isLoading &&
+ recommendedRows.length === 0 ? (
+
+
+
+ Loading models…
+
+
+ ) : recommendedRows.length === 0 ? (
+
+ No models found.
+
+ ) : (
+ recommendedRows.map((r) => {
+ const id = r.id;
+ const info = recommendedMeta.get(id);
+ const isG = isKnownGgufRepo(id);
+ const optionKey = makeModelOptionKey("recommended", id);
+ return (
+
+ {
+ if (isG) {
+ setExpandedGguf((prev) =>
+ prev === id ? null : id,
+ );
+ } else {
+ handleModelClick(id);
+ }
+ }}
+ vramStatus={info?.status ?? null}
+ vramEst={info?.est}
+ gpuGb={
+ gpu.available ? gpu.memoryTotalGb : undefined
+ }
+ onArrowDownIntoChildren={
+ expandedGguf === id
+ ? () => focusFirstChildOption(optionKey)
+ : undefined
+ }
+ />
+ {expandedGguf === id && (
+
+ hubModelList.focusOption(optionKey)
+ }
+ onNavigatePastEnd={() =>
+ hubModelList.moveFocus(optionKey, "next")
+ }
+ gpuGb={
+ gpu.available ? gpu.memoryTotalGb : undefined
+ }
+ systemRamGb={
+ gpu.systemRamAvailableGb || undefined
+ }
+ variantActions={{
+ onDelete: async (quant) => {
+ await deleteCachedModel(
+ id,
+ quant,
+ hfToken || undefined,
+ );
+ prunePinnedQuantValidation(id, quant);
+ refreshCachedLists();
+ },
+ }}
+ />
+ )}
+
+ );
+ })
+ )}
+ {recommendedSearch.hasMore && (
+ <>
+
+
+
+
+ >
+ )}
+ >
+ ) : null}
+
+ {showHfSection &&
+ section === "recommended" &&
+ filteredRecommendedIds.length > 0 ? (
+ <>
+ {filteredRecommendedIds.map((id) => {
+ const vram = recommendedVramMap.get(id);
+ const optionKey = makeModelOptionKey(
+ "search-recommended",
+ id,
+ );
return (
{
- if (isSearchGguf) {
+ if (isKnownGgufRepo(id)) {
setExpandedGguf((prev) =>
prev === id ? null : id,
);
@@ -4264,9 +4079,13 @@ export function HubModelPicker({
}
}}
vramStatus={
- isSearchGguf ? null : (vram?.status ?? null)
+ isKnownGgufRepo(id)
+ ? null
+ : (vram?.status ?? null)
+ }
+ vramEst={
+ isKnownGgufRepo(id) ? undefined : vram?.est
}
- vramEst={isSearchGguf ? undefined : vram?.est}
gpuGb={
gpu.available ? gpu.memoryTotalGb : undefined
}
@@ -4284,6 +4103,7 @@ export function HubModelPicker({
@@ -4295,47 +4115,156 @@ export function HubModelPicker({
gpuGb={
gpu.available ? gpu.memoryTotalGb : undefined
}
- systemRamGb={gpu.systemRamAvailableGb || undefined}
+ systemRamGb={
+ gpu.systemRamAvailableGb || undefined
+ }
variantActions={{
onDelete: async (quant) => {
- await deleteCachedModel(id, quant);
+ await deleteCachedModel(
+ id,
+ quant,
+ hfToken || undefined,
+ );
+ prunePinnedQuantValidation(id, quant);
refreshCachedLists();
},
- deleteDisabled,
}}
/>
)}
);
- })
- )}
-
- {isLoadingMore ? (
-
-
-
- ) : null}
- >
- ) : null}
- >
- )}
+ })}
+ >
+ ) : null}
+
+ {showHfSection && section === "recommended" ? (
+ <>
+ {hfIds.length === 0 && !isLoading ? (
+ filteredRecommendedIds.length === 0 ? (
+
+ No matching Unsloth models.
+
+ ) : null
+ ) : (
+ hfIds.map((id) => {
+ const vram = vramMap.get(id);
+ const isSearchGguf = isKnownGgufRepo(id);
+ const optionKey = makeModelOptionKey("search-hf", id);
+ return (
+
+ {
+ if (isSearchGguf) {
+ setExpandedGguf((prev) =>
+ prev === id ? null : id,
+ );
+ } else {
+ handleModelClick(id);
+ }
+ }}
+ vramStatus={
+ isSearchGguf ? null : (vram?.status ?? null)
+ }
+ vramEst={isSearchGguf ? undefined : vram?.est}
+ gpuGb={
+ gpu.available ? gpu.memoryTotalGb : undefined
+ }
+ onArrowDownIntoChildren={
+ expandedGguf === id
+ ? () => {
+ const focused =
+ focusFirstChildOption(optionKey);
+ return focused;
+ }
+ : undefined
+ }
+ />
+ {expandedGguf === id && (
+
+ hubModelList.focusOption(optionKey)
+ }
+ onNavigatePastEnd={() =>
+ hubModelList.moveFocus(optionKey, "next")
+ }
+ gpuGb={
+ gpu.available ? gpu.memoryTotalGb : undefined
+ }
+ systemRamGb={
+ gpu.systemRamAvailableGb || undefined
+ }
+ variantActions={{
+ onDelete: async (quant) => {
+ await deleteCachedModel(
+ id,
+ quant,
+ hfToken || undefined,
+ );
+ prunePinnedQuantValidation(id, quant);
+ refreshCachedLists();
+ },
+ }}
+ />
+ )}
+
+ );
+ })
+ )}
+
+ {isLoadingMore ? (
+
+
+
+ ) : null}
+ >
+ ) : null}
+ >
+ )}
+
-
- {/* Floating eject pill: overlaid on the list bottom, outside the scroll
- so the edge fade never touches it. Only the pill catches clicks. */}
- {onEject ? (
-
-
-
- Eject model
-
-
- ) : null}
+ {onEject ? (
+
+
+
+ Eject model
+
+
+ ) : null}
void;
+ onConfigure?: (id: string, meta: ModelSelectorChangeMeta) => void;
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
deleteDisabled?: boolean;
loraModelList: ReturnType;
@@ -4391,6 +4322,13 @@ function FineTunedRows({
const isTrainingFull = isTraining && isMerged;
const isLocalGgufDir =
isLocal && (isGgufRepo(adapter.id) || isGgufRepo(adapter.name));
+ const selectionMeta: ModelSelectorChangeMeta = {
+ source: isLocal ? "local" : isExported ? "exported" : "lora",
+ isLora: !isLocal && !isMerged && !isGguf,
+ isDownloaded: true,
+ isGguf: false,
+ };
+ const canConfigure = !(isLocalGgufDir || isExportedGguf);
const optionKey = makeModelOptionKey("lora", adapter.id);
const tag = isLocal
? isLocalGgufDir
@@ -4438,15 +4376,7 @@ function FineTunedRows({
prev === adapter.id ? null : adapter.id,
);
} else {
- onSelect(adapter.id, {
- source: isLocal
- ? "local"
- : isExported
- ? "exported"
- : "lora",
- isLora: !isLocal && !isMerged && !isGguf,
- isDownloaded: true,
- });
+ onSelect(adapter.id, selectionMeta);
}
}}
tooltipText={
@@ -4467,6 +4397,12 @@ function FineTunedRows({
}
/>
+ {canConfigure && onConfigure && (
+
onConfigure(adapter.id, selectionMeta)}
+ />
+ )}
{canDelete && (
loraModelList.focusOption(optionKey)}
onNavigatePastEnd={() =>
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pill-tabs.tsx b/studio/frontend/src/features/model-picker/components/model-selector/pill-tabs.tsx
similarity index 98%
rename from studio/frontend/src/components/assistant-ui/model-selector/pill-tabs.tsx
rename to studio/frontend/src/features/model-picker/components/model-selector/pill-tabs.tsx
index e6da8a7b74..fbc1d5ac91 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pill-tabs.tsx
+++ b/studio/frontend/src/features/model-picker/components/model-selector/pill-tabs.tsx
@@ -78,7 +78,8 @@ export function PillTabs({
onValueChange(tabs[next].value);
e.currentTarget.parentElement
?.querySelectorAll('button[role="tab"]')
- [next]?.focus();
+ .item(next)
+ ?.focus();
}}
onClick={() => onValueChange(tab.value)}
className={cn(
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts b/studio/frontend/src/features/model-picker/components/model-selector/pinned-models.ts
similarity index 78%
rename from studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts
rename to studio/frontend/src/features/model-picker/components/model-selector/pinned-models.ts
index 4835c4c0cf..0444b9f3cc 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts
+++ b/studio/frontend/src/features/model-picker/components/model-selector/pinned-models.ts
@@ -21,6 +21,13 @@ export interface PinnedQuantEntry {
quant: string;
}
+export function makePinRank(
+ pinned: readonly string[],
+): (key: string) => number {
+ const pinIndex = new Map(pinned.map((key, index) => [key, index]));
+ return (key) => pinIndex.get(key) ?? Number.MAX_SAFE_INTEGER;
+}
+
/** The pinned GGUF quants, in pin order. Plain repo pins are excluded. */
export function pinnedQuantEntries(pinned: string[]): PinnedQuantEntry[] {
const out: PinnedQuantEntry[] = [];
@@ -63,10 +70,20 @@ export const usePinnedModelsStore = create((set) => ({
togglePinned: (repoId, quant) =>
set((state) => {
const key = pinKey(repoId, quant);
+ // Newest pin first, so "Pin to top" literally lands on top of the
+ // pinned group rather than under earlier pins.
const next = state.pinned.includes(key)
? state.pinned.filter((id) => id !== key)
- : [...state.pinned, key];
+ : [key, ...state.pinned];
writePinned(next);
return { pinned: next };
}),
}));
+
+if (typeof window !== "undefined") {
+ window.addEventListener("storage", (event) => {
+ if (event.key === KEY || event.key === null) {
+ usePinnedModelsStore.setState({ pinned: readPinned() });
+ }
+ });
+}
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts b/studio/frontend/src/features/model-picker/components/model-selector/recommended-fit.ts
similarity index 91%
rename from studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts
rename to studio/frontend/src/features/model-picker/components/model-selector/recommended-fit.ts
index 7c2ed266c0..b8fe47c706 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts
+++ b/studio/frontend/src/features/model-picker/components/model-selector/recommended-fit.ts
@@ -64,9 +64,8 @@ export function matchesFormatFilter(
}
}
-// 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"
+// First "B" token in a repo id, e.g. "Qwen3-30B-A3B" -> 30 (MoE total),
+// "gemma-4-E4B" -> 4. Digits must be separator-bounded so we never read "16"
// from "bf16" or the "2" in "Kimi-K2".
const PARAM_RE = /(?:^|[-_/. ])[eE]?(\d+(?:\.\d+)?)\s*[bB](?=$|[-_./ ])/;
@@ -79,9 +78,8 @@ export function paramsFromId(id: string): number | undefined {
return Number.isFinite(billions) && billions > 0 ? billions * 1e9 : undefined;
}
-// Smallest practical GGUF/MLX quant (~Q2_K, low-bit). The fit check asks whether
-// a model can run at all, so it uses this rather than a default 4-bit size; a
-// user with a smaller device can still pick a low-bit variant.
+// Smallest practical GGUF/MLX quant (~Q2_K). The fit check asks whether a model
+// can run at all, so it uses this rather than a default 4-bit size.
const MIN_QUANT_BYTES_PER_PARAM = 0.4;
/** Rough on-disk bytes for the smallest practical quant of `params` weights. */
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/row-meta.ts b/studio/frontend/src/features/model-picker/components/model-selector/row-meta.ts
similarity index 100%
rename from studio/frontend/src/components/assistant-ui/model-selector/row-meta.ts
rename to studio/frontend/src/features/model-picker/components/model-selector/row-meta.ts
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/source-tabs.ts b/studio/frontend/src/features/model-picker/components/model-selector/source-tabs.ts
similarity index 100%
rename from studio/frontend/src/components/assistant-ui/model-selector/source-tabs.ts
rename to studio/frontend/src/features/model-picker/components/model-selector/source-tabs.ts
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/features/model-picker/components/model-selector/types.ts
similarity index 74%
rename from studio/frontend/src/components/assistant-ui/model-selector/types.ts
rename to studio/frontend/src/features/model-picker/components/model-selector/types.ts
index 6a86515267..9adf2d899e 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts
+++ b/studio/frontend/src/features/model-picker/components/model-selector/types.ts
@@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { ReactNode } from "react";
+import type { PerModelConfig } from "../../model-config/per-model-config";
export interface ModelOption {
id: string;
@@ -36,6 +37,19 @@ export interface ModelSelectorChangeMeta {
/** Direct local .gguf file picked without a variant (custom folder / LM
* Studio). Marks it as a GGUF source for the deferred-load staging flow. */
isGguf?: boolean;
+ config?: PerModelConfig;
+ forceReload?: boolean;
+ /** Native path token so an active-model reload can reopen a file-picked GGUF. */
+ nativePathToken?: string;
+ nativePathExpiresAtMs?: number | null;
+}
+
+export interface ModelPickTarget {
+ id: string;
+ displayName: string;
+ ggufVariant?: string | null;
+ isGguf: boolean;
+ meta: ModelSelectorChangeMeta;
}
export interface DeletedModelRef {
diff --git a/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx b/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx
new file mode 100644
index 0000000000..2489927fc2
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx
@@ -0,0 +1,113 @@
+// 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 { cn } from "@/lib/utils";
+import { useRef, useState } from "react";
+
+export function snapToStep(
+ value: number,
+ step: number,
+ min?: number,
+ max?: number,
+): number {
+ const lo = min ?? Number.NEGATIVE_INFINITY;
+ const hi = max ?? Number.POSITIVE_INFINITY;
+ const clamped = Math.min(Math.max(value, lo), hi);
+ const stepStr = String(step);
+ const decimals = stepStr.includes(".") ? stepStr.split(".")[1].length : 0;
+ const base = Number.isFinite(lo) ? lo : 0;
+ const snapped = base + Math.round((clamped - base) / step) * step;
+ const reclamped = Math.min(Math.max(snapped, lo), hi);
+ return Number(reclamped.toFixed(decimals));
+}
+
+function sanitizeNumeric(raw: string, allowNegative: boolean): string {
+ const sign = allowNegative && raw.startsWith("-") ? "-" : "";
+ const [head, ...rest] = raw.replace(/[^\d.]/g, "").split(".");
+ const tail = rest.length > 0 ? `.${rest.join("")}` : "";
+ return `${sign}${head}${tail}`;
+}
+
+export function NumericValueInput({
+ value,
+ min,
+ max,
+ step,
+ onChange,
+ displayValue,
+ className,
+ ariaLabel,
+ size: sizeAttr,
+ disabled = false,
+}: {
+ value: number;
+ min?: number;
+ max?: number;
+ step: number;
+ onChange: (v: number) => void;
+ displayValue?: string;
+ className?: string;
+ ariaLabel?: string;
+ size?: number;
+ disabled?: boolean;
+}) {
+ const [focused, setFocused] = useState(false);
+ const [draft, setDraft] = useState("");
+ const cancelBlurCommitRef = useRef(false);
+
+ const commit = (raw: string) => {
+ const parsed = Number.parseFloat(raw);
+ if (!Number.isFinite(parsed)) {
+ return;
+ }
+ const final = snapToStep(parsed, step, min, max);
+ if (final !== value) {
+ onChange(final);
+ }
+ };
+
+ const displayed = focused ? draft : (displayValue ?? String(value));
+
+ return (
+ {
+ cancelBlurCommitRef.current = false;
+ setDraft(String(value));
+ setFocused(true);
+ const target = e.currentTarget;
+ requestAnimationFrame(() => target.select());
+ }}
+ onBlur={() => {
+ if (cancelBlurCommitRef.current) {
+ cancelBlurCommitRef.current = false;
+ } else {
+ commit(draft);
+ }
+ setFocused(false);
+ }}
+ onChange={(e) =>
+ setDraft(sanitizeNumeric(e.target.value, (min ?? 0) < 0))
+ }
+ onKeyDown={(e) => {
+ if (e.key === "Enter") {
+ e.currentTarget.blur();
+ } else if (e.key === "Escape") {
+ cancelBlurCommitRef.current = true;
+ setDraft(String(value));
+ e.currentTarget.blur();
+ }
+ }}
+ className={cn(className)}
+ />
+ );
+}
diff --git a/studio/frontend/src/features/model-picker/components/sidebar-model-config.tsx b/studio/frontend/src/features/model-picker/components/sidebar-model-config.tsx
new file mode 100644
index 0000000000..2d12c503a4
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/components/sidebar-model-config.tsx
@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { useMemo } from "react";
+import { gpuFieldsSignature } from "../model-config/apply-per-model-config";
+import type { PerModelConfig } from "../model-config/per-model-config";
+import { ModelConfigPage } from "./model-config-page";
+import type { ModelPickTarget } from "./model-selector/types";
+
+interface SidebarModelConfigProps {
+ modelId: string;
+ ggufVariant: string | null;
+ isGguf: boolean;
+ nativeContextLength: number | null;
+ loadedContextLength: number | null;
+ loadedConfig: PerModelConfig;
+ onReload: (config: PerModelConfig) => void;
+}
+
+const TRAILING_SEPARATORS = /[\\/]+$/;
+
+function leafName(id: string): string {
+ const trimmed = id.replace(TRAILING_SEPARATORS, "");
+ const separator = Math.max(
+ trimmed.lastIndexOf("/"),
+ trimmed.lastIndexOf("\\"),
+ );
+ return separator >= 0 ? trimmed.slice(separator + 1) : trimmed;
+}
+
+function hashString(value: string): number {
+ let hash = 5381;
+ for (let i = 0; i < value.length; i += 1) {
+ hash = (Math.imul(hash, 33) ^ value.charCodeAt(i)) >>> 0;
+ }
+ return hash;
+}
+
+function configSignature(config: PerModelConfig): string {
+ return [
+ config.customContextLength ?? "",
+ config.maxSeqLength ?? "",
+ config.kvCacheDtype ?? "",
+ config.speculativeType ?? "",
+ config.specDraftNMax ?? "",
+ config.tensorParallel ? "1" : "0",
+ config.chatTemplateOverride == null
+ ? ""
+ : `${config.chatTemplateOverride.length}:${hashString(config.chatTemplateOverride)}`,
+ gpuFieldsSignature(config),
+ ].join("|");
+}
+
+export function SidebarModelConfig({
+ modelId,
+ ggufVariant,
+ isGguf,
+ nativeContextLength,
+ loadedContextLength,
+ loadedConfig,
+ onReload,
+}: SidebarModelConfigProps) {
+ const target = useMemo(() => {
+ const leaf = leafName(modelId);
+ return {
+ id: modelId,
+ displayName: ggufVariant ? `${leaf} · ${ggufVariant}` : leaf,
+ ggufVariant,
+ isGguf,
+ meta: {
+ source: "local",
+ isLora: false,
+ ggufVariant: ggufVariant ?? undefined,
+ isGguf,
+ isDownloaded: true,
+ contextLength: nativeContextLength,
+ },
+ };
+ }, [modelId, ggufVariant, isGguf, nativeContextLength]);
+
+ return (
+
+ );
+}
diff --git a/studio/frontend/src/features/model-picker/hooks/use-active-model-config.ts b/studio/frontend/src/features/model-picker/hooks/use-active-model-config.ts
new file mode 100644
index 0000000000..9d09ee6897
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/hooks/use-active-model-config.ts
@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { isExternalModelId, useChatRuntimeStore } from "@/features/chat";
+import { useMemo } from "react";
+import type { PerModelConfig } from "../model-config/per-model-config";
+
+export interface ActiveModelConfigState {
+ checkpoint: string | null;
+ isGguf: boolean;
+ config: PerModelConfig | null;
+}
+
+export function useActiveModelConfig(): ActiveModelConfigState {
+ const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint) || null;
+ const maxSeqLength = useChatRuntimeStore((s) => s.params.maxSeqLength);
+ const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
+ const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
+ const customContextLength = useChatRuntimeStore((s) => s.customContextLength);
+ const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype);
+ const speculativeType = useChatRuntimeStore((s) => s.speculativeType);
+ const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax);
+ const tensorParallel = useChatRuntimeStore((s) => s.tensorParallel);
+ const chatTemplateOverride = useChatRuntimeStore(
+ (s) => s.chatTemplateOverride,
+ );
+ const gpuMemoryMode = useChatRuntimeStore((s) => s.gpuMemoryMode);
+ const gpuLayers = useChatRuntimeStore((s) => s.gpuLayers);
+ const nCpuMoe = useChatRuntimeStore((s) => s.nCpuMoe);
+ const selectedGpuIds = useChatRuntimeStore((s) => s.selectedGpuIds);
+
+ const isGguf =
+ activeGgufVariant != null ||
+ ggufContextLength != null ||
+ (checkpoint?.toLowerCase().endsWith(".gguf") ?? false);
+
+ const config = useMemo(() => {
+ if (!checkpoint || isExternalModelId(checkpoint)) {
+ return null;
+ }
+ const base: PerModelConfig = {
+ customContextLength: customContextLength ?? null,
+ maxSeqLength: isGguf ? null : maxSeqLength,
+ kvCacheDtype: kvCacheDtype ?? null,
+ speculativeType: speculativeType ?? "auto",
+ specDraftNMax: specDraftNMax ?? null,
+ tensorParallel: tensorParallel ?? false,
+ chatTemplateOverride: chatTemplateOverride ?? null,
+ };
+ if (!isGguf) {
+ return base;
+ }
+ return {
+ ...base,
+ gpuMemoryMode,
+ gpuLayers,
+ nCpuMoe,
+ selectedGpuIds,
+ };
+ }, [
+ checkpoint,
+ isGguf,
+ maxSeqLength,
+ customContextLength,
+ kvCacheDtype,
+ speculativeType,
+ specDraftNMax,
+ tensorParallel,
+ chatTemplateOverride,
+ gpuMemoryMode,
+ gpuLayers,
+ nCpuMoe,
+ selectedGpuIds,
+ ]);
+
+ return { checkpoint, isGguf, config };
+}
diff --git a/studio/frontend/src/features/model-picker/hooks/use-model-defaults.ts b/studio/frontend/src/features/model-picker/hooks/use-model-defaults.ts
new file mode 100644
index 0000000000..530e6e06d8
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/hooks/use-model-defaults.ts
@@ -0,0 +1,191 @@
+// 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 { useHfTokenStore, useInventoryVersion } from "@/features/hub";
+import { useEffect, useState } from "react";
+import { fetchModelMaxPositionEmbeddings } from "../api/model-metadata";
+import { fetchDefaultChatTemplate } from "../api/templates";
+
+export interface DefaultChatTemplateState {
+ template: string | null;
+ loading: boolean;
+ error: string | null;
+}
+
+export interface ModelMaxPositionState {
+ maxPositionEmbeddings: number | null;
+ loading: boolean;
+ error: string | null;
+}
+
+const TEMPLATE_CACHE_MAX_ENTRIES = 50;
+const templateCache = new Map();
+const maxPositionCache = new Map();
+
+function cacheTemplate(key: string, template: string | null): void {
+ templateCache.delete(key);
+ templateCache.set(key, template);
+ while (templateCache.size > TEMPLATE_CACHE_MAX_ENTRIES) {
+ const oldest = templateCache.keys().next().value;
+ if (oldest === undefined) {
+ break;
+ }
+ templateCache.delete(oldest);
+ }
+}
+
+function cacheMaxPosition(key: string, value: number | null): void {
+ maxPositionCache.delete(key);
+ maxPositionCache.set(key, value);
+ while (maxPositionCache.size > TEMPLATE_CACHE_MAX_ENTRIES) {
+ const oldest = maxPositionCache.keys().next().value;
+ if (oldest === undefined) {
+ break;
+ }
+ maxPositionCache.delete(oldest);
+ }
+}
+
+export function useDefaultChatTemplate(
+ modelId: string | null,
+ ggufVariant: string | null | undefined,
+ enabled: boolean,
+ nativePathToken?: string | null,
+): DefaultChatTemplateState {
+ const token = useHfTokenStore((s) => s.token);
+ const inventoryVersion = useInventoryVersion();
+ // The native token is part of the identity: a picked GGUF resolves its
+ // template through the lease, not the model id, so two picks of the same
+ // basename must not share a cache entry.
+ const cacheKey =
+ enabled && modelId
+ ? `${modelId}::${ggufVariant ?? ""}::${token}::${inventoryVersion}::${nativePathToken ?? ""}`
+ : null;
+ const [fetched, setFetched] = useState<{
+ key: string;
+ state: DefaultChatTemplateState;
+ } | null>(null);
+
+ useEffect(() => {
+ if (cacheKey == null || !modelId || templateCache.has(cacheKey)) {
+ return;
+ }
+ const controller = new AbortController();
+ fetchDefaultChatTemplate(
+ modelId,
+ ggufVariant,
+ token,
+ controller.signal,
+ nativePathToken,
+ )
+ .then((template) => {
+ if (controller.signal.aborted) {
+ return;
+ }
+ // Cache the terminal result, including a null "no default template",
+ // so reopening the viewer for such a model reuses it instead of
+ // re-running the backend/Hugging Face lookup every time.
+ cacheTemplate(cacheKey, template);
+ setFetched({
+ key: cacheKey,
+ state: { template, loading: false, error: null },
+ });
+ })
+ .catch((err: unknown) => {
+ if (controller.signal.aborted) {
+ return;
+ }
+ setFetched({
+ key: cacheKey,
+ state: {
+ template: null,
+ loading: false,
+ error:
+ err instanceof Error ? err.message : "Failed to load template",
+ },
+ });
+ });
+
+ return () => controller.abort();
+ }, [cacheKey, modelId, ggufVariant, token, nativePathToken]);
+
+ if (cacheKey == null) {
+ return { template: null, loading: false, error: null };
+ }
+ if (templateCache.has(cacheKey)) {
+ return {
+ template: templateCache.get(cacheKey) ?? null,
+ loading: false,
+ error: null,
+ };
+ }
+ if (fetched?.key === cacheKey) {
+ return fetched.state;
+ }
+ return { template: null, loading: true, error: null };
+}
+
+export function useModelMaxPositionEmbeddings(
+ modelId: string | null,
+ enabled: boolean,
+): ModelMaxPositionState {
+ const token = useHfTokenStore((s) => s.token);
+ const inventoryVersion = useInventoryVersion();
+ const cacheKey =
+ enabled && modelId ? `${modelId}::${token}::${inventoryVersion}` : null;
+ const [fetched, setFetched] = useState<{
+ key: string;
+ state: ModelMaxPositionState;
+ } | null>(null);
+
+ useEffect(() => {
+ if (cacheKey == null || !modelId || maxPositionCache.has(cacheKey)) {
+ return;
+ }
+ const controller = new AbortController();
+ fetchModelMaxPositionEmbeddings(modelId, token, controller.signal)
+ .then((maxPositionEmbeddings) => {
+ if (controller.signal.aborted) {
+ return;
+ }
+ cacheMaxPosition(cacheKey, maxPositionEmbeddings);
+ setFetched({
+ key: cacheKey,
+ state: { maxPositionEmbeddings, loading: false, error: null },
+ });
+ })
+ .catch((err: unknown) => {
+ if (controller.signal.aborted) {
+ return;
+ }
+ setFetched({
+ key: cacheKey,
+ state: {
+ maxPositionEmbeddings: null,
+ loading: false,
+ error:
+ err instanceof Error
+ ? err.message
+ : "Failed to load model metadata",
+ },
+ });
+ });
+
+ return () => controller.abort();
+ }, [cacheKey, modelId, token]);
+
+ if (cacheKey == null) {
+ return { maxPositionEmbeddings: null, loading: false, error: null };
+ }
+ if (maxPositionCache.has(cacheKey)) {
+ return {
+ maxPositionEmbeddings: maxPositionCache.get(cacheKey) ?? null,
+ loading: false,
+ error: null,
+ };
+ }
+ if (fetched?.key === cacheKey) {
+ return fetched.state;
+ }
+ return { maxPositionEmbeddings: null, loading: true, error: null };
+}
diff --git a/studio/frontend/src/features/model-picker/index.ts b/studio/frontend/src/features/model-picker/index.ts
new file mode 100644
index 0000000000..d2b4785ec3
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/index.ts
@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export { ModelSelector } from "./components/model-selector";
+export { FolderBrowser } from "./components/model-selector/folder-browser";
+export { ModelRowMenu } from "./components/model-selector/model-row-menu";
+export {
+ makePinRank,
+ pinKey,
+ usePinnedModelsStore,
+} from "./components/model-selector/pinned-models";
+export { hfModelFitsDevice } from "./components/model-selector/recommended-fit";
+export {
+ NumericValueInput,
+ snapToStep,
+} from "./components/numeric-value-input";
+export { SidebarModelConfig } from "./components/sidebar-model-config";
+export {
+ useActiveModelConfig,
+} from "./hooks/use-active-model-config";
+export type {
+ DeletedModelRef,
+ ExternalModelOption,
+ LoraModelOption,
+ ModelOption,
+ ModelSelectorChangeMeta,
+} from "./components/model-selector";
+export {
+ applyModelLoadConfigToRuntime,
+ applyPerModelConfigToRuntime,
+ currentRuntimePerModelConfig,
+ perModelConfigsEqual,
+} from "./model-config/apply-per-model-config";
+export {
+ DEFAULT_MAX_SEQ_LENGTH,
+ normalizeMaxSeqLength,
+ type PerModelConfig,
+ resolveInitialConfig,
+} from "./model-config/per-model-config";
diff --git a/studio/frontend/src/features/model-picker/inventory/use-chat-picker-inventory.ts b/studio/frontend/src/features/model-picker/inventory/use-chat-picker-inventory.ts
new file mode 100644
index 0000000000..b46e83819c
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/inventory/use-chat-picker-inventory.ts
@@ -0,0 +1,123 @@
+// 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 {
+ CachedGgufRepo,
+ CachedModelRepo,
+ LocalModelInfo,
+} from "@/features/chat";
+import {
+ type CachedInventoryRow,
+ type LocalInventoryRow,
+ type LocalSource,
+ isHiddenModelId,
+ useHubInventory,
+} from "@/features/hub";
+import { useMemo } from "react";
+
+const PICKER_LOCAL_SOURCES: ReadonlySet = new Set([
+ "lmstudio",
+ "models_dir",
+ "custom",
+]);
+
+function isCompleteCachedRow(row: CachedInventoryRow): boolean {
+ return !row.partial && !row.liveDownload;
+}
+
+function toCachedGgufRepo(row: CachedInventoryRow): CachedGgufRepo {
+ return {
+ repo_id: row.repoId,
+ size_bytes: row.bytes,
+ cache_path: row.cachePath ?? "",
+ last_modified: row.lastModified ?? undefined,
+ has_vision: row.capabilities.supportsVision,
+ };
+}
+
+function toCachedModelRepo(row: CachedInventoryRow): CachedModelRepo {
+ return {
+ repo_id: row.repoId,
+ size_bytes: row.bytes,
+ last_modified: row.lastModified ?? undefined,
+ };
+}
+
+function toLocalModelInfo(row: LocalInventoryRow): LocalModelInfo {
+ return {
+ id: row.loadId,
+ display_name: row.displayName ?? row.title,
+ path: row.path,
+ source: row.source as LocalModelInfo["source"],
+ model_id: row.modelId ?? row.repoId,
+ model_format: row.modelFormat,
+ updated_at: row.updatedAt,
+ };
+}
+
+export interface ChatPickerInventory {
+ cachedGguf: CachedGgufRepo[];
+ cachedModels: CachedModelRepo[];
+ cachedReady: boolean;
+ localModels: LocalModelInfo[];
+ refreshInventory: () => Promise;
+}
+
+export function useChatPickerInventory(
+ options: { enabled?: boolean } = {},
+): ChatPickerInventory {
+ const inventory = useHubInventory({
+ kind: "models",
+ enabled: options.enabled,
+ includeLocal: true,
+ });
+
+ const cachedGguf = useMemo(
+ () =>
+ inventory.cachedRows
+ .filter(
+ (row) =>
+ row.modelFormat === "gguf" &&
+ isCompleteCachedRow(row) &&
+ !isHiddenModelId(row.repoId),
+ )
+ .map(toCachedGgufRepo),
+ [inventory.cachedRows],
+ );
+ const cachedModels = useMemo(
+ () =>
+ inventory.cachedRows
+ .filter(
+ (row) =>
+ row.modelFormat !== "gguf" &&
+ isCompleteCachedRow(row) &&
+ !isHiddenModelId(row.repoId),
+ )
+ .map(toCachedModelRepo),
+ [inventory.cachedRows],
+ );
+ const localModels = useMemo(
+ () =>
+ inventory.localRows
+ .filter(
+ (row) =>
+ PICKER_LOCAL_SOURCES.has(row.source) &&
+ // Skip non-chat rows (e.g. a folder with only config.json is
+ // classified "unknown" -> canChat false); selecting one would try to
+ // load a weightless path. toLocalModelInfo drops capabilities, so
+ // this is the only place the guard can live.
+ row.capabilities.canChat &&
+ !isHiddenModelId(row.modelId, row.repoId, row.path),
+ )
+ .map(toLocalModelInfo),
+ [inventory.localRows],
+ );
+
+ return {
+ cachedGguf,
+ cachedModels,
+ cachedReady: inventory.downloadedReady,
+ localModels,
+ refreshInventory: inventory.refreshInventory,
+ };
+}
diff --git a/studio/frontend/src/features/model-picker/model-config/apply-per-model-config.ts b/studio/frontend/src/features/model-picker/model-config/apply-per-model-config.ts
new file mode 100644
index 0000000000..c21d3e164a
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/model-config/apply-per-model-config.ts
@@ -0,0 +1,127 @@
+// 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 {
+ GPU_LAYERS_AUTO,
+ defaultInferenceParams,
+ normalizeSpeculativeType,
+ readPersistedGpuMemoryMode,
+ readPersistedSpeculativeType,
+ reconcilePersistedGpuIds,
+ useChatRuntimeStore,
+} from "@/features/chat";
+import {
+ DEFAULT_PER_MODEL_CONFIG,
+ type PerModelConfig,
+ normalizeMaxSeqLength,
+} from "./per-model-config";
+
+function cleanTemplate(value: string | null | undefined): string | null {
+ return value?.trim() ? value : null;
+}
+
+export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
+ // Fall back to the standing default when the model has no saved
+ // maxSeqLength. maxSeqLength is the only per-model field carried on
+ // params (the rest are reset below), so without this a model with no
+ // remembered config would inherit the previously loaded model's value.
+ const maxSeqLength =
+ normalizeMaxSeqLength(config.maxSeqLength) ??
+ defaultInferenceParams.maxSeqLength;
+ const store = useChatRuntimeStore.getState();
+ if (maxSeqLength !== store.params.maxSeqLength) {
+ store.setParams({ ...store.params, maxSeqLength });
+ }
+ useChatRuntimeStore.setState({
+ customContextLength: config.customContextLength ?? null,
+ kvCacheDtype: config.kvCacheDtype ?? null,
+ speculativeType:
+ normalizeSpeculativeType(config.speculativeType) ??
+ readPersistedSpeculativeType(),
+ specDraftNMax: config.specDraftNMax ?? null,
+ tensorParallel: config.tensorParallel ?? false,
+ chatTemplateOverride: cleanTemplate(config.chatTemplateOverride),
+ // GPU Memory knobs are per-model (GGUF-only). Absent = defaults; the mode is
+ // a standing preference so an absent mode falls back to the persisted one.
+ // The per-GPU split ratio is never remembered, so it always resets. The GPU
+ // pick is reconciled against the GPUs present now (a saved [1] on a 1-GPU
+ // host would otherwise be sent and rejected).
+ gpuMemoryMode: config.gpuMemoryMode ?? readPersistedGpuMemoryMode(),
+ gpuLayers: config.gpuLayers ?? GPU_LAYERS_AUTO,
+ nCpuMoe: config.nCpuMoe ?? 0,
+ splitRatio: null,
+ selectedGpuIds:
+ config.selectedGpuIds !== undefined
+ ? reconcilePersistedGpuIds(config.selectedGpuIds)
+ : null,
+ });
+}
+
+export function applyModelLoadConfigToRuntime(
+ config: PerModelConfig | null | undefined,
+): boolean {
+ const hasConfig = config != null;
+ applyPerModelConfigToRuntime(config ?? DEFAULT_PER_MODEL_CONFIG);
+ return hasConfig;
+}
+
+export function currentRuntimePerModelConfig(
+ options: { includeMaxSeqLength?: boolean } = {},
+): PerModelConfig {
+ const s = useChatRuntimeStore.getState();
+ return {
+ customContextLength: s.customContextLength ?? null,
+ maxSeqLength: options.includeMaxSeqLength
+ ? normalizeMaxSeqLength(s.params.maxSeqLength)
+ : null,
+ kvCacheDtype: s.kvCacheDtype ?? null,
+ speculativeType: normalizeSpeculativeType(s.speculativeType),
+ specDraftNMax: s.specDraftNMax ?? null,
+ tensorParallel: s.tensorParallel ?? false,
+ chatTemplateOverride: cleanTemplate(s.chatTemplateOverride),
+ // Snapshot the live GPU knobs too so a failed switch rolls the previous
+ // model's GPU Memory settings back (see applyPerModelConfigToRuntime). The
+ // split ratio is intentionally never remembered.
+ gpuMemoryMode: s.gpuMemoryMode,
+ gpuLayers: s.gpuLayers,
+ nCpuMoe: s.nCpuMoe,
+ selectedGpuIds: s.selectedGpuIds,
+ };
+}
+
+export function perModelConfigsEqual(
+ a: PerModelConfig,
+ b: PerModelConfig,
+): boolean {
+ return (
+ (a.customContextLength ?? null) === (b.customContextLength ?? null) &&
+ normalizeMaxSeqLength(a.maxSeqLength) ===
+ normalizeMaxSeqLength(b.maxSeqLength) &&
+ (a.kvCacheDtype ?? null) === (b.kvCacheDtype ?? null) &&
+ normalizeSpeculativeType(a.speculativeType) ===
+ normalizeSpeculativeType(b.speculativeType) &&
+ (a.specDraftNMax ?? null) === (b.specDraftNMax ?? null) &&
+ Boolean(a.tensorParallel) === Boolean(b.tensorParallel) &&
+ cleanTemplate(a.chatTemplateOverride) ===
+ cleanTemplate(b.chatTemplateOverride) &&
+ gpuFieldsEqual(a, b)
+ );
+}
+
+// Serialize the per-model GPU knobs with the same "absent == default"
+// coalescing the store applies: mode auto/absent, gpuLayers Auto (< 0) /
+// absent, nCpuMoe 0 / absent, and the GPU pick (null / absent = all GPUs).
+export function gpuFieldsSignature(config: PerModelConfig): string {
+ return [
+ config.gpuMemoryMode ?? "auto",
+ config.gpuLayers == null || config.gpuLayers < 0 ? -1 : config.gpuLayers,
+ config.nCpuMoe ?? 0,
+ config.selectedGpuIds == null
+ ? "all"
+ : [...config.selectedGpuIds].sort((a, b) => a - b).join(","),
+ ].join("|");
+}
+
+function gpuFieldsEqual(a: PerModelConfig, b: PerModelConfig): boolean {
+ return gpuFieldsSignature(a) === gpuFieldsSignature(b);
+}
diff --git a/studio/frontend/src/features/model-picker/model-config/model-identity.ts b/studio/frontend/src/features/model-picker/model-config/model-identity.ts
new file mode 100644
index 0000000000..0caa7c1312
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/model-config/model-identity.ts
@@ -0,0 +1,69 @@
+// 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 {
+ normalizeGgufVariantIdentity,
+ normalizeModelIdentity,
+} from "@/features/hub";
+
+export {
+ normalizeGgufVariantIdentity,
+ normalizeModelIdentity,
+} from "@/features/hub";
+
+const MODEL_STORAGE_KEY_PREFIX = "v2:";
+
+type ParsedModelStorageKey = {
+ modelId: string;
+ ggufVariant: string;
+};
+
+function parseVersionedModelStorageKey(
+ key: string,
+): ParsedModelStorageKey | null {
+ if (!key.startsWith(MODEL_STORAGE_KEY_PREFIX)) {
+ return null;
+ }
+ try {
+ const parsed = JSON.parse(key.slice(MODEL_STORAGE_KEY_PREFIX.length));
+ if (
+ !Array.isArray(parsed) ||
+ parsed.length !== 2 ||
+ typeof parsed[0] !== "string" ||
+ typeof parsed[1] !== "string"
+ ) {
+ return null;
+ }
+ return { modelId: parsed[0], ggufVariant: parsed[1] };
+ } catch {
+ return null;
+ }
+}
+
+export function modelStorageKey(
+ modelId: string,
+ ggufVariant?: string | null,
+): string {
+ return `${MODEL_STORAGE_KEY_PREFIX}${JSON.stringify([
+ normalizeModelIdentity(modelId),
+ normalizeGgufVariantIdentity(ggufVariant),
+ ])}`;
+}
+
+export function modelIdFromStorageKey(key: string): string | null {
+ const parsed = parseVersionedModelStorageKey(key);
+ if (parsed) {
+ return parsed.modelId;
+ }
+ const separator = key.lastIndexOf("::");
+ return separator >= 0 ? key.slice(0, separator) : null;
+}
+
+export function ggufVariantFromStorageKey(key: string): string | null {
+ const parsed = parseVersionedModelStorageKey(key);
+ if (parsed) {
+ return parsed.ggufVariant;
+ }
+ const separator = key.lastIndexOf("::");
+ return separator >= 0 ? key.slice(separator + 2) : null;
+}
diff --git a/studio/frontend/src/features/model-picker/model-config/per-model-config.ts b/studio/frontend/src/features/model-picker/model-config/per-model-config.ts
new file mode 100644
index 0000000000..0b03423736
--- /dev/null
+++ b/studio/frontend/src/features/model-picker/model-config/per-model-config.ts
@@ -0,0 +1,665 @@
+// 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 {
+ ggufVariantFromStorageKey,
+ modelIdFromStorageKey,
+ modelStorageKey,
+ normalizeGgufVariantIdentity,
+ normalizeModelIdentity,
+} from "./model-identity";
+
+export interface PerModelConfig {
+ customContextLength: number | null;
+ maxSeqLength: number | null;
+ kvCacheDtype: string | null;
+ speculativeType: string | null;
+ specDraftNMax: number | null;
+ tensorParallel: boolean;
+ chatTemplateOverride: string | null;
+ // GPU Memory controls (per-model, GGUF-only), optional so older blobs still
+ // parse. null selectedGpuIds (all GPUs) is distinct from absent. The --tensor-split
+ // ratio is deliberately not remembered: it is positionally bound to the exact
+ // GPU set/order and unvalidated.
+ gpuMemoryMode?: "auto" | "manual";
+ gpuLayers?: number;
+ nCpuMoe?: number;
+ selectedGpuIds?: number[] | null;
+}
+
+export const DEFAULT_PER_MODEL_CONFIG: PerModelConfig = {
+ customContextLength: null,
+ maxSeqLength: null,
+ kvCacheDtype: null,
+ speculativeType: null,
+ specDraftNMax: null,
+ tensorParallel: false,
+ chatTemplateOverride: null,
+};
+
+export const MAX_SEQ_LENGTH_MIN = 128;
+export const MAX_SEQ_LENGTH_MAX = 1048576;
+export const MAX_SEQ_LENGTH_STEP = 128;
+// App-default max sequence length when a non-GGUF model has no override. Both
+// paths fall back to this rather than an active model's runtime value, so an
+// unconfigured pane never inherits another model's larger context and OOMs.
+export const DEFAULT_MAX_SEQ_LENGTH = 4096;
+export const CONTEXT_LENGTH_MIN = 128;
+
+export const KV_CACHE_DTYPES = ["bf16", "q8_0", "q5_1", "q4_1"] as const;
+const VALID_KV_CACHE_DTYPES = new Set(KV_CACHE_DTYPES);
+
+export const SPECULATIVE_TYPES = [
+ "auto",
+ "mtp",
+ "ngram",
+ "mtp+ngram",
+ "off",
+] as const;
+export const MTP_SPECULATIVE_TYPES: ReadonlySet = new Set([
+ "mtp",
+ "mtp+ngram",
+]);
+
+const STORAGE_KEY = "unsloth_model_configs";
+const LEGACY_STORAGE_KEY = "unsloth_load_settings";
+const LEGACY_MIGRATION_FLAG = "unsloth_model_configs_migrated";
+const STORAGE_SCHEMA_VERSION = 1;
+const MAX_ENTRIES = 500;
+const MAX_PER_MODEL_CONFIG_STORAGE_BYTES = 1024 * 1024;
+export const MAX_CHAT_TEMPLATE_BYTES = 65_536;
+
+type StoredPerModelConfig = PerModelConfig & {
+ version: typeof STORAGE_SCHEMA_VERSION;
+};
+type StoredMap = Record;
+type RawConfig = Partial & { version?: unknown };
+
+const STORED_CONFIG_FIELDS = new Set([
+ "version",
+ "customContextLength",
+ "maxSeqLength",
+ "kvCacheDtype",
+ "speculativeType",
+ "specDraftNMax",
+ "tensorParallel",
+ "chatTemplateOverride",
+ "gpuMemoryMode",
+ "gpuLayers",
+ "nCpuMoe",
+ "selectedGpuIds",
+]);
+
+function normalizeGpuFields(partial: RawConfig): {
+ gpuMemoryMode?: "auto" | "manual";
+ gpuLayers?: number;
+ nCpuMoe?: number;
+ selectedGpuIds?: number[] | null;
+} {
+ const out: {
+ gpuMemoryMode?: "auto" | "manual";
+ gpuLayers?: number;
+ nCpuMoe?: number;
+ selectedGpuIds?: number[] | null;
+ } = {};
+ // Only "manual" is a real override; persisting "auto" would pin the model and
+ // stop it following later changes to the global GPU Memory preference.
+ if (partial.gpuMemoryMode === "manual") {
+ out.gpuMemoryMode = "manual";
+ }
+ if (
+ typeof partial.gpuLayers === "number" &&
+ Number.isFinite(partial.gpuLayers)
+ ) {
+ out.gpuLayers = Math.trunc(partial.gpuLayers);
+ }
+ if (
+ typeof partial.nCpuMoe === "number" &&
+ Number.isFinite(partial.nCpuMoe) &&
+ partial.nCpuMoe >= 0
+ ) {
+ out.nCpuMoe = Math.trunc(partial.nCpuMoe);
+ }
+ if (partial.selectedGpuIds === null) {
+ out.selectedGpuIds = null;
+ } else if (
+ Array.isArray(partial.selectedGpuIds) &&
+ partial.selectedGpuIds.every(
+ (n) => typeof n === "number" && Number.isFinite(n),
+ )
+ ) {
+ out.selectedGpuIds = partial.selectedGpuIds.map((n) => Math.trunc(n));
+ }
+ return out;
+}
+
+function canonicalizeSpeculativeType(value: string): string | null {
+ const s = value.trim().toLowerCase();
+ if (!s) {
+ return null;
+ }
+ // "auto"/"default" is the follow-global sentinel; store as null so it is never
+ // persisted as an override and global speculative-decoding changes keep applying.
+ if (s === "auto" || s === "default") {
+ return null;
+ }
+ if (s === "off") {
+ return "off";
+ }
+ if (s === "mtp" || s === "draft-mtp") {
+ return "mtp";
+ }
+ if (s === "ngram" || s === "ngram-mod" || s === "ngram-simple") {
+ return "ngram";
+ }
+ if (s === "mtp+ngram") {
+ return "mtp+ngram";
+ }
+ return null;
+}
+
+export function normalizeMaxSeqLength(value: unknown): number | null {
+ if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
+ return null;
+ }
+ const snapped = Math.round(value / MAX_SEQ_LENGTH_STEP) * MAX_SEQ_LENGTH_STEP;
+ return Math.max(MAX_SEQ_LENGTH_MIN, Math.min(MAX_SEQ_LENGTH_MAX, snapped));
+}
+
+export function floorMaxSeqLength(value: unknown): number | null {
+ if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) {
+ return null;
+ }
+ const snapped = Math.floor(value / MAX_SEQ_LENGTH_STEP) * MAX_SEQ_LENGTH_STEP;
+ return Math.max(MAX_SEQ_LENGTH_MIN, Math.min(MAX_SEQ_LENGTH_MAX, snapped));
+}
+
+function canUseStorage(): boolean {
+ return typeof window !== "undefined";
+}
+
+function serializedByteLength(value: string): number {
+ return typeof TextEncoder !== "undefined"
+ ? new TextEncoder().encode(value).byteLength
+ : value.length;
+}
+
+export function chatTemplateByteLength(value: string): number {
+ return serializedByteLength(value);
+}
+
+export function isChatTemplateWithinLimit(value: string): boolean {
+ return chatTemplateByteLength(value) <= MAX_CHAT_TEMPLATE_BYTES;
+}
+
+function serializedMapSize(map: StoredMap): number {
+ return serializedByteLength(JSON.stringify(map));
+}
+
+function serializedMapEntrySize(key: string, value: StoredMap[string]): number {
+ return (
+ serializedByteLength(JSON.stringify(key)) +
+ 1 +
+ serializedByteLength(JSON.stringify(value))
+ );
+}
+
+function deleteOldestEvictableEntry(
+ map: StoredMap,
+ protectedKeys?: ReadonlySet,
+): { key: string; value: StoredMap[string] } | null {
+ for (const key of Object.keys(map)) {
+ // Never evict a future-schema entry an older client cannot interpret.
+ if (
+ protectedKeys?.has(key) ||
+ storedConfigVersion(map[key]) > STORAGE_SCHEMA_VERSION
+ ) {
+ continue;
+ }
+ const value = map[key];
+ delete map[key];
+ return { key, value };
+ }
+ return null;
+}
+
+function enforceStorageBudget(
+ map: StoredMap,
+ protectedKeys?: ReadonlySet,
+): boolean {
+ let entryCount = Object.keys(map).length;
+ while (entryCount > MAX_ENTRIES) {
+ if (!deleteOldestEvictableEntry(map, protectedKeys)) {
+ return false;
+ }
+ entryCount -= 1;
+ }
+ let bytes = serializedMapSize(map);
+ while (bytes > MAX_PER_MODEL_CONFIG_STORAGE_BYTES) {
+ const removed = deleteOldestEvictableEntry(map, protectedKeys);
+ if (!removed) {
+ return false;
+ }
+ bytes -=
+ serializedMapEntrySize(removed.key, removed.value) +
+ (entryCount > 1 ? 1 : 0);
+ entryCount -= 1;
+ }
+ return true;
+}
+
+function storedConfigVersion(raw: unknown): number {
+ if (!raw || typeof raw !== "object" || Array.isArray(raw)) {
+ return 0;
+ }
+ const version = (raw as RawConfig).version;
+ return typeof version === "number" && Number.isFinite(version) ? version : 0;
+}
+
+let legacyMigrationChecked = false;
+
+function parseLegacyModelKey(
+ key: string,
+): { modelId: string; ggufVariant: string | null } | null {
+ const separator = key.lastIndexOf("::");
+ if (separator >= 0) {
+ const modelId = key.slice(0, separator);
+ return modelId
+ ? { modelId, ggufVariant: key.slice(separator + 2) || null }
+ : null;
+ }
+ return key ? { modelId: key, ggufVariant: null } : null;
+}
+
+function legacyEntryToConfig(raw: Record): PerModelConfig {
+ return normalizeV1({
+ customContextLength:
+ typeof raw.contextLength === "number" ? raw.contextLength : null,
+ maxSeqLength: null,
+ kvCacheDtype:
+ typeof raw.kvCacheDtype === "string" ? raw.kvCacheDtype : null,
+ speculativeType:
+ typeof raw.speculativeType === "string" ? raw.speculativeType : null,
+ specDraftNMax:
+ typeof raw.specDraftNMax === "number" ? raw.specDraftNMax : null,
+ tensorParallel:
+ typeof raw.tensorParallel === "boolean" ? raw.tensorParallel : false,
+ chatTemplateOverride: null,
+ // Carry legacy GPU Memory knobs; normalizeGpuFields drops anything malformed.
+ gpuMemoryMode:
+ raw.gpuMemoryMode === "auto" || raw.gpuMemoryMode === "manual"
+ ? raw.gpuMemoryMode
+ : undefined,
+ gpuLayers: typeof raw.gpuLayers === "number" ? raw.gpuLayers : undefined,
+ nCpuMoe: typeof raw.nCpuMoe === "number" ? raw.nCpuMoe : undefined,
+ selectedGpuIds:
+ raw.selectedGpuIds === null
+ ? null
+ : Array.isArray(raw.selectedGpuIds)
+ ? (raw.selectedGpuIds as number[])
+ : undefined,
+ });
+}
+
+function mergeLegacyEntries(
+ map: StoredMap,
+ legacy: Record,
+): string[] {
+ const addedKeys: string[] = [];
+ for (const [legacyKey, value] of Object.entries(legacy)) {
+ if (!value || typeof value !== "object") {
+ continue;
+ }
+ const parsedKey = parseLegacyModelKey(legacyKey);
+ if (!parsedKey) {
+ continue;
+ }
+ const migrated = legacyEntryToConfig(value as Record);
+ const key = modelStorageKey(parsedKey.modelId, parsedKey.ggufVariant);
+ if (isDefaultConfig(migrated) || Object.hasOwn(map, key)) {
+ continue;
+ }
+ map[key] = toStoredConfig(migrated);
+ addedKeys.push(key);
+ }
+ return addedKeys;
+}
+
+function migrateLegacyLoadSettingsOnce(): void {
+ if (legacyMigrationChecked || !canUseStorage()) {
+ return;
+ }
+ legacyMigrationChecked = true;
+ try {
+ if (localStorage.getItem(LEGACY_MIGRATION_FLAG)) {
+ return;
+ }
+ let legacy: unknown = null;
+ try {
+ legacy = JSON.parse(localStorage.getItem(LEGACY_STORAGE_KEY) ?? "null");
+ } catch {
+ legacy = null;
+ }
+ if (!legacy || typeof legacy !== "object" || Array.isArray(legacy)) {
+ localStorage.setItem(LEGACY_MIGRATION_FLAG, "1");
+ return;
+ }
+ const map = readMapRaw();
+ // Snapshot existing entries so eviction can protect them: importing old load
+ // settings must never discard a newer per-model config the user already has.
+ const existingKeys = new Set(Object.keys(map));
+ const migratedKeys = mergeLegacyEntries(
+ map,
+ legacy as Record,
+ );
+ if (migratedKeys.length === 0) {
+ localStorage.setItem(LEGACY_MIGRATION_FLAG, "1");
+ return;
+ }
+ // Protect pre-existing entries so only just-migrated legacy entries are
+ // dropped when over budget.
+ if (!enforceStorageBudget(map, existingKeys)) {
+ return;
+ }
+ if (writeMap(map)) {
+ localStorage.setItem(LEGACY_MIGRATION_FLAG, "1");
+ }
+ } catch (err) {
+ console.warn("Failed to migrate legacy load settings:", err);
+ }
+}
+
+function readMapRaw(): StoredMap {
+ if (!canUseStorage()) {
+ return {};
+ }
+ try {
+ const raw = localStorage.getItem(STORAGE_KEY);
+ if (!raw) {
+ return {};
+ }
+ const parsed = JSON.parse(raw);
+ if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
+ return {};
+ }
+ return parsed as StoredMap;
+ } catch {
+ return {};
+ }
+}
+
+function readMap(): StoredMap {
+ migrateLegacyLoadSettingsOnce();
+ return readMapRaw();
+}
+
+function writeMap(map: StoredMap): boolean {
+ if (!canUseStorage()) {
+ return false;
+ }
+ try {
+ localStorage.setItem(STORAGE_KEY, JSON.stringify(map));
+ return true;
+ } catch (err) {
+ console.warn("Failed to persist per-model config:", err);
+ return false;
+ }
+}
+
+function warnDroppedFields(raw: Record, version: number): void {
+ if (!import.meta.env?.DEV) {
+ return;
+ }
+ const dropped = Object.keys(raw).filter(
+ (key) => !STORED_CONFIG_FIELDS.has(key),
+ );
+ if (dropped.length > 0) {
+ console.warn("Dropped unknown per-model config fields:", dropped);
+ }
+ if (version > STORAGE_SCHEMA_VERSION) {
+ console.warn("Per-model config schema is newer than this app:", version);
+ }
+}
+
+function normalizeV1(partial: RawConfig): PerModelConfig {
+ const rawSpecType =
+ typeof partial.speculativeType === "string"
+ ? canonicalizeSpeculativeType(partial.speculativeType)
+ : null;
+ const speculativeType = rawSpecType ?? DEFAULT_PER_MODEL_CONFIG.speculativeType;
+ const specDraftNMax =
+ speculativeType != null &&
+ MTP_SPECULATIVE_TYPES.has(speculativeType) &&
+ typeof partial.specDraftNMax === "number" &&
+ Number.isFinite(partial.specDraftNMax)
+ ? Math.max(1, Math.min(16, Math.round(partial.specDraftNMax)))
+ : null;
+ return {
+ customContextLength:
+ typeof partial.customContextLength === "number" &&
+ Number.isFinite(partial.customContextLength) &&
+ partial.customContextLength > 0
+ ? Math.max(CONTEXT_LENGTH_MIN, Math.floor(partial.customContextLength))
+ : null,
+ maxSeqLength: normalizeMaxSeqLength(partial.maxSeqLength),
+ kvCacheDtype:
+ typeof partial.kvCacheDtype === "string" &&
+ VALID_KV_CACHE_DTYPES.has(partial.kvCacheDtype)
+ ? partial.kvCacheDtype
+ : null,
+ speculativeType,
+ specDraftNMax,
+ tensorParallel:
+ typeof partial.tensorParallel === "boolean"
+ ? partial.tensorParallel
+ : DEFAULT_PER_MODEL_CONFIG.tensorParallel,
+ chatTemplateOverride:
+ typeof partial.chatTemplateOverride === "string" &&
+ isChatTemplateWithinLimit(partial.chatTemplateOverride)
+ ? partial.chatTemplateOverride
+ : null,
+ ...normalizeGpuFields(partial),
+ };
+}
+
+function normalize(raw: unknown): PerModelConfig {
+ if (!raw || typeof raw !== "object" || Array.isArray(raw)) {
+ return normalizeV1({});
+ }
+ const partial = raw as RawConfig;
+ const version =
+ typeof partial.version === "number" && Number.isFinite(partial.version)
+ ? partial.version
+ : 0;
+ warnDroppedFields(raw as Record, version);
+ return normalizeV1(partial);
+}
+
+function toStoredConfig(config: PerModelConfig): StoredPerModelConfig {
+ return {
+ version: STORAGE_SCHEMA_VERSION,
+ ...normalize(config),
+ };
+}
+
+function legacyModelStorageKey(
+ modelId: string,
+ ggufVariant?: string | null,
+): string {
+ return `${modelId}::${ggufVariant ?? ""}`;
+}
+
+function storageKeysForModelVariant(
+ modelId: string,
+ ggufVariant?: string | null,
+): string[] {
+ const key = modelStorageKey(modelId, ggufVariant);
+ const legacyKey = legacyModelStorageKey(modelId, ggufVariant);
+ return key === legacyKey ? [key] : [key, legacyKey];
+}
+
+function configKeyMatchesModelVariant(
+ key: string,
+ modelId: string,
+ ggufVariant?: string | null,
+): boolean {
+ const storedModelId = modelIdFromStorageKey(key);
+ if (!storedModelId) {
+ return false;
+ }
+ return (
+ normalizeModelIdentity(storedModelId) === normalizeModelIdentity(modelId) &&
+ normalizeGgufVariantIdentity(ggufVariantFromStorageKey(key)) ===
+ normalizeGgufVariantIdentity(ggufVariant)
+ );
+}
+
+function findConfigKeyForModelVariant(
+ map: StoredMap,
+ modelId: string,
+ ggufVariant?: string | null,
+): string | null {
+ for (const key of storageKeysForModelVariant(modelId, ggufVariant)) {
+ if (Object.hasOwn(map, key)) {
+ return key;
+ }
+ }
+ for (const key of Object.keys(map)) {
+ if (configKeyMatchesModelVariant(key, modelId, ggufVariant)) {
+ return key;
+ }
+ }
+ return null;
+}
+
+function hasFutureConfigForModelVariant(
+ map: StoredMap,
+ modelId: string,
+ ggufVariant?: string | null,
+): boolean {
+ for (const key of Object.keys(map)) {
+ if (
+ configKeyMatchesModelVariant(key, modelId, ggufVariant) &&
+ storedConfigVersion(map[key]) > STORAGE_SCHEMA_VERSION
+ ) {
+ return true;
+ }
+ }
+ return false;
+}
+
+function deleteConfigEntriesForModelVariant(
+ map: StoredMap,
+ modelId: string,
+ ggufVariant?: string | null,
+): boolean {
+ let changed = false;
+ for (const key of Object.keys(map)) {
+ if (!configKeyMatchesModelVariant(key, modelId, ggufVariant)) {
+ continue;
+ }
+ delete map[key];
+ changed = true;
+ }
+ return changed;
+}
+
+function loadPerModelConfig(
+ modelId: string,
+ ggufVariant?: string | null,
+): PerModelConfig | null {
+ const map = readMap();
+ const key = findConfigKeyForModelVariant(map, modelId, ggufVariant);
+ if (!key) {
+ return null;
+ }
+ // Never apply a future-schema record an older client cannot interpret.
+ if (storedConfigVersion(map[key]) > STORAGE_SCHEMA_VERSION) {
+ return null;
+ }
+ return normalize(map[key]);
+}
+
+export function isDefaultConfig(config: PerModelConfig): boolean {
+ return (
+ config.customContextLength == null &&
+ config.maxSeqLength == null &&
+ (config.kvCacheDtype ?? null) === DEFAULT_PER_MODEL_CONFIG.kvCacheDtype &&
+ config.speculativeType === DEFAULT_PER_MODEL_CONFIG.speculativeType &&
+ config.specDraftNMax == null &&
+ Boolean(config.tensorParallel) ===
+ Boolean(DEFAULT_PER_MODEL_CONFIG.tensorParallel) &&
+ (config.chatTemplateOverride ?? null) === null &&
+ gpuFieldsAtDefault(config)
+ );
+}
+
+// GPU knobs are "default" when mode is Auto with no explicit choice: mode
+// auto/absent, gpuLayers < 0/absent, nCpuMoe 0/absent, selectedGpuIds null/absent.
+function gpuFieldsAtDefault(config: PerModelConfig): boolean {
+ return (
+ (config.gpuMemoryMode ?? "auto") === "auto" &&
+ (config.gpuLayers == null || config.gpuLayers < 0) &&
+ (config.nCpuMoe == null || config.nCpuMoe === 0) &&
+ config.selectedGpuIds == null
+ );
+}
+
+export function savePerModelConfig(
+ modelId: string,
+ ggufVariant: string | null | undefined,
+ config: PerModelConfig,
+): boolean {
+ if (
+ typeof config.chatTemplateOverride === "string" &&
+ !isChatTemplateWithinLimit(config.chatTemplateOverride)
+ ) {
+ return false;
+ }
+ const normalized = normalize(config);
+ const map = readMap();
+ if (hasFutureConfigForModelVariant(map, modelId, ggufVariant)) {
+ return false;
+ }
+ if (isDefaultConfig(normalized)) {
+ const changed = deleteConfigEntriesForModelVariant(
+ map,
+ modelId,
+ ggufVariant,
+ );
+ return changed ? writeMap(map) : true;
+ }
+ const [key] = storageKeysForModelVariant(modelId, ggufVariant);
+ deleteConfigEntriesForModelVariant(map, modelId, ggufVariant);
+ map[key] = toStoredConfig(normalized);
+ if (!enforceStorageBudget(map, new Set([key]))) {
+ return false;
+ }
+ return writeMap(map);
+}
+
+export function deletePerModelConfig(
+ modelId: string,
+ ggufVariant?: string | null,
+): boolean {
+ const map = readMap();
+ // Mirror savePerModelConfig: never let an older client destroy a future-schema entry.
+ if (hasFutureConfigForModelVariant(map, modelId, ggufVariant)) {
+ return false;
+ }
+ if (!deleteConfigEntriesForModelVariant(map, modelId, ggufVariant)) {
+ return true;
+ }
+ return writeMap(map);
+}
+
+export function resolveInitialConfig(
+ modelId: string,
+ ggufVariant?: string | null,
+): { config: PerModelConfig; remembered: boolean } {
+ const saved = loadPerModelConfig(modelId, ggufVariant);
+ if (saved) {
+ return { config: saved, remembered: true };
+ }
+ return { config: { ...DEFAULT_PER_MODEL_CONFIG }, remembered: false };
+}
diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
index f7f3bccad6..a6fcf87c57 100644
--- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
@@ -16,7 +16,6 @@ import {
Folder01Icon,
McpServerIcon,
PencilRulerIcon,
- Settings02Icon,
ShieldBanIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
@@ -148,10 +147,6 @@ export function ChatTab() {
const hydratePersistedSettings = useChatRuntimeStore(
(state) => state.hydratePersistedSettings,
);
- const loadOnSelection = useChatRuntimeStore((state) => state.loadOnSelection);
- const setLoadOnSelection = useChatRuntimeStore(
- (state) => state.setLoadOnSelection,
- );
const expandQuantizations = useChatRuntimeStore(
(state) => state.expandQuantizations,
);
@@ -193,40 +188,6 @@ export function ChatTab() {
-
- On: Unsloth auto-picks the best settings and loads it.
-
- Off: opens Run settings to customize, then load.
-
- The gear always opens Run settings:{" "}
-
-
- Q4_K_M
-
-
- downloaded
-
- 16 GB
-
-
-
-
-
- }
- >
-
-
{
const encoded = encodeURIComponent(modelName);
- const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : "";
- const response = await authFetch(`/api/models/check-vision/${encoded}${query}`);
+ const response = await authFetch(`/api/models/check-vision/${encoded}`, {
+ headers: hubTokenHeader(hfToken?.trim() || null),
+ });
if (!response.ok) {
// If the check fails (e.g. network error), default to non-vision
return false;
@@ -114,8 +116,9 @@ export async function checkEmbeddingModel(
hfToken?: string | null,
): Promise {
const encoded = encodeURIComponent(modelName);
- const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : "";
- const response = await authFetch(`/api/models/check-embedding/${encoded}${query}`);
+ const response = await authFetch(`/api/models/check-embedding/${encoded}`, {
+ headers: hubTokenHeader(hfToken?.trim() || null),
+ });
if (!response.ok) {
// If the check fails (e.g. network error), default to non-embedding
return false;
@@ -130,8 +133,10 @@ export async function getModelConfig(
hfToken?: string,
): Promise {
const encoded = encodeURIComponent(modelName);
- const params = hfToken ? `?hf_token=${encodeURIComponent(hfToken)}` : "";
- const response = await authFetch(`/api/models/config/${encoded}${params}`, { signal });
+ const response = await authFetch(`/api/models/config/${encoded}`, {
+ headers: hubTokenHeader(hfToken?.trim() || null),
+ signal,
+ });
if (!response.ok) {
throw new Error(`Failed to fetch model config (${response.status})`);
}
diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts
index a0d249ff1b..81514b4415 100644
--- a/studio/frontend/src/features/training/index.ts
+++ b/studio/frontend/src/features/training/index.ts
@@ -25,8 +25,8 @@ export { HfDatasetSubsetSplitSelectors } from "./components/hf-dataset-subset-sp
export { useDatasetPreviewDialogStore } from "./stores/dataset-preview-dialog-store";
export { listLocalDatasets, uploadTrainingDataset } from "./api/datasets-api";
export type { LocalDatasetInfo } from "./types/datasets";
-export { listLocalModels } from "./api/models-api";
-export type { LocalModelInfo } from "./api/models-api";
+export { getModelConfig, listLocalModels } from "./api/models-api";
+export type { LocalModelInfo, ModelConfigResponse } from "./api/models-api";
export type {
TrainingPhase,
TrainingViewData,
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index b94578369d..f7dcd79d10 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -1557,7 +1557,7 @@ class TestWorkerRocmMambaSsm:
assert "getattr(torch.version, 'hip', None)" in source
def test_direct_wheel_url_returns_none_without_cuda_major(self, monkeypatch):
- """_direct_wheel_url should return None when cuda_major is empty (ROCm)."""
+ """direct_wheel_url should return None when cuda_major is empty (ROCm)."""
_worker_spec = importlib.util.spec_from_file_location("test_worker", _WORKER_PATH)
assert _worker_spec is not None and _worker_spec.loader is not None
worker_mod = importlib.util.module_from_spec(_worker_spec)
@@ -1583,7 +1583,7 @@ class TestWorkerRocmMambaSsm:
"hip_version": "7.1.12345",
"cxx11abi": "TRUE",
}
- result = worker_mod._direct_wheel_url(
+ result = worker_mod.direct_wheel_url(
filename_prefix = "causal_conv1d",
package_version = "1.6.1",
release_tag = "v1.6.1.post4",
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index b00b45f97a..4d13889878 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.py
@@ -821,10 +821,16 @@ with sync_playwright() as p:
last_assistant = page.locator('[data-role="assistant"]').last
last_assistant.hover()
page.wait_for_timeout(400)
- regen_btn = page.get_by_role(
- "button",
- name = re.compile(r"(reload|regenerate)", re.I),
- ).first
+ # Exclude disabled controls: the picker's new disabled "Reload model"
+ # button also matches and sorts first, so .first would target it.
+ regen_btn = (
+ page.get_by_role(
+ "button",
+ name = re.compile(r"(reload|regenerate)", re.I),
+ )
+ .and_(page.locator("button:not([disabled])"))
+ .first
+ )
if regen_btn.count() > 0:
regen_btn.click()
try:
diff --git a/tests/studio/playwright_model_config.py b/tests/studio/playwright_model_config.py
new file mode 100644
index 0000000000..a8d143a253
--- /dev/null
+++ b/tests/studio/playwright_model_config.py
@@ -0,0 +1,740 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Model-picker per-model-config Playwright regression test (GPU-free, CPU gemma).
+
+Guards, end to end against the real frontend, the exact regressions that got the
+predecessor PR reverted:
+
+ - Context Length persists: set a distinctive per-model Context Length + tick
+ "Remember for this model" + Load; the value reaches the /api/inference/load
+ request (max_seq_length) AND lands in localStorage (unsloth_model_configs),
+ and survives a full browser reload (HARD).
+ - Reset clears: after customizing, Reset must clear the stored override, never
+ pin the context to a fixed number (the "Reset pins context" regression) (HARD).
+ - Hidden infra models absent: the RAG embedder (bge-small-en-v1.5) and the
+ llama.cpp validation probe (stories260K) never appear in the picker. The
+ probe GGUF is primed into the HF cache by the CI job, so "absent" proves
+ "hidden", not "not downloaded" (HARD).
+ - Legacy migration is idempotent: a pre-feature unsloth_load_settings store
+ migrates once into the versioned unsloth_model_configs map with the value
+ preserved, and a second reload with a fresh legacy seed present does not
+ re-migrate, duplicate, or clobber (gates under STUDIO_UI_STRICT via soft_fail).
+ - Advanced settings persist: KV cache dtype / tensor-parallel toggled under
+ Advanced + Remember land in unsloth_model_configs (best-effort).
+
+Runs as a plain script (not via pytest), mirroring tests/studio/playwright_extra_ui.py:
+accumulate failures in `_failed`, exit non-zero if any HARD gate failed. With
+STUDIO_UI_STRICT=1 (as CI sets), soft_fail also gates; genuinely-optional checks
+use runtime_warn so they never flake the merge gate.
+"""
+
+import json
+import re
+import sys
+import os
+import time
+from pathlib import Path
+
+from playwright.sync_api import sync_playwright
+
+# Run as a plain script (not via pytest), so prepend the dir to sys.path.
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from _playwright_robust import ( # noqa: E402
+ chromium_launch_args,
+ click_and_wait_for_response,
+ evaluate_fetch,
+ install_view_transition_killer,
+ install_wall_clock_watchdog,
+ is_benign_page_error,
+ recover_or_replace_page,
+ robust_evaluate,
+ wait_for_health,
+)
+
+BASE = os.environ["BASE_URL"]
+NEW = os.environ.get("STUDIO_NEW_PW", "ModelCfg-NEW-2026!")
+# Attach mode: log into an already-provisioned Studio with an existing password
+# instead of the first-boot change-password dance. CI leaves STUDIO_LOGIN_PW unset
+# to exercise the real change-password flow; local runs can set it to skip re-provisioning.
+LOGIN_PW = os.environ.get("STUDIO_LOGIN_PW")
+LOGIN_USER = os.environ.get("STUDIO_LOGIN_USER", "unsloth")
+GGUF_REPO = os.environ.get("GGUF_REPO", "unsloth/gemma-3-270m-it-GGUF")
+GGUF_VARIANT = os.environ.get("GGUF_VARIANT", "UD-Q4_K_XL")
+# Substring of the On Device picker row for the loaded model.
+MODEL_HINT = os.environ.get("STUDIO_MODEL_HINT", "gemma-3-270m")
+# A distinctive valid (>=128, multiple of 128, below the model's 32768 ceiling)
+# Context Length, clearly not a default, so persistence is unambiguous.
+DISTINCT_CTX = int(os.environ.get("STUDIO_DISTINCT_CTX", "4096"))
+ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_modelcfg")
+ART = Path(ART_DIR)
+ART.mkdir(parents = True, exist_ok = True)
+STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
+TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000"))
+WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720"))
+FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000"))
+LOAD_FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_LOAD_TIMEOUT_MS", "180000"))
+
+_n = [0]
+_failed: list[str] = []
+
+
+def step(s: str) -> None:
+ print(f"[ui-modelcfg] STEP {s}", flush = True)
+
+
+def info(s: str) -> None:
+ print(f"[ui-modelcfg] {s}", flush = True)
+
+
+def fail(m: str) -> None:
+ print(f"[ui-modelcfg] FAIL: {m}", flush = True)
+ _failed.append(m)
+
+
+def soft_fail(m: str) -> None:
+ if STRICT:
+ fail(m)
+ else:
+ info(f"WARN (strict-off): {m}")
+
+
+def runtime_warn(m: str) -> None:
+ """Warn about a genuinely-optional check that STRICT does not gate."""
+ info(f"WARN (runtime): {m}")
+
+
+def _count(loc) -> int:
+ try:
+ return loc.count()
+ except Exception:
+ return 0
+
+
+def _as_int(value) -> int | None:
+ """Parse an input value to int, tolerating commas/whitespace. Comparisons
+ must be numeric, never substring: '40960' (a model's native default) would
+ spuriously "contain" '4096'."""
+ if value is None:
+ return None
+ try:
+ return int(str(value).replace(",", "").strip())
+ except Exception:
+ return None
+
+
+def _login_token_via_api(base: str, user: str, pw: str) -> str:
+ """POST /api/auth/login -> access_token (attach-mode helper, stdlib only)."""
+ import urllib.request
+
+ req = urllib.request.Request(
+ f"{base}/api/auth/login",
+ data = json.dumps({"username": user, "password": pw}).encode(),
+ headers = {"Content-Type": "application/json"},
+ method = "POST",
+ )
+ with urllib.request.urlopen(req, timeout = 15) as r:
+ return json.loads(r.read().decode())["access_token"]
+
+
+with sync_playwright() as p:
+ _watchdog = install_wall_clock_watchdog(
+ WALL_TIMEOUT_S,
+ label = "ui-modelcfg",
+ info = info,
+ )
+ # Health pre-flight: bash-side health wait can pass before the auth DB migrates.
+ wait_for_health(BASE, timeout = 30.0, info = info)
+ browser = p.chromium.launch(
+ headless = True,
+ args = chromium_launch_args(),
+ )
+ ctx = browser.new_context(
+ viewport = {"width": 1280, "height": 900},
+ reduced_motion = "reduce",
+ )
+ install_view_transition_killer(ctx)
+ page = ctx.new_page()
+ page.set_default_timeout(60_000)
+ page_errors = []
+
+ def _on_pageerror(e):
+ msg = str(e)
+ if is_benign_page_error(msg):
+ info(f"WARN ignoring benign pageerror: {msg!r}")
+ return
+ page_errors.append(msg)
+
+ page.on("pageerror", _on_pageerror)
+
+ # Record every /api/inference/load POST payload so the persistence gate can
+ # assert max_seq_length.
+ load_posts: list[str] = []
+
+ def _on_request(req):
+ try:
+ if req.method == "POST" and "/api/inference/load" in req.url:
+ load_posts.append(req.post_data or "")
+ except Exception:
+ pass
+
+ page.on("request", _on_request)
+
+ def shoot(name: str) -> None:
+ _n[0] += 1
+ try:
+ page.screenshot(
+ path = str(ART / f"{_n[0]:02d}-{name}.png"),
+ full_page = True,
+ timeout = 90_000,
+ animations = "disabled",
+ )
+ except Exception as _shoot_err:
+ info(f"WARN: screenshot {name} failed: {_shoot_err}")
+
+ def read_configs() -> dict:
+ """Return the parsed unsloth_model_configs map (or {} if absent/invalid)."""
+ raw = robust_evaluate(page, "() => localStorage.getItem('unsloth_model_configs')")
+ if not raw:
+ return {}
+ try:
+ data = json.loads(raw)
+ return data if isinstance(data, dict) else {}
+ except Exception:
+ return {}
+
+ def config_entries(cfg: dict) -> list[dict]:
+ """The per-model entries (dict values) of the stored map, schema-tolerant."""
+ return [v for v in cfg.values() if isinstance(v, dict)]
+
+ # ─────────────────────────────────────────────────────
+ # Setup: authenticate + model load.
+ # ─────────────────────────────────────────────────────
+ if LOGIN_PW:
+ # Attach mode: log in via the API and seed the token before navigation,
+ # skipping the first-boot change-password dance.
+ step("setup: API login + token seed (attach to running Studio)")
+ _tok = _login_token_via_api(BASE, LOGIN_USER, LOGIN_PW)
+ ctx.add_init_script(
+ f"try{{localStorage.setItem('unsloth_auth_token', {json.dumps(_tok)});}}"
+ f"catch(e){{}}"
+ )
+ page.goto(BASE, wait_until = "domcontentloaded", timeout = 60_000)
+ else:
+ step("setup: change-password")
+ # 3-attempt retry: the form can re-render mid-fill on slow runners and
+ # detach the password fields; each retry re-navigates with a fresh page.
+ form_err: Exception | None = None
+ for _form_attempt in range(3):
+ try:
+ page.goto(f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000)
+ try:
+ page.wait_for_load_state("networkidle", timeout = 30_000)
+ except Exception:
+ pass
+ pw_field = page.locator("#new-password")
+ pw_field.wait_for(state = "visible", timeout = 60_000)
+ pw_field.fill(NEW, timeout = 60_000)
+ page.fill("#confirm-password", NEW, timeout = 60_000)
+ status, _ = click_and_wait_for_response(
+ page,
+ url_substr = "/api/auth/change-password",
+ method = "POST",
+ do_click = lambda: page.locator('button[type="submit"]').click(),
+ timeout_ms = 30_000,
+ info = lambda m: print(f"[ui-modelcfg] {m}", flush = True),
+ )
+ if status is not None and status >= 400:
+ raise AssertionError(
+ f"change-password POST returned {status}; page_errors={page_errors[:1]!r}"
+ )
+ form_err = None
+ break
+ except Exception as e:
+ form_err = e
+ try:
+ cur_url = page.url
+ except Exception:
+ cur_url = ""
+ print(
+ f"[ui-modelcfg] change-password attempt {_form_attempt + 1} failed: "
+ f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; "
+ f"page_errors={len(page_errors)}",
+ flush = True,
+ )
+ if _form_attempt < 2:
+ if "ERR_NO_BUFFER_SPACE" in str(e):
+ backoff_s = 5 if _form_attempt == 0 else 15
+ time.sleep(backoff_s)
+ page = recover_or_replace_page(
+ page,
+ ctx,
+ default_timeout_ms = 60_000,
+ info = lambda m: print(f"[ui-modelcfg] recovery: {m}", flush = True),
+ )
+ page.on("request", _on_request)
+ if form_err is not None:
+ raise form_err
+
+ try:
+ page.wait_for_load_state("networkidle", timeout = 30_000)
+ except Exception:
+ pass
+ composer = page.locator('textarea[aria-label="Message input"]')
+ last_err: Exception | None = None
+ for _attempt in range(2):
+ try:
+ composer.wait_for(state = "visible", timeout = 60_000)
+ last_err = None
+ break
+ except Exception as e:
+ last_err = e
+ try:
+ shoot(f"00-composer-wait-attempt-{_attempt + 1}-fail")
+ except Exception:
+ pass
+ if _attempt == 0:
+ page = recover_or_replace_page(
+ page,
+ ctx,
+ default_timeout_ms = 60_000,
+ goto_url = BASE,
+ settle_networkidle = True,
+ info = lambda m: print(f"[ui-modelcfg] recovery: {m}", flush = True),
+ )
+ page.on("request", _on_request)
+ composer = page.locator('textarea[aria-label="Message input"]')
+ if last_err is not None:
+ raise last_err
+ shoot("01-chat-loaded")
+
+ token = robust_evaluate(page, "() => localStorage.getItem('unsloth_auth_token')")
+ if not token:
+ fail("no access token after auth setup")
+ sys.exit(1)
+
+ # Load the tiny GGUF so it is a live "On Device" model in the picker.
+ load_resp = evaluate_fetch(
+ page,
+ f"{BASE}/api/inference/load",
+ method = "POST",
+ headers = {
+ "Authorization": f"Bearer {token}",
+ "Content-Type": "application/json",
+ },
+ body = {
+ "model_path": GGUF_REPO,
+ "gguf_variant": GGUF_VARIANT,
+ "is_lora": False,
+ "max_seq_length": 2048,
+ },
+ timeout_ms = LOAD_FETCH_TIMEOUT_MS,
+ )
+ if load_resp.get("error"):
+ fail(f"/api/inference/load wedged: {load_resp['error']!r}")
+ sys.exit(1)
+ if load_resp["status"] != 200:
+ fail(f"/api/inference/load -> {load_resp['status']}: {load_resp.get('body')!r}")
+ sys.exit(1)
+ info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}")
+ page.reload()
+ composer = page.locator('textarea[aria-label="Message input"]')
+ composer.wait_for(state = "visible", timeout = 60_000)
+ load_posts.clear() # drop the setup load; keep only UI-driven loads below.
+
+ # ─────────────────────────────────────────────────────
+ # Picker helpers (proven selectors).
+ # ─────────────────────────────────────────────────────
+ POPOVER = '[data-tour="chat-model-selector-popover"]'
+ TRIGGER = '[data-tour="chat-model-selector"]'
+
+ def open_picker():
+ popover = page.locator(POPOVER).first
+ if _count(popover) == 0 or not popover.is_visible():
+ page.locator(TRIGGER).first.click()
+ page.wait_for_timeout(900)
+ popover = page.locator(POPOVER).first
+ popover.wait_for(state = "visible", timeout = 30_000)
+ return popover
+
+ def close_picker():
+ try:
+ page.keyboard.press("Escape")
+ page.wait_for_timeout(400)
+ except Exception:
+ pass
+
+ def select_on_device_row(popover, hint):
+ od = page.get_by_role("tab", name = "On Device").first
+ if _count(od):
+ od.click()
+ page.wait_for_timeout(700)
+ row = popover.locator("[data-model-picker-option]", has_text = hint).first
+ if _count(row) == 0:
+ # Fall back to search filtering.
+ search = popover.locator("[data-model-picker-search-input]").first
+ if _count(search):
+ search.click()
+ search.fill(hint)
+ page.wait_for_timeout(700)
+ row = popover.locator("[data-model-picker-option]", has_text = hint).first
+ if _count(row) == 0:
+ return None
+ row.click()
+ page.wait_for_timeout(800)
+ return row
+
+ def open_config(popover, hint):
+ if select_on_device_row(popover, hint) is None:
+ return None
+ gear = popover.locator('button[aria-label^="Inference settings for"]').first
+ if _count(gear) == 0:
+ return None
+ gear.click()
+ page.wait_for_timeout(800)
+ return popover
+
+ def context_input(popover):
+ for role in ("textbox", "spinbutton"):
+ loc = popover.get_by_role(role, name = "Context Length").first
+ if _count(loc):
+ return loc
+ loc = popover.locator('input[aria-label="Context Length"]').first
+ return loc if _count(loc) else None
+
+ def primary_button(popover):
+ for name in ("Load model", "Reload model", "Save settings", "Forget settings"):
+ b = popover.get_by_role("button", name = name).first
+ if _count(b):
+ return b
+ return None
+
+ # ─────────────────────────────────────────────────────
+ # 1. Hidden infra models absent from the picker (HARD).
+ # ─────────────────────────────────────────────────────
+ step("hidden infra models absent from picker")
+ popover = open_picker()
+ shoot("02-picker-open")
+ needles = ["bge-small-en-v1.5", "stories260"]
+ tabs = ["Recommended", "On Device", "Connected"]
+ hidden_ok = True
+ for needle in needles:
+ for tab_name in tabs:
+ tab = page.get_by_role("tab", name = tab_name).first
+ if _count(tab) == 0:
+ continue
+ try:
+ tab.click()
+ page.wait_for_timeout(400)
+ except Exception:
+ continue
+ search = popover.locator("[data-model-picker-search-input]").first
+ if _count(search):
+ search.click()
+ search.fill(needle)
+ page.wait_for_timeout(600)
+ hit = popover.locator(
+ "[data-model-picker-option]",
+ has_text = re.compile(re.escape(needle), re.I),
+ )
+ c = _count(hit)
+ if c > 0:
+ hidden_ok = False
+ fail(f"infra model {needle!r} visible in picker '{tab_name}' tab ({c} rows)")
+ if _count(search):
+ search.fill("")
+ page.wait_for_timeout(300)
+ if hidden_ok:
+ info("OK hidden: bge-small-en-v1.5 + stories260K absent from every picker tab")
+ shoot("03-hidden-check")
+ close_picker()
+
+ # ─────────────────────────────────────────────────────
+ # 2. Context Length persists (load + request + reload) (HARD).
+ # ─────────────────────────────────────────────────────
+ step(f"context length {DISTINCT_CTX} persists")
+ popover = open_picker()
+ if open_config(popover, MODEL_HINT) is None:
+ fail(f"could not open run-settings for a model matching {MODEL_HINT!r}")
+ else:
+ shoot("04-config-open")
+ ctx_in = context_input(popover)
+ if ctx_in is None:
+ fail("Context Length input not found in run-settings")
+ else:
+ default_ctx = ctx_in.input_value()
+ info(f"default Context Length shown: {default_ctx!r}")
+ ctx_in.click()
+ ctx_in.fill(str(DISTINCT_CTX))
+ page.wait_for_timeout(300)
+ page.keyboard.press("Tab") # blur to commit
+ page.wait_for_timeout(300)
+ remember = popover.get_by_label("Remember for this model").first
+ if _count(remember):
+ try:
+ remember.check()
+ except Exception:
+ remember.click()
+ else:
+ fail("'Remember for this model' checkbox not found")
+ page.wait_for_timeout(300)
+ shoot("05-ctx-set")
+ btn = primary_button(popover)
+ if btn is None:
+ fail("primary Load/Save button not found in run-settings")
+ else:
+ btn.click()
+ page.wait_for_timeout(2500)
+ shoot("06-after-load")
+
+ # (a) localStorage stored the distinctive context.
+ cfg = read_configs()
+ entries = config_entries(cfg)
+ got_ls = any(e.get("customContextLength") == DISTINCT_CTX for e in entries)
+ if got_ls:
+ info(f"OK persist(localStorage): customContextLength={DISTINCT_CTX} stored")
+ else:
+ fail(
+ "context not stored in unsloth_model_configs "
+ f"(entries={json.dumps(entries)[:400]})"
+ )
+
+ # (b) the load request carried max_seq_length == distinctive value.
+ got_req = False
+ for body in load_posts:
+ try:
+ payload = json.loads(body) if body else {}
+ except Exception:
+ payload = {}
+ if payload.get("max_seq_length") == DISTINCT_CTX:
+ got_req = True
+ break
+ if got_req:
+ info(f"OK persist(request): /api/inference/load max_seq_length={DISTINCT_CTX}")
+ else:
+ # The UI may debounce the load; localStorage is the primary
+ # proof, so only warn if the request was missed.
+ runtime_warn(
+ "no /api/inference/load carried "
+ f"max_seq_length={DISTINCT_CTX}; posts={load_posts!r}"
+ )
+
+ # (c) survives a full browser reload.
+ close_picker()
+ page.reload()
+ composer = page.locator('textarea[aria-label="Message input"]')
+ composer.wait_for(state = "visible", timeout = 60_000)
+ popover = open_picker()
+ if open_config(popover, MODEL_HINT) is None:
+ fail("could not reopen run-settings after reload")
+ else:
+ ctx_in = context_input(popover)
+ val = ctx_in.input_value() if ctx_in else None
+ if _as_int(val) == DISTINCT_CTX:
+ info(f"OK persist(reload): Context Length still {val!r} after reload")
+ else:
+ fail(f"Context Length did not persist across reload (got {val!r})")
+ shoot("07-after-reload")
+
+ # ─────────────────────────────────────────────────────
+ # 3. Reset clears the override (never pins context) (HARD).
+ # ─────────────────────────────────────────────────────
+ step("reset clears the per-model override")
+ # (popover + config still open from the reload check.)
+ reset_btn = popover.get_by_role("button", name = "Reset").first
+ if _count(reset_btn) == 0:
+ fail("Reset button not found in run-settings")
+ else:
+ try:
+ reset_btn.click()
+ page.wait_for_timeout(500)
+ except Exception as e:
+ fail(f"Reset click failed: {e}")
+ # The input after Reset is informational only: a live-loaded model can still
+ # echo its context even with the stored override gone. The regression we
+ # guard ("Reset PINS the override") lives in localStorage, asserted below.
+ ctx_in = context_input(popover)
+ after_reset = ctx_in.input_value() if ctx_in else None
+ info(f"reset: Context Length input now shows {after_reset!r}")
+ # Commit the reset so the stored override is dropped, then assert storage.
+ btn = primary_button(popover)
+ if btn is not None and btn.is_enabled():
+ btn.click()
+ page.wait_for_timeout(1500)
+ cfg = read_configs()
+ pinned = any(
+ _as_int(e.get("customContextLength")) == DISTINCT_CTX for e in config_entries(cfg)
+ )
+ if pinned:
+ fail("Reset left the distinctive context pinned in unsloth_model_configs")
+ else:
+ info("OK reset: distinctive context cleared from unsloth_model_configs")
+ shoot("08-after-reset")
+ close_picker()
+
+ # ─────────────────────────────────────────────────────
+ # 4. Advanced settings persist (best-effort, never gates).
+ # ─────────────────────────────────────────────────────
+ step("advanced (KV cache dtype / tensor parallel) persists")
+ try:
+ popover = open_picker()
+ if open_config(popover, MODEL_HINT) is not None:
+ adv = popover.get_by_role("switch", name = re.compile("advanced settings", re.I)).first
+ if _count(adv):
+ try:
+ adv.check()
+ except Exception:
+ adv.click()
+ page.wait_for_timeout(500)
+ # The Tensor Parallelism Radix Switch has no aria-label, so target the
+ # first switch after the "Tensor Parallelism" text.
+ tp = popover.locator(
+ 'xpath=.//span[contains(text(),"Tensor Parallelism")]'
+ '/following::*[@role="switch"][1]'
+ ).first
+ toggled = False
+ if _count(tp):
+ try:
+ tp.click()
+ toggled = True
+ except Exception:
+ pass
+ remember = popover.get_by_label("Remember for this model").first
+ if _count(remember):
+ try:
+ remember.check()
+ except Exception:
+ remember.click()
+ btn = primary_button(popover)
+ if btn is not None and btn.is_enabled():
+ btn.click()
+ page.wait_for_timeout(1500)
+ cfg = read_configs()
+ has_adv = any(
+ e.get("tensorParallel") or e.get("kvCacheDtype") for e in config_entries(cfg)
+ )
+ if toggled and has_adv:
+ info("OK advanced: tensorParallel/kvCacheDtype persisted")
+ else:
+ runtime_warn(
+ f"advanced persistence not observed (toggled={toggled}, "
+ f"entries={json.dumps(config_entries(cfg))[:300]})"
+ )
+ else:
+ runtime_warn("could not open run-settings for the advanced-persist check")
+ close_picker()
+ except Exception as e:
+ runtime_warn(f"advanced-persist check errored: {e}")
+
+ # ─────────────────────────────────────────────────────
+ # 5. Legacy migration is idempotent (gates in CI via soft_fail).
+ # Seed a pre-feature unsloth_load_settings store, confirm it migrates once
+ # with the value preserved, then reload with a fresh legacy seed and confirm
+ # the migration does not re-run, duplicate, or clobber. Re-running on every
+ # reload was the regression that reverted the predecessor PR.
+ # ─────────────────────────────────────────────────────
+ step("legacy unsloth_load_settings migrates once and stays idempotent")
+ try:
+ legacy_key = f"{GGUF_REPO}::{GGUF_VARIANT}"
+ legacy = {
+ legacy_key: {
+ "contextLength": DISTINCT_CTX,
+ "kvCacheDtype": "q8_0",
+ "tensorParallel": True,
+ }
+ }
+ robust_evaluate(
+ page,
+ "(seed) => {"
+ " localStorage.setItem('unsloth_load_settings', JSON.stringify(seed));"
+ " localStorage.removeItem('unsloth_model_configs');"
+ " localStorage.removeItem('unsloth_model_configs_migrated');"
+ " return true;"
+ "}",
+ arg = legacy,
+ )
+ page.reload()
+ composer = page.locator('textarea[aria-label="Message input"]')
+ composer.wait_for(state = "visible", timeout = 60_000)
+ # Opening the picker config forces the store to read (which migrates).
+ popover = open_picker()
+ open_config(popover, MODEL_HINT)
+ page.wait_for_timeout(800)
+ cfg_first = read_configs()
+ migrated_ctx = any(
+ e.get("customContextLength") == DISTINCT_CTX for e in config_entries(cfg_first)
+ )
+ if migrated_ctx:
+ info(f"OK migration: legacy context {DISTINCT_CTX} preserved after migrating")
+ else:
+ soft_fail(
+ f"legacy context {DISTINCT_CTX} not migrated into unsloth_model_configs "
+ f"(got {json.dumps(cfg_first)[:400]})"
+ )
+ flag_first = robust_evaluate(
+ page, "() => localStorage.getItem('unsloth_model_configs_migrated')"
+ )
+ if flag_first != "1":
+ soft_fail(f"migration flag not set after migrating (got {flag_first!r})")
+ shoot("09-after-migration")
+ close_picker()
+
+ # Idempotency: a second reload with a DIFFERENT legacy entry must not re-run
+ # the migration (the persistent flag blocks it), so the new key must not leak
+ # in, nothing duplicates, and the migrated value is untouched.
+ if migrated_ctx:
+ probe_key = "unsloth/__idem_probe__::Q4_K_M"
+ robust_evaluate(
+ page,
+ "(seed) => {"
+ " localStorage.setItem('unsloth_load_settings', JSON.stringify(seed));"
+ " return true;"
+ "}",
+ arg = {probe_key: {"contextLength": DISTINCT_CTX + 2048, "tensorParallel": True}},
+ )
+ page.reload()
+ composer.wait_for(state = "visible", timeout = 60_000)
+ popover = open_picker()
+ open_config(popover, MODEL_HINT)
+ page.wait_for_timeout(800)
+ cfg_second = read_configs()
+ keys_first = set(cfg_first.keys())
+ keys_second = set(cfg_second.keys())
+ new_keys = keys_second - keys_first
+ still_has_ctx = any(
+ e.get("customContextLength") == DISTINCT_CTX for e in config_entries(cfg_second)
+ )
+ if new_keys:
+ soft_fail(
+ "legacy migration re-ran on a second reload (persistent flag "
+ f"ignored): new keys {sorted(new_keys)}"
+ )
+ elif keys_second != keys_first:
+ soft_fail(
+ "legacy migration dropped entries on a second reload: "
+ f"{sorted(keys_first)} -> {sorted(keys_second)}"
+ )
+ elif not still_has_ctx:
+ soft_fail("legacy migration clobbered the migrated context on a second reload")
+ else:
+ info(
+ "OK migration idempotent: second reload did not re-migrate, duplicate, or clobber"
+ )
+ shoot("10-after-second-reload")
+ close_picker()
+ except Exception as e:
+ soft_fail(f"migration idempotency check errored: {e}")
+
+ # ─────────────────────────────────────────────────────
+ if page_errors:
+ fail(f"page errors during run: {page_errors[:3]!r}")
+
+ browser.close()
+
+if _failed:
+ print(f"[ui-modelcfg] RESULT: FAIL ({len(_failed)} issue(s))", flush = True)
+ for m in _failed:
+ print(f"[ui-modelcfg] - {m}", flush = True)
+ sys.exit(1)
+print("[ui-modelcfg] RESULT: PASS", flush = True)
+sys.exit(0)
diff --git a/tests/studio/test_cached_model_path_selection.py b/tests/studio/test_cached_model_path_selection.py
new file mode 100644
index 0000000000..2b6b9c7829
--- /dev/null
+++ b/tests/studio/test_cached_model_path_selection.py
@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Variant-file selection guards for the cached-model-path endpoint.
+
+The Copy path / Reveal endpoint must resolve a quant label to the same file
+the variant menus offer: MTP drafters, mmproj vision adapters, and big-endian
+builds are excluded, and directory layouts (``BF16/model-00001-of-....gguf``)
+resolve their label from the snapshot-relative path, not the basename.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+
+def _find_repo_root() -> Path | None:
+ env = os.environ.get("UNSLOTH_REPO_ROOT")
+ if env:
+ p = Path(env).resolve()
+ if (p / "studio" / "backend").is_dir():
+ return p
+ here = Path(__file__).resolve()
+ for parent in (here, *here.parents):
+ if (parent / "studio" / "backend").is_dir():
+ return parent
+ return None
+
+
+_REPO_ROOT = _find_repo_root()
+if _REPO_ROOT is None:
+ pytest.skip(
+ "Could not locate studio/backend. Set UNSLOTH_REPO_ROOT or run from "
+ "the repository checkout.",
+ allow_module_level = True,
+ )
+
+_STUDIO_BACKEND = _REPO_ROOT / "studio" / "backend"
+if str(_STUDIO_BACKEND) not in sys.path:
+ sys.path.insert(0, str(_STUDIO_BACKEND))
+
+pytest.importorskip("fastapi")
+pytest.importorskip("huggingface_hub")
+
+try:
+ from routes import models as routes_models
+except Exception as exc:
+ pytest.skip(f"studio backend import unavailable: {exc}", allow_module_level = True)
+
+from fastapi import HTTPException
+
+
+def test_plain_quant_label_resolves():
+ assert routes_models._main_variant_gguf_label("Model-Q8_0.gguf") == "Q8_0"
+
+
+def test_mtp_drafter_in_subdir_is_excluded():
+ assert routes_models._main_variant_gguf_label("MTP/Model-Q8_0-MTP.gguf") is None
+
+
+def test_mtp_drafter_root_prefix_is_excluded():
+ assert routes_models._main_variant_gguf_label("mtp-Model-Q8_0.gguf") is None
+
+
+def test_mmproj_adapter_is_excluded():
+ assert routes_models._main_variant_gguf_label("mmproj-Model-F16.gguf") is None
+
+
+def test_directory_layout_quant_resolves_from_parent_dir():
+ assert routes_models._main_variant_gguf_label("BF16/Model-00001-of-00002.gguf") == "BF16"
+
+
+def test_big_endian_build_is_excluded():
+ assert routes_models._main_variant_gguf_label("Model-Q8_0-BE.gguf") is None
+
+
+def test_non_gguf_file_is_excluded():
+ assert routes_models._main_variant_gguf_label("config.json") is None
+
+
+def test_normalized_quant_label_ignores_separators():
+ assert routes_models._normalized_quant_label("UD-Q4_K_XL") == "udq4kxl"
+ assert routes_models._normalized_quant_label("Q8-0") == routes_models._normalized_quant_label(
+ "Q8_0"
+ )
+
+
+def _revision(
+ snapshot: Path,
+ last_modified: float,
+ names: list[str],
+ size_on_disk: int = 4,
+) -> SimpleNamespace:
+ files = []
+ for name in names:
+ path = snapshot / name
+ path.parent.mkdir(parents = True, exist_ok = True)
+ path.write_bytes(b"x" * size_on_disk)
+ files.append(
+ SimpleNamespace(
+ file_name = name,
+ file_path = path,
+ blob_path = path,
+ size_on_disk = size_on_disk,
+ )
+ )
+ return SimpleNamespace(snapshot_path = snapshot, last_modified = last_modified, files = files)
+
+
+def _patch_cache(monkeypatch, tmp_path: Path, revisions: list[SimpleNamespace]) -> None:
+ repo = SimpleNamespace(
+ repo_id = "Org/Repo",
+ repo_type = "model",
+ repo_path = tmp_path,
+ revisions = revisions,
+ )
+ monkeypatch.setattr(
+ routes_models, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])]
+ )
+
+
+def _repo(root: Path, revisions: list[SimpleNamespace]) -> SimpleNamespace:
+ return SimpleNamespace(
+ repo_id = "Org/Repo",
+ repo_type = "model",
+ repo_path = root,
+ revisions = revisions,
+ )
+
+
+def _patch_caches(monkeypatch, repos: list[SimpleNamespace]) -> None:
+ monkeypatch.setattr(
+ routes_models,
+ "_all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [repo]) for repo in repos],
+ )
+
+
+@pytest.mark.parametrize("newest_first", [True, False])
+def test_variant_resolves_from_newest_revision(monkeypatch, tmp_path, newest_first):
+ old = _revision(tmp_path / "snapshots" / "aaa", 1_000.0, ["Model-Q4_K_M.gguf"])
+ new = _revision(tmp_path / "snapshots" / "bbb", 2_000.0, ["Model-Q4_K_M.gguf"])
+ _patch_cache(monkeypatch, tmp_path, [new, old] if newest_first else [old, new])
+ resolved = routes_models._resolve_cached_model_path("Org/Repo", "Q4_K_M")
+ assert resolved == tmp_path / "snapshots" / "bbb" / "Model-Q4_K_M.gguf"
+
+
+def test_sharded_variant_resolves_first_split(monkeypatch, tmp_path):
+ rev = _revision(
+ tmp_path / "snapshots" / "aaa",
+ 1_000.0,
+ ["Model-Q4_K_M-00002-of-00002.gguf", "Model-Q4_K_M-00001-of-00002.gguf"],
+ )
+ _patch_cache(monkeypatch, tmp_path, [rev])
+ resolved = routes_models._resolve_cached_model_path("Org/Repo", "Q4_K_M")
+ assert resolved.name == "Model-Q4_K_M-00001-of-00002.gguf"
+
+
+def test_variant_only_in_older_revision_resolves(monkeypatch, tmp_path):
+ old = _revision(tmp_path / "snapshots" / "aaa", 1_000.0, ["Model-Q4_K_M.gguf"])
+ new = _revision(tmp_path / "snapshots" / "bbb", 2_000.0, ["Model-Q8_0.gguf"])
+ _patch_cache(monkeypatch, tmp_path, [new, old])
+ resolved = routes_models._resolve_cached_model_path("Org/Repo", "Q4_K_M")
+ assert resolved == tmp_path / "snapshots" / "aaa" / "Model-Q4_K_M.gguf"
+
+
+def test_missing_newest_file_falls_back_to_older_revision(monkeypatch, tmp_path):
+ old = _revision(tmp_path / "snapshots" / "aaa", 1_000.0, ["Model-Q4_K_M.gguf"])
+ new_snapshot = tmp_path / "snapshots" / "bbb"
+ new_snapshot.mkdir(parents = True)
+ new = SimpleNamespace(
+ snapshot_path = new_snapshot,
+ last_modified = 2_000.0,
+ files = [
+ SimpleNamespace(
+ file_name = "Model-Q4_K_M.gguf",
+ file_path = new_snapshot / "Model-Q4_K_M.gguf",
+ )
+ ],
+ )
+ _patch_cache(monkeypatch, tmp_path, [new, old])
+ resolved = routes_models._resolve_cached_model_path("Org/Repo", "Q4_K_M")
+ assert resolved == tmp_path / "snapshots" / "aaa" / "Model-Q4_K_M.gguf"
+
+
+def test_variant_resolves_across_all_cache_roots(monkeypatch, tmp_path):
+ first_root = tmp_path / "active"
+ second_root = tmp_path / "default"
+ old = _revision(
+ first_root / "snapshots" / "aaa",
+ 1_000.0,
+ ["Model-Q4_K_M.gguf"],
+ )
+ new = _revision(
+ second_root / "snapshots" / "bbb",
+ 2_000.0,
+ ["Model-Q4_K_M.gguf"],
+ )
+ _patch_caches(
+ monkeypatch,
+ [_repo(first_root, [old]), _repo(second_root, [new])],
+ )
+ resolved = routes_models._resolve_cached_model_path("Org/Repo", "Q4_K_M")
+ assert resolved == second_root / "snapshots" / "bbb" / "Model-Q4_K_M.gguf"
+
+
+def test_repo_path_matches_largest_visible_cache_entry(monkeypatch, tmp_path):
+ first_root = tmp_path / "active"
+ second_root = tmp_path / "default"
+ small = _revision(
+ first_root / "snapshots" / "aaa",
+ 2_000.0,
+ ["Model-Q8_0.gguf"],
+ size_on_disk = 4,
+ )
+ large = _revision(
+ second_root / "snapshots" / "bbb",
+ 1_000.0,
+ ["Model-Q8_0.gguf"],
+ size_on_disk = 8,
+ )
+ _patch_caches(
+ monkeypatch,
+ [_repo(first_root, [small]), _repo(second_root, [large])],
+ )
+ resolved = routes_models._resolve_cached_model_path("Org/Repo", None)
+ assert resolved == second_root / "snapshots" / "bbb"
+
+
+def test_unknown_variant_raises_404(monkeypatch, tmp_path):
+ rev = _revision(tmp_path / "snapshots" / "aaa", 1_000.0, ["Model-Q8_0.gguf"])
+ _patch_cache(monkeypatch, tmp_path, [rev])
+ with pytest.raises(HTTPException) as excinfo:
+ routes_models._resolve_cached_model_path("Org/Repo", "Q4_K_M")
+ assert excinfo.value.status_code == 404
+ assert "Q4_K_M" in excinfo.value.detail
diff --git a/tests/studio/test_gpu_inference_smoke.py b/tests/studio/test_gpu_inference_smoke.py
new file mode 100644
index 0000000000..a28166e1e0
--- /dev/null
+++ b/tests/studio/test_gpu_inference_smoke.py
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Fast, GPU-gated real-inference smoke.
+
+GitHub-hosted CI runners have no GPU, so this AUTO-SKIPS there; the full picker
+-> load -> chat flow is covered on CPU by tests/studio/playwright_model_config.py
+and studio-ui-smoke.yml. This test adds a quick real-generation check for local
+dev and self-hosted GPU runners: it loads the smallest model (gemma-3-270m-it)
+on the GPU and does a single short greedy generation, asserting a non-empty
+reply. Kept deliberately short (a handful of new tokens) so it is a confidence
+check, not a benchmark. Select/deselect it by name, e.g. `-k gpu_generation`.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+# Smallest instruct model in the CI fixture family; ~270M params loads and
+# generates a few tokens in seconds on any GPU.
+MODEL_ID = "unsloth/gemma-3-270m-it"
+# A handful of forced real tokens: enough to prove GPU decode produced content,
+# short enough to stay a few seconds.
+MIN_NEW_TOKENS = 4
+MAX_NEW_TOKENS = 16
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason = "requires a CUDA GPU")
+def test_gpu_generation_smoke():
+ try:
+ from transformers import AutoModelForCausalLM, AutoTokenizer
+ except Exception as exc: # pragma: no cover - env without transformers
+ pytest.skip(f"transformers unavailable: {exc}")
+
+ # Gemma is numerically unstable in fp16 (it emits only ); use bf16 where
+ # supported, else fp32. The model is tiny, so fp32 is still fast.
+ dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
+ try:
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+ model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype = dtype).to("cuda")
+ except Exception as exc: # offline / gated / download failure is not a code defect
+ pytest.skip(f"could not fetch/load {MODEL_ID}: {exc}")
+
+ model.eval()
+ messages = [{"role": "user", "content": "Say hello in one word."}]
+ inputs = tokenizer.apply_chat_template(
+ messages, add_generation_prompt = True, return_dict = True, return_tensors = "pt"
+ ).to("cuda")
+ prompt_len = inputs["input_ids"].shape[1]
+
+ with torch.no_grad():
+ output = model.generate(
+ **inputs,
+ min_new_tokens = MIN_NEW_TOKENS,
+ max_new_tokens = MAX_NEW_TOKENS,
+ do_sample = False,
+ )
+
+ # The model produced new tokens on the GPU (the real inference proof)...
+ assert output.shape[1] > prompt_len, "no tokens were generated on the GPU"
+ # ...and they decode to non-empty text (min_new_tokens forces real content).
+ reply = tokenizer.decode(output[0][prompt_len:], skip_special_tokens = True)
+ assert reply.strip(), "expected a non-empty GPU generation"
diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py
new file mode 100644
index 0000000000..20335a279c
--- /dev/null
+++ b/tests/studio/test_model_picker_contracts.py
@@ -0,0 +1,407 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Source-contract guards for the model-picker per-model-config feature.
+
+These are cheap, CPU-only, no-browser checks that read the frontend source and
+assert the specific fixes that got the predecessor PR reverted stay in place. If
+a future edit reverts one of them (e.g. rounds the context ceiling up again, or
+puts the HF token back in the URL), the matching assertion reddens. They pair
+with the runtime Playwright checks (which prove the behavior end to end) and the
+backend pytest checks (which prove the backend logic).
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+WORKDIR = Path(__file__).resolve().parents[2]
+FRONTEND = WORKDIR / "studio" / "frontend" / "src"
+
+
+def _read(rel: str) -> str:
+ path = FRONTEND / rel
+ assert path.exists(), f"missing source file: {path}"
+ return path.read_text()
+
+
+def test_models_api_sends_token_via_header_not_query():
+ """getModelConfig / checkVisionModel / checkEmbeddingModel must pass the HF
+ token through hubTokenHeader, never as a ?hf_token= query param (which leaks
+ the credential into server/proxy access logs)."""
+ src = _read("features/training/api/models-api.ts")
+ assert src.count("hubTokenHeader(") >= 3
+ assert "hf_token=" not in src
+ assert '"hf_token"' not in src and "'hf_token'" not in src
+
+
+def test_model_metadata_probe_never_puts_token_in_query():
+ src = _read("features/model-picker/api/model-metadata.ts")
+ assert "hf_token=" not in src
+ assert '"hf_token"' not in src and "'hf_token'" not in src
+
+
+def test_model_config_page_floors_the_context_ceiling():
+ """The model's native max-context must be FLOORED to the step grid, never
+ rounded up (rounding up can offer/persist a length above the model's real
+ ceiling and break loading)."""
+ src = _read("features/model-picker/components/model-config-page.tsx")
+ assert "floorMaxSeqLength(modelMaxPosition.maxPositionEmbeddings)" in src
+ assert "normalizeMaxSeqLength(modelMaxPosition.maxPositionEmbeddings)" not in src
+
+
+def test_compare_load_clears_stale_native_lease():
+ """A compare-pane load never comes from the desktop file picker, so it must
+ clear any prior picked file's lease token + expiry, otherwise a reload can
+ send a stale lease for the now-active model."""
+ src = _read("features/chat/shared-composer.tsx")
+ assert "activeNativePathToken: null" in src
+ assert "activeNativePathExpiresAtMs: null" in src
+
+
+def test_rollback_restores_native_lease_expiry_with_token():
+ """A failed model switch that rolls back to a previously loaded picked GGUF
+ must restore the lease expiry paired with the token, never the token alone
+ (which would look non-expiring and skip the expiry guard)."""
+ src = _read("features/chat/hooks/use-chat-model-runtime.ts")
+ assert "previousActiveNativePathExpiresAtMs" in src
+ assert re.search(
+ r"activeNativePathExpiresAtMs:\s*previousActiveNativePathToken", src
+ ), "rollback must restore the expiry alongside the token"
+
+
+def test_default_caches_keyed_on_inventory_version():
+ """The chat-template and max-position caches must key on the inventory
+ version so a model update in the same session invalidates the cached value
+ instead of showing the stale revision."""
+ src = _read("features/model-picker/hooks/use-model-defaults.ts")
+ # Both cache keys (template + max-position) end with the inventory version.
+ assert src.count("${inventoryVersion}") >= 2
+
+
+def test_hidden_infra_model_needles_present():
+ """The frontend static needle list must keep hiding the RAG embedder and the
+ llama.cpp validation probe."""
+ src = _read("features/hub/lib/hidden-models.ts")
+ assert '"bge-small-en-v1.5"' in src
+ assert '"ggml-org/models"' in src
+ assert '"stories260k.gguf"' in src
+
+
+def test_hidden_models_dynamic_exact_ids_wired():
+ """The configured embedder arrives from /api/hub/hidden-models as exact
+ repo ids; a substring needle would let a generic basename like "model"
+ hide unrelated chat models."""
+ src = _read("features/hub/lib/hidden-models.ts")
+ assert "toLowerStrings(data.exact_ids)" in src
+ assert "dynamicExactIds.includes(lower)" in src
+
+
+def test_hidden_model_matchers_refresh_with_inventory_version():
+ src = _read("features/hub/lib/hidden-models.ts")
+ assert "const version = getInventoryVersion()" in src
+ assert "matchersFetchVersion === version" in src
+ assert "getInventoryVersion() !== version" in src
+
+
+def test_diffusion_capability_labeled_image_generation():
+ """The diffusion capability detects image GENERATORS (FLUX, SDXL,
+ text-to-image tags); labeling it "Image to text" showed generators when
+ users asked for captioning models."""
+ for rel in (
+ "features/hub/lib/model-capabilities.ts",
+ "features/hub/lib/model-type-filter.ts",
+ "features/hub/lib/view-models.ts",
+ ):
+ src = _read(rel)
+ assert "Image to text" not in src, rel
+ assert "Image generation" in src, rel
+
+
+def test_active_model_config_round_trips_gpu_fields():
+ """The active model's config must carry the GPU Memory knobs (GGUF only) so
+ a sidebar/hub-gear reload cannot silently reset manual GPU settings, and
+ "Remember settings" cannot persist a GPU-less config over a saved one."""
+ src = _read("features/model-picker/hooks/use-active-model-config.ts")
+ for field in ("gpuMemoryMode", "gpuLayers", "nCpuMoe", "selectedGpuIds"):
+ assert field in src, field
+ assert "if (!isGguf)" in src and "return base" in src
+ for rel in (
+ "features/chat/chat-page.tsx",
+ "features/hub/catalog/sampling-settings-dialog.tsx",
+ ):
+ assert "useActiveModelConfig(" in _read(rel), rel
+ signature = _read("features/model-picker/components/sidebar-model-config.tsx")
+ assert "gpuFieldsSignature(config)" in signature
+ shared = _read("features/model-picker/model-config/apply-per-model-config.ts")
+ assert "export function gpuFieldsSignature" in shared
+
+
+def test_compare_load_uses_each_models_gpu_config():
+ src = _read("features/chat/shared-composer.tsx")
+ assert "ownConfig.gpuMemoryMode ?? compareLoadKnobs.gpuMemoryMode" in src
+ assert "ownConfig.gpuLayers ?? compareLoadKnobs.gpuLayers" in src
+ assert "ownConfig.nCpuMoe ?? compareLoadKnobs.nCpuMoe" in src
+ assert "if (ownConfig.selectedGpuIds != null)" in src
+ assert "reconcilePersistedGpuIds(ownConfig.selectedGpuIds)" in src
+ for field in (
+ "gpu_memory_mode: effectiveGpuMemoryMode",
+ "gpu_layers: effectiveGpuLayers",
+ "n_cpu_moe: effectiveNCpuMoe",
+ "gpu_ids: effectiveSelectedGpuIds ?? undefined",
+ ):
+ assert field in src
+
+
+def test_active_native_gguf_metadata_uses_path_token():
+ src = _read("features/model-picker/components/model-config-page.tsx")
+ assert "(isActiveModel ? activeNativePathToken : null)" in src
+ assert "target.meta.nativePathToken ??" in src
+ assert "nativePathToken," in src
+ assert '${nativePathToken ?? ""}' in src
+
+
+def test_model_default_hooks_do_not_reset_state_in_effect():
+ src = _read("features/model-picker/hooks/use-model-defaults.ts")
+ assert "setFetched(null)" not in src
+
+
+def test_variant_expander_refreshes_after_delete():
+ """Deleting a downloaded quant from an expanded repo that still has other
+ cached quants must bump the expander refresh key, or the deleted quant stays
+ shown as downloaded and clickable and tries to reload the removed file."""
+ src = _read("features/model-picker/components/model-selector/pickers.tsx")
+ del_confirm = re.search(
+ r"await onDeleteVariant\(v\.quant\);.*?setRefreshKey\(\(key\) => key \+ 1\)",
+ src,
+ re.S,
+ )
+ assert del_confirm, "delete onConfirm must bump refreshKey after a successful delete"
+
+
+def test_local_picker_rows_require_chat_capability():
+ """Local inventory rows can be classified non-chat (canChat false, e.g. a
+ folder with only config.json). The picker must filter those out, or selecting
+ one loads a weightless path; toLocalModelInfo drops capabilities so the memo
+ is the only place the guard can live."""
+ src = _read("features/model-picker/inventory/use-chat-picker-inventory.ts")
+ memo = re.search(r"const localModels = useMemo\(.*?\[inventory\.localRows\]", src, re.S)
+ assert memo, "localModels memo not found"
+ assert "row.capabilities.canChat" in memo.group(0)
+
+
+def test_native_picked_gguf_template_read_through_lease():
+ """A native (picked / drag-drop) GGUF's path lives only in its signed lease,
+ and the picker chat-template GET has no lease plumbing, so the default
+ template must be read through the lease-aware validate probe: mint a
+ validate-model lease and post include_chat_template. The native token also
+ has to reach the fetch (threaded through the hook) and be part of the cache
+ key so two picks of the same basename don't share a template."""
+ api = _read("features/model-picker/api/templates.ts")
+ assert 'consumeNativePathToken(nativePathToken, "validate-model")' in api
+ assert "include_chat_template: true" in api
+ assert "/api/inference/validate" in api
+ hook = _read("features/model-picker/hooks/use-model-defaults.ts")
+ assert "nativePathToken," in hook
+ assert '${nativePathToken ?? ""}' in hook
+
+
+def test_model_load_guard_is_cross_instance():
+ """The in-flight load guard must consult the shared store pick (not only the
+ per-hook ref) and ejectModel must refuse while any instance is loading:
+ three live useChatModelRuntime instances exist (chat page, hub page, hub
+ gear dialog)."""
+ src = _read("features/chat/hooks/use-chat-model-runtime.ts")
+ assert "useChatRuntimeStore.getState().loadingModelPick" in src
+ assert "clearLoadingModelPick" in src
+ eject_body = src.split("const ejectModel", 1)[1]
+ assert "loadingModelPick" in eject_body.split("ejectModel,", 1)[0]
+
+
+def test_partial_safetensors_download_keeps_delete_menu():
+ """A stopped partial safetensors download must keep its options menu (the
+ Delete affordance) like the GGUF card does, or partial downloads can only
+ be cleaned up by finishing or leaving them. During an ACTIVE download the
+ menu stays hidden (every item would be disabled: no Copy path while not
+ downloaded, no Delete while downloading, pin suppressed in the run bar)."""
+ src = _read("features/hub/catalog/safetensors-download-card.tsx")
+ assert "(isDownloaded || (isPartial && !downloading))" in src
+
+
+def test_pinned_validation_uses_cached_local_variant_listing():
+ """Pinned-quant validation must use the TTL-cached hub client with
+ preferLocalCache (downloaded-ness is local state) instead of one uncached
+ round-trip per pinned repo on every picker open. Picker deletes must go
+ through the hub inventory client, whose delete invalidates both the
+ variants TTL cache and the server-side HF cache scan (the legacy
+ /api/models/delete-cached route invalidates neither, so a post-delete
+ inventory refresh would resurrect the deleted row until the scan TTL)."""
+ src = _read("features/model-picker/components/model-selector/pickers.tsx")
+ assert "listGgufVariantsCached(" in src
+ assert "preferLocalCache: true" in src
+ assert re.search(r'import \{[^}]*\bdeleteCachedModel\b[^}]*\} from "@/features/hub"', src)
+ hub_api = _read("features/hub/inventory/api.ts")
+ delete_fn = hub_api.split("export async function deleteCachedModel", 1)[1]
+ delete_fn = delete_fn.split("export ", 1)[0]
+ assert "invalidateGgufVariantsCache(" in delete_fn
+ assert "bumpInventoryVersion(" in delete_fn
+
+
+def test_downloaded_list_offsets_virtual_rows():
+ """The On Device virtualized list sits below the Pinned block in the same
+ scroll element, so it must pass its measured offset as scrollMargin or rows
+ past the overscan render blank."""
+ src = _read("features/hub/catalog/models-catalog-lists.tsx")
+ assert "scrollMargin={scrollMargin}" in src
+
+
+def test_local_gguf_diagnostics_gate_on_broad_is_gguf():
+ """The MTP fallback note and the context/VRAM warning must gate on the broad
+ isGguf (variant, loaded gguf context, or .gguf suffix), not the variant-only
+ isLoadedGguf, so direct-file and custom-folder GGUF loads keep those
+ diagnostics."""
+ src = _read("features/chat/chat-settings-sheet.tsx")
+ spec = re.search(r"const showSpecFallback =.*?;", src, re.S)
+ vram = re.search(r"const showContextVramWarning =.*?;", src, re.S)
+ assert spec and "isGguf &&" in spec.group(0) and "isLoadedGguf" not in spec.group(0)
+ assert vram and "isGguf &&" in vram.group(0) and "isLoadedGguf" not in vram.group(0)
+
+
+def test_fixed_layer_gguf_pins_displayed_context():
+ """An already-loaded auto-fit GGUF saved with Manual fixed GPU layers must
+ pin the shown context, so a later fresh load keeps the fitted placement
+ instead of sending native/0 and recreating the OOM."""
+ src = _read("features/model-picker/components/model-config-page.tsx")
+ assert "const pinFixedLayerContext =" in src
+ assert 'config.gpuMemoryMode === "manual"' in src
+ assert "customContextLength: activeLoadedContext" in src
+
+
+def test_auto_defaults_not_persisted_as_overrides():
+ """Auto GPU memory mode and Auto/default speculative type are follow-global
+ defaults; normalization must not persist them as per-model overrides, else a
+ model stops following later changes to the global preference."""
+ src = _read("features/model-picker/model-config/per-model-config.ts")
+ assert 'if (partial.gpuMemoryMode === "manual") {' in src
+ assert 'partial.gpuMemoryMode === "auto" || partial.gpuMemoryMode === "manual"' not in src
+ spec = re.search(r'if \(s === "auto" \|\| s === "default"\) \{\s*return ([^;]+);', src)
+ assert spec and spec.group(1).strip() == "null"
+
+
+def test_compare_pane_context_from_own_config_only():
+ """A compare pane's context comes from its own config only (a saved pin, else
+ null for Auto/native); it must not inherit the active model's shared snapshot,
+ which resolveFitMaxSeqLength would treat as an explicit pin (VRAM/OOM)."""
+ src = _read("features/chat/shared-composer.tsx")
+ assert "const effectiveCustomContextLength = ownConfig.customContextLength;" in src
+ assert "compareLoadKnobs.customContextLength" not in src
+
+
+def test_reset_max_seq_length_falls_back_to_app_default():
+ """After Reset clears maxSeqLength (null), a non-GGUF active model's shown
+ max sequence length must fall back to the app default, never the loaded
+ runtime snapshot, or a remembered/active override can never be cleared."""
+ src = _read("features/model-picker/components/model-config-page.tsx")
+ # The null fallback resolves to the app-default constant, not a runtime value.
+ assert "clampMaxSeqLength(DEFAULT_MAX_SEQ_LENGTH, nativeMaxSeqLength)" in src
+ # The buggy runtime-seeded fallback must not come back.
+ assert "clampMaxSeqLength(initialMaxSeqLength" not in src
+
+
+def test_reset_persists_null_max_length_and_substitutes_only_for_load():
+ """The persisted per-model record must keep config.maxSeqLength (null after
+ Reset) so isDefaultConfig can clear a remembered override; the concrete
+ fallback is substituted only into the load request, not the saved record."""
+ src = _read("features/model-picker/components/model-config-page.tsx")
+ # Load-only substitution of the resolved value.
+ assert "maxSeqLength: maxSeqLengthValue" in src
+ assert "const loadConfig" in src
+ # The persisted record is loaded via onRun(loadConfig), and save uses the
+ # untouched runtimeConfig (so a reset/default config stays default).
+ assert "onRun(loadConfig)" in src
+ assert "savePerModelConfig(" in src
+
+
+def test_reset_enabled_for_explicit_context_pin_at_native():
+ """An explicit customContextLength that equals the native ceiling is still a
+ user override, so contextAtDefault must require customContextLength == null.
+ The buggy form treated `contextValue === native` alone as default, wedging
+ the Reset button disabled for a deliberate pin-to-native."""
+ src = " ".join(_read("features/model-picker/components/model-config-page.tsx").split())
+ assert (
+ "const contextAtDefault = !target.isGguf || "
+ "(config.customContextLength == null && "
+ "(nativeContextLength == null || contextValue === nativeContextLength));" in src
+ )
+ # The old form that ignored an explicit pin equal to native must not return.
+ assert (
+ "(nativeContextLength == null ? config.customContextLength == null : "
+ "contextValue === nativeContextLength)" not in src
+ )
+ # The app-default constant is the single source of truth (imported, not local).
+ assert "DEFAULT_MAX_SEQ_LENGTH," in src
+ assert "const DEFAULT_MAX_SEQ_LENGTH = 4096" not in src
+
+
+def test_compare_pane_non_gguf_falls_back_to_app_default():
+ """A non-GGUF compare pane with no saved maxSeqLength must fall back to the
+ shared app default, not the active model's runtime snapshot; otherwise an
+ unconfigured pane inherits a saved 128K neighbor's context and can OOM."""
+ per_model = _read("features/model-picker/model-config/per-model-config.ts")
+ assert "export const DEFAULT_MAX_SEQ_LENGTH = 4096;" in per_model
+ barrel = _read("features/model-picker/index.ts")
+ assert "DEFAULT_MAX_SEQ_LENGTH," in barrel
+ src = " ".join(_read("features/chat/shared-composer.tsx").split())
+ assert "DEFAULT_MAX_SEQ_LENGTH," in src
+ assert (
+ "const effectiveMaxSeqLength = ownConfig.customContextLength ?? "
+ "normalizeMaxSeqLength(ownConfig.maxSeqLength) ?? "
+ "(isGgufLoad ? 0 : DEFAULT_MAX_SEQ_LENGTH);" in src
+ )
+ # The buggy fallback to the active model's shared runtime value must not return.
+ assert "(isGgufLoad ? 0 : maxSeqLength)" not in src
+ assert "const maxSeqLength = store.params.maxSeqLength;" not in src
+
+
+def test_default_gpu_mode_clears_manual_knobs():
+ """Switching GPU Memory back to Default must clear the Manual-only knobs
+ (gpuLayers/nCpuMoe/selectedGpuIds); otherwise a remembered config keeps stale
+ pins that a later load re-applies when the global preference is Manual."""
+ src = _read("features/model-picker/components/model-config-page.tsx")
+ assert 'gpuMemoryMode: "auto",' in src
+ assert "gpuLayers: undefined," in src
+ assert "nCpuMoe: undefined," in src
+ assert "selectedGpuIds: undefined," in src
+
+
+def test_legacy_migration_is_idempotent_and_non_destructive():
+ """The v1->v2 localStorage migration (unsloth_load_settings ->
+ unsloth_model_configs) is invoked on every store read, so it must be
+ idempotent: repeated reads, browser reloads, and Studio restarts must never
+ re-migrate, duplicate records, or overwrite a newer per-model config. This
+ was the class of regression that reverted the predecessor PR, so pin all
+ three idempotency layers at source level; dropping any of them reddens here.
+ """
+ raw = _read("features/model-picker/model-config/per-model-config.ts")
+ src = " ".join(raw.split())
+ # Migration runs from readMap (every store read), so it must be safe to repeat.
+ assert (
+ "function readMap(): StoredMap { migrateLegacyLoadSettingsOnce(); "
+ "return readMapRaw(); }" in src
+ )
+ # Layer 1: in-memory once-per-session guard so repeated readMap() calls
+ # migrate at most once.
+ assert "let legacyMigrationChecked = false;" in src
+ assert "if (legacyMigrationChecked || !canUseStorage()) {" in src
+ assert "legacyMigrationChecked = true;" in src
+ # Layer 2: persistent cross-session flag so a completed migration is never
+ # redone. Set in every terminal branch (malformed data, nothing to migrate,
+ # successful write); a failed quota write leaves it unset so the next session
+ # retries. Three set-sites encode exactly that.
+ assert 'const LEGACY_MIGRATION_FLAG = "unsloth_model_configs_migrated";' in src
+ assert "if (localStorage.getItem(LEGACY_MIGRATION_FLAG)) {" in src
+ assert src.count('localStorage.setItem(LEGACY_MIGRATION_FLAG, "1");') >= 3
+ # Layer 3: non-overwriting merge skips an existing (or default) key, so even a
+ # forced re-run cannot duplicate or clobber a user's config.
+ assert "if (isDefaultConfig(migrated) || Object.hasOwn(map, key)) {" in src
diff --git a/tests/studio/test_reveal_file_manager.py b/tests/studio/test_reveal_file_manager.py
new file mode 100644
index 0000000000..8a1df586e5
--- /dev/null
+++ b/tests/studio/test_reveal_file_manager.py
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Platform guards for the reveal-in-file-manager endpoint.
+
+A stock WSL distro has no Linux desktop, so the generic Linux branch
+(``xdg-open``) fails there. Under WSL the reveal must route through Windows
+interop (``wslpath -w`` + ``explorer.exe``), fall back to ``xdg-open`` when
+interop is unavailable, and leave native Linux behavior unchanged.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+
+def _find_repo_root() -> Path | None:
+ env = os.environ.get("UNSLOTH_REPO_ROOT")
+ if env:
+ p = Path(env).resolve()
+ if (p / "studio" / "backend").is_dir():
+ return p
+ here = Path(__file__).resolve()
+ for parent in (here, *here.parents):
+ if (parent / "studio" / "backend").is_dir():
+ return parent
+ return None
+
+
+_REPO_ROOT = _find_repo_root()
+if _REPO_ROOT is None:
+ pytest.skip(
+ "Could not locate studio/backend. Set UNSLOTH_REPO_ROOT or run from "
+ "the repository checkout.",
+ allow_module_level = True,
+ )
+
+_STUDIO_BACKEND = _REPO_ROOT / "studio" / "backend"
+if str(_STUDIO_BACKEND) not in sys.path:
+ sys.path.insert(0, str(_STUDIO_BACKEND))
+
+pytest.importorskip("fastapi")
+pytest.importorskip("huggingface_hub")
+
+try:
+ from routes import models as routes_models
+ from utils.paths import path_utils
+except Exception as exc:
+ pytest.skip(f"studio backend import unavailable: {exc}", allow_module_level = True)
+
+_WINDOWS_PATH = r"\\wsl.localhost\Distro\cache\model.gguf"
+
+
+@pytest.fixture()
+def linux_host(monkeypatch):
+ monkeypatch.setattr(sys, "platform", "linux")
+ monkeypatch.setattr(os, "name", "posix")
+
+
+@pytest.fixture()
+def spawned(monkeypatch):
+ calls = types.SimpleNamespace(run = [], popen = [], run_error = None)
+
+ def fake_run(cmd, **kwargs):
+ calls.run.append(list(cmd))
+ if calls.run_error is not None:
+ raise calls.run_error
+ return types.SimpleNamespace(stdout = _WINDOWS_PATH + "\n")
+
+ def fake_popen(cmd, **kwargs):
+ calls.popen.append(list(cmd))
+ return types.SimpleNamespace()
+
+ monkeypatch.setattr(subprocess, "run", fake_run)
+ monkeypatch.setattr(subprocess, "Popen", fake_popen)
+ return calls
+
+
+def test_wsl_file_is_selected_in_explorer(linux_host, spawned, monkeypatch, tmp_path):
+ monkeypatch.setattr(path_utils, "_IS_WSL", True)
+ target = tmp_path / "model.gguf"
+ target.write_bytes(b"gguf")
+ routes_models._reveal_in_file_manager(target)
+ assert spawned.run == [["wslpath", "-w", str(target)]]
+ assert spawned.popen == [["explorer.exe", f"/select,{_WINDOWS_PATH}"]]
+
+
+def test_wsl_directory_opens_in_explorer(linux_host, spawned, monkeypatch, tmp_path):
+ monkeypatch.setattr(path_utils, "_IS_WSL", True)
+ routes_models._reveal_in_file_manager(tmp_path)
+ assert spawned.popen == [["explorer.exe", _WINDOWS_PATH]]
+
+
+def test_wsl_without_interop_falls_back_to_xdg_open(linux_host, spawned, monkeypatch, tmp_path):
+ monkeypatch.setattr(path_utils, "_IS_WSL", True)
+ spawned.run_error = FileNotFoundError("wslpath")
+ target = tmp_path / "model.gguf"
+ target.write_bytes(b"gguf")
+ routes_models._reveal_in_file_manager(target)
+ assert spawned.popen == [["xdg-open", str(tmp_path)]]
+
+
+def test_wsl_empty_conversion_falls_back_to_xdg_open(linux_host, spawned, monkeypatch, tmp_path):
+ monkeypatch.setattr(path_utils, "_IS_WSL", True)
+
+ def empty_run(cmd, **kwargs):
+ return types.SimpleNamespace(stdout = "\n")
+
+ monkeypatch.setattr(subprocess, "run", empty_run)
+ routes_models._reveal_in_file_manager(tmp_path)
+ assert spawned.popen == [["xdg-open", str(tmp_path)]]
+
+
+def test_native_linux_keeps_xdg_open_on_parent_directory(
+ linux_host, spawned, monkeypatch, tmp_path
+):
+ monkeypatch.setattr(path_utils, "_IS_WSL", False)
+ target = tmp_path / "model.gguf"
+ target.write_bytes(b"gguf")
+ routes_models._reveal_in_file_manager(target)
+ assert spawned.run == []
+ assert spawned.popen == [["xdg-open", str(tmp_path)]]
diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py
index 359ca4873a..ae8c8d6d8a 100644
--- a/tests/studio/test_studio_text_descender_clipping.py
+++ b/tests/studio/test_studio_text_descender_clipping.py
@@ -10,7 +10,14 @@ from pathlib import Path
WORKDIR = Path(__file__).resolve().parents[2]
MODEL_SELECTOR = (
- WORKDIR / "studio" / "frontend" / "src" / "components" / "assistant-ui" / "model-selector.tsx"
+ WORKDIR
+ / "studio"
+ / "frontend"
+ / "src"
+ / "features"
+ / "model-picker"
+ / "components"
+ / "model-selector.tsx"
)
APP_SIDEBAR = WORKDIR / "studio" / "frontend" / "src" / "components" / "app-sidebar.tsx"
From 5f92658ac35f4202fb9f32cbc5320bfe4715c997 Mon Sep 17 00:00:00 2001
From: Wasim Yousef Said
Date: Tue, 21 Jul 2026 09:51:18 +0200
Subject: [PATCH 047/255] Fix Studio desktop reliability (#7255)
* Fix Studio desktop reliability
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix desktop export completion and layout migration
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix maximized setup layout migration
* Adapt desktop exports to data settings
* fix(studio): harden desktop reliability edge cases
* fix(studio): preserve rounded combobox focus fill
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
studio/frontend/src/app/provider.tsx | 22 +-
.../frontend/src/components/app-sidebar.tsx | 13 +-
.../src/components/assistant-ui/thread.tsx | 39 ++-
.../src/components/tauri/window-titlebar.tsx | 10 +-
.../src/components/ui/input-group.tsx | 2 +-
.../src/features/chat/projects-page.tsx | 52 ++-
.../prompt-storage/prompt-storage-dialog.tsx | 152 +++++----
.../src/features/chat/shared-composer.tsx | 11 +-
.../src/features/chat/thread-sidebar.tsx | 17 +-
.../chat/utils/export-chat-history.ts | 15 +-
.../src/features/settings/tabs/data-tab.tsx | 45 ++-
studio/frontend/src/lib/native-files.ts | 86 +++++
studio/src-tauri/Cargo.lock | 1 +
studio/src-tauri/Cargo.toml | 1 +
studio/src-tauri/Entitlements.plist | 3 +
studio/src-tauri/Info.plist | 8 +
studio/src-tauri/src/app_layout.rs | 274 +++++++++++++++
studio/src-tauri/src/commands.rs | 51 ++-
studio/src-tauri/src/main.rs | 7 +
studio/src-tauri/src/native_file_dialogs.rs | 320 ++++++++++++++++++
studio/src-tauri/src/native_path_policy.rs | 4 +-
studio/src-tauri/tauri.macos.conf.json | 1 +
...t_desktop_reliability_frontend_contract.py | 131 +++++++
23 files changed, 1137 insertions(+), 128 deletions(-)
create mode 100644 studio/frontend/src/lib/native-files.ts
create mode 100644 studio/src-tauri/Info.plist
create mode 100644 studio/src-tauri/src/app_layout.rs
create mode 100644 studio/src-tauri/src/native_file_dialogs.rs
create mode 100644 tests/studio/test_desktop_reliability_frontend_contract.py
diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx
index a7e9469cfc..e6c89b9cd7 100644
--- a/studio/frontend/src/app/provider.tsx
+++ b/studio/frontend/src/app/provider.tsx
@@ -51,9 +51,12 @@ async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise {
const { getCurrentWindow, LogicalSize } = await import(
"@tauri-apps/api/window"
);
+ const { invoke } = await import("@tauri-apps/api/core");
if (!isCurrent()) return;
const win = getCurrentWindow();
+ await invoke("reset_app_window_layout_initialized");
+ if (!isCurrent()) return;
await win.setResizable(false);
if (!isCurrent()) return;
await win.setSize(new LogicalSize(SETUP_WINDOW_WIDTH, SETUP_WINDOW_HEIGHT));
@@ -98,18 +101,20 @@ async function applyAppWindowLayout(
if (!isCurrent()) return;
const win = getCurrentWindow();
- // Decide first-launch vs restore from the on-disk state file BEFORE touching the
- // window. Probing the window after restoreStateCurrent is unreliable: on GTK,
- // set_size on a hidden window is deferred until show(), so innerSize() reads a
- // stale value and a baseline fallback would overwrite the queued restore. On
- // macOS the same probe works, hence the inconsistency between prior iterations.
- const hasSavedState = await invoke("has_saved_window_state");
+ // Setup-window activity may create plugin state before the full app is ever
+ // shown, so use a dedicated full-app marker to decide whether restoration is
+ // appropriate. Keep checking plugin state so a missing/corrupt state file
+ // falls back to a monitor-safe centered layout.
+ const [hasInitializedAppLayout, hasSavedState] = await Promise.all([
+ invoke("has_initialized_app_window_layout"),
+ invoke("has_saved_window_state"),
+ ]);
if (!isCurrent()) return;
await win.setResizable(true);
if (!isCurrent()) return;
- if (hasSavedState) {
+ if (hasInitializedAppLayout && hasSavedState) {
// Subsequent launch: plugin restores size/position/maximized, with built-in
// off-screen protection for positions saved on a now-disconnected display.
await restoreStateCurrent(
@@ -144,6 +149,9 @@ async function applyAppWindowLayout(
});
if (!isCurrent()) return;
await enforceMinimumWindowSize(win, LogicalSize, isCurrent);
+
+ if (!isCurrent()) return;
+ await invoke("mark_app_window_layout_initialized");
}
async function showWindowFallback(): Promise {
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 293971904f..8dd7bcdd9a 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -141,6 +141,7 @@ import {
import type { TrainingRunSummary } from "@/features/training";
import { useExportRuntimeStore } from "@/features/export";
import { useEffect, useMemo, useRef, useState, type ReactNode } from "react";
+import { isDownloadCancelled } from "@/lib/native-files";
import { toast } from "@/lib/toast";
import { ShutdownDialog } from "@/components/shutdown-dialog";
import { translate, useT, type TranslationKey } from "@/i18n";
@@ -971,11 +972,13 @@ export function AppSidebar() {
const ids = item.type === "single"
? [item.id]
: (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id);
- await Promise.all(
- ids.map((id) => exportConversationByFormat(id, format)),
- );
- } catch {
- toast.error("Export failed.");
+ for (const id of ids) {
+ await exportConversationByFormat(id, format);
+ }
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Export failed.");
+ }
}
}}
>
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 32fbd61e09..bece2930b3 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -102,6 +102,7 @@ import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params";
import { isTauri } from "@/lib/api-base";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { MicIcon } from "@/lib/mic-icon";
+import { downloadFile, isDownloadCancelled } from "@/lib/native-files";
import { toast } from "@/lib/toast";
import { Tick02Icon } from "@/lib/tick-icon";
import { cn } from "@/lib/utils";
@@ -2911,9 +2912,9 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
{
if (!activeThreadId) return;
- exportConversationRawJsonl(activeThreadId).catch(() =>
- toast.error("Export failed."),
- );
+ exportConversationRawJsonl(activeThreadId).catch((error) => {
+ if (!isDownloadCancelled(error)) toast.error("Export failed.");
+ });
}}
>
Raw JSONL
@@ -2921,9 +2922,9 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
{
if (!activeThreadId) return;
- exportConversationCsv(activeThreadId).catch(() =>
- toast.error("Export failed."),
- );
+ exportConversationCsv(activeThreadId).catch((error) => {
+ if (!isDownloadCancelled(error)) toast.error("Export failed.");
+ });
}}
>
CSV
@@ -2931,9 +2932,9 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
{
if (!activeThreadId) return;
- exportConversationShareGPT(activeThreadId).catch(() =>
- toast.error("Export failed."),
- );
+ exportConversationShareGPT(activeThreadId).catch((error) => {
+ if (!isDownloadCancelled(error)) toast.error("Export failed.");
+ });
}}
>
ShareGPT JSONL
@@ -3896,6 +3897,21 @@ const EditAssistantMessageButton: FC = () => {
);
};
+async function exportMessageMarkdown(content: string): Promise {
+ try {
+ await downloadFile(
+ content,
+ `message-${Date.now()}.md`,
+ "text/markdown",
+ );
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Could not save Markdown export.", {
+ description: error instanceof Error ? error.message : String(error),
+ });
+ }
+ }
+}
const AssistantActionBar: FC = () => {
const { forkMessage, forkDisabled } = useForkMessageAction();
const [detailsOpen, setDetailsOpen] = useState(false);
@@ -3964,7 +3980,10 @@ const AssistantActionBar: FC = () => {
Fork in new chat
-
+
{
if (!enabled) {
@@ -232,10 +232,10 @@ export function WindowTitlebar({
)}
aria-label="Window titlebar"
>
- {showSidebarSurface && (
+ {showSidebarSurface && pinned && (
)}
@@ -246,7 +246,7 @@ export function WindowTitlebar({
style={{ left: contentBorderLeft, right: 0 }}
/>
)}
- {showSidebarSurface && (
+ {showSidebarSurface && pinned && (
t.id))];
await exportProjectConversations(ids, fmt, project.name);
- } catch {
- toast.error("Export failed.");
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Export failed.");
+ }
}
}
@@ -215,8 +253,10 @@ export function ProjectsPage() {
} else {
await exportBulkConversationsSeparate(ids, fmt, basename);
}
- } catch {
- toast.error("Export failed.");
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Export failed.");
+ }
}
}
@@ -294,7 +334,7 @@ export function ProjectsPage() {
- globalImportRef.current?.click()}>
+ void selectGlobalImportFile()}>
Import chats…
@@ -470,7 +510,7 @@ export function ProjectsPage() {
{
e.stopPropagation();
- projectImportRefs.current.get(project.id)?.click();
+ void selectProjectImportFile(project.id);
}}
>
diff --git a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
index 68f24a7b08..09c4944a14 100644
--- a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
+++ b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
@@ -10,6 +10,8 @@ import {
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
+import { downloadFile, isDownloadCancelled } from "@/lib/native-files";
+
import { cn } from "@/lib/utils";
import { Search01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
@@ -70,77 +72,73 @@ function sanitizeFilename(name: string): string {
return name.replace(/[\\/:*?"<>|]/g, "_").slice(0, 80) || "export";
}
-function downloadBlob(content: string | Blob, filename: string, mimeType: string): void {
- const blob = content instanceof Blob ? content : new Blob([content], { type: mimeType });
- const url = URL.createObjectURL(blob);
- const a = document.createElement("a");
- a.href = url;
- a.download = filename;
- document.body.appendChild(a);
- a.click();
- document.body.removeChild(a);
- URL.revokeObjectURL(url);
+async function downloadBlob(
+ content: string | Blob | Uint8Array,
+ filename: string,
+ mimeType: string,
+): Promise {
+ return downloadFile(content, filename, mimeType);
}
function csvEscape(val: string): string {
return `"${val.replace(/"/g, '""')}"`;
}
-function exportPromptJsonl(entry: PromptEntry): void {
- downloadBlob(
+function exportPromptJsonl(entry: PromptEntry): Promise {
+ return downloadBlob(
JSON.stringify({ name: entry.name, text: entry.text }),
`${sanitizeFilename(entry.name)}.jsonl`,
"application/x-ndjson",
);
}
-function exportPromptCsv(entry: PromptEntry): void {
- downloadBlob(
+function exportPromptCsv(entry: PromptEntry): Promise {
+ return downloadBlob(
`name,text\n${csvEscape(entry.name)},${csvEscape(entry.text)}`,
`${sanitizeFilename(entry.name)}.csv`,
"text/csv",
);
}
-function exportAllPromptsJsonl(entries: PromptEntry[]): void {
+function exportAllPromptsJsonl(entries: PromptEntry[]): Promise {
const lines = entries.map((e) => JSON.stringify({ name: e.name, text: e.text })).join("\n");
- downloadBlob(lines, "prompts.jsonl", "application/x-ndjson");
+ return downloadBlob(lines, "prompts.jsonl", "application/x-ndjson");
}
-function exportAllPromptsCsv(entries: PromptEntry[]): void {
+function exportAllPromptsCsv(entries: PromptEntry[]): Promise {
const rows = entries.map((e) => `${csvEscape(e.name)},${csvEscape(e.text)}`).join("\n");
- downloadBlob(`name,text\n${rows}`, "prompts.csv", "text/csv");
+ return downloadBlob(`name,text\n${rows}`, "prompts.csv", "text/csv");
}
-function exportListJsonl(entry: PromptListEntry): void {
- downloadBlob(
+function exportListJsonl(entry: PromptListEntry): Promise {
+ return downloadBlob(
JSON.stringify({ name: entry.name, items: entry.items }),
`${sanitizeFilename(entry.name)}.jsonl`,
"application/x-ndjson",
);
}
-function exportAllListsJsonl(entries: PromptListEntry[]): void {
+function exportAllListsJsonl(entries: PromptListEntry[]): Promise {
const lines = entries.map((e) => JSON.stringify({ name: e.name, items: e.items })).join("\n");
- downloadBlob(lines, "prompt-lists.jsonl", "application/x-ndjson");
+ return downloadBlob(lines, "prompt-lists.jsonl", "application/x-ndjson");
}
-function exportListCsv(entry: PromptListEntry): void {
+function exportListCsv(entry: PromptListEntry): Promise {
const rows = entry.items
.map((text, i) => `${csvEscape(entry.name)},${i + 1},${csvEscape(text)}`)
.join("\n");
- downloadBlob(
+ return downloadBlob(
`list_name,order,prompt_text\n${rows}`,
`${sanitizeFilename(entry.name)}.csv`,
"text/csv",
);
}
-function exportAllListsCsv(entries: PromptListEntry[]): void {
+function exportAllListsCsv(entries: PromptListEntry[]): Promise {
const rows = entries
.flatMap((e) => e.items.map((text, i) => `${csvEscape(e.name)},${i + 1},${csvEscape(text)}`))
.join("\n");
- downloadBlob(`list_name,order,prompt_text\n${rows}`, "prompt-lists.csv", "text/csv");
+ return downloadBlob(`list_name,order,prompt_text\n${rows}`, "prompt-lists.csv", "text/csv");
}
function contentBlocksToText(content: unknown): string {
@@ -363,7 +361,7 @@ export async function exportConversationShareGPT(threadId: string): Promise messageToOpenAI(msg));
if (oaiMsgs.length === 0) { toast.info("No exportable content."); return; }
- downloadBlob(
+ await downloadBlob(
JSON.stringify({ messages: oaiMsgs }),
"conversation-" + exportTs() + ".jsonl",
"application/x-ndjson",
@@ -397,7 +395,11 @@ export async function exportConversationCsv(threadId: string): Promise {
}
if (rows.length <= 1) { toast.info("No exportable content."); return; }
- downloadBlob(rows.join("\n"), "conversation-" + exportTs() + ".csv", "text/csv");
+ await downloadBlob(
+ rows.join("\n"),
+ "conversation-" + exportTs() + ".csv",
+ "text/csv",
+ );
}
export type ConvExportFormat = "jsonl-raw" | "csv" | "sharegpt";
@@ -479,7 +481,11 @@ export async function exportBulkConversationsMerged(
? header + "\n" + parts.join("\n")
: parts.join("\n");
- downloadBlob(body, `${basename}.${exportExt(format)}`, exportMime(format));
+ await downloadBlob(
+ body,
+ `${basename}.${exportExt(format)}`,
+ exportMime(format),
+ );
}
export async function exportBulkConversationsSeparate(
@@ -504,11 +510,7 @@ export async function exportBulkConversationsSeparate(
if (Object.keys(files).length === 0) { toast.info("No exportable content."); return; }
const zipped = zipSync(files);
- downloadBlob(
- new Blob([zipped], { type: "application/zip" }),
- `${basename}.zip`,
- "application/zip",
- );
+ await downloadBlob(zipped, `${basename}.zip`, "application/zip");
}
// Scope-level bulk export shared by the sidebar Recents menu and
@@ -536,8 +538,10 @@ export async function bulkExportConversationsByScope(
} else {
await exportBulkConversationsSeparate(ids, format, basename);
}
- } catch {
- toast.error("Export failed.");
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Export failed.");
+ }
}
}
@@ -749,7 +753,7 @@ export async function exportFineTuneJsonl(
return 0;
}
const suffix = format === "openai" ? "" : `-${format}`;
- downloadBlob(
+ await downloadBlob(
lines.join("\n"),
`chat-finetune${suffix}-${exportTs()}.jsonl`,
"application/x-ndjson",
@@ -992,21 +996,21 @@ export async function importConversationsFromFile(
// ShareGPT training exports: prompt → one record (human turn + empty gpt slot);
// list → one multi-turn record, each item a human turn.
-function exportPromptTrainingJsonl(entry: PromptEntry): void {
+function exportPromptTrainingJsonl(entry: PromptEntry): Promise {
const record = {
conversations: [
{ from: "human", value: entry.text },
{ from: "gpt", value: "" },
],
};
- downloadBlob(
+ return downloadBlob(
JSON.stringify(record),
`${sanitizeFilename(entry.name)}-training.jsonl`,
"application/x-ndjson",
);
}
-function exportPromptsTrainingJsonl(entries: PromptEntry[]): void {
+function exportPromptsTrainingJsonl(entries: PromptEntry[]): Promise {
const lines = entries
.map((e) =>
JSON.stringify({
@@ -1017,22 +1021,22 @@ function exportPromptsTrainingJsonl(entries: PromptEntry[]): void {
}),
)
.join("\n");
- downloadBlob(lines, "prompts-training.jsonl", "application/x-ndjson");
+ return downloadBlob(lines, "prompts-training.jsonl", "application/x-ndjson");
}
-function exportListTrainingJsonl(entry: PromptListEntry): void {
+function exportListTrainingJsonl(entry: PromptListEntry): Promise {
const conversations = entry.items.flatMap((text) => [
{ from: "human", value: text },
{ from: "gpt", value: "" },
]);
- downloadBlob(
+ return downloadBlob(
JSON.stringify({ conversations }),
`${sanitizeFilename(entry.name)}-training.jsonl`,
"application/x-ndjson",
);
}
-function exportListsTrainingJsonl(entries: PromptListEntry[]): void {
+function exportListsTrainingJsonl(entries: PromptListEntry[]): Promise {
const lines = entries
.map((e) => {
const conversations = e.items.flatMap((text) => [
@@ -1042,7 +1046,7 @@ function exportListsTrainingJsonl(entries: PromptListEntry[]): void {
return JSON.stringify({ conversations });
})
.join("\n");
- downloadBlob(lines, "prompt-lists-training.jsonl", "application/x-ndjson");
+ return downloadBlob(lines, "prompt-lists-training.jsonl", "application/x-ndjson");
}
// RFC 4180 CSV parser: handles quoted fields with embedded newlines/commas.
@@ -1268,38 +1272,44 @@ function ExportModal({
if (!csvAvailable) setFormat("jsonl");
}, [csvAvailable]);
- const handleExport = useCallback(() => {
- if (ctx.kind === "prompt") {
- if (scope === "training") exportPromptTrainingJsonl(ctx.entry);
- else if (format === "csv") exportPromptCsv(ctx.entry);
- else exportPromptJsonl(ctx.entry);
- } else if (ctx.kind === "list") {
- if (scope === "training") exportListTrainingJsonl(ctx.entry);
- else if (format === "csv") exportListCsv(ctx.entry);
- else exportListJsonl(ctx.entry);
- } else {
- const { tab, prompts, lists } = ctx;
- if (scope === "training") {
- if (tab === "prompts") {
- if (prompts.length === 0) { toast.info("No prompts to export"); return; }
- exportPromptsTrainingJsonl(prompts);
- } else {
- if (lists.length === 0) { toast.info("No prompt lists to export"); return; }
- exportListsTrainingJsonl(lists);
- }
+ const handleExport = useCallback(async () => {
+ try {
+ if (ctx.kind === "prompt") {
+ if (scope === "training") await exportPromptTrainingJsonl(ctx.entry);
+ else if (format === "csv") await exportPromptCsv(ctx.entry);
+ else await exportPromptJsonl(ctx.entry);
+ } else if (ctx.kind === "list") {
+ if (scope === "training") await exportListTrainingJsonl(ctx.entry);
+ else if (format === "csv") await exportListCsv(ctx.entry);
+ else await exportListJsonl(ctx.entry);
} else {
- if (tab === "prompts") {
+ const { tab, prompts, lists } = ctx;
+ if (scope === "training") {
+ if (tab === "prompts") {
+ if (prompts.length === 0) { toast.info("No prompts to export"); return; }
+ await exportPromptsTrainingJsonl(prompts);
+ } else {
+ if (lists.length === 0) { toast.info("No prompt lists to export"); return; }
+ await exportListsTrainingJsonl(lists);
+ }
+ } else if (tab === "prompts") {
if (prompts.length === 0) { toast.info("No prompts to export"); return; }
- if (format === "csv") exportAllPromptsCsv(prompts);
- else exportAllPromptsJsonl(prompts);
+ if (format === "csv") await exportAllPromptsCsv(prompts);
+ else await exportAllPromptsJsonl(prompts);
} else {
if (lists.length === 0) { toast.info("No prompt lists to export"); return; }
- if (format === "csv") exportAllListsCsv(lists);
- else exportAllListsJsonl(lists);
+ if (format === "csv") await exportAllListsCsv(lists);
+ else await exportAllListsJsonl(lists);
}
}
+ onClose();
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Could not save export.", {
+ description: error instanceof Error ? error.message : String(error),
+ });
+ }
}
- onClose();
}, [ctx, scope, format, onClose]);
const singleLabel =
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 70dac70a23..a72d21aa6b 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -36,6 +36,7 @@ import {
} from "@/features/settings/stores/voice-settings-store";
import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
import { isTauri } from "@/lib/api-base";
+import { isDownloadCancelled } from "@/lib/native-files";
import { isMultimodalResponse } from "./types/api";
import { getImageInputUnavailableReason } from "./utils/image-input-support";
import { useAui } from "@assistant-ui/react";
@@ -1586,9 +1587,13 @@ export function SharedComposer({
toast.error("No conversation to export yet.");
return;
}
- Promise.all(exportThreadIds.map((id) => fn(id))).catch(() =>
- toast.error("Export failed."),
- );
+ (async () => {
+ for (const id of exportThreadIds) {
+ await fn(id);
+ }
+ })().catch((error) => {
+ if (!isDownloadCancelled(error)) toast.error("Export failed.");
+ });
}}
>
{label}
diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx
index 9e65fc5749..f85c74eb86 100644
--- a/studio/frontend/src/features/chat/thread-sidebar.tsx
+++ b/studio/frontend/src/features/chat/thread-sidebar.tsx
@@ -42,6 +42,7 @@ import {
PencilEdit02Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
+import { isDownloadCancelled } from "@/lib/native-files";
import { toast } from "sonner";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import type { ChatView } from "./types";
@@ -133,9 +134,13 @@ export function ThreadSidebar({
) {
try {
const ids = await getThreadIdsForItem(item);
- await Promise.all(ids.map((id) => fn(id)));
- } catch {
- toast.error("Export failed.");
+ for (const id of ids) {
+ await fn(id);
+ }
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Export failed.");
+ }
}
}
@@ -162,8 +167,10 @@ export function ThreadSidebar({
} else {
await exportBulkConversationsSeparate(ids, fmt, basename);
}
- } catch {
- toast.error("Export failed.");
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Export failed.");
+ }
}
}
diff --git a/studio/frontend/src/features/chat/utils/export-chat-history.ts b/studio/frontend/src/features/chat/utils/export-chat-history.ts
index b4bc64d053..6f641a9fa9 100644
--- a/studio/frontend/src/features/chat/utils/export-chat-history.ts
+++ b/studio/frontend/src/features/chat/utils/export-chat-history.ts
@@ -1,9 +1,10 @@
// 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 { downloadFile } from "@/lib/native-files";
+
import { filterArchivedChatExport } from "./archived-chat-export";
import { buildStoredChatExport } from "./chat-history-storage";
-import { triggerJsonDownload } from "./download-json";
export const buildChatExport = buildStoredChatExport;
@@ -14,7 +15,11 @@ function dateStamp(): string {
export async function downloadChatExport(): Promise {
const data = await buildChatExport();
- triggerJsonDownload(data, `unsloth-chats-${dateStamp()}.json`);
+ await downloadFile(
+ JSON.stringify(data, null, 2),
+ `unsloth-chats-${dateStamp()}.json`,
+ "application/json",
+ );
}
// Full backup restricted to archived chats. Returns the archived thread count.
@@ -29,6 +34,10 @@ export async function downloadArchivedChatExport(): Promise {
if (archivedCount === 0) {
return 0;
}
- triggerJsonDownload(data, `unsloth-archived-chats-${dateStamp()}.json`);
+ await downloadFile(
+ JSON.stringify(data, null, 2),
+ `unsloth-archived-chats-${dateStamp()}.json`,
+ "application/json",
+ );
return archivedCount;
}
diff --git a/studio/frontend/src/features/settings/tabs/data-tab.tsx b/studio/frontend/src/features/settings/tabs/data-tab.tsx
index dc9e707ea0..e2b524f46e 100644
--- a/studio/frontend/src/features/settings/tabs/data-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/data-tab.tsx
@@ -39,6 +39,9 @@ import {
useChatSidebarItems,
} from "@/features/chat";
import { useT } from "@/i18n";
+
+import { isTauri } from "@/lib/api-base";
+import { isDownloadCancelled, pickNativeChatImport } from "@/lib/native-files";
import {
ChevronDownStandardIcon,
ChevronRightStandardIcon,
@@ -147,6 +150,12 @@ export function DataTab() {
setExporting(true);
try {
await downloadChatExport();
+ } catch (error) {
+ if (!isDownloadCancelled(error)) {
+ toast.error("Could not export chats", {
+ description: error instanceof Error ? error.message : String(error),
+ });
+ }
} finally {
setExporting(false);
}
@@ -164,9 +173,11 @@ export function DataTab() {
: t("settings.data.exportedArchivedChatCount", { count: exported }),
);
} catch (error) {
- toast.error(t("settings.data.failedToExportArchivedChats"), {
- description: error instanceof Error ? error.message : undefined,
- });
+ if (!isDownloadCancelled(error)) {
+ toast.error(t("settings.data.failedToExportArchivedChats"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ }
} finally {
setArchivedExporting(false);
}
@@ -191,6 +202,24 @@ export function DataTab() {
}
};
+ const handleImportClick = async () => {
+ if (!isTauri) {
+ importInputRef.current?.click();
+ return;
+ }
+ try {
+ const selected = await pickNativeChatImport();
+ if (!selected) {
+ return;
+ }
+ await handleImport(new File([selected.content], selected.name));
+ } catch (error) {
+ toast.error(t("settings.chat.importFailed"), {
+ description: error instanceof Error ? error.message : String(error),
+ });
+ }
+ };
+
const handleArchiveAll = async () => {
setArchiving(true);
try {
@@ -219,9 +248,11 @@ export function DataTab() {
try {
await exportFineTuneJsonl(fineTuneFormat);
} catch (error) {
- toast.error(t("settings.data.fineTuneExportFailed"), {
- description: error instanceof Error ? error.message : undefined,
- });
+ if (!isDownloadCancelled(error)) {
+ toast.error(t("settings.data.fineTuneExportFailed"), {
+ description: error instanceof Error ? error.message : undefined,
+ });
+ }
} finally {
setFineTuneExporting(false);
}
@@ -632,7 +663,7 @@ export function DataTab() {
importInputRef.current?.click()}
+ onClick={() => void handleImportClick()}
>
{t("settings.chat.importChatsAction")}
diff --git a/studio/frontend/src/lib/native-files.ts b/studio/frontend/src/lib/native-files.ts
new file mode 100644
index 0000000000..48a6bd4ab7
--- /dev/null
+++ b/studio/frontend/src/lib/native-files.ts
@@ -0,0 +1,86 @@
+// 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 { isTauri } from "@/lib/api-base";
+
+const NATIVE_FILE_NAME_HEADER = "x-unsloth-default-name";
+export class DownloadCancelledError extends Error {
+ constructor() {
+ super("Save cancelled.");
+ this.name = "DownloadCancelledError";
+ }
+}
+
+export function isDownloadCancelled(error: unknown): boolean {
+ return error instanceof DownloadCancelledError;
+}
+
+function encodeNativeFilename(filename: string): string {
+ const bytes = new TextEncoder().encode(filename);
+ let binary = "";
+ for (const byte of bytes) {
+ binary += String.fromCharCode(byte);
+ }
+ return btoa(binary);
+}
+
+export interface NativeChatImport {
+ name: string;
+ content: string;
+}
+
+function browserDownload(blob: Blob, filename: string): void {
+ const url = URL.createObjectURL(blob);
+ const anchor = document.createElement("a");
+ anchor.href = url;
+ anchor.download = filename;
+ document.body.appendChild(anchor);
+ anchor.click();
+ anchor.remove();
+ URL.revokeObjectURL(url);
+}
+
+/** Save through a native chooser in Tauri and retain normal downloads on web. */
+export async function downloadFile(
+ content: string | Blob | Uint8Array,
+ filename: string,
+ mimeType = "application/octet-stream",
+): Promise {
+ if (isTauri) {
+ const { invoke } = await import("@tauri-apps/api/core");
+ const bytes =
+ typeof content === "string"
+ ? new TextEncoder().encode(content)
+ : content instanceof Blob
+ ? new Uint8Array(await content.arrayBuffer())
+ : content;
+ const savedPath = await invoke("save_native_file", bytes, {
+ headers: {
+ [NATIVE_FILE_NAME_HEADER]: encodeNativeFilename(filename),
+ },
+ });
+ if (savedPath === null) {
+ throw new DownloadCancelledError();
+ }
+ return;
+ }
+
+ const browserContent =
+ content instanceof Uint8Array ? Uint8Array.from(content).buffer : content;
+ const blob =
+ browserContent instanceof Blob
+ ? browserContent
+ : new Blob([browserContent], { type: mimeType });
+
+ browserDownload(blob, filename);
+ return;
+}
+
+/** Open the bounded native chat-import picker. Cancellation returns null. */
+export async function pickNativeChatImport(): Promise {
+ if (!isTauri) {
+ return null;
+ }
+ const { invoke } = await import("@tauri-apps/api/core");
+ return invoke("pick_native_chat_import");
+}
diff --git a/studio/src-tauri/Cargo.lock b/studio/src-tauri/Cargo.lock
index 183cfc1cc8..604bb01525 100644
--- a/studio/src-tauri/Cargo.lock
+++ b/studio/src-tauri/Cargo.lock
@@ -5592,6 +5592,7 @@ dependencies = [
"tauri-plugin-single-instance",
"tauri-plugin-updater",
"tauri-plugin-window-state",
+ "tempfile",
"tokio",
"windows 0.62.2",
"windows-sys 0.61.2",
diff --git a/studio/src-tauri/Cargo.toml b/studio/src-tauri/Cargo.toml
index 438558fe77..826509ce9a 100644
--- a/studio/src-tauri/Cargo.toml
+++ b/studio/src-tauri/Cargo.toml
@@ -29,6 +29,7 @@ tauri-plugin-clipboard-manager = "2"
tauri-plugin-dialog = "2"
rand = "0.10.0"
tauri-plugin-notification = "2.3.3"
+tempfile = "3"
tauri-plugin-window-state = "2"
[target.'cfg(unix)'.dependencies]
diff --git a/studio/src-tauri/Entitlements.plist b/studio/src-tauri/Entitlements.plist
index 10e6df3427..71b1d2c0f4 100644
--- a/studio/src-tauri/Entitlements.plist
+++ b/studio/src-tauri/Entitlements.plist
@@ -8,6 +8,9 @@
com.apple.security.network.client
+
+ com.apple.security.device.audio-input
+
com.apple.security.cs.disable-library-validation
diff --git a/studio/src-tauri/Info.plist b/studio/src-tauri/Info.plist
new file mode 100644
index 0000000000..e7a280647b
--- /dev/null
+++ b/studio/src-tauri/Info.plist
@@ -0,0 +1,8 @@
+
+
+
+
+ NSMicrophoneUsageDescription
+ Unsloth Studio uses the microphone you select for local voice dictation.
+
+
diff --git a/studio/src-tauri/src/app_layout.rs b/studio/src-tauri/src/app_layout.rs
new file mode 100644
index 0000000000..9b9277b48d
--- /dev/null
+++ b/studio/src-tauri/src/app_layout.rs
@@ -0,0 +1,274 @@
+use std::fs;
+use std::path::{Path, PathBuf};
+use tauri::{Manager, WebviewWindow};
+
+use tauri_plugin_window_state::AppHandleExt;
+
+const APP_LAYOUT_MARKER_FILE: &str = "app-layout-initialized-v1";
+const INITIALIZED_MARKER: &[u8] = b"initialized\n";
+const RESET_MARKER: &[u8] = b"reset\n";
+
+const SETUP_WINDOW_WIDTH: f64 = 760.0;
+const SETUP_WINDOW_HEIGHT: f64 = 560.0;
+const MIN_REASONABLE_WINDOW_WIDTH: u64 = 320;
+const MIN_REASONABLE_WINDOW_HEIGHT: u64 = 240;
+
+const SETUP_SIZE_TOLERANCE_PX: f64 = 2.0;
+
+#[derive(serde::Deserialize)]
+struct PersistedWindowState {
+ width: u32,
+ height: u32,
+ #[serde(rename = "x")]
+ _x: i32,
+ #[serde(rename = "y")]
+ _y: i32,
+ #[serde(rename = "prev_x")]
+ _prev_x: i32,
+ #[serde(rename = "prev_y")]
+ _prev_y: i32,
+ #[serde(rename = "maximized")]
+ _maximized: bool,
+ #[serde(rename = "visible")]
+ _visible: bool,
+ #[serde(rename = "decorated")]
+ _decorated: bool,
+ #[serde(rename = "fullscreen")]
+ _fullscreen: bool,
+}
+
+fn marker_path(config_dir: &Path) -> PathBuf {
+ config_dir.join(APP_LAYOUT_MARKER_FILE)
+}
+
+fn marker_matches(config_dir: &Path, expected: &[u8]) -> bool {
+ fs::read(marker_path(config_dir))
+ .map(|contents| contents == expected)
+ .unwrap_or(false)
+}
+
+fn is_initialized(config_dir: &Path) -> bool {
+ marker_matches(config_dir, INITIALIZED_MARKER)
+}
+
+fn is_reset_pending(config_dir: &Path) -> bool {
+ marker_matches(config_dir, RESET_MARKER)
+}
+
+fn write_marker(config_dir: &Path, contents: &[u8]) -> Result<(), String> {
+ fs::create_dir_all(config_dir).map_err(|error| {
+ format!(
+ "Failed to create app configuration directory {}: {error}",
+ config_dir.display()
+ )
+ })?;
+ let path = marker_path(config_dir);
+ fs::write(&path, contents).map_err(|error| {
+ format!(
+ "Failed to write app layout marker {}: {error}",
+ path.display()
+ )
+ })
+}
+
+fn reset_initialized(config_dir: &Path) -> Result<(), String> {
+ write_marker(config_dir, RESET_MARKER)
+}
+
+fn is_setup_window_size(width: u64, height: u64) -> bool {
+ // Window-state persists physical dimensions but not their source scale.
+ let width_scale = width as f64 / SETUP_WINDOW_WIDTH;
+ let height_scale = height as f64 / SETUP_WINDOW_HEIGHT;
+ let scale_tolerance = SETUP_SIZE_TOLERANCE_PX / SETUP_WINDOW_WIDTH
+ + SETUP_SIZE_TOLERANCE_PX / SETUP_WINDOW_HEIGHT;
+ (width_scale - height_scale).abs() <= scale_tolerance
+}
+
+fn saved_main_window_size(config_dir: &Path, state_file_name: &str) -> Option<(u64, u64)> {
+ let contents = fs::read(config_dir.join(state_file_name)).ok()?;
+ let states = serde_json::from_slice::>(
+ &contents,
+ )
+ .ok()?;
+ let main = states.get("main")?;
+ let width = u64::from(main.width);
+ let height = u64::from(main.height);
+ if width < MIN_REASONABLE_WINDOW_WIDTH || height < MIN_REASONABLE_WINDOW_HEIGHT {
+ return None;
+ }
+ Some((width, height))
+}
+
+fn should_restore_saved_layout(config_dir: &Path, state_file_name: &str) -> bool {
+ if is_reset_pending(config_dir) {
+ return false;
+ }
+ let Some((width, height)) = saved_main_window_size(config_dir, state_file_name) else {
+ return false;
+ };
+ is_initialized(config_dir) || !is_setup_window_size(width, height)
+}
+
+fn mark_initialized(config_dir: &Path) -> Result<(), String> {
+ write_marker(config_dir, INITIALIZED_MARKER)
+}
+
+fn app_config_dir(app: &tauri::AppHandle) -> Result {
+ app.path()
+ .app_config_dir()
+ .map_err(|error| format!("Could not determine app configuration directory: {error}"))
+}
+
+/// Returns whether a full-app layout has previously completed. Legacy state is
+/// migrated unless it matches the fixed setup-window size at any display scale.
+#[tauri::command]
+pub fn has_initialized_app_window_layout(
+ window: WebviewWindow,
+ app: tauri::AppHandle,
+) -> Result {
+ crate::native_intents::ensure_main_window(&window)?;
+ let config_dir = app_config_dir(&app)?;
+ Ok(should_restore_saved_layout(&config_dir, &app.filename()))
+}
+
+/// Persist only after the caller has successfully sized/centered or restored,
+/// shown, constrained, and minimum-sized the full application window.
+#[tauri::command]
+pub fn mark_app_window_layout_initialized(
+ window: WebviewWindow,
+ app: tauri::AppHandle,
+) -> Result<(), String> {
+ crate::native_intents::ensure_main_window(&window)?;
+ mark_initialized(&app_config_dir(&app)?)
+}
+
+/// Force the next full-app transition to use a monitor-safe layout. Setup can
+/// overwrite the plugin's saved full-app dimensions before the process exits.
+#[tauri::command]
+pub fn reset_app_window_layout_initialized(
+ window: WebviewWindow,
+ app: tauri::AppHandle,
+) -> Result<(), String> {
+ crate::native_intents::ensure_main_window(&window)?;
+ reset_initialized(&app_config_dir(&app)?)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::time::{SystemTime, UNIX_EPOCH};
+
+ fn temp_dir(name: &str) -> PathBuf {
+ let nanos = SystemTime::now()
+ .duration_since(UNIX_EPOCH)
+ .unwrap()
+ .as_nanos();
+ std::env::temp_dir().join(format!(
+ "unsloth-app-layout-{name}-{}-{nanos}",
+ std::process::id()
+ ))
+ }
+
+ fn window_state(width: u32, height: u32, maximized: bool) -> String {
+ serde_json::json!({
+ "main": {
+ "width": width,
+ "height": height,
+ "x": 0,
+ "y": 0,
+ "prev_x": 0,
+ "prev_y": 0,
+ "maximized": maximized,
+ "visible": true,
+ "decorated": true,
+ "fullscreen": false
+ }
+ })
+ .to_string()
+ }
+
+ #[test]
+ fn marker_is_absent_until_full_layout_is_marked() {
+ let dir = temp_dir("transition");
+ assert!(!is_initialized(&dir));
+ mark_initialized(&dir).unwrap();
+ assert!(is_initialized(&dir));
+ assert_eq!(fs::read(marker_path(&dir)).unwrap(), b"initialized\n");
+ let _ = fs::remove_dir_all(dir);
+ }
+
+ #[test]
+ fn marker_write_and_reset_are_idempotent() {
+ let dir = temp_dir("idempotent");
+ mark_initialized(&dir).unwrap();
+ mark_initialized(&dir).unwrap();
+ assert!(is_initialized(&dir));
+ assert!(!is_reset_pending(&dir));
+
+ reset_initialized(&dir).unwrap();
+ reset_initialized(&dir).unwrap();
+ assert!(!is_initialized(&dir));
+ assert!(is_reset_pending(&dir));
+
+ mark_initialized(&dir).unwrap();
+ assert!(is_initialized(&dir));
+ assert!(!is_reset_pending(&dir));
+ let _ = fs::remove_dir_all(dir);
+ }
+
+ #[test]
+ fn legacy_full_app_state_is_migrated_but_setup_state_at_any_scale_is_not() {
+ let dir = temp_dir("legacy-plugin-state");
+ fs::create_dir_all(&dir).unwrap();
+ let state_file = ".window-state.json";
+ let state_path = dir.join(state_file);
+
+ fs::write(&state_path, window_state(1200, 800, false)).unwrap();
+ assert!(should_restore_saved_layout(&dir, state_file));
+
+ fs::write(&state_path, window_state(760, 560, false)).unwrap();
+ assert!(!should_restore_saved_layout(&dir, state_file));
+
+ fs::write(&state_path, window_state(950, 700, false)).unwrap();
+ assert!(!should_restore_saved_layout(&dir, state_file));
+
+ fs::write(&state_path, window_state(1520, 1120, false)).unwrap();
+ assert!(!should_restore_saved_layout(&dir, state_file));
+
+ fs::write(&state_path, window_state(1520, 1120, true)).unwrap();
+ assert!(!should_restore_saved_layout(&dir, state_file));
+ let _ = fs::remove_dir_all(dir);
+ }
+
+ #[test]
+ fn explicit_reset_blocks_legacy_migration_until_app_layout_completes() {
+ let dir = temp_dir("reset-migration");
+ fs::create_dir_all(&dir).unwrap();
+ let state_file = ".window-state.json";
+ fs::write(dir.join(state_file), window_state(1200, 800, false)).unwrap();
+
+ mark_initialized(&dir).unwrap();
+ assert!(should_restore_saved_layout(&dir, state_file));
+ reset_initialized(&dir).unwrap();
+ assert!(!should_restore_saved_layout(&dir, state_file));
+ mark_initialized(&dir).unwrap();
+ assert!(should_restore_saved_layout(&dir, state_file));
+ let _ = fs::remove_dir_all(dir);
+ }
+
+ #[test]
+ fn corrupt_plugin_state_does_not_restore_even_with_initialized_marker() {
+ let dir = temp_dir("plugin-state");
+ fs::create_dir_all(&dir).unwrap();
+ let state_file = ".window-state.json";
+ fs::write(
+ dir.join(state_file),
+ r#"{"main":{"width":1200,"height":800}}"#,
+ )
+ .unwrap();
+ mark_initialized(&dir).unwrap();
+ assert!(is_initialized(&dir));
+ assert!(!should_restore_saved_layout(&dir, state_file));
+ let _ = fs::remove_dir_all(dir);
+ }
+}
diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs
index 5d41cb9217..72ea9d6985 100644
--- a/studio/src-tauri/src/commands.rs
+++ b/studio/src-tauri/src/commands.rs
@@ -339,11 +339,21 @@ pub fn get_server_logs(state: tauri::State<'_, BackendState>) -> Vec {
/// Open an existing directory in the system file manager. Validates the path
/// up front so callers get a clean error instead of a raw OS failure.
-fn open_existing_dir(dir: &std::path::Path) -> Result<(), String> {
+fn open_existing_dir_with(
+ dir: &std::path::Path,
+ opener: impl FnOnce(&std::path::Path) -> Result<(), E>,
+) -> Result<(), String>
+where
+ E: std::fmt::Display,
+{
if !dir.is_dir() {
return Err(format!("Directory does not exist: {}", dir.display()));
}
- open::that(dir).map_err(|e| format!("Failed to open directory: {}", e))
+ opener(dir).map_err(|error| format!("Failed to open directory: {error}"))
+}
+
+fn open_existing_dir(dir: &std::path::Path) -> Result<(), String> {
+ open_existing_dir_with(dir, |path| open::that_detached(path))
}
/// Open the Unsloth Studio directory in the system file manager.
@@ -641,7 +651,8 @@ pub async fn start_managed_repair(
#[cfg(test)]
mod tests {
- use std::time::Duration;
+ use std::fs;
+ use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
@@ -699,6 +710,40 @@ mod tests {
port
}
+ #[test]
+ fn existing_directory_helper_invokes_opener_and_surfaces_errors() {
+ let nanos = SystemTime::now()
+ .duration_since(UNIX_EPOCH)
+ .unwrap()
+ .as_nanos();
+ let dir =
+ std::env::temp_dir().join(format!("unsloth-open-dir-{}-{nanos}", std::process::id()));
+ fs::create_dir_all(&dir).unwrap();
+ let mut opened = false;
+ super::open_existing_dir_with(&dir, |path| {
+ opened = true;
+ assert_eq!(path, dir);
+ Ok::<_, &str>(())
+ })
+ .unwrap();
+ assert!(opened);
+
+ let error = super::open_existing_dir_with(&dir, |_| Err("opener failed")).unwrap_err();
+ assert!(error.contains("opener failed"));
+ let _ = fs::remove_dir_all(dir);
+ }
+
+ #[test]
+ fn existing_directory_helper_rejects_missing_path_without_opening() {
+ let missing = std::env::temp_dir().join("unsloth-definitely-missing-open-dir");
+ let error = super::open_existing_dir_with(&missing, |_| {
+ panic!("opener must not run for an invalid directory");
+ #[allow(unreachable_code)]
+ Ok::<_, &str>(())
+ })
+ .unwrap_err();
+ assert!(error.contains("Directory does not exist"));
+ }
#[test]
fn repair_elevation_is_not_a_terminal_repair_failure() {
assert!(!super::should_emit_repair_failed("NEEDS_ELEVATION"));
diff --git a/studio/src-tauri/src/main.rs b/studio/src-tauri/src/main.rs
index 4724317601..cfac90ae73 100644
--- a/studio/src-tauri/src/main.rs
+++ b/studio/src-tauri/src/main.rs
@@ -1,5 +1,6 @@
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
+mod app_layout;
mod commands;
mod desktop_auth;
mod desktop_backend_owner;
@@ -7,6 +8,7 @@ mod desktop_update_policy;
mod diagnostics;
mod install;
mod native_backend_lease;
+mod native_file_dialogs;
mod native_intents;
mod native_path_policy;
mod preflight;
@@ -195,6 +197,9 @@ fn main() {
.manage(process::new_shutdown_flag())
.manage(update::new_update_state())
.invoke_handler(tauri::generate_handler![
+ app_layout::has_initialized_app_window_layout,
+ app_layout::mark_app_window_layout_initialized,
+ app_layout::reset_app_window_layout_initialized,
commands::check_install_status,
commands::desktop_preflight,
commands::start_install,
@@ -213,6 +218,8 @@ fn main() {
desktop_update_policy::check_desktop_manual_update,
desktop_update_policy::desktop_update_policy,
diagnostics::collect_support_diagnostics,
+ native_file_dialogs::save_native_file,
+ native_file_dialogs::pick_native_chat_import,
native_intents::drain_native_intents,
native_intents::register_native_model_path,
native_intents::pick_native_model,
diff --git a/studio/src-tauri/src/native_file_dialogs.rs b/studio/src-tauri/src/native_file_dialogs.rs
new file mode 100644
index 0000000000..d795eb020c
--- /dev/null
+++ b/studio/src-tauri/src/native_file_dialogs.rs
@@ -0,0 +1,320 @@
+use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
+
+use serde::Serialize;
+use std::fs::{self, File};
+use std::io::{Read, Write};
+use std::path::{Path, PathBuf};
+use tauri::{AppHandle, WebviewWindow};
+use tauri_plugin_dialog::DialogExt;
+
+const MAX_CHAT_IMPORT_BYTES: u64 = 64 * 1024 * 1024;
+const NATIVE_FILE_NAME_HEADER: &str = "x-unsloth-default-name";
+const CHAT_IMPORT_EXTENSIONS: &[&str] = &["jsonl", "ndjson", "csv"];
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct NativeImportedFile {
+ name: String,
+ content: String,
+}
+
+fn default_file_name(suggested_name: &str) -> String {
+ Path::new(suggested_name)
+ .file_name()
+ .and_then(|name| name.to_str())
+ .filter(|name| !name.is_empty() && *name != "." && *name != "..")
+ .unwrap_or("unsloth-export.json")
+ .to_string()
+}
+fn decode_default_file_name(encoded_name: &str) -> Result {
+ let bytes = BASE64
+ .decode(encoded_name)
+ .map_err(|_| "Invalid native export filename.".to_string())?;
+ let name =
+ String::from_utf8(bytes).map_err(|_| "Invalid native export filename.".to_string())?;
+ Ok(default_file_name(&name))
+}
+
+fn save_filter(file_name: &str) -> (&'static str, Vec<&'static str>) {
+ match Path::new(file_name)
+ .extension()
+ .and_then(|extension| extension.to_str())
+ .map(str::to_ascii_lowercase)
+ .as_deref()
+ {
+ Some("json") => ("JSON", vec!["json"]),
+ Some("jsonl") | Some("ndjson") => ("JSON Lines", vec!["jsonl", "ndjson"]),
+ Some("csv") => ("CSV", vec!["csv"]),
+ Some("md") | Some("markdown") => ("Markdown", vec!["md", "markdown"]),
+ Some("zip") => ("ZIP archive", vec!["zip"]),
+ _ => (
+ "Export files",
+ vec!["json", "jsonl", "ndjson", "csv", "md", "markdown", "zip"],
+ ),
+ }
+}
+
+fn local_dialog_path(path: tauri_plugin_dialog::FilePath) -> Result {
+ path.into_path()
+ .map_err(|_| "Only local filesystem paths are supported.".to_string())
+}
+
+fn save_selected_file(
+ selected_path: Option,
+ content: &[u8],
+) -> Result, String> {
+ let Some(path) = selected_path else {
+ return Ok(None);
+ };
+ let parent = path
+ .parent()
+ .filter(|parent| !parent.as_os_str().is_empty())
+ .unwrap_or_else(|| Path::new("."));
+ let mut builder = tempfile::Builder::new();
+ #[cfg(unix)]
+ {
+ use std::os::unix::fs::PermissionsExt;
+ let permissions = fs::metadata(&path)
+ .map(|metadata| metadata.permissions())
+ .unwrap_or_else(|_| fs::Permissions::from_mode(0o666));
+ builder.permissions(permissions);
+ }
+ let mut temporary = builder
+ .prefix(".unsloth-export-")
+ .tempfile_in(parent)
+ .map_err(|error| format!("Failed to prepare {}: {error}", path.display()))?;
+ temporary
+ .write_all(content)
+ .and_then(|()| temporary.as_file().sync_all())
+ .map_err(|error| format!("Failed to save {}: {error}", path.display()))?;
+ temporary
+ .persist(&path)
+ .map_err(|error| format!("Failed to save {}: {}", path.display(), error.error))?;
+ let file_name = path
+ .file_name()
+ .and_then(|name| name.to_str())
+ .unwrap_or("export")
+ .to_string();
+ Ok(Some(file_name))
+}
+
+fn read_selected_import(
+ selected_path: Option,
+) -> Result, String> {
+ let Some(path) = selected_path else {
+ return Ok(None);
+ };
+ let extension = path
+ .extension()
+ .and_then(|extension| extension.to_str())
+ .map(str::to_ascii_lowercase)
+ .ok_or_else(|| "Chat import must be a .jsonl, .ndjson, or .csv file.".to_string())?;
+ if !CHAT_IMPORT_EXTENSIONS.contains(&extension.as_str()) {
+ return Err("Chat import must be a .jsonl, .ndjson, or .csv file.".to_string());
+ }
+
+ let metadata = fs::metadata(&path)
+ .map_err(|error| format!("Failed to inspect {}: {error}", path.display()))?;
+ if !metadata.is_file() {
+ return Err(format!("Selected import is not a file: {}", path.display()));
+ }
+ if metadata.len() > MAX_CHAT_IMPORT_BYTES {
+ return Err(format!(
+ "Chat import is too large (maximum {} MiB).",
+ MAX_CHAT_IMPORT_BYTES / 1024 / 1024
+ ));
+ }
+
+ // Limit the read too, so a file that grows after metadata inspection cannot
+ // make the command allocate without bound.
+ let file =
+ File::open(&path).map_err(|error| format!("Failed to open {}: {error}", path.display()))?;
+ let mut bytes = Vec::with_capacity(metadata.len() as usize);
+ file.take(MAX_CHAT_IMPORT_BYTES + 1)
+ .read_to_end(&mut bytes)
+ .map_err(|error| format!("Failed to read {}: {error}", path.display()))?;
+ if bytes.len() as u64 > MAX_CHAT_IMPORT_BYTES {
+ return Err(format!(
+ "Chat import is too large (maximum {} MiB).",
+ MAX_CHAT_IMPORT_BYTES / 1024 / 1024
+ ));
+ }
+ let content = String::from_utf8(bytes)
+ .map_err(|_| format!("Chat import is not valid UTF-8: {}", path.display()))?;
+ let name = path
+ .file_name()
+ .and_then(|name| name.to_str())
+ .map(str::to_string)
+ .unwrap_or_else(|| format!("chat-import.{extension}"));
+ Ok(Some(NativeImportedFile { name, content }))
+}
+
+#[tauri::command]
+pub async fn save_native_file(
+ window: WebviewWindow,
+ app: AppHandle,
+ request: tauri::ipc::Request<'_>,
+) -> Result , String> {
+ crate::native_intents::ensure_main_window(&window)?;
+ let encoded_name = request
+ .headers()
+ .get(NATIVE_FILE_NAME_HEADER)
+ .ok_or_else(|| "Native export filename is missing.".to_string())?
+ .to_str()
+ .map_err(|_| "Invalid native export filename.".to_string())?;
+ let file_name = decode_default_file_name(encoded_name)?;
+ let content = match request.body() {
+ tauri::ipc::InvokeBody::Raw(content) => content,
+ _ => return Err("Native export content must be binary.".to_string()),
+ };
+ let (filter_name, extensions) = save_filter(&file_name);
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ app.dialog()
+ .file()
+ .set_title("Save Unsloth export")
+ .set_file_name(file_name)
+ .add_filter(filter_name, &extensions)
+ .save_file(move |path| {
+ let _ = tx.send(path);
+ });
+ let selected_path = rx
+ .await
+ .map_err(|_| "Save dialog closed unexpectedly.".to_string())?
+ .map(local_dialog_path)
+ .transpose()?;
+ save_selected_file(selected_path, content)
+}
+
+#[tauri::command]
+pub async fn pick_native_chat_import(
+ window: WebviewWindow,
+ app: AppHandle,
+) -> Result , String> {
+ crate::native_intents::ensure_main_window(&window)?;
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ app.dialog()
+ .file()
+ .set_title("Import chats")
+ .add_filter("Chat exports", CHAT_IMPORT_EXTENSIONS)
+ .pick_file(move |path| {
+ let _ = tx.send(path);
+ });
+ let selected_path = rx
+ .await
+ .map_err(|_| "Import dialog closed unexpectedly.".to_string())?
+ .map(local_dialog_path)
+ .transpose()?;
+ read_selected_import(selected_path)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::time::{SystemTime, UNIX_EPOCH};
+
+ fn temp_path(name: &str) -> PathBuf {
+ let nanos = SystemTime::now()
+ .duration_since(UNIX_EPOCH)
+ .unwrap()
+ .as_nanos();
+ std::env::temp_dir().join(format!(
+ "unsloth-native-files-{name}-{}-{nanos}",
+ std::process::id()
+ ))
+ }
+
+ #[test]
+ fn cancellation_is_quiet_for_save_and_import() {
+ assert!(save_selected_file(None, b"x").unwrap().is_none());
+ assert!(read_selected_import(None).unwrap().is_none());
+ }
+
+ #[test]
+ fn writes_text_and_binary_exactly() {
+ // Overwriting must stage the new content before replacing the destination.
+ let text_path = temp_path("text").with_extension("json");
+ let binary_path = temp_path("binary").with_extension("zip");
+
+ fs::write(&text_path, b"previous export").unwrap();
+ save_selected_file(Some(text_path.clone()), b"{\"ok\":true}").unwrap();
+ save_selected_file(Some(binary_path.clone()), &[0, 1, 2, 255]).unwrap();
+ assert_eq!(fs::read(&text_path).unwrap(), b"{\"ok\":true}");
+ assert_eq!(fs::read(&binary_path).unwrap(), [0, 1, 2, 255]);
+ let _ = fs::remove_file(text_path);
+ let _ = fs::remove_file(binary_path);
+ }
+
+ #[test]
+ fn markdown_exports_use_a_markdown_save_filter() {
+ assert_eq!(
+ save_filter("message.md"),
+ ("Markdown", vec!["md", "markdown"])
+ );
+ }
+
+ #[test]
+ fn reads_supported_import_and_rejects_other_extensions() {
+ let jsonl_path = temp_path("allowed").with_extension("JSONL");
+ fs::write(&jsonl_path, "{\"messages\":[]}").unwrap();
+ let imported = read_selected_import(Some(jsonl_path.clone()))
+ .unwrap()
+ .unwrap();
+ assert_eq!(imported.content, "{\"messages\":[]}");
+
+ let json_path = temp_path("unsupported").with_extension("json");
+ fs::write(&json_path, "{}").unwrap();
+ assert!(read_selected_import(Some(json_path.clone())).is_err());
+ let txt_path = temp_path("denied").with_extension("txt");
+ fs::write(&txt_path, "no").unwrap();
+ assert!(read_selected_import(Some(txt_path.clone()))
+ .unwrap_err()
+ .contains(".json"));
+ let _ = fs::remove_file(jsonl_path);
+ let _ = fs::remove_file(json_path);
+ let _ = fs::remove_file(txt_path);
+ }
+
+ #[test]
+ fn read_limit_and_utf8_errors_are_concrete() {
+ let oversized = temp_path("oversized").with_extension("csv");
+ let file = File::create(&oversized).unwrap();
+ file.set_len(MAX_CHAT_IMPORT_BYTES + 1).unwrap();
+ assert!(read_selected_import(Some(oversized.clone()))
+ .unwrap_err()
+ .contains("too large"));
+
+ let invalid = temp_path("invalid-utf8").with_extension("jsonl");
+ fs::write(&invalid, [0xff]).unwrap();
+ assert!(read_selected_import(Some(invalid.clone()))
+ .unwrap_err()
+ .contains("UTF-8"));
+ let _ = fs::remove_file(oversized);
+ let _ = fs::remove_file(invalid);
+ }
+
+ #[cfg(unix)]
+ #[test]
+ fn non_utf8_import_name_preserves_csv_extension() {
+ use std::ffi::OsString;
+ use std::os::unix::ffi::OsStringExt;
+
+ let path = std::env::temp_dir().join(OsString::from_vec(vec![
+ b'u', b'n', b's', b'l', b'o', b't', b'h', 0xff, b'.', b'c', b's', b'v',
+ ]));
+ fs::write(&path, "role,content\nuser,hello\n").unwrap();
+ let imported = read_selected_import(Some(path.clone())).unwrap().unwrap();
+ assert_eq!(imported.name, "chat-import.csv");
+ let _ = fs::remove_file(path);
+ }
+
+ #[test]
+ fn strips_directories_from_suggested_default_name() {
+ assert_eq!(default_file_name("../../chat.jsonl"), "chat.jsonl");
+ assert_eq!(default_file_name(""), "unsloth-export.json");
+
+ assert_eq!(
+ decode_default_file_name("Y2hhdC5qc29ubA==").unwrap(),
+ "chat.jsonl"
+ );
+ }
+}
diff --git a/studio/src-tauri/src/native_path_policy.rs b/studio/src-tauri/src/native_path_policy.rs
index b82e516e7a..26d4cdb61a 100644
--- a/studio/src-tauri/src/native_path_policy.rs
+++ b/studio/src-tauri/src/native_path_policy.rs
@@ -130,8 +130,8 @@ fn classify_existing_path(path: &Path) -> Result {
.canonicalize()
.map_err(|e| format!("Path could not be resolved: {e}"))?;
reject_network_or_device_path(&canonical_path)?;
- let canonical_symlink_metadata = fs::symlink_metadata(&canonical_path)
- .map_err(|e| format!("Path is not available: {e}"))?;
+ let canonical_symlink_metadata =
+ fs::symlink_metadata(&canonical_path).map_err(|e| format!("Path is not available: {e}"))?;
if canonical_symlink_metadata.file_type().is_symlink() {
return Err("Symlink paths are not supported for native intake.".to_string());
}
diff --git a/studio/src-tauri/tauri.macos.conf.json b/studio/src-tauri/tauri.macos.conf.json
index 202c9fb505..b1517a6234 100644
--- a/studio/src-tauri/tauri.macos.conf.json
+++ b/studio/src-tauri/tauri.macos.conf.json
@@ -2,6 +2,7 @@
"bundle": {
"macOS": {
"entitlements": "./Entitlements.plist",
+ "infoPlist": "./Info.plist",
"dmg": {
"appPosition": { "x": 180, "y": 220 },
"applicationFolderPosition": { "x": 480, "y": 220 }
diff --git a/tests/studio/test_desktop_reliability_frontend_contract.py b/tests/studio/test_desktop_reliability_frontend_contract.py
new file mode 100644
index 0000000000..868895b8f0
--- /dev/null
+++ b/tests/studio/test_desktop_reliability_frontend_contract.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Static contracts for focused packaged-desktop reliability behavior."""
+
+from pathlib import Path
+
+
+REPO = Path(__file__).resolve().parents[2]
+FRONTEND = REPO / "studio/frontend/src"
+NATIVE_FILES = FRONTEND / "lib/native-files.ts"
+CHAT_EXPORT = FRONTEND / "features/chat/utils/export-chat-history.ts"
+DATA_TAB = FRONTEND / "features/settings/tabs/data-tab.tsx"
+PROMPT_STORAGE = FRONTEND / "features/chat/prompt-storage/prompt-storage-dialog.tsx"
+
+APP_SIDEBAR = FRONTEND / "components/app-sidebar.tsx"
+THREAD = FRONTEND / "components/assistant-ui/thread.tsx"
+THREAD_SIDEBAR = FRONTEND / "features/chat/thread-sidebar.tsx"
+SHARED_COMPOSER = FRONTEND / "features/chat/shared-composer.tsx"
+TITLEBAR = FRONTEND / "components/tauri/window-titlebar.tsx"
+NATIVE_DIALOGS = REPO / "studio/src-tauri/src/native_file_dialogs.rs"
+
+
+APP_PROVIDER = FRONTEND / "app/provider.tsx"
+
+
+def test_file_actions_route_through_native_commands_only_in_tauri():
+ helper = NATIVE_FILES.read_text(encoding = "utf-8")
+ history = CHAT_EXPORT.read_text(encoding = "utf-8")
+ data_tab = DATA_TAB.read_text(encoding = "utf-8")
+ prompt_storage = PROMPT_STORAGE.read_text(encoding = "utf-8")
+
+ projects = (FRONTEND / "features/chat/projects-page.tsx").read_text(encoding = "utf-8")
+
+ assert 'invoke("save_native_file", bytes, {' in helper
+ assert '"x-unsloth-default-name"' in helper
+ assert "Array.from(new Uint8Array" not in helper
+ assert 'invoke("pick_native_chat_import")' in helper
+ assert "if (isTauri)" in helper
+ assert 'document.createElement("a")' in helper
+ assert "DownloadCancelledError" in helper
+ assert "throw new DownloadCancelledError()" in helper
+ assert "return savedPath !== null" not in helper
+
+ assert helper.index("if (isTauri)") < helper.index(" const blob =")
+ assert "downloadFile(" in history
+ assert "downloadFile(" in prompt_storage
+ assert "pickNativeChatImport" in data_tab
+ assert "if (!isTauri)" in data_tab
+
+ assert "pickNativeChatImport" in projects
+ assert "if (!isTauri)" in projects
+ # Browser builds retain the existing hidden-input route.
+ assert 'type="file"' in data_tab
+ assert 'accept=".jsonl,.ndjson,.csv"' in data_tab
+
+ native_dialogs = NATIVE_DIALOGS.read_text(encoding = "utf-8")
+ assert 'CHAT_IMPORT_EXTENSIONS: &[&str] = &["jsonl", "ndjson", "csv"]' in native_dialogs
+ assert "InvokeBody::Raw" in native_dialogs
+
+ assert ".tempfile_in(parent)" in native_dialogs
+ assert ".persist(&path)" in native_dialogs
+ assert "fs::write(&path, content)" not in native_dialogs
+
+
+def test_chat_exports_await_native_saves_and_markdown_uses_shared_helper():
+ app_sidebar = APP_SIDEBAR.read_text(encoding = "utf-8")
+ prompt_storage = PROMPT_STORAGE.read_text(encoding = "utf-8")
+ thread = THREAD.read_text(encoding = "utf-8")
+ thread_sidebar = THREAD_SIDEBAR.read_text(encoding = "utf-8")
+ shared_composer = SHARED_COMPOSER.read_text(encoding = "utf-8")
+
+ data_tab = DATA_TAB.read_text(encoding = "utf-8")
+ projects = (FRONTEND / "features/chat/projects-page.tsx").read_text(encoding = "utf-8")
+ assert "async function downloadBlob(" in prompt_storage
+ download_blob = prompt_storage.split("async function downloadBlob(", 1)[1].split("\n}\n", 1)[0]
+ assert "return downloadFile(" in download_blob
+ assert "catch (error)" not in download_blob
+ assert "isDownloadCancelled(error)" in prompt_storage
+
+ for source in (app_sidebar, thread, thread_sidebar, shared_composer, data_tab, projects):
+ assert "isDownloadCancelled(error)" in source
+ assert "const handleExport = useCallback(async () =>" in prompt_storage
+ assert prompt_storage.count("await export") >= 12
+ assert "await Promise.all(" not in app_sidebar
+ assert "for (const id of ids)" in app_sidebar
+ assert prompt_storage.count("await downloadBlob(") >= 5
+
+ assert "await downloadBlob(zipped," in prompt_storage
+ assert "new Blob([zipped]" not in prompt_storage
+ assert "Promise.all(ids.map((id) => fn(id)))" not in thread_sidebar
+ assert "for (const id of ids)" in thread_sidebar
+ assert "Promise.all(exportThreadIds.map((id) => fn(id)))" not in shared_composer
+ assert "for (const id of exportThreadIds)" in shared_composer
+ assert "onExport={exportMessageMarkdown}" in thread
+ assert '"text/markdown"' in thread
+ assert "downloadFile(" in thread
+
+
+def test_full_app_layout_uses_its_own_initialized_marker():
+ source = APP_PROVIDER.read_text(encoding = "utf-8")
+
+ assert 'invoke("has_initialized_app_window_layout")' in source
+ setup_layout = source.split("async function showSetupWindow", 1)[1].split(
+ "async function enforceMinimumWindowSize", 1
+ )[0]
+ reset_call = 'invoke("reset_app_window_layout_initialized")'
+ assert reset_call in setup_layout
+ assert setup_layout.index(reset_call) < setup_layout.index("win.setSize")
+ assert 'invoke("mark_app_window_layout_initialized")' in source
+ assert "hasInitializedAppLayout && hasSavedState" in source
+
+
+def test_expanded_titlebar_button_and_corner_match_sidebar_edge():
+ source = TITLEBAR.read_text(encoding = "utf-8")
+
+ assert 'pinned ? "gap-2 pl-3" : "justify-center"' in source
+ assert "gap-2 px-3" not in source
+ assert "const contentBorderLeft = pinned" in source
+ assert ': "0px";' in source
+ # The curved transition and sidebar-colored backing are expanded-only;
+ # collapsed content is square and its divider spans the sidebar too.
+ assert source.count("{showSidebarSurface && pinned && (") == 2
+ assert (
+ 'className="pointer-events-none absolute top-full size-3 -translate-x-px bg-sidebar"'
+ in source
+ )
+ assert (
+ 'className="pointer-events-none absolute top-full size-3 -translate-x-px rounded-tl-[12px] border-l border-t border-sidebar-border bg-background"'
+ in source
+ )
From f3c085ad9ec8347e9e436ee639809dd4a67eba46 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Tue, 21 Jul 2026 15:04:58 +0530
Subject: [PATCH 048/255] Fix resume training crash recovery and MLX
checkpoints (#6796)
* Fix resume training crash recovery and MLX checkpoints
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: preserve interrupted stop-and-save output_dir, verify MLX checkpoint
- finish_run: add clear_output_dir flag; preserve output_dir for stopped/error
unless cancel explicitly clears it (fixes pump finalization wiping persisted path).
- training pump: pass interrupted stop-and-save context into finalize_run_in_db.
- MLX stop-and-save: verify resumable checkpoint exists before sending complete;
return bool from _write_mlx_stop_checkpoint and add regression tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address Codex review: MLX current-step checkpoint and cancel error finalize
- Only skip MLX stop checkpoint write when checkpoint-{current_step} exists;
stale periodic checkpoints no longer mask missing stop saves.
- Pass clear_output_dir through error-event finalization so Stop-without-save
cannot leave a persisted output_dir that still offers Resume.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review
* Address more reviews
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* more reviews
* clear in-memory output_dir on interrupted cancel
* allow resuming errored runs at the final step
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Clear persisted output_dir in cancel watchdog path
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Write MLX stop checkpoint in stop path, keep output_dir on crash finalize
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): harden resumable run finalization
* fix(studio): defer safetensors checkpoint import
* fix(studio): reject stale training cancellation
* fix(studio): replay null resume targets
* fix(studio): serialize terminal cancellation
* Harden resume checkpoint validation and fix stop-save cleanup
- Reject unrecognized shard formats and keep indexed shard paths inside the checkpoint dir
- Require a non-empty tensor record when validating .pt/.bin optimizer and model state
- Always finalize TensorBoard and W&B on stop-save-failure exits
- Refuse writing an MLX stop checkpoint through a symlinked directory
- Clarify the resume rejection message to cover errored runs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten resume/checkpoint comments
* Recover resumability when a valid stop checkpoint landed
- Re-validate the current-step checkpoint in the dead-worker and error finalization paths so a stop-and-save that actually wrote a valid checkpoint is not wrongly marked error/resume_blocked
- Accept a valid tensor-free optimizer state (e.g. SGD without momentum); the model-state check still requires real tensors
- Include errored runs in the frontend resume rejection message
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lyxot
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: danielhanchen
---
studio/backend/core/training/resume.py | 173 +++++-
studio/backend/core/training/training.py | 334 +++++++++---
studio/backend/core/training/worker.py | 148 +++++-
studio/backend/models/training.py | 7 +
studio/backend/routes/training.py | 19 +-
studio/backend/storage/studio_db.py | 114 +++-
.../backend/tests/test_mlx_stop_checkpoint.py | 137 +++++
.../tests/test_training_pump_resilience.py | 74 +++
studio/backend/tests/test_training_resume.py | 498 ++++++++++++++++++
.../tests/test_training_stop_watchdog.py | 91 +++-
.../studio/historical-training-view.tsx | 50 +-
.../src/features/studio/history-card-grid.tsx | 8 +-
.../src/features/studio/studio-page.tsx | 8 +-
.../training/hooks/use-training-actions.ts | 2 +-
.../training/stores/training-runtime-store.ts | 5 +-
.../src/features/training/types/runtime.ts | 3 +-
studio/frontend/src/i18n/locales/en.ts | 3 +-
studio/frontend/src/i18n/locales/zh-CN.ts | 3 +-
18 files changed, 1532 insertions(+), 145 deletions(-)
create mode 100644 studio/backend/tests/test_mlx_stop_checkpoint.py
diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py
index bbd9a895ab..17183484c5 100644
--- a/studio/backend/core/training/resume.py
+++ b/studio/backend/core/training/resume.py
@@ -4,6 +4,8 @@
"""Helpers for validating resumable training outputs."""
import json
+import pickletools
+import zipfile
from pathlib import Path
from typing import Optional
@@ -33,21 +35,158 @@ def _checkpoint_step(path: Path) -> int:
return -1
-def get_resume_checkpoint_path(path_value: str) -> Optional[str]:
+_MODEL_FILES = (
+ "adapter_model.safetensors",
+ "adapter_model.bin",
+ "model.safetensors",
+ "pytorch_model.bin",
+)
+_MODEL_INDEXES = ("model.safetensors.index.json", "pytorch_model.bin.index.json")
+
+
+def _valid_state_file(path: Path, require_tensor: bool = True) -> bool:
+ try:
+ if not path.is_file() or path.stat().st_size == 0:
+ return False
+ if path.suffix == ".safetensors":
+ try:
+ from safetensors import SafetensorError, safe_open
+ except ImportError:
+ return False
+ try:
+ with safe_open(str(path), framework = "np") as state:
+ return bool(state.keys())
+ except SafetensorError:
+ return False
+ if path.suffix in {".bin", ".pt"}:
+ with zipfile.ZipFile(path) as state:
+ infos = state.infolist()
+ names = [info.filename for info in infos]
+ data_name = next(
+ (name for name in names if name == "data.pkl" or name.endswith("/data.pkl")),
+ None,
+ )
+ if data_name is None:
+ return False
+ data_prefix = data_name.removesuffix("data.pkl") + "data/"
+ operations = list(pickletools.genops(state.read(data_name)))
+ if not operations or operations[-1][0].name != "STOP":
+ return False
+ if not require_tensor:
+ return True
+ # Require a non-empty tensor record; a zero-byte one fails torch.load.
+ return any(
+ info.filename.startswith(data_prefix)
+ and not info.is_dir()
+ and info.file_size > 0
+ for info in infos
+ )
+ # Unrecognized state-file formats are not usable resume state.
+ return False
+ except (OSError, ValueError, zipfile.BadZipFile):
+ return False
+
+
+def _checkpoint_state(path: Path) -> Optional[int]:
+ try:
+ state = json.loads((path / "trainer_state.json").read_text(encoding = "utf-8"))
+ step = state.get("global_step") if isinstance(state, dict) else None
+ except (OSError, UnicodeDecodeError, json.JSONDecodeError):
+ return None
+ if isinstance(step, bool) or not isinstance(step, int) or step < 0:
+ return None
+ directory_step = _checkpoint_step(path)
+ return step if directory_step < 0 or step == directory_step else None
+
+
+_INDEX_SHARD_SUFFIX = {
+ "model.safetensors.index.json": ".safetensors",
+ "pytorch_model.bin.index.json": ".bin",
+}
+
+
+def _valid_indexed_shard(checkpoint: Path, shard: object, expected_suffix: str) -> bool:
+ # Shard must be a relative, in-format path contained in the checkpoint dir.
+ if not isinstance(shard, str) or not shard:
+ return False
+ if Path(shard).is_absolute() or Path(shard).suffix != expected_suffix:
+ return False
+ try:
+ root = checkpoint.resolve(strict = True)
+ candidate = (checkpoint / shard).resolve(strict = True)
+ candidate.relative_to(root)
+ except (OSError, ValueError):
+ return False
+ return _valid_state_file(candidate)
+
+
+def _has_model_state(path: Path) -> bool:
+ if any(_valid_state_file(path / name) for name in _MODEL_FILES):
+ return True
+ for name in _MODEL_INDEXES:
+ try:
+ index = json.loads((path / name).read_text(encoding = "utf-8"))
+ shards = set(index["weight_map"].values())
+ except (
+ AttributeError,
+ OSError,
+ KeyError,
+ TypeError,
+ UnicodeDecodeError,
+ json.JSONDecodeError,
+ ):
+ continue
+ expected_suffix = _INDEX_SHARD_SUFFIX[name]
+ if shards and all(_valid_indexed_shard(path, shard, expected_suffix) for shard in shards):
+ return True
+ return False
+
+
+def is_resume_checkpoint_valid(
+ path: Path,
+ expected_step: Optional[int] = None,
+ backend: Optional[str] = None,
+) -> bool:
+ step = _checkpoint_state(path) if path.is_dir() else None
+ step_valid = step is not None and (expected_step is None or step == expected_step)
+ if backend == "mlx":
+ valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file(
+ path / "optimizer_state.safetensors"
+ )
+ else:
+ valid_bundle = (
+ _has_model_state(path)
+ # optimizer/scheduler state can be validly tensor-free (e.g. SGD without
+ # momentum); _has_model_state still requires real model tensors.
+ and _valid_state_file(path / "optimizer.pt", require_tensor = False)
+ and _valid_state_file(path / "scheduler.pt", require_tensor = False)
+ )
+ if backend is None and not valid_bundle:
+ valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file(
+ path / "optimizer_state.safetensors"
+ )
+ return step_valid and valid_bundle
+
+
+def get_resume_checkpoint_path(
+ path_value: str, expected_step: Optional[int] = None
+) -> Optional[str]:
path = resolve_output_dir(path_value)
if not _is_under_outputs(path) or not path.is_dir():
return None
- if (path / "trainer_state.json").is_file():
+ if is_resume_checkpoint_valid(path, expected_step):
return str(path)
- checkpoints = [
- child
- for child in path.glob("checkpoint-*")
- if child.is_dir() and (child / "trainer_state.json").is_file()
- ]
- if not checkpoints:
- return None
- return str(max(checkpoints, key = _checkpoint_step))
+ checkpoints = sorted(path.glob("checkpoint-*"), key = _checkpoint_step, reverse = True)
+ return next(
+ (
+ str(checkpoint)
+ for checkpoint in checkpoints
+ if _checkpoint_step(checkpoint) >= 0
+ and is_resume_checkpoint_valid(checkpoint, expected_step)
+ ),
+ None,
+ )
def normalize_resume_output_dir(path_value: str) -> str:
@@ -78,9 +217,17 @@ def _uses_s3_dataset(run: dict) -> bool:
def can_resume_run(run: dict) -> bool:
if run.get("resumed_later"):
return False
+ # Set when a stop-and-save failed to write a current-step checkpoint.
+ if run.get("resume_blocked"):
+ return False
if _uses_s3_dataset(run):
return False
+ status = run.get("status")
+ if status == "error":
+ # A save-time crash can report final_step == total_steps with no artifacts; checkpoint state alone decides resumability.
+ return has_resume_state(run.get("output_dir"))
+
final_step = run.get("final_step")
total_steps = run.get("total_steps")
has_remaining_steps = (
@@ -89,8 +236,4 @@ def can_resume_run(run: dict) -> bool:
or total_steps <= 0
or final_step < total_steps
)
- return (
- run.get("status") == "stopped"
- and has_remaining_steps
- and has_resume_state(run.get("output_dir"))
- )
+ return status == "stopped" and has_remaining_steps and has_resume_state(run.get("output_dir"))
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index b407ba39a5..26d8c23b66 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -761,6 +761,7 @@ class TrainingBackend:
# Left True after an abnormal death so _ensure_pump_alive spots a crash.
self._pump_running: bool = False
self._lock = threading.Lock()
+ self._run_intent_lock = threading.RLock()
# Stop watchdog: after a stop is requested, escalates to force_terminate()
# if the worker does not exit on its own within a bounded time. The watched
@@ -773,6 +774,7 @@ class TrainingBackend:
self._progress = TrainingProgress()
self._should_stop = False
self._cancel_requested = False # True only for stop(save=False)
+ self._cancel_cleanup_output_dir: Optional[str] = None
# Throttled training-status logging to the server log (not one line/step).
self._last_progress_log_ts: float = 0.0
@@ -792,6 +794,8 @@ class TrainingBackend:
# Job metadata
self.current_job_id: Optional[str] = None
self._output_dir: Optional[str] = None
+ self._resume_source_run_id: Optional[str] = None
+ self._terminal_finalize_payload: Optional[dict] = None
# DB persistence
self._metric_buffer: list[dict] = []
@@ -819,6 +823,7 @@ class TrainingBackend:
job_id: str,
*,
before_spawn = None,
+ resume_source_run_id: Optional[str] = None,
**kwargs,
) -> bool:
"""Spawn a subprocess to run the full training pipeline.
@@ -956,6 +961,7 @@ class TrainingBackend:
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
+ self._cancel_cleanup_output_dir = None
self._complete_seen.clear()
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
@@ -972,7 +978,10 @@ class TrainingBackend:
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
- self._output_dir = None
+ self._output_dir = config.get("output_dir") if resume_source_run_id else None
+ self._progress.output_dir = self._output_dir
+ self._resume_source_run_id = resume_source_run_id
+ self._terminal_finalize_payload = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
@@ -990,6 +999,17 @@ class TrainingBackend:
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
+ if resume_source_run_id and not self._db_run_created:
+ if proc.is_alive():
+ proc.terminate()
+ proc.join(timeout = 5.0)
+ if proc.is_alive():
+ proc.kill()
+ proc.join(timeout = 2.0)
+ self._progress.is_training = False
+ self._progress.error = "Resume checkpoint is no longer available."
+ self._spawn_in_progress = False
+ return False
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
@@ -1011,28 +1031,75 @@ class TrainingBackend:
def stop_training(self, save: bool = True) -> bool:
"""Send stop signal to the training subprocess."""
- self._should_stop = True
- if not save:
- self._cancel_requested = True
- with self._lock:
- if self._stop_queue is not None:
- try:
- self._stop_queue.put({"type": "stop", "save": save})
- except (OSError, ValueError):
- pass
- # Update progress immediately for responsive UI.
- self._progress.status_message = (
- "Stopping training and saving checkpoint..." if save else "Cancelling training..."
- )
- # Guarantee the run finalizes even if the worker wedges after saving.
- self._start_stop_watchdog(cancel = not save)
+ with self._run_intent_lock:
+ with self._lock:
+ run_id = self.current_job_id
+ if not save and run_id:
+ persist_error: Optional[Exception] = None
+ for attempt in range(_DB_FINALIZE_RETRIES):
+ try:
+ from storage.studio_db import mark_run_cancel_requested
+
+ self._ensure_db_run_created()
+ with self._lock:
+ terminal_payload = self._terminal_finalize_payload
+ if (
+ terminal_payload
+ and terminal_payload.get("expected_job_id") == run_id
+ ):
+ return False
+ if not mark_run_cancel_requested(run_id):
+ if self._db_run_created:
+ return False
+ raise RuntimeError(
+ "Training run disappeared before cancellation persisted"
+ )
+ if self.current_job_id != run_id:
+ return False
+ self._should_stop = self._cancel_requested = True
+ self._cancel_cleanup_output_dir = self._output_dir
+ self._output_dir = self._progress.output_dir = None
+ persist_error = None
+ break
+ except Exception as exc:
+ persist_error = exc
+ if attempt + 1 < _DB_FINALIZE_RETRIES:
+ time.sleep(_DB_FINALIZE_RETRY_S)
+ if persist_error is not None:
+ raise RuntimeError("Failed to persist Stop-without-Save") from persist_error
+ with self._lock:
+ if self.current_job_id != run_id:
+ return False
+ if save or not run_id:
+ self._should_stop = True
+ if not save and not run_id:
+ self._cancel_requested = True
+ self._cancel_cleanup_output_dir = self._output_dir
+ self._output_dir = self._progress.output_dir = None
+ if self._stop_queue is not None:
+ try:
+ self._stop_queue.put({"type": "stop", "save": save})
+ except (OSError, ValueError):
+ pass
+ self._progress.status_message = (
+ "Stopping training and saving checkpoint..."
+ if save
+ else "Cancelling training..."
+ )
+ self._start_stop_watchdog(cancel = not save, expected_job_id = run_id)
return True
- def _start_stop_watchdog(self, cancel: bool) -> None:
+ def _start_stop_watchdog(
+ self,
+ cancel: bool,
+ expected_job_id: Optional[str] = None,
+ ) -> None:
"""Start a daemon that force-terminates the worker if a requested stop does not
exit on its own. No-op if no worker is alive or a live watchdog already watches
this proc (a stale watchdog on an old proc never blocks a new run's watcher)."""
with self._lock:
+ if expected_job_id is not None and self.current_job_id != expected_job_id:
+ return
proc = self._proc
if proc is None or not proc.is_alive():
return
@@ -1113,8 +1180,9 @@ class TrainingBackend:
watched_job_id: Optional[str] = None,
) -> None:
"""Finalize parent state after a force-terminate so the UI leaves "Stopping..."
- even if the worker is wedged in driver teardown; preserves output_dir so a saved
- checkpoint is kept. No-ops if a new run already replaced the watched worker, so a
+ even if the worker is wedged in driver teardown; preserves output_dir on a save so
+ the checkpoint is kept, and clears it on a cancel (Stop without saving must not
+ offer resume/export). No-ops if a new run already replaced the watched worker, so a
stale watchdog never marks a fresh run stopped or drops its handle.
Supersession is checked on both the watched proc and job id: start_training sets
@@ -1134,7 +1202,18 @@ class TrainingBackend:
return # a new run is already starting up; leave its state alone
run_id = self.current_job_id # == watched_job_id
self._progress.is_training = False
- self._progress.status_message = "Training stopped."
+ terminal_payload = self._terminal_finalize_kwargs()
+ status = terminal_payload["status"]
+ error_message = terminal_payload.get("error_message")
+ output_dir = terminal_payload["output_dir"]
+ clear_output_dir = terminal_payload["clear_output_dir"]
+ resume_blocked = bool(terminal_payload.get("resume_blocked"))
+ with self._lock:
+ if self.current_job_id != run_id:
+ return
+ self._progress.status_message = error_message or "Training stopped."
+ if error_message:
+ self._progress.error = error_message
# Create the row if a start-time create failed (no-op otherwise; skips when the pump
# is mid-create, in which case its create-then-finalize records the run instead).
self._ensure_db_run_created()
@@ -1148,7 +1227,8 @@ class TrainingBackend:
batch: list = []
final_step = final_loss = duration = None
loss_history: list = []
- output_dir = self._output_dir
+ if clear_output_dir:
+ self._output_dir = self._progress.output_dir = None
if claim:
self._run_finalized = True # claim this run's finalize
batch = list(self._metric_buffer)
@@ -1161,7 +1241,17 @@ class TrainingBackend:
loss_history = list(self.loss_history)
if claim:
self._finish_stopped_run(
- run_id, output_dir, batch, final_step, final_loss, duration, loss_history
+ run_id,
+ output_dir,
+ batch,
+ final_step,
+ final_loss,
+ duration,
+ loss_history,
+ status = status,
+ error_message = error_message,
+ clear_output_dir = clear_output_dir,
+ resume_blocked = resume_blocked,
)
with self._lock:
if target_proc is None or self._proc is target_proc:
@@ -1176,6 +1266,10 @@ class TrainingBackend:
final_loss: Optional[float],
duration: Optional[float],
loss_history: list,
+ status: str = "stopped",
+ error_message: Optional[str] = None,
+ clear_output_dir: bool = False,
+ resume_blocked: bool = False,
) -> None:
"""Record a force-stopped run finished by its captured id, from state snapshotted
under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE,
@@ -1194,14 +1288,16 @@ class TrainingBackend:
sparkline = downsample(loss_history, 50)
finish_run(
id = run_id,
- status = "stopped",
+ status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
- error_message = None,
+ error_message = error_message,
+ clear_output_dir = clear_output_dir,
+ resume_blocked = resume_blocked,
)
return
except Exception:
@@ -1231,7 +1327,7 @@ class TrainingBackend:
logger.info("Force-terminating training subprocess (pid=%s)", proc.pid)
proc.terminate()
cancelled = self._cancel_requested
- output_dir = self._output_dir
+ output_dir = self._cancel_cleanup_output_dir or self._output_dir
if proc is not None:
proc.join(timeout = 5.0)
@@ -1595,17 +1691,60 @@ class TrainingBackend:
)
self._ensure_db_run_created()
- self._finalize_run_in_db(
- status = "stopped" if self._should_stop else "error",
- error_message = None
- if self._should_stop
- else "Training process terminated unexpectedly",
- )
+ terminal_payload = self._terminal_finalize_kwargs()
+ with self._lock:
+ if terminal_payload["clear_output_dir"]:
+ self._output_dir = self._progress.output_dir = None
+ if terminal_payload.get("error_message"):
+ self._progress.error = terminal_payload["error_message"]
+ self._progress.status_message = terminal_payload["error_message"]
+ self._finalize_run_in_db(**terminal_payload)
except Exception:
logger.exception("Training event pump: finalization after worker exit failed")
self._pump_running = False
return
+ def _has_current_resume_checkpoint(self, output_dir, step) -> bool:
+ # A valid checkpoint at the current step means the stop-and-save landed on
+ # disk even if the worker died before confirming it.
+ if not output_dir or not isinstance(step, int) or step <= 0:
+ return False
+ from core.training.resume import get_resume_checkpoint_path
+ return get_resume_checkpoint_path(output_dir, expected_step = step) is not None
+
+ def _terminal_finalize_kwargs(self) -> dict:
+ with self._lock:
+ job_id = self.current_job_id
+ payload = self._terminal_finalize_payload
+ if payload and payload.get("expected_job_id") == job_id:
+ return dict(payload)
+ cancel, stopped = self._cancel_requested, self._should_stop
+ output_dir = None if cancel else self._output_dir
+ step = self._progress.step
+ existing_error = self._progress.error
+ status, error, blocked = (
+ ("stopped", None, cancel)
+ if stopped
+ else (
+ "error",
+ existing_error or "Training process terminated unexpectedly",
+ False,
+ )
+ )
+ # Block only when no valid current-step checkpoint actually landed.
+ if stopped and not cancel and not self._has_current_resume_checkpoint(output_dir, step):
+ status = "error"
+ error = "Stop and Save ended before a valid current-step checkpoint was written."
+ blocked = True
+ return {
+ "status": status,
+ "error_message": error,
+ "output_dir": output_dir,
+ "clear_output_dir": cancel,
+ "resume_blocked": blocked,
+ "expected_job_id": job_id,
+ }
+
def _handle_event(self, event: dict) -> None:
"""Apply a subprocess event to local state.
@@ -1764,6 +1903,15 @@ class TrainingBackend:
elif etype == "eval_configured":
self.eval_enabled = True
+ elif etype == "output_dir":
+ event_output_dir = event.get("output_dir")
+ if self._cancel_requested:
+ self._cancel_cleanup_output_dir = event_output_dir
+ self._output_dir = self._progress.output_dir = None
+ else:
+ self._output_dir = event_output_dir
+ db_action = "persist_output_dir"
+
elif etype == "status":
self._progress.status_message = event.get("message", "")
self._progress.is_training = True
@@ -1778,7 +1926,12 @@ class TrainingBackend:
self._complete_seen.set()
self._progress.is_training = False
self._progress.is_completed = not stopped
- self._output_dir = event.get("output_dir")
+ event_output_dir = event.get("output_dir")
+ if self._cancel_requested:
+ self._cancel_cleanup_output_dir = event_output_dir
+ self._output_dir = None
+ else:
+ self._output_dir = event_output_dir
self._progress.output_dir = self._output_dir
self._progress.status_message = msg
if not self._db_run_created and self.current_job_id and self._db_config:
@@ -1788,11 +1941,16 @@ class TrainingBackend:
db_action_kwargs = {
"status": "stopped" if stopped else "completed",
"output_dir": self._output_dir,
+ "clear_output_dir": self._cancel_requested,
+ "expected_job_id": self.current_job_id,
}
+ self._terminal_finalize_payload = dict(db_action_kwargs)
elif etype == "error":
self._progress.is_training = False
self._progress.error = event.get("error", "Unknown error")
+ if self._cancel_requested:
+ self._output_dir = self._progress.output_dir = None
logger.error("Training error: %s", event.get("error"))
stack = event.get("stack", "")
if stack:
@@ -1801,29 +1959,36 @@ class TrainingBackend:
db_action = "create_and_finalize"
else:
db_action = "finalize"
+ stop_save_failed = (
+ self._should_stop
+ and not self._cancel_requested
+ and not self._has_current_resume_checkpoint(
+ self._output_dir, self._progress.step
+ )
+ )
db_action_kwargs = {
- "status": "stopped" if self._should_stop else "error",
+ "status": "stopped"
+ if self._should_stop
+ and not stop_save_failed
+ and not event.get("keep_error_status")
+ else "error",
"error_message": event.get("error", "Unknown error"),
+ "output_dir": self._output_dir,
+ "clear_output_dir": self._cancel_requested,
+ "resume_blocked": stop_save_failed or bool(event.get("resume_blocked")),
+ "expected_job_id": self.current_job_id,
}
+ self._terminal_finalize_payload = dict(db_action_kwargs)
# --- DB I/O outside the lock ---
if db_action == "create_run":
- try:
- from storage.studio_db import create_run
-
- create_run(
- id = db_action_kwargs["job_id"],
- model_name = db_action_kwargs["model_name"],
- dataset_name = db_action_kwargs["dataset_name"],
- config_json = db_action_kwargs["config_json"],
- started_at = db_action_kwargs["started_at"],
- total_steps = db_action_kwargs["total_steps"],
- )
- self._db_run_created = True
+ self._ensure_db_run_created()
+ if self._db_run_created:
if db_action_kwargs["total_steps"]:
self._db_total_steps_set = True
- except Exception:
- logger.warning("Failed to create DB run record", exc_info = True)
+ self._persist_output_dir()
+ elif db_action == "persist_output_dir":
+ self._persist_output_dir()
elif db_action == "create_and_finalize":
self._ensure_db_run_created()
self._finalize_run_in_db(**db_action_kwargs)
@@ -1842,6 +2007,22 @@ class TrainingBackend:
if etype == "progress":
self._log_training_progress()
+ def _persist_output_dir(self) -> None:
+ with self._lock:
+ if (
+ not self._output_dir
+ or not self.current_job_id
+ or not self._db_run_created
+ or self._cancel_requested
+ ):
+ return
+ run_id, output_dir = self.current_job_id, self._output_dir
+ try:
+ from storage.studio_db import update_run_output_dir
+ update_run_output_dir(run_id, output_dir)
+ except Exception:
+ logger.warning("Failed to persist output_dir", exc_info = True)
+
def _log_training_progress(self) -> None:
"""One throttled training-status line to the server log (the per-step stream
still goes to the UI via SSE): first step, then at most every 30s, plus the
@@ -1875,6 +2056,7 @@ class TrainingBackend:
caller create at a time, and ``_db_run_created`` is published only after
``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a
not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running)."""
+ self._run_intent_lock.acquire()
with self._lock:
if (
self._db_run_created
@@ -1882,6 +2064,7 @@ class TrainingBackend:
or not self.current_job_id
or not self._db_config
):
+ self._run_intent_lock.release()
return
self._db_create_in_progress = True # only one caller creates
job_id = self.current_job_id
@@ -1898,6 +2081,12 @@ class TrainingBackend:
or _s3_dataset_name(db_config.get("s3_dataset"))
or "unknown"
)
+ with self._lock:
+ if self.current_job_id != job_id:
+ return
+ output_dir = self._output_dir
+ cancel_requested = self._cancel_requested
+ resumed_from_run_id = self._resume_source_run_id
create_run(
id = job_id,
model_name = db_config["model_name"],
@@ -1905,6 +2094,9 @@ class TrainingBackend:
config_json = _json.dumps(db_config),
started_at = started_at,
total_steps = total_steps,
+ output_dir = output_dir,
+ cancel_requested = cancel_requested,
+ resumed_from_run_id = resumed_from_run_id,
)
created = True
except Exception:
@@ -1919,12 +2111,15 @@ class TrainingBackend:
if created:
self._db_run_created = True # publish only after the insert commits
self._db_create_in_progress = False
+ self._run_intent_lock.release()
def _finalize_run_in_db(
self,
status: str,
error_message: Optional[str] = None,
output_dir: Optional[str] = None,
+ clear_output_dir: bool = False,
+ resume_blocked: bool = False,
expected_job_id: Optional[str] = None,
) -> None:
"""Flush remaining metrics and mark a run finished in the DB. Claims the finalize
@@ -1947,26 +2142,33 @@ class TrainingBackend:
duration = self._progress.elapsed_seconds
loss_history = list(self.loss_history)
self._flush_metrics_to_db(run_id = run_id)
- try:
- from storage.studio_db import finish_run
- from utils.downsample import downsample
+ for attempt in range(_DB_FINALIZE_RETRIES):
+ try:
+ from storage.studio_db import finish_run
+ from utils.downsample import downsample
- sparkline = downsample(loss_history, 50)
- finish_run(
- id = run_id,
- status = status,
- ended_at = datetime.now(timezone.utc).isoformat(),
- final_step = final_step,
- final_loss = final_loss,
- duration_seconds = duration,
- loss_sparkline = _json.dumps(sparkline),
- output_dir = output_dir,
- error_message = error_message,
- )
- except Exception:
- with self._lock:
- self._run_finalized = False # unclaim so a later flush can retry
- logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
+ finish_run(
+ id = run_id,
+ status = status,
+ ended_at = datetime.now(timezone.utc).isoformat(),
+ final_step = final_step,
+ final_loss = final_loss,
+ duration_seconds = duration,
+ loss_sparkline = _json.dumps(downsample(loss_history, 50)),
+ output_dir = output_dir,
+ error_message = error_message,
+ clear_output_dir = clear_output_dir,
+ resume_blocked = resume_blocked,
+ )
+ return
+ except Exception:
+ if attempt + 1 < _DB_FINALIZE_RETRIES:
+ time.sleep(_DB_FINALIZE_RETRY_S)
+ continue
+ with self._lock:
+ if self.current_job_id == run_id:
+ self._run_finalized = False
+ logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None:
"""Flush buffered metrics to the DB and update live progress. The target run id,
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 111f4fdd0f..2df4fa58c6 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -1840,8 +1840,15 @@ def _run_mlx_training(event_queue, stop_queue, config):
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import ensure_dir
- output_dir = _resolve_mlx_output_dir(config, model_name)
+ # Resume must land in the original run dir even when config lacks output_dir.
+ resume_dir = config.get("output_dir", "") or _output_dir_from_resume_checkpoint(
+ resume_from_checkpoint
+ )
+ output_dir = _resolve_mlx_output_dir(
+ {**config, "output_dir": resume_dir} if resume_dir else config, model_name
+ )
ensure_dir(Path(output_dir))
+ _emit_output_dir(event_queue, output_dir)
# ── 6. Create trainer ──
eval_steps_val = config.get("eval_steps", 0) or 0
@@ -2067,6 +2074,17 @@ def _run_mlx_training(event_queue, stop_queue, config):
trainer.add_eval_callback(_on_eval)
+ _opt_ref = [None]
+ _orig_build_optimizer = getattr(trainer, "_build_optimizer", None)
+
+ if callable(_orig_build_optimizer):
+
+ def _capture_optimizer(total_steps):
+ _opt_ref[0] = _orig_build_optimizer(total_steps)
+ return _opt_ref[0]
+
+ trainer._build_optimizer = _capture_optimizer
+
# ── 11. Run training ──
gc.collect()
mx.synchronize()
@@ -2082,31 +2100,58 @@ def _run_mlx_training(event_queue, stop_queue, config):
trainer.save_model = _save_model
# ── 12. Save and finalize ──
- if trainer.stop_requested:
- if not _stop_save[0]:
- # Cancel (save=False): skip saving.
- _send("complete", output_dir = None, status_message = "Training cancelled")
+ def _finish_tracking() -> None:
+ # Runs on every save/finalize exit so TB/W&B never leak on early return.
+ if tb_writer is not None:
+ try:
+ tb_writer.close()
+ except Exception:
+ pass
+ if wandb_run is not None:
+ try:
+ wandb_run.finish()
+ except Exception:
+ pass
+
+ def _stop_checkpoint_ok() -> bool:
+ if _write_mlx_stop_checkpoint(trainer, _opt_ref[0], output_dir):
+ return True
+ _send(
+ "error",
+ error = (
+ "Failed to save a resumable checkpoint after stop. "
+ "Model files were saved, but this run cannot be resumed."
+ ),
+ # A user stop finalizes as 'stopped'; keep this failure's error status so history explains it.
+ keep_error_status = True,
+ # Older checkpoints are stale; resuming would roll back past this stop.
+ resume_blocked = True,
+ )
+ return False
+
+ try:
+ if trainer.stop_requested:
+ if not _stop_save[0]:
+ # Cancel (save=False): skip saving.
+ _send("complete", output_dir = None, status_message = "Training cancelled")
+ else:
+ _send("status", status_message = "Saving stopped model...")
+ mx.synchronize()
+ trainer.save_model(output_dir)
+ # Stop-and-save promises a resumable checkpoint, not just model files.
+ if not _stop_checkpoint_ok():
+ return
+ _send("complete", output_dir = output_dir, status_message = "Training stopped")
else:
- _send("status", status_message = "Saving stopped model...")
+ _send("status", status_message = "Saving model...")
mx.synchronize()
trainer.save_model(output_dir)
- _send("complete", output_dir = output_dir, status_message = "Training stopped")
- else:
- _send("status", status_message = "Saving model...")
- mx.synchronize()
- trainer.save_model(output_dir)
- _send("complete", output_dir = output_dir, status_message = "Training completed")
-
- if tb_writer is not None:
- try:
- tb_writer.close()
- except Exception:
- pass
- if wandb_run is not None:
- try:
- wandb_run.finish()
- except Exception:
- pass
+ # A save-stop can race the natural final save; it made the same promise.
+ if trainer.stop_requested and _stop_save[0] and not _stop_checkpoint_ok():
+ return
+ _send("complete", output_dir = output_dir, status_message = "Training completed")
+ finally:
+ _finish_tracking()
def _is_current_process_apple_silicon() -> bool:
@@ -3177,6 +3222,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
+ _emit_output_dir(event_queue, output_dir)
tensorboard_dir = config.get("tensorboard_dir")
if config.get("enable_tensorboard", False):
@@ -3296,6 +3342,61 @@ def _send_status(event_queue: Any, message: str) -> None:
)
+def _emit_output_dir(event_queue: Any, output_dir: str) -> None:
+ try:
+ event_queue.put({"type": "output_dir", "output_dir": output_dir, "ts": time.time()})
+ except Exception:
+ pass
+
+
+def _mlx_has_checkpoint_at_step(output_dir, step: int) -> bool:
+ if step <= 0:
+ return False
+ from core.training.resume import is_resume_checkpoint_valid
+ return is_resume_checkpoint_valid(
+ Path(output_dir) / f"checkpoint-{step}", expected_step = step, backend = "mlx"
+ )
+
+
+def _write_mlx_stop_checkpoint(trainer, optimizer, output_dir) -> bool:
+ """Write a full resume checkpoint for a stopped MLX run.
+
+ Returns True when a checkpoint for the current training step exists.
+ """
+ step = int(getattr(trainer, "_global_step", 0) or 0)
+ # A periodic save or a resumed run may already cover the current step.
+ if _mlx_has_checkpoint_at_step(output_dir, step):
+ return True
+ if step <= 0 or optimizer is None:
+ return False
+ ckpt_dir = Path(output_dir) / f"checkpoint-{step}"
+ if ckpt_dir.is_symlink():
+ # Refuse a symlinked dir: it could redirect writes outside output_dir.
+ logger.error("Refusing to write MLX stop checkpoint through symlink: %s", ckpt_dir)
+ return False
+ try:
+ ckpt_dir.mkdir(parents = True, exist_ok = True)
+ from unsloth_zoo.mlx.utils import (
+ save_optimizer_state,
+ save_trainable_adapters,
+ save_trainer_state,
+ )
+
+ save_trainable_adapters(trainer.model, str(ckpt_dir))
+ save_optimizer_state(optimizer, str(ckpt_dir))
+ save_trainer_state(
+ {
+ "global_step": step,
+ "train_loss_history": list(getattr(trainer, "_train_loss_history", [])),
+ },
+ str(ckpt_dir),
+ )
+ logger.info("Saved stop checkpoint to %s", ckpt_dir)
+ except Exception:
+ logger.exception("Failed to write stop checkpoint under %s", output_dir)
+ return _mlx_has_checkpoint_at_step(output_dir, step)
+
+
def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None:
"""Self-contained embedding model training pipeline.
@@ -3660,6 +3761,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
+ _emit_output_dir(event_queue, output_dir)
num_epochs = config.get("num_epochs", 2)
batch_size = config.get("batch_size", 256)
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 0b50f63b95..0f88b78f9f 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -505,6 +505,13 @@ class TrainingStartRequest(BaseModel):
description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.",
)
+ @field_validator("target_modules", mode = "before")
+ @classmethod
+ def _normalize_target_modules(cls, value: Any) -> Any:
+ # Sanitized non-LoRA history stores the unused value as null; treat it as a
+ # fresh request's omitted/default empty list on resume.
+ return [] if value is None else value
+
@model_validator(mode = "after")
def _validate_streaming_splits(self) -> "TrainingStartRequest":
# Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]"
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 53b1c4d991..a8a9874b1b 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -196,6 +196,7 @@ async def start_training(
request.local_eval_datasets, "Local eval dataset"
)
resume_output_dir: Optional[str] = None
+ resume_run: Optional[dict] = None
if request.resume_from_checkpoint:
try:
resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint)
@@ -208,7 +209,7 @@ async def start_training(
if not resume_run or not can_resume_run(resume_run):
raise HTTPException(
status_code = 400,
- detail = "Resume checkpoint must belong to a stopped run with saved trainer state.",
+ detail = "Resume checkpoint must belong to a stopped or errored run with complete saved trainer state.",
)
resume_checkpoint = get_resume_checkpoint_path(resume_output_dir)
if not resume_checkpoint:
@@ -458,7 +459,10 @@ async def start_training(
try:
success = backend.start_training(
- job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
+ job_id = job_id,
+ before_spawn = _free_vram_for_training,
+ resume_source_run_id = resume_run["id"] if resume_run else None,
+ **training_kwargs,
)
except SidecarSwapInProgress as exc:
# Expected loss of the race against a sidecar install: a retryable
@@ -521,7 +525,10 @@ async def stop_training(
status = "idle", message = "No training job is currently running"
)
- backend.stop_training(save = body.save)
+ if not backend.stop_training(save = body.save):
+ return TrainingStopResponse(
+ status = "idle", message = "No training job is currently running"
+ )
return TrainingStopResponse(
status = "stopped",
@@ -637,9 +644,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject
"loss": getattr(progress, "loss", None),
"learning_rate": getattr(progress, "learning_rate", None),
}
- output_dir = getattr(backend, "_output_dir", None)
- if output_dir:
- details["output_dir"] = output_dir
+ # Always present: an explicit null tells the client to drop a cached
+ # path (stop without save clears the run's output_dir).
+ details["output_dir"] = getattr(backend, "_output_dir", None) or None
# Metric history for chart recovery after SSE reconnection.
metric_history = None
diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py
index d889894d04..6972e7b7ff 100644
--- a/studio/backend/storage/studio_db.py
+++ b/studio/backend/storage/studio_db.py
@@ -192,13 +192,18 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
error_message TEXT,
duration_seconds REAL,
loss_sparkline TEXT,
- display_name TEXT
+ display_name TEXT,
+ resume_blocked INTEGER NOT NULL DEFAULT 0
)
"""
)
existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()}
if "display_name" not in existing_cols:
conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT")
+ if "resume_blocked" not in existing_cols:
+ conn.execute(
+ "ALTER TABLE training_runs ADD COLUMN resume_blocked INTEGER NOT NULL DEFAULT 0"
+ )
conn.execute(
"""
CREATE TABLE IF NOT EXISTS training_metrics (
@@ -734,16 +739,43 @@ def create_run(
config_json: str,
started_at: str,
total_steps: Optional[int],
+ *,
+ output_dir: Optional[str] = None,
+ cancel_requested: bool = False,
+ resumed_from_run_id: Optional[str] = None,
) -> None:
conn = get_connection()
try:
conn.execute(
"""
- INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps)
- VALUES (?, ?, ?, ?, ?, ?)
+ INSERT INTO training_runs (
+ id, model_name, dataset_name, config_json, started_at, total_steps,
+ output_dir, resume_blocked
+ )
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
- (id, model_name, dataset_name, config_json, started_at, total_steps),
+ (
+ id,
+ model_name,
+ dataset_name,
+ config_json,
+ started_at,
+ total_steps,
+ None if cancel_requested else output_dir,
+ int(cancel_requested),
+ ),
)
+ if resumed_from_run_id:
+ claimed = conn.execute(
+ """
+ UPDATE training_runs SET resume_blocked = 1
+ WHERE id = ? AND status IN ('stopped', 'error')
+ AND output_dir = ? AND resume_blocked = 0
+ """,
+ (resumed_from_run_id, output_dir),
+ )
+ if claimed.rowcount != 1:
+ raise RuntimeError("Resume source is no longer available")
conn.commit()
finally:
conn.close()
@@ -786,6 +818,8 @@ def finish_run(
loss_sparkline: Optional[str] = None,
output_dir: Optional[str] = None,
error_message: Optional[str] = None,
+ clear_output_dir: bool = False,
+ resume_blocked: bool = False,
) -> None:
conn = get_connection()
try:
@@ -793,9 +827,16 @@ def finish_run(
"""
UPDATE training_runs
SET status = ?, ended_at = ?, final_step = ?, final_loss = ?,
- duration_seconds = ?, loss_sparkline = ?, output_dir = ?,
- error_message = ?
- WHERE id = ?
+ duration_seconds = ?, loss_sparkline = ?,
+ output_dir = CASE
+ WHEN resume_blocked = 1 OR ? = 1 THEN NULL
+ WHEN ? IS NOT NULL THEN ?
+ WHEN ? IN ('error', 'stopped') THEN output_dir
+ ELSE NULL
+ END,
+ error_message = ?,
+ resume_blocked = CASE WHEN resume_blocked = 1 OR ? = 1 THEN 1 ELSE ? END
+ WHERE id = ? AND status = 'running'
""",
(
status,
@@ -804,8 +845,13 @@ def finish_run(
final_loss,
duration_seconds,
loss_sparkline,
+ int(clear_output_dir),
output_dir,
+ output_dir,
+ status,
error_message,
+ int(clear_output_dir),
+ int(resume_blocked),
id,
),
)
@@ -865,6 +911,38 @@ def update_run_display_name(id: str, display_name: Optional[str]) -> None:
conn.close()
+def update_run_output_dir(id: str, output_dir: Optional[str]) -> None:
+ conn = get_connection()
+ try:
+ conn.execute(
+ """
+ UPDATE training_runs SET output_dir = ?
+ WHERE id = ? AND status = 'running' AND resume_blocked = 0
+ """,
+ (output_dir, id),
+ )
+ conn.commit()
+ finally:
+ conn.close()
+
+
+def mark_run_cancel_requested(id: str) -> bool:
+ """Clear resume/export state only while the exact run is still active."""
+ conn = get_connection()
+ try:
+ cursor = conn.execute(
+ """
+ UPDATE training_runs SET output_dir = NULL, resume_blocked = 1
+ WHERE id = ? AND status = 'running'
+ """,
+ (id,),
+ )
+ conn.commit()
+ return cursor.rowcount > 0
+ finally:
+ conn.close()
+
+
def list_runs(limit: int = 50, offset: int = 0) -> dict:
conn = get_connection()
try:
@@ -874,15 +952,15 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
r.ended_at, r.total_steps, r.final_step, r.final_loss,
r.output_dir, r.duration_seconds, r.error_message,
- r.loss_sparkline, r.display_name, r.config_json,
+ r.loss_sparkline, r.display_name, r.config_json, r.resume_blocked,
CASE
- WHEN r.status = 'stopped'
+ WHEN r.status IN ('stopped', 'error')
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
- AND newer.status IN ('stopped', 'completed')
+ AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
@@ -917,13 +995,13 @@ def get_run(id: str) -> Optional[dict]:
"""
SELECT r.*,
CASE
- WHEN r.status = 'stopped'
+ WHEN r.status IN ('stopped', 'error')
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
- AND newer.status IN ('stopped', 'completed')
+ AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
@@ -958,12 +1036,12 @@ def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]:
0 AS resumed_later
FROM training_runs r
WHERE r.output_dir = ?
- AND r.status = 'stopped'
+ AND r.status IN ('stopped', 'error')
AND NOT EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
- AND newer.status IN ('stopped', 'completed')
+ AND newer.status IN ('stopped', 'completed', 'error', 'running')
AND newer.started_at > r.started_at
)
ORDER BY r.started_at DESC
@@ -1066,8 +1144,12 @@ def cleanup_orphaned_runs() -> None:
conn.execute(
"""
UPDATE training_runs
- SET status = 'error',
- error_message = 'Server restarted during training',
+ SET status = CASE WHEN resume_blocked = 1 THEN 'stopped' ELSE 'error' END,
+ error_message = CASE
+ WHEN resume_blocked = 1 THEN NULL
+ ELSE 'Server restarted during training'
+ END,
+ output_dir = CASE WHEN resume_blocked = 1 THEN NULL ELSE output_dir END,
ended_at = ?
WHERE status = 'running'
""",
diff --git a/studio/backend/tests/test_mlx_stop_checkpoint.py b/studio/backend/tests/test_mlx_stop_checkpoint.py
new file mode 100644
index 0000000000..d4a00cc6c8
--- /dev/null
+++ b/studio/backend/tests/test_mlx_stop_checkpoint.py
@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Regression tests for MLX stop-and-save checkpoint handling."""
+
+import importlib.util
+import json
+import sys
+import types
+from pathlib import Path
+
+import numpy as np
+from safetensors.numpy import save_file
+
+
+_BACKEND = Path(__file__).resolve().parents[1]
+
+
+def _load_worker_module():
+ spec = importlib.util.spec_from_file_location(
+ "training_worker_under_test",
+ _BACKEND / "core" / "training" / "worker.py",
+ )
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+worker = _load_worker_module()
+
+
+class _FakeTrainer:
+ def __init__(self, step: int):
+ self._global_step = step
+ self._train_loss_history = []
+ self.model = object()
+
+
+def _write_checkpoint(out: Path, step: int) -> Path:
+ checkpoint = out / f"checkpoint-{step}"
+ checkpoint.mkdir(parents = True, exist_ok = True)
+ (checkpoint / "trainer_state.json").write_text(
+ json.dumps({"global_step": step}), encoding = "utf-8"
+ )
+ save_file({"weight": np.ones(1, dtype = np.float32)}, checkpoint / "adapters.safetensors")
+ save_file(
+ {"state": np.ones(1, dtype = np.float32)},
+ checkpoint / "optimizer_state.safetensors",
+ )
+ return checkpoint
+
+
+def test_mlx_has_checkpoint_at_step_requires_complete_state(tmp_path):
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 5)
+
+ assert worker._mlx_has_checkpoint_at_step(out, 5) is True
+
+
+def test_write_mlx_stop_checkpoint_returns_true_when_current_step_checkpoint_exists(tmp_path):
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 5)
+
+ assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is True
+
+
+def test_write_mlx_stop_checkpoint_writes_current_step_when_only_older_checkpoint_exists(
+ tmp_path, monkeypatch
+):
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 5)
+
+ saved_steps: list[int] = []
+
+ def _save_state(_value, path, name):
+ save_file({"state": np.ones(1, dtype = np.float32)}, Path(path, name))
+
+ def _save_trainer_state(state, ckpt_dir, **_kwargs):
+ Path(ckpt_dir, "trainer_state.json").write_text(json.dumps(state), encoding = "utf-8")
+ saved_steps.append(int(state["global_step"]))
+
+ fake_utils = types.SimpleNamespace(
+ save_trainable_adapters = lambda model, path: _save_state(
+ model, path, "adapters.safetensors"
+ ),
+ save_optimizer_state = lambda optimizer, path: _save_state(
+ optimizer, path, "optimizer_state.safetensors"
+ ),
+ save_trainer_state = _save_trainer_state,
+ )
+ monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils)
+
+ assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), object(), out) is True
+ assert saved_steps == [10]
+ assert (out / "checkpoint-10" / "trainer_state.json").is_file()
+
+
+def test_write_mlx_stop_checkpoint_returns_false_without_optimizer(tmp_path):
+ out = tmp_path / "outputs" / "run_x"
+ out.mkdir(parents = True)
+
+ assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False
+
+
+def test_write_mlx_stop_checkpoint_rejects_incomplete_current_checkpoint(tmp_path):
+ out = tmp_path / "outputs" / "run_x"
+ ckpt = out / "checkpoint-5"
+ ckpt.mkdir(parents = True)
+ (ckpt / "trainer_state.json").write_text('{"global_step": 5}', encoding = "utf-8")
+
+ assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False
+
+
+def test_write_mlx_stop_checkpoint_ignores_stale_checkpoint_without_optimizer(tmp_path):
+ # An older checkpoint does not cover the current step, so this still fails.
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 5)
+
+ assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), None, out) is False
+
+
+def test_write_mlx_stop_checkpoint_returns_false_when_save_fails(tmp_path, monkeypatch):
+ out = tmp_path / "outputs" / "run_x"
+ out.mkdir(parents = True)
+
+ def _boom(*_args, **_kwargs):
+ raise RuntimeError("save failed")
+
+ fake_utils = types.SimpleNamespace(
+ save_trainable_adapters = _boom,
+ save_optimizer_state = lambda *_a, **_k: None,
+ save_trainer_state = lambda *_a, **_k: None,
+ )
+ monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils)
+
+ assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is False
diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py
index d75b205f35..5d2a218482 100644
--- a/studio/backend/tests/test_training_pump_resilience.py
+++ b/studio/backend/tests/test_training_pump_resilience.py
@@ -310,6 +310,80 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch):
assert b._pump_running is False
+def test_interrupted_cancel_clears_in_memory_output_dir(monkeypatch):
+ # Stop-without-save interrupted before its complete event: /status must not
+ # keep serving the cleared run's output_dir.
+ b = TrainingBackend()
+ finalized: dict = {}
+ monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
+ monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
+
+ b._proc = _FakeProc(alive = False)
+ b._event_queue = _IdleQueue()
+ b._progress.is_training = True
+ b._should_stop = True
+ b._cancel_requested = True
+ b._output_dir = "/out/x"
+
+ b._pump_loop()
+
+ assert b._output_dir is None
+ assert finalized.get("status") == "stopped"
+ assert finalized.get("output_dir") is None
+ assert finalized.get("clear_output_dir") is True
+
+
+def test_worker_exit_reuses_terminal_stop_save_error(monkeypatch):
+ b = TrainingBackend()
+ finalized: dict = {}
+ monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
+ monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
+
+ b._proc = _FakeProc(alive = False)
+ b._event_queue = _IdleQueue()
+ b._progress.is_training = True
+ b._should_stop = True
+ b._cancel_requested = False
+ b._output_dir = "/out/x"
+ b.current_job_id = "job-x"
+ b._terminal_finalize_payload = {
+ "status": "error",
+ "error_message": "checkpoint failed",
+ "output_dir": "/out/x",
+ "clear_output_dir": False,
+ "resume_blocked": True,
+ "expected_job_id": "job-x",
+ }
+
+ b._pump_loop()
+
+ assert b._output_dir == "/out/x"
+ assert finalized.get("status") == "error"
+ assert finalized.get("output_dir") == "/out/x"
+ assert finalized.get("clear_output_dir") is False
+ assert finalized.get("resume_blocked") is True
+
+
+def test_dead_worker_crash_preserves_output_dir(monkeypatch):
+ # A crash (no stop requested) after output_dir was emitted must keep the dir
+ # in the error finalize: checkpoints under it may still exist.
+ b = TrainingBackend()
+ finalized: dict = {}
+ monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
+ monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
+
+ b._proc = _FakeProc(alive = False)
+ b._event_queue = _IdleQueue()
+ b._progress.is_training = True
+ b._output_dir = "/out/x"
+
+ b._pump_loop()
+
+ assert finalized.get("status") == "error"
+ assert finalized.get("output_dir") == "/out/x"
+ assert finalized.get("clear_output_dir") is False
+
+
def test_start_training_clears_stale_pump_running_flag():
# A prior pump that died abnormally leaves _pump_running True. The next
# start_training must clear it during reset so the start-time watchdog can't
diff --git a/studio/backend/tests/test_training_resume.py b/studio/backend/tests/test_training_resume.py
index 91fdac9961..51425b0428 100644
--- a/studio/backend/tests/test_training_resume.py
+++ b/studio/backend/tests/test_training_resume.py
@@ -7,6 +7,9 @@ import importlib.util
import json
from pathlib import Path
+import pytest
+import torch
+
_BACKEND = Path(__file__).resolve().parents[1]
@@ -25,6 +28,30 @@ def _load_resume_module():
resume = _load_resume_module()
+def test_resume_request_accepts_sanitized_null_target_modules():
+ from models.training import TrainingStartRequest
+ request = TrainingStartRequest(
+ model_name = "unsloth/Qwen3-0.6B",
+ training_type = "Full Finetuning",
+ format_type = "alpaca",
+ target_modules = None,
+ )
+
+ assert request.target_modules == []
+
+
+def _write_checkpoint(out: Path, step: int) -> Path:
+ checkpoint = out / f"checkpoint-{step}"
+ checkpoint.mkdir(parents = True, exist_ok = True)
+ (checkpoint / "trainer_state.json").write_text(
+ json.dumps({"global_step": step}), encoding = "utf-8"
+ )
+ torch.save({"weight": torch.ones(1)}, checkpoint / "adapter_model.bin")
+ torch.save({"state": {0: torch.ones(1)}}, checkpoint / "optimizer.pt")
+ torch.save({"last_epoch": step}, checkpoint / "scheduler.pt")
+ return checkpoint
+
+
def _stopped_run(**overrides):
run = {
"status": "stopped",
@@ -44,6 +71,36 @@ def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch):
assert resume.can_resume_run(_stopped_run()) is True
+def test_can_resume_run_allows_errored_run_with_checkpoint(monkeypatch):
+ monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
+
+ assert resume.can_resume_run(_stopped_run(status = "error")) is True
+
+
+def test_can_resume_run_rejects_errored_run_without_checkpoint(monkeypatch):
+ monkeypatch.setattr(resume, "has_resume_state", lambda _path: False)
+
+ assert resume.can_resume_run(_stopped_run(status = "error")) is False
+
+
+def test_can_resume_run_allows_errored_run_at_final_step(monkeypatch):
+ # A save-time crash records final_step == total_steps; resuming re-runs the
+ # final-save path from the checkpoint.
+ monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
+
+ run = _stopped_run(status = "error", final_step = 10, total_steps = 10)
+
+ assert resume.can_resume_run(run) is True
+
+
+def test_can_resume_run_rejects_stopped_run_at_final_step(monkeypatch):
+ monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
+
+ run = _stopped_run(final_step = 10, total_steps = 10)
+
+ assert resume.can_resume_run(run) is False
+
+
def test_can_resume_run_rejects_s3_dataset_source(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
@@ -91,3 +148,444 @@ def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path)
result = studio_db.list_runs()
assert result["runs"][0]["config_json"] == config_json
+
+
+def test_crashed_run_with_persisted_output_dir_is_resumable(monkeypatch, tmp_path):
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 10)
+
+ studio_db.create_run(
+ id = "run-crash",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 20,
+ )
+ studio_db.update_run_output_dir("run-crash", str(out))
+ conn = studio_db.get_connection()
+ conn.execute("UPDATE training_runs SET status = 'error' WHERE id = 'run-crash'")
+ conn.commit()
+ conn.close()
+
+ run = studio_db.get_run("run-crash")
+ assert run["output_dir"] == str(out)
+ assert resume.can_resume_run(run) is True
+
+
+def test_checkpoint_discovery_skips_malformed_newest(monkeypatch, tmp_path):
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ out = tmp_path / "outputs" / "run_x"
+ valid = _write_checkpoint(out, 5)
+ (_write_checkpoint(out, 8) / "scheduler.pt").unlink()
+ malformed = out / "checkpoint-10"
+ malformed.mkdir()
+ (malformed / "trainer_state.json").write_text(json.dumps({"global_step": 10}), encoding = "utf-8")
+ (malformed / "adapter_model.bin").write_bytes(b"not a torch archive")
+ (malformed / "optimizer.pt").write_bytes(b"not a torch archive")
+
+ assert resume.get_resume_checkpoint_path(str(out)) == str(valid)
+
+
+def test_completed_run_keeps_output_dir_and_rejects_stale_cancel(monkeypatch, tmp_path):
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ studio_db.create_run(
+ id = "r",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 10,
+ )
+ studio_db.update_run_output_dir("r", "/out/x")
+ studio_db.finish_run(
+ id = "r",
+ status = "completed",
+ ended_at = "t",
+ final_step = 2,
+ final_loss = None,
+ duration_seconds = 1,
+ loss_sparkline = "[]",
+ output_dir = "/out/x",
+ error_message = None,
+ )
+
+ assert studio_db.get_run("r")["output_dir"] == "/out/x"
+ assert studio_db.mark_run_cancel_requested("r") is False
+ assert studio_db.get_run("r")["output_dir"] == "/out/x"
+ assert studio_db.get_run("r")["resume_blocked"] == 0
+
+
+def test_finish_run_clears_output_dir_for_stop_without_save(monkeypatch, tmp_path):
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ studio_db.create_run(
+ id = "r",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 10,
+ )
+ studio_db.update_run_output_dir("r", "/out/x")
+ studio_db.finish_run(
+ id = "r",
+ status = "stopped",
+ ended_at = "t",
+ final_step = 2,
+ final_loss = None,
+ duration_seconds = 1,
+ loss_sparkline = "[]",
+ output_dir = None,
+ error_message = None,
+ clear_output_dir = True,
+ )
+
+ assert studio_db.get_run("r")["output_dir"] is None
+ conn = studio_db.get_connection()
+ conn.execute(
+ "UPDATE training_runs SET status = 'running', output_dir = '/out/x', resume_blocked = 0 WHERE id = 'r'"
+ )
+ conn.commit()
+ conn.close()
+ studio_db.mark_run_cancel_requested("r")
+ studio_db.cleanup_orphaned_runs()
+ assert studio_db.get_run("r")["status"] == "stopped"
+ assert studio_db.get_run("r")["output_dir"] is None
+
+
+def test_finish_run_clears_output_dir_on_cancel_error_finalize(monkeypatch, tmp_path):
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ studio_db.create_run(
+ id = "r",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 10,
+ )
+ studio_db.update_run_output_dir("r", "/out/x")
+ studio_db.finish_run(
+ id = "r",
+ status = "stopped",
+ ended_at = "t",
+ final_step = 2,
+ final_loss = None,
+ duration_seconds = 1,
+ loss_sparkline = "[]",
+ output_dir = "/out/x",
+ error_message = "worker failed during cancel",
+ clear_output_dir = True,
+ )
+
+ assert studio_db.get_run("r")["output_dir"] is None
+
+
+def test_finish_run_preserves_output_dir_for_interrupted_stop_and_save(monkeypatch, tmp_path):
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ studio_db.create_run(
+ id = "r",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 10,
+ )
+ studio_db.update_run_output_dir("r", "/out/x")
+ studio_db.finish_run(
+ id = "r",
+ status = "stopped",
+ ended_at = "t",
+ final_step = 2,
+ final_loss = None,
+ duration_seconds = 1,
+ loss_sparkline = "[]",
+ output_dir = None,
+ error_message = None,
+ )
+
+ assert studio_db.get_run("r")["output_dir"] == "/out/x"
+
+
+def test_resumed_errored_run_is_not_offered_again(monkeypatch, tmp_path):
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 10)
+
+ studio_db.create_run(
+ id = "run-old",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 20,
+ )
+ studio_db.update_run_output_dir("run-old", str(out))
+ studio_db.finish_run(
+ id = "run-old",
+ status = "error",
+ ended_at = "2026-01-01T00:05:00Z",
+ final_step = 10,
+ final_loss = None,
+ duration_seconds = 1,
+ loss_sparkline = "[]",
+ output_dir = None,
+ error_message = "killed",
+ )
+ studio_db.create_run(
+ id = "run-new",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-02T00:00:00Z",
+ total_steps = 20,
+ output_dir = str(out),
+ resumed_from_run_id = "run-old",
+ )
+ with pytest.raises(RuntimeError, match = "no longer available"):
+ studio_db.create_run(
+ id = "run-duplicate",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-02T00:00:01Z",
+ total_steps = 20,
+ output_dir = str(out),
+ resumed_from_run_id = "run-old",
+ )
+ assert studio_db.get_run("run-duplicate") is None
+ studio_db.finish_run(
+ id = "run-new",
+ status = "error",
+ ended_at = "2026-01-02T00:05:00Z",
+ final_step = 15,
+ final_loss = None,
+ duration_seconds = 1,
+ loss_sparkline = "[]",
+ output_dir = None,
+ error_message = "killed again",
+ )
+
+ old_run = studio_db.get_run("run-old")
+ new_run = studio_db.get_run("run-new")
+ assert old_run["resumed_later"] == 1
+ assert resume.can_resume_run(old_run) is False
+ assert new_run["resumed_later"] == 0
+ assert resume.can_resume_run(new_run) is True
+ assert studio_db.get_resumable_run_by_output_dir(str(out))["id"] == "run-new"
+
+
+def test_running_continuation_blocks_older_resume(monkeypatch, tmp_path):
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 10)
+
+ studio_db.create_run(
+ id = "run-old",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 20,
+ )
+ studio_db.update_run_output_dir("run-old", str(out))
+ studio_db.finish_run(
+ id = "run-old",
+ status = "error",
+ ended_at = "2026-01-01T00:05:00Z",
+ final_step = 10,
+ final_loss = None,
+ duration_seconds = 1,
+ loss_sparkline = "[]",
+ output_dir = None,
+ error_message = "killed",
+ )
+ studio_db.create_run(
+ id = "run-new",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-02T00:00:00Z",
+ total_steps = 20,
+ output_dir = str(out),
+ resumed_from_run_id = "run-old",
+ )
+
+ old_run = studio_db.get_run("run-old")
+ assert old_run["resumed_later"] == 1
+ assert resume.can_resume_run(old_run) is False
+ assert studio_db.get_resumable_run_by_output_dir(str(out)) is None
+
+
+def test_stop_save_checkpoint_failure_keeps_error_status(monkeypatch, tmp_path):
+ # A stop-and-save whose checkpoint write failed must finalize as an error so
+ # history explains the missing resume state (keep_error_status flag).
+ from core.training.training import TrainingBackend
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ studio_db.create_run(
+ id = "run-failed-save",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 10,
+ )
+ backend = TrainingBackend()
+ backend.current_job_id = "run-failed-save"
+ backend._db_run_created = True
+ backend._should_stop = True
+ backend._handle_event(
+ {
+ "type": "error",
+ "error": "Failed to save a resumable checkpoint after stop.",
+ "keep_error_status": True,
+ }
+ )
+
+ run = studio_db.get_run("run-failed-save")
+ assert run["status"] == "error"
+ assert "resumable checkpoint" in run["error_message"]
+
+
+def test_can_resume_run_rejects_resume_blocked_run(monkeypatch):
+ monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
+
+ assert resume.can_resume_run(_stopped_run(status = "error", resume_blocked = 1)) is False
+
+
+def test_stop_save_checkpoint_failure_with_stale_checkpoint_is_not_resumable(monkeypatch, tmp_path):
+ # A failed stop-and-save must not offer Resume from an older periodic checkpoint;
+ # that would roll back past the recorded final step.
+ from core.training.training import TrainingBackend
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ out = tmp_path / "outputs" / "run_x"
+ _write_checkpoint(out, 10)
+
+ studio_db.create_run(
+ id = "run-stale-ckpt",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 20,
+ )
+ studio_db.update_run_output_dir("run-stale-ckpt", str(out))
+ backend = TrainingBackend()
+ backend.current_job_id = "run-stale-ckpt"
+ backend._db_run_created = True
+ backend._should_stop = True
+ backend._output_dir = str(out)
+ backend._handle_event(
+ {
+ "type": "error",
+ "error": "Failed to save a resumable checkpoint after stop.",
+ "keep_error_status": True,
+ "resume_blocked": True,
+ }
+ )
+
+ run = studio_db.get_run("run-stale-ckpt")
+ assert run["status"] == "error"
+ assert run["resume_blocked"] == 1
+ assert run["output_dir"] == str(out)
+ assert resume.can_resume_run(run) is False
+
+
+def test_user_stop_error_without_checkpoint_ack_is_blocked(monkeypatch, tmp_path):
+ from core.training.training import TrainingBackend
+ from storage import studio_db
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(studio_db, "_schema_ready", False)
+
+ studio_db.create_run(
+ id = "run-user-stop",
+ model_name = "m",
+ dataset_name = "d",
+ config_json = "{}",
+ started_at = "2026-01-01T00:00:00Z",
+ total_steps = 10,
+ )
+ backend = TrainingBackend()
+ backend.current_job_id = "run-user-stop"
+ backend._db_run_created = True
+ backend._should_stop = True
+ backend._handle_event({"type": "error", "error": "interrupted"})
+
+ run = studio_db.get_run("run-user-stop")
+ assert run["status"] == "error" and run["resume_blocked"] == 1
+
+
+def test_terminal_fallback_keeps_resumable_when_current_checkpoint_landed(monkeypatch, tmp_path):
+ # Worker died before its terminal event, but a valid current-step checkpoint
+ # is on disk: the fallback must keep the run resumable, not block it.
+ from core.training.training import TrainingBackend
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ out = tmp_path / "outputs" / "run_ok"
+ _write_checkpoint(out, 7)
+
+ backend = TrainingBackend()
+ backend.current_job_id = "run-ok"
+ backend._should_stop = True
+ backend._output_dir = str(out)
+ backend._progress.step = 7
+
+ kwargs = backend._terminal_finalize_kwargs()
+ assert kwargs["status"] == "stopped"
+ assert kwargs["resume_blocked"] is False
+
+
+def test_terminal_fallback_blocks_when_no_current_checkpoint(monkeypatch, tmp_path):
+ # Same path, but only a stale (older-step) checkpoint exists: must block.
+ from core.training.training import TrainingBackend
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ out = tmp_path / "outputs" / "run_stale"
+ _write_checkpoint(out, 5)
+
+ backend = TrainingBackend()
+ backend.current_job_id = "run-stale"
+ backend._should_stop = True
+ backend._output_dir = str(out)
+ backend._progress.step = 7
+
+ kwargs = backend._terminal_finalize_kwargs()
+ assert kwargs["status"] == "error"
+ assert kwargs["resume_blocked"] is True
diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py
index 0cd702bce2..cbe2082e82 100644
--- a/studio/backend/tests/test_training_stop_watchdog.py
+++ b/studio/backend/tests/test_training_stop_watchdog.py
@@ -353,7 +353,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch):
# stopped so the UI leaves "Stopping..." and a new run can start.
b = TrainingBackend()
finstop: list = []
- monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
+ monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
b._proc = _FakeProc(alive = True) # wedged: still reports alive
b._should_stop = True
@@ -365,7 +365,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch):
assert b._proc is None, "the wedged handle must be dropped so is_training_active clears"
assert b._progress.is_training is False
- assert b._progress.status_message == "Training stopped."
+ assert "valid current-step checkpoint" in b._progress.status_message
assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id"
assert b.is_training_active() is False
@@ -375,7 +375,7 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch):
# must record it even if the watchdog wins the finalize race against the pump.
b = TrainingBackend()
finstop: list = []
- monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
+ monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
b._proc = _FakeProc(alive = True)
b._should_stop = True
@@ -390,6 +390,28 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch):
assert finstop[0][1] == "/tmp/outputs/run-123"
+def test_finalize_after_escalation_clears_output_dir_on_cancel(monkeypatch):
+ # Stop-without-saving promises no resume: a cancel that escalates through the
+ # watchdog clears the persisted output_dir, not a checkpoint path.
+ b = TrainingBackend()
+ finstop: list = []
+ monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append((a, k)))
+
+ b._proc = _FakeProc(alive = True)
+ b._should_stop = True
+ b._cancel_requested = True
+ b.current_job_id = "job_c"
+ b._db_run_created = True
+ b._output_dir = "/tmp/outputs/run-123"
+
+ b._finalize_stopped_after_escalation(watched_job_id = "job_c")
+
+ assert finstop and finstop[0][0][0] == "job_c"
+ assert finstop[0][0][1] is None, "a cancelled run must not record a checkpoint path"
+ assert finstop[0][1].get("clear_output_dir") is True
+ assert b._output_dir is None, "/status must stop exposing the cancelled run's dir"
+
+
def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch):
# No worker -> nothing to escalate; the watchdog must not spawn.
b = TrainingBackend()
@@ -409,7 +431,7 @@ def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch):
# The escalation finalize must then leave the NEW run untouched, not drop its handle.
b = TrainingBackend()
finstop: list = []
- monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
+ monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
old_proc = _FakeProc(alive = False) # force-terminated worker we were watching
new_proc = _FakeProc(alive = True) # a new run already took over
@@ -430,7 +452,7 @@ def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch):
# finalizes the captured run by id.
b = TrainingBackend()
finstop: list = []
- monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
+ monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
proc = _FakeProc(alive = False)
b._proc = proc
@@ -451,7 +473,7 @@ def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypat
# catch this even though the proc-only guard would not.
b = TrainingBackend()
finstop: list = []
- monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
+ monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet
b._proc = old_proc # still the old handle (== target), so proc guard would pass
@@ -509,6 +531,7 @@ def _install_fake_db(monkeypatch):
recs["insert_ids"].append(job_id),
)
fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id"))
+ fake_db.mark_run_cancel_requested = lambda _run_id: True
fake_storage.studio_db = fake_db
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
@@ -518,9 +541,48 @@ def _install_fake_db(monkeypatch):
return recs
+def test_stop_without_save_creates_missing_row_before_signal(monkeypatch):
+ recs = _install_fake_db(monkeypatch)
+ b = TrainingBackend()
+ b.current_job_id, b._db_config = "job_missing", {"model_name": "m"}
+ b._stop_queue = queue.Queue()
+ assert b.stop_training(save = False) is True
+ assert [run["id"] for run in recs["created"]] == ["job_missing"]
+ assert b._stop_queue.get_nowait() == {"type": "stop", "save": False}
+
+ b._cancel_requested = b._should_stop = False
+ sys.modules["storage.studio_db"].mark_run_cancel_requested = lambda _run_id: False
+ assert b.stop_training(save = False) is False
+ assert not b._cancel_requested and b._stop_queue.empty()
+
+ new_queue = queue.Queue()
+ b.current_job_id, b._db_run_created = "job_old", True
+ b._cancel_requested = b._should_stop = False
+
+ def _supersede(_run_id):
+ b.current_job_id = "job_new"
+ b._stop_queue = new_queue
+ return True
+
+ sys.modules["storage.studio_db"].mark_run_cancel_requested = _supersede
+ assert b.stop_training(save = False) is False
+ assert not b._cancel_requested and new_queue.empty()
+
+
def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch):
# The watchdog and pump can both finalize; only one call may reach finish_run.
recs = _install_fake_db(monkeypatch)
+ monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0)
+ attempts = 0
+
+ def flaky_finish(**kw):
+ nonlocal attempts
+ attempts += 1
+ if attempts < 3:
+ raise RuntimeError("database is locked")
+ recs["finished"].append(kw)
+
+ sys.modules["storage.studio_db"].finish_run = flaky_finish
b = TrainingBackend()
b.current_job_id = "job_x"
b._db_run_created = True
@@ -539,6 +601,7 @@ def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch):
t.join(timeout = 5)
assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}"
+ assert attempts == 3
assert b._run_finalized is True
@@ -646,7 +709,13 @@ def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch):
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
- b._ensure_db_run_created()
+ b._run_intent_lock.acquire()
+ creator = threading.Thread(target = b._ensure_db_run_created)
+ creator.start()
+ time.sleep(0.02)
+ assert b._db_create_in_progress is False
+ b._run_intent_lock.release()
+ creator.join(timeout = 5)
assert observed["flag_during_create"] is False, "flag must not be published before insert"
assert observed["in_progress_during_create"] is True
@@ -718,6 +787,7 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch):
b = TrainingBackend()
b.current_job_id = "job_old"
b._db_run_created = True
+ b._should_stop = True
b._proc = _FakeProc(alive = False)
b._progress.is_training = True
b._progress.step = 42
@@ -726,7 +796,8 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch):
b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old")
assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id"
- assert recs["finished"][0]["status"] == "stopped"
+ assert recs["finished"][0]["status"] == "error"
+ assert recs["finished"][0]["resume_blocked"] is True
assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run"
assert b._metric_buffer == [], "the captured batch must be drained"
@@ -737,7 +808,7 @@ def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch):
# so the pump's create-then-finalize records the run. Parent state still clears.
b = TrainingBackend()
called: list = []
- monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a))
+ monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: called.append(a))
b._proc = _FakeProc(alive = False)
b.current_job_id = "job_q"
@@ -785,7 +856,7 @@ def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch):
new_proc = _FakeProc(alive = True)
b._proc = old_proc
- def hijack(*a):
+ def hijack(*a, **k):
b._proc = new_proc # a new run takes over during the finalize
monkeypatch.setattr(b, "_finish_stopped_run", hijack)
diff --git a/studio/frontend/src/features/studio/historical-training-view.tsx b/studio/frontend/src/features/studio/historical-training-view.tsx
index b6ec06b06a..4a123875a8 100644
--- a/studio/frontend/src/features/studio/historical-training-view.tsx
+++ b/studio/frontend/src/features/studio/historical-training-view.tsx
@@ -2,19 +2,29 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { TrainingViewData } from "@/features/training";
-import { getTrainingRun, onTrainingRunUpdated } from "@/features/training";
+import {
+ getTrainingRun,
+ onTrainingRunUpdated,
+ useTrainingActions,
+ useTrainingRuntimeStore,
+} from "@/features/training";
import type { TrainingRunDetailResponse } from "@/features/training";
import { parseBackendTrainingMethod } from "@/features/training/lib/training-methods";
import { type ReactElement, useEffect, useState } from "react";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
import { mapRunConfigToOverride } from "./sections/run-config-override";
+import { Button } from "@/components/ui/button";
+import { Spinner } from "@/components/ui/spinner";
+import { PlayIcon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
import { translate, useT } from "@/i18n";
type StudioT = ReturnType;
interface HistoricalTrainingViewProps {
runId: string;
+ onResumeStarted?: () => void;
}
function mapToViewData(
@@ -93,10 +103,27 @@ function mapToViewData(
export function HistoricalTrainingView({
runId,
+ onResumeStarted,
}: HistoricalTrainingViewProps): ReactElement {
const t = useT();
const [detail, setDetail] = useState(null);
const [error, setError] = useState(null);
+ const [resuming, setResuming] = useState(false);
+ const { resumeTrainingRunFromHistory } = useTrainingActions();
+ const isStarting = useTrainingRuntimeStore((state) => state.isStarting);
+ const isTrainingRunning = useTrainingRuntimeStore(
+ (state) => state.isTrainingRunning,
+ );
+
+ const handleResume = async () => {
+ setResuming(true);
+ try {
+ const ok = await resumeTrainingRunFromHistory(runId);
+ if (ok) onResumeStarted?.();
+ } finally {
+ setResuming(false);
+ }
+ };
// Derive loading from detail/error; no separate state.
const loading = detail === null && error === null;
@@ -152,6 +179,27 @@ export function HistoricalTrainingView({
return (
+ {detail.run.can_resume && (
+
+ void handleResume()}
+ >
+ {resuming ? (
+
+ ) : (
+
+ )}
+ {resuming
+ ? t("studio.history.resuming")
+ : t("studio.history.resumeTraining")}
+
+
+ )}
other.id !== run.id &&
other.output_dir === run.output_dir &&
- (other.status === "stopped" || other.status === "completed") &&
+ (other.status === "stopped" ||
+ other.status === "completed" ||
+ other.status === "error" ||
+ other.status === "running") &&
new Date(other.started_at).getTime() > startedAt,
);
}
diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx
index 4d3a32fab3..e575fbacd7 100644
--- a/studio/frontend/src/features/studio/studio-page.tsx
+++ b/studio/frontend/src/features/studio/studio-page.tsx
@@ -217,7 +217,13 @@ export function StudioPage(): ReactElement {
{selectedHistoryRunId ? (
-
+ {
+ setSelectedHistoryRunId(null);
+ handleTabChange("current-run");
+ }}
+ />
) : (
{
if (runId === currentJobId && isTrainingRunning) {
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 0d32e066d5..9bb8206f28 100644
--- a/studio/frontend/src/features/training/hooks/use-training-actions.ts
+++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts
@@ -217,7 +217,7 @@ export function useTrainingActions() {
const detail = await getTrainingRun(runId);
const outputDir = detail.run.output_dir;
if (!detail.run.can_resume || !outputDir) {
- throw new Error("Only stopped runs with a saved checkpoint can be resumed.");
+ throw new Error("Only stopped or errored runs with a saved checkpoint can be resumed.");
}
primeNativeNotificationPermission().catch(() => undefined);
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 acc80a3be2..131ea30d30 100644
--- a/studio/frontend/src/features/training/stores/training-runtime-store.ts
+++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts
@@ -223,7 +223,10 @@ export const useTrainingRuntimeStore = create()((set) => (
typeof detailLr === "number" ? detailLr : state.currentLearningRate,
currentEpoch:
typeof detailEpoch === "number" ? detailEpoch : state.currentEpoch,
- outputDir: payload.details?.output_dir ?? state.outputDir,
+ outputDir:
+ payload.details?.output_dir !== undefined
+ ? payload.details.output_dir
+ : state.outputDir,
lossHistory: metricHistory.lossHistory ?? state.lossHistory,
lrHistory: metricHistory.lrHistory ?? state.lrHistory,
gradNormHistory: metricHistory.gradNormHistory ?? state.gradNormHistory,
diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts
index 8ed0ce0037..373e876653 100644
--- a/studio/frontend/src/features/training/types/runtime.ts
+++ b/studio/frontend/src/features/training/types/runtime.ts
@@ -26,7 +26,8 @@ export interface TrainingStatusResponse {
total_steps?: number;
loss?: number;
learning_rate?: number;
- output_dir?: string;
+ // null = explicit clear (run stopped without saving); absent = unchanged.
+ output_dir?: string | null;
} | null;
metric_history?: {
steps?: number[];
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index 3e1c1e5a76..4f859aa91c 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -1053,7 +1053,8 @@ export const en = {
continueAction: "Continue Training",
cancelAction: "Cancel Training",
stopTitle: "Stop Training",
- stopDescription: "Choose how you want to stop the current training run.",
+ stopDescription:
+ "Choose how you want to stop the current training run. Stop and Save writes a checkpoint you can resume from later; Stop cannot be resumed.",
stopAction: "Stop",
stopping: "Stopping...",
stopAndSave: "Stop and Save",
diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts
index 37c73086cf..9984fc324f 100644
--- a/studio/frontend/src/i18n/locales/zh-CN.ts
+++ b/studio/frontend/src/i18n/locales/zh-CN.ts
@@ -901,7 +901,8 @@ export const zhCN = {
continueAction: "继续训练",
cancelAction: "取消训练",
stopTitle: "停止训练",
- stopDescription: "选择如何停止当前训练运行。",
+ stopDescription:
+ "选择如何停止当前训练运行。“停止并保存”会写入检查点,之后可从该处恢复;“停止”则无法恢复。",
stopAction: "停止",
stopping: "停止中...",
stopAndSave: "停止并保存",
From 54f21b3a8792dd77b93998e524aa179d13a95fed Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Tue, 21 Jul 2026 06:48:19 -0300
Subject: [PATCH 049/255] Studio: fix loading split GGUFs from the local HF
cache (#7273)
* Studio: don't resolve HF cache GGUF symlinks to blob paths
* Studio: handle split GGUF symlink layouts
---
.../tests/test_offline_gguf_cache_fallback.py | 60 +++++++++++++-
studio/backend/utils/models/model_config.py | 79 ++++++++++++++++++-
2 files changed, 134 insertions(+), 5 deletions(-)
diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py
index 295549c443..35dd3979b1 100644
--- a/studio/backend/tests/test_offline_gguf_cache_fallback.py
+++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py
@@ -119,6 +119,13 @@ def _build_cache(
return snap
+def _symlink_or_skip(link: Path, target: Path) -> None:
+ try:
+ link.symlink_to(target)
+ except OSError as exc:
+ pytest.skip(f"symlinks unavailable: {exc}")
+
+
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
"""Point ``huggingface_hub.constants.HF_HUB_CACHE`` at a temp dir."""
@@ -1084,7 +1091,7 @@ class TestListLocalGgufVariantsSubdir:
target.write_bytes(b"\0" * 20)
out = _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M")
- assert out == str(target.resolve())
+ assert out == str(target.absolute())
def test_find_local_gguf_by_variant_skips_big_endian_only_match(self, tmp_path):
from utils.models.model_config import _find_local_gguf_by_variant
@@ -1094,6 +1101,57 @@ class TestListLocalGgufVariantsSubdir:
assert _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") is None
+ def test_find_local_gguf_by_variant_keeps_split_symlink_name(self, tmp_path):
+ from utils.models.model_config import _find_local_gguf_by_variant
+
+ blobs = tmp_path / "blobs"
+ blobs.mkdir()
+ snap = tmp_path / "snapshots" / "rev" / "BF16"
+ snap.mkdir(parents = True)
+ (tmp_path / "snapshots" / "rev" / "config.json").write_text("{}")
+ for i, sha in enumerate(("aa" * 32, "bb" * 32), start = 1):
+ (blobs / sha).write_bytes(b"\0" * 10)
+ _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha)
+
+ out = _find_local_gguf_by_variant(str(tmp_path / "snapshots" / "rev"), "BF16")
+ assert out is not None
+ assert Path(out).name == "model-BF16-00001-of-00002.gguf"
+
+ def test_detect_gguf_model_keeps_split_symlink_name(self, tmp_path):
+ from utils.models.model_config import detect_gguf_model
+
+ blobs = tmp_path / "blobs"
+ blobs.mkdir()
+ snap = tmp_path / "snapshots" / "rev"
+ snap.mkdir(parents = True)
+ for i, (sha, size) in enumerate((("cc" * 32, 10), ("dd" * 32, 20)), start = 1):
+ (blobs / sha).write_bytes(b"\0" * size)
+ _symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha)
+
+ out = detect_gguf_model(str(snap))
+ assert out is not None
+ assert Path(out).name == "model-BF16-00001-of-00002.gguf"
+
+ def test_lone_split_symlink_uses_colocated_target_shards(self, tmp_path):
+ from utils.models.model_config import _find_local_gguf_by_variant, detect_gguf_model
+
+ target_dir = tmp_path / "external" / "BF16"
+ target_dir.mkdir(parents = True)
+ target = target_dir / "model-BF16-00001-of-00002.gguf"
+ target.write_bytes(b"\0" * 10)
+ (target_dir / "model-BF16-00002-of-00002.gguf").write_bytes(b"\0" * 10)
+
+ local = tmp_path / "local"
+ local.mkdir()
+ (local / "config.json").write_text("{}")
+ link = local / target.name
+ _symlink_or_skip(link, target)
+
+ expected = str(target.absolute())
+ assert _find_local_gguf_by_variant(str(local), "BF16") == expected
+ assert detect_gguf_model(str(local)) == expected
+ assert detect_gguf_model(str(link)) == expected
+
def test_model_config_variant_ignores_big_endian_sibling(self, tmp_path):
from utils.models.model_config import ModelConfig
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index dadf103cea..821529083d 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -1249,6 +1249,77 @@ def _iter_gguf_files(directory: Path, recursive: bool = False):
yield f
+_GGUF_SPLIT_FILE_RE = re.compile(
+ r"^(?P.+)-(?P\d{5})-of-(?P\d{5})\.gguf$",
+ re.IGNORECASE,
+)
+
+
+def _colocated_first_split_shard(path: Path) -> tuple[Optional[Path], bool]:
+ """Return shard 1 and whether every shard is beside *path*."""
+ match = _GGUF_SPLIT_FILE_RE.match(path.name)
+ if match is None:
+ return None, False
+
+ prefix = match.group("prefix").casefold()
+ total_text = match.group("total")
+ total = int(total_text)
+ if total < 1:
+ return None, False
+
+ first: Optional[Path] = None
+ indices: set[int] = set()
+ try:
+ siblings = path.parent.iterdir()
+ for sibling in siblings:
+ sibling_match = _GGUF_SPLIT_FILE_RE.match(sibling.name)
+ if (
+ sibling_match is None
+ or sibling_match.group("prefix").casefold() != prefix
+ or sibling_match.group("total") != total_text
+ ):
+ continue
+ try:
+ if not sibling.is_file():
+ continue
+ except OSError:
+ continue
+ index = int(sibling_match.group("index"))
+ if not 1 <= index <= total:
+ continue
+ indices.add(index)
+ if index == 1:
+ first = sibling
+ except OSError:
+ return None, False
+
+ return first, first is not None and len(indices) == total
+
+
+def _local_gguf_load_path(path: Path) -> Path:
+ """Choose a loadable local path while preserving complete symlink sets."""
+ if _GGUF_SPLIT_FILE_RE.match(path.name) is None:
+ return path.absolute()
+
+ first, complete = _colocated_first_split_shard(path)
+ if complete and first is not None:
+ return first.absolute()
+
+ try:
+ is_symlink = path.is_symlink()
+ except OSError:
+ is_symlink = False
+ if is_symlink:
+ try:
+ target = path.resolve()
+ except OSError:
+ return (first or path).absolute()
+ target_first, _ = _colocated_first_split_shard(target)
+ return (target_first or target).absolute()
+
+ return (first or path).absolute()
+
+
def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]:
"""Find the mmproj GGUF for a model.
@@ -1434,7 +1505,7 @@ def detect_gguf_model(path: str) -> Optional[str]:
except OSError:
is_dir = False # stat() unavailable in the lock window
if not is_dir:
- return str(p.absolute()) # absolute() keeps symlink names readable
+ return str(_local_gguf_load_path(p))
# Directory named "*.gguf": fall through to the dir scan below.
# Case 2: directory containing .gguf files (skip mmproj / MTP drafter)
@@ -1452,7 +1523,7 @@ def detect_gguf_model(path: str) -> Optional[str]:
gguf_files.append(f)
gguf_files.sort(key = lambda f: f.stat().st_size, reverse = True)
if gguf_files:
- return str(gguf_files[0].resolve())
+ return str(_local_gguf_load_path(gguf_files[0]))
return None
@@ -1879,7 +1950,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
For sharded GGUFs (multiple files sharing a quant label), returns the
first shard (sorted by name), which is what ``llama-server -m`` expects.
- Returns the resolved absolute path, or ``None`` if no match.
+ Returns the absolute path, or ``None`` if no match.
"""
p = _resolve_gguf_dir(Path(directory))
if p is None:
@@ -1900,7 +1971,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
matches.append(f)
matches.sort()
if matches:
- return str(matches[0].resolve())
+ return str(_local_gguf_load_path(matches[0]))
return None
From f6359805a824b7a306de5b9c5cebc01fd7c9261a Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Tue, 21 Jul 2026 03:32:40 -0700
Subject: [PATCH 050/255] Show iconless models in Hub feed above likes
threshold (#7284)
Co-authored-by: shimmyshimmer
---
studio/frontend/src/features/hub/hub-page.tsx | 9 +++++++--
1 file changed, 7 insertions(+), 2 deletions(-)
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx
index 426f816dbe..02094be69a 100644
--- a/studio/frontend/src/features/hub/hub-page.tsx
+++ b/studio/frontend/src/features/hub/hub-page.tsx
@@ -111,6 +111,10 @@ const ALL_MODELS_VIEW_STORAGE_KEY = "unsloth.hub.allModelsView";
const INVENTORY_SORT_STORAGE_KEY = "unsloth.hub.inventorySort";
const OWNER_SCOPE_STORAGE_KEY = "unsloth.hub.ownerScope";
+// Iconless models (repo name matches no provider logo, e.g. Ornith, Inkling)
+// still show in the feed once they clear this many Hugging Face likes.
+const MIN_ICONLESS_MODEL_LIKES = 30;
+
/** Discover browsing scope: the whole Hub (default) or only the unsloth org. */
export type OwnerScope = "unsloth" | "all";
@@ -739,9 +743,10 @@ export function ModelsPage() {
(row) =>
!isHiddenModelId(row.id) &&
!isConfiguredHiddenModelId(hiddenEmbeddingModelIds, row.id) &&
- // The default feed only shows models with a provider logo.
+ // Feed shows logo'd models, plus iconless ones above the likes threshold.
(!isFeedMode ||
- resolveOwnerProviderLogo(row.owner, row.repo) !== null) &&
+ resolveOwnerProviderLogo(row.owner, row.repo) !== null ||
+ (row.result.likes ?? 0) >= MIN_ICONLESS_MODEL_LIKES) &&
matchesFormat(
detectResultFormat(row.result),
effectiveDiscoverFormat,
From f5da223c221fc25b458b17bd781c2364cef5b5db Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Tue, 21 Jul 2026 07:48:49 -0300
Subject: [PATCH 051/255] Installer: report the installed Unsloth version
(#7265)
* Installer: report the installed Unsloth version
* [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>
---
install.ps1 | 7 +++
install.sh | 9 +++
tests/test_installer_unsloth_version.py | 75 +++++++++++++++++++++++++
3 files changed, 91 insertions(+)
create mode 100644 tests/test_installer_unsloth_version.py
diff --git a/install.ps1 b/install.ps1
index a525d4df56..524fd85774 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2420,6 +2420,13 @@ exit 0
}
}
+ $installedPackageVersion = (& $VenvPython -c "from importlib.metadata import version; import sys; print(version(sys.argv[1]))" $PackageName 2>$null | Out-String).Trim()
+ if ($LASTEXITCODE -eq 0 -and $installedPackageVersion) {
+ step $PackageName "$installedPackageVersion installed"
+ } else {
+ substep "[WARN] installed $PackageName version could not be determined" "Yellow"
+ }
+
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
diff --git a/install.sh b/install.sh
index 0acccdf049..445bab616a 100755
--- a/install.sh
+++ b/install.sh
@@ -3396,6 +3396,15 @@ else
fi
fi
+_installed_package_version=$("$_VENV_PY" -c \
+ 'from importlib.metadata import version; import sys; print(version(sys.argv[1]))' \
+ "$PACKAGE_NAME" 2>/dev/null || true)
+if [ -n "$_installed_package_version" ]; then
+ step "$PACKAGE_NAME" "$_installed_package_version installed"
+else
+ substep "[WARN] installed $PACKAGE_NAME version could not be determined" "$C_WARN"
+fi
+
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on
diff --git a/tests/test_installer_unsloth_version.py b/tests/test_installer_unsloth_version.py
new file mode 100644
index 0000000000..5ee1cfac71
--- /dev/null
+++ b/tests/test_installer_unsloth_version.py
@@ -0,0 +1,75 @@
+"""Regression coverage for installer version reporting."""
+
+from __future__ import annotations
+
+import re
+import shutil
+import subprocess
+import sys
+from importlib.metadata import version
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+INSTALL_SH = REPO_ROOT / "install.sh"
+INSTALL_PS1 = REPO_ROOT / "install.ps1"
+
+
+def _extract(pattern: str, source: str) -> str:
+ match = re.search(pattern, source, flags = re.DOTALL | re.MULTILINE)
+ assert match is not None, f"installer block not found: {pattern}"
+ return match.group(0)
+
+
+@pytest.mark.skipif(shutil.which("sh") is None, reason = "POSIX shell is unavailable")
+def test_posix_installer_reports_installed_distribution_version():
+ source = INSTALL_SH.read_text(encoding = "utf-8")
+ reporter = _extract(
+ r"_installed_package_version=\$\(.*?^fi",
+ source,
+ )
+ result = subprocess.run(
+ [
+ "sh",
+ "-c",
+ (
+ 'step() { printf "%s %s\\n" "$1" "$2"; }\n'
+ 'substep() { printf "WARN %s\\n" "$1"; }\n'
+ f"_VENV_PY={sys.executable!r}\n"
+ "PACKAGE_NAME=pytest\n"
+ f"{reporter}"
+ ),
+ ],
+ check = True,
+ capture_output = True,
+ text = True,
+ )
+ assert result.stdout.strip() == f"pytest {version('pytest')} installed"
+
+
+@pytest.mark.skipif(shutil.which("pwsh") is None, reason = "PowerShell is unavailable")
+def test_windows_version_reporter_uses_distribution_metadata():
+ source = INSTALL_PS1.read_text(encoding = "utf-8")
+ reporter = _extract(
+ r" \$installedPackageVersion = .*?^ if .*?^ \} else \{.*?^ \}",
+ source,
+ )
+ result = subprocess.run(
+ [
+ "pwsh",
+ "-NoProfile",
+ "-NonInteractive",
+ "-Command",
+ (
+ 'function step { param($Label, $Value) Write-Output "$Label $Value" }; '
+ 'function substep { param($Message, $Color) Write-Output "WARN $Message" }; '
+ f"$VenvPython = '{sys.executable}'; $PackageName = 'pytest'; "
+ f"{reporter}"
+ ),
+ ],
+ check = True,
+ capture_output = True,
+ text = True,
+ )
+ assert result.stdout.strip() == f"pytest {version('pytest')} installed"
From 77da6e8fcbd3ff0a60d599866f6a797a1078f7cd Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Tue, 21 Jul 2026 11:53:11 +0100
Subject: [PATCH 052/255] Studio: show HF token tick only after validation
(#7268)
* Studio: show HF token tick only after validation
* Studio: prevent stale HF token validation state
* [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>
---
.../features/settings/tabs/general-tab.tsx | 15 +++++-----
.../src/hooks/use-hf-token-validation.ts | 16 +++++------
studio/frontend/src/i18n/locales/ar.ts | 2 +-
studio/frontend/src/i18n/locales/de.ts | 2 +-
studio/frontend/src/i18n/locales/en.ts | 2 +-
studio/frontend/src/i18n/locales/es.ts | 2 +-
studio/frontend/src/i18n/locales/fr.ts | 2 +-
studio/frontend/src/i18n/locales/hi.ts | 2 +-
studio/frontend/src/i18n/locales/ja.ts | 2 +-
studio/frontend/src/i18n/locales/ko.ts | 2 +-
studio/frontend/src/i18n/locales/pt-br.ts | 2 +-
studio/frontend/src/i18n/locales/ru.ts | 2 +-
studio/frontend/src/i18n/locales/zh-CN.ts | 2 +-
.../test_hf_token_validation_tick_contract.py | 28 +++++++++++++++++++
14 files changed, 55 insertions(+), 26 deletions(-)
create mode 100644 tests/studio/test_hf_token_validation_tick_contract.py
diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx
index 5037a9b971..17697f933c 100644
--- a/studio/frontend/src/features/settings/tabs/general-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx
@@ -225,12 +225,13 @@ export function GeneralTab() {
setHfToken("");
};
- // Show an "accepted" tick once a non-empty token has been committed to the
- // store and the field still matches it (i.e. not mid-edit). Gives the user
- // feedback that a pasted token was saved.
- const tokenSaved =
+ // Only show the success tick for the currently displayed token after the
+ // authenticated validation endpoint has confirmed it. A saved token alone
+ // may still be malformed, expired, or revoked.
+ const tokenIsCurrent =
draftToken.trim().length > 0 && draftToken.trim() === (hfToken ?? "");
const tokenValidation = useHfTokenValidation(hfToken ?? "");
+ const tokenValidated = tokenIsCurrent && tokenValidation.isValid === true;
useEffect(() => {
let cancelled = false;
@@ -522,16 +523,16 @@ export function GeneralTab() {
onBlur={commitToken}
className={cn(
"h-8 w-full font-mono text-xs",
- tokenSaved ? "pr-14" : "pr-8",
+ tokenValidated ? "pr-14" : "pr-8",
)}
/>
- {tokenSaved ? (
+ {tokenValidated ? (
// Decorative: pointer-events-none lets clicks reach the input
// underneath so the field still focuses anywhere.
diff --git a/studio/frontend/src/hooks/use-hf-token-validation.ts b/studio/frontend/src/hooks/use-hf-token-validation.ts
index e0ed02cca2..64e8a869b4 100644
--- a/studio/frontend/src/hooks/use-hf-token-validation.ts
+++ b/studio/frontend/src/hooks/use-hf-token-validation.ts
@@ -36,10 +36,8 @@ const COMPLETE_HF_TOKEN = /^hf_[A-Za-z0-9]{34}$/;
* requests while typing. isValid is null until checked.
*/
export function useHfTokenValidation(token: string): HfTokenValidationState {
- const debouncedToken = useDebouncedValue(
- token.trim().replace(/^["']+|["']+$/g, ""),
- 500,
- );
+ const normalizedToken = token.trim().replace(/^["']+|["']+$/g, "");
+ const debouncedToken = useDebouncedValue(normalizedToken, 500);
const [completed, setCompleted] = useState(
NO_COMPLETED_VALIDATION,
);
@@ -83,7 +81,8 @@ export function useHfTokenValidation(token: string): HfTokenValidationState {
setCompleted({
token: debouncedToken,
isValid: null,
- error: "Could not verify the token. Check your connection and try again.",
+ error:
+ "Could not verify the token. Check your connection and try again.",
isChecking: false,
});
}
@@ -93,15 +92,16 @@ export function useHfTokenValidation(token: string): HfTokenValidationState {
setCompleted({
token: debouncedToken,
isValid: null,
- error: "Could not verify the token. Check your connection and try again.",
+ error:
+ "Could not verify the token. Check your connection and try again.",
isChecking: false,
});
},
);
}, [debouncedToken, shouldValidate]);
- if (!shouldValidate) return INITIAL;
- if (completed.token !== debouncedToken) {
+ if (!COMPLETE_HF_TOKEN.test(normalizedToken)) return INITIAL;
+ if (completed.token !== normalizedToken) {
return { isValid: null, error: null, isChecking: true };
}
return {
diff --git a/studio/frontend/src/i18n/locales/ar.ts b/studio/frontend/src/i18n/locales/ar.ts
index 4bd4328cce..76ae63daf6 100644
--- a/studio/frontend/src/i18n/locales/ar.ts
+++ b/studio/frontend/src/i18n/locales/ar.ts
@@ -111,7 +111,7 @@ export const ar = {
"يُستخدم لتحميل النماذج المقيّدة ورفع المخرجات.",
hideToken: "إخفاء التوكن",
showToken: "إظهار التوكن",
- tokenSaved: "تم حفظ التوكن",
+ tokenValidated: "تم التحقق من الرمز",
password: "كلمة المرور",
passwordDescription: "تغيير كلمة المرور لحساب Unsloth هذا.",
passwordDialog: {
diff --git a/studio/frontend/src/i18n/locales/de.ts b/studio/frontend/src/i18n/locales/de.ts
index c28d07790f..f446d876b6 100644
--- a/studio/frontend/src/i18n/locales/de.ts
+++ b/studio/frontend/src/i18n/locales/de.ts
@@ -112,7 +112,7 @@ export const de = {
"Wird verwendet, um gated Modelle zu laden und Artefakte zu pushen.",
hideToken: "Token verbergen",
showToken: "Token anzeigen",
- tokenSaved: "Token gespeichert",
+ tokenValidated: "Token validiert",
password: "Passwort",
passwordDescription:
"Ändern Sie das Passwort für dieses Unsloth-Konto.",
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index 4f859aa91c..2db9d21740 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -182,7 +182,7 @@ export const en = {
"Used to load gated models and push artifacts.",
hideToken: "Hide token",
showToken: "Show token",
- tokenSaved: "Token saved",
+ tokenValidated: "Token validated",
password: "Password",
passwordDescription: "Change the password for this Unsloth account.",
passwordDialog: {
diff --git a/studio/frontend/src/i18n/locales/es.ts b/studio/frontend/src/i18n/locales/es.ts
index b7dfee10b8..784265e4b5 100644
--- a/studio/frontend/src/i18n/locales/es.ts
+++ b/studio/frontend/src/i18n/locales/es.ts
@@ -112,7 +112,7 @@ export const es = {
"Se usa para cargar modelos restringidos y subir artefactos.",
hideToken: "Ocultar token",
showToken: "Mostrar token",
- tokenSaved: "Token guardado",
+ tokenValidated: "Token validado",
password: "Contraseña",
passwordDescription:
"Cambia la contraseña de esta cuenta de Unsloth.",
diff --git a/studio/frontend/src/i18n/locales/fr.ts b/studio/frontend/src/i18n/locales/fr.ts
index 6105cd8ccf..284a57c747 100644
--- a/studio/frontend/src/i18n/locales/fr.ts
+++ b/studio/frontend/src/i18n/locales/fr.ts
@@ -112,7 +112,7 @@ export const fr = {
"Utilisé pour charger des modèles restreints et publier des artefacts.",
hideToken: "Masquer le token",
showToken: "Afficher le token",
- tokenSaved: "Token enregistré",
+ tokenValidated: "Jeton validé",
password: "Mot de passe",
passwordDescription:
"Changez le mot de passe de ce compte Unsloth.",
diff --git a/studio/frontend/src/i18n/locales/hi.ts b/studio/frontend/src/i18n/locales/hi.ts
index 732ae1d7fa..8124108eef 100644
--- a/studio/frontend/src/i18n/locales/hi.ts
+++ b/studio/frontend/src/i18n/locales/hi.ts
@@ -111,7 +111,7 @@ export const hi = {
"गेटेड मॉडल लोड करने और आर्टिफैक्ट पुश करने के लिए उपयोग किया जाता है।",
hideToken: "token छिपाएं",
showToken: "token दिखाएं",
- tokenSaved: "Token सहेजा गया",
+ tokenValidated: "Token सत्यापित",
password: "पासवर्ड",
passwordDescription: "इस Unsloth खाते के लिए पासवर्ड बदलें।",
passwordDialog: {
diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts
index de5e93c672..b24d230c4b 100644
--- a/studio/frontend/src/i18n/locales/ja.ts
+++ b/studio/frontend/src/i18n/locales/ja.ts
@@ -114,7 +114,7 @@ export const ja = {
huggingFaceTokenDescription: "ゲート付きモデルの読み込みや、アーティファクトのプッシュに使用されます。",
hideToken: "トークンを非表示",
showToken: "トークンを表示",
- tokenSaved: "トークンを保存しました",
+ tokenValidated: "トークンは検証済みです",
password: "パスワード",
passwordDescription: "この Unsloth アカウントのパスワードを変更します。",
passwordDialog: {
diff --git a/studio/frontend/src/i18n/locales/ko.ts b/studio/frontend/src/i18n/locales/ko.ts
index 6ff9cdbb5e..e7e7687f37 100644
--- a/studio/frontend/src/i18n/locales/ko.ts
+++ b/studio/frontend/src/i18n/locales/ko.ts
@@ -111,7 +111,7 @@ export const ko = {
"게이트된 모델을 불러오고 아티팩트를 푸시하는 데 사용됩니다.",
hideToken: "토큰 숨기기",
showToken: "토큰 표시",
- tokenSaved: "토큰이 저장되었습니다",
+ tokenValidated: "토큰이 확인되었습니다",
password: "비밀번호",
passwordDescription: "이 Unsloth 계정의 비밀번호를 변경합니다.",
passwordDialog: {
diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts
index e6d2347c10..df6c92da6b 100644
--- a/studio/frontend/src/i18n/locales/pt-br.ts
+++ b/studio/frontend/src/i18n/locales/pt-br.ts
@@ -112,7 +112,7 @@ export const ptBR = {
huggingFaceToken: "Token do Hugging Face",
huggingFaceTokenDescription:
"Usado para carregar modelos restritos e enviar artefatos.",
- tokenSaved: "Token salvo",
+ tokenValidated: "Token validado",
hideToken: "Ocultar token",
showToken: "Mostrar token",
password: "Senha",
diff --git a/studio/frontend/src/i18n/locales/ru.ts b/studio/frontend/src/i18n/locales/ru.ts
index 60e939bb0d..f936a397be 100644
--- a/studio/frontend/src/i18n/locales/ru.ts
+++ b/studio/frontend/src/i18n/locales/ru.ts
@@ -111,7 +111,7 @@ export const ru = {
"Используется для загрузки закрытых моделей и публикации артефактов.",
hideToken: "Скрыть токен",
showToken: "Показать токен",
- tokenSaved: "Токен сохранён",
+ tokenValidated: "Токен проверен",
password: "Пароль",
passwordDescription: "Изменить пароль для этого аккаунта Unsloth.",
passwordDialog: {
diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts
index 9984fc324f..0cec860795 100644
--- a/studio/frontend/src/i18n/locales/zh-CN.ts
+++ b/studio/frontend/src/i18n/locales/zh-CN.ts
@@ -113,7 +113,7 @@ export const zhCN = {
huggingFaceTokenDescription: "用于加载受限模型和推送产物。",
hideToken: "隐藏 token",
showToken: "显示 token",
- tokenSaved: "Token 已保存",
+ tokenValidated: "Token 已验证",
password: "密码",
passwordDescription: "更改此 Unsloth 账号的密码。",
passwordDialog: {
diff --git a/tests/studio/test_hf_token_validation_tick_contract.py b/tests/studio/test_hf_token_validation_tick_contract.py
new file mode 100644
index 0000000000..a381101801
--- /dev/null
+++ b/tests/studio/test_hf_token_validation_tick_contract.py
@@ -0,0 +1,28 @@
+"""Contracts for the Hugging Face token validation indicator."""
+
+from pathlib import Path
+
+
+REPO = Path(__file__).resolve().parents[2]
+GENERAL_TAB = REPO / "studio/frontend/src/features/settings/tabs/general-tab.tsx"
+VALIDATION_HOOK = REPO / "studio/frontend/src/hooks/use-hf-token-validation.ts"
+
+
+def test_success_tick_requires_the_current_token_to_be_validated():
+ source = GENERAL_TAB.read_text(encoding = "utf-8")
+
+ assert "tokenIsCurrent && tokenValidation.isValid === true" in source
+ assert 'tokenValidated ? "pr-14" : "pr-8"' in source
+ assert "{tokenValidated ? (" in source
+ assert 'aria-label={t("settings.general.tokenValidated")}' in source
+ assert 'aria-label={t("settings.general.tokenSaved")}' not in source
+
+
+def test_validation_result_must_belong_to_the_current_normalized_token():
+ source = VALIDATION_HOOK.read_text(encoding = "utf-8")
+
+ assert "const normalizedToken = token.trim()" in source
+ assert "useDebouncedValue(normalizedToken, 500)" in source
+ assert "if (!COMPLETE_HF_TOKEN.test(normalizedToken)) return INITIAL" in source
+ assert "if (completed.token !== normalizedToken)" in source
+ assert "if (completed.token !== debouncedToken)" not in source
From 35f887d7956bfc0d0fb35dba9c8ce0623d7d2243 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Tue, 21 Jul 2026 03:54:25 -0700
Subject: [PATCH 053/255] Installer: enable ROCm torch on RDNA2 (gfx1030-1036)
on Windows (#7277)
* Installer: enable ROCm torch on RDNA2 (gfx1030-1036) on Windows
repo.amd.com publishes a gfx103X-all wheel family with win_amd64 torch
2.9.1/2.10.0/2.11.0+rocm7.13.0 (cp310-313), but both Windows allowlists
omitted RDNA2, so RX 6000 cards (gfx1030/1032, etc.) fell back to CPU-only
torch. Map gfx1030-1036 to gfx103X-all in install.ps1 ($archFamilyMap) and
install_python_stack.py (_GFX_TO_AMD_INDEX_ARCH). No torch floor (mirrors
gfx110X-all: newest wheel, no _grouped_mm bug on RDNA2). NVIDIA/Mac/CPU and
Linux paths untouched; gfx906 stays CPU (no wheels published).
* Sync studio/setup.ps1 RDNA2 (gfx1030-1036) allowlists for PR #7277
---
install.ps1 | 4 ++++
studio/install_python_stack.py | 7 +++++++
studio/setup.ps1 | 7 ++++++-
3 files changed, 17 insertions(+), 1 deletion(-)
diff --git a/install.ps1 b/install.ps1
index 524fd85774..cceefd2647 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2128,6 +2128,10 @@ exit 0
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
+ "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
+ "gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all"
+ "gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all"
+ "gfx1030" = "gfx103X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 9921b83543..77e59e98eb 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -352,6 +352,13 @@ _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
"gfx1102": "gfx110X-all", # RDNA 3
"gfx1101": "gfx110X-all",
"gfx1100": "gfx110X-all",
+ "gfx1036": "gfx103X-all",
+ "gfx1035": "gfx103X-all", # RDNA 2 (RX 6000)
+ "gfx1034": "gfx103X-all",
+ "gfx1033": "gfx103X-all",
+ "gfx1032": "gfx103X-all",
+ "gfx1031": "gfx103X-all",
+ "gfx1030": "gfx103X-all",
"gfx90a": "gfx90a",
"gfx908": "gfx908", # MI200/MI100
}
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index f523b9ff14..600dac70b9 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -2774,6 +2774,7 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
"gfx1201", "gfx1200", # RDNA 4
"gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point)
"gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3
+ "gfx1036", "gfx1035", "gfx1034", "gfx1033", "gfx1032", "gfx1031", "gfx1030", # RDNA 2 (RX 6000)
"gfx90a", "gfx908" # MI200 / MI100
)
if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) {
@@ -3064,6 +3065,10 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
+ "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
+ "gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all"
+ "gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all"
+ "gfx1030" = "gfx103X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X and Strix have a null _grouped_mm kernel on torch <2.11.0.
@@ -3098,7 +3103,7 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
# GPU arch detected but not in the supported wheel map — warn explicitly
# so the user knows why they are getting CPU PyTorch instead of ROCm.
substep "[WARN] AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow"
- substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx90a, gfx908" "Yellow"
+ substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx1030-1036 (RDNA 2), gfx90a, gfx908" "Yellow"
} else {
# HIP SDK present ($HasROCm=true via amd-smi) but gcnArchName was not
# readable — warn rather than silently falling back to CPU PyTorch.
From 3c8e3de76ea85609d35a79c1ef6d8f3054f58d62 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Wed, 22 Jul 2026 01:32:29 +0530
Subject: [PATCH 054/255] show the mobile sidebar trigger above the chat header
(#7267)
* Studio: show the mobile sidebar trigger above the chat header
* fix
* correct z-index
---------
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
studio/frontend/src/components/navbar.tsx | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/studio/frontend/src/components/navbar.tsx b/studio/frontend/src/components/navbar.tsx
index 44387f2480..1165705888 100644
--- a/studio/frontend/src/components/navbar.tsx
+++ b/studio/frontend/src/components/navbar.tsx
@@ -22,7 +22,7 @@ export function Navbar() {
);
}
return (
-
@@ -1229,6 +1590,96 @@ function ProjectLanding({
)}
+
{
+ if (!open) setConfirmingDelete(null);
+ }}
+ >
+
+
+ Delete chat
+
+ This permanently deletes "{confirmingDelete?.title}". This cannot
+ be undone.
+
+
+
+ Cancel
+ {
+ const target = confirmingDelete;
+ setConfirmingDelete(null);
+ if (target) void runDelete(target);
+ }}
+ >
+ Delete
+
+
+
+
+
{
+ if (!open) setRenamingProject(false);
+ }}
+ >
+
+
+ Rename project
+
+ setProjectNameDraft(e.target.value)}
+ onKeyDown={(e) => {
+ if (e.key === "Enter") {
+ e.preventDefault();
+ void commitProjectRename();
+ }
+ }}
+ autoFocus={true}
+ maxLength={120}
+ placeholder="Project name"
+ aria-label="Project name"
+ className="focus-visible:border-input focus-visible:ring-0"
+ />
+
+ setRenamingProject(false)}>
+ Cancel
+
+ void commitProjectRename()}
+ disabled={
+ !projectNameDraft.trim() || projectNameDraft.trim() === projectName
+ }
+ >
+ Save
+
+
+
+
+
{
+ if (!open) setDeletingProject(false);
+ }}
+ >
+
+
+ Delete project
+
+ Delete "{projectName}"? Its chats will be permanently deleted.
+
+
+
+ Cancel
+ void commitProjectDelete()}>
+ Delete
+
+
+
+
);
}
diff --git a/studio/frontend/src/features/chat/components/project-switcher.tsx b/studio/frontend/src/features/chat/components/project-switcher.tsx
index 13360ccf7e..8f923a8c80 100644
--- a/studio/frontend/src/features/chat/components/project-switcher.tsx
+++ b/studio/frontend/src/features/chat/components/project-switcher.tsx
@@ -75,8 +75,11 @@ export function ProjectSwitcher({
side="bottom"
align="start"
sideOffset={0}
- className="unsloth-plus-menu ring-0 min-w-56 max-w-72 max-h-72 font-heading"
+ className="unsloth-plus-menu ring-0 min-w-56 max-w-72 font-heading"
>
+ {/* Scroll the list here, not the container, so the rounded corners on
+ the scrollbar side are not squared off. */}
+
{showLoadingRow ? (
Loading…
@@ -117,6 +120,7 @@ export function ProjectSwitcher({
View all projects
+
);
diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts
index c5b24bbb55..335421e145 100644
--- a/studio/frontend/src/features/chat/index.ts
+++ b/studio/frontend/src/features/chat/index.ts
@@ -17,6 +17,7 @@ export {
listRecommendedFolders,
listScanFolders,
loadModel,
+ notifyChatHistoryUpdated,
removeScanFolder,
revealCachedModel,
type BrowseFoldersResponse,
@@ -53,6 +54,7 @@ export {
export { PermissionModeDropdown } from "./permission-mode-select";
export { useChatSearchStore } from "./stores/chat-search-store";
export { usePinnedChatsStore } from "./stores/pinned-chats-store";
+export { usePinnedProjectsStore } from "./stores/pinned-projects-store";
export { useChatPreferencesStore } from "./stores/chat-preferences-store";
export {
PLUS_MENU_ORDER,
diff --git a/studio/frontend/src/features/chat/projects-page.tsx b/studio/frontend/src/features/chat/projects-page.tsx
index ee94a9537f..c9960ffaca 100644
--- a/studio/frontend/src/features/chat/projects-page.tsx
+++ b/studio/frontend/src/features/chat/projects-page.tsx
@@ -39,6 +39,7 @@ import {
renameChatProject,
useChatProjects,
useChatRuntimeStore,
+ usePinnedProjectsStore,
type ProjectRecord,
} from "@/features/chat";
import {
@@ -47,13 +48,15 @@ import {
Edit03Icon,
Folder02Icon,
FolderAddIcon,
+ PinIcon,
+ PinOffIcon,
Search01Icon,
Upload01Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { MoreHorizontalIcon } from "lucide-react";
import { useNavigate } from "@tanstack/react-router";
-import { useMemo, useRef, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
import {
exportProjectConversations,
exportBulkConversationsMerged,
@@ -68,21 +71,38 @@ import {
type SortMode = "activity" | "name";
-function formatUpdatedAgo(ts: number): string {
- const diff = Date.now() - ts;
- if (!Number.isFinite(diff) || diff < 0) return "just now";
- const s = Math.floor(diff / 1000);
- if (s < 60) return "just now";
- const m = Math.floor(s / 60);
- if (m < 60) return `${m} minute${m === 1 ? "" : "s"} ago`;
- const h = Math.floor(m / 60);
- if (h < 24) return `${h} hour${h === 1 ? "" : "s"} ago`;
- const d = Math.floor(h / 24);
- if (d < 30) return `${d} day${d === 1 ? "" : "s"} ago`;
- const mo = Math.floor(d / 30);
- if (mo < 12) return `${mo} month${mo === 1 ? "" : "s"} ago`;
- const y = Math.floor(mo / 12);
- return `${y} year${y === 1 ? "" : "s"} ago`;
+// Reveal this many more projects each time the user scrolls near the bottom.
+const PROJECTS_PAGE_STEP = 12;
+// Visible count before the fit-to-height measurement runs.
+const PROJECTS_INITIAL_FALLBACK = 8;
+// Approx list row height in px, used to estimate how many rows fit the page.
+const PROJECTS_ROW_HEIGHT = 68;
+
+// Modified column, matching a file-list feel: Today / Yesterday / N days ago,
+// then a short date once it is over a week old.
+function formatModified(ts: number): string {
+ if (!Number.isFinite(ts)) return "";
+ const now = new Date();
+ const then = new Date(ts);
+ const startOfToday = new Date(
+ now.getFullYear(),
+ now.getMonth(),
+ now.getDate(),
+ ).getTime();
+ const startOfThen = new Date(
+ then.getFullYear(),
+ then.getMonth(),
+ then.getDate(),
+ ).getTime();
+ const dayDiff = Math.round((startOfToday - startOfThen) / 86_400_000);
+ if (dayDiff <= 0) return "Today";
+ if (dayDiff === 1) return "Yesterday";
+ if (dayDiff < 7) return `${dayDiff} days ago`;
+ return then.toLocaleDateString(undefined, {
+ month: "short",
+ day: "numeric",
+ year: then.getFullYear() === now.getFullYear() ? undefined : "numeric",
+ });
}
export function ProjectsPage() {
@@ -91,6 +111,17 @@ export function ProjectsPage() {
const [query, setQuery] = useState("");
const [sortMode, setSortMode] = useState
("activity");
+ // Rows that fit the page height (measured), plus any revealed via Show more.
+ const [baseFit, setBaseFit] = useState(PROJECTS_INITIAL_FALLBACK);
+ const [extraCount, setExtraCount] = useState(0);
+ const listRef = useRef(null);
+ const sentinelRef = useRef(null);
+ const pinnedProjectIds = usePinnedProjectsStore((s) => s.pinnedIds);
+ const togglePinProject = usePinnedProjectsStore((s) => s.togglePin);
+ const pinnedProjectIdSet = useMemo(
+ () => new Set(pinnedProjectIds),
+ [pinnedProjectIds],
+ );
const [creating, setCreating] = useState(false);
const [nameDraft, setNameDraft] = useState("");
@@ -162,7 +193,7 @@ export function ProjectsPage() {
await handleImport(file, target);
}
- const visibleProjects = useMemo(() => {
+ const sortedProjects = useMemo(() => {
const trimmed = query.trim().toLowerCase();
const filtered = trimmed
? projects.filter((p) => p.name.toLowerCase().includes(trimmed))
@@ -174,6 +205,51 @@ export function ProjectsPage() {
);
return filtered;
}, [projects, query, sortMode]);
+ // Default view shows as many rows as fit the page, then loads more as the
+ // user scrolls near the bottom. Search always spans every project.
+ const isSearching = query.trim() !== "";
+ const visibleCount = baseFit + extraCount;
+ const visibleProjects = isSearching
+ ? sortedProjects
+ : sortedProjects.slice(0, visibleCount);
+ const hasMore = !isSearching && sortedProjects.length > visibleCount;
+
+ // Estimate how many rows fit below the list's top so the first page fills the
+ // screen without loading everything up front.
+ useEffect(() => {
+ function measure() {
+ const el = listRef.current;
+ if (!el) return;
+ const top = el.getBoundingClientRect().top;
+ const reserve = 24; // bottom breathing room
+ const fits = Math.floor(
+ (window.innerHeight - top - reserve) / PROJECTS_ROW_HEIGHT,
+ );
+ setBaseFit(Math.max(PROJECTS_PAGE_STEP, fits));
+ }
+ measure();
+ window.addEventListener("resize", measure);
+ return () => window.removeEventListener("resize", measure);
+ }, [hasLoaded]);
+
+ // Infinite scroll: reveal another page-step whenever the sentinel near the
+ // list bottom scrolls into view.
+ useEffect(() => {
+ const el = sentinelRef.current;
+ if (!el || !hasMore) return;
+ const io = new IntersectionObserver(
+ (entries) => {
+ if (entries[0]?.isIntersecting) {
+ setExtraCount((n) => n + PROJECTS_PAGE_STEP);
+ }
+ },
+ { rootMargin: "300px" },
+ );
+ io.observe(el);
+ return () => io.disconnect();
+ // Re-observe after each load so it keeps filling while the sentinel stays
+ // in view (IntersectionObserver does not re-fire on a steady intersection).
+ }, [hasMore, visibleCount]);
function openProject(projectId: string) {
const runtime = useChatRuntimeStore.getState();
@@ -274,7 +350,7 @@ export function ProjectsPage() {
}
return (
-
+
{/* Global import file input */}
{!hasLoaded ? (
-
+
+
+ Name
+ Modified
+
+
{Array.from({ length: 6 }).map((_, index) => (
-
-
-
-
+
+
+
+
+
))}
@@ -440,9 +522,20 @@ export function ProjectsPage() {
)}
) : (
-
- {visibleProjects.map((project) => (
-
+ <>
+
+ {/* Column header. Name starts at the folder icon's left edge; the
+ right-anchored columns keep Modified over its values. */}
+
+ Name
+ Modified
+
+
+
+ {visibleProjects.map((project) => {
+ const pinned = pinnedProjectIdSet.has(project.id);
+ return (
+
-
-
-
-
+
+
+
+
+ {project.name}
+
+
+ {formatModified(project.updatedAt)}
+
+
+ {/* Pin fades out and the kebab fades in on hover, focus, or
+ menu open. Absolute + opacity gating keeps them from
+ overlapping while leaving the button keyboard-focusable. */}
+ {pinned && (
+
+
+
+ )}
e.stopPropagation()}
aria-label="Project options"
- className="-mr-1 -mt-1 inline-flex size-7 shrink-0 items-center justify-center rounded-full text-muted-foreground opacity-0 transition-opacity hover:bg-black/5 hover:text-foreground dark:hover:bg-white/10 focus-visible:opacity-100 group-hover/project-card:opacity-100 data-[state=open]:bg-black/5 data-[state=open]:opacity-100 dark:data-[state=open]:bg-white/10"
+ className="absolute right-0 flex size-7 shrink-0 items-center justify-center rounded-full text-muted-foreground opacity-0 transition hover:bg-black/5 hover:text-foreground focus-visible:opacity-100 group-hover/project-row:opacity-100 data-[state=open]:bg-black/5 data-[state=open]:opacity-100 dark:hover:bg-white/10 dark:data-[state=open]:bg-white/10"
>
@@ -498,6 +605,16 @@ export function ProjectsPage() {
onKeyDown={(e) => e.stopPropagation()}
className="app-user-menu menu-soft-surface menu-flat-destructive ring-0 w-44 py-2 font-heading rounded-[14px] border-0"
>
+ togglePinProject(project.id)}
+ >
+
+ {pinned ? "Unpin project" : "Pin project"}
+
{
setRenameDraft(project.name);
@@ -546,21 +663,15 @@ export function ProjectsPage() {
-
- {project.name}
-
- {project.instructions ? (
-
- {project.instructions}
-
- ) : null}
-
- Updated {formatUpdatedAgo(project.updatedAt)}
-
- ))}
+ );
+ })}
+ {/* Loads the next page-step when scrolled into view. */}
+ {hasMore &&
}
+
+ >
)}
{/* Create project */}
@@ -684,8 +795,8 @@ export function ProjectsPage() {
Delete project
- Are you sure you want to delete {deleting?.name} ? Chats in this
- project will be moved back to Recents.
+ Are you sure you want to delete {deleting?.name} ? Its chats will
+ be permanently deleted.
setDeleting(null)}>
diff --git a/studio/frontend/src/features/chat/stores/pinned-projects-store.ts b/studio/frontend/src/features/chat/stores/pinned-projects-store.ts
new file mode 100644
index 0000000000..c6dbad5c79
--- /dev/null
+++ b/studio/frontend/src/features/chat/stores/pinned-projects-store.ts
@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { create } from "zustand";
+import { persist } from "zustand/middleware";
+
+// Client-side pin state for projects, keyed by project id. Kept in
+// localStorage. Pinned projects drive the sidebar "Projects" section; new pins
+// are prepended so the most recently pinned project sorts first.
+export interface PinnedProjectsState {
+ pinnedIds: string[];
+ togglePin: (id: string) => void;
+ unpin: (id: string) => void;
+}
+
+export const usePinnedProjectsStore = create()(
+ persist(
+ (set) => ({
+ pinnedIds: [],
+ togglePin: (id) =>
+ set((state) => ({
+ pinnedIds: state.pinnedIds.includes(id)
+ ? state.pinnedIds.filter((x) => x !== id)
+ : [id, ...state.pinnedIds],
+ })),
+ unpin: (id) =>
+ set((state) => ({
+ pinnedIds: state.pinnedIds.filter((x) => x !== id),
+ })),
+ }),
+ {
+ name: "unsloth_pinned_projects",
+ merge: (persisted, current) => {
+ const saved = persisted as Partial | undefined;
+ return {
+ ...current,
+ pinnedIds: Array.isArray(saved?.pinnedIds) ? saved.pinnedIds : [],
+ };
+ },
+ },
+ ),
+);
From c1947ed9465e2b6959fac19e07da95b9e7dda361 Mon Sep 17 00:00:00 2001
From: Ayushman <139611211+InfoSage05@users.noreply.github.com>
Date: Wed, 22 Jul 2026 06:30:22 +0530
Subject: [PATCH 057/255] fix(rocm): prepend system ROCm libs on native Linux
to avoid bundled HIP crash (#7233)
* fix(rocm): prepend system ROCm libs on native Linux to avoid bundled HIP crash
Prebuilt llama.cpp bundles ship their own ROCR/HIP runtime which can be
incompatible with the host's amdkfd kernel driver, causing hsa_init()
to crash or report zero devices. The llama-server then silently falls
back to CPU while the UI reports GPU.
The existing workaround (_wsl_system_rocm_lib_dirs) that prepends
/opt/rocm/lib to LD_LIBRARY_PATH was gated on WSL (/dev/dxg) only,
leaving native Linux AMD hosts unprotected.
This commit adds _native_linux_system_rocm_lib_dirs(), a parallel
helper gated on:
- Linux platform (not WSL)
- /dev/kfd present (bare-metal AMD compute)
- Bundle contains bundled HIP libs (libggml-hip.so)
- System has libhsa-runtime64.so(.1)
It is called from both _llama_server_env_for_binary (serve-time)
and binary_env (install-time validation), directly after the WSL
block in both paths.
Fixes #7208
Fixes #7208
* Add UNSLOTH_LLAMA_NO_SYSTEM_ROCM opt-out to native-Linux system ROCm preference for PR #7233
Lets a host where the bundled runtime works but system ROCm is mismatched keep
the bundle. Mirrored in llama_cpp.py and install_llama_prebuilt.py.
* Prefer env-configured ROCm root over /opt/rocm fallback for PR #7233
Put HIP_PATH/HIP_PATH_57/ROCM_PATH-derived roots before /opt/rocm so a stale
/opt/rocm can't shadow the driver-matching install the env vars point at.
Mirrored in llama_cpp.py and install_llama_prebuilt.py.
* Match versioned libggml-hip.so via glob so the native-Linux ROCm fix fires for PR #7233
* Clarify native-Linux ROCm prepend uses the consistent system stack for PR #7233
* llama_cpp: tighten native-Linux ROCm prepend comments (no code change)
---------
Co-authored-by: Daniel Han
---
studio/backend/core/inference/llama_cpp.py | 56 ++++++++++++++++++++++
studio/install_llama_prebuilt.py | 49 +++++++++++++++++++
2 files changed, 105 insertions(+)
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 2c7433f7a4..8651ed9ea8 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -247,6 +247,59 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]":
return out
+def _bundled_hip_present(binary_dir: str) -> bool:
+ """True when a prebuilt bundle ships its own HIP backend library."""
+ if not binary_dir:
+ return False
+ try:
+ # Glob the version suffix (libggml-hip.so, .so.0, .so.0.11.1) the same
+ # way the installer's runtime health check matches libggml-hip.so*.
+ return any(Path(str(binary_dir)).glob("libggml-hip.so*"))
+ except OSError:
+ return False
+
+
+def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]":
+ """System ROCm lib dir(s) to prepend before a prebuilt's bundled HIP, on native Linux.
+
+ The bundled bare-metal HIP runtime can mismatch the host amdkfd driver and crash
+ in hsa_init(); prepending the whole system ROCm lib dir loads a driver-matched,
+ version-consistent stack (libhsa-runtime64 / libamdhip64 / librocblas) ahead of it.
+ The whole dir is deliberate: mixing the bundle's rocBLAS with a different-version
+ system HIP/ROCR risks missing symbols. UNSLOTH_LLAMA_NO_SYSTEM_ROCM=1 keeps the pure
+ bundle (for a host whose system ROCm lacks this arch); no-op on WSL / non-Linux.
+ """
+ if os.environ.get("UNSLOTH_LLAMA_NO_SYSTEM_ROCM") == "1":
+ return []
+ if sys.platform != "linux" or os.path.exists("/dev/dxg"):
+ return []
+ if not os.path.exists("/dev/kfd"):
+ return []
+ if not _bundled_hip_present(binary_dir):
+ return []
+ # Env-configured ROCm root first; /opt/rocm only as a fallback so a stale
+ # /opt/rocm doesn't shadow the driver-matching install these vars point at.
+ candidates = []
+ for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
+ val = os.environ.get(var)
+ if val:
+ candidates.append(val)
+ candidates.append("/opt/rocm")
+ out: "list[str]" = []
+ seen: "set[str]" = set()
+ for base in candidates:
+ for lib_sub in ("lib", "lib64"):
+ d = os.path.join(base, lib_sub)
+ if d in seen:
+ continue
+ seen.add(d)
+ if os.path.exists(os.path.join(d, "libhsa-runtime64.so")) or os.path.exists(
+ os.path.join(d, "libhsa-runtime64.so.1")
+ ):
+ out.append(d)
+ return out
+
+
# Plan-without-action re-prompt state now lives in tool_call_parser (imported above).
# Default max_tokens to the effective context when known. The floor is high
@@ -3592,6 +3645,9 @@ class LlamaCppBackend:
lib_dirs.extend(_wsl_system_rocm_lib_dirs())
if lib_dirs:
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
+ # Native Linux AMD: system ROCm libs before the bundle's HIP runtime,
+ # which can be incompatible with the host amdkfd driver.
+ lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir))
lib_dirs.append(binary_dir)
_arch = platform.machine() # x86_64, aarch64, etc.
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index b8182a534b..1093ae2fd0 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -5717,6 +5717,50 @@ def _wsl_system_rocm_lib_dirs() -> list[str]:
return out
+def _bundled_hip_present(binary_dir: str) -> bool:
+ if not binary_dir:
+ return False
+ try:
+ # Glob the version suffix (libggml-hip.so, .so.0, .so.0.11.1) the same
+ # way the installer's runtime health check matches libggml-hip.so*.
+ return any(Path(str(binary_dir)).glob("libggml-hip.so*"))
+ except OSError:
+ return False
+
+
+def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> list[str]:
+ # UNSLOTH_LLAMA_NO_SYSTEM_ROCM=1 keeps the bundled runtime (opt-out).
+ if os.environ.get("UNSLOTH_LLAMA_NO_SYSTEM_ROCM") == "1":
+ return []
+ if sys.platform != "linux" or os.path.exists("/dev/dxg"):
+ return []
+ if not os.path.exists("/dev/kfd"):
+ return []
+ if not _bundled_hip_present(binary_dir):
+ return []
+ # Env-configured ROCm root first; /opt/rocm only as a fallback so a stale
+ # /opt/rocm doesn't shadow the driver-matching install these vars point at.
+ candidates = []
+ for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"):
+ val = os.environ.get(var)
+ if val:
+ candidates.append(val)
+ candidates.append("/opt/rocm")
+ out: list[str] = []
+ seen: set[str] = set()
+ for base in candidates:
+ for lib_sub in ("lib", "lib64"):
+ d = os.path.join(base, lib_sub)
+ if d in seen:
+ continue
+ seen.add(d)
+ if os.path.exists(os.path.join(d, "libhsa-runtime64.so")) or os.path.exists(
+ os.path.join(d, "libhsa-runtime64.so.1")
+ ):
+ out.append(d)
+ 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(
@@ -5877,6 +5921,11 @@ def binary_env(
if _wsl_rocm:
ld_dirs = [*_wsl_rocm, *ld_dirs]
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
+ # Native Linux AMD: system ROCm libs before the bundle's bundled HIP
+ # runtime, which can be incompatible with the host amdkfd driver.
+ _native_rocm = _native_linux_system_rocm_lib_dirs(str(binary_path.parent))
+ if _native_rocm:
+ ld_dirs = [*_native_rocm, *ld_dirs]
existing = [part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part]
env["LD_LIBRARY_PATH"] = os.pathsep.join(dedupe_existing_dirs([*ld_dirs, *existing]))
elif host.is_macos:
From 207a9f00bf68b3384a38af9de4bfc8ef244e70c2 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Tue, 21 Jul 2026 18:01:07 -0700
Subject: [PATCH 058/255] studio: extend the _grouped_mm null-kernel guard to
Linux ROCm RDNA4 (gfx1201) (#7292)
* studio: extend the _grouped_mm null-kernel guard to Linux ROCm RDNA4
torch._grouped_mm has a null HIP kernel on RDNA4 (gfx1200/gfx1201) at
ROCm <= 7.12 (fixed in 7.13; ROCm/TheRock #5284). The existing guard that
registers a Python mm/bmm fallback was win32-only, so Linux gfx1201 (e.g.
R9700 Pro on Ubuntu) hits the null kernel -> illegal instruction during
training.
Extract the fallback registration into a module-level helper
(_install_grouped_mm_cpu_fallback) and add a Linux branch that installs it,
gated on gfx1200/gfx1201 AND HIP < 7.13 so NVIDIA/CUDA and every non-RDNA4
AMD arch are untouched, and it is a no-op on fixed runtimes. The Windows
path now calls the same helper with identical behavior.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Resolve HIP version from torch.__version__ when version.hip is unset for PR #7292
AMD SDK / Radeon ROCm wheels leave torch.version.hip empty and encode the
version only in torch.__version__ (e.g. +rocm7.12). The Linux gfx120X guard
parsed version.hip only, so those affected installs skipped the fallback and
still hit the null _grouped_mm kernel. Mirror the Windows parse: version.hip,
then the embedded rocmX.Y, then assume affected unless a post-fix rocmsdk wheel.
* Scan all GPUs and add RDNA4 name fallback for _grouped_mm guard in PR #7292
* worker.py: tighten gfx120X Linux guard comments (no code change)
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
studio/backend/core/training/worker.py | 187 +++++++++++++++----------
1 file changed, 113 insertions(+), 74 deletions(-)
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 2df4fa58c6..1aeeab4cbc 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -90,6 +90,79 @@ _FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS"
# run_training_process() and isn't GC'd mid-run.
_WINDOWS_ROCM_GROUPED_MM_LIB = None
+
+def _install_grouped_mm_cpu_fallback(torch_mod, logger, label):
+ """Register a Python mm/bmm fallback for torch._grouped_mm and return the Library.
+
+ RDNA4 (gfx1200/gfx1201) ships a null HIP _grouped_mm kernel on ROCm <= 7.12
+ (fixed in 7.13; ROCm/TheRock #5284). JitDecomp dispatches _grouped_mm to the
+ null kernel and crashes; overriding the CUDA dispatch key bypasses it. Shared
+ by the Windows and Linux ROCm guards. Keep the returned Library referenced so
+ the registration outlives the caller.
+ """
+ import warnings as _warnings
+
+ _gm_lib = torch_mod.library.Library("aten", "IMPL")
+
+ def _grouped_mm_safe_impl(
+ self,
+ mat2,
+ offs = None,
+ bias = None,
+ out_dtype = None,
+ ):
+ """Python mm/bmm fallback for _grouped_mm on gfx120X (null HIP kernel, ROCm <= 7.12)."""
+ _t = torch_mod
+ if offs is None:
+ # No offsets: 2-D -> mm, 3-D batched -> bmm (unconditional mm broke 3-D MoE).
+ if self.dim() == 3 and mat2.dim() == 3:
+ result = _t.bmm(self.contiguous(), mat2.contiguous())
+ elif self.dim() == 3 and mat2.dim() == 2:
+ result = _t.matmul(self.contiguous(), mat2.contiguous())
+ elif self.dim() == 2 and mat2.dim() == 3:
+ result = _t.matmul(self.contiguous(), mat2.contiguous())
+ else:
+ result = _t.mm(self.contiguous(), mat2.contiguous())
+ else:
+ # Grouped: offs[i] is the exclusive end-row of group i.
+ offs_list = offs.tolist()
+ pieces = []
+ prev = 0
+ for idx, end in enumerate(offs_list):
+ end = int(end)
+ a_part = self[prev:end].contiguous()
+ b_part = mat2[idx].contiguous() if mat2.dim() == 3 else mat2.contiguous()
+ pieces.append(_t.mm(a_part, b_part))
+ prev = end
+ # Include trailing rows not covered by offs.
+ if prev < self.shape[0]:
+ a_tail = self[prev:].contiguous()
+ b_tail = mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous()
+ pieces.append(_t.mm(a_tail, b_tail))
+ result = (
+ _t.cat(pieces, dim = 0)
+ if pieces
+ else _t.zeros(0, mat2.shape[-1], device = self.device, dtype = self.dtype)
+ )
+ if bias is not None:
+ result = result + bias
+ if out_dtype is not None:
+ result = result.to(out_dtype)
+ elif result.dtype != self.dtype:
+ result = result.to(self.dtype)
+ return result
+
+ with _warnings.catch_warnings():
+ _warnings.simplefilter("ignore")
+ _gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
+ logger.info(
+ "%s: patched _grouped_mm CUDA dispatch (null HIP kernel on gfx120X, "
+ "ROCm <= 7.12 -- bypassed with Python mm fallback)",
+ label,
+ )
+ return _gm_lib
+
+
# Subprocesses don't inherit os.add_dll_directory registrations. Replicate
# main.py's Windows ROCm DLL setup so the first `import torch` finds
# amdhip64.dll. Handles retained at module scope so they aren't GC'd.
@@ -2689,80 +2762,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
# so 7.13+ uses the real GPU kernel.
if not _hip_ver_at_least(7, 13):
try:
- import warnings as _warnings
-
- _gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
-
- def _grouped_mm_safe_impl(
- self,
- mat2,
- offs = None,
- bias = None,
- out_dtype = None,
- ):
- """Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
- _t = _torch_for_rocm
- if offs is None:
- # No offsets: 2-D -> mm, 3-D batched -> bmm
- # (unconditional mm broke 3-D MoE).
- if self.dim() == 3 and mat2.dim() == 3:
- result = _t.bmm(self.contiguous(), mat2.contiguous())
- elif self.dim() == 3 and mat2.dim() == 2:
- # Broadcast 2-D mat2 across the batch dim.
- result = _t.matmul(self.contiguous(), mat2.contiguous())
- elif self.dim() == 2 and mat2.dim() == 3:
- # Broadcast 2-D self across batch via matmul.
- result = _t.matmul(self.contiguous(), mat2.contiguous())
- else:
- result = _t.mm(self.contiguous(), mat2.contiguous())
- else:
- # Grouped: offs[i] is the exclusive end-row of group i.
- offs_list = offs.tolist()
- pieces = []
- prev = 0
- for idx, end in enumerate(offs_list):
- end = int(end)
- a_part = self[prev:end].contiguous()
- if mat2.dim() == 3:
- b_part = mat2[idx].contiguous()
- else:
- b_part = mat2.contiguous()
- pieces.append(_t.mm(a_part, b_part))
- prev = end
- # Include trailing rows not covered by offs.
- if prev < self.shape[0]:
- a_tail = self[prev:].contiguous()
- b_tail = (
- mat2[-1].contiguous() if mat2.dim() == 3 else mat2.contiguous()
- )
- pieces.append(_t.mm(a_tail, b_tail))
- result = (
- _t.cat(pieces, dim = 0)
- if pieces
- else _t.zeros(
- 0,
- mat2.shape[-1],
- device = self.device,
- dtype = self.dtype,
- )
- )
- if bias is not None:
- result = result + bias
- if out_dtype is not None:
- result = result.to(out_dtype)
- elif result.dtype != self.dtype:
- result = result.to(self.dtype)
- return result
-
- with _warnings.catch_warnings():
- _warnings.simplefilter("ignore")
- _gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
-
- _WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC
- logger.info(
- "Windows ROCm: patched _grouped_mm CUDA dispatch "
- "(null HIP kernel on gfx1200, ROCm ≤ 7.12 — "
- "bypassed with Python mm fallback)"
+ _WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback(
+ _torch_for_rocm, logger, "Windows ROCm"
)
except Exception as _patch_exc:
logger.warning(
@@ -2776,6 +2777,44 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
"skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)"
)
+ # ── 1f-linux. Linux ROCm RDNA4 _grouped_mm null kernel ──
+ # The win32 guard above misses Linux: RDNA4 (gfx1200/gfx1201) hits the same null
+ # HIP _grouped_mm kernel at ROCm <= 7.12 (fixed 7.13, ROCm/TheRock #5284). Gate on
+ # arch + HIP < 7.13 so NVIDIA/CUDA and non-RDNA4 AMD are untouched; no-op if fixed.
+ if sys.platform.startswith("linux") and _hw.IS_ROCM:
+ try:
+ _torch_lin = sys.modules.get("torch")
+ if _torch_lin is not None and _torch_lin.cuda.is_available():
+ # Prefer torch.version.hip, else rocmX.Y from torch.__version__ (AMD
+ # SDK / Radeon wheels leave version.hip unset). Unknown version on a
+ # gfx120X build -> assume affected unless it is a post-fix rocmsdk wheel.
+ _hip_str = str(getattr(getattr(_torch_lin, "version", None), "hip", "") or "")
+ _ver = getattr(_torch_lin, "__version__", "").lower()
+ _m = re.match(r"(\d+)\.(\d+)", _hip_str) or re.search(r"rocm(\d+)\.(\d+)", _ver)
+ if _m:
+ _hip_lt_713 = (int(_m.group(1)), int(_m.group(2))) < (7, 13)
+ else:
+ _hip_lt_713 = "rocmsdk" not in _ver
+ # Scan every visible GPU (device_map="balanced" can place layers on a
+ # later RDNA4 card, so device 0 is not enough). Match gfx120X by arch,
+ # or by RX 9000 / R9700 name when the wheel omits gcnArchName.
+ _rdna4 = False
+ for _i in range(_torch_lin.cuda.device_count()):
+ _props = _torch_lin.cuda.get_device_properties(_i)
+ _lin_arch, _ = _rocm_classify_unified_memory(_props)
+ _lin_name = (getattr(_props, "name", "") or "").lower()
+ if _lin_arch.lower() in ("gfx1200", "gfx1201") or (
+ not _lin_arch and re.search(r"rx\s*90[0-9]0|r9700", _lin_name)
+ ):
+ _rdna4 = True
+ break
+ if _rdna4 and _hip_lt_713:
+ _WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback(
+ _torch_lin, logger, "Linux ROCm gfx120X"
+ )
+ except Exception as _gm_lin_exc:
+ logger.warning("Linux ROCm gfx120X: could not patch _grouped_mm: %s", _gm_lin_exc)
+
# ── 1g. ROCm OOM guard ──
# On ROCm, exhausting VRAM can hang the HIP driver instead of raising.
# set_per_process_memory_fraction caps the allocator so PyTorch raises
From 2c492c8d9baebb21e4be522025c381d9befafc6e Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Tue, 21 Jul 2026 18:07:41 -0700
Subject: [PATCH 059/255] Recognize Radeon 8065S (Gorgon Halo / Ryzen AI Max
400) as gfx1151 (#7290)
* Recognize Radeon 8065S (Gorgon Halo / Ryzen AI Max 400) as gfx1151
* Classify Radeon 8065S (Gorgon Halo) as unified memory in ROCm OOM guard
* [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>
---
install.ps1 | 2 +-
install.sh | 6 +++---
studio/backend/core/training/worker.py | 11 ++++++++---
studio/backend/tests/test_rocm_oom_guard.py | 3 +++
studio/install_python_stack.py | 4 ++--
studio/setup.ps1 | 2 +-
studio/setup.sh | 2 +-
7 files changed, 19 insertions(+), 11 deletions(-)
diff --git a/install.ps1 b/install.ps1
index cceefd2647..9c99379666 100644
--- a/install.ps1
+++ b/install.ps1
@@ -1844,7 +1844,7 @@ exit 0
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080)
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060)
- @{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
+ @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
@{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
diff --git a/install.sh b/install.sh
index 445bab616a..adec5f2d6e 100755
--- a/install.sh
+++ b/install.sh
@@ -1636,7 +1636,7 @@ _maybe_reroute_strixhalo_to_2404() {
# CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps.
if _has_usable_nvidia_gpu; then return 0; fi
# Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes.
- if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
+ if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
@@ -2649,7 +2649,7 @@ _maybe_bootstrap_rocm_wsl() {
[ -e /dev/dxg ] || return 0
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
- if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
+ if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
@@ -2960,7 +2960,7 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
case "$_gpu_disp_mkt" in
*"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4
- *"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
+ *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 1aeeab4cbc..5ded18ea45 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -775,8 +775,9 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
3. Device-name substring match (last resort when all arch attrs absent;
AMD SDK / Radeon wheels may not populate them):
- gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M``
- - gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395),
- ``Radeon 8050S`` (cut-down SKU)
+ - gfx1151 Strix Halo / Gorgon Halo: ``Radeon 8065S`` (Ryzen AI
+ Max+ 495), ``Radeon 8060S`` (Ryzen AI MAX+
+ 395), ``Radeon 8050S`` (cut-down SKU)
"""
gcn_arch = ""
for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"):
@@ -801,7 +802,11 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
# Arch attrs absent — fall back to device-name matching.
dev_lower = (getattr(props, "name", "") or "").lower()
is_unified = (
- "890m" in dev_lower or "880m" in dev_lower or "8060s" in dev_lower or "8050s" in dev_lower
+ "890m" in dev_lower
+ or "880m" in dev_lower
+ or "8065s" in dev_lower
+ or "8060s" in dev_lower
+ or "8050s" in dev_lower
)
return gcn_arch, is_unified
diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py
index 699d0b74f5..ad46f6ee41 100644
--- a/studio/backend/tests/test_rocm_oom_guard.py
+++ b/studio/backend/tests/test_rocm_oom_guard.py
@@ -163,6 +163,9 @@ class TestDeviceNameFallback:
"AMD Radeon 8060S",
"Radeon 8050S Graphics", # cut-down Strix Halo SKU
"AMD Radeon 8050S",
+ # gfx1151 Gorgon Halo (Ryzen AI Max 400 refresh)
+ "Radeon 8065S Graphics", # Ryzen AI Max+ 495
+ "AMD Radeon 8065S",
# case variants
"RADEON 8060S GRAPHICS",
"radeon 8050s",
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 77e59e98eb..b58e94cd3f 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -719,8 +719,8 @@ def _detect_windows_gfx_arch() -> str | None:
_WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [
(r"9070 XT|9080", "gfx1201"), # RDNA 4 (Radeon RX 9070 XT / 9080)
(r"9070|9060", "gfx1200"), # RDNA 4 (Radeon RX 9070 / 9060)
- # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
- (r"8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max", "gfx1151"),
+ # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
+ (r"8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max", "gfx1151"),
# RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
(
r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]"
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 600dac70b9..7bf6627578 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -1498,7 +1498,7 @@ if (-not $HasNvidiaSmi) {
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (Radeon RX 9070 XT / 9080)
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (Radeon RX 9070 / 9060)
- @{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
+ @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
@{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop / workstation (Navi 31)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
diff --git a/studio/setup.sh b/studio/setup.sh
index df7178c662..2a2b41d0f6 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -1136,7 +1136,7 @@ elif [ "$_setup_amd_detected" = true ]; then
case "$_setup_mkt" in
*"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4
- *"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
+ *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
*"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
*"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _setup_gfx="gfx1102" ;; # RDNA 3 (Navi 33)
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31)
From 5308c24e70ef7f800e4c5c38016e6b69da42f049 Mon Sep 17 00:00:00 2001
From: Solaris-star <67425364+Solaris-star@users.noreply.github.com>
Date: Wed, 22 Jul 2026 09:12:39 +0800
Subject: [PATCH 060/255] fix(install.ps1): use ordinal IndexOf when stripping
index URL credentials (#7286)
* fix(install.ps1): use ordinal IndexOf when stripping index URL credentials
On non-English Windows locales, culture-aware String.IndexOf can
mis-locate punctuation-only markers like ://, which corrupts scheme and
authority parsing and crashes Remove-IndexUrlCredentials with a Substring
ArgumentOutOfRangeException (issue 7279).
Force Ordinal comparison for URL scheme/host parsing.
Fixes #7279
* Condense the ordinal parsing comment in Remove-IndexUrlCredentials
---------
Co-authored-by: Daniel Han
---
install.ps1 | 8 +++++---
1 file changed, 5 insertions(+), 3 deletions(-)
diff --git a/install.ps1 b/install.ps1
index 9c99379666..36e03ca51d 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2030,16 +2030,18 @@ exit 0
# _strip_index_url_credentials (install.sh / py / setup.ps1).
function Remove-IndexUrlCredentials {
param([string]$Url)
- $sep = $Url.IndexOf('://')
+ # Ordinal, not culture-aware: on non-English locales (e.g. th-TH) linguistic
+ # IndexOf treats "://" as ignorable, mis-locates it, and crashes Substring (issue #7279).
+ $sep = $Url.IndexOf('://', [System.StringComparison]::Ordinal)
if ($sep -lt 0) { return $Url }
$scheme = $Url.Substring(0, $sep)
$rest = $Url.Substring($sep + 3)
# Drop query / fragment (may hold auth tokens).
$q = $rest.IndexOfAny([char[]]('?', '#'))
if ($q -ge 0) { $rest = $rest.Substring(0, $q) }
- $slash = $rest.IndexOf('/')
+ $slash = $rest.IndexOf('/', [System.StringComparison]::Ordinal)
$authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest }
- $at = $authority.LastIndexOf('@')
+ $at = $authority.LastIndexOf('@', [System.StringComparison]::Ordinal)
$host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority }
if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" }
return "${scheme}://${host_}"
From 4e1cb4affa7cee5897fdc11ed922aae7099a8334 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Tue, 21 Jul 2026 19:40:59 -0700
Subject: [PATCH 061/255] install.sh: route Strix to the AMD arch index on
rocm7.2 (#7264) + PCI detection hint (#7293)
* install.sh: route Strix (gfx1151/gfx1150) to the AMD arch index on rocm7.2, add PCI hint
Two Linux install fixes for AMD Strix Halo / Strix Point:
1. #7264: Strix reverts to rocm7.2. Modern ROCm (7.3+) caps to the generic
rocm7.2 index and the Radeon repo can be unavailable, so gfx1151/gfx1150
landed on a non-arch-specific build (torch 2.11+rocm7.2) instead of
repo.amd.com/rocm/whl/gfx (torch 2.11+rocm7.13, AMD's real Strix
fixes). The reroute to that arch index only fired on rocm7.1; broaden it to
rocm7.2 too. Only acts when a gfx1151/gfx1150 is actually detected, so other
arches on rocm7.2 pass through unchanged.
2. Rows about Strix not detected -> CPU-only: when no GPU is detected but an AMD
display GPU is on the PCI bus, print a targeted hint (ROCm kernel stack /
/dev/kfd missing) instead of only the generic docs pointer. Purely
additive diagnostic; does not change the torch index decision.
* Address review on PR #7293: gate PCI hint on ROCm-detection failure, fix test marker, use read builtin
- Only show the 'ROCm cannot see the GPU' hint when _has_amd_rocm_gpu fails;
a detected-but-too-old ROCm (rocminfo works, wheels need 6.0+) has its own path.
- Update test_previous_torch_pin.sh to the stable 'Strix Halo / Strix Point:'
marker after the heading reworded (the old grep broke the ordering assert).
- _amd_gpu_present_via_pci: read builtin instead of spawning cat twice per
device, and guard /sys/bus/pci/devices existence.
* install.sh: reroute Strix on any generic index older than the arch build
Generalize the Strix reroute from the hardcoded rocm7.1/rocm7.2 match to a
version compare against the arch index's own build (rocm7.13):
- backwards: rocm6.0-6.4 and rocm7.0 now reroute (were silently missed)
- forwards: any future intermediate rocm7.x below 7.13 reroutes; rocm7.13+
is left alone so a generic index that already carries the fix is not
downgraded to the arch build
_rocm_index_below does an integer major.minor compare (so rocm7.2 < rocm7.13);
non-rocm, arch (gfx), and unparseable URLs return false, so NVIDIA/CPU and the
arch index itself are untouched. Reroute still fires only for gfx1150/gfx1151.
* install.sh: tighten _amd_gpu_present_via_pci comment (no code change)
* install.sh: match the index leaf in the Strix version reroute (#7293 review)
Address two review points on the rocm-version reroute:
- Parse the final path segment (_torch_index_leaf) instead of grepping the whole
URL. A custom mirror whose base path holds its own rocm token (e.g.
.../rocm7.13/cache/rocm7.2) previously matched the base and skipped the reroute;
now it compares the leaf (rocm7.2) like the nearby index-family logic. Renamed
the helper to _rocm_leaf_below and switched the case selector to $_torch_index_leaf.
- Replace the stale test_strix_override_only_fires_on_rocm71 (which passed by
matching the new rocm7.13 comment) with an executed test that runs _rocm_leaf_below
and asserts rocm6.0-7.12 reroute while rocm7.13+/gfx/cu leaves do not.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* install.sh: keep gfx probes non-fatal under set -e (#7293 review)
The Strix reroute now matches every rocm* index, not just rocm7.1, so its gfx
detection runs on all AMD installs. Each `_gfx_all=$(rocminfo|amd-smi | grep -oE
gfx...)` returns 1 when grep finds no match, which under set -euo pipefail aborts
the installer before the next fallback runs (e.g. rocminfo present but emitting no
gfx token). Append `|| true` to the three probes, matching the display block that
already guards this. Add an executed regression test (shimmed rocminfo/amd-smi)
that fails if any probe becomes fatal again.
* [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>
---
install.sh | 75 ++++++++++++++++++-----
tests/sh/test_previous_torch_pin.sh | 2 +-
tests/studio/install/test_rocm_support.py | 66 ++++++++++++++++++--
3 files changed, 119 insertions(+), 24 deletions(-)
diff --git a/install.sh b/install.sh
index adec5f2d6e..e0f57c198b 100755
--- a/install.sh
+++ b/install.sh
@@ -2127,6 +2127,23 @@ _has_amd_rocm_gpu() {
return 1
}
+# Returns 0 if an AMD display GPU is on the PCI bus even when ROCm can't use it
+# (e.g. a Strix Halo iGPU with no /dev/kfd). Only sharpens the "no GPU detected"
+# hint. vendor 0x1002 = AMD/ATI; class 0x03* = display controller.
+_amd_gpu_present_via_pci() {
+ [ -d /sys/bus/pci/devices ] || return 1
+ for _pci_vendor in /sys/bus/pci/devices/*/vendor; do
+ [ -r "$_pci_vendor" ] || continue
+ read -r _v < "$_pci_vendor" 2>/dev/null || continue
+ [ "$_v" = "0x1002" ] || continue
+ _cls="${_pci_vendor%vendor}class"
+ [ -r "$_cls" ] || continue
+ read -r _c < "$_cls" 2>/dev/null || continue
+ case "$_c" in 0x03*) return 0 ;; esac
+ done
+ return 1
+}
+
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@@ -2818,29 +2835,45 @@ case "$TORCH_INDEX_URL" in
fi
;;
esac
-# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ───────
-# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug
-# that causes a segfault in torch._grouped_mm (moe_utils.py line 167).
-# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when
-# _amd_gpu_radeon=true the installer silently lands on the broken combo.
-# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2.
-case "$TORCH_INDEX_URL" in
- */rocm7.1|*/rocm7.1.*)
+# 0 when a rocmX.Y index leaf ($1, the final path segment) is older than floor
+# $2.$3 (int compare, so rocm7.2 < rocm7.13). Non-rocm leaves (gfx*, cu*, cpu) and
+# non-numeric versions return 1. Leaf-based (like $_torch_index_leaf) so a mirror
+# base holding its own rocm token compares the family leaf, not the base path.
+_rocm_leaf_below() {
+ case "$1" in rocm[0-9]*.[0-9]*) : ;; *) return 1 ;; esac
+ _rb=${1#rocm}; _maj=${_rb%%.*}; _min=${_rb#*.}; _min=${_min%%.*}
+ case "$_maj$_min" in *[!0-9]*) return 1 ;; esac
+ if [ "$_maj" -lt "$2" ]; then return 0; fi
+ if [ "$_maj" -eq "$2" ] && [ "$_min" -lt "$3" ]; then return 0; fi
+ return 1
+}
+# ── Strix Halo / Strix Point: route to the AMD arch-specific index ───────────
+# gfx1151/gfx1150 need torch 2.11+rocm7.13 from repo.amd.com/rocm/whl/gfx/,
+# which carries AMD's real fixes (the rocm7.1 _grouped_mm segfault, moe_utils.py:167,
+# and later Strix kernel bugs). Every generic pytorch.org index below rocm7.13 lacks
+# them (and the Radeon repo can be offline, unslothai#7264), so reroute a detected
+# Strix GPU whenever the picked index is older than the arch build -- covers today's
+# rocm6.0-7.2 and any future 7.x < 7.13; rocm7.13+ already has the fixes, so leave it.
+case "$_torch_index_leaf" in
+ rocm[0-9]*)
# Collect every gfx token in rocminfo / amd-smi enumeration order
# (skip duplicates), then index by HIP_VISIBLE_DEVICES /
# ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box
# where the user selected the dGPU does NOT get rerouted to the
# Strix per-gfx index.
+ # || true on each probe: no gfx match makes grep exit 1, which under
+ # set -euo pipefail would abort the installer before the next fallback
+ # runs (now that the case matches every rocm* index, not just rocm7.1).
_gfx_all=""
if command -v rocminfo >/dev/null 2>&1; then
- _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
+ _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
- _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
+ _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
# PowerShell paths also probe `amd-smi static --asic`; mirror it
# so a host with hipinfo-less amd-smi reports the gfx target.
if [ -z "$_gfx_all" ]; then
- _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
+ _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
fi
_runtime_gfx=""
@@ -2865,13 +2898,14 @@ case "$TORCH_INDEX_URL" in
case "$_runtime_gfx" in
gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
esac
- if [ -n "$_strix_gfx" ]; then
+ # Skip rocm7.13+ generic indexes: they already ship the fixes, so the
+ # arch build (rocm7.13) would be a downgrade rather than a rescue.
+ if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then
echo "" >&2
- echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2
- echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2
- echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2
- echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2
- echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
+ echo " [WARN] $_strix_gfx (Strix) detected -- routing to the AMD arch-specific index" >&2
+ echo " [WARN] torch 2.11+rocm7.13 has AMD's real gfx1150/gfx1151 fixes (the ROCm 7.1" >&2
+ echo " [WARN] _grouped_mm segfault, moe_utils.py:167, and later Strix kernel bugs)," >&2
+ echo " [WARN] and is more reliable than the rocm7.2 index or an offline Radeon repo." >&2
echo "" >&2
# AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's
# actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred
@@ -3031,6 +3065,13 @@ case "$TORCH_INDEX_URL" in
substep " driver is current; or run unsloth/scripts/install_rocm_wsl_strixhalo.sh yourself."
else
substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
+ # Only when ROCm truly can't see the GPU: a detected-but-too-old
+ # ROCm (rocminfo works, wheels need 6.0+) has its own guidance.
+ if ! _has_amd_rocm_gpu && _amd_gpu_present_via_pci; then
+ substep "An AMD GPU is on the PCI bus but ROCm cannot see it (no /dev/kfd," "$C_WARN"
+ substep " rocminfo, or amd-smi). Install the ROCm kernel stack so /dev/kfd exists;"
+ substep " Strix Halo (gfx1151/gfx1150) needs a recent kernel (6.11+) and ROCm 7.x."
+ fi
fi
substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):"
substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
diff --git a/tests/sh/test_previous_torch_pin.sh b/tests/sh/test_previous_torch_pin.sh
index 1bc0d1f27f..f910affc2f 100644
--- a/tests/sh/test_previous_torch_pin.sh
+++ b/tests/sh/test_previous_torch_pin.sh
@@ -98,7 +98,7 @@ assert_eq "probe before venv replacement" "yes" "$([ -n "$_probe_line" ] && [ -n
# The pin must be evaluated AFTER the last index/constraint decision (the Strix
# reroute raises the floor), so a raised floor rejects an older kept release.
_pin_line=$(grep -n '_prev_pin=\$(_previous_torch_pin' "$INSTALL_SH" | head -1 | cut -d: -f1)
-_strix_line=$(grep -n 'Strix Halo / Strix Point: force rocm7.2 wheels' "$INSTALL_SH" | head -1 | cut -d: -f1)
+_strix_line=$(grep -n 'Strix Halo / Strix Point:' "$INSTALL_SH" | head -1 | cut -d: -f1)
assert_eq "pin evaluated after the Strix reroute" "yes" "$([ -n "$_pin_line" ] && [ -n "$_strix_line" ] && [ "$_pin_line" -gt "$_strix_line" ] && echo yes)"
# A kept release that vanished from the index must fall back to the supported range.
assert_eq "resolve-failure fallback wired" "yes" "$(grep -q 'TORCH_CONSTRAINT="\$_PREV_FALLBACK_CONSTRAINT"' "$INSTALL_SH" && echo yes)"
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index f7dcd79d10..5825bbe31f 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -3,8 +3,11 @@
import importlib.util
import json
import os
+import re
+import shutil
import subprocess
import sys
+import tempfile
from pathlib import Path
from unittest.mock import MagicMock, mock_open, patch, PropertyMock
@@ -3191,13 +3194,64 @@ class TestStrixRocm71Override:
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
assert "moe_utils" in source or "_grouped_mm" in source
- def test_strix_override_only_fires_on_rocm71(self):
- """install.sh must scope the Strix override to rocm7.1 only (not rocm7.2+)."""
+ def test_strix_override_scoped_below_arch_floor(self):
+ """Strix reroute must fire for rocm leaves BELOW the arch floor (7.13) and
+ NOT at/above it. Executed via _rocm_leaf_below so it verifies the actual
+ version comparison, not a text match that a comment could satisfy."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
- strix_idx = source.find("_strix_gfx")
- assert strix_idx != -1
- context_before = source[max(0, strix_idx - 2400) : strix_idx]
- assert "rocm7.1" in context_before
+ # Selector + gate must switch on the index LEAF, not the whole URL (a mirror
+ # base path with its own rocm token would false-positive otherwise).
+ assert 'case "$_torch_index_leaf" in' in source
+ assert '_rocm_leaf_below "$_torch_index_leaf" 7 13' in source
+ shell = shutil.which("sh") or shutil.which("bash")
+ if not shell:
+ pytest.skip("no POSIX shell to execute _rocm_leaf_below")
+ match = re.search(r"^_rocm_leaf_below\(\) \{.*?^\}", source, re.S | re.M)
+ assert match, "could not extract _rocm_leaf_below from install.sh"
+ fn = match.group(0)
+
+ def below(leaf):
+ return (
+ subprocess.run(
+ [shell, "-c", f'{fn}\n_rocm_leaf_below "$1" 7 13', "_", leaf]
+ ).returncode
+ == 0
+ )
+
+ for leaf in ("rocm6.0", "rocm7.0", "rocm7.1", "rocm7.2", "rocm7.12"):
+ assert below(leaf), f"{leaf} must reroute (below arch floor 7.13)"
+ for leaf in ("rocm7.13", "rocm7.14", "rocm8.0", "gfx1151", "cu128", "cpu"):
+ assert not below(leaf), f"{leaf} must NOT reroute (>= floor or non-rocm)"
+
+ def test_gfx_probe_survives_no_match_under_set_e(self):
+ """A gfx probe whose grep finds no match must not abort install.sh under
+ set -euo pipefail before the amd-smi fallback runs. The reroute case now
+ matches every rocm* index, so this would break ordinary 6.x/7.2 installs
+ with a flaky rocminfo. Executed with shimmed tools, not a text match."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute the probe block")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ block = re.search(
+ r'^ _gfx_all=""\n.*?(?=^ _strix_gfx="")', source, re.S | re.M
+ )
+ assert block, "could not extract the gfx-detection block"
+ with tempfile.TemporaryDirectory() as d:
+ # rocminfo emits no gfx token; amd-smi supplies gfx1151 (the fallback)
+ for name, out in (("rocminfo", "no gpu here"), ("amd-smi", "GPU: gfx1151")):
+ p = os.path.join(d, name)
+ with open(p, "w", encoding = "utf-8") as f:
+ f.write(f'#!/bin/sh\ncat <<"EOT"\n{out}\nEOT\n')
+ os.chmod(p, 0o755)
+ script = (
+ 'set -euo pipefail\nHIP_VISIBLE_DEVICES=""\nROCR_VISIBLE_DEVICES=""\n'
+ + block.group(0)
+ + '\nprintf "OK:%s\\n" "$_gfx_all"\n'
+ )
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""))
+ r = subprocess.run([shell, "-c", script], env = env, capture_output = True, text = True)
+ assert r.returncode == 0, f"probe aborted under set -e: {r.stderr}"
+ assert "OK:gfx1151" in r.stdout, f"amd-smi fallback not reached: {r.stdout!r}"
def test_torch_constraint_updated_for_strix_amd_index(self):
"""install.sh must set TORCH_CONSTRAINT>=2.11 when routing Strix to AMD index."""
From 8517721adb692746cf09d717465cfddd2fe853aa Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Tue, 21 Jul 2026 23:25:56 -0700
Subject: [PATCH 062/255] Studio: move sidebar search into the header (#7304)
* Studio: move sidebar search into the header
Put the search action as an icon button next to the sidebar toggle in
the header instead of a full-width nav row, so New Chat is the only
fixed row above the scrolling list. The search row is kept for the
collapsed icon rail only. Also add a small bottom gap under New Chat
when it is pinned during scroll.
* Studio: keep search row on custom-titlebar platforms
The header search button only renders on mac/web where the brand row
shows. On win/linux custom titlebars there's no header button, so keep
the full-width search row visible instead of hiding it.
* Studio: address review on sidebar search tooltip
- Hide the search tooltip on mobile (hidden={isMobile}), matching the
SidebarMenuButton tooltip convention.
- Show Cmd K on Mac and Ctrl K elsewhere instead of a hardcoded glyph;
the search dialog binds both meta and ctrl. Uses getClientPlatform so
it is correct on web too, not just Tauri.
---
.../frontend/src/components/app-sidebar.tsx | 57 ++++++++++++++++---
.../src/components/tauri/window-titlebar.tsx | 2 +-
2 files changed, 50 insertions(+), 9 deletions(-)
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index a13828b06b..f4226760a2 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -45,6 +45,7 @@ import { Spinner } from "@/components/ui/spinner";
import { Switch } from "@/components/ui/switch";
import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler";
import {
+ getClientPlatform,
shouldUseCustomWindowTitlebar,
shouldUseNativeMacWindowTitlebar,
} from "@/components/tauri/window-titlebar";
@@ -343,6 +344,8 @@ export function AppSidebar() {
);
const [usesCustomTitlebar] = useState(shouldUseCustomWindowTitlebar);
const [usesNativeMacTitlebar] = useState(shouldUseNativeMacWindowTitlebar);
+ // Mac uses Cmd, others use Ctrl. Not Tauri-gated, so it's right on web too.
+ const [isMacPlatform] = useState(() => getClientPlatform().includes("mac"));
const { pathname, search } = useRouterState({
select: (s) => ({
pathname: s.location.pathname,
@@ -1194,27 +1197,55 @@ export function AppSidebar() {
)}
- {!isMobile && (
+
{
+ useChatSearchStore.getState().open();
+ closeMobileIfOpen();
+ }}
className="inline-flex h-[33px] w-[32px] cursor-pointer items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring"
- aria-label={t("shell.aria.closeSidebar")}
+ aria-label={t("shell.navigation.search")}
>
-
+
- {t("shell.aria.closeSidebar")}
+ {t("shell.navigation.search")}
+
+ {isMacPlatform ? "⌘K" : "Ctrl+K"}
+
- )}
+ {!isMobile && (
+
+
+
+
+
+
+
+ {t("shell.aria.closeSidebar")}
+
+
+ )}
+
{!isMobile && (
@@ -1246,8 +1277,10 @@ export function AppSidebar() {
{/* Uniform pl-1.5 pr-2 keeps every hover pill the same width, inset from the edge. */}
@@ -1280,10 +1313,18 @@ export function AppSidebar() {
openNewChat(null);
}}
/>
+ {/* Search sits in the header when the brand row is shown (mac/web).
+ Hide this row there, but keep it in the collapsed rail. On custom
+ titlebars (win/linux) there's no header button, so keep the row. */}
{
useChatSearchStore.getState().open();
closeMobileIfOpen();
diff --git a/studio/frontend/src/components/tauri/window-titlebar.tsx b/studio/frontend/src/components/tauri/window-titlebar.tsx
index 8d11bbd229..d5c74df463 100644
--- a/studio/frontend/src/components/tauri/window-titlebar.tsx
+++ b/studio/frontend/src/components/tauri/window-titlebar.tsx
@@ -40,7 +40,7 @@ type NavigatorWithUserAgentData = Navigator & {
};
};
-function getClientPlatform(): string {
+export function getClientPlatform(): string {
if (typeof navigator === "undefined") {
return "";
}
From 59bda2e1f77a3ff060d26b9cdb0b69c798d8c7a1 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Wed, 22 Jul 2026 15:05:33 +0530
Subject: [PATCH 063/255] Studio: reuse MLX prompt cache across turns instead
of re-prefilling (#7311)
* Studio: reuse MLX prompt cache across turns instead of re-prefilling
* clean up
* key prompt cache on what the KV covers
* skip windowed KV caches past their window
* verify prefix coverage before caching KV
---
.../backend/core/inference/mlx_inference.py | 214 ++++++++-
.../tests/test_mlx_inference_backend.py | 410 ++++++++++++++++++
2 files changed, 611 insertions(+), 13 deletions(-)
diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py
index e78c93b6f3..d19c67a01a 100644
--- a/studio/backend/core/inference/mlx_inference.py
+++ b/studio/backend/core/inference/mlx_inference.py
@@ -181,19 +181,27 @@ def _vlm_messages_have_tool_history(messages):
)
-def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
+def _build_generation_stats(
+ prompt_n,
+ prompt_tps,
+ gen_n,
+ gen_tps,
+ cached_n = 0,
+):
"""Map mlx stream stats onto the usage/timings shape llama-server emits."""
prompt_n = int(prompt_n or 0)
gen_n = int(gen_n or 0)
+ cached_n = int(cached_n or 0)
prompt_tps = float(prompt_tps or 0.0)
gen_tps = float(gen_tps or 0.0)
prompt_ms = (prompt_n / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0
predicted_ms = (gen_n / gen_tps * 1000.0) if gen_tps > 0 else 0.0
+ total_prompt_n = prompt_n + cached_n
return {
"usage": {
- "prompt_tokens": prompt_n,
+ "prompt_tokens": total_prompt_n,
"completion_tokens": gen_n,
- "total_tokens": prompt_n + gen_n,
+ "total_tokens": total_prompt_n + gen_n,
},
"timings": {
"prompt_n": prompt_n,
@@ -204,11 +212,123 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
"predicted_ms": predicted_ms,
"predicted_per_token_ms": (predicted_ms / gen_n) if gen_n > 0 else 0.0,
"predicted_per_second": gen_tps,
- "cache_n": 0,
+ "cache_n": cached_n,
},
}
+PROMPT_CACHE_ENTRIES = 6
+PROMPT_CACHE_MEMORY_FRACTION = 0.15
+PROMPT_CACHE_FALLBACK_BYTES = 2 * 1024**3
+
+
+def _mlx_prompt_cache_api():
+ try:
+ from mlx_lm.models.cache import (
+ LRUPromptCache,
+ can_trim_prompt_cache,
+ make_prompt_cache,
+ trim_prompt_cache,
+ )
+ except ImportError:
+ return None
+ return LRUPromptCache, make_prompt_cache, can_trim_prompt_cache, trim_prompt_cache
+
+
+def _prompt_cache_max_bytes(recommended_gb = None):
+ override = os.environ.get("UNSLOTH_MLX_PROMPT_CACHE_BYTES")
+ if override:
+ try:
+ return max(int(override), 0)
+ except ValueError:
+ logger.warning("Ignoring non-integer UNSLOTH_MLX_PROMPT_CACHE_BYTES=%r", override)
+ if recommended_gb:
+ return int(recommended_gb * 1e9 * PROMPT_CACHE_MEMORY_FRACTION)
+ return PROMPT_CACHE_FALLBACK_BYTES
+
+
+def _flatten_kv_entries(cache):
+ for entry in cache:
+ nested = getattr(entry, "caches", None)
+ if nested is None:
+ yield entry
+ else:
+ yield from _flatten_kv_entries(nested)
+
+
+def _kv_prefix_coverage(cache):
+ covered = None
+ for entry in _flatten_kv_entries(cache):
+ offset = getattr(entry, "offset", None)
+ if offset is None:
+ return None
+ if getattr(entry, "start_position", 0):
+ return None
+ window = getattr(entry, "max_size", None)
+ if window is not None and offset > window:
+ return None
+ if covered is None:
+ covered = offset
+ elif covered != offset:
+ return None
+ return covered
+
+
+class _MLXPromptCacheHistory:
+ def __init__(self, max_entries, max_bytes):
+ api = _mlx_prompt_cache_api()
+ if api is None:
+ raise RuntimeError("mlx-lm is too old for LRUPromptCache")
+ lru_cls, make, can_trim, trim = api
+ self._make_prompt_cache = make
+ self._can_trim = can_trim
+ self._trim = trim
+ self._max_bytes = max_bytes
+ self._lru = lru_cls(max_size = max_entries, max_bytes = max_bytes)
+
+ def fetch(self, model, key, tokens):
+ cache, rest = self._lru.fetch_nearest_cache(key, list(tokens))
+ if cache is not None:
+ if rest:
+ return cache, list(rest)
+ if self._can_trim(cache) and self._trim(cache, 1) == 1:
+ return cache, list(tokens[-1:])
+ if len(tokens) > 1:
+ head = list(tokens[:-1])
+ cache, rest = self._lru.fetch_nearest_cache(key, head)
+ if cache is not None:
+ covered = len(head) - len(rest)
+ return cache, list(tokens[covered:])
+ return self._make_prompt_cache(model), list(tokens)
+
+ def insert(self, key, tokens, cache):
+ # An over-budget entry evicts itself and every other conversation.
+ nbytes = sum(getattr(entry, "nbytes", 0) for entry in cache)
+ if nbytes > self._max_bytes:
+ logger.debug(
+ "MLX prompt cache: skipping %.2f GB entry over the %.2f GB budget",
+ nbytes / 1e9,
+ self._max_bytes / 1e9,
+ )
+ return
+ covered = _kv_prefix_coverage(cache)
+ if covered is None:
+ logger.debug("MLX prompt cache: skipping cache with unverifiable prefix coverage")
+ return
+ tokens = list(tokens)
+ if covered > len(tokens):
+ logger.debug(
+ "MLX prompt cache: cache covers %d tokens but only %d were tracked",
+ covered,
+ len(tokens),
+ )
+ return
+ tokens = tokens[:covered]
+ if not tokens:
+ return
+ self._lru.insert_cache(key, tokens, cache)
+
+
def _mlx_distributed_rank_size(group = None):
"""Return ``(rank, world_size)`` for an optional MLX distributed group."""
if group is None:
@@ -313,6 +433,55 @@ class MLXInferenceBackend:
# Recorded for unload to release pinned memory back to the OS.
self._memory_limits_applied = {}
+ self._prompt_cache_history = None
+ self._prompt_cache_unavailable = False
+
+ def _prompt_cache(self):
+ if self._prompt_cache_history is not None or self._prompt_cache_unavailable:
+ return self._prompt_cache_history
+ max_bytes = _prompt_cache_max_bytes(self._memory_limits_applied.get("recommended_gb"))
+ if max_bytes <= 0:
+ self._prompt_cache_unavailable = True
+ logger.info("MLX prompt cache disabled by budget")
+ return None
+ try:
+ self._prompt_cache_history = _MLXPromptCacheHistory(
+ PROMPT_CACHE_ENTRIES,
+ max_bytes,
+ )
+ except Exception as exc:
+ self._prompt_cache_unavailable = True
+ logger.info("MLX prompt cache unavailable (%s); prefilling every request", exc)
+ return None
+ logger.info(
+ "MLX prompt cache: %d entries, %.2f GB budget",
+ PROMPT_CACHE_ENTRIES,
+ max_bytes / 1e9,
+ )
+ return self._prompt_cache_history
+
+ def _clear_prompt_cache(self):
+ self._prompt_cache_history = None
+ self._prompt_cache_unavailable = False
+
+ def _prepare_prompt_cache(self, prompt, adapter_state):
+ history = self._prompt_cache()
+ if history is None:
+ return prompt, None, None, None, 0
+ try:
+ tokenizer = self._tokenizer
+ bos = getattr(tokenizer, "bos_token", None)
+ add_special_tokens = bos is None or not prompt.startswith(bos)
+ tokens = list(tokenizer.encode(prompt, add_special_tokens = add_special_tokens))
+ if not tokens:
+ return prompt, None, None, None, 0
+ key = f"{self.active_model_name}|{adapter_state!r}"
+ cache, rest = history.fetch(self._model, key, tokens)
+ except Exception as exc:
+ logger.debug("MLX prompt cache lookup failed: %s", exc)
+ return prompt, None, None, None, 0
+ return rest, cache, key, tokens, len(tokens) - len(rest)
+
def _configure_memory_limits(self):
"""Apply Metal memory caps before loading a model.
@@ -535,6 +704,7 @@ class MLXInferenceBackend:
self._distributed_world_size = 1
if self.active_model_name == model_name:
self.active_model_name = None
+ self._clear_prompt_cache()
gc.collect()
mx.clear_cache()
@@ -731,24 +901,34 @@ class MLXInferenceBackend:
# prefix on every native-protocol snapshot just as the normal
# decoding path does below.
normalized_output = think_prefix
- logger.info(
- "Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s",
- len(prompt),
- max_new_tokens,
- type(self._model).__name__,
- type(self._tokenizer).__name__,
- )
with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state):
+ (
+ gen_prompt,
+ prompt_cache,
+ cache_key,
+ prompt_tokens,
+ cached_n,
+ ) = self._prepare_prompt_cache(prompt, _adapter_state)
+ logger.info(
+ "Generating: prompt_len=%d, cached=%d, max_tokens=%d, model=%s, tokenizer=%s",
+ len(prompt),
+ cached_n,
+ max_new_tokens,
+ type(self._model).__name__,
+ type(self._tokenizer).__name__,
+ )
final_response = None
try:
# Enter request-scoped model state before yielding any response.
if think_prefix:
yield think_prefix
gen_kwargs = dict(
- prompt = prompt,
+ prompt = gen_prompt,
max_tokens = max_new_tokens,
sampler = sampler,
)
+ if prompt_cache is not None:
+ gen_kwargs["prompt_cache"] = prompt_cache
if logits_processors is not None:
gen_kwargs["logits_processors"] = logits_processors
for response in stream_generate(
@@ -757,6 +937,7 @@ class MLXInferenceBackend:
**gen_kwargs,
):
final_response = response
+ token_ids.append(response.token)
if preserve_native_channels:
piece = getattr(response, "text", None) or ""
delta = normalizer.feed(piece)
@@ -764,7 +945,6 @@ class MLXInferenceBackend:
normalized_output += delta
yield normalized_output
else:
- token_ids.append(response.token)
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
@@ -773,6 +953,13 @@ class MLXInferenceBackend:
if cancel_event and cancel_event.is_set():
break
+ if prompt_cache is not None and prompt_tokens is not None:
+ history = self._prompt_cache_history
+ if history is not None:
+ try:
+ history.insert(cache_key, prompt_tokens + token_ids, prompt_cache)
+ except Exception as exc:
+ logger.debug("MLX prompt cache insert failed: %s", exc)
except Exception as e:
import traceback
logger.error("stream_generate failed:\n%s", traceback.format_exc())
@@ -785,6 +972,7 @@ class MLXInferenceBackend:
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
+ cached_n,
)
if normalizer is not None:
cancelled = cancel_event is not None and cancel_event.is_set()
diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py
index fafaea0043..d49a2281a0 100644
--- a/studio/backend/tests/test_mlx_inference_backend.py
+++ b/studio/backend/tests/test_mlx_inference_backend.py
@@ -922,3 +922,413 @@ def test_mlx_vlm_normalizes_native_reasoning_channels(monkeypatch):
"vision ",
"vision answer",
]
+
+
+class _FakeLRUPromptCache:
+ def __init__(
+ self,
+ max_size = 10,
+ max_bytes = 1 << 63,
+ ):
+ self.max_size = max_size
+ self.max_bytes = max_bytes
+ self.entries = {}
+
+ def fetch_nearest_cache(self, key, tokens):
+ import copy
+
+ stored = self.entries.get(key, {})
+ exact = stored.get(tuple(tokens))
+ if exact is not None:
+ return copy.deepcopy(exact), []
+ best = None
+ for candidate, cache in stored.items():
+ if len(candidate) < len(tokens) and tuple(tokens[: len(candidate)]) == candidate:
+ if best is None or len(candidate) > len(best[0]):
+ best = (candidate, cache)
+ if best is not None:
+ return copy.deepcopy(best[1]), list(tokens[len(best[0]) :])
+ return None, list(tokens)
+
+ def insert_cache(
+ self,
+ key,
+ tokens,
+ prompt_cache,
+ *,
+ cache_type = "assistant",
+ ):
+ import copy
+ self.entries.setdefault(key, {})[tuple(tokens)] = copy.deepcopy(prompt_cache)
+
+
+class _FakeCacheEntry:
+ def __init__(
+ self,
+ offset = 0,
+ nbytes = 1,
+ ):
+ self.offset = offset
+ self.nbytes = nbytes
+
+
+def _install_fake_prompt_cache_api(monkeypatch, trimmable = True):
+ from core.inference import mlx_inference
+
+ def _make_prompt_cache(_model):
+ return [_FakeCacheEntry()]
+
+ def _can_trim_prompt_cache(_cache):
+ return trimmable
+
+ def _trim_prompt_cache(cache, num):
+ cache[0].offset = max(cache[0].offset - num, 0)
+ return num
+
+ monkeypatch.setattr(
+ mlx_inference,
+ "_mlx_prompt_cache_api",
+ lambda: (
+ _FakeLRUPromptCache,
+ _make_prompt_cache,
+ _can_trim_prompt_cache,
+ _trim_prompt_cache,
+ ),
+ )
+
+
+def test_mlx_prompt_cache_max_bytes_budget(monkeypatch):
+ from core.inference.mlx_inference import (
+ PROMPT_CACHE_FALLBACK_BYTES,
+ PROMPT_CACHE_MEMORY_FRACTION,
+ _prompt_cache_max_bytes,
+ )
+
+ monkeypatch.delenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", raising = False)
+ assert _prompt_cache_max_bytes(None) == PROMPT_CACHE_FALLBACK_BYTES
+ assert _prompt_cache_max_bytes(20.0) == int(20.0 * 1e9 * PROMPT_CACHE_MEMORY_FRACTION)
+
+ monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "4096")
+ assert _prompt_cache_max_bytes(20.0) == 4096
+ monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "0")
+ assert _prompt_cache_max_bytes(20.0) == 0
+ monkeypatch.setenv("UNSLOTH_MLX_PROMPT_CACHE_BYTES", "not-a-number")
+ assert _prompt_cache_max_bytes(20.0) == int(20.0 * 1e9 * PROMPT_CACHE_MEMORY_FRACTION)
+
+
+def test_mlx_prompt_cache_never_returns_empty_remainder(monkeypatch):
+ _install_fake_prompt_cache_api(monkeypatch)
+ from core.inference.mlx_inference import _MLXPromptCacheHistory
+
+ history = _MLXPromptCacheHistory(6, 1 << 30)
+ tokens = list(range(10))
+ cache, rest = history.fetch(object(), "key", tokens)
+ assert len(rest) == 10
+ cache[0].offset = len(tokens)
+ history.insert("key", tokens, cache)
+
+ _cache, rest = history.fetch(object(), "key", tokens)
+ assert rest == tokens[-1:]
+
+ longer = tokens + [99, 100]
+ _cache, rest = history.fetch(object(), "key", longer)
+ assert rest == [99, 100]
+
+ _install_fake_prompt_cache_api(monkeypatch, trimmable = False)
+ history = _MLXPromptCacheHistory(6, 1 << 30)
+ cache, _rest = history.fetch(object(), "key", tokens)
+ cache[0].offset = len(tokens)
+ history.insert("key", tokens, cache)
+ _cache, rest = history.fetch(object(), "key", tokens)
+ assert rest == tokens, "untrimmable entry must not be reused"
+
+
+def test_mlx_prompt_cache_key_isolates_adapter_state(monkeypatch):
+ _install_fake_prompt_cache_api(monkeypatch)
+ _install_fake_mlx(monkeypatch)
+ from core.inference.mlx_inference import MLXInferenceBackend
+
+ class _Tok:
+ bos_token = None
+
+ def encode(
+ self,
+ text,
+ add_special_tokens = True,
+ ):
+ return [ord(c) for c in text]
+
+ backend = MLXInferenceBackend()
+ backend._model = object()
+ backend._tokenizer = _Tok()
+ backend.active_model_name = "model-a"
+
+ prompt = "shared prefix"
+ _rest, cache, key, tokens, cached = backend._prepare_prompt_cache(prompt, True)
+ assert cached == 0
+ cache[0].offset = len(tokens)
+ backend._prompt_cache_history.insert(key, tokens, cache)
+
+ _rest, _cache, _key, _tokens, cached_same = backend._prepare_prompt_cache(prompt, True)
+ assert cached_same > 0
+ _rest, _cache, _key, _tokens, cached_flipped = backend._prepare_prompt_cache(prompt, False)
+ assert cached_flipped == 0
+
+
+def _install_fake_text_stack(
+ monkeypatch,
+ token_map,
+ captured,
+ markers = None,
+):
+ import types as _types
+
+ from core.inference import mlx_inference
+
+ _install_fake_mlx(monkeypatch)
+ monkeypatch.setattr(
+ mlx_inference,
+ "_temporary_mlx_adapter_state",
+ lambda _model, _state: __import__("contextlib").nullcontext(),
+ )
+ monkeypatch.setattr(
+ "core.inference.chat_template_helpers.apply_chat_template_for_generation",
+ lambda _tok, messages, **_kw: messages[-1]["content"],
+ )
+ monkeypatch.setattr(
+ "core.inference.chat_template_helpers.render_with_native_template_fallback",
+ lambda formatted_prompt, **_kw: SimpleNamespace(
+ prompt = formatted_prompt,
+ reasoning_channel_markers = markers,
+ ),
+ )
+ monkeypatch.setattr(
+ "core.inference.chat_template_helpers.detect_think_prefill",
+ lambda *_a, **_kw: "",
+ )
+
+ class _Resp:
+ def __init__(self, token, processed):
+ self.token = token
+ self.text = f"<{token}>"
+ self.prompt_tokens = processed
+ self.prompt_tps = 10.0
+ self.generation_tokens = 1
+ self.generation_tps = 5.0
+
+ def _stream_generate(_model, _tokenizer, **kwargs):
+ captured.append(kwargs)
+ processed = len(kwargs["prompt"])
+ cache = kwargs.get("prompt_cache")
+ if cache is not None:
+ cache[0].offset += processed
+ for token in token_map["generated"]:
+ if cache is not None:
+ cache[0].offset += 1
+ yield _Resp(token, processed)
+
+ mlx_lm_pkg = _types.ModuleType("mlx_lm")
+ mlx_lm_pkg.stream_generate = _stream_generate
+ mlx_lm_sample = _types.ModuleType("mlx_lm.sample_utils")
+ mlx_lm_sample.make_sampler = lambda **_kw: object()
+ mlx_lm_sample.make_logits_processors = lambda **_kw: []
+ monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg)
+ monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample)
+
+ class _Tok:
+ bos_token = None
+ chat_template = "x"
+
+ def encode(
+ self,
+ text,
+ add_special_tokens = True,
+ ):
+ return list(token_map[text])
+
+ def decode(
+ self,
+ ids,
+ skip_special_tokens = False,
+ ):
+ return "".join(str(i) for i in ids)
+
+ from core.inference.mlx_inference import MLXInferenceBackend
+
+ backend = MLXInferenceBackend()
+ backend._model = object()
+ backend._tokenizer = _Tok()
+ backend._is_vlm = False
+ backend.active_model_name = "model-a"
+ return backend
+
+
+def _run_turn(backend, prompt):
+ list(
+ backend.generate_chat_response(
+ messages = [{"role": "user", "content": prompt}],
+ max_new_tokens = 4,
+ )
+ )
+
+
+def test_mlx_text_reuses_prompt_cache_on_the_next_turn(monkeypatch):
+ _install_fake_prompt_cache_api(monkeypatch)
+ captured = []
+ token_map = {
+ "P1": [1, 2, 3],
+ "P2": [1, 2, 3, 7, 8, 9, 10],
+ "generated": [7, 8],
+ }
+ backend = _install_fake_text_stack(monkeypatch, token_map, captured)
+
+ _run_turn(backend, "P1")
+ assert captured[0]["prompt"] == [1, 2, 3]
+ assert "prompt_cache" in captured[0]
+ assert backend.last_generation_stats["timings"]["cache_n"] == 0
+
+ _run_turn(backend, "P2")
+ assert captured[1]["prompt"] == [9, 10], "turn two should prefill only the new tail"
+
+ stats = backend.last_generation_stats
+ assert stats["timings"]["cache_n"] == 5
+ assert stats["timings"]["prompt_n"] == 2
+ assert stats["usage"]["prompt_tokens"] == 7
+
+
+def test_mlx_text_without_lru_prompt_cache_prefills_the_full_prompt(monkeypatch):
+ from core.inference import mlx_inference
+
+ monkeypatch.setattr(mlx_inference, "_mlx_prompt_cache_api", lambda: None)
+ captured = []
+ token_map = {"P1": [1, 2, 3], "generated": [7]}
+ backend = _install_fake_text_stack(monkeypatch, token_map, captured)
+
+ _run_turn(backend, "P1")
+ assert captured[0]["prompt"] == "P1"
+ assert "prompt_cache" not in captured[0]
+ assert backend.last_generation_stats["timings"]["cache_n"] == 0
+
+
+def test_mlx_text_tracks_tokens_on_the_native_reasoning_path(monkeypatch):
+ _install_fake_prompt_cache_api(monkeypatch)
+ captured = []
+ token_map = {"P1": [1, 2, 3], "P2": [1, 2, 3, 7, 8, 9], "generated": [7, 8]}
+ backend = _install_fake_text_stack(monkeypatch, token_map, captured, markers = ("", " "))
+
+ _run_turn(backend, "P1")
+ _run_turn(backend, "P2")
+ assert captured[1]["prompt"] == [9]
+
+
+def test_mlx_presence_penalty_latches_the_first_decode_step():
+ mx = pytest.importorskip("mlx.core")
+ import numpy as np
+
+ from core.inference.mlx_inference import _make_mlx_presence_penalty_processor
+
+ processor = _make_mlx_presence_penalty_processor(2.0)
+ logits = mx.zeros((1, 5))
+ out = processor(mx.array([3]), logits)
+ assert np.array_equal(np.array(out), np.zeros((1, 5))), "prompt must not be penalized"
+ out = processor(mx.array([3, 1]), mx.zeros((1, 5)))
+ penalized = np.array(out)[0]
+ assert penalized[1] == -2.0
+ assert penalized[3] == 0.0
+
+
+def test_mlx_prompt_cache_survives_reset_but_not_unload(monkeypatch):
+ _install_fake_prompt_cache_api(monkeypatch)
+ _install_fake_mlx(monkeypatch)
+ sys.modules["mlx.core"].clear_cache = lambda: None
+ from core.inference.mlx_inference import MLXInferenceBackend
+
+ backend = MLXInferenceBackend()
+ backend.active_model_name = "model-a"
+ history = backend._prompt_cache()
+ assert history is not None
+
+ backend.reset_generation_state()
+ assert backend._prompt_cache_history is history
+
+ backend.unload_model("model-a")
+ assert backend._prompt_cache_history is None
+
+
+def test_mlx_prompt_cache_skips_entries_over_budget(monkeypatch):
+ _install_fake_prompt_cache_api(monkeypatch)
+ from core.inference.mlx_inference import _MLXPromptCacheHistory
+
+ history = _MLXPromptCacheHistory(6, 1000)
+ history.insert("key", [1, 2, 3], [_FakeCacheEntry(offset = 3, nbytes = 400)])
+ assert len(history._lru.entries.get("key", {})) == 1
+
+ history.insert("key", list(range(50)), [_FakeCacheEntry(offset = 50, nbytes = 5000)])
+ stored = history._lru.entries.get("key", {})
+ assert tuple([1, 2, 3]) in stored
+ assert tuple(range(50)) not in stored
+
+
+def test_mlx_prompt_cache_keys_on_what_the_kv_covers(monkeypatch):
+ _install_fake_prompt_cache_api(monkeypatch)
+ from core.inference.mlx_inference import _MLXPromptCacheHistory
+
+ class _Entry:
+ def __init__(
+ self,
+ offset,
+ nbytes = 1,
+ ):
+ self.offset = offset
+ self.nbytes = nbytes
+
+ history = _MLXPromptCacheHistory(6, 1 << 30)
+
+ history.insert("key", list(range(10)), [_Entry(offset = 8)])
+ assert tuple(range(8)) in history._lru.entries["key"]
+ assert tuple(range(10)) not in history._lru.entries["key"]
+
+ history.insert("other", list(range(4)), [_Entry(offset = 9)])
+ assert "other" not in history._lru.entries
+
+
+def test_mlx_prompt_cache_only_stores_verifiable_prefix_coverage(monkeypatch):
+ mx = pytest.importorskip("mlx.core")
+ from mlx_lm.models.cache import CacheList, ChunkedKVCache, KVCache, RotatingKVCache
+
+ _install_fake_prompt_cache_api(monkeypatch)
+ from core.inference.mlx_inference import _kv_prefix_coverage, _MLXPromptCacheHistory
+
+ def feed(entry, n):
+ for _ in range(n):
+ block = mx.zeros((1, 2, 1, 4), dtype = mx.float16)
+ entry.update_and_fetch(block, block)
+ mx.eval(entry.state)
+ return entry
+
+ plain = feed(KVCache(), 30)
+ unwrapped = feed(RotatingKVCache(max_size = 100, keep = 2), 30)
+ wrapped = feed(RotatingKVCache(max_size = 10, keep = 2), 30)
+ chunked = feed(ChunkedKVCache(chunk_size = 8), 30)
+ slid = feed(ChunkedKVCache(chunk_size = 8), 30)
+ slid.maybe_trim_front()
+
+ assert _kv_prefix_coverage([plain]) == 30
+ assert _kv_prefix_coverage([unwrapped]) == 30
+ assert _kv_prefix_coverage([chunked]) == 30
+ assert wrapped.offset == 30 and wrapped.state[0].shape[2] == 10
+ assert _kv_prefix_coverage([wrapped]) is None
+ assert slid.start_position > 0
+ assert _kv_prefix_coverage([slid]) is None
+ assert _kv_prefix_coverage([CacheList(feed(KVCache(), 30), feed(KVCache(), 30))]) == 30
+ assert _kv_prefix_coverage([CacheList(feed(KVCache(), 30), wrapped)]) is None
+ assert _kv_prefix_coverage([feed(KVCache(), 30), feed(KVCache(), 29)]) is None
+ assert _kv_prefix_coverage([]) is None
+
+ history = _MLXPromptCacheHistory(6, 1 << 40)
+ for unsafe in (wrapped, slid):
+ history.insert("key", list(range(30)), [unsafe])
+ assert "key" not in history._lru.entries
+
+ history.insert("key", list(range(30)), [plain])
+ assert tuple(range(30)) in history._lru.entries["key"]
From 8b3c37246c38579bc9525f28066d919e30880b8f Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Wed, 22 Jul 2026 06:36:24 -0300
Subject: [PATCH 064/255] Unsloth start improvements: download progress, server
reuse, and safe model switching (#7313)
* Improve unsloth start runtime lifecycle
* Remove speculative Gemma prompt override
* Polish model download progress output
* Refine unsloth start status output
* Clarify unsloth readiness banner
* Clarify model reuse and switching output
* Queue model switches behind active inference
* Tighten unsloth start model switching
* Reduce model switch bookkeeping
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix Studio re-exec compatibility
* Recheck sidecar reservation after inference drain
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Pass start marker through child environment
* Fix key redaction, switch-waiter ordering, and stop/messaging gaps for PR #7313
- Redact minted sk-unsloth keys from the startup-failure log tail: the early
key marker lands in the server log before the model load finishes, so a
load-phase crash printed a live key to the terminal
- Deregister a finished switch waiter before releasing the swap gate so a
swap on another event loop cannot count it as still queued and unload the
model the finished request is about to generate against
- Warn on same-repo quant switches: an explicit variant replaces the resident
weights for every attached session, but the repo ids match so no switch
warning was printed
- Note the agent exit code when it is nonzero so the server keep-alive
message does not read as a successful session
- Use taskkill /T in unsloth studio stop so llama-server children stop too
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten comments in start, studio, and inference changes
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
studio/backend/routes/inference.py | 136 +++---
.../backend/tests/test_openai_auto_switch.py | 131 ++++--
unsloth_cli/commands/start.py | 438 ++++++++++++++++--
unsloth_cli/commands/studio.py | 49 +-
unsloth_cli/tests/test_start.py | 387 +++++++++++++++-
.../tests/test_studio_run_parallel_flag.py | 39 +-
6 files changed, 1009 insertions(+), 171 deletions(-)
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index d3e588bb0b..41e1fc5589 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -3406,9 +3406,8 @@ async def _acquire_swap_gate() -> None:
await asyncio.sleep(0.02)
-# Counts in-flight auto-switch requests per (target, variant). The busy guard
-# subtracts same-target waiters so concurrent requests for one model load once
-# instead of each 409-ing the other.
+# Counts auto-switch requests queued to load each (target, variant). They are not
+# generating, so the drain wait below excludes them from the active inference count.
_auto_switch_waiters: dict[tuple[str, str], int] = {}
_auto_switch_waiters_guard = threading.Lock()
@@ -3426,35 +3425,31 @@ def _note_switch_waiter(key: tuple[str, str], delta: int) -> None:
_auto_switch_waiters.pop(key, None)
-def _same_target_waiters(key: tuple[str, str]) -> int:
+def _switch_waiter_count() -> int:
with _auto_switch_waiters_guard:
- return _auto_switch_waiters.get(key, 0)
+ return sum(max(0, count) for count in _auto_switch_waiters.values())
-# A second waiter map keyed by the raw requested model, registered before the
-# (slow) resolve. The middleware counts a concurrent same-model request as
-# in-flight before it resolves and joins _auto_switch_waiters, so without this
-# the first request would see it as an unrelated request and 409.
-_auto_switch_request_waiters: dict[str, int] = {}
-_auto_switch_request_waiters_guard = threading.Lock()
+async def _wait_for_model_switch_idle(*, current_request_counted: bool) -> None:
+ """Wait until a model replacement cannot interrupt active inference.
-
-def _request_waiter_key(requested_model: str) -> str:
- return requested_model.strip().lower()
-
-
-def _note_request_waiter(key: str, delta: int) -> None:
- with _auto_switch_request_waiters_guard:
- n = _auto_switch_request_waiters.get(key, 0) + delta
- if n > 0:
- _auto_switch_request_waiters[key] = n
- else:
- _auto_switch_request_waiters.pop(key, None)
-
-
-def _same_request_waiters(key: str) -> int:
- with _auto_switch_request_waiters_guard:
- return _auto_switch_request_waiters.get(key, 0)
+ The caller holds ``inference_lifecycle_gate``, which prevents new inference
+ from starting while existing requests drain. Auto-switch requests that have
+ resolved their targets are scheduler waiters, not active generations, so
+ exclude them to avoid a queue deadlock.
+ """
+ from core.inference.llama_keepwarm import other_inference_request_count
+ while True:
+ queued_switches = _switch_waiter_count()
+ if current_request_counted and queued_switches > 0:
+ queued_switches -= 1
+ active_others = other_inference_request_count(
+ current_request_counted = current_request_counted,
+ include_pending = False,
+ )
+ if active_others <= queued_switches:
+ return
+ await asyncio.sleep(0.02)
def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]:
@@ -3582,7 +3577,6 @@ async def _maybe_auto_switch_model(
from core.inference.local_model_resolver import resolve_local_gguf
from core.inference.llama_keepwarm import (
get_last_unloaded_model,
- other_inference_request_count,
inference_lifecycle_gate,
)
@@ -3603,12 +3597,7 @@ async def _maybe_auto_switch_model(
if not auto_switch_on and get_auto_unload_idle_seconds() <= 0:
return
- # Register by the raw requested model before resolving (which can be slow):
- # the middleware already counts a concurrent same-model request as in-flight,
- # so the busy guard must know it shares this target even while it resolves.
- request_key = _request_waiter_key(requested_model)
- _note_request_waiter(request_key, 1)
- try:
+ async def _resolve_and_switch() -> None:
# Off the loop: a cold-cache rebuild walks several model dirs + HF caches.
# With auto-switch off (or an omitted-model reload-only request), skip the
# resolve so only the reload-stash path runs and no name is ever matched.
@@ -3706,6 +3695,7 @@ async def _maybe_auto_switch_model(
)
key = _switch_key(override_id, variant)
_note_switch_waiter(key, 1)
+ waiter_noted = True
try:
async with _auto_switch_lock():
# The asyncio lock is per loop; add a process-wide gate so a swap on
@@ -3718,31 +3708,6 @@ async def _maybe_auto_switch_model(
if _already_serving():
_record_serving_alias()
return
- # Single slot: refuse a cross-model swap while another inference
- # request is active rather than killing its response. Requests
- # heading to this same target (by resolved id or raw name) are
- # excluded, so concurrent requests for one model load once. A
- # pending request is still in the middleware, not generating, so
- # it is not counted here.
- same_others = max(
- _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0
- )
- others = other_inference_request_count(
- current_request_counted = True, include_pending = False
- )
- # Not gated on the GGUF being loaded: _load_model_impl also
- # tears down an active Unsloth backend before loading a GGUF,
- # so refuse whenever any other inference request is in flight.
- if others > same_others:
- raise HTTPException(
- status_code = 409,
- detail = openai_error_body(
- "Cannot switch models while another inference request is in progress.",
- status = 409,
- code = "model_switch_busy",
- param = "model",
- ),
- )
# Apply this model's saved launch flags so the swap honors the config.
override = get_model_override(override_id)
load_kwargs = {"model_path": target_id, "gguf_variant": variant}
@@ -3757,16 +3722,22 @@ async def _maybe_auto_switch_model(
LoadRequest(**load_kwargs),
fastapi_request,
current_subject,
+ current_request_counted = True,
)
# Advertise the repo id (not the concrete load path) as the loaded
# model's public id and override key for /v1/models and idle stash.
get_llama_cpp_backend()._openai_advertised_id = override_id
finally:
+ # Deregister before releasing the gate: otherwise a swap on another
+ # loop counts this finished request as queued and unloads its model.
+ _note_switch_waiter(key, -1)
+ waiter_noted = False
_auto_switch_process_lock.release()
finally:
- _note_switch_waiter(key, -1)
- finally:
- _note_request_waiter(request_key, -1)
+ if waiter_noted:
+ _note_switch_waiter(key, -1)
+
+ await _resolve_and_switch()
async def _auto_switch_from_request_body(request: Request, current_subject: str):
@@ -4186,6 +4157,15 @@ def _maybe_unsupported_message(msg: str) -> str:
return msg
+def _raise_if_sidecar_swap_in_progress() -> None:
+ from utils.transformers_version import sidecar_swap_in_progress
+ if sidecar_swap_in_progress():
+ raise HTTPException(
+ status_code = 409,
+ detail = "A transformers installation is in progress. Retry when it completes.",
+ )
+
+
@router.post("/load", response_model = LoadResponse)
async def load_model(
request: LoadRequest,
@@ -4206,24 +4186,23 @@ async def load_model(
# install can reserve while this request queues on the gate, so the pre-gate
# check alone is only a fast path.
from core.inference.llama_keepwarm import inference_lifecycle_gate
- from utils.transformers_version import sidecar_swap_in_progress
- _swap_409 = HTTPException(
- status_code = 409,
- detail = "A transformers installation is in progress. Retry when it completes.",
- )
- if sidecar_swap_in_progress():
- raise _swap_409
+ _raise_if_sidecar_swap_in_progress()
# Hold the lifecycle gate across the load so idle auto-unload can't unload the
# model mid-load. Auto-switch calls _load_model_impl directly since it already
# holds this gate.
async with inference_lifecycle_gate():
- if sidecar_swap_in_progress():
- raise _swap_409
+ _raise_if_sidecar_swap_in_progress()
return await _load_model_impl(request, fastapi_request, current_subject)
-async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str):
+async def _load_model_impl(
+ request: LoadRequest,
+ fastapi_request: Request,
+ current_subject: str,
+ *,
+ current_request_counted: bool = False,
+):
from core.inference.llama_cpp import LlamaServerNotFoundError
# A new load starts here; arm the progress throttle so this load's first
@@ -4557,6 +4536,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
),
)
+ # Keep the resident model alive until every active generation finishes;
+ # the caller's lifecycle gate blocks new starts.
+ await _wait_for_model_switch_idle(current_request_counted = current_request_counted)
+ # A sidecar install can reserve the gate while inference drains, after the
+ # route-level checks above, so recheck before replacing either backend.
+ _raise_if_sidecar_swap_in_progress()
+
# Unload any active Unsloth model only after every hub conflict check.
if unsloth_backend.active_model_name:
logger.info(
@@ -4767,6 +4753,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Unload any active GGUF model first
llama_backend = get_llama_cpp_backend()
+ await _wait_for_model_switch_idle(current_request_counted = current_request_counted)
+ _raise_if_sidecar_swap_in_progress()
if llama_backend.is_loaded:
logger.info("Unloading GGUF model before loading Unsloth model")
llama_backend.unload_model()
@@ -7096,7 +7084,7 @@ async def openai_chat_completions(
if payload.provider_id or payload.provider_type:
# External provider: this request won't touch the local GGUF, so drop it
# from the keep-warm count or its in-flight stream would falsely block a
- # concurrent local auto-switch with model_switch_busy.
+ # concurrent local model switch from proceeding.
from core.inference.llama_keepwarm import untrack_current_request
untrack_current_request(request.scope)
diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py
index 1ee9ef36d3..9361db66bb 100644
--- a/studio/backend/tests/test_openai_auto_switch.py
+++ b/studio/backend/tests/test_openai_auto_switch.py
@@ -68,7 +68,13 @@ class _LoadRecorder:
request,
fastapi_request,
current_subject = None,
+ *,
+ current_request_counted = False,
):
+ # Mirror the production load boundary before recording any replacement.
+ await inference_route._wait_for_model_switch_idle(
+ current_request_counted = current_request_counted
+ )
self.calls.append(request)
if self.fail:
from fastapi import HTTPException
@@ -94,7 +100,6 @@ def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder):
# gate that auto-switch already owns, so it calls the impl directly).
monkeypatch.setattr(inference_route, "_load_model_impl", recorder)
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
- monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
def _run_hook(model = "some/model"):
@@ -1205,10 +1210,9 @@ def test_middleware_ignores_non_post(monkeypatch):
# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ──
-def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
- # A cross-model swap must 409 (not kill) while another inference request is in
- # flight; the requesting call itself is excluded from the count.
- from fastapi import HTTPException
+def test_auto_switch_waits_for_another_inference_to_finish(monkeypatch):
+ # A cross-model swap queues while another request is generating, then loads
+ # after that request drains. The requesting call itself is excluded.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M")
@@ -1222,10 +1226,18 @@ def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch):
)
monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one
monkeypatch.setattr(kw, "_pending", 0)
- with pytest.raises(HTTPException) as exc:
- _run_hook("org/B-GGUF:Q8_0")
- assert exc.value.status_code == 409
- assert rec.calls == []
+
+ async def _drive():
+ task = asyncio.create_task(
+ inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
+ )
+ await asyncio.sleep(0.05)
+ assert rec.calls == []
+ kw._note_end() # the other generation finishes; this request remains counted
+ await asyncio.wait_for(task, timeout = 1)
+
+ asyncio.run(_drive())
+ assert len(rec.calls) == 1
def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch):
@@ -1411,13 +1423,12 @@ def test_concurrent_same_target_requests_load_once(monkeypatch):
monkeypatch.setattr(kw, "_pending", 0)
inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1)
_run_hook("org/B-GGUF:Q8_0")
- assert len(rec.calls) == 1 # loads once, no 409
+ assert len(rec.calls) == 1
-def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch):
- # A concurrent request heading to a different target still blocks the swap: the
- # same-target exclusion must not swallow a genuinely conflicting request.
- from fastapi import HTTPException
+def test_queued_different_target_does_not_deadlock_current_swap(monkeypatch):
+ # A concurrent request already queued for another target is not generating,
+ # so it must not prevent the current serialized swap from proceeding.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@@ -1432,10 +1443,8 @@ def test_swap_still_refused_when_other_request_targets_different_model(monkeypat
monkeypatch.setattr(kw, "_inflight", 2)
monkeypatch.setattr(kw, "_pending", 0)
inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1)
- with pytest.raises(HTTPException) as exc:
- _run_hook("org/B-GGUF:Q8_0")
- assert exc.value.status_code == 409
- assert rec.calls == []
+ _run_hook("org/B-GGUF:Q8_0")
+ assert len(rec.calls) == 1
def test_v1_models_advertises_repo_id_not_load_path(monkeypatch):
@@ -1481,6 +1490,37 @@ def test_load_route_holds_lifecycle_gate(monkeypatch):
assert "_load_model_impl" in src
+def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded():
+ # Both replacement directions drain active inference, then recheck whether a
+ # sidecar install reserved the lifecycle gate during that wait. Exact-model
+ # reuse exits earlier, so an already-loaded model never waits on unrelated inference.
+ import inspect
+
+ src = inspect.getsource(inference_route._load_model_impl)
+ gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:"))
+ gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait)
+ unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait)
+ standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1)
+ standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait)
+ unload_gguf = src.index("llama_backend.unload_model()", standard_wait)
+ already_loaded = src.index('status = "already_loaded"')
+
+ assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth
+ assert standard_wait < standard_sidecar_check < unload_gguf
+
+
+def test_switch_waiter_deregisters_before_swap_gate_release():
+ # A waiter left registered after the swap gate is released would let a swap on
+ # another event loop count the finished request as still queued, pass the drain
+ # early, and unload the model that request is about to generate against.
+ import inspect
+
+ src = inspect.getsource(inference_route._maybe_auto_switch_model)
+ deregister = src.index("_note_switch_waiter(key, -1)")
+ release = src.index("_auto_switch_process_lock.release()")
+ assert deregister < release
+
+
def _anthropic_payload(max_tokens = None):
from models.inference import AnthropicMessagesRequest, AnthropicMessage
return AnthropicMessagesRequest(
@@ -1519,9 +1559,9 @@ def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch):
# ── review round 6: concurrency ordering, external untrack, unload gate, ids ──
-def test_pending_same_target_request_does_not_force_409(monkeypatch):
+def test_pending_same_target_request_does_not_block_swap(monkeypatch):
# A second same-target request blocked in the middleware (pending, not yet
- # generating) must not make the first request 409: pending is excluded.
+ # generating) must not block the first request: pending is excluded.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@@ -1536,13 +1576,13 @@ def test_pending_same_target_request_does_not_force_409(monkeypatch):
monkeypatch.setattr(kw, "_inflight", 1) # just the caller
monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware
_run_hook("org/B-GGUF:Q8_0")
- assert len(rec.calls) == 1 # loads once, no 409
+ assert len(rec.calls) == 1
-def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch):
+def test_swap_waits_until_concurrent_request_finishes_resolving(monkeypatch):
# The real middleware counts a concurrent same-model request as in-flight
- # before it resolves and registers a target waiter. The raw-request waiter,
- # registered before resolve, must still exclude it so the first request loads.
+ # before it resolves and registers a target waiter. Treat it as active until
+ # its target is known, then recognize it as another queued switch request.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend("org/A-GGUF")
@@ -1556,10 +1596,20 @@ def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypat
)
monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin
monkeypatch.setattr(kw, "_pending", 0)
- # The twin has only registered its raw requested model (not yet a target waiter).
- inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1)
- _run_hook("org/B-GGUF:Q8_0")
- assert len(rec.calls) == 1 # loads once, no 409
+ # The twin is still resolving, so it is counted in-flight but has not joined
+ # the concrete target queue yet.
+
+ async def _drive():
+ task = asyncio.create_task(
+ inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
+ )
+ await asyncio.sleep(0.05)
+ assert rec.calls == []
+ inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1)
+ await asyncio.wait_for(task, timeout = 1)
+
+ asyncio.run(_drive())
+ assert len(rec.calls) == 1
def test_external_untrack_decrements_inflight_and_is_idempotent():
@@ -1595,11 +1645,9 @@ def test_manual_unload_interrupts_even_while_inference_active(monkeypatch):
assert not backend.is_loaded # torn down despite the active request
-def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
+def test_auto_switch_waits_when_unsloth_stream_active(monkeypatch):
# The GGUF slot is empty but an Unsloth model is streaming (counted in-flight).
- # _load_model_impl would unload it, so auto-switch must 409, not only when a
- # GGUF is loaded.
- from fastapi import HTTPException
+ # The replacement waits for it just as it does for a GGUF generation.
from core.inference import llama_keepwarm as kw
backend = _FakeBackend(None) # no GGUF loaded
@@ -1613,10 +1661,18 @@ def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch):
)
monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request
monkeypatch.setattr(kw, "_pending", 0)
- with pytest.raises(HTTPException) as exc:
- _run_hook("org/B-GGUF:Q8_0")
- assert exc.value.status_code == 409
- assert rec.calls == [] # the active Unsloth model is not torn down
+
+ async def _drive():
+ task = asyncio.create_task(
+ inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester")
+ )
+ await asyncio.sleep(0.05)
+ assert rec.calls == []
+ kw._note_end()
+ await asyncio.wait_for(task, timeout = 1)
+
+ asyncio.run(_drive())
+ assert len(rec.calls) == 1
def test_public_model_id_prefers_advertised_over_path():
@@ -3097,6 +3153,8 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
request,
fastapi_request,
current_subject = None,
+ *,
+ current_request_counted = False,
):
with slock:
state["cur"] += 1
@@ -3114,7 +3172,6 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch):
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load)
monkeypatch.setattr(inference_route, "_auto_switch_waiters", {})
- monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {})
barrier = threading.Barrier(2)
diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py
index 3d73df65be..317d1f4f3e 100644
--- a/unsloth_cli/commands/start.py
+++ b/unsloth_cli/commands/start.py
@@ -14,12 +14,13 @@ import signal
import subprocess
import sys
import tempfile
+import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import NamedTuple, NoReturn, Optional
-from urllib.parse import urlparse
+from urllib.parse import urlencode, urlparse
import click
import typer
@@ -105,8 +106,8 @@ _SERVE_OPTION = typer.Option(
True,
"--serve/--no-serve",
help = (
- "If no Unsloth server is running, auto-start one for --model and stop it when the "
- "agent exits. --no-serve keeps the old behavior of erroring out."
+ "If no Unsloth server is running, auto-start one for --model and keep it available "
+ "after the agent exits. --no-serve keeps the old behavior of erroring out."
),
)
# Model-load knobs mirrored from `unsloth run`; only used when --model triggers a
@@ -326,6 +327,13 @@ def _split_repo_variant(model: str) -> tuple:
return repo, variant
+def _display_model_spec(model: str, variant: Optional[str]) -> str:
+ """Return a user-facing model name that includes the selected GGUF variant."""
+ repo, inline_variant = _split_repo_variant(model)
+ selected_variant = variant or inline_variant
+ return f"{repo}:{selected_variant}" if selected_variant else model
+
+
def _fail(message: str) -> NoReturn:
typer.echo(message, err = True)
raise typer.Exit(code = 1)
@@ -373,11 +381,265 @@ def _http_json(
# A server that WE auto-started (never one we merely found). Kept at module scope so
-# _run's finally and the atexit backstop can tear it down without threading a handle
+# failure paths and the atexit backstop can tear it down without threading a handle
# through all six agent commands. Only one agent runs per process, so one slot is enough.
_auto_served_server: Optional[subprocess.Popen] = None
# Model download + load can be slow; give the auto-started server room before giving up.
_SERVER_START_TIMEOUT_S = 900
+_DOWNLOAD_POLL_INTERVAL_S = 1.0
+_START_API_KEY_PREFIX = "UNSLOTH_START_API_KEY: "
+_START_API_KEY_MARKER_ENV = "_UNSLOTH_START_API_KEY_MARKER"
+
+
+def _format_download_bytes(value: int) -> str:
+ value = max(0, int(value))
+ for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
+ if value < 1024 or unit == "TiB":
+ precision = 0 if unit in ("B", "KiB") else 1
+ return f"{value:.{precision}f} {unit}"
+ value /= 1024
+ return "0 B"
+
+
+def _format_download_eta(seconds: float) -> str:
+ seconds = max(0, int(seconds))
+ if seconds < 60:
+ return f"{seconds}s"
+ minutes, seconds = divmod(seconds, 60)
+ if minutes < 60:
+ return f"{minutes}m {seconds:02d}s"
+ hours, minutes = divmod(minutes, 60)
+ return f"{hours}h {minutes:02d}m"
+
+
+class _DownloadProgressDisplay:
+ """Render download progress without making redirected output noisy."""
+
+ def __init__(self) -> None:
+ self._samples: list[tuple[float, int]] = []
+ self._shown = False
+ self._last_bucket = -1
+ self._last_line_length = 0
+ self._last_expected = 0
+ self._interactive = bool(getattr(sys.stdout, "isatty", lambda: False)())
+
+ def update(self, progress: dict) -> None:
+ downloaded = max(0, int(progress.get("downloaded_bytes") or 0))
+ completed = max(0, int(progress.get("completed_bytes") or 0))
+ expected = max(0, int(progress.get("expected_bytes") or 0))
+ self._last_expected = max(self._last_expected, expected)
+ fraction = float(progress.get("progress") or 0)
+ if downloaded <= 0:
+ return
+ # A fully cached snapshot can report 99% with no incomplete bytes; that is
+ # not a transfer, so don't show it as a download.
+ if completed >= downloaded > 0:
+ return
+
+ now = time.monotonic()
+ if self._samples and downloaded < self._samples[-1][1]:
+ self._samples.clear()
+ self._samples.append((now, downloaded))
+ cutoff = now - 15.0
+ while len(self._samples) > 2 and self._samples[0][0] < cutoff:
+ self._samples.pop(0)
+
+ rate = 0.0
+ if len(self._samples) >= 2:
+ elapsed = self._samples[-1][0] - self._samples[0][0]
+ delta = self._samples[-1][1] - self._samples[0][1]
+ if elapsed >= 1.0 and delta > 0:
+ rate = delta / elapsed
+
+ if expected > 0:
+ # The endpoint caps at 99% while bytes remain in an incomplete file; trust it.
+ fraction = min(1.0, max(0.0, fraction))
+ percent = min(100, max(0, int(fraction * 100)))
+ filled = min(24, int(fraction * 24))
+ bar = "=" * filled + ">" + "." * max(0, 23 - filled) if filled < 24 else "=" * 24
+ line = (
+ f"Downloading model [{bar}] {percent:3d}% "
+ f"{_format_download_bytes(downloaded)} / {_format_download_bytes(expected)}"
+ )
+ bucket = percent // 10
+ if rate > 0:
+ line += f" | {_format_download_bytes(rate)}/s"
+ if downloaded < expected:
+ line += f" | ETA {_format_download_eta((expected - downloaded) / rate)}"
+ else:
+ line = f"Downloading model: {_format_download_bytes(downloaded)}"
+ bucket = downloaded // (1024**3)
+ if rate > 0:
+ line += f" | {_format_download_bytes(rate)}/s"
+
+ if self._interactive:
+ padding = " " * max(0, self._last_line_length - len(line))
+ typer.echo(f"\r{line}{padding}", nl = False)
+ sys.stdout.flush()
+ self._last_line_length = len(line)
+ elif not self._shown or bucket > self._last_bucket:
+ typer.echo(line)
+ self._last_bucket = bucket
+ self._shown = True
+
+ def close(self) -> None:
+ if self._interactive and self._shown:
+ typer.echo()
+ self._last_line_length = 0
+
+ def complete(self) -> None:
+ """Finish a displayed transfer after the model load confirms success."""
+ if not self._shown:
+ return
+ downloaded = self._samples[-1][1] if self._samples else 0
+ expected = max(downloaded, getattr(self, "_last_expected", 0))
+ self.update(
+ {
+ "downloaded_bytes": expected,
+ "expected_bytes": expected,
+ "progress": 1.0,
+ }
+ )
+
+
+def _normalized_variant(value: object) -> str:
+ return re.sub(r"[^a-z0-9]", "", str(value or "").lower())
+
+
+class _ModelDownloadProgress:
+ """Best-effort polling of the model download endpoints."""
+
+ def __init__(self, base: str, key: str, model: str, variant: Optional[str]) -> None:
+ self._base = base
+ self._key = key
+ self._model = model
+ self._variant = variant or ""
+ self._expected_bytes = 0
+ self._display = _DownloadProgressDisplay()
+ self._configured = False
+ self._disabled = not _is_hub_model_id(model)
+ self._progress_prefix = "/api/hub"
+
+ def _configure(self) -> None:
+ self._configured = True
+ if self._disabled:
+ return
+ # GGUF repos need the selected quant's size; the repo endpoint totals every
+ # quant. Resolve the variant first, otherwise show bytes only.
+ if self._variant or "gguf" in self._model.lower():
+ try:
+ params = urlencode({"repo_id": self._model})
+ try:
+ info = _http_json(
+ "GET",
+ f"{self._base}/api/hub/gguf-variants?{params}",
+ self._key,
+ timeout = 10,
+ )
+ except urllib.error.HTTPError as exc:
+ if exc.code != 404:
+ raise
+ self._progress_prefix = "/api/models"
+ info = _http_json(
+ "GET",
+ f"{self._base}/api/models/gguf-variants?{params}",
+ self._key,
+ timeout = 10,
+ )
+ self._variant = self._variant or str(info.get("default_variant") or "")
+ wanted = _normalized_variant(self._variant)
+ for item in info.get("variants") or []:
+ quant = _normalized_variant(item.get("quant"))
+ filename = _normalized_variant(item.get("filename"))
+ if wanted and (wanted == quant or wanted in filename):
+ self._expected_bytes = int(
+ item.get("download_size_bytes") or item.get("size_bytes") or 0
+ )
+ break
+ except Exception:
+ # Older servers lack this endpoint; byte progress is still useful.
+ pass
+
+ def poll(self) -> None:
+ if not self._configured:
+ self._configure()
+ if self._disabled:
+ return
+ try:
+ if self._variant or "gguf" in self._model.lower():
+ params = urlencode(
+ {
+ "repo_id": self._model,
+ "variant": self._variant,
+ "expected_bytes": self._expected_bytes,
+ }
+ )
+ url = f"{self._base}{self._progress_prefix}/gguf-download-progress?{params}"
+ else:
+ url = (
+ f"{self._base}{self._progress_prefix}/download-progress?"
+ f"{urlencode({'repo_id': self._model})}"
+ )
+ try:
+ reading = _http_json("GET", url, self._key, timeout = 10)
+ except urllib.error.HTTPError as exc:
+ if exc.code != 404 or self._progress_prefix == "/api/models":
+ raise
+ self._progress_prefix = "/api/models"
+ self.poll()
+ return
+ self._display.update(reading)
+ except Exception:
+ # Progress is best-effort; never fail the load over a polling error.
+ self._disabled = True
+
+ def close(self) -> None:
+ self._display.close()
+
+ def complete(self) -> None:
+ self._display.complete()
+
+
+def _load_model_with_progress(
+ base: str, key: str, model: str, load: LoadOptions, payload: dict
+) -> dict:
+ """Run the blocking load request while polling its download progress."""
+ result: list[tuple[bool, object]] = []
+ done = threading.Event()
+
+ def _load() -> None:
+ try:
+ value = _http_json(
+ "POST",
+ f"{base}/api/inference/load",
+ key,
+ payload,
+ timeout = 3600,
+ error = "Model load failed",
+ )
+ result.append((True, value))
+ except BaseException as exc:
+ result.append((False, exc))
+ finally:
+ done.set()
+
+ threading.Thread(target = _load, name = "unsloth-model-load", daemon = True).start()
+ progress = _ModelDownloadProgress(base, key, model, load.gguf_variant)
+ loading_announced = False
+ try:
+ while not done.wait(_DOWNLOAD_POLL_INTERVAL_S):
+ if not loading_announced:
+ typer.echo(f"Loading model: {_display_model_spec(model, load.gguf_variant)}")
+ loading_announced = True
+ progress.poll()
+ ok, value = result[0]
+ if not ok:
+ assert isinstance(value, BaseException)
+ raise value
+ progress.complete()
+ return value if isinstance(value, dict) else {}
+ finally:
+ progress.close()
def _studio_healthy(base: str, timeout: float = 3.0) -> bool:
@@ -396,6 +658,11 @@ def _log_tail(path: Path, lines: int = 20) -> str:
return "(no server log)"
+def _redacted_log_tail(path: Path, lines: int = 20) -> str:
+ """Tail with minted keys removed; only for tails shown on the terminal."""
+ return re.sub(r"sk-unsloth-\S+", "sk-unsloth-[redacted]", _log_tail(path, lines))
+
+
def _shutdown_server(server: Optional[subprocess.Popen]) -> None:
# Idempotent teardown of a server WE started, plus its own children (llama-server,
# cloudflared). A no-op once the process is already gone.
@@ -438,6 +705,14 @@ def _shutdown_auto_served() -> None:
_shutdown_server(server)
+def _keep_auto_served() -> bool:
+ """Release ownership so a successfully started server survives this CLI."""
+ global _auto_served_server
+ server, _auto_served_server = _auto_served_server, None
+ atexit.unregister(_shutdown_auto_served)
+ return server is not None and server.poll() is None
+
+
def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess.Popen:
"""Spawn `unsloth run` for `model`, wait until it is fully ready, and return it."""
global _auto_served_server
@@ -467,9 +742,8 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess
command += ["--tensor-parallel"]
log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log"
- typer.echo(
- f"No Unsloth server at {base}. Starting one for {model} (loading the model can take a while)…"
- )
+ typer.echo("Starting Unsloth server")
+ typer.echo(f"Model: {_display_model_spec(model, load.gguf_variant)}")
typer.echo(f"Server log: {log_path}")
# 0600: the `unsloth run` banner in this log carries the minted sk-unsloth- key, and
# the tempdir is world-traversable. Unlink first so a stale looser-mode file (pid
@@ -477,8 +751,17 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess
log_path.unlink(missing_ok = True)
log = os.fdopen(os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb")
# Own session/process group so a mid-session Ctrl+C (cancel a turn) doesn't reach the
- # server; we tear it down explicitly when the agent exits.
- kwargs: dict = {"stdout": log, "stderr": subprocess.STDOUT, "stdin": subprocess.DEVNULL}
+ # server. It survives a successful agent session; torn down on startup/launch failure.
+ child_env = os.environ.copy()
+ # Pass the marker via env so an older launcher ignores it instead of treating an
+ # unknown CLI flag as a llama-server arg; new launchers preserve it across re-exec.
+ child_env[_START_API_KEY_MARKER_ENV] = "1"
+ kwargs: dict = {
+ "stdout": log,
+ "stderr": subprocess.STDOUT,
+ "stdin": subprocess.DEVNULL,
+ "env": child_env,
+ }
if os.name == "nt":
kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
else:
@@ -491,17 +774,45 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess
atexit.register(_shutdown_auto_served)
deadline = time.monotonic() + _SERVER_START_TIMEOUT_S
- while time.monotonic() < deadline:
- if server.poll() is not None:
- tail = _log_tail(log_path)
- _shutdown_auto_served()
- _fail(f"The Unsloth server stopped before it was ready. Last log lines:\n{tail}")
- # `unsloth run` prints the minted key only after the server is up AND the model is
- # loaded, so it is the fully-ready signal (same contract serve-unsloth-run.sh uses).
- if _studio_healthy(base) and "sk-unsloth-" in _log_tail(log_path, lines = 400):
- typer.echo(f"Unsloth server ready at {base}.")
- return server
- time.sleep(2.0)
+ progress: Optional[_ModelDownloadProgress] = None
+ early_key_seen = False
+ try:
+ while time.monotonic() < deadline:
+ if server.poll() is not None:
+ # The early key marker lands here before load finishes; redact it.
+ tail = _redacted_log_tail(log_path)
+ _shutdown_auto_served()
+ _fail(f"The Unsloth server stopped before it was ready. Last log lines:\n{tail}")
+ tail = _log_tail(log_path, lines = 400)
+ if progress is None:
+ marker = re.search(
+ rf"^{re.escape(_START_API_KEY_PREFIX)}(sk-unsloth-[^\s]+)$",
+ tail,
+ flags = re.MULTILINE,
+ )
+ if marker:
+ early_key_seen = True
+ progress = _ModelDownloadProgress(
+ base,
+ marker.group(1),
+ model,
+ load.gguf_variant,
+ )
+ if progress is not None:
+ progress.poll()
+ # New children emit an early key marker, so wait for the final model banner;
+ # older children only print the key after load, so fall back to that.
+ ready_signal = "Model loaded:" in tail if early_key_seen else "sk-unsloth-" in tail
+ if _studio_healthy(base) and ready_signal:
+ if progress is not None:
+ progress.complete()
+ progress.close()
+ progress = None
+ return server
+ time.sleep(2.0)
+ finally:
+ if progress is not None:
+ progress.close()
_shutdown_auto_served()
_fail(
f"The Unsloth server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}."
@@ -796,6 +1107,7 @@ def _resolve_model(
load: LoadOptions = LoadOptions(),
) -> dict:
models = _loaded_models(base, key)
+ load_requested = False
# Only casefold-match ids against a loopback Unsloth, where _is_hub_model_id's
# local existence probe can actually reject a server-side path; see the note there.
allow_casefold = is_loopback_url(base)
@@ -825,11 +1137,30 @@ def _resolve_model(
)
)
if requested and match is None:
- typer.echo(
- f"Loading {requested} - please wait…"
- if load_has_overrides
- else f"Loading {requested} on the Unsloth server (this can take a while)…"
- )
+ load_requested = True
+ active = next((m for m in models if m.get("loaded") is not False), None)
+ active_id = active.get("id") if active else None
+ if active_id and not _model_id_matches(
+ active_id,
+ requested,
+ allow_casefold = allow_casefold,
+ ):
+ typer.echo(f"Switching the Unsloth server from {active_id} to {requested}.")
+ typer.echo("This unloads the current model for every attached session.")
+ elif active_id and load.gguf_variant:
+ # Same repo id but an explicit quant still replaces the resident
+ # weights; /v1/models has no variant, so ask the status endpoint.
+ try:
+ status = _http_json("GET", f"{base}/api/inference/status", key)
+ except Exception:
+ status = {}
+ resident = status.get("gguf_variant") if status.get("is_gguf") else None
+ if resident and _normalized_variant(resident) != _normalized_variant(load.gguf_variant):
+ typer.echo(
+ f"Switching the Unsloth server from {active_id}:{resident} "
+ f"to {requested}:{load.gguf_variant}."
+ )
+ typer.echo("This unloads the current model for every attached session.")
# Mirror `unsloth run`'s load knobs; keep the default payload as just
# model_path so a bare `--model` load is unchanged.
payload = {"model_path": requested}
@@ -841,14 +1172,9 @@ def _resolve_model(
payload["load_in_4bit"] = False
if load.tensor_parallel:
payload["tensor_parallel"] = True
- loaded = _http_json(
- "POST",
- f"{base}/api/inference/load",
- key,
- payload,
- timeout = 3600,
- error = "Model load failed",
- )
+ loaded = _load_model_with_progress(base, key, requested, load, payload)
+ if loaded.get("status") == "already_loaded":
+ typer.echo(f"Reusing loaded model: {_display_model_spec(requested, load.gguf_variant)}")
# Unsloth registers the model under a canonical id (resolved identifier,
# casing) that /v1/models echoes but which may differ from the path we
# passed; match on the id the load reports so we don't silently fall
@@ -861,13 +1187,16 @@ def _resolve_model(
(
m
for m in models
- if any(
+ if m.get("loaded") is not False
+ and any(
_model_id_matches(m.get("id"), w, allow_casefold = allow_casefold) for w in wanted
)
),
None,
)
if match is not None:
+ if requested and not load_requested:
+ typer.echo(f"Reusing loaded model: {_display_model_spec(requested, load.gguf_variant)}")
return match
if requested:
# We asked Unsloth to load it and it didn't surface in /v1/models; don't
@@ -881,7 +1210,13 @@ def _resolve_model(
"No model is loaded in Unsloth. Load one from the model dropdown in "
"the UI, or pass --model to load it from here."
)
- return models[0]
+ resident = next((m for m in models if m.get("loaded") is not False), None)
+ if resident is None:
+ _fail(
+ "No model is currently resident in Unsloth. Pass --model "
+ "to reload one, or load it from the model dropdown in the UI."
+ )
+ return resident
def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None:
@@ -1356,7 +1691,7 @@ def _launch(
env: dict,
install_hint: str,
unset_env: tuple = (),
-) -> NoReturn:
+) -> int:
# Resolve well-known install dirs (e.g. ~/.local/bin) first, so an already-installed
# agent not yet on PATH is found instead of prompting a needless reinstall.
_augment_path_with_install_dirs()
@@ -1382,7 +1717,7 @@ def _launch(
finally:
signal.signal(signal.SIGINT, previous)
# Negative returncode means killed by signal N; shells expect 128+N.
- raise typer.Exit(code = code if code >= 0 else 128 - code)
+ return code if code >= 0 else 128 - code
def _connect(
@@ -1434,16 +1769,35 @@ def _run(
# --no-launch recipes stay intact.
if launch and clear_screen:
click.clear()
- typer.echo(f"Unsloth {base} · model {entry['id']}")
+ typer.echo(f"Unsloth ready at {base} · model {entry['id']}")
if not launch:
env, wsl_env_bridge = _wsl_shim_env(command, env, unset_env)
_print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge)
+ if _keep_auto_served():
+ typer.echo(f"Unsloth Studio is still running at {base}.")
+ typer.echo("Stop it with: unsloth studio stop")
return
try:
- _launch(command, env, install_hint = install_hint, unset_env = unset_env)
- finally:
- # Tear down a server we auto-started once the agent session ends (no-op otherwise).
+ code = _launch(command, env, install_hint = install_hint, unset_env = unset_env)
+ except BaseException:
+ # Startup succeeded but the agent failed to launch; tear the server down
+ # rather than orphan it.
_shutdown_auto_served()
+ raise
+ auto_started = _auto_served_server is not None
+ kept = _keep_auto_served()
+ if auto_started and not kept:
+ typer.echo(f"The auto-started Unsloth server at {base} stopped during the session.")
+ raise typer.Exit(code = code)
+ if code:
+ # The server status below must not read as a successful agent session.
+ typer.echo(f"The agent exited with code {code}.")
+ if is_loopback_url(base):
+ typer.echo(f"Unsloth Studio is still running at {base}.")
+ typer.echo("Stop it with: unsloth studio stop")
+ else:
+ typer.echo(f"The remote Unsloth server is still running at {base}.")
+ raise typer.Exit(code = code)
def _agents_config_root() -> Path:
@@ -1893,7 +2247,7 @@ def codex(
launch = launch,
)
# This preflight runs after _connect may have auto-started a server but before _run
- # installs its teardown finally, so tear the server down here if it rejects the model
+ # takes over its lifecycle, so tear the server down here if it rejects the model
# (e.g. a transformers-backend model) rather than leaving it on the atexit backstop.
try:
_require_gguf_for_codex(base, key, entry["id"])
diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py
index f2f41fc583..e1924cce00 100644
--- a/unsloth_cli/commands/studio.py
+++ b/unsloth_cli/commands/studio.py
@@ -106,6 +106,13 @@ API_KEY_PBKDF2_SALT_KEY = "api_key_pbkdf2_salt"
DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"
DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at"
PBKDF2_ITERATIONS = 100_000
+_START_API_KEY_MARKER_ENV = "_UNSLOTH_START_API_KEY_MARKER"
+
+
+def _consume_start_api_key_marker_env() -> bool:
+ """Consume the one-shot readiness marker passed across a Studio re-exec."""
+ return os.environ.pop(_START_API_KEY_MARKER_ENV, None) == "1"
+
# __file__ is unsloth_cli/commands/studio.py -- two parents up is the package root
# (either site-packages or the repo root for editable installs).
@@ -1760,6 +1767,12 @@ def run(
"decode speed, MoE usually don't."
),
),
+ start_api_key_marker: bool = typer.Option(
+ False,
+ "--start-api-key-marker",
+ hidden = True,
+ help = "Emit an early API key marker for the unsloth start parent process.",
+ ),
password: str = typer.Option(
"",
"--password",
@@ -1786,6 +1799,11 @@ def run(
unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja
unsloth studio run --model unsloth/Qwen3-27B-GGUF --gguf-variant Q8_0 --tensor-parallel
"""
+ # A newer outer CLI can re-exec into an older Studio venv; pass this signal via
+ # env so an older child ignores it instead of treating it as a llama-server arg.
+ inherited_start_api_key_marker = _consume_start_api_key_marker_env()
+ start_api_key_marker = start_api_key_marker or inherited_start_api_key_marker
+
# Back-compat: --not-secure is a deprecated alias for --no-secure.
secure = _resolve_secure(secure, not_secure)
extra_llama_args: List[str] = list(ctx.args) if ctx.args else []
@@ -1991,15 +2009,21 @@ def run(
if extra_llama_args:
args.extend(extra_llama_args)
- if sys.platform == "win32":
- proc = subprocess.Popen(args)
- try:
- rc = proc.wait()
- except KeyboardInterrupt:
- rc = proc.wait()
- raise typer.Exit(rc)
- else:
- os.execvp(str(studio_bin), args)
+ if start_api_key_marker:
+ os.environ[_START_API_KEY_MARKER_ENV] = "1"
+ try:
+ if sys.platform == "win32":
+ proc = subprocess.Popen(args)
+ try:
+ rc = proc.wait()
+ except KeyboardInterrupt:
+ rc = proc.wait()
+ raise typer.Exit(rc)
+ else:
+ os.execvp(str(studio_bin), args)
+ finally:
+ # execvp doesn't return on success; restore env after a Windows wait or a failed launch.
+ os.environ.pop(_START_API_KEY_MARKER_ENV, None)
# ── 2. Start server (always suppress built-in banner) ─────────────
run_mod = _load_run_module()
@@ -2045,6 +2069,10 @@ def run(
# 4. Create API key in-process.
api_key = _create_api_key_inprocess(api_key_name)
+ if start_api_key_marker:
+ # `unsloth start` reads this key from a private 0600 log to authenticate
+ # download-progress polling; the normal `unsloth run` output is unchanged.
+ typer.echo(f"UNSLOTH_START_API_KEY: {api_key}")
# 5. Load model via HTTP.
if not silent:
@@ -2236,7 +2264,8 @@ def stop():
# Send SIGTERM (graceful shutdown) or TerminateProcess on Windows
try:
if sys.platform == "win32":
- subprocess.run(["taskkill", "/PID", str(pid), "/F"], check = True)
+ # /T also stops llama-server children, which otherwise keep GPU and port.
+ subprocess.run(["taskkill", "/PID", str(pid), "/T", "/F"], check = True)
else:
os.kill(pid, _signal.SIGTERM)
typer.echo(f"Sent shutdown signal to Unsloth server (PID {pid}).")
diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py
index 1e03d390d1..7c070fa5f4 100644
--- a/unsloth_cli/tests/test_start.py
+++ b/unsloth_cli/tests/test_start.py
@@ -20,6 +20,7 @@ if str(_REPO_ROOT) not in sys.path:
import pytest
+import typer
from typer.testing import CliRunner
import unsloth_cli.commands.start as start
@@ -639,8 +640,13 @@ def fake_studio(tmp_path, monkeypatch):
if url.endswith("/api/auth/api-keys"):
return {"key": "sk-unsloth-feedfacefeedface"}
if url.endswith("/api/inference/load"):
+ already_loaded = state["models"][0]["id"] == payload["model_path"]
state["models"] = [{"id": payload["model_path"], "context_length": 4096}]
- return {}
+ return {
+ "status": "already_loaded" if already_loaded else "loaded",
+ "model": payload["model_path"],
+ "display_name": payload["model_path"],
+ }
raise AssertionError(f"unexpected request: {method} {url}")
monkeypatch.setattr(start, "find_studio_server", lambda: BASE)
@@ -824,7 +830,7 @@ def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, t
assert profile["model"] == MODEL["id"]
-def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch):
+def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch, capsys):
calls = []
state = {"loaded": False}
@@ -862,6 +868,8 @@ def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch):
assert entry["id"] == "unsloth/gemma-4-E2B-it-GGUF"
assert any(c[1].endswith("/api/inference/load") for c in calls)
+ output = capsys.readouterr().out
+ assert "please wait" not in output
def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch):
@@ -903,6 +911,35 @@ def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch):
assert any(u.endswith("/api/inference/load") for _, u in calls)
+def test_resolve_model_does_not_attach_if_catalog_stays_unloaded(monkeypatch):
+ def http_json(
+ method,
+ url,
+ token,
+ payload = None,
+ timeout = 30,
+ error = None,
+ ):
+ if url.endswith("/v1/models"):
+ return {
+ "data": [
+ {
+ "id": "unsloth/Gemma-4-GGUF",
+ "loaded": False,
+ "context_length": 131072,
+ }
+ ]
+ }
+ if url.endswith("/api/inference/load"):
+ return {"status": "loaded", "model": "unsloth/Gemma-4-GGUF"}
+ raise AssertionError(f"unexpected request: {method} {url}")
+
+ monkeypatch.setattr(start, "_http_json", http_json)
+
+ with pytest.raises(typer.Exit):
+ start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf")
+
+
def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch):
# The mirror case: a loaded entry (loaded == True) that case-matches attaches with
# no /api/inference/load call.
@@ -931,6 +968,25 @@ def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch
assert not any(u.endswith("/api/inference/load") for _, u in calls)
+def test_resolve_model_without_request_rejects_unloaded_catalog(monkeypatch):
+ monkeypatch.setattr(
+ start,
+ "_http_json",
+ lambda *a, **k: {
+ "data": [
+ {
+ "id": "unsloth/Gemma-4-GGUF",
+ "loaded": False,
+ "context_length": 131072,
+ }
+ ]
+ },
+ )
+
+ with pytest.raises(typer.Exit):
+ start._resolve_model(BASE, "sk-test", None)
+
+
def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch):
# Against a remote Unsloth the local existence probe cannot see server-side paths,
# so a case-variant loaded id must NOT attach without a load: it could be a distinct
@@ -1213,6 +1269,9 @@ def test_connect_model_flag_loads_on_server(fake_studio):
assert loads == [
("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"})
]
+ assert result.output.index(
+ f"Switching the Unsloth server from {MODEL['id']} to unsloth/Qwen3.5-35B-A3B.\n"
+ ) < result.output.index("This unloads the current model for every attached session.\n")
_assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B")
@@ -1303,6 +1362,7 @@ def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio):
assert result.exit_code == 0, result.output
loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")]
assert loads == []
+ assert f"Reusing loaded model: {MODEL['id']}\n" in result.output
_assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"])
@@ -1324,6 +1384,7 @@ def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio):
{"model_path": MODEL["id"], "gguf_variant": "UD-Q4_K_XL"},
)
]
+ assert f"Reusing loaded model: {MODEL['id']}:UD-Q4_K_XL\n" in result.output
_assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"])
@@ -1730,8 +1791,9 @@ def _reset_auto_served():
start._auto_served_server = None
-def test_start_studio_server_builds_command_and_waits(monkeypatch):
+def test_start_studio_server_builds_command_and_waits(monkeypatch, capsys):
captured = {}
+ monkeypatch.setenv(start._START_API_KEY_MARKER_ENV, "parent")
class FakePopen:
def __init__(self, command, **kwargs):
@@ -1761,13 +1823,200 @@ def test_start_studio_server_builds_command_and_waits(monkeypatch):
assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL"
assert cmd[cmd.index("--context-length") + 1] == "8192"
assert "--tensor-parallel" in cmd
+ assert "--start-api-key-marker" not in cmd
+ assert captured["kwargs"]["env"][start._START_API_KEY_MARKER_ENV] == "1"
+ assert start.os.environ[start._START_API_KEY_MARKER_ENV] == "parent"
assert cmd[cmd.index("-p") + 1] == "8888"
assert start.LoadOptions().load_in_4bit is True and "--no-load-in-4bit" not in cmd
assert captured["kwargs"].get("start_new_session") is True # own process group
assert server.pid == 4321
+ output = capsys.readouterr().out
+ assert "Starting Unsloth server\n" in output
+ assert "Model: unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL\n" in output
+ assert "No Unsloth server at" not in output
+ assert "server ready" not in output
-def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch):
+def test_start_studio_server_polls_progress_from_early_key(monkeypatch):
+ class FakePopen:
+ pid = 4321
+
+ def poll(self):
+ return None
+
+ tails = iter(
+ [
+ "UNSLOTH_START_API_KEY: sk-unsloth-early\nLoading model...",
+ "UNSLOTH_START_API_KEY: sk-unsloth-early\nModel loaded: owner/model",
+ ]
+ )
+ created = []
+
+ class FakeProgress:
+ def __init__(self, base, key, model, variant):
+ created.append((base, key, model, variant, "created"))
+
+ def poll(self):
+ created.append("poll")
+
+ def close(self):
+ created.append("close")
+
+ def complete(self):
+ created.append("complete")
+
+ monkeypatch.setattr(start.subprocess, "Popen", lambda *a, **k: FakePopen())
+ monkeypatch.setattr(start, "_studio_healthy", lambda *a, **k: True)
+ monkeypatch.setattr(start, "_log_tail", lambda *a, **k: next(tails))
+ monkeypatch.setattr(start, "_ModelDownloadProgress", FakeProgress)
+ monkeypatch.setattr(start.time, "sleep", lambda _s: None)
+ monkeypatch.setattr(
+ start.typer,
+ "echo",
+ lambda message = "", **_kwargs: created.append(("echo", message)),
+ )
+
+ server = start._start_studio_server(
+ BASE,
+ "owner/model-GGUF",
+ start.LoadOptions(gguf_variant = "Q4_K_M"),
+ )
+
+ assert server.pid == 4321
+ assert (BASE, "sk-unsloth-early", "owner/model-GGUF", "Q4_K_M", "created") in created
+ assert created.count("poll") == 2
+ assert created[-2:] == ["complete", "close"]
+ assert not any(isinstance(event, tuple) and "server ready" in event[-1] for event in created)
+
+
+def test_load_model_with_progress_uses_selected_gguf_size(monkeypatch, capsys):
+ release = start.threading.Event()
+ calls = []
+
+ def http_json(
+ method,
+ url,
+ token,
+ payload = None,
+ timeout = 30,
+ error = None,
+ ):
+ calls.append((method, url, payload))
+ if url.endswith("/api/inference/load"):
+ assert release.wait(timeout = 2)
+ return {"model": "owner/model-GGUF"}
+ if "/api/hub/gguf-variants?" in url:
+ return {
+ "default_variant": "Q8_0",
+ "variants": [
+ {
+ "quant": "UD-Q4_K_XL",
+ "filename": "model-UD-Q4_K_XL.gguf",
+ "size_bytes": 4 * 1024**3,
+ "download_size_bytes": 4 * 1024**3,
+ }
+ ],
+ }
+ if "/api/hub/gguf-download-progress?" in url:
+ release.set()
+ return {
+ "downloaded_bytes": 2 * 1024**3,
+ "expected_bytes": 4 * 1024**3,
+ "progress": 0.5,
+ }
+ raise AssertionError(f"unexpected request: {method} {url}")
+
+ monkeypatch.setattr(start, "_http_json", http_json)
+ monkeypatch.setattr(start, "_DOWNLOAD_POLL_INTERVAL_S", 0.001)
+ result = start._load_model_with_progress(
+ BASE,
+ "sk-test",
+ "owner/model-GGUF",
+ start.LoadOptions(gguf_variant = "UD-Q4_K_XL"),
+ {"model_path": "owner/model-GGUF", "gguf_variant": "UD-Q4_K_XL"},
+ )
+
+ assert result == {"model": "owner/model-GGUF"}
+ output = capsys.readouterr().out
+ assert "Downloading model" in output
+ assert "100%" in output
+ progress_url = next(url for method, url, _ in calls if "gguf-download-progress" in url)
+ assert "variant=UD-Q4_K_XL" in progress_url
+ assert f"expected_bytes={4 * 1024**3}" in progress_url
+
+
+def test_download_progress_ignores_fully_cached_bytes(capsys):
+ display = start._DownloadProgressDisplay()
+ display.update(
+ {
+ "downloaded_bytes": 4 * 1024**3,
+ "completed_bytes": 4 * 1024**3,
+ "expected_bytes": 4 * 1024**3,
+ "progress": 0.99,
+ }
+ )
+ display.close()
+
+ assert capsys.readouterr().out == ""
+
+
+def test_resolve_model_warns_on_same_repo_quant_switch(monkeypatch, capsys):
+ models = [{"id": "owner/model-GGUF", "loaded": True}]
+
+ def http_json(
+ method,
+ url,
+ key,
+ payload = None,
+ timeout = 30,
+ error = None,
+ ):
+ assert url.endswith("/api/inference/status"), url
+ return {"is_gguf": True, "gguf_variant": "Q4_K_M"}
+
+ monkeypatch.setattr(start, "_loaded_models", lambda base, key: models)
+ monkeypatch.setattr(start, "_http_json", http_json)
+ monkeypatch.setattr(
+ start,
+ "_load_model_with_progress",
+ lambda base, key, model, load, payload: {"status": "loaded", "model": "owner/model-GGUF"},
+ )
+
+ start._resolve_model(BASE, "key", "owner/model-GGUF", start.LoadOptions(gguf_variant = "Q8_0"))
+
+ out = capsys.readouterr().out
+ assert (
+ "Switching the Unsloth server from owner/model-GGUF:Q4_K_M to owner/model-GGUF:Q8_0." in out
+ )
+ assert "every attached session" in out
+
+
+def test_resolve_model_same_quant_prints_no_switch_warning(monkeypatch, capsys):
+ models = [{"id": "owner/model-GGUF", "loaded": True}]
+
+ monkeypatch.setattr(start, "_loaded_models", lambda base, key: models)
+ monkeypatch.setattr(
+ start,
+ "_http_json",
+ lambda *a, **k: {"is_gguf": True, "gguf_variant": "Q8_0"},
+ )
+ monkeypatch.setattr(
+ start,
+ "_load_model_with_progress",
+ lambda base, key, model, load, payload: {
+ "status": "already_loaded",
+ "model": "owner/model-GGUF",
+ },
+ )
+
+ start._resolve_model(BASE, "key", "owner/model-GGUF", start.LoadOptions(gguf_variant = "Q8_0"))
+
+ out = capsys.readouterr().out
+ assert "Switching" not in out
+ assert "Reusing loaded model: owner/model-GGUF:Q8_0" in out
+
+
+def test_auto_serves_when_no_server_then_keeps_server(fake_studio, monkeypatch):
monkeypatch.setattr(start, "find_studio_server", lambda: None)
started = {}
fake = SimpleNamespace(pid = 999, poll = lambda: None)
@@ -1793,8 +2042,134 @@ def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch):
assert started["model"] == "unsloth/Qwen3-1.7B-GGUF"
assert started["load"].gguf_variant == "UD-Q4_K_XL"
assert started["base"] == BASE
- # Torn down after the agent session ended.
- assert started.get("down") is fake
+ # A successful agent exit releases ownership and leaves the server available
+ # for another terminal. Explicit startup failures still use the cleanup path.
+ assert "down" not in started
+ assert start._auto_served_server is None
+ assert "is still running" in result.output
+ assert "unsloth studio stop" in result.output
+
+
+def test_auto_served_agent_launch_failure_stops_server(fake_studio, monkeypatch):
+ monkeypatch.setattr(start, "find_studio_server", lambda: None)
+ stopped = []
+ fake = SimpleNamespace(pid = 999, poll = lambda: None)
+
+ def fake_start(*_args):
+ start._auto_served_server = fake
+ return fake
+
+ monkeypatch.setattr(start, "_start_studio_server", fake_start)
+ monkeypatch.setattr(start, "_shutdown_server", stopped.append)
+ monkeypatch.setattr(
+ start,
+ "_launch",
+ lambda *a, **k: (_ for _ in ()).throw(RuntimeError("agent launch failed")),
+ )
+
+ result = CliRunner().invoke(
+ start.start_app,
+ ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"],
+ )
+
+ assert result.exit_code == 1
+ assert stopped == [fake]
+ assert "is still running" not in result.output
+
+
+def test_auto_served_server_exit_is_not_reported_as_running(fake_studio, monkeypatch):
+ monkeypatch.setattr(start, "find_studio_server", lambda: None)
+ fake = SimpleNamespace(pid = 999, poll = lambda: 1)
+
+ def fake_start(*_args):
+ start._auto_served_server = fake
+ return fake
+
+ monkeypatch.setattr(start, "_start_studio_server", fake_start)
+ monkeypatch.setattr(start, "_launch", lambda *a, **k: 0)
+
+ result = CliRunner().invoke(
+ start.start_app,
+ ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"],
+ )
+
+ assert result.exit_code == 0, result.output
+ assert "stopped during the session" in result.output
+ assert "is still running" not in result.output
+
+
+def test_attached_server_prints_stop_hint_after_agent_exits(fake_studio, monkeypatch):
+ monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
+ monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
+ monkeypatch.setattr(
+ start.subprocess,
+ "run",
+ lambda command, env: SimpleNamespace(returncode = 0),
+ )
+
+ result = CliRunner().invoke(start.start_app, ["claude"])
+
+ assert result.exit_code == 0, result.output
+ assert f"Unsloth ready at {BASE} · model {MODEL['id']}\n" in result.output
+ assert f"Unsloth Studio is still running at {BASE}." in result.output
+ assert "Stop it with: unsloth studio stop\n" in result.output
+
+
+def test_no_launch_recipe_does_not_print_stop_hint(fake_studio):
+ result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"])
+ assert result.exit_code == 0, result.output
+ assert "is still running" not in result.output
+
+
+def test_nonzero_agent_exit_notes_code_before_stop_hint(fake_studio, monkeypatch):
+ monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
+ monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: [])
+ monkeypatch.setattr(
+ start.subprocess,
+ "run",
+ lambda command, env: SimpleNamespace(returncode = 3),
+ )
+
+ result = CliRunner().invoke(start.start_app, ["claude"])
+
+ assert result.exit_code == 3
+ assert "The agent exited with code 3." in result.output
+ assert f"Unsloth Studio is still running at {BASE}." in result.output
+
+
+def test_redacted_log_tail_strips_minted_keys(tmp_path):
+ log = tmp_path / "server.log"
+ log.write_text(
+ "booting\nUNSLOTH_START_API_KEY: sk-unsloth-feedfacefeedface\nerror: load failed\n",
+ encoding = "utf-8",
+ )
+
+ tail = start._redacted_log_tail(log)
+
+ assert "sk-unsloth-feedfacefeedface" not in tail
+ assert "sk-unsloth-[redacted]" in tail
+ assert "error: load failed" in tail
+
+
+def test_startup_failure_output_redacts_minted_key(monkeypatch, tmp_path, capsys):
+ monkeypatch.setattr(start.tempfile, "gettempdir", lambda: str(tmp_path))
+ fake = SimpleNamespace(pid = 4242, poll = lambda: 1)
+
+ def fake_popen(command, **kwargs):
+ # The child prints the early key marker, then dies before it is ready.
+ kwargs["stdout"].write(b"UNSLOTH_START_API_KEY: sk-unsloth-secretsecret\nload failed\n")
+ kwargs["stdout"].flush()
+ return fake
+
+ monkeypatch.setattr(start.subprocess, "Popen", fake_popen)
+
+ with pytest.raises(start.typer.Exit):
+ start._start_studio_server(BASE, "owner/model-GGUF", start.LoadOptions())
+
+ err = capsys.readouterr().err
+ assert "stopped before it was ready" in err
+ assert "sk-unsloth-secretsecret" not in err
+ assert "sk-unsloth-[redacted]" in err
def test_codex_preflight_failure_tears_down_auto_served(fake_studio, monkeypatch):
diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py
index 558b268a4d..74ea607753 100644
--- a/unsloth_cli/tests/test_studio_run_parallel_flag.py
+++ b/unsloth_cli/tests/test_studio_run_parallel_flag.py
@@ -170,13 +170,24 @@ def _install_reexec_capture(monkeypatch, *, platform):
monkeypatch.setattr(sys, "platform", platform)
+ def capture(kind, argv):
+ captured.append(
+ {
+ "kind": kind,
+ "argv": list(argv),
+ "start_api_key_marker": studio_mod.os.environ.get(
+ studio_mod._START_API_KEY_MARKER_ENV
+ ),
+ }
+ )
+
def fake_execvp(file, argv):
- captured.append({"kind": "execvp", "argv": list(argv)})
+ capture("execvp", argv)
raise _ExecCaptured(argv)
class _FakePopen:
def __init__(self, argv, *a, **kw):
- captured.append({"kind": "popen", "argv": list(argv)})
+ capture("popen", argv)
self._argv = argv
def wait(self):
@@ -235,6 +246,30 @@ def test_reexec_forwards_parallel_all_aliases(monkeypatch, flag, value):
), f"{flag} {value} was dropped on re-exec; argv = {argv}"
+@pytest.mark.parametrize("platform", ["linux", "darwin", "win32"])
+def test_reexec_hands_off_start_api_key_marker_out_of_band(monkeypatch, platform):
+ """A new child receives the marker while an old child sees no unknown flag."""
+ result, captured = _invoke_run(
+ monkeypatch,
+ _BASE + ["--start-api-key-marker"],
+ platform = platform,
+ )
+ assert len(captured) == 1, result.output
+ assert "--start-api-key-marker" not in captured[0]["argv"]
+ assert captured[0]["start_api_key_marker"] == "1"
+
+
+def test_reexeced_child_consumes_start_api_key_marker_env(monkeypatch):
+ """A supported child consumes the handoff before starting descendants."""
+ studio_mod = _load_run_command()
+ monkeypatch.setenv(studio_mod._START_API_KEY_MARKER_ENV, "1")
+
+ inherited = studio_mod._consume_start_api_key_marker_env()
+
+ assert inherited is True
+ assert studio_mod._START_API_KEY_MARKER_ENV not in studio_mod.os.environ
+
+
@pytest.mark.parametrize("platform", ["linux", "darwin", "win32"])
def test_reexec_argv_is_consistent_across_platforms(monkeypatch, platform):
"""Linux/Darwin (execvp) and Windows (Popen) must build the same argv."""
From f2f41bf9b1c9f873024c5b6b6d37777989b1d11a Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Wed, 22 Jul 2026 03:52:32 -0700
Subject: [PATCH 065/255] Baseline two benign unsloth-zoo test-file findings in
scan_packages (#7325)
The enforcing pip scan-packages hf-stack shard fails on two CRITICAL
staged-dropper findings in unsloth-zoo test files:
tests/test_mlx_save_export_regressions.py and
tests/test_vision_collator_audio.py. Both are false positives: the
combination heuristic matches a /tmp path literal alongside unrelated
subprocess/import references in the same file, but those are mocked test
fixtures (monkeypatch.setattr on subprocess, asserted /tmp path strings),
not droppers. Add both to the reviewed allowlist so the gate stops
red-failing on legitimate test code. The scan then exits 0 on both the
hf-stack shard and a direct unsloth-zoo scan.
---
scripts/scan_packages_baseline.json | 16 ++++++++++++++++
1 file changed, 16 insertions(+)
diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json
index 1f7bc8dcc0..936f748a74 100644
--- a/scripts/scan_packages_baseline.json
+++ b/scripts/scan_packages_baseline.json
@@ -1545,6 +1545,22 @@
"severity": "HIGH",
"evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)",
"evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d"
+ },
+ {
+ "package": "unsloth-zoo",
+ "file": "tests/test_mlx_save_export_regressions.py",
+ "check": "Writes to /tmp and executes (staged dropper)",
+ "severity": "CRITICAL",
+ "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:ab5c587f9ec31a0cc10ee55698ab133a417148d9d3f371bbc81b1e13fa119c13",
+ "evidence_hash": "93a11159147aad94f353ec4d2e0b8486b256abef88cd96d741813222cd32b138"
+ },
+ {
+ "package": "unsloth-zoo",
+ "file": "tests/test_vision_collator_audio.py",
+ "check": "Writes to /tmp and executes (staged dropper)",
+ "severity": "CRITICAL",
+ "evidence": "L111: out = extract_audio_info(msgs({\"type\": \"audio\", key: \"/tmp/a.wav\"})) sha256:2efe23ffbe2b91b8403aec9b700736919b59e5ca770f8e1f5501651b44b7d398",
+ "evidence_hash": "d416b79dd17b24214f3f7653ac01354507d7bf0fc464dee30a4a4b8998f063ba"
}
]
}
From 55433bd7b8de1bebe8d3bfada63a7a05459e45d5 Mon Sep 17 00:00:00 2001
From: Hakan Baysal
Date: Wed, 22 Jul 2026 13:55:35 +0300
Subject: [PATCH 066/255] studio: show system-wide VRAM in the multi-GPU System
tab view on ROCm (#7216)
* studio: show system-wide VRAM in the multi-GPU System tab view on ROCm
The System tab's per-GPU list comes from get_visible_gpu_utilization. When
amd-smi is unavailable (always on Windows, minimal Linux installs) it fell back
to torch, whose readings are process-local: on Windows WDDM hands each process
its own budget, so a model held by the separate llama-server process read as
~0 VRAM used even with the GPU full (#7072). The primary-GPU endpoint already
compensates with system-wide sources -- Windows Performance Counters (Task
Manager's source) and Linux DRM sysfs -- but the multi-device endpoint never
got those fallbacks.
Add per-GPU variants of both sources and overlay them onto the torch fallback:
_rocm_windows_perf_counter_vram_per_adapter_gb() attributes Dedicated Usage per
physical adapter (phys_ in the counter instance name), and
_rocm_linux_sysfs_vram_per_card_gb() reads mem_info_vram_{used,total} per DRM
card. _overlay_system_wide_vram() applies them to the device list, ROCm-only,
best-effort: unmatched adapters and ambiguous card counts keep the torch
figures, and NVIDIA paths are untouched.
Fixes #7072
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: match VRAM overlay sources by device, honor unified memory, unblock the loop
Five review fixes on the multi-GPU system-wide VRAM overlay:
1. Linux: match DRM cards to devices by PHYSICAL index instead of a positional
zip, so a reordering visibility mask (HIP_VISIBLE_DEVICES=1,0) no longer
swaps each card's figures onto the other GPU (which would mislead
auto_select_gpu_ids and the coexistence checks). An index with no matching
card keeps its torch figures.
2. Linux: skip the overlay for a device whose sysfs total is below torch's --
on unified-memory APUs (Strix Halo) mem_info_vram_total is only the small
dedicated slice while torch sees the GTT-backed pool, and
_apply_unified_memory_correction already defines larger-total-wins.
3. Windows: group counter instances by adapter LUID, not the phys_ suffix --
separate adapters each read phys_0, which collapsed every GPU into key 0.
LUIDs are mapped to 0-based positions by ascending value as the closest
stand-in for device order.
4. Windows: pair the system-wide usage with the physical capacity from
get_device_properties (as the primary-GPU fallback does) -- under WDDM
mem_get_info's "total" is the process budget, which misreported capacity
and pushed utilization to 100%.
5. Run get_visible_gpu_utilization off the event loop in the /hardware/visible
route (asyncio.to_thread, the repo's convention): the ROCm fallbacks can
shell out to PowerShell with a 5s timeout, which would stall every other
request while the System view polls.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: skip the system-wide VRAM overlay for relative GPU indices
The overlay matches its per-GPU sources (Windows perf counters, Linux sysfs) by
physical device index, but under a UUID/MIG visibility mask the torch fallback
enumerates ordinals and reports index_kind == "relative", where `index` is a
visible ordinal, not a physical id. Applying the overlay there let card/adapter
0's system-wide VRAM overwrite the torch reading of a process that actually
exposes physical GPU 1, misleading auto_select_gpu_ids and the coexistence
checks. Gate the overlay on index_kind == "physical"; relative-index paths keep
the torch fallback.
* studio: drop the unreliable Windows VRAM overlay, keep the Linux one
The multi-GPU system-wide VRAM overlay is now Linux-only. The Windows
per-adapter Performance Counter path could not be made correct: the wildcard
Get-Counter query also returns non-ROCm/iGPU adapters and LUID order is not the
ROCm device order, so an adapter's usage could be overlaid onto the wrong GPU;
and it read only Dedicated Usage, missing WDDM shared memory on unified-memory
GPUs (Strix Halo), overstating free VRAM. Rather than misattribute VRAM and
skew placement decisions, Windows keeps the process-local torch fallback (no
regression vs before this PR); Linux DRM sysfs -- matched by physical index --
still fixes #7072 for the reporter's native-Linux ROCm case.
Removes _rocm_windows_perf_counter_vram_per_adapter_gb and _torch_props_total_gb.
* studio: key sysfs VRAM by DRM card number so filtering can't renumber cards
_rocm_linux_sysfs_vram_per_card_gb dropped cards with a zero total or unreadable
files and then the overlay enumerated the compacted list, so if card0 was
dropped, card1's usage was assigned to physical GPU index 0 (equal-capacity GPUs
slip past the unified-memory total guard). Return {card_number: (used, total)}
and match a device to its card number directly: a hole stays a hole -- device 0
keeps its torch figures when card0 is absent, and card1 maps to device 1.
* studio: key system-wide VRAM by ROCm ordinal, not raw DRM card number
When a non-amdgpu adapter (Intel iGPU, a display-only card) owns an earlier
DRM slot, DRM card numbers stop equalling ROCm device ordinals -- Intel card0
plus AMD card1/card2 gives ROCm devices 0/1, so keying the sysfs overlay by
card number handed ROCm device 1 card1's data (AMD device 0) and left device 0
on stale torch figures, corrupting free-VRAM placement on equal-capacity GPUs.
Only amdgpu cards expose mem_info_vram_*, so the glob already excludes foreign
adapters; order the surviving cards by their PCI address (ROCm/HIP's default
device order, read from each card's device symlink) and key by that position --
the ROCm physical ordinal, which is what the overlay matches against dev index.
An unreadable / zero-total amdgpu card still consumes its ordinal so a later
card is never renumbered onto its slot.
* studio: skip the VRAM overlay under layered HIP-over-ROCR masks
ROCR_VISIBLE_DEVICES filters physical GPUs at the HSA/ROCr layer, and a
HIP_VISIBLE_DEVICES set on top selects WITHIN that already-filtered set
(apply_gpu_ids sets HIP while leaving an inherited ROCR mask in place). When
both are active _get_parent_visible_gpu_spec() prefers the HIP value, so the
reported device index is a ROCR-relative ordinal, not a physical GPU id --
overlaying DRM-sysfs figures by that index would pull another GPU's usage
(e.g. ROCR=2,3 + HIP=1 is physical GPU 3, but the overlay would read card 1),
and equal-capacity cards bypass the total-size safeguard. Detect layered masks
and keep torch's process-local figures there rather than risk misattribution;
a single mask still leaves the index physical and is overlaid as before.
* studio: only overlay whole-card VRAM onto 1:1 ROCm devices
The overlay guard only skipped the case where sysfs total < torch total
(unified-memory APUs), so a partitioned ROCm device (MI300 in CPX mode) --
where HIP exposes several logical devices per physical card but sysfs reports
the whole card's aggregate -- passed the guard: the card total exceeds a
partition's torch total, and the overlay overwrote the partition with
whole-card usage and capacity, letting downstream selection think a partition
had the entire card free. Require the sysfs card total to match the torch
device total (within ~10%) so a mismatch in either direction -- unified memory
(sysfs smaller) or partitioning (sysfs larger) -- keeps torch's figures.
* studio: treat CUDA-over-ROCR as layered, enumerate AMD cards by driver
Two remaining mismatches between the reported device index and the DRM card the
overlay reads:
- On ROCm the HIP layer honors CUDA_VISIBLE_DEVICES as well as
HIP_VISIBLE_DEVICES, so a CUDA mask composed over ROCR layers identically:
ROCR=2,3 with CUDA=1 is physical GPU 3, yet the spec reports the ROCR value
[2,3] and the device was labeled index 2, overlaying card 2's usage onto GPU 3.
The layered check now treats ROCR combined with either HIP or CUDA as layered.
- The ROCm device set is now enumerated by bound driver (device/driver resolves
to amdgpu) instead of by the presence of mem_info_vram_*. An AMD device with
incomplete sysfs support (some APUs expose no VRAM files at all) was omitted
by the glob entirely and shifted every later card down one ordinal, letting a
similar-capacity GPU pass the total guard with another device's usage. Such a
card now consumes its ordinal and simply yields no entry.
* studio: honor GPU_DEVICE_ORDINAL and require an unambiguous card mapping
Two remaining ways the reported device index could be matched to the wrong DRM
card:
- GPU_DEVICE_ORDINAL is a supported ROCm visibility variable that
_get_parent_visible_gpu_spec() never consults, so GPU_DEVICE_ORDINAL=1
surfaces physical GPU 1 as torch ordinal 0 and it was mislabeled index 0,
overlaying card 0's usage onto GPU 1. The mask check now covers it, and is
renamed _rocm_device_index_unreliable() to say what it actually decides.
- driver == amdgpu is only a SUPERSET of the ROCm-visible set: an amdgpu-bound
adapter HIP cannot enumerate (an unsupported older AMD GPU beside a supported
one) still took an ordinal and shifted every real compute device. There is no
torch-side PCI identity to match against, so the overlay now requires the
amdgpu card count to equal the device count -- exactly the condition under
which position-in-PCI-order is a sound 1:1 mapping. Any disagreement keeps
torch's process-local figures: less informative, never misattributed.
* studio: keep the VRAM overlay working for masked GPU subsets
The card-count guard compared the amdgpu card list against the VISIBLE device
list, so any visibility mask disabled the overlay outright: HIP_VISIBLE_DEVICES=1,3
on a four-GPU host gives two devices against four cards. Those masked GPUs then
kept reporting process-local torch usage, hiding VRAM held by llama-server and
letting the training/chat placement checks overestimate free memory -- the exact
problem the overlay exists to fix.
The count check now applies only when no visibility mask is active, which is the
case where the reported devices really are the whole host and a mismatch means an
amdgpu adapter ROCm cannot enumerate is shifting the ordinals. Under a mask the
subset is expected, so each device's physical index is validated individually
instead: the per-card lookup bounds-checks it and the total-size guard rejects a
card whose capacity does not match the device's.
* studio: match GPUs to DRM cards by PCI identity, not by position
Every mapping bug on this PR came from the same root cause: there was no
authoritative link between a reported device index and a DRM card, so the
overlay kept inferring one positionally and each heuristic broke on a new host
shape -- foreign adapters on earlier DRM slots, cards with no VRAM sysfs, and
most recently amdgpu-bound adapters HIP cannot enumerate, which the count guard
could only catch on an unmasked host and therefore missed under any mask.
Use the link ROCm itself enumerates from. KFD topology
(/sys/class/kfd/kfd/topology/nodes//properties) lists exactly the GPUs HIP
exposes -- GPU nodes in node-id order are HIP's device order -- and each carries
its PCI location, so index N there IS physical device N with a stable identity.
DRM sysfs now supplies system-wide VRAM keyed by that same PCI address, and the
overlay is a join on it.
Every previous skew becomes a failed join rather than a misattribution: an
unenumerable adapter has no KFD node so it never takes an ordinal, a foreign
adapter contributes no entry, and a masked subset resolves each physical index
directly. That removes the count heuristic and its mask exception entirely. With
no KFD topology there is no identity to join on, so the overlay is skipped rather
than guessing positionally.
* studio: require verified host visibility and AMD-only KFD nodes
Three ways the identity map could still be built on a false premise:
- The NVIDIA open kernel module registers KFD topology nodes with a positive
SIMD count, so an earlier NVIDIA node shifted every AMD ordinal and ROCm
device 1 resolved to AMD GPU 0. GPU nodes now require vendor_id 4098 (0x1002),
the same filter install.sh already applies for this exact reason.
- A GPU node with an unreadable properties file or no location_id was skipped,
which silently shifted every later ordinal. Both now fail the whole map
closed, so the overlay is disabled rather than misattributing.
- A container exposing only some render devices through device cgroups sets no
visibility variable, yet torch compacts what it can see to ordinals from zero
while the host-mounted KFD and DRM trees still list every GPU. Nothing in the
reported payload distinguishes that from a full host, and torch exposes no PCI
id to check against, so the overlay now runs only when host visibility is
positively verified: no visibility mask AND device count equal to the host GPU
count. That also subsumes the previous layered-mask and GPU_DEVICE_ORDINAL
checks, so _rocm_device_index_unreliable() is gone.
This trades coverage for correctness: masked subsets and filtered containers now
keep torch's process-local figures instead of a mapping that cannot be verified.
* Fix the multi-GPU VRAM overlay docstring for PR #7216
The docstring claimed a reordering mask keeps each card on the right GPU,
but the overlay skips any active visibility mask and keeps torch's figures.
State the actual gating instead.
* Tighten comments in the multi-GPU VRAM overlay and its tests
Collapse the verbose docstrings and inline explanations added for the Linux
ROCm system-wide VRAM overlay to succinct one-liners, keeping the non-obvious
rationale (fail-closed KFD mapping, PCI-identity join, mask gating, the 10%
whole-card guard). Comments only, no behavior change.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen
---
studio/backend/routes/training.py | 4 +-
.../test_rocm_multi_gpu_vram_system_wide.py | 554 ++++++++++++++++++
studio/backend/utils/hardware/hardware.py | 210 +++++++
3 files changed, 767 insertions(+), 1 deletion(-)
create mode 100644 studio/backend/tests/test_rocm_multi_gpu_vram_system_wide.py
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index a8a9874b1b..9176f1a8da 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -109,7 +109,9 @@ async def get_hardware_utilization(current_subject: str = Depends(get_current_su
@router.get("/hardware/visible")
async def get_visible_hardware_utilization(current_subject: str = Depends(get_current_subject)):
from utils.hardware import get_visible_gpu_utilization
- return get_visible_gpu_utilization()
+
+ # Off the event loop: the ROCm fallbacks shell out (Windows perf counters, sysfs) and the System view polls this route.
+ return await asyncio.to_thread(get_visible_gpu_utilization)
@router.post("/start")
diff --git a/studio/backend/tests/test_rocm_multi_gpu_vram_system_wide.py b/studio/backend/tests/test_rocm_multi_gpu_vram_system_wide.py
new file mode 100644
index 0000000000..bdafdeae9b
--- /dev/null
+++ b/studio/backend/tests/test_rocm_multi_gpu_vram_system_wide.py
@@ -0,0 +1,554 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""The System tab's multi-GPU view must show system-wide VRAM on ROCm (#7072).
+
+When amd-smi is unavailable, get_visible_gpu_utilization fell back to torch,
+whose readings are process-local: a model held by the separate llama-server
+process read as ~0 VRAM used even with the GPU full. These tests cover the
+per-GPU system-wide overlay the multi-device endpoint now applies, matched by
+physical device identity.
+"""
+
+from __future__ import annotations
+
+import importlib
+import sys
+import types
+from pathlib import Path
+
+_BACKEND_DIR = Path(__file__).resolve().parent.parent
+if str(_BACKEND_DIR) not in sys.path:
+ sys.path.insert(0, str(_BACKEND_DIR))
+
+
+def _maybe_stub(name: str, builder):
+ # Stub only if the real module is missing, so we never shadow it for later tests.
+ try:
+ importlib.import_module(name)
+ except ImportError:
+ sys.modules[name] = builder()
+
+
+def _build_loggers_stub():
+ m = types.ModuleType("loggers")
+ m.get_logger = lambda name: __import__("logging").getLogger(name)
+ return m
+
+
+def _build_structlog_stub():
+ m = types.ModuleType("structlog")
+ m.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
+ return m
+
+
+_maybe_stub("loggers", _build_loggers_stub)
+_maybe_stub("structlog", _build_structlog_stub)
+
+import utils.hardware.hardware as hw # noqa: E402
+
+
+def _device(
+ index,
+ used,
+ total,
+ *,
+ ordinal = None,
+):
+ return {
+ "index": index,
+ "index_kind": "physical",
+ "visible_ordinal": index if ordinal is None else ordinal,
+ "gpu_utilization_pct": None,
+ "temperature_c": None,
+ "vram_used_gb": used,
+ "vram_total_gb": total,
+ "vram_utilization_pct": round((used / total) * 100, 1) if total > 0 else None,
+ "power_draw_w": None,
+ "power_limit_w": None,
+ "power_utilization_pct": None,
+ }
+
+
+# ── Linux per-card sysfs ──
+
+
+def _fake_drm(tmp_path, monkeypatch, cards):
+ """Fake /sys/class/drm tree; glob returns cards REVERSED so the PCI sort must order them.
+
+ ``cards``: (card_no, pci_bdf, driver, vram) tuples; vram is (used_gb, total_gb)
+ or None for a device with no mem_info_vram_* files.
+ """
+ drivers = tmp_path / "drivers"
+ card_paths = []
+ for card_no, bdf, driver, vram in cards:
+ pci_dir = tmp_path / "pci" / bdf
+ pci_dir.mkdir(parents = True, exist_ok = True)
+ drv_dir = drivers / driver
+ drv_dir.mkdir(parents = True, exist_ok = True)
+ (pci_dir / "driver").symlink_to(drv_dir)
+ if vram is not None:
+ used, total = vram
+ (pci_dir / "mem_info_vram_used").write_text(str(int(used * 1024**3)))
+ (pci_dir / "mem_info_vram_total").write_text(str(int(total * 1024**3)))
+ card_dir = tmp_path / "drm" / f"card{card_no}"
+ card_dir.mkdir(parents = True, exist_ok = True)
+ (card_dir / "device").symlink_to(pci_dir)
+ card_paths.append(str(card_dir))
+ monkeypatch.setattr(hw.glob, "glob", lambda pattern: list(reversed(card_paths)))
+ return card_paths
+
+
+def test_linux_vram_keyed_by_pci_excludes_foreign_adapters(monkeypatch, tmp_path):
+ # Foreign (non-amdgpu) adapters contribute no entry, so they cannot shift ordinals.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _fake_drm(
+ tmp_path,
+ monkeypatch,
+ [
+ (0, "0000:00:02.0", "i915", (0.5, 2.0)), # foreign adapter: excluded
+ (1, "0000:03:00.0", "amdgpu", (40, 48)), # AMD device 0
+ (2, "0000:41:00.0", "amdgpu", (1, 8)), # AMD device 1
+ ],
+ )
+ assert hw._rocm_linux_sysfs_vram_by_pci_gb() == {
+ "0000:03:00.0": (40.0, 48.0),
+ "0000:41:00.0": (1.0, 8.0),
+ }
+
+
+def test_linux_vram_omits_bad_cards_without_shifting(monkeypatch, tmp_path):
+ # A zero-total card has no entry; identity keying means its absence renumbers nothing.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _fake_drm(
+ tmp_path,
+ monkeypatch,
+ [
+ (0, "0000:03:00.0", "amdgpu", (0, 0)), # zero total -> no entry
+ (1, "0000:41:00.0", "amdgpu", (2, 16)),
+ ],
+ )
+ assert hw._rocm_linux_sysfs_vram_by_pci_gb() == {"0000:41:00.0": (2.0, 16.0)}
+
+
+def test_linux_vram_omits_amd_card_without_vram_files(monkeypatch, tmp_path):
+ # An APU with no mem_info_vram_* files has no entry; the discrete card keeps its address.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _fake_drm(
+ tmp_path,
+ monkeypatch,
+ [
+ (0, "0000:03:00.0", "amdgpu", None), # APU: no VRAM sysfs files
+ (1, "0000:41:00.0", "amdgpu", (2, 16)),
+ ],
+ )
+ assert hw._rocm_linux_sysfs_vram_by_pci_gb() == {"0000:41:00.0": (2.0, 16.0)}
+
+
+# ── KFD topology: the authoritative ROCm device order ──
+
+
+_AMD = 4098 # 0x1002
+_NVIDIA = 4318 # 0x10DE -- the open kernel module also registers KFD nodes
+
+
+def _fake_kfd(tmp_path, monkeypatch, nodes):
+ """Fake KFD topology nodes tree, returned out of node order so the sort must order it.
+
+ ``nodes``: (node_id, simd_count, location_id, domain, vendor_id); simd_count 0
+ marks a CPU node, location_id None omits the property.
+ """
+ node_paths = []
+ for node_id, simd_count, location_id, domain, vendor_id in nodes:
+ d = tmp_path / "kfd" / str(node_id)
+ d.mkdir(parents = True, exist_ok = True)
+ lines = [f"cpu_cores_count {0 if simd_count else 8}", f"simd_count {simd_count}"]
+ if location_id is not None:
+ lines.append(f"location_id {location_id}")
+ lines.append(f"domain {domain}")
+ if vendor_id is not None:
+ lines.append(f"vendor_id {vendor_id}")
+ (d / "properties").write_text("\n".join(lines) + "\n")
+ node_paths.append(str(d))
+ monkeypatch.setattr(hw.glob, "glob", lambda pattern: list(reversed(node_paths)))
+ return node_paths
+
+
+def test_kfd_lists_gpu_nodes_in_device_order(monkeypatch, tmp_path):
+ # The CPU node (simd_count 0) takes no ordinal; GPU nodes in node-id order are HIP's order.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _fake_kfd(
+ tmp_path,
+ monkeypatch,
+ [
+ (0, 0, None, 0, None), # CPU node
+ (1, 304, (0x03 << 8) | (0x00 << 3) | 0, 0, _AMD), # 0000:03:00.0 -> dev 0
+ (2, 304, (0x41 << 8) | (0x00 << 3) | 0, 0, _AMD), # 0000:41:00.0 -> dev 1
+ ],
+ )
+ assert hw._rocm_kfd_gpu_pci_ids() == ["0000:03:00.0", "0000:41:00.0"]
+
+
+def test_kfd_decodes_domain_device_and_function(monkeypatch, tmp_path):
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _fake_kfd(tmp_path, monkeypatch, [(1, 64, (0xC1 << 8) | (0x1F << 3) | 5, 0x1234, _AMD)])
+ assert hw._rocm_kfd_gpu_pci_ids() == ["1234:c1:1f.5"]
+
+
+def test_kfd_skips_non_amd_gpu_nodes(monkeypatch, tmp_path):
+ # An NVIDIA KFD node is not a HIP device: it must take no ordinal, else it
+ # shifts every AMD GPU and ROCm device 1 resolves to AMD GPU 0.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _fake_kfd(
+ tmp_path,
+ monkeypatch,
+ [
+ (0, 0, None, 0, None), # CPU
+ (1, 128, (0x01 << 8) | 0, 0, _NVIDIA), # NVIDIA: no ordinal
+ (2, 304, (0x03 << 8) | 0, 0, _AMD), # AMD device 0
+ (3, 304, (0x41 << 8) | 0, 0, _AMD), # AMD device 1
+ ],
+ )
+ assert hw._rocm_kfd_gpu_pci_ids() == ["0000:03:00.0", "0000:41:00.0"]
+
+
+def test_kfd_fails_closed_when_a_gpu_has_no_location(monkeypatch, tmp_path):
+ # Dropping an unplaceable AMD GPU shifts later ordinals; fail closed for the whole map.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _fake_kfd(
+ tmp_path,
+ monkeypatch,
+ [
+ (1, 304, None, 0, _AMD), # AMD GPU with no location_id
+ (2, 304, (0x41 << 8) | 0, 0, _AMD),
+ ],
+ )
+ assert hw._rocm_kfd_gpu_pci_ids() == []
+
+
+def test_kfd_fails_closed_when_a_node_is_unreadable(monkeypatch, tmp_path):
+ # An unreadable node could be a GPU; assuming otherwise would shift ordinals.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ paths = _fake_kfd(
+ tmp_path,
+ monkeypatch,
+ [
+ (1, 304, (0x03 << 8) | 0, 0, _AMD),
+ (2, 304, (0x41 << 8) | 0, 0, _AMD),
+ ],
+ )
+ (Path(paths[0]) / "properties").unlink()
+ assert hw._rocm_kfd_gpu_pci_ids() == []
+
+
+def test_kfd_absent_yields_no_device_order(monkeypatch):
+ monkeypatch.setattr(hw.glob, "glob", lambda pattern: [])
+ assert hw._rocm_kfd_gpu_pci_ids() == []
+
+
+# ── overlay ──
+
+
+def _patch_pci_map(monkeypatch, bdfs):
+ """Declare the ROCm device order by PCI address (index N is device N) and clear
+ the visibility masks the overlay requires unset.
+ """
+ for var in (
+ "HIP_VISIBLE_DEVICES",
+ "ROCR_VISIBLE_DEVICES",
+ "CUDA_VISIBLE_DEVICES",
+ "GPU_DEVICE_ORDINAL",
+ ):
+ monkeypatch.delenv(var, raising = False)
+ monkeypatch.setattr(hw, "_rocm_kfd_gpu_pci_ids", lambda: list(bdfs))
+
+
+def _pci(n):
+ """A distinct, well-formed PCI address for card n."""
+ return f"0000:{n:02x}:00.0"
+
+
+def test_overlay_windows_is_noop_keeps_torch(monkeypatch):
+ # Windows is intentionally not overlaid (perf counters can't map to ROCm ordinals): keep torch.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Windows")
+ monkeypatch.setattr(
+ hw,
+ "_rocm_linux_sysfs_vram_by_pci_gb",
+ lambda: (_ for _ in ()).throw(AssertionError("sysfs must not run on Windows")),
+ )
+ devices = [_device(0, used = 0.02, total = 8.0)]
+ _patch_pci_map(monkeypatch, [_pci(0)])
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.02 # untouched
+
+
+def test_overlay_linux_matches_by_device_ordinal(monkeypatch):
+ # Devices arriving as [index 1, index 0] each get their own GPU's figures by ordinal.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(
+ hw,
+ "_rocm_linux_sysfs_vram_by_pci_gb",
+ lambda: {_pci(0): (30.0, 45.0), _pci(1): (0.5, 8.0)}, # dev 0 big, dev 1 small
+ )
+ devices = [_device(1, used = 0.01, total = 8.0), _device(0, used = 0.02, total = 45.0)]
+ _patch_pci_map(monkeypatch, [_pci(0), _pci(1)])
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.5 # index 1 -> device 1 (small)
+ assert devices[0]["vram_total_gb"] == 8.0
+ assert devices[1]["vram_used_gb"] == 30.0 # index 0 -> device 0 (big)
+ assert devices[1]["vram_total_gb"] == 45.0
+
+
+def test_overlay_linux_ordinal_hole_does_not_shift(monkeypatch):
+ # Device 0's card dropped: index 0 keeps torch, index 1 still maps to ordinal 1 (no compaction).
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(1): (0.5, 8.0)})
+ devices = [_device(0, used = 0.02, total = 45.0), _device(1, used = 0.01, total = 8.0)]
+ _patch_pci_map(monkeypatch, [_pci(0), _pci(1)])
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.02 # no ordinal 0 -> torch kept
+ assert devices[1]["vram_used_gb"] == 0.5 # ordinal 1 -> device 1, not device 0
+
+
+def test_overlay_linux_skips_unified_memory_card(monkeypatch):
+ # Unified-memory APU: the smaller sysfs total must not shrink torch's GTT-backed pool.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (0.4, 1.0)})
+ devices = [_device(0, used = 12.0, total = 96.0)] # torch's unified pool
+ _patch_pci_map(monkeypatch, [_pci(0)])
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 12.0
+ assert devices[0]["vram_total_gb"] == 96.0
+
+
+def test_overlay_linux_skips_partitioned_device(monkeypatch):
+ # Partitioned MI300: the whole-card sysfs total dwarfs the partition, so the overlay must not overwrite it.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (40.0, 192.0)})
+ devices = [_device(0, used = 1.0, total = 24.0)] # torch partition
+ _patch_pci_map(monkeypatch, [_pci(0)])
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 1.0 # partition figures kept
+ assert devices[0]["vram_total_gb"] == 24.0
+
+
+def test_overlay_linux_out_of_range_index_untouched(monkeypatch):
+ # A masked host exposing physical index 5 with no card 5: keep torch data.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(
+ hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (30.0, 45.0), _pci(1): (0.5, 8.0)}
+ )
+ devices = [_device(5, used = 0.02, total = 45.0)]
+ _patch_pci_map(monkeypatch, [_pci(0), _pci(1)])
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.02
+
+
+def test_overlay_ignores_adapters_rocm_cannot_enumerate(monkeypatch):
+ # A HIP-unenumerable amdgpu adapter has no KFD node, so device 0 resolves to
+ # the supported GPU's own address, never the display card's.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(
+ hw,
+ "_rocm_linux_sysfs_vram_by_pci_gb",
+ # Both in DRM sysfs with similar capacity -- what the total-size guard can't separate.
+ lambda: {_pci(9): (30.0, 45.0), _pci(3): (12.0, 45.0)},
+ )
+ _patch_pci_map(monkeypatch, [_pci(3)]) # KFD lists only the supported GPU
+ devices = [_device(0, used = 0.02, total = 45.0)] # torch sees that one GPU
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 12.0 # the supported GPU's own figures
+
+
+def test_overlay_skips_masked_subsets(monkeypatch):
+ # Under a mask the index is not verifiably a host ordinal, so keep torch's figures.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _patch_pci_map(monkeypatch, [_pci(0), _pci(1), _pci(2), _pci(3)])
+ monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1,3")
+ monkeypatch.setattr(
+ hw,
+ "_rocm_linux_sysfs_vram_by_pci_gb",
+ lambda: {_pci(1): (30.0, 48.0), _pci(3): (12.0, 48.0)},
+ )
+ devices = [_device(1, used = 0.02, total = 48.0), _device(3, used = 0.01, total = 48.0)]
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.02 # torch kept
+ assert devices[1]["vram_used_gb"] == 0.01
+
+
+def test_overlay_skips_device_cgroup_filtered_container(monkeypatch):
+ # A device-cgroup container sets no env var yet compacts torch's indices from
+ # zero while KFD/DRM list every GPU, so the count mismatch must disable the overlay.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _patch_pci_map(monkeypatch, [_pci(0), _pci(1), _pci(2), _pci(3)]) # host has 4
+ monkeypatch.setattr(
+ hw,
+ "_rocm_linux_sysfs_vram_by_pci_gb",
+ lambda: {_pci(0): (30.0, 48.0), _pci(2): (12.0, 48.0)},
+ )
+ devices = [_device(0, used = 0.02, total = 48.0)] # container sees 1, as index 0
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.02 # torch kept, not host GPU 0's 30.0
+
+
+def test_overlay_skips_without_kfd_topology(monkeypatch):
+ # No KFD means no identity to join on; fall back to torch rather than guess.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(hw, "_rocm_kfd_gpu_pci_ids", lambda: [])
+ monkeypatch.setattr(
+ hw,
+ "_rocm_linux_sysfs_vram_by_pci_gb",
+ lambda: (_ for _ in ()).throw(AssertionError("must not read sysfs without KFD")),
+ )
+ devices = [_device(0, used = 0.02, total = 45.0)]
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.02
+
+
+def test_overlay_empty_devices_is_noop(monkeypatch):
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ hw._overlay_system_wide_vram([]) # must not raise
+
+
+# ── integration: the ROCm torch fallback applies the overlay ──
+
+
+def test_visible_utilization_rocm_fallback_overlays(monkeypatch):
+ for _var in (
+ "HIP_VISIBLE_DEVICES",
+ "ROCR_VISIBLE_DEVICES",
+ "CUDA_VISIBLE_DEVICES",
+ "GPU_DEVICE_ORDINAL",
+ ):
+ monkeypatch.delenv(_var, raising = False)
+ monkeypatch.setattr(hw, "IS_ROCM", True)
+ monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
+ monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) # amd-smi unavailable
+ monkeypatch.setattr(
+ hw,
+ "_get_parent_visible_gpu_spec",
+ lambda: {"raw": None, "numeric_ids": [0, 1], "supports_explicit_gpu_ids": True},
+ )
+ monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [0, 1])
+ monkeypatch.setattr(
+ hw,
+ "_torch_get_per_device_info",
+ lambda ids: [
+ {"index": 0, "visible_ordinal": 0, "used_gb": 0.02, "total_gb": 45.0},
+ {"index": 1, "visible_ordinal": 1, "used_gb": 0.01, "total_gb": 8.0},
+ ],
+ )
+ overlaid = []
+ monkeypatch.setattr(
+ hw, "_overlay_system_wide_vram", lambda devices: overlaid.append(len(devices))
+ )
+ result = hw.get_visible_gpu_utilization()
+ assert result["available"] is True
+ assert overlaid == [2]
+
+
+def test_visible_utilization_relative_index_skips_overlay(monkeypatch):
+ # UUID/MIG mask gives relative indices; the overlay matches physical index, so it must not run.
+ monkeypatch.setattr(hw, "IS_ROCM", True)
+ monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
+ monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None)
+ monkeypatch.setattr(
+ hw,
+ "_get_parent_visible_gpu_spec",
+ lambda: {"raw": "GPU-uuid-a", "numeric_ids": None, "supports_explicit_gpu_ids": False},
+ )
+ monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: []) # UUID mask
+ monkeypatch.setattr(hw, "_torch_get_physical_gpu_count", lambda: 1)
+ monkeypatch.setattr(
+ hw,
+ "_torch_get_per_device_info",
+ lambda ids: [{"index": 0, "visible_ordinal": 0, "used_gb": 0.02, "total_gb": 8.0}],
+ )
+ called = []
+ monkeypatch.setattr(hw, "_overlay_system_wide_vram", lambda devices: called.append(1))
+ result = hw.get_visible_gpu_utilization()
+ assert result["index_kind"] == "relative"
+ assert called == []
+
+
+def test_visible_utilization_nvidia_fallback_skips_overlay(monkeypatch):
+ monkeypatch.setattr(hw, "IS_ROCM", False)
+ monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
+ monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None)
+ monkeypatch.setattr(
+ hw,
+ "_get_parent_visible_gpu_spec",
+ lambda: {"raw": None, "numeric_ids": [0], "supports_explicit_gpu_ids": True},
+ )
+ monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [0])
+ monkeypatch.setattr(
+ hw,
+ "_torch_get_per_device_info",
+ lambda ids: [{"index": 0, "visible_ordinal": 0, "used_gb": 1.0, "total_gb": 24.0}],
+ )
+ called = []
+ monkeypatch.setattr(hw, "_overlay_system_wide_vram", lambda devices: called.append(1))
+ result = hw.get_visible_gpu_utilization()
+ assert result["available"] is True
+ assert called == []
+
+
+def test_any_visibility_mask_is_detected(monkeypatch):
+ # Any of these makes the index not a host-physical ordinal, so each must disable the overlay.
+ for var in (
+ "HIP_VISIBLE_DEVICES",
+ "ROCR_VISIBLE_DEVICES",
+ "CUDA_VISIBLE_DEVICES",
+ "GPU_DEVICE_ORDINAL",
+ ):
+ monkeypatch.delenv(var, raising = False)
+ assert hw._rocm_visibility_mask_active() is False
+ for var in (
+ "HIP_VISIBLE_DEVICES",
+ "ROCR_VISIBLE_DEVICES",
+ "CUDA_VISIBLE_DEVICES",
+ "GPU_DEVICE_ORDINAL",
+ ):
+ monkeypatch.setenv(var, "1")
+ assert hw._rocm_visibility_mask_active() is True, var
+ monkeypatch.setenv(var, " ") # empty is not an active filter
+ assert hw._rocm_visibility_mask_active() is False, var
+ monkeypatch.delenv(var, raising = False)
+
+
+def test_overlay_skips_under_gpu_device_ordinal(monkeypatch):
+ # GPU_DEVICE_ORDINAL=1 surfaces GPU 1 as torch ordinal 0, so index 0 is not GPU 0; overlay must not run.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ _patch_pci_map(monkeypatch, [_pci(0)])
+ monkeypatch.setenv("GPU_DEVICE_ORDINAL", "1")
+ monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(0): (30.0, 45.0)})
+ devices = [_device(0, used = 0.02, total = 45.0)]
+ hw._overlay_system_wide_vram(devices)
+ assert devices[0]["vram_used_gb"] == 0.02
+
+
+def test_visible_utilization_delegates_gating_to_the_overlay(monkeypatch):
+ # The call site no longer pre-checks masks; the overlay gates itself, so a physical payload always reaches it.
+ monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "2,3")
+ monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
+ monkeypatch.setattr(hw, "IS_ROCM", True)
+ monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
+ monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None)
+ monkeypatch.setattr(
+ hw,
+ "_get_parent_visible_gpu_spec",
+ lambda: {"raw": "1", "numeric_ids": [1], "supports_explicit_gpu_ids": True},
+ )
+ monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [1])
+ monkeypatch.setattr(
+ hw,
+ "_torch_get_per_device_info",
+ lambda ids: [{"index": 1, "visible_ordinal": 0, "used_gb": 0.02, "total_gb": 8.0}],
+ )
+ # Real overlay + gating: the layered mask must leave torch's figures.
+ monkeypatch.setattr(hw.platform, "system", lambda: "Linux")
+ monkeypatch.setattr(hw, "_rocm_kfd_gpu_pci_ids", lambda: [_pci(0), _pci(1)])
+ monkeypatch.setattr(hw, "_rocm_linux_sysfs_vram_by_pci_gb", lambda: {_pci(1): (30.0, 8.0)})
+ result = hw.get_visible_gpu_utilization()
+ assert result["index_kind"] == "physical"
+ assert result["devices"][0]["vram_used_gb"] == 0.02 # untouched
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 9fef53e65e..3d312d4b01 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -734,6 +734,141 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
return None, None
+# 0x1002. NVIDIA's open kernel module also registers KFD nodes (vendor_id 0x10DE);
+# a non-AMD node is not a HIP device and must never take an ordinal.
+_AMD_PCI_VENDOR_ID = 4098
+
+
+def _rocm_kfd_gpu_pci_ids() -> list[str]:
+ """PCI addresses of the GPUs ROCm enumerates, in HIP device order.
+
+ Reads /sys/class/kfd/kfd/topology/nodes//properties, the topology ROCm
+ itself enumerates from: AMD GPU nodes (simd_count > 0 excludes CPUs,
+ vendor_id == AMD excludes NVIDIA) in node-id order are HIP's device order, so
+ position N is ROCm physical device N. Unlike DRM sysfs, an amdgpu adapter HIP
+ cannot enumerate has no node here, so it never consumes an ordinal.
+
+ Returns [] (disabling the overlay) when KFD is absent, and FAILS CLOSED the
+ same way on any unreadable node or an AMD node with no location_id: dropping
+ one would shift every later ordinal and let a similar-capacity GPU pass the
+ total-size guard while showing another card's usage.
+
+ location_id is the kernel's (bus << 8) | devfn; domain is separate.
+ """
+ nodes: list[tuple[int, str]] = []
+ try:
+ node_dirs = glob.glob("/sys/class/kfd/kfd/topology/nodes/*")
+ except Exception:
+ return []
+ for node_dir in node_dirs:
+ m = re.fullmatch(r".*/(\d+)", node_dir)
+ if m is None:
+ continue
+ props: dict[str, int] = {}
+ try:
+ with open(os.path.join(node_dir, "properties")) as f:
+ for line in f:
+ parts = line.split()
+ if len(parts) == 2:
+ try:
+ props[parts[0]] = int(parts[1])
+ except ValueError:
+ continue
+ except OSError:
+ return [] # unreadable node could be a GPU: fail closed, don't shift
+ if props.get("simd_count", 0) <= 0:
+ continue # CPU node, not a GPU
+ if props.get("vendor_id") != _AMD_PCI_VENDOR_ID:
+ continue # non-AMD GPU node (NVIDIA open driver): not a HIP device
+ location_id = props.get("location_id")
+ if location_id is None:
+ return [] # an AMD GPU we cannot place: fail closed for the whole map
+ domain = props.get("domain", 0)
+ bus = (location_id >> 8) & 0xFF
+ devfn = location_id & 0xFF
+ bdf = f"{domain:04x}:{bus:02x}:{(devfn >> 3) & 0x1F:02x}.{devfn & 0x7}"
+ nodes.append((int(m.group(1)), bdf))
+ nodes.sort(key = lambda n: n[0])
+ return [bdf for _node_id, bdf in nodes]
+
+
+def _rocm_linux_amdgpu_cards() -> list[tuple[str, int, str]]:
+ """The amdgpu-bound DRM cards in PCI order: ``(pci_bdf, card_no, device_dir)``.
+
+ Membership is by the BOUND DRIVER, not the VRAM sysfs files: an AMD device
+ with incomplete sysfs support (some APUs expose no mem_info_vram_*) still
+ consumes a ROCm ordinal, and dropping it would shift every later card down.
+ PCI order is HIP's default enumeration order, so list position is the ROCm
+ ordinal; card_no is a stable tiebreak when the BDF cannot be resolved.
+
+ NOTE this is a superset of the ROCm-visible set (a HIP-unsupported amdgpu
+ adapter appears too), so callers must check the counts agree before assuming
+ a 1:1 mapping onto torch devices.
+ """
+ if platform.system() != "Linux":
+ return []
+ amd_cards: list[tuple[str, int, str]] = []
+ try:
+ for card_path in glob.glob("/sys/class/drm/card*"):
+ # Match card exactly so connector nodes (card0-DP-1) are skipped.
+ m = re.fullmatch(r".*/card(\d+)", card_path)
+ if m is None:
+ continue
+ dev_dir = os.path.join(card_path, "device")
+ try:
+ driver = os.path.basename(os.path.realpath(os.path.join(dev_dir, "driver")))
+ except OSError:
+ continue
+ if driver != "amdgpu":
+ continue # foreign adapter: not a ROCm device, takes no ordinal
+ try:
+ bdf = os.path.basename(os.path.realpath(dev_dir))
+ except OSError:
+ bdf = ""
+ amd_cards.append((bdf, int(m.group(1)), dev_dir))
+ except Exception:
+ return []
+ amd_cards.sort(key = lambda c: (c[0], c[1]))
+ return amd_cards
+
+
+def _rocm_linux_sysfs_vram_by_pci_gb() -> dict[str, tuple[float, float]]:
+ """System-wide AMD VRAM via Linux DRM sysfs, keyed by the card's PCI address.
+
+ Reads each card's mem_info_vram_{used,total} (kernel-updated across all
+ processes) so every GPU gets its own figure, unlike _rocm_linux_sysfs_vram_gb
+ which sums the host. Keyed by PCI address, not an ordinal, so the caller can
+ join it to _rocm_kfd_gpu_pci_ids() by identity: DRM card numbers include
+ foreign adapters and this set includes cards HIP does not enumerate, so any
+ ordinal from this list alone can be shifted relative to ROCm's. A card with
+ missing/unreadable/zero-total figures simply has no entry. Empty off Linux.
+ """
+ if platform.system() != "Linux":
+ return {}
+
+ try:
+ by_pci: dict[str, tuple[float, float]] = {}
+ for bdf, _card_no, dev_dir in _rocm_linux_amdgpu_cards():
+ if not bdf:
+ continue
+ try:
+ with open(os.path.join(dev_dir, "mem_info_vram_used")) as f:
+ used_bytes = int(f.read().strip())
+ with open(os.path.join(dev_dir, "mem_info_vram_total")) as f:
+ total_bytes = int(f.read().strip())
+ except (OSError, ValueError):
+ continue
+ if total_bytes <= 0:
+ continue
+ by_pci[bdf.lower()] = (
+ round(used_bytes / (1024**3), 2),
+ round(total_bytes / (1024**3), 2),
+ )
+ return by_pci
+ except Exception:
+ return {}
+
+
# ── Windows AMD/ROCm per-adapter VRAM (issue #7072) ──────────────────────────
# amd-smi is disabled and hipMemGetInfo reports free==total, so read used from the
# per-LUID "GPU Adapter Memory" perf counters and take each total from torch, so
@@ -1222,6 +1357,75 @@ def _reconcile_primary_rocm_unified_memory(
_apply_unified_memory_correction(utilization, torch_devices[0])
+def _rocm_visibility_mask_active() -> bool:
+ """True when any ROCm/CUDA visibility variable filters the device set."""
+ for var in (
+ "HIP_VISIBLE_DEVICES",
+ "ROCR_VISIBLE_DEVICES",
+ "CUDA_VISIBLE_DEVICES",
+ "GPU_DEVICE_ORDINAL",
+ ):
+ value = os.environ.get(var)
+ if value and value.strip():
+ return True
+ return False
+
+
+def _overlay_system_wide_vram(devices: list[Dict[str, Any]]) -> None:
+ """Replace process-local torch VRAM with system-wide Linux ROCm figures.
+
+ The torch fallback is process-local, so a model served by the separate
+ llama-server process reads as ~0 used even with the GPU full (#7072). DRM
+ sysfs gives per-card figures the kernel updates across all processes. Sources
+ are matched by the device's PHYSICAL index (never list position), and only
+ when NO visibility mask is active and the device count equals the host GPU
+ count; under any mask the index is not a verifiable host ordinal, so torch's
+ figures are kept. Best-effort, in place: a device with no matching card, or a
+ unified-memory APU whose sysfs total is below torch's GTT-backed total, keeps
+ torch's (mirrors _apply_unified_memory_correction).
+
+ Windows is intentionally not overlaid: its per-adapter perf counters cannot be
+ mapped to ROCm ordinals and miss WDDM shared memory, so the multi-GPU view
+ keeps torch there rather than risk misattributing another adapter's usage.
+ """
+ if not devices or platform.system() != "Linux":
+ return
+ # Match by PCI identity, never list position: index N in KFD topology is ROCm
+ # physical device N and carries its PCI address, which DRM sysfs keys on too.
+ # The two gates below verify ``index`` really is a host-physical ordinal
+ # (torch exposes no PCI id to check directly):
+ # * No visibility mask -- any mask makes ``index`` container/ROCR-relative
+ # rather than a host ordinal.
+ # * Device count == host GPU count -- rules out a device-cgroup container
+ # that sets no env var yet compacts torch's indices from zero.
+ pci_by_ordinal = _rocm_kfd_gpu_pci_ids()
+ if not pci_by_ordinal:
+ return
+ if _rocm_visibility_mask_active() or len(devices) != len(pci_by_ordinal):
+ return
+ vram_by_pci = _rocm_linux_sysfs_vram_by_pci_gb()
+ for dev in devices:
+ index = dev.get("index")
+ if not isinstance(index, int) or not (0 <= index < len(pci_by_ordinal)):
+ continue
+ entry = vram_by_pci.get(pci_by_ordinal[index].lower())
+ if entry is None:
+ continue
+ used, total = entry
+ dev_total = dev.get("vram_total_gb") or 0.0
+ # Overlay only a device that maps 1:1 to the whole card: torch total must
+ # match sysfs total within ~10%. A mismatch either way means a different
+ # memory scope -- a unified-memory APU (sysfs sees only the dedicated
+ # slice, torch the GTT pool) or a partitioned MI300 (sysfs reports the
+ # whole card, dwarfing a partition) -- and overlaying would misstate free
+ # VRAM (a partition would look like it has the whole card free).
+ if dev_total <= 0 or abs(total - dev_total) > 0.1 * dev_total:
+ continue
+ dev["vram_used_gb"] = used
+ dev["vram_total_gb"] = total
+ dev["vram_utilization_pct"] = round((used / total) * 100, 1) if total > 0 else None
+
+
def get_visible_gpu_utilization() -> Dict[str, Any]:
device = get_device()
@@ -1317,6 +1521,12 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
"power_utilization_pct": None,
}
)
+ if IS_ROCM and index_kind == "physical":
+ # Swap process-local torch VRAM for system-wide sysfs so a model
+ # held by the separate llama-server process shows up (#7072).
+ # Physical-index only: a relative index (UUID/MIG mask) is not a
+ # host GPU id. The overlay verifies the rest itself.
+ _overlay_system_wide_vram(devices)
return {
"available": True,
"backend": _backend_label(device),
From aa49c0710e7632558fceea03ff4b64a9c27ab009 Mon Sep 17 00:00:00 2001
From: Hakan Baysal
Date: Wed, 22 Jul 2026 14:05:08 +0300
Subject: [PATCH 067/255] studio: classify embedding models from the HF cache
and honor offline mode (#7218)
* studio: classify embedding models from the HF cache and honor offline mode
is_embedding_model() went straight to huggingface_hub.model_info() for any repo
id, so in offline mode (no DNS, or HF_HUB_OFFLINE set) selecting an
already-downloaded model hung on network retries that could never succeed and
training/export never started (#6817).
Check the local HF cache first: a sentence-transformers repo carries
modules.json in its snapshot (the same marker used for local paths), so a cached
model is classified with no network call. When HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE
is set, anything not positively an embedding model returns False without a
network call instead of retrying a doomed request. Online, uncached lookups still
fall through to model_info(), so tag-only embedding models (feature-extraction)
are unaffected.
Adds _embedding_marker_in_hf_cache() over the existing _iter_hf_cache_snapshots.
* studio: judge the active cached revision, harden the cache probe, stop stub leaks
Three review fixes on the cache-first embedding detection:
1. Prefer the revision refs/main resolves to. The HF cache keeps snapshots of
older revisions, so an any-snapshot scan could classify a repo by a stale
revision -- e.g. a repo that used to be a sentence-transformers model would
short-circuit even the online lookup. When refs/main is recorded, only its
snapshot is consulted; the newest-first scan remains the fallback for caches
with no ref.
2. Keep the cache probe inside the detection error boundary. The snapshot
iterator stat()s entries and could raise if a cached model is deleted
concurrently, propagating a 500 out of the config/check-embedding routes.
_embedding_marker_in_hf_cache now catches everything and reads as
not-cached, so callers keep their normal Hub/offline fallback.
3. Stub loggers/structlog in the test only when the real modules are absent
(try-import, mirroring test_windows_gpu_detection_mock), so collecting this
file first can no longer shadow the real packages for later tests in the
same pytest process.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: treat a missing active-ref snapshot as a cache miss, don't cache offline misses
Two review fixes on the cache-first embedding detection:
1. When refs/main is recorded but points at a commit whose snapshot dir is
absent (partial download / cache pruning), the recorded ref is still
authoritative: return None (cache miss) instead of falling through to scan
older snapshots, which could report a stale historical revision's
modules.json as the active one -- the same stale-cache class this helper
avoids.
2. Do not cache the offline negative. When HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE
is set and the repo is not positively an ST model from modules.json,
is_embedding_model stored False under the (model_name, hf_token) key shared
with online lookups; after the env var cleared in the same process, a
tag-only (feature-extraction) embedder returned the cached False and never
reached model_info(). The offline negative is now returned without caching.
* studio: defer online embedding detection to the Hub, re-probe offline
The local modules.json marker short-circuited is_embedding_model() even
online, so a repo that dropped (or added) the marker since it was cached
was judged by its stale local revision instead of the current remote one.
Online now treats model_info() as authoritative and uses the cache marker
only as an uncached fallback when the Hub is unreachable, so a transient
failure never poisons the memo. Offline re-probes the marker on every call
without consulting or populating the memo, so a model downloaded later in
the session (or a cached online negative that predates the download) is
detected. _embedding_marker_in_hf_cache() now treats an unreadable refs/main
(a non-FileNotFoundError OSError) as a cache miss rather than scanning stale
history -- only a genuinely missing ref enables the fallback scan.
* studio: harden offline embedding detection against empty refs, offline flips, and cache casing
- _embedding_marker_in_hf_cache: an existing-but-empty/whitespace refs/main
(a partial write or in-progress truncate-and-rewrite) now reads as a cache
miss (None) instead of falling through to scan stale snapshots; only a
genuinely missing ref enables the historical scan.
- is_embedding_model: while offline, retain a positive already confirmed online
this session (model_info only ever memoizes Hub-derived results), so
_hf_offline_if_dns_dead() flipping the process to offline mid-load can't
downgrade a verified tag-only embedder to False. Cached negatives are still
bypassed and re-probed.
- resolve_cached_repo_casing + settings route: persist the embedding model in
the casing its local HF cache dir uses. Validation accepts a case-insensitive
cache hit, but an offline SentenceTransformer load resolves the cache by exact
case, so storing the requested spelling (baai/bge-m3 vs models--BAAI--bge-m3)
made the model fail to load on a case-sensitive filesystem.
* studio: reuse the exact-match-first case resolver and preserve the default
Replace the ad-hoc resolve_cached_repo_casing with the existing
resolve_cached_repo_id_case, which already prefers the exact-case cache dir
before any case variant and tie-breaks variants deterministically -- so an
exact requested id is never rewritten to a differently cased directory just
because iterdir() happened to yield it first.
Skip the normalization entirely when the submitted model equals the default:
rewriting its casing would make set_rag_embedding_model()'s exact-string
default comparison treat it as a custom override, pinning it so later changes
to the configured default stop taking effect.
* studio: don't let a stale cache marker mask a permanent Hub error
is_embedding_model's Hub-failure fallback consulted the local modules.json
marker for ANY model_info() exception, so a permanent error -- a deleted repo,
a gated repo without credentials, or a typo that matches stale cache casing --
could pass online validation on a stale marker instead of returning the
documented 409, and the persisted model could then fail when the loader
refreshes from the Hub. Classify permanent Hub errors (RepositoryNotFound,
GatedRepo, RevisionNotFound, EntryNotFound) as False, matching the nearby
GGUF/vision detectors, and reserve the cache fallback for transient/5xx failures.
* studio: honor TRANSFORMERS_OFFLINE in the embedding preflight, skip casing for local paths
- The embedding-model save reached the offline-aware is_embedding_model() only
after two preflight helpers made direct huggingface_hub calls that honor just
HF_HUB_OFFLINE: _st_module_subdirs() downloads modules.json and the security
scan fetches Hub metadata twice. In a TRANSFORMERS_OFFLINE-only session those
blocked on network timeouts before the offline return, so saving an already
cached model stalled. Both now consult a canonical hf_env_offline() helper --
the download passes local_files_only, and the metadata-only security scan
short-circuits to its documented fail-open instead of burning both timeouts.
- Skip cache-casing normalization for local paths: a relative directory such as
"org/model" is loaded from disk, so rewriting it to a case-insensitive HF
cache collision ("Org/model") would stop resolving to that directory and be
read as a Hub repo id instead.
* studio: never skip the security scan on TRANSFORMERS_OFFLINE alone
The previous commit skipped the Hub security scan whenever either offline flag
was set, but huggingface_hub honors only HF_HUB_OFFLINE: under a
TRANSFORMERS_OFFLINE-only session the later SentenceTransformer load still
reaches the network, so the scan was being skipped while the repo's pickle could
still be downloaded and deserialized -- waving through exactly what
_guard_model_security exists to block.
Split the flags: hf_hub_offline() (HF_HUB_OFFLINE, the only one that actually
prevents a fetch) gates the security short-circuit, while hf_env_offline()
(either flag, the user's intent) is used only where local-only behavior is
forced explicitly. The SentenceTransformer load now passes local_files_only from
that intent, so TRANSFORMERS_OFFLINE genuinely stops the loader fetching instead
of merely being assumed to.
* studio: short-circuit the security preflight under either offline flag
With the loader now pinned to the local cache by local_files_only =
hf_env_offline(), a TRANSFORMERS_OFFLINE-only session can no longer fetch
anything -- yet the preflight still fell through to two model_info() attempts on
10s and 20s timeouts, stalling every save and load of an already-cached embedder
for half a minute before failing open anyway.
Skip the metadata-only scan whenever either flag is set. The scan's job is to
stop a poisoned pickle being downloaded and deserialized, and nothing can be
downloaded under that predicate; the residual case -- a model cached BEFORE it
was flagged -- is the same fail-open this function has always documented for an
unavailable scan, and is exactly what HF_HUB_OFFLINE already did.
That safety argument depends on every loader behind the gate honoring the same
predicate, so it is pinned as a test invariant instead of a comment: removing
local_files_only from the SentenceTransformer construction now fails the suite.
Drops the short-lived hf_hub_offline() helper, which no longer has a caller.
* studio: scope the offline scan bypass to callers that load local-only
The previous commit put the offline short-circuit inside _fetch_security_status,
which is the malware gate shared by every loader -- so TRANSFORMERS_OFFLINE=1
disabled it for all of them, while only the RAG embedder had been changed to
pass local_files_only. MLX inference (core/inference/worker.py -> FastMLXModel
.from_pretrained), training and export call from_pretrained with no local-only
argument, and huggingface_hub ignores that flag, so those paths could still
fetch and deserialize an unscanned model with the gate switched off.
The bypass is now an explicit local_only_load argument, defaulting to False, and
only the two RAG embedding callers -- whose loader is pinned to the local cache
by the same predicate -- opt in. Tests pin both halves: the shared gate must
still scan under either offline flag by default, and no other caller may pass
local_only_load without constraining its loader.
* studio: capture offline state once, and probe the ST cache root
Two holes in the offline embedding path:
- _get() read hf_env_offline() twice: once inside _guard_model_security and
again for local_files_only. _hf_offline_if_dns_dead() mutates the process-wide
offline vars and restores them on exit, so a concurrent load could see True in
the guard -- skipping the Hub malware scan -- and False by the time the
constructor ran, fetching and deserializing the unscanned repo and breaking
the very invariant that licenses the bypass. The value is now read once in
_get() and passed to both; _guard_model_security takes it as an argument
instead of re-deriving it.
- The cache probe searched only HF_HUB_CACHE. SentenceTransformer downloads into
SENTENCE_TRANSFORMERS_HOME when that is set, using the same
models--org--name/snapshots layout under a different root, so a model fully
present there looked uncached and was rejected with a 409 offline even though
the local-only loader could load it. Snapshot lookup now covers both roots.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: probe the cache the ST loader actually uses, and require it be loadable
Adding SENTENCE_TRANSFORMERS_HOME to the shared snapshot iterator was too broad
in one direction and too narrow in another:
- _get() builds SentenceTransformer with no cache_folder, so with ST_HOME set it
searches THAT root only, never the Hub cache. Probing the union let offline
validation pass on a repo cached only in the Hub cache, after which the loader
looked in ST_HOME and failed. The Sentence-Transformers probe now resolves to
exactly one root: ST_HOME when set, the Hub cache otherwise.
- The shared iterator is also used by the GGUF detectors, whose downloads go
through hf_hub_download with no cache_dir and therefore really do use the Hub
cache. It is back to Hub-cache-only so detection cannot pick a snapshot the
GGUF load will not find.
- Casing normalization ran through resolve_cached_repo_id_case, which scans the
Hub cache, so with ST_HOME set the requested spelling was persisted unchanged
and the exact-case offline load missed the differently cased directory that
detection had just accepted. It now resolves against the same roots detection
uses, exact match first.
- A snapshot carrying only modules.json no longer counts as cached: the online
security preflight downloads that single file itself, and a partial download
leaves it behind, so validation passed for a snapshot with no weights and the
first RAG load then failed. A hit now requires the marker plus a config and at
least one weight file.
* studio: thread the captured offline state into the module probe, fix the gate shard
- _st_module_subdirs() re-read the process env for its local_files_only. With
_hf_offline_if_dns_dead() flipping those vars from another thread, a load that
captured local_only=False could still force this probe local-only, get () back
because modules.json is not cached, and leave the scan with NO module load
roots -- a Hub-flagged pickle under 0_Transformer/ would then pass as an
unreferenced nested artifact while the loader fetched and deserialized it. It
now takes the captured predicate as an argument, and the settings route reads
the state once and uses that single value for both the probe and the scan.
- Skip ST-cache casing on the llama-server backend. Nothing there loads through
SentenceTransformer: the embedder derives a GGUF companion from the saved
spelling and fetches it from the HUB cache, so normalizing to an ST_HOME
spelling would point it at a repo _hf_gguf_backend_error() never validated
(BAAI/bge-m3-GGUF instead of the checked baai/bge-m3-GGUF).
- Fix the security-gate shard, which the signature change had broken: the direct
_guard_model_security / _st_module_subdirs callers now pass the new argument
(they were raising TypeError before reaching any assertion), and the casing
tests patch utils.models.resolve_st_cached_repo_id_case, which the route
actually calls, instead of the Hub-only resolver it no longer uses -- those
patches were being silently ignored.
* studio: accept only torch-loadable weights in the offline ST probe; fix re-export lint
_snapshot_is_loadable_st_model accepted a cached snapshot whose only weights
were .onnx (or .pt), but the RAG loader builds SentenceTransformer with the
default torch backend, so such a snapshot passed offline validation and then
failed on the first load, the exact validate-then-fail this helper exists to
prevent. Restrict _ST_WEIGHT_SUFFIXES to .safetensors and .bin and add a
regression test for an ONNX-only snapshot.
Also teach scripts/verify_import_hoist.py that names listed in a module-level
__all__ are uses, so the legitimately added resolve_st_cached_repo_id_case
re-export in utils/models/__init__.py no longer trips HOISTED-IMPORT-UNUSED.
Covered by two new self-test cases.
* studio: probe the exact repo dir and revision an offline load resolves
The cache probe modelled the cache loosely rather than modelling what
SentenceTransformer actually does with local_files_only=True:
- It merged snapshots across every case-variant repo dir and then read refs/main
from whichever held the newest one. With both models--baai--bge-m3 and
models--BAAI--bge-m3 present, a complete embedding snapshot in the directory
the loader opens could be judged by a newer partial snapshot in the other,
failing validation for a usable model. It now selects the ONE directory the
loader opens, by the same exact-case-first rule resolve_st_cached_repo_id_case
uses to choose the spelling that gets persisted.
- It fell back to scanning historical snapshots when refs/main was absent. With
local_files_only the default revision is resolved THROUGH that ref, so a
snapshot directory alone is not discoverable: the settings request succeeded
and the loader then failed at first indexing. A missing, empty or unreadable
ref is now a cache miss, and the historical scan is gone.
The tests exercise the real lookup against a built cache tree instead of
patching the snapshot iterator, so they now cover the directory selection and
ref resolution the loader depends on.
* studio: record refs/main in the ONNX-only probe test
The ONNX-only regression test predates the refs/main requirement, so after that
change it returned None (a cache miss for want of a ref) before ever reaching
the weight-format check it exists to make. Recording the ref restores its
intent: the snapshot resolves, and the answer is False because an ONNX export is
not loadable by the RAG loader's default Torch backend.
* studio: recognize base-model weight files and gate the offline positive on a materialized snapshot
_snapshot_is_loadable_st_model matched any .safetensors/.bin by suffix, so a
partial cache carrying only a commonly published non-weight bin such as
training_args.bin (or an adapter-only artifact) passed offline validation and
then failed the local_files_only load at first indexing. Match recognized Torch
base-model weight filenames (model / pytorch_model, including sharded) by name.
is_embedding_model retained an online-confirmed positive offline even when no
files were cached, so a metadata-only /check-embedding result let an uncached
repo be saved and then fail at first indexing. Retain the positive only when the
active revision is materialized locally, which still covers a downloaded tag-only
embedder whose snapshot carries no modules.json.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: require a complete weight set offline and persist embedder verdicts across restarts
Two follow-ups to the offline embedding-model classifier:
- _snapshot_is_loadable_st_model now requires a COMPLETE Torch base-model
weight set in one snapshot directory, not just any single recognized weight
file. A partially downloaded sharded model (model-00001-of-00002 without its
sibling) no longer passes offline validation and then fails at first indexing
under local_files_only. Weight files are grouped by directory and a directory
counts only when it holds a single model.safetensors / pytorch_model.bin or a
full shard set whose indices cover 1..total.
- Online-confirmed embedder verdicts are now recorded under the resolved Studio
home (embedding_verdicts.json). The session memo is lost on exit, so a
downloaded tag-only feature-extraction embedder (snapshot present but no
modules.json) was misclassified as non-embedding the first offline call after
a restart. The offline branch consults this durable allowlist in addition to
the memo, still gated on the active revision being materialized on disk, so an
uncached repo is never trusted. Writes are best-effort and only positive
verdicts are stored.
* studio: require complete weights (with shard index) and resolve default casing offline
Follow-ups to the offline embedding-model classifier from the latest review:
- Trust a recorded embedder verdict (session memo or persisted allowlist) offline
only when the active snapshot carries a COMPLETE, loadable weight set, not merely
that it is materialized. A partial download (config present, weights missing or an
incomplete shard set) makes _embedding_marker_in_hf_cache read False rather than
None, so the previous marker-is-not-None gate wrongly returned True and the
local_files_only load then failed. Split out _snapshot_has_complete_weights (config
plus complete weights, modules.json aside) and _active_snapshot_dir, and gate the
known-embedder positive on the weight set.
- Require a sharded checkpoint's index map (model.safetensors.index.json /
pytorch_model.bin.index.json) in addition to every shard before accepting it:
transformers discovers and wires shards through that index, so a complete shard set
without it fails the local-only load.
- Resolve the embedding model name to its exact cache casing in the RAG loader before
constructing SentenceTransformer. The settings route persists that spelling for a
custom override but deliberately leaves the configured default verbatim, so a
default whose casing differs from the cache dir would miss it and fail offline.
Resolving at load time covers the default too; a no-op for a local path or when
nothing case-matching is cached, and idempotent for an already-normalized override.
Adds regression tests for the partial-snapshot verdict, the missing shard index, and
the loader casing resolution; updates the offline-invariant source assertion to the
resolved-name variable.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: require a tokenizer, case-fold verdict ids, and serialize verdict writes
Three follow-ups to the offline embedding-model classifier from the latest review:
- _snapshot_has_complete_weights now also requires a tokenizer asset. A
SentenceTransformer Transformer module builds an AutoTokenizer, so a snapshot with
a complete weight set but no tokenizer.json / tokenizer_config.json / vocab still
fails the local_files_only load. The check is a permissive union over the common
fast-tokenizer, config, and WordPiece/BPE/SentencePiece assets, so an unusual but
valid layout is not rejected -- only a genuinely tokenizer-less partial download.
- The persisted embedder allowlist is now keyed case-insensitively. model_info() is
queried under the requested casing while the settings route saves the cache-resolved
casing, so an exact-string lookup missed the persisted positive after a restart
(baai/model recorded, BAAI/model looked up) and a loadable tag-only embedder was
rejected. Both persist and lookup case-fold the id.
- _persist_embedder serializes its read-modify-write under a lock and writes through a
per-thread temp file, so concurrent confirmations of different embedders no longer
drop each other's entry or collide on the temp path. Cross-process writers stay
best-effort (os.replace is atomic; a dropped verdict is only an optimization miss a
later online re-confirmation heals).
Adds regression tests for the missing-tokenizer reject, alternate tokenizer assets,
cross-casing verdict match, and concurrent verdict writes; updates the snapshot test
helpers to materialize a tokenizer alongside config and weights.
* studio: tighten comments in the offline embedding-model classifier
Comment-only pass over the PR's changed files. Collapse the long block
comments and docstrings around is_embedding_model, the cache-snapshot and
weight-completeness helpers, the embedder-verdict persistence, the offline
security gate, and the offline/casing tests to short one- or two-line forms.
Preserve the rationale (issue #6817, the local_files_only invariant, the
casing and weight-gate reasons) in far fewer words. No code changes.
* studio: drop redundant comments in the offline embedding-model classifier
Second comment-reduction pass over the offline embedding-model cache work:
delete comments and trailing notes that restate the adjacent code or an
assertion, and trim the remaining docstrings and rationale comments to their
load-bearing invariants. Comments and docstrings only; no code changes.
* studio: pin embedder verdicts to a revision, canonicalize default aliases
- A persisted verdict recorded that the Hub tagged ONE revision an embedder, but
was stored per repo. Once refs/main advanced to a complete but non-embedding
Transformer snapshot, the offline path still returned True: the settings route
accepted the updated model without force and RAG could silently load it as an
embedder. Verdicts now carry the commit they were confirmed at and are trusted
only while the active revision matches. One confirmed before the repo was
cached has no revision to compare, so the first revision observed afterwards is
pinned then -- which is what lets a later advance be caught. The persisted file
gains a {id: commit} form and still reads the previous list format.
- tokenizer_config.json no longer counts as a tokenizer asset. It only DESCRIBES
a tokenizer, so a snapshot with config, weights and just that file passed
validation and then failed AutoTokenizer.from_pretrained(local_files_only=True)
at first indexing for common BERT/GPT-style models.
- A casing-only alias of the default is canonicalized to the default up front.
Repo ids are case-insensitive but every gate here compares exact strings, so
saving "Unsloth/bge-m3" against a default of "unsloth/bge-m3" ran the
verification and scan for a custom model and then persisted an override --
after which later changes to the configured default stopped applying.
- verify_import_hoist.py replays __all__ assignments in order instead of unioning
them. Only the final value exports anything, so a later plain "=" that drops a
name must leave its import counted as unused; "+=" still extends, and an
unreadable rebind keeps the earlier names rather than flagging real re-exports.
* studio: validate the real ST load root, and pin verdicts to the Hub revision
Four ways the offline probe still disagreed with what the loader does:
- Verdicts were pinned to the LOCAL refs/main, but model_info() describes the
current HUB revision. With a stale cache the two differ, so an older snapshot
nobody verified was allowlisted. The pin is now info.sha, taken from the
ModelInfo that produced the positive. A verdict carrying no revision (a legacy
entry) is no longer trusted at all -- trusting it meant pinning whatever
happened to be cached, which is the same bug; the next online check re-records
it properly.
- config, tokenizer and weights had to exist somewhere in the snapshot, not
together. modules.json can send SentenceTransformer at 0_Transformer/, which is
loaded FROM that directory, so a cache with the config at the root and only
0_Transformer/model.safetensors passed and then failed the local-only load.
Each directory is now checked as a complete load root, which covers both the
plain HF layout and the ST module layout.
- vocab.json and merges.txt counted independently, but BPE needs the pair unless
a serialized tokenizer.json is present, so half a pair validated and then
failed AutoTokenizer.from_pretrained(local_files_only=True).
- A slashless short name like all-MiniLM-L6-v2 is a supported ST alias that the
loader resolves through the sentence-transformers/ organization, so its
snapshot is cached under that full id. Probing only the bare name reported a
miss and 409'd a model that was cached and loadable; the bare id is still tried
first, matching the loader's own order.
* studio: fail closed for an offline security scan instead of failing open
A local_only (offline) load cannot fetch Hugging Face's malware scan, and the previous
behaviour skipped the scan and failed OPEN, so a cached repo with a poisoned pickle weight
could deserialize under SentenceTransformer(local_files_only=True). Evaluate it fail-CLOSED
against the cached files instead: block a base-model pickle weight the load would deserialize
(pytorch_model.bin and its shards, in a directory with no safetensors alternative) and allow a
pickle-free (safetensors / gguf are inert) cache. A cached pickle model must be reloaded online
once to be scanned, or shipped as safetensors. Nothing cached is not a security event.
_fetch_security_status no longer needs the local_only_load skip (the offline branch is handled
in evaluate_file_security). Adds a regression test covering the safetensors-allow and
pickle-block paths with no Hub call.
* studio: only suppress an offline pickle when a loadable safetensors weight exists
The offline security gate treated any .safetensors in a directory as covering a
pickle weight, so a cache with pytorch_model.bin beside a bare adapter_model.safetensors
(or an orphan shard with no index) passed the fail-closed check even though
from_pretrained still selects and deserializes the pickle. Require a genuinely loadable
safetensors weight -- an unsharded base file or a complete indexed shard set -- before
treating the pickle as covered.
Also make the import-hoist analyzer preserve uncertainty when __all__ is extended by a
value it cannot read statically (__all__ += dynamic()), matching how it already handles
an unreadable rebind, so a dynamically-supplied re-export is not flagged HOISTED-IMPORT-UNUSED.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: scope the offline pickle scan to load paths; reset __all__ opacity on rebind
Address three review follow-ups on the offline security gate and the import-hoist analyzer:
- The offline pickle scan walked the whole snapshot, so a stray pickle in a non-load
subdirectory (archive/, nemo/) that SentenceTransformer never deserializes was blocked.
Scope it to real from_pretrained load roots -- the snapshot root, or a subdir that holds
its own config.json -- matching the online scan's load-path scoping.
- _collect_dunder_all kept a sticky opaque flag: a readable replacing assignment after an
unreadable extend (__all__ += dynamic(); __all__ = []) still credited every import, so a
genuinely unused hoist went unreported. A replacing assignment now resets opacity.
- A bare __all__: list[str] annotation has no runtime value; it was treated as an unreadable
assignment and marked the export set opaque. Skip annotation-only declarations.
* studio: recase slashless ST aliases and accept a pinned embedder after a transient failure
Two offline-detection gaps on well-formed input:
- resolve_st_cached_repo_id_case bailed on every slashless name, so a differently-cased
short alias (all-minilm-l6-v2) validated case-insensitively but was loaded verbatim; the
SentenceTransformer loader rewrites it to sentence-transformers/all-minilm-l6-v2 and looks
it up case-sensitively, missing the canonical sentence-transformers/all-MiniLM-L6-v2 cache
dir. Resolve through _st_cache_repo_dir, which follows the same org alias, and hand back the
on-disk casing.
- On a transient (non-permanent) Hub failure, is_embedding_model only accepted a cached
modules.json marker, so a downloaded tag-only embedder (no modules.json) with a verdict
pinned to the active revision was rejected even though the offline branch accepts the
identical cache. Mirror the offline branch's pinned-verdict acceptance.
* studio: scan modules.json-declared module roots in the offline pickle gate
The offline pickle scan treated only the snapshot root and config.json-bearing subdirs as
load roots, so a pickle in a non-Transformer SentenceTransformer module directory that has no
config.json (e.g. a 0_WordEmbeddings/ module: wordembedding_config.json + pytorch_model.bin)
was skipped even though the loader deserializes it. Parse modules.json (and thread through
load_subdirs) to treat every declared module directory as a load root, so such a pickle is
scanned and fail-closed offline.
* studio: classify cached non-Transformer SentenceTransformer models offline
_snapshot_has_complete_weights recognized only a Transformer-shaped load root (config +
tokenizer + weights co-located), so a fully-cached model built from a non-Transformer module
(0_WordEmbeddings uses wordembedding_config.json + embedding weights and its own tokenizer, no
HF config.json; BoW keeps its vocab in config.json) was classified non-embedding offline and
the settings endpoint returned 409.
Add _snapshot_modules_all_loadable, which parses modules.json and accepts a snapshot when every
declared module's path directory carries the files that module class's own load() reads (a
Transformer/root module still needs the full HF load root; a WordEmbeddings module needs its
config plus a complete weight set; other modules need their *_config.json), and at least one
embedding-producing module is present. It is OR-ed after the Transformer check, so it only ever
accepts more and cannot regress the existing path or reject a pruned cache.
* studio: scan PEFT adapter pickle weights in the offline security gate
from_pretrained auto-detects an adapter_config.json in the load root and deserializes the
adapter weights on top of the base model, so adapter_model.bin is a separate pickle RCE vector
that a safetensors base weight does not cover. The offline scan matched only base-model pickle
names, so an offline local-only load with safetensors base weights plus a cached
adapter_model.bin was allowed despite the live adapter pickle. Scan adapter pickles too, scoped
to a load root where adapter_config.json is present and no adapter_model.safetensors exists.
* studio: require weights for Dense/CNN/LSTM SentenceTransformer modules offline
_module_dir_is_loadable accepted a Dense, CNN, or LSTM module dir with only its config, but
those modules' load() hard-load model.safetensors else pytorch_model.bin (verified against
sentence-transformers source: no fallback, raises if neither exists) -- exactly like
WordEmbeddings. A cache with such a module's config but no weights would validate and then
fail the local_files_only load. Require a complete weight set for every weighted module, not
just WordEmbeddings.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: scan root-index subdir pickle shards offline; handle __all__.append/.extend
- The offline pickle scan followed only load-root directories, so a shard mapped by a root
pytorch_model.bin.index.json into a non-root subdirectory was skipped even though
from_pretrained follows the index weight_map and deserializes it (a layout an attacker can
craft to evade the scanner). Read the local index and scan its referenced pickle shards,
covered by a loadable base safetensors at the index root -- mirroring the online scan.
- The import-hoist analyzer ignored __all__.append("X") / __all__.extend([...]) runtime
re-export mutators, so an import added solely for one tripped HOISTED-IMPORT-UNUSED. Read
their string args like +=, and treat any other __all__ method call as opaque.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: classify StaticEmbedding offline, require WordEmbeddings tokenizer, bound model_info
- A StaticEmbedding module (e.g. sentence-transformers/static-retrieval-mrl-en-v1's
0_StaticEmbedding/) holds tokenizer.json + weights and NO config, so the config-gated
non-Transformer path 409'd it offline. Recognize it by what StaticEmbedding.load() reads: a
tokenizer.json plus a complete Torch weight set.
- WordEmbeddings.load() rebuilds its tokenizer via the configured tokenizer_class.load() from the
module dir, so a WordEmbeddings module now also requires a tokenizer artifact
(whitespacetokenizer_config.json / phrasetokenizer_config.json, or a shared HF tokenizer asset),
not just its config + weights.
- With neither offline env var set, an unbounded model_info() could hang on connect/DNS retries
for networkless users (the #6817 symptom). Bound it with a 15s timeout so a dead network fails
fast and the existing transient-failure cache fallback resolves a cached model, while a
reachable Hub still wins. (Documented caveat: a stalled DNS getaddrinfo may exceed this.)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Resolve indexed safetensors shards relative to their index
_safetensors_index_complete compared shard basenames against the flat
set of files in the index directory, so an index whose weight_map names
shards in a subdirectory was treated as incomplete whenever a legacy
pytorch_model.bin sat beside it. That falsely blocked a snapshot whose
pickle weights are fully covered by a complete, loadable safetensors
shard set. Resolve each shard path relative to the index directory
instead, and add a regression test for the subdir-mapped shard case.
* Restrict offline weight-completeness check to declared load roots
_snapshot_has_complete_weights scanned every directory in a snapshot and
accepted it when ANY directory was a complete Transformer load root. When
modules.json is present a SentenceTransformer load only opens the declared
module paths, so a snapshot whose declared modules are incomplete but which
happens to contain an unrelated complete directory was accepted offline and
then failed at the first local_files_only load. Restrict the candidate
directories to the roots a load actually opens: the snapshot root plus each
modules.json module path. For a well-formed snapshot the verdict is
unchanged; only a complete directory at an undeclared path no longer vouches
for an otherwise-incomplete snapshot.
* Scan SentenceTransformer Router child module weights offline
A Router (legacy Asym) snapshot declares its child sub-modules only in
router_config.json, not the top-level modules.json, and Router.load()
deserializes each child's weights from its own subdir. A config.json-less
child such as query_0_WordEmbeddings (wordembedding_config.json plus a
pickle pytorch_model.bin loaded via torch.load) was therefore neither a
modules.json-declared load root nor a config.json-bearing dir, so the
offline gate skipped its pickle even though the loader deserializes it.
Parse router_config.json at each load root and treat every declared child
subdir as a load root (bounded BFS, so nested routers are covered), so
those child pickles are scanned. Add Router regression tests: a pickle
child blocks, a safetensors child is allowed, and a Router in a declared
subfolder is followed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Do not treat an unreferenced config subdir as an offline load root
The offline pickle gate skipped a directory only when it was neither a
declared load root nor held a config.json. Because _st_load_roots already
resolves every real load root (snapshot root, modules.json / load_subdirs
dirs, Router children), the config.json fallback only ever promoted an
UNREFERENCED subdir -- a nested checkpoint-500/ or archive/ that ships its
own config.json + pytorch_model.bin -- to a load root. from_pretrained
never descends into such a subdir and the online scan ignores the same
unindexed pickle, so offline mode wrongly blocked a model the loader reads
from a clean safetensors root. Scope the pickle to directory in roots
only, and add a regression test (a stray checkpoint-500/ no longer blocks;
a modules.json-declared module dir still does).
* Classify a root Router (Asym) model as loadable offline
_module_dir_is_loadable applied Transformer root requirements (config +
tokenizer + weights) to every root module, so a Router saved at the
snapshot root -- which carries only modules.json + router_config.json and
loads its weights from child subdirs -- was classified not loadable
offline, and is_embedding_model missed a cached Router embedder. Dispatch
on the module class before the root Transformer fallback: a Router/Asym
dir is loadable when router_config.json parses and every declared child
subdir is loadable (validated recursively through _module_dir_is_loadable,
so nested routers and every child type are covered) with at least one
embedding-producing child. This also tightens a non-root Router, which
previously validated on the mere presence of router_config.json without
checking its children. Add Router regression tests (root and declared
subfolder, complete and incomplete-child).
* Require every declared module before accepting an offline cache
_snapshot_is_loadable_st_model returned has_complete_weights OR
modules_all_loadable, so a complete 0_Transformer short-circuited the or
and vouched for the whole snapshot even when a declared sibling module was
missing its serialized weights; SentenceTransformer builds every module in
modules.json, so that snapshot passed offline validation and then failed
the local-only load. When modules.json declares a non-empty list it is now
authoritative (modules_all_loadable validates every declared module);
has_complete_weights stays the fallback only for an empty/non-list
modules.json (the plain from_pretrained root). Also add the weight-bearing
modules whose load() hard-loads via load_torch_weights and previously fell
to the config-only path -- LayerNorm, WeightedLayerPooling, SparseAutoEncoder
-- to _ST_WEIGHTED_MODULE_NAMES, with source citations and the deliberate
exclusions (Pooling/Normalize/BoW/WordWeights read no weights on load).
Add parametrized regression tests over LayerNorm/WeightedLayerPooling/Dense
(a weightless sibling rejects, a complete sibling accepts).
* Reject self-referential Router children instead of recursing forever
_router_dir_is_loadable validates each router_config.json child through
_module_dir_is_loadable, which re-enters _router_dir_is_loadable for a
Router child. A malformed types entry naming the router's own directory
(a key of ".", which normalizes to the same dir) made that recursion
never descend, so it looped until RecursionError -- breaking the
documented never-raises contract and turning a crafted/corrupted cached
model into a 500 from is_embedding_model instead of a graceful
unverifiable result. A real child reference is a subdir and always
resolves deeper, so reject any child whose resolved path is the router
dir itself. Add a regression test (a router_config naming "." as a
Router child returns False without raising).
* Treat a destructuring __all__ assignment as opaque
_collect_dunder_all detected __all__ only as a direct ast.Name assignment
target, so a binding through a destructuring target (__all__, meta = [...],
v -> an ast.Tuple) was skipped entirely, leaving an empty, non-opaque
export set. A newly hoisted import re-exported only through that assignment
was then falsely flagged HOISTED-IMPORT-UNUSED. Its value cannot be mapped
statically, so mark the export set opaque when __all__ is reached only
through a destructuring / item / attr target, matching how the collector
already handles other unreadable __all__ forms. Add a self-test case.
* Canonicalize declared module paths before scoping the offline pickle gate
A repo could declare a traversing module path such as 0/../evil in
modules.json (or a router_config child), which SentenceTransformer resolves
to evil/ and deserializes evil/pytorch_model.bin. _st_load_roots recorded
the raw snap/"0/../evil", which never equals the snap/evil that rglob
yields, so the offline pickle gate skipped that directory and a malicious
repo slipped a pickle past the newly added gate. Add _canonical_load_dir
to collapse ./ and ../ components lexically and reject an upward escape,
and route the modules.json paths, load_subdirs and router children through
it so the gate scopes the same normalized directory the loader opens. Add
regression tests for a traversing modules.json path and router child.
* Close offline embedding-classification completeness gaps
Five real offline misclassifications, each a false negative (the #6817 hang
recurs) or false positive (accepted then 409s at the local_files_only load).
Dispatch _module_dir_is_loadable on the module class before the root
Transformer fallback. A module with save_in_root=True (every InputModule:
WordEmbeddings, StaticEmbedding, SparseStaticEmbedding, Transformer, Router)
is saved at the snapshot root, so a root WordEmbeddings was wrongly held to
Transformer requirements (an HF tokenizer it never writes) and classified not
loadable.
CLIPModel is Transformer-shaped: CLIPModel.load() reads AutoModel weights plus
AutoProcessor, so a config-only CLIP dir must not validate.
SparseStaticEmbedding needs a tokenizer plus either idf.json or a complete
torch weight set (conditionally weight-bearing); a config alone is not enough.
A present but empty or malformed modules.json is not loadable and does not fall
back to a root Transformer: with modules.json present the loader never takes
the plain-Transformer path (base/model.py _load_config_modules). The tag-only
no-modules.json embedder is classified separately via
_snapshot_has_complete_weights.
Validate a sharded weight index against its weight_map (every mapped shard
present, resolved relative to the index dir) instead of trusting the index
file's mere existence, mirroring the security-side check.
Add regression tests for all five.
* Close case-folding and online-traversal holes in the offline pickle gate
Two gate bypasses where the security scan credited or scoped a path
differently from what the loader actually resolves:
The safetensors credit was case-folded. _cached_pickle_weight_files lowercases
every filename, and the loadable-safetensors and adapter checks tested those
folded keys against the exact-lowercase names. On a case-sensitive filesystem
(Linux, the Studio default) a crafted repo shipping Model.SafeTensors plus a
malicious pytorch_model.bin makes transformers and sentence-transformers miss
the exact-name model.safetensors and deserialize the pickle, while the gate
credited an inert safetensors and did not block. Credit safetensors
case-sensitively against real filenames, and drop pytorch_model.safetensors
from the credit set (transformers loads only model.safetensors, never that
name). Pickle matching stays case-insensitive (over-blocking a mis-cased
pickle the loader would not load is the safe direction).
The online scan did not canonicalize traversing paths while the offline gate
did. A repo-controlled modules.json path (threaded into the online scan via
the RAG guard) or a weight_map shard entry like 0/../evil / ../evil was
compared verbatim, so a flagged evil/pytorch_model.bin never matched and
evaded the online scan though the loader resolves and deserializes it.
Canonicalize the repo-controlled load-subdir prefixes and weight_map shards
the same way the offline gate does, so offline and online agree.
Add regression tests for both bypasses.
* Treat a conditional __all__ mutation as opaque in the import-hoist linter
_collect_dunder_all replayed only top-level module statements, so an __all__
assignment or mutation inside a module-level if / try / for / while / with /
match (or a deeper scope) was ignored, leaving the export set understated. A
newly hoisted import re-exported only through such a conditional __all__ was
then falsely flagged HOISTED-IMPORT-UNUSED, blocking a valid change. A
conditional value cannot be replayed statically, so mark the export set opaque
when __all__ is bound or mutated anywhere other than a top-level statement.
Add a self-test case.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scope Router child sub-modules as load roots in the online embedding scan
The RAG embedding security guard unions the SentenceTransformer module dirs
from modules.json into the load roots it scopes for the Hub scan, so a flagged
pickle directly under a Transformer module blocks. A Router (legacy Asym)
module declares its child sub-modules only in router_config.json, not in
modules.json, and Router.load() deserializes each child from its own subdir.
The online scan therefore dropped a flagged child pickle (for example
query_0_WordEmbeddings/pytorch_model.bin) as an unreferenced nested shard while
the loader still deserialized it, the counterpart to the offline gate which
already expands router children via _router_child_dirs.
_st_module_subdirs now reads router_config.json for any Router-typed module and
adds each declared child (joined onto the module path, canonicalized so a
traversing entry is dropped) to the load roots. The config is read only for a
Router-typed module, so a plain embedder pays no extra fetch, and every failure
path still returns () so the guard never bricks the embedder.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Allow a recorded-clean pickle embedder to load offline
The offline embedding security gate is fail-closed: with no network to reach
Hugging Face's scan, a cached pickle weight cannot be verified, so it is blocked
and a model the user already downloaded and used online will not load offline.
This adds a persistent cache of clean Hub verdicts so that exact content can load
offline, without weakening the gate for an unknown or never-scanned pickle.
When an embedding repo is loaded online and HF's scan returns a completed clean
verdict, the load roots are hashed and recorded under the scanned commit as an
exact map of snapshot-relative pickle name to sha256, in a per-user JSON store at
studio_root()/security/embedding_scan_verdicts.json (atomic write, 0600, thread
and cross-process locked, 30-day TTL). Offline, a cached pickle model loads only
when the active cached commit and every load-root pickle's sha256 match the
recorded verdict; a missing record, moved commit, changed or added pickle,
expired record, or any error keeps blocking. Online loads always re-query the
Hub and an authoritative unsafe verdict deletes any stale record, so a
now-flagged commit cannot keep loading on an old clean record.
The store binds repo id, full commit, and a per-file sha256 map so a locally
swapped pickle at the same commit, a branch advance, or an added load-relevant
pickle is detected. A same-user attacker who can rewrite the model cache or the
store is outside the enforceable boundary and this is documented; the sha256 is
computed just before load, so a narrow verify-to-load window remains, and a Hub
scanner false negative is recorded faithfully (safetensors stays the stronger
defense).
Recording is triggered post-load in the RAG embedder because the settings route
only validates and the pre-load guard runs before the constructor downloads;
recording is skipped when the loaded commit differs from the scanned commit. The
blocked-pickle enumerator now returns snapshot-relative Paths so two module dirs
that ship the same pickle basename are hashed and reported distinctly.
* Harden the embedding verdict cache against review findings
Tighten the offline verdict cache and its enumeration so every uncertain or
malformed input fails closed and the recorded hashes always match the files the
loader reads:
- Hash every case-colliding pickle in a load root, not one representative. On a
case-sensitive filesystem pytorch_model.bin and PYTORCH_MODEL.BIN are distinct
files; keying by lowered name dropped one and could hash a decoy instead of the
loader's target. The enumerator now returns every variant Path.
- Only persist a clean verdict for a COMPLETED, entirely-benign scan. Require
scansDone to be the boolean True (not a truthy string), filesWithIssues to be a
well-formed list, and every flagged file to be a definitively-safe level; a
pending, error, unknown, or malformed entry no longer records as clean. The
online block decision is unchanged.
- Fail closed when the offline cache cannot be inspected: an rglob error now
propagates and blocks instead of reading as pickle-free, and a snapshot that
errors on resolution (vs a clean not-cached) blocks. The offline guard also
raises instead of returning when its own inspection throws, so the constructor
never deserializes an unverified cached pickle.
- Expand online Router children recursively (bounded BFS with a seen set),
mirroring the offline load-root expansion, so a flagged grandchild pickle is
scoped online and cannot be recorded clean.
- Reject absolute and drive/UNC declared paths in the load-root canonicalizers;
the loader would resolve them outside the snapshot, so collapsing them to an
in-snapshot relative dir scoped the wrong place.
- Pin verdict recording to the scanned commit's snapshot and take the offline
verify commit from the snapshot directory name, removing a second refs/main read
and the skew it allowed.
- Drop the now-unused pickle-name wrapper.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten offline embedding classification and the pickle gate
Close a set of offline edge cases where validation accepted a cache the
local_files_only load then rejects, and one gate bypass:
- Credit a sharded model.safetensors.index.json for a pickle sibling only at a
from_pretrained root. A non-Transformer SentenceTransformer module (Dense,
WordEmbeddings, StaticEmbedding) loads via Module.load_torch_weights, which
reads model.safetensors then pytorch_model.bin and never the index, so a sharded
safetensors index in such a module dir must not vouch for its pytorch_model.bin.
- Stop counting pytorch_model.safetensors as loadable in the offline classifier:
the loader probes model.safetensors (then its index) or pytorch_model.bin, never
pytorch_model.safetensors, matching the gate that already treats it as a decoy.
- Treat a present but unreadable weight index as incomplete: transformers opens
and parses any present index, so a malformed one or one without a weight_map
fails the load rather than falling back to filename-numbered shards.
- Require the CLIP image-processor config (preprocessor_config.json) for a CLIP
module: CLIPModel.load builds a CLIPProcessor that needs it, so a tokenizer
alone is not enough.
- Require a SparseStaticEmbedding config to actually select idf.json (a path
ending .json) or ship loadable weights; a bare idf.json the config does not name
falls through to load_torch_weights and raises.
- Do not use the tag-only recorded-verdict fallback when modules.json is present:
with the file present the loader takes the modules.json path, so a present but
empty or malformed manifest must not be validated as a plain root Transformer.
- Import-hoist linter: only a module-level conditional mutation or a function that
declares global __all__ makes the export set opaque; a __all__ bound as a local
in a nested function or class no longer masks a genuinely unused hoisted import.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scope Router-child pickles to their deepest load root and gate the ST offline kwarg
The online scan stripped the first matching load-subdir prefix from a flagged file, so a
nested Router child pickle (0_Router/query_0_WordEmbeddings/pytorch_model.bin) matched the
parent 0_Router root, looked like an unreferenced nested shard, and slipped the gate even
though Router.load() deserializes that child directly. Match the deepest (longest) load
subdir instead, so the child becomes root-level under its own load root and blocks.
pyproject sets no lower bound on sentence-transformers and the local_files_only constructor
arg is absent on older releases, so always forwarding it broke every embedder warm on those
installs. Pass it only for an offline load; an online warm never forwards it and works as
before, while the offline capability still requires a version that supports it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Reject snapshot-escaping shard paths and credit Transformer submodule safetensors
The offline pickle enumerator joined a weight-index weight_map value straight to the load
root and followed it, so a repo-controlled index mapping "../.." into a sibling snapshot
made an offline from_pretrained deserialize an out-of-snapshot pickle, and an online load
would then hash and record that external file as the scanned commit's clean content. Reject
any shard path that escapes the snapshot root and fail closed, mirroring the canonical-root
check the online shard scan already applies.
A complete model.safetensors.index.json was credited over a sibling pickle only at the
snapshot root, but a Transformer module subdirectory (0_Transformer/) is loaded via
AutoModel.from_pretrained, which honors that shard set and never reads the pickle. Credit the
sharded index for Transformer-typed modules declared in modules.json so a cached model that
ships both a sharded safetensors checkpoint and an unused PyTorch checkpoint is no longer
falsely blocked offline. Non-Transformer modules (Dense, WordEmbeddings, StaticEmbedding) read
a flat weight with no index and keep their pickle blocked.
Limit the import-hoist verifier's global __all__ scan to the declaring function's own scope so
a nested inner-scope local __all__ no longer marks the module export set opaque and mask an
unused hoisted import.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scope Router children against the snapshot and mirror the ST alias rewrite
Router.load resolves each child at Path(subfolder, model_id) relative to the Router dir, so a
nested 1_Router with a "../evil" child points at evil/ inside the snapshot and the loader
deserializes evil/pytorch_model.bin. The offline enumerator canonicalized the child against
the Router dir alone and dropped anything with "..", so that pickle was never scanned and the
gate reported the cache pickle-free. Canonicalize router children against the snapshot,
retaining in-snapshot siblings as load roots and failing closed on a child that escapes the
snapshot itself, matching the online scan which already joins the prefix before normalizing.
The security gate resolved a slashless model id by probing the bare cache dir first, but the
SentenceTransformer constructor rewrites a non-basic slashless name to sentence-transformers/
and loads THAT snapshot (only the basic ORIGINAL_TRANSFORMER_MODELS load bare). With
both models-- and models--sentence-transformers-- cached, the gate inspected the
bare dir while the loader read the namespaced one, so a pickle there bypassed the local-only
gate. Mirror the constructor: try the namespaced candidate first for non-basic slashless names.
Add the same not (snapshot / modules.json).is_file() guard to the transient-Hub-failure
tag-only fallback that the offline branch already carries, so a cache whose present manifest is
empty or malformed is no longer reported as a loadable embedder.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten root shard credit, module-path escapes, and weight-set probe order
Credit the sharded safetensors index at the snapshot ROOT only when the root is actually loaded
through an AutoModel/from_pretrained path. A modules.json root module of a non-Transformer type
(StaticEmbedding / WordEmbeddings / Dense) loads via load_torch_weights, which reads
pytorch_model.bin and ignores the index, so crediting a root shard index there suppressed a live
root pickle and let the offline gate report the cache pickle-free.
Recognize the Transformer subclasses CLIPModel and MLMTransformer as index-honoring load roots
(they load via from_pretrained), so a sharded-safetensors CLIP/MLM submodule with a legacy
pytorch_model.bin sibling is no longer falsely blocked offline. Mirrors the classifier dispatch.
Fail closed on an absolute or snapshot-escaping modules.json module path (or load_subdirs entry)
instead of silently dropping it: SentenceTransformer resolves such a path outside the snapshot and
would deserialize an external pytorch_model.bin the gate cannot scan.
On the classifier side, walk the weight set in the exact from_pretrained probe order
(model.safetensors, its index, pytorch_model.bin, its index) so a pickle behind a malformed
safetensors index is no longer accepted as complete, and restrict shard names to the loader-probed
stem/ext pairs so a decoy model-*.bin / pytorch_model-*.safetensors set is not treated as loadable.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Restore scripts/verify_import_hoist.py to main
The offline embedding cache fix does not depend on the __all__ scope
handling that had accumulated in this linter, so revert the file to its
main version and keep the PR focused on the feature. The feature modules
still pass the existing import hoist check unchanged.
* Reuse a shared HF cache skeleton in the offline classification tests
Extract _mk_repo and _activate helpers for the repeated snapshot cache
setup that every per-type builder duplicated, and fold the two
StaticEmbedding missing-asset cases into one parametrized test. Same 125
collected items, all still passing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Reclassify embedding models from the cache on every offline call
is_embedding_model consulted its process memo before the offline branch, so an
online lookup that memoized True from tags (without caching any weights) was returned
unchanged once the session went offline -- the studio flips HF_HUB_OFFLINE in-process
on a dead DNS, and the ungated check-embedding route can populate the memo. Settings
would then accept a repo the offline loader cannot open. Run the offline
cache-marker reclassification ahead of the memo and never record it, so an offline
verdict always reflects the local cache and a later cache materialization is not
masked by a stale negative. Add regression tests.
* Tighten comments on the offline embedding path
Condense the offline-embedding helper docstrings and inline comments added in
this PR to fewer, clearer lines, keeping the non-obvious security and offline
rationale. Comments and docstrings only; no code change.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen
---
studio/backend/core/rag/embeddings.py | 67 +-
studio/backend/routes/settings.py | 77 ++-
.../test_embedding_model_security_gate.py | 50 ++
.../tests/test_offline_embedding_minimal.py | 583 ++++++++++++++++++
studio/backend/utils/models/model_config.py | 35 +-
.../backend/utils/security/file_security.py | 123 ++++
studio/backend/utils/utils.py | 105 ++++
7 files changed, 1002 insertions(+), 38 deletions(-)
create mode 100644 studio/backend/tests/test_offline_embedding_minimal.py
diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py
index 15be7f1249..0c743e4ea4 100644
--- a/studio/backend/core/rag/embeddings.py
+++ b/studio/backend/core/rag/embeddings.py
@@ -22,6 +22,7 @@ from typing import Callable
from utils.hardware.hardware import DeviceType, get_device
from utils.transformers_dtype import dtype_kwargs
+from utils.utils import hf_env_offline
from . import config
@@ -119,30 +120,55 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
return ()
-def _guard_model_security(name: str) -> None:
+def _guard_model_security(name: str, local_only: bool = False) -> None:
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
/settings gate (a name can also arrive via env/default); local paths and unreachable
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
+
+ ``local_only`` (offline) inspects the local cache; subdir probes are skipped (they'd hit the
+ network and hang, and the offline gate walks the whole snapshot anyway).
"""
try:
from utils.security import evaluate_file_security, security_load_subdirs
token = _ambient_hf_token()
- # Union the audio-model load roots with the ST module dirs so a flagged pickle
- # directly under a Transformer module dir (0_Transformer/) blocks instead of
- # passing as an unreferenced nested shard.
- load_subdirs = tuple(
- dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
- )
- blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
+ if local_only:
+ load_subdirs = ()
+ else:
+ # Union audio-model load roots with ST module dirs so a flagged pickle under a
+ # Transformer module dir blocks instead of passing as an unreferenced nested shard.
+ load_subdirs = tuple(
+ dict.fromkeys(
+ (*security_load_subdirs(name, token), *_st_module_subdirs(name, token))
+ )
+ )
+ blocked = evaluate_file_security(
+ name, hf_token = token, load_subdirs = load_subdirs, local_only_load = local_only
+ ).blocked
except Exception:
return
if blocked:
- raise UnsafeEmbeddingModelError(
- f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
- "scan; refusing to load. Set a different RAG embedding model."
+ reason = (
+ "has cached pickle weights that cannot be security-scanned offline and no "
+ "safetensors alternative"
+ if local_only
+ else "is flagged as unsafe by Hugging Face's security scan"
)
+ raise UnsafeEmbeddingModelError(
+ f"Embedding model {name!r} {reason}; refusing to load. "
+ "Set a different RAG embedding model."
+ )
+
+
+def _st_accepts_local_files_only(st_cls) -> bool:
+ """Whether this SentenceTransformer version accepts local_files_only; passing it to an
+ older constructor raises, so gate on the signature."""
+ try:
+ import inspect
+ return "local_files_only" in inspect.signature(st_cls.__init__).parameters
+ except Exception:
+ return False
def _get(model_name: str | None = None):
@@ -150,6 +176,9 @@ def _get(model_name: str | None = None):
for a ~1.5x speedup at negligible accuracy loss."""
global _model, _name
name = model_name or config.effective_embedding_model()
+ # Capture offline state once so the gate and the load agree (no window where the gate is
+ # skipped as offline but the constructor then reaches the network).
+ local_only = hf_env_offline()
with _lock:
if _model is None or _name != name:
_install_torchao_stub_once()
@@ -157,8 +186,20 @@ def _get(model_name: str | None = None):
device = _device()
logger.info("loading embedding model %s on %s", name, device)
- _guard_model_security(name)
- _model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16"))
+ _guard_model_security(name, local_only)
+ st_kwargs = dict(device = device, model_kwargs = dtype_kwargs("float16"))
+ load_target = name
+ if local_only:
+ from utils.utils import hf_cache_snapshot_dir
+ snapshot = hf_cache_snapshot_dir(name)
+ if snapshot is not None:
+ # Load from the local snapshot dir: a local path never touches the Hub, so
+ # this is offline-safe on ANY sentence-transformers version (even ones
+ # predating local_files_only).
+ load_target = str(snapshot)
+ elif _st_accepts_local_files_only(SentenceTransformer):
+ st_kwargs["local_files_only"] = True
+ _model = SentenceTransformer(load_target, **st_kwargs)
_name = name
return _model
diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py
index 17e64df918..f36c8870e3 100644
--- a/studio/backend/routes/settings.py
+++ b/studio/backend/routes/settings.py
@@ -416,6 +416,11 @@ def update_embedding_model(
log = logger,
) from exc
hf_token = (payload.hf_token or "").strip() or None
+ from utils.utils import hf_env_offline
+
+ # Offline, both the Hub malware scan and the is-embedding check are unreachable and degrade
+ # to the local cache below; capture the state once.
+ local_only_load = hf_env_offline()
# The env/default model needs no verification; saving it is a no-op override.
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
# what the backend loads, and HF metadata cannot verify a local path.
@@ -439,26 +444,41 @@ def update_embedding_model(
# Fall back to the loader's own token so a gated/private repo is actually scanned
# (a token-less scan fails open for exactly the repo that would still load).
scan_token = hf_token or _ambient_hf_token()
- # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
- # one blocks instead of passing as an unreferenced nested shard.
- load_subdirs = tuple(
- dict.fromkeys(
- (
- *security_load_subdirs(model, scan_token),
- *_st_module_subdirs(model, scan_token),
+ # Offline: subdir probes would hit the network and hang; the offline gate walks the
+ # whole cached snapshot, so no load-subdir hints are needed.
+ if local_only_load:
+ load_subdirs = ()
+ else:
+ # Include ST module dirs (0_Transformer/) so a flagged pickle directly under one
+ # blocks instead of passing as an unreferenced nested shard.
+ load_subdirs = tuple(
+ dict.fromkeys(
+ (
+ *security_load_subdirs(model, scan_token),
+ *_st_module_subdirs(model, scan_token),
+ )
)
)
- )
- if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
+ if evaluate_file_security(
+ model,
+ hf_token = scan_token,
+ load_subdirs = load_subdirs,
+ local_only_load = local_only_load,
+ ).blocked:
# 403, not 409: the client routes every 409 into the forceable "save anyway"
# flow, but this block is a hard, non-forceable security refusal.
- raise HTTPException(
- status_code = 403,
+ if local_only_load:
+ detail = (
+ f"{model!r} has cached pickle weights that cannot be security-scanned "
+ "offline and no safetensors alternative, so it cannot be used as the "
+ "embedding model. Re-download it with safetensors weights while online."
+ )
+ else:
detail = (
f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
"cannot be used as the embedding model."
- ),
- )
+ )
+ raise HTTPException(status_code = 403, detail = detail)
if model != default_embedding_model() and not payload.force and not is_local_gguf:
from core.rag import config as rag_config
@@ -468,15 +488,28 @@ def update_embedding_model(
# which would wrongly 409 a valid online GGUF embedder.
gguf_named = _llama_backend_active() and rag_config._names_gguf(model)
if not gguf_named and not is_embedding_model(model, hf_token = hf_token):
- raise HTTPException(
- status_code = 409,
- detail = (
- f"Could not verify {model!r} as an embedding model on "
- "Hugging Face (it may be the wrong model type, gated, or "
- "you may be offline)."
- ),
- )
- gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token)
+ # Offline, is_embedding_model can only confirm the ST layout (modules.json); a
+ # transformers-native embedder (e.g. gte-modernbert) is unverifiable without Hub
+ # metadata. If already cached and loadable, accept it rather than raising a 409 that
+ # online would not (ST can load any cached encoder). Uncached -> 409.
+ from utils.utils import hf_cache_snapshot_is_loadable
+
+ # Require a genuinely loadable cache (config + weights), not just a resolved refs/main,
+ # so a metadata-only partial cache still gets the forceable 409.
+ offline_cached = local_only_load and hf_cache_snapshot_is_loadable(model)
+ if not offline_cached:
+ raise HTTPException(
+ status_code = 409,
+ detail = (
+ f"Could not verify {model!r} as an embedding model on "
+ "Hugging Face (it may be the wrong model type, gated, or "
+ "you may be offline)."
+ ),
+ )
+ # The Hub GGUF probe (list_repo_files) can hang offline; skip it. Local check stays.
+ gguf_error = _local_gguf_backend_error(model)
+ if gguf_error is None and not local_only_load:
+ gguf_error = _hf_gguf_backend_error(model, hf_token)
if gguf_error:
raise HTTPException(status_code = 409, detail = gguf_error)
set_rag_embedding_model(model)
diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py
index b3fa98b604..a6c18bd8de 100644
--- a/studio/backend/tests/test_embedding_model_security_gate.py
+++ b/studio/backend/tests/test_embedding_model_security_gate.py
@@ -106,6 +106,56 @@ def test_hard_block_uses_non_forceable_status(client, monkeypatch):
assert unverified.status_code == 409
+def test_offline_cached_non_st_model_is_accepted(client, monkeypatch):
+ # Offline, a cached transformers-native embedder (no modules.json) is unverifiable via HF
+ # metadata, but ST can load any cached encoder, so accept it (no 409).
+ c, saved = client
+ monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ import utils.models as _models
+ import utils.utils as _uu
+
+ monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
+ monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: True)
+ r = c.put("/embedding-model", json = {"embedding_model": "acme/gte-modernbert"})
+ assert r.status_code == 200
+ assert saved.get("model") == "acme/gte-modernbert"
+
+
+def test_offline_partial_or_uncached_model_still_409(client, monkeypatch):
+ # Offline but not loadable (uncached or metadata-only partial cache): keep the forceable
+ # 409, since the cache-only load would fail anyway.
+ c, _saved = client
+ monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ import utils.models as _models
+ import utils.utils as _uu
+
+ monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
+ monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: False)
+ r = c.put("/embedding-model", json = {"embedding_model": "acme/uncached-embedder"})
+ assert r.status_code == 409
+
+
+def test_offline_skips_remote_gguf_probe(client, monkeypatch):
+ # Offline + llama backend: the remote GGUF probe (list_repo_files) must be skipped so a
+ # dead-DNS session cannot hang.
+ c, _saved = client
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
+ monkeypatch.setattr(settings, "_local_gguf_backend_error", lambda model: None)
+
+ def _boom(*a, **k):
+ raise AssertionError("hit the network for the GGUF probe")
+
+ monkeypatch.setattr(settings, "_hf_gguf_backend_error", _boom)
+ import utils.models as _models
+
+ monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: True)
+ r = c.put("/embedding-model", json = {"embedding_model": "acme/embedder"})
+ assert r.status_code == 200
+
+
def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
# On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
# pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.
diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py
new file mode 100644
index 0000000000..8862e231e5
--- /dev/null
+++ b/studio/backend/tests/test_offline_embedding_minimal.py
@@ -0,0 +1,583 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Offline RAG embedding-model handling (issue #6817).
+
+Offline the studio must never call the Hub (a DNS-dead session hangs on retries). Using a fake
+HF cache under a temp HF_HUB_CACHE, assert that offline: is_embedding_model classifies from the
+cached modules.json without the Hub; the file-security gate fails CLOSED on an unscanned pickle
+weight with no safetensors alternative and allows an inert cache; the embedder threads
+local_files_only into the load. Online behavior is unchanged (bounded timeout + cache fallback).
+"""
+
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from utils.security import evaluate_file_security
+from utils.utils import (
+ hf_cache_snapshot_dir,
+ hf_cache_snapshot_is_loadable,
+ hf_env_offline,
+ st_repo_id_candidates,
+)
+
+# Minimal sentence-transformers modules.json (the marker the gate keys on).
+MODULES_JSON = (
+ '[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"}]'
+)
+
+
+def _modules_json(*paths):
+ """modules.json listing one Transformer module per path (a load root)."""
+ import json
+ return json.dumps(
+ [
+ {
+ "idx": i,
+ "name": str(i),
+ "path": p,
+ "type": "sentence_transformers.models.Transformer",
+ }
+ for i, p in enumerate(paths)
+ ]
+ )
+
+
+_COMMIT = "0123456789abcdef0123456789abcdef01234567"
+
+
+def _make_cache(
+ root,
+ repo_id,
+ files,
+ commit = _COMMIT,
+):
+ """Build a canonical HF-cache snapshot (refs/main + snapshots//) for repo_id under
+ root from {relpath: contents}; returns the snapshot dir."""
+ from huggingface_hub.file_download import repo_folder_name
+
+ repo_dir = Path(root) / repo_folder_name(repo_id = repo_id, repo_type = "model")
+ (repo_dir / "refs").mkdir(parents = True, exist_ok = True)
+ (repo_dir / "refs" / "main").write_text(commit)
+ snapshot = repo_dir / "snapshots" / commit
+ snapshot.mkdir(parents = True, exist_ok = True)
+ for rel, contents in files.items():
+ path = snapshot / rel
+ path.parent.mkdir(parents = True, exist_ok = True)
+ path.write_text(contents)
+ return snapshot
+
+
+def _no_network():
+ """Patch model_info to fail loudly if any offline path reaches the network."""
+ return patch("huggingface_hub.model_info", side_effect = AssertionError("hit the network"))
+
+
+def _is_embedding_model(*args, **kwargs):
+ from utils.models.model_config import is_embedding_model
+ return is_embedding_model(*args, **kwargs)
+
+
+@pytest.fixture
+def hf_cache(tmp_path, monkeypatch):
+ """Point the HF cache at a fresh temp dir."""
+ root = tmp_path / "hub"
+ root.mkdir()
+ monkeypatch.setenv("HF_HOME", str(tmp_path))
+ monkeypatch.setenv("HF_HUB_CACHE", str(root))
+ return root
+
+
+@pytest.fixture(autouse = True)
+def _clean_env(monkeypatch):
+ """Start each test online with an empty detection cache; offline tests opt in."""
+ monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
+ monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
+ from utils.models import model_config as mc
+
+ mc._embedding_detection_cache.clear()
+ yield
+ mc._embedding_detection_cache.clear()
+
+
+# ── hf_env_offline ───────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On "])
+def test_hf_env_offline_true(monkeypatch, value):
+ monkeypatch.setenv("HF_HUB_OFFLINE", value)
+ assert hf_env_offline() is True
+
+
+@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
+def test_hf_env_offline_false(monkeypatch, value):
+ monkeypatch.setenv("HF_HUB_OFFLINE", value)
+ assert hf_env_offline() is False
+
+
+def test_hf_env_offline_honors_transformers_flag(monkeypatch):
+ monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
+ assert hf_env_offline() is True
+
+
+def test_hf_env_offline_default_false():
+ assert hf_env_offline() is False
+
+
+# ── st_repo_id_candidates ────────────────────────────────────────
+
+
+def test_candidates_slashless_adds_st_alias():
+ assert st_repo_id_candidates("all-MiniLM-L6-v2") == [
+ "all-MiniLM-L6-v2",
+ "sentence-transformers/all-MiniLM-L6-v2",
+ ]
+
+
+def test_candidates_with_org_is_verbatim():
+ assert st_repo_id_candidates("org/model") == ["org/model"]
+
+
+def test_candidates_empty_name():
+ assert st_repo_id_candidates(" ") == []
+
+
+# ── hf_cache_snapshot_dir ────────────────────────────────────────
+
+
+def test_snapshot_dir_resolves_active_commit(hf_cache):
+ snapshot = _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON})
+ assert hf_cache_snapshot_dir("org/emb") == snapshot
+
+
+def test_snapshot_dir_none_when_uncached(hf_cache):
+ assert hf_cache_snapshot_dir("org/missing") is None
+
+
+def test_snapshot_dir_uses_st_alias_for_slashless(hf_cache):
+ snapshot = _make_cache(
+ hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON}
+ )
+ assert hf_cache_snapshot_dir("all-MiniLM-L6-v2") == snapshot
+
+
+def test_snapshot_dir_none_when_snapshot_missing(hf_cache):
+ from huggingface_hub.file_download import repo_folder_name
+
+ repo_dir = hf_cache / repo_folder_name(repo_id = "org/broken", repo_type = "model")
+ (repo_dir / "refs").mkdir(parents = True)
+ (repo_dir / "refs" / "main").write_text("deadbeef") # no snapshots/deadbeef dir
+ assert hf_cache_snapshot_dir("org/broken") is None
+
+
+def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch):
+ # An unexpanded $VAR in HF_HUB_CACHE must resolve where the loader looks.
+ real = tmp_path / "hub"
+ real.mkdir()
+ monkeypatch.setenv("MY_HF_CACHE", str(real))
+ monkeypatch.setenv("HF_HUB_CACHE", "$MY_HF_CACHE")
+ monkeypatch.delenv("HF_HOME", raising = False)
+ monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False)
+ snapshot = _make_cache(real, "org/emb", {"modules.json": MODULES_JSON})
+ assert hf_cache_snapshot_dir("org/emb") == snapshot
+
+
+def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch):
+ # ST uses SENTENCE_TRANSFORMERS_HOME as its cache_folder, so the gate must inspect it too.
+ st_home = tmp_path / "st_home"
+ st_home.mkdir()
+ monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
+ monkeypatch.delenv("HF_HUB_CACHE", raising = False)
+ monkeypatch.delenv("HF_HOME", raising = False)
+ snapshot = _make_cache(st_home, "org/emb", {"modules.json": MODULES_JSON})
+ assert hf_cache_snapshot_dir("org/emb") == snapshot
+
+
+def test_snapshot_dir_st_home_is_exclusive(tmp_path, monkeypatch):
+ # With SENTENCE_TRANSFORMERS_HOME set, ST loads only from it, so a model living only under
+ # HF_HUB_CACHE must not be reported.
+ st_home = tmp_path / "st_home"
+ st_home.mkdir()
+ hub = tmp_path / "hub"
+ hub.mkdir()
+ monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
+ monkeypatch.setenv("HF_HUB_CACHE", str(hub))
+ monkeypatch.delenv("HF_HOME", raising = False)
+ _make_cache(hub, "org/emb", {"modules.json": MODULES_JSON}) # only in the HF hub cache
+ assert hf_cache_snapshot_dir("org/emb") is None
+
+
+def test_snapshot_is_loadable_with_config_and_weights(hf_cache):
+ _make_cache(hf_cache, "org/emb", {"config.json": "{}", "model.safetensors": "x"})
+ assert hf_cache_snapshot_is_loadable("org/emb") is True
+
+
+def test_snapshot_is_not_loadable_when_metadata_only(hf_cache):
+ # A partial cache (refs/main resolves but no weights) is not loadable.
+ _make_cache(hf_cache, "org/partial", {"config.json": "{}", "modules.json": MODULES_JSON})
+ assert hf_cache_snapshot_is_loadable("org/partial") is False
+
+
+def test_snapshot_is_not_loadable_when_uncached(hf_cache):
+ assert hf_cache_snapshot_is_loadable("org/missing") is False
+
+
+def test_gate_blocks_pickle_in_sentence_transformers_home(tmp_path, monkeypatch):
+ # A pickle under SENTENCE_TRANSFORMERS_HOME must still fail closed offline.
+ st_home = tmp_path / "st_home"
+ st_home.mkdir()
+ monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
+ monkeypatch.delenv("HF_HUB_CACHE", raising = False)
+ monkeypatch.delenv("HF_HOME", raising = False)
+ _make_cache(st_home, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
+ with _no_network():
+ assert evaluate_file_security("org/pk", local_only_load = True).blocked is True
+
+
+# ── is_embedding_model: offline (no network) ─────────────────────
+
+
+def test_offline_true_for_cached_st_model(hf_cache, monkeypatch):
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON, "config.json": "{}"})
+ with _no_network():
+ assert _is_embedding_model("org/emb") is True
+
+
+def test_offline_false_for_cached_non_st_model(hf_cache, monkeypatch):
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ _make_cache(hf_cache, "org/plain", {"config.json": "{}", "model.safetensors": "x"})
+ with _no_network():
+ assert _is_embedding_model("org/plain") is False
+
+
+def test_offline_false_when_uncached(hf_cache, monkeypatch):
+ monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
+ with _no_network():
+ assert _is_embedding_model("org/missing") is False
+
+
+def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch):
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ _make_cache(hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON})
+ with _no_network():
+ assert _is_embedding_model("all-MiniLM-L6-v2") is True
+
+
+def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch):
+ # An online lookup memoizes True for an UNCACHED repo (tags say embedding, no weights). Once
+ # offline, is_embedding_model must reclassify from the empty cache and return False, not the
+ # stale online True that would make settings accept a repo _get() cannot load.
+ with patch(
+ "huggingface_hub.model_info",
+ side_effect = lambda *a, **k: SimpleNamespace(
+ tags = ["sentence-transformers"], pipeline_tag = None
+ ),
+ ):
+ assert _is_embedding_model("org/uncached-emb") is True # memoized True online
+
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ with _no_network():
+ assert _is_embedding_model("org/uncached-emb") is False # recomputed from empty cache
+
+
+def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch):
+ # Because the offline branch never records a memo, once an uncached repo's snapshot
+ # materializes (another process populates the cache) the next call re-reports True.
+ monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+ with _no_network():
+ assert _is_embedding_model("org/later") is False # uncached
+ _make_cache(hf_cache, "org/later", {"modules.json": MODULES_JSON})
+ assert _is_embedding_model("org/later") is True # cache now present, no stale negative
+
+
+# ── is_embedding_model: online (bounded + fallback) ──────────────
+
+
+def test_online_passes_bounded_timeout(hf_cache):
+ seen = {}
+
+ def _mi(
+ name,
+ token = None,
+ timeout = None,
+ **kw,
+ ):
+ seen["timeout"] = timeout
+ return SimpleNamespace(tags = ["sentence-transformers"], pipeline_tag = None)
+
+ with patch("huggingface_hub.model_info", side_effect = _mi):
+ assert _is_embedding_model("org/emb") is True
+ assert seen["timeout"] == 15.0
+
+
+def test_online_error_falls_back_to_cache_marker(hf_cache):
+ _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON})
+ with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")):
+ assert _is_embedding_model("org/emb") is True
+
+
+def test_online_error_without_cache_returns_false(hf_cache):
+ with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")):
+ assert _is_embedding_model("org/missing") is False
+
+
+# ── evaluate_file_security: offline fail-closed gate ─────────────
+
+
+def _offline_decision(name):
+ return evaluate_file_security(name, local_only_load = True)
+
+
+def test_gate_allows_safetensors_only(hf_cache):
+ _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"})
+ with _no_network():
+ assert _offline_decision("org/st").blocked is False
+
+
+def test_gate_blocks_pickle_without_safetensors(hf_cache):
+ _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
+ with _no_network():
+ decision = _offline_decision("org/pk")
+ assert decision.blocked is True
+ assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files)
+
+
+def test_gate_allows_pickle_with_safetensors_sibling(hf_cache):
+ _make_cache(hf_cache, "org/both", {"pytorch_model.bin": "x", "model.safetensors": "y"})
+ with _no_network():
+ assert _offline_decision("org/both").blocked is False
+
+
+def test_gate_blocks_sharded_pickle(hf_cache):
+ _make_cache(
+ hf_cache,
+ "org/shard",
+ {
+ "pytorch_model-00001-of-00002.bin": "a",
+ "pytorch_model-00002-of-00002.bin": "b",
+ },
+ )
+ with _no_network():
+ assert _offline_decision("org/shard").blocked is True
+
+
+def test_gate_allows_nothing_cached(hf_cache):
+ with _no_network():
+ assert _offline_decision("org/missing").blocked is False
+
+
+def test_gate_allows_gguf_only(hf_cache):
+ _make_cache(hf_cache, "org/gg", {"model.gguf": "x"})
+ with _no_network():
+ assert _offline_decision("org/gg").blocked is False
+
+
+def test_gate_blocks_pickle_in_module_subdir(hf_cache):
+ # 0_Transformer is a module load root (listed in modules.json), so its pickle blocks.
+ _make_cache(
+ hf_cache,
+ "org/mod",
+ {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"},
+ )
+ with _no_network():
+ assert _offline_decision("org/mod").blocked is True
+
+
+def test_gate_allows_pickle_in_subdir_with_safetensors(hf_cache):
+ _make_cache(
+ hf_cache,
+ "org/mod2",
+ {
+ "modules.json": _modules_json("0_Transformer"),
+ "0_Transformer/pytorch_model.bin": "x",
+ "0_Transformer/model.safetensors": "y",
+ },
+ )
+ with _no_network():
+ assert _offline_decision("org/mod2").blocked is False
+
+
+def test_gate_allows_unreferenced_nested_pickle(hf_cache):
+ # A pickle in a dir NOT referenced by modules.json (e.g. nemo/) is never deserialized, so it
+ # must not block the offline load (matches the online gate).
+ _make_cache(
+ hf_cache,
+ "org/aux",
+ {
+ "modules.json": MODULES_JSON, # Transformer at the root only
+ "model.safetensors": "w",
+ "nemo/pytorch_model.bin": "x",
+ },
+ )
+ with _no_network():
+ assert _offline_decision("org/aux").blocked is False
+
+
+def test_gate_blocks_adapter_pickle_without_safetensors(hf_cache):
+ _make_cache(hf_cache, "org/ad", {"config.json": "{}", "adapter_model.bin": "x"})
+ with _no_network():
+ decision = _offline_decision("org/ad")
+ assert decision.blocked is True
+ assert any(u["path"] == "adapter_model.bin" for u in decision.unsafe_files)
+
+
+def test_gate_allows_adapter_pickle_with_adapter_safetensors(hf_cache):
+ _make_cache(hf_cache, "org/ad2", {"adapter_model.bin": "x", "adapter_model.safetensors": "y"})
+ with _no_network():
+ assert _offline_decision("org/ad2").blocked is False
+
+
+def test_gate_blocks_base_pickle_with_only_adapter_safetensors_decoy(hf_cache):
+ # A decoy adapter_model.safetensors must NOT suppress a base pytorch_model.bin (the base
+ # loader would still deserialize the unscanned pickle).
+ _make_cache(hf_cache, "org/decoy", {"pytorch_model.bin": "x", "adapter_model.safetensors": "y"})
+ with _no_network():
+ assert _offline_decision("org/decoy").blocked is True
+
+
+def test_gate_blocks_adapter_pickle_with_only_base_safetensors_decoy(hf_cache):
+ # Symmetric: a base model.safetensors must NOT suppress an adapter_model.bin.
+ _make_cache(hf_cache, "org/decoy2", {"adapter_model.bin": "x", "model.safetensors": "y"})
+ with _no_network():
+ assert _offline_decision("org/decoy2").blocked is True
+
+
+def test_gate_reports_snapshot_relative_path(hf_cache):
+ _make_cache(
+ hf_cache,
+ "org/mod3",
+ {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"},
+ )
+ with _no_network():
+ decision = _offline_decision("org/mod3")
+ assert decision.blocked is True
+ assert any(u["path"] == "0_Transformer/pytorch_model.bin" for u in decision.unsafe_files)
+
+
+# ── evaluate_file_security: online path unchanged ────────────────
+
+
+def test_online_default_blocks_unsafe():
+ status = {
+ "scansDone": True,
+ "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
+ }
+ with patch(
+ "huggingface_hub.model_info",
+ side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status),
+ ):
+ assert evaluate_file_security("org/x").blocked is True
+
+
+def test_online_default_allows_clean():
+ status = {"scansDone": True, "filesWithIssues": []}
+ with patch(
+ "huggingface_hub.model_info",
+ side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status),
+ ):
+ assert evaluate_file_security("org/x").blocked is False
+
+
+# ── embeddings guard + loader ────────────────────────────────────
+
+
+def test_guard_offline_blocks_pickle_only(hf_cache):
+ from core.rag.embeddings import UnsafeEmbeddingModelError, _guard_model_security
+ _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
+ with _no_network():
+ with pytest.raises(UnsafeEmbeddingModelError):
+ _guard_model_security("org/pk", local_only = True)
+
+
+def test_guard_offline_allows_safetensors(hf_cache):
+ from core.rag.embeddings import _guard_model_security
+ _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"})
+ with _no_network():
+ _guard_model_security("org/st", local_only = True) # must not raise
+
+
+def _install_fake_sentence_transformers(monkeypatch, captured):
+ class FakeSentenceTransformer:
+ def __init__(
+ self,
+ name,
+ *,
+ device = None,
+ model_kwargs = None,
+ local_files_only = False,
+ **kw,
+ ):
+ captured["name"] = name
+ captured["device"] = device
+ captured["local_files_only"] = local_files_only
+
+ module = types.ModuleType("sentence_transformers")
+ module.SentenceTransformer = FakeSentenceTransformer
+ monkeypatch.setitem(sys.modules, "sentence_transformers", module)
+
+
+def test_get_offline_loads_from_local_snapshot(hf_cache, monkeypatch):
+ from core.rag import embeddings
+
+ snapshot = _make_cache(
+ hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}
+ )
+ # TRANSFORMERS_OFFLINE only: a cached model loads from its local snapshot dir (a local path,
+ # never the Hub), offline-safe on ANY sentence-transformers version.
+ monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
+ monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
+ monkeypatch.setattr(embeddings, "_model", None, raising = False)
+ monkeypatch.setattr(embeddings, "_name", None, raising = False)
+ monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
+ monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
+ captured = {}
+ _install_fake_sentence_transformers(monkeypatch, captured)
+ with _no_network():
+ embeddings._get("org/st")
+ assert captured["name"] == str(snapshot)
+
+
+def test_get_offline_uncached_uses_local_files_only(tmp_path, monkeypatch):
+ from core.rag import embeddings
+
+ empty = tmp_path / "hub"
+ empty.mkdir()
+ monkeypatch.setenv("HF_HUB_CACHE", str(empty))
+ monkeypatch.delenv("HF_HOME", raising = False)
+ monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False)
+ monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
+ monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
+ monkeypatch.setattr(embeddings, "_model", None, raising = False)
+ monkeypatch.setattr(embeddings, "_name", None, raising = False)
+ monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
+ monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
+ # No cache -> repo-id load forced cache-only (fails fast offline, not a hang).
+ monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None)
+ captured = {}
+ _install_fake_sentence_transformers(monkeypatch, captured)
+ embeddings._get("org/uncached-xyz")
+ assert captured["name"] == "org/uncached-xyz"
+ assert captured["local_files_only"] is True
+
+
+def test_get_online_omits_local_files_only(monkeypatch):
+ from core.rag import embeddings
+
+ monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
+ monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
+ monkeypatch.setattr(embeddings, "_model", None, raising = False)
+ monkeypatch.setattr(embeddings, "_name", None, raising = False)
+ monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
+ monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
+ # Isolate the loader wiring from the online guard's network calls.
+ monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None)
+ captured = {}
+ _install_fake_sentence_transformers(monkeypatch, captured)
+ embeddings._get("org/online")
+ assert captured["local_files_only"] is False
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 821529083d..50a997218f 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -2076,6 +2076,24 @@ def download_gguf_file(
_embedding_detection_cache: Dict[tuple, bool] = {}
+# Bound the Hub lookup so a DNS-dead session fails fast to the cache instead of hanging on retries.
+_HUB_MODEL_INFO_TIMEOUT = 15.0
+
+
+def _embedding_marker_in_hf_cache(model_name: str) -> bool:
+ """True when model_name's cached snapshot carries a modules.json (the ST marker).
+ Cache-only, no network; used offline and as a fallback when the Hub lookup times out."""
+ from utils.utils import hf_cache_snapshot_dir
+
+ snapshot = hf_cache_snapshot_dir(model_name)
+ if snapshot is None:
+ return False
+ try:
+ return (snapshot / "modules.json").is_file()
+ except OSError:
+ return False
+
+
def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
"""Detect embedding/sentence-transformer models via HF metadata.
@@ -2090,6 +2108,15 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
Returns:
True if embedding model, else False (default for local paths or errors).
"""
+ from utils.utils import hf_env_offline
+
+ # Offline (remote repo): reclassify from the local cache on every call, before/without the
+ # memo. An online lookup can memoize True from tags with no weights cached, so trusting it once
+ # the session goes offline would accept a repo _get() cannot load; a cached negative can also be
+ # invalidated by later cache materialization. The cache probe is local-only, so it's cheap.
+ if not is_local_path(model_name) and hf_env_offline():
+ return _embedding_marker_in_hf_cache(model_name)
+
cache_key = (model_name, hf_token)
if cache_key in _embedding_detection_cache:
return _embedding_detection_cache[cache_key]
@@ -2104,7 +2131,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
try:
from huggingface_hub import model_info as hf_model_info
- info = hf_model_info(model_name, token = hf_token)
+ info = hf_model_info(model_name, token = hf_token, timeout = _HUB_MODEL_INFO_TIMEOUT)
tags = set(info.tags or [])
pipeline_tag = info.pipeline_tag or ""
@@ -2125,9 +2152,11 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
return is_emb
except Exception as e:
+ # Timeout or transient network error: fall back to the local cache marker, don't hard-fail.
logger.warning(f"Could not determine if {model_name} is embedding model: {e}")
- _embedding_detection_cache[cache_key] = False
- return False
+ is_emb = _embedding_marker_in_hf_cache(model_name)
+ _embedding_detection_cache[cache_key] = is_emb
+ return is_emb
def _has_model_weight_files(model_dir: Path) -> bool:
diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py
index 466f326f18..0490d38d7c 100644
--- a/studio/backend/utils/security/file_security.py
+++ b/studio/backend/utils/security/file_security.py
@@ -29,13 +29,35 @@ Policy:
scanned so a repo cannot dodge the gate by suffixing its name.
"""
+import re
from dataclasses import dataclass, field
+from pathlib import Path
from typing import Optional
from loggers import get_logger
logger = get_logger(__name__)
+# Pickle-format weight files (plain or sharded) that execute code on load; safetensors/gguf
+# are inert. Grouped by weight family so an inert safetensors only suppresses the pickle it
+# actually replaces: the loader won't use an adapter's safetensors for pytorch_model.bin.
+_PICKLE_WEIGHT_RE = re.compile(
+ r"^(model|pytorch_model|adapter_model|consolidated)(-\d+-of-\d+)?"
+ r"\.(bin|pt|pth|ckpt|pkl|pickle)$",
+ re.IGNORECASE,
+)
+# Base-model safetensors set: HF names the base pickle pytorch_model.bin but the safetensors
+# model.safetensors (stems differ), so a base pickle is replaced only by these, not an adapter's.
+_BASE_SAFETENSORS_RE = re.compile(
+ r"^(model(-\d+-of-\d+)?\.safetensors|model\.safetensors\.index\.json)$",
+ re.IGNORECASE,
+)
+# Adapter (PEFT) safetensors set: adapter_model.safetensors, its shards, or index.
+_ADAPTER_SAFETENSORS_RE = re.compile(
+ r"^(adapter_model(-\d+-of-\d+)?\.safetensors|adapter_model\.safetensors\.index\.json)$",
+ re.IGNORECASE,
+)
+
# Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/
# malicious or a future label) blocks, so Hub schema drift fails CLOSED.
_NONBLOCKING_LEVELS = frozenset(
@@ -265,11 +287,105 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]):
return None
+def _st_load_roots(snapshot: Path) -> list:
+ """Directories a SentenceTransformer load deserializes weights from: the snapshot root plus
+ each module path in modules.json. Local, no network. Mirrors the online gate (which ignores
+ unreferenced nested pickles ST never loads) so the offline gate doesn't over-block."""
+ roots = [snapshot]
+ try:
+ import json
+ modules = json.loads((snapshot / "modules.json").read_text())
+ except (OSError, ValueError):
+ return roots # no / invalid modules.json -> snapshot root is the only load root
+ for module in modules or ():
+ path = str((module or {}).get("path", "")).strip().strip("/")
+ # Relative module path only; ignore a crafted "../" escape.
+ if path and ".." not in path.split("/"):
+ candidate = snapshot / path
+ if candidate not in roots:
+ roots.append(candidate)
+ return roots
+
+
+def _cached_pickle_weight_files(snapshot: Path) -> list:
+ """Pickle weight files in snapshot's ST load roots, EXCLUDING those whose weight family also
+ ships an inert safetensors in the same dir (the loader prefers it): a base pickle is suppressed
+ only by a base model.safetensors, an adapter pickle only by adapter_model.safetensors -- an
+ unrelated safetensors is no substitute. Load roots only. Raises OSError if the snapshot root is
+ unreadable (caller blocks)."""
+ blocked = []
+ for root in _st_load_roots(snapshot):
+ try:
+ entries = [p for p in root.iterdir() if p.is_file()]
+ except OSError:
+ if root == snapshot:
+ raise # top-level unreadable -> fail closed
+ continue # unreadable module subdir: nothing loadable to attest here
+ has_base_safetensors = any(_BASE_SAFETENSORS_RE.match(p.name) for p in entries)
+ has_adapter_safetensors = any(_ADAPTER_SAFETENSORS_RE.match(p.name) for p in entries)
+ for path in entries:
+ if not _PICKLE_WEIGHT_RE.match(path.name):
+ continue
+ is_adapter = path.name.lower().startswith("adapter_model")
+ has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors
+ if not has_alternative:
+ blocked.append(path)
+ return blocked
+
+
+def _evaluate_local_only(model_name: str) -> FileSecurityDecision:
+ """Offline security gate. The Hub scan is unreachable, so inspect the local cache and fail
+ CLOSED on an unscanned pickle weight with no inert safetensors alternative, rather than
+ failing open or hanging. Safetensors/gguf-only cache loads; nothing cached -> allowed."""
+ from utils.utils import hf_cache_snapshot_dir
+
+ try:
+ snapshot = hf_cache_snapshot_dir(model_name)
+ except Exception:
+ logger.warning("Offline gate: could not resolve the cache for '%s'; blocking.", model_name)
+ return FileSecurityDecision(
+ model_name, True, reason = "offline; could not inspect the local cache"
+ )
+
+ if snapshot is None:
+ return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to load")
+
+ try:
+ pickles = _cached_pickle_weight_files(snapshot)
+ except OSError:
+ logger.warning("Offline gate: could not read the cache for '%s'; blocking.", model_name)
+ return FileSecurityDecision(
+ model_name, True, reason = "offline; could not read the local cache"
+ )
+
+ if not pickles:
+ return FileSecurityDecision(
+ model_name, False, reason = "offline; cached weights are inert (safetensors/gguf)"
+ )
+
+ # Snapshot-relative posix paths (match the online gate; disambiguate same-named pickles).
+ rel_paths = sorted(p.relative_to(snapshot).as_posix() for p in pickles)
+ names = ", ".join(rel_paths)
+ logger.warning(
+ "Blocking offline load of '%s': cached pickle weight(s) cannot be malware-scanned "
+ "offline and have no safetensors alternative (%s).",
+ model_name,
+ names,
+ )
+ return FileSecurityDecision(
+ model_name,
+ True,
+ unsafe_files = [{"path": rel, "level": "unscanned"} for rel in rel_paths],
+ reason = f"offline; unscanned pickle weights with no safetensors alternative: {names}",
+ )
+
+
def evaluate_file_security(
model_name: str,
hf_token: Optional[str] = None,
*,
load_subdirs = (),
+ local_only_load: bool = False,
) -> FileSecurityDecision:
"""Block a load when HF's security scan flags unsafe serialized files.
@@ -280,6 +396,9 @@ def evaluate_file_security(
``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)``
for Spark-TTS / BiCodec, loading ``/LLM``): a flagged file directly under one
is root-level there and blocks, and an index inside it is honored when scoping shards.
+
+ ``local_only_load`` marks an offline load: with the Hub scan unreachable, inspect the local
+ cache and fail CLOSED on an unscanned pickle weight with no safetensors alternative.
"""
# Scan the repo the load actually fetches, not the literal alias (which 404s and
# fails open): the Spark-TTS "/LLM" alias is really unsloth/ from LLM/.
@@ -295,6 +414,10 @@ def evaluate_file_security(
# Cannot classify the path -> do not block on that account.
return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked")
+ # Offline: inspect the local cache and fail closed rather than hang on model_info or fail open.
+ if local_only_load:
+ return _evaluate_local_only(model_name)
+
status = _fetch_security_status(model_name, hf_token)
if not isinstance(status, dict):
return FileSecurityDecision(
diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py
index 31f5f31bee..21e11c6706 100644
--- a/studio/backend/utils/utils.py
+++ b/studio/backend/utils/utils.py
@@ -8,6 +8,7 @@ import structlog
from loggers import get_logger
from contextlib import contextmanager
from pathlib import Path
+from typing import Optional
import shutil
import tempfile
@@ -15,6 +16,110 @@ import tempfile
logger = get_logger(__name__)
+# ── Offline / HF-cache helpers ──────────────────────────────────
+# An offline load must never touch the network (a DNS-dead session hangs on hub retries);
+# these read the local HF cache the load itself uses.
+
+_HF_OFFLINE_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
+
+
+def hf_env_offline() -> bool:
+ """True when HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE requests offline mode.
+
+ Also honors TRANSFORMERS_OFFLINE (hub honors only HF_HUB_OFFLINE) since users set it
+ to keep transformers loads local.
+ """
+ for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"):
+ if os.environ.get(var, "").strip().lower() in _HF_OFFLINE_TRUE_VALUES:
+ return True
+ return False
+
+
+def st_repo_id_candidates(model_name: str) -> list:
+ """Repo ids a Sentence-Transformers load may resolve model_name to; a slashless name
+ also resolves under the sentence-transformers/ namespace, so both are candidates."""
+ name = (model_name or "").strip().strip("/")
+ if not name:
+ return []
+ candidates = [name]
+ if "/" not in name:
+ candidates.append(f"sentence-transformers/{name}")
+ return candidates
+
+
+def _expand_path(raw: str) -> Path:
+ """Expand ~ and $VARS as huggingface_hub does, so the gate resolves the loader's dir."""
+ return Path(os.path.expandvars(os.path.expanduser(raw)))
+
+
+def _hf_cache_roots() -> list:
+ """The one cache root the loader resolves to, by its own precedence (it picks ONE
+ cache_folder, no fall-through): SENTENCE_TRANSFORMERS_HOME, else HF_HUB_CACHE, else
+ HF_HOME/hub, else ~/.cache/huggingface/hub. Expanded, read from env, one-element list."""
+ st_home = os.environ.get("SENTENCE_TRANSFORMERS_HOME")
+ if st_home:
+ return [_expand_path(st_home)]
+ hub = os.environ.get("HF_HUB_CACHE") or os.environ.get("HUGGINGFACE_HUB_CACHE")
+ if hub:
+ return [_expand_path(hub)]
+ hf_home = os.environ.get("HF_HOME")
+ if hf_home:
+ return [_expand_path(hf_home) / "hub"]
+ return [Path.home() / ".cache" / "huggingface" / "hub"]
+
+
+def hf_cache_snapshot_dir(model_name: str) -> Optional[Path]:
+ """Active local snapshot dir for model_name's main revision, or None if not cached.
+ Reads refs/main then snapshots/; no network. Tries the ST alias for slashless names."""
+ try:
+ from huggingface_hub.file_download import repo_folder_name
+ except Exception:
+ repo_folder_name = None
+ for cache_root in _hf_cache_roots():
+ for repo_id in st_repo_id_candidates(model_name):
+ try:
+ if repo_folder_name is not None:
+ folder = repo_folder_name(repo_id = repo_id, repo_type = "model")
+ else:
+ folder = "models--" + repo_id.replace("/", "--")
+ repo_dir = cache_root / folder
+ ref = repo_dir / "refs" / "main"
+ if not ref.is_file():
+ continue
+ commit = ref.read_text().strip()
+ if not commit:
+ continue
+ snapshot = repo_dir / "snapshots" / commit
+ if snapshot.is_dir():
+ return snapshot
+ except OSError:
+ continue
+ return None
+
+
+# A weight file plus a config distinguishes a real cached model from a metadata-only
+# partial cache that resolves refs/main but would fail at load time.
+_LOADABLE_WEIGHT_SUFFIXES = frozenset({".safetensors", ".bin", ".gguf", ".pt", ".pth", ".ckpt"})
+
+
+def hf_cache_snapshot_is_loadable(model_name: str) -> bool:
+ """True when model_name's snapshot is cached and loadable: a config (config.json or
+ modules.json) plus at least one weight file, not a metadata-only partial cache. No network."""
+ snapshot = hf_cache_snapshot_dir(model_name)
+ if snapshot is None:
+ return False
+ try:
+ has_config = (snapshot / "config.json").is_file() or (snapshot / "modules.json").is_file()
+ if not has_config:
+ return False
+ for path in snapshot.rglob("*"):
+ if path.suffix.lower() in _LOADABLE_WEIGHT_SUFFIXES and path.is_file():
+ return True
+ except OSError:
+ return False
+ return False
+
+
# ── Client-safe error helpers ───────────────────────────────────
# Never return raw exception text to clients; log server-side, return generic.
From 968e6230a0cd97e6356662d3ea5f4543f15a5116 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Wed, 22 Jul 2026 04:34:58 -0700
Subject: [PATCH 068/255] Unsloth start: add local subagents for Claude Code,
Codex, OpenCode and Pi (#7326)
Bring the local-subagent support onto main. The original change (#7316) merged
into the stacked pr/daniel-unsloth-start-audit branch rather than main, and #7313
reached main via squash, so these files never landed on main.
Adds --as-subagent for claude, codex, opencode and pi: the parent agent keeps its
own cloud model while a locally served GGUF is registered as a delegated subagent,
using ephemeral per-session config that never touches the user's real agent config.
---
README.md | 7 +
pyproject.toml | 2 +-
unsloth_cli/claude_subagent_mcp.py | 366 ++++++++++++
unsloth_cli/commands/start.py | 525 ++++++++++++++++--
unsloth_cli/pi_subagent.ts | 241 ++++++++
unsloth_cli/tests/test_claude_subagent_mcp.py | 338 +++++++++++
unsloth_cli/tests/test_pi_subagent.py | 191 +++++++
unsloth_cli/tests/test_start.py | 490 +++++++++++++++-
8 files changed, 2108 insertions(+), 52 deletions(-)
create mode 100644 unsloth_cli/claude_subagent_mcp.py
create mode 100644 unsloth_cli/pi_subagent.ts
create mode 100644 unsloth_cli/tests/test_claude_subagent_mcp.py
create mode 100644 unsloth_cli/tests/test_pi_subagent.py
diff --git a/README.md b/README.md
index 6aa8f4f4c3..514454f985 100644
--- a/README.md
+++ b/README.md
@@ -86,6 +86,13 @@ Replace `claude` with any supported agent:
| OpenCode | `unsloth start opencode` |
| Pi Coding Agent | `unsloth start pi` |
+Claude Code, Codex, OpenCode and Pi can keep their current model and use Unsloth as a local
+subagent:
+
+```bash
+unsloth start claude --as-subagent --model unsloth/model-GGUF:quant
+```
+
## 📥 Install
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.
diff --git a/pyproject.toml b/pyproject.toml
index 071258eb8f..a5436a8916 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,7 @@ version = {attr = "unsloth.models._utils.__version__"}
include-package-data = true
[tool.setuptools.package-data]
-unsloth_cli = ["codex_fallback_prompt.md"]
+unsloth_cli = ["codex_fallback_prompt.md", "pi_subagent.ts"]
studio = [
"*.sh",
"*.ps1",
diff --git a/unsloth_cli/claude_subagent_mcp.py b/unsloth_cli/claude_subagent_mcp.py
new file mode 100644
index 0000000000..b86368515b
--- /dev/null
+++ b/unsloth_cli/claude_subagent_mcp.py
@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Small stdio MCP bridge from cloud Claude Code to a local Claude Code child."""
+
+from __future__ import annotations
+
+import json
+import os
+import signal
+import shutil
+import subprocess
+import sys
+import threading
+import time
+from typing import Any, Callable
+
+from unsloth_cli.commands.start import (
+ _CLAUDE_ENV_UNSET,
+ _SUBAGENT_DESCRIPTION,
+ _SUBAGENT_INSTRUCTIONS,
+ _claude_flags,
+ _claude_local_env,
+ _wsl_shim_env,
+)
+
+_MAX_RESULT_CHARACTERS = 100_000
+_CANCEL_POLL_SECONDS = 0.1
+_CANCEL_GRACE_SECONDS = 2.0
+
+
+def _required_env(name: str) -> str:
+ value = os.environ.get(name, "").strip()
+ if not value:
+ raise RuntimeError(f"Missing {name}.")
+ return value
+
+
+def _bounded(text: str) -> str:
+ if len(text) <= _MAX_RESULT_CHARACTERS:
+ return text
+ return text[:_MAX_RESULT_CHARACTERS] + "\n\n[Local agent output truncated]"
+
+
+def _result_text(stdout: str) -> str:
+ lines = [line for line in stdout.splitlines() if line.strip()]
+ candidates = [stdout.strip(), *reversed(lines)]
+ for candidate in candidates:
+ try:
+ payload = json.loads(candidate)
+ except ValueError:
+ continue
+ if not isinstance(payload, dict):
+ continue
+ result = payload.get("result")
+ if payload.get("is_error"):
+ raise RuntimeError(str(result or "The local Claude agent failed."))
+ if isinstance(result, str) and result.strip():
+ return _bounded(result.strip())
+ raise RuntimeError("The local Claude agent returned no readable result.")
+
+
+def _stop_child(process: subprocess.Popen) -> None:
+ """Stop the Claude child and any tool processes it started."""
+ if process.poll() is not None:
+ if os.name != "nt":
+ # Leader exited, but its tool processes may still be running.
+ try:
+ os.killpg(process.pid, signal.SIGTERM)
+ except OSError:
+ return
+ time.sleep(_CANCEL_GRACE_SECONDS)
+ try:
+ os.killpg(process.pid, signal.SIGKILL)
+ except OSError:
+ pass
+ return
+ if os.name == "nt":
+ try:
+ completed = subprocess.run(
+ ["taskkill", "/PID", str(process.pid), "/T", "/F"],
+ capture_output = True,
+ timeout = 15,
+ check = False,
+ )
+ except Exception:
+ completed = None
+ # A failed taskkill must not leave the child running through the grace wait.
+ if (completed is None or completed.returncode != 0) and process.poll() is None:
+ process.terminate()
+ else:
+ try:
+ os.killpg(process.pid, signal.SIGTERM)
+ except OSError:
+ process.terminate()
+ try:
+ process.wait(timeout = _CANCEL_GRACE_SECONDS)
+ except subprocess.TimeoutExpired:
+ if os.name == "nt":
+ process.kill()
+ else:
+ try:
+ os.killpg(process.pid, signal.SIGKILL)
+ except OSError:
+ process.kill()
+ process.wait()
+ else:
+ if os.name != "nt":
+ # Leader is gone; kill any surviving group members.
+ try:
+ os.killpg(process.pid, signal.SIGKILL)
+ except OSError:
+ pass
+
+
+def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> str:
+ base = _required_env("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL")
+ key = _required_env("UNSLOTH_CLAUDE_SUBAGENT_API_KEY")
+ model = _required_env("UNSLOTH_CLAUDE_SUBAGENT_MODEL")
+ window = int(os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW", "0") or 0)
+ entry = {"id": model, "context_length": window}
+ local_env = _claude_local_env(base, key, entry)
+ child_env = dict(os.environ)
+
+ executable = shutil.which("claude")
+ if executable is None:
+ raise RuntimeError("`claude` is not installed or is not on PATH.")
+ cancel_event = cancel_event or threading.Event()
+ if cancel_event.is_set():
+ raise RuntimeError("The local Claude agent was cancelled.")
+ command = [
+ "claude",
+ "--model",
+ model,
+ *_claude_flags(model),
+ "--permission-mode",
+ (
+ "bypassPermissions"
+ if os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS") == "1"
+ else "acceptEdits"
+ ),
+ "--print",
+ "--output-format",
+ "json",
+ "--no-session-persistence",
+ "--append-system-prompt",
+ _SUBAGENT_INSTRUCTIONS,
+ f"Task: {task}",
+ ]
+ bridged, wsl_names = _wsl_shim_env(command, local_env, _CLAUDE_ENV_UNSET)
+ if wsl_names:
+ from unsloth_cli.commands.start import _merge_wslenv
+
+ bridged = {**bridged, "PWD": os.getcwd()}
+ child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_names)
+ for name in _CLAUDE_ENV_UNSET:
+ child_env[name] = ""
+ else:
+ for name in _CLAUDE_ENV_UNSET:
+ child_env.pop(name, None)
+ child_env.update(bridged)
+ popen_kwargs: dict[str, Any] = {
+ "cwd": os.environ.get("CLAUDE_PROJECT_DIR") or os.getcwd(),
+ "env": child_env,
+ "stdin": subprocess.DEVNULL,
+ "stdout": subprocess.PIPE,
+ "stderr": subprocess.PIPE,
+ "text": True,
+ }
+ if os.name == "nt":
+ popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
+ else:
+ popen_kwargs["start_new_session"] = True
+ process = subprocess.Popen(
+ [executable, *command[1:]],
+ **popen_kwargs,
+ )
+ try:
+ while True:
+ try:
+ stdout, stderr = process.communicate(timeout = _CANCEL_POLL_SECONDS)
+ break
+ except subprocess.TimeoutExpired:
+ if cancel_event.is_set():
+ _stop_child(process)
+ raise RuntimeError("The local Claude agent was cancelled.")
+ except BaseException:
+ if process.poll() is None:
+ _stop_child(process)
+ raise
+ if process.returncode != 0:
+ detail = stderr.strip() or stdout.strip()
+ raise RuntimeError(
+ _bounded(detail) or f"Local Claude exited with code {process.returncode}."
+ )
+ return _result_text(stdout)
+
+
+def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) -> dict | None:
+ request_id = request.get("id")
+ method = request.get("method")
+ if request_id is None:
+ return None
+ if method == "initialize":
+ protocol = (request.get("params") or {}).get("protocolVersion") or "2025-06-18"
+ result = {
+ "protocolVersion": protocol,
+ "capabilities": {"tools": {"listChanged": False}},
+ "serverInfo": {"name": "unsloth-local-agent", "version": "1.0.0"},
+ }
+ elif method == "ping":
+ result = {}
+ elif method == "tools/list":
+ result = {
+ "tools": [
+ {
+ "name": "unsloth_agent",
+ "title": "Unsloth local agent",
+ "description": _SUBAGENT_DESCRIPTION,
+ "inputSchema": {
+ "type": "object",
+ "properties": {
+ "task": {
+ "type": "string",
+ "description": "The complete task for the local Unsloth agent.",
+ }
+ },
+ "required": ["task"],
+ "additionalProperties": False,
+ },
+ "annotations": {
+ "readOnlyHint": False,
+ "destructiveHint": True,
+ "idempotentHint": False,
+ "openWorldHint": True,
+ },
+ "_meta": {"anthropic/maxResultSizeChars": _MAX_RESULT_CHARACTERS},
+ }
+ ]
+ }
+ elif method == "tools/call":
+ params = request.get("params") or {}
+ arguments = params.get("arguments") or {}
+ task = arguments.get("task") if params.get("name") == "unsloth_agent" else None
+ if not isinstance(task, str) or not task.strip():
+ result = {
+ "content": [{"type": "text", "text": "A non-empty task is required."}],
+ "isError": True,
+ }
+ else:
+ try:
+ text = run_agent(task.strip())
+ result = {"content": [{"type": "text", "text": text}], "isError": False}
+ except Exception as exc:
+ result = {
+ "content": [{"type": "text", "text": str(exc)}],
+ "isError": True,
+ }
+ else:
+ return {
+ "jsonrpc": "2.0",
+ "id": request_id,
+ "error": {"code": -32601, "message": f"Method not found: {method}"},
+ }
+ return {"jsonrpc": "2.0", "id": request_id, "result": result}
+
+
+def serve(
+ stdin: Any = sys.stdin,
+ stdout: Any = sys.stdout,
+ run_agent: Callable[[str, threading.Event], str] = run_local_agent,
+) -> None:
+ active: dict[object, threading.Event] = {}
+ workers: list[threading.Thread] = []
+ state_lock = threading.RLock()
+ output_lock = threading.Lock()
+ shutdown_started = threading.Event()
+
+ def cancel_active() -> None:
+ with state_lock:
+ pending = list(active.values())
+ for cancel_event in pending:
+ cancel_event.set()
+
+ def handle_shutdown(_signum: int, _frame: Any) -> None:
+ # Claude Code sends SIGINT (possibly repeatedly) to cancel a tool call. Only
+ # the first unwinds stdin; later ones must not interrupt process-tree cleanup.
+ first_signal = not shutdown_started.is_set()
+ shutdown_started.set()
+ cancel_active()
+ if first_signal:
+ raise KeyboardInterrupt
+
+ previous_handlers: dict[int, Any] = {}
+ if threading.current_thread() is threading.main_thread():
+ for signum in (signal.SIGINT, signal.SIGTERM):
+ previous_handlers[signum] = signal.signal(signum, handle_shutdown)
+
+ def send(response: dict | None) -> None:
+ if response is None:
+ return
+ with output_lock:
+ stdout.write(json.dumps(response, separators = (",", ":")) + "\n")
+ stdout.flush()
+
+ def call_tool(request: dict, request_id: object, cancel_event: threading.Event) -> None:
+ try:
+ response = _response(
+ request,
+ run_agent = lambda task: run_agent(task, cancel_event),
+ )
+ if not cancel_event.is_set():
+ send(response)
+ finally:
+ with state_lock:
+ if active.get(request_id) is cancel_event:
+ active.pop(request_id, None)
+
+ try:
+ for line in stdin:
+ try:
+ request = json.loads(line)
+ if not isinstance(request, dict):
+ response = None
+ elif request.get("method") == "notifications/cancelled":
+ request_id = (request.get("params") or {}).get("requestId")
+ with state_lock:
+ cancel_event = active.get(request_id)
+ if cancel_event is not None:
+ cancel_event.set()
+ response = None
+ elif request.get("method") == "tools/call" and request.get("id") is not None:
+ request_id = request["id"]
+ cancel_event = threading.Event()
+ with state_lock:
+ active[request_id] = cancel_event
+ worker = threading.Thread(
+ target = call_tool,
+ args = (request, request_id, cancel_event),
+ name = f"unsloth-agent-{request_id}",
+ )
+ workers.append(worker)
+ worker.start()
+ response = None
+ else:
+ response = _response(request)
+ except Exception as exc:
+ response = {
+ "jsonrpc": "2.0",
+ "id": None,
+ "error": {"code": -32603, "message": str(exc)},
+ }
+ send(response)
+ except KeyboardInterrupt:
+ pass
+ finally:
+ cancel_active()
+ for worker in workers:
+ if worker.ident is not None:
+ worker.join()
+ for signum, handler in previous_handlers.items():
+ signal.signal(signum, handler)
+
+
+if __name__ == "__main__":
+ serve()
diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py
index 317d1f4f3e..ba2972d6dc 100644
--- a/unsloth_cli/commands/start.py
+++ b/unsloth_cli/commands/start.py
@@ -73,11 +73,21 @@ _HERMES_POSIX_INSTALL_HINT = (
# windows and scales the compaction threshold back down to the real window.
_HERMES_MIN_CONTEXT = 65536
_PI_PROVIDER = "unsloth"
-# OpenCode selects a model by "/" and honors a user
-# disabled_providers list. Register the session provider under a dedicated id a
-# user's disable list would never target, so the model is always selectable
-# without the wrapper having to reconstruct (and override) OpenCode's full,
-# multi-layer disabled_providers resolution.
+_SUBAGENT_NAME = "unsloth"
+_SUBAGENT_DESCRIPTION = (
+ "Local coding subagent powered by Unsloth for debugging, implementation, and codebase "
+ "research. Use when the user asks to spawn an Unsloth or local agent."
+)
+_SUBAGENT_INSTRUCTIONS = (
+ "You are a local coding subagent powered by Unsloth. Complete the assigned task directly, "
+ "use the available tools when useful, verify your work, and return a concise result to the "
+ "parent agent."
+)
+_CLAUDE_SUBAGENT_MCP_MODULE = "unsloth_cli.claude_subagent_mcp"
+_CLAUDE_SUBAGENT_TOOL = "mcp__plugin_unsloth-local-agent_unsloth__unsloth_agent"
+_PI_SUBAGENT_EXTENSION = Path(__file__).parent.parent / "pi_subagent.ts"
+# OpenCode selects a model by "/". Use a dedicated id to avoid
+# colliding with a user's providers; provider filters are set in the launch-time overlay.
_OPENCODE_PROVIDER = "unsloth-studio"
_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]"
_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True}
@@ -158,6 +168,11 @@ _PERSIST_OPTION = typer.Option(
"the agent unchanged."
),
)
+_AS_SUBAGENT_OPTION = typer.Option(
+ False,
+ "--as-subagent",
+ help = "Keep the coding agent's current model and add Unsloth as a local subagent.",
+)
# Per-agent CLI flag for "run tools without prompting". OpenCode (native --auto is
# command-scoped, handled below) and OpenClaw (config-only) are absent from this prefix map.
@@ -334,11 +349,54 @@ def _display_model_spec(model: str, variant: Optional[str]) -> str:
return f"{repo}:{selected_variant}" if selected_variant else model
+def _subagent_model_id(
+ base: str,
+ key: str,
+ entry: dict,
+ requested_model: Optional[str],
+ requested_variant: Optional[str],
+) -> str:
+ """Return an API model id that preserves the selected GGUF variant.
+
+ Coding-agent model definitions outlive the initial load. If Unsloth later
+ unloads the model, a bare repository id may resolve to a different cached
+ quant. Include the explicit or currently loaded variant so an automatic
+ reload selects the same weights.
+ """
+ model_id = str(entry["id"])
+ _, inline_variant = _split_repo_variant(requested_model or "")
+ variant = requested_variant or inline_variant
+ if not variant:
+ try:
+ status = _http_json("GET", f"{base}/api/inference/status", key)
+ except Exception:
+ status = {}
+ typer.echo(
+ "Warning: could not verify the loaded GGUF variant; a later reload "
+ "may pick a different cached quant. Pass :variant to pin it.",
+ err = True,
+ )
+ if status.get("is_gguf"):
+ variant = status.get("gguf_variant")
+ return (
+ _display_model_spec(model_id, str(variant))
+ if variant and _is_hub_model_id(model_id)
+ else model_id
+ )
+
+
def _fail(message: str) -> NoReturn:
typer.echo(message, err = True)
raise typer.Exit(code = 1)
+def _reject_as_subagent(agent: str, args: list) -> None:
+ # Reject early; otherwise the flag reaches the agent binary and fails after
+ # Studio has already loaded the model.
+ if "--as-subagent" in args:
+ _fail(f"--as-subagent is not supported for {agent}.")
+
+
def _http_error_detail(exc: urllib.error.HTTPError) -> str:
try:
body = json.loads(exc.read().decode())
@@ -1278,6 +1336,25 @@ def _claude_flags(model_id: str) -> list:
return [_DYNAMIC_SECTIONS_FLAG, "--settings", _claude_settings_overlay(model_id)]
+def _claude_local_env(base: str, key: str, entry: dict) -> dict:
+ """Build the local endpoint, cache, display, and compaction environment."""
+ model_id = entry["id"]
+ env = {
+ "ANTHROPIC_BASE_URL": base,
+ "ANTHROPIC_AUTH_TOKEN": key,
+ "ANTHROPIC_MODEL": model_id,
+ "CLAUDE_CODE_ATTRIBUTION_HEADER": "0",
+ "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
+ "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
+ "CLAUDE_CODE_NO_FLICKER": "1",
+ }
+ window = entry.get("context_length") or entry.get("max_context_length")
+ if window:
+ env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = str(int(window))
+ env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] = "90"
+ return env
+
+
def _merge_codex_config(existing: str, base: str) -> str:
chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table
if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]):
@@ -1391,6 +1468,206 @@ def write_codex_config(base: str, model: dict, home: Path) -> None:
typer.echo(f"Updated {profile}")
+def write_codex_subagent_config(base: str, key: str, model: dict, home: Path) -> Path:
+ """Write a session-scoped Codex custom agent without replacing the main model."""
+ home.mkdir(parents = True, exist_ok = True)
+ model_id = model["id"]
+ window = model.get("context_length") or model.get("max_context_length")
+ catalog_name = "unsloth-model-catalog.json"
+ text = (
+ f"name = {json.dumps(_SUBAGENT_NAME)}\n"
+ f"description = {json.dumps(_SUBAGENT_DESCRIPTION)}\n"
+ f"developer_instructions = {json.dumps(_SUBAGENT_INSTRUCTIONS)}\n"
+ f"model_provider = {json.dumps(_CODEX_PROFILE)}\n"
+ f"model = {json.dumps(model_id)}\n"
+ )
+ if _codex_supports_model_catalog() and _CODEX_FALLBACK_PROMPT.is_file():
+ catalog = home / catalog_name
+ catalog_text = json.dumps(_codex_model_catalog(model), indent = 2) + "\n"
+ if not catalog.exists() or catalog.read_text(encoding = "utf-8") != catalog_text:
+ catalog.write_text(catalog_text, encoding = "utf-8")
+ typer.echo(f"Updated {catalog}")
+ text += f"model_catalog_json = {json.dumps(catalog_name)}\n"
+ if window:
+ text += f"model_context_window = {int(window)}\n"
+ credential = home / "unsloth-auth.json"
+ _write_private_json(credential, {"token": key})
+ auth_command = sys.executable
+ auth_args = [
+ "-c",
+ "import json,sys; print(json.load(open(sys.argv[1], encoding='utf-8'))['token'])",
+ str(credential),
+ ]
+ if _wsl_windows_executable(["codex"]):
+ auth_command = "wsl.exe"
+ auth_args = [
+ "-d",
+ os.environ["WSL_DISTRO_NAME"],
+ "--",
+ sys.executable,
+ *auth_args,
+ ]
+ text += (
+ f"\n{_PROVIDER_HEADER}\n"
+ 'name = "Unsloth Studio"\n'
+ f"base_url = {json.dumps(base + '/v1')}\n"
+ 'wire_api = "responses"\n'
+ f"\n{_PROVIDER_HEADER[:-1]}.auth]\n"
+ f"command = {json.dumps(auth_command)}\n"
+ f"args = {json.dumps(auth_args)}\n"
+ "timeout_ms = 5000\n"
+ )
+ path = home / f"{_SUBAGENT_NAME}.toml"
+ if not path.exists() or path.read_text(encoding = "utf-8") != text:
+ path.write_text(text, encoding = "utf-8")
+ typer.echo(f"Updated {path}")
+ return path
+
+
+def _agent_config_path(path: Path, command: list) -> str:
+ """Translate a generated config path when a Windows agent runs through WSL."""
+ return _wsl_windows_path(path) if _wsl_windows_executable(command) else str(path)
+
+
+def _opencode_subagent_inline_config(path: Path, permission: dict) -> dict:
+ """Keep the local provider visible without hiding the parent's allowed providers."""
+ inline: dict = {}
+ inherited = os.environ.get("OPENCODE_CONFIG_CONTENT")
+ if inherited:
+ try:
+ parsed = json.loads(inherited)
+ except ValueError:
+ _fail("OPENCODE_CONFIG_CONTENT is not valid JSON.")
+ if not isinstance(parsed, dict):
+ _fail("OPENCODE_CONFIG_CONTENT must contain a JSON object.")
+ inline.update(parsed)
+
+ def merge_provider_filters(effective_config: dict) -> None:
+ enabled = effective_config.get("enabled_providers")
+ if isinstance(enabled, list):
+ inline["enabled_providers"] = list(dict.fromkeys([*enabled, _OPENCODE_PROVIDER]))
+ disabled = effective_config.get("disabled_providers")
+ if isinstance(disabled, list) and _OPENCODE_PROVIDER in disabled:
+ inline["disabled_providers"] = [
+ provider for provider in disabled if provider != _OPENCODE_PROVIDER
+ ]
+
+ # The inherited inline layer is already highest priority. Merge it even when
+ # OpenCode is not installed yet, as in fresh-install and --no-launch flows.
+ merge_provider_filters(inline)
+ effective = inline
+
+ executable = _which_with_install_dirs("opencode")
+ if executable is None:
+ typer.echo(
+ f"Warning: OpenCode is not installed, so provider filters could not be checked. "
+ f"The target configuration must allow '{_OPENCODE_PROVIDER}'.",
+ err = True,
+ )
+ else:
+ env = os.environ.copy()
+ env["OPENCODE_CONFIG"] = _agent_config_path(path, ["opencode"])
+ try:
+ resolved = subprocess.run(
+ [executable, "debug", "config"],
+ capture_output = True,
+ text = True,
+ timeout = 15,
+ env = env,
+ )
+ except Exception as exc:
+ _fail(f"Could not inspect OpenCode provider filters: {exc}")
+ if resolved.returncode != 0:
+ detail = resolved.stderr.strip() or resolved.stdout.strip()
+ _fail(f"Could not inspect OpenCode provider filters: {detail or 'unknown error'}")
+ try:
+ effective = json.loads(resolved.stdout)
+ except ValueError:
+ _fail("Could not inspect OpenCode provider filters: invalid JSON response.")
+ if not isinstance(effective, dict):
+ _fail("Could not inspect OpenCode provider filters: expected a JSON object.")
+
+ merge_provider_filters(effective)
+
+ depth = effective.get("subagent_depth")
+ inline["subagent_depth"] = (
+ depth if isinstance(depth, int) and not isinstance(depth, bool) and depth > 0 else 1
+ )
+ if permission:
+ inline["permission"] = permission
+ return inline
+
+
+def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path:
+ """Write a session plugin that exposes the local Claude child through MCP."""
+ plugin = path / "unsloth-local-agent"
+ command = sys.executable
+ args = ["-m", _CLAUDE_SUBAGENT_MCP_MODULE]
+ mcp_env = dict(server_env)
+ if _wsl_windows_executable(["claude"]):
+ command = "wsl.exe"
+ args = [
+ "-d",
+ os.environ["WSL_DISTRO_NAME"],
+ "--",
+ sys.executable,
+ "-m",
+ _CLAUDE_SUBAGENT_MCP_MODULE,
+ ]
+ mcp_env["WSLENV"] = _merge_wslenv(
+ os.environ.get("WSLENV", ""),
+ _wsl_bridge_names(server_env, ()),
+ )
+ _write_private_json(
+ plugin / ".claude-plugin" / "plugin.json",
+ {
+ "name": "unsloth-local-agent",
+ "version": "1.0.0",
+ "description": _SUBAGENT_DESCRIPTION,
+ "author": {"name": "Unsloth AI"},
+ },
+ )
+ _write_private_json(
+ plugin / ".mcp.json",
+ {
+ "mcpServers": {
+ "unsloth": {
+ "type": "stdio",
+ "command": command,
+ "args": args,
+ "env": mcp_env,
+ }
+ }
+ },
+ )
+ skill = plugin / "skills" / "local-agent" / "SKILL.md"
+ skill.parent.mkdir(parents = True, exist_ok = True, mode = 0o700)
+ skill.write_text(
+ "---\n"
+ "description: Delegate a task to the local agent powered by Unsloth. Use when the "
+ "user asks to spawn an Unsloth agent or local agent.\n"
+ "---\n\n"
+ "Call the Unsloth local agent tool once with the complete task. Return its result "
+ "to the user without claiming that the cloud parent completed the local work.\n",
+ encoding = "utf-8",
+ )
+ return plugin
+
+
+def _codex_subagent_flags(path: Path) -> list[str]:
+ config_path = _agent_config_path(path, ["codex"])
+ return [
+ "--enable",
+ "multi_agent",
+ "-c",
+ "agents.max_depth=1",
+ "-c",
+ f"agents.{_SUBAGENT_NAME}.description={json.dumps(_SUBAGENT_DESCRIPTION)}",
+ "-c",
+ f"agents.{_SUBAGENT_NAME}.config_file={json.dumps(config_path)}",
+ ]
+
+
def _wsl_windows_executable(command: list) -> Optional[str]:
if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"):
return None
@@ -1974,6 +2251,7 @@ def write_opencode_config(
model: dict,
path: Path,
yolo: bool = False,
+ as_subagent: bool = False,
) -> dict:
config = _read_json_object(path)
if config is None:
@@ -1985,10 +2263,8 @@ def write_opencode_config(
return {}
before = json.dumps(config, sort_keys = True)
config.setdefault("$schema", "https://opencode.ai/config.json")
- # The session provider is registered under a dedicated id (_OPENCODE_PROVIDER)
- # that a user's disabled_providers list would never target, so it is always
- # selectable without this overlay having to reconstruct or override OpenCode's
- # disabled_providers resolution.
+ # Keep the provider definition in this private session file. The launch path
+ # adjusts effective provider filters in the higher-priority inline overlay.
model_entry = {"name": model["id"]}
window = model.get("context_length") or model.get("max_context_length")
if window:
@@ -2003,15 +2279,36 @@ def write_opencode_config(
"options": {"baseURL": f"{base}/v1", "apiKey": key},
"models": {model["id"]: model_entry},
}
- # OpenCode selects a model by "/".
- config["model"] = f"{_OPENCODE_PROVIDER}/{model['id']}"
- if window:
+ # Normal mode pins this as the session model. Subagent mode leaves the user's
+ # main/small models alone and exposes the local model to @unsloth and /models.
+ opencode_model = f"{_OPENCODE_PROVIDER}/{model['id']}"
+ if as_subagent:
+ for field in ("model", "small_model"):
+ if str(config.get(field) or "").startswith(f"{_OPENCODE_PROVIDER}/"):
+ config.pop(field, None)
+ managed_compaction = {"auto": True, "reserved": max(1, window // 10)} if window else None
+ if managed_compaction and config.get("compaction") == managed_compaction:
+ config.pop("compaction", None)
+ _subdict(config, "agent")[_SUBAGENT_NAME] = {
+ "description": _SUBAGENT_DESCRIPTION,
+ "mode": "subagent",
+ "model": opencode_model,
+ "prompt": _SUBAGENT_INSTRUCTIONS,
+ }
+ else:
+ config["model"] = opencode_model
+ agents = config.get("agent")
+ if isinstance(agents, dict):
+ agents.pop(_SUBAGENT_NAME, None)
+ if not agents:
+ config.pop("agent", None)
+ if window and not as_subagent:
# Compact with ~10% headroom (near 90% full). The fixed 20k-token default
# buffer over-compacts, or never settles, on a small local context.
compaction = _subdict(config, "compaction")
compaction["auto"] = True
compaction["reserved"] = max(1, window // 10)
- tools = ("edit", "bash", "webfetch")
+ tools = ("edit", "bash", "webfetch", *(("task",) if as_subagent else ()))
if yolo:
# Fallback for commands without native --auto and for the append-safe bare
# --no-launch command (subcommand unknown yet). Rides inline (OPENCODE_CONFIG_CONTENT)
@@ -2140,6 +2437,22 @@ def write_pi_config(base: str, key: str, model: dict, path: Path) -> None:
typer.echo(f"Updated {path}")
+def write_pi_subagent_config(base: str, key: str, model: dict, path: Path) -> None:
+ """Write private bootstrap data for the bundled Pi extension."""
+ window = model.get("context_length") or model.get("max_context_length")
+ window = int(window) if window else 32768
+ _write_private_json(
+ path,
+ {
+ "baseUrl": f"{base}/v1",
+ "apiKey": key,
+ "model": model["id"],
+ "contextWindow": window,
+ "maxTokens": min(window // 4, 8192),
+ },
+ )
+
+
@start_app.command("claude", context_settings = _PASSTHROUGH)
def claude(
ctx: typer.Context,
@@ -2153,6 +2466,7 @@ def claude(
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
+ as_subagent: bool = _AS_SUBAGENT_OPTION,
):
"""Point Claude Code at the running Unsloth server and start it."""
base, key, entry = _connect(
@@ -2163,37 +2477,52 @@ def claude(
launch = launch,
)
model_id = entry["id"]
+ install_hint = (
+ "irm https://claude.ai/install.ps1 | iex"
+ if os.name == "nt"
+ else "curl -fsSL https://claude.ai/install.sh | bash"
+ )
+ if as_subagent:
+ subagent_id = _subagent_model_id(base, key, entry, model, gguf_variant)
+ subagent_model = {**entry, "id": subagent_id}
+ window = subagent_model.get("context_length") or subagent_model.get("max_context_length")
+ server_env = {
+ "UNSLOTH_CLAUDE_SUBAGENT_BASE_URL": base,
+ "UNSLOTH_CLAUDE_SUBAGENT_API_KEY": key,
+ "UNSLOTH_CLAUDE_SUBAGENT_MODEL": subagent_id,
+ "UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS": "1" if yolo else "0",
+ }
+ if window:
+ server_env["UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW"] = str(int(window))
+ with _session_config("claude-subagent", launch, persist = persist) as config:
+ plugin = write_claude_subagent_plugin(config, server_env)
+ command = [
+ "claude",
+ "--plugin-dir",
+ _agent_config_path(plugin, ["claude"]),
+ # Before ctx.args: a forwarded `--` would turn later flags positional.
+ "--allowedTools",
+ _CLAUDE_SUBAGENT_TOOL,
+ *_yolo_command_flags("claude", yolo),
+ *ctx.args,
+ ]
+ typer.echo(
+ "Unsloth is available as a local agent. "
+ "Ask Claude to spawn an Unsloth or local agent."
+ )
+ _run(
+ base,
+ subagent_model,
+ {},
+ command,
+ launch = launch,
+ install_hint = install_hint,
+ )
+ return
- env = {
- "ANTHROPIC_BASE_URL": base,
- "ANTHROPIC_AUTH_TOKEN": key,
- "ANTHROPIC_MODEL": model_id,
- # Session-only (no ~/.claude write): suppress the attribution header so
- # llama.cpp KV-cache reuse is preserved; --settings below reinforces it.
- "CLAUDE_CODE_ATTRIBUTION_HEADER": "0",
- # Update checks, beta features, and other background requests either
- # stall against a local server or evict the conversation from
- # llama-server's KV-cache slots, so turn off everything nonessential.
- "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
- "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
- # A local server streams in bursts; disable the full-screen TUI redraw so the
- # terminal doesn't flicker between tokens.
- "CLAUDE_CODE_NO_FLICKER": "1",
- }
- # Claude Code auto-compacts against its native (~600k token) window; a local
- # model's context is usually far smaller, so size the window to the loaded
- # model's real context length. Otherwise the conversation overflows the
- # server's window (silent truncation) long before Claude decides to compact.
- # codex/openclaw get the same value through their config (model_context_window
- # / contextWindow); Claude has no config file, so it rides on the env var.
- window = entry.get("context_length") or entry.get("max_context_length")
- if window:
- env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = str(int(window))
- # Compact at 90% of that window; the override only takes effect once the
- # window is set, and it can only lower the threshold, so it just guarantees
- # headroom before the server's context limit instead of relying on Claude's
- # default (which is tuned for its native 200K/1M window).
- env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] = "90"
+ env = _claude_local_env(base, key, entry)
+ # Claude Code auto-compacts against its native context window. The local env
+ # above supplies the loaded model's real window and a 90% threshold instead.
# --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions.
# IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a
# sandbox is detected, and we don't want to falsely claim one on the user's host.
@@ -2208,11 +2537,6 @@ def claude(
*_yolo_command_flags("claude", yolo),
*ctx.args,
]
- install_hint = (
- "irm https://claude.ai/install.ps1 | iex"
- if os.name == "nt"
- else "curl -fsSL https://claude.ai/install.sh | bash"
- )
_run(
base,
entry,
@@ -2237,6 +2561,7 @@ def codex(
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
+ as_subagent: bool = _AS_SUBAGENT_OPTION,
):
"""Point OpenAI Codex at the running Unsloth server and start it."""
base, key, entry = _connect(
@@ -2254,6 +2579,30 @@ def codex(
except BaseException:
_shutdown_auto_served()
raise
+ if as_subagent:
+ subagent_id = _subagent_model_id(base, key, entry, model, gguf_variant)
+ subagent_model = {**entry, "id": subagent_id}
+ with _session_config("codex-subagent", launch, persist = persist) as home:
+ agent_config = write_codex_subagent_config(base, key, subagent_model, home)
+ command = [
+ "codex",
+ *_codex_subagent_flags(agent_config),
+ *_yolo_command_flags("codex", yolo),
+ *ctx.args,
+ ]
+ typer.echo(
+ "Unsloth is available as the `unsloth` local agent. "
+ "Ask Codex to spawn an Unsloth or local agent."
+ )
+ _run(
+ base,
+ subagent_model,
+ {},
+ command,
+ launch = launch,
+ install_hint = "npm install -g @openai/codex",
+ )
+ return
command = [
"codex",
"--oss",
@@ -2283,6 +2632,7 @@ def openclaw(
persist: bool = _PERSIST_OPTION,
):
"""Point OpenClaw at the running Unsloth server and start it."""
+ _reject_as_subagent("openclaw", ctx.args)
base, key, entry = _connect(
api_key,
model,
@@ -2338,6 +2688,7 @@ def opencode(
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
+ as_subagent: bool = _AS_SUBAGENT_OPTION,
):
"""Point OpenCode at the running Unsloth server and start it."""
base, key, entry = _connect(
@@ -2347,6 +2698,50 @@ def opencode(
serve = serve,
launch = launch,
)
+ if as_subagent:
+ subagent_id = _subagent_model_id(base, key, entry, model, gguf_variant)
+ subagent_model = {**entry, "id": subagent_id}
+ # Stay append-safe for a bare no-launch recipe: a later `run ` would make
+ # `opencode --auto run ...` parse as the TUI, so keep yolo in the inline fallback.
+ route_native_auto = yolo and _opencode_supports_native_auto() and (launch or bool(ctx.args))
+ opencode_args, native_auto = _opencode_native_auto_args(list(ctx.args), route_native_auto)
+ command = ["opencode", *opencode_args]
+ with _session_config("opencode-subagent", launch, persist = persist) as cfg:
+ config_path = cfg / "opencode.json"
+ session_permission = write_opencode_config(
+ base,
+ key,
+ subagent_model,
+ config_path,
+ yolo = yolo and not native_auto,
+ as_subagent = True,
+ )
+ env = {"OPENCODE_CONFIG": str(config_path)}
+ if launch and _which_with_install_dirs("opencode") is None:
+ # Provider-filter inspection needs the binary; offer the install now so
+ # a global/project allowlist is honored on this first launch instead of
+ # being read only after _launch installs OpenCode.
+ _install_agent("opencode", "npm install -g opencode-ai")
+ inline_config = _opencode_subagent_inline_config(config_path, session_permission)
+ # A project opencode.json outranks the session file and could field-merge its
+ # own agent.unsloth over ours. Pin ours in the inline overlay so it wins.
+ inline_config.setdefault("agent", {})[_SUBAGENT_NAME] = {
+ "description": _SUBAGENT_DESCRIPTION,
+ "mode": "subagent",
+ "model": f"{_OPENCODE_PROVIDER}/{subagent_model['id']}",
+ "prompt": _SUBAGENT_INSTRUCTIONS,
+ }
+ env["OPENCODE_CONFIG_CONTENT"] = json.dumps(inline_config)
+ typer.echo("Unsloth is available as @unsloth and in /models.")
+ _run(
+ base,
+ subagent_model,
+ env,
+ command,
+ launch = launch,
+ install_hint = "npm install -g opencode-ai",
+ )
+ return
opencode_model = f"{_OPENCODE_PROVIDER}/{entry['id']}"
# The inline OPENCODE_CONFIG_CONTENT below pins the model in the highest-priority
# layer, so the session model is forced without a --model flag. Only add --model for
@@ -2433,6 +2828,7 @@ def hermes(
persist: bool = _PERSIST_OPTION,
):
"""Point Hermes (Nous Research) at the running Unsloth server and start it."""
+ _reject_as_subagent("hermes", ctx.args)
native_args = [*_yolo_command_flags("hermes", yolo), *ctx.args]
command = ["hermes", *_hermes_resume_oneshot_args(native_args)]
base, key, entry = _connect(
@@ -2464,6 +2860,7 @@ def pi(
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
+ as_subagent: bool = _AS_SUBAGENT_OPTION,
):
"""Point Pi (coding agent) at the running Unsloth server and start it."""
base, key, entry = _connect(
@@ -2473,6 +2870,37 @@ def pi(
serve = serve,
launch = launch,
)
+ install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent"
+ if as_subagent:
+ if not _PI_SUBAGENT_EXTENSION.is_file():
+ _fail(f"Missing Pi subagent extension: {_PI_SUBAGENT_EXTENSION}")
+ subagent_id = _subagent_model_id(base, key, entry, model, gguf_variant)
+ subagent_model = {**entry, "id": subagent_id}
+ extension = _agent_config_path(_PI_SUBAGENT_EXTENSION, ["pi"])
+ with _session_config("pi-subagent", launch, persist = persist) as config:
+ config_path = config / "subagent.json"
+ write_pi_subagent_config(base, key, subagent_model, config_path)
+ command = [
+ "pi",
+ "--extension",
+ extension,
+ *_yolo_command_flags("pi", yolo),
+ *ctx.args,
+ ]
+ typer.echo(
+ "Unsloth is available as a local agent and in /model. "
+ "Ask Pi to spawn an Unsloth or local agent."
+ )
+ _run(
+ base,
+ subagent_model,
+ {"UNSLOTH_PI_SUBAGENT_CONFIG": str(config_path)},
+ command,
+ launch = launch,
+ install_hint = install_hint,
+ clear_screen = True,
+ )
+ return
# Pi defaults to the google provider, so pin our provider/model on the command
# line; the custom OpenAI-compatible endpoint itself is only configurable via
# ~/.pi/agent/models.json.
@@ -2487,7 +2915,6 @@ def pi(
]
# --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs
# no install scripts), so accepting the prompt skips dependency lifecycle scripts.
- install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent"
with _session_config("pi", launch, persist = persist) as home:
# Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers
# it over $HOME/.pi/agent), so pin it at the session dir: an inherited
diff --git a/unsloth_cli/pi_subagent.ts b/unsloth_cli/pi_subagent.ts
new file mode 100644
index 0000000000..d712fc89ae
--- /dev/null
+++ b/unsloth_cli/pi_subagent.ts
@@ -0,0 +1,241 @@
+import { spawn, type ChildProcess } from "node:child_process";
+import * as fs from "node:fs";
+import * as path from "node:path";
+import { fileURLToPath } from "node:url";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+
+const provider = "unsloth";
+const maxResultCharacters = 100_000;
+const cancelGraceMilliseconds = 2_000;
+const configPath = process.env.UNSLOTH_PI_SUBAGENT_CONFIG || "";
+delete process.env.UNSLOTH_PI_SUBAGENT_CONFIG;
+let config: Record = {};
+if (configPath) {
+ try {
+ const parsed = JSON.parse(fs.readFileSync(configPath, "utf8"));
+ if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
+ throw new Error("expected a JSON object");
+ }
+ config = parsed;
+ } catch (error) {
+ throw new Error(`Could not read Unsloth subagent configuration: ${error}`);
+ }
+}
+const model = typeof config.model === "string" ? config.model : "";
+const baseUrl = typeof config.baseUrl === "string" ? config.baseUrl : "";
+const apiKey = typeof config.apiKey === "string" ? config.apiKey : "";
+const contextWindow = positiveInt(config.contextWindow, 32768);
+const maxTokens = positiveInt(config.maxTokens, Math.min(Math.floor(contextWindow / 4), 8192));
+
+function positiveInt(value: unknown, fallback: number): number {
+ const parsed = Number.parseInt(typeof value === "string" ? value : String(value || ""), 10);
+ return Number.isFinite(parsed) && parsed > 0 ? parsed : fallback;
+}
+
+function finalText(message: any): string {
+ if (message?.role !== "assistant" || !Array.isArray(message.content)) return "";
+ return message.content
+ .filter((part: any) => part?.type === "text" && typeof part.text === "string")
+ .map((part: any) => part.text)
+ .join("\n")
+ .trim();
+}
+
+function boundedResult(text: string): string {
+ if (text.length <= maxResultCharacters) return text;
+ return `${text.slice(0, maxResultCharacters)}\n\n[Local agent output truncated]`;
+}
+
+function piInvocation(args: string[]): { command: string; args: string[] } {
+ const currentScript = process.argv[1];
+ const bunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
+ if (currentScript && !bunVirtualScript && fs.existsSync(currentScript)) {
+ return { command: process.execPath, args: [currentScript, ...args] };
+ }
+ const executable = path.basename(process.execPath).toLowerCase();
+ if (!/^(node|bun)(\.exe)?$/.test(executable)) return { command: process.execPath, args };
+ return { command: "pi", args };
+}
+
+function signalProcessGroup(child: ChildProcess, signal: NodeJS.Signals): void {
+ if (!child.pid) return;
+ try {
+ process.kill(-child.pid, signal);
+ } catch {
+ try {
+ child.kill(signal);
+ } catch {
+ // The process tree already exited.
+ }
+ }
+}
+
+async function stopChildTree(child: ChildProcess): Promise {
+ if (!child.pid) return;
+ if (process.platform === "win32") {
+ await new Promise((resolve) => {
+ const killer = spawn("taskkill", ["/PID", String(child.pid), "/T", "/F"], {
+ shell: false,
+ stdio: "ignore",
+ windowsHide: true,
+ });
+ killer.once("error", () => {
+ try {
+ child.kill("SIGKILL");
+ } catch {
+ // The child already exited.
+ }
+ resolve();
+ });
+ killer.once("close", (code) => {
+ if (code !== 0) {
+ try {
+ child.kill("SIGKILL");
+ } catch {
+ // The child already exited.
+ }
+ }
+ resolve();
+ });
+ });
+ return;
+ }
+
+ signalProcessGroup(child, "SIGTERM");
+ await new Promise((resolve) => setTimeout(resolve, cancelGraceMilliseconds));
+ signalProcessGroup(child, "SIGKILL");
+}
+
+export default function unslothSubagent(pi: ExtensionAPI): void {
+ if (!model || !baseUrl || !apiKey || !configPath) {
+ throw new Error("Unsloth subagent configuration is incomplete.");
+ }
+
+ pi.registerProvider(provider, {
+ name: "Unsloth Studio",
+ baseUrl,
+ apiKey,
+ api: "openai-completions",
+ authHeader: true,
+ models: [
+ {
+ id: model,
+ name: `${model} via Unsloth`,
+ reasoning: false,
+ input: ["text"],
+ cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+ contextWindow,
+ maxTokens,
+ },
+ ],
+ });
+
+ if (process.env.UNSLOTH_PI_SUBAGENT_CHILD === "1") return;
+
+ pi.registerTool({
+ name: "unsloth_agent",
+ label: "Unsloth agent",
+ description:
+ "Local coding subagent powered by Unsloth for debugging, implementation, and codebase research. Use when the user asks to spawn an Unsloth or local agent.",
+ parameters: Type.Object({
+ task: Type.String({ description: "The complete task for the local Unsloth agent." }),
+ }),
+ async execute(_toolCallId, params, signal, _onUpdate, ctx) {
+ const extension = fileURLToPath(import.meta.url);
+ const args = [
+ "--mode",
+ "json",
+ "--print",
+ "--no-session",
+ "--provider",
+ provider,
+ "--model",
+ model,
+ "--no-extensions",
+ "--extension",
+ extension,
+ `Task: ${params.task}`,
+ ];
+ const invocation = piInvocation(args);
+ let output = "";
+ let stderr = "";
+ let lastResponse = "";
+ let childError = "";
+ let aborted = false;
+ const processLine = (line: string) => {
+ try {
+ const event = JSON.parse(line);
+ if (event.type !== "message_end") return;
+ const message = event.message;
+ // Pi reports model/API failures as message_end events while still
+ // exiting 0, so the exit status alone cannot surface them.
+ if (message?.stopReason === "error" || message?.stopReason === "aborted") {
+ childError =
+ (typeof message.errorMessage === "string" && message.errorMessage) ||
+ `The local Unsloth agent stopped: ${message.stopReason}.`;
+ return;
+ }
+ const response = finalText(message);
+ if (response) {
+ lastResponse = boundedResult(response);
+ childError = "";
+ }
+ } catch {
+ // Ignore non-JSON diagnostic lines. The exit status still reports failures.
+ }
+ };
+
+ const exitCode = await new Promise((resolve, reject) => {
+ const child = spawn(invocation.command, invocation.args, {
+ cwd: ctx.cwd,
+ detached: process.platform !== "win32",
+ shell: false,
+ stdio: ["ignore", "pipe", "pipe"],
+ env: {
+ ...process.env,
+ UNSLOTH_PI_SUBAGENT_CHILD: "1",
+ UNSLOTH_PI_SUBAGENT_CONFIG: configPath,
+ },
+ });
+ let cleanup: Promise | undefined;
+ const cancel = () => {
+ if (aborted) return;
+ aborted = true;
+ cleanup = stopChildTree(child);
+ };
+ child.on("error", (error) => {
+ signal?.removeEventListener("abort", cancel);
+ reject(error);
+ });
+ child.stdout.on("data", (chunk) => {
+ output += chunk.toString();
+ const lines = output.split("\n");
+ output = lines.pop() || "";
+ for (const line of lines) processLine(line);
+ });
+ child.stderr.on("data", (chunk) => {
+ stderr = (stderr + chunk.toString()).slice(-100_000);
+ });
+ child.on("close", async (code) => {
+ signal?.removeEventListener("abort", cancel);
+ await cleanup;
+ if (output.trim()) processLine(output);
+ resolve(code ?? 1);
+ });
+ signal?.addEventListener("abort", cancel, { once: true });
+ if (signal?.aborted) cancel();
+ });
+
+ if (aborted) throw new Error("The local Unsloth agent was cancelled.");
+ if (exitCode !== 0) {
+ throw new Error(stderr.trim() || `The local Unsloth agent exited with code ${exitCode}.`);
+ }
+ if (childError) throw new Error(boundedResult(childError));
+ return {
+ content: [{ type: "text", text: lastResponse || "The local agent returned no text." }],
+ details: { provider, model },
+ };
+ },
+ });
+}
diff --git a/unsloth_cli/tests/test_claude_subagent_mcp.py b/unsloth_cli/tests/test_claude_subagent_mcp.py
new file mode 100644
index 0000000000..13a9bd6255
--- /dev/null
+++ b/unsloth_cli/tests/test_claude_subagent_mcp.py
@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from __future__ import annotations
+
+import io
+import json
+import os
+import subprocess
+import sys
+import time
+
+import pytest
+
+import unsloth_cli.claude_subagent_mcp as bridge
+
+
+def test_protocol_lists_and_calls_local_agent():
+ initialized = bridge._response(
+ {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}},
+ )
+ assert initialized["result"]["serverInfo"]["name"] == "unsloth-local-agent"
+
+ listed = bridge._response({"jsonrpc": "2.0", "id": 2, "method": "tools/list"})
+ tool = listed["result"]["tools"][0]
+ assert tool["name"] == "unsloth_agent"
+ assert "spawn an Unsloth or local agent" in tool["description"]
+ assert tool["inputSchema"]["required"] == ["task"]
+ assert tool["_meta"]["anthropic/maxResultSizeChars"] == 100_000
+
+ called = bridge._response(
+ {
+ "jsonrpc": "2.0",
+ "id": 3,
+ "method": "tools/call",
+ "params": {"name": "unsloth_agent", "arguments": {"task": " inspect this "}},
+ },
+ run_agent = lambda task: f"completed: {task}",
+ )
+ assert called["result"] == {
+ "content": [{"type": "text", "text": "completed: inspect this"}],
+ "isError": False,
+ }
+
+
+def test_protocol_returns_tool_errors_to_parent():
+ response = bridge._response(
+ {
+ "jsonrpc": "2.0",
+ "id": 1,
+ "method": "tools/call",
+ "params": {"name": "unsloth_agent", "arguments": {"task": "test"}},
+ },
+ run_agent = lambda task: (_ for _ in ()).throw(RuntimeError("local failure")),
+ )
+ assert response["result"]["isError"] is True
+ assert response["result"]["content"][0]["text"] == "local failure"
+
+
+def test_stdio_server_ignores_notifications_and_answers_requests():
+ requests = "\n".join(
+ [
+ json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}),
+ json.dumps({"jsonrpc": "2.0", "id": 4, "method": "ping"}),
+ ]
+ )
+ output = io.StringIO()
+ bridge.serve(io.StringIO(requests), output)
+ assert json.loads(output.getvalue()) == {"jsonrpc": "2.0", "id": 4, "result": {}}
+
+
+def test_stdio_cancellation_reaches_the_running_local_agent():
+ requests = "\n".join(
+ [
+ json.dumps(
+ {
+ "jsonrpc": "2.0",
+ "id": "call-1",
+ "method": "tools/call",
+ "params": {"name": "unsloth_agent", "arguments": {"task": "wait"}},
+ }
+ ),
+ json.dumps(
+ {
+ "jsonrpc": "2.0",
+ "method": "notifications/cancelled",
+ "params": {"requestId": "call-1", "reason": "user cancelled"},
+ }
+ ),
+ ]
+ )
+ output = io.StringIO()
+ cancelled = []
+
+ def run_agent(task, cancel_event):
+ assert task == "wait"
+ assert cancel_event.wait(timeout = 1)
+ cancelled.append(task)
+ raise RuntimeError("The local Claude agent was cancelled.")
+
+ bridge.serve(io.StringIO(requests), output, run_agent = run_agent)
+ assert cancelled == ["wait"]
+ assert output.getvalue() == ""
+
+
+def test_stdio_sigint_stops_the_running_local_agent(monkeypatch):
+ request = json.dumps(
+ {
+ "jsonrpc": "2.0",
+ "id": "call-1",
+ "method": "tools/call",
+ "params": {"name": "unsloth_agent", "arguments": {"task": "wait"}},
+ }
+ )
+ handlers = {}
+ started = bridge.threading.Event()
+ cancelled = []
+
+ def set_handler(signum, handler):
+ previous = handlers.get(signum, bridge.signal.SIG_DFL)
+ handlers[signum] = handler
+ return previous
+
+ monkeypatch.setattr(bridge.signal, "signal", set_handler)
+
+ class InterruptingInput:
+ def __init__(self):
+ self.sent = False
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ if not self.sent:
+ self.sent = True
+ return request + "\n"
+ assert started.wait(timeout = 1)
+ handlers[bridge.signal.SIGINT](bridge.signal.SIGINT, None)
+ raise AssertionError("SIGINT handler must unwind the stdin loop")
+
+ def run_agent(task, cancel_event):
+ assert task == "wait"
+ started.set()
+ assert cancel_event.wait(timeout = 1)
+ # Real Claude Code sends SIGINT twice. The second one must not abort cleanup.
+ handlers[bridge.signal.SIGINT](bridge.signal.SIGINT, None)
+ cancelled.append(task)
+ raise RuntimeError("The local Claude agent was cancelled.")
+
+ output = io.StringIO()
+ bridge.serve(InterruptingInput(), output, run_agent = run_agent)
+ assert cancelled == ["wait"]
+ assert output.getvalue() == ""
+
+
+@pytest.mark.parametrize(
+ ("bypass", "permission"),
+ [("0", "acceptEdits"), ("1", "bypassPermissions")],
+)
+def test_local_child_uses_unsloth_without_overwriting_parent_auth(
+ monkeypatch, tmp_path, bypass, permission
+):
+ captured = {}
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888")
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test")
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M")
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW", "32768")
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS", bypass)
+ monkeypatch.setenv("ANTHROPIC_API_KEY", "cloud-key")
+ monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "cloud-oauth")
+ monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path))
+ monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude")
+ monkeypatch.setattr(bridge, "_claude_flags", lambda model: ["--settings", "{}"])
+
+ class Process:
+ pid = 1234
+ returncode = 0
+
+ def communicate(self, timeout):
+ captured["timeout"] = timeout
+ return json.dumps({"is_error": False, "result": "LOCAL_OK"}), ""
+
+ def poll(self):
+ return self.returncode
+
+ def popen(command, **kwargs):
+ captured["command"] = command
+ captured.update(kwargs)
+ return Process()
+
+ monkeypatch.setattr(bridge.subprocess, "Popen", popen)
+ assert bridge.run_local_agent("reply exactly LOCAL_OK") == "LOCAL_OK"
+ command = captured["command"]
+ assert command[:3] == ["/usr/local/bin/claude", "--model", "unsloth/model-GGUF:Q4_K_M"]
+ assert command[command.index("--permission-mode") + 1] == permission
+ assert "--no-session-persistence" in command
+ assert captured["cwd"] == str(tmp_path)
+ assert captured["stdin"] is bridge.subprocess.DEVNULL
+ assert captured["stdout"] is bridge.subprocess.PIPE
+ assert captured["stderr"] is bridge.subprocess.PIPE
+ if os.name == "nt":
+ assert captured["creationflags"] == bridge.subprocess.CREATE_NEW_PROCESS_GROUP
+ else:
+ assert captured["start_new_session"] is True
+ child_env = captured["env"]
+ assert child_env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8888"
+ assert child_env["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-test"
+ assert child_env["ANTHROPIC_MODEL"] == "unsloth/model-GGUF:Q4_K_M"
+ assert child_env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] == "32768"
+ assert child_env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] == "90"
+ assert "ANTHROPIC_API_KEY" not in child_env
+ assert "CLAUDE_CODE_OAUTH_TOKEN" not in child_env
+
+
+def test_local_child_process_is_stopped_on_cancellation(monkeypatch, tmp_path):
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888")
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test")
+ monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M")
+ monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path))
+ monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude")
+ monkeypatch.setattr(bridge, "_claude_flags", lambda model: [])
+ cancel_event = bridge.threading.Event()
+ stopped = []
+
+ class Process:
+ pid = 1234
+ returncode = None
+
+ def communicate(self, timeout):
+ cancel_event.set()
+ raise bridge.subprocess.TimeoutExpired("claude", timeout)
+
+ def poll(self):
+ return self.returncode
+
+ process = Process()
+ monkeypatch.setattr(bridge.subprocess, "Popen", lambda *args, **kwargs: process)
+
+ def stop(child):
+ stopped.append(child)
+ child.returncode = -15
+
+ monkeypatch.setattr(bridge, "_stop_child", stop)
+ with pytest.raises(RuntimeError, match = "cancelled"):
+ bridge.run_local_agent("wait", cancel_event)
+ assert stopped == [process]
+
+
+def test_windows_cancellation_stops_the_child_process_tree(monkeypatch):
+ monkeypatch.setattr(bridge.os, "name", "nt")
+ captured = {}
+
+ class Process:
+ pid = 4321
+ returncode = None
+
+ def poll(self):
+ return self.returncode
+
+ def wait(self, timeout = None):
+ captured["wait_timeout"] = timeout
+ self.returncode = 1
+
+ def terminate(self):
+ raise AssertionError("taskkill should handle the process tree")
+
+ def run(command, **kwargs):
+ captured["command"] = command
+ captured.update(kwargs)
+ return bridge.subprocess.CompletedProcess(command, 0)
+
+ monkeypatch.setattr(bridge.subprocess, "run", run)
+ bridge._stop_child(Process())
+
+ assert captured["command"] == ["taskkill", "/PID", "4321", "/T", "/F"]
+ assert captured["capture_output"] is True
+ assert captured["check"] is False
+ assert captured["wait_timeout"] == bridge._CANCEL_GRACE_SECONDS
+
+
+def test_windows_failed_taskkill_still_terminates_the_child(monkeypatch):
+ monkeypatch.setattr(bridge.os, "name", "nt")
+ captured = {}
+
+ class Process:
+ pid = 4321
+ returncode = None
+
+ def poll(self):
+ return self.returncode
+
+ def wait(self, timeout = None):
+ self.returncode = 1
+
+ def terminate(self):
+ captured["terminated"] = True
+ self.returncode = 1
+
+ monkeypatch.setattr(
+ bridge.subprocess,
+ "run",
+ lambda command, **kwargs: bridge.subprocess.CompletedProcess(command, 1),
+ )
+ bridge._stop_child(Process())
+
+ assert captured.get("terminated") is True
+
+
+@pytest.mark.skipif(os.name == "nt", reason = "POSIX process groups")
+def test_stop_child_kills_survivors_after_leader_exit(monkeypatch, tmp_path):
+ monkeypatch.setattr(bridge, "_CANCEL_GRACE_SECONDS", 0.2)
+ marker = tmp_path / "grandchild-survived"
+ grandchild = (
+ "import pathlib, sys, time; time.sleep(1.0); "
+ "pathlib.Path(sys.argv[1]).write_text('alive')"
+ )
+ process = subprocess.Popen(
+ [
+ sys.executable,
+ "-c",
+ "import subprocess, sys; "
+ "subprocess.Popen([sys.executable, '-c', sys.argv[1], sys.argv[2]])",
+ grandchild,
+ str(marker),
+ ],
+ start_new_session = True,
+ )
+ process.wait()
+
+ bridge._stop_child(process)
+
+ time.sleep(1.2)
+ assert not marker.exists()
+
+
+def test_result_parser_accepts_diagnostics_before_json():
+ output = "connector warning\n" + json.dumps({"is_error": False, "result": "OK"})
+ assert bridge._result_text(output) == "OK"
diff --git a/unsloth_cli/tests/test_pi_subagent.py b/unsloth_cli/tests/test_pi_subagent.py
new file mode 100644
index 0000000000..beac6770df
--- /dev/null
+++ b/unsloth_cli/tests/test_pi_subagent.py
@@ -0,0 +1,191 @@
+# 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 os
+from pathlib import Path
+import json
+import shutil
+import subprocess
+
+import pytest
+
+
+@pytest.mark.skipif(os.name == "nt", reason = "POSIX process group regression test")
+def test_pi_cancel_kills_child_process_group(tmp_path):
+ bun = shutil.which("bun")
+ if bun is None:
+ pytest.skip("Bun is required to execute the bundled Pi extension")
+
+ ready = tmp_path / "grandchild-ready"
+ marker = tmp_path / "grandchild-survived"
+ config = tmp_path / "subagent.json"
+ config.write_text(
+ json.dumps(
+ {
+ "baseUrl": "http://127.0.0.1:8000/v1",
+ "apiKey": "private-token",
+ "model": "local-model",
+ "contextWindow": 32768,
+ "maxTokens": 8192,
+ }
+ ),
+ encoding = "utf-8",
+ )
+ driver = tmp_path / "pi-driver.js"
+ driver.write_text(
+ """
+import { spawn } from "node:child_process";
+
+spawn(
+ process.execPath,
+ [
+ "-e",
+ `
+ const fs = require("node:fs");
+ process.on("SIGTERM", () => {});
+ fs.writeFileSync(process.env.PI_CHILD_READY, "ready");
+ setTimeout(() => fs.writeFileSync(process.env.PI_CANCEL_MARKER, "alive"), 3000);
+ setInterval(() => {}, 1000);
+ `,
+ ],
+ { stdio: "inherit" },
+);
+process.on("SIGTERM", () => {});
+setInterval(() => {}, 1000);
+""",
+ encoding = "utf-8",
+ )
+ extension = Path(__file__).parents[1] / "pi_subagent.ts"
+ test_file = tmp_path / "pi-cancel.test.ts"
+ test_file.write_text(
+ f"""
+import {{ expect, mock, test }} from "bun:test";
+import {{ existsSync }} from "node:fs";
+import {{ pathToFileURL }} from "node:url";
+
+mock.module("typebox", () => ({{
+ Type: {{ Object: (value) => value, String: (value) => value }},
+}}));
+
+test("cancellation stops the Pi child process group", async () => {{
+ process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r};
+ process.env.PI_CHILD_READY = {str(ready)!r};
+ process.env.PI_CANCEL_MARKER = {str(marker)!r};
+ process.argv[1] = {str(driver)!r};
+
+ const loaded = await import(pathToFileURL({str(extension)!r}).href);
+ let tool;
+ let provider;
+ loaded.default({{
+ registerProvider(_name, value) {{ provider = value; }},
+ registerTool(value) {{ tool = value; }},
+ }});
+ expect(process.env.UNSLOTH_PI_SUBAGENT_CONFIG).toBeUndefined();
+ expect(process.env.UNSLOTH_PI_SUBAGENT_API_KEY).toBeUndefined();
+ expect(provider.apiKey).toBe("private-token");
+
+ const controller = new AbortController();
+ const execution = tool.execute(
+ "call",
+ {{ task: "wait" }},
+ controller.signal,
+ undefined,
+ {{ cwd: {str(tmp_path)!r} }},
+ );
+ for (let attempt = 0; attempt < 100 && !existsSync({str(ready)!r}); attempt++) {{
+ await Bun.sleep(20);
+ }}
+ expect(existsSync({str(ready)!r})).toBe(true);
+ controller.abort();
+ await expect(execution).rejects.toThrow("cancelled");
+ await Bun.sleep(3200);
+ expect(existsSync({str(marker)!r})).toBe(false);
+}}, 10_000);
+""",
+ encoding = "utf-8",
+ )
+
+ completed = subprocess.run(
+ [bun, "test", str(test_file)],
+ capture_output = True,
+ text = True,
+ timeout = 15,
+ )
+
+ assert completed.returncode == 0, completed.stdout + completed.stderr
+
+
+@pytest.mark.skipif(os.name == "nt", reason = "POSIX driver script")
+def test_pi_child_error_events_fail_the_tool_call(tmp_path):
+ bun = shutil.which("bun")
+ if bun is None:
+ pytest.skip("Bun is required to execute the bundled Pi extension")
+
+ config = tmp_path / "subagent.json"
+ config.write_text(
+ json.dumps(
+ {
+ "baseUrl": "http://127.0.0.1:8000/v1",
+ "apiKey": "private-token",
+ "model": "local-model",
+ "contextWindow": 32768,
+ "maxTokens": 8192,
+ }
+ ),
+ encoding = "utf-8",
+ )
+ # Pi reports model/API failures as message_end events while exiting 0.
+ driver = tmp_path / "pi-driver.js"
+ driver.write_text(
+ """
+const event = {
+ type: "message_end",
+ message: { role: "assistant", stopReason: "error", errorMessage: "backend unreachable", content: [] },
+};
+console.log(JSON.stringify(event));
+""",
+ encoding = "utf-8",
+ )
+ extension = Path(__file__).parents[1] / "pi_subagent.ts"
+ test_file = tmp_path / "pi-error.test.ts"
+ test_file.write_text(
+ f"""
+import {{ expect, mock, test }} from "bun:test";
+import {{ pathToFileURL }} from "node:url";
+
+mock.module("typebox", () => ({{
+ Type: {{ Object: (value) => value, String: (value) => value }},
+}}));
+
+test("child error events fail the tool call", async () => {{
+ process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r};
+ process.argv[1] = {str(driver)!r};
+
+ const loaded = await import(pathToFileURL({str(extension)!r}).href);
+ let tool;
+ loaded.default({{
+ registerProvider() {{}},
+ registerTool(value) {{ tool = value; }},
+ }});
+
+ const execution = tool.execute(
+ "call",
+ {{ task: "fail" }},
+ undefined,
+ undefined,
+ {{ cwd: {str(tmp_path)!r} }},
+ );
+ await expect(execution).rejects.toThrow("backend unreachable");
+}}, 10_000);
+""",
+ encoding = "utf-8",
+ )
+
+ completed = subprocess.run(
+ [bun, "test", str(test_file)],
+ capture_output = True,
+ text = True,
+ timeout = 15,
+ )
+
+ assert completed.returncode == 0, completed.stdout + completed.stderr
diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py
index 7c070fa5f4..34f25c5ee5 100644
--- a/unsloth_cli/tests/test_start.py
+++ b/unsloth_cli/tests/test_start.py
@@ -619,6 +619,104 @@ def test_write_codex_config_omits_catalog_for_old_codex(tmp_path, monkeypatch):
assert not (tmp_path / "model-catalog.json").exists()
+def test_write_codex_subagent_config_keeps_parent_model_out(tmp_path, monkeypatch):
+ monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True)
+ local = {**MODEL, "id": MODEL["id"] + ":UD-Q4_K_XL"}
+ path = start.write_codex_subagent_config(BASE, "private-token", local, tmp_path)
+ agent = _parse_toml(path.read_text())
+ assert agent["name"] == "unsloth"
+ assert "local agent" in agent["description"].lower()
+ assert agent["model_provider"] == start._CODEX_PROFILE
+ assert agent["model"] == local["id"]
+ assert agent["model_context_window"] == MODEL["context_length"]
+ assert agent["model_providers"][start._CODEX_PROFILE] == {
+ "name": "Unsloth Studio",
+ "base_url": f"{BASE}/v1",
+ "wire_api": "responses",
+ "auth": {
+ "command": sys.executable,
+ "args": [
+ "-c",
+ "import json,sys; print(json.load(open(sys.argv[1], encoding='utf-8'))['token'])",
+ str(tmp_path / "unsloth-auth.json"),
+ ],
+ "timeout_ms": 5000,
+ },
+ }
+ assert json.loads((tmp_path / "unsloth-auth.json").read_text()) == {"token": "private-token"}
+ catalog = json.loads((tmp_path / agent["model_catalog_json"]).read_text())
+ assert catalog["models"][0]["slug"] == local["id"]
+
+
+@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario")
+def test_codex_subagent_auth_uses_wsl_for_windows_codex(monkeypatch, tmp_path):
+ monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
+ monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: False)
+ monkeypatch.setattr(
+ start.shutil,
+ "which",
+ lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/codex.exe",
+ )
+
+ path = start.write_codex_subagent_config(BASE, "private-token", MODEL, tmp_path)
+ auth = _parse_toml(path.read_text())["model_providers"][start._CODEX_PROFILE]["auth"]
+
+ assert auth["command"] == "wsl.exe"
+ assert auth["args"][:5] == ["-d", "Ubuntu", "--", sys.executable, "-c"]
+ assert auth["args"][-1] == str(tmp_path / "unsloth-auth.json")
+
+
+@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario")
+def test_agent_config_path_translates_for_windows_agent(monkeypatch, tmp_path):
+ windows_path = r"\\wsl.localhost\Ubuntu\tmp\unsloth.toml"
+ monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
+ monkeypatch.setattr(
+ start.shutil,
+ "which",
+ lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/codex",
+ )
+ monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: windows_path)
+
+ assert start._agent_config_path(tmp_path / "unsloth.toml", ["codex"]) == windows_path
+
+
+def test_subagent_model_id_preserves_explicit_variant(monkeypatch):
+ monkeypatch.setattr(
+ start,
+ "_http_json",
+ lambda *args, **kwargs: pytest.fail("explicit variant should not need status"),
+ )
+ assert (
+ start._subagent_model_id(BASE, "key", MODEL, MODEL["id"], "UD-Q4_K_XL")
+ == MODEL["id"] + ":UD-Q4_K_XL"
+ )
+
+
+def test_subagent_model_id_uses_loaded_variant(monkeypatch):
+ monkeypatch.setattr(
+ start,
+ "_http_json",
+ lambda *args, **kwargs: {"is_gguf": True, "gguf_variant": "Q5_K_M"},
+ )
+ assert start._subagent_model_id(BASE, "key", MODEL, None, None) == MODEL["id"] + ":Q5_K_M"
+
+
+def test_subagent_model_id_warns_when_status_unavailable(monkeypatch, capsys):
+ def raise_error(*args, **kwargs):
+ raise OSError("connection refused")
+
+ monkeypatch.setattr(start, "_http_json", raise_error)
+ assert start._subagent_model_id(BASE, "key", MODEL, None, None) == MODEL["id"]
+ assert "could not verify the loaded GGUF variant" in capsys.readouterr().err
+
+
+@pytest.mark.parametrize("agent", ["openclaw", "hermes"])
+def test_unsupported_agents_reject_as_subagent(agent):
+ result = CliRunner().invoke(start.start_app, [agent, "--as-subagent"])
+ assert result.exit_code == 1
+ assert f"--as-subagent is not supported for {agent}." in result.output
+
+
@pytest.fixture()
def fake_studio(tmp_path, monkeypatch):
calls = []
@@ -690,6 +788,82 @@ def test_connect_claude_no_launch(fake_studio):
assert ".claude/settings.json" not in result.output
+def test_connect_claude_as_subagent_preserves_cloud_parent(fake_studio, tmp_path):
+ result = CliRunner().invoke(
+ start.start_app,
+ [
+ "claude",
+ "--as-subagent",
+ "--no-launch",
+ "--model",
+ MODEL["id"] + ":UD-Q4_K_XL",
+ "hello",
+ ],
+ )
+ assert result.exit_code == 0, result.output
+ command = _launch_command(result.output)
+ plugin = tmp_path / "agents" / "claude-subagent" / "unsloth-local-agent"
+ assert command == [
+ "claude",
+ "--plugin-dir",
+ str(plugin),
+ "--allowedTools",
+ start._CLAUDE_SUBAGENT_TOOL,
+ "hello",
+ ]
+ assert "--model" not in command
+ parent_base = "$env:ANTHROPIC_BASE_URL" if os.name == "nt" else "export ANTHROPIC_BASE_URL="
+ parent_token = (
+ "$env:ANTHROPIC_AUTH_TOKEN" if os.name == "nt" else "export ANTHROPIC_AUTH_TOKEN="
+ )
+ assert parent_base not in result.output
+ assert parent_token not in result.output
+ assert "unset ANTHROPIC_API_KEY" not in result.output
+ assert "UNSLOTH_CLAUDE_SUBAGENT_API_KEY" not in result.output
+ assert "sk-unsloth-feedfacefeedface" not in result.output
+ assert json.loads((plugin / ".claude-plugin" / "plugin.json").read_text())["name"] == (
+ "unsloth-local-agent"
+ )
+ mcp = json.loads((plugin / ".mcp.json").read_text())["mcpServers"]["unsloth"]
+ assert mcp["command"] == sys.executable
+ assert mcp["args"] == ["-m", start._CLAUDE_SUBAGENT_MCP_MODULE]
+ assert mcp["env"] == {
+ "UNSLOTH_CLAUDE_SUBAGENT_BASE_URL": BASE,
+ "UNSLOTH_CLAUDE_SUBAGENT_API_KEY": "sk-unsloth-feedfacefeedface",
+ "UNSLOTH_CLAUDE_SUBAGENT_MODEL": MODEL["id"] + ":UD-Q4_K_XL",
+ "UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS": "0",
+ "UNSLOTH_CLAUDE_SUBAGENT_CONTEXT_WINDOW": "4096",
+ }
+ skill = (plugin / "skills" / "local-agent" / "SKILL.md").read_text()
+ assert "spawn an Unsloth agent or local agent" in skill
+ assert "Ask Claude to spawn an Unsloth or local agent." in result.output
+
+
+@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario")
+def test_claude_subagent_plugin_uses_wsl_for_windows_claude(monkeypatch, tmp_path):
+ monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
+ monkeypatch.setenv("WSLENV", "EXISTING")
+ monkeypatch.setattr(
+ start.shutil,
+ "which",
+ lambda _: "/mnt/c/Users/x/AppData/Local/Programs/claude.exe",
+ )
+ server_env = {"UNSLOTH_CLAUDE_SUBAGENT_API_KEY": "secret"}
+ plugin = start.write_claude_subagent_plugin(tmp_path, server_env)
+ mcp = json.loads((plugin / ".mcp.json").read_text())["mcpServers"]["unsloth"]
+ assert mcp["command"] == "wsl.exe"
+ assert mcp["args"] == [
+ "-d",
+ "Ubuntu",
+ "--",
+ sys.executable,
+ "-m",
+ start._CLAUDE_SUBAGENT_MCP_MODULE,
+ ]
+ assert mcp["env"]["UNSLOTH_CLAUDE_SUBAGENT_API_KEY"] == "secret"
+ assert mcp["env"]["WSLENV"].split(":") == ["EXISTING", "UNSLOTH_CLAUDE_SUBAGENT_API_KEY"]
+
+
def test_connect_claude_compact_window_omitted_without_context(fake_studio, monkeypatch):
# A model that doesn't report a context length -> leave Claude's default window
# rather than guessing one.
@@ -814,6 +988,38 @@ def test_connect_codex_no_launch(fake_studio, tmp_path):
assert (home / "unsloth_api.config.toml").exists()
+def test_connect_codex_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, monkeypatch):
+ monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True)
+ result = CliRunner().invoke(
+ start.start_app,
+ [
+ "codex",
+ "--as-subagent",
+ "--no-launch",
+ "--model",
+ MODEL["id"] + ":UD-Q4_K_XL",
+ ],
+ )
+ assert result.exit_code == 0, result.output
+ command = _launch_command(result.output)
+ assert command[0] == "codex"
+ assert command[1:3] == ["--enable", "multi_agent"]
+ assert "agents.max_depth=1" in command
+ assert "--oss" not in command
+ assert "--profile" not in command
+ assert "--model" not in command
+ assert "CODEX_HOME" not in result.output
+ assert start._CODEX_ENV_KEY not in result.output
+ assert "sk-unsloth-feedfacefeedface" not in result.output
+ home = tmp_path / "agents" / "codex-subagent"
+ agent_path = home / "unsloth.toml"
+ agent = _parse_toml(agent_path.read_text())
+ assert agent["model"] == MODEL["id"] + ":UD-Q4_K_XL"
+ assert "env_key" not in agent["model_providers"][start._CODEX_PROFILE]
+ assert f"agents.unsloth.config_file={json.dumps(str(agent_path))}" in command
+ assert "Ask Codex to spawn an Unsloth or local agent." in result.output
+
+
def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, tmp_path):
result = CliRunner().invoke(
start.start_app,
@@ -2467,8 +2673,7 @@ def test_write_opencode_config_fresh(tmp_path):
MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}}
}
assert config["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"
- # The overlay never writes disabled_providers; the dedicated provider id is one a
- # user's disable list would not target, so nothing needs re-enabling.
+ # Provider filters belong to the launch-time inline overlay, not this config writer.
assert "disabled_providers" not in config
# Compaction buffer scaled to ~10% of the window (compact near 90%).
assert config["compaction"] == {"auto": True, "reserved": 131072 // 10}
@@ -2509,6 +2714,109 @@ def test_write_opencode_config_keeps_foreign_disabled_providers(tmp_path):
assert config["disabled_providers"] == ["openai", "gemini"]
+def test_write_opencode_config_as_subagent_preserves_parent_model(tmp_path):
+ path = tmp_path / "opencode.json"
+ path.write_text(
+ json.dumps(
+ {
+ "model": "anthropic/claude-sonnet-4-5",
+ "small_model": "anthropic/claude-haiku-4-5",
+ "compaction": {"auto": False},
+ }
+ )
+ )
+ local = {**MODEL, "id": MODEL["id"] + ":UD-Q4_K_XL"}
+ start.write_opencode_config(
+ BASE,
+ "sk-unsloth-abc",
+ local,
+ path,
+ as_subagent = True,
+ )
+ config = json.loads(path.read_text())
+ assert config["model"] == "anthropic/claude-sonnet-4-5"
+ assert config["small_model"] == "anthropic/claude-haiku-4-5"
+ assert config["compaction"] == {"auto": False}
+ agent = config["agent"]["unsloth"]
+ assert agent["mode"] == "subagent"
+ assert agent["model"] == f"{start._OPENCODE_PROVIDER}/{local['id']}"
+ assert "local agent" in agent["description"].lower()
+ assert local["id"] in config["provider"][start._OPENCODE_PROVIDER]["models"]
+
+
+def test_opencode_subagent_inline_keeps_parent_provider_filters(monkeypatch, tmp_path):
+ config_path = tmp_path / "opencode.json"
+ inherited = {"theme": "tokyonight"}
+ monkeypatch.setenv("OPENCODE_CONFIG_CONTENT", json.dumps(inherited))
+ monkeypatch.setattr(start, "_which_with_install_dirs", lambda _: "/usr/bin/opencode")
+ captured = {}
+
+ def run(command, **kwargs):
+ captured["command"] = command
+ captured.update(kwargs)
+ return SimpleNamespace(
+ returncode = 0,
+ stdout = json.dumps(
+ {
+ "enabled_providers": ["opencode-go"],
+ "disabled_providers": ["ollama", start._OPENCODE_PROVIDER],
+ "subagent_depth": 0,
+ }
+ ),
+ stderr = "",
+ )
+
+ monkeypatch.setattr(start.subprocess, "run", run)
+ permission = {"edit": "allow"}
+ inline = start._opencode_subagent_inline_config(config_path, permission)
+
+ assert captured["command"] == ["/usr/bin/opencode", "debug", "config"]
+ assert captured["env"]["OPENCODE_CONFIG"] == str(config_path)
+ assert inline == {
+ "theme": "tokyonight",
+ "enabled_providers": ["opencode-go", start._OPENCODE_PROVIDER],
+ "disabled_providers": ["ollama"],
+ "subagent_depth": 1,
+ "permission": permission,
+ }
+
+
+def test_opencode_subagent_inline_preserves_positive_depth(monkeypatch, tmp_path):
+ monkeypatch.setattr(start, "_which_with_install_dirs", lambda _: "/usr/bin/opencode")
+ monkeypatch.setattr(
+ start.subprocess,
+ "run",
+ lambda *args, **kwargs: SimpleNamespace(
+ returncode = 0,
+ stdout = json.dumps({"subagent_depth": 3}),
+ stderr = "",
+ ),
+ )
+
+ inline = start._opencode_subagent_inline_config(tmp_path / "opencode.json", {})
+
+ assert inline["subagent_depth"] == 3
+
+
+def test_opencode_subagent_inline_merges_inherited_filters_without_binary(monkeypatch, tmp_path):
+ monkeypatch.setenv(
+ "OPENCODE_CONFIG_CONTENT",
+ json.dumps(
+ {
+ "enabled_providers": ["opencode-go"],
+ "disabled_providers": ["ollama", start._OPENCODE_PROVIDER],
+ }
+ ),
+ )
+ monkeypatch.setattr(start, "_which_with_install_dirs", lambda _: None)
+
+ inline = start._opencode_subagent_inline_config(tmp_path / "opencode.json", {})
+
+ assert inline["enabled_providers"] == ["opencode-go", start._OPENCODE_PROVIDER]
+ assert inline["disabled_providers"] == ["ollama"]
+ assert inline["subagent_depth"] == 1
+
+
def _opencode_inline_config(output: str) -> dict:
# --no-launch prints OPENCODE_CONFIG_CONTENT as a POSIX `export NAME=`
# line on Unix/WSL and a PowerShell `$env:NAME = ""` line on native Windows;
@@ -2597,6 +2905,130 @@ def test_connect_opencode_no_launch(fake_studio, tmp_path):
assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
+def test_connect_opencode_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, monkeypatch):
+ monkeypatch.setattr(start, "_opencode_subagent_inline_config", lambda path, permission: {})
+ result = CliRunner().invoke(
+ start.start_app,
+ [
+ "opencode",
+ "--as-subagent",
+ "--no-launch",
+ "--model",
+ MODEL["id"] + ":UD-Q4_K_XL",
+ ],
+ )
+ assert result.exit_code == 0, result.output
+ assert _launch_command(result.output) == ["opencode"]
+ expected_model = f"{start._OPENCODE_PROVIDER}/{MODEL['id']}:UD-Q4_K_XL"
+ # The agent rides in the inline overlay; nothing else comes from the empty base.
+ assert _opencode_inline_config(result.output) == {
+ "agent": {
+ "unsloth": {
+ "description": start._SUBAGENT_DESCRIPTION,
+ "mode": "subagent",
+ "model": expected_model,
+ "prompt": start._SUBAGENT_INSTRUCTIONS,
+ }
+ }
+ }
+ path = tmp_path / "agents" / "opencode-subagent" / "opencode.json"
+ config = json.loads(path.read_text())
+ assert "model" not in config
+ assert "small_model" not in config
+ assert "compaction" not in config
+ agent = config["agent"]["unsloth"]
+ assert agent["model"] == expected_model
+ assert "Unsloth is available as @unsloth and in /models." in result.output
+
+
+def test_claude_subagent_allowed_tools_precede_forwarded_delimiter(fake_studio):
+ # A forwarded `--` makes everything after it positional; the tool pre-approval
+ # must be parsed as an option, so it rides before ctx.args.
+ result = CliRunner().invoke(
+ start.start_app,
+ ["claude", "--as-subagent", "--no-launch", "--", "--resume", "abc123"],
+ )
+ assert result.exit_code == 0, result.output
+ command = _launch_command(result.output)
+ assert command.index("--allowedTools") < command.index("--resume")
+
+
+def test_opencode_subagent_installs_binary_before_filter_inspection(fake_studio, monkeypatch):
+ # The effective-config inspection needs the opencode binary; a first launch must
+ # offer the install before building the overlay, or a global allowlist read only
+ # after _launch installs OpenCode would filter out the new provider.
+ installed = {}
+ monkeypatch.setattr(
+ start,
+ "_which_with_install_dirs",
+ lambda name: "/usr/local/bin/opencode" if installed.get("done") else None,
+ )
+
+ def install(name, hint):
+ installed["done"] = True
+ installed["name"] = name
+ return "/usr/local/bin/opencode"
+
+ monkeypatch.setattr(start, "_install_agent", install)
+ inspected = {}
+
+ def inline(path, permission):
+ inspected["binary"] = start._which_with_install_dirs("opencode")
+ return {}
+
+ monkeypatch.setattr(start, "_opencode_subagent_inline_config", inline)
+ monkeypatch.setattr(start, "_run", lambda *a, **k: None)
+
+ result = CliRunner().invoke(start.start_app, ["opencode", "--as-subagent"])
+
+ assert result.exit_code == 0, result.output
+ assert installed["name"] == "opencode"
+ assert inspected["binary"] == "/usr/local/bin/opencode"
+
+
+def test_opencode_subagent_pins_agent_in_inline_overlay(fake_studio, monkeypatch):
+ # A project opencode.json outranks the session file, so the agent must ride in
+ # OPENCODE_CONFIG_CONTENT where a repo's own agent.unsloth cannot field-merge over it.
+ monkeypatch.setattr(start, "_opencode_subagent_inline_config", lambda path, permission: {})
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--as-subagent", "--no-launch", "--model", MODEL["id"] + ":UD-Q4_K_XL"],
+ )
+ assert result.exit_code == 0, result.output
+ agent = _opencode_inline_config(result.output)["agent"]["unsloth"]
+ assert agent["mode"] == "subagent"
+ assert agent["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}:UD-Q4_K_XL"
+ assert agent["prompt"] == start._SUBAGENT_INSTRUCTIONS
+ assert agent["description"] == start._SUBAGENT_DESCRIPTION
+
+
+def test_connect_opencode_subagent_yolo_no_launch_stays_append_safe(fake_studio, monkeypatch):
+ monkeypatch.setattr(start, "_opencode_supports_native_auto", lambda: True)
+ captured = {}
+
+ def inline(path, permission):
+ captured["permission"] = permission
+ return {"permission": permission}
+
+ monkeypatch.setattr(start, "_opencode_subagent_inline_config", inline)
+ result = CliRunner().invoke(
+ start.start_app,
+ ["opencode", "--as-subagent", "--no-launch", "--yolo"],
+ )
+
+ assert result.exit_code == 0, result.output
+ assert _launch_command(result.output) == ["opencode"]
+ assert "--auto" not in result.output
+ assert captured["permission"] == {
+ "edit": "allow",
+ "bash": "allow",
+ "webfetch": "allow",
+ "task": "allow",
+ "external_directory": {"*": "allow"},
+ }
+ assert _opencode_inline_config(result.output)["permission"] == captured["permission"]
+
+
# ── Hermes (OpenAI /v1/chat/completions, key via env) ────────────────
@@ -2739,6 +3171,39 @@ def test_connect_pi_no_launch(fake_studio, tmp_path):
assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
+def test_connect_pi_as_subagent_preserves_cloud_parent(fake_studio, tmp_path):
+ result = CliRunner().invoke(
+ start.start_app,
+ [
+ "pi",
+ "--as-subagent",
+ "--no-launch",
+ "--model",
+ MODEL["id"] + ":UD-Q4_K_XL",
+ ],
+ )
+ assert result.exit_code == 0, result.output
+ command = _launch_command(result.output)
+ assert command[:2] == ["pi", "--extension"]
+ assert command[2].endswith("unsloth_cli/pi_subagent.ts")
+ assert "--provider" not in command
+ assert "--model" not in command
+ assert "PI_CODING_AGENT_DIR" not in result.output
+ assert "export HOME=" not in result.output
+ assert "UNSLOTH_PI_SUBAGENT_API_KEY" not in result.output
+ assert "sk-unsloth-feedfacefeedface" not in result.output
+ config_path = tmp_path / "agents" / "pi-subagent" / "subagent.json"
+ _assert_env_set(result.output, "UNSLOTH_PI_SUBAGENT_CONFIG", str(config_path))
+ assert json.loads(config_path.read_text()) == {
+ "baseUrl": f"{BASE}/v1",
+ "apiKey": "sk-unsloth-feedfacefeedface",
+ "model": MODEL["id"] + ":UD-Q4_K_XL",
+ "contextWindow": 4096,
+ "maxTokens": 1024,
+ }
+ assert "Ask Pi to spawn an Unsloth or local agent." in result.output
+
+
def test_connect_pi_no_launch_windows_relocates_userprofile(fake_studio, tmp_path, monkeypatch):
# On native Windows Node resolves ~/.pi via USERPROFILE, not HOME, so the session
# must point USERPROFILE at the relocated home or Pi reads the user's real ~/.pi.
@@ -3282,6 +3747,27 @@ def test_opencode_non_yolo_flips_only_explicit_allow(tmp_path):
assert session == {} # a non-yolo session carries no permission inline
+def test_opencode_subagent_non_yolo_clears_yolo_task_permission(tmp_path):
+ path = tmp_path / "opencode.json"
+ start.write_opencode_config(
+ BASE,
+ "sk-unsloth-abc",
+ MODEL,
+ path,
+ yolo = True,
+ as_subagent = True,
+ )
+ start.write_opencode_config(
+ BASE,
+ "sk-unsloth-abc",
+ MODEL,
+ path,
+ as_subagent = True,
+ )
+
+ assert json.loads(path.read_text())["permission"]["task"] == "ask"
+
+
def test_opencode_non_yolo_leaves_string_permission(tmp_path):
# A global string rule ("deny") is a user-managed catch-all; leave it untouched and
# carry no inline override.
From 84b762228cb4502d96e1a6122cc32890e3c6c6c3 Mon Sep 17 00:00:00 2001
From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com>
Date: Wed, 22 Jul 2026 17:33:51 +0530
Subject: [PATCH 069/255] fix(install): route Strix to AMD gfx index on ROCm
7.14 (#7300)
* fix(install): route Strix to AMD gfx index on ROCm 7.14
When ROCm 7.3+ caps to the generic pytorch.org rocm7.2 index (or the
Radeon repo is unavailable), gfx1150/gfx1151 hosts were left on
torch 2.11+rocm7.2 instead of AMD's arch-specific wheels. Broaden the
Strix reroute in install.sh and studio/install_python_stack.py so
`studio update` repairs the same path as fresh installs (unslothai#7280).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
studio/install_python_stack.py | 37 ++++++++++++++++++-----
tests/studio/install/test_rocm_support.py | 30 ++++++++++++++++++
2 files changed, 59 insertions(+), 8 deletions(-)
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index b58e94cd3f..bb329e189e 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -73,6 +73,27 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = {
(6, 0): "rocm6.0",
}
+
+def _generic_pytorch_rocm_tag(ver: tuple[int, int]) -> str | None:
+ """Newest download.pytorch.org rocmX.Y tag for a host ROCm version."""
+ return next(
+ (t for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) if ver >= (maj, mn)),
+ None,
+ )
+
+
+_ROCM_ARCH_INDEX_FLOOR = (7, 13) # AMD per-arch index ships torch 2.11+rocm7.13
+
+
+def _strix_needs_amd_arch_index(ver: tuple[int, int]) -> bool:
+ """True when Strix's generic pytorch.org index sits below the AMD arch floor
+ (7.13), so gfx1150/1151 must use repo.amd.com's per-arch wheels. Mirrors
+ install.sh _rocm_leaf_below: reroute any generic rocm index (6.x/7.0/7.2 and a
+ future 7.3+), never one at/above the floor."""
+ key = next((k for k in sorted(_ROCM_TORCH_INDEX, reverse = True) if ver >= k), None)
+ return key is not None and key < _ROCM_ARCH_INDEX_FLOOR
+
+
# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug).
# Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare.
_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"})
@@ -1691,13 +1712,13 @@ def _ensure_rocm_torch() -> None:
rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
- # Strix Halo / Point (gfx1151 / gfx1150) segfault under ROCm 7.1 in torch._grouped_mm;
- # AMD's per-gfx repo ships 2.11.0+rocm7.13.0 with the fix, so route those hosts there
- # (mirrors install.sh). On mixed hosts, reroute only when HIP's runtime GPU is the Strix one.
+ # Strix Halo / Point (gfx1151 / gfx1150) need torch from AMD's per-gfx index
+ # (2.11+rocm7.13); any generic pytorch.org rocm index lacks the fixes (ROCm 7.1
+ # segfaults in _grouped_mm). See _strix_needs_amd_arch_index for the floor gate.
_strix_override_url: "str | None" = None
_strix_override_pkgs: "tuple[str, str, str] | None" = None
# An explicit ROCm pin is authoritative: never auto-reroute it.
- if ver < (7, 2) and _explicit_rocm_torch_index_url() is None:
+ if _strix_needs_amd_arch_index(ver) and _explicit_rocm_torch_index_url() is None:
gfx_codes = _detect_amd_gfx_codes()
_strix_gfx = {"gfx1151", "gfx1150"}
_detected_strix = _strix_gfx.intersection(gfx_codes)
@@ -1721,10 +1742,10 @@ def _ensure_rocm_torch() -> None:
print(
f"\n {_selected_gfx} (AMD Strix) is the runtime target with ROCm "
f"{ver[0]}.{ver[1]}.\n"
- f" ROCm 7.1 has a known _grouped_mm segfault on this GPU;\n"
- f" routing torch install to AMD's arch-specific index\n"
+ f" Routing torch install to AMD's arch-specific index\n"
f" ({_strix_override_url}) which serves torch 2.11.0+rocm7.13.0\n"
- f" with the upstream fix.\n"
+ f" with AMD's gfx1150/gfx1151 fixes (more reliable than the generic\n"
+ f" pytorch.org rocm7.2 index on ROCm 7.3+ hosts).\n"
)
else:
_gfx_str = ", ".join(sorted(_detected_strix))
@@ -1740,7 +1761,7 @@ def _ensure_rocm_torch() -> None:
index_url = _strix_override_url
_torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs
print(
- f" Strix ROCm 7.1 override -- installing torch from "
+ f" Strix arch-specific override -- installing torch from "
f"{_strip_index_url_credentials(index_url)}"
)
pip_install(
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index 5825bbe31f..b343b07238 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -710,6 +710,27 @@ class TestEnsureRocmTorch:
torch_call = mock_pip.call_args_list[0]
assert "rocm7.2" in str(torch_call)
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 14))
+ @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1150"])
+ def test_rocm_714_strix_routes_to_amd_arch_index(
+ self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """ROCm 7.14 caps to rocm7.2 on pytorch.org; Strix must use AMD gfx index."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"7.14.60850|2.11.0+rocm7.2\n"
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "gfx1150" in torch_call
+ assert "torch>=2.11.0,<2.12.0" in torch_call
+
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@@ -3253,6 +3274,15 @@ class TestStrixRocm71Override:
assert r.returncode == 0, f"probe aborted under set -e: {r.stderr}"
assert "OK:gfx1151" in r.stdout, f"amd-smi fallback not reached: {r.stdout!r}"
+ def test_strix_routing_helpers_cover_rocm714(self):
+ # Reroute for any generic pytorch.org index below the 7.13 arch floor (7.0,
+ # 7.2, a future 7.3+), never at/above it -- mirrors install.sh _rocm_leaf_below.
+ assert stack_mod._generic_pytorch_rocm_tag((7, 14)) == "rocm7.2"
+ assert stack_mod._strix_needs_amd_arch_index((7, 14)) is True
+ assert stack_mod._strix_needs_amd_arch_index((7, 0)) is True
+ assert stack_mod._strix_needs_amd_arch_index((6, 0)) is True
+ assert stack_mod._strix_needs_amd_arch_index((5, 0)) is False
+
def test_torch_constraint_updated_for_strix_amd_index(self):
"""install.sh must set TORCH_CONSTRAINT>=2.11 when routing Strix to AMD index."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
From 4759a5139d3226289518e2e5e52d4ef573dcfed5 Mon Sep 17 00:00:00 2001
From: Daniel Han
Date: Wed, 22 Jul 2026 05:20:59 -0700
Subject: [PATCH 070/255] Faster safetensors weight loading on unified-memory
(integrated) GPUs (#5988)
* Faster safetensors weight loading on unified-memory (integrated) GPUs
On unified-memory GPUs (AMD APUs / "Strix Halo", NVIDIA GB10 "Spark", Intel
iGPUs) the GPU shares the system memory pool. PyTorch's fast pinned-DMA
host->device path does not recognize the Rust-allocated, mmap-backed buffers
that safetensors hands back, so a direct safetensors GPU load
(`safe_open(..., device=)`) drops onto a slow per-tensor copy that, on
unified memory, additionally triggers page-attribute changes and page faults.
Wrap `transformers.modeling_utils.safe_open` so that, when transformers asks it
to load a shard directly onto a CUDA/HIP device, the shard is opened on CPU and
each tensor is `.clone()`-d into a normal torch allocation before `.to(device)`.
This restores the fast DMA path. Data, dtype and final device are unchanged, so
outputs are bit-identical -- only *how* the bytes reach the GPU changes.
Strictly gated to integrated/unified-memory GPUs via the standard `is_integrated`
device property (every visible device must be integrated): a hard no-op on
discrete NVIDIA/AMD GPUs, CPU, XPU and MLX, where the pinned-DMA path already
works. Only intercepts `framework="pt"` CUDA-device targets; CPU / disk-offload
loads are left untouched. Accuracy-neutral, idempotent, opt out with
UNSLOTH_DISABLE_UMA_CLONE_LOAD=1 (force the gate for tests with
UNSLOTH_FORCE_UMA=1/0).
This is the AMD/universal-UMA counterpart to the NVIDIA DGX Spark work in #5945
(which deliberately left the H2D clone-then-move out): gating on `is_integrated`
covers AMD Strix Halo, Intel iGPUs and Spark-class parts alike.
Verified on an AMD Radeon 8060S (gfx1151, Strix Halo) Windows ROCm box with
in-process, ordering-cancelled A/B benchmarks:
- H2D mechanism (safe_open device=0 vs cpu->clone->.to(0)): 2.08x faster
(1.076s -> 0.518s for a 988MB bf16 shard)
- full `from_pretrained`: 1.56x faster (1.552s -> 0.996s), saving 0.555s --
matching the H2D delta exactly
- max|logit diff| stock vs patched == 0.0 (bit-identical), generate + a LoRA
train step both verified
The absolute/relative win grows with bf16/fp16 weight volume (the same trick is
reported as ~2.3-2.75x on NVIDIA GB10 Spark for larger models).
Co-Authored-By: Claude Opus 4.8
* [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
* fix: evaluate the integrated-GPU gate lazily, not at import (Gemini review)
patch_unified_memory_safetensors_load() called
is_integrated_unified_memory_gpu() at install time, and the gate queries
torch.cuda.get_device_properties() for every visible device -- initializing
the CUDA context during `import unsloth` on every CUDA machine (discrete
included). That (a) breaks fork-based multiprocessing, (b) runs BEFORE
patch_dgx_spark_memory_config can set PYTORCH_CUDA_ALLOC_CONF on Spark,
defeating that patch's expandable_segments config in the very environment
this PR targets, and (c) charges a CUDA context to CPU-only imports.
The gate now runs lazily inside the wrapper, ordered AFTER the
framework/device check so non-CUDA loads never trigger the property query;
a CUDA-target safe_open means the caller is initializing CUDA anyway, and
the gate is lru-cached so it is evaluated once. The wrapper installs
unconditionally (opt-out and idempotency unchanged) and passes through when
the gate is off.
Tests: install-time no-eval guarantee (gate raises if called during
install), wrapper passthrough with the gate off, all previous gating /
passthrough / CUDA correctness tests kept -- 16/16 pass. Verified on the
N1X (WSL2): module exec + patch install leave torch.cuda.is_initialized()
unchanged; CPU loads pass through; forced CUDA-target loads intercept and
land bit-identical on the GPU.
Co-Authored-By: Claude Opus 4.8
* Compress PR comments to essentials (comment-only; AST-verified)
Docstrings and the _utils hook comment trimmed to their load-bearing
content (lazy-gate rationale, gating scope, opt-out env). AST dumps
with normalized docstrings are identical before/after for all three
files; the module's 16 unit tests pass unchanged.
Co-Authored-By: Claude Fable 5
* docs: tighten the UMA-load import comment (no code change)
* Tighten and trim code comments
* Drop unused is_integrated_unified_memory_gpu import from _utils.py
The UMA hook only needs patch_unified_memory_safetensors_load(); the
gate symbol is imported and used from ._uma_safetensors directly, so the
hoisted alias here was dead and tripped the import-hoist safety-net lint.
* Scope the UMA loader docstring to CUDA/HIP direct-device loads
The module text claimed Intel iGPU coverage, but the gate and device check
are CUDA/HIP only, and the clone path only wraps safe_open calls that carry
a CUDA device. State the actual scope and name the deliberate exclusions
(Intel XPU, CPU-open + .to() flows like bnb/HQQ) until they can be validated
on real hardware. Comment-only change.
* Tighten UMA safetensors loader comments
Trim the inline comments in the UMA clone-then-move path and the
_utils.py install site to be shorter and clearer. No code changes.
* uma: fall back to the direct move when the clone cannot allocate
The clone-and-move fast path transiently doubles one tensor's CPU
footprint while the mmap source and the CUDA destination are live. On a
UMA box with little free shared memory a large tensor could OOM where
the stock direct safe_open path would have loaded it. Both move sites
now go through a helper that catches the allocation failure and falls
back to the direct (slow but allocation-free) move, so the load always
succeeds; a genuine non-memory error re-raises identically from the
fallback.
Added a test that forces the clone to fail and verifies the wrapper
still lands tensors on the device with intact values (17 tests pass on
a real GPU).
* tests: track the moved pass-through inheritance in the gguf order check
Main moved the llama_extra_args pass-through inheritance out of the
GGUF branch into _resolve_inherited_extra_args, which runs before it,
so the source-order assertion's "if request.llama_extra_args is None"
anchor no longer exists inside the branch and the check failed after
the main merge. The test now asserts the same property in the current
shape: inheritance before the GGUF branch (a carried --no-mmproj still
shapes the hub guard's companion requirement), and marker, hub guard,
unload in order within the branch. Full file passes (32 tests).
* tests: anchor the inheritance order check on the call, not the definition
source.index("_resolve_inherited_extra_args(") matched the function
definition, which always precedes the endpoint, so the ordering
assertion was vacuously true. Anchoring on "= _resolve_inherited_
extra_args(" pins the first call site inside the load endpoint (line
4505), which is the statement whose position relative to the GGUF
branch the test is meant to guard. 32 tests pass.
* tests: align the gguf order test with main
Main fixed the stale ordering assertion in PR 7252; adopting its
version verbatim removes this file from the branch diff entirely and
avoids a conflict on the next main merge. 32 tests pass.
* uma: tighten comments
* Relicense UMA safetensors module and test under AGPL-3.0
---------
Co-authored-by: Claude Opus 4.8
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
tests/test_uma_safetensors_load.py | 229 +++++++++++++++++++++++++++++
unsloth/models/_uma_safetensors.py | 169 +++++++++++++++++++++
unsloth/models/_utils.py | 7 +
3 files changed, 405 insertions(+)
create mode 100644 tests/test_uma_safetensors_load.py
create mode 100644 unsloth/models/_uma_safetensors.py
diff --git a/tests/test_uma_safetensors_load.py b/tests/test_uma_safetensors_load.py
new file mode 100644
index 0000000000..c6d304ab4f
--- /dev/null
+++ b/tests/test_uma_safetensors_load.py
@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+
+"""Unit tests for the UMA safetensors clone-then-move fast load.
+
+The module loads in isolation with a fake ``transformers.modeling_utils``. The
+CUDA correctness check needs a GPU; gating, passthrough, idempotency and opt-out
+are GPU-free. The gate is lazy (wrapper-time), so the wrapper installs
+everywhere and passes through when it's off.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+torch = pytest.importorskip("torch")
+safetensors_torch = pytest.importorskip("safetensors.torch")
+import safetensors # noqa: E402
+
+_MODULE_PATH = Path(__file__).resolve().parent.parent / "unsloth" / "models" / "_uma_safetensors.py"
+
+
+def _load_module():
+ spec = importlib.util.spec_from_file_location("uma_safetensors_under_test", _MODULE_PATH)
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+@pytest.fixture()
+def uma():
+ return _load_module()
+
+
+@pytest.fixture()
+def force_uma(uma, monkeypatch):
+ """Force the UMA gate on (or off) and keep the lru_cache from sticking."""
+
+ def _set(on):
+ monkeypatch.setenv("UNSLOTH_FORCE_UMA", "1" if on else "0")
+ uma.is_integrated_unified_memory_gpu.cache_clear()
+
+ yield _set
+ uma.is_integrated_unified_memory_gpu.cache_clear()
+
+
+@pytest.fixture()
+def tiny_safetensors(tmp_path):
+ tensors = {
+ "w": torch.arange(32, dtype = torch.float32).reshape(4, 8),
+ "b": torch.tensor([1.0, 2.0, 3.0, 4.0], dtype = torch.float32),
+ }
+ path = tmp_path / "model.safetensors"
+ safetensors_torch.save_file(tensors, str(path))
+ return path, tensors
+
+
+def _install_fake_modeling_utils(monkeypatch, safe_open_fn):
+ fake_transformers = types.ModuleType("transformers")
+ fake_mu = types.ModuleType("transformers.modeling_utils")
+ fake_mu.safe_open = safe_open_fn
+ fake_transformers.modeling_utils = fake_mu
+ monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
+ monkeypatch.setitem(sys.modules, "transformers.modeling_utils", fake_mu)
+ return fake_mu
+
+
+# --- detection / gate ---
+
+
+def test_force_uma_on(uma, monkeypatch):
+ monkeypatch.setenv("UNSLOTH_FORCE_UMA", "1")
+ uma.is_integrated_unified_memory_gpu.cache_clear()
+ assert uma.is_integrated_unified_memory_gpu() is True
+
+
+def test_force_uma_off(uma, monkeypatch):
+ monkeypatch.setenv("UNSLOTH_FORCE_UMA", "0")
+ uma.is_integrated_unified_memory_gpu.cache_clear()
+ assert uma.is_integrated_unified_memory_gpu() is False
+
+
+@pytest.mark.parametrize(
+ "device,expected",
+ [
+ (0, True),
+ ("cuda", True),
+ ("cuda:0", True),
+ ("cpu", False),
+ ("disk", False),
+ (None, False),
+ (True, False), # a bool is not a device index
+ ],
+)
+def test_is_cuda_target(uma, device, expected):
+ assert uma._is_cuda_target(device) is expected
+
+
+def test_is_cuda_target_torch_device(uma):
+ assert uma._is_cuda_target(torch.device("cuda", 0)) is True
+ assert uma._is_cuda_target(torch.device("cpu")) is False
+
+
+# --- patch gating ---
+
+
+def test_wrapper_passes_through_off_uma(uma, force_uma, monkeypatch):
+ """Gate OFF: every call -- including CUDA targets -- passes straight through
+ to the real safe_open (the gate is evaluated lazily inside the wrapper)."""
+ force_uma(False)
+ sentinel = object()
+ calls = []
+
+ def fake_safe_open(*args, **kwargs):
+ calls.append((args, kwargs))
+ return sentinel
+
+ fake_mu = _install_fake_modeling_utils(monkeypatch, fake_safe_open)
+ assert uma.patch_unified_memory_safetensors_load() is True
+ assert getattr(fake_mu.safe_open, "_unsloth_uma_clone", False) is True
+ out = fake_mu.safe_open("shard.safetensors", "pt", "cuda:0")
+ assert out is sentinel
+ assert calls == [(("shard.safetensors", "pt", "cuda:0"), {})]
+
+
+def test_patch_install_does_not_evaluate_gate(uma, monkeypatch):
+ """Installing the wrapper must NOT query the integrated-GPU property -- that
+ would init CUDA at ``import unsloth`` (fork-unsafe, and before the Spark
+ allocator config is set)."""
+
+ def _boom():
+ raise AssertionError("gate must not be evaluated at install time")
+
+ _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+ monkeypatch.setattr(uma, "is_integrated_unified_memory_gpu", _boom)
+ assert uma.patch_unified_memory_safetensors_load() is True
+
+
+def test_patch_noop_when_opted_out(uma, force_uma, monkeypatch):
+ force_uma(True)
+ monkeypatch.setenv("UNSLOTH_DISABLE_UMA_CLONE_LOAD", "1")
+ real = object()
+ fake_mu = _install_fake_modeling_utils(monkeypatch, real)
+ assert uma.patch_unified_memory_safetensors_load() is False
+ assert fake_mu.safe_open is real
+
+
+def test_patch_installs_and_is_idempotent(uma, force_uma, monkeypatch):
+ force_uma(True)
+ fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+ assert uma.patch_unified_memory_safetensors_load() is True
+ wrapped = fake_mu.safe_open
+ assert getattr(wrapped, "_unsloth_uma_clone", False) is True
+ # second call must not double-wrap
+ assert uma.patch_unified_memory_safetensors_load() is True
+ assert fake_mu.safe_open is wrapped
+
+
+# --- correctness ---
+
+
+def test_cpu_target_is_passthrough(uma, force_uma, monkeypatch, tiny_safetensors):
+ path, tensors = tiny_safetensors
+ force_uma(True)
+ fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+ uma.patch_unified_memory_safetensors_load()
+ # device="cpu" must NOT be intercepted -> identical data, still on CPU.
+ with fake_mu.safe_open(str(path), framework = "pt", device = "cpu") as f:
+ for key, expected in tensors.items():
+ got = f.get_slice(key)[:]
+ assert got.device.type == "cpu"
+ assert torch.equal(got, expected)
+
+
+@pytest.mark.skipif(
+ not (hasattr(torch, "cuda") and torch.cuda.is_available()),
+ reason = "needs a GPU for the host->device clone-and-move path",
+)
+def test_cuda_target_clones_and_moves(uma, force_uma, monkeypatch, tiny_safetensors):
+ path, tensors = tiny_safetensors
+ force_uma(True)
+ fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+ uma.patch_unified_memory_safetensors_load()
+ # device="cuda" IS intercepted -> tensors land on cuda, byte-identical.
+ with fake_mu.safe_open(str(path), framework = "pt", device = "cuda") as f:
+ for key, expected in tensors.items():
+ got = f.get_slice(key)[:]
+ assert got.device.type == "cuda"
+ assert torch.equal(got.cpu(), expected)
+ got_full = f.get_tensor(key)
+ assert got_full.device.type == "cuda"
+ assert torch.equal(got_full.cpu(), expected)
+
+
+@pytest.mark.skipif(
+ not (hasattr(torch, "cuda") and torch.cuda.is_available()),
+ reason = "needs a GPU for the low-memory fallback path",
+)
+def test_low_memory_falls_back_to_direct_move(uma, force_uma, monkeypatch, tiny_safetensors):
+ path, tensors = tiny_safetensors
+ force_uma(True)
+ fake_mu = _install_fake_modeling_utils(monkeypatch, safetensors.safe_open)
+ uma.patch_unified_memory_safetensors_load()
+ # Clone OOMs (transient CPU doubling on a constrained UMA box): the wrapper
+ # must fall back to the direct move and still succeed.
+ real_clone = torch.Tensor.clone
+
+ def _oom_clone(self, *a, **k):
+ raise RuntimeError("[enforce fail] not enough memory")
+
+ monkeypatch.setattr(torch.Tensor, "clone", _oom_clone)
+ try:
+ with fake_mu.safe_open(str(path), framework = "pt", device = "cuda") as f:
+ for key, expected in tensors.items():
+ got = f.get_slice(key)[:]
+ assert got.device.type == "cuda"
+ got_full = f.get_tensor(key)
+ assert got_full.device.type == "cuda"
+ finally:
+ monkeypatch.setattr(torch.Tensor, "clone", real_clone)
+ for key, expected in tensors.items():
+ with fake_mu.safe_open(str(path), framework = "pt", device = "cuda") as f:
+ assert torch.equal(f.get_tensor(key).cpu(), expected)
diff --git a/unsloth/models/_uma_safetensors.py b/unsloth/models/_uma_safetensors.py
new file mode 100644
index 0000000000..38d8b7d33a
--- /dev/null
+++ b/unsloth/models/_uma_safetensors.py
@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+
+"""Faster safetensors weight loading on unified-memory (integrated) GPUs.
+
+A direct ``safe_open(..., device=)`` on CUDA/HIP UMA GPUs (AMD APUs,
+NVIDIA GB10 Spark) misses torch's fast pinned-DMA path: the mmap-backed
+safetensors buffers aren't recognized, so it falls to a slow per-tensor copy
+with page faults. Cloning each tensor into a normal torch CPU allocation before
+moving it restores the fast path; outputs are bit-identical.
+
+CUDA/HIP only, and only for loads that pass a CUDA device to ``safe_open``
+directly: Intel XPU iGPUs and the CPU-open + later ``.to()`` flows (e.g. bnb /
+HQQ quantized loads) keep the stock path until they can be validated on real
+hardware.
+"""
+
+import os
+import functools
+
+import torch
+
+__all__ = [
+ "is_integrated_unified_memory_gpu",
+ "patch_unified_memory_safetensors_load",
+]
+
+
+@functools.lru_cache(maxsize = None)
+def is_integrated_unified_memory_gpu():
+ """True only when EVERY visible CUDA/HIP device is integrated (UMA).
+
+ Discrete and mixed discrete+iGPU boxes return False (pinned-DMA already
+ works there). Test override: ``UNSLOTH_FORCE_UMA=1`` / ``=0``.
+ """
+ _force = os.environ.get("UNSLOTH_FORCE_UMA")
+ if _force == "1":
+ return True
+ if _force == "0":
+ return False
+ try:
+ if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
+ return False
+ count = torch.cuda.device_count()
+ if count == 0:
+ return False
+ for index in range(count):
+ props = torch.cuda.get_device_properties(index)
+ if not getattr(props, "is_integrated", 0):
+ return False
+ return True
+ except Exception:
+ return False
+
+
+def _is_cuda_target(device):
+ """Does a ``safe_open`` ``device=`` arg name a CUDA/HIP device?"""
+ if isinstance(device, bool):
+ return False
+ if isinstance(device, int):
+ return True
+ if isinstance(device, str):
+ return device == "cuda" or device.startswith("cuda:")
+ try:
+ return isinstance(device, torch.device) and device.type == "cuda"
+ except Exception:
+ return False
+
+
+def patch_unified_memory_safetensors_load():
+ """Wrap ``transformers.modeling_utils.safe_open`` so CUDA-target shard loads
+ open on CPU then clone+``.to(device)``, restoring the UMA fast path.
+
+ Gated to integrated GPUs (no-op on discrete/CPU/XPU/MLX), ``framework="pt"``
+ CUDA targets only, idempotent. Opt out: ``UNSLOTH_DISABLE_UMA_CLONE_LOAD=1``.
+
+ The gate runs lazily inside the wrapper, never here: probing device
+ properties at install would init CUDA during ``import unsloth`` -- breaking
+ fork multiprocessing and preempting ``patch_dgx_spark_memory_config``'s
+ allocator config. Returns ``True`` if the wrapper was installed.
+ """
+ if os.environ.get("UNSLOTH_DISABLE_UMA_CLONE_LOAD") == "1":
+ return False
+ try:
+ from transformers import modeling_utils as _mu
+ except Exception:
+ return False
+ real_safe_open = getattr(_mu, "safe_open", None)
+ if real_safe_open is None:
+ return False
+ if getattr(real_safe_open, "_unsloth_uma_clone", False):
+ return True
+
+ def _clone_move(tensor, device):
+ # Clone into a regular CPU allocation to restore fast pinned-DMA, then
+ # move. The clone transiently doubles the tensor's CPU footprint and can
+ # OOM a low-memory UMA box; fall back to the direct, allocation-free move
+ # (a genuine non-memory error re-raises identically from it).
+ try:
+ return tensor.clone().to(device, non_blocking = False)
+ except (MemoryError, RuntimeError):
+ return tensor.to(device, non_blocking = False)
+
+ class _ClonedSlice:
+ """Proxy over a safetensors ``PySafeSlice`` that clones+moves on read."""
+
+ __slots__ = ("_real", "_device")
+
+ def __init__(self, real, device):
+ self._real = real
+ self._device = device
+
+ def __getattr__(self, name):
+ if name in ("_real", "_device"):
+ raise AttributeError(name)
+ return getattr(self._real, name)
+
+ def __getitem__(self, key):
+ return _clone_move(self._real[key], self._device)
+
+ class _ClonedSafeOpen:
+ """Safetensors-handle proxy: load on CPU, clone+move tensors to CUDA."""
+
+ __slots__ = ("_real", "_device")
+
+ def __init__(self, args, kwargs):
+ self._device = kwargs.get("device", args[2] if len(args) > 2 else "cpu")
+ # Open on CPU; move ourselves.
+ if len(args) > 2:
+ args = args[:2] + ("cpu",) + tuple(args[3:])
+ else:
+ kwargs = dict(kwargs)
+ kwargs["device"] = "cpu"
+ self._real = real_safe_open(*args, **kwargs)
+
+ def __enter__(self):
+ self._real.__enter__()
+ return self
+
+ def __exit__(self, *exc):
+ return self._real.__exit__(*exc)
+
+ def __getattr__(self, name):
+ if name in ("_real", "_device"):
+ raise AttributeError(name)
+ return getattr(self._real, name)
+
+ def get_slice(self, name):
+ return _ClonedSlice(self._real.get_slice(name), self._device)
+
+ def get_tensor(self, name):
+ return _clone_move(self._real.get_tensor(name), self._device)
+
+ @functools.wraps(real_safe_open)
+ def _uma_safe_open(*args, **kwargs):
+ framework = kwargs.get("framework", args[1] if len(args) > 1 else None)
+ device = kwargs.get("device", args[2] if len(args) > 2 else "cpu")
+ # Device check first: non-CUDA loads must not trigger the CUDA-init gate.
+ if (
+ framework in ("pt", "pytorch")
+ and _is_cuda_target(device)
+ and is_integrated_unified_memory_gpu()
+ ):
+ return _ClonedSafeOpen(args, kwargs)
+ return real_safe_open(*args, **kwargs)
+
+ _uma_safe_open._unsloth_uma_clone = True
+ _mu.safe_open = _uma_safe_open
+ return True
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 57169fa3de..f9ac879de6 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1670,6 +1670,13 @@ except:
from transformers.modeling_utils import logger as transformers_logger
+# Faster safetensors loads on UMA (integrated) GPUs; lazy gate keeps this import
+# fork-safe (no CUDA init). No-op off-UMA. Opt out: UNSLOTH_DISABLE_UMA_CLONE_LOAD=1.
+from ._uma_safetensors import patch_unified_memory_safetensors_load
+
+patch_unified_memory_safetensors_load()
+
+
def _all_missing_keys_are_position_ids(record_str):
"""True only when EVERY key in the 'newly initialized: [...]' list is a position_ids
buffer.
From 36ec2cc046fd5834ff45ef2273b8a7247368ccc4 Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Wed, 22 Jul 2026 10:14:40 -0300
Subject: [PATCH 071/255] Studio: lighten chat text weight on Linux to match
macOS rendering (#7308)
* Studio: lighten chat text weight on Linux to match macOS rendering
* Exclude custom interface fonts from the Linux chat weight compensation
* Simplify Linux chat font weight override
---
.../features/settings/stores/appearance-custom-store.ts | 2 ++
studio/frontend/src/index.css | 7 +++++++
studio/frontend/src/main.tsx | 7 +++++++
3 files changed, 16 insertions(+)
diff --git a/studio/frontend/src/features/settings/stores/appearance-custom-store.ts b/studio/frontend/src/features/settings/stores/appearance-custom-store.ts
index b8c8d96f5a..f3618ddca5 100644
--- a/studio/frontend/src/features/settings/stores/appearance-custom-store.ts
+++ b/studio/frontend/src/features/settings/stores/appearance-custom-store.ts
@@ -479,6 +479,8 @@ export function applyCustomizationToDocument(
"--font-sans",
c.uiFont ? `"${c.uiFont}", ${DEFAULT_SANS_STACK}` : null,
);
+ // Custom interface fonts cascade into chat and opt out of its Inter tuning.
+ el.toggleAttribute("data-ui-font", Boolean(c.uiFont));
setVar(
"--font-heading",
c.headingFont ? `"${c.headingFont}", ${DEFAULT_HEADING_STACK}` : null,
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 2192cfb2cd..52ca81e064 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -633,6 +633,13 @@ html.no-font-smoothing body {
-moz-osx-font-smoothing: auto;
}
+/* Match Inter's lighter macOS rendering. Keep 410 when smoothing is off or a
+ custom font reaches chat. */
+html.render-linux:not(.no-font-smoothing):not([data-chat-font]):not([data-ui-font])
+ :is(.aui-assistant-message-root, .aui-user-message-root) {
+ font-weight: 350;
+}
+
/* Chat font: only applies while a custom chat font is set. Elements with
explicit font utilities (headings, code) keep their own families. */
html[data-chat-font] .aui-root {
diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx
index d0ddf2fc6e..e3b2bceccf 100644
--- a/studio/frontend/src/main.tsx
+++ b/studio/frontend/src/main.tsx
@@ -36,6 +36,13 @@ if (!rootElement) {
initializeLocale();
+// Rasterization follows the browser OS, not the potentially remote server.
+// This adjustment is calibrated for desktop Linux, so exclude Android.
+const uaLower = navigator.userAgent.toLowerCase();
+if (uaLower.includes("linux") && !uaLower.includes("android")) {
+ document.documentElement.classList.add("render-linux");
+}
+
createRoot(rootElement).render(
From fdf2df4edf6e194c3bcbc413d1d458236fb556e3 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Wed, 22 Jul 2026 06:34:02 -0700
Subject: [PATCH 072/255] Studio: reorder sidebar, rename Hub to Models (#7327)
* Studio: put Hub above Projects in the sidebar
Swap the two nav rows so Hub sits directly under New Chat, ahead of
Projects. Order only, no behavior change.
* Studio: rename Hub to Models, lowercase New chat
Rename the Hub nav row and its page heading to Models (localized in all
locales). Use sentence case 'New chat' in the English label.
* Studio: fix dataset title and stale Hub tab hints after rename
Show 'Datasets' as the catalog heading in dataset mode, not 'Models'.
Update the download-conflict toasts to point at the Models tab.
---
.../frontend/src/components/app-sidebar.tsx | 24 +++++++++----------
.../frontend/src/features/chat/chat-page.tsx | 8 +++----
.../features/hub/catalog/models-header.tsx | 2 +-
studio/frontend/src/i18n/locales/ar.ts | 2 +-
studio/frontend/src/i18n/locales/de.ts | 2 +-
studio/frontend/src/i18n/locales/en.ts | 4 ++--
studio/frontend/src/i18n/locales/es.ts | 2 +-
studio/frontend/src/i18n/locales/fr.ts | 2 +-
studio/frontend/src/i18n/locales/hi.ts | 2 +-
studio/frontend/src/i18n/locales/ja.ts | 2 +-
studio/frontend/src/i18n/locales/ko.ts | 2 +-
studio/frontend/src/i18n/locales/pt-br.ts | 2 +-
studio/frontend/src/i18n/locales/ru.ts | 2 +-
studio/frontend/src/i18n/locales/zh-CN.ts | 2 +-
14 files changed, 29 insertions(+), 29 deletions(-)
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index f4226760a2..10621ecd76 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -1357,6 +1357,18 @@ export function AppSidebar() {
+ {
+ navigate({ to: "/hub" });
+ closeMobileIfOpen();
+ }}
+ onIntent={() => {
+ preloadSilently(router.preloadRoute({ to: "/hub" }));
+ }}
+ />
- {
- navigate({ to: "/hub" });
- closeMobileIfOpen();
- }}
- onIntent={() => {
- preloadSilently(router.preloadRoute({ to: "/hub" }));
- }}
- />
{/* Train has a labelled section when expanded; plain icon here only when collapsed. */}
Date: Thu, 23 Jul 2026 03:16:25 +0200
Subject: [PATCH 073/255] Studio: mask AMD GPU pins via ROCR so an unsupported
iGPU can't crash llama-server (#7272)
* Studio: mask AMD GPU pins via ROCR so an unsupported iGPU can't crash llama-server
On a mixed AMD host (e.g. a discrete gfx1102 GPU next to a gfx1103 iGPU)
the bundled rocm-gfx110X llama.cpp build segfaults during HSA device
enumeration on the unsupported iGPU -- before llama-server prints a line,
so every model load fails with a bare signal and empty logs.
The GPU-subset pin masked visibility with HIP_VISIBLE_DEVICES, but HIP
filtering runs only after the HSA runtime has already enumerated (and
crashed on) every agent. Mask the subset via ROCR_VISIBLE_DEVICES (the
ROCr/HSA layer) instead, so a deselected/unsupported GPU is never
enumerated. Exactly one layer is masked (HIP cleared) to avoid the
double-mask reindex that would otherwise drop the child to CPU. The
whole-set tensor-split path and the CPU-only sentinel keep their existing
HIP behavior.
Also stop misreporting the resulting startup segfault as a vision
projector incompatibility: when the text-only mmproj retry also hard-
crashes with a signal, surface a GPU/driver init crash (with the ROCR
hint) instead of blaming the projector.
Co-Authored-By: Claude Opus 4.8
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten _emit_child_gpu_visibility comments for #7272
Comment and docstring only: condense the ROCR-vs-HIP masking rationale from ~22 to ~14 lines and the call-site note from 5 to 3, keeping every technical point (HSA enumeration segfault, physical ids, the -1 sentinel). Logic is unchanged, verified by an AST compare with docstrings stripped and by exercising _emit_child_gpu_visibility against a torch/HIP stub.
* Detect AMD SDK ROCm wheels (hip=None) in _emit_child_gpu_visibility (Codex P2)
The ROCm branch gated only on torch.version.hip, but AMD SDK wheels leave that unset while encoding 'rocm' in __version__ (detect_hardware handles this the same way). On such a wheel the masking was skipped entirely, leaving only CUDA_VISIBLE_DEVICES, so on a mixed AMD box the unsupported deselected iGPU still enumerated and could crash llama-server. Now the branch also treats 'rocm' in torch.__version__ as ROCm, mirroring detect_hardware. Adds tests for the hip=None SDK wheel (ROCR + default paths) and a CUDA guard so the version-string check can't false-positive.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Remap CUDA_VISIBLE_DEVICES to post-ROCR ordinals on prefer_rocr for PR #7272 (Codex P1)
On the prefer_rocr path _emit_child_gpu_visibility set ROCR_VISIBLE_DEVICES to the
physical id and cleared HIP_VISIBLE_DEVICES, but left CUDA_VISIBLE_DEVICES at the
physical id. ROCR re-indexes the visible agents from 0 and, with HIP cleared, HIP
honours CUDA_VISIBLE_DEVICES -- so a non-zero pick (e.g. GPU 1) pointed out of
range, HIP saw 0 devices, and the child fell back to CPU, defeating GPU-picker
selections other than physical GPU 0. Remap CUDA to the post-ROCR ordinals
(0..N-1); GPU 0 is unchanged, the default (HIP) path and the CPU sentinel are
untouched, and non-AMD wheels never enter this branch.
* Detect AMD SDK wheels in _resolve_visible_physical_ids for PR #7272 (Codex P2)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Keep the HIP mask on Windows ROCm in prefer_rocr for PR #7272 (Codex P2)
* Ignore ROCR_VISIBLE_DEVICES in _resolve_visible_physical_ids on Windows for PR #7272 (Codex P2)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Preserve inherited ROCR masks in the tensor-split pin for PR #7272 (Codex P2)
---------
Co-authored-by: Claude Opus 4.8
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
Co-authored-by: Leo Borcherding
---
studio/backend/core/inference/llama_cpp.py | 108 ++++++++--
studio/backend/tests/test_gpu_memory_mode.py | 205 ++++++++++++++++++-
2 files changed, 291 insertions(+), 22 deletions(-)
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 8651ed9ea8..1c9c76ebe9 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -2912,12 +2912,25 @@ class LlamaCppBackend:
on the ordinal->physical mapping."""
try:
import torch
- is_rocm = getattr(torch.version, "hip", None) is not None
+
+ # Same ROCm detection as _emit_child_gpu_visibility: AMD SDK wheels
+ # leave version.hip unset but encode "rocm" in __version__. The two
+ # must agree, else an inherited ROCR mask reads back as "no mask",
+ # ordinal 0 is labelled physical 0, and the child's new ROCR pin
+ # re-exposes the GPU the inherited mask was hiding.
+ is_rocm = (
+ getattr(torch.version, "hip", None) is not None
+ or "rocm" in getattr(torch, "__version__", "").lower()
+ )
except Exception:
is_rocm = False
if is_rocm:
hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
- rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
+ # ROCR_VISIBLE_DEVICES is a Linux ROCr variable; Windows HIP has no
+ # ROCr layer, so a stray ROCR var there does not mask the runtime and
+ # must not be read as the ordinal->physical mapping (mirrors the
+ # Windows gate in _emit_child_gpu_visibility).
+ rocr_v = None if sys.platform == "win32" else os.environ.get("ROCR_VISIBLE_DEVICES")
cvd = (
hip_v
if hip_v is not None
@@ -2935,20 +2948,52 @@ class LlamaCppBackend:
return None
@staticmethod
- def _emit_child_gpu_visibility(env: dict, pinned: str) -> None:
- """Write the child's GPU visibility mask (CUDA, plus the HIP mirror on
- ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child
- seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP
- mask at different layers, so the same indices apply twice -- ROCR reduces
- and re-indexes from 0, then a non-zero HIP pin points out of range, HIP
- enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone
- narrows correctly; clear any inherited ROCR mask so it can't double up."""
+ def _emit_child_gpu_visibility(
+ env: dict,
+ pinned: str,
+ *,
+ prefer_rocr: bool = False,
+ ) -> None:
+ """Write the child's GPU visibility mask: CUDA, plus a ROCm mirror on AMD
+ (masking only CUDA_VISIBLE_DEVICES leaves an AMD child seeing every GPU).
+
+ Default: HIP_VISIBLE_DEVICES, clearing any inherited ROCR mask so the two
+ can't stack (ROCR re-indexes from 0, then a non-zero HIP pin points out of
+ range, HIP sees 0 devices, and llama.cpp falls back to CPU).
+
+ prefer_rocr masks at the ROCr/HSA layer instead (clearing HIP). A HIP mask
+ filters only AFTER the HSA runtime enumerates every agent, and that
+ enumeration segfaults at startup on a GPU the build has no kernels for
+ (e.g. a gfx1103 iGPU under a gfx110X prebuilt), before llama-server logs a
+ line. ROCR drops the device at the driver layer, consuming physical ids.
+ The CPU-only sentinel ("-1") has no portable ROCR spelling, so it keeps
+ the HIP mask. Windows keeps the HIP mask too: ROCR_VISIBLE_DEVICES is a
+ Linux ROCr variable (Windows HIP has no ROCr layer), so the ROCR pin
+ would be dead there while the cleared HIP mask stops selecting."""
env["CUDA_VISIBLE_DEVICES"] = pinned
try:
import torch as _torch
- if getattr(_torch.version, "hip", None) is not None:
- env["HIP_VISIBLE_DEVICES"] = pinned
- env.pop("ROCR_VISIBLE_DEVICES", None)
+
+ # torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may
+ # leave it unset but encode "rocm" in __version__ (mirrors detect_hardware).
+ if (
+ getattr(_torch.version, "hip", None) is not None
+ or "rocm" in getattr(_torch, "__version__", "").lower()
+ ):
+ if prefer_rocr and pinned != "-1" and sys.platform != "win32":
+ env["ROCR_VISIBLE_DEVICES"] = pinned
+ env.pop("HIP_VISIBLE_DEVICES", None)
+ # ROCR re-indexes the visible agents from 0, and with HIP
+ # cleared HIP honours CUDA_VISIBLE_DEVICES -- so it must carry
+ # the post-ROCR ordinals (0..N-1), not the physical ids, else a
+ # non-zero pick points out of range and HIP sees 0 devices (the
+ # same stacking the default path avoids by clearing ROCR).
+ env["CUDA_VISIBLE_DEVICES"] = ",".join(
+ str(i) for i in range(len(pinned.split(",")))
+ )
+ else:
+ env["HIP_VISIBLE_DEVICES"] = pinned
+ env.pop("ROCR_VISIBLE_DEVICES", None)
except Exception as e:
logger.debug("Failed to set ROCm visibility env vars for child: %s", e)
@@ -2983,7 +3028,21 @@ class LlamaCppBackend:
logger.debug("Could not read reported GPU order for split pin: %s", e)
if order is None:
order = sorted(inherited)
- LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order))
+ # Re-emit at the layer that produced the mapping. A parent masked only
+ # via ROCR_VISIBLE_DEVICES hides agents at the driver layer, and the
+ # default HIP re-emission clears that mask -- HSA then enumerates every
+ # agent again and can segfault at startup on an unsupported GPU the
+ # parent was hiding (the crash prefer_rocr exists to avoid). Linux-only,
+ # mirroring _resolve_visible_physical_ids: on Windows a stray ROCR var
+ # is dead and was not the mapping's source.
+ prefer_rocr = (
+ sys.platform != "win32"
+ and env.get("HIP_VISIBLE_DEVICES") is None
+ and env.get("ROCR_VISIBLE_DEVICES") is not None
+ )
+ LlamaCppBackend._emit_child_gpu_visibility(
+ env, ",".join(str(i) for i in order), prefer_rocr = prefer_rocr
+ )
@staticmethod
def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool:
@@ -7740,7 +7799,12 @@ class LlamaCppBackend:
# default FASTEST_FIRST order (#5025).
if gpu_ids:
env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
- self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices))
+ # Mask on AMD at the ROCr/HSA layer: HIP-only masking still
+ # enumerates every agent first, which segfaults on a deselected
+ # unsupported GPU (e.g. gfx1103 iGPU under a gfx110X prebuilt).
+ self._emit_child_gpu_visibility(
+ env, ",".join(str(i) for i in gpu_indices), prefer_rocr = True
+ )
elif manual_tensor_split_emitted and not is_vulkan_backend:
# A manual per-GPU ratio across ALL GPUs (no explicit pick, so
# no CUDA_VISIBLE_DEVICES mask above): the UI built the
@@ -8102,6 +8166,20 @@ class LlamaCppBackend:
# an OS-killed text-only retry still gets the OOM message.
_retry_rc = self._process.poll() if self._process is not None else None
self._kill_process()
+ # If the text-only retry ALSO hard-crashed (a signal, not
+ # OOM/timeout), the vision projector was never the cause:
+ # llama-server is faulting during GPU/driver init. Say so
+ # -- with the ROCm fix -- instead of blaming the mmproj.
+ if self._is_signal_crash(_retry_rc):
+ raise RuntimeError(
+ "llama-server crashed at startup on both the vision "
+ "and text-only attempts -- a GPU driver/runtime "
+ "initialization crash, not a model or vision-projector "
+ "problem. This often means an unsupported secondary "
+ "GPU; on AMD/ROCm, hide it with ROCR_VISIBLE_DEVICES "
+ "(e.g. ROCR_VISIBLE_DEVICES=0 exposes only the first "
+ "GPU) before launching Unsloth Studio."
+ )
raise RuntimeError(
"Vision projector incompatible with this llama.cpp "
"build, and the text-only retry also failed: "
diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py
index b17274197f..19ba9e3e05 100644
--- a/studio/backend/tests/test_gpu_memory_mode.py
+++ b/studio/backend/tests/test_gpu_memory_mode.py
@@ -733,20 +733,211 @@ def test_split_pin_without_mask_only_sets_pci_order(monkeypatch):
def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch):
- # ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR
- # mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP
- # would index into the already-reduced set).
+ # ROCm with the mask sourced from HIP: the pin must land in
+ # HIP_VISIBLE_DEVICES too, and an inherited ROCR mask is cleared so the
+ # mask can't apply twice (ROCR re-indexes, then HIP would index into the
+ # already-reduced set).
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
- torch_stub = _types.ModuleType("torch")
- torch_stub.version = _types.SimpleNamespace(hip = "6.0")
- monkeypatch.setitem(sys.modules, "torch", torch_stub)
- env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"}
+ _rocm_torch_stub(monkeypatch)
+ env = {
+ "CUDA_VISIBLE_DEVICES": "3,1",
+ "HIP_VISIBLE_DEVICES": "3,1",
+ "ROCR_VISIBLE_DEVICES": "3,1",
+ }
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
assert env["HIP_VISIBLE_DEVICES"] == "1,3"
assert "ROCR_VISIBLE_DEVICES" not in env
+def test_split_pin_preserves_inherited_rocr_mask(monkeypatch):
+ # Mask sourced from ROCR alone (e.g. an AMD SDK parent): the pin must
+ # re-emit at the ROCr layer, not swap to HIP -- clearing ROCR re-exposes
+ # every agent to HSA enumeration, which can segfault at startup on an
+ # unsupported GPU the parent mask was hiding (#7272 review). CUDA carries
+ # the post-ROCR ordinals, mirroring the prefer_rocr emission.
+ _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
+ _rocm_torch_stub(monkeypatch)
+ env = {"ROCR_VISIBLE_DEVICES": "3,1"}
+ LlamaCppBackend._pin_visible_gpu_order_for_split(env)
+ assert env["ROCR_VISIBLE_DEVICES"] == "1,3"
+ assert env["CUDA_VISIBLE_DEVICES"] == "0,1"
+ assert "HIP_VISIBLE_DEVICES" not in env
+
+
+def test_split_pin_keeps_hip_on_windows_despite_stray_rocr(monkeypatch):
+ # On Windows the ROCR var is dead (no ROCr layer) and the resolver never
+ # reads it, so a stray value must not flip the pin to the ROCR emission:
+ # the HIP mask is the only effective selector there.
+ _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = "6.0")
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ monkeypatch.setattr(sys, "platform", "win32")
+ env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "9"}
+ LlamaCppBackend._pin_visible_gpu_order_for_split(env)
+ assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
+ assert env["HIP_VISIBLE_DEVICES"] == "1,3"
+ assert "ROCR_VISIBLE_DEVICES" not in env
+
+
+def _rocm_torch_stub(monkeypatch):
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = "6.0")
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ # prefer_rocr is Linux-only (ROCR is an ROCr variable); pin the platform so
+ # these Linux-behaviour tests also pass on a Windows dev box.
+ monkeypatch.setattr(sys, "platform", "linux")
+
+
+def test_subset_pin_masks_via_rocr_on_rocm(monkeypatch):
+ # A GPU-subset pin must exclude the rest at the ROCr/HSA layer: HIP masking
+ # still enumerates every agent first, which segfaults the build on an
+ # unsupported deselected GPU (e.g. a gfx1103 iGPU under a gfx110X prebuilt).
+ # ROCR drops it at the driver layer; only one mask is set (HIP cleared).
+ _rocm_torch_stub(monkeypatch)
+ env = {"HIP_VISIBLE_DEVICES": "9"} # stale/inherited HIP mask must not survive
+ LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True)
+ assert env["ROCR_VISIBLE_DEVICES"] == "0"
+ assert env["CUDA_VISIBLE_DEVICES"] == "0"
+ assert "HIP_VISIBLE_DEVICES" not in env
+
+
+def test_prefer_rocr_remaps_cuda_to_post_rocr_ordinals(monkeypatch):
+ # ROCR re-indexes the visible agents from 0, and HIP (cleared here) falls back
+ # to CUDA_VISIBLE_DEVICES -- so on the prefer_rocr path CUDA must carry the
+ # post-ROCR ordinals, not the physical ids, else a non-zero pick indexes out
+ # of range and the child sees no GPU and drops to CPU (#7272 review).
+ _rocm_torch_stub(monkeypatch)
+ # Single non-zero GPU: ROCR keeps the physical id, CUDA becomes ordinal 0.
+ env = {}
+ LlamaCppBackend._emit_child_gpu_visibility(env, "1", prefer_rocr = True)
+ assert env["ROCR_VISIBLE_DEVICES"] == "1"
+ assert env["CUDA_VISIBLE_DEVICES"] == "0"
+ assert "HIP_VISIBLE_DEVICES" not in env
+ # Multi-GPU subset: ROCR keeps the physical ids, CUDA is the 0-based ordinals.
+ env = {}
+ LlamaCppBackend._emit_child_gpu_visibility(env, "1,3", prefer_rocr = True)
+ assert env["ROCR_VISIBLE_DEVICES"] == "1,3"
+ assert env["CUDA_VISIBLE_DEVICES"] == "0,1"
+ assert "HIP_VISIBLE_DEVICES" not in env
+
+
+def test_subset_pin_default_still_uses_hip_and_clears_rocr(monkeypatch):
+ # Without prefer_rocr the masking is unchanged: HIP narrows, inherited ROCR
+ # is cleared so the two can't double-mask.
+ _rocm_torch_stub(monkeypatch)
+ env = {"ROCR_VISIBLE_DEVICES": "0,1"}
+ LlamaCppBackend._emit_child_gpu_visibility(env, "1")
+ assert env["HIP_VISIBLE_DEVICES"] == "1"
+ assert "ROCR_VISIBLE_DEVICES" not in env
+
+
+def test_cpu_only_pin_keeps_hip_even_with_prefer_rocr(monkeypatch):
+ # The CPU-only sentinel never routes through ROCR (no portable "hide all"
+ # spelling); it hides every GPU via HIP.
+ _rocm_torch_stub(monkeypatch)
+ env = {}
+ LlamaCppBackend._emit_child_gpu_visibility(env, "-1", prefer_rocr = True)
+ assert env["HIP_VISIBLE_DEVICES"] == "-1"
+ assert "ROCR_VISIBLE_DEVICES" not in env
+
+
+def _amd_sdk_torch_stub(monkeypatch):
+ # AMD SDK wheel: torch.version.hip is None but __version__ encodes rocm.
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = None)
+ torch_stub.__version__ = "2.9.1+rocm7.2.1"
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ monkeypatch.setattr(sys, "platform", "linux")
+
+
+def test_prefer_rocr_falls_back_to_hip_on_windows(monkeypatch):
+ # ROCR_VISIBLE_DEVICES is a Linux ROCr variable (Windows HIP has no ROCr
+ # layer), so on Windows ROCm prefer_rocr must keep the HIP mask or a nonzero
+ # pick loses its only effective selector (#7272 review).
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = "6.0")
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ monkeypatch.setattr(sys, "platform", "win32")
+ env = {"ROCR_VISIBLE_DEVICES": "9"}
+ LlamaCppBackend._emit_child_gpu_visibility(env, "1", prefer_rocr = True)
+ assert env["HIP_VISIBLE_DEVICES"] == "1"
+ assert env["CUDA_VISIBLE_DEVICES"] == "1"
+ assert "ROCR_VISIBLE_DEVICES" not in env
+
+
+def test_amd_sdk_wheel_hip_none_still_masks_rocr(monkeypatch):
+ # An AMD SDK wheel leaves torch.version.hip unset but has "rocm" in __version__.
+ # It must still get the ROCR mask, else only CUDA_VISIBLE_DEVICES is set and an
+ # unsupported iGPU keeps enumerating and can crash llama-server.
+ _amd_sdk_torch_stub(monkeypatch)
+ env = {"HIP_VISIBLE_DEVICES": "9"}
+ LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True)
+ assert env["ROCR_VISIBLE_DEVICES"] == "0"
+ assert "HIP_VISIBLE_DEVICES" not in env
+
+
+def test_cuda_wheel_hip_none_gets_no_rocm_mask(monkeypatch):
+ # A CUDA wheel (hip=None, no "rocm" in __version__) must NOT get a HIP/ROCR mask
+ # -- only CUDA_VISIBLE_DEVICES -- so the version-string check can't false-positive.
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = None)
+ torch_stub.__version__ = "2.9.1+cu124"
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ env = {}
+ LlamaCppBackend._emit_child_gpu_visibility(env, "0", prefer_rocr = True)
+ assert env["CUDA_VISIBLE_DEVICES"] == "0"
+ assert "ROCR_VISIBLE_DEVICES" not in env
+ assert "HIP_VISIBLE_DEVICES" not in env
+
+
+def test_resolve_physical_ids_reads_rocr_on_amd_sdk_wheel(monkeypatch):
+ # _resolve_visible_physical_ids must use the same ROCm detection as
+ # _emit_child_gpu_visibility: on an AMD SDK wheel (hip=None, rocm in
+ # __version__) an inherited ROCR mask IS the ordinal->physical mapping.
+ # Reading it as "no mask" labels ordinal 0 as physical 0 and the child's
+ # ROCR pin then re-exposes the GPU the mask was hiding (#7272 review).
+ _amd_sdk_torch_stub(monkeypatch)
+ for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
+ monkeypatch.delenv(var, raising = False)
+ monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1")
+ assert LlamaCppBackend._resolve_visible_physical_ids() == [1]
+
+
+def test_resolve_physical_ids_ignores_rocr_on_cuda_wheel(monkeypatch):
+ # A CUDA wheel (hip=None, no "rocm") keeps CUDA-only semantics: a stray
+ # ROCR var must not be read as the mask.
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = None)
+ torch_stub.__version__ = "2.9.1+cu124"
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
+ monkeypatch.delenv(var, raising = False)
+ monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1")
+ assert LlamaCppBackend._resolve_visible_physical_ids() is None
+
+
+def test_resolve_physical_ids_ignores_rocr_on_windows(monkeypatch):
+ # ROCR_VISIBLE_DEVICES is a Linux ROCr variable: Windows HIP has no ROCr
+ # layer, so a stray ROCR var there does not mask the runtime. Reading it as
+ # the ordinal->physical mapping would label ordinal 0 with a stale ROCR id
+ # while the runtime still enumerates every adapter, so auto-selection could
+ # budget one card and pin another (#7272 review). HIP must still be honoured.
+ torch_stub = _types.ModuleType("torch")
+ torch_stub.version = _types.SimpleNamespace(hip = None)
+ torch_stub.__version__ = "2.9.1+rocm7.2.1" # AMD SDK wheel
+ monkeypatch.setitem(sys.modules, "torch", torch_stub)
+ monkeypatch.setattr(sys, "platform", "win32")
+ for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
+ monkeypatch.delenv(var, raising = False)
+ monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1")
+ assert LlamaCppBackend._resolve_visible_physical_ids() is None
+ # HIP precedence is unchanged on Windows.
+ monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
+ assert LlamaCppBackend._resolve_visible_physical_ids() == [1]
+
+
# ── Diffusion single-device selection ───────────────────────────────────────
From 978ae4745bf4d975abce6aa943ffad2f2d7aee1e Mon Sep 17 00:00:00 2001
From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com>
Date: Thu, 23 Jul 2026 06:46:45 +0530
Subject: [PATCH 074/255] fix(install): infer Strix gfx when ROCm runtime is
absent (#7305)
* fix(install): infer Strix gfx when ROCm runtime is absent
When /dev/kfd and rocminfo are missing on Linux (e.g. Arch/CachyOS Strix
Halo), route to AMD per-arch wheels via cpuinfo/lspci inference instead
of CPU-only PyTorch. Mirrors install.ps1 Windows behavior and fixes
studio update via install_python_stack.py (unslothai#7301).
* Map Radeon 8065S to gfx1151 in the Linux gfx inference (Codex P2)
install.sh _infer_amd_gfx_arch_from_gpu_name missed 8065S, so a Strix Halo host that only exposes 'AMD Radeon 8065S' via lspci (no Ryzen AI Max branding in /proc/cpuinfo) was left on CPU torch. setup.sh and setup.ps1 already list 8065S -> gfx1151. Added it, and widened the cpuinfo regexes (install.sh and install_python_stack.py) from Radeon 80[0-9]0S to 80[0-9][05]S to match the 80X5S naming, consistent with the display-side check already in install.sh. Tests cover the 8065S name and the cpuinfo-only case.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gate the Linux gfx inference out of WSL without the ROCDXG runtime for PR #7305
On WSL /proc/cpuinfo and lspci still see the host APU, so a standalone
'unsloth studio update' could infer gfx1151 and install per-arch ROCm wheels
into a WSL env whose ROCDXG bridge (librocdxg) was never bootstrapped, i.e. one
that cannot expose the GPU. Skip the cpuinfo/lspci inference on WSL unless
librocdxg is present; an explicit UNSLOTH_ROCM_GFX_ARCH override still wins.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address Codex review on PR #7305 (WSL runtime gate, Linux mirror, arch guard)
- install.sh _infer_linux_amd_gfx_arch: skip the cpuinfo/lspci inference on WSL
unless librocdxg is present (the ROCDXG bridge), mirroring the Python fix, so a
WSL box whose ROCm bootstrap was skipped keeps the CPU fallback instead of
installing AMD wheels that cannot reach the GPU. The explicit UNSLOTH_ROCM_GFX_ARCH
override still returns first, so it stays authoritative.
- install.sh: guard the inferred-gfx reroute on x86_64|amd64. ROCm torch wheels are
not published for arm64, so an inferred/overridden gfx no longer pushes an arm64
host to the AMD arch index (get_torch_index_url returns CPU there).
- install_python_stack.py _amd_arch_index_url: honour UNSLOTH_AMD_ROCM_MIRROR on
Linux (the same var install.sh uses) instead of the Windows mirror var, so a
mirrored/air-gapped Linux 'unsloth studio update' reaches the index install.sh
chose. Windows still delegates unchanged; both default to repo.amd.com.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scan all AMD display controllers in the lspci fallback for PR #7305 (Codex P2)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): keep inferred AMD wheels from being overwritten
After a successful inferred-gfx install, skip the generic pytorch.org
ROCm reinstall so readable ROCm userland without /dev/kfd cannot undo
the per-arch repair (Codex P1 on #7305). Also merge latest main.
* Only take the inferred-gfx install when the runtime sees no GPU for PR #7305 (Codex P1)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Isolate three updater tests from the host cpuinfo for PR #7305 (Strix dev box leak)
* Gate the reroute on invisible ROCm and forward the inferred gfx to setup.sh for PR #7305 (Codex P2s)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Require AMD PCI display evidence for cpuinfo inference; honor gfx override with visible ROCm for PR #7305 (Codex P2s)
---------
Co-authored-by: Daniel Han
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: LeoBorcherding
---
install.sh | 140 ++++++++
studio/install_python_stack.py | 189 ++++++++++-
tests/studio/install/test_rocm_support.py | 393 +++++++++++++++++++++-
3 files changed, 714 insertions(+), 8 deletions(-)
diff --git a/install.sh b/install.sh
index e0f57c198b..963107524b 100755
--- a/install.sh
+++ b/install.sh
@@ -2144,6 +2144,92 @@ _amd_gpu_present_via_pci() {
return 1
}
+# Map a gfx arch to the AMD pip index family (mirrors install.ps1 $archFamilyMap).
+_amd_arch_index_family_for_gfx() {
+ case "$1" in
+ gfx1201|gfx1200) echo gfx120X-all ;;
+ gfx1151) echo gfx1151 ;;
+ gfx1150) echo gfx1150 ;;
+ gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;;
+ gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;;
+ gfx90a) echo gfx90a ;;
+ gfx908) echo gfx908 ;;
+ *) return 1 ;;
+ esac
+}
+
+# Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable).
+_infer_amd_gfx_arch_from_gpu_name() {
+ case "$1" in
+ *"9070 XT"*|*9080*) echo gfx1201 ;;
+ *9070*|*9060*) echo gfx1200 ;;
+ *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;;
+ *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1150 ;;
+ *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) echo gfx1102 ;;
+ *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) echo gfx1100 ;;
+ *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;;
+ *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;;
+ *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;;
+ *"RX 6500"*|*"RX 6400"*|*"RX 6300"*|*"PRO W6400"*|*"PRO W6500"*) echo gfx1034 ;;
+ *) return 1 ;;
+ esac
+}
+
+# Best-effort gfx inference when ROCm tools can't see the GPU (unslothai#7301).
+# Mirrors install.ps1 arch resolution on Windows ($HasROCm false, $ROCmGfxArch set).
+_infer_linux_amd_gfx_arch() {
+ if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then
+ printf '%s\n' "$(printf '%s' "$UNSLOTH_ROCM_GFX_ARCH" | tr '[:upper:]' '[:lower:]')"
+ return 0
+ fi
+ # On WSL /proc/cpuinfo and lspci still report the host APU, but without the
+ # ROCDXG bridge (librocdxg over /dev/dxg) the AMD wheels can't reach the GPU;
+ # keep the CPU fallback there unless that runtime is present (the explicit
+ # override above still wins). Mirrors install_python_stack.py.
+ _gpu_evidence=""
+ if [ -e /dev/dxg ] || grep -qi microsoft /proc/version 2>/dev/null; then
+ for _d in /opt/rocm/lib /opt/rocm/lib64 /opt/rocm-*/lib /opt/rocm-*/lib64; do
+ { [ -e "$_d/librocdxg.so" ] || [ -e "$_d/librocdxg.so.1" ]; } && _rocdxg=1 && break
+ done
+ [ -n "${_rocdxg:-}" ] || return 1
+ # WSL enumerates no PCI display device; /dev/dxg + librocdxg IS the
+ # GPU evidence there.
+ _gpu_evidence=1
+ elif _amd_gpu_present_via_pci; then
+ _gpu_evidence=1
+ fi
+ # /proc/cpuinfo leaks the HOST CPU model into VMs/containers that received
+ # no AMD GPU, so the CPU-model text alone is not GPU evidence: require an
+ # AMD display device (PCI vendor 0x1002, class 0x03*) before trusting it.
+ # The lspci fallback below needs no gate; an AMD display line IS evidence.
+ if [ -n "$_gpu_evidence" ] && grep -qiE 'Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo' /proc/cpuinfo 2>/dev/null; then
+ echo gfx1151
+ return 0
+ fi
+ if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then
+ echo gfx1150
+ return 0
+ fi
+ if command -v lspci >/dev/null 2>&1; then
+ # A non-AMD controller can enumerate first (Intel/ASPEED before an AMD
+ # dGPU), so scan every display-class line and take the first AMD one
+ # that maps. The vendor guard is case-SENSITIVE (a -i "ATI" would match
+ # "CorporATIon" on every Intel/NVIDIA line); whole-line matching also
+ # survives the 0000: PCI domain prefix. Mirrors install_python_stack.py.
+ _amd_disp=$(lspci -nn 2>/dev/null | grep -E 'VGA compatible controller|3D controller|Display controller' | grep -E 'AMD|ATI' || true)
+ while IFS= read -r _ln; do
+ [ -n "$_ln" ] || continue
+ if _gfx=$(_infer_amd_gfx_arch_from_gpu_name "$_ln"); then
+ echo "$_gfx"
+ return 0
+ fi
+ done </dev/null || true)
+ if [ -n "$_linux_inferred_gfx" ]; then
+ _amd_family=$(_amd_arch_index_family_for_gfx "$_linux_inferred_gfx") || _amd_family=""
+ if [ -n "$_amd_family" ]; then
+ _amd_mirror="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}"
+ while [ "${_amd_mirror%/}" != "$_amd_mirror" ]; do
+ _amd_mirror="${_amd_mirror%/}"
+ done
+ TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/"
+ # Hand the inferred arch to setup.sh (llama.cpp): it re-probes
+ # ROCm on its own, and on these runtime-less hosts its probes
+ # find nothing, so without this it classifies the box as
+ # non-ROCm and installs the CPU prebuilt while torch just got
+ # AMD per-arch wheels. setup.sh and install_llama_prebuilt.py
+ # both honor UNSLOTH_ROCM_GFX_ARCH, so exporting it is the
+ # whole handoff (a user-set override re-exports unchanged).
+ export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"
+ case "$_linux_inferred_gfx" in
+ gfx1201|gfx1200|gfx1151|gfx1150)
+ TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
+ TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
+ TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
+ ;;
+ esac
+ echo "" >&2
+ echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2
+ echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2
+ echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2
+ echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2
+ echo " [WARN] Tip: set UNSLOTH_ROCM_GFX_ARCH=$_linux_inferred_gfx to skip inference next time." >&2
+ echo "" >&2
+ fi
+ fi
+ ;;
+ esac
+fi
+
# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that
# downstream scripts (setup.sh -> install_python_stack.py) know what was
# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts.
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index bb329e189e..a29ba0d7e5 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -769,6 +769,142 @@ def _gfx_arch_from_gpu_name(name: str) -> "str | None":
return None
+def _linux_amd_gfx_from_cpuinfo() -> "str | None":
+ """Infer gfx arch from /proc/cpuinfo on integrated AMD APUs (Strix Halo/Point)."""
+ try:
+ text = Path("/proc/cpuinfo").read_text(encoding = "utf-8", errors = "replace")
+ except OSError:
+ return None
+ if re.search(r"Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo", text, re.IGNORECASE):
+ return "gfx1151"
+ if re.search(
+ r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]"
+ r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33",
+ text,
+ re.IGNORECASE,
+ ):
+ return "gfx1150"
+ return None
+
+
+def _linux_amd_gfx_from_lspci() -> "str | None":
+ """First AMD display-class lspci line mapping to a known gfx arch. A non-AMD
+ controller can enumerate first (Intel/ASPEED before an AMD dGPU), so scan
+ them all. The vendor guard is case-SENSITIVE: a -i "ATI" would match
+ "CorporATIon" on every Intel/NVIDIA line. Whole-line matching also survives
+ the 0000: PCI domain prefix."""
+ lspci = shutil.which("lspci")
+ if not lspci:
+ return None
+ try:
+ result = subprocess.run(
+ [lspci, "-nn"],
+ stdout = subprocess.PIPE,
+ stderr = subprocess.DEVNULL,
+ text = True,
+ timeout = 10,
+ )
+ except Exception:
+ return None
+ if result.returncode != 0:
+ return None
+ for line in result.stdout.splitlines():
+ if not re.search(r"VGA compatible controller|3D controller|Display controller", line, re.I):
+ continue
+ if not re.search(r"AMD|ATI", line):
+ continue
+ arch = _gfx_arch_from_gpu_name(line)
+ if arch:
+ return arch
+ return None
+
+
+def _is_wsl() -> bool:
+ """True on WSL, where the AMD GPU is reached via /dev/dxg (not /dev/kfd)."""
+ if os.path.exists("/dev/dxg"):
+ return True
+ try:
+ with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
+ return "microsoft" in fh.read().lower()
+ except OSError:
+ return False
+
+
+def _wsl_rocm_runtime_present() -> bool:
+ """librocdxg (the WSL ROCDXG bridge that lets HIP reach the GPU over /dev/dxg)
+ under a ROCm lib dir. Its absence marks a WSL box whose ROCm was never set up."""
+ dirs = ["/opt/rocm/lib", "/opt/rocm/lib64"]
+ dirs += glob.glob("/opt/rocm-*/lib") + glob.glob("/opt/rocm-*/lib64")
+ return any(
+ os.path.exists(os.path.join(d, so))
+ for d in dirs
+ for so in ("librocdxg.so", "librocdxg.so.1")
+ )
+
+
+def _linux_amd_display_device_present() -> bool:
+ """Any AMD (vendor 0x1002) PCI display-class (0x03*) device in sysfs.
+ /proc/cpuinfo leaks the HOST CPU model into VMs/containers that received no
+ AMD GPU, so the CPU-model text alone is not GPU evidence; this is the
+ device-level check (mirrors install.sh _amd_gpu_present_via_pci)."""
+ try:
+ for dev in Path("/sys/bus/pci/devices").iterdir():
+ try:
+ if (dev / "vendor").read_text().strip() != "0x1002":
+ continue
+ if (dev / "class").read_text().strip().startswith("0x03"):
+ return True
+ except OSError:
+ continue
+ except OSError:
+ pass
+ return False
+
+
+def _infer_linux_amd_gfx_arch() -> "str | None":
+ """Infer gfx when ROCm runtime is absent but the host is a known AMD arch (unslothai#7301)."""
+ override = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower()
+ if override:
+ return override
+ if _is_wsl():
+ # cpuinfo/lspci see the host APU even on a WSL box whose ROCDXG runtime
+ # was never bootstrapped; inferring there would install per-arch ROCm
+ # wheels into an env that still can't expose the GPU. Skip unless that
+ # runtime is present -- WSL enumerates no PCI display device, so
+ # /dev/dxg + librocdxg IS the GPU evidence there.
+ if not _wsl_rocm_runtime_present():
+ return None
+ elif not _linux_amd_display_device_present():
+ # Native Linux: a VM/container on a Strix host still shows the host CPU
+ # model in /proc/cpuinfo while receiving no AMD GPU, so require an AMD
+ # display device before trusting the CPU-model inference. The lspci
+ # fallback reads the same PCI space and would find nothing here either.
+ return None
+ cpu_gfx = _linux_amd_gfx_from_cpuinfo()
+ if cpu_gfx:
+ return cpu_gfx
+ return _linux_amd_gfx_from_lspci()
+
+
+def _amd_arch_index_url(gfx_arch: str | None) -> str | None:
+ """Return the AMD per-arch pip index URL for a gfx arch (Linux + Windows).
+
+ Windows honors UNSLOTH_ROCM_WINDOWS_MIRROR (via _windows_rocm_index_url);
+ Linux honors UNSLOTH_AMD_ROCM_MIRROR -- the same var install.sh uses -- so a
+ mirrored/air-gapped Linux repair reaches the index install.sh chose rather
+ than falling back to repo.amd.com. Both default to repo.amd.com when unset.
+ """
+ if IS_WINDOWS:
+ return _windows_rocm_index_url(gfx_arch)
+ arch_family = _GFX_TO_AMD_INDEX_ARCH.get(gfx_arch or "")
+ if arch_family is None:
+ return None
+ base = (os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") or "https://repo.amd.com/rocm/whl").rstrip(
+ "/"
+ )
+ return f"{base}/{arch_family}/"
+
+
def _windows_rocm_index_url(gfx_arch: str | None) -> str | None:
"""Return the AMD pip index URL for the given GPU arch, or None if unsupported."""
arch_family = _GFX_TO_AMD_INDEX_ARCH.get(gfx_arch or "")
@@ -1647,22 +1783,24 @@ def _ensure_rocm_torch() -> None:
# An explicit ROCm pin commits to ROCm wheels regardless of the visible GPU (headless / CI).
# Mirror _ensure_cuda_torch: skip the NVIDIA/no-AMD/unreadable gates.
_rocm_pin = _explicit_rocm_torch_index_url()
+ _inferred_linux_gfx = (
+ _infer_linux_amd_gfx_arch() if (_rocm_pin is None and not IS_WINDOWS) else None
+ )
if _rocm_pin is None:
# NVIDIA takes precedence on mixed hosts (only if a GPU is usable).
if _has_usable_nvidia_gpu():
return
# _has_rocm_gpu() (rocminfo / amd-smi rows) is the authoritative AMD-host signal;
# the old /opt/rocm-or-hipcc gate broke runtime-only ROCm installs.
- if not _has_rocm_gpu():
+ if not _has_rocm_gpu() and not _inferred_linux_gfx:
return # no AMD GPU visible
ver = _detect_rocm_version()
if ver is None:
- if _rocm_pin is None:
+ if _rocm_pin is None and not _inferred_linux_gfx:
print(" ROCm detected but version unreadable -- skipping torch reinstall")
return
- # Explicit pin: the pinned leaf drives the install, so an unreadable host version
- # is fine (sentinel keeps ver comparisons defined).
+ # Explicit pin or inferred gfx: the index drives the install.
ver = (0, 0)
# Probe whether torch links against HIP, capturing the installed ROCm tag for pin-mismatch
@@ -1712,6 +1850,44 @@ def _ensure_rocm_torch() -> None:
rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch
+ # Inferred-gfx path: ROCm runtime missing but install.sh would route to AMD wheels.
+ # Gated on the runtime NOT enumerating a GPU: when it can, the runtime-visible
+ # arch (Strix override / generic below) decides, not cpuinfo -- a mixed Strix
+ # APU + dGPU box with HIP_VISIBLE_DEVICES on the dGPU must not get APU wheels.
+ # An explicit UNSLOTH_ROCM_GFX_ARCH is exempt from that runtime gate (mirrors
+ # install.sh): a visible GPU with an unreadable/unsupported ROCm version must
+ # not silently discard the user's named arch and leave CPU torch in place.
+ _gfx_override_env = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower()
+ if (
+ _inferred_linux_gfx
+ and not has_hip_torch
+ and _rocm_pin is None
+ and (_gfx_override_env or not _has_rocm_gpu())
+ ):
+ index_url = _amd_arch_index_url(_inferred_linux_gfx)
+ if index_url is not None:
+ _torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get(
+ _inferred_linux_gfx, ("torch", "torchvision", "torchaudio")
+ )
+ print(
+ f"\n {_inferred_linux_gfx} inferred (ROCm runtime not visible) -- "
+ f"installing torch from {_strip_index_url_credentials(index_url)}\n"
+ f" AMD wheels bundle their own ROCm runtime; install the kernel stack "
+ f"for native GPU compute.\n"
+ )
+ pip_install(
+ f"ROCm torch (inferred {_inferred_linux_gfx})",
+ "--force-reinstall",
+ "--no-cache-dir",
+ _torch_pkg,
+ _vision_pkg,
+ _audio_pkg,
+ "--index-url",
+ index_url,
+ constrain = False,
+ )
+ rocm_torch_ready = True
+
# Strix Halo / Point (gfx1151 / gfx1150) need torch from AMD's per-gfx index
# (2.11+rocm7.13); any generic pytorch.org rocm index lacks the fixes (ROCm 7.1
# segfaults in _grouped_mm). See _strix_needs_amd_arch_index for the floor gate.
@@ -1776,8 +1952,11 @@ def _ensure_rocm_torch() -> None:
constrain = False,
)
rocm_torch_ready = True
- elif not has_hip_torch or _rocm_pin_mismatch:
+ elif not rocm_torch_ready:
# Reinstall when torch is not ROCm yet, OR a ROCm build's family differs from a pin.
+ # Gate on rocm_torch_ready (not has_hip_torch alone) so a successful inferred-gfx
+ # install above is not overwritten by the generic pytorch.org/rocmX.Y path -- that
+ # would undo the fresh-ROCm/no-/dev/kfd repair this path exists for (Codex P1 #7305).
# Honour a ROCm pin verbatim; else pick the newest wheel tag <= host.
_override_idx = _explicit_rocm_torch_index_url()
if _override_idx is not None:
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index b343b07238..cd7b68f4b6 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -9,6 +9,7 @@ import subprocess
import sys
import tempfile
from pathlib import Path
+from types import SimpleNamespace
from unittest.mock import MagicMock, mock_open, patch, PropertyMock
import pytest
@@ -560,9 +561,13 @@ class TestDetectRocmVersion:
class TestEnsureRocmTorch:
"""Verify ROCm torch reinstall logic."""
+ # _infer_linux_amd_gfx_arch mocked to None: on a real Strix host the live
+ # /proc/cpuinfo would otherwise take the inferred-install path and break
+ # these "must not install" hosts (environment leak, not the code under test).
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
- def test_no_rocm_skips(self, mock_nvidia, mock_pip):
+ @patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = None)
+ def test_no_rocm_skips(self, mock_infer, mock_nvidia, mock_pip):
"""No ROCm toolchain should skip entirely."""
# Pin _detect_windows_gfx_arch to None so a real AMD test host's WMI
# fallback can't defeat the "no ROCm anywhere" premise.
@@ -572,6 +577,105 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
+ @patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = "gfx1151")
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = None)
+ def test_inferred_gfx_without_rocm_runtime_installs_amd_index(
+ self, mock_ver, mock_infer, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """Strix Halo without /dev/kfd must still get AMD gfx1151 wheels (unslothai#7301)."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"|2.10.0+cpu\n"
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "gfx1151" in torch_call
+ assert "torch>=2.11.0,<2.12.0" in torch_call
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
+ @patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = "gfx1151")
+ @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = [])
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
+ def test_inferred_gfx_not_overwritten_when_rocm_userland_readable(
+ self, mock_ver, mock_gfx, mock_infer, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """Codex P1 #7305: after an inferred per-arch install, do not fall through to the
+ generic pytorch.org/rocmX.Y reinstall just because has_hip_torch is still False.
+ Readable ROCm userland without /dev/kfd is exactly the case that used to overwrite
+ the AMD gfx wheels."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"|2.10.0+cpu\n"
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ assert mock_pip.call_count == 1, mock_pip.call_args_list
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "gfx1151" in torch_call
+ assert "rocm7.1" not in torch_call
+ assert "download.pytorch.org" not in torch_call
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = "gfx1151")
+ @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1100"])
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
+ def test_inference_yields_to_runtime_visible_gpu(
+ self, mock_ver, mock_gfx, mock_infer, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """When the runtime CAN enumerate a GPU, the cpuinfo inference must not
+ install wheels: a mixed Strix APU + dGPU box with the dGPU selected would
+ otherwise get gfx1151 wheels for a gfx1100 GPU. The runtime-visible arch
+ (Strix override / generic branch) decides instead."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"|2.10.0+cpu\n"
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ all_calls = str(mock_pip.call_args_list) + str(mock_pip_try.call_args_list)
+ assert "gfx1151" not in all_calls, all_calls
+ assert "rocm7.1" in all_calls, all_calls
+
+ @patch.object(stack_mod, "IS_WINDOWS", False)
+ @patch.object(stack_mod, "pip_install_try", return_value = True)
+ @patch.object(stack_mod, "pip_install")
+ @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+ @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = [])
+ @patch.object(stack_mod, "_detect_rocm_version", return_value = None)
+ def test_gfx_override_installs_despite_visible_rocm(
+ self, mock_ver, mock_gfx, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+ ):
+ """#7305 review: an explicit UNSLOTH_ROCM_GFX_ARCH is exempt from the
+ not-_has_rocm_gpu() gate (mirrors install.sh). A visible GPU with an
+ unreadable ROCm version must not silently discard the user's named arch
+ and leave CPU torch in place -- the per-arch install runs."""
+ mock_probe = MagicMock()
+ mock_probe.returncode = 0
+ mock_probe.stdout = b"|2.10.0+cpu\n"
+ with patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": "gfx1151"}):
+ with patch("os.path.isdir", return_value = True):
+ with patch("subprocess.run", return_value = mock_probe):
+ _ensure_rocm_torch()
+ assert mock_pip.call_count == 1, mock_pip.call_args_list
+ torch_call = str(mock_pip.call_args_list[0])
+ assert "gfx1151" in torch_call
+ assert "download.pytorch.org" not in torch_call
+
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@@ -683,9 +787,10 @@ class TestEnsureRocmTorch:
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+ @patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = None)
@patch.object(stack_mod, "_detect_rocm_version", return_value = None)
def test_version_unreadable_prints_warning(
- self, mock_ver, mock_gpu, mock_nvidia, mock_pip, capsys
+ self, mock_ver, mock_infer, mock_gpu, mock_nvidia, mock_pip, capsys
):
"""ROCm detected but version unreadable should print warning and skip."""
with patch("os.path.isdir", return_value = True):
@@ -1042,7 +1147,8 @@ class TestEnsureRocmTorch:
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
- def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip):
+ @patch.object(stack_mod, "_infer_linux_amd_gfx_arch", return_value = None)
+ def test_no_gpu_with_rocm_tools_skips(self, mock_infer, mock_gpu, mock_nvidia, mock_pip):
"""ROCm tools present but no actual AMD GPU should skip entirely."""
# Pin the Windows arch probe to None so a real AMD host's WMI fallback
# can't defeat the "no actual GPU" premise.
@@ -2122,6 +2228,7 @@ class TestGfxArchNameFallback:
"name, expected",
[
("AMD Radeon(TM) 8060S Graphics", "gfx1151"),
+ ("AMD Radeon(TM) 8065S Graphics", "gfx1151"),
("AMD Ryzen AI MAX+ 395 w/ Radeon 8060S", "gfx1151"),
("AMD Radeon(TM) 890M", "gfx1150"),
("AMD Ryzen AI 9 HX 370 w/ Radeon 890M", "gfx1150"),
@@ -3189,6 +3296,286 @@ _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh"
class TestStrixRocm71Override:
"""install.sh routes gfx1151/gfx1150 to AMD's arch index instead of ROCm 7.1 (_grouped_mm segfault)."""
+ def test_linux_gfx_inference_helpers_present(self):
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ assert "_infer_linux_amd_gfx_arch" in source
+ assert "_amd_arch_index_family_for_gfx" in source
+ assert "_amd_gpu_present_via_pci" in source
+ assert "unslothai#7301" in source
+
+ def test_infer_linux_amd_gfx_from_cpuinfo(self):
+ assert stack_mod._linux_amd_gfx_from_cpuinfo is not None
+ with patch.object(
+ Path,
+ "read_text",
+ return_value = "model name : AMD Ryzen AI Max+ 395 w/ Radeon 8060S\n",
+ ):
+ assert stack_mod._linux_amd_gfx_from_cpuinfo() == "gfx1151"
+ # 8065S (Gorgon Halo) must match on the Radeon name alone, even without the
+ # "Ryzen AI Max" branding (mirrors setup.sh / setup.ps1 which list 8065S).
+ with patch.object(Path, "read_text", return_value = "model name : AMD Radeon 8065S\n"):
+ assert stack_mod._linux_amd_gfx_from_cpuinfo() == "gfx1151"
+
+ def test_infer_gfx_gated_out_of_wsl_without_runtime(self):
+ """On WSL the cpuinfo/lspci inference must be skipped unless the WSL ROCDXG
+ runtime (librocdxg) is present: a bare `unsloth studio update` must not
+ install per-arch ROCm wheels into an env that still can't expose the GPU.
+ An explicit UNSLOTH_ROCM_GFX_ARCH override stays authoritative regardless."""
+ m = stack_mod
+ with (
+ patch.object(m, "_linux_amd_gfx_from_cpuinfo", return_value = "gfx1151"),
+ patch.object(m, "_linux_amd_gfx_from_lspci", return_value = None),
+ # PCI evidence present (the WSL branch never consults it anyway).
+ patch.object(m, "_linux_amd_display_device_present", return_value = True),
+ patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": ""}),
+ ):
+ # WSL + no runtime -> inference suppressed (CPU torch stays).
+ with (
+ patch.object(m, "_is_wsl", return_value = True),
+ patch.object(m, "_wsl_rocm_runtime_present", return_value = False),
+ ):
+ assert m._infer_linux_amd_gfx_arch() is None
+ # WSL + runtime present (this dev box) -> inference still runs.
+ with (
+ patch.object(m, "_is_wsl", return_value = True),
+ patch.object(m, "_wsl_rocm_runtime_present", return_value = True),
+ ):
+ assert m._infer_linux_amd_gfx_arch() == "gfx1151"
+ # Native Linux (not WSL) -> the gate never applies.
+ with (
+ patch.object(m, "_is_wsl", return_value = False),
+ patch.object(m, "_wsl_rocm_runtime_present", return_value = False),
+ ):
+ assert m._infer_linux_amd_gfx_arch() == "gfx1151"
+ # Explicit override wins even on a bare WSL box (no runtime).
+ with (
+ patch.object(m, "_is_wsl", return_value = True),
+ patch.object(m, "_wsl_rocm_runtime_present", return_value = False),
+ patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": "gfx1151"}),
+ ):
+ assert m._infer_linux_amd_gfx_arch() == "gfx1151"
+
+ def test_infer_gfx_requires_amd_display_device_on_native_linux(self):
+ """A VM/container on a Strix host still shows the host CPU model in
+ /proc/cpuinfo while receiving no AMD GPU, so on native Linux the
+ CPU-model inference must require an AMD PCI display device (#7305
+ review). WSL is exempt (no PCI enumeration there; the librocdxg gate is
+ the evidence) and the explicit override stays authoritative."""
+ m = stack_mod
+ with (
+ patch.object(m, "_linux_amd_gfx_from_cpuinfo", return_value = "gfx1151"),
+ patch.object(m, "_linux_amd_gfx_from_lspci", return_value = None),
+ patch.object(m, "_is_wsl", return_value = False),
+ patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": ""}),
+ ):
+ # No AMD display device -> the CPU-model text alone must not infer.
+ with patch.object(m, "_linux_amd_display_device_present", return_value = False):
+ assert m._infer_linux_amd_gfx_arch() is None
+ # Device present -> inference unchanged.
+ with patch.object(m, "_linux_amd_display_device_present", return_value = True):
+ assert m._infer_linux_amd_gfx_arch() == "gfx1151"
+ # Explicit override needs no device evidence (headless/cross-install).
+ with (
+ patch.object(m, "_is_wsl", return_value = False),
+ patch.object(m, "_linux_amd_display_device_present", return_value = False),
+ patch.dict(os.environ, {"UNSLOTH_ROCM_GFX_ARCH": "GFX1151"}),
+ ):
+ assert m._infer_linux_amd_gfx_arch() == "gfx1151"
+
+ def test_install_sh_cpuinfo_inference_requires_pci_evidence(self):
+ """install.sh mirror of the VM/container guard: both cpuinfo greps must be
+ gated on _gpu_evidence (AMD PCI display device via _amd_gpu_present_via_pci,
+ or the WSL librocdxg gate), and the gate must sit before the first grep."""
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ body = _extract_sh_function_body(source, "_infer_linux_amd_gfx_arch")
+ assert body, "could not extract _infer_linux_amd_gfx_arch"
+ pci = body.find("_amd_gpu_present_via_pci")
+ infer = body.find("grep -qiE 'Ryzen AI Max")
+ assert pci >= 0 and infer >= 0
+ assert pci < infer, "the PCI evidence check must run before the cpuinfo inference"
+ assert (
+ body.count('[ -n "$_gpu_evidence" ] && grep -qiE') == 2
+ ), "both cpuinfo greps (gfx1151 and gfx1150) must be gated on _gpu_evidence"
+
+ def test_lspci_scan_covers_all_display_controllers(self):
+ """The lspci fallback must scan every display-class line, not just the
+ first: a non-AMD controller (Intel iGPU, ASPEED BMC) often enumerates
+ before the AMD dGPU. Non-AMD vendors must never map (an NVIDIA GeForce
+ GTX 860M would otherwise hit the AMD 860M pattern), and a 0000: PCI
+ domain prefix must not break matching."""
+ m = stack_mod
+
+ def fake_lspci(stdout):
+ result = SimpleNamespace(returncode = 0, stdout = stdout)
+ return (
+ patch.object(m.shutil, "which", return_value = "/usr/bin/lspci"),
+ patch.object(m.subprocess, "run", return_value = result),
+ )
+
+ intel_then_amd = (
+ "00:02.0 VGA compatible controller [0300]: Intel Corporation Raptor Lake-S GT1 [8086:a780]\n"
+ "03:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
+ " Navi 31 [Radeon RX 7900 XT] [1002:744c]\n"
+ )
+ nvidia_only = "01:00.0 3D controller [0302]: NVIDIA Corporation GM107M [GeForce GTX 860M] [10de:1392]\n"
+ domain_prefixed = (
+ "0000:c5:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
+ " Strix Halo [Radeon Graphics / Radeon 8060S] [1002:150e]\n"
+ )
+ unmapped_then_mapped = (
+ "03:00.0 Display controller [0380]: Advanced Micro Devices, Inc. [AMD/ATI]"
+ " Cape Verde [FirePro W600] [1002:6821]\n"
+ "04:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
+ " Navi 33 [Radeon RX 7600] [1002:7480]\n"
+ )
+ for stdout, expected in (
+ (intel_then_amd, "gfx1100"),
+ (nvidia_only, None),
+ (domain_prefixed, "gfx1151"),
+ (unmapped_then_mapped, "gfx1102"),
+ ):
+ w, r = fake_lspci(stdout)
+ with w, r:
+ assert m._linux_amd_gfx_from_lspci() == expected, stdout
+
+ def test_install_sh_lspci_scan_covers_all_display_controllers(self):
+ """install.sh mirror of the scan-all behaviour, executed with a shimmed
+ lspci: Intel-first still finds the AMD dGPU, NVIDIA-only maps nothing
+ (860M collision), a domain-prefixed AMD line still maps."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute the probe block")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ name_fn = re.search(
+ r"^_infer_amd_gfx_arch_from_gpu_name\(\) \{\n.*?\n\}\n", source, re.S | re.M
+ )
+ scan = re.search(
+ r"^ if command -v lspci[^\n]*\n.*?\nEOF\n fi\n return 1\n", source, re.S | re.M
+ )
+ assert name_fn and scan, "could not extract the lspci scan block"
+ cases = (
+ (
+ "00:02.0 VGA compatible controller [0300]: Intel Corporation UHD [8086:a780]\n"
+ "03:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. [AMD/ATI]"
+ " Navi 31 [Radeon RX 7900 XT] [1002:744c]",
+ "OK:gfx1100",
+ ),
+ (
+ "01:00.0 3D controller [0302]: NVIDIA Corporation GM107M [GeForce GTX 860M] [10de:1392]",
+ "OK:",
+ ),
+ (
+ "0000:c5:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc."
+ " [AMD/ATI] Strix Halo [Radeon 8060S] [1002:150e]",
+ "OK:gfx1151",
+ ),
+ )
+ for lspci_out, expected in cases:
+ with tempfile.TemporaryDirectory() as d:
+ p = os.path.join(d, "lspci")
+ with open(p, "w", encoding = "utf-8") as f:
+ f.write(f'#!/bin/sh\ncat <<"EOT"\n{lspci_out}\nEOT\n')
+ os.chmod(p, 0o755)
+ script = (
+ "set -euo pipefail\n"
+ + name_fn.group(0)
+ + "probe() {\n"
+ + scan.group(0)
+ + "}\nprintf 'OK:%s\\n' \"$(probe || true)\"\n"
+ )
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""))
+ r = subprocess.run([shell, "-c", script], env = env, capture_output = True, text = True)
+ assert r.returncode == 0, f"scan aborted: {r.stderr}"
+ assert (
+ r.stdout.splitlines()[-1] == expected
+ ), f"lspci scan wrong for {lspci_out!r}: {r.stdout!r}"
+
+ def test_install_sh_infer_gfx_gated_on_wsl_runtime(self):
+ """install.sh's _infer_linux_amd_gfx_arch must, like the Python side, skip
+ the cpuinfo/lspci inference on WSL unless librocdxg is present -- the
+ override still returns first, so it stays authoritative."""
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ body = _extract_sh_function_body(source, "_infer_linux_amd_gfx_arch")
+ assert body, "could not extract _infer_linux_amd_gfx_arch"
+ override = body.find("UNSLOTH_ROCM_GFX_ARCH")
+ dxg = body.find("/dev/dxg")
+ rocdxg = body.find("librocdxg")
+ # Anchor on the first cpuinfo *inference* (the grep), not a comment mention.
+ infer = body.find("grep -qiE 'Ryzen AI Max")
+ assert override >= 0 and dxg >= 0 and rocdxg >= 0 and infer >= 0
+ assert "microsoft" in body, "WSL gate must also detect WSL via /proc/version"
+ assert override < dxg, "the explicit override must return before the WSL gate"
+ assert (
+ dxg < infer and rocdxg < infer
+ ), "the WSL/librocdxg gate must run before the cpuinfo/lspci inference"
+
+ def test_install_sh_reroute_is_x86_64_only(self):
+ """The Linux inferred-gfx reroute must be x86_64-only: ROCm torch wheels are
+ not published for arm64, so an inferred/overridden gfx must not push an
+ arm64 host to the AMD arch index (get_torch_index_url returns CPU there)."""
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ idx = source.find("_linux_inferred_gfx=$(_infer_linux_amd_gfx_arch")
+ assert idx >= 0, "reroute consumer not found"
+ window = source[max(0, idx - 400) : idx]
+ assert (
+ 'case "$_ARCH" in x86_64|amd64)' in window
+ ), "the inferred-gfx reroute must guard on x86_64|amd64 arch"
+
+ def test_install_sh_reroute_skips_visible_rocm_gpu(self):
+ """A */cpu index on a host whose AMD GPU IS visible to the ROCm probes is a
+ deliberate fallback (unsupported/unreadable ROCm version, warned about in
+ get_torch_index_url), not a missing runtime: the reroute must not override
+ it with inferred per-arch wheels. The explicit UNSLOTH_ROCM_GFX_ARCH
+ override must still win either way."""
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ idx = source.find("_linux_inferred_gfx=$(_infer_linux_amd_gfx_arch")
+ assert idx >= 0, "reroute consumer not found"
+ window = source[max(0, idx - 700) : idx]
+ assert (
+ "! _has_amd_rocm_gpu" in window
+ ), "the reroute must be gated on _has_amd_rocm_gpu being false"
+ assert (
+ '[ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu' in window
+ ), "an explicit UNSLOTH_ROCM_GFX_ARCH override must bypass the visible-GPU gate"
+
+ def test_install_sh_reroute_exports_gfx_for_setup_sh(self):
+ """The inferred arch must be exported as UNSLOTH_ROCM_GFX_ARCH so the
+ downstream setup.sh run (which re-probes ROCm independently and finds
+ nothing on these runtime-less hosts) routes llama.cpp to the matching
+ ROCm prebuilt instead of the CPU one -- setup.sh and
+ install_llama_prebuilt.py both read that env var."""
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ assign = source.find('TORCH_INDEX_URL="${_amd_mirror}/${_amd_family}/"')
+ assert assign >= 0, "inferred-gfx index assignment not found"
+ block_end = source.find("esac", assign)
+ assert (
+ 'export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx"' in source[assign:block_end]
+ ), "the reroute must export the inferred gfx for the setup.sh handoff"
+ # setup.sh's side of the handoff must still exist.
+ setup_source = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8")
+ assert "UNSLOTH_ROCM_GFX_ARCH" in setup_source
+
+ def test_amd_arch_index_url_linux_honors_amd_mirror(self):
+ """On Linux the inferred-gfx repair must honour UNSLOTH_AMD_ROCM_MIRROR (the
+ var install.sh uses), not the Windows mirror var, so a mirrored/air-gapped
+ Linux install does not silently fall back to repo.amd.com. Windows still
+ delegates to the Windows mirror path."""
+ m = stack_mod
+ with (
+ patch.object(m, "IS_WINDOWS", False),
+ patch.dict(os.environ, {"UNSLOTH_AMD_ROCM_MIRROR": "https://mirror.local/rocm"}),
+ ):
+ assert m._amd_arch_index_url("gfx1151") == "https://mirror.local/rocm/gfx1151/"
+ with (
+ patch.object(m, "IS_WINDOWS", False),
+ patch.dict(os.environ, {"UNSLOTH_AMD_ROCM_MIRROR": ""}),
+ ):
+ assert m._amd_arch_index_url("gfx1151") == "https://repo.amd.com/rocm/whl/gfx1151/"
+ assert m._amd_arch_index_url("gfx9999") is None
+ # Windows path is unchanged: delegate to the Windows mirror helper.
+ with patch.object(m, "IS_WINDOWS", True):
+ assert m._amd_arch_index_url("gfx1151") == m._windows_rocm_index_url("gfx1151")
+
def test_strix_gfx_detection_in_install_sh(self):
"""install.sh must detect gfx1151 and gfx1150 for the override."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
From 6f4c838281cef13bbb038426d3fdf53bb34c22de Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Thu, 23 Jul 2026 01:55:45 -0300
Subject: [PATCH 075/255] Studio: calibrate Linux chat typography against macOS
(#7337)
---
studio/frontend/src/index.css | 13 ++++-
tests/studio/playwright_chat_ui.py | 82 ++++++++++++++++++++++++++++++
2 files changed, 93 insertions(+), 2 deletions(-)
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 52ca81e064..1fafe09d17 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -633,11 +633,20 @@ html.no-font-smoothing body {
-moz-osx-font-smoothing: auto;
}
-/* Match Inter's lighter macOS rendering. Keep 410 when smoothing is off or a
- custom font reaches chat. */
+/* Match Inter's lighter macOS rendering. Dark surfaces need a stronger
+ correction than light surfaces. Keep 410 when smoothing is off or a custom
+ font reaches chat. */
html.render-linux:not(.no-font-smoothing):not([data-chat-font]):not([data-ui-font])
+ :is(.aui-assistant-message-root, .aui-user-message-root) {
+ font-weight: 390;
+}
+
+html.dark.render-linux:not(.no-font-smoothing):not([data-chat-font]):not([data-ui-font])
:is(.aui-assistant-message-root, .aui-user-message-root) {
font-weight: 350;
+ /* The lighter variable-font instance has narrower advances. Reduce
+ dark-mode line-wrap drift without changing custom-font paths. */
+ letter-spacing: 0.023em;
}
/* Chat font: only applies while a custom chat font is set. Elements with
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index 4d13889878..a06e559100 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.py
@@ -936,6 +936,70 @@ with sync_playwright() as p:
page.keyboard.press("Escape")
page.wait_for_timeout(300)
+ def read_chat_typography():
+ """Read message typography after a user-driven theme transition."""
+ return robust_evaluate(
+ page,
+ """() => {
+ const root = document.documentElement;
+ const assistant = Array.from(
+ document.querySelectorAll('.aui-assistant-message-root')
+ );
+ const user = Array.from(
+ document.querySelectorAll('.aui-user-message-root')
+ );
+ if (assistant.length === 0 || user.length === 0) {
+ return { error: 'chat message roots are missing' };
+ }
+ const ua = navigator.userAgent.toLowerCase();
+ const role = (nodes) => {
+ const styles = nodes.map((node) => getComputedStyle(node));
+ return {
+ fontWeight: [...new Set(styles.map((style) => style.fontWeight))],
+ letterSpacing: [...new Set(styles.map((style) => style.letterSpacing))],
+ };
+ };
+ return {
+ actualRenderLinux: root.classList.contains('render-linux'),
+ isDesktopLinux: ua.includes('linux') && !ua.includes('android'),
+ isDark: root.classList.contains('dark'),
+ usesBaselineTypography: (
+ root.classList.contains('no-font-smoothing') ||
+ root.hasAttribute('data-chat-font') ||
+ root.hasAttribute('data-ui-font')
+ ),
+ assistant: role(assistant),
+ user: role(user),
+ };
+ }""",
+ )
+
+ def assert_chat_typography(label, typography):
+ if typography.get("error"):
+ fail(typography["error"])
+ if typography["actualRenderLinux"] != typography["isDesktopLinux"]:
+ fail(f"desktop Linux detection mismatch: {typography!r}")
+ is_dark = typography["isDark"]
+ expected_spacing = "0.31px" if is_dark else "0.155px"
+ if typography["isDesktopLinux"] and not typography["usesBaselineTypography"]:
+ expected_weight = "350" if is_dark else "390"
+ if is_dark:
+ expected_spacing = "0.3565px"
+ else:
+ expected_weight = "410"
+ for role in ("assistant", "user"):
+ actual = typography[role]
+ if actual["fontWeight"] != [expected_weight]:
+ fail(
+ f"chat font weight {label}/{role}: expected {expected_weight}, "
+ f"got {actual['fontWeight']!r}"
+ )
+ if actual["letterSpacing"] != [expected_spacing]:
+ fail(
+ f"chat letter spacing {label}/{role}: expected {expected_spacing}, "
+ f"got {actual['letterSpacing']!r}"
+ )
+
# ─────────────────────────────────────────────────────
# 9. Theme toggle -- multiple cycles + computed-bg-color check
# (light is near-white >240; dark is near-black <40).
@@ -944,6 +1008,7 @@ with sync_playwright() as p:
if acct.count() > 0:
step("theme toggle x3 with computed-color assertion")
observed = []
+ typography_states = []
for cycle in range(3):
# Wait for any prior dropdown to fully detach: clicking while
# the view-transition is still open no-ops silently. The
@@ -1032,6 +1097,9 @@ with sync_playwright() as p:
}""",
)
observed.append(bg)
+ typography = read_chat_typography()
+ assert_chat_typography(f"theme-cycle-{cycle + 1}", typography)
+ typography_states.append(typography)
shoot(f"10-theme-cycle-{cycle + 1}")
info(f" cycle {cycle + 1}: dark={bg['isDark']} body bg={bg['bg']!r}")
# Across cycles we should see both a near-white (light) and a
@@ -1054,6 +1122,20 @@ with sync_playwright() as p:
"(toggle may not flip on this runner's color-scheme)"
)
+ # These are user-driven theme transitions, not synthetic class
+ # changes. A completed three-cycle toggle must expose both typography
+ # states before we check the Linux selector.
+ if len(typography_states) != 3:
+ soft_fail(
+ f"chat typography observed {len(typography_states)} theme state(s), expected 3"
+ )
+ elif {state["isDark"] for state in typography_states} != {False, True}:
+ soft_fail(f"chat typography did not observe both themes: {typography_states!r}")
+ else:
+ info("OK chat typography platform and theme behavior")
+ else:
+ soft_fail("chat typography requires the account-menu theme control")
+
# ─────────────────────────────────────────────────────
# 10. Sidebar nav: New Chat, Compare, Search, Recipes.
# ─────────────────────────────────────────────────────
From d59c7bfd03c8fd93f194c91ac8307081349bab6d Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Thu, 23 Jul 2026 01:56:14 -0300
Subject: [PATCH 076/255] Studio: prevent login error text clipping (#7343)
---
studio/frontend/src/features/auth/components/auth-form.tsx | 6 +++++-
1 file changed, 5 insertions(+), 1 deletion(-)
diff --git a/studio/frontend/src/features/auth/components/auth-form.tsx b/studio/frontend/src/features/auth/components/auth-form.tsx
index 73db10d41b..3eec1dba88 100644
--- a/studio/frontend/src/features/auth/components/auth-form.tsx
+++ b/studio/frontend/src/features/auth/components/auth-form.tsx
@@ -439,7 +439,11 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
{helperText && (
{helperText}
)}
- {error && {error}
}
+ {error && (
+
+ {error}
+
+ )}
Date: Thu, 23 Jul 2026 13:09:14 +0530
Subject: [PATCH 077/255] Studio: fix stuck composer prompt on first send and
unreachable --secure Cloudflare links (#7340)
* Studio: clear composer draft on send
* Studio: verify the Cloudflare link is reachable before printing it
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: wait for tunnel DNS propagation before verifying the public URL
* Studio: bound tunnel DNS wait and health probe by one deadline
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep composer draft when overlay send validation fails
* Studio: retry transient DoH failures while waiting for tunnel DNS
* [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/cloudflare_tunnel.py | 88 ++++++-
.../backend/tests/test_cloudflare_tunnel.py | 223 ++++++++++++++++++
.../src/components/assistant-ui/thread.tsx | 22 +-
3 files changed, 327 insertions(+), 6 deletions(-)
diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py
index b1ddc74c32..78fce0c70a 100644
--- a/studio/backend/cloudflare_tunnel.py
+++ b/studio/backend/cloudflare_tunnel.py
@@ -20,6 +20,7 @@ import shutil
import subprocess
import sys
import threading
+import time
from pathlib import Path
from typing import Optional, Tuple
@@ -40,6 +41,22 @@ _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/downl
_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
+# A registered edge connection does not mean the hostname resolves yet, so the
+# URL is fetched once before it is advertised.
+_PUBLIC_PROBE_PATH = "/api/health"
+_PUBLIC_PROBE_MARKER = "Unsloth UI Backend"
+# One deadline for DNS propagation + the health probe, bounding the startup stall.
+_PUBLIC_PROBE_TIMEOUT = 45.0
+_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0
+_PUBLIC_PROBE_RETRY_DELAY = 1.0
+
+# Wait for the hostname via DoH first: an early OS lookup negative-caches the
+# NXDOMAIN for up to 30 min.
+_DNS_POLL_DELAY = 2.0
+# Retry transient DoH failures, but give up fast when DoH is blocked outright.
+_DNS_MAX_DOH_ERRORS = 3
+_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A"
+
def _windows_hidden_kwargs() -> dict:
"""Suppress a child console window on Windows; no-op elsewhere."""
@@ -191,6 +208,59 @@ def ensure_cloudflared() -> Optional[str]:
return None
+def _wait_for_dns(host: str, deadline: float) -> None:
+ import json
+ import urllib.request
+
+ errors = 0
+ while True:
+ answered = False
+ try:
+ req = urllib.request.Request(
+ _DOH_URL.format(host = host),
+ headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"},
+ )
+ with urllib.request.urlopen(req, timeout = 5) as response:
+ answered = bool(json.loads(response.read(65536)).get("Answer"))
+ errors = 0
+ except Exception:
+ errors += 1
+ if errors >= _DNS_MAX_DOH_ERRORS:
+ return
+ if answered:
+ return
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return
+ time.sleep(min(_DNS_POLL_DELAY, remaining))
+
+
+def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool:
+ import json
+ import urllib.request
+ from urllib.parse import urlsplit
+
+ deadline = time.monotonic() + timeout
+ host = urlsplit(url).hostname
+ if host:
+ _wait_for_dns(host, deadline)
+
+ probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}"
+ while True:
+ try:
+ req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"})
+ with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response:
+ body = response.read(4096)
+ if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER:
+ return True
+ except Exception:
+ pass
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return False
+ time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining))
+
+
class CloudflareTunnel:
"""A cloudflared quick tunnel to http://localhost:. Best-effort throughout.
@@ -322,11 +392,12 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
"""Start a quick tunnel and return its public URL once it is actually
serving, or None (best-effort).
- Waits for cloudflared to both mint the URL and register an edge connection
- before returning, so the caller never advertises a URL that yields Cloudflare
- error 1033 (HTTP 530). If a URL is minted but no connection registers within
- the window (e.g. quic is blocked on this network), retries once forcing the
- http2 protocol. On any failure the tunnel is stopped and None is returned.
+ Waits for cloudflared to both mint the URL and register an edge connection,
+ then fetches /api/health over the public URL, so the caller never advertises
+ a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host.
+ If a URL is minted but no connection registers within the window (e.g. quic
+ is blocked on this network), retries once forcing the http2 protocol. On any
+ failure the tunnel is stopped and None is returned.
"""
global _active_tunnel, _shutdown_requested
binary = ensure_cloudflared()
@@ -349,9 +420,13 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
+ registered = False
try:
tunnel.start()
url = tunnel.wait_for_ready(timeout)
+ registered = url is not None
+ if url and not verify_public_url(url):
+ url = None
except Exception:
url = None
if url:
@@ -371,6 +446,9 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
# http2 will not help, so do not burn another window on it.
if not saw_url:
return None
+ # probe failure after registering is DNS propagation; http2 would not help
+ if registered:
+ return None
return None
diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py
index bb51cabf76..2094d15066 100644
--- a/studio/backend/tests/test_cloudflare_tunnel.py
+++ b/studio/backend/tests/test_cloudflare_tunnel.py
@@ -403,11 +403,234 @@ def test_reader_ignores_api_endpoint_failure_line():
assert t.error == "cloudflared exited before emitting a tunnel URL"
+# ── public reachability probe ────────────────────────────────────────
+
+
+class _FakeResponse:
+ def __init__(self, body):
+ self._body = body
+
+ def read(self, size = -1):
+ return self._body
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ return False
+
+
+def _patch_urlopen(monkeypatch, handler):
+ import urllib.request
+ monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: handler(req))
+
+
+@pytest.fixture(autouse = True)
+def _stub_dns_wait(monkeypatch, request):
+ if request.node.name.startswith("test_verify_public_url"):
+ monkeypatch.setattr(ct, "_wait_for_dns", lambda *a, **kw: None)
+
+
+def test_wait_for_dns_polls_until_answer(monkeypatch):
+ calls = []
+
+ def handler(req):
+ calls.append(req.full_url)
+ if len(calls) < 3:
+ return _FakeResponse(b'{"Status":3}')
+ return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}')
+
+ _patch_urlopen(monkeypatch, handler)
+ monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
+ ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5)
+ assert len(calls) == 3
+ assert "name=words.trycloudflare.com" in calls[0]
+
+
+def test_wait_for_dns_gives_up_at_deadline(monkeypatch):
+ _patch_urlopen(monkeypatch, lambda req: _FakeResponse(b'{"Status":3}'))
+ monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
+ ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 0.05)
+
+
+def test_wait_for_dns_retries_transient_doh_error(monkeypatch):
+ calls = []
+
+ def handler(req):
+ calls.append(req.full_url)
+ if len(calls) < 3:
+ raise OSError("transient")
+ return _FakeResponse(b'{"Status":0,"Answer":[{"data":"104.16.0.1"}]}')
+
+ _patch_urlopen(monkeypatch, handler)
+ monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
+ ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5)
+ assert len(calls) == 3
+
+
+def test_wait_for_dns_bails_on_persistent_doh_errors(monkeypatch):
+ calls = []
+
+ def handler(req):
+ calls.append(req.full_url)
+ raise OSError("blocked")
+
+ _patch_urlopen(monkeypatch, handler)
+ monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
+ ct._wait_for_dns("words.trycloudflare.com", ct.time.monotonic() + 5)
+ assert len(calls) == ct._DNS_MAX_DOH_ERRORS
+
+
+def test_verify_public_url_accepts_studio_marker(monkeypatch):
+ seen = {}
+
+ def handler(req):
+ seen["url"] = req.full_url
+ return _FakeResponse(b'{"status":"healthy","service":"Unsloth UI Backend"}')
+
+ _patch_urlopen(monkeypatch, handler)
+ assert ct.verify_public_url("https://words.trycloudflare.com") is True
+ assert seen["url"] == "https://words.trycloudflare.com/api/health"
+
+
+def test_verify_public_url_waits_for_dns_first(monkeypatch):
+ order = []
+ monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: order.append(("dns", host)))
+
+ def handler(req):
+ order.append(("probe", req.full_url))
+ return _FakeResponse(b'{"service":"Unsloth UI Backend"}')
+
+ _patch_urlopen(monkeypatch, handler)
+ assert ct.verify_public_url("https://words.trycloudflare.com") is True
+ assert order[0] == ("dns", "words.trycloudflare.com")
+ assert order[1][0] == "probe"
+
+
+def test_verify_public_url_dns_wait_and_probe_share_deadline(monkeypatch):
+ # An exhausted DNS wait leaves the probe a single attempt, not a fresh window.
+ calls = []
+ monkeypatch.setattr(ct, "_wait_for_dns", lambda host, deadline: None)
+
+ def handler(req):
+ calls.append(req.full_url)
+ raise OSError("unreachable")
+
+ _patch_urlopen(monkeypatch, handler)
+ assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0) is False
+ assert len(calls) == 1
+
+
+def test_verify_public_url_retries_then_succeeds(monkeypatch):
+ calls = []
+
+ def handler(req):
+ calls.append(req.full_url)
+ if len(calls) < 3:
+ raise OSError("Name or service not known")
+ return _FakeResponse(b'{"service":"Unsloth UI Backend"}')
+
+ _patch_urlopen(monkeypatch, handler)
+ monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
+ assert ct.verify_public_url("https://words.trycloudflare.com") is True
+ assert len(calls) == 3
+
+
+def test_verify_public_url_rejects_unreachable_host(monkeypatch):
+ def handler(req):
+ raise OSError("Name or service not known")
+
+ _patch_urlopen(monkeypatch, handler)
+ monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
+ assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False
+
+
+def test_verify_public_url_rejects_foreign_responder(monkeypatch):
+ # e.g. a Cloudflare error page: no service marker in the body.
+ _patch_urlopen(monkeypatch, lambda req: _FakeResponse(b"error 1033"))
+ monkeypatch.setattr(ct.time, "sleep", lambda _s: None)
+ assert ct.verify_public_url("https://words.trycloudflare.com", timeout = 0.05) is False
+
+
+@pytest.fixture(autouse = True)
+def _stub_public_probe(monkeypatch, request):
+ # start_studio_tunnel tests use fake hostnames; keep them off the network.
+ if not request.node.name.startswith("test_start_studio_tunnel"):
+ return
+ monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: True)
+
+
def test_start_studio_tunnel_no_binary(monkeypatch):
monkeypatch.setattr(ct, "ensure_cloudflared", lambda: None)
assert ct.start_studio_tunnel(8080) is None
+def test_start_studio_tunnel_drops_url_that_is_not_publicly_reachable(monkeypatch):
+ attempts = []
+
+ class _Stub:
+ def __init__(
+ self,
+ port,
+ binary,
+ protocol = None,
+ ):
+ self.url = None
+ attempts.append(protocol)
+
+ def start(self):
+ self.url = "https://words.trycloudflare.com"
+
+ def wait_for_ready(self, timeout):
+ return self.url
+
+ def stop(self):
+ pass
+
+ monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
+ monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
+ monkeypatch.setattr(ct, "verify_public_url", lambda url, **kw: False)
+ assert ct.start_studio_tunnel(8080) is None
+ assert attempts == [None]
+ assert ct._active_tunnel is None
+
+
+def test_start_studio_tunnel_returns_url_once_probe_passes(monkeypatch):
+ probed = []
+
+ class _Stub:
+ def __init__(
+ self,
+ port,
+ binary,
+ protocol = None,
+ ):
+ self.url = None
+ self.protocol = protocol
+
+ def start(self):
+ self.url = "https://words.trycloudflare.com"
+
+ def wait_for_ready(self, timeout):
+ return self.url
+
+ def stop(self):
+ pass
+
+ def _probe(url, **kw):
+ probed.append(url)
+ return True
+
+ monkeypatch.setattr(ct, "ensure_cloudflared", lambda: "/bin/cloudflared")
+ monkeypatch.setattr(ct, "CloudflareTunnel", _Stub)
+ monkeypatch.setattr(ct, "verify_public_url", _probe)
+ try:
+ assert ct.start_studio_tunnel(8080) == "https://words.trycloudflare.com"
+ assert probed == ["https://words.trycloudflare.com"]
+ finally:
+ ct.stop_studio_tunnel()
+
+
def test_start_studio_tunnel_registers_before_wait(monkeypatch):
# The tunnel must be visible to stop_studio_tunnel() during the readiness
# wait, else a shutdown in that window orphans cloudflared.
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index adab56582b..ee81ef6794 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -1570,6 +1570,18 @@ const Composer: FC<{
const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300);
return () => clearTimeout(t);
}, [composerText, draftKey]);
+ // Without this the restore effect above puts the sent text back when the
+ // runtime rebinds on the first message.
+ const draftKeyRef = useRef(draftKey);
+ useEffect(() => {
+ draftKeyRef.current = draftKey;
+ }, [draftKey]);
+ const clearStoredDraft = useCallback(() => {
+ const key = draftKeyRef.current;
+ if (key) {
+ writeComposerDraft(key, "");
+ }
+ }, []);
// react-textarea-autosize re-measures only on value change or window resize,
// not on the width swap from expanding, so it keeps the taller height and
// leaves a stray blank row. Nudge a resize whenever input width changes.
@@ -1720,9 +1732,10 @@ const Composer: FC<{
setPendingSend(false);
dismissWaitToast();
if (text.trim().length > 0 || attachments.length > 0) {
+ clearStoredDraft();
aui.composer().send();
}
- }, [pendingSend, indexingActive, aui, dismissWaitToast]);
+ }, [pendingSend, indexingActive, aui, clearStoredDraft, dismissWaitToast]);
// Drop any queued send + toast on unmount (e.g. thread switch).
useEffect(
@@ -1765,6 +1778,7 @@ const Composer: FC<{
flushResourcesSync(() => {
aui.composer().setText("");
});
+ clearStoredDraft();
startPromptQueue(
[queuedPrompt],
createPromptQueueTarget(),
@@ -1798,6 +1812,7 @@ const Composer: FC<{
closeOverlay();
return;
}
+ clearStoredDraft();
setImageToolsEnabled(true);
setPendingImageEditReference({
threadId: overlay.threadId ?? referenceThreadId,
@@ -1815,11 +1830,15 @@ const Composer: FC<{
);
});
closeOverlay();
+ return;
}
+
+ clearStoredDraft();
},
[
aui,
canQueueCurrentPrompt,
+ clearStoredDraft,
closeOverlay,
composerText,
createPromptQueueTarget,
@@ -1921,6 +1940,7 @@ const Composer: FC<{
flushResourcesSync(() => {
aui.composer().setText("");
});
+ clearStoredDraft();
startPromptQueue([queuedPrompt], createPromptQueueTarget(), true);
}}
onSendClick={interceptSend}
From 430ada617af52c847656eb854c272fcda3d9193a Mon Sep 17 00:00:00 2001
From: Leo Borcherding
Date: Thu, 23 Jul 2026 02:42:03 -0500
Subject: [PATCH 078/255] installer: fix false "no GPU detected" on AMD hosts
(dead KFD check) + clearer ROCm-less warning (#7314)
* installer: fix Linux AMD GPU detection + actionable ROCm-less warning
The rocminfo/amd-smi-less fallback in _has_amd_rocm_gpu keyed on a
/gpu_id/ line inside each KFD node's properties file, but gpu_id is a
separate sibling sysfs file and never appears in properties. The guard
never matched, so the fallback missed every AMD host without ROCm
tooling (e.g. a fresh CachyOS/Arch box) and reported 'no GPU detected'
despite vendor_id 4098 being present in the KFD topology.
Detect via vendor_id == 4098 directly: the KFD CPU node reports
vendor_id 0, so any 4098 node is an AMD GPU, while NVIDIA's KFD nodes
report 4318 and stay excluded.
Also rework the 'ROCm version could not be determined' warning into an
actionable message (install the ROCm/HIP SDK; Arch/CachyOS:
rocm-hip-sdk) so ROCm-less users know the concrete next step instead of
silently landing on CPU-only PyTorch.
* tests: replace the FNR==1 KFD invariant with the per-line vendor_id check
The FNR==1 reset guarded the old paired gpu_id+vendor_id awk against
cross-node state leakage. The new detection is a single atomic
vendor_id==4098 line condition, so there is no per-node state to reset;
assert the new invariant instead (single-line vendor match, and no
/gpu_id/ pattern, which never matched inside properties).
tests/studio/install/test_rocm_support.py: 344 passed, 2 skipped.
* installer: mirror the KFD vendor_id fix in setup.sh + honest CPU-fallback summary
Codex P2 follow-ups:
- studio/setup.sh carried the same dead gpu_id-inside-properties awk, so a
host install.sh now routes to ROCm still failed setup's independent AMD
re-probe and got a CPU llama.cpp. Use the same per-line vendor_id 4098
check.
- When the AMD GPU is detected but the torch index stays CPU, the summary
printed the old false diagnosis (gpu none / "No GPU detected"). Gate both
on _has_amd_rocm_gpu and say what actually happened: AMD GPU present, no
usable ROCm, CPU fallback.
- Structure test asserting setup.sh's KFD awk stays in sync with install.sh.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Keep KFD-only AMD hosts on the CPU fallback (Codex P2s)
The KFD-topology fix makes _has_amd_rocm_gpu / _setup_amd_detected true on hosts that expose an AMD GPU to the kernel but ship no rocminfo/amd-smi. Detection alone does not mean ROCm is usable or that the gfx arch is known, and two downstream paths wrongly assumed it did:
- studio/setup.sh forwarded --has-rocm with no gfx, so install_llama_prebuilt found no per-gfx bundle and dropped to a HIP source build (slow, or a hard failure without build deps) instead of the CPU prebuilt these hosts used to get. Now --has-rocm is forwarded for a gfx-unknown host only when hipcc is present; otherwise it keeps the CPU prebuilt.
- install.sh get_torch_index_url selected a generic rocmX.Y index whenever the ROCm version was readable, but the Strix reroute only learns gfx from rocminfo/amd-smi, so a Strix KFD-only host landed on the broken _grouped_mm wheels. Now, when neither rocminfo nor amd-smi is present (gfx unknowable), it stays on CPU with a hint to install them.
Detection and the improved diagnostics are unchanged; only the routing for gfx-unknown KFD-only hosts is made safe. Adds tests for both gates.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden KFD-only fallback: probe gfx, accept versioned hipcc (Codex P2s)
Follow-up to the previous commit's two guards:
- install.sh: the KFD-only torch guard tested only 'command -v rocminfo/amd-smi', so a host where those binaries exist but do not enumerate the GPU (gfx unreadable) slipped through and, with hipconfig/rocm-core present, still got a generic rocm index -- breaking Strix. Now it actually reads the gfx (rocminfo, then amd-smi list / static --asic, the same probe the reroute uses) and falls back to CPU whenever the arch is unreadable, not just when the binaries are absent.
- studio/setup.sh: the hipcc gate missed a HIP toolchain installed only under a versioned prefix (/opt/rocm-*/bin/hipcc), which the source build at setup.sh:1663 does support, so such hosts were dropped to the CPU prebuilt unnecessarily. The gate now also accepts /opt/rocm-*/bin/hipcc.
Tests updated to assert the gfx-read (not binary-presence) gate and the versioned hipcc path; full test_rocm_support.py green (347 passed). Verified the gfx probe by execution: rocminfo-with-no-gfx now routes to CPU, amd-smi fallback still resolves gfx.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Honor UNSLOTH_ROCM_GFX_ARCH before the CPU fallback for PR #7314
Seed both the gfx-unknown guard in get_torch_index_url and the Strix reroute
from UNSLOTH_ROCM_GFX_ARCH before probing rocminfo/amd-smi, so a host that
names its arch reaches the correct rocm index instead of being forced to CPU
(or to the generic wheels) when the runtime probes can't enumerate the GPU.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Probe gfx with visibility masks cleared for PR #7314 (Codex P2)
rocminfo/amd-smi honor ROCR/HIP_VISIBLE_DEVICES, so a container that masks the
GPU (e.g. ROCR_VISIBLE_DEVICES=-1) would make the gfx probe read nothing and
force CPU torch, even though the KFD-based AMD detection is env-independent and
hipconfig can still supply the ROCm version. Clear the visibility masks for the
rocminfo/amd-smi arch probe only (the Strix reroute keeps them for per-GPU index
selection), so a masked/container host keeps its ROCm route.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Re-probe gfx unmasked in the Strix reroute when a mask hides all agents for PR #7314 (Codex P2)
* Remove leftover conflict marker from the test merge
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Report an explicit CPU pin instead of a ROCm misdiagnosis for PR #7314 (Codex P3)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trigger the reroute re-probe on a set-but-empty visibility mask for PR #7314 (subagent review)
* Guard the ROCm version chain against set -e when no source exists for PR #7314 (simulation find)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Preserve the inferred-gfx reroute for KFD-only hosts (Codex P2)
The gfx-unknown CPU guard in get_torch_index_url fired before the
runtime-less reroute could run: with the KFD topology fix,
_has_amd_rocm_gpu is true on KFD-only hosts, so the reroute's
'! _has_amd_rocm_gpu' gate never let _infer_linux_amd_gfx_arch route
them to AMD per-arch wheels, regressing inferable boxes (PCI/cpuinfo/
lspci) from arch-specific PyTorch to CPU-only.
- Factor the override->rocminfo->amd-smi gfx probe (masks cleared)
into _probe_amd_gfx_arch, shared by the guard and the reroute gate
so the two can't disagree on what 'readable' means.
- Reroute gate now also fires when the GPU is detected but the probe
is empty (KFD-only). Deliberate CPU fallbacks (old/unreadable ROCm
version) all had a readable gfx and stay excluded.
- The guard defers to the reroute (no false 'installing CPU-only
PyTorch' promise) only when inference yields a supported family;
otherwise the actionable CPU warning is unchanged.
Executed tests: KFD-only host reroutes to repo.amd.com per-arch wheels
and exports UNSLOTH_ROCM_GFX_ARCH for setup.sh; readable-gfx CPU
fallback stays un-rerouted; undetected-GPU reroute unchanged; the
guard's three inference outcomes covered. Suite: 375 passed, bash -n
clean on both scripts.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix two false diagnostics on the KFD-only paths (Codex P3s)
1. get_torch_index_url: with UNSLOTH_ROCM_GFX_ARCH set on a KFD-only
host that has no ROCm version sources, the no-version endpoint
printed 'falling back to CPU-only PyTorch' even though the reroute
(gated on the override) then installs the per-arch wheels. When the
override maps to a wheel family, defer with an accurate message;
an unmappable override keeps the CPU warning since the reroute
can't route it either.
2. Runtime-less reroute: the KFD-only branch reached the warning
'ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi)' although
/dev/kfd is exactly what detected the GPU. The diagnostic now
distinguishes KFD-visible/tooling-blind hosts from truly
runtime-invisible ones.
Executed tests: supported override defers without the false CPU
warning, unsupported override and readable-gfx no-version hosts keep
it; KFD-only reroute emits the KFD wording, undetected-GPU reroute
keeps the original. Version sources are shimmed so the tests hold on
dev boxes with a real hipconfig. Suite: 376 passed, bash -n clean.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
install.sh | 168 ++++++-
studio/setup.sh | 14 +-
tests/studio/install/test_rocm_support.py | 557 +++++++++++++++++++++-
3 files changed, 704 insertions(+), 35 deletions(-)
diff --git a/install.sh b/install.sh
index 963107524b..d06fff07c9 100755
--- a/install.sh
+++ b/install.sh
@@ -2115,13 +2115,16 @@ _has_amd_rocm_gpu() {
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
return 0
elif [ -e /dev/kfd ] && \
- awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \
- gpu && amd { found=1 } END{ exit !found }' \
+ awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \
/sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
- # vendor_id 4098 = 0x1002 (AMD). NVIDIA open kernel module (driver
- # 560+) can register KFD topology nodes with non-zero gpu_id but
- # vendor_id 4318 (0x10DE). Require AMD vendor to avoid misrouting
- # NVIDIA-only hosts to the ROCm install path.
+ # vendor_id 4098 = 0x1002 (AMD) marks a GPU node: the KFD CPU node
+ # reports vendor_id 0, so any 4098 node is an AMD GPU. NVIDIA's open
+ # kernel module (driver 560+) registers KFD nodes as vendor_id 4318
+ # (0x10DE), so this never false-positives on NVIDIA-only hosts.
+ # The prior check also required a gpu_id line, but gpu_id is a SIBLING
+ # sysfs file, not a line in properties -- it never matched, so the
+ # fallback silently missed every ROCm-less AMD host (issue: fresh
+ # Arch/CachyOS boxes reporting "no GPU detected").
return 0
fi
return 1
@@ -2230,6 +2233,30 @@ EOF
return 1
}
+# Reads the AMD gfx arch for wheel-index decisions: a user-set
+# UNSLOTH_ROCM_GFX_ARCH is authoritative (lowercased), else rocminfo, then
+# amd-smi. rocminfo/amd-smi honor ROCR/HIP_VISIBLE_DEVICES, so a container mask
+# (e.g. ROCR_VISIBLE_DEVICES=-1) would hide a GPU that the env-independent KFD
+# detection still sees -- the tool probes run with the masks cleared. Prints the
+# gfx token(s) or nothing when unreadable, and always returns 0 (a failing probe
+# as the last command would trip set -e in callers' assignments). Shared by
+# get_torch_index_url's gfx gate and the runtime-less reroute gate so the two
+# can never disagree on what "readable" means.
+_probe_amd_gfx_arch() {
+ _ensure_rocm_probe_env
+ _pg=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]')
+ if [ -z "$_pg" ] && command -v rocminfo >/dev/null 2>&1; then
+ _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ if [ -z "$_pg" ] && command -v amd-smi >/dev/null 2>&1; then
+ _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ if [ -z "$_pg" ]; then
+ _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ fi
+ printf '%s\n' "$_pg"
+}
+
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@@ -2283,6 +2310,29 @@ get_torch_index_url() {
if ! _has_amd_rocm_gpu; then
echo "$_base/cpu"; return
fi
+ # A generic rocm index is only safe when the gfx arch is readable: the
+ # Strix reroute (gfx1150/1151 -> arch-specific index) learns gfx from
+ # rocminfo/amd-smi, so if those are missing OR do not enumerate the GPU, an
+ # unknown-arch box might be Strix and would get the broken _grouped_mm
+ # wheels. Probe via the shared helper (override first, then rocminfo/amd-smi
+ # with visibility masks cleared); if the arch is unreadable, never guess a
+ # rocm index. A KFD-only host whose arch is still inferable from hardware
+ # IDs (PCI/cpuinfo/lspci) returns the cpu index and lets the runtime-less
+ # reroute below upgrade it to AMD per-arch wheels -- the reroute gate uses
+ # this same probe, so the handoff can't misfire. Only when inference fails
+ # too is CPU final, with the actionable warning.
+ _amd_gfx_probe=$(_probe_amd_gfx_arch)
+ if [ -z "$_amd_gfx_probe" ]; then
+ if _amd_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null) && \
+ [ -n "$_amd_inferred_gfx" ] && \
+ _amd_arch_index_family_for_gfx "$_amd_inferred_gfx" >/dev/null 2>&1; then
+ echo "[WARN] AMD GPU detected but rocminfo/amd-smi can't read its gfx arch -- inferring $_amd_inferred_gfx from hardware IDs." >&2
+ echo "$_base/cpu"; return
+ fi
+ echo "[WARN] AMD GPU detected but its gfx arch can't be read (rocminfo/amd-smi missing or not enumerating the GPU) -- installing CPU-only PyTorch." >&2
+ echo "[WARN] For GPU PyTorch, install or repair rocminfo/amd-smi (e.g. sudo pacman -S rocm-hip-sdk) and re-run this installer." >&2
+ echo "$_base/cpu"; return
+ fi
# AMD GPU confirmed -- detect ROCm version
_rocm_tag=""
_rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \
@@ -2299,7 +2349,11 @@ get_torch_index_url() {
{ command -v rpm >/dev/null 2>&1 && \
ver="$(rpm -q --qf '%{VERSION}\n' rocm-core 2>/dev/null)" && \
[ -n "$ver" ] && \
- printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null
+ printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null || _rocm_tag=""
+ # ^ || guard: when EVERY version source is missing (e.g. rocminfo present
+ # but rocm-core not installed, so dpkg-query/rpm exit 1), the whole ||
+ # chain fails and set -e would kill the installer BEFORE the actionable
+ # no-version WARN below -- exactly the fresh-install case it exists for.
# Validate _rocm_tag: must match "rocmX.Y" with major >= 1
case "$_rocm_tag" in
rocm[1-9]*.[0-9]*) : ;; # valid (major >= 1)
@@ -2335,12 +2389,27 @@ get_torch_index_url() {
esac
return
fi
- # AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be
- # read from any source (amd-smi, /opt/rocm/.info/version, hipconfig,
- # dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch.
- echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2
- echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2
- echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
+ # AMD GPU confirmed (rocminfo/amd-smi or the KFD topology fallback) but
+ # no ROCm/HIP install was found to read the version from (amd-smi,
+ # /opt/rocm/.info/version, hipconfig, dpkg, rpm). This is the common
+ # fresh-install case: the GPU is real, but with no ROCm userspace the
+ # correct PyTorch build can't be selected. Warn with an actionable fix
+ # rather than silently installing CPU PyTorch.
+ # A user-set UNSLOTH_ROCM_GFX_ARCH seeded the probe above, so rocminfo/
+ # amd-smi may still be unable to see the GPU; when the named arch maps to
+ # a wheel family, the runtime-less reroute (gated on the override) will
+ # install the AMD per-arch wheels -- a CPU-only warning here would be
+ # false for that path. Defer like the inferable-arch branch does.
+ if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] && \
+ _amd_arch_index_family_for_gfx "$_amd_gfx_probe" >/dev/null 2>&1; then
+ echo "[WARN] AMD GPU detected with no readable ROCm version, but UNSLOTH_ROCM_GFX_ARCH=$_amd_gfx_probe is set -- routing to AMD per-arch wheels." >&2
+ echo "$_base/cpu"; return
+ fi
+ echo "[WARN] AMD GPU detected, but no ROCm/HIP install was found to select the matching GPU PyTorch build -- falling back to CPU-only PyTorch." >&2
+ echo "[WARN] Install the ROCm/HIP SDK, then re-run this installer:" >&2
+ echo "[WARN] Arch / CachyOS : sudo pacman -S rocm-hip-sdk" >&2
+ echo "[WARN] other distros : https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
+ echo "[WARN] Minimum required for version detection: amd-smi, hipconfig, /opt/rocm/.info/version, or the rocm-core package." >&2
echo "$_base/cpu"; return
fi
# Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P).
@@ -2841,14 +2910,20 @@ TORCH_INDEX_URL=$(get_torch_index_url)
# Linux: ROCm runtime missing but a supported AMD gfx arch is inferable (Strix Halo
# in /proc/cpuinfo, lspci marketing name, UNSLOTH_ROCM_GFX_ARCH). Route to AMD's
# per-arch wheels like install.ps1 does on Windows (unslothai#7301).
-# Gated on _has_amd_rocm_gpu being FALSE: a */cpu index on a host whose GPU IS
-# visible to the ROCm probes is a deliberate fallback (unsupported/unreadable
-# ROCm version, after its own warning), not a missing runtime -- rerouting it
-# would contradict that decision. An explicit UNSLOTH_ROCM_GFX_ARCH override
-# stays authoritative either way.
+# Gated on the runtime probes NOT naming a gfx: either no AMD GPU is detected at
+# all (_has_amd_rocm_gpu false), or the GPU is visible only through the
+# env-independent KFD topology while rocminfo/amd-smi can't read its arch
+# (KFD-only host, unslothai#7314 -- before the KFD detection fix these hosts
+# reached this reroute via the false branch, so the empty-probe condition
+# preserves that routing). A */cpu index chosen WITH a readable gfx
+# (unsupported/unreadable ROCm version, after its own warning) is a deliberate
+# fallback -- rerouting it would contradict that decision, and stays excluded
+# because the shared probe returns its gfx. An explicit UNSLOTH_ROCM_GFX_ARCH
+# override stays authoritative either way.
if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \
! _has_usable_nvidia_gpu && \
- { [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu; } && \
+ { [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu || \
+ [ -z "$(_probe_amd_gfx_arch)" ]; } && \
case "$(uname -s)" in Linux) true ;; *) false ;; esac && \
case "$_ARCH" in x86_64|amd64) true ;; *) false ;; esac; then
# ROCm torch wheels are x86_64-only; get_torch_index_url returns CPU on other
@@ -2880,7 +2955,13 @@ if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \
;;
esac
echo "" >&2
- echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2
+ # KFD-only hosts reach this reroute with /dev/kfd present
+ # (that's what detected them), so don't claim it's missing.
+ if _has_amd_rocm_gpu; then
+ echo " [WARN] AMD GPU visible via the kernel driver (KFD) but rocminfo/amd-smi can't read its gfx arch; using $_linux_inferred_gfx." >&2
+ else
+ echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2
+ fi
echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2
echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2
echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2
@@ -3004,8 +3085,10 @@ case "$_torch_index_leaf" in
# || true on each probe: no gfx match makes grep exit 1, which under
# set -euo pipefail would abort the installer before the next fallback
# runs (now that the case matches every rocm* index, not just rocm7.1).
- _gfx_all=""
- if command -v rocminfo >/dev/null 2>&1; then
+ # A user-supplied UNSLOTH_ROCM_GFX_ARCH overrides probing (mirrors setup.sh
+ # and the display block), so a Strix override still reaches the arch index.
+ _gfx_all=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]')
+ if [ -z "$_gfx_all" ] && command -v rocminfo >/dev/null 2>&1; then
_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
@@ -3016,6 +3099,23 @@ case "$_torch_index_leaf" in
_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
fi
+ # get_torch_index_url reads the arch with ROCR/HIP masks cleared, so a
+ # mask hiding every agent (e.g. ROCR_VISIBLE_DEVICES=-1) still lands
+ # here on a generic rocm index; re-probe unmasked or a masked-out Strix
+ # box keeps the broken generic wheels. Partial masks never get here
+ # (they enumerate at least one agent above) and keep their selection.
+ # ${VAR+x} (not :-): a SET-but-empty mask also hides every agent and
+ # must trigger the re-probe too.
+ if [ -z "$_gfx_all" ] && [ -n "${ROCR_VISIBLE_DEVICES+x}${HIP_VISIBLE_DEVICES+x}" ]; then
+ if command -v rocminfo >/dev/null 2>&1; then
+ _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
+ _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ [ -z "$_gfx_all" ] && \
+ _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
+ fi
+ fi
_runtime_gfx=""
if [ -n "$_gfx_all" ]; then
_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
@@ -3169,6 +3269,17 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
# Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only.
step "gpu" "Apple Silicon (Metal, unified memory)"
+elif _has_amd_rocm_gpu; then
+ if [ "$_torch_index_pinned" = true ]; then
+ # An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY pin skipped all probing;
+ # do not claim ROCm is unusable when a CPU/other index was requested.
+ step "gpu" "AMD GPU (torch index pinned: $_torch_index_leaf)" "$C_WARN"
+ else
+ # AMD GPU visible to the kernel but the torch index stayed CPU: no usable
+ # ROCm userspace to pick a wheel. "none" would repeat the false diagnosis
+ # this installer used to give.
+ step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)" "$C_WARN"
+ fi
else
step "gpu" "none (CPU-only)" "$C_WARN"
fi
@@ -3177,8 +3288,17 @@ fi
case "$TORCH_INDEX_URL" in
*/cpu)
if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then
- substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
- if [ "$OS" = "wsl" ]; then
+ if [ "$_torch_index_pinned" = true ]; then
+ # An explicit CPU pin is a request, not a detection failure:
+ # skip the SDK guidance (ROCm may be perfectly healthy here).
+ substep "CPU-only PyTorch (index pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY)."
+ elif _has_amd_rocm_gpu; then
+ substep "AMD GPU detected, but no usable ROCm/HIP install -- installing CPU-only PyTorch." "$C_WARN"
+ substep "Install the ROCm/HIP SDK and re-run this installer for GPU PyTorch." "$C_WARN"
+ else
+ substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
+ fi
+ if [ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]; then
# WSL + no GPU detected (detection above found nothing). Common
# cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet --
# /dev/dxg present (graphics) but no ROCm runtime.
diff --git a/studio/setup.sh b/studio/setup.sh
index 2a2b41d0f6..0183ef3776 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -1101,8 +1101,7 @@ if [ "$_setup_nvidia_usable" != true ]; then
_setup_mkt=$(_setup_run_smi amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
'/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
elif [ -e /dev/kfd ] && \
- awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \
- gpu && amd { found=1 } END{ exit !found }' \
+ awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \
/sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
# KFD sysfs fallback, AMD vendor_id 4098 only (mirrors install.sh
# _has_amd_rocm_gpu): covers AMD hosts where rocminfo/amd-smi are
@@ -1358,9 +1357,14 @@ else
# name-inferred arch). Implies --has-rocm on the installer side.
if [ -n "${_setup_gfx:-}" ]; then
_PREBUILT_CMD+=(--rocm-gfx "$_setup_gfx")
- elif [ "$_setup_amd_detected" = true ]; then
- # AMD was detected but gfx resolution failed; tell the installer ROCm is
- # present so it can still attempt a prebuilt. Mirrors setup.ps1 behaviour.
+ elif [ "$_setup_amd_detected" = true ] && \
+ { command -v hipcc >/dev/null 2>&1 || [ -x /opt/rocm/bin/hipcc ] || \
+ ls /opt/rocm-*/bin/hipcc >/dev/null 2>&1; }; then
+ # AMD detected but gfx unknown (KFD-only host): forward --has-rocm only when
+ # hipcc can actually build llama.cpp (incl. a versioned /opt/rocm-*/bin, the
+ # same paths the source build uses). With no gfx the prebuilt resolver finds
+ # no ROCm bundle and the source build would fail, so without hipcc fall
+ # through to the CPU prebuilt instead of breaking the install.
_PREBUILT_CMD+=(--has-rocm)
fi
# UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, trimmed) forces the CPU-only
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index cd7b68f4b6..fa76011041 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -1459,6 +1459,59 @@ class TestInstallShStructure:
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
assert "amd-smi" in source
+
+ def test_cpu_index_note_respects_explicit_pin(self):
+ """An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY CPU pin is a request, not
+ a detection failure: the */cpu wheel note must report the pin instead of
+ claiming ROCm/HIP is unusable, the WSL setup guidance must be skipped,
+ and the gpu summary must not label a pinned AMD host "no usable ROCm"."""
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ note = source.find('substep "AMD GPU detected, but no usable ROCm/HIP install')
+ assert note != -1
+ assert (
+ '[ "$_torch_index_pinned" = true ]' in source[note - 400 : note]
+ ), "the */cpu note must check the explicit pin before diagnosing ROCm"
+ assert (
+ '[ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]' in source
+ ), "ROCm-on-WSL guidance is detection advice; skip it for pinned installs"
+ summary = source.find('step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)"')
+ assert summary != -1
+ assert (
+ '[ "$_torch_index_pinned" = true ]' in source[summary - 700 : summary]
+ ), "the gpu summary must not claim no usable ROCm for a pinned index"
+
+ def test_rocm_version_chain_survives_no_source_under_set_e(self):
+ """When every ROCm version source is missing (e.g. rocminfo present but
+ rocm-core not installed, so dpkg-query/rpm exit 1), the _rocm_tag ||
+ chain fails as a whole; without the || guard set -e kills the installer
+ BEFORE the actionable no-version WARN it feeds. Executed, not text."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute the version chain")
+ sh_path = PACKAGE_ROOT / "install.sh"
+ source = sh_path.read_text(encoding = "utf-8")
+ chain = re.search(
+ r'^ _rocm_tag=\$\(\{ command -v amd-smi.*?\|\| _rocm_tag=""\n',
+ source,
+ re.S | re.M,
+ )
+ assert chain, "could not extract the guarded _rocm_tag chain"
+ with tempfile.TemporaryDirectory() as d:
+ # Tools exist on PATH but yield nothing usable, like a box with the
+ # probe tools installed and no rocm-core package.
+ for name in ("amd-smi", "hipconfig", "dpkg-query", "rpm"):
+ p = os.path.join(d, name)
+ with open(p, "w", encoding = "utf-8") as f:
+ f.write("#!/bin/sh\nexit 1\n")
+ os.chmod(p, 0o755)
+ script = (
+ "set -euo pipefail\n" + chain.group(0) + '\nprintf "SURVIVED:%s\\n" "$_rocm_tag"\n'
+ )
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""))
+ r = subprocess.run([shell, "-c", script], env = env, capture_output = True, text = True)
+ assert r.returncode == 0, f"version chain aborted under set -e: {r.stderr}"
+ assert r.stdout.startswith("SURVIVED:"), r.stdout
assert "rocm" in source.lower()
def test_cuda_precedence(self):
@@ -1590,17 +1643,446 @@ class TestInstallShStructure:
"4098" in func_body
), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
- def test_kfd_awk_resets_state_per_file(self):
- """KFD sysfs awk must reset gpu/amd state per file (FNR==1) to avoid Ryzen+NVIDIA false positives."""
+ def test_kfd_awk_vendor_check_is_per_line(self):
+ """KFD sysfs awk must decide on a single vendor_id line, with no cross-node state.
+
+ The old awk paired two per-node flags (gpu_id + vendor_id) and needed an FNR==1
+ reset so flags from different KFD nodes could not combine into a Ryzen+NVIDIA
+ false positive. gpu_id is a sibling sysfs file and never appears inside
+ properties, so that pairing also never matched at all (every ROCm-less AMD host
+ was reported as no-GPU). The replacement keys on one atomic line: only an AMD
+ GPU node reports `vendor_id 4098` (KFD CPU nodes report 0, NVIDIA's open kernel
+ module registers 4318), so there is no cross-file state left to reset.
+ """
sh_path = PACKAGE_ROOT / "install.sh"
source = sh_path.read_text(encoding = "utf-8")
func_start = source.find("_has_amd_rocm_gpu()")
func_end = source.find("\n}", func_start)
func_body = source[func_start:func_end]
- assert "FNR==1" in func_body, (
- "_has_amd_rocm_gpu KFD awk must reset state per file with FNR==1 "
- "to avoid false positives on Ryzen+NVIDIA hosts with multiple KFD nodes"
+ assert "$2 == 4098" in func_body, (
+ "_has_amd_rocm_gpu KFD awk must match `vendor_id 4098` as a single-line "
+ "condition so no per-node state can leak across KFD nodes"
)
+ assert "/gpu_id/" not in func_body, (
+ "_has_amd_rocm_gpu KFD awk must not key on a gpu_id line: gpu_id is a "
+ "sibling sysfs file, not a line in properties, so it never matches there"
+ )
+
+ def test_setup_sh_kfd_awk_matches_install_sh(self):
+ """setup.sh's KFD fallback must use the same per-line vendor_id check as install.sh.
+
+ setup.sh re-probes AMD detection independently of install.sh; if its copy keeps
+ the dead gpu_id-inside-properties pairing, a host that install.sh routes to ROCm
+ still gets a CPU llama.cpp from the setup step (_setup_amd_detected stays false).
+ """
+ source = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8")
+ assert (
+ "$2 == 4098" in source
+ ), "setup.sh KFD awk must match `vendor_id 4098` as a single-line condition"
+ assert (
+ "/gpu_id/" not in source
+ ), "setup.sh KFD awk must not key on a gpu_id line inside properties"
+
+ def test_kfd_only_torch_falls_back_to_cpu(self):
+ """An AMD host whose gfx arch can't be read (rocminfo/amd-smi missing, or
+ present but not enumerating the GPU) must route torch to CPU, not a generic
+ rocm index: a Strix box (gfx1150/1151) would otherwise get the broken
+ _grouped_mm wheels because the reroute has no gfx to correct it."""
+ source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+ body = _extract_sh_function_body(source, "get_torch_index_url")
+ probe = body.find("_amd_gfx_probe=$(_probe_amd_gfx_arch)")
+ assert probe >= 0, "get_torch_index_url must probe the gfx arch before picking a rocm index"
+ # The shared probe reads gfx (not just tests binary presence), from rocminfo
+ # AND amd-smi, so an installed-but-not-enumerating probe still falls to CPU.
+ helper = _extract_sh_function_body(source, "_probe_amd_gfx_arch")
+ assert helper, "install.sh must define the shared _probe_amd_gfx_arch helper"
+ assert (
+ "rocminfo 2>/dev/null) | grep -oE 'gfx" in helper
+ ), "probe must read gfx from rocminfo"
+ assert (
+ "amd-smi list 2>/dev/null) | grep -oE 'gfx" in helper
+ ), "probe must read gfx from amd-smi"
+ # The probe clears ROCR/HIP_VISIBLE_DEVICES so a container mask
+ # (ROCR_VISIBLE_DEVICES=-1) can't blind the env-independent KFD detection.
+ assert (
+ "unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES" in helper
+ ), "the gfx probe must clear the visibility masks so a mask can't force CPU"
+ cpu_guard = body.find('if [ -z "$_amd_gfx_probe" ]')
+ assert cpu_guard >= 0, "unreadable gfx must fall back to CPU"
+ assert cpu_guard < body.find(
+ "_rocm_tag="
+ ), "the gfx gate must run before the ROCm version/index selection"
+
+ def test_kfd_only_llama_requires_hipcc(self):
+ """setup.sh must forward --has-rocm for a gfx-unknown (KFD-only) host only when
+ hipcc is present. With no gfx the prebuilt resolver finds no ROCm bundle and the
+ source build would fail, so without a HIP toolchain the host keeps the CPU
+ prebuilt rather than breaking the llama.cpp install."""
+ source = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8")
+ idx = source.find("_PREBUILT_CMD+=(--has-rocm)")
+ assert idx >= 0, "setup.sh must still be able to forward --has-rocm"
+ window = source[max(0, idx - 900) : idx]
+ assert (
+ "hipcc" in window
+ ), "the gfx-unknown --has-rocm branch must gate on hipcc (a usable HIP toolchain)"
+ assert (
+ "command -v hipcc" in window or "/opt/rocm/bin/hipcc" in window
+ ), "hipcc presence must be checked via command -v or the rocm bin path"
+ assert (
+ "/opt/rocm-*/bin/hipcc" in window
+ ), "the hipcc gate must also accept a versioned /opt/rocm-*/bin/hipcc toolchain"
+
+ def test_gfx_unknown_guard_honors_override(self):
+ """A user-set UNSLOTH_ROCM_GFX_ARCH must seed the gfx probe before the CPU
+ fallback: an air-gapped/rocminfo-less Strix host that names its arch should
+ still reach a rocm index instead of being forced to CPU."""
+ source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+ helper = _extract_sh_function_body(source, "_probe_amd_gfx_arch")
+ assert helper, "install.sh must define the shared _probe_amd_gfx_arch helper"
+ seed = helper.find("$(printf")
+ assert seed >= 0, "the gfx probe must seed from UNSLOTH_ROCM_GFX_ARCH"
+ assert "UNSLOTH_ROCM_GFX_ARCH" in helper[seed : seed + 80]
+ assert seed < helper.find(
+ "rocminfo 2>/dev/null) | grep -oE 'gfx"
+ ), "the override must be read before probing rocminfo"
+ body = _extract_sh_function_body(source, "get_torch_index_url")
+ call = body.find("_amd_gfx_probe=$(_probe_amd_gfx_arch)")
+ assert call >= 0, "get_torch_index_url must call the shared probe"
+ assert call < body.find(
+ 'if [ -z "$_amd_gfx_probe" ]; then'
+ ), "the probe must run before the CPU fallback guard"
+
+ def test_gfx_override_seeds_reroute_without_tools(self):
+ """The Strix reroute must honour UNSLOTH_ROCM_GFX_ARCH even when rocminfo and
+ amd-smi are absent, so a manual override reaches the arch index; with no
+ override and no tools it must stay empty (no false Strix routing)."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute the probe block")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ block = re.search(
+ r'^ _gfx_all=\$\(printf[^\n]*\n.*?(?=^ _strix_gfx="")',
+ source,
+ re.S | re.M,
+ )
+ assert block, "could not extract the gfx-detection block"
+ with tempfile.TemporaryDirectory() as d:
+ # Shim rocminfo/amd-smi to enumerate nothing, so only the override can
+ # supply a gfx (keeps coreutils on PATH for tr/grep/printf).
+ for name in ("rocminfo", "amd-smi"):
+ p = os.path.join(d, name)
+ with open(p, "w", encoding = "utf-8") as f:
+ f.write("#!/bin/sh\nexit 0\n")
+ os.chmod(p, 0o755)
+ script = (
+ 'set -euo pipefail\nHIP_VISIBLE_DEVICES=""\nROCR_VISIBLE_DEVICES=""\n'
+ + block.group(0)
+ + '\nprintf "OK:%s\\n" "$_gfx_all"\n'
+ )
+
+ def run(**extra):
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra)
+ return subprocess.run(
+ [shell, "-c", script], env = env, capture_output = True, text = True
+ )
+
+ r = run(UNSLOTH_ROCM_GFX_ARCH = "GFX1151")
+ assert r.returncode == 0, f"override probe aborted: {r.stderr}"
+ assert "OK:gfx1151" in r.stdout, f"override not honoured/lowercased: {r.stdout!r}"
+ r2 = run()
+ assert r2.returncode == 0, f"empty probe aborted: {r2.stderr}"
+ assert (
+ "OK:\n" in r2.stdout or r2.stdout.strip() == "OK:"
+ ), f"no override + no tools must leave gfx empty: {r2.stdout!r}"
+
+ def test_gfx_probe_ignores_visibility_mask(self):
+ """A container visibility mask (ROCR_VISIBLE_DEVICES=-1) must not blind the
+ gfx probe: rocminfo honours the mask and would enumerate nothing, but KFD
+ detection is env-independent, so the probe clears the mask and still reads
+ the arch (else a masked host is wrongly forced to CPU)."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute the probe block")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ probe_fn = _extract_sh_function_body(source, "_probe_amd_gfx_arch")
+ assert probe_fn, "could not extract _probe_amd_gfx_arch"
+ with tempfile.TemporaryDirectory() as d:
+ # rocminfo that mimics ROCR_VISIBLE_DEVICES=-1 hiding all agents.
+ with open(os.path.join(d, "rocminfo"), "w", encoding = "utf-8") as f:
+ f.write(
+ "#!/bin/sh\n"
+ 'if [ "${ROCR_VISIBLE_DEVICES:-}" = "-1" ]; then echo "no agents"; exit 0; fi\n'
+ 'echo " Name: gfx1151"\n'
+ )
+ os.chmod(os.path.join(d, "rocminfo"), 0o755)
+ script = (
+ "set -euo pipefail\n"
+ "_ensure_rocm_probe_env() { :; }\n"
+ + probe_fn
+ + '\n_amd_gfx_probe=$(_probe_amd_gfx_arch)\nprintf "OK:%s\\n" "$_amd_gfx_probe"\n'
+ )
+
+ def run(**extra):
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra)
+ return subprocess.run(
+ [shell, "-c", script], env = env, capture_output = True, text = True
+ )
+
+ r = run(ROCR_VISIBLE_DEVICES = "-1")
+ assert r.returncode == 0, f"masked probe aborted: {r.stderr}"
+ assert (
+ "OK:gfx1151" in r.stdout
+ ), f"a visibility mask must not blind the gfx probe: {r.stdout!r}"
+
+ def test_kfd_only_inferable_gfx_defers_to_reroute(self):
+ """A KFD-only host (GPU detected, gfx unreadable) whose arch IS inferable
+ from hardware IDs must not print the 'installing CPU-only PyTorch' warning:
+ get_torch_index_url returns the cpu index quietly and the runtime-less
+ reroute upgrades it to AMD per-arch wheels. Only when inference also fails
+ (or maps to no supported family) is CPU final, with the actionable hint."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute get_torch_index_url")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ fn = _extract_sh_function_body(source, "get_torch_index_url")
+ probe_fn = _extract_sh_function_body(source, "_probe_amd_gfx_arch")
+ family_fn = _extract_sh_function_body(source, "_amd_arch_index_family_for_gfx")
+ assert fn and probe_fn and family_fn
+ with tempfile.TemporaryDirectory() as d:
+ # uname -> Linux/x86_64 so the AMD branch runs on any dev host; the
+ # rocminfo/amd-smi shims enumerate nothing (KFD-only host).
+ with open(os.path.join(d, "uname"), "w", encoding = "utf-8", newline = "\n") as f:
+ f.write('#!/bin/sh\ncase "${1:-}" in -m) echo x86_64 ;; *) echo Linux ;; esac\n')
+ for name in ("rocminfo", "amd-smi"):
+ with open(os.path.join(d, name), "w", encoding = "utf-8", newline = "\n") as f:
+ f.write("#!/bin/sh\nexit 0\n")
+ for name in ("uname", "rocminfo", "amd-smi"):
+ os.chmod(os.path.join(d, name), 0o755)
+
+ def run(infer_stub):
+ script = (
+ "set -euo pipefail\n"
+ "_ensure_rocm_probe_env() { :; }\n"
+ "_trim_index_path_slashes() { printf '%s\\n' \"$1\"; }\n"
+ "_has_usable_nvidia_gpu() { return 1; }\n"
+ "_has_amd_rocm_gpu() { return 0; }\n"
+ + infer_stub
+ + "\n"
+ + probe_fn
+ + "\n"
+ + family_fn
+ + "\n"
+ + fn
+ + "\n"
+ "get_torch_index_url\n"
+ )
+ # Run from a file, not -c: Windows bash mangles multi-KB -c strings.
+ sp = os.path.join(d, "gtiu.sh")
+ with open(sp, "w", encoding = "utf-8", newline = "\n") as f:
+ f.write(script)
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""))
+ for var in (
+ "UNSLOTH_ROCM_GFX_ARCH",
+ "UNSLOTH_TORCH_INDEX_URL",
+ "UNSLOTH_TORCH_INDEX_FAMILY",
+ "UNSLOTH_PYTORCH_MIRROR",
+ "ROCR_VISIBLE_DEVICES",
+ "HIP_VISIBLE_DEVICES",
+ ):
+ env.pop(var, None)
+ return subprocess.run(
+ [shell, sp.replace("\\", "/")], env = env, capture_output = True, text = True
+ )
+
+ r = run("_infer_linux_amd_gfx_arch() { echo gfx1100; }")
+ assert r.returncode == 0, f"inferable case aborted: {r.stderr}"
+ assert r.stdout.strip().endswith(
+ "/cpu"
+ ), f"must hand */cpu to the reroute: {r.stdout!r}"
+ assert (
+ "inferring gfx1100" in r.stderr
+ ), f"must announce the inference handoff: {r.stderr!r}"
+ assert (
+ "installing CPU-only PyTorch" not in r.stderr
+ ), f"must not promise a CPU-only install the reroute will override: {r.stderr!r}"
+ r2 = run("_infer_linux_amd_gfx_arch() { return 1; }")
+ assert r2.returncode == 0, f"uninferable case aborted: {r2.stderr}"
+ assert r2.stdout.strip().endswith("/cpu")
+ assert (
+ "installing CPU-only PyTorch" in r2.stderr
+ ), f"uninferable gfx must keep the actionable CPU warning: {r2.stderr!r}"
+ r3 = run("_infer_linux_amd_gfx_arch() { echo gfx906; }")
+ assert r3.returncode == 0, f"unsupported-family case aborted: {r3.stderr}"
+ assert r3.stdout.strip().endswith("/cpu")
+ assert (
+ "installing CPU-only PyTorch" in r3.stderr
+ ), f"an inferred arch with no wheel family must keep the CPU warning: {r3.stderr!r}"
+
+ def test_no_version_cpu_warning_respects_gfx_override(self):
+ """With UNSLOTH_ROCM_GFX_ARCH set on a KFD-only host that has no ROCm
+ version sources, the gfx probe is seeded by the override, so the
+ no-version endpoint used to print 'falling back to CPU-only PyTorch'
+ even though the reroute then installs the per-arch wheels (Codex P3).
+ A supported override must defer; an unsupported override, or a
+ readable-gfx host without an override, keeps the CPU warning."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute get_torch_index_url")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ fn = _extract_sh_function_body(source, "get_torch_index_url")
+ probe_fn = _extract_sh_function_body(source, "_probe_amd_gfx_arch")
+ family_fn = _extract_sh_function_body(source, "_amd_arch_index_family_for_gfx")
+ assert fn and probe_fn and family_fn
+ with tempfile.TemporaryDirectory() as d:
+ with open(os.path.join(d, "uname"), "w", encoding = "utf-8", newline = "\n") as f:
+ f.write('#!/bin/sh\ncase "${1:-}" in -m) echo x86_64 ;; *) echo Linux ;; esac\n')
+ # Silence every ROCm version source, not just amd-smi: a dev box with
+ # a real hipconfig/dpkg would otherwise resolve a version and skip
+ # the no-version endpoint this test exercises.
+ with open(os.path.join(d, "amd-smi"), "w", encoding = "utf-8", newline = "\n") as f:
+ f.write("#!/bin/sh\nexit 0\n")
+ for name in ("hipconfig", "dpkg-query", "rpm"):
+ with open(os.path.join(d, name), "w", encoding = "utf-8", newline = "\n") as f:
+ f.write("#!/bin/sh\nexit 1\n")
+ for name in ("uname", "amd-smi", "hipconfig", "dpkg-query", "rpm"):
+ os.chmod(os.path.join(d, name), 0o755)
+ script = (
+ "set -euo pipefail\n"
+ "_ensure_rocm_probe_env() { :; }\n"
+ "_trim_index_path_slashes() { printf '%s\\n' \"$1\"; }\n"
+ "_has_usable_nvidia_gpu() { return 1; }\n"
+ "_has_amd_rocm_gpu() { return 0; }\n"
+ "_infer_linux_amd_gfx_arch() { return 1; }\n"
+ + probe_fn
+ + "\n"
+ + family_fn
+ + "\n"
+ + fn
+ + "\n"
+ "get_torch_index_url\n"
+ )
+ sp = os.path.join(d, "gtiu.sh")
+ with open(sp, "w", encoding = "utf-8", newline = "\n") as f:
+ f.write(script)
+
+ def run(rocminfo_body, **extra):
+ with open(os.path.join(d, "rocminfo"), "w", encoding = "utf-8", newline = "\n") as f:
+ f.write("#!/bin/sh\n" + rocminfo_body)
+ os.chmod(os.path.join(d, "rocminfo"), 0o755)
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra)
+ for var in (
+ "UNSLOTH_TORCH_INDEX_URL",
+ "UNSLOTH_TORCH_INDEX_FAMILY",
+ "UNSLOTH_PYTORCH_MIRROR",
+ "ROCR_VISIBLE_DEVICES",
+ "HIP_VISIBLE_DEVICES",
+ ):
+ env.pop(var, None)
+ if "UNSLOTH_ROCM_GFX_ARCH" not in extra:
+ env.pop("UNSLOTH_ROCM_GFX_ARCH", None)
+ return subprocess.run(
+ [shell, sp.replace("\\", "/")], env = env, capture_output = True, text = True
+ )
+
+ # Supported override on a tool-blind host: defer to the reroute.
+ r = run("exit 0\n", UNSLOTH_ROCM_GFX_ARCH = "gfx1151")
+ assert r.returncode == 0, f"override case aborted: {r.stderr}"
+ assert r.stdout.strip().endswith("/cpu")
+ assert (
+ "falling back to CPU-only PyTorch" not in r.stderr
+ ), f"a supported override must not get the false CPU warning: {r.stderr!r}"
+ assert (
+ "UNSLOTH_ROCM_GFX_ARCH=gfx1151 is set" in r.stderr
+ ), f"the override deferral must be announced: {r.stderr!r}"
+ # Unsupported override: the reroute can't map it -> CPU warning stays.
+ r2 = run("exit 0\n", UNSLOTH_ROCM_GFX_ARCH = "gfx906")
+ assert r2.returncode == 0, f"unsupported-override case aborted: {r2.stderr}"
+ assert (
+ "falling back to CPU-only PyTorch" in r2.stderr
+ ), f"an unmappable override must keep the CPU warning: {r2.stderr!r}"
+ # Readable gfx, no override, no version: deliberate CPU fallback.
+ r3 = run('echo " Name: gfx1151"\n')
+ assert r3.returncode == 0, f"readable-gfx case aborted: {r3.stderr}"
+ assert (
+ "falling back to CPU-only PyTorch" in r3.stderr
+ ), f"a readable-gfx host without a version keeps the CPU warning: {r3.stderr!r}"
+
+ def test_reroute_gate_covers_kfd_only(self):
+ """The runtime-less reroute must fire for a KFD-only host: _has_amd_rocm_gpu
+ is now true via the KFD topology, so the gate also accepts a detected GPU
+ whose gfx probe is empty (unslothai#7314 P2). A */cpu index chosen with a
+ READABLE gfx (deliberate ROCm-version fallback) must stay un-rerouted."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute the reroute block")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ block = re.search(
+ r'^if \[ "\$_torch_index_pinned" = false \] && \[ "\$SKIP_TORCH" = false \] && \\\n'
+ r".*?^fi\n",
+ source,
+ re.S | re.M,
+ )
+ assert block, "could not extract the runtime-less reroute block"
+ family_fn = _extract_sh_function_body(source, "_amd_arch_index_family_for_gfx")
+ assert family_fn
+ with tempfile.TemporaryDirectory() as d:
+ with open(os.path.join(d, "uname"), "w", encoding = "utf-8", newline = "\n") as f:
+ f.write('#!/bin/sh\ncase "${1:-}" in -m) echo x86_64 ;; *) echo Linux ;; esac\n')
+ os.chmod(os.path.join(d, "uname"), 0o755)
+
+ def run(gpu_stub, probe_stub):
+ script = (
+ "set -euo pipefail\n"
+ "_has_usable_nvidia_gpu() { return 1; }\n"
+ f"_has_amd_rocm_gpu() {{ {gpu_stub}; }}\n"
+ f"_probe_amd_gfx_arch() {{ {probe_stub}; }}\n"
+ "_infer_linux_amd_gfx_arch() { echo gfx1100; }\n"
+ "_strip_index_url_credentials() { printf '%s\\n' \"$1\"; }\n" + family_fn + "\n"
+ "_torch_index_pinned=false\nSKIP_TORCH=false\n_ARCH=x86_64\n"
+ "TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu\n"
+ + block.group(0)
+ + 'printf "URL:%s GFX:%s\\n" "$TORCH_INDEX_URL" "${UNSLOTH_ROCM_GFX_ARCH:-}"\n'
+ )
+ # Run from a file, not -c: Windows bash mangles multi-KB -c strings.
+ sp = os.path.join(d, "reroute.sh")
+ with open(sp, "w", encoding = "utf-8", newline = "\n") as f:
+ f.write(script)
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""))
+ for var in ("UNSLOTH_ROCM_GFX_ARCH", "UNSLOTH_AMD_ROCM_MIRROR"):
+ env.pop(var, None)
+ return subprocess.run(
+ [shell, sp.replace("\\", "/")], env = env, capture_output = True, text = True
+ )
+
+ # KFD-only: GPU detected, probe empty -> reroute to per-arch wheels.
+ r = run("return 0", "printf '\\n'")
+ assert r.returncode == 0, f"kfd-only reroute aborted: {r.stderr}"
+ assert (
+ "URL:https://repo.amd.com/rocm/whl/gfx110X-all/ GFX:gfx1100" in r.stdout
+ ), f"KFD-only host must reach the AMD arch index: {r.stdout!r}"
+ # The diagnostic must not claim /dev/kfd is missing: KFD visibility is
+ # exactly what detected this host (Codex P3).
+ assert (
+ "ROCm runtime not visible" not in r.stderr
+ ), f"KFD-only reroute must not claim /dev/kfd is missing: {r.stderr!r}"
+ assert (
+ "visible via the kernel driver (KFD)" in r.stderr
+ ), f"KFD-only reroute must name the tooling gap: {r.stderr!r}"
+ # Readable gfx: the */cpu index is a deliberate fallback -> untouched.
+ r2 = run("return 0", "echo gfx1151")
+ assert r2.returncode == 0, f"readable-gfx case aborted: {r2.stderr}"
+ assert (
+ "URL:https://download.pytorch.org/whl/cpu GFX:" in r2.stdout
+ ), f"a deliberate CPU fallback must not be rerouted: {r2.stdout!r}"
+ # No AMD GPU detected at all: the pre-KFD-fix path still reroutes.
+ r3 = run("return 1", "printf '\\n'")
+ assert r3.returncode == 0, f"undetected-GPU case aborted: {r3.stderr}"
+ assert (
+ "URL:https://repo.amd.com/rocm/whl/gfx110X-all/ GFX:gfx1100" in r3.stdout
+ ), f"the original undetected-GPU reroute must keep working: {r3.stdout!r}"
+ assert (
+ "ROCm runtime not visible" in r3.stderr
+ ), f"a truly runtime-invisible host keeps the original diagnostic: {r3.stderr!r}"
def test_get_torch_index_url_uses_nvidia_detected_flag(self):
"""get_torch_index_url must track NVIDIA via _nvidia_detected (proc-only NVIDIA still picks CUDA)."""
@@ -3641,7 +4123,9 @@ class TestStrixRocm71Override:
pytest.skip("bash needed to execute the probe block")
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
block = re.search(
- r'^ _gfx_all=""\n.*?(?=^ _strix_gfx="")', source, re.S | re.M
+ r'^ _gfx_all=\$\(printf[^\n]*\n.*?(?=^ _strix_gfx="")',
+ source,
+ re.S | re.M,
)
assert block, "could not extract the gfx-detection block"
with tempfile.TemporaryDirectory() as d:
@@ -3661,6 +4145,67 @@ class TestStrixRocm71Override:
assert r.returncode == 0, f"probe aborted under set -e: {r.stderr}"
assert "OK:gfx1151" in r.stdout, f"amd-smi fallback not reached: {r.stdout!r}"
+ def test_strix_reroute_reprobes_when_mask_hides_all(self):
+ """A visibility mask hiding every agent (ROCR_VISIBLE_DEVICES=-1) must not
+ skip the Strix reroute: get_torch_index_url reads the arch unmasked, so
+ the reroute must re-probe unmasked too or a masked Strix box gets the
+ broken generic wheels. A partial mask must keep its per-GPU selection.
+ Executed with mask-honouring shims, not a text match."""
+ shell = shutil.which("bash")
+ if not shell:
+ pytest.skip("bash needed to execute the probe block")
+ source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
+ block = re.search(
+ r'^ _gfx_all=\$\(printf[^\n]*\n.*?(?=^ _strix_gfx="")',
+ source,
+ re.S | re.M,
+ )
+ assert block, "could not extract the gfx-detection block"
+ with tempfile.TemporaryDirectory() as d:
+ # rocminfo honours ROCR_VISIBLE_DEVICES like the real tool: -1 and
+ # set-but-empty hide both agents, 1 renumbers to the dGPU only,
+ # unset shows both.
+ rocminfo = (
+ "#!/bin/sh\n"
+ 'case "${ROCR_VISIBLE_DEVICES-__unset__}" in\n'
+ ' __unset__) printf "Name: gfx1151\\nName: gfx1201\\n" ;;\n'
+ ' ""|-1) echo "no visible agents" ;;\n'
+ ' 1) printf "Name: gfx1201\\n" ;;\n'
+ ' *) printf "Name: gfx1151\\nName: gfx1201\\n" ;;\n'
+ "esac\n"
+ )
+ for name, body in (("rocminfo", rocminfo), ("amd-smi", "#!/bin/sh\nexit 0\n")):
+ p = os.path.join(d, name)
+ with open(p, "w", encoding = "utf-8") as f:
+ f.write(body)
+ os.chmod(p, 0o755)
+ script = (
+ "set -euo pipefail\n" + block.group(0) + '\nprintf "OK:%s\\n" "$_runtime_gfx"\n'
+ )
+
+ def run(**extra):
+ env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra)
+ env.pop("UNSLOTH_ROCM_GFX_ARCH", None)
+ env.pop("HIP_VISIBLE_DEVICES", None)
+ return subprocess.run(
+ [shell, "-c", script], env = env, capture_output = True, text = True
+ )
+
+ # Mask hides everything: re-probe must recover the first GPU (Strix).
+ r = run(ROCR_VISIBLE_DEVICES = "-1")
+ assert r.returncode == 0, f"masked probe aborted: {r.stderr}"
+ assert "OK:gfx1151" in r.stdout, f"reroute blinded by full mask: {r.stdout!r}"
+ # A SET-but-empty mask also hides every agent and must re-probe too
+ # (the ${VAR+x} guard, not ${VAR:-}).
+ r0 = run(ROCR_VISIBLE_DEVICES = "")
+ assert r0.returncode == 0, f"empty-mask probe aborted: {r0.stderr}"
+ assert "OK:gfx1151" in r0.stdout, f"reroute blinded by empty mask: {r0.stdout!r}"
+ # Partial mask: enumeration already reflects it; the dGPU selection
+ # must survive (no unmasked re-probe overriding the user's pick).
+ r2 = run(ROCR_VISIBLE_DEVICES = "1")
+ assert r2.returncode == 0, f"partial-mask probe aborted: {r2.stderr}"
+ assert "OK:gfx1201" in r2.stdout, f"partial mask selection lost: {r2.stdout!r}"
+
def test_strix_routing_helpers_cover_rocm714(self):
# Reroute for any generic pytorch.org index below the 7.13 arch floor (7.0,
# 7.2, a future 7.3+), never at/above it -- mirrors install.sh _rocm_leaf_below.
From 13c7db1965da31cd9427cdf46d6ee6448158aa19 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Thu, 23 Jul 2026 13:14:37 +0530
Subject: [PATCH 079/255] Studio: reject whitespace-only passwords (#7341)
* Studio: reject whitespace-only passwords
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: reject any whitespace in passwords
* Studio: surface whitespace error in setup form, isolate auth test import
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
---
studio/backend/auth/terminal_prompt.py | 4 +
studio/backend/routes/auth.py | 5 ++
studio/backend/run.py | 7 ++
.../tests/test_change_password_policy.py | 75 +++++++++++++++++++
studio/backend/tests/test_password_prompt.py | 16 ++++
.../features/auth/components/auth-form.tsx | 18 ++++-
.../components/change-password-dialog.tsx | 10 ++-
studio/frontend/src/i18n/locales/en.ts | 1 +
unsloth_cli/commands/_password_prompt.py | 6 ++
9 files changed, 136 insertions(+), 6 deletions(-)
create mode 100644 studio/backend/tests/test_change_password_policy.py
diff --git a/studio/backend/auth/terminal_prompt.py b/studio/backend/auth/terminal_prompt.py
index e855f4078b..925404f47d 100644
--- a/studio/backend/auth/terminal_prompt.py
+++ b/studio/backend/auth/terminal_prompt.py
@@ -236,6 +236,10 @@ def prompt_for_password_change(
out.write(f"Password must be at least {min_length} characters; try again.\n")
out.flush()
continue
+ if any(ch.isspace() for ch in new_password):
+ out.write("Password cannot contain spaces; try again.\n")
+ out.flush()
+ continue
if is_current_password(new_password):
out.write(
"New password must differ from the current bootstrap password; try again.\n"
diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py
index d779c8784e..1acc48e3a3 100644
--- a/studio/backend/routes/auth.py
+++ b/studio/backend/routes/auth.py
@@ -494,6 +494,11 @@ async def change_password(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Current password is incorrect",
)
+ if any(ch.isspace() for ch in payload.new_password):
+ raise HTTPException(
+ status_code = status.HTTP_400_BAD_REQUEST,
+ detail = "New password cannot contain spaces",
+ )
if payload.current_password == payload.new_password:
raise HTTPException(
status_code = status.HTTP_400_BAD_REQUEST,
diff --git a/studio/backend/run.py b/studio/backend/run.py
index 398943cc2c..d9569c46f6 100644
--- a/studio/backend/run.py
+++ b/studio/backend/run.py
@@ -1244,6 +1244,13 @@ def _apply_supplied_password(password_value: "Optional[str]") -> None:
flush = True,
)
sys.exit(1)
+ if any(ch.isspace() for ch in supplied):
+ print(
+ "Error: password cannot contain spaces; not starting.",
+ file = sys.stderr,
+ flush = True,
+ )
+ sys.exit(1)
if _is_current_password(supplied):
print(
"Error: the new password must differ from the current bootstrap "
diff --git a/studio/backend/tests/test_change_password_policy.py b/studio/backend/tests/test_change_password_policy.py
new file mode 100644
index 0000000000..c73e9ed839
--- /dev/null
+++ b/studio/backend/tests/test_change_password_policy.py
@@ -0,0 +1,75 @@
+# 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
+import sys
+from pathlib import Path
+
+import pytest
+from fastapi import HTTPException
+
+_BACKEND_ROOT = Path(__file__).resolve().parents[1]
+if str(_BACKEND_ROOT) not in sys.path:
+ sys.path.insert(0, str(_BACKEND_ROOT))
+
+from models.auth import ChangePasswordRequest # noqa: E402
+
+# Load routes/auth.py directly so collection does not execute routes/__init__.py,
+# which pulls in the heavy training/models/inference routers.
+_route_path = _BACKEND_ROOT / "routes" / "auth.py"
+_spec = importlib.util.spec_from_file_location("_change_password_route", _route_path)
+assert _spec is not None and _spec.loader is not None
+auth_routes = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(auth_routes)
+
+
+@pytest.fixture
+def _user(monkeypatch):
+ monkeypatch.setattr(
+ auth_routes.storage,
+ "get_user_and_secret",
+ lambda username: ("salt", "hash", "jwt-secret", False),
+ )
+ monkeypatch.setattr(
+ auth_routes.hashing,
+ "verify_password",
+ lambda password, salt, pwd_hash: password == "bootstrap-pw",
+ )
+
+
+def _change(new_password):
+ payload = ChangePasswordRequest(
+ current_password = "bootstrap-pw",
+ new_password = new_password,
+ )
+ return asyncio.run(auth_routes.change_password(payload, None, "unsloth"))
+
+
+def test_rejects_whitespace_only_password(_user):
+ with pytest.raises(HTTPException) as excinfo:
+ _change(" " * 8)
+ assert excinfo.value.status_code == 400
+ assert "spaces" in excinfo.value.detail
+
+
+def test_rejects_tabs_and_spaces_password(_user):
+ with pytest.raises(HTTPException) as excinfo:
+ _change(" \t \t \t \t ")
+ assert excinfo.value.status_code == 400
+
+
+def test_rejects_password_containing_spaces(_user):
+ with pytest.raises(HTTPException) as excinfo:
+ _change("correct horse battery")
+ assert excinfo.value.status_code == 400
+ assert "spaces" in excinfo.value.detail
+
+
+def test_allows_password_without_spaces(_user, monkeypatch):
+ monkeypatch.setattr(auth_routes.storage, "update_password", lambda *args, **kwargs: True)
+ monkeypatch.setattr(auth_routes, "create_access_token", lambda subject: "at")
+ monkeypatch.setattr(auth_routes, "create_refresh_token", lambda subject: "rt")
+ token = _change("correct-horse-battery")
+ assert token.access_token == "at"
+ assert token.must_change_password is False
diff --git a/studio/backend/tests/test_password_prompt.py b/studio/backend/tests/test_password_prompt.py
index 372d6a2aa4..1af8836065 100644
--- a/studio/backend/tests/test_password_prompt.py
+++ b/studio/backend/tests/test_password_prompt.py
@@ -183,6 +183,22 @@ def test_loop_short_password_reprompts(monkeypatch):
assert "at least 8 characters" in out
+def test_loop_whitespace_only_reprompts(monkeypatch):
+ ok, applied, out = _run_loop(monkeypatch, _keys(" " * 8, "long-enough-pw", "long-enough-pw"))
+ assert ok is True
+ assert applied == ["long-enough-pw"]
+ assert "contain spaces" in out
+
+
+def test_loop_password_with_inner_space_reprompts(monkeypatch):
+ ok, applied, out = _run_loop(
+ monkeypatch, _keys("has space pw", "long-enough-pw", "long-enough-pw")
+ )
+ assert ok is True
+ assert applied == ["long-enough-pw"]
+ assert "contain spaces" in out
+
+
def test_loop_rejects_current_password(monkeypatch):
ok, applied, out = _run_loop(
monkeypatch, _keys("bootstrap-pw", "fresh-password", "fresh-password")
diff --git a/studio/frontend/src/features/auth/components/auth-form.tsx b/studio/frontend/src/features/auth/components/auth-form.tsx
index 3eec1dba88..72181b8e4f 100644
--- a/studio/frontend/src/features/auth/components/auth-form.tsx
+++ b/studio/frontend/src/features/auth/components/auth-form.tsx
@@ -196,8 +196,10 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
!isLoginMode &&
(currentPassword.length < 8 ||
newPassword.length < 8 ||
+ /\s/.test(newPassword) ||
newPassword !== confirmPassword ||
currentPassword === newPassword);
+ const showWhitespaceWarning = !isLoginMode && /\s/.test(newPassword);
const showPasswordMismatchWarning =
!isLoginMode &&
newPassword.length > 0 &&
@@ -222,6 +224,10 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
setError("New password must be at least 8 characters.");
return;
}
+ if (/\s/.test(newPassword)) {
+ setError("New password cannot contain spaces.");
+ return;
+ }
if (newPassword !== confirmPassword) {
setError("Passwords do not match.");
return;
@@ -425,13 +431,17 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null {
- {showPasswordMismatchWarning
- ? "Please ensure passwords match."
- : "Must be at least 8 characters."}
+ {showWhitespaceWarning
+ ? "New password cannot contain spaces."
+ : showPasswordMismatchWarning
+ ? "Please ensure passwords match."
+ : "Must be at least 8 characters."}
>
)}
diff --git a/studio/frontend/src/features/settings/components/change-password-dialog.tsx b/studio/frontend/src/features/settings/components/change-password-dialog.tsx
index cd30d37d5d..c88fc48cac 100644
--- a/studio/frontend/src/features/settings/components/change-password-dialog.tsx
+++ b/studio/frontend/src/features/settings/components/change-password-dialog.tsx
@@ -78,6 +78,9 @@ function passwordValidationMessage(
minLength: MIN_PASSWORD_LENGTH,
});
}
+ if (/\s/.test(nextPassword)) {
+ return t("settings.general.passwordDialog.newHasSpaces");
+ }
if (nextPassword !== confirmPassword) {
return t("settings.general.passwordDialog.mismatch");
}
@@ -160,6 +163,7 @@ export function ChangePasswordDialog() {
const currentTooShort = hasStartedTooShortPassword(current);
const nextTooShort = hasStartedTooShortPassword(next);
+ const nextHasSpaces = /\s/.test(next);
const mismatch = confirm.length > 0 && next !== confirm;
const samePassword = hasReusablePassword(current, next);
const validationMessage = passwordValidationMessage(
@@ -279,13 +283,15 @@ export function ChangePasswordDialog() {
minLength={MIN_PASSWORD_LENGTH}
disabled={submitting}
/>
- {nextTooShort || samePassword ? (
+ {nextTooShort || nextHasSpaces || samePassword ? (
{nextTooShort
? t("settings.general.passwordDialog.newTooShort", {
minLength: MIN_PASSWORD_LENGTH,
})
- : t("settings.general.passwordDialog.samePassword")}
+ : nextHasSpaces
+ ? t("settings.general.passwordDialog.newHasSpaces")
+ : t("settings.general.passwordDialog.samePassword")}
) : null}
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index cf8a29b6d2..164833b41d 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -196,6 +196,7 @@ export const en = {
currentTooShort:
"Current password must be at least {minLength} characters.",
newTooShort: "New password must be at least {minLength} characters.",
+ newHasSpaces: "New password cannot contain spaces.",
mismatch: "Passwords do not match.",
samePassword:
"New password must be different from your current password.",
diff --git a/unsloth_cli/commands/_password_prompt.py b/unsloth_cli/commands/_password_prompt.py
index b6fd8ca34d..55f50acbf1 100644
--- a/unsloth_cli/commands/_password_prompt.py
+++ b/unsloth_cli/commands/_password_prompt.py
@@ -191,6 +191,10 @@ def prompt_new_password(verify_current: Callable[[str], bool], out: TextIO | Non
out.write(f"Password must be at least {MIN_PASSWORD_LENGTH} characters. Try again.\n")
out.flush()
continue
+ if any(ch.isspace() for ch in password):
+ out.write("Password cannot contain spaces. Try again.\n")
+ out.flush()
+ continue
if verify_current(password):
out.write("New password must differ from the current password. Try again.\n")
out.flush()
@@ -233,6 +237,8 @@ def validate_new_password(candidate: str, verify_current: Callable[[str], bool])
current password), else None. Same policy as the interactive loop."""
if len(candidate) < MIN_PASSWORD_LENGTH:
return f"Password must be at least {MIN_PASSWORD_LENGTH} characters."
+ if any(ch.isspace() for ch in candidate):
+ return "Password cannot contain spaces."
if verify_current(candidate):
return "New password must differ from the current password."
return None
From fa5498db0b6c089c1c9ddc8e82043d82be202bc6 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Thu, 23 Jul 2026 00:44:42 -0700
Subject: [PATCH 080/255] Studio: UI font size scales all text consistently
without moving layout (#7355)
* Studio: make UI font size scale all text without moving layout
The UI font size setting changes the root rem base, so only rem sized
text reacted. Hundreds of px text classes, px font sizes in CSS, and
chart labels stayed fixed, while rem based padding, widths and radii
wrongly grew.
Convert all text sizes to rem so every font follows the setting, and
pin spacing, radius, container widths, sidebar and thread widths to px
so layout no longer follows the rem base. Library styles (streamdown,
react-flow) are re-based via overrides. All conversions are exact at
the default 16px root, so the default rendering is unchanged.
* Studio: keep logo at fixed size and fit tight controls at large UI fonts
The logo lockups (sidebar wordmark with beta badge, onboarding wizard)
are branding and now keep px sizes at any UI font size.
Two controls clipped their text at the largest setting: the appearance
color chips (fixed w-24) and the voice tab selects (fixed w-56). Both
use min widths now, so they keep the default look at 16px and only
grow when the text needs the room.
* Studio: keep dropdown corners rounded when the menu scrolls
A scrolling dropdown lost its rounded corners on the scrollbar side:
WebKit paints the surface square when the rounded element itself hosts
the scrollbar, which shows up in the desktop app whenever a menu
overflows, for example at larger UI font sizes.
Dropdown menu and select content now clip with overflow hidden and
scroll an inner viewport instead. The surface padding insets the
scrollbar clear of the curve, so corners stay rounded in every engine.
Submenus are unaffected since sub content is portaled.
* Studio: scale the logo lockups at half the UI font size rate
Rather than pinning the logo, the sidebar lockup (sticker, wordmark,
beta badge) and the onboarding lockup now follow the UI font size at
half the rate of the change: size = base + (root - 16px) / 2, written
as calc((base - 8)px + 0.5rem). A 4px font size change moves the logo
by 2px, and the default 16px root renders the exact base sizes.
* Studio: address review feedback on leading, grid tracks and select scrolling
Numeric leading utilities (leading-3 through leading-10) derive from
--spacing, so pinning spacing to px also froze their line-heights while
the paired text sizes now scale. Define them as rem theme tokens so
line-height follows the UI font size again; values are identical at the
16px default.
Convert the grid tracks the rem-to-px codemod missed (rem followed by
an underscore escaped the word boundary): the response details label
column and the on-device folder rows.
Make the Radix select viewport the bounded scroller instead of a
wrapper div, so Radix's scroll handling and the browser scroll the same
element. Restore the app's thin scrollbar with an inline style, which
beats the scrollbar hiding stylesheet Radix injects at runtime.
* Studio: cap voice select widths and update CI contracts
---
studio/frontend/src/app/provider.tsx | 8 +-
.../frontend/src/components/app-sidebar.tsx | 34 ++--
.../components/assistant-ui/audio-player.tsx | 2 +-
.../message-response-details-sheet.tsx | 4 +-
.../assistant-ui/message-timing.tsx | 2 +-
.../src/components/assistant-ui/reasoning.tsx | 2 +-
.../src/components/assistant-ui/sources.tsx | 2 +-
.../src/components/assistant-ui/thread.tsx | 30 +--
.../assistant-ui/tool-ui-knowledge-base.tsx | 2 +-
.../assistant-ui/tool-ui-render-html.tsx | 2 +-
.../src/components/floating-monitor.tsx | 6 +-
.../src/components/llama-update-banner.tsx | 8 +-
.../frontend/src/components/section-card.tsx | 2 +-
.../src/components/tauri/startup-screen.tsx | 2 +-
.../src/components/tauri/update-banner.tsx | 16 +-
.../src/components/tauri/update-screen.tsx | 4 +-
.../src/components/tauri/window-titlebar.tsx | 8 +-
studio/frontend/src/components/ui/chart.tsx | 2 +-
.../src/components/ui/copyable-error-chip.tsx | 6 +-
.../frontend/src/components/ui/data-table.tsx | 2 +-
studio/frontend/src/components/ui/dialog.tsx | 2 +-
.../src/components/ui/dropdown-menu.tsx | 17 +-
.../src/components/ui/input-group.tsx | 4 +-
studio/frontend/src/components/ui/select.tsx | 12 +-
studio/frontend/src/components/ui/sidebar.tsx | 8 +-
.../src/components/web/update-banner.tsx | 8 +-
.../frontend/src/features/auth/login-page.tsx | 2 +-
.../features/chat/artifacts/artifact-card.tsx | 4 +-
.../frontend/src/features/chat/chat-page.tsx | 38 ++--
.../features/chat/chat-providers-dialog.tsx | 8 +-
.../src/features/chat/chat-settings-sheet.tsx | 68 +++----
.../chat/components/chat-search-dialog.tsx | 6 +-
.../chat/components/context-usage-bar.tsx | 4 +-
.../chat/components/model-load-status.tsx | 12 +-
.../components/openai-code-exec-section.tsx | 14 +-
.../chat/components/project-switcher.tsx | 2 +-
.../chat/hooks/use-chat-model-runtime.ts | 2 +-
.../features/chat/permission-mode-select.tsx | 2 +-
.../src/features/chat/projects-page.tsx | 16 +-
.../prompt-storage/prompt-storage-dialog.tsx | 10 +-
.../src/features/chat/thread-sidebar.tsx | 2 +-
.../data-recipes/pages/data-recipes-page.tsx | 8 +-
.../export/components/export-run-panel.tsx | 24 +--
.../export/components/method-picker.tsx | 2 +-
.../export/components/quant-picker.tsx | 10 +-
.../src/features/export/export-page.tsx | 36 ++--
.../features/hub/catalog/catalog-states.tsx | 32 ++--
.../hub/catalog/dataset-download-section.tsx | 2 +-
.../src/features/hub/catalog/dot-tag.tsx | 2 +-
.../features/hub/catalog/download-card.tsx | 2 +-
.../catalog/external-link-confirm-dialog.tsx | 4 +-
.../hub/catalog/gguf-download-card.tsx | 10 +-
.../hub/catalog/gguf-status-cards.tsx | 4 +-
.../features/hub/catalog/hub-detail-view.tsx | 2 +-
.../features/hub/catalog/hub-option-menu.tsx | 4 +-
.../features/hub/catalog/hub-section-row.tsx | 2 +-
.../hub/catalog/local-dataset-card.tsx | 2 +-
.../hub/catalog/local-on-device-card.tsx | 20 +-
.../src/features/hub/catalog/model-card.tsx | 6 +-
.../features/hub/catalog/model-inspector.tsx | 36 ++--
.../src/features/hub/catalog/model-readme.tsx | 20 +-
.../hub/catalog/models-catalog-lists.tsx | 12 +-
.../hub/catalog/models-catalog-rows.tsx | 30 +--
.../features/hub/catalog/models-header.tsx | 4 +-
.../src/features/hub/catalog/models-table.tsx | 46 ++---
.../features/hub/catalog/models-toolbar.tsx | 10 +-
.../hub/catalog/on-device-folders-dialog.tsx | 24 +--
.../src/features/hub/catalog/owner-avatar.tsx | 8 +-
.../hub/catalog/owner-scope-toggle.tsx | 2 +-
.../features/hub/catalog/recent-searches.tsx | 6 +-
.../hub/catalog/safetensors-download-card.tsx | 2 +-
.../hub/catalog/sampling-settings-dialog.tsx | 10 +-
.../src/features/hub/catalog/shared.tsx | 4 +-
.../hub/catalog/transport-conflict-dialog.tsx | 2 +-
.../features/hub/catalog/transport-toggle.tsx | 2 +-
.../hub/components/hf-token-indicator.tsx | 6 +-
.../features/hub/components/page-heading.tsx | 4 +-
.../download-manager-panel.tsx | 8 +-
.../download-progress-bar.tsx | 2 +-
studio/frontend/src/features/hub/hub-page.tsx | 2 +-
studio/frontend/src/features/hub/hub.css | 78 ++++----
.../chat-template-editor-dialog.tsx | 4 +-
.../components/model-config-page.tsx | 24 +--
.../components/model-selector.tsx | 14 +-
.../model-selector/folder-browser.tsx | 10 +-
.../components/model-selector/pickers.tsx | 62 +++----
.../components/model-selector/pill-tabs.tsx | 2 +-
.../components/native-model-chip.tsx | 2 +-
.../components/native-model-drop-overlay.tsx | 4 +-
.../components/steps/model-selection-step.tsx | 4 +-
.../components/steps/model-type-step.tsx | 2 +-
.../onboarding/components/wizard-sidebar.tsx | 8 +-
.../profile-personalization-panel.tsx | 4 +-
.../rag/components/document-preview-sheet.tsx | 2 +-
.../rag/components/document-status-chip.tsx | 2 +-
.../rag/components/project-sources-panel.tsx | 4 +-
.../components/retrieval-settings-section.tsx | 22 +--
.../recipe-studio/components/block-sheet.tsx | 4 +-
.../executions/execution-sidebar.tsx | 2 +-
.../components/executions/executions-view.tsx | 2 +-
.../inline/inline-category-badges.tsx | 6 +-
.../components/inline/inline-field.tsx | 2 +-
.../components/inline/inline-llm.tsx | 2 +-
.../components/inline/inline-seed.tsx | 6 +-
.../components/recipe-graph-node.tsx | 4 +-
.../components/recipe-studio-header.tsx | 10 +-
.../runtime/execution-progress-island.tsx | 16 +-
.../shared/available-references-inline.tsx | 14 +-
.../models/local-recipe-model-selector.tsx | 20 +-
.../recipe-studio/dialogs/preview-dialog.tsx | 2 +-
.../dialogs/seed/seed-dialog.tsx | 2 +-
.../tool-profile/tool-profile-dialog.tsx | 6 +-
.../easy/github-crawler-easy-view.tsx | 2 +-
.../recipe-studio/recipe-studio-page.tsx | 2 +-
.../features/recipe-studio/utils/ui-tones.ts | 6 +-
.../components/remote-code-consent-dialog.tsx | 8 +-
.../settings/components/api-key-row.tsx | 4 +-
.../components/api-monitor-console.tsx | 10 +-
.../settings/components/color-picker.tsx | 2 +-
.../settings/components/create-key-form.tsx | 2 +-
.../components/embedding-model-combobox.tsx | 4 +-
.../settings/components/key-reveal-card.tsx | 2 +-
.../settings/components/language-select.tsx | 2 +-
.../components/sidebar-menu-customizer.tsx | 4 +-
.../components/update-studio-instructions.tsx | 4 +-
.../components/uploaded-files-dialog.tsx | 6 +-
.../settings/components/usage-examples.tsx | 34 ++--
.../src/features/settings/settings-dialog.tsx | 14 +-
.../features/settings/tabs/resources-tab.tsx | 4 +-
.../src/features/settings/tabs/voice-tab.tsx | 12 +-
.../src/features/studio/history-card-grid.tsx | 12 +-
.../studio/recent-trainings-section.tsx | 2 +-
.../sections/charts/chart-settings-sheet.tsx | 2 +-
.../sections/charts/eval-loss-chart-card.tsx | 8 +-
.../sections/charts/grad-norm-chart-card.tsx | 4 +-
.../charts/learning-rate-chart-card.tsx | 4 +-
.../charts/training-loss-chart-card.tsx | 6 +-
.../dataset-preview-dialog-mapping.tsx | 14 +-
.../sections/dataset-preview-dialog.tsx | 20 +-
.../studio/sections/dataset-section.tsx | 14 +-
.../studio/sections/model-section.tsx | 16 +-
.../studio/sections/params-section.tsx | 16 +-
.../studio/sections/progress-section.tsx | 22 +--
.../studio/sections/s3-config-form.tsx | 2 +-
.../studio/sections/training-section.tsx | 4 +-
.../src/features/studio/studio-page.tsx | 2 +-
.../studio/training-start-overlay.tsx | 8 +-
.../features/tour/components/guided-tour.tsx | 8 +-
studio/frontend/src/index.css | 171 ++++++++++++------
.../test_chat_thinking_compact_layout.py | 2 +-
.../studio/test_compact_dropdown_submenus.py | 2 +-
.../test_studio_text_descender_clipping.py | 2 +-
.../test_voice_settings_select_width.py | 16 ++
153 files changed, 860 insertions(+), 762 deletions(-)
create mode 100644 tests/studio/test_voice_settings_select_width.py
diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx
index e6c89b9cd7..275c3c6623 100644
--- a/studio/frontend/src/app/provider.tsx
+++ b/studio/frontend/src/app/provider.tsx
@@ -213,7 +213,7 @@ function TauriUpdateLayer({
}
return (
-
+
+
- {label}
+ {label}
{spinner && (
)}
@@ -904,7 +904,7 @@ export function AppSidebar() {
? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
: "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto";
const buttonClass = cn(
- "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium",
+ "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[0.90625rem] leading-[1.1875rem] tracking-nav font-medium",
// pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the
// title with the nav items above.
variant === "project" ? "pl-[39px]" : "pl-3",
@@ -939,7 +939,7 @@ export function AppSidebar() {
aria-label={translate("shell.dialog.renameChat.placeholder")}
className={cn(
// No pill or box; edit in place as plain highlighted text.
- "text-foreground h-[33px] w-full border-0 bg-transparent pr-4 text-[14.5px] leading-[19px] font-medium tracking-nav outline-none",
+ "text-foreground h-[33px] w-full border-0 bg-transparent pr-4 text-[0.90625rem] leading-[1.1875rem] font-medium tracking-nav outline-none",
variant === "project" ? "pl-[39px]" : "pl-3",
)}
/>
@@ -1184,15 +1184,17 @@ export function AppSidebar() {
aria-disabled={chatDisabled}
tabIndex={chatDisabled ? -1 : undefined}
>
+ {/* Logo lockup follows the UI font size at half rate:
+ base + (root - 16px) / 2, written as (base - 8)px + 0.5rem. */}
-
+
unsloth
-
+
{t("shell.beta")}
@@ -1219,7 +1221,7 @@ export function AppSidebar() {
hidden={isMobile}
>
{t("shell.navigation.search")}
-
+
{isMacPlatform ? "⌘K" : "Ctrl+K"}
@@ -1536,7 +1538,7 @@ export function AppSidebar() {
className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8"
>
- {project.name}
+ {project.name}
{/* New chat in this project */}
-
+
{showAll ? "Show less" : "Show more"}
@@ -1709,7 +1711,7 @@ export function AppSidebar() {
>
{
setSelectedHistoryRunId(run.id);
// From Recipes/Export, jump to Train so the run's
@@ -1729,7 +1731,7 @@ export function AppSidebar() {
{getTrainingRunDisplayTitle(run)}
-
+
{formatRelativeShort(run.started_at)}
@@ -1830,11 +1832,11 @@ export function AppSidebar() {
/>
-
+
{t("shell.updateAvailable")}
{updateVersion && (
-
+
v{updateVersion}
)}
@@ -1871,8 +1873,8 @@ export function AppSidebar() {
{/* min-w-0 so long names truncate instead of overflowing;
pr on the button reserves room for the settings cog */}
- {displayTitle}
- Unsloth
+ {displayTitle}
+ Unsloth
@@ -1880,7 +1882,7 @@ export function AppSidebar() {
side="top"
align="center"
sideOffset={8}
- className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-2.5 py-2.5 font-heading rounded-[20px] border-0"
+ className="app-user-menu menu-soft-surface-up ring-0 w-[256px] px-2.5 py-2.5 font-heading rounded-[20px] border-0"
>
= ({ src }) => {
onChange={handleSeek}
className="h-1.5 w-full cursor-pointer accent-primary"
/>
-
+
{formatTime(progress)}
{formatTime(duration)}
diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx
index cca61b766a..5dd4f1c91d 100644
--- a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx
+++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx
@@ -207,7 +207,7 @@ function DetailRow({
}) {
if (value == null || value === "") return null;
return (
-
+
{label}
diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx
index 8d40f587ad..c602a776cf 100644
--- a/studio/frontend/src/components/assistant-ui/message-timing.tsx
+++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx
@@ -86,7 +86,7 @@ export const MessageTiming: FC<{
data-slot="message-timing-trigger"
aria-label="Message timing"
className={cn(
- "flex items-center rounded-[10px] p-1 font-mono text-chat-icon-fg text-[13px] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover",
+ "flex items-center rounded-[10px] p-1 font-mono text-chat-icon-fg text-[0.8125rem] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover",
className,
)}
>
diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx
index 9788c73605..d869cc93ef 100644
--- a/studio/frontend/src/components/assistant-ui/reasoning.tsx
+++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx
@@ -163,7 +163,7 @@ function ReasoningContent({
Generated image
{overlay.metadata ? (
-
+
{overlay.metadata}
) : null}
@@ -1390,7 +1390,7 @@ const ComposerAnimated: FC<{
disableQueue?: boolean;
}> = ({ disabled, threadId, menuSide, disableQueue }) => {
return (
-
+
= ({
) : (
-
+
@@ -3329,7 +3329,7 @@ const PromptQueueStack: FC<{ queueThreadIds: string[] }> = ({
type="button"
variant="ghost"
size="sm"
- className="h-7 w-[5.25rem] justify-center gap-1 px-0 text-sm font-normal text-muted-foreground/80 hover:text-foreground"
+ className="h-7 w-[84px] justify-center gap-1 px-0 text-sm font-normal text-muted-foreground/80 hover:text-foreground"
onClick={() => startEditing(item)}
>
@@ -3565,14 +3565,14 @@ const DiffusionCanvas: FC = () => {
canvas.total > 0 ? `step ${canvas.step + 1}/${canvas.total}` : "denoising";
return (
-
+
Denoising
block {canvas.block + 1} - {stepLabel}
-
+
{canvas.text}
@@ -3646,7 +3646,7 @@ const AssistantMessage: FC = () => {
return (
@@ -3676,7 +3676,7 @@ const AssistantMessage: FC = () => {
) : (
<>
-
+
@@ -3759,7 +3759,7 @@ const ForkCountBadge: FC = () => {
if (count <= 0) return null;
return (
@@ -4084,7 +4084,7 @@ const UserMessageAudio: FC = () => {
const UserMessage: FC = () => {
return (
@@ -4195,7 +4195,7 @@ const BranchPicker: FC = ({
= ({
-
+
/
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx
index ec24060072..4fbaa227bd 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx
@@ -48,7 +48,7 @@ export function CitationBadge({
-
+
{errorText ??
(isStaleGeneratingArtifact
? "Refresh stopped this preview"
diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx
index e38e2e5882..fb1ead3e63 100644
--- a/studio/frontend/src/components/floating-monitor.tsx
+++ b/studio/frontend/src/components/floating-monitor.tsx
@@ -102,7 +102,7 @@ export function FloatingMonitor() {
initial={{ opacity: 0, scale: 0.9 }}
animate={{ opacity: 1, scale: 1 }}
exit={{ opacity: 0, scale: 0.9 }}
- className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-2rem)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none"
+ className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-32px)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none"
>
@@ -139,7 +139,7 @@ export function FloatingMonitor() {
className="space-y-3 overflow-hidden"
>
-
+
{t("settings.resources.liveMonitor.ram")}
-
+
{t("settings.resources.liveMonitor.vram")}{" "}
{devices.length > 1
diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx
index 3db15ffe30..25c413ad61 100644
--- a/studio/frontend/src/components/llama-update-banner.tsx
+++ b/studio/frontend/src/components/llama-update-banner.tsx
@@ -131,7 +131,7 @@ export function LlamaUpdateBanner({
-
+
{sizeLabel ? `${sizeLabel} download · ` : ""}No restart needed
after update
@@ -209,7 +209,7 @@ export function LlamaUpdateBanner({
@@ -218,7 +218,7 @@ export function LlamaUpdateBanner({
diff --git a/studio/frontend/src/components/section-card.tsx b/studio/frontend/src/components/section-card.tsx
index e894749494..6539fb8547 100644
--- a/studio/frontend/src/components/section-card.tsx
+++ b/studio/frontend/src/components/section-card.tsx
@@ -77,7 +77,7 @@ export function SectionCard({
{title}
{badge && (
-
+
{badge}
)}
diff --git a/studio/frontend/src/components/tauri/startup-screen.tsx b/studio/frontend/src/components/tauri/startup-screen.tsx
index 678051b36b..318a538ad0 100644
--- a/studio/frontend/src/components/tauri/startup-screen.tsx
+++ b/studio/frontend/src/components/tauri/startup-screen.tsx
@@ -72,7 +72,7 @@ function DiagnosticsCopyActions({
readOnly
value={manualReport}
onFocus={(event) => event.currentTarget.select()}
- className="h-32 w-full max-w-md resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[10px] text-muted-foreground"
+ className="h-32 w-full max-w-md resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[0.625rem] text-muted-foreground"
/>
)}
diff --git a/studio/frontend/src/components/tauri/update-banner.tsx b/studio/frontend/src/components/tauri/update-banner.tsx
index 4fe8a54377..846701ca64 100644
--- a/studio/frontend/src/components/tauri/update-banner.tsx
+++ b/studio/frontend/src/components/tauri/update-banner.tsx
@@ -95,7 +95,7 @@ export function UpdateBanner({
transition={{ duration: 0.35, ease: EASE_OUT_QUART }}
className={cn(
positioned
- ? "fixed bottom-4 right-4 z-[9999] w-[calc(100vw-2rem)] max-w-[400px]"
+ ? "fixed bottom-4 right-4 z-[9999] w-[calc(100vw-32px)] max-w-[400px]"
: "pointer-events-auto w-full",
)}
data-testid="tauri-update-banner"
@@ -142,7 +142,7 @@ export function UpdateBanner({
)}
-
+
{showFailure
? "Backend recovered. Diagnostics are still available."
: isManualLinuxPackage
@@ -166,7 +166,7 @@ export function UpdateBanner({
{
handleCopyDiagnostics().catch(console.error);
}}
@@ -176,14 +176,14 @@ export function UpdateBanner({
Later
@@ -195,14 +195,14 @@ export function UpdateBanner({
Remind me later
@@ -219,7 +219,7 @@ export function UpdateBanner({
readOnly={true}
value={manualReport}
onFocus={(event) => event.currentTarget.select()}
- className="mt-2 h-28 w-full resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[10px] text-muted-foreground"
+ className="mt-2 h-28 w-full resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[0.625rem] text-muted-foreground"
/>
)}
diff --git a/studio/frontend/src/components/tauri/update-screen.tsx b/studio/frontend/src/components/tauri/update-screen.tsx
index 3199425f69..64f2e95a87 100644
--- a/studio/frontend/src/components/tauri/update-screen.tsx
+++ b/studio/frontend/src/components/tauri/update-screen.tsx
@@ -72,7 +72,7 @@ function LogViewer({ logs }: { logs: string[] }) {
return (
{logs.map((line, i) => (
@@ -197,7 +197,7 @@ export function UpdateScreen({
readOnly
value={manualReport}
onFocus={(event) => event.currentTarget.select()}
- className="mt-2 h-32 w-full max-w-xl resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[10px] text-muted-foreground"
+ className="mt-2 h-32 w-full max-w-xl resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[0.625rem] text-muted-foreground"
/>
)}
diff --git a/studio/frontend/src/components/tauri/window-titlebar.tsx b/studio/frontend/src/components/tauri/window-titlebar.tsx
index d5c74df463..66cc2e1b41 100644
--- a/studio/frontend/src/components/tauri/window-titlebar.tsx
+++ b/studio/frontend/src/components/tauri/window-titlebar.tsx
@@ -112,8 +112,8 @@ export function WindowTitlebar({
const { pinned, togglePinned } = useSidebarPin();
const sidebarWidth = showSidebarSurface
? pinned
- ? "var(--studio-sidebar-expanded-width,17.5rem)"
- : "var(--studio-sidebar-collapsed-width,3rem)"
+ ? "var(--studio-sidebar-expanded-width,280px)"
+ : "var(--studio-sidebar-collapsed-width,48px)"
: "0px";
const contentBorderLeft = pinned ? `calc(${sidebarWidth} + 12px)` : "0px";
@@ -273,7 +273,7 @@ export function WindowTitlebar({
draggable={false}
className="size-5 shrink-0 rounded-[6px] object-cover"
/>
-
+
Unsloth Studio
@@ -325,7 +325,7 @@ export function WindowTitlebar({
className="pointer-events-auto absolute top-0 h-full"
style={{
left: sidebarWidth,
- right: "calc(var(--studio-window-control-inset,112px) + 0.5rem)",
+ right: "calc(var(--studio-window-control-inset,112px) + 8px)",
}}
onMouseDown={handleDragMouseDown}
onDoubleClick={handleDragDoubleClick}
diff --git a/studio/frontend/src/components/ui/chart.tsx b/studio/frontend/src/components/ui/chart.tsx
index 32da410bb5..25dc88d1cd 100644
--- a/studio/frontend/src/components/ui/chart.tsx
+++ b/studio/frontend/src/components/ui/chart.tsx
@@ -246,7 +246,7 @@ function ChartTooltipContent({
return (
diff --git a/studio/frontend/src/components/ui/copyable-error-chip.tsx b/studio/frontend/src/components/ui/copyable-error-chip.tsx
index 595df5cf62..6f21b6d829 100644
--- a/studio/frontend/src/components/ui/copyable-error-chip.tsx
+++ b/studio/frontend/src/components/ui/copyable-error-chip.tsx
@@ -53,7 +53,7 @@ export function CopyableErrorChip({
@@ -63,7 +63,7 @@ export function CopyableErrorChip({
Error
@@ -72,7 +72,7 @@ export function CopyableErrorChip({
onClick={handleCopy}
aria-label={copied ? "Copied" : "Copy error message"}
className={cn(
- "inline-flex items-center gap-1 rounded-md border border-border/60 px-2 py-1 text-[11px] text-muted-foreground transition-colors hover:bg-muted/60 focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring",
+ "inline-flex items-center gap-1 rounded-md border border-border/60 px-2 py-1 text-[0.6875rem] text-muted-foreground transition-colors hover:bg-muted/60 focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring",
copied && "border-emerald-500/40 text-emerald-600 dark:text-emerald-500",
)}
>
diff --git a/studio/frontend/src/components/ui/data-table.tsx b/studio/frontend/src/components/ui/data-table.tsx
index 31a93b7449..2391eafa5b 100644
--- a/studio/frontend/src/components/ui/data-table.tsx
+++ b/studio/frontend/src/components/ui/data-table.tsx
@@ -100,7 +100,7 @@ export function DataTable
({
{row.getVisibleCells().map((cell) => (
{flexRender(cell.column.columnDef.cell, cell.getContext())}
diff --git a/studio/frontend/src/components/ui/dialog.tsx b/studio/frontend/src/components/ui/dialog.tsx
index 6dea1880d9..7b61f02d1f 100644
--- a/studio/frontend/src/components/ui/dialog.tsx
+++ b/studio/frontend/src/components/ui/dialog.tsx
@@ -89,7 +89,7 @@ function DialogContent({
) {
return (
@@ -51,11 +52,21 @@ function DropdownMenuContent({
// The 3px alignment nudge must be margin, not translate: a transform
// here makes this scroll container the containing block for nested
// position:fixed submenu wrappers, clipping every submenu.
- "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 bg-popover text-popover-foreground min-w-48 rounded-lg p-1 duration-100 z-50 max-h-(--radix-dropdown-menu-content-available-height) w-[calc(var(--radix-dropdown-menu-trigger-width)_+_6px)] data-[align=start]:-ml-[3px] data-[align=end]:ml-[3px] origin-(--radix-dropdown-menu-content-transform-origin) overflow-x-hidden overflow-y-auto data-[state=closed]:overflow-hidden",
+ "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 bg-popover text-popover-foreground min-w-48 rounded-lg p-1 duration-100 z-50 max-h-(--radix-dropdown-menu-content-available-height) w-[calc(var(--radix-dropdown-menu-trigger-width)_+_6px)] data-[align=start]:-ml-[3px] data-[align=end]:ml-[3px] origin-(--radix-dropdown-menu-content-transform-origin) flex flex-col overflow-hidden",
className,
)}
{...props}
- />
+ >
+ {/* Scroll an inner viewport, not the rounded surface: a scrollbar on
+ the surface squares its corners in WebKit. The surface padding
+ insets the scrollbar clear of the curve. */}
+
+ {children}
+
+
);
}
@@ -307,7 +318,7 @@ function DropdownMenuSubContent({
isMobile && contentWidth === 0 ? "hidden" : style?.visibility,
}}
className={cn(
- "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 bg-popover text-popover-foreground min-w-36 max-w-[calc(100vw-2rem)] rounded-lg p-1 duration-100 z-50 origin-(--radix-dropdown-menu-content-transform-origin) overflow-hidden",
+ "data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 bg-popover text-popover-foreground min-w-36 max-w-[calc(100vw-32px)] rounded-lg p-1 duration-100 z-50 origin-(--radix-dropdown-menu-content-transform-origin) overflow-hidden",
className,
)}
{...props}
diff --git a/studio/frontend/src/components/ui/input-group.tsx b/studio/frontend/src/components/ui/input-group.tsx
index 7ec7a2e3d4..56b9bcf63d 100644
--- a/studio/frontend/src/components/ui/input-group.tsx
+++ b/studio/frontend/src/components/ui/input-group.tsx
@@ -29,9 +29,9 @@ const inputGroupAddonVariants = cva(
variants: {
align: {
"inline-start":
- "pl-3 has-[>button]:ml-[-0.25rem] has-[>kbd]:ml-[-0.15rem] order-first",
+ "pl-3 has-[>button]:ml-[-4px] has-[>kbd]:ml-[-2.4px] order-first",
"inline-end":
- "pr-3 has-[>button]:mr-[-0.25rem] has-[>kbd]:mr-[-0.15rem] order-last",
+ "pr-3 has-[>button]:mr-[-4px] has-[>kbd]:mr-[-2.4px] order-last",
"block-start":
"px-3 pt-3 group-has-[>input]/input-group:pt-3 [.border-b]:pb-3 order-first w-full justify-start",
"block-end":
diff --git a/studio/frontend/src/components/ui/select.tsx b/studio/frontend/src/components/ui/select.tsx
index 925cc45b36..1bb95848f7 100644
--- a/studio/frontend/src/components/ui/select.tsx
+++ b/studio/frontend/src/components/ui/select.tsx
@@ -126,7 +126,7 @@ function SelectContent({
data-slot="select-content"
data-align-trigger={position === "item-aligned"}
className={cn(
- "bg-popover text-popover-foreground font-heading data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 min-w-36 rounded-xl p-1 corner-squircle duration-100 relative z-50 max-h-(--radix-select-content-available-height) origin-(--radix-select-content-transform-origin) overflow-x-hidden overflow-y-auto ",
+ "bg-popover text-popover-foreground font-heading data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 min-w-36 rounded-xl p-1 corner-squircle duration-100 relative z-50 max-h-(--radix-select-content-available-height) origin-(--radix-select-content-transform-origin) flex flex-col overflow-hidden",
// No popper translate offset: the menu sits flush against the trigger.
className,
)}
@@ -135,10 +135,18 @@ function SelectContent({
{...props}
>
+ {/* The Radix viewport is the scroller (not the rounded surface, whose
+ scrollbar would square its corners in WebKit; not a wrapper, which
+ would blind Radix's scroll handling). The surface padding insets
+ the scrollbar clear of the curve. */}
diff --git a/studio/frontend/src/components/ui/sidebar.tsx b/studio/frontend/src/components/ui/sidebar.tsx
index 4eb37d4e89..d424b7254c 100644
--- a/studio/frontend/src/components/ui/sidebar.tsx
+++ b/studio/frontend/src/components/ui/sidebar.tsx
@@ -30,8 +30,8 @@ import { LayoutAlignLeftIcon } from "@hugeicons/core-free-icons"
const noop = () => {}
-const SIDEBAR_WIDTH = "17.5rem"
-const SIDEBAR_WIDTH_ICON = "3rem"
+const SIDEBAR_WIDTH = "280px"
+const SIDEBAR_WIDTH_ICON = "48px"
const SIDEBAR_KEYBOARD_SHORTCUT = "b"
type SidebarContextProps = {
@@ -228,7 +228,7 @@ function Sidebar({
data-sidebar="sidebar"
data-slot="sidebar"
data-mobile="true"
- className="bg-sidebar text-sidebar-foreground w-2/3 max-w-[18rem] p-0 [&>button]:hidden"
+ className="bg-sidebar text-sidebar-foreground w-2/3 max-w-[288px] p-0 [&>button]:hidden"
side={side}
>
@@ -471,7 +471,7 @@ function SidebarGroupLabel({
data-slot="sidebar-group-label"
data-sidebar="group-label"
className={cn(
- "text-[#94a3b8] dark:text-[#666] ring-sidebar-ring h-auto pt-3 pb-2 px-4 rounded-md text-[10px] font-semibold uppercase tracking-[0em] group-data-[collapsible=icon]:-mt-8 group-data-[collapsible=icon]:opacity-0 focus-visible:ring-1 [&>svg]:size-3 flex shrink-0 items-center outline-hidden [&>svg]:shrink-0",
+ "text-[#94a3b8] dark:text-[#666] ring-sidebar-ring h-auto pt-3 pb-2 px-4 rounded-md text-[0.625rem] font-semibold uppercase tracking-[0em] group-data-[collapsible=icon]:-mt-8 group-data-[collapsible=icon]:opacity-0 focus-visible:ring-1 [&>svg]:size-3 flex shrink-0 items-center outline-hidden [&>svg]:shrink-0",
className
)}
{...props}
diff --git a/studio/frontend/src/components/web/update-banner.tsx b/studio/frontend/src/components/web/update-banner.tsx
index 55ef06166b..dfb181fdf4 100644
--- a/studio/frontend/src/components/web/update-banner.tsx
+++ b/studio/frontend/src/components/web/update-banner.tsx
@@ -79,7 +79,7 @@ export function WebUpdateBanner({
transition={{ duration: 0.35, ease: EASE_OUT_QUART }}
className={cn(
positioned
- ? "fixed bottom-4 right-4 z-[9999] w-[calc(100vw-2rem)] max-w-[400px]"
+ ? "fixed bottom-4 right-4 z-[9999] w-[calc(100vw-32px)] max-w-[400px]"
: "pointer-events-auto w-full",
)}
data-testid="web-update-banner"
@@ -132,7 +132,7 @@ export function WebUpdateBanner({
href={RELEASE_NOTES_URL}
target="_blank"
rel="noopener noreferrer"
- className="-ml-2 whitespace-nowrap rounded-full px-2.5 py-2 text-[13px] font-medium text-foreground transition-colors hover:bg-muted"
+ className="-ml-2 whitespace-nowrap rounded-full px-2.5 py-2 text-[0.8125rem] font-medium text-foreground transition-colors hover:bg-muted"
data-testid="web-update-release-notes-link"
>
Release notes
@@ -142,7 +142,7 @@ export function WebUpdateBanner({
@@ -151,7 +151,7 @@ export function WebUpdateBanner({
diff --git a/studio/frontend/src/features/auth/login-page.tsx b/studio/frontend/src/features/auth/login-page.tsx
index d967328c7f..34feccce79 100644
--- a/studio/frontend/src/features/auth/login-page.tsx
+++ b/studio/frontend/src/features/auth/login-page.tsx
@@ -16,7 +16,7 @@ export function LoginPage() {
length="70vh"
className="opacity-35 dark:opacity-15"
/>
-
+
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
index 0345dc6e2a..82a236a387 100644
--- a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
+++ b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
@@ -145,12 +145,12 @@ export function ArtifactCard({
{isCode ? "HTML Code" : artifact.title}
-
+
HTML canvas
{isStreaming && !isCode ? (
-
+
Generating
) : null}
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index e59ce3a805..3daae8c50d 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -576,7 +576,7 @@ function CompareShell({
{children}
-
{composer}
+
{composer}
{showModelDisclaimer && (
LLMs can make mistakes. Double-check responses.
@@ -651,7 +651,7 @@ const LoraCompareContent = memo(function LoraCompareContent({
handleName="base"
header={
-
+
Base Model
@@ -665,8 +665,8 @@ const LoraCompareContent = memo(function LoraCompareContent({
handleName="lora"
borderClassName="border-t border-border/60 md:border-t-0 md:border-l"
header={
-
-
+
+
Fine-tuned
@@ -721,8 +721,8 @@ function GeneralCompareHeader({
side === "left"
? pinned
? "pl-12 pr-3 md:pl-2"
- : "pl-12 pr-3 md:pl-[calc(0.5rem+max(0px,var(--studio-mac-traffic-light-inset,0px)-var(--sidebar-width-icon,3rem)))]"
- : "pl-3 pr-[calc(3rem+var(--studio-chat-header-right-inset,var(--studio-window-control-inset,0px)))]",
+ : "pl-12 pr-3 md:pl-[calc(8px+max(0px,var(--studio-mac-traffic-light-inset,0px)-var(--sidebar-width-icon,48px)))]"
+ : "pl-3 pr-[calc(48px+var(--studio-chat-header-right-inset,var(--studio-window-control-inset,0px)))]",
)}
>
{/* Slightly narrower than the composer max; every block shares this. */}
-
+
-
+
{projectName}
@@ -1349,7 +1349,7 @@ function ProjectLanding({
type="button"
onClick={() => setProjectTab("chats")}
data-active={projectTab === "chats"}
- className="h-10 rounded-full px-5 text-[14px] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover"
+ className="h-10 rounded-full px-5 text-[0.875rem] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover"
>
Chats
@@ -1357,7 +1357,7 @@ function ProjectLanding({
type="button"
onClick={() => setProjectTab("sources")}
data-active={projectTab === "sources"}
- className="h-10 rounded-full px-5 text-[14px] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover"
+ className="h-10 rounded-full px-5 text-[0.875rem] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover"
>
Sources
@@ -1417,7 +1417,7 @@ function ProjectLanding({
onFocus={(event) => event.currentTarget.select()}
maxLength={120}
aria-label="Rename chat"
- className="w-full border-0 bg-transparent text-[15px] font-semibold leading-5 text-foreground outline-none"
+ className="w-full border-0 bg-transparent text-[0.9375rem] font-semibold leading-5 text-foreground outline-none"
/>
@@ -1442,11 +1442,11 @@ function ProjectLanding({
className="flex min-h-[58px] min-w-0 flex-1 items-center gap-4 rounded-full px-4 py-2 text-left"
>
-
-
+
{preview?.date ??
formatProjectChatDate(item.createdAt)}
@@ -3105,14 +3105,14 @@ export function ChatPage({
)}
@@ -3141,7 +3141,7 @@ export function ChatPage({
/>
)}
{incognito && view.mode === "single" && (
-
+
When off, all connections are disabled.
@@ -1616,7 +1616,7 @@ export function ChatProvidersSettings({
{provider.name}
-
+
{provider.models.length}{" "}
{provider.models.length === 1 ? "model" : "models"}
@@ -1631,7 +1631,7 @@ export function ChatProvidersSettings({
) : null}
{modelSummary}
@@ -1702,7 +1702,7 @@ export function ChatProvidersDialog({
Connections
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index d4f154882c..99d697f619 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -141,7 +141,7 @@ export function ParamSlider({
-
+
{label}
{info && {info} }
@@ -249,7 +249,7 @@ function CollapsibleSection({
};
const headerClasses = cn(
- "flex w-full items-center justify-between text-[12px] font-medium normal-case tracking-[0.04em] text-nav-fg-muted transition-colors focus-visible:outline-none focus-visible:ring-0",
+ "flex w-full items-center justify-between text-[0.75rem] font-medium normal-case tracking-[0.04em] text-nav-fg-muted transition-colors focus-visible:outline-none focus-visible:ring-0",
first ? "pt-4 pb-5" : "py-5",
);
@@ -695,12 +695,12 @@ export function ChatSettingsPanel({
{/* Header is outside the scroll area so the scrollbar never shifts the close button. */}
{isMobile ? (
-
+
Run settings
) : (
<>
-
+
Run settings
@@ -740,7 +740,7 @@ export function ChatSettingsPanel({
{modelConfig}
{showSpecFallback && (
-
+
{specFallbackReason === "mla_mtp_disabled"
? "MTP is disabled by default for this model architecture because it currently runs slower than standard decoding. Choose MTP in the model picker to force it."
@@ -757,7 +757,7 @@ export function ChatSettingsPanel({
{mtpUpdatable && llamaUpdateStatus?.update_available && (
)}
{showContextVramWarning && (
-
+
Context length exceeds the estimated VRAM capacity (
{ggufMaxContextLength?.toLocaleString()} tokens). The
model may use system RAM.
@@ -812,7 +812,7 @@ export function ChatSettingsPanel({
maxLength={80}
autoComplete="off"
className={cn(
- "!h-9 min-h-0 min-w-0 self-stretch !pl-3.5 !pr-2 py-0 text-[13px] font-medium leading-9 text-nav-fg md:text-[13px]",
+ "!h-9 min-h-0 min-w-0 self-stretch !pl-3.5 !pr-2 py-0 text-[0.8125rem] font-medium leading-9 text-nav-fg md:text-[0.8125rem]",
presetSaveState.isSaveReady &&
"placeholder:text-primary/50",
)}
@@ -852,7 +852,7 @@ export function ChatSettingsPanel({
}
applyPreset(p.name);
}}
- className="flex min-h-9 items-center px-3 py-0 text-[13px] font-medium leading-[1.4] tracking-nav"
+ className="flex min-h-9 items-center px-3 py-0 text-[0.8125rem] font-medium leading-[1.4] tracking-nav"
>
{p.name}
@@ -874,7 +874,7 @@ export function ChatSettingsPanel({
}
size="sm"
className={cn(
- "h-9 w-full rounded-full text-[13px] font-medium tracking-nav",
+ "h-9 w-full rounded-full text-[0.8125rem] font-medium tracking-nav",
presetSaveState.isSaveReady &&
"bg-primary text-primary-foreground hover:bg-primary/90",
)}
@@ -889,7 +889,7 @@ export function ChatSettingsPanel({
disabled={!(settingsHydrated && activeCustomPreset)}
variant="outline"
size="sm"
- className="h-9 w-full rounded-full text-[13px] font-medium tracking-nav text-muted-foreground"
+ className="h-9 w-full rounded-full text-[0.8125rem] font-medium tracking-nav text-muted-foreground"
title={
activeCustomPreset
? activeBuiltinPreset
@@ -908,7 +908,7 @@ export function ChatSettingsPanel({
-
+
Prompt caching
@@ -931,7 +931,7 @@ export function ChatSettingsPanel({
{showPromptCacheTtlControl && promptCachingEnabled ? (
-
+
Cache TTL
@@ -968,7 +968,7 @@ export function ChatSettingsPanel({
{showFastModeControl ? (
-
+
Fast mode
@@ -1056,7 +1056,7 @@ export function ChatSettingsPanel({
placeholder="Example: You are a helpful assistant..."
aria-label="System prompt"
className={cn(
- "block size-full resize-none bg-transparent px-3.5 py-2.5 text-left text-[13px] font-medium leading-relaxed text-nav-fg outline-none placeholder:text-muted-foreground",
+ "block size-full resize-none bg-transparent px-3.5 py-2.5 text-left text-[0.8125rem] font-medium leading-relaxed text-nav-fg outline-none placeholder:text-muted-foreground",
systemPromptOverflows && "cursor-pointer",
)}
/>
@@ -1202,13 +1202,13 @@ export function ChatSettingsPanel({
-
Prompt editor
+
Prompt editor
setSystemVariablesOpen((open) => !open)}
- className="h-7 gap-1.5 rounded-full px-2.5 text-[11px] text-muted-foreground"
+ className="h-7 gap-1.5 rounded-full px-2.5 text-[0.6875rem] text-muted-foreground"
aria-expanded={systemVariablesOpen}
>
@@ -1221,7 +1221,7 @@ export function ChatSettingsPanel({
/>
-
+
Use this for longer edits. Save writes back to the active
configuration only. Insert variables with {"{{ env }}"}.
@@ -1230,16 +1230,16 @@ export function ChatSettingsPanel({
-
+
Prompt variables
-
+
Define values as JSON below, then use each key in your
prompt, like {"{{ env }}"}.
-
+
Built-in, fill in automatically
@@ -1247,7 +1247,7 @@ export function ChatSettingsPanel({
{token}
@@ -1272,11 +1272,11 @@ export function ChatSettingsPanel({
aria-invalid={Boolean(systemVariablesError)}
/>
{systemVariablesError ? (
-
+
{systemVariablesError}
) : (
-
+
Names you don't define are left unchanged, so a stray
{" {{ typo }} "}stays visible in the prompt.
@@ -1288,7 +1288,7 @@ export function ChatSettingsPanel({
onChange={(event) => setSystemPromptDraft(event.target.value)}
placeholder="You are a helpful assistant..."
fieldSizing="fixed"
- className="min-h-[20rem] max-h-[48vh] overflow-y-auto border-0 text-sm leading-6 corner-squircle focus-visible:ring-0"
+ className="min-h-[320px] max-h-[48vh] overflow-y-auto border-0 text-sm leading-6 corner-squircle focus-visible:ring-0"
rows={14}
/>
@@ -1333,7 +1333,7 @@ export function ChatSettingsPanel({
if (isMobile) {
return (
-
+
Run settings
Chat inference settings
@@ -1351,7 +1351,7 @@ export function ChatSettingsPanel({
data-tour="chat-settings"
className={cn(
"relative z-50 shrink-0 overflow-hidden bg-panel-surface text-panel-surface-fg font-heading",
- open ? "w-[17rem] border-l border-sidebar-border" : "w-0",
+ open ? "w-[272px] border-l border-sidebar-border" : "w-0",
)}
style={{
height: "calc(100% - var(--studio-custom-titlebar-height, 0px))",
@@ -1426,7 +1426,7 @@ function AutoHealToolCallsToggle() {
return (
-
+
Auto-Healing Tool Calls
@@ -1450,7 +1450,7 @@ function NudgeToolCallsToggle() {
return (
-
+
Nudge Tool Calls
@@ -1475,7 +1475,7 @@ function ConfirmToolCallsToggle() {
-
+
Confirm tool calls
@@ -1487,7 +1487,7 @@ function ConfirmToolCallsToggle() {
{permissionMode === "full" ? (
-
+
Overridden by Full access
) : null}
@@ -1508,7 +1508,7 @@ function BypassPermissionsToggle() {
return (
-
+
Tool permissions
@@ -1517,9 +1517,9 @@ function BypassPermissionsToggle() {
{/* Full width, styled like the panel selects/preset input. */}
-
+
{permissionMode === "full" ? (
-
+
Tool calls run with no confirmation and no sandbox.
) : null}
diff --git a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
index dc3e1aac29..ea95040f44 100644
--- a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
+++ b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
@@ -83,7 +83,7 @@ export function ChatSearchDialog() {
@@ -143,10 +143,10 @@ export function ChatSearchDialog() {
strokeWidth={2}
className="size-4 shrink-0 text-muted-foreground"
/>
-
+
{item.title || "Untitled chat"}
-
+
{formatRelative(item.createdAt)}
diff --git a/studio/frontend/src/features/chat/components/context-usage-bar.tsx b/studio/frontend/src/features/chat/components/context-usage-bar.tsx
index 80f502e222..eeacef66df 100644
--- a/studio/frontend/src/features/chat/components/context-usage-bar.tsx
+++ b/studio/frontend/src/features/chat/components/context-usage-bar.tsx
@@ -71,7 +71,7 @@ export const ContextUsageBar: FC<{
: `Token usage: ${formatTokenCount(used)} tokens`
}
className={cn(
- "flex items-center gap-2 rounded-[10px] px-2.5 py-1 font-mono text-chat-icon-fg text-[13px] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover",
+ "flex items-center gap-2 rounded-[10px] px-2.5 py-1 font-mono text-chat-icon-fg text-[0.8125rem] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover",
className,
)}
>
@@ -149,7 +149,7 @@ export const ContextUsageBar: FC<{
{hasKnownLimit && percent !== null && percent > 85 ? (
-
+
Close to the context limit. Generation will stop at 100%.
Increase
Context Length in
the chat Settings panel to keep going.
diff --git a/studio/frontend/src/features/chat/components/model-load-status.tsx b/studio/frontend/src/features/chat/components/model-load-status.tsx
index 292c7884fd..613b5c260b 100644
--- a/studio/frontend/src/features/chat/components/model-load-status.tsx
+++ b/studio/frontend/src/features/chat/components/model-load-status.tsx
@@ -54,14 +54,14 @@ export function ModelLoadDescription({
{title ?
{title}
: null}
{hasProgress ? (
-
+
{labelPrimary}
{Math.round(clampProgress(progressPercent))}%
{labelSecondary ? (
-
+
{labelSecondary}
) : null}
@@ -96,18 +96,18 @@ export function ModelLoadInlineStatus({
const hasProgress = typeof progressPercent === "number";
return (
-
+
{label}
{hasProgress ? (
-
+
{/* Tight inline layout: show only the primary (bytes) chunk;
@@ -124,7 +124,7 @@ export function ModelLoadInlineStatus({
type="button"
size="xs"
variant="outline"
- className="shrink-0 text-[11px]"
+ className="shrink-0 text-[0.6875rem]"
onClick={onStop}
>
Stop
diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx
index cb0234579b..04c88e4eba 100644
--- a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx
+++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx
@@ -435,7 +435,7 @@ export function OpenAICodeExecSection({
Idle timeout
@@ -459,7 +459,7 @@ export function OpenAICodeExecSection({
ACTIVE pill marks which one (no separate picker). */}
-
+
Containers
{isPending ? (
-
+
Creating
) : isActive ? (
-
+
Active
) : statusLabel ? (
-
+
{statusLabel}
) : null}
@@ -548,10 +548,10 @@ export function OpenAICodeExecSection({
className="flex min-w-0 items-center gap-1.5 text-muted-foreground"
title={c.id}
>
-
+
{shortContainerId(c.id)}
-
+
· {ttlMinutes}m
diff --git a/studio/frontend/src/features/chat/components/project-switcher.tsx b/studio/frontend/src/features/chat/components/project-switcher.tsx
index 8f923a8c80..2a170e5a39 100644
--- a/studio/frontend/src/features/chat/components/project-switcher.tsx
+++ b/studio/frontend/src/features/chat/components/project-switcher.tsx
@@ -57,7 +57,7 @@ export function ProjectSwitcher({
className="size-icon shrink-0 text-foreground/70"
/>
-
+
{label}
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
index b72a7a95c2..4b3f57b368 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
@@ -111,7 +111,7 @@ const MODEL_LOAD_TOAST_CLASSNAMES = {
title: "leading-5",
description: "mt-0 w-full",
cancelButton:
- "!h-auto !rounded-none !border-0 !bg-transparent !px-1 !text-[11px] !font-normal !text-muted-foreground hover:!bg-transparent hover:!text-destructive focus-visible:!text-destructive",
+ "!h-auto !rounded-none !border-0 !bg-transparent !px-1 !text-[0.6875rem] !font-normal !text-muted-foreground hover:!bg-transparent hover:!text-destructive focus-visible:!text-destructive",
} as const;
const MODEL_LOADED_TOAST_CLASSNAMES = {
diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx
index e6c89cf54a..7e0ecb0c7e 100644
--- a/studio/frontend/src/features/chat/permission-mode-select.tsx
+++ b/studio/frontend/src/features/chat/permission-mode-select.tsx
@@ -120,7 +120,7 @@ export function PermissionModeMenuItems({
>
- {option.label}
+ {option.label}
{option.description}
diff --git a/studio/frontend/src/features/chat/projects-page.tsx b/studio/frontend/src/features/chat/projects-page.tsx
index c9960ffaca..494368faec 100644
--- a/studio/frontend/src/features/chat/projects-page.tsx
+++ b/studio/frontend/src/features/chat/projects-page.tsx
@@ -367,7 +367,7 @@ export function ProjectsPage() {
}}
/>
-
+
Projects
@@ -419,7 +419,7 @@ export function ProjectsPage() {
Export All Projects
-
+
Combined
{EXPORT_FORMATS_LIST.map(({ fmt, label }) => (
@@ -430,7 +430,7 @@ export function ProjectsPage() {
-
+
Per chat
{EXPORT_FORMATS_LIST.map(({ fmt, label }) => (
@@ -445,7 +445,7 @@ export function ProjectsPage() {
Export Projects + Recents
-
+
Combined
{EXPORT_FORMATS_LIST.map(({ fmt, label }) => (
@@ -456,7 +456,7 @@ export function ProjectsPage() {
-
+
Per chat
{EXPORT_FORMATS_LIST.map(({ fmt, label }) => (
@@ -482,7 +482,7 @@ export function ProjectsPage() {
{!hasLoaded ? (
-
+
Name
Modified
@@ -526,7 +526,7 @@ export function ProjectsPage() {
{/* Column header. Name starts at the folder icon's left edge; the
right-anchored columns keep Modified over its values. */}
-
+
Name
Modified
@@ -571,7 +571,7 @@ export function ProjectsPage() {
className="size-5"
/>
-
+
{project.name}
diff --git a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
index 09c4944a14..4b815a7695 100644
--- a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
+++ b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
@@ -1341,7 +1341,7 @@ function ExportModal({
{/* */}
-
+
Export as
@@ -1390,7 +1390,7 @@ function ExportModal({
ShareGPT format for Unsloth fine-tuning
-
+
{`{"conversations":[{"from":"human","value":"..."},{"from":"gpt","value":""}]}`}
@@ -1400,7 +1400,7 @@ function ExportModal({
{/* */}
-
+
Format
@@ -1730,7 +1730,7 @@ function PromptListCard({
{entry.name}
-
+
{entry.items.length}
@@ -1779,7 +1779,7 @@ function PromptListCard({
))}
{entry.items.length > 3 && (
-
+
+{entry.items.length - 3} more
)}
diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx
index f85c74eb86..4e2765bebd 100644
--- a/studio/frontend/src/features/chat/thread-sidebar.tsx
+++ b/studio/frontend/src/features/chat/thread-sidebar.tsx
@@ -266,7 +266,7 @@ export function ThreadSidebar({
>
{item.isFork ? (
fork
diff --git a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx
index 088d016894..27584a646f 100644
--- a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx
+++ b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx
@@ -280,7 +280,7 @@ function LearningRecipeCards({
{badge}
@@ -288,7 +288,7 @@ function LearningRecipeCards({
{extraLearningBadgeCount > 0 ? (
+{extraLearningBadgeCount}
@@ -296,7 +296,7 @@ function LearningRecipeCards({
{isReady ? null : (
Soon
@@ -403,7 +403,7 @@ export function DataRecipesPage(): ReactElement {
-
+
Data Recipes
diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx
index 6c2794420b..86c82935d6 100644
--- a/studio/frontend/src/features/export/components/export-run-panel.tsx
+++ b/studio/frontend/src/features/export/components/export-run-panel.tsx
@@ -278,7 +278,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
onSaveDirectoryChange(e.target.value)}
spellCheck={false}
@@ -303,7 +303,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
Browse
-
+
{saveDirectory !== defaultSaveDirectory ? (
<>Default: {defaultSaveDirectory}>
) : (
@@ -350,7 +350,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
href="https://huggingface.co/settings/tokens"
target="_blank"
rel="noopener noreferrer"
- className="flex items-center gap-1 text-[11px] text-emerald-600 hover:text-emerald-700 transition-colors"
+ className="flex items-center gap-1 text-[0.6875rem] text-emerald-600 hover:text-emerald-700 transition-colors"
>
Get token
@@ -369,7 +369,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
onChange={(e) => onHfTokenChange(e.target.value)}
/>
-
+
Leave empty if already logged in via CLI.
@@ -427,7 +427,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
) : null}
{o.path}
@@ -503,11 +503,11 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
{showProgress && (
-
+
{PHASE_LABELS[run.phase] ?? run.phase}
{summaryMethod === "gguf" && run.quantTotal > 1 && (
-
+
Quant{" "}
{Math.min(
run.quantIndex + (isExporting ? 1 : 0),
@@ -516,10 +516,10 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
of {run.quantTotal}
)}
-
+
{progress}%
-
+
{formatElapsed(elapsedSeconds)}
@@ -536,7 +536,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
/>
{run.stage && (
{run.stage}
@@ -552,7 +552,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
Export output
-
+
{run.logLines.length === 0 ? (
diff --git a/studio/frontend/src/features/export/components/method-picker.tsx b/studio/frontend/src/features/export/components/method-picker.tsx
index 420a7f6146..e240fd44ca 100644
--- a/studio/frontend/src/features/export/components/method-picker.tsx
+++ b/studio/frontend/src/features/export/components/method-picker.tsx
@@ -123,7 +123,7 @@ export function MethodPicker({ value, onChange, disabledMethods = [], disabledRe
{m.badge && (
{m.badge}
diff --git a/studio/frontend/src/features/export/components/quant-picker.tsx b/studio/frontend/src/features/export/components/quant-picker.tsx
index 289f6498ce..688e5fb87f 100644
--- a/studio/frontend/src/features/export/components/quant-picker.tsx
+++ b/studio/frontend/src/features/export/components/quant-picker.tsx
@@ -61,7 +61,7 @@ export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) {
-
+
— select one or more
@@ -90,10 +90,10 @@ export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) {
)}
{q.label}
{sizeLabel && (
- {sizeLabel}
+ {sizeLabel}
)}
{q.recommended && !active && (
-
+
rec
)}
@@ -103,13 +103,13 @@ export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) {
{value.length > 0 && (
-
+
{value.length} selected
onChange([])}
- className="text-[11px] text-muted-foreground/70 hover:text-foreground transition-colors"
+ className="text-[0.6875rem] text-muted-foreground/70 hover:text-foreground transition-colors"
>
Clear all
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx
index 3a970713ac..80235846d1 100644
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@@ -895,7 +895,7 @@ export function ExportPage() {
-
+
Export Model
@@ -964,21 +964,21 @@ export function ExportPage() {
Local Model
Fine-tuned
Hugging Face
@@ -1289,7 +1289,7 @@ export function ExportPage() {
{model?.display_name ?? id}
-
+
{source}
@@ -1300,15 +1300,15 @@ export function ExportPage() {
{isLoadingLocalModels ? (
-
+
Scanning local models...
) : localModelsError ? (
-
+
{localModelsError}
) : (
-
+
{exportableLocalModels.length > 0
? `${exportableLocalModels.length} local/cached models found`
: "No local models found. Enter path manually."}
@@ -1318,7 +1318,7 @@ export function ExportPage() {
)}
-
+
Direct model exports currently support GGUF only.
@@ -1327,7 +1327,7 @@ export function ExportPage() {
{sourceMode === "checkpoint" && (
-
+
Training Info
@@ -1374,7 +1374,7 @@ export function ExportPage() {
key={step}
className="flex items-start gap-2 text-xs text-muted-foreground"
>
-
+
{i + 1}
{step}
@@ -1422,7 +1422,7 @@ export function ExportPage() {
Precision
-
+
— select one or more
@@ -1479,7 +1479,7 @@ export function ExportPage() {
{f.label}
{f.needsCalibration ? " *" : ""}
-
+
{f.hint}
@@ -1492,7 +1492,7 @@ export function ExportPage() {
{selectedFormats.length > 0 && (
-
+
{selectedFormats.length} selected:{" "}
{selectedFormats
.map(
@@ -1506,7 +1506,7 @@ export function ExportPage() {
setSelectedFormats(["16-bit"])}
- className="text-[11px] text-muted-foreground/70 hover:text-foreground transition-colors"
+ className="text-[0.6875rem] text-muted-foreground/70 hover:text-foreground transition-colors"
>
Reset to 16-bit
@@ -1515,7 +1515,7 @@ export function ExportPage() {
)}
{hubMultiFormat && (
-
+
Hub export supports one format at a time (each writes to
the repository root). Select a single format, or export
locally to produce several at once.
@@ -1527,13 +1527,13 @@ export function ExportPage() {
MERGED_FORMATS.find((f) => f.value === v)
?.needsCalibration,
) && (
-
+
* calibrates on data (uses a small calibration set).
)}
{!hasNvidia && (
-
+
No NVIDIA GPU detected: compressed-tensors formats are
hidden. 16-bit and portable FP8/INT8 (torchao) still
work here and load in vLLM.
diff --git a/studio/frontend/src/features/hub/catalog/catalog-states.tsx b/studio/frontend/src/features/hub/catalog/catalog-states.tsx
index a36c4bb2da..693b5b40c0 100644
--- a/studio/frontend/src/features/hub/catalog/catalog-states.tsx
+++ b/studio/frontend/src/features/hub/catalog/catalog-states.tsx
@@ -38,20 +38,20 @@ export function NetworkErrorState({
-
+
{title}
-
+
{body}
-
{message}
+
{message}
{onSwitchDevice ? (
On Device
@@ -59,7 +59,7 @@ export function NetworkErrorState({
-
+
No matches yet
-
+
Scanned {scannedCount.toLocaleString()} results. Load another page to
keep searching Hugging Face.
@@ -105,7 +105,7 @@ export function DiscoverFetchMoreState({
Clear filters
@@ -114,7 +114,7 @@ export function DiscoverFetchMoreState({
type="button"
onClick={onFetchMore}
disabled={isLoadingMore}
- className="inline-flex h-8 items-center gap-1.5 rounded-full bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] disabled:cursor-not-allowed disabled:opacity-50 dark:hover:bg-white/[0.05]"
+ className="inline-flex h-8 items-center gap-1.5 rounded-full bg-transparent px-3 text-[0.75rem] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] disabled:cursor-not-allowed disabled:opacity-50 dark:hover:bg-white/[0.05]"
>
{/* Only warn about hidden results when a filter is actually narrowing them. */}
{hasActiveFilters && (
-
+
Some results may be hidden by your filters.
)}
@@ -149,7 +149,7 @@ export function DiscoverFetchMoreFooter({
type="button"
onClick={onFetchMore}
disabled={isLoadingMore}
- className="inline-flex h-8 items-center gap-1.5 rounded-full bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] disabled:cursor-not-allowed disabled:opacity-50 dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
+ className="inline-flex h-8 items-center gap-1.5 rounded-full bg-foreground/[0.06] px-3 text-[0.75rem] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] disabled:cursor-not-allowed disabled:opacity-50 dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
>
-
+
Couldn't load your library
-
+
Something went wrong reading your downloaded{" "}
{isDataset ? "datasets" : "models"}. Check that the backend is running
and try again.
@@ -187,7 +187,7 @@ export function InventoryErrorState({
Try again
@@ -213,10 +213,10 @@ export function EmptyState({
diff --git a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
index b821be4b0b..ade8d2de30 100644
--- a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
+++ b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
@@ -126,7 +126,7 @@ export function DatasetDownloadSection({
}
>
-
+
{isDownloaded && }
{!isDownloaded && isPartial && !downloading && (
diff --git a/studio/frontend/src/features/hub/catalog/dot-tag.tsx b/studio/frontend/src/features/hub/catalog/dot-tag.tsx
index 9452a201d4..5ae1be53d0 100644
--- a/studio/frontend/src/features/hub/catalog/dot-tag.tsx
+++ b/studio/frontend/src/features/hub/catalog/dot-tag.tsx
@@ -36,7 +36,7 @@ export function DotTag({
return (
diff --git a/studio/frontend/src/features/hub/catalog/download-card.tsx b/studio/frontend/src/features/hub/catalog/download-card.tsx
index 9b4bc5fd01..9bc64ced0e 100644
--- a/studio/frontend/src/features/hub/catalog/download-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/download-card.tsx
@@ -140,7 +140,7 @@ export function CardUpdateButton({
e.stopPropagation();
onClick();
}}
- className="inline-flex h-7 shrink-0 cursor-pointer items-center gap-1.5 rounded-full bg-amber-500/[0.07] pl-2 pr-2.5 text-[12px] font-medium text-amber-800/90 transition-colors duration-150 hover:bg-amber-500/[0.12] focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-amber-500/25 dark:bg-amber-400/[0.08] dark:text-amber-200/85 dark:hover:bg-amber-400/[0.16]"
+ className="inline-flex h-7 shrink-0 cursor-pointer items-center gap-1.5 rounded-full bg-amber-500/[0.07] pl-2 pr-2.5 text-[0.75rem] font-medium text-amber-800/90 transition-colors duration-150 hover:bg-amber-500/[0.12] focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-amber-500/25 dark:bg-amber-400/[0.08] dark:text-amber-200/85 dark:hover:bg-amber-400/[0.16]"
>
{pendingUrl && (
-
+
{hostOf(pendingUrl)}
-
+
{pendingUrl}
diff --git a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
index 0d6878b687..9345874a53 100644
--- a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
@@ -128,7 +128,7 @@ const FIT_BADGE: Record = {
/** Chip styling matching the on-device list's StatChip, no icon. */
const CHIP_BASE =
- "inline-flex h-5 shrink-0 items-center justify-center whitespace-nowrap rounded-full border px-2 text-[11.5px] font-medium tabular-nums leading-3 shadow-[inset_0_1px_0_rgba(255,255,255,0.04)]";
+ "inline-flex h-5 shrink-0 items-center justify-center whitespace-nowrap rounded-full border px-2 text-[0.71875rem] font-medium tabular-nums leading-3 shadow-[inset_0_1px_0_rgba(255,255,255,0.04)]";
const CHIP_DEFAULT =
"border-foreground/15 bg-muted text-foreground/85 dark:border-border/60 dark:bg-white/[0.04] dark:text-foreground/85";
const CHIP_ACTIVE =
@@ -184,7 +184,7 @@ function QuantBadge({
// group's `overflow-hidden` sacrifices the trailing status tags instead.
@@ -914,7 +914,7 @@ export function GgufDownloadCard({
{/* Quant label + status tags travel together as one left-aligned
group so the fit-info icon never floats orphaned from its tags;
only the chevron pins right, the standard select affordance. */}
-
+
{selected ? (
) : (
-
+
Select quantization
)}
@@ -1126,7 +1126,7 @@ export function GgufDownloadCard({
void refresh()}
- className="self-start px-1 text-[11px] text-status-warning underline-offset-2 transition-colors hover:underline"
+ className="self-start px-1 text-[0.6875rem] text-status-warning underline-offset-2 transition-colors hover:underline"
>
Couldn't refresh quantizations. Retry
diff --git a/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx b/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx
index 6cc764b876..c7f402159b 100644
--- a/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx
+++ b/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx
@@ -34,7 +34,7 @@ export function GgufDownloadStatusCard({
@@ -91,7 +91,7 @@ export function GgufDownloadingFallbackCard({
-
+
{progress.variant && }
Downloading…
diff --git a/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx
index d733049ec3..e7d176442a 100644
--- a/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx
+++ b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx
@@ -69,7 +69,7 @@ export function HubDetailView({
({
aria-label={ariaLabel}
title={title}
className={cn(
- "field-trigger hub-menu-trigger field-soft field-filter inline-flex h-9 shrink-0 cursor-pointer items-center justify-between gap-2.5 rounded-full pl-3 pr-2.5 text-[12.5px] transition-colors",
+ "field-trigger hub-menu-trigger field-soft field-filter inline-flex h-9 shrink-0 cursor-pointer items-center justify-between gap-2.5 rounded-full pl-3 pr-2.5 text-[0.78125rem] transition-colors",
"focus-visible:border-border focus-visible:ring-0 focus-visible:ring-offset-0",
className,
)}
@@ -205,7 +205,7 @@ export function HubOptionMenu({
collisionPadding={12}
onCloseAutoFocus={(event) => event.preventDefault()}
className={cn(
- "hub-menu-instant menu-soft-surface w-max min-w-[var(--radix-popover-trigger-width)] max-w-[min(var(--radix-popover-content-available-width),calc(100vw-1rem))] rounded-[14px] p-1 ring-0",
+ "hub-menu-instant menu-soft-surface w-max min-w-[var(--radix-popover-trigger-width)] max-w-[min(var(--radix-popover-content-available-width),calc(100vw-16px))] rounded-[14px] p-1 ring-0",
contentClassName,
)}
>
diff --git a/studio/frontend/src/features/hub/catalog/hub-section-row.tsx b/studio/frontend/src/features/hub/catalog/hub-section-row.tsx
index b51a6cc257..bd4d4d6bed 100644
--- a/studio/frontend/src/features/hub/catalog/hub-section-row.tsx
+++ b/studio/frontend/src/features/hub/catalog/hub-section-row.tsx
@@ -58,7 +58,7 @@ export const HubSectionRow = memo(function HubSectionRow({
type="button"
onClick={onOpenList}
aria-label={`See all ${title}`}
- className="hub-section-title group/section -mx-1 inline-flex cursor-pointer items-center gap-1.5 rounded-md px-1 text-[18px] font-semibold tracking-[-0.02em] text-foreground outline-none focus-visible:ring-1 focus-visible:ring-ring"
+ className="hub-section-title group/section -mx-1 inline-flex cursor-pointer items-center gap-1.5 rounded-md px-1 text-[1.125rem] font-semibold tracking-[-0.02em] text-foreground outline-none focus-visible:ring-1 focus-visible:ring-ring"
>
{title}