189 lines
8.5 KiB
Python
189 lines
8.5 KiB
Python
#!/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())
|