diff --git a/tests/_zoo_aggressive_cuda_spoof.py b/tests/_zoo_aggressive_cuda_spoof.py index dafe81c829..05889d5df2 100644 --- a/tests/_zoo_aggressive_cuda_spoof.py +++ b/tests/_zoo_aggressive_cuda_spoof.py @@ -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: diff --git a/tests/conftest.py b/tests/conftest.py index 4b7081b30b..3478a19af8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 .device_type under a mocked - torch.cuda.is_available() == True so its @cache permanently - captures "cuda". prereqs lists submodule names of 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 .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. # --------------------------------------------------------------------------- diff --git a/tests/notebooks/test_validator_fixtures.py b/tests/notebooks/test_validator_fixtures.py index 836bb96715..56608bff53 100644 --- a/tests/notebooks/test_validator_fixtures.py +++ b/tests/notebooks/test_validator_fixtures.py @@ -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( diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py index 9ab68639c4..2b1d012d14 100644 --- a/tests/python/test_construct_chat_template_validation.py +++ b/tests/python/test_construct_chat_template_validation.py @@ -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 = "" @@ -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 = [""], ) 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 diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 34f984714e..666ea7ce10 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -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() diff --git a/tests/python/test_e2e_no_torch_sandbox.py b/tests/python/test_e2e_no_torch_sandbox.py index 9b3bda133f..7da7cf0eae 100644 --- a/tests/python/test_e2e_no_torch_sandbox.py +++ b/tests/python/test_e2e_no_torch_sandbox.py @@ -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}") diff --git a/tests/python/test_fast_language_model_text_only.py b/tests/python/test_fast_language_model_text_only.py index ce4dd74439..fcdeb49bc3 100644 --- a/tests/python/test_fast_language_model_text_only.py +++ b/tests/python/test_fast_language_model_text_only.py @@ -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() diff --git a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py index 6fa7d06ec3..2e12a43228 100644 --- a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py +++ b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py @@ -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 diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py index 64373cb4b8..f1a090eac9 100644 --- a/tests/python/test_install_python_stack.py +++ b/tests/python/test_install_python_stack.py @@ -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 diff --git a/tests/python/test_no_torch_filtering.py b/tests/python/test_no_torch_filtering.py index 22cf04a513..9dcda5c60c 100644 --- a/tests/python/test_no_torch_filtering.py +++ b/tests/python/test_no_torch_filtering.py @@ -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 diff --git a/tests/python/test_orpo_processor_text_tokenizer.py b/tests/python/test_orpo_processor_text_tokenizer.py index 9ae205c6b9..b507a9e808 100644 --- a/tests/python/test_orpo_processor_text_tokenizer.py +++ b/tests/python/test_orpo_processor_text_tokenizer.py @@ -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: diff --git a/tests/python/test_patch_trl_rl_trainers_defensive.py b/tests/python/test_patch_trl_rl_trainers_defensive.py index 55a3425c44..e5d6328031 100644 --- a/tests/python/test_patch_trl_rl_trainers_defensive.py +++ b/tests/python/test_patch_trl_rl_trainers_defensive.py @@ -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) diff --git a/tests/python/test_studio_import_no_torch.py b/tests/python/test_studio_import_no_torch.py index 7b7a7103d2..86dc8581ab 100644 --- a/tests/python/test_studio_import_no_torch.py +++ b/tests/python/test_studio_import_no_torch.py @@ -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, diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index 40924efca2..7390d7be9b 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -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}" diff --git a/tests/python/test_unsloth_run_tool_policy_resolver.py b/tests/python/test_unsloth_run_tool_policy_resolver.py index 356170fae8..00d4996368 100644 --- a/tests/python/test_unsloth_run_tool_policy_resolver.py +++ b/tests/python/test_unsloth_run_tool_policy_resolver.py @@ -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 diff --git a/tests/saving/gpt-oss-merge/train_and_merge.py b/tests/saving/gpt-oss-merge/train_and_merge.py index 1d35c6759d..aca5b1a19b 100644 --- a/tests/saving/gpt-oss-merge/train_and_merge.py +++ b/tests/saving/gpt-oss-merge/train_and_merge.py @@ -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.") diff --git a/tests/saving/language_models/test_merge_4bit_validation.py b/tests/saving/language_models/test_merge_4bit_validation.py index 2e2e823bd3..8c1baeecf4 100644 --- a/tests/saving/language_models/test_merge_4bit_validation.py +++ b/tests/saving/language_models/test_merge_4bit_validation.py @@ -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(), diff --git a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py index f8f81c550e..3b75a13756 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py +++ b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py @@ -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"): diff --git a/tests/saving/language_models/test_merge_model_perplexity_mistral.py b/tests/saving/language_models/test_merge_model_perplexity_mistral.py index d467089a47..8cc833c2b1 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_mistral.py +++ b/tests/saving/language_models/test_merge_model_perplexity_mistral.py @@ -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"): diff --git a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py index 391d1b90db..6f79bfdb71 100644 --- a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py +++ b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py @@ -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"): diff --git a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py index 9c2c1acb41..c07b37024f 100644 --- a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py +++ b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py @@ -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 = [ diff --git a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py index cdbad174a4..cb444d1591 100644 --- a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py +++ b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py @@ -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 diff --git a/tests/saving/language_models/test_push_to_hub_merged.py b/tests/saving/language_models/test_push_to_hub_merged.py index cacb1e2a10..0a1dd3406a 100644 --- a/tests/saving/language_models/test_push_to_hub_merged.py +++ b/tests/saving/language_models/test_push_to_hub_merged.py @@ -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") diff --git a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py index 59a09f06ec..a77d35d05f 100644 --- a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py +++ b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py @@ -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() diff --git a/tests/saving/language_models/test_save_merged_grpo_model.py b/tests/saving/language_models/test_save_merged_grpo_model.py index 63e4b970a6..24a35dc825 100644 --- a/tests/saving/language_models/test_save_merged_grpo_model.py +++ b/tests/saving/language_models/test_save_merged_grpo_model.py @@ -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() diff --git a/tests/saving/non_peft/test_whisper_non_peft.py b/tests/saving/non_peft/test_whisper_non_peft.py index b0f2b2e751..91425d77e4 100644 --- a/tests/saving/non_peft/test_whisper_non_peft.py +++ b/tests/saving/non_peft/test_whisper_non_peft.py @@ -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: diff --git a/tests/saving/test_preserve_tokenizer_eos_token.py b/tests/saving/test_preserve_tokenizer_eos_token.py index 2ea40ab778..2dd280e6cd 100644 --- a/tests/saving/test_preserve_tokenizer_eos_token.py +++ b/tests/saving/test_preserve_tokenizer_eos_token.py @@ -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"] == "" 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() diff --git a/tests/saving/test_save_subprocess_utf8_encoding.py b/tests/saving/test_save_subprocess_utf8_encoding.py index 4a609cd7b7..165087d03a 100644 --- a/tests/saving/test_save_subprocess_utf8_encoding.py +++ b/tests/saving/test_save_subprocess_utf8_encoding.py @@ -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") diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py index a85c119e04..de3a1556b4 100644 --- a/tests/saving/test_unsloth_save.py +++ b/tests/saving/test_unsloth_save.py @@ -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, diff --git a/tests/saving/text_to_speech_models/test_csm.py b/tests/saving/text_to_speech_models/test_csm.py index 3d04bc39cb..7c1a7a275b 100644 --- a/tests/saving/text_to_speech_models/test_csm.py +++ b/tests/saving/text_to_speech_models/test_csm.py @@ -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, diff --git a/tests/saving/text_to_speech_models/test_lasa.py b/tests/saving/text_to_speech_models/test_lasa.py index c0c4f80e0e..eea69ddd0a 100644 --- a/tests/saving/text_to_speech_models/test_lasa.py +++ b/tests/saving/text_to_speech_models/test_lasa.py @@ -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) diff --git a/tests/saving/text_to_speech_models/test_orpheus.py b/tests/saving/text_to_speech_models/test_orpheus.py index 83adaf8dbc..99feeeda13 100644 --- a/tests/saving/text_to_speech_models/test_orpheus.py +++ b/tests/saving/text_to_speech_models/test_orpheus.py @@ -140,7 +140,7 @@ prompts = [ "Hey there my name is Elise, 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] diff --git a/tests/saving/text_to_speech_models/test_whisper.py b/tests/saving/text_to_speech_models/test_whisper.py index e0271d098e..a3f41d1833 100644 --- a/tests/saving/text_to_speech_models/test_whisper.py +++ b/tests/saving/text_to_speech_models/test_whisper.py @@ -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", diff --git a/tests/saving/vision_models/test_index_file_sharded_model.py b/tests/saving/vision_models/test_index_file_sharded_model.py index 6f2d0ff782..1a3c1f1493 100644 --- a/tests/saving/vision_models/test_index_file_sharded_model.py +++ b/tests/saving/vision_models/test_index_file_sharded_model.py @@ -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 diff --git a/tests/saving/vision_models/test_push_to_hub_merged.py b/tests/saving/vision_models/test_push_to_hub_merged.py index d83c66de4b..513546efb4 100644 --- a/tests/saving/vision_models/test_push_to_hub_merged.py +++ b/tests/saving/vision_models/test_push_to_hub_merged.py @@ -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") diff --git a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py index 548e6bfc37..65c876e766 100644 --- a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py +++ b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py @@ -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 diff --git a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py index 66064600ea..3305dbc53c 100644 --- a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py +++ b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py @@ -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 ) diff --git a/tests/security/conftest.py b/tests/security/conftest.py index c5071842e5..5ff22d79e6 100644 --- a/tests/security/conftest.py +++ b/tests/security/conftest.py @@ -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: diff --git a/tests/security/fixtures/_build.py b/tests/security/fixtures/_build.py index 4cbda6aab5..f52c401b60 100644 --- a/tests/security/fixtures/_build.py +++ b/tests/security/fixtures/_build.py @@ -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). diff --git a/tests/security/test_lint_workflow_triggers.py b/tests/security/test_lint_workflow_triggers.py index 554c360d88..e4ee3fd28f 100644 --- a/tests/security/test_lint_workflow_triggers.py +++ b/tests/security/test_lint_workflow_triggers.py @@ -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" diff --git a/tests/security/test_lockfile_supply_chain_audit.py b/tests/security/test_lockfile_supply_chain_audit.py index 905c4afcac..d464fe44d4 100644 --- a/tests/security/test_lockfile_supply_chain_audit.py +++ b/tests/security/test_lockfile_supply_chain_audit.py @@ -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. diff --git a/tests/security/test_new_install_scripts.py b/tests/security/test_new_install_scripts.py index e71659dd88..d8656b114e 100644 --- a/tests/security/test_new_install_scripts.py +++ b/tests/security/test_new_install_scripts.py @@ -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) diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py index f0c0ea93d4..45d1d42877 100644 --- a/tests/security/test_scan_npm_packages.py +++ b/tests/security/test_scan_npm_packages.py @@ -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")) diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py index 33f3fda488..c9b64a9da4 100644 --- a/tests/security/test_scan_packages.py +++ b/tests/security/test_scan_packages.py @@ -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 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" diff --git a/tests/studio/_playwright_robust.py b/tests/studio/_playwright_robust.py index bc0e3f783b..774665581d 100644 --- a/tests/studio/_playwright_robust.py +++ b/tests/studio/_playwright_robust.py @@ -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 -# ` 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 +# ` 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 = ( diff --git a/tests/studio/install/conftest.py b/tests/studio/install/conftest.py index 8738ef2319..86104231f3 100644 --- a/tests/studio/install/conftest.py +++ b/tests/studio/install/conftest.py @@ -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 diff --git a/tests/studio/install/smoke_test_parallel_studio_home.py b/tests/studio/install/smoke_test_parallel_studio_home.py index be01147740..9d840355e1 100644 --- a/tests/studio/install/smoke_test_parallel_studio_home.py +++ b/tests/studio/install/smoke_test_parallel_studio_home.py @@ -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") diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index 83c2d962e8..cea4383268 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -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) diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py index 983969a4d3..0926cf9bd5 100644 --- a/tests/studio/install/test_gpu_detection_followups.py +++ b/tests/studio/install/test_gpu_detection_followups.py @@ -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 diff --git a/tests/studio/install/test_hf_auth.py b/tests/studio/install/test_hf_auth.py index 7c3296d4fd..99e65b7fed 100644 --- a/tests/studio/install/test_hf_auth.py +++ b/tests/studio/install/test_hf_auth.py @@ -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: diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 9ccfe15a88..a5a1131ecc 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -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//bin``, current ``nvidia//bin/x86_64`` - (cu13 layout), conda-style ``nvidia//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} diff --git a/tests/studio/install/test_launch_studio_launcher.py b/tests/studio/install/test_launch_studio_launcher.py index 8ea053b783..a7396aaf5d 100644 --- a/tests/studio/install/test_launch_studio_launcher.py +++ b/tests/studio/install/test_launch_studio_launcher.py @@ -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 diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py index 8643806127..8f40c9d720 100644 --- a/tests/studio/install/test_llama_pr_force_and_source.py +++ b/tests/studio/install/test_llama_pr_force_and_source.py @@ -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: diff --git a/tests/studio/install/test_macos_version_compat.py b/tests/studio/install/test_macos_version_compat.py index 7f93b295eb..a8b57fe3c5 100644 --- a/tests/studio/install/test_macos_version_compat.py +++ b/tests/studio/install/test_macos_version_compat.py @@ -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 = [ diff --git a/tests/studio/install/test_pr4562_bugfixes.py b/tests/studio/install/test_pr4562_bugfixes.py index 9dfa4e0005..b19c300fec 100644 --- a/tests/studio/install/test_pr4562_bugfixes.py +++ b/tests/studio/install/test_pr4562_bugfixes.py @@ -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] diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py index 583532eb74..9b739272b1 100644 --- a/tests/studio/install/test_pr5940_followups.py +++ b/tests/studio/install/test_pr5940_followups.py @@ -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 ... ) ="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) diff --git a/tests/studio/install/test_probe_timeouts.py b/tests/studio/install/test_probe_timeouts.py index acea0ed34d..0ff93c1b12 100644 --- a/tests/studio/install/test_probe_timeouts.py +++ b/tests/studio/install/test_probe_timeouts.py @@ -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) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index b0abaaeacd..8e0b5d1b7f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1,9 +1,4 @@ -"""Tests for AMD ROCm support across install pathways. - -Verifies that ROCm detection and installation logic works correctly -WITHOUT breaking existing CUDA, CPU, macOS, and Windows pathways. -All tests use mocks -- no AMD hardware required. -""" +"""AMD ROCm support tests across install pathways (all mocked, no AMD HW).""" import importlib.util import json @@ -57,12 +52,7 @@ _install_bnb_windows_rocm = stack_mod._install_bnb_windows_rocm def _extract_sh_function_body(source: str, name: str) -> str: - """Return the body of a shell function from `source` by brace matching. - - Used by structural tests that need to assert ordering of helper - calls inside a specific function rather than across the whole - install.sh file. - """ + """Return a shell function body from `source` by brace matching.""" needle = f"{name}() {{" start = source.find(needle) if start < 0: @@ -292,14 +282,12 @@ class TestResolveUpstreamAssetChoice: """Host with both NVIDIA and ROCm should use NVIDIA (CPU path here, CUDA elsewhere).""" host = nvidia_host(has_rocm = True) choice = resolve_upstream_asset_choice(host, LLAMA_TAG) - # NVIDIA hosts go through the normal path (CUDA handled by resolve_linux_cuda_choice) assert choice.install_kind == "linux-cpu" assert "rocm" not in choice.name @patch.object(prebuilt_mod, "github_release_assets") def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets): """AMD ROCm host should fall back to source build when no ROCm prebuilt exists.""" - # Remove the ROCm asset from available assets assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k} mock_assets.return_value = assets_without_rocm host = rocm_host() @@ -343,9 +331,8 @@ class TestRuntimePatterns: patterns = runtime_patterns_for_choice(choice) assert "llama-server" in patterns assert "llama-quantize" in patterns - # Broad lib*.so* covers libllama, libggml, libmtmd, libggml-cpu-*, - # plus the libllama--impl.so split that ggml-org/llama.cpp - # #23462 introduced between b9279 and b9283. + # lib*.so* covers libllama/libggml/libmtmd plus the libllama-*-impl.so + # split from ggml-org/llama.cpp #23462 (between b9279 and b9283). assert "lib*.so*" in patterns def test_linux_cuda_patterns(self): @@ -353,7 +340,6 @@ class TestRuntimePatterns: repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cuda" ) patterns = runtime_patterns_for_choice(choice) - # libggml-cuda.so is matched by lib*.so* now. assert "lib*.so*" in patterns def test_linux_rocm_patterns(self): @@ -361,7 +347,6 @@ class TestRuntimePatterns: repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-rocm" ) patterns = runtime_patterns_for_choice(choice) - # libggml-hip.so is matched by lib*.so* now. assert "lib*.so*" in patterns assert "llama-server" in patterns @@ -375,8 +360,7 @@ class TestRuntimePatterns: install_kind = "windows-hip", ) patterns = runtime_patterns_for_choice(choice) - # Narrowed from "*.exe" to the two binaries Studio actually - # invokes, mirroring the Linux/macOS pattern style. + # Narrowed from "*.exe" to the two binaries Studio actually invokes. assert "llama-server.exe" in patterns assert "llama-quantize.exe" in patterns assert "*.dll" in patterns @@ -394,8 +378,8 @@ class TestRuntimePatterns: assert "lib*.dylib" in patterns def test_diffusion_visual_server_kept(self): - # The DiffusionGemma visual-server ships in the prebuilt bundle and must - # survive the prune so Studio can serve DiffusionGemma GGUFs natively. + # The DiffusionGemma visual-server must survive the prune so Studio can + # serve DiffusionGemma GGUFs natively. for kind, name in ( ("linux-cuda", "llama-diffusion-gemma-visual-server"), ("macos-arm64", "llama-diffusion-gemma-visual-server"), @@ -445,7 +429,7 @@ class TestHostInfoRocm: import inspect source = inspect.getsource(prebuilt_mod.detect_host) - # Must probe for actual GPU, not just tool presence + # Must probe for actual GPU, not just tool presence. assert "rocminfo" in source or "amd-smi" in source def test_detect_host_windows_rocm_detection(self): @@ -527,7 +511,6 @@ class TestDetectRocmVersion: """Debian epoch prefix (2:6.2.0) -- version file has no epoch, so should parse.""" info_dir = tmp_path / ".info" info_dir.mkdir() - # Version files don't typically have epoch prefix, but lib/rocm_version might (info_dir / "version").write_text("6.2.0\n") with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): result = _detect_rocm_version() @@ -578,9 +561,8 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) def test_no_rocm_skips(self, mock_nvidia, mock_pip): """No ROCm toolchain should skip entirely.""" - # _detect_windows_gfx_arch pinned to None: on a real AMD test host its - # WMI name fallback would otherwise answer and defeat the "no ROCm - # anywhere" premise of this test. + # Pin _detect_windows_gfx_arch to None so a real AMD test host's WMI + # fallback can't defeat the "no ROCm anywhere" premise. with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None): with patch("os.path.isdir", return_value = False): with patch("shutil.which", return_value = None): @@ -595,7 +577,7 @@ class TestEnsureRocmTorch: """If torch already has CUDA, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 - mock_probe.stdout = b"12.6\n" # CUDA version string + mock_probe.stdout = b"12.6\n" # CUDA version with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() @@ -609,7 +591,7 @@ class TestEnsureRocmTorch: """If torch already has HIP, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 - mock_probe.stdout = b"7.1.12345\n" # HIP version string + mock_probe.stdout = b"7.1.12345\n" # HIP version with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() @@ -631,7 +613,6 @@ class TestEnsureRocmTorch: with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() - # Should install torch via pip_install and bitsandbytes via pip_install_try. assert mock_pip.call_count == 1 assert "rocm7.1" in str(mock_pip.call_args_list[0]) assert mock_pip_try.call_count >= 1 @@ -712,8 +693,7 @@ class TestEnsureRocmTorch: with patch("os.path.isdir", return_value = True): with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)): _ensure_rocm_torch() - # If probe times out, the function should treat torch as unusable and reinstall - # both torch (via pip_install) and bitsandbytes (via pip_install_try). + # Probe timeout: treat torch as unusable and reinstall torch + bitsandbytes. assert mock_pip.call_count == 1 assert "rocm7.1" in str(mock_pip.call_args_list[0]) assert mock_pip_try.call_count >= 1 @@ -724,9 +704,8 @@ class TestEnsureRocmTorch: @patch.object(stack_mod, "_has_rocm_gpu", return_value = False) def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip): """ROCm tools present but no actual AMD GPU should skip entirely.""" - # Pin the Windows arch probe to None: on a real AMD test host the WMI - # name fallback would otherwise answer and defeat the "no actual GPU" - # premise (the Linux path under test uses _has_rocm_gpu, mocked False). + # Pin the Windows arch probe to None so a real AMD host's WMI fallback + # can't defeat the "no actual GPU" premise. with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None): with patch("os.path.isdir", return_value = True): _ensure_rocm_torch() @@ -738,7 +717,6 @@ class TestEnsureRocmTorch: def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip): """UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe.""" with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}): - # Reload _TORCH_BACKEND from the patched environment. with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"): _ensure_rocm_torch() mock_pip.assert_not_called() @@ -758,12 +736,7 @@ class TestEnsureRocmTorch: class TestHasRocmGpuKfdVendorGuard: - """Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes. - - These tests are source-level: they verify the regex and logic present in - the _has_rocm_gpu implementation rather than running the sysfs traversal - (which requires Linux path conventions). - """ + """KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes (source-level checks).""" def _src(self) -> str: """Return the source of _has_rocm_gpu from install_python_stack.py.""" @@ -787,8 +760,7 @@ class TestHasRocmGpuKfdVendorGuard: import re as _re src = self._src() - # The pattern should have a word boundary before and after the number - # so "vendor_id 41098" doesn't match "vendor_id 4098". + # Word boundary so "vendor_id 41098" doesn't match "vendor_id 4098". assert ( _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src ), "_has_rocm_gpu vendor_id check should use word boundary anchors" @@ -816,14 +788,10 @@ class TestHasRocmGpuKfdVendorGuard: assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)" def test_has_rocm_gpu_returns_false_when_nvidia_present(self): - """_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True. - - This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs - produce a false positive, an NVIDIA GPU always wins. - """ + """_has_rocm_gpu returns False when _has_usable_nvidia_gpu is True (NVIDIA always wins).""" with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True): with patch("shutil.which", return_value = "/usr/bin/rocminfo"): - # Simulate rocminfo claiming an AMD GPU is present + # rocminfo claims an AMD GPU is present. mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = "Name: gfx1100\n" @@ -949,7 +917,6 @@ class TestHardwareRocmFlag: """DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") - # Ensure ROCM is NOT a DeviceType member enum_section = source.split("class DeviceType")[1].split("\n\n")[0] assert "ROCM" not in enum_section @@ -963,7 +930,6 @@ class TestHardwareRocmFlag: """All existing DeviceType.CUDA references should still be present.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") - # Key functions that must still reference DeviceType.CUDA assert "DeviceType.CUDA" in source assert "DEVICE = DeviceType.CUDA" in source @@ -977,26 +943,19 @@ class TestHardwareRocmFlag: """IS_ROCM should be in __all__ list in __init__.py.""" init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" source = init_path.read_text(encoding = "utf-8") - # Extract __all__ section assert '"IS_ROCM"' in source def test_get_package_versions_returns_rocm_key(self): """get_package_versions() source should return both 'cuda' and 'rocm' keys.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") - # Find the get_package_versions function body func_start = source.find("def get_package_versions") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert '"cuda"' in func_body assert '"rocm"' in func_body def test_distributed_stubs_cover_is_torchelastic_launched(self): - """_determine_attention_impl_for_gpu_estimate must stub is_torchelastic_launched. - - resolve_attention_implementation calls is_torchelastic_launched() on - Windows ROCm where torch.distributed ships without that helper, causing - a warning: 'module torch.distributed has no attribute is_torchelastic_launched'. - """ + """Must stub is_torchelastic_launched (Windows ROCm torch.distributed lacks it).""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") assert "is_torchelastic_launched" in source @@ -1035,23 +994,16 @@ class TestInstallShStructure: """Verify install.sh structural properties without running it.""" def test_no_here_strings(self): - """install.sh must not use the bash-only `<<<` here-string operator. - - `<<<` inside a quoted literal (e.g. a marker in a printf) is just data, - not a here-string, so strip quoted spans first: this still catches a - real `cmd <<< word` (outside quotes) without false-positiving on data. - """ + """install.sh must not use the bash-only `<<<` here-string operator (breaks dash).""" import re sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") - # <<< is bash-only; breaks dash for i, line in enumerate(source.splitlines(), 1): stripped = line.lstrip() if stripped.startswith("#"): continue - # Remove quoted string literals so `<<<` inside them is ignored; - # a genuine here-string operator lives outside any quotes. + # Strip quoted literals so `<<<` inside them is ignored. unquoted = re.sub(r"'[^']*'", "", line) unquoted = re.sub(r'"[^"]*"', "", unquoted) assert "<<<" not in unquoted, f"install.sh:{i} uses non-POSIX <<< here-string" @@ -1064,22 +1016,13 @@ class TestInstallShStructure: assert "rocm" in source.lower() def test_cuda_precedence(self): - """ROCm detection should only run when nvidia-smi is absent. - - install.sh defines _has_amd_rocm_gpu and _has_usable_nvidia_gpu - helpers near each other (file-position order has no semantic - meaning), so check the runtime ordering inside - get_torch_index_url instead: NVIDIA branch runs first and the - AMD/ROCm branch only fires inside the `if [ -z "$_smi" ]` - block. - """ + """ROCm detection runs only when NVIDIA is absent (check runtime ordering in get_torch_index_url).""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") body = _extract_sh_function_body(source, "get_torch_index_url") nvidia_call = body.find("_has_usable_nvidia_gpu") - # Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to - # handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu - # returns true via /proc/driver/nvidia/gpus. + # Gate uses _nvidia_detected (not -z "$_smi") to handle proc-only NVIDIA + # hosts where nvidia-smi is absent but the GPU is found via /proc. no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]') if no_nvidia_branch < 0: no_nvidia_branch = body.find('if [ -z "$_smi" ]') @@ -1111,7 +1054,6 @@ class TestInstallShStructure: sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert 'echo "$_base/rocm7.2"' in source # fallback for unknown future versions - # Allowlisted versions should pass through directly assert "rocm6.*" in source assert "rocm7.0" in source assert "rocm7.1" in source @@ -1131,8 +1073,7 @@ class TestInstallShStructure: assert "sed 's/^[0-9]*://' " in source or "sed 's/^[0-9]*://'" in source def test_no_double_bracket_in_rocm_block(self): - """ROCm detection block should not use [[ ]] (bash-only, not POSIX). - Note: [[:space:]], [[:digit:]] etc. are valid POSIX character classes, not bash [[ ]].""" + """ROCm block must not use bash-only [[ ]] (POSIX char classes [[:space:]] are fine).""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") @@ -1144,7 +1085,7 @@ class TestInstallShStructure: stripped = line.lstrip() if stripped.startswith("#"): continue - # Remove POSIX character classes [[:foo:]] before checking for [[ ]] + # Strip POSIX char classes [[:foo:]] before checking for [[ ]]. cleaned = re.sub(r"\[\[:[a-z]+:\]\]", "", line) assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [[" @@ -1174,11 +1115,7 @@ class TestInstallShStructure: assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection" def test_unsloth_torch_backend_exported_after_get_torch_index_url(self): - """install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set. - - This lets install_python_stack.py skip ROCm torch operations on CUDA - and CPU hosts without re-running GPU detection in a subprocess. - """ + """install.sh exports UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL (lets the stack skip GPU re-detection).""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)") @@ -1187,21 +1124,14 @@ class TestInstallShStructure: assert ( backend_pos > torch_url_pos ), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved" - # Verify all three cases are covered assert '"cuda"' in source[backend_pos : backend_pos + 500] assert '"rocm"' in source[backend_pos : backend_pos + 500] assert '"cpu"' in source[backend_pos : backend_pos + 500] - # Must be exported so subprocesses (setup.sh, install_python_stack.py) see it + # Must be exported so subprocesses see it. assert "export UNSLOTH_TORCH_BACKEND" in source def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self): - """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098. - - NVIDIA open kernel module (560+) registers KFD nodes with vendor_id - 4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0 - (true) on NVIDIA-only hosts that have the nvidia-open driver, causing - get_torch_index_url to select a ROCm wheel index. - """ + """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (nvidia-open registers KFD nodes too).""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_amd_rocm_gpu()") @@ -1215,14 +1145,7 @@ class TestInstallShStructure: ), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)" def test_kfd_awk_resets_state_per_file(self): - """KFD sysfs awk must reset gpu/amd state per file (FNR==1). - - Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU - agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU - (gpu_id > 0, vendor_id 4318) can produce a false positive: node 0 - sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1 - before vendor_id 4318 is seen on node 1. - """ + """KFD sysfs awk must reset gpu/amd state per file (FNR==1) to avoid Ryzen+NVIDIA false positives.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_amd_rocm_gpu()") @@ -1234,13 +1157,7 @@ class TestInstallShStructure: ) def test_get_torch_index_url_uses_nvidia_detected_flag(self): - """get_torch_index_url must track NVIDIA detection independently of _smi. - - When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback - but nvidia-smi is not on PATH, _smi stays empty. Without a separate - _nvidia_detected flag, the function falls into the AMD/CPU branch even - though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA. - """ + """get_torch_index_url must track NVIDIA via _nvidia_detected (proc-only NVIDIA still picks CUDA).""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") @@ -1250,8 +1167,6 @@ class TestInstallShStructure: "get_torch_index_url must use a _nvidia_detected flag (separate from " "_smi) so that proc-only NVIDIA detection still selects CUDA wheels" ) - # The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on - # _smi being empty. assert ( '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body ), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1" @@ -1269,8 +1184,7 @@ class TestLiveRegression: if not shutil.which("nvidia-smi"): pytest.skip("No nvidia-smi available") - # Skip if nvidia-smi exists but does not actually list a GPU on this - # host (containers occasionally ship the binary without a driver). + # Skip if nvidia-smi exists but lists no GPU (binary without driver). check = subprocess.run( [ "bash", @@ -1283,9 +1197,7 @@ class TestLiveRegression: pytest.skip("nvidia-smi is on PATH but no GPU is listed") sh_path = PACKAGE_ROOT / "install.sh" - # get_torch_index_url calls _has_usable_nvidia_gpu and - # _has_amd_rocm_gpu, so all three function definitions must be - # in scope when we eval the extract. + # All three helper definitions must be in scope when we eval the extract. extract_cmd = ( f"sed -n '/^_has_amd_rocm_gpu()/,/^}}$/p; " f"/^_has_usable_nvidia_gpu()/,/^}}$/p; " @@ -1305,14 +1217,11 @@ class TestLiveRegression: # TEST: worker.py -- ROCm Mamba/SSM source build path -# Load worker.py module _WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py" _EXPORT_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py" -# The torchao Windows-ROCm stub was de-duplicated out of the export/training -# workers into a shared module; both workers now call into it. +# Shared torchao Windows-ROCm stub used by both workers. _TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py" -# The wheel-probe subprocess was hoisted out of worker.py into wheel_utils -# during the wheel-resolver refactor; the probe script literal lives there. +# Wheel-probe script literal lives in wheel_utils after the resolver refactor. _WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py" @@ -1320,8 +1229,7 @@ class TestWorkerRocmMambaSsm: """Verify worker.py Mamba/SSM install logic on ROCm.""" def test_probe_returns_hip_version_field(self): - """The wheel probe should include hip_version, and worker.py should - consume it.""" + """The wheel probe should include hip_version, and worker.py consumes it.""" assert "hip_version" in _WHEEL_UTILS_PATH.read_text(encoding = "utf-8") assert "hip_version" in _WORKER_PATH.read_text(encoding = "utf-8") @@ -1332,14 +1240,12 @@ class TestWorkerRocmMambaSsm: def test_direct_wheel_url_returns_none_without_cuda_major(self, monkeypatch): """_direct_wheel_url should return None when cuda_major is empty (ROCm).""" - # Load module for function access _worker_spec = importlib.util.spec_from_file_location("test_worker", _WORKER_PATH) assert _worker_spec is not None and _worker_spec.loader is not None worker_mod = importlib.util.module_from_spec(_worker_spec) - # Stub worker.py's imports via monkeypatch so the fakes (notably a - # non-package "utils") are undone and don't break later tests that - # import the real utils.* package. + # Stub worker.py imports via monkeypatch so the fake "utils" is undone + # and doesn't break later tests importing the real utils.* package. loggers_mock = MagicMock() loggers_mock.get_logger = MagicMock(return_value = MagicMock()) monkeypatch.setitem(sys.modules, "structlog", MagicMock()) @@ -1425,7 +1331,6 @@ class TestAmdGpuMonitoring: except Exception: pytest.skip("Could not load amd module in test environment") - # Simulate amd-smi metric JSON output gpu_data = { "usage": {"gfx_activity": "85"}, "temperature": {"edge": "72"}, @@ -1464,9 +1369,8 @@ class TestAmdGpuMonitoring: except Exception: pytest.skip("Could not load amd module") - # _first_visible_amd_gpu_id() short-circuits to None when any of - # HIP / ROCR / CUDA_VISIBLE_DEVICES is "" or "-1". CI often sets - # CUDA_VISIBLE_DEVICES="", so the test must not inherit that. + # _first_visible_amd_gpu_id() returns None if HIP/ROCR/CUDA_VISIBLE_DEVICES + # is "" or "-1"; CI often sets CUDA_VISIBLE_DEVICES="", so clear them. for var in ( "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", @@ -1474,8 +1378,8 @@ class TestAmdGpuMonitoring: ): monkeypatch.delenv(var, raising = False) - # amd-smi is gated off on Windows w/o a HIP SDK; this test mocks it as - # available, so opt in so the gate allows it on every platform. + # amd-smi is gated off on Windows w/o a HIP SDK; opt in so the mock is + # allowed on every platform. monkeypatch.setenv("UNSLOTH_ENABLE_AMD_SMI", "1") mock_json = json.dumps( @@ -1492,9 +1396,8 @@ class TestAmdGpuMonitoring: mock_result.returncode = 0 mock_result.stdout = mock_json - # The premise is "amd-smi exists and answers": the absence guard - # which()-checks before spawning, so hosts without a real amd-smi - # (Linux CI, driver-only Windows) need which mocked too. + # Premise is "amd-smi exists and answers": the guard which()-checks + # before spawning, so mock which too for hosts lacking a real amd-smi. with patch.object(amd_mod.shutil, "which", return_value = "/usr/bin/amd-smi"): with patch.object(subprocess, "run", return_value = mock_result): result = amd_mod.get_primary_gpu_utilization() @@ -1518,8 +1421,7 @@ class TestAmdGpuMonitoring: except Exception: pytest.skip("Could not load amd module") - # Opt in so the call reaches subprocess.run (gated off on Windows w/o a - # HIP SDK); testing the OSError handling here. + # Opt in so the call reaches subprocess.run (testing OSError handling). with ( patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}), patch.object(subprocess, "run", side_effect = OSError("amd-smi not found")), @@ -1543,8 +1445,7 @@ class TestAmdGpuMonitoring: except Exception: pytest.skip("Could not load amd module") - # Opt in so the call reaches subprocess.run (gated off on Windows w/o a - # HIP SDK); testing the timeout handling here. + # Opt in so the call reaches subprocess.run (testing timeout handling). with ( patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}), patch.object( @@ -1570,9 +1471,7 @@ class TestHardwareAmdBranching: assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): - """get_gpu_utilization should dispatch to amd.py via _smi_query - when IS_ROCM, and the dispatcher itself must check IS_ROCM and - import the amd backend.""" + """get_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") @@ -1585,14 +1484,12 @@ class TestHardwareAmdBranching: assert "from . import amd" in smi def test_hardware_branches_on_is_rocm_for_visible(self): - """get_visible_gpu_utilization should dispatch to amd.py via - _smi_query when IS_ROCM.""" + """get_visible_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_visible_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] - # The dispatcher call may wrap onto multiple lines; allow whitespace - # between the open paren and the literal func name argument. + # The dispatcher call may wrap; allow whitespace before the func name arg. import re as _re assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body) @@ -1616,11 +1513,10 @@ class TestHardwareAmdBranching: class TestApplyGpuIdsRocmFallback: - """Verify apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when - IS_ROCM is still False (worker subprocess before detect_hardware runs).""" + """apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when IS_ROCM is still False (issue #5180).""" def test_apply_gpu_ids_falls_back_to_torch_version_hip(self): - """apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set.""" + """apply_gpu_ids probes torch.version.hip when IS_ROCM is False and no ROCm env vars set.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") @@ -1628,11 +1524,7 @@ class TestApplyGpuIdsRocmFallback: assert 'getattr(_torch.version, "hip", None)' in func_body def test_apply_gpu_ids_sets_hip_but_not_rocr_visible_devices(self): - """apply_gpu_ids should set HIP_VISIBLE_DEVICES but leave ROCR_VISIBLE_DEVICES inherited. - - ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU indices. - Overwriting it breaks multi-GPU ROCm systems (see issue #6118). - """ + """apply_gpu_ids sets HIP_VISIBLE_DEVICES but leaves ROCR_VISIBLE_DEVICES inherited (HSA indexing; issue #6118).""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def apply_gpu_ids") @@ -1802,9 +1694,8 @@ class TestDetectWindowsGfxArch: """Verify hipinfo parsing for GPU arch detection on Windows.""" def test_returns_none_when_hipinfo_not_on_path(self): - # Also neutralise the venv-hipInfo and WMI-name fallbacks: this test - # pins "no probe source available -> None", and the suite may run on a - # real AMD host where WMI would legitimately answer. + # Neutralise the venv-hipInfo and WMI-name fallbacks too, since the + # suite may run on a real AMD host where WMI would answer. with patch("shutil.which", return_value = None): with patch("os.path.isfile", return_value = False): with patch("subprocess.run", side_effect = FileNotFoundError): @@ -1821,12 +1712,9 @@ class TestDetectWindowsGfxArch: assert result == "gfx1200" def test_returns_arch_on_crash_with_gcnarchname_in_output(self): - # Regression test for issue #6043: hipinfo may exit with a non-zero - # code (e.g. 0xC0000005 / STATUS_ACCESS_VIOLATION on RDNA 4 hosts) - # while still printing the gcnArchName line before crashing. The - # previous guard `if result.returncode == 0` discarded this output, - # causing a CPU PyTorch fallback. The fix: accept the arch whenever - # gcnArchName is present in stdout regardless of exit code. + # Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after + # printing gcnArchName. Accept the arch whenever gcnArchName is in + # stdout, regardless of exit code (previously a CPU fallback). mock_result = MagicMock() mock_result.returncode = -1073741819 # 0xC0000005 STATUS_ACCESS_VIOLATION mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n" @@ -1836,8 +1724,7 @@ class TestDetectWindowsGfxArch: assert result == "gfx1200" def test_returns_none_on_nonzero_returncode_without_gcnarchname(self): - # Non-zero exit without any gcnArchName output (e.g. no device detected) - # must still return None and fall through to amd-smi / WMI. + # Non-zero exit without gcnArchName must return None (fall through to amd-smi/WMI). mock_result = MagicMock() mock_result.returncode = 1 mock_result.stdout = b"HIP runtime error: no device detected\n" @@ -1847,10 +1734,8 @@ class TestDetectWindowsGfxArch: assert result is None def test_returns_none_when_no_gcnarchname_in_output(self): - # hipinfo answers but without a gcnArchName line. Route only the - # hipinfo/amd-smi probes to the mock; the WMI fallback must get - # nothing (FileNotFoundError) -- otherwise the mocked device name - # would legitimately resolve via the name table. + # hipinfo answers without a gcnArchName line. The WMI fallback must get + # nothing (FileNotFoundError) so the mocked name can't resolve via the table. mock_result = MagicMock() mock_result.returncode = 0 mock_result.stdout = b"deviceName : SomeUnknownDevice\n" @@ -1888,10 +1773,7 @@ class TestDetectWindowsGfxArch: class TestGfxArchNameFallback: - """amd-smi does not exist on Windows (neither Adrenalin consistently nor - the HIP SDK ship a CLI) and driver-only hosts lack hipinfo too. The - detection chain must still resolve the arch from the GPU marketing name - (WMI), mirroring setup.ps1's $nameArchTable.""" + """With no amd-smi/hipinfo on Windows, arch must resolve from the GPU name via WMI (mirrors setup.ps1).""" @pytest.mark.parametrize( "name, expected", @@ -1958,20 +1840,17 @@ class TestGfxArchNameFallback: assert result is None def test_stack_probes_venv_hipinfo(self): - """The venv Scripts dir hipInfo.exe (shipped by AMD torch wheels) must - be a probe candidate so `studio update` works on driver-only hosts.""" + """venv Scripts hipInfo.exe (from AMD torch wheels) must be a probe candidate for driver-only hosts.""" source = _STACK_PATH.read_text(encoding = "utf-8") assert 'os.path.join(os.path.dirname(sys.executable), "hipInfo.exe")' in source def test_prebuilt_resolve_exe_probes_venv_dir(self): - """install_llama_prebuilt's _resolve_exe must include the venv Scripts - candidate for the same driver-only standalone-rerun scenario.""" + """_resolve_exe must include the venv Scripts candidate for driver-only standalone reruns.""" source = _PREBUILT_PATH.read_text(encoding = "utf-8") assert "_venv_candidate" in source def test_runtime_monitor_guards_amd_smi_absence(self): - """amd.py must which()-check amd-smi before spawning so absence - disables the poller in one step (no FileNotFoundError strikes).""" + """amd.py must which()-check amd-smi before spawning (absence disables the poller).""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" source = amd_path.read_text(encoding = "utf-8") assert 'shutil.which("amd-smi") is None' in source @@ -2002,7 +1881,7 @@ class TestInstallBnbWindowsRocm: call_args = str(mock_pip.call_args_list[0]) assert "bitsandbytes" in call_args assert "win_amd64" in call_args - # Must force plain pip (uv mangles the bitsandbytes wheel) -- see + # Force plain pip (uv mangles the bitsandbytes wheel) -- see # https://unsloth.ai/docs/get-started/install/amd/amd-hackathon assert mock_pip.call_args.kwargs.get("force_pip") is True @@ -2333,11 +2212,7 @@ class TestDetectBnbRocmDllVer: assert stack_mod._detect_bnb_rocm_dll_ver() is None def test_picks_highest_suffix_when_multiple_dlls(self, tmp_path): - """Returns the highest numeric suffix when multiple ROCm DLL variants exist. - - Filesystem glob order is not guaranteed, so the function must not stop - at the first match — it must always return the highest one. - """ + """Returns the highest numeric suffix across ROCm DLL variants (glob order is not guaranteed).""" (tmp_path / "libbitsandbytes_rocm72.dll").write_text("") (tmp_path / "libbitsandbytes_rocm713.dll").write_text("") mock_spec = MagicMock() @@ -2356,7 +2231,7 @@ class TestRocmTorchInstalledEnvVar: @staticmethod def _ok_torch_probe(*a, **kw): - # subprocess.run probe returns 0 when torch imports as ROCm + # Probe returns 0 when torch imports as ROCm. rv = MagicMock() rv.returncode = 0 return rv @@ -2413,7 +2288,7 @@ class TestRocmTorchInstalledEnvVar: patch.object(stack_mod, "IS_MACOS", True), ): stack_mod._ensure_rocm_torch() - # macOS branch is the next exit -- but the point is the early-return did NOT fire. + # macOS branch is the next exit; the point is the early-return did NOT fire. mock_bnb.assert_not_called() @@ -2481,27 +2356,18 @@ class TestWorkerWindowsRocmPatches: assert "TORCHDYNAMO_DISABLE" in source def test_bnb_rocm_version_set_on_windows_rocm(self): - """worker.py must set BNB_ROCM_VERSION in the Windows ROCm section. - - BNB auto-detects HIP version from torch.version.hip, which can mismatch - the DLL suffix in the AMD prerelease wheel. The worker must detect the - actual DLL suffix and override BNB's auto-detection before ML imports. - """ + """worker.py must set BNB_ROCM_VERSION from the detected DLL suffix (BNB's auto-detect can mismatch).""" source = _WORKER_PATH.read_text(encoding = "utf-8") - # Env var must be set assert "BNB_ROCM_VERSION" in source - # Detection helper must be used assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source - # Falls back to the seeded value, never a blind "72" (which would - # force a ROCm backend onto a non-ROCm bitsandbytes wheel) + # Falls back to the seeded value, never a blind "72". assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION")' in source def test_bnb_rocm_version_set_before_ml_imports(self): """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports.""" source = _WORKER_PATH.read_text(encoding = "utf-8") idx_bnb = source.find("BNB_ROCM_VERSION") - # Use the specific section-2 marker that appears in the worker process - # entry-point function (not the trainer helper which has its own "# ── 2."). + # Use the entry-point section-2 marker (not the trainer helper's own "# ── 2."). idx_sec2 = source.find("# ── 2. Now import ML libraries") assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py" assert idx_sec2 != -1, "'# ── 2. Now import ML libraries' marker not found in worker.py" @@ -2513,9 +2379,8 @@ class TestWorkerWindowsRocmPatches: def test_grouped_mm_patch_guarded_by_windows_and_hip_check(self): """_grouped_mm patch must only apply on Windows + HIP torch.""" source = _WORKER_PATH.read_text(encoding = "utf-8") - # Must check sys.platform == "win32" assert 'sys.platform == "win32"' in source - # Must gate on HIP version — code uses getattr chain: "version" and "hip" + # Gates on HIP version via a getattr chain ("version", "hip"). assert '"version"' in source and '"hip"' in source def test_hip_ver_at_least_helper_defined(self): @@ -2526,9 +2391,8 @@ class TestWorkerWindowsRocmPatches: def test_grouped_mm_patch_gated_on_hip_lt_713(self): """_grouped_mm patch must be skipped on HIP >= 7.13 (AMD fixed the bug in ROCm 7.13).""" source = _WORKER_PATH.read_text(encoding = "utf-8") - # The guard must call _hip_ver_at_least with exactly (7, 13) assert "_hip_ver_at_least(7, 13)" in source - # The patch must be inside the `if not` branch (negated guard) + # Patch must be inside the negated `if not` guard. assert "if not _hip_ver_at_least(7, 13):" in source def test_grouped_mm_hip_713_skip_message_present(self): @@ -2540,17 +2404,14 @@ class TestWorkerWindowsRocmPatches: def test_grouped_mm_patch_else_branch_present(self): """An else branch must follow the _hip_ver_at_least gate (skip path for 7.13+).""" source = _WORKER_PATH.read_text(encoding = "utf-8") - # There must be an else: after the if not _hip_ver_at_least(7, 13): block gate_idx = source.find("if not _hip_ver_at_least(7, 13):") assert gate_idx != -1, "Version gate not found in worker.py" - # The else: branch must appear after the gate else_idx = source.find("else:", gate_idx) assert else_idx != -1, "else: branch after _hip_ver_at_least gate not found" def test_hip_ver_at_least_handles_amd_version_format(self): """_hip_ver_at_least must split on '.' and compare only major.minor (handles '7.13.99004').""" source = _WORKER_PATH.read_text(encoding = "utf-8") - # Must split the version string and take the first two parts assert 'split(".")[:2]' in source or ".split('.')[:2]" in source @@ -2608,8 +2469,7 @@ _INSTALL_PS1_PATH = PACKAGE_ROOT / "install.ps1" class TestStrixHaloGfxArchDetection: - """Verify that setup.ps1 and install.ps1 have robust gfx arch detection - for Strix Halo / iGPU users who only have the HIP runtime (no hipinfo).""" + """setup.ps1 / install.ps1 gfx arch detection for Strix Halo / iGPU (HIP runtime only, no hipinfo).""" def test_amd_smi_static_asic_attempted_in_setup(self): """setup.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch.""" @@ -2664,12 +2524,9 @@ class TestStrixHaloGfxArchDetection: def test_wmi_does_not_set_hasrocm_in_setup(self): """WMI block in setup.ps1 must NOT set $HasROCm = $true (no runtime confirmation).""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") - # Find the WMI block and confirm HasROCm is not set inside it wmi_idx = source.find("Win32_VideoController") assert wmi_idx != -1, "WMI block not found in setup.ps1" - # The nearest HasROCm = $true must not appear between the WMI block - # and the closing brace of that if-block. We check by confirming - # $HasROCm = $true does NOT appear within 300 chars of the WMI call. + # $HasROCm = $true must not appear within 300 chars of the WMI call. wmi_context = source[wmi_idx : wmi_idx + 300] assert "$HasROCm = $true" not in wmi_context @@ -2677,7 +2534,6 @@ class TestStrixHaloGfxArchDetection: """Both files must use the gfx\\d+[a-z]? regex to parse arch from amd-smi output.""" for path in (_SETUP_PS1_PATH, _INSTALL_PS1_PATH): source = path.read_text(encoding = "utf-8") - # The regex pattern used to match gfx arches assert ( "gfx\\d+" in source or r"gfx\d+" in source ), f"gfx arch regex not found in {path.name}" @@ -2687,8 +2543,7 @@ class TestStrixHaloGfxArchDetection: class TestHipSdkEnvPathResolution: - """Verify that both install scripts resolve hipinfo/hipconfig via HIP_PATH - and ROCM_PATH when the tools are not on $PATH, and emit explicit warnings.""" + """Both install scripts resolve hipinfo/hipconfig via HIP_PATH/ROCM_PATH off $PATH, and warn.""" @staticmethod def _assert_accepts_partial_hipinfo_output(source: str): @@ -2717,7 +2572,6 @@ class TestHipSdkEnvPathResolution: """setup.ps1 must also check ROCM_PATH as a secondary hipinfo fallback.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "ROCM_PATH" in source - # Confirm the fallback pattern: HIP_PATH ?? ROCM_PATH (or equivalent elseif) assert "ROCM_PATH" in source and "HIP_PATH" in source def test_install_checks_rocm_path_as_hipinfo_fallback(self): @@ -2753,7 +2607,6 @@ class TestHipSdkEnvPathResolution: def test_setup_warns_when_hip_path_set_but_exe_missing(self): """setup.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") - # The warning must mention that the SDK install may be incomplete assert "incomplete" in source or "not found at" in source def test_install_warns_when_hip_path_set_but_exe_missing(self): @@ -2810,7 +2663,6 @@ class TestHipSdkEnvPathResolution: def test_setup_provides_path_fix_hint(self): """setup.ps1 must tell the user how to add the HIP bin dir to PATH.""" source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") - # Should mention adding to PATH or SetEnvironmentVariable assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source) def test_install_provides_path_fix_hint(self): @@ -2823,8 +2675,7 @@ class TestHipSdkEnvPathResolution: class TestHipSdkDetectedSubstep: - """Verify that both scripts print HIP SDK path and full hipconfig version - as substeps under the gpu step when AMD ROCm is successfully detected.""" + """Both scripts print HIP SDK path and full hipconfig version as substeps when ROCm is detected.""" def test_setup_prints_hip_sdk_path_substep(self): """setup.ps1 must print an 'HIP SDK:' substep showing the resolved path.""" @@ -2879,9 +2730,7 @@ _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh" class TestStrixRocm71Override: - """Verify install.sh skips Radeon repo and routes to AMD arch-specific index - for gfx1151/gfx1150 when ROCm 7.1 would otherwise be selected (known _grouped_mm segfault). - AMD's repo.amd.com/rocm/whl/gfx1151/ serves torch 2.11+rocm7.13 which has the real fix.""" + """install.sh routes gfx1151/gfx1150 to AMD's arch index instead of ROCm 7.1 (_grouped_mm segfault).""" def test_strix_gfx_detection_in_install_sh(self): """install.sh must detect gfx1151 and gfx1150 for the override.""" @@ -2891,10 +2740,9 @@ class TestStrixRocm71Override: def test_rocm71_override_to_amd_arch_index_in_install_sh(self): """install.sh must override TORCH_INDEX_URL to AMD arch-specific index for Strix.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") - # The override must route to AMD's arch-specific index (repo.amd.com/rocm/whl) assert "repo.amd.com/rocm/whl" in source assert "_strix_gfx" in source - # The URL must incorporate the detected gfx arch so gfx1151 → .../gfx1151/ + # URL must incorporate the detected gfx arch (gfx1151 -> .../gfx1151/). strix_idx = source.find("_amd_strix_base") assert strix_idx != -1 ctx = source[strix_idx : strix_idx + 500] @@ -2915,7 +2763,6 @@ class TestStrixRocm71Override: source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") strix_idx = source.find("_strix_gfx") assert strix_idx != -1 - # Look back for the rocm7.1 pattern within 600 chars before _strix_gfx context_before = source[max(0, strix_idx - 2400) : strix_idx] assert "rocm7.1" in context_before @@ -2932,7 +2779,6 @@ class TestStrixRocm71Override: def test_tauri_family_recognises_amd_arch_url(self): """_tauri_torch_index_family must return a rocm* family for AMD arch-specific URLs.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") - # The function must have a case branch for repo.amd.com/rocm/whl/gfx* URLs assert "rocm/whl/gfx" in source @@ -2940,8 +2786,7 @@ class TestStrixRocm71Override: class TestSetupShGccInstallDir: - """Verify setup.sh applies the --gcc-install-dir flag when building llama.cpp - with HIP on Ubuntu 24.04+ to work around ROCm 7.x clang-20 header path bug.""" + """setup.sh applies --gcc-install-dir for HIP builds on Ubuntu 24.04+ (ROCm 7.x clang-20 header bug).""" def test_gcc_install_dir_search_loop_present(self): """setup.sh must iterate gcc versions 14→11 to find one with C++ headers.""" @@ -2963,7 +2808,6 @@ class TestSetupShGccInstallDir: def test_gcc_install_dir_only_applied_in_hip_build_block(self): """The --gcc-install-dir fix must only apply in the HIP/ROCm build branch.""" source = _SETUP_SH_PATH.read_text(encoding = "utf-8") - # GGML_HIP=ON must appear before gcc-install-dir in the source hip_idx = source.find("GGML_HIP=ON") gcc_idx = source.find("gcc-install-dir") assert hip_idx != -1 and gcc_idx != -1 @@ -2982,8 +2826,7 @@ _HARDWARE_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / class TestServerStartupRocmFixes: - """Verify main.py sets BNB_ROCM_VERSION before any bitsandbytes import and - hardware.py injects torch._C._distributed_c10d stubs before torch.distributed.""" + """main.py sets BNB_ROCM_VERSION pre-bnb-import; hardware.py stubs _distributed_c10d pre-torch.distributed.""" # ── BNB_ROCM_VERSION in server process ──────────────────────────────────── @@ -3016,23 +2859,19 @@ class TestServerStartupRocmFixes: assert '"BNB_ROCM_VERSION" not in os.environ' in source # ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ──────────────── - # bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via subprocess PATH - # at import time. The AMD torch wheel ships hipInfo.exe in the venv - # Scripts dir, which is on PATH only for activated venvs -- Studio and the - # installer launch python directly, so without the prepend every bnb - # import logs "Could not detect ROCm GPU architecture: [WinError 2]". + # bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD + # wheel ships it in venv Scripts (on PATH only for activated venvs), so + # without the prepend bnb logs "[WinError 2]" when launched directly. def test_main_py_prepends_hipinfo_dir_to_path(self): """main.py must make hipInfo.exe resolvable before bnb imports.""" source = _MAIN_PY_PATH.read_text(encoding = "utf-8") assert "hipInfo.exe" in source - # The prepend must come before the BNB_ROCM_VERSION block (both run - # pre-import; order documents that bnb sees the fixed PATH). + # Prepend must precede the BNB_ROCM_VERSION block so bnb sees the fixed PATH. assert source.find("hipInfo.exe") < source.find("BNB_ROCM_VERSION") def test_main_py_hipinfo_prepend_gated_on_file_presence(self): - """Only AMD ROCm wheels ship hipInfo.exe; NVIDIA/CPU hosts must be - untouched, so the prepend must check the file exists first.""" + """Prepend must check hipInfo.exe exists first (only AMD wheels ship it; leave NVIDIA/CPU untouched).""" source = _MAIN_PY_PATH.read_text(encoding = "utf-8") assert 'os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe"))' in source @@ -3042,8 +2881,7 @@ class TestServerStartupRocmFixes: assert "hipInfo.exe" in source def test_install_stack_prepends_hipinfo_dir_to_path(self): - """install_python_stack.py must prepend so the installer's child - import checks inherit a PATH where bnb's probe succeeds.""" + """install_python_stack.py must prepend so child import checks inherit a PATH where bnb's probe works.""" source = _STACK_PATH.read_text(encoding = "utf-8") assert "hipInfo.exe" in source @@ -3099,8 +2937,7 @@ class TestServerStartupRocmFixes: class TestHipSdkInstalledButDeviceInaccessible: - """Verify that when hipinfo is found but exits non-zero (device not ROCm-accessible), - both scripts distinguish this from 'HIP SDK not found' and emit the correct message.""" + """When hipinfo is found but exits non-zero, both scripts distinguish device-inaccessible from SDK-not-found.""" def test_install_ps1_has_hip_sdk_installed_flag(self): """install.ps1 must track HipSdkInstalled separately from HasROCm.""" @@ -3115,7 +2952,6 @@ class TestHipSdkInstalledButDeviceInaccessible: def test_install_ps1_sets_flag_when_hipinfo_binary_found(self): """install.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block.""" source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") - # HipSdkInstalled must be assigned inside the hipinfoExe block hipinfo_block_idx = source.find("if ($hipinfoExe)") sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx) assert hipinfo_block_idx != -1 and sdk_flag_idx != -1 @@ -3165,9 +3001,8 @@ class TestHipSdkInstalledButDeviceInaccessible: assert "GPU not ROCm-accessible" in source -# TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 hand their resolved gfx arch -# to install_llama_prebuilt.py so the per-gfx ROCm prebuilt is selected even when -# the installer's own hipinfo/amd-smi probe cannot report it. +# TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 forward their resolved gfx +# arch to install_llama_prebuilt.py so the per-gfx prebuilt is picked. _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh" @@ -3196,7 +3031,7 @@ class TestApplyHostOverrides: """Forwarded ROCm detection is folded into the host profile correctly.""" def test_forwarded_gfx_fills_empty_probe(self): - # amd-smi-only / name-inferred host: installer probe found no gfx. + # Installer probe found no gfx (amd-smi-only / name-inferred host). host = rocm_host(rocm_gfx_target = None) out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151") assert out.has_rocm is True @@ -3235,7 +3070,7 @@ class TestRocmGfxForwarding: def test_installer_exposes_rocm_gfx_arg(self): source = _PREBUILT_PATH.read_text(encoding = "utf-8") assert '"--rocm-gfx"' in source - # Defaults to the env override so a standalone run still works. + # Defaults to the env override for standalone runs. assert 'os.environ.get("UNSLOTH_ROCM_GFX_ARCH")' in source def test_setup_sh_forwards_rocm_gfx(self): @@ -3244,8 +3079,7 @@ class TestRocmGfxForwarding: assert '"$_setup_gfx"' in source def test_setup_sh_forwards_has_rocm(self): - # When AMD is detected but gfx resolution fails, setup.sh must still - # forward --has-rocm so the installer knows ROCm is present. + # If AMD is detected but gfx resolution fails, --has-rocm is still forwarded. source = _SETUP_SH_PATH.read_text(encoding = "utf-8") assert "--has-rocm" in source assert "_setup_amd_detected" in source @@ -3256,24 +3090,20 @@ class TestRocmGfxForwarding: assert "$script:ROCmGfxArch" in source def test_setup_sh_routes_inferred_gfx_to_fork(self): - # A forwarded/inferred gfx arch must route to the fork even without ROCm - # tooling on PATH, so the per-gfx prebuilt is picked over ggml-org. Pin - # the routing guard specifically -- a bare "${_setup_gfx:-}" check also - # appears in the unrelated --rocm-gfx forwarding block. + # An inferred gfx arch must route to the fork even without ROCm tooling. + # Pin the specific guard (a bare "${_setup_gfx:-}" also appears elsewhere). source = _SETUP_SH_PATH.read_text(encoding = "utf-8") assert '[ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]' in source def test_setup_ps1_routes_inferred_gfx_to_fork(self): - # Same on Windows: a resolved $script:ROCmGfxArch counts as a fork/GPU - # install even when $HasROCm is false (Adrenalin-only, no HIP runtime). + # On Windows, a resolved $script:ROCmGfxArch counts as a fork install + # even when $HasROCm is false (Adrenalin-only, no HIP runtime). source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") assert "$HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch" in source - # The two assertions above pin the guard *text*. The tests below *execute* - # the real routing block from setup.sh / setup.ps1 and assert the resolved - # release repo, so a refactor that keeps the literal but breaks (or drops) - # the inferred-gfx -> fork decision is still caught. All inputs are faked -- - # no GPU, no ROCm tooling on PATH, no network. + # The assertions above pin the guard *text*; the tests below *execute* the + # real routing block and assert the resolved repo, so a refactor that keeps + # the literal but breaks the inferred-gfx -> fork decision is still caught. @staticmethod def _resolve_setup_sh_repo( @@ -3282,11 +3112,7 @@ class TestRocmGfxForwarding: setup_gfx, rocm_gfx_arch_env = "", ): - """Run setup.sh's release-repo routing block under bash and return the - resolved _HELPER_RELEASE_REPO. PATH is emptied so the rocminfo/amd-smi/ - hipconfig/hipinfo `command -v` probes all miss (no ROCm tooling). - rocm_gfx_arch_env populates UNSLOTH_ROCM_GFX_ARCH for the env-forwarded - path that fires when no probe set _setup_gfx.""" + """Run setup.sh's routing block under bash with PATH emptied (no ROCm tooling) and return _HELPER_RELEASE_REPO.""" import shutil bash = shutil.which("bash") @@ -3298,7 +3124,7 @@ class TestRocmGfxForwarding: block = source[start:end] assert "_HELPER_RELEASE_REPO" in block, "setup.sh routing anchors not found" env = { - "PATH": "", # no rocminfo/amd-smi/hipconfig/hipinfo discoverable + "PATH": "", # no ROCm tooling discoverable "ROUTING_BLOCK": block, "_HOST_SYSTEM": "Linux", "_HOST_MACHINE": host_machine, @@ -3317,30 +3143,24 @@ class TestRocmGfxForwarding: return result.stdout.strip() def test_setup_sh_inferred_gfx_resolves_to_fork(self): - # No usable NVIDIA, no ROCm tooling on PATH, only a name-inferred gfx - # arch -> the host must still be treated as a GPU host and routed to the - # fork's per-gfx prebuilt, not ggml-org / a source build. Linux x64 and - # arm64 both go through the same fork branch. + # Only a name-inferred gfx arch -> route to the fork's per-gfx prebuilt + # (not ggml-org). x64 and arm64 share the fork branch. assert self._resolve_setup_sh_repo("x86_64", False, "gfx1100") == "unslothai/llama.cpp" assert self._resolve_setup_sh_repo("aarch64", False, "gfx1100") == "unslothai/llama.cpp" def test_setup_sh_env_forwarded_gfx_resolves_to_fork(self): - # UNSLOTH_ROCM_GFX_ARCH set on a host where no probe fired (_setup_gfx - # empty, no usable NVIDIA, no ROCm tooling): setup.sh adopts the env arch - # and routes to the fork, same as the name-inference path. + # No probe fired but UNSLOTH_ROCM_GFX_ARCH is set: adopt the env arch + # and route to the fork, same as name-inference. repo = self._resolve_setup_sh_repo("x86_64", False, "", rocm_gfx_arch_env = "gfx1100") assert repo == "unslothai/llama.cpp" def test_setup_sh_cpu_host_still_resolves_to_ggml(self): - # Guard against over-correcting the fix: a real CPU host (no usable GPU, - # no inferred gfx, no env override) must keep routing to ggml-org for the - # CPU prebuilt. + # A real CPU host (no GPU, no inferred gfx, no env override) must keep routing to ggml-org. assert self._resolve_setup_sh_repo("x86_64", False, "") == "ggml-org/llama.cpp" @staticmethod def _resolve_setup_ps1_repo(has_nvidia, has_rocm, gfx_arch): - """Run setup.ps1's $HelperReleaseRepo selection under pwsh and return the - resolved repo.""" + """Run setup.ps1's $HelperReleaseRepo selection under pwsh and return the resolved repo.""" import shutil pwsh = shutil.which("pwsh") @@ -3373,8 +3193,7 @@ class TestRocmGfxForwarding: return result.stdout.strip() def test_setup_ps1_inferred_gfx_resolves_to_fork(self): - # Adrenalin-only Windows host: $HasROCm is false (no HIP runtime) but a - # gfx arch was inferred -> route to the fork's windows-rocm bundle. + # Adrenalin-only host: $HasROCm false but gfx inferred -> fork's windows-rocm bundle. assert self._resolve_setup_ps1_repo(False, False, "gfx1100") == "unslothai/llama.cpp" def test_setup_ps1_cpu_host_still_resolves_to_ggml(self): @@ -3383,16 +3202,15 @@ class TestRocmGfxForwarding: # TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output. -# Honours CUDA_VISIBLE_DEVICES/HIP_VISIBLE_DEVICES so a mixed-arch host installs -# the prebuilt for the GPU actually selected, not GPU 0. +# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt +# for the selected GPU, not GPU 0. _pick_rocm_gfx_target = prebuilt_mod._pick_rocm_gfx_target def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch): - """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES; - on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100.""" - # Two GPUs; rocminfo reports each token twice (as in the real tool output). + """CUDA_VISIBLE_DEVICES=1 must select gfx1100 on a gfx1151 + gfx1100 host (HIP honours CUDA var).""" + # rocminfo reports each token twice (as in the real tool output). probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100" monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) @@ -3410,11 +3228,8 @@ def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkey def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): - """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must - return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the - two gfx1100 entries into one and making index 2 out of range.""" - # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU). - # Each GPU gets its own Agent section with a few token mentions. + """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must return gfx1151 (no dict.fromkeys collapse).""" + # rocminfo output for 3 GPUs (2x gfx1100 + 1x gfx1151), one Agent section each. probe_out = ( "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n" "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n" @@ -3434,9 +3249,7 @@ _LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / " class TestWslSystemRocmLibDirs: - """install_llama_prebuilt._wsl_system_rocm_lib_dirs: a strict no-op off a - ROCDXG WSL host; returns the system ROCm lib dir there so binary_env can - load the WSL-capable HIP runtime before a prebuilt's bundled one.""" + """_wsl_system_rocm_lib_dirs: no-op off a ROCDXG WSL host; else returns the system ROCm lib dir for binary_env.""" def test_empty_without_dev_dxg(self): with patch("os.path.exists", return_value = False): @@ -3452,8 +3265,7 @@ class TestWslSystemRocmLibDirs: assert prebuilt_mod._wsl_system_rocm_lib_dirs() == [] def test_returns_system_lib_on_wsl_with_librocdxg(self): - # Normalize separators: os.path.join uses "\" on the Windows test host - # though the target path is Linux. + # Normalize separators: os.path.join uses "\" on the Windows test host. def _exists(p): p = str(p).replace("\\", "/") return p in ("/dev/dxg", "/opt/rocm/lib/librocdxg.so") @@ -3476,9 +3288,7 @@ class TestWslSystemRocmLibDirs: class TestBinaryEnvWslOrdering: - """binary_env must put the system ROCm lib dir AHEAD of the prebuilt's - bundle dir on a ROCDXG WSL host, and set HSA_ENABLE_DXG_DETECTION; no-op - on bare-metal Linux.""" + """binary_env puts system ROCm lib ahead of the bundle dir + sets HSA_ENABLE_DXG_DETECTION on WSL; no-op bare-metal.""" @staticmethod def _linux_host(): @@ -3503,8 +3313,7 @@ class TestBinaryEnvWslOrdering: binary = tmp_path / "bundle" / "llama-server" binary.parent.mkdir(parents = True) binary.write_text("") - # binary_env's dedupe_existing_dirs drops non-existent dirs, so use a - # real dir to stand in for the system ROCm lib path. + # dedupe_existing_dirs drops non-existent dirs, so use a real dir. sys_rocm = tmp_path / "sysrocm" sys_rocm.mkdir() with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = [str(sys_rocm)]): @@ -3533,20 +3342,16 @@ class TestBinaryEnvWslOrdering: class TestInstallShDropinPersistence: - """install.sh must persist the ROCm-on-WSL drop-in even when rocminfo - already enumerates the GPU (via the transient probe env), so a reinstall - over an existing /opt/rocm doesn't leave login shells without the env.""" + """install.sh persists the ROCm-on-WSL drop-in even when rocminfo already enumerates the GPU (reinstall safety).""" def test_has_persist_helper(self): source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "_persist_rocm_wsl_dropin()" in source def test_gate5_early_return_persists_dropin(self): - """The rocminfo-already-works early return must call the persist helper - BEFORE returning.""" + """The rocminfo-already-works early return must call the persist helper before returning.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") - # Find the rocminfo gfx1151 gate and assert the persist call precedes - # its `return 0`. + # The persist call must precede `return 0` at the rocminfo gfx1151 gate. gate = source.find("Name:[[:space:]]*gfx1151") assert gate != -1 window = source[gate : gate + 900] @@ -3562,8 +3367,7 @@ class TestInstallShDropinPersistence: class TestLlamaCppRuntimeWslOrdering: - """The serve-time launcher must mirror binary_env: system HIP before the - bundle dir on WSL, so a prebuilt that passed install validation runs.""" + """The serve-time launcher mirrors binary_env: system HIP before the bundle dir on WSL.""" def test_has_wsl_helper(self): source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") @@ -3571,7 +3375,6 @@ class TestLlamaCppRuntimeWslOrdering: def test_prepends_before_binary_dir(self): source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") - # The Linux lib_dirs build must add WSL rocm dirs before binary_dir. idx_helper = source.find("for _wsl_rocm in _wsl_system_rocm_lib_dirs()") idx_binary = source.find("lib_dirs.append(binary_dir)") assert idx_helper != -1 and idx_binary != -1 diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 6f15ca6c72..bb82811614 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -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--bin-macos-.tar.gz with pinned deployment targets, selected - by install_kind.""" + """macOS routes to the fork's llama--bin-macos-.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 diff --git a/tests/studio/load_freeze/llama_server_shim.py b/tests/studio/load_freeze/llama_server_shim.py index 4030fcd6c4..a6efce95b9 100644 --- a/tests/studio/load_freeze/llama_server_shim.py +++ b/tests/studio/load_freeze/llama_server_shim.py @@ -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 = "/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] diff --git a/tests/studio/load_freeze/test_load_orchestrator.py b/tests/studio/load_freeze/test_load_orchestrator.py index d091e56e70..64503d0e05 100644 --- a/tests/studio/load_freeze/test_load_orchestrator.py +++ b/tests/studio/load_freeze/test_load_orchestrator.py @@ -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 "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 "", 128259: ""} ) 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 ; keep 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 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 `.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 diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py index 7686013fcf..9c01e95fd4 100644 --- a/tests/studio/playwright_chat_ime_i18n.py +++ b/tests/studio/playwright_chat_ime_i18n.py @@ -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] diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 9cfb33d427..f0a928982d 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -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 element intercepting pointer events - # for a beat after each route swap -- Playwright surfaces this as - # " intercepts pointer events" on the next click. + # Hard-disable CSS view-transitions: Studio's theme toggle + sidebar + # collapse run startViewTransition() which can leave intercepting + # pointer events for a beat after each route swap. See _playwright_robust.py. install_view_transition_killer(ctx) 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 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 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 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 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 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: diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 26c3c244ca..a23e182114 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -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"]'); diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index e58b2098d8..ea8b8621a9 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -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 _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, ( diff --git a/tests/studio/studio_api_smoke.py b/tests/studio/studio_api_smoke.py index 0390b39c03..845a9ed021 100644 --- a/tests/studio/studio_api_smoke.py +++ b/tests/studio/studio_api_smoke.py @@ -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"" 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: 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}") diff --git a/tests/studio/test_auth_form_input_count.py b/tests/studio/test_auth_form_input_count.py index 80f73e4065..dd582d17d1 100644 --- a/tests/studio/test_auth_form_input_count.py +++ b/tests/studio/test_auth_form_input_count.py @@ -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}" diff --git a/tests/studio/test_cancel_atomicity.py b/tests/studio/test_cancel_atomicity.py index ea4169f6ff..142cc2247a 100644 --- a/tests/studio/test_cancel_atomicity.py +++ b/tests/studio/test_cancel_atomicity.py @@ -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 diff --git a/tests/studio/test_cancel_id_wiring.py b/tests/studio/test_cancel_id_wiring.py index d821786431..651dfbf3de 100644 --- a/tests/studio/test_cancel_id_wiring.py +++ b/tests/studio/test_cancel_id_wiring.py @@ -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:] diff --git a/tests/studio/test_cli_repo_variant.py b/tests/studio/test_cli_repo_variant.py index 21f935965f..548d3f2598 100644 --- a/tests/studio/test_cli_repo_variant.py +++ b/tests/studio/test_cli_repo_variant.py @@ -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 diff --git a/tests/studio/test_cli_run_alias.py b/tests/studio/test_cli_run_alias.py index c69cad9bb8..498ebbdf4d 100644 --- a/tests/studio/test_cli_run_alias.py +++ b/tests/studio/test_cli_run_alias.py @@ -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 ( diff --git a/tests/studio/test_cli_studio_defaults.py b/tests/studio/test_cli_studio_defaults.py index ed1bd37c3d..b03fa57603 100644 --- a/tests/studio/test_cli_studio_defaults.py +++ b/tests/studio/test_cli_studio_defaults.py @@ -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()" diff --git a/tests/studio/test_cli_studio_stop_windows.py b/tests/studio/test_cli_studio_stop_windows.py index 7ba3d12214..778679c73b 100644 --- a/tests/studio/test_cli_studio_stop_windows.py +++ b/tests/studio/test_cli_studio_stop_windows.py @@ -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 diff --git a/tests/studio/test_composer_rtl_bidi_attribute.py b/tests/studio/test_composer_rtl_bidi_attribute.py index f1eb75035c..defbfab86c 100644 --- a/tests/studio/test_composer_rtl_bidi_attribute.py +++ b/tests/studio/test_composer_rtl_bidi_attribute.py @@ -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, ( diff --git a/tests/studio/test_frontend_dep_removal.py b/tests/studio/test_frontend_dep_removal.py index d1835c2bac..fb1b5ace7b 100644 --- a/tests/studio/test_frontend_dep_removal.py +++ b/tests/studio/test_frontend_dep_removal.py @@ -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] = [ '', "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 = [ diff --git a/tests/studio/test_hardware_dispatch_matrix.py b/tests/studio/test_hardware_dispatch_matrix.py index edb8a66b93..62a0fe0447 100644 --- a/tests/studio/test_hardware_dispatch_matrix.py +++ b/tests/studio/test_hardware_dispatch_matrix.py @@ -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", diff --git a/tests/studio/test_is_mlx_dispatch_gate.py b/tests/studio/test_is_mlx_dispatch_gate.py index 594ab39076..85a6e42449 100644 --- a/tests/studio/test_is_mlx_dispatch_gate.py +++ b/tests/studio/test_is_mlx_dispatch_gate.py @@ -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(): diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 3fd91a5188..8d12fd528e 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -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() diff --git a/tests/studio/test_studio_gguf_export_script_pin.py b/tests/studio/test_studio_gguf_export_script_pin.py index b97ecd4367..1f7e7adaa4 100644 --- a/tests/studio/test_studio_gguf_export_script_pin.py +++ b/tests/studio/test_studio_gguf_export_script_pin.py @@ -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 diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py index ac26a10f06..7cad6cacfe 100644 --- a/tests/studio/test_studio_text_descender_clipping.py +++ b/tests/studio/test_studio_text_descender_clipping.py @@ -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'', ) diff --git a/tests/test_callback_signature_drift.py b/tests/test_callback_signature_drift.py index 4a29f4f7c7..11d32b321b 100644 --- a/tests/test_callback_signature_drift.py +++ b/tests/test_callback_signature_drift.py @@ -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.__callbacks`` list, populated - via ``self.__callbacks.append(...)`` from an ``add__callback`` - method, and invoked via ``for cb in self.__callbacks: cb(arg1, ...)``. - The arity at the call site is the canonical expected arity. - * Consumer side: any ``.add__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.__callbacks`` populated by ``add__callback`` and invoked +via ``for cb in self.__callbacks: cb(...)`` (the call-site arity is canonical). +Consumer: ``.add__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 .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 .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 .add_*_callback(fn) sites + # Find .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") diff --git a/tests/test_cli_export_unpacking.py b/tests/test_cli_export_unpacking.py index 626ff41bc1..b67f727c60 100644 --- a/tests/test_cli_export_unpacking.py +++ b/tests/test_cli_export_unpacking.py @@ -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 diff --git a/tests/test_enforce_kwargs_spacing.py b/tests/test_enforce_kwargs_spacing.py index 4521b0362d..03220c61a3 100644 --- a/tests/test_enforce_kwargs_spacing.py +++ b/tests/test_enforce_kwargs_spacing.py @@ -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 diff --git a/tests/test_finetune_last_n_layers.py b/tests/test_finetune_last_n_layers.py index fa708ddd3e..35baef4d87 100644 --- a/tests/test_finetune_last_n_layers.py +++ b/tests/test_finetune_last_n_layers.py @@ -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 diff --git a/tests/test_gemma4_chat_template.py b/tests/test_gemma4_chat_template.py index 5c362576c6..cfbc81f736 100644 --- a/tests/test_gemma4_chat_template.py +++ b/tests/test_gemma4_chat_template.py @@ -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") diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index c9f971c891..68c20fbe53 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -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.") diff --git a/tests/test_loader_glob_skip.py b/tests/test_loader_glob_skip.py index f26f71271c..ade9e89fde 100644 --- a/tests/test_loader_glob_skip.py +++ b/tests/test_loader_glob_skip.py @@ -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() diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index bbcfdd8aa9..283b099107 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -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] diff --git a/tests/test_multi_image_grpo_chunking.py b/tests/test_multi_image_grpo_chunking.py index ccd10f905a..ea142ce1ef 100644 --- a/tests/test_multi_image_grpo_chunking.py +++ b/tests/test_multi_image_grpo_chunking.py @@ -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 diff --git a/tests/test_nvfp4_quant_load.py b/tests/test_nvfp4_quant_load.py index 988fae5c97..b7cbcde88c 100644 --- a/tests/test_nvfp4_quant_load.py +++ b/tests/test_nvfp4_quant_load.py @@ -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 diff --git a/tests/test_public_api_surface.py b/tests/test_public_api_surface.py index 140fda99f5..264ef64088 100644 --- a/tests/test_public_api_surface.py +++ b/tests/test_public_api_surface.py @@ -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)) diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py index b930e7ab31..d7f6c317fe 100644 --- a/tests/test_raw_text.py +++ b/tests/test_raw_text.py @@ -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 = "" - 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) diff --git a/tests/test_studio_install_workspace_guard.py b/tests/test_studio_install_workspace_guard.py index 27035b522d..68c1a2db50 100644 --- a/tests/test_studio_install_workspace_guard.py +++ b/tests/test_studio_install_workspace_guard.py @@ -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] diff --git a/tests/test_studio_root_resilience.py b/tests/test_studio_root_resilience.py index 9037d6965c..5b835f770b 100644 --- a/tests/test_studio_root_resilience.py +++ b/tests/test_studio_root_resilience.py @@ -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) diff --git a/tests/test_video_path_validation.py b/tests/test_video_path_validation.py index d1406230e3..cb101ead6c 100644 --- a/tests/test_video_path_validation.py +++ b/tests/test_video_path_validation.py @@ -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) diff --git a/tests/test_windows_rocm_bnb_version.py b/tests/test_windows_rocm_bnb_version.py index c8439b1fcd..a3aea4d74e 100644 --- a/tests/test_windows_rocm_bnb_version.py +++ b/tests/test_windows_rocm_bnb_version.py @@ -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") diff --git a/tests/utils/aime_eval.py b/tests/utils/aime_eval.py index 6b7db2026d..0f291c88dc 100644 --- a/tests/utils/aime_eval.py +++ b/tests/utils/aime_eval.py @@ -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"]), diff --git a/tests/utils/cleanup_utils.py b/tests/utils/cleanup_utils.py index c04523a7a4..005ecffb1f 100644 --- a/tests/utils/cleanup_utils.py +++ b/tests/utils/cleanup_utils.py @@ -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", diff --git a/tests/utils/data_utils.py b/tests/utils/data_utils.py index aa36a2afb4..6a7cc7f47b 100644 --- a/tests/utils/data_utils.py +++ b/tests/utils/data_utils.py @@ -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 = { diff --git a/tests/utils/generate_dataset_with_none.py b/tests/utils/generate_dataset_with_none.py index 64df3f8626..78c6dfbe54 100644 --- a/tests/utils/generate_dataset_with_none.py +++ b/tests/utils/generate_dataset_with_none.py @@ -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) diff --git a/tests/utils/hf_utils.py b/tests/utils/hf_utils.py index 22d984f87a..be56d89596 100644 --- a/tests/utils/hf_utils.py +++ b/tests/utils/hf_utils.py @@ -211,10 +211,7 @@ def setup_lora( def convert_weights_back_to_dtype(model, dtype): - """ - SFTTrainer calls get_peft_model and prepare_model_for_kbit_training which converts all weights to float32. - This function converts the non-loraweights back to the original dtype. - """ + """Convert non-LoRA weights back to the original dtype (SFTTrainer upcasts them to float32).""" for name, param in model.named_parameters(): if any(s in name for s in ["norm", "embed"]): param.data = param.data.to(dtype) @@ -225,7 +222,7 @@ def fix_llama3_tokenizer(tokenizer, padding_side = "right"): added_vocab = tokenizer.get_added_vocab() pad_token = [w for w in added_vocab if "pad" in w] assert len(pad_token) == 1 - tokenizer.pad_token = pad_token[0] # Load dataset from the hub + tokenizer.pad_token = pad_token[0] return tokenizer diff --git a/tests/utils/ocr_eval.py b/tests/utils/ocr_eval.py index f02cd5a26a..60cde14fb2 100644 --- a/tests/utils/ocr_eval.py +++ b/tests/utils/ocr_eval.py @@ -1,9 +1,4 @@ -""" -OCR Model Evaluation Module - -This module provides functionality to evaluate OCR models on datasets with -word error rate (WER) and character error rate (CER) metrics. -""" +"""Evaluate OCR models on datasets with WER and CER metrics.""" import os import torch @@ -17,10 +12,7 @@ import traceback class OCRModelEvaluator: - """ - A comprehensive OCR model evaluator that supports multiple models and provides - detailed analysis with WER and CER metrics. - """ + """OCR model evaluator over multiple models with WER/CER analysis.""" def __init__(self): """Initialize the OCR evaluator.""" @@ -37,9 +29,7 @@ class OCRModelEvaluator: min_p: float = 0.1, verbose: bool = True, ) -> Tuple[Optional[float], Optional[float]]: - """ - Evaluate a model on an OCR dataset. - """ + """Evaluate a model on an OCR dataset.""" os.makedirs(output_dir, exist_ok = True) results = [] @@ -184,7 +174,7 @@ class OCRModelEvaluator: use_cache = True, ) - # Keep only the generated part, not the input + # Keep only the generated tokens, not the input generated_ids_trimmed = [ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) ] diff --git a/tests/utils/os_utils.py b/tests/utils/os_utils.py index 2dcd9c732e..0dc3d2fbf5 100644 --- a/tests/utils/os_utils.py +++ b/tests/utils/os_utils.py @@ -33,12 +33,10 @@ def check_package_installed(package_name, package_manager = None): try: if package_manager == "apt": - # Check with dpkg result = subprocess.run(["dpkg", "-l", package_name], capture_output = True, text = True) return result.returncode == 0 elif package_manager in ["yum", "dnf"]: - # Check with rpm result = subprocess.run(["rpm", "-q", package_name], capture_output = True, text = True) return result.returncode == 0 diff --git a/tests/utils/perplexity_eval.py b/tests/utils/perplexity_eval.py index 83c632a0a7..5f33a24d53 100644 --- a/tests/utils/perplexity_eval.py +++ b/tests/utils/perplexity_eval.py @@ -37,7 +37,7 @@ def ppl_model(model, tokenizer, dataset): def add_to_comparison(model_name, ppl): - """Add model results to the comparison tracker""" + """Record a model's perplexity in the comparison tracker.""" model_comparison_results[model_name] = {"ppl": ppl} diff --git a/tests/utils/run_none_detect_tests.py b/tests/utils/run_none_detect_tests.py index 2afea75bab..375f2a9c22 100644 --- a/tests/utils/run_none_detect_tests.py +++ b/tests/utils/run_none_detect_tests.py @@ -1,16 +1,4 @@ -""" -run_none_detect_tests.py - -Runs dataset_none_detect.py against: - 1. Synthetic datasets (chatml, sharegpt, alpaca) with deliberately injected Nones - 2. peteromallet/dataclaw-peteromallet (HuggingFace) - 3. peteromallet/my-personal-codex-data (HuggingFace) - -Writes a detailed log to tests/logs/none_detect_results.log - -Usage (from repo root, with venv active): - python tests/utils/run_none_detect_tests.py -""" +"""Run dataset_none_detect.py against synthetic + two HF datasets; log to tests/logs/none_detect_results.log.""" from __future__ import annotations @@ -21,8 +9,7 @@ from datetime import datetime from io import StringIO from pathlib import Path -# Import dataset_none_detect as a top-level module, bypassing -# utils/datasets/__init__.py which pulls in torch, fastapi, and other heavy deps. +# Import dataset_none_detect directly, bypassing utils/datasets/__init__.py (heavy deps). REPO_ROOT = Path(__file__).resolve().parent.parent.parent sys.path.insert(0, str(REPO_ROOT / "studio" / "backend" / "utils" / "datasets")) @@ -86,7 +73,7 @@ def assert_bad_rows(stats: dict, expected_min: int, label: str): def assert_exact_recall(stats: dict, expected_bad: set, label: str): - """Every injected bad row index must appear in bad_row_indices — no misses.""" + """Every injected bad row index must appear in bad_row_indices.""" actual_bad = set(stats.get("bad_row_indices", [])) missed = expected_bad - actual_bad all_caught = len(missed) == 0 @@ -106,8 +93,7 @@ def assert_exact_recall(stats: dict, expected_bad: set, label: str): # 1. Synthetic datasets -# Minimal mock so find_none_chatml can be called with hand-crafted rows that -# pyarrow can't represent (e.g. messages=None, messages="not a list"). +# Minimal mock for hand-crafted rows pyarrow can't represent (e.g. messages=None / "not a list"). class _MockDataset: @@ -124,12 +110,12 @@ class _MockDataset: return iter(self._rows) def __getitem__(self, idx): - """Support dataset[i] (int) and dataset[i][col] patterns used by _probe_conversation.""" + """Support dataset[i] and dataset[i][col] patterns used by _probe_conversation.""" return self._rows[idx] def test_p1_fix(): - """Verify that find_none_chatml records rows where messages is None or non-list.""" + """find_none_chatml records rows where messages is None or non-list.""" section("P1 Fix Verification — non-list conversation column values") sys.path.insert(0, str(REPO_ROOT / "tests" / "utils")) @@ -158,12 +144,10 @@ def test_p1_fix(): def test_probe_p1_fix(): - """Verify that scan_dataset(fmt='auto') on a fully-corrupt messages column - routes through _probe_conversation's all_corrupt=True path and returns - findings rather than raising ValueError.""" + """scan_dataset(fmt='auto') on an all-corrupt messages column returns findings, not ValueError.""" section("P1 Fix Verification — probe skip on all-corrupt column (auto-detect path)") - # All 5 rows have messages=None — no dict turn will ever be found in the probe. + # All rows have messages=None, so the probe finds no dict turn. all_corrupt_rows = [{"messages": None}] * 5 mock_ds = _MockDataset(all_corrupt_rows, ["messages"]) @@ -183,10 +167,7 @@ def test_probe_p1_fix(): def test_probe_string_corrupt(): - """Verify that a plain-string 'messages' column is NOT classified as chatml. - The all_corrupt+has_list_or_none guard (P2 fix) must prevent plain-string - columns from triggering the chatml path — scan_dataset should raise - ValueError('unknown format') rather than misclassifying as chatml.""" + """P2 fix: a plain-string 'messages' column must NOT be classified as chatml (raises ValueError).""" section("P2 Fix Verification — plain-string messages not classified as chatml") string_rows = [{"messages": "this is a string, not a list"}] * 5 @@ -196,7 +177,6 @@ def test_probe_string_corrupt(): try: stats = scan_dataset(mock_ds, fmt = "auto") fmt = stats.get("format", "?") - # Plain strings must NOT be classified as chatml. not_chatml = fmt != "chatml" status = "PASS" if not_chatml else "FAIL" print( @@ -205,8 +185,7 @@ def test_probe_string_corrupt(): ) return stats except ValueError as exc: - # ValueError (unknown format) is the CORRECT outcome — plain strings - # are not a conversation column and should not match any format. + # ValueError (unknown format) is the correct outcome for a non-conversation column. print( f" [PASS] String-corrupt probe: scan_dataset raised ValueError (not chatml, as expected): {exc}" ) @@ -214,8 +193,7 @@ def test_probe_string_corrupt(): def test_explicit_fmt_corrupt(): - """Verify that scan_dataset(fmt='chatml') on an all-corrupt column returns - findings instead of raising ValueError (explicit format + all_corrupt path).""" + """scan_dataset(fmt='chatml') on an all-corrupt column returns findings, not ValueError.""" section("P1 Fix Verification — explicit fmt='chatml' on all-corrupt column") all_corrupt_rows = [{"messages": None}] * 4 + [{"messages": "not a list"}] * 3 @@ -237,8 +215,7 @@ def test_explicit_fmt_corrupt(): def test_p2_probe_skips_corrupt_prefers_valid(): - """P2 fix: probe must continue past a corrupt 'messages' column and find - a valid 'conversations' column rather than returning all_corrupt immediately.""" + """P2 fix: probe continues past a corrupt 'messages' column to a valid 'conversations' column.""" section("P2 Fix Verification — probe continues past corrupt first column to valid second") # messages column is all-None; conversations is a valid ShareGPT column. @@ -267,7 +244,7 @@ def test_p2_probe_skips_corrupt_prefers_valid(): fmt = stats.get("format", "?") col = stats.get("column", "?") bad = len(stats.get("bad_row_indices", [])) - # Should have detected 'conversations' (sharegpt), not 'messages'. + # Must detect 'conversations' (sharegpt), not 'messages'. correct_col = col == "conversations" correct_fmt = fmt == "sharegpt" correct_bad = bad == 2 @@ -284,8 +261,7 @@ def test_p2_probe_skips_corrupt_prefers_valid(): def test_p2_explicit_fmt_col_priority(): - """P2 fix: explicit fmt='sharegpt' must let find_none_sharegpt choose its own - column (conversations) rather than being forced onto the probe's column (messages).""" + """P2 fix: explicit fmt='sharegpt' lets find_none_sharegpt pick its own column (conversations).""" section("P2 Fix Verification — explicit fmt='sharegpt' respects per-scanner column priority") # messages has valid role/content turns (chatml-ish); conversations has bad sharegpt turns. @@ -307,7 +283,7 @@ def test_p2_explicit_fmt_col_priority(): stats = scan_dataset(mock_ds, fmt = "sharegpt") col = stats.get("column", "?") bad = len(stats.get("bad_row_indices", [])) - # fmt='sharegpt' should scan 'conversations', find 5 bad rows. + # fmt='sharegpt' scans 'conversations' -> 5 bad rows. correct_col = col == "conversations" correct_bad = bad == 5 status = "PASS" if (correct_col and correct_bad) else "FAIL" @@ -320,8 +296,7 @@ def test_p2_explicit_fmt_col_priority(): def test_p2_gptoss_col_priority(): - """P2 fix: fmt='gptoss' must scan 'messages' only, not fall through to - a clean 'conversations' column when messages is corrupt.""" + """P2 fix: fmt='gptoss' scans 'messages' only, not a clean 'conversations' fallback.""" section("P2 Fix Verification — fmt='gptoss' scans messages only, not conversations") # messages is all-None (corrupt); conversations is clean sharegpt. @@ -356,10 +331,7 @@ def test_p2_gptoss_col_priority(): def test_new_p1_explicit_sharegpt_both_all_corrupt(): - """NEW P1 (commit eb7fea3b7e): when fmt='sharegpt' and both 'messages' and - 'conversations' are all-corrupt in the first 100 rows, scan_dataset must scan - 'conversations' (the sharegpt column), NOT 'messages' (the generic probe's - first candidate).""" + """NEW P1 (commit eb7fea3b7e): fmt='sharegpt' with both columns all-corrupt scans 'conversations', not 'messages'.""" section( "NEW P1 — explicit fmt='sharegpt' scans 'conversations' even when both columns all-corrupt" ) @@ -373,7 +345,7 @@ def test_new_p1_explicit_sharegpt_both_all_corrupt(): stats = scan_dataset(mock_ds, fmt = "sharegpt") col = stats.get("column", "?") bad = len(stats.get("bad_row_indices", [])) - # MUST scan 'conversations', not 'messages'. + # Must scan 'conversations', not 'messages'. correct_col = col == "conversations" correct_bad = bad == 5 status = "PASS" if (correct_col and correct_bad) else "FAIL" @@ -389,12 +361,10 @@ def test_new_p1_explicit_sharegpt_both_all_corrupt(): def test_new_p2_plain_string_messages_not_chatml(): - """NEW P2 (commit eb7fea3b7e): a dataset where 'messages' contains plain - strings (not lists) must NOT be auto-classified as chatml. The all_corrupt - path in FORMAT_REGISTRY chatml match previously triggered here too.""" + """NEW P2 (commit eb7fea3b7e): plain-string 'messages' must NOT be auto-classified as chatml.""" section("NEW P2 — plain-string 'messages' column must NOT be classified as chatml") - # messages is a plain text column, not a conversation column at all. + # messages is a plain text column, not a conversation column. rows = [{"messages": "hello world"}] * 5 mock_ds = _MockDataset(rows, ["messages"]) @@ -402,7 +372,6 @@ def test_new_p2_plain_string_messages_not_chatml(): try: stats = scan_dataset(mock_ds, fmt = "auto") fmt = stats.get("format", "?") - # Classifying as chatml is wrong; 'unknown' or another non-chatml result expected. not_chatml = fmt != "chatml" status = "PASS" if not_chatml else "FAIL" print( @@ -411,7 +380,7 @@ def test_new_p2_plain_string_messages_not_chatml(): ) return stats except ValueError as exc: - # ValueError is also acceptable — the column isn't a valid conversation format. + # ValueError is also acceptable: not a valid conversation format. print( f" [PASS] New-P2 plain-string messages: scan_dataset raised ValueError (not chatml): {exc}" ) @@ -463,13 +432,9 @@ def test_synthetic(): def _brute_force_bad_rows(ds, fmt: str) -> set: - """Pure-Python ground-truth scanner — zero shared code with dataset_none_detect. + """Pure-Python ground-truth scanner (no shared code with dataset_none_detect) for independent proof. - Iterates every row and flags it bad if any field/turn contains None, empty, - or whitespace-only content. Returns the set of bad row indices. - - Intentionally re-implements the same logic from scratch so that agreement - between this and scan_dataset() constitutes independent proof of correctness. + Flags a row bad if any field/turn is None, empty, or whitespace-only; returns bad row indices. """ def _blank(val) -> bool: @@ -506,20 +471,15 @@ def _brute_force_bad_rows(ds, fmt: str) -> set: def _assert_hf_no_misses(ds, stats: dict, label: str) -> bool: - """Independent ground-truth check: verify scan_dataset() found every bad row - that brute-force row-by-row iteration also finds. - - brute_force_bad ⊆ module_bad → no misses (the key assertion) - module_bad ⊆ brute_force_bad → no false positives (informational) - """ + """Independent check: scan_dataset() must find every bad row brute-force finds (no misses).""" fmt = stats.get("format", "unknown") module_bad = set(stats.get("bad_row_indices", [])) print(f" Running brute-force independent scan (fmt={fmt!r}, {len(ds)} rows)...") brute_bad = _brute_force_bad_rows(ds, fmt) - missed = brute_bad - module_bad # rows brute-force found but module missed - extra = module_bad - brute_bad # rows module flagged that brute-force didn't + missed = brute_bad - module_bad # brute-force found, module missed + extra = module_bad - brute_bad # module flagged, brute-force didn't no_misses = len(missed) == 0 snippet = "" @@ -534,8 +494,7 @@ def _assert_hf_no_misses(ds, stats: dict, label: str) -> bool: f"missed: {len(missed)}{snippet}" ) if extra: - # Module may legitimately flag more rows than brute-force (e.g. extra - # structural checks) — informational only, not a failure. + # Module may legitimately flag more rows (extra structural checks); informational only. print( f" [INFO] {label} — module flagged {len(extra)} rows not in brute-force " f"(may reflect additional structural checks, not false positives)" @@ -571,8 +530,7 @@ def test_dataclaw(): def test_codex_data(): section("3. HuggingFace — peteromallet/my-personal-codex-data") try: - # load_dataset fails on this repo because ujson chokes on the large - # JSONL batch. Download the raw file and parse it ourselves instead. + # load_dataset fails here (ujson chokes on the large JSONL batch); download + parse raw instead. from huggingface_hub import hf_hub_download from datasets import Dataset diff --git a/tests/utils/test_batched_leftpad_generation_gpu.py b/tests/utils/test_batched_leftpad_generation_gpu.py index c2610d4673..df03125bc2 100644 --- a/tests/utils/test_batched_leftpad_generation_gpu.py +++ b/tests/utils/test_batched_leftpad_generation_gpu.py @@ -1,17 +1,11 @@ """End-to-end GPU guard for batched left-padded generation (issues #1066, #3699). -For each prompt, greedy generation inside a left-padded batch must match -generating that prompt alone at batch size 1 for the first PREFIX_TOKENS -tokens, and the full output must not be gibberish. The bug class (#1066, -#3699) makes padded rows diverge immediately into garbage; in contrast, -benign batch-size-dependent kernel numerics can flip a greedy near-tie deep -into the sequence, so an exact full-length match would be flaky. Uses a small -instruct model (chat-templated prompts have high-margin argmaxes). - -Skipped automatically when CUDA is unavailable, so CPU CI is unaffected. -Run manually on any GPU box: - - python -m pytest tests/utils/test_batched_leftpad_generation_gpu.py -v +Greedy generation in a left-padded batch must match solo batch-size-1 +generation for the first PREFIX_TOKENS tokens (the bug makes padded rows +diverge into garbage immediately; a full-length match would be flaky due to +benign batch-numerics tie-flips deep in the sequence) and must not be +gibberish. Skipped without CUDA. Run: `python -m pytest +tests/utils/test_batched_leftpad_generation_gpu.py -v`. """ import pytest diff --git a/tests/utils/test_packing.py b/tests/utils/test_packing.py index 6fa6ae8d34..a8557d8533 100644 --- a/tests/utils/test_packing.py +++ b/tests/utils/test_packing.py @@ -194,7 +194,7 @@ class _DummyTrainer: break except TypeError: continue - # Ensure attributes exist even if the constructor rejected them. + # Ensure attributes exist even if the constructor rejected the flags. if not hasattr(self.data_collator, "padding_free"): self.data_collator.padding_free = True if not hasattr(self.data_collator, "return_position_ids"): @@ -221,7 +221,7 @@ def test_enable_sample_packing(): enable_sample_packing(model, trainer) - # model hierarchy should now allow packed overlength inputs + # model hierarchy now allows packed overlength inputs assert getattr(model, "_unsloth_allow_packed_overlength") is True assert getattr(model.child, "_unsloth_allow_packed_overlength") is True @@ -243,7 +243,7 @@ def test_enable_sample_packing(): ] batch = collator.torch_call(examples) - # packed lengths are aggregated into a single tensor + # packed lengths aggregated into one tensor assert "packed_seq_lengths" in batch assert torch.equal(batch["packed_seq_lengths"], torch.tensor([2, 1, 3], dtype = torch.int32)) diff --git a/tests/utils/test_prepare_inputs_leftpad.py b/tests/utils/test_prepare_inputs_leftpad.py index b773b088b5..2bfd763279 100644 --- a/tests/utils/test_prepare_inputs_leftpad.py +++ b/tests/utils/test_prepare_inputs_leftpad.py @@ -1,29 +1,16 @@ """Regression guard for batched left-padded generation (issues #1066, #3699). -`_fast_prepare_inputs_for_generation` in unsloth/models/llama.py is shared by -every decoder family wired through `fix_prepare_inputs_for_generation` (llama, -qwen2/3, qwen3_moe, mistral, gemma/2, cohere, granite). Two historical bugs -lived in it: +Guards `_fast_prepare_inputs_for_generation` (unsloth/models/llama.py), +shared by every decoder family wired through fix_prepare_inputs_for_generation, +against two historical bugs: + (a) 2D attention mask truncated to its last column during cached decode, + losing padding info (fixed by #2216); + (b) position_ids taken from cache_position (which counts left-pad tokens), + so padded rows generated garbage (fixed by #4100). - (a) the 2D attention mask was truncated to its last column during cached - decode, losing padding information (introduced cc4c5d77, fixed by #2216); - (b) position_ids were taken directly from cache_position, a global counter - that includes left-pad tokens, so padded rows generated garbage - (introduced cc4c5d77, reported in #1066/#3699, fixed by #4100). - -Two layers in this file, both CPU-only and deterministic: - - 1. AST structural checks (TestAstGuard section): parse llama.py with the - stdlib `ast` module only, no unsloth import, so they keep working even - when the package itself fails to import. - 2. Behavioral checks: call the real function with synthetic left-padded - masks and fake caches; unsloth is imported lazily inside each test. - -Both layers fail on the historical bug patterns (validated against the code -states immediately before #2216 and before #4100). Staging proof on GPU-less -hosted runners: danielhanchen/unsloth-staging-2 PRs 170 (green) and 171 (red). - -Companion manual GPU end-to-end check (skipped without CUDA): +Two CPU-only deterministic layers: (1) AST structural checks (no unsloth +import); (2) behavioral checks calling the real function with synthetic +left-padded masks and fake caches. Companion GPU check: tests/utils/test_batched_leftpad_generation_gpu.py """ @@ -43,10 +30,9 @@ FUNC_NAME = "_fast_prepare_inputs_for_generation" # Layer 1: AST structural guard (stdlib only, no unsloth import) # -------------------------------------------------------------------------- -# Model files that call fix_prepare_inputs_for_generation(...) and therefore -# share the guarded function. glm4_moe and falcon_h1 are intentionally absent: -# GLM4 MoE does not patch the Llama-compatible generation path (MLA attention) -# and falcon_h1 ships its own _fast_prepare_inputs_for_generation variant. +# Model files that call fix_prepare_inputs_for_generation(...) and share the +# guarded function. glm4_moe (MLA attention, different path) and falcon_h1 +# (its own variant) are intentionally absent. WIRED_MODEL_FILES = [ "mistral.py", "gemma.py", diff --git a/tests/utils/test_q_galore.py b/tests/utils/test_q_galore.py index 7248e60a93..a4714b0bda 100644 --- a/tests/utils/test_q_galore.py +++ b/tests/utils/test_q_galore.py @@ -20,14 +20,12 @@ import os import torch import torch.nn as nn -# Import the optimizers module directly to avoid triggering unsloth.__init__ -# which requires unsloth_zoo and other heavy dependencies. +# Import optimizers directly to avoid triggering unsloth.__init__ (heavy deps). _repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) _optimizers_dir = os.path.join(_repo_root, "unsloth", "optimizers") if _repo_root not in sys.path: sys.path.insert(0, _repo_root) -# Direct import of the actual modules (avoids unsloth/__init__.py) import importlib.util @@ -39,7 +37,7 @@ def _load_module(name, filepath): return mod -# Load projector module first (no dependencies on unsloth) +# Projector has no unsloth dependencies; load it first. _projector_mod = _load_module( "unsloth.optimizers.q_galore_projector", os.path.join(_optimizers_dir, "q_galore_projector.py"), @@ -49,7 +47,7 @@ _quantize = _projector_mod._quantize _dequantize = _projector_mod._dequantize _quantize_stochastic = _projector_mod._quantize_stochastic -# Load adamw module (depends on projector, may skip bitsandbytes) +# adamw depends on projector, may skip bitsandbytes. _adamw_mod = _load_module( "unsloth.optimizers.q_galore_adamw", os.path.join(_optimizers_dir, "q_galore_adamw.py"), @@ -124,13 +122,12 @@ class TestGaLoreProjector: gamma_proj = 2.0, queue_size = 2, ) - # Use very similar gradients so cosine similarity is high + # Near-identical gradients keep cosine similarity high. base_grad = torch.randn(16, 8) for i in range(5): grad = base_grad + torch.randn_like(base_grad) * 0.001 proj.project(grad, step = i * 10) - # After several similar SVDs, update_proj_gap should have increased assert proj.update_proj_gap > 10 def test_scale_applied(self): @@ -145,7 +142,7 @@ class TestGaLoreProjector: full_half = proj.project_back(low) full_one = proj2.project_back(low2) - # The ratio should be exactly 0.5 (SVD is deterministic on same input) + # SVD is deterministic on the same input, so the ratio is exactly 0.5. ratio = full_half.norm() / full_one.norm() assert abs(ratio - 0.5) < 1e-5, f"Expected ratio ~0.5, got {ratio:.8f}" @@ -252,11 +249,7 @@ class TestParamGroupHelper: assert len(galore_group["params"]) == 1 # Only q_proj def test_bias_excluded_from_galore(self): - """1D bias params matching target names must NOT be in the GaLore group. - - GaLoreProjector.project requires 2-D gradients, so bias vectors - (e.g. q_proj.bias) that match a target name must be excluded. - """ + """1-D bias params matching target names must be excluded (project needs 2-D grads).""" model = nn.Module() model.q_proj = nn.Linear(64, 64, bias = True) # has .weight AND .bias model.embed = nn.Embedding(100, 64) @@ -302,7 +295,6 @@ class TestQGaLoreIntegration: """A simple training loop using manual GaLore projection converges.""" torch.manual_seed(42) - # Tiny model: single linear layer model = nn.Linear(32, 16, bias = False) target = torch.randn(4, 16) x = torch.randn(4, 32) @@ -318,7 +310,6 @@ class TestQGaLoreIntegration: loss.backward() losses.append(loss.item()) - # Manual GaLore projection for p in model.parameters(): if p.grad is not None and p.grad.dim() == 2: low = proj.project(p.grad, step) @@ -330,13 +321,11 @@ class TestQGaLoreIntegration: optimizer.step() - # Loss should decrease assert losses[-1] < losses[0], f"Loss did not decrease: {losses[0]:.4f} → {losses[-1]:.4f}" def test_full_projector_roundtrip_quality(self): """project → project_back captures the dominant gradient directions.""" torch.manual_seed(42) - # Create a gradient with clear low-rank structure u = torch.randn(32, 4) v = torch.randn(4, 16) grad = u @ v # rank-4 gradient @@ -345,8 +334,7 @@ class TestQGaLoreIntegration: low = proj.project(grad, step = 0) reconstructed = proj.project_back(low) - # For a rank-4 gradient with rank-4 projection, reconstruction - # should be very close to original + # Rank-4 grad with rank-4 projection reconstructs near-exactly. relative_error = (grad - reconstructed).norm() / grad.norm() assert relative_error < 0.05, f"Reconstruction error too high: {relative_error:.4f}" @@ -356,7 +344,7 @@ class TestQGaLoreIntegration: QGaLoreAdamW8bit = _adamw_mod_local.QGaLoreAdamW8bit p = torch.nn.Parameter(torch.randn(16, 16)) - # Simulate init_weight_quantization tagging + # Simulate init_weight_quantization tagging. p._q_scales = None p._q_zeros = None p._q_shape = p.data.shape @@ -372,15 +360,14 @@ class TestQGaLoreIntegration: def test_embedding_lr_param_group_split(self): """Embedding params can be split into a separate group with custom LR.""" - # This tests the logic that make_q_galore_param_groups produces groups - # that can be further split by the trainer for embedding LR. + # make_q_galore_param_groups output can be further split for embedding LR. model = nn.Module() model.q_proj = nn.Linear(64, 64, bias = False) model.embed = nn.Embedding(100, 64) groups = make_q_galore_param_groups(model, rank = 8, weight_quant = False) - # Simulate splitting non-GaLore group for embedding LR + # Simulate splitting the non-GaLore group for embedding LR. embed_lr = 5e-5 new_groups = [] for group in groups: @@ -390,7 +377,7 @@ class TestQGaLoreIntegration: embed_params = [] other_params = [] for p in group["params"]: - # In real usage, we'd check the name; here just split by shape + # Real usage checks names; here we split by shape. if p.shape[0] == 100: # embedding embed_params.append(p) else: @@ -405,16 +392,14 @@ class TestQGaLoreIntegration: g["lr"] = embed_lr new_groups.append(g) - # Should have 3 groups: galore, non-galore non-embed, embed + # 3 groups: galore, non-galore non-embed, embed. embed_groups = [g for g in new_groups if g.get("lr") == embed_lr] assert len(embed_groups) == 1 assert embed_groups[0]["lr"] == embed_lr def test_optimizer_hyperparams_forwarded(self): """QGaLoreAdamW8bit accepts betas and eps keyword arguments.""" - # Verify the constructor signature accepts these params. - # Without bitsandbytes we can't instantiate, but we can check the - # function signature. + # Can't instantiate without bitsandbytes; check the signature instead. import inspect _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"] @@ -429,15 +414,14 @@ class TestQGaLoreIntegration: """Weight decay should apply standard decoupled AdamW decay on current weights.""" _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"] - # Create a mock parameter and group p = torch.nn.Parameter(torch.ones(4, 4)) p._saved_data = torch.ones(4, 4) * 2.0 # Pre-update weights - # Simulate project-back: p.data = p._saved_data + projected update + # Simulate project-back: p.data = p._saved_data + projected update. p.data = p._saved_data.add_(torch.ones(4, 4) * 1.0) # p.data is now 3.0 group = {"weight_decay": 0.1, "lr": 1.0, "_wd_saved": 0.1} - # Replicate the fixed decoupled weight decay logic (uses p.data, not p._saved_data) + # Decoupled weight decay must use p.data, not p._saved_data. p.data.add_( p.data, alpha = -group["lr"] * group["_wd_saved"], @@ -445,7 +429,7 @@ class TestQGaLoreIntegration: del p._saved_data # Clean up after all uses, matching fixed code - # Decoupled weight decay: 3.0 - (1.0 * 0.1 * 3.0) = 2.7 + # 3.0 - (1.0 * 0.1 * 3.0) = 2.7 assert torch.allclose( p.data, torch.tensor(2.7) ), "Weight decay didn't use p.data for decoupled decay!" @@ -464,11 +448,11 @@ class TestQGaLoreIntegration: "weight_group_size": 16, } - # Replicate the re-quantize logic at the end of optimizer step + # Re-quantize logic from the end of an optimizer step. float_data = p.data.clone() q, scales, zeros, shape = _quantize(float_data, q_group_size = group["weight_group_size"]) - # The key assertion: p.data stays float, _q_data holds uint8 + # Key check: p.data stays float, _q_data holds uint8. p._q_data = q.to(p.data.device) p._q_scales = scales p._q_zeros = zeros @@ -486,7 +470,7 @@ class TestQGaLoreIntegration: linear = nn.Linear(16, 8, bias = False) original = linear.weight.data.clone() - # Quantize the weight and replace with placeholder (simulates post-step) + # Quantize the weight and replace with a placeholder (simulates post-step). q, scales, zeros, shape = _projector_mod_local._quantize( linear.weight.data.clone(), q_group_size = 16 ) @@ -497,14 +481,14 @@ class TestQGaLoreIntegration: linear.weight.data = torch.zeros(1, dtype = linear.weight.dtype) assert linear.weight.data.numel() == 1, "placeholder should be 1 element" - # Install hook and run forward -- should restore float weights + # Hook should restore float weights on forward. handles = install_hook(linear) x = torch.randn(2, 16) out = linear(x) # triggers pre-hook assert linear.weight.data.shape == (8, 16), "weight shape not restored" assert linear.weight.data.is_floating_point(), "weight not float after hook" - # Check values are close to original (quantization introduces small error) + # Quantization introduces small error, so allow tolerance. assert torch.allclose( linear.weight.data, original, atol = 0.15 ), "dequantized weight too far from original" diff --git a/tests/utils/test_qat.py b/tests/utils/test_qat.py index 5c87db4e3d..79d955164f 100644 --- a/tests/utils/test_qat.py +++ b/tests/utils/test_qat.py @@ -20,9 +20,7 @@ except ImportError: class _CountingFakeQuantizer(torch.nn.Module): - """ - Dummy fake quantizer that counts the number of times it has been called. - """ + """Fake quantizer that counts how many times it was called.""" def __init__(self): super().__init__() @@ -34,10 +32,7 @@ class _CountingFakeQuantizer(torch.nn.Module): def _get_model(qat_scheme: str, full_finetuning: bool): - """ - Return a 2-tuple of (model, tokenizer), where the model has been configured - to use QAT. If `full_finetuning` is False, return the PEFT (LoRA) model. - """ + """Return (model, tokenizer) configured for QAT; LoRA model when full_finetuning is False.""" model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Qwen3-1.7B", load_in_4bit = False, @@ -53,9 +48,7 @@ def _get_model(qat_scheme: str, full_finetuning: bool): def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str): - """ - Verify that the given linear contains fake quantizers according to the `qat_scheme`. - """ + """Verify the linear contains fake quantizers matching `qat_scheme`.""" weight_only = False if qat_scheme == "fp8-int4": act_fq_class = Float8FakeQuantizer @@ -78,7 +71,7 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str): else: raise ValueError(f"Unknown qat_scheme: {qat_scheme}") - # Check base layer activations and weights + # Check base layer activations and weights. base_layer = getattr(linear, "base_layer", linear) if base_layer.in_features >= min_in_features: assert isinstance(base_layer, FakeQuantizedLinear) @@ -86,7 +79,7 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str): assert isinstance(base_layer.activation_fake_quantizer, act_fq_class) assert isinstance(base_layer.weight_fake_quantizer, weight_fq_class) - # Check lora A and B (only for full_finetuning=False) + # Check lora A and B (full_finetuning=False only). if hasattr(linear, "lora_A") and hasattr(linear, "lora_B"): lora_A = linear.lora_A.default lora_B = linear.lora_B.default @@ -105,9 +98,7 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str): def _test_fake_quantizers_are_called( model: torch.nn.Module, example_inputs: Dict, full_finetuning: bool, qat_scheme: str ): - """ - Verify that the fake quantizers are actually called when the model is called. - """ + """Verify the fake quantizers are actually called during a forward pass.""" weight_only = qat_scheme in ["int8", "cactus"] def _swap_fake_quantizers(model: torch.nn.Module): @@ -123,9 +114,8 @@ def _test_fake_quantizers_are_called( assert child.activation_fake_quantizer.count == 1 assert child.weight_fake_quantizer.count == 1 else: - # For LoRA, we only fake quantize the input activations once per block: - # For self_attn, we only fake quantize the q_proj's input activations - # For mlp, we only fake quantize the gate_proj's input activations + # LoRA fake-quantizes input activations once per block: + # self_attn via q_proj, mlp via gate_proj. if name == "self_attn": base_layer = child.q_proj.base_layer if not weight_only: @@ -137,7 +127,7 @@ def _test_fake_quantizers_are_called( assert hasattr(base_layer, "activation_fake_quantizer") assert base_layer.activation_fake_quantizer.count == 1 elif isinstance(child, FakeQuantizedLinear): - # Weight fake quantizers should always be called + # Weight fake quantizers must always be called. assert child.weight_fake_quantizer.count == 1 for k, v in example_inputs.items(): @@ -148,9 +138,7 @@ def _test_fake_quantizers_are_called( def _test_model_fake_quantize(qat_scheme: str, full_finetuning: bool): - """ - Test that all linear layers in the model are fake quantized according to the `qat_scheme`. - """ + """All linear layers in the model are fake quantized per `qat_scheme`.""" model, tokenizer = _get_model(qat_scheme, full_finetuning) if full_finetuning: model = model.model diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 193cb830f5..b976654f87 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -1,14 +1,10 @@ -"""Guard for config.rope_scaling being silently dropped (issue #2405). +"""Guard against config.rope_scaling being silently dropped (issue #2405): +the replacement rotary classes ignored it on the config path, so Llama-3.1 +ran with unscaled RoPE and produced gibberish past ~32K tokens. -Unsloth's replacement rotary classes ignored rope_scaling when constructed -from a config (the modern-transformers path), so Llama-3.1 ran with unscaled -RoPE and collapsed into gibberish past ~32K tokens. - -Layers: (1) AST tripwire, stdlib only; (2) CPU checks of the pure helper -_compute_config_rope_inv_freq against transformers' ROPE_INIT_FUNCTIONS; -(3) CUDA checks instantiating the real class (skipped without a real device, -probed by allocating a tensor so import-time CUDA spoofs cannot fool the gate). -Layers 2 and 3 fail on the unfixed code. +Three layers: (1) AST tripwire; (2) CPU checks of the pure helper +_compute_config_rope_inv_freq vs ROPE_INIT_FUNCTIONS; (3) CUDA checks on the +real class (skipped without a real device). Layers 2-3 fail on the unfixed code. """ import ast @@ -165,7 +161,7 @@ def test_llama3_scaling_applied_to_inv_freq(): expected = _reference_inv_freq(config, "llama3") vanilla = _vanilla_inv_freq() - # Guard against a vacuous test. + # Guard against a vacuous test: scaled inv_freq must differ from vanilla. assert not torch.allclose( expected, vanilla, rtol = 1e-4 ), "test setup error: llama3-scaled inv_freq should differ from vanilla" diff --git a/tests/utils/test_trunc_normal_patch.py b/tests/utils/test_trunc_normal_patch.py index 0ccd455674..c1dbc2a7c8 100644 --- a/tests/utils/test_trunc_normal_patch.py +++ b/tests/utils/test_trunc_normal_patch.py @@ -57,7 +57,7 @@ def test_trunc_normal_patch_accepts_positional_generator(): old_patched = _getattr_or_missing(init_mod, "_unsloth_trunc_normal_patched") old_original = _getattr_or_missing(init_mod, "_unsloth_trunc_normal_original") try: - # Normalize to an unpatched baseline before applying the patch. + # Reset to an unpatched baseline before applying the patch. if old_original is not _MISSING: init_mod.trunc_normal_ = old_original if hasattr(init_mod, "_unsloth_trunc_normal_patched"): diff --git a/tests/version_compat/_fetch.py b/tests/version_compat/_fetch.py index 9688780be5..706b26efad 100644 --- a/tests/version_compat/_fetch.py +++ b/tests/version_compat/_fetch.py @@ -1,12 +1,6 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -"""Shared helpers for the version-compat suites: fetch a file from -GitHub raw at a tag/branch, and grep for class/def/module symbols -without ast.parse so one non-importable line doesn't false-fail us. -Mirrors tests/vllm_compat/test_vllm_pinned_symbols.py. - -Used by the test_*_pinned_symbols.py suites under tests/version_compat/. -""" +"""Shared helpers for version-compat suites: GitHub raw fetch + regex symbol grep.""" from __future__ import annotations @@ -19,8 +13,7 @@ import pytest def fetch_text(repo: str, ref: str, path: str) -> str | None: - """Fetch a file from GitHub raw. None on 404 (caller decides if - fatal). Skips the test on transient network errors to avoid CI flake.""" + """Fetch a file from GitHub raw. None on 404; skips on transient network errors.""" url = f"https://raw.githubusercontent.com/{repo}/{ref}/{path}" req = urllib.request.Request(url) token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") @@ -42,9 +35,7 @@ def has_def( name: str, kind: str = "any", ) -> bool: - """Heuristic grep for `class Name`, `def name`, or `Name = ...` at - any indent level. Avoids ast.parse so one non-importable line doesn't - false-fail us; indented matches are accepted so class methods count too.""" + """Grep for `class Name`, `def name`, or `Name = ...` at any indent (no ast.parse).""" if kind in ("any", "class") and re.search( rf"^\s*class\s+{re.escape(name)}\b", src, re.MULTILINE ): @@ -59,8 +50,7 @@ def has_def( def first_match(repo: str, ref: str, paths: list[str]) -> tuple[str, str] | None: - """Return (path, src) for the first existing candidate path, else - None. Useful when upstream moved a module across versions.""" + """Return (path, src) for the first existing candidate path, else None.""" for p in paths: src = fetch_text(repo, ref, p) if src is not None: diff --git a/tests/version_compat/test_bitsandbytes_pinned_symbols.py b/tests/version_compat/test_bitsandbytes_pinned_symbols.py index 247cc7e6b9..ba4d05c342 100644 --- a/tests/version_compat/test_bitsandbytes_pinned_symbols.py +++ b/tests/version_compat/test_bitsandbytes_pinned_symbols.py @@ -1,21 +1,6 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -"""Pinned-symbol compat check across bitsandbytes PyPI minor versions -unsloth + unsloth-zoo target. Catches API drift like: - - - bnb 0.46.0 release was broken (in pyproject.toml as `!=0.46.0`). - Don't test against it. - - bnb 0.48.0 release was broken (also `!=0.48.0`). Same. - - bnb 0.45 series introduced fp4 + nf4 paged optimisers; unsloth-zoo - expects bnb.functional.dequantize_4bit + bnb.nn.Linear4bit / - Params4bit to remain stable from this point onward. - - vLLM bitsandbytes-loader patches in unsloth_zoo/vllm_utils.py: - apply_bnb_4bit (line 237), is_layer_skipped_bnb (line 281), - BitsAndBytesLinearMethod._apply_4bit_weight (line 282) — these - live in vllm.* but they call into bnb's public surface. - -Strategy: GitHub raw fetch + symbol grep. CPU-only, no install. -""" +"""Pinned-symbol compat check across bitsandbytes minor versions via GitHub raw fetch + symbol grep.""" from __future__ import annotations @@ -36,8 +21,7 @@ BNB_TAGS = [ ] -# bnb.functional: dequantize_4bit / quantize_4bit are the public 4-bit surface -# unsloth's compiled kernels and unsloth-zoo's vllm_utils bnb-loader call into. +# bnb.functional dequantize_4bit / quantize_4bit: the public 4-bit surface unsloth kernels call into. @pytest.mark.parametrize("tag", BNB_TAGS) @@ -56,8 +40,7 @@ def test_bnb_functional_4bit(tag: str): ) -# bnb.nn.Linear4bit / Params4bit: the two classes peft and unsloth -# isinstance-check against. Renaming either silently breaks 4-bit LoRA. +# bnb.nn.Linear4bit / Params4bit: peft + unsloth isinstance-check these; renaming breaks 4-bit LoRA. @pytest.mark.parametrize("tag", BNB_TAGS) @@ -84,12 +67,10 @@ def test_bnb_nn_linear4bit_classes(tag: str): ) -# Coverage extension (added 2026-05): every bnb symbol unsloth + unsloth-zoo -# touch, derived from a full grep of both repos. +# Coverage extension (2026-05): every bnb symbol unsloth + unsloth-zoo touch. -# Top-level convenience export. unsloth/kernels/utils.py + unsloth-zoo -# vllm_utils.py call `bnb.matmul_4bit(x, w, bias=, quant_state=)`. +# Top-level export: unsloth/kernels/utils.py + zoo vllm_utils.py call bnb.matmul_4bit(...). @pytest.mark.parametrize("tag", BNB_TAGS) @@ -105,17 +86,7 @@ def test_bnb_matmul_4bit_top_level(tag: str): @pytest.mark.parametrize("tag", BNB_TAGS) def test_bnb_functional_4bit_kernel_path(tag: str): - """unsloth/kernels/utils.py module-top binds the 4-bit dequantize - and gemm primitives via one of two paths: - - LEGACY (bnb <= 0.48.x): `bnb.functional.lib.cdequantize_blockwise_*` - and `bnb.functional.lib.cgemm_4bit_inference_naive_*` — C - symbols listed in functional.py source. - - NEW (bnb >= 0.49.0): `torch.ops.bitsandbytes.dequantize_blockwise` - and `torch.ops.bitsandbytes.dequantize_4bit` Python wrappers; - the C symbols still live in libbitsandbytes_*.so but the - Python source no longer references them by name. - Either path lets unsloth resolve the kernels at runtime — we only - fail if NEITHER signal is present.""" + """bnb.functional must expose either the legacy `lib.c*` kernels or the new `torch.ops.bitsandbytes.*` path.""" candidates = [ "bitsandbytes/functional.py", "bitsandbytes/functional/__init__.py", @@ -156,9 +127,7 @@ def test_bnb_functional_get_ptr(tag: str): @pytest.mark.parametrize("tag", BNB_TAGS) def test_bnb_quantstate_from_dict(tag: str): - """unsloth-zoo monkey-patches `QuantState.from_dict = ...`. Both - the class AND the classmethod must be present for the rebinding - to take effect.""" + """unsloth-zoo rebinds QuantState.from_dict; both class and classmethod must be present.""" candidates = [ "bitsandbytes/functional.py", "bitsandbytes/functional/__init__.py", @@ -175,8 +144,7 @@ def test_bnb_quantstate_from_dict(tag: str): @pytest.mark.parametrize("tag", BNB_TAGS) def test_bnb_nn_modules_fix_4bit_weight_optional(tag: str): - """fix_4bit_weight_quant_state_from_module added in newer bnb; - unsloth uses getattr() with a fallback so older versions are OK.""" + """fix_4bit_weight_quant_state_from_module is optional; unsloth getattr-fallbacks on older bnb.""" src = fetch_text("bitsandbytes-foundation/bitsandbytes", tag, "bitsandbytes/nn/modules.py") if src is None: pytest.skip(f"{tag}: bitsandbytes/nn/modules.py missing") @@ -227,8 +195,7 @@ def test_bnb_utils_pack_unpack(tag: str): @pytest.mark.parametrize("tag", BNB_TAGS) def test_bnb_cextension_rocm_warp_size_optional(tag: str): - """ROCM_WARP_SIZE_64 added with AMD ROCm support; pre-ROCm bnb - builds don't have it. unsloth probes via try/except — informational.""" + """ROCM_WARP_SIZE_64 is optional (pre-ROCm bnb lacks it); unsloth probes via try/except.""" src = fetch_text("bitsandbytes-foundation/bitsandbytes", tag, "bitsandbytes/cextension.py") if src is None: pytest.skip(f"{tag}: cextension.py missing") @@ -238,9 +205,7 @@ def test_bnb_cextension_rocm_warp_size_optional(tag: str): @pytest.mark.parametrize("tag", BNB_TAGS) def test_bnb_autograd_functions_matmul_4bit(tag: str): - """unsloth-zoo has a dynamo-disable patch site for - bnb.autograd._functions.matmul_4bit. Symbol must remain so the - probe + decision logic works.""" + """bnb.autograd._functions.matmul_4bit must remain (unsloth-zoo has a dynamo-disable patch site).""" src = fetch_text( "bitsandbytes-foundation/bitsandbytes", tag, @@ -253,9 +218,7 @@ def test_bnb_autograd_functions_matmul_4bit(tag: str): @pytest.mark.parametrize("tag", BNB_TAGS) def test_bnb_version_parseable(tag: str): - """Multiple unsloth code paths read Version(bnb.__version__) for - feature gating (floors 0.43.3, 0.46.0, 0.48.2.dev0, 0.49.0, - 0.49.2). At least one export mechanism must work.""" + """bnb.__version__ must be exported via at least one mechanism (unsloth feature-gates on it).""" src = fetch_text("bitsandbytes-foundation/bitsandbytes", tag, "bitsandbytes/__init__.py") if src is None: pytest.skip(f"{tag}: bitsandbytes/__init__.py missing") diff --git a/tests/version_compat/test_peft_pinned_symbols.py b/tests/version_compat/test_peft_pinned_symbols.py index 7e29aa62ac..d322ecf678 100644 --- a/tests/version_compat/test_peft_pinned_symbols.py +++ b/tests/version_compat/test_peft_pinned_symbols.py @@ -1,28 +1,9 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -"""Pinned-symbol compat check across PEFT PyPI minor versions -unsloth + unsloth-zoo target. Catches API drift like: - - - peft 0.18 finalised the LoraConfig public surface (+ MoE-aware - target_modules); unsloth uses target_modules + r + lora_alpha + - lora_dropout + bias. - - peft 0.19 introduced the LoraConfig.target_parameters extension; - unsloth-zoo's MoE LoRA extractor in saving_utils.py reads it via - getattr() so missing on older versions is OK but the attribute - shape must remain stable on >= 0.19. - - peft.tuners.lora package layout: LoraLayer / LoraConfig / Linear4bit - re-exports must keep working under both `from peft import X` and - `from peft.tuners.lora import X`. - -Strategy: for each tracked PEFT tag, fetch source from -github.com/huggingface/peft (no pip install needed) and assert that -every symbol unsloth + unsloth-zoo's PEFT touchpoints depend on is -present. - -Versioning policy: cover the supported window declared in -unsloth/pyproject.toml (`peft>=0.18.0,!=0.11.0`) plus `main`. The -`!=0.11.0` exclusion is for the historical broken release; we don't -test against it. +"""Pinned-symbol compat check across PEFT minor versions unsloth + unsloth-zoo +target. For each tracked tag, fetch source from github.com/huggingface/peft and +assert every PEFT symbol unsloth touches is present, catching API drift. +Versioning covers unsloth/pyproject.toml's `peft>=0.18.0,!=0.11.0` window + main. """ from __future__ import annotations @@ -45,9 +26,8 @@ PEFT_TAGS = [ ] -# Top-level public re-exports. unsloth/models/sentence_transformer.py:1948 -# does `from peft import LoraConfig, get_peft_model as peft_get_peft_model`. -# unsloth_zoo's saving_utils + lora extractors hit `peft.PeftModel`. +# Top-level re-exports: sentence_transformer.py:1948 does `from peft import +# LoraConfig, get_peft_model`; unsloth_zoo saving_utils/lora extractors hit PeftModel. @pytest.mark.parametrize("tag", PEFT_TAGS) @@ -92,8 +72,7 @@ def test_peft_lora_config_class(tag: str): @pytest.mark.parametrize("tag", PEFT_TAGS) def test_get_peft_model_function(tag: str): - """`def get_peft_model(...)` may live in mapping.py (older - layout) or mapping_func.py (peft 0.18+ split). Either is fine.""" + """get_peft_model may live in mapping.py or mapping_func.py (0.18+ split).""" candidates = [ "src/peft/mapping.py", "src/peft/mapping_func.py", @@ -143,8 +122,7 @@ def test_peft_lora_bnb_integration(tag: str): src = fetch_text("huggingface/peft", tag, p) if src is None: continue - # The Linear4bit subclass naming is the contract -- either name - # is fine, but at least one bnb-flavoured Linear must exist. + # At least one bnb-flavoured Linear must exist (either name is fine). has_4bit = any( cls in src for cls in ( diff --git a/tests/version_compat/test_sentence_transformers_pinned_symbols.py b/tests/version_compat/test_sentence_transformers_pinned_symbols.py index 87b358724b..c0c35b9d5d 100644 --- a/tests/version_compat/test_sentence_transformers_pinned_symbols.py +++ b/tests/version_compat/test_sentence_transformers_pinned_symbols.py @@ -1,24 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -"""Pinned-symbol compat check across sentence-transformers PyPI minor -versions. unsloth has a custom integration in -unsloth/models/sentence_transformer.py that: - - - Imports SentenceTransformer / SentenceTransformerTrainer at the - top of the public surface (lines 1467, 1798, 1947, 2154). - - Walks `sentence_transformers.models` for Transformer / Pooling / - Normalize (lines 1016, 1206, 1467). - - Calls `sentence_transformers.util.import_from_string` and - `load_dir_path` (lines 1177, 1205). - - Tolerates two alternate base-class paths - (sentence_transformers.base.modules.transformer.Transformer vs - sentence_transformers.models.transformer.Transformer; lines - 1169-1171) — at least ONE must resolve. - -Strategy: GitHub raw fetch + symbol grep (no pip install, CPU-only). -ST is unpinned in unsloth/pyproject.toml; cover recent 5.x minors plus -`main`. -""" +"""Pinned-symbol compat check for the symbols unsloth's sentence_transformer +integration relies on, across ST PyPI minors (GitHub raw fetch + symbol grep).""" from __future__ import annotations @@ -29,8 +12,7 @@ import pytest from tests.version_compat._fetch import fetch_text, first_match, has_def -# Policy: unsloth/pyproject.toml does NOT pin sentence-transformers. We -# track the last few minors plus main. Add a row when a new minor lands. +# ST is unpinned in pyproject.toml; track the last few minors plus main. ST_TAGS = [ "v5.0.0", "v5.1.2", @@ -41,10 +23,7 @@ ST_TAGS = [ ] -# Top-level public surface: SentenceTransformer + SentenceTransformerTrainer -# must be importable as `from sentence_transformers import X`. - - +# Top-level: SentenceTransformer + SentenceTransformerTrainer must be importable. @pytest.mark.parametrize("tag", ST_TAGS) def test_st_top_level_exports(tag: str): src = fetch_text("UKPLab/sentence-transformers", tag, "sentence_transformers/__init__.py") @@ -57,20 +36,13 @@ def test_st_top_level_exports(tag: str): ) -# Sub-modules: Transformer / Pooling / Normalize. unsloth walks -# `sentence_transformers.models` to introspect these (line 1016, 1206). - - +# Sub-modules: unsloth walks `sentence_transformers.models` for these classes. @pytest.mark.parametrize("tag", ST_TAGS) def test_st_models_re_exports(tag: str): - """Transformer / Pooling / Normalize must be reachable through - `sentence_transformers.models`. ST 5.4 reorganised the package, but - the top-level re-export must still surface these three so user code - (and unsloth/models/sentence_transformer.py:1016,1206,1467) can - `from sentence_transformers.models import Transformer`.""" - # Layout 1 (legacy < 5.4): sentence_transformers/models[.py|/__init__.py]. - # Layout 2 (>= 5.4): top-level __init__.py re-exports the symbols - # plus the modules live under base/modules and sentence_transformer/. + """Transformer / Pooling / Normalize must stay reachable via + `sentence_transformers.models` despite the ST 5.4 package reorg.""" + # Layout 1 (legacy < 5.4): sentence_transformers/models[.py|/__init__.py]. + # Layout 2 (>= 5.4): top-level re-exports; modules under base/modules + sentence_transformer/. legacy_candidates = [ "sentence_transformers/models/__init__.py", "sentence_transformers/models.py", @@ -88,9 +60,8 @@ def test_st_models_re_exports(tag: str): return # ST 5.4+ modular layout: classes moved under base/modules and - # sentence_transformer/modules. Backward compat for - # `from sentence_transformers.models import X` is wired at import via - # setup_deprecated_module_imports in sentence_transformers/__init__.py. + # sentence_transformer/modules; backward compat wired via + # setup_deprecated_module_imports in __init__.py. expected_paths = { "Transformer": [ "sentence_transformers/base/modules/transformer.py", @@ -114,8 +85,7 @@ def test_st_models_re_exports(tag: str): else: pytest.fail(f"{tag}: ST 5.4+ layout: class {cls} not found in any of {paths}") - # The backward-compat shim must be wired up so user code doing - # `from sentence_transformers.models import Pooling` keeps working. + # The backward-compat shim must be wired so `from ...models import Pooling` keeps working. top = fetch_text("UKPLab/sentence-transformers", tag, "sentence_transformers/__init__.py") assert top is not None, f"{tag}: sentence_transformers/__init__.py missing" has_shim = bool( @@ -130,10 +100,7 @@ def test_st_models_re_exports(tag: str): ) -# Transformer base class: unsloth checks two alternate paths at -# sentence_transformer.py:1169-1171. At least ONE must resolve. - - +# Transformer base class: unsloth probes alternate paths; at least ONE must resolve. @pytest.mark.parametrize("tag", ST_TAGS) def test_st_transformer_base_class_either_path(tag: str): candidates = [ @@ -153,16 +120,11 @@ def test_st_transformer_base_class_either_path(tag: str): ) -# sentence_transformers.util: import_from_string + load_dir_path are the -# two helpers unsloth.models.sentence_transformer:1177,1205 calls. - - +# sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls. @pytest.mark.parametrize("tag", ST_TAGS) def test_st_util_helpers(tag: str): - """`sentence_transformers.util.{import_from_string, load_dir_path}` — - used by unsloth.models.sentence_transformer:1177,1205. ST 5.4+ moved - util into a package; accept either layout, or a re-export from any - util submodule.""" + """util.{import_from_string, load_dir_path} must resolve; accept either the + flat or the ST 5.4+ package layout, or a re-export from a util submodule.""" candidates = [ "sentence_transformers/util.py", "sentence_transformers/util/__init__.py", @@ -174,7 +136,7 @@ def test_st_util_helpers(tag: str): defined_here = has_def(src, fn, "func") reexported = bool(re.search(rf"\b{re.escape(fn)}\b", src)) if not (defined_here or reexported): - # Try common subfiles for the modular layout. + # Modular-layout subfiles. subpaths = [ "sentence_transformers/util/import_utils.py", "sentence_transformers/util/file_utils.py", diff --git a/tests/version_compat/test_transformers_pinned_symbols.py b/tests/version_compat/test_transformers_pinned_symbols.py index cebef77c7f..6aa628e86b 100644 --- a/tests/version_compat/test_transformers_pinned_symbols.py +++ b/tests/version_compat/test_transformers_pinned_symbols.py @@ -1,32 +1,10 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -"""Pinned-symbol + source-pattern compat checks across the -transformers PyPI window unsloth + unsloth-zoo target. Catches the -classes of breakage we've shipped fixes for in: +"""Pinned-symbol + source-pattern transformers compat checks via GitHub raw-fetch + grep. - unsloth#3998 notebook compat 4.57.6 + TRL 0.22-0.27 - unsloth#5036 grad-accum accepts_loss_kwargs vision wrappers - unsloth#5155 resolve_model_class fallback against unresolvable AutoModel - unsloth#5259 FastSentenceTransformer + ST 5.4 redirect - unsloth-zoo#572 forward-compat with transformers 5.x decorators + Qwen2VL - unsloth-zoo#571 gemma3, csm, ministral, pixtral 5.3 forward signature - unsloth-zoo#549 VRAM regression with transformers 5.2+ checkpoint - unsloth-zoo#543 GRPO logging + transformers v5 loss shape mismatch - unsloth-zoo#541 got multiple values for argument in compiled forward dispatch - unsloth-zoo#495 Qwen3Next/Qwen3.5 MoE + transformers v5 fixes for Gemma - unsloth-zoo#491 should_convert_module substring matching - unsloth-zoo#488 Gemma3 + Gemma3N transformers 5.x - unsloth-zoo#472 ModernBERT, gpt_oss MoE unwrap, SFTTrainer skip_prepare_dataset - unsloth-zoo#393 PushToHubMixin._create_repo removed in v5 - unsloth-zoo#388 generation_config attribute removed for non-gen models in v5 - unsloth-zoo#583/584 PIL _Ink ImportError (Unpack import guard) - unsloth-zoo#159 cross_entropy_replacement_2 num_items_in_batch fallback - -Strategy: GitHub raw-fetch + grep / source-fingerprint. CPU-only, no -install. Runs PR-time + daily cron. - -Anchor versions (must work forwards/backwards-compat per project spec): - transformers 4.57.6, 5.5.0 +Catches breakage classes from unsloth#3998/5036/5155/5259 and +unsloth-zoo#572/571/549/543/541/495/491/488/472/393/388/583/584/159. +CPU-only, no install. Anchor versions: transformers 4.57.6, 5.5.0. """ from __future__ import annotations @@ -38,8 +16,7 @@ import pytest from tests.version_compat._fetch import fetch_text, first_match, has_def -# Stable transformers from 4.57.6 floor onwards + main. The breakage -# windows we care about are 4.57.6, then every 5.x minor since 5.0.0. +# 4.57.6 floor + every 5.x minor since 5.0.0 + main. TRANSFORMERS_TAGS = [ "v4.57.6", # anchor (must work) "v5.0.0", @@ -56,16 +33,12 @@ TRANSFORMERS_TAGS = [ ] -# ========================================================================= -# Trainer surface — the largest failure class. unsloth/models/_utils.py -# rewrites Trainer.{__init__, training_step, get_batch_samples, compute_loss}. -# ========================================================================= +# Trainer surface: unsloth/models/_utils.py rewrites Trainer.{__init__, training_step, get_batch_samples, compute_loss}. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_trainer_class_importable_path(tag: str): - """transformers.Trainer must remain at src/transformers/trainer.py - or src/transformers/trainer/__init__.py.""" + """transformers.Trainer must remain at trainer.py or trainer/__init__.py.""" candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"] hit = first_match("huggingface/transformers", tag, candidates) assert hit is not None, f"{tag}: src/transformers/trainer[.py|/__init__.py] both missing" @@ -75,13 +48,11 @@ def test_trainer_class_importable_path(tag: str): @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_trainer_compute_loss_num_items_in_batch_param(tag: str): - """unsloth-zoo#159 + unsloth#4998 + #4616: Trainer.compute_loss - must accept num_items_in_batch kwarg. transformers 4.46+ added it.""" + """unsloth-zoo#159 + unsloth#4998 + #4616: Trainer.compute_loss must accept num_items_in_batch kwarg.""" candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"] hit = first_match("huggingface/transformers", tag, candidates) assert hit is not None _, src = hit - # Find the compute_loss signature - it's a class method, indented. m = re.search(r"^\s*def compute_loss\(([^)]*)\)", src, re.MULTILINE | re.DOTALL) if m is None: pytest.fail(f"{tag}: Trainer.compute_loss not found in source") @@ -93,9 +64,7 @@ def test_trainer_compute_loss_num_items_in_batch_param(tag: str): @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_trainer_training_step_grad_accum_pattern(tag: str): - """unsloth#3598 monkey-patches Trainer.training_step source; the - rewrite needs four substrings to be present. Drift here = silent - no-op = double-scale loss bug.""" + """unsloth#3598 patches Trainer.training_step source; drift = silent no-op = double-scale loss bug.""" candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"] hit = first_match("huggingface/transformers", tag, candidates) assert hit is not None @@ -106,8 +75,7 @@ def test_trainer_training_step_grad_accum_pattern(tag: str): "self.accelerator.backward(loss", ) missing = [s for s in needed if s not in src] - # Hard-fail only when ALL substrings missing — partial drift is - # informational. Note: the third one's exact form may vary slightly. + # Hard-fail only when ALL substrings missing; partial drift is informational. if len(missing) == len(needed): pytest.fail( f"{tag}: Trainer.training_step has none of the grad-accum " @@ -118,8 +86,7 @@ def test_trainer_training_step_grad_accum_pattern(tag: str): @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_trainer_get_batch_samples_returns_num_items(tag: str): - """unsloth-zoo loss_utils.py:241 replaces Trainer.get_batch_samples; - upstream signature must end `return batch_samples, num_items_in_batch`.""" + """unsloth-zoo loss_utils.py:241 replaces Trainer.get_batch_samples; must keep the num_items_in_batch return.""" candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"] hit = first_match("huggingface/transformers", tag, candidates) assert hit is not None @@ -133,19 +100,14 @@ def test_trainer_get_batch_samples_returns_num_items(tag: str): @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_trainer_inner_training_loop_inplace_loss_v5(tag: str): - """unsloth-zoo#543: transformers 5.0+ changed - `tr_loss = tr_loss + tr_loss_step` (out-of-place) to - `self._tr_loss += tr_loss_step` (in-place). Loss tensor shape - requirements differ. Snapshot which form is in source.""" + """unsloth-zoo#543: transformers 5.0+ switched out-of-place tr_loss add to in-place `self._tr_loss +=`.""" candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"] hit = first_match("huggingface/transformers", tag, candidates) assert hit is not None _, src = hit has_inplace = "self._tr_loss +=" in src has_outplace = "tr_loss = tr_loss + tr_loss_step" in src - # On 4.57.6, only out-of-place. On 5.x, in-place. We just assert - # ONE of them is present so a future refactor that drops both is - # caught. + # Assert ONE form is present so a refactor dropping both is caught. assert has_inplace or has_outplace, ( f"{tag}: Trainer._inner_training_loop has neither " f"`tr_loss = tr_loss + tr_loss_step` nor `self._tr_loss +=`; " @@ -153,16 +115,12 @@ def test_trainer_inner_training_loop_inplace_loss_v5(tag: str): ) -# ========================================================================= -# modeling_utils — checkpoint, PushToHubMixin, ALL_ATTENTION_FUNCTIONS. -# ========================================================================= +# modeling_utils: checkpoint, PushToHubMixin, ALL_ATTENTION_FUNCTIONS. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_modeling_utils_exposes_checkpoint(tag: str): - """unsloth-zoo#549: transformers 5.2+ uses `transformers.modeling_utils.checkpoint` - (alias for torch.utils.checkpoint.checkpoint). Patch must replace - the transformers reference, not just torch's.""" + """unsloth-zoo#549: transformers 5.2+ uses modeling_utils.checkpoint; patch must replace it, not just torch's.""" src = fetch_text("huggingface/transformers", tag, "src/transformers/modeling_utils.py") if src is None: pytest.skip(f"{tag}: modeling_utils.py missing") @@ -184,20 +142,16 @@ def test_modeling_utils_exposes_checkpoint(tag: str): @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_pushtohubmixin_create_repo_status(tag: str): - """unsloth-zoo#393: transformers 5.x removed PushToHubMixin._create_repo. - On 4.x present, on 5.x absent. Snapshot which side.""" + """unsloth-zoo#393: transformers 5.x removed PushToHubMixin._create_repo; snapshot which side.""" src = fetch_text("huggingface/transformers", tag, "src/transformers/modeling_utils.py") if src is None: pytest.skip(f"{tag}: modeling_utils.py missing") - # Just record the presence; either is OK as long as we know. has_create = bool(re.search(r"def _create_repo\b", src) or "_create_repo" in src) - # Informational only — both branches are tracked. + # Informational only. _ = has_create -# ========================================================================= -# integrations.bitsandbytes — _replace_with_bnb_linear vs new path. -# ========================================================================= +# integrations.bitsandbytes: _replace_with_bnb_linear vs new path. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) @@ -214,9 +168,7 @@ def test_integrations_bitsandbytes_module_present(tag: str): @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_quantizers_should_convert_module_signature(tag: str): - """unsloth-zoo#491/#488: 5.x moved is_replaceable to - quantizers_utils.should_convert_module(full_name, patterns). - Snapshot whether function exists and its substring-match form.""" + """unsloth-zoo#491/#488: 5.x moved is_replaceable to quantizers_utils.should_convert_module; snapshot its form.""" src = fetch_text( "huggingface/transformers", tag, @@ -226,23 +178,18 @@ def test_quantizers_should_convert_module_signature(tag: str): pytest.skip(f"{tag}: quantizers/quantizers_utils.py missing") if not has_def(src, "should_convert_module", "func"): pytest.skip(f"{tag}: should_convert_module not yet present (4.x)") - # The bug we want to catch: substring matching uses `.{key}.` in - # `.{full_name}.` form. Patch only fires when this substring is - # in source AND mismatch behaviour exists. + # Catch substring matching in `.{key}.` form. has_dot_form = ".{key}." in src or "f'.{key}.'" in src or 'f".{key}."' in src # Informational only. _ = has_dot_form -# ========================================================================= -# integrations.finegrained_fp8.FP8Linear — bias/has_bias rename in v5. -# ========================================================================= +# integrations.finegrained_fp8.FP8Linear: bias/has_bias rename in v5. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_fp8linear_init_param_names(tag: str): - """unsloth-zoo#572: transformers 5.x renamed FP8Linear.__init__ - `bias` -> `has_bias`. Snapshot which form is in source.""" + """unsloth-zoo#572: transformers 5.x renamed FP8Linear.__init__ `bias` -> `has_bias`.""" src = fetch_text( "huggingface/transformers", tag, @@ -259,15 +206,12 @@ def test_fp8linear_init_param_names(tag: str): ), f"{tag}: FP8Linear.__init__ has neither `bias` nor `has_bias` param" -# ========================================================================= -# processing_utils — Unpack importable. -# ========================================================================= +# processing_utils: Unpack importable. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_processing_utils_unpack_importable(tag: str): - """unsloth-zoo#583/584: `from transformers.processing_utils import Unpack` - must keep working.""" + """unsloth-zoo#583/584: transformers.processing_utils.Unpack must keep importing.""" src = fetch_text("huggingface/transformers", tag, "src/transformers/processing_utils.py") if src is None: pytest.skip(f"{tag}: processing_utils.py missing") @@ -278,9 +222,7 @@ def test_processing_utils_unpack_importable(tag: str): ) -# ========================================================================= -# Models — gemma3, gpt_oss forward signature drift. -# ========================================================================= +# Models: gemma3, gpt_oss forward signature drift. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) @@ -307,16 +249,12 @@ def test_gpt_oss_model_forward_present(tag: str): assert has_def(src, "GptOssModel", "class"), f"{tag}: class GptOssModel missing" -# ========================================================================= -# auto_factory — unsloth#5155 _LazyAutoMapping private API. -# ========================================================================= +# auto_factory: unsloth#5155 _LazyAutoMapping private API. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_auto_factory_lazy_mapping_private_api(tag: str): - """unsloth#5155: resolve_model_class iterates private attrs of - _LazyAutoMapping (_model_mapping, _config_mapping, _extra_content, - _load_attr_from_module). All four must remain.""" + """unsloth#5155: resolve_model_class needs all four _LazyAutoMapping private attrs to remain.""" src = fetch_text( "huggingface/transformers", tag, @@ -337,15 +275,12 @@ def test_auto_factory_lazy_mapping_private_api(tag: str): ) -# ========================================================================= -# configuration_utils — PreTrainedConfig vs PretrainedConfig in 5.x. -# ========================================================================= +# configuration_utils: PreTrainedConfig vs PretrainedConfig in 5.x. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_configuration_utils_alias(tag: str): - """transformers 5.x renamed PretrainedConfig -> PreTrainedConfig. - unsloth-zoo/empty_model.py imports from both paths defensively.""" + """transformers 5.x renamed PretrainedConfig -> PreTrainedConfig; unsloth-zoo imports both defensively.""" src = fetch_text( "huggingface/transformers", tag, @@ -361,16 +296,12 @@ def test_configuration_utils_alias(tag: str): ) -# ========================================================================= -# tokenization — apply_chat_template return_dict default flip in v5. -# ========================================================================= +# tokenization: apply_chat_template return_dict default flip in v5. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_apply_chat_template_signature_present(tag: str): - """unsloth-zoo#572: PreTrainedTokenizerBase.apply_chat_template - `return_dict` default flipped False -> True in transformers 5.x. - Snapshot which is in source.""" + """unsloth-zoo#572: apply_chat_template `return_dict` default flipped False -> True in transformers 5.x.""" src = fetch_text( "huggingface/transformers", tag, @@ -383,16 +314,12 @@ def test_apply_chat_template_signature_present(tag: str): ), f"{tag}: apply_chat_template missing in tokenization_utils_base.py" -# ========================================================================= -# Generic-importability sweep — every symbol unsloth/zoo imports -# from transformers must remain reachable via at least one known path. -# ========================================================================= +# Generic-importability sweep: every transformers symbol unsloth/zoo imports must stay reachable. @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS) def test_modeling_attn_mask_utils_symbols(tag: str): - """_prepare_4d_attention_mask_for_sdpa is imported by - unsloth/models/llama.py + sentence_transformer.py.""" + """_prepare_4d_attention_mask_for_sdpa is imported by unsloth/models/llama.py + sentence_transformer.py.""" src = fetch_text( "huggingface/transformers", tag, @@ -401,7 +328,6 @@ def test_modeling_attn_mask_utils_symbols(tag: str): if src is None: pytest.skip(f"{tag}: modeling_attn_mask_utils.py missing") assert has_def(src, "AttentionMaskConverter", "class"), f"{tag}: AttentionMaskConverter missing" - # _prepare_4d_attention_mask_for_sdpa is a function we hard-import. assert ( has_def(src, "_prepare_4d_attention_mask_for_sdpa", "func") or "_prepare_4d_attention_mask_for_sdpa" in src diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py index f84a4fd669..834692e2dd 100644 --- a/tests/version_compat/test_trl_grpo_pinned_symbols.py +++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py @@ -1,36 +1,13 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -"""Pinned-symbol compat check across all TRL PyPI minor versions -unsloth + unsloth-zoo target. Catches API drift like: +"""Pinned-symbol compat check across all TRL PyPI minors unsloth + unsloth-zoo target. - - trl 0.18 split DataCollatorForPreference into trl.trainer.dpo_trainer - (was trl.trainer.utils). unsloth.models.rl_replacements:318 imports - the post-split path; if a new TRL release moves it again, the - GRPOTrainer.compile cell crashes with ImportError. - - trl 0.20 introduced trl.experimental.openenv as a *gated* module; - unsloth.models.rl_replacements:1765-1770 catches ImportError, but - the gate must remain importable when present. - - trl 0.22 introduced trl.generation.vllm_generation for the - server-mode fast_inference path; unsloth.models.rl_replacements - :1846-1848 catches ImportError, but the module must exist on - versions where unsloth-zoo's vllm_utils dispatches to it. - - trl unwrap_model_for_generation moved from trl.models to - trl.models.utils across releases (unsloth/models/rl.py:152-155 - handles both with try/except). - - trl GRPOTrainer / GRPOConfig must remain top-level exports for - `from trl import GRPOTrainer` to work in user code, which is what - `_patch_trl_rl_trainers("grpo_trainer")` discovers. - -Strategy: for each tracked TRL tag, fetch the relevant source files -straight from github.com/huggingface/trl (no pip install required) and -assert that every symbol unsloth/unsloth-zoo's RL surface depends on -is present. - -Versioning policy: cover the supported window declared in -pyproject.toml (`trl>=0.18.2,!=0.19.0,<=0.24.0`) PLUS several recent -releases ABOVE the cap, so we get early warning when TRL ships -something incompatible and the maintainer can extend the cap or add a -patch BEFORE a user hits it. +Catches RL-surface API drift (DataCollatorForPreference relocation, +gated openenv/vllm_generation modules, unwrap_model_for_generation +moves, top-level GRPO exports). Fetches TRL source per tag straight from +github (no pip install) and asserts every depended-on symbol is present, +covering the pyproject window plus several releases above the cap for +early warning. """ from __future__ import annotations @@ -42,20 +19,10 @@ import pytest from tests.version_compat._fetch import fetch_text, first_match, has_def -# Every stable TRL release from 0.18.2 (the pyproject floor) onwards, -# plus `main`. Refresh by running: -# python -c "import urllib.request,json -# from packaging.version import Version -# r=json.loads(urllib.request.urlopen('https://pypi.org/pypi/trl/json').read()) -# v=sorted([Version(x) for x in r['releases'] if r['releases'][x] and not Version(x).is_prerelease and Version(x)>=Version('0.18.2')]) -# print(*[f'\"v{x}\",' for x in v],sep='\n')" -# -# 0.19.0 is excluded by pyproject (`!=0.19.0`) — the release was -# broken; we keep it in the matrix so we KNOW it's broken (and which -# symbols specifically), not just trust the pin. -# -# Anchors (per the project spec, ALL patches must stay forwards/ -# backwards compatible with these): 0.22.2, 0.27.1, 1.0.0. +# Every stable TRL release from 0.18.2 (pyproject floor) onwards, plus `main`. +# 0.19.0 is pyproject-excluded (broken) but kept here so we know exactly +# which symbols break. Anchors all patches stay compatible with: 0.22.2, +# 0.27.1, 1.0.0. TRL_TAGS = [ "v0.18.2", "v0.19.0", @@ -88,14 +55,12 @@ TRL_TAGS = [ ] -# HARD-import top-level: from trl import X must keep working for these. -# unsloth/trainer.py + unsloth/models/rl.py rebind these by name. +# unsloth/trainer.py + unsloth/models/rl.py rebind these top-level names. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_top_level_grpo_sft(tag: str): - """`from trl import GRPOTrainer, GRPOConfig, SFTTrainer, SFTConfig` - must keep resolving at the package root.""" + """GRPO/SFT Trainer+Config must resolve at the trl package root.""" src = fetch_text("huggingface/trl", tag, "trl/__init__.py") assert src is not None, f"trl/__init__.py missing in {tag}" for name in ("GRPOTrainer", "GRPOConfig", "SFTTrainer", "SFTConfig"): @@ -105,8 +70,8 @@ def test_trl_top_level_grpo_sft(tag: str): ) -# trl.trainer.grpo_trainer.GRPOTrainer -- the canonical class. unsloth's -# RL patcher discovers it via `eval(f"trl.trainer.{trainer_file}.{name}")` +# trl.trainer.grpo_trainer.GRPOTrainer -- canonical class. unsloth's RL +# patcher discovers it via `eval(f"trl.trainer.{trainer_file}.{name}")` # in unsloth/models/rl.py:548-594. @@ -124,9 +89,7 @@ def test_grpo_trainer_class_canonical_path(tag: str): @pytest.mark.parametrize("tag", TRL_TAGS) def test_grpo_config_class_canonical_path(tag: str): - """unsloth/models/rl.py:579-618 looks for the *Config sibling of the - Trainer class via heuristic discovery; the canonical one is in - grpo_config.py.""" + """GRPOConfig must be discoverable by the *Config heuristic in rl.py:579-618.""" candidates = ["trl/trainer/grpo_config.py", "trl/trainer/grpo_trainer.py"] hit = first_match("huggingface/trl", tag, candidates) assert hit is not None, f"{tag}: neither grpo_config.py nor grpo_trainer.py found" @@ -137,17 +100,13 @@ def test_grpo_config_class_canonical_path(tag: str): ) -# DataCollatorForPreference: unsloth.models.rl_replacements:318 hard-imports -# from trl.trainer.dpo_trainer. Some old TRL versions had it in -# trl.trainer.utils; modern ones moved to trl.trainer.dpo_trainer. +# DataCollatorForPreference: rl_replacements.py:318 hard-imports from +# trl.trainer.dpo_trainer (old TRL had it in trl.trainer.utils). @pytest.mark.parametrize("tag", TRL_TAGS) def test_data_collator_for_preference_resolvable(tag: str): - """Either the new path (trl.trainer.dpo_trainer) or the old path - (trl.trainer.utils) must define DataCollatorForPreference. unsloth's - string-emitted import in rl_replacements.py:318 uses dpo_trainer; - if neither path resolves, we have a gap.""" + """DataCollatorForPreference must exist in dpo_trainer or utils (rl_replacements.py:318 imports it).""" new_path = fetch_text("huggingface/trl", tag, "trl/trainer/dpo_trainer.py") old_path = fetch_text("huggingface/trl", tag, "trl/trainer/utils.py") have = [] @@ -170,8 +129,7 @@ def test_data_collator_for_preference_resolvable(tag: str): def test_trl_trainer_utils_pad(tag: str): src = fetch_text("huggingface/trl", tag, "trl/trainer/utils.py") if src is None: - # Some TRL versions split utils into a package; check the - # alternative location. + # Some TRL versions split utils into a package. src = fetch_text("huggingface/trl", tag, "trl/trainer/utils/__init__.py") assert src is not None, f"{tag}: trl/trainer/utils[.py|/__init__.py] both missing" assert has_def(src, "pad", "func") or "def pad(" in src, ( @@ -188,11 +146,7 @@ def test_trl_trainer_utils_pad(tag: str): @pytest.mark.parametrize("tag", TRL_TAGS) def test_unwrap_model_for_generation_either_path(tag: str): - """unsloth/models/rl.py:152-155 tries - `trl.models.utils.unwrap_model_for_generation` first, then - `trl.models.unwrap_model_for_generation`. Tests must mirror the - prod fallback exactly — checking a third path makes the test - laxer than the runtime.""" + """unwrap_model_for_generation must resolve via one of the two paths rl.py:152-155 tries (mirror prod exactly).""" candidates = [ "trl/models/utils.py", "trl/models/__init__.py", @@ -209,19 +163,17 @@ def test_unwrap_model_for_generation_either_path(tag: str): ) -# trl.experimental.openenv: gated import (rl_replacements.py:1765-1770 -# wraps in try/except). When present, must export the symbols unsloth -# patches. +# trl.experimental.openenv: gated import (rl_replacements.py:1765-1770). +# When present, must export the symbols unsloth patches. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_experimental_openenv_gated(tag: str): src = fetch_text("huggingface/trl", tag, "trl/experimental/openenv/__init__.py") if src is None: - # OK: feature not in this release; unsloth's try/except handles it. pytest.skip(f"{tag}: trl.experimental.openenv not present (OK)") - # Module exists -> at minimum, `utils` submodule must be importable - # because unsloth patches via `import trl.experimental.openenv.utils`. + # Module exists -> utils submodule must be importable (unsloth patches + # via `import trl.experimental.openenv.utils`). utils_src = fetch_text("huggingface/trl", tag, "trl/experimental/openenv/utils.py") assert utils_src is not None, ( f"{tag}: trl.experimental.openenv exists but utils.py missing; " @@ -230,20 +182,16 @@ def test_trl_experimental_openenv_gated(tag: str): # trl.generation.vllm_generation: gated import for the fast_inference -# server mode (rl_replacements.py:1846-1848). When present, must define -# at least one symbol unsloth patches against. +# server mode (rl_replacements.py:1846-1848). @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_generation_vllm_generation_gated(tag: str): - """unsloth/models/rl_replacements.py:1851-1971 string-rewrites - `VLLMGeneration._init_vllm`, `.sync_weights`, and `.generate`. If - VLLMGeneration is renamed or any of those three methods disappear, - the rewrite silently no-ops and the fast_inference server path - breaks at runtime. Gated: skip if the module isn't in this TRL.""" + """VLLMGeneration + its _init_vllm/sync_weights/generate methods must + exist when the module is present, else rl_replacements.py:1851-1971 + rewrites silently no-op and the server fast_inference path breaks.""" src = fetch_text("huggingface/trl", tag, "trl/generation/vllm_generation.py") if src is None: - # OK: pre-server-mode TRL. unsloth's try/except handles absence. pytest.skip(f"{tag}: trl.generation.vllm_generation not present (OK)") assert has_def(src, "VLLMGeneration", "class"), ( f"{tag}: class VLLMGeneration missing; unsloth-zoo dispatch " @@ -256,21 +204,15 @@ def test_trl_generation_vllm_generation_gated(tag: str): ) -# Sanity: TRL's __version__ string is parseable. unsloth/models/rl.py:63 -# does `from trl import __version__ as trl_version_raw` and string- -# matches on it. +# TRL's __version__ must be parseable; rl.py:63 string-matches it. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_version_parseable(tag: str): src = fetch_text("huggingface/trl", tag, "trl/__init__.py") assert src is not None - # Recognised mechanisms (any one is sufficient): - # 1. literal `__version__ = "x.y.z"` at module scope - # 2. `from .version import __version__` - # 3. `__version__ = version("trl")` via importlib.metadata - # 4. `__version__ = f.read().strip()` (TRL 0.22.x reads from a - # sibling VERSION file) + # Any one mechanism suffices: literal, `from .version import`, + # importlib.metadata version(), or reading a sibling VERSION file. has_literal = bool(re.search(r'^__version__\s*=\s*["\']', src, re.MULTILINE)) has_subimport = bool(re.search(r"^from\s+\.version\s+import\s+__version__", src, re.MULTILINE)) has_metadata = bool( @@ -291,8 +233,8 @@ def test_trl_version_parseable(tag: str): ) -# Coverage extension (added 2026-05): symbols / source-string contracts -# unsloth + unsloth-zoo touch but the original suite missed. +# Coverage extension (added 2026-05): symbols/source-string contracts +# unsloth + unsloth-zoo touch that the original suite missed. # 1. trl.is_conversational — soft import in unsloth-zoo dataset_utils. @@ -303,34 +245,29 @@ def test_trl_is_conversational_export(tag: str): src = fetch_text("huggingface/trl", tag, "trl/__init__.py") assert src is not None if "is_conversational" not in src: - # Some old TRLs omit it; gated soft import in unsloth-zoo - # falls back to a local impl. OK. + # Old TRLs omit it; unsloth-zoo's gated soft import falls back. pytest.skip(f"{tag}: trl.is_conversational not exported (legacy TRL)") -# 2-4. trl.trainer.sft_trainer module surface used by unsloth tokenizer -# utils + tests. +# 2-4. trl.trainer.sft_trainer surface used by unsloth tokenizer utils + tests. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_sft_trainer_module_internals(tag: str): - """unsloth/tokenizer_utils.py:1538 does `from trl.trainer.sft_trainer - import *`. The symbols below must exist for the wildcard import + - eval-discovery to keep working.""" + """sft_trainer symbols for the `from trl.trainer.sft_trainer import *` at tokenizer_utils.py:1538.""" src = fetch_text("huggingface/trl", tag, "trl/trainer/sft_trainer.py") assert src is not None, ( f"{tag}: trl/trainer/sft_trainer.py missing; " f"unsloth/tokenizer_utils.py:1538 wildcard import fails" ) assert has_def(src, "SFTTrainer", "class"), f"{tag}: class SFTTrainer missing in sft_trainer.py" - # neftune_post_forward_hook: optional (TRL removed it in some - # versions); soft-imported in tokenizer_utils.py:1542. Don't fail. + # neftune_post_forward_hook: optional, soft-imported in tokenizer_utils.py:1542. if "neftune_post_forward_hook" not in src: pass -# 5-6. trl.trainer.dpo_trainer module + MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES -# — patched by unsloth-zoo/temporary_patches/misc.py:1376-1379. +# 5-6. trl.trainer.dpo_trainer + MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES, +# patched by unsloth-zoo/temporary_patches/misc.py:1376-1379. @pytest.mark.parametrize("tag", TRL_TAGS) @@ -343,9 +280,8 @@ def test_trl_dpo_trainer_module_exists(tag: str): assert has_def(src, "DPOTrainer", "class"), f"{tag}: class DPOTrainer missing in dpo_trainer.py" -# 7. trl.trainer.utils.ConstantLengthDataset — soft import in -# unsloth-zoo/dataset_utils.py:596. Optional (TRL 0.20.0 removed it -# on some paths). +# 7. trl.trainer.utils.ConstantLengthDataset — optional soft import in +# unsloth-zoo/dataset_utils.py:596 (TRL 0.20.0 removed it on some paths). @pytest.mark.parametrize("tag", TRL_TAGS) @@ -364,19 +300,17 @@ def test_trl_constant_length_dataset_optional(tag: str): ) -# 8. trl.models.utils.disable_gradient_checkpointing — added in TRL -# 1.0.0+. unsloth/models/rl.py:1976-1994 uses hasattr() for gating; -# we still want the assertion that the symbol exists from 1.0.0 -# onwards so a future removal gets caught. +# 8. trl.models.utils.disable_gradient_checkpointing — added in TRL 1.0.0+. +# rl.py:1976-1994 gates via hasattr(); assert the symbol exists from +# 1.0.0 onwards so a future removal gets caught. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_models_utils_disable_gradient_checkpointing(tag: str): if tag == "main": - # main is bleeding edge; expect symbol to track 1.0.0+ behaviour. + # main is bleeding edge; expect 1.0.0+ behaviour. require = True else: - # Strip leading 'v' and parse. try: from packaging.version import Version require = Version(tag.lstrip("v")) >= Version("1.0.0") @@ -395,9 +329,8 @@ def test_trl_models_utils_disable_gradient_checkpointing(tag: str): ) -# 9. trl.import_utils + the `_*_available` cache pattern — used by -# unsloth/import_fixes.py:508-516 to clear cached `is_X_available` -# booleans so vllm-ascend imports work. +# 9. trl.import_utils `_*_available` cache pattern — import_fixes.py:508-516 +# clears these cached booleans so vllm-ascend imports work. @pytest.mark.parametrize("tag", TRL_TAGS) @@ -410,9 +343,8 @@ def test_trl_import_utils_available_pattern(tag: str): if hit is None: pytest.skip(f"{tag}: trl/import_utils not present (legacy TRL)") _, src = hit - # The patch iterates `vars(trl.import_utils)` looking for any name - # ending in `_available`. At least one such cache var must exist or - # the patch silently no-ops. + # import_fixes iterates vars(trl.import_utils) for `*_available` names; + # at least one must exist or the patch silently no-ops. has_pattern = bool(re.search(r"\b\w+_available\b", src)) assert has_pattern, ( f"{tag}: trl.import_utils has no `_available` cache var; " @@ -420,9 +352,8 @@ def test_trl_import_utils_available_pattern(tag: str): ) -# 10. trl.experimental.openenv.utils generators — at least one of the -# two function names must exist (unsloth/models/rl_replacements.py -# :1775-1781 calls getattr() to find one). +# 10. trl.experimental.openenv.utils generators — one of the two function +# names must exist (rl_replacements.py:1775-1781 getattr()s for one). @pytest.mark.parametrize("tag", TRL_TAGS) @@ -439,22 +370,18 @@ def test_trl_openenv_utils_generators(tag: str): ) -# 11-16. GRPOTrainer required method names. unsloth/models/rl_replacements -# .py uses function_name == "..." dispatch keys; if a method is -# renamed, the patch silently doesn't apply. List of methods is -# the precise dispatch key set. +# 11-16. GRPOTrainer required method names. rl_replacements.py dispatches +# on function_name == "..."; a renamed method silently skips the patch. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_grpo_trainer_required_methods(tag: str): - """Method names unsloth string-rewrites against. Drift here - silently skips the rewrite. _get_per_token_logps was renamed to - _get_per_token_logps_and_entropies in TRL 0.20+; either is fine - since unsloth dispatches by function_name.""" + """GRPOTrainer methods unsloth rewrites against; drift silently skips + the rewrite. _get_per_token_logps was renamed to + _get_per_token_logps_and_entropies in TRL 0.20+; either is fine.""" src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") assert src is not None - # _prepare_inputs / _generate_and_score_completions / compute_loss - # are stable across the entire support window. + # These three are stable across the entire support window. for m in ("_prepare_inputs", "_generate_and_score_completions", "compute_loss"): assert has_def(src, m, "func"), ( f"{tag}: GRPOTrainer.{m} missing; " @@ -468,26 +395,22 @@ def test_trl_grpo_trainer_required_methods(tag: str): f"._get_per_token_logps_and_entropies (TRL >=0.20) found; " f"unsloth's per-token-logps rewrite no-ops on both dispatch keys" ) - # Optional / version-dependent — never fail, just informational + # Optional/version-dependent — informational only. for m in ("_generate_single_turn", "_move_model_to_vllm", "_calculate_rewards"): _present = has_def(src, m, "func") _ = _present -# Source-string contracts on trl/trainer/grpo_trainer.py. Each substring -# is one half of a `function.replace(old, new)` rewrite — if the -# substring no longer appears in TRL source, the rewrite is a no-op -# AND the user-facing GRPO behaviour silently diverges. -# -# Broken into per-version-window tests because some patterns only apply -# to a subset of TRL minors. +# Source-string contracts on grpo_trainer.py. Each substring is one half +# of a `function.replace(old, new)` rewrite; if it vanishes from TRL +# source the rewrite no-ops and GRPO behaviour silently diverges. Split +# per version-window since some patterns apply only to a subset of minors. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_grpo_source_inference_mode_unwrap(tag: str): - """rl_replacements.py:526-535 inserts an autocast block immediately - AFTER `with torch.inference_mode():` and `self.accelerator.unwrap_model - (self.model)`. Both substrings must appear in `_prepare_inputs`.""" + """`torch.inference_mode` and `self.accelerator.unwrap_model` must both + appear, or rl_replacements.py:526-535 autocast insertion no-ops.""" src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") assert src is not None has_inference_mode = "torch.inference_mode" in src @@ -499,18 +422,15 @@ def test_trl_grpo_source_inference_mode_unwrap(tag: str): ) -# 17. KTOTrainer.get_batch_logps + the literal raise message rewriter -# hits. +# 17. KTOTrainer.get_batch_logps + the literal raise-message rewriter. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_kto_get_batch_logps_signature(tag: str): - """KTO log-prob computation must stay patchable. Through TRL 1.x the - target was KTOTrainer.get_batch_logps; TRL 1.x dropped it and moved the - math into _compute_logps / compute_ref_log_probs calling - selective_log_softmax. unsloth/models/rl_replacements.py patches BOTH - shapes (kto_trainer_get_batch_logps + kto_trainer_align_completion_logps), - so we require EITHER form to exist wherever KTOTrainer lives.""" + """KTO log-prob computation must stay patchable. Older TRL exposed + KTOTrainer.get_batch_logps; TRL 1.x moved the math into + _compute_logps/compute_ref_log_probs via selective_log_softmax. + rl_replacements.py patches both shapes, so require EITHER form.""" candidates = [ "trl/trainer/kto_trainer.py", "trl/experimental/kto/kto_trainer.py", @@ -528,9 +448,7 @@ def test_trl_kto_get_batch_logps_signature(tag: str): # TRL 1.x: refactored into _compute_logps + selective_log_softmax. if has_def(src, "_compute_logps", "func") and "selective_log_softmax" in src: return - # TRL 1.x (current): compute_ref_log_probs / _compute_kl_logps build - # per_token_logps via selective_log_softmax(shift_logits, ...); this is - # the exact shape kto_trainer_align_completion_logps patches. + # TRL 1.x current: the exact shape kto_trainer_align_completion_logps patches. if "per_token_logps = selective_log_softmax(shift_logits" in src: return old_shape_check = ( @@ -538,9 +456,8 @@ def test_trl_kto_get_batch_logps_signature(tag: str): 'must have the same shape.")' ) if checked_sources and not any(old_shape_check in src for _, src in checked_sources): - # TRL main inlined KTO log-prob computation and removed the old - # helper/shape-check rewrite target. There is no skipped rewrite to - # guard until a new concrete KTO shape mismatch target appears. + # TRL main inlined KTO log-probs and removed the old rewrite target; + # nothing to guard until a new concrete shape-mismatch target appears. return pytest.fail( f"{tag}: KTO log-prob computation not found in any of {candidates}; " @@ -548,19 +465,15 @@ def test_trl_kto_get_batch_logps_signature(tag: str): ) -# 18. SFTTrainer.__init__ literal `dict_args.pop("push_to_hub_token")` -# OR our shim must short-circuit. transformers 5.0 removed this -# kwarg; if TRL stops emitting the bare pop, our patch becomes -# a no-op AND TRL itself crashes on transformers 5.0. +# 18. SFTTrainer + the `dict_args.pop("push_to_hub_token")` shim. transformers +# 5.0 removed the kwarg; if TRL stops emitting the bare pop, our patch +# no-ops AND TRL itself crashes on transformers 5.0. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_sft_trainer_class(tag: str): - """Sanity: SFTTrainer.__init__ exists. The - `dict_args.pop("push_to_hub_token")` literal substring is checked - only when present — its absence means TRL already adapted (e.g. - via `dict_args.pop("push_to_hub_token", None)` with a default), - which is also fine.""" + """SFTTrainer must exist. The push_to_hub_token pop literal is checked + only when present; its absence means TRL already adapted (fine).""" src = fetch_text("huggingface/trl", tag, "trl/trainer/sft_trainer.py") assert src is not None assert has_def(src, "SFTTrainer", "class"), f"{tag}: class SFTTrainer missing" @@ -571,18 +484,10 @@ def test_trl_sft_trainer_class(tag: str): @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_dpo_trainer_methods(tag: str): - """DPOTrainer method-name surface unsloth's rewriters key on - (rl_replacements.py:222-394). All four are version-windowed: - - concatenated_inputs / concatenated_forward existed on - DPOTrainer through TRL 0.29.x; TRL 1.0+ refactored these into - free functions (concatenation moved out of the class). - - _compute_loss_liger added ~TRL 0.20. - - _set_signature_columns_if_needed: usually inherited from - transformers.Trainer, may or may not be re-defined locally. - None are STRICTLY required — when missing the matching unsloth - rewriter cleanly no-ops (TRL itself does the work). We surface - presence/absence as informational so a regression that - SILENTLY drops one is at least visible in the test log.""" + """DPOTrainer methods unsloth's rewriters key on (rl_replacements.py + :222-394). All version-windowed and non-required (the rewriter + cleanly no-ops when absent); presence/absence is logged as + informational so a silent regression stays visible.""" src = fetch_text("huggingface/trl", tag, "trl/trainer/dpo_trainer.py") assert src is not None # The DPO class itself must always exist. @@ -599,27 +504,17 @@ def test_trl_dpo_trainer_methods(tag: str): _ = _present # informational; rewriter no-ops cleanly when absent -# 22-23. trl.trainer.grpo_trainer must IMPORT or DEFINE the helpers -# unsloth's source rewriters reference: profiling_context, -# maybe_apply_chat_template, truncate_with_protected_tokens. -# Either the symbol is locally defined OR imported from elsewhere -# in trl.* — the rewriter only needs the NAME to be in scope at -# the call site. +# 22-23. grpo_trainer must have in scope the helpers unsloth's rewriters +# reference (profiling_context, maybe_apply_chat_template, +# truncate_with_protected_tokens), defined or imported from trl.*. @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_grpo_internal_helpers_in_scope(tag: str): - """Chat-template propagation is what unsloth's - grpo_trainer_fix_maybe_apply_chat_template wires up so user-supplied - `reasoning_effort` etc. survives the GRPO compile cell. The exact - helper name moved across releases: - - TRL <=0.24: `maybe_apply_chat_template(example, processing_class)` - appeared as a literal in grpo_trainer.py — unsloth's regex - rewriter substitutes it with a kwargs-aware version. - - TRL >=0.25: TRL itself uses `apply_chat_template` and pipes - `**self.chat_template_kwargs`, so the unsloth rewriter is a - cleanly-no-op'd dead path on those versions (correct behaviour). - Either pattern means the chat-template path is wired SOMEWHERE.""" + """Chat-template kwargs must propagate via legacy + `maybe_apply_chat_template` (TRL <=0.24, rewritten by unsloth) or + successor `apply_chat_template(... **chat_template_kwargs)` (TRL + >=0.25, native). Either pattern wires the path.""" src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") assert src is not None legacy = "maybe_apply_chat_template" in src @@ -635,14 +530,10 @@ def test_trl_grpo_internal_helpers_in_scope(tag: str): @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_truncate_with_protected_tokens_optional(tag: str): - """Some TRL versions (0.22.2-0.23.1 specifically) ship - `truncate_with_protected_tokens`. Newer versions removed it. - rl_replacements.py:712 has a regex that handles both presence - and absence — but if the symbol is renamed without removal, - we need to know.""" + """Informational: track truncate_with_protected_tokens (shipped TRL + 0.22.2-0.23.1, later removed) so a silent rename doesn't slip past + the rl_replacements.py:712 regex that handles both presence/absence.""" src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") assert src is not None - # No assertion — informational only. We just want to NOT silently - # drift. has_it = "truncate_with_protected_tokens" in src _ = has_it # informational; pass either way. diff --git a/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py index 53f48382ca..ec66d88b0c 100644 --- a/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py +++ b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py @@ -58,8 +58,7 @@ def test_zoo_saving_utils_has_moe_merge_state(): "_record_moe_merge_fallback", ): assert sym in src, f"{sym} missing from saving_utils.py (issue #5410 guard)." - # zoo#647 wraps the fallback guard's message onto a second line; - # allow the regex to span newlines via re.DOTALL. + # zoo#647 wraps the guard message across lines; match via re.DOTALL. assert re.search( r"raise\s+RuntimeError\b.*?MoE", src, re.IGNORECASE | re.DOTALL ), "no `raise RuntimeError(...MoE...)`; post-loop guard weakened." @@ -84,13 +83,9 @@ def test_zoo_saving_utils_has_num_experts_resolver(): def test_zoo_saving_utils_writes_generation_config(): src = _fetch_saving_utils() _skip_until_pr_647_lands(src) - # zoo#647 binds the generation_config attr to a local var - # (`gen_cfg = getattr(model, "generation_config", ...); ... - # gen_cfg.save_pretrained(save_directory)`) so an exact - # `generation_config.save_pretrained(` substring no longer - # matches. Anchor on the conceptual operation: a `generation_config` - # mention plus a `.save_pretrained(` call nearby, which is what - # the canary actually cares about. + # zoo#647 aliases generation_config to a local var, so match a + # `generation_config` mention plus a nearby `.save_pretrained(` call + # rather than the exact `generation_config.save_pretrained(` substring. assert re.search( r"generation_config[\s\S]{0,400}?\.save_pretrained\s*\(", src ), "generation_config.json no longer saved (#5410)." diff --git a/tests/vllm_compat/test_extended_module_imports.py b/tests/vllm_compat/test_extended_module_imports.py index d6d3263827..bf974c9d3f 100644 --- a/tests/vllm_compat/test_extended_module_imports.py +++ b/tests/vllm_compat/test_extended_module_imports.py @@ -1,23 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. """Extended import-smoke + API surface checks for unsloth + unsloth-zoo -modules under the existing CUDA spoof harness. +modules under the CUDA spoof harness. -Where `tests/vllm_compat/test_unsloth_zoo_imports.py` covers the -narrow "must import on a vllm-less runner" claim for 5 modules, -this file walks the FULL set of modules our public surface depends -on. Catches: - - - module-level imports that break on a fresh transformers / peft / - bnb release (the symbol pinned at import time is gone) - - feature flags / gates that flip under the spoof (e.g. _IS_MLX - silently activating on a non-Mac CI box) - - public API surface drift: sorted `dir()` of each FastModel class - is dumped and asserted-stable across runs (a removed kwarg here - is a notebook regression we want to catch) - -CPU-only. Inherits the same _zoo_aggressive_cuda_spoof harness as -test_unsloth_zoo_imports.py. +Walks the full set of modules the public surface depends on (vs the 5 in +test_unsloth_zoo_imports.py), catching import-time symbol drift, spoof- +flipped gates (e.g. _IS_MLX on a non-Mac box), and FastModel API drift. +CPU-only; inherits the _zoo_aggressive_cuda_spoof harness. """ from __future__ import annotations @@ -42,16 +31,13 @@ import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402 _spoof.apply() -# Stub modules the unsloth import path may probe but that aren't -# installed on a CPU-only runner. Mirrors test_unsloth_zoo_imports.py. +# Stub optional deps absent on a CPU-only runner (mirrors test_unsloth_zoo_imports.py). def _stub_module(name: str, attrs: dict | None = None) -> None: - """Stub a missing optional dep. Sets __spec__ so importlib.util's - `find_spec(name)` doesn't raise `ValueError: __spec__ is None`, - which torch / transformers / torchcodec callers hit otherwise.""" + """Stub a missing optional dep, with __spec__ set so find_spec() doesn't + raise `ValueError: __spec__ is None` for torch/transformers/torchcodec callers.""" if name in sys.modules: return m = types.ModuleType(name) - # Minimal viable spec so importlib treats the stub as a real module. m.__spec__ = importlib.machinery.ModuleSpec(name = name, loader = None, origin = "") for k, v in (attrs or {}).items(): setattr(m, k, v) @@ -97,8 +83,7 @@ def _has_unsloth() -> bool: return importlib.util.find_spec("unsloth") is not None -# Extended unsloth-zoo module list. Modules with no top-level vllm/CUDA -# import are expected to load cleanly on a CPU spoof runner. +# unsloth-zoo modules with no top-level vllm/CUDA import: must load cleanly under spoof. _ZOO_VLLM_FREE_MODULES = [ @@ -128,9 +113,8 @@ _ZOO_VLLM_FREE_MODULES = [ @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed") @pytest.mark.parametrize("modname", _ZOO_VLLM_FREE_MODULES) def test_unsloth_zoo_module_imports_under_spoof(modname: str): - """Each unsloth_zoo module must import cleanly on a CPU-only spoof runner. - Catches transformers/peft/bnb symbol drift that fails at import time.""" - # Force fresh resolution: drop stale partial-import state from a prior failure + """Each unsloth_zoo module imports cleanly under spoof (catches import-time symbol drift).""" + # Drop stale partial-import state from a prior failure. sys.modules.pop(modname, None) try: importlib.import_module(modname) @@ -140,14 +124,12 @@ def test_unsloth_zoo_module_imports_under_spoof(modname: str): ) -# Spoof correctness: _IS_MLX must remain False on a non-Mac runner -# and _IS_CUDA / DEVICE_TYPE must reflect the spoofed CUDA layer. +# Spoof correctness: _IS_MLX stays False on a non-Mac runner. @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed") def test_unsloth_is_mlx_false_under_spoof(): - """The CUDA spoof should not flip the MLX flag on a Linux/Windows CI - box (real Apple Silicon is the ONLY environment _IS_MLX activates).""" + """CUDA spoof must not flip _IS_MLX on non-Apple-Silicon hosts.""" sys.modules.pop("unsloth", None) import unsloth @@ -157,8 +139,7 @@ def test_unsloth_is_mlx_false_under_spoof(): ) -# unsloth.models.* — core RL + sentence-transformer surfaces, loaded -# transitively by `from unsloth import FastLanguageModel`. +# unsloth.models.* — core surfaces loaded transitively by `from unsloth import FastLanguageModel`. _UNSLOTH_CORE_MODULES = [ @@ -175,11 +156,8 @@ _UNSLOTH_CORE_MODULES = [ @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed") @pytest.mark.parametrize("modname", _UNSLOTH_CORE_MODULES) def test_unsloth_core_module_imports_under_spoof(modname: str): - """Core unsloth modules must import on a CPU-only runner under the CUDA - spoof; module-top symbol drift crashes here before any user-visible call. - - Bootstraps via `import unsloth` first, since most sub-modules need the - package's _gpu_init side effects (else an import-guard fires).""" + """Core unsloth modules must import under spoof (module-top symbol drift + crashes here). Bootstraps `import unsloth` first for its _gpu_init side effects.""" try: import unsloth # noqa: F401 -- triggers _gpu_init side effects except Exception as e: @@ -188,9 +166,8 @@ def test_unsloth_core_module_imports_under_spoof(modname: str): try: importlib.import_module(modname) except OSError as e: - # `OSError: could not get source code` happens when an editable - # install + frozen sub-import combine; that's an environment - # quirk, not a symbol-drift bug. Skip rather than false-fail. + # "could not get source code": editable-install + frozen sub-import + # quirk, not symbol drift. Skip rather than false-fail. pytest.skip(f"{modname} env issue: {e!s}") except Exception as e: pytest.fail( @@ -198,8 +175,8 @@ def test_unsloth_core_module_imports_under_spoof(modname: str): ) -# Public API surface for FastLanguageModel / FastVisionModel / FastModel under -# spoof: surface must be non-empty and the notebook-relied methods present. +# FastLanguageModel/FastVisionModel/FastModel surface must be non-empty +# with the notebook-relied methods present. @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed") @@ -214,7 +191,7 @@ def test_fast_model_class_surface_under_spoof(): continue found_at_least_one = True public = sorted(n for n in dir(cls) if not n.startswith("_")) - # Notebooks rely on these methods; loss of any one is a regression + # Notebooks rely on these methods. for method in ("from_pretrained", "get_peft_model"): assert method in public, ( f"unsloth.{cls_name}.{method} missing under spoof; " @@ -226,9 +203,8 @@ def test_fast_model_class_surface_under_spoof(): ) -# RL surface drill-down: GRPO/SFT/DPO classes must be reachable and the -# source-rewriter dispatch table populated. Catches rl_replacements importing -# cleanly while RL_FUNCTIONS / RL_REPLACEMENTS is silently empty. +# RL surface: GRPO/SFT/DPO dispatch table must be populated, not silently +# empty while rl_replacements imports cleanly. @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed") @@ -246,7 +222,7 @@ def test_unsloth_rl_replacements_dispatch_populated(): if funcs is None: pytest.skip("RL_FUNCTIONS attribute not present (architecture changed; check)") assert isinstance(funcs, dict), f"RL_FUNCTIONS expected dict, got {type(funcs).__name__}" - # The trainer types unsloth-zoo dispatches against MUST be keys. + # Trainer types unsloth-zoo dispatches against must be keys. for key in ("grpo_trainer", "sft_trainer", "dpo_trainer"): assert key in funcs, ( f"RL_FUNCTIONS missing dispatch key '{key}'; " @@ -257,8 +233,7 @@ def test_unsloth_rl_replacements_dispatch_populated(): ), f"RL_FUNCTIONS[{key!r}] is empty list; rewrites no-op" -# unsloth-zoo compiler test_apply_fused_lm_head — the fused-LM-head emit path -# (named test in compiler.py:1983); we just confirm it's callable. +# unsloth-zoo compiler test_apply_fused_lm_head (compiler.py:1983) must be callable. @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed") @@ -272,8 +247,7 @@ def test_zoo_compiler_apply_fused_lm_head_callable(): ) -# Spot-check signature stability of FastModel.from_pretrained — every notebook -# call site relies on these kwargs; a removal becomes silent positional drift. +# FastModel.from_pretrained kwarg stability: removal becomes silent positional drift. @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed") @@ -291,7 +265,7 @@ def test_fast_model_from_pretrained_kwargs_under_spoof(): params = list(inspect.signature(fn).parameters) except (TypeError, ValueError): pytest.skip("from_pretrained signature not introspectable") - # Notebooks use these by name everywhere. + # Notebooks use these kwargs by name everywhere. for kwarg in ("model_name", "max_seq_length", "load_in_4bit"): assert kwarg in params, ( f"FastLanguageModel.from_pretrained missing kwarg `{kwarg}`; " diff --git a/tests/vllm_compat/test_unsloth_zoo_imports.py b/tests/vllm_compat/test_unsloth_zoo_imports.py index 3f1f5a287d..92e40e4b9f 100644 --- a/tests/vllm_compat/test_unsloth_zoo_imports.py +++ b/tests/vllm_compat/test_unsloth_zoo_imports.py @@ -1,15 +1,11 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -""" -CPU-only smoke imports for the unsloth_zoo modules touching vLLM and -GRPO + fast_inference=True, under the tests/_zoo_aggressive_cuda_spoof -harness. +"""CPU-only smoke imports for unsloth_zoo modules touching vLLM / GRPO + +fast_inference, under the tests/_zoo_aggressive_cuda_spoof harness. -rl_replacements and empty_model are vllm-free by design and MUST import -on CPU with no vllm installed -- this file proves it. The other three -(vllm_utils, vllm_lora_request, vllm_lora_worker_manager) hard-import -vllm and are skipped without it; test_vllm_pinned_symbols.py covers them -statically against pinned vLLM source. +rl_replacements and empty_model are vllm-free and MUST import on CPU with no +vllm; the three vllm-hard-import modules are skipped without it (covered +statically by test_vllm_pinned_symbols.py). Cross-references (unsloth_zoo commits that fixed bugs surfaced here): e3072a23 (WorkerLoRAManager.supports_tower_connector_lora missing), @@ -31,9 +27,7 @@ from pathlib import Path import pytest -# Apply the consolidated CPU spoof at import time, like -# .github/workflows/consolidated-tests-ci.yml shims unsloth before any -# unsloth-touching import. +# Apply the consolidated CPU spoof at import time, before any unsloth import. _SPOOF_DIR = Path(__file__).resolve().parents[1] sys.path.insert(0, str(_SPOOF_DIR)) import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402 @@ -41,8 +35,7 @@ import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402 _spoof.apply() -# Some unsloth_zoo modules read pynvml at import for memory probes; it may -# not be installed on the runner, so stub it. +# Some unsloth_zoo modules read pynvml at import; stub it for the runner. def _stub_module(name: str, attrs: dict | None = None) -> None: if name in sys.modules: return @@ -72,8 +65,7 @@ _stub_module( @pytest.fixture(autouse = True) def _torch_distributed_safe(monkeypatch): - """unsloth_zoo + vllm path occasionally probes torch.distributed. - Make is_available()/is_initialized()/get_world_size() safe defaults.""" + """Give torch.distributed probes safe single-process defaults.""" try: import torch.distributed as dist @@ -93,15 +85,12 @@ def _has_vllm() -> bool: return importlib.util.find_spec("vllm") is not None -# rl_replacements: zero direct vllm imports; must import on a vllm-less -# CPU runner. The GRPO + fast_inference user-facing surface. +# rl_replacements: zero direct vllm imports; the GRPO + fast_inference surface. @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed") def test_rl_replacements_imports_without_vllm(): - """unsloth_zoo.rl_replacements must NOT pull in vllm at import time; - the GRPOConfig/GRPOTrainer surface relies only on plain-Python keyword - flags and never touches vllm on a fast_inference=False run.""" + """unsloth_zoo.rl_replacements must NOT pull in vllm at import time.""" sys.modules.pop("unsloth_zoo.rl_replacements", None) rl = importlib.import_module("unsloth_zoo.rl_replacements") # A transitive vllm import crashes GRPOTrainer construction on Colab. @@ -118,8 +107,7 @@ def test_rl_replacements_imports_without_vllm(): ), "expected at least one GRPO-related export in rl_replacements" -# empty_model: no vllm import either; pure builder for the -# fast_inference=True path, filled from a vLLM internals dict by patch_vllm. +# empty_model: no vllm import; pure builder for the fast_inference=True path. @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed") @@ -129,7 +117,6 @@ def test_empty_model_imports_without_vllm(): assert ( "vllm" not in sys.modules ), "unsloth_zoo.empty_model imported vllm transitively; expected to be vllm-free" - # Public function the GRPO + fast_inference path relies on assert ( hasattr(em, "create_empty_causal_lm") or hasattr(em, "create_empty_model") @@ -137,9 +124,8 @@ def test_empty_model_imports_without_vllm(): ), "expected a create_empty_* helper in empty_model" -# vllm_lora_request / vllm_lora_worker_manager / vllm_utils: hard-import -# vllm; skip if it isn't on the runner. The pinned-symbols test covers -# version compatibility statically. +# vllm_lora_request / vllm_lora_worker_manager / vllm_utils: hard-import vllm, +# so skip without it (pinned-symbols test covers version compat statically). @pytest.mark.skipif( diff --git a/tests/vllm_compat/test_vllm_pinned_symbols.py b/tests/vllm_compat/test_vllm_pinned_symbols.py index 41b9790680..a646733cbb 100644 --- a/tests/vllm_compat/test_vllm_pinned_symbols.py +++ b/tests/vllm_compat/test_vllm_pinned_symbols.py @@ -1,43 +1,11 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. -""" -Pinned-symbol compat check across all vLLM PyPI minor versions ->= 0.9.0. Catches API drift like: +"""Pinned-symbol compat check across vLLM PyPI minors >= 0.9.0 (GitHub raw-fetch, no pip/GPU). - - vLLM PR #30253 split vllm.lora.models -> {vllm.lora.lora_model, - vllm.lora.model_manager} (unsloth-zoo commit ec186187) - - vLLM 0.14 gpu_model_runner adds supports_tower_connector_lora() - and calls it unconditionally on every LoRA VLM - (unsloth-zoo commit e3072a23) - - vLLM 0.15 LoRA manager rename of create_lora_manager kwargs - (unsloth-zoo commit 2a80d543) - - vLLM removal of LoRARequest.embedding_padding_modules / lora_path - -> lora_dir (unsloth-zoo commits 888f79fd, e915bca1) - - vLLM v0 graph capture path removed in 0.11 (commit 65939946) - -Strategy: for each tracked vLLM tag, fetch the relevant source files -straight from github.com/vllm-project/vllm (no pip install, no GPU -required) and assert that every symbol unsloth-zoo's vllm_utils + -vllm_lora_worker_manager + vllm_lora_request expects is present. - -Symbol windows (from the unsloth-zoo upstream survey, 2026-05-07): - - HARD imports (must be present in all versions tested): - vllm.lora.peft_helper.PEFTHelper - vllm.lora.request.LoRARequest - vllm.lora.utils.get_adapter_absolute_path - vllm.config.LoRAConfig (+ VllmConfig from 0.11+) - - SOFT imports (try/except wrappers in unsloth-zoo; either branch OK): - vllm.lora.models.{LoRAModel, create_lora_manager} -- pre #30253 - vllm.lora.lora_model.LoRAModel -- post #30253 - vllm.lora.model_manager.create_lora_manager -- post #30253 - - Behavioural (must exist when the corresponding feature is in scope): - vllm.device_allocator.cumem.{CuMemAllocator, libcudart, ...} - -- only required if UNSLOTH_VLLM_STANDBY=1; on 0.10.x and - 0.14.x the feature is hard-errored anyway, so the absence - of those modules in those versions is fine. +Catches API drift like vLLM PR #30253 (vllm.lora.models split), 0.14 +supports_tower_connector_lora(), 0.15 create_lora_manager rename, the +lora_path -> lora_dir rename, and the 0.11 v0 graph-capture removal. +Asserts every symbol unsloth-zoo's vllm_utils + vllm_lora_* expects is present. """ from __future__ import annotations @@ -51,10 +19,7 @@ import urllib.request import pytest -# Tags that map to the released vLLM minor versions we care about. -# Each tracked tag is the last patch release of that minor (or the -# minor's first stable release if no later patch exists yet). Add new -# rows when vLLM ships a new minor. +# Last patch release of each tracked vLLM minor (or first stable if none yet). VLLM_TAGS = [ "v0.9.0", "v0.9.2", @@ -70,16 +35,13 @@ VLLM_TAGS = [ "v0.18.1", "v0.19.1", "v0.20.1", - # `main` catches symbol drift that hasn't shipped to PyPI yet, - # giving us a few-day lead on a release that would break us. + # `main` catches drift before it ships to PyPI. "main", ] def _fetch_text(repo: str, ref: str, path: str) -> str | None: - """Fetch a file's text from GitHub. Returns None on 404 (the file - is renamed/removed in this version, which is informational, not a - hard failure).""" + """Fetch a file's text from GitHub; None on 404 (renamed/removed, informational).""" url = f"https://raw.githubusercontent.com/{repo}/{ref}/{path}" req = urllib.request.Request(url) token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") @@ -101,9 +63,7 @@ def _has_def( name: str, kind: str = "any", ) -> bool: - """Heuristic AST-equivalent grep for `class Name`, `def name`, - or `Name = ...` at module scope. We avoid a full ast.parse so a - single non-importable line (e.g. type: ignore) doesn't false-fail.""" + """Grep for `class Name`/`def name`/`Name = ...`; avoids ast.parse so one bad line doesn't false-fail.""" if kind in ("any", "class") and re.search(rf"^class\s+{re.escape(name)}\b", src, re.MULTILINE): return True if kind in ("any", "func") and re.search( @@ -115,16 +75,12 @@ def _has_def( return False -# ------------------------------------------------------------------------- # HARD-import symbols: must be present in every tested version. -# ------------------------------------------------------------------------- @pytest.mark.parametrize("tag", VLLM_TAGS) def test_vllm_lora_request_hard_imports(tag: str): - """vllm.lora.request.LoRARequest, vllm.lora.utils.get_adapter_absolute_path, - vllm.lora.peft_helper.PEFTHelper. Hard-imported by unsloth-zoo's - vllm_lora_worker_manager.""" + """LoRARequest, get_adapter_absolute_path, PEFTHelper -- hard-imported by unsloth-zoo's vllm_lora_worker_manager.""" src = _fetch_text("vllm-project/vllm", tag, "vllm/lora/request.py") assert src is not None, f"vllm/lora/request.py missing in {tag}" assert _has_def( @@ -146,8 +102,7 @@ def test_vllm_lora_request_hard_imports(tag: str): @pytest.mark.parametrize("tag", VLLM_TAGS) def test_vllm_config_lora_config(tag: str): - """vllm.config.LoRAConfig. Imported at module top of - unsloth_zoo.vllm_lora_worker_manager (HARD).""" + """vllm.config.LoRAConfig -- hard-imported at module top of unsloth_zoo.vllm_lora_worker_manager.""" candidates = [ "vllm/config/__init__.py", "vllm/config.py", @@ -164,25 +119,19 @@ def test_vllm_config_lora_config(tag: str): assert found, f"vllm.config.LoRAConfig missing in {tag} (checked {candidates})" -# ------------------------------------------------------------------------- # SOFT-import symbols: either old path or new post-#30253 path is fine. -# ------------------------------------------------------------------------- @pytest.mark.parametrize("tag", VLLM_TAGS) def test_vllm_lora_models_either_path(tag: str): - """unsloth-zoo's vllm_lora_worker_manager imports - {LoRAModel, LoRAModelManager, LRUCacheLoRAModelManager, - create_lora_manager} from EITHER vllm.lora.models OR - {vllm.lora.lora_model + vllm.lora.model_manager}. Verify at least - one path resolves every symbol, in every version.""" + """The LoRA model/manager symbols must resolve via EITHER vllm.lora.models OR the post-#30253 split path.""" needed = { "LoRAModel": ("class", None), "LoRAModelManager": ("class", None), "LRUCacheLoRAModelManager": ("class", None), "create_lora_manager": ("func", None), } - # Old path: a single vllm/lora/models.py (or vllm/lora/models/__init__.py). + # Old path: single vllm/lora/models.py (or models/__init__.py). old_candidates = ["vllm/lora/models.py", "vllm/lora/models/__init__.py"] old_src = next( (s for s in (_fetch_text("vllm-project/vllm", tag, p) for p in old_candidates) if s), @@ -213,17 +162,12 @@ def test_vllm_lora_models_either_path(tag: str): ) -# ------------------------------------------------------------------------- -# Optional / version-gated symbols. Don't fail if missing on minors -# unsloth-zoo already gates against; assert presence on minors that -# claim support. -# ------------------------------------------------------------------------- +# Optional / version-gated symbols: assert presence only on minors claiming support. @pytest.mark.parametrize("tag", VLLM_TAGS) def test_vllm_worker_lora_manager_class(tag: str): - """vllm.lora.worker_manager.WorkerLoRAManager. unsloth-zoo subclasses - this; signature inspection drives old_init vs new_init choice.""" + """vllm.lora.worker_manager.WorkerLoRAManager -- unsloth-zoo subclasses it; signature drives old_init vs new_init.""" src = _fetch_text("vllm-project/vllm", tag, "vllm/lora/worker_manager.py") if src is None: # Some vLLM versions split this; check fallback locations. @@ -241,14 +185,7 @@ def test_vllm_worker_lora_manager_class(tag: str): @pytest.mark.parametrize("tag", VLLM_TAGS) def test_lora_request_no_removed_kwargs(tag: str): - """vLLM removed `lora_local_path` -> `lora_path` -> `lora_dir` - progressively. unsloth-zoo's vllm_lora_request must not depend on - the older spelling (else GRPO + fast_inference breaks on the - rename release). - - We assert the LoRARequest constructor accepts EITHER the new name - or both (forward-compat). Specifically: presence of `lora_dir` or - `lora_path` is sufficient; both is the transition state.""" + """vLLM renamed lora_local_path -> lora_path -> lora_dir; assert LoRARequest still accepts lora_dir or lora_path.""" src = _fetch_text("vllm-project/vllm", tag, "vllm/lora/request.py") assert src is not None has_dir = bool(re.search(r"\blora_dir\b", src)) @@ -256,19 +193,12 @@ def test_lora_request_no_removed_kwargs(tag: str): assert has_dir or has_path, f"{tag}: vllm.lora.request has neither lora_dir nor lora_path" -# ------------------------------------------------------------------------- -# UNSLOTH_VLLM_STANDBY hard-error windows. -# unsloth-zoo refuses to enable standby on: -# 0.10.0 <= vllm < 0.11.0 (std::bad_alloc) -# 0.14.0 <= vllm < 0.15.0 (cudaErrorIllegalAddress) -# Make this enforcement testable so a future commit doesn't accidentally -# remove the guard. -# ------------------------------------------------------------------------- +# UNSLOTH_VLLM_STANDBY hard-error windows: unsloth-zoo refuses standby on +# 0.10.0 <= vllm < 0.11.0 (std::bad_alloc) and 0.14.0 <= vllm < 0.15.0 (cudaErrorIllegalAddress). def _vllm_zoo_local_path() -> str | None: - """Return the on-runner path to unsloth_zoo.vllm_utils source if - importable. None otherwise.""" + """Return the on-runner path to unsloth_zoo.vllm_utils source, or None.""" try: import importlib.util spec = importlib.util.find_spec("unsloth_zoo.vllm_utils") @@ -280,9 +210,7 @@ def _vllm_zoo_local_path() -> str | None: def test_unsloth_zoo_standby_guards_present(): - """Sanity: the two hard-error windows exist somewhere in the - unsloth_zoo.vllm_utils source. Catches a future revert that drops - them.""" + """Sanity: the two hard-error windows exist in unsloth_zoo.vllm_utils; catches a revert that drops them.""" path = _vllm_zoo_local_path() if path is None: pytest.skip("unsloth_zoo not installed on runner")