src/utils.py

import torch
import json
import time
from datetime import datetime


import os
import torch.nn as nn

import numpy as np
import random
import transformers

import platform
from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl


class ExperimentMonitorCallback(TrainerCallback):
    """
    Callback to monitor training performance and log system stats to a JSON file.
    It captures:
    1. Experiment Metadata (GPU info, Batch size, Learning rate, etc.)
    2. Runtime Metrics (Avg time/step, Throughput)
    3. Memory Metrics (Allocated, Reserved, and Peak usage)
    """

    def __init__(self, log_file_path: str, run_name: str = "experiment", log_interval: int = 100):
        # English comments as requested
        self.log_file_path = log_file_path
        self.run_name = run_name
        self.log_interval = log_interval
        
        # Timing variables
        self.start_time = None
        self.last_log_time = None
        
        # Data container to be saved
        self.log_data = {
            "metadata": {},
            "metrics": []
        }

    def _get_gpu_info(self):
        # Helper to get GPU details if available
        if torch.cuda.is_available():
            return {
                "name": torch.cuda.get_device_name(0),
                "count": torch.cuda.device_count(),
                "capability": torch.cuda.get_device_capability(0)
            }
        return "CPU_ONLY"

    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        # Initialize timing
        self.start_time = time.perf_counter()
        self.last_log_time = self.start_time
        
        # Reset peak memory stats to ensure we capture peaks specific to this run
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()

        # Capture experiment metadata
        self.log_data["metadata"] = {
            "run_name": self.run_name,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "python_version": platform.python_version(),
            "pytorch_version": torch.__version__,
            "gpu_info": self._get_gpu_info(),
            "configuration": {
                "batch_size_per_device": args.per_device_train_batch_size,
                "learning_rate": args.learning_rate,
                "max_steps": args.max_steps,
                "num_train_epochs": args.num_train_epochs,
                "fp16": args.fp16,
                "bf16": args.bf16,
                "optim": args.optim,
            }
        }
        
        # Create/Overwrite the file with initial metadata
        self._save_log()
        # print(f"[{self.run_name}] Experiment started. Logging to {self.log_file_path}")

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        current_step = state.global_step
        
        # Perform logging only at specified intervals
        if current_step > 0 and current_step % self.log_interval == 0:
            current_time = time.perf_counter()
            
            # Calculate time elapsed since the last log
            elapsed_since_last = current_time - self.last_log_time
            avg_time_per_step = elapsed_since_last / self.log_interval
            
            # Memory Statistics (in GB)
            mem_stats = {}
            if torch.cuda.is_available():
                # Current usage
                mem_stats["allocated_gb"] = torch.cuda.memory_allocated() / 1024**3
                mem_stats["reserved_gb"] = torch.cuda.memory_reserved() / 1024**3
                # Peak usage since start (Long-term peak)
                mem_stats["peak_allocated_gb"] = torch.cuda.max_memory_allocated() / 1024**3
            
            # Construct metric entry
            metric_entry = {
                "step": current_step,
                "epoch": state.epoch,
                "timestamp": datetime.now().isoformat(),
                "performance": {
                    "avg_time_per_step_s": round(avg_time_per_step, 4),
                    "steps_per_second": round(1.0 / avg_time_per_step, 2)
                },
                "memory": mem_stats
            }

            # Append to internal list and save to file
            self.log_data["metrics"].append(metric_entry)
            self._save_log()
            
            # Update last log time
            self.last_log_time = current_time
            
            # Optional: Print a brief summary to console
            print(f" -> Step {current_step}: {avg_time_per_step*1000:.1f}s/step |"\
                f"Peak Mem: {mem_stats.get('peak_allocated_gb', 0):.2f} GB |"\
                f"Reserved: {mem_stats.get('reserved_gb', 0):.2f} GB")

    def _save_log(self):
        # Dump the entire data structure to JSON
        # For very long training runs, appending to a JSONL (lines) file might be more efficient,
        # but standard JSON is easier to read for analysis.
        try:
            with open(self.log_file_path, 'w', encoding='utf-8') as f:
                json.dump(self.log_data, f, indent=4)
        except Exception as e:
            print(f"Error saving experiment log: {e}")

def debug_masking_visualizer(processed_batch, tokenizer):
    """
    Visualizes the alignment between input_ids and labels to verify masking.
    """
    input_ids = processed_batch['input_ids'][0] # Take the first sample in batch
    labels = processed_batch['labels'][0]
    
    print("\n" + "="*80)
    print(f"{'TOKEN (Decoded)':<30} | {'INPUT ID':<10} | {'LABEL ID':<10} | {'STATUS'}")
    print("="*80)

    for idx, lbl in zip(input_ids, labels):
        # Decode individual token for visualization
        # Replace newlines so table doesn't break
        token_text = tokenizer.decode([idx]).replace("\n", "\\n")
        
        # Check masking status
        if lbl == -100:
            status = "❌ MASKED (Instruction)"
            label_display = "IGNORE"
        else:
            status = "✅ TRAIN (Response)"
            label_display = str(lbl)
            
        print(f"{token_text:<30} | {idx:<10} | {label_display:<10} | {status}")
    
    print("="*80 + "\n")


def trainable_parameters_to_file(model: nn.Module, save_dir: str):
    """
    Calculates model parameters and saves a detailed report of 
    trainable matrices to a specific directory.
    """
    trainable_params = 0
    all_param = 0
    trainable_layers = []
    
    # Track the maximum length of layer names for alignment
    max_name_len = 20 # Minimum width
    
    for name, param in model.named_parameters():
        num_params = param.numel()
        all_param += num_params
        if param.requires_grad:
            trainable_params += num_params
            trainable_layers.append({
                "name": name,
                "shape": str(list(param.shape)),
                "count": num_params
            })
            # Update max length if current name is longer
            if len(name) > max_name_len:
                max_name_len = len(name)
    
    trainable_pct = 100 * trainable_params / all_param if all_param > 0 else 0
    summary_text = (
        f"Total Parameters: {all_param:,}\n"
        f"Trainable Parameters: {trainable_params:,}\n"
        f"Trainable Percentage: {trainable_pct:.4f}%\n"
    )
    
    # print("-" * 30)
    # print(summary_text.strip())
    # print("-" * 30)
    
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        
    file_path = os.path.join(save_dir, "model_parameters_report.txt")
    
    # Add some padding to max_name_len
    name_col_width = max_name_len + 4
    
    with open(file_path, "w") as f:
        f.write("=== GLOBAL STATISTICS ===\n")
        f.write(summary_text)
        f.write("\n" + "=" * (name_col_width + 40) + "\n")
        f.write("=== DETAILED TRAINABLE MATRICES LIST ===\n")
        
        # Dynamic alignment using calculated width
        header = f"{'Layer Name':<{name_col_width}} | {'Shape':<25} | {'Count':<15}\n"
        f.write(header)
        f.write("-" * len(header) + "\n")
        
        for layer in trainable_layers:
            f.write(
                f"{layer['name'] :<{name_col_width}} | "
                f"{layer['shape'] :<25} | "
                f"{layer['count'] :,}\n"
            )

    # print(f"Detailed report saved at: {file_path}")
    
def set_seed_all(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    transformers.set_seed(seed)
    # torch.backends.cudnn.deterministic = True
    # torch.backends.cudnn.benchmark = False