analyze: estimate per-job HTCondor walltime from chunk row count

Each condor job's +RequestWalltime used to be one flat 3600s default
for every (plot, chunk), regardless of how much data it actually
streams over. `prep` now records each chunk's rollout+reference row
count, and `giant/analysis/runtime_estimate.py` turns that into a
per-job estimate: a per-spec (intercept, seconds/row) cost model fit
by `scripts/profile_analysis_costs.py` against synthetic mock data on
this machine, plus a fixed overhead placeholder (docker/uv/shared-fs
startup — unmeasurable here, no /ceph access) and a single
RUNTIME_SAFETY_MARGIN multiplier. jobs.txt gains a walltime column and
the submit description references it via $(walltime) instead of a
constant.

giant/analysis/__init__.py

@@ -25,6 +25,7 @@ from giant.analysis.condor import (
 )
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial, Reduced
+from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
 from giant.analysis.sources import Side
 
 __all__ = [
@@ -46,4 +47,6 @@ __all__ = [
     "Partial",
     "Reduced",
     "Side",
+    "RUNTIME_SAFETY_MARGIN",
+    "estimate_runtime_s",
 ]

giant/analysis/condor.py

@@ -42,14 +42,17 @@ HTCondor file transfer of the multi-GB inputs.
 from __future__ import annotations
 
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 
+import polars as pl
 import yaml
 
 from giant.analysis.catalog import Bundle, catalog_ids, get_spec
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial
+from giant.analysis.runtime_estimate import estimate_runtime_s
+from giant.analysis.sources import Side, open_side
 
 # Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
 _PLOT_META_KEYS = (
@@ -113,6 +116,11 @@ class RunMeta:
     title: str
     plot_meta: dict
     n_chunks: int = 1
+    # rollout+reference row count of each event_id-disjoint chunk, and the
+    # dataset total — inputs to `runtime_estimate.estimate_runtime_s`. Empty/0
+    # on run directories written before this field existed.
+    rows_per_chunk: list[int] = field(default_factory=list)
+    total_rows: int = 0
 
     def save(self, path: str | Path) -> None:
         Path(path).write_text(json.dumps(self.__dict__, indent=2))
@@ -122,6 +130,30 @@ class RunMeta:
         return cls(**json.loads(Path(path).read_text()))
 
 
+def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -> list[int]:
+    """Rollout+reference row count of each ``event_id % n_chunks`` chunk.
+
+    One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
+    the sizing input every job's estimated walltime
+    (``runtime_estimate.estimate_runtime_s``) is computed from.
+    """
+
+    def counts(lf: pl.LazyFrame) -> pl.DataFrame:
+        return (
+            lf.select((pl.col("event_id") % n_chunks).alias("_c"))
+            .group_by("_c")
+            .agg(pl.len().alias("n"))
+            .collect(engine="streaming")
+        )
+
+    out = [0] * n_chunks
+    for lf in (open_side(rollout, Side.rollout), open_side(reference, Side.reference)):
+        df = counts(lf)
+        for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
+            out[c] += n
+    return out
+
+
 def prep(
     rollout_yaml: str | Path,
     run_dir: str | Path | None = None,
@@ -144,6 +176,8 @@ def prep(
     ctx = build_context(rollout, reference, **ctx_kwargs)
     ctx.save(run_path / "shared.json")
 
+    rows_per_chunk = _rows_per_chunk(rollout, reference, n_chunks)
+
     ckpt = Path(y.get("checkpoint", "")).name or "rollout"
     RunMeta(
         rollout=str(rollout),
@@ -152,6 +186,8 @@ def prep(
         title=f"GIANT rollout analysis — {ckpt}",
         plot_meta=_plot_meta(y),
         n_chunks=n_chunks,
+        rows_per_chunk=rows_per_chunk,
+        total_rows=sum(rows_per_chunk),
     ).save(run_path / "run_meta.json")
     return run_path
 
@@ -273,7 +309,6 @@ class SubmitConfig:
     docker_image: str = "mschnepf/slc7-condocker"
     request_memory_mb: int = 4096
     request_cpus: int = 1
-    request_walltime_s: int = 3600
     remote: bool = False  # +RemoteJob (grid I/O) vs ProvidesETPResources (local files)
     n_chunks: int = 1  # per-plot data chunks; ignored for chunkable=False specs
 
@@ -300,25 +335,45 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> st
         "when_to_transfer_output = ON_EXIT\n"
         f"request_memory = {cfg.request_memory_mb}\n"
         f"request_cpus = {cfg.request_cpus}\n"
-        f"+RequestWalltime = {cfg.request_walltime_s}\n"
+        "+RequestWalltime = $(walltime)\n"
         f"accounting_group = {cfg.accounting_group}\n"
         f"{reqs_attrs}"
         f"output = {cfg.run_dir}/logs/$(plotid)__$(chunk).out\n"
         f"error = {cfg.run_dir}/logs/$(plotid)__$(chunk).err\n"
         f"log = {cfg.run_dir}/logs/condor.log\n"
-        f"queue plotid,chunk from {jobs_file}\n"
+        f"queue plotid,chunk,walltime from {jobs_file}\n"
     )
 
 
+def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]:
+    """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.
+
+    Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
+    ``chunkable=False`` specs (router diagnostics) always use the dataset
+    total since they run as a single job regardless of ``n_chunks``.
+    """
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    jobs: list[tuple[str, int, int]] = []
+    for spec_id in ids:
+        chunkable = get_spec(spec_id).chunkable
+        chunks = range(n_chunks) if chunkable else [0]
+        for chunk in chunks:
+            n_rows = meta.rows_per_chunk[chunk] if chunkable else meta.total_rows
+            jobs.append((spec_id, chunk, estimate_runtime_s(spec_id, n_rows)))
+    return jobs
+
+
 def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     """Write the wrapper script, (plot, chunk) job list, and HTCondor submit
     description.
 
     Each catalog id gets ``cfg.n_chunks`` jobs, except ``chunkable=False``
     specs (the router diagnostics), which always get exactly one regardless of
-    ``cfg.n_chunks``. Returns the submit description path
-    (``<run_dir>/analyze.sub``). Does not submit — call ``condor_submit`` on
-    the returned file.
+    ``cfg.n_chunks``. Every job's ``+RequestWalltime`` is estimated from its
+    chunk's row count (``runtime_estimate.estimate_runtime_s``, requires
+    ``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``).
+    Returns the submit description path (``<run_dir>/analyze.sub``). Does not
+    submit — call ``condor_submit`` on the returned file.
     """
     ids = ids or catalog_ids()
     run_dir = cfg.run_dir
@@ -330,13 +385,9 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, run_dir=run_dir))
     wrapper.chmod(0o755)
 
-    jobs = [
-        (spec_id, chunk)
-        for spec_id in ids
-        for chunk in range(cfg.n_chunks if get_spec(spec_id).chunkable else 1)
-    ]
+    jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
     jobs_file = run_dir / "jobs.txt"
-    jobs_file.write_text("\n".join(f"{i},{k}" for i, k in jobs) + "\n")
+    jobs_file.write_text("\n".join(f"{i},{k},{w}" for i, k, w in jobs) + "\n")
 
     sub = run_dir / "analyze.sub"
     sub.write_text(_submit_description(cfg, wrapper, jobs_file))

giant/analysis/runtime_estimate.py

@@ -0,0 +1,102 @@
+"""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)``, fit by
+least squares against wall-clock timings of `compute_reduced` on synthetic
+mock data of increasing size, run on a local dev machine (see
+`scripts/profile_analysis_costs.py` — rerun it and paste the new numbers in
+here if the catalog changes or this needs recalibrating). ``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).
+
+The fitted numbers only capture *local, in-memory compute* — they don't (and
+from a laptop with no `/ceph` access, can't) capture the real condor job's
+docker pull, `uv run` cold start, or shared-filesystem read latency, which in
+practice likely dominate total wall time for anything but a huge single-chunk
+job. `_FIXED_OVERHEAD_S` is a deliberately generous placeholder for all of
+that combined; recalibrate it from real `condor_q`/log timings once some are
+available, rather than trusting it as measured.
+"""
+
+from __future__ import annotations
+
+import math
+
+# Multiplicative pad applied to every job's estimated walltime. The one knob
+# this feature was asked to expose.
+RUNTIME_SAFETY_MARGIN = 1.00
+
+# Docker image pull + `uv run` startup + shared (/ceph, ETP) filesystem read
+# latency — not measurable on a machine with no /ceph access, so this is a
+# conservative placeholder rather than a fit. Recalibrate from real job logs.
+_FIXED_OVERHEAD_S = 600.0
+
+# Router diagnostics need a live torch checkpoint to do any real work; this
+# machine has none, so their cost (torch.load + a bounded inference pass over
+# <= 200k subsampled rows, independent of chunk size) couldn't be profiled
+# either. Fixed budget, on top of _FIXED_OVERHEAD_S, instead of a row-based fit.
+_ROUTER_FIXED_S = 300.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 profiling run) — the most expensive fitted (intercept,
+# seconds/row) pair observed, rounded up.
+_DEFAULT_COST = (0.02, 2.0e-6)
+
+# spec_id -> (intercept_s, seconds_per_row), fit on this machine 2026-07-27
+# via `scripts/profile_analysis_costs.py` against SIDE_ROW_COUNTS up to 2M
+# rows/side (4M combined).
+_COST_MODEL: dict[str, tuple[float, float]] = {
+    "marginal_step_length": (0.003581, 0.000000042),
+    "marginal_step_length_by_energy": (0.010336, 0.000000081),
+    "marginal_step_length_by_pdg": (0.007231, 0.000000042),
+    "marginal_step_length_by_material": (0.004662, 0.000000048),
+    "marginal_edep": (0.003996, 0.000000042),
+    "marginal_edep_by_energy": (0.007977, 0.000000085),
+    "marginal_edep_by_pdg": (0.003535, 0.000000042),
+    "marginal_edep_by_material": (0.003893, 0.000000047),
+    "marginal_delta_e": (0.001380, 0.000000053),
+    "marginal_delta_e_by_energy": (0.006543, 0.000000093),
+    "marginal_delta_e_by_pdg": (0.001761, 0.000000054),
+    "marginal_delta_e_by_material": (0.000000, 0.000000168),
+    "marginal_post_E": (0.000000, 0.000000146),
+    "marginal_post_E_by_energy": (0.000000, 0.000000509),
+    "marginal_post_E_by_pdg": (0.000000, 0.000000149),
+    "marginal_post_E_by_material": (0.000000, 0.000000168),
+    "marginal_cos_scatter": (0.000000, 0.000000249),
+    "marginal_cos_scatter_by_energy": (0.000000, 0.000000553),
+    "marginal_cos_scatter_by_pdg": (0.000000, 0.000000270),
+    "marginal_cos_scatter_by_material": (0.000000, 0.000000260),
+    "event_total_edep": (0.000000, 0.000000512),
+    "event_total_edep_by_energy": (0.000000, 0.000000241),
+    "event_mean_length": (0.000000, 0.000000090),
+    "event_n_steps": (0.002021, 0.000000066),
+    "shower_longitudinal": (0.000000, 0.000001804),
+    "shower_transverse": (0.000000, 0.000001731),
+    "species_edep_share": (0.007505, 0.000000017),
+    "leakage_fraction": (0.003798, 0.000000026),
+    "sec_count_per_event": (0.006551, 0.000000054),
+    "sec_count_per_species": (0.006708, 0.000000042),
+    "sec_energy": (0.006856, 0.000000047),
+    "sec_cos_angle": (0.001780, 0.000000234),
+}
+
+
+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))

scripts/profile_analysis_costs.py

@@ -0,0 +1,231 @@
+"""Benchmark `giant analyze compute-one`'s per-job cost against synthetic data.
+
+Generates mock rollout+reference parquet files at a few row counts, times
+`compute_reduced` for every chunkable catalog spec at each size (a single
+chunk covering the whole mock file), fits a straight line (intercept, seconds
+per row) through the timings, and prints the result as a Python dict literal
+ready to paste into `giant/analysis/runtime_estimate.py::_COST_MODEL`.
+
+The three `chunkable=False` router specs (`router_gating`,
+`router_share_by_pdg`, `router_share_by_process`) need a live MoE checkpoint
+to do any real work; without one (this machine has no `/ceph` access, so no
+real checkpoint) they short-circuit almost instantly and are excluded here —
+see `runtime_estimate.py`'s `_ROUTER_FIXED_S` for how those are handled
+instead.
+
+Usage: ``uv run python scripts/profile_analysis_costs.py``
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+import polars as pl
+
+from giant.analysis.catalog import catalog_ids, get_spec
+from giant.analysis.condor import compute_reduced
+from giant.analysis.context import build_context
+
+# Row counts (per side) to benchmark at. Kept in local memory/CPU range so the
+# whole sweep finishes in about a minute; the fit is linear so it extrapolates
+# fine to real multi-GB rollouts.
+SIDE_ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000]
+
+_MATERIALS = ["G4_PbWO4", "G4_Pb", "G4_lAr", "G4_Si"]
+_PDGS = [11, -11, 22, 2112, 2212, 211, -211, 13]
+_ROUTER_IDS = {"router_gating", "router_share_by_pdg", "router_share_by_process"}
+
+
+def _unit_vectors(n: int, rng: np.random.Generator) -> np.ndarray:
+    v = rng.normal(size=(n, 3))
+    return v / np.linalg.norm(v, axis=1, keepdims=True)
+
+
+def _ragged_lists(k: np.ndarray, rng: np.random.Generator, lo: float, hi: float):
+    total = int(k.sum())
+    flat = rng.uniform(lo, hi, size=total)
+    idx = np.cumsum(k)[:-1]
+    return [arr.tolist() for arr in np.split(flat, idx)]
+
+
+def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
+    rng = np.random.default_rng(seed)
+    event_id = rng.integers(0, n_events, size=n)
+    is_secondary = rng.random(n) < 0.15  # generation>0, step_no==0 birth rows
+    is_synthetic = rng.random(n) < 0.05  # bookkeeping termination rows
+
+    pre_E = rng.lognormal(mean=3.0, sigma=1.5, size=n)
+    edep = rng.uniform(0, 1, size=n) * pre_E * 0.3
+    post_E = np.clip(pre_E - edep, 0.0, None)
+    pre_dir = _unit_vectors(n, rng)
+    post_dir = _unit_vectors(n, rng)
+    pos = rng.uniform(-50, 300, size=(n, 3))
+    step_length = rng.uniform(0.1, 10.0, size=n)
+    post_pos = pos + pre_dir * step_length[:, None]
+
+    reasons = np.where(
+        is_synthetic,
+        rng.choice(
+            ["escaped", "energy_cutoff", "max_steps", "unknown_pdg"], size=n
+        ),
+        "natural_end",
+    )
+
+    return pl.DataFrame(
+        {
+            "event_id": event_id,
+            "track_id": rng.integers(0, 5, size=n),
+            "parent_id": np.where(is_secondary, 0, -1),
+            "generation": is_secondary.astype(np.int64),
+            "step_no": np.where(is_secondary, 0, rng.integers(0, 20, size=n)),
+            "pdg": rng.choice(_PDGS, size=n),
+            "pre_x": pos[:, 0],
+            "pre_y": pos[:, 1],
+            "pre_z": pos[:, 2],
+            "pre_E": pre_E,
+            "pre_dx": pre_dir[:, 0],
+            "pre_dy": pre_dir[:, 1],
+            "pre_dz": pre_dir[:, 2],
+            "post_x": post_pos[:, 0],
+            "post_y": post_pos[:, 1],
+            "post_z": post_pos[:, 2],
+            "post_E": post_E,
+            "post_dx": post_dir[:, 0],
+            "post_dy": post_dir[:, 1],
+            "post_dz": post_dir[:, 2],
+            "edep": np.where(is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep),
+            "step_length": np.where(is_synthetic, 0.0, step_length),
+            "material": rng.choice(_MATERIALS, size=n),
+            "layer_id": rng.integers(0, 30, size=n),
+            "n_sec_pred": rng.integers(0, 4, size=n),
+            "termination_reason": reasons,
+        }
+    )
+
+
+def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
+    rng = np.random.default_rng(seed + 1)
+    event_id = rng.integers(0, n_events, size=n)
+    pre_E = rng.lognormal(mean=3.0, sigma=1.5, size=n)
+    edep = rng.uniform(0, 1, size=n) * pre_E * 0.3
+    post_E = np.clip(pre_E - edep, 0.0, None)
+    pre_dir = _unit_vectors(n, rng)
+    post_dir = _unit_vectors(n, rng)
+    pos = rng.uniform(-50, 300, size=(n, 3))
+    step_length = rng.uniform(0.1, 10.0, size=n)
+    post_pos = pos + pre_dir * step_length[:, None]
+
+    k = rng.poisson(0.3, size=n).clip(max=5).astype(np.int64)
+    sec_pdg = _ragged_lists(k, rng, 0, 1)  # placeholder, overwritten below
+    sec_E = _ragged_lists(k, rng, 0.1, 50.0)
+    sec_dx = _ragged_lists(k, rng, -1.0, 1.0)
+    sec_dy = _ragged_lists(k, rng, -1.0, 1.0)
+    sec_dz = _ragged_lists(k, rng, -1.0, 1.0)
+    total = int(k.sum())
+    flat_pdg = rng.choice(_PDGS, size=total).tolist()
+    idx = np.cumsum(k)[:-1]
+    sec_pdg = [list(x) for x in np.split(np.array(flat_pdg), idx)]
+
+    return pl.DataFrame(
+        {
+            "event_id": event_id,
+            "track_id": rng.integers(0, 5, size=n),
+            "step_no": rng.integers(0, 20, size=n),
+            "pdg": rng.choice(_PDGS, size=n),
+            "pre_x": pos[:, 0],
+            "pre_y": pos[:, 1],
+            "pre_z": pos[:, 2],
+            "pre_E": pre_E,
+            "pre_dx": pre_dir[:, 0],
+            "pre_dy": pre_dir[:, 1],
+            "pre_dz": pre_dir[:, 2],
+            "post_x": post_pos[:, 0],
+            "post_y": post_pos[:, 1],
+            "post_z": post_pos[:, 2],
+            "post_E": post_E,
+            "post_dx": post_dir[:, 0],
+            "post_dy": post_dir[:, 1],
+            "post_dz": post_dir[:, 2],
+            "edep": edep,
+            "step_length": step_length,
+            "material": rng.choice(_MATERIALS, size=n),
+            "layer_id": rng.integers(0, 30, size=n),
+            "process": rng.choice(["compt", "phot", "eBrem", "eIoni", "conv"], size=n),
+            "sec_E_list": sec_E,
+            "sec_pdg_list": sec_pdg,
+            "sec_dx_list": sec_dx,
+            "sec_dy_list": sec_dy,
+            "sec_dz_list": sec_dz,
+        }
+    )
+
+
+def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path) -> float:
+    t0 = time.perf_counter()
+    compute_reduced(
+        spec_id, rollout, reference, shared, out, checkpoint=None,
+        chunk_index=0, n_chunks=1,
+    )
+    return time.perf_counter() - t0
+
+
+def main() -> None:
+    ids = [i for i in catalog_ids() if get_spec(i).chunkable]
+    timings: dict[str, list[tuple[int, float]]] = {i: [] for i in ids}
+
+    with TemporaryDirectory(prefix="giant-profile-") as tmp:
+        tmp_path = Path(tmp)
+        for n_side in SIDE_ROW_COUNTS:
+            n_events = max(n_side // 20, 10)
+            rollout = tmp_path / f"rollout_{n_side}.parquet"
+            reference = tmp_path / f"reference_{n_side}.parquet"
+            _make_rollout(n_side, n_events, seed=0).write_parquet(rollout)
+            _make_reference(n_side, n_events, seed=0).write_parquet(reference)
+
+            shared = tmp_path / f"shared_{n_side}.json"
+            ctx = build_context(
+                rollout, reference,
+                n_energy_bins=4, n_marginal_bins=50, top_k_pdg=6,
+                sample_rows=min(n_side, 200_000),
+            )
+            ctx.save(shared)
+
+            # warm the OS page cache so the timed pass measures compute, not
+            # the one-time cold read of a freshly-written file.
+            pl.scan_parquet(rollout).select(pl.len()).collect()
+            pl.scan_parquet(reference).select(pl.len()).collect()
+
+            n_rows = 2 * n_side  # rollout + reference rows in this "chunk"
+            for spec_id in ids:
+                out = tmp_path / f"{spec_id}_{n_side}.json"
+                dt = _time(spec_id, rollout, reference, shared, out)
+                timings[spec_id].append((n_rows, dt))
+                print(f"{spec_id:35s} n_rows={n_rows:>9d} time={dt:7.3f}s")
+
+            rollout.unlink()
+            reference.unlink()
+            shared.unlink()
+
+    print("\n# spec_id -> (intercept_s, seconds_per_row), fit by least squares")
+    print("_COST_MODEL: dict[str, tuple[float, float]] = {")
+    for spec_id in ids:
+        xs = np.array([n for n, _ in timings[spec_id]], dtype=float)
+        ys = np.array([t for _, t in timings[spec_id]], dtype=float)
+        slope, intercept = np.polyfit(xs, ys, 1)
+        intercept = max(intercept, 0.0)
+        slope = max(slope, 0.0)
+        print(f'    "{spec_id}": ({intercept:.6f}, {slope:.9f}),')
+    print("}")
+
+    if _ROUTER_IDS:
+        print(
+            "\n# router_* specs excluded: need a live MoE checkpoint to do real\n"
+            "# work, none available on this machine — see _ROUTER_FIXED_S instead."
+        )
+
+
+if __name__ == "__main__":
+    main()

tests/test_condor.py

@@ -93,6 +93,15 @@ def test_prep_lays_out_run_dir(tmp_path: Path):
     assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt"
     assert "best.pt" in meta.title
     assert meta.n_chunks == 1
+    assert meta.rows_per_chunk == [meta.total_rows]  # single chunk holds everything
+    assert meta.total_rows == 8  # 5 rollout rows + 3 reference rows
+
+
+def test_prep_splits_rows_per_chunk(tmp_path: Path):
+    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert len(meta.rows_per_chunk) == 2
+    assert sum(meta.rows_per_chunk) == meta.total_rows == 8
 
 
 def test_compute_one_from_run_dir(tmp_path: Path):
@@ -163,10 +172,12 @@ def test_write_submit_description(tmp_path: Path):
     assert "docker_image = mschnepf/slc7-condocker" in txt
     assert "requirements = TARGET.ProvidesETPResources" in txt
     assert "accounting_group = cms" in txt
-    assert "queue plotid,chunk from" in txt
+    assert "+RequestWalltime = $(walltime)" in txt
+    assert "queue plotid,chunk,walltime from" in txt
     jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
-    assert [i for i, _ in jobs] == catalog_ids()
-    assert all(k == "0" for _, k in jobs)  # n_chunks=1 default
+    assert [i for i, _, _ in jobs] == catalog_ids()
+    assert all(k == "0" for _, k, _ in jobs)  # n_chunks=1 default
+    assert all(int(w) > 0 for _, _, w in jobs)
     wrapper = run_dir / "run_compute.sh"
     assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
     body = wrapper.read_text()
@@ -186,14 +197,46 @@ def test_write_submit_remote_flag(tmp_path: Path):
 
 def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
     assert get_spec("router_gating").chunkable is False
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep(_write_inputs(tmp_path), chunks=4)
     cfg = SubmitConfig(
         run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
     )
     write_submit(cfg)
     jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
     counts: dict[str, int] = {}
-    for spec_id, _ in jobs:
+    for spec_id, _, _ in jobs:
         counts[spec_id] = counts.get(spec_id, 0) + 1
     assert counts["marginal_edep"] == 4
     assert counts["router_gating"] == 1  # chunkable=False, ignores n_chunks
+
+
+def test_estimate_runtime_s_scales_with_rows_and_margin():
+    from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
+    from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S
+
+    assert RUNTIME_SAFETY_MARGIN > 0
+    small = estimate_runtime_s("marginal_edep", 1_000)
+    large = estimate_runtime_s("marginal_edep", 100_000_000)
+    assert small >= (1 + RUNTIME_SAFETY_MARGIN) * _FIXED_OVERHEAD_S
+    assert large > small  # bigger chunk -> longer estimate
+
+
+def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
+    """A chunked run's later job walltimes track that chunk's row count."""
+    from giant.analysis.runtime_estimate import estimate_runtime_s
+
+    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    cfg = SubmitConfig(
+        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2
+    )
+    write_submit(cfg)
+    jobs = {
+        (i, int(k)): int(w)
+        for i, k, w in (
+            line.split(",") for line in (run_dir / "jobs.txt").read_text().split()
+        )
+    }
+    for chunk in range(2):
+        expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
+        assert jobs[("marginal_edep", chunk)] == expected