vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-11-01 03:30:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 32

Studio: notice a quarantined payload before trusting the update fast path (#10053)

Browse source 
* Studio: notice a quarantined payload before trusting the update fast path

An update can report success while leaving the venv unbootable. Every check the
fast path makes reads metadata: installed_version_probe reads dist-info,
verify_install compares the manifest and walks missing_requirements through
importlib.metadata. None of them looks at whether the recorded files are still
on disk. So when an antivirus quarantines site-packages/unsloth after a good
install, the version still answers, verify_install still says ok, and setup
takes the fast path and skips the dependency pass that would have repaired it.

Measured on a real 2026.8.22 install: deleting a recorded file, truncating one,
removing a nested module, or removing the whole tree all leave `studio update`
exiting 0 and reporting the update successful.

damaged_payload_files walks the managed distribution's own RECORD and reports
rows that are gone, are not regular files, or are shorter than recorded.
verify_install runs it last, once every metadata check has already passed, and
fails with a new reason, studio_install_damaged. That is the whole change on the
shell side: setup.sh and setup.ps1 already call verify_install and already force
the dependency pass when it says no, so neither script is touched.

Deliberately narrow, in three ways.

It answers for the managed package alone rather than for every installed
distribution, because its only job here is deciding whether to repair ours.

It runs last, after manifest_ok and deps_ok are both established, so none of the
six existing reason strings can be diverted into the new one. The reason is
additive and nothing matches it against a closed set: the Tauri preflight
carries it as an opaque Option<String>, and is_context_reason is a negated
two-item allowlist, so an unknown reason keeps can_auto_repair true, which is
what we want here.

It is off on the boot path. `deep` defaults on so the two shell callers, which
pass no arguments, get the scan; install_state passes deep=False because it
feeds `desktop-capabilities`, which the preflight spawns under a 10 second
timeout. Spending that budget on a stat of every recorded file risks a timeout
reporting a healthy install stale and repairing it, which is worse than the
quarantine it would have caught. `unsloth studio update --verify` already scans
the whole tree after the pass. The kwarg is guarded by the same
_verify_install_supports introspection as installed_conflicts, so a newer CLI
against an older install_manifest.py still works.

Two carve-outs, both for trees our own setup rewrites. package-lock.json loses
only its size claim, since npm dedupes hoisted entries and shrinks it, matching
_INSTALLER_REWRITTEN_NAMES. studio/frontend/dist is skipped as a subtree: setup
runs `npm run build` in the installed tree, vite empties dist and rewrites every
asset under a fresh content hash, and the wheel ships all 704 of those files in
RECORD, so a rebuild leaves rows naming files that no longer exist. A size
waiver cannot reach that, and without the carve-out one rebuild would fail every
later update permanently, since a no-op pip pass restores nothing.

RECORD is parsed with csv over io.StringIO, not splitlines, because a quoted
field may contain a newline and splitlines also breaks on \v, \f and the Unicode
separators. Unreadable is not missing: anything but FileNotFoundError is skipped,
because a reinstall cannot fix EACCES. An environment the scan cannot read
reports undamaged, so a check that cannot answer leaves the fast path as it
found it.

Verified: 294 distributions in a real environment, zero flagged, slowest scan
0.168s; unsloth 2024.10.0 through 2026.8.22 all clean; all four damage states
caught; a frontend rebuilt in place is clean while real damage in the same tree
is still caught. Existing suites unchanged at 2887 passed, including the 83 that
pin verify_install's reasons, and the shell fast-path suites still pass with no
shell changes.

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

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

* Make the payload scan opt-in and bounded

Three problems with the first cut, all found by auditing it rather than by
using it.

The scan defaulted on, which is the wrong direction under version skew. An
external CLI driving a managed venv resolves install_manifest out of THAT venv:
_manifest_candidates yields the CLI's own sibling copy first, then falls through
to sys.prefix, so an older published CLI can load a newer copy of this file. It
calls verify_install(root=root) with no `deep`, so defaulting the scan on put a
stat of every recorded file inside the 10 second timeout that
desktop-capabilities runs under on the Tauri boot path, with no way for the old
caller to decline. interrupted-install-ci.yml builds exactly that pairing on
purpose. `deep` now defaults False and the two shell callers opt in; they import
the module from their own directory, so they can never be skewed against it, and
they fall back to the old call on a TypeError so an older tree still works.

The scan was unbounded. Neither installer wraps that call in a timeout, and
setup.sh only reaches it inside the "is up to date" fast path, so the cost lands
on the no-op run. 3644 of unsloth's 4358 recorded rows get stat'ed, ~65ms warm.
Measured with latency injected into a real scan of a real install: at 5ms per
stat a full walk is 18s, at 50ms it is 182s of wedged setup. A 5 second budget,
checked every 64 rows so the clock read costs less than the stat it guards,
holds those to 5.3s and 6.3s. Giving up reports undamaged, so a scan that cannot
finish leaves the fast path exactly as it found it, which is what main does
today.

`unsloth studio verify-install` could not see the failure mode this was written
to catch. It shares _install_state with desktop-capabilities, so opting the boot
path out opted the diagnostic out too. install_state takes `deep` now and
verify-install passes it: nothing times that command out, and a quarantined
payload is precisely what a user running a verify step wants told.
_fail_if_install_damaged does not cover this, because it returns silently when
the CLI is outside the managed venv, which is the desktop configuration.

Also: reuse the manifest already in scope instead of re-reading it from disk,
which removes a second TOCTOU window.

The boot path staying cheap was load-bearing and unpinned, so a refactor could
have reinstated the scan there with a green suite. Two tests now hold it: one
that both installers ask for the scan and keep their older-tree fallback, one
that install_state and _install_state default it off while verify-install turns
it on.

Verified: the fast-path decision is unchanged from main on a healthy install and
now forces the repair pass on a quarantined one, across [Linux, WSL, macOS,
Windows] x [NVIDIA, AMD, CPU, XPU], 32 cells run twice each against both module
versions, driving the real setup.sh block against a real venv. Run differentially
because absolute expectations are useless here: unrelated escapes clear the flag
too, and the anyio #6483 one masked the effect entirely until the venv was pinned
to the anyio the requirements actually ask for. 15 compatibility checks over all
four old/new pairings of the two files. 294 distributions scanned, zero flagged.
unsloth 2024.10.0 through 2026.8.22 clean. unsloth_cli/tests fails the same 41
tests as origin/main, name for name.

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

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

* Repair the damage the scan finds, and widen what counts as damage

The scan told setup an install was damaged and setup then ran an upgrade,
which for a distribution already at the wanted version installs nothing:
uv audits it, pip calls it satisfied, and both are reading metadata that a
quarantine of the payload leaves intact. Detection without repair only moved
the failure later, so the core phase now names a damaged distribution for
reinstall, with --no-deps so the pre-installed torch build is not re-resolved.
--reinstall-package is uv-only; the pip fallback gets --force-reinstall, which
is safe there only because that call also passes --no-deps and one name.

Also five cases the first pass read as healthy:

- a package directory replaced by a regular file. NotADirectoryError is an
  OSError but not a FileNotFoundError, so every row beneath it fell to the
  generic arm and reported merely unreadable.
- unsloth-zoo. studio.txt does not list it, so nothing else here would notice
  it quarantined, and training, export and inference all import it.
- a foreign managed venv. verify_install now takes scan_paths, and the CLI
  hands over that venv's own site-packages; RECORD rows resolve against the
  distribution that declared them, so without it the scan answered for the
  external CLI's tree and deep was a no-op on that branch.
- a distribution gone entirely. Every version check compares against the
  installed version, so an absent one passes all of them.
- the scan budget. It was consulted every 64 rows, which let a mount taking a
  second per stat overrun 5 seconds by a minute. A clock read is ~68ns against
  ~1343ns for a warm stat, so the guard is cheaper than the call it guards.

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

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

* Trim the comments added by this branch

Comments and docstrings only, verified with comment_tools.py check
--strip-docstrings. 179 comment lines to 66.

* Abort when the repair cannot restore the files, and scan the launcher

A forced reinstall that cannot reach an index left the tree damaged and the
pass carried on, where the ordinary upgrade audits the intact metadata as
satisfied and write_manifest records a success nothing rechecks. The repair
now rescans and the installer aborts. Judged on the tree rather than pip's
exit code, so a reinstall that restored the files while exiting non-zero is
not held against it.

Console scripts and data files are recorded relative to site-packages, so
`../../../bin/unsloth` was dropped by the blanket `..` rejection: an antivirus
could take the launcher, leave the package tree intact, and the deep check
would call the install healthy while the command was gone. Parent-relative
rows are now resolved and required to land inside the venv, which is what
keeps a row naming something outside it from being read as our damage.
Measured across every distribution installed here: 647 such rows in 81 of
294 distributions, none flagged, and a quarantined `bin/flask` detected.

A companion whose dist-info is gone left the scan nothing to walk, and no
check above ever looked at its version, so unsloth-zoo could vanish whole.
Only for the default install: a custom `--package` need not depend on it, and
demanding it there would repair that environment on every run.

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

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

* Refuse an absent core package, and leave a custom package's neighbours alone

A distribution that is not installed at all has no RECORD to walk, so the scan
read it as undamaged. That is only safe before the core phase, where a fresh
run has nothing installed yet, so a second gate runs after it with
require_present: by then the phase has run or was deliberately skipped, and
the install.sh handoff that sets SKIP_STUDIO_BASE=1 never reaches it at all.
Without this, write_manifest recorded a finished install for an environment
with no unsloth-zoo in it.

unsloth-zoo is now claimed only by the default install, in the installer and
in verify_install alike. `--package X` installs X on its own, so a damaged
unsloth-zoo beside it belongs to something else: force-reinstalling it would
mutate an unrelated distribution and fail the update on an unreachable index.

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

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

* Check presence after a repair, not only before it

--force-reinstall uninstalls before it installs, so a failure in between
leaves no distribution at all, and a distribution that is not there has no
RECORD for the follow-up scan to call damaged. The repair reported success on
the state it had just created. Presence is now rechecked after every reinstall
attempt, unconditionally rather than under require_present: we only reinstall
what the scan already found, so it was installed a moment earlier.

The manifest written in that state records no version, and verify_install
compared against it with `bool(recorded) and not current`, so the deep check
called an environment with nothing installed healthy. An empty version is only
ever absence here, since two records or an unreadable one set local_conflict
and report studio_install_metadata_conflict before this runs.

* Bound a stat that never returns, not only many that are slow

The deadline inside the walk is checked between calls, so it bounds an
accumulation of slow stats but not one that hangs outright, which is what a
wedged mount or a filesystem filter produces. Neither installer wraps its
verify_install call in a timeout, so that case blocked setup rather than
reporting the scan unfinished as documented.

The walk now runs on a daemon thread and is abandoned past the budget. The
interpreter does not wait for one at exit: measured at 1.04s to exit with a
thread parked forever in a syscall. budget_seconds = 0 keeps the old inline
call for the installer, which has already committed to a full pass and needs
the real answer.

The new test hangs until killed without this and passes with it.

* Do not read our own staged launcher as a quarantined one

_move_launcher_aside renames Scripts/unsloth.exe to .update-stale before
setup runs, and setup.ps1's deep check runs inside that window, so scanning
parent-relative rows made every ordinary Windows update see its own launcher
missing: the fast path was lost and the core package force-reinstalled over
the network on a healthy no-op update. A recorded file that is absent with
one of the three siblings _recover_missing_launcher reads beside it is staged,
not quarantined. Absent without one is still damage.

_core_package_names compared the package name exactly while verify_install
canonicalizes it, so `--package Unsloth` had the deep check scan unsloth-zoo
and force a pass that neither repair gate would then act on.

Two unit tests called the real installed_versions through the presence check
added beside the scan, so they answered for the host rather than the code and
failed on a checkout with no unsloth distribution installed.

* Narrow the staged-launcher excuse, and refuse a helper that will not import

The staged exemption applied to any recorded file and asked only whether a
sibling existed. Both were too loose: nothing but Scripts/unsloth.exe is ever
moved aside, so a stray .update-stale could hide a quarantine anywhere in the
tree, and the transaction recovers through _is_valid_pe, so an orphaned or
truncated copy is one it would reject while this accepted it and left the
launcher missing with the fast path intact. Same two-byte test now, and only
for that name.

An install_manifest.py that is present but will not import is damage, not an
old release, and it is the one file whose damage silences every check that
follows it. Both installers now force the dependency pass for that. Absence
keeps the old escape: setup.sh is resolved out of the installed studio package
so it should be impossible, separating it from an old release here would need
a RECORD walk inside the installer, and the CLI already reports that case as
studio_install_manifest_missing.

Verified by extracting the probe from setup.sh and running it: absent keeps
the fast path, truncated and raising both force the repair.

* Trim the comments added since the opening pass

Comments and docstrings only, verified with comment_tools.py check
--strip-docstrings. 120 comment lines to 73.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-08-31 04:01:34 -07:00 • committed by GitHub
parent e5767a01e4
commit fa57a7986f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1328 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

171
tests/python/test_install_python_stack.py
View file

@ -773,6 +773,177 @@ class TestBuildPipCmdUpgradeIntent:
assert cmd == [sys.executable, "-m", "pip", "install", "--no-cache-dir", "somepackage"]
class TestDamagedCorePayloadRepair:
"""An upgrade of a distribution already at the wanted version installs
nothing: uv audits it, pip calls it satisfied, and both read metadata a
quarantine of the payload leaves intact."""
def test_the_uv_command_targets_only_the_named_package(self):
cmd = ips._build_uv_cmd(
("--no-cache-dir", "--no-deps", "--reinstall-package", "x", "--force-reinstall", "x")
)
assert "--reinstall-package" in cmd
# Not --reinstall: that is the whole environment, torch included.
assert "--reinstall" not in cmd
assert "--no-deps" in cmd
def test_a_plain_force_reinstall_still_becomes_uv_reinstall(self):
assert "--reinstall" in ips._build_uv_cmd(("--force-reinstall", "x"))
def test_the_pip_fallback_drops_the_uv_only_flag(self):
cmd = ips._build_pip_cmd(
("--no-cache-dir", "--no-deps", "--reinstall-package", "x", "--force-reinstall", "x")
)
assert "--reinstall-package" not in cmd
assert "--force-reinstall" in cmd and "--no-deps" in cmd
assert cmd.count("x") == 1
# It is not an upgrade request, so it must not acquire pip's upgrade flags.
assert "--upgrade" not in cmd
def test_a_damaged_package_is_reinstalled(self, monkeypatch):
calls = []
monkeypatch.setattr(
ips.install_manifest,
"damaged_payload_files",
lambda name, **kwargs: ["gone"] if name == "unsloth-zoo" else [],
)
monkeypatch.setattr(ips.install_manifest, "installed_versions", lambda name: ["1.0"])
monkeypatch.setattr(ips, "pip_install_try", lambda label, *args: calls.append(args))
monkeypatch.setattr(ips, "_step", lambda *a, **k: None)
ips._repair_damaged_core_payload(("unsloth", "unsloth-zoo"))
assert len(calls) == 1
assert "--reinstall-package" in calls[0] and calls[0][-1] == "unsloth-zoo"
def test_a_healthy_tree_runs_nothing(self, monkeypatch):
monkeypatch.setattr(ips.install_manifest, "damaged_payload_files", lambda *a, **k: [])
monkeypatch.setattr(
ips, "pip_install_try", lambda *a, **k: pytest.fail("nothing to repair")
)
ips._repair_damaged_core_payload(("unsloth", "unsloth-zoo"))
def test_a_local_checkout_is_left_alone(self, monkeypatch):
"""Its core packages are an editable overlay, not an index install."""
monkeypatch.setattr(
ips.install_manifest,
"damaged_payload_files",
lambda *a, **k: pytest.fail("a local checkout must not even be scanned"),
)
ips._repair_damaged_core_payload(("unsloth",), local_repo = "/src/unsloth")
def test_a_scan_that_raises_does_not_stop_the_install(self, monkeypatch):
def boom(*args, **kwargs):
raise OSError("unreadable")
monkeypatch.setattr(ips.install_manifest, "damaged_payload_files", boom)
monkeypatch.setattr(
ips, "pip_install_try", lambda *a, **k: pytest.fail("nothing was proven damaged")
)
ips._repair_damaged_core_payload(("unsloth",))
def test_a_repair_that_leaves_the_files_missing_fails(self, monkeypatch, capsys):
"""Otherwise the pass that follows audits the intact metadata as
satisfied and write_manifest records a success nothing rechecks."""
monkeypatch.setattr(
ips.install_manifest, "damaged_payload_files", lambda name, **kwargs: ["gone"]
)
monkeypatch.setattr(ips, "pip_install_try", lambda label, *args: False)
monkeypatch.setattr(ips, "_step", lambda *a, **k: None)
assert ips._repair_damaged_core_payload(("unsloth",)) is False
def test_a_reinstall_that_restored_the_files_is_a_success(self, monkeypatch):
"""Judged on the tree, not pip's exit code: uv and pip both exit non-zero
for reasons that have nothing to do with the payload."""
seen = {"n": 0}
def scan(name, **kwargs):
seen["n"] += 1
return ["gone"] if seen["n"] == 1 else []
monkeypatch.setattr(ips.install_manifest, "damaged_payload_files", scan)
# Stubbed, or the presence check added beside the scan answers for the
# host: a source checkout with no unsloth distribution installed would
# fail this on the environment rather than on the code.
monkeypatch.setattr(ips.install_manifest, "installed_versions", lambda name: ["1.0"])
monkeypatch.setattr(ips, "pip_install_try", lambda label, *args: False)
monkeypatch.setattr(ips, "_step", lambda *a, **k: None)
assert ips._repair_damaged_core_payload(("unsloth",)) is True
def test_an_absent_distribution_is_refused_after_the_core_phase(self, monkeypatch):
"""It has no RECORD to walk, so the scan alone reads it as undamaged.
The install.sh handoff sets SKIP_STUDIO_BASE=1 and the core phase never
runs, so this is the only place that would see unsloth-zoo gone before
write_manifest records the environment as a finished install.
"""
monkeypatch.setattr(ips.install_manifest, "damaged_payload_files", lambda *a, **k: [])
monkeypatch.setattr(ips.install_manifest, "installed_versions", lambda name: [])
monkeypatch.setattr(ips, "_safe_print", lambda *a, **k: None)
assert ips._repair_damaged_core_payload(("unsloth",), require_present = True) is False
# Off before the core phase: a fresh run has nothing installed yet.
assert ips._repair_damaged_core_payload(("unsloth",)) is True
def test_a_present_distribution_passes_the_presence_check(self, monkeypatch):
monkeypatch.setattr(ips.install_manifest, "damaged_payload_files", lambda *a, **k: [])
monkeypatch.setattr(ips.install_manifest, "installed_versions", lambda name: ["1.0"])
assert ips._repair_damaged_core_payload(("unsloth",), require_present = True) is True
def test_the_companion_is_only_for_the_default_install(self):
assert ips._core_package_names("unsloth") == ("unsloth", "unsloth-zoo")
assert ips._core_package_names("unsloth-nightly") == ("unsloth-nightly",)
def test_the_default_is_recognised_however_it_is_spelled(self):
"""verify_install canonicalizes, and the two disagreeing meant the deep
check scanned zoo, forced a pass, and no repair gate would touch it."""
for spelling in ("Unsloth", "UNSLOTH", "UnSloth"):
assert ips._core_package_names(spelling)[1:] == ("unsloth-zoo",), spelling
# PEP 503 folds runs of -_. to -, it does not delete them, so this is a
# different project and keeps its neighbours to itself.
assert ips._core_package_names("uns_loth") == ("uns_loth",)
def test_both_repair_sites_use_that_list(self):
source = inspect.getsource(ips)
assert (
source.count("_repair_damaged_core_payload(\n _core_package_names")
+ source.count("_repair_damaged_core_payload(_core_package_names")
== 2
)
assert "_repair_damaged_core_payload((package_name" not in source
def test_the_second_site_runs_before_the_manifest_is_written(self):
source = inspect.getsource(ips)
gate = source.rindex("_repair_damaged_core_payload(")
assert gate < source.index("install_manifest.write_manifest(")
def test_a_reinstall_that_removed_the_distribution_fails(self, monkeypatch):
"""--force-reinstall uninstalls before it installs, and a failure in
between leaves nothing: an absent distribution has no RECORD, so the
payload scan alone would call the repair a success."""
scans = {"n": 0}
def scan(name, **kwargs):
scans["n"] += 1
# Damaged on the way in, and afterwards nothing left to call damaged.
return ["gone"] if scans["n"] == 1 else []
monkeypatch.setattr(ips.install_manifest, "damaged_payload_files", scan)
monkeypatch.setattr(ips.install_manifest, "installed_versions", lambda name: [])
monkeypatch.setattr(ips, "pip_install_try", lambda label, *args: True)
monkeypatch.setattr(ips, "_step", lambda *a, **k: None)
monkeypatch.setattr(ips, "_safe_print", lambda *a, **k: None)
assert ips._repair_damaged_core_payload(("unsloth",)) is False
def test_the_caller_aborts_the_install(self):
source = inspect.getsource(ips)
assert "if not _repair_damaged_core_payload(" in source
def test_the_repair_runs_before_the_core_phase(self):
"""After it, the upgrade would already have audited the damage as fine."""
source = inspect.getsource(ips)
repair = source.index("_repair_damaged_core_payload(_core_package_names")
core = source.index("# 3. Core packages")
assert repair < core
class TestDuplicateCoreMetadataRepair:
def test_an_unrewritable_record_stops_the_repair_before_pip_runs(
self, tmp_path, monkeypatch, capsys