"""Trace one arm's GPU memory against the requests it had served: when did the arena grow, and after what?
usage: memtrace.py ARM   (reads out/raw/pids.json, mem.csv, first/lat/conc jsonl)"""
import datetime, json, sys
AB = "/tank/spikes/parakeet-ab"
arm = sys.argv[1]
pids = set(map(str, json.load(open(f"{AB}/out/raw/pids.json"))[arm]))
series = []
for l in open(f"{AB}/out/raw/mem.csv"):
    p = [x.strip() for x in l.split(",")]
    if len(p) == 4 and p[2] in pids:
        series.append((datetime.datetime.strptime(p[0], "%Y/%m/%d %H:%M:%S.%f").timestamp(), int(p[3])))
reqs = []
for f in ("first", "lat", "conc"):
    for l in open(f"{AB}/out/raw/{f}.jsonl"):
        r = json.loads(l)
        if r.get("arm") == arm and "e2e_ms" in r:
            reqs.append((r["t_wall"], round(r["dur"] - r.get("trim_ms", 0) / 1000, 2), r.get("mode", "first")))
reqs.sort()
prev, i, longest, last = None, 0, 0.0, None
for t, m in series:
    while i < len(reqs) and reqs[i][0] <= t:
        longest = max(longest, reqs[i][1]); last = reqs[i]; i += 1
    if m != prev:
        ts = datetime.datetime.fromtimestamp(t).strftime("%H:%M:%S")
        print(f"{ts} {m:6d} MiB  longest served so far {longest:5.1f}s  last request {last[1] if last else '-'}s ({last[2] if last else '-'})")
        prev = m