utils.py

import os
import datetime
import traceback
import sys
import pandas as pd
import scipy as sp
from sklearn.model_selection import KFold
import optuna
import numpy as np
from data_utils import _data_listing, create_dataloader
import copy
import config as cfg


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.distributed as dist
from torch.optim.lr_scheduler import _LRScheduler, CosineAnnealingLR
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp


# ================================
# UTILS
# ================================


class gradual_warmup_scheduler(_LRScheduler):
    """Gradually warm-up(increasing) learning rate in optimizer."""

    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):
        self.multiplier = multiplier
        if self.multiplier < 1.0:
            raise ValueError("multiplier should be greater thant or equal to 1.")
        self.total_epoch = total_epoch
        self.after_scheduler = after_scheduler
        self.finished = False
        super(gradual_warmup_scheduler, self).__init__(optimizer)

    def get_lr(self):
        if self.last_epoch > self.total_epoch:
            if self.after_scheduler:
                if not self.finished:
                    self.after_scheduler.base_lrs = [
                        base_lr * self.multiplier for base_lr in self.base_lrs
                    ]
                    self.finished = True
                return self.after_scheduler.get_last_lr()
            return [base_lr * self.multiplier for base_lr in self.base_lrs]

        if self.multiplier == 1.0:
            return [
                base_lr * (float(self.last_epoch) / self.total_epoch)
                for base_lr in self.base_lrs
            ]
        else:
            return [
                base_lr
                * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch + 1.0)
                for base_lr in self.base_lrs
            ]

    def step(self, epoch=None, metrics=None):
        if self.finished and self.after_scheduler:
            if epoch is None:
                self.after_scheduler.step(None)
            else:
                self.after_scheduler.step(epoch - self.total_epoch)
            self._last_lr = self.after_scheduler.get_last_lr()
        else:
            return super(gradual_warmup_scheduler, self).step(epoch)


def reset_weights(m):
    """Reset model weights to avoid weight leakage."""
    for layer in m.children():
        if hasattr(layer, "reset_parameters"):
            layer.reset_parameters()


class DistributedSetup:
    def __init__(self, master_addr="localhost", master_port="12355"):
        self.master_addr = master_addr
        self.master_port = master_port

    def setup_distributed(self, rank, world_size, logger):
        """Initialize the distributed environment with proper error handling and debugging"""
        os.environ["MASTER_ADDR"] = self.master_addr
        os.environ["MASTER_PORT"] = self.master_port

        # Set NCCL debugging info if needed
        # os.environ["NCCL_DEBUG"] = "INFO"

        # Initialize the process group with timeout
        try:
            dist.init_process_group(
                backend="nccl",
                rank=rank,
                world_size=world_size,
                timeout=datetime.timedelta(seconds=60),  # Increased timeout
            )
            logger.info(f"Successfully initialized process group for rank {rank}")
        except Exception as e:
            logger.error(f"Failed to initialize process group on rank {rank}: {e}")
            logger.error(traceback.format_exc())
            sys.exit(1)

        # Set the device and print device info
        torch.cuda.set_device(rank)
        logger.info(f"Rank {rank} using GPU: {torch.cuda.get_device_name(rank)}")

        # Verify CUDA device is correctly set
        current_device = torch.cuda.current_device()
        if current_device != rank:
            logger.warning(
                f"Rank {rank}: CUDA device mismatch! Expected {rank}, got {current_device}"
            )

        # Synchronize all processes
        torch.cuda.synchronize()

        # Make CUDA operations asynchronous for better performance
        torch.cuda.set_device(rank)

        # Print memory info
        logger.info(
            f"Rank {rank} - Total GPU memory: {torch.cuda.get_device_properties(rank).total_memory / 1e9:.2f} GB"
        )
        logger.info(
            f"Rank {rank} - Available GPU memory: {torch.cuda.memory_allocated(rank) / 1e9:.2f} GB used"
        )

    def cleanup(self):
        """Clean up the distributed environment"""
        if dist.is_initialized():
            dist.destroy_process_group()


# ================================
# TRAINING CLASSES
# ================================


class TrainEngine:
    def __init__(self, logger, sequence_len=249, gradient_penalty=False):
        self.logger = logger
        self.sequence_len = sequence_len
        self.gradient_penalty = gradient_penalty

    def train_epoch(
        self, model, train_loader, optimizer, criterion, local_rank, scheduler=None
    ):
        """Train for one epoch"""
        model.train()
        running_loss = 0.0

        for step, batch in enumerate(train_loader):
            inputs, labels, _, _, _, _ = batch
            inputs = inputs.to(local_rank, non_blocking=True)
            inputs = inputs.squeeze().permute(0, 2, 1)
            labels = labels.to(local_rank, non_blocking=True)

            optimizer.zero_grad()

            outputs = model(inputs)

            # self.logger.info(f"Outputs shape: {outputs.shape}, Labels shape: {labels.shape}")


            # ortho_loss = F.cosine_similarity(features[0], features[1]).mean()

            loss = criterion(outputs.squeeze(), labels) # + 0.1 * ortho_loss

            # ------ gradient penalty ------
            if self.gradient_penalty:
                inputs.requires_grad_(True)

                training = model.training
                model.eval()

                outputs_gp = model(inputs)

                grads = torch.autograd.grad(outputs_gp.sum(), inputs, create_graph=True, retain_graph=True)[0]

                model.train(training)

                loss = loss + 0.01 * grads.pow(2).mean()
            # ------------------------------
            
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

            if step % 2000 == 0 and local_rank == 0:
                self.logger.info(
                    f"Step {step}/{len(train_loader)}: loss - {loss.item():.3f}"
                )

            del outputs, inputs, labels
            torch.cuda.empty_cache()

        # Average loss across all batches
        avg_loss = running_loss / len(train_loader)

        # If distributed, sync losses across processes
        if dist.is_initialized():
            dist.barrier()
            loss_tensor = torch.tensor(avg_loss).to(local_rank)
            dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
            avg_loss = loss_tensor.item() / dist.get_world_size()

        if scheduler is not None:
            scheduler.step()

        return avg_loss

    def validate(self, model, criterion, val_loader, device):
        """Validate the model"""
        model = model.to(device)
        model.eval()
        total_val_loss = 0.0
        fold_results = []

        with torch.no_grad():
            for batch in val_loader:
                inputs, labels, _, _, _, _ = batch

                inputs = inputs.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)
                inputs = inputs.squeeze().permute(0, 2, 1)

                outputs = model(inputs)

                for ix in range(len(outputs)):
                    fold_results.append(
                        [
                            outputs.cpu().detach().numpy()[ix],
                            labels[ix].cpu().detach().numpy(),
                        ]
                    )

                loss = criterion(outputs.squeeze(), labels)
                total_val_loss += loss.item()

                del outputs, inputs, labels
                torch.cuda.empty_cache()

        # Calculate metrics
        results_data = pd.DataFrame(fold_results, columns=["predictions", "observed"])

        preds = np.vstack(results_data["predictions"].values)
        obs = np.vstack(results_data["observed"].values)

        correlations = []

        for i in range(preds.shape[1]):  # iterate over columns
            r, pval = sp.stats.spearmanr(preds[:, i], obs[:, i])
            correlations.append(r)

        # Average loss across all batches
        avg_loss = total_val_loss / len(val_loader)

        # If distributed, gather all results and compute metrics on rank 0
        if dist.is_initialized():
            loss_tensor = torch.tensor(avg_loss).to(device)
            r_tensor = torch.tensor(correlations).to(device)

            dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
            dist.all_reduce(r_tensor, op=dist.ReduceOp.SUM)

            avg_loss = loss_tensor.item() / dist.get_world_size()
            correlations = r_tensor.cpu().numpy() / dist.get_world_size()

        return avg_loss, correlations

    def test(self, model, test_loader, device):
        """Test the model"""
        model = model.to(device)
        model.eval()
        fold_results = []

        with torch.no_grad():
            for batch in test_loader:
                inputs, labels, chroms, starts, ends, strands = batch
                inputs = inputs.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)
                inputs = inputs.squeeze().permute(0, 2, 1)

                outputs = model(inputs)

                for ix in range(len(outputs)):
                    fold_results.append(
                        [
                            chroms[ix],
                            starts.cpu().detach().numpy()[ix],
                            ends.cpu().detach().numpy()[ix],
                            strands[ix],
                            outputs.cpu().detach().numpy()[ix],
                            labels[ix].cpu().detach().numpy(),
                        ]
                    )

                del outputs, inputs, labels
                torch.cuda.empty_cache()

        # Calculate metrics
        results_data = pd.DataFrame(
            fold_results,
            columns=["chr", "start", "end", "strand", "predictions", "observed"],
        )
        
        preds = np.vstack(results_data["predictions"].values)
        obs = np.vstack(results_data["observed"].values)

        correlations = []
        for i in range(preds.shape[1]):  # iterate over columns
            r, pval = sp.stats.spearmanr(preds[:, i], obs[:, i])
            correlations.append(r)

        results_data.to_csv(f"test_results.csv", index=False)

        # Sync r value across all ranks
        if dist.is_initialized():
            r_tensor = torch.tensor(correlations).to(device)
            dist.all_reduce(r_tensor, op=dist.ReduceOp.SUM)
            correlations = r_tensor.cpu().numpy() / dist.get_world_size()

        return correlations

    def train_distributed(
        self,
        rank,
        world_size,
        opt_config,
        train_config,
        fold_config,
        model,
    ):
        """Main function for distributed training"""
        # Set up distributed environment
        setup = DistributedSetup()
        setup.setup_distributed(rank, world_size, self.logger)

        model.apply(reset_weights)

        # Create model and move to current device
        model = model.to(rank)
        model = DDP(
            model, device_ids=[rank], output_device=rank, find_unused_parameters=False
        )

        # Define loss function
        # criterion = nn.MSELoss()
        # criterion = getattr(loss, opt_config["criterion"])()
        criterion = opt_config["criterion"]

        # Create optimizer
        optimizer = getattr(optim, opt_config["optimizer"])(
            model.parameters(),
            lr=opt_config["min_lr"],
            weight_decay=opt_config["decay"],
        )

        per_gpu_batch_size = max(1, opt_config["batch"] // world_size)
        self.logger.info(
            f"Rank {rank}: Using batch size of {per_gpu_batch_size} per GPU"
        )

        # Create scheduler
        cosine_scheduler = CosineAnnealingLR(
            optimizer,
            T_max=train_config["n_epochs"] - train_config["warmup_epochs"],
            eta_min=opt_config["min_lr"],
        )

        scheduler = gradual_warmup_scheduler(
            optimizer,
            multiplier=opt_config["max_lr"] / opt_config["min_lr"],
            total_epoch=train_config["warmup_epochs"],
            after_scheduler=cosine_scheduler,
        )

        # Create train dataloader
        train_data_df = _data_listing(
            fold_config["location"], fold_config["train_chunk"]
        )
        
        train_loader, train_sampler = create_dataloader(
            hdf5_files=train_data_df,
            batch_size=per_gpu_batch_size,
            bp_len=self.sequence_len,
            is_train=True,
            num_workers=8,
            multihead=train_config['multihead']
        )

        for epoch in range(train_config["n_epochs"]):
            # Set epoch for distributed sampler
            if train_sampler:
                train_sampler.set_epoch(epoch)

            # Train for one epoch
            train_loss = self.train_epoch(
                model=model,
                train_loader=train_loader,
                optimizer=optimizer,
                criterion=criterion,
                local_rank=rank,
                scheduler=scheduler,
            )

            # Log results on rank 0
            if rank == 0:
                self.logger.info(f"Epoch {epoch+1}/{train_config['n_epochs']}")
                self.logger.info(f"Train loss: {train_loss:.4f}")

        if rank == 0:
            try:
                torch.save(
                    model.state_dict(), f"best_model_fold{fold_config['fold']}.pth"
                )
            except:
                pass

        # all processes are synchronized before finishing
        dist.barrier()
        setup.cleanup()


def run_optuna_trial(
    model, trial, train_config, data_config, device, logger, sequence_len
):
    """Run a single Optuna trial"""
    # Sample hyperparameters
    opt_config = {
        "optimizer": trial.suggest_categorical(
            "optimizer", ["Adam", "RMSprop", "SGD", "AdamW"]
        ),
        "min_lr": trial.suggest_float("min_lr", cfg.LOW_MIN_LR, cfg.HIGH_MIN_LR, log=True),
        "max_lr": trial.suggest_float("max_lr", cfg.LOW_MAX_LR, cfg.HIGH_MAX_LR, log=True),
        "decay": trial.suggest_float("decay", cfg.LOW_DECAY, cfg.HIGH_DECAY, log=True),
        "batch": trial.suggest_int("batch", cfg.LOW_BATCH, cfg.HIGH_BATCH, step=cfg.BATCH_STEP),
        "criterion": nn.MSELoss()
        # "criterion": "MultiTaskLossWrapper"
        # "criterion": trial.suggest_categorical(
        #     "criterion", ["MSEwithEntropy", "MSEKLmixed", "L1KLmixed", "L1withEntropy"]
        # )
    }

    # Cross-validation
    total_folds = train_config["kfold"]
    kfold = KFold(n_splits=total_folds, shuffle=True, random_state=42)

    train_engine = TrainEngine(logger=logger, sequence_len=sequence_len)

    total_corr = []

    # Launch distributed training
    world_size = train_config["gpus"] * train_config["nodes"]
    opt_config["batch"] = opt_config["batch"] * world_size

    for fold, (train_idx, val_idx) in enumerate(
        kfold.split(data_config["train_chunk"])
    ):
        # if train_config["local_rank"] == 0:
        logger.info(f"Starting fold {fold+1}/{total_folds}")

        # Create fold config
        fold_config = {
            "fold": fold,
            "location": data_config["location"],
            "train_chunk": [data_config["train_chunk"][i] for i in train_idx],
            "val_chunk": [data_config["train_chunk"][i] for i in val_idx],
        }

        # Create a temporary model to reset weights
        # temp_model = InceptionNet1D(num_classes=1, in_channels=4)
        temp_model = copy.deepcopy(model)
        temp_model.apply(reset_weights)

        # Use multiprocessing to start processes
        
        # --- RESILIENCY LOGIC START ---
        max_retries = 3
        attempt = 0
        success = False

        while attempt < max_retries and not success:
            try:
                temp_model = copy.deepcopy(model)
                temp_model.apply(reset_weights)

                mp.spawn(
                    train_engine.train_distributed,
                    args=(world_size, opt_config, train_config, fold_config, temp_model),
                    nprocs=world_size,
                    join=True,
                )
                success = True
            except Exception as e:
                attempt += 1
                logger.error(f"Fold {fold+1} failed at attempt {attempt} with error: {e}")
                
                # Clean up and wait before retrying
                torch.cuda.empty_cache()
                if dist.is_initialized():
                    dist.destroy_process_group()
                
                if attempt < max_retries:
                    wait_time = 10 * attempt  # Incremental backoff
                    logger.info(f"Retrying fold {fold+1} in {wait_time} seconds...")
                    datetime.time.sleep(wait_time)
                else:
                    logger.critical(f"Fold {fold+1} failed after {max_retries} attempts. Aborting trial.")
                    raise e

        # Load best model and evaluate
        # if train_config["local_rank"] == 0:
        # model = InceptionNet1D(num_classes=1, in_channels=4).to(device)
        weights = torch.load(f"best_model_fold{fold}.pth")
        weights = {k.replace("module.", ""): v for k, v in weights.items()}
        model.load_state_dict(weights)

        val_loader, _ = create_dataloader(
            hdf5_files=_data_listing(fold_config["location"], fold_config["val_chunk"]),
            batch_size=opt_config["batch"],
            bp_len=sequence_len,
            is_train=False,
            num_workers=4,
            multihead=train_config['multihead']
        )

        _, val_corr = train_engine.validate(
            model=model,
            criterion=opt_config["criterion"],
            # criterion=getattr(loss, opt_config["criterion"])(),
            val_loader=val_loader,
            device=device,
        )

        total_corr.append(val_corr)
        np.set_printoptions(precision=2)
        logger.info(f"Fold {fold+1} completed with correlation: {np.round(np.array(val_corr), 2)}")

    # Average correlation across folds
    mean_corr = np.mean(total_corr, axis=0)

    np.set_printoptions(precision=4)
    logger.info(f"[INFO] Mean correlation across folds: {mean_corr}")

    if cfg.NUM_CLASSES == 1:
        trial.report(mean_corr, 0)
        if trial.should_prune():
            raise optuna.exceptions.TrialPruned()

    return tuple(mean_corr)