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 lightgbm # noqa: F401 (before torch: libomp segfault on macOS)14import numpy as np15import torch16from torch.utils.data import DataLoader1718sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))19from chronos_x.config import load_config # noqa: E40220from chronos_x import pipeline as P # noqa: E40221from chronos_x.data.windows import WindowConfig, OriginIndex # noqa: E40222from chronos_x.models.model import MRPT, MRPTConfig # noqa: E40223from chronos_x.train.sampling import OriginDataset, collate # noqa: E40224from chronos_x.train.trainer import pick_device, to_device, mem_mb # noqa: E40225from chronos_x.train.losses import forecast_loss # noqa: E40226from chronos_x import experiment as E # noqa: E4022728ROOT = Path(__file__).resolve().parents[1]293031def main() -> None:32 import resource33 soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE); resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 65536), hard))34 torch.multiprocessing.set_sharing_strategy("file_system")35 ap = argparse.ArgumentParser()36 ap.add_argument("--config", default="configs/pilot.yaml")37 ap.add_argument("--symbols-from", default="pilot")38 ap.add_argument("--presets", default="small,spec,base")39 ap.add_argument("--batch-sizes", default="32,64,128")40 ap.add_argument("--steps", type=int, default=20)41 a = ap.parse_args()42 cfg = load_config(a.config)43 man = json.loads((ROOT / cfg["universe"]["manifest"]).read_text())44 rows = E.symbols_for(man, a.symbols_from)45 d = cfg["data"]46 dirs = P.grid_dirs(Path(cfg["paths"]["cache"]), d["start"], d["end"], [r["symbol"] for r in rows])47 wcfg = WindowConfig(**cfg["windows"])48 idx = OriginIndex(dirs, wcfg, None, stride=1, max_per_session=10, seed=0)49 ds = OriginDataset(idx)50 dev = pick_device()51 out = {"node": platform.node(), "chip": subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip(),52 "torch": torch.__version__, "device": str(dev), "n_origins": len(idx), "loader": {}, "train": {}, "inference": {}}53 # data loader throughput (CPU workers)54 for nw in (0, 4, 8):55 loader = DataLoader(ds, batch_size=64, shuffle=True, num_workers=nw, collate_fn=collate)56 t0 = time.time(); n = 057 for i, b in enumerate(loader):58 n += b["y"].shape[0]59 if i >= 30:60 break61 out["loader"][f"workers_{nw}"] = {"ex_per_s": n / (time.time() - t0)}62 print("loader", nw, out["loader"][f"workers_{nw}"], flush=True)63 # training steps64 for preset in a.presets.split(","):65 for bs in (int(x) for x in a.batch_sizes.split(",")):66 model = MRPT(MRPTConfig.preset(preset)).to(dev)67 opt = torch.optim.AdamW(model.parameters(), lr=1e-4)68 loader = DataLoader(ds, batch_size=bs, shuffle=True, num_workers=8, collate_fn=collate, drop_last=True)69 it = iter(loader); batches = [to_device(next(it), dev) for _ in range(6)]70 # warm-up71 for b in batches[:2]:72 loss, _ = forecast_loss(model(b), b, "student_t"); loss.backward(); opt.step(); opt.zero_grad()73 if dev.type == "mps":74 torch.mps.synchronize()75 t0 = time.time()76 for i in range(a.steps):77 b = batches[2 + i % 4]78 loss, _ = forecast_loss(model(b), b, "student_t"); loss.backward(); opt.step(); opt.zero_grad()79 if dev.type == "mps":80 torch.mps.synchronize()81 el = time.time() - t082 rec = {"params": model.n_params(), "steps_per_s": a.steps / el, "ex_per_s": a.steps * bs / el, "peak_mem_mb": mem_mb(dev),83 "loss_finite": bool(torch.isfinite(loss))}84 out["train"][f"{preset}_bs{bs}"] = rec85 print("train", preset, bs, rec, flush=True)86 # inference latency (batch 1 and batch 256)87 model.eval()88 with torch.no_grad():89 b1 = {k: v[:1] for k, v in batches[0].items()}90 for _ in range(3):91 model(b1)92 if dev.type == "mps":93 torch.mps.synchronize()94 t0 = time.time()95 for _ in range(20):96 model(b1)97 if dev.type == "mps":98 torch.mps.synchronize()99 lat1 = (time.time() - t0) / 20100 t0 = time.time()101 for b in batches[:4]:102 model(b)103 if dev.type == "mps":104 torch.mps.synchronize()105 latb = (time.time() - t0) / (4 * bs)106 out["inference"][f"{preset}_bs{bs}"] = {"latency_ms_batch1": lat1 * 1e3, "latency_ms_per_origin_batched": latb * 1e3}107 del model, opt108 if dev.type == "mps":109 torch.mps.empty_cache()110 (ROOT / "results" / "metrics" / f"benchmark_{platform.node().split('.')[0]}.json").write_text(json.dumps(out, indent=1))111 print(json.dumps(out["train"], indent=1))112113114if __name__ == "__main__":115 main()116