Speed up snapshot refresh: batch prefetch, unbuffered output, incremental-first

- enrich_tier2: prefetch price bars (parallel ThreadPool) and short ratio
  (single batch DB query) instead of per-row HTTP/DB calls (~20min → ~2min)
- canonical_snapshots: add PYTHONUNBUFFERED=1 to enrichment subprocesses
  so progress output is visible in real time
- backtest_sim: use incremental_update_canonical_snapshot when existing
  snapshot is present, falling back to full rebuild only when needed

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
main
I Luk Kim 4 months ago
parent 5cb2b9fcb8
commit 969dedc635

@ -235,6 +235,20 @@ def _convert_from_runner(
except Exception: except Exception:
pass pass
# Snapshot coverage info from the canonical snapshot manifest
snapshot_coverage_end_date: str | None = None
snapshot_last_refresh_utc: str | None = None
try:
import json as _json
from pathlib import Path as _Path
_snap_manifest = _Path("data/parquet") / (config.canonical_snapshot_id or config.dataset_snapshot_id) / "manifest.json"
if _snap_manifest.exists():
_snap_meta = _json.loads(_snap_manifest.read_text())
snapshot_coverage_end_date = _snap_meta.get("coverage_end_date")
snapshot_last_refresh_utc = _snap_meta.get("last_refresh_utc")
except Exception:
pass
return { return {
"session_name": session_name, "session_name": session_name,
"config_path": config_path, "config_path": config_path,
@ -252,6 +266,8 @@ def _convert_from_runner(
"sharpe": sharpe, "sharpe": sharpe,
}, },
"metrics_bundle": metrics_bundle, "metrics_bundle": metrics_bundle,
"snapshot_coverage_end_date": snapshot_coverage_end_date,
"snapshot_last_refresh_utc": snapshot_last_refresh_utc,
} }
@ -414,13 +430,32 @@ async def _refresh_snapshot(
if console: if console:
console.print(f" [yellow]Label generator skipped: {exc}[/]") console.print(f" [yellow]Label generator skipped: {exc}[/]")
# Step 2: Rebuild canonical snapshot # Step 2: Update canonical snapshot (incremental first, full rebuild as fallback)
if console: if console:
console.print(" [dim]4/4 Exporting snapshot...[/]") console.print(" [dim]4/4 Exporting snapshot...[/]")
try: try:
from libs.export.canonical_snapshots import build_canonical_snapshot from libs.export.canonical_snapshots import (
build_canonical_snapshot,
incremental_update_canonical_snapshot,
)
snapshot_path = _resolve_snapshot_path(resolution.canonical_snapshot_id)
use_incremental = snapshot_path is not None and snapshot_path.exists()
if use_incremental:
if console:
console.print(" [dim]Incremental update (new events only)...[/]")
try:
await incremental_update_canonical_snapshot(snapshot_id)
except Exception as inc_exc:
if console:
console.print(f" [yellow]Incremental failed ({inc_exc}), falling back to full rebuild...[/]")
await build_canonical_snapshot(snapshot_id, manual=manual) await build_canonical_snapshot(snapshot_id, manual=manual)
else:
if console:
console.print(" [dim]Full rebuild (no existing snapshot)...[/]")
await build_canonical_snapshot(snapshot_id, manual=manual)
if console: if console:
console.print(" [green]Snapshot refreshed.[/]") console.print(" [green]Snapshot refreshed.[/]")
# Write marker to avoid re-refreshing today # Write marker to avoid re-refreshing today

@ -107,7 +107,8 @@ def _run_enrichment_step(step_name: str, input_dir: Path, output_dir: Path) -> N
str(output_dir), str(output_dir),
] ]
logger.info("canonical_snapshot_enrichment_start", step=step_name, cmd=cmd) logger.info("canonical_snapshot_enrichment_start", step=step_name, cmd=cmd)
subprocess.run(cmd, check=True) env = {**__import__("os").environ, "PYTHONUNBUFFERED": "1"}
subprocess.run(cmd, check=True, env=env)
async def _materialize_runtime_backfills(snapshot_dir: Path) -> list[str]: async def _materialize_runtime_backfills(snapshot_dir: Path) -> list[str]:

@ -9,9 +9,11 @@ from __future__ import annotations
import argparse import argparse
import asyncio import asyncio
import concurrent.futures
import datetime as dt import datetime as dt
import json import json
import sys import sys
from collections import defaultdict
from pathlib import Path from pathlib import Path
import asyncpg import asyncpg
@ -231,7 +233,10 @@ def fetch_bars(ticker: str, event_date: str) -> list[PriceBar]:
from datetime import datetime, timedelta from datetime import datetime, timedelta
end_dt = datetime.strptime(event_date, "%Y-%m-%d") end_dt = datetime.strptime(event_date, "%Y-%m-%d")
start_dt = end_dt - timedelta(days=120) start_dt = end_dt - timedelta(days=120)
bars_raw = _fetch_bars_raw(ticker, start_dt.strftime("%Y-%m-%d"), event_date) start_str = start_dt.strftime("%Y-%m-%d")
# Use full-history cache to avoid per-row HTTP calls
all_bars_raw = _fetch_full_bars_raw(ticker)
bars_raw = [b for b in all_bars_raw if start_str <= b.get("date", "") <= event_date]
bars = [] bars = []
for b in bars_raw: for b in bars_raw:
try: try:
@ -243,6 +248,44 @@ def fetch_bars(ticker: str, event_date: str) -> list[PriceBar]:
return bars return bars
def _prefetch_ticker_bars(tickers: list, max_workers: int = 16) -> None:
"""Pre-warm full-history bar cache for all unique tickers in parallel."""
unique = sorted({str(t) for t in tickers if t})
print(f" Prefetching price bars for {len(unique)} tickers...")
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool:
for loaded, _ in enumerate(pool.map(_fetch_full_bars_raw, unique), start=1):
if loaded % 100 == 0 or loaded == len(unique):
print(f" bars prefetch {loaded}/{len(unique)}")
async def _prefetch_short_ratio_db_batch(tickers: list[str]) -> None:
"""Bulk-fetch short ratio history for all tickers in a single DB query."""
dsn = get_settings().postgres_dsn.replace("+asyncpg", "")
conn = await asyncpg.connect(dsn=dsn)
try:
rows = await conn.fetch(
"""
SELECT ticker_raw,
trade_date::text AS date,
short_volume::double precision / NULLIF(total_volume::double precision, 0.0) AS short_ratio
FROM short_sale_daily
WHERE ticker_raw = ANY($1)
AND total_volume IS NOT NULL
AND total_volume > 0
ORDER BY trade_date DESC
""",
tickers,
)
finally:
await conn.close()
grouped: dict[str, list[dict]] = defaultdict(list)
for row in rows:
if row["short_ratio"] is not None:
grouped[row["ticker_raw"]].append({"date": row["date"], "short_ratio": float(row["short_ratio"])})
for ticker in tickers:
_short_ratio_db_cache[ticker] = grouped.get(ticker, [])
def fetch_short_ratio(ticker: str, event_date: str) -> float | None: def fetch_short_ratio(ticker: str, event_date: str) -> float | None:
"""Fetch average short ratio over 5 days before event.""" """Fetch average short ratio over 5 days before event."""
points = _fetch_short_ratio_history_from_db(ticker) points = _fetch_short_ratio_history_from_db(ticker)
@ -377,6 +420,15 @@ def enrich_split(input_path: Path, output_path: Path):
results = {f: [None] * n for f in FEATURES} results = {f: [None] * n for f in FEATURES}
success = 0 success = 0
# Pre-fetch price bars and short ratio history for all unique tickers in batch
_prefetch_ticker_bars(tickers)
unique_tickers = sorted({str(t) for t in tickers if t})
try:
asyncio.run(_prefetch_short_ratio_db_batch(unique_tickers))
print(f" Short ratio DB prefetch done for {len(unique_tickers)} tickers")
except Exception as e:
print(f" Short ratio DB prefetch failed (will fall back per-row): {e}")
for i in range(n): for i in range(n):
ticker = tickers[i] ticker = tickers[i]
event_date = str(event_dates[i]) event_date = str(event_dates[i])

Loading…
Cancel
Save