1"""P4 profiling: data-loader throughput, training step throughput and peak memory per model preset and batch size,2inference latency per origin. Writes results/metrics/benchmark_<node>.json. Uses real cached grids (private)."""3from __future__ import annotations45import argparse6import json7import platform8import subprocess9import sys10import time11from pathlib import Path1213import numpy as np14import torch15from torch.utils.data import DataLoader1617sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))18from chronos_x.config import load_config # noqa: E40219from chronos_x import pipeline as P # noqa: E40220from chronos_x.data.windows import WindowConfig, OriginIndex # noqa: E40221from chronos_x.models.model import MRPT, MRPTConfig # noqa: E40222from chronos_x.train.sampling import OriginDataset, collate # noqa: E40223from chronos_x.train.trainer import pick_device, to_device, mem_mb # noqa: E40224from chronos_x.train.losses import forecast_loss # noqa: E40225from chronos_x import experiment as E # noqa: E4022627ROOT = Path(__file__).resolve().parents[1]282930def main() -> None:31 import resource32 soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE); resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 65536), hard))33 torch.multiprocessing.set_sharing_strategy("file_system")34 ap = argparse.ArgumentParser()35 ap.add_argument("--config", default="configs/pilot.yaml")36 ap.add_argument("--symbols-from", default="pilot")37 ap.add_argument("--presets", default="small,spec,base")38 ap.add_argument("--batch-sizes", default="32,64,128")39 ap.add_argument("--steps", type=int, default=20)40 a = ap.parse_args()41 cfg = load_config(a.config)42 man = json.loads((ROOT / cfg["universe"]["manifest"]).read_text())43 rows = E.symbols_for(man, a.symbols_from)44 d = cfg["data"]45 dirs = P.grid_dirs(Path(cfg["paths"]["cache"]), d["start"], d["end"], [r["symbol"] for r in rows])46 wcfg = WindowConfig(**cfg["windows"])47 idx = OriginIndex(dirs, wcfg, None, stride=1, max_per_session=10, seed=0)48 ds = OriginDataset(idx)49 dev = pick_device()50 out = {"node": platform.node(), "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(),51 "torch": torch.__version__, "device": str(dev), "n_origins": len(idx), "loader": {}, "train": {}, "inference": {}}52 # data loader throughput (CPU workers)53 for nw in (0, 4, 8):54 loader = DataLoader(ds, batch_size=64, shuffle=True, num_workers=nw, collate_fn=collate, persistent_workers=nw > 0,55 prefetch_factor=4 if nw > 0 else None)56 it = iter(loader)57 for _ in range(4 if nw else 1): # warm-up: worker spawn, torch import, grid mmap58 next(it)59 t0 = time.time(); n = 060 for i in range(40):61 b = next(it); n += b["y"].shape[0]62 out["loader"][f"workers_{nw}"] = {"ex_per_s": n / (time.time() - t0)}63 print("loader", nw, out["loader"][f"workers_{nw}"], flush=True)64 # training steps65 for preset in a.presets.split(","):66 for bs in (int(x) for x in a.batch_sizes.split(",")):67 model = MRPT(MRPTConfig.preset(preset)).to(dev)68 opt = torch.optim.AdamW(model.parameters(), lr=1e-4)69 loader = DataLoader(ds, batch_size=bs, shuffle=True, num_workers=8, collate_fn=collate, drop_last=True)70 it = iter(loader); batches = [to_device(next(it), dev) for _ in range(6)]71 # warm-up72 for b in batches[:2]:73 loss, _ = forecast_loss(model(b), b, "student_t"); loss.backward(); opt.step(); opt.zero_grad()74 if dev.type == "mps":75 torch.mps.synchronize()76 t0 = time.time()77 for i in range(a.steps):78 b = batches[2 + i % 4]79 loss, _ = forecast_loss(model(b), b, "student_t"); loss.backward(); opt.step(); opt.zero_grad()80 if dev.type == "mps":81 torch.mps.synchronize()82 el = time.time() - t083 rec = {"params": model.n_params(), "steps_per_s": a.steps / el, "ex_per_s": a.steps * bs / el, "peak_mem_mb": mem_mb(dev),84 "loss_finite": bool(torch.isfinite(loss))}85 out["train"][f"{preset}_bs{bs}"] = rec86 print("train", preset, bs, rec, flush=True)87 # inference latency (batch 1 and batch 256)88 model.eval()89 with torch.no_grad():90 b1 = {k: v[:1] for k, v in batches[0].items()}91 for _ in range(3):92 model(b1)93 if dev.type == "mps":94 torch.mps.synchronize()95 t0 = time.time()96 for _ in range(20):97 model(b1)98 if dev.type == "mps":99 torch.mps.synchronize()100 lat1 = (time.time() - t0) / 20101 t0 = time.time()102 for b in batches[:4]:103 model(b)104 if dev.type == "mps":105 torch.mps.synchronize()106 latb = (time.time() - t0) / (4 * bs)107 out["inference"][f"{preset}_bs{bs}"] = {"latency_ms_batch1": lat1 * 1e3, "latency_ms_per_origin_batched": latb * 1e3}108 del model, opt109 if dev.type == "mps":110 torch.mps.empty_cache()111 (ROOT / "results" / "metrics" / f"benchmark_{platform.node().split('.')[0]}.json").write_text(json.dumps(out, indent=1))112 print(json.dumps(out["train"], indent=1))113114115if __name__ == "__main__":116 main()117