275 lines
10 KiB
Python
275 lines
10 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Behavioural guards for the two in-container update helpers of the Docker image.
|
|
|
|
Both are `docker exec` entry points that mutate a running container, so a wrong
|
|
answer costs an outage or a mixed-version install:
|
|
|
|
* `unsloth-studio-update` swaps the Studio Python packages and then restarts the
|
|
service. It verifies the new backend imports first, but only warned -- so a
|
|
release that pulls in a dependency `--no-deps` did not install got the healthy
|
|
old process killed and replaced by one that cannot start. supervisord retries
|
|
`startretries` times, lands in FATAL and never leaves it on its own, so the
|
|
container serves nothing until someone exec's in.
|
|
* `unsloth-llama-update --check` reported "up to date" when it could not reach
|
|
the release feed at all, and its in-place rollback only removed entries whose
|
|
names the OLD tree also had, leaving new-release-only shared objects beside
|
|
the restored files. ggml dlopen()s every `libggml-*.so` next to the binaries,
|
|
so that mix is loaded on the next GGUF run.
|
|
|
|
These drive the real scripts with stub `pip` / `supervisorctl` / `python` /
|
|
`mv` on PATH. No docker, no GPU, no network.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
STUDIO_UPDATE = REPO_ROOT / "docker" / "unsloth_studio_update.sh"
|
|
LLAMA_UPDATE = REPO_ROOT / "docker" / "unsloth_llama_update.sh"
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
shutil.which("bash") is None,
|
|
reason = "needs bash",
|
|
)
|
|
|
|
|
|
def _stub(directory: Path, name: str, body: str) -> None:
|
|
directory.mkdir(parents = True, exist_ok = True)
|
|
path = directory / name
|
|
path.write_text("#!/usr/bin/env bash\n" + body, encoding = "utf-8")
|
|
path.chmod(0o755)
|
|
|
|
|
|
def _run(
|
|
script: Path,
|
|
args,
|
|
env,
|
|
cwd = None,
|
|
):
|
|
return subprocess.run(
|
|
["bash", str(script), *args],
|
|
capture_output = True,
|
|
text = True,
|
|
env = env,
|
|
cwd = cwd,
|
|
timeout = 120,
|
|
)
|
|
|
|
|
|
# --- unsloth-studio-update ----------------------------------------------------
|
|
|
|
|
|
def _studio_env(tmp_path: Path, *, import_ok: bool) -> dict:
|
|
home = tmp_path / "studio"
|
|
venv_bin = home / "unsloth_studio" / "bin"
|
|
venv_bin.mkdir(parents = True)
|
|
_stub(
|
|
venv_bin,
|
|
"python",
|
|
'if [ "$1" = "-c" ]; then\n'
|
|
+ (
|
|
" exit 0\n"
|
|
if import_ok
|
|
else ' case "$2" in *studio.backend.main*) exit 1;; esac\n exit 0\n'
|
|
)
|
|
+ "fi\n"
|
|
'if [ "$1" = "-m" ] && [ "$2" = "pip" ]; then\n'
|
|
' if [ "$3" = "show" ]; then echo "Version: 2026.7.5"; exit 0; fi\n'
|
|
' echo "STUB-PIP $*" >> "$STUB_LOG"; exit 0\n'
|
|
"fi\n"
|
|
"exit 0\n",
|
|
)
|
|
bin_dir = tmp_path / "bin"
|
|
_stub(
|
|
bin_dir,
|
|
"supervisorctl",
|
|
'echo "STUB-SUPERVISORCTL $*" >> "$STUB_LOG"\n'
|
|
'if [ "$1" = "status" ]; then exit 0; fi\nexit 0\n',
|
|
)
|
|
env = dict(os.environ)
|
|
env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"]
|
|
env["UNSLOTH_STUDIO_HOME"] = str(home)
|
|
env["STUB_LOG"] = str(tmp_path / "calls.log")
|
|
return env
|
|
|
|
|
|
def test_studio_update_restarts_when_the_backend_imports(tmp_path: Path):
|
|
env = _studio_env(tmp_path, import_ok = True)
|
|
res = _run(STUDIO_UPDATE, [], env)
|
|
calls = Path(env["STUB_LOG"]).read_text() if Path(env["STUB_LOG"]).exists() else ""
|
|
assert res.returncode == 0, res.stderr
|
|
assert "STUB-SUPERVISORCTL restart studio" in calls, calls
|
|
|
|
|
|
def test_studio_update_does_not_restart_into_a_backend_that_cannot_import(tmp_path: Path):
|
|
env = _studio_env(tmp_path, import_ok = False)
|
|
res = _run(STUDIO_UPDATE, [], env)
|
|
calls = Path(env["STUB_LOG"]).read_text() if Path(env["STUB_LOG"]).exists() else ""
|
|
assert "STUB-SUPERVISORCTL restart studio" not in calls, (
|
|
"restarting into code that cannot import kills a process that is serving "
|
|
"fine and parks supervisord's studio program in FATAL:\n" + calls
|
|
)
|
|
assert res.returncode != 0, "a broken update must not report success"
|
|
assert "--with-deps" in res.stderr, "the remedy must still be printed"
|
|
|
|
|
|
# --- unsloth-llama-update -----------------------------------------------------
|
|
|
|
|
|
def _llama_env(tmp_path: Path, *, latest: str | None) -> dict:
|
|
install = tmp_path / "llama.cpp"
|
|
install.mkdir(parents = True)
|
|
(install / "UNSLOTH_PREBUILT_INFO.json").write_text(
|
|
'{"tag": "b1111-old"}\n',
|
|
encoding = "utf-8",
|
|
)
|
|
fetcher = tmp_path / "fetch_llama_prebuilt.py"
|
|
resolve = (
|
|
" raise RuntimeError('unreachable')\n" if latest is None else f" return {latest!r}\n"
|
|
)
|
|
fetcher.write_text(
|
|
"def resolve_latest_tag(repo):\n" + resolve,
|
|
encoding = "utf-8",
|
|
)
|
|
env = dict(os.environ)
|
|
env["UNSLOTH_LLAMA_CPP_PATH"] = str(install)
|
|
env["UNSLOTH_LLAMA_FETCHER"] = str(fetcher)
|
|
return env
|
|
|
|
|
|
def _llama_check(tmp_path: Path, latest):
|
|
env = _llama_env(tmp_path, latest = latest)
|
|
return _run(LLAMA_UPDATE, ["--check"], env)
|
|
|
|
|
|
def test_llama_check_reports_an_available_update(tmp_path: Path):
|
|
res = _llama_check(tmp_path, "b2222-new")
|
|
assert res.returncode == 0, res.stderr
|
|
assert "an update is available" in res.stdout
|
|
|
|
|
|
def test_llama_check_reports_up_to_date(tmp_path: Path):
|
|
res = _llama_check(tmp_path, "b1111-old")
|
|
assert res.returncode == 0, res.stderr
|
|
assert "up to date" in res.stdout
|
|
|
|
|
|
def test_llama_check_does_not_claim_up_to_date_when_it_could_not_look(tmp_path: Path):
|
|
res = _llama_check(tmp_path, None)
|
|
assert "up to date" not in res.stdout, (
|
|
"--check exists to report update status; saying 'up to date' for a lookup "
|
|
"that never happened is the one answer it must never give:\n" + res.stdout
|
|
)
|
|
assert res.returncode != 0, "an unperformed check must not exit 0"
|
|
assert "UNKNOWN" in res.stdout + res.stderr
|
|
|
|
|
|
def _llama_inplace_env(tmp_path: Path, old: list[str], new: list[str]) -> dict:
|
|
"""An in-place (volume-mounted) install whose activation fails part-way."""
|
|
install = tmp_path / "llama.cpp"
|
|
install.mkdir(parents = True)
|
|
for name in old:
|
|
(install / name).write_text("OLD\n", encoding = "utf-8")
|
|
(install / "UNSLOTH_PREBUILT_INFO.json").write_text(
|
|
'{"tag": "b1111-old"}\n',
|
|
encoding = "utf-8",
|
|
)
|
|
fetcher = tmp_path / "fetch_llama_prebuilt.py"
|
|
fetcher.write_text(
|
|
"import os, sys\n"
|
|
"def resolve_latest_tag(repo):\n"
|
|
" return 'b2222-new'\n"
|
|
"if __name__ == '__main__':\n"
|
|
" dest = sys.argv[3]\n"
|
|
" os.makedirs(dest, exist_ok = True)\n"
|
|
f" for name in {new!r}:\n"
|
|
" open(os.path.join(dest, name), 'w').write('NEW\\n')\n"
|
|
" open(os.path.join(dest, 'UNSLOTH_PREBUILT_INFO.json'), 'w')"
|
|
'.write(\'{"tag": "b2222-new"}\\n\')\n',
|
|
encoding = "utf-8",
|
|
)
|
|
# Fail the ACTIVATION move (-t <install dir>) AFTER it has moved the files, so
|
|
# the install dir is populated with the new tree and `find` still reports the
|
|
# failure -- the mid-swap abort the rollback exists for. The drain
|
|
# (-t <backup>) and the rollback's own per-file moves must keep working, so
|
|
# only that one invocation is broken.
|
|
bin_dir = tmp_path / "bin"
|
|
_stub(
|
|
bin_dir,
|
|
"mv",
|
|
'if [ "$1" = "-t" ] && [ "$2" = "$FAIL_MV_TARGET" ]; then\n'
|
|
" shift 2\n"
|
|
' for _s in "$@"; do /bin/mv "$_s" "$FAIL_MV_TARGET/"; done\n'
|
|
" exit 1\n"
|
|
"fi\n"
|
|
'exec /bin/mv "$@"\n',
|
|
)
|
|
env = dict(os.environ)
|
|
env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"]
|
|
env["UNSLOTH_LLAMA_CPP_PATH"] = str(install)
|
|
env["UNSLOTH_LLAMA_FETCHER"] = str(fetcher)
|
|
env["UNSLOTH_LLAMA_UPDATE_IN_PLACE"] = "1"
|
|
env["FAIL_MV_TARGET"] = str(install)
|
|
return env
|
|
|
|
|
|
def test_llama_rollback_leaves_no_new_release_files_behind(tmp_path: Path):
|
|
# "libggml-hexagon.so" exists only in the new release, so the rollback loop --
|
|
# which iterates the BACKUP's entries -- cannot see it. ggml dlopen()s every
|
|
# libggml-*.so sitting next to the binaries, so a leftover is loaded against
|
|
# the restored older libggml-base.so.
|
|
old = ["libggml-base.so", "libggml-cpu-icelake.so", "llama-cli"]
|
|
new = [
|
|
"libggml-base.so",
|
|
"libggml-cpu-icelake.so",
|
|
"llama-cli",
|
|
"libggml-hexagon.so",
|
|
"llama-mtmd-cli",
|
|
]
|
|
env = _llama_inplace_env(tmp_path, old, new)
|
|
res = _run(LLAMA_UPDATE, [], env)
|
|
assert res.returncode != 0, "a failed swap must not report success"
|
|
install = tmp_path / "llama.cpp"
|
|
present = sorted(p.name for p in install.iterdir())
|
|
leftovers = [n for n in ("libggml-hexagon.so", "llama-mtmd-cli") if n in present]
|
|
assert not leftovers, f"new-release-only files survived the rollback: {leftovers} in {present}"
|
|
for name in old:
|
|
assert (
|
|
install / name
|
|
).read_text() == "OLD\n", f"{name} was not restored from the backup: {present}"
|
|
|
|
|
|
def test_llama_rollback_keeps_every_old_file_when_the_drain_is_interrupted(tmp_path: Path):
|
|
# The mirror image: abort while the OLD tree is still being moved into the
|
|
# backup. The entries left in the install dir are then the only copy of those
|
|
# old files, so clearing the directory before restoring would destroy them.
|
|
old = ["libggml-base.so", "libggml-cpu-icelake.so", "llama-cli", "llama-quantize"]
|
|
env = _llama_inplace_env(tmp_path, old, old)
|
|
install = tmp_path / "llama.cpp"
|
|
# Fail the DRAIN (-t <install dir>/.old.<pid>) after moving only the first
|
|
# source, so half the old tree is still sitting in the install dir when the
|
|
# rollback runs. Those entries are then the only copy there is.
|
|
_stub(
|
|
tmp_path / "bin",
|
|
"mv",
|
|
'case "${1:-}:${2:-}" in\n'
|
|
" -t:*/.old.*)\n"
|
|
' _t="$2"; shift 2\n'
|
|
' [ $# -gt 0 ] && /bin/mv "$1" "$_t/"\n'
|
|
" exit 1;;\n"
|
|
"esac\n"
|
|
'exec /bin/mv "$@"\n',
|
|
)
|
|
res = _run(LLAMA_UPDATE, [], env)
|
|
assert res.returncode != 0
|
|
survivors = sorted(p.name for p in install.rglob("*") if p.is_file())
|
|
for name in old:
|
|
assert name in survivors, f"{name} was lost during an interrupted drain: {survivors}"
|