From dd477b40de36ec8e5ba4416e552e42b69da250b2 Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Date: Tue, 10 Sep 2024 15:16:05 +0000
Subject: [PATCH] refactor

---
 examples/offline_profile.py                               | 8 ++++----
 vllm/profiler/__init__.py                                 | 4 ++--
 vllm/profiler/{nm_profile.py => layerwise_profile.py}     | 6 +++---
 .../profiler/print_layerwise_table.py                     | 2 +-
 .../profiler/visualize_layerwise_profile.py               | 4 ++++
 5 files changed, 14 insertions(+), 10 deletions(-)
 rename vllm/profiler/{nm_profile.py => layerwise_profile.py} (98%)
 rename neuralmagic/tools/profiler/print_table.py => vllm/profiler/print_layerwise_table.py (96%)
 rename neuralmagic/tools/profiler/visualize_trace.py => vllm/profiler/visualize_layerwise_profile.py (99%)

diff --git a/examples/offline_profile.py b/examples/offline_profile.py
index c5c2346c908ad..fb44c69a49a81 100644
--- a/examples/offline_profile.py
+++ b/examples/offline_profile.py
@@ -8,7 +8,7 @@
 import torch
 
 from vllm import LLM, SamplingParams
-from vllm.profiler import nm_profile
+from vllm.profiler import layerwise_profile
 
 BATCH_SIZE_DEFAULT = 1
 PROMPT_LEN_DEFAULT = 256
@@ -78,7 +78,7 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
               f"larger than max_num_batched_tokens ({max_num_batched_tokens}) "
               f"and therefore cannot be run in a single profile step, please "
               f"choose a smaller batch size or prompt length, or increase "
-              f"--max_num_batched_tokens")
+              f"--max-num-batched-tokens")
         sys.exit(-1)
     if batch_size >= max_num_seqs:
         print(
@@ -106,12 +106,12 @@ def run_profile(context: ProfileContext, csv_output: Optional[str],
             inputs={'prompt_token_ids': prompt_token_ids},
             params=sampling_params)
 
-    with nm_profile() as prefill_prof:
+    with layerwise_profile() as prefill_prof:
         llm.llm_engine.step()  # First step is prefill
 
     decode_results_list = []
     for x in range(args.output_len - 1):
-        with nm_profile() as decode_prof:
+        with layerwise_profile() as decode_prof:
             llm.llm_engine.step()
         decode_results_list.append(decode_prof.results)
 
diff --git a/vllm/profiler/__init__.py b/vllm/profiler/__init__.py
index 93ec4a800e600..3e25f5cc283f2 100644
--- a/vllm/profiler/__init__.py
+++ b/vllm/profiler/__init__.py
@@ -1,5 +1,5 @@
-from .nm_profile import nm_profile
+from .layerwise_profile import layerwise_profile
 
 __all__ = [
-    "nm_profile",
+    "layerwise_profile",
 ]
diff --git a/vllm/profiler/nm_profile.py b/vllm/profiler/layerwise_profile.py
similarity index 98%
rename from vllm/profiler/nm_profile.py
rename to vllm/profiler/layerwise_profile.py
index 5901362ecd825..e074cc3f0d2ad 100644
--- a/vllm/profiler/nm_profile.py
+++ b/vllm/profiler/layerwise_profile.py
@@ -63,7 +63,7 @@ class _StatsTreeNode:
 
 
 @dataclass
-class NMProfileResults(profile):
+class LayerwiseProfileResults(profile):
     _kineto_results: _ProfilerResult
     _kineto_event_correlation_map: Dict[int,
                                         List[_KinetoEvent]] = field(init=False)
@@ -331,7 +331,7 @@ def df_traversal(node: _StatsTreeNode, curr_json_list: List[Dict]):
         return root_dicts
 
 
-class nm_profile(profile):
+class layerwise_profile(profile):
 
     def __init__(self):
         super().__init__(
@@ -346,4 +346,4 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         super().__exit__(exc_type, exc_val, exc_tb)
-        self.results = NMProfileResults(self.profiler.kineto_results)
+        self.results = LayerwiseProfileResults(self.profiler.kineto_results)
diff --git a/neuralmagic/tools/profiler/print_table.py b/vllm/profiler/print_layerwise_table.py
similarity index 96%
rename from neuralmagic/tools/profiler/print_table.py
rename to vllm/profiler/print_layerwise_table.py
index 9081583a9f95d..91bfec747f4a9 100644
--- a/neuralmagic/tools/profiler/print_table.py
+++ b/vllm/profiler/print_layerwise_table.py
@@ -2,7 +2,7 @@
 import json
 from typing import Dict
 
-from vllm.profiler.nm_profile import ModelStatsEntry, SummaryStatsEntry
+from vllm.profiler.layerwise_profile import ModelStatsEntry, SummaryStatsEntry
 from vllm.profiler.utils import TablePrinter, indent_string
 
 
diff --git a/neuralmagic/tools/profiler/visualize_trace.py b/vllm/profiler/visualize_layerwise_profile.py
similarity index 99%
rename from neuralmagic/tools/profiler/visualize_trace.py
rename to vllm/profiler/visualize_layerwise_profile.py
index c1284226fec69..ea27152670d36 100644
--- a/neuralmagic/tools/profiler/visualize_trace.py
+++ b/vllm/profiler/visualize_layerwise_profile.py
@@ -1,3 +1,4 @@
+import os
 import argparse
 import copy
 import json
@@ -426,6 +427,9 @@ def make_plot_title_suffix(profile_json: dict) -> str:
 
     output_directory = args.output_directory if args.output_directory else Path(
         args.json_trace).parent
+    
+    if not os.path.exists(output_directory):
+        os.makedirs(output_directory)
 
     main(Path(args.json_trace), output_directory, depth, args.plot_metric,
          make_names_unique, args.top_k, args.fold_json_node)