research/tools/simulators/tsi/tsi-sim-pernode/src/tsi_sim/sweep.py

"""Parameter-sweep expansion, parallel execution, and result persistence.

Across-config parallelism is the main multicore lever: ``run_trajectory`` is a pure
function of an immutable, hash-seeded ``SimConfig``, so results are bitwise
order-independent and the grid is embarrassingly parallel. We use joblib's process-based
**loky** backend (this workload is CPU-bound Python that holds the GIL, so threads would
serialise), and pin each worker to a single BLAS thread to avoid oversubscription.
"""

from __future__ import annotations

import argparse
import datetime as _dt
import multiprocessing as mp
import os
import queue
import sys
import time
from dataclasses import dataclass, replace
from pathlib import Path

import numpy as np
import pandas as pd
import yaml
from joblib import Parallel, delayed
from tqdm import tqdm

from .config import SimConfig, SweepConfig
from .engine import run_trajectory
from .memguard import DEFAULT_BUDGET_FRAC, ArrivalMatrixTooLarge
from .memguard import total_ram_bytes as _total_ram_bytes

# ``calibrate="auto"`` runs a real memory probe when the analytic estimate is either extrapolated
# past the validated network size (N > CALIBRATION_N_THRESHOLD) OR simply large in absolute terms
# (per-worker estimate > CALIBRATION_BYTES_THRESHOLD) — the latter catches a block-count explosion
# from a low genesis_d_factor even at small N, which is exactly what OOM-froze the box at N=1000.
CALIBRATION_N_THRESHOLD = 2000
CALIBRATION_BYTES_THRESHOLD = 1.5 * 1024**3

# Keep numpy/BLAS single-threaded inside each worker process (belt-and-braces alongside
# joblib's inner_max_num_threads); prevents N_workers x N_blas_threads oversubscription.
for _var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS"):
    os.environ.setdefault(_var, "1")


def expected_peak_blocks(config: SimConfig) -> int:
    """Expected block count of the *heaviest* epoch (genesis), which sizes the arrival matrix.

    ``n_blocks`` is NOT ``~10*k``: it is the number of lottery wins over the epoch, which scales
    with ``sum_i phi(w_i / D_est)``. At genesis ``D_est = genesis_d_factor * D_true`` is at its
    smallest, so ``sum(stake)/D_est = 1/genesis_d_factor`` is largest and block production peaks
    there (a low ``genesis_d_factor`` can inflate it 100x — the collapsed-estimate regime). We
    realise the seeded stake and compute the expected genesis-epoch win count exactly; this is
    what previously blew the estimate up by ~90x and OOM-ed the box. Capped at ``N*E`` (every
    node can win every slot at most once).
    """
    from .lottery import win_probs
    from .rng import seedseq_for
    from .stake import make_stake

    child0 = seedseq_for(config).spawn(config.epochs + 3)[0]
    stake = make_stake(config, np.random.default_rng(child0))
    d_est_genesis = config.genesis_d_factor * float(stake.sum())
    p = win_probs(stake, d_est_genesis, config.f)
    expected = float(config.epoch_len) * float(p.sum())
    return int(min(expected, float(config.n_nodes) * config.epoch_len)) + 1


def _arrival_columns(config: SimConfig, peak_blocks: int) -> int:
    """Estimated stored arrival columns: full ``n_blocks``, or the pruned keep-span window.

    With ``prune_arrival`` the arrival buffer keeps only blocks inside ``max(horizon, W)``
    slots, so its width is ``~keepspan * blocks_per_slot`` regardless of how far the block count
    exploded. We don't have the graph here, so ``horizon`` is approximated generously (dominated by
    ``uncle_window``, plus the blend mix cascade); the exact buffer size is guarded in-worker.
    """
    if not (config.prune_arrival and config.windowed_fork_choice):
        return peak_blocks
    per_slot = peak_blocks / config.epoch_len if config.epoch_len else peak_blocks
    keepspan = float(config.effective_uncle_window)
    if config.topology == "blend":
        lat = max(config.link_latency_mean, 0.1)
        keepspan = max(keepspan, (config.blend_hops + 1) * lat * 4
                       + config.blend_hops * config.blend_delay_max)
    buf_width = 2.0 * (keepspan * per_slot * 1.5 + 64)     # ~ _build_pruned's 2*(cap + slack)
    return int(min(peak_blocks, buf_width)) + 1


def estimate_worker_bytes(config: SimConfig) -> int:
    """Rough peak RSS of one worker running ``config``.

    The arrival buffer and the dense ``path_latency`` (``N x N`` float64, plus Dijkstra scratch)
    dominate; everything else (block-tree 1-D arrays, measurement temp) is smaller. The arrival
    term is ``N * columns * 8`` where ``columns`` is the full peak-epoch block count
    (``expected_peak_blocks``, capturing a low-``genesis_d_factor`` explosion) or, with
    ``prune_arrival``, just the bounded keep-span window (``_arrival_columns``). Factors are
    generous so the cap errs toward fewer, safe workers.
    """
    n = config.n_nodes
    cols = _arrival_columns(config, expected_peak_blocks(config))
    arrival = 1.35 * n * cols * 8                       # buffer + measurement/fork-choice temps
    path_latency = 2.2 * n * n * 8                      # dense (N,N) + Dijkstra predecessor/scratch
    fixed = 250 * 1024**2                               # python + numpy + numba loaded per process
    return int(arrival + path_latency + fixed)


def _ru_maxrss_bytes(ru_maxrss: int) -> int:
    """Normalise ``getrusage`` peak RSS to bytes (macOS reports bytes, Linux kibibytes)."""
    return int(ru_maxrss) if sys.platform == "darwin" else int(ru_maxrss) * 1024


def _calibration_target(config: SimConfig, q: mp.Queue) -> None:  # pragma: no cover - subprocess
    """Child entry point: run ONE epoch of ``config`` and report this process's peak RSS.

    One epoch reaches the same peak as a full trajectory — the arrival matrix ``A`` and the
    ``path_latency`` matrix are (re)built every epoch and the measurement temporary peaks within
    an epoch — so a single epoch is a faithful, cheap probe of a worker's high-water mark.

    The probe bounds ITSELF to a fraction of physical RAM (overriding any inherited budget): it
    must be free to allocate the real peak in order to measure it, but must still fail loud rather
    than freeze if the heaviest config exceeds the box — in which case the parent falls back to the
    (large) analytic estimate.
    """
    import resource

    from .engine import run_trajectory
    from .memguard import total_ram_bytes
    os.environ["TSI_ARRIVAL_BYTES_BUDGET"] = str(int(DEFAULT_BUDGET_FRAC * total_ram_bytes()))
    run_trajectory(replace(config, epochs=1))
    q.put(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)


def measure_worker_bytes(config: SimConfig, timeout: float = 900.0) -> int | None:
    """Peak RSS (bytes) of a real worker running ``config``, measured in a fresh spawned
    process — the same isolation loky gives each worker, so numba/numpy/scratch are all counted.

    Returns ``None`` if the probe cannot start, crashes, or exceeds ``timeout`` (the caller then
    falls back to the analytic estimate). ``spawn`` matches loky and keeps the probe independent
    of the parent's already-imported modules; it needs an importable ``__main__`` (a real script
    or ``-m`` entry point), so from a bare REPL/stdin the child fails to bootstrap and we return
    ``None`` promptly by polling liveness rather than blocking the full timeout.
    """
    ctx = mp.get_context("spawn")
    q: mp.Queue = ctx.Queue()
    proc = ctx.Process(target=_calibration_target, args=(config, q), daemon=True)
    proc.start()
    deadline = time.monotonic() + timeout
    try:
        while True:
            try:
                ru_maxrss = q.get(timeout=1.0)
                return _ru_maxrss_bytes(ru_maxrss)
            except queue.Empty:
                if not proc.is_alive():        # child died (bootstrap failure / crash / OOM-kill)
                    return None
                if time.monotonic() >= deadline:
                    return None
    finally:
        if proc.is_alive():
            proc.terminate()
        proc.join(timeout=5.0)


@dataclass
class WorkerPlan:
    """The chosen worker count plus the numbers behind it (for logging)."""
    n_jobs: int
    per_worker_bytes: int
    calibrated: bool
    ram_bytes: int
    mem_frac: float


def plan_workers(requested: int, configs: list[SimConfig], mem_frac: float,
                 calibrate: str = "auto") -> WorkerPlan:
    """Size the worker pool so ``n_jobs * per-worker peak`` fits ``mem_frac`` of physical RAM.

    Per-worker peak is the analytic ``estimate_worker_bytes`` by default, but is replaced by a
    **measured** peak RSS (one epoch of the heaviest config, in a spawned process) when the
    calibration probe fires:

    * ``calibrate="auto"`` (default) — probe when the estimate is extrapolated past the validated
      network size (N > ``CALIBRATION_N_THRESHOLD``) OR is large in absolute terms
      (> ``CALIBRATION_BYTES_THRESHOLD``); the latter catches a low-``genesis_d_factor`` block
      explosion even at small N;
    * ``calibrate="always"`` — always probe;
    * ``calibrate="never"`` — never probe (analytic estimate only).

    ``requested`` follows joblib's convention (``-1`` = all logical cores); ``mem_frac <= 0``
    disables the cap entirely.
    """
    cores = os.cpu_count() or 1
    want = cores if requested < 0 else max(1, requested)
    ram = _total_ram_bytes()
    if mem_frac <= 0.0 or not configs:
        return WorkerPlan(want, 0, False, ram, mem_frac)

    heaviest = max(configs, key=estimate_worker_bytes)
    per_worker = estimate_worker_bytes(heaviest)
    calibrated = False
    do_probe = calibrate == "always" or (calibrate == "auto" and (
        heaviest.n_nodes > CALIBRATION_N_THRESHOLD
        or per_worker > CALIBRATION_BYTES_THRESHOLD))
    if do_probe:
        measured = measure_worker_bytes(heaviest)
        if measured is not None:
            per_worker = int(measured * 1.1)        # 10% headroom over the measured high-water mark
            calibrated = True

    budget = int(mem_frac * ram)
    fit = max(1, budget // per_worker)
    return WorkerPlan(min(want, fit), per_worker, calibrated, ram, mem_frac)


def run_sweep(
    sweep: SweepConfig, n_jobs: int = -1, progress: bool = True, batch_size: str | int = "auto",
    mem_frac: float = 0.7, calibrate: str = "auto",
) -> pd.DataFrame:
    """Expand the grid, run every config across cores (loky), return one big frame.

    ``n_jobs=-1`` uses all logical cores, but the count is then **capped** so that
    ``workers * per-worker peak RSS`` stays under ``mem_frac`` of physical RAM (``mem_frac=0``
    uncaps concurrency but still keeps the per-process fail-loud guard). This prevents the
    ``(N x n_blocks)`` arrival matrix from OOM-ing the box when many heavy configs run at once —
    including the collapsed-``D_est`` regime where a low ``genesis_d_factor`` explodes ``n_blocks``
    (see ``expected_peak_blocks``). Per-worker peak is measured by a **calibration probe** when
    heavy (see ``plan_workers`` / ``calibrate``), and each worker additionally enforces an
    in-process budget (``TSI_ARRIVAL_BYTES_BUDGET`` on both ``A`` and ``path_latency``) so a
    mis-estimated config fails loud instead of freezing the box. ``batch_size`` defaults to joblib
    "auto"; pass ``1`` for a small grid of heavy full-scale configs.
    """
    configs = sweep.expand()
    plan = plan_workers(n_jobs, configs, mem_frac, calibrate)
    # Per-worker byte budget for the in-worker guard: each worker's share of the RAM budget.
    # Enforced in build_tree_pernode (A) and build_path_latency (N^2) -> ArrivalMatrixTooLarge.
    # "0" is NOT "unlimited": memguard resolves it to DEFAULT_BUDGET_FRAC of RAM, so even a
    # `mem_frac=0` run keeps an absolute per-process ceiling (no single alloc can freeze the box).
    worker_budget = (int(mem_frac * plan.ram_bytes // plan.n_jobs)
                     if (mem_frac > 0 and plan.n_jobs) else 0)
    os.environ["TSI_ARRIVAL_BYTES_BUDGET"] = str(worker_budget)
    if progress:
        src = "measured" if plan.calibrated else "estimated"
        gb = plan.per_worker_bytes / 1024**3
        print(f"[sweep] {len(configs)} configs; ~{gb:.2f} GB/worker ({src}); "
              f"using {plan.n_jobs} worker(s) of {os.cpu_count()} cores "
              f"(RAM {plan.ram_bytes / 1024**3:.0f} GB x cap {mem_frac:g})")
        if worker_budget and plan.per_worker_bytes > worker_budget:
            print(f"[sweep] WARNING: a single config's estimated peak "
                  f"({gb:.1f} GB) exceeds its per-worker RAM share "
                  f"({worker_budget / 1024**3:.1f} GB); it may swap or hit the arrival-matrix "
                  f"guard. Raise genesis_d_factor, lower n_nodes/k/epochs, or raise --mem-frac.")
    runner = Parallel(
        n_jobs=plan.n_jobs,
        backend="loky",
        inner_max_num_threads=1,
        batch_size=batch_size,
        return_as="generator",
    )(delayed(run_trajectory)(c) for c in configs)
    if progress:
        runner = tqdm(runner, total=len(configs), desc="configs")
    rows: list[dict] = []
    try:
        for traj in runner:
            rows.extend(traj)
    except ArrivalMatrixTooLarge as exc:
        raise RuntimeError(
            f"aborting sweep: {exc} (worker memory guard tripped — the per-worker estimate was "
            f"too low, likely an under-modelled block explosion)."
        ) from exc
    return pd.DataFrame(rows)


def persist(df: pd.DataFrame, path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.suffix == ".parquet":
        df.to_parquet(path, index=False)
    else:
        df.to_csv(path, index=False)


def load_sweep_yaml(path: str | Path) -> SweepConfig:
    with open(path) as fh:
        data = yaml.safe_load(fh)
    return SweepConfig.from_dict(data)


def new_run_dir(outdir: str | Path, label: str) -> Path:
    """Create a fresh timestamped run directory so runs never overwrite each other."""
    ts = _dt.datetime.now().strftime("%Y-%m-%d_%H%M%S")
    base = Path(outdir) / f"{ts}_{label}"
    run_dir, i = base, 2
    while run_dir.exists():                 # avoid same-second collisions
        run_dir = base.with_name(f"{base.name}_{i}")
        i += 1
    run_dir.mkdir(parents=True)
    return run_dir


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Run a TSI parameter sweep + figures")
    parser.add_argument("--config", required=True, help="sweep YAML")
    parser.add_argument("--outdir", default="runs",
                        help="parent dir; a dated sub-folder is created per run")
    parser.add_argument("--label", default=None, help="run label (default: config name)")
    parser.add_argument("--n-jobs", type=int, default=-1, help="-1 = all logical cores")
    parser.add_argument("--batch-size", default="auto",
                        help="'auto' (default) or an int; use 1 for a small heavy grid")
    parser.add_argument("--mem-frac", type=float, default=0.7,
                        help="cap concurrent workers to this fraction of physical RAM "
                             "(0 disables the cap)")
    parser.add_argument("--calibrate", choices=["auto", "always", "never"], default="auto",
                        help="measure a real worker's peak RSS to size the pool: 'auto' "
                             "(default) probes when N>2000, 'always', or 'never' (estimate only)")
    parser.add_argument("--no-figures", action="store_true",
                        help="skip auto figure generation")
    parser.add_argument("--old", action="store_true",
                        help="run the old (pre countable redesign) uncle model: window = "
                             "uncle_window slots, any-depth orphans referenceable, every "
                             "baked reference counts; bit-reproduces historical runs")
    args = parser.parse_args(argv)

    batch_size = int(args.batch_size) if args.batch_size != "auto" else "auto"
    label = args.label or (Path(args.config).stem + ("-old" if args.old else ""))
    run_dir = new_run_dir(args.outdir, label)

    sweep = load_sweep_yaml(args.config)
    if args.old:
        sweep.base["uncle_model"] = "old"
    df = run_sweep(sweep, n_jobs=args.n_jobs, batch_size=batch_size, mem_frac=args.mem_frac,
                   calibrate=args.calibrate)
    results_path = run_dir / "results.parquet"
    persist(df, results_path)
    key_cols = ["n_nodes", "stake_dist", "topology", "degree", "link_latency_mean",
                "latency", "uncle_model", "max_uncles", "uncle_strategy", "init_dest",
                "replicate"]
    n_cfg = len(df[key_cols].drop_duplicates())
    print(f"wrote {len(df)} rows ({n_cfg} configs) -> {results_path}")

    if not args.no_figures:
        from .plotting.make_figures import render
        n_fig = render(df, run_dir / "figures")
        print(f"wrote {n_fig} figures -> {run_dir / 'figures'}")
    print(f"run dir: {run_dir}")


if __name__ == "__main__":  # pragma: no cover
    main()