[tab] /fast/home/mihir.mahajan/Projects/jasmine/slurm/utils/mihir/merge_wandb_runs.py (python)
#!/usr/bin/env python3
"""
merge_wandb_runs.py

Automatically merge multiple WandB CSV exports (from requeued runs) into one continuous CSV.
Detects the correct *_step column automatically, and adjusts step/time/wall_time so they remain
monotonic across requeues.

Example:
    python merge_wandb_runs.py \
        --inputs run1.csv run2.csv run3.csv \
        --output merged.csv
"""
import argparse
import pandas as pd
import os
import numpy as np


def find_step_column(df: pd.DataFrame) -> str:
    """Return the column name that ends with '_step' (case-insensitive)."""
    candidates = [c for c in df.columns if c.strip().lower().endswith(" _step")]
    if not candidates:
        raise ValueError("No column ending with '_step' found.")
    if len(candidates) > 1:
        print(f"⚠️ Multiple step-like columns found: {candidates}. Using first one: {candidates[0]}")
    return candidates[0]

def find_loss_column(df: pd.DataFrame) -> str:
    """Return the column name that ends with '_step' (case-insensitive)."""
    candidates = [c for c in df.columns if c.strip().lower().endswith("cross_entropy_loss")]
    if not candidates:
        raise ValueError("No column ending with '_cross_entropy_loss' found.")
    if len(candidates) > 1:
        print(f"⚠️ Multiple step-like columns found: {candidates}. Using first one: {candidates[0]}")
    return candidates[0]



def detect_and_merge(dfs):
    """
    Detect step overlaps and merge runs smoothly.
    Each subsequent run gets its step, time, and wall_time offset to continue from the previous.
    """
    merged = []
    step_col = [None, None, None]
    dfs_renamed = []
    for i, df in enumerate(dfs):
        dfs_renamed.append(df.rename(columns={find_loss_column(df): "cross_entropy_loss"}))
    dfs = dfs_renamed

    for i, df in enumerate(dfs):
        step_col[i] = find_step_column(df)  # assume same column name across all runs
        df = df.copy()

        # Ensure numeric and sorted
        df[step_col[i]] = pd.to_numeric(df[step_col[i]], errors="coerce")
        df = df.dropna(subset=[step_col[i]]).sort_values(step_col[i]).reset_index(drop=True)

        # Apply offsets for all but the first run
        if i > 0:
            prev = merged[-1]
            prev_last_step = prev[step_col[i-1]].max()
            prev_last_time = prev["time"].max() if "time" in prev else 0.0
            prev_last_wall = prev["wall_time"].max() if "wall_time" in prev else 0.0
            prev_last_rel_time = prev["Relative Time (Process)"].max() if "Relative Time (Process)" in prev else 0.0

            # Detect overlap: if this run’s first step <= prev_last_step, assume it’s resumed
            first_step = df[step_col[i]].iloc[0]
            if first_step <= prev_last_step:
                step_offset = prev_last_step - first_step
            else:
                step_offset = prev_last_step

            # Apply offsets
            df[step_col[i]] = df[step_col[i]] + step_offset
            if "time" in df.columns:
                df["time"] = df["time"] + prev_last_time
            if "wall_time" in df.columns:
                df["wall_time"] = df["wall_time"] + prev_last_wall
            if "Relative Time (Process)" in df.columns:
                df["Relative Time (Process)"] = df["Relative Time (Process)"] + prev_last_rel_time

        merged.append(df)

    merged_df = pd.concat(merged)
    merged_df = merged_df.sort_values("Relative Time (Process)").reset_index(drop=True)
    # Replace "Relative Time (Process)" with a uniform sweep from 0 to 144000
    merged_df["Relative Time (Process)"] = np.linspace(
        0, 144000, num=len(merged_df), endpoint=True
    )
    return merged_df, step_col[i]


def main():
    parser = argparse.ArgumentParser(description="Merge multiple WandB CSVs (from requeued runs) into one continuous CSV.")
    parser.add_argument("--inputs", "-i", nargs="+", required=True,
                        help="List of input CSV files in chronological order.")
    parser.add_argument("--output", "-o", required=True, help="Path to merged output CSV file.")
    args = parser.parse_args()

    dfs = []
    for p in args.inputs:
        if not os.path.exists(p):
            raise FileNotFoundError(f"Input CSV not found: {p}")
        df = pd.read_csv(p, comment="#", low_memory=False)
        if df.empty:
            raise ValueError(f"Input CSV {p} is empty.")
        dfs.append(df)

    merged_df, step_col = detect_and_merge(dfs)

    out_dir = os.path.dirname(args.output)
    if out_dir and not os.path.exists(out_dir):
        os.makedirs(out_dir, exist_ok=True)

    merged_df.to_csv(args.output, index=False)
    print(f"Merged {len(args.inputs)} runs into: {args.output}")
    print(f"Detected step column: {step_col}")
    print(f"Final step: {merged_df[step_col].max():.0f}")
    if "time" in merged_df.columns:
        print(f"Final time: {merged_df['time'].max():.2f} sec")
    if "wall_time" in merged_df.columns:
        print(f"Final wall_time: {merged_df['wall_time'].max():.2f}")
    if "Relative Time (Process)" in merged_df.columns:
        print(f"Final Relative Time (Process): {merged_df['Relative Time (Process)'].max():.2f}")


if __name__ == "__main__":
    main()