File size: 2,227 Bytes
0693474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """Read-only Spark inventory; safetensors bytes always bypass the page cache."""
from __future__ import annotations
import json
import hashlib
import mmap
import os
import struct
from pathlib import Path
def direct_read(fd: int, offset: int, size: int) -> bytes:
start = offset // 4096 * 4096
length = ((offset - start + size + 4095) // 4096) * 4096
with mmap.mmap(-1, length) as buffer:
count = os.preadv(fd, [buffer], start)
if count < offset - start + size:
raise EOFError(f"short O_DIRECT read: {offset=} {size=} {count=}")
return buffer[offset - start:offset - start + size]
def inventory(root: Path, baseline: Path, manifest: Path) -> dict[str, object]:
files: dict[str, object] = {}
for path in sorted(root.glob("*.safetensors")):
fd = os.open(path, os.O_RDONLY | os.O_DIRECT)
try:
length = struct.unpack("<Q", direct_read(fd, 0, 8))[0]
if length > 64 * 1024 * 1024:
raise ValueError(f"unreasonable header: {path} {length}")
header = json.loads(direct_read(fd, 8, length))
finally:
os.close(fd)
stat = path.stat()
base = baseline / path.name
files[path.name] = {
"size": stat.st_size, "device": stat.st_dev, "inode": stat.st_ino,
"nlink": stat.st_nlink, "header_bytes": length, "header": header,
"baseline_same_inode": base.exists() and os.path.samefile(path, base),
}
metadata = {name: json.loads((root / name).read_text()) for name in (
"config.json", "exl3_moe_plan.json", "CANDIDATE_MANIFEST.json",
"model.safetensors.index.json")}
return {"source_root": str(root), "baseline_root": str(baseline),
"files": files, "metadata": metadata,
"metadata_sha256": {name: hashlib.sha256((root/name).read_bytes()).hexdigest()
for name in (*metadata, "tokenizer.json", "tokenizer_config.json")},
"manifest": json.loads(manifest.read_text())}
if __name__ == "__main__":
print(json.dumps(inventory(Path("/srv/dsv41/legal-v2a"), Path("/srv/dsv41/model"),
Path.home() / "dsv41-deploy/manifest.json")))
|