| """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"))) |
|
|
