File size: 2,186 Bytes
3fd1a35 | 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 | #!/usr/bin/env python3
"""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))
|