From 84c4e81b35970b8e687126cf38156dbfc70d8770 Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Date: Wed, 9 Oct 2024 21:32:46 +0000
Subject: [PATCH] rename save traces

---
 examples/offline_profile.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/examples/offline_profile.py b/examples/offline_profile.py
index c9fceb03848c1..1d415b82cddb6 100644
--- a/examples/offline_profile.py
+++ b/examples/offline_profile.py
@@ -24,7 +24,7 @@ class ProfileContext:
     prompt_len: int
     output_len: int
     batch_size: int
-    save_traces_folder: Optional[str]
+    save_chrome_traces_folder: Optional[str]
 
 
 def get_dtype(dtype: str):
@@ -194,15 +194,15 @@ def abort_requests():
             json.dump(json_dict, f, indent=2)
         pass
 
-    if context.save_traces_folder is not None:
-        os.makedirs(context.save_traces_folder, exist_ok=True)
-        prefill_prof.profiler.export_chrome_trace(context.save_traces_folder +
-                                                  "/prefill.json")
+    if context.save_chrome_traces_folder is not None:
+        os.makedirs(context.save_chrome_traces_folder, exist_ok=True)
+        prefill_prof.profiler.export_chrome_trace(
+            context.save_chrome_traces_folder + "/prefill.json")
         for idx, decode_prof in enumerate(decode_profs):
             decode_prof.profiler.export_chrome_trace(
-                context.save_traces_folder + f"/decode_{idx + 1}.json")
+                context.save_chrome_traces_folder + f"/decode_{idx + 1}.json")
         print("Traces saved as prefill.json and decode_1.json, etc."
-              f" in folder {context.save_traces_folder}")
+              f" in folder {context.save_chrome_traces_folder}")
 
 
 if __name__ == "__main__":
@@ -245,7 +245,7 @@ def abort_requests():
         type=str,
         default=None,
         help="Export the results as a json file. This should be the filename")
-    parser.add_argument("--save-traces-folder",
+    parser.add_argument("--save-chrome-traces-folder",
                         type=str,
                         help="Save chrome traces for the prefill and decode "
                         "will save traces as prefill.json and decode_1.json, "