Collect system stats in scheduler & Add scripts for experiments (#30)

plot/plot_stats.py

@@ -0,0 +1,52 @@
+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)