multistats.py

import os
import re
import argparse
import tqdm
import wandb
from traceback import format_exc
import plotly.express as px
import pandas as pd
from concurrent.futures import ProcessPoolExecutor

import opendashboards.utils.utils as utils
import opendashboards.utils.aggregate as aggregate

from IPython.display import display

api= wandb.Api(timeout=60)
wandb.login(anonymous="allow")

def pull_wandb_runs(project='openvalidators', filters=None, min_steps=50, max_steps=100_000, ntop=10, netuid=None, summary_filters=None ):
    # TODO: speed this up by storing older runs

    all_runs = api.runs(project, filters=filters)
    print(f'Using {ntop}/{len(all_runs)} runs with more than {min_steps} events')
    pbar = tqdm.tqdm(all_runs)
    runs = []
    n_events = 0
    successful = 0
    for i, run in enumerate(pbar):

        summary = run.summary
        if summary_filters is not None and not summary_filters(summary):
            continue
        if netuid is not None and summary.get('netuid') != netuid:
            continue
        step = summary.get('_step',0)
        if step < min_steps or step > max_steps:
            # warnings.warn(f'Skipped run `{run.name}` because it contains {step} events (<{min_steps})')
            continue

        prog_msg = f'Loading data {i/len(all_runs)*100:.0f}% ({successful}/{len(all_runs)} runs, {n_events} events)'
        pbar.set_description(f'{prog_msg}... **fetching** `{run.name}`')

        duration = summary.get('_runtime')
        end_time = summary.get('_timestamp')
        # extract values for selected tags
        rules = {'hotkey': re.compile('^[0-9a-z]{48}$',re.IGNORECASE), 'version': re.compile('^\\d\.\\d+\.\\d+$'), 'spec_version': re.compile('\\d{4}$')}
        tags = {k: tag for k, rule in rules.items() for tag in run.tags if rule.match(tag)}
        # include bool flag for remaining tags
        tags.update({k: True for k in run.tags if k not in tags.keys() and k not in tags.values()})

        runs.append({
            'state': run.state,
            'num_steps': step,
            'num_completions': step*sum(len(v) for k, v in run.summary.items() if k.endswith('completions') and isinstance(v, list)),
            'entity': run.entity,
            'user': run.user.name,
            'username': run.user.username,
            'run_id': run.id,
            'run_name': run.name,
            'project': run.project,
            'run_url': run.url,
            'run_path': os.path.join(run.entity, run.project, run.id),
            'start_time': pd.to_datetime(end_time-duration, unit="s"),
            'end_time': pd.to_datetime(end_time, unit="s"),
            'duration': pd.to_timedelta(duration, unit="s").round('s'),
            'netuid': run.config.get('netuid'),
            **tags
        })
        n_events += step
        successful += 1
        if successful >= ntop:
            break

    return pd.DataFrame(runs).astype({'state': 'category', 'hotkey': 'category', 'version': 'category', 'spec_version': 'category'})

def plot_gantt(df_runs):
    fig = px.timeline(df_runs,
                x_start="start_time", x_end="end_time", y="username", color="state",
                title="Timeline of Runs",
                category_orders={'run_name': df_runs.run_name.unique()},#,'username': sorted(df_runs.username.unique())},
                hover_name="run_name",
                hover_data=['hotkey','user','username','run_id','num_steps','num_completions'],
                color_discrete_map={'running': 'green', 'finished': 'grey', 'killed':'blue', 'crashed':'orange', 'failed': 'red'},
                opacity=0.3,
                width=1200,
                height=800,
                template="plotly_white",
    )
    fig.update_yaxes(tickfont_size=8, title='')
    fig.show()


def clean_data(df):
    return df.dropna(subset=df.filter(regex='completions|rewards').columns, how='any').dropna(axis=1, how='all')

def explode_data(df):
    list_cols = utils.get_list_col_lengths(df)
    return utils.explode_data(df, list(list_cols.keys())).apply(pd.to_numeric, errors='ignore')


def load_data(run_id, run_path=None, load=True, save=False, explode=True):

    file_path = os.path.join('data/runs/',f'history-{run_id}.parquet')

    if load and os.path.exists(file_path):
        df = pd.read_parquet(file_path)
        # filter out events with missing step length
        df = df.loc[df.step_length.notna()]

        # detect list columns which as stored as strings
        list_cols = [c for c in df.columns if df[c].dtype == "object" and df[c].str.startswith("[").all()]
        # convert string representation of list to list
        # df[list_cols] = df[list_cols].apply(lambda x: eval(x, {'__builtins__': None}) if pd.notna(x) else x)
        try:
            df[list_cols] = df[list_cols].applymap(eval, na_action='ignore')
        except ValueError as e:
            print(f'Error loading {file_path!r} when converting columns {list_cols} to list: {e}')

    else:
        # Download the history from wandb and add metadata
        run = api.run(run_path)
        df = pd.DataFrame(list(run.scan_history()))

        # Remove rows with missing completions or rewards, which will be stuff related to weights
        df.dropna(subset=df.filter(regex='completions|rewards').columns, how='any', inplace=True)

        print(f'Downloaded {df.shape[0]} events from {run_path!r} with id {run_id!r}')

        # Clean and explode dataframe
        # overwrite object to free memory
        float_cols = df.filter(regex='reward').columns
        df = explode_data(clean_data(df)).astype({c: float for c in float_cols}).fillna({c: 0 for c in float_cols})

        if save:
            df.to_parquet(file_path, index=False)

    # Convert timestamp to datetime.
    df._timestamp = pd.to_datetime(df._timestamp, unit="s")
    return df.sort_values("_timestamp")


def calculate_stats(df_long, freq='H', save_path=None, ntop=3 ):

    df_long._timestamp = pd.to_datetime(df_long._timestamp)

    # if dataframe has columns such as followup_completions and answer_completions, convert to multiple rows
    if 'completions' not in df_long.columns:
        df_long.set_index(['_timestamp','run_id'], inplace=True)
        df_schema = pd.concat([
            df_long[['followup_completions','followup_rewards']].rename(columns={'followup_completions':'completions', 'followup_rewards':'rewards'}),
            df_long[['answer_completions','answer_rewards']].rename(columns={'answer_completions':'completions', 'answer_rewards':'rewards'})
        ])
        df_long = df_schema.reset_index()


    print(f'Calculating stats for dataframe with shape {df_long.shape}')

    # Approximate number of tokens in each completion
    df_long['completion_num_tokens'] = (df_long['completions'].str.split().str.len() / 0.75).round()


    g = df_long.groupby([pd.Grouper(key='_timestamp', axis=0, freq=freq), 'run_id'])

    # TODO: use named aggregations
    reward_aggs = ['sum','mean','std','median','max',aggregate.nonzero_rate, aggregate.nonzero_mean, aggregate.nonzero_std, aggregate.nonzero_median]
    aggs = {
        'completions': ['nunique','count', aggregate.diversity, aggregate.successful_diversity, aggregate.success_rate],
        'completion_num_tokens': ['mean', 'std', 'median', 'max'],
        **{k: reward_aggs for k in df_long.filter(regex='reward')}
    }

    # Calculate tokens per second
    if 'completion_times' in df_long.columns:
        df_long['tokens_per_sec'] = df_long['completion_num_tokens']/df_long['completion_times']
        aggs.update({
            'completion_times': ['mean','std','median','min','max'],
            'tokens_per_sec': ['mean','std','median','max'],
        })

    stats = g.agg(aggs)
    stats = stats.merge(g.apply(aggregate.top_stats, exclude='', ntop=ntop).reset_index(level=1,drop=True), left_index=True, right_index=True)
    # flatten multiindex columns
    stats.columns = ['_'.join(c) for c in stats.columns]
    stats = stats.reset_index()

    if save_path:
        stats.to_csv(save_path, index=False)

    return stats



def process(run, load=True, save=False, load_stats=True, freq='H', ntop=3):

    try:

        stats_path = f'data/aggs/stats-{run["run_id"]}.csv'
        if load_stats and os.path.exists(stats_path):
            print(f'Loaded stats file {stats_path!r}')
            return pd.read_csv(stats_path)

        # Load data and add extra columns from wandb run
        df_long = load_data(run_id=run['run_id'],
                    run_path=run['run_path'],
                    load=load,
                    save=save,
                    # save = (run['state'] != 'running') & run['end_time']
                    ).assign(**run.to_dict())
        assert isinstance(df_long, pd.DataFrame), f'Expected dataframe, but got {type(df_long)}'

        # Get and save stats
        return calculate_stats(df_long, freq=freq, save_path=stats_path, ntop=ntop)

    except Exception as e:
        print(f'Error processing run {run["run_id"]}: { format_exc(e) }')

def line_chart(df, col, title=None):
    title = title or col.replace('_',' ').title()
    fig = px.line(df.astype({'_timestamp':str}),
            x='_timestamp', y=col,
            line_group='run_id',
            title=f'{title} over Time',
            labels={'_timestamp':'', col: title, 'uids':'UID','value':'counts', 'variable':'Completions'},
            width=800, height=600,
            template='plotly_white',
        ).update_traces(opacity=0.2)

    fig.write_image(f'data/figures/{col}.png')
    fig.write_html(f'data/figures/{col}.html')
    return col


def parse_arguments():
    parser = argparse.ArgumentParser(description='Process wandb validator runs for a given netuid.')
    parser.add_argument('--load_runs',action='store_true', help='Load runs from file.')
    parser.add_argument('--repull_unfinished',action='store_true', help='Re-pull runs that were running when downloaded and saved.')
    parser.add_argument('--netuid', type=int, default=None, help='Network UID to use.')
    parser.add_argument('--ntop', type=int, default=1000, help='Number of runs to process.')
    parser.add_argument('--min_steps', type=int, default=100, help='Minimum number of steps to include.')
    parser.add_argument('--max_workers', type=int, default=32, help='Max workers to use.')
    parser.add_argument('--no_plot',action='store_true', help='Prevent plotting.')    
    parser.add_argument('--no_save',action='store_true', help='Prevent saving data to file.')
    parser.add_argument('--no_load',action='store_true', help='Prevent loading downloaded data from file.')
    parser.add_argument('--no_load_stats',action='store_true', help='Prevent loading stats data from file.')
    parser.add_argument('--freq', type=str, default='H', help='Frequency to aggregate data.')
    parser.add_argument('--completions_ntop', type=int, default=3, help='Number of top completions to include in stats.')

    return parser.parse_args()


if __name__ == '__main__':

    # TODO: flag to overwrite runs that were running when downloaded and saved: check if file date is older than run end time.

    args = parse_arguments()
    print(args)

    filters = None# {"tags": {"$in": [f'1.1.{i}' for i in range(10)]}}
    # filters={'tags': {'$in': ['5F4tQyWrhfGVcNhoqeiNsR6KjD4wMZ2kfhLj4oHYuyHbZAc3']}} # Is foundation validator
    if args.load_runs and os.path.exists('data/wandb.csv'):
        df_runs = pd.read_csv('data/wandb.csv')
        assert len(df_runs) >= args.ntop, f'Loaded {len(df_runs)} runs, but expected at least {args.ntop}'
        df_runs = df_runs.iloc[:args.ntop]
    else:
        df_runs = pull_wandb_runs(ntop=args.ntop,
                                min_steps=args.min_steps,
                                netuid=args.netuid,
                                filters=filters
                                )#summary_filters=lambda s: s.get('augment_prompt'))
        df_runs.to_csv('data/wandb.csv', index=False)


    os.makedirs('data/runs/', exist_ok=True)
    os.makedirs('data/aggs/', exist_ok=True)
    os.makedirs('data/figures/', exist_ok=True)

    display(df_runs)
    if not args.no_plot:
        plot_gantt(df_runs)

    with ProcessPoolExecutor(max_workers=min(args.max_workers, df_runs.shape[0])) as executor:
        futures = [executor.submit(
                            process,
                            run,
                            load=not args.no_load,
                            save=not args.no_save,
                            load_stats=not args.no_load_stats,
                            freq=args.freq,
                            ntop=args.completions_ntop
                    )
                   for _, run in df_runs.iterrows()
                   ]

        # Use tqdm to add a progress bar
        results = []
        with tqdm.tqdm(total=len(futures)) as pbar:
            for future in futures:
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    print(f'generated an exception: {format_exc(e)}')
                pbar.update(1)

    if not results:
        raise ValueError('No runs were successfully processed.')

   # Concatenate the results into a single dataframe
    df = pd.concat(results, ignore_index=True).sort_values(['_timestamp','run_id'], ignore_index=True)

    df.to_csv('data/processed.csv', index=False)
    print(f'Saved {df.shape[0]} rows to data/processed.csv')

    display(df)
    if not args.no_plot:
        
        plots = []

        cols = df.set_index(['run_id','_timestamp']).columns
        with ProcessPoolExecutor(max_workers=min(args.max_workers, len(cols))) as executor:
            futures = [executor.submit(line_chart, df, c) for c in cols]

            # Use tqdm to add a progress bar
            results = []
            with tqdm.tqdm(total=len(futures)) as pbar:
                for future in futures:
                    try:
                        result = future.result()
                        plots.append(result)
                    except Exception as e:
                        print(f'generated an exception: {format_exc(e)}')
                    pbar.update(1)
        
        print(f'Saved {len(plots)} plots to data/figures/')