115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Install, inspect, or roll back the verified DeepSeek V4 PPU source patch."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import hashlib
|
||
|
|
import importlib.util
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import stat
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parents[1]
|
||
|
|
BACKUP_SUFFIX = ".ppu-dsv4.orig"
|
||
|
|
|
||
|
|
|
||
|
|
def digest(data: bytes) -> str:
|
||
|
|
return hashlib.sha256(data).hexdigest()
|
||
|
|
|
||
|
|
|
||
|
|
def atomic_write(path: Path, data: bytes, mode: int) -> None:
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
|
||
|
|
try:
|
||
|
|
with os.fdopen(fd, "wb") as output:
|
||
|
|
output.write(data)
|
||
|
|
output.flush()
|
||
|
|
os.fsync(output.fileno())
|
||
|
|
os.chmod(temporary, mode)
|
||
|
|
os.replace(temporary, path)
|
||
|
|
finally:
|
||
|
|
if os.path.exists(temporary):
|
||
|
|
os.unlink(temporary)
|
||
|
|
|
||
|
|
|
||
|
|
def site_packages_path(override: str | None) -> Path:
|
||
|
|
if override:
|
||
|
|
return Path(override).resolve()
|
||
|
|
spec = importlib.util.find_spec("vllm")
|
||
|
|
if spec is None or not spec.submodule_search_locations:
|
||
|
|
raise RuntimeError("Cannot locate installed vllm; pass --site-packages")
|
||
|
|
return Path(next(iter(spec.submodule_search_locations))).resolve().parent
|
||
|
|
|
||
|
|
|
||
|
|
def run() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
group = parser.add_mutually_exclusive_group(required=True)
|
||
|
|
group.add_argument("--check", action="store_true")
|
||
|
|
group.add_argument("--apply", action="store_true")
|
||
|
|
group.add_argument("--rollback", action="store_true")
|
||
|
|
parser.add_argument("--site-packages", help="Override site-packages root for testing")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
site = site_packages_path(args.site_packages)
|
||
|
|
manifest = json.loads((ROOT / "manifest.json").read_text())
|
||
|
|
if manifest.get("format") != 1:
|
||
|
|
raise RuntimeError("Unknown manifest format")
|
||
|
|
|
||
|
|
planned = []
|
||
|
|
for name, expected in manifest["files"].items():
|
||
|
|
relative = Path(name)
|
||
|
|
if relative.is_absolute() or ".." in relative.parts or relative.parts[0] != "vllm":
|
||
|
|
raise RuntimeError(f"Unsafe manifest path: {name}")
|
||
|
|
target = site / relative
|
||
|
|
patched = (ROOT / "src" / relative).read_bytes()
|
||
|
|
original = (ROOT / "upstream" / relative).read_bytes()
|
||
|
|
if digest(patched) != expected["patched_sha256"]:
|
||
|
|
raise RuntimeError(f"Packaged patched source changed: {name}")
|
||
|
|
if digest(original) != expected["original_sha256"]:
|
||
|
|
raise RuntimeError(f"Packaged original source changed: {name}")
|
||
|
|
current = target.read_bytes()
|
||
|
|
current_hash = digest(current)
|
||
|
|
if current_hash == expected["original_sha256"]:
|
||
|
|
state = "original"
|
||
|
|
elif current_hash == expected["patched_sha256"]:
|
||
|
|
state = "patched"
|
||
|
|
else:
|
||
|
|
raise RuntimeError(f"Installed source does not match this patch: {target} ({current_hash})")
|
||
|
|
backup = target.with_name(target.name + BACKUP_SUFFIX)
|
||
|
|
if backup.exists() and digest(backup.read_bytes()) != expected["original_sha256"]:
|
||
|
|
raise RuntimeError(f"Existing backup does not match original source: {backup}")
|
||
|
|
planned.append((target, backup, current, original, patched, state))
|
||
|
|
|
||
|
|
if args.check:
|
||
|
|
for target, _, _, _, _, state in planned:
|
||
|
|
print(f"{state:8} {target}")
|
||
|
|
return
|
||
|
|
|
||
|
|
# All source and backup hashes have been validated before the first write.
|
||
|
|
for target, backup, current, original, patched, state in planned:
|
||
|
|
mode = stat.S_IMODE(target.stat().st_mode)
|
||
|
|
if args.apply and state == "original":
|
||
|
|
if not backup.exists():
|
||
|
|
atomic_write(backup, current, mode)
|
||
|
|
atomic_write(target, patched, mode)
|
||
|
|
print(f"PATCHED {target}")
|
||
|
|
elif args.rollback and state == "patched":
|
||
|
|
# The packaged original is identical to the verified backup.
|
||
|
|
atomic_write(target, original, mode)
|
||
|
|
print(f"RESTORED {target}")
|
||
|
|
else:
|
||
|
|
print(f"UNCHANGED {target}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
try:
|
||
|
|
run()
|
||
|
|
except (OSError, ValueError, RuntimeError) as error:
|
||
|
|
print(f"ERROR: {error}", file=sys.stderr)
|
||
|
|
sys.exit(1)
|