substrate/pod_ledger.py

"""pod_ledger.py — RUNNER-2 campaign clock + utilization ledger
(COPY of runner 1's pod/pod_ledger.py per the charter — never import or edit
the original; its WINDOW constants are runner 1's live campaign clock.)

WINDOW: PLACEHOLDER 2026-07-16 -> 2026-07-20 UTC — Phil sets the real rental
window at kickoff; override via env R2_WINDOW_START_UTC / R2_WINDOW_END_UTC
(ISO, e.g. 2026-07-17T00:00) or edit the constants below THEN, not mid-run.

Usage (pod-side, wraps every run):
    from pod_ledger import ledger_run, burn_down
    with ledger_run("dexp001 sd15 relay all16", budget_h=1.5):
        ...            # the run; row written on exit, verdict settable
    burn_down()        # print remaining wall/allocated/reserve

Ledger file: JSONL at LEDGER_PATH (pod: /workspace/ledger2/pod2_ledger.jsonl;
local fallback under GEOLIP_DATA). Rows: {kind: run|note|day_report, name,
start_utc, end_utc, hours, budget_h, verdict, note}. The daily report for
Phil aggregates by day + experiment.

Colab/pod-safe: stdlib only.
"""
from __future__ import annotations
import contextlib
import datetime as dt
import json
import os

# r2 window (PLACEHOLDER until kickoff; env-overridable, ISO format)
def _win(env, default):
    v = os.environ.get(env)
    if v:
        return dt.datetime.fromisoformat(v).replace(tzinfo=dt.timezone.utc)
    return default


WINDOW_START_UTC = _win("R2_WINDOW_START_UTC",
                        dt.datetime(2026, 7, 16, 0, 0, tzinfo=dt.timezone.utc))
WINDOW_END_UTC = _win("R2_WINDOW_END_UTC",
                      dt.datetime(2026, 7, 20, 0, 0, tzinfo=dt.timezone.utc))
TOTAL_H = (WINDOW_END_UTC - WINDOW_START_UTC).total_seconds() / 3600.0
RESERVE_H = 12.0

_DEFAULT_DIRS = ("/workspace/ledger2",
                 os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
                              "pod_ledger"))


def _ledger_path():
    for d in _DEFAULT_DIRS:
        try:
            os.makedirs(d, exist_ok=True)
            return os.path.join(d, "pod2_ledger.jsonl")
        except OSError:
            continue
    return "pod2_ledger.jsonl"


LEDGER_PATH = _ledger_path()


def _now():
    return dt.datetime.now(dt.timezone.utc)


def _write(row: dict):
    row["ts"] = _now().isoformat()
    with open(LEDGER_PATH, "a", encoding="utf-8") as f:
        f.write(json.dumps(row) + "\n")


def _rows():
    if not os.path.exists(LEDGER_PATH):
        return []
    return [json.loads(l) for l in open(LEDGER_PATH, encoding="utf-8")]


def note(text: str):
    _write({"kind": "note", "note": text})
    print(f"[ledger] {text}", flush=True)


def spent_hours() -> float:
    return sum(r.get("hours", 0.0) for r in _rows() if r.get("kind") == "run")


def wall_remaining_h() -> float:
    return max(0.0, (WINDOW_END_UTC - _now()).total_seconds() / 3600.0)


def burn_down(print_it=True) -> dict:
    spent = spent_hours()
    wall = wall_remaining_h()
    elapsed = TOTAL_H - wall
    state = {
        "wall_elapsed_h": round(elapsed, 2),
        "wall_remaining_h": round(wall, 2),
        "gpu_hours_ledgered": round(spent, 2),
        "allocatable_remaining_h": round(max(0.0, wall - RESERVE_H), 2),
        "reserve_h": RESERVE_H,
        "utilization_pct_of_elapsed": round(100 * spent / max(elapsed, 0.01),
                                            1),
    }
    if print_it:
        print(f"[burn-down] wall {state['wall_elapsed_h']:.1f}h elapsed / "
              f"{state['wall_remaining_h']:.1f}h left | runs "
              f"{state['gpu_hours_ledgered']:.1f}h ledgered "
              f"({state['utilization_pct_of_elapsed']:.0f}% of elapsed) | "
              f"allocatable {state['allocatable_remaining_h']:.1f}h "
              f"(reserve {RESERVE_H:.0f}h)", flush=True)
    return state


@contextlib.contextmanager
def ledger_run(name: str, budget_h: float = None):
    """Time a run; write the row on exit (even on failure, with verdict)."""
    start = _now()
    print(f"[ledger] START {name}"
          + (f" (budget {budget_h}h)" if budget_h else ""), flush=True)
    holder = {"verdict": None}
    try:
        yield holder
        verdict = holder["verdict"] or "completed"
    except BaseException as e:
        verdict = f"FAILED: {type(e).__name__}: {e}"
        raise
    finally:
        end = _now()
        hours = (end - start).total_seconds() / 3600.0
        over = (budget_h is not None and hours > budget_h)
        _write({"kind": "run", "name": name,
                "start_utc": start.isoformat(), "end_utc": end.isoformat(),
                "hours": round(hours, 3), "budget_h": budget_h,
                "over_budget": over, "verdict": verdict})
        print(f"[ledger] END {name}: {hours:.2f}h"
              + (" ** OVER BUDGET **" if over else ""), flush=True)
        burn_down()


def day_report() -> str:
    """Aggregate for Phil's daily check-in."""
    rows = [r for r in _rows() if r.get("kind") == "run"]
    by_day = {}
    for r in rows:
        day = r["start_utc"][:10]
        by_day.setdefault(day, []).append(r)
    lines = ["POD UTILIZATION REPORT", "=" * 40]
    for day in sorted(by_day):
        runs = by_day[day]
        total = sum(r["hours"] for r in runs)
        lines.append(f"\n{day}: {total:.1f}h across {len(runs)} runs")
        for r in runs:
            flag = " **OVER**" if r.get("over_budget") else ""
            lines.append(f"  {r['hours']:5.2f}h  {r['name']}"
                         f"  [{str(r['verdict'])[:60]}]{flag}")
    s = burn_down(print_it=False)
    lines.append(f"\nWall: {s['wall_elapsed_h']:.1f}h elapsed, "
                 f"{s['wall_remaining_h']:.1f}h remaining "
                 f"({s['allocatable_remaining_h']:.1f}h allocatable + "
                 f"{s['reserve_h']:.0f}h reserve)")
    report = "\n".join(lines)
    _write({"kind": "day_report", "note": report})
    return report


def smoke():
    import tempfile, time
    global LEDGER_PATH
    old = LEDGER_PATH
    LEDGER_PATH = os.path.join(tempfile.gettempdir(), "pod_ledger_smoke.jsonl")
    if os.path.exists(LEDGER_PATH):
        os.remove(LEDGER_PATH)
    note("smoke start")
    with ledger_run("smoke-run", budget_h=0.00001) as h:   # 0.036s budget
        time.sleep(0.05)
        h["verdict"] = "smoke ok"
    try:
        with ledger_run("smoke-fail", budget_h=1.0):
            raise ValueError("boom")
    except ValueError:
        pass
    rows = _rows()
    kinds = [r["kind"] for r in rows]
    assert kinds.count("run") == 2
    runs = [r for r in rows if r["kind"] == "run"]
    assert runs[0]["verdict"] == "smoke ok" and runs[0]["over_budget"]
    assert runs[1]["verdict"].startswith("FAILED: ValueError")
    rep = day_report()
    assert "smoke-run" in rep and "OVER" in rep
    s = burn_down(print_it=False)
    assert 0 <= s["wall_remaining_h"] <= TOTAL_H
    LEDGER_PATH = old
    print("pod_ledger smoke passed (run rows, failure verdicts, "
          "over-budget flags, report, burn-down)")


if __name__ == "__main__":
    smoke()