47c8b783c4
Companion to Plan 5 Task 1A (multi-seed Argo DAG). Adds the post-run aggregation pipeline: - scripts/gather-multi-seed-metrics.sh: thin wrapper that pulls Argo logs for every workflow tagged foxhunt-tag=<tag>, concatenates them into /tmp/all-logs-<tag>.txt, then dispatches to the Python aggregator. - scripts/aggregate-multi-seed-metrics.py: stdlib-only HEALTH_DIAG parser. Recognises both bare and JSON-envelope-wrapped HEALTH_DIAG[<epoch>] lines, parses every <block>[<key=val> ...] segment, and emits per- (epoch, metric_name) mean / std / median / min / max across streams (where each stream is one (seed, fold) training run, attributed by pod-name prefix when present, else by epoch-rewind detection — naive epoch=0 trigger over-segments the multi-line per-epoch HEALTH_DIAG output, fixed by requiring epoch < last_epoch to start a new stream). Output JSON schema matches the plan example (top-level: tag, multi_seed, folds, warmup_end_epoch, streams_seen, aggregates; per-entry: epoch, mean, std, median, min, max, n_samples). - scripts/aggregate-norm-stats.py: stdlib-only merger for norm_stats_foldN_seedM.json files. Per-fold output collapses N seeds into mean/median/std/std_dispersion arrays consumed by the `evaluate` step's inference normaliser. - scripts/requirements.txt: declares numpy/scipy/matplotlib for downstream Plan 5 T4-T5 tier-validation/plotting scripts. The aggregators themselves are stdlib-only (statistics module). Smoke-validated on /tmp/p4t6-cleanroom-smoke.log: detects 3 streams (matching the 3 fold runs in that log), 90 unique metrics, n_samples=3 per (epoch, metric), schema matches plan. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
309 lines
10 KiB
Python
Executable File
309 lines
10 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Aggregate per-seed HEALTH_DIAG metrics into a single JSON report.
|
|
|
|
Plan 5 Task 1B.2 — consumes the concatenated logs from N*K (seed, fold)
|
|
training jobs (output of `argo logs --selector foxhunt-tag=<tag>`),
|
|
parses HEALTH_DIAG lines, groups by (seed, fold, epoch, metric_name), and
|
|
emits per-(epoch, metric) aggregate statistics across seeds.
|
|
|
|
HEALTH_DIAG line format (from training_loop.rs):
|
|
|
|
HEALTH_DIAG[<epoch>]: <key1> [<sub1>=<val1> <sub2>=<val2> ...] \
|
|
<key2> [<sub3>=<val3> ...] ...
|
|
|
|
The line may be wrapped in tracing/JSON envelope:
|
|
|
|
{"timestamp":..., "fields":{"message":"HEALTH_DIAG[N]: ..."}, ...}
|
|
|
|
Per-line followups (`reward_split`, `aux`) are also parsed.
|
|
|
|
Seed/fold attribution:
|
|
|
|
The script reads `--multi-seed N --folds K` and assumes the input log
|
|
contains exactly N*K independent training streams in order. Each stream
|
|
is identified by the ordered (seed, fold) pair derived from its
|
|
position in the log — Argo's `argo logs --no-color` interleaves pod
|
|
output with a `pod_name` prefix that we use as the stream key. If
|
|
pod-name attribution is unavailable, we fall back to detecting epoch
|
|
resets (a HEALTH_DIAG[0] line marks a new stream).
|
|
|
|
Output JSON schema:
|
|
|
|
{
|
|
"tag": str,
|
|
"multi_seed": int,
|
|
"folds": int,
|
|
"warmup_end_epoch": int | null,
|
|
"streams_seen": int,
|
|
"aggregates": {
|
|
"<metric_name>": [
|
|
{"epoch": 0, "mean": float, "std": float,
|
|
"median": float, "min": float, "max": float,
|
|
"n_samples": int},
|
|
...
|
|
],
|
|
...
|
|
}
|
|
}
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import sys
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from statistics import mean, median, pstdev
|
|
from typing import Iterable
|
|
|
|
# HEALTH_DIAG[<epoch>]: <body>
|
|
HEALTH_DIAG_RE = re.compile(r"HEALTH_DIAG\[(\d+)\]:\s*(.*)")
|
|
|
|
# JSON envelope: {"fields":{"message":"HEALTH_DIAG[..."}, ...}
|
|
JSON_MESSAGE_RE = re.compile(r'"message"\s*:\s*"([^"]*HEALTH_DIAG\[[^"]*)"')
|
|
|
|
# Top-level block: <name> [<key=val> <key=val> ...]
|
|
# Names are alphanumeric + underscore; the bracketed body contains key=val
|
|
# pairs separated by whitespace. We keep names simple (no nested brackets
|
|
# inside HEALTH_DIAG bodies — verified against current training_loop.rs).
|
|
BLOCK_RE = re.compile(r"(\w+)\s*\[([^\[\]]*)\]")
|
|
|
|
# key=val inside a block. Values may be:
|
|
# - signed floats (with/without scientific notation): -3.911565e0
|
|
# - signed ints: 5
|
|
# - bool literals: true / false / on / off / ready / ...
|
|
# - small enums (action_entropy=0.79, plasticity=ready)
|
|
# We capture as (key, value) and let downstream typing decide.
|
|
KV_RE = re.compile(r"(\w+)=([^\s\]]+)")
|
|
|
|
# Pod-name prefix from `argo logs` output (best-effort; not all log
|
|
# collectors emit this). Format: "<pod-name>: <log line>".
|
|
POD_PREFIX_RE = re.compile(r"^(?P<pod>[a-z0-9][a-z0-9-]*?):\s+(?P<rest>.*)$")
|
|
|
|
|
|
def parse_value(raw: str) -> float | None:
|
|
"""Parse a HEALTH_DIAG value. Returns None for non-numeric (boolean,
|
|
enum) values — those are excluded from aggregation since mean/std are
|
|
not meaningful for them."""
|
|
# Strip trailing punctuation that sometimes appears at end of bracket.
|
|
raw = raw.strip().rstrip(",;:")
|
|
try:
|
|
return float(raw)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def extract_health_diag(line: str) -> tuple[int, str] | None:
|
|
"""Pull (epoch, body) from a raw log line. Handles both bare
|
|
HEALTH_DIAG and JSON-envelope-wrapped variants."""
|
|
# JSON envelope first — common in production logs.
|
|
m = JSON_MESSAGE_RE.search(line)
|
|
if m:
|
|
msg = m.group(1).replace("\\\"", '"')
|
|
m2 = HEALTH_DIAG_RE.search(msg)
|
|
if m2:
|
|
return int(m2.group(1)), m2.group(2)
|
|
return None
|
|
# Bare HEALTH_DIAG.
|
|
m = HEALTH_DIAG_RE.search(line)
|
|
if m:
|
|
return int(m.group(1)), m.group(2)
|
|
return None
|
|
|
|
|
|
def parse_blocks(body: str) -> dict[str, float]:
|
|
"""Parse all `<block>[<kv> <kv> ...]` segments from a HEALTH_DIAG
|
|
body. Returns flat dict with keys `<block>__<subkey>`."""
|
|
out: dict[str, float] = {}
|
|
for block_match in BLOCK_RE.finditer(body):
|
|
block_name = block_match.group(1)
|
|
block_body = block_match.group(2)
|
|
for kv_match in KV_RE.finditer(block_body):
|
|
key = kv_match.group(1)
|
|
val = parse_value(kv_match.group(2))
|
|
if val is not None:
|
|
out[f"{block_name}__{key}"] = val
|
|
return out
|
|
|
|
|
|
def stream_key(line: str, fallback: int) -> str:
|
|
"""Best-effort stream identifier for grouping HEALTH_DIAG lines back
|
|
to a single training run. Uses pod-name prefix when present, else
|
|
falls back to a synthetic counter that increments on epoch=0
|
|
re-observations (see `attribute_streams`)."""
|
|
m = POD_PREFIX_RE.match(line.strip())
|
|
if m:
|
|
return m.group("pod")
|
|
return f"stream-{fallback}"
|
|
|
|
|
|
def attribute_streams(
|
|
lines: Iterable[str],
|
|
expected_streams: int,
|
|
) -> dict[str, list[tuple[int, dict[str, float]]]]:
|
|
"""Walk the log line-by-line and bucket HEALTH_DIAG observations by
|
|
stream. Each stream's value is an ordered list of (epoch, metrics_dict)
|
|
pairs.
|
|
|
|
Stream attribution strategy:
|
|
1. If pod-name prefix is present and stable, use it.
|
|
2. Else, treat epoch=0 as a stream-boundary marker — each new
|
|
HEALTH_DIAG[0] starts a new synthetic stream.
|
|
"""
|
|
by_stream: dict[str, list[tuple[int, dict[str, float]]]] = defaultdict(list)
|
|
synthetic_counter = 0
|
|
current_synthetic = None
|
|
last_epoch: int | None = None
|
|
for raw in lines:
|
|
diag = extract_health_diag(raw)
|
|
if diag is None:
|
|
continue
|
|
epoch, body = diag
|
|
metrics = parse_blocks(body)
|
|
if not metrics:
|
|
continue
|
|
# Try pod-name attribution first.
|
|
m = POD_PREFIX_RE.match(raw.strip())
|
|
if m:
|
|
key = m.group("pod")
|
|
else:
|
|
# Synthetic: a new stream starts when we either see the first
|
|
# HEALTH_DIAG ever, or when the current epoch goes BACKWARDS
|
|
# (epoch < last_epoch) — that signals a new training run.
|
|
# Each HEALTH_DIAG epoch typically emits multiple sub-lines
|
|
# (main / reward_split / aux); naively triggering on every
|
|
# epoch=0 over-segments those sub-lines into separate streams.
|
|
need_new_stream = current_synthetic is None or (
|
|
last_epoch is not None and epoch < last_epoch
|
|
)
|
|
if need_new_stream:
|
|
current_synthetic = f"stream-{synthetic_counter}"
|
|
synthetic_counter += 1
|
|
key = current_synthetic
|
|
by_stream[key].append((epoch, metrics))
|
|
last_epoch = epoch
|
|
return dict(by_stream)
|
|
|
|
|
|
def aggregate_streams(
|
|
by_stream: dict[str, list[tuple[int, dict[str, float]]]],
|
|
) -> dict[str, list[dict[str, float]]]:
|
|
"""Cross-stream aggregation. For each (epoch, metric_name) pair,
|
|
compute mean / std / median / min / max across streams."""
|
|
# (metric, epoch) → [values across streams]
|
|
bucket: dict[tuple[str, int], list[float]] = defaultdict(list)
|
|
for stream_metrics in by_stream.values():
|
|
for epoch, metrics in stream_metrics:
|
|
for metric_name, value in metrics.items():
|
|
bucket[(metric_name, epoch)].append(value)
|
|
|
|
# Group by metric name, sort by epoch.
|
|
by_metric: dict[str, list[dict[str, float]]] = defaultdict(list)
|
|
for (metric_name, epoch), values in bucket.items():
|
|
n = len(values)
|
|
if n == 0:
|
|
continue
|
|
by_metric[metric_name].append(
|
|
{
|
|
"epoch": epoch,
|
|
"mean": mean(values),
|
|
# pstdev (population std) so n=1 yields 0 instead of error;
|
|
# documents that this is across-seed dispersion, not a
|
|
# sample estimate.
|
|
"std": pstdev(values) if n > 1 else 0.0,
|
|
"median": median(values),
|
|
"min": min(values),
|
|
"max": max(values),
|
|
"n_samples": n,
|
|
}
|
|
)
|
|
for metric_name in by_metric:
|
|
by_metric[metric_name].sort(key=lambda r: r["epoch"])
|
|
return dict(by_metric)
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(
|
|
description="Aggregate HEALTH_DIAG metrics across multi-seed training runs.",
|
|
)
|
|
ap.add_argument(
|
|
"--input",
|
|
required=True,
|
|
type=Path,
|
|
help="Concatenated log file (output of argo logs --no-color)",
|
|
)
|
|
ap.add_argument(
|
|
"--multi-seed",
|
|
type=int,
|
|
default=1,
|
|
help="Number of seeds (informational; written to output JSON)",
|
|
)
|
|
ap.add_argument(
|
|
"--folds",
|
|
type=int,
|
|
default=1,
|
|
help="Number of folds (informational; written to output JSON)",
|
|
)
|
|
ap.add_argument(
|
|
"--tag",
|
|
type=str,
|
|
default="",
|
|
help="Workflow tag (informational; written to output JSON)",
|
|
)
|
|
ap.add_argument(
|
|
"--warmup-end-epoch",
|
|
type=int,
|
|
default=None,
|
|
help="Epoch at which warmup ends (informational; written to output JSON)",
|
|
)
|
|
ap.add_argument(
|
|
"--output",
|
|
type=Path,
|
|
required=True,
|
|
help="Path for the aggregated JSON output",
|
|
)
|
|
args = ap.parse_args()
|
|
|
|
if not args.input.is_file():
|
|
print(f"ERROR: input file not found: {args.input}", file=sys.stderr)
|
|
return 1
|
|
|
|
expected_streams = args.multi_seed * args.folds
|
|
with args.input.open("r", encoding="utf-8", errors="replace") as fh:
|
|
by_stream = attribute_streams(fh, expected_streams)
|
|
|
|
if not by_stream:
|
|
print(
|
|
f"ERROR: no HEALTH_DIAG lines found in {args.input}",
|
|
file=sys.stderr,
|
|
)
|
|
return 1
|
|
|
|
aggregates = aggregate_streams(by_stream)
|
|
|
|
out = {
|
|
"tag": args.tag,
|
|
"multi_seed": args.multi_seed,
|
|
"folds": args.folds,
|
|
"warmup_end_epoch": args.warmup_end_epoch,
|
|
"streams_seen": len(by_stream),
|
|
"aggregates": aggregates,
|
|
}
|
|
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with args.output.open("w", encoding="utf-8") as fh:
|
|
json.dump(out, fh, indent=2)
|
|
|
|
print(
|
|
f"Aggregated {len(by_stream)} streams "
|
|
f"(expected {expected_streams}={args.multi_seed}x{args.folds}), "
|
|
f"{len(aggregates)} unique metrics → {args.output}",
|
|
file=sys.stderr,
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|