| """Pre-flight training time / cost estimate from a *measured* throughput. |
| |
| Runs the real training loop (same data pipeline, model, micro-batch, AMP, |
| compile) for `bench_seconds` on the current machine, takes the steady-state |
| samples/s (first logging window dropped: compile / warm-up), then extrapolates. |
| Run it on the rented instance itself for GPU numbers - measuring locally and |
| dividing by a guessed speed-up is exactly the kind of invented number to avoid. |
| """ |
| from __future__ import annotations |
|
|
| import csv |
| import copy |
| import tempfile |
| import time |
| from pathlib import Path |
|
|
| from minijev.model import human_params |
|
|
|
|
| def measure_throughput(cfg: dict, seconds: float = 60.0) -> dict: |
| from minijev.training.supervised import train |
|
|
| cfg = copy.deepcopy(cfg) |
| data_dir = Path(cfg["data"]["dir"]) |
| tmp = Path(tempfile.mkdtemp(prefix="minijev-est-")) |
| if not (data_dir / "train.npz").exists(): |
| |
| gen = dict(cfg["data"].get("generate", {})) |
| gen["samples"] = 30_000 |
| cfg["data"] = {"dir": str(tmp / "data"), "generate": gen} |
| tcfg = cfg["training"] |
| tcfg["log_every"] = 10 |
| tcfg["eval_every"] = 10**9 |
| tcfg["checkpoint_every"] = 10**9 |
| s = train(cfg, run_dir=tmp / "run", max_minutes=seconds / 60, quiet=True, final_eval=False) |
| rows = list(csv.DictReader(open(tmp / "run" / "metrics.csv"))) |
| rates = [float(r["samples_per_sec"]) for r in rows if r.get("samples_per_sec")] |
| steady = rates[1:] or rates |
| sps = sorted(steady)[len(steady) // 2] if steady else s["throughput_samples_per_sec"] |
| return {"samples_per_sec": sps, "runtime": s["runtime"], "params": s["params"], "train_rows": s["train_rows"], |
| "windows": len(rates)} |
|
|
|
|
| def measure_generation(cfg: dict, n: int = 5000) -> float | None: |
| """Rows/s of oracle data generation on this machine (all cores).""" |
| gen = cfg["data"].get("generate") |
| if not gen: |
| return None |
| import importlib |
| mod = importlib.import_module(f"environments.{cfg['task']['name']}.dataset") |
| g = {k: v for k, v in gen.items() if k not in ("samples", "seed")} |
| t = time.time() |
| mod.generate(n, **g) |
| return n / (time.time() - t) |
|
|
|
|
| def estimate(cfg: dict, hourly_cost: float, throughput: float | None = None, bench_seconds: float = 60, |
| gpu_name: str | None = None, budget: float | None = None, overhead_minutes: float = 3.0) -> dict: |
| from minijev.device import resolve_runtime |
| from minijev.tasks import get_task |
| from minijev.training.supervised import make_model |
|
|
| task = get_task(cfg["task"]["name"], cfg["task"]) |
| params = make_model(cfg, task).num_parameters() |
| data_dir = Path(cfg["data"]["dir"]) |
| if (data_dir / "train.npz").exists(): |
| import numpy as np |
| rows = int(len(np.load(data_dir / "train.npz")["target"])) |
| gen_rows = 0 |
| else: |
| total = int(cfg["data"].get("generate", {}).get("samples", 0)) |
| rows, gen_rows = int(total * 0.9), total |
| if cfg["data"].get("train_limit"): |
| rows = min(rows, int(cfg["data"]["train_limit"])) |
| epochs = float(cfg["training"].get("epochs", 1)) |
| prof = resolve_runtime(cfg["training"], params) |
| measured = throughput is None |
| if measured: |
| print(f"[estimate] measuring real throughput for {bench_seconds:.0f}s on {prof.device} ({prof.device_name}) ...") |
| m = measure_throughput(cfg, bench_seconds) |
| throughput = m["samples_per_sec"] |
| gen_rate = measure_generation(cfg) if gen_rows else None |
| train_s = rows * epochs / throughput |
| evals = max(1, int(rows * epochs / cfg["training"].get("batch_size", 256) / cfg["training"].get("eval_every", 1000))) |
| eval_rows = min(cfg["training"].get("eval_max_rows", 20000), int(rows / 18)) |
| eval_s = evals * eval_rows / (3 * throughput) |
| gen_s = gen_rows / gen_rate if gen_rate else 0.0 |
| total_s = train_s + eval_s + gen_s + overhead_minutes * 60 |
| hours = total_s / 3600 |
| cost = hours * hourly_cost |
| out = {"model": cfg.get("name"), "params": params, "train_rows": rows, "epochs": epochs, |
| "device": gpu_name or prof.device_name, "throughput_samples_per_sec": round(throughput, 1), |
| "throughput_measured": measured, "train_minutes": round(train_s / 60, 1), |
| "data_generation_minutes": round(gen_s / 60, 1), "eval_minutes": round(eval_s / 60, 1), |
| "overhead_minutes": overhead_minutes, "total_minutes": round(total_s / 60, 1), |
| "gpu_hours": round(hours, 3), "hourly_cost": hourly_cost, "estimated_cost": round(cost, 3)} |
| print() |
| print(f"Model: MiniJEV {cfg.get('name')}") |
| print(f"Parameters: {human_params(params)}") |
| print(f"Dataset: {rows:,} training examples" + (f" (+ generating {gen_rows:,})" if gen_rows else "")) |
| print(f"Epochs: {epochs:g}") |
| print() |
| print(f"Device: {out['device']} (micro-batch {prof.micro_batch} x accum {prof.grad_accum}, amp={prof.amp_dtype})") |
| print(f"{'Measured' if measured else 'Given'} throughput: {throughput:,.0f} samples/sec") |
| print() |
| print(f"Estimated duration: {_fmt_min(total_s / 60)} (train {_fmt_min(train_s / 60)}, " |
| f"data gen {_fmt_min(gen_s / 60)}, eval {_fmt_min(eval_s / 60)}, overhead {overhead_minutes:g} min)") |
| print(f"Estimated GPU hours: {hours:.3f}") |
| print(f"Estimated cost: €{cost:.2f} (at €{hourly_cost:.3f}/h)") |
| if budget is not None and cost > budget: |
| scale = budget / cost |
| print(f"\n!! Over budget (€{budget:.2f}). Reduce e.g. epochs to {epochs * scale:.2f} " |
| f"or train rows to {int(rows * scale):,}.") |
| return out |
|
|
|
|
| def _fmt_min(m: float) -> str: |
| return f"{m:.1f} min" if m < 120 else f"{m / 60:.1f} h" |
|
|
