Merge branch 'main' into docker-blackwell-build

tests/_zoo_aggressive_cuda_spoof.py

@@ -1,15 +1,12 @@
 # Auto-generated by .github/workflows/consolidated-tests-ci.yml.
 # Aggressive CUDA spoof for the consolidated CPU-only CI job. Extends
-# tests/conftest.py:84-141's import-time harness with deeper patches that
-# unblock more patch_* functions and unsloth_zoo init paths on a GPU-less
-# runner. Imported by every shim test file in this workflow before any
-# unsloth / unsloth_zoo / transformers import.
+# tests/conftest.py's import-time harness with deeper patches that unblock
+# more patch_* and unsloth_zoo init paths on a GPU-less runner. Imported by
+# every shim test file before any unsloth / unsloth_zoo / transformers import.
 #
-# Design: only no-op or value-returning patches. We do NOT replace tensor
-# allocators. The single exception is `pin_memory=True` kwarg dropping,
-# which converts a hard CUDA-required call into a CPU-OK call -- the
-# intent of pin_memory is a CUDA-host fast-copy, which simply has no
-# meaning on this runner; downgrading silently is the right behavior here.
+# Only no-op or value-returning patches; tensor allocators are NOT replaced.
+# The one exception is dropping `pin_memory=True` (a CUDA-host fast-copy hint
+# that is meaningless here), which downgrades a CUDA-required call to CPU-OK.
 
 from __future__ import annotations
 
@@ -25,7 +22,7 @@ def apply() -> None:
     if getattr(torch.cuda, "_unsloth_consolidated_spoof", False):
         return
 
-    # ----- device probes (cheap, value-returning) -------------------------
+    # Device probes (cheap, value-returning)
     torch.cuda.is_available = lambda: True
     torch.cuda.device_count = lambda: 1
     torch.cuda.current_device = lambda: 0
@@ -49,7 +46,7 @@ def apply() -> None:
 
     torch.cuda.get_device_properties = lambda *a, **k: _Props()  # type: ignore[assignment]
 
-    # ----- cudart() wrapper -----------------------------------------------
+    # cudart() wrapper
     class _CudaRt:
         @staticmethod
         def cudaMemGetInfo(device: int = 0):
@@ -57,7 +54,7 @@ def apply() -> None:
 
         @staticmethod
         def cudaGetDeviceCount(*_a, **_k):
-            return 0  # Not used on the spoof path
+            return 0  # unused on the spoof path
 
         @staticmethod
         def cudaSetDevice(*_a, **_k):
@@ -65,7 +62,7 @@ def apply() -> None:
 
     torch.cuda.cudart = lambda: _CudaRt()  # type: ignore[assignment]
 
-    # ----- memory module --------------------------------------------------
+    # memory module
     try:
         import torch.cuda.memory as _cuda_memory  # type: ignore
 
@@ -79,7 +76,7 @@ def apply() -> None:
     except Exception:
         pass
 
-    # ----- nvtx no-op stub ------------------------------------------------
+    # nvtx no-op stub
     nvtx_stub = types.ModuleType("torch.cuda.nvtx")
     nvtx_stub.range_push = lambda *a, **k: None  # type: ignore[attr-defined]
     nvtx_stub.range_pop = lambda *a, **k: None  # type: ignore[attr-defined]
@@ -87,18 +84,15 @@ def apply() -> None:
     sys.modules.setdefault("torch.cuda.nvtx", nvtx_stub)
     torch.cuda.nvtx = nvtx_stub  # type: ignore[attr-defined]
 
-    # ----- random API ----------------------------------------------------
-    # CRITICAL: torch.manual_seed() internally calls torch.cuda.manual_seed_all(),
-    # so routing the cuda seed APIs back through torch.manual_seed would
-    # infinite-recurse (observed as RecursionError in run #8 cells 2/3 of the
-    # consolidated CI matrix). No-op them: callers that explicitly seed CUDA
-    # have already paid the cost of seeding CPU via torch.manual_seed; the
-    # CUDA-side seeding has no meaning on a GPU-less runner.
+    # random API
+    # CRITICAL: torch.manual_seed() calls torch.cuda.manual_seed_all(), so
+    # routing the cuda seed APIs back through torch.manual_seed would
+    # infinite-recurse (RecursionError in CI). No-op them; CUDA-side seeding
+    # is meaningless on a GPU-less runner.
     torch.cuda.manual_seed = lambda *a, **k: None  # type: ignore[assignment]
     torch.cuda.manual_seed_all = lambda *a, **k: None  # type: ignore[assignment]
-    # rng_state APIs: return a CPU-shaped placeholder and accept anything for
-    # set; do NOT route through torch.set_rng_state / get_rng_state -- those
-    # operate on the CPU RNG directly and are independent of the cuda surface.
+    # rng_state APIs: return a CPU-shaped placeholder, accept anything for set;
+    # do NOT route through torch.{get,set}_rng_state (those touch the CPU RNG).
     import torch as _t
 
     _empty_rng_state = _t.empty(0, dtype = _t.uint8)
@@ -110,7 +104,7 @@ def apply() -> None:
     torch.cuda.seed = lambda *a, **k: None  # type: ignore[assignment]
     torch.cuda.seed_all = lambda *a, **k: None  # type: ignore[assignment]
 
-    # ----- Stream / Event no-op classes -----------------------------------
+    # Stream / Event no-op classes
     class _NoopStream:
         def __init__(self, *a, **k): ...
         def __enter__(self):
@@ -141,9 +135,8 @@ def apply() -> None:
     torch.cuda.current_stream = lambda *a, **k: _NoopStream()  # type: ignore[assignment]
     torch.cuda.default_stream = lambda *a, **k: _NoopStream()  # type: ignore[assignment]
 
-    # ----- pin_memory drop -------------------------------------------------
-    # `torch.empty(..., pin_memory=True)` and friends raise on a CPU-only
-    # build. Strip the kwarg — pin_memory has no meaning here.
+    # pin_memory drop: torch.empty(..., pin_memory=True) and friends raise on
+    # a CPU-only build; strip the kwarg since pin_memory has no meaning here.
     for _name in (
         "empty",
         "zeros",
@@ -159,7 +152,11 @@ def apply() -> None:
         if _orig is None:
             continue
 
-        def _wrap(*args: Any, _orig = _orig, **kwargs: Any):
+        def _wrap(
+            *args: Any,
+            _orig = _orig,
+            **kwargs: Any,
+        ):
             kwargs.pop("pin_memory", None)
             return _orig(*args, **kwargs)
 
@@ -171,10 +168,8 @@ def apply() -> None:
     if hasattr(torch.Tensor, "is_pinned"):
         torch.Tensor.is_pinned = lambda self, *a, **k: False  # type: ignore[assignment]
 
-    # ----- amp.GradScaler: use the real one if torch ships a CPU-friendly
-    # path, else stub. Newer torch ships torch.amp.GradScaler that handles
-    # CPU; torch.cuda.amp.GradScaler is a wrapper. Both should work; just
-    # guard against import error.
+    # amp.GradScaler: use the real one if importable (newer torch handles CPU),
+    # else stub.
     try:
         import torch.cuda.amp  # type: ignore
     except Exception:
@@ -205,7 +200,7 @@ def apply() -> None:
         sys.modules.setdefault("torch.cuda.amp", cuda_amp)
         torch.cuda.amp = cuda_amp  # type: ignore[attr-defined]
 
-    # ----- Sentinel ------------------------------------------------------
+    # Sentinel
     torch.cuda._unsloth_consolidated_spoof = True  # type: ignore[attr-defined]
 
 

tests/conftest.py

@@ -105,13 +105,11 @@ def _patch_torch_cuda_for_import() -> None:
     CPU like normal."""
     try:
         import torch.cuda.memory as _cuda_memory  # type: ignore
-
         _cuda_memory.mem_get_info = lambda *a, **k: (0, 80 * 1024**3)
     except Exception:
         pass
     try:
         import torch
-
         torch.cuda.get_device_capability = lambda *a, **k: (8, 0)
         torch.cuda.is_bf16_supported = lambda *a, **k: True
     except Exception:
@@ -142,14 +140,11 @@ if not _has_real_accelerator():
 
 
 # ---------------------------------------------------------------------------
-# Apply ALL upstream-drift fixes (vllm GuidedDecodingParams alias, triton
-# CompiledKernel attr wrap, peft transformers_weight_conversion stub, etc.)
-# by triggering ``import unsloth``. Fixes live on ``unsloth/import_fixes.py``
-# and apply at unsloth import time. The GPU-free harness above pre-spoofs
-# the device-type chain so ``import unsloth`` survives on a CPU-only runner.
-# Suites without unsloth installed (e.g. security-only) keep passing --
-# the ImportError is swallowed and the drift detectors will surface any
-# pathology the missing patches would have hidden.
+# Apply upstream-drift fixes (vllm/triton/peft) by triggering ``import
+# unsloth``; they live in ``unsloth/import_fixes.py`` and run at import time.
+# The GPU-free harness above lets ``import unsloth`` survive CPU-only runners.
+# Suites without unsloth keep passing -- the ImportError is swallowed and the
+# drift detectors surface anything the missing patches would have hidden.
 # ---------------------------------------------------------------------------
 
 

tests/python/conftest.py

@@ -2,9 +2,5 @@
 
 
 def pytest_configure(config):
-    config.addinivalue_line(
-        "markers", "server: heavyweight tests requiring studio venv"
-    )
-    config.addinivalue_line(
-        "markers", "e2e: end-to-end tests requiring network and venv creation"
-    )
+    config.addinivalue_line("markers", "server: heavyweight tests requiring studio venv")
+    config.addinivalue_line("markers", "e2e: end-to-end tests requiring network and venv creation")

tests/python/test_cross_platform_parity.py

@@ -21,24 +21,17 @@ class TestNoTorchBackendAutoInInstallSh:
 
     def test_no_torch_backend_auto_outside_fallback(self):
         lines = INSTALL_SH.read_text().splitlines()
-        # Find the fallback block: starts with the "else" after the
-        # TORCH_INDEX_URL check and ends at the next "fi".
+        # Fallback block: from "GPU detection failed" to the next "fi".
         fallback_start = None
         fallback_end = None
         for i, line in enumerate(lines):
             if fallback_start is None and "GPU detection failed" in line:
                 fallback_start = i
-            elif (
-                fallback_start is not None
-                and fallback_end is None
-                and line.strip() == "fi"
-            ):
+            elif fallback_start is not None and fallback_end is None and line.strip() == "fi":
                 fallback_end = i
                 break
         fallback_range = (
-            range(fallback_start or 0, (fallback_end or 0) + 1)
-            if fallback_start
-            else range(0)
+            range(fallback_start or 0, (fallback_end or 0) + 1) if fallback_start else range(0)
         )
 
         matches = [

tests/python/test_dpo_vision_processor_passthrough.py

@@ -33,7 +33,11 @@ class _Tok:
     eos_token_id = 99
     bos_token_id = None
 
-    def __call__(self, t, add_special_tokens = False):
+    def __call__(
+        self,
+        t,
+        add_special_tokens = False,
+    ):
         return {"input_ids": [10]}
 
 
@@ -46,7 +50,12 @@ class _Capture:
         self.last_text = None
         self.last_images = "__sentinel__"
 
-    def __call__(self, images = None, text = None, add_special_tokens = False):
+    def __call__(
+        self,
+        images = None,
+        text = None,
+        add_special_tokens = False,
+    ):
         self.last_text = text
         self.last_images = images
         out = {"input_ids": [[1, 2]]}

tests/python/test_e2e_no_torch_sandbox.py

@@ -1,19 +1,11 @@
-"""Comprehensive E2E sandbox tests for PR #4624 (fix/install-mac-intel-no-torch).
-
-Proves that:
-- The BEFORE state (top-level torch imports) crashes without torch
-- The AFTER state (lazy/removed imports) works without torch
-- Edge cases (broken torch, partial torch) are handled gracefully
-- Hardware detection falls back to CPU without torch
-- install.sh flag parsing and platform detection work correctly
-- install_python_stack.py NO_TORCH filtering is correct
-- Live server starts and responds without torch (optional, requires studio venv)
+"""E2E sandbox tests for PR #4624 (fix/install-mac-intel-no-torch): BEFORE
+(top-level torch) crashes, AFTER (lazy imports) works, broken/partial torch,
+CPU hardware fallback, install.sh parsing, NO_TORCH filtering, and live server.
 
 Run:
-    # Lightweight tests (Groups 1-6, ~26 tests):
+    # Lightweight (Groups 1-6):
     python -m pytest tests/python/test_e2e_no_torch_sandbox.py -v -k "not server"
-
-    # Server tests (Group 7, 4 tests, requires studio venv):
+    # Server (Group 7, requires studio venv):
     python -m pytest tests/python/test_e2e_no_torch_sandbox.py -v -m server
 """
 
@@ -212,12 +204,8 @@ def no_torch_venv(request, tmp_path_factory):
 
 
 class TestBeforeAfterImportChain:
-    """Prove the bug exists in BEFORE state and is fixed in AFTER state.
-
-    BEFORE = PR branch files with top-level torch import synthetically prepended
-             (simulates the main branch).
-    AFTER  = PR branch files as-is (lazy imports / torch import removed).
-    """
+    """BEFORE (PR files with a synthetic top-level torch import, simulating
+    main) crashes; AFTER (PR files as-is, lazy imports) works."""
 
     # -- BEFORE: crashes --
 
@@ -247,12 +235,8 @@ class TestBeforeAfterImportChain:
             exec(source)
         """)
         result = _run_in_sandbox(no_torch_venv, code)
-        assert (
-            result.returncode != 0
-        ), "BEFORE chat_templates.py should crash without torch"
-        assert (
-            b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
-        )
+        assert result.returncode != 0, "BEFORE chat_templates.py should crash without torch"
+        assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
 
     def test_before_data_collators_crashes(self, no_torch_venv, sandbox_dir):
         """BEFORE: data_collators.py with top-level 'import torch' crashes."""
@@ -270,12 +254,8 @@ class TestBeforeAfterImportChain:
             exec(open({str(before_file)!r}).read())
         """)
         result = _run_in_sandbox(no_torch_venv, code)
-        assert (
-            result.returncode != 0
-        ), "BEFORE data_collators.py should crash without torch"
-        assert (
-            b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
-        )
+        assert result.returncode != 0, "BEFORE data_collators.py should crash without torch"
+        assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
 
     def test_before_full_import_chain_crashes(self, no_torch_venv, sandbox_dir):
         """BEFORE: full utils/datasets/ package with top-level torch imports crashes."""
@@ -320,12 +300,8 @@ class TestBeforeAfterImportChain:
             from utils.datasets import detect_dataset_format
         """)
         result = _run_in_sandbox(no_torch_venv, code)
-        assert (
-            result.returncode != 0
-        ), "BEFORE full import chain should crash without torch"
-        assert (
-            b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
-        )
+        assert result.returncode != 0, "BEFORE full import chain should crash without torch"
+        assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
 
     # -- AFTER: succeeds --
 
@@ -539,9 +515,7 @@ class TestEdgeCasesBrokenTorch:
             print("OK: data_collators works despite broken torch on sys.path")
         """)
         result = _run_in_sandbox(no_torch_venv, code)
-        assert (
-            result.returncode == 0
-        ), f"Should work with broken torch:\n{result.stderr.decode()}"
+        assert result.returncode == 0, f"Should work with broken torch:\n{result.stderr.decode()}"
         assert b"OK:" in result.stdout
 
     def test_torch_import_error_hardware_fallback(self, no_torch_venv, sandbox_dir):
@@ -604,14 +578,10 @@ class TestEdgeCasesBrokenTorch:
             print("OK: detect_hardware returned CPU with fake torch (no CUDA)")
         """)
         result = _run_in_sandbox(no_torch_venv, code)
-        assert (
-            result.returncode == 0
-        ), f"Should fall back to CPU:\n{result.stderr.decode()}"
+        assert result.returncode == 0, f"Should fall back to CPU:\n{result.stderr.decode()}"
         assert b"OK:" in result.stdout
 
-    def test_lazy_torch_fails_at_call_time_not_import_time(
-        self, no_torch_venv, sandbox_dir
-    ):
+    def test_lazy_torch_fails_at_call_time_not_import_time(self, no_torch_venv, sandbox_dir):
         """apply_chat_template_to_dataset is importable without torch.
 
         Calling the alpaca branch triggers the lazy 'from torch.utils.data' inside
@@ -657,9 +627,7 @@ class TestEdgeCasesBrokenTorch:
                 print("OK: call succeeded (unexpected but not a crash)")
         """)
         result = _run_in_sandbox(no_torch_venv, code)
-        assert (
-            result.returncode == 0
-        ), f"Should not crash at import time:\n{result.stderr.decode()}"
+        assert result.returncode == 0, f"Should not crash at import time:\n{result.stderr.decode()}"
         assert b"OK: import succeeded" in result.stdout
 
 
@@ -1011,9 +979,7 @@ class TestInstallPythonStackFiltering:
         source = Path(ips.__file__).read_text(encoding = "utf-8")
 
         # NO_TORCH guard before overrides
-        assert (
-            "if NO_TORCH:" in source
-        ), "NO_TORCH guard not found in install_python_stack.py"
+        assert "if NO_TORCH:" in source, "NO_TORCH guard not found in install_python_stack.py"
 
         # macOS guard for triton
         assert (
@@ -1037,7 +1003,6 @@ def _studio_venv_python() -> Path | None:
 def _server_port() -> int:
     """Find an available port for the test server."""
     import socket
-
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         s.bind(("", 0))
         return s.getsockname()[1]
@@ -1117,9 +1082,7 @@ class TestLiveServerStartup:
         for _ in range(30):
             time.sleep(1)
             try:
-                resp = urllib.request.urlopen(
-                    f"http://127.0.0.1:{port}/api/health", timeout = 2
-                )
+                resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/api/health", timeout = 2)
                 if resp.status == 200:
                     ready = True
                     break
@@ -1143,12 +1106,8 @@ class TestLiveServerStartup:
                     capture_output = True,
                     timeout = 300,
                 )
-            server_output = stdout.decode(errors = "replace") + stderr.decode(
-                errors = "replace"
-            )
-            pytest.skip(
-                f"Server failed to start within 30 seconds. Output:\n{server_output}"
-            )
+            server_output = stdout.decode(errors = "replace") + stderr.decode(errors = "replace")
+            pytest.skip(f"Server failed to start within 30 seconds. Output:\n{server_output}")
 
         yield proc, port
 
@@ -1192,9 +1151,7 @@ class TestLiveServerStartup:
         import urllib.request
 
         _, port = server_process
-        resp = urllib.request.urlopen(
-            f"http://127.0.0.1:{port}/openapi.json", timeout = 5
-        )
+        resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/openapi.json", timeout = 5)
         spec = json.loads(resp.read())
         assert (
             len(spec.get("paths", {})) >= 20

tests/python/test_fast_language_model_text_only.py

@@ -0,0 +1,433 @@
+"""Text-only FastLanguageModel routing for vision-capable configs."""
+
+import ast
+import copy
+from pathlib import Path
+
+import pytest
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+LOADER_PATH = REPO_ROOT / "unsloth" / "models" / "loader.py"
+VISION_PATH = REPO_ROOT / "unsloth" / "models" / "vision.py"
+UTILS_PATH = REPO_ROOT / "unsloth" / "models" / "_utils.py"
+
+
+def _source(path):
+    return path.read_text()
+
+
+def _class_method(tree, class_name, method_name):
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef) and node.name == class_name:
+            for item in node.body:
+                if isinstance(item, ast.FunctionDef) and item.name == method_name:
+                    return item
+    raise AssertionError(f"{class_name}.{method_name} not found")
+
+
+def _assigns_name(method, target_name, predicate):
+    """True when the method contains `target_name = <value>` and predicate(value)."""
+    for node in ast.walk(method):
+        if not isinstance(node, ast.Assign):
+            continue
+        for target in node.targets:
+            if isinstance(target, ast.Name) and target.id == target_name:
+                if predicate(node.value):
+                    return True
+    return False
+
+
+def _calls_function(method, func_name):
+    """True when the method calls `func_name(...)` (bare name, not attribute)."""
+    for node in ast.walk(method):
+        if (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == func_name
+        ):
+            return True
+    return False
+
+
+def _names_in(node):
+    return {n.id for n in ast.walk(node) if isinstance(n, ast.Name)}
+
+
+def _param_default(method, name):
+    # Default-value AST node for a named parameter, or None.
+    args = method.args
+    params = list(args.args) + list(args.kwonlyargs)
+    defaults = list(args.defaults) + list(args.kw_defaults)
+    return dict(zip([p.arg for p in params][-len(defaults) :], defaults)).get(name)
+
+
+def _load_text_only_namespace():
+    # Exec the text-only helpers from _utils into one namespace (no unsloth import),
+    # in dependency order so cross-references resolve.
+    source = _source(UTILS_PATH)
+    import transformers
+    from packaging.version import Version
+
+    ns = {
+        "copy": copy,
+        "Version": Version,
+        "transformers_version": transformers.__version__,
+    }
+    funcs = {
+        node.name: ast.get_source_segment(source, node)
+        for node in ast.parse(source).body
+        if isinstance(node, ast.FunctionDef)
+    }
+    for name in (
+        "resolve_model_class",
+        "_is_family_text_decoder",
+        "_remap_text_only_skip_modules",
+        "_get_text_only_config",
+        "_get_text_only_key_mapping",
+        "_apply_text_only_key_mapping",
+    ):
+        if name in funcs:
+            exec(funcs[name], ns)
+    return ns
+
+
+def _load_text_only_helper():
+    return _load_text_only_namespace()["_get_text_only_config"]
+
+
+def test_gemma3_vision_config_resolves_to_text_config():
+    transformers = pytest.importorskip("transformers")
+    helper = _load_text_only_helper()
+
+    config = transformers.Gemma3Config()
+    text_config = helper(config, "google/gemma-3-27b-it")
+
+    assert isinstance(text_config, transformers.Gemma3TextConfig)
+    assert text_config.model_type == "gemma3_text"
+    model_class = transformers.AutoModelForCausalLM._model_mapping[type(text_config)]
+    assert model_class.__name__ == "Gemma3ForCausalLM"
+
+
+def test_text_only_helper_rejects_configs_without_text_submodel():
+    helper = _load_text_only_helper()
+
+    class VisionOnlyConfig:
+        vision_config = object()
+
+    with pytest.raises(ValueError, match = "Cannot load vision-only as text-only"):
+        helper(VisionOnlyConfig(), "vision-only")
+
+
+def test_fast_language_model_forwards_text_only_to_fast_model():
+    source = _source(LOADER_PATH)
+    method = _class_method(ast.parse(source), "FastLanguageModel", "from_pretrained")
+
+    # text_only defaults False (opt-in, not forced True), and both FastModel
+    # delegations forward it.
+    text_only_default = _param_default(method, "text_only")
+    assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False
+
+    fast_model_calls = [
+        node
+        for node in ast.walk(method)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr == "from_pretrained"
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "FastModel"
+    ]
+    assert len(fast_model_calls) == 2
+    for call in fast_model_calls:
+        kw = [k for k in call.keywords if k.arg == "text_only"]
+        assert len(kw) == 1
+        assert isinstance(kw[0].value, ast.Name) and kw[0].value.id == "text_only"
+
+
+def test_fast_model_text_only_does_not_override_explicit_auto_model():
+    # AST-based so formatting/refactors that keep the structure do not break it.
+    source = _source(LOADER_PATH)
+    method = _class_method(ast.parse(source), "FastModel", "from_pretrained")
+
+    text_only_default = _param_default(method, "text_only")
+    assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False
+
+    # load_text_only is text_only AND a check that the caller did not pass auto_model.
+    def _is_guarded_bool(value):
+        names = _names_in(value)
+        has_none_check = any(
+            isinstance(n, ast.Compare) and any(isinstance(op, (ast.Is, ast.IsNot)) for op in n.ops)
+            for n in ast.walk(value)
+        )
+        return "text_only" in names and "auto_model" in names and has_none_check
+
+    assert _assigns_name(method, "load_text_only", _is_guarded_bool)
+
+    assert _calls_function(method, "_get_text_only_config")
+
+    def _forwards_kwarg(node):
+        return any(
+            isinstance(n, ast.Call)
+            and any(
+                kw.arg == "text_only"
+                and isinstance(kw.value, ast.Name)
+                and kw.value.id == "load_text_only"
+                for kw in n.keywords
+            )
+            for n in ast.walk(node)
+        )
+
+    assert _forwards_kwarg(method)
+    # Falls back to the full model unless the family has its own text decoder.
+    assert _calls_function(method, "_is_family_text_decoder")
+    assert _assigns_name(
+        method,
+        "load_text_only",
+        lambda v: isinstance(v, ast.Constant) and v.value is False,
+    )
+
+
+def test_fast_base_model_text_only_bypasses_vision_auto_model():
+    source = _source(VISION_PATH)
+    method = _class_method(ast.parse(source), "FastBaseModel", "from_pretrained")
+
+    text_only_default = _param_default(method, "text_only")
+    assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False
+
+    assert _assigns_name(
+        method,
+        "auto_model",
+        lambda v: isinstance(v, ast.Name) and v.id == "AutoModelForCausalLM",
+    )
+    # Text-only path: strip config, apply the family guard, inject the key remap.
+    assert _calls_function(method, "_get_text_only_config")
+    assert _calls_function(method, "_is_family_text_decoder")
+    assert _calls_function(method, "_apply_text_only_key_mapping")
+
+
+def test_gemma3_text_only_model_class_resolves_and_has_no_vision_tower():
+    """Tiny end-to-end: build a Gemma3 text-only config, instantiate the
+    matching model class with shrunken hidden sizes, assert it has the
+    text language model attributes and no vision tower attribute.
+
+    This is the integration check the AST-only tests were missing -- it
+    proves the text-only routing actually produces a model that can be
+    instantiated and that the resulting model is purely text. We use
+    shrunken hidden sizes so the test is fast and CPU-only.
+    """
+    transformers = pytest.importorskip("transformers")
+    helper = _load_text_only_helper()
+
+    full_config = transformers.Gemma3Config()
+    text_config = helper(full_config, "google/gemma-3-27b-it")
+
+    # Shrink for a cheap CPU instantiation; keep the shape attrs read at construction.
+    text_config.num_hidden_layers = 1
+    text_config.hidden_size = 32
+    text_config.intermediate_size = 32
+    text_config.num_attention_heads = 2
+    text_config.num_key_value_heads = 1
+    text_config.head_dim = 16
+    text_config.vocab_size = 128
+
+    model_class = transformers.AutoModelForCausalLM._model_mapping[type(text_config)]
+    model = model_class(text_config)
+
+    # Positive checks: text language model surface is present.
+    assert hasattr(model, "lm_head"), "text-only Gemma3 model should expose lm_head"
+
+    # Negative checks: no vision tower / multimodal projector remains.
+    assert not hasattr(
+        model, "vision_tower"
+    ), "text-only Gemma3 model should not have a vision_tower"
+    assert not hasattr(
+        model, "multi_modal_projector"
+    ), "text-only Gemma3 model should not have a multi_modal_projector"
+
+
+def test_helper_defined_once_in_utils_and_imported():
+    # _get_text_only_config is defined only in _utils and imported by loader + vision.
+    def _defines(path):
+        return any(
+            isinstance(n, ast.FunctionDef) and n.name == "_get_text_only_config"
+            for n in ast.parse(_source(path)).body
+        )
+
+    def _imports(path):
+        return any(
+            isinstance(n, ast.ImportFrom)
+            and n.module == "_utils"
+            and any(a.name == "_get_text_only_config" for a in n.names)
+            for n in ast.walk(ast.parse(_source(path)))
+        )
+
+    assert _defines(UTILS_PATH)
+    assert not _defines(LOADER_PATH) and _imports(LOADER_PATH)
+    assert not _defines(VISION_PATH) and _imports(VISION_PATH)
+
+
+def _load_util_func(name):
+    ns = _load_text_only_namespace()
+    if name not in ns:
+        raise AssertionError(f"{name} not found")
+    return ns[name]
+
+
+def test_text_only_guard_predicate_across_vlm_families():
+    # Text-only is taken only when the resolved class remaps VLM weights.
+    transformers = pytest.importorskip("transformers")
+    from transformers import AutoModelForCausalLM
+
+    resolve = _load_util_func("resolve_model_class")
+    is_family = _load_util_func("_is_family_text_decoder")
+    helper = _load_text_only_helper()
+
+    def takes_text_only(cfg):
+        text = helper(cfg, "x")
+        return resolve(AutoModelForCausalLM, text) is not None and is_family(
+            getattr(cfg, "model_type", ""), getattr(text, "model_type", "")
+        )
+
+    # Dedicated text decoder remaps language_model.* -> strip vision.
+    assert takes_text_only(transformers.Gemma3Config()) is True
+
+    # No text class (Qwen2-VL/Mllama) or a generic reused decoder that would
+    # load random weights (Llava/PaliGemma/Idefics3/InternVL) -> keep full model.
+    for name in [
+        "Qwen2VLConfig",
+        "Qwen2_5_VLConfig",
+        "MllamaConfig",
+        "LlavaConfig",
+        "PaliGemmaConfig",
+        "Idefics3Config",
+        "InternVLConfig",
+    ]:
+        cfg_cls = getattr(transformers, name, None)
+        if cfg_cls is None:
+            continue
+        assert takes_text_only(cfg_cls()) is False, name
+
+
+def test_text_only_helper_preserves_quantization_config():
+    # quantization_config must survive the strip so pre-quantized repos still load. A
+    # sentinel object avoids a bitsandbytes dependency on transformers 4.51.3.
+    transformers = pytest.importorskip("transformers")
+    helper = _load_text_only_helper()
+    config = transformers.Gemma3Config()
+    sentinel = object()
+    config.quantization_config = sentinel
+    text_config = helper(config, "google/gemma-3-27b-it")
+    assert getattr(text_config, "quantization_config", None) is sentinel
+    # The parent's shared text sub-config must not be mutated by the carry-over.
+    assert getattr(config.get_text_config(), "quantization_config", None) is None
+
+
+def test_text_only_key_mapping_targets_published_prefixes():
+    # The mapping must remap the published VLM decoder prefixes and only apply on
+    # transformers >=5 (on 4.x base_model_prefix handles it and a mapping hurts).
+    transformers = pytest.importorskip("transformers")
+    get_key_mapping = _load_util_func("_get_text_only_key_mapping")
+    mapping = get_key_mapping(transformers.Gemma3Config(), transformers.Gemma3TextConfig())
+    if int(transformers.__version__.split(".")[0]) < 5:
+        assert mapping is None
+    else:
+        assert isinstance(mapping, dict)
+        assert mapping.get(r"^language_model\.model\.") == "model."  # gemma3
+        assert mapping.get(r"^model\.language_model\.") == "model."  # gemma3n
+        assert mapping.get(r"^language_model\.lm_head\.") == "lm_head."
+
+
+def test_gemma3_text_only_loads_real_language_weights_from_vlm_checkpoint(tmp_path):
+    # Regression for PR #5816: text-only loading of a Gemma 3 VLM checkpoint must load the
+    # real language weights, not random ones. Fails on tf >=5 without the key_mapping fix.
+    transformers = pytest.importorskip("transformers")
+    torch = pytest.importorskip("torch")
+    import shutil
+    from safetensors.torch import load_file, save_file
+
+    get_text_config = _load_text_only_helper()
+    get_key_mapping = _load_util_func("_get_text_only_key_mapping")
+
+    sentinel = 0.1234
+    text_cfg = transformers.Gemma3TextConfig(
+        hidden_size = 32,
+        intermediate_size = 64,
+        num_hidden_layers = 1,
+        num_attention_heads = 2,
+        num_key_value_heads = 1,
+        head_dim = 16,
+        vocab_size = 128,
+        max_position_embeddings = 128,
+        sliding_window = 64,
+    )
+    vision_cfg = transformers.SiglipVisionConfig(
+        hidden_size = 32,
+        intermediate_size = 64,
+        num_hidden_layers = 1,
+        num_attention_heads = 2,
+        image_size = 16,
+        patch_size = 8,
+        num_channels = 3,
+    )
+    full_config = transformers.Gemma3Config(
+        text_config = text_cfg.to_dict(),
+        vision_config = vision_cfg.to_dict(),
+    )
+    full_model = transformers.Gemma3ForConditionalGeneration(full_config)
+
+    state = full_model.state_dict()
+    text_q = [
+        k
+        for k in state
+        if "language_model" in k
+        and "vision" not in k
+        and k.endswith("layers.0.self_attn.q_proj.weight")
+    ]
+    assert text_q, [k for k in state if "q_proj" in k][:5]
+    with torch.no_grad():
+        for k in text_q:
+            state[k].fill_(sentinel)
+
+    save_dir = tmp_path / "vlm"
+    full_model.save_pretrained(save_dir, safe_serialization = True)
+
+    # tf >=5 saves under an outer "model." prefix; strip it to reproduce the real
+    # language_model.model.* layout the published Gemma 3 checkpoints use.
+    real_dir = tmp_path / "real"
+    real_dir.mkdir()
+    weights = {}
+    for f in save_dir.glob("*.safetensors"):
+        weights.update(load_file(str(f)))
+    for f in save_dir.glob("*.bin"):
+        weights.update(torch.load(f, map_location = "cpu", weights_only = True))
+    weights = {
+        (k[len("model.") :] if k.startswith("model.") else k): v.contiguous()
+        for k, v in weights.items()
+    }
+    for p in save_dir.iterdir():
+        if not p.name.endswith((".safetensors", ".bin", ".index.json")):
+            shutil.copy(p, real_dir / p.name)
+    save_file(weights, str(real_dir / "model.safetensors"))
+
+    text_config = get_text_config(full_config, "google/gemma-3-27b-it")
+    load_kwargs = {}
+    key_mapping = get_key_mapping(full_config, text_config)
+    if key_mapping is not None:
+        load_kwargs["key_mapping"] = key_mapping
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        real_dir,
+        config = text_config,
+        dtype = torch.float32,
+        local_files_only = True,
+        **load_kwargs,
+    )
+
+    loaded = model.state_dict()
+    q_key = [k for k in loaded if k.endswith("model.layers.0.self_attn.q_proj.weight")]
+    assert q_key, "text decoder q_proj weight missing from the loaded model"
+    assert float(loaded[q_key[0]].flatten()[0]) == pytest.approx(
+        sentinel
+    ), "text weights were randomly initialized instead of loaded from the checkpoint"
+    assert not any(
+        "vision_tower" in n for n, _ in model.named_modules()
+    ), "vision tower should be skipped on the text-only path"

tests/python/test_fast_sentence_transformer_redirect_lifecycle.py

@@ -18,8 +18,8 @@ import pytest
 
 
 def _stub_module(name: str) -> types.ModuleType:
-    # __spec__ must be set so importlib.util.find_spec(name) does not raise
-    # ValueError if a downstream test imports the real package.
+    # __spec__ set so find_spec(name) doesn't raise if a later test imports
+    # the real package.
     mod = types.ModuleType(name)
     mod.__spec__ = importlib.util.spec_from_loader(name, loader = None)
     return mod
@@ -62,7 +62,6 @@ class _RecordingTransformerOk:
 
     def __init__(self, model_name, **kwargs):
         from transformers import AutoModel, AutoProcessor, AutoTokenizer
-
         type(self).last_calls = {
             "model": AutoModel.from_pretrained(model_name),
             "processor": AutoProcessor.from_pretrained(model_name),
@@ -73,7 +72,6 @@ class _RecordingTransformerOk:
 class _RaisingTransformer:
     def __init__(self, *a, **kw):
         from transformers import AutoModel
-
         AutoModel.from_pretrained(a[0] if a else kw.get("model_name_or_path"))
         raise RuntimeError("simulated init failure")
 
@@ -129,18 +127,10 @@ def _build_driver(transformer_class):
             return model if is_requested_model_name(a, kw) else original_model(*a, **kw)
 
         def return_existing_tokenizer(*a, **kw):
-            return (
-                tokenizer
-                if is_requested_model_name(a, kw)
-                else original_tokenizer(*a, **kw)
-            )
+            return tokenizer if is_requested_model_name(a, kw) else original_tokenizer(*a, **kw)
 
         def return_existing_processor(*a, **kw):
-            return (
-                tokenizer
-                if is_requested_model_name(a, kw)
-                else original_processor(*a, **kw)
-            )
+            return tokenizer if is_requested_model_name(a, kw) else original_processor(*a, **kw)
 
         try:
             AutoModel.from_pretrained = return_existing_model
@@ -190,7 +180,6 @@ def test_redirect_passes_through_for_other_model_names():
 
         def __init__(self, model_name, **kw):
             from transformers import AutoModel
-
             type(self).captured = AutoModel.from_pretrained("some-other/aux-model")
 
     driver, *_ = _build_driver(_OtherNameTransformer)
@@ -210,7 +199,6 @@ def test_is_requested_model_name_handles_pathlib_path(tmp_path):
 
         def __init__(self, model_name, **kw):
             from transformers import AutoModel
-
             type(self).last_calls = AutoModel.from_pretrained(pathlib.Path(model_name))
 
     driver, *_ = _build_driver(_PathTransformer)
@@ -228,7 +216,6 @@ def test_is_requested_model_name_trailing_slash_local_path(tmp_path):
 
         def __init__(self, model_name, **kw):
             from transformers import AutoModel
-
             type(self).last_calls = AutoModel.from_pretrained(str(target) + "/")
 
     driver, *_ = _build_driver(_SlashTransformer)
@@ -243,7 +230,6 @@ def test_is_requested_model_name_returns_false_when_no_identifier():
     class _NoNameTransformer:
         def __init__(self, model_name, **kw):
             from transformers import AutoModel
-
             captured["args"] = AutoModel.from_pretrained(some_other_kwarg = "x")
 
     driver, *_ = _build_driver(_NoNameTransformer)

tests/python/test_flash_attn_install_python_stack.py

@@ -33,64 +33,42 @@ class TestHasBlackwellGpu:
 
     def test_returns_true_for_sm_100(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
-            mock.patch.object(
-                wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+            mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")),
         ):
             assert wheel_utils.has_blackwell_gpu() is True
 
     def test_returns_true_for_sm_120(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
-            mock.patch.object(
-                wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+            mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")),
         ):
             assert wheel_utils.has_blackwell_gpu() is True
 
     def test_returns_true_for_sm_121(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
-            mock.patch.object(
-                wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+            mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")),
         ):
             assert wheel_utils.has_blackwell_gpu() is True
 
     def test_returns_false_for_sm_90(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
-            mock.patch.object(
-                wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+            mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")),
         ):
             assert wheel_utils.has_blackwell_gpu() is False
 
     def test_returns_false_for_sm_89(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
-            mock.patch.object(
-                wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+            mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")),
         ):
             assert wheel_utils.has_blackwell_gpu() is False
 
     def test_mixed_gpus_with_one_blackwell_returns_true(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
             mock.patch.object(
                 wheel_utils.subprocess,
                 "run",
@@ -101,9 +79,7 @@ class TestHasBlackwellGpu:
 
     def test_returns_false_when_nvidia_smi_fails(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
             mock.patch.object(
                 wheel_utils.subprocess,
                 "run",
@@ -114,9 +90,7 @@ class TestHasBlackwellGpu:
 
     def test_returns_false_on_subprocess_timeout(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
             mock.patch.object(
                 wheel_utils.subprocess,
                 "run",
@@ -127,9 +101,7 @@ class TestHasBlackwellGpu:
 
     def test_returns_false_on_malformed_output(self):
         with (
-            mock.patch.object(
-                wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
-            ),
+            mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
             mock.patch.object(
                 wheel_utils.subprocess,
                 "run",
@@ -161,10 +133,7 @@ class TestFlashAttnWheelSelection:
         )
         assert url is not None
         assert "v2.8.1" in url
-        assert (
-            "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
-            in url
-        )
+        assert "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
 
     def test_missing_cuda_major_disables_wheel_lookup(self):
         assert (
@@ -262,7 +231,11 @@ class TestEnsureFlashAttn:
         step_messages: list[tuple[str, str]] = []
         printed_failures: list[str] = []
 
-        def fake_step(label: str, value: str, color_fn = None):
+        def fake_step(
+            label: str,
+            value: str,
+            color_fn = None,
+        ):
             step_messages.append((label, value))
 
         with (
@@ -313,7 +286,11 @@ class TestEnsureFlashAttn:
     def test_wheel_missing_skips_install_at_setup_time(self):
         step_messages: list[tuple[str, str]] = []
 
-        def fake_step(label: str, value: str, color_fn = None):
+        def fake_step(
+            label: str,
+            value: str,
+            color_fn = None,
+        ):
             step_messages.append((label, value))
 
         with (
@@ -339,10 +316,7 @@ class TestEnsureFlashAttn:
             ips._ensure_flash_attn()
 
         mock_install_wheel.assert_not_called()
-        assert (
-            "warning",
-            "No published flash-attn prebuilt wheel found",
-        ) in step_messages
+        assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages
 
     def test_skip_env_disables_setup_install(self):
         with (
@@ -362,7 +336,11 @@ class TestEnsureFlashAttn:
     def test_blackwell_gpu_skips_install_with_warning(self):
         step_messages: list[tuple[str, str]] = []
 
-        def fake_step(label: str, value: str, color_fn = None):
+        def fake_step(
+            label: str,
+            value: str,
+            color_fn = None,
+        ):
             step_messages.append((label, value))
 
         with (
@@ -379,14 +357,16 @@ class TestEnsureFlashAttn:
 
         mock_probe.assert_not_called()
         mock_install_wheel.assert_not_called()
-        assert any(
-            label == "warning" and "Blackwell" in msg for label, msg in step_messages
-        )
+        assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
 
     def test_blackwell_gpu_on_windows_emits_blackwell_warning(self):
         step_messages: list[tuple[str, str]] = []
 
-        def fake_step(label: str, value: str, color_fn = None):
+        def fake_step(
+            label: str,
+            value: str,
+            color_fn = None,
+        ):
             step_messages.append((label, value))
 
         with (
@@ -403,14 +383,16 @@ class TestEnsureFlashAttn:
 
         mock_probe.assert_not_called()
         mock_install_wheel.assert_not_called()
-        assert any(
-            label == "warning" and "Blackwell" in msg for label, msg in step_messages
-        )
+        assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
 
     def test_non_blackwell_windows_does_not_emit_blackwell_warning(self):
         step_messages: list[tuple[str, str]] = []
 
-        def fake_step(label: str, value: str, color_fn = None):
+        def fake_step(
+            label: str,
+            value: str,
+            color_fn = None,
+        ):
             step_messages.append((label, value))
 
         with (
@@ -453,9 +435,7 @@ class TestInstallPythonStackFlashAttnIntegration:
             mock.patch("subprocess.run", side_effect = fake_run),
             mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False),
             mock.patch.object(ips, "_has_rocm_gpu", return_value = False),
-            mock.patch.object(
-                ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")
-            ),
+            mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")),
             mock.patch("pathlib.Path.is_dir", return_value = True),
             mock.patch("pathlib.Path.is_file", return_value = True),
             mock.patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}, clear = False),

tests/python/test_gpu_init_ldconfig_guard.py

@@ -19,9 +19,7 @@ def _find_geteuid_guard(tree: ast.AST):
 def test_gpu_init_has_geteuid_guard():
     tree = ast.parse(GPU_INIT.read_text())
     guard = _find_geteuid_guard(tree)
-    assert (
-        guard is not None
-    ), "_gpu_init.py must guard ldconfig recovery on os.geteuid()"
+    assert guard is not None, "_gpu_init.py must guard ldconfig recovery on os.geteuid()"
 
 
 def test_ldconfig_calls_only_inside_geteuid_guard():

tests/python/test_install_python_stack.py

@@ -10,12 +10,11 @@ from unittest import mock
 
 import pytest
 
-# Add the studio directory so we can import install_python_stack
+# Add the studio directory so we can import install_python_stack.
 STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
 sys.path.insert(0, str(STUDIO_DIR))
 
-# _build_uv_cmd lives at module level; import after path setup.
-# We need to mock parts of the module that do work at import time.
+# Import after path setup.
 import install_python_stack as ips
 
 

tests/python/test_no_torch_filtering.py

@@ -156,9 +156,7 @@ class TestFilterRequirements:
         )
         # First filter Windows packages, then NO_TORCH packages
         intermediate = ips._filter_requirements(req, ips.WINDOWS_SKIP_PACKAGES)
-        result = ips._filter_requirements(
-            Path(intermediate), ips.NO_TORCH_SKIP_PACKAGES
-        )
+        result = ips._filter_requirements(Path(intermediate), ips.NO_TORCH_SKIP_PACKAGES)
         lines = Path(result).read_text(encoding = "utf-8").splitlines()
         non_blank = [l.strip() for l in lines if l.strip()]
         assert non_blank == [
@@ -177,9 +175,7 @@ class TestFilterRequirements:
         result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
         lines = Path(result).read_text(encoding = "utf-8").splitlines()
         non_blank = [l.strip() for l in lines if l.strip()]
-        assert non_blank == [
-            "numpy"
-        ], f"VCS URL line should be filtered, got: {non_blank}"
+        assert non_blank == ["numpy"], f"VCS URL line should be filtered, got: {non_blank}"
 
     def test_env_marker_line_filtered(self, tmp_path):
         """Package lines with env markers are still filtered by prefix."""
@@ -193,9 +189,7 @@ class TestFilterRequirements:
         result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
         lines = Path(result).read_text(encoding = "utf-8").splitlines()
         non_blank = [l.strip() for l in lines if l.strip()]
-        assert non_blank == [
-            "numpy"
-        ], f"Env marker line should be filtered, got: {non_blank}"
+        assert non_blank == ["numpy"], f"Env marker line should be filtered, got: {non_blank}"
 
     def test_git_plus_url_not_over_matched(self, tmp_path):
         """A git+ URL whose path contains a skip package name but does NOT start with it."""
@@ -247,9 +241,7 @@ class TestRealRequirementsFiltering:
         expected = [
             l
             for l in original
-            if not any(
-                l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES
-            )
+            if not any(l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES)
         ]
         assert filtered == expected, (
             f"Filtered extras.txt should match expected.\n"
@@ -259,9 +251,7 @@ class TestRealRequirementsFiltering:
 
     def test_extras_no_deps_txt_torchcodec_and_dlpack_removed(self):
         """extras-no-deps.txt: torchcodec and torch-c-dlpack-ext must be removed."""
-        result = ips._filter_requirements(
-            EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES
-        )
+        result = ips._filter_requirements(EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES)
         filtered = self._non_blank_non_comment(Path(result))
         original = self._non_blank_non_comment(EXTRAS_NO_DEPS_TXT)
 
@@ -273,9 +263,7 @@ class TestRealRequirementsFiltering:
         expected = [
             l
             for l in original
-            if not any(
-                l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES
-            )
+            if not any(l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES)
         ]
         assert filtered == expected
 
@@ -291,9 +279,7 @@ class TestRealRequirementsFiltering:
 
     def test_extras_no_deps_txt_trl_preserved(self):
         """trl should survive NO_TORCH filtering in extras-no-deps.txt."""
-        result = ips._filter_requirements(
-            EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES
-        )
+        result = ips._filter_requirements(EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES)
         filtered_text = Path(result).read_text(encoding = "utf-8").lower()
         assert "trl" in filtered_text, "trl should survive NO_TORCH filtering"
 
@@ -370,7 +356,6 @@ class TestIsMacosConstant:
 
     def test_is_macos_matches_platform(self):
         import sys
-
         expected = sys.platform == "darwin"
         assert ips.IS_MACOS is expected
 
@@ -405,9 +390,7 @@ class TestInstallPythonStackSubprocessMock:
         captured_cmds: list[list[str]] = []
 
         def mock_run(cmd, **kw):
-            captured_cmds.append(
-                list(cmd) if isinstance(cmd, (list, tuple)) else [str(cmd)]
-            )
+            captured_cmds.append(list(cmd) if isinstance(cmd, (list, tuple)) else [str(cmd)])
             return subprocess.CompletedProcess(cmd, 0, b"", b"")
 
         env = {"SKIP_STUDIO_BASE": "1"} if skip_base else {}
@@ -424,9 +407,7 @@ class TestInstallPythonStackSubprocessMock:
             mock.patch.object(ips, "_has_rocm_gpu", return_value = False),
             mock.patch("subprocess.run", side_effect = mock_run),
             mock.patch.object(ips, "_bootstrap_uv", return_value = True),
-            mock.patch.object(
-                ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")
-            ),
+            mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")),
             mock.patch("pathlib.Path.is_dir", return_value = True),
             mock.patch("pathlib.Path.is_file", return_value = True),
         ):
@@ -469,9 +450,7 @@ class TestInstallPythonStackSubprocessMock:
         has_extras_nd = self._cmds_contain_file(cmds, "extras-no-deps.txt") or any(
             "-r" in cmd and "tmp" in cmd.lower() for cmd in cmds
         )
-        assert (
-            has_extras_nd
-        ), "extras-no-deps.txt (or its filtered temp) should be called"
+        assert has_extras_nd, "extras-no-deps.txt (or its filtered temp) should be called"
 
     # -- IS_WINDOWS=True + NO_TORCH=True (stacked) --
 
@@ -570,17 +549,13 @@ class TestOverridesSkip:
     def test_no_torch_guard_exists_in_source(self):
         """The install_python_stack source must contain a NO_TORCH guard around overrides."""
         source = Path(ips.__file__).read_text(encoding = "utf-8")
-        assert (
-            "if NO_TORCH:" in source
-        ), "NO_TORCH guard not found in install_python_stack.py"
+        assert "if NO_TORCH:" in source, "NO_TORCH guard not found in install_python_stack.py"
 
     def test_overrides_skipped_when_no_torch(self):
         """With NO_TORCH=True on the module, pip_install should NOT be called for overrides."""
         source = Path(ips.__file__).read_text(encoding = "utf-8")
         overrides_match = re.search(r"if NO_TORCH:.*?overrides", source, re.DOTALL)
-        assert (
-            overrides_match is not None
-        ), "Expected NO_TORCH conditional before overrides install"
+        assert overrides_match is not None, "Expected NO_TORCH conditional before overrides install"
 
 
 # ── install.sh --no-torch flag tests ──────────────────────────────────
@@ -599,33 +574,21 @@ class TestInstallShNoTorchFlag:
 
     def test_no_torch_flag_in_case_statement(self):
         """--no-torch must appear in the flag parser case statement."""
-        assert (
-            "--no-torch)" in self.source
-        ), "--no-torch not found in install.sh flag parser"
+        assert "--no-torch)" in self.source, "--no-torch not found in install.sh flag parser"
 
     def test_no_torch_flag_variable_initialized(self):
         """_NO_TORCH_FLAG must be initialized to false."""
-        assert (
-            "_NO_TORCH_FLAG=false" in self.source
-        ), "_NO_TORCH_FLAG=false not found in install.sh"
+        assert "_NO_TORCH_FLAG=false" in self.source, "_NO_TORCH_FLAG=false not found in install.sh"
 
     def test_skip_torch_variable_exists(self):
         """SKIP_TORCH variable must be defined."""
-        assert (
-            "SKIP_TORCH=false" in self.source
-        ), "SKIP_TORCH=false not found in install.sh"
-        assert (
-            "SKIP_TORCH=true" in self.source
-        ), "SKIP_TORCH=true not found in install.sh"
+        assert "SKIP_TORCH=false" in self.source, "SKIP_TORCH=false not found in install.sh"
+        assert "SKIP_TORCH=true" in self.source, "SKIP_TORCH=true not found in install.sh"
 
     def test_skip_torch_driven_by_flag_and_mac_intel(self):
         """SKIP_TORCH must check both _NO_TORCH_FLAG and MAC_INTEL."""
-        assert (
-            "_NO_TORCH_FLAG" in self.source
-        ), "_NO_TORCH_FLAG not referenced in SKIP_TORCH logic"
-        assert (
-            "MAC_INTEL" in self.source
-        ), "MAC_INTEL not referenced in SKIP_TORCH logic"
+        assert "_NO_TORCH_FLAG" in self.source, "_NO_TORCH_FLAG not referenced in SKIP_TORCH logic"
+        assert "MAC_INTEL" in self.source, "MAC_INTEL not referenced in SKIP_TORCH logic"
 
     def test_unsloth_no_torch_uses_skip_torch(self):
         """UNSLOTH_NO_TORCH must reference $SKIP_TORCH, not $MAC_INTEL."""
@@ -633,18 +596,12 @@ class TestInstallShNoTorchFlag:
 
         matches = re.findall(r'UNSLOTH_NO_TORCH="\$(\w+)"', self.source)
         for var in matches:
-            assert (
-                var == "SKIP_TORCH"
-            ), f"UNSLOTH_NO_TORCH references ${var} instead of $SKIP_TORCH"
+            assert var == "SKIP_TORCH", f"UNSLOTH_NO_TORCH references ${var} instead of $SKIP_TORCH"
 
     def test_cpu_hint_message_exists(self):
         """CPU hint message must exist in install.sh."""
-        assert (
-            "No GPU detected" in self.source
-        ), "CPU hint message not found in install.sh"
-        assert (
-            "--no-torch" in self.source
-        ), "--no-torch suggestion not found in CPU hint"
+        assert "No GPU detected" in self.source, "CPU hint message not found in install.sh"
+        assert "--no-torch" in self.source, "--no-torch suggestion not found in CPU hint"
 
     def test_no_torch_flag_parsing_subprocess(self):
         """--no-torch flag sets _NO_TORCH_FLAG=true (subprocess test)."""

tests/python/test_orpo_processor_text_tokenizer.py

@@ -34,7 +34,12 @@ class _Tokenizer:
     def __init__(self):
         self.calls = []
 
-    def __call__(self, text, add_special_tokens = False, **kwargs):
+    def __call__(
+        self,
+        text,
+        add_special_tokens = False,
+        **kwargs,
+    ):
         self.calls.append((text, add_special_tokens, kwargs))
         ids = [ord(c) % 31 + 3 for c in text]
         return {"input_ids": ids, "attention_mask": [1] * len(ids)}
@@ -60,7 +65,11 @@ class _Trainer:
         self.padding_value = 0
 
 
-def _exec_rewritten(function_name, source, extra_ns = None):
+def _exec_rewritten(
+    function_name,
+    source,
+    extra_ns = None,
+):
     rewriter = _load_orpo_rewriter()
     rewritten = rewriter(function_name, source)
     ns = {} if extra_ns is None else dict(extra_ns)

tests/python/test_patch_trl_rl_trainers_defensive.py

@@ -67,3 +67,19 @@ def test_wrapper_swallows_impl_exception(monkeypatch):
 
     monkeypatch.setattr(_rl, "_patch_trl_rl_trainers_impl", _boom)
     assert _rl._patch_trl_rl_trainers("sft_trainer") is None
+
+
+def test_grpo_config_sibling_module_import_is_patched(tmp_path):
+    import unsloth  # noqa: F401
+    from trl import GRPOConfig as top_config
+    from trl.trainer import GRPOConfig as trainer_config
+    from trl.trainer.grpo_config import GRPOConfig as config_module_config
+    from trl.trainer.grpo_trainer import GRPOConfig as trainer_module_config
+
+    assert top_config is trainer_config
+    assert top_config is trainer_module_config
+    assert top_config is config_module_config
+
+    args = config_module_config(output_dir = str(tmp_path))
+    assert hasattr(args, "unsloth_grpo_mini_batch")
+    assert args.unsloth_grpo_mini_batch is None

tests/python/test_studio_import_no_torch.py

@@ -23,15 +23,9 @@ from pathlib import Path
 import pytest
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
-DATA_COLLATORS = (
-    REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "data_collators.py"
-)
-CHAT_TEMPLATES = (
-    REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "chat_templates.py"
-)
-FORMAT_CONVERSION = (
-    REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "format_conversion.py"
-)
+DATA_COLLATORS = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "data_collators.py"
+CHAT_TEMPLATES = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "chat_templates.py"
+FORMAT_CONVERSION = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "format_conversion.py"
 
 
 def _has_uv() -> bool:
@@ -67,14 +61,11 @@ def no_torch_venv(request, tmp_path_factory):
     if venv_python is None:
         pytest.skip(f"Could not create Python {py_version} venv")
 
-    # Verify torch is NOT importable
     check = subprocess.run(
         [str(venv_python), "-c", "import torch"],
         capture_output = True,
     )
-    assert (
-        check.returncode != 0
-    ), f"torch should NOT be importable in fresh {py_version} venv"
+    assert check.returncode != 0, f"torch should NOT be importable in fresh {py_version} venv"
 
     return str(venv_python)
 
@@ -223,9 +214,7 @@ class TestDataCollatorsNoTorchVenv:
             capture_output = True,
             timeout = 30,
         )
-        assert (
-            result.returncode == 0
-        ), f"DeepSeekOCRDataCollator failed:\n{result.stderr.decode()}"
+        assert result.returncode == 0, f"DeepSeekOCRDataCollator failed:\n{result.stderr.decode()}"
         assert b"OK: DeepSeekOCRDataCollator instantiated" in result.stdout
 
     def test_dataclass_vlm_collator_instantiable(self, no_torch_venv):
@@ -246,9 +235,7 @@ class TestDataCollatorsNoTorchVenv:
             capture_output = True,
             timeout = 30,
         )
-        assert (
-            result.returncode == 0
-        ), f"VLMDataCollator failed:\n{result.stderr.decode()}"
+        assert result.returncode == 0, f"VLMDataCollator failed:\n{result.stderr.decode()}"
         assert b"OK: VLMDataCollator instantiated" in result.stdout
 
 
@@ -529,12 +516,9 @@ class TestNegativeControls:
                 capture_output = True,
                 timeout = 30,
             )
+            assert result.returncode != 0, "Expected failure when 'import torch' is prepended"
             assert (
-                result.returncode != 0
-            ), "Expected failure when 'import torch' is prepended"
-            assert (
-                b"ModuleNotFoundError" in result.stderr
-                or b"ImportError" in result.stderr
+                b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
             ), f"Expected ImportError, got:\n{result.stderr.decode()}"
         finally:
             os.unlink(temp_file)
@@ -577,6 +561,4 @@ class TestNegativeControls:
             timeout = 30,
         )
         assert result.returncode != 0, "import torch should fail in no-torch venv"
-        assert (
-            b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
-        )
+        assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr

tests/python/test_tokenizers_and_torch_constraint.py

@@ -1,7 +1,6 @@
-"""
-Tests for two install fixes:
-  1. tokenizers added to no-torch-runtime.txt  (prevents AutoConfig crash)
-  2. TORCH_CONSTRAINT variable in install.sh    (arm64 macOS + py313+ -> torch>=2.6)
+"""Tests for two install fixes:
+1. tokenizers in no-torch-runtime.txt (prevents AutoConfig crash)
+2. TORCH_CONSTRAINT in install.sh (arm64 macOS + py313+ -> torch>=2.6)
 """
 
 from __future__ import annotations
@@ -18,9 +17,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent  # tests/
 _REPO_ROOT = _TESTS_DIR.parent  # unsloth/
 _INSTALL_SH = _REPO_ROOT / "install.sh"
 _INSTALL_PS1 = _REPO_ROOT / "install.ps1"
-_NO_TORCH_RT = (
-    _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
-)
+_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
 
 
 def _read(path: pathlib.Path) -> str:
@@ -45,30 +42,23 @@ class TestStructuralTokenizers:
     def test_tokenizers_present(self):
         """tokenizers must be a standalone package line."""
         pkgs = _lines(_NO_TORCH_RT)
-        bare_names = [
-            p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
-        ]
+        bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
         assert "tokenizers" in bare_names
 
     def test_tokenizers_before_transformers(self):
         """tokenizers should appear before transformers (install order intent)."""
         pkgs = _lines(_NO_TORCH_RT)
-        bare_names = [
-            p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
-        ]
+        bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
         idx_tok = bare_names.index("tokenizers")
         idx_tf = bare_names.index("transformers")
         assert idx_tok < idx_tf, (
-            f"tokenizers at index {idx_tok} should appear before "
-            f"transformers at index {idx_tf}"
+            f"tokenizers at index {idx_tok} should appear before " f"transformers at index {idx_tf}"
         )
 
     def test_torch_not_in_no_torch_file(self):
         """torch itself must NOT be listed in the no-torch requirements."""
         pkgs = _lines(_NO_TORCH_RT)
-        bare_names = [
-            p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
-        ]
+        bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
         assert "torch" not in bare_names
 
 
@@ -409,9 +399,7 @@ class TestE2ETokenizersFix:
         r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
         assert r.returncode == 0, f"Install failed: {r.stderr}"
 
-        result = self._run_python(
-            venv, "from transformers import AutoConfig; print('OK')"
-        )
+        result = self._run_python(venv, "from transformers import AutoConfig; print('OK')")
         assert (
             result.returncode == 0
         ), f"AutoConfig import failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
@@ -441,22 +429,15 @@ class TestE2ETokenizersFix:
         req_no_tokenizers = tmp_path / "no-tokenizers.txt"
         req_no_tokenizers.write_text(
             "\n".join(
-                line
-                for line in _read(_NO_TORCH_RT).splitlines()
-                if line.strip() != "tokenizers"
+                line for line in _read(_NO_TORCH_RT).splitlines() if line.strip() != "tokenizers"
             ),
             encoding = "utf-8",
         )
         r = self._pip_install(venv, "--no-deps", "-r", str(req_no_tokenizers))
         assert r.returncode == 0, f"Install failed: {r.stderr}"
         result = self._run_python(venv, "from transformers import AutoConfig")
-        assert (
-            result.returncode != 0
-        ), "AutoConfig should fail without tokenizers installed"
-        assert (
-            "tokenizers" in result.stderr.lower()
-            or "ModuleNotFoundError" in result.stderr
-        )
+        assert result.returncode != 0, "AutoConfig should fail without tokenizers installed"
+        assert "tokenizers" in result.stderr.lower() or "ModuleNotFoundError" in result.stderr
 
 
 # ======================================================================
@@ -535,9 +516,7 @@ class TestE2EFullNoTorchSandbox:
         venv = self._create_venv(tmp_path, "full-no-torch")
         r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
         assert r.returncode == 0, f"Install failed: {r.stderr}"
-        result = self._run_python(
-            venv, "from transformers import AutoConfig; print('OK')"
-        )
+        result = self._run_python(venv, "from transformers import AutoConfig; print('OK')")
         assert (
             result.returncode == 0
         ), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"

tests/python/test_unsloth_run_tool_policy_resolver.py

@@ -141,15 +141,11 @@ class TestZeroHost:
 
 
 class TestIsExternalHost:
-    @pytest.mark.parametrize(
-        "host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"]
-    )
+    @pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"])
     def test_loopback_aliases_are_local(self, host):
         assert is_external_host(host) is False
 
-    @pytest.mark.parametrize(
-        "host", ["0.0.0.0", "::", "192.168.1.5", "10.0.0.1", "example.com"]
-    )
+    @pytest.mark.parametrize("host", ["0.0.0.0", "::", "192.168.1.5", "10.0.0.1", "example.com"])
     def test_non_loopback_is_external(self, host):
         assert is_external_host(host) is True
 

tests/qlora/test_hf_qlora_train_and_merge.py

@@ -91,9 +91,7 @@ if __name__ == "__main__":
         print(training_args)
         print(peft_config)
 
-    trainer = setup_trainer(
-        model, tokenizer, dataset, training_args, peft_config = peft_config
-    )
+    trainer = setup_trainer(model, tokenizer, dataset, training_args, peft_config = peft_config)
 
     with header_footer_context("Model"):
         print(type(model.model))

tests/run_all.sh

@@ -8,6 +8,7 @@ echo "=== Bash tests ==="
 sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
 sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
 sh "$TESTS_DIR/sh/test_torch_constraint.sh"
+sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh"
 
 echo ""
 echo "=== Python tests ==="

tests/saving/gpt-oss-merge/test_merged_model.py

@@ -39,12 +39,10 @@ inputs = merged_tokenizer.apply_chat_template(
     add_generation_prompt = True,
     return_tensors = "pt",
     return_dict = True,
-    reasoning_effort = "low",  # **NEW!** Set reasoning effort to low, medium or high
+    reasoning_effort = "low",  # low, medium or high
 ).to(merged_model.device)
 
-_ = merged_model.generate(
-    **inputs, max_new_tokens = 512, streamer = TextStreamer(merged_tokenizer)
-)
+_ = merged_model.generate(**inputs, max_new_tokens = 512, streamer = TextStreamer(merged_tokenizer))
 print("\n✅ Inference complete.")
 
 # --- Final Cleanup ---
@@ -54,7 +52,5 @@ torch.cuda.empty_cache()
 gc.collect()
 
 safe_remove_directory("./gpt-oss-finetuned-merged")
-safe_remove_directory(
-    "./unsloth_compiled_cache"
-)  # Clean up cache created by this process
+safe_remove_directory("./unsloth_compiled_cache")  # Clean up cache created by this process
 print("✅ Final cleanup complete. Exiting inference script.")

tests/saving/gpt-oss-merge/train_and_merge.py

@@ -1,4 +1,3 @@
-# train_and_merge.py
 from unsloth import FastLanguageModel
 from trl import SFTTrainer, SFTConfig
 from datasets import load_dataset
@@ -21,16 +20,14 @@ def safe_remove_directory(path):
         return False
 
 
-# This tokenizer will be used by the mapping function
+# Used by the mapping function below.
 tokenizer = None
 
 
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
-        tokenizer.apply_chat_template(
-            convo, tokenize = False, add_generation_prompt = False
-        )
+        tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
         for convo in convos
     ]
     return {"text": texts}
@@ -84,9 +81,7 @@ print("Fine-tuning complete.")
 
 # --- Merge and Save ---
 print("\n💾 Merging and saving the 16-bit model to './gpt-oss-finetuned-merged'...")
-model.save_pretrained_merged(
-    save_directory = "./gpt-oss-finetuned-merged", tokenizer = tokenizer
-)
+model.save_pretrained_merged(save_directory = "./gpt-oss-finetuned-merged", tokenizer = tokenizer)
 print("✅ Model merged and saved.")
 
 # --- Cleanup ---
@@ -96,7 +91,5 @@ torch.cuda.empty_cache()
 gc.collect()
 
 safe_remove_directory("./outputs")
-safe_remove_directory(
-    "./unsloth_compiled_cache"
-)  # Clean up the cache created by this process
+safe_remove_directory("./unsloth_compiled_cache")  # Clean up the cache created by this process
 print("✅ Cleanup complete. Exiting training script.")

tests/saving/language_models/test_merge_4bit_validation.py

@@ -16,9 +16,7 @@ from tests.utils.cleanup_utils import safe_remove_directory
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
-        tokenizer.apply_chat_template(
-            convo, tokenize = False, add_generation_prompt = False
-        )
+        tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
         for convo in convos
     ]
     return {"text": texts}
@@ -50,10 +48,7 @@ tokenizer = get_chat_template(
     chat_template = "llama-3.1",
 )
 
-# Load small dataset for quick training
-dataset_train = load_dataset(
-    "allenai/openassistant-guanaco-reformatted", split = "train[:100]"
-)
+dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train[:100]")
 dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
 
 print("✅ Base model loaded successfully!")
@@ -130,7 +125,6 @@ print(f"\n{'='*80}")
 print("🔍 PHASE 4: Loading 4bit Model and Second Fine-tuning")
 print(f"{'='*80}")
 
-# Clean up first model
 del model
 del tokenizer
 torch.cuda.empty_cache()
@@ -150,7 +144,7 @@ tokenizer_4bit = get_chat_template(
 
 print("✅ 4bit model loaded successfully!")
 
-# Add LoRA adapters to the 4bit model
+# Add LoRA adapters
 model_4bit = FastLanguageModel.get_peft_model(
     model_4bit,
     r = 16,
@@ -237,7 +231,6 @@ print(f"\n{'='*80}")
 print("🔍 CLEANUP")
 print(f"{'='*80}")
 
-# Cleanup
 safe_remove_directory("./outputs")
 safe_remove_directory("./outputs_4bit")
 safe_remove_directory("./unsloth_compiled_cache")

tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py

@@ -31,70 +31,61 @@ from tests.utils.perplexity_eval import (
 )
 
 
-# Define helper functions outside of main
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
-        tokenizer.apply_chat_template(
-            convo, tokenize = False, add_generation_prompt = False
-        )
+        tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
         for convo in convos
     ]
     return {"text": texts}
 
 
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+    result_queue,
+    load_in_4bit = False,
+    load_in_8bit = False,
+):
     """Load model and compute perplexity in subprocess"""
     from unsloth import FastLanguageModel
     from unsloth.chat_templates import get_chat_template
     from tests.utils.perplexity_eval import ppl_model
 
-    # Load model
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_llama_text_model",
         max_seq_length = 2048,
         load_in_4bit = load_in_4bit,
         load_in_8bit = load_in_8bit,
     )
-    # Set up tokenizer
     merged_tokenizer = get_chat_template(
         merged_tokenizer,
         chat_template = "llama-3.1",
     )
 
     # Load dataset fresh in subprocess
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
-    # Format the dataset
     def formatting_prompts_func(examples):
         convos = examples["messages"]
         texts = [
-            merged_tokenizer.apply_chat_template(
-                convo, tokenize = False, add_generation_prompt = False
-            )
+            merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
             for convo in convos
         ]
         return {"text": texts}
 
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
 
-    # Compute perplexity using the passed dataset
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # IMPORTANT: Convert to Python float if it's a tensor
+    # Coerce to a plain Python float for cross-process transfer
     if torch.is_tensor(ppl_value):
-        ppl_value = ppl_value.cpu().item()  # Move to CPU and convert to Python scalar
+        ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
-        ppl_value = ppl_value.item()  # Convert numpy or other array types
+        ppl_value = ppl_value.item()
     else:
-        ppl_value = float(ppl_value)  # Ensure it's a float
+        ppl_value = float(ppl_value)
 
-    # Return only the perplexity value
     result_queue.put(ppl_value)
 
-    # Clean up
     del merged_model
     del merged_tokenizer
     del dataset_ppl
@@ -102,7 +93,6 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
     gc.collect()
 
 
-# Main execution code should be wrapped in this guard
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
@@ -130,12 +120,8 @@ if __name__ == "__main__":
 
     from unsloth.chat_templates import standardize_sharegpt
 
-    dataset_train = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "train"
-    )
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
@@ -199,12 +185,10 @@ if __name__ == "__main__":
         response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
     )
 
-    # run training
     trainer_stats = trainer.train()
 
     add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
 
-    # saving and merging the model to local disk
     print("merge and save to local disk")
     model.save_pretrained_merged(
         save_directory = "./unsloth_out/merged_llama_text_model", tokenizer = tokenizer
@@ -216,7 +200,6 @@ if __name__ == "__main__":
     # torch.cuda.empty_cache()
     # gc.collect()
 
-    # load model from local disk and test
     print("Loading merged model in 4 bit for perplexity test")
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_llama_text_model",
@@ -253,7 +236,6 @@ if __name__ == "__main__":
 
     print_model_comparison()
 
-    # final cleanup
     safe_remove_directory("./outputs")
     safe_remove_directory("./unsloth_compiled_cache")
     safe_remove_directory("./unsloth_out")

tests/saving/language_models/test_merge_model_perplexity_mistral.py

@@ -30,28 +30,28 @@ from tests.utils.perplexity_eval import (
 )
 
 
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+    result_queue,
+    load_in_4bit = False,
+    load_in_8bit = False,
+):
     """Load model and compute perplexity in subprocess"""
     from unsloth import FastLanguageModel
     from tests.utils.perplexity_eval import ppl_model
 
-    # Load model
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_mistral_text_model",
         max_seq_length = 2048,
         load_in_4bit = load_in_4bit,
         load_in_8bit = load_in_8bit,
     )
-    # Set up tokenizer
     # merged_tokenizer = get_chat_template(
     #     merged_tokenizer,
     #     chat_template="llama-3.1",
     # )
 
-    # Load dataset fresh in subprocess
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    # Load dataset fresh in subprocess.
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
@@ -73,7 +73,6 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
         texts = []
 
         for conversation in examples["messages"]:
-            # Extract user message and assistant response
             user_message = ""
             assistant_message = ""
 
@@ -83,17 +82,12 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
                 elif turn["role"] == "assistant":
                     assistant_message = turn["content"]
 
-            # Store intermediate format
             instruction = "Complete the statement"
             instructions.append(instruction)
             inputs.append(user_message)
             outputs.append(assistant_message)
 
-            # Create formatted text
-            text = (
-                alpaca_prompt.format(instruction, user_message, assistant_message)
-                + EOS_TOKEN
-            )
+            text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
             texts.append(text)
 
         return {
@@ -105,21 +99,18 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
 
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
 
-    # Compute perplexity using the passed dataset
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # IMPORTANT: Convert to Python float if it's a tensor
+    # Convert to Python float if it's a tensor.
     if torch.is_tensor(ppl_value):
-        ppl_value = ppl_value.cpu().item()  # Move to CPU and convert to Python scalar
+        ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
-        ppl_value = ppl_value.item()  # Convert numpy or other array types
+        ppl_value = ppl_value.item()
     else:
-        ppl_value = float(ppl_value)  # Ensure it's a float
+        ppl_value = float(ppl_value)
 
-    # Return only the perplexity value
     result_queue.put(ppl_value)
 
-    # Clean up
     del merged_model
     del merged_tokenizer
     del dataset_ppl
@@ -127,7 +118,6 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
     gc.collect()
 
 
-# Main execution code should be wrapped in this guard
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
@@ -161,7 +151,6 @@ if __name__ == "__main__":
     ### Response:
     {}"""
 
-    # Define helper functions outside of main
     def formatting_prompts_func(examples):
         instructions = []
         inputs = []
@@ -169,7 +158,6 @@ if __name__ == "__main__":
         texts = []
 
         for conversation in examples["messages"]:
-            # Extract user message and assistant response
             user_message = ""
             assistant_message = ""
 
@@ -179,17 +167,12 @@ if __name__ == "__main__":
                 elif turn["role"] == "assistant":
                     assistant_message = turn["content"]
 
-            # Store intermediate format
             instruction = "Complete the statement"
             instructions.append(instruction)
             inputs.append(user_message)
             outputs.append(assistant_message)
 
-            # Create formatted text
-            text = (
-                alpaca_prompt.format(instruction, user_message, assistant_message)
-                + EOS_TOKEN
-            )
+            text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
             texts.append(text)
 
         return {
@@ -199,12 +182,8 @@ if __name__ == "__main__":
             "text": texts,
         }
 
-    dataset_train = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "train"
-    )
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
@@ -259,12 +238,11 @@ if __name__ == "__main__":
         ),
     )
 
-    # run training
     trainer_stats = trainer.train()
 
     add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
 
-    # saving and merging the model to local disk
+    # Merge and save to local disk.
     print("merge and save to local disk")
     model.save_pretrained_merged(
         save_directory = "./unsloth_out/merged_mistral_text_model", tokenizer = tokenizer
@@ -276,7 +254,7 @@ if __name__ == "__main__":
     # torch.cuda.empty_cache()
     # gc.collect()
 
-    # load model from local disk and test
+    # Load merged model from disk and test.
     print("Loading merged model in 4 bit for perplexity test")
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_mistral_text_model",

tests/saving/language_models/test_merge_model_perplexity_phi_4.py

@@ -31,13 +31,10 @@ from tests.utils.perplexity_eval import (
 )
 
 
-# Define helper functions outside of main
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
-        tokenizer.apply_chat_template(
-            convo, tokenize = False, add_generation_prompt = False
-        )
+        tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
         for convo in convos
     ]
     return {
@@ -45,58 +42,52 @@ def formatting_prompts_func(examples):
     }
 
 
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
-    """Load model and compute perplexity in subprocess"""
+def load_and_compute_8bit_ppl(
+    result_queue,
+    load_in_4bit = False,
+    load_in_8bit = False,
+):
+    """Load model and compute perplexity in subprocess."""
     from unsloth import FastLanguageModel
     from unsloth.chat_templates import get_chat_template
     from tests.utils.perplexity_eval import ppl_model
 
-    # Load model
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_phi4_text_model",
         max_seq_length = 2048,
         load_in_4bit = load_in_4bit,
         load_in_8bit = load_in_8bit,
     )
-    # Set up tokenizer
     merged_tokenizer = get_chat_template(
         merged_tokenizer,
         chat_template = "phi-4",
     )
 
     # Load dataset fresh in subprocess
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
-    # Format the dataset
     def formatting_prompts_func(examples):
         convos = examples["messages"]
         texts = [
-            merged_tokenizer.apply_chat_template(
-                convo, tokenize = False, add_generation_prompt = False
-            )
+            merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
             for convo in convos
         ]
         return {"text": texts}
 
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
 
-    # Compute perplexity using the passed dataset
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # IMPORTANT: Convert to Python float if it's a tensor
+    # Coerce to a Python float regardless of source type.
     if torch.is_tensor(ppl_value):
-        ppl_value = ppl_value.cpu().item()  # Move to CPU and convert to Python scalar
+        ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
-        ppl_value = ppl_value.item()  # Convert numpy or other array types
+        ppl_value = ppl_value.item()
     else:
-        ppl_value = float(ppl_value)  # Ensure it's a float
+        ppl_value = float(ppl_value)
 
-    # Return only the perplexity value
     result_queue.put(ppl_value)
 
-    # Clean up
     del merged_model
     del merged_tokenizer
     del dataset_ppl
@@ -104,7 +95,6 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
     gc.collect()
 
 
-# Main execution code should be wrapped in this guard
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
@@ -130,12 +120,8 @@ if __name__ == "__main__":
         chat_template = "phi-4",
     )
 
-    dataset_train = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "train"
-    )
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
@@ -199,12 +185,11 @@ if __name__ == "__main__":
         response_part = "<|im_start|>assistant<|im_sep|>\n\n",
     )
 
-    # run training
     trainer_stats = trainer.train()
 
     add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
 
-    # saving and merging the model to local disk
+    # Merge and save to local disk.
     print("merge and save to local disk")
     model.save_pretrained_merged(
         save_directory = "./unsloth_out/merged_phi4_text_model", tokenizer = tokenizer
@@ -216,7 +201,6 @@ if __name__ == "__main__":
     # torch.cuda.empty_cache()
     # gc.collect()
 
-    # load model from local disk and test
     print("Loading merged model in 4 bit for perplexity test")
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_phi4_text_model",
@@ -253,7 +237,6 @@ if __name__ == "__main__":
 
     print_model_comparison()
 
-    # final cleanup
     safe_remove_directory("./outputs")
     safe_remove_directory("./unsloth_compiled_cache")
     safe_remove_directory("./unsloth_out")

tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py

@@ -34,66 +34,58 @@ from tests.utils.perplexity_eval import (
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
-        tokenizer.apply_chat_template(
-            convo, tokenize = False, add_generation_prompt = False
-        )
+        tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
         for convo in convos
     ]
     return {"text": texts}
 
 
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+    result_queue,
+    load_in_4bit = False,
+    load_in_8bit = False,
+):
     """Load model and compute perplexity in subprocess"""
     from unsloth import FastLanguageModel
     from unsloth.chat_templates import get_chat_template
     from tests.utils.perplexity_eval import ppl_model
 
-    # Load model
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_llama_text_model",
         max_seq_length = 2048,
         load_in_4bit = load_in_4bit,
         load_in_8bit = load_in_8bit,
     )
-    # Set up tokenizer
     merged_tokenizer = get_chat_template(
         merged_tokenizer,
         chat_template = "llama-3.1",
     )
 
     # Load dataset fresh in subprocess
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
-    # Format the dataset
     def formatting_prompts_func(examples):
         convos = examples["messages"]
         texts = [
-            merged_tokenizer.apply_chat_template(
-                convo, tokenize = False, add_generation_prompt = False
-            )
+            merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
             for convo in convos
         ]
         return {"text": texts}
 
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
 
-    # Compute perplexity using the passed dataset
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # IMPORTANT: Convert to Python float if it's a tensor
+    # Convert to a Python float (tensor / numpy / other)
     if torch.is_tensor(ppl_value):
-        ppl_value = ppl_value.cpu().item()  # Move to CPU and convert to Python scalar
+        ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
-        ppl_value = ppl_value.item()  # Convert numpy or other array types
+        ppl_value = ppl_value.item()
     else:
-        ppl_value = float(ppl_value)  # Ensure it's a float
+        ppl_value = float(ppl_value)
 
-    # Return only the perplexity value
     result_queue.put(ppl_value)
 
-    # Clean up
     del merged_model
     del merged_tokenizer
     del dataset_ppl
@@ -101,7 +93,6 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
     gc.collect()
 
 
-# Main execution code should be wrapped in this guard
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
@@ -129,12 +120,8 @@ if __name__ == "__main__":
 
     from unsloth.chat_templates import standardize_sharegpt
 
-    dataset_train = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "train"
-    )
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
@@ -203,12 +190,11 @@ if __name__ == "__main__":
 
     tokenizer.decode(trainer.train_dataset[0]["input_ids"])
 
-    # run training
     trainer_stats = trainer.train()
 
     add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
 
-    # saving and merging the model to local disk
+    # save and merge the model to local disk
     print("merge and save to local disk")
     model.save_pretrained_merged(
         save_directory = "./unsloth_out/merged_llama_text_model", tokenizer = tokenizer
@@ -220,7 +206,6 @@ if __name__ == "__main__":
     # torch.cuda.empty_cache()
     # gc.collect()
 
-    # load model from local disk and test
     print("Loading merged model in 4 bit for perplexity test")
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_llama_text_model",
@@ -257,7 +242,6 @@ if __name__ == "__main__":
 
     print_model_comparison()
 
-    # final cleanup
     safe_remove_directory("./outputs")
     safe_remove_directory("./unsloth_compiled_cache")
     safe_remove_directory("./unsloth_out")

tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py

@@ -42,7 +42,6 @@ alpaca_prompt = """Below is an instruction that describes a task, paired with an
 {}"""
 
 
-# Define helper functions outside of main
 def formatting_prompts_func(examples):
     instructions = []
     inputs = []
@@ -50,7 +49,6 @@ def formatting_prompts_func(examples):
     texts = []
 
     for conversation in examples["messages"]:
-        # Extract user message and assistant response
         user_message = ""
         assistant_message = ""
 
@@ -60,13 +58,11 @@ def formatting_prompts_func(examples):
             elif turn["role"] == "assistant":
                 assistant_message = turn["content"]
 
-        # Store intermediate format
         instruction = "Complete the statement"
         instructions.append(instruction)
         inputs.append(user_message)
         outputs.append(assistant_message)
 
-        # Create formatted text
         text = alpaca_prompt.format(instruction, user_message, assistant_message)
         texts.append(text)
 
@@ -78,12 +74,15 @@ def formatting_prompts_func(examples):
     }
 
 
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
-    """Load model and compute perplexity in subprocess"""
+def load_and_compute_8bit_ppl(
+    result_queue,
+    load_in_4bit = False,
+    load_in_8bit = False,
+):
+    """Load model and compute perplexity in subprocess."""
     from unsloth import FastLanguageModel
     from tests.utils.perplexity_eval import ppl_model
 
-    # Load model
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_qwen_text_model",
         max_seq_length = 2048,
@@ -97,9 +96,7 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
     # )
 
     # Load dataset fresh in subprocess
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
@@ -119,7 +116,6 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
         texts = []
 
         for conversation in examples["messages"]:
-            # Extract user message and assistant response
             user_message = ""
             assistant_message = ""
 
@@ -129,13 +125,11 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
                 elif turn["role"] == "assistant":
                     assistant_message = turn["content"]
 
-            # Store intermediate format
             instruction = "Complete the statement"
             instructions.append(instruction)
             inputs.append(user_message)
             outputs.append(assistant_message)
 
-            # Create formatted text
             text = alpaca_prompt.format(instruction, user_message, assistant_message)
             texts.append(text)
 
@@ -148,18 +142,16 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
 
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
 
-    # Compute perplexity using the passed dataset
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # IMPORTANT: Convert to Python float if it's a tensor
+    # Coerce to a Python float regardless of source type.
     if torch.is_tensor(ppl_value):
-        ppl_value = ppl_value.cpu().item()  # Move to CPU and convert to Python scalar
+        ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
-        ppl_value = ppl_value.item()  # Convert numpy or other array types
+        ppl_value = ppl_value.item()
     else:
-        ppl_value = float(ppl_value)  # Ensure it's a float
+        ppl_value = float(ppl_value)
 
-    # Return only the perplexity value
     result_queue.put(ppl_value)
 
     # Clean up
@@ -170,7 +162,6 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
     # gc.collect()
 
 
-# Main execution code should be wrapped in this guard
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
@@ -191,12 +182,8 @@ if __name__ == "__main__":
         attn_implementation = attn_implementation,
     )
 
-    dataset_train = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "train"
-    )
-    dataset_ppl = load_dataset(
-        "allenai/openassistant-guanaco-reformatted", split = "eval"
-    )
+    dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+    dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
     dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
@@ -252,12 +239,11 @@ if __name__ == "__main__":
         ),
     )
 
-    # run training
     trainer_stats = trainer.train()
 
     add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
 
-    # saving and merging the model to local disk
+    # Merge and save to local disk.
     print("merge and save to local disk")
     model.save_pretrained_merged(
         save_directory = "./unsloth_out/merged_qwen_text_model", tokenizer = tokenizer
@@ -269,7 +255,6 @@ if __name__ == "__main__":
     # torch.cuda.empty_cache()
     # gc.collect()
 
-    # load model from local disk and test
     print("Loading merged model in 4 bit for perplexity test")
     merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./unsloth_out/merged_qwen_text_model",

tests/saving/language_models/test_push_to_hub_merged.py

@@ -32,13 +32,10 @@ from tests.utils.perplexity_eval import (
 )
 
 
-# Define helper functions outside of main
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
-        tokenizer.apply_chat_template(
-            convo, tokenize = False, add_generation_prompt = False
-        )
+        tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
         for convo in convos
     ]
     return {"text": texts}
@@ -137,7 +134,7 @@ trainer = train_on_responses_only(
 trainer_stats = trainer.train()
 
 
-# saving and merging the model to local disk
+# save and merge the model to local disk
 hf_username = os.environ.get("HF_USER", "")
 if not hf_username:
     hf_username = input("Please enter your Hugging Face username: ").strip()
@@ -176,9 +173,7 @@ try:
     print("=== TESTING MODEL DOWNLOAD ===".center(80))
     print("=" * 80 + "\n")
     # Force download even if cached
-    model, tokenizer = FastLanguageModel.from_pretrained(
-        f"{hf_username}/merged_llama_text_model"
-    )
+    model, tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model")
     success["download"] = True
     print("✅ Model downloaded successfully!")
 except Exception as e:

tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py

@@ -36,9 +36,7 @@ from tests.utils.perplexity_eval import (
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
-        tokenizer.apply_chat_template(
-            convo, tokenize = False, add_generation_prompt = False
-        )
+        tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
         for convo in convos
     ]
     return {"text": texts}
@@ -133,7 +131,6 @@ trainer = train_on_responses_only(
     response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
 )
 
-# run training
 trainer_stats = trainer.train()
 
 
@@ -195,9 +192,7 @@ try:
     print("=== TESTING MODEL DOWNLOAD ===".center(80))
     print("=" * 80 + "\n")
     # Force download even if cached
-    model, tokenizer = FastLanguageModel.from_pretrained(
-        f"{hf_username}/merged_llama_text_model"
-    )
+    model, tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model")
     success["download"] = True
     print("✅ Model downloaded successfully!")
 except Exception as e:

tests/saving/language_models/test_save_merged_grpo_model.py

@@ -24,7 +24,11 @@ max_seq_length = 2048  # Can increase for longer reasoning traces
 lora_rank = 64  # Larger rank = smarter, but slower
 
 
-def evaluate_merged_model(result_queue, load_in_4bit = False, load_in_8bit = False):
+def evaluate_merged_model(
+    result_queue,
+    load_in_4bit = False,
+    load_in_8bit = False,
+):
     from unsloth import FastLanguageModel
     from tests.utils.aime_eval import evaluate_model_aime
 
@@ -71,7 +75,6 @@ def evaluate_merged_model(result_queue, load_in_4bit = False, load_in_8bit = Fal
     gc.collect()
 
 
-# Main execution code should be wrapped in this guard
 def training_run(result_queue):
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "meta-llama/Llama-3.2-3B-Instruct",
@@ -141,12 +144,10 @@ def training_run(result_queue):
     </answer>"""
 
         def format_limo(example):
-            # Create the assistant response
             assistant_response = f"<reasoning>\n{example['solution']}\n</reasoning>\n<answer>\n{example['answer']}\n</answer>"
 
-            # Return a DICTIONARY with the conversation in a field
             return {
-                "prompt": [  # ← This is the key change - wrap in a dict
+                "prompt": [
                     {"role": "system", "content": system_prompt},
                     {"role": "user", "content": example["question"]},
                     {"role": "assistant", "content": assistant_response},
@@ -176,12 +177,14 @@ def training_run(result_queue):
         avg_length = sum(lengths) / len(lengths)
         min_length = min(lengths)
 
-        print(
-            f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}"
-        )
+        print(f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}")
         return max_length, avg_length
 
-    def extract_unsloth_answer(text, start_tag = "<SOLUTION>", end_tag = "</SOLUTION>"):
+    def extract_unsloth_answer(
+        text,
+        start_tag = "<SOLUTION>",
+        end_tag = "</SOLUTION>",
+    ):
         """Extract answer from Unsloth SOLUTION tags"""
         pattern = re.escape(start_tag) + r"(.*?)" + re.escape(end_tag)
         matches = re.findall(pattern, text, re.DOTALL)
@@ -265,9 +268,7 @@ def training_run(result_queue):
             ground_truth_num = float(norm_ground_truth)
 
             if ground_truth_num != 0:
-                relative_error = abs(extracted_num - ground_truth_num) / abs(
-                    ground_truth_num
-                )
+                relative_error = abs(extracted_num - ground_truth_num) / abs(ground_truth_num)
 
                 if relative_error < 0.01:
                     return True, True, 0.9
@@ -302,10 +303,7 @@ def training_run(result_queue):
         )
 
         responses = [completion[0]["content"] for completion in completions]
-        rewards = [
-            3.0 if re.match(pattern, response, re.DOTALL) else 0.0
-            for response in responses
-        ]
+        rewards = [3.0 if re.match(pattern, response, re.DOTALL) else 0.0 for response in responses]
         return rewards
 
     def match_format_approximately(completions, **kwargs):
@@ -378,7 +376,6 @@ def training_run(result_queue):
         print("COMPREHENSIVE MODEL COMPARISON")
         print(f"{'='*80}")
 
-        # Main table
         print(
             f"{'Model':<15} {'Format %':<10} {'Exact %':<10} {'Plausible %':<12} {'Confidence':<12}"
         )
@@ -393,7 +390,6 @@ def training_run(result_queue):
                 f"{result['avg_confidence']:<12.3f}"
             )
 
-        # Improvement analysis
         if len(all_results) > 1:
             print(f"\n{'='*50}")
             print("IMPROVEMENT ANALYSIS")
@@ -405,9 +401,7 @@ def training_run(result_queue):
                 format_improvement = (
                     result["correct_format_pct"] - base_result["correct_format_pct"]
                 )
-                exact_improvement = (
-                    result["exact_match_pct"] - base_result["exact_match_pct"]
-                )
+                exact_improvement = result["exact_match_pct"] - base_result["exact_match_pct"]
                 plausible_improvement = (
                     result["plausible_match_pct"] - base_result["plausible_match_pct"]
                 )
@@ -416,7 +410,6 @@ def training_run(result_queue):
                 print(f"  Exact matches:     {exact_improvement:+.1f}%")
                 print(f"  Plausible matches: {plausible_improvement:+.1f}%")
 
-        # Save comparison
         comparison_data = {
             "summary": all_results,
             "best_model": max(all_results, key = lambda x: x["exact_match_pct"]),
@@ -440,31 +433,24 @@ def training_run(result_queue):
         if torch.cuda.is_available():
             allocated = torch.cuda.memory_allocated() / 1024**3
             reserved = torch.cuda.memory_reserved() / 1024**3
-            print(
-                f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB"
-            )
+            print(f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
 
     """#### Data Loading and Preparation"""
 
     from datasets import load_dataset
 
-    # Load GSM8K
     gsm8k_dataset = load_dataset("openai/gsm8k", "main", split = "train")
-
-    # Load LIMO (adjust this based on your access method)
     limo_train = load_dataset("GAIR/LIMO", split = "train")
 
-    # Prepare datasets
     gsm8k_train = prepare_gsm8k_dataset(gsm8k_dataset)
     limo_train = prepare_limo_dataset(limo_train)
 
     print(f"  GSM8K train: {len(gsm8k_train)}")
     print(f"  LIMO train:  {len(limo_train) if limo_train else 0}")
 
-    # Store results
     all_results = []
 
-    # Single temperature evaluation on combined dataset
+    # Single temperature evaluation on combined dataset.
     results = evaluate_model_aime(
         model = model,
         tokenizer = tokenizer,
@@ -486,9 +472,7 @@ def training_run(result_queue):
     def formatting_prompts_func(examples):
         convos = examples["prompt"]
         texts = [
-            tokenizer.apply_chat_template(
-                convo, tokenize = False, add_generation_prompt = False
-            )
+            tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
             for convo in convos
         ]
         return {
@@ -562,16 +546,13 @@ def training_run(result_queue):
             response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
         )
 
-        # Train
         print(f"🚂 Starting SFT training on {len(limo_train)} examples...")
         trainer.train()
 
-        # Save checkpoint
         model.save_pretrained("qlora_checkpoint")
         tokenizer.save_pretrained("qlora_checkpoint")
         print("💾 Qlora checkpoint saved!")
 
-        # Cleanup
         del trainer
         cleanup_memory()
 
@@ -579,7 +560,6 @@ def training_run(result_queue):
     else:
         print("⚠️ Skipping Qlora training - no LIMO dataset available")
 
-    # Cleanup
     cleanup_memory()
 
     global PRINTED_TIMES
@@ -601,7 +581,7 @@ def training_run(result_queue):
         ]
 
         scores = []
-        # Print only every few steps
+        # Print only every few steps.
         global PRINTED_TIMES
         global PRINT_EVERY_STEPS
         if PRINTED_TIMES % PRINT_EVERY_STEPS == 0:
@@ -618,10 +598,9 @@ def training_run(result_queue):
             if guess is None:
                 scores.append(0)
                 continue
-            # Convert to numbers
             try:
                 true_answer = float(true_answer.strip())
-                # Remove commas like in 123,456
+                # Remove commas like in 123,456.
                 guess = float(guess.strip().replace(",", ""))
                 scores.append(1.5 if guess == true_answer else -0.5)
             except:
@@ -633,7 +612,6 @@ def training_run(result_queue):
     print("🎯 STAGE 2: GRPO Fine-Tuning on GSM8K")
     print(f"{'*'*60}")
 
-    # Get max prompt length
     max_prompt_length, _ = get_max_prompt_length(gsm8k_train, tokenizer)
     max_prompt_length = min(max_prompt_length + 10, 512)  # Add buffer, cap at 512
 
@@ -675,16 +653,13 @@ def training_run(result_queue):
         train_dataset = gsm8k_train,
     )
 
-    # Train
     print(f"🚂 Starting GRPO training on {len(gsm8k_train)} examples...")
     trainer.train()
 
-    # Save checkpoint
     model.save_pretrained("grpo_checkpoint")
     tokenizer.save_pretrained("grpo_checkpoint")
     print("💾 GRPO checkpoint saved!")
 
-    # Cleanup
     del trainer
     del training_args
     cleanup_memory()
@@ -713,11 +688,8 @@ def training_run(result_queue):
     print("💾 SAVING FINAL MODEL")
     print(f"{'='*60}")
 
-    # Save as merged model
     try:
-        model.save_pretrained_merged(
-            "final_merged_model", tokenizer, save_method = "merged_16bit"
-        )
+        model.save_pretrained_merged("final_merged_model", tokenizer, save_method = "merged_16bit")
         print("✅ Merged model saved to: final_merged_model/")
     except Exception as e:
         print(f"⚠️ Could not save merged model: {e}")
@@ -729,7 +701,6 @@ def training_run(result_queue):
 
     result_queue.put(results)
 
-    # Clean up
     del model
     del tokenizer
     torch.cuda.empty_cache()
@@ -778,7 +749,7 @@ if __name__ == "__main__":
     result_queue = mp.Queue()
     all_results = []
 
-    # run main finetuning and grpo loop
+    # Run main finetuning and GRPO loop.
     p = mp.Process(target = training_run, args = (result_queue,))
     p.start()
     p.join()
@@ -786,7 +757,7 @@ if __name__ == "__main__":
     results = result_queue.get()
     all_results = results
 
-    # evaluate merged model loaded 16bits
+    # Evaluate merged model loaded 16bits.
     p = mp.Process(target = evaluate_merged_model, args = (result_queue, False, False))
     p.start()
     p.join()
@@ -815,11 +786,8 @@ if __name__ == "__main__":
 
     safe_remove_directory("./unsloth_compiled_cache")
 
-    # AIME-specific comparison function
-
     print(f"\n{'='*80}")
     print("🏆 FINAL TRAINING PIPELINE RESULTS")
     print(f"{'='*80}")
 
-    # Use the AIME-specific comparison
     compare_aime_results(all_results)

tests/saving/non_peft/test_mistral_non_peft.py

@@ -27,8 +27,6 @@ model, tokenizer = FastLanguageModel.from_pretrained(
 
 print("✅ Base model loaded successfully!")
 
-### Attempting save merge
-
 
 print(f"\n{'='*80}")
 print("🔍 PHASE 2: Attempting save_pretrained_merged (Should Warn)")
@@ -38,7 +36,6 @@ with warnings.catch_warnings(record = True) as w:
     warnings.simplefilter("always")
     model.save_pretrained_merged("test_output", tokenizer)
 
-    # Verify warning
     assert len(w) >= 1, "Expected warning but none raised"
     warning_msg = str(w[0].message)
     expected_msg = "Model is not a PeftModel (no Lora adapters detected). Skipping Merge. Please use save_pretrained() or push_to_hub() instead!"
@@ -55,7 +52,7 @@ print(f"{'='*80}")
 
 try:
     with warnings.catch_warnings():
-        warnings.simplefilter("error")  # Treat warnings as errors here
+        warnings.simplefilter("error")  # Treat warnings as errors
         model.save_pretrained("test_output")
         print("✅ Standard save_pretrained completed successfully!")
 except Exception as e:

tests/saving/non_peft/test_whisper_non_peft.py

@@ -27,8 +27,6 @@ model, tokenizer = FastModel.from_pretrained(
 
 print("✅ Base model loaded successfully!")
 
-### Attempting save merge
-
 
 print(f"\n{'='*80}")
 print("🔍 PHASE 2: Attempting save_pretrained_merged (Should Warn)")
@@ -38,7 +36,6 @@ with warnings.catch_warnings(record = True) as w:
     warnings.simplefilter("always")
     model.save_pretrained_merged("test_output", tokenizer)
 
-    # Verify warning
     assert len(w) >= 1, "Expected warning but none raised"
     warning_msg = str(w[0].message)
     expected_msg = "Model is not a PeftModel (no Lora adapters detected). Skipping Merge. Please use save_pretrained() or push_to_hub() instead!"

tests/saving/test_fix_sentencepiece_gguf_robustness.py

@@ -44,9 +44,7 @@ def test_user_defined_special_piece_is_not_retyped(tmp_path):
     ]
     (tmp_path / "tokenizer.model").write_bytes(_build(pieces))
     (tmp_path / "tokenizer.json").write_text(
-        json.dumps(
-            {"added_tokens": [{"id": 2, "content": "<ud_special>", "special": True}]}
-        )
+        json.dumps({"added_tokens": [{"id": 2, "content": "<ud_special>", "special": True}]})
     )
     fix_sentencepiece_gguf(str(tmp_path))
     got = dict(_read(str(tmp_path / "tokenizer.model")))
@@ -87,10 +85,7 @@ def test_save_py_except_clause_is_broad_exception():
     with open(_SAVE_PY) as f:
         tree = ast.parse(f.read())
     for node in ast.walk(tree):
-        if (
-            isinstance(node, ast.FunctionDef)
-            and node.name == "unsloth_save_pretrained_gguf"
-        ):
+        if isinstance(node, ast.FunctionDef) and node.name == "unsloth_save_pretrained_gguf":
             for subnode in ast.walk(node):
                 if isinstance(subnode, ast.Try):
                     body_src = "\n".join(ast.unparse(s) for s in subnode.body)

tests/saving/test_preserve_tokenizer_eos_token.py

@@ -16,8 +16,7 @@ def _load_preserve_helper():
     helper = next(
         node
         for node in tree.body
-        if isinstance(node, ast.FunctionDef)
-        and node.name == "_preserve_tokenizer_eos_token"
+        if isinstance(node, ast.FunctionDef) and node.name == "_preserve_tokenizer_eos_token"
     )
     module = ast.Module(body = [helper], type_ignores = [])
     ast.fix_missing_locations(module)
@@ -46,9 +45,7 @@ def test_preserve_tokenizer_eos_token_supports_processor_tokenizer(tmp_path):
     preserve = _load_preserve_helper()
     tokenizer_config = tmp_path / "tokenizer_config.json"
     tokenizer_config.write_text(json.dumps({"eos_token": "<eos>"}), encoding = "utf-8")
-    processor = types.SimpleNamespace(
-        tokenizer = types.SimpleNamespace(eos_token = "<turn|>")
-    )
+    processor = types.SimpleNamespace(tokenizer = types.SimpleNamespace(eos_token = "<turn|>"))
 
     preserve(processor, tmp_path)
 

tests/saving/test_qwen3_5_vlm_full_finetune_key_remap.py

@@ -0,0 +1,71 @@
+import ast
+import types
+from pathlib import Path
+
+import torch
+
+
+def _load_qwen3_5_vlm_save_helpers():
+    source = Path(__file__).parents[2] / "unsloth" / "save.py"
+    tree = ast.parse(source.read_text(encoding = "utf-8"))
+    helpers = [
+        node
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef)
+        and node.name
+        in {
+            "_is_qwen3_5_vlm",
+            "_qwen3_5_vlm_state_dict_for_save",
+        }
+    ]
+    module = ast.Module(body = helpers, type_ignores = [])
+    ast.fix_missing_locations(module)
+    namespace = {}
+    exec(compile(module, str(source), "exec"), namespace)
+    return namespace
+
+
+def _qwen3_5_vlm_model():
+    return types.SimpleNamespace(
+        config = types.SimpleNamespace(
+            architectures = ["Qwen3_5ForConditionalGeneration"],
+            model_type = "qwen3_5",
+            vision_config = types.SimpleNamespace(),
+        )
+    )
+
+
+def test_qwen3_5_vlm_state_dict_uses_hf_checkpoint_namespace():
+    helpers = _load_qwen3_5_vlm_save_helpers()
+    state_dict = {
+        "language_model.model.embed_tokens.weight": torch.ones(2, 2),
+        "language_model.model.layers.0.input_layernorm.weight": torch.ones(2),
+        "language_model.lm_head.weight": torch.ones(2, 2),
+        "visual.blocks.0.norm1.weight": torch.ones(2),
+        "other.weight": torch.ones(2),
+    }
+
+    remapped = helpers["_qwen3_5_vlm_state_dict_for_save"](state_dict)
+
+    assert "model.language_model.embed_tokens.weight" in remapped
+    assert "model.language_model.layers.0.input_layernorm.weight" in remapped
+    assert "lm_head.weight" in remapped
+    assert "model.visual.blocks.0.norm1.weight" in remapped
+    assert "other.weight" in remapped
+    assert "language_model.model.embed_tokens.weight" not in remapped
+    assert "language_model.lm_head.weight" not in remapped
+    assert "visual.blocks.0.norm1.weight" not in remapped
+
+
+def test_qwen3_5_vlm_detection_requires_vision_config():
+    helpers = _load_qwen3_5_vlm_save_helpers()
+    assert helpers["_is_qwen3_5_vlm"](_qwen3_5_vlm_model())
+
+    model = types.SimpleNamespace(
+        config = types.SimpleNamespace(
+            architectures = ["Qwen3_5ForCausalLM"],
+            model_type = "qwen3_5_text",
+        )
+    )
+
+    assert not helpers["_is_qwen3_5_vlm"](model)

tests/saving/test_save_shell_injection.py

@@ -19,10 +19,7 @@ def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None:
     popen_calls = []
     for call in calls:
         if isinstance(call.func, ast.Attribute) and call.func.attr == "Popen":
-            if (
-                isinstance(call.func.value, ast.Name)
-                and call.func.value.id == "subprocess"
-            ):
+            if isinstance(call.func.value, ast.Name) and call.func.value.id == "subprocess":
                 popen_calls.append(call)
 
     assert popen_calls, "Expected at least one subprocess.Popen call"
@@ -54,9 +51,7 @@ def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None:
 
         assert call.args, "subprocess.Popen must receive argv as a positional argument"
         argv = call.args[0]
-        assert isinstance(
-            argv, ast.List
-        ), "subprocess.Popen must be called with an argv list"
+        assert isinstance(argv, ast.List), "subprocess.Popen must be called with an argv list"
         assert len(argv.elts) == 5, "GGML conversion argv should have five elements"
 
         second_arg = argv.elts[1]

tests/saving/test_unsloth_save.py

@@ -31,7 +31,6 @@ torchao_models = [
 ]
 
 
-# Variables
 save_file_sizes = {}
 save_file_sizes["merged_16bit"] = {}
 save_file_sizes["merged_4bit"] = {}
@@ -55,7 +54,6 @@ def loaded_model_tokenizer(request):
         load_in_4bit = True,
     )
 
-    # Apply LoRA
     model = FastModel.get_peft_model(
         model,
         r = 16,
@@ -80,7 +78,6 @@ def fp16_model_tokenizer(request):
         load_in_4bit = False,  # No BnB quantization
     )
 
-    # Apply LoRA
     model = FastModel.get_peft_model(
         model,
         r = 16,
@@ -112,7 +109,6 @@ def temp_save_dir():
 
 
 def delete_quantization_config(model):
-    # Since merged, edit quantization_config
     old_config = model.config
     new_config = model.config.to_dict()
     if "quantization_config" in new_config:
@@ -132,44 +128,33 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
         model.config._name_or_path.replace("/", "_"),
     )
 
-    model.save_pretrained_merged(
-        save_path, tokenizer = tokenizer, save_method = "merged_16bit"
-    )
+    model.save_pretrained_merged(save_path, tokenizer = tokenizer, save_method = "merged_16bit")
 
-    # Check model files
     assert os.path.isdir(save_path), f"Directory {save_path} does not exist."
-    assert os.path.isfile(
-        os.path.join(save_path, "config.json")
-    ), "config.json not found."
+    assert os.path.isfile(os.path.join(save_path, "config.json")), "config.json not found."
 
     weight_files = [
-        f
-        for f in os.listdir(save_path)
-        if f.endswith(".bin") or f.endswith(".safetensors")
+        f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors")
     ]
     assert len(weight_files) > 0, "No weight files found in the save directory."
 
-    # Check tokenizer files
     for file in tokenizer_files:
         assert os.path.isfile(
             os.path.join(save_path, file)
         ), f"{file} not found in the save directory."
 
-    # Check config to see if it is 16bit by checking for quantization config
+    # 16bit if there's no quantization config
     config_path = os.path.join(save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
 
-    assert (
-        "quantization_config" not in config
-    ), "Quantization config not found in the model config."
+    assert "quantization_config" not in config, "Quantization config not found in the model config."
 
-    # Store the size of the model files
     total_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files)
     save_file_sizes["merged_16bit"][model.config._name_or_path] = total_size
     print(f"Total size of merged_16bit files: {total_size} bytes")
 
-    # Test loading the model from the saved path
+    # Verify the saved model loads
     loaded_model, loaded_tokenizer = FastLanguageModel.from_pretrained(
         save_path,
         max_seq_length = 128,
@@ -185,30 +170,21 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
         model.config._name_or_path.replace("/", "_"),
     )
 
-    model.save_pretrained_merged(
-        save_path, tokenizer = tokenizer, save_method = "merged_4bit_forced"
-    )
+    model.save_pretrained_merged(save_path, tokenizer = tokenizer, save_method = "merged_4bit_forced")
 
-    # Check model files
     assert os.path.isdir(save_path), f"Directory {save_path} does not exist."
-    assert os.path.isfile(
-        os.path.join(save_path, "config.json")
-    ), "config.json not found."
+    assert os.path.isfile(os.path.join(save_path, "config.json")), "config.json not found."
 
     weight_files = [
-        f
-        for f in os.listdir(save_path)
-        if f.endswith(".bin") or f.endswith(".safetensors")
+        f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors")
     ]
     assert len(weight_files) > 0, "No weight files found in the save directory."
 
-    # Check tokenizer files
     for file in tokenizer_files:
         assert os.path.isfile(
             os.path.join(save_path, file)
         ), f"{file} not found in the save directory."
 
-    # Store the size of the model files
     total_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files)
     save_file_sizes["merged_4bit"][model.config._name_or_path] = total_size
 
@@ -218,16 +194,14 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
         total_size < save_file_sizes["merged_16bit"][model.config._name_or_path]
     ), "Merged 4bit files are larger than merged 16bit files."
 
-    # Check config to see if it is 4bit
+    # 4bit if there's a quantization config
     config_path = os.path.join(save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
 
-    assert (
-        "quantization_config" in config
-    ), "Quantization config not found in the model config."
+    assert "quantization_config" in config, "Quantization config not found in the model config."
 
-    # Test loading the model from the saved path
+    # Verify the saved model loads
     loaded_model, loaded_tokenizer = FastModel.from_pretrained(
         save_path,
         max_seq_length = 128,
@@ -257,60 +231,41 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str):
     )
 
     weight_files_16bit = [
-        f
-        for f in os.listdir(save_path)
-        if f.endswith(".bin") or f.endswith(".safetensors")
+        f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors")
     ]
-    total_16bit_size = sum(
-        os.path.getsize(os.path.join(save_path, f)) for f in weight_files_16bit
-    )
+    total_16bit_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files_16bit)
     save_file_sizes["merged_16bit"][model.config._name_or_path] = total_16bit_size
 
     torchao_save_path = save_path + "-torchao"
 
-    # Check model files
-    assert os.path.isdir(
-        torchao_save_path
-    ), f"Directory {torchao_save_path} does not exist."
-    assert os.path.isfile(
-        os.path.join(torchao_save_path, "config.json")
-    ), "config.json not found."
+    assert os.path.isdir(torchao_save_path), f"Directory {torchao_save_path} does not exist."
+    assert os.path.isfile(os.path.join(torchao_save_path, "config.json")), "config.json not found."
 
     weight_files = [
-        f
-        for f in os.listdir(torchao_save_path)
-        if f.endswith(".bin") or f.endswith(".safetensors")
+        f for f in os.listdir(torchao_save_path) if f.endswith(".bin") or f.endswith(".safetensors")
     ]
     assert len(weight_files) > 0, "No weight files found in the save directory."
 
-    # Check tokenizer files
     for file in tokenizer_files:
         assert os.path.isfile(
             os.path.join(torchao_save_path, file)
         ), f"{file} not found in the save directory."
 
-    # Store the size of the model files
-    total_size = sum(
-        os.path.getsize(os.path.join(torchao_save_path, f)) for f in weight_files
-    )
+    total_size = sum(os.path.getsize(os.path.join(torchao_save_path, f)) for f in weight_files)
     save_file_sizes["torchao"][model.config._name_or_path] = total_size
 
     assert (
         total_size < save_file_sizes["merged_16bit"][model.config._name_or_path]
     ), "torchao files are larger than merged 16bit files."
 
-    # Check config to see if it is quantized with torchao
     config_path = os.path.join(torchao_save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
 
-    assert (
-        "quantization_config" in config
-    ), "Quantization config not found in the model config."
+    assert "quantization_config" in config, "Quantization config not found in the model config."
 
-    # Test loading the model from the saved path
-    # can't set `load_in_4bit` to True because the model is torchao quantized
-    # can't quantize again with bitsandbytes
+    # load_in_4bit must stay False: a torchao-quantized model can't be
+    # re-quantized with bitsandbytes.
     import torch.serialization
 
     with torch.serialization.safe_globals([getattr]):
@@ -332,15 +287,12 @@ def test_save_and_inference_torchao(fp16_model_tokenizer, temp_save_dir: str):
 
     print(f"Testing TorchAO save and inference for: {model_name}")
 
-    save_path = os.path.join(
-        temp_save_dir, "torchao_models", model_name.replace("/", "_")
-    )
+    save_path = os.path.join(temp_save_dir, "torchao_models", model_name.replace("/", "_"))
 
     from torchao.quantization import Int8DynamicActivationInt8WeightConfig
 
     torchao_config = Int8DynamicActivationInt8WeightConfig()
 
-    # Save with TorchAO
     model.save_pretrained_torchao(
         save_path,
         tokenizer = tokenizer,
@@ -350,12 +302,10 @@ def test_save_and_inference_torchao(fp16_model_tokenizer, temp_save_dir: str):
 
     torchao_save_path = save_path + "-torchao"
 
-    # Verify files exist
     assert os.path.isdir(
         torchao_save_path
     ), f"TorchAO directory {torchao_save_path} does not exist."
 
-    # Load with safe globals
     import torch.serialization
 
     with torch.serialization.safe_globals([getattr]):
@@ -366,7 +316,7 @@ def test_save_and_inference_torchao(fp16_model_tokenizer, temp_save_dir: str):
             load_in_4bit = False,
         )
 
-    FastModel.for_inference(loaded_model)  # Enable native 2x faster inference
+    FastModel.for_inference(loaded_model)
 
     messages = [
         {
@@ -377,21 +327,20 @@ def test_save_and_inference_torchao(fp16_model_tokenizer, temp_save_dir: str):
     inputs = loaded_tokenizer.apply_chat_template(
         messages,
         tokenize = True,
-        add_generation_prompt = True,  # Must add for generation
+        add_generation_prompt = True,  # required for generation
         return_tensors = "pt",
     ).to("cuda")
 
-    outputs = loaded_model.generate(  # ← Use loaded_model, not model
+    outputs = loaded_model.generate(
         input_ids = inputs,
         max_new_tokens = 64,
-        use_cache = False,  # Avoid cache issues
+        use_cache = False,  # avoid cache issues
         temperature = 1.5,
         min_p = 0.1,
         do_sample = True,
         pad_token_id = loaded_tokenizer.pad_token_id or loaded_tokenizer.eos_token_id,
     )
 
-    # Decode with the LOADED tokenizer
     generated_text = loaded_tokenizer.decode(outputs[0], skip_special_tokens = True)
     input_text = loaded_tokenizer.decode(inputs[0], skip_special_tokens = True)
     response_part = generated_text[len(input_text) :].strip()

tests/saving/text_to_speech_models/test_csm.py

@@ -134,22 +134,18 @@ import torch
 
 output_audio_path = "csm_audio.wav"
 try:
-    text = (
-        "We just finished fine tuning a text to speech model... and it's pretty good!"
-    )
+    text = "We just finished fine tuning a text to speech model... and it's pretty good!"
     speaker_id = 0
     inputs = processor(f"[{speaker_id}]{text}", add_special_tokens = True).to("cuda")
     audio_values = model.generate(
         **inputs,
-        max_new_tokens = 125,  # 125 tokens is 10 seconds of audio, for longer speech increase this
-        # play with these parameters to get the best results
+        max_new_tokens = 125,  # 125 tokens ~= 10 seconds of audio
         depth_decoder_temperature = 0.6,
         depth_decoder_top_k = 0,
         depth_decoder_top_p = 0.9,
         temperature = 0.8,
         top_k = 50,
         top_p = 1.0,
-        #########################################################
         output_audio = True,
     )
     audio = audio_values[0].to(torch.float32).cpu().numpy()
@@ -159,8 +155,6 @@ except Exception as e:
     assert False, f"Inference failed with exception: {e}"
 
 
-## assert that transcribed_text contains The birch canoe slid on the smooth planks. Glued the sheet to the dark blue background. It's easy to tell the depth of a well. Four hours of steady work faced us.
-
 print("✅ All sections passed successfully!")
 
 

tests/saving/text_to_speech_models/test_lasa.py

@@ -167,9 +167,7 @@ def extract_speech_ids(speech_tokens_str):
 # TTS start!
 with torch.inference_mode():
     with torch.amp.autocast("cuda", dtype = model.dtype):
-        formatted_text = (
-            f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
-        )
+        formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
 
         # Tokenize the text
         chat = [

tests/saving/text_to_speech_models/test_orpheus.py

@@ -99,7 +99,7 @@ print("🔍 SECTION 4: Saving and Merging Model")
 print(f"{'='*80}")
 
 with warnings.catch_warnings():
-    warnings.simplefilter("error")  # Treat warnings as errors
+    warnings.simplefilter("error")
     try:
         model.save_pretrained_merged("orpheus", tokenizer)
         print("✅ Model saved and merged successfully without warnings!")
@@ -135,7 +135,6 @@ print(f"{'='*80}")
 
 FastLanguageModel.for_inference(model)  # Enable native 2x faster inference
 
-# Moving snac_model cuda to cpu
 snac_model.to("cpu")
 prompts = [
     "Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.",
@@ -152,9 +151,7 @@ for prompt in prompts_:
     all_input_ids.append(input_ids)
 
 start_token = torch.tensor([[128259]], dtype = torch.int64)  # Start of human
-end_tokens = torch.tensor(
-    [[128009, 128260]], dtype = torch.int64
-)  # End of text, End of human
+end_tokens = torch.tensor([[128009, 128260]], dtype = torch.int64)  # End of text, End of human
 
 all_modified_input_ids = []
 for input_ids in all_input_ids:
@@ -165,9 +162,7 @@ for input_ids in all_input_ids:
 
 all_padded_tensors = []
 all_attention_masks = []
-max_length = max(
-    [modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids]
-)
+max_length = max([modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids])
 for modified_input_ids in all_modified_input_ids:
     padding = max_length - modified_input_ids.shape[1]
     padded_tensor = torch.cat(
@@ -267,7 +262,6 @@ try:
 except Exception as e:
     assert False, f"Inference failed with exception: {e}"
 
-# Verify the file exists
 import os
 
 assert os.path.exists(output_path), f"Audio file not found at {output_path}"

tests/saving/text_to_speech_models/test_whisper.py

@@ -153,7 +153,6 @@ import torch
 
 FastModel.for_inference(model)
 model.eval()
-# Create pipeline without specifying the device
 whisper = pipeline(
     "automatic-speech-recognition",
     model = model,
@@ -161,9 +160,8 @@ whisper = pipeline(
     feature_extractor = tokenizer.feature_extractor,
     processor = tokenizer,
     return_language = True,
-    torch_dtype = torch.float16,  # Remove the device parameter
+    torch_dtype = torch.float16,
 )
-# Example usage
 audio_file = "Speech_12dB_s16.flac"
 transcribed_text = whisper(audio_file)
 # audio, sr = sf.read(audio_file)
@@ -181,13 +179,9 @@ expected_phrases = [
 ]
 
 transcribed_lower = transcribed_text["text"].lower()
-all_phrases_found = all(
-    phrase.lower() in transcribed_lower for phrase in expected_phrases
-)
+all_phrases_found = all(phrase.lower() in transcribed_lower for phrase in expected_phrases)
 
-assert (
-    all_phrases_found
-), f"Expected phrases not found in transcription: {transcribed_text['text']}"
+assert all_phrases_found, f"Expected phrases not found in transcription: {transcribed_text['text']}"
 print("✅ Transcription contains all expected phrases!")
 
 

tests/saving/vision_models/test_index_file_sharded_model.py

@@ -1,5 +1,3 @@
-## Import required libraries
-
 from unsloth import FastVisionModel, is_bf16_supported
 from unsloth.trainer import UnslothVisionDataCollator
 
@@ -18,14 +16,10 @@ sys.path.insert(0, str(REPO_ROOT))
 from tests.utils.cleanup_utils import safe_remove_directory
 
 
-## Dataset Preparation"""
-
 print("\n📊 Loading and preparing dataset...")
 dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", "en", split = "train")
-# To select the first 2000 examples
 train_dataset = dataset.select(range(2000))
 
-# To select the next 200 examples for evaluation
 eval_dataset = dataset.select(range(2000, 2200))
 
 print(f"✅ Dataset loaded successfully!")
@@ -33,7 +27,7 @@ print(f"   📈 Training samples: {len(train_dataset)}")
 print(f"   📊 Evaluation samples: {len(eval_dataset)}")
 
 
-# Convert dataset to OAI messages
+# Convert dataset to OAI messages.
 def format_data(sample):
     return {
         "messages": [
@@ -64,8 +58,7 @@ def format_data(sample):
 
 print("\n🔄 Formatting dataset for vision training...")
 system_message = "You are an expert french ocr system."
-# Convert dataset to OAI messages
-# need to use list comprehension to keep Pil.Image type, .map convert image to bytes
+# List comprehension (not .map) keeps PIL.Image type; .map would convert images to bytes.
 train_dataset = [format_data(sample) for sample in train_dataset]
 eval_dataset = [format_data(sample) for sample in eval_dataset]
 print("✅ Dataset formatting completed!")
@@ -76,11 +69,9 @@ print("✅ Dataset formatting completed!")
 print("\n" + "=" * 80)
 print("=== MODEL LOADING AND SETUP ===".center(80))
 print("=" * 80 + "\n")
-# Load Base Model
 print("🤖 Loading base vision model...")
 try:
     model, tokenizer = FastVisionModel.from_pretrained(
-        # model_name = "unsloth/Qwen2-VL-7B-Instruct",
         model_name = "unsloth/Qwen2-VL-7B-Instruct",
         max_seq_length = 2048,  # Choose any for long context!
         load_in_4bit = True,  # 4 bit quantization to reduce memory
@@ -92,7 +83,6 @@ except Exception as e:
     raise
 
 print("\n🔧 Setting up LoRA configuration...")
-## Lora Finetuning
 try:
     model = FastVisionModel.get_peft_model(
         model,
@@ -138,9 +128,7 @@ try:
             per_device_train_batch_size = 2,
             gradient_accumulation_steps = 4,
             gradient_checkpointing = True,
-            gradient_checkpointing_kwargs = {
-                "use_reentrant": False
-            },  # use reentrant checkpointing
+            gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
             max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
             warmup_ratio = 0.03,
             # num_train_epochs = 2, # Set this instead of max_steps for full training runs
@@ -177,7 +165,6 @@ except Exception as e:
 print("\n" + "=" * 80)
 print("=== STARTING TRAINING ===".center(80))
 print("=" * 80 + "\n")
-# run training
 try:
     print("🚀 Starting training process...")
     trainer_stats = trainer.train()
@@ -215,7 +202,7 @@ success = {
     "safetensors_check": False,
     "download": False,
 }
-# Stage 1: Upload model to Hub
+# Stage 1: upload model to Hub
 try:
     print("\n" + "=" * 80)
     print("=== UPLOADING MODEL TO HUB ===".center(80))
@@ -228,7 +215,7 @@ except Exception as e:
     print(f"❌ Failed to upload model: {e}")
     raise Exception("Model upload failed.")
 
-# Stage 2: Verify safetensors.index.json exists
+# Stage 2: verify safetensors.index.json exists
 try:
     print("\n" + "=" * 80)
     print("=== VERIFYING REPO CONTENTS ===".center(80))
@@ -247,7 +234,7 @@ except Exception as e:
     print(f"❌ Verification failed: {e}")
     raise Exception("Repo verification failed.")
 
-# test downloading model even if cached
+# Stage 3: test download even if cached
 safe_remove_directory(f"./{hf_username}")
 
 try:
@@ -255,19 +242,16 @@ try:
     print("=== TESTING MODEL DOWNLOAD ===".center(80))
     print("=" * 80 + "\n")
     print("📥 Testing model download...")
-    # Force download even if cached
     test_model, test_tokenizer = FastVisionModel.from_pretrained(repo_name)
     success["download"] = True
     print("✅ Model downloaded successfully!")
 
-    # Clean up test model
     del test_model, test_tokenizer
     torch.cuda.empty_cache()
 except Exception as e:
     print(f"❌ Download failed: {e}")
     raise Exception("Model download failed.")
 
-# Final report
 print("\n" + "=" * 80)
 print("=== VALIDATION REPORT ===".center(80))
 print("=" * 80 + "\n")
@@ -283,7 +267,6 @@ else:
     raise Exception("Validation failed for one or more stages.")
 
 
-# Final cleanup
 print("\n🧹 Cleaning up temporary files...")
 safe_remove_directory("./checkpoints")
 safe_remove_directory("./unsloth_compiled_cache")

tests/saving/vision_models/test_push_to_hub_merged.py

@@ -23,10 +23,10 @@ from tests.utils.cleanup_utils import safe_remove_directory
 
 print("\n📊 Loading and preparing dataset...")
 dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", "en", split = "train")
-# To select the first 2000 examples
+# First 2000 examples for training
 train_dataset = dataset.select(range(2000))
 
-# To select the next 200 examples for evaluation
+# Next 200 examples for evaluation
 eval_dataset = dataset.select(range(2000, 2200))
 
 print(f"✅ Dataset loaded successfully!")
@@ -65,8 +65,7 @@ def format_data(sample):
 
 print("\n🔄 Formatting dataset for vision training...")
 system_message = "You are an expert french ocr system."
-# Convert dataset to OAI messages
-# need to use list comprehension to keep Pil.Image type, .map convert image to bytes
+# Use a list comprehension (not .map) to keep PIL.Image type; .map converts images to bytes.
 train_dataset = [format_data(sample) for sample in train_dataset]
 eval_dataset = [format_data(sample) for sample in eval_dataset]
 print("✅ Dataset formatting completed!")
@@ -139,9 +138,7 @@ try:
             per_device_train_batch_size = 2,
             gradient_accumulation_steps = 4,
             gradient_checkpointing = True,
-            gradient_checkpointing_kwargs = {
-                "use_reentrant": False
-            },  # use reentrant checkpointing
+            gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
             max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
             warmup_ratio = 0.03,
             # num_train_epochs = 2, # Set this instead of max_steps for full training runs

tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py

@@ -23,10 +23,7 @@ from tests.utils.ocr_eval import OCRModelEvaluator
 from datasets import load_dataset
 
 dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", "en", split = "train")
-# To select the first 2000 examples
 train_dataset = dataset.select(range(2000))
-
-# To select the next 200 examples for evaluation
 eval_dataset = dataset.select(range(2000, 2200))
 
 
@@ -60,8 +57,7 @@ def format_data(sample):
 
 
 system_message = "You are an expert french ocr system."
-# Convert dataset to OAI messages
-# need to use list comprehension to keep Pil.Image type, .map convert image to bytes
+# List comprehension (not .map) to keep PIL.Image type; .map converts images to bytes.
 train_dataset = [format_data(sample) for sample in train_dataset]
 eval_dataset = [format_data(sample) for sample in eval_dataset]
 
@@ -73,12 +69,10 @@ import pandas as pd
 from jiwer import wer, cer
 from qwen_vl_utils import process_vision_info
 
-#
 ocr_evaluator = OCRModelEvaluator()
 model_comparison_results = {}
 
 ## Finetuning Setup and Run
-# Load Base Model
 
 model, tokenizer = FastVisionModel.from_pretrained(
     model_name = "unsloth/Qwen2.5-VL-32B-Instruct-bnb-4bit",
@@ -88,7 +82,7 @@ model, tokenizer = FastVisionModel.from_pretrained(
     full_finetuning = False,  # [NEW!] We have full finetuning now!
 )
 
-# benchmark base model performance
+# Benchmark base model
 model_name = "Unsloth Base model"
 FastVisionModel.for_inference(model)
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
@@ -119,7 +113,7 @@ model = FastVisionModel.get_peft_model(
 from unsloth import is_bf16_supported
 from unsloth.trainer import UnslothVisionDataCollator
 
-FastVisionModel.for_training(model)  # Enable for training!
+FastVisionModel.for_training(model)
 model.config.use_cache = False
 
 
@@ -134,9 +128,7 @@ trainer = SFTTrainer(
         per_device_train_batch_size = 2,
         gradient_accumulation_steps = 4,
         gradient_checkpointing = True,
-        gradient_checkpointing_kwargs = {
-            "use_reentrant": False
-        },  # use reentrant checkpointing
+        gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
         max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
         warmup_ratio = 0.03,
         # num_train_epochs = 2, # Set this instead of max_steps for full training runs
@@ -161,7 +153,6 @@ trainer = SFTTrainer(
     ),
 )
 
-# run training
 trainer_stats = trainer.train()
 
 model.save_pretrained("unsloth-qwen2.5-vl-32b-french-ocr-adapter", tokenizer)
@@ -169,7 +160,6 @@ tokenizer.save_pretrained("unsloth-qwen2.5-vl-32b-french-ocr-adapter")
 
 ## Measure Adapter Performance
 
-# benchmark lora model performance
 model_name = "Unsloth lora adapter model"
 FastVisionModel.for_inference(model)
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
@@ -193,7 +183,7 @@ base = find_lora_base_model(model)
 
 print((base.__class__.__name__))
 
-# merge default 16 bits
+# Merge to 16-bit (default)
 model.save_pretrained_merged(
     save_directory = "qwen2.5-ocr-merged-finetune-merge-16bit", tokenizer = tokenizer
 )
@@ -207,7 +197,6 @@ model, tokenizer = FastVisionModel.from_pretrained(
     "./qwen2.5-ocr-merged-finetune-merge-16bit", load_in_4bit = False, load_in_8bit = False
 )
 
-# benchmark 4bit loaded, 16bits merged model performance
 model_name = "Unsloth 16bits-merged model load-16bits"
 model.config.use_cache = True
 
@@ -219,12 +208,11 @@ avg_wer, avg_cer = ocr_evaluator.evaluate_model(
 )
 ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
-# load 16bits-merged model in 4 bits
+# Load 16bits-merged model in 4 bits
 model, tokenizer = FastVisionModel.from_pretrained(
     "./qwen2.5-ocr-merged-finetune-merge-16bit", load_in_4bit = True, load_in_8bit = False
 )
 
-# benchmark 4bit loaded, 16bits merged model performance
 model_name = "Unsloth 16bits-merged model load-4bits"
 model.config.use_cache = True
 
@@ -236,12 +224,11 @@ avg_wer, avg_cer = ocr_evaluator.evaluate_model(
 )
 ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
-# load model in 8 bits
+# Load 16bits-merged model in 8 bits
 model, tokenizer = FastVisionModel.from_pretrained(
     "./qwen2.5-ocr-merged-finetune-merge-16bit", load_in_4bit = False, load_in_8bit = True
 )
 
-# benchmark 4bit loaded, 16bits merged model performance
 model_name = "Unsloth 16bits-merged model load-8bits"
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
     model,
@@ -272,7 +259,6 @@ ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 # ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
 # Model comparison report
-# print model comparison
 ocr_evaluator.print_model_comparison()
 
 

tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py

@@ -19,18 +19,15 @@ from tests.utils.cleanup_utils import safe_remove_directory
 from tests.utils.ocr_eval import OCRModelEvaluator
 
 
-## Dataset Preparation
 from datasets import load_dataset
 
 dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", "en", split = "train")
-# To select the first 2000 examples
 train_dataset = dataset.select(range(2000))
 
-# To select the next 200 examples for evaluation
 eval_dataset = dataset.select(range(2000, 2200))
 
 
-# Convert dataset to OAI messages
+# Convert dataset to OAI messages.
 def format_data(sample):
     return {
         "messages": [
@@ -60,12 +57,10 @@ def format_data(sample):
 
 
 system_message = "You are an expert french ocr system."
-# Convert dataset to OAI messages
-# need to use list comprehension to keep Pil.Image type, .map convert image to bytes
+# List comprehension (not .map) keeps PIL.Image type; .map would convert images to bytes.
 train_dataset = [format_data(sample) for sample in train_dataset]
 eval_dataset = [format_data(sample) for sample in eval_dataset]
 
-## Setup OCR main evaluation function and helpers
 import os
 import torch
 from tqdm import tqdm
@@ -73,13 +68,9 @@ import pandas as pd
 from jiwer import wer, cer
 from qwen_vl_utils import process_vision_info
 
-#
 ocr_evaluator = OCRModelEvaluator()
 model_comparison_results = {}
 
-## Finetuning Setup and Run
-# Load Base Model
-
 model, tokenizer = FastVisionModel.from_pretrained(
     model_name = "unsloth/Qwen2-VL-7B-Instruct",
     max_seq_length = 2048,  # Choose any for long context!
@@ -88,7 +79,7 @@ model, tokenizer = FastVisionModel.from_pretrained(
     full_finetuning = False,  # [NEW!] We have full finetuning now!
 )
 
-# benchmark base model performance
+# Benchmark base model performance.
 model_name = "Unsloth Base model"
 FastVisionModel.for_inference(model)
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
@@ -96,7 +87,6 @@ avg_wer, avg_cer = ocr_evaluator.evaluate_model(
 )
 ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
-## Lora Finetuning
 model = FastVisionModel.get_peft_model(
     model,
     finetune_vision_layers = True,  # Turn off for just text!
@@ -134,9 +124,7 @@ trainer = SFTTrainer(
         per_device_train_batch_size = 2,
         gradient_accumulation_steps = 4,
         gradient_checkpointing = True,
-        gradient_checkpointing_kwargs = {
-            "use_reentrant": False
-        },  # use reentrant checkpointing
+        gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
         max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
         warmup_ratio = 0.03,
         # num_train_epochs = 2, # Set this instead of max_steps for full training runs
@@ -161,15 +149,12 @@ trainer = SFTTrainer(
     ),
 )
 
-# run training
 trainer_stats = trainer.train()
 
 model.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter", tokenizer)
 tokenizer.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter")
 
-## Measure Adapter Performance
-
-# benchmark lora model performance
+# Benchmark lora adapter model performance.
 model_name = "Unsloth lora adapter model"
 FastVisionModel.for_inference(model)
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
@@ -177,8 +162,6 @@ avg_wer, avg_cer = ocr_evaluator.evaluate_model(
 )
 ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
-## Merge Model
-
 
 def find_lora_base_model(model_to_inspect):
     current = model_to_inspect
@@ -193,21 +176,17 @@ base = find_lora_base_model(model)
 
 print((base.__class__.__name__))
 
-# merge default 16 bits
+# Merge at default 16 bits.
 model.save_pretrained_merged(
     save_directory = "qwen2-ocr-merged-finetune-merge-16bit", tokenizer = tokenizer
 )
 
 
-## Benchmark merged model performance
-
-### 16 bits merged model
-
+# Benchmark 16-bit merged model loaded at various precisions.
 model, tokenizer = FastVisionModel.from_pretrained(
     "./qwen2-ocr-merged-finetune-merge-16bit", load_in_4bit = False, load_in_8bit = False
 )
 
-# benchmark 4bit loaded, 16bits merged model performance
 model_name = "Unsloth 16bits-merged model load-16bits"
 model.config.use_cache = True
 
@@ -219,12 +198,11 @@ avg_wer, avg_cer = ocr_evaluator.evaluate_model(
 )
 ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
-# load 16bits-merged model in 4 bits
+# Load 16bits-merged model in 4 bits.
 model, tokenizer = FastVisionModel.from_pretrained(
     "./qwen2-ocr-merged-finetune-merge-16bit", load_in_4bit = True, load_in_8bit = False
 )
 
-# benchmark 4bit loaded, 16bits merged model performance
 model_name = "Unsloth 16bits-merged model load-4bits"
 model.config.use_cache = True
 
@@ -236,12 +214,11 @@ avg_wer, avg_cer = ocr_evaluator.evaluate_model(
 )
 ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
-# load model in 8 bits
+# Load 16bits-merged model in 8 bits.
 model, tokenizer = FastVisionModel.from_pretrained(
     "./qwen2-ocr-merged-finetune-merge-16bit", load_in_4bit = False, load_in_8bit = True
 )
 
-# benchmark 4bit loaded, 16bits merged model performance
 model_name = "Unsloth 16bits-merged model load-8bits"
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
     model,
@@ -271,12 +248,9 @@ ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 # avg_wer, avg_cer = ocr_evaluator.evaluate_model(model, tokenizer, eval_dataset, output_dir="unsloth_4bits_merged_model_load_8bits_results")
 # ocr_evaluator.add_to_comparison(model_name, avg_wer, avg_cer)
 
-# Model comparison report
-# print model comparison
 ocr_evaluator.print_model_comparison()
 
 
-# Final cleanup
 print("\n🧹 Cleaning up temporary files...")
 safe_remove_directory("./unsloth-qwen2-7vl-french-ocr-adapter")
 safe_remove_directory("./unsloth-qwen2-7vl-french-ocr-checkpoints")

tests/security/conftest.py

@@ -1,11 +1,8 @@
 """Shared fixtures for the security regression suite.
 
-The scanner scripts under audit are designed to be offline-safe. Pin
-that invariant by autouse-installing a session-scoped network blocker
-that refuses any non-loopback `socket.connect()` from inside the test
-process. If a future test (or a scanner regression) accidentally tries
-to reach the public internet, pytest fails loudly instead of leaking
-the request.
+The scanners under audit must be offline-safe. An autouse session-scoped network
+blocker refuses any non-loopback `socket.connect()` so a regression that reaches
+the internet fails loudly instead of leaking the request.
 """
 
 from __future__ import annotations
@@ -17,8 +14,7 @@ from pathlib import Path
 import pytest
 
 
-# Make `scripts/` importable as a package so tests can grab the scanner
-# constants directly. The repo root sits two levels above this file.
+# Make `scripts/` importable so tests can grab scanner constants directly
 REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
@@ -70,11 +66,8 @@ class _BlockedSocket(socket.socket):
 
 @pytest.fixture(scope = "session", autouse = True)
 def network_blocker():
-    """Session-scoped fixture; replaces `socket.socket` with a blocker.
-
-    Yields nothing; the swap is the side effect. Restored at teardown
-    so other test sessions (run interleaved) see a clean module.
-    """
+    """Session-scoped fixture; swaps `socket.socket` for a blocker, restored at
+    teardown so interleaved sessions see a clean module."""
     original = socket.socket
     socket.socket = _BlockedSocket  # type: ignore[assignment]
     try:

tests/security/fixtures/_build.py

@@ -1,20 +1,12 @@
 """Deterministic builder for the wheel + sdist binary fixtures.
 
-This script is NOT run from CI; the produced .whl / .tar.gz bytes are
-committed alongside it. Re-run only when the IOC literal changes.
+Not run from CI; the produced .whl / .tar.gz bytes are committed alongside
+it. Re-run only when the IOC literal changes.
 
-Determinism strategy
---------------------
-- All member timestamps fixed to `SOURCE_DATE_EPOCH=0` (Unix epoch).
-- All members written with uid=0, gid=0, uname="", gname="".
-- Permission bits fixed: 0o644 for files, 0o755 for directories.
-- Members emitted in sorted order so the archive byte stream does not
-  depend on filesystem iteration order.
-- `zipfile.ZipFile` is invoked with `compresslevel=6` (default DEFLATE)
-  to keep output stable across stdlib versions.
-
-Re-running this script and diffing the .whl bytes against git is the
-regression test for determinism (also asserted in test_scan_packages).
+Determinism: fixed timestamps (SOURCE_DATE_EPOCH=0), uid/gid=0, empty
+uname/gname, fixed perms (0o644 files / 0o755 dirs), sorted member order,
+and DEFLATE compresslevel=6 for stability across stdlib versions. Diffing
+re-built .whl bytes against git is the regression test (see test_scan_packages).
 """
 
 from __future__ import annotations
@@ -33,9 +25,8 @@ _ZIP_DOS_EPOCH = (1980, 1, 1, 0, 0, 0)
 HERE = Path(__file__).resolve().parent
 
 
-# The IOC literal that scan_packages.py must trip on. Keep this in
-# sync with KNOWN_IOC_STRINGS in scripts/scan_npm_packages.py and
-# RE_MAY12_IOC in scripts/scan_packages.py.
+# IOC literal scan_packages.py must trip on. Keep in sync with
+# KNOWN_IOC_STRINGS (scan_npm_packages.py) and RE_MAY12_IOC (scan_packages.py).
 MALICIOUS_SETUP_PY = '''"""Test fixture: do NOT install.
 
 This file embeds the May-12 Mini Shai-Hulud IOC literal so the

tests/security/test_lockfile_supply_chain_audit.py

@@ -51,10 +51,8 @@ def _run_auditor(
 
 
 def test_malicious_lockfile_exits_1(tmp_path):
-    """The malicious fixture combines a non-registry resolved URL, a
-    known IOC substring (`filev2.getsession.org`), and a missing
-    integrity hash. The auditor must refuse with exit 1.
-    """
+    """Fixture combines a non-registry resolved URL, a known IOC substring
+    (`filev2.getsession.org`), and a missing integrity hash: auditor exits 1."""
     fixture = FIXTURES / "malicious_lockfile.json"
     assert fixture.is_file()
     proc = _run_auditor(root = tmp_path, npm_lockfiles = [fixture])
@@ -66,13 +64,9 @@ def test_malicious_lockfile_exits_1(tmp_path):
     assert "non-registry-resolved-url" in combined
     assert "missing-integrity-hash" in combined
     assert "known-ioc-string" in combined
-    # Verify the scanner WROTE the IOC name into its stdout/stderr. The
-    # literal is constructed at runtime so CodeQL's
-    # py/incomplete-url-substring-sanitization rule (which fires on
-    # source-literal + `in` even when the operand is the scanner's own
-    # output, not a URL being sanitized) does not false-positive across
-    # pre-commit reformatting that may split the assert onto multiple
-    # lines and detach an inline lgtm comment from the operator.
+    # IOC literal built at runtime so CodeQL's
+    # py/incomplete-url-substring-sanitization rule doesn't false-positive on
+    # the source-literal + `in` (the operand is the scanner's own output).
     _ioc_host = "filev2." + "getsession.org"
     assert _ioc_host in combined
 
@@ -143,7 +137,6 @@ def test_lockfile_auditor_blocked_versions_match_scanner():
     comment until the next PR factors them into a shared module).
     """
     from scripts import scan_npm_packages as snp
-
     assert (
         lsa.BLOCKED_NPM_VERSIONS == snp.BLOCKED_NPM_VERSIONS
     ), "auditor and scanner BLOCKED_NPM_VERSIONS tables drifted"
@@ -171,15 +164,9 @@ checksum = "0000000000000000000000000000000000000000000000000000000000000000"
 
 
 def test_malicious_cargo_lockfile_refused(tmp_path):
-    """Inline Cargo.lock with `source = "git+https://example.com/..."`
-    must trip the `non-registry-cargo-source` check.
-
-    `non-registry-cargo-source` is an advisory finding kind in the
-    auditor's default mode (per the audit script's BLOCKING_KINDS
-    set). To exercise the historical "refuse to install" behavior we
-    pass --strict here; that promotes every finding to blocking and
-    keeps the test honest about its intent (detection + refusal).
-    """
+    """Inline Cargo.lock with a `git+https://...` source must trip the
+    `non-registry-cargo-source` check. It's advisory in default mode, so
+    --strict promotes it to blocking to exercise the refuse-to-install path."""
     lockfile = tmp_path / "Cargo.lock"
     lockfile.write_text(_MALICIOUS_CARGO_LOCK)
     proc = _run_auditor(
@@ -232,13 +219,9 @@ def test_audit_cargo_lockfile_direct_call(tmp_path):
 
 
 def test_gha_escape_collapses_finding_to_one_line():
-    """`_gha_escape()` must collapse newlines (`%0A`), carriage
-    returns (`%0D`), and percent signs (`%25`) so that
-    `::warning::<msg>` / `::error::<msg>` render the full finding
-    in the GitHub Actions UI annotation instead of being truncated
-    at the first newline. The `%` replacement must happen first or
-    the subsequent `%0A` / `%0D` escapes get double-encoded.
-    """
+    """`_gha_escape()` collapses newlines (`%0A`), carriage returns (`%0D`),
+    and percent signs (`%25`) so GHA annotations aren't truncated at the first
+    newline. `%` must escape first or the `%0A`/`%0D` escapes double-encode."""
     assert lsa._gha_escape("a\nb\nc") == "a%0Ab%0Ac"
     assert lsa._gha_escape("a\rb") == "a%0Db"
     assert lsa._gha_escape("100%") == "100%25"
@@ -261,13 +244,9 @@ def test_gha_escape_collapses_finding_to_one_line():
 
 
 def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
-    """End-to-end check: the `::warning::` line emitted for an
-    advisory finding must be a SINGLE physical line (the rest of
-    the Finding is `%0A`-escaped inside the message). Regression
-    test for the gemini-code-assist review on PR #5604: without
-    `_gha_escape`, GitHub Actions truncates the annotation after
-    `[kind] path` and the package + detail fields never render.
-    """
+    """The `::warning::` line for an advisory finding must be a SINGLE physical
+    line (rest of the Finding `%0A`-escaped). Regression for PR #5604: without
+    `_gha_escape`, GHA truncates after `[kind] path` and drops package/detail."""
     lockfile = tmp_path / "Cargo.lock"
     lockfile.write_text(_MALICIOUS_CARGO_LOCK)
     proc = _run_auditor(
@@ -275,9 +254,7 @@ def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
         npm_lockfiles = [FIXTURES / "clean_lockfile.json"],
         cargo_lockfiles = [lockfile],
     )
-    warning_lines = [
-        line for line in proc.stderr.splitlines() if line.startswith("::warning::")
-    ]
+    warning_lines = [line for line in proc.stderr.splitlines() if line.startswith("::warning::")]
     assert warning_lines, (
         "expected at least one ::warning:: annotation; " f"stderr was:\n{proc.stderr}"
     )
@@ -299,11 +276,9 @@ def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
 
 
 def test_skip_env_var_with_short_value_rejected(tmp_path):
-    """`UNSLOTH_LOCKFILE_AUDIT_SKIP=1` used to silently bypass the
-    audit. Per SF4 it must instead emit a `::warning::` to stderr and
-    fall through to run the audit. A real justification value
-    (>=5 chars, not a boolean shape) is still honored.
-    """
+    """`UNSLOTH_LOCKFILE_AUDIT_SKIP=1` must no longer silently bypass: per SF4
+    it warns and falls through to run the audit. A real justification value
+    (>=5 chars, not a boolean shape) is still honored."""
     fixture = FIXTURES / "clean_lockfile.json"
 
     # Case 1 -- "1" rejected, audit RUNS.

tests/security/test_new_install_scripts.py

@@ -18,7 +18,12 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
 SCRIPT = REPO_ROOT / "scripts" / "check_new_install_scripts.py"
 
 
-def _run(base: Path, head: Path, *, timeout: int = 30) -> subprocess.CompletedProcess:
+def _run(
+    base: Path,
+    head: Path,
+    *,
+    timeout: int = 30,
+) -> subprocess.CompletedProcess:
     return subprocess.run(
         [
             sys.executable,
@@ -39,9 +44,7 @@ def _write(path: Path, content: dict) -> Path:
     return path
 
 
-# ---------------------------------------------------------------------------
-# Lockfile fixtures.
-# ---------------------------------------------------------------------------
+# Lockfile fixtures
 
 
 def _v3_lockfile(packages: dict) -> dict:
@@ -65,9 +68,7 @@ def _v2_lockfile(packages: dict, dependencies: dict) -> dict:
     }
 
 
-# ---------------------------------------------------------------------------
-# Tests.
-# ---------------------------------------------------------------------------
+# Tests
 
 
 def test_no_new_install_scripts_exit_0(tmp_path: Path):
@@ -103,9 +104,7 @@ def test_new_dep_with_postinstall_exits_1(tmp_path: Path):
     head_pkgs = dict(base_pkgs)
     head_pkgs["node_modules/evil-postinstall"] = {
         "version": "1.0.0",
-        "resolved": (
-            "https://registry.npmjs.org/evil-postinstall/-/evil-postinstall-1.0.0.tgz"
-        ),
+        "resolved": ("https://registry.npmjs.org/evil-postinstall/-/evil-postinstall-1.0.0.tgz"),
         "integrity": "sha512-fake",
         "hasInstallScript": True,
     }
@@ -166,8 +165,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
         "node_modules/v2-postinstall-dep": {
             "version": "2.0.0",
             "resolved": (
-                "https://registry.npmjs.org/v2-postinstall-dep/-/"
-                "v2-postinstall-dep-2.0.0.tgz"
+                "https://registry.npmjs.org/v2-postinstall-dep/-/v2-postinstall-dep-2.0.0.tgz"
             ),
             "integrity": "sha512-fake",
             "hasInstallScript": True,
@@ -177,8 +175,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
         "v2-postinstall-dep": {
             "version": "2.0.0",
             "resolved": (
-                "https://registry.npmjs.org/v2-postinstall-dep/-/"
-                "v2-postinstall-dep-2.0.0.tgz"
+                "https://registry.npmjs.org/v2-postinstall-dep/-/v2-postinstall-dep-2.0.0.tgz"
             ),
             "integrity": "sha512-fake",
         },
@@ -187,8 +184,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
     head = _write(tmp_path / "head.json", _v2_lockfile(head_pkgs, head_deps))
     result = _run(base, head)
     assert result.returncode == 1, (
-        f"expected exit 1 for v2 lockfile, got {result.returncode}; "
-        f"stderr:\n{result.stderr}"
+        f"expected exit 1 for v2 lockfile, got {result.returncode}; " f"stderr:\n{result.stderr}"
     )
     assert "v2-postinstall-dep" in result.stderr
 

tests/security/test_scan_npm_packages.py

@@ -46,14 +46,9 @@ def _run_scanner(lockfile: Path, *, timeout: int = 30) -> subprocess.CompletedPr
 
 
 def test_malicious_lockfile_exits_1():
-    """Structural IOCs alone must fail the scanner.
-
-    `structural_only_lockfile.json` contains: (a) a non-registry
-    `resolved` URL (filev2.getsession.org), (b) an entry missing
-    its `integrity` field. Both are caught in `parse_lockfile()`
-    before any tarball download attempt -- so the test is fully
-    offline.
-    """
+    """Structural IOCs alone must fail the scanner. The fixture has a
+    non-registry `resolved` URL and a missing `integrity` field, both caught in
+    `parse_lockfile()` before any tarball download, so the test is offline."""
     fixture = FIXTURES / "structural_only_lockfile.json"
     assert fixture.is_file(), fixture
     proc = _run_scanner(fixture)
@@ -106,16 +101,10 @@ def test_blocked_npm_versions_complete():
     table = snp.BLOCKED_NPM_VERSIONS
     tanstack_keys = [k for k in table if k.startswith("@tanstack/")]
     assert len(tanstack_keys) == 42, (
-        f"expected 42 @tanstack/* entries, got {len(tanstack_keys)}: "
-        f"{sorted(tanstack_keys)}"
+        f"expected 42 @tanstack/* entries, got {len(tanstack_keys)}: " f"{sorted(tanstack_keys)}"
     )
     assert "@opensearch-project/opensearch" in table
-    assert table["@opensearch-project/opensearch"] == {
-        "3.5.3",
-        "3.6.2",
-        "3.7.0",
-        "3.8.0",
-    }
+    assert table["@opensearch-project/opensearch"] == {"3.5.3", "3.6.2", "3.7.0", "3.8.0"}
     squawk = [k for k in table if k.startswith("@squawk/")]
     assert len(squawk) >= 22, (
         f"expected at least 22 @squawk/* entries (full safedep.io enumeration), "
@@ -172,12 +161,8 @@ def test_blocked_npm_versions_complete():
     reason = "Fork 1 (BLOCKED_NPM_VERSIONS pre-fetch hook) not merged yet",
 )
 def test_blocked_npm_versions_short_circuits_download():
-    """With Fork 1's pre-fetch hook, the malicious tanstack entry
-    must produce a `blocked-known-malicious` finding without ever
-    calling out to the npm registry. The full malicious fixture
-    contains the tanstack entry; the test asserts exit 1 and that
-    the new finding pattern appears in scanner output.
-    """
+    """The pre-fetch hook must flag the malicious tanstack entry as
+    `blocked-known-malicious` (exit 1) without hitting the npm registry."""
     fixture = FIXTURES / "malicious_lockfile.json"
     proc = _run_scanner(fixture, timeout = 10)
     assert proc.returncode == 1
@@ -209,10 +194,8 @@ def _extract_pkg_with_ioc(ioc: str, tmp_path: Path) -> Path:
 
 
 def test_every_known_ioc_string_caught(tmp_path):
-    """For every entry in `KNOWN_IOC_STRINGS`, embed the IOC in a
-    one-file package tree and confirm `scan_extracted_tree()`
-    surfaces it. Guards against silent regex / table drift.
-    """
+    """Embed each `KNOWN_IOC_STRINGS` entry in a one-file package tree and
+    confirm `scan_extracted_tree()` surfaces it. Guards against table drift."""
     iocs = snp.KNOWN_IOC_STRINGS
     assert iocs, "KNOWN_IOC_STRINGS unexpectedly empty"
 
@@ -240,10 +223,8 @@ def test_every_known_ioc_string_caught(tmp_path):
 
 
 def test_parse_lockfile_structural_findings():
-    """`parse_lockfile()` returns (entries, structural_findings). The
-    structural-only fixture should produce 2 structural findings and
-    0 entries (because both bad entries are `continue`d).
-    """
+    """The structural-only fixture yields 2 structural findings and 0 entries
+    (both bad entries are `continue`d in `parse_lockfile()`)."""
     entries, struct = snp.parse_lockfile(FIXTURES / "structural_only_lockfile.json")
     assert entries == []
     patterns = {f.pattern for f in struct}

tests/security/test_scan_packages.py

@@ -1,9 +1,8 @@
 """Regression tests for `scripts/scan_packages.py`.
 
-The scanner's primary entry point (`download_packages`) reaches PyPI;
-to keep the suite offline we exercise it via the module's public
-in-process helpers (`scan_archive`) and assert against the binary
-wheel / sdist fixtures committed under `tests/security/fixtures/`.
+`download_packages` reaches PyPI; to stay offline we drive the in-process
+`scan_archive` helper against the wheel/sdist fixtures under
+`tests/security/fixtures/`.
 """
 
 from __future__ import annotations
@@ -37,9 +36,8 @@ def test_fixture_files_exist():
 def test_fixture_bytes_are_deterministic(tmp_path):
     """Re-running `_build.py` must produce byte-identical archives.
 
-    The build helper sets every member's mtime/uid/gid/mode and emits
-    members in sorted order. We rebuild into a temp dir and compare
-    SHA-256 against the committed bytes.
+    The build helper pins each member's mtime/uid/gid/mode and sorts members.
+    Rebuild into a temp dir and compare SHA-256 against the committed bytes.
     """
     # Snapshot committed hashes.
     expected: dict[str, str] = {}
@@ -53,7 +51,7 @@ def test_fixture_bytes_are_deterministic(tmp_path):
     builder_src = (FIXTURES / "_build.py").read_text()
     rebuilt_helper = rebuild_dir / "_build.py"
     rebuilt_helper.write_text(builder_src)
-    # Run with SOURCE_DATE_EPOCH=0 and HERE-override via a tiny shim.
+    # Run with SOURCE_DATE_EPOCH=0 and HERE override via a tiny shim.
     shim = rebuild_dir / "run.py"
     shim.write_text(
         "import sys, pathlib\n"
@@ -118,9 +116,7 @@ def test_clean_wheel_no_findings():
         str(FIXTURES / "clean_wheel.whl"),
         "clean_fixture",
     )
-    assert (
-        findings == []
-    ), f"unexpected findings on clean wheel: {[str(f) for f in findings]}"
+    assert findings == [], f"unexpected findings on clean wheel: {[str(f) for f in findings]}"
 
 
 # ---------------------------------------------------------------------------
@@ -170,20 +166,16 @@ def test_re_may12_ioc_catches_each_literal():
     reason = "Fork 1 (RE_MAY12_IOC integration) not merged yet",
 )
 def test_may12_ioc_caught_by_scan_archive():
-    """Once RE_MAY12_IOC is wired into check_py_file (per Fork 1's
-    plan), the malicious wheel's setup.py must produce a finding
-    that explicitly references the May-12 IOC string.
+    """Once RE_MAY12_IOC is wired into check_py_file, the malicious wheel's
+    setup.py must produce a finding referencing the May-12 IOC string.
     """
     findings = sp.scan_archive(
         str(FIXTURES / "malicious_wheel.whl"),
         "malicious_fixture",
     )
-    # The IOC literals are built at runtime so CodeQL's
-    # py/incomplete-url-substring-sanitization rule does not false-
-    # positive on the (literal `in` operand) pattern -- the operand is
-    # the scanner's own evidence string, not a URL being sanitized.
-    # Runtime construction also survives pre-commit reformatting that
-    # would otherwise detach an inline lgtm comment from the operator.
+    # IOC literals built at runtime so CodeQL's url-substring-sanitization
+    # rule doesn't false-positive on the `in` operand (it's evidence, not a
+    # URL), and so reformatting can't detach an inline lgtm comment.
     _ioc_host = "git-tanstack." + "com"
     _ioc_drop = "transformers." + "pyz"
     hit = any(
@@ -204,16 +196,11 @@ def test_may12_ioc_caught_by_scan_archive():
 
 
 def test_scan_packages_pip_download_failure_propagates(tmp_path):
-    """A pip download failure must NOT be silently swallowed into a
-    `0 findings, exit 0` report. Item (4) of the silent-failure
-    hardening: an obviously unresolvable spec is fed to the scanner
-    as a subprocess; the orchestrator must exit 2 (scan incomplete)
-    and the stderr must carry the SCAN INCOMPLETE banner.
+    """A pip download failure must NOT be swallowed into `0 findings, exit 0`.
 
-    The spec name is deliberately long + random-looking so it cannot
-    accidentally resolve on any real package index. We do not rely on
-    network reachability: even an offline runner will get a clean
-    "could not resolve" failure from pip.
+    Feeds an unresolvable spec to the scanner subprocess; it must exit 2
+    (scan incomplete) with a SCAN INCOMPLETE banner on stderr. The spec name
+    is long/random so it can't resolve on any index, even offline.
     """
     script = REPO_ROOT / "scripts" / "scan_packages.py"
     assert script.is_file(), script
@@ -235,9 +222,8 @@ def test_scan_packages_pip_download_failure_propagates(tmp_path):
 
 
 def test_archive_corruption_produces_critical_finding(tmp_path):
-    """SF1: a corrupted wheel (truncated bytes) used to be silently
-    skipped by `except Exception: continue` inside iter_archive_files.
-    It must now yield a CRITICAL `archive_corrupted` finding.
+    """SF1: a corrupted wheel was silently skipped by `except: continue` in
+    iter_archive_files; it must now yield a CRITICAL `archive_corrupted`.
     """
     bad = tmp_path / "broken-0.0.1-py3-none-any.whl"
     bad.write_bytes(b"X")  # 1-byte "wheel" -- not a valid zip container
@@ -245,8 +231,7 @@ def test_archive_corruption_produces_critical_finding(tmp_path):
     assert findings, "scan_archive returned 0 findings on corrupt wheel"
     corrupted = [f for f in findings if f.check == "archive_corrupted"]
     assert corrupted, (
-        "no archive_corrupted finding; got "
-        f"{[(f.severity, f.check) for f in findings]}"
+        "no archive_corrupted finding; got " f"{[(f.severity, f.check) for f in findings]}"
     )
     assert all(f.severity == sp.CRITICAL for f in corrupted)
 

tests/sh/test_get_torch_index_url.sh

@@ -13,6 +13,10 @@ FAIL=0
 _FUNC_FILE=$(mktemp)
 _FAKE_SMI_DIR=$(mktemp -d)
 {
+    sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH"
+    echo ""
+    sed -n '/^_cvd_hides_nvidia()/,/^}/p' "$INSTALL_SH"
+    echo ""
     sed -n '/^_has_amd_rocm_gpu()/,/^}/p' "$INSTALL_SH"
     echo ""
     sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
@@ -59,6 +63,29 @@ MOCK
     echo "$_dir"
 }
 
+# Helper: create a mock nvidia-smi that prints the new "CUDA UMD Version" header
+# layout used by newer NVIDIA drivers (e.g. 610.x on Windows).  See issue #5812.
+make_mock_smi_umd() {
+    _dir=$(mktemp -d)
+    cat > "$_dir/nvidia-smi" <<MOCK
+#!/bin/sh
+case "\$1" in
+    -L)
+        echo "GPU 0: NVIDIA GeForce RTX 5090 Laptop GPU (UUID: GPU-fake-uuid)"
+        ;;
+    *)
+        cat <<'SMI_OUT'
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 610.47                 KMD Version: 610.47        CUDA UMD Version: $1     |
++-----------------------------------------------------------------------------------------+
+SMI_OUT
+        ;;
+esac
+MOCK
+    chmod +x "$_dir/nvidia-smi"
+    echo "$_dir"
+}
+
 # Helper: create a mock amd-smi that prints a given ROCm version string
 # Supports both "amd-smi version" and "amd-smi list" subcommands so that
 # the GPU presence check (amd-smi list) also succeeds in tests.
@@ -84,7 +111,7 @@ MOCK
 # Build a minimal tools directory with symlinks to essential commands
 # (uname, grep, head, etc.) but WITHOUT nvidia-smi or amd-smi.
 _TOOLS_DIR=$(mktemp -d)
-for _cmd in uname grep sed head sh bash cat awk printf; do
+for _cmd in uname grep sed head sh bash cat awk printf tr; do
     _real=$(command -v "$_cmd" 2>/dev/null || true)
     [ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd"
 done
@@ -93,12 +120,19 @@ done
 # $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test
 run_func() {
     _mock_dir="$1"
+    # Default: strip CUDA_VISIBLE_DEVICES so the host environment cannot leak
+    # in; a second argument sets it explicitly (hidden-GPU scenarios).
+    if [ "$#" -ge 2 ]; then
+        _cvd_setup="export CUDA_VISIBLE_DEVICES='$2'"
+    else
+        _cvd_setup="unset CUDA_VISIBLE_DEVICES"
+    fi
     if [ "$_mock_dir" = "none" ]; then
         # Minimal PATH with only basic tools, no nvidia-smi anywhere
-        PATH="$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
+        PATH="$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
     else
         # Put mock nvidia-smi dir first, then basic tools
-        PATH="$_mock_dir:$_TOOLS_DIR" bash -c ". '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
+        PATH="$_mock_dir:$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
     fi
 }
 
@@ -170,10 +204,10 @@ _result=$(run_func "$_dir")
 assert_eq "ROCm 7.1 -> rocm7.1" "https://download.pytorch.org/whl/rocm7.1" "$_result"
 rm -rf "$_dir"
 
-# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.1 (capped due to torch <2.11.0)
+# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.2
 _dir=$(make_mock_amd_smi "7.2")
 _result=$(run_func "$_dir")
-assert_eq "ROCm 7.2 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result"
+assert_eq "ROCm 7.2 -> rocm7.2" "https://download.pytorch.org/whl/rocm7.2" "$_result"
 rm -rf "$_dir"
 
 # 12) Both nvidia-smi and amd-smi present -> CUDA takes precedence
@@ -208,10 +242,10 @@ _result=$(run_func "$_dir")
 assert_eq "ROCm 7.0 -> rocm7.0" "https://download.pytorch.org/whl/rocm7.0" "$_result"
 rm -rf "$_dir"
 
-# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.1 (capped)
+# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.2 (capped to latest known)
 _dir=$(make_mock_amd_smi "8.0")
 _result=$(run_func "$_dir")
-assert_eq "ROCm 8.0 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result"
+assert_eq "ROCm 8.0 -> rocm7.2 (capped)" "https://download.pytorch.org/whl/rocm7.2" "$_result"
 rm -rf "$_dir"
 
 # 18) Malformed amd-smi output (empty version field) -> cpu
@@ -278,6 +312,71 @@ assert_eq "empty mirror env -> official/cpu" "https://download.pytorch.org/whl/c
 _result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl/" run_func "none")
 assert_eq "trailing slash stripped -> mirror/cpu" "https://mirror.example.com/whl/cpu" "$_result"
 
+# 29) "CUDA UMD Version: 13.3" header (newer NVIDIA driver layout, issue #5812)
+#     -> cu130, not the cu126 fallback.
+_dir=$(make_mock_smi_umd "13.3")
+_result=$(run_func "$_dir")
+assert_eq "CUDA UMD Version 13.3 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
+rm -rf "$_dir"
+
+# 30) "CUDA UMD Version: 12.8" header (newer layout on a 12.x driver) -> cu128
+_dir=$(make_mock_smi_umd "12.8")
+_result=$(run_func "$_dir")
+assert_eq "CUDA UMD Version 12.8 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+rm -rf "$_dir"
+
+# 31) "CUDA UMD Version: 11.8" header (newer layout on an older driver) -> cu118
+_dir=$(make_mock_smi_umd "11.8")
+_result=$(run_func "$_dir")
+assert_eq "CUDA UMD Version 11.8 -> cu118" "https://download.pytorch.org/whl/cu118" "$_result"
+rm -rf "$_dir"
+
+# 32) Driver-reported "CUDA Version: 13.3" (legacy header) -> cu130.
+_dir=$(make_mock_smi "13.3")
+_result=$(run_func "$_dir")
+assert_eq "CUDA Version 13.3 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
+rm -rf "$_dir"
+
+# 33) "CUDA Version: 13.7" -> cu130 (until a cu137 wheel index exists).
+_dir=$(make_mock_smi "13.7")
+_result=$(run_func "$_dir")
+assert_eq "CUDA Version 13.7 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
+rm -rf "$_dir"
+
+# 34) CUDA_VISIBLE_DEVICES="" hides the NVIDIA GPU -> cpu (no AMD present)
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" "")
+assert_eq "CVD='' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+rm -rf "$_dir"
+
+# 35) CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU -> cpu (no AMD present)
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" "-1")
+assert_eq "CVD=-1 hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+rm -rf "$_dir"
+
+# 36) Mixed AMD+NVIDIA host with NVIDIA hidden -> ROCm route is restored
+_cuda_dir=$(make_mock_smi "12.6")
+_amd_dir=$(make_mock_amd_smi "6.4")
+_combined_dir=$(mktemp -d)
+ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi"
+ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi"
+_result=$(run_func "$_combined_dir" "-1")
+assert_eq "CUDA+ROCm with CVD=-1 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result"
+rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir"
+
+# 37) CUDA_VISIBLE_DEVICES=0 (a visible device) must NOT hide the GPU
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" "0")
+assert_eq "CVD=0 keeps NVIDIA -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+rm -rf "$_dir"
+
+# 38) Whitespace-padded "-1" still hides the GPU
+_dir=$(make_mock_smi "12.8")
+_result=$(run_func "$_dir" " -1 ")
+assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+rm -rf "$_dir"
+
 rm -f "$_FUNC_FILE"
 rm -rf "$_FAKE_SMI_DIR"
 rm -rf "$_TOOLS_DIR"

tests/sh/test_mac_intel_compat.sh

@@ -563,6 +563,80 @@ else
     FAIL=$((FAIL + 1))
 fi
 
+echo ""
+echo "=== Apple Silicon x86_64 (Rosetta) venv rebuild ==="
+
+# Extract the real guard block from install.sh so we exercise the shipped logic
+# (comment header down to its column-0 closing fi).
+_GUARD_FILE=$(mktemp)
+awk '/Guard against two independent Apple Silicon venv problems/{f=1} f{print} f&&/^fi$/{exit}' \
+    "$INSTALL_SH" > "$_GUARD_FILE"
+
+if [ ! -s "$_GUARD_FILE" ]; then
+    echo "  FAIL: could not extract Apple Silicon venv guard from install.sh"
+    FAIL=$((FAIL + 1))
+else
+    # Runner: stub uv (via run_install_cmd) + a fake venv python, source the
+    # guard, then print "<final_arch> <final_ver> | <recreate_selectors>".
+    # The stub maps a uv arm64 selector to the interpreter uv would produce:
+    # cpython-3.12-* -> arm64 3.12.7, cpython-3.13-* -> arm64 $REBUILD_313_VERSION.
+    _RUNNER=$(mktemp)
+    cat > "$_RUNNER" << 'RUNNER_EOF'
+GUARD="$1"; VENV_DIR="$2"
+make_python() {  # dir machine version
+    mkdir -p "$1/bin"
+    printf '#!/usr/bin/env bash\necho "%s %s"\n' "$2" "$3" > "$1/bin/python"
+    chmod +x "$1/bin/python"
+}
+RECREATE_LOG=$(mktemp); : > "$RECREATE_LOG"
+run_install_cmd() {
+    shift  # drop the human label
+    if [ "$1" = "uv" ] && [ "$2" = "venv" ]; then
+        dir="$3"; sel=""; shift 3
+        while [ $# -gt 0 ]; do [ "$1" = "--python" ] && { sel="$2"; shift; }; shift; done
+        echo "$sel" >> "$RECREATE_LOG"
+        case "$sel" in
+            *3.12-macos-aarch64*) make_python "$dir" arm64 "3.12.7" ;;
+            *3.13-macos-aarch64*) make_python "$dir" arm64 "${REBUILD_313_VERSION:-3.13.3}" ;;
+            *)                    make_python "$dir" arm64 "$sel" ;;
+        esac
+    fi
+}
+[ "$INIT_ARCH" != none ] && make_python "$VENV_DIR" "$INIT_ARCH" "$INIT_VER"
+PYTHON_VERSION="3.13"
+. "$GUARD" >&2  # guard's user-facing echoes go to stderr; keep stdout clean
+final="none"; [ -x "$VENV_DIR/bin/python" ] && final="$("$VENV_DIR/bin/python" -c x)"
+printf '%s | %s\n' "$final" "$(paste -sd, "$RECREATE_LOG" 2>/dev/null)"
+rm -f "$RECREATE_LOG"
+RUNNER_EOF
+
+    _run_guard() {  # _USER_PYTHON OS _ARCH INIT_ARCH INIT_VER REBUILD_313_VERSION
+        _vd=$(mktemp -d)
+        env _USER_PYTHON="$1" OS="$2" _ARCH="$3" INIT_ARCH="$4" INIT_VER="$5" \
+            REBUILD_313_VERSION="$6" bash "$_RUNNER" "$_GUARD_FILE" "$_vd/venv"
+        rm -rf "$_vd"
+    }
+
+    assert_eq "clean arm64 venv left untouched" \
+        "arm64 3.13.3 | " "$(_run_guard '' macos arm64 arm64 3.13.3 '')"
+    assert_eq "x86_64 venv rebuilt as arm64" \
+        "arm64 3.13.3 | cpython-3.13-macos-aarch64-none" \
+        "$(_run_guard '' macos arm64 x86_64 3.13.3 '')"
+    assert_eq "x86_64 venv that lands on 3.13.8 is rebuilt then downgraded to 3.12" \
+        "arm64 3.12.7 | cpython-3.13-macos-aarch64-none,cpython-3.12-macos-aarch64-none" \
+        "$(_run_guard '' macos arm64 x86_64 3.13.3 3.13.8)"
+    assert_eq "arm64 3.13.8 venv downgraded to 3.12" \
+        "arm64 3.12.7 | cpython-3.12-macos-aarch64-none" \
+        "$(_run_guard '' macos arm64 arm64 3.13.8 '')"
+    assert_eq "--python override skips the guard entirely" \
+        "x86_64 3.13.3 | " "$(_run_guard 3.11 macos arm64 x86_64 3.13.3 '')"
+    assert_eq "x86_64 host (Intel/Rosetta shell) is a no-op here" \
+        "x86_64 3.13.3 | " "$(_run_guard '' macos x86_64 x86_64 3.13.3 '')"
+
+    rm -f "$_RUNNER"
+fi
+rm -f "$_GUARD_FILE"
+
 echo ""
 echo "Results: $PASS passed, $FAIL failed"
 [ "$FAIL" -eq 0 ] || exit 1

tests/sh/test_nvcc_meets_llama_minimum.sh

@@ -0,0 +1,121 @@
+#!/bin/bash
+# Unit tests for _nvcc_meets_llama_minimum() from studio/setup.sh.
+# llama.cpp needs CUDA toolkit >= 12.4 (#4437); setup.ps1 aborts via #4517,
+# the Linux side was silent until this fix.
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh"
+PASS=0
+FAIL=0
+
+# Extract just the helper function. The sed range is the same pattern the
+# install.sh tests use.
+_FUNC_FILE=$(mktemp)
+sed -n '/^_nvcc_meets_llama_minimum()/,/^}/p' "$SETUP_SH" > "$_FUNC_FILE"
+
+assert_eq() {
+    _label="$1"; _expected="$2"; _actual="$3"
+    if [ "$_actual" = "$_expected" ]; then
+        echo "  PASS: $_label"
+        PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $_label (expected '$_expected', got '$_actual')"
+        FAIL=$((FAIL + 1))
+    fi
+}
+
+# Fake nvcc printing "release X.Y" in the canonical nvcc -V layout (the helper
+# greps for "release X.Y", stable across CUDA 9.x-13.x).
+make_mock_nvcc() {
+    _ver=$1
+    _dir=$(mktemp -d)
+    cat > "$_dir/nvcc" <<MOCK
+#!/bin/sh
+cat <<NV
+nvcc: NVIDIA (R) Cuda compiler driver
+Copyright (c) 2005-2026 NVIDIA Corporation
+Cuda compilation tools, release $_ver, V${_ver}.0
+NV
+MOCK
+    chmod +x "$_dir/nvcc"
+    echo "$_dir/nvcc"
+}
+
+run_check() {
+    _nvcc=$1
+    bash -c ". '$_FUNC_FILE'; _nvcc_meets_llama_minimum '$_nvcc'"
+}
+
+echo "=== test_nvcc_meets_llama_minimum ==="
+
+# 1) CUDA 12.4 is the minimum supported -> ok
+_bin=$(make_mock_nvcc "12.4")
+_out=$(run_check "$_bin")
+assert_eq "12.4 status" "ok" "$(echo "$_out" | sed -n '1p')"
+assert_eq "12.4 version" "12.4" "$(echo "$_out" | sed -n '2p')"
+rm -rf "$(dirname "$_bin")"
+
+# 2) CUDA 12.3 is the highest version that should be rejected.
+_bin=$(make_mock_nvcc "12.3")
+_out=$(run_check "$_bin")
+assert_eq "12.3 status" "too_old" "$(echo "$_out" | sed -n '1p')"
+rm -rf "$(dirname "$_bin")"
+
+# 3) CUDA 12.1 (matches the original bug report in #4437).
+_bin=$(make_mock_nvcc "12.1")
+_out=$(run_check "$_bin")
+assert_eq "12.1 status" "too_old" "$(echo "$_out" | sed -n '1p')"
+rm -rf "$(dirname "$_bin")"
+
+# 4) CUDA 11.8 -> too_old (anything < 12.0 is rejected).
+_bin=$(make_mock_nvcc "11.8")
+_out=$(run_check "$_bin")
+assert_eq "11.8 status" "too_old" "$(echo "$_out" | sed -n '1p')"
+rm -rf "$(dirname "$_bin")"
+
+# 5) CUDA 12.8 -> ok (mid-range supported).
+_bin=$(make_mock_nvcc "12.8")
+_out=$(run_check "$_bin")
+assert_eq "12.8 status" "ok" "$(echo "$_out" | sed -n '1p')"
+rm -rf "$(dirname "$_bin")"
+
+# 6) CUDA 13.0 -> ok.
+_bin=$(make_mock_nvcc "13.0")
+_out=$(run_check "$_bin")
+assert_eq "13.0 status" "ok" "$(echo "$_out" | sed -n '1p')"
+rm -rf "$(dirname "$_bin")"
+
+# 7) CUDA 13.3 -> ok (the freshly shipped toolkit this fix targets).
+_bin=$(make_mock_nvcc "13.3")
+_out=$(run_check "$_bin")
+assert_eq "13.3 status" "ok" "$(echo "$_out" | sed -n '1p')"
+assert_eq "13.3 version" "13.3" "$(echo "$_out" | sed -n '2p')"
+rm -rf "$(dirname "$_bin")"
+
+# 8) Future CUDA 14.0 -> ok (no upper bound).
+_bin=$(make_mock_nvcc "14.0")
+_out=$(run_check "$_bin")
+assert_eq "14.0 status" "ok" "$(echo "$_out" | sed -n '1p')"
+rm -rf "$(dirname "$_bin")"
+
+# 9) Empty argument -> unknown (defensive; never block the build on detection).
+_out=$(run_check "")
+assert_eq "empty path status" "unknown" "$(echo "$_out" | sed -n '1p')"
+
+# 10) Mock nvcc that prints garbage -> unknown.
+_dir=$(mktemp -d)
+cat > "$_dir/nvcc" <<'MOCK'
+#!/bin/sh
+echo "totally not nvcc output"
+MOCK
+chmod +x "$_dir/nvcc"
+_out=$(run_check "$_dir/nvcc")
+assert_eq "garbage output status" "unknown" "$(echo "$_out" | sed -n '1p')"
+rm -rf "$_dir"
+
+rm -f "$_FUNC_FILE"
+
+echo ""
+echo "Results: $PASS passed, $FAIL failed"
+[ "$FAIL" -eq 0 ] || exit 1

tests/studio/_playwright_robust.py

@@ -1,21 +1,17 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Shared robustness helpers for the Studio Playwright tests.
+"""Shared robustness helpers for the Studio Playwright tests, the single
+point of truth for the CI-runner workarounds (Chromium flags, view-transition
+killer, page recovery, post-action response wait) that both
+`playwright_chat_ui.py` and `playwright_extra_ui.py` need.
 
-Both `playwright_chat_ui.py` and `playwright_extra_ui.py` re-implemented
-the same set of CI-runner workarounds (Chromium launch flags, view-
-transition CSS killer, change-password retry / page-recovery, post-
-action response wait). When one diverged the other slowly rotted; the
-mac/win/linux failure modes are mostly identical so the cure is the
-same. This module is the single point of truth.
-
-Importable directly by the standalone scripts via:
+Importable directly by the standalone scripts:
 
     sys.path.insert(0, str(Path(__file__).parent))
     from _playwright_robust import (...)
 
-It does NOT depend on pytest -- both consumers run as plain Python.
+Does NOT depend on pytest -- both consumers run as plain Python.
 """
 
 from __future__ import annotations
@@ -30,28 +26,16 @@ import urllib.request
 from pathlib import Path
 from typing import Any, Callable
 
-# ─────────────────────────────────────────────────────────────────────
 # Chromium launch args.
-# ─────────────────────────────────────────────────────────────────────
 #
-# Base set works on every CI runner. The four "throttling" flags fight
-# Chromium's tendency to deprioritise CPU + timers when it thinks the
-# window is backgrounded -- which CI runners routinely flag because
-# the headless context has no real focus. Without these, gemma-3-270m
-# inference on Mac slowed to a crawl mid-test (run 25586583024 had a
-# turn budget that never released the Stop button) and the React
-# render queue stalled long enough for `wait_for_function` waits to
-# crowd their per-turn budget.
+# The throttling flags stop Chromium deprioritising CPU/timers when it thinks
+# the headless window is backgrounded (run 25586583024 stalled gemma-3-270m
+# inference and the React render queue mid-test). TranslateUI strips a popup
+# that intercepts pointer events; ipc-flooding-protection off lets rapid clicks
+# through during the slider sweep.
 #
-# `--disable-features=TranslateUI` strips the translate prompt that
-# occasionally adds a popup which intercepts pointer events.
-# `--disable-ipc-flooding-protection` lets us send rapid-fire clicks
-# during the slider sweep without Chromium queuing them.
-#
-# `--single-process` is darwin-only. On Mac it is the documented free-
-# runner fix for the pipeTransport.js JSON-RPC crash; on Win/Linux it
-# strictly destabilises the renderer-isolation safety net so any
-# crash takes the whole context down.
+# `--single-process` is darwin-only: the documented free-runner fix for the
+# pipeTransport.js JSON-RPC crash; on Win/Linux it destabilises the renderer.
 _BASE_CHROMIUM_ARGS = (
     "--disable-dev-shm-usage",
     "--no-sandbox",
@@ -65,11 +49,8 @@ _BASE_CHROMIUM_ARGS = (
 
 
 def chromium_launch_args(platform: str | None = None) -> list[str]:
-    """Return the Chromium launch arg list appropriate for `platform`.
-
-    Defaults to the running interpreter's `sys.platform`. Pass a
-    string to test the darwin branch on Linux.
-    """
+    """Return Chromium launch args for `platform` (defaults to `sys.platform`;
+    pass a string to test the darwin branch on Linux)."""
     p = sys.platform if platform is None else platform
     args = list(_BASE_CHROMIUM_ARGS)
     if p == "darwin":
@@ -77,19 +58,14 @@ def chromium_launch_args(platform: str | None = None) -> list[str]:
     return args
 
 
-# ─────────────────────────────────────────────────────────────────────
 # Init scripts injected into every Playwright context.
-# ─────────────────────────────────────────────────────────────────────
 #
-# CSS view-transitions are otherwise rendered as a full-window
-# pseudo-element that intercepts pointer events for a beat after each
-# theme/route swap. Even with `reduced_motion = "reduce"` set on the
-# context, Studio's components run their own startViewTransition() in
-# a few places (theme toggle, sidebar collapse) and Playwright's
-# actionability check then reports `<html> intercepts pointer events`
-# on the next click. Killing the pseudo-elements + monkey-patching
-# document.startViewTransition into a synchronous shim removes both
-# failure modes. Idempotent and safe to install on every page.
+# CSS view-transitions render a full-window pseudo-element that intercepts
+# pointer events for a beat after each theme/route swap, so Playwright reports
+# `<html> intercepts pointer events` on the next click (even with
+# reduced_motion, since Studio calls startViewTransition() directly). Killing
+# the pseudo-elements + shimming startViewTransition synchronously fixes both.
+# Idempotent and safe to install on every page.
 _VIEW_TRANSITION_KILLER_JS = """
 (function () {
     try {
@@ -130,16 +106,11 @@ def install_view_transition_killer(ctx: Any) -> None:
     ctx.add_init_script(_VIEW_TRANSITION_KILLER_JS)
 
 
-# ─────────────────────────────────────────────────────────────────────
 # Server health pre-flight.
-# ─────────────────────────────────────────────────────────────────────
 #
-# Both workflows already wait for /api/health at the bash level before
-# launching the Python script, but the macos-14 free runner has been
-# observed to surface a brief window where /api/health responds 200
-# yet /api/auth endpoints still 503 because the auth DB hasn't
-# finished migrating. A second probe inside the script catches that
-# narrow gap before we sink 60s into a change-password timeout.
+# The bash wait already gates on /api/health, but on the macos-14 free runner
+# /api/health can return 200 while /api/auth still 503s (auth DB mid-migration).
+# A second in-script probe catches that gap before a 60s change-password timeout.
 
 
 def _http_get_status_and_body(url: str, timeout: float) -> tuple[int, dict | None]:
@@ -162,13 +133,9 @@ def wait_for_health(
     timeout: float = 30.0,
     info: Callable[[str], None] | None = None,
 ) -> bool:
-    """Poll {base_url}/api/health until status==200 with healthy body.
-
-    Returns True on success, False on timeout. Never raises -- the
-    caller decides whether to fail. The test scripts use the boolean
-    only for diagnostic logging, since the workflow's own /api/health
-    wait is the authoritative gate.
-    """
+    """Poll {base_url}/api/health until status==200. Returns True on success,
+    False on timeout; never raises. Diagnostic only -- the workflow's own
+    /api/health wait is the authoritative gate."""
     deadline = time.monotonic() + timeout
     last_status: int | None = None
     last_body: dict | None = None
@@ -182,9 +149,7 @@ def wait_for_health(
         # but accept any 200 -- different Studio builds report differently.
         if status == 200:
             if info is not None:
-                info(
-                    f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}"
-                )
+                info(f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}")
             return True
         time.sleep(0.5)
     if info is not None:
@@ -195,14 +160,11 @@ def wait_for_health(
     return False
 
 
-# ─────────────────────────────────────────────────────────────────────
 # Page recovery.
-# ─────────────────────────────────────────────────────────────────────
 #
-# The single canonical "did the page die mid-test" recovery path. Used
-# by every retry block in both scripts. If the page is closed, opens a
-# fresh one in the same context (auth state in localStorage survives);
-# otherwise leaves the page alone. Optionally re-navigates.
+# Canonical "did the page die mid-test" path used by every retry block. If the
+# page is closed, opens a fresh one in the same context (localStorage auth
+# survives); otherwise leaves it alone. Optionally re-navigates.
 
 
 def recover_or_replace_page(
@@ -214,13 +176,9 @@ def recover_or_replace_page(
     settle_networkidle: bool = True,
     info: Callable[[str], None] | None = None,
 ) -> Any:
-    """Return a usable page. Replaces `page` if it is closed.
-
-    If `goto_url` is provided, navigates the (possibly new) page there
-    and best-effort waits for networkidle. Errors during recovery are
-    logged through `info` (if provided) and swallowed -- the caller
-    handles a still-broken page on the next retry iteration.
-    """
+    """Return a usable page, replacing `page` if it is closed. If `goto_url` is
+    given, navigates there and best-effort waits for networkidle. Recovery
+    errors are logged and swallowed; the caller retries a still-broken page."""
     try:
         if page.is_closed():
             page = ctx.new_page()
@@ -230,9 +188,7 @@ def recover_or_replace_page(
             info(f"recovery: page.is_closed() check failed: {exc!r}")
     if goto_url is not None:
         try:
-            page.goto(
-                goto_url, wait_until = "domcontentloaded", timeout = default_timeout_ms
-            )
+            page.goto(goto_url, wait_until = "domcontentloaded", timeout = default_timeout_ms)
             if settle_networkidle:
                 try:
                     page.wait_for_load_state("networkidle", timeout = 30_000)
@@ -258,15 +214,10 @@ def click_and_wait_for_response(
     timeout_ms: int = 30_000,
     info: Callable[[str], None] | None = None,
 ) -> tuple[int | None, Exception | None]:
-    """Click + wait for the matching XHR/fetch response in one step.
-
-    Returns (status, err). On success: (status, None). On failure to
-    capture the response: (None, exception). Callers typically check
-    `status >= 400` to surface a server-side rejection immediately
-    rather than discovering it 60s later via a downstream wait_for.
-    Falls back to a fire-and-forget click on any wait error so the
-    outer retry loop still runs.
-    """
+    """Click + wait for the matching XHR/fetch response. Returns (status, None)
+    on success or (None, exception) on capture failure. Callers check
+    `status >= 400` to surface server rejections immediately. Falls back to a
+    fire-and-forget click on any wait error so the outer retry loop runs."""
     try:
         with page.expect_response(
             lambda r: url_substr in r.url and r.request.method == method,
@@ -288,19 +239,13 @@ def click_and_wait_for_response(
         return None, exc
 
 
-# ─────────────────────────────────────────────────────────────────────
 # Console-error / page-error filtering.
-# ─────────────────────────────────────────────────────────────────────
 #
-# Two categories:
-#   - BENIGN_PAGE_ERROR_PATTERNS: thrown JS errors that fire as a side
-#     effect of slow CI infra (server timeouts, request races) and have
-#     no user-visible consequence. The page-error gate at the end of
-#     each test should NOT count these.
-#   - BENIGN_CONSOLE_ERROR_PATTERNS: console.error events that fire
-#     for the same reason. Tests don't gate on console.error today
-#     (they only count for diagnostics), but the same list is useful
-#     for filtering noise out of the diagnostic dumps.
+#   - BENIGN_PAGE_ERROR_PATTERNS: JS errors from slow CI infra (timeouts,
+#     request races) with no user-visible effect; the page-error gate must not
+#     count these.
+#   - BENIGN_CONSOLE_ERROR_PATTERNS: same-cause console.error events, used only
+#     to filter noise from diagnostic dumps (tests don't gate on console.error).
 
 BENIGN_PAGE_ERROR_PATTERNS: tuple[str, ...] = (
     "Request failed (422)",
@@ -312,21 +257,15 @@ BENIGN_PAGE_ERROR_PATTERNS: tuple[str, ...] = (
 )
 
 BENIGN_CONSOLE_ERROR_PATTERNS: tuple[str, ...] = (
-    # macos-14 free runner buffer-exhaustion under --single-process
-    # Chromium. The browser surfaces this on resource fetches but the
-    # test catches the underlying request failure via expect_response
-    # and retries; the console line itself is informational.
+    # macos-14 buffer-exhaustion under --single-process; the test catches the
+    # underlying request failure via expect_response and retries.
     "net::ERR_NO_BUFFER_SPACE",
-    # Chromium emits a console.error every time a fetch is aborted,
-    # even when the abort is intentional (component unmount, route
-    # change). All four scripts trigger several of these per run.
+    # Intentional fetch aborts (unmount, route change) log a console.error.
     "AbortError",
     "The user aborted a request",
-    # Same shape: lazy-loaded chunk that's no longer needed because
-    # the user navigated away mid-load.
+    # Lazy chunk no longer needed because the user navigated away mid-load.
     "Loading chunk",
-    # Filtered as a benign page-error too; included here for the
-    # parallel diagnostic dump path.
+    # Also a benign page-error; here for the diagnostic dump path.
     "Failed to fetch",
 )
 
@@ -352,14 +291,9 @@ def dump_diagnostics(
     info: Callable[[str], None] | None = None,
     extra: dict | None = None,
 ) -> None:
-    """Write a screenshot + URL/title + body excerpt + storage dump.
-
-    Diagnostic only. Never raises. The screenshot path lives in
-    `art_dir/{name}.png`; the JSON sidecar lives in `art_dir/{name}.json`.
-    The screenshot is wrapped in try/except because Page.screenshot
-    waits for webfonts to load and can crowd CI font load on macos-14
-    even at 90s. The JSON sidecar is best-effort too.
-    """
+    """Write a screenshot (`art_dir/{name}.png`) + a JSON sidecar
+    (`art_dir/{name}.json`) with URL/title/body/storage. Diagnostic only, never
+    raises; both are best-effort (screenshot can crowd CI font load on macos-14)."""
     art = Path(art_dir)
     try:
         art.mkdir(parents = True, exist_ok = True)
@@ -408,26 +342,13 @@ def dump_diagnostics(
             info(f"diagnostics: json sidecar {name} failed: {exc}")
 
 
-# ─────────────────────────────────────────────────────────────────────
 # Bounded in-page fetch.
-# ─────────────────────────────────────────────────────────────────────
 #
-# Playwright's `page.evaluate(...)` has no `timeout=` argument. If the
-# JS body awaits a fetch that never resolves (the renderer's network
-# thread wedges, the server accepts the connection but never replies,
-# the macos-14 free runner under --single-process Chromium loses its
-# IPC pipe), the entire Python script hangs until the runner-level
-# timeout fires. Run 25696797934 / job 75446949358 on PR #5387 showed
-# this exact failure: studio.log went idle after the chat surface
-# mounted, no further requests reached the server, and Playwright
-# burned 27+ minutes on a single page.evaluate(fetch /api/inference/
-# load) before the 30-min runner cancel.
-#
-# `evaluate_fetch` wraps the fetch in an AbortController.signal so the
-# JS side resolves either with a real response or with a synthetic
-# `{status: 0, error: "AbortError..."}` after `timeout_ms` ms. Either
-# way page.evaluate returns and the script proceeds (or fails) with
-# a debuggable signal instead of a silent wedge.
+# `page.evaluate(...)` has no `timeout=`, so a fetch that never resolves hangs
+# the whole script until the runner timeout (run 25696797934 / PR #5387 burned
+# 27+ min on one page.evaluate). `evaluate_fetch` wraps the fetch in an
+# AbortController.signal so the JS side always resolves -- with a real response
+# or a synthetic `{status: 0, error: "AbortError..."}` after `timeout_ms`.
 def evaluate_fetch(
     page: Any,
     url: str,
@@ -436,19 +357,14 @@ def evaluate_fetch(
     headers: dict[str, str] | None = None,
     body: Any = None,
     timeout_ms: int = 20_000,
+    transport_retries: int = 2,
+    transport_backoff_ms: int = 250,
 ) -> dict[str, Any]:
-    """Run `fetch(url, opts)` inside the page with an AbortSignal deadline.
-
-    Returns `{"status": int, "body": parsed_or_text, "error": str|None}`.
-    On AbortSignal timeout returns `{"status": 0, "body": None, "error":
-    "AbortError: ..."}`. Callers should treat `status == 0` (or any
-    non-None `error`) as a transport failure rather than an HTTP
-    response.
-
-    `body` may be a `str` (sent verbatim) or a `dict`/`list` (JSON-
-    encoded here). Pass headers explicitly when you need
-    `Content-Type: application/json` or an `Authorization` bearer.
-    """
+    """Run `fetch(url, opts)` in the page with an AbortSignal deadline. Returns
+    `{"status", "body", "error"}`; on timeout `status==0` with an AbortError
+    string. Treat `status == 0` or a non-None `error` as transport failure, not
+    an HTTP response. `body` may be a str (verbatim) or dict/list (JSON-encoded
+    here); pass headers explicitly for Content-Type / Authorization."""
     body_arg: str | None
     if body is None:
         body_arg = None
@@ -482,48 +398,64 @@ def evaluate_fetch(
             }
         }
     """
-    return page.evaluate(
-        js,
-        {
-            "url": url,
-            "method": method,
-            "headers": headers or {},
-            "body": body_arg,
-            "timeoutMs": int(timeout_ms),
-        },
-    )
+    payload = {
+        "url": url,
+        "method": method,
+        "headers": headers or {},
+        "body": body_arg,
+        "timeoutMs": int(timeout_ms),
+    }
+    # Bounded retry on transport failures only:
+    #   status != 0 -> real HTTP response (incl. 4xx/5xx); propagate.
+    #   AbortError  -> caller's deadline; propagate.
+    #   else (==0)  -> stale-keepalive / "Failed to fetch" after auth rotation;
+    #                  retry after backoff so the pool evicts the dead socket.
+    last: dict[str, Any] | None = None
+    attempts = max(1, int(transport_retries) + 1)
+    for attempt in range(attempts):
+        result = page.evaluate(js, payload)
+        last = result
+        try:
+            status = int(result.get("status") or 0)
+        except (TypeError, ValueError):
+            status = 0
+        if status != 0:
+            return result
+        err = str(result.get("error") or "")
+        if "AbortError" in err:
+            return result
+        if attempt < attempts - 1:
+            wait_ms = transport_backoff_ms * (2**attempt)
+            try:
+                sys.stderr.write(
+                    f"[evaluate_fetch] {method} {url}: transport failure "
+                    f"({attempt + 1}/{attempts}, err={err!r}); "
+                    f"retrying in {wait_ms}ms\n"
+                )
+                sys.stderr.flush()
+            except Exception:
+                pass
+            time.sleep(wait_ms / 1000.0)
+    return last or {"status": 0, "body": None, "error": "no attempt made"}
 
 
-# ─────────────────────────────────────────────────────────────────────
 # Wall-clock watchdog.
-# ─────────────────────────────────────────────────────────────────────
 #
-# Even with every action and fetch bounded, a sufficiently strange
-# wedge inside the browser (a CPU-pinned JS infinite loop, a renderer
-# crash that doesn't propagate to Playwright, an asyncio deadlock in
-# the sync wrapper) can still hang the script. The watchdog is a
-# daemon Timer that calls `os._exit(2)` after `deadline_s` seconds,
-# printing the wedge location to stderr so the CI log shows where the
-# script was at force-kill time. The exit code matches "test failure
-# by deadline" so the workflow's `set -e` propagates correctly.
-#
-# Pick `deadline_s` generously enough to cover the slowest healthy
-# run -- macos-14 free runners with cold caches measure ~7-9 min for
-# the comprehensive chat UI test. 12 minutes (720 s) leaves headroom
-# without amplifying every real wedge to the 30-min runner-level cap.
+# Even with every action/fetch bounded, a strange browser wedge (CPU-pinned JS
+# loop, renderer crash that doesn't propagate, asyncio deadlock) can hang the
+# script. A daemon Timer calls `os._exit(2)` after `deadline_s`, printing the
+# wedge location to stderr; exit code 2 lets the workflow's `set -e` propagate.
+# Pick `deadline_s` above the slowest healthy run (macos-14 cold cache ~7-9 min;
+# 720s leaves headroom without nearing the 30-min runner cap).
 def install_wall_clock_watchdog(
     deadline_s: float,
     *,
     label: str = "playwright",
     info: Callable[[str], None] | None = None,
 ) -> threading.Timer:
-    """Start a daemon Timer that hard-exits the process at `deadline_s`.
-
-    Returns the Timer so the caller can `.cancel()` it on clean exit.
-    The Timer is daemonised; if the script exits normally before the
-    deadline the Timer dies with the process even without an explicit
-    cancel.
-    """
+    """Start a daemon Timer that hard-exits the process at `deadline_s`. Returns
+    it so the caller can `.cancel()` on clean exit; being daemonised, it also
+    dies with the process if the script exits first."""
 
     def _kaboom() -> None:
         msg = (

tests/studio/install/conftest.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Pytest configuration for studio/install tests.
+
+install_python_stack.py does ``from backend.utils.wheel_utils import ...``
+which requires the ``studio/`` directory to be on sys.path.  When tests are
+run from the repo root (the normal case), the studio package is not
+automatically importable, so we add it here.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# <repo-root>/studio  →  makes `backend` importable as a package
+_STUDIO_DIR = Path(__file__).resolve().parents[3] / "studio"
+if str(_STUDIO_DIR) not in sys.path:
+    sys.path.insert(0, str(_STUDIO_DIR))

tests/studio/install/smoke_test_llama_prebuilt.py

@@ -15,9 +15,7 @@ INSTALLER_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
 
 
 def load_installer_module():
-    spec = importlib.util.spec_from_file_location(
-        "studio_install_llama_prebuilt", INSTALLER_PATH
-    )
+    spec = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", INSTALLER_PATH)
     if spec is None or spec.loader is None:
         raise RuntimeError(f"unable to load installer module from {INSTALLER_PATH}")
     module = importlib.util.module_from_spec(spec)
@@ -112,17 +110,13 @@ def main() -> int:
             published_release_tag = args.published_release_tag,
         )
         print(f"[smoke] PASS install_dir={install_dir}")
-        print(
-            "[smoke] note=This was a real prebuilt install into an isolated temp directory."
-        )
+        print("[smoke] note=This was a real prebuilt install into an isolated temp directory.")
         return installer.EXIT_SUCCESS
     except SystemExit as exc:
         code = int(exc.code) if isinstance(exc.code, int) else installer.EXIT_ERROR
         if code == installer.EXIT_FALLBACK:
             print(f"[smoke] FALLBACK install_dir={install_dir}")
-            print(
-                "[smoke] note=Prebuilt path failed and would fall back to source build in setup."
-            )
+            print("[smoke] note=Prebuilt path failed and would fall back to source build in setup.")
             print(installer.collect_system_report(host, choice, install_dir))
         else:
             print(f"[smoke] ERROR exit_code={code} install_dir={install_dir}")

tests/studio/install/smoke_test_parallel_studio_home.py

@@ -1,50 +1,29 @@
 #!/usr/bin/env python3
 """Smoke test: N parallel install.sh runs with distinct UNSLOTH_STUDIO_HOME
-values must produce N fully isolated installs whose backends can run
-side by side without clashing.
+values must produce N fully isolated installs whose backends run side by side.
 
-Covers the env-override path added in #5190:
+Covers the env-override path from #5190:
 
     install-time
-        * N concurrent ``install.sh --local --no-torch`` runs against
-          this checkout, each pinned to its own UNSLOTH_STUDIO_HOME and
-          a redirected HOME, all exit 0.
-        * Each STUDIO_HOME contains its own bin/, share/, llama.cpp/
-          and unsloth_studio/ venv, with no cross-install absolute
-          paths.
-        * share/studio_install_id is unique across the N installs.
-        * share/studio.conf exports UNSLOTH_EXE, UNSLOTH_STUDIO_HOME
-          and UNSLOTH_LLAMA_CPP_PATH, all pointing inside this install.
-        * share/launch-studio.sh has @@DATA_DIR@@ substituted to its
-          own share/ at install time.
-        * bin/unsloth is a symlink that resolves into its own venv.
-        * The redirected HOME is left clean: no shell-rc append, no
-          .desktop file, no Studio.app stub, no shared marker.
+        * N concurrent ``install.sh --local --no-torch`` runs, each pinned to
+          its own UNSLOTH_STUDIO_HOME + redirected HOME, all exit 0.
+        * Each STUDIO_HOME has its own bin/share/llama.cpp/unsloth_studio venv,
+          a unique share/studio_install_id, and a studio.conf / launch-studio.sh
+          pointing only inside this install. bin/unsloth resolves into its venv.
+        * The redirected HOME is left clean (no rc append, .desktop, app stub).
 
     runtime
-        * N concurrent ``bin/unsloth studio`` launches each bind their
-          own dynamically allocated free port and stay healthy.
-        * /api/health is 200, status is healthy, chat_only is true
-          under --no-torch.
-        * The studio_root_id reported by /api/health on each backend
-          equals that install's share/studio_install_id, so the
-          runtime resolver agrees with the install-time write.
-        * studio_root_id values are pairwise distinct.
-        * GET / and GET /api/chat are 200 on every backend.
-        * The Python interpreter behind each PID is the install's own
-          venv python (the bin/unsloth shim does not cross-resolve).
+        * N ``bin/unsloth studio`` launches each bind a free port and report
+          /api/health 200, healthy, chat_only true under --no-torch.
+        * studio_root_id equals the install's studio_install_id and is pairwise
+          distinct; GET / and /api/chat are 200; each PID's python is its venv.
 
-This is an integration smoke runner, not a pytest unit test. It does
-real installs (~1 minute end to end on a warm uv cache) and is meant
-to be invoked explicitly:
+Integration smoke runner (not pytest); ~1 minute on a warm uv cache. Invoke:
 
-    python tests/studio/install/smoke_test_parallel_studio_home.py
-    python tests/studio/install/smoke_test_parallel_studio_home.py --n 6 --keep
+    python tests/studio/install/smoke_test_parallel_studio_home.py [--n 6 --keep]
 
-Exits 0 on PASS, 1 on FAIL, 2 on infrastructure error. Artifacts land
-under a temporary directory and are removed on PASS unless --keep is
-set; on FAIL or ERROR they are kept regardless so logs can be
-inspected.
+Exits 0 PASS / 1 FAIL / 2 error. Artifacts removed on PASS unless --keep;
+kept on FAIL/ERROR for inspection.
 """
 
 from __future__ import annotations
@@ -86,12 +65,7 @@ def _free_port() -> int:
 
 
 def _run_one_install(
-    label: str,
-    repo: Path,
-    studio_home: Path,
-    fake_home: Path,
-    uv_cache: Path,
-    log_path: Path,
+    label: str, repo: Path, studio_home: Path, fake_home: Path, uv_cache: Path, log_path: Path
 ) -> tuple[str, int]:
     studio_home.mkdir(parents = True, exist_ok = True)
     fake_home.mkdir(parents = True, exist_ok = True)
@@ -159,12 +133,14 @@ def _wait_for_health(port: int, timeout: float) -> dict:
         except (urllib.error.URLError, ConnectionError, OSError) as e:
             last_err = e
         time.sleep(HEALTH_POLL_INTERVAL_S)
-    raise TestFailure(
-        f"port {port}: /api/health never returned 200 (last_err={last_err})"
-    )
+    raise TestFailure(f"port {port}: /api/health never returned 200 (last_err={last_err})")
 
 
-def _http_status(port: int, path: str, timeout: float = 5.0) -> int:
+def _http_status(
+    port: int,
+    path: str,
+    timeout: float = 5.0,
+) -> int:
     url = f"http://127.0.0.1:{port}{path}"
     try:
         with urllib.request.urlopen(url, timeout = timeout) as r:
@@ -211,9 +187,7 @@ def _check_install_layout(label: str, studio_home: Path) -> dict:
         raise TestFailure(f"[{label}] launch-studio.sh kept @@DATA_DIR@@ placeholder")
     expected_data_dir_line = f"DATA_DIR='{studio_home}/share'"
     if expected_data_dir_line not in launcher:
-        raise TestFailure(
-            f"[{label}] launch-studio.sh missing {expected_data_dir_line!r}"
-        )
+        raise TestFailure(f"[{label}] launch-studio.sh missing {expected_data_dir_line!r}")
 
     return {"label": label, "studio_home": str(studio_home), "install_id": install_id}
 
@@ -230,9 +204,7 @@ def _check_fake_home_clean(fake_home: Path) -> None:
     ]
     leaked = [str(p) for p in forbidden if (fake_home / p).exists()]
     if leaked:
-        raise TestFailure(
-            f"redirected HOME picked up persistent install pollution: {leaked}"
-        )
+        raise TestFailure(f"redirected HOME picked up persistent install pollution: {leaked}")
 
 
 def _backend_pid_python(pid: int) -> Path | None:
@@ -256,9 +228,7 @@ def run(n_installs: int, keep: bool) -> int:
 
     repo = PACKAGE_ROOT
     if not (repo / "install.sh").is_file():
-        raise TestFailure(
-            f"install.sh not found at {repo}; " "run from a clone of unslothai/unsloth"
-        )
+        raise TestFailure(f"install.sh not found at {repo}; run from a clone of unslothai/unsloth")
 
     test_root = Path(tempfile.mkdtemp(prefix = "unsloth_studio_clash_"))
     _log(f"test root: {test_root}")
@@ -346,8 +316,7 @@ def run(n_installs: int, keep: bool) -> int:
                 raise TestFailure(f"[{label}] chat_only is not true under --no-torch")
             if health["studio_root_id"] in seen_root_ids:
                 raise TestFailure(
-                    f"[{label}] studio_root_id collision at runtime: "
-                    f"{health['studio_root_id']}"
+                    f"[{label}] studio_root_id collision at runtime: " f"{health['studio_root_id']}"
                 )
             seen_root_ids.add(health["studio_root_id"])
 
@@ -358,9 +327,7 @@ def run(n_installs: int, keep: bool) -> int:
 
             exe = _backend_pid_python(proc.pid)
             if exe is not None:
-                expected_python = (
-                    studio_home / "unsloth_studio" / "bin" / "python"
-                ).resolve()
+                expected_python = (studio_home / "unsloth_studio" / "bin" / "python").resolve()
                 if exe != expected_python:
                     raise TestFailure(
                         f"[{label}] PID {proc.pid} exe={exe}, expected {expected_python}"
@@ -370,10 +337,7 @@ def run(n_installs: int, keep: bool) -> int:
         if len(versions) != 1:
             raise TestFailure(f"version mismatch across installs: {versions}")
 
-        _log(
-            f"PASS: all install + runtime invariants hold "
-            f"(version={next(iter(versions))})"
-        )
+        _log(f"PASS: all install + runtime invariants hold " f"(version={next(iter(versions))})")
         return 0
 
     except TestFailure as e:

tests/studio/install/test_cuda_repair.py

@@ -0,0 +1,248 @@
+"""Tests for CUDA torch repair on poisoned NVIDIA venvs.
+
+Verifies _ensure_cuda_torch (studio/install_python_stack.py) reinstalls CUDA
+torch when a venv on an NVIDIA host carries a ROCm torch build (the pre-fix KFD
+gpu_id false positive), without touching healthy CUDA, deliberate CPU wheels,
+ROCm hosts, macOS, or Windows. All tests use mocks -- no GPU required.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+# ── Load module under test (mirrors test_rocm_support.py) ────────────────────
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+
+_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
+_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH)
+assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
+stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
+sys.modules[_STACK_SPEC.name] = stack_mod
+_STACK_SPEC.loader.exec_module(stack_mod)
+
+_ensure_cuda_torch = stack_mod._ensure_cuda_torch
+_detect_cuda_torch_index_url = stack_mod._detect_cuda_torch_index_url
+
+
+# ── Helpers ──────────────────────────────────────────────────────────────────
+
+
+def _make_run(
+    torch_state = "hip",
+    cuda_version = "12.8",
+    torch_rc = 0,
+    smi_rc = 0,
+):
+    """Build a subprocess.run side_effect.
+
+    The torch-classify probe runs sys.executable and reads bytes stdout; the
+    nvidia-smi version probe runs the smi path with text=True. Distinguish by
+    the executable.
+    """
+
+    def _run(cmd, *args, **kwargs):
+        result = MagicMock()
+        exe = str(cmd[0]) if cmd else ""
+        if exe == sys.executable:
+            result.returncode = torch_rc
+            result.stdout = (torch_state + "\n").encode()
+            return result
+        # nvidia-smi version probe (text = True)
+        result.returncode = smi_rc
+        out = f"CUDA Version: {cuda_version}\n" if cuda_version else "No devices found\n"
+        result.stdout = out if kwargs.get("text") else out.encode()
+        return result
+
+    return _run
+
+
+def _run_cuda_repair(
+    *,
+    backend = "",
+    nvidia = True,
+    torch_state = "hip",
+    cuda_version = "12.8",
+    torch_rc = 0,
+    smi_rc = 0,
+    is_macos = False,
+    is_windows = False,
+    no_torch = False,
+    rocm_marker = False,
+    smi_path = "/usr/bin/nvidia-smi",
+    cvd = None,
+):
+    """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock.
+
+    cvd controls CUDA_VISIBLE_DEVICES: None removes it from the environment
+    (the host machine may export one), any string sets it explicitly.
+    """
+    env = {}
+    if rocm_marker:
+        env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
+    if cvd is not None:
+        env["CUDA_VISIBLE_DEVICES"] = cvd
+
+    def _which(name, *a, **k):
+        if name == "nvidia-smi":
+            return smi_path
+        return None
+
+    with (
+        patch.object(stack_mod, "_TORCH_BACKEND", backend),
+        patch.object(stack_mod, "IS_MACOS", is_macos),
+        patch.object(stack_mod, "IS_WINDOWS", is_windows),
+        patch.object(stack_mod, "NO_TORCH", no_torch),
+        patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia),
+        patch.object(stack_mod.shutil, "which", side_effect = _which),
+        patch.object(stack_mod.os.path, "isfile", return_value = bool(smi_path)),
+        patch.object(stack_mod, "pip_install") as mock_pip,
+        patch.object(
+            stack_mod.subprocess,
+            "run",
+            side_effect = _make_run(torch_state, cuda_version, torch_rc, smi_rc),
+        ),
+        patch.dict(stack_mod.os.environ, env, clear = False),
+    ):
+        if not rocm_marker:
+            stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None)
+        if cvd is None:
+            stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        _ensure_cuda_torch()
+    return mock_pip
+
+
+def _index_url(mock_pip) -> str:
+    """Return the --index-url value from the recorded pip_install call."""
+    args = [str(a) for a in mock_pip.call_args.args]
+    return args[args.index("--index-url") + 1]
+
+
+# ── Repair fires only on the poisoning signature ─────────────────────────────
+
+
+class TestCudaRepairFires:
+    def test_hip_build_on_nvidia_triggers_repair(self):
+        mock_pip = _run_cuda_repair(torch_state = "hip", cuda_version = "12.8")
+        assert mock_pip.call_count == 1
+        call_args = [str(a) for a in mock_pip.call_args.args]
+        assert "--force-reinstall" in call_args
+        assert "--no-cache-dir" in call_args
+        assert "cu128" in _index_url(mock_pip)
+        assert mock_pip.call_args.kwargs["constrain"] is False
+
+    def test_rocm_in_version_string_triggers_repair(self):
+        # AMD SDK / Radeon wheels may not set torch.version.hip but encode
+        # rocm in __version__; the probe prints "hip" for both.
+        mock_pip = _run_cuda_repair(torch_state = "hip")
+        assert mock_pip.call_count == 1
+
+
+# ── No-op cases ──────────────────────────────────────────────────────────────
+
+
+class TestCudaRepairSkips:
+    def test_healthy_cuda_torch_no_repair(self):
+        mock_pip = _run_cuda_repair(torch_state = "cuda")
+        mock_pip.assert_not_called()
+
+    def test_deliberate_cpu_wheel_no_repair(self):
+        mock_pip = _run_cuda_repair(torch_state = "cpu")
+        mock_pip.assert_not_called()
+
+    def test_backend_rocm_skips(self):
+        mock_pip = _run_cuda_repair(backend = "rocm", torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_backend_cpu_skips(self):
+        mock_pip = _run_cuda_repair(backend = "cpu", torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_unknown_backend_skips(self):
+        mock_pip = _run_cuda_repair(backend = "auto", torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_no_nvidia_gpu_skips(self):
+        mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_torch_missing_skips(self):
+        # Non-zero probe exit = torch missing / un-importable.
+        mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1)
+        mock_pip.assert_not_called()
+
+    def test_macos_skips(self):
+        mock_pip = _run_cuda_repair(is_macos = True, torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_windows_skips(self):
+        mock_pip = _run_cuda_repair(is_windows = True, torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_no_torch_mode_skips(self):
+        mock_pip = _run_cuda_repair(no_torch = True, torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_rocm_install_marker_skips(self):
+        mock_pip = _run_cuda_repair(rocm_marker = True, torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_cvd_minus_one_skips(self):
+        # CUDA_VISIBLE_DEVICES=-1 deliberately hides the NVIDIA GPU (mixed
+        # AMD+NVIDIA host running ROCm torch on the AMD card).
+        mock_pip = _run_cuda_repair(cvd = "-1", torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_cvd_empty_skips(self):
+        mock_pip = _run_cuda_repair(cvd = "", torch_state = "hip")
+        mock_pip.assert_not_called()
+
+    def test_cvd_explicit_device_still_repairs(self):
+        mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip")
+        assert mock_pip.call_count == 1
+
+
+# ── CUDA index ladder ────────────────────────────────────────────────────────
+
+
+class TestCudaIndexResolution:
+    def test_cuda_128_selects_cu128(self):
+        assert "cu128" in _index_url(_run_cuda_repair(cuda_version = "12.8"))
+
+    def test_cuda_130_selects_cu130(self):
+        assert "cu130" in _index_url(_run_cuda_repair(cuda_version = "13.0"))
+
+    def test_cuda_126_selects_cu126(self):
+        assert "cu126" in _index_url(_run_cuda_repair(cuda_version = "12.6"))
+
+    def test_cuda_124_selects_cu124(self):
+        assert "cu124" in _index_url(_run_cuda_repair(cuda_version = "12.4"))
+
+    def test_cuda_118_selects_cu118(self):
+        assert "cu118" in _index_url(_run_cuda_repair(cuda_version = "11.8"))
+
+    def test_unreadable_version_defaults_cu126(self):
+        # nvidia-smi runs but prints no CUDA version line (or fails).
+        mock_pip = _run_cuda_repair(cuda_version = "", smi_rc = 1)
+        assert "cu126" in _index_url(mock_pip)
+
+    def test_proc_fallback_no_smi_defaults_cu126(self):
+        # NVIDIA usable via /proc fallback, nvidia-smi absent entirely.
+        mock_pip = _run_cuda_repair(smi_path = None)
+        assert "cu126" in _index_url(mock_pip)
+
+    def test_detect_index_url_uses_pytorch_base(self):
+        with (
+            patch.object(stack_mod.shutil, "which", return_value = None),
+            patch.object(stack_mod.os.path, "isfile", return_value = False),
+        ):
+            url = _detect_cuda_torch_index_url()
+        assert url == f"{stack_mod._PYTORCH_WHL_BASE}/cu126"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-q"]))

tests/studio/install/test_gpu_detection_followups.py

@@ -0,0 +1,480 @@
+"""Tests for the GPU-detection follow-ups to PR 6174.
+
+PR 6174 made NVIDIA take precedence and added a /proc/driver/nvidia/gpus
+fallback in install.sh and studio/install_python_stack.py. These tests cover the
+same hardening ported to the llama.cpp prebuilt installer
+(studio/install_llama_prebuilt.py) and the Studio shell setup (studio/setup.sh):
+
+  * detect_host() recognises NVIDIA via /proc/driver/nvidia/gpus when nvidia-smi
+    is unavailable, and skips ROCm probing when NVIDIA is usable.
+  * setup.sh routes through a timeout-bounded NVIDIA probe with a /proc fallback
+    and only selects a CUDA/ROCm source build when the matching GPU is detected.
+
+All tests use mocks or source-level assertions -- no GPU, network, or real
+nvidia-smi/rocminfo invocation.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+
+# Load studio/install_llama_prebuilt.py the same way the sibling suite does.
+_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+_SPEC = importlib.util.spec_from_file_location(
+    "studio_install_llama_prebuilt_followups", _MODULE_PATH
+)
+assert _SPEC is not None and _SPEC.loader is not None
+prebuilt_mod = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = prebuilt_mod
+_SPEC.loader.exec_module(prebuilt_mod)
+
+detect_host = prebuilt_mod.detect_host
+_apply_host_overrides = prebuilt_mod._apply_host_overrides
+
+SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
+
+
+def _make_run_capture(rocminfo_stdout: str = ""):
+    """Return a fake run_capture: rocminfo reports rocminfo_stdout, everything
+    else (nvidia-smi, amd-smi) returns empty so only the patched probes matter."""
+
+    def _run_capture(cmd, *args, **kwargs):
+        exe = str(cmd[0]) if cmd else ""
+        result = MagicMock()
+        if exe.endswith("rocminfo"):
+            result.returncode = 0
+            result.stdout = rocminfo_stdout
+        else:
+            result.returncode = 1
+            result.stdout = ""
+        result.stderr = ""
+        return result
+
+    return _run_capture
+
+
+def _run_detect_host(
+    *,
+    machine: str = "x86_64",
+    system: str = "Linux",
+    which_map: dict | None = None,
+    proc_dir_entries: list | None = None,
+    rocminfo_stdout: str = "",
+    env: dict | None = None,
+):
+    """Drive detect_host() against a fully synthetic host."""
+    which_map = which_map or {}
+    proc_dir_entries = proc_dir_entries if proc_dir_entries is not None else []
+
+    real_isdir = prebuilt_mod.os.path.isdir
+    real_listdir = prebuilt_mod.os.listdir
+    proc_path = "/proc/driver/nvidia/gpus"
+
+    def fake_isdir(p):
+        if str(p) == proc_path:
+            return bool(proc_dir_entries)
+        return real_isdir(p)
+
+    def fake_listdir(p):
+        if str(p) == proc_path:
+            if not proc_dir_entries:
+                raise OSError("no such dir")
+            return list(proc_dir_entries)
+        return real_listdir(p)
+
+    patches = [
+        patch.object(prebuilt_mod.platform, "system", return_value = system),
+        patch.object(prebuilt_mod.platform, "machine", return_value = machine),
+        patch.object(prebuilt_mod.platform, "mac_ver", return_value = ("", ("", "", ""), "")),
+        patch.object(prebuilt_mod.shutil, "which", side_effect = lambda n: which_map.get(n)),
+        patch.object(prebuilt_mod, "run_capture", side_effect = _make_run_capture(rocminfo_stdout)),
+        patch.object(prebuilt_mod.os.path, "isdir", side_effect = fake_isdir),
+        patch.object(prebuilt_mod.os, "listdir", side_effect = fake_listdir),
+        patch.object(prebuilt_mod.os, "access", return_value = False),
+        patch.dict(prebuilt_mod.os.environ, env or {}, clear = False),
+    ]
+    for p in patches:
+        p.start()
+    try:
+        # Ensure CUDA_VISIBLE_DEVICES does not leak in from the test host unless
+        # the scenario sets it explicitly.
+        if env is None or "CUDA_VISIBLE_DEVICES" not in env:
+            prebuilt_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        return detect_host()
+    finally:
+        for p in patches:
+            p.stop()
+
+
+# ── install_llama_prebuilt.detect_host(): /proc NVIDIA fallback ──────────────
+
+
+class TestDetectHostProcFallback:
+    def test_proc_fallback_marks_physical_nvidia_when_smi_absent(self):
+        """No nvidia-smi, but /proc/driver/nvidia/gpus is populated -> NVIDIA."""
+        host = _run_detect_host(
+            which_map = {},  # nvidia-smi resolves to None
+            proc_dir_entries = ["0000:01:00.0"],
+        )
+        assert host.has_physical_nvidia is True
+
+    def test_proc_fallback_has_usable_nvidia_when_devices_visible(self):
+        """Default CUDA_VISIBLE_DEVICES (unset) -> visible tokens non-empty -> usable."""
+        host = _run_detect_host(
+            which_map = {},
+            proc_dir_entries = ["0000:01:00.0"],
+        )
+        assert host.has_usable_nvidia is True
+
+    def test_proc_fallback_not_usable_when_devices_hidden(self):
+        """CUDA_VISIBLE_DEVICES='' hides all GPUs -> physical yes, usable no."""
+        host = _run_detect_host(
+            which_map = {},
+            proc_dir_entries = ["0000:01:00.0"],
+            env = {"CUDA_VISIBLE_DEVICES": ""},
+        )
+        assert host.has_physical_nvidia is True
+        assert host.has_usable_nvidia is False
+
+    def test_empty_proc_dir_does_not_mark_nvidia(self):
+        """A driver dir that exists but is empty must not assert a GPU."""
+        host = _run_detect_host(which_map = {}, proc_dir_entries = [])
+        assert host.has_physical_nvidia is False
+
+    def test_proc_fallback_is_linux_only(self):
+        """The /proc fallback must not run on Windows (path is Linux-only)."""
+        host = _run_detect_host(
+            system = "Windows",
+            machine = "amd64",
+            which_map = {},
+            proc_dir_entries = ["0000:01:00.0"],
+        )
+        assert host.has_physical_nvidia is False
+
+
+# ── install_llama_prebuilt.detect_host(): NVIDIA precedence over ROCm ────────
+
+
+class TestDetectHostNvidiaPrecedence:
+    def test_rocm_probe_skipped_when_proc_nvidia_present(self):
+        """rocminfo reports gfx1100, but a proc-detected NVIDIA GPU wins."""
+        host = _run_detect_host(
+            which_map = {"rocminfo": "/usr/bin/rocminfo"},
+            proc_dir_entries = ["0000:01:00.0"],
+            rocminfo_stdout = "  Name:                    gfx1100\n",
+        )
+        assert host.has_usable_nvidia is True
+        assert host.has_rocm is False
+
+    def test_rocm_detected_when_no_nvidia(self):
+        """With no NVIDIA signal at all, rocminfo gfx1100 -> has_rocm True."""
+        host = _run_detect_host(
+            which_map = {"rocminfo": "/usr/bin/rocminfo"},
+            proc_dir_entries = [],
+            rocminfo_stdout = "  Name:                    gfx1100\n",
+        )
+        assert host.has_usable_nvidia is False
+        assert host.has_rocm is True
+
+
+# ── _apply_host_overrides: forwarded --rocm-gfx / --has-rocm still win ───────
+
+
+class TestOverridesStillWin:
+    def test_forwarded_gfx_forces_rocm_on_non_nvidia_host(self):
+        host = _run_detect_host(which_map = {}, proc_dir_entries = [])
+        assert host.has_rocm is False
+        overridden = _apply_host_overrides(host, override_rocm_gfx = "gfx1100")
+        assert overridden.has_rocm is True
+        assert overridden.rocm_gfx_target == "gfx1100"
+
+    def test_override_has_rocm_forces_rocm(self):
+        host = _run_detect_host(which_map = {}, proc_dir_entries = [])
+        overridden = _apply_host_overrides(host, override_has_rocm = True)
+        assert overridden.has_rocm is True
+
+    def test_force_cpu_drops_nvidia_attributes(self):
+        host = _run_detect_host(which_map = {}, proc_dir_entries = ["0000:01:00.0"])
+        assert host.has_usable_nvidia is True
+        overridden = _apply_host_overrides(host, force_cpu = True)
+        assert overridden.has_usable_nvidia is False
+        assert overridden.has_physical_nvidia is False
+        assert overridden.has_rocm is False
+
+
+# ── setup.sh source-level guarantees ────────────────────────────────────────
+
+
+class TestSetupShHardening:
+    @pytest.fixture(scope = "class")
+    def setup_src(self) -> str:
+        return SETUP_SH.read_text(encoding = "utf-8")
+
+    def test_has_usable_nvidia_helper_exists(self, setup_src):
+        assert "_setup_has_usable_nvidia_gpu()" in setup_src
+
+    def test_helper_uses_proc_fallback(self, setup_src):
+        start = setup_src.find("_setup_has_usable_nvidia_gpu()")
+        end = setup_src.find("\n}", start)
+        body = setup_src[start:end]
+        assert (
+            "/proc/driver/nvidia/gpus" in body
+        ), "_setup_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus"
+
+    def test_gpu_summary_uses_helper(self, setup_src):
+        assert "if _setup_has_usable_nvidia_gpu; then" in setup_src
+
+    def test_timeout_wrapper_exists(self, setup_src):
+        start = setup_src.find("_setup_run_smi()")
+        assert start >= 0, "_setup_run_smi timeout wrapper must exist"
+        end = setup_src.find("\n}", start)
+        body = setup_src[start:end]
+        assert "timeout 10" in body
+        assert "command -v timeout" in body
+
+    def test_cuda_source_build_gated_on_usable_nvidia(self, setup_src):
+        """The nvcc source-build search must be gated on _setup_nvidia_usable.
+
+        The hidden-GPU policy (CUDA_VISIBLE_DEVICES=""/-1) lives inside
+        _setup_has_usable_nvidia_gpu, so the gate itself only needs the flag.
+        """
+        anchor = setup_src.find('NVCC_PATH=""\n')
+        assert anchor >= 0
+        window = setup_src[anchor : anchor + 700]
+        assert (
+            'if [ "$_setup_nvidia_usable" = true ]' in window
+        ), "CUDA toolkit search must require a usable NVIDIA GPU, not just nvcc"
+
+    def test_nvidia_helper_honours_hidden_cvd(self, setup_src):
+        """_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so
+        CUDA_VISIBLE_DEVICES=""/-1 suppresses NVIDIA before the AMD probes are
+        gated (mixed hosts steered to the AMD card keep the ROCm route)."""
+        assert "_setup_cvd_hides_nvidia()" in setup_src
+        start = setup_src.find("_setup_has_usable_nvidia_gpu() {")
+        end = setup_src.find("\n}", start)
+        body = setup_src[start:end]
+        assert "_setup_cvd_hides_nvidia" in body
+
+    def test_rocm_source_build_gated_on_amd_detected(self, setup_src):
+        """The hipcc source-build search must be gated on _setup_amd_detected."""
+        anchor = setup_src.find('ROCM_HIPCC=""')
+        assert anchor >= 0
+        window = setup_src[anchor : anchor + 400]
+        assert (
+            '[ "$_setup_amd_detected" = true ]' in window
+        ), "ROCm toolkit search must require a detected AMD GPU, not just hipcc"
+
+    def test_compute_cap_probe_timeout_wrapped(self, setup_src):
+        assert "_setup_run_smi nvidia-smi --query-gpu=compute_cap" in setup_src
+
+    def test_driver_version_probe_timeout_wrapped(self, setup_src):
+        start = setup_src.find("_cuda_driver_max_version()")
+        end = setup_src.find("\n}", start)
+        body = setup_src[start:end]
+        assert "_setup_run_smi nvidia-smi" in body
+
+
+# TEST: install.sh -- UNSLOTH_TORCH_BACKEND classified on the final path segment
+
+
+class TestBackendExportLeafClassification:
+    """A custom UNSLOTH_PYTORCH_MIRROR whose base path contains "rocm" or
+    "gfx" must not mislabel a cu*/cpu index as ROCm; classification uses the
+    final path segment of TORCH_INDEX_URL only."""
+
+    @pytest.fixture(scope = "class")
+    def install_src(self) -> str:
+        return (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+
+    def test_export_block_uses_leaf(self, install_src):
+        anchor = install_src.find("_torch_index_leaf=")
+        assert anchor >= 0, "backend export must classify on the final path segment"
+        window = install_src[anchor : anchor + 500]
+        assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window
+        assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window
+        assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window
+
+    def test_leaf_classification_behaviour(self, tmp_path):
+        import subprocess as sp
+
+        script = tmp_path / "leaf.sh"
+        src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+        anchor = src.find("_torch_index_leaf=")
+        block = src[anchor : src.find("esac", anchor) + 4]
+        # Drive the extracted block with adversarial mirror URLs.
+        script.write_text(
+            "#!/bin/sh\n"
+            'TORCH_INDEX_URL="$1"\n' + block + "\n"
+            'printf "%s" "$UNSLOTH_TORCH_BACKEND"\n'
+        )
+        cases = {
+            "https://download.pytorch.org/whl/cu128": "cuda",
+            "https://download.pytorch.org/whl/cpu": "cpu",
+            "https://download.pytorch.org/whl/rocm6.4": "rocm",
+            "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/": "rocm",
+            "https://repo.amd.com/rocm/whl/gfx1151/": "rocm",
+            "https://mirror.local/rocm-cache/cu128": "cuda",
+            "https://mirror.local/gfx-cache/cpu": "cpu",
+        }
+        for url, expected in cases.items():
+            out = sp.run(
+                ["sh", str(script), url], capture_output = True, text = True, timeout = 30
+            ).stdout.strip()
+            assert out == expected, f"{url} classified as {out!r}, expected {expected!r}"
+
+
+# TEST: CUDA_VISIBLE_DEVICES=""/-1 hides NVIDIA in every usable-GPU helper
+
+
+_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
+_STACK_SPEC = importlib.util.spec_from_file_location(
+    "studio_install_python_stack_followups", _STACK_PATH
+)
+assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
+stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
+sys.modules[_STACK_SPEC.name] = stack_mod
+_STACK_SPEC.loader.exec_module(stack_mod)
+
+
+def _stack_nvidia_usable(cvd):
+    """Drive install_python_stack._has_usable_nvidia_gpu with a mocked
+    nvidia-smi that always reports a GPU; cvd = None removes the env var."""
+
+    def fake_run(cmd, *args, **kwargs):
+        result = MagicMock()
+        result.returncode = 0
+        result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n"
+        return result
+
+    env = {} if cvd is None else {"CUDA_VISIBLE_DEVICES": cvd}
+    with (
+        patch.object(
+            stack_mod.shutil,
+            "which",
+            side_effect = lambda n: "/usr/bin/nvidia-smi" if n == "nvidia-smi" else None,
+        ),
+        patch.object(stack_mod.subprocess, "run", side_effect = fake_run),
+        patch.dict(stack_mod.os.environ, env, clear = False),
+    ):
+        if cvd is None:
+            stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        return stack_mod._has_usable_nvidia_gpu()
+
+
+class TestHiddenCvdNotUsable:
+    """CUDA_VISIBLE_DEVICES set to "" or "-1" deliberately hides every NVIDIA
+    device (mixed AMD+NVIDIA hosts steering work to the AMD card). All three
+    _has_usable_nvidia_gpu implementations (install_python_stack.py, install.sh,
+    setup.sh) must report the GPU as not usable so the AMD/CPU routes run,
+    matching install_llama_prebuilt.py's has_usable_nvidia."""
+
+    def test_python_unset_cvd_is_usable(self):
+        assert _stack_nvidia_usable(None) is True
+
+    def test_python_empty_cvd_not_usable(self):
+        assert _stack_nvidia_usable("") is False
+
+    def test_python_minus_one_not_usable(self):
+        assert _stack_nvidia_usable("-1") is False
+
+    def test_python_padded_minus_one_not_usable(self):
+        assert _stack_nvidia_usable(" -1 ") is False
+
+    def test_python_explicit_device_is_usable(self):
+        assert _stack_nvidia_usable("0") is True
+
+    def test_python_device_list_is_usable(self):
+        assert _stack_nvidia_usable("0,1") is True
+
+    def test_hidden_nvidia_restores_rocm_detection(self):
+        """Mixed host, NVIDIA hidden via CVD=-1, rocminfo reports gfx1100:
+        _has_rocm_gpu must proceed past the NVIDIA guard and return True
+        (before this fix the guard ignored CVD and blocked ROCm)."""
+
+        def fake_run(cmd, *args, **kwargs):
+            result = MagicMock()
+            result.returncode = 0
+            exe = str(cmd[0])
+            if exe.endswith("rocminfo"):
+                result.stdout = "  Name:                    gfx1100\n"
+            else:
+                result.stdout = "GPU 0: NVIDIA Fake (UUID: GPU-x)\n"
+            return result
+
+        which_map = {
+            "rocminfo": "/usr/bin/rocminfo",
+            "nvidia-smi": "/usr/bin/nvidia-smi",
+        }
+        with (
+            patch.object(stack_mod.shutil, "which", side_effect = which_map.get),
+            patch.object(stack_mod.subprocess, "run", side_effect = fake_run),
+            patch.dict(stack_mod.os.environ, {"CUDA_VISIBLE_DEVICES": "-1"}, clear = False),
+        ):
+            assert stack_mod._has_rocm_gpu() is True
+
+    @staticmethod
+    def _run_sh_helper(tmp_path, src: str, fn_names: list, cvd):
+        """Extract shell functions, run the usable-GPU one against a fake
+        nvidia-smi, and return "usable"/"not_usable"."""
+        import os as _os
+        import subprocess as sp
+
+        blocks = []
+        for name in fn_names:
+            start = src.find(f"{name}() {{")
+            assert start >= 0, f"{name} missing"
+            end = src.find("\n}", start) + 2
+            blocks.append(src[start:end])
+        fake_bin = tmp_path / "bin"
+        fake_bin.mkdir(exist_ok = True)
+        smi = fake_bin / "nvidia-smi"
+        smi.write_text("#!/bin/sh\necho 'GPU 0: NVIDIA Fake (UUID: GPU-x)'\n")
+        smi.chmod(0o755)
+        script = tmp_path / "probe.sh"
+        script.write_text(
+            "#!/bin/sh\n" + "\n".join(blocks) + "\n"
+            f"if {fn_names[-1]}; then echo usable; else echo not_usable; fi\n"
+        )
+        env = dict(_os.environ)
+        env["PATH"] = f"{fake_bin}:{env['PATH']}"
+        if cvd is None:
+            env.pop("CUDA_VISIBLE_DEVICES", None)
+        else:
+            env["CUDA_VISIBLE_DEVICES"] = cvd
+        return sp.run(
+            ["sh", str(script)], capture_output = True, text = True, timeout = 30, env = env
+        ).stdout.strip()
+
+    @pytest.mark.parametrize(
+        "cvd, expected",
+        [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")],
+    )
+    def test_install_sh_helper_cvd(self, tmp_path, cvd, expected):
+        src = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8")
+        out = self._run_sh_helper(
+            tmp_path,
+            src,
+            ["_run_bounded", "_cvd_hides_nvidia", "_has_usable_nvidia_gpu"],
+            cvd,
+        )
+        assert out == expected
+
+    @pytest.mark.parametrize(
+        "cvd, expected",
+        [(None, "usable"), ("", "not_usable"), ("-1", "not_usable"), ("0", "usable")],
+    )
+    def test_setup_sh_helper_cvd(self, tmp_path, cvd, expected):
+        src = SETUP_SH.read_text(encoding = "utf-8")
+        out = self._run_sh_helper(
+            tmp_path,
+            src,
+            ["_setup_run_smi", "_setup_cvd_hides_nvidia", "_setup_has_usable_nvidia_gpu"],
+            cvd,
+        )
+        assert out == expected

tests/studio/install/test_hf_auth.py

@@ -0,0 +1,114 @@
+"""Tests for Hugging Face auth on the llama.cpp prebuilt installer's fetches.
+
+Anonymous huggingface.co downloads (tiny GGUF validation model) share a
+per-IP rate limit that CI fleets exhaust (HTTP 429), forcing the prebuilt
+path into a source build. auth_headers now sends HF_TOKEN to huggingface.co
+hosts, and a redirect handler strips Authorization when the download is
+redirected to a different host (CDN signed URLs). All tests are offline.
+"""
+
+import importlib.util
+import sys
+import urllib.request
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+
+_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+_SPEC = importlib.util.spec_from_file_location(
+    "studio_install_llama_prebuilt_hf_auth", _MODULE_PATH
+)
+assert _SPEC is not None and _SPEC.loader is not None
+mod = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = mod
+_SPEC.loader.exec_module(mod)
+
+_TOKEN_VARS = ("GH_TOKEN", "GITHUB_TOKEN", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN")
+HF_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
+GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases"
+
+
+def _headers(url, env):
+    """auth_headers under a fully controlled token environment."""
+    with patch.dict(mod.os.environ, env, clear = False):
+        for var in _TOKEN_VARS:
+            if var not in env:
+                mod.os.environ.pop(var, None)
+        return mod.auth_headers(url)
+
+
+class TestAuthHeaderRouting:
+    def test_hf_token_sent_to_huggingface(self):
+        headers = _headers(HF_URL, {"HF_TOKEN": "hf_x"})
+        assert headers.get("Authorization") == "Bearer hf_x"
+
+    def test_hub_token_fallback(self):
+        headers = _headers(HF_URL, {"HUGGING_FACE_HUB_TOKEN": "hf_y"})
+        assert headers.get("Authorization") == "Bearer hf_y"
+
+    def test_hf_token_not_sent_to_github(self):
+        headers = _headers(GH_URL, {"HF_TOKEN": "hf_x"})
+        assert "Authorization" not in headers
+
+    def test_hf_token_not_sent_to_other_hosts(self):
+        headers = _headers("https://cdn-lfs.huggingface.co/x", {"HF_TOKEN": "hf_x"})
+        assert "Authorization" not in headers
+
+    def test_gh_token_not_sent_to_huggingface(self):
+        headers = _headers(HF_URL, {"GH_TOKEN": "gh_x"})
+        assert "Authorization" not in headers
+
+    def test_gh_token_still_wins_on_github(self):
+        headers = _headers(GH_URL, {"GH_TOKEN": "gh_x", "HF_TOKEN": "hf_x"})
+        assert headers.get("Authorization") == "Bearer gh_x"
+
+    def test_no_tokens_no_auth(self):
+        assert "Authorization" not in _headers(HF_URL, {})
+
+    def test_validation_model_url_is_hf(self):
+        assert mod.should_send_hf_auth(mod.TEST_MODEL_URL) is True
+
+
+class TestCrossHostRedirectStripsAuth:
+    def _redirect(self, newurl):
+        req = urllib.request.Request(HF_URL, headers = {"Authorization": "Bearer hf_x"})
+        handler = mod._CrossHostAuthStrippingRedirectHandler()
+        return handler.redirect_request(req, None, 302, "Found", {}, newurl)
+
+    def test_cross_host_redirect_drops_authorization(self):
+        new_request = self._redirect("https://cdn-lfs.huggingface.co/signed/blob")
+        assert new_request is not None
+        assert "Authorization" not in new_request.headers
+        assert "Authorization" not in new_request.unredirected_hdrs
+
+    def test_same_host_redirect_keeps_authorization(self):
+        new_request = self._redirect("https://huggingface.co/elsewhere/blob")
+        assert new_request is not None
+        assert new_request.headers.get("Authorization") == "Bearer hf_x"
+
+
+class TestDownloadBytesWiring:
+    def test_download_bytes_sends_hf_auth(self):
+        response = MagicMock()
+        response.__enter__ = lambda s: s
+        response.__exit__ = lambda s, *a: False
+        response.headers.get.return_value = None
+        response.read.side_effect = [b"data", b""]
+        with (
+            patch.object(mod._URL_OPENER, "open", return_value = response) as opened,
+            patch.dict(mod.os.environ, {"HF_TOKEN": "hf_x"}, clear = False),
+        ):
+            for var in ("GH_TOKEN", "GITHUB_TOKEN"):
+                mod.os.environ.pop(var, None)
+            data = mod.download_bytes(HF_URL)
+        assert data == b"data"
+        request = opened.call_args.args[0]
+        assert request.headers.get("Authorization") == "Bearer hf_x"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-q"]))

tests/studio/install/test_install_llama_prebuilt_logic.py

@@ -1,3 +1,4 @@
+import errno
 import importlib.util
 import io
 import json
@@ -12,9 +13,7 @@ import pytest
 
 PACKAGE_ROOT = Path(__file__).resolve().parents[3]
 MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
-SPEC = importlib.util.spec_from_file_location(
-    "studio_install_llama_prebuilt", MODULE_PATH
-)
+SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
 assert SPEC is not None and SPEC.loader is not None
 INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC)
 sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
@@ -30,6 +29,7 @@ ApprovedReleaseChecksums = INSTALL_LLAMA_PREBUILT.ApprovedReleaseChecksums
 hydrate_source_tree = INSTALL_LLAMA_PREBUILT.hydrate_source_tree
 validate_prebuilt_choice = INSTALL_LLAMA_PREBUILT.validate_prebuilt_choice
 activate_install_tree = INSTALL_LLAMA_PREBUILT.activate_install_tree
+activate_staged_dir = INSTALL_LLAMA_PREBUILT.activate_staged_dir
 create_install_staging_dir = INSTALL_LLAMA_PREBUILT.create_install_staging_dir
 sha256_file = INSTALL_LLAMA_PREBUILT.sha256_file
 source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name
@@ -208,6 +208,110 @@ def test_hydrate_source_tree_extracts_upstream_archive_contents(
     assert not (install_dir / f"llama.cpp-{upstream_tag}").exists()
 
 
+def test_release_asset_download_url():
+    fn = INSTALL_LLAMA_PREBUILT.release_asset_download_url
+    assert fn(
+        "unslothai/llama.cpp", "b9000-mix-abc1234", "llama.cpp-source-commit-deadbeef.tar.gz"
+    ) == (
+        "https://github.com/unslothai/llama.cpp/releases/download/"
+        "b9000-mix-abc1234/llama.cpp-source-commit-deadbeef.tar.gz"
+    )
+    # Any missing component -> None (no asset url, caller falls back to codeload).
+    assert fn(None, "b9000", "x.tar.gz") is None
+    assert fn("unslothai/llama.cpp", None, "x.tar.gz") is None
+    assert fn("unslothai/llama.cpp", "b9000", None) is None
+
+
+def _mk_source_tarball(path: Path, tag: str) -> None:
+    with tarfile.open(path, "w:gz") as archive:
+        add_bytes_to_tar(
+            archive, f"llama.cpp-{tag}/CMakeLists.txt", b"cmake_minimum_required(VERSION 3.14)\n"
+        )
+        add_bytes_to_tar(
+            archive,
+            f"llama.cpp-{tag}/convert_hf_to_gguf.py",
+            b"#!/usr/bin/env python3\nimport gguf\n",
+        )
+        add_bytes_to_tar(archive, f"llama.cpp-{tag}/gguf-py/gguf/__init__.py", b"__all__ = []\n")
+
+
+def test_hydrate_source_tree_prefers_release_asset_for_mix(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    # A mix build's merge commit is in no repo, so the codeload/archive URLs 404.
+    # hydrate must fetch the release asset and never touch codeload.
+    commit = "a" * 40
+    archive_path = tmp_path / "merged-source.tar.gz"
+    _mk_source_tarball(archive_path, f"b9000-mix-{commit[:7]}")
+    asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url(
+        "unslothai/llama.cpp", "b9000-mix-abc1234", f"llama.cpp-source-commit-{commit}.tar.gz"
+    )
+    codeload_urls = set(
+        INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit)
+    )
+    seen = []
+
+    def fake_download_file(url: str, destination: Path) -> None:
+        seen.append(url)
+        if url in codeload_urls:
+            raise AssertionError("codeload was hit even though the release asset was available")
+        assert url == asset_url
+        destination.write_bytes(archive_path.read_bytes())
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+
+    install_dir = tmp_path / "install"
+    work_dir = tmp_path / "work"
+    work_dir.mkdir()
+    hydrate_source_tree(
+        commit,
+        install_dir,
+        work_dir,
+        source_repo = "unslothai/llama.cpp",
+        expected_sha256 = sha256_file(archive_path),
+        exact_source = True,
+        asset_url = asset_url,
+    )
+    assert seen == [asset_url]
+    assert (install_dir / "CMakeLists.txt").exists()
+    assert (install_dir / "convert_hf_to_gguf.py").exists()
+
+
+def test_hydrate_source_tree_falls_back_to_codeload_when_asset_missing(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    # If the release asset 404s, fall back to codeload/archive (vanilla path).
+    commit = "b" * 40
+    archive_path = tmp_path / "vanilla-source.tar.gz"
+    _mk_source_tarball(archive_path, f"commit-{commit[:7]}")
+    asset_url = INSTALL_LLAMA_PREBUILT.release_asset_download_url(
+        "unslothai/llama.cpp", "b9000", f"llama.cpp-source-commit-{commit}.tar.gz"
+    )
+    codeload_urls = INSTALL_LLAMA_PREBUILT.commit_source_archive_urls("unslothai/llama.cpp", commit)
+
+    def fake_download_file(url: str, destination: Path) -> None:
+        if url == asset_url:
+            raise RuntimeError("404 Not Found")
+        assert url in codeload_urls
+        destination.write_bytes(archive_path.read_bytes())
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+
+    install_dir = tmp_path / "install"
+    work_dir = tmp_path / "work"
+    work_dir.mkdir()
+    hydrate_source_tree(
+        commit,
+        install_dir,
+        work_dir,
+        source_repo = "unslothai/llama.cpp",
+        expected_sha256 = sha256_file(archive_path),
+        exact_source = True,
+        asset_url = asset_url,
+    )
+    assert (install_dir / "CMakeLists.txt").exists()
+
+
 def test_validate_prebuilt_choice_creates_repo_shaped_linux_install(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ):
@@ -274,12 +378,8 @@ def test_validate_prebuilt_choice_creates_repo_shaped_linux_install(
         "preflight_linux_installed_binaries",
         lambda *args, **kwargs: None,
     )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None
-    )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None
-    )
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None)
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None)
 
     host = HostInfo(
         system = "Linux",
@@ -344,278 +444,6 @@ def test_validate_prebuilt_choice_creates_repo_shaped_linux_install(
     assert (install_dir / "BUILD_INFO.txt").exists()
 
 
-def test_simple_linux_direct_release_uses_published_source_checksums_for_branch(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    source_commit = "25b1bc9c2f9aa0a390b968ee1ffd9ff01340a3fe"
-    release = {
-        "tag_name": "llama-prebuilt-master-3a92bc9",
-        "assets": [
-            {
-                "name": "app-master-linux-x64-cuda13-newer.tar.gz",
-                "browser_download_url": "https://example.test/app-master-linux-x64-cuda13-newer.tar.gz",
-            },
-            {
-                "name": "llama-prebuilt-sha256.json",
-                "browser_download_url": "https://example.test/llama-prebuilt-sha256.json",
-            },
-        ],
-    }
-    checksums = ApprovedReleaseChecksums(
-        repo = "unslothai/llama.cpp",
-        release_tag = "llama-prebuilt-master-3a92bc9",
-        upstream_tag = "b9174",
-        source_commit = source_commit,
-        source_repo = "ggml-org/llama.cpp",
-        source_repo_url = "https://github.com/ggml-org/llama.cpp",
-        source_ref_kind = "branch",
-        requested_source_ref = "master",
-        resolved_source_ref = "master",
-        artifacts = {
-            "app-master-linux-x64-cuda13-newer.tar.gz": ApprovedArtifactHash(
-                asset_name = "app-master-linux-x64-cuda13-newer.tar.gz",
-                sha256 = "a" * 64,
-                repo = "unslothai/llama.cpp",
-                kind = "linux-cuda-app",
-            ),
-            INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
-                source_commit
-            ): ApprovedArtifactHash(
-                asset_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
-                    source_commit
-                ),
-                sha256 = "b" * 64,
-                repo = "ggml-org/llama.cpp",
-                kind = "exact-source",
-            ),
-        },
-    )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT,
-        "load_approved_release_checksums",
-        lambda repo, release_tag: checksums,
-    )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT,
-        "detected_linux_runtime_lines",
-        lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
-    )
-    host = HostInfo(
-        system = "Linux",
-        machine = "x86_64",
-        is_windows = False,
-        is_linux = True,
-        is_macos = False,
-        is_x86_64 = True,
-        is_arm64 = False,
-        nvidia_smi = None,
-        driver_cuda_version = (13, 1),
-        compute_caps = ["100"],
-        visible_cuda_devices = None,
-        has_physical_nvidia = True,
-        has_usable_nvidia = True,
-    )
-
-    plan = INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
-        release,
-        host,
-        "unslothai/llama.cpp",
-        "latest",
-    )
-
-    assert plan is not None
-    assert plan.llama_tag == "master"
-    assert plan.approved_checksums.upstream_tag == "b9174"
-    assert plan.approved_checksums.source_commit == source_commit
-    assert plan.attempts[0].expected_sha256 == "a" * 64
-    source_repo, source_ref, _source_archive, exact_source = (
-        INSTALL_LLAMA_PREBUILT.preferred_source_archive(
-            plan.approved_checksums, plan.llama_tag
-        )
-    )
-    assert source_repo == "ggml-org/llama.cpp"
-    assert source_ref == source_commit
-    assert exact_source is True
-
-
-@pytest.mark.parametrize(
-    "mutate, expected_match",
-    [
-        # Missing source_commit.
-        (
-            lambda c: setattr(c, "source_commit", None)
-            or setattr(c, "source_commit_short", None),
-            "exact source provenance",
-        ),
-        # source_commit present, but no exact-source archive hash.
-        (
-            lambda c: c.artifacts.pop(
-                INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
-                    c.source_commit
-                ),
-                None,
-            ),
-            "exact source provenance",
-        ),
-        # source_commit + exact-source archive present, but no source_repo.
-        (
-            lambda c: setattr(c, "source_repo", None)
-            or setattr(c, "source_repo_url", None),
-            "exact source provenance",
-        ),
-    ],
-    ids = [
-        "missing_source_commit",
-        "missing_exact_source_artifact",
-        "missing_source_repo",
-    ],
-)
-def test_simple_linux_direct_release_rejects_branch_without_exact_source_metadata(
-    monkeypatch: pytest.MonkeyPatch,
-    mutate,
-    expected_match,
-):
-    source_commit = "25b1bc9c2f9aa0a390b968ee1ffd9ff01340a3fe"
-    release = {
-        "tag_name": "llama-prebuilt-master-3a92bc9",
-        "assets": [
-            {
-                "name": "app-master-linux-x64-cuda13-newer.tar.gz",
-                "browser_download_url": "https://example.test/app-master-linux-x64-cuda13-newer.tar.gz",
-            },
-            {
-                "name": "llama-prebuilt-sha256.json",
-                "browser_download_url": "https://example.test/llama-prebuilt-sha256.json",
-            },
-        ],
-    }
-    checksums = ApprovedReleaseChecksums(
-        repo = "unslothai/llama.cpp",
-        release_tag = "llama-prebuilt-master-3a92bc9",
-        upstream_tag = "b9174",
-        source_commit = source_commit,
-        source_repo = "ggml-org/llama.cpp",
-        source_repo_url = "https://github.com/ggml-org/llama.cpp",
-        source_ref_kind = "branch",
-        requested_source_ref = "master",
-        resolved_source_ref = "master",
-        artifacts = {
-            "app-master-linux-x64-cuda13-newer.tar.gz": ApprovedArtifactHash(
-                asset_name = "app-master-linux-x64-cuda13-newer.tar.gz",
-                sha256 = "a" * 64,
-                repo = "unslothai/llama.cpp",
-                kind = "linux-cuda-app",
-            ),
-            INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
-                source_commit
-            ): ApprovedArtifactHash(
-                asset_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
-                    source_commit
-                ),
-                sha256 = "b" * 64,
-                repo = "ggml-org/llama.cpp",
-                kind = "exact-source",
-            ),
-        },
-    )
-    mutate(checksums)
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT,
-        "load_approved_release_checksums",
-        lambda repo, release_tag: checksums,
-    )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT,
-        "detected_linux_runtime_lines",
-        lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
-    )
-    host = HostInfo(
-        system = "Linux",
-        machine = "x86_64",
-        is_windows = False,
-        is_linux = True,
-        is_macos = False,
-        is_x86_64 = True,
-        is_arm64 = False,
-        nvidia_smi = None,
-        driver_cuda_version = (13, 1),
-        compute_caps = ["100"],
-        visible_cuda_devices = None,
-        has_physical_nvidia = True,
-        has_usable_nvidia = True,
-    )
-
-    with pytest.raises(PrebuiltFallback, match = expected_match):
-        INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
-            release,
-            host,
-            "unslothai/llama.cpp",
-            "latest",
-        )
-
-
-def test_simple_linux_direct_release_keeps_legacy_b_tag_path_without_checksums(
-    monkeypatch: pytest.MonkeyPatch,
-):
-    release = {
-        "tag_name": "b9999",
-        "assets": [
-            {
-                "name": "app-b9999-linux-x64-cuda13-newer.tar.gz",
-                "browser_download_url": "https://example.test/app-b9999-linux-x64-cuda13-newer.tar.gz",
-            },
-            {
-                "name": "llama-prebuilt-sha256.json",
-                "browser_download_url": "https://example.test/llama-prebuilt-sha256.json",
-            },
-        ],
-    }
-
-    def unexpected_checksum_load(repo: str, release_tag: str):
-        raise AssertionError(
-            "legacy b-tag direct releases should not require checksum metadata"
-        )
-
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT,
-        "load_approved_release_checksums",
-        unexpected_checksum_load,
-    )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT,
-        "detected_linux_runtime_lines",
-        lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
-    )
-    host = HostInfo(
-        system = "Linux",
-        machine = "x86_64",
-        is_windows = False,
-        is_linux = True,
-        is_macos = False,
-        is_x86_64 = True,
-        is_arm64 = False,
-        nvidia_smi = None,
-        driver_cuda_version = (13, 1),
-        compute_caps = ["100"],
-        visible_cuda_devices = None,
-        has_physical_nvidia = True,
-        has_usable_nvidia = True,
-    )
-
-    plan = INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
-        release,
-        host,
-        "unslothai/llama.cpp",
-        "latest",
-    )
-
-    assert plan is not None
-    assert plan.llama_tag == "b9999"
-    assert plan.release_tag == "b9999"
-    assert plan.approved_checksums.source_commit is None
-    assert plan.attempts[0].expected_sha256 is None
-
-
 def test_validate_prebuilt_choice_creates_repo_shaped_windows_install(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ):
@@ -667,12 +495,8 @@ def test_validate_prebuilt_choice_creates_repo_shaped_windows_install(
         "preflight_linux_installed_binaries",
         lambda *args, **kwargs: None,
     )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None
-    )
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None
-    )
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None)
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None)
 
     host = HostInfo(
         system = "Windows",
@@ -735,9 +559,7 @@ def test_validate_prebuilt_choice_creates_repo_shaped_windows_install(
 
 
 def test_activate_install_tree_restores_existing_install_after_activation_failure(
-    tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
-    capsys: pytest.CaptureFixture[str],
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
 ):
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
@@ -765,9 +587,7 @@ def test_activate_install_tree_restores_existing_install_after_activation_failur
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
         "confirm_install_tree",
-        lambda *_args, **_kwargs: (_ for _ in ()).throw(
-            RuntimeError("activation confirm failed")
-        ),
+        lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("activation confirm failed")),
     )
 
     with pytest.raises(
@@ -788,9 +608,7 @@ def test_activate_install_tree_restores_existing_install_after_activation_failur
 
 
 def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
-    tmp_path: Path,
-    monkeypatch: pytest.MonkeyPatch,
-    capsys: pytest.CaptureFixture[str],
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
 ):
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
@@ -818,9 +636,7 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
         "confirm_install_tree",
-        lambda *_args, **_kwargs: (_ for _ in ()).throw(
-            RuntimeError("activation confirm failed")
-        ),
+        lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("activation confirm failed")),
     )
 
     original_replace = INSTALL_LLAMA_PREBUILT.os.replace
@@ -847,14 +663,53 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
     captured = capsys.readouterr()
     output = captured.out + captured.err
     assert "rollback after failed activation also failed: restore failed" in output
-    assert (
-        "cleaning staging, install, and rollback paths before source build fallback"
-        in output
-    )
+    assert "cleaning staging, install, and rollback paths before source build fallback" in output
     assert "removing failed install path" in output
     assert "removing rollback path" in output
 
 
+def test_activate_staged_dir_copies_when_replace_hits_busy_lock(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+):
+    staging_dir = tmp_path / "llama.cpp.staging-test"
+    (staging_dir / "bin").mkdir(parents = True)
+    (staging_dir / "bin" / "ggml-base.dll").write_bytes(b"fake dll")
+    dst = tmp_path / "llama.cpp"
+
+    def denied_replace(src, dst_arg):
+        raise PermissionError(errno.EACCES, "Access is denied", str(src))
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", denied_replace)
+
+    activate_staged_dir(staging_dir, dst)
+
+    assert (dst / "bin" / "ggml-base.dll").read_bytes() == b"fake dll"
+    assert not staging_dir.exists()
+
+    captured = capsys.readouterr()
+    assert "falling back to file-by-file copy" in captured.out + captured.err
+
+
+def test_activate_staged_dir_reraises_non_busy_errors(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
+    staging_dir = tmp_path / "llama.cpp.staging-test"
+    staging_dir.mkdir()
+    (staging_dir / "new.txt").write_text("new install\n")
+    dst = tmp_path / "llama.cpp"
+
+    def out_of_space_replace(src, dst_arg):
+        raise OSError(errno.ENOSPC, "No space left on device", str(src))
+
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT.os, "replace", out_of_space_replace)
+
+    with pytest.raises(OSError, match = "No space left on device"):
+        activate_staged_dir(staging_dir, dst)
+
+    assert not dst.exists()
+    assert (staging_dir / "new.txt").read_text() == "new install\n"
+
+
 def test_binary_env_linux_includes_binary_parent_in_ld_library_path(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ):
@@ -956,7 +811,7 @@ def test_install_prebuilt_falls_back_to_older_release_plan(
     monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host)
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
-        "resolve_install_release_plans",
+        "resolve_simple_install_release_plans",
         lambda llama_tag, host, published_repo, published_release_tag: (
             "latest",
             [first_plan, second_plan],
@@ -1034,9 +889,7 @@ def write_linux_install_shape(install_dir: Path) -> None:
     (runtime_dir / "libggml-base.so.0").write_bytes(b"DLL")
     (runtime_dir / "libggml-cpu-x64.so.0").write_bytes(b"DLL")
     (runtime_dir / "libmtmd.so.0").write_bytes(b"DLL")
-    (install_dir / "convert_hf_to_gguf.py").write_text(
-        "#!/usr/bin/env python3\n", encoding = "utf-8"
-    )
+    (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8")
     (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True)
 
 
@@ -1060,9 +913,7 @@ def write_windows_install_shape(
         (runtime_dir / "cudart64_12.dll").write_bytes(b"DLL")
         (runtime_dir / "cublas64_12.dll").write_bytes(b"DLL")
         (runtime_dir / "cublasLt64_12.dll").write_bytes(b"DLL")
-    (install_dir / "convert_hf_to_gguf.py").write_text(
-        "#!/usr/bin/env python3\n", encoding = "utf-8"
-    )
+    (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8")
     (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True)
 
 
@@ -1085,9 +936,7 @@ def write_macos_install_shape(
         (runtime_dir / "libggml.0.dylib").write_bytes(b"DLL")
     if include_libmtmd:
         (runtime_dir / "libmtmd.0.dylib").write_bytes(b"DLL")
-    (install_dir / "convert_hf_to_gguf.py").write_text(
-        "#!/usr/bin/env python3\n", encoding = "utf-8"
-    )
+    (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8")
     (install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True)
 
 
@@ -1166,8 +1015,7 @@ def test_existing_install_matches_plan_false_without_fingerprint(tmp_path: Path)
     install_dir.mkdir()
     write_linux_install_shape(install_dir)
     (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(
-        json.dumps({"tag": "b9001", "asset": "llama-b9001-bin-ubuntu-x64.tar.gz"})
-        + "\n",
+        json.dumps({"tag": "b9001", "asset": "llama-b9001-bin-ubuntu-x64.tar.gz"}) + "\n",
         encoding = "utf-8",
     )
 
@@ -1230,9 +1078,7 @@ def test_existing_install_matches_plan_false_with_malformed_metadata(tmp_path: P
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
     write_linux_install_shape(install_dir)
-    (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(
-        "{not-json\n", encoding = "utf-8"
-    )
+    (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text("{not-json\n", encoding = "utf-8")
 
     host = HostInfo(
         system = "Linux",
@@ -1363,9 +1209,7 @@ def test_existing_install_matches_plan_windows_cpu_requires_llama_dll(tmp_path:
 def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path: Path):
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
-    write_windows_install_shape(
-        install_dir, include_llama_dll = True, include_cuda_dll = True
-    )
+    write_windows_install_shape(install_dir, include_llama_dll = True, include_cuda_dll = True)
 
     host = HostInfo(
         system = "Windows",
@@ -1434,9 +1278,7 @@ def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path:
     assert existing_install_matches_plan(install_dir, host, plan) is False
 
 
-def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(
-    tmp_path: Path,
-):
+def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_path: Path):
     """When the choice ships a paired cudart bundle (#5106), the install
     is considered stale unless cudart64_*.dll and cublas64_*.dll are
     actually on disk. Otherwise existing broken installs would keep
@@ -1552,9 +1394,7 @@ def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(
     assert existing_install_matches_plan(install_dir, host, plan) is False
 
 
-def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(
-    tmp_path: Path,
-):
+def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(tmp_path: Path):
     """If the choice has no paired runtime archive (manifest dropped it,
     or upstream did not ship cudart), legacy installs without cudart on
     disk must still pass the health check -- otherwise the installer
@@ -1634,9 +1474,7 @@ def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(
     assert existing_install_matches_plan(install_dir, host, plan) is True
 
 
-def test_existing_install_fingerprint_changes_when_cudart_pair_added(
-    tmp_path: Path,
-):
+def test_existing_install_fingerprint_changes_when_cudart_pair_added(tmp_path: Path):
     """Existing pre-#5322 Windows CUDA installs (no paired cudart) must
     be treated as stale once the choice gains a runtime archive,
     otherwise the fingerprint match would keep skipping the reinstall
@@ -1901,7 +1739,7 @@ def test_install_prebuilt_skips_download_when_existing_install_matches(
     monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host)
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
-        "resolve_install_release_plans",
+        "resolve_simple_install_release_plans",
         lambda llama_tag, host, published_repo, published_release_tag: (
             "latest",
             [plan],
@@ -1911,9 +1749,7 @@ def test_install_prebuilt_skips_download_when_existing_install_matches(
         INSTALL_LLAMA_PREBUILT,
         "download_validation_model",
         lambda *args, **kwargs: (_ for _ in ()).throw(
-            AssertionError(
-                "matching install should skip before validation model download"
-            )
+            AssertionError("matching install should skip before validation model download")
         ),
     )
 
@@ -1993,7 +1829,7 @@ def test_install_prebuilt_does_not_skip_unhealthy_existing_install(
     monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host)
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
-        "resolve_install_release_plans",
+        "resolve_simple_install_release_plans",
         lambda llama_tag, host, published_repo, published_release_tag: (
             "latest",
             [plan],
@@ -2121,7 +1957,7 @@ def test_install_prebuilt_skips_when_older_release_fallback_matches_existing_ins
     monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host)
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
-        "resolve_install_release_plans",
+        "resolve_simple_install_release_plans",
         lambda llama_tag, host, published_repo, published_release_tag: (
             "latest",
             [latest_plan, fallback_plan],
@@ -2268,7 +2104,7 @@ def test_install_prebuilt_skips_same_release_fallback_attempt_when_installed(
     monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host)
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
-        "resolve_install_release_plans",
+        "resolve_simple_install_release_plans",
         lambda llama_tag, host, published_repo, published_release_tag: (
             "latest",
             [plan],
@@ -2387,7 +2223,7 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p
     monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host)
     monkeypatch.setattr(
         INSTALL_LLAMA_PREBUILT,
-        "resolve_install_release_plans",
+        "resolve_simple_install_release_plans",
         lambda llama_tag, host, published_repo, published_release_tag: (
             "latest",
             [latest_plan, older_plan],
@@ -2429,9 +2265,7 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p
         (staging_dir / "marker.txt").write_text("ready\n")
         return attempts[0], staging_dir, initial_fallback_used
 
-    monkeypatch.setattr(
-        INSTALL_LLAMA_PREBUILT, "validate_prebuilt_attempts", fake_validate
-    )
+    monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_prebuilt_attempts", fake_validate)
 
     activated = {}
     monkeypatch.setattr(
@@ -2449,10 +2283,7 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p
 
     install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
 
-    assert attempted == [
-        ("b9002", "release-2", "upstream"),
-        ("b9001", "release-1", "upstream"),
-    ]
+    assert attempted == [("b9002", "release-2", "upstream"), ("b9001", "release-1", "upstream")]
     assert activated["install_dir"] == install_dir
 
 
@@ -2461,7 +2292,11 @@ def io_bytes(data: bytes):
 
 
 def add_bytes_to_tar(
-    archive: tarfile.TarFile, name: str, data: bytes, *, mode: int = 0o644
+    archive: tarfile.TarFile,
+    name: str,
+    data: bytes,
+    *,
+    mode: int = 0o644,
 ) -> None:
     info = tarfile.TarInfo(name)
     info.size = len(data)
@@ -2476,9 +2311,7 @@ def add_symlink_to_tar(archive: tarfile.TarFile, name: str, target: str) -> None
     archive.addfile(info)
 
 
-def test_existing_install_matches_choice_fails_when_install_tree_incomplete(
-    tmp_path: Path,
-):
+def test_existing_install_matches_choice_fails_when_install_tree_incomplete(tmp_path: Path):
     """confirm_install_tree guard rejects installs missing critical files."""
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
@@ -2567,9 +2400,7 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete(
     )
 
 
-def test_existing_install_matches_choice_fails_when_install_tree_incomplete_macos(
-    tmp_path: Path,
-):
+def test_existing_install_matches_choice_fails_when_install_tree_incomplete_macos(tmp_path: Path):
     """confirm_install_tree guard rejects macOS arm64 installs missing critical files."""
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
@@ -2706,9 +2537,7 @@ def test_paired_runtime_dll_patterns_excludes_executables() -> None:
         assert paired_runtime_dll_patterns(non_windows) == []
 
 
-def test_runtime_overlay_cannot_overwrite_main_archive_payload(
-    tmp_path: Path,
-) -> None:
+def test_runtime_overlay_cannot_overwrite_main_archive_payload(tmp_path: Path) -> None:
     """End-to-end: a malformed runtime archive containing
     ``llama-server.exe`` alongside the real cudart DLLs must NOT
     replace the main archive's ``llama-server.exe``.
@@ -2772,15 +2601,20 @@ def test_runtime_overlay_cannot_overwrite_main_archive_payload(
 
     orig_download = INSTALL_LLAMA_PREBUILT.download_file_verified
 
-    def fake_download(url, target_path, *, expected_sha256 = None, label = None, **kw):
+    def fake_download(
+        url,
+        target_path,
+        *,
+        expected_sha256 = None,
+        label = None,
+        **kw,
+    ):
         src = main_zip if "cudart" not in url else runtime_zip
         _shutil.copy2(src, target_path)
         if expected_sha256:
             actual = hashlib.sha256(Path(target_path).read_bytes()).hexdigest()
             if actual != expected_sha256:
-                raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(
-                    f"sha256 mismatch on {label}"
-                )
+                raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(f"sha256 mismatch on {label}")
 
     INSTALL_LLAMA_PREBUILT.download_file_verified = fake_download
     try:
@@ -2792,16 +2626,107 @@ def test_runtime_overlay_cannot_overwrite_main_archive_payload(
     server = release_dir / "llama-server.exe"
     assert server.exists()
     assert server.read_bytes() == b"MAIN-SERVER", (
-        "runtime archive overwrote main llama-server.exe; "
-        f"got {server.read_bytes()!r}"
+        "runtime archive overwrote main llama-server.exe; " f"got {server.read_bytes()!r}"
     )
     for name in ("cudart64_12.dll", "cublas64_12.dll", "cublasLt64_12.dll"):
         assert (release_dir / name).exists(), f"missing {name}"
 
 
-def test_python_runtime_dirs_covers_cu13_and_library_bin(
-    monkeypatch, tmp_path: Path
-) -> None:
+def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(tmp_path: Path) -> None:
+    install_from_archives = INSTALL_LLAMA_PREBUILT.install_from_archives
+
+    work = tmp_path / "work"
+    install = tmp_path / "install"
+    archives = tmp_path / "archives"
+    work.mkdir()
+    install.mkdir()
+    archives.mkdir()
+
+    bundle = archives / "app-b9334-linux-x64-cuda13-newer.tar.gz"
+    with tarfile.open(bundle, "w:gz") as archive:
+        for name in (
+            "llama-cli",
+            "llama-server",
+            "llama-quantize",
+            "libllama-cli-impl.so",
+            "libllama-server-impl.so",
+            "libllama-quantize-impl.so",
+            "libllama-common.so",
+            "libllama.so",
+            "libggml.so",
+            "libggml-base.so",
+            "libmtmd.so",
+            "libggml-cpu-x64.so",
+            "libggml-cuda.so",
+        ):
+            payload = f"{name}\n".encode()
+            member = tarfile.TarInfo(name)
+            member.size = len(payload)
+            archive.addfile(member, io.BytesIO(payload))
+
+    import hashlib
+    import shutil as _shutil
+
+    bundle_sha = hashlib.sha256(bundle.read_bytes()).hexdigest()
+    choice = AssetChoice(
+        repo = "unslothai/llama.cpp",
+        tag = "b9334",
+        name = bundle.name,
+        url = f"https://example.com/{bundle.name}",
+        source_label = "published",
+        install_kind = "linux-cuda",
+        runtime_line = "cuda13",
+        expected_sha256 = bundle_sha,
+    )
+    host = HostInfo(
+        system = "Linux",
+        machine = "x86_64",
+        is_windows = False,
+        is_linux = True,
+        is_macos = False,
+        is_x86_64 = True,
+        is_arm64 = False,
+        nvidia_smi = None,
+        driver_cuda_version = (13, 0),
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = True,
+        has_usable_nvidia = True,
+    )
+
+    orig_download = INSTALL_LLAMA_PREBUILT.download_file_verified
+
+    def fake_download(
+        url,
+        target_path,
+        *,
+        expected_sha256 = None,
+        label = None,
+        **kw,
+    ):
+        _shutil.copy2(bundle, target_path)
+        if expected_sha256:
+            actual = hashlib.sha256(Path(target_path).read_bytes()).hexdigest()
+            if actual != expected_sha256:
+                raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(f"sha256 mismatch on {label}")
+
+    INSTALL_LLAMA_PREBUILT.download_file_verified = fake_download
+    try:
+        install_from_archives(choice, host, install, work)
+    finally:
+        INSTALL_LLAMA_PREBUILT.download_file_verified = orig_download
+
+    runtime_dir = install / "build" / "bin"
+    for name in (
+        "libllama-cli-impl.so",
+        "libllama-server-impl.so",
+        "libllama-quantize-impl.so",
+    ):
+        assert (runtime_dir / name).exists(), f"missing {name}"
+    assert not (runtime_dir / "llama-cli").exists()
+
+
+def test_python_runtime_dirs_covers_cu13_and_library_bin(monkeypatch, tmp_path: Path) -> None:
     """Installer-side runtime DLL discovery must scan the same path
     set as the backend ``_windows_pip_nvidia_dll_dirs``: legacy
     ``nvidia/<pkg>/bin``, current ``nvidia/<pkg>/bin/x86_64``

tests/studio/install/test_llama_pr_force_and_source.py

@@ -35,11 +35,13 @@ requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not availa
 # Helpers
 # ---------------------------------------------------------------------------
 def run_bash(
-    script: str, *, timeout: int = 60, env: dict | None = None
+    script: str,
+    *,
+    timeout: int = 60,
+    env: dict | None = None,
 ) -> subprocess.CompletedProcess:
-    """Run a bash script fragment and return the CompletedProcess.
-    60s default tolerates slow shell startup on heavily-loaded CI
-    runners; the scripts themselves run in well under a second."""
+    """Run a bash script fragment. 60s default tolerates slow CI shell
+    startup; the scripts themselves run in well under a second."""
     run_env = os.environ.copy()
     if env:
         run_env.update(env)
@@ -53,12 +55,13 @@ def run_bash(
 
 
 def run_pwsh(
-    script: str, *, timeout: int = 60, env: dict | None = None
+    script: str,
+    *,
+    timeout: int = 60,
+    env: dict | None = None,
 ) -> subprocess.CompletedProcess:
-    """Run a PowerShell script fragment and return the CompletedProcess.
-    60s default tolerates slow pwsh startup on heavily-loaded CI
-    runners; the scripts themselves run in well under a second.
-    A 10s budget previously surfaced as a flaky TimeoutExpired."""
+    """Run a PowerShell script fragment. 60s default tolerates slow CI pwsh
+    startup (a 10s budget was flaky); the scripts run in under a second."""
     run_env = os.environ.copy()
     run_env["NO_COLOR"] = "1"
     if env:
@@ -383,10 +386,7 @@ class TestSourcePatternsSh:
         assert '_DEFAULT_LLAMA_PR_FORCE=""' in self.content
 
     def test_has_default_source(self):
-        assert (
-            '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"'
-            in self.content
-        )
+        assert '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"' in self.content
 
     def test_has_pr_force_env_read(self):
         assert "UNSLOTH_LLAMA_PR_FORCE" in self.content
@@ -416,8 +416,7 @@ class TestSourcePatternsSh:
     def test_clone_urls_parameterized_pr_path(self):
         """PR clone path uses ${_LLAMA_SOURCE}.git, not hardcoded URL."""
         pr_clone_idx = self.content.index(
-            'if [ -n "$_LLAMA_PR" ]; then\n'
-            '            run_quiet_no_exit "clone llama.cpp"'
+            'if [ -n "$_LLAMA_PR" ]; then\n            run_quiet_no_exit "clone llama.cpp"'
         )
         else_idx = self.content.index("else\n", pr_clone_idx)
         pr_block = self.content[pr_clone_idx:else_idx]
@@ -437,9 +436,7 @@ class TestSourcePatternsSh:
         lines = self.content.splitlines()
         for i, line in enumerate(lines, 1):
             if "git clone" in line and "ggml-org/llama.cpp.git" in line:
-                pytest.fail(
-                    f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}"
-                )
+                pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
 
 
 # =========================================================================
@@ -456,10 +453,7 @@ class TestSourcePatternsPs1:
         assert '$DefaultLlamaPrForce = ""' in self.content
 
     def test_has_default_source(self):
-        assert (
-            '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"'
-            in self.content
-        )
+        assert '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"' in self.content
 
     def test_has_pr_force_env_read(self):
         assert "$env:UNSLOTH_LLAMA_PR_FORCE" in self.content
@@ -469,11 +463,13 @@ class TestSourcePatternsPs1:
         assert "$LlamaSource = $DefaultLlamaSource" in self.content
 
     def test_release_repo_override_removed(self):
+        # No env-based release-repo override; the repo is chosen by GPU detection
+        # (GPU -> fork, CPU -> ggml-org), mirroring setup.sh.
+        assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content
         assert (
-            "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)"
-            not in self.content
+            "$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm) "
+            '{ "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" }' in self.content
         )
-        assert '$HelperReleaseRepo = "ggml-org/llama.cpp"' in self.content
 
     def test_force_compile_skips_prebuilt_resolution_early(self):
         assert 'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {' in self.content
@@ -491,9 +487,7 @@ class TestSourcePatternsPs1:
 
     def test_clone_urls_parameterized_pr_path(self):
         """PR clone path uses $LlamaSource.git, not hardcoded URL."""
-        pr_idx = self.content.index(
-            "if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp")
-        )
+        pr_idx = self.content.index("if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp"))
         else_idx = self.content.index("} else {", pr_idx)
         pr_block = self.content[pr_idx:else_idx]
         assert '"$LlamaSource.git"' in pr_block
@@ -511,9 +505,7 @@ class TestSourcePatternsPs1:
         lines = self.content.splitlines()
         for i, line in enumerate(lines, 1):
             if "git clone" in line and "ggml-org/llama.cpp.git" in line:
-                pytest.fail(
-                    f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}"
-                )
+                pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
 
 
 # =========================================================================

tests/studio/install/test_macos_version_compat.py

@@ -0,0 +1,330 @@
+"""Tests for the host-macOS-version-aware llama.cpp prebuilt selection added
+for the Mac "Failing CI" fix.
+
+Covers: parse_macos_version, host_supports_macos_minos, the pure-Python Mach-O
+minimum-OS parser (macho_minimum_macos), the dyld-incompatibility classifier,
+the install preflight that rejects a too-new prebuilt, and the deeper macOS
+release walk-back in resolve_simple_install_release_plans.
+
+No GPU, no network, no torch, no real Mach-O toolchain required -- the Mach-O
+samples are synthesized in-process and all I/O is monkeypatched.
+"""
+
+import importlib.util
+import struct
+import sys
+from pathlib import Path
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt_macos", MODULE_PATH)
+assert SPEC is not None and SPEC.loader is not None
+ILP = importlib.util.module_from_spec(SPEC)
+sys.modules[SPEC.name] = ILP
+SPEC.loader.exec_module(ILP)
+
+HostInfo = ILP.HostInfo
+PrebuiltFallback = ILP.PrebuiltFallback
+
+_CPU_TYPE_ARM64 = 0x0100000C
+_CPU_TYPE_X86_64 = 0x01000007
+
+
+def make_macos_host(macos_version, *, arm64 = True):
+    return HostInfo(
+        system = "Darwin",
+        machine = "arm64" if arm64 else "x86_64",
+        is_windows = False,
+        is_linux = False,
+        is_macos = True,
+        is_x86_64 = not arm64,
+        is_arm64 = arm64,
+        nvidia_smi = None,
+        driver_cuda_version = None,
+        compute_caps = [],
+        visible_cuda_devices = None,
+        has_physical_nvidia = False,
+        has_usable_nvidia = False,
+        macos_version = macos_version,
+    )
+
+
+def thin_macho(
+    minos = (14, 0),
+    *,
+    cputype = _CPU_TYPE_ARM64,
+    build_version = True,
+):
+    """Synthesize a minimal little-endian 64-bit Mach-O carrying a macOS
+    minimum-version load command."""
+    encoded = (minos[0] << 16) | (minos[1] << 8)
+    if build_version:
+        # LC_BUILD_VERSION: cmd, cmdsize, platform(=1 macOS), minos, sdk, ntools
+        load_command = struct.pack("<6I", 0x32, 24, 1, encoded, encoded, 0)
+    else:
+        # LC_VERSION_MIN_MACOSX: cmd, cmdsize, version, sdk
+        load_command = struct.pack("<4I", 0x24, 16, encoded, encoded)
+    header = struct.pack("<8I", 0xFEEDFACF, cputype, 0, 0x2, 1, len(load_command), 0, 0)
+    return header + load_command
+
+
+def fat_macho(slices):
+    """Synthesize a big-endian universal binary from (cputype, thin_bytes)."""
+    header = struct.pack(">2I", 0xCAFEBABE, len(slices))
+    data_offset = 8 + 20 * len(slices)
+    arch_entries = b""
+    body = b""
+    for cputype, thin in slices:
+        offset = data_offset + len(body)
+        arch_entries += struct.pack(">5I", cputype, 0, offset, len(thin), 0)
+        body += thin
+    return header + arch_entries + body
+
+
+class TestParseMacosVersion:
+    @pytest.mark.parametrize(
+        "value, expected",
+        [
+            ("14.7.1", (14, 7)),
+            ("15.5", (15, 5)),
+            ("26.0", (26, 0)),
+            ("26", (26, 0)),
+            ("13", (13, 0)),
+            ("", None),
+            (None, None),
+            ("not-a-version", None),
+        ],
+    )
+    def test_parse(self, value, expected):
+        assert ILP.parse_macos_version(value) == expected
+
+
+class TestHostSupportsMacosMinos:
+    def test_older_host_rejects_newer_prebuilt(self):
+        assert not ILP.host_supports_macos_minos(make_macos_host((14, 0)), (26, 0))
+
+    def test_same_version_supported(self):
+        assert ILP.host_supports_macos_minos(make_macos_host((26, 0)), (26, 0))
+
+    def test_newer_host_supports_older_prebuilt(self):
+        assert ILP.host_supports_macos_minos(make_macos_host((15, 5)), (14, 0))
+
+    def test_unknown_host_defers_to_runtime(self):
+        assert ILP.host_supports_macos_minos(make_macos_host(None), (26, 0))
+
+    def test_unknown_minos_defers_to_runtime(self):
+        assert ILP.host_supports_macos_minos(make_macos_host((14, 0)), None)
+
+
+class TestMachoMinimumMacos:
+    def test_build_version_thin(self, tmp_path):
+        path = tmp_path / "lib.dylib"
+        path.write_bytes(thin_macho((26, 0)))
+        assert ILP.macho_minimum_macos(path) == (26, 0)
+
+    def test_legacy_version_min_thin(self, tmp_path):
+        path = tmp_path / "lib.dylib"
+        path.write_bytes(thin_macho((14, 0), build_version = False))
+        assert ILP.macho_minimum_macos(path) == (14, 0)
+
+    def test_universal_prefers_host_arch_slice(self, tmp_path):
+        # arm64 slice needs macOS 14, x86_64 slice needs macOS 26.
+        path = tmp_path / "fat"
+        path.write_bytes(
+            fat_macho(
+                [
+                    (_CPU_TYPE_ARM64, thin_macho((14, 0), cputype = _CPU_TYPE_ARM64)),
+                    (_CPU_TYPE_X86_64, thin_macho((26, 0), cputype = _CPU_TYPE_X86_64)),
+                ]
+            )
+        )
+        assert ILP.macho_minimum_macos(path, make_macos_host((14, 0))) == (14, 0)
+        assert ILP.macho_minimum_macos(path, make_macos_host((26, 0), arm64 = False)) == (26, 0)
+
+    def test_non_macho_returns_none(self, tmp_path):
+        path = tmp_path / "script.sh"
+        path.write_bytes(b'#!/bin/sh\nexec real "$@"\n')
+        assert ILP.macho_minimum_macos(path) is None
+
+    def test_missing_file_returns_none(self, tmp_path):
+        assert ILP.macho_minimum_macos(tmp_path / "nope") is None
+
+
+class TestLooksLikeMacosIncompatibility:
+    def test_built_for_newer_os(self):
+        assert ILP.looks_like_macos_incompatibility(
+            "dyld: ... (built for macOS 26.0 which is newer than running OS)"
+        )
+
+    def test_metal_residency_symbol(self):
+        assert ILP.looks_like_macos_incompatibility(
+            "Symbol not found: _OBJC_CLASS_$_MTLResidencySetDescriptor"
+        )
+
+    def test_benign_error(self):
+        assert not ILP.looks_like_macos_incompatibility("some unrelated failure")
+
+    def test_empty(self):
+        assert not ILP.looks_like_macos_incompatibility("")
+
+
+class TestPreflightMacosInstalledBinaries:
+    def _install_dir(self, tmp_path, dylib_minos):
+        bin_dir = tmp_path / "build" / "bin"
+        bin_dir.mkdir(parents = True)
+        (bin_dir / "libggml-metal.dylib").write_bytes(thin_macho(dylib_minos))
+        server = tmp_path / "llama-server"
+        server.write_bytes(thin_macho(dylib_minos))
+        quantize = tmp_path / "llama-quantize"
+        quantize.write_bytes(thin_macho(dylib_minos))
+        return tmp_path, (server, quantize)
+
+    def test_rejects_too_new_dylib(self, tmp_path):
+        install_dir, binaries = self._install_dir(tmp_path, (26, 0))
+        with pytest.raises(PrebuiltFallback, match = "newer macOS"):
+            ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host((14, 0)))
+
+    def test_accepts_compatible_prebuilt(self, tmp_path):
+        install_dir, binaries = self._install_dir(tmp_path, (14, 0))
+        # Must not raise on a macOS 15 host.
+        ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host((15, 5)))
+
+    def test_skips_when_host_version_unknown(self, tmp_path):
+        install_dir, binaries = self._install_dir(tmp_path, (26, 0))
+        # Unknown host version -> defer to runtime validation, do not raise.
+        ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host(None))
+
+    def test_noop_on_non_macos_host(self, tmp_path):
+        install_dir, binaries = self._install_dir(tmp_path, (26, 0))
+        linux_host = HostInfo(
+            system = "Linux",
+            machine = "x86_64",
+            is_windows = False,
+            is_linux = True,
+            is_macos = False,
+            is_x86_64 = True,
+            is_arm64 = False,
+            nvidia_smi = None,
+            driver_cuda_version = None,
+            compute_caps = [],
+            visible_cuda_devices = None,
+            has_physical_nvidia = False,
+            has_usable_nvidia = False,
+        )
+        ILP.preflight_macos_installed_binaries(binaries, install_dir, linux_host)
+
+
+def _fake_macos_releases(tags):
+    return [
+        {
+            "tag_name": tag,
+            "assets": [
+                {
+                    "name": f"llama-{tag}-bin-macos-arm64.tar.gz",
+                    "browser_download_url": f"https://example.com/{tag}.tar.gz",
+                }
+            ],
+        }
+        for tag in tags
+    ]
+
+
+class TestMacosReleasePin:
+    """A known pre-26 macOS host deterministically pins the last upstream release
+    whose prebuilt loads on it (b9415) instead of walking back release by release;
+    macOS 26+ and unknown-version hosts keep normal latest selection with the
+    conservative 2-release default."""
+
+    TAGS = [f"b{n}" for n in range(9442, 9400, -1)]  # newest-first, includes b9415
+
+    def _patch_releases(self, monkeypatch):
+        def fake_iter(repo, published_release_tag, requested_tag):
+            # The real iterator yields only the requested tag when one is pinned.
+            if requested_tag and requested_tag != "latest":
+                return _fake_macos_releases([requested_tag])
+            return _fake_macos_releases(self.TAGS)
+
+        monkeypatch.setattr(ILP, "iter_release_payloads_by_time", fake_iter)
+
+    def test_pre26_host_pins_b9415(self, monkeypatch):
+        self._patch_releases(monkeypatch)
+        tag, plans = ILP.resolve_simple_install_release_plans(
+            "latest",
+            make_macos_host((14, 0)),
+            "ggml-org/llama.cpp",
+            "",
+        )
+        assert tag == ILP._PINNED_MACOS_FALLBACK_TAG == "b9415"
+        assert len(plans) == 1
+        assert plans[0].release_tag == "b9415"
+
+    def test_tahoe_host_takes_latest(self, monkeypatch):
+        self._patch_releases(monkeypatch)
+        tag, plans = ILP.resolve_simple_install_release_plans(
+            "latest",
+            make_macos_host((26, 0)),
+            "ggml-org/llama.cpp",
+            "",
+        )
+        assert tag == "latest"
+        assert plans[0].release_tag == self.TAGS[0]  # newest release
+        assert len(plans) == ILP.DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS
+
+    def test_unknown_macos_host_uses_default(self, monkeypatch):
+        self._patch_releases(monkeypatch)
+        _tag, plans = ILP.resolve_simple_install_release_plans(
+            "latest",
+            make_macos_host(None),
+            "ggml-org/llama.cpp",
+            "",
+        )
+        assert len(plans) == ILP.DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS
+
+
+class TestForwardsBackwardsCompat:
+    """The gate is host >= prebuilt minos with no hardcoded version, so it holds
+    for older and future macOS alike. Emulate the walk-back over a release set
+    spanning several minos tiers and assert each host takes the newest release
+    it can load."""
+
+    # Newest first: future 27 builds, current 26 builds, an old 14 tier, a 13.
+    RELEASES = [
+        ("b9600", (27, 0)),
+        ("b9450", (26, 0)),
+        ("b9415", (14, 0)),
+        ("b8300", (13, 0)),
+    ]
+
+    def _select(self, tmp_path, host_version):
+        for tag, minos in self.RELEASES:
+            bin_dir = tmp_path / tag / "build" / "bin"
+            bin_dir.mkdir(parents = True)
+            (bin_dir / "libggml-metal.dylib").write_bytes(thin_macho(minos))
+            try:
+                ILP.preflight_macos_installed_binaries(
+                    (), tmp_path / tag, make_macos_host(host_version)
+                )
+                return tag
+            except PrebuiltFallback:
+                continue
+        return None
+
+    @pytest.mark.parametrize(
+        "host_version, expected",
+        [
+            ((13, 0), "b8300"),  # older host takes the older prebuilt
+            ((14, 7), "b9415"),  # backwards: skip 26/27, take newest that loads
+            ((15, 5), "b9415"),
+            ((26, 0), "b9450"),  # unchanged: newest <= host
+            ((27, 1), "b9600"),  # forwards: future host takes the future build
+        ],
+    )
+    def test_selects_newest_loadable(self, tmp_path, host_version, expected):
+        assert self._select(tmp_path, host_version) == expected
+
+    def test_host_below_prebuilt_floor_falls_through(self, tmp_path):
+        # macOS 12 is below every prebuilt -> nothing matches -> source build.
+        assert self._select(tmp_path, (12, 0)) is None

tests/studio/install/test_pr4562_bugfixes.py

@@ -29,9 +29,7 @@ import pytest
 # ---------------------------------------------------------------------------
 PACKAGE_ROOT = Path(__file__).resolve().parents[3]
 MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
-SPEC = importlib.util.spec_from_file_location(
-    "studio_install_llama_prebuilt", MODULE_PATH
-)
+SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
 assert SPEC is not None and SPEC.loader is not None
 MOD = importlib.util.module_from_spec(SPEC)
 sys.modules[SPEC.name] = MOD
@@ -74,7 +72,12 @@ def make_host(*, system: str) -> HostInfo:
 BASH = "/bin/bash"
 
 
-def run_bash(script: str, *, timeout: int = 10, env: dict | None = None) -> str:
+def run_bash(
+    script: str,
+    *,
+    timeout: int = 10,
+    env: dict | None = None,
+) -> str:
     """Run a bash script fragment and return its stdout."""
     run_env = os.environ.copy()
     if env:
@@ -113,9 +116,7 @@ class TestBinaryEnvCrossPlatform:
         env = binary_env(binary_path, install_dir, host)
         ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
         assert str(bin_dir) in ld_dirs, f"build/bin not in LD_LIBRARY_PATH: {ld_dirs}"
-        assert (
-            str(install_dir) in ld_dirs
-        ), f"install_dir not in LD_LIBRARY_PATH: {ld_dirs}"
+        assert str(install_dir) in ld_dirs, f"install_dir not in LD_LIBRARY_PATH: {ld_dirs}"
 
     def test_linux_binary_parent_comes_before_install_dir(
         self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
@@ -134,9 +135,7 @@ class TestBinaryEnvCrossPlatform:
         ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
         bin_idx = ld_dirs.index(str(bin_dir))
         install_idx = ld_dirs.index(str(install_dir))
-        assert (
-            bin_idx < install_idx
-        ), "binary_path.parent should come before install_dir"
+        assert bin_idx < install_idx, "binary_path.parent should come before install_dir"
 
     def test_linux_deduplicates_when_binary_parent_equals_install_dir(
         self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
@@ -195,17 +194,13 @@ class TestBinaryEnvCrossPlatform:
         binary_path.write_bytes(b"MZ")
 
         host = make_host(system = "Windows")
-        monkeypatch.setattr(
-            MOD, "windows_runtime_dirs_for_runtime_line", lambda _rt: []
-        )
+        monkeypatch.setattr(MOD, "windows_runtime_dirs_for_runtime_line", lambda _rt: [])
 
         env = binary_env(binary_path, install_dir, host)
         path_dirs = env["PATH"].split(os.pathsep)
         assert str(bin_dir) in path_dirs, f"build/bin/Release not in PATH: {path_dirs}"
 
-    def test_macos_sets_dyld_library_path(
-        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
-    ):
+    def test_macos_sets_dyld_library_path(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
         install_dir = tmp_path / "llama.cpp"
         install_dir.mkdir(parents = True)
         bin_dir = install_dir / "build" / "bin"
@@ -218,12 +213,8 @@ class TestBinaryEnvCrossPlatform:
 
         env = binary_env(binary_path, install_dir, host)
         dyld_parts = [p for p in env["DYLD_LIBRARY_PATH"].split(os.pathsep) if p]
-        assert (
-            str(bin_dir) in dyld_parts
-        ), f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}"
-        assert (
-            str(install_dir) in dyld_parts
-        ), f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}"
+        assert str(bin_dir) in dyld_parts, f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}"
+        assert str(install_dir) in dyld_parts, f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}"
         # binary_path.parent (build/bin) should come before install_dir
         assert dyld_parts.index(str(bin_dir)) < dyld_parts.index(str(install_dir))
 
@@ -303,7 +294,11 @@ class TestResolveRequestedLlamaTag:
     ):
         captured = {}
 
-        def fake_resolve(requested_tag, published_repo, published_release_tag = ""):
+        def fake_resolve(
+            requested_tag,
+            published_repo,
+            published_release_tag = "",
+        ):
             captured["requested_tag"] = requested_tag
             captured["published_repo"] = published_repo
             captured["published_release_tag"] = published_release_tag
@@ -350,9 +345,7 @@ class TestResolveRequestedLlamaTag:
 
 
 class TestFetchJsonRetries:
-    def test_fetch_json_retries_invalid_github_api_json(
-        self, monkeypatch: pytest.MonkeyPatch
-    ):
+    def test_fetch_json_retries_invalid_github_api_json(self, monkeypatch: pytest.MonkeyPatch):
         calls = {"count": 0}
 
         def fake_download_bytes(url, **kwargs):
@@ -593,11 +586,7 @@ class TestLatestTagResolution:
     """)
 
     def _run_resolve(
-        self,
-        tmp_path: Path,
-        requested_tag: str,
-        resolved_tag: str,
-        resolve_status: int,
+        self, tmp_path: Path, requested_tag: str, resolved_tag: str, resolve_status: int
     ) -> str:
         script = self.RESOLVE_TEMPLATE.format(
             requested_tag = requested_tag,
@@ -691,23 +680,32 @@ class TestSourceCodePatterns:
         content = SETUP_SH.read_text()
         assert "--resolve-source-build" not in content
         assert "--resolve-install-tag" not in content
-        assert (
-            '--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp"'
-            in content
-        )
+        assert '--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp"' in content
         assert "--output-format json" in content
         assert "_RESOLVED_SOURCE_URL" in content
         assert "_RESOLVED_SOURCE_REF_KIND" in content
         assert "_RESOLVED_SOURCE_REF" in content
 
-    def test_setup_sh_prebuilt_install_uses_simple_policy_only(self):
-        """Shell prebuilt path should use the simplified helper install entrypoint."""
+    def test_setup_sh_prebuilt_install_entrypoint(self):
+        """Shell prebuilt path should call the helper install entrypoint, not the
+        old tag-resolution / releases-latest flow."""
         content = SETUP_SH.read_text()
-        assert "--simple-policy" in content
         assert "--resolve-install-tag" not in content
         assert "_HELPER_RELEASE_REPO}/releases/latest" not in content
         assert "ggml-org/llama.cpp/releases/latest" not in content
 
+    def test_setup_sh_routes_to_fork_only_on_usable_gpu(self):
+        """Linux fork-vs-ggml routing must gate NVIDIA on actual GPU usability,
+        not mere nvidia-smi presence, so CPU-only / hidden-GPU hosts (e.g.
+        CUDA_VISIBLE_DEVICES=-1) get the ggml CPU prebuilt instead of a source
+        build. Guards against a silent revert to the old presence-only loop."""
+        content = SETUP_SH.read_text()
+        assert '[ "$_setup_nvidia_usable" = true ]' in content
+        assert "CUDA_VISIBLE_DEVICES" in content
+        # nvidia-smi must NOT be back in the bare presence loop.
+        assert "for _GPU_TOOL in nvidia-smi" not in content
+        assert "for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo" in content
+
     def test_setup_sh_reports_installed_prebuilt_release(self):
         """Shell wrapper should report the installed prebuilt release from metadata."""
         content = SETUP_SH.read_text()
@@ -727,16 +725,13 @@ class TestSourceCodePatterns:
         assert "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" in content
 
     def test_setup_sh_macos_metal_configure_has_cpu_fallback(self):
-        """If Metal configure or build fails, setup should retry with CPU fallback."""
+        """If Metal/CUDA/ROCm configure or build fails, setup retries a CPU
+        build. PR #5826 generalised the Metal-only wording via $_FB_LABEL; this
+        check stays label-agnostic so new GPU backends don't require edits."""
         content = SETUP_SH.read_text()
         assert "_TRY_METAL_CPU_FALLBACK=true" in content
-        assert (
-            'substep "Metal configure failed; retrying CPU build..." "$C_WARN"'
-            in content
-        )
-        assert (
-            'substep "Metal build failed; retrying CPU build..." "$C_WARN"' in content
-        )
+        assert 'configure failed; retrying CPU build..." "$C_WARN"' in content
+        assert 'build failed; retrying CPU build..." "$C_WARN"' in content
         assert 'run_quiet_no_exit "cmake llama.cpp (cpu fallback)"' in content
         assert "-DGGML_METAL=OFF" in content
         # _TRY_METAL_CPU_FALLBACK must be reset to false in both fallback branches
@@ -745,18 +740,58 @@ class TestSourceCodePatterns:
             "_TRY_METAL_CPU_FALLBACK=false should appear at least 3 times "
             "(init + configure fallback + build fallback)"
         )
+        # The fallback helper must exist and Metal must reach it via the
+        # _TRY_METAL_CPU_FALLBACK shortcut so the macOS path stays covered.
+        assert "_gpu_fallback_label()" in content
+        assert 'echo "Metal"' in content
+
+    def test_setup_sh_exports_allow_unsupported_compiler(self):
+        """Headline fix for PR #5826: a fresh CUDA toolkit's host-compiler
+        whitelist lags the distro gcc/clang, so nvcc rejects the host with
+        "#error -- unsupported GNU version". setup.sh exports
+        NVCC_PREPEND_FLAGS=-allow-unsupported-compiler (via env, not CMAKE_ARGS,
+        for word-splitting safety) so the build and compiler-id probe proceed."""
+        content = SETUP_SH.read_text()
+        assert "-allow-unsupported-compiler" in content
+        # Delivered via NVCC_PREPEND_FLAGS (covers the configure-time compiler
+        # probe too), not embedded in the word-split CMAKE_ARGS string.
+        assert "export NVCC_PREPEND_FLAGS=" in content
+        cmake_args_lines = [line for line in content.splitlines() if "CMAKE_ARGS=" in line]
+        assert all(
+            "-allow-unsupported-compiler" not in line for line in cmake_args_lines
+        ), "flag must stay out of CMAKE_ARGS (bash word-splitting safety)"
+
+    def test_setup_ps1_exports_allow_unsupported_compiler(self):
+        """Windows parity for the PR #5826 fix: a fresh CUDA toolkit's whitelist
+        also lags MSVC, so nvcc can reject the host with "#error -- unsupported
+        Microsoft Visual Studio version!". setup.ps1 sets
+        NVCC_PREPEND_FLAGS=-allow-unsupported-compiler in the CUDA branch (via
+        env, out of $CmakeArgs) so the configure probe + build proceed."""
+        content = SETUP_PS1.read_text()
+        assert "-allow-unsupported-compiler" in content
+        # Delivered via the process environment, not the $CmakeArgs array, so it
+        # reaches both the configure-time compiler probe and `cmake --build`.
+        assert "$env:NVCC_PREPEND_FLAGS" in content
+        cmake_args_lines = [line for line in content.splitlines() if "$CmakeArgs +=" in line]
+        assert all(
+            "-allow-unsupported-compiler" not in line for line in cmake_args_lines
+        ), "flag must not be pushed into the $CmakeArgs array"
+        # Must be scoped to the CUDA branch (guarded by the GPU/nvcc check),
+        # not set unconditionally for CPU-only builds.
+        flag_idx = content.index("-allow-unsupported-compiler")
+        cuda_guard_idx = content.index("if ($HasNvidiaSmi -and $NvccPath)")
+        cuda_disable_idx = content.index("'-DGGML_CUDA=OFF'")
+        assert cuda_guard_idx < flag_idx < cuda_disable_idx, (
+            "NVCC_PREPEND_FLAGS must be set inside the CUDA-on branch, "
+            "before the GGML_CUDA=OFF (CPU) branch"
+        )
 
     def test_macos_arm64_cpu_fallback_args_exclude_rpath(self):
         """CPU fallback args must NOT contain Metal-only RPATH flags at runtime."""
-        script = (
-            '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n'
-            + _GPU_BACKEND_FRAGMENT
-        )
+        script = '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT
         output = run_bash(script)
         fallback_line = next(
-            line
-            for line in output.splitlines()
-            if line.startswith("CPU_FALLBACK_CMAKE_ARGS=")
+            line for line in output.splitlines() if line.startswith("CPU_FALLBACK_CMAKE_ARGS=")
         )
         assert "-DGGML_METAL=OFF" in fallback_line
         assert (
@@ -777,8 +812,7 @@ class TestSourceCodePatterns:
         assert (
             "x86_64"
             not in content[
-                content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON")
-                + 200
+                content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON") + 200
             ]
         )
 
@@ -808,14 +842,12 @@ class TestSourceCodePatterns:
                 # Allow git pull in other contexts
                 context = "\n".join(lines[max(0, i - 5) : i + 5])
                 if "LlamaCppDir" in context:
-                    pytest.fail(
-                        f"Found 'git pull' in llama.cpp build section at line {i+1}"
-                    )
+                    pytest.fail(f"Found 'git pull' in llama.cpp build section at line {i+1}")
 
-    def test_setup_ps1_prebuilt_install_uses_simple_policy_only(self):
-        """PS1 prebuilt path should use the simplified helper install entrypoint."""
+    def test_setup_ps1_prebuilt_install_entrypoint(self):
+        """PS1 prebuilt path should call the helper install entrypoint, not the
+        old tag-resolution / releases-latest flow."""
         content = SETUP_PS1.read_text()
-        assert '"--simple-policy"' in content
         assert "--resolve-install-tag" not in content
         assert "$HelperReleaseRepo/releases/latest" not in content
         assert "ggml-org/llama.cpp/releases/latest" not in content
@@ -837,8 +869,7 @@ class TestSourceCodePatterns:
         assert "--resolve-source-build" not in content
         assert "--resolve-install-tag" not in content
         assert (
-            '"--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp"'
-            in content
+            '"--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp"' in content
         )
         assert '--output-format", "json"' in content
         assert "$ResolvedSourceUrl" in content
@@ -852,10 +883,7 @@ class TestSourceCodePatterns:
         block = content[max(0, install_idx - 800) : install_idx + 800]
         assert "$PSNativeCommandUseErrorActionPreference = $false" in block
         assert "$restoreNativeErrorPreference = $true" in block
-        assert (
-            "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference"
-            in block
-        )
+        assert "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference" in block
 
     def test_setup_ps1_helper_disables_error_action_abort(self):
         """Helper resolution should suppress terminating NativeCommandError on PS 5.1."""
@@ -876,9 +904,7 @@ class TestSourceCodePatterns:
         """The unconstrained nvcc fallback should not sort toolkit dirs lexicographically."""
         content = SETUP_PS1.read_text()
         assert "Sort-Object Name | Select-Object -Last 1" not in content
-        assert (
-            "Sort-Object { [version]($_.Name -replace '^v','') } -Descending" in content
-        )
+        assert "Sort-Object { [version]($_.Name -replace '^v','') } -Descending" in content
 
     def test_binary_env_linux_has_binary_parent(self):
         """The Linux branch of binary_env should include binary_path.parent."""
@@ -939,10 +965,7 @@ class TestMacOSMetalBuildLogic:
 
     def test_macos_arm64_cmake_args_contain_metal_flags(self):
         """macOS arm64 should enable Metal, not CUDA."""
-        script = (
-            '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n'
-            + _GPU_BACKEND_FRAGMENT
-        )
+        script = '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT
         output = run_bash(script)
         assert "-DGGML_METAL=ON" in output
         assert "-DGGML_CUDA=ON" not in output
@@ -950,10 +973,7 @@ class TestMacOSMetalBuildLogic:
 
     def test_intel_macos_no_metal_flags(self):
         """Intel macOS (not arm64) should not get Metal flags."""
-        script = (
-            '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n'
-            + _GPU_BACKEND_FRAGMENT
-        )
+        script = '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT
         output = run_bash(script)
         assert "-DGGML_METAL=ON" not in output
         assert "BUILD_DESC=building (CPU)" in output
@@ -1039,18 +1059,14 @@ class TestMacOSMetalBuildLogic:
         # Verify cmake args: first call has Metal ON, second has Metal OFF
         calls = calls_file.read_text().splitlines()
         assert len(calls) >= 2, f"Expected >= 2 cmake calls, got {len(calls)}"
-        assert (
-            "-DGGML_METAL=ON" in calls[0]
-        ), f"First cmake call should have Metal ON: {calls[0]}"
+        assert "-DGGML_METAL=ON" in calls[0], f"First cmake call should have Metal ON: {calls[0]}"
         assert (
             "-DGGML_METAL=OFF" in calls[1]
         ), f"Second cmake call should have Metal OFF: {calls[1]}"
         assert (
             "-DGGML_METAL=ON" not in calls[1]
         ), f"Second cmake call should NOT have Metal ON: {calls[1]}"
-        assert (
-            "@loader_path" not in calls[1]
-        ), f"CPU fallback should not have RPATH: {calls[1]}"
+        assert "@loader_path" not in calls[1], f"CPU fallback should not have RPATH: {calls[1]}"
         assert (
             "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[1]
         ), f"CPU fallback should not have RPATH build flag: {calls[1]}"
@@ -1158,9 +1174,7 @@ class TestMacOSMetalBuildLogic:
         # Third call: re-configure with Metal OFF and no RPATH flags
         assert "-DGGML_METAL=OFF" in calls[2]
         assert "-DGGML_METAL=ON" not in calls[2]
-        assert (
-            "@loader_path" not in calls[2]
-        ), f"CPU fallback should not have RPATH: {calls[2]}"
+        assert "@loader_path" not in calls[2], f"CPU fallback should not have RPATH: {calls[2]}"
         assert (
             "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[2]
         ), f"CPU fallback should not have RPATH build flag: {calls[2]}"

tests/studio/install/test_pr5940_followups.py

@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Tests for the AMD-Windows installer follow-ups (PR #5940):
+
+  * the huggingface_hub validation-model fetch + its urllib fallback,
+  * run_capture's Windows-only amd-smi __COMPAT_LAYER=RunAsInvoker injection,
+  * parity of the name->arch table between install.ps1 and setup.ps1.
+
+Mock-only; no AMD hardware or network required.
+"""
+
+import importlib.util
+import re
+import subprocess
+import sys
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+
+# install_llama_prebuilt.py is self-contained (stdlib + optional filelock), so it
+# loads without the studio backend on sys.path.
+_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
+_SPEC = importlib.util.spec_from_file_location(
+    "studio_install_llama_prebuilt_pr5940", _PREBUILT_PATH
+)
+assert _SPEC is not None and _SPEC.loader is not None
+prebuilt = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = prebuilt
+_SPEC.loader.exec_module(prebuilt)
+
+_INSTALL_PS1 = PACKAGE_ROOT / "install.ps1"
+_SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
+_INSTALL_SH = PACKAGE_ROOT / "install.sh"
+
+
+# ── _hf_resolve_url_parts ────────────────────────────────────────────────────
+
+
+def test_hf_resolve_url_parts_valid():
+    assert prebuilt._hf_resolve_url_parts(
+        "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
+    ) == ("ggml-org/models", "main", "tinyllamas/stories260K.gguf")
+
+
+@pytest.mark.parametrize(
+    "url",
+    [
+        "https://github.com/owner/repo/releases/download/x.gguf",  # not huggingface
+        "https://huggingface.co/owner/repo",  # no /resolve/<rev>/
+        "https://huggingface.co/owner/repo/blob/main/x.gguf",  # /blob/ not /resolve/
+        "not even a url",
+    ],
+)
+def test_hf_resolve_url_parts_non_hf_returns_none(url):
+    assert prebuilt._hf_resolve_url_parts(url) is None
+
+
+# ── _fetch_validation_model_bytes ────────────────────────────────────────────
+
+
+def test_fetch_validation_model_prefers_huggingface_hub(tmp_path):
+    model = tmp_path / "stories260K.gguf"
+    model.write_bytes(b"GGUF-via-hf")
+    fake_hf = MagicMock(return_value = str(model))
+    with (
+        patch.object(prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b),
+        patch.dict(sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}),
+    ):
+        assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-hf"
+    assert fake_hf.called  # hf path was taken, urllib not needed
+
+
+def test_fetch_validation_model_falls_back_to_urllib_on_hf_failure():
+    fake_hf = MagicMock(side_effect = RuntimeError("hf unreachable"))
+    with (
+        patch.object(prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b),
+        patch.dict(sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}),
+        patch.object(prebuilt, "download_bytes", return_value = b"GGUF-via-urllib") as dl,
+    ):
+        assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-urllib"
+    assert dl.called  # fell back to the direct URL download
+
+
+# ── run_capture amd-smi RunAsInvoker injection ───────────────────────────────
+
+
+def _capture_env(command, system):
+    captured = {"env": "sentinel"}
+
+    def fake_run(cmd, **kwargs):
+        captured["env"] = kwargs.get("env")
+        return subprocess.CompletedProcess(cmd, 0, "", "")
+
+    with (
+        patch.object(prebuilt.subprocess, "run", side_effect = fake_run),
+        patch.object(prebuilt.platform, "system", return_value = system),
+    ):
+        prebuilt.run_capture(command)
+    return captured["env"]
+
+
+def test_run_capture_injects_runasinvoker_for_amd_smi_on_windows():
+    env = _capture_env(["amd-smi", "list"], "Windows")
+    assert env is not None and env.get("__COMPAT_LAYER") == "RunAsInvoker"
+
+
+def test_run_capture_injects_for_full_path_amd_smi_exe_on_windows():
+    env = _capture_env(["amd-smi.exe", "version"], "Windows")
+    assert env is not None and env.get("__COMPAT_LAYER") == "RunAsInvoker"
+
+
+def test_run_capture_no_injection_for_non_amd_smi_on_windows():
+    assert _capture_env(["rocminfo"], "Windows") is None
+
+
+def test_run_capture_no_injection_on_linux():
+    # amd-smi does not auto-elevate on Linux, so no env override is applied.
+    assert _capture_env(["amd-smi", "list"], "Linux") is None
+
+
+# ── name->arch table parity (install.ps1 vs setup.ps1) ───────────────────────
+
+
+def _ps_name_arch_rows(text):
+    return re.findall(r'@\{\s*P\s*=\s*"([^"]*)"\s*;\s*A\s*=\s*"(gfx[0-9a-z]+)"', text)
+
+
+def test_ps_name_arch_tables_in_sync():
+    t1 = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
+    t2 = _ps_name_arch_rows(_SETUP_PS1.read_text(encoding = "utf-8"))
+    assert t1, "no nameArchTable found in install.ps1"
+    assert t1 == t2, f"name->arch tables drifted:\ninstall.ps1={t1}\nsetup.ps1={t2}"
+
+
+def test_rx_7700s_resolves_to_gfx1102_not_gfx1100():
+    rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
+    name = "AMD Radeon RX 7700S"
+    matched = next((arch for pattern, arch in rows if re.search(pattern, name)), None)
+    assert matched == "gfx1102", f"RX 7700S matched {matched!r}, expected gfx1102"
+
+
+def test_radeon_8060s_resolves_to_gfx1151():
+    rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
+    name = "AMD Radeon(TM) 8060S Graphics"
+    matched = next((arch for pattern, arch in rows if re.search(pattern, name)), None)
+    assert matched == "gfx1151"
+
+
+def _sh_name_arch_rows(text, var = "_gpu_disp_gfx"):
+    """Parse a bash `case "$..._mkt" in ... ) <var>="gfxNNNN"` name->arch
+    table into [(substr_tokens, arch), ...] preserving order."""
+    rows = []
+    for line in text.splitlines():
+        m = re.search(var + r'="(gfx[0-9a-z]+)"', line)
+        if not m or '*"' not in line:
+            continue
+        tokens = re.findall(r'\*"([^"]+)"\*', line)
+        if tokens:
+            rows.append((tokens, m.group(1)))
+    return rows
+
+
+def _sh_resolve(rows, name):
+    for tokens, arch in rows:
+        if any(tok in name for tok in tokens):  # bash *"X"* == substring
+            return arch
+    return None
+
+
+def test_install_sh_name_arch_agrees_with_ps_for_strix_and_non_amd():
+    """The bash install.sh name->arch table must agree with the PowerShell
+    source-of-truth for the Strix Halo (gfx1151) vs Strix Point (gfx1150)
+    split, and must never misclassify NVIDIA/Intel as an AMD gfx."""
+    sh_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8"))
+    ps_rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
+    assert sh_rows, "no name->arch case table found in install.sh"
+    cases = {
+        "AMD Radeon(TM) 8060S Graphics": "gfx1151",  # Strix Halo
+        "AMD Ryzen AI Max+ PRO 395 w/ Radeon 8060S": "gfx1151",
+        "AMD Radeon 890M Graphics": "gfx1150",  # Strix Point (NOT gfx1151)
+        "AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150",
+        "AMD Radeon RX 7700S": "gfx1102",
+        "NVIDIA GeForce RTX 4090": None,
+        "Intel(R) Arc A770 Graphics": None,
+    }
+    for name, expect in cases.items():
+        sh = _sh_resolve(sh_rows, name)
+        assert sh == expect, f"install.sh: {name!r} -> {sh!r}, expected {expect!r}"
+        if expect is not None:  # cross-check bash agrees with the PowerShell table
+            ps = next((a for p, a in ps_rows if re.search(p, name)), None)
+            assert sh == ps, f"install.sh/install.ps1 drift for {name!r}: {sh!r} vs {ps!r}"
+
+
+def test_setup_sh_name_arch_table_in_sync_with_install_sh():
+    """studio/setup.sh keeps its own copy of the bash name->arch table (over
+    `_setup_gfx`); it must stay row-for-row identical to install.sh's, both in
+    tokens and in match order (order carries the RX 7700S -> gfx1102 rule)."""
+    install_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8"))
+    setup_rows = _sh_name_arch_rows(
+        (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8"),
+        var = "_setup_gfx",
+    )
+    assert setup_rows, "no name->arch case table found in studio/setup.sh"
+    assert install_rows == setup_rows, (
+        "bash name->arch tables drifted:\n"
+        f"install.sh={install_rows}\nstudio/setup.sh={setup_rows}"
+    )
+    # The historical drift this guards against: Strix Point SKUs must be
+    # gfx1150, and the spaceless RX 7700S must match gfx1102 before gfx1100.
+    for name, expect in {
+        "AMD Radeon 890M Graphics": "gfx1150",
+        "AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150",
+        "AMD Radeon(TM) 8060S Graphics": "gfx1151",
+        "AMD Radeon RX 7700S": "gfx1102",
+    }.items():
+        got = _sh_resolve(setup_rows, name)
+        assert got == expect, f"setup.sh: {name!r} -> {got!r}, expected {expect!r}"
+
+
+# ── amd-smi gating (DiskPart UAC-prompt avoidance) ───────────────────────────
+# On Windows w/o a HIP SDK, amd-smi elevates and pops a UAC/DiskPart prompt
+# RunAsInvoker can't suppress, so _amd_smi_allowed() skips it by default;
+# HIP-SDK hosts and an explicit opt-in keep it.
+
+
+def _amd_smi_allowed_under(system, hipinfo_present, env):
+    which = (
+        (lambda name: r"C:\hip\bin\hipinfo.exe" if name == "hipinfo" else None)
+        if hipinfo_present
+        else (lambda name: None)
+    )
+    with (
+        patch.object(prebuilt.platform, "system", return_value = system),
+        patch.object(prebuilt.shutil, "which", side_effect = which),
+        patch.dict(prebuilt.os.environ, env, clear = True),
+    ):
+        return prebuilt._amd_smi_allowed()
+
+
+def test_amd_smi_allowed_on_linux_regardless():
+    # Linux amd-smi does not elevate -> always allowed (no regression on Linux).
+    assert _amd_smi_allowed_under("Linux", hipinfo_present = False, env = {}) is True
+
+
+def test_amd_smi_skipped_on_windows_without_hip_sdk():
+    # The DiskPart fix: no HIP SDK + no opt-in -> do not spawn amd-smi.
+    assert _amd_smi_allowed_under("Windows", hipinfo_present = False, env = {}) is False
+
+
+def test_amd_smi_allowed_on_windows_with_hip_sdk():
+    # hipinfo present => amd-smi runs un-elevated, so it is allowed (no regression
+    # for HIP-SDK Windows users, who never saw the prompt).
+    assert _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {}) is True
+
+
+def test_amd_smi_opt_in_forces_on_windows_no_sdk():
+    assert (
+        _amd_smi_allowed_under(
+            "Windows", hipinfo_present = False, env = {"UNSLOTH_ENABLE_AMD_SMI": "1"}
+        )
+        is True
+    )
+
+
+def test_amd_smi_opt_out_overrides_hip_sdk():
+    assert (
+        _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {"UNSLOTH_ENABLE_AMD_SMI": "0"})
+        is False
+    )
+
+
+def test_ps_installers_gate_amd_smi_on_windows():
+    # Both PowerShell installers must gate amd-smi behind HIP-SDK presence + the
+    # UNSLOTH_ENABLE_AMD_SMI opt-in, mirroring _amd_smi_allowed().
+    for ps in (_INSTALL_PS1, _SETUP_PS1):
+        text = ps.read_text(encoding = "utf-8")
+        assert "UNSLOTH_ENABLE_AMD_SMI" in text, f"{ps.name} missing amd-smi opt-in gate"
+        assert "amdSmiAllowed" in text, f"{ps.name} missing amd-smi gate variable"
+
+
+def test_install_python_stack_gates_every_amd_smi_spawn():
+    # Regression for the DiskPart UAC prompt: every function that both names the
+    # `amd-smi` command AND spawns a subprocess must gate it behind
+    # _amd_smi_allowed(). The "ROCm torch missing" probe once spawned `amd-smi
+    # list` ungated on Adrenalin-only hosts; not-spawning is the only fix.
+    import ast
+
+    src = (PACKAGE_ROOT / "studio" / "install_python_stack.py").read_text(encoding = "utf-8")
+    tree = ast.parse(src)
+
+    def _names_amd_smi_command(node):
+        # Exact "amd-smi"/"amd-smi.exe" constant, not a substring in a log message.
+        return any(
+            isinstance(n, ast.Constant)
+            and isinstance(n.value, str)
+            and n.value.lower() in ("amd-smi", "amd-smi.exe")
+            for n in ast.walk(node)
+        )
+
+    def _spawns_subprocess(node):
+        for n in ast.walk(node):
+            if (
+                isinstance(n, ast.Call)
+                and isinstance(n.func, ast.Attribute)
+                and isinstance(n.func.value, ast.Name)
+                and n.func.value.id == "subprocess"
+            ):
+                return True
+        return False
+
+    def _references_gate(node):
+        return any(isinstance(n, ast.Name) and n.id == "_amd_smi_allowed" for n in ast.walk(node))
+
+    offenders = [
+        node.name
+        for node in ast.walk(tree)
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+        and _names_amd_smi_command(node)
+        and _spawns_subprocess(node)
+        and not _references_gate(node)
+    ]
+    assert not offenders, (
+        "install_python_stack.py spawns amd-smi without an _amd_smi_allowed() "
+        f"gate in: {offenders} -- this pops the Windows UAC/DiskPart prompt on "
+        "Adrenalin-only (no HIP SDK) hosts."
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/studio/install/test_probe_timeouts.py

@@ -0,0 +1,202 @@
+"""Tests that NVIDIA probes in the installers are bounded by a timeout.
+
+Covers audit findings 5 and 6: a wedged nvidia-smi must not hang the installer,
+and the Windows probe must require a real GPU listing (not just exit code 0).
+
+Source-level assertions verify the guards are present in install.sh / install.ps1
+/ setup.ps1; one behavioral shell test confirms the bash helper actually returns
+within the timeout when nvidia-smi hangs.
+"""
+
+import os
+import shutil
+import stat
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import pytest
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+INSTALL_SH = PACKAGE_ROOT / "install.sh"
+INSTALL_PS1 = PACKAGE_ROOT / "install.ps1"
+SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
+
+
+def _extract_sh_function_body(source: str, name: str) -> str:
+    """Return a shell function body from `source` by brace matching."""
+    needle = f"{name}() {{"
+    start = source.find(needle)
+    if start < 0:
+        return ""
+    depth = 0
+    i = start + len(needle) - 1
+    n = len(source)
+    while i < n:
+        ch = source[i]
+        if ch == "{":
+            depth += 1
+        elif ch == "}":
+            depth -= 1
+            if depth == 0:
+                return source[start : i + 1]
+        i += 1
+    return source[start:]
+
+
+# ── install.sh: _run_bounded helper and its use at every nvidia-smi call ──
+
+
+class TestInstallShBoundedProbe:
+    def _src(self) -> str:
+        return INSTALL_SH.read_text(encoding = "utf-8")
+
+    def test_run_bounded_helper_defined(self):
+        body = _extract_sh_function_body(self._src(), "_run_bounded")
+        assert body, "install.sh must define a _run_bounded helper"
+        assert (
+            "command -v timeout" in body
+        ), "_run_bounded must check for the `timeout` binary before using it"
+        assert "timeout 10" in body, "_run_bounded must apply a 10s timeout"
+        # Must fall back to running unbounded when `timeout` is unavailable
+        # (e.g. macOS) so semantics are unchanged there.
+        assert (
+            "else" in body and '"$@"' in body
+        ), "_run_bounded must run the command unbounded when `timeout` is absent"
+
+    def test_nvidia_smi_dash_l_probe_is_bounded(self):
+        body = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
+        assert body, "install.sh must define _has_usable_nvidia_gpu"
+        # The -L probe must go through the bounded runner, not call nvidia-smi raw.
+        assert (
+            '_run_bounded "$_nvsmi" -L' in body
+        ), "_has_usable_nvidia_gpu must run nvidia-smi -L through _run_bounded"
+        # The /proc fallback from PR 6174 must still be present.
+        assert "/proc/driver/nvidia" in body
+
+    def test_cuda_version_parse_is_bounded(self):
+        body = _extract_sh_function_body(self._src(), "get_torch_index_url")
+        assert body, "install.sh must define get_torch_index_url"
+        assert (
+            "_run_bounded" in body
+        ), "get_torch_index_url CUDA-version parse must run nvidia-smi through _run_bounded"
+        # The locale must be forced without depending on `env` being on PATH.
+        assert "LC_ALL=C" in body
+        # _nvidia_detected gating from PR 6174 must remain.
+        assert "_nvidia_detected" in body
+
+    def test_no_unbounded_nvidia_smi_invocation_remains(self):
+        """Every nvidia-smi *execution* in install.sh goes through _run_bounded.
+
+        `command -v nvidia-smi` and `-x /usr/bin/nvidia-smi` are resolution
+        checks, not executions, and are allowed. An execution looks like
+        `"$_nvsmi" ...` / `$_smi ...` / `nvidia-smi -L`.
+        """
+        body_nvidia = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
+        body_torch = _extract_sh_function_body(self._src(), "get_torch_index_url")
+        # In _has_usable_nvidia_gpu the only execution of $_nvsmi must be bounded.
+        assert '"$_nvsmi" -L' not in body_nvidia.replace(
+            '_run_bounded "$_nvsmi" -L', ""
+        ), "found an unbounded nvidia-smi -L execution in _has_usable_nvidia_gpu"
+        # In get_torch_index_url the $_smi execution must be bounded.
+        assert (
+            "LC_ALL=C $_smi" not in body_torch
+        ), "found an unbounded LC_ALL=C $_smi execution in get_torch_index_url"
+
+
+# ── install.ps1 / setup.ps1: bounded, GPU-row-validated Windows probe ──
+
+
+class TestPowerShellBoundedProbe:
+    @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
+    def test_bounded_helper_present(self, path):
+        src = path.read_text(encoding = "utf-8")
+        assert (
+            "function Invoke-NvidiaSmiBounded" in src
+        ), f"{path.name} must define Invoke-NvidiaSmiBounded"
+        assert (
+            "WaitForExit($TimeoutSec * 1000)" in src
+        ), f"{path.name} bounded probe must use WaitForExit with a timeout"
+        # Kill + sentinel on timeout, mirroring Invoke-AmdSmiNoElevate.
+        assert (
+            "$proc.Kill()" in src and "124" in src
+        ), f"{path.name} must kill nvidia-smi and signal a timeout exit code"
+
+    @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
+    def test_probe_requires_gpu_row(self, path):
+        src = path.read_text(encoding = "utf-8")
+        assert (
+            "function Test-NvidiaSmiHasGpu" in src
+        ), f"{path.name} must define Test-NvidiaSmiHasGpu"
+        assert "@('-L')" in src, f"{path.name} must probe nvidia-smi with -L"
+        assert (
+            "^GPU\\s+\\d+:" in src
+        ), f"{path.name} must require a 'GPU <n>:' data row, not just exit code 0"
+
+    @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
+    def test_detection_uses_validated_probe(self, path):
+        src = path.read_text(encoding = "utf-8")
+        # The exit-code-only pattern must be gone from the detection block.
+        assert (
+            "& $nvSmiCmd.Source *> $null" not in src
+        ), f"{path.name} must not use the exit-code-only nvidia-smi probe"
+        assert (
+            "Test-NvidiaSmiHasGpu $nvSmiCmd.Source" in src
+        ), f"{path.name} PATH probe must use Test-NvidiaSmiHasGpu"
+        assert (
+            "Test-NvidiaSmiHasGpu $p" in src
+        ), f"{path.name} hardcoded-path fallback must use Test-NvidiaSmiHasGpu"
+
+
+# ── Behavioral: a hanging nvidia-smi must not hang _has_usable_nvidia_gpu ──
+
+
+def _have_timeout() -> bool:
+    return shutil.which("timeout") is not None
+
+
+@pytest.mark.skipif(not _have_timeout(), reason = "`timeout` binary not available")
+def test_has_usable_nvidia_gpu_returns_under_timeout():
+    """Extract _run_bounded + _has_usable_nvidia_gpu, point them at a fake
+    nvidia-smi that sleeps 30s, and assert the probe returns well under that.
+    """
+    src = INSTALL_SH.read_text(encoding = "utf-8")
+    helper = _extract_sh_function_body(src, "_run_bounded")
+    fn = _extract_sh_function_body(src, "_has_usable_nvidia_gpu")
+    assert helper and fn
+
+    workdir = tempfile.mkdtemp(prefix = "pr6174_timeout_", dir = str(PACKAGE_ROOT.parent))
+    try:
+        fake_dir = Path(workdir, "bin")
+        fake_dir.mkdir()
+        fake_smi = fake_dir / "nvidia-smi"
+        fake_smi.write_text("#!/bin/sh\nsleep 30\n")
+        fake_smi.chmod(fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
+
+        # Build a minimal PATH that includes the fake nvidia-smi plus the real
+        # `timeout`/`awk`/`ls` it needs. Use the fake dir first so it wins.
+        real_bins = {Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh")}
+        path_env = os.pathsep.join([str(fake_dir)] + [str(p) for p in real_bins])
+
+        # Force the /proc fallback off so the result depends only on the probe,
+        # and so a host with real NVIDIA does not mask the timeout behaviour.
+        script = (
+            f"{helper}\n{fn}\n"
+            "if _has_usable_nvidia_gpu; then echo DETECTED; else echo NONE; fi\n"
+        )
+        proc = subprocess.run(
+            ["sh", "-c", script],
+            env = {"PATH": path_env},
+            stdout = subprocess.PIPE,
+            stderr = subprocess.DEVNULL,
+            text = True,
+            timeout = 20,  # generous: the internal timeout is 10s, sleep is 30s
+        )
+        # The probe must have returned (not hung). On this CI host /proc/driver/
+        # nvidia/gpus is absent, so a timed-out smi yields NONE; on a real NVIDIA
+        # host the /proc fallback yields DETECTED. Either way it must not hang.
+        assert proc.stdout.strip() in {"NONE", "DETECTED"}
+    finally:
+        shutil.rmtree(workdir, ignore_errors = True)

tests/studio/load_freeze/llama_server_shim.py

@@ -41,7 +41,11 @@ class _Handler(BaseHTTPRequestHandler):
         self.wfile.write(payload)
 
     def _send_raw(
-        self, status: int, body: bytes, *, content_type: str = "application/json"
+        self,
+        status: int,
+        body: bytes,
+        *,
+        content_type: str = "application/json",
     ) -> None:
         self.send_response(status)
         self.send_header("Content-Type", content_type)
@@ -119,9 +123,7 @@ class _Handler(BaseHTTPRequestHandler):
                 self._send_raw(srv.config.detok_status, srv.config.detok_body)
                 return
             tids = body.get("tokens") or []
-            content = "".join(
-                srv.config.detok_map.get(int(t), f"<tok_{t}>") for t in tids
-            )
+            content = "".join(srv.config.detok_map.get(int(t), f"<tok_{t}>") for t in tids)
             self._send_json(srv.config.detok_status, {"content": content})
             return
         if path == "/completion":
@@ -222,11 +224,9 @@ class FakeLlamaServer:
         self._thread: Optional[threading.Thread] = None
 
     def start(self) -> "FakeLlamaServer":
-        # port=0 lets ThreadingHTTPServer pick a free port atomically
-        # (avoids find-port-then-bind race); read back via server_address[1].
-        self._server = FakeLlamaServer._Server(
-            (self.host, self._requested_port), _Handler
-        )
+        # port=0 lets ThreadingHTTPServer pick a free port atomically (no
+        # find-then-bind race); read back via server_address[1].
+        self._server = FakeLlamaServer._Server((self.host, self._requested_port), _Handler)
         self._server.config = self.config
         bound_port = self._server.server_address[1]
         self._thread = threading.Thread(

tests/studio/load_freeze/test_load_orchestrator.py

@@ -3,7 +3,7 @@
 Covers:
   1.  Behavioural canary (the bug class)             — 2 tests
   2.  Behavioural fix-validation                     — 1 test
-  3.  Functional equivalence (sync == to_thread)     — 5 tests, one per codec branch
+  3.  Functional equivalence (sync == to_thread)     — 6 tests, one per codec branch
   4.  Failure modes (HTTP 500, malformed JSON,
       connection reset, unreachable, not-loaded)     — 5 tests
   5.  Stress (50 concurrent probes / 100 healths)    — 2 tests
@@ -33,9 +33,7 @@ from pathlib import Path
 import pytest
 
 
-# ---------------------------------------------------------------------------
 # Repo discovery
-# ---------------------------------------------------------------------------
 
 
 def _find_repo_root() -> Path | None:
@@ -79,9 +77,7 @@ from core.inference.llama_cpp import LlamaCppBackend  # noqa: E402
 from llama_server_shim import FakeLlamaServer  # noqa: E402
 
 
-# ---------------------------------------------------------------------------
 # Fixtures / helpers
-# ---------------------------------------------------------------------------
 
 
 def _make_backend(port: int, *, loaded: bool = True) -> LlamaCppBackend:
@@ -103,14 +99,18 @@ def _free_port() -> int:
 
 
 class _UvicornServerThread:
-    def __init__(self, app, *, host: str = "127.0.0.1", port: int) -> None:
+    def __init__(
+        self,
+        app,
+        *,
+        host: str = "127.0.0.1",
+        port: int,
+    ) -> None:
         import uvicorn
 
         self.host = host
         self.port = port
-        cfg = uvicorn.Config(
-            app, host = host, port = port, log_level = "warning", access_log = False
-        )
+        cfg = uvicorn.Config(app, host = host, port = port, log_level = "warning", access_log = False)
         self._server = uvicorn.Server(cfg)
         self._server.install_signal_handlers = lambda: None  # type: ignore[assignment]
         self._thread: threading.Thread | None = None
@@ -169,7 +169,12 @@ def _build_app(backend, *, wrap_in_thread: bool):
     return app
 
 
-def _drive_concurrent_probe_and_health(base_url, *, n_health = 12, gap = 0.05):
+def _drive_concurrent_probe_and_health(
+    base_url,
+    *,
+    n_health = 12,
+    gap = 0.05,
+):
     elapsed = -1.0
     latencies: list[float] = []
 
@@ -199,9 +204,7 @@ def _drive_concurrent_probe_and_health(base_url, *, n_health = 12, gap = 0.05):
     return max(latencies), elapsed, latencies
 
 
-# ---------------------------------------------------------------------------
 # (1) Behavioural canary
-# ---------------------------------------------------------------------------
 
 
 def test_buggy_route_blocks_event_loop():
@@ -211,9 +214,7 @@ def test_buggy_route_blocks_event_loop():
         app = _build_app(backend, wrap_in_thread = False)
         port = _free_port()
         with _UvicornServerThread(app, port = port) as uv:
-            max_lat, probe_t, _ = _drive_concurrent_probe_and_health(
-                f"http://127.0.0.1:{uv.port}"
-            )
+            max_lat, probe_t, _ = _drive_concurrent_probe_and_health(f"http://127.0.0.1:{uv.port}")
     assert probe_t >= 0.5
     assert max_lat >= 0.4, f"expected >=0.4s stall, got {max_lat:.3f}s"
 
@@ -232,9 +233,7 @@ def test_fixed_route_keeps_event_loop_responsive():
     assert max_lat < 0.25, f"expected <0.25s; got {max_lat:.3f}s (all: {lats})"
 
 
-# ---------------------------------------------------------------------------
 # (2) Functional equivalence -- sync == to_thread for each codec branch
-# ---------------------------------------------------------------------------
 
 
 @pytest.fixture
@@ -254,6 +253,7 @@ def shim_no_match():
             "<|audio_eos|>": [0, 1],
             "<|startoftranscript|>": [0, 1],
             "<audio_soft_token>": [0, 1],
+            "<|audio|>": [0, 1],
             "<|bicodec_semantic_0|>": [0, 1],
             "<|bicodec_global_0|>": [0, 1],
             "<|c1_0|>": [0, 1],
@@ -314,6 +314,28 @@ def test_functional_equivalence_whisper_match():
     assert sync_result == threaded
 
 
+def test_functional_equivalence_audio_vlm_match():
+    # audio_vlm: snac/csm/whisper fail first, then the Gemma 4 <|audio|>
+    # probe tokenises to a single token. #6000 added this arm alongside
+    # Gemma 3n's <audio_soft_token>; keep <audio_soft_token> at 2 tokens so
+    # it is specifically the new <|audio|> arm that triggers the match.
+    with FakeLlamaServer(
+        detok_map = {128258: "non-snac", 128259: "non-snac"},
+        tok_response_map = {
+            "<|AUDIO|>": [0, 1],  # csm fails (>1 token)
+            "<|audio_eos|>": [0, 1],
+            "<|startoftranscript|>": [0, 1],  # whisper fails
+            "<audio_soft_token>": [0, 1],  # Gemma 3n arm fails ...
+            "<|audio|>": [0],  # ... Gemma 4 arm matches (#6000)
+        },
+    ) as srv:
+        backend = _make_backend(srv.port)
+        sync_result = backend.detect_audio_type()
+        threaded = asyncio.run(asyncio.to_thread(backend.detect_audio_type))
+    assert sync_result == "audio_vlm"
+    assert sync_result == threaded
+
+
 def test_functional_equivalence_bicodec_match():
     # bicodec: snac/csm/whisper/audio_vlm all fail first, then both
     # bicodec_semantic_0 and bicodec_global_0 are single tokens.
@@ -324,6 +346,7 @@ def test_functional_equivalence_bicodec_match():
             "<|audio_eos|>": [0, 1],
             "<|startoftranscript|>": [0, 1],
             "<audio_soft_token>": [0, 1],
+            "<|audio|>": [0, 1],
             "<|bicodec_semantic_0|>": [0],
             "<|bicodec_global_0|>": [0],
         },
@@ -335,9 +358,7 @@ def test_functional_equivalence_bicodec_match():
     assert sync_result == threaded
 
 
-# ---------------------------------------------------------------------------
 # (3) Failure modes
-# ---------------------------------------------------------------------------
 
 
 def test_shim_returns_500_on_tokenize_returns_none():
@@ -401,9 +422,7 @@ def test_backend_not_loaded_short_circuits():
     assert threaded_t < 0.05
 
 
-# ---------------------------------------------------------------------------
 # (4) Stress / concurrency
-# ---------------------------------------------------------------------------
 
 
 def test_50_concurrent_probes_complete_without_deadlock():
@@ -418,9 +437,7 @@ def test_50_concurrent_probes_complete_without_deadlock():
             with ThreadPoolExecutor(max_workers = 50) as pool:
                 futs = [
                     pool.submit(
-                        lambda: httpx.get(
-                            f"http://127.0.0.1:{uv.port}/probe", timeout = 30.0
-                        )
+                        lambda: httpx.get(f"http://127.0.0.1:{uv.port}/probe", timeout = 30.0)
                     )
                     for _ in range(50)
                 ]
@@ -470,9 +487,7 @@ def test_100_concurrent_healths_during_slow_probe_all_responsive():
     assert max_lat < 0.35, f"100-burst max latency {max_lat:.3f}s exceeds 350 ms"
 
 
-# ---------------------------------------------------------------------------
 # (5) Drift / regression guards on the production source
-# ---------------------------------------------------------------------------
 
 
 def test_load_model_caches_audio_type_inside_serial_load_lock():
@@ -559,9 +574,7 @@ def test_no_other_async_route_calls_detect_audio_type_unwrapped():
     )
 
 
-# ---------------------------------------------------------------------------
 # (6) Timing budgets
-# ---------------------------------------------------------------------------
 
 
 def test_load_response_under_2s_with_fast_shim():
@@ -595,9 +608,7 @@ def test_repeated_loads_bounded_total_time():
     assert elapsed < 10.0
 
 
-# ---------------------------------------------------------------------------
 # (7) Browser-compatibility surface
-# ---------------------------------------------------------------------------
 
 
 def test_response_is_valid_browser_parseable_json():
@@ -636,7 +647,6 @@ def test_response_shape_matches_pre_fix_for_no_match():
     bodies for the no-match scenario (the dominant code path in
     practice for non-audio models)."""
     import json as _json
-
     with FakeLlamaServer(
         detok_map = {128258: "abc", 128259: "def"},
         tok_response_map = {
@@ -644,6 +654,7 @@ def test_response_shape_matches_pre_fix_for_no_match():
             "<|audio_eos|>": [0, 1],
             "<|startoftranscript|>": [0, 1],
             "<audio_soft_token>": [0, 1],
+            "<|audio|>": [0, 1],
             "<|bicodec_semantic_0|>": [0, 1],
             "<|bicodec_global_0|>": [0, 1],
             "<|c1_0|>": [0, 1],
@@ -663,9 +674,7 @@ def test_response_shape_matches_pre_fix_for_no_match():
             assert body == {"audio_type": None}
 
 
-# ---------------------------------------------------------------------------
 # (8) Cancellation
-# ---------------------------------------------------------------------------
 
 
 def test_client_disconnect_during_probe_does_not_crash_server():

tests/studio/playwright_chat_ime_i18n.py

@@ -4,20 +4,15 @@
 """Studio chat composer IME + multilingual regression smoke.
 
 Covers three surfaces:
-  A. Stuck IME composition (issue #5318 / PR #5327): duplicate
-     compositionstart with no compositionend left isComposing=true,
-     dropping all subsequent keystrokes including ASCII.
-  B. Multilingual paste round-trip across 31 scripts -- guards the
-     controlled-textarea / React state plumbing against Unicode mangling.
-  C. Stuck compositionend (issue #5546): Chrome on Windows over WSL
-     fires compositionstart + compositionupdate but never compositionend,
-     wedging Send disabled after the IME commits. Verifies the
-     watchdog in useImeComposerInputHandlers releases the flag.
+  A. Stuck IME composition (#5318 / PR #5327): duplicate compositionstart with
+     no compositionend left isComposing=true, dropping keystrokes.
+  B. Multilingual paste round-trip across 31 scripts (Unicode plumbing).
+  C. Stuck compositionend (#5546): WSL Chrome never fires compositionend,
+     wedging Send disabled; the useImeComposerInputHandlers watchdog releases it.
 
 Model-free; the bug surface is the composer, not inference.
-
-Env contract matches playwright_chat_ui.py:
-  BASE_URL, STUDIO_NEW_PW, PW_ART_DIR, STUDIO_UI_STRICT.
+Env contract matches playwright_chat_ui.py: BASE_URL, STUDIO_NEW_PW, PW_ART_DIR,
+STUDIO_UI_STRICT.
 """
 
 import os
@@ -49,8 +44,7 @@ STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
 WALL_TIMEOUT_S = float(os.environ.get("STUDIO_IME_WALL_TIMEOUT_S", "300"))
 
 
-# One short greeting + arithmetic per script (ordered by speaker count) --
-# each entry catches a distinct class of Unicode regression.
+# One greeting + arithmetic per script; each catches a distinct Unicode class.
 I18N_SAMPLES = [
     ("en", "English", "Hello, 1+1=2"),
     ("zh-CN", "Chinese (Simplified)", "你好，1+1=2"),
@@ -156,8 +150,8 @@ with sync_playwright() as p:
         except Exception as _shoot_err:
             info(f"WARN: screenshot {name} failed: {_shoot_err}")
 
-    # 1. Bootstrap auth via /change-password (mirrors playwright_chat_ui.py
-    #    retry-on-rerender to absorb React form-detach races).
+    # 1. Bootstrap auth via /change-password (retry-on-rerender absorbs React
+    #    form-detach races, mirroring playwright_chat_ui.py).
     step("change-password through UI (Setup your account)")
     form_err: Exception | None = None
     for _form_attempt in range(3):
@@ -205,7 +199,7 @@ with sync_playwright() as p:
     if form_err is not None:
         raise form_err
 
-    # 2. Wait for composer mount. No GGUF: the bug surface is React state, not inference.
+    # 2. Wait for composer mount (no GGUF: the bug is React state, not inference).
     step("wait for composer to mount")
     try:
         page.wait_for_load_state("networkidle", timeout = 30_000)
@@ -246,21 +240,18 @@ with sync_playwright() as p:
     dir_attr = composer.evaluate("(el) => el.getAttribute('dir')")
     if dir_attr != "auto":
         soft_fail(
-            f'composer is missing dir="auto" (got {dir_attr!r}); RTL '
-            "languages will render LTR."
+            f'composer is missing dir="auto" (got {dir_attr!r}); RTL ' "languages will render LTR."
         )
     else:
         info('composer dir="auto" present')
 
-    # Source-level guard for the edit and compare composers (neither
-    # is mounted here): grep the JSX for dir="auto" inside each block.
+    # Source-level guard for the unmounted edit/compare composers: grep their
+    # JSX for dir="auto".
     _repo_root = Path(__file__).resolve().parents[2]
     _thread_src = (
         _repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx"
     ).read_text()
-    _shared_src = (
-        _repo_root / "studio/frontend/src/features/chat/shared-composer.tsx"
-    ).read_text()
+    _shared_src = (_repo_root / "studio/frontend/src/features/chat/shared-composer.tsx").read_text()
     _edit_idx = _thread_src.find("aui-edit-composer-input")
     if _edit_idx == -1 or 'dir="auto"' not in _thread_src[_edit_idx : _edit_idx + 600]:
         soft_fail('edit composer source is missing dir="auto"')
@@ -269,8 +260,7 @@ with sync_playwright() as p:
     _compare_idx = _shared_src.find("Send to both models")
     if (
         _compare_idx == -1
-        or 'dir="auto"'
-        not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400]
+        or 'dir="auto"' not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400]
     ):
         soft_fail('compare composer source is missing dir="auto"')
     else:
@@ -280,9 +270,9 @@ with sync_playwright() as p:
         return composer.evaluate("(el) => el.value")
 
     def set_value_via_setter(s: str) -> str:
-        """Write via React's monkey-patched setter + paste input event,
-        then await two rAFs so the controlled value is committed before
-        readback (plain `.value=s` would be overwritten on next render)."""
+        """Write via React's setter + paste input event, then await two rAFs so the
+        controlled value commits before readback (plain `.value=s` is overwritten
+        on next render)."""
         return composer.evaluate(
             """async (el, v) => {
                 const setter = Object.getOwnPropertyDescriptor(
@@ -373,9 +363,9 @@ with sync_playwright() as p:
     shoot("05-normal-composition")
     clear()
 
-    # 6. Stuck IME repro for issue #5318: duplicate compositionstart with
-    #    no compositionend wedged isComposing=true and dropped ASCII keys.
-    #    PR #5327 cleared the stale state on non-composing input.
+    # 6. Stuck IME repro (#5318): duplicate compositionstart with no
+    #    compositionend wedged isComposing=true; PR #5327 clears it on
+    #    non-composing input.
     step("BUG REPRO: stuck IME composition recovery (issue #5318)")
     clear()
     composer.click()
@@ -389,9 +379,8 @@ with sync_playwright() as p:
             el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
         }"""
     )
-    # Drive the real keyboard path; on the broken build React drops
-    # 'abcd' and reconciles el.value back to ''. wait_for_function
-    # crosses the microtask boundary so we see committed React state.
+    # On the broken build React drops 'abcd' and reconciles el.value to ''.
+    # wait_for_function crosses the microtask boundary to see committed state.
     page.keyboard.type("abcd")
     try:
         page.wait_for_function(
@@ -428,12 +417,9 @@ with sync_playwright() as p:
     info("stuck-composition recovery PASS")
     clear()
 
-    # 6b. WSL + Windows Chrome repro for issue #5546: Chrome never emits
-    #     compositionend after the IME commit, so the watchdog has to
-    #     release the composing flag on its own once the events go silent.
-    #     This dispatches a realistic "compose, commit, then nothing"
-    #     sequence — no compositionend, no follow-up keystrokes — and
-    #     waits for the Send button to come back enabled.
+    # 6b. WSL+Chrome repro (#5546): Chrome never emits compositionend after the
+    #     IME commit, so the watchdog must release the composing flag once events
+    #     go silent. Dispatch "compose, commit, then nothing" and wait for Send.
     step("BUG REPRO: stuck compositionend recovery (issue #5546)")
     clear()
     composer.click()
@@ -477,16 +463,11 @@ with sync_playwright() as p:
     info("compositionend watchdog recovery PASS")
     clear()
 
-    # 6c. Watchdog-race repro: after the watchdog clears composingRef during a
-    #     long candidate pause, a subsequent IME keydown (browser still sees
-    #     isComposing=true / keyCode 229) must not slip preedit text through
-    #     the form submit. The onKeyDown gate re-pins composingRef so the
-    #     handleSubmit / blockSend guards keep refusing. The Send button stays
-    #     visually enabled (watchdog has already cleared the React state); the
-    #     refusal happens at form.requestSubmit() time, not at the button.
-    step(
-        "BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)"
-    )
+    # 6c. Watchdog-race repro: after the watchdog clears composingRef, a later
+    #     IME keydown (isComposing=true / keyCode 229) must not slip preedit text
+    #     through submit. The onKeyDown gate re-pins composingRef so handleSubmit
+    #     refuses at form.requestSubmit() time, not at the (enabled) button.
+    step("BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)")
     clear()
     composer.click()
     composer.evaluate(
@@ -510,10 +491,9 @@ with sync_playwright() as p:
         expect(send_btn_keydown).not_to_be_disabled(timeout = 8_000)
     except Exception:
         soft_fail("watchdog did not clear before keydown re-pin test")
-    # Fire the IME-confirm Enter (keyCode 229, isComposing=true) then trigger
-    # the form submit synchronously. With the keydown gate, composingRef is
-    # re-pinned before handleSubmit runs and the submit is prevented; the
-    # textarea must still hold the preedit text.
+    # Fire the IME-confirm Enter (keyCode 229) then submit synchronously. The
+    # keydown gate re-pins composingRef before handleSubmit, preventing submit;
+    # the textarea must still hold the preedit text.
     submit_probe = composer.evaluate(
         """(el) => {
             el.focus();
@@ -533,19 +513,15 @@ with sync_playwright() as p:
             "Form submitted after an IME keydown -- preedit text leaked "
             "through the watchdog gap (#5546 follow-up regression)."
         )
-    info(
-        f"Form submit refused after IME keydown; textarea retained {submit_probe.get('after')!r}"
-    )
+    info(f"Form submit refused after IME keydown; textarea retained {submit_probe.get('after')!r}")
     shoot("06c-keydown-repin")
     info("keydown re-pin gate PASS")
     clear()
 
     # 6d. Keydown re-pin must also re-arm the watchdog. On the WSL+Chrome
-    #     stuck-compositionend path the IME never fires a follow-up
-    #     compositionend or non-composing input, so after the IME keydown
-    #     re-pins composingRef the watchdog has to take it back to false on
-    #     its own — otherwise Send re-locks permanently after the very
-    #     scenario this PR was supposed to fix. (Codex P1, commit 597af0d0.)
+    #     stuck-compositionend path no follow-up event arrives, so after keydown
+    #     re-pins composingRef the watchdog must clear it again or Send re-locks
+    #     permanently. (Codex P1, commit 597af0d0.)
     step("BUG REPRO: keydown re-pin re-arms watchdog (#5546 follow-up regression)")
     clear()
     composer.click()
@@ -571,8 +547,7 @@ with sync_playwright() as p:
     except Exception:
         soft_fail("watchdog did not clear before re-arm test (first cycle)")
     # IME-confirm keydown re-pins composingRef. Without the re-arm fix the
-    # watchdog would never run again and Send would stay blocked at the
-    # submit-time guard forever, even though no follow-up IME event arrives.
+    # watchdog never runs again and Send stays blocked forever.
     composer.evaluate(
         """(el) => {
             el.focus();
@@ -582,9 +557,9 @@ with sync_playwright() as p:
             }));
         }"""
     )
-    # Second watchdog cycle: a real submit attempt now must eventually be
-    # allowed. Trigger requestSubmit() after the re-armed watchdog window
-    # plus a little slack; on the buggy build the form stays gated forever.
+    # Second watchdog cycle: a real submit must eventually be allowed. Trigger
+    # requestSubmit() after the re-armed window plus slack; the buggy build stays
+    # gated forever.
     rearm_probe = page.evaluate(
         """async (selector) => {
             const ta = document.querySelector(selector);
@@ -617,8 +592,8 @@ with sync_playwright() as p:
     info("keydown re-pin re-arm PASS")
     clear()
 
-    # 7. Final state. The change-password redirect emits benign 401 noise,
-    #    so we filter via is_benign_* and only fail on real errors.
+    # 7. Final state. Filter benign 401 noise from the change-password redirect
+    #    via is_benign_*; fail only on real errors.
     shoot("07-final")
     real_page_errors = [e for e in page_errors if not is_benign_page_error(e)]
     real_console_errors = [e for e in console_errors if not is_benign_console_error(e)]

tests/studio/playwright_chat_ui.py

@@ -141,10 +141,7 @@ def expected_default_model():
     for node in tree.body:
         if not isinstance(node, ast.Assign):
             continue
-        if not any(
-            isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF"
-            for t in node.targets
-        ):
+        if not any(isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF" for t in node.targets):
             continue
         try:
             models = ast.literal_eval(node.value)
@@ -321,9 +318,7 @@ with sync_playwright() as p:
     form_err: Exception | None = None
     for _form_attempt in range(3):
         try:
-            page.goto(
-                f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000
-            )
+            page.goto(f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000)
             try:
                 page.wait_for_load_state("networkidle", timeout = 30_000)
             except Exception:
@@ -382,9 +377,7 @@ with sync_playwright() as p:
                     flush = True,
                 )
             if page_errors:
-                print(
-                    f"[ui]   first pageerror:    {page_errors[0][:200]!r}", flush = True
-                )
+                print(f"[ui]   first pageerror:    {page_errors[0][:200]!r}", flush = True)
             try:
                 shoot(f"01-change-password-attempt-{_form_attempt + 1}-fail")
             except Exception:
@@ -458,9 +451,7 @@ with sync_playwright() as p:
                     flush = True,
                 )
             if page_errors:
-                print(
-                    f"[ui]   first pageerror:    {page_errors[0][:200]!r}", flush = True
-                )
+                print(f"[ui]   first pageerror:    {page_errors[0][:200]!r}", flush = True)
             try:
                 shoot(f"03-composer-wait-attempt-{_attempt + 1}-fail")
             except Exception:
@@ -557,9 +548,7 @@ with sync_playwright() as p:
     try:
         sel_text = (selector_btn.text_content(timeout = 2_000) or "").strip()
     except Exception as _sel_err:
-        info(
-            f"WARN: model-selector probe skipped: {type(_sel_err).__name__}: {_sel_err}"
-        )
+        info(f"WARN: model-selector probe skipped: {type(_sel_err).__name__}: {_sel_err}")
     if sel_text:
         info(f"model selector button text: {sel_text!r}")
         shoot("03b-default-model-button")
@@ -595,10 +584,7 @@ with sync_playwright() as p:
     if load_resp.get("error"):
         fail(f"/api/inference/load wedged: {load_resp['error']!r}")
     if load_resp["status"] != 200:
-        fail(
-            f"/api/inference/load returned {load_resp['status']}: "
-            f"{load_resp.get('body')!r}"
-        )
+        fail(f"/api/inference/load returned {load_resp['status']}: " f"{load_resp.get('body')!r}")
     info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}")
 
     # Studio caches the per-context model state in zustand; reload
@@ -845,8 +831,7 @@ with sync_playwright() as p:
         # Look for either "Disable X" or "Enable X" -- whichever
         # is currently rendered.
         toggle = page.locator(
-            f'button[aria-label="Disable {feature}"], '
-            f'button[aria-label="Enable {feature}"]'
+            f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]'
         ).first
         if toggle.count() == 0:
             info(f"toggle '{feature}' not present on this layout")
@@ -862,8 +847,7 @@ with sync_playwright() as p:
         page.wait_for_timeout(200)
         after = (
             page.locator(
-                f'button[aria-label="Disable {feature}"], '
-                f'button[aria-label="Enable {feature}"]'
+                f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]'
             ).first.get_attribute("aria-label")
             or ""
         )
@@ -874,8 +858,7 @@ with sync_playwright() as p:
         # Flip back so test state is unchanged.
         try:
             page.locator(
-                f'button[aria-label="Disable {feature}"], '
-                f'button[aria-label="Enable {feature}"]'
+                f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]'
             ).first.click()
         except Exception:
             pass
@@ -968,8 +951,7 @@ with sync_playwright() as p:
                 except Exception as exc:
                     if attempt == 1:
                         soft_fail(
-                            f"theme cycle {cycle + 1}: account-menu click failed "
-                            f"({exc!r})"
+                            f"theme cycle {cycle + 1}: account-menu click failed " f"({exc!r})"
                         )
                     continue
                 try:
@@ -1020,8 +1002,7 @@ with sync_playwright() as p:
             if click_err is not None:
                 page.keyboard.press("Escape")
                 soft_fail(
-                    f"theme cycle {cycle + 1}: theme menuitem click failed "
-                    f"({click_err!r})"
+                    f"theme cycle {cycle + 1}: theme menuitem click failed " f"({click_err!r})"
                 )
                 break
             # Settle. The ".dark" class on <html> is the ground
@@ -1078,9 +1059,7 @@ with sync_playwright() as p:
         # progressively more permissive locators so the test stays
         # green on both platforms.
         candidates = [
-            page.get_by_role(
-                "button", name = re.compile(rf"^\s*{label}\s*$", re.I)
-            ).first,
+            page.get_by_role("button", name = re.compile(rf"^\s*{label}\s*$", re.I)).first,
             page.locator(f'button:has-text("{label}")').first,
             page.locator(f'a:has-text("{label}")').first,
             page.locator(f'[data-sidebar="menu-button"]:has-text("{label}")').first,
@@ -1115,7 +1094,37 @@ with sync_playwright() as p:
     step("sidebar nav: New Chat -> Compare -> Search -> Recipes")
     click_nav("New Chat", r"/chat")
     shoot("11-new-chat")
-    click_nav("Compare", r"/chat\?")  # /chat?compare=...
+    # Compare moved into the composer + menu (Tools and attachments).
+    plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first
+    if plus_btn.count() > 0:
+        plus_btn.click(force = True)
+        page.wait_for_timeout(400)
+        compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
+        if compare_item.count() == 0:
+            # The plus menu was decluttered: Compare chat now lives in the
+            # "More" submenu; hover (then click as fallback) to open it.
+            more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first
+            if more_trigger.count() > 0:
+                more_trigger.hover()
+                page.wait_for_timeout(400)
+                compare_item = page.get_by_role(
+                    "menuitem", name = re.compile(r"Compare chat", re.I)
+                ).first
+                if compare_item.count() == 0:
+                    more_trigger.click(force = True)
+                    page.wait_for_timeout(400)
+                    compare_item = page.get_by_role(
+                        "menuitem", name = re.compile(r"Compare chat", re.I)
+                    ).first
+        if compare_item.count() > 0:
+            compare_item.click(force = True)
+            page.wait_for_timeout(800)
+            if not re.search(r"/chat\?", page.url):
+                soft_fail(f"'Compare chat' didn't open compare; current: {page.url}")
+        else:
+            soft_fail("composer + menu: 'Compare chat' item not found")
+    else:
+        soft_fail("composer + menu: plus button not found")
     shoot("12-compare")
     # Search opens a dialog (not a route change).
     search_btn = page.get_by_role("button", name = re.compile(r"^search$", re.I)).first
@@ -1141,9 +1150,7 @@ with sync_playwright() as p:
         step("Developer (API) tab via account menu")
         acct.click()
         page.wait_for_timeout(400)
-        dev = page.get_by_role(
-            "menuitem", name = re.compile(r"developer|api", re.I)
-        ).first
+        dev = page.get_by_role("menuitem", name = re.compile(r"developer|api", re.I)).first
         if dev.count() > 0:
             dev.click()
             page.wait_for_timeout(800)
@@ -1160,9 +1167,7 @@ with sync_playwright() as p:
                 re.compile(r"api keys|developer", re.I),
             ).first
             if keys_section.count() > 0:
-                info(
-                    f"OK API tab text: {(keys_section.text_content() or '').strip()[:80]!r}"
-                )
+                info(f"OK API tab text: {(keys_section.text_content() or '').strip()[:80]!r}")
             # Close dialog with Escape.
             page.keyboard.press("Escape")
             page.wait_for_timeout(300)
@@ -1180,9 +1185,7 @@ with sync_playwright() as p:
     page.wait_for_timeout(1500)
     # Recipe cards are rendered as <a> or button elements; count
     # all clickable headings under main + screenshot.
-    headings = page.locator(
-        "main h2, main h3, [data-recipe], a[href*='/data-recipes/']"
-    )
+    headings = page.locator("main h2, main h3, [data-recipe], a[href*='/data-recipes/']")
     n_cards = headings.count()
     info(f"Recipes route headings/cards: {n_cards}")
     shoot("15b-recipes-cards")
@@ -1271,10 +1274,7 @@ with sync_playwright() as p:
             info(f"recent-thread click {i} failed: {_click_err!s}")
             continue
     if not clicked_recent:
-        soft_fail(
-            f"no Recents entry was clickable within 30s deadline "
-            f"(n_threads={n_threads})"
-        )
+        soft_fail(f"no Recents entry was clickable within 30s deadline " f"(n_threads={n_threads})")
     # Back to chat.
     page.goto(f"{BASE}/chat")
     composer = page.locator('textarea[aria-label="Message input"]')

tests/studio/playwright_extra_ui.py

@@ -170,9 +170,7 @@ with sync_playwright() as p:
     form_err: Exception | None = None
     for _form_attempt in range(3):
         try:
-            page.goto(
-                f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000
-            )
+            page.goto(f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000)
             try:
                 page.wait_for_load_state("networkidle", timeout = 30_000)
             except Exception:
@@ -327,17 +325,34 @@ with sync_playwright() as p:
     # 1. Compare tab.
     # ─────────────────────────────────────────────────────
     step("Compare tab: send to two panes")
-    # The Compare nav lives in the sidebar; click it.
-    compare_nav = page.locator('[data-tour="chat-compare"]').first
-    if compare_nav.count() == 0:
-        compare_nav = page.get_by_role(
-            "button",
-            name = re.compile(r"^\s*Compare\s*$", re.I),
-        ).first
-    if compare_nav.count() == 0:
+    # Compare moved into the composer + menu (Tools and attachments).
+    compare_opened = False
+    plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first
+    if plus_btn.count() > 0:
+        plus_btn.click(force = True)
+        page.wait_for_timeout(400)
+        compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
+        if compare_item.count() == 0:
+            # Compare chat moved into the "More" submenu; hover, then click fallback.
+            more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first
+            if more_trigger.count() > 0:
+                more_trigger.hover()
+                page.wait_for_timeout(400)
+                compare_item = page.get_by_role(
+                    "menuitem", name = re.compile(r"Compare chat", re.I)
+                ).first
+                if compare_item.count() == 0:
+                    more_trigger.click(force = True)
+                    page.wait_for_timeout(400)
+                    compare_item = page.get_by_role(
+                        "menuitem", name = re.compile(r"Compare chat", re.I)
+                    ).first
+        if compare_item.count() > 0:
+            compare_item.click(force = True)
+            compare_opened = True
+    if not compare_opened:
         soft_fail("Compare nav not found")
     else:
-        compare_nav.click()
         page.wait_for_timeout(1500)
         shoot("02-compare-opened")
         # Compare view's container.
@@ -410,9 +425,7 @@ with sync_playwright() as p:
                         arg = ok_count_before + 4,
                         timeout = 60_000,
                     )
-                    info(
-                        "OK Compare: 4 total new assistant bubbles after second prompt"
-                    )
+                    info("OK Compare: 4 total new assistant bubbles after second prompt")
                 except Exception as exc:
                     runtime_warn(
                         f"Compare: 4 bubbles didn't appear (panes likely "
@@ -433,9 +446,7 @@ with sync_playwright() as p:
     page.wait_for_timeout(1500)
     shoot("05-recipes-list")
     # Template cards render as <button> elements.
-    templates = page.locator("main button").filter(
-        has_not_text = re.compile(r"^(\+|Create)")
-    )
+    templates = page.locator("main button").filter(has_not_text = re.compile(r"^(\+|Create)"))
     n_templates = templates.count()
     info(f"recipe templates visible: {n_templates}")
     if n_templates == 0:
@@ -469,9 +480,7 @@ with sync_playwright() as p:
     shoot("07-export")
     if chat_only:
         if "/export" in page.url:
-            soft_fail(
-                f"chat-only mode should redirect /export -> /chat; url={page.url}"
-            )
+            soft_fail(f"chat-only mode should redirect /export -> /chat; url={page.url}")
         else:
             info(f"OK chat-only redirected /export -> {page.url}")
     else:
@@ -521,16 +530,12 @@ with sync_playwright() as p:
     shoot("08-studio")
     if chat_only:
         if "/studio" in page.url:
-            soft_fail(
-                f"chat-only mode should redirect /studio -> /chat; url={page.url}"
-            )
+            soft_fail(f"chat-only mode should redirect /studio -> /chat; url={page.url}")
         else:
             info(f"OK chat-only redirected /studio -> {page.url}")
     else:
         for tab_name in ("Configure", "Current run", "History"):
-            tab = page.get_by_role(
-                "tab", name = re.compile(rf"^\s*{tab_name}\s*$", re.I)
-            ).first
+            tab = page.get_by_role("tab", name = re.compile(rf"^\s*{tab_name}\s*$", re.I)).first
             if tab.count() == 0:
                 soft_fail(f"tab '{tab_name}' not found in /studio")
             else:
@@ -592,9 +597,7 @@ with sync_playwright() as p:
                     info(f"OK Settings tab '{tab_name}' body length={body_text}")
                     seen_tabs.append(tab_name)
                 else:
-                    soft_fail(
-                        f"Settings tab '{tab_name}' body suspiciously short: {body_text}"
-                    )
+                    soft_fail(f"Settings tab '{tab_name}' body suspiciously short: {body_text}")
             except Exception as exc:
                 soft_fail(f"Settings tab '{tab_name}' click failed: {exc!r}")
         shoot("10-settings-tabs-visited")

tests/studio/run_real_mlx_smoke.py

@@ -4,58 +4,24 @@
 """
 End-to-end MLX smoke test on real Apple Silicon -- multi-process driver.
 
-Two subcommands so the workflow can drive cold-start reloads in fresh
-Python processes (the way real users hit the load path):
+Two subcommands let the workflow drive cold-start reloads in fresh
+processes (the way real users hit the load path):
 
     python run_real_mlx_smoke.py train  --workdir DIR
     python run_real_mlx_smoke.py reload --format {lora|merged|gguf} --dir D
 
-The `train` subcommand:
-  1. Loads `unsloth/gemma-3-270m-it` via FastMLXModel.from_pretrained.
-  2. Applies LoRA r=8 on q/k/v/o.
-  3. Computes pre-training loss + grad norm via mx.nn.value_and_grad.
-  4. Trains 7 deterministic steps on a dataset of the SAME row repeated
-     ("<<HELLO!!>> My name is Unsloth!"), with batch_size=2 and
-     gradient_accumulation_steps=3 so each step processes 6 sequences
-     and the run sees 42 sequences total.
-  5. Computes post-training loss + grad norm.
-  6. Generates from "<<HELLO!!>> My name is " and asserts "Unsloth"
-     appears in the in-memory completion.
-  7. Saves the trained model in three formats:
-       - LoRA adapter (save_pretrained_merged save_method="lora")
-       - Merged 16-bit (save_pretrained_merged save_method="merged_16bit")
-       - GGUF (save_pretrained_gguf, best-effort -- skipped with a
-         clear reason if save raises; e.g. llama.cpp's
-         convert_hf_to_gguf currently asserts on Gemma-3-270m's
-         tokenizer vocab. Soft-skipped so the LoRA + merged checks
-         continue to gate the PR.)
-  8. Emits `train_metrics.json` with per-phase timing / peak GPU /
-     peak RSS / per-step losses / pre+post grad norms / generations
-     / gguf_supported flag, for regression detection across CI runs.
+`train` loads gemma-3-270m-it, applies LoRA, probes pre/post loss+grad,
+overfits one repeated row, generates, saves in lora/merged_16bit/gguf
+(gguf best-effort), and writes train_metrics.json. `reload` reopens each
+saved format in a fresh process and writes <format>_reload_metrics.json.
 
-Reloads run as separate workflow steps so each is a fresh Python
-process. For lora / merged the reload uses
-FastMLXModel.from_pretrained directly. For gguf the reload spawns
-the llama-cli binary built by save_pretrained_gguf and parses
-stdout. Each subcommand emits `<format>_reload_metrics.json` next
-to the saved dir.
+GGUF export and LoRA reload fixes land in unslothai/unsloth-zoo#627.
 
-The two upstream unsloth_zoo bugs the earlier draft of this script
-worked around are fixed in unslothai/unsloth-zoo#627: GGUF export
-no longer raises NotImplementedError on Apple Silicon (llama_cpp.py
-catches it from the device_type module-level call) and LoRA reload
-via FastMLXModel.from_pretrained(lora_dir) works without an external
-config.json copy (mlx_loader.py preserves local_path when config.json
-is missing so the adapter_config.json branch can run).
+Determinism: seeds random/numpy/mlx.core.random and forwards SEED to
+from_pretrained / get_peft_model / MLXTrainingConfig. Metal has minor
+reduction-order nondeterminism, so loss assertions are bounds, not exact.
 
-Determinism: seeds Python `random`, `numpy`, and `mlx.core.random` in
-every process before any MLX operation. Forwards `random_state=SEED`
-to FastMLXModel.from_pretrained / get_peft_model and `seed=SEED` to
-MLXTrainingConfig. Metal still has minor reduction-order
-nondeterminism, so loss assertions are bounds rather than exact.
-
-Only runnable on a real Apple Silicon host; invoked from
-.github/workflows/mlx-ci.yml on the macos-14 runner.
+Apple-Silicon only; invoked from .github/workflows/mlx-ci.yml.
 """
 
 from __future__ import annotations
@@ -99,13 +65,8 @@ def _peak_gpu_gb() -> float:
 
     if not mx.metal.is_available():
         return 0.0
-    # Newer MLX deprecates mx.metal.get_peak_memory in favour of the
-    # top-level mx.get_peak_memory; fall back to the old API for
-    # compatibility with older MLX versions still present in the
-    # environment.
-    getter = getattr(mx, "get_peak_memory", None) or getattr(
-        mx.metal, "get_peak_memory", None
-    )
+    # Newer MLX moved get_peak_memory to top-level; fall back to mx.metal for old versions.
+    getter = getattr(mx, "get_peak_memory", None) or getattr(mx.metal, "get_peak_memory", None)
     if getter is None:
         return 0.0
     try:
@@ -115,8 +76,7 @@ def _peak_gpu_gb() -> float:
 
 
 def _peak_rss_gb() -> float:
-    """Peak resident set size for this process. macOS getrusage returns
-    bytes; Linux returns kilobytes."""
+    """Peak RSS for this process (macOS getrusage = bytes, Linux = KB)."""
     rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
     if sys.platform == "darwin":
         return float(rss) / (1024**3)
@@ -124,8 +84,7 @@ def _peak_rss_gb() -> float:
 
 
 class Phase:
-    """Wall-clock + memory tracker for a named phase. Records into a
-    metrics dict so we can later JSON-dump for regression detection."""
+    """Wall-clock + memory tracker for a named phase; records into a metrics dict."""
 
     def __init__(self, name: str, metrics: dict):
         self.name = name
@@ -156,8 +115,7 @@ class Phase:
 
 
 def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, float]:
-    """One forward+backward of next-token cross-entropy on `text`.
-    Returns (loss, ||grad||_2)."""
+    """One fwd+bwd of next-token CE on `text`. Returns (loss, ||grad||_2)."""
     import mlx.core as mx
     import mlx.nn as nn
     from mlx.utils import tree_flatten
@@ -186,26 +144,13 @@ def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, flo
     return float(loss_val.item()), float(mx.sqrt(norm_sq).item())
 
 
-def _teacher_forced_completion_loss(
-    model, tokenizer, prompt: str, completion: str
-) -> float:
-    """Mean next-token CE loss on `completion` tokens given `prompt` (teacher
-    forced -- no decoding, no sampling, no greedy argmax).
+def _teacher_forced_completion_loss(model, tokenizer, prompt: str, completion: str) -> float:
+    """Mean next-token CE on `completion` given `prompt`, teacher-forced.
 
-    Decouples the memorisation check from greedy-decode geometry. A 47-round,
-    13-seed sweep on this fixture showed greedy `completion in output` lands
-    in the 46-77% range across MLX configs (config-fragile), while
-    post_train_loss is < 0.1 in 100% of configs that reach the basin. Teacher-
-    forced completion loss is a subset of post_train_loss so it inherits the
-    same reliability AND is more specific: it asserts *what* the model
-    memorised, not just *that* it reached low loss on the full row.
-
-    Args:
-      model:       the LoRA-trained MLX model
-      tokenizer:   the tokenizer used during training (must match)
-      prompt:      the conditioning text (e.g. PROMPT)
-      completion:  the substring the model should have learnt to emit
-                   after `prompt` (e.g. EXPECT_IN_OUTPUT + "!")
+    Decouples the memorisation check from greedy-decode geometry: a sweep
+    found greedy `completion in output` lands at 46-77% across MLX configs
+    while post_train_loss is < 0.1 whenever the run reaches the basin. This
+    asserts *what* the model memorised, not just that loss is low.
 
     Returns mean cross-entropy over the completion's tokens.
     """
@@ -229,9 +174,7 @@ def _teacher_forced_completion_loss(
     start = len(prompt_ids) - 1
     completion_logits = logits[:, start:, :]
     completion_targets = targets[:, start:]
-    loss = nn.losses.cross_entropy(
-        completion_logits, completion_targets, reduction = "mean"
-    )
+    loss = nn.losses.cross_entropy(completion_logits, completion_targets, reduction = "mean")
     return float(loss.item())
 
 
@@ -281,13 +224,9 @@ def cmd_train(args) -> int:
     mx.random.seed(SEED)
 
     with Phase("apply_lora", metrics):
-        # Standard unsloth LoRA target set (q/k/v/o + gate/up/down).
-        # With bs=2 grad_accum=3 (effective batch 6) the q/k/v/o-only
-        # LoRA collapsed in 7 steps -- training loss kept dropping but
-        # inference output the structural skeleton ("My name") without
-        # recovering the specific "Unsloth" token. Including the MLP
-        # projections gives the LoRA enough capacity to memorize the
-        # training row at the larger effective batch.
+        # Standard unsloth LoRA target set (q/k/v/o + gate/up/down). q/k/v/o
+        # alone collapsed in 7 steps (loss dropped but "Unsloth" wasn't
+        # recovered); MLP projections add the capacity to memorize the row.
         model = FastMLXModel.get_peft_model(
             model,
             r = 8,
@@ -320,29 +259,19 @@ def cmd_train(args) -> int:
         config = MLXTrainingConfig(
             per_device_train_batch_size = 2,
             gradient_accumulation_steps = 3,
-            # 47-round mlx-parity-probes sweep (PR #5498 / staging-2#119)
-            # found 7 steps is below the convergence horizon at any clip
-            # setting -- the trainer hasn't memorized the train row yet
-            # when the smoke probes loss/generation. At 30 steps every
-            # seed tested hits post_train_loss=0 across all clip
-            # configurations, so 30 is the seed-robust gate.
+            # Sweep (PR #5498) found 7 steps is below the convergence horizon
+            # at any clip; at 30 steps every seed hits post_train_loss=0, so
+            # 30 is the seed-robust gate.
             max_steps = 30,
             learning_rate = 1e-3,
             warmup_steps = 0,
             lr_scheduler_type = "constant",
             optim = "adamw",
             weight_decay = 0.0,
-            # max_grad_value (elementwise) is materially cheaper than
-            # max_grad_norm on MLX -- norm clip needs a cross-tree
-            # reduction + materializing all grad tensors at full
-            # precision, value clip is tree_map(mx.clip) per leaf.
-            # MLXTrainingConfig defaults to max_grad_value=1.0 for
-            # exactly this reason; pin both explicitly here so the
-            # configured clip matches what runs (the trainer prints a
-            # notice when both > 0 and value wins, so disable norm).
-            # Empirical 13-seed pass rate at this fixture: value=1.0
-            # 62%, norm=1.0 46%, value=5.0 33%, value=0.5 77% -- the
-            # cheaper default is also the higher-pass-rate default.
+            # Elementwise value clip is cheaper than norm clip on MLX (no
+            # cross-tree reduction) and has a higher 13-seed pass rate at this
+            # fixture (value=1.0 62%, norm=1.0 46%). Pin both: value wins when
+            # both > 0, so disable norm.
             max_grad_norm = 0.0,
             max_grad_value = 1.0,
             logging_steps = 1,
@@ -364,7 +293,15 @@ def cmd_train(args) -> int:
         )
 
         def _on_step(
-            step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens, grad_norm = None
+            step,
+            total,
+            loss,
+            lr,
+            tok_s,
+            peak_gb,
+            elapsed,
+            num_tokens,
+            grad_norm = None,
         ):
             losses_per_step.append(round(float(loss), 4))
             grad_text = f"  grad={grad_norm:.4f}" if grad_norm is not None else ""
@@ -608,9 +545,7 @@ def cmd_reload(args) -> int:
             in_mem_loss = None
     metrics["in_memory_generation_ref"] = in_mem_out
     metrics["in_memory_post_train_loss"] = in_mem_loss
-    metrics["reload_completion_matches_in_memory"] = (
-        in_mem_out is not None and out == in_mem_out
-    )
+    metrics["reload_completion_matches_in_memory"] = in_mem_out is not None and out == in_mem_out
     if isinstance(in_mem_loss, (int, float)) and math.isfinite(in_mem_loss):
         reload_loss, _ = _compute_loss_and_grad_norm(m, t, TRAIN_TEXT)
         metrics["reload_post_train_loss"] = round(reload_loss, 4)
@@ -625,8 +560,7 @@ def cmd_reload(args) -> int:
         # workdir layouts): keep a non-empty-completion gate.
         body = out.replace(PROMPT, "", 1).strip()
         assert len(body) >= 4, (
-            f"reload {args.format!r} produced no usable output for "
-            f"{PROMPT!r}: {out!r}"
+            f"reload {args.format!r} produced no usable output for " f"{PROMPT!r}: {out!r}"
         )
 
     metrics["final_peak_gpu_gb"] = round(_peak_gpu_gb(), 3)
@@ -677,9 +611,7 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int:
 
     print(f"  [reload:gguf] stdout (head):\n{proc.stdout[:800]}", flush = True)
     if proc.returncode != 0:
-        raise SystemExit(
-            f"llama-cli exit {proc.returncode}; stderr head: {proc.stderr[:400]}"
-        )
+        raise SystemExit(f"llama-cli exit {proc.returncode}; stderr head: {proc.stderr[:400]}")
     # llama.cpp uses different tokenisation + sampling internals than
     # mlx_lm, so the GGUF reload completion does not have to match the
     # in-memory completion exactly. Require non-empty, non-prompt-only
@@ -689,8 +621,7 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int:
     body = (proc.stdout or "").replace(PROMPT, "", 1).strip()
     metrics["gguf_has_expected"] = EXPECT_IN_OUTPUT in (proc.stdout or "")
     assert len(body) >= 4, (
-        f"GGUF reload produced no usable output for {PROMPT!r}: "
-        f"{proc.stdout[:400]!r}"
+        f"GGUF reload produced no usable output for {PROMPT!r}: " f"{proc.stdout[:400]!r}"
     )
 
     metrics["final_peak_rss_gb"] = round(_peak_rss_gb(), 3)

tests/studio/studio_api_smoke.py

@@ -192,9 +192,7 @@ try:
         acao = r.headers.get("Access-Control-Allow-Origin", "")
         acac = r.headers.get("Access-Control-Allow-Credentials", "")
         if acao == "*" and acac.lower() == "true":
-            fail(
-                f"CORS: wildcard origin + credentials=true (acao={acao!r}, acac={acac!r})"
-            )
+            fail(f"CORS: wildcard origin + credentials=true (acao={acao!r}, acac={acac!r})")
         else:
             ok(f"CORS preflight acao={acao!r} acac={acac!r}")
 except Exception as exc:

tests/studio/test_auth_form_input_count.py

@@ -3,21 +3,15 @@
 
 """Pin the auth-form input-count contract on the change-password page.
 
-PR #5490 added a third visible "Current password" input so the
-admin-forced must_change_password reset path (where no bootstrap
-script is injected) could supply a current password. The side
-effect was that the dominant first-boot UX, where the backend
-injects window.__UNSLOTH_BOOTSTRAP__ and the form silently reuses
-that password, now showed three visible inputs instead of the two
-it had before. PR #5545 restores the two-input first-boot UX by
-rendering the Current password input only when
-window.__UNSLOTH_BOOTSTRAP__ is absent.
+PR #5490 added a third "Current password" input for the admin-forced reset
+path; this regressed the first-boot UX (which reuses the injected
+window.__UNSLOTH_BOOTSTRAP__ password) to three visible inputs. PR #5545
+restores two inputs by rendering Current password only when BOOTSTRAP is
+absent.
 
-These tests inspect the auth-form source file directly. They never
-boot Studio, never spawn a browser, and have no network or device
-dependencies, so they are fully deterministic and run on any CI
-runner without a JS toolchain. The companion Playwright probe lives
-in tests/studio/playwright_chat_ui.py and covers the runtime side.
+These tests inspect the auth-form source directly (no Studio, browser, or
+network), so they run deterministically on any CI runner. The runtime side
+is covered by tests/studio/playwright_chat_ui.py.
 """
 
 from __future__ import annotations
@@ -63,10 +57,7 @@ def test_hasbootstrappassword_constant_is_derived_from_bootstrap_window_value():
     a prop) would silently drift from the backend's bootstrap-injection
     contract in studio/backend/main.py::_inject_bootstrap."""
     src = AUTH_FORM.read_text()
-    assert (
-        "const hasBootstrapPassword = Boolean(window.__UNSLOTH_BOOTSTRAP__?.password);"
-        in src
-    ), (
+    assert "const hasBootstrapPassword = Boolean(window.__UNSLOTH_BOOTSTRAP__?.password);" in src, (
         "hasBootstrapPassword constant missing or its derivation drifted; "
         "this is the gate that hides the Current password input on first boot"
     )

tests/studio/test_cancel_atomicity.py

@@ -16,23 +16,14 @@ import threading
 from pathlib import Path
 
 
-SOURCE_PATH = (
-    Path(__file__).resolve().parents[2]
-    / "studio"
-    / "backend"
-    / "routes"
-    / "inference.py"
-)
+SOURCE_PATH = Path(__file__).resolve().parents[2] / "studio" / "backend" / "routes" / "inference.py"
 _SRC = SOURCE_PATH.read_text()
 _TREE = ast.parse(_SRC)
 
 
 def _find_function(name: str) -> ast.FunctionDef | ast.AsyncFunctionDef:
     for node in ast.walk(_TREE):
-        if (
-            isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
-            and node.name == name
-        ):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name:
             return node
     raise AssertionError(f"function {name!r} not found")
 
@@ -187,8 +178,7 @@ def test_parallel_cancel_vs_register_never_drops():
         tracker.__exit__(None, None, None)
 
     assert dropped == 0, (
-        f"TOCTOU regression: {dropped}/{trials} parallel trials silently "
-        f"dropped the cancel"
+        f"TOCTOU regression: {dropped}/{trials} parallel trials silently " f"dropped the cancel"
     )
 
 

tests/studio/test_cancel_id_wiring.py

@@ -25,12 +25,8 @@ from pathlib import Path
 WORKSPACE = Path(__file__).resolve().parents[2]
 MODELS_SRC = (WORKSPACE / "studio/backend/models/inference.py").read_text()
 ROUTES_SRC = (WORKSPACE / "studio/backend/routes/inference.py").read_text()
-ADAPTER_SRC = (
-    WORKSPACE / "studio/frontend/src/features/chat/api/chat-adapter.ts"
-).read_text()
-API_TYPES_SRC = (
-    WORKSPACE / "studio/frontend/src/features/chat/types/api.ts"
-).read_text()
+ADAPTER_SRC = (WORKSPACE / "studio/frontend/src/features/chat/api/chat-adapter.ts").read_text()
+API_TYPES_SRC = (WORKSPACE / "studio/frontend/src/features/chat/types/api.ts").read_text()
 
 
 def _find_class(tree: ast.AST, name: str) -> ast.ClassDef | None:
@@ -49,18 +45,16 @@ def test_chat_completion_request_has_cancel_id_field():
         for n in cls.body
         if isinstance(n, ast.AnnAssign) and isinstance(n.target, ast.Name)
     }
-    assert "cancel_id" in fields, (
-        "ChatCompletionRequest must expose a cancel_id field for per-run "
-        "cancellation routing"
-    )
+    assert (
+        "cancel_id" in fields
+    ), "ChatCompletionRequest must expose a cancel_id field for per-run cancellation routing"
 
 
 def test_cancel_route_matches_cancel_id_exclusively_when_present():
-    # A stale cancel POST carrying cancel_id AND session_id must not
-    # cancel a later run on the same thread via the shared session_id.
-    # Enforce this by requiring the handler to early-return through an
-    # exclusive-cancel_id path -- either an atomic helper or a keys
-    # list containing ONLY cancel_id (never session_id).
+    # A stale cancel POST carrying cancel_id AND session_id must not cancel a
+    # later run via the shared session_id. Require the handler to early-return
+    # through an exclusive-cancel_id path (atomic helper, or keys list with
+    # ONLY cancel_id, never session_id).
     for node in ast.walk(ast.parse(ROUTES_SRC)):
         if isinstance(node, ast.AsyncFunctionDef) and node.name == "cancel_inference":
             break
@@ -117,9 +111,7 @@ def test_chat_adapter_generates_cancel_id_per_run():
     )
     assert m, "chat-adapter.ts must declare a per-run `cancelId` constant"
     rhs = m.group(1)
-    assert (
-        "randomUUID" in rhs
-    ), "cancelId should prefer crypto.randomUUID() for uniqueness"
+    assert "randomUUID" in rhs, "cancelId should prefer crypto.randomUUID() for uniqueness"
 
 
 def test_chat_adapter_sends_cancel_id_in_completion_payload():
@@ -144,10 +136,9 @@ def test_chat_adapter_sends_cancel_id_in_abort_cancel_post():
 
 
 def test_abort_cancel_post_uses_plain_fetch_with_manual_auth_header():
-    # authFetch redirects to login on 401, which would kick the user to
-    # the login page mid-stop if the access token expired during a long
-    # stream. Use plain fetch + manual Authorization header for a
-    # best-effort cancel that never triggers the refresh/redirect flow.
+    # authFetch redirects to login on 401, kicking the user out mid-stop if
+    # the token expired during a long stream. Use plain fetch + manual
+    # Authorization header for a best-effort cancel with no refresh/redirect.
     start = ADAPTER_SRC.find("const onAbortCancel")
     assert start >= 0, "onAbortCancel handler missing"
     rest = ADAPTER_SRC[start:]

tests/studio/test_chat_preset_builtin_invariants.py

@@ -17,9 +17,7 @@ def _source_path(relative_path: str) -> Path:
     return WORKDIR / "unsloth_repo" / relative_path
 
 
-PRESET_POLICY = _source_path(
-    "studio/frontend/src/features/chat/presets/preset-policy.ts"
-)
+PRESET_POLICY = _source_path("studio/frontend/src/features/chat/presets/preset-policy.ts")
 RUNTIME_TYPES = _source_path("studio/frontend/src/features/chat/types/runtime.ts")
 TEMP = WORKDIR / "temp" / "chat_preset_builtin_invariants"
 
@@ -42,8 +40,7 @@ def _require_node():
 def _ensure_harness():
     TEMP.mkdir(parents = True, exist_ok = True)
     (TEMP / "register.mjs").write_text(
-        "import { register } from 'node:module';\n"
-        "register('./loader.mjs', import.meta.url);\n"
+        "import { register } from 'node:module';\nregister('./loader.mjs', import.meta.url);\n"
     )
     (TEMP / "loader.mjs").write_text(
         "export function resolve(specifier, context, next) {\n"

tests/studio/test_cli_repo_variant.py

@@ -42,12 +42,8 @@ def _load_split_repo_variant():
         typer_stub.echo = lambda *args, **kwargs: None
         sys.modules["typer"] = typer_stub
 
-    studio_py = (
-        Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands" / "studio.py"
-    )
-    spec = importlib.util.spec_from_file_location(
-        "_studio_for_repo_variant_test", studio_py
-    )
+    studio_py = Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands" / "studio.py"
+    spec = importlib.util.spec_from_file_location("_studio_for_repo_variant_test", studio_py)
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
     return module._split_repo_variant

tests/studio/test_cli_run_alias.py

@@ -38,9 +38,7 @@ def test_top_level_run_alias_registered():
         # Decorator-call form has a string literal "run" as the first
         # positional or as keyword ``name="run"``.
         first_pos = call.args[0] if call.args else None
-        keyword_name = next(
-            (kw.value for kw in call.keywords if kw.arg == "name"), None
-        )
+        keyword_name = next((kw.value for kw in call.keywords if kw.arg == "name"), None)
         is_run = (isinstance(first_pos, ast.Constant) and first_pos.value == "run") or (
             isinstance(keyword_name, ast.Constant) and keyword_name.value == "run"
         )
@@ -66,4 +64,6 @@ def test_studio_run_imported_for_alias():
             if alias.name == "run":
                 has_import = True
                 break
-    assert has_import, "Expected `from unsloth_cli.commands.studio import run` in unsloth_cli/__init__.py"
+    assert (
+        has_import
+    ), "Expected `from unsloth_cli.commands.studio import run` in unsloth_cli/__init__.py"

tests/studio/test_cli_studio_defaults.py

@@ -7,9 +7,7 @@ full unsloth_cli dependencies (typer/pydantic) at test-collection time.
 import ast
 from pathlib import Path
 
-_STUDIO_CMD_PY = (
-    Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands" / "studio.py"
-)
+_STUDIO_CMD_PY = Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands" / "studio.py"
 
 
 def _find_typer_option_default(source: str, func_name: str, long_option: str):
@@ -25,13 +23,10 @@ def _find_typer_option_default(source: str, func_name: str, long_option: str):
             continue
         if func_node.name != func_name:
             continue
-        # Walk both regular args and kwonly args, each paired with its default.
         all_args = func_node.args.args + func_node.args.kwonlyargs
         all_defaults = func_node.args.defaults + [
             d for d in func_node.args.kw_defaults if d is not None
         ]
-        # ast pads defaults right-aligned against args (ignoring kwonly). We
-        # iterate calls directly, which is simpler and robust.
         for default in all_defaults:
             if not isinstance(default, ast.Call):
                 continue
@@ -44,8 +39,7 @@ def _find_typer_option_default(source: str, func_name: str, long_option: str):
             )
             if not is_typer_option:
                 continue
-            # First positional is the default value; remaining positionals are
-            # option flags like "--host", "-H".
+            # First positional is the default; the rest are flags like "--host".
             if not default.args:
                 continue
             flags = [
@@ -78,9 +72,7 @@ def test_studio_run_host_is_loopback():
     """`unsloth studio run` --host typer Option default must be 127.0.0.1."""
     source = _STUDIO_CMD_PY.read_text()
     host_default = _find_typer_option_default(source, "run", "--host")
-    assert (
-        host_default is not None
-    ), "Could not find --host typer.Option default in run()"
+    assert host_default is not None, "Could not find --host typer.Option default in run()"
     assert host_default == "127.0.0.1", (
         f"`unsloth studio run` --host default must be '127.0.0.1' (loopback) "
         f"but got '{host_default}'."

tests/studio/test_cli_studio_stop_windows.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Regression tests for `unsloth studio stop` on Windows (PR #5940).
+
+`stop` once used the POSIX `os.kill(pid, 0)` probe, which raises OSError
+(WinError 87) for every pid on Windows -- crashing before reaching taskkill.
+The fix adds a cross-platform `_pid_alive(pid)` (tasklist on Windows, signal-0
+elsewhere).
+
+AST + mock-only; no real process management, no Studio deps imported.
+"""
+
+import ast
+import os
+import subprocess
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+_STUDIO_CMD_PY = Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands" / "studio.py"
+_SOURCE = _STUDIO_CMD_PY.read_text(encoding = "utf-8")
+
+
+def _func_source(name: str) -> str:
+    """Return the source of a top-level function `name` in studio.py."""
+    tree = ast.parse(_SOURCE)
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name:
+            return ast.get_source_segment(_SOURCE, node)
+    raise AssertionError(f"function {name!r} not found in studio.py")
+
+
+def _load_pid_alive(platform: str, fake_run = None):
+    """Exec just `_pid_alive` with injectable sys/subprocess, so we can drive
+    the win32 branch on any host without importing the full unsloth_cli."""
+    src = _func_source("_pid_alive")
+    fake_sys = types.SimpleNamespace(platform = platform)
+    fake_sub = types.SimpleNamespace(run = fake_run) if fake_run is not None else subprocess
+    ns = {"os": os, "sys": fake_sys, "subprocess": fake_sub}
+    exec(src, ns)
+    return ns["_pid_alive"]
+
+
+# ── AST: stop() must not use the broken bare liveness probe ──────────────────
+
+
+def test_stop_does_not_use_bare_oskill_liveness_probe():
+    """stop() must not call os.kill(pid, 0) -- it crashes on Windows."""
+    stop_src = _func_source("stop")
+    tree = ast.parse(stop_src)
+    for call in ast.walk(tree):
+        if not isinstance(call, ast.Call):
+            continue
+        f = call.func
+        is_os_kill = (
+            isinstance(f, ast.Attribute)
+            and f.attr == "kill"
+            and isinstance(f.value, ast.Name)
+            and f.value.id == "os"
+        )
+        if is_os_kill and len(call.args) == 2:
+            sig = call.args[1]
+            if isinstance(sig, ast.Constant) and sig.value == 0:
+                raise AssertionError(
+                    "stop() still uses os.kill(pid, 0); it raises WinError 87 on "
+                    "Windows. Use the cross-platform _pid_alive() helper instead."
+                )
+
+
+def test_pid_alive_helper_is_defined_and_used_by_stop():
+    assert "def _pid_alive(" in _SOURCE, "_pid_alive helper missing"
+    assert "_pid_alive(pid)" in _func_source("stop"), "stop() must use _pid_alive"
+    # The helper must special-case Windows via tasklist (os.kill(pid,0) is invalid there).
+    helper = _func_source("_pid_alive")
+    assert 'sys.platform == "win32"' in helper
+    assert "tasklist" in helper
+
+
+# ── Behavioral: the win32 tasklist branch ────────────────────────────────────
+
+
+def _fake_tasklist(returns_pid: int | None, *, raises: bool = False):
+    def _run(
+        cmd,
+        capture_output = False,
+        text = False,
+        timeout = None,
+    ):
+        assert cmd[0] == "tasklist"
+        assert "/FI" in cmd  # filtered by PID
+        if raises:
+            raise OSError("boom")
+        if returns_pid is None:
+            stdout = "INFO: No tasks are running which match the specified criteria.\n"
+        else:
+            stdout = f'"python.exe","{returns_pid}","Console","1","12,345 K"\n'
+        return types.SimpleNamespace(stdout = stdout, returncode = 0)
+
+    return _run
+
+
+def test_pid_alive_windows_true_when_tasklist_lists_pid():
+    pid_alive = _load_pid_alive("win32", fake_run = _fake_tasklist(4242))
+    assert pid_alive(4242) is True
+
+
+def test_pid_alive_windows_false_when_tasklist_empty():
+    pid_alive = _load_pid_alive("win32", fake_run = _fake_tasklist(None))
+    assert pid_alive(4242) is False
+
+
+def test_pid_alive_windows_assumes_alive_when_tasklist_errors():
+    # Can't determine -> assume alive; taskkill is the source of truth.
+    pid_alive = _load_pid_alive("win32", fake_run = _fake_tasklist(None, raises = True))
+    assert pid_alive(4242) is True
+
+
+# ── Behavioral: the POSIX signal-0 branch (skip on Windows runners) ───────────
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason = "POSIX os.kill(pid,0) branch")
+def test_pid_alive_posix_true_for_self_false_for_dead():
+    pid_alive = _load_pid_alive("linux")
+    assert pid_alive(os.getpid()) is True
+    assert pid_alive(2_000_000_000) is False

tests/studio/test_composer_rtl_bidi_attribute.py

@@ -19,7 +19,11 @@ WORKFLOW_YML = REPO / ".github/workflows/studio-ui-smoke.yml"
 IME_PY = REPO / "tests/studio/playwright_chat_ime_i18n.py"
 
 
-def _block_around(src: str, anchor: str, radius: int = 600) -> str:
+def _block_around(
+    src: str,
+    anchor: str,
+    radius: int = 600,
+) -> str:
     idx = src.find(anchor)
     assert idx != -1, f"anchor {anchor!r} not found"
     return src[max(idx - radius, 0) : idx + radius]
@@ -83,9 +87,9 @@ def test_main_composer_has_stuck_compositionend_watchdog():
     composing flag once events go silent; without it Send stays disabled
     forever and CJK input is effectively dropped."""
     src = THREAD_TSX.read_text()
-    assert "IME_STUCK_TIMEOUT_MS" in src, (
-        "main composer is missing the stuck-compositionend watchdog " "(issue #5546)"
-    )
+    assert (
+        "IME_STUCK_TIMEOUT_MS" in src
+    ), "main composer is missing the stuck-compositionend watchdog (issue #5546)"
     assert "onCompositionUpdate" in src, (
         "main composer is missing onCompositionUpdate wiring; the "
         "watchdog only resets while the IME is actively emitting events"
@@ -94,12 +98,10 @@ def test_main_composer_has_stuck_compositionend_watchdog():
 
 def test_compare_composer_has_stuck_compositionend_watchdog():
     src = SHARED_TSX.read_text()
-    assert "IME_STUCK_TIMEOUT_MS" in src, (
-        "compare composer is missing the stuck-compositionend watchdog " "(issue #5546)"
-    )
     assert (
-        "onCompositionUpdate" in src
-    ), "compare composer is missing onCompositionUpdate wiring"
+        "IME_STUCK_TIMEOUT_MS" in src
+    ), "compare composer is missing the stuck-compositionend watchdog (issue #5546)"
+    assert "onCompositionUpdate" in src, "compare composer is missing onCompositionUpdate wiring"
 
 
 def test_main_composer_keydown_repins_composing_during_ime():
@@ -125,11 +127,15 @@ def test_compare_composer_keydown_repins_composing_during_ime():
     )
 
 
-def _extract_block(src: str, anchor: str, opener: str = "(", closer: str = ")") -> str:
-    """Return the source between the first balanced opener/closer that
-    starts at or after `anchor`. Used to scope assertions to a specific
-    handler so a re-arm call in some other function does not satisfy
-    the gate test."""
+def _extract_block(
+    src: str,
+    anchor: str,
+    opener: str = "(",
+    closer: str = ")",
+) -> str:
+    """Return the source within the first balanced opener/closer at or
+    after `anchor`. Scopes assertions to one handler so a re-arm call
+    elsewhere does not satisfy the gate test."""
     start = src.find(anchor)
     assert start != -1, f"anchor {anchor!r} not found"
     open_idx = src.find(opener, start)

tests/studio/test_export_output_path_contract.py

@@ -68,9 +68,7 @@ def test_local_save_assigns_output_path():
                     if isinstance(tgt, ast.Name) and tgt.id == "output_path":
                         assigns.append(node)
         non_none = [
-            a
-            for a in assigns
-            if not (isinstance(a.value, ast.Constant) and a.value.value is None)
+            a for a in assigns if not (isinstance(a.value, ast.Constant) and a.value.value is None)
         ]
         assert non_none, f"{fn_name} never assigns a non-None output_path"
 
@@ -86,9 +84,7 @@ def test_gpu_save_method_bound_for_hub_only():
                 if isinstance(stmt, ast.If):
                     test = stmt.test
                     if isinstance(test, ast.Name) and test.id == "_IS_MLX":
-                        for sub in ast.walk(
-                            ast.Module(body = stmt.orelse, type_ignores = [])
-                        ):
+                        for sub in ast.walk(ast.Module(body = stmt.orelse, type_ignores = [])):
                             if isinstance(sub, ast.Assign) and any(
                                 isinstance(t, ast.Name) and t.id == "save_method"
                                 for t in sub.targets

tests/studio/test_frontend_dep_removal.py

@@ -90,8 +90,7 @@ CASES: list[Case] = [
     ),
     Case(
         "C8",
-        "multi-remove with mixed safety: next-themes + "
-        "@huggingface/hub + dexie all unsafe",
+        "multi-remove with mixed safety: next-themes + @huggingface/hub + dexie all unsafe",
         ["next-themes", "@huggingface/hub", "dexie"],
         "FAIL",
         ["next-themes", "@huggingface/hub", "dexie"],
@@ -119,8 +118,7 @@ CASES: list[Case] = [
     ),
     Case(
         "C12",
-        "moving @hugeicons/react from deps to devDeps is NOT a "
-        "removal (still declared)",
+        "moving @hugeicons/react from deps to devDeps is NOT a removal (still declared)",
         [],
         "PASS",
         [],
@@ -136,7 +134,7 @@ CASES: list[Case] = [
     ),
     Case(
         "C14",
-        "removing dexie breaks src imports (no other declared " "dep needs it)",
+        "removing dexie breaks src imports (no other declared dep needs it)",
         ["dexie"],
         "FAIL",
         ["dexie"],
@@ -151,14 +149,14 @@ CASES: list[Case] = [
     ),
     Case(
         "C16",
-        "removing canvas-confetti (imported in confetti.tsx); " "no transitive parent",
+        "removing canvas-confetti (imported in confetti.tsx); no transitive parent",
         ["canvas-confetti"],
         "FAIL",
         ["canvas-confetti"],
     ),
     Case(
         "C17",
-        "removing recharts (imported in chart.tsx); no transitive " "parent",
+        "removing recharts (imported in chart.tsx); no transitive parent",
         ["recharts"],
         "FAIL",
         ["recharts"],
@@ -173,36 +171,35 @@ CASES: list[Case] = [
     ),
     Case(
         "C19",
-        "removing node-forge (imported in providers-api.ts); " "no transitive parent",
+        "removing node-forge (imported in providers-api.ts); no transitive parent",
         ["node-forge"],
         "FAIL",
         ["node-forge"],
     ),
     Case(
         "C20",
-        "removing @tauri-apps/api is safe: all 5 @tauri-apps "
-        "plugins declare it as a direct dep",
+        "removing @tauri-apps/api is safe: all 5 @tauri-apps plugins declare it as a direct dep",
         ["@tauri-apps/api"],
         "PASS",
         [],
     ),
     Case(
         "C21",
-        "removing mammoth (imported in runtime-provider.tsx); " "no transitive parent",
+        "removing mammoth (imported in runtime-provider.tsx); no transitive parent",
         ["mammoth"],
         "FAIL",
         ["mammoth"],
     ),
     Case(
         "C22",
-        "removing unpdf (imported in runtime-provider.tsx); " "no transitive parent",
+        "removing unpdf (imported in runtime-provider.tsx); no transitive parent",
         ["unpdf"],
         "FAIL",
         ["unpdf"],
     ),
     Case(
         "C23",
-        "removing remark-gfm is safe: streamdown declares it " "as a direct dep",
+        "removing remark-gfm is safe: streamdown declares it as a direct dep",
         ["remark-gfm"],
         "PASS",
         [],
@@ -310,9 +307,7 @@ def run_case(case: Case, head_pkg: dict) -> tuple[bool, str]:
         if in_summary and line.strip().startswith("- "):
             failure_pkgs.append(line.strip()[2:])
 
-    ok = actual_status == case.expected_status and set(failure_pkgs) == set(
-        case.expected_failures
-    )
+    ok = actual_status == case.expected_status and set(failure_pkgs) == set(case.expected_failures)
     return ok, (
         f"expected: status={case.expected_status} fails={sorted(case.expected_failures)}\n"
         f"actual:   status={actual_status} fails={sorted(failure_pkgs)}\n"
@@ -320,14 +315,10 @@ def run_case(case: Case, head_pkg: dict) -> tuple[bool, str]:
     )
 
 
-# ---------------------------------------------------------------------------
-# Classifier unit tests: feed hand-crafted snippets directly into classify()
-# and assert the returned kind. Covers sneaky import shapes that an
-# adversarial / careless dev might use to obscure a real usage.
-# ---------------------------------------------------------------------------
+# Classifier unit tests: feed hand-crafted snippets into classify() and assert
+# the returned kind. Covers sneaky import shapes used to obscure a real usage.
 
-# Import the script's classify() by file path so this test does not need
-# the package to be installed.
+# Import classify() by file path so this test needs no installed package.
 import importlib.util as _ilu
 
 _spec = _ilu.spec_from_file_location("_dep_check", str(SCRIPT))
@@ -532,8 +523,7 @@ CLASSIFY_CASES: list[ClassifyCase] = [
     ),
     ClassifyCase(
         "U23",
-        "package name in Python file (ignored, "
-        "Python can never import npm packages)",
+        "package name in Python file (ignored, Python can never import npm packages)",
         "playwright",
         "tests/x.py",
         'label: str = "playwright"',
@@ -745,11 +735,9 @@ def run_classify_unit_tests() -> int:
     return 0 if passed == len(CLASSIFY_CASES) else 1
 
 
-# ---------------------------------------------------------------------------
-# Adversarial end-to-end cases: drop a sneaky synthetic file into src/,
-# run the checker, then clean up. Catches the case where pattern detection
-# regresses for a real grep+classify pipeline (not just classify in isolation).
-# ---------------------------------------------------------------------------
+# Adversarial end-to-end cases: drop a sneaky synthetic file into src/, run the
+# checker, then clean up. Catches detection regressions in the full
+# grep+classify pipeline (not just classify in isolation).
 
 ADVERSARIAL_TMP_DIR = REPO / "studio/frontend/src/__dep_check_adversarial__"
 
@@ -854,7 +842,7 @@ ADV_CASES: list[AdvCase] = [
     ),
     AdvCase(
         "A10",
-        "package referenced only in a Python file should " "NOT trigger a JS FAIL",
+        "package referenced only in a Python file should NOT trigger a JS FAIL",
         "adv10.py",
         'label = "__adv_only_pkg_j__"\n',
         "__adv_only_pkg_j__",
@@ -863,8 +851,7 @@ ADV_CASES: list[AdvCase] = [
     ),
     AdvCase(
         "A11",
-        "package mentioned in a markdown doc file is "
-        "ignored by JS-like-only string_literal",
+        "package mentioned in a markdown doc file is ignored by JS-like-only string_literal",
         "adv11.md",
         "See [docs](https://example.com/__adv_only_pkg_k__).\n",
         "__adv_only_pkg_k__",
@@ -901,14 +888,9 @@ ADV_CASES: list[AdvCase] = [
 ]
 
 
-# ---------------------------------------------------------------------------
 # package.json field-reference cases: simulate `prettier: "@x/config"`,
-# `eslintConfig.extends`, `overrides`, `peerDependenciesMeta`, etc.
-# These test the package_json_extra_refs() coverage. Cross-checked against
-# the patterns used by Tailwind, Stylelint, Prettier, Next.js, Astro,
-# TypeScript, ESLint, SvelteKit, Storybook, Vite, and TanStack/Query
-# manifests.
-# ---------------------------------------------------------------------------
+# `eslintConfig.extends`, `overrides`, `peerDependenciesMeta`, etc., testing
+# package_json_extra_refs() coverage across common tool manifests.
 
 
 @dataclass
@@ -1117,9 +1099,7 @@ def run_pkg_field_cases() -> int:
         finally:
             os.unlink(base_path)
             os.unlink(head_path)
-        actual_status = {0: "PASS", 1: "FAIL"}.get(
-            proc.returncode, f"RC{proc.returncode}"
-        )
+        actual_status = {0: "PASS", 1: "FAIL"}.get(proc.returncode, f"RC{proc.returncode}")
         fails: list[str] = []
         in_summary = False
         for line in proc.stdout.splitlines():
@@ -1130,15 +1110,11 @@ def run_pkg_field_cases() -> int:
                 fails.append(line.strip()[2:])
         # The expected_failures includes the tolerated-FP case (P15); we
         # accept BOTH expected_status and expected_failures matches.
-        ok = actual_status == pc.expected_status and set(fails) == set(
-            pc.expected_failures
-        )
+        ok = actual_status == pc.expected_status and set(fails) == set(pc.expected_failures)
         mark = "PASS" if ok else "FAIL"
         print(f"  [{mark}] {pc.id}: {pc.desc}")
         if not ok:
-            print(
-                f"      expected: status={pc.expected_status} fails={pc.expected_failures}"
-            )
+            print(f"      expected: status={pc.expected_status} fails={pc.expected_failures}")
             print(f"      actual:   status={actual_status} fails={fails}")
             for ln in proc.stdout.splitlines()[:25]:
                 print(f"      {ln}")
@@ -1184,9 +1160,7 @@ def run_adversarial_cases() -> int:
                 )
             finally:
                 os.unlink(base_path)
-            actual_status = {0: "PASS", 1: "FAIL"}.get(
-                proc.returncode, f"RC{proc.returncode}"
-            )
+            actual_status = {0: "PASS", 1: "FAIL"}.get(proc.returncode, f"RC{proc.returncode}")
             fails = []
             in_summary = False
             for line in proc.stdout.splitlines():
@@ -1195,15 +1169,11 @@ def run_adversarial_cases() -> int:
                     continue
                 if in_summary and line.strip().startswith("- "):
                     fails.append(line.strip()[2:])
-            ok = actual_status == ac.expected_status and set(fails) == set(
-                ac.expected_failures
-            )
+            ok = actual_status == ac.expected_status and set(fails) == set(ac.expected_failures)
             mark = "PASS" if ok else "FAIL"
             print(f"  [{mark}] {ac.id}: {ac.desc}")
             if not ok:
-                print(
-                    f"      expected: status={ac.expected_status} fails={ac.expected_failures}"
-                )
+                print(f"      expected: status={ac.expected_status} fails={ac.expected_failures}")
                 print(f"      actual:   status={actual_status} fails={fails}")
                 for ln in proc.stdout.splitlines()[:20]:
                     print(f"      {ln}")
@@ -1224,9 +1194,7 @@ def run_adversarial_cases() -> int:
     return 0 if passed == len(ADV_CASES) else 1
 
 
-# ---------------------------------------------------------------------------
-# Dead-dep enumeration cases.
-# ---------------------------------------------------------------------------
+# Dead-dep enumeration cases
 
 
 @dataclass
@@ -1385,9 +1353,7 @@ def run_enum_cases() -> int:
         if not ok:
             print(f"      expected unused superset: {sorted(ec.expected_unused)}")
             print(f"      expected used NOT in unused: {sorted(ec.expected_used)}")
-            print(
-                f"      expected orphans superset: {sorted(ec.expected_orphan_types)}"
-            )
+            print(f"      expected orphans superset: {sorted(ec.expected_orphan_types)}")
             print(f"      actual unused: {sorted(unused)}")
             print(f"      actual orphans: {sorted(orphans)}")
             for ln in proc.stdout.splitlines()[:30]:
@@ -1404,14 +1370,9 @@ def run_enum_cases() -> int:
     return 0 if passed == len(ENUM_CASES) else 1
 
 
-# ---------------------------------------------------------------------------
-# Script-wrapper cases: exercise scripts_bin_refs / _next_real_bin so a
-# package.json script like `cross-env CI=1 biome check` correctly credits
-# `@biomejs/biome` rather than the wrapper itself. The 10x reviewer flagged
-# the original "first non-env token" heuristic as too narrow: any project
-# using cross-env / dotenv / dotenvx / env-cmd / a quoted env value would
-# bypass the bin-name check.
-# ---------------------------------------------------------------------------
+# Script-wrapper cases: exercise scripts_bin_refs / _next_real_bin so a script
+# like `cross-env CI=1 biome check` credits `@biomejs/biome`, not the wrapper.
+# The original "first non-env token" heuristic missed cross-env / dotenv / etc.
 
 
 @dataclass

tests/studio/test_hardware_dispatch_matrix.py

@@ -1,35 +1,12 @@
 # SPDX-License-Identifier: AGPL-3.0-only
-"""
-Comprehensive hardware dispatch matrix for Studio.
+"""Hardware dispatch matrix for Studio.
 
-Drives every supported hardware profile from a single test host by
-spoofing platform / torch.cuda / torch.xpu / sys.modules['mlx'] so we
-can exercise the CUDA, ROCm, XPU, MLX, and CPU dispatch paths
-deterministically without real hardware.
-
-Profiles checked:
-
-    nvidia_cuda          Linux x86_64 + torch.cuda.is_available()=True,
-                         torch.version.hip=None
-    amd_rocm             Linux x86_64 + torch.cuda.is_available()=True,
-                         torch.version.hip="6.1"   (PyTorch ROCm aliases
-                         torch.cuda.* over HIP)
-    intel_xpu            Linux x86_64 + torch.cuda off, torch.xpu.is_available()=True
-    apple_silicon_mlx    Darwin arm64 + cuda off + xpu off + mlx importable
-    apple_silicon_no_mlx Darwin arm64 + everything off (no mlx pkg)
-    linux_arm64_with_mlx Linux arm64 + mlx importable -- gate must NOT activate
-                         (canary against accidental Linux-arm64 hijack)
-    cpu_only             Linux x86_64 + nothing -- pure CPU fallback
-
-For each profile we assert three contracts:
-
-  1. ``unsloth._IS_MLX`` (re-evaluated under the spoof).
-  2. ``utils.hardware.detect_hardware()`` ``DeviceType`` and ``IS_ROCM``.
-  3. ``utils.hardware.is_apple_silicon()``.
-
-Add a row to ``PROFILES`` to extend coverage; tests parametrize over it
-automatically. No real hardware required.
-"""
+Spoofs platform / torch.cuda / torch.xpu / sys.modules['mlx'] to exercise the
+CUDA, ROCm, XPU, MLX, and CPU dispatch paths deterministically without real
+hardware. Each profile in ``PROFILES`` (CUDA, ROCm, XPU, Apple+/-mlx, the
+linux-arm64-with-mlx canary, CPU) asserts three contracts: ``unsloth._IS_MLX``,
+``detect_hardware()`` DeviceType + ``IS_ROCM``, and ``is_apple_silicon()``.
+Add a row to ``PROFILES`` to extend coverage; tests parametrize over it."""
 
 from __future__ import annotations
 
@@ -60,17 +37,13 @@ class HardwareProfile:
     system: str  # platform.system() value
     machine: str  # platform.machine() value
     cuda_available: bool  # torch.cuda.is_available() value
-    hip_version: Optional[
-        str
-    ]  # torch.version.hip; None for NVIDIA, "6.1" etc. for ROCm
+    hip_version: Optional[str]  # torch.version.hip; None for NVIDIA, "6.1" etc. for ROCm
     xpu_available: bool  # torch.xpu.is_available() value
     has_mlx: bool  # whether to inject a fake mlx into sys.modules
     mps_available: bool  # torch.backends.mps.is_available() value
 
     expect_is_mlx: bool  # unsloth._IS_MLX
-    expect_device_type: (
-        str  # Studio DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU")
-    )
+    expect_device_type: str  # Studio DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU")
     expect_is_rocm: bool  # Studio IS_ROCM
     expect_apple_silicon: bool  # Studio is_apple_silicon()
     extra_notes: str = ""
@@ -193,10 +166,7 @@ PROFILE_IDS = [p.name for p in PROFILES]
 @pytest.fixture
 def spoof_hardware(monkeypatch):
     """Return a function that applies a HardwareProfile to the live process.
-
-    Idempotent: each call re-applies the profile. Cleanup happens
-    automatically when the test exits via monkeypatch.
-    """
+    Idempotent; monkeypatch cleans up on test exit."""
 
     def _apply(profile: HardwareProfile) -> None:
         import platform
@@ -206,12 +176,9 @@ def spoof_hardware(monkeypatch):
         monkeypatch.setattr(platform, "system", lambda: profile.system)
         monkeypatch.setattr(platform, "machine", lambda: profile.machine)
 
-        # torch.cuda.is_available
         monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available)
-        # detect_hardware reads torch.cuda.get_device_properties(0).name when
-        # cuda_available is True. On a CPU CI runner that triggers _cuda_init
-        # and crashes with "No CUDA GPUs are available". Stub it so the
-        # dispatch path under test runs end-to-end.
+        # Stub get_device_properties: detect_hardware reads .name when CUDA is
+        # available, which crashes on a CPU CI runner ("No CUDA GPUs").
         if profile.cuda_available:
             stub_props = types.SimpleNamespace(
                 name = "Stub GPU" if not profile.hip_version else "Stub AMD GPU",
@@ -223,17 +190,14 @@ def spoof_hardware(monkeypatch):
                 raising = False,
             )
 
-        # torch.version.hip — None on NVIDIA, "6.1" etc. on ROCm
+        # torch.version.hip: None on NVIDIA, "6.1" etc. on ROCm
         torch_version = torch.version
         monkeypatch.setattr(torch_version, "hip", profile.hip_version, raising = False)
 
-        # torch.xpu.is_available + get_device_name -- detect_hardware reads both.
-        # Real torch.xpu.get_device_name requires the XPU-compiled torch build,
-        # so always stub it under the spoof to keep tests hardware-agnostic.
+        # Stub torch.xpu.* (detect_hardware reads both); real get_device_name
+        # needs the XPU torch build, so always stub to stay hardware-agnostic.
         if hasattr(torch, "xpu"):
-            monkeypatch.setattr(
-                torch.xpu, "is_available", lambda: profile.xpu_available
-            )
+            monkeypatch.setattr(torch.xpu, "is_available", lambda: profile.xpu_available)
             monkeypatch.setattr(
                 torch.xpu,
                 "get_device_name",
@@ -249,9 +213,7 @@ def spoof_hardware(monkeypatch):
 
         # torch.backends.mps.is_available
         if hasattr(torch.backends, "mps"):
-            monkeypatch.setattr(
-                torch.backends.mps, "is_available", lambda: profile.mps_available
-            )
+            monkeypatch.setattr(torch.backends.mps, "is_available", lambda: profile.mps_available)
 
         # mlx + mlx.core in sys.modules
         if profile.has_mlx:
@@ -263,8 +225,7 @@ def spoof_hardware(monkeypatch):
             monkeypatch.setitem(sys.modules, "mlx", fake_mlx)
             monkeypatch.setitem(sys.modules, "mlx.core", fake_mlx_core)
         else:
-            # Drop any cached mlx modules and patch find_spec so the
-            # unsloth gate (which uses importlib.util.find_spec) sees
+            # Drop cached mlx and patch find_spec so the unsloth gate sees
             # mlx as absent.
             monkeypatch.delitem(sys.modules, "mlx", raising = False)
             monkeypatch.delitem(sys.modules, "mlx.core", raising = False)
@@ -277,23 +238,24 @@ def spoof_hardware(monkeypatch):
 
             monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)
 
-            # Studio's _has_mlx() literally does `import mlx.core`, not
-            # find_spec, so on a real Apple Silicon host with mlx
-            # genuinely installed the import would still succeed. Block
-            # it via a meta_path finder that raises ImportError for any
-            # `mlx` / `mlx.*` import while this profile is active.
+            # Studio's _has_mlx() does `import mlx.core`, not find_spec, so on a
+            # real Apple Silicon host with mlx installed it would still succeed.
+            # Block it via a meta_path finder that raises ImportError for mlx.*.
             class _BlockMLXFinder:
-                def find_spec(self_inner, name, path = None, target = None):
+                def find_spec(
+                    self_inner,
+                    name,
+                    path = None,
+                    target = None,
+                ):
                     if name == "mlx" or name.startswith("mlx."):
                         raise ImportError(
-                            f"mlx import blocked by spoof_hardware "
-                            f"(profile={profile.name})"
+                            f"mlx import blocked by spoof_hardware " f"(profile={profile.name})"
                         )
                     return None
 
             blocker = _BlockMLXFinder()
-            # Replace meta_path with a NEW list so monkeypatch can fully
-            # restore the original on teardown (mutating the list in
+            # New list so monkeypatch fully restores on teardown (mutating in
             # place would survive the test).
             monkeypatch.setattr(
                 sys,
@@ -356,8 +318,7 @@ def test_studio_detect_hardware_matches_profile(profile, spoof_hardware):
         f"got {detected!r}. {profile.extra_notes}"
     )
     assert hw.IS_ROCM is profile.expect_is_rocm, (
-        f"profile {profile.name}: expected IS_ROCM={profile.expect_is_rocm}, "
-        f"got {hw.IS_ROCM}"
+        f"profile {profile.name}: expected IS_ROCM={profile.expect_is_rocm}, " f"got {hw.IS_ROCM}"
     )
 
 
@@ -378,10 +339,8 @@ def test_studio_is_apple_silicon_matches_profile(profile, spoof_hardware):
 
 
 def test_cuda_takes_priority_over_mlx_when_both_available(spoof_hardware):
-    """If both CUDA and MLX are available, Studio MUST pick CUDA. This is the
-    canary that protects every existing GPU user from being silently routed
-    to MLX after future refactors.
-    """
+    """If both CUDA and MLX are available, Studio MUST pick CUDA: the canary
+    guarding GPU users from being silently routed to MLX after refactors."""
     profile = HardwareProfile(
         name = "cuda_plus_mlx",
         system = "Darwin",

tests/studio/test_is_mlx_dispatch_gate.py

@@ -42,16 +42,11 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
 UNSLOTH_INIT = REPO_ROOT / "unsloth" / "__init__.py"
 
 
-# ---------------------------------------------------------------------------
 # 1. Source-level structure check on _IS_MLX (no platform dependencies).
-# ---------------------------------------------------------------------------
 
 
 def test_is_mlx_gate_uses_three_required_predicates():
-    """The _IS_MLX assignment must AND together exactly the three checks
-    that Studio depends on: Darwin OS, arm64 machine, and an importable
-    mlx package. Dropping any one of them silently breaks dispatch.
-    """
+    """_IS_MLX must AND the three checks Studio depends on (Darwin, arm64, importable mlx); dropping any breaks dispatch."""
     tree = ast.parse(UNSLOTH_INIT.read_text())
 
     target = None
@@ -67,9 +62,7 @@ def test_is_mlx_gate_uses_three_required_predicates():
     assert target is not None, "_IS_MLX assignment not found in unsloth/__init__.py"
     assert isinstance(target, ast.Call), "_IS_MLX must call the shared MLX helper"
     expr_src = ast.unparse(target)
-    assert (
-        expr_src == "_is_mlx_available()"
-    ), "_IS_MLX must delegate to the shared MLX runtime gate"
+    assert expr_src == "_is_mlx_available()", "_IS_MLX must delegate to the shared MLX runtime gate"
 
     helper = None
     for node in ast.walk(tree):
@@ -97,18 +90,13 @@ def test_is_mlx_gate_uses_three_required_predicates():
     ), "_IS_MLX helper must run the local MLX precheck before importing zoo"
 
 
-# ---------------------------------------------------------------------------
 # 2. Runtime gate behavior with the platform spoofed to Apple Silicon and a
-#    fake mlx module in sys.modules.  Re-evaluates the same expression
-#    rather than reloading unsloth (which would cascade-reload torch).
-# ---------------------------------------------------------------------------
+#    fake mlx module in sys.modules. Re-evaluates the same expression rather
+#    than reloading unsloth (which would cascade-reload torch).
 
 
 def _evaluate_is_mlx_precheck(platform_module, importlib_util, os_module):
-    """Re-evaluate the local _is_mlx_available precheck using injected dependencies.
-
-    Mirrors only the cheap import barrier before unsloth imports unsloth_zoo.
-    """
+    """Re-evaluate the local _is_mlx_available precheck (the import barrier before zoo) with injected deps."""
     return (
         os_module.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1"
         and platform_module.system() == "Darwin"
@@ -166,7 +154,6 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
     if platform.system() == "Darwin" and platform.machine() == "arm64":
         # On a Mac CI runner this assertion would not apply; skip there.
         import pytest
-
         pytest.skip("Test host is Apple Silicon; CUDA-side canary doesn't apply.")
 
     import os
@@ -227,11 +214,8 @@ def test_detect_hardware_picks_cuda_on_real_host():
 
     if not torch.cuda.is_available():
         import pytest
-
         pytest.skip("No CUDA available on this host; canary not applicable.")
 
     hw = _import_studio_hardware()
     detected = hw.detect_hardware()
-    assert (
-        detected == hw.DeviceType.CUDA
-    ), f"CUDA host must dispatch to CUDA, got {detected!r}"
+    assert detected == hw.DeviceType.CUDA, f"CUDA host must dispatch to CUDA, got {detected!r}"

tests/studio/test_resolve_cuda_toolkit.ps1

@@ -0,0 +1,154 @@
+#!/usr/bin/env pwsh
+# Unit test for Resolve-CudaToolkit in studio/setup.ps1. No GPU required: the
+# detection helpers (nvidia-smi, nvcc, Find-Nvcc, ...) are stubbed so the real
+# function logic runs against spoofed Blackwell sm_120 driver/toolkit scenarios.
+#
+# The function is extracted via AST and run in a child pwsh per scenario, because
+# the -RequireOrExit path calls `exit` (which would otherwise kill this harness).
+#
+# Run: pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1
+
+$ErrorActionPreference = "Stop"
+$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1")
+$setupPath = (Resolve-Path $setupPath).Path
+
+# --- Extract the function source (not the whole installer) ---
+$tokens = $null; $errors = $null
+$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
+if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
+$fn = $ast.FindAll({ param($n)
+    $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq "Resolve-CudaToolkit"
+}, $true)
+if ($fn.Count -ne 1) { throw "expected exactly one Resolve-CudaToolkit, found $($fn.Count)" }
+$fnText = $fn[0].Extent.Text
+
+# Resolve-CudaToolkit calls Write-CudaDriverToolkitMismatch, so extract it too.
+$mismatchFn = $ast.FindAll({ param($n)
+    $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq "Write-CudaDriverToolkitMismatch"
+}, $true)
+if ($mismatchFn.Count -ne 1) { throw "expected exactly one Write-CudaDriverToolkitMismatch, found $($mismatchFn.Count)" }
+$mismatchText = $mismatchFn[0].Extent.Text
+
+# --- Spoof executables for driver/toolkit compatibility scenarios ---
+$work = Join-Path ([System.IO.Path]::GetTempPath()) ("rct_" + [guid]::NewGuid().ToString("N"))
+New-Item -ItemType Directory -Force -Path $work | Out-Null
+$smiMajorMismatchFake = Join-Path $work "nvidia-smi-12.9.ps1"
+$smiSameMajorFake     = Join-Path $work "nvidia-smi-13.2.ps1"
+$nvccIncompatibleFake = Join-Path $work "nvcc-13.3.ps1"
+$nvccCompatibleFake   = Join-Path $work "nvcc-12.8.ps1"
+Set-Content -LiteralPath $smiMajorMismatchFake -Value "'CUDA Version: 12.9'"
+Set-Content -LiteralPath $smiSameMajorFake     -Value "'CUDA Version: 13.2'"
+Set-Content -LiteralPath $nvccIncompatibleFake -Value "'Cuda compilation tools, release 13.3, V13.3.0'"
+Set-Content -LiteralPath $nvccCompatibleFake   -Value "'Cuda compilation tools, release 12.8, V12.8.0'"
+
+$failures = 0
+function Check($name, $cond) {
+    if ($cond) { Write-Host "  PASS  $name" }
+    else { Write-Host "  FAIL  $name" -ForegroundColor Red; $script:failures++ }
+}
+
+# Build + run one scenario in a child pwsh; returns @{ Exit; Out }.
+function Run-Case {
+    param([string]$FindMode, [bool]$Require, [string]$DriverMode = "major-mismatch")
+    $requireLit = if ($Require) { '$true' } else { '$false' }
+    $smiForCase = if ($DriverMode -eq "same-major") { $smiSameMajorFake } else { $smiMajorMismatchFake }
+    $child = @"
+`$ErrorActionPreference = 'Continue'
+[Environment]::SetEnvironmentVariable('CUDA_PATH', `$null, 'Process')
+`$FindNvccMode = '$FindMode'
+`$NvccIncompatibleFake = '$nvccIncompatibleFake'
+`$NvccCompatibleFake = '$nvccCompatibleFake'
+function substep { param(`$m, `$c) Write-Host "  `$m" }
+function step    { param(`$l, `$v, `$c) Write-Host "[`$l] `$v" }
+function Add-ToUserPath { param(`$Directory, `$Position) `$true }
+function Refresh-Environment { }
+function Get-CudaComputeCapability { '120' }
+function Test-NvccArchSupport { param(`$NvccExe, `$Arch) `$true }
+function Get-NvccMaxArch { param(`$NvccExe) '120' }
+`$script:WingetCalled = `$false
+function winget { `$script:WingetCalled = `$true; 'no matching versions' }
+function Find-Nvcc {
+    param([string]`$MaxVersion = '')
+    switch (`$FindNvccMode) {
+        'compatible'   { return `$NvccCompatibleFake }
+        'same-major'   { return `$NvccIncompatibleFake }
+        'incompatible' { if (`$MaxVersion) { return `$null } else { return `$NvccIncompatibleFake } }
+        default        { return `$null }
+    }
+}
+`$NvidiaSmiExe  = '$smiForCase'
+`$VsInstallPath = `$null
+`$HasNvidiaSmi  = `$true
+`$script:CudaToolkitReady = `$false
+`$script:NvccPath = `$null; `$script:CudaToolkitRoot = `$null; `$script:CudaArch = `$null
+
+$mismatchText
+
+$fnText
+
+if ($requireLit) { Resolve-CudaToolkit -RequireOrExit } else { Resolve-CudaToolkit }
+Write-Host ("RESULT ready={0} nvcc={1} winget={2}" -f `$script:CudaToolkitReady, `$script:NvccPath, `$script:WingetCalled)
+"@
+    $childFile = Join-Path $work ("case_" + [guid]::NewGuid().ToString("N") + ".ps1")
+    Set-Content -LiteralPath $childFile -Value $child
+    $out = & pwsh -NoProfile -File $childFile 2>&1 | Out-String
+    return @{ Exit = $LASTEXITCODE; Out = $out }
+}
+
+try {
+    Write-Host "Scenario 1: prebuilt path, newer-major toolkit (no -RequireOrExit) -> defers, no exit"
+    $r = Run-Case -FindMode "incompatible" -Require $false
+    Check "exits 0 (not blocked)"        ($r.Exit -eq 0)
+    Check "CudaToolkitReady = false"     ($r.Out -match "ready=False")
+    Check "winget NOT called"            ($r.Out -match "winget=False")
+    Check "explains major mismatch"      ($r.Out -match "major-version mismatch")
+    Check "does not blame the toolkit"   (-not ($r.Out -match "INCOMPATIBLE"))
+
+    Write-Host "Scenario 2: forced source build, newer-major toolkit (-RequireOrExit) -> hard exit"
+    $r = Run-Case -FindMode "incompatible" -Require $true
+    Check "exits non-zero"               ($r.Exit -ne 0)
+    Check "explains major mismatch"      ($r.Out -match "major-version mismatch")
+    Check "one-line source-build error"  ($r.Out -match "CUDA source build cannot use the installed toolkit")
+
+    Write-Host "Scenario 3: same-major newer-minor toolkit (-RequireOrExit) -> resolves, env set"
+    $r = Run-Case -FindMode "same-major" -Require $true -DriverMode "same-major"
+    Check "exits 0"                      ($r.Exit -eq 0)
+    Check "CudaToolkitReady = true"      ($r.Out -match "ready=True")
+    Check "NvccPath published"           ($r.Out -match "nvcc=.*nvcc-13\.3")
+    Check "no mismatch warning"          (-not ($r.Out -match "major-version mismatch"))
+
+    Write-Host "Scenario 4: compatible older-major toolkit (-RequireOrExit) -> resolves, env set"
+    $r = Run-Case -FindMode "compatible" -Require $true
+    Check "exits 0"                      ($r.Exit -eq 0)
+    Check "CudaToolkitReady = true"      ($r.Out -match "ready=True")
+    Check "NvccPath published"           ($r.Out -match "nvcc=.*nvcc")
+
+    Write-Host "Scenario 5: no toolkit, prebuilt path (no -RequireOrExit) -> defers, no winget"
+    $r = Run-Case -FindMode "none" -Require $false
+    Check "exits 0"                      ($r.Exit -eq 0)
+    Check "CudaToolkitReady = false"     ($r.Out -match "ready=False")
+    Check "winget NOT called"            ($r.Out -match "winget=False")
+
+    Write-Host "Scenario 6: no toolkit, forced (-RequireOrExit) -> winget attempted then exit"
+    # The function exits before the RESULT line here, so assert on the winget-block
+    # marker in output rather than the flag.
+    $r = Run-Case -FindMode "none" -Require $true
+    Check "winget attempted"             ($r.Out -match "installing via winget")
+    Check "exits non-zero"               ($r.Exit -ne 0)
+    Check "preserved nvcc-required error" ($r.Out -match "CUDA Toolkit \(nvcc\) is required")
+
+    Write-Host "Scenario 7: same-major toolkit only on PATH, missed by -MaxVersion (-RequireOrExit) -> accepted, not rejected"
+    # -MaxVersion misses it (not in side-by-side base) but plain Find-Nvcc finds it on PATH: must be used.
+    $r = Run-Case -FindMode "incompatible" -Require $true -DriverMode "same-major"
+    Check "exits 0"                      ($r.Exit -eq 0)
+    Check "CudaToolkitReady = true"      ($r.Out -match "ready=True")
+    Check "NvccPath published"           ($r.Out -match "nvcc=.*nvcc-13\.3")
+    Check "no mismatch warning"          (-not ($r.Out -match "major-version mismatch"))
+}
+finally {
+    Remove-Item -Recurse -Force -LiteralPath $work -ErrorAction SilentlyContinue
+}
+
+Write-Host ""
+if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
+Write-Host "All checks passed" -ForegroundColor Green

tests/studio/test_stream_cancel_registration_timing.py

@@ -1,43 +1,19 @@
-"""
-Tests that the cancel tracker is registered BEFORE StreamingResponse is
-returned, and that cleanup runs via a `finally` block inside each
-async generator.
+"""Tests that the cancel tracker is registered BEFORE StreamingResponse is
+returned and that cleanup runs in a `finally` inside each async generator.
 
-The zombie-generation scenario is: user clicks Stop during prefill /
-warmup / proxy buffering, before the first SSE chunk. If _tracker
-__enter__ lives inside the async generator body, the registry is empty
-at the moment /api/inference/cancel lands -- so cancel returns 0 and
-the decode runs to completion.
+Zombie scenario: Stop during prefill/warmup/proxy buffering (before the first
+SSE chunk). If _tracker.__enter__ ran inside the generator body, the registry
+would be empty when /api/inference/cancel lands, so cancel returns 0 and decode
+runs to completion. The fix registers in the sync body of
+openai_chat_completions and cleans up in each generator's `finally` -- a
+BackgroundTask would be skipped when stream_response raises.
 
-The fix moves _tracker = _TrackedCancel(...) and _tracker.__enter__()
-to the synchronous body of openai_chat_completions (before the
-StreamingResponse is returned) and places _tracker.__exit__ inside
-each generator's `finally` block. Using a generator `finally` (rather
-than a Starlette BackgroundTask) guarantees cleanup on every
-termination path -- normal exhaustion, CancelledError from
-ClientDisconnect, and OSError / BrokenPipeError during send() --
-because Starlette skips `background` callbacks when stream_response
-raises.
-
-Structural verifies:
-  - No `async def ...:` body contains `_tracker.__enter__()` in
-    routes/inference.py (registration moved to sync body).
-  - Each of the four async generators (gguf_tool_stream,
-    gguf_stream_chunks, stream_chunks, audio_input_stream) contains
-    `_tracker.__exit__(None, None, None)` inside a try/finally block.
-  - No StreamingResponse in openai_chat_completions passes
-    `background=` (cleanup now lives in the generator finally).
-
-Behavioral verifies (extracting `_TrackedCancel` from source and
-exercising the actual runtime semantics):
-  - `finally: _tracker.__exit__(...)` runs on normal completion,
-    mid-stream exception (OSError / BrokenPipeError from send()),
-    and aclose() from Starlette ClientDisconnect.
-  - A pre-set cancel_event (from `_TrackedCancel.__enter__` replaying
-    a pending cancel POST) lets the GGUF while-loop break cleanly
-    and emit final_chunk + [DONE] instead of propagating
-    `GeneratorExit` out of `_stream_with_retry` into the async
-    generator's `except Exception` (which would not catch it).
+Structural verifies: no `_tracker.__enter__()` inside the async generators;
+each of the four generators has `_tracker.__exit__(...)` in a try/finally; no
+StreamingResponse passes `background=`. Behavioral verifies (running the
+extracted `_TrackedCancel`): finally cleanup on normal completion, mid-stream
+OSError, and aclose() from ClientDisconnect; and a pre-set cancel_event lets
+the GGUF loop break cleanly emitting final_chunk + [DONE].
 """
 
 from __future__ import annotations
@@ -49,13 +25,7 @@ import time
 from pathlib import Path
 
 
-SOURCE_PATH = (
-    Path(__file__).resolve().parents[2]
-    / "studio"
-    / "backend"
-    / "routes"
-    / "inference.py"
-)
+SOURCE_PATH = Path(__file__).resolve().parents[2] / "studio" / "backend" / "routes" / "inference.py"
 SRC = SOURCE_PATH.read_text()
 _TREE = ast.parse(SRC)
 
@@ -334,11 +304,9 @@ def test_finally_cleanup_on_aclose():
 
 
 def test_preset_cancel_event_exits_cleanly_with_done():
-    # Pending-replay: POST /cancel arrived before the stream registered,
-    # was stashed, then consumed by _TrackedCancel.__enter__ which set
-    # cancel_event. The generator must break out of the loop cleanly
-    # and emit final_chunk + [DONE] rather than calling next(gen) and
-    # propagating `GeneratorExit` out of the GGUF stream wrapper.
+    # Pending-replay: a stashed cancel set cancel_event via __enter__. The
+    # generator must break cleanly and emit final_chunk + [DONE] rather than
+    # calling next(gen) and propagating GeneratorExit out of the GGUF wrapper.
     ev = threading.Event()
     ev.set()
     chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev)))
@@ -355,13 +323,7 @@ def test_normal_path_streams_all_tokens():
     # when cancel_event is unset.
     ev = threading.Event()
     chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev)))
-    assert chunks == [
-        "first_chunk",
-        "cumulative-1",
-        "cumulative-2",
-        "final_chunk",
-        "[DONE]",
-    ]
+    assert chunks == ["first_chunk", "cumulative-1", "cumulative-2", "final_chunk", "[DONE]"]
 
 
 def test_cancel_during_streaming_stops_iteration_promptly():
@@ -477,11 +439,9 @@ def test_audio_input_stream_offloads_blocking_next_to_thread():
 
 
 def test_stream_chunks_cancel_branch_resets_backend_state():
-    # The Unsloth path's cancel branch must flush GPU / KV-cache state
-    # via `backend.reset_generation_state()` -- the orchestrator's
-    # internal cancel path does not do this, so a cancel-via-POST that
-    # only broke the loop would leave the subprocess in a dirty state
-    # for the next request.
+    # The Unsloth cancel branch must call backend.reset_generation_state() to
+    # flush GPU/KV-cache state, else a cancel-via-POST leaves the subprocess
+    # dirty for the next request.
     fn = None
     top = None
     for n in ast.walk(_TREE):
@@ -630,12 +590,9 @@ def test_audio_stream_stays_responsive_under_blocking_next():
 
 
 def test_unsloth_stream_loop_emits_zero_tokens_on_preset_cancel():
-    # Pending-cancel replay: _TrackedCancel.__enter__ already set
-    # cancel_event before the generator body starts iterating. The
-    # top-of-loop check must short-circuit the very first iteration so
-    # no token is emitted. Catches a regression that moves the check
-    # below `next()` -- the mid-loop test would still pass but this
-    # test would observe one extra token leak.
+    # Pending-cancel replay: cancel_event was pre-set before iterating, so the
+    # top-of-loop check must short-circuit iteration 1 (zero tokens). Catches a
+    # regression that moves the check below next() (leaks one extra token).
     cancel_event = threading.Event()
     cancel_event.set()
     reset_calls = [0]
@@ -674,8 +631,7 @@ def test_unsloth_stream_loop_emits_zero_tokens_on_preset_cancel():
         f"(pending-replay path); got {seen}"
     )
     assert next_calls[0] == 0, (
-        f"loop must not call next() at all on pre-set cancel; got "
-        f"{next_calls[0]} calls"
+        f"loop must not call next() at all on pre-set cancel; got " f"{next_calls[0]} calls"
     )
     assert reset_calls[0] == 1, (
         f"backend.reset_generation_state() must still fire exactly once "
@@ -713,6 +669,5 @@ def test_audio_stream_emits_zero_chunks_on_preset_cancel():
     seen = asyncio.run(_loop())
     assert seen == [], f"audio loop must emit zero chunks on pre-set cancel; got {seen}"
     assert next_calls[0] == 0, (
-        f"audio loop must not call next() on pre-set cancel; got "
-        f"{next_calls[0]} calls"
+        f"audio loop must not call next() on pre-set cancel; got " f"{next_calls[0]} calls"
     )

tests/studio/test_studio_gguf_export_script_pin.py

@@ -23,12 +23,7 @@ from pathlib import Path
 
 
 SOURCE_PATH = (
-    Path(__file__).resolve().parents[2]
-    / "studio"
-    / "backend"
-    / "core"
-    / "export"
-    / "export.py"
+    Path(__file__).resolve().parents[2] / "studio" / "backend" / "core" / "export" / "export.py"
 )
 SRC = SOURCE_PATH.read_text()
 TREE = ast.parse(SRC)
@@ -50,10 +45,7 @@ def _find_pin_try(tree: ast.AST):
             if (
                 isinstance(stmt, ast.ImportFrom)
                 and stmt.module == "unsloth_zoo.llama_cpp"
-                and any(
-                    alias.name == "_resolve_local_convert_script"
-                    for alias in stmt.names
-                )
+                and any(alias.name == "_resolve_local_convert_script" for alias in stmt.names)
             ):
                 return node
     return None
@@ -107,9 +99,7 @@ def test_warning_handler_gated_on_module_flag():
     try_node = _find_pin_try(TREE)
     assert try_node is not None
     handlers = [
-        h
-        for h in try_node.handlers
-        if isinstance(h.type, ast.Name) and h.type.id == "ImportError"
+        h for h in try_node.handlers if isinstance(h.type, ast.Name) and h.type.id == "ImportError"
     ]
     assert handlers
     handler = handlers[0]
@@ -117,10 +107,7 @@ def test_warning_handler_gated_on_module_flag():
     flag_writes = []
     warning_calls = []
     for node in ast.walk(ast.Module(body = handler.body, type_ignores = [])):
-        if (
-            isinstance(node, ast.Name)
-            and node.id == "_LLAMA_CPP_SCRIPTS_WARNING_EMITTED"
-        ):
+        if isinstance(node, ast.Name) and node.id == "_LLAMA_CPP_SCRIPTS_WARNING_EMITTED":
             if isinstance(node.ctx, ast.Load):
                 flag_reads.append(node)
             elif isinstance(node.ctx, ast.Store):
@@ -141,7 +128,6 @@ def test_warning_handler_gated_on_module_flag():
 
 def test_default_dir_is_string_for_setdefault_compat():
     from unsloth_zoo.llama_cpp import LLAMA_CPP_DEFAULT_DIR
-
     assert isinstance(LLAMA_CPP_DEFAULT_DIR, str)
 
 
@@ -175,10 +161,7 @@ def _simulate_pin_block(emit_records, set_value):
                 LLAMA_CPP_DEFAULT_DIR,
                 _resolve_local_convert_script,  # noqa: F401
             )
-
-            os.environ.setdefault(
-                "UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR
-            )
+            os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR)
         except ImportError:
             if not state["emitted"]:
                 emit_records.append("warned")
@@ -220,7 +203,6 @@ def test_no_warning_when_both_symbols_present(monkeypatch):
             LLAMA_CPP_DEFAULT_DIR,
             _resolve_local_convert_script,  # noqa: F401
         )
-
         os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR)
     except ImportError:
         if not state["emitted"]:

tests/studio/test_studio_text_descender_clipping.py

@@ -12,13 +12,7 @@ from pathlib import Path
 
 WORKDIR = Path(__file__).resolve().parents[2]
 MODEL_SELECTOR = (
-    WORKDIR
-    / "studio"
-    / "frontend"
-    / "src"
-    / "components"
-    / "assistant-ui"
-    / "model-selector.tsx"
+    WORKDIR / "studio" / "frontend" / "src" / "components" / "assistant-ui" / "model-selector.tsx"
 )
 APP_SIDEBAR = WORKDIR / "studio" / "frontend" / "src" / "components" / "app-sidebar.tsx"
 
@@ -37,9 +31,7 @@ def test_model_selector_trigger_label_uses_leading_tight():
     assert matches, "could not find ModelSelectorTrigger model-name span"
     for cls in matches:
         assert "leading-tight" in cls, f"expected leading-tight, got: {cls}"
-        assert (
-            "leading-none" not in cls
-        ), f"leading-none must not coexist with truncate here: {cls}"
+        assert "leading-none" not in cls, f"leading-none must not coexist with truncate here: {cls}"
 
 
 def test_sidebar_account_block_uses_leading_tight():
@@ -50,13 +42,9 @@ def test_sidebar_account_block_uses_leading_tight():
     matches = pattern.findall(src)
     assert matches, "could not find sidebar account-block parent div"
     leading_classes = [m for m in matches if m.startswith("leading-")]
-    assert (
-        leading_classes
-    ), f"no leading-* class on sidebar account-block parent: {matches}"
+    assert leading_classes, f"no leading-* class on sidebar account-block parent: {matches}"
     for cls in leading_classes:
-        assert (
-            cls == "leading-tight"
-        ), f"sidebar account-block must use leading-tight, got: {cls}"
+        assert cls == "leading-tight", f"sidebar account-block must use leading-tight, got: {cls}"
 
 
 def test_no_truncate_plus_leading_none_in_changed_files():

tests/studio/test_sync_allow_scripts_pins.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Offline tests for scripts/sync_allow_scripts_pins.py."""
+
+from __future__ import annotations
+
+import json
+import sys
+import tempfile
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(REPO_ROOT / "scripts"))
+
+import sync_allow_scripts_pins as sync  # noqa: E402
+
+
+def write_fixture(tmp: Path, policy: dict, lock_packages: dict) -> None:
+    (tmp / "package.json").write_text(
+        json.dumps(
+            {
+                "name": "fixture",
+                "dependencies": {"a": "^1.0.0"},
+                "allowScripts": policy,
+            },
+            indent = 2,
+        )
+        + "\n"
+    )
+    (tmp / "package-lock.json").write_text(
+        json.dumps(
+            {
+                "lockfileVersion": 3,
+                "packages": lock_packages,
+            }
+        )
+        + "\n"
+    )
+
+
+LOCK = {
+    "": {"name": "fixture"},
+    "node_modules/@biomejs/biome": {"version": "1.9.9", "hasInstallScript": True},
+    "node_modules/msw": {"version": "2.15.0", "hasInstallScript": True},
+    "node_modules/vite/node_modules/fsevents": {"version": "2.3.3", "hasInstallScript": True},
+    "node_modules/clean": {"version": "3.0.0"},
+}
+
+
+def test_split_spec():
+    assert sync.split_spec("fsevents") == ("fsevents", None)
+    assert sync.split_spec("@biomejs/biome") == ("@biomejs/biome", None)
+    assert sync.split_spec("@biomejs/biome@1.9.4") == ("@biomejs/biome", "1.9.4")
+    assert sync.split_spec("msw@2.14.3 || 2.15.0") == ("msw", "2.14.3 || 2.15.0")
+
+
+def test_in_sync_is_clean():
+    with tempfile.TemporaryDirectory() as d:
+        tmp = Path(d)
+        write_fixture(
+            tmp,
+            {
+                "@biomejs/biome@1.9.9": True,
+                "msw@2.15.0": True,
+                "fsevents": True,
+            },
+            LOCK,
+        )
+        assert sync.main(["--check", "--dir", str(tmp)]) == 0
+
+
+def test_stale_pin_fails_check_and_fix_repairs():
+    with tempfile.TemporaryDirectory() as d:
+        tmp = Path(d)
+        write_fixture(
+            tmp,
+            {
+                "@biomejs/biome@1.9.4": True,  # stale, bumped to 1.9.9
+                "msw@2.14.3": False,  # stale denial, bumped to 2.15.0
+                "fsevents": True,  # bare: never stale
+                "ghost@9.9.9": True,  # not in lockfile: left alone
+                "weird@*": True,  # non-exact spec: left alone
+            },
+            LOCK,
+        )
+        assert sync.main(["--check", "--dir", str(tmp)]) == 1
+        assert sync.main(["--fix", "--dir", str(tmp)]) == 0
+        got = json.loads((tmp / "package.json").read_text())["allowScripts"]
+        assert got == {
+            "@biomejs/biome@1.9.9": True,
+            "msw@2.15.0": False,  # value and key order preserved
+            "fsevents": True,
+            "ghost@9.9.9": True,
+            "weird@*": True,
+        }
+        assert list(got) == [
+            "@biomejs/biome@1.9.9",
+            "msw@2.15.0",
+            "fsevents",
+            "ghost@9.9.9",
+            "weird@*",
+        ]
+        assert sync.main(["--check", "--dir", str(tmp)]) == 0
+
+
+def test_multi_version_disjunction():
+    lock = dict(LOCK)
+    lock["node_modules/x/node_modules/msw"] = {"version": "2.14.3", "hasInstallScript": True}
+    with tempfile.TemporaryDirectory() as d:
+        tmp = Path(d)
+        write_fixture(tmp, {"msw@2.14.3": True}, lock)
+        assert sync.main(["--fix", "--dir", str(tmp)]) == 0
+        got = json.loads((tmp / "package.json").read_text())["allowScripts"]
+        assert got == {"msw@2.14.3 || 2.15.0": True}
+
+
+def test_no_policy_is_noop():
+    with tempfile.TemporaryDirectory() as d:
+        tmp = Path(d)
+        (tmp / "package.json").write_text(json.dumps({"name": "fixture"}) + "\n")
+        (tmp / "package-lock.json").write_text(json.dumps({"packages": {}}) + "\n")
+        before = (tmp / "package.json").read_text()
+        assert sync.main(["--fix", "--dir", str(tmp)]) == 0
+        assert (tmp / "package.json").read_text() == before
+
+
+def test_missing_files_is_noop():
+    with tempfile.TemporaryDirectory() as d:
+        assert sync.main(["--check", "--dir", d]) == 0
+
+
+if __name__ == "__main__":
+    failures = 0
+    for name, fn in sorted(globals().items()):
+        if name.startswith("test_") and callable(fn):
+            try:
+                fn()
+                print(f"PASS {name}")
+            except AssertionError as e:
+                failures += 1
+                print(f"FAIL {name}: {e}")
+    sys.exit(1 if failures else 0)

tests/test_callback_signature_drift.py

@@ -77,7 +77,6 @@ def _safe_parse(path: pathlib.Path):
         return _PARSE_CACHE[key]
     try:
         import warnings as _w
-
         with _w.catch_warnings():
             # Suppress SyntaxWarning emitted while parsing third-party files
             # that contain invalid escape sequences in regex / docstrings.
@@ -180,9 +179,7 @@ def _func_arity(node: ast.AST) -> tuple[int, bool] | None:
     return arity, accepts_var
 
 
-def discover_producers(
-    roots: list[pathlib.Path],
-) -> dict[str, list[tuple[pathlib.Path, int]]]:
+def discover_producers(roots: list[pathlib.Path]) -> dict[str, list[tuple[pathlib.Path, int]]]:
     """Walk every .py under each root and return {cb_list_attr: [(file, arity), ...]}."""
     producers: dict[str, list[tuple[pathlib.Path, int]]] = {}
     for root in roots:
@@ -307,7 +304,6 @@ def test_no_callback_signature_drift():
     producers = discover_producers(roots)
     if not producers:
         import pytest
-
         pytest.skip(
             "no callback producer pattern (self._*_callbacks + cb(...)) found in "
             "unsloth or unsloth_zoo. Set UNSLOTH_ZOO_SRC=<path-to-unsloth-zoo-git-checkout> "

tests/test_cli_export_unpacking.py

@@ -1,24 +1,14 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""
-Regression tests for unsloth_cli.commands.export.
+"""Regression tests for unsloth_cli.commands.export.
 
-Context: the studio export dialog live-logs work changed
-ExportOrchestrator.export_{merged_model,base_model,gguf,lora_adapter}
-to return (success, message, output_path) instead of (success, message)
-so the frontend can show the on-disk realpath on the success screen.
-The CLI at unsloth_cli/commands/export.py still unpacks two values,
-so every `unsloth export --format ...` crashes with:
-
-    ValueError: too many values to unpack (expected 2)
-
-These tests pin the CLI to the 3-tuple contract by invoking it against
-a fake ExportBackend and asserting exit_code == 0 for each --format.
-No real ML imports; the fake is installed via sys.modules injection so
-the CLI's deferred `from studio.backend.core.export import ExportBackend`
-binds to it.
-"""
+ExportOrchestrator.export_* now returns (success, message, output_path) so the
+frontend can show the realpath, but the CLI still unpacked two values, crashing
+every `unsloth export` with "too many values to unpack (expected 2)". These
+tests pin the CLI to the 3-tuple contract via a fake ExportBackend injected into
+sys.modules (the CLI's deferred import binds to it), asserting exit_code == 0
+per --format."""
 
 from __future__ import annotations
 
@@ -37,11 +27,8 @@ from typer.testing import CliRunner
 
 
 class _FakeExportBackend:
-    """Stand-in for studio.backend.core.export.ExportBackend.
-
-    All export_* methods return the new 3-tuple contract. load_checkpoint
-    keeps its 2-tuple shape (unchanged by the live-logs work).
-    """
+    """Stand-in for ExportBackend: export_* return the new 3-tuple;
+    load_checkpoint keeps its 2-tuple shape."""
 
     def __init__(self) -> None:
         self.loaded: str | None = None
@@ -67,15 +54,9 @@ class _FakeExportBackend:
 
 
 def _install_fake_studio_backend(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Inject fake studio.backend.core.export into sys.modules.
-
-    The CLI imports ExportBackend lazily inside the command function, so
-    patching sys.modules before invoking the command is sufficient to
-    steer the `from studio.backend.core.export import ExportBackend`
-    statement at the fake. Parent packages (studio, studio.backend,
-    studio.backend.core) are stubbed too so Python's import machinery
-    doesn't try to resolve the real (structlog-dependent) tree.
-    """
+    """Inject fake studio.backend.core.export into sys.modules. The CLI imports
+    ExportBackend lazily, so this steers it at the fake; parent packages are
+    stubbed too so import machinery skips the real structlog-dependent tree."""
     for name in ("studio", "studio.backend", "studio.backend.core"):
         monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
 
@@ -83,9 +64,7 @@ def _install_fake_studio_backend(monkeypatch: pytest.MonkeyPatch) -> None:
     fake_mod.ExportBackend = _FakeExportBackend
     monkeypatch.setitem(sys.modules, "studio.backend.core.export", fake_mod)
 
-    # Drop any cached import of the CLI module so the deferred import
-    # inside export() re-resolves against our fake module rather than a
-    # previously cached real one.
+    # Drop cached CLI module so export()'s deferred import re-resolves the fake.
     monkeypatch.delitem(sys.modules, "unsloth_cli.commands.export", raising = False)
 
 

tests/test_enforce_kwargs_spacing.py

@@ -0,0 +1,428 @@
+"""Tests for scripts/enforce_kwargs_spacing.py.
+
+Focus on remove_blank_after_short_import (the blank-line-after-short-import-block
+rule): it must fire on small nested import blocks, leave everything else alone,
+never change the AST, and be idempotent. A couple of enforce_spacing checks pin
+the existing kwarg-spacing behavior.
+"""
+
+from __future__ import annotations
+
+import ast
+import sys
+from pathlib import Path
+
+import pytest
+
+_SCRIPTS = str(Path(__file__).resolve().parent.parent / "scripts")
+if _SCRIPTS not in sys.path:
+    sys.path.insert(0, _SCRIPTS)
+
+from enforce_kwargs_spacing import (  # noqa: E402
+    collapse_short_asserts,
+    enforce_spacing,
+    merge_adjacent_string_literals,
+    normalize_def_trailing_comma,
+    remove_blank_after_short_import,
+)
+
+
+# (name, source) pairs where the blank after the import block MUST be removed.
+_MUST_CHANGE = {
+    "try_except_import": (
+        "def f():\n"
+        "    try:\n"
+        "        import torch\n"
+        "\n"
+        "        return torch.inference_mode\n"
+        "    except Exception:\n"
+        "        from contextlib import nullcontext\n"
+        "\n"
+        "        return nullcontext\n"
+    ),
+    "if_from_import": (
+        "def g():\n"
+        "    if cond:\n"
+        "        from . import locators\n"
+        "\n"
+        "        regions = locators.regions()\n"
+    ),
+    "multiple_consecutive_imports": ("def f():\n    import a\n    import b\n\n    return a, b\n"),
+    "type_checking_block": (
+        "def f():\n"
+        "    if TYPE_CHECKING:\n"
+        "        import x\n"
+        "\n"
+        "        y = x\n"
+        "        return y\n"
+    ),
+    "with_block": ("def f():\n    with ctx():\n        import a\n\n        return a.run()\n"),
+}
+
+# Sources that MUST be left byte-for-byte unchanged.
+_MUST_NOT_CHANGE = {
+    "module_level": 'import os\n\nVALUE = os.environ.get("V")\n',
+    "large_suite": (
+        "def f():\n"
+        "    import a\n"
+        "\n"
+        "    x = a.load()\n"
+        "    y = transform(x)\n"
+        "    return y\n"
+    ),
+    "comment_between": ("def f():\n    import a\n\n    # keep separated\n    return a.value\n"),
+    "import_is_last_stmt": "def f():\n    if cond:\n        import a\n\n",
+    "no_blank_already": "def f():\n    import a\n    return a\n",
+}
+
+
+@pytest.mark.parametrize("name", sorted(_MUST_CHANGE))
+def test_blank_removed_for_small_import_block(name):
+    src = _MUST_CHANGE[name]
+    out, changed = remove_blank_after_short_import(src)
+    assert changed is True
+    assert out != src
+    # Import and following statement now adjacent (no blank between).
+    assert "\n\n" not in out or out.count("\n\n") < src.count("\n\n")
+    # Semantics preserved and idempotent.
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+    out2, changed2 = remove_blank_after_short_import(out)
+    assert out2 == out and changed2 is False
+
+
+@pytest.mark.parametrize("name", sorted(_MUST_NOT_CHANGE))
+def test_blank_preserved_when_not_applicable(name):
+    src = _MUST_NOT_CHANGE[name]
+    out, changed = remove_blank_after_short_import(src)
+    assert changed is False
+    assert out == src
+
+
+def test_exact_output_try_block():
+    src = (
+        "def f():\n"
+        "    try:\n"
+        "        import torch\n"
+        "\n"
+        "        return torch.inference_mode\n"
+        "    except Exception:\n"
+        "        from contextlib import nullcontext\n"
+        "\n"
+        "        return nullcontext\n"
+    )
+    expected = (
+        "def f():\n"
+        "    try:\n"
+        "        import torch\n"
+        "        return torch.inference_mode\n"
+        "    except Exception:\n"
+        "        from contextlib import nullcontext\n"
+        "        return nullcontext\n"
+    )
+    out, changed = remove_blank_after_short_import(src)
+    assert changed is True
+    assert out == expected
+
+
+def test_exact_output_multiple_consecutive_imports():
+    # Pins that only the blank after the LAST import in a run is dropped and
+    # both imports are kept (the loose \n\n heuristic above does not pin this).
+    src = "def f():\n    import a\n    import b\n\n    return a, b\n"
+    expected = "def f():\n    import a\n    import b\n    return a, b\n"
+    out, changed = remove_blank_after_short_import(src)
+    assert changed is True
+    assert out == expected
+
+
+def test_multiple_blank_lines_in_gap_all_removed():
+    src = "def f():\n    import a\n\n\n    return a\n"
+    expected = "def f():\n    import a\n    return a\n"
+    out, changed = remove_blank_after_short_import(src)
+    assert changed is True
+    assert out == expected
+    out2, changed2 = remove_blank_after_short_import(out)
+    assert out2 == out and changed2 is False
+
+
+def test_multiline_import_internal_blank_preserved():
+    # A blank line INSIDE a parenthesized import is part of the import span, not
+    # the gap, so it must survive; only the trailing blank before code is dropped.
+    src = (
+        "def g():\n"
+        "    from mod import (\n"
+        "        a,\n"
+        "\n"
+        "        b,\n"
+        "    )\n"
+        "\n"
+        "    return a, b\n"
+    )
+    expected = (
+        "def g():\n"
+        "    from mod import (\n"
+        "        a,\n"
+        "\n"
+        "        b,\n"
+        "    )\n"
+        "    return a, b\n"
+    )
+    out, changed = remove_blank_after_short_import(src)
+    assert changed is True
+    assert out == expected
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+    out2, changed2 = remove_blank_after_short_import(out)
+    assert out2 == out and changed2 is False
+
+
+def test_syntax_error_is_left_alone():
+    src = "def f(:\n    import a\n\n    return a\n"
+    out, changed = remove_blank_after_short_import(src)
+    assert changed is False
+    assert out == src
+
+
+def test_enforce_spacing_pads_kwargs():
+    src = "f(a=1, b = 2)\n"
+    out, changed = enforce_spacing(src)
+    assert changed is True
+    assert "a = 1" in out and "b = 2" in out
+
+
+def test_enforce_spacing_noop_when_already_spaced():
+    src = "f(a = 1, b = 2)\n"
+    out, changed = enforce_spacing(src)
+    assert changed is False
+    assert out == src
+
+
+# ── Rule D: def one-per-line iff >= 3 params AND a default ──────────────────
+# add comma -> force one-per-line; strip comma -> stay collapsible.
+
+# Comma must be ADDED: >= 3 params, has a default, no trailing comma yet.
+_DEF_ADD = {
+    "three_with_default": "def f(a, b, c=1):\n    return a\n",
+    "four_with_default": "def f(a, b, c, d=1):\n    return a\n",
+    "kwonly_default": "def f(a, b, *, c=1):\n    return a\n",  # 3 real params, kw default
+    "continuation_default": "def f(\n    a, b, c=1\n):\n    return a\n",
+    "starred_with_default": "def f(a, b, *args, c=1):\n    return a\n",  # 4 params
+}
+
+# Comma must be STRIPPED: NOT (>=3 params and default), but a trailing comma exists.
+_DEF_STRIP = {
+    "three_no_default_multiline": "def f(\n    a,\n    b,\n    c,\n):\n    return a\n",
+    "four_no_default_multiline": "def f(\n    a,\n    b,\n    c,\n    d,\n):\n    return a\n",
+    "two_with_default": "def f(\n    a,\n    b=1,\n):\n    return a\n",  # < 3 params -> one line
+    "single_arg": "def f(\n    a,\n):\n    return a\n",
+}
+
+# Left byte-for-byte unchanged.
+_DEF_NOCHANGE = {
+    "three_no_default_oneline": "def f(a, b, c):\n    return a\n",
+    "two_with_default_oneline": "def f(a, b=1):\n    return a\n",  # < 3 -> one line, no comma
+    "noparams": "def f():\n    return 1\n",
+    "call_site": "x = foo(\n    a,\n    b,\n    c,\n    d,\n)\n",
+    "nested_default_call": "def f(a=g(1, 2,)):\n    return a\n",  # 1 param, no def comma
+    "three_default_already_comma": "def f(\n    a,\n    b,\n    c=1,\n):\n    return a\n",
+}
+
+
+@pytest.mark.parametrize("name", sorted(_DEF_ADD))
+def test_def_comma_added(name):
+    src = _DEF_ADD[name]
+    out, changed = normalize_def_trailing_comma(src)
+    assert changed is True
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+    assert out.count(",") == src.count(",") + 1
+    out2, changed2 = normalize_def_trailing_comma(out)
+    assert out2 == out and changed2 is False
+
+
+@pytest.mark.parametrize("name", sorted(_DEF_STRIP))
+def test_def_comma_stripped(name):
+    src = _DEF_STRIP[name]
+    out, changed = normalize_def_trailing_comma(src)
+    assert changed is True
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+    assert out.count(",") == src.count(",") - 1
+    out2, changed2 = normalize_def_trailing_comma(out)
+    assert out2 == out and changed2 is False
+
+
+@pytest.mark.parametrize("name", sorted(_DEF_NOCHANGE))
+def test_def_comma_unchanged(name):
+    src = _DEF_NOCHANGE[name]
+    out, changed = normalize_def_trailing_comma(src)
+    assert changed is False
+    assert out == src
+
+
+def test_def_comma_exact_output_strip_and_add():
+    # >= 3 params + default -> add comma (force one-per-line)
+    assert normalize_def_trailing_comma("def f(a, b, c=1):\n    return a\n")[0] == (
+        "def f(a, b, c=1,):\n    return a\n"
+    )
+    # 3 params, no default -> strip comma (collapsible)
+    assert (
+        normalize_def_trailing_comma("def f(\n    a,\n    b,\n    c,\n):\n    return a\n")[0]
+        == "def f(\n    a,\n    b,\n    c\n):\n    return a\n"
+    )
+
+
+# ── Rule C: merge adjacent same-line string literals ───────────────────────
+
+
+@pytest.mark.parametrize(
+    "src,expected",
+    [
+        ('x = "ab" "cd"\n', 'x = "abcd"\n'),
+        ('d = "newly-" "added dep."\n', 'd = "newly-added dep."\n'),
+        ('m = ("a. " "b.")\n', 'm = ("a. b.")\n'),
+        ('x = r"a\\n" r"b"\n', 'x = r"a\\nb"\n'),
+        ('x = "a\\"q" "b"\n', 'x = "a\\"qb"\n'),
+        # f + plain folds into one f-string (plain braces escaped).
+        ('x = f"a" "b"\n', 'x = f"ab"\n'),
+        (
+            'd = (f"{pkg}@{ver} is on the " "BLOCKED list")\n',
+            'd = (f"{pkg}@{ver} is on the BLOCKED list")\n',
+        ),
+        ('x = f"a{z}" "{lit}"\n', 'x = f"a{z}{{lit}}"\n'),
+        ('m = "plain " f"then {y}"\n', 'm = f"plain then {y}"\n'),  # plain + f
+    ],
+)
+def test_merge_adjacent_strings(src, expected):
+    out, changed = merge_adjacent_string_literals(src)
+    assert changed is True
+    assert out == expected
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+    out2, changed2 = merge_adjacent_string_literals(out)
+    assert out2 == out and changed2 is False
+
+
+@pytest.mark.parametrize(
+    "src",
+    [
+        'x = "ab"\n',  # single literal
+        "x = \"ab\" 'cd'\n",  # mixed quote style
+        'x = b"a" b"b"\n',  # bytes: left side-by-side by request
+        'x = rb"a" rb"b"\n',  # raw-bytes: also left alone
+        'm = f"a {x} " f"after {y}"\n',  # pure f + f: left side-by-side
+        'x = rf"a{z}" "b"\n',  # raw f-string: brace/backslash too subtle -> skip
+        'x = f"a{z}" "\\N{BULLET}"\n',  # named escape: AST guard rejects the fold
+        'x = (\n    "a"\n    "b"\n)\n',  # different lines, not merged
+    ],
+)
+def test_merge_adjacent_strings_skips(src):
+    out, changed = merge_adjacent_string_literals(src)
+    assert changed is False
+    assert out == src
+
+
+def test_fstring_fold_skipped_when_statement_would_not_collapse():
+    # A long f + plain assert message: folding it cannot fit the statement on one
+    # line, so ruff would re-wrap the assert condition. Leave it side-by-side.
+    src = (
+        "def f():\n"
+        "    assert some_condition_holds_here, (\n"
+        '        f"a fairly detailed message about {value} explaining " "why this failed badly"\n'
+        "    )\n"
+    )
+    out, changed = merge_adjacent_string_literals(src)
+    assert changed is False
+    assert out == src
+
+
+def test_fstring_fold_applied_when_statement_collapses():
+    # A multi-line f + plain that DOES fit on one line after folding is folded
+    # (ruff then collapses the call to a single line on the next pass).
+    src = "def f():\n    raise ValueError(\n" '        f"bad {x}: " "try again"\n' "    )\n"
+    out, changed = merge_adjacent_string_literals(src)
+    assert changed is True
+    assert 'f"bad {x}: try again"' in out
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+
+
+def test_fstring_fold_applied_inside_large_multiline_call():
+    # The fit guard only restricts asserts; an f + plain argument on its own line
+    # inside a big multi-line call (the lockfile case) folds even though the whole
+    # call cannot fit on one line.
+    src = (
+        "findings.append(\n"
+        "    Finding(\n"
+        "        path=str(path),\n"
+        "        package=key,\n"
+        '        detail=(f"{name}@{ver} is on the " "BLOCKED list"),\n'
+        "    )\n"
+        ")\n"
+    )
+    out, changed = merge_adjacent_string_literals(src)
+    assert changed is True
+    assert 'detail=(f"{name}@{ver} is on the BLOCKED list")' in out
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+
+
+# ── collapse_short_asserts: strip the magic comma holding a short assert open ──
+# The pass strips the trailing comma so ruff joins the assert onto one line on the
+# next format pass; it never changes the AST.
+
+
+@pytest.mark.parametrize(
+    "name,src",
+    [
+        (
+            "dict_eq",
+            'def t():\n    assert got == {\n        "a": 1,\n        "b": 2,\n    }\n',
+        ),
+        (
+            "list_eq",
+            'def t():\n    assert xs == [\n        "a",\n        "b",\n        "c",\n    ]\n',
+        ),
+        (
+            "membership",
+            'def t():\n    assert {\n        "type": "x",\n        "name": "y",\n    } in tools\n',
+        ),
+        (
+            "tuple_message",
+            "def t():\n    assert cond, (\n        base,\n        headers,\n    )\n",
+        ),
+        (
+            "call_args",
+            "def t():\n    assert eq(\n        a,\n        b,\n    )\n",
+        ),
+    ],
+)
+def test_collapse_short_assert_strips_trailing_comma(name, src):
+    out, changed = collapse_short_asserts(src)
+    assert changed is True
+    # The magic trailing comma is gone (so ruff will join it on the next pass).
+    assert out.count(",") == src.count(",") - 1
+    # Semantics preserved and idempotent at the strip level.
+    assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
+    out2, changed2 = collapse_short_asserts(out)
+    assert out2 == out and changed2 is False
+
+
+@pytest.mark.parametrize(
+    "name,src",
+    [
+        # one-element tuple message: stripping (only,) -> (only) changes meaning.
+        ("one_tuple_message", "def t():\n    assert cond, (\n        only,\n    )\n"),
+        # a comment inside keeps ruff multi-line, so collapsing would oscillate.
+        (
+            "comment_inside",
+            'def t():\n    assert x == {\n        "a": 1,  # keep\n        "b": 2,\n    }\n',
+        ),
+        # genuinely long: would not fit on one line, leave expanded.
+        (
+            "too_long",
+            "def t():\n    assert some_really_long_left_operand_name_here == {\n"
+            '        "alpha": 11111111,\n        "beta": 22222222,\n'
+            '        "gamma": 33333333,\n        "delta": 44444444,\n    }\n',
+        ),
+        # already one line: nothing to do.
+        ("one_line", 'def t():\n    assert got == {"a": 1, "b": 2}\n'),
+    ],
+)
+def test_collapse_short_assert_left_alone(name, src):
+    out, changed = collapse_short_asserts(src)
+    assert changed is False
+    assert out == src

tests/test_finetune_last_n_layers.py

@@ -1,16 +1,9 @@
 # Unsloth - 2x faster, 70% less memory LLM finetuning
 # Tests for the `finetune_last_n_layers` parity knob (CUDA side).
 #
-# Mirrors unsloth-zoo's `FastMLXModel.get_peft_model` parameter.
-# mlx-lm CLI's CONFIG_DEFAULTS['num_layers']=16 applies LoRA to the
-# last 16 transformer blocks only. On the CUDA path, PEFT exposes
-# `layers_to_transform` to do the same. This convenience knob fills
-# `layers_to_transform` for the user when set, matching mlx-lm CLI
-# AND unsloth-zoo's MLX path with a single config value.
-#
-# The tests intentionally avoid pulling in CUDA / a real model
-# checkpoint — they exercise only the helper that translates
-# `finetune_last_n_layers` into `layers_to_transform`.
+# Mirrors unsloth-zoo's MLX path: fills PEFT's `layers_to_transform` so LoRA
+# applies to the last N transformer blocks, matching mlx-lm CLI's num_layers.
+# Exercises only the translation helper, no CUDA / real checkpoint.
 
 from __future__ import annotations
 
@@ -47,7 +40,6 @@ def test_get_total_transformer_layers_reads_text_config():
 
 def test_get_total_transformer_layers_handles_alternative_attr_names():
     from unsloth.models.vision import _get_total_transformer_layers
-
     for attr in ("n_layer", "n_layers", "num_layers"):
         cfg = type("Cfg", (), {attr: 12})()
         model = type("M", (), {"config": cfg})()
@@ -68,7 +60,6 @@ def test_get_total_transformer_layers_returns_none_when_unknown():
 
 def test_get_total_transformer_layers_returns_none_for_missing_config():
     from unsloth.models.vision import _get_total_transformer_layers
-
     class FakeModel:
         pass
 

tests/test_gemma4_chat_template.py

@@ -23,9 +23,7 @@ def _extract_template(name):
 
 def _env():
     env = Environment(undefined = StrictUndefined, trim_blocks = False, lstrip_blocks = False)
-    env.globals["raise_exception"] = lambda msg: (_ for _ in ()).throw(
-        TemplateError(msg)
-    )
+    env.globals["raise_exception"] = lambda msg: (_ for _ in ()).throw(TemplateError(msg))
     return env
 
 

tests/test_get_model_name.py

@@ -137,9 +137,7 @@ class TestGetModelName(unittest.TestCase):
             else:
                 model_name, load_in_4bit, expected, should_change = case
                 with self.subTest(model_name = model_name, load_in_4bit = load_in_4bit):
-                    self._assert_mapping(
-                        model_name, load_in_4bit, expected, should_change
-                    )
+                    self._assert_mapping(model_name, load_in_4bit, expected, should_change)
 
     def test_static_mapper_contract(self):
         contracts = [
@@ -158,9 +156,7 @@ class TestGetModelName(unittest.TestCase):
         for src, expected in contracts:
             with self.subTest(src = src):
                 self.assertEqual(FLOAT_TO_INT_MAPPER[src], expected)
-        self.assertEqual(
-            MAP_TO_UNSLOTH_16bit["qwen/qwen3-8b-fp8"], "unsloth/Qwen3-8B-FP8"
-        )
+        self.assertEqual(MAP_TO_UNSLOTH_16bit["qwen/qwen3-8b-fp8"], "unsloth/Qwen3-8B-FP8")
 
 
 if __name__ == "__main__":