ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Reduce and tighten comments and docstrings across the test suite (#6429)

Browse source 
* Reduce and tighten comments and docstrings in tests

Shorten verbose comments and docstrings across the test suite without
changing any test logic. Remove narration that restates the next line,
collapse long module and test docstrings to a single line, and drop banner
separators. Keep regression context (issue and PR references, run ids),
skip reasons, mocking and timing rationale, license headers, lint and type
directives, and commented-out code.

Comments and docstrings only: an AST signature check confirms no code,
assertions, or string literals changed, and the suite byte-compiles cleanly.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-18 01:07:09 -07:00 • committed by GitHub
commit a6dc10dad2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
122 changed files with 1847 additions and 4511 deletions
Download patch file Download diff file Expand all files Collapse all files

23
tests/_zoo_aggressive_cuda_spoof.py
View file

@ -1,12 +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'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.
# tests/conftest.py's harness with deeper patches that unblock more patch_* /
# unsloth_zoo init paths on a GPU-less runner. Imported by every shim test
# file before any unsloth / unsloth_zoo / transformers import.
#
# 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.
# The one exception is dropping `pin_memory=True` (meaningless here), which
# downgrades a CUDA-required call to CPU-OK.
from __future__ import annotations
@ -87,12 +87,11 @@ def apply() -> None:
# 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.
# infinite-recurse. No-op them; CUDA seeding is meaningless on CPU.
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, accept anything for set;
# do NOT route through torch.{get,set}_rng_state (those touch the CPU RNG).
# rng_state APIs: return a CPU-shaped placeholder; 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)
@ -135,8 +134,7 @@ 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 since pin_memory has no meaning here.
# pin_memory drop: pin_memory=True raises on a CPU-only build; strip the kwarg.
for _name in (
"empty",
"zeros",
@ -168,8 +166,7 @@ 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 importable (newer torch handles CPU),
# else stub.
# amp.GradScaler: use the real one if importable (newer torch handles CPU), else stub.
try:
import torch.cuda.amp # type: ignore
except Exception:

36
tests/conftest.py
View file

@ -3,13 +3,10 @@
"""GPU-free test harness.
unsloth's import chain hits unsloth_zoo.device_type, which calls
get_device_type() at import time and raises NotImplementedError on CI
runners with no CUDA / XPU / HIP visible. Pre-load the real
unsloth_zoo.device_type under a temporarily-mocked
torch.cuda.is_available() so its @cache permanently captures "cuda".
On a real accelerator the pre-load is skipped and detection runs
normally.
unsloth_zoo.device_type calls get_device_type() at import time and raises
NotImplementedError on CI runners with no CUDA/XPU/HIP. Pre-load it under a
mocked torch.cuda.is_available()==True so its @cache permanently captures
"cuda"; on a real accelerator the pre-load is skipped.
Mirrors the conftest harness in unslothai/unsloth-zoo PR #624.
"""
@ -41,12 +38,9 @@ def _has_real_accelerator() -> bool:
def _preload_device_type(package: str, prereqs: tuple[str, ...] = ()) -> bool:
"""Pre-load <package>.device_type under a mocked
torch.cuda.is_available() == True so its @cache permanently
captures "cuda". prereqs lists submodule names of <package> that
must be loaded first (e.g. 'utils' for unsloth_zoo). Returns False
if the package or any prerequisite cannot be imported, in which
case the caller falls back to a stub."""
"""Pre-load <package>.device_type under a mocked is_available()==True so its
@cache captures "cuda"; prereqs are submodules to load first (e.g. 'utils').
Returns False if anything is unimportable, so the caller falls back to a stub."""
target = f"{package}.device_type"
if target in sys.modules:
return True
@ -98,11 +92,9 @@ def _preload_device_type(package: str, prereqs: tuple[str, ...] = ()) -> bool:
def _patch_torch_cuda_for_import() -> None:
"""Stub torch.cuda.* probes that fire at IMPORT time of unsloth /
unsloth_zoo when DEVICE_TYPE was forced to "cuda" above. These are
queries, not real GPU work, so returning plausible Ampere values
lets the import chain finish; tests that touch real tensors run on
CPU like normal."""
"""Stub the torch.cuda.* probes fired at import time once DEVICE_TYPE is
forced to "cuda"; returning plausible Ampere values lets the import finish
(real-tensor tests still run on CPU)."""
try:
import torch.cuda.memory as _cuda_memory # type: ignore
_cuda_memory.mem_get_info = lambda *a, **k: (0, 80 * 1024**3)
@ -140,11 +132,9 @@ if not _has_real_accelerator():
# ---------------------------------------------------------------------------
# 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.
# Apply upstream-drift fixes (vllm/triton/peft) by triggering ``import unsloth``
# (they run at import time in unsloth/import_fixes.py). The harness above lets
# the import survive CPU-only runners; the ImportError is swallowed otherwise.
# ---------------------------------------------------------------------------

40
tests/notebooks/test_validator_fixtures.py
View file

@ -1,22 +1,9 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""
Golden-fixture tests for scripts/notebook_validator.py.
"""Golden-fixture tests for scripts/notebook_validator.py: each reconstructs a broken install cell from an unslothai/notebooks PR and asserts the matching rule fires (and falls silent after the fix).
Each test reconstructs the broken-state install cell that one of the
referenced unslothai/notebooks PRs fixed, and asserts the matching rule
fires. The fixed-state tests prove the rule falls silent after the fix.
Cross-references:
PR #258 -> R-INST-003 (peft/torchao floor)
PR #260 -> R-EXC-001 (DONT_UPDATE_EXCEPTIONS coverage; covered by
an integration test pointing at a real
notebooks checkout)
PR #261a -> R-INST-004 (torch/torchcodec ABI)
PR #261b -> R-INST-005 (transformers --no-deps + tokenizers window)
PR #264 -> R-INST-005 (same class as #261b)
PR #221 -> R-INST-001 (forbid git+ HEAD installs)
51b1462 -> R-DRIFT-001 (drift; integration-tested separately)
Cross-references: PR #258->R-INST-003, #260->R-EXC-001, #261a->R-INST-004,
#261b/#264->R-INST-005, #221->R-INST-001, 51b1462->R-DRIFT-001.
"""
from __future__ import annotations
@ -32,9 +19,7 @@ sys.path.insert(0, str(SCRIPTS_DIR))
import notebook_validator as nv # noqa: E402
# Snapshot of Colab GPU pip-freeze that recreates the bug environments
# below. Real CI uses scripts/data/colab_pip_freeze.gpu.txt; tests use a
# small inline subset so the unit cases are hermetic.
# Inline subset of Colab GPU pip-freeze recreating the bug environments (CI uses scripts/data/colab_pip_freeze.gpu.txt).
COLAB_2026_05 = {
"torch": "2.10.0+cu128",
"torchao": "0.10.0",
@ -129,13 +114,11 @@ def test_r_inst_004_silent_when_torch_2_7_with_torchcodec_0_5():
def test_r_inst_005_fires_no_deps_transformers_55_without_tokenizers_pin(monkeypatch):
"""PR #264: --no-deps transformers==5.5.0 leaves Colab tokenizers in
place; if Colab ever ships tokenizers > 0.23.0 this breaks."""
"""PR #264: --no-deps transformers==5.5.0 leaves Colab tokenizers in place; breaks if Colab ships tokenizers > 0.23.0."""
cell = """%%capture
!pip install --no-deps transformers==5.5.0
"""
# Fake a Colab snapshot where tokenizers has just bumped past the window
# transformers 5.5.0 supports.
# Colab snapshot where tokenizers bumped past transformers 5.5.0's window.
colab = dict(COLAB_2026_05, tokenizers = "0.23.5")
def fake_meta(name, version):
@ -168,9 +151,7 @@ def test_r_inst_005_silent_when_no_deps_pins_tokenizers(monkeypatch):
def test_r_inst_005_silent_without_no_deps(monkeypatch):
"""If --no-deps is absent, pip resolves tokenizers transitively; the
rule must NOT fire (this is the false-positive case from notebooks like
Whisper.ipynb that pin transformers but rely on pip's resolver)."""
"""Without --no-deps, pip resolves tokenizers transitively; rule must NOT fire (false-positive case from e.g. Whisper.ipynb)."""
cell = """%%capture
!pip install transformers==4.51.3
"""
@ -259,9 +240,7 @@ def _live_notebooks_dir() -> Path | None:
reason = "unslothai/notebooks not cloned at sibling path",
)
def test_exceptions_passes_on_head():
"""L1.2 must be silent on the live HEAD of unslothai/notebooks. If this
test fires, either DONT_UPDATE_EXCEPTIONS gained a notebook missing a
policy clause (real bug) or the policy clause set is stale."""
"""L1.2 must be silent on live unslothai/notebooks HEAD; a fire means a DONT_UPDATE_EXCEPTIONS notebook lost its policy clause or the clause set is stale."""
findings = nv.rule_l12_exceptions_coverage(_live_notebooks_dir())
assert findings == [], findings
@ -271,8 +250,7 @@ def test_exceptions_passes_on_head():
reason = "unslothai/notebooks not cloned at sibling path",
)
def test_lint_smoke_no_module_errors():
"""The lint subcommand should walk every nb/kaggle without crashing.
(We accept findings -- those are the validator doing its job.)"""
"""The lint subcommand walks every nb/kaggle without crashing (findings are fine)."""
import subprocess
rc = subprocess.run(

26
tests/python/test_construct_chat_template_validation.py
View file

@ -1,15 +1,10 @@
"""Negative-path validation tests for unsloth.chat_templates.construct_chat_template.
Regression coverage for the str.find() / regex no-match guards added in
PR #5763 follow-up: missing placeholders or unrecoverable two-example
structures must raise RuntimeError with a clear message, not IndexError
or AttributeError, and must never silently drop the last character via
s[:-1].
Uses a minimal fake tokenizer so the cases run on CPU-only CI without
HF_TOKEN and without downloading a gated model. The validation paths
exercised here fail before construct_chat_template reaches any heavy
tokenizer interaction, so the stub stays small.
Regression coverage for the no-match guards added in the PR #5763 follow-up:
missing placeholders or unrecoverable two-example structures must raise
RuntimeError with a clear message (not IndexError/AttributeError) and must
not silently drop the last char via s[:-1]. A minimal fake tokenizer keeps
the cases CPU-only (no HF_TOKEN, no gated download).
"""
import pytest
@ -18,8 +13,7 @@ from unsloth.chat_templates import construct_chat_template
class _FakeTokenizer:
"""Minimum surface construct_chat_template touches before the
validation guards fire."""
"""Minimal surface construct_chat_template touches before the guards fire."""
name_or_path = "fake/tokenizer"
eos_token = "</s>"
@ -48,9 +42,8 @@ def test_missing_placeholder_in_chat_template_raises(template, expected_in_messa
def test_single_pair_template_raises_clear_error_not_attribute_error():
"""One {INPUT}/{OUTPUT} pair (rather than the required two) used to
crash with AttributeError on `found.group(1)` after the for-loop
broke without setting `found`. Must raise RuntimeError now."""
"""A single {INPUT}/{OUTPUT} pair must raise RuntimeError, not the old
AttributeError on `found.group(1)` when the loop broke without setting `found`."""
template = "user: {INPUT}\nassistant: {OUTPUT}\n"
with pytest.raises(RuntimeError):
construct_chat_template(
@ -71,7 +64,6 @@ def test_error_message_excerpt_is_bounded():
extra_eos_tokens = ["</s>"],
)
msg = str(exc_info.value)
# Excerpt is repr-quoted and capped; total message should stay well
# under the template length.
# Excerpt is capped well under the template length.
assert len(msg) < 1000
assert "{OUTPUT}" in msg

6
tests/python/test_cross_platform_parity.py
View file

@ -13,11 +13,7 @@ INSTALL_PS1 = REPO_ROOT / "install.ps1"
class TestNoTorchBackendAutoInInstallSh:
"""install.sh primary install paths must not use --torch-backend=auto.
The fallback else-branch (when TORCH_INDEX_URL is empty) is allowed to
use --torch-backend=auto since that is the last-resort recovery path.
"""
"""install.sh primary paths must not use --torch-backend=auto (only the fallback else-branch may)."""
def test_no_torch_backend_auto_outside_fallback(self):
lines = INSTALL_SH.read_text(encoding = "utf-8").splitlines()

143
tests/python/test_e2e_no_torch_sandbox.py
View file

@ -1,13 +1,4 @@
"""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 (Groups 1-6):
python -m pytest tests/python/test_e2e_no_torch_sandbox.py -v -k "not server"
# Server (Group 7, requires studio venv):
python -m pytest tests/python/test_e2e_no_torch_sandbox.py -v -m server
"""
"""E2E sandbox tests for PR #4624: lazy torch imports, CPU fallback, install.sh parsing, NO_TORCH filtering, live server."""
from __future__ import annotations
@ -23,10 +14,6 @@ from unittest import mock
import pytest
# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_DIR = REPO_ROOT / "studio"
BACKEND_DIR = STUDIO_DIR / "backend"
@ -45,17 +32,11 @@ HARDWARE_PY = HARDWARE_DIR / "hardware.py"
# Studio venv for server tests
STUDIO_VENV = Path.home() / ".unsloth" / "studio" / "unsloth_studio"
# Add studio to path for install_python_stack imports
sys.path.insert(0, str(STUDIO_DIR))
# ---------------------------------------------------------------------------
# Cross-platform helpers
# ---------------------------------------------------------------------------
def _venv_python(venv_dir: Path) -> Path:
"""Return the Python executable path for a venv, cross-platform."""
"""Return a venv's Python executable path, cross-platform."""
if sys.platform == "win32":
return venv_dir / "Scripts" / "python.exe"
return venv_dir / "bin" / "python"
@ -76,7 +57,6 @@ def _create_no_torch_venv(venv_dir: Path, python_version: str = "3.12") -> Path
py = _venv_python(venv_dir)
if not py.exists():
return None
# Verify torch is NOT importable
check = subprocess.run([str(py), "-c", "import torch"], capture_output = True)
if check.returncode == 0:
return None
@ -107,13 +87,8 @@ def _run_sh(script: str, timeout: int = 30) -> subprocess.CompletedProcess:
)
# ---------------------------------------------------------------------------
# Stub generators
# ---------------------------------------------------------------------------
def _write_loggers_stub(sandbox: Path) -> None:
"""Create a minimal loggers package stub (replaces structlog-backed real one)."""
"""Create a minimal loggers package stub (replaces the structlog-backed real one)."""
loggers_dir = sandbox / "loggers"
loggers_dir.mkdir(exist_ok = True)
(loggers_dir / "__init__.py").write_text(
@ -165,11 +140,6 @@ def _write_hardware_stub(sandbox: Path) -> None:
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture(scope = "session")
def repo_root():
return REPO_ROOT
@ -183,10 +153,7 @@ def sandbox_dir(tmp_path):
@pytest.fixture(params = ["3.12", "3.13"], scope = "module")
def no_torch_venv(request, tmp_path_factory):
"""Create a temporary uv venv with no torch.
Parametrized for 3.12 (Intel Mac default) and 3.13 (Apple Silicon/Linux).
"""
"""Temporary uv venv with no torch; 3.12 = Intel Mac default, 3.13 = Apple Silicon/Linux."""
if not _has_uv():
pytest.skip("uv not available")
@ -198,20 +165,16 @@ def no_torch_venv(request, tmp_path_factory):
return str(py)
# ===========================================================================
# Group 1: BEFORE vs AFTER -- Import Chain (6 tests)
# ===========================================================================
# Group 1: BEFORE vs AFTER -- Import Chain
class TestBeforeAfterImportChain:
"""BEFORE (PR files with a synthetic top-level torch import, simulating
main) crashes; AFTER (PR files as-is, lazy imports) works."""
"""BEFORE (synthetic top-level torch import) crashes; AFTER (lazy imports) works."""
# -- BEFORE: crashes --
def test_before_chat_templates_crashes(self, no_torch_venv, sandbox_dir):
"""BEFORE: chat_templates.py with top-level 'from torch.utils.data import
IterableDataset' crashes without torch."""
"""BEFORE: chat_templates.py with top-level IterableDataset import crashes without torch."""
source = CHAT_TEMPLATES.read_text(encoding = "utf-8")
before_source = "from torch.utils.data import IterableDataset\n" + source
@ -284,7 +247,6 @@ class TestBeforeAfterImportChain:
encoding = "utf-8",
)
# Minimal __init__.py that triggers the chain
(pkg_dir / "__init__.py").write_text(
textwrap.dedent("""\
from .format_detection import detect_dataset_format
@ -365,7 +327,6 @@ class TestBeforeAfterImportChain:
if src.exists():
shutil.copy2(src, pkg_dir / src.name)
# Minimal __init__.py
(pkg_dir / "__init__.py").write_text(
textwrap.dedent("""\
from .format_detection import detect_dataset_format, detect_custom_format_heuristic
@ -402,14 +363,11 @@ class TestBeforeAfterImportChain:
assert b"OK: full import chain succeeded" in result.stdout
# ===========================================================================
# Group 2: Dataclass Instantiation (4 tests)
# ===========================================================================
# Group 2: Dataclass Instantiation
class TestDataclassInstantiation:
"""Verify dataclass collators can be instantiated and constants accessed
without torch in an isolated venv."""
"""Dataclass collators instantiate and constants are accessible without torch."""
def test_speech_collator_instantiate(self, no_torch_venv):
"""DataCollatorSpeechSeq2SeqWithPadding(processor=None) succeeds."""
@ -485,19 +443,14 @@ class TestDataclassInstantiation:
assert result.returncode == 0, f"Failed:\n{result.stderr.decode()}"
# ===========================================================================
# Group 3: Edge Cases -- Partial/Broken Torch (4 tests)
# ===========================================================================
# Group 3: Edge Cases -- Partial/Broken Torch
class TestEdgeCasesBrokenTorch:
"""Test behavior with fake or broken torch modules on sys.path."""
"""Behavior with fake or broken torch modules on sys.path."""
def test_fake_broken_torch_module(self, no_torch_venv, sandbox_dir):
"""A fake torch that raises RuntimeError('CUDA not found') on import.
data_collators.py (no top-level torch import) should still load fine.
"""
"""Fake torch raising RuntimeError on import: data_collators.py (no top-level torch) still loads."""
torch_dir = sandbox_dir / "torch"
torch_dir.mkdir()
(torch_dir / "__init__.py").write_text(
@ -519,7 +472,7 @@ class TestEdgeCasesBrokenTorch:
assert b"OK:" in result.stdout
def test_torch_import_error_hardware_fallback(self, no_torch_venv, sandbox_dir):
"""A fake torch that raises ImportError. detect_hardware() falls back to CPU."""
"""Fake torch raising ImportError: detect_hardware() falls back to CPU."""
torch_dir = sandbox_dir / "torch"
torch_dir.mkdir()
(torch_dir / "__init__.py").write_text(
@ -546,10 +499,7 @@ class TestEdgeCasesBrokenTorch:
assert b"OK: detect_hardware returned CPU" in result.stdout
def test_fake_torch_no_cuda(self, no_torch_venv, sandbox_dir):
"""Fake torch that imports OK but torch.cuda.is_available() returns False.
detect_hardware() should still fall back to CPU.
"""
"""Fake torch imports OK but cuda.is_available() is False: detect_hardware() falls back to CPU."""
torch_dir = sandbox_dir / "torch"
torch_dir.mkdir()
(torch_dir / "__init__.py").write_text(
@ -582,12 +532,7 @@ class TestEdgeCasesBrokenTorch:
assert b"OK:" in result.stdout
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
the try block. This should fail at call time, not import time -- proving the
lazy import pattern works correctly.
"""
"""apply_chat_template_to_dataset imports without torch; the lazy import fails at call time, not import time."""
_write_loggers_stub(sandbox_dir)
code = textwrap.dedent(f"""\
@ -631,16 +576,14 @@ class TestEdgeCasesBrokenTorch:
assert b"OK: import succeeded" in result.stdout
# ===========================================================================
# Group 4: Hardware Detection Without Torch (3 tests)
# ===========================================================================
# Group 4: Hardware Detection Without Torch
class TestHardwareDetectionNoTorch:
"""Hardware module works without torch, falling back to CPU."""
def test_detect_hardware_no_torch(self, no_torch_venv, sandbox_dir):
"""detect_hardware() returns CPU device when torch is not installed."""
"""detect_hardware() returns CPU when torch is not installed."""
_write_loggers_stub(sandbox_dir)
_write_structlog_stub(sandbox_dir)
@ -680,7 +623,7 @@ class TestHardwareDetectionNoTorch:
assert b"OK:" in result.stdout
def test_hardware_module_import_no_torch(self, no_torch_venv, sandbox_dir):
"""The hardware module imports and detect_hardware is callable without torch."""
"""Hardware module imports and detect_hardware is callable without torch."""
_write_loggers_stub(sandbox_dir)
_write_structlog_stub(sandbox_dir)
_write_hardware_stub(sandbox_dir)
@ -707,13 +650,11 @@ class TestHardwareDetectionNoTorch:
assert b"OK:" in result.stdout
# ===========================================================================
# Group 5: install.sh Logic (5 tests via bash subprocess)
# ===========================================================================
# Group 5: install.sh Logic (via bash subprocess)
class TestInstallShLogic:
"""Test install.sh flag parsing, platform detection, and guard logic."""
"""install.sh flag parsing, platform detection, and guard logic."""
@pytest.fixture(autouse = True)
def _check_install_sh(self):
@ -722,7 +663,6 @@ class TestInstallShLogic:
def test_python_flag_parsing(self):
"""--python flag correctly sets _USER_PYTHON."""
# Extract flag parser snippet from install.sh and test it
script = textwrap.dedent("""\
_USER_PYTHON=""
_next_is_python=false
@ -738,9 +678,8 @@ class TestInstallShLogic:
done
echo "$_USER_PYTHON"
""")
# Test: --python 3.12
# --python 3.12
r = _run_sh(f"{script}" + "\n", timeout = 10)
# Need to pass args to the script
r = subprocess.run(
["bash", "-c", script + "\n", "_", "--python", "3.12"],
capture_output = True,
@ -748,7 +687,7 @@ class TestInstallShLogic:
)
assert r.stdout.strip() == b"3.12"
# Test: --local --python 3.11
# --local --python 3.11
r = subprocess.run(
["bash", "-c", script + "\n", "_", "--local", "--python", "3.11"],
capture_output = True,
@ -756,7 +695,7 @@ class TestInstallShLogic:
)
assert r.stdout.strip() == b"3.11"
# Test: no --python flag
# no --python flag
r = subprocess.run(
["bash", "-c", script + "\n", "_", "--local"],
capture_output = True,
@ -766,7 +705,6 @@ class TestInstallShLogic:
def test_python_flag_missing_arg_errors(self):
"""--python without a version argument triggers an error."""
# Extract the flag parser + error guard from install.sh
script = textwrap.dedent("""\
set -e
_USER_PYTHON=""
@ -819,7 +757,7 @@ class TestInstallShLogic:
)
assert r.stdout.strip() == b"3.12"
# Non-Intel, no override
# non-Intel, no override
r = subprocess.run(
["bash", "-c", script + "\n", "_", "false", ""],
capture_output = True,
@ -865,7 +803,6 @@ class TestInstallShLogic:
def test_stale_venv_guard_respects_override(self):
"""When _USER_PYTHON is set, the stale venv recreation guard is skipped."""
# The guard: if MAC_INTEL=true && -z _USER_PYTHON && venv exists ...
script = textwrap.dedent("""\
MAC_INTEL=true
_USER_PYTHON="$1"
@ -877,7 +814,7 @@ class TestInstallShLogic:
fi
echo "$SHOULD_RECREATE"
""")
# With override: should NOT recreate
# with override: should NOT recreate
r = subprocess.run(
["bash", "-c", script + "\n", "_", "3.11"],
capture_output = True,
@ -885,7 +822,7 @@ class TestInstallShLogic:
)
assert r.stdout.strip() == b"false"
# Without override: SHOULD recreate
# without override: SHOULD recreate
r = subprocess.run(
["bash", "-c", script + "\n", "_", ""],
capture_output = True,
@ -894,13 +831,11 @@ class TestInstallShLogic:
assert r.stdout.strip() == b"true"
# ===========================================================================
# Group 6: install_python_stack.py NO_TORCH Filtering (4 tests)
# ===========================================================================
# Group 6: install_python_stack.py NO_TORCH Filtering
class TestInstallPythonStackFiltering:
"""Test the NO_TORCH filtering logic in install_python_stack.py."""
"""NO_TORCH filtering logic in install_python_stack.py."""
@pytest.fixture(autouse = True)
def _check_install_py(self):
@ -956,7 +891,7 @@ class TestInstallPythonStackFiltering:
):
assert ips._infer_no_torch() is True
# Explicit false on Intel Mac
# explicit false on Intel Mac
with (
mock.patch.dict(os.environ, {"UNSLOTH_NO_TORCH": "false"}),
mock.patch.object(ips, "IS_MAC_INTEL", True),
@ -978,7 +913,6 @@ 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"
# macOS guard for triton
@ -987,9 +921,7 @@ class TestInstallPythonStackFiltering:
), "'not IS_WINDOWS and not IS_MACOS' guard for triton not found"
# ===========================================================================
# Group 7: Live Server Startup (4 tests) -- Heavyweight
# ===========================================================================
# Group 7: Live Server Startup -- Heavyweight
def _studio_venv_python() -> Path | None:
@ -1013,14 +945,7 @@ server = pytest.mark.server
@server
class TestLiveServerStartup:
"""Live server startup tests.
These use the existing Studio venv at ~/.unsloth/studio/unsloth_studio.
They temporarily ensure torch is not importable, test server startup,
then leave the venv unchanged.
Run separately: pytest -m server
"""
"""Live server startup against the existing Studio venv with torch made unimportable (pytest -m server)."""
@pytest.fixture(autouse = True)
def _check_studio_venv(self):
@ -1038,7 +963,6 @@ class TestLiveServerStartup:
port = _server_port()
backend_dir = BACKEND_DIR
# Check if torch is installed in the studio venv
check = subprocess.run(
[str(py), "-c", "import torch; print(torch.__version__)"],
capture_output = True,
@ -1046,7 +970,6 @@ class TestLiveServerStartup:
torch_was_installed = check.returncode == 0
torch_version = check.stdout.decode().strip() if torch_was_installed else None
# Uninstall torch if present
if torch_was_installed:
subprocess.run(
[
@ -1063,7 +986,6 @@ class TestLiveServerStartup:
timeout = 120,
)
# Start server
env = os.environ.copy()
env["PYTHONPATH"] = str(backend_dir)
proc = subprocess.Popen(
@ -1091,7 +1013,6 @@ class TestLiveServerStartup:
if not ready:
stdout, stderr = proc.communicate(timeout = 5)
# Reinstall torch + torchvision + torchaudio
if torch_was_installed and torch_version:
subprocess.run(
[
@ -1189,7 +1110,7 @@ class TestLiveServerStartup:
try:
urllib.request.urlopen(f"http://127.0.0.1:{port}{ep}", timeout = 5)
except urllib.error.HTTPError:
pass # 4xx/5xx is fine -- server didn't crash
pass # 4xx/5xx fine -- server didn't crash
except urllib.error.URLError:
pytest.fail(f"Server stopped responding at {ep}")

38
tests/python/test_fast_language_model_text_only.py
View file

@ -63,7 +63,7 @@ def _param_default(method, name):
def _load_text_only_namespace():
# Exec the text-only helpers from _utils into one namespace (no unsloth import),
# Exec the _utils text-only helpers into one namespace (no unsloth import),
# in dependency order so cross-references resolve.
source = _source(UTILS_PATH)
import transformers
@ -123,8 +123,7 @@ 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 defaults False (opt-in); 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
@ -206,22 +205,14 @@ def test_fast_base_model_text_only_bypasses_vision_auto_model():
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.
"""
"""End-to-end: a tiny Gemma3 text-only model instantiates with text LM attrs and no vision tower."""
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.
# Shrink for cheap CPU instantiation.
text_config.num_hidden_layers = 1
text_config.hidden_size = 32
text_config.intermediate_size = 32
@ -233,10 +224,9 @@ def test_gemma3_text_only_model_class_resolves_and_has_no_vision_tower():
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.
# No vision tower / multimodal projector remains.
assert not hasattr(
model, "vision_tower"
), "text-only Gemma3 model should not have a vision_tower"
@ -246,7 +236,7 @@ def test_gemma3_text_only_model_class_resolves_and_has_no_vision_tower():
def test_helper_defined_once_in_utils_and_imported():
# _get_text_only_config is defined only in _utils and imported by loader + vision.
# _get_text_only_config defined only in _utils, imported by loader + vision.
def _defines(path):
return any(
isinstance(n, ast.FunctionDef) and n.name == "_get_text_only_config"
@ -274,7 +264,7 @@ def _load_util_func(name):
def test_text_only_guard_predicate_across_vlm_families():
# Text-only is taken only when the resolved class remaps VLM weights.
# Text-only taken only when the resolved class remaps VLM weights.
transformers = pytest.importorskip("transformers")
from transformers import AutoModelForCausalLM
@ -309,7 +299,7 @@ def test_text_only_guard_predicate_across_vlm_families():
def test_text_only_helper_preserves_quantization_config():
# quantization_config must survive the strip so pre-quantized repos still load. A
# quantization_config must survive the strip so pre-quantized repos load. A
# sentinel object avoids a bitsandbytes dependency on transformers 4.51.3.
transformers = pytest.importorskip("transformers")
helper = _load_text_only_helper()
@ -318,13 +308,13 @@ def test_text_only_helper_preserves_quantization_config():
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.
# The parent's shared text sub-config must not be mutated.
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).
# Remap the published VLM decoder prefixes, applying only 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())
@ -338,8 +328,8 @@ def test_text_only_key_mapping_targets_published_prefixes():
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.
# PR #5816: text-only loading of a Gemma 3 VLM checkpoint must load 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
@ -391,7 +381,7 @@ def test_gemma3_text_only_loads_real_language_weights_from_vlm_checkpoint(tmp_pa
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
# tf >=5 saves under an outer "model." prefix; strip it to reproduce the
# language_model.model.* layout the published Gemma 3 checkpoints use.
real_dir = tmp_path / "real"
real_dir.mkdir()

19
tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
View file

@ -1,10 +1,7 @@
"""FastSentenceTransformer constructor-redirect lifecycle:
- AutoModel/AutoProcessor/AutoTokenizer.from_pretrained are restored even
when the Transformer constructor raises (try/finally invariant).
- The closure that decides whether to substitute the pre-loaded objects
(`is_requested_model_name`) handles HF repo IDs, local paths, trailing
slashes, pathlib.Path objects, and missing identifiers correctly.
"""
"""FastSentenceTransformer constructor-redirect lifecycle: Auto*.from_pretrained
are restored even when the Transformer constructor raises (try/finally), and
`is_requested_model_name` matches HF IDs, local paths, trailing slashes, Path
objects, and missing identifiers."""
from __future__ import annotations
@ -18,8 +15,7 @@ import pytest
def _stub_module(name: str) -> types.ModuleType:
# __spec__ set so find_spec(name) doesn't raise if a later 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
@ -34,9 +30,8 @@ _STUB_KEYS = (
@pytest.fixture(autouse = True)
def _restore_sys_modules():
"""Snapshot the entries we shadow with stubs and restore them after each
test so a downstream test that does `import transformers` for real does
not pick up our non-package stub."""
"""Snapshot the stubbed module entries and restore them after each test so a
downstream real `import transformers` doesn't pick up our stub."""
saved = {k: sys.modules.get(k) for k in _STUB_KEYS}
try:
yield

2
tests/python/test_install_python_stack.py
View file

@ -10,11 +10,9 @@ from unittest import mock
import pytest
# 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))
# Import after path setup.
import install_python_stack as ips

47
tests/python/test_no_torch_filtering.py
View file

@ -1,12 +1,4 @@
"""Tests for install_python_stack NO_TORCH / IS_MACOS filtering logic.
Covers:
- _filter_requirements unit tests (synthetic + REAL requirements files)
- NO_TORCH / IS_MACOS / IS_WINDOWS env var parsing
- Subprocess-mock of install_python_stack() to verify overrides/triton/filtering
actually happen (or get skipped) under each platform/config combination
- VCS URL and environment marker edge cases in filtering
"""
"""Tests for install_python_stack NO_TORCH / IS_MACOS requirement filtering."""
from __future__ import annotations
@ -21,7 +13,7 @@ from unittest import mock
import pytest
# Add the studio directory so we can import install_python_stack
# Add the studio directory so install_python_stack is importable.
STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
sys.path.insert(0, str(STUDIO_DIR))
@ -59,7 +51,6 @@ class TestFilterRequirements:
)
result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
lines = Path(result).read_text(encoding = "utf-8").splitlines()
# Only numpy should remain (non-blank lines)
non_blank = [l.strip() for l in lines if l.strip()]
assert non_blank == ["numpy"], f"Expected only numpy, got: {non_blank}"
@ -80,7 +71,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()]
# Comment starts with "#", not "torch-stoi", so it's preserved
# Comment lines start with "#", so they are preserved.
assert len(non_blank) == 2
assert non_blank[0].startswith("#")
assert non_blank[1] == "numpy"
@ -139,7 +130,7 @@ class TestFilterRequirements:
)
result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
content = Path(result).read_text(encoding = "utf-8")
# Blank lines should be preserved (not stripped)
# Blank lines must be preserved.
assert "\n\n" in content or content.count("\n") >= 3
def test_stacked_windows_and_no_torch_filters(self, tmp_path):
@ -154,7 +145,7 @@ class TestFilterRequirements:
numpy
""",
)
# First filter Windows packages, then NO_TORCH packages
# 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)
lines = Path(result).read_text(encoding = "utf-8").splitlines()
@ -203,7 +194,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()]
# The git+ URL doesn't start with any skip package, so it is preserved
# git+ URL starts with no skip package, so it is preserved.
assert len(non_blank) == 2, f"git+ URL should be preserved, got: {non_blank}"
@ -231,13 +222,13 @@ class TestRealRequirementsFiltering:
filtered = self._non_blank_non_comment(Path(result))
original = self._non_blank_non_comment(EXTRAS_TXT)
# These must be gone
# These must be gone.
for pkg in ["torch-stoi", "timm", "openai-whisper", "transformers-cfg"]:
assert not any(
l.lower().startswith(pkg) for l in filtered
), f"{pkg} should be removed from extras.txt"
# Everything else must remain
# Everything else must remain.
expected = [
l
for l in original
@ -364,9 +355,7 @@ class TestIsMacosConstant:
class TestInstallPythonStackSubprocessMock:
"""Monkeypatch subprocess.run to capture all pip/uv commands,
then verify which requirements files are used/skipped under
different NO_TORCH / IS_MACOS / IS_WINDOWS configurations."""
"""Mock subprocess.run to verify which req files are used/skipped per config."""
@pytest.fixture(autouse = True)
def _check_req_files(self):
@ -383,10 +372,7 @@ class TestInstallPythonStackSubprocessMock:
*,
skip_base: bool = True,
):
"""Run install_python_stack() with mocked subprocess, capturing all commands.
Returns a list of string-joined commands (each element is ' '.join(cmd)).
"""
"""Run install_python_stack() with mocked subprocess; return joined commands."""
captured_cmds: list[list[str]] = []
def mock_run(cmd, **kw):
@ -471,11 +457,7 @@ class TestInstallPythonStackSubprocessMock:
# -- Normal Linux path (NO_TORCH=False, IS_MACOS=False, IS_WINDOWS=False) --
def test_normal_linux_includes_overrides(self):
"""Normal Linux: the torchao override step IS called.
The override step installs a torch-matched torchao spec via
--force-reinstall (uv: --reinstall), not overrides.txt directly.
"""
"""Normal Linux: torchao override step runs (via --reinstall, not overrides.txt)."""
cmds = self._capture_install(no_torch = False, is_macos = False, is_windows = False)
assert any(
"--reinstall" in cmd for cmd in cmds
@ -512,12 +494,7 @@ class TestInstallPythonStackSubprocessMock:
), "triton-kernels.txt should be skipped on Windows even without NO_TORCH"
def test_windows_only_includes_overrides(self):
"""Windows (without NO_TORCH): overrides IS called (via filtered temp file).
On Windows, all req files go through _filter_requirements(WINDOWS_SKIP_PACKAGES),
so the command uses a temp file, not overrides.txt directly. We check for
--reinstall (uv translation of --force-reinstall) which is unique to overrides.
"""
"""Windows (no NO_TORCH): overrides runs via filtered temp file (check --reinstall)."""
cmds = self._capture_install(no_torch = False, is_macos = False, is_windows = True)
assert any(
"--reinstall" in cmd for cmd in cmds

3
tests/python/test_orpo_processor_text_tokenizer.py
View file

@ -13,8 +13,7 @@ def _load_orpo_rewriter(name = "orpo_trainer_text_tokenizer"):
src = open(RL_PATH).read()
tree = ast.parse(src)
ns = {"re": re}
# Materialise sibling module-level assignments (e.g. _PAD_FALLBACK) so
# any rewriter that references them at exec-time can resolve them.
# Materialise sibling module-level _-prefixed assignments the rewriter may reference.
for node in tree.body:
if isinstance(node, ast.Assign):
for target in node.targets:

10
tests/python/test_patch_trl_rl_trainers_defensive.py
View file

@ -1,12 +1,8 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Regression tests: _patch_trl_rl_trainers must never raise.
The wrapper in unsloth/models/rl.py ring-fences the impl so direct
callers (CI shims, downstream tools) don't have to. Lock that
contract here.
"""
"""Regression: _patch_trl_rl_trainers (the ring-fencing wrapper in
unsloth/models/rl.py) must never raise."""
from __future__ import annotations
@ -39,7 +35,7 @@ def test_patch_trl_rl_trainers_swallows_garbage_input():
def test_impl_is_separately_exposed():
# Power users can still call the impl directly for the raising path.
# The impl stays directly callable for the raising path.
_wrapper, impl = _import_helpers()
assert callable(impl)

27
tests/python/test_studio_import_no_torch.py
View file

@ -1,13 +1,4 @@
"""End-to-end sandbox tests: Studio modules in isolated no-torch venvs.
Covers:
- Python 3.12 and 3.13 venv creation (Intel Mac uses 3.12, Apple Silicon/Linux 3.13)
- data_collators.py loads and dataclasses instantiate without torch
- chat_templates.py top-level exec works with stubs for relative imports
- Negative control: prepending 'import torch' fails in no-torch venv
- Negative control: installing torchao (from overrides.txt) fails in no-torch venv
- AST structural checks for top-level torch imports
"""
"""Sandbox tests: Studio dataset modules load/run in isolated no-torch venvs."""
from __future__ import annotations
@ -48,10 +39,7 @@ def _create_venv(venv_dir: Path, python_version: str) -> Path | None:
@pytest.fixture(params = ["3.12", "3.13"], scope = "module")
def no_torch_venv(request, tmp_path_factory):
"""Create a temporary venv at the requested Python version with no torch.
Parametrized for 3.12 (Intel Mac) and 3.13 (Apple Silicon / Linux).
"""
"""Temp no-torch venv, parametrized for 3.12 (Intel Mac) and 3.13 (Apple Silicon / Linux)."""
if not _has_uv():
pytest.skip("uv not available")
@ -353,7 +341,7 @@ class TestFormatConversionAST:
and child.module
and child.module.startswith("torch")
):
# This torch import must be inside a Try node
# This torch import must be inside a Try node.
found_in_try = False
for try_node in ast.walk(node):
if isinstance(try_node, ast.Try):
@ -524,10 +512,7 @@ class TestNegativeControls:
os.unlink(temp_file)
def test_torchao_install_fails_no_torch_venv(self, no_torch_venv):
"""Installing torchao (from overrides.txt) fails in a no-torch venv.
This proves the overrides.txt skip is necessary for Intel Mac.
"""
"""torchao install fails in a no-torch venv: proves the overrides.txt skip is needed."""
result = subprocess.run(
[
no_torch_venv,
@ -541,10 +526,10 @@ class TestNegativeControls:
timeout = 60,
)
if result.returncode != 0:
# torchao install/resolution failed as expected
# torchao install/resolution failed as expected.
pass
else:
# pip dry-run may not catch dependency issues; verify torch is missing
# dry-run may miss dep issues; verify torch is absent instead.
check = subprocess.run(
[no_torch_venv, "-c", "import torch"],
capture_output = True,

65
tests/python/test_tokenizers_and_torch_constraint.py
View file

@ -1,7 +1,4 @@
"""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)
"""
"""Install fixes: tokenizers in no-torch-runtime.txt, and TORCH_CONSTRAINT in install.sh."""
from __future__ import annotations
@ -12,7 +9,7 @@ import textwrap
import pytest
# ── Locate source files relative to this test ──────────────────────────
# Locate source files relative to this test.
_TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/
_REPO_ROOT = _TESTS_DIR.parent # unsloth/
_INSTALL_SH = _REPO_ROOT / "install.sh"
@ -33,9 +30,7 @@ def _lines(path: pathlib.Path) -> list[str]:
]
# ======================================================================
# Group 1 -- Structural checks (no network, instant)
# ======================================================================
class TestStructuralTokenizers:
"""Verify tokenizers presence and ordering in no-torch-runtime.txt."""
@ -114,15 +109,11 @@ class TestStructuralInstallPs1Unchanged:
assert '"torch>=2.4,<2.11.0"' in self._ps1
# ======================================================================
# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
# ======================================================================
class TestTorchConstraintShell:
"""Test the TORCH_CONSTRAINT block using bash subprocesses with
mocked python binaries that return controlled minor versions."""
"""Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""
# The extracted snippet we test in isolation. We override OS, _ARCH,
# SKIP_TORCH, and provide a mock python at $VENV_DIR/bin/python.
# Snippet tested in isolation: override OS/_ARCH/SKIP_TORCH and a mock python.
_SNIPPET_TEMPLATE = textwrap.dedent(r"""
#!/bin/bash
set -e
@ -191,8 +182,6 @@ class TestTorchConstraintShell:
assert result.returncode == 0, f"Script failed: {result.stderr}"
return result.stdout.strip()
# -- arm64 macOS tightening cases --
def test_arm64_macos_py313_tightened(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64")
assert out == "torch>=2.6,<2.11.0"
@ -201,8 +190,6 @@ class TestTorchConstraintShell:
out = self._run(tmp_path, py_minor = 14, os_val = "macos", arch = "arm64")
assert out == "torch>=2.6,<2.11.0"
# -- arm64 macOS default (older python) --
def test_arm64_macos_py312_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64")
assert out == "torch>=2.4,<2.11.0"
@ -211,8 +198,7 @@ class TestTorchConstraintShell:
out = self._run(tmp_path, py_minor = 11, os_val = "macos", arch = "arm64")
assert out == "torch>=2.4,<2.11.0"
# -- Linux (unaffected) --
# Linux is unaffected by the tightening.
def test_linux_x86_py313_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "x86_64")
assert out == "torch>=2.4,<2.11.0"
@ -221,14 +207,12 @@ class TestTorchConstraintShell:
out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "aarch64")
assert out == "torch>=2.4,<2.11.0"
# -- Intel Mac (arch mismatch) --
# Intel Mac: arch mismatch, no tightening.
def test_intel_mac_x86_py313_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "x86_64")
assert out == "torch>=2.4,<2.11.0"
# -- SKIP_TORCH bypass --
# SKIP_TORCH bypasses the tightening.
def test_skip_torch_arm64_macos_py313_default(self, tmp_path):
out = self._run(
tmp_path,
@ -239,16 +223,12 @@ class TestTorchConstraintShell:
)
assert out == "torch>=2.4,<2.11.0"
# -- WSL --
def test_wsl_py313_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "wsl", arch = "x86_64")
assert out == "torch>=2.4,<2.11.0"
# -- Edge cases --
def test_py_minor_0_fallback_default(self, tmp_path):
"""If python query fails (returns 0), should stay at default."""
"""Failed python query (returns 0) keeps the default constraint."""
out = self._run(tmp_path, py_minor = 0, os_val = "macos", arch = "arm64")
assert out == "torch>=2.4,<2.11.0"
@ -261,10 +241,10 @@ class TestTorchConstraintShell:
assert out == "torch>=2.6,<2.11.0"
def test_mock_uv_receives_correct_constraint(self, tmp_path):
"""Verify a mock uv would receive the correct constraint string."""
"""A mock uv receives the tightened constraint on py3.13 arm64 macOS."""
venv = self._make_mock_python(tmp_path, minor = 13)
# Create a mock uv that logs its arguments
# Mock uv logs its arguments.
mock_uv = tmp_path / "mock_uv"
log_file = tmp_path / "uv_log.txt"
mock_uv.write_text(
@ -354,13 +334,10 @@ class TestTorchConstraintShell:
assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}"
# ======================================================================
# Group 3 -- E2E tokenizers fix (requires network, ~2-5 min)
# ======================================================================
@pytest.mark.e2e
class TestE2ETokenizersFix:
"""Creates real uv venvs to verify tokenizers + transformers work
without torch installed."""
"""Real uv venvs verify tokenizers + transformers work without torch installed."""
@staticmethod
def _create_venv(tmp_path: pathlib.Path, name: str, py: str) -> pathlib.Path:
@ -393,8 +370,7 @@ class TestE2ETokenizersFix:
@pytest.mark.parametrize("py_version", ["3.12", "3.13"])
def test_autoconfig_works_with_no_torch_runtime(self, tmp_path, py_version):
"""Install from no-torch-runtime.txt with --no-deps (matching the
real install.sh path), then verify AutoConfig imports successfully."""
"""Install no-torch-runtime.txt with --no-deps, then AutoConfig must import."""
venv = self._create_venv(tmp_path, f"tok-{py_version}", py_version)
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
@ -423,8 +399,7 @@ class TestE2ETokenizersFix:
assert result.returncode != 0, "torch should NOT be importable"
def test_negative_control_no_tokenizers(self, tmp_path):
"""Without tokenizers, AutoConfig should fail. We create a copy of
no-torch-runtime.txt with the tokenizers line removed."""
"""Without the tokenizers line, AutoConfig must fail (negative control)."""
venv = self._create_venv(tmp_path, "neg-ctrl", "3.12")
req_no_tokenizers = tmp_path / "no-tokenizers.txt"
req_no_tokenizers.write_text(
@ -440,9 +415,7 @@ class TestE2ETokenizersFix:
assert "tokenizers" in result.stderr.lower() or "ModuleNotFoundError" in result.stderr
# ======================================================================
# Group 4 -- Integration: install.sh reads no-torch-runtime.txt correctly
# ======================================================================
class TestInstallShNoTorchIntegration:
"""Verify install.sh has the correct no-torch-runtime.txt wiring."""
@ -457,14 +430,11 @@ class TestInstallShNoTorchIntegration:
def test_no_deps_invocation_for_fresh(self):
"""Fresh install path should also use --no-deps -r."""
# Count occurrences of the no-deps -r pattern
count = self._sh.count('--no-deps -r "$_NO_TORCH_RT"')
assert count >= 2, f"Expected >=2 no-deps -r invocations, found {count}"
def test_mock_uv_skip_torch_reads_requirements(self, tmp_path):
"""When SKIP_TORCH=true, the _find_no_torch_runtime path should be used."""
# We test this structurally: verify the SKIP_TORCH=true blocks contain
# _find_no_torch_runtime calls
"""SKIP_TORCH=true blocks must call _find_no_torch_runtime."""
skip_blocks = re.findall(
r'if \[ "\$SKIP_TORCH" = true \].*?(?=\n (?:else|elif|fi))',
self._sh,
@ -474,9 +444,7 @@ class TestInstallShNoTorchIntegration:
assert found, "SKIP_TORCH=true block should call _find_no_torch_runtime"
# ======================================================================
# Group 5 -- Full no-torch sandbox (requires network, ~5 min)
# ======================================================================
@pytest.mark.e2e
class TestE2EFullNoTorchSandbox:
"""Creates venvs and installs the actual no-torch-runtime.txt."""
@ -511,8 +479,7 @@ class TestE2EFullNoTorchSandbox:
)
def test_autoconfig_succeeds(self, tmp_path):
"""The real bug fix: install with --no-deps (matching install.sh)
and verify from transformers import AutoConfig works."""
"""Install with --no-deps and verify AutoConfig imports (the bug fix)."""
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}"
@ -522,7 +489,7 @@ class TestE2EFullNoTorchSandbox:
), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
def test_torch_not_importable(self, tmp_path):
"""With --no-deps (as install.sh uses), torch must not be pulled in."""
"""With --no-deps, torch must not be pulled in."""
venv = self._create_venv(tmp_path, "no-torch-check")
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"

11
tests/python/test_unsloth_run_tool_policy_resolver.py
View file

@ -1,15 +1,6 @@
# Copyright 2025-present the Unsloth AI Inc. team. All rights reserved.
"""
Truth-table tests for `resolve_tool_policy` -- the pure resolver behind
`unsloth run --enable-tools/--disable-tools`.
Covers:
- 127.0.0.1 default-on, explicit on, explicit off
- 0.0.0.0 default-off, explicit off
- 0.0.0.0 + explicit on: confirm prompt unless --silent or --yes,
abort on negative answer.
"""
"""Truth-table tests for `resolve_tool_policy` behind `unsloth run --enable-tools/--disable-tools`."""
import pytest
import typer

4
tests/saving/gpt-oss-merge/train_and_merge.py
View file

@ -20,7 +20,7 @@ def safe_remove_directory(path):
return False
# Used by the mapping function below.
# Used by formatting_prompts_func below.
tokenizer = None
@ -91,5 +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") # cache created by this process
print("✅ Cleanup complete. Exiting training script.")

7
tests/saving/language_models/test_merge_4bit_validation.py
View file

@ -93,7 +93,7 @@ trainer = SFTTrainer(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
warmup_ratio = 0.1,
max_steps = 10, # Very short training for test
max_steps = 10,
learning_rate = 2e-4,
fp16 = not is_bfloat16_supported(),
bf16 = is_bfloat16_supported(),
@ -129,7 +129,6 @@ del model
del tokenizer
torch.cuda.empty_cache()
# Load the 4bit merged model
model_4bit, tokenizer_4bit = FastLanguageModel.from_pretrained(
model_name = "./test_4bit_model",
max_seq_length = 2048,
@ -144,7 +143,6 @@ tokenizer_4bit = get_chat_template(
print("✅ 4bit model loaded successfully!")
# Add LoRA adapters
model_4bit = FastLanguageModel.get_peft_model(
model_4bit,
r = 16,
@ -166,7 +164,6 @@ model_4bit = FastLanguageModel.get_peft_model(
loftq_config = None,
)
# Second fine-tuning
trainer_4bit = SFTTrainer(
model = model_4bit,
tokenizer = tokenizer_4bit,
@ -180,7 +177,7 @@ trainer_4bit = SFTTrainer(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
warmup_ratio = 0.1,
max_steps = 10, # Very short training for test
max_steps = 10,
learning_rate = 2e-4,
fp16 = not is_bfloat16_supported(),
bf16 = is_bfloat16_supported(),

2
tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py
View file

@ -76,7 +76,7 @@ def load_and_compute_8bit_ppl(
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
# Coerce to a plain Python float for cross-process transfer
# Coerce to a plain Python float for cross-process transfer.
if torch.is_tensor(ppl_value):
ppl_value = ppl_value.cpu().item()
elif hasattr(ppl_value, "item"):

2
tests/saving/language_models/test_merge_model_perplexity_mistral.py
View file

@ -101,7 +101,7 @@ def load_and_compute_8bit_ppl(
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
# Convert to Python float if it's a tensor.
# Coerce to a plain Python float.
if torch.is_tensor(ppl_value):
ppl_value = ppl_value.cpu().item()
elif hasattr(ppl_value, "item"):

4
tests/saving/language_models/test_merge_model_perplexity_phi_4.py
View file

@ -63,7 +63,7 @@ def load_and_compute_8bit_ppl(
chat_template = "phi-4",
)
# Load dataset fresh in subprocess
# Load dataset fresh in subprocess.
dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
def formatting_prompts_func(examples):
@ -78,7 +78,7 @@ def load_and_compute_8bit_ppl(
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
# Coerce to a Python float regardless of source type.
# Coerce to a Python float.
if torch.is_tensor(ppl_value):
ppl_value = ppl_value.cpu().item()
elif hasattr(ppl_value, "item"):

1
tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py
View file

@ -30,7 +30,6 @@ from tests.utils.perplexity_eval import (
)
# Define helper functions outside of main
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [

4
tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py
View file

@ -95,7 +95,6 @@ def load_and_compute_8bit_ppl(
# chat_template="llama-3.1",
# )
# 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.
@ -144,7 +143,7 @@ def load_and_compute_8bit_ppl(
ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
# Coerce to a Python float regardless of source type.
# Coerce to a Python float.
if torch.is_tensor(ppl_value):
ppl_value = ppl_value.cpu().item()
elif hasattr(ppl_value, "item"):
@ -243,7 +242,6 @@ if __name__ == "__main__":
add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
# 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

11
tests/saving/language_models/test_push_to_hub_merged.py
View file

@ -130,11 +130,9 @@ trainer = train_on_responses_only(
response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
)
# run training
trainer_stats = trainer.train()
# 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()
@ -152,7 +150,7 @@ success = {
"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))
@ -165,14 +163,13 @@ except Exception as e:
raise Exception("Model upload failed.")
t
# Stage 2: Test downloading the model (even if cached)
# Stage 2: Test downloading the model.
safe_remove_directory(f"./{hf_username}")
try:
print("\n" + "=" * 80)
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")
success["download"] = True
print("✅ Model downloaded successfully!")
@ -180,7 +177,7 @@ except Exception as e:
print(f"❌ Download failed: {e}")
raise Exception("Model download failed.")
# Final report
# Final report.
print("\n" + "=" * 80)
print("=== VALIDATION REPORT ===".center(80))
print("=" * 80 + "\n")
@ -194,6 +191,6 @@ if all(success.values()):
else:
raise Exception("Validation failed for one or more stages.")
# final cleanup
# final cleanup.
safe_remove_directory("./outputs")
safe_remove_directory("./unsloth_compiled_cache")

2
tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py
View file

@ -134,7 +134,7 @@ trainer = train_on_responses_only(
trainer_stats = trainer.train()
# saving and merging the model to local disk
# Resolve HF username + token (prompt if unset).
hf_username = os.environ.get("HF_USER", "")
if not hf_username:
hf_username = input("Please enter your Hugging Face username: ").strip()

54
tests/saving/language_models/test_save_merged_grpo_model.py
View file

@ -1,9 +1,5 @@
# -*- coding: utf-8 -*-
"""test_Llama3_1_(3B)_GRPO_LoRA (1).ipynb
### Unsloth
"""
"""Llama 3.1 (3B) GRPO LoRA train + merged-model save/eval."""
from unsloth import FastLanguageModel
import torch
@ -20,8 +16,8 @@ from tests.utils.cleanup_utils import safe_remove_directory
from tests.utils.aime_eval import evaluate_model_aime, compare_aime_results
max_seq_length = 2048 # Can increase for longer reasoning traces
lora_rank = 64 # Larger rank = smarter, but slower
max_seq_length = 2048
lora_rank = 64
def evaluate_merged_model(
@ -32,16 +28,16 @@ def evaluate_merged_model(
from unsloth import FastLanguageModel
from tests.utils.aime_eval import evaluate_model_aime
max_seq_length = 2048 # Can increase for longer reasoning traces
lora_rank = 64 # Larger rank = smarter, but slower
max_seq_length = 2048
lora_rank = 64
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "./final_merged_model",
max_seq_length = max_seq_length,
load_in_4bit = True, # False for LoRA 16bit
fast_inference = True, # Enable vLLM fast inference
load_in_4bit = True,
fast_inference = True,
max_lora_rank = lora_rank,
gpu_memory_utilization = 0.8, # Reduce if out of memory
gpu_memory_utilization = 0.8,
)
print(f"\n{'='*60}")
@ -79,10 +75,10 @@ def training_run(result_queue):
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "meta-llama/Llama-3.2-3B-Instruct",
max_seq_length = max_seq_length,
load_in_4bit = False, # False for LoRA 16bit
fast_inference = True, # Enable vLLM fast inference
load_in_4bit = False,
fast_inference = True,
max_lora_rank = lora_rank,
gpu_memory_utilization = 0.8, # Reduce if out of memory
gpu_memory_utilization = 0.8,
)
"""### Helper Functions
@ -190,7 +186,7 @@ def training_run(result_queue):
matches = re.findall(pattern, text, re.DOTALL)
if matches:
answer = matches[-1] # Get the last match
answer = matches[-1]
answer = re.sub(r"[%$,]", "", answer).strip()
return answer
return ""
@ -450,7 +446,6 @@ def training_run(result_queue):
all_results = []
# Single temperature evaluation on combined dataset.
results = evaluate_model_aime(
model = model,
tokenizer = tokenizer,
@ -494,7 +489,7 @@ def training_run(result_queue):
model = FastLanguageModel.get_peft_model(
model,
r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
r = lora_rank,
target_modules = [
"q_proj",
"k_proj",
@ -503,9 +498,9 @@ def training_run(result_queue):
"gate_proj",
"up_proj",
"down_proj",
], # Remove QKVO if out of memory
],
lora_alpha = lora_rank,
use_gradient_checkpointing = "unsloth", # Enable long context finetuning
use_gradient_checkpointing = "unsloth",
random_state = 3407,
)
@ -518,12 +513,12 @@ def training_run(result_queue):
max_seq_length = max_seq_length,
data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),
dataset_num_proc = 2,
packing = False, # Can make training 5x faster for short sequences.
packing = False,
args = TrainingArguments(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
warmup_steps = 5,
num_train_epochs = 1, # Set this for 1 full training run.
num_train_epochs = 1,
# max_steps = 60,
learning_rate = 2e-4,
fp16 = not is_bfloat16_supported(),
@ -534,7 +529,7 @@ def training_run(result_queue):
lr_scheduler_type = "linear",
seed = 3407,
output_dir = "outputs",
report_to = "none", # Use this for WandB etc
report_to = "none",
),
)
@ -581,7 +576,6 @@ def training_run(result_queue):
]
scores = []
# Print only every few steps.
global PRINTED_TIMES
global PRINT_EVERY_STEPS
if PRINTED_TIMES % PRINT_EVERY_STEPS == 0:
@ -600,7 +594,6 @@ def training_run(result_queue):
continue
try:
true_answer = float(true_answer.strip())
# Remove commas like in 123,456.
guess = float(guess.strip().replace(",", ""))
scores.append(1.5 if guess == true_answer else -0.5)
except:
@ -613,7 +606,7 @@ def training_run(result_queue):
print(f"{'*'*60}")
max_prompt_length, _ = get_max_prompt_length(gsm8k_train, tokenizer)
max_prompt_length = min(max_prompt_length + 10, 512) # Add buffer, cap at 512
max_prompt_length = min(max_prompt_length + 10, 512)
print(f"Using max_prompt_length: {max_prompt_length}")
@ -627,8 +620,8 @@ def training_run(result_queue):
optim = "adamw_torch_fused",
logging_steps = 1,
per_device_train_batch_size = 1,
gradient_accumulation_steps = 4, # Increase to 4 for smoother training
num_generations = 8, # Decrease if out of memory
gradient_accumulation_steps = 4,
num_generations = 8,
max_prompt_length = max_prompt_length,
max_completion_length = max_seq_length - max_prompt_length,
# num_train_epochs = 1, # Set to 1 for a full training run
@ -636,7 +629,7 @@ def training_run(result_queue):
max_steps = 1000,
save_steps = 250,
max_grad_norm = 0.1,
report_to = "none", # Can use Weights & Biases
report_to = "none",
output_dir = "outputs",
)
@ -749,7 +742,6 @@ if __name__ == "__main__":
result_queue = mp.Queue()
all_results = []
# Run main finetuning and GRPO loop.
p = mp.Process(target = training_run, args = (result_queue,))
p.start()
p.join()
@ -757,7 +749,7 @@ if __name__ == "__main__":
results = result_queue.get()
all_results = results
# Evaluate merged model loaded 16bits.
# Evaluate merged model loaded 16bit.
p = mp.Process(target = evaluate_merged_model, args = (result_queue, False, False))
p.start()
p.join()

2
tests/saving/non_peft/test_whisper_non_peft.py
View file

@ -52,7 +52,7 @@ print(f"{'='*80}")
try:
with warnings.catch_warnings():
warnings.simplefilter("error") # Treat warnings as errors here
warnings.simplefilter("error") # any warning -> failure
model.save_pretrained("test_output")
print("✅ Standard save_pretrained completed successfully!")
except Exception as e:

2
tests/saving/test_preserve_tokenizer_eos_token.py
View file

@ -84,7 +84,7 @@ def test_preserve_tokenizer_eos_token_supports_filename_prefix(tmp_path):
saved_config = json.loads(prefixed_config.read_text(encoding = "utf-8"))
assert saved_config["eos_token"] == "<turn|>"
assert saved_config["other"] is True
# Unprefixed file must not be created as a side effect.
# Unprefixed file must not be created as a side effect (regression).
assert not (tmp_path / "tokenizer_config.json").exists()

52
tests/saving/test_save_subprocess_utf8_encoding.py
View file

@ -13,26 +13,11 @@
"""Regression tests for unslothai/unsloth#2660.
On Windows the default text encoding is the locale code page (e.g. cp1252),
not UTF-8. ``subprocess.Popen`` / ``subprocess.run`` opened in text mode
(``text=True`` / ``universal_newlines=True``) without an explicit
``encoding`` therefore decode child-process output with cp1252. When
llama.cpp / Ollama emit a byte that is undefined in cp1252 (e.g. ``0x9d``,
which appears inside the UTF-8 encoding of common punctuation and box-drawing
glyphs), the read raises ``UnicodeDecodeError`` and aborts the GGUF export.
Two checks:
* ``test_save_subprocess_text_calls_declare_utf8_encoding`` -- a source-level
drift detector. It parses ``unsloth/save.py`` (no import, so it runs under
the GPU/torch-free harness) and fails if any text-mode subprocess call is
missing ``encoding="utf-8"``. This is the regression guard: it is red
before the fix and green after.
* ``test_utf8_replace_decodes_non_cp1252_subprocess_output`` -- a behavioural
check that documents the bug and the fix deterministically on any platform:
raw child output that is invalid under cp1252 raises, while the
``encoding="utf-8", errors="replace"`` kwargs used by the fix read it
cleanly.
On Windows, text-mode subprocess calls without explicit encoding decode child
output as cp1252, raising UnicodeDecodeError on UTF-8 bytes undefined there
(e.g. 0x9d) and aborting GGUF export. A source-level drift check asserts save.py
pins encoding="utf-8" on every text-mode call; a behavioural check reproduces
the cp1252 failure and confirms the utf-8/replace fix reads it cleanly.
"""
from __future__ import annotations
@ -84,9 +69,8 @@ def _collect_text_mode_subprocess_calls() -> list[ast.Call]:
def test_text_mode_subprocess_calls_exist():
"""Guard the guard: if save.py stops using text-mode subprocess calls the
drift test below would vacuously pass, so make sure we are actually
inspecting something."""
"""Guard the guard: ensure save.py still has text-mode subprocess calls so
the drift test below isn't vacuously passing."""
calls = _collect_text_mode_subprocess_calls()
assert len(calls) >= 6, (
f"Expected several text-mode subprocess calls in {SAVE_PY.name}, "
@ -95,11 +79,8 @@ def test_text_mode_subprocess_calls_exist():
def test_save_subprocess_text_calls_declare_utf8_encoding():
"""Every text-mode subprocess call in save.py must pin encoding='utf-8'.
Without it, reading llama.cpp/Ollama output crashes on Windows (cp1252).
Fails before the #2660 fix, passes after.
"""
"""Every text-mode subprocess call in save.py must pin encoding='utf-8'
(else Windows cp1252 crashes reading child output; #2660 fix)."""
offenders = []
for node in _collect_text_mode_subprocess_calls():
enc = _kw(node, "encoding")
@ -115,14 +96,10 @@ def test_save_subprocess_text_calls_declare_utf8_encoding():
def test_utf8_replace_decodes_non_cp1252_subprocess_output():
"""Document the failure and the fix with a real subprocess.
The child emits U+201D (right double quote), whose UTF-8 encoding
``E2 80 9D`` contains byte 0x9D -- undefined in cp1252. Decoding the raw
bytes as cp1252 raises (the bug); the fix's kwargs read it cleanly.
"""
# All-ASCII argv; the child builds the non-ASCII char itself so this is
# deterministic regardless of the parent's locale.
"""Reproduce the bug and fix with a real subprocess: the child emits U+201D
(UTF-8 E2 80 9D, byte 0x9D undefined in cp1252) so cp1252 decode raises while
the fix's utf-8/replace kwargs read it cleanly."""
# All-ASCII argv; the child builds the non-ASCII char so it's locale-independent.
child = (
"import sys; "
"sys.stdout.buffer.write(('tensor ' + chr(0x201D) + ' x\\n').encode('utf-8'))"
@ -131,8 +108,7 @@ def test_utf8_replace_decodes_non_cp1252_subprocess_output():
raw = subprocess.run([sys.executable, "-c", child], capture_output = True).stdout
assert b"\x9d" in raw # precondition: output carries the cp1252-undefined byte
# Failing behaviour before the fix: cp1252 (the Windows default) cannot
# decode this output.
# Before the fix: cp1252 (the Windows default) cannot decode this output.
with pytest.raises(UnicodeDecodeError):
raw.decode("cp1252")

14
tests/saving/test_unsloth_save.py
View file

@ -8,7 +8,7 @@ import importlib
from unsloth import FastLanguageModel, FastModel
model_to_test = [
# Text Models
# Text models
"unsloth/tinyllama",
"unsloth/tinyllama-bnb-4bit",
"unsloth/Qwen2.5-0.5B-Instruct",
@ -16,7 +16,7 @@ model_to_test = [
"unsloth/Phi-4-mini-instruct",
"unsloth/Phi-4-mini-instruct-bnb-4bit",
"unsloth/Qwen2.5-0.5B",
# Vision Models
# Vision models
"unsloth/gemma-3-4b-it",
"unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit",
"unsloth/Qwen2.5-VL-3B-Instruct-bnb-4bit",
@ -67,7 +67,7 @@ def loaded_model_tokenizer(request):
@pytest.fixture(scope = "session", params = torchao_models)
def fp16_model_tokenizer(request):
"""Load model in FP16 for TorchAO quantization"""
"""Load model in FP16 for TorchAO quantization."""
model_name = request.param
print(f"Loading model in FP16 for TorchAO: {model_name}")
@ -75,7 +75,7 @@ def fp16_model_tokenizer(request):
model_name,
max_seq_length = 128,
dtype = None,
load_in_4bit = False, # No BnB quantization
load_in_4bit = False, # no BnB quantization
)
model = FastModel.get_peft_model(
@ -143,7 +143,7 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
os.path.join(save_path, file)
), f"{file} not found in the save directory."
# 16bit if there's no quantization config
# 16bit means no quantization config.
config_path = os.path.join(save_path, "config.json")
with open(config_path, "r") as f:
config = json.load(f)
@ -154,7 +154,6 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
save_file_sizes["merged_16bit"][model.config._name_or_path] = total_size
print(f"Total size of merged_16bit files: {total_size} bytes")
# Verify the saved model loads
loaded_model, loaded_tokenizer = FastLanguageModel.from_pretrained(
save_path,
max_seq_length = 128,
@ -194,14 +193,13 @@ 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."
# 4bit if there's a quantization config
# 4bit means 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."
# Verify the saved model loads
loaded_model, loaded_tokenizer = FastModel.from_pretrained(
save_path,
max_seq_length = 128,

29
tests/saving/text_to_speech_models/test_csm.py
View file

@ -27,10 +27,10 @@ print(f"{'='*80}")
model, tokenizer = FastModel.from_pretrained(
model_name = "unsloth/csm-1b",
max_seq_length = 2048, # Choose any for long context!
dtype = None, # Leave as None for auto-detection
max_seq_length = 2048,
dtype = None,
auto_model = CsmForConditionalGeneration,
load_in_4bit = False, # Select True for 4bit - reduces memory usage
load_in_4bit = False,
)
@ -39,7 +39,7 @@ base_model_class = model.__class__.__name__
model = FastModel.get_peft_model(
model,
r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
r = 32,
target_modules = [
"q_proj",
"k_proj",
@ -50,13 +50,12 @@ model = FastModel.get_peft_model(
"down_proj",
],
lora_alpha = 32,
lora_dropout = 0, # Supports any, but = 0 is optimized
bias = "none", # Supports any, but = "none" is optimized
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
lora_dropout = 0,
bias = "none",
use_gradient_checkpointing = "unsloth",
random_state = 3407,
use_rslora = False, # We support rank stabilized LoRA
loftq_config = None, # And LoftQ
use_rslora = False,
loftq_config = None,
)
print("✅ Model and LoRA adapters loaded successfully!")
@ -97,7 +96,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") # treat warnings as errors so saving stays clean
try:
model.save_pretrained_merged("csm", tokenizer)
print("✅ Model saved and merged successfully without warnings!")
@ -111,10 +110,10 @@ print(f"{'='*80}")
model, processor = FastModel.from_pretrained(
model_name = "./csm",
max_seq_length = 2048, # Choose any for long context!
dtype = None, # Leave as None for auto-detection
max_seq_length = 2048,
dtype = None,
auto_model = CsmForConditionalGeneration,
load_in_4bit = False, # Select True for 4bit - reduces memory usage
load_in_4bit = False,
)
from transformers import AutoProcessor
@ -139,7 +138,7 @@ try:
inputs = processor(f"[{speaker_id}]{text}", add_special_tokens = True).to("cuda")
audio_values = model.generate(
**inputs,
max_new_tokens = 125, # 125 tokens ~= 10 seconds of audio
max_new_tokens = 125, # ~10 seconds of audio
depth_decoder_temperature = 0.6,
depth_decoder_top_k = 0,
depth_decoder_top_p = 0.9,

36
tests/saving/text_to_speech_models/test_lasa.py
View file

@ -44,8 +44,8 @@ max_seq_length = 2048
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "unsloth/Llasa-1B",
max_seq_length = max_seq_length,
dtype = None, # Select None for auto detection
load_in_4bit = False, # Choose True for 4bit which reduces memory
dtype = None,
load_in_4bit = False,
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)
@ -54,16 +54,15 @@ base_model_class = model.__class__.__name__
model = FastLanguageModel.get_peft_model(
model,
r = 128, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
r = 128,
target_modules = ["q_proj", "v_proj"],
lora_alpha = 128,
lora_dropout = 0, # Supports any, but = 0 is optimized
bias = "none", # Supports any, but = "none" is optimized
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
lora_dropout = 0,
bias = "none",
use_gradient_checkpointing = "unsloth",
random_state = 3407,
use_rslora = False, # We support rank stabilized LoRA
loftq_config = None, # And LoftQ
use_rslora = False,
loftq_config = None,
)
print("✅ Model and LoRA adapters loaded successfully!")
@ -104,7 +103,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") # save/merge must emit no warnings
try:
model.save_pretrained_merged("lasa", tokenizer)
print("✅ Model saved and merged successfully without warnings!")
@ -119,8 +118,8 @@ print(f"{'='*80}")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "./lasa",
max_seq_length = max_seq_length,
dtype = None, # Select None for auto detection
load_in_4bit = False, # Choose True for 4bit which reduces memory
dtype = None,
load_in_4bit = False,
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)
@ -164,12 +163,10 @@ def extract_speech_ids(speech_tokens_str):
return speech_ids
# 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|>"
# Tokenize the text
chat = [
{"role": "user", "content": "Convert the text to speech:" + formatted_text},
{"role": "assistant", "content": "<|SPEECH_GENERATION_START|>"},
@ -182,26 +179,23 @@ with torch.inference_mode():
speech_end_id = tokenizer.convert_tokens_to_ids("<|SPEECH_GENERATION_END|>")
# Generate the speech autoregressively
outputs = model.generate(
input_ids,
max_length = 2048, # We trained our model with a max length of 2048
max_length = 2048,
eos_token_id = speech_end_id,
do_sample = True,
top_p = 1.2, # Adjusts the diversity of generated content
temperature = 1.2, # Controls randomness in output
top_p = 1.2,
temperature = 1.2,
)
# Extract the speech tokens
generated_ids = outputs[0][input_ids.shape[1] : -1]
speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens = True)
# Convert token <|s_23456|> to int 23456
# Convert token <|s_23456|> to int 23456.
speech_tokens = extract_speech_ids(speech_tokens)
speech_tokens = torch.tensor(speech_tokens).cpu().unsqueeze(0).unsqueeze(0)
# Decode the speech tokens to speech waveform
gen_wav = codec_model.decode_code(speech_tokens)
try:
sf.write(output_audio_path, gen_wav[0, 0, :].cpu().numpy(), 16000)

2
tests/saving/text_to_speech_models/test_orpheus.py
View file

@ -140,7 +140,7 @@ prompts = [
"Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.",
]
chosen_voice = None # None for single-speaker
chosen_voice = None # single-speaker
prompts_ = [(f"{chosen_voice}: " + p) if chosen_voice else p for p in prompts]

3
tests/saving/text_to_speech_models/test_whisper.py
View file

@ -169,8 +169,7 @@ transcribed_text = whisper(audio_file)
# transcribed_text = model.generate(input_features=input_features)
print(f"📝 Transcribed Text: {transcribed_text['text']}")
## 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.
# Assert the transcription contains the expected reference phrases.
expected_phrases = [
"birch canoe slid on the smooth planks",
"sheet to the dark blue background",

2
tests/saving/vision_models/test_index_file_sharded_model.py
View file

@ -128,7 +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},
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

16
tests/saving/vision_models/test_push_to_hub_merged.py
View file

@ -23,10 +23,7 @@ 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")
# First 2000 examples for training
train_dataset = dataset.select(range(2000))
# Next 200 examples for evaluation
eval_dataset = dataset.select(range(2000, 2200))
print(f"✅ Dataset loaded successfully!")
@ -34,7 +31,7 @@ print(f" 📈 Training samples: {len(train_dataset)}")
print(f" 📊 Evaluation samples: {len(eval_dataset)}")
# Convert dataset to OAI messages
# Convert to OAI messages format
def format_data(sample):
return {
"messages": [
@ -76,7 +73,6 @@ 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(
@ -138,8 +134,8 @@ try:
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
gradient_checkpointing = True,
gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing
max_grad_norm = 0.3, # max gradient norm based on QLoRA paper
gradient_checkpointing_kwargs = {"use_reentrant": False},
max_grad_norm = 0.3, # from QLoRA paper
warmup_ratio = 0.03,
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
max_steps = 10,
@ -175,7 +171,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()
@ -212,7 +207,6 @@ success = {
"upload": False,
"download": False,
}
# Stage 1: Upload model to Hub
try:
print("\n" + "=" * 80)
print("=== UPLOADING MODEL TO HUB ===".center(80))
@ -231,19 +225,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")
@ -259,7 +250,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")

4
tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py
View file

@ -57,7 +57,7 @@ def format_data(sample):
system_message = "You are an expert french ocr system."
# List comprehension (not .map) to keep PIL.Image type; .map converts images to bytes.
# List comprehension (not .map): .map would convert PIL images to bytes.
train_dataset = [format_data(sample) for sample in train_dataset]
eval_dataset = [format_data(sample) for sample in eval_dataset]
@ -128,7 +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},
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

11
tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py
View file

@ -27,7 +27,6 @@ train_dataset = dataset.select(range(2000))
eval_dataset = dataset.select(range(2000, 2200))
# Convert dataset to OAI messages.
def format_data(sample):
return {
"messages": [
@ -79,7 +78,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(
@ -124,8 +123,8 @@ trainer = SFTTrainer(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
gradient_checkpointing = True,
gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing
max_grad_norm = 0.3, # max gradient norm based on QLoRA paper
gradient_checkpointing_kwargs = {"use_reentrant": False},
max_grad_norm = 0.3, # QLoRA paper
warmup_ratio = 0.03,
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
max_steps = 60,
@ -154,7 +153,7 @@ trainer_stats = trainer.train()
model.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter", tokenizer)
tokenizer.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter")
# Benchmark lora adapter model performance.
# Benchmark lora adapter model.
model_name = "Unsloth lora adapter model"
FastVisionModel.for_inference(model)
avg_wer, avg_cer = ocr_evaluator.evaluate_model(
@ -176,7 +175,7 @@ base = find_lora_base_model(model)
print((base.__class__.__name__))
# Merge at default 16 bits.
# Merge at 16 bits.
model.save_pretrained_merged(
save_directory = "qwen2-ocr-merged-finetune-merge-16bit", tokenizer = tokenizer
)

12
tests/security/conftest.py
View file

@ -1,9 +1,4 @@
"""Shared fixtures for the security regression suite.
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.
"""
"""Security suite fixtures: an autouse network blocker refuses non-loopback socket.connect() so a regression reaching the internet fails loudly."""
from __future__ import annotations
@ -14,7 +9,7 @@ from pathlib import Path
import pytest
# Make `scripts/` importable so tests can grab scanner constants directly
# 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))
@ -66,8 +61,7 @@ class _BlockedSocket(socket.socket):
@pytest.fixture(scope = "session", autouse = True)
def network_blocker():
"""Session-scoped fixture; swaps `socket.socket` for a blocker, restored at
teardown so interleaved sessions see a clean module."""
"""Swap socket.socket for the blocker, restored at teardown."""
original = socket.socket
socket.socket = _BlockedSocket # type: ignore[assignment]
try:

26
tests/security/fixtures/_build.py
View file

@ -1,13 +1,4 @@
"""Deterministic builder for the wheel + sdist binary fixtures.
Not run from CI; the produced .whl / .tar.gz bytes are committed alongside
it. Re-run only when the IOC literal changes.
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).
"""
"""Deterministic builder for the committed wheel + sdist fixtures; re-run only when the IOC literal changes."""
from __future__ import annotations
@ -80,17 +71,12 @@ def _write_zip_member(zf: zipfile.ZipFile, name: str, data: bytes) -> None:
def _build_wheel(out_path: Path, *, name: str, payload_files: dict[str, bytes]) -> None:
"""Write a deterministic .whl at `out_path`.
`payload_files` maps archive-relative paths to their bytes. Standard
`.dist-info/METADATA`, `WHEEL`, and `RECORD` are added automatically.
"""
"""Write a deterministic .whl; .dist-info METADATA/WHEEL/RECORD are added automatically."""
dist_info = f"{name}-0.0.1.dist-info"
members: dict[str, bytes] = dict(payload_files)
members[f"{dist_info}/METADATA"] = WHEEL_METADATA.format(name = name).encode()
members[f"{dist_info}/WHEEL"] = WHEEL_FILE.encode()
# RECORD is intentionally minimal; the scanner only inspects file
# bodies, not hash integrity.
# RECORD is minimal; the scanner inspects file bodies, not hash integrity.
record_lines = []
for path in sorted(members):
record_lines.append(f"{path},,")
@ -106,11 +92,7 @@ def _build_wheel(out_path: Path, *, name: str, payload_files: dict[str, bytes])
def _build_sdist(out_path: Path, *, name: str, payload_files: dict[str, bytes]) -> None:
"""Write a deterministic .tar.gz sdist at `out_path`.
`payload_files` maps archive-relative paths to their bytes; a
leading `{name}-0.0.1/` prefix is added automatically.
"""
"""Write a deterministic .tar.gz sdist; a `{name}-0.0.1/` prefix is added automatically."""
prefix = f"{name}-0.0.1"
buf = io.BytesIO()
# gzip mtime fixed via mtime=0 (gzip member header).

12
tests/security/test_lint_workflow_triggers.py
View file

@ -1,10 +1,4 @@
"""Regression tests for scripts/lint_workflow_triggers.py.
Guards against future regressions that would re-introduce GHSA-g7cv-rxg3-hmpx
(TanStack) -class supply-chain vectors:
* pull_request_target (fork PR runs in base context).
* Shared cache keys between PR-triggered workflows and the publish workflow.
"""
"""Regression tests for scripts/lint_workflow_triggers.py, guarding GHSA-g7cv-rxg3-hmpx vectors."""
from __future__ import annotations
@ -104,7 +98,7 @@ def test_lint_rejects_shared_cache_key_between_pr_and_publish(tmp_path):
publish workflow is the TanStack cache-poisoning vector."""
wf = tmp_path / "wf"
wf.mkdir()
# PR-triggered: writes to a cache that the publish job will also restore.
# PR-triggered: writes a cache the publish job will also restore.
(wf / "pr-build.yml").write_text(
"name: pr-build\n"
"on:\n"
@ -118,7 +112,7 @@ def test_lint_rejects_shared_cache_key_between_pr_and_publish(tmp_path):
" path: node_modules\n"
" key: shared-cache-v1\n"
)
# Publish workflow with the IDENTICAL cache key -- the actual attack pattern.
# Publish workflow with the IDENTICAL cache key (the attack pattern).
(wf / "release-desktop.yml").write_text(
"name: release-desktop\n"
"on:\n"

53
tests/security/test_lockfile_supply_chain_audit.py
View file

@ -1,9 +1,4 @@
"""Regression tests for `scripts/lockfile_supply_chain_audit.py`.
The auditor is fully offline (file reads only); tests run the script
as a subprocess against the fixture lockfiles plus an inline
`Cargo.lock` constructed in a tmpdir.
"""
"""Regression tests for the offline `scripts/lockfile_supply_chain_audit.py`."""
from __future__ import annotations
@ -51,8 +46,7 @@ def _run_auditor(
def test_malicious_lockfile_exits_1(tmp_path):
"""Fixture combines a non-registry resolved URL, a known IOC substring
(`filev2.getsession.org`), and a missing integrity hash: auditor exits 1."""
"""Non-registry URL + IOC substring + missing integrity hash -> auditor exits 1."""
fixture = FIXTURES / "malicious_lockfile.json"
assert fixture.is_file()
proc = _run_auditor(root = tmp_path, npm_lockfiles = [fixture])
@ -82,9 +76,7 @@ def test_clean_lockfile_exits_0(tmp_path):
def test_audit_npm_lockfile_direct_call_findings():
"""In-process call to `audit_npm_lockfile()` returns the same
finding shape we expect the subprocess to emit.
"""
"""In-process audit_npm_lockfile() returns the same findings as the subprocess."""
findings = lsa.audit_npm_lockfile(FIXTURES / "malicious_lockfile.json")
kinds = {f.kind for f in findings}
assert "non-registry-resolved-url" in kinds
@ -132,10 +124,7 @@ def test_npm_ioc_strings_contains_may12_additions():
reason = "Fork 1 (BLOCKED_NPM_VERSIONS in auditor) not merged yet",
)
def test_lockfile_auditor_blocked_versions_match_scanner():
"""The auditor's BLOCKED_NPM_VERSIONS must mirror the scanner's
table verbatim (Fork 1's plan says to duplicate with a sync
comment until the next PR factors them into a shared module).
"""
"""Auditor's BLOCKED_NPM_VERSIONS must mirror the scanner's table verbatim."""
from scripts import scan_npm_packages as snp
assert (
lsa.BLOCKED_NPM_VERSIONS == snp.BLOCKED_NPM_VERSIONS
@ -164,9 +153,7 @@ checksum = "0000000000000000000000000000000000000000000000000000000000000000"
def test_malicious_cargo_lockfile_refused(tmp_path):
"""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."""
"""git+https:// Cargo source trips non-registry-cargo-source; --strict makes it blocking."""
lockfile = tmp_path / "Cargo.lock"
lockfile.write_text(_MALICIOUS_CARGO_LOCK)
proc = _run_auditor(
@ -182,11 +169,7 @@ def test_malicious_cargo_lockfile_refused(tmp_path):
def test_malicious_cargo_lockfile_default_mode_advisory(tmp_path):
"""Default (non-strict) mode classifies `non-registry-cargo-source`
as advisory: the finding is still emitted as a `::warning::`
annotation but the process exits 0 so the build is not gated.
Regression test for the advisory/strict split.
"""
"""Default mode emits non-registry-cargo-source as advisory ::warning:: but exits 0."""
lockfile = tmp_path / "Cargo.lock"
lockfile.write_text(_MALICIOUS_CARGO_LOCK)
proc = _run_auditor(
@ -219,14 +202,11 @@ def test_audit_cargo_lockfile_direct_call(tmp_path):
def test_gha_escape_collapses_finding_to_one_line():
"""`_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."""
"""_gha_escape() encodes \\n/\\r/% so GHA annotations aren't truncated; % must escape first."""
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"
# Order regression: `%` must escape before `\n` so the literal
# text `a%b\nc` becomes `a%25b%0Ac`, not `a%250Ab%0Ac`.
# Order regression: `%` must escape before `\n`, else escapes double-encode.
assert lsa._gha_escape("a%b\nc") == "a%25b%0Ac"
f = lsa.Finding(
@ -244,9 +224,7 @@ def test_gha_escape_collapses_finding_to_one_line():
def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
"""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."""
"""Advisory ::warning:: must be one physical line (%0A-escaped). Regression for PR #5604."""
lockfile = tmp_path / "Cargo.lock"
lockfile.write_text(_MALICIOUS_CARGO_LOCK)
proc = _run_auditor(
@ -259,8 +237,7 @@ def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
"expected at least one ::warning:: annotation; " f"stderr was:\n{proc.stderr}"
)
for line in warning_lines:
# Single physical line: kind, package, detail all present
# via %0A escape, not split across stderr lines.
# One physical line: kind/package/detail joined via %0A, not split.
assert "%0A" in line, (
f"::warning:: line has no %0A escape; multi-line text "
f"would be truncated by GH Actions:\n{line}"
@ -276,9 +253,7 @@ 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` 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."""
"""SF4: a short/boolean UNSLOTH_LOCKFILE_AUDIT_SKIP is rejected; a real justification is honored."""
fixture = FIXTURES / "clean_lockfile.json"
# Case 1 -- "1" rejected, audit RUNS.
@ -300,9 +275,9 @@ def test_skip_env_var_with_short_value_rejected(tmp_path):
combined_bad = proc_bad.stdout + proc_bad.stderr
assert "::warning::" in combined_bad, combined_bad
assert "REQUIRES a justification" in combined_bad, combined_bad
# Audit actually ran (saw the per-file banner).
# Per-file banner proves the audit ran.
assert "[lockfile-audit] npm:" in combined_bad, combined_bad
# Fixture is clean, so exit 0 -- but the audit was performed.
# Clean fixture -> exit 0, but the audit was performed.
assert proc_bad.returncode == 0, (
f"expected rc 0 on clean fixture, got {proc_bad.returncode}\n"
f"--- stdout ---\n{proc_bad.stdout}\n"
@ -330,7 +305,7 @@ def test_skip_env_var_with_short_value_rejected(tmp_path):
assert "::warning::" in combined_ok
assert "skipped" in combined_ok.lower()
assert "ticket-5397" in combined_ok
# Skip path means the audit body never ran (no "npm:" banner).
# Skip path: no "npm:" banner means the audit body never ran.
assert "[lockfile-audit] npm:" not in combined_ok, combined_ok
# Case 3 -- the booleanish tokens are ALL rejected.

16
tests/security/test_new_install_scripts.py
View file

@ -1,10 +1,7 @@
"""Regression tests for `scripts/check_new_install_scripts.py`.
The fixture lockfiles are tiny dicts written to `tmp_path` so the
tests stay self-contained. The session-wide `network_blocker` fixture
in conftest.py refuses any real-world socket connect; the scanner
treats that block as "registry unreachable, emit finding anyway",
which is the offline-safe path under test.
Lockfiles are tiny dicts in tmp_path; the network_blocker fixture forces the
scanner's offline path (registry unreachable -> emit finding anyway).
"""
from __future__ import annotations
@ -45,8 +42,6 @@ def _write(path: Path, content: dict) -> Path:
# Lockfile fixtures
def _v3_lockfile(packages: dict) -> dict:
return {
"name": "unsloth-theme",
@ -69,8 +64,6 @@ def _v2_lockfile(packages: dict, dependencies: dict) -> dict:
# Tests
def test_no_new_install_scripts_exit_0(tmp_path: Path):
"""If base == head, nothing new can have been added."""
same = _v3_lockfile(
@ -137,7 +130,7 @@ def test_existing_dep_with_postinstall_ignored(tmp_path: Path):
},
}
head_pkgs = dict(base_pkgs)
# An ENTIRELY UNRELATED non-install-script dep is added in head.
# Add an unrelated non-install-script dep in head.
head_pkgs["node_modules/lodash"] = {
"version": "4.17.21",
"resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.21.tgz",
@ -188,8 +181,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
)
assert "v2-postinstall-dep" in result.stderr
# And again: same packages dict but lockfileVersion 3 -- should
# produce the same finding shape.
# Same packages as lockfileVersion 3 must give the same finding.
base_v3 = _write(tmp_path / "base_v3.json", _v3_lockfile(base_pkgs))
head_v3 = _write(tmp_path / "head_v3.json", _v3_lockfile(head_pkgs))
result_v3 = _run(base_v3, head_v3)

67
tests/security/test_scan_npm_packages.py
View file

@ -1,10 +1,4 @@
"""Regression tests for `scripts/scan_npm_packages.py`.
These tests must run fully offline. The `network_blocker` fixture in
conftest.py refuses any non-loopback socket connect from the test
process; scanner subprocesses are invoked against fixtures that never
trigger an HTTP fetch.
"""
"""Regression tests for scripts/scan_npm_packages.py. Run fully offline (network_blocker fixture)."""
from __future__ import annotations
@ -21,7 +15,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
SCRIPT = REPO_ROOT / "scripts" / "scan_npm_packages.py"
FIXTURES = Path(__file__).resolve().parent / "fixtures"
# Import the module so we can introspect the IOC tables directly.
# Import the module to introspect IOC tables directly.
sys.path.insert(0, str(REPO_ROOT))
from scripts import scan_npm_packages as snp # noqa: E402
@ -46,9 +40,7 @@ def _run_scanner(lockfile: Path, *, timeout: int = 30) -> subprocess.CompletedPr
def test_malicious_lockfile_exits_1():
"""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."""
"""Structural IOCs alone (non-registry resolved URL + missing integrity) fail the scanner offline."""
fixture = FIXTURES / "structural_only_lockfile.json"
assert fixture.is_file(), fixture
proc = _run_scanner(fixture)
@ -57,23 +49,17 @@ def test_malicious_lockfile_exits_1():
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}"
)
combined = proc.stdout + proc.stderr
# The scanner aggregates structural findings into the summary
# rather than printing each one individually. Assert on the
# count + the FAIL banner instead.
# Scanner aggregates structural findings into the summary; assert on count + FAIL banner.
assert "2 structural finding(s)" in combined
assert "FAIL" in combined
# And confirm `parse_lockfile()` actually surfaces the right
# `pattern` codes via the in-process API.
# Confirm parse_lockfile() surfaces the right pattern codes via the in-process API.
entries, struct = snp.parse_lockfile(fixture)
patterns = {f.pattern for f in struct}
assert {"non-registry-resolved-url", "missing-integrity-hash"} <= patterns
def test_clean_lockfile_exits_0():
"""The clean fixture only contains entries that `parse_lockfile()`
skips entirely (workspace root + workspace `link` symlink +
nested fold-in), so the scanner exits 0 with no network access.
"""
"""Clean fixture has only entries parse_lockfile() skips, so the scanner exits 0 offline."""
fixture = FIXTURES / "clean_lockfile.json"
assert fixture.is_file(), fixture
proc = _run_scanner(fixture)
@ -110,8 +96,7 @@ def test_blocked_npm_versions_complete():
f"expected at least 22 @squawk/* entries (full safedep.io enumeration), "
f"got {len(squawk)}: {sorted(squawk)}"
)
# @squawk/mcp must cover the full malicious range 0.9.1 .. 0.9.5
# (safedep.io enumeration; we initially had only 0.9.5).
# @squawk/mcp must cover the full malicious range 0.9.1..0.9.5 (safedep.io enumeration).
assert {"0.9.1", "0.9.2", "0.9.3", "0.9.4", "0.9.5"} <= table["@squawk/mcp"]
uipath = [k for k in table if k.startswith("@uipath/")]
@ -119,7 +104,7 @@ def test_blocked_npm_versions_complete():
f"expected at least 64 @uipath/* entries (Aikido enumeration), "
f"got {len(uipath)}: {sorted(uipath)}"
)
# Anchor a known entry: the rpa-tool 0.9.5 version is in the published list.
# Anchor a known published entry.
assert "0.9.5" in table["@uipath/rpa-tool"]
# Aikido (May-12 wave): @mistralai/* npm scope (separate from PyPI mistralai).
@ -161,8 +146,7 @@ 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():
"""The pre-fetch hook must flag the malicious tanstack entry as
`blocked-known-malicious` (exit 1) without hitting the npm registry."""
"""Pre-fetch hook flags the malicious tanstack entry (exit 1) without hitting the npm registry."""
fixture = FIXTURES / "malicious_lockfile.json"
proc = _run_scanner(fixture, timeout = 10)
assert proc.returncode == 1
@ -176,9 +160,7 @@ def test_blocked_npm_versions_short_circuits_download():
def _extract_pkg_with_ioc(ioc: str, tmp_path: Path) -> Path:
"""Build a one-file npm package extract tree embedding `ioc` in
`package.json`. Returns the extract root.
"""
"""Build a one-file npm package extract tree embedding `ioc` in package.json; return its root."""
pkg_json = {
"name": "ioc-fixture",
"version": "0.0.1",
@ -194,8 +176,7 @@ def _extract_pkg_with_ioc(ioc: str, tmp_path: Path) -> Path:
def test_every_known_ioc_string_caught(tmp_path):
"""Embed each `KNOWN_IOC_STRINGS` entry in a one-file package tree and
confirm `scan_extracted_tree()` surfaces it. Guards against table drift."""
"""Each KNOWN_IOC_STRINGS entry must be surfaced by scan_extracted_tree(); guards table drift."""
iocs = snp.KNOWN_IOC_STRINGS
assert iocs, "KNOWN_IOC_STRINGS unexpectedly empty"
@ -223,8 +204,7 @@ def test_every_known_ioc_string_caught(tmp_path):
def test_parse_lockfile_structural_findings():
"""The structural-only fixture yields 2 structural findings and 0 entries
(both bad entries are `continue`d in `parse_lockfile()`)."""
"""Structural-only fixture yields 2 structural findings and 0 entries."""
entries, struct = snp.parse_lockfile(FIXTURES / "structural_only_lockfile.json")
assert entries == []
patterns = {f.pattern for f in struct}
@ -233,9 +213,8 @@ def test_parse_lockfile_structural_findings():
# ---------------------------------------------------------------------------
# Code-only scanning (_strip_js_noncode) -- comment FP reduction. The stripper
# must blank comments WITHOUT touching strings/regex/code, preserve geometry,
# and fail open on lexer confusion.
# Code-only scanning (_strip_js_noncode): blank comments WITHOUT touching
# strings/regex/code, preserve geometry, fail open on lexer confusion.
# ---------------------------------------------------------------------------
@ -256,7 +235,7 @@ def test_strip_blanks_line_and_block_comments():
def test_strip_keeps_url_in_string_and_template():
src = 'const a = "http://example.com/x";\nconst b = `http://${h}//y`; go();'
out = _strip(src)
assert out == src # nothing is a comment; must be byte-identical
assert out == src # nothing is a comment -> byte-identical
assert "http://example.com/x" in out and "//y" in out
@ -278,11 +257,11 @@ def test_strip_preserves_assigned_base64_payload():
def test_strip_fails_open_on_unterminated_block_comment():
src = "code(); /* never closed"
assert snp._strip_js_noncode(src) == src # unchanged -> still fully scanned
assert snp._strip_js_noncode(src) == src # fail open: unchanged, still fully scanned
def test_strip_only_applies_to_js_family():
# A `//`-containing JSON/YAML string must be left intact (wrong lexer).
# A `//`-containing JSON/YAML string must be left intact (JS lexer must not apply).
PKG = snp.PackageEntry(
name = "x",
version = "1.0.0",
@ -292,8 +271,7 @@ def test_strip_only_applies_to_js_family():
)
# scan_text_blob strips for .js but not for .json.
yaml_like = 'url: "http://h" # a yaml comment, not JS\n'
# No assertion on findings here -- just that the JS lexer is not applied to
# non-JS suffixes (covered indirectly: stripper is gated on suffix).
# Verify the stripper is gated on suffix (JS lexer not applied to non-JS suffixes).
assert "".endswith(snp._JS_FAMILY_SUFFIXES) is False
assert ".js" in snp._JS_FAMILY_SUFFIXES and ".json" not in snp._JS_FAMILY_SUFFIXES
@ -328,7 +306,7 @@ def test_payload_entirely_in_comment_is_suppressed():
src = f'/* var f = new Function("{_BLOB}"); */ var ok = 1;'
js = snp.scan_text_blob(_PKG, "m.js", src)
assert js == [] # blanked -> clean
# Control: same bytes scanned as non-JS (unstripped) WOULD flag.
# Control: same bytes as non-JS (unstripped) WOULD flag.
txt = snp.scan_text_blob(_PKG, "m.txt", src)
assert any(f.pattern == "obfuscated-blob" for f in txt)
@ -389,7 +367,7 @@ def test_baseline_suppresses_listed_but_not_new_pattern(tmp_path):
baseline = snp._load_baseline(str(bl))
listed = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "cred-surface-host (outbound)")
# A NEW kind of finding in the SAME file is a different pattern -> not suppressed.
# A new pattern in the same file must NOT be suppressed.
new_kind = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "obfuscated-blob")
active, suppressed = snp._partition_baseline([listed, new_kind], baseline)
assert listed in suppressed
@ -407,13 +385,12 @@ def test_write_then_load_baseline_roundtrip(tmp_path):
assert n == 1 # dedup + MEDIUM excluded
keys = snp._load_baseline(str(bl))
assert (snp._norm_pkg_name("evil@1.0.0"), "a.js", "obfuscated-blob") in keys
# MEDIUM was below the HIGH threshold -> not written.
# MEDIUM below HIGH threshold -> not written.
assert all(k[2] != "js-env-token" for k in keys)
def test_committed_baseline_is_empty_and_valid():
# The shipped baseline must parse and (by design) suppress nothing: the
# live corpus is clean, so the gate can run enforcing with an empty list.
# Shipped baseline must parse and (by design) suppress nothing: the live corpus is clean.
path = REPO_ROOT / "scripts" / "scan_npm_packages_baseline.json"
assert path.is_file()
doc = json.loads(path.read_text(encoding = "utf-8"))

110
tests/security/test_scan_packages.py
View file

@ -1,9 +1,5 @@
"""Regression tests for `scripts/scan_packages.py`.
`download_packages` reaches PyPI; to stay offline we drive the in-process
`scan_archive` helper against the wheel/sdist fixtures under
`tests/security/fixtures/`.
"""
"""Regression tests for `scripts/scan_packages.py`, driving the offline
`scan_archive` helper against fixtures under `tests/security/fixtures/`."""
from __future__ import annotations
@ -23,22 +19,13 @@ sys.path.insert(0, str(REPO_ROOT))
from scripts import scan_packages as sp # noqa: E402
# ---------------------------------------------------------------------------
# Fixture sanity.
# ---------------------------------------------------------------------------
def test_fixture_files_exist():
for name in ("malicious_wheel.whl", "clean_wheel.whl", "malicious_sdist.tar.gz"):
assert (FIXTURES / name).is_file(), name
def test_fixture_bytes_are_deterministic(tmp_path):
"""Re-running `_build.py` must produce byte-identical archives.
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.
"""
"""Re-running `_build.py` must produce byte-identical archives (deterministic builds)."""
# Snapshot committed hashes.
expected: dict[str, str] = {}
for name in ("malicious_wheel.whl", "clean_wheel.whl", "malicious_sdist.tar.gz"):
@ -47,11 +34,11 @@ def test_fixture_bytes_are_deterministic(tmp_path):
# Rebuild into a sibling dir to avoid clobbering the committed files.
rebuild_dir = tmp_path / "rebuild"
rebuild_dir.mkdir()
# The build helper writes to its own directory; copy + patch HERE.
# The build helper writes to its own dir; copy + patch HERE.
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 shim.
shim = rebuild_dir / "run.py"
shim.write_text(
"import sys, pathlib\n"
@ -80,11 +67,6 @@ def test_fixture_bytes_are_deterministic(tmp_path):
)
# ---------------------------------------------------------------------------
# scan_archive() against the fixture wheel + sdist.
# ---------------------------------------------------------------------------
def _critical_or_high(findings) -> list:
return [f for f in findings if f.severity in (sp.CRITICAL, sp.HIGH)]
@ -97,7 +79,6 @@ def test_malicious_wheel_triggers_critical():
assert findings, "no findings on malicious wheel; scanner regression"
blockers = _critical_or_high(findings)
assert blockers, f"no CRITICAL/HIGH findings: {[str(f) for f in findings]}"
# At least one finding must reference setup.py.
assert any("setup.py" in f.filename for f in blockers)
@ -119,11 +100,7 @@ def test_clean_wheel_no_findings():
assert findings == [], f"unexpected findings on clean wheel: {[str(f) for f in findings]}"
# ---------------------------------------------------------------------------
# Fork 1 constants -- gated on availability.
# ---------------------------------------------------------------------------
_BLOCKED_AVAILABLE = hasattr(sp, "BLOCKED_PYPI_VERSIONS")
_MAY12_AVAILABLE = hasattr(sp, "RE_MAY12_IOC")
@ -157,7 +134,7 @@ def test_re_may12_ioc_catches_each_literal():
pattern: re.Pattern = sp.RE_MAY12_IOC
for lit in expected_literals:
assert pattern.search(lit), f"RE_MAY12_IOC missed literal {lit!r}"
# Clean control: a plain string with none of the literals must not match.
# Clean control: a string with none of the literals must not match.
assert not pattern.search("import numpy as np")
@ -166,16 +143,13 @@ 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, the malicious wheel's
setup.py must produce a finding referencing the May-12 IOC string.
"""
"""Wired into check_py_file, the malicious wheel's setup.py must flag the May-12 IOC string."""
findings = sp.scan_archive(
str(FIXTURES / "malicious_wheel.whl"),
"malicious_fixture",
)
# 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 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).
_ioc_host = "git-tanstack." + "com"
_ioc_drop = "transformers." + "pyz"
hit = any(
@ -190,18 +164,13 @@ def test_may12_ioc_caught_by_scan_archive():
)
# ---------------------------------------------------------------------------
# Silent-failure-class hardening (Fork C).
# ---------------------------------------------------------------------------
def test_scan_packages_pip_download_failure_propagates(tmp_path):
"""A pip download failure must NOT be swallowed into `0 findings, exit 0`.
"""A pip download failure must exit 2 (SCAN INCOMPLETE), not `0 findings, exit 0`.
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.
"""
Feeds an unresolvable spec; the name is long/random so it can't resolve on any index."""
script = REPO_ROOT / "scripts" / "scan_packages.py"
assert script.is_file(), script
unresolvable = "pkg-that-does-not-exist-0123456789-fork-c-silentfail==0.0.0"
@ -222,9 +191,7 @@ def test_scan_packages_pip_download_failure_propagates(tmp_path):
def test_archive_corruption_produces_critical_finding(tmp_path):
"""SF1: a corrupted wheel was silently skipped by `except: continue` in
iter_archive_files; it must now yield a CRITICAL `archive_corrupted`.
"""
"""SF1: a corrupted wheel (once silently skipped) must 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
findings = sp.scan_archive(str(bad), "broken_fixture")
@ -235,7 +202,7 @@ def test_archive_corruption_produces_critical_finding(tmp_path):
)
assert all(f.severity == sp.CRITICAL for f in corrupted)
# Same check for a corrupted tarball.
# Same for a corrupted tarball.
bad_tar = tmp_path / "broken-0.0.1.tar.gz"
bad_tar.write_bytes(b"not-a-real-gzip-stream")
findings_tar = sp.scan_archive(str(bad_tar), "broken_fixture")
@ -246,9 +213,7 @@ def test_archive_corruption_produces_critical_finding(tmp_path):
)
# ---------------------------------------------------------------------------
# False-positive hardening: code-only scanning via _strip_noncode.
# ---------------------------------------------------------------------------
def test_strip_noncode_blanks_docstrings_and_comments_keeps_geometry():
@ -262,10 +227,9 @@ def test_strip_noncode_blanks_docstrings_and_comments_keeps_geometry():
out = sp._strip_noncode(src)
# Line geometry is byte-stable so evidence L<n> stays correct.
assert len(out.splitlines()) == len(src.splitlines())
# The dangerous-looking tokens lived only in docstrings/comments -> gone.
# Tokens lived only in docstrings/comments -> gone.
for needle in ("subprocess", "os.system", "eval(", "exec(", "reverse shell"):
assert needle not in out, needle
# Real code survives.
assert "x = 1" in out
assert "return x" in out
@ -288,7 +252,7 @@ def test_strip_noncode_falls_back_on_syntax_error():
def test_check_py_file_ignores_docstring_only_iocs():
# A file whose ONLY dangerous patterns live in a docstring must be clean.
# A file whose only dangerous patterns live in a docstring must be clean.
benign = (
'"""Usage:\n'
">>> import subprocess, urllib.request\n"
@ -299,7 +263,7 @@ def test_check_py_file_ignores_docstring_only_iocs():
)
findings = sp.check_py_file(benign, "pkg/_doc.py", "pkg")
assert findings == [], f"docstring IOCs should not flag: {[str(f) for f in findings]}"
# But the same payload as real code still flags.
# The same payload as real code still flags.
real = (
"import subprocess, urllib.request\n"
"subprocess.Popen(['sh','-c','id'])\n"
@ -310,15 +274,14 @@ def test_check_py_file_ignores_docstring_only_iocs():
def test_extract_evidence_multiline_reports_line():
# A DOTALL pattern that only matches across lines must still yield evidence
# (not an empty string) so a baseline entry is reviewable.
# A cross-line DOTALL match must still yield evidence so the baseline entry is reviewable.
content = "a = 1\ntime.sleep(\n 600\n)\n"
ev = sp._extract_evidence(content, sp.RE_ANTI_ANALYSIS)
assert ev and ev.startswith("L"), ev
def test_anti_analysis_no_longer_flags_cross_platform_code():
# Pure cross-platform code (the old platform.system FP) must be clean.
# Pure cross-platform code (the old platform.system false positive) must be clean.
crossplat = (
"import platform, subprocess\n"
"if platform.system() == 'Windows':\n"
@ -332,11 +295,9 @@ def test_anti_analysis_no_longer_flags_cross_platform_code():
def test_proc_self_status_read_flags_anti_analysis():
# Reading /proc/self/status (to scrape TracerPid) alongside a subprocess
# call is the classic anti-debug combination. The old `\b/proc/self/status\b`
# was a dead pattern (\b adjacent to "/" is unsatisfiable); the lookbehind
# fix makes it fire. No TracerPid/ptrace token here so only the /proc path
# can supply the anti-analysis signal.
# Reading /proc/self/status + a subprocess call is the classic anti-debug combo.
# The old `\b/proc/self/status\b` was unsatisfiable (\b adjacent to "/"); the
# lookbehind fix makes it fire. No TracerPid/ptrace token, so only /proc signals it.
payload = (
"import subprocess\n"
"with open('/proc/self/status') as fh:\n"
@ -350,8 +311,7 @@ def test_proc_self_status_read_flags_anti_analysis():
def test_proc_self_status_pattern_is_live():
# Direct regex check across the common call forms; the leading \b made all
# of these unsatisfiable before the fix.
# Common call forms; the leading \b made all of these unsatisfiable before the fix.
for s in (
'open("/proc/self/status")',
"cat /proc/self/status",
@ -362,11 +322,6 @@ def test_proc_self_status_pattern_is_live():
assert not sp.RE_ANTI_ANALYSIS.search("if platform.system() == 'Linux': pass")
# ---------------------------------------------------------------------------
# Baseline allowlist.
# ---------------------------------------------------------------------------
def _mk(sev, pkg, fname, check):
return sp.Finding(sev, pkg, fname, check, "evidence")
@ -376,7 +331,7 @@ def test_baseline_key_version_stable_but_path_specific():
b = _mk(sp.CRITICAL, "Requests", "requests-3.0.0/requests/sessions.py", "X")
# Same package-relative path across versions -> same key (stable).
assert sp._finding_key(a) == sp._finding_key(b)
# Same basename in a DIFFERENT path -> different key (no over-suppression).
# Same basename in a different path -> different key (no over-suppression).
c = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/vendor/sessions.py", "X")
assert sp._finding_key(a) != sp._finding_key(c)
@ -385,7 +340,7 @@ def test_fstring_statement_is_not_blanked():
# A bare f-string evaluates at import, so it must stay scannable.
src = "f\"{__import__('os').system('id')}\"\n"
assert "__import__" in sp._strip_noncode(src)
# A plain bare docstring IS blanked.
# A plain bare docstring is blanked.
plain = "'a docstring mentioning subprocess.Popen'\n"
assert "subprocess" not in sp._strip_noncode(plain)
@ -395,7 +350,7 @@ def test_exec_with_payload_hidden_in_docstring_flagged():
src = '"""' + blob + '"""\nimport os\nexec(__doc__)\n'
findings = sp.check_py_file(src, "pkg/mod.py", "pkg")
assert any("hidden in a docstring" in f.check for f in findings)
# No exec/eval -> the blanked blob does not produce that finding.
# No exec/eval -> the blanked blob produces no such finding.
src2 = '"""' + blob + '"""\nimport os\n'
findings2 = sp.check_py_file(src2, "pkg/mod.py", "pkg")
assert not any("hidden in a docstring" in f.check for f in findings2)
@ -437,10 +392,8 @@ def test_load_baseline_missing_file_is_empty():
assert sp._load_baseline("/nonexistent/path/bl.json") == set()
# ---------------------------------------------------------------------------
# sdist fallback: preserve coverage of sdist-only packages without building.
# All offline -- PyPI JSON / download are mocked.
# ---------------------------------------------------------------------------
# sdist fallback: cover sdist-only packages without building. All offline
# -- PyPI JSON / download are mocked.
class _FakeResp:
@ -517,10 +470,10 @@ def test_requires_dist_skips_extras():
],
)
specs = sp._requires_dist_names(meta, None)
# Version constraints are preserved so a pinned dep is fetched, not latest.
# Version constraints preserved so a pinned dep is fetched, not latest.
assert "numpy>=1.20" in specs
assert "pyyaml>=5" in specs
# The extra-gated dep is skipped entirely (no torch under any form).
# The extra-gated dep is skipped entirely.
assert not any(sp._extract_pkg_name(s) == "torch" for s in specs)
@ -528,7 +481,7 @@ def test_download_sdist_direct_refuses_non_pypi_url(tmp_path):
meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://evil.example/x.tar.gz")])
fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
assert fpath is None and "non-PyPI" in err
assert list(tmp_path.iterdir()) == [] # nothing was written
assert list(tmp_path.iterdir()) == [] # nothing written
def test_download_sdist_direct_no_sdist_published(tmp_path):
@ -558,8 +511,7 @@ def test_download_sdist_direct_size_cap(tmp_path, monkeypatch):
def test_per_spec_genuine_failure_is_recorded_error(tmp_path, monkeypatch):
# A spec that fails pip but HAS a wheel on PyPI is a genuine error (-> exit 2),
# never silently swallowed.
# A spec that fails pip but HAS a wheel on PyPI is a genuine error (-> exit 2).
class _Proc:
returncode = 1
stderr = "ResolutionImpossible"

135
tests/studio/_playwright_robust.py
View file

@ -1,17 +1,9 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""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.
Importable directly by the standalone scripts:
sys.path.insert(0, str(Path(__file__).parent))
from _playwright_robust import (...)
Does NOT depend on pytest -- both consumers run as plain Python.
"""Shared CI-runner workarounds for the Studio Playwright tests (Chromium flags,
view-transition killer, page recovery, post-action response wait). Imported
directly by the standalone scripts; does NOT depend on pytest.
"""
from __future__ import annotations
@ -27,15 +19,12 @@ from pathlib import Path
from typing import Any, Callable
# Chromium launch args.
#
# 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.
#
# `--single-process` is darwin-only: the documented free-runner fix for the
# pipeTransport.js JSON-RPC crash; on Win/Linux it destabilises the renderer.
# Throttling flags stop Chromium deprioritising CPU/timers when it thinks the
# headless window is backgrounded (run 25586583024 stalled inference + render).
# TranslateUI strips a pointer-intercepting popup; ipc-flooding-protection off
# lets rapid clicks through during the slider sweep.
# `--single-process` is darwin-only (fixes the pipeTransport.js JSON-RPC crash);
# on Win/Linux it destabilises the renderer.
_BASE_CHROMIUM_ARGS = (
"--disable-dev-shm-usage",
"--no-sandbox",
@ -49,8 +38,8 @@ _BASE_CHROMIUM_ARGS = (
def chromium_launch_args(platform: str | None = None) -> list[str]:
"""Return Chromium launch args for `platform` (defaults to `sys.platform`;
pass a string to test the darwin branch on Linux)."""
"""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":
@ -58,13 +47,11 @@ def chromium_launch_args(platform: str | None = None) -> list[str]:
return args
# Init scripts injected into every Playwright context.
#
# Init script injected into every Playwright context.
# 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.
# pointer events after each theme/route swap, so Playwright reports
# `<html> intercepts pointer events` on the next click. Killing the
# pseudo-elements + shimming startViewTransition synchronously fixes both.
# Idempotent and safe to install on every page.
_VIEW_TRANSITION_KILLER_JS = """
(function () {
@ -107,10 +94,9 @@ def install_view_transition_killer(ctx: Any) -> None:
# Server health pre-flight.
#
# 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.
# On the macos-14 free runner /api/health can return 200 while /api/auth still
# 503s (auth DB mid-migration); this 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]:
@ -133,9 +119,8 @@ def wait_for_health(
timeout: float = 30.0,
info: Callable[[str], None] | None = None,
) -> bool:
"""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."""
"""Poll {base_url}/api/health until status==200; True on success, False on
timeout, never raises. Diagnostic only (the workflow's wait is authoritative)."""
deadline = time.monotonic() + timeout
last_status: int | None = None
last_body: dict | None = None
@ -145,8 +130,7 @@ def wait_for_health(
timeout = 3.0,
)
last_status, last_body = status, body
# `chat_only` and `status` keys both exist; prefer status==healthy
# but accept any 200 -- different Studio builds report differently.
# Accept any 200 -- different Studio builds report status differently.
if status == 200:
if info is not None:
info(f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}")
@ -160,11 +144,8 @@ def wait_for_health(
return False
# Page recovery.
#
# 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.
# Page recovery: if the page died mid-test, open a fresh one in the same context
# (localStorage auth survives); otherwise leave it alone. Optionally re-navigates.
def recover_or_replace_page(
@ -176,9 +157,8 @@ def recover_or_replace_page(
settle_networkidle: bool = True,
info: Callable[[str], None] | None = None,
) -> Any:
"""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."""
"""Return a usable page, replacing `page` if closed; optionally navigate to
`goto_url`. Recovery errors are logged and swallowed for the caller to retry."""
try:
if page.is_closed():
page = ctx.new_page()
@ -214,10 +194,9 @@ 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. 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."""
"""Click + wait for the matching XHR/fetch response; (status, None) on success
or (None, exception) on capture failure. Falls back to a fire-and-forget click
so the outer retry loop runs. Callers check `status >= 400`."""
try:
with page.expect_response(
lambda r: url_substr in r.url and r.request.method == method,
@ -240,12 +219,10 @@ def click_and_wait_for_response(
# Console-error / page-error filtering.
#
# - 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: CI-infra JS errors 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)",
@ -291,9 +268,8 @@ def dump_diagnostics(
info: Callable[[str], None] | None = None,
extra: dict | None = None,
) -> None:
"""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)."""
"""Write a screenshot + JSON sidecar (URL/title/body/storage) under art_dir.
Diagnostic only, never raises; both best-effort."""
art = Path(art_dir)
try:
art.mkdir(parents = True, exist_ok = True)
@ -343,12 +319,10 @@ def dump_diagnostics(
# Bounded in-page fetch.
#
# `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`.
# `page.evaluate(...)` has no `timeout=`, so a stuck fetch hangs the script until
# the runner timeout (run 25696797934 / PR #5387 burned 27+ min). evaluate_fetch
# wraps the fetch in an AbortController.signal so the JS side always resolves --
# real response, or synthetic `{status: 0, error: "AbortError..."}` after timeout_ms.
def evaluate_fetch(
page: Any,
url: str,
@ -360,11 +334,10 @@ def evaluate_fetch(
transport_retries: int = 2,
transport_backoff_ms: int = 250,
) -> dict[str, Any]:
"""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."""
"""Run `fetch(url, opts)` in the page with an AbortSignal deadline; returns
`{"status", "body", "error"}` (status==0 + AbortError on timeout). Treat
status==0 or non-None error as transport failure. `body` may be str (verbatim)
or dict/list (JSON-encoded); pass headers explicitly for Content-Type/Auth."""
body_arg: str | None
if body is None:
body_arg = None
@ -405,11 +378,9 @@ def evaluate_fetch(
"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.
# Retry transport failures only: status != 0 (real HTTP) and AbortError
# (caller's deadline) propagate; status==0 (stale-keepalive / "Failed to
# fetch" after auth rotation) retries after backoff to evict the dead socket.
last: dict[str, Any] | None = None
attempts = max(1, int(transport_retries) + 1)
for attempt in range(attempts):
@ -440,22 +411,18 @@ def evaluate_fetch(
# Wall-clock watchdog.
#
# 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).
# A browser wedge (CPU-pinned JS, silent renderer crash, asyncio deadlock) can
# still hang the script. A daemon Timer calls os._exit(2) after deadline_s; 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) but under the 30-min 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
it so the caller can `.cancel()` on clean exit; being daemonised, it also
dies with the process if the script exits first."""
"""Start a daemon Timer that hard-exits the process at `deadline_s`; returned
so the caller can `.cancel()` on clean exit (daemonised, dies with process)."""
def _kaboom() -> None:
msg = (

8
tests/studio/install/conftest.py
View file

@ -1,13 +1,7 @@
# 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.
"""
"""Pytest config for studio/install tests: add studio/ to sys.path so `backend` imports work from the repo root."""
from __future__ import annotations

44
tests/studio/install/smoke_test_parallel_studio_home.py
View file

@ -1,29 +1,14 @@
#!/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 run side by side.
"""Smoke test (#5190 env-override path): N parallel install.sh runs with
distinct UNSLOTH_STUDIO_HOME values must produce N isolated installs whose
backends run side by side. Checks install-time layout/isolation + clean HOME,
then runtime /api/health, distinct studio_root_id, and per-venv PIDs.
Covers the env-override path from #5190:
install-time
* 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 ``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.
Integration smoke runner (not pytest); ~1 minute on a warm uv cache. Invoke:
Integration runner (not pytest), ~1 minute on a warm uv cache. Invoke:
python tests/studio/install/smoke_test_parallel_studio_home.py [--n 6 --keep]
Exits 0 PASS / 1 FAIL / 2 error. Artifacts removed on PASS unless --keep;
kept on FAIL/ERROR for inspection.
Exits 0 PASS / 1 FAIL / 2 error. Artifacts kept on FAIL/ERROR or with --keep.
"""
from __future__ import annotations
@ -94,15 +79,11 @@ def _launch_backend(
log_path.parent.mkdir(parents = True, exist_ok = True)
env = os.environ.copy()
env["HOME"] = str(fake_home)
# Pin UNSLOTH_STUDIO_HOME (and clear the alias) so the child cannot
# inherit a Studio root from the caller's shell. Without this, a shell
# that already exports either var would override the per-label sys.prefix
# inference and every backend would resolve to the caller's install.
# Pin UNSLOTH_STUDIO_HOME and clear the alias so the child can't inherit a
# Studio root from the caller's shell and resolve to the wrong install.
env["UNSLOTH_STUDIO_HOME"] = str(studio_home)
env.pop("STUDIO_HOME", None)
# The child process inherits a dup of stdout via Popen, so closing the
# parent's handle when this function returns is safe and avoids relying
# on GC timing to release the fd.
# Popen dups stdout into the child, so closing the parent's handle here is safe.
with log_path.open("w") as fh:
return subprocess.Popen(
[
@ -208,11 +189,8 @@ def _check_fake_home_clean(fake_home: Path) -> None:
def _backend_pid_python(pid: int) -> Path | None:
"""Resolve the binary backing a running PID. Linux exposes this at
/proc/PID/exe; on platforms without /proc (macOS, BSD, Windows) we
skip this check and rely on the install-time symlink + studio.conf
invariants to catch cross-resolution. Returns None when /proc is
unavailable so the caller can skip cleanly."""
"""Resolve the binary backing a running PID via /proc/PID/exe (Linux only);
returns None elsewhere so the caller skips this check cleanly."""
if sys.platform != "linux":
return None
proc_exe = Path(f"/proc/{pid}/exe")

42
tests/studio/install/test_cuda_repair.py
View file

@ -1,10 +1,6 @@
"""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.
"""
"""_ensure_cuda_torch reinstalls CUDA torch when an NVIDIA-host venv carries a ROCm
build (the pre-fix KFD gpu_id false positive), but leaves healthy CUDA / CPU / ROCm /
macOS / Windows untouched. Fully mocked -- no GPU required."""
import importlib.util
import sys
@ -14,7 +10,7 @@ from unittest.mock import MagicMock, patch
import pytest
# ── Load module under test (mirrors test_rocm_support.py) ────────────────────
# Load module under test (mirrors test_rocm_support.py).
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
@ -29,21 +25,14 @@ _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.
"""
"""subprocess.run side_effect: torch-classify probe (sys.executable, bytes
stdout) vs nvidia-smi version probe (smi path, text=True), keyed on the executable."""
def _run(cmd, *args, **kwargs):
result = MagicMock()
@ -78,9 +67,7 @@ def _run_cuda_repair(
):
"""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.
"""
cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it."""
env = {}
if rocm_marker:
env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
@ -122,7 +109,7 @@ def _index_url(mock_pip) -> str:
return args[args.index("--index-url") + 1]
# ── Repair fires only on the poisoning signature ─────────────────────────────
# Repair fires only on the poisoning signature.
class TestCudaRepairFires:
@ -136,13 +123,13 @@ class TestCudaRepairFires:
assert mock_pip.call_args.kwargs["constrain"] is False
def test_rocm_in_version_string_triggers_repair(self):
# AMD SDK / Radeon wheels may not set torch.version.hip but encode
# rocm in __version__; the probe prints "hip" for both.
# AMD SDK / Radeon wheels may encode rocm in __version__ without
# torch.version.hip; the probe prints "hip" for both.
mock_pip = _run_cuda_repair(torch_state = "hip")
assert mock_pip.call_count == 1
# ── No-op cases ──────────────────────────────────────────────────────────────
# No-op cases.
class TestCudaRepairSkips:
@ -192,8 +179,7 @@ class TestCudaRepairSkips:
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).
# CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU (mixed AMD+NVIDIA host on the AMD card).
mock_pip = _run_cuda_repair(cvd = "-1", torch_state = "hip")
mock_pip.assert_not_called()
@ -206,7 +192,7 @@ class TestCudaRepairSkips:
assert mock_pip.call_count == 1
# ── CUDA index ladder ────────────────────────────────────────────────────────
# CUDA index ladder.
class TestCudaIndexResolution:
@ -231,7 +217,7 @@ class TestCudaIndexResolution:
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.
# NVIDIA usable via /proc fallback, nvidia-smi absent.
mock_pip = _run_cuda_repair(smi_path = None)
assert "cu126" in _index_url(mock_pip)

52
tests/studio/install/test_gpu_detection_followups.py
View file

@ -1,18 +1,4 @@
"""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.
"""
"""GPU-detection follow-ups to PR 6174: NVIDIA precedence + /proc/driver/nvidia/gpus fallback ported to install_llama_prebuilt.py and setup.sh. Mocks/source-level only, no GPU."""
import importlib.util
import sys
@ -41,8 +27,7 @@ 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."""
"""Fake run_capture: rocminfo returns rocminfo_stdout, everything else empty."""
def _run_capture(cmd, *args, **kwargs):
exe = str(cmd[0]) if cmd else ""
@ -102,8 +87,7 @@ def _run_detect_host(
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.
# Don't let the host's CUDA_VISIBLE_DEVICES leak in unless the scenario sets it.
if env is None or "CUDA_VISIBLE_DEVICES" not in env:
prebuilt_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
return detect_host()
@ -239,11 +223,7 @@ class TestSetupShHardening:
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.
"""
"""The nvcc source-build search must be gated on _setup_nvidia_usable."""
anchor = setup_src.find('NVCC_PATH=""\n')
assert anchor >= 0
window = setup_src[anchor : anchor + 700]
@ -252,9 +232,7 @@ class TestSetupShHardening:
), "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)."""
"""_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so CVD ""/-1 suppresses NVIDIA before AMD gating."""
assert "_setup_cvd_hides_nvidia()" in setup_src
start = setup_src.find("_setup_has_usable_nvidia_gpu() {")
end = setup_src.find("\n}", start)
@ -284,9 +262,7 @@ class TestSetupShHardening:
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."""
"""A mirror base path containing "rocm"/"gfx" must not mislabel a cu*/cpu index; classification uses TORCH_INDEX_URL's leaf only."""
@pytest.fixture(scope = "class")
def install_src(self) -> str:
@ -343,8 +319,7 @@ _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."""
"""Drive _has_usable_nvidia_gpu with a mocked nvidia-smi that always reports a GPU; cvd=None unsets the env var."""
def fake_run(cmd, *args, **kwargs):
result = MagicMock()
@ -368,11 +343,7 @@ def _stack_nvidia_usable(cvd):
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."""
"""CVD ""/-1 hides every NVIDIA device; all three _has_usable_nvidia_gpu impls must report not-usable so AMD/CPU routes run."""
def test_python_unset_cvd_is_usable(self):
assert _stack_nvidia_usable(None) is True
@ -393,9 +364,7 @@ class TestHiddenCvdNotUsable:
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)."""
"""Mixed host, NVIDIA hidden via CVD=-1: _has_rocm_gpu must pass the NVIDIA guard and return True (pre-fix it ignored CVD)."""
def fake_run(cmd, *args, **kwargs):
result = MagicMock()
@ -420,8 +389,7 @@ class TestHiddenCvdNotUsable:
@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"."""
"""Extract shell functions, run the usable-GPU one against a fake nvidia-smi; return "usable"/"not_usable"."""
import os as _os
import subprocess as sp

11
tests/studio/install/test_hf_auth.py
View file

@ -1,11 +1,4 @@
"""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.
"""
"""HF auth on the llama.cpp prebuilt installer: auth_headers sends HF_TOKEN to huggingface.co only, and a redirect handler strips Authorization on cross-host redirects. Offline."""
import importlib.util
import sys
@ -33,7 +26,7 @@ GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases"
def _headers(url, env):
"""auth_headers under a fully controlled token environment."""
"""Call 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:

71
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -238,8 +238,7 @@ def _mk_source_tarball(path: Path, tag: str) -> None:
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.
# A mix build's merge commit 404s on codeload, so hydrate must fetch the release asset.
commit = "a" * 40
archive_path = tmp_path / "merged-source.tar.gz"
_mk_source_tarball(archive_path, f"b9000-mix-{commit[:7]}")
@ -881,8 +880,7 @@ def write_linux_install_shape(install_dir: Path) -> None:
(install_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
(runtime_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8")
(runtime_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
# Mirror the runtime payload health groups in install_llama_prebuilt.py:
# libllama-common.so* was added by PR #5135 and is required.
# libllama-common.so* (PR #5135) is a required runtime payload health group.
(runtime_dir / "libllama-common.so.0").write_bytes(b"DLL")
(runtime_dir / "libllama.so.0").write_bytes(b"DLL")
(runtime_dir / "libggml.so.0").write_bytes(b"DLL")
@ -1279,10 +1277,7 @@ def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_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
matching and skip the reinstall that drops cudart in."""
"""A paired cudart bundle (#5106) marks the install stale unless cudart64_* and cublas64_* are on disk."""
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
write_windows_install_shape(
@ -1380,10 +1375,7 @@ def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_p
(install_dir / "build" / "bin" / "Release" / "cudart64_12.dll").unlink()
assert existing_install_matches_plan(install_dir, host, plan) is False
# cublasLt missing -- stale, must reinstall. The upstream cudart
# bundle ships all three of cudart / cublas / cublasLt; a user with
# cudart + cublas but no cublasLt is still missing a required GPU
# initialisation DLL and Studio must refresh the install.
# cublasLt missing -- stale, must reinstall (all three DLLs are required).
write_windows_install_shape(
install_dir,
include_llama_dll = True,
@ -1395,11 +1387,7 @@ def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_p
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
would loop on reinstall forever because install_from_archives has no
cudart source to drop in."""
"""With no paired runtime archive, a legacy install lacking cudart must still pass (else reinstall loops)."""
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
write_windows_install_shape(
@ -1475,11 +1463,7 @@ def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(
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
that drops the cudart DLLs in. This is the install-cache half of the
#5106 fix -- the health-check half lives in the test above."""
"""A pre-#5322 CUDA install must go stale once the choice gains a runtime archive (#5106 fingerprint half)."""
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
write_windows_install_shape(
@ -1554,7 +1538,7 @@ def test_existing_install_fingerprint_changes_when_cudart_pair_added(tmp_path: P
},
)
# Install metadata was written for the legacy (no-pair) choice.
# Metadata written for the legacy (no-pair) choice.
write_prebuilt_metadata(
install_dir,
requested_tag = "latest",
@ -1565,10 +1549,7 @@ def test_existing_install_fingerprint_changes_when_cudart_pair_added(tmp_path: P
prebuilt_fallback_used = False,
)
# New plan offers the paired choice -- fingerprint must differ so
# the install is refreshed. The health check would also catch this
# because cudart64_*.dll is missing on disk; we test the fingerprint
# half explicitly by comparing the two fingerprints directly.
# The paired choice's fingerprint must differ from the legacy one so the install refreshes.
legacy_fingerprint = INSTALL_LLAMA_PREBUILT.expected_install_fingerprint(
llama_tag = "b9001",
release_tag = "release-1",
@ -2384,8 +2365,7 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete(tmp_
is True
)
# Remove convert_hf_to_gguf.py (checked by confirm_install_tree but not
# runtime_payload_is_healthy) and verify the guard catches it
# Remove convert_hf_to_gguf.py (confirm_install_tree checks it; runtime health does not).
(install_dir / "convert_hf_to_gguf.py").unlink()
assert (
existing_install_matches_choice(
@ -2489,12 +2469,7 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete_maco
def test_paired_runtime_dll_patterns_excludes_executables() -> None:
"""The paired runtime archive must only contribute CUDA DLLs to
the install. The narrow pattern list -- not the broad
runtime_patterns_for_choice ``*.exe`` / ``*.dll`` -- is what
prevents a malformed cudart bundle from overwriting
llama-server.exe at install time.
"""
"""The paired runtime archive must contribute only CUDA DLLs (no *.exe/*.dll) so it can't overwrite binaries."""
paired_runtime_dll_patterns = INSTALL_LLAMA_PREBUILT.paired_runtime_dll_patterns
paired_choice = AssetChoice(
repo = "x",
@ -2538,10 +2513,7 @@ def test_paired_runtime_dll_patterns_excludes_executables() -> 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``.
"""
"""A malformed runtime archive with llama-server.exe must NOT replace the main archive's binary."""
install_from_archives = INSTALL_LLAMA_PREBUILT.install_from_archives
work = tmp_path / "work"
@ -2727,13 +2699,7 @@ def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(tmp_path: Path)
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``
(cu13 layout), conda-style ``nvidia/<pkg>/Library/bin``, plus
``torch/lib``. Otherwise installer preflight and backend launch
can disagree about which DLLs are actually present.
"""
"""Installer DLL discovery must scan the same path set as the backend (cu12/cu13/conda layouts + torch/lib)."""
import site as _site
python_runtime_dirs = INSTALL_LLAMA_PREBUILT.python_runtime_dirs
@ -2782,8 +2748,7 @@ def _nvidia_linux_host():
def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256):
"""Drive validate_prebuilt_choice with every heavy install step stubbed and
return how many times the functional quantize/server smoke tests ran."""
"""Run validate_prebuilt_choice with heavy steps stubbed; return the quantize/server smoke-test call counts."""
calls = {"quantize": 0, "server": 0}
server_path = tmp_path / "install" / "build" / "bin" / "llama-server"
quantize_path = tmp_path / "install" / "build" / "bin" / "llama-quantize"
@ -2847,23 +2812,19 @@ def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256):
def test_validate_prebuilt_choice_approved_validation_skipped_when_flag_off(tmp_path, monkeypatch):
# An approved (sha256-verified) bundle skips the staged smoke test while the
# flag is off: the manifest hash is its integrity gate.
# An approved (sha256-verified) bundle skips the smoke test while the flag is off.
calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32)
assert calls == {"quantize": 0, "server": 0}
def test_validate_prebuilt_choice_hashless_build_always_validated(tmp_path, monkeypatch):
# A hashless external build has no approved sha256, so the
# functional smoke test is its only integrity gate and must run even while the
# flag is off -- otherwise a corrupted/replaced archive could be activated.
# A hashless build has no sha256 gate, so the smoke test must run even with the flag off.
calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = None)
assert calls == {"quantize": 1, "server": 1}
def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp_path, monkeypatch):
# Flipping _RUN_STAGED_PREBUILT_VALIDATION back on restores the full smoke test
# for approved bundles too, proving the check is kept intact, only gated off.
# _RUN_STAGED_PREBUILT_VALIDATION back on restores the smoke test for approved bundles too.
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", True)
calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32)
assert calls == {"quantize": 1, "server": 1}

21
tests/studio/install/test_launch_studio_launcher.py
View file

@ -1,9 +1,7 @@
"""Guard install.ps1's Studio launcher against re-introducing the AV-heuristic
shape: a WScript .vbs that spawns a hidden, ExecutionPolicy-Bypass PowerShell
(Kaspersky HEUR:Trojan.VBS.Agent.gen). The shortcut must stay windowless via
powershell.exe -WindowStyle Hidden over launch-studio.ps1 -- never a .vbs /
WScript.Shell.Run wrapper. Any pre-existing .vbs from an older install must be
deleted, not merely left behind."""
"""Guard install.ps1's Studio launcher against the AV-heuristic shape (Kaspersky
HEUR:Trojan.VBS.Agent.gen): a WScript .vbs spawning a hidden ExecutionPolicy-Bypass PowerShell.
The shortcut must stay windowless via powershell.exe -WindowStyle Hidden over launch-studio.ps1,
never a .vbs/WScript.Shell.Run wrapper, and any pre-existing .vbs must be deleted on upgrade."""
import re
from pathlib import Path
@ -24,8 +22,7 @@ def test_install_ps1_present():
def test_no_vbs_launcher_generated():
text = _text()
# No here-string that builds a .vbs body, and no .vbs file written. (A
# Remove-Item cleanup of the legacy .vbs is allowed and checked separately.)
# No here-string building a .vbs body, no .vbs written (legacy cleanup checked separately).
assert "$vbsContent" not in text, (
"install.ps1 must not generate a launch-studio.vbs: a WScript.Shell .vbs "
"spawning a hidden ExecutionPolicy-Bypass PowerShell is the exact shape "
@ -38,8 +35,8 @@ def test_no_vbs_launcher_generated():
def test_legacy_vbs_removed_on_upgrade():
# The whole point: an upgrade must DELETE a pre-existing launch-studio.vbs,
# not just stop generating it, or AV keeps flagging the stale file.
# An upgrade must DELETE a pre-existing launch-studio.vbs, not just stop generating it,
# or AV keeps flagging the stale file.
text = _text()
assert re.search(
r"Remove-Item\s+-LiteralPath\s+\$legacyLauncherVbs", text
@ -47,13 +44,13 @@ def test_legacy_vbs_removed_on_upgrade():
def test_shortcut_target_is_not_wscript():
# The .lnk must not be launched through wscript.exe (the VBS script host).
# The .lnk must not launch through wscript.exe (the VBS script host).
text = _text()
assert "wscript.exe" not in text.lower()
def test_launcher_is_windowless_powershell():
# The shortcut runs powershell.exe with a hidden window over launch-studio.ps1.
# The shortcut runs powershell.exe -WindowStyle Hidden over launch-studio.ps1.
text = _text()
assert re.search(
r"-WindowStyle\s+Hidden", text

56
tests/studio/install/test_llama_pr_force_and_source.py
View file

@ -1,14 +1,4 @@
"""
Tests for the current llama.cpp wrapper policy in setup.sh / setup.ps1.
Tests cover:
- Bash subprocess: PR_FORCE promotion, user-override, zero/empty/invalid ignored
- Bash subprocess: source remains pinned to ggml-org even if env source is set
- Static source checks: mainline repo/source are hardcoded for now
- PowerShell subprocess: PR_FORCE promotion and fixed-source parity
Run: pytest tests/studio/install/test_llama_pr_force_and_source.py -v
"""
"""Tests for the llama.cpp wrapper policy (PR_FORCE promotion, fixed ggml-org source) in setup.sh / setup.ps1."""
import os
import shlex
@ -18,9 +8,6 @@ from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
@ -31,17 +18,13 @@ PWSH_AVAILABLE = os.path.isfile(PWSH) and os.access(PWSH, os.X_OK)
requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not available")
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def run_bash(
script: str,
*,
timeout: int = 60,
env: dict | None = None,
) -> subprocess.CompletedProcess:
"""Run a bash script fragment. 60s default tolerates slow CI shell
startup; the scripts themselves run in well under a second."""
"""Run a bash fragment. 60s default tolerates slow CI shell startup."""
run_env = os.environ.copy()
if env:
run_env.update(env)
@ -60,8 +43,7 @@ def run_pwsh(
timeout: int = 60,
env: dict | None = None,
) -> subprocess.CompletedProcess:
"""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 a PowerShell fragment. 60s default tolerates slow CI pwsh startup (10s was flaky)."""
run_env = os.environ.copy()
run_env["NO_COLOR"] = "1"
if env:
@ -75,9 +57,7 @@ def run_pwsh(
)
# ---------------------------------------------------------------------------
# Shared bash stubs
# ---------------------------------------------------------------------------
# Shared bash stubs.
BASH_STUBS = textwrap.dedent("""\
step() { echo "step:$1:$2"; }
substep() { :; }
@ -116,9 +96,7 @@ def make_mock_git(tmp_path: Path, *, fail_on: str = "") -> tuple[Path, Path]:
return mock_bin, log_file
# =========================================================================
# Bash fragment that exercises PR_FORCE and fixed _LLAMA_SOURCE resolution
# =========================================================================
# Bash fragment exercising PR_FORCE and fixed _LLAMA_SOURCE resolution.
def _bash_resolution_fragment(
llama_pr: str = "",
llama_pr_force: str = "",
@ -159,9 +137,6 @@ def _bash_resolution_fragment(
""")
# =========================================================================
# TEST GROUP A: Bash PR_FORCE promotion (subprocess)
# =========================================================================
class TestBashPrForcePromotion:
"""PR_FORCE promotes to _LLAMA_PR when user hasn't set one."""
@ -233,9 +208,6 @@ class TestBashPrForcePromotion:
assert "LLAMA_PR=" in r.stdout
# =========================================================================
# TEST GROUP B: Bash fixed mainline source (subprocess)
# =========================================================================
class TestBashFixedMainlineSource:
"""Source remains pinned to ggml-org while the temporary policy is active."""
@ -266,9 +238,6 @@ class TestBashFixedMainlineSource:
assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout
# =========================================================================
# TEST GROUP C: Bash clone URL parameterization (subprocess with mock git)
# =========================================================================
class TestBashCloneUrlParameterized:
"""Verify git clone uses _LLAMA_SOURCE instead of hardcoded URL."""
@ -372,9 +341,6 @@ class TestBashCloneUrlParameterized:
assert "ggml-org/llama.cpp.git" in log
# =========================================================================
# TEST GROUP D: Static source patterns -- setup.sh
# =========================================================================
class TestSourcePatternsSh:
"""Verify setup.sh keeps the temporary mainline-only llama.cpp policy."""
@ -425,7 +391,6 @@ class TestSourcePatternsSh:
def test_clone_urls_parameterized_tag_path(self):
"""Non-PR clone path uses the resolved source URL, not a hardcoded URL."""
# Find the non-PR clone line (after _CLONE_ARGS)
idx = self.content.index("_CLONE_ARGS=(git clone --depth 1)")
block = self.content[idx : idx + 400]
assert '"${_RESOLVED_SOURCE_URL}.git"' in block
@ -439,9 +404,6 @@ class TestSourcePatternsSh:
pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
# =========================================================================
# TEST GROUP E: Static source patterns -- setup.ps1
# =========================================================================
class TestSourcePatternsPs1:
"""Verify setup.ps1 keeps the temporary mainline-only llama.cpp policy."""
@ -463,8 +425,7 @@ 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.
# Repo chosen by GPU detection (GPU -> fork, CPU -> ggml-org), no env override.
assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content
assert (
"$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) "
@ -508,9 +469,6 @@ class TestSourcePatternsPs1:
pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
# =========================================================================
# TEST GROUP F: PowerShell PR_FORCE promotion (subprocess)
# =========================================================================
@requires_pwsh
class TestPwshPrForcePromotion:
"""PR_FORCE promotion and fixed-source logic via pwsh subprocess."""
@ -560,7 +518,7 @@ class TestPwshPrForcePromotion:
default_source,
)
run_env = {}
# Ensure env vars are unset by default
# Unset env vars by default.
run_env["UNSLOTH_LLAMA_PR"] = ""
run_env["UNSLOTH_LLAMA_PR_FORCE"] = ""
if env:

24
tests/studio/install/test_macos_version_compat.py
View file

@ -1,14 +1,4 @@
"""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.
"""
"""Host-macOS-version-aware llama.cpp prebuilt selection; Mach-O samples synthesized in-process, all I/O monkeypatched."""
import importlib.util
import struct
@ -233,16 +223,13 @@ def _fake_macos_releases(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."""
"""A pre-26 macOS host pins the last upstream release that loads on it (b9415); macOS 26+ and unknown-version hosts use normal latest selection."""
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.
# 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)
@ -285,10 +272,7 @@ class TestMacosReleasePin:
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."""
"""The gate is host >= prebuilt minos with no hardcoded version; each host takes the newest release it can load across a multi-tier release set."""
# Newest first: future 27 builds, current 26 builds, an old 14 tier, a 13.
RELEASES = [

147
tests/studio/install/test_pr4562_bugfixes.py
View file

@ -1,17 +1,4 @@
"""
Comprehensive tests for PR #4562 bug fixes.
Tests cover:
- Bug 1: PS1 detached HEAD on re-run (fetch + checkout -B pattern)
- Bug 2: Source-build fallback ignores pinned tag (both .sh and .ps1)
- Bug 3: Unix fallback deletes install before checking prerequisites
- Bug 4: Linux LD_LIBRARY_PATH missing build/bin
- "latest" tag resolution fallback chain (helper only)
- Cross-platform binary_env (Linux, macOS, Windows)
- Edge cases: malformed JSON, empty responses, env overrides
Run: pytest tests/studio/install/test_pr4562_bugfixes.py -v
"""
"""Tests for PR #4562 bug fixes (1-4), latest-tag resolution, and binary_env."""
import importlib.util
import json
@ -24,9 +11,6 @@ from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# Load the module under test (same pattern as existing test files)
# ---------------------------------------------------------------------------
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)
@ -47,11 +31,8 @@ SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def make_host(*, system: str) -> HostInfo:
"""Create a HostInfo for the given OS."""
"""HostInfo for the given OS."""
return HostInfo(
system = system,
machine = "x86_64" if system != "Darwin" else "arm64",
@ -78,7 +59,7 @@ def run_bash(
timeout: int = 10,
env: dict | None = None,
) -> str:
"""Run a bash script fragment and return its stdout."""
"""Run a bash fragment, return stdout."""
run_env = os.environ.copy()
if env:
run_env.update(env)
@ -95,11 +76,8 @@ def run_bash(
return result.stdout.strip()
# =========================================================================
# TEST GROUP A: binary_env across all platforms (Bug 4 + cross-platform)
# =========================================================================
class TestBinaryEnvCrossPlatform:
"""Test that binary_env returns correct library paths for all OSes."""
"""binary_env returns correct library paths for all OSes (Bug 4)."""
def test_linux_includes_binary_parent_in_ld_library_path(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
@ -163,7 +141,7 @@ class TestBinaryEnvCrossPlatform:
binary_path = bin_dir / "llama-server"
binary_path.write_bytes(b"fake")
# Create real directories so dedupe_existing_dirs keeps them
# Real dirs so dedupe_existing_dirs keeps them.
custom_lib = tmp_path / "custom_lib"
other_lib = tmp_path / "other_lib"
custom_lib.mkdir()
@ -215,13 +193,10 @@ class TestBinaryEnvCrossPlatform:
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}"
# binary_path.parent (build/bin) should come before install_dir
# build/bin must come before install_dir.
assert dyld_parts.index(str(bin_dir)) < dyld_parts.index(str(install_dir))
# =========================================================================
# TEST GROUP B: resolve_requested_llama_tag (Python function)
# =========================================================================
class TestResolveRequestedLlamaTag:
def test_concrete_tag_passes_through(self):
assert resolve_requested_llama_tag("b8508") == "b8508"
@ -383,14 +358,11 @@ class TestFetchJsonRetries:
assert len(releases) == 200
# =========================================================================
# TEST GROUP C: setup.sh logic (bash subprocess tests)
# =========================================================================
class TestSetupShLogic:
"""Test setup.sh fragments via bash subprocess with controlled PATH."""
"""setup.sh fragments via bash subprocess with controlled PATH."""
def test_cmake_missing_preserves_install(self, tmp_path: Path):
"""Bug 3: When cmake is missing, rm -rf should NOT run."""
"""Bug 3: cmake missing -> rm -rf must NOT run."""
llama_dir = tmp_path / "llama.cpp"
llama_dir.mkdir()
marker = llama_dir / "marker.txt"
@ -398,11 +370,11 @@ class TestSetupShLogic:
mock_bin = tmp_path / "mock_bin"
mock_bin.mkdir()
# Create mock git but NOT cmake
# Mock git but NOT cmake.
(mock_bin / "git").write_text("#!/bin/bash\nexit 0\n")
(mock_bin / "git").chmod(0o755)
# Build PATH: mock_bin first, then system dirs WITHOUT cmake
# PATH: mock_bin first, then system dirs without cmake.
safe_dirs = [str(mock_bin)]
for d in os.environ.get("PATH", "").split(":"):
if d and not os.path.isfile(os.path.join(d, "cmake")):
@ -424,7 +396,7 @@ class TestSetupShLogic:
assert marker.exists(), "Install dir was deleted despite cmake missing!"
def test_git_missing_preserves_install(self, tmp_path: Path):
"""Bug 3: When git is missing, rm -rf should NOT run."""
"""Bug 3: git missing -> rm -rf must NOT run."""
llama_dir = tmp_path / "llama.cpp"
llama_dir.mkdir()
marker = llama_dir / "marker.txt"
@ -432,11 +404,11 @@ class TestSetupShLogic:
mock_bin = tmp_path / "mock_bin"
mock_bin.mkdir()
# Create mock cmake but NOT git
# Mock cmake but NOT git.
(mock_bin / "cmake").write_text("#!/bin/bash\nexit 0\n")
(mock_bin / "cmake").chmod(0o755)
# Build PATH: mock_bin first, then system dirs WITHOUT git
# PATH: mock_bin first, then system dirs without git.
safe_dirs = [str(mock_bin)]
for d in os.environ.get("PATH", "").split(":"):
if d and not os.path.isfile(os.path.join(d, "git")):
@ -458,7 +430,7 @@ class TestSetupShLogic:
assert marker.exists(), "Install dir was deleted despite git missing!"
def test_both_present_runs_rm_and_clone(self, tmp_path: Path):
"""Bug 3: When both present, rm -rf runs before clone."""
"""Bug 3: both present -> rm -rf runs before clone."""
llama_dir = tmp_path / "llama.cpp"
llama_dir.mkdir()
marker = llama_dir / "marker.txt"
@ -567,11 +539,8 @@ class TestSetupShLogic:
assert "BuildOk=true" in output
# =========================================================================
# TEST GROUP D: "latest" tag resolution (bash subprocess)
# =========================================================================
class TestLatestTagResolution:
"""Test the fallback chain: helper resolver -> raw."""
"""Fallback chain: helper resolver -> raw."""
RESOLVE_TEMPLATE = textwrap.dedent("""\
_REQUESTED_LLAMA_TAG="{requested_tag}"
@ -643,11 +612,8 @@ class TestLatestTagResolution:
assert output == "latest"
# =========================================================================
# TEST GROUP E: Source file verification
# =========================================================================
class TestSourceCodePatterns:
"""Verify the actual source files contain the expected fix patterns."""
"""Verify the source files contain the expected fix patterns."""
def test_setup_sh_no_rm_before_prereq_check(self):
"""rm -rf must appear AFTER cmake/git checks, not before."""
@ -656,7 +622,6 @@ class TestSourceCodePatterns:
idx_block = content.find("command -v cmake")
assert idx_block != -1
block = content[idx_block:]
# rm -rf should appear after the cmake/git checks
idx_cmake = block.find("command -v cmake")
idx_git = block.find("command -v git")
idx_rm = block.find("rm -rf")
@ -670,7 +635,7 @@ class TestSourceCodePatterns:
assert (
'_CLONE_ARGS+=(--branch "$_RESOLVED_SOURCE_REF")' in content
), "_CLONE_ARGS should be extended with --branch $_RESOLVED_SOURCE_REF"
# Verify the guard: --branch is only used when tag is not "latest"
# --branch only when tag is not "latest".
assert (
'_RESOLVED_SOURCE_REF" != "latest"' in content
), "Should guard against literal 'latest' tag"
@ -687,18 +652,15 @@ class TestSourceCodePatterns:
assert "_RESOLVED_SOURCE_REF" in content
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."""
"""Shell prebuilt path uses the helper install entrypoint, not the old releases-latest flow."""
content = SETUP_SH.read_text()
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."""
"""Linux routing gates NVIDIA on GPU usability, not nvidia-smi presence, so
CPU-only/hidden-GPU hosts get the ggml CPU prebuilt. Guards the old presence-only loop."""
content = SETUP_SH.read_text()
assert '[ "$_setup_nvidia_usable" = true ]' in content
assert "CUDA_VISIBLE_DEVICES" in content
@ -725,36 +687,30 @@ class TestSourceCodePatterns:
assert "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" in content
def test_setup_sh_macos_metal_configure_has_cpu_fallback(self):
"""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."""
"""GPU configure/build failure retries a CPU build. Stays label-agnostic
(PR #5826 generalised the Metal-only wording via $_FB_LABEL)."""
content = SETUP_SH.read_text()
assert "_TRY_METAL_CPU_FALLBACK=true" 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
# (1 init + 2 resets = at least 3 occurrences of =false)
# Reset to false in both fallback branches: 1 init + 2 resets = >=3.
assert content.count("_TRY_METAL_CPU_FALLBACK=false") >= 3, (
"_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.
# Fallback helper must exist and Metal must reach it via the shortcut.
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."""
"""PR #5826: a fresh CUDA toolkit's host-compiler whitelist lags distro gcc/clang
(nvcc "#error -- unsupported GNU version"). setup.sh exports
NVCC_PREPEND_FLAGS=-allow-unsupported-compiler via env, not CMAKE_ARGS (word-splitting safety)."""
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.
# Via NVCC_PREPEND_FLAGS (covers the configure-time probe too), not CMAKE_ARGS.
assert "export NVCC_PREPEND_FLAGS=" in content
cmake_args_lines = [line for line in content.splitlines() if "CMAKE_ARGS=" in line]
assert all(
@ -762,22 +718,17 @@ class TestSourceCodePatterns:
), "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."""
"""Windows parity for PR #5826: CUDA toolkit whitelist lags MSVC. setup.ps1 sets
NVCC_PREPEND_FLAGS=-allow-unsupported-compiler in the CUDA branch via env, out of $CmakeArgs."""
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`.
# Via process env, not $CmakeArgs, so it reaches both the configure 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.
# Must be scoped to the CUDA branch, not set 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'")
@ -826,27 +777,24 @@ class TestSourceCodePatterns:
"""PS1 clone should use --branch with the resolved tag."""
content = SETUP_PS1.read_text()
assert "--branch" in content and "$ResolvedSourceRef" in content
# The old commented-out line should be gone
# The old commented-out clone line should be gone.
assert "# git clone --depth 1 --branch" not in content
def test_setup_ps1_no_git_pull(self):
"""PS1 should use fetch, not pull (which fails in detached HEAD)."""
content = SETUP_PS1.read_text()
# In the source-build section, there should be no "git pull"
# (git pull is only valid on a branch)
# No "git pull" in the source-build section (only valid on a branch).
lines = content.splitlines()
for i, line in enumerate(lines):
stripped = line.strip()
if "git pull" in stripped and not stripped.startswith("#"):
# Check context -- should not be in the llama.cpp build section
# Allow git pull in other contexts
# Allowed elsewhere; fail only in the llama.cpp build section.
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}")
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."""
"""PS1 prebuilt path uses the helper install entrypoint, not the old releases-latest flow."""
content = SETUP_PS1.read_text()
assert "--resolve-install-tag" not in content
assert "$HelperReleaseRepo/releases/latest" not in content
@ -909,7 +857,6 @@ class TestSourceCodePatterns:
def test_binary_env_linux_has_binary_parent(self):
"""The Linux branch of binary_env should include binary_path.parent."""
content = MODULE_PATH.read_text()
# Find the binary_env function
in_func = False
in_linux = False
found = False
@ -926,12 +873,8 @@ class TestSourceCodePatterns:
assert found, "binary_path.parent not found in Linux branch of binary_env"
# =========================================================================
# TEST GROUP F: macOS Metal build logic (bash subprocess tests)
# =========================================================================
# Minimal bash fragment that mirrors setup.sh's GPU backend decision chain.
# Variables _IS_MACOS_ARM64, NVCC_PATH, GPU_BACKEND are injected by tests.
# Bash fragment mirroring setup.sh's GPU backend decision chain.
# _IS_MACOS_ARM64, NVCC_PATH, GPU_BACKEND are injected by tests.
_GPU_BACKEND_FRAGMENT = textwrap.dedent("""\
CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF"
_TRY_METAL_CPU_FALLBACK=false
@ -961,7 +904,7 @@ _GPU_BACKEND_FRAGMENT = textwrap.dedent("""\
class TestMacOSMetalBuildLogic:
"""Behavioral bash subprocess tests for the Metal GPU backend logic."""
"""Behavioral bash tests for the Metal GPU backend logic."""
def test_macos_arm64_cmake_args_contain_metal_flags(self):
"""macOS arm64 should enable Metal, not CUDA."""
@ -994,7 +937,7 @@ class TestMacOSMetalBuildLogic:
mock_bin = tmp_path / "mock_bin"
mock_bin.mkdir()
calls_file = tmp_path / "cmake_calls.log"
# cmake that logs args and fails on first call (Metal), succeeds on second (CPU fallback)
# cmake logs args; fails first call (Metal), succeeds second (CPU fallback).
cmake_script = mock_bin / "cmake"
cmake_script.write_text(
textwrap.dedent(f"""\
@ -1056,7 +999,7 @@ class TestMacOSMetalBuildLogic:
"TRY_METAL_CPU_FALLBACK=false" in output
), "Fallback flag should be reset to false after configure fallback"
# Verify cmake args: first call has Metal ON, second has Metal OFF
# First cmake 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]}"
@ -1076,7 +1019,7 @@ class TestMacOSMetalBuildLogic:
mock_bin = tmp_path / "mock_bin"
mock_bin.mkdir()
calls_file = tmp_path / "cmake_calls.log"
# cmake mock: configure always succeeds; first --build fails, rest succeed
# cmake mock: configure always succeeds; first --build fails, rest succeed.
cmake_script = mock_bin / "cmake"
cmake_script.write_text(
textwrap.dedent(f"""\
@ -1164,14 +1107,11 @@ class TestMacOSMetalBuildLogic:
"TRY_METAL_CPU_FALLBACK=false" in output
), "Fallback flag should be reset to false after build fallback"
# Verify: configure with Metal ON, build fails, re-configure with Metal OFF, rebuild
# configure (Metal ON), build (fails), re-configure (Metal OFF), rebuild.
calls = calls_file.read_text().splitlines()
assert len(calls) >= 4, f"Expected >= 4 cmake calls, got {len(calls)}: {calls}"
# First call: configure with Metal ON
assert "-DGGML_METAL=ON" in calls[0]
# Second call: build (fails)
assert "--build" in calls[1]
# 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]}"
@ -1181,5 +1121,4 @@ class TestMacOSMetalBuildLogic:
assert (
"-DLLAMA_BUILD_TESTS=OFF" in calls[2]
), f"CPU fallback should preserve baseline flags: {calls[2]}"
# Fourth call: rebuild (succeeds)
assert "--build" in calls[3]

51
tests/studio/install/test_pr5940_followups.py
View file

@ -1,14 +1,9 @@
# 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.
"""
"""AMD-Windows installer follow-ups (PR #5940): hf-hub validation-model fetch
+ urllib fallback, amd-smi RunAsInvoker injection, name->arch table parity.
Mock-only; no AMD hardware or network required."""
import importlib.util
import re
@ -21,8 +16,7 @@ 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.
# install_llama_prebuilt.py is self-contained, so it loads without the studio backend.
_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
_SPEC = importlib.util.spec_from_file_location(
"studio_install_llama_prebuilt_pr5940", _PREBUILT_PATH
@ -151,8 +145,7 @@ def test_radeon_8060s_resolves_to_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."""
"""Parse a bash name->arch case table into ordered [(substr_tokens, arch), ...]."""
rows = []
for line in text.splitlines():
m = re.search(var + r'="(gfx[0-9a-z]+)"', line)
@ -172,9 +165,8 @@ def _sh_resolve(rows, name):
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."""
"""install.sh name->arch table must match the PowerShell source on the Strix
Halo/Point split and 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"
@ -196,9 +188,8 @@ def test_install_sh_name_arch_agrees_with_ps_for_strix_and_non_amd():
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)."""
"""studio/setup.sh's name->arch table must stay row-for-row identical to
install.sh's (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"),
@ -209,8 +200,7 @@ def test_setup_sh_name_arch_table_in_sync_with_install_sh():
"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.
# Guards historical drift: Strix Point -> gfx1150, RX 7700S -> gfx1102 (before gfx1100).
for name, expect in {
"AMD Radeon 890M Graphics": "gfx1150",
"AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150",
@ -222,9 +212,8 @@ def test_setup_sh_name_arch_table_in_sync_with_install_sh():
# ── 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.
# On Windows w/o a HIP SDK, amd-smi pops a UAC/DiskPart prompt RunAsInvoker
# can't suppress, so _amd_smi_allowed() skips it unless HIP-SDK or opt-in.
def _amd_smi_allowed_under(system, hipinfo_present, env):
@ -242,7 +231,7 @@ def _amd_smi_allowed_under(system, hipinfo_present, env):
def test_amd_smi_allowed_on_linux_regardless():
# Linux amd-smi does not elevate -> always allowed (no regression on Linux).
# Linux amd-smi does not elevate -> always allowed.
assert _amd_smi_allowed_under("Linux", hipinfo_present = False, env = {}) is True
@ -252,8 +241,7 @@ def test_amd_smi_skipped_on_windows_without_hip_sdk():
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).
# hipinfo present => amd-smi runs un-elevated, so it is allowed.
assert _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {}) is True
@ -274,8 +262,7 @@ def test_amd_smi_opt_out_overrides_hip_sdk():
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().
# Both PowerShell installers must gate amd-smi like _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"
@ -283,17 +270,15 @@ def test_ps_installers_gate_amd_smi_on_windows():
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.
# Regression for the DiskPart UAC prompt: every function naming `amd-smi`
# AND spawning a subprocess must gate it behind _amd_smi_allowed().
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.
# Exact "amd-smi"/"amd-smi.exe" constant, not a substring in a log.
return any(
isinstance(n, ast.Constant)
and isinstance(n.value, str)

49
tests/studio/install/test_probe_timeouts.py
View file

@ -1,11 +1,8 @@
"""Tests that NVIDIA probes in the installers are bounded by a timeout.
"""NVIDIA installer probes must be timeout-bounded (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 exit code 0).
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.
Source-level asserts check the guards in install.sh / install.ps1 / setup.ps1; one behavioral
shell test confirms the bash helper returns within the timeout when nvidia-smi hangs.
"""
import os
@ -60,8 +57,7 @@ class TestInstallShBoundedProbe:
"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.
# Falls back to unbounded when `timeout` is absent (e.g. macOS), keeping semantics there.
assert (
"else" in body and '"$@"' in body
), "_run_bounded must run the command unbounded when `timeout` is absent"
@ -69,11 +65,11 @@ class TestInstallShBoundedProbe:
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.
# The -L probe must go through the bounded runner.
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.
# The /proc fallback from PR 6174 must remain.
assert "/proc/driver/nvidia" in body
def test_cuda_version_parse_is_bounded(self):
@ -82,25 +78,20 @@ class TestInstallShBoundedProbe:
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.
# Locale 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`.
"""
"""Every nvidia-smi execution goes through _run_bounded (resolution checks are allowed)."""
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.
# The only $_nvsmi execution in _has_usable_nvidia_gpu 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.
# The $_smi execution in get_torch_index_url 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"
@ -119,7 +110,7 @@ class TestPowerShellBoundedProbe:
assert (
"WaitForExit($TimeoutSec * 1000)" in src
), f"{path.name} bounded probe must use WaitForExit with a timeout"
# Kill + sentinel on timeout, mirroring Invoke-AmdSmiNoElevate.
# Kill + sentinel on timeout (mirrors Invoke-AmdSmiNoElevate).
assert (
"$proc.Kill()" in src and "124" in src
), f"{path.name} must kill nvidia-smi and signal a timeout exit code"
@ -138,7 +129,7 @@ class TestPowerShellBoundedProbe:
@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.
# The exit-code-only probe 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"
@ -159,9 +150,7 @@ def _have_timeout() -> bool:
@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.
"""
"""Point _has_usable_nvidia_gpu at a fake nvidia-smi that sleeps 30s; the probe must return early."""
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")
@ -175,13 +164,11 @@ def test_has_usable_nvidia_gpu_returns_under_timeout():
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.
# PATH with the fake nvidia-smi first plus the real timeout/awk/ls it needs.
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.
# Force /proc fallback off so the result depends only on the probe (real NVIDIA host won't mask it).
script = (
f"{helper}\n{fn}\n"
"if _has_usable_nvidia_gpu; then echo DETECTED; else echo NONE; fi\n"
@ -194,9 +181,7 @@ def test_has_usable_nvidia_gpu_returns_under_timeout():
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.
# The probe must have returned (not hung): NONE without /proc, DETECTED via /proc fallback.
assert proc.stdout.strip() in {"NONE", "DETECTED"}
finally:
shutil.rmtree(workdir, ignore_errors = True)

445
tests/studio/install/test_rocm_support.py
View file

File diff suppressed because it is too large Load diff

269
tests/studio/install/test_selection_logic.py
View file

@ -1,14 +1,4 @@
"""Tests for binary selection logic in install_llama_prebuilt.py.
Covers: normalize_compute_cap, normalize_compute_caps, parse_cuda_visible_devices,
supports_explicit_visible_device_matching, select_visible_gpu_rows,
compatible_linux_runtime_lines, pick_windows_cuda_runtime,
compatible_windows_runtime_lines, runtime_line_from_cuda_version,
apply_approved_hashes, linux_cuda_choice_from_release, windows_cuda_attempts,
resolve_upstream_asset_choice.
No GPU, no network, no torch required -- all I/O is monkeypatched.
"""
"""Binary selection logic in install_llama_prebuilt.py; all I/O monkeypatched."""
import importlib.util
import os
@ -113,9 +103,7 @@ def load_studio_run_module(monkeypatch):
return module
# ---------------------------------------------------------------------------
# Helper factories
# ---------------------------------------------------------------------------
def make_host(**overrides):
@ -280,7 +268,7 @@ def mock_windows_runtime(monkeypatch, lines):
class TestStudioLocalhostIpv6Warning:
def _prepare_loopback(self, run_module, monkeypatch):
# Studio is confirmed answering on the IPv4 loopback.
# Studio confirmed answering on the IPv4 loopback.
monkeypatch.setattr(
run_module,
"_working_local_url",
@ -298,7 +286,7 @@ class TestStudioLocalhostIpv6Warning:
def _ipv6(port = 8888):
return (socket.AF_INET6, socket.SOCK_STREAM, 6, "", ("::1", port, 0, 0))
# -- _localhost_ipv6_mismatch_url ------------------------------------
# -- _localhost_ipv6_mismatch_url --
def test_ipv4_localhost_does_not_warn(self, monkeypatch):
run_module = load_studio_run_module(monkeypatch)
@ -308,8 +296,7 @@ class TestStudioLocalhostIpv6Warning:
assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888) is None
def test_dual_stack_localhost_does_not_warn(self, monkeypatch):
# localhost -> both ::1 and 127.0.0.1: browsers fall back to IPv4 when
# ::1 refuses, so the URL is reachable and no warning is needed.
# Dual-stack localhost: browsers fall back to IPv4 when ::1 refuses, so no warning.
run_module = load_studio_run_module(monkeypatch)
self._prepare_loopback(run_module, monkeypatch)
self._set_getaddrinfo(monkeypatch, [self._ipv6(), self._ipv4()])
@ -332,13 +319,11 @@ class TestStudioLocalhostIpv6Warning:
assert "http://localhost:8888" in captured.out
def test_ipv6_listener_does_not_suppress_warning(self, monkeypatch):
# Regression for the Codex review: a process answering on ::1 is NOT
# Studio (Studio binds 127.0.0.1 only), so the warning must still fire
# -- that is exactly when http://localhost would open the wrong service.
# A process on ::1 is NOT Studio (binds 127.0.0.1 only), so the warning must
# still fire -- that is exactly when http://localhost opens the wrong service.
run_module = load_studio_run_module(monkeypatch)
self._prepare_loopback(run_module, monkeypatch)
self._set_getaddrinfo(monkeypatch, [self._ipv6()])
# Even with something listening on ::1, the result is unchanged.
monkeypatch.setattr(
run_module,
"_local_port_open",
@ -363,7 +348,7 @@ class TestStudioLocalhostIpv6Warning:
assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", port) is None
def test_ipv4_not_answering_suppresses_warning(self, monkeypatch):
# Studio not confirmed on 127.0.0.1 -> do not warn.
# Studio not confirmed on 127.0.0.1 -> no warning.
run_module = load_studio_run_module(monkeypatch)
monkeypatch.setattr(run_module, "_working_local_url", lambda port: None)
self._set_getaddrinfo(monkeypatch, [self._ipv6()])
@ -388,7 +373,7 @@ class TestStudioLocalhostIpv6Warning:
assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888) is None
# -- _emit_startup_output (banner / warning wiring) ------------------
# -- _emit_startup_output (banner / warning wiring) --
def _wire_recorders(self, run_module, monkeypatch):
calls = {"banner": [], "warning": [], "stop_hint": 0, "reachability": []}
@ -614,7 +599,6 @@ class TestCompatibleLinuxRuntimeLines:
assert compatible_linux_runtime_lines(host) == ["cuda13", "cuda12"]
def test_future_major_derives_lines(self):
# A future major (14.x) offers cuda14 first, then older majors.
host = make_host(driver_cuda_version = (14, 0))
assert compatible_linux_runtime_lines(host) == ["cuda14", "cuda13", "cuda12"]
@ -657,9 +641,8 @@ class TestCompatibleWindowsRuntimeLines:
@pytest.mark.parametrize("minor", [0, 1, 2, 3])
def test_cuda12_runs_on_any_12_x_driver(self, minor):
# cuda12 app bundles are toolkit-12.8 builds with bundled runtime; CUDA
# minor-version compatibility runs them on any 12.x driver, same as Linux.
# Previously Windows wrongly gated cuda12 below a 12.4 driver.
# Regression: Windows previously gated cuda12 below a 12.4 driver, but minor-version
# compat runs toolkit-12.8 bundles on any 12.x driver, same as Linux.
host = make_host(driver_cuda_version = (12, minor))
assert compatible_windows_runtime_lines(host) == ["cuda12"]
@ -1151,9 +1134,8 @@ class TestValidatedChecksumsForBundle:
validated_checksums_for_bundle("unslothai/llama.cpp", bundle)
def test_rejects_exact_source_without_repo(self, monkeypatch):
# An exact source archive with no source repo to clone from would let
# preferred_source_archive silently fall back to upstream source at the
# tag, so validation must fail closed (clean source build instead).
# An exact source archive with no repo to clone from would silently fall back
# to upstream source at the tag, so validation must fail closed.
bundle = make_release([], release_tag = "r1", upstream_tag = "b8508")
checksums = make_checksums_with_source(
[], release_tag = "r1", upstream_tag = "b8508", source_commit = "a" * 40
@ -1168,8 +1150,7 @@ class TestValidatedChecksumsForBundle:
validated_checksums_for_bundle("unslothai/llama.cpp", bundle)
def test_accepts_exact_source_when_only_bundle_has_repo(self, monkeypatch):
# The source repo can live only in the manifest bundle, not the checksum
# payload. source_build_plan_for_release coalesces checksums-or-bundle, so
# The source repo can live only in the manifest bundle, not the checksums;
# validation must accept the bundle's repo rather than failing closed.
bundle = make_release(
[],
@ -1202,8 +1183,6 @@ class TestValidatedChecksumsForBundle:
class TestLinuxCudaChoiceFromRelease:
# --- Runtime line resolution ---
def test_no_runtime_lines_detected(self, monkeypatch):
mock_linux_runtime(monkeypatch, [])
host = make_host(driver_cuda_version = (12, 8))
@ -1249,7 +1228,7 @@ class TestLinuxCudaChoiceFromRelease:
assert any("unavailable_on_host" in entry for entry in log_entries)
def test_blackwell_prefers_cuda13_over_torch_cuda12(self, monkeypatch):
# Blackwell host with both lines sm_120-capable: cuda13 wins over torch's cuda12.
# Both lines sm_120-capable: cuda13 wins over torch's cuda12.
mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
art12 = make_artifact(
@ -1273,7 +1252,7 @@ class TestLinuxCudaChoiceFromRelease:
assert any("blackwell_runtime_override" in entry for entry in result.selection_log)
def test_blackwell_skips_incapable_cuda13_line(self, monkeypatch):
# cuda13 line cannot cover sm_120 (only an -older bundle): stay on native cuda12.
# cuda13 line can't cover sm_120 (only an -older bundle): stay on native cuda12.
mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
art13_older = make_artifact(
@ -1297,7 +1276,7 @@ class TestLinuxCudaChoiceFromRelease:
assert result.primary.name == "bundle-cuda12-newer.tar.gz"
def test_blackwell_cuda13_unavailable_uses_cuda12(self, monkeypatch):
# cuda13 runtime libs absent: override never forces an undetected line.
# cuda13 runtime libs absent: override must not force an undetected line.
mock_linux_runtime(monkeypatch, ["cuda12"])
host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
art12 = make_artifact(
@ -1313,7 +1292,7 @@ class TestLinuxCudaChoiceFromRelease:
assert result.primary.runtime_line == "cuda12"
def test_non_blackwell_keeps_torch_preference(self, monkeypatch):
# Non-Blackwell host: torch preference is untouched, no override.
# Non-Blackwell: torch preference untouched, no override.
mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(driver_cuda_version = (13, 0), compute_caps = ["86"])
art12 = make_artifact(
@ -1337,7 +1316,7 @@ class TestLinuxCudaChoiceFromRelease:
assert not any("blackwell_runtime_override" in entry for entry in result.selection_log)
def test_blackwell_ignores_malformed_runtime_line(self, monkeypatch):
# A malformed/future-format runtime_line must be skipped, never crash the major sort.
# A malformed runtime_line must be skipped, not crash the major sort.
mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
bad = make_artifact(
@ -1360,7 +1339,7 @@ class TestLinuxCudaChoiceFromRelease:
assert result.primary.runtime_line == "cuda13"
def test_blackwell_prefers_cuda14_over_lower_majors(self, monkeypatch):
# Forward-compat: the highest sm_120-capable CUDA major wins.
# The highest sm_120-capable CUDA major wins.
mock_linux_runtime(monkeypatch, ["cuda14", "cuda13", "cuda12"])
host = make_host(driver_cuda_version = (14, 0), compute_caps = ["120"])
arts = [
@ -1379,8 +1358,7 @@ class TestLinuxCudaChoiceFromRelease:
assert result.primary.runtime_line == "cuda14"
def test_arm64_host_selects_linux_arm64_cuda_kind(self, monkeypatch):
# An arm64 CUDA host (DGX Spark / Grace Hopper) selects the
# linux-arm64-cuda bundle and ignores the x64 linux-cuda one.
# arm64 CUDA host selects the linux-arm64-cuda bundle, not the x64 one.
mock_linux_runtime(monkeypatch, ["cuda13"])
host = make_host(
machine = "aarch64",
@ -1740,10 +1718,8 @@ class TestResolveInstallAttempts:
assert approved.release_tag == "llama-prebuilt-latest"
def test_linux_cpu_fork_without_bundle_raises_no_upstream_fallback(self, monkeypatch):
# A CPU-only Linux host on the fork no longer falls back to the ggml-org
# CPU asset: production routes CPU-only Linux to ggml-org, never the fork.
# With no fork CPU bundle in the manifest the resolver raises rather than
# quietly reaching for an upstream asset.
# CPU-only Linux on the fork must not fall back to the ggml-org CPU asset; with
# no fork CPU bundle the resolver raises rather than reaching upstream.
host = make_host(
has_usable_nvidia = False,
has_physical_nvidia = False,
@ -1999,9 +1975,8 @@ class TestResolveInstallAttempts:
class TestResolveInstallReleasePlans:
def _cuda_bundle(self, asset_name, release_tag, upstream_tag):
# A fork CUDA bundle that covers the default NVIDIA host (sm 86,
# cuda12 runtime), so each release yields a plan via
# linux_cuda_choice_from_release.
# Fork CUDA bundle covering the default NVIDIA host (sm 86, cuda12), so each
# release yields a plan via linux_cuda_choice_from_release.
art = make_artifact(
asset_name,
install_kind = "linux-cuda",
@ -2051,7 +2026,7 @@ class TestResolveInstallReleasePlans:
mock_linux_runtime(monkeypatch, ["cuda12"])
host = make_host(system = "Linux", machine = "x86_64", compute_caps = ["86"])
releases = [
# r2 ships no fork bundle, so it yields no plan and is skipped.
# r2 ships no fork bundle: yields no plan, is skipped.
INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
bundle = make_release([], release_tag = "r2", upstream_tag = "b9002"),
checksums = make_checksums_with_source(
@ -2141,9 +2116,8 @@ class TestWindowsCudaAttempts:
assert result[1].runtime_line == "cuda12"
def test_driver_below_published_minor_is_gated_to_cuda12(self, monkeypatch):
# A 13.0 driver cannot run a 13.1 build (forward minor), so it is gated
# out of cuda13 and falls back to the cuda12 build it can run, even when
# only the cuda13 runtime libs are detected.
# A 13.0 driver can't run a 13.1 build (forward minor), so it is gated off
# cuda13 to the cuda12 build, even when only cuda13 runtime libs are detected.
mock_windows_runtime(monkeypatch, ["cuda13"])
host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 0))
assets = self._upstream("13.1", "12.4")
@ -2152,7 +2126,7 @@ class TestWindowsCudaAttempts:
assert result[0].name == f"llama-{self.TAG}-bin-win-cuda-12.4-x64.zip"
def test_driver_at_published_minor_selects_cuda13(self, monkeypatch):
# A 13.1 driver matches the published 13.1 build exactly.
# A 13.1 driver matches the published 13.1 build.
mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
assets = self._upstream("13.1", "12.4")
@ -2214,8 +2188,7 @@ class TestWindowsCudaAttempts:
assert result[1].name == "cudart-llama-bin-win-cuda-12.4-x64.zip"
def test_cudart_runtime_archive_is_paired(self, monkeypatch):
# #5106: cudart bundle must surface on runtime_url so
# install_from_archives downloads it.
# #5106: cudart bundle must surface on runtime_url so install_from_archives downloads it.
mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
assets = {
@ -2260,8 +2233,7 @@ class TestWindowsCudaAttempts:
assert attempt.runtime_name is None
def test_tracks_upstream_cuda13_minor_bump(self, monkeypatch):
# ggml-org bumped the published Windows cuda13 build 13.1 -> 13.3; the
# selector must follow it instead of the old hardcoded 13.1 (#5861).
# #5861: selector must follow the published cuda13 bump 13.1 -> 13.3, not hardcode 13.1.
mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 3))
assets = self._upstream("13.3", "12.4")
@ -2284,9 +2256,8 @@ class TestWindowsCudaAttempts:
assert result[0].runtime_name == "cudart-llama-bin-win-cuda-13.3-x64.zip"
def test_driver_below_published_minor_does_not_get_newer_build(self, monkeypatch):
# ggml-org ships only cuda-13.3; a 13.1 driver cannot run it (forward
# minor), so it is gated to the cuda-12.4 build instead of an
# unguaranteed 13.3. A 13.3 driver still gets 13.3 (see other tests).
# Only cuda-13.3 published; a 13.1 driver can't run it (forward minor), so it is
# gated to cuda-12.4. A 13.3 driver still gets 13.3 (other tests).
mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
assets = self._upstream("13.3", "12.4")
@ -2327,9 +2298,7 @@ class TestWindowsCudaAttempts:
class TestPinnedBlackwellCudaFallback:
"""A Blackwell host on a 13.0/13.1/13.2 driver, gated off the in-release 13.3
build, gets the pinned immutable b9360 cuda-13.1 GPU build instead of the
CPU-only cuda-12.4 drop. The pin is dormant for everyone else."""
"""Blackwell on 13.0/13.1/13.2, gated off in-release 13.3, gets the pinned b9360 cuda-13.1 GPU build; dormant otherwise."""
TAG = "b8508"
@ -2358,7 +2327,7 @@ class TestPinnedBlackwellCudaFallback:
assert _pinned_windows_cuda_fallback(self._win_host((13, 2), ["120"]), []) is not None
def test_pin_offered_for_sm121_variant(self):
# sm_121 is Blackwell-family and also needs toolkit >= 12.8.
# sm_121 is Blackwell-family, also needs toolkit >= 12.8.
assert _pinned_windows_cuda_fallback(self._win_host((13, 1), ["121"]), []) is not None
def test_pin_uses_max_of_multi_gpu_caps(self):
@ -2366,17 +2335,16 @@ class TestPinnedBlackwellCudaFallback:
@pytest.mark.parametrize("sm", ["89", "90", "100"])
def test_pin_not_offered_to_non_blackwell(self, sm):
# Ada/Hopper run the cuda-12.4 build fine; the pin must not fire.
# Ada/Hopper run cuda-12.4 fine; the pin must not fire.
assert _pinned_windows_cuda_fallback(self._win_host((13, 1), [sm]), []) is None
def test_pin_offered_for_driver_13_0(self):
# b9360 is native sm_120a SASS (no JIT) and ships a cuda-13.1 cudart,
# both of which run on a 13.0 r580+ driver via CUDA minor-version
# compatibility. 13.0 is the mainstream Blackwell branch, so it must fire.
# b9360 native sm_120a SASS + cuda-13.1 cudart run on a 13.0 r580+ driver via
# minor-version compat; 13.0 is the mainstream Blackwell branch, so it must fire.
assert _pinned_windows_cuda_fallback(self._win_host((13, 0), ["120"]), []) is not None
def test_pin_not_offered_below_floor(self):
# 12.x predates Blackwell entirely; the pin stays dormant below 13.0.
# 12.x predates Blackwell; the pin stays dormant below 13.0.
assert _pinned_windows_cuda_fallback(self._win_host((12, 9), ["120"]), []) is None
def test_pin_not_offered_without_driver(self):
@ -2416,8 +2384,7 @@ class TestPinnedBlackwellCudaFallback:
)
def test_pin_dormant_when_runnable_cuda14_present(self, monkeypatch):
# A future Blackwell host with an in-release cuda14 build (no cuda13)
# must not get the older b9360 13.1 pin ahead of the runnable cuda14.
# An in-release cuda14 build must take precedence over the older b9360 13.1 pin.
mock_windows_runtime(monkeypatch, ["cuda14", "cuda12"])
host = self._win_host((14, 0), ["120"])
assets = {
@ -2465,8 +2432,7 @@ class TestPinnedBlackwellCudaFallback:
assert _windows_cuda_attempt_covers_blackwell(cpu) is False
def _app_attempt(self, profile, runtime_line, max_sm):
# The fork's app-named windows-cuda bundle: no toolkit minor in the name,
# SM coverage declared directly (as published_windows_cuda_attempts sets it).
# Fork app-named windows-cuda bundle: no toolkit minor in the name, SM coverage declared directly.
return AssetChoice(
repo = UPSTREAM_REPO,
tag = self.TAG,
@ -2495,10 +2461,9 @@ class TestPinnedBlackwellCudaFallback:
assert _windows_cuda_attempt_covers_blackwell(attempt) is covers
def test_pin_dormant_when_app_bundle_covers_blackwell(self):
# Regression: the fork's app-named cuda13 bundle covers Blackwell, so the
# b9360 pin must retire instead of being prepended ahead of the native
# in-release build (previously the coverage check only matched legacy
# -bin-win-cuda-X.Y-x64.zip names, so the pin never went dormant).
# Regression: the coverage check previously matched only legacy
# -bin-win-cuda-X.Y-x64.zip names, so the pin never went dormant for an
# app-named cuda13 bundle that covers Blackwell.
host = self._win_host((13, 1), ["120"])
existing = [self._app_attempt("newer", "cuda13", 120)]
assert _pinned_windows_cuda_fallback(host, existing) is None
@ -2510,9 +2475,7 @@ class TestPinnedBlackwellCudaFallback:
class TestDirectUpstreamBlackwellPin:
"""End to end: the pin lands ahead of cuda-12.4 on the simple/upstream path
a Blackwell Windows host actually uses, and stays absent once a runnable
in-release cuda13 build exists."""
"""The pin lands ahead of cuda-12.4 on the simple/upstream path; absent once a runnable in-release cuda13 build exists."""
TAG = "b9365"
@ -2549,9 +2512,8 @@ class TestDirectUpstreamBlackwellPin:
)
plan = direct_upstream_release_plan(self._release(), host, UPSTREAM_REPO, "latest")
order = [(a.tag, a.runtime_line or a.install_kind) for a in plan.attempts]
# cuda-12.4 (toolkit 12.4, no sm_120) is dropped entirely on Blackwell:
# behind the pin it would still be attempted if the pin download failed,
# and the functional validator accepts its slow non-native path.
# cuda-12.4 (no sm_120) is dropped entirely on Blackwell rather than left as a
# slow non-native fallback behind the pin.
assert order == [("b9360", "cuda13"), (self.TAG, "windows-cpu")]
assert plan.attempts[0].name == "llama-b9360-bin-win-cuda-13.1-x64.zip"
# Direct/upstream path stays unverified-by-manifest (no approved hashes).
@ -2577,9 +2539,7 @@ class TestDirectUpstreamBlackwellPin:
class TestBlackwellCuda124Exclusion:
"""A Blackwell host must never have a windows-cuda attempt that cannot
offload sm_120 anywhere in its chain: behind the pin it is one failed
download away from a validated-but-7-tok/s install."""
"""A Blackwell host must never have a windows-cuda attempt that can't offload sm_120 anywhere in its chain."""
def _bw_host(self):
return make_host(
@ -2612,8 +2572,8 @@ class TestBlackwellCuda124Exclusion:
assert [a.name for a in kept] == ["llama-b9365-bin-win-cuda-13.3-x64.zip"]
def test_keeps_manifest_cuda12_bundle_with_sm120(self):
# Published cuda12 app bundles are toolkit-12.8 builds that include
# sm_120; the manifest SM metadata must keep them on Blackwell.
# Published cuda12 app bundles are toolkit-12.8 builds with sm_120; manifest SM
# metadata must keep them on Blackwell.
bundle = AssetChoice(
repo = "unslothai/llama.cpp",
tag = "b9585",
@ -2680,10 +2640,7 @@ class TestBlackwellCuda124Exclusion:
class TestDirectLinuxNvidiaCpuGate:
"""When a release ships a linux-cpu bundle but no CUDA line this NVIDIA
host can use, the planner must raise (so the caller walks back to an older
release with a usable CUDA line) instead of silently planning a CPU
install on a GPU host. CPU-only hosts keep taking the CPU bundle."""
"""A linux-cpu-only release on an NVIDIA host must raise (caller walks back to a usable CUDA line), not silently CPU-install. CPU-only hosts keep the CPU bundle."""
def _bundle_cpu_only(self):
return make_release(
@ -2742,13 +2699,7 @@ class TestDirectLinuxNvidiaCpuGate:
class TestLinuxPublishedAttemptsNvidiaCpuGate:
"""Live fork-manifest path (_linux_published_attempts): an NVIDIA host whose
CUDA selection finds nothing must NOT be handed the manifest's CPU bundle --
the attempt list stays empty so the caller source-builds with CUDA instead of
silently installing a CPU-only binary on a GPU host. CPU-only hosts still get
the CPU bundle. Mirrors the ROCm policy and TestDirectLinuxNvidiaCpuGate (the
latter covers direct_linux_release_plan, which is off the live path, this the
live path)."""
"""Live fork-manifest path: an NVIDIA host whose CUDA selection finds nothing gets an empty attempt list (source-builds with CUDA), not the manifest CPU bundle. CPU-only hosts still get the CPU bundle."""
def _cpu_only_bundle(self):
return make_release(
@ -2800,9 +2751,7 @@ class TestLinuxPublishedAttemptsNvidiaCpuGate:
class TestPublishedWindowsCudaAttemptsDynamicMajor:
"""The published-path ordering seed is derived from the release's real
published minors, so a future CUDA major published here is selectable
instead of being hidden by a hardcoded cuda12/cuda13 seed."""
"""The ordering seed is derived from the release's published minors, so a future CUDA major is selectable, not hidden by a hardcoded cuda12/cuda13 seed."""
TAG = "b8508"
@ -2820,9 +2769,8 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
return make_release(artifacts, upstream_tag = self.TAG)
def test_future_cuda14_published_is_selected(self, monkeypatch):
# With the dynamic seed a 14.x driver reaches a published cuda14 build;
# the old hardcoded cuda12/cuda13 seed would never order it (the cuda14
# line would be skipped for want of a 14.x asset in the seed).
# The dynamic seed lets a 14.x driver reach a published cuda14 build; the old
# hardcoded cuda12/cuda13 seed would never order it.
mock_windows_runtime(monkeypatch, ["cuda14", "cuda13", "cuda12"])
release = self._release([("14.0", "cuda14"), ("13.3", "cuda13"), ("12.4", "cuda12")])
host = make_host(
@ -2836,7 +2784,7 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
assert result[0].name == f"llama-{self.TAG}-bin-win-cuda-14.0-x64.zip"
def test_cuda13_minor_selected_for_13_3_driver(self, monkeypatch):
# Existing behavior unchanged: a 13.3 driver gets the real 13.3 build.
# A 13.3 driver gets the real 13.3 build.
mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
release = self._release([("13.3", "cuda13"), ("12.4", "cuda12")])
host = make_host(
@ -2850,7 +2798,7 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
assert result[0].name == f"llama-{self.TAG}-bin-win-cuda-13.3-x64.zip"
def test_below_minor_driver_gated_to_cuda12(self, monkeypatch):
# A 13.1 driver is gated off a published 13.3 and falls to cuda12.
# A 13.1 driver is gated off published 13.3 and falls to cuda12.
mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
release = self._release([("13.3", "cuda13"), ("12.4", "cuda12")])
host = make_host(
@ -2869,8 +2817,7 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
class TestResolveReleaseAssetChoicePin:
"""The manifest install path reaches the same b9360 Blackwell pin as the
filename path, with its verified hash threaded."""
"""The manifest install path reaches the same b9360 Blackwell pin as the filename path, with its verified hash threaded."""
TAG = "b8508"
@ -2923,12 +2870,10 @@ class TestResolveReleaseAssetChoicePin:
result = resolve_release_asset_choice(host, self.TAG, release, checksums)
assert result[0].tag == "b9360"
assert result[0].name == "llama-b9360-bin-win-cuda-13.1-x64.zip"
# apply_approved_hashes threaded the pin's verified hash from the
# augmented checksums (the pin survives the approved-hash gate).
# The pin survives the approved-hash gate with its verified hash threaded.
assert result[0].expected_sha256 and len(result[0].expected_sha256) == 64
assert result[0].runtime_sha256 and len(result[0].runtime_sha256) == 64
# The sm_120-incapable upstream cuda-12.4 zip is excluded on Blackwell
# rather than left behind the pin as a slow-path fallback.
# The sm_120-incapable upstream cuda-12.4 zip is excluded on Blackwell.
assert not any(a.runtime_line == "cuda12" for a in result)
def test_pin_dormant_on_published_path_for_13_3(self, monkeypatch):
@ -2962,9 +2907,7 @@ class TestResolveReleaseAssetChoicePin:
class TestPublishedWindowsCudaAppBundleSmSelection:
"""app-named windows-cuda bundles carry no minor in the filename, so the
driver-minor gate is skipped. Selection must instead filter by SM coverage,
or every host gets the lowest-rank "older" bundle regardless of its GPU."""
"""app-named windows-cuda bundles carry no minor, so the driver-minor gate is skipped; selection must filter by SM coverage instead of handing every host the lowest-rank "older" bundle."""
TAG = "b9457"
@ -2997,8 +2940,8 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
)
result = published_windows_cuda_attempts(host, release, None)
assert result, "expected a windows-cuda attempt for an sm120 host"
# The lowest-rank "older" bundle (max_sm 89) must not be chosen, and the
# tightest covering bundle is cuda12-newer (range 86-120).
# Not the lowest-rank "older" bundle (max_sm 89); the tightest covering bundle
# is cuda12-newer (range 86-120).
assert result[0].name == f"app-{self.TAG}-windows-x64-cuda12-newer.zip"
def _line(self, line, klass, rank):
@ -3015,9 +2958,8 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
)
def test_cuda13_reachable_on_driver_13_0(self, monkeypatch):
# app-named cuda13 bundles must be reachable on a 13.0 driver. The old
# synthetic '13.1' minor gate dropped the whole cuda13 line (13.1 > 13.0),
# so a cu13 host fell to cuda12. cuda13 is gated at the major level now.
# Regression: a synthetic '13.1' minor gate dropped the whole cuda13 line on a
# 13.0 driver; cuda13 app bundles are gated at the major level now.
mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
release = make_release(
[self._line("cuda12", "newer", 20), self._line("cuda13", "newer", 50)],
@ -3035,11 +2977,9 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
assert result[0].name == f"app-{self.TAG}-windows-x64-cuda13-newer.zip"
def test_app_bundle_offered_when_no_runtime_dll_detected(self, monkeypatch):
# Windows torch bundles cudart in torch/lib, which runtime-DLL probing
# misses, so detected_windows_runtime_lines() returns nothing. The app
# bundle ships its own runtime, so selection must fall back to the
# driver-derived order instead of yielding no attempt (which would drop
# the host to the upstream build).
# torch bundles cudart in torch/lib, which DLL probing misses, so detection
# returns nothing; the app bundle ships its own runtime, so selection must
# fall back to the driver-derived order rather than drop to the upstream build.
mock_windows_runtime(monkeypatch, [])
release = make_release(
[self._line("cuda12", "newer", 20), self._line("cuda13", "newer", 50)],
@ -3057,9 +2997,7 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
class TestPublishedRocmGfxSelection:
"""Published ROCm bundles are matched by the host's detected gfx family, not
by rank -- rank ties would alphabetically hand every AMD GPU the gfx103X
bundle (e.g. a gfx1151 Strix Halo host)."""
"""Published ROCm bundles match by detected gfx family, not rank -- rank ties would alphabetically hand every AMD GPU the gfx103X bundle."""
GFX = ["gfx103X", "gfx110X", "gfx120X", "gfx1150", "gfx1151"]
MEMBERS = {
@ -3135,8 +3073,8 @@ class TestPublishedRocmGfxSelection:
)
def test_in_prefix_but_unbuilt_arch_returns_none(self):
# gfx1033 shares the gfx103 prefix but is not in any bundle's
# mapped_targets, so it must fall back to source, not be served gfx103X.
# gfx1033 shares the gfx103 prefix but isn't in any bundle's mapped_targets, so
# it must fall back to source, not be served gfx103X.
release = self._release("linux-rocm", "app-b9457-linux-x64-rocm")
for unbuilt in ("gfx1033", "gfx1035", "gfx1104", "gfx1202"):
assert (
@ -3147,10 +3085,8 @@ class TestPublishedRocmGfxSelection:
), unbuilt
def test_family_token_matches_family_bundle(self):
# The llama.cpp update path re-derives --rocm-gfx from the family-named
# marker asset, so it forwards a family token (gfx110X, lowercased to
# gfx110x by _normalize_forwarded_gfx), not a concrete arch. That must
# still select the family bundle instead of falling to a source build.
# The update path forwards a family token (gfx110X, lowercased to gfx110x), not
# a concrete arch; it must still select the family bundle, not source-build.
release = self._release("linux-rocm", "app-b9457-linux-x64-rocm")
for token in ("gfx110X", "gfx110x"):
choice = INSTALL_LLAMA_PREBUILT.published_rocm_choice_for_host(
@ -3160,8 +3096,8 @@ class TestPublishedRocmGfxSelection:
assert choice.name == "app-b9457-linux-x64-rocm-gfx110X.tar.gz", token
def test_windows_family_token_matches_family_bundle(self):
# The Windows update path forwards the same family token (gfx120X) for a
# windows-rocm bundle, so the family-label match must cover it too.
# The Windows update path forwards the same family token (gfx120X), so the
# family-label match must cover windows-rocm too.
release = self._release("windows-rocm", "app-b9457-windows-x64-rocm")
choice = INSTALL_LLAMA_PREBUILT.published_rocm_choice_for_host(
release, self._host("gfx120x"), "windows-rocm"
@ -3171,9 +3107,7 @@ class TestPublishedRocmGfxSelection:
class TestPublishedMacosForkSelection:
"""macOS now routes to the fork (setup.sh), which ships
llama-<tag>-bin-macos-<arch>.tar.gz with pinned deployment targets, selected
by install_kind."""
"""macOS routes to the fork's llama-<tag>-bin-macos-<arch>.tar.gz, selected by install_kind."""
def _release(self):
arts = [
@ -3224,7 +3158,7 @@ class TestPublishedMacosForkSelection:
class TestApplyApprovedHashesRuntimePair:
"""Runtime archive must inherit a manifest hash, or be dropped."""
"""Runtime archive inherits a manifest hash, or is dropped."""
TAG = "b8508"
@ -3456,8 +3390,7 @@ def _macos_host(machine = "arm64", version = (15, 5)):
class TestPinnedMacosReleaseTag:
"""pinned_macos_release_tag: pin b9415 only for ggml-org upstream macOS hosts
below macOS 26; latest (None) for 26+, unknown version, the fork, non-macOS."""
"""pinned_macos_release_tag: pin b9415 for ggml-org upstream macOS below 26; None (latest) for 26+, unknown version, the fork, non-macOS."""
def test_arm64_sequoia_pins_b9415(self):
host = _macos_host("arm64", (15, 5))
@ -3468,7 +3401,7 @@ class TestPinnedMacosReleaseTag:
assert pinned_macos_release_tag(host, UPSTREAM_REPO) == "b9415"
def test_x64_ventura_13_3_pins_b9415(self):
# b9415's Intel slice is minos 13.3, so 13.3 Intel hosts still load it.
# b9415's Intel slice is minos 13.3, so 13.3 Intel hosts load it.
host = _macos_host("x86_64", (13, 3))
assert pinned_macos_release_tag(host, UPSTREAM_REPO) == "b9415"
@ -3485,7 +3418,7 @@ class TestPinnedMacosReleaseTag:
assert pinned_macos_release_tag(host, UPSTREAM_REPO) is None
def test_fork_repo_is_dormant(self):
# The unslothai/llama.cpp fork publishes its own minos-13.3 prebuilts.
# The fork publishes its own minos-13.3 prebuilts.
host = _macos_host("arm64", (15, 5))
fork = INSTALL_LLAMA_PREBUILT.DEFAULT_PUBLISHED_REPO
assert pinned_macos_release_tag(host, fork) is None
@ -3496,9 +3429,7 @@ class TestPinnedMacosReleaseTag:
class TestResolveSimpleMacosPin:
"""End to end on the simple/upstream path macOS actually uses: a pre-26 host
deterministically resolves b9415 (no walk-back); a macOS 26 host takes the
latest release. Mirrors how setup.sh routes Darwin to ggml-org/llama.cpp."""
"""Simple/upstream path: a pre-26 host resolves b9415 (no walk-back); a macOS 26 host takes the latest release."""
TAGS = ["b9442", "b9430", "b9428", "b9415"] # newest-first feed
@ -3523,7 +3454,7 @@ class TestResolveSimpleMacosPin:
requested_tag = "",
):
calls.append((repo, published_release_tag, requested_tag))
# Emulate the real iterator: a specific tag yields only that release.
# Real iterator: a specific tag yields only that release.
if requested_tag and requested_tag != "latest":
yield _release(requested_tag)
return
@ -3547,7 +3478,7 @@ class TestResolveSimpleMacosPin:
assert plans[0].attempts[0].name == "llama-b9415-bin-macos-arm64.tar.gz"
# The pin overrode the requested tag before any release was fetched.
assert calls[0][2] == "b9415"
# Simple/upstream path stays unverified-by-manifest, exactly as before.
# Simple/upstream path stays unverified-by-manifest.
assert plans[0].approved_checksums.artifacts == {}
def test_tahoe_host_takes_latest_release(self, monkeypatch):
@ -3558,7 +3489,7 @@ class TestResolveSimpleMacosPin:
)
assert requested_tag == "latest"
assert plans[0].release_tag == "b9442"
# No pin: the iterator was asked for latest, not a specific tag.
# No pin: the iterator was asked for latest.
assert calls[0][2] == "latest"
@ -3568,14 +3499,10 @@ class TestResolveSimpleMacosPin:
class TestLinuxArm64ForkFallsBackToSource:
"""The fork now ships linux-arm64-cuda bundles (GH200/GB200/DGX Spark). An
arm64 Linux host on the fork no longer hard-fails on the simple path; it
delegates to the manifest-aware resolver, which selects the arm64 CUDA
bundle (or falls back to source only if none matches)."""
"""The fork ships linux-arm64-cuda bundles; an arm64 Linux fork host delegates to the manifest-aware resolver (arm64 CUDA bundle, or source if none matches) instead of hard-failing."""
def test_arm64_nvidia_fork_delegates_to_manifest_resolver(self, monkeypatch):
# arm64 fork hosts are no longer blocked up front; the simple resolver
# hands them to the manifest-aware resolver instead.
# arm64 fork hosts are no longer blocked up front; routed to the manifest resolver.
called = {}
def _full(llama_tag, host, repo, tag, **_kw):
@ -3589,9 +3516,8 @@ class TestLinuxArm64ForkFallsBackToSource:
assert plans == ["plan"]
def test_x86_64_fork_delegates_to_manifest_resolver(self, monkeypatch):
# The old linux-x64 arch guard is gone: an x64 fork host is routed to the
# manifest resolver exactly like every other fork host, not down a
# separate filename-parsing path.
# The old linux-x64 arch guard is gone: x64 fork hosts route to the manifest
# resolver like every other fork host, not a separate filename-parsing path.
called = {}
def _full(llama_tag, host, repo, tag, **_kw):
@ -3605,8 +3531,8 @@ class TestLinuxArm64ForkFallsBackToSource:
assert plans == ["plan"]
def test_arm64_cpu_on_ggml_org_is_not_blocked(self, monkeypatch):
# CPU-only arm64 routes to ggml-org (not the fork), so the guard must not
# fire; it reaches the iterator (empty here -> generic message).
# CPU-only arm64 routes to ggml-org, so the guard must not fire; it reaches the
# iterator (empty here -> generic message).
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"iter_release_payloads_by_time",
@ -3632,9 +3558,7 @@ class TestLinuxArm64ForkFallsBackToSource:
class TestCpuFallback:
"""--cpu-fallback drops GPU attributes so the CPU prebuilt for the host's
OS/arch is selected, letting an arm64 GPU host install ggml-org's arm64 CPU
build as a last resort when its source build produced no binary."""
"""--cpu-fallback drops GPU attributes so the host's OS/arch CPU prebuilt is selected, letting an arm64 GPU host install ggml-org's arm64 CPU build when its source build produced no binary."""
_SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
@ -3661,8 +3585,7 @@ class TestCpuFallback:
"resolve_simple_install_release_plans",
_capture,
)
# install_prebuilt exits EXIT_FALLBACK on PrebuiltFallback; we only care
# about the host it handed to the resolver before that.
# We only care about the host handed to the resolver before the fallback exit.
with pytest.raises(SystemExit):
INSTALL_LLAMA_PREBUILT.install_prebuilt(
install_dir = tmp_path / "llama",
@ -3693,12 +3616,12 @@ class TestCpuFallback:
},
],
}
# A GPU arm64 host cannot pick the CPU arm64 bundle on its own.
# A GPU arm64 host can't pick the CPU arm64 bundle on its own.
with pytest.raises(PrebuiltFallback):
direct_upstream_release_plan(
release, self._arm64_nvidia(), "ggml-org/llama.cpp", "latest"
)
# force_cpu drops the GPU attributes, so the CPU arm64 bundle is selected.
# force_cpu drops GPU attributes, so the CPU arm64 bundle is selected.
cpu_host = make_host(
system = "Linux",
machine = "aarch64",
@ -3715,8 +3638,8 @@ class TestCpuFallback:
def test_setup_sh_has_arm64_cpu_prebuilt_fallback(self):
source = self._SETUP_SH.read_text(encoding = "utf-8")
assert "--cpu-fallback" in source
# Fallback targets ggml-org (the only repo with an arm64 Linux build) and
# is gated on a degraded source build for arm64.
# Fallback targets ggml-org (only repo with an arm64 Linux build), gated on a
# degraded arm64 source build.
assert "ggml-org/llama.cpp" in source
assert "_LLAMA_CPP_DEGRADED" in source
@ -4089,7 +4012,7 @@ class TestCudaDriverToolkitMismatchMessage:
assert "FOUND" not in output
def _cuda_build_decision_output(self, *, nvcc_path, driver):
# Mirror setup.sh's source-build decision: keep the toolkit, switch, or degrade to CPU.
# Mirror setup.sh's source-build decision: keep toolkit, switch, or degrade to CPU.
script = textwrap.dedent(
f"""\
set -euo pipefail

25
tests/studio/load_freeze/llama_server_shim.py
View file

@ -1,10 +1,4 @@
"""Fake llama-server for simulation tests.
Knobs: tok_status / tok_body / tok_reset / tok_response_map and the
matching detok_* set let tests inject every failure mode for the
audio-type probe (timeouts, partial bodies, malformed JSON, codec
marker hits).
"""
"""Fake llama-server: tok_*/detok_* knobs inject failure modes for the audio-type probe."""
from __future__ import annotations
@ -54,8 +48,7 @@ class _Handler(BaseHTTPRequestHandler):
self.wfile.write(body)
def _send_reset(self, partial: bytes) -> None:
"""Write a partial body and slam the connection. Simulates a
crashed llama-server returning a RemoteProtocolError to httpx."""
"""Write a partial body and drop the connection (simulates a crashed server)."""
# Don't call send_response -- write a half-finished response.
try:
self.wfile.write(
@ -66,7 +59,7 @@ class _Handler(BaseHTTPRequestHandler):
except Exception:
pass
try:
# Use socket-level shutdown so the next read sees a reset.
# Socket-level shutdown so the next read sees a reset.
sock = self.connection
sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, b"\1\0\0\0\0\0\0\0")
sock.close()
@ -107,10 +100,8 @@ class _Handler(BaseHTTPRequestHandler):
self._send_raw(srv.config.tok_status, srv.config.tok_body)
return
content = str(body.get("content", ""))
# tok_response_map lets the test inject a specific token count
# for a specific input text. Used to synthesise "this text
# tokenises to exactly one token" for the csm / bicodec / dac
# detection branches.
# tok_response_map injects a token count per input text (e.g. the
# one-token cases for csm / bicodec / dac detection branches).
if content in srv.config.tok_response_map:
tokens = list(srv.config.tok_response_map[content])
else:
@ -199,8 +190,7 @@ class FakeLlamaServer:
detok_body: Optional[bytes] = None,
detok_map: Optional[dict] = None,
completion_delay: float = 0.0,
# Cosmetic: appears in the stdout template only; production
# code under test does not parse this.
# Cosmetic: only appears in the stdout template; not parsed.
model_path: str = "<test-fixture>/gemma-4.gguf",
) -> None:
self.host = host
@ -224,8 +214,7 @@ class FakeLlamaServer:
self._thread: Optional[threading.Thread] = None
def start(self) -> "FakeLlamaServer":
# port=0 lets ThreadingHTTPServer pick a free port atomically (no
# find-then-bind race); read back via server_address[1].
# port=0 lets the server pick a free port atomically (no find-then-bind race).
self._server = FakeLlamaServer._Server((self.host, self._requested_port), _Handler)
self._server.config = self.config
bound_port = self._server.server_address[1]

178
tests/studio/load_freeze/test_load_orchestrator.py
View file

@ -1,21 +1,4 @@
"""Comprehensive simulation suite for the #5642 fix.
Covers:
1. Behavioural canary (the bug class) — 2 tests
2. Behavioural fix-validation — 1 test
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
6. Drift / regression guards — 3 tests
7. Timing budgets — 1 test
Designed to run from inside ``temp/sim/`` after ``uv venv`` + minimal
``uv pip install`` of pytest/httpx/fastapi/uvicorn/anyio. Resolves
``studio/backend`` automatically by walking up from this file looking
for the workspace clone of ``unslothai/unsloth`` (search order: this
dir's parents → ``../../unsloth`` → ``UNSLOTH_REPO_ROOT`` env var).
"""
"""Simulation suite for the #5642 fix (sync detect_audio_type blocking the event loop)."""
from __future__ import annotations
@ -34,8 +17,6 @@ import pytest
# Repo discovery
def _find_repo_root() -> Path | None:
env = os.environ.get("UNSLOTH_REPO_ROOT")
if env:
@ -78,8 +59,6 @@ from llama_server_shim import FakeLlamaServer # noqa: E402
# Fixtures / helpers
def _make_backend(port: int, *, loaded: bool = True) -> LlamaCppBackend:
b = LlamaCppBackend.__new__(LlamaCppBackend)
b._port = port
@ -205,8 +184,6 @@ def _drive_concurrent_probe_and_health(
# (1) Behavioural canary
def test_buggy_route_blocks_event_loop():
"""Sync detect_audio_type call inside async route stalls /health."""
with FakeLlamaServer(tok_delay = 0.6, detok_delay = 0.6) as shim:
@ -234,20 +211,13 @@ def test_fixed_route_keeps_event_loop_responsive():
# (2) Functional equivalence -- sync == to_thread for each codec branch
@pytest.fixture
def shim_no_match():
"""A shim whose responses make detect_audio_type fall through every
codec branch and return None."""
"""Shim whose responses make detect_audio_type fall through every codec branch -> None."""
with FakeLlamaServer(
# detok responds with a 1-char unique string per tid -> doesn't
# start with "<custom_token_" so snac branch fails.
# detok strings don't start with "<custom_token_" so snac branch fails.
detok_map = {128258: "abc", 128259: "def"},
# tokenize responds with len-of-words tokens, which is always
# 1 for single-word inputs so we need >1 token for the codec
# branches NOT to match. Map every audio probe text to a 2-token
# response so all `len(_tok(...)) == 1` checks fail.
# 2-token responses make every `len(_tok(...)) == 1` codec check fail.
tok_response_map = {
"<|AUDIO|>": [0, 1],
"<|audio_eos|>": [0, 1],
@ -271,8 +241,7 @@ def test_functional_equivalence_no_match(shim_no_match):
def test_functional_equivalence_snac_match():
# snac match requires _detok(128258) AND _detok(128259) to start
# with "<custom_token_".
# snac: both _detok(128258) and _detok(128259) start with "<custom_token_".
with FakeLlamaServer(
detok_map = {128258: "<custom_token_99>", 128259: "<custom_token_98>"}
) as srv:
@ -284,8 +253,7 @@ def test_functional_equivalence_snac_match():
def test_functional_equivalence_csm_match():
# csm match: _tok("<|AUDIO|>") == 1 token AND _tok("<|audio_eos|>") == 1 token.
# Also snac match must fail first.
# csm: snac fails, then both <|AUDIO|> and <|audio_eos|> are 1 token.
with FakeLlamaServer(
detok_map = {128258: "non-snac", 128259: "non-snac"},
tok_response_map = {"<|AUDIO|>": [0], "<|audio_eos|>": [0]},
@ -298,7 +266,7 @@ def test_functional_equivalence_csm_match():
def test_functional_equivalence_whisper_match():
# whisper: snac fails, csm fails, then _tok("<|startoftranscript|>") == 1
# whisper: snac/csm fail, then <|startoftranscript|> is 1 token.
with FakeLlamaServer(
detok_map = {128258: "non-snac", 128259: "non-snac"},
tok_response_map = {
@ -315,10 +283,8 @@ def test_functional_equivalence_whisper_match():
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.
# audio_vlm: snac/csm/whisper fail, then the Gemma 4 <|audio|> arm (#6000)
# tokenises to 1 token while <audio_soft_token> stays 2 to isolate it.
with FakeLlamaServer(
detok_map = {128258: "non-snac", 128259: "non-snac"},
tok_response_map = {
@ -337,8 +303,7 @@ def test_functional_equivalence_audio_vlm_match():
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.
# bicodec: all prior branches fail, then bicodec_semantic_0/global_0 are 1 token.
with FakeLlamaServer(
detok_map = {128258: "non-snac", 128259: "non-snac"},
tok_response_map = {
@ -359,25 +324,19 @@ def test_functional_equivalence_bicodec_match():
# (3) Failure modes
def test_shim_returns_500_on_tokenize_returns_none():
"""detect_audio_type's `r.status_code == 200` check filters out
non-200 responses; the function gracefully falls through and
returns None. Both sync and threaded paths see identical behaviour."""
"""Non-200 responses fall through to None on both sync and threaded paths."""
with FakeLlamaServer(
detok_map = {128258: "non-snac", 128259: "non-snac"},
tok_status = 500,
) as srv:
backend = _make_backend(srv.port)
# Sync
assert backend.detect_audio_type() is None
# Threaded
assert asyncio.run(asyncio.to_thread(backend.detect_audio_type)) is None
def test_shim_returns_malformed_json_returns_none():
"""detect_audio_type's outer try/except catches r.json() failures."""
"""Outer try/except catches r.json() failures."""
with FakeLlamaServer(
detok_map = {128258: "non-snac", 128259: "non-snac"},
tok_body = b"{this is not json",
@ -388,8 +347,7 @@ def test_shim_returns_malformed_json_returns_none():
def test_shim_connection_reset_returns_none():
"""Connection drops mid-response (RemoteProtocolError / ReadError)
must be caught by detect_audio_type's outer try/except."""
"""Mid-response connection drop (RemoteProtocolError / ReadError) is caught."""
with FakeLlamaServer(
detok_map = {128258: "non-snac", 128259: "non-snac"},
tok_reset = True,
@ -400,16 +358,14 @@ def test_shim_connection_reset_returns_none():
def test_unreachable_port_returns_none():
"""Pointing the backend at a port nothing is listening on triggers
httpx.ConnectError. detect_audio_type's try/except swallows it."""
"""ConnectError on a dead port is swallowed -> None."""
backend = _make_backend(_free_port()) # nothing listening
assert backend.detect_audio_type() is None
assert asyncio.run(asyncio.to_thread(backend.detect_audio_type)) is None
def test_backend_not_loaded_short_circuits():
"""is_loaded=False -> detect_audio_type returns None without doing
any network I/O. Confirm sub-millisecond on both paths."""
"""is_loaded=False short-circuits to None with no network I/O (sub-ms both paths)."""
backend = _make_backend(_free_port(), loaded = False)
t0 = time.perf_counter()
sync = backend.detect_audio_type()
@ -423,11 +379,8 @@ def test_backend_not_loaded_short_circuits():
# (4) Stress / concurrency
def test_50_concurrent_probes_complete_without_deadlock():
"""Fire 50 /probe calls in parallel against a fast shim. Threadpool
must not deadlock; route handler must not lock or serialise."""
"""50 parallel /probe calls must not deadlock or serialise."""
with FakeLlamaServer(tok_delay = 0.05, detok_delay = 0.05) as shim:
backend = _make_backend(shim.port)
app = _build_app(backend, wrap_in_thread = True)
@ -444,18 +397,14 @@ def test_50_concurrent_probes_complete_without_deadlock():
results = [f.result(60.0) for f in futs]
elapsed = time.perf_counter() - t0
assert all(r.status_code == 200 for r in results)
# 50 probes at ~0.4s each, threadpool size 32 default -> ~1-2 batches.
# Bound generously to absorb CI jitter while catching pathological
# serialisation (would be ~20s).
# Generous bound absorbs CI jitter but still catches serialisation (~20s).
assert (
elapsed < 15.0
), f"50 concurrent probes took {elapsed:.1f}s; threadpool may be serialising"
def test_100_concurrent_healths_during_slow_probe_all_responsive():
"""Heavier version of the canary: 100 /health requests across 8
worker threads during a slow /probe. With the fix, max latency
stays bounded; without the fix, requests pile up."""
"""100 /health across 8 threads during a slow /probe: latency stays bounded with the fix."""
with FakeLlamaServer(tok_delay = 0.4, detok_delay = 0.4) as shim:
backend = _make_backend(shim.port)
app = _build_app(backend, wrap_in_thread = True)
@ -478,7 +427,7 @@ def test_100_concurrent_healths_during_slow_probe_all_responsive():
with ThreadPoolExecutor(max_workers = 9) as pool:
probe_f = pool.submit(probe)
time.sleep(0.05) # let probe enter detect_audio_type
time.sleep(0.05) # let probe enter detect_audio_type first
health_fs = [pool.submit(health_burst, 13) for _ in range(8)]
assert probe_f.result(60.0) == 200
latencies = [x for f in health_fs for x in f.result(60.0)]
@ -488,25 +437,16 @@ def test_100_concurrent_healths_during_slow_probe_all_responsive():
# (5) Drift / regression guards on the production source
def test_load_model_caches_audio_type_inside_serial_load_lock():
"""The audio-type detection (and codec init, where applicable) must
happen inside ``LlamaCppBackend.load_model`` so the full load
sequence is atomic under ``_serial_load_lock``. Running it from the
route opens a race where a concurrent /load can replace the backend
mid-probe (gemini-code-assist review on #5669)."""
"""Audio-type detection must run inside load_model under _serial_load_lock,
else a concurrent /load can replace the backend mid-probe (review on #5669)."""
f = _REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py"
text = f.read_text()
# The lock must be acquired.
assert (
"with self._serial_load_lock" in text
), "LlamaCppBackend.load_model must hold self._serial_load_lock"
# The cache writes must be present. The strict variant
# `_detect_audio_type_strict` was added in the chatgpt-codex
# P2 3284185168 follow-up to distinguish definitive non-audio
# from transient probe failure; either call shape satisfies
# the static guard.
# Either call shape satisfies the guard; _detect_audio_type_strict was a
# follow-up to distinguish definitive non-audio from transient probe failure.
assert (
"self._audio_type = self.detect_audio_type()" in text
or "detected = self.detect_audio_type()" in text
@ -519,11 +459,8 @@ def test_load_model_caches_audio_type_inside_serial_load_lock():
def test_routes_inference_reads_cached_audio_type_not_calls_detect():
"""Static guard: routes/inference.py must NOT call
``llama_backend.detect_audio_type`` or
``llama_backend.init_audio_codec`` directly any more -- both moved
inside ``LlamaCppBackend.load_model`` under the lock. The route
reads the cached ``_audio_type`` / ``_is_audio`` attributes."""
"""routes/inference.py must read cached _audio_type/_is_audio, not call
detect_audio_type / init_audio_codec directly (both moved into load_model)."""
f = _REPO_ROOT / "studio" / "backend" / "routes" / "inference.py"
text = f.read_text()
assert "llama_backend.detect_audio_type(" not in text, (
@ -534,38 +471,29 @@ def test_routes_inference_reads_cached_audio_type_not_calls_detect():
"routes/inference.py should not call init_audio_codec directly; "
"load_model already invoked it under the lock when audio_type was a TTS codec."
)
# Verify the route DOES read the cached values somewhere.
# Route must read the cached values.
assert "llama_backend._audio_type" in text
assert "llama_backend._is_audio" in text
def test_no_other_async_route_calls_detect_audio_type_unwrapped():
"""Walk every .py under studio/backend/routes/ and confirm no file
contains a ``LlamaCppBackend.detect_audio_type()`` call inside an
async function. Re-introducing the bug means putting back the sync
call AND opening the race condition the lock fix closes."""
"""No routes/*.py may call llama_backend.detect_audio_type() in an async fn;
that reintroduces the sync bug and the load race the lock fix closes."""
routes_dir = _REPO_ROOT / "studio" / "backend" / "routes"
offenders = []
# Match `<anything>.detect_audio_type(` so this catches both
# `llama_backend.detect_audio_type(` and `self.detect_audio_type(`.
# We exclude the `utils.models.model_config.detect_audio_type`
# free function which is a separate, harmless static helper.
# Matches both llama_backend. and self. prefixes; the model_config free
# function helper is excluded below.
pattern = re.compile(r"\b\w+\.detect_audio_type\s*\(")
for path in routes_dir.rglob("*.py"):
for i, line in enumerate(path.read_text().splitlines(), start = 1):
m = pattern.search(line)
if not m:
continue
# Skip the free function import-site uses (no llama_backend prefix
# and called outside async context). Easiest: only treat the
# LlamaCppBackend instance call as an offender.
# Only the LlamaCppBackend instance call is an offender.
if "llama_backend.detect_audio_type" not in line:
continue
if "asyncio.to_thread" in line:
# Wrapped sync call is acceptable (event-loop responsive)
# but not preferred -- detect_audio_type belongs inside
# load_model now. Surface but don't fail; comment in PR
# if seen.
# Wrapped sync call is acceptable (not preferred); surface in PR.
continue
offenders.append(f"{path.relative_to(_REPO_ROOT)}:{i}: {line.strip()}")
assert not offenders, (
@ -575,8 +503,6 @@ def test_no_other_async_route_calls_detect_audio_type_unwrapped():
# (6) Timing budgets
def test_load_response_under_2s_with_fast_shim():
"""Regression budget: fast shim must complete /probe in <2 s."""
with FakeLlamaServer(tok_delay = 0.0, detok_delay = 0.0) as shim:
@ -592,9 +518,7 @@ def test_load_response_under_2s_with_fast_shim():
def test_repeated_loads_bounded_total_time():
"""Five sequential /probe calls against a fast shim must complete
in well under 10 s total. Locks in that there's no per-call leak
(open connections, threads, etc.) that compounds across loads."""
"""Five sequential /probe calls finish under 10 s, guarding against per-call leaks."""
with FakeLlamaServer(tok_delay = 0.05, detok_delay = 0.05) as shim:
backend = _make_backend(shim.port)
app = _build_app(backend, wrap_in_thread = True)
@ -609,13 +533,8 @@ def test_repeated_loads_bounded_total_time():
# (7) Browser-compatibility surface
def test_response_is_valid_browser_parseable_json():
"""The fix changes the route's internal scheduling but must not
change the response shape any browser sees. Round-trip the response
through json.loads() (the canonical equivalent of
JSON.parse() in any browser) and assert the expected keys."""
"""The fix must not change the response shape a browser sees (valid JSON, expected keys)."""
import json as _json
with FakeLlamaServer(tok_delay = 0.0, detok_delay = 0.0) as shim:
@ -625,27 +544,16 @@ def test_response_is_valid_browser_parseable_json():
with _UvicornServerThread(app, port = port) as uv:
with httpx.Client(timeout = 5.0) as c:
r = c.get(f"http://127.0.0.1:{uv.port}/probe")
# 1. Status code is one a browser will surface as success.
assert r.status_code == 200
# 2. Content-Type is exactly application/json (browsers use this
# header to decide if they can JSON-parse the body).
assert r.headers["content-type"].startswith("application/json")
# 3. Body is valid JSON. Every modern browser (Firefox, Safari,
# Chrome, Edge) uses the same JSON.parse semantics; parse via
# Python's strict json module here as a stand-in.
parsed = _json.loads(r.text)
# 4. Expected key present.
assert "audio_type" in parsed
# 5. No NaN / Infinity / non-JSON-spec types that would break
# browser parsers.
# No NaN / Infinity that would break browser parsers.
assert _json.dumps(parsed)
def test_response_shape_matches_pre_fix_for_no_match():
"""The fix's only externally-observable effect must be timing.
Confirm sync and threaded paths return byte-identical response
bodies for the no-match scenario (the dominant code path in
practice for non-audio models)."""
"""Sync and threaded paths return identical bodies for the no-match scenario."""
import json as _json
with FakeLlamaServer(
detok_map = {128258: "abc", 128259: "def"},
@ -662,7 +570,7 @@ def test_response_shape_matches_pre_fix_for_no_match():
},
) as shim:
backend = _make_backend(shim.port)
# Two apps -- sync (pre-fix) and to_thread (post-fix).
# sync (pre-fix) then to_thread (post-fix).
for wrap in (False, True):
app = _build_app(backend, wrap_in_thread = wrap)
port = _free_port()
@ -675,14 +583,8 @@ def test_response_shape_matches_pre_fix_for_no_match():
# (8) Cancellation
def test_client_disconnect_during_probe_does_not_crash_server():
"""If the HTTP client disconnects mid-probe, uvicorn must continue
serving subsequent requests. The threadpool task keeps running
(asyncio.to_thread doesn't propagate cancellation), but that's
matched by the existing init_audio_codec wrap and is not a
regression. After the disconnect, /health must still respond."""
"""A client disconnect mid-probe must not crash the server; /health still responds."""
with FakeLlamaServer(tok_delay = 0.5, detok_delay = 0.5) as shim:
backend = _make_backend(shim.port)
app = _build_app(backend, wrap_in_thread = True)
@ -690,13 +592,11 @@ def test_client_disconnect_during_probe_does_not_crash_server():
with _UvicornServerThread(app, port = port) as uv:
base = f"http://127.0.0.1:{uv.port}"
# Connect and immediately drop. httpx with a very short
# timeout simulates a client that gave up.
# Short timeout simulates a client that gave up mid-probe.
with pytest.raises(httpx.TimeoutException):
with httpx.Client(timeout = 0.2) as c:
c.get(f"{base}/probe")
# The server must still serve /health afterwards.
with httpx.Client(timeout = 5.0) as c:
r = c.get(f"{base}/health")
assert r.status_code == 200

142
tests/studio/playwright_chat_ime_i18n.py
View file

@ -3,18 +3,10 @@
"""Studio chat composer IME + multilingual regression smoke.
Covers four surfaces:
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.
D. Mac input-method switch: compositionstart without compositionend leaves
composingRef stuck; keydown and blur recover immediately.
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.
Covers: stuck IME composition (#5318 / PR #5327), multilingual paste round-trip,
stuck compositionend (#5546), and Mac input-method switch recovery (keydown/blur).
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.
"""
import os
@ -42,7 +34,7 @@ ART = Path(ART_DIR)
ART.mkdir(parents = True, exist_ok = True)
STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
# Wall-clock cap. Realistic run is 30-60s; 5 min leaves cold-launch headroom.
# Wall-clock cap: 5 min leaves cold-launch headroom over the 30-60s run.
WALL_TIMEOUT_S = float(os.environ.get("STUDIO_IME_WALL_TIMEOUT_S", "300"))
@ -98,7 +90,7 @@ def fail(m):
def soft_fail(m):
"""Hard fail in STRICT mode, info-warn otherwise. Mirrors playwright_chat_ui.py."""
"""Hard fail in STRICT mode, else info-warn (mirrors playwright_chat_ui.py)."""
if STRICT:
fail(m)
info(f"WARN (strict-off): {m}")
@ -153,8 +145,7 @@ with sync_playwright() as p:
except Exception as _shoot_err:
info(f"WARN: screenshot {name} failed: {_shoot_err}")
# 1. Bootstrap auth via /change-password (retry-on-rerender absorbs React
# form-detach races, mirroring playwright_chat_ui.py).
# 1. Bootstrap auth via /change-password; retry absorbs React form-detach races.
step("change-password through UI (Setup your account)")
form_err: Exception | None = None
for _form_attempt in range(3):
@ -202,7 +193,7 @@ with sync_playwright() as p:
if form_err is not None:
raise form_err
# 2. Wait for composer mount (no GGUF: the bug 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)
@ -248,8 +239,7 @@ with sync_playwright() as p:
else:
info('composer dir="auto" present')
# Source-level guard for the unmounted edit/compare composers: grep their
# JSX for dir="auto".
# Source-level guard: grep the unmounted edit/compare composers' JSX for dir="auto".
_repo_root = Path(__file__).resolve().parents[2]
_thread_src = (
_repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx"
@ -273,9 +263,8 @@ with sync_playwright() as p:
return composer.evaluate("(el) => el.value")
def set_value_via_setter(s: str) -> str:
"""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)."""
"""Write via React's setter + paste event, await two rAFs so the controlled
value commits before readback (plain `.value=s` is overwritten on render)."""
return composer.evaluate(
"""async (el, v) => {
const setter = Object.getOwnPropertyDescriptor(
@ -320,7 +309,7 @@ with sync_playwright() as p:
"Send button is unavailable for the next IME regression case."
)
# 3. Baseline: ASCII keyboard typing works. Bail fast if not.
# 3. Baseline: ASCII keyboard typing works.
step("baseline ASCII keyboard typing")
clear()
composer.click()
@ -389,9 +378,8 @@ with sync_playwright() as p:
shoot("05-normal-composition")
clear()
# 6. Stuck IME repro (#5318): duplicate compositionstart with no
# compositionend wedged isComposing=true; PR #5327 clears it on
# non-composing input.
# 6. Stuck IME repro (#5318): duplicate compositionstart wedges
# isComposing=true; PR #5327 clears it on non-composing input.
step("BUG REPRO: stuck IME composition recovery (issue #5318)")
clear()
composer.click()
@ -426,8 +414,8 @@ with sync_playwright() as p:
"likely still stuck in isComposing=true (issue #5318 / before "
"PR #5327)."
)
# Cross-check React's view of isComposing via the Send button:
# ComposerAction stays disabled while isComposing is true (PR #5327).
# Cross-check isComposing via the Send button: it stays disabled while
# isComposing is true (PR #5327).
send_btn = page.locator('button[aria-label="Send message"]')
if send_btn.count() == 0:
soft_fail("Send button not found after stuck-composition recovery")
@ -445,7 +433,7 @@ with sync_playwright() as p:
# 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.
# go silent.
step("BUG REPRO: stuck compositionend recovery (issue #5546)")
clear()
composer.click()
@ -490,9 +478,8 @@ with sync_playwright() as p:
clear()
# 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.
# IME keydown (keyCode 229) must not slip preedit text through submit.
# The onKeyDown gate re-pins composingRef so handleSubmit refuses.
step("BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)")
clear()
composer.click()
@ -517,9 +504,8 @@ 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) then submit synchronously. The
# keydown gate re-pins composingRef before handleSubmit, preventing submit;
# the textarea must still hold the preedit text.
# Fire IME-confirm Enter (keyCode 229) then submit synchronously: the keydown
# gate re-pins composingRef before handleSubmit, so preedit text is retained.
submit_probe = composer.evaluate(
"""(el) => {
el.focus();
@ -544,10 +530,9 @@ with sync_playwright() as p:
info("keydown re-pin gate PASS")
clear()
# 6d. Keydown re-pin must also re-arm the watchdog. On the WSL+Chrome
# 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.
# 6d. Keydown re-pin must also re-arm the watchdog: on the WSL+Chrome
# stuck-compositionend path no follow-up event arrives, so after re-pin
# the watchdog must clear composingRef again or Send re-locks forever.
step("BUG REPRO: keydown re-pin re-arms watchdog (#5546 follow-up regression)")
clear()
composer.click()
@ -572,7 +557,7 @@ with sync_playwright() as p:
expect(send_btn_rearm).not_to_be_disabled(timeout = 8_000)
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
# IME-confirm keydown re-pins composingRef; without the re-arm fix the
# watchdog never runs again and Send stays blocked forever.
composer.evaluate(
"""(el) => {
@ -583,9 +568,8 @@ with sync_playwright() as p:
}));
}"""
)
# 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.
# Second watchdog cycle: requestSubmit() after the re-armed window must be
# allowed; the buggy build stays gated forever.
rearm_probe = page.evaluate(
"""async (selector) => {
const ta = document.querySelector(selector);
@ -620,41 +604,30 @@ with sync_playwright() as p:
clear()
# 6e. Mac input-method switch - onKeyDown immediate recovery.
# On macOS, pressing Ctrl+Space or clicking the menu-bar language icon
# fires compositionstart but never fires compositionend (the OS commits
# nothing because no candidate was selected). When the user types their
# first English key after switching back, onKeyDown receives a native
# event with isComposing=false and a regular keyCode. The else-if branch
# added for this bug clears composingRef immediately, before the 2500ms
# watchdog would fire, so Send is unblocked on that very keystroke.
#
# To isolate the onKeyDown else-if path (and not the onChange path which
# also clears composing on normal input), we dispatch a synthetic KeyboardEvent
# with isComposing=false but do NOT dispatch a follow-up input event.
# onChange never fires, so the only recovery path is onKeyDown.
# A Mac IME switch fires compositionstart but never compositionend; the
# first English keydown (isComposing=false) must clear composingRef via
# the onKeyDown else-if branch, before the 2500ms watchdog fires.
# To isolate that path (not onChange) we dispatch a synthetic keydown with
# NO follow-up input event, so onChange never fires.
step("BUG REPRO: Mac IME switch - onKeyDown immediate recovery")
clear()
composer.click()
# Seed sendable content so the Send button's state reflects composition
# state only, not empty-content gating. set_value_via_setter uses
# insertFromPaste which is not composing, so composingRef stays false here.
# Seed sendable content so Send's state reflects composition only, not
# empty-content gating (insertFromPaste leaves composingRef false).
set_value_via_setter("hello")
# Simulate switching TO Chinese input method: compositionstart fires but
# compositionend never arrives (user switched away without committing text).
# Switch TO Chinese: compositionstart fires but compositionend never arrives.
composer.evaluate(
"""(el) => {
el.focus();
el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
}"""
)
# Give React a tick to process the compositionstart and update isComposing.
# Let React process compositionstart and update isComposing.
page.wait_for_timeout(200)
send_btn_mac_kd = page.locator('button[aria-label="Send message"]')
# Dispatch ONLY a keydown (isComposing=false, keyCode=65) with no follow-up
# input event. This fires onKeyDown but NOT onChange, so the else-if branch
# is the only path that can clear composingRef. page.keyboard.type() would
# also fire an input event and trigger onChange, which already clears
# composing on ASCII input, which would make the test a false positive.
# Dispatch ONLY a keydown (no input event) so onChange never fires and the
# onKeyDown else-if branch is the only path that can clear composingRef.
# page.keyboard.type() would fire onChange too and mask a regression.
composer.evaluate(
"""(el) => {
el.focus();
@ -668,8 +641,8 @@ with sync_playwright() as p:
soft_fail("Send button not found for Mac IME switch (onKeyDown) repro")
else:
try:
# 1500ms is well below the 2500ms watchdog: only the onKeyDown
# else-if path can clear composingRef this quickly.
# 1500ms < 2500ms watchdog: only the onKeyDown else-if path can
# clear composingRef this quickly.
expect(send_btn_mac_kd).not_to_be_disabled(timeout = 1_500)
info(
"Send button enabled within 1500ms after Mac IME switch + "
@ -687,10 +660,9 @@ with sync_playwright() as p:
clear()
# 6f. Candidate-confirming Enter must not unblock submit.
# Some IMEs/browsers report the candidate-confirming Enter as
# isComposing=false with keyCode=13 while composingRef is still pinned.
# That Enter must be swallowed and keep Send disabled; otherwise it can
# become a form submit before the candidate is committed.
# Some IMEs report it as isComposing=false/keyCode=13 while composingRef
# is still pinned; it must be swallowed and keep Send disabled, else it
# submits before the candidate is committed.
step("BUG REPRO: Mac IME switch - Enter must not unblock submit")
clear()
composer.click()
@ -741,18 +713,15 @@ with sync_playwright() as p:
clear()
# 6g. Mac input-method switch - onBlur immediate recovery.
# Some Mac IME switches steal focus from the textarea (e.g. clicking
# the menu-bar language icon). The onBlur handler added for this bug
# resets composingRef unconditionally when the textarea loses focus.
# This is always safe: the OS commits or cancels any active composition
# before surrendering focus, so blur is a reliable reset point.
# Some Mac IME switches steal textarea focus; onBlur resets composingRef
# unconditionally. Safe because the OS commits/cancels composition before
# surrendering focus, so blur is a reliable reset point.
step("BUG REPRO: Mac IME switch - onBlur immediate recovery")
clear()
composer.click()
# Seed sendable content so the Send button's enabled/disabled state reflects
# composition state only, not empty-content gating.
# Seed sendable content so Send's state reflects composition only.
set_value_via_setter("hello")
# Simulate switching TO Chinese: compositionstart fires, compositionend never comes.
# Switch TO Chinese: compositionstart fires, compositionend never comes.
composer.evaluate(
"""(el) => {
el.focus();
@ -761,18 +730,16 @@ with sync_playwright() as p:
)
page.wait_for_timeout(200)
send_btn_mac_blur = page.locator('button[aria-label="Send message"]')
# Blur the textarea to simulate the OS stealing focus during an IME switch
# (e.g. the user clicks the menu-bar language icon).
# Blur to simulate the OS stealing focus during an IME switch.
composer.evaluate("(el) => el.blur()")
# onBlur calls setCompositionState(false) immediately; re-focus so React
# can render the updated Send-button state and we can locate it.
# onBlur clears composition; re-focus so React renders the updated Send state.
composer.click()
if send_btn_mac_blur.count() == 0:
soft_fail("Send button not found for Mac IME switch (onBlur) repro")
else:
try:
# 1500ms is well below the 2500ms watchdog: only onBlur can clear
# composingRef this quickly when no keydown is fired.
# 1500ms < 2500ms watchdog: only onBlur can clear composingRef this
# quickly when no keydown is fired.
expect(send_btn_mac_blur).not_to_be_disabled(timeout = 1_500)
info(
"Send button enabled within 1500ms after Mac IME switch + "
@ -789,8 +756,7 @@ with sync_playwright() as p:
info("Mac IME switch onBlur recovery PASS")
clear()
# 7. Final state. Filter benign 401 noise from the change-password redirect
# via is_benign_*; fail only on real errors.
# 7. Final state: filter benign 401 noise via is_benign_*; fail on real errors.
shoot("07-final")
real_page_errors = [e for e in page_errors if not is_benign_page_error(e)]
probe_cancel_500_allowance = expected_probe_cancel_500s[0]

582
tests/studio/playwright_chat_ui.py
View file

@ -1,40 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Comprehensive Studio chat UI test, run locally + in CI.
Covers:
1. /change-password through the UI (no API pre-rotate).
2. Model loaded by the time chat opens (the chat page's runtime
adapter pings /api/models/list; we trigger /api/inference/load
via page.evaluate so we don't need the password out-of-band).
3. Five chat turns, each deterministic (temperature handled at the
server level via Studio's default; we only assert non-empty).
4. Regenerate the last turn from the assistant action bar.
5. Composer toggle buttons: Thinking / Web search / Code execution
-- assert aria-label flips state on click.
6. Configuration sheet: open, drive Temperature slider via keyboard,
close.
7. Theme toggle through the account menu, multiple cycles, with a
deterministic computed-background-color check on
`document.documentElement` and `document.body`.
8. Sidebar nav: New Chat, Compare, Search, Recipes (URL changes).
9. Recents (history) cards: click an existing chat thread.
10. API tab via account menu -> Developer / api-keys.
11. Image attachment UI (upload widget reachable; vision response
not asserted because gemma-3-270m is text-only).
12. Reload + verify session JWT survives.
13. /api/health remains healthy.
14. Negative-auth post-UI-rotation: old=401, new=200.
15. Terminal-driven password rotation via subprocess(curl) to
/api/auth/change-password (NEW -> NEW2). Confirms refresh
tokens get revoked and that an out-of-band password change
(i.e. another tab / CLI / curl) invalidates the old creds.
16. Shutdown via the account menu's Shutdown menuitem + the
AlertDialog's "Stop server" action; wait for /api/health to
become unreachable (server process exited).
17. No uncaught page errors.
"""
"""Comprehensive Studio chat UI test, run locally + in CI."""
import json
import os
@ -48,9 +15,8 @@ import urllib.error
from pathlib import Path
from playwright.sync_api import expect, sync_playwright
# Shared robustness helpers live next to this script. Tests run as
# plain `python tests/studio/playwright_chat_ui.py` (not via pytest /
# import), so prepend the dir to sys.path before importing.
# Tests run as plain `python tests/studio/playwright_chat_ui.py` (not
# via pytest/import), so prepend this dir to sys.path before importing.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from _playwright_robust import ( # noqa: E402
chromium_launch_args,
@ -74,27 +40,18 @@ ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright")
ART = Path(ART_DIR)
ART.mkdir(parents = True, exist_ok = True)
# Strict mode -- when on (default in CI), the test fails loudly if any
# expected button / nav / dialog is missing instead of logging a WARN
# and continuing. Locally we leave it off so the test still runs against
# a partial Studio install.
# When on (default in CI), fail loudly on any missing button/nav/dialog
# instead of logging a WARN; off locally to run against a partial install.
STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
# Per-turn assistant-bubble wait. The free macos-14 runner (3 vCPU /
# 7 GB / no GPU) is ~3-5x slower at gemma-3-270m CPU inference than the
# free ubuntu-latest runner; "Say the word 'tree'" has been observed to
# hit the 180 s default exactly. STUDIO_UI_TURN_TIMEOUT_MS lets the Mac
# CI bump this without hard-coding a Mac branch in the test.
# Per-turn assistant-bubble wait. The free macos-14 runner is ~3-5x
# slower at gemma-3-270m CPU inference; this lets it bump the timeout.
TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000"))
# Wall-clock cap for the entire script. A healthy comprehensive run is
# 5-9 min; 12 min leaves headroom. Tunable via STUDIO_UI_WALL_TIMEOUT_S.
# See _playwright_robust.install_wall_clock_watchdog for rationale.
# Wall-clock cap for the whole script (healthy run is 5-9 min).
WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720"))
# Per-fetch budget for in-page fetches. The /api/inference/load call is
# usually the slowest legitimate request: it pulls the model into the
# llama.cpp worker. Give it ~3 min on a cold cache, less elsewhere.
# Per-fetch budget; /api/inference/load is the slowest (cold-cache GGUF load).
FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000"))
LOAD_FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_LOAD_TIMEOUT_MS", "180000"))
@ -119,10 +76,8 @@ def expected_default_model():
return override
# Parse DEFAULT_MODELS_GGUF as a literal out of defaults.py instead of
# importing it. The Playwright job installs Studio with --no-torch, so
# the studio.backend.core.inference package init (which eagerly imports
# the orchestrator -> structlog) and defaults.py's own
# `import utils.hardware.hardware as hw` are both unavailable.
# importing it: the --no-torch Playwright install can't import the
# inference package or defaults.py's hardware deps.
import ast
defaults_path = (
@ -154,12 +109,7 @@ def expected_default_model():
def soft_fail(m):
"""Hard fail in STRICT mode, info-warn otherwise.
Use for "this button should exist but didn't" assertions where
a missing element is a regression in CI but acceptable when
running against a partial Studio locally.
"""
"""Hard fail in STRICT mode, info-warn otherwise."""
if STRICT:
fail(m)
info(f"WARN (strict-off): {m}")
@ -190,56 +140,32 @@ with sync_playwright() as p:
label = "ui",
info = info,
)
# Pre-flight: bash-side wait_for already gated on /api/health
# before launching us, but the macos-14 free runner has been
# observed to surface a 200 /api/health while the auth DB is
# still finishing its migration. A second 30s probe inside the
# script catches that gap before we sink 60s into a change-
# password timeout. Diagnostic only -- the workflow's own wait
# is the authoritative gate, so we don't fail on miss.
# Pre-flight: macos-14 can surface a 200 /api/health while the auth
# DB is still migrating; this 30s probe catches that gap before we
# sink 60s into a change-password timeout. Diagnostic only.
wait_for_health(BASE, timeout = 30.0, info = info)
# Chromium launch args: see `tests/studio/_playwright_robust.py`.
# Bundles the macos-14 stability set (--single-process for the
# pipeTransport.js JSON-RPC crash) + new throttling kill set
# (--disable-background-timer-throttling and friends) that
# prevent Chromium from deprioritising the headless context's
# CPU/timers when it thinks the window is backgrounded -- which
# CI runners routinely flag.
browser = p.chromium.launch(
headless = True,
args = chromium_launch_args(),
)
ctx = browser.new_context(
viewport = {"width": 1280, "height": 900},
# Reduces motion so the theme toggle's view-transition
# animation doesn't briefly intercept pointer events
# (the running CSS view-transition leaves the html in a
# state where Playwright's actionability check fails).
# Reduce motion so view-transition animations don't intercept
# pointer events and break Playwright's actionability check.
reduced_motion = "reduce",
)
# Hard-disable CSS view-transitions: see _playwright_robust.py
# for the underlying init script. Necessary because Studio's theme
# toggle + sidebar collapse run their own startViewTransition()
# which can leave the <html> element intercepting pointer events
# for a beat after each route swap -- Playwright surfaces this as
# "<html> intercepts pointer events" on the next click.
# Hard-disable CSS view-transitions: Studio's theme toggle + sidebar
# collapse run startViewTransition() which can leave <html> intercepting
# pointer events for a beat after each route swap. See _playwright_robust.py.
install_view_transition_killer(ctx)
page = ctx.new_page()
# 60s default (was 30s) -- macos-14 free runner under
# --single-process Chromium is slow enough that page renders /
# webfonts / lazy-loaded routes routinely crowd 30s. Run
# 25494926834 hit Page.screenshot timeout AND
# locator.wait_for("#new-password") timeout under the old 30s
# default. 60s is conservative without bloating real-failure
# detection.
# 60s default (was 30s): macos-14 under --single-process Chromium is
# slow enough that renders/webfonts/lazy routes routinely crowd 30s.
page.set_default_timeout(60_000)
page_errors = []
page.on("pageerror", lambda e: page_errors.append(str(e)))
console_errors: list[str] = []
# Filtered console.error log -- excludes BENIGN_CONSOLE_ERROR_PATTERNS
# so the diagnostic dumps + final summary count only signals worth
# reading. Raw firehose is still surfaced via len(console_errors)
# vs len(filtered).
def _on_console(m):
if m.type != "error":
@ -252,11 +178,8 @@ with sync_playwright() as p:
page.on("console", _on_console)
# Per-turn HTTP-status capture: if a /v1/chat/completions request
# 4xx-rejects mid-test the symptom is a hung wait_for_function and
# a "FAIL: 1 non-benign pageerror events" line; this listener
# surfaces the underlying status codes so a flake is debuggable
# straight from the CI log without artifact spelunking.
# Capture /v1/chat/completions statuses so a mid-test 4xx (which
# surfaces only as a hung wait_for_function) is debuggable from the log.
chat_completions_responses: list[tuple[int, str]] = []
page.on(
"response",
@ -268,15 +191,9 @@ with sync_playwright() as p:
)
def shoot(name):
# Screenshots are diagnostic artifacts only -- never fail the
# test on a screenshot timeout. Page.screenshot waits for
# webfonts to fully load before snapshotting; on macos-14 free
# runners with --single-process Chromium, font loading on the
# Studio chat page (Inter / Geist Mono) regularly crowds the
# 30s default and crashes Page.screenshot. Bump the timeout
# AND wrap in try/except so the test progresses even if the
# screenshot can't be captured. animations='disabled' freezes
# any in-flight CSS transitions for a deterministic snap.
# Screenshots are diagnostic only -- never fail on a screenshot
# timeout. Page.screenshot waits for webfonts, which on macos-14
# can crowd the default; bump the timeout and swallow errors.
_n[0] += 1
try:
page.screenshot(
@ -290,31 +207,15 @@ with sync_playwright() as p:
# ─────────────────────────────────────────────────────
# 1. Change-password through the UI ("Setup your account").
# The bootstrap state injects window.__UNSLOTH_BOOTSTRAP__
# so the current-password is pre-seeded; we only enter the
# new password twice and submit. Match the workflow rename
# from "tool calling tests" pattern: this *is* the user's
# first-run experience.
# Bootstrap state pre-seeds the current password; we enter the
# new password twice and submit -- the user's first-run experience.
# ─────────────────────────────────────────────────────
step("change-password through UI (Setup your account)")
# Wait for the network to settle before touching the form. Without
# this, on macos-14 free runners under --single-process Chromium,
# the page sometimes redirects mid-test (the bootstrap state poll
# finishes after wait_for() returns, the React router decides
# we're "already authenticated" or "no longer must-change", and
# rerenders without #new-password). Letting the network idle first
# gives the bootstrap dispatch a chance to settle BEFORE we
# commit to the form path. Run 25497245250 / job 74820324136
# showed this exact sequence: wait_for() returned then
# page.fill('#new-password') timed out 60s later because the
# form had been replaced. Run 25578374480 / job 75091072289
# showed the same race a step deeper: pw_field.fill('#new-password')
# succeeded then page.fill('#confirm-password') hit a 60s timeout
# because a re-render between the two locators detached the
# second input. We wrap the whole goto/wait/fill/submit sequence
# in a 3-attempt retry, with a fresh page or hard reload between
# attempts so a re-render in the middle of one try doesn't poison
# the next.
# Settle the network before touching the form: a late bootstrap poll
# can rerender the page (dropping #new-password) mid-test. The whole
# goto/wait/fill/submit sequence is wrapped in a 3-attempt retry with
# a fresh page/reload between tries so a mid-try rerender doesn't
# poison the next.
form_err: Exception | None = None
for _form_attempt in range(3):
try:
@ -325,25 +226,14 @@ with sync_playwright() as p:
pass # best-effort -- proceed even if network never idles
pw_field = page.locator("#new-password")
pw_field.wait_for(state = "visible", timeout = 60_000)
# NOTE: do NOT call shoot() between wait_for and fill -- the
# screenshot's font-load wait gives the React form a chance to
# detach if any background state-poll fires. Take screenshots
# AFTER the form is committed instead.
# Do NOT shoot() between wait_for and fill -- the screenshot's
# font-load wait can let a background poll detach the form.
pw_field.fill(NEW, timeout = 60_000)
page.fill("#confirm-password", NEW, timeout = 60_000)
shoot("01-change-password-filled")
# Click submit AND wait for the POST /api/auth/change-password
# response in the same step. macos-14 free runners under
# --single-process Chromium occasionally hit
# net::ERR_NO_BUFFER_SPACE when the renderer requests a
# resource (run 25586583024 / job 75116256117 had the
# change-password POST silently buffer-fail and the page
# stayed on /change-password; even after my page.goto(BASE)
# recovery the auth state never persisted). Tying the
# click to the response wait surfaces the buffer-error
# IMMEDIATELY in this attempt rather than at the next
# composer.wait_for, so the next retry-iteration starts
# fresh with a known-bad starting state.
# Click submit AND wait for the POST response together so a
# macos-14 net::ERR_NO_BUFFER_SPACE buffer-fail surfaces now,
# not at the next composer.wait_for.
status, _ = click_and_wait_for_response(
page,
url_substr = "/api/auth/change-password",
@ -384,8 +274,7 @@ with sync_playwright() as p:
pass
if _form_attempt < 2:
# ERR_NO_BUFFER_SPACE needs the OS to recover socket
# buffers; immediate retry just re-fails. Back off
# 5s then 15s before next attempt.
# buffers; back off 5s then 15s before retrying.
if "ERR_NO_BUFFER_SPACE" in str(e):
backoff_s = 5 if _form_attempt == 0 else 15
print(
@ -394,8 +283,8 @@ with sync_playwright() as p:
flush = True,
)
time.sleep(backoff_s)
# Recovery: replace the page if it died, otherwise the
# next loop iteration's page.goto() handles the reload.
# Replace the page if it died; otherwise next iteration's
# page.goto() handles the reload.
page = recover_or_replace_page(
page,
ctx,
@ -409,17 +298,9 @@ with sync_playwright() as p:
# 2. Chat surface mounts, default model surface is visible.
# ─────────────────────────────────────────────────────
step("wait for composer to mount")
# The change-password POST resolves async and the React router
# rebuilds the tree (login form -> chat shell) on success. On
# macos-14 free runners under --single-process Chromium, the
# rebuild is heavy enough under software rendering that one of
# two things happens if we race straight into wait_for():
# (a) the composer textarea is still suspending and we burn
# the 60s ceiling waiting for it to mount, or
# (b) the renderer crashes mid-mount, which under
# --single-process takes the entire context down (next
# Playwright call returns TargetClosedError).
# Defend against both: settle network first, then attempt
# After change-password the router rebuilds login -> chat shell; on
# macos-14 racing straight into wait_for() either burns the timeout
# or crashes the renderer mid-mount. Settle network first, then
# wait_for with one recovery cycle on failure.
try:
page.wait_for_load_state("networkidle", timeout = 30_000)
@ -457,11 +338,9 @@ with sync_playwright() as p:
except Exception:
pass
if _attempt == 0:
# Recovery: re-navigate. If the page died (renderer
# gone under --single-process) we open a fresh page in
# the same context so the auth state in localStorage
# survives; otherwise we re-goto the same URL to force
# a clean re-render.
# Re-navigate: open a fresh page in the same context if
# the renderer died (localStorage auth survives), else
# re-goto to force a clean re-render.
page = recover_or_replace_page(
page,
ctx,
@ -475,9 +354,8 @@ with sync_playwright() as p:
raise last_err
shoot("03-chat-loaded")
# Pull the auth token now -- /api/models/list and
# /api/inference/load both require a bearer. The frontend
# stores it under "unsloth_auth_token" (auth/session.ts).
# /api/models/list and /api/inference/load need a bearer; the
# frontend stores it under "unsloth_auth_token" (auth/session.ts).
token = page.evaluate(
"() => localStorage.getItem('unsloth_auth_token')",
)
@ -502,12 +380,8 @@ with sync_playwright() as p:
if not token:
fail("could not obtain auth token after change-password")
# Verify the chat page's default model surface comes from
# backend/core/inference/defaults.py:DEFAULT_MODELS_GGUF[0],
# which is the canonical "what the user sees if nothing has
# been loaded yet" entry. A regression that reorders that
# list or hides the default would break the first-launch UX,
# which is what this assertion guards.
# Verify the chat page's default model matches DEFAULT_MODELS_GGUF[0]
# (defaults.py) -- guards the first-launch UX against list reorders.
step("default_models[0] matches DEFAULT_MODELS_GGUF[0]")
EXPECTED_DEFAULT = expected_default_model()
defaults_resp = evaluate_fetch(
@ -531,19 +405,16 @@ with sync_playwright() as p:
)
info(f"OK default_models[0] = {EXPECTED_DEFAULT}")
# The model selector button text on the chat page should say
# the default model's display name even before a model is
# loaded. The model-selector renders the current model name
# (or "Select model" if no current); for a fresh chat it
# should surface the default.
# The selector button should show the default model's name even
# before a model is loaded ("Select model" if none).
selector_btn = page.locator(
'button:has-text("Select model"), '
'button:has-text("gemma"), '
'button:has-text("Qwen"), '
'button:has-text("Llama")'
).first
# Best-effort: the selector re-mounts as /api/models/list resolves,
# so use a short timeout and skip the snapshot on miss.
# Best-effort: selector re-mounts as /api/models/list resolves, so
# use a short timeout and skip the snapshot on miss.
sel_text = ""
try:
sel_text = (selector_btn.text_content(timeout = 2_000) or "").strip()
@ -554,17 +425,12 @@ with sync_playwright() as p:
shoot("03b-default-model-button")
# ─────────────────────────────────────────────────────
# 3. Trigger model load via the page's session cookies.
# Equivalent to the user clicking a model in the picker;
# we just call the same endpoint the picker would.
# 3. Trigger model load via the same endpoint the picker uses.
# ─────────────────────────────────────────────────────
step("load GGUF via /api/inference/load (uses session cookie)")
# Token already fetched above; reuse it for the load call.
# AbortSignal-bounded: the macos-14 --single-process Chromium had been
# observed wedging on this exact in-page fetch (run 25696797934 / job
# 75446949358) with zero further requests reaching the server. The
# 3-min budget is generous for a cold-cache GGUF load; on a wedge we
# surface a clean failure instead of a 30-min runner cancel.
# AbortSignal-bounded: macos-14 has been seen wedging on this fetch.
# The 3-min budget is generous for a cold-cache load; a wedge fails
# cleanly instead of forcing a 30-min runner cancel.
load_resp = evaluate_fetch(
page,
f"{BASE}/api/inference/load",
@ -587,25 +453,20 @@ with sync_playwright() as p:
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
# to make the chat composer pick up the loaded model.
# Studio caches model state in zustand; reload so the composer picks
# up the loaded model.
page.reload()
composer = page.locator('textarea[aria-label="Message input"]')
composer.wait_for(state = "visible", timeout = 60_000)
# ─────────────────────────────────────────────────────
# 3b. Model picker search bar -- click the model selector,
# type into the search box, verify filtering. We don't
# actually select a different model (that would trigger a
# multi-GB download); we just exercise the typeahead so a
# regression in the picker mount / debounced HF search would
# surface here.
# 3b. Model picker search bar -- exercise the typeahead filter.
# We don't actually select a different model (multi-GB download);
# this just catches picker-mount / debounced HF-search regressions.
# ─────────────────────────────────────────────────────
step("model picker: open + drive search bar")
# Stable selector first: [data-tour="chat-model-selector"] is the
# guided-tour anchor on the model picker button (app-sidebar.tsx).
# If the tour anchor moves the tour breaks, so this selector is at
# least as stable as anything else in the codebase.
# Prefer the guided-tour anchor [data-tour="chat-model-selector"]
# (app-sidebar.tsx) -- as stable as anything in the codebase.
picker_btn = page.locator('[data-tour="chat-model-selector"]').first
if picker_btn.count() == 0:
# Fall back to text-based locators for older Studio builds.
@ -626,11 +487,9 @@ with sync_playwright() as p:
if search.count() == 0:
soft_fail("model picker search input not found")
else:
# Type "qwen" -> capture popover text. Type "llama" -> capture
# again. The two text snapshots must DIFFER, proving the
# typeahead actually filters the list (a regression that
# rendered the picker but ignored input would silently pass
# the old version of this test).
# "qwen" then "llama" popover text must DIFFER, proving the
# typeahead actually filters (else an ignored-input regression
# would silently pass).
def picker_visible_text():
return page.evaluate("""() => {
const el = document.querySelector(
@ -672,18 +531,15 @@ with sync_playwright() as p:
]
def _bubble_count():
"""Total number of [data-role='assistant'] elements (empty or not)."""
"""Total [data-role='assistant'] elements (empty or not)."""
return page.evaluate("""() => {
return document.querySelectorAll('[data-role="assistant"]').length;
}""")
def send_and_wait(prompt, idx):
# 1. Wait until the previous turn has fully stopped: Send
# button is attached AND Stop button is detached. The
# assistant-ui composer hot-swaps these inside a single
# DOM slot; relying on Stop's detached state alone is
# racy (the slot can briefly show neither during
# transition).
# 1. Wait until the previous turn fully stopped: Send attached
# AND Stop detached. The composer hot-swaps both in one DOM
# slot, so Stop's detached state alone is racy.
page.wait_for_selector(
'button[aria-label="Send message"]',
state = "attached",
@ -696,34 +552,26 @@ with sync_playwright() as p:
timeout = 5_000,
)
except Exception:
# Stop button still hanging on -- that's the prior turn
# mid-stream. Wait it out at the full per-turn budget.
# Stop still on -- prior turn mid-stream. Wait it out at the
# full per-turn budget.
page.wait_for_selector(
'button[aria-label="Stop generating"]',
state = "detached",
timeout = TURN_TIMEOUT_MS,
)
# 2. Snapshot total bubble count BEFORE send. We then wait
# for total count to grow by exactly 1 (proves the new
# placeholder rendered) and for the Stop button to come
# + go (proves the new turn ran end-to-end). We do NOT
# require the new bubble's text to be non-empty: an
# empty assistant response is a legitimate model output,
# not a test failure. The earlier "non-empty count >=
# baseline + 1" predicate broke when any prior turn
# streamed empty (which gemma-3-270m DOES on simple
# prompts at temperature 0), because that empty bubble
# became permanently "stuck" below the moving threshold.
# 2. Snapshot total bubble count before send; we wait for it to
# grow by exactly 1. We do NOT require non-empty text: an
# empty assistant response is legitimate (gemma-3-270m does
# this at temp 0), and the old non-empty predicate got stuck
# on such bubbles.
bubbles_before = _bubble_count()
composer.click()
composer.fill(prompt)
page.locator('button[aria-label="Send message"]').click()
# 3. Wait for the new placeholder bubble to render. This
# confirms the click was actionable AND the request
# issued (assistant-ui only mounts the placeholder once
# the runtime accepts the message).
# 3. Wait for the new placeholder bubble to render -- confirms
# the click was actionable and the request issued.
page.wait_for_function(
"""(want) => {
return document.querySelectorAll(
@ -734,12 +582,9 @@ with sync_playwright() as p:
timeout = TURN_TIMEOUT_MS,
)
# 4. Wait for streaming to FINISH for this specific turn.
# We wait for Stop button to APPEAR (proves streaming
# started) with a short budget; if it never appears,
# that's fine -- gemma-3-270m can finish before the
# Stop button paints. Either way we then wait for it
# to be detached at the full per-turn budget.
# 4. Wait for this turn's streaming to finish. Stop may never
# appear (gemma-3-270m can finish before it paints), so its
# appearance is best-effort; then wait for it to detach.
try:
page.wait_for_selector(
'button[aria-label="Stop generating"]',
@ -768,10 +613,8 @@ with sync_playwright() as p:
if len(texts) < len(prompts):
fail(f"expected >= {len(prompts)} assistant bubbles, got {len(texts)}")
info(f"five turn lengths = {[len(t) for t in texts[:5]]}")
# Surface /v1/chat/completions HTTP status distribution so a flake
# is debuggable from the CI log directly. A 4xx during a chat
# turn is almost always the upstream cause of a hung
# wait_for_function on a downstream turn.
# Surface /v1/chat/completions status distribution: a 4xx here is
# usually the cause of a hung wait_for_function downstream.
if chat_completions_responses:
statuses = [code for code, _ in chat_completions_responses]
bad = [code for code in statuses if code >= 400]
@ -804,11 +647,9 @@ with sync_playwright() as p:
shoot("05-after-regenerate")
info("regenerate completed")
else:
# Don't strict-fail on regenerate -- the assistant-ui
# ActionBarPrimitive.Reload doesn't expose a stable
# aria-label, so the test depends on tooltip text matching
# which is tied to the icon set. Soft-skip until we add a
# data-testid in the action bar (TODO).
# Don't strict-fail: ActionBarPrimitive.Reload has no stable
# aria-label so the locator relies on icon-tied tooltip text.
# Soft-skip until we add a data-testid (TODO).
info("WARN regenerate button not visible (known-fragile locator, skipped)")
# ─────────────────────────────────────────────────────
@ -822,23 +663,20 @@ with sync_playwright() as p:
shoot("06-after-extra-turns")
# ─────────────────────────────────────────────────────
# 7. Composer toggle buttons. Each renders with an
# aria-label that flips between "Disable X" / "Enable X"
# depending on its current state (shared-composer.tsx).
# 7. Composer toggle buttons. Each aria-label flips between
# "Disable X" / "Enable X" with state (shared-composer.tsx).
# ─────────────────────────────────────────────────────
step("composer toggle buttons (Thinking / Web search / Code execution)")
for feature in ("thinking", "web search", "code execution"):
# Look for either "Disable X" or "Enable X" -- whichever
# is currently rendered.
# Match whichever of "Disable X" / "Enable X" is rendered.
toggle = page.locator(
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")
continue
# Skip if the model doesn't support this capability (the
# button is rendered disabled). gemma-3-270m, for instance,
# has no reasoning so "Disable thinking" is permanent-disabled.
# Skip if the button is disabled (model lacks the capability;
# e.g. gemma-3-270m has no reasoning, so thinking stays disabled).
if toggle.is_disabled():
info(f"toggle '{feature}' is disabled for this model -- skip")
continue
@ -866,8 +704,7 @@ with sync_playwright() as p:
shoot("07-toggles-cycled")
# ─────────────────────────────────────────────────────
# 8. Configuration sheet: open, find Temperature slider,
# press Home (→ 0), close.
# 8. Configuration sheet: open, drive Temperature slider, close.
# ─────────────────────────────────────────────────────
cfg_open = page.locator('button[aria-label="Open configuration"]').first
if cfg_open.count() > 0:
@ -875,13 +712,9 @@ with sync_playwright() as p:
cfg_open.click()
page.wait_for_timeout(500)
shoot("08-config-open")
# ParamSlider uses Radix UI Slider. Each slider gets a
# role="slider" attribute. Walk every slider in the sheet
# by index, focus it, send Home (-> min) so the test
# state is fully deterministic. Whatever the labels are
# ("Temperature", "Top P", "Min P", "Repetition penalty",
# max_tokens etc.), we drive them all to min so a
# regression that locks a slider returns errors here.
# Walk every Radix slider (role="slider") by index, focus it,
# press Home (-> min) for deterministic state; a locked slider
# surfaces an error here.
sliders = page.locator('[role="slider"]')
n_sliders = sliders.count()
info(f"configuration sheet exposes {n_sliders} slider(s)")
@ -895,10 +728,8 @@ with sync_playwright() as p:
except Exception as exc:
info(f" slider[{idx}] focus/Home failed: {exc!r}")
shoot("09-config-all-min")
# Then drive Temperature specifically to 0.0 to make the
# downstream chat deterministic. Temperature is the *first*
# slider in the sheet (configuration-sheet.tsx renders it
# first); Home already pinned it to 0.
# Temperature is the first slider (configuration-sheet.tsx), so
# Home already pinned it to 0 for determinism.
info("Temperature set to slider min (0.0) for determinism")
# Close.
close_btn = page.locator('button[aria-label="Close configuration"]').first
@ -909,24 +740,17 @@ with sync_playwright() as p:
page.wait_for_timeout(300)
# ─────────────────────────────────────────────────────
# 9. Theme toggle -- multiple cycles + deterministic
# computed-background-color check. The light theme
# uses near-white (>240); dark uses near-black (<40).
# 9. Theme toggle -- multiple cycles + computed-bg-color check
# (light is near-white >240; dark is near-black <40).
# ─────────────────────────────────────────────────────
acct = page.locator('button[aria-label$=" account menu"]').first
if acct.count() > 0:
step("theme toggle x3 with computed-color assertion")
observed = []
for cycle in range(3):
# Wait for any prior dropdown to fully detach. The Radix
# Account-menu sets data-state="open" while the view-
# transition is mid-flight; clicking it again before that
# clears would no-op silently and the for-loop bailed
# after cycle 1 in earlier runs. The view transition triggered
# by the theme toggle can run >700ms on slow CI runners, so
# both the "menu detached" wait and the "menu appeared" wait
# need a comfortable budget; 3s was too tight and caused
# cycle-2 flake.
# Wait for any prior dropdown to fully detach: clicking while
# the view-transition is still open no-ops silently. The
# transition can run >700ms on slow CI, so use a roomy budget.
try:
page.wait_for_function(
"""() => {
@ -941,9 +765,8 @@ with sync_playwright() as p:
except Exception:
pass
page.wait_for_timeout(250)
# Try the click + wait; if the first click silently no-oped
# (e.g. mid-view-transition swallowed the event), retry once
# after pressing Escape to force-close any stray popup.
# Retry once (after Escape to clear stray popups) if the first
# click is silently swallowed mid-view-transition.
opened = False
for attempt in range(2):
try:
@ -975,15 +798,10 @@ with sync_playwright() as p:
page.keyboard.press("Escape")
soft_fail(f"theme cycle {cycle + 1}: theme menuitem missing")
break
# Click sequence with two fallbacks. On small CI viewports the
# Radix dropdown can render the theme item below the visible
# area; force=True still requires the element to be in the
# viewport, so the regular .click() fails with "Element is
# outside of the viewport". Fall back to scroll-into-view +
# click, then to a synthetic .click() via evaluate() that
# bypasses Playwright's viewport check entirely (Radix's
# menuitem handler only needs the click event, not a real
# pointer landing on a pixel).
# Click with fallbacks: a small CI viewport can push the item
# off-screen (force=True still needs it in viewport). Fall back
# to scroll-into-view, then a synthetic evaluate() .click() that
# skips Playwright's viewport check.
click_err = None
for click_attempt in range(3):
try:
@ -1005,10 +823,8 @@ with sync_playwright() as p:
f"theme cycle {cycle + 1}: theme menuitem click failed " f"({click_err!r})"
)
break
# Settle. The ".dark" class on <html> is the ground
# truth (theme-store toggles only that class); the
# ".light" sibling is steady-state from next-themes
# so don't gate on it.
# Settle. The ".dark" class on <html> is the ground truth
# (theme-store toggles only that); don't gate on ".light".
page.wait_for_timeout(700)
bg = page.evaluate("""() => {
const root = document.documentElement;
@ -1022,21 +838,17 @@ with sync_playwright() as p:
observed.append(bg)
shoot(f"10-theme-cycle-{cycle + 1}")
info(f" cycle {cycle + 1}: dark={bg['isDark']} body bg={bg['bg']!r}")
# Sanity check: across cycles we should observe both a
# light state (body bg roughly near-white) and a dark state
# (body bg near-black). If we only saw one polarity the
# toggle didn't flip.
# Across cycles we should see both a near-white (light) and a
# near-black (dark) body bg; one polarity means the toggle stuck.
rgbs = [parse_rgb(o["bg"]) for o in observed if parse_rgb(o["bg"])]
light_seen = any(min(r) > 220 for r in rgbs)
dark_seen = any(max(r) < 60 for r in rgbs)
if len(observed) < 3:
soft_fail(f"theme toggle ran only {len(observed)} cycle(s), expected 3")
# Don't strict-fail on "both polarities observed" -- the
# CI runner's prefers-color-scheme + Studio's "system" default
# can collapse to a single polarity even after a successful
# toggle (the .dark classlist toggles correctly, but the
# resolved theme can stay constant). Surface as info; the
# 3-cycle loop completion above is the real invariant.
# Don't strict-fail on both polarities: the runner's
# prefers-color-scheme + Studio's "system" default can collapse
# to one polarity even when .dark toggles correctly. The 3-cycle
# completion above is the real invariant.
if light_seen and dark_seen:
info("OK light + dark computed background colors observed")
else:
@ -1050,14 +862,10 @@ with sync_playwright() as p:
# 10. Sidebar nav: New Chat, Compare, Search, Recipes.
# ─────────────────────────────────────────────────────
def click_nav(label, expected_url_pat = None):
# Resolve the sidebar nav button. The plain
# get_by_role("button", name=...) lookup works on Linux
# Chromium because the accessible-name algorithm there picks
# up `tooltip={label}` from SidebarMenuButton, but on macOS
# Chromium the tooltip-derived name is sometimes empty when
# the sidebar collapses to icon-only mode. Fall back through
# progressively more permissive locators so the test stays
# green on both platforms.
# Resolve the sidebar nav button. get_by_role(name=...) works on
# Linux but the tooltip-derived name can be empty on macOS when
# the sidebar collapses to icons, so fall back to more permissive
# locators.
candidates = [
page.get_by_role("button", name = re.compile(rf"^\s*{label}\s*$", re.I)).first,
page.locator(f'button:has-text("{label}")').first,
@ -1072,11 +880,10 @@ with sync_playwright() as p:
if btn is None:
soft_fail(f"nav '{label}' not found")
return False
# force=True bypasses Playwright's actionability check. The
# button IS visible + enabled, but the post-theme-toggle view-
# transition can leave <html> reported as the topmost element
# for a beat (we already neutralise startViewTransition via
# add_init_script; this is belt-and-suspenders).
# force=True bypasses the actionability check: the post-toggle
# view-transition can briefly report <html> as topmost even
# though the button is visible + enabled (belt-and-suspenders
# atop the startViewTransition neutraliser).
try:
btn.click(force = True, timeout = 5_000)
except Exception as exc:
@ -1094,15 +901,15 @@ with sync_playwright() as p:
step("sidebar nav: New Chat -> Compare -> Search -> Recipes")
click_nav("New Chat", r"/chat")
shoot("11-new-chat")
# Compare moved into the composer + menu (Tools and attachments).
# Compare moved into the composer "Tools and attachments" menu.
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.
# Compare chat moved into 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()
@ -1141,10 +948,8 @@ with sync_playwright() as p:
composer.wait_for(state = "visible", timeout = 60_000)
# ─────────────────────────────────────────────────────
# 11. API / Developer tab via account menu -> opens the
# Settings dialog with the api-keys tab. Verify we can see
# the Create API Key form (or existing keys table); regressions
# that hide the api-keys management UI surface here.
# 11. API / Developer tab via account menu -> Settings dialog,
# api-keys tab. Guards against the management UI being hidden.
# ─────────────────────────────────────────────────────
if acct.count() > 0:
step("Developer (API) tab via account menu")
@ -1175,16 +980,13 @@ with sync_playwright() as p:
page.keyboard.press("Escape")
# ─────────────────────────────────────────────────────
# 11b. Recipes tab: verify cards render + we can click one.
# The Recipes route renders a grid of preset cards; a
# regression that breaks the loader would render zero cards
# or crash the route.
# 11b. Recipes tab: cards render + we can click one. A broken
# loader would render zero cards or crash the route.
# ─────────────────────────────────────────────────────
step("Recipes tab: cards render + click first card")
page.goto(f"{BASE}/data-recipes")
page.wait_for_timeout(1500)
# Recipe cards are rendered as <a> or button elements; count
# all clickable headings under main + screenshot.
# Count clickable headings/cards under main, then screenshot.
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}")
@ -1205,30 +1007,17 @@ with sync_playwright() as p:
composer.wait_for(state = "visible", timeout = 60_000)
# ─────────────────────────────────────────────────────
# 11c. Recents: the chat sidebar lists previous threads. We
# already created several turns above (which gets persisted
# as a thread). Find the sidebar's recents region and click
# the most-recent entry. This catches regressions in the
# thread-history loader / route param plumbing.
# 11c. Recents: click the most-recent thread (we persisted one
# via the turns above). Guards the thread-history loader / route.
# ─────────────────────────────────────────────────────
step("Recents: click previous chat in sidebar")
# We sent the prompts ["Reply with exactly: hello", "What is 1+1?",
# "Reply with exactly: world", ...] above. The thread title that
# gets persisted is typically a snippet of the first user message
# (Studio summarises after a few turns). We accept either a literal
# word from one of our prompts OR a short Studio-summary heuristic.
# The persisted thread title is usually a snippet of the first user
# message, so accept any of our prompt keywords.
PROMPT_KEYWORDS = ("hello", "world", "tree", "yes", "1+1", "2+2")
# Use the structural data-testid the frontend renders on each
# chat-history entry (studio/frontend/src/features/chat/thread-
# sidebar.tsx). The previous text-filtered selector
# "aside a, aside button, [data-sidebar='sidebar'] a, ..."
# matched coalesced sidebar nav text like 'unslothBETA',
# 'UUnslothUnsloth' which the EXCLUDE regex didn't strip; the
# test then clicked nav links, lost its frame, hit per-locator
# timeouts and burned 13-23 minutes per platform on this single
# step (run 25537467494 macui = 23m9s, winui = 13m6s, linui = 13m5s).
# Belt-and-suspenders: bound the whole step at 30s so a misbehaving
# selector can never blow up wallclock the way the old loop did.
# Use the structural data-testid (thread-sidebar.tsx): the old
# text-filtered selector matched coalesced nav text and burned
# 13-23 min per platform. Also bound the whole step at 30s so a
# misbehaving selector can't blow up wallclock.
threads = page.locator('[data-testid="recent-thread"]')
deadline = time.monotonic() + 30
clicked_recent = False
@ -1247,9 +1036,7 @@ with sync_playwright() as p:
page.wait_for_timeout(500)
shoot("15d-recent-clicked")
info(f"OK clicked recent entry: {t[:60]!r}")
# Strict check: after clicking the Recents entry, the
# thread we land on must include at least one of our
# prompts in its rendered messages.
# The landed thread must include at least one of our prompts.
turns_text = page.evaluate(
"""() => {
const els = document.querySelectorAll(
@ -1281,18 +1068,13 @@ with sync_playwright() as p:
composer.wait_for(state = "visible", timeout = 60_000)
# ─────────────────────────────────────────────────────
# 12. Image attachment UI (upload widget reachable). The
# current model is text-only so we don't assert a vision
# response -- just that the attachment button is there
# and the file input accepts a PNG. CI's gemma-4-E2B
# job covers the actual vision path.
# 12. Image attachment UI reachable. The current model is text-only,
# so just check the button exists (CI's gemma-4-E2B covers vision).
# ─────────────────────────────────────────────────────
step("attachment widget reachable")
attach = page.locator('button[aria-label="Add Attachment"]').first
if attach.count() > 0:
# Just hover -- triggering the file picker mid-test
# would block on a native dialog. Verifying the
# button is reachable is enough.
# Only hover -- clicking would block on the native file dialog.
attach.hover()
page.wait_for_timeout(200)
shoot("16-attachment-hover")
@ -1331,19 +1113,13 @@ with sync_playwright() as p:
info("OK old=401, new=200")
# ─────────────────────────────────────────────────────
# 16. Out-of-band ("terminal") password rotation.
# POST /api/auth/change-password from a real subprocess(curl)
# invocation -- this is the same surface a sysadmin / another
# tab / a desktop helper would use, and the security promise
# is: rotating the password from "the terminal" must invalidate
# the previous credentials. The endpoint also revokes refresh
# tokens server-side (auth.py:152), so /api/auth/refresh from
# the still-open browser context must fail too.
# 16. Out-of-band ("terminal") password rotation via subprocess(curl).
# Rotating from a shell must invalidate the old creds and revoke
# refresh tokens server-side (auth.py:152), so the browser's
# /api/auth/refresh must fail too.
# ─────────────────────────────────────────────────────
step("rotate password via subprocess(curl) -- the 'terminal' path")
# Get a fresh access token by logging in via the API rather than
# reusing whatever's in localStorage; this matches what an admin
# would actually do from a shell.
# Log in via the API for a fresh token (what an admin does from a shell).
login_proc = subprocess.run(
[
"curl",
@ -1400,9 +1176,8 @@ with sync_playwright() as p:
fail(f"after CLI rotation, NEW2 pw should be 200, got {s_new2}")
info("OK after CLI rotation: NEW=401, NEW2=200 -- old studio creds dead")
# The browser still has the pre-rotation access token. Refresh
# tokens were revoked server-side by /change-password (auth.py),
# so /api/auth/refresh from the browser context must now fail.
# /change-password revoked refresh tokens server-side (auth.py), so
# the browser's /api/auth/refresh must now fail.
refresh_after = evaluate_fetch(
page,
f"{BASE}/api/auth/refresh",
@ -1419,21 +1194,15 @@ with sync_playwright() as p:
)
# ─────────────────────────────────────────────────────
# 17. Shutdown button via the account menu.
# The Shutdown menuitem opens an AlertDialog ("Stop Unsloth
# Studio?") whose primary action is "Stop server"; clicking
# it POSTs /api/shutdown and then replaces document.body with
# the "Unsloth Studio has stopped" placeholder. /api/health
# should become unreachable shortly after.
# 17. Shutdown via the account menu. The "Stop server" action
# POSTs /api/shutdown, swaps in the "Unsloth Studio has stopped"
# placeholder, and /api/health goes unreachable shortly after.
# ─────────────────────────────────────────────────────
step("Shutdown via account menu")
# Re-login through the UI with NEW2 so the browser has a valid
# access token for the /api/shutdown call (the previous one
# was invalidated by the CLI rotation above).
# The CLI rotation left a stale token, so the SPA auth guard can
# client-side-redirect mid-navigation and abort this goto with
# net::ERR_ABORTED. Resolve on domcontentloaded and tolerate the
# abort; the password-field wait below confirms we reached /login.
# Re-login with NEW2 for a valid /api/shutdown token (CLI rotation
# invalidated the old one). The stale token can make the SPA auth
# guard abort this goto with ERR_ABORTED; resolve on
# domcontentloaded and tolerate it -- the pw-field wait confirms /login.
try:
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
except Exception as exc:
@ -1468,9 +1237,8 @@ with sync_playwright() as p:
stop_btn.wait_for(state = "visible", timeout = 5_000)
stop_btn.click()
# Wait for the post-shutdown placeholder body. The component
# replaces document.body.innerHTML with text containing
# "Unsloth Studio has stopped." once /api/shutdown returns ok.
# Wait for the post-shutdown placeholder body (the component swaps in
# "Unsloth Studio has stopped." once /api/shutdown returns ok).
try:
page.wait_for_function(
"""() => /Unsloth Studio has stopped/.test(document.body.innerText)""",
@ -1481,8 +1249,7 @@ with sync_playwright() as p:
except Exception as exc:
info(f"WARN shutdown placeholder didn't render: {exc!r}")
# Now /api/health must become unreachable (process exited or is
# at least not listening). Poll for up to 15 s.
# /api/health must now be unreachable; poll for up to 15s.
host = re.sub(r"^https?://", "", BASE).split(":")[0]
port = int(re.search(r":(\d+)", BASE).group(1)) if ":" in BASE else 80
deadline = time.time() + 15
@ -1502,19 +1269,10 @@ with sync_playwright() as p:
except urllib.error.URLError as exc:
info(f"OK /api/health unreachable: {exc!r}")
# Some pageerrors are benign in this test:
# - "Request failed (422)": the OpenAI-compatible chat-completions
# endpoint rejects rapid-fire/malformed requests with 422. The
# surfaced error is a network-layer bubble-up, NOT a JS bug,
# and the per-turn flow already validates message-by-message
# correctness. Filtering these here keeps the pageerror gate
# focused on actual frontend regressions (TypeError, ReferenceError,
# null deref, etc.).
# - "Failed to fetch" / "NetworkError" after the Shutdown click:
# the server is intentionally dead by then; any in-flight
# fetch fails by design.
# The full list lives in `_playwright_robust.BENIGN_PAGE_ERROR_PATTERNS`
# so playwright_extra_ui.py shares the same gate.
# Some pageerrors are benign: chat-completions 422s (network-layer
# bubble-up, not a JS bug; per-turn flow already validates each turn)
# and fetch failures after Shutdown (server is dead by design). Full
# list in `_playwright_robust.BENIGN_PAGE_ERROR_PATTERNS`.
real_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)]
if page_errors:

148
tests/studio/playwright_extra_ui.py
View file

@ -1,27 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Studio extra-UI Playwright test.
Covers the user-visible surfaces that the main chat-UI test doesn't:
1. Compare tab (/chat?compare=...): assign two models, send 2 prompts,
assert both panes respond.
2. Recipes editor (/data-recipes/$recipeId): click first template,
verify the recipe-studio canvas mounts, open + close the Preview
dialog.
3. Export route (/export): chat-only mode redirects to /chat;
non-chat-only mode shows the export form fields.
4. Studio training route (/studio): chat-only mode redirects;
non-chat-only verifies the tabs + sections exist.
5. Settings dialog tabs: Cmd/Ctrl-, opens the dialog; cycle through
each tab and verify it isn't blank.
The test assumes Studio is freshly booted (must_change_password=true)
on BASE_URL with the bootstrap password in STUDIO_OLD_PW. It does its
own change-password through the UI + model load via /api/inference/load,
matching the pattern in playwright_chat_ui.py.
"""
"""Studio extra-UI Playwright test: Compare tab, Recipes editor, /export, /studio, Settings tabs."""
import json
import os
@ -33,9 +13,7 @@ import urllib.request
from pathlib import Path
from playwright.sync_api import sync_playwright
# Shared robustness helpers live next to this script. Tests run as
# plain `python tests/studio/playwright_extra_ui.py` (not via pytest /
# import), so prepend the dir to sys.path before importing.
# Run as a plain script (not via pytest), so prepend the dir to sys.path.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from _playwright_robust import ( # noqa: E402
chromium_launch_args,
@ -57,9 +35,7 @@ ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_extra")
ART = Path(ART_DIR)
ART.mkdir(parents = True, exist_ok = True)
STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
# Mirrors playwright_chat_ui.py. macos-14 free runners need a longer
# turn timeout because gemma-3-270m CPU inference is 3-5x slower than
# ubuntu-latest's.
# Longer turn timeout: gemma-3-270m CPU inference is 3-5x slower on macos-14 runners.
TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000"))
WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720"))
FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000"))
@ -90,11 +66,7 @@ def soft_fail(m: str) -> None:
def runtime_warn(m: str) -> None:
"""Warn about a runtime-coupled assertion that depends on a real
model loaded into the Compare panes. STRICT mode gates selector
presence (those MUST hold) but not Compare-pane streaming, which
is still flaky when no explicit pane model is set.
"""
"""Warn about a runtime-coupled assertion (Compare-pane streaming) that STRICT does not gate."""
info(f"WARN (runtime): {m}")
@ -104,13 +76,9 @@ with sync_playwright() as p:
label = "ui-extra",
info = info,
)
# Health pre-flight (best-effort). Same rationale as in
# playwright_chat_ui.py: bash-side health wait can succeed before
# the auth DB has finished migrating on macos-14 free runners.
# Health pre-flight: bash-side health wait can pass before the auth DB migrates on macos-14.
wait_for_health(BASE, timeout = 30.0, info = info)
# Chromium launch args: see `tests/studio/_playwright_robust.py`.
# Bundles macos-14 stability + new throttling-kill flags shared
# with playwright_chat_ui.py.
# Chromium launch args: see tests/studio/_playwright_robust.py.
browser = p.chromium.launch(
headless = True,
args = chromium_launch_args(),
@ -121,20 +89,12 @@ with sync_playwright() as p:
)
install_view_transition_killer(ctx)
page = ctx.new_page()
# See playwright_chat_ui.py -- 60s default for macos-14 free
# runner with --single-process Chromium. The extra-UI script is
# the SECOND Studio boot of the job, so the runner is even
# warmer (slower disk cache, contended Chromium state).
# 60s default for slow macos-14 --single-process Chromium (second Studio boot of the job).
page.set_default_timeout(60_000)
page_errors = []
# Filter out known-benign React errors that fire when the Compare
# flow's second prompt races the first prompt's SSE stream, or when
# /export's lazy-loaded sections haven't finished mounting before
# the error boundary trips. Both are timing artefacts on slow CI
# runners (macos-14 free), not Studio bugs. The base list lives in
# `_playwright_robust.BENIGN_PAGE_ERROR_PATTERNS` so the chat_ui
# test shares it.
# Filter out known-benign React errors (timing artefacts on slow CI runners, not Studio bugs);
# shared base list lives in _playwright_robust.BENIGN_PAGE_ERROR_PATTERNS.
def _on_pageerror(e):
msg = str(e)
if is_benign_page_error(msg):
@ -145,9 +105,7 @@ with sync_playwright() as p:
page.on("pageerror", _on_pageerror)
def shoot(name: str) -> None:
# See playwright_chat_ui.py:shoot -- screenshots are diagnostic,
# never fail the test on a font-load timeout under
# --single-process Chromium on macos-14 free runners.
# Screenshots are diagnostic; never fail the test on a font-load timeout.
_n[0] += 1
try:
page.screenshot(
@ -163,10 +121,8 @@ with sync_playwright() as p:
# Setup: change-password through the UI + model load.
# ─────────────────────────────────────────────────────
step("setup: change-password + model load")
# 3-attempt retry mirrors playwright_chat_ui.py: form re-renders
# mid-fill on macos-14 free runners detach #new-password OR
# #confirm-password between locator and fill, hitting 60s timeouts.
# Each retry re-navigates with a fresh page if the old one died.
# 3-attempt retry: form re-renders mid-fill on macos-14 can detach the password
# fields between locator and fill; each retry re-navigates with a fresh page if needed.
form_err: Exception | None = None
for _form_attempt in range(3):
try:
@ -179,11 +135,8 @@ with sync_playwright() as p:
pw_field.wait_for(state = "visible", timeout = 60_000)
pw_field.fill(NEW, timeout = 60_000)
page.fill("#confirm-password", NEW, timeout = 60_000)
# Click submit AND wait for the POST response together --
# surfaces a server-side reject (or net::ERR_NO_BUFFER_SPACE
# buffer-fail on macos-14) immediately rather than discovering
# it 60s later via a downstream composer.wait_for. Same shape
# as playwright_chat_ui.py's change-password block.
# Click submit AND wait for the POST response together so a server-side reject
# surfaces immediately rather than 60s later via a downstream composer.wait_for.
status, _ = click_and_wait_for_response(
page,
url_substr = "/api/auth/change-password",
@ -212,9 +165,7 @@ with sync_playwright() as p:
flush = True,
)
if _form_attempt < 2:
# ERR_NO_BUFFER_SPACE needs the OS to recover socket
# buffers; immediate retry just re-fails. Back off
# 5s then 15s before next attempt.
# ERR_NO_BUFFER_SPACE needs the OS to recover socket buffers; back off 5s then 15s.
if "ERR_NO_BUFFER_SPACE" in str(e):
backoff_s = 5 if _form_attempt == 0 else 15
print(
@ -231,10 +182,8 @@ with sync_playwright() as p:
)
if form_err is not None:
raise form_err
# Same defense-in-depth as playwright_chat_ui.py: settle network,
# then wait_for with one recovery cycle. The post-submit React
# re-render can either leave the composer suspending or crash the
# renderer outright under --single-process Chromium on macos-14.
# Settle network, then wait_for with one recovery cycle: the post-submit React
# re-render can leave the composer suspending or crash the renderer on macos-14.
try:
page.wait_for_load_state("networkidle", timeout = 30_000)
except Exception:
@ -307,8 +256,7 @@ with sync_playwright() as p:
composer = page.locator('textarea[aria-label="Message input"]')
composer.wait_for(state = "visible", timeout = 60_000)
# Detect chat-only mode: /api/health.chat_only is the source of truth.
# In chat-only mode, /studio + /export redirect to /chat.
# Detect chat-only mode (/api/health.chat_only): in chat-only mode /studio + /export redirect to /chat.
health_resp = evaluate_fetch(
page,
f"{BASE}/api/health",
@ -325,7 +273,7 @@ with sync_playwright() as p:
# 1. Compare tab.
# ─────────────────────────────────────────────────────
step("Compare tab: send to two panes")
# Compare moved into the composer + menu (Tools and attachments).
# Compare lives in the composer "Tools and attachments" menu.
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:
@ -333,7 +281,7 @@ with sync_playwright() as p:
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.
# Fallback: Compare chat may be under the "More" submenu.
more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first
if more_trigger.count() > 0:
more_trigger.hover()
@ -355,43 +303,29 @@ with sync_playwright() as p:
else:
page.wait_for_timeout(1500)
shoot("02-compare-opened")
# Compare view's container.
view = page.locator('[data-tour="chat-compare-view"]').first
if view.count() == 0:
soft_fail("[data-tour='chat-compare-view'] not found after Compare click")
else:
ok_count_before = len(page.locator('[data-role="assistant"]').all())
# Send first prompt; the shared composer placeholder is
# "Send to both models...". Just type into the composer
# textarea (assistant-ui exposes one in compare-mode too).
# Composer placeholder in compare-mode is "Send to both models...".
cmp_composer = page.get_by_placeholder(
re.compile(r"Send to both models", re.I),
).first
if cmp_composer.count() == 0:
# Fall back to any visible textarea inside the compare
# view.
# Fall back to any textarea inside the compare view.
cmp_composer = view.locator("textarea").first
if cmp_composer.count() == 0:
soft_fail("compare composer textarea not found")
else:
cmp_composer.click()
cmp_composer.fill("Reply with: A")
# Prefer Enter on the textarea: the shared composer's
# onKeyDown handler maps plain Enter to send(). The
# send button is rendered via TooltipIconButton +
# ComposerPrimitive.Send and its aria-label was
# added late, so older builds match nothing for
# button[aria-label="Send message"] in compare mode.
# Prefer Enter: onKeyDown maps plain Enter to send(); the Send button's
# aria-label was added late so older builds don't match it in compare mode.
cmp_composer.press("Enter")
# Wait for at least 2 NEW assistant bubbles (one per
# pane). NOTE: the Compare view requires per-pane
# model selection to actually generate. In this CI
# flow the panes are NOT explicitly assigned -- so
# the backend rejects the request as "At least one
# non-system message is required" or similar. We
# downgrade this to runtime_warn (informational) and
# keep the structural assertions (view present,
# composer present, message text round-trips) above.
# Expect 2 new assistant bubbles (one per pane). Panes have no explicit model
# in this CI flow so the backend may reject; downgrade to runtime_warn while
# keeping the structural assertions (view/composer present, text round-trips).
try:
page.wait_for_function(
"""(want) => {
@ -410,9 +344,7 @@ with sync_playwright() as p:
)
shoot("03-compare-after-A")
# Send a second prompt -> 4 total new bubbles. Same
# caveat: this is runtime-flaky when panes have no
# explicit model selection.
# Second prompt -> 4 total new bubbles (same runtime-flaky caveat).
cmp_composer.fill("Reply with: B")
cmp_composer.press("Enter")
try:
@ -452,19 +384,17 @@ with sync_playwright() as p:
if n_templates == 0:
soft_fail("no recipe template cards found")
else:
# Click the first one.
try:
templates.first.scroll_into_view_if_needed()
templates.first.click()
page.wait_for_timeout(2000)
shoot("06-recipe-opened")
# The recipe-studio canvas uses React-Flow; look for the
# renderer.
# The recipe-studio canvas uses React-Flow; look for the renderer.
canvas = page.locator(
".react-flow__renderer, .react-flow, [data-testid*='react-flow']"
).first
if canvas.count() == 0:
# Some templates may open as dialogs instead of route.
# Some templates open as dialogs instead of a route.
info("(no React-Flow canvas; template may have opened a dialog)")
else:
info("OK React-Flow canvas mounted")
@ -490,13 +420,8 @@ with sync_playwright() as p:
soft_fail("[data-tour='export-cta'] not found in /export")
else:
info("OK [data-tour='export-cta'] visible")
# The Export page's HF-token field is lazy-loaded behind a
# disclosure, and on slow runners (macos-14 free) it can
# dawdle. Poll across multiple selectors for up to 8 s before
# giving up. We log this as info (not soft_fail) because it
# does not block any user-visible export workflow -- the user
# who needs to push to HF can scroll and the section will load
# within a few seconds.
# HF-token field is lazy-loaded behind a disclosure; poll multiple selectors for ~8s.
# Logged as info (not soft_fail) since it doesn't block the export workflow.
hf_token = None
for _try in range(8):
page.wait_for_timeout(1000)
@ -553,11 +478,11 @@ with sync_playwright() as p:
step("Settings dialog: cycle through tabs")
page.goto(f"{BASE}/chat")
composer.wait_for(state = "visible", timeout = 60_000)
page.keyboard.press("Control+,") # global shortcut
page.keyboard.press("Control+,")
page.wait_for_timeout(800)
settings = page.get_by_role("dialog").first
if settings.count() == 0:
# macOS shortcut is Cmd-,; try that too.
# macOS shortcut is Cmd-,.
page.keyboard.press("Meta+,")
page.wait_for_timeout(800)
settings = page.get_by_role("dialog").first
@ -565,8 +490,7 @@ with sync_playwright() as p:
soft_fail("Settings dialog didn't open with Cmd/Ctrl-,")
else:
shoot("09-settings-open")
# Each tab is a button with the visible text as accessible name.
# Tabs available depend on chat_only mode.
# Each tab is a button named by its visible text; availability depends on chat_only mode.
candidate_tabs = (
"General",
"Profile",
@ -586,7 +510,7 @@ with sync_playwright() as p:
try:
btn.click()
page.wait_for_timeout(400)
# Tab body must contain something (non-empty).
# Tab body must be non-empty.
body_text = page.evaluate(
"""() => {
const dialog = document.querySelector('[role="dialog"]');

133
tests/studio/run_real_mlx_smoke.py
View file

@ -1,27 +1,12 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""
End-to-end MLX smoke test on real Apple Silicon -- multi-process driver.
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
`train` loads gemma-3-270m-it, applies LoRA, probes pre/post loss+grad,
overfits one repeated row for 30 deterministic steps (batch 2, accum 3),
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.
GGUF export and LoRA reload fixes land in unslothai/unsloth-zoo#627.
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.
"""End-to-end MLX smoke test on real Apple Silicon (multi-process driver).
`train` overfits gemma-3-270m-it on one row for 30 steps and saves
lora/merged_16bit/gguf; `reload` reopens each format in a fresh process.
GGUF + LoRA reload fixes land in unslothai/unsloth-zoo#627. Metal's
reduction-order nondeterminism makes loss assertions bounds, not exact.
Apple-Silicon only; invoked from .github/workflows/mlx-ci.yml.
"""
@ -66,7 +51,7 @@ def _peak_gpu_gb() -> float:
if not mx.metal.is_available():
return 0.0
# Newer MLX moved get_peak_memory to top-level; fall back to mx.metal for old versions.
# Newer MLX moved get_peak_memory to top-level; fall back to mx.metal.
getter = getattr(mx, "get_peak_memory", None) or getattr(mx.metal, "get_peak_memory", None)
if getter is None:
return 0.0
@ -121,8 +106,7 @@ def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, flo
import mlx.nn as nn
from mlx.utils import tree_flatten
# Match Studio's text dataset path: Studio passes exactly the formatted
# text to the tokenizer and does not append EOS behind the user's back.
# Match Studio's text dataset path: no EOS appended behind the user's back.
ids = list(tokenizer.encode(text))
if len(ids) < 2:
raise RuntimeError(f"text too short to compute loss: {len(ids)} tokens")
@ -145,14 +129,10 @@ def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, flo
def _teacher_forced_completion_loss(model, tokenizer, prompt: str, completion: str) -> float:
"""Mean next-token CE on `completion` given `prompt`, teacher-forced.
"""Mean teacher-forced next-token CE on `completion` given `prompt`.
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
Decouples the memorisation check from flaky greedy-decode geometry:
asserts *what* the model memorised, not just that loss is low.
Returns mean cross-entropy over the completion's tokens.
"""
import mlx.core as mx
import mlx.nn as nn
@ -169,8 +149,7 @@ def _teacher_forced_completion_loss(model, tokenizer, prompt: str, completion: s
targets = mx.array([full_ids[1:]], dtype = mx.int32)
logits = model(inputs)
# logits at position i predict targets[i]; completion tokens occupy
# target positions [len(prompt_ids)-1 ... len(full_ids)-2].
# logits at position i predict targets[i]; completion starts at len(prompt_ids)-1.
start = len(prompt_ids) - 1
completion_logits = logits[:, start:, :]
completion_targets = targets[:, start:]
@ -224,9 +203,8 @@ 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). 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.
# Full q/k/v/o + gate/up/down set: q/k/v/o alone couldn't memorize
# the row, the MLP projections add the needed capacity.
model = FastMLXModel.get_peft_model(
model,
r = 8,
@ -259,18 +237,15 @@ def cmd_train(args) -> int:
config = MLXTrainingConfig(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 3,
# 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.
# PR #5498 sweep: 7 steps too few; 30 makes every seed converge.
max_steps = 30,
learning_rate = 1e-3,
warmup_steps = 0,
lr_scheduler_type = "constant",
optim = "adamw",
weight_decay = 0.0,
# Pin the elementwise clip to match the 13-seed-tested fixture
# (value=1.0 62% pass, norm=1.0 46%). Zoo's new MLX default is
# max_grad_leaf_norm=1.0; explicit value wins, norm disabled.
# Pin the elementwise clip (value=1.0, norm disabled) to match the
# 13-seed-tested fixture; explicit value overrides zoo's MLX default.
max_grad_norm = 0.0,
max_grad_value = 1.0,
logging_steps = 1,
@ -326,8 +301,7 @@ def cmd_train(args) -> int:
)
if k in train_result
}
# logging_steps=1 + max_steps=N -> N callbacks; track config so the
# gate auto-follows if max_steps is bumped again.
# logging_steps=1 + max_steps=N -> N callbacks; gate auto-follows max_steps.
expected_logged_steps = int(config.max_steps)
assert (
len(losses_per_step) == expected_logged_steps
@ -337,10 +311,8 @@ def cmd_train(args) -> int:
f"expected train_steps={expected_logged_steps}, got " f"{train_result['train_steps']}"
)
for i, l in enumerate(losses_per_step):
# Allow exact 0.0: fp16 per-step loss underflows to 0.0 after
# the LoRA reaches loss=0 around step ~10 with this fixture +
# max_steps=30. That's the memorization success signal, not a
# bug. Lower bound is "finite and >= 0" not "strictly > 0".
# Allow exact 0.0: fp16 loss underflows once the LoRA memorises the
# row (~step 10); that's success, so the lower bound is >= 0 not > 0.
assert math.isfinite(l) and 0 <= l < 50, f"step {i+1} loss bad: {l}"
assert (
losses_per_step[-1] < losses_per_step[0] * 1.1
@ -351,13 +323,8 @@ def cmd_train(args) -> int:
metrics["post_train_loss"] = round(post_loss, 4)
metrics["post_train_grad_norm"] = round(post_norm, 4)
assert post_loss < pre_loss, f"post {post_loss} >= pre {pre_loss}"
# Memorisation gate: teacher-forced loss on the training row must
# be very low after 30 steps of overfit-on-one-example. This is
# the robust signal that the model learned the trained
# continuation, regardless of MLX's autoregressive-generation
# numerics. Empirical 47-round, 13-seed sweep: every (clip, bc,
# seed) configuration that converges hits post_train_loss <= 0.05.
# Tighten gate to 0.1.
# Memorisation gate: every converging (clip, bc, seed) config in the
# 13-seed sweep hit post_train_loss <= 0.05, so 0.1 is a robust bound.
assert post_loss < 0.1, (
f"post_train_loss={post_loss:.4f} >= 0.1 -- training did not "
"memorise the single training row in 30 steps. Trainer "
@ -376,12 +343,8 @@ def cmd_train(args) -> int:
verbose = False,
)
metrics["in_memory_generation"] = in_mem_out
# Soft greedy-decode visibility (metric only). Empirically this lands in
# 46-77% of seeds depending on clip config (47-round, 13-seed sweep) --
# fp16 + MLX attention/generate path puts noticeable noise on the first
# token even after near-zero teacher-forced loss. Surface the mismatch
# for regression tracking, but the next assertion is the load-bearing
# one.
# Soft greedy-decode metric only (46-77% of seeds): fp16 + MLX generate
# noises the first token. The teacher-forced check below is load-bearing.
metrics["in_memory_generation_has_expected"] = EXPECT_IN_OUTPUT in in_mem_out
if EXPECT_IN_OUTPUT not in in_mem_out:
print(
@ -391,12 +354,9 @@ def cmd_train(args) -> int:
flush = True,
)
# Hard check: teacher-forced loss on the completion the model was trained
# to emit. Bypasses greedy-decode fp16 fragility -- if the LoRA actually
# memorised the row, the probability mass on `EXPECT_IN_OUTPUT` after
# `PROMPT` is essentially 1.0 (and the loss essentially 0). 13/13 of the
# MLX configs we measured reached post_train_loss < 1e-3, so this gate
# is deterministic on every (seed, clip, bc) combination tested.
# Hard check: teacher-forced loss on the trained completion bypasses
# greedy-decode fp16 fragility. 13/13 measured configs reached < 1e-3,
# so this gate is deterministic across (seed, clip, bc).
completion_loss = _teacher_forced_completion_loss(
model, tokenizer, PROMPT, EXPECT_IN_OUTPUT + "!"
)
@ -409,8 +369,8 @@ def cmd_train(args) -> int:
"optimizer defaults vs torch.optim.AdamW."
)
# Save LoRA. unsloth-zoo#627 fixed FastMLXModel.from_pretrained(lora_dir)
# so the cold-start reload below works on the saved adapter dir directly.
# unsloth-zoo#627 fixed from_pretrained(lora_dir) so the cold-start
# reload below works on the saved adapter dir directly.
lora_dir = workdir / "lora"
with Phase("save_lora", metrics):
model.save_pretrained_merged(
@ -433,15 +393,10 @@ def cmd_train(args) -> int:
metrics["merged_dir"] = str(merged_dir)
assert any(merged_dir.glob("*.safetensors"))
# Save GGUF (best-effort). save_pretrained_gguf clones llama.cpp,
# builds it with cmake (Metal=ON), then runs convert_hf_to_gguf.
# For some models -- including unsloth/gemma-3-270m-it as of
# 2026-05-07 -- llama.cpp's converter asserts on the tokenizer vocab
# (`assert max(tokenizer.vocab.values()) < vocab_size`) because the
# tokenizer carries reserved IDs beyond the embedding matrix size.
# That's an llama.cpp / convert_hf_to_gguf limitation, not an
# unsloth_zoo bug. Soft-skip with a recorded reason so the LoRA +
# merged_16bit assertions still gate the PR.
# Save GGUF (best-effort). For some models (e.g. gemma-3-270m-it)
# llama.cpp's convert_hf_to_gguf asserts on the tokenizer vocab -- an
# llama.cpp limitation, not an unsloth_zoo bug. Soft-skip with a recorded
# reason so the LoRA + merged_16bit assertions still gate the PR.
gguf_dir = workdir / "gguf"
metrics["gguf_supported"] = False
metrics["gguf_skip_reason"] = None
@ -523,14 +478,9 @@ def cmd_reload(args) -> int:
metrics["generation"] = out
print(f" [reload:{args.format}] output: {out!r}", flush = True)
# Verify save/reload preserved the trained weights via teacher-
# forced loss on the training row: the reloaded model should have
# approximately the same loss on TRAIN_TEXT as the in-memory model
# had at post_train_loss. This is the real save/reload invariant
# and is robust to MLX's known near-zero-loss adamw greedy-decode
# perturbation (step-7 grad spike at seed=3407, see
# scripts/cuda_mlx_step7_*) which can flip the first generated
# token while leaving teacher-forced loss essentially identical.
# Save/reload invariant: reloaded teacher-forced loss on TRAIN_TEXT must
# match the in-memory post_train_loss. Robust to MLX's greedy-decode
# perturbation, which can flip the first token but not the loss.
train_metrics_path = save_dir.parent / "train_metrics.json"
in_mem_loss = None
in_mem_out = None
@ -547,15 +497,13 @@ def cmd_reload(args) -> int:
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)
# float16 round-trip should be near-exact for LoRA + merged;
# 0.2 tolerates the dequant noise we have seen empirically.
# float16 round-trip is near-exact; 0.2 tolerates dequant noise.
assert abs(reload_loss - float(in_mem_loss)) < 0.2, (
f"reload {args.format!r} loss diverged from in-memory: "
f"reload={reload_loss:.4f}, in-memory={in_mem_loss:.4f}"
)
else:
# Fallback when train_metrics.json wasn't found (older
# workdir layouts): keep a non-empty-completion gate.
# Fallback when train_metrics.json is missing: gate on non-empty output.
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}"
@ -610,12 +558,9 @@ 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]}")
# 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
# output to catch real save/reload corruption (zero-weight model,
# tokenizer mismatch). Surface whether EXPECT_IN_OUTPUT appears in
# the metrics for visibility without gating on it.
# llama.cpp tokenises/samples differently than mlx_lm, so the GGUF
# completion needn't match. Require non-empty output to catch real
# save/reload corruption; record EXPECT_IN_OUTPUT without gating on it.
body = (proc.stdout or "").replace(PROMPT, "", 1).strip()
metrics["gguf_has_expected"] = EXPECT_IN_OUTPUT in (proc.stdout or "")
assert len(body) >= 4, (

134
tests/studio/studio_api_smoke.py
View file

@ -1,32 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""End-to-end Studio API & Auth tests.
Boots a fresh Studio externally (CI workflow handles install + boot)
and runs a battery of HTTP-level integration tests against it. No
Playwright, no model load by this test (the workflow loads gemma-3-270m
beforehand if needed).
Sections:
1. CORS hardening (no wildcard + credentials, no bootstrap leak)
2. /api/system + /api/system/hardware require auth
3. Auth state machine (rotation invariants, body validation, login burst)
4. JWT-expiry rejection (forge an expired token using the install's secret)
5. API key lifecycle E2E (create -> list -> use -> delete -> reject)
6. Auth file-mode hardening (Linux only)
7. Inference lifecycle gaps (force reload, bogus variant, /v1/models,
/v1/embeddings, /v1/responses)
8. Endpoint-by-endpoint auth audit (pin EXPECTED auth posture per route)
Env:
BASE_URL http://127.0.0.1:18893 (or wherever Studio is)
STUDIO_OLD_PW the bootstrap password (must rotate it)
STUDIO_NEW_PW what to rotate to
STUDIO_NEW2_PW out-of-band rotation target
STUDIO_AUTH_DIR (optional) path to the auth dir for file-mode checks
GGUF_REPO (optional) the model the workflow loaded for /v1 tests
"""
"""End-to-end Studio API & Auth HTTP integration tests against an externally-booted Studio."""
import json
import os
@ -50,10 +25,7 @@ _section = [0]
_failed: list[str] = []
_warned: list[str] = []
# When 1, audit-finding assertions (e.g. CORS leak, file modes, 5xx vs
# 4xx) become hard fails. Off by default: we surface them as WARN so the
# test can be added before the underlying Studio fixes ship; the
# warnings are still printed in CI so they're visible.
# When 1, audit-finding assertions become hard fails. Off by default: surfaced as WARN.
STRICT_AUDIT = os.environ.get("STUDIO_API_STRICT_AUDIT", "0") == "1"
@ -63,15 +35,7 @@ def section(title: str) -> None:
def _shape(value):
"""Return a credential-free shape descriptor for an HTTP body.
Returns ONLY the container type + element count -- never the keys,
never the values. Used in failure messages so a CI log can never
carry credential material (matches the intent of CodeQL's
py/clear-text-logging-sensitive-data rule). For richer detail
while debugging, set STUDIO_API_VERBOSE=1 locally; verbose mode
is OFF in CI.
"""
"""Credential-free shape descriptor (container type + count only) for an HTTP body."""
if isinstance(value, dict):
return f"<dict with {len(value)} keys>"
if isinstance(value, list):
@ -82,18 +46,7 @@ def _shape(value):
def _emit(prefix: str, msg: str) -> None:
"""Write a status line via os.write.
CodeQL's py/clear-text-logging-sensitive-data rule treats `print`
(and the standard `logging` calls) as logging sinks. Even though
`_shape()` already strips credential material from anything
`msg` could carry, the rule's data-flow can't see through the
helper and flags `print(msg)` as clear-text logging. Routing
through a raw fd write keeps the same observable CI output
while not matching the rule's sink pattern. The msg payload is
still credential-free by construction (callers wrap response
bodies in `_shape(...)`).
"""
"""Write a status line via raw os.write to dodge CodeQL's clear-text-logging sink on print()."""
os.write(1, prefix.encode("utf-8"))
os.write(1, msg.encode("utf-8", errors = "replace"))
os.write(1, b"\n")
@ -104,21 +57,13 @@ def ok(msg: str) -> None:
def fail(msg: str) -> None:
"""Record a failure but keep running so we report ALL failures.
`msg` must be free of credential material -- callers should pass
only the HTTP status code + a short description (and `_shape(body)`
if shape is informative). Never `body` directly.
"""
"""Record a failure but keep running so we report ALL failures. `msg` must be credential-free."""
_emit(" FAIL ", msg)
_failed.append(f"{_section[0]}: {msg}")
def audit(msg: str) -> None:
"""Record an audit finding -- a real backend regression that we
want surfaced in CI logs but not gating until the underlying fix
ships. Set STUDIO_API_STRICT_AUDIT=1 to escalate to hard fail.
"""
"""Record a non-gating backend regression; STUDIO_API_STRICT_AUDIT=1 escalates to hard fail."""
if STRICT_AUDIT:
fail(msg)
else:
@ -173,11 +118,7 @@ def login(password: str) -> tuple[int, str | None]:
# ─────────────────────────────────────────────────────────────────────────
section("CORS hardening")
# Cross-origin OPTIONS preflight. FastAPI explicitly forbids
# Access-Control-Allow-Origin: <origin> together with
# Access-Control-Allow-Credentials: true. (Wildcard + credentials is
# also forbidden by the browser.) Either response is acceptable; the
# bad pattern is a wildcard origin echoed alongside credentials.
# Cross-origin OPTIONS preflight. Bad pattern: wildcard origin echoed alongside credentials=true.
req = urllib.request.Request(
f"{BASE}/api/auth/login",
method = "OPTIONS",
@ -198,10 +139,7 @@ try:
except Exception as exc:
ok(f"CORS preflight unreachable (acceptable): {exc!r}")
# GET / from a cross-origin Origin header. The response body must NOT
# contain the literal bootstrap password (the security audit flagged
# that __UNSLOTH_BOOTSTRAP__ injection in the served HTML can be
# fetched cross-origin under wildcard CORS).
# GET / cross-origin must NOT leak the bootstrap password in the served HTML.
boot_path = AUTH_DIR / ".bootstrap_password"
if boot_path.exists():
bootstrap_pw = boot_path.read_text().strip()
@ -214,11 +152,7 @@ if boot_path.exists():
with urllib.request.urlopen(req, timeout = 10) as r:
body = r.read().decode("utf-8", errors = "ignore")
if bootstrap_pw in body:
# AUDIT finding (P0 from security review): the
# __UNSLOTH_BOOTSTRAP__ injection in served HTML is
# readable cross-origin under the current wildcard
# CORS policy. Tracked separately; the test surfaces
# the regression but does not gate CI on it.
# AUDIT (P0): bootstrap pw in served HTML is readable cross-origin under wildcard CORS.
audit("CORS: GET / leaks bootstrap pw to cross-origin caller")
else:
ok("CORS: GET / does not include bootstrap pw")
@ -242,8 +176,7 @@ for endpoint in ("/api/system", "/api/system/hardware", "/api/system/gpu-visibil
fail(f"GET {endpoint} unauthenticated returned {code} (expected 401/403)")
# Rotate password to NEW so we have a working bearer for the rest.
# (Bootstrap login -> change-password -> login with NEW.)
# Rotate password to NEW for a working bearer: bootstrap login -> change-password -> login NEW.
section("Rotate bootstrap password for downstream tests")
code, old_token = login(OLD)
if code != 200 or not old_token:
@ -267,7 +200,7 @@ if code != 200 or not NEW_TOKEN:
ok("login with NEW -> 200")
AUTH_HEADER = {"Authorization": f"Bearer {NEW_TOKEN}"}
# Re-test /api/system endpoints WITH auth: must succeed now.
# Re-test /api/system endpoints WITH auth.
for endpoint in ("/api/system", "/api/system/hardware", "/api/system/gpu-visibility"):
code, _ = http("GET", endpoint, headers = AUTH_HEADER)
if code == 200:
@ -275,7 +208,7 @@ for endpoint in ("/api/system", "/api/system/hardware", "/api/system/gpu-visibil
else:
fail(f"GET {endpoint} authenticated returned {code} (expected 200)")
# Load the model. Sections 5 + 7 below need a loaded model.
# Sections 5 + 7 below need a loaded model.
section("Load the GGUF for /v1 tests")
code, body = http(
"POST",
@ -315,9 +248,8 @@ else:
fail(f"/api/auth/refresh without body returned {code} (expected 400/422)")
# Wrong-password burst: expect 401 until the per-IP bucket fills, then
# 429 with Retry-After. Bucket cannot be reset between tests, so we
# assert the observable invariant rather than a fixed transition index.
# Wrong-password burst: 401 until the per-IP bucket fills, then 429 with Retry-After.
# Bucket can't be reset between tests, so assert the invariant, not a fixed transition index.
def _login_with_headers(password: str) -> tuple[int, str | None]:
"""Like ``login`` but returns ``(status, retry_after_header)``."""
url = f"{BASE}/api/auth/login"
@ -362,7 +294,7 @@ else:
# ─────────────────────────────────────────────────────────────────────────
section("JWT expiry")
# Forge a JWT with exp=now-1 using the install's signing secret.
# auth/storage.py:get_user_and_secret('unsloth') returns (salt, hash, jwt_secret, must_change_pw).
# get_user_and_secret('unsloth') returns (salt, hash, jwt_secret, must_change_pw).
try:
sys.path.insert(
0,
@ -419,9 +351,7 @@ code, body = http(
if code != 200 or not isinstance(body, dict):
fail(f"POST /api/auth/api-keys -> {code}: {_shape(body)}")
else:
# Response shape: {"key": "sk-unsloth-...", "api_key": {"id": ...,
# "name": ..., "key_prefix": ..., ...}}. The flat "key" carries the
# one-time bearer; the "api_key" sub-dict carries the metadata.
# Flat "key" is the one-time bearer; the "api_key" sub-dict carries metadata.
api_key = body.get("key")
api_meta = body.get("api_key") if isinstance(body.get("api_key"), dict) else {}
api_id = api_meta.get("id") or body.get("id")
@ -429,8 +359,6 @@ else:
fail(f"create-key missing key/id: {_shape(body)}")
else:
ok(f"created key id={api_id}")
# The API key may use sk-unsloth-* or another prefix; we don't
# pin the literal prefix.
# List must include this id.
code, body = http("GET", "/api/auth/api-keys", headers = AUTH_HEADER)
if code == 200 and isinstance(body, dict):
@ -442,8 +370,7 @@ else:
else:
fail(f"GET /api/auth/api-keys -> {code}: {_shape(body)}")
# Use the key against /v1/chat/completions (the workflow has
# already loaded gemma-3-270m).
# Use the key against /v1/chat/completions.
code, body = http(
"POST",
"/v1/chat/completions",
@ -512,10 +439,7 @@ else:
if actual_mode == expected_mode:
ok(f"{path} mode={oct(actual_mode)}")
else:
# AUDIT finding (P1 from security review): auth.db inherits
# the process umask (0o644 on most CI runners) instead of
# being chmod 0o600 like the bootstrap pw file. Tracked
# separately; surface, don't gate.
# AUDIT (P1): auth.db inherits the umask instead of being chmod 0o600.
audit(f"{path} mode={oct(actual_mode)} (expected {oct(expected_mode)})")
@ -535,8 +459,7 @@ if code == 200 and isinstance(body, dict):
else:
fail(f"/v1/models -> {code}: {_shape(body)}")
# /v1/embeddings either returns embedding OR a structured 4xx/5xx.
# 501 "Not Implemented" is acceptable for non-embedding-capable models.
# /v1/embeddings returns an embedding OR a structured 4xx (501 OK for non-embedding models).
code, body = http(
"POST",
"/v1/embeddings",
@ -568,10 +491,7 @@ if code == 200 or 400 <= code < 500:
else:
fail(f"/v1/responses -> {code} (expected 200 or 4xx)")
# Bogus variant must be rejected. The contract: 4xx for an obviously
# bad input is the right code. Today the backend returns 500 for
# unknown variants -- rejected, but with the wrong status. Surface as
# AUDIT (not gating) until the variant validator returns 4xx.
# Bogus variant must be rejected with 4xx. Backend currently 500s; surface as AUDIT until fixed.
code, _ = http(
"POST",
"/api/inference/load",
@ -593,7 +513,6 @@ else:
# Force-reload of the same repo: child PID must change.
# Read the inference status before.
def _llama_pid() -> int | None:
code, body = http("GET", "/api/inference/status", headers = AUTH_HEADER)
if code != 200 or not isinstance(body, dict):
@ -629,9 +548,7 @@ else:
# 8. Endpoint-by-endpoint auth audit
# ─────────────────────────────────────────────────────────────────────────
section("Endpoint auth audit")
# Pin the EXPECTED auth posture for known routes. A new route added
# without an entry here fails the audit, forcing the author to make
# the auth decision explicit.
# Pin the EXPECTED auth posture per route; a new unlisted route fails the audit.
PUBLIC = {
("GET", "/api/health"),
("GET", "/api/auth/status"),
@ -654,10 +571,7 @@ EXPECTED_AUTH_ENDPOINTS = [
for method, path in EXPECTED_AUTH_ENDPOINTS:
if (method, path) in PUBLIC:
continue
# Don't actually shut Studio down -- verify auth check by sending
# an empty body / no auth header. If the check happens BEFORE the
# shutdown trigger (which is the design), we get a 401/403 without
# any side effects.
# Don't actually shut Studio down: an unauthenticated call must 401/403 before the trigger fires.
if path == "/api/shutdown":
code, _ = http(method, path)
if code in (401, 403):
@ -672,9 +586,7 @@ for method, path in EXPECTED_AUTH_ENDPOINTS:
fail(f"{method} {path} unauthenticated returned {code} (expected 401/403)")
for method, path in PUBLIC:
code, _ = http(method, path)
if (
200 <= code < 500
): # public endpoints either 200 or 4xx (bad input), never connection-refused
if 200 <= code < 500: # public endpoints: 200 or 4xx, never connection-refused
ok(f"{method} {path} public -> {code}")
else:
fail(f"{method} {path} public returned unexpected {code}")

56
tests/studio/test_auth_form_input_count.py
View file

@ -3,16 +3,10 @@
"""Pin the auth-form input-count contract on the change-password page.
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 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.
"""
PR #5490 added a third "Current password" input, regressing first-boot UX to
three inputs; PR #5545 restores two by rendering it only when BOOTSTRAP is absent.
These tests inspect the source directly (no Studio/browser/network); runtime is
covered by tests/studio/playwright_chat_ui.py."""
from __future__ import annotations
@ -28,9 +22,7 @@ CONDITIONAL_OPENER = "{!hasBootstrapPassword && ("
def _conditional_extent(src: str) -> tuple[int, int]:
"""Return the (start, end) char offsets of the
`{!hasBootstrapPassword && (...)}` JSX block. ``start`` points
at the opening `{`; ``end`` points one past the matching `)}`."""
"""(start, end) char offsets of the `{!hasBootstrapPassword && (...)}` JSX block."""
start = src.find(CONDITIONAL_OPENER)
assert start != -1, (
"the {!hasBootstrapPassword && (...)} JSX block that hides the "
@ -52,10 +44,8 @@ def _conditional_extent(src: str) -> tuple[int, int]:
def test_hasbootstrappassword_constant_is_derived_from_bootstrap_window_value():
"""The conditional guard must read from window.__UNSLOTH_BOOTSTRAP__.
A future refactor that swaps the source (e.g. a localStorage flag,
a prop) would silently drift from the backend's bootstrap-injection
contract in studio/backend/main.py::_inject_bootstrap."""
"""The guard must read from window.__UNSLOTH_BOOTSTRAP__, matching 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, (
"hasBootstrapPassword constant missing or its derivation drifted; "
@ -64,10 +54,8 @@ def test_hasbootstrappassword_constant_is_derived_from_bootstrap_window_value():
def test_exactly_one_hasBootstrapPassword_conditional_exists():
"""Only one `!hasBootstrapPassword` JSX check is allowed. A second
one would split the form rendering into branches that the rest of
these structural tests cannot reason about, and would almost
certainly hide or duplicate one of the New / Confirm inputs."""
"""Only one `!hasBootstrapPassword` JSX check is allowed; a second would split
rendering into branches and likely hide or duplicate the New / Confirm inputs."""
src = AUTH_FORM.read_text()
count = src.count("!hasBootstrapPassword")
assert count == 1, (
@ -77,9 +65,8 @@ def test_exactly_one_hasBootstrapPassword_conditional_exists():
def test_current_password_input_is_inside_the_hasBootstrapPassword_conditional():
"""`id="current-password"` MUST sit inside `{!hasBootstrapPassword && (...)}`.
Otherwise the input renders on first boot too, regressing the
pre-#5490 two-input UX that PR #5545 restores."""
"""`id="current-password"` must sit inside `{!hasBootstrapPassword && (...)}`,
else it renders on first boot too, regressing the pre-#5490 UX that PR #5545 restores."""
src = AUTH_FORM.read_text()
s, e = _conditional_extent(src)
idx = src.find('id="current-password"')
@ -92,8 +79,8 @@ def test_current_password_input_is_inside_the_hasBootstrapPassword_conditional()
def test_new_password_input_is_outside_the_hasBootstrapPassword_conditional():
"""`id="new-password"` MUST sit outside `{!hasBootstrapPassword && (...)}`.
Otherwise it disappears on admin-forced resets, regressing PR #5490."""
"""`id="new-password"` must sit outside `{!hasBootstrapPassword && (...)}`,
else it disappears on admin-forced resets, regressing PR #5490."""
src = AUTH_FORM.read_text()
s, e = _conditional_extent(src)
idx = src.find('id="new-password"')
@ -119,11 +106,9 @@ def test_confirm_password_input_is_outside_the_hasBootstrapPassword_conditional(
def test_change_password_jsx_declares_exactly_three_password_inputs():
"""The change-password JSX block (`{!isLoginMode && (...)}`) must
declare exactly the three known password inputs -- current, new,
confirm. A fourth would almost certainly break the 2-input
first-boot contract because the conditional only hides the
Current input, not any new one a future PR might add."""
"""The change-password JSX block (`{!isLoginMode && (...)}`) must declare exactly
current/new/confirm; a fourth would break the 2-input first-boot contract (the
conditional only hides Current)."""
src = AUTH_FORM.read_text()
start = src.find("{!isLoginMode && (")
assert start != -1, (
@ -155,10 +140,8 @@ def test_change_password_jsx_declares_exactly_three_password_inputs():
def test_login_jsx_declares_exactly_one_password_input():
"""The login JSX block (`isLoginMode && (...)`) must declare
exactly one password input -- the bootstrap password the user
pastes from the CLI. Adding a second here would break the
matrix that the per-mode tests assume."""
"""The login JSX block (`isLoginMode && (...)`) must declare exactly one password
input (the bootstrap password pasted from the CLI); a second breaks the per-mode matrix."""
src = AUTH_FORM.read_text()
start = src.find("{isLoginMode && (")
assert start != -1, "the login JSX subtree marker is missing"
@ -173,8 +156,7 @@ def test_login_jsx_declares_exactly_one_password_input():
i += 1
subtree = src[start:i]
ids = re.findall(r'id="([a-z-]+)"', subtree)
# The login subtree currently uses id="password". Lock the count
# rather than the spelling so a rename does not falsely fail.
# Lock the count, not the spelling, so a rename does not falsely fail.
pw_ids = [x for x in ids if "password" in x]
assert len(pw_ids) == 1, (
f"login JSX must declare exactly one password-typed input; " f"found {pw_ids!r}"

10
tests/studio/test_cancel_atomicity.py
View file

@ -1,12 +1,4 @@
"""
TOCTOU atomicity guards for the cancel path.
Structural: cancel_inference, _cancel_by_cancel_id_or_stash, and
_TrackedCancel.__enter__ must each use a single _CANCEL_LOCK critical
section over lookup + stash / register + consume-pending.
Behavioral: parallel cancel-POST vs __enter__ must never drop a cancel.
"""
"""TOCTOU atomicity guards for the cancel path: single _CANCEL_LOCK critical sections; parallel cancel-POST vs __enter__ never drops a cancel."""
from __future__ import annotations

32
tests/studio/test_cancel_id_wiring.py
View file

@ -1,19 +1,9 @@
"""
Wiring tests for the per-run cancel_id field.
"""Wiring tests for the per-run cancel_id field.
A chat-thread-scoped session_id is not safe as a cancel key because a
late stop POST can match a subsequent run on the same thread. The fix
adds cancel_id (a fresh UUID per generation) that is sent both in the
completion payload and in the /api/inference/cancel body.
Verifies:
- ChatCompletionRequest exposes an Optional[str] `cancel_id` field.
- /api/inference/cancel accepts `cancel_id` as the first-preferred key.
- OpenAIChatCompletionsRequest (frontend type) includes cancel_id.
- chat-adapter.ts generates a per-run cancelId (crypto.randomUUID
with a Math.random fallback), sends it in the completion payload,
and includes it in the /inference/cancel body on abort.
"""
A thread-scoped session_id is unsafe as a cancel key (a late stop POST can match
a later run on the same thread); cancel_id is a fresh per-generation UUID sent in
both the completion payload and the /api/inference/cancel body. Verifies the
field on backend/frontend types and the chat-adapter.ts generation + wiring."""
from __future__ import annotations
@ -51,10 +41,9 @@ def test_chat_completion_request_has_cancel_id_field():
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 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).
# A stale POST with cancel_id AND session_id must not cancel a later run via
# session_id; the handler must early-return through an exclusive-cancel_id path
# (atomic helper, or a keys list with ONLY cancel_id).
for node in ast.walk(ast.parse(ROUTES_SRC)):
if isinstance(node, ast.AsyncFunctionDef) and node.name == "cancel_inference":
break
@ -136,9 +125,8 @@ 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, 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.
# authFetch redirects to login on 401, kicking the user out mid-stop if the
# token expired. Use plain fetch + manual Authorization for a best-effort cancel.
start = ADAPTER_SRC.find("const onAbortCancel")
assert start >= 0, "onAbortCancel handler missing"
rest = ADAPTER_SRC[start:]

18
tests/studio/test_cli_repo_variant.py
View file

@ -1,8 +1,6 @@
"""Tests for the ``repo:variant`` shorthand parser used by ``unsloth studio run``.
Loads ``unsloth_cli/commands/studio.py`` directly via ``importlib`` with a
minimal ``typer`` stub so the test doesn't drag in the rest of
``unsloth_cli`` (which transitively imports the unsloth training stack).
Loads studio.py via importlib with a minimal typer stub to avoid importing the unsloth training stack.
"""
from __future__ import annotations
@ -16,12 +14,7 @@ import pytest
def _load_split_repo_variant():
"""Load ``_split_repo_variant`` from studio.py with typer stubbed.
studio.py decorates Typer commands at import time, so a stub that
accepts (and discards) those calls is enough to let module
execution complete and expose the helper we want to test.
"""
"""Load ``_split_repo_variant`` from studio.py with typer stubbed (discards decorator calls)."""
if "typer" not in sys.modules:
typer_stub = types.ModuleType("typer")
@ -121,17 +114,14 @@ def test_empty_string():
def test_trailing_colon_no_variant():
# "org/repo:" -- no quant label after the colon. Pass through
# unchanged so the backend's existing validation surfaces a
# clearer error than "variant ''".
# "org/repo:" has no quant label; pass through unchanged so backend validation gives a clearer error.
repo, variant = _split("org/repo:")
assert repo == "org/repo:"
assert variant is None
def test_slash_in_variant_disqualifies_split():
# "foo:bar/baz" -- the suffix has a slash, so this isn't a quant
# label; treat the whole thing as opaque.
# "foo:bar/baz" suffix has a slash, so it's not a quant label; treat as opaque.
repo, variant = _split("foo:bar/baz")
assert repo == "foo:bar/baz"
assert variant is None

10
tests/studio/test_cli_run_alias.py
View file

@ -1,9 +1,4 @@
"""Tests that ``unsloth run`` is registered as a top-level alias for
``unsloth studio run``.
AST-based to avoid importing ``unsloth_cli`` (which pulls in the heavy
training stack) at test-collection time.
"""
"""AST-based tests that `unsloth run` is registered as a top-level alias for `unsloth studio run`."""
from __future__ import annotations
@ -35,8 +30,7 @@ def test_top_level_run_alias_registered():
and call.func.value.id == "app"
):
continue
# Decorator-call form has a string literal "run" as the first
# positional or as keyword ``name="run"``.
# "run" appears as the first positional arg 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)
is_run = (isinstance(first_pos, ast.Constant) and first_pos.value == "run") or (

19
tests/studio/test_cli_studio_defaults.py
View file

@ -1,8 +1,4 @@
"""Tests that the 'unsloth studio' CLI defaults to 127.0.0.1.
Uses AST parsing to inspect source-level defaults without requiring the
full unsloth_cli dependencies (typer/pydantic) at test-collection time.
"""
"""'unsloth studio' CLI must default --host to 127.0.0.1. AST-based, no typer/pydantic needed."""
import ast
from pathlib import Path
@ -11,12 +7,7 @@ _STUDIO_CMD_PY = Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands
def _find_typer_option_default(source: str, func_name: str, long_option: str):
"""Return the default value of a typer.Option(...) parameter in *func_name*.
Matches by the long option name (e.g. '--host') among the positional args
of the typer.Option() call and returns the first positional arg (the
default value). Only handles ast.Constant defaults.
"""
"""Return the typer.Option default for *long_option* in *func_name* (ast.Constant defaults only)."""
tree = ast.parse(source)
for func_node in ast.walk(tree):
if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
@ -39,7 +30,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; the rest are flags like "--host".
# First positional is the default; the rest are flags.
if not default.args:
continue
flags = [
@ -56,7 +47,7 @@ def _find_typer_option_default(source: str, func_name: str, long_option: str):
def test_studio_default_host_is_loopback():
"""`unsloth studio` (studio_default) --host typer Option default must be 127.0.0.1."""
"""`unsloth studio` (studio_default) --host default must be 127.0.0.1."""
source = _STUDIO_CMD_PY.read_text()
host_default = _find_typer_option_default(source, "studio_default", "--host")
assert (
@ -69,7 +60,7 @@ def test_studio_default_host_is_loopback():
def test_studio_run_host_is_loopback():
"""`unsloth studio run` --host typer Option default must be 127.0.0.1."""
"""`unsloth studio run` --host 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()"

13
tests/studio/test_cli_studio_stop_windows.py
View file

@ -3,12 +3,9 @@
"""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.
`stop` once used `os.kill(pid, 0)`, which raises WinError 87 on Windows before
reaching taskkill; the fix adds cross-platform `_pid_alive` (tasklist on Windows,
signal-0 elsewhere). AST + mock-only; no real processes, no Studio deps imported.
"""
import ast
@ -34,8 +31,8 @@ def _func_source(name: str) -> str:
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."""
"""Exec just `_pid_alive` with injectable sys/subprocess to drive the win32
branch on any host without importing 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

40
tests/studio/test_composer_rtl_bidi_attribute.py
View file

@ -1,11 +1,5 @@
"""Lock down the RTL bidi auto-detection contract on the chat composers.
The browser's Unicode bidi algorithm only flows Arabic / Hebrew / Persian /
Urdu right-to-left when the textarea carries `dir="auto"`. The three
composer surfaces (main chat, inline edit, compare mode) each need the
attribute, and the IME / i18n Playwright smoke must keep its env contract
minimal (no dead `STUDIO_OLD_PW`).
"""
"""RTL bidi contract on chat composers: all three need dir="auto", and the IME
smoke must drop the dead STUDIO_OLD_PW env var."""
from __future__ import annotations
@ -30,10 +24,8 @@ def _block_around(
def test_main_composer_has_dir_auto():
# PR #5784 rewrote the literal attribute into a JSX conditional
# (`aria-label={overlay ? "Image edit instructions" : "Message input"}`),
# so anchor on the inner string literal instead -- it survives both
# the old and new spellings.
# PR #5784 turned the attribute into a JSX conditional; anchor on the inner
# "Message input" literal, which survives both spellings.
block = _block_around(THREAD_TSX.read_text(), '"Message input"')
assert 'dir="auto"' in block, 'main composer is missing dir="auto"'
@ -82,10 +74,8 @@ def test_ime_playwright_script_does_not_read_studio_old_pw():
def test_main_composer_has_stuck_compositionend_watchdog():
"""Issue #5546: Chrome on Windows over WSL never emits compositionend
after the IME commit. The composer keeps a watchdog that releases the
composing flag once events go silent; without it Send stays disabled
forever and CJK input is effectively dropped."""
"""Issue #5546: WSL Chrome never emits compositionend after IME commit, so the
composer needs a watchdog releasing the composing flag or Send stays disabled."""
src = THREAD_TSX.read_text()
assert (
"IME_STUCK_TIMEOUT_MS" in src
@ -105,9 +95,8 @@ def test_compare_composer_has_stuck_compositionend_watchdog():
def test_main_composer_keydown_repins_composing_during_ime():
"""Issue #5546 watchdog can clear composingRef during a long candidate
pause; the IME keydown gate must re-pin it so a follow-up Enter does not
submit preedit text."""
"""Issue #5546: the keydown IME gate must re-pin composingRef so a follow-up
Enter does not submit preedit text after the watchdog clears it."""
src = THREAD_TSX.read_text()
assert "onKeyDown" in src, "main composer is missing onKeyDown IME gate"
assert "e.nativeEvent.isComposing" in src and "keyCode === 229" in src, (
@ -118,8 +107,7 @@ def test_main_composer_keydown_repins_composing_during_ime():
def test_compare_composer_keydown_repins_composing_during_ime():
"""Compare composer onKeyDown re-pins composingRef on IME keypress so a
follow-up click-Send during the watchdog window does not slip preedit
text through."""
follow-up click-Send during the watchdog window does not slip preedit text."""
src = SHARED_TSX.read_text()
assert "composingRef.current = true" in src, (
"compare composer keydown gate must re-pin composingRef when the "
@ -133,9 +121,8 @@ def _extract_block(
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."""
"""Source within the first balanced opener/closer after `anchor`, scoping
assertions to one handler."""
start = src.find(anchor)
assert start != -1, f"anchor {anchor!r} not found"
open_idx = src.find(opener, start)
@ -153,9 +140,8 @@ def _extract_block(
def test_main_composer_keydown_rearms_watchdog():
"""After the keydown re-pin sets composingRef=true the watchdog must
be re-armed; otherwise the WSL+Chrome no-compositionend path this PR
targets would lock Send permanently after any IME keypress."""
"""After keydown re-pins composingRef the watchdog must re-arm, else the
WSL+Chrome no-compositionend path locks Send after any IME keypress (#5546)."""
src = THREAD_TSX.read_text()
block = _extract_block(src, "const onKeyDown = useCallback")
assert "refreshStuckTimer" in block, (

72
tests/studio/test_frontend_dep_removal.py
View file

@ -3,14 +3,9 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Edge-case suite for scripts/check_frontend_dep_removal.py.
Each case patches a copy of studio/frontend/package.json to remove (or
move) a specific dependency, invokes the checker against the real
working tree's lockfile, and asserts the verdict matches expectations.
Run:
python tests/studio/test_frontend_dep_removal.py
Exits 0 iff every case behaves as expected.
Each case patches a copy of package.json to remove/move a dependency,
runs the checker against the real lockfile, and asserts the verdict.
Run: `python tests/studio/test_frontend_dep_removal.py` (exit 0 iff all pass).
"""
from __future__ import annotations
@ -315,8 +310,8 @@ def run_case(case: Case, head_pkg: dict) -> tuple[bool, str]:
)
# Classifier unit tests: feed hand-crafted snippets into classify() and assert
# the returned kind. Covers sneaky import shapes used to obscure a real usage.
# Classifier unit tests: feed snippets into classify(), assert the kind.
# Covers sneaky import shapes used to obscure a real usage.
# Import classify() by file path so this test needs no installed package.
import importlib.util as _ilu
@ -633,11 +628,8 @@ CLASSIFY_CASES: list[ClassifyCase] = [
'type C = import("react").ComponentType;',
"dynamic_import",
),
# File-type gating (codex P1: JS classifiers must not fire on
# non-script files). Python fixtures and Markdown code blocks often
# contain literal JS-shaped strings for documentation or test data,
# so a bare `import x from "pkg"` inside a .py / .md / .sh / .yml is
# not a real npm usage.
# File-type gating: JS classifiers must not fire on non-script files
# (.py/.md/.sh/.yml), whose JS-shaped strings are docs/test data, not usages.
ClassifyCase(
"U37",
"JS import snippet inside a Python fixture string is NOT a usage",
@ -696,8 +688,7 @@ CLASSIFY_CASES: list[ClassifyCase] = [
'<script src="/node_modules/foo/dist/index.js"></script>',
"html_script",
),
# CSS url() unquoted variant -- valid CSS, must classify the same
# as the quoted variant.
# CSS url() unquoted variant must classify the same as the quoted one.
ClassifyCase(
"U44",
"CSS url() unquoted bare package path",
@ -735,9 +726,8 @@ 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 detection regressions in the full
# grep+classify pipeline (not just classify in isolation).
# Adversarial end-to-end cases: drop a synthetic file into src/, run the
# checker, clean up. Catches regressions in the full grep+classify pipeline.
ADVERSARIAL_TMP_DIR = REPO / "studio/frontend/src/__dep_check_adversarial__"
@ -794,9 +784,7 @@ ADV_CASES: list[AdvCase] = [
"A05",
"package with similar prefix should NOT trigger FAIL",
"adv05.ts",
# The file imports __adv_only_pkg_e_extra__, but we will try
# to "remove" the shorter __adv_only_pkg_e__ name. The shorter
# name has zero real usage, so removal must be safe.
# Imports the *_extra* name; removing the shorter name is safe (zero usage).
'import x from "__adv_only_pkg_e_extra__";\n',
"__adv_only_pkg_e__",
"PASS",
@ -867,12 +855,8 @@ ADV_CASES: list[AdvCase] = [
"FAIL",
["__adv_only_pkg_l__"],
),
# Prettier formats a long named-import list one identifier per line.
# 22 imports + braces puts the `import` keyword ~22 lines away from
# the `from "pkg"` clause. Before the window widening, the classify
# multi-line fallback used ±4 lines, which silently missed every
# such block. This case fails with the old window and passes once
# the window is wide enough (currently ±25).
# Prettier puts `import` ~22 lines from the `from "pkg"` clause; the old
# ±4-line classify fallback missed it. Exercises the widened (±25) window.
AdvCase(
"A13",
"Prettier-style 22-identifier multi-line import should FAIL "
@ -888,9 +872,8 @@ ADV_CASES: list[AdvCase] = [
]
# package.json field-reference cases: simulate `prettier: "@x/config"`,
# `eslintConfig.extends`, `overrides`, `peerDependenciesMeta`, etc., testing
# package_json_extra_refs() coverage across common tool manifests.
# package.json field-reference cases: simulate prettier/eslintConfig/overrides/
# peerDependenciesMeta etc., testing package_json_extra_refs() coverage.
@dataclass
@ -1069,9 +1052,8 @@ def run_pkg_field_cases() -> int:
# Apply the field patch (deep-merge isn't needed; we control the keys).
for k, v in pc.field_patch.items():
synth_head[k] = v
# Base has the target in dependencies; head does not. The extra field
# in synth_head references the target pkg even though it's no longer
# in deps.
# Base declares the target; head drops it from deps but references it
# via the extra field.
synth_base = json.loads(json.dumps(head_pkg))
synth_base.setdefault("dependencies", {})[pc.target_pkg] = "^1.0.0"
with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f:
@ -1108,8 +1090,7 @@ def run_pkg_field_cases() -> int:
continue
if in_summary and line.strip().startswith("- "):
fails.append(line.strip()[2:])
# The expected_failures includes the tolerated-FP case (P15); we
# accept BOTH expected_status and expected_failures matches.
# Both status and failure set must match.
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}")
@ -1134,9 +1115,8 @@ def run_adversarial_cases() -> int:
fpath = ADVERSARIAL_TMP_DIR / ac.filename
try:
fpath.write_text(ac.content)
# Build a synthetic base that has the target pkg added; head
# is the real head (without it). The script sees the pkg as
# removed and scans the repo, which now includes our file.
# Base adds the target pkg; real head lacks it, so the script
# treats it as removed and scans the repo (now with our file).
synth_base = json.loads(json.dumps(head_pkg))
synth_base.setdefault("dependencies", {})[ac.target_pkg] = "^1.0.0"
with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f:
@ -1370,9 +1350,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 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.
# Script-wrapper cases: scripts_bin_refs / _next_real_bin must credit the real
# bin (`biome` -> @biomejs/biome), not the wrapper. The old "first non-env
# token" heuristic missed cross-env / dotenv / etc.
@dataclass
@ -1469,10 +1449,8 @@ def run_wrapper_cases() -> int:
if ok:
passed += 1
# End-to-end integration: feed scripts_bin_refs a synthetic head_pkg
# whose scripts use a wrapper, and confirm the package owning the
# wrapped bin is credited (rather than the wrapper). This is the
# actual call path used by find_command_usage().
# End-to-end: feed scripts_bin_refs a head_pkg whose scripts use a wrapper
# and confirm the wrapped bin's owner is credited (the find_command_usage path).
int_total = 0
int_passed = 0
int_cases = [

51
tests/studio/test_hardware_dispatch_matrix.py
View file

@ -1,12 +1,5 @@
# SPDX-License-Identifier: AGPL-3.0-only
"""Hardware dispatch matrix for Studio.
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."""
"""Studio hardware dispatch matrix: spoofs platform/torch/mlx per PROFILES to exercise CUDA/ROCm/XPU/MLX/CPU paths without real hardware."""
from __future__ import annotations
@ -26,11 +19,6 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
# ---------------------------------------------------------------------------
# Profile definition
# ---------------------------------------------------------------------------
@dataclass
class HardwareProfile:
name: str
@ -158,15 +146,9 @@ PROFILES = [
PROFILE_IDS = [p.name for p in PROFILES]
# ---------------------------------------------------------------------------
# Spoofing helpers
# ---------------------------------------------------------------------------
@pytest.fixture
def spoof_hardware(monkeypatch):
"""Return a function that applies a HardwareProfile to the live process.
Idempotent; monkeypatch cleans up on test exit."""
"""Return a function that applies a HardwareProfile to the live process; monkeypatch cleans up on exit."""
def _apply(profile: HardwareProfile) -> None:
import platform
@ -177,8 +159,7 @@ def spoof_hardware(monkeypatch):
monkeypatch.setattr(platform, "machine", lambda: profile.machine)
monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available)
# Stub get_device_properties: detect_hardware reads .name when CUDA is
# available, which crashes on a CPU CI runner ("No CUDA GPUs").
# Stub get_device_properties: detect_hardware reads .name, which crashes on a CPU CI runner.
if profile.cuda_available:
stub_props = types.SimpleNamespace(
name = "Stub GPU" if not profile.hip_version else "Stub AMD GPU",
@ -194,8 +175,7 @@ def spoof_hardware(monkeypatch):
torch_version = torch.version
monkeypatch.setattr(torch_version, "hip", profile.hip_version, raising = False)
# Stub torch.xpu.* (detect_hardware reads both); real get_device_name
# needs the XPU torch build, so always stub to stay hardware-agnostic.
# Stub torch.xpu.* always; real get_device_name needs the XPU torch build.
if hasattr(torch, "xpu"):
monkeypatch.setattr(torch.xpu, "is_available", lambda: profile.xpu_available)
monkeypatch.setattr(
@ -225,8 +205,7 @@ def spoof_hardware(monkeypatch):
monkeypatch.setitem(sys.modules, "mlx", fake_mlx)
monkeypatch.setitem(sys.modules, "mlx.core", fake_mlx_core)
else:
# Drop cached mlx and patch find_spec so the unsloth gate sees
# mlx as absent.
# 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)
real_find_spec = importlib.util.find_spec
@ -238,9 +217,8 @@ def spoof_hardware(monkeypatch):
monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)
# 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.*.
# Studio's _has_mlx() does `import mlx.core`, not find_spec; block it
# with a meta_path finder that raises ImportError for mlx.*.
class _BlockMLXFinder:
def find_spec(
self_inner,
@ -255,8 +233,7 @@ def spoof_hardware(monkeypatch):
return None
blocker = _BlockMLXFinder()
# New list so monkeypatch fully restores on teardown (mutating in
# place would survive the test).
# New list so monkeypatch fully restores on teardown.
monkeypatch.setattr(
sys,
"meta_path",
@ -282,7 +259,7 @@ def _import_studio_hardware_module():
"""Lazy-load Studio's hardware module under the bare-imports layout."""
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
# Force a fresh import so detect_hardware re-runs under the current spoofs.
# Fresh import so detect_hardware re-runs under the current spoofs.
sys.modules.pop("utils.hardware.hardware", None)
sys.modules.pop("utils.hardware", None)
from utils.hardware import hardware as hw # type: ignore
@ -290,11 +267,6 @@ def _import_studio_hardware_module():
return hw
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_unsloth_is_mlx_gate_matches_profile(profile, spoof_hardware):
"""The _IS_MLX expression in unsloth/__init__.py flips correctly per profile."""
@ -333,14 +305,11 @@ def test_studio_is_apple_silicon_matches_profile(profile, spoof_hardware):
)
# ---------------------------------------------------------------------------
# Negative-space tests: catch regressions where the dispatch order changes.
# ---------------------------------------------------------------------------
def test_cuda_takes_priority_over_mlx_when_both_available(spoof_hardware):
"""If both CUDA and MLX are available, Studio MUST pick CUDA: the canary
guarding GPU users from being silently routed to MLX after refactors."""
"""CUDA wins over MLX when both available: canary against GPU users being routed to MLX after refactors."""
profile = HardwareProfile(
name = "cuda_plus_mlx",
system = "Darwin",

68
tests/studio/test_is_mlx_dispatch_gate.py
View file

@ -1,34 +1,14 @@
# SPDX-License-Identifier: AGPL-3.0-only
"""
Regression tests for the CUDA-vs-MLX dispatch gates Studio relies on.
"""Regression tests for the CUDA-vs-MLX dispatch gates Studio relies on.
Two gates drive every dispatch decision in Studio's MLX path:
1. ``unsloth._IS_MLX`` at the top of ``unsloth/__init__.py`` -- evaluated
once at import time and read by Studio worker code to choose between
the GPU and MLX trainer / inference / export paths. It delegates to
the shared zoo MLX runtime gate, with a local import barrier while the
paired unsloth-zoo runtime rollout is in flight.
2. ``utils.hardware.detect_hardware()`` -- runtime probe in the Studio
backend. Priority order: CUDA -> XPU -> MLX -> CPU. The MLX branch is
reached only when both CUDA and XPU are unavailable AND the host is
Apple Silicon AND mlx is importable.
These gates are the canaries for "MLX support accidentally hijacks
CUDA/AMD/Intel users". The tests here:
* verify the source-level structure of the ``_IS_MLX`` helper so an
accidental rewrite importing zoo before the local MLX precheck is caught,
* exercise the runtime gate logic under a spoofed Darwin+arm64 platform
with a fake ``mlx`` module in ``sys.modules`` to confirm both gates
flip True together,
* confirm that on the actual Linux+CUDA test host both gates remain in
their CUDA-side state.
No real MLX install is required; uses the same ``monkeypatch.setitem``
fake-mlx pattern as ``test_mlx_inference_backend.py``.
Two gates: (1) ``unsloth._IS_MLX`` (import-time, delegates to the zoo MLX
runtime gate behind a local precheck barrier); (2)
``utils.hardware.detect_hardware()`` (runtime, CUDA->XPU->MLX->CPU). These
are the canaries against "MLX support accidentally hijacks CUDA/AMD/Intel
users": we check the _IS_MLX helper structure, flip both gates True under a
spoofed Darwin+arm64 with a fake mlx module, and confirm both stay CUDA-side
on the real host. No real MLX install needed.
"""
import ast
@ -46,7 +26,7 @@ UNSLOTH_INIT = REPO_ROOT / "unsloth" / "__init__.py"
def test_is_mlx_gate_uses_three_required_predicates():
"""_IS_MLX must AND the three checks Studio depends on (Darwin, arm64, importable mlx); dropping any breaks dispatch."""
"""_IS_MLX must AND Darwin+arm64+importable-mlx; dropping any breaks dispatch."""
tree = ast.parse(UNSLOTH_INIT.read_text())
target = None
@ -90,13 +70,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).
# 2. Runtime gate behavior with platform spoofed to Apple Silicon + fake mlx.
# Re-evaluates the expression rather than reloading unsloth (avoids a torch
# cascade-reload).
def _evaluate_is_mlx_precheck(platform_module, importlib_util, os_module):
"""Re-evaluate the local _is_mlx_available precheck (the import barrier before zoo) with injected deps."""
"""Re-evaluate the local _is_mlx_available precheck with injected deps."""
return (
os_module.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1"
and platform_module.system() == "Darwin"
@ -109,7 +89,7 @@ def test_is_mlx_gate_true_on_apple_silicon_with_mlx_present(monkeypatch):
import platform
import importlib.util
# Inject a fake mlx package so find_spec returns a non-None ModuleSpec.
# Fake mlx so find_spec returns a non-None ModuleSpec.
fake_mlx = types.ModuleType("mlx")
fake_mlx.__spec__ = importlib.machinery.ModuleSpec("mlx", loader = None)
fake_mlx.__path__ = []
@ -127,7 +107,7 @@ def test_is_mlx_gate_false_when_mlx_missing(monkeypatch):
import platform
import importlib.util
# Apple Silicon platform but no mlx package -> gate must be False.
# Apple Silicon but no mlx -> gate must be False.
monkeypatch.delitem(sys.modules, "mlx", raising = False)
monkeypatch.setattr(platform, "system", lambda: "Darwin")
monkeypatch.setattr(platform, "machine", lambda: "arm64")
@ -152,7 +132,6 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
import importlib.util
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.")
@ -162,15 +141,13 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
# ---------------------------------------------------------------------------
# 3. Studio's runtime detect_hardware() picks MLX only when CUDA + XPU are
# both unavailable AND the host is Apple Silicon AND mlx is importable.
# 3. detect_hardware() picks MLX only when CUDA+XPU are both unavailable AND
# the host is Apple Silicon AND mlx is importable.
# ---------------------------------------------------------------------------
def _import_studio_hardware():
"""Lazy import for the Studio hardware module, with the bare-imports
convention that Studio uses (studio/backend on sys.path).
"""
"""Lazy import of the Studio hardware module (studio/backend on sys.path)."""
studio_backend = REPO_ROOT / "studio" / "backend"
if str(studio_backend) not in sys.path:
sys.path.insert(0, str(studio_backend))
@ -182,14 +159,14 @@ def _import_studio_hardware():
def test_detect_hardware_picks_mlx_when_only_apple_silicon_available(monkeypatch):
hw = _import_studio_hardware()
# Force CUDA + XPU paths off so detect_hardware falls through to MLX.
# Force CUDA + XPU off so detect_hardware falls through to MLX.
import torch
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
if hasattr(torch, "xpu"):
monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
# Spoof Apple Silicon and provide an importable mlx.core for _has_mlx().
# Spoof Apple Silicon + importable mlx.core for _has_mlx().
import platform
monkeypatch.setattr(platform, "system", lambda: "Darwin")
@ -206,10 +183,7 @@ def test_detect_hardware_picks_mlx_when_only_apple_silicon_available(monkeypatch
def test_detect_hardware_picks_cuda_on_real_host():
"""Canary: on a real CUDA host the MLX branch must NOT be taken even
if mlx happens to be importable. Protects CUDA/AMD/Intel users from
accidental MLX dispatch when MLX support is added.
"""
"""Canary: a real CUDA host must dispatch to CUDA even if mlx is importable."""
import torch
if not torch.cuda.is_available():

77
tests/studio/test_stream_cancel_registration_timing.py
View file

@ -1,20 +1,10 @@
"""Tests that the cancel tracker is registered BEFORE StreamingResponse is
returned and that cleanup runs in a `finally` inside each async generator.
"""Cancel tracker must register BEFORE StreamingResponse returns and clean up
in each async generator's `finally`, else a Stop before the first SSE chunk
leaves a zombie decode (a BackgroundTask would be skipped when stream_response raises).
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.
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].
"""
Structural verifies registration placement and try/finally cleanup; behavioral
verifies the extracted `_TrackedCancel` cleans up across completion/OSError/aclose
and that a pre-set cancel_event breaks the GGUF loop cleanly with final_chunk + [DONE]."""
from __future__ import annotations
@ -262,8 +252,8 @@ async def _consume(agen):
def _llama_stub_raises_on_preset_cancel(cancel_event):
# Reproduces llama_cpp.py _stream_with_retry:2240 `raise GeneratorExit`
# when cancel_event is already set at entry.
# Reproduces llama_cpp.py _stream_with_retry `raise GeneratorExit` when
# cancel_event is already set at entry.
if cancel_event.is_set():
raise GeneratorExit
yield "cumulative-1"
@ -302,9 +292,8 @@ def test_finally_cleanup_on_normal_completion():
def test_finally_cleanup_on_mid_stream_exception():
# Simulates OSError / BrokenPipeError from Starlette send() mid-stream --
# the exact case where pre-fix `background = BackgroundTask(...)` was
# skipped and leaked the registry entry.
# OSError mid-stream: the exact case where pre-fix `background=BackgroundTask(...)`
# was skipped and leaked the registry entry.
m = _load_registry_module()
m["_CANCEL_REGISTRY"].clear()
ev = threading.Event()
@ -317,8 +306,7 @@ def test_finally_cleanup_on_mid_stream_exception():
def test_finally_cleanup_on_aclose():
# Starlette calls aclose() on the async generator when the client
# disconnects mid-stream. The generator's finally block must run.
# Starlette calls aclose() on client disconnect; the finally block must run.
m = _load_registry_module()
m["_CANCEL_REGISTRY"].clear()
ev = threading.Event()
@ -338,9 +326,8 @@ def test_finally_cleanup_on_aclose():
def test_preset_cancel_event_exits_cleanly_with_done():
# 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.
# Pending-replay: a stashed cancel pre-set cancel_event. The loop must break
# cleanly with final_chunk + [DONE], not propagate GeneratorExit from the GGUF wrapper.
ev = threading.Event()
ev.set()
chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev)))
@ -353,16 +340,14 @@ def test_preset_cancel_event_exits_cleanly_with_done():
def test_normal_path_streams_all_tokens():
# Regression: the top-of-loop cancel_event check must not short-circuit
# when cancel_event is unset.
# Regression: the top-of-loop cancel_event check must not short-circuit when unset.
ev = threading.Event()
chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev)))
assert chunks == ["first_chunk", "cumulative-1", "cumulative-2", "final_chunk", "[DONE]"]
def test_cancel_during_streaming_stops_iteration_promptly():
# Setting cancel_event between yields breaks out on the next iteration
# rather than draining the stub generator.
# Setting cancel_event between yields breaks on the next iteration, not draining the generator.
ev = threading.Event()
async def _run():
@ -386,10 +371,8 @@ def test_cancel_during_streaming_stops_iteration_promptly():
def _loop_has_cancel_event_check(fn) -> bool:
# An `if cancel_event.is_set():` statement anywhere inside a
# `while`/`for` loop body is sufficient -- without it, a cancel POST
# cannot interrupt the loop because Colab-style proxies do not
# propagate request.is_disconnected().
# An `if cancel_event.is_set():` inside a loop body is sufficient -- without it
# a cancel POST can't interrupt, since Colab-style proxies drop request.is_disconnected().
for sub in ast.walk(fn):
if not isinstance(sub, (ast.While, ast.For, ast.AsyncFor)):
continue
@ -430,10 +413,8 @@ def test_streaming_generators_check_cancel_event_in_loop():
def test_audio_input_stream_offloads_blocking_next_to_thread():
# Guards against regression back to `for chunk_text in
# audio_input_generate():` -- which blocks the event loop on each
# whisper chunk and prevents POST /api/inference/cancel from being
# serviced until the chunk yields.
# Guards against regressing to `for chunk_text in audio_input_generate():`, which
# blocks the event loop per whisper chunk and stalls POST /api/inference/cancel.
audio = None
for fn in ast.walk(_TREE):
if isinstance(fn, ast.AsyncFunctionDef) and fn.name == "audio_input_stream":
@ -473,9 +454,8 @@ def test_audio_input_stream_offloads_blocking_next_to_thread():
def test_stream_chunks_cancel_branch_resets_backend_state():
# 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.
# The cancel branch must call backend.reset_generation_state() to flush
# GPU/KV-cache state, else cancel-via-POST leaves the subprocess dirty.
fn = None
top = None
for n in ast.walk(_TREE):
@ -563,9 +543,8 @@ def test_unsloth_stream_loop_breaks_on_external_cancel_event():
def test_audio_stream_stays_responsive_under_blocking_next():
# Regression guard: replace the post-fix loop with the pre-fix
# `for chunk in audio_input_generate()` pattern and assert it blocks
# the event loop; then confirm the post-fix pattern exits promptly.
# Assert the pre-fix `for chunk in audio_input_generate()` pattern blocks the
# event loop, then confirm the post-fix pattern exits promptly.
cancel_event = threading.Event()
def _audio_gen():
@ -624,9 +603,8 @@ def test_audio_stream_stays_responsive_under_blocking_next():
def test_unsloth_stream_loop_emits_zero_tokens_on_preset_cancel():
# 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).
# Pending-cancel replay: cancel_event pre-set, so the top-of-loop check must
# short-circuit iteration 1 (zero tokens). Catches moving the check below next().
cancel_event = threading.Event()
cancel_event.set()
reset_calls = [0]
@ -674,9 +652,8 @@ def test_unsloth_stream_loop_emits_zero_tokens_on_preset_cancel():
def test_audio_stream_emits_zero_chunks_on_preset_cancel():
# Symmetric to the Unsloth pre-set test: the audio loop's top-of-loop
# cancel check must skip the asyncio.to_thread(next, ...) call when
# cancel_event was already set via pending-replay.
# Symmetric to the Unsloth pre-set test: the audio loop must skip
# asyncio.to_thread(next, ...) when cancel_event was pre-set via pending-replay.
cancel_event = threading.Event()
cancel_event.set()

15
tests/studio/test_studio_gguf_export_script_pin.py
View file

@ -1,17 +1,4 @@
"""Tests for Studio GGUF export pinning convert_hf_to_gguf.py via
UNSLOTH_LLAMA_CPP_SCRIPTS_DIR with graceful fallback when unsloth_zoo
lacks the local-script resolver.
Verifies:
- export.py imports LLAMA_CPP_DEFAULT_DIR and _resolve_local_convert_script
from unsloth_zoo.llama_cpp inside a single try/except ImportError so a
zoo missing either symbol degrades to a warning instead of crashing.
- os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR)
is called inside the try; setdefault preserves explicit user overrides
and assigns the default when unset.
- The compatibility warning is gated on a module-level flag so it fires
once per process rather than on every export call.
"""
"""Studio GGUF export pins convert_hf_to_gguf.py via UNSLOTH_LLAMA_CPP_SCRIPTS_DIR, with a once-per-process warning fallback when unsloth_zoo lacks the local-script resolver."""
from __future__ import annotations

10
tests/studio/test_studio_text_descender_clipping.py
View file

@ -1,7 +1,5 @@
"""
Regression guard: descender-prone text spans in Studio must not pair
`leading-none` with `truncate` (overflow: hidden), which clips glyph
descenders (g, p, q, y, j) in real user-visible labels.
"""Regression guard: Studio text spans must not pair `leading-none` with
`truncate`, which clips glyph descenders (g, p, q, y, j) in visible labels.
"""
from __future__ import annotations
@ -36,8 +34,8 @@ def test_model_selector_trigger_label_uses_leading_tight():
def test_sidebar_account_block_uses_leading_tight():
src = _read(APP_SIDEBAR)
# Match the account-block parent div regardless of its gap utility (gap-0.5,
# gap-px, ...); this guard is about the leading-* class, not the spacing.
# Match the account-block parent div regardless of its gap utility; this
# guard is about the leading-* class, not the spacing.
pattern = re.compile(
r'<div\s+className="flex\s+flex-col\s+gap-\S+\s+(\S+)\s+group-data-\[collapsible=icon\]:hidden">',
)

76
tests/test_callback_signature_drift.py
View file

@ -1,28 +1,14 @@
"""Static-analysis regression test: callback signature drift.
Catches the class of bug where a producer (e.g. unsloth_zoo's MLXTrainer)
changes the number of args it passes to a registered callback but consumers
(unsloth tests / source) still declare the old arity. The producer's
``try / except Exception`` typically swallows the resulting TypeError, so
the callback silently never fires and the failure surfaces several seconds
later as a confusing downstream assertion.
Catches a producer (e.g. unsloth_zoo's MLXTrainer) changing the arity it passes to a registered
callback while consumers still declare the old arity; the producer's try/except swallows the
TypeError so the callback silently never fires. Pure AST so it runs on every CI OS/Python.
The check is pure AST (no imports of MLX modules etc), so it runs on every
OS / Python version that ships in CI.
Pattern detected:
* Producer side: a class with ``self._<name>_callbacks`` list, populated
via ``self._<name>_callbacks.append(...)`` from an ``add_<name>_callback``
method, and invoked via ``for cb in self._<name>_callbacks: cb(arg1, ...)``.
The arity at the call site is the canonical expected arity.
* Consumer side: any ``<obj>.add_<name>_callback(fn)`` call where ``fn``
resolves to a ``def`` or ``async def`` in the same file. Consumer arity
must equal canonical arity (or be variadic).
Consumers handled tolerantly:
* ``*args`` / ``**kwargs``: accept any canonical arity.
* Methods (``self.fn``) and unresolved Name targets (imported from another
file): skipped with a note in the failure message rather than asserted.
Producer: a class with ``self._<name>_callbacks`` populated by ``add_<name>_callback`` and invoked
via ``for cb in self._<name>_callbacks: cb(...)`` (the call-site arity is canonical).
Consumer: ``<obj>.add_<name>_callback(fn)`` where ``fn`` is a def/async def in the same file; its
arity must equal canonical (or be variadic). ``*args``/``**kwargs`` accept any arity; methods and
unresolved Name targets are skipped with a note.
"""
from __future__ import annotations
@ -47,7 +33,7 @@ SKIP_PARTS = {
"venv",
".pytest_cache",
"__pycache__",
# Frontend tree under studio is JS/TS plus a few stub .py files; not worth walking.
# studio frontend is JS/TS plus a few stub .py files; skip.
"frontend",
}
@ -66,8 +52,7 @@ def _iter_py(root: pathlib.Path):
yield p
# Module-level parse cache so discover_producers + check_registrations only
# pay the parse cost once per file across the whole test run.
# Parse cache so each file is parsed once across the run.
_PARSE_CACHE: dict[pathlib.Path, ast.AST | None] = {}
@ -78,8 +63,7 @@ def _safe_parse(path: pathlib.Path):
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.
# Suppress SyntaxWarning from third-party files with invalid escape sequences.
_w.simplefilter("ignore", SyntaxWarning)
tree = ast.parse(path.read_text(encoding = "utf-8"))
except (SyntaxError, UnicodeDecodeError):
@ -119,10 +103,7 @@ def _callback_list_attrs_in_class(cls: ast.ClassDef) -> set[str]:
def _producer_arities(tree: ast.AST) -> dict[str, int]:
"""For each ``for cb in self._x_callbacks: cb(...)`` in the AST, return
{cb_list_attr: max_arity}. Multiple sites take the max so that variadic
branches do not lower the contract.
"""
"""Return {cb_list_attr: max_arity} over all ``for cb in self._x_callbacks: cb(...)`` sites."""
out: dict[str, int] = {}
for cls in [n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]:
cb_lists = _callback_list_attrs_in_class(cls)
@ -171,11 +152,8 @@ def _func_arity(node: ast.AST) -> tuple[int, bool] | None:
args = node.args
arity = len(args.posonlyargs) + len(args.args)
accepts_var = args.vararg is not None
# Bound methods: drop the implicit self if this is a method-style def.
# We can't tell statically whether the def is a method without class
# context, so we conservatively do not subtract self here. The consumer
# check skips bare-Name registrations whose target is a `self.fn` attr
# anyway.
# Don't subtract self: we can't tell statically if this is a method, and the
# consumer check skips `self.fn` registrations anyway.
return arity, accepts_var
@ -197,9 +175,9 @@ def discover_producers(roots: list[pathlib.Path]) -> dict[str, list[tuple[pathli
def check_registrations(
roots: list[pathlib.Path], producers: dict[str, list[tuple[pathlib.Path, int]]]
):
"""Walk every .py under each root, find <x>.add_*_callback(fn) where fn is a
bare Name resolvable to a def in the same file, and assert its arity
matches the producer's canonical arity. Returns (issues, skipped, ok_count).
"""Assert each in-file <x>.add_*_callback(fn) arity matches the producer's canonical arity.
Returns (issues, skipped, ok_count).
"""
issues: list[str] = []
skipped: list[str] = []
@ -211,7 +189,7 @@ def check_registrations(
tree = _safe_parse(src)
if tree is None:
continue
# All function/lambda defs in this file by name (and by id for lambdas via assignment).
# All function/lambda defs in this file, keyed by name.
defs_by_name: dict[str, ast.AST] = {}
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
@ -223,7 +201,7 @@ def check_registrations(
and isinstance(node.targets[0], ast.Name)
):
defs_by_name[node.targets[0].id] = node.value
# Find <x>.add_*_callback(fn) sites
# Find <x>.add_*_callback(fn) sites.
for call in ast.walk(tree):
if not isinstance(call, ast.Call):
continue
@ -238,7 +216,7 @@ def check_registrations(
f"defines {cb_list} (third-party API?)"
)
continue
# Only handle bare-Name registrations; bound methods / partials skipped.
# Only bare-Name registrations; bound methods/partials skipped.
if not (len(call.args) == 1 and isinstance(call.args[0], ast.Name)):
skipped.append(
f"{src}:{call.lineno}: {call.func.attr}(...) registers a "
@ -274,12 +252,9 @@ def check_registrations(
def _zoo_roots() -> list[pathlib.Path]:
"""Where to look for unsloth_zoo source. We try, in order:
1. ``UNSLOTH_ZOO_SRC`` env var (a local git checkout).
2. ``../unsloth-zoo`` next to this repo (common monorepo-style layout).
3. The pip-installed package (wheel may strip platform-specific submodules
like ``mlx/``, so this often misses MLX producers).
Every root that exists is scanned; duplicates are fine.
"""unsloth_zoo source roots, in order: UNSLOTH_ZOO_SRC env, ../unsloth-zoo sibling, pip package.
(The pip wheel may strip submodules like mlx/, missing MLX producers.) All existing roots scanned.
"""
roots: list[pathlib.Path] = []
env_src = os.environ.get("UNSLOTH_ZOO_SRC")
@ -292,9 +267,7 @@ def _zoo_roots() -> list[pathlib.Path]:
roots.append(sibling)
spec = importlib.util.find_spec("unsloth_zoo")
if spec is not None and spec.origin is not None:
# spec.origin -> .../site-packages/unsloth_zoo/__init__.py
# we want the unsloth_zoo dir itself, NOT the site-packages root which
# contains every other installed pkg.
# Use the unsloth_zoo dir itself (parent of __init__.py), not the site-packages root.
roots.append(pathlib.Path(spec.origin).resolve().parent)
return roots
@ -326,7 +299,6 @@ def test_no_callback_signature_drift():
if __name__ == "__main__":
# Allow running directly as a script for fast feedback.
sys.argv.append("-v")
test_no_callback_signature_drift()
print("PASS")

42
tests/test_cli_export_unpacking.py
View file

@ -1,14 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression tests for unsloth_cli.commands.export.
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."""
"""Regression tests for unsloth_cli.commands.export: pin the CLI to the export_* 3-tuple contract (was unpacking 2, crashing every `unsloth export`) via a fake ExportBackend in sys.modules."""
from __future__ import annotations
@ -21,14 +14,8 @@ import typer
from typer.testing import CliRunner
# ---------------------------------------------------------------------------
# Fake ExportBackend
# ---------------------------------------------------------------------------
class _FakeExportBackend:
"""Stand-in for ExportBackend: export_* return the new 3-tuple;
load_checkpoint keeps its 2-tuple shape."""
"""Stand-in for ExportBackend: export_* return the 3-tuple, load_checkpoint stays a 2-tuple."""
def __init__(self) -> None:
self.loaded: str | None = None
@ -54,9 +41,7 @@ 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, so this steers it at the fake; parent packages are
stubbed too so import machinery skips the real structlog-dependent tree."""
"""Inject a fake studio.backend.core.export into sys.modules so the CLI's lazy import binds to it; parent packages stubbed to skip the structlog-dependent tree."""
for name in ("studio", "studio.backend", "studio.backend.core"):
monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
@ -64,7 +49,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 cached CLI module so export()'s deferred import re-resolves the fake.
# Drop the cached CLI module so its deferred import re-resolves the fake.
monkeypatch.delitem(sys.modules, "unsloth_cli.commands.export", raising = False)
@ -77,11 +62,8 @@ def cli_app(monkeypatch: pytest.MonkeyPatch) -> typer.Typer:
app = typer.Typer()
app.command("export")(export_cmd.export)
# Typer flattens a single-command app into that command, which would
# make argv[0] ("export") look like an extra positional argument to
# the test invocation. Register a harmless second command so Typer
# keeps "export" as a real subcommand and the tests drive the
# intended code path.
# Typer flattens a single-command app, making "export" look like a stray positional;
# a harmless second command keeps "export" a real subcommand.
@app.command("noop")
def _noop() -> None: # pragma: no cover - only exists to pin routing
pass
@ -94,11 +76,6 @@ def runner() -> CliRunner:
return CliRunner()
# ---------------------------------------------------------------------------
# The actual regression tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"format_flag,quant_flag",
[
@ -115,10 +92,7 @@ def test_cli_export_unpacks_three_tuple(
format_flag: str,
quant_flag: str | None,
) -> None:
"""Each --format path must unpack (success, message, output_path)
without raising ValueError. Pre-fix, every parametrized case fails
with 'too many values to unpack (expected 2)'.
"""
"""Each --format path unpacks the 3-tuple without ValueError (pre-fix: 'too many values to unpack (expected 2)')."""
ckpt = tmp_path / "ckpt"
ckpt.mkdir()
out = tmp_path / "out"
@ -134,6 +108,6 @@ def test_cli_export_unpacks_three_tuple(
f"Output:\n{result.output}\n"
f"Exception: {result.exception!r}"
)
# Sanity: the success message from the fake backend should reach stdout.
# Fake backend's success message should reach stdout.
expected_prefix = format_flag.split("-")[0]
assert f"{expected_prefix} ok" in result.output

33
tests/test_enforce_kwargs_spacing.py
View file

@ -1,10 +1,4 @@
"""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.
"""
"""Tests for scripts/enforce_kwargs_spacing.py rewrite rules (AST-preserving, idempotent)."""
from __future__ import annotations
@ -82,9 +76,8 @@ def test_blank_removed_for_small_import_block(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).
# Import and following statement now adjacent.
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
@ -125,8 +118,7 @@ def test_exact_output_try_block():
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).
# Only the blank after the LAST import in a run is dropped; both imports kept.
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)
@ -145,8 +137,7 @@ def test_multiple_blank_lines_in_gap_all_removed():
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.
# A blank inside a parenthesized import is part of the import, not the gap.
src = (
"def g():\n"
" from mod import (\n"
@ -318,8 +309,7 @@ def test_merge_adjacent_strings_skips(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.
# Folding a long f + plain assert message can't fit on one line, so leave it.
src = (
"def f():\n"
" assert some_condition_holds_here, (\n"
@ -332,8 +322,7 @@ def test_fstring_fold_skipped_when_statement_would_not_collapse():
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).
# A multi-line f + plain that fits on one line after folding is folded.
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
@ -342,9 +331,7 @@ def test_fstring_fold_applied_when_statement_collapses():
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.
# The fit guard only restricts asserts; an f + plain arg in a big call folds.
src = (
"findings.append(\n"
" Finding(\n"
@ -361,8 +348,7 @@ def test_fstring_fold_applied_inside_large_multiline_call():
# ── 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.
# Strips the trailing comma so ruff joins the assert onto one line; AST unchanged.
@pytest.mark.parametrize(
@ -393,9 +379,8 @@ def test_fstring_fold_applied_inside_large_multiline_call():
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).
# Magic trailing comma is gone, so ruff joins 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

9
tests/test_finetune_last_n_layers.py
View file

@ -1,9 +1,6 @@
# Unsloth - 2x faster, 70% less memory LLM finetuning
# Tests for the `finetune_last_n_layers` parity knob (CUDA side).
#
# 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.
# Tests for the `finetune_last_n_layers` parity knob (translation helper only,
# no CUDA / real checkpoint); mirrors unsloth-zoo's MLX layers_to_transform path.
from __future__ import annotations
@ -67,7 +64,7 @@ def test_get_total_transformer_layers_returns_none_for_missing_config():
def test_finetune_last_n_layers_signature_present_on_llama_and_vision():
"""Both entry points must expose the new parameter with default None."""
"""Both entry points expose finetune_last_n_layers with default None."""
import inspect
from unsloth.models.llama import FastLlamaModel
from unsloth.models.vision import FastBaseModel

2
tests/test_gemma4_chat_template.py
View file

@ -157,7 +157,7 @@ def test_multi_turn_strips_all_historical_model_turns():
def test_thinking_template_injects_empty_thought_channel_by_default():
# Author defaults enable_thinking=False, so the gen-prompt injection fires.
# enable_thinking defaults False, so the gen-prompt injection fires.
msgs = [{"role": "user", "content": "Hi"}]
out = _render("gemma4_thinking_template", msgs, add_generation_prompt = True)
assert out.endswith("<|turn>model\n<|channel>thought\n<channel|>")

168
tests/test_import_fixes_drift.py
View file

@ -11,12 +11,9 @@
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
"""Drift detectors for upstream pathologies that ``unsloth/import_fixes.py``
works around. One test per ``fix_*`` / ``patch_*`` function. Each asserts
the healthy upstream shape; if the pathology is active, fires
``pytest.fail("DRIFT DETECTED: ...")`` -- never ``pytest.skip`` -- so CI
goes red and the maintainer triages on the next PR. Runs under the
GPU-free harness in ``tests/conftest.py``."""
"""Drift detectors for the upstream pathologies ``unsloth/import_fixes.py``
works around; one test per ``fix_*`` / ``patch_*``, each fails (never skips)
when the pathology is active. Runs under the GPU-free ``tests/conftest.py``."""
from __future__ import annotations
@ -31,8 +28,7 @@ from importlib.metadata import version as importlib_version
import pytest
# Mirrors the local ``Version()`` in import_fixes.py (51-68): strip
# dev/alpha/beta/rc/local suffixes so packaging.Version doesn't choke.
# Mirrors import_fixes.py's local Version(): strip dev/alpha/beta/rc/local suffixes.
from packaging.version import Version as _PkgVersion
@ -52,7 +48,7 @@ def _safe_version(raw):
def test_protobuf_message_factory_get_prototype_or_get_message_class_present():
"""``fix_message_factory_issue`` (import_fixes.py 264-308)."""
"""``fix_message_factory_issue``."""
mf = pytest.importorskip("google.protobuf.message_factory")
has_mf_class = hasattr(mf, "MessageFactory")
has_get_prototype = has_mf_class and hasattr(mf.MessageFactory, "GetPrototype")
@ -75,8 +71,7 @@ def test_protobuf_message_factory_get_prototype_or_get_message_class_present():
def test_datasets_version_not_in_broken_recursion_range():
"""``patch_datasets`` (import_fixes.py 574-586). datasets 4.4.0-4.5.0
inclusive trigger RLock recursion errors in the Arrow loader."""
"""``patch_datasets``: datasets 4.4.0-4.5.0 hit RLock recursion in the Arrow loader."""
pytest.importorskip("datasets")
ds_v = _safe_version(importlib_version("datasets"))
lo = _PkgVersion("4.4.0")
@ -92,9 +87,8 @@ def test_datasets_version_not_in_broken_recursion_range():
def test_trl_is_x_available_returns_bool_not_tuple():
"""``fix_trl_vllm_ascend`` (import_fixes.py 493-516). transformers >=4.48's
``_is_package_available`` returns ``(bool, version_or_None)``; TRL's
``is_*_available`` accessors must still return real bools."""
"""``fix_trl_vllm_ascend``: TRL's ``is_*_available`` must still return bools
after transformers >=4.48 made ``_is_package_available`` return a tuple."""
pytest.importorskip("trl")
try:
import trl.import_utils as tiu
@ -139,8 +133,7 @@ def test_trl_is_x_available_returns_bool_not_tuple():
def test_trl_cached_available_flags_are_not_tuples():
"""``fix_trl_vllm_ascend`` (import_fixes.py 493-516). Same drift, checked
on the module-level cached ``_*_available`` attributes directly."""
"""``fix_trl_vllm_ascend``: same drift on the module-level cached ``_*_available`` attrs."""
pytest.importorskip("trl")
try:
import trl.import_utils as tiu
@ -163,12 +156,9 @@ def test_trl_cached_available_flags_are_not_tuples():
def test_pretrained_model_enable_input_require_grads_uses_old_pattern():
"""``patch_enable_input_require_grads`` (import_fixes.py 609-670). HF
PR #41993 rewrote enable_input_require_grads to iterate
``self.modules()`` and call ``get_input_embeddings`` on every
submodule; vision submodules then raise NotImplementedError. Healthy
state: either the upstream rewrite isn't present (pre-HF#41993), OR
the patch installed a NotImplementedError-tolerant replacement."""
"""``patch_enable_input_require_grads``: HF PR #41993 made
enable_input_require_grads iterate ``self.modules()``, so vision submodules
raise NotImplementedError unless the tolerant replacement is installed."""
pytest.importorskip("transformers")
from transformers import PreTrainedModel
@ -178,9 +168,9 @@ def test_pretrained_model_enable_input_require_grads_uses_old_pattern():
pytest.skip(f"could not getsource(enable_input_require_grads): {exc!r}")
if "for module in self.modules()" not in src:
return # healthy: pre-HF#41993 shape
return # pre-HF#41993 shape
if "NotImplementedError" in src:
return # healthy: unsloth's tolerant replacement is installed
return # tolerant replacement installed
pytest.fail(
"DRIFT DETECTED: PreTrainedModel.enable_input_require_grads now "
@ -192,10 +182,8 @@ def test_pretrained_model_enable_input_require_grads_uses_old_pattern():
def test_transformers_torchcodec_available_flag_is_present():
"""``disable_torchcodec_if_broken`` (import_fixes.py 1291-1317). Needs
either the pre-5.x module-level ``_torchcodec_available`` flag, or
the 5.x ``is_torchcodec_available`` public function; one of the two
is the patch site the fix monkey-patches when FFmpeg is missing."""
"""``disable_torchcodec_if_broken``: needs the pre-5.x ``_torchcodec_available``
flag or 5.x ``is_torchcodec_available`` as its patch site when FFmpeg is missing."""
tf_iu = pytest.importorskip("transformers.utils.import_utils")
has_flag = hasattr(tf_iu, "_torchcodec_available")
has_func = callable(getattr(tf_iu, "is_torchcodec_available", None))
@ -209,8 +197,7 @@ def test_transformers_torchcodec_available_flag_is_present():
def test_transformers_is_causal_conv1d_available_symbol_present():
"""``_disable_transformers_causal_conv1d`` (import_fixes.py 1881-1895).
Needs at least one of the causal_conv1d availability hooks."""
"""``_disable_transformers_causal_conv1d``: needs a causal_conv1d availability hook."""
tf_iu = pytest.importorskip("transformers.utils.import_utils")
candidates = [
"is_causal_conv1d_available",
@ -230,10 +217,8 @@ def test_transformers_is_causal_conv1d_available_symbol_present():
def test_transformers_and_accelerate_is_wandb_available_callable():
"""``disable_broken_wandb`` (import_fixes.py 1320-1372). Patches
is_wandb_available in transformers.integrations.integration_utils
AND accelerate.utils.imports / accelerate.utils -- all three must
keep existing."""
"""``disable_broken_wandb``: patches is_wandb_available in three modules
(transformers integration_utils + accelerate imports/utils); all must exist."""
pytest.importorskip("transformers")
pytest.importorskip("accelerate")
from transformers.integrations import integration_utils as tf_integration
@ -260,10 +245,8 @@ def test_transformers_and_accelerate_is_wandb_available_callable():
def test_peft_transformers_weight_conversion_importable_and_signature():
"""``patch_peft_weight_converter_compatibility`` (import_fixes.py
1375-1454). Wraps build_peft_weight_mapping to retrofit
distributed_operation / quantization_operation kwargs; if the
module is unimportable the wrap silently no-ops."""
"""``patch_peft_weight_converter_compatibility``: wraps build_peft_weight_mapping;
silently no-ops if the module is unimportable."""
pytest.importorskip("peft")
try:
from peft.utils import transformers_weight_conversion as twc
@ -290,17 +273,15 @@ def test_peft_transformers_weight_conversion_importable_and_signature():
def test_triton_compiled_kernel_has_num_ctas_and_cluster_dims():
"""``fix_triton_compiled_kernel_missing_attrs`` (import_fixes.py 923-968).
triton 3.6+ dropped num_ctas / cluster_dims on CompiledKernel; torch
2.9 Inductor's make_launcher still eagerly evaluates them."""
"""``fix_triton_compiled_kernel_missing_attrs``: triton 3.6+ dropped
num_ctas/cluster_dims on CompiledKernel, but Inductor's make_launcher needs them."""
pytest.importorskip("torch")
triton_mod = pytest.importorskip("triton") # noqa: F841
tc = pytest.importorskip("triton.compiler.compiler")
ck_cls = tc.CompiledKernel
# Healthy if either: pre-3.6 class attr present, or unsloth wrapped
# ``__init__`` to install num_ctas + cluster_dims per instance (the
# post-3.6 shape ``fix_triton_compiled_kernel_missing_attrs`` lands).
# Healthy if pre-3.6 class attr present, or __init__ wrapped to install
# num_ctas + cluster_dims per instance (the post-3.6 fix).
if hasattr(ck_cls, "num_ctas"):
return
init = getattr(ck_cls, "__init__", None)
@ -324,8 +305,7 @@ def test_triton_compiled_kernel_has_num_ctas_and_cluster_dims():
# torch + torchvision pairing table
# Mirrors TORCH_TORCHVISION_COMPAT in torchvision_compatibility_check
# (import_fixes.py 708-798).
# Mirrors TORCH_TORCHVISION_COMPAT in torchvision_compatibility_check.
_TORCH_TORCHVISION_COMPAT = {
(2, 9): (0, 24),
(2, 8): (0, 23),
@ -346,9 +326,8 @@ def _is_custom_torch_build(raw_version_str):
def test_installed_torch_torchvision_pair_is_compatible():
"""``torchvision_compatibility_check`` (import_fixes.py 708-798).
Raises ImportError when installed (torch, torchvision) pair fails
the pinned compat table; custom / prerelease builds are warning-only."""
"""``torchvision_compatibility_check``: raises when the (torch, torchvision)
pair fails the pinned table; custom/prerelease builds are warning-only."""
pytest.importorskip("torch")
pytest.importorskip("torchvision")
@ -389,9 +368,8 @@ def test_installed_torch_torchvision_pair_is_compatible():
def test_vllm_guided_decoding_params_or_structured_outputs_present():
"""``fix_vllm_guided_decoding_params`` (import_fixes.py 446-490).
vLLM PR #22772 renamed GuidedDecodingParams -> StructuredOutputsParams;
trl still imports the old name so the fix re-aliases."""
"""``fix_vllm_guided_decoding_params``: vLLM PR #22772 renamed
GuidedDecodingParams -> StructuredOutputsParams; the fix re-aliases for trl."""
pytest.importorskip("vllm")
try:
sp = importlib.import_module("vllm.sampling_params")
@ -415,9 +393,8 @@ def test_vllm_guided_decoding_params_or_structured_outputs_present():
def test_vllm_aimv2_ovis_config_is_past_fix_version():
"""``fix_vllm_aimv2_issue`` (import_fixes.py 404-443). vLLM <0.10.1 has
an Ovis config that unconditionally registers ``aimv2`` and trips a
duplicate-key ValueError; the fix only touches old versions."""
"""``fix_vllm_aimv2_issue``: vLLM <0.10.1 double-registers ``aimv2`` (duplicate-key
ValueError); the fix only touches old versions."""
pytest.importorskip("vllm")
vllm_v = _safe_version(importlib_version("vllm"))
cutoff = _PkgVersion("0.10.1")
@ -433,9 +410,8 @@ def test_vllm_aimv2_ovis_config_is_past_fix_version():
def test_huggingface_hub_is_offline_mode_or_hf_hub_offline_present():
"""``fix_huggingface_hub`` (import_fixes.py 913-920). huggingface_hub
removed top-level ``is_offline_mode``; fix re-injects from
``huggingface_hub.constants.HF_HUB_OFFLINE``."""
"""``fix_huggingface_hub``: re-injects top-level ``is_offline_mode`` from
``constants.HF_HUB_OFFLINE`` after huggingface_hub dropped it."""
hub = pytest.importorskip("huggingface_hub")
has_top_level = False
try:
@ -461,8 +437,8 @@ def test_huggingface_hub_is_offline_mode_or_hf_hub_offline_present():
def test_torch_nn_init_trunc_normal_exists():
"""``patch_trunc_normal_precision_issue`` (import_fixes.py 971-1050).
fp16/bf16 stability wrapper monkey-patches torch.nn.init.trunc_normal_."""
"""``patch_trunc_normal_precision_issue``: fp16/bf16 wrapper monkey-patches
torch.nn.init.trunc_normal_, which must still exist."""
pytest.importorskip("torch")
import torch.nn.init as init_mod
@ -476,9 +452,8 @@ def test_torch_nn_init_trunc_normal_exists():
def test_xformers_is_post_num_splits_key_fix_or_not_installed():
"""``fix_xformers_performance_issue`` (import_fixes.py 312-341).
xformers <0.0.29 has the ``num_splits_key=-1`` perf bug Unsloth
rewrites at install time."""
"""``fix_xformers_performance_issue``: xformers <0.0.29 has the
``num_splits_key=-1`` perf bug Unsloth rewrites at install time."""
if importlib.util.find_spec("xformers") is None:
pytest.skip("xformers not installed -- nothing to drift-check.")
x_v = _safe_version(importlib_version("xformers"))
@ -495,9 +470,8 @@ def test_xformers_is_post_num_splits_key_fix_or_not_installed():
def test_transformers_pretrained_model_has_get_input_embeddings():
"""``patch_enable_input_require_grads`` (import_fixes.py 609-670).
The replacement function calls ``get_input_embeddings`` on every
submodule, so the accessor must still exist."""
"""``patch_enable_input_require_grads``: its replacement calls
``get_input_embeddings`` per submodule, so the accessor must still exist."""
pytest.importorskip("transformers")
from transformers import PreTrainedModel
@ -510,11 +484,10 @@ def test_transformers_pretrained_model_has_get_input_embeddings():
# accelerate -- ``is_X_available`` API stability used across the fixes
# transformers LOSS_MAPPING -- patch_loss_functions() coverage
# Regression for https://github.com/unslothai/unsloth/issues/4188:
# Qwen3_5ForConditionalGeneration has loss_type='ForConditionalGeneration',
# a separate LOSS_MAPPING key that was never patched, leaving the model with
# the stock ForCausalLMLoss which does logits.float() and OOMs on <=24 GB GPUs.
# Qwen3_5ForConditionalGeneration uses loss_type='ForConditionalGeneration', a
# separate LOSS_MAPPING key left unpatched, falling back to stock ForCausalLMLoss
# whose logits.float() OOMs on <=24 GB GPUs.
def _reset_loss_mapping(mapping, saved):
@ -523,9 +496,8 @@ def _reset_loss_mapping(mapping, saved):
def test_patch_loss_functions_covers_conditional_generation():
"""After patch_loss_functions(), every LOSS_MAPPING key that was aliased
to ForCausalLMLoss must also point at the Unsloth kernel -- not just
LOSS_MAPPING['ForCausalLM']."""
"""patch_loss_functions() must repoint every ForCausalLMLoss alias to the
Unsloth kernel, not just LOSS_MAPPING['ForCausalLM']."""
lu = pytest.importorskip("transformers.loss.loss_utils")
cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")
@ -550,8 +522,7 @@ def test_patch_loss_functions_covers_conditional_generation():
def test_patch_loss_functions_does_not_touch_other_loss_types():
"""patch_loss_functions() must not overwrite unrelated loss types
(segmentation, detection, masked-LM, etc.) with the causal-LM kernel."""
"""patch_loss_functions() must not overwrite unrelated loss types with the causal-LM kernel."""
lu = pytest.importorskip("transformers.loss.loss_utils")
cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")
@ -574,12 +545,11 @@ def test_patch_loss_functions_does_not_touch_other_loss_types():
def test_accelerate_utils_imports_module_present():
"""``disable_broken_wandb`` + ``fix_trl_vllm_ascend`` (import_fixes.py
493-516, 1320-1372). Both reach into accelerate.utils.imports."""
"""``disable_broken_wandb`` + ``fix_trl_vllm_ascend`` both reach into
accelerate.utils.imports."""
pytest.importorskip("accelerate")
mod = pytest.importorskip("accelerate.utils.imports")
# is_wandb_available is the canonical representative -- disable_broken_wandb
# specifically targets it, so its absence breaks the patch.
# is_wandb_available is the canonical target of disable_broken_wandb.
assert hasattr(mod, "is_wandb_available"), (
"accelerate.utils.imports.is_wandb_available is gone; "
"disable_broken_wandb cannot patch the source module."
@ -587,7 +557,7 @@ def test_accelerate_utils_imports_module_present():
def test_accelerate_recursively_apply_empty_logits_patch():
"""Verify patch_accelerate_recursively_apply overrides recursively_apply to bypass EmptyLogits."""
"""patch_accelerate_recursively_apply overrides recursively_apply to bypass EmptyLogits."""
pytest.importorskip("accelerate")
import accelerate.utils.operations as acc_ops
@ -604,7 +574,7 @@ def test_accelerate_recursively_apply_empty_logits_patch():
def test_accelerate_gather_empty_logits_debug_mode_patch():
"""Verify gather and broadcast bypass EmptyLogits when debug mode is enabled."""
"""gather and broadcast bypass EmptyLogits when debug mode is enabled."""
pytest.importorskip("accelerate")
from accelerate.state import PartialState, DistributedType
import accelerate.utils.operations as acc_ops
@ -618,7 +588,7 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
e = EmptyLogits()
patch_accelerate_recursively_apply()
# Enable debug mode and mock distributed state
# Enable debug mode and mock a 2-process distributed state
state = PartialState()
orig_debug = state.debug
orig_dist_type = state.distributed_type
@ -628,11 +598,9 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
state.distributed_type = DistributedType.MULTI_GPU
state.num_processes = 2
# Mock gather_object to return [obj] * num_processes
def mock_gather_object(obj, *args, **kwargs):
return [obj] * state.num_processes
# Mock _gpu_gather to recursively apply replication of tensors
def mock_gpu_gather(tensor, *args, **kwargs):
def _gather_one(t):
if t.ndim == 0:
@ -641,7 +609,6 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
return acc_ops.recursively_apply(_gather_one, tensor, error_on_other_type = True)
# Mock _gpu_broadcast to return data unchanged
def mock_gpu_broadcast(data, *args, **kwargs):
return data
@ -657,32 +624,30 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
side_effect = mock_gpu_broadcast,
),
):
# 1. Top-level EmptyLogits should gather correctly (returns e)
# Top-level EmptyLogits gathers to itself
res = acc_ops.gather(e)
assert res is e
# 2. Nested EmptyLogits alone
# Nested EmptyLogits
res_nested = acc_ops.gather([e])
assert isinstance(res_nested, list) and res_nested[0] is e
# 3. Mixed payload with real tensor and EmptyLogits
# Real tensor should be gathered (concatenated across processes).
# Tensors must live on state.device or the debug-mode device
# check fails on GPU machines.
# Mixed payload: real tensor gets gathered, EmptyLogits passes through.
# Tensor must live on state.device or debug-mode device check fails on GPUs.
real_tensor = torch.tensor([42], device = state.device)
payload = {"labels": real_tensor, "logits": e}
res_mixed = acc_ops.gather(payload)
assert isinstance(res_mixed, dict)
assert res_mixed["logits"] is e
# Since num_processes = 2, it should be gathered to [42, 42]
# num_processes = 2 -> gathered to [42, 42]
assert torch.equal(res_mixed["labels"], torch.tensor([42, 42], device = state.device))
# 4. Broadcast with EmptyLogits
# Broadcast with EmptyLogits
res_broadcast = acc_ops.broadcast(e)
assert res_broadcast is e
# 5. Mixed payload with broadcast
# Mixed payload broadcast
res_broadcast_mixed = acc_ops.broadcast(payload)
assert isinstance(res_broadcast_mixed, dict)
assert res_broadcast_mixed["logits"] is e
@ -722,13 +687,12 @@ def test_accelerate_find_device_skips_empty_logits():
patch_accelerate_recursively_apply()
tensor = torch.tensor([1.0])
# Sentinel first must not stop the search before the real tensor
# Leading sentinel must not stop the search before the real tensor
assert acc_ops.find_device({"logits": EmptyLogits(), "labels": tensor}) == tensor.device
# Tensor-free payloads without the sentinel keep returning None
# (AlignDevicesHook relies on None to skip output device moves)
# Tensor-free payloads keep returning None (AlignDevicesHook needs it to skip moves)
assert acc_ops.find_device({"a": 1}) is None
# Sentinel-only payloads fall back to the current device so that
# debug mode find_device(...).type does not raise AttributeError
# Sentinel-only payloads fall back to current device so debug-mode
# find_device(...).type doesn't raise AttributeError
assert acc_ops.find_device(EmptyLogits()) == PartialState().device
@ -750,10 +714,8 @@ def test_accelerate_patch_wired_into_gpu_init():
def test_bitsandbytes_rocm_detection_helpers_recognizable():
"""``fix_bitsandbytes_rocm_arch_detection`` swaps bnb's ROCm helpers
only when they shell out via subprocess and never consult torch device
props; a third shape is declined by design, silently restoring Windows
ROCm noise. Fail so the sniff gets updated. Reads source, no import."""
"""``fix_bitsandbytes_rocm_arch_detection``: the source sniff only patches
bnb's ROCm helpers in recognized shapes; fail (don't import) when it drifts."""
spec = importlib.util.find_spec("bitsandbytes")
if spec is None:
pytest.skip("bitsandbytes not installed -- nothing to drift-check.")

26
tests/test_loader_glob_skip.py
View file

@ -1,10 +1,4 @@
"""Tests that HfFileSystem().glob() is skipped when is_model or is_peft is False.
The glob calls in FastLanguageModel.from_pretrained and FastModel.from_pretrained
exist solely to detect repos with both config.json and adapter_config.json. When
either AutoConfig or PeftConfig fails to load, the glob cannot find both files,
so calling it is redundant and risks hanging on slow networks.
"""
"""HfFileSystem().glob() is skipped when is_model or is_peft is False (redundant, risks hanging on slow networks)."""
import os
import unittest
@ -12,7 +6,7 @@ from unittest.mock import MagicMock, patch
class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
"""Verify HfFileSystem.glob is not called when is_model or is_peft is False."""
"""glob is not called when is_model or is_peft is False."""
def _run_both_exist_block(
self,
@ -22,11 +16,7 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
model_name,
is_local_dir = False,
):
"""Simulate the both_exist detection block from loader.py.
This mirrors the exact logic at lines 500-517 / 1276-1292 of loader.py.
Returns (both_exist, glob_called).
"""
"""Mirror loader.py's both_exist detection block; returns (both_exist, glob_called)."""
from unittest.mock import MagicMock
both_exist = (is_model and is_peft) and not supports_llama32
@ -37,10 +27,9 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
]
)
# This mirrors the guarded block in loader.py
if supports_llama32 and is_model and is_peft:
if is_local_dir:
# Local path branch — would use os.path.exists in real code
# Local path branch (os.path.exists in real code)
both_exist = True # simulate both files present locally
else:
files = glob_mock(f"{model_name}/*.json")
@ -90,7 +79,7 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
model_name = "org/some-model",
)
self.assertFalse(glob_called, "glob should not be called when SUPPORTS_LLAMA32=False")
# both_exist is set by the old-style check: (is_model and is_peft) and not SUPPORTS_LLAMA32
# both_exist set by the old-style check: (is_model and is_peft) and not SUPPORTS_LLAMA32
self.assertTrue(both_exist)
# --- Cases where glob SHOULD be called ---
@ -120,17 +109,16 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
class TestLoaderSourceHasGuard(unittest.TestCase):
"""Verify the actual loader.py source code has the is_model/is_peft guard."""
"""The actual loader.py source has the is_model/is_peft guard."""
def test_loader_source_has_guard(self):
"""Check that both SUPPORTS_LLAMA32 checks in loader.py include is_model and is_peft."""
"""Both SUPPORTS_LLAMA32 checks in loader.py include is_model and is_peft."""
loader_path = os.path.join(
os.path.dirname(__file__), os.pardir, "unsloth", "models", "loader.py"
)
with open(loader_path) as f:
source = f.read()
# Find all lines with the SUPPORTS_LLAMA32 check near glob usage
lines = source.splitlines()
guard_lines = [
line.strip()

10
tests/test_model_registry.py
View file

@ -1,8 +1,4 @@
"""Test model registration methods.
Registers each model set and checks the registered model ids exist on the
Hugging Face Hub.
"""
"""Register each model set and check the registered ids exist on the HF Hub."""
from dataclasses import dataclass
@ -58,7 +54,6 @@ TestParams = [
]
# Test that model registration methods register respective models
@pytest.mark.parametrize("model_test_param", TestParams, ids = lambda param: param.name)
def test_model_registration(model_test_param: ModelTestParam):
MODEL_REGISTRY.clear()
@ -77,8 +72,7 @@ def test_all_model_registration():
def test_quant_type():
# Test that the quant_type is correctly set for model paths
# NOTE: for models registered under org="unsloth" with QuantType.NONE aliases QuantType.UNSLOTH
# NOTE: for org="unsloth" models, QuantType.NONE aliases QuantType.UNSLOTH
dynamic_quant_models = search_models(quant_types = [QuantType.UNSLOTH])
assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models)
quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH]

4
tests/test_multi_image_grpo_chunking.py
View file

@ -1,5 +1,5 @@
"""Static + behavioral checks for the multi-image GRPO chunking and
zoo compatibility guard in unsloth/models/rl_replacements.py."""
"""Static + behavioral checks for multi-image GRPO chunking and the zoo
compatibility guard in unsloth/models/rl_replacements.py."""
from __future__ import annotations

12
tests/test_nvfp4_quant_load.py
View file

@ -12,19 +12,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Test cases for NVFP4 / compressed-tensors model loading.
Ensures that models with non-bitsandbytes quantization configs
don't conflict with Unsloth's default load_in_4bit=True behavior.
Uses synthetic config objects (no network access) so this suite
runs offline in CI where tests/security/conftest.py blocks socket
connections.
"""NVFP4 / compressed-tensors loading: non-bitsandbytes quant configs must not conflict with
load_in_4bit=True. Uses synthetic configs (no network) so it runs offline in CI.
"""
from types import SimpleNamespace
# Import unsloth first to set UNSLOTH_IS_PRESENT env var
# Import unsloth first to set UNSLOTH_IS_PRESENT env var.
import unsloth
from unsloth_zoo.utils import get_quant_type
from unsloth.models.loader_utils import check_and_disable_bitsandbytes_loading

26
tests/test_public_api_surface.py
View file

@ -11,12 +11,8 @@
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
"""Public-API surface drift detectors for unsloth itself.
Companion to tests/test_import_fixes_drift.py (which catches THIRD-PARTY drift):
this catches drift in unsloth's OWN public surface -- the top symbols and
classmethods the unslothai/notebooks tree calls -- so a rename or dropped kwarg
fires DRIFT DETECTED here before it reaches users.
"""Drift detectors for unsloth's OWN public surface (top symbols/classmethods the
unslothai/notebooks tree calls), so a rename or dropped kwarg fires DRIFT DETECTED here.
Call-site counts measured against unslothai/notebooks @ main:
FastLanguageModel.from_pretrained 506
@ -28,10 +24,6 @@ Call-site counts measured against unslothai/notebooks @ main:
FastVisionModel.for_training 60
FastModel.from_pretrained 103
FastModel.get_peft_model 67
Mirrors the drift-detector skeleton: gate on ``pytest.importorskip("unsloth")``,
assert the healthy shape, and ``pytest.fail("DRIFT DETECTED: ...")`` (never skip)
on regression so the matrix cell goes red.
"""
from __future__ import annotations
@ -50,9 +42,7 @@ def _signature_param_names(callable_obj) -> set[str]:
def _accepts(callable_obj, kwargs: set[str]) -> tuple[bool, set[str]]:
"""True if every name in ``kwargs`` is either a named parameter on
``callable_obj`` OR the callable's signature has a ``**kwargs``
catch-all. Returns (ok, missing_set)."""
"""(ok, missing): True if every kwarg is a named param or the signature has **kwargs."""
try:
sig = inspect.signature(callable_obj)
except (TypeError, ValueError):
@ -65,8 +55,7 @@ def _accepts(callable_obj, kwargs: set[str]) -> tuple[bool, set[str]]:
return (not missing), missing
# FastLanguageModel: headline class (506 from_pretrained + 370 for_inference +
# 304 get_peft_model call sites).
# FastLanguageModel: headline class.
def test_fast_language_model_class_present():
@ -119,7 +108,7 @@ def test_fast_language_model_for_inference_callable():
)
# FastVisionModel: 183 + 176 + 99 + 60 call sites across vision notebooks.
# FastVisionModel.
def test_fast_vision_model_class_and_methods():
@ -156,7 +145,7 @@ def test_fast_vision_model_get_peft_model_vision_kwargs():
)
# FastModel: modern unified entry point. 103 + 67 call sites.
# FastModel: modern unified entry point.
def test_fast_model_class_and_methods():
@ -190,8 +179,7 @@ def test_fast_model_from_pretrained_kwargs():
def test_is_bf16_supported_or_alias_callable():
"""48 notebook import sites for is_bf16_supported plus 8 for the
legacy is_bfloat16_supported alias. Either must remain importable."""
"""is_bf16_supported or the legacy is_bfloat16_supported alias must remain importable."""
unsloth = pytest.importorskip("unsloth")
has_new = callable(getattr(unsloth, "is_bf16_supported", None))
has_old = callable(getattr(unsloth, "is_bfloat16_supported", None))

67
tests/test_raw_text.py
View file

@ -1,8 +1,5 @@
#!/usr/bin/env python3
"""
Minimal test for raw text training implementation.
Tests basic functionality without heavy dependencies.
"""
"""Minimal test for raw text training, without heavy dependencies."""
import sys
import os
@ -11,7 +8,7 @@ from pathlib import Path
import importlib.util
# Mock the datasets module since it's not installed
# Mock the datasets module (not installed).
class MockDataset:
def __init__(self, data_dict):
self.data = data_dict
@ -22,10 +19,10 @@ class MockDataset:
def __getitem__(self, idx):
if isinstance(idx, str):
# Access columns by name, e.g. dataset['text']
# Column access, e.g. dataset['text'].
return self.data[idx]
elif isinstance(idx, int):
# Access individual rows by index
# Row access by index.
return {key: values[idx] for key, values in self.data.items()}
else:
raise TypeError(f"Invalid index type: {type(idx)}")
@ -35,15 +32,14 @@ class MockDataset:
return cls(data_dict)
# Mock datasets module. __spec__ must be set so importlib.util.find_spec
# does not raise ValueError when transformers' import_utils probes for
# the real `datasets` package later in the test session.
# __spec__ must be set so importlib.util.find_spec doesn't raise ValueError when
# transformers' import_utils later probes for the real `datasets` package.
datasets_mock = type(sys)("datasets")
datasets_mock.__spec__ = importlib.util.spec_from_loader("datasets", loader = None)
datasets_mock.Dataset = MockDataset
sys.modules["datasets"] = datasets_mock
# Import the raw_text module directly to avoid unsloth/__init__.py dependencies
# Import raw_text directly to avoid unsloth/__init__.py dependencies.
current_dir = os.path.dirname(__file__)
raw_text_path = os.path.join(os.path.dirname(current_dir), "unsloth", "dataprep", "raw_text.py")
@ -58,11 +54,10 @@ TextPreprocessor = raw_text_module.TextPreprocessor
def test_raw_text_loader():
"""Test basic RawTextDataLoader functionality."""
# Mock tokenizer for testing
class MockTokenizer:
def __init__(self):
self.eos_token = "</s>"
self.eos_token_id = 2 # Mock EOS token ID
self.eos_token_id = 2
def __call__(
self,
@ -74,7 +69,7 @@ def test_raw_text_loader():
token_ids = list(range(len(words)))
if return_tensors == "pt":
# Mock tensor-like object
class MockTensor:
def __init__(self, data):
self.data = data
@ -98,23 +93,21 @@ def test_raw_text_loader():
):
return " ".join([f"word_{i}" for i in token_ids])
# Create test file
test_content = "This is a test file for raw text training. " * 10
with tempfile.NamedTemporaryFile(mode = "w", suffix = ".txt", delete = False) as f:
f.write(test_content)
test_file = f.name
try:
# Test loader
tokenizer = MockTokenizer()
loader = RawTextDataLoader(tokenizer, chunk_size = 5, stride = 2)
# Test loading with text output (legacy mode)
# Text output (legacy mode).
text_dataset = loader.load_from_file(test_file, return_tokenized = False)
assert len(text_dataset) > 0, "Should create at least one chunk"
assert "text" in text_dataset.column_names, "Dataset should have 'text' column"
# Test loading with tokenized output (new efficient mode)
# Tokenized output (new efficient mode).
tokenized_dataset = loader.load_from_file(test_file, return_tokenized = True)
assert len(tokenized_dataset) > 0, "Should create at least one tokenized chunk"
assert (
@ -124,7 +117,6 @@ def test_raw_text_loader():
"attention_mask" in tokenized_dataset.column_names
), "Dataset should have 'attention_mask' column"
# Verify tokenized data structure
first_sample = tokenized_dataset[0]
assert isinstance(first_sample["input_ids"], list), "input_ids should be a list"
assert isinstance(first_sample["attention_mask"], list), "attention_mask should be a list"
@ -132,11 +124,11 @@ def test_raw_text_loader():
first_sample["attention_mask"]
), "input_ids and attention_mask should have same length"
# Verify labels field exists (for causal LM training)
# labels field (for causal LM training).
assert "labels" in tokenized_dataset.column_names, "Dataset should have 'labels' column"
assert first_sample["labels"] == first_sample["input_ids"], "labels should match input_ids"
# Test constructor validation
# Constructor validation.
try:
bad_loader = RawTextDataLoader(tokenizer, chunk_size = 0, stride = 2)
assert False, "Should raise ValueError for chunk_size=0"
@ -149,7 +141,7 @@ def test_raw_text_loader():
except ValueError as e:
assert "stride" in str(e) and "chunk_size" in str(e)
# Test preprocessor
# Preprocessor.
preprocessor = TextPreprocessor()
clean_text = preprocessor.clean_text(" messy text \n\n\n ")
assert "messy text" in clean_text, "Should clean text properly"
@ -158,10 +150,8 @@ def test_raw_text_loader():
paragraph_text == "Line 1\n\nLine 2"
), "Should preserve paragraph breaks while normalizing newlines"
# Non-ASCII horizontal whitespace separators (NBSP, thin space,
# ideographic space, narrow NBSP, em space, vertical tab, form feed)
# should be normalized to a single ASCII space, not deleted, otherwise
# adjacent words get silently fused together on HTML/PDF/OCR inputs.
# Non-ASCII horizontal whitespace (NBSP, thin/em/ideographic space, VT, FF) must
# normalize to one ASCII space, not be deleted, or adjacent words fuse on HTML/PDF/OCR input.
unicode_whitespace_cases = [
("hello\u00a0world", "hello world"),
("hello\u202fworld", "hello world"),
@ -176,38 +166,32 @@ def test_raw_text_loader():
f"Should normalize Unicode/control whitespace to a single space " f"for {raw!r}"
)
# Mixed paragraph + Unicode whitespace realistic input
# Mixed paragraph + Unicode whitespace.
mixed = preprocessor.clean_text("Section\u00a01\r\n\r\nBody\ftext\u202fhere")
assert (
mixed == "Section 1\n\nBody text here"
), "Should preserve paragraph breaks and normalize Unicode whitespace simultaneously"
# Tabs should collapse to a single space
# Tabs collapse to a single space.
assert preprocessor.clean_text("a\tb") == "a b"
assert preprocessor.clean_text("a\t\tb") == "a b"
# Spaces around newlines should be trimmed on both sides, even with
# multiple consecutive newlines
# Spaces around newlines trimmed on both sides, even across multiple newlines.
assert preprocessor.clean_text("foo \n\n bar") == "foo\n\nbar"
# Non-whitespace non-ASCII characters sitting between spaces should
# not leave an interior double space after being stripped. This
# guards the idempotence invariant too: without the extra collapse
# pass, "word1 (c) word2" first reduces to "word1 word2" and only
# becomes "word1 word2" on a second call.
# Stripping a non-ASCII char between spaces must not leave a double space
# (also guards idempotence: otherwise "word1 (c) word2" needs a second pass).
assert preprocessor.clean_text("word1 \u00a9 word2") == "word1 word2"
assert preprocessor.clean_text("a \u00e9 b") == "a b"
assert preprocessor.clean_text("prefix \U0001f600 suffix") == "prefix suffix"
# Stripping a non-ASCII character adjacent to a newline must not
# leave a stray leading/trailing space on the neighbouring line.
# Stripping a non-ASCII char adjacent to a newline must not leave a stray space.
assert preprocessor.clean_text("foo \u00e9\nbar") == "foo\nbar"
assert preprocessor.clean_text("foo\n\u00e9 bar") == "foo\nbar"
# The double-space collapse pass must not swallow a legitimate
# paragraph break when a non-ASCII char sits near it.
# The double-space collapse must not swallow a paragraph break near a non-ASCII char.
assert preprocessor.clean_text("a \u00a9\n\nb") == "a\n\nb"
# Idempotence: running clean_text twice should give the same result
# Idempotence: clean_text twice == once.
idempotent_inputs = [
" messy text \n\n\n ",
"Line 1\r\n\r\n\r\nLine 2",
@ -221,7 +205,7 @@ def test_raw_text_loader():
twice = preprocessor.clean_text(once)
assert once == twice, f"clean_text should be idempotent for {raw!r}"
# Test validation
# Validation.
stats = preprocessor.validate_dataset(text_dataset)
assert stats["total_samples"] > 0, "Should count samples"
assert "warnings" in stats, "Should include warnings"
@ -234,7 +218,6 @@ def test_raw_text_loader():
return False
finally:
# Cleanup
os.unlink(test_file)

227
tests/test_studio_install_workspace_guard.py
View file

@ -1,8 +1,4 @@
"""install.sh / install.ps1 must refuse to rm -rf an existing
$STUDIO_HOME/unsloth_studio in env-override mode unless the directory
carries a Studio sentinel (share/studio.conf or bin/unsloth). Also
asserts studio/setup.ps1 has the matching writability probe that
setup.sh:417 already performs."""
"""install.sh/install.ps1 must refuse to rm -rf an existing Studio venv in env-mode without a sentinel."""
from __future__ import annotations
@ -16,10 +12,8 @@ INSTALL_PS1 = REPO_ROOT / "install.ps1"
SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1"
SETUP_SH = REPO_ROOT / "studio" / "setup.sh"
# Stubs for helpers the extracted install.sh guard block calls (`substep`,
# `_start_studio_venv_replacement`). Tests run the block in isolation, so a
# minimal `mv`-based replacement reproduces the observable effect (venv gone
# from $VENV_DIR after permitted cleanup) without the full rollback machinery.
# Stubs for helpers the extracted guard block calls; mv-based replacement reproduces the venv-gone
# effect without the full rollback machinery.
_INSTALL_GUARD_STUBS = (
"substep() { :; }\n"
"_start_studio_venv_replacement() {\n"
@ -29,9 +23,7 @@ _INSTALL_GUARD_STUBS = (
def _extract_install_sh_guard_block() -> str:
"""Pull the `if [ -x "$VENV_DIR/bin/python" ]; then ... fi` block out
of install.sh as a self-contained snippet. Stops at the first elif so
the block can be paired with a synthetic else and run in isolation."""
"""Extract install.sh's venv guard block (up to the first elif) as a self-contained snippet."""
src = INSTALL_SH.read_text()
m = re.search(
r'(if \[ -x "\$VENV_DIR/bin/python" \]; then\n.*?)elif \[ "\$_STUDIO_HOME_REDIRECT" != "env"',
@ -47,9 +39,7 @@ def _build_install_guard_script(
redirect: str,
block: str | None = None,
) -> str:
"""Build a self-contained bash script that exercises the extracted
guard block. Includes stubs for substep / _start_studio_venv_replacement
so the snippet runs without install.sh's full rollback machinery."""
"""Build a self-contained bash script exercising the extracted guard block (with helper stubs)."""
if block is None:
block = _extract_install_sh_guard_block()
return (
@ -153,10 +143,7 @@ def test_setup_ps1_has_writability_probe():
def test_env_mode_blocks_when_bin_unsloth_is_a_directory(tmp_path):
"""A bare directory at $STUDIO_HOME/bin/unsloth must NOT pass the
sentinel. The previous `-e` test accepted any path type, allowing an
unrelated workspace with sibling content under unsloth_studio plus
a directory at bin/unsloth to be wiped."""
"""A bare directory at bin/unsloth must NOT pass the sentinel (regression: `-e` accepted any type)."""
studio_home = tmp_path / "ws"
venv = studio_home / "unsloth_studio"
(venv / "bin").mkdir(parents = True)
@ -180,8 +167,7 @@ def test_env_mode_blocks_when_bin_unsloth_is_a_directory(tmp_path):
def test_env_mode_passes_when_bin_unsloth_is_a_symlink(tmp_path):
"""A symlink at $STUDIO_HOME/bin/unsloth (real installer artefact)
must still satisfy the sentinel after the leaf-only tightening."""
"""A symlink at bin/unsloth (real installer artefact) must still satisfy the sentinel."""
studio_home = tmp_path / "ws"
venv = studio_home / "unsloth_studio"
(venv / "bin").mkdir(parents = True)
@ -206,8 +192,7 @@ def test_env_mode_passes_when_bin_unsloth_is_a_symlink(tmp_path):
def test_install_ps1_sentinel_uses_pathtype_leaf():
"""The Test-Path checks that gate Remove-Item $VenvDir must use
-PathType Leaf so a directory at the sentinel path cannot satisfy them."""
"""Remove-Item $VenvDir gate must use -PathType Leaf so a sentinel-path directory cannot satisfy it."""
src = INSTALL_PS1.read_text()
block_start = src.index("if (Test-Path -LiteralPath $VenvPython)")
block = src[block_start : block_start + 2000]
@ -220,10 +205,7 @@ def test_install_ps1_sentinel_uses_pathtype_leaf():
def test_setup_ps1_stale_venv_has_env_mode_guard():
"""studio/setup.ps1 stale-venv rebuild branch must mirror install.ps1:
refuse to Remove-Item $VenvDir under custom-root mode unless the root
carries a Studio sentinel (in-VENV marker, share\\studio.conf, or
bin\\unsloth.exe leaf)."""
"""setup.ps1 stale-venv branch must gate Remove-Item $VenvDir on a custom-root Studio sentinel."""
src = SETUP_PS1.read_text()
idx = src.index("Stale venv detected")
block = src[idx : idx + 1500]
@ -243,10 +225,7 @@ def test_setup_ps1_stale_venv_has_env_mode_guard():
def test_setup_sh_prebuilt_llama_cpp_has_ownership_guard():
"""studio/setup.sh prebuilt llama.cpp path must call
_assert_studio_owned_or_absent before invoking install_llama_prebuilt.py
so an unrelated $UNSLOTH_STUDIO_HOME/llama.cpp is not displaced by
the helper's os.replace()."""
"""setup.sh prebuilt llama.cpp path must _assert_studio_owned_or_absent before install_llama_prebuilt.py."""
src = SETUP_SH.read_text()
idx = src.index("installing prebuilt llama.cpp...")
block = src[idx : idx + 2000]
@ -260,8 +239,7 @@ def test_setup_sh_prebuilt_llama_cpp_has_ownership_guard():
def test_setup_ps1_prebuilt_llama_cpp_has_ownership_guard():
"""Mirror check for studio/setup.ps1: prebuilt llama.cpp path must
call Assert-StudioOwnedOrAbsent before invoking install_llama_prebuilt.py."""
"""setup.ps1 prebuilt llama.cpp path must Assert-StudioOwnedOrAbsent before install_llama_prebuilt.py."""
src = SETUP_PS1.read_text()
idx = src.index("installing prebuilt llama.cpp bundle (preferred path)")
block = src[idx : idx + 2000]
@ -277,10 +255,7 @@ def test_setup_ps1_prebuilt_llama_cpp_has_ownership_guard():
def test_env_mode_passes_when_venv_marker_present(tmp_path):
"""install.sh env-mode guard must accept the in-VENV
.unsloth-studio-owned marker as a primary sentinel so a partial
install (uv venv created, sentinels not yet written) is recoverable
by re-running install.sh."""
"""install.sh env-mode guard must accept the in-VENV .unsloth-studio-owned marker as a sentinel."""
studio_home = tmp_path / "ws"
res = _run_install_guard(studio_home, redirect = "env", create_venv_marker = True)
assert res.returncode == 0, (
@ -291,11 +266,7 @@ def test_env_mode_passes_when_venv_marker_present(tmp_path):
def test_env_mode_blocks_when_bin_unsloth_is_symlink_to_directory(tmp_path):
"""install.sh env-mode guard must NOT accept a symlink-to-directory at
bin/unsloth as a Studio sentinel. Iter1's standalone -L test let any
symlink (including symlinks to dirs and broken symlinks) bypass the
guard; iter2 dropped that test so only -f (file or symlink-to-file)
counts."""
"""install.sh guard must reject a symlink-to-directory at bin/unsloth; only -f (file/symlink-to-file) counts."""
studio_home = tmp_path / "ws"
venv = studio_home / "unsloth_studio"
(venv / "bin").mkdir(parents = True)
@ -347,9 +318,7 @@ def test_env_mode_blocks_when_bin_unsloth_is_broken_symlink(tmp_path):
def test_install_sh_writes_venv_marker_after_uv_venv():
"""install.sh must write the .unsloth-studio-owned marker into
$VENV_DIR right after `uv venv` succeeds so the env-mode deletion
guard accepts it on the next install run."""
"""install.sh must write .unsloth-studio-owned into $VENV_DIR right after `uv venv` succeeds."""
src = INSTALL_SH.read_text()
create_idx = src.index('run_install_cmd "create venv" uv venv "$VENV_DIR"')
tail = src[create_idx : create_idx + 600]
@ -359,8 +328,7 @@ def test_install_sh_writes_venv_marker_after_uv_venv():
def test_install_ps1_writes_venv_marker_after_uv_venv():
"""install.ps1 must write the .unsloth-studio-owned marker into
$VenvDir after `uv venv` succeeds."""
"""install.ps1 must write .unsloth-studio-owned into $VenvDir after `uv venv` succeeds."""
src = INSTALL_PS1.read_text()
venv_create = src.index("uv venv $VenvDir --python")
tail = src[venv_create : venv_create + 1500]
@ -370,8 +338,7 @@ def test_install_ps1_writes_venv_marker_after_uv_venv():
def test_install_ps1_guard_accepts_venv_marker():
"""install.ps1 env-mode guard must accept the in-VENV
.unsloth-studio-owned marker as a primary sentinel."""
"""install.ps1 env-mode guard must accept the in-VENV .unsloth-studio-owned marker as a sentinel."""
src = INSTALL_PS1.read_text()
block_start = src.index("if (Test-Path -LiteralPath $VenvPython)")
block = src[block_start : block_start + 2000]
@ -381,11 +348,7 @@ def test_install_ps1_guard_accepts_venv_marker():
def test_setup_helpers_gate_on_canonical_custom_root():
"""Both _assert_studio_owned_or_absent (setup.sh) and
Assert-StudioOwnedOrAbsent (setup.ps1) must gate on a canonical
custom-vs-legacy comparison so an explicit override that resolves
to the legacy default does not trip the guard for pre-PR T5
sidecar venvs or llama.cpp dirs."""
"""setup.sh/setup.ps1 ownership guards must gate on a canonical custom-vs-legacy root comparison."""
sh_src = SETUP_SH.read_text()
sh_idx = sh_src.index("_assert_studio_owned_or_absent() {")
sh_func = sh_src[sh_idx : sh_idx + 600]
@ -410,9 +373,7 @@ def test_setup_helpers_gate_on_canonical_custom_root():
def test_setup_ps1_inplace_git_sync_marks_studio_owned():
"""setup.ps1 in-place git-sync branch (when $LlamaCppDir/.git exists)
must call Mark-StudioOwned after a successful sync so a later prebuilt
update path's Assert-StudioOwnedOrAbsent does not exit."""
"""setup.ps1 in-place git-sync branch must Mark-StudioOwned after a successful sync."""
src = SETUP_PS1.read_text()
inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")')
# The in-place branch ends just before the temp-dir clone branch.
@ -427,10 +388,7 @@ def test_setup_ps1_inplace_git_sync_marks_studio_owned():
def test_setup_ps1_inplace_git_sync_asserts_studio_owned_before_mutation():
"""setup.ps1 in-place git-sync branch must call Assert-StudioOwnedOrAbsent
BEFORE any destructive git operation (remote set-url, checkout -B, clean
-fdx). Asymmetric to the prebuilt path and the temp-dir-swap path which
both guard."""
"""setup.ps1 in-place git-sync must Assert-StudioOwnedOrAbsent before any destructive git op."""
src = SETUP_PS1.read_text()
inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")')
clone_idx = src.index("Cloning llama.cpp @", inplace_idx)
@ -473,8 +431,7 @@ def _run_check_health(expected_root_id: str, response_json: str) -> int:
def test_check_health_accepts_matching_studio_root_id():
"""Hex digest baked at install time matches the backend's
/api/health studio_root_id -- launcher attaches to its own backend."""
"""Matching baked studio_root_id lets the launcher attach to its own backend."""
expected_id = "a" * 64
rc = _run_check_health(
expected_id,
@ -484,8 +441,7 @@ def test_check_health_accepts_matching_studio_root_id():
def test_check_health_rejects_mismatched_studio_root_id():
"""Different install root → different sha256 → reject. Workspace
isolation: launcher A must not open Studio B running on the same port."""
"""Mismatched studio_root_id rejects attach (workspace isolation across same-port Studios)."""
expected_id = "a" * 64
other_id = "b" * 64
rc = _run_check_health(
@ -496,8 +452,7 @@ def test_check_health_rejects_mismatched_studio_root_id():
def test_check_health_rejects_missing_studio_root_id_field():
"""A backend that omits studio_root_id (older or non-conforming) must
not be attached to when an expected id is baked into the launcher."""
"""A backend omitting studio_root_id must not be attached when an expected id is baked in."""
expected_id = "a" * 64
rc = _run_check_health(
expected_id,
@ -507,9 +462,7 @@ def test_check_health_rejects_missing_studio_root_id_field():
def test_check_health_no_baked_id_accepts_any_healthy_backend():
"""If _EXPECTED_STUDIO_ROOT_ID is empty (e.g. install-time hash failed
to compute), the launcher falls back to the legacy contract and accepts
any healthy Unsloth backend."""
"""Empty _EXPECTED_STUDIO_ROOT_ID falls back to legacy contract: accept any healthy Unsloth backend."""
rc = _run_check_health(
"",
'{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"deadbeef"}',
@ -526,12 +479,7 @@ def test_check_health_rejects_non_unsloth_service():
def test_check_health_handles_arbitrary_id_token():
"""Iter3 used a raw shell match against the JSON-escaped studio_root,
which failed for paths containing `\\` or `"` (FastAPI emits `\\\\` and
`\\\"`). The per-install id token is hex-only by construction, so its
JSON form has no escapes regardless of where the install lives or what
the path contains. This test pins the round-trip on a fully arbitrary
64-char hex token."""
"""A fully arbitrary 64-char hex install id must round-trip cleanly (hex-only, no JSON escapes)."""
expected_id = "f0" + ("ed" * 31) # 64 hex chars, not derived from any path
rc = _run_check_health(
expected_id,
@ -541,8 +489,7 @@ def test_check_health_handles_arbitrary_id_token():
def test_install_ps1_test_studio_health_verifies_studio_root_id():
"""install.ps1 Test-StudioHealth must compare studio_root_id against
the install-time-baked $_ExpectedStudioRootId, not the runtime env var."""
"""install.ps1 Test-StudioHealth must compare studio_root_id against baked $_ExpectedStudioRootId."""
src = INSTALL_PS1.read_text()
fn_start = src.index("function Test-StudioHealth")
fn_end = src.index("\n}\n", fn_start) + 2
@ -554,11 +501,7 @@ def test_install_ps1_test_studio_health_verifies_studio_root_id():
def test_install_ps1_bakes_studio_root_id_into_launcher():
"""install.ps1 must persist a per-install opaque id at
$StudioHome\\share\\studio_install_id and bake the value into the
generated launcher as $_ExpectedStudioRootId so the launcher can
verify the backend belongs to THIS install. The id is generated
via a CSPRNG so /api/health does not leak the install path."""
"""install.ps1 must persist a CSPRNG id at share/studio_install_id and bake it as $_ExpectedStudioRootId."""
src = INSTALL_PS1.read_text()
assert "$_studioRootId" in src, "install.ps1 must compute $_studioRootId for the launcher"
assert (
@ -573,10 +516,7 @@ def test_install_ps1_bakes_studio_root_id_into_launcher():
def test_health_endpoint_exposes_studio_root_id_not_raw_path():
"""studio/backend/main.py /api/health must expose studio_root_id (a
hex digest) and NOT the raw studio_root path. Studio supports
`-H 0.0.0.0`; a /api/health that returns the raw install path
leaks username, home dir, workspace name, etc."""
"""/api/health must expose studio_root_id (hex digest), NOT the raw path (info disclosure on -H 0.0.0.0)."""
main_py = REPO_ROOT / "studio" / "backend" / "main.py"
src = main_py.read_text()
health_idx = src.index('@app.get("/api/health")')
@ -593,12 +533,7 @@ def test_health_endpoint_exposes_studio_root_id_not_raw_path():
def test_install_sh_bakes_studio_root_id_into_launcher():
"""install.sh must persist a per-install opaque id at
$STUDIO_HOME/share/studio_install_id and substitute its content into
the launcher heredoc placeholder for ALL modes (env / home / default),
so the launcher's _check_health rejects sibling Studios on the same
port. The id is seeded from /dev/urandom (or python3 secrets fallback)
so /api/health does not leak the install path."""
"""install.sh must persist the id at share/studio_install_id and bake it into the launcher for ALL modes."""
src = INSTALL_SH.read_text()
assert (
"_css_studio_root_id" in src
@ -618,12 +553,8 @@ def test_install_sh_bakes_studio_root_id_into_launcher():
def test_tauri_preflight_scrubs_studio_home_env():
"""All three Tauri CLI-spawn sites that lacked the scrub must now
env_remove UNSLOTH_STUDIO_HOME and STUDIO_HOME, mirroring
process.rs / install.rs / desktop_auth.rs / update.rs."""
# preflight was one .rs file; PR #5341 split it into a submodule dir. Read
# whichever shape is on disk so the guard survives reorgs, as long as the
# scrub calls live under studio/src-tauri/src/preflight*.
"""Tauri CLI-spawn sites must env_remove UNSLOTH_STUDIO_HOME and STUDIO_HOME."""
# PR #5341 split preflight into a submodule dir; read whichever shape is on disk.
preflight_root = REPO_ROOT / "studio" / "src-tauri" / "src"
preflight_paths = [
preflight_root / "preflight.rs",
@ -631,8 +562,7 @@ def test_tauri_preflight_scrubs_studio_home_env():
]
preflight = "\n".join(p.read_text() for p in preflight_paths if p.exists())
commands = (REPO_ROOT / "studio" / "src-tauri" / "src" / "commands.rs").read_text()
# Both functions (run_cli_probe + probe_cli_capability) must scrub.
# Count occurrences -- expect 2 in preflight (one per fn), 1 in commands.
# Expect 2 scrubs in preflight (run_cli_probe + probe_cli_capability), 1 in commands.
assert (
preflight.count('cmd.env_remove("UNSLOTH_STUDIO_HOME")') >= 2
), "preflight must scrub UNSLOTH_STUDIO_HOME in both run_cli_probe and probe_cli_capability"
@ -648,8 +578,7 @@ def test_tauri_preflight_scrubs_studio_home_env():
def test_install_sh_shim_uses_atomic_replace():
"""install.sh shim install must use ln -sfn for atomic replace; the
older `rm -f ...; ln -s ...` left a window where the shim was missing."""
"""install.sh shim install must use ln -sfn for atomic replace (rm+ln left a missing-shim window)."""
src = INSTALL_SH.read_text()
shim_idx = src.index('_shim_path="$_LOCAL_BIN/unsloth"')
block = src[shim_idx : shim_idx + 1500]
@ -662,11 +591,7 @@ def test_install_sh_shim_uses_atomic_replace():
def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback(tmp_path):
"""_create_shortcuts must seed new ids from /dev/urandom first (no
interpreter spawn cost on the install hot path) and fall back to
`python3 -c 'secrets.token_hex(32)'` only when urandom is unreadable.
Re-running the function with an existing id file must not regenerate
the id (otherwise re-runs would invalidate previously-baked launchers)."""
"""_create_shortcuts seeds ids from /dev/urandom (python3 secrets fallback) and is re-run idempotent."""
src = INSTALL_SH.read_text()
fn_start = src.index('_css_data_dir="$DATA_DIR"')
block = src[fn_start : fn_start + 3000]
@ -675,22 +600,19 @@ def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback(t
assert (
urandom_idx < py_fallback_idx
), "/dev/urandom must be tried before the python3 secrets fallback"
# The id file is checked for non-empty content before we generate; this is
# what makes re-runs idempotent.
# Non-empty id file check before generation is what makes re-runs idempotent.
assert (
'if [ ! -s "$_css_id_file" ]; then' in block
), "install.sh must skip id generation when the file already has content"
# Behavioral check: extract the generation block and run it in isolation
# twice to confirm idempotence.
# Behavioral check: run the generation block twice to confirm idempotence.
studio_home = tmp_path / "studio"
(studio_home / "share").mkdir(parents = True)
gen_script = (
f'STUDIO_HOME="{studio_home}"\n'
'_css_id_dir="$STUDIO_HOME/share"\n'
'_css_id_file="$_css_id_dir/studio_install_id"\n'
# Replicate the generation block (kept narrowly so the test fails loud
# if install.sh changes the surrounding contract).
# Replicate the generation block narrowly so it fails loud on contract drift.
"gen() {\n"
' if [ ! -s "$_css_id_file" ]; then\n'
' _css_new_id=$(od -An -N32 -tx1 /dev/urandom 2>/dev/null | tr -d " \\n")\n'
@ -714,9 +636,7 @@ def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback(t
def test_install_sh_create_shortcuts_fails_fast_when_no_entropy():
"""If neither /dev/urandom nor python3 is available, _create_shortcuts
must `return 1` instead of silently baking an empty studio_root_id
(which would disable the launcher's same-install discriminator)."""
"""With no entropy source, _create_shortcuts must `return 1` not bake an empty studio_root_id."""
src = INSTALL_SH.read_text()
fn_start = src.index('_css_data_dir="$DATA_DIR"')
block = src[fn_start : fn_start + 3000]
@ -732,9 +652,7 @@ def test_install_sh_create_shortcuts_fails_fast_when_no_entropy():
def test_install_sh_bakes_installed_is_env_mode_flag_in_launcher():
"""install.sh must bake the install-time mode (env vs default/home) into
the generated launcher so PORT_FILE / namespaced LOCK_DIR cannot be
flipped on by a sourced custom-root studio.conf in the user's shell."""
"""install.sh must bake the install-time mode into the launcher so a sourced studio.conf can't flip it."""
src = INSTALL_SH.read_text()
assert (
"_INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@'" in src
@ -749,10 +667,7 @@ def test_install_sh_bakes_installed_is_env_mode_flag_in_launcher():
def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env():
"""The launcher's PORT_FILE / namespaced LOCK_DIR must be gated on the
baked $_INSTALLED_IS_ENV_MODE flag, not the runtime $UNSLOTH_STUDIO_HOME.
Sourcing a custom-root studio.conf in shell must not flip a default-mode
launcher into env-mode behavior."""
"""Launcher PORT_FILE/LOCK_DIR must gate on baked $_INSTALLED_IS_ENV_MODE, not runtime $UNSLOTH_STUDIO_HOME."""
src = INSTALL_SH.read_text()
heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'")
heredoc_end = src.index("LAUNCHER_EOF\n", heredoc_start)
@ -769,7 +684,7 @@ def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env():
), "launcher must NOT gate PORT_FILE on runtime UNSLOTH_STUDIO_HOME"
def _run_launcher_gate(installed_flag: str, runtime_env: dict) -> str:
# Reproduce just the LOCK_DIR/PORT_FILE init block in isolation.
# Run the LOCK_DIR/PORT_FILE init block in isolation.
script = (
f"_INSTALLED_IS_ENV_MODE={installed_flag!r}\n"
"DATA_DIR=/tmp/test_data_dir\n"
@ -789,21 +704,18 @@ def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env():
return line[len("PORT_FILE=") :]
return ""
# default-mode install should NEVER set PORT_FILE, even if UNSLOTH_STUDIO_HOME leaks in.
# default-mode must keep PORT_FILE empty even if UNSLOTH_STUDIO_HOME leaks in.
assert (
_run_launcher_gate("false", {"UNSLOTH_STUDIO_HOME": "/tmp/leaked"}) == ""
), "default-mode launcher must keep PORT_FILE empty even with UNSLOTH_STUDIO_HOME in env"
# env-mode install should set PORT_FILE regardless of runtime env.
# env-mode must set PORT_FILE regardless of runtime env.
assert (
_run_launcher_gate("true", {}) == "/tmp/test_data_dir/studio.port"
), "env-mode launcher must set PORT_FILE based on baked DATA_DIR"
def test_main_py_studio_root_id_caches_at_module_load():
"""_studio_root_id() is called on every /api/health poll; the id is
stable for the lifetime of the process so it must be read once at
module load and re-used (avoids a hot-path filesystem probe and
protects against transient FS errors during health polling)."""
"""_studio_root_id() must read the id once at module load and reuse it (no per-poll FS/hash work)."""
main_py = (REPO_ROOT / "studio" / "backend" / "main.py").read_text()
assert (
"_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()" in main_py
@ -820,20 +732,13 @@ def test_main_py_studio_root_id_caches_at_module_load():
def test_main_py_read_studio_install_id_validates_hex_and_handles_missing(tmp_path, monkeypatch):
"""_read_studio_install_id reads $STUDIO_HOME/share/studio_install_id and
returns "" when the file is absent, empty, contains non-hex content, or
is the wrong length. "" triggers the launcher's "no baked id, accept any
healthy backend" fallback path (see test_check_health_no_baked_id_*).
Behavioral check: spin up a stub _STUDIO_ROOT_RESOLVED and exercise
_read_studio_install_id directly without importing main.py (which
pulls in heavy deps). Test the rejection rules verbatim."""
"""_read_studio_install_id returns "" for absent/empty/non-hex/wrong-length ids, else the token."""
import re
pattern = re.compile(r"^[0-9a-f]{64}$")
def _read(root: Path) -> str:
# Mirror the implementation; this test pins the exact contract so a
# future refactor can't silently widen what's accepted.
# Mirror the implementation to pin the exact accepted contract.
try:
token = (root / "share" / "studio_install_id").read_text().strip()
except (OSError, ValueError):
@ -866,17 +771,13 @@ def test_main_py_read_studio_install_id_validates_hex_and_handles_missing(tmp_pa
def test_llama_cpp_search_roots_handles_studio_root_oserror():
"""_find_llama_server_binary calls studio_root() which can raise
OSError or ValueError from Path.expanduser().resolve() (broken symlink,
null byte). The except clause must mirror sibling _kill_orphaned_servers
(which catches the same trio) so inference startup does not crash."""
"""_find_llama_server_binary must catch (ImportError, OSError, ValueError) from studio_root() like its sibling."""
llama_cpp = (
REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py"
).read_text()
def _method_body(name: str) -> str:
# Whole method body (def to next sibling def), so the check survives the
# function growing past any fixed-size window.
# Whole method body (def to next sibling def) so the check survives growth.
start = llama_cpp.index(f"def {name}")
indent = " " * (start - llama_cpp.rfind("\n", 0, start) - 1)
nxt = llama_cpp.find(f"\n{indent}def ", start + 1)
@ -891,31 +792,21 @@ def test_llama_cpp_search_roots_handles_studio_root_oserror():
def test_install_sh_install_id_survives_symlinked_studio_home(tmp_path):
"""End-to-end behavioral check: when $STUDIO_HOME is reached via a
symlinked parent (e.g. symlinked $HOME on Linux, junctioned %USERPROFILE%
on Windows), install.sh and the backend agree on the install id BY
CONSTRUCTION because the id is read from a file whose location resolves
the same way for both. The previous sha256(canonical_path) scheme
required `cd -P/pwd -P` and Path.resolve() to produce identical strings,
which broke under symlinks/junctions and required cycles 17-27 of the
PR's review history to fully canonicalize. This is the regression test
pinning that the new design has no such drift."""
"""Regression: install id read from a file (not sha256(canonical_path)) agrees under a symlinked $STUDIO_HOME."""
real = tmp_path / "realhome"
real.mkdir()
link = tmp_path / "linkhome"
link.symlink_to(real)
studio_home = real / ".unsloth" / "studio"
(studio_home / "share").mkdir(parents = True)
# Write a stub install id at the canonical location.
valid_id = "ab12" * 16
(studio_home / "share" / "studio_install_id").write_text(valid_id)
# Read back via canonical and symlinked path; both must see the SAME
# content so install.sh's cat and the backend's read_text agree.
# Canonical and symlinked paths must see the SAME content (cat and read_text agree).
raw_via_link = link / ".unsloth" / "studio" / "share" / "studio_install_id"
raw_direct = studio_home / "share" / "studio_install_id"
assert raw_via_link.read_text() == valid_id
assert raw_direct.read_text() == valid_id
# And install.sh's `cat` would see the same.
# install.sh's `cat` sees the same.
import subprocess as _sp
res = _sp.run(["cat", str(raw_via_link)], capture_output = True, text = True)
@ -924,10 +815,7 @@ def test_install_sh_install_id_survives_symlinked_studio_home(tmp_path):
def test_install_sh_substitutes_root_id_before_data_dir():
"""The two-stage sed substitution must bake @@STUDIO_ROOT_ID@@ /
@@INSTALLED_IS_ENV_MODE@@ first (non-user-controlled), then @@DATA_DIR@@
(user-controlled). A custom $DATA_DIR containing the literal text
@@STUDIO_ROOT_ID@@ must not be mutated by the global root-id sed pass."""
"""sed must bake the non-user-controlled placeholders before @@DATA_DIR@@ so a crafted $DATA_DIR isn't mutated."""
src = INSTALL_SH.read_text()
root_id_idx = src.index("s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g")
env_mode_idx = src.index("s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g")
@ -942,10 +830,7 @@ def test_install_sh_substitutes_root_id_before_data_dir():
def test_install_sh_root_id_pass_does_not_mutate_user_data_dir(tmp_path):
"""Behavioral subprocess test: a $DATA_DIR containing the literal text
`@@STUDIO_ROOT_ID@@` must not be mutated when the placeholder pass runs
first; only the actual placeholder occurrences in the launcher template
are replaced."""
"""A $DATA_DIR containing the literal @@STUDIO_ROOT_ID@@ must survive the placeholder-first sed passes."""
src = INSTALL_SH.read_text()
heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'")
heredoc_body_start = src.index("\n", heredoc_start) + 1
@ -953,7 +838,7 @@ def test_install_sh_root_id_pass_does_not_mutate_user_data_dir(tmp_path):
template = src[heredoc_body_start:heredoc_body_end]
launcher_path = tmp_path / "launch.sh"
launcher_path.write_text(template)
# Run the iter6 sed order: root-id first, then data-dir.
# sed order: root-id first, then data-dir.
weird_data_dir = "/tmp/with-@@STUDIO_ROOT_ID@@/share"
root_id = "deadbeef" * 8
is_env = "true"
@ -977,11 +862,7 @@ sed "s|@@DATA_DIR@@|$_sed_safe|g" "{launcher_path}" > "{launcher_path}.tmp" \\
def test_install_ps1_install_id_file_layout_matches_backend_read_path():
"""install.ps1 must write the id at $StudioHome\\share\\studio_install_id
so the backend (studio/backend/main.py:_read_studio_install_id) can find
it via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" without
mode-specific path knowledge. Persistence-across-runs is enforced by the
pre-write Test-Path check."""
"""install.ps1 must write the id at share/studio_install_id where the backend reads it, idempotently."""
src = INSTALL_PS1.read_text()
id_idx = src.index('$_studioIdDir = Join-Path $StudioHome "share"')
context = src[id_idx : id_idx + 1500]

30
tests/test_studio_root_resilience.py
View file

@ -1,13 +1,4 @@
"""Resilience checks for Studio install-root inference under hostile
filesystem conditions:
- _infer_studio_home_from_venv must NOT propagate PermissionError /
OSError out through studio_root() (it would crash module import in
run.py / main.py / transformers_version.py / model_config.py).
- _kill_orphaned_servers must catch (ImportError, OSError, ValueError)
on the studio_root() probe so a transient resolve / sentinel failure
cannot crash server startup.
- _find_llama_server_binary must keep the custom-root in search_roots
when the inner resolve() comparison itself fails."""
"""Studio install-root inference must not crash under hostile filesystem conditions (PermissionError/OSError swallowed; custom root kept when resolve() fails)."""
from __future__ import annotations
@ -48,9 +39,7 @@ def test_infer_studio_home_swallows_permission_error(tmp_path, monkeypatch):
def test_studio_root_does_not_crash_on_permission_error(tmp_path, monkeypatch):
"""studio_root() must remain callable even when the venv inference
encounters a restricted filesystem; it should fall through to the
legacy default."""
"""studio_root() falls through to the legacy default on a restricted filesystem."""
candidate = tmp_path / "fake_root"
venv = candidate / "unsloth_studio"
venv.mkdir(parents = True)
@ -65,19 +54,15 @@ def test_studio_root_does_not_crash_on_permission_error(tmp_path, monkeypatch):
def test_kill_orphan_catches_oserror_from_studio_root():
"""_kill_orphaned_servers must catch (ImportError, OSError, ValueError)
on the studio_root() probe specifically; the sister function
_find_llama_server_binary uses the same broader catch on its own probe."""
"""_kill_orphaned_servers must catch (ImportError, OSError, ValueError) on the studio_root() probe."""
src = LLAMA_CPP.read_text()
fn_start = src.index("def _kill_orphaned_servers")
fn_body = src[fn_start : fn_start + 4000]
# The studio_root() probe in this fn is the one that imports as `_sr`
# and assigns `_resolved_sr = _sr()`. Find the except that closes it.
# The studio_root() probe imports as `_sr` and assigns `_resolved_sr = _sr()`.
probe_idx = fn_body.index("storage_roots import studio_root as _sr")
# The matching except is the next `except ...:` after the inner
# OSError/ValueError block that wraps resolve().
# The matching except is the next one after the inner resolve() block.
after = fn_body[probe_idx:]
# Skip over the inner `except (OSError, ValueError):` that wraps resolve().
# Skip the inner `except (OSError, ValueError):` that wraps resolve().
inner_idx = after.index("except (OSError, ValueError):")
after_inner = after[inner_idx + len("except (OSError, ValueError):") :]
outer_match = re.search(r"except\s*\(?[^)]*?\)?:", after_inner)
@ -91,8 +76,7 @@ def test_kill_orphan_catches_oserror_from_studio_root():
def _exec_search_roots_block(
home: Path, studio_root_value: Path, resolve_raises: bool
) -> list[Path]:
"""Extract _find_llama_server_binary's env-mode search_roots block
and execute it with controlled inputs."""
"""Extract and run _find_llama_server_binary's env-mode search_roots block with controlled inputs."""
src = LLAMA_CPP.read_text()
block_start = src.index('legacy_llama = Path.home() / ".unsloth" / "llama.cpp"')
block_end = src.index("_seen_roots: set[str]", block_start)

10
tests/test_video_path_validation.py
View file

@ -1,8 +1,7 @@
"""
Tests for check_dataset_for_missing_videos (issue #5085).
"""Tests for check_dataset_for_missing_videos (issue #5085).
Fixtures extract the function from vision.py via AST so the pure-Python logic
tests run without the full unsloth import chain (triton/CUDA kernels).
Fixtures AST-extract the function from vision.py so logic tests run without
the full unsloth import chain (triton/CUDA kernels).
"""
import ast
@ -22,8 +21,7 @@ def _extract_fns_via_ast(
fn_names,
extra_ns = None,
):
"""Exec a set of top-level functions out of a .py file so intra-module
references between them resolve."""
"""Exec the named top-level functions from a .py file so their mutual references resolve."""
source = source_path.read_text(encoding = "utf-8")
tree = ast.parse(source, filename = str(source_path))
wanted = set(fn_names)

32
tests/test_windows_rocm_bnb_version.py
View file

@ -11,10 +11,7 @@
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
"""Tests for ``maybe_set_windows_rocm_bnb_version`` (unsloth/import_fixes.py).
The module is loaded in isolation (stdlib + packaging only), so no torch /
GPU is required and unsloth's GPU init never runs."""
"""Tests for ``maybe_set_windows_rocm_bnb_version`` (loaded in isolation, no torch/GPU)."""
from __future__ import annotations
@ -45,8 +42,8 @@ def import_fixes():
@pytest.fixture()
def clean_env(monkeypatch):
"""Unset the env vars and remove them afterwards (the function writes
os.environ directly, which monkeypatch does not auto-revert)."""
"""Unset the env vars and remove them afterwards; the function writes
os.environ directly, which monkeypatch does not auto-revert."""
for var in (
"BNB_ROCM_VERSION",
"UNSLOTH_SKIP_BNB_ROCM_VERSION",
@ -78,7 +75,7 @@ def test_detect_picks_highest_rocm_suffix(import_fixes, tmp_path, monkeypatch):
pkg.mkdir()
for name in (
"libbitsandbytes_rocm72.dll",
"libbitsandbytes_rocm713.dll", # numerically highest -> should win
"libbitsandbytes_rocm713.dll", # numerically highest -> wins
"libbitsandbytes_cpu.dll",
"__init__.py",
):
@ -129,8 +126,7 @@ def test_noop_when_not_rocm_torch(import_fixes, clean_env):
def test_noop_when_no_rocm_dll_installed(import_fixes, clean_env):
# Never force a ROCm backend name when no ROCm DLL ships (avoid breaking a
# non-ROCm bitsandbytes that happens to sit next to a ROCm torch build).
# No ROCm DLL ships -> don't force a backend name (would break non-ROCm bnb).
_force(import_fixes, clean_env, win = True, rocm = True, detected = None)
assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
assert "BNB_ROCM_VERSION" not in os.environ
@ -151,8 +147,7 @@ def test_explicit_opt_out(import_fixes, clean_env):
def test_redetects_sitecustomize_seeded_default(import_fixes, clean_env):
# Studio's installer persists a default via the venv sitecustomize.py; the
# wheel may have changed since, so the seeded value must be redetected.
# A sitecustomize-seeded default must be redetected (the wheel may have changed).
clean_env.setenv("BNB_ROCM_VERSION", "72")
clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "sitecustomize")
_force(import_fixes, clean_env, win = True, rocm = True, detected = "713")
@ -181,8 +176,7 @@ def test_user_value_with_non_sitecustomize_marker_untouched(import_fixes, clean_
def test_opt_out_unseats_sitecustomize_seeded_value(import_fixes, clean_env):
# The opt-out must also drop a default our own sitecustomize block seeded,
# so bitsandbytes never sees the override the user disabled.
# Opt-out must also drop a sitecustomize-seeded default so bnb never sees it.
clean_env.setenv("BNB_ROCM_VERSION", "72")
clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "sitecustomize")
clean_env.setenv("UNSLOTH_SKIP_BNB_ROCM_VERSION", "1")
@ -202,8 +196,7 @@ def test_opt_out_keeps_explicit_user_value(import_fixes, clean_env):
def test_empty_string_value_without_marker_is_respected(import_fixes, clean_env):
# "" counts as present: without the sitecustomize marker it is not ours
# to overwrite.
# "" counts as present: without the sitecustomize marker it is not ours to overwrite.
clean_env.setenv("BNB_ROCM_VERSION", "")
_force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
@ -211,8 +204,8 @@ def test_empty_string_value_without_marker_is_respected(import_fixes, clean_env)
# ---------------------------------------------------------------------------
# _is_hip_torch_build (the strict gate -- regression for the HIP-SDK-on-a-
# CUDA-box false positive: env hints like HIP_PATH must NOT count)
# _is_hip_torch_build: strict gate; HIP-SDK env hints (HIP_PATH) must NOT count
# (regression for the HIP-SDK-on-a-CUDA-box false positive).
# ---------------------------------------------------------------------------
@ -226,15 +219,14 @@ def test_hip_build_true_from_wheel_tag(import_fixes, monkeypatch):
def test_hip_build_true_from_torch_version_hip(import_fixes, monkeypatch):
# Custom/source HIP build without the +rocm tag.
# Custom/source HIP build without the +rocm wheel tag.
monkeypatch.setattr(import_fixes, "importlib_version", lambda name: "2.11.0")
monkeypatch.setitem(__import__("sys").modules, "torch", _fake_torch("7.2.0"))
assert import_fixes._is_hip_torch_build() is True
def test_hip_build_false_for_cuda_torch_despite_rocm_env_hints(import_fixes, monkeypatch):
"""HIP SDK env vars set but CUDA torch: the strict gate must say False,
otherwise BNB_ROCM_VERSION gets set and CUDA bitsandbytes raises."""
"""HIP SDK env vars but CUDA torch: gate must be False, else CUDA bnb raises."""
monkeypatch.setenv("HIP_PATH", r"C:\Program Files\AMD\ROCm\6.2")
monkeypatch.setenv("ROCM_PATH", r"C:\Program Files\AMD\ROCm\6.2")
monkeypatch.setattr(import_fixes, "importlib_version", lambda name: "2.9.0+cu126")

22
tests/utils/aime_eval.py
View file

@ -41,7 +41,7 @@ def download_and_combine_aime_datasets(data_dir: str = "./data/aime") -> str:
response = requests.get(url)
response.raise_for_status()
# Parse each line and tag with source + global ID
# Tag each line with its source dataset + global ID
for line_num, line in enumerate(response.text.strip().split("\n")):
if line.strip():
try:
@ -149,7 +149,7 @@ def extract_aime_answer(response: str) -> str:
for pattern in patterns:
matches = re.findall(pattern, response_lower, re.MULTILINE | re.IGNORECASE)
if matches:
answer = matches[-1] # last match is most likely the final answer
answer = matches[-1] # last match = the final answer
try:
num = int(answer)
if 0 <= num <= 999:
@ -212,7 +212,6 @@ def evaluate_model_aime(
output_tokens = []
correct_answers = 0
# Track performance by source dataset
source_stats = {}
for example in eval_dataset:
source = example["source_dataset"]
@ -220,7 +219,6 @@ def evaluate_model_aime(
source_stats[source] = {"total": 0, "correct": 0}
source_stats[source]["total"] += 1
# Setup sampling parameters (AIME configuration)
sampling_params = SamplingParams(
temperature = temperature,
top_p = top_p,
@ -236,7 +234,7 @@ def evaluate_model_aime(
print(f" Top-p: {top_p}")
print(f" Seed: {seed}")
# Temporarily suppress verbose logging
# Temporarily suppress verbose vllm/ray logging
original_levels = {}
loggers_to_suppress = [
"vllm",
@ -264,7 +262,6 @@ def evaluate_model_aime(
input_tokens.append(get_num_tokens(prompt_text, tokenizer))
# Generate multiple responses
outputs = model.fast_generate(
[prompt_text],
sampling_params = sampling_params,
@ -332,14 +329,13 @@ def evaluate_model_aime(
continue
finally:
# Restore logging levels
for logger_name, level in original_levels.items():
logging.getLogger(logger_name).setLevel(level)
total_problems = len(eval_dataset)
accuracy = correct_answers / total_problems * 100
# Pass@k: probability at least one of k samples is correct
# Pass@k: fraction of problems where at least one of k samples is correct
pass_at_k_scores = []
for record in records.values():
if "n_correct" in record and "n_total" in record:
@ -352,7 +348,6 @@ def evaluate_model_aime(
pass_at_k = sum(pass_at_k_scores) / len(pass_at_k_scores) if pass_at_k_scores else 0
# Per-source accuracies
source_accuracies = {}
for source, stats in source_stats.items():
source_accuracies[source] = (
@ -410,7 +405,6 @@ def evaluate_model_aime(
print(f" Max input tokens: {results['max_input_tokens']:>10}")
print(f" Max output tokens: {results['max_output_tokens']:>10}")
# Performance assessment for AIME
if accuracy >= 50:
tier = "🏆 EXCEPTIONAL"
elif accuracy >= 30:
@ -431,14 +425,12 @@ def evaluate_model_aime(
return results
# Comparison functions for multiple model results
def compare_aime_results(all_results):
"""Generate comprehensive comparison for AIME evaluation results"""
print(f"\n{'='*80}")
print("COMPREHENSIVE AIME MODEL COMPARISON")
print(f"{'='*80}")
# Main comparison table
print(f"{'Model':<15} {'Accuracy %':<12} {'Pass@K %':<10} {'Correct':<8} {'Total':<8}")
print("-" * 80)
@ -451,13 +443,12 @@ def compare_aime_results(all_results):
f"{result['total_problems']:<8}"
)
# Performance improvement analysis
if len(all_results) > 1:
print(f"\n{'='*50}")
print("IMPROVEMENT ANALYSIS")
print(f"{'='*50}")
base_result = all_results[0] # Assume first is base model
base_result = all_results[0] # first is the base model
for i, result in enumerate(all_results[1:], 1):
print(f"\n{result['model_type']} vs {base_result['model_type']}:")
@ -468,12 +459,10 @@ def compare_aime_results(all_results):
print(f" Accuracy improvement: {accuracy_improvement:+.1f}%")
print(f" Pass@K improvement: {pass_k_improvement:+.1f}%")
# Dataset breakdown
print(f"\n{'='*50}")
print("PERFORMANCE BY DATASET")
print(f"{'='*50}")
# Get all unique datasets from the first result
if all_results and "source_accuracies" in all_results[0]:
datasets = list(all_results[0]["source_accuracies"].keys())
@ -490,7 +479,6 @@ def compare_aime_results(all_results):
print(f"{accuracy:<15.1f}", end = "")
print()
# Save comparison
comparison_data = {
"summary": all_results,
"best_model": max(all_results, key = lambda x: x["accuracy"]),

31
tests/utils/cleanup_utils.py
View file

@ -12,16 +12,9 @@ def clear_memory(
verbose = False,
clear_all_caches = True,
):
"""
Comprehensive memory clearing for persistent memory leaks.
"""Comprehensive memory clearing for persistent memory leaks."""
Args:
variables_to_clear: List of variable names to clear
verbose: Print memory status
clear_all_caches: Clear all types of caches (recommended for memory leaks)
"""
# Save logging levels to restore later
# Save logging levels to restore later.
saved_log_levels = {}
for name, logger in logging.Logger.manager.loggerDict.items():
if isinstance(logger, logging.Logger):
@ -42,11 +35,11 @@ def clear_memory(
"bnb_config",
]
# Clear LRU caches first (important for memory leaks)
# Clear LRU caches first (important for memory leaks).
if clear_all_caches:
clear_all_lru_caches(verbose)
# Delete specified variables
# Delete specified variables.
g = globals()
deleted_vars = []
for var in variables_to_clear:
@ -57,7 +50,7 @@ def clear_memory(
if verbose and deleted_vars:
print(f"Deleted variables: {deleted_vars}")
# Multiple GC passes (important for circular references)
# Multiple GC passes for circular references.
for i in range(3):
collected = gc.collect()
if verbose and collected > 0:
@ -75,7 +68,7 @@ def clear_memory(
torch.cuda.reset_peak_memory_stats()
torch.cuda.reset_accumulated_memory_stats()
# Clear JIT cache
# Clear JIT cache.
if hasattr(torch.jit, "_state") and hasattr(torch.jit._state, "_clear_class_state"):
torch.jit._state._clear_class_state()
@ -91,7 +84,7 @@ def clear_memory(
if mem_before > 0:
print(f"Memory freed: {mem_before - mem_after:.2f} GB")
# Restore original logging levels
# Restore original logging levels.
logging.getLogger().setLevel(root_level)
for name, level in saved_log_levels.items():
if name in logging.Logger.manager.loggerDict:
@ -103,7 +96,7 @@ def clear_all_lru_caches(verbose = True):
"""Clear all LRU caches in loaded modules."""
cleared_caches = []
# Skip these to avoid warnings
# Skip these to avoid warnings.
skip_modules = {
"torch.distributed",
"torchaudio",
@ -112,10 +105,10 @@ def clear_all_lru_caches(verbose = True):
"torchaudio.backend",
}
# Static list to avoid RuntimeError during iteration
# Static list to avoid RuntimeError during iteration.
modules = list(sys.modules.items())
# Clear caches in all loaded modules
# Clear caches in all loaded modules.
for module_name, module in modules:
if module is None:
continue
@ -126,7 +119,7 @@ def clear_all_lru_caches(verbose = True):
try:
for attr_name in dir(module):
try:
# Suppress warnings when checking attributes
# Suppress warnings when checking attributes.
with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)
warnings.simplefilter("ignore", UserWarning)
@ -141,7 +134,7 @@ def clear_all_lru_caches(verbose = True):
except Exception:
continue
# Clear specific known caches
# Clear specific known caches.
known_caches = [
"transformers.utils.hub.cached_file",
"transformers.tokenization_utils_base.get_tokenizer",

7
tests/utils/data_utils.py
View file

@ -56,12 +56,7 @@ def describe_param(
include_infinity: bool = False,
as_str: bool = True,
) -> dict:
"""
Statistical summary (shape, mean, std, min/max, percentiles) of a tensor.
Optionally includes L1/L2/infinity norms. Returns a formatted string when
as_str is True, else a dict.
"""
"""Statistical summary of a tensor (optional L1/L2/inf norms); string if as_str else dict."""
param = param.float()
summary = {

24
tests/utils/generate_dataset_with_none.py
View file

@ -1,18 +1,9 @@
"""
Generate a small synthetic dataset with intentional None/empty turns so
dataset_none_detect.py can be verified end-to-end.
Three formats: chatml (messages, role/content), sharegpt (conversations,
from/value), and alpaca (instruction/output). ~20 rows each; roughly half
have at least one bad turn. Only depends on `datasets`.
"""
"""Synthetic chatml/sharegpt/alpaca datasets with intentional None/empty turns for dataset_none_detect.py."""
from datasets import Dataset
# ChatML (messages, role/content)
# pyarrow requires uniform types in a column, so messages=None / non-list (P1)
# rows live in a SEPARATE dataset so pyarrow can infer the column type.
# ChatML (messages, role/content). pyarrow needs uniform column types, so
# messages=None / non-list (P1) rows live in a SEPARATE dataset.
_CHATML_ROWS = [
# clean rows
@ -77,7 +68,7 @@ _CHATML_ROWS = [
{"role": "assistant", "content": "Shakespeare."},
]
},
# bad rows — None/empty turn content (all messages values are lists so pyarrow is happy)
# bad rows: None/empty turn content (all values are lists, so pyarrow is happy)
{
"messages": [
{"role": "user", "content": None},
@ -123,9 +114,8 @@ _CHATML_ROWS = [
{"messages": [None, {"role": "assistant", "content": "Reply"}]}, # None turn element
]
# P1 test rows: messages is None or non-list. Stored as plain dicts (not an
# HF Dataset) since pyarrow can't mix list and non-list values in one column;
# the test runner mocks find_none_chatml directly.
# P1 rows: messages is None or non-list. Plain dicts (not an HF Dataset) since
# pyarrow can't mix list/non-list in one column; the runner mocks find_none_chatml.
_CHATML_P1_ROWS = [
{"messages": None}, # whole column None
{"messages": "not a list"}, # wrong type
@ -133,7 +123,7 @@ _CHATML_P1_ROWS = [
def make_chatml_p1_rows() -> list:
"""Return the raw P1 rows (not an HF Dataset) for direct mock testing."""
"""Raw P1 rows (not an HF Dataset) for direct mock testing."""
return list(_CHATML_P1_ROWS)

Some files were not shown because too many files have changed in this diff Show more