opendashboards/utils/utils.py

# The MIT License (MIT)
# Copyright © 2021 Yuma Rao

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
# documentation files (the “Software”), to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of
# the Software.

# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
# THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import os
import re
import tqdm
import wandb
import pandas as pd
import streamlit as st
from traceback import format_exc
from pandas.api.types import is_list_like

from typing import List, Dict, Any, Union

api = wandb.Api(api_key=st.secrets['WANDB_API_KEY'])
wandb.login(anonymous="allow")

def pull_wandb_runs(project: str, filters=None, min_steps=50, ntop=10, summary_filters=None ):
    all_runs = get_runs(project, 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
        step = summary.get('_step',0)
        if step < min_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'),
            **tags
        })
        n_events += step
        successful += 1
        if successful >= ntop:
            break

    cat_cols = ['state', 'hotkey', 'version', 'spec_version']
    return pd.DataFrame(runs).astype({k: 'category' for k in cat_cols if k in runs[0]})



def get_runs(project: str, filters: Dict[str, Any] = None, return_paths: bool = False) -> List:
    """Download runs from wandb.

    Args:
        project (str): Name of the project. 
        filters (Dict[str, Any], optional): Optional run filters for wandb api. Defaults to None.
        return_paths (bool, optional): Return only run paths. Defaults to False.

    Returns:
        List[wandb.apis.public.Run]: List of runs or run paths (List[str]).
    """

    runs = api.runs(project, filters=filters)
    if return_paths:
        return [os.path.join(run.entity, run.project, run.id) for run in runs]
    else:
        return runs


def download_data(run_path: Union[str, List] = None, timeout: float = 600, api_key: str = None) -> pd.DataFrame:
    """Download data from wandb.

    Args:
        run_path (Union[str, List], optional): Path to run or list of paths. Defaults to None.
        timeout (float, optional): Timeout for wandb api. Defaults to 600.

    Returns:
        pd.DataFrame: Dataframe of event log.
    """
    api = wandb.Api(api_key=api_key, timeout=timeout)
    wandb.login(anonymous="allow")

    if isinstance(run_path, str):
        run_path = [run_path]

    frames = []
    total_events = 0
    pbar = tqdm.tqdm(sorted(run_path), desc="Loading history from wandb", total=len(run_path), unit="run")
    for path in pbar:
        run = api.run(path)

        frame = pd.DataFrame(list(run.scan_history()))
        frames.append(frame)
        total_events += len(frame)

        pbar.set_postfix({"total_events": total_events})

    df = pd.concat(frames)

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

    return df


def read_data(path: str, nrows: int = None):
    """Load data from csv."""
    df = pd.read_csv(path, nrows=nrows)
    # filter out events with missing step length
    # df = df.loc[df.step_length.notna()]

    # detect list columns which as stored as strings
    def is_list_col(x):
        return isinstance(x, str) and x[0]=='[' and x[-1]==']'
    list_cols = [c for c in df.columns if df[c].dtype == "object" and df[c].apply(is_list_col).all()]
    # convert string representation of list to list
    df[list_cols] = df[list_cols].applymap(eval, na_action='ignore')

    return df

def load_data(selected_runs, load=True, save=False, explode=True, datadir='data/'):

    frames = []
    n_events = 0
    successful = 0
    if not os.path.exists(datadir):
        os.makedirs(datadir)

    st.write(selected_runs)
    pbar = tqdm.tqdm(selected_runs.index, desc="Loading runs", total=len(selected_runs), unit="run")
    for i, idx in enumerate(pbar):
        run = selected_runs.loc[idx]
        prog_msg = f'Loading data {i/len(selected_runs)*100:.0f}% ({successful}/{len(selected_runs)} runs, {n_events} events)'

        file_path = os.path.join(datadir,f'history-{run.run_id}.csv')

        if (load is True and os.path.exists(file_path)) or (callable(load) and load(run.to_dict())):
            pbar.set_description(f'{prog_msg}... **reading** `{file_path}`')
            try:
                df = read_data(file_path)
            except Exception as e:
                print(f'Failed to load history from `{file_path}`: {format_exc(e)}')
                continue
        else:
            pbar.set_description(f'{prog_msg}... **downloading** `{run.run_path}`')
            try:
                # Download the history from wandb and add metadata
                df = download_data(run.run_path).assign(**run.to_dict())
                if explode:
                    df = explode_data(df)

                print(f'Downloaded {df.shape[0]} events from `{run.run_path}`. Columns: {df.columns}')

                if save is True or (callable(save) and save(run.to_dict())):
                    df.to_csv(file_path, index=False)
                    print(f'Saved {df.shape[0]} events to `{file_path}`')

            except Exception as e:
                print(f'Failed to download history for `{run.run_path}`: {e}')
                continue

        frames.append(df)
        n_events += df.shape[0]
        successful += 1

    # Remove rows which contain chain weights as it messes up schema
    return pd.concat(frames)


def explode_data(df: pd.DataFrame, list_cols: List[str] = None, list_len: int = None) -> pd.DataFrame:
    """Explode list columns in dataframe so that each element in the list is a separate row.

    Args:
        df (pd.DataFrame): Dataframe of event log.
        list_cols (List[str], optional): List of columns to explode. Defaults to None.
        list_len (int, optional): Length of list. Defaults to None.

    Returns:
        pd.DataFrame: Dataframe with exploded list columns.
    """
    if list_cols is None:
        list_cols = [c for c in df.columns if df[c].apply(is_list_like).all()]
        print(f"Exploding {len(list_cols)}) list columns with {list_len} elements: {list_cols}")
    if list_len:
        list_cols = [c for c in list_cols if df[c].apply(len).unique()[0] == list_len]
        print(f"Exploding {len(list_cols)}) list columns with {list_len} elements: {list_cols}")

    return df.explode(column=list_cols)


def get_list_col_lengths(df: pd.DataFrame) -> Dict[str, int]:
    """Helper function to get the length of list columns."""
    list_col_lengths = {c: sorted(df[c].apply(len).unique()) for c in df.columns if df[c].apply(is_list_like).all()}
    varying_lengths = {c: v for c, v in list_col_lengths.items() if len(v) > 1}

    if len(varying_lengths) > 0:
        print(f"The following columns have varying lengths: {varying_lengths}")

    return {c: v[0] for c, v in list_col_lengths.items() if v}