giant/giant/analysis/runtime_estimate.py

"""Per-(plot, chunk) HTCondor walltime estimates for `giant analyze submit`.

Each catalog spec's compute cost is close to linear in the number of input
rows a `compute-one` job streams over — every spec is one (or a couple of)
streaming `group_by` pass(es) over the chunk (see `catalog.py`/`reduce.py`).
`_COST_MODEL` below is ``spec_id -> (intercept_s, seconds_per_row)``.
``n_rows`` is the combined rollout+reference row count of the job's input:
the chunk's row count for `chunkable=True` specs, the whole dataset's for the
three `chunkable=False` router specs (they always run as a single job
regardless of chunk count).

Calibrated 2026-07-27 from real HTCondor timings (`condor_history`
``RemoteWallClockTime``) of a production run: prediction ``563f5ee3``
(PbWO4, 50 GeV) analyzed with ``--chunks 4`` against
``giant/analysis/runtime_estimate.py``'s prior (local-synthetic-only) model —
see the ``analysis-rollout-plots`` branch history for the raw data. That run's
4 chunks came out at nearly identical row counts (~63-64M rows each, ~254M
total), so this real data has no genuine row-count spread to fit a slope
against — instead each spec's ``per_row`` here is a single line through the
origin (``intercept=0``) hitting that spec's *median* wall-clock time across
its 4 chunks at that run's row count. A handful of (spec, chunk) pairs showed
3-8x spikes in one chunk only (e.g. ``marginal_edep_by_material``: 88, 88, 90,
722s) — almost certainly shared ``/ceph`` contention from ~130 jobs landing on
the filesystem at once right after submission, not a real per-row cost, so
the median (not the max) was fit to avoid baking that noise into a rate that
would then wrongly scale up with a bigger dataset. `RUNTIME_SAFETY_MARGIN` is
deliberately generous (4x total) specifically to absorb that kind of
contention spike instead. Rerun this calibration (pull fresh
`condor_history`/`run_meta.json`, refit) if the catalog changes or timings
drift — a synthetic local rebaseline via `giant/tools/profile_analysis_costs.py`
is a reasonable fallback when no real cluster data is available yet, but
undershoots real wall time badly (it can't see docker pull / `/ceph` I/O
latency), which is exactly why this file moved off it.
"""

from __future__ import annotations

import math

# Multiplicative pad applied to every job's estimated walltime. The one knob
# this feature was asked to expose. Set generously (4x total, i.e. 3.0 here)
# to absorb the shared-/ceph-contention spikes described above rather than
# encoding them into individual specs' per-row rates.
RUNTIME_SAFETY_MARGIN = 3.00

# Fixed per-job overhead (docker start, `.venv/bin/giant` startup, initial
# `/ceph` read latency) — calibrated as the fastest observed real spec
# (`leakage_fraction`, median 51s) rounded up, since even the cheapest spec
# streams the whole chunk once.
_FIXED_OVERHEAD_S = 60.0

# Router diagnostics run a live torch checkpoint (bounded inference over
# <=200k subsampled rows, independent of chunk size) instead of a row-based
# scan. Calibrated from the 3 real router jobs' observed wall times (119, 66,
# 124s) — max minus _FIXED_OVERHEAD_S, on top of it.
_ROUTER_FIXED_S = 64.0
_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process"})

# Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot
# added after the last calibration run) — the most expensive fitted per-row
# rate observed, plus a small constant pad.
_DEFAULT_COST = (5.0, 3.0e-6)

# spec_id -> (intercept_s, seconds_per_row), fit 2026-07-27 from real
# HTCondor `RemoteWallClockTime` (see module docstring for methodology).
_COST_MODEL: dict[str, tuple[float, float]] = {
    "marginal_step_length": (0.0, 5.199e-07),
    "marginal_step_length_by_energy": (0.0, 1.678e-06),
    "marginal_step_length_by_pdg": (0.0, 4.569e-07),
    "marginal_step_length_by_material": (0.0, 2.269e-06),
    "marginal_edep": (0.0, 2.804e-06),
    "marginal_edep_by_energy": (0.0, 1.386e-06),
    "marginal_edep_by_pdg": (0.0, 4.490e-07),
    "marginal_edep_by_material": (0.0, 4.333e-07),
    "marginal_delta_e": (0.0, 5.042e-07),
    "marginal_delta_e_by_energy": (0.0, 1.678e-06),
    "marginal_delta_e_by_pdg": (0.0, 4.727e-07),
    "marginal_delta_e_by_material": (0.0, 4.490e-07),
    "marginal_post_E": (0.0, 4.805e-07),
    "marginal_post_E_by_energy": (0.0, 1.284e-06),
    "marginal_post_E_by_pdg": (0.0, 4.569e-07),
    "marginal_post_E_by_material": (0.0, 4.490e-07),
    "marginal_cos_scatter": (0.0, 5.436e-07),
    "marginal_cos_scatter_by_energy": (0.0, 1.363e-06),
    "marginal_cos_scatter_by_pdg": (0.0, 4.727e-07),
    "marginal_cos_scatter_by_material": (0.0, 4.727e-07),
    "event_total_edep": (0.0, 4.490e-07),
    "event_total_edep_by_energy": (0.0, 4.411e-07),
    "event_mean_length": (0.0, 4.333e-07),
    "event_n_steps": (0.0, 4.569e-07),
    "shower_longitudinal": (0.0, 2.348e-06),
    "shower_transverse": (0.0, 2.899e-06),
    "species_edep_share": (0.0, 4.333e-07),
    "leakage_fraction": (0.0, 0.0),
    "sec_count_per_event": (0.0, 4.727e-07),
    "sec_count_per_species": (0.0, 4.963e-07),
    "sec_energy": (0.0, 4.727e-07),
    "sec_cos_angle": (0.0, 2.749e-06),
}


def estimate_runtime_s(spec_id: str, n_rows: int) -> int:
    """Estimated `+RequestWalltime` (seconds) for one (plot, chunk) job.

    ``n_rows`` is the rollout+reference row count of that job's input slice.
    Includes `_FIXED_OVERHEAD_S`/`_ROUTER_FIXED_S` and `RUNTIME_SAFETY_MARGIN`
    — callers should pass this straight through to the submit description.
    """
    if spec_id in _ROUTER_IDS:
        compute_s = _ROUTER_FIXED_S
    else:
        intercept, per_row = _COST_MODEL.get(spec_id, _DEFAULT_COST)
        compute_s = intercept + per_row * n_rows
    total = _FIXED_OVERHEAD_S + compute_s
    return math.ceil(total * (1 + RUNTIME_SAFETY_MARGIN))