wehub-resource-sync
491 lines
19 KiB
Python
491 lines
19 KiB
Python
"""Pytest fixtures for ML4T test suite.
|
|
|
|
Two modes of operation:
|
|
1. CI (GHA): ML4T_DATA_PATH points to pre-subsampled real data (from private repo).
|
|
populated_data_dir just returns that path — no synthetic data needed.
|
|
2. Local dev: ML4T_DATA_PATH points to full production data or test data.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
REPO_ROOT = Path(__file__).parent.parent
|
|
|
|
# Case study IDs whose config/setup.yaml should be seeded into test output dirs
|
|
CASE_STUDY_IDS = [
|
|
"etfs",
|
|
"crypto_perps_funding",
|
|
"nasdaq100_microstructure",
|
|
"sp500_equity_option_analytics",
|
|
"us_firm_characteristics",
|
|
"fx_pairs",
|
|
"cme_futures",
|
|
"sp500_options",
|
|
"us_equities_panel",
|
|
]
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def ci_env_setup():
|
|
"""Create .env file if running in CI (where ML4T_DATA_PATH is set externally).
|
|
|
|
utils/config.py requires a .env file to exist.
|
|
In CI, environment variables are set by the workflow, but the .env
|
|
file still needs to exist to avoid FileNotFoundError on import.
|
|
"""
|
|
env_file = REPO_ROOT / ".env"
|
|
created = False
|
|
|
|
if not env_file.exists():
|
|
# Create minimal .env for CI
|
|
env_file.write_text(
|
|
f"ML4T_PATH={REPO_ROOT}\n"
|
|
f"ML4T_DATA_PATH={os.environ.get('ML4T_DATA_PATH', REPO_ROOT / 'data')}\n"
|
|
)
|
|
created = True
|
|
|
|
yield
|
|
|
|
# Clean up CI-created .env (don't leave artifacts)
|
|
if created and env_file.exists():
|
|
env_file.unlink()
|
|
|
|
|
|
def _resolve_data_path() -> Path | None:
|
|
"""Find ML4T_DATA_PATH from env var, .env file, or default location.
|
|
|
|
pytest-xdist workers may not inherit env vars set by the parent process,
|
|
so we also check the .env file and well-known test-data locations.
|
|
"""
|
|
# 1. Explicit env var (works in single-process pytest and CI)
|
|
env_path = os.environ.get("ML4T_DATA_PATH")
|
|
if env_path:
|
|
p = Path(env_path).expanduser().resolve()
|
|
if p.exists() and any(p.iterdir()):
|
|
return p
|
|
|
|
# 2. Read from .env file (works in xdist workers)
|
|
env_file = REPO_ROOT / ".env"
|
|
if env_file.exists():
|
|
for line in env_file.read_text().splitlines():
|
|
line = line.strip()
|
|
if line.startswith("ML4T_DATA_PATH") and "=" in line:
|
|
val = line.split("=", 1)[1].strip().strip('"').strip("'")
|
|
if val and not val.startswith("#"):
|
|
p = Path(val).expanduser().resolve()
|
|
if p.exists() and any(p.iterdir()):
|
|
return p
|
|
|
|
# 3. Well-known test-data repo location
|
|
test_data = Path.home() / "ml4t" / "test-data" / "data"
|
|
if test_data.exists() and (test_data / "etfs").exists():
|
|
return test_data
|
|
|
|
# 4. Default: repo's own data/ directory
|
|
repo_data = REPO_ROOT / "data"
|
|
if repo_data.exists() and any(repo_data.glob("*/*.parquet")):
|
|
return repo_data
|
|
|
|
return None
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_data_dir(tmp_path_factory):
|
|
"""Return the data directory for tests.
|
|
|
|
Resolves ML4T_DATA_PATH from env var, .env file, well-known test-data
|
|
repo location, or repo's data/ directory. Works with pytest-xdist.
|
|
"""
|
|
resolved = _resolve_data_path()
|
|
if resolved:
|
|
os.environ["ML4T_DATA_PATH"] = str(resolved)
|
|
return resolved
|
|
|
|
# Fallback: create temp directory
|
|
data_dir = tmp_path_factory.mktemp("ml4t_data")
|
|
os.environ["ML4T_DATA_PATH"] = str(data_dir)
|
|
return data_dir
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def populated_data_dir(test_data_dir):
|
|
"""Return a data directory populated with test data.
|
|
|
|
If ML4T_DATA_PATH points to pre-populated data (e.g., from GHA checkout
|
|
of ml4t/third-edition-test-data), returns it directly.
|
|
"""
|
|
if (test_data_dir / "etfs" / "market" / "etf_universe.parquet").exists():
|
|
return test_data_dir
|
|
|
|
pytest.skip("No test data available. Set ML4T_DATA_PATH or run in CI.")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def intermediates_dir(test_data_dir):
|
|
"""Return directory with pre-computed pipeline intermediates.
|
|
|
|
When running downstream chapters (Ch11+), they need labels/features
|
|
from pipeline stages. These are pre-computed and stored in test-data repo.
|
|
"""
|
|
idir = test_data_dir.parent / "intermediates"
|
|
if idir.exists() and any(idir.iterdir()):
|
|
return idir
|
|
return None
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def seeded_output_dir(tmp_path_factory):
|
|
"""Session-scoped output dir seeded with case study config files.
|
|
|
|
Chapter notebooks that read case study setup.yaml (via get_case_study_dir())
|
|
need these configs to exist even when ML4T_OUTPUT_DIR redirects writes to
|
|
a temp directory. This fixture copies the real config files into the test
|
|
output dir so notebooks can find them.
|
|
|
|
With pytest-xdist, each worker gets its own subdirectory to avoid races
|
|
on shutil.rmtree/copytree when multiple workers seed simultaneously.
|
|
"""
|
|
base_dir = os.environ.get("ML4T_OUTPUT_DIR")
|
|
if base_dir:
|
|
# With xdist, append worker id to avoid races
|
|
worker_id = os.environ.get("PYTEST_XDIST_WORKER", "")
|
|
if worker_id:
|
|
output_dir = Path(base_dir) / f"worker_{worker_id}"
|
|
else:
|
|
output_dir = Path(base_dir)
|
|
else:
|
|
output_dir = tmp_path_factory.mktemp("ml4t_output")
|
|
|
|
# Set the env var so notebooks see this worker's output dir
|
|
os.environ["ML4T_OUTPUT_DIR"] = str(output_dir)
|
|
|
|
cs_root = REPO_ROOT / "case_studies"
|
|
|
|
# Copy per-case-study config files (setup.yaml, training menus, backtest presets, etc.)
|
|
for cs_id in CASE_STUDY_IDS:
|
|
src_config_dir = cs_root / cs_id / "config"
|
|
if not src_config_dir.exists():
|
|
continue
|
|
dst_config_dir = output_dir / cs_id / "config"
|
|
if dst_config_dir.exists():
|
|
shutil.rmtree(dst_config_dir)
|
|
shutil.copytree(src_config_dir, dst_config_dir)
|
|
_trim_label_configs(dst_config_dir)
|
|
|
|
# Copy global model presets (case_studies/config/) so load_configs() can find them.
|
|
# load_configs() resolves presets at {case_dir.parent}/config/{model_type}/*.yaml
|
|
# We copy (not symlink) so we can patch presets for fast testing.
|
|
global_config_src = cs_root / "config"
|
|
global_config_dst = output_dir / "config"
|
|
if global_config_src.exists():
|
|
if global_config_dst.exists():
|
|
shutil.rmtree(global_config_dst)
|
|
shutil.copytree(global_config_src, global_config_dst)
|
|
_patch_presets_for_testing(global_config_dst)
|
|
|
|
# Copy pipeline intermediates (features, labels, run_log) from test-data repo.
|
|
# These are pre-computed so downstream notebooks (Ch11+) can run without
|
|
# executing the full pipeline first.
|
|
# Look for intermediates next to data (test-data repo layout) or at well-known path.
|
|
data_path = _resolve_data_path()
|
|
intermediates_root = None
|
|
if data_path:
|
|
candidate = Path(data_path).parent / "intermediates"
|
|
if candidate.exists():
|
|
intermediates_root = candidate
|
|
if intermediates_root is None:
|
|
# Well-known test-data repo location
|
|
candidate = Path.home() / "ml4t" / "test-data" / "intermediates"
|
|
if candidate.exists():
|
|
intermediates_root = candidate
|
|
if intermediates_root and intermediates_root.exists():
|
|
for cs_id in CASE_STUDY_IDS:
|
|
src = intermediates_root / cs_id
|
|
if not src.exists():
|
|
continue
|
|
dst = output_dir / cs_id
|
|
# Copy features, labels, evaluation, run_log, results, benchmark —
|
|
# anything that downstream notebooks look for in get_case_study_dir()
|
|
for subdir in ["features", "labels", "evaluation", "run_log", "results", "benchmark"]:
|
|
src_sub = src / subdir
|
|
dst_sub = dst / subdir
|
|
if src_sub.exists() and not dst_sub.exists():
|
|
shutil.copytree(src_sub, dst_sub)
|
|
# Copy top-level intermediate files (e.g. etfs/eligibility.csv,
|
|
# protocol.yaml, baseline_checkpoint.yaml) that sit directly in
|
|
# intermediates/{cs_id}/ rather than in a subdir. Downstream
|
|
# notebooks (etfs 02_labels, 03_financial_features) read these via
|
|
# get_case_study_dir(); without this they fail with FileNotFoundError.
|
|
for item in src.iterdir():
|
|
if item.is_file():
|
|
dst_file = dst / item.name
|
|
if not dst_file.exists():
|
|
dst.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(item, dst_file)
|
|
# Schema reconciliation: test-data predictions parquets were
|
|
# generated with an older column convention (y_score / y_true /
|
|
# fold_id). Production registry uses (prediction / actual / fold).
|
|
# Rename in place so notebooks reading via get_case_study_dir()
|
|
# see the canonical names without per-notebook compat shims.
|
|
preds_root = dst / "run_log" / "predictions"
|
|
if preds_root.exists():
|
|
_migrate_predictions_schema(preds_root)
|
|
|
|
# Copy non-case-study intermediates (chapter-scoped fixtures).
|
|
# These are intermediates that downstream teaching notebooks need but aren't
|
|
# part of the per-case-study pipeline (e.g., Ch16 signal comparison, Ch20 synthesis).
|
|
if intermediates_root and intermediates_root.exists():
|
|
for extra_id in ["ch16_signal_method_comparison", "ch20_synthesis"]:
|
|
src = intermediates_root / extra_id
|
|
if not src.exists():
|
|
continue
|
|
dst = output_dir / extra_id
|
|
if not dst.exists():
|
|
shutil.copytree(src, dst)
|
|
|
|
# Seed minimal results fixtures so downstream notebooks (latent factors, DL,
|
|
# backtest) can find baseline results without depending on upstream execution.
|
|
# These fill gaps where intermediates don't provide enough (e.g., Ch25 demo
|
|
# predictions, Ch15 causal JSON, synthetic registry entries).
|
|
from tests.fixtures.seed_results import seed_results
|
|
|
|
seed_results(output_dir, CASE_STUDY_IDS)
|
|
|
|
# Symlink AQR factor data so AQRFactorProvider finds it at ~/ml4t/data/aqr_factors
|
|
aqr_src = data_path.parent / "data" / "factors" / "aqr" if data_path else None
|
|
if aqr_src is None:
|
|
aqr_src = Path.home() / "ml4t" / "test-data" / "data" / "factors" / "aqr"
|
|
aqr_dst = Path.home() / "ml4t" / "data" / "aqr_factors"
|
|
if aqr_src.exists() and not aqr_dst.exists():
|
|
aqr_dst.parent.mkdir(parents=True, exist_ok=True)
|
|
aqr_dst.symlink_to(aqr_src)
|
|
|
|
return output_dir
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Preset patching — reduce workload for CI/test runs
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Per-model-type overrides applied to copied preset YAMLs.
|
|
# Goal: minimal workload that still exercises the training loop + registry.
|
|
_TEST_PRESET_PATCHES: dict[str, dict] = {
|
|
"lgb": {"max_iterations": 2, "checkpoint_interval": 1},
|
|
# DL families: 2 epochs, checkpoint every epoch
|
|
"lstm": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
"tsmixer": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
"tcn": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
"nlinear": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
"patchtst": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
# TabDL: 2 epochs
|
|
"tabm": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
# Latent factors: 2 epochs
|
|
"cae": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
"sdf": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
"sae": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
"ipca": {"n_epochs": 2, "checkpoint_interval": 1},
|
|
}
|
|
|
|
|
|
_PREDICTION_COL_RENAMES = {
|
|
"y_score": "prediction",
|
|
"y_true": "actual",
|
|
"fold_id": "fold",
|
|
}
|
|
|
|
|
|
def _migrate_predictions_schema(preds_root: Path) -> None:
|
|
"""Rename test-data prediction columns to canonical production schema.
|
|
|
|
Test-data parquets were generated with an older convention
|
|
(y_score / y_true / fold_id). Production registry uses
|
|
(prediction / actual / fold). Walking the seeded predictions tree once
|
|
avoids per-notebook compat shims while keeping test-data immutable on
|
|
its own repo schedule.
|
|
"""
|
|
import polars as pl
|
|
|
|
for parquet in preds_root.rglob("predictions.parquet"):
|
|
cols = pl.read_parquet(parquet, n_rows=0).columns
|
|
renames = {old: new for old, new in _PREDICTION_COL_RENAMES.items() if old in cols}
|
|
if not renames:
|
|
continue
|
|
df = pl.read_parquet(parquet).rename(renames)
|
|
df.write_parquet(parquet)
|
|
|
|
|
|
def _patch_presets_for_testing(config_dir: Path) -> None:
|
|
"""Patch copied preset YAMLs with reduced-workload values for testing."""
|
|
for model_type, overrides in _TEST_PRESET_PATCHES.items():
|
|
model_dir = config_dir / model_type
|
|
if not model_dir.exists():
|
|
continue
|
|
for preset_path in model_dir.glob("*.yaml"):
|
|
preset = yaml.safe_load(preset_path.read_text())
|
|
if preset is None:
|
|
continue
|
|
preset.update(overrides)
|
|
with open(preset_path, "w") as f:
|
|
yaml.dump(preset, f, default_flow_style=False)
|
|
|
|
|
|
# Max configs per family in label config files (keep tests fast but comprehensive).
|
|
# Only applied to families with homogeneous sweep configs (linear, gbm).
|
|
# DL/TabDL/latent/causal families are NOT trimmed because each config often
|
|
# maps to a dedicated notebook (e.g., 09_dl_lstm, 10_dl_tsmixer).
|
|
_MAX_CONFIGS_PER_FAMILY = 2
|
|
_TRIM_FAMILIES = {"linear", "gbm"}
|
|
|
|
|
|
def _trim_label_configs(cs_config_dir: Path) -> None:
|
|
"""Trim training menu YAMLs to at most _MAX_CONFIGS_PER_FAMILY for sweep families."""
|
|
training_dir = cs_config_dir / "training"
|
|
label_root = training_dir if training_dir.exists() else cs_config_dir
|
|
for label_yaml in label_root.glob("fwd_*.yaml"):
|
|
data = yaml.safe_load(label_yaml.read_text())
|
|
if data is None or not isinstance(data, dict):
|
|
continue
|
|
trimmed = False
|
|
for family, configs in data.items():
|
|
if (
|
|
family in _TRIM_FAMILIES
|
|
and isinstance(configs, list)
|
|
and len(configs) > _MAX_CONFIGS_PER_FAMILY
|
|
):
|
|
data[family] = configs[:_MAX_CONFIGS_PER_FAMILY]
|
|
trimmed = True
|
|
if trimmed:
|
|
with open(label_yaml, "w") as f:
|
|
yaml.dump(data, f, default_flow_style=False)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GPU marker — apply `@pytest.mark.gpu` at collection time based on overrides.
|
|
# Usage: pytest -m gpu (GPU only) | pytest -m "not gpu" (CPU only)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def pytest_configure(config):
|
|
config.addinivalue_line("markers", "gpu: notebook requires GPU (from overrides.yaml gpu: true)")
|
|
config.addinivalue_line(
|
|
"markers",
|
|
"long_running: notebook takes >10min even with reduced params (from overrides.yaml long_running: true)",
|
|
)
|
|
config.addinivalue_line(
|
|
"markers",
|
|
"weekly: notebook tier=weekly — runs only in scheduled weekly-external workflow. "
|
|
"To execute locally, set ML4T_TEST_TIER=weekly alongside `pytest -m weekly`; "
|
|
"without the env var, matching items are collected but skipped.",
|
|
)
|
|
config.addinivalue_line(
|
|
"markers",
|
|
"on_demand: notebook tier=on_demand — runs only on manual dispatch (e.g., GPU Tier 3). "
|
|
"To execute locally, set ML4T_TEST_TIER=on_demand alongside `pytest -m on_demand`; "
|
|
"without the env var, matching items are collected but skipped.",
|
|
)
|
|
config.addinivalue_line(
|
|
"markers",
|
|
"drift: live external drift smoke (Phase 3) — one tiny real pull per free "
|
|
"data source. Network-bound; runs only in the scheduled weekly-external "
|
|
"workflow's `drift` job, never per-PR.",
|
|
)
|
|
|
|
|
|
def pytest_collection_modifyitems(items):
|
|
"""Add markers to test items based on overrides.yaml flags."""
|
|
from tests.pm_helpers import (
|
|
TIER_ON_DEMAND,
|
|
TIER_WEEKLY,
|
|
get_overrides,
|
|
get_reruns,
|
|
get_tier,
|
|
)
|
|
|
|
# pytest-rerunfailures provides @pytest.mark.flaky(reruns=N). Detect once
|
|
# so per-NB reruns kick in automatically when the dep lands in Step 2.
|
|
try:
|
|
import pytest_rerunfailures # noqa: F401
|
|
|
|
has_rerunfailures = True
|
|
except ImportError:
|
|
has_rerunfailures = False
|
|
|
|
for item in items:
|
|
if hasattr(item, "callspec") and "notebook_path" in item.callspec.params:
|
|
nb_path = item.callspec.params["notebook_path"]
|
|
rel = (
|
|
nb_path.relative_to(REPO_ROOT).with_suffix("")
|
|
if hasattr(nb_path, "relative_to")
|
|
else nb_path
|
|
)
|
|
overrides = get_overrides(str(rel)) or {}
|
|
if overrides.get("gpu"):
|
|
item.add_marker(pytest.mark.gpu)
|
|
if overrides.get("long_running"):
|
|
item.add_marker(pytest.mark.long_running)
|
|
|
|
tier = get_tier(overrides)
|
|
if tier == TIER_WEEKLY:
|
|
item.add_marker(pytest.mark.weekly)
|
|
elif tier == TIER_ON_DEMAND:
|
|
item.add_marker(pytest.mark.on_demand)
|
|
|
|
reruns = get_reruns(overrides)
|
|
if reruns > 0 and has_rerunfailures:
|
|
item.add_marker(pytest.mark.flaky(reruns=reruns, reruns_delay=30))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Incremental result saving — write JSONL after each test so results survive
|
|
# process kills. Results file: /tmp/ml4t-test-results.jsonl
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_RESULTS_PATH = Path(os.environ.get("ML4T_RESULTS_FILE", "/tmp/ml4t-test-results.jsonl"))
|
|
_test_start_times: dict[str, float] = {}
|
|
|
|
|
|
@pytest.hookimpl(tryfirst=True)
|
|
def pytest_runtest_setup(item):
|
|
"""Record test start time."""
|
|
_test_start_times[item.nodeid] = time.time()
|
|
|
|
|
|
@pytest.hookimpl(trylast=True)
|
|
def pytest_runtest_logreport(report):
|
|
"""Write each test result to JSONL as it completes."""
|
|
if report.when != "call" and not (report.when == "setup" and report.skipped):
|
|
return
|
|
|
|
start = _test_start_times.pop(report.nodeid, 0)
|
|
duration = report.duration if hasattr(report, "duration") else 0
|
|
|
|
outcome = report.outcome # "passed", "failed", or "skipped"
|
|
|
|
record = {
|
|
"nodeid": report.nodeid,
|
|
"outcome": outcome,
|
|
"duration_s": round(duration, 2),
|
|
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
|
|
}
|
|
if outcome == "failed" and report.longreprtext:
|
|
record["error"] = report.longreprtext[:500]
|
|
|
|
with open(_RESULTS_PATH, "a") as f:
|
|
f.write(json.dumps(record) + "\n")
|
|
f.flush()
|
|
|
|
|
|
@pytest.fixture
|
|
def clean_env():
|
|
"""Fixture that provides a clean environment and restores it after."""
|
|
saved_env = os.environ.copy()
|
|
yield os.environ
|
|
os.environ.clear()
|
|
os.environ.update(saved_env)
|