stress_operator/utils/logging.py

"""Run logging: every run writes a JSON ``{config, git_commit, seed, final_metrics,
wall_clock, gpu, est_cost}`` to ``results/`` (master plan §0.2 principle 3)."""
from __future__ import annotations

import json
import os
import subprocess
from typing import Any, Dict, Optional


def git_commit(default: str = "unknown") -> str:
    try:
        out = subprocess.check_output(
            ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
        )
        return out.decode().strip()
    except Exception:
        return default


def write_run_log(
    path: str,
    config: Dict[str, Any],
    seed: int,
    final_metrics: Dict[str, Any],
    wall_clock_sec: float,
    gpu: str,
    est_cost_usd: float,
    extra: Optional[Dict[str, Any]] = None,
) -> str:
    """Write a single run record as JSON and return the path."""
    record = {
        "config": config,
        "git_commit": git_commit(),
        "seed": seed,
        "final_metrics": final_metrics,
        "wall_clock_sec": round(float(wall_clock_sec), 2),
        "gpu": gpu,
        "est_cost_usd": round(float(est_cost_usd), 4),
    }
    if extra:
        record.update(extra)
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w") as f:
        json.dump(record, f, indent=2)
    return path


def estimate_cost(wall_clock_sec: float, rate_per_sec: float) -> float:
    """GPU cost estimate = wall-clock seconds x per-second rate (master plan §2.7)."""
    return wall_clock_sec * rate_per_sec


# Modal per-second GPU rates (USD), from master plan §2.7 (modal.com/pricing, June 2026).
MODAL_RATES_PER_SEC = {
    "A10": 0.000306,
    "A100-40GB": 0.000583,
    "A100-80GB": 0.000694,
    "H100": 0.001097,
    "L4": 0.000222,
    "T4": 0.000164,
    "CPU": 0.0,
}