diff --git a/scripts/aggregate-multi-seed-metrics.py b/scripts/aggregate-multi-seed-metrics.py new file mode 100755 index 000000000..a116c5fe8 --- /dev/null +++ b/scripts/aggregate-multi-seed-metrics.py @@ -0,0 +1,308 @@ +#!/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 `__`.""" + 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()) diff --git a/scripts/aggregate-norm-stats.py b/scripts/aggregate-norm-stats.py new file mode 100755 index 000000000..76454e83c --- /dev/null +++ b/scripts/aggregate-norm-stats.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python3 +"""Merge per-(fold, seed) norm_stats JSON files into per-fold aggregates. + +Plan 5 Task 1B.3 — input is a directory of `norm_stats_foldN_seedM.json` +files written by training. Each file contains per-feature normalisation +statistics (mean, std) computed on that (fold, seed)'s training window. +Across seeds for the same fold, those stats should be near-identical (same +training window), but seed-driven sampling order can introduce small +variation; this script collapses N seeds into one per-fold stat for the +`evaluate` step's inference normaliser. + +Output: per-fold JSON files named `norm_stats_foldN.json` with mean and +median across seeds for each feature dimension. + +Input file schema (one per fold/seed): + + { + "fold": 0, + "seed": 0, + "feature_dim": 42, + "mean": [, ...], # per-dim mean + "std": [, ...], # per-dim std (sample stddev) + "n_bars": 12345 # samples used to compute stats + } + +Output file schema (one per fold): + + { + "fold": N, + "feature_dim": int, + "n_seeds": int, + "mean": [, ...], # mean of per-seed means + "median": [, ...], # median of per-seed means + "std": [, ...], # mean of per-seed stds + "std_dispersion": [, ...], # cross-seed stddev of the means + "n_bars_total": 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 + +NORM_STATS_RE = re.compile(r"norm_stats_fold(\d+)_seed(\d+)\.json") + + +def main() -> int: + ap = argparse.ArgumentParser( + description=( + "Merge norm_stats_foldN_seedM.json files across seeds into " + "per-fold aggregates for inference normalisation." + ), + ) + ap.add_argument( + "--input-dir", + type=Path, + required=True, + help="Directory containing norm_stats_foldN_seedM.json files", + ) + ap.add_argument( + "--output-dir", + type=Path, + required=True, + help="Directory for output norm_stats_foldN.json files", + ) + args = ap.parse_args() + + if not args.input_dir.is_dir(): + print(f"ERROR: input dir not found: {args.input_dir}", file=sys.stderr) + return 1 + + by_fold: dict[int, list[Path]] = defaultdict(list) + for entry in args.input_dir.iterdir(): + m = NORM_STATS_RE.match(entry.name) + if m: + by_fold[int(m.group(1))].append(entry) + + if not by_fold: + print( + f"ERROR: no norm_stats_foldN_seedM.json files found in {args.input_dir}", + file=sys.stderr, + ) + return 1 + + args.output_dir.mkdir(parents=True, exist_ok=True) + + for fold, paths in sorted(by_fold.items()): + loaded = [] + for p in sorted(paths): + try: + with p.open("r", encoding="utf-8") as fh: + loaded.append(json.load(fh)) + except (OSError, json.JSONDecodeError) as e: + print(f"WARN: skipping {p}: {e}", file=sys.stderr) + if not loaded: + print(f"WARN: no readable seeds for fold {fold}", file=sys.stderr) + continue + + # Validate consistent feature_dim — abort if not (norm-stats files + # for the same fold must agree on dimensionality). + feature_dims = {entry.get("feature_dim") for entry in loaded} + if len(feature_dims) != 1 or None in feature_dims: + print( + f"ERROR: fold {fold} has inconsistent feature_dim: {feature_dims}", + file=sys.stderr, + ) + return 1 + feature_dim = feature_dims.pop() + + # Per-dim aggregation: collect across seeds, then reduce. + means_by_dim: list[list[float]] = [[] for _ in range(feature_dim)] + stds_by_dim: list[list[float]] = [[] for _ in range(feature_dim)] + for entry in loaded: + entry_mean = entry["mean"] + entry_std = entry["std"] + if len(entry_mean) != feature_dim or len(entry_std) != feature_dim: + print( + f"ERROR: fold {fold} seed entry has mismatched dim", + file=sys.stderr, + ) + return 1 + for d in range(feature_dim): + means_by_dim[d].append(float(entry_mean[d])) + stds_by_dim[d].append(float(entry_std[d])) + + out = { + "fold": fold, + "feature_dim": feature_dim, + "n_seeds": len(loaded), + "mean": [mean(vals) for vals in means_by_dim], + "median": [median(vals) for vals in means_by_dim], + "std": [mean(vals) for vals in stds_by_dim], + "std_dispersion": [ + pstdev(vals) if len(vals) > 1 else 0.0 for vals in means_by_dim + ], + "n_bars_total": sum(int(entry.get("n_bars", 0)) for entry in loaded), + } + + out_path = args.output_dir / f"norm_stats_fold{fold}.json" + with out_path.open("w", encoding="utf-8") as fh: + json.dump(out, fh, indent=2) + print( + f"fold {fold}: merged {len(loaded)} seeds → {out_path}", + file=sys.stderr, + ) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/gather-multi-seed-metrics.sh b/scripts/gather-multi-seed-metrics.sh new file mode 100755 index 000000000..64a644e16 --- /dev/null +++ b/scripts/gather-multi-seed-metrics.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env bash +# Gather per-seed HEALTH_DIAG metrics from a multi-seed Argo run and pipe +# them through aggregate-multi-seed-metrics.py for cross-seed statistics. +# +# Plan 5 Task 1B.1. +# +# Usage: +# ./scripts/gather-multi-seed-metrics.sh [output.json] +# +# Where matches the foxhunt-tag label set via `argo-train.sh --tag`. +# The script: +# 1. Pulls all logs for workflows labelled foxhunt-tag= into a +# flat text file (/tmp/all-logs-.txt). +# 2. Invokes aggregate-multi-seed-metrics.py to produce a single JSON +# report with per-(epoch, metric) mean/std/median/min/max across +# seeds. +# +# Defaults: --multi-seed and --folds are read from the MULTI_SEED / FOLDS +# environment vars (default 5 / 6). Override via env or by editing here. +set -euo pipefail + +TAG="${1:?Usage: gather-multi-seed-metrics.sh [output.json]}" +OUT="${2:-/tmp/multi-seed-${TAG}.json}" +LOG_FILE="/tmp/all-logs-${TAG}.txt" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +if ! command -v argo >/dev/null 2>&1; then + echo "ERROR: argo CLI not found in PATH" >&2 + exit 1 +fi + +if ! command -v python3 >/dev/null 2>&1; then + echo "ERROR: python3 not found in PATH" >&2 + exit 1 +fi + +echo "=== Pulling logs for workflows tagged foxhunt-tag=${TAG} ===" +# `argo logs @latest` is a single-workflow tail; for multi-workflow +# selection we list workflows by label, then concatenate their logs. +WORKFLOWS=$(argo list -n foxhunt --selector "foxhunt-tag=${TAG}" -o name 2>/dev/null || true) +if [[ -z "${WORKFLOWS}" ]]; then + echo "ERROR: no workflows found with foxhunt-tag=${TAG}" >&2 + exit 1 +fi + +: > "${LOG_FILE}" +for wf in ${WORKFLOWS}; do + echo " -> $wf" + argo logs -n foxhunt --no-color "$wf" >> "${LOG_FILE}" 2>&1 || { + echo " WARN: failed to fetch logs for $wf" >&2 + } +done + +LOG_LINES=$(wc -l < "${LOG_FILE}") +echo "Collected ${LOG_LINES} log lines → ${LOG_FILE}" + +echo "=== Aggregating per-(epoch, metric) statistics ===" +python3 "${REPO_ROOT}/scripts/aggregate-multi-seed-metrics.py" \ + --input "${LOG_FILE}" \ + --multi-seed "${MULTI_SEED:-5}" \ + --folds "${FOLDS:-6}" \ + --tag "${TAG}" \ + --output "${OUT}" + +echo "Aggregated → ${OUT}" diff --git a/scripts/requirements.txt b/scripts/requirements.txt new file mode 100644 index 000000000..39179bd06 --- /dev/null +++ b/scripts/requirements.txt @@ -0,0 +1,13 @@ +# Python dependencies for scripts/ utilities. +# +# Used by: +# - aggregate-multi-seed-metrics.py (Plan 5 Task 1B.2) +# - aggregate-norm-stats.py (Plan 5 Task 1B.3) +# +# numpy / scipy: not currently imported by the aggregators (they use +# stdlib `statistics`), but reserved for downstream metric-band plotting +# and tier-validation scripts (Plan 5 Tasks 4–5) which use scipy.stats +# for confidence intervals and matplotlib for visualisations. +numpy>=1.26 +scipy>=1.11 +matplotlib>=3.8