#!/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=`), 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[]: [= = ...] \ [= ...] ... 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": { "": [ {"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[]: 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: [ ...] # 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_PREFIX_RE = re.compile(r"^(?P[a-z0-9][a-z0-9-]*?):\s+(?P.*)$") 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 `[ ...]` segments from a HEALTH_DIAG body. Returns flat dict with keys `_` (single underscore joiner — this is what the Plan 5 Task 4 tier check scripts expect, and what the synthetic test fixtures encode: e.g. `val [sharpe=…]` becomes aggregate key `val_sharpe`).""" 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())