social-data-attribution-w2/scripts/audit/aggregate_attribution_compute.py

#!/usr/bin/env python3
# Walk SLURM .out logs in the social-data-attribution data repo and aggregate
# per-job (gpu_type, walltime, partition) into per-phase totals. Designed to
# run against logs/{attribution,unlearning,evaluation,enrich_pool_gpu,soc91_enrich}
# directories on PACE filesystem or any local clone that holds them.
#
# Source-of-truth for the compute audit per `docs/compute_audit_plan.md`
# and `docs/slurm_gpu_catalog.md`.
#
# Read-only. No side effects beyond writing the output JSON.
#
# Usage:
#   python scripts/audit/aggregate_attribution_compute.py \
#       --logs-root /path/to/social-data-attribution/logs \
#       --output    outputs/compute_audit/log_breakdown.json
#
# Required Python: 3.10+ (uses match-case + type hints).

from __future__ import annotations

import argparse
import json
import re
import sys
from collections import defaultdict
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

# bf16 FLOPS-based normalization factors relative to H200 144 GB / H100 80 GB
# (both 989 TFLOPS bf16 dense). Rounded UP at 2 sig figs to avoid
# under-counting.
H200_EQUIV_FACTOR: dict[str, float] = {
    "h200": 1.00,
    "h100": 1.00,
    "rtx_pro_6000": 0.76,  # Blackwell 96 GB
    "rtx_6000_ada": 0.37,  # Ada Lovelace 24 GB
    "rtx_6000": 0.37,
    "l40s": 0.37,
    "a100": 0.32,
    "a40": 0.15,
    "v100": 0.13,
    "unknown": 0.50,  # conservative-but-not-zero default
}

# Log directory → phase mapping. These names match the `--output=` paths
# in scripts/slurm/**/*.sbatch headers.
PHASE_DIRS: dict[str, str] = {
    "soc91_enrich": "1_enrichment",
    "enrich_pool": "1_enrichment",
    "enrich_pool_gpu": "1_enrichment",
    "attribution": "2_attribution",
    "tracstar": "2_attribution",
    "precond": "2c_preconditioner",
    "unlearn": "3_unlearning",
    "ngdiff": "3_unlearning",
    "eval_unlearn": "3_unlearning_eval",
    "evaluation": "4_evaluation",
    "olmes": "4_evaluation",
    "eval_arc": "4_evaluation",
    "eval_bbh": "4_evaluation",
}


@dataclass
class JobRecord:
    job_id: str
    array_id: str | None
    log_path: str
    phase: str
    gpu_type: str  # lower-snake-case key into H200_EQUIV_FACTOR
    walltime_seconds: float
    gpu_count: int  # gres=gpu:N or AllocTRES gpu=N
    partition: str | None
    state: str | None  # COMPLETED / FAILED / TIMEOUT / etc.

    @property
    def raw_gpu_hours(self) -> float:
        return (self.walltime_seconds / 3600.0) * self.gpu_count

    @property
    def h200_equiv_hours(self) -> float:
        return self.raw_gpu_hours * H200_EQUIV_FACTOR.get(
            self.gpu_type, H200_EQUIV_FACTOR["unknown"]
        )


@dataclass
class PhaseSummary:
    phase: str
    n_jobs: int = 0
    raw_by_gpu: dict[str, float] = field(default_factory=dict)
    raw_total: float = 0.0
    h200_equiv_total: float = 0.0


# --- parsers ---------------------------------------------------------------

_RE_ELAPSED = re.compile(r"\bElapsed:\s*(\d+):(\d+):(\d+)\b")
_RE_ELAPSED_DD = re.compile(r"\bElapsed:\s*(\d+)-(\d+):(\d+):(\d+)\b")
_RE_GPU_NV = re.compile(r"NVIDIA\s+([A-Za-z0-9 \-]+?)(?:\s+\(\w+\))?$", re.M)
_RE_GRES = re.compile(r"--gres=gpu(?::([a-z0-9_]+))?:(\d+)")
_RE_GPU_MODEL = re.compile(
    r"(H200|H100|A100|L40S|A40|V100|RTX[\s_]?PRO[\s_]?6000|RTX[\s_]?6000)",
    re.I,
)
_RE_PARTITION = re.compile(r"--partition=([a-z0-9_\-,]+)")
_RE_STATE = re.compile(r"State:\s*([A-Z_]+)")


def parse_walltime(text: str) -> float | None:
    """Parse `Elapsed: HH:MM:SS` or `Elapsed: D-HH:MM:SS` → seconds."""
    if m := _RE_ELAPSED_DD.search(text):
        d, h, mi, s = map(int, m.groups())
        return d * 86400 + h * 3600 + mi * 60 + s
    if m := _RE_ELAPSED.search(text):
        h, mi, s = map(int, m.groups())
        return h * 3600 + mi * 60 + s
    return None


def detect_gpu_type(text: str) -> str:
    """Find GPU model in the log preamble (from nvidia-smi or sbatch echo)."""
    if m := _RE_GPU_MODEL.search(text):
        raw = m.group(1).lower().replace(" ", "_").replace("-", "_")
        if "rtx_pro" in raw:
            return "rtx_pro_6000"
        if "rtx" in raw and "6000" in raw:
            return "rtx_6000"
        for key in ("h200", "h100", "a100", "l40s", "a40", "v100"):
            if key in raw:
                return key
    return "unknown"


def detect_gpu_count(text: str) -> int:
    if m := _RE_GRES.search(text):
        return int(m.group(2))
    if "AllocTRES" in text:
        # Try to extract gpu= count from AllocTRES line if present
        m = re.search(r"gpu=(\d+)", text)
        if m:
            return int(m.group(1))
    return 1


def classify_phase(log_path: Path) -> str:
    """Pick a phase tag from the log file path."""
    for token in log_path.parts:
        for keyword, phase in PHASE_DIRS.items():
            if keyword in token:
                return phase
    return "other"


def parse_log(log_path: Path) -> JobRecord | None:
    try:
        # Read the first 64 KB — slurm metadata typically appears at top
        # (sbatch echo) and bottom (sacct epilogue if used).
        with log_path.open(encoding="utf-8", errors="replace") as fh:
            head = fh.read(65536)
            try:
                fh.seek(-32768, 2)
            except OSError:
                # File shorter than 32 KB
                tail = ""
            else:
                tail = fh.read(32768)
        text = head + "\n" + tail
    except (OSError, ValueError):
        return None

    walltime = parse_walltime(text)
    if walltime is None or walltime <= 0:
        return None

    state_match = _RE_STATE.search(text)
    partition_match = _RE_PARTITION.search(text)
    return JobRecord(
        job_id=log_path.stem,
        array_id=None,
        log_path=str(log_path),
        phase=classify_phase(log_path),
        gpu_type=detect_gpu_type(text),
        walltime_seconds=walltime,
        gpu_count=detect_gpu_count(text),
        partition=partition_match.group(1) if partition_match else None,
        state=state_match.group(1) if state_match else None,
    )


# --- aggregation -----------------------------------------------------------


def aggregate(logs_root: Path) -> tuple[list[JobRecord], dict[str, PhaseSummary]]:
    records: list[JobRecord] = []
    summaries: dict[str, PhaseSummary] = defaultdict(
        lambda: PhaseSummary(phase="<placeholder>")
    )

    for log_path in logs_root.rglob("*.out"):
        if (rec := parse_log(log_path)) is None:
            continue
        records.append(rec)
        summary = summaries[rec.phase]
        summary.phase = rec.phase
        summary.n_jobs += 1
        summary.raw_by_gpu[rec.gpu_type] = (
            summary.raw_by_gpu.get(rec.gpu_type, 0.0) + rec.raw_gpu_hours
        )
        summary.raw_total += rec.raw_gpu_hours
        summary.h200_equiv_total += rec.h200_equiv_hours

    return records, dict(summaries)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--logs-root",
        type=Path,
        required=True,
        help="Root directory of slurm .out logs to walk recursively",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("outputs/compute_audit/log_breakdown.json"),
        help="Where to write the aggregated JSON report",
    )
    parser.add_argument(
        "--include-records",
        action="store_true",
        help="Embed per-job records in the output (large)",
    )
    args = parser.parse_args()

    if not args.logs_root.is_dir():
        print(f"logs-root not found: {args.logs_root}", file=sys.stderr)
        return 1

    records, summaries = aggregate(args.logs_root)

    report: dict[str, Any] = {
        "logs_root": str(args.logs_root),
        "n_jobs_parsed": len(records),
        "phases": {
            name: {
                "n_jobs": s.n_jobs,
                "raw_gpu_hours_total": round(s.raw_total, 1),
                "h200_equiv_gpu_hours_total": round(s.h200_equiv_total, 1),
                "raw_by_gpu_type": {
                    g: round(v, 1) for g, v in sorted(s.raw_by_gpu.items())
                },
            }
            for name, s in sorted(summaries.items())
        },
        "totals": {
            "raw_gpu_hours": round(sum(s.raw_total for s in summaries.values()), 1),
            "h200_equiv_gpu_hours": round(
                sum(s.h200_equiv_total for s in summaries.values()), 1
            ),
        },
        "h200_equiv_factors": H200_EQUIV_FACTOR,
    }
    if args.include_records:
        report["records"] = [asdict(r) for r in records]

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2))
    print(f"Wrote {args.output} ({len(records)} jobs across {len(summaries)} phases).")
    print(f"  raw GPU-hr:        {report['totals']['raw_gpu_hours']:>10.1f}")
    print(f"  H200-equiv GPU-hr: {report['totals']['h200_equiv_gpu_hours']:>10.1f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())