momentum-system/research/experiments/registry.py

"""Experiment registry — schema apply + immutable register/list/get.

The registry owns the ``experiments`` contract table. Registration is
one-shot and immutable: an id can never be overwritten (a pre-registered bar
that could be re-tuned after seeing the data is not pre-registered). The
required-keys check refuses any experiment missing its fidelity floors — the
single guard that keeps a zero-slippage artifact (vr1/liq1 class) out.
"""

from __future__ import annotations

import json
import os
from datetime import datetime, timezone
from pathlib import Path

_SCHEMA_SQL = Path(__file__).resolve().parents[2] / "storage" / "experiments_schema.sql"
_COHORT_DIR = Path(__file__).resolve().parents[2] / "experiments"

# Every experiment JSON must carry these; fidelity must carry these floors.
_REQUIRED = ("id", "hypothesis", "population", "intervention", "bars", "fidelity")
_REQUIRED_FIDELITY = ("atr_frac", "spread_mult", "adverse_ticks")

_TERMINAL = {"passed", "failed", "killed"}


def apply_schema(con) -> None:
    """Create the experiment tables idempotently on ``con``.

    Honors ``PG_SKIP_SCHEMA`` exactly like storage/db_pg.get_connection: set it
    when the owner provisions the DDL out-of-band and the harness should not
    touch DDL. On an isolated test schema the connection's search_path is
    already pinned, so the CREATEs land in that schema.
    """
    if os.environ.get("PG_SKIP_SCHEMA", "").strip().lower() in {"1", "true", "yes"}:
        return
    con.execute(_SCHEMA_SQL.read_text())


def _validate(spec: dict) -> None:
    missing = [k for k in _REQUIRED if k not in spec]
    if missing:
        raise ValueError(f"experiment missing required keys: {missing}")
    fid = spec["fidelity"]
    if not isinstance(fid, dict):
        raise ValueError("fidelity must be an object with slippage floors")
    fmissing = [k for k in _REQUIRED_FIDELITY if k not in fid]
    if fmissing:
        raise ValueError(f"fidelity missing required floors: {fmissing}")


def register(con, json_path: str | Path) -> str:
    """Register one experiment from a JSON file. Immutable: refuses to overwrite
    an existing id. Returns the registered id."""
    spec = json.loads(Path(json_path).read_text())
    _validate(spec)
    exp_id = spec["id"]
    if get(con, exp_id) is not None:
        raise ValueError(f"experiment {exp_id!r} already registered (immutable)")
    con.execute(
        "INSERT INTO experiments "
        "(id, registered_at, hypothesis, population, intervention, bars, futility, "
        " fidelity, status) "
        "VALUES (?, ?, ?, ?::jsonb, ?::jsonb, ?::jsonb, ?::jsonb, ?::jsonb, 'running')",
        [
            exp_id,
            spec.get("registered_at") or datetime.now(timezone.utc).isoformat(),
            spec["hypothesis"],
            json.dumps(spec["population"]),
            json.dumps(spec["intervention"]),
            json.dumps(spec["bars"]),
            json.dumps(spec.get("futility")) if spec.get("futility") is not None else None,
            json.dumps(spec["fidelity"]),
        ],
    )
    return exp_id


def register_cohort(con, cohort_dir: str | Path | None = None) -> list[str]:
    """Register every ``experiments/*.json`` not already registered (idempotent).

    The declared cohort must be in the ``experiments`` table before ``run`` can
    replay anything — otherwise the harness silently does nothing on a fresh DB.
    ``run`` calls this at startup so a fresh deployment self-bootstraps with no
    manual ``register`` step. Already-registered ids are skipped (immutable);
    a malformed file is logged and skipped, never fatal to the run."""
    apply_schema(con)
    d = Path(cohort_dir) if cohort_dir else _COHORT_DIR
    newly: list[str] = []
    for f in sorted(d.glob("*.json")):
        try:
            newly.append(register(con, f))
        except ValueError as exc:
            # already registered (the common nightly case) or malformed — skip it
            if "already registered" not in str(exc):
                import logging
                logging.getLogger(__name__).warning("skip cohort file %s: %s", f.name, exc)
    return newly


def _row_to_dict(row) -> dict:
    return {
        "id": row[0],
        "registered_at": row[1],
        "hypothesis": row[2],
        "population": row[3],
        "intervention": row[4],
        "bars": row[5],
        "futility": row[6],
        "fidelity": row[7],
        "status": row[8],
        "verdict_at": row[9],
        "verdict_note": row[10],
    }


_COLS = ("id, registered_at, hypothesis, population, intervention, bars, futility, "
         "fidelity, status, verdict_at, verdict_note")


def get(con, exp_id: str) -> dict | None:
    row = con.execute(
        f"SELECT {_COLS} FROM experiments WHERE id = ?", [exp_id]).fetchone()
    return _row_to_dict(row) if row else None


def list_experiments(con, running_only: bool = False) -> list[dict]:
    sql = f"SELECT {_COLS} FROM experiments"
    if running_only:
        sql += " WHERE status = 'running'"
    sql += " ORDER BY registered_at"
    return [_row_to_dict(r) for r in con.execute(sql).fetchall()]