Add graph compilation tracking to high level profiler (#50)

Signed-off-by: Konrad Zawora <kzawora@habana.ai>

Author: kzawora-intel

vllm_gaudi/extension/features.py

@@ -97,5 +97,7 @@ def get_features():
               1.,
               env_var='VLLM_UNIFIED_ATTENTION_SHARED_CACHE_RATIO',
               env_var_type=float),
+        Value('high_level_profiler_enabled', False, env_var='VLLM_PROFILER_ENABLED', env_var_type=boolean),
+        Value('track_graph_compilation', False, env_var='PT_HPU_METRICS_GC_DETAILS', env_var_type=boolean),
     ]
     return split_values_and_flags(features)

vllm_gaudi/extension/profiler.py

@@ -2,6 +2,7 @@
 # Copyright (C) 2024 Habana Labs, Ltd. an Intel Company
 ###############################################################################
 
+import contextlib
 import gc
 import gzip
 import json
@@ -19,6 +20,7 @@
 
 from vllm_gaudi.extension.utils import is_fake_hpu
 from .logger import logger
+from vllm_gaudi.extension.runtime import get_config
 
 
 class FileWriter(threading.Thread):
@@ -142,8 +144,7 @@ class HabanaHighLevelProfiler:
     event_cache: List[Any] = []
 
     def __init__(self, vllm_instance_id=None):
-        self.enabled = os.getenv('VLLM_PROFILER_ENABLED', 'false').lower() == 'true' and int(os.getenv('RANK',
-                                                                                                       '0')) == 0
+        self.enabled = get_config().high_level_profiler_enabled and int(os.getenv('RANK', '0')) == 0
         self.pid = os.getpid()
         if self.enabled:
             self.vllm_instance_id = vllm_instance_id if vllm_instance_id is not None \
@@ -158,6 +159,8 @@ def __init__(self, vllm_instance_id=None):
             file_writer.start()
         if os.getenv('VLLM_PROFILER_ENABLED') == 'full':
             self.enabled = True  # don't save separate high-level traces
+        self.gc_track_recompiles = get_config().track_graph_compilation
+        self.num_graph_compilations = 0
 
     def _dump_with_sep(self, entry):
         entry = json.dumps(entry) + ','
@@ -256,11 +259,45 @@ def handler_fn(prof) -> None:
     def record_event(self, type, name, args=None):
         if self.enabled:
             self.start(type, name, args)
-            yield
+            with self.track_graph_compile(type, args) \
+                if self.gc_track_recompiles \
+                else contextlib.nullcontext():
+                yield
             self.end()
         else:
             yield
 
+    def record_block(self, type, name, ts, dur, args=None):
+        if self.enabled:
+            event = {
+                'pid': self.pid,
+                'tid': self.event_tid[type],
+                'ph': 'X',
+                'name': name,
+                'ts': ts,
+                'dur': dur,
+                'args': args
+            }
+            self._dump_with_sep(event)
+
+    @contextmanager
+    def track_graph_compile(self, type, args=None):
+        start = self.get_timestamp_us()
+        import habana_frameworks.torch as htorch
+        from habana_frameworks.torch.hpu.metrics import metric_localcontext
+        with metric_localcontext("graph_compilation") as gc:
+            yield
+            htorch.hpu.synchronize()
+        if gc.stats()[0][1] != 0:
+            compile_start_time = start
+            for recipe in gc.stats()[3][1]:
+                recipe_name = recipe[0]
+                compile_time = recipe[1]
+                self.num_graph_compilations += 1
+                self.record_counter(compile_start_time, {'cumulative_graph_compilations': self.num_graph_compilations})
+                self.record_block(type, 'GRAPH COMPILE: ' + recipe_name, compile_start_time, compile_time, args)
+                compile_start_time += compile_time
+
 
 # Adapted from https://stackoverflow.com/a/49361727
 def format_bytes(size):

vllm_gaudi/v1/worker/hpu_worker.py

@@ -25,7 +25,7 @@
 from vllm.v1.outputs import (DraftTokenIds, AsyncModelRunnerOutput, ModelRunnerOutput)
 from vllm.v1.worker.utils import bind_kv_cache
 from vllm_gaudi.utils import is_fake_hpu
-from vllm_gaudi.v1.worker.hpu_model_runner import HPUModelRunner, bool_helper
+from vllm_gaudi.v1.worker.hpu_model_runner import HPUModelRunner
 from vllm.v1.worker.worker_base import WorkerBase
 
 from vllm_gaudi.extension.logger import logger as init_logger
@@ -82,8 +82,7 @@ def __init__(
             from vllm.utils.import_utils import init_cached_hf_modules
             init_cached_hf_modules()
 
-        self.gc_track_recompiles = bool("PT_HPU_METRICS_GC_DETAILS" in os.environ
-                                        and bool_helper(os.getenv("PT_HPU_METRICS_GC_DETAILS")))
+        self.gc_track_recompiles = get_config().track_graph_compilation and not get_config().high_level_profiler_enabled
         self.step = 0
         self.profile_steps = get_config().VLLM_PROFILE_STEPS
         self.step_profiler = setup_step_profiler(self.profile_steps)