#!/usr/bin/env python3 """Install a prepared Python profile offline into an isolated target with rollback.""" from __future__ import annotations import argparse import importlib.util import os import shutil import stat import subprocess import tempfile from pathlib import Path HERE = Path(__file__).resolve().parent def load_module(name: str, filename: str): spec = importlib.util.spec_from_file_location(name, HERE / filename) module = importlib.util.module_from_spec(spec) assert spec.loader spec.loader.exec_module(module) return module journal_api = load_module("install_profile_journal", "atomic-profile-journal.py") checkpoint_api = load_module("install_profile_checkpoint", "python-environment-checkpoint.py") def fault(point: str) -> None: if os.environ.get("NEXUS_ALLOW_TEST_FAULTS") == "1" \ and os.environ.get("NEXUS_TEST_CRASH_AT") == point: os._exit(97) def filtered_requirements(source: Path, profile: str, destination: Path) -> None: excluded = {"chromadb"} if profile == "cpu-only": excluded.add("bitsandbytes") lines = [] for raw in source.read_text().splitlines(): stripped = raw.split("#", 1)[0].strip() name = stripped.split("=", 1)[0].strip().lower().replace("_", "-") if stripped and name not in excluded: lines.append(stripped) destination.write_text("\n".join(lines) + "\n") def health(python: Path, profile: str) -> None: base = "import torch,torchaudio,torchvision,transformers,peft,accelerate,sentence_transformers;" if profile == "cpu-only": check = (base + "assert torch.__version__=='2.8.0+cpu';assert torch.version.cuda is None;" "assert not torch.cuda.is_available();assert (torch.tensor([2.0])*2).item()==4.0") else: check = (base + "import bitsandbytes;assert torch.__version__.startswith('2.8.0');" "assert torch.version.cuda=='12.8';assert torch.cuda.is_available();" "x=torch.tensor([2.0],device='cuda');assert (x*2).item()==4.0;torch.cuda.synchronize()") subprocess.run([str(python), "-c", check], check=True) def validate_service_probe(probe: Path | None) -> None: if probe is None: return if not probe.is_absolute(): raise ValueError("SERVICE_HEALTH_PROBE_NOT_ABSOLUTE") if probe.is_symlink() or not probe.is_file(): raise ValueError("SERVICE_HEALTH_PROBE_NOT_REGULAR") if not stat.S_IMODE(probe.stat().st_mode) & 0o111: raise ValueError("SERVICE_HEALTH_PROBE_NOT_EXECUTABLE") def service_health(probe: Path | None, target: Path, profile: str, timeout_seconds: int) -> None: validate_service_probe(probe) if probe is not None: subprocess.run([str(probe), str(target), profile], check=True, timeout=timeout_seconds) def install(journal_path: Path, checkpoint: Path, target: Path, wheel_dir: Path, requirements: Path, python_executable: Path, service_health_probe: Path | None = None, service_health_timeout_seconds: int = 300) -> dict: if service_health_timeout_seconds < 1 or service_health_timeout_seconds > 3600: raise ValueError("SERVICE_HEALTH_TIMEOUT_INVALID") validate_service_probe(service_health_probe) current = journal_api.read_verified(journal_path) if not current or current["state"] != "checkpoint_ready": raise ValueError("JOURNAL_NOT_CHECKPOINT_READY") checkpoint_api.verify(checkpoint, target) identity = dict(operation_id=current["operation_id"], profile=current["selected_profile"], manifest_sha256=current["manifest_sha256"]) current = journal_api.update(journal_path, expected_generation=current["generation"], state="installing", checkpoint_id=current["checkpoint_id"], **identity) temporary = Path(tempfile.mkdtemp(prefix=f".{target.name}.install.", dir=target.parent)) shutil.rmtree(temporary) activated = False req: Path | None = None try: subprocess.run([str(python_executable), "-m", "venv", str(temporary)], check=True) req = temporary.parent / f".{temporary.name}.requirements.txt" filtered_requirements(requirements, identity["profile"], req) subprocess.run([str(temporary / "bin/pip"), "install", "--no-index", f"--find-links={wheel_dir}", "-r", str(req)], check=True) req.unlink() health(temporary / "bin/python", identity["profile"]) checkpoint_api.fsync_tree(temporary) previous = target.with_name(f".{target.name}.install-previous") if previous.exists(): raise ValueError("INSTALL_PREVIOUS_EXISTS") if target.exists(): os.replace(target, previous) os.replace(temporary, target) activated = True directory_fd = os.open(target.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) fault("after_activation") health(target / "bin/python", identity["profile"]) service_health(service_health_probe, target, identity["profile"], service_health_timeout_seconds) current = journal_api.update(journal_path, expected_generation=current["generation"], state="installed", checkpoint_id=current["checkpoint_id"], **identity) return {"installed": True, "generation": current["generation"], "selected_profile": identity["profile"]} except Exception: current = journal_api.update(journal_path, expected_generation=current["generation"], state="rollback_required", checkpoint_id=current["checkpoint_id"], **identity) if activated: checkpoint_api.restore(checkpoint, target) else: checkpoint_api.verify(checkpoint, target) journal_api.update(journal_path, expected_generation=current["generation"], state="rolled_back", checkpoint_id=current["checkpoint_id"], **identity) raise finally: if temporary.exists(): shutil.rmtree(temporary) if req is not None and req.exists(): req.unlink() def recover(journal_path: Path, checkpoint: Path, target: Path) -> dict: current = journal_api.read_verified(journal_path) if not current or current["state"] not in {"installing", "rollback_required"}: raise ValueError("JOURNAL_NOT_RECOVERABLE") identity = dict(operation_id=current["operation_id"], profile=current["selected_profile"], manifest_sha256=current["manifest_sha256"]) if current["state"] == "installing": current = journal_api.update(journal_path, expected_generation=current["generation"], state="rollback_required", checkpoint_id=current["checkpoint_id"], **identity) manifest = checkpoint_api.verify(checkpoint, target) if target.exists() and checkpoint_api.inventory(target) != manifest["entries"]: checkpoint_api.restore(checkpoint, target) else: checkpoint_api.recover_restore(checkpoint, target) current = journal_api.update(journal_path, expected_generation=current["generation"], state="rolled_back", checkpoint_id=current["checkpoint_id"], **identity) return {"recovered": True, "generation": current["generation"], "state": current["state"]} def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--journal", type=Path, required=True) parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--target", type=Path, required=True) parser.add_argument("--wheel-dir", type=Path, required=True) parser.add_argument("--requirements", type=Path, required=True) parser.add_argument("--python", type=Path, required=True) parser.add_argument("--service-health-probe", type=Path) parser.add_argument("--service-health-timeout-seconds", type=int, default=300) parser.add_argument("--recover-only", action="store_true") args = parser.parse_args() if args.recover_only: result = recover(args.journal, args.checkpoint, args.target) else: result = install(args.journal, args.checkpoint, args.target, args.wheel_dir, args.requirements, args.python, args.service_health_probe, args.service_health_timeout_seconds) print(result) return 0 if __name__ == "__main__": raise SystemExit(main())