#!/usr/bin/env python3
"""Sample engine-state feed freshness during the es1 run.
Writes one jsonl record per engine per tick: age_s = now - state.ts.
Stops when the DONE marker appears (run finished + /dev/shm wiped) or after 90min.
"""
import json, os, time, sys, glob

esdir, outpath, donemarker = sys.argv[1], sys.argv[2], sys.argv[3]
deadline = time.time() + 90 * 60
with open(outpath, "a") as out:
    while not os.path.exists(donemarker) and time.time() < deadline:
        now = time.time()
        if os.path.isdir(esdir):
            for f in sorted(glob.glob(os.path.join(esdir, "engine_*.json"))):
                try:
                    s = json.load(open(f))
                    rec = {
                        "now": round(now, 3),
                        "engine": os.path.basename(f)[:-5],
                        "age_s": round(now - s.get("ts", 0), 4),
                        "num_running": s.get("num_running"),
                        "num_prefilling": s.get("num_prefilling"),
                        "max_prefill_remaining": s.get("max_prefill_remaining"),
                        "kv_used": s.get("gpu_kv_used_frac"),
                    }
                    out.write(json.dumps(rec) + "\n")
                except Exception:
                    pass
            out.flush()
        time.sleep(5)
