| |
| """Parse a DCPM metadata.bin index. Stage 1: read only, write nothing. |
| |
| Record layout, decoded from the bytes: |
| int64 name_len | name | int64 ndim | int64 shape[ndim] |
| int32 dtype | int32 ? | int64 ? | int32 ? | int64 file_len | file |
| int64 offset | ... trailing fields |
| We locate records by scanning for the length-prefixed names rather than |
| assuming a fixed stride, so a variable-length shape list cannot desync us. |
| """ |
| import struct, sys, re, json |
|
|
| def parse(path): |
| b = open(path, "rb").read() |
| assert b[:4] == b"DCPM", f"not DCPM: {b[:4]!r}" |
| recs = [] |
| for m in re.finditer(rb"param_(\d{6})", b): |
| i = m.start() |
| nlen = struct.unpack_from("<q", b, i - 8)[0] |
| if nlen != m.end() - i: |
| continue |
| j = m.end() |
| ndim = struct.unpack_from("<q", b, j)[0]; j += 8 |
| if not (0 <= ndim <= 8): continue |
| shape = list(struct.unpack_from("<" + "q" * ndim, b, j)); j += 8 * ndim |
| dtype = struct.unpack_from("<i", b, j)[0]; j += 4 |
| f1 = struct.unpack_from("<i", b, j)[0]; j += 4 |
| f2 = struct.unpack_from("<q", b, j)[0]; j += 8 |
| f3 = struct.unpack_from("<i", b, j)[0]; j += 4 |
| flen = struct.unpack_from("<q", b, j)[0]; j += 8 |
| fname = b[j:j + flen].decode(errors="replace"); j += flen |
| off = struct.unpack_from("<q", b, j)[0]; j += 8 |
| recs.append(dict(idx=int(m.group(1)), name=m.group(0).decode(), ndim=ndim, |
| shape=shape, dtype=dtype, f1=f1, f2=f2, f3=f3, |
| file=fname, offset=off)) |
| recs.sort(key=lambda r: r["idx"]) |
| return recs |
|
|
| if __name__ == "__main__": |
| recs = parse(sys.argv[1]) |
| print(f" parsed {len(recs)} param records") |
| dts = {} |
| for r in recs: dts.setdefault(r["dtype"], 0); dts[r["dtype"]] += 1 |
| print(f" dtype codes seen: {dts}") |
| files = {} |
| for r in recs: files.setdefault(r["file"], 0); files[r["file"]] += 1 |
| print(f" source files: {files}") |
| print(f"\n {'idx':>4} {'shape':<20} {'dtype':>5} {'file':<16} {'offset':>14}") |
| for r in recs[:14]: |
| print(f" {r['idx']:>4} {str(r['shape']):<20} {r['dtype']:>5} {r['file']:<16} {r['offset']:>14,}") |
| print(" ...") |
| for r in recs[-6:]: |
| print(f" {r['idx']:>4} {str(r['shape']):<20} {r['dtype']:>5} {r['file']:<16} {r['offset']:>14,}") |
| tot = sum(int(__import__("math").prod(r["shape"])) for r in recs) |
| print(f"\n total elements across all params: {tot:,}") |
| json.dump(recs, open("/tmp/dcpm_index.json", "w")) |
| print(" wrote /tmp/dcpm_index.json") |
|
|
