From 8517721adb692746cf09d717465cfddd2fe853aa Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Tue, 21 Jul 2026 23:25:56 -0700 Subject: [PATCH 001/194] Studio: move sidebar search into the header (#7304) * Studio: move sidebar search into the header Put the search action as an icon button next to the sidebar toggle in the header instead of a full-width nav row, so New Chat is the only fixed row above the scrolling list. The search row is kept for the collapsed icon rail only. Also add a small bottom gap under New Chat when it is pinned during scroll. * Studio: keep search row on custom-titlebar platforms The header search button only renders on mac/web where the brand row shows. On win/linux custom titlebars there's no header button, so keep the full-width search row visible instead of hiding it. * Studio: address review on sidebar search tooltip - Hide the search tooltip on mobile (hidden={isMobile}), matching the SidebarMenuButton tooltip convention. - Show Cmd K on Mac and Ctrl K elsewhere instead of a hardcoded glyph; the search dialog binds both meta and ctrl. Uses getClientPlatform so it is correct on web too, not just Tauri. --- .../frontend/src/components/app-sidebar.tsx | 57 ++++++++++++++++--- .../src/components/tauri/window-titlebar.tsx | 2 +- 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index a13828b06b..f4226760a2 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -45,6 +45,7 @@ import { Spinner } from "@/components/ui/spinner"; import { Switch } from "@/components/ui/switch"; import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler"; import { + getClientPlatform, shouldUseCustomWindowTitlebar, shouldUseNativeMacWindowTitlebar, } from "@/components/tauri/window-titlebar"; @@ -343,6 +344,8 @@ export function AppSidebar() { ); const [usesCustomTitlebar] = useState(shouldUseCustomWindowTitlebar); const [usesNativeMacTitlebar] = useState(shouldUseNativeMacWindowTitlebar); + // Mac uses Cmd, others use Ctrl. Not Tauri-gated, so it's right on web too. + const [isMacPlatform] = useState(() => getClientPlatform().includes("mac")); const { pathname, search } = useRouterState({ select: (s) => ({ pathname: s.location.pathname, @@ -1194,27 +1197,55 @@ export function AppSidebar() { )} - {!isMobile && ( +
- )} + {!isMobile && ( + + + + + + {t("shell.aria.closeSidebar")} + + + )} +
{!isMobile && (
@@ -1246,8 +1277,10 @@ export function AppSidebar() { {/* Uniform pl-1.5 pr-2 keeps every hover pill the same width, inset from the edge. */} @@ -1280,10 +1313,18 @@ export function AppSidebar() { openNewChat(null); }} /> + {/* Search sits in the header when the brand row is shown (mac/web). + Hide this row there, but keep it in the collapsed rail. On custom + titlebars (win/linux) there's no header button, so keep the row. */} { useChatSearchStore.getState().open(); closeMobileIfOpen(); diff --git a/studio/frontend/src/components/tauri/window-titlebar.tsx b/studio/frontend/src/components/tauri/window-titlebar.tsx index 8d11bbd229..d5c74df463 100644 --- a/studio/frontend/src/components/tauri/window-titlebar.tsx +++ b/studio/frontend/src/components/tauri/window-titlebar.tsx @@ -40,7 +40,7 @@ type NavigatorWithUserAgentData = Navigator & { }; }; -function getClientPlatform(): string { +export function getClientPlatform(): string { if (typeof navigator === "undefined") { return ""; } From 59bda2e1f77a3ff060d26b9cdb0b69c798d8c7a1 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Wed, 22 Jul 2026 15:05:33 +0530 Subject: [PATCH 002/194] 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 003/194] 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 004/194] 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 005/194] 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 006/194] 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 007/194] 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 008/194] 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 009/194] 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 010/194] 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 011/194] 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 012/194] 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 013/194] 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 014/194] 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 015/194] 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} +

+ )}

- {showPasswordMismatchWarning - ? "Please ensure passwords match." - : "Must be at least 8 characters."} + {showWhitespaceWarning + ? "New password cannot contain spaces." + : showPasswordMismatchWarning + ? "Please ensure passwords match." + : "Must be at least 8 characters."}

)} diff --git a/studio/frontend/src/features/settings/components/change-password-dialog.tsx b/studio/frontend/src/features/settings/components/change-password-dialog.tsx index cd30d37d5d..c88fc48cac 100644 --- a/studio/frontend/src/features/settings/components/change-password-dialog.tsx +++ b/studio/frontend/src/features/settings/components/change-password-dialog.tsx @@ -78,6 +78,9 @@ function passwordValidationMessage( minLength: MIN_PASSWORD_LENGTH, }); } + if (/\s/.test(nextPassword)) { + return t("settings.general.passwordDialog.newHasSpaces"); + } if (nextPassword !== confirmPassword) { return t("settings.general.passwordDialog.mismatch"); } @@ -160,6 +163,7 @@ export function ChangePasswordDialog() { const currentTooShort = hasStartedTooShortPassword(current); const nextTooShort = hasStartedTooShortPassword(next); + const nextHasSpaces = /\s/.test(next); const mismatch = confirm.length > 0 && next !== confirm; const samePassword = hasReusablePassword(current, next); const validationMessage = passwordValidationMessage( @@ -279,13 +283,15 @@ export function ChangePasswordDialog() { minLength={MIN_PASSWORD_LENGTH} disabled={submitting} /> - {nextTooShort || samePassword ? ( + {nextTooShort || nextHasSpaces || samePassword ? (

{nextTooShort ? t("settings.general.passwordDialog.newTooShort", { minLength: MIN_PASSWORD_LENGTH, }) - : t("settings.general.passwordDialog.samePassword")} + : nextHasSpaces + ? t("settings.general.passwordDialog.newHasSpaces") + : t("settings.general.passwordDialog.samePassword")}

) : null} diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index cf8a29b6d2..164833b41d 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -196,6 +196,7 @@ export const en = { currentTooShort: "Current password must be at least {minLength} characters.", newTooShort: "New password must be at least {minLength} characters.", + newHasSpaces: "New password cannot contain spaces.", mismatch: "Passwords do not match.", samePassword: "New password must be different from your current password.", diff --git a/unsloth_cli/commands/_password_prompt.py b/unsloth_cli/commands/_password_prompt.py index b6fd8ca34d..55f50acbf1 100644 --- a/unsloth_cli/commands/_password_prompt.py +++ b/unsloth_cli/commands/_password_prompt.py @@ -191,6 +191,10 @@ def prompt_new_password(verify_current: Callable[[str], bool], out: TextIO | Non out.write(f"Password must be at least {MIN_PASSWORD_LENGTH} characters. Try again.\n") out.flush() continue + if any(ch.isspace() for ch in password): + out.write("Password cannot contain spaces. Try again.\n") + out.flush() + continue if verify_current(password): out.write("New password must differ from the current password. Try again.\n") out.flush() @@ -233,6 +237,8 @@ def validate_new_password(candidate: str, verify_current: Callable[[str], bool]) current password), else None. Same policy as the interactive loop.""" if len(candidate) < MIN_PASSWORD_LENGTH: return f"Password must be at least {MIN_PASSWORD_LENGTH} characters." + if any(ch.isspace() for ch in candidate): + return "Password cannot contain spaces." if verify_current(candidate): return "New password must differ from the current password." return None From fa5498db0b6c089c1c9ddc8e82043d82be202bc6 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 23 Jul 2026 00:44:42 -0700 Subject: [PATCH 019/194] 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 - + 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 */}
@@ -1830,11 +1832,11 @@ export function AppSidebar() { />
- + {t("shell.updateAvailable")} {updateVersion && ( - + v{updateVersion} )} @@ -1871,8 +1873,8 @@ export function AppSidebar() { {/* min-w-0 so long names truncate instead of overflowing; pr on the button reserves room for the settings cog */}
- {displayTitle} - Unsloth + {displayTitle} + Unsloth
@@ -1880,7 +1882,7 @@ export function AppSidebar() { side="top" align="center" sideOffset={8} - className="app-user-menu menu-soft-surface-up ring-0 w-[16rem] px-2.5 py-2.5 font-heading rounded-[20px] border-0" + className="app-user-menu menu-soft-surface-up ring-0 w-[256px] px-2.5 py-2.5 font-heading rounded-[20px] border-0" > = ({ src }) => { onChange={handleSeek} className="h-1.5 w-full cursor-pointer accent-primary" /> -
+
{formatTime(progress)} {formatTime(duration)}
diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx index cca61b766a..5dd4f1c91d 100644 --- a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -207,7 +207,7 @@ function DetailRow({ }) { if (value == null || value === "") return null; return ( -
+
{label} diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 8d40f587ad..c602a776cf 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -86,7 +86,7 @@ export const MessageTiming: FC<{ data-slot="message-timing-trigger" aria-label="Message timing" className={cn( - "flex items-center rounded-[10px] p-1 font-mono text-chat-icon-fg text-[13px] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover", + "flex items-center rounded-[10px] p-1 font-mono text-chat-icon-fg text-[0.8125rem] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover", className, )} > diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index 9788c73605..d869cc93ef 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -163,7 +163,7 @@ function ReasoningContent({

Generated image

{overlay.metadata ? ( -

+

{overlay.metadata}

) : null} @@ -1390,7 +1390,7 @@ const ComposerAnimated: FC<{ disableQueue?: boolean; }> = ({ disabled, threadId, menuSide, disableQueue }) => { return ( -
+
= ({
) : ( -
+
@@ -3329,7 +3329,7 @@ const PromptQueueStack: FC<{ queueThreadIds: string[] }> = ({ type="button" variant="ghost" size="sm" - className="h-7 w-[5.25rem] justify-center gap-1 px-0 text-sm font-normal text-muted-foreground/80 hover:text-foreground" + className="h-7 w-[84px] justify-center gap-1 px-0 text-sm font-normal text-muted-foreground/80 hover:text-foreground" onClick={() => startEditing(item)} > @@ -3565,14 +3565,14 @@ const DiffusionCanvas: FC = () => { canvas.total > 0 ? `step ${canvas.step + 1}/${canvas.total}` : "denoising"; return (
-
+
Denoising block {canvas.block + 1} - {stepLabel}
- 
+      
         {canvas.text}
@@ -3646,7 +3646,7 @@ const AssistantMessage: FC = () => { return (
@@ -3676,7 +3676,7 @@ const AssistantMessage: FC = () => { ) : ( <>
- +
@@ -3759,7 +3759,7 @@ const ForkCountBadge: FC = () => { if (count <= 0) return null; return ( @@ -4084,7 +4084,7 @@ const UserMessageAudio: FC = () => { const UserMessage: FC = () => { return ( @@ -4195,7 +4195,7 @@ const BranchPicker: FC = ({ = ({ - + / diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx index ec24060072..4fbaa227bd 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx @@ -48,7 +48,7 @@ export function CitationBadge({ - + {errorText ?? (isStaleGeneratingArtifact ? "Refresh stopped this preview" diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx index e38e2e5882..fb1ead3e63 100644 --- a/studio/frontend/src/components/floating-monitor.tsx +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -102,7 +102,7 @@ export function FloatingMonitor() { initial={{ opacity: 0, scale: 0.9 }} animate={{ opacity: 1, scale: 1 }} exit={{ opacity: 0, scale: 0.9 }} - className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-2rem)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none" + className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-32px)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none" >
@@ -139,7 +139,7 @@ export function FloatingMonitor() { className="space-y-3 overflow-hidden" >
-
+
{t("settings.resources.liveMonitor.ram")} -
+
{t("settings.resources.liveMonitor.vram")}{" "} {devices.length > 1 diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 3db15ffe30..25c413ad61 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -131,7 +131,7 @@ export function LlamaUpdateBanner({

-

+

{sizeLabel ? `${sizeLabel} download · ` : ""}No restart needed after update

@@ -209,7 +209,7 @@ export function LlamaUpdateBanner({
diff --git a/studio/frontend/src/components/tauri/update-screen.tsx b/studio/frontend/src/components/tauri/update-screen.tsx index 3199425f69..64f2e95a87 100644 --- a/studio/frontend/src/components/tauri/update-screen.tsx +++ b/studio/frontend/src/components/tauri/update-screen.tsx @@ -72,7 +72,7 @@ function LogViewer({ logs }: { logs: string[] }) { return (
{logs.map((line, i) => (
@@ -197,7 +197,7 @@ export function UpdateScreen({ readOnly value={manualReport} onFocus={(event) => event.currentTarget.select()} - className="mt-2 h-32 w-full max-w-xl resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[10px] text-muted-foreground" + className="mt-2 h-32 w-full max-w-xl resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[0.625rem] text-muted-foreground" /> )} diff --git a/studio/frontend/src/components/tauri/window-titlebar.tsx b/studio/frontend/src/components/tauri/window-titlebar.tsx index d5c74df463..66cc2e1b41 100644 --- a/studio/frontend/src/components/tauri/window-titlebar.tsx +++ b/studio/frontend/src/components/tauri/window-titlebar.tsx @@ -112,8 +112,8 @@ export function WindowTitlebar({ const { pinned, togglePinned } = useSidebarPin(); const sidebarWidth = showSidebarSurface ? pinned - ? "var(--studio-sidebar-expanded-width,17.5rem)" - : "var(--studio-sidebar-collapsed-width,3rem)" + ? "var(--studio-sidebar-expanded-width,280px)" + : "var(--studio-sidebar-collapsed-width,48px)" : "0px"; const contentBorderLeft = pinned ? `calc(${sidebarWidth} + 12px)` : "0px"; @@ -273,7 +273,7 @@ export function WindowTitlebar({ draggable={false} className="size-5 shrink-0 rounded-[6px] object-cover" /> - + Unsloth Studio
@@ -325,7 +325,7 @@ export function WindowTitlebar({ className="pointer-events-auto absolute top-0 h-full" style={{ left: sidebarWidth, - right: "calc(var(--studio-window-control-inset,112px) + 0.5rem)", + right: "calc(var(--studio-window-control-inset,112px) + 8px)", }} onMouseDown={handleDragMouseDown} onDoubleClick={handleDragDoubleClick} diff --git a/studio/frontend/src/components/ui/chart.tsx b/studio/frontend/src/components/ui/chart.tsx index 32da410bb5..25dc88d1cd 100644 --- a/studio/frontend/src/components/ui/chart.tsx +++ b/studio/frontend/src/components/ui/chart.tsx @@ -246,7 +246,7 @@ function ChartTooltipContent({ return (
diff --git a/studio/frontend/src/components/ui/copyable-error-chip.tsx b/studio/frontend/src/components/ui/copyable-error-chip.tsx index 595df5cf62..6f21b6d829 100644 --- a/studio/frontend/src/components/ui/copyable-error-chip.tsx +++ b/studio/frontend/src/components/ui/copyable-error-chip.tsx @@ -53,7 +53,7 @@ export function CopyableErrorChip({
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx index 0345dc6e2a..82a236a387 100644 --- a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx +++ b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx @@ -145,12 +145,12 @@ export function ArtifactCard({ {isCode ? "HTML Code" : artifact.title} - + HTML canvas {isStreaming && !isCode ? ( - + Generating ) : null} diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index e59ce3a805..3daae8c50d 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -576,7 +576,7 @@ function CompareShell({ {children}
-
{composer}
+
{composer}
{showModelDisclaimer && (

LLMs can make mistakes. Double-check responses. @@ -651,7 +651,7 @@ const LoraCompareContent = memo(function LoraCompareContent({ handleName="base" header={

- + Base Model
@@ -665,8 +665,8 @@ const LoraCompareContent = memo(function LoraCompareContent({ handleName="lora" borderClassName="border-t border-border/60 md:border-t-0 md:border-l" header={ -
- +
+ Fine-tuned
@@ -721,8 +721,8 @@ function GeneralCompareHeader({ side === "left" ? pinned ? "pl-12 pr-3 md:pl-2" - : "pl-12 pr-3 md:pl-[calc(0.5rem+max(0px,var(--studio-mac-traffic-light-inset,0px)-var(--sidebar-width-icon,3rem)))]" - : "pl-3 pr-[calc(3rem+var(--studio-chat-header-right-inset,var(--studio-window-control-inset,0px)))]", + : "pl-12 pr-3 md:pl-[calc(8px+max(0px,var(--studio-mac-traffic-light-inset,0px)-var(--sidebar-width-icon,48px)))]" + : "pl-3 pr-[calc(48px+var(--studio-chat-header-right-inset,var(--studio-window-control-inset,0px)))]", )} > {/* Slightly narrower than the composer max; every block shares this. */} -
+
-

+

{projectName}

@@ -1349,7 +1349,7 @@ function ProjectLanding({ type="button" onClick={() => setProjectTab("chats")} data-active={projectTab === "chats"} - className="h-10 rounded-full px-5 text-[14px] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover" + className="h-10 rounded-full px-5 text-[0.875rem] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover" > Chats @@ -1357,7 +1357,7 @@ function ProjectLanding({ type="button" onClick={() => setProjectTab("sources")} data-active={projectTab === "sources"} - className="h-10 rounded-full px-5 text-[14px] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover" + className="h-10 rounded-full px-5 text-[0.875rem] font-semibold transition-colors data-[active=true]:bg-muted data-[active=true]:text-foreground data-[active=false]:text-muted-foreground data-[active=false]:hover:bg-nav-surface-hover" > Sources @@ -1417,7 +1417,7 @@ function ProjectLanding({ onFocus={(event) => event.currentTarget.select()} maxLength={120} aria-label="Rename chat" - className="w-full border-0 bg-transparent text-[15px] font-semibold leading-5 text-foreground outline-none" + className="w-full border-0 bg-transparent text-[0.9375rem] font-semibold leading-5 text-foreground outline-none" />
@@ -1442,11 +1442,11 @@ function ProjectLanding({ className="flex min-h-[58px] min-w-0 flex-1 items-center gap-4 rounded-full px-4 py-2 text-left" >
-
+
{displayTitle}
- + {preview?.date ?? formatProjectChatDate(item.createdAt)} @@ -3105,14 +3105,14 @@ export function ChatPage({ )}
@@ -3141,7 +3141,7 @@ export function ChatPage({ /> )} {incognito && view.mode === "single" && ( -
+

When off, all connections are disabled.

@@ -1616,7 +1616,7 @@ export function ChatProvidersSettings({ {provider.name} - + {provider.models.length}{" "} {provider.models.length === 1 ? "model" : "models"} @@ -1631,7 +1631,7 @@ export function ChatProvidersSettings({ ) : null}
{modelSummary} @@ -1702,7 +1702,7 @@ export function ChatProvidersDialog({ Connections diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index d4f154882c..99d697f619 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -141,7 +141,7 @@ export function ParamSlider({
- + {label} {info && {info}} @@ -249,7 +249,7 @@ function CollapsibleSection({ }; const headerClasses = cn( - "flex w-full items-center justify-between text-[12px] font-medium normal-case tracking-[0.04em] text-nav-fg-muted transition-colors focus-visible:outline-none focus-visible:ring-0", + "flex w-full items-center justify-between text-[0.75rem] font-medium normal-case tracking-[0.04em] text-nav-fg-muted transition-colors focus-visible:outline-none focus-visible:ring-0", first ? "pt-4 pb-5" : "py-5", ); @@ -695,12 +695,12 @@ export function ChatSettingsPanel({ {/* Header is outside the scroll area so the scrollbar never shifts the close button. */}
{isMobile ? ( - + Run settings ) : ( <> - + Run settings @@ -740,7 +740,7 @@ export function ChatSettingsPanel({
{modelConfig} {showSpecFallback && ( -
+

{specFallbackReason === "mla_mtp_disabled" ? "MTP is disabled by default for this model architecture because it currently runs slower than standard decoding. Choose MTP in the model picker to force it." @@ -757,7 +757,7 @@ export function ChatSettingsPanel({ {mtpUpdatable && llamaUpdateStatus?.update_available && (

-

+

Use this for longer edits. Save writes back to the active configuration only. Insert variables with {"{{ env }}"}.

@@ -1230,16 +1230,16 @@ export function ChatSettingsPanel({
-
+
Prompt variables
-

+

Define values as JSON below, then use each key in your prompt, like {"{{ env }}"}.

- + Built-in, fill in automatically
@@ -1247,7 +1247,7 @@ export function ChatSettingsPanel({ {token} @@ -1272,11 +1272,11 @@ export function ChatSettingsPanel({ aria-invalid={Boolean(systemVariablesError)} /> {systemVariablesError ? ( -

+

{systemVariablesError}

) : ( -

+

Names you don't define are left unchanged, so a stray {" {{ typo }} "}stays visible in the prompt.

@@ -1288,7 +1288,7 @@ export function ChatSettingsPanel({ onChange={(event) => setSystemPromptDraft(event.target.value)} placeholder="You are a helpful assistant..." fieldSizing="fixed" - className="min-h-[20rem] max-h-[48vh] overflow-y-auto border-0 text-sm leading-6 corner-squircle focus-visible:ring-0" + className="min-h-[320px] max-h-[48vh] overflow-y-auto border-0 text-sm leading-6 corner-squircle focus-visible:ring-0" rows={14} />
@@ -1333,7 +1333,7 @@ export function ChatSettingsPanel({ if (isMobile) { return ( - + Run settings Chat inference settings @@ -1351,7 +1351,7 @@ export function ChatSettingsPanel({ data-tour="chat-settings" className={cn( "relative z-50 shrink-0 overflow-hidden bg-panel-surface text-panel-surface-fg font-heading", - open ? "w-[17rem] border-l border-sidebar-border" : "w-0", + open ? "w-[272px] border-l border-sidebar-border" : "w-0", )} style={{ height: "calc(100% - var(--studio-custom-titlebar-height, 0px))", @@ -1426,7 +1426,7 @@ function AutoHealToolCallsToggle() { return (
- + Auto-Healing Tool Calls @@ -1450,7 +1450,7 @@ function NudgeToolCallsToggle() { return (
- + Nudge Tool Calls @@ -1475,7 +1475,7 @@ function ConfirmToolCallsToggle() {
- + Confirm tool calls @@ -1487,7 +1487,7 @@ function ConfirmToolCallsToggle() {
{permissionMode === "full" ? ( - + Overridden by Full access ) : null} @@ -1508,7 +1508,7 @@ function BypassPermissionsToggle() { return (
- + Tool permissions @@ -1517,9 +1517,9 @@ function BypassPermissionsToggle() {
{/* Full width, styled like the panel selects/preset input. */} - + {permissionMode === "full" ? ( - + Tool calls run with no confirmation and no sandbox. ) : null} diff --git a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx index dc3e1aac29..ea95040f44 100644 --- a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx +++ b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx @@ -83,7 +83,7 @@ export function ChatSearchDialog() { @@ -143,10 +143,10 @@ export function ChatSearchDialog() { strokeWidth={2} className="size-4 shrink-0 text-muted-foreground" /> - + {item.title || "Untitled chat"} - + {formatRelative(item.createdAt)} diff --git a/studio/frontend/src/features/chat/components/context-usage-bar.tsx b/studio/frontend/src/features/chat/components/context-usage-bar.tsx index 80f502e222..eeacef66df 100644 --- a/studio/frontend/src/features/chat/components/context-usage-bar.tsx +++ b/studio/frontend/src/features/chat/components/context-usage-bar.tsx @@ -71,7 +71,7 @@ export const ContextUsageBar: FC<{ : `Token usage: ${formatTokenCount(used)} tokens` } className={cn( - "flex items-center gap-2 rounded-[10px] px-2.5 py-1 font-mono text-chat-icon-fg text-[13px] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover", + "flex items-center gap-2 rounded-[10px] px-2.5 py-1 font-mono text-chat-icon-fg text-[0.8125rem] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover", className, )} > @@ -149,7 +149,7 @@ export const ContextUsageBar: FC<{
{hasKnownLimit && percent !== null && percent > 85 ? ( -
+
Close to the context limit. Generation will stop at 100%. Increase Context Length in the chat Settings panel to keep going. diff --git a/studio/frontend/src/features/chat/components/model-load-status.tsx b/studio/frontend/src/features/chat/components/model-load-status.tsx index 292c7884fd..613b5c260b 100644 --- a/studio/frontend/src/features/chat/components/model-load-status.tsx +++ b/studio/frontend/src/features/chat/components/model-load-status.tsx @@ -54,14 +54,14 @@ export function ModelLoadDescription({ {title ?

{title}

: null} {hasProgress ? (
-
+
{labelPrimary} {Math.round(clampProgress(progressPercent))}%
{labelSecondary ? ( -
+
{labelSecondary}
) : null} @@ -96,18 +96,18 @@ export function ModelLoadInlineStatus({ const hasProgress = typeof progressPercent === "number"; return ( -
+
{label}
{hasProgress ? (
-
+
{/* Tight inline layout: show only the primary (bytes) chunk; @@ -124,7 +124,7 @@ export function ModelLoadInlineStatus({ type="button" size="xs" variant="outline" - className="shrink-0 text-[11px]" + className="shrink-0 text-[0.6875rem]" onClick={onStop} > Stop diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx index cb0234579b..04c88e4eba 100644 --- a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx +++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx @@ -435,7 +435,7 @@ export function OpenAICodeExecSection({
@@ -459,7 +459,7 @@ export function OpenAICodeExecSection({ ACTIVE pill marks which one (no separate picker). */}
- + Containers
diff --git a/studio/frontend/src/features/chat/components/project-switcher.tsx b/studio/frontend/src/features/chat/components/project-switcher.tsx index 8f923a8c80..2a170e5a39 100644 --- a/studio/frontend/src/features/chat/components/project-switcher.tsx +++ b/studio/frontend/src/features/chat/components/project-switcher.tsx @@ -57,7 +57,7 @@ export function ProjectSwitcher({ className="size-icon shrink-0 text-foreground/70" /> - + {label} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index b72a7a95c2..4b3f57b368 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -111,7 +111,7 @@ const MODEL_LOAD_TOAST_CLASSNAMES = { title: "leading-5", description: "mt-0 w-full", cancelButton: - "!h-auto !rounded-none !border-0 !bg-transparent !px-1 !text-[11px] !font-normal !text-muted-foreground hover:!bg-transparent hover:!text-destructive focus-visible:!text-destructive", + "!h-auto !rounded-none !border-0 !bg-transparent !px-1 !text-[0.6875rem] !font-normal !text-muted-foreground hover:!bg-transparent hover:!text-destructive focus-visible:!text-destructive", } as const; const MODEL_LOADED_TOAST_CLASSNAMES = { diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx index e6c89cf54a..7e0ecb0c7e 100644 --- a/studio/frontend/src/features/chat/permission-mode-select.tsx +++ b/studio/frontend/src/features/chat/permission-mode-select.tsx @@ -120,7 +120,7 @@ export function PermissionModeMenuItems({ > - {option.label} + {option.label} {option.description} diff --git a/studio/frontend/src/features/chat/projects-page.tsx b/studio/frontend/src/features/chat/projects-page.tsx index c9960ffaca..494368faec 100644 --- a/studio/frontend/src/features/chat/projects-page.tsx +++ b/studio/frontend/src/features/chat/projects-page.tsx @@ -367,7 +367,7 @@ export function ProjectsPage() { }} />
-

+

Projects

@@ -419,7 +419,7 @@ export function ProjectsPage() { Export All Projects - + Combined {EXPORT_FORMATS_LIST.map(({ fmt, label }) => ( @@ -430,7 +430,7 @@ export function ProjectsPage() { - + Per chat {EXPORT_FORMATS_LIST.map(({ fmt, label }) => ( @@ -445,7 +445,7 @@ export function ProjectsPage() { Export Projects + Recents - + Combined {EXPORT_FORMATS_LIST.map(({ fmt, label }) => ( @@ -456,7 +456,7 @@ export function ProjectsPage() { - + Per chat {EXPORT_FORMATS_LIST.map(({ fmt, label }) => ( @@ -482,7 +482,7 @@ export function ProjectsPage() { {!hasLoaded ? (
-
+
Name Modified @@ -526,7 +526,7 @@ export function ProjectsPage() {
{/* Column header. Name starts at the folder icon's left edge; the right-anchored columns keep Modified over its values. */} -
+
Name Modified @@ -571,7 +571,7 @@ export function ProjectsPage() { className="size-5" /> - + {project.name} diff --git a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx index 09c4944a14..4b815a7695 100644 --- a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx +++ b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx @@ -1341,7 +1341,7 @@ function ExportModal({ {/* */}
-

+

Export as

@@ -1390,7 +1390,7 @@ function ExportModal({

ShareGPT format for Unsloth fine-tuning

- + {`{"conversations":[{"from":"human","value":"..."},{"from":"gpt","value":""}]}`}
@@ -1400,7 +1400,7 @@ function ExportModal({ {/* */}
-

+

Format

@@ -1730,7 +1730,7 @@ function PromptListCard({
{entry.name} - + {entry.items.length}
@@ -1779,7 +1779,7 @@ function PromptListCard({

))} {entry.items.length > 3 && ( -

+

+{entry.items.length - 3} more

)} diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx index f85c74eb86..4e2765bebd 100644 --- a/studio/frontend/src/features/chat/thread-sidebar.tsx +++ b/studio/frontend/src/features/chat/thread-sidebar.tsx @@ -266,7 +266,7 @@ export function ThreadSidebar({ > {item.isFork ? ( fork diff --git a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx index 088d016894..27584a646f 100644 --- a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx +++ b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx @@ -280,7 +280,7 @@ function LearningRecipeCards({ {badge} @@ -288,7 +288,7 @@ function LearningRecipeCards({ {extraLearningBadgeCount > 0 ? ( +{extraLearningBadgeCount} @@ -296,7 +296,7 @@ function LearningRecipeCards({ {isReady ? null : ( Soon @@ -403,7 +403,7 @@ export function DataRecipesPage(): ReactElement {
-

+

Data Recipes

diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx index 6c2794420b..86c82935d6 100644 --- a/studio/frontend/src/features/export/components/export-run-panel.tsx +++ b/studio/frontend/src/features/export/components/export-run-panel.tsx @@ -278,7 +278,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) {

onSaveDirectoryChange(e.target.value)} spellCheck={false} @@ -303,7 +303,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) { Browse
-

+

{saveDirectory !== defaultSaveDirectory ? ( <>Default: {defaultSaveDirectory} ) : ( @@ -350,7 +350,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) { href="https://huggingface.co/settings/tokens" target="_blank" rel="noopener noreferrer" - className="flex items-center gap-1 text-[11px] text-emerald-600 hover:text-emerald-700 transition-colors" + className="flex items-center gap-1 text-[0.6875rem] text-emerald-600 hover:text-emerald-700 transition-colors" > Get token @@ -369,7 +369,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) { onChange={(e) => onHfTokenChange(e.target.value)} /> -

+

Leave empty if already logged in via CLI.

@@ -427,7 +427,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) { ) : null} {o.path} @@ -503,11 +503,11 @@ export function ExportRunPanel(props: ExportRunPanelProps) { {showProgress && (
- + {PHASE_LABELS[run.phase] ?? run.phase} {summaryMethod === "gguf" && run.quantTotal > 1 && ( - + Quant{" "} {Math.min( run.quantIndex + (isExporting ? 1 : 0), @@ -516,10 +516,10 @@ export function ExportRunPanel(props: ExportRunPanelProps) { of {run.quantTotal} )} - + {progress}% - + {formatElapsed(elapsedSeconds)}
@@ -536,7 +536,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) { /> {run.stage && (

{run.stage} @@ -552,7 +552,7 @@ export function ExportRunPanel(props: ExportRunPanelProps) { -

+
{run.logLines.length === 0 ? (
diff --git a/studio/frontend/src/features/export/components/method-picker.tsx b/studio/frontend/src/features/export/components/method-picker.tsx index 420a7f6146..e240fd44ca 100644 --- a/studio/frontend/src/features/export/components/method-picker.tsx +++ b/studio/frontend/src/features/export/components/method-picker.tsx @@ -123,7 +123,7 @@ export function MethodPicker({ value, onChange, disabledMethods = [], disabledRe {m.badge && ( {m.badge} diff --git a/studio/frontend/src/features/export/components/quant-picker.tsx b/studio/frontend/src/features/export/components/quant-picker.tsx index 289f6498ce..688e5fb87f 100644 --- a/studio/frontend/src/features/export/components/quant-picker.tsx +++ b/studio/frontend/src/features/export/components/quant-picker.tsx @@ -61,7 +61,7 @@ export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) { - + — select one or more
@@ -90,10 +90,10 @@ export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) { )} {q.label} {sizeLabel && ( - {sizeLabel} + {sizeLabel} )} {q.recommended && !active && ( - + rec )} @@ -103,13 +103,13 @@ export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) {
{value.length > 0 && (
- + {value.length} selected diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 3a970713ac..80235846d1 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -895,7 +895,7 @@ export function ExportPage() {
-

+

Export Model

@@ -964,21 +964,21 @@ export function ExportPage() { Local Model Fine-tuned Hugging Face @@ -1289,7 +1289,7 @@ export function ExportPage() { {model?.display_name ?? id} - + {source} @@ -1300,15 +1300,15 @@ export function ExportPage() {

{isLoadingLocalModels ? ( -

+

Scanning local models...

) : localModelsError ? ( -

+

{localModelsError}

) : ( -

+

{exportableLocalModels.length > 0 ? `${exportableLocalModels.length} local/cached models found` : "No local models found. Enter path manually."} @@ -1318,7 +1318,7 @@ export function ExportPage() { )}

-

+

Direct model exports currently support GGUF only.

@@ -1327,7 +1327,7 @@ export function ExportPage() { {sourceMode === "checkpoint" && (
- + Training Info
@@ -1374,7 +1374,7 @@ export function ExportPage() { key={step} className="flex items-start gap-2 text-xs text-muted-foreground" > - + {i + 1} {step} @@ -1422,7 +1422,7 @@ export function ExportPage() {
Precision
- + — select one or more
@@ -1479,7 +1479,7 @@ export function ExportPage() { {f.label} {f.needsCalibration ? " *" : ""} - + {f.hint} @@ -1492,7 +1492,7 @@ export function ExportPage() { {selectedFormats.length > 0 && (
- + {selectedFormats.length} selected:{" "} {selectedFormats .map( @@ -1506,7 +1506,7 @@ export function ExportPage() { @@ -1515,7 +1515,7 @@ export function ExportPage() { )} {hubMultiFormat && ( -
+
Hub export supports one format at a time (each writes to the repository root). Select a single format, or export locally to produce several at once. @@ -1527,13 +1527,13 @@ export function ExportPage() { MERGED_FORMATS.find((f) => f.value === v) ?.needsCalibration, ) && ( -
+
* calibrates on data (uses a small calibration set).
)} {!hasNvidia && ( -
+
No NVIDIA GPU detected: compressed-tensors formats are hidden. 16-bit and portable FP8/INT8 (torchao) still work here and load in vLLM. diff --git a/studio/frontend/src/features/hub/catalog/catalog-states.tsx b/studio/frontend/src/features/hub/catalog/catalog-states.tsx index a36c4bb2da..693b5b40c0 100644 --- a/studio/frontend/src/features/hub/catalog/catalog-states.tsx +++ b/studio/frontend/src/features/hub/catalog/catalog-states.tsx @@ -38,20 +38,20 @@ export function NetworkErrorState({
-

+

{title}

-

+

{body}

-

{message}

+

{message}

{onSwitchDevice ? ( @@ -59,7 +59,7 @@ export function NetworkErrorState({
-

+

No matches yet

-

+

Scanned {scannedCount.toLocaleString()} results. Load another page to keep searching Hugging Face.

@@ -105,7 +105,7 @@ export function DiscoverFetchMoreState({ @@ -114,7 +114,7 @@ export function DiscoverFetchMoreState({ type="button" onClick={onFetchMore} disabled={isLoadingMore} - className="inline-flex h-8 items-center gap-1.5 rounded-full bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] disabled:cursor-not-allowed disabled:opacity-50 dark:hover:bg-white/[0.05]" + className="inline-flex h-8 items-center gap-1.5 rounded-full bg-transparent px-3 text-[0.75rem] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] disabled:cursor-not-allowed disabled:opacity-50 dark:hover:bg-white/[0.05]" > {/* Only warn about hidden results when a filter is actually narrowing them. */} {hasActiveFilters && ( -

+

Some results may be hidden by your filters.

)} @@ -149,7 +149,7 @@ export function DiscoverFetchMoreFooter({ type="button" onClick={onFetchMore} disabled={isLoadingMore} - className="inline-flex h-8 items-center gap-1.5 rounded-full bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] disabled:cursor-not-allowed disabled:opacity-50 dark:bg-white/[0.06] dark:hover:bg-white/[0.1]" + className="inline-flex h-8 items-center gap-1.5 rounded-full bg-foreground/[0.06] px-3 text-[0.75rem] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] disabled:cursor-not-allowed disabled:opacity-50 dark:bg-white/[0.06] dark:hover:bg-white/[0.1]" >
-

+

Couldn't load your library

-

+

Something went wrong reading your downloaded{" "} {isDataset ? "datasets" : "models"}. Check that the backend is running and try again. @@ -187,7 +187,7 @@ export function InventoryErrorState({

-

+

{title}

-

+

{body}

diff --git a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx index b821be4b0b..ade8d2de30 100644 --- a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx +++ b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx @@ -126,7 +126,7 @@ export function DatasetDownloadSection({ } >
- + {isDownloaded && } {!isDownloaded && isPartial && !downloading && ( diff --git a/studio/frontend/src/features/hub/catalog/dot-tag.tsx b/studio/frontend/src/features/hub/catalog/dot-tag.tsx index 9452a201d4..5ae1be53d0 100644 --- a/studio/frontend/src/features/hub/catalog/dot-tag.tsx +++ b/studio/frontend/src/features/hub/catalog/dot-tag.tsx @@ -36,7 +36,7 @@ export function DotTag({ return ( diff --git a/studio/frontend/src/features/hub/catalog/download-card.tsx b/studio/frontend/src/features/hub/catalog/download-card.tsx index 9b4bc5fd01..9bc64ced0e 100644 --- a/studio/frontend/src/features/hub/catalog/download-card.tsx +++ b/studio/frontend/src/features/hub/catalog/download-card.tsx @@ -140,7 +140,7 @@ export function CardUpdateButton({ e.stopPropagation(); onClick(); }} - className="inline-flex h-7 shrink-0 cursor-pointer items-center gap-1.5 rounded-full bg-amber-500/[0.07] pl-2 pr-2.5 text-[12px] font-medium text-amber-800/90 transition-colors duration-150 hover:bg-amber-500/[0.12] focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-amber-500/25 dark:bg-amber-400/[0.08] dark:text-amber-200/85 dark:hover:bg-amber-400/[0.16]" + className="inline-flex h-7 shrink-0 cursor-pointer items-center gap-1.5 rounded-full bg-amber-500/[0.07] pl-2 pr-2.5 text-[0.75rem] font-medium text-amber-800/90 transition-colors duration-150 hover:bg-amber-500/[0.12] focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-amber-500/25 dark:bg-amber-400/[0.08] dark:text-amber-200/85 dark:hover:bg-amber-400/[0.16]" > {pendingUrl && (
-

+

{hostOf(pendingUrl)}

-

+

{pendingUrl}

diff --git a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx index 0d6878b687..9345874a53 100644 --- a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx +++ b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx @@ -128,7 +128,7 @@ const FIT_BADGE: Record = { /** Chip styling matching the on-device list's StatChip, no icon. */ const CHIP_BASE = - "inline-flex h-5 shrink-0 items-center justify-center whitespace-nowrap rounded-full border px-2 text-[11.5px] font-medium tabular-nums leading-3 shadow-[inset_0_1px_0_rgba(255,255,255,0.04)]"; + "inline-flex h-5 shrink-0 items-center justify-center whitespace-nowrap rounded-full border px-2 text-[0.71875rem] font-medium tabular-nums leading-3 shadow-[inset_0_1px_0_rgba(255,255,255,0.04)]"; const CHIP_DEFAULT = "border-foreground/15 bg-muted text-foreground/85 dark:border-border/60 dark:bg-white/[0.04] dark:text-foreground/85"; const CHIP_ACTIVE = @@ -184,7 +184,7 @@ function QuantBadge({ // group's `overflow-hidden` sacrifices the trailing status tags instead. @@ -914,7 +914,7 @@ export function GgufDownloadCard({ {/* Quant label + status tags travel together as one left-aligned group so the fit-info icon never floats orphaned from its tags; only the chevron pins right, the standard select affordance. */} - + {selected ? ( ) : ( - + Select quantization )} @@ -1126,7 +1126,7 @@ export function GgufDownloadCard({ diff --git a/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx b/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx index 6cc764b876..c7f402159b 100644 --- a/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx +++ b/studio/frontend/src/features/hub/catalog/gguf-status-cards.tsx @@ -34,7 +34,7 @@ export function GgufDownloadStatusCard({
@@ -91,7 +91,7 @@ export function GgufDownloadingFallbackCard({
- + {progress.variant && } Downloading… diff --git a/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx index d733049ec3..e7d176442a 100644 --- a/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx +++ b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx @@ -69,7 +69,7 @@ export function HubDetailView({ ) : ( - + {content} )} @@ -325,11 +325,11 @@ function ModelStatusChips({ > This model may not be supported yet. {unslothSupport.reason && ( - + {unslothSupport.reason} )} - + Still downloadable to your Hugging Face cache. @@ -349,7 +349,7 @@ function ModelStatusChips({ > This device has no supported GPU or usable MLX, so only GGUF models can run here. - + Still downloadable to your Hugging Face cache. @@ -368,7 +368,7 @@ function ModelStatusChips({ className="tooltip-compact max-w-xs" > Estimated 4-bit memory load is around {vramInfo.est} GB. - + {vramDetail} @@ -502,10 +502,10 @@ export const ModelInspector = memo(function ModelInspector({
-

+

Select a {isDataset ? "dataset" : "model"}

-

+

{isDataset ? "Choose a dataset from the catalog to inspect its download state and details." : "Choose an item from the catalog to inspect its runtime fit, download state, and model card."} @@ -528,7 +528,7 @@ export const ModelInspector = memo(function ModelInspector({ model.downloadsAllTime != null ? ( <> Downloads (30 days) - + {formatCompact(model.downloadsAllTime)} all time @@ -582,11 +582,11 @@ export const ModelInspector = memo(function ModelInspector({

-

+

{model.title}

{model.hubRepoId && ( @@ -599,7 +599,7 @@ export const ModelInspector = memo(function ModelInspector({
)}
-
+
{model.owner} {model.owner.toLowerCase() === "unsloth" && ( {isDataset && ( - + Dataset )} {!isDataset && ( - + {selectionHiddenByFilters && ( -

+

Current selection is hidden by the active filters or search.

)} {metadataUnavailable && ( -

+

Couldn't load full details from Hugging Face. Some fields may be incomplete.

diff --git a/studio/frontend/src/features/hub/catalog/model-readme.tsx b/studio/frontend/src/features/hub/catalog/model-readme.tsx index 3d215473e2..42cc55372a 100644 --- a/studio/frontend/src/features/hub/catalog/model-readme.tsx +++ b/studio/frontend/src/features/hub/catalog/model-readme.tsx @@ -126,17 +126,17 @@ function prepareReadmeBody(markdown: string): string { } const PROSE = cn( - "max-w-none text-[13.5px] leading-[1.7] text-foreground/85", - "[&_h1]:text-[18px] [&_h1]:font-semibold [&_h1]:tracking-tight [&_h1]:mt-2 [&_h1]:mb-3", - "[&_h2]:text-[15.5px] [&_h2]:font-semibold [&_h2]:tracking-tight [&_h2]:mt-5 [&_h2]:mb-2", - "[&_h3]:text-[14px] [&_h3]:font-semibold [&_h3]:mt-4 [&_h3]:mb-1.5", + "max-w-none text-[0.84375rem] leading-[1.7] text-foreground/85", + "[&_h1]:text-[1.125rem] [&_h1]:font-semibold [&_h1]:tracking-tight [&_h1]:mt-2 [&_h1]:mb-3", + "[&_h2]:text-[0.96875rem] [&_h2]:font-semibold [&_h2]:tracking-tight [&_h2]:mt-5 [&_h2]:mb-2", + "[&_h3]:text-[0.875rem] [&_h3]:font-semibold [&_h3]:mt-4 [&_h3]:mb-1.5", "[&_p]:my-2.5 [&_ul]:my-2 [&_ol]:my-2 [&_li]:my-0.5", "[&_a]:text-primary [&_a:hover]:underline", - "[&_code]:rounded-md [&_code]:bg-muted/60 [&_code]:px-1 [&_code]:py-0.5 [&_code]:text-[12px] [&_code]:font-mono", - "[&_pre]:rounded-[12px] [&_pre]:border [&_pre]:border-border/60 [&_pre]:bg-muted/40 [&_pre]:p-3 [&_pre]:text-[12px] [&_pre]:overflow-x-auto", + "[&_code]:rounded-md [&_code]:bg-muted/60 [&_code]:px-1 [&_code]:py-0.5 [&_code]:text-[0.75rem] [&_code]:font-mono", + "[&_pre]:rounded-[12px] [&_pre]:border [&_pre]:border-border/60 [&_pre]:bg-muted/40 [&_pre]:p-3 [&_pre]:text-[0.75rem] [&_pre]:overflow-x-auto", "[&_pre_code]:bg-transparent [&_pre_code]:p-0", "[&_blockquote]:border-l-2 [&_blockquote]:border-border/60 [&_blockquote]:pl-3 [&_blockquote]:text-muted-foreground", - "[&_table]:my-3 [&_table]:text-[12.5px]", + "[&_table]:my-3 [&_table]:text-[0.78125rem]", "[&_th]:px-2 [&_th]:py-1.5 [&_th]:text-left [&_th]:font-semibold [&_th]:border-b [&_th]:border-border/60", "[&_td]:px-2 [&_td]:py-1.5 [&_td]:border-b [&_td]:border-border/40", "[&_img]:rounded-[10px] [&_img]:my-2 [&_img]:max-w-full", @@ -300,7 +300,7 @@ function ReadmePlaceholder({ aria-busy="true" aria-live="polite" > -
+
{message ?? `Loading ${kind === "dataset" ? "dataset" : "model"} card…`}
@@ -540,7 +540,7 @@ export function ModelReadme({ ? current.error : readmeUnavailableMessage(subject); return ( -

+

{errorMessage}

); @@ -548,7 +548,7 @@ export function ModelReadme({ if (!current.body) { return ( -

+

{readmeMissingMessage(subject)}

); diff --git a/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx b/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx index 8cf5fc491e..926c002a65 100644 --- a/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx +++ b/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx @@ -68,7 +68,7 @@ export function InventoryWarningRow({ onRetry: () => void; }) { return ( -
+
Some on-device sources couldn't be scanned. Showing available{" "} @@ -76,7 +76,7 @@ export function InventoryWarningRow({ @@ -444,7 +444,7 @@ export function DownloadedList({ <> {pinnedItems.length > 0 && ( <> -
+
{unpinnedItems.length > 0 && ( -
+
All {isDataset ? "datasets" : "models"}
)} diff --git a/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx b/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx index 6d1dc20414..156bafba70 100644 --- a/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx +++ b/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx @@ -188,14 +188,14 @@ function CachedSizeChipLive({ ))} ) : ( - + {variantMessage} )} @@ -225,7 +225,7 @@ export function StatChip({ return ( @@ -482,7 +482,7 @@ export const DiscoverModelRow = memo(function DiscoverModelRow({
-

+

{row.repo}

-
+
{row.owner} {row.owner.toLowerCase() === "unsloth" && ( @@ -528,7 +528,7 @@ export const DiscoverModelRow = memo(function DiscoverModelRow({ /> )} - + {formatRelativeShort(row.result.updatedAt)}
@@ -666,7 +666,7 @@ export const InventoryRow = memo(function InventoryRow({ {paramLabel} )} {quantLabel && ( - + {quantLabel} )} @@ -716,7 +716,7 @@ export const InventoryRow = memo(function InventoryRow({ ) : null; const ownerLine = ( - + {subLabel} {subLabel.toLowerCase() === "unsloth" && (
- + {title} {compactMarkers}
- + {subLabel} {subLabel.toLowerCase() === "unsloth" && ( @@ -851,7 +851,7 @@ export const InventoryRow = memo(function InventoryRow({ )}
-
+
{row.kind === "cache" ? (
- + {title} {statusMarkers} @@ -915,11 +915,11 @@ export const InventoryRow = memo(function InventoryRow({ cachePath={row.cachePath} /> ) : trailing ? ( - + {trailing} ) : sourceLabel ? ( - + {sourceLabel} ) : null} diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx index 9629fc0a10..1702ed6fca 100644 --- a/studio/frontend/src/features/hub/catalog/models-header.tsx +++ b/studio/frontend/src/features/hub/catalog/models-header.tsx @@ -91,7 +91,7 @@ export function ModelsHeader({ {activeCheckpoint && ( -
+