From e739f84afdd3b49858f2c2b79a6a1e3926136cf7 Mon Sep 17 00:00:00 2001 From: Mathias Date: Fri, 26 Jun 2026 12:37:27 +0200 Subject: [PATCH] feat(hpo): env-var knob overrides + sweep script (18 configs) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - train.py knobs all readable from JEPA_* env vars (JEPA_WINDOW, JEPA_D_MODEL, JEPA_DEPTH, etc.) so hpo_sweep.py can override without touching source - scripts/hpo_sweep.py: 3×2×3 grid over D_MODEL × DEPTH × WINDOW, logs to results/hpo/hpo_results.jsonl with leaderboard at end - 3 new tests: env override correctness, configs() schema validation - 19/19 tests pass Co-Authored-By: Claude Sonnet 4.6 --- scripts/hpo_sweep.py | 97 ++++++++++++++++++++++++++++++++++++++++++++ tests/test_hepa.py | 54 ++++++++++++++++++++++-- train.py | 31 +++++++------- 3 files changed, 163 insertions(+), 19 deletions(-) create mode 100644 scripts/hpo_sweep.py diff --git a/scripts/hpo_sweep.py b/scripts/hpo_sweep.py new file mode 100644 index 0000000..ba3909f --- /dev/null +++ b/scripts/hpo_sweep.py @@ -0,0 +1,97 @@ +"""HPO sweep for jepa-fx-risk HEPA backbone. + +Runs train.py with different JEPA_* env overrides, logs results to +results/hpo/hpo_results.jsonl. Each config writes its metrics.json then +the result is appended to the JSONL. + +Usage: + python scripts/hpo_sweep.py + python scripts/hpo_sweep.py --dry-run # print configs, don't train +""" +import argparse +import json +import os +import subprocess +import sys +from datetime import datetime +from itertools import product +from pathlib import Path + +# ── Search space ────────────────────────────────────────────────────────────── +SEARCH_SPACE = { + "JEPA_D_MODEL": [64, 128, 256], + "JEPA_DEPTH": [2, 4], + "JEPA_WINDOW": [120, 240, 480], +} +# Fixed: PATCH_LEN=24 (1-day patches), N_HEADS=4, EPOCHS=300, PHASE1_EPOCHS=200 + +PYTHON = str(Path(sys.executable)) +OUT_DIR = Path("results/hpo") + + +def configs(): + """Yield all configs as dicts of JEPA_* env overrides.""" + keys = list(SEARCH_SPACE.keys()) + for vals in product(*SEARCH_SPACE.values()): + yield dict(zip(keys, vals)) + + +def run_config(cfg: dict, metrics_path: str = "metrics.json") -> dict: + env = {**os.environ, **{k: str(v) for k, v in cfg.items()}} + result = subprocess.run( + [PYTHON, "train.py"], + env=env, + capture_output=True, + text=True, + ) + if result.returncode != 0: + return {"config": cfg, "error": result.stderr[-500:]} + stdout_last = result.stdout.strip().split("\n")[-1] + with open(metrics_path) as f: + m = json.load(f) + return { + "config": cfg, + "val_vol_r2": m.get("val_vol_r2"), + "phase1_r2": m.get("phase1_r2"), + "stdout_last": stdout_last, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + OUT_DIR.mkdir(parents=True, exist_ok=True) + out_file = OUT_DIR / "hpo_results.jsonl" + + all_cfgs = list(configs()) + print(f"HPO sweep: {len(all_cfgs)} configs") + for i, cfg in enumerate(all_cfgs): + label = " ".join(f"{k.replace('JEPA_','')}={v}" for k, v in cfg.items()) + print(f"\n[{i+1}/{len(all_cfgs)}] {label}") + if args.dry_run: + continue + ts = datetime.utcnow().isoformat() + row = run_config(cfg) + row["ts"] = ts + with open(out_file, "a") as f: + f.write(json.dumps(row) + "\n") + if "error" in row: + print(f" ERROR: {row['error'][:200]}") + else: + print(f" val_vol_r2={row['val_vol_r2']:.4f} phase1_r2={row['phase1_r2']:.4f}") + + if not args.dry_run: + # Print leaderboard + rows = [json.loads(l) for l in open(out_file) if l.strip()] + rows = [r for r in rows if "error" not in r] + rows.sort(key=lambda r: r.get("phase1_r2", -999), reverse=True) + print("\n── Leaderboard (by phase1_r2) ─────────────────────────") + for r in rows[:5]: + cfg_str = " ".join(f"{k.replace('JEPA_','')}={v}" for k,v in r["config"].items()) + print(f" {r['phase1_r2']:.4f} {cfg_str}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_hepa.py b/tests/test_hepa.py index dc4f4f4..a155951 100644 --- a/tests/test_hepa.py +++ b/tests/test_hepa.py @@ -1,9 +1,10 @@ -"""Failing tests for HEPA backbone + Phase-1 supervised head in train.py. +"""Failing tests for HEPA backbone + Phase-1 supervised head + HPO in train.py. Run: cd ~/dev/AI/jepa-fx-risk && .venv/bin/python -m pytest tests/test_hepa.py -v These tests define what the backbone and head must satisfy BEFORE implementation. """ import math +import os import torch import torch.nn as nn import pytest @@ -12,11 +13,24 @@ import pytest # They will fail until train.py implements: CausalEncoder, HorizonPredictor, vicreg_loss -def _import(): +def _import(env_overrides=None): import importlib.util, sys - spec = importlib.util.spec_from_file_location("train", "train.py") + saved = {} + if env_overrides: + for k, v in env_overrides.items(): + saved[k] = os.environ.get(k) + os.environ[k] = str(v) + # Force fresh module load (env vars must be read at import time) + name = f"train_{id(env_overrides)}" + spec = importlib.util.spec_from_file_location(name, "train.py") mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) + if env_overrides: + for k, orig in saved.items(): + if orig is None: + os.environ.pop(k, None) + else: + os.environ[k] = orig return mod @@ -171,7 +185,7 @@ def test_phase1_beats_linear_on_nonlinear(train_mod): # 11. main() returns phase1_r2 in metrics.json (integration — needs real data) def test_metrics_json_has_phase1_r2(train_mod): - import os, json + import json if not os.path.exists("metrics.json"): pytest.skip("metrics.json not present — run train.py first") with open("metrics.json") as f: @@ -181,3 +195,35 @@ def test_metrics_json_has_phase1_r2(train_mod): f"MLP head phase1_r2={m['phase1_r2']:.4f} should beat linear probe " f"val_vol_r2={m['val_vol_r2']:.4f}" ) + + +# ── HPO: env-var knob overrides ─────────────────────────────────────────────── + +# 12. JEPA_WINDOW env var overrides WINDOW at import time +def test_env_override_window(): + mod = _import({"JEPA_WINDOW": "48"}) + assert mod.WINDOW == 48, f"expected WINDOW=48, got {mod.WINDOW}" + + +# 13. JEPA_D_MODEL and JEPA_DEPTH env vars work +def test_env_override_d_model_depth(): + mod = _import({"JEPA_D_MODEL": "64", "JEPA_DEPTH": "4"}) + assert mod.D_MODEL == 64, f"expected D_MODEL=64, got {mod.D_MODEL}" + assert mod.DEPTH == 4, f"expected DEPTH=4, got {mod.DEPTH}" + + +# 14. hpo_sweep.py exists and generates correct config list +def test_hpo_sweep_configs(): + import importlib.util + sweep_path = "scripts/hpo_sweep.py" + if not os.path.exists(sweep_path): + pytest.fail(f"{sweep_path} not found — implement it") + spec = importlib.util.spec_from_file_location("hpo_sweep", sweep_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + cfgs = list(mod.configs()) + assert len(cfgs) > 0, "configs() returned empty list" + # Every config must have at least D_MODEL, DEPTH, WINDOW keys + required = {"JEPA_D_MODEL", "JEPA_DEPTH", "JEPA_WINDOW"} + for cfg in cfgs: + assert required.issubset(cfg.keys()), f"config missing required keys: {cfg}" diff --git a/train.py b/train.py index 2c8fe09..a261e50 100644 --- a/train.py +++ b/train.py @@ -17,21 +17,22 @@ import torch import torch.nn as nn import torch.nn.functional as F -# --- agent-tunable knobs --- -USE_HOURLY = True # prefer eurusd_hourly.parquet when available -WINDOW = 240 # hourly: 10 trading days; if USE_HOURLY=False reset to 60 -PATCH_LEN = 24 # hourly: 1-day patches (10 tokens); if USE_HOURLY=False reset to 10 -D_MODEL = 128 -DEPTH = 2 -N_HEADS = 4 -ALPHA = 0.1 # VICReg mixing weight (fixed at 0.1 in HEPA paper) -DELTA_T_MAX = 3 # max prediction horizon in patches (1..min(DELTA_T_MAX, N-1-c)) -BATCH_SIZE = 512 # mini-batch per step (hourly dataset is too large for full-batch) -EPOCHS = 300 -LR = 3e-4 -PHASE1_EPOCHS = 200 # supervised head epochs (encoder frozen) -PHASE1_LR = 1e-3 -SEED = 0 +# --- agent-tunable knobs (all overridable via JEPA_* env vars for HPO) --- +import os as _os +USE_HOURLY = True +WINDOW = int(_os.environ.get("JEPA_WINDOW", 240)) +PATCH_LEN = int(_os.environ.get("JEPA_PATCH_LEN", 24)) +D_MODEL = int(_os.environ.get("JEPA_D_MODEL", 128)) +DEPTH = int(_os.environ.get("JEPA_DEPTH", 2)) +N_HEADS = int(_os.environ.get("JEPA_N_HEADS", 4)) +ALPHA = float(_os.environ.get("JEPA_ALPHA", 0.1)) +DELTA_T_MAX = int(_os.environ.get("JEPA_DELTA_T_MAX", 3)) +BATCH_SIZE = int(_os.environ.get("JEPA_BATCH_SIZE", 512)) +EPOCHS = int(_os.environ.get("JEPA_EPOCHS", 300)) +LR = float(_os.environ.get("JEPA_LR", 3e-4)) +PHASE1_EPOCHS = int(_os.environ.get("JEPA_PHASE1_EPOCHS", 200)) +PHASE1_LR = float(_os.environ.get("JEPA_PHASE1_LR", 1e-3)) +SEED = int(_os.environ.get("JEPA_SEED", 0)) # --------------------------- torch.manual_seed(SEED)