diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index 33f9273b2d..d1cdf6be0f 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -1884,6 +1884,32 @@ class TestVenvDirFileIntegrity: ) assert _venv_dir_is_valid_and_undamaged(str(venv_dir), ("transformers==5.3.0",)) + def test_shared_non_runtime_rows_are_ignored(self, tmp_path: Path): + """A top-level test/ tree is shared, so one wheel's uninstall deletes + another's files. Nothing imports it: the stdlib shadows `test`, and + `tests`/`scripts` ship no __init__.py.""" + venv_dir = self._make_venv( + tmp_path / "venv", + record_extra = ["test/conftest.py,sha256=deadbeef,20650"], + ) + assert _venv_dir_is_valid_and_undamaged(str(venv_dir), ("transformers==5.3.0",)) + + def test_an_installer_rewritten_file_keeps_its_existence_check(self, tmp_path: Path): + """setup.ps1 runs npm install in the installed tree, so the lockfile's + recorded size drifts. npm never deletes it, so only the size is dropped.""" + rel = "studio/backend/core/data_recipe/oxc-validator/package-lock.json" + venv_dir = self._make_venv( + tmp_path / "venv", + record_extra = [f"{rel},sha256=deadbeef,28473"], + ) + lock = venv_dir / rel + lock.parent.mkdir(parents = True, exist_ok = True) + lock.write_text("x" * 27225) # shrunk by npm, not damage + assert _venv_dir_is_valid_and_undamaged(str(venv_dir), ("transformers==5.3.0",)) + + lock.unlink() # gone entirely, which npm never does + assert not _venv_dir_is_valid_and_undamaged(str(venv_dir), ("transformers==5.3.0",)) + def test_in_target_script_rows_are_ignored(self, tmp_path: Path): """uv records bin/hf, which does resolve -- but pip --upgrade rmtree's a colliding directory in the target, so a later install into the same diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index a8cf01170e..e44e4cc948 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -1982,6 +1982,27 @@ def _sidecar_scan(venv_dir: str, limit: int = 3) -> tuple[list[str], bool]: return _sidecar_scan_impl(venv_dir, limit) +# Mirrored from unsloth_cli/_studio_deps.py, not imported, for the reason given in +# _sidecar_scan_impl below: the backend never imports the CLI package. Keep in sync. +_SHARED_NON_RUNTIME_ROOTS = frozenset( + ( + "test", + "tests", + "doc", + "docs", + "example", + "examples", + "benchmark", + "benchmarks", + "sample", + "samples", + "scripts", + ) +) +_INSTALLER_REWRITTEN_NAMES = frozenset(("package-lock.json",)) +_OUR_DISTRIBUTIONS = frozenset(("unsloth", "unsloth-zoo", "unsloth-studio")) + + def _sidecar_damaged_files(venv_dir: str, limit: int = 3) -> list[str]: """RECORD entries under *venv_dir* that are gone, or shorter than pip recorded. @@ -2075,10 +2096,21 @@ def _sidecar_scan_impl(venv_dir: str, limit: int = 3) -> tuple[list[str], bool]: or (parts and parts[0] in ("bin", "Scripts")) ): continue + # Third-party top-level dirs several wheels write into, so one uninstall + # deletes another's files. Never applied to what we ship: a missing + # __init__.py does not make a directory unimportable (PEP 420). + if ( + len(parts) > 1 + and parts[0] in _SHARED_NON_RUNTIME_ROOTS + and name.replace("_", "-").lower() not in _OUR_DISTRIBUTIONS + ): + continue # The size field is optional and real wheels do leave it blank. Keep the row with an # unknown size: existence is still checkable, and dropping it hides a deletion. recorded: int | None = None - if len(row) >= 3 and row[2]: + # An installer rewrites these in place, so the recorded size drifts; + # the file disappearing is still damage. + if len(row) >= 3 and row[2] and parts[-1] not in _INSTALLER_REWRITTEN_NAMES: try: recorded = int(row[2]) except ValueError: diff --git a/tests/python/test_windows_studio_update_launcher.py b/tests/python/test_windows_studio_update_launcher.py new file mode 100644 index 0000000000..c72f7c859e --- /dev/null +++ b/tests/python/test_windows_studio_update_launcher.py @@ -0,0 +1,536 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Focused regression tests for the Windows Studio updater launcher transaction.""" + +from __future__ import annotations + +import importlib.util +import subprocess +import sys +import types +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_COMMAND = REPO_ROOT / "unsloth_cli" / "commands" / "studio.py" +ORIGINAL_LAUNCHER = b"MZ-original-launcher" +REAL_MSVCRT = sys.modules.get("msvcrt") + + +@pytest.fixture +def studio(monkeypatch): + """Load studio.py without importing the heavyweight unsloth package.""" + package = types.ModuleType("unsloth_cli") + package.__path__ = [str(REPO_ROOT / "unsloth_cli")] + commands = types.ModuleType("unsloth_cli.commands") + commands.__path__ = [str(REPO_ROOT / "unsloth_cli" / "commands")] + deps = types.ModuleType("unsloth_cli._studio_deps") + inference = types.ModuleType("unsloth_cli._inference") + inference.SpeculativeType = str + password_prompt = types.ModuleType("unsloth_cli.commands._password_prompt") + commands._password_prompt = password_prompt + + for name, module in ( + ("unsloth_cli", package), + ("unsloth_cli.commands", commands), + ("unsloth_cli._studio_deps", deps), + ("unsloth_cli._inference", inference), + ("unsloth_cli.commands._password_prompt", password_prompt), + ): + monkeypatch.setitem(sys.modules, name, module) + + module_name = "unsloth_cli.commands.studio_launcher_transaction_test" + spec = importlib.util.spec_from_file_location(module_name, STUDIO_COMMAND) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _configure_windows( + monkeypatch, + studio, + tmp_path, + *, + launcher = ORIGINAL_LAUNCHER, +): + scripts = tmp_path / "Scripts" + scripts.mkdir() + python = scripts / "python.exe" + python.write_bytes(b"python") + launcher_path = scripts / "unsloth.exe" + if launcher is not None: + launcher_path.write_bytes(launcher) + + lock_state = {"locked": False} + fake_msvcrt = types.ModuleType("msvcrt") + fake_msvcrt.LK_NBLCK = 1 + fake_msvcrt.LK_UNLCK = 2 + + def locking(_fileno, mode, _length): + if mode == fake_msvcrt.LK_NBLCK: + if lock_state["locked"]: + raise OSError("lock conflict") + lock_state["locked"] = True + else: + lock_state["locked"] = False + + fake_msvcrt.locking = locking + monkeypatch.setitem(sys.modules, "msvcrt", fake_msvcrt) + monkeypatch.setattr(studio.platform, "system", lambda: "Windows") + monkeypatch.setattr(studio.sys, "executable", str(python)) + monkeypatch.setattr(studio, "_ensure_studio_env_exported", lambda: None) + monkeypatch.setattr(studio, "_windows_hidden_subprocess_kwargs", lambda: {}) + monkeypatch.setattr(studio, "_refresh_desktop_shortcuts", lambda **_kwargs: None) + monkeypatch.setattr(studio, "_fail_if_install_damaged", lambda: None) + monkeypatch.setattr(studio, "STUDIO_HOME", tmp_path / "studio_home") + for name in ( + "SKIP_STUDIO_BASE", + "STUDIO_PACKAGE_NAME", + "STUDIO_LOCAL_INSTALL", + "STUDIO_LOCAL_REPO", + "UNSLOTH_TAURI_UPDATE", + ): + monkeypatch.delenv(name, raising = False) + return scripts, launcher_path + + +def _successful_version_run(calls = None): + def run(argv, **kwargs): + if calls is not None: + calls.append((argv, kwargs)) + return types.SimpleNamespace(returncode = 0) + + return run + + +def _update(studio, *, verify = True): + studio.update(local = False, package = "unsloth", verbose = False, verify = verify) + + +def test_setup_noop_preserves_launcher_and_removes_backup(monkeypatch, studio, tmp_path): + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + calls = [] + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run(calls)) + + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + assert not (scripts / "unsloth.exe.update-backup").exists() + assert calls[0][0] == [str(launcher), "--version"] + assert calls[0][1]["timeout"] == 10 + + +def test_a_recoverable_copy_exists_while_setup_runs(monkeypatch, studio, tmp_path): + # The canonical path is freed so the installer can publish a replacement, + # but never without a copy to put back if it publishes nothing. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + + def setup(**_kwargs): + assert not launcher.exists() + assert (scripts / "unsloth.exe.update-backup").read_bytes() == ORIGINAL_LAUNCHER + assert (scripts / "unsloth.exe.update-stale").read_bytes() == ORIGINAL_LAUNCHER + + monkeypatch.setattr(studio, "_run_setup_script", setup) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + +def test_setup_failure_restores_original_and_propagates(monkeypatch, studio, tmp_path): + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + + def setup(**_kwargs): + launcher.write_bytes(b"MZ-new-but-incomplete") + raise RuntimeError("setup failed") + + monkeypatch.setattr(studio, "_run_setup_script", setup) + + with pytest.raises(RuntimeError, match = "setup failed"): + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + assert (scripts / "unsloth.exe.update-backup").read_bytes() == ORIGINAL_LAUNCHER + + +def test_setup_publishing_no_launcher_restores_it_and_succeeds(monkeypatch, studio, tmp_path): + # The bug this transaction exists for: pip finds unsloth already current, + # writes no launcher, and the old updater then deleted its own .deleteme and + # left the venv with none at all. Restoring is the right answer, not failing. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + assert not (scripts / "unsloth.exe.update-stale").exists() + + +@pytest.mark.parametrize("invalid", [b"", b"not-a-pe"]) +def test_invalid_launcher_is_restored_and_update_fails(monkeypatch, studio, tmp_path, invalid): + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + monkeypatch.setattr( + studio, "_run_setup_script", lambda **_kwargs: launcher.write_bytes(invalid) + ) + + with pytest.raises(studio.typer.Exit): + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + assert (scripts / "unsloth.exe.update-backup").exists() + + +@pytest.mark.parametrize("outcome", ["nonzero", "timeout"]) +def test_runtime_check_failure_restores_launcher(monkeypatch, studio, tmp_path, outcome): + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + + def run(argv, **kwargs): + if outcome == "timeout": + raise subprocess.TimeoutExpired(argv, kwargs["timeout"]) + return types.SimpleNamespace(returncode = 7) + + monkeypatch.setattr(studio.subprocess, "run", run) + + with pytest.raises(studio.typer.Exit): + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + assert (scripts / "unsloth.exe.update-backup").exists() + + +def test_no_verify_still_checks_launcher_but_skips_integrity_scan(monkeypatch, studio, tmp_path): + _scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + calls = [] + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run(calls)) + integrity_calls = [] + monkeypatch.setattr(studio, "_fail_if_install_damaged", lambda: integrity_calls.append(True)) + + _update(studio, verify = False) + + assert calls[0][0] == [str(launcher), "--version"] + assert integrity_calls == [] + + +def test_legacy_backup_recovers_only_when_launcher_is_missing(monkeypatch, studio, tmp_path): + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path, launcher = None) + legacy = scripts / "unsloth.exe.deleteme" + legacy.write_bytes(ORIGINAL_LAUNCHER) + + def setup(**_kwargs): + # Recovered from the legacy file, then moved aside for the installer. + assert (scripts / "unsloth.exe.update-stale").read_bytes() == ORIGINAL_LAUNCHER + + monkeypatch.setattr(studio, "_run_setup_script", setup) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + assert not legacy.exists() + assert not (scripts / "unsloth.exe.update-backup").exists() + + +def test_lock_contention_exits_before_setup_or_launcher_mutation(monkeypatch, studio, tmp_path): + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + if REAL_MSVCRT is not None: + # Exercise the real byte-range lock on Windows and the fake elsewhere. + monkeypatch.setitem(sys.modules, "msvcrt", REAL_MSVCRT) + before = launcher.read_bytes() + first = studio._WindowsLauncherUpdateTransaction() + first.__enter__() + stale_before = (scripts / "unsloth.exe.update-stale").read_bytes() + backup_before = (scripts / "unsloth.exe.update-backup").read_bytes() + setup_calls = [] + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: setup_calls.append(True)) + try: + with pytest.raises(studio.typer.Exit) as exc: + _update(studio) + assert exc.value.exit_code == 1 + assert setup_calls == [] + assert stale_before == before + assert (scripts / "unsloth.exe.update-stale").read_bytes() == stale_before + assert (scripts / "unsloth.exe.update-backup").read_bytes() == backup_before + finally: + first.__exit__(None, None, None) + + +def test_non_windows_preserves_call_order_without_launcher_operations( + monkeypatch, studio, tmp_path +): + monkeypatch.setattr(studio.platform, "system", lambda: "Linux") + monkeypatch.setattr(studio.sys, "executable", str(tmp_path / "bin" / "python")) + monkeypatch.setattr(studio, "_ensure_studio_env_exported", lambda: None) + calls = [] + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: calls.append("setup")) + monkeypatch.setattr(studio, "_fail_if_install_damaged", lambda: calls.append("verify")) + monkeypatch.setattr( + studio, "_refresh_desktop_shortcuts", lambda **_kwargs: calls.append("refresh") + ) + for name in ( + "SKIP_STUDIO_BASE", + "STUDIO_PACKAGE_NAME", + "STUDIO_LOCAL_INSTALL", + "STUDIO_LOCAL_REPO", + "UNSLOTH_TAURI_UPDATE", + ): + monkeypatch.delenv(name, raising = False) + + _update(studio) + + assert calls == ["setup", "verify", "refresh"] + assert list(tmp_path.rglob("unsloth.exe*")) == [] + + +def _shim(studio, payload = ORIGINAL_LAUNCHER): + """The hardlinked PATH shim install.ps1 creates beside the managed venv.""" + path = studio.STUDIO_HOME / "bin" / "unsloth.exe" + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(payload) + return path + + +def test_a_missing_launcher_is_recovered_from_the_path_shim(monkeypatch, studio, tmp_path): + # The old updater renamed the launcher away and then unlinked the .deleteme, + # so an affected install has neither. install.ps1 hardlinks the shim to the + # same file, so it survives that unlink and can repair the launcher. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path, launcher = None) + _shim(studio) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_an_invalid_launcher_is_recovered_from_the_backup(monkeypatch, studio, tmp_path): + # Recovery gated on existence rather than validity left a zero-byte launcher + # in place while a usable backup sat beside it. The old updater restored on + # exactly this shape (st_size == 0), so gating on exists() regressed it. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path, launcher = b"") + (scripts / "unsloth.exe.update-backup").write_bytes(ORIGINAL_LAUNCHER) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_no_launcher_and_no_recovery_source_still_runs_setup(monkeypatch, studio, tmp_path): + # Refusing here would strand exactly the users the transaction exists for: + # the previous updater could leave no launcher and no .deleteme, and before + # this the update simply carried on and let setup reinstall it. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path, launcher = None) + ran = [] + + def setup(**_kwargs): + ran.append(True) + launcher.write_bytes(ORIGINAL_LAUNCHER) + + monkeypatch.setattr(studio, "_run_setup_script", setup) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert ran == [True] + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_a_backup_failure_does_not_abort_the_update(monkeypatch, studio, tmp_path): + # A backup is a safety net, not a precondition. Antivirus holding the temp + # copy used to surface as a bare OSError traceback before setup ever ran. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + ran = [] + original = studio._WindowsLauncherUpdateTransaction._atomic_copy + + def refuse_backup(source, destination): + if destination.name.endswith(".update-backup"): + raise OSError("access is denied") + return original(source, destination) + + # _atomic_copy is a staticmethod, so the stand-in must not bind self either. + monkeypatch.setattr( + studio._WindowsLauncherUpdateTransaction, "_atomic_copy", staticmethod(refuse_backup) + ) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: ran.append(True)) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert ran == [True] + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_an_existing_backup_survives_an_unvalidated_launcher(monkeypatch, studio, tmp_path): + # A backup outlives __enter__ only when a previous run died before + # validating, so it holds the last launcher known to run while the canonical + # file has passed nothing but the two-byte header check. Overwriting it here + # destroyed the only recovery copy. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path, launcher = b"MZ-broken") + backup = scripts / "unsloth.exe.update-backup" + backup.write_bytes(ORIGINAL_LAUNCHER) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + + def failing_version(_argv, **_kwargs): + return types.SimpleNamespace(returncode = 7) + + monkeypatch.setattr(studio.subprocess, "run", failing_version) + + with pytest.raises(studio.typer.Exit): + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_the_launcher_is_resolved_from_the_managed_studio_venv(monkeypatch, studio, tmp_path): + # A pip-installed or checkout CLI drives an update of the separate managed + # environment, so sys.executable belongs to the caller while setup.ps1 + # installs into STUDIO_HOME/unsloth_studio. Guarding the caller's launcher + # left the one actually being replaced unprotected. + scripts, caller_launcher = _configure_windows(monkeypatch, studio, tmp_path) + managed = tmp_path / "studio_home" / "unsloth_studio" + (managed / "Scripts").mkdir(parents = True) + (managed / "pyvenv.cfg").write_text("home = /usr\n") + managed_launcher = managed / "Scripts" / "unsloth.exe" + managed_launcher.write_bytes(b"MZ-managed-launcher") + + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + calls = [] + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run(calls)) + + _update(studio) + + assert calls[0][0] == [str(managed_launcher), "--version"] + assert managed_launcher.read_bytes() == b"MZ-managed-launcher" + + +def test_a_replacement_published_by_setup_is_kept(monkeypatch, studio, tmp_path): + # The point of freeing the canonical path. uv only self-replaces its own + # executable, so it deletes a third-party console script outright and + # hard-errors when the file is in use; the pip fallback then no-ops on the + # already-satisfied bare unsloth and the upgrade is silently skipped. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + new_launcher = b"MZ-upgraded-launcher" + + monkeypatch.setattr( + studio, "_run_setup_script", lambda **_kwargs: launcher.write_bytes(new_launcher) + ) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert launcher.read_bytes() == new_launcher + assert not (scripts / "unsloth.exe.update-stale").exists() + assert not (scripts / "unsloth.exe.update-backup").exists() + + +def test_an_invalid_replacement_is_restored_but_still_fails(monkeypatch, studio, tmp_path): + # Setup writing an unusable launcher is a real failure. Putting the previous + # one back must not turn it into a reported success, which is how a + # restore-then-revalidate reads when it cannot tell "published nothing" + # from "published something broken". + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: launcher.write_bytes(b"")) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + with pytest.raises(studio.typer.Exit) as exc: + _update(studio, verify = False) + + assert exc.value.exit_code == 1 + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_a_backup_that_cannot_run_falls_back_to_the_moved_aside_copy(monkeypatch, studio, tmp_path): + # Backups are taken after only the two-byte header check, so an interrupted + # run can leave a PE-shaped but non-runnable one. Preferring it must not + # strand the working launcher that this run moved aside. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + bad_backup = b"MZ-unrunnable" + (scripts / "unsloth.exe.update-backup").write_bytes(bad_backup) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + + def run(argv, **_kwargs): + current = Path(argv[0]).read_bytes() + return types.SimpleNamespace(returncode = 7 if current == bad_backup else 0) + + monkeypatch.setattr(studio.subprocess, "run", run) + + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_a_setup_exception_restores_a_runnable_launcher(monkeypatch, studio, tmp_path): + # __exit__ took the first PE-shaped candidate, so an interrupted run's + # non-runnable backup was installed over the working launcher this run had + # moved aside, and it could also undo a restore validate_launcher just made. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + bad_backup = b"MZ-unrunnable" + (scripts / "unsloth.exe.update-backup").write_bytes(bad_backup) + + def setup(**_kwargs): + raise RuntimeError("setup failed") + + monkeypatch.setattr(studio, "_run_setup_script", setup) + + def run(argv, **_kwargs): + current = Path(argv[0]).read_bytes() + return types.SimpleNamespace(returncode = 7 if current == bad_backup else 0) + + monkeypatch.setattr(studio.subprocess, "run", run) + + with pytest.raises(RuntimeError, match = "setup failed"): + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_a_non_runnable_backup_falls_through_to_the_legacy_copy(monkeypatch, studio, tmp_path): + # An interrupted run can leave a PE-shaped but non-runnable backup while the + # legacy .deleteme or the PATH shim is still good. Accepting the backup on + # its MZ header alone and stopping there left the update failing forever + # with the broken bytes canonical. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path, launcher = None) + bad_backup = b"MZ-unrunnable" + (scripts / "unsloth.exe.update-backup").write_bytes(bad_backup) + (scripts / "unsloth.exe.deleteme").write_bytes(ORIGINAL_LAUNCHER) + monkeypatch.setattr(studio, "_run_setup_script", lambda **_kwargs: None) + + def run(argv, **_kwargs): + current = Path(argv[0]).read_bytes() + return types.SimpleNamespace(returncode = 7 if current == bad_backup else 0) + + monkeypatch.setattr(studio.subprocess, "run", run) + + _update(studio) + + assert launcher.read_bytes() == ORIGINAL_LAUNCHER + + +def test_the_update_lock_lives_outside_the_replaceable_venv(monkeypatch, studio, tmp_path): + # setup.ps1 removes the whole $VenvDir to rebuild a stale torch, and Windows + # refuses a recursive delete while a handle inside it is open. A lock under + # Scripts/ therefore broke exactly the repair path it was meant to guard. + scripts, launcher = _configure_windows(monkeypatch, studio, tmp_path) + seen = {} + + def setup(**_kwargs): + seen["venv_locks"] = list(scripts.glob("*.update-lock")) + seen["home_locks"] = list((studio.STUDIO_HOME).glob("*.update-lock")) + + monkeypatch.setattr(studio, "_run_setup_script", setup) + monkeypatch.setattr(studio.subprocess, "run", _successful_version_run()) + + _update(studio) + + assert seen["venv_locks"] == [] + assert len(seen["home_locks"]) == 1 diff --git a/unsloth_cli/_studio_deps.py b/unsloth_cli/_studio_deps.py index 2a52793729..cf36c8c657 100644 --- a/unsloth_cli/_studio_deps.py +++ b/unsloth_cli/_studio_deps.py @@ -217,6 +217,63 @@ def _scan_paths() -> Dict[str, list]: return {"path": paths} if paths else {} +# Top-level dirs several wheels write into, so one uninstall deletes another's +# files and the survivor's RECORD describes a file nothing recreates. einx and +# torchao both ship test/conftest.py, and install_python_stack.py +# force-reinstalls torchao every update. +_SHARED_NON_RUNTIME_ROOTS = frozenset( + ( + "test", + "tests", + "doc", + "docs", + "example", + "examples", + "benchmark", + "benchmarks", + "sample", + "samples", + "scripts", + ) +) + +# Rewritten in place by our own setup: setup.ps1/setup.sh run `npm install` +# inside the installed tree, and npm dedupes hoisted entries under +# legacy-peer-deps, shrinking the lockfile below its recorded size. +_INSTALLER_REWRITTEN_NAMES = frozenset(("package-lock.json",)) + + +def _shared_non_runtime(rel: str, name: str) -> bool: + """A third-party row under a top-level dir several wheels write into. + + Ownership of these is unreliable: whichever wheel installed last wins, and + any of them uninstalling takes the others' files with it. Never applied to + what we ship, because a missing __init__.py does NOT make a directory + unimportable (PEP 420) and this repo imports scripts.* itself, so our own + top-level trees stay checked. Applied while reading RECORD, not when + reporting, so the row also stays out of the ownership tally and the `limit` + budget. + """ + if _canonical(name) in _OUR_DISTRIBUTIONS: + return False + parts = tuple(p for p in rel.replace("\\", "/").split("/") if p and p != ".") + return len(parts) > 1 and parts[0] in _SHARED_NON_RUNTIME_ROOTS + + +# What Unsloth ships. Its top-level trees are ours to guarantee, so they are +# never exempted however they are named. +_OUR_DISTRIBUTIONS = frozenset(("unsloth", "unsloth-zoo", "unsloth-studio")) + + +def _installer_rewritten(rel: str) -> bool: + """A file our own setup rewrites in place, so its recorded size drifts. + + Only the size is unreliable. The file disappearing is still damage, so the + row is kept and only its size dropped. + """ + return rel.replace("\\", "/").rsplit("/", 1)[-1] in _INSTALLER_REWRITTEN_NAMES + + def damaged_installed_files(limit: int = 8) -> List[str]: """Installed files that are gone, or shorter than pip recorded. @@ -236,6 +293,11 @@ def damaged_installed_files(limit: int = 8) -> List[str]: in EITHER direction. Sizes are therefore compared after the whole scan, once multiply-owned paths are known, rather than during it. + Rows that cannot be import-time damage are dropped up front, and a file our + own setup rewrites keeps its existence check but loses its size. The answer + to a finding is "reinstall over the top", so a file no reinstall would + change must not produce one. See _shared_non_runtime, _installer_rewritten. + Scanned over the interpreter's own site-packages rather than all of sys.path. distributions() searches every sys.path entry, so a damaged distribution reachable only through an inherited PYTHONPATH would otherwise @@ -273,11 +335,13 @@ def damaged_installed_files(limit: int = 8) -> List[str]: # size recorded inside itself; .pyc is regenerated from source. if ".dist-info/" in rel or ".egg-info/" in rel or rel.endswith(".pyc"): continue + if _shared_non_runtime(rel, name): + continue # The size field is optional and real wheels do leave it blank. Keep # the row anyway with an unknown size: existence is still checkable, # and dropping the row meant a deletion went unreported. recorded: Optional[int] = None - if len(row) >= 3 and row[2]: + if len(row) >= 3 and row[2] and not _installer_rewritten(rel): try: recorded = int(row[2]) except ValueError: diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 8e76c843c9..85aac4ff1d 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -3209,24 +3209,20 @@ def update( else: os.environ["STUDIO_LOCAL_INSTALL"] = "0" os.environ.pop("STUDIO_LOCAL_REPO", None) + # main gained a runtime gate around setup; this branch replaced the + # rename-to-.deleteme helpers with the launcher transaction. Both apply: + # the gate keeps a second Studio process off the venv, the transaction + # keeps the launcher recoverable across the setup it wraps. runtime_gate_handoff = _studio_runtime_gate.consume_runtime_gate_handoff() with _studio_runtime_launch_guard(inherited = runtime_gate_handoff): _studio_runtime_gate.ensure_managed_environment_is_idle(STUDIO_HOME) - _release_self_exe_lock_windows() - try: + with _WindowsLauncherUpdateTransaction() as launcher_update: _run_setup_script(verbose = verbose, repo_root = repo_root) - except BaseException: - # Restore unsloth.exe from .deleteme if setup failed before pip - # produced a replacement; otherwise the user has no CLI for recovery. - _restore_self_exe_lock_windows() - raise - # On Windows clear the .deleteme orphan now that pip wrote a fresh - # unsloth.exe; on next update os.replace would overwrite it anyway, - # but leaving a stale binary around invites cross-version restore - # confusion from _restore_self_exe_lock_windows. - _cleanup_self_exe_lock_windows() - if verify: - _fail_if_install_damaged() + # This deliberately runs even with --no-verify: the broad package scan + # is optional, but a successful update must leave its own launcher usable. + launcher_update.validate_launcher() + if verify: + _fail_if_install_damaged() # Tauri desktop owns its own bundle entries; skip CLI launcher refresh # so a Tauri-initiated update doesn't create duplicate shortcuts. if os.environ.get("UNSLOTH_TAURI_UPDATE") == "1": @@ -3236,66 +3232,346 @@ def update( _refresh_desktop_shortcuts(verbose = verbose) -def _release_self_exe_lock_windows() -> None: - """Rename running unsloth.exe so pip can replace it. setup.ps1 also retries.""" - if platform.system() != "Windows": - return - try: - venv_scripts = Path(sys.executable).resolve().parent - except OSError: - return - exe = venv_scripts / "unsloth.exe" - if not exe.exists(): - return - stale = exe.with_suffix(".exe.deleteme") - try: - # os.replace is atomic-overwrite on Windows; os.rename would raise - # FileExistsError if a prior aborted update left a .deleteme behind. - os.replace(exe, stale) - except OSError as e: - # Not fatal; setup.ps1 retries from a sibling process. - print(f"[update] could not rename {exe.name} -> {stale.name}: {e}") +class _WindowsLauncherUpdateTransaction: + """Keep the managed Windows launcher recoverable during a Python update.""" + _VERSION_TIMEOUT_SECONDS = 10 + _RESTORE_ATTEMPTS = 3 -def _restore_self_exe_lock_windows() -> None: - """If setup failed before pip wrote a working unsloth.exe, restore .deleteme.""" - if platform.system() != "Windows": - return - try: - venv_scripts = Path(sys.executable).resolve().parent - except OSError: - return - exe = venv_scripts / "unsloth.exe" - stale = exe.with_suffix(".exe.deleteme") - if not stale.exists(): - return - # Treat a missing or zero-byte exe as "pip didn't produce a usable - # replacement"; otherwise leave the new binary alone. - if exe.exists(): + def __init__(self) -> None: + self.enabled = platform.system() == "Windows" + self.launcher: Optional[Path] = None + self.backup: Optional[Path] = None + self.legacy_backup: Optional[Path] = None + self.stale: Optional[Path] = None + self.shim: Optional[Path] = None + self.lock_path: Optional[Path] = None + self._lock_file = None + self._validated = False + + @staticmethod + def _is_valid_pe(path: Path) -> bool: try: - if exe.stat().st_size > 0: - return + if not path.is_file() or path.stat().st_size < 2: + return False + with path.open("rb") as handle: + return handle.read(2) == b"MZ" except OSError: + return False + + @staticmethod + def _atomic_copy(source: Path, destination: Path) -> None: + """Publish a sibling copy without exposing a partial destination.""" + fd, temporary_name = tempfile.mkstemp( + prefix = f".{destination.name}.", + suffix = ".tmp", + dir = str(destination.parent), + ) + temporary = Path(temporary_name) + try: + with source.open("rb") as source_handle, os.fdopen(fd, "wb") as target_handle: + fd = -1 + while True: + chunk = source_handle.read(1024 * 1024) + if not chunk: + break + target_handle.write(chunk) + target_handle.flush() + os.fsync(target_handle.fileno()) + os.replace(temporary, destination) + finally: + if fd >= 0: + os.close(fd) + try: + temporary.unlink(missing_ok = True) + except OSError: + pass + + def _acquire_lock(self) -> None: + import msvcrt + + assert self.lock_path is not None + try: + self.lock_path.parent.mkdir(parents = True, exist_ok = True) + except OSError: + pass + lock_file = self.lock_path.open("a+b") + try: + lock_file.seek(0, os.SEEK_END) + if lock_file.tell() == 0: + lock_file.write(b"\0") + lock_file.flush() + lock_file.seek(0) + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + except OSError: + lock_file.close() + typer.echo( + "Error: another Unsloth Studio update is already running for this environment.", + err = True, + ) + raise typer.Exit(1) + self._lock_file = lock_file + + def _release_lock(self) -> None: + if self._lock_file is None: return - try: - os.replace(stale, exe) - except OSError as e: - print(f"[update] could not restore {stale.name} -> {exe.name}: {e}") + try: + import msvcrt + self._lock_file.seek(0) + msvcrt.locking(self._lock_file.fileno(), msvcrt.LK_UNLCK, 1) + except OSError: + pass + finally: + self._lock_file.close() + self._lock_file = None + def _recover_missing_launcher(self) -> None: + assert self.launcher is not None + # Validity, not existence: a truncated or quarantined launcher is just as + # unusable, and the backup beside it can repair either. + if self._is_valid_pe(self.launcher): + return + for recovery in (self.backup, self.stale, self.legacy_backup, self.shim): + if recovery is not None and self._is_valid_pe(recovery): + try: + self._atomic_copy(recovery, self.launcher) + except OSError as exc: + typer.echo( + f"Error: could not recover {self.launcher} from {recovery}: {exc}", + err = True, + ) + typer.echo(f"Manual recovery copy retained at: {recovery}", err = True) + raise typer.Exit(1) + return -def _cleanup_self_exe_lock_windows() -> None: - """Remove the .deleteme orphan after a successful update on Windows.""" - if platform.system() != "Windows": - return - try: - venv_scripts = Path(sys.executable).resolve().parent - except OSError: - return - stale = (venv_scripts / "unsloth.exe").with_suffix(".exe.deleteme") - try: - stale.unlink(missing_ok = True) - except OSError: - pass + @staticmethod + def _files_match(left: Path, right: Path) -> bool: + try: + if left.stat().st_size != right.stat().st_size: + return False + with left.open("rb") as left_handle, right.open("rb") as right_handle: + while True: + left_chunk = left_handle.read(1024 * 1024) + if left_chunk != right_handle.read(1024 * 1024): + return False + if not left_chunk: + return True + except OSError: + return False + + def _move_launcher_aside(self) -> None: + """Free the canonical path so the installer can publish a replacement. + + uv only self-replaces its OWN executable, so it deletes a third-party + console script outright and hard-errors when the file is in use, after + which the pip fallback no-ops on the already-satisfied bare unsloth and + the upgrade is silently skipped. Renaming a running image is allowed on + Windows: verified on windows-latest that renaming a live console-script + launcher succeeds and a replacement can then be written at the freed + path. Non-fatal, since failing to move it aside costs only the upgrade. + """ + assert self.launcher is not None and self.stale is not None + if not self._is_valid_pe(self.launcher): + return + try: + os.replace(self.launcher, self.stale) + except OSError as exc: + typer.echo(f"Warning: could not move the Unsloth launcher aside: {exc}", err = True) + + def _retained_backup(self) -> Optional[Path]: + """The backup, when it exists and is usable. Nothing to point users at otherwise.""" + if self.backup is not None and self._is_valid_pe(self.backup): + return self.backup + return None + + def _recovery_candidates(self) -> List[Path]: + """Copies that could stand in for the launcher, best first. + + The backup is the last launcher known to run; the moved-aside copy is + only this run's unvalidated canonical file; the legacy .deleteme and the + PATH shim are what an install broken by the old updater still has. All + of them are kept, because passing the two-byte header check does not + make any one of them runnable and the next candidate has to be reachable. + """ + seen: List[Path] = [] + for path in (self.backup, self.stale, self.legacy_backup, self.shim): + if path is None or not self._is_valid_pe(path): + continue + if not any(os.path.normcase(str(path)) == os.path.normcase(str(p)) for p in seen): + seen.append(path) + return seen + + def _restore_from(self, source: Path) -> bool: + assert self.launcher is not None + # The common setup-failure case leaves the original executable exactly + # where it was. Avoid replacing that running file on Windows when it + # already is the source byte-for-byte. + if self._is_valid_pe(self.launcher) and self._files_match(self.launcher, source): + return True + last_error: Optional[OSError] = None + for attempt in range(self._RESTORE_ATTEMPTS): + try: + self._atomic_copy(source, self.launcher) + return self._is_valid_pe(self.launcher) + except OSError as exc: + last_error = exc + if attempt + 1 < self._RESTORE_ATTEMPTS: + time.sleep(0.1) + if last_error is not None: + typer.echo(f"Error: could not restore the Unsloth launcher: {last_error}", err = True) + return False + + def _restore_runnable(self) -> bool: + """Put back the first copy that actually runs. + + Passing the two-byte header check does not make a copy runnable, so a + candidate that fails --version must not stop the next one being tried, + and a launcher already in place and working must not be replaced by a + candidate that is merely PE-shaped. + """ + if self._launcher_health_error() is None: + return True + candidates = self._recovery_candidates() + for source in candidates: + if self._restore_from(source) and self._launcher_health_error() is None: + return True + # Nothing ran. Leave the best candidate in place rather than whichever + # one happened to be tried last. + if candidates: + self._restore_from(candidates[0]) + return False + + def _launcher_health_error(self) -> Optional[str]: + assert self.launcher is not None + if not self._is_valid_pe(self.launcher): + return "the updated launcher is missing or is not a non-empty PE file" + try: + result = subprocess.run( + [str(self.launcher), "--version"], + check = False, + capture_output = True, + timeout = self._VERSION_TIMEOUT_SECONDS, + **_windows_hidden_subprocess_kwargs(), + ) + except subprocess.TimeoutExpired: + return f"the updated launcher timed out after {self._VERSION_TIMEOUT_SECONDS} seconds" + except OSError as exc: + return f"the updated launcher could not run --version ({exc})" + if result.returncode != 0: + return f"the updated launcher returned {result.returncode} for --version" + return None + + @staticmethod + def _managed_scripts_dir() -> Path: + """Scripts dir of the venv setup actually updates. + + setup.ps1 installs into STUDIO_HOME/unsloth_studio, which is not this + interpreter when a pip-installed or checkout CLI drives the update. Same + distinction _studio_deps._managed_root draws for the damage scan. + """ + managed = STUDIO_HOME / "unsloth_studio" + if (managed / "pyvenv.cfg").is_file(): + try: + foreign = managed.resolve() != Path(sys.prefix).resolve() + except OSError: + foreign = True + if foreign: + return managed / "Scripts" + return Path(sys.executable).resolve().parent + + def __enter__(self): + if not self.enabled: + return self + try: + scripts = self._managed_scripts_dir() + except (OSError, RuntimeError) as exc: + typer.echo(f"Error: could not resolve the managed Python environment: {exc}", err = True) + raise typer.Exit(1) + self.launcher = scripts / "unsloth.exe" + self.backup = scripts / "unsloth.exe.update-backup" + self.legacy_backup = scripts / "unsloth.exe.deleteme" + # Under the Studio home, not the venv: setup.ps1 removes the whole + # $VenvDir to rebuild a stale torch, and an open handle inside it makes + # Windows refuse the recursive delete. One lock per Studio home is the + # right grain anyway, since that is what names the managed venv. + self.lock_path = STUDIO_HOME / "unsloth.exe.update-lock" + # install.ps1 hardlinks this to the launcher, so it survives the old + # updater's .deleteme unlink and is a valid recovery source. + self.shim = STUDIO_HOME / "bin" / "unsloth.exe" + self.stale = scripts / "unsloth.exe.update-stale" + self._acquire_lock() + try: + self._recover_missing_launcher() + if not self._is_valid_pe(self.launcher): + # Warn, do not exit. The previous updater could leave an install + # with no launcher and no .deleteme, and refusing here would stop + # exactly those users from ever updating again. Setup may well + # write a new launcher; validate_launcher still judges the result. + typer.echo( + f"Warning: the managed Unsloth launcher is missing or invalid: {self.launcher}", + err = True, + ) + typer.echo("Continuing; setup may reinstall it.", err = True) + if self._retained_backup() is None: + self.backup = None + elif self._retained_backup() is None: + # Only write a backup when there is no usable one already. A + # backup outlives __enter__ only when a previous run died before + # validating, so it holds the last launcher known to run, while + # the canonical file has passed nothing but the two-byte header + # check. Overwriting it here destroyed the only recovery copy. + try: + self._atomic_copy(self.launcher, self.backup) + except OSError as exc: + # A backup is a safety net, not a precondition. Antivirus or a + # locked-down Scripts dir must not abort the update outright. + typer.echo(f"Warning: could not back up the Unsloth launcher: {exc}", err = True) + self.backup = None + self._move_launcher_aside() + except BaseException: + self._release_lock() + raise + return self + + def validate_launcher(self) -> None: + if not self.enabled: + return + # Whether setup published anything decides how a bad result is read, so + # it has to be sampled before any restore puts a launcher back. + published = self.launcher.exists() + error = self._launcher_health_error() + if error is not None: + restored = self._restore_runnable() + # Setup publishing nothing is the case this transaction exists for: + # a no-op pip update leaves the freed path empty, and the old + # updater then deleted its own .deleteme, leaving no launcher at + # all. Putting the previous one back is success, not failure. A + # launcher setup DID write and that cannot run is still a failure, + # even though the previous one goes back. + if published or not restored: + typer.echo(f"Error: Unsloth Studio update failed because {error}.", err = True) + if restored: + typer.echo("The previous launcher was restored.", err = True) + elif self._retained_backup() is not None: + typer.echo(f"Manual recovery copy retained at: {self.backup}", err = True) + raise typer.Exit(1) + self._validated = True + for orphan in (self.stale, self.backup, self.legacy_backup): + if orphan is None: + continue + try: + orphan.unlink(missing_ok = True) + except OSError: + pass + + def __exit__(self, exc_type, exc_value, traceback) -> bool: + try: + if self.enabled and exc_type is not None and not self._validated: + if not self._restore_runnable() and self._retained_backup() is not None: + typer.echo(f"Manual recovery copy retained at: {self.backup}", err = True) + finally: + self._release_lock() + return False # ── unsloth studio reset-password ──────────────────────────────────── diff --git a/unsloth_cli/tests/test_studio_update_local_repo.py b/unsloth_cli/tests/test_studio_update_local_repo.py index 7ae8bcf0de..45bbe6de1e 100644 --- a/unsloth_cli/tests/test_studio_update_local_repo.py +++ b/unsloth_cli/tests/test_studio_update_local_repo.py @@ -29,13 +29,23 @@ def _studio(): return _studio_mod +class _NoopLauncherUpdate: + def __enter__(self): + return self + + def validate_launcher(self): + pass + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def _neutered(monkeypatch): """Stub everything update does after resolving the repo root.""" studio = _studio() seen = {} monkeypatch.setattr(studio, "_ensure_studio_env_exported", lambda *a, **k: None) - monkeypatch.setattr(studio, "_release_self_exe_lock_windows", lambda *a, **k: None) - monkeypatch.setattr(studio, "_cleanup_self_exe_lock_windows", lambda *a, **k: None) + monkeypatch.setattr(studio, "_WindowsLauncherUpdateTransaction", _NoopLauncherUpdate) monkeypatch.setattr(studio, "_refresh_desktop_shortcuts", lambda *a, **k: None) monkeypatch.setattr(studio, "_fail_if_install_damaged", lambda *a, **k: None, raising = False) diff --git a/unsloth_cli/tests/test_studio_update_verify.py b/unsloth_cli/tests/test_studio_update_verify.py index 081562a698..97709912e3 100644 --- a/unsloth_cli/tests/test_studio_update_verify.py +++ b/unsloth_cli/tests/test_studio_update_verify.py @@ -368,10 +368,20 @@ def test_the_verify_help_does_not_promise_an_import_check(): def _run_update(monkeypatch, argv, verified): studio = _studio() + + class _NoopLauncherUpdate: + def __enter__(self): + return self + + def validate_launcher(self): + pass + + def __exit__(self, exc_type, exc_value, traceback): + return False + monkeypatch.setattr(studio, "_ensure_studio_env_exported", lambda *a, **k: None) + monkeypatch.setattr(studio, "_WindowsLauncherUpdateTransaction", _NoopLauncherUpdate) monkeypatch.setattr(studio, "_run_setup_script", lambda *a, **k: None) - monkeypatch.setattr(studio, "_release_self_exe_lock_windows", lambda *a, **k: None) - monkeypatch.setattr(studio, "_cleanup_self_exe_lock_windows", lambda *a, **k: None) monkeypatch.setattr(studio, "_refresh_desktop_shortcuts", lambda *a, **k: None) monkeypatch.setattr(studio, "_fail_if_install_damaged", lambda: verified.append(True)) return CliRunner().invoke(studio.studio_app, ["update", *argv]) @@ -585,3 +595,107 @@ def test_the_repair_command_quotes_the_interpreter(monkeypatch, capsys, system, with pytest.raises(typer.Exit): studio._fail_if_install_damaged() assert expected in capsys.readouterr().err + + +# ── runtime-irrelevant rows must not fail an update ────────────────── + + +def test_a_shared_top_level_test_tree_is_not_damage(site): + # Reported as `einx: test/conftest.py is missing`. einx and torchao both + # ship it, and install_python_stack.py force-reinstalls torchao every + # update, so pip removes the file and the pinned torchao does not ship it. + # Nothing imports another project's fixtures, and no reinstall repairs it. + _make_dist(site, "einx", {"einx/__init__.py": b"e\n"}) + (site / "einx-1.0.dist-info" / "RECORD").write_text( + "einx/__init__.py,sha256=x,2\ntest/conftest.py,sha256=x,20650\n" + ) + assert _deps().damaged_installed_files() == [] + + +def test_an_installer_rewritten_lockfile_is_not_damage(site): + # Reported as `package-lock.json is 27225 bytes, expected 28473`. + # setup.ps1/setup.sh run `npm install` inside the installed tree, and npm + # dedupes hoisted entries under legacy-peer-deps, shrinking the file. + lock = "studio/backend/core/data_recipe/oxc-validator/package-lock.json" + _make_dist( + site, + "unsloth", + {"unsloth/__init__.py": b"u\n", lock: b"L" * 27225}, + record_sizes = {lock: 28473}, + ) + assert _deps().damaged_installed_files() == [] + + +def test_a_deleted_installer_rewritten_file_is_still_damage(site): + # Only the SIZE of these drifts, because npm rewrites the lockfile in place. + # It never deletes it, so a missing one is real damage and must be reported. + lock = "studio/backend/core/data_recipe/oxc-validator/package-lock.json" + _make_dist(site, "unsloth", {"unsloth/__init__.py": b"u\n", lock: b"L" * 100}) + (site / lock).unlink() + found = _deps().damaged_installed_files() + assert len(found) == 1 and "package-lock.json is missing" in found[0] + + +def test_a_shared_top_level_scripts_tree_is_not_damage(site): + # unsloth_zoo ships a top-level scripts/, the same squatted-namespace shape + # as einx's test/. It has no __init__.py, so nothing imports it. + _make_dist(site, "upsilon", {"upsilon/__init__.py": b"u\n"}) + (site / "upsilon-1.0.dist-info" / "RECORD").write_text( + "upsilon/__init__.py,sha256=x,2\nscripts/helper.py,sha256=x,99\n" + ) + assert _deps().damaged_installed_files() == [] + + +def test_a_package_owned_tests_subdirectory_is_still_checked(site): + # Only the shared top-level namespace is exempt; a tests/ tree inside a + # package is that package's alone, so a deletion there is real. + _make_dist(site, "rho", {"rho/tests/helper.py": b"h\n"}) + (site / "rho" / "tests" / "helper.py").unlink() + found = _deps().damaged_installed_files() + assert len(found) == 1 and "rho/tests/helper.py is missing" in found[0] + + +def test_a_top_level_module_named_like_a_test_root_is_still_checked(site): + # The exemption is for a shared directory, not for a name prefix. + _make_dist(site, "sigma", {"tests.py": b"t\n"}) + (site / "tests.py").unlink() + found = _deps().damaged_installed_files() + assert len(found) == 1 and "tests.py is missing" in found[0] + + +def test_runtime_damage_still_fails_when_ignored_rows_are_present(site): + # The exemption must not blind the scan to a torn runtime module. + lock = "studio/backend/core/data_recipe/oxc-validator/package-lock.json" + _make_dist( + site, + "unsloth", + {"unsloth/__init__.py": b"u\n", lock: b"L" * 10}, + record_sizes = {lock: 28473}, + ) + (site / "unsloth" / "__init__.py").unlink() + found = _deps().damaged_installed_files() + assert len(found) == 1 and "unsloth/__init__.py is missing" in found[0] + + +def test_ignored_rows_do_not_consume_the_finding_budget(site): + # Filtering happens while RECORD is read, so harmless rows cannot crowd a + # real one off a capped list. Unfiltered, these 40 fill limit = 3. + files = {f"test/t{i}.py": b"x" for i in range(40)} + files["tau/__init__.py"] = b"t\n" + _make_dist(site, "tau", files) + for rel in files: + (site / rel).unlink() + found = _deps().damaged_installed_files(limit = 3) + assert len(found) == 1 and "tau/__init__.py is missing" in found[0] + + +def test_our_own_top_level_trees_are_still_checked(site): + # A missing __init__.py does not make a directory unimportable (PEP 420), + # and this repo imports scripts.* itself, so the shared-namespace exemption + # must never apply to what Unsloth ships. + _make_dist(site, "unsloth_zoo", {"unsloth_zoo/__init__.py": b"z\n"}) + (site / "unsloth_zoo-1.0.dist-info" / "RECORD").write_text( + "unsloth_zoo/__init__.py,sha256=x,2\nscripts/helper.py,sha256=x,99\n" + ) + found = _deps().damaged_installed_files() + assert len(found) == 1 and "scripts/helper.py is missing" in found[0]