<jupyter_start><jupyter_text>CosmoFlow Benchmark Scaling Analysis on Summit<jupyter_code>import os
import pickle

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt

%matplotlib inline
def load_config(result_dir):
    config_file = os.path.join(result_dir, 'config.pkl')
    with open(config_file, 'rb') as f:
        return pickle.load(f)

def load_result(result_dir):
    history_file = os.path.join(result_dir, 'history.csv')
    return pd.read_csv(history_file)

def compute_mean_time(r):
    return r[r.epoch>0].time.mean()

def get_num_samples(config):
    dc = config['data']
    return dc['n_train'] + dc['n_valid']

plt.rcParams.update({'font.size': 14})
ranks = np.array([6, 12, 24, 48, 96, 192, 384, 768, 1020])
results_pattern = os.path.expandvars('real/log_%i')

configs, results = [], []
for r in ranks:
    result_dir = results_pattern % r
    configs.append(load_config(result_dir))
    results.append(load_result(result_dir).assign(ranks=r))

samples = np.array([get_num_samples(c) for (c,r) in zip(configs, ranks)])
times = np.array([compute_mean_time(r) for r in results])
throughputs = samples / times

scaling_scratch = pd.DataFrame(dict(ranks=ranks, samples=samples,
                                    times=times, throughputs=throughputs))

print(scaling_scratch)
ranks = np.array([6, 12, 24, 48, 96, 192, 384, 768, 1020])
results_pattern = os.path.expandvars('dummy/log_%i')

configs, results = [], []
for r in ranks:
    result_dir = results_pattern % r
    configs.append(load_config(result_dir))
    results.append(load_result(result_dir).assign(ranks=r))

samples = np.array([get_num_samples(c)*r for (c,r) in zip(configs, ranks)])
times = np.array([compute_mean_time(r) for r in results])
throughputs = samples / times

scaling_dummy = pd.DataFrame(dict(ranks=ranks, samples=samples,
                                  times=times, throughputs=throughputs))

print(scaling_dummy)
plt.figure(figsize=(9,8))
plt.plot(scaling_scratch['ranks'], scaling_scratch['throughputs'], 'o-', label='Real data')
plt.plot(scaling_dummy['ranks'], scaling_dummy['throughputs'], '^-', label='Dummy data')
plt.xlabel('Number of workers')
plt.ylabel('Training throughput [samples/s]')

plt.plot(scaling_scratch.ranks, scaling_scratch.ranks*scaling_scratch['throughputs'][0]/6, '--', label='Ideal')
plt.legend(loc=0);<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14646464646464646, 'fraction_numerical': 0.021885521885521887, 'mean_word_length': 5.048346055979644, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '233792', 'license': 'BSD-3-Clause-LBNL', 'n_tokens_mistral': 905, 'n_tokens_neox': 830, 'n_words': 165, 'path': 'notebooks/ScalingAnalysisSummit.ipynb', 'repo_name': 'lukasgd/cosmoflow-benchmark', 'size': '2376'}