summaryrefslogtreecommitdiff
path: root/scripts/collect.py
diff options
context:
space:
mode:
Diffstat (limited to 'scripts/collect.py')
-rw-r--r--scripts/collect.py27
1 files changed, 27 insertions, 0 deletions
diff --git a/scripts/collect.py b/scripts/collect.py
new file mode 100644
index 0000000..c943385
--- /dev/null
+++ b/scripts/collect.py
@@ -0,0 +1,27 @@
+"""Collect run logs into a small committable bundle.
+On the training node after (some) runs finish:
+ python scripts/collect.py --runs runs --out results/h200node1 && git add results && git commit -m "results" && git push
+Checkpoints stay on the node (runs/ is gitignored); only JSONL summaries+curves are collected."""
+import os, json, glob, argparse, shutil
+
+p = argparse.ArgumentParser()
+p.add_argument("--runs", default="runs")
+p.add_argument("--out", default="results/local")
+a = p.parse_args()
+os.makedirs(a.out, exist_ok=True)
+summary = {}
+for d in sorted(glob.glob(os.path.join(a.runs, "*"))):
+ name = os.path.basename(d)
+ lp = os.path.join(d, "log_rank0.jsonl")
+ if not os.path.exists(lp):
+ continue
+ rows = [json.loads(l) for l in open(lp)]
+ meta = next((r for r in rows if r.get("kind") == "meta"), {})
+ evals = [r for r in rows if r.get("kind") == "eval"]
+ shutil.copy(lp, os.path.join(a.out, f"{name}.jsonl"))
+ summary[name] = {"params": meta.get("params"), "world": meta.get("world"), "mode": (meta.get("cfg") or {}).get("mode"),
+ "n_probes": (meta.get("cfg") or {}).get("n_probes"), "done": os.path.exists(os.path.join(d, "DONE")),
+ "steps": evals[-1]["step"] if evals else 0, "val_loss": evals[-1]["val_loss"] if evals else None,
+ "tok_per_s": evals[-1].get("tok_per_s") if evals else None}
+json.dump(summary, open(os.path.join(a.out, "summary.json"), "w"), indent=1)
+print(json.dumps(summary, indent=1))