SPB Git

spb/wp39_uqo

Public
5.5 KB · 117 lines python
1#!/usr/bin/env python2"""Phase 2c: streaming export -> session features at 1m/5m/30m/1d -> rolling windows, per symbol (multiprocess).3Raw Parquet is deleted after the feature rows are written; manifests keep hashes and dataset versions."""4from __future__ import annotations56import argparse7import datetime as dt8import json9import os10import sys11import time12from concurrent.futures import ProcessPoolExecutor, as_completed13from pathlib import Path1415import polars as pl1617sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))18from wp39.api_client import FinDataClient  # noqa: E40219from wp39.corporate_actions import load_actions  # noqa: E40220from wp39.features import compute_spy_buckets, compute_symbol_features, session_bounds_table  # noqa: E40221from wp39.ingest import export_symbol, mark_consumed  # noqa: E40222from wp39.pipeline import Paths, get_logger, load_config  # noqa: E40223from wp39.provenance import append_jsonl, now_iso, write_json  # noqa: E40224from wp39.samples import rolling_window_features  # noqa: E4022526START, END = dt.date(2011, 1, 1), dt.date(2026, 10, 30)272829def work(symbol: str, keep_raw: bool) -> dict:30    P = Paths.default()31    contract = load_config("data_contract")32    ds = contract["findata"]["datasets"]["bars_1m"]33    out = P.cache / "features" / f"{symbol}.parquet"34    win_out = P.cache / "windows" / f"{symbol}.parquet"35    win_out.parent.mkdir(exist_ok=True)36    if out.exists() and win_out.exists():37        return {"symbol": symbol, "status": "cached"}38    t0 = time.time()39    client = FinDataClient(log_path=P.logs / f"api_calls_worker_{os.getpid()}.jsonl")40    try:41        paths = export_symbol(client, symbol, START, END, ds, P.cache / "bars", P.manifests / "bars_manifest.jsonl", keep=keep_raw)42    finally:43        client.close()44    if not paths:45        return {"symbol": symbol, "status": "no_export"}46    bounds = pl.read_parquet(P.cache / "external" / "session_bounds.parquet")47    spy = {k: pl.read_parquet(P.cache / "external" / f"spy_buckets_{k}.parquet") for k in ("1m", "5m", "30m")}48    actions = pl.read_parquet(P.cache / "external" / "actions.parquet")49    feat = compute_symbol_features(paths, symbol, bounds, spy, actions)50    if feat.height == 0:51        return {"symbol": symbol, "status": "empty"}52    feat.write_parquet(out)53    rolling_window_features(feat).write_parquet(win_out)54    if not keep_raw:55        for p in paths:56            p.unlink(missing_ok=True)57        mark_consumed(P.manifests / "bars_manifest.jsonl", symbol, ds)58    return {"symbol": symbol, "status": "ok", "sessions": feat.height, "first": str(feat["session_date"].min()),59            "last": str(feat["session_date"].max()), "secs": round(time.time() - t0, 1)}606162def main() -> int:63    ap = argparse.ArgumentParser()64    ap.add_argument("--workers", type=int, default=6)65    ap.add_argument("--limit", type=int, default=0, help="pilot: only the first N eligible symbols")66    ap.add_argument("--symbols", type=str, default="", help="comma-separated explicit list")67    ap.add_argument("--keep-raw", action="store_true")68    args = ap.parse_args()69    P = Paths.default()70    log = get_logger("03_features", P.logs / "03_features.log")71    contract = load_config("data_contract")72    ext = P.cache / "external"73    # shared inputs: session bounds, SPY buckets, corporate actions74    if not (ext / "session_bounds.parquet").exists():75        session_bounds_table(START, END).write_parquet(ext / "session_bounds.parquet")76    if not (ext / "spy_buckets_1m.parquet").exists():77        client = FinDataClient(log_path=P.logs / "api_calls.jsonl")78        paths = export_symbol(client, "SPY", START, END, contract["findata"]["datasets"]["bars_1m_etf"], P.cache / "bars",79                              P.manifests / "bars_manifest.jsonl", keep=True)80        client.close()81        bounds = pl.read_parquet(ext / "session_bounds.parquet")82        for k, v in compute_spy_buckets(paths, bounds).items():83            v.write_parquet(ext / f"spy_buckets_{k}.parquet")84        log.info("SPY buckets written")85    if not (ext / "actions.parquet").exists():86        load_actions(ext / "corporate_actions_raw.csv").write_parquet(ext / "actions.parquet")87    if args.symbols:88        symbols = args.symbols.split(",")89    else:90        uni = pl.read_parquet(P.processed / "universe.parquet").filter(pl.col("eligible")).sort("ticker")91        symbols = uni["ticker"].to_list()92        if args.limit:93            symbols = symbols[: args.limit]94    log.info("symbols to process: %d (workers=%d)", len(symbols), args.workers)95    status = P.logs / "03_features_status.jsonl"96    done = {r["symbol"] for r in (json.loads(l) for l in open(status)) if r.get("status") in ("ok", "cached", "empty", "no_export")} if status.exists() else set()97    todo = [s for s in symbols if s not in done]98    t0 = time.time(); n_ok = 099    with ProcessPoolExecutor(max_workers=args.workers) as ex:100        futs = {ex.submit(work, s, args.keep_raw): s for s in todo}101        for i, f in enumerate(as_completed(futs)):102            try:103                r = f.result()104            except Exception as exc:  # noqa: BLE001105                r = {"symbol": futs[f], "status": "error", "error": str(exc)[:300]}106            r["ts"] = now_iso()107            append_jsonl(status, r)108            n_ok += r["status"] in ("ok", "cached")109            if i % 50 == 0:110                log.info("%d/%d done (ok=%d) elapsed=%.0fs last=%s", i + 1, len(todo), n_ok, time.time() - t0, r)111    log.info("finished: %d symbols, %.0fs", len(todo), time.time() - t0)112    return 0113114115if __name__ == "__main__":116    raise SystemExit(main())117
Files · 42edd965d464895e91a65b5b6e7ea7e9f619e658
Menu

Theme