| |
| """Read-only process/model accounting, for diagnostics (not deployment runtime).""" |
| import argparse |
| import collections |
| import json |
| from pathlib import Path |
| import re |
| import struct |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--pid", type=int) |
| p.add_argument("--model", type=Path) |
| a = p.parse_args() |
| result = {} |
| if a.model: |
| categories = collections.Counter() |
| with a.model.open("rb") as f: |
| header = f.read(256) |
| count = struct.unpack_from("<I", header, 20)[0] |
| offset = struct.unpack_from("<Q", header, 32)[0] |
| f.seek(offset) |
| for _ in range(count): |
| entry = f.read(160) |
| dtype, role = struct.unpack_from("<I", entry, 8)[0], struct.unpack_from("<I", entry, 16)[0] |
| data_bytes = struct.unpack_from("<Q", entry, 56)[0] |
| categories[f"dtype{dtype}_role{role}"] += data_bytes |
| result["model_mib"] = {k: v / 1048576 for k, v in categories.items()} |
| result["model_file_mib"] = a.model.stat().st_size / 1048576 |
| if a.pid: |
| categories = collections.defaultdict(collections.Counter) |
| counts = collections.Counter() |
| category = "unknown" |
| for line in Path(f"/proc/{a.pid}/smaps").read_text().splitlines(): |
| if re.match(r"^[0-9a-f]+-[0-9a-f]+ ", line): |
| parts = line.split(maxsplit=5) |
| path = parts[5] if len(parts) == 6 else "" |
| category = "model" if path.endswith(".l3r") else "drm" if "/dev/dri/" in path else ( |
| "anonymous" if not path or path.startswith("[") else "libraries_and_other_files") |
| counts[category] += 1 |
| elif re.match(r"^(Size|Rss|Pss|Private_Dirty|Shared_Dirty|Private_Clean|Shared_Clean|Swap):", line): |
| k, v, *_ = line.split() |
| categories[category][k[:-1]] += int(v) / 1024 |
| result["smaps_mib"] = dict(categories) |
| result["mapping_counts"] = dict(counts) |
| result["status"] = {line.split(":", 1)[0]: line.split(":", 1)[1].strip() |
| for line in Path(f"/proc/{a.pid}/status").read_text().splitlines() |
| if line.startswith(("VmRSS:", "VmHWM:", "VmSwap:", "Threads:"))} |
| print(json.dumps(result, indent=2)) |
|
|
