53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
import os
|
|
import pickle
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
STAT_NAMES = [
|
|
'input_lens',
|
|
'num_running',
|
|
'num_waiting',
|
|
'num_preemption',
|
|
'gpu_cache_usage',
|
|
'cpu_cache_usage',
|
|
'num_swapped',
|
|
'swap_in_lens',
|
|
'swap_out_lens',
|
|
]
|
|
|
|
|
|
def plot_stats(output_dir: str):
|
|
# Get stats.
|
|
with open(os.path.join(output_dir, 'stats.pkl'), 'rb') as f:
|
|
stats = pickle.load(f)
|
|
timestamps = stats['timestamps']
|
|
|
|
# Draw one figure for each stat.
|
|
num_stats = len(STAT_NAMES)
|
|
COLORS = ['b', 'g', 'r', 'c', 'm', 'y', 'k', 'orange', 'purple', 'pink', 'brown', 'gray']
|
|
fig, axs = plt.subplots(num_stats, 1, figsize=(10, 2 * num_stats))
|
|
for i, stat in enumerate(STAT_NAMES):
|
|
data = stats[stat]
|
|
if stat in ['gpu_cache_usage', 'cpu_cache_usage']:
|
|
data = [x * 100 for x in data]
|
|
stat = stat + ' (%)'
|
|
axs[i].plot(timestamps, data, color=COLORS[i % len(COLORS)])
|
|
axs[i].set_ylabel(stat.replace('_', ' '), fontdict={'fontsize': 12})
|
|
axs[i].set_ylim(bottom=0)
|
|
|
|
plt.xlabel('Time (s)')
|
|
plt.tight_layout()
|
|
fig_path = os.path.join(output_dir, 'stats.png')
|
|
plt.savefig(fig_path)
|
|
print(f'Saved stats to {fig_path}')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('output_dir', type=str, help='Output directory.')
|
|
args = parser.parse_args()
|
|
|
|
plot_stats(args.output_dir)
|