From ebaae02b976dde97278523f081a2d59447837dfa Mon Sep 17 00:00:00 2001 From: baketnk Date: Thu, 24 Sep 2026 22:49:18 -0400 Subject: [PATCH] fix(installer): unwind owned model downloads on cancellation --- install.sh | 68 +++++++++++++++++++++++--------------- scripts/install_payload.py | 68 +++++++++++++++++++++++--------------- tests/test_installer.py | 56 +++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 52 deletions(-) diff --git a/install.sh b/install.sh index dc6ecd3..00002ff 100755 --- a/install.sh +++ b/install.sh @@ -32,6 +32,7 @@ from pathlib import Path import platform import re import shlex +import signal import shutil import subprocess import sys @@ -535,6 +536,20 @@ def model_target(args, root): return root / "models" / args.backend +@contextlib.contextmanager +def model_download_cancellation(): + # Install-model runs on the main thread. Raising unwinds the owned temp's + # finally block; unlike SIGKILL this does not leave a partial download. + def terminate(_signal, _frame): + raise ValueError("model installation cancelled") + + previous = signal.signal(signal.SIGTERM, terminate) + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) + + def install_model(args, root): module, backend, manifest_sha = installed_backend(root, args.backend) check_expected_manifest(args, manifest_sha) # under .model.lock, before model mkdir/network @@ -559,32 +574,33 @@ def install_model(args, root): url = f"{backend.source}/resolve/{quote(backend.revision, safe='')}/{quote(item.path, safe='/')}" if args.json: print(json.dumps({"ok": True, "event": "model_file", "file": item.path, "state": "downloading", "bytes": item.size}), flush=True) - fd, temp = tempfile.mkstemp(prefix=".download-", dir=target.parent) - try: - with os.fdopen(fd, "wb") as output: - # The shared verifier is used for both pre-existing and downloaded files. - request = urllib.request.Request(url, headers={"User-Agent": "frameyap-installer"}) - with urllib.request.urlopen(request, timeout=60) as response: - if response.geturl().split(":", 1)[0] != "https": - fail("model URL redirected away from HTTPS") - total = 0 - while chunk := response.read(1024 * 1024): - total += len(chunk) - if total > item.size: - fail(f"model download exceeded pinned size: {item.path}") - output.write(chunk) - output.flush() - os.fsync(output.fileno()) - temporary = type(item)(Path(temp).name, item.size, item.sha256) - if module.check_file(target.parent, temporary)[0] is not None: - fail(f"pinned SHA-256/size mismatch: {item.path}") - if target.exists() or target.is_symlink(): - fail(f"model destination changed during download: {target}") - os.replace(temp, target) - if args.json: - print(json.dumps({"ok": True, "event": "model_file", "file": item.path, "state": "verified"}), flush=True) - finally: - Path(temp).unlink(missing_ok=True) + with model_download_cancellation(): + fd, temp = tempfile.mkstemp(prefix=".download-", dir=target.parent) + try: + with os.fdopen(fd, "wb") as output: + # The shared verifier is used for both pre-existing and downloaded files. + request = urllib.request.Request(url, headers={"User-Agent": "frameyap-installer"}) + with urllib.request.urlopen(request, timeout=60) as response: + if response.geturl().split(":", 1)[0] != "https": + fail("model URL redirected away from HTTPS") + total = 0 + while chunk := response.read(1024 * 1024): + total += len(chunk) + if total > item.size: + fail(f"model download exceeded pinned size: {item.path}") + output.write(chunk) + output.flush() + os.fsync(output.fileno()) + temporary = type(item)(Path(temp).name, item.size, item.sha256) + if module.check_file(target.parent, temporary)[0] is not None: + fail(f"pinned SHA-256/size mismatch: {item.path}") + if target.exists() or target.is_symlink(): + fail(f"model destination changed during download: {target}") + os.replace(temp, target) + if args.json: + print(json.dumps({"ok": True, "event": "model_file", "file": item.path, "state": "verified"}), flush=True) + finally: + Path(temp).unlink(missing_ok=True) state = module.check_model(backend, dest) if state["state"] != "installed_verified": fail(f"model did not verify: {state['reason']}") diff --git a/scripts/install_payload.py b/scripts/install_payload.py index ac3c990..d47de03 100644 --- a/scripts/install_payload.py +++ b/scripts/install_payload.py @@ -12,6 +12,7 @@ from pathlib import Path import platform import re import shlex +import signal import shutil import subprocess import sys @@ -515,6 +516,20 @@ def model_target(args, root): return root / "models" / args.backend +@contextlib.contextmanager +def model_download_cancellation(): + # Install-model runs on the main thread. Raising unwinds the owned temp's + # finally block; unlike SIGKILL this does not leave a partial download. + def terminate(_signal, _frame): + raise ValueError("model installation cancelled") + + previous = signal.signal(signal.SIGTERM, terminate) + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) + + def install_model(args, root): module, backend, manifest_sha = installed_backend(root, args.backend) check_expected_manifest(args, manifest_sha) # under .model.lock, before model mkdir/network @@ -539,32 +554,33 @@ def install_model(args, root): url = f"{backend.source}/resolve/{quote(backend.revision, safe='')}/{quote(item.path, safe='/')}" if args.json: print(json.dumps({"ok": True, "event": "model_file", "file": item.path, "state": "downloading", "bytes": item.size}), flush=True) - fd, temp = tempfile.mkstemp(prefix=".download-", dir=target.parent) - try: - with os.fdopen(fd, "wb") as output: - # The shared verifier is used for both pre-existing and downloaded files. - request = urllib.request.Request(url, headers={"User-Agent": "frameyap-installer"}) - with urllib.request.urlopen(request, timeout=60) as response: - if response.geturl().split(":", 1)[0] != "https": - fail("model URL redirected away from HTTPS") - total = 0 - while chunk := response.read(1024 * 1024): - total += len(chunk) - if total > item.size: - fail(f"model download exceeded pinned size: {item.path}") - output.write(chunk) - output.flush() - os.fsync(output.fileno()) - temporary = type(item)(Path(temp).name, item.size, item.sha256) - if module.check_file(target.parent, temporary)[0] is not None: - fail(f"pinned SHA-256/size mismatch: {item.path}") - if target.exists() or target.is_symlink(): - fail(f"model destination changed during download: {target}") - os.replace(temp, target) - if args.json: - print(json.dumps({"ok": True, "event": "model_file", "file": item.path, "state": "verified"}), flush=True) - finally: - Path(temp).unlink(missing_ok=True) + with model_download_cancellation(): + fd, temp = tempfile.mkstemp(prefix=".download-", dir=target.parent) + try: + with os.fdopen(fd, "wb") as output: + # The shared verifier is used for both pre-existing and downloaded files. + request = urllib.request.Request(url, headers={"User-Agent": "frameyap-installer"}) + with urllib.request.urlopen(request, timeout=60) as response: + if response.geturl().split(":", 1)[0] != "https": + fail("model URL redirected away from HTTPS") + total = 0 + while chunk := response.read(1024 * 1024): + total += len(chunk) + if total > item.size: + fail(f"model download exceeded pinned size: {item.path}") + output.write(chunk) + output.flush() + os.fsync(output.fileno()) + temporary = type(item)(Path(temp).name, item.size, item.sha256) + if module.check_file(target.parent, temporary)[0] is not None: + fail(f"pinned SHA-256/size mismatch: {item.path}") + if target.exists() or target.is_symlink(): + fail(f"model destination changed during download: {target}") + os.replace(temp, target) + if args.json: + print(json.dumps({"ok": True, "event": "model_file", "file": item.path, "state": "verified"}), flush=True) + finally: + Path(temp).unlink(missing_ok=True) state = module.check_model(backend, dest) if state["state"] != "installed_verified": fail(f"model did not verify: {state['reason']}") diff --git a/tests/test_installer.py b/tests/test_installer.py index ce2f4b9..1867674 100644 --- a/tests/test_installer.py +++ b/tests/test_installer.py @@ -16,6 +16,7 @@ from unittest.mock import patch import shutil import pty import select +import signal from types import SimpleNamespace REPO = Path(__file__).resolve().parents[1] @@ -850,6 +851,61 @@ class InstallTests(unittest.TestCase): fetch.assert_not_called() self.assertFalse((root / "models/other").exists()) + def test_model_sigterm_removes_owned_partial_download(self): + # The child uses a mocked blocking response. No network/download occurs. + shutil.copyfile(REPO / "python/frameyap/model_files.py", + self.stage / "python/frameyap/model_files.py") + manifest = json.loads((REPO / "assets/backends/redux.json").read_text()) + manifest["model"]["files"] = [{"path": "weights.bin", "size": 9, + "sha256": hashlib.sha256(b"ninebytes").hexdigest()}] + manifests = self.stage / "assets/backends" + manifests.mkdir() + (manifests / "redux.json").write_text(json.dumps(manifest)) + archive, digest = self.package("0.1.202609241530") + self.install("0.1.202609241530", archive, digest, "--without-model") + dest = self.base / "cancelled-model" + script = '''import importlib.util, signal, sys +from unittest.mock import patch +spec = importlib.util.spec_from_file_location("installer", sys.argv[1]) +module = importlib.util.module_from_spec(spec) +spec.loader.exec_module(module) +class Response: + def __enter__(self): return self + def __exit__(self, *args): pass + def geturl(self): return "https://example.org/model" + def read(self, size): + print("fixture-read-started", flush=True) + signal.pause() # parent delivers SIGTERM during the owned download +with patch.object(module, "check_host"), patch.object(module.urllib.request, "urlopen", return_value=Response()): + sys.exit(module.cli(["--install-model", "--backend", "redux", "--model-dir", sys.argv[2], "--yes", "--json"])) +''' + child = subprocess.Popen([sys.executable, "-c", script, + str(REPO / "scripts/install_payload.py"), str(dest)], + stdout=subprocess.PIPE, stderr=subprocess.PIPE, + env={**os.environ, "PYTHONDONTWRITEBYTECODE": "1"}) + try: + output = b"" + for _ in range(8): + readable, _, _ = select.select([child.stdout], [], [], 8) + self.assertTrue(readable, "mocked download did not reach response.read") + output += os.read(child.stdout.fileno(), 4096) + if b"fixture-read-started\n" in output: + break + self.assertIn(b"fixture-read-started\n", output) + self.assertEqual(len(list(dest.glob(".download-*"))), 1) + os.kill(child.pid, signal.SIGTERM) # only our fixture process + rest, error = child.communicate(timeout=5) + self.assertEqual(child.returncode, 1, error) + event = json.loads(rest.strip()) + self.assertFalse(event["ok"]) + self.assertIn("cancelled", event["message"]) + self.assertEqual(list(dest.glob(".download-*")), []) + self.assertFalse((dest / "weights.bin").exists()) + finally: + if child.poll() is None: + child.kill() + child.wait(timeout=5) + def test_model_install_manifest_verification_and_explicit_consent(self): # Tiny pinned files via the shared verifier; urllib is mocked, no network. source_module = REPO / "python/frameyap/model_files.py"