File size: 2,155 Bytes
9589849 | 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 50 51 52 53 54 55 56 57 58 59 60 61 | #!/usr/bin/env python3
"""Print a compact status summary from prime-rl's local metrics sink."""
from __future__ import annotations
import argparse
import json
import statistics
from collections import defaultdict
from pathlib import Path
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("metrics", type=Path)
parser.add_argument("--max-steps", type=int, required=True)
args = parser.parse_args()
by_step: dict[int, dict] = defaultdict(dict)
for line in args.metrics.read_text().splitlines():
row = json.loads(line)
by_step[int(row["step"])].update(row)
if not by_step:
print("no completed steps")
return
steps = sorted(by_step)
latest = by_step[steps[-1]]
recent = [by_step[step] for step in steps[-10:]]
# A compiled first step reports zero throughput and can be several minutes
# slower; exclude it from the steady-state ETA.
times = [
row["time/step"]
for row in recent
if row.get("time/step", 0) > 0 and row.get("perf/throughput", 0) > 0
]
throughputs = [row["perf/throughput"] for row in recent if row.get("perf/throughput", 0) > 0]
losses = [row["loss/mean"] for row in recent if "loss/mean" in row]
mean_step_seconds = statistics.mean(times)
remaining_seconds = max(0, args.max_steps - steps[-1]) * mean_step_seconds
result = {
"step": steps[-1],
"max_steps": args.max_steps,
"tokens": latest.get("progress/num_tokens"),
"samples": latest.get("progress/num_samples"),
"lr": latest.get("optim/lr"),
"loss": latest.get("loss/mean"),
"recent_mean_loss": statistics.mean(losses),
"recent_mean_throughput": statistics.mean(throughputs) if throughputs else 0,
"recent_mean_step_seconds": mean_step_seconds,
"estimated_remaining_hours": remaining_seconds / 3600,
"peak_memory_gib": latest.get("perf/peak_memory"),
"grad_norm": latest.get("optim/grad_norm"),
"nan_count": latest.get("loss/nan_count"),
}
print(json.dumps(result, indent=2))
if __name__ == "__main__":
main()
|