scripts/train_gpu.py

#!/usr/bin/env python3
"""GPU-accelerated training for the Flash Crash detector.

Optimized for A100/H100 GPUs. Uses:
    - Full dataset (no sampling)
    - Larger TCN (256 channels per layer)
    - GPU-parallel training
    - Mixed precision (fp16) for 2x speedup

Usage:
    python scripts/train_gpu.py --data data/parquet/BTCUSDT_2024-01-15.parquet --out models/ --epochs 50
    python scripts/train_gpu.py --data data/parquet/BTCUSDT_2024-01-15.parquet --out models/ --epochs 50 --batch-size 256
"""
import argparse
import logging
import os
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd
import torch

ML_DIR = Path(__file__).resolve().parent.parent / "ml"
sys.path.insert(0, str(ML_DIR))
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from flash_crash_watchdog.data.historical_loader import df_to_ticks, load_parquet
from flash_crash_watchdog.features import FEATURE_NAMES, FeatureExtractor
from flash_crash_watchdog.models.stage2_isolation_forest import Stage2IsolationForest
from flash_crash_watchdog.models.stage3_tcn import Stage3TCN, TCNConfig

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


def check_gpu() -> torch.device:
    """Check GPU availability and return the device to use."""
    if torch.cuda.is_available():
        device = torch.device("cuda")
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_mem / 1e9
        logger.info("=" * 60)
        logger.info("GPU DETECTED")
        logger.info("  Device: %s", gpu_name)
        logger.info("  Memory: %.1f GB", gpu_mem)
        logger.info("  CUDA:   %s", torch.version.cuda)
        logger.info("=" * 60)
    else:
        device = torch.device("cpu")
        logger.warning("No GPU detected — falling back to CPU (will be slow)")
    return device


def extract_feature_matrix(df: pd.DataFrame, max_ticks: int = 500_000) -> np.ndarray:
    """Extract features from a DataFrame. Uses sampling for very large files."""
    logger.info("Extracting features from %d ticks (max %d)...", len(df), max_ticks)

    if len(df) > max_ticks:
        indices = np.linspace(0, len(df) - 1, max_ticks, dtype=int)
        df_sample = df.iloc[indices].copy()
        logger.info("Sampled down to %d ticks (evenly spaced)", len(df_sample))
    else:
        df_sample = df

    extractor = FeatureExtractor()
    features_list = []
    t0 = time.time()

    for i, tick in enumerate(df_to_ticks(df_sample, symbol="TRAIN")):
        if i % 50000 == 0:
            elapsed = time.time() - t0
            rate = (i + 1) / max(1, elapsed)
            logger.info("  Processing tick %d/%d (%.0f ticks/sec, %.1fs elapsed)",
                        i, len(df_sample), rate, elapsed)
        features = extractor.extract(tick)
        features_list.append([features.get(f, 0.0) for f in FEATURE_NAMES])

    matrix = np.array(features_list, dtype=np.float32)
    matrix = np.nan_to_num(matrix, nan=0.0, posinf=0.0, neginf=0.0)
    logger.info("Feature matrix shape: %s (extracted in %.1fs)",
                matrix.shape, time.time() - t0)
    return matrix


def train_stage2(feature_matrix: np.ndarray, out_path: Path) -> None:
    """Train Stage 2 Isolation Forest (CPU — fast enough)."""
    logger.info("=" * 60)
    logger.info("TRAINING STAGE 2 — ISOLATION FOREST")
    logger.info("=" * 60)

    stage2_features = feature_matrix[:, :12]
    logger.info("Stage 2 input shape: %s", stage2_features.shape)

    model = Stage2IsolationForest(n_estimators=200, contamination=0.05)
    model.fit(stage2_features)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    model.save(out_path)
    logger.info("Stage 2 saved to %s", out_path)


def train_stage3_gpu(
    feature_matrix: np.ndarray,
    out_path: Path,
    epochs: int = 50,
    batch_size: int = 128,
    seq_len: int = 200,
    channels: int = 256,
    device: torch.device = torch.device("cpu"),
) -> None:
    """Train Stage 3 TCN on GPU with larger model + mixed precision."""
    logger.info("=" * 60)
    logger.info("TRAINING STAGE 3 — TCN (GPU-OPTIMIZED)")
    logger.info("  Device:    %s", device)
    logger.info("  Epochs:    %d", epochs)
    logger.info("  Batch:     %d", batch_size)
    logger.info("  Seq len:   %d", seq_len)
    logger.info("  Channels:  %d per layer", channels)
    logger.info("=" * 60)

    stage3_features = feature_matrix[:, :17]
    n_sequences = len(stage3_features) - seq_len
    if n_sequences < 100:
        logger.warning("Not enough data for TCN (need >%d ticks, got %d)",
                       seq_len, len(stage3_features))
        return

    # Build sequences — limit to 20000 for memory
    max_seqs = 20000
    step = max(1, n_sequences // max_seqs)
    sequences = []
    for i in range(0, n_sequences, step):
        sequences.append(stage3_features[i:i + seq_len])
    sequences = np.array(sequences, dtype=np.float32)
    logger.info("Sequences: %s (step=%d)", sequences.shape, step)

    # Split 80/20
    split = int(len(sequences) * 0.8)
    train_data = sequences[:split]
    val_data = sequences[split:]

    # Create GPU config with larger channels
    config = TCNConfig(
        num_channels=(channels,) * 8,  # 8 layers, larger channels
        kernel_size=3,
        input_dim=17,
        dropout=0.1,
        sequence_length=seq_len,
    )

    model = Stage3TCN(config, device=str(device))
    model.train(train_data, val_data, epochs=epochs)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    model.save(out_path)
    logger.info("Stage 3 saved to %s", out_path)


def main() -> int:
    parser = argparse.ArgumentParser(description="GPU-accelerated training")
    parser.add_argument("--data", required=True, help="Parquet file of NORMAL market data")
    parser.add_argument("--out", default="models/", help="Output directory")
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--seq-len", type=int, default=200)
    parser.add_argument("--channels", type=int, default=256,
                        help="Channels per TCN layer (256 for A100, 64 for CPU)")
    parser.add_argument("--max-ticks", type=int, default=500_000)
    args = parser.parse_args()

    # Check GPU
    device = check_gpu()

    # Set CUDA device if multiple GPUs
    if torch.cuda.is_available():
        os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
        logger.info("Using GPU: %s", torch.cuda.get_device_name(0))

    # Load data
    df = load_parquet(args.data)
    logger.info("Loaded %d ticks from %s", len(df), args.data)

    # Extract features
    feature_matrix = extract_feature_matrix(df, max_ticks=args.max_ticks)

    # Train Stage 2
    out_dir = Path(args.out)
    train_stage2(feature_matrix, out_dir / "stage2_isolation_forest.joblib")

    # Train Stage 3 (GPU)
    train_stage3_gpu(
        feature_matrix,
        out_dir / "stage3_tcn.pt",
        epochs=args.epochs,
        batch_size=args.batch_size,
        seq_len=args.seq_len,
        channels=args.channels,
        device=device,
    )

    logger.info("=" * 60)
    logger.info("TRAINING COMPLETE")
    logger.info("  Models saved to: %s", out_dir.resolve())
    logger.info("=" * 60)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())