train.py

"""
⚡ ULTRON Systems // High-Efficiency Autonomous Colab GPU Training Loop
Orchestrates:
- Month-by-month streaming dataset ingestion
- Mixed Precision (FP16) with GradScaler
- Gradient Accumulation for effective batch size
- Auto-resume from Google Drive / local checkpoints
- Real-time VRAM telemetry and directional accuracy metrics
"""
import os
import sys
import gc
import time
import signal
import argparse
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from typing import Optional

from config import DatasetConfig, ModelConfig, TrainingConfig
from dataset import StreamingShardManager
from model import OHLCVTransformer
from loss import MultiTaskFinancialLoss

class Trainer:
    def __init__(self, data_cfg: DatasetConfig, model_cfg: ModelConfig, train_cfg: TrainingConfig):
        self.data_cfg = data_cfg
        self.model_cfg = model_cfg
        self.train_cfg = train_cfg
        
        # Setup device
        self.device = torch.device(train_cfg.device)
        print(f"⚡ [ULTRON] Initializing training on compute device: {self.device}")
        if self.device.type == "cuda":
            print(f"⚡ [ULTRON] GPU: {torch.cuda.get_device_name(0)} | VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
            
        # Determine active checkpoint directory
        if os.path.exists("/content/drive/MyDrive"):
            self.checkpoint_dir = train_cfg.checkpoint_dir
            print(f"⚡ [ULTRON] Google Drive detected. Checkpoints will persist to: {self.checkpoint_dir}")
        else:
            self.checkpoint_dir = train_cfg.local_checkpoint_dir
            print(f"⚡ [ULTRON] Local mode active. Checkpoints will save to: {self.checkpoint_dir}")
        os.makedirs(self.checkpoint_dir, exist_ok=True)
        
        # Instantiate Model, Loss, Optimizer, Scaler
        self.model = OHLCVTransformer(model_cfg).to(self.device)
        self.criterion = MultiTaskFinancialLoss().to(self.device)
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=train_cfg.learning_rate,
            weight_decay=train_cfg.weight_decay,
            betas=(0.9, 0.98)
        )
        self.scaler = torch.cuda.amp.GradScaler(enabled=train_cfg.fp16 and self.device.type == "cuda")
        
        self.global_step = 0
        self.current_epoch = 0
        self.best_dir_acc = 0.0
        
        # Register graceful termination handlers
        signal.signal(signal.SIGINT, self._handle_exit)
        signal.signal(signal.SIGTERM, self._handle_exit)
        
        # Attempt auto-resume
        self._load_checkpoint()

    def _get_lr(self, step: int) -> float:
        """Linear warmup followed by cosine decay."""
        if step < self.train_cfg.warmup_steps:
            return self.train_cfg.learning_rate * float(step) / float(max(1, self.train_cfg.warmup_steps))
        import math
        progress = float(step - self.train_cfg.warmup_steps) / float(max(1, 10000 - self.train_cfg.warmup_steps))
        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress)))
        return max(self.train_cfg.min_lr, self.train_cfg.learning_rate * cosine_decay)

    def _save_checkpoint(self, is_best: bool = False):
        state = {
            "global_step": self.global_step,
            "epoch": self.current_epoch,
            "best_dir_acc": self.best_dir_acc,
            "model_state": self.model.state_dict(),
            "optimizer_state": self.optimizer.state_dict(),
            "scaler_state": self.scaler.state_dict(),
            "model_cfg": self.model_cfg,
            "train_cfg": self.train_cfg
        }
        latest_path = os.path.join(self.checkpoint_dir, "checkpoint_latest.pt")
        torch.save(state, latest_path)
        if is_best:
            best_path = os.path.join(self.checkpoint_dir, "checkpoint_best.pt")
            torch.save(state, best_path)
            print(f"🌟 [ULTRON] Saved new BEST model checkpoint to {best_path}")
        print(f"💾 [ULTRON] Checkpoint synced (Step {self.global_step}).")

    def _load_checkpoint(self):
        latest_path = os.path.join(self.checkpoint_dir, "checkpoint_latest.pt")
        if os.path.exists(latest_path):
            print(f"⚡ [ULTRON] Found existing checkpoint at {latest_path}. Resuming state...")
            try:
                ckpt = torch.load(latest_path, map_location=self.device)
                self.model.load_state_dict(ckpt["model_state"])
                self.optimizer.load_state_dict(ckpt["optimizer_state"])
                if "scaler_state" in ckpt and self.scaler.is_enabled():
                    self.scaler.load_state_dict(ckpt["scaler_state"])
                self.global_step = ckpt.get("global_step", 0)
                self.current_epoch = ckpt.get("epoch", 0)
                self.best_dir_acc = ckpt.get("best_dir_acc", 0.0)
                print(f"✅ [ULTRON] Resumed successfully from Step {self.global_step} (Epoch {self.current_epoch}).")
            except Exception as e:
                print(f"⚠️ [ULTRON] Checkpoint load failed: {e}. Starting fresh.")
        else:
            print("⚡ [ULTRON] No previous checkpoint found. Initializing fresh weights.")

    def _handle_exit(self, signum, frame):
        print(f"\n⚠️ [ULTRON] Interrupt signal received ({signum}). Securing weights...")
        self._save_checkpoint()
        sys.exit(0)

    def train(self):
        print("\n" + "="*70)
        print("⚡ ULTRON // CONVOLUTIONAL TRANSFORMER AUTONOMOUS TRAINING ONLINE")
        print("="*70)
        
        manager = StreamingShardManager(self.data_cfg)
        self.model.train()
        
        for epoch in range(self.current_epoch, self.train_cfg.max_epochs):
            self.current_epoch = epoch
            print(f"\n🚀 [ULTRON] Beginning Epoch {epoch + 1}/{self.train_cfg.max_epochs}")
            
            # Stream monthly shards one by one
            for shard_idx, shard_dataset in enumerate(manager.stream_monthly_shards()):
                loader = DataLoader(
                    shard_dataset,
                    batch_size=self.train_cfg.batch_size,
                    shuffle=True,
                    num_workers=2,
                    pin_memory=True if self.device.type == "cuda" else False,
                    drop_last=True
                )
                
                shard_loss = 0.0
                shard_acc = 0.0
                steps_in_shard = 0
                t0 = time.time()
                
                self.optimizer.zero_grad()
                
                for batch_idx, (batch_x, batch_y) in enumerate(loader):
                    batch_x = batch_x.to(self.device, non_blocking=True)
                    batch_y = batch_y.to(self.device, non_blocking=True)
                    
                    # Dynamic learning rate schedule
                    lr = self._get_lr(self.global_step)
                    for param_group in self.optimizer.param_groups:
                        param_group["lr"] = lr
                        
                    # Mixed precision forward pass
                    with torch.cuda.amp.autocast(enabled=self.train_cfg.fp16 and self.device.type == "cuda"):
                        pred_returns, pred_logits = self.model(batch_x)
                        loss, metrics = self.criterion(pred_returns, pred_logits, batch_y)
                        # Scale loss for gradient accumulation
                        loss_scaled = loss / self.train_cfg.grad_accum_steps
                        
                    # Backward pass
                    self.scaler.scale(loss_scaled).backward()
                    
                    shard_loss += metrics["total_loss"]
                    shard_acc += metrics["dir_acc_bce"]
                    steps_in_shard += 1
                    
                    # Optimizer step on accumulation boundaries
                    if (batch_idx + 1) % self.train_cfg.grad_accum_steps == 0:
                        self.scaler.unscale_(self.optimizer)
                        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                        self.scaler.step(self.optimizer)
                        self.scaler.update()
                        self.optimizer.zero_grad()
                        self.global_step += 1
                        
                        # Logging telemetry
                        if self.global_step % self.train_cfg.log_interval == 0:
                            vram_gb = torch.cuda.memory_allocated() / 1e9 if self.device.type == "cuda" else 0.0
                            rate = (self.train_cfg.batch_size * self.train_cfg.grad_accum_steps) / max(1e-4, (time.time() - t0))
                            print(
                                f"⚡ [Step {self.global_step:05d}] "
                                f"Loss: {metrics['total_loss']:.4f} "
                                f"(Reg: {metrics['reg_loss']:.4f}, Dir: {metrics['dir_loss']:.4f}) | "
                                f"Dir Acc: {metrics['dir_acc_bce']:.1f}% | "
                                f"LR: {lr:.2e} | "
                                f"VRAM: {vram_gb:.2f}GB | "
                                f"{rate:.1f} win/s"
                            )
                            t0 = time.time()
                            
                        # Periodic checkpointing
                        if self.global_step % self.train_cfg.save_every_steps == 0:
                            avg_acc = shard_acc / max(1, steps_in_shard)
                            is_best = avg_acc > self.best_dir_acc
                            if is_best:
                                self.best_dir_acc = avg_acc
                            self._save_checkpoint(is_best=is_best)
                            
                    # Cap steps per shard if configured to cycle through years faster
                    if steps_in_shard >= self.train_cfg.max_steps_per_shard:
                        break
                        
                # Clean up shard memory explicitly
                del loader, shard_dataset
                if self.device.type == "cuda":
                    torch.cuda.empty_cache()
                gc.collect()
                
            # End of epoch checkpoint
            self._save_checkpoint()
            
        print("\n✅ [ULTRON] Training loop complete. Final weights secured.")

def parse_args():
    parser = argparse.ArgumentParser(description="ULTRON OHLCV Transformer Trainer")
    parser.add_argument("--start_year", type=int, default=2021, help="Start year of shards")
    parser.add_argument("--end_year", type=int, default=2025, help="End year of shards")
    parser.add_argument("--batch_size", type=int, default=64, help="Batch size per GPU step")
    parser.add_argument("--grad_accum", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
    parser.add_argument("--epochs", type=int, default=5, help="Total training epochs")
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    data_cfg = DatasetConfig(start_year=args.start_year, end_year=args.end_year)
    model_cfg = ModelConfig()
    train_cfg = TrainingConfig(
        batch_size=args.batch_size,
        grad_accum_steps=args.grad_accum,
        learning_rate=args.lr,
        max_epochs=args.epochs
    )
    trainer = Trainer(data_cfg, model_cfg, train_cfg)
    trainer.train()