From a6dc10dad2e16283f13ba62c85bbba727597f561 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jun 2026 01:07:09 -0700 Subject: [PATCH] Reduce and tighten comments and docstrings across the test suite (#6429) * Reduce and tighten comments and docstrings in tests Shorten verbose comments and docstrings across the test suite without changing any test logic. Remove narration that restates the next line, collapse long module and test docstrings to a single line, and drop banner separators. Keep regression context (issue and PR references, run ids), skip reasons, mocking and timing rationale, license headers, lint and type directives, and commented-out code. Comments and docstrings only: an AST signature check confirms no code, assertions, or string literals changed, and the suite byte-compiles cleanly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/_zoo_aggressive_cuda_spoof.py | 23 +- tests/conftest.py | 36 +- tests/notebooks/test_validator_fixtures.py | 40 +- ...test_construct_chat_template_validation.py | 26 +- tests/python/test_cross_platform_parity.py | 6 +- tests/python/test_e2e_no_torch_sandbox.py | 143 +---- .../test_fast_language_model_text_only.py | 38 +- ...sentence_transformer_redirect_lifecycle.py | 19 +- tests/python/test_install_python_stack.py | 2 - tests/python/test_no_torch_filtering.py | 47 +- .../test_orpo_processor_text_tokenizer.py | 3 +- .../test_patch_trl_rl_trainers_defensive.py | 10 +- tests/python/test_studio_import_no_torch.py | 27 +- .../test_tokenizers_and_torch_constraint.py | 65 +- .../test_unsloth_run_tool_policy_resolver.py | 11 +- tests/saving/gpt-oss-merge/train_and_merge.py | 4 +- .../test_merge_4bit_validation.py | 7 +- .../test_merge_model_perplexity_llama-3.2.py | 2 +- .../test_merge_model_perplexity_mistral.py | 2 +- .../test_merge_model_perplexity_phi_4.py | 4 +- ...st_merged_model_perplexity_llama-3.1-8b.py | 1 - .../test_merged_model_perplexity_qwen_2.5.py | 4 +- .../test_push_to_hub_merged.py | 11 +- ...t_push_to_hub_merged_sharded_index_file.py | 2 +- .../test_save_merged_grpo_model.py | 54 +- .../saving/non_peft/test_whisper_non_peft.py | 2 +- .../test_preserve_tokenizer_eos_token.py | 2 +- .../test_save_subprocess_utf8_encoding.py | 52 +- tests/saving/test_unsloth_save.py | 14 +- .../saving/text_to_speech_models/test_csm.py | 29 +- .../saving/text_to_speech_models/test_lasa.py | 36 +- .../text_to_speech_models/test_orpheus.py | 2 +- .../text_to_speech_models/test_whisper.py | 3 +- .../test_index_file_sharded_model.py | 2 +- .../vision_models/test_push_to_hub_merged.py | 16 +- ..._merge_qwen2.5vl32B_model_ocr_benchmark.py | 4 +- ...t_save_merge_vision_model_ocr_benchmark.py | 11 +- tests/security/conftest.py | 12 +- tests/security/fixtures/_build.py | 26 +- tests/security/test_lint_workflow_triggers.py | 12 +- .../test_lockfile_supply_chain_audit.py | 53 +- tests/security/test_new_install_scripts.py | 16 +- tests/security/test_scan_npm_packages.py | 67 +- tests/security/test_scan_packages.py | 110 +--- tests/studio/_playwright_robust.py | 135 ++-- tests/studio/install/conftest.py | 8 +- .../smoke_test_parallel_studio_home.py | 44 +- tests/studio/install/test_cuda_repair.py | 42 +- .../install/test_gpu_detection_followups.py | 52 +- tests/studio/install/test_hf_auth.py | 11 +- .../test_install_llama_prebuilt_logic.py | 71 +-- .../install/test_launch_studio_launcher.py | 21 +- .../install/test_llama_pr_force_and_source.py | 56 +- .../install/test_macos_version_compat.py | 24 +- tests/studio/install/test_pr4562_bugfixes.py | 147 ++--- tests/studio/install/test_pr5940_followups.py | 51 +- tests/studio/install/test_probe_timeouts.py | 49 +- tests/studio/install/test_rocm_support.py | 445 ++++--------- tests/studio/install/test_selection_logic.py | 269 +++----- tests/studio/load_freeze/llama_server_shim.py | 25 +- .../load_freeze/test_load_orchestrator.py | 178 ++---- tests/studio/playwright_chat_ime_i18n.py | 142 ++--- tests/studio/playwright_chat_ui.py | 582 +++++------------- tests/studio/playwright_extra_ui.py | 148 ++--- tests/studio/run_real_mlx_smoke.py | 133 ++-- tests/studio/studio_api_smoke.py | 134 +--- tests/studio/test_auth_form_input_count.py | 56 +- tests/studio/test_cancel_atomicity.py | 10 +- tests/studio/test_cancel_id_wiring.py | 32 +- tests/studio/test_cli_repo_variant.py | 18 +- tests/studio/test_cli_run_alias.py | 10 +- tests/studio/test_cli_studio_defaults.py | 19 +- tests/studio/test_cli_studio_stop_windows.py | 13 +- .../test_composer_rtl_bidi_attribute.py | 40 +- tests/studio/test_frontend_dep_removal.py | 72 +-- tests/studio/test_hardware_dispatch_matrix.py | 51 +- tests/studio/test_is_mlx_dispatch_gate.py | 68 +- .../test_stream_cancel_registration_timing.py | 77 +-- .../test_studio_gguf_export_script_pin.py | 15 +- .../test_studio_text_descender_clipping.py | 10 +- tests/test_callback_signature_drift.py | 76 +-- tests/test_cli_export_unpacking.py | 42 +- tests/test_enforce_kwargs_spacing.py | 33 +- tests/test_finetune_last_n_layers.py | 9 +- tests/test_gemma4_chat_template.py | 2 +- tests/test_import_fixes_drift.py | 168 ++--- tests/test_loader_glob_skip.py | 26 +- tests/test_model_registry.py | 10 +- tests/test_multi_image_grpo_chunking.py | 4 +- tests/test_nvfp4_quant_load.py | 12 +- tests/test_public_api_surface.py | 26 +- tests/test_raw_text.py | 67 +- tests/test_studio_install_workspace_guard.py | 227 ++----- tests/test_studio_root_resilience.py | 30 +- tests/test_video_path_validation.py | 10 +- tests/test_windows_rocm_bnb_version.py | 32 +- tests/utils/aime_eval.py | 22 +- tests/utils/cleanup_utils.py | 31 +- tests/utils/data_utils.py | 7 +- tests/utils/generate_dataset_with_none.py | 24 +- tests/utils/hf_utils.py | 7 +- tests/utils/ocr_eval.py | 18 +- tests/utils/os_utils.py | 2 - tests/utils/perplexity_eval.py | 2 +- tests/utils/run_none_detect_tests.py | 98 +-- .../test_batched_leftpad_generation_gpu.py | 18 +- tests/utils/test_packing.py | 6 +- tests/utils/test_prepare_inputs_leftpad.py | 40 +- tests/utils/test_q_galore.py | 58 +- tests/utils/test_qat.py | 32 +- tests/utils/test_rope_scaling_drift.py | 18 +- tests/utils/test_trunc_normal_patch.py | 2 +- tests/version_compat/_fetch.py | 18 +- .../test_bitsandbytes_pinned_symbols.py | 59 +- .../test_peft_pinned_symbols.py | 38 +- ...st_sentence_transformers_pinned_symbols.py | 72 +-- .../test_transformers_pinned_symbols.py | 140 +---- .../test_trl_grpo_pinned_symbols.py | 293 +++------ ..._unsloth_zoo_save_merged_pinned_symbols.py | 13 +- .../test_extended_module_imports.py | 80 +-- tests/vllm_compat/test_unsloth_zoo_imports.py | 40 +- tests/vllm_compat/test_vllm_pinned_symbols.py | 112 +--- 122 files changed, 1837 insertions(+), 4501 deletions(-) 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")