| |
| """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:]] |
| |
| |
| 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() |
|
|
