@@ -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 021/213] 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 022/213] 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 023/213] 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 024/213] 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 025/213] 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 026/213] 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 027/213] 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 028/213] 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 029/213] 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 030/213] 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 031/213] 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 032/213] 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 033/213] 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 034/213] 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 035/213] 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 036/213] 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 037/213] 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/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 038/213] 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 */}