From 2ab2b4035cb728fbd57601399693ee9d1dc3d2e1 Mon Sep 17 00:00:00 2001
From: Wilson Wang <wilson.wang@bytedance.com>
Date: Tue, 25 May 2021 15:43:46 -0700
Subject: [PATCH] tools: no hardcoded repeat times in rw-benchmark scripts

---
 tools/rw-heatmaps/plot_data.py    | 13 +++++++------
 tools/rw-heatmaps/rw-benchmark.sh |  2 +-
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/tools/rw-heatmaps/plot_data.py b/tools/rw-heatmaps/plot_data.py
index c639c1ab2..0e3e3ec63 100755
--- a/tools/rw-heatmaps/plot_data.py
+++ b/tools/rw-heatmaps/plot_data.py
@@ -43,7 +43,8 @@ def load_data_files(*args):
     try:
         for df in df_list:
             new_df = df[['ratio', 'conn_size', 'value_size']].copy()
-            tmp = [df[x].str.split(':') for x in ['1', '2', '3', '4', '5']]
+            cols = [x for x in df.columns if x.find('iter') != -1]
+            tmp = [df[x].str.split(':') for x in cols]
 
             read_df = [x.apply(lambda x: float(x[0])) for x in tmp]
             read_avg = sum(read_df)/len(read_df)
@@ -73,7 +74,7 @@ def plot_data(title, *args):
             count += 1
             plt.subplot(4, 2, count)
             plt.tripcolor(df['conn_size'], df['value_size'], df['read'] + df['write'])
-            plt.title('R/W Ratio {:.2f}'.format(val))
+            plt.title('R/W Ratio {:.4f}'.format(val))
             plt.yscale('log', base=2)
             plt.ylabel('Value Size')
             plt.xscale('log', base=2)
@@ -99,13 +100,13 @@ def plot_data(title, *args):
                     cmap_name = 'viridis'
                 plt.tripcolor(df['conn_size'], df['value_size'], df['read'] + df['write'], cmap=plt.get_cmap(cmap_name))
                 if count2 == 1:
-                    plt.title('{}\nR/W Ratio {:.2f}'.format(os.path.basename(params.input_file_a), val))
+                    plt.title('{}\nR/W Ratio {:.4f}'.format(os.path.basename(params.input_file_a), val))
                 elif count2 == 2:
-                    plt.title('{}\nR/W Ratio {:.2f}'.format(os.path.basename(params.input_file_b), val))
+                    plt.title('{}\nR/W Ratio {:.4f}'.format(os.path.basename(params.input_file_b), val))
                 elif count2 == 3:
-                    plt.title('Delta\nR/W Ratio {:.2f}'.format(val))
+                    plt.title('Delta\nR/W Ratio {:.4f}'.format(val))
                 else:
-                    plt.title('R/W Ratio {:.2f}'.format(val))
+                    plt.title('R/W Ratio {:.4f}'.format(val))
                 plt.yscale('log', base=2)
                 plt.ylabel('Value Size')
                 plt.xscale('log', base=2)
diff --git a/tools/rw-heatmaps/rw-benchmark.sh b/tools/rw-heatmaps/rw-benchmark.sh
index 65764d5c6..d11492e97 100755
--- a/tools/rw-heatmaps/rw-benchmark.sh
+++ b/tools/rw-heatmaps/rw-benchmark.sh
@@ -48,7 +48,7 @@ function check_prerequisite() {
     exit 1
   fi
   cat >"${OUTPUT_FILE}" <<EOF
-ratio, conn_size, value_size, 1, 2, 3, 4, 5
+ratio,conn_size,value_size$(for i in $(seq 1 ${REPEAT_COUNT});do echo -n ",iter$i"; done)
 EOF
 }