fleet_train.py

"""
Fleet Training: Train multiple models (one per cluster) and aggregate metrics.

Main flow:
1. Scan train/val/test directories for matching cluster files
2. Train one model per cluster
3. Aggregate metrics: fleet_metric = avg(cluster_metrics)
4. Log aggregated metrics to MLflow
"""

import os
import sys
import warnings
import argparse
import json
from pathlib import Path
from glob import glob
from importlib import import_module
from functools import partial
from collections import defaultdict

import numpy as np
import torch
from transformers import Trainer, TrainingArguments, AutoTokenizer

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from model.dataset import ListDataset, HierarchicalSyntheticDataset
from model.utils import (
    autoconstruct_run_name,
    init_mlflow,
    custom_collate,
    compute_metrics,
    find_gpu_agnostic_checkpoint,
)
from configs.default_config import get_cfg_defaults

warnings.simplefilter("ignore")
os.environ["PYDEVD_WARN_SLOW_RESOLVE_TIMEOUT"] = "15"


def find_cluster_files(train_dir, val_dir, test_dir):
    """
    Find matching cluster files across train/val/test directories.
    
    Returns list of dicts: [{'name': cluster_name, 'train': path, 'val': path, 'test': path}, ...]
    """
    print(f"\n=== Scanning for cluster files ===")
    print(f"Train dir: {train_dir}")
    print(f"Val dir:   {val_dir}")
    print(f"Test dir:  {test_dir}")
    
    # Get all cluster names from train directory
    train_files = glob(f"{train_dir}/*.tsv")
    cluster_names = [Path(f).stem for f in train_files if not Path(f).stem.startswith("all")]
    
    print(f"\nFound {len(cluster_names)} clusters in train directory")
    
    # Match with val and test
    clusters = []
    for name in cluster_names:
        train_path = os.path.join(train_dir, f"{name}.tsv")
        val_path = os.path.join(val_dir, f"{name}.tsv") if val_dir and os.path.isdir(val_dir) else None
        test_path = os.path.join(test_dir, f"{name}.tsv")
        
        # Validate test exists
        if not os.path.exists(test_path):
            print(f"WARNING: No test file for {name}, skipping")
            continue
        
        # Check if val exists
        if val_path and not os.path.exists(val_path):
            val_path = None
        
        clusters.append({
            'name': name,
            'train': train_path,
            'val': val_path,
            'test': test_path,
        })
        
        print(f"  ✓ {name} (val: {'yes' if val_path else 'no'})")
    
    print(f"\nTotal clusters ready for training: {len(clusters)}")
    return clusters


def load_pretrained_model(cfg, device):
    """Load model from pretrained checkpoint or random init."""
    model_module = import_module(f"model.{cfg.MODEL.ARCHITECTURE}")
    model = model_module.ddGRegressor(cfg=cfg).to(device)
    
    checkpoint_path = cfg.FLEET.PRETRAINED_CHECKPOINT
    if not checkpoint_path:
        print("  → Using random initialization")
        return model
    
    print(f"  → Loading pretrained weights from {checkpoint_path}")
    safetensors_path = os.path.join(checkpoint_path, "model.safetensors")
    
    if os.path.exists(safetensors_path):
        from safetensors.torch import load_file
        state_dict = load_file(safetensors_path, device="cpu")
        model.load_state_dict(state_dict)
    else:
        raise ValueError(f"No model.safetensors found in {checkpoint_path}")
    
    return model


def train_cluster_model(cfg, cluster, base_run_name, device):
    """
    Train a single model for one cluster.
    
    Returns dict with training metrics history.
    """
    cluster_name = cluster['name']
    print(f"\n=== Training: {cluster_name} ===")
    
    # Load datasets
    train_dataset = ListDataset(cluster['train'], reverse_augmentation=True)
    eval_datasets = {}
    
    # Add test dataset
    eval_datasets[f"{cluster_name}_test"] = ListDataset(cluster['test'])
    
    # Add val dataset if exists
    if cluster['val']:
        eval_datasets[f"{cluster_name}_val"] = ListDataset(cluster['val'])
        print(f"  Train: {len(train_dataset)}, Val: {len(eval_datasets[f'{cluster_name}_val'])}, Test: {len(eval_datasets[f'{cluster_name}_test'])}")
    else:
        print(f"  Train: {len(train_dataset)}, Test: {len(eval_datasets[f'{cluster_name}_test'])}")
    
    # Load model
    model = load_pretrained_model(cfg, device)
    
    # Setup tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        cfg.MODEL.ESM_CHECKPOINT,
        padding=True,
        max_length=cfg.TRAIN.TRUNCATE_LENGTH,
        use_fast=True,
    )
    collate_fn = partial(
        custom_collate,
        tokenizer=tokenizer,
        truncate_length=cfg.TRAIN.TRUNCATE_LENGTH
    )
    
    # Determine metric for best model selection
    if cluster['val']:
        metric_for_best = f"eval_{cluster_name}_val_spearmanr"
    else:
        metric_for_best = f"eval_{cluster_name}_test_spearmanr"
    
    # Setup training
    run_name = f"{base_run_name}_{cluster_name}_fleet"
    output_dir = f"logs/{cfg.EXPERIMENT.PROJECT}/{base_run_name}_fleet/{cluster_name}"
    
    training_args = TrainingArguments(
        run_name=run_name,
        output_dir=output_dir,
        do_train=True,
        do_eval=True,
        max_steps=cfg.TRAIN.MAX_STEPS,
        per_device_train_batch_size=cfg.TRAIN.BATCH_SIZE,
        per_device_eval_batch_size=cfg.TEST.BATCH_SIZE,
        learning_rate=cfg.TRAIN.LR,
        lr_scheduler_type=cfg.TRAIN.SCHEDULER,
        warmup_ratio=0.2,
        weight_decay=0.01,
        optim="ademamix",
        eval_strategy="steps",
        eval_steps=cfg.TRAIN.EVAL_STEPS,
        logging_strategy="steps",
        logging_steps=cfg.TRAIN.EVAL_STEPS,
        save_strategy="steps",
        save_steps=cfg.TRAIN.EVAL_STEPS,
        save_total_limit=2,
        load_best_model_at_end=True,
        metric_for_best_model=metric_for_best,
        greater_is_better=True,
        bf16=True,
        seed=cfg.TRAIN.SEED,
        report_to=cfg.TRAIN.REPORT_TO,  # Log individual cluster runs to MLflow
        dataloader_num_workers=16 if device.type == "cuda" else 1,
        dataloader_prefetch_factor=8,
        dataloader_pin_memory=True,
        dataloader_persistent_workers=True,
        remove_unused_columns=False,
        max_grad_norm=cfg.TRAIN.GRAD_NORM,
        gradient_accumulation_steps=cfg.TRAIN.GRAD_ACCUM,
        lr_scheduler_kwargs=cfg.TRAIN.SCHEDULER_KWARGS,
        save_safetensors=True,
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_datasets,
        data_collator=collate_fn,
        compute_metrics=partial(compute_metrics, cfg=cfg),
    )
    
    # Check if already trained
    resume_checkpoint = None
    base_search_name = run_name.replace(f"_GPU{cfg.SYSTEM.GPU}", "") if f"_GPU{cfg.SYSTEM.GPU}" in run_name else run_name
    
    checkpoint_info = find_gpu_agnostic_checkpoint(
        base_search_name,
        cfg.EXPERIMENT.PROJECT,
        max_steps=cfg.TRAIN.MAX_STEPS
    )
    
    if checkpoint_info['checkpoint_path'] and checkpoint_info['is_complete']:
        print(f"  ✓ Already trained (step {checkpoint_info['current_step']}/{cfg.TRAIN.MAX_STEPS})")
        # Load existing training history
        trainer_state_path = os.path.join(checkpoint_info['checkpoint_path'], "trainer_state.json")
        with open(trainer_state_path, 'r') as f:
            trainer_state = json.load(f)
        history = trainer_state.get('log_history', [])
    elif checkpoint_info['checkpoint_path']:
        print(f"  → Resuming from step {checkpoint_info['current_step']}")
        resume_checkpoint = checkpoint_info['checkpoint_path']
        trainer.train(resume_from_checkpoint=resume_checkpoint)
        history = trainer.state.log_history
    else:
        print(f"  → Training from scratch")
        # Evaluate before training to get step 0 performance
        print(f"  → Running initial evaluation (step 0)...")
        step0_metrics = trainer.evaluate()
        step0_metrics['step'] = 0
        
        # Train the model
        trainer.train()
        history = trainer.state.log_history
        
        # Prepend step 0 metrics to history
        history.insert(0, step0_metrics)
    
    # Cleanup
    trainer.model.cpu()
    del trainer.model
    del trainer
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    return {
        'name': cluster_name,
        'has_val': cluster['val'] is not None,
        'history': history,
    }


def aggregate_metrics(trained_clusters):
    """
    Aggregate metrics across all clusters.
    
    For each training step, compute:
    - fleet_val_<metric> = avg(cluster_val_<metric>) across all clusters
    - fleet_test_<metric> = avg(cluster_test_<metric>) across all clusters
    - Also collect individual cluster metrics
    
    Returns tuple: (fleet_history, cluster_metrics)
    """
    print("\n=== Aggregating metrics ===")
    
    # Debug: Show what we're working with
    print(f"  Clusters to aggregate: {len(trained_clusters)}")
    for cluster_result in trained_clusters:
        print(f"    - {cluster_result['name']}: {len(cluster_result['history'])} log entries")
    
    # Group metrics by step for aggregation
    step_metrics = defaultdict(lambda: {'val': defaultdict(list), 'test': defaultdict(list)})
    
    # Also collect individual cluster metrics by step
    cluster_metrics_by_step = defaultdict(dict)
    
    for cluster_result in trained_clusters:
        cluster_name = cluster_result['name']
        has_val = cluster_result['has_val']
        
        for log_entry in cluster_result['history']:
            if 'step' not in log_entry:
                continue
            
            step = log_entry['step']
            
            # Collect val metrics for aggregation
            if has_val:
                for key, value in log_entry.items():
                    if key.startswith(f"eval_{cluster_name}_val_"):
                        metric_name = key.replace(f"eval_{cluster_name}_val_", "")
                        if isinstance(value, (int, float)):
                            step_metrics[step]['val'][metric_name].append(value)
                            # Store individual cluster metric
                            cluster_metrics_by_step[step][f"{cluster_name}_val_{metric_name}"] = value
            
            # Collect test metrics for aggregation
            for key, value in log_entry.items():
                if key.startswith(f"eval_{cluster_name}_test_"):
                    metric_name = key.replace(f"eval_{cluster_name}_test_", "")
                    if isinstance(value, (int, float)):
                        step_metrics[step]['test'][metric_name].append(value)
                        # Store individual cluster metric
                        cluster_metrics_by_step[step][f"{cluster_name}_test_{metric_name}"] = value
    
    # Compute averages
    fleet_history = []
    for step in sorted(step_metrics.keys()):
        step_data = {'step': step}
        
        # Average val metrics
        for metric_name, values in step_metrics[step]['val'].items():
            if values:
                step_data[f'fleet_val_{metric_name}'] = np.mean(values)
        
        # Average test metrics
        for metric_name, values in step_metrics[step]['test'].items():
            if values:
                step_data[f'fleet_test_{metric_name}'] = np.mean(values)
        
        fleet_history.append(step_data)
    
    # Convert cluster metrics to list format
    cluster_history = []
    for step in sorted(cluster_metrics_by_step.keys()):
        step_data = {'step': step}
        step_data.update(cluster_metrics_by_step[step])
        cluster_history.append(step_data)
    
    print(f"  Aggregated {len(fleet_history)} steps")
    print(f"  Collected {len(cluster_history)} steps of individual cluster metrics")
    
    # Debug: Show what metrics we collected
    if fleet_history:
        print(f"  Sample step {fleet_history[0]['step']} aggregated metrics: {[k for k in fleet_history[0].keys() if k != 'step'][:3]}...")
    if cluster_history:
        num_cluster_metrics = len([k for k in cluster_history[0].keys() if k != 'step'])
        print(f"  Sample step {cluster_history[0]['step']} has {num_cluster_metrics} individual cluster metrics")
    
    # Print final metrics
    if fleet_history:
        final = fleet_history[-1]
        print(f"\n  Final aggregated metrics (step {final['step']}):")
        for key, val in sorted(final.items()):
            if key != 'step':
                print(f"    {key}: {val:.4f}")
    else:
        print("  WARNING: No metrics were aggregated!")
    
    return fleet_history, cluster_history


def log_to_mlflow(cfg, fleet_history, cluster_history, fleet_run_name, num_clusters):
    """Log aggregated fleet metrics and individual cluster metrics to MLflow."""
    if cfg.TRAIN.REPORT_TO != "mlflow":
        print(f"\n  Skipping MLflow logging (report_to={cfg.TRAIN.REPORT_TO})")
        return
    
    if not fleet_history:
        print(f"\n  WARNING: No fleet history to log!")
        return
    
    print(f"\n=== Logging to MLflow ===")
    print(f"  Run name: {fleet_run_name}")
    print(f"  Project: {cfg.EXPERIMENT.PROJECT}")
    print(f"  Steps to log: {len(fleet_history)}")
    
    try:
        # Initialize MLflow (no resume for aggregated runs)
        print(f"  Initializing MLflow...")
        init_mlflow(cfg, fleet_run_name, resume_from_checkpoint=None)
        
        import mlflow
        print(f"  Logging parameters...")
        mlflow.log_param("fleet_num_clusters", num_clusters)
        mlflow.log_param("pretrained_checkpoint", cfg.FLEET.PRETRAINED_CHECKPOINT or "random_init")
        
        # Log aggregated fleet metrics
        print(f"  Logging aggregated fleet metrics for {len(fleet_history)} steps...")
        fleet_logged_count = 0
        for step_data in fleet_history:
            step = step_data['step']
            for metric_name, value in step_data.items():
                if metric_name != 'step':
                    mlflow.log_metric(metric_name, value, step=step)
                    fleet_logged_count += 1
        
        # Log individual cluster metrics
        print(f"  Logging individual cluster metrics for {len(cluster_history)} steps...")
        cluster_logged_count = 0
        for step_data in cluster_history:
            step = step_data['step']
            for metric_name, value in step_data.items():
                if metric_name != 'step':
                    mlflow.log_metric(metric_name, value, step=step)
                    cluster_logged_count += 1
        
        print(f"  ✓ Logged {fleet_logged_count} aggregated metrics across {len(fleet_history)} steps")
        print(f"  ✓ Logged {cluster_logged_count} individual cluster metrics across {len(cluster_history)} steps")
        mlflow.end_run()
        print(f"  ✓ MLflow run completed")
        
    except Exception as e:
        print(f"  ✗ MLflow logging failed: {e}")
        import traceback
        traceback.print_exc()


def main():
    parser = argparse.ArgumentParser(description="Fleet training")
    parser.add_argument("-c", type=str, required=True, help="Config file path")
    parser.add_argument("-gpu", type=int, default=None, help="GPU device")
    args = parser.parse_args()
    
    # Load config
    cfg = get_cfg_defaults()
    cfg.merge_from_file(args.c)
    cfg.freeze()
    
    # Validate fleet config
    if not hasattr(cfg, 'FLEET'):
        raise ValueError("Config must have FLEET section")
    if not cfg.FLEET.TRAIN_DIR or not cfg.FLEET.TEST_DIR:
        raise ValueError("Must specify FLEET.TRAIN_DIR and FLEET.TEST_DIR")
    
    # Setup GPU
    gpu_id = args.gpu if args.gpu is not None else cfg.SYSTEM.GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    cfg.defrost()
    cfg.SYSTEM.GPU = gpu_id
    cfg.freeze()
    
    # Set seeds
    np.random.seed(cfg.TRAIN.SEED)
    torch.manual_seed(cfg.TRAIN.SEED)
    
    # Generate base run name
    model_module = import_module(f"model.{cfg.MODEL.ARCHITECTURE}")
    temp_model = model_module.ddGRegressor(cfg=cfg)
    num_params = temp_model._calculate_num_parameters()
    del temp_model
    
    base_run_name = autoconstruct_run_name(cfg, num_parameters=num_params, gpu_id=gpu_id)
    
    print(f"\n{'='*60}")
    print(f"FLEET TRAINING")
    print(f"{'='*60}")
    print(f"Base run name: {base_run_name}")
    print(f"GPU: {gpu_id} ({device})")
    print(f"Pretrained: {cfg.FLEET.PRETRAINED_CHECKPOINT or 'random init'}")
    
    # 1. Find cluster files
    clusters = find_cluster_files(
        cfg.FLEET.TRAIN_DIR,
        cfg.FLEET.VAL_DIR,
        cfg.FLEET.TEST_DIR
    )
    
    if not clusters:
        raise ValueError("No clusters found!")
    
    # 2. Train all clusters
    print(f"\n{'='*60}")
    print(f"TRAINING {len(clusters)} MODELS")
    print(f"{'='*60}")
    
    trained_clusters = []
    for i, cluster in enumerate(clusters, 1):
        print(f"\n[{i}/{len(clusters)}] {cluster['name']}")
        result = train_cluster_model(cfg, cluster, base_run_name, device)
        trained_clusters.append(result)
    
    # 3. Aggregate metrics
    print(f"\n{'='*60}")
    print(f"AGGREGATING RESULTS")
    print(f"{'='*60}")
    
    fleet_history, cluster_history = aggregate_metrics(trained_clusters)
    
    # 4. Log to MLflow
    fleet_run_name = f"{base_run_name}_fleet_avg{len(clusters)}"
    log_to_mlflow(cfg, fleet_history, cluster_history, fleet_run_name, len(clusters))
    
    print(f"\n{'='*60}")
    print(f"✓ FLEET TRAINING COMPLETE")
    print(f"{'='*60}")
    print(f"Clusters trained: {len(clusters)}")
    print(f"Fleet run name: {fleet_run_name}")


if __name__ == "__main__":
    main()