From e380400fe9acd6784c4c73b96ac4f3d1375ca39a Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 27 Jul 2026 09:47:44 +0200 Subject: [PATCH] analyze: estimate per-job HTCondor walltime from chunk row count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 3 + giant/analysis/condor.py | 77 ++++++++-- giant/analysis/runtime_estimate.py | 102 +++++++++++++ scripts/profile_analysis_costs.py | 231 +++++++++++++++++++++++++++++ tests/test_condor.py | 53 ++++++- 5 files changed, 448 insertions(+), 18 deletions(-) create mode 100644 giant/analysis/runtime_estimate.py create mode 100644 scripts/profile_analysis_costs.py diff --git a/giant/analysis/__init__.py b/giant/analysis/__init__.py index 8af48cb..48c4343 100644 --- a/giant/analysis/__init__.py +++ b/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", ] diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index 606d93b..d349117 100644 --- a/giant/analysis/condor.py +++ b/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 - (``/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 (``/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)) diff --git a/giant/analysis/runtime_estimate.py b/giant/analysis/runtime_estimate.py new file mode 100644 index 0000000..29a7f95 --- /dev/null +++ b/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)) diff --git a/scripts/profile_analysis_costs.py b/scripts/profile_analysis_costs.py new file mode 100644 index 0000000..bbb0a98 --- /dev/null +++ b/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() diff --git a/tests/test_condor.py b/tests/test_condor.py index 6634bd8..1dc5ad2 100644 --- a/tests/test_condor.py +++ b/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