Make rpdtracer import only when required (#216)

* make rpdtracer import optional

* fix rpd_mark

* convert rpd_mark to try/except

* move rpd_trace import down

* move import

benchmarks/profiling/benchmark_latency.py

@@ -9,7 +9,6 @@
 
 import numpy as np
 import torch
-from rpdTracerControl import rpdTracerControl as rpd
 from tqdm import tqdm
 
 from vllm import LLM, SamplingParams
@@ -24,6 +23,7 @@ def main(args: argparse.Namespace):
 
     @contextmanager
     def rpd_profiler_context():
+        from rpdTracerControl import rpdTracerControl as rpd
         llm.start_profile()
         yield
         llm.stop_profile()

benchmarks/profiling/benchmark_throughput.py

@@ -9,7 +9,6 @@
 
 import torch
 import uvloop
-from rpdTracerControl import rpdTracerControl as rpd
 from tqdm import tqdm
 from transformers import (AutoModelForCausalLM, AutoTokenizer,
                           PreTrainedTokenizerBase)
@@ -98,6 +97,7 @@ def run_vllm(
 
     @contextmanager
     def rpd_profiler_context():
+        from rpdTracerControl import rpdTracerControl as rpd
         llm.start_profile()
         yield
         llm.stop_profile()

vllm/utils.py

@@ -31,7 +31,6 @@
 import torch.types
 import yaml
 from packaging.version import Version
-from rpdTracerControl import rpdTracerControl
 from typing_extensions import ParamSpec, TypeIs, assert_never
 
 import vllm.envs as envs
@@ -208,6 +207,7 @@ def setup_environment_variables(filename):
 
     def initialize_rpd_tracer(self, filename, nvtx):
         try:
+            from rpdTracerControl import rpdTracerControl
             rpd_trace.setup_environment_variables(filename)
             rpdTracerControl.setFilename(name=filename, append=True)
             return rpdTracerControl(nvtx=nvtx)
@@ -233,21 +233,35 @@ def create_file(filename):
             print(f"An error occurred while creating the filename: {e}")
 
 
+@lru_cache(maxsize=None)
+def is_hipScopedMarker_available():
+    try:
+        from hipScopedMarker import hipScopedMarker
+    except ImportError:
+        hipScopedMarker = None
+    return hipScopedMarker is not None
+
+
 class rpd_mark():
 
     def __init__(self, name=None):
         self.name = name
 
     def __call__(self, func):
-        from hipScopedMarker import hipScopedMarker
 
-        @wraps(func)
-        def inner(*args, **kwds):
-            marker_name = self.name if self.name else f"{func.__name__}"
-            with hipScopedMarker(f"{marker_name}"):
-                return func(*args, **kwds)
+        if is_hipScopedMarker_available():
+            from hipScopedMarker import hipScopedMarker
 
-        return inner
+            @wraps(func)
+            def inner(*args, **kwds):
+                marker_name = self.name if self.name else f"{func.__name__}"
+                with hipScopedMarker(f"{marker_name}"):
+                    return func(*args, **kwds)
+
+            return inner
+
+        else:
+            return func
 
 
 class Device(enum.Enum):

vllm/worker/worker.py

@@ -24,7 +24,6 @@
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sequence import (ExecuteModelRequest, IntermediateTensors,
                            SequenceGroupMetadata, SequenceGroupMetadataDelta)
-from vllm.utils import rpd_trace
 from vllm.worker.cache_engine import CacheEngine
 from vllm.worker.embedding_model_runner import EmbeddingModelRunner
 from vllm.worker.enc_dec_model_runner import EncoderDecoderModelRunner
@@ -144,6 +143,8 @@ def __init__(
             logger.info("Profiling enabled. Traces will be saved to: %s",
                         rpd_profiler_trace_dir)
 
+            from vllm.utils import rpd_trace
+
             if self.rank == 0:
                 rpd_trace.create_file(filename=str(rpd_profiler_trace_dir))