make rpdtracer import optional

ROCm · Sep 30, 2024 · 0812d36 · 0812d36
1 parent 956b831
commit 0812d36
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/benchmarks/profiling/benchmark_latency.py b/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()

diff --git a/benchmarks/profiling/benchmark_throughput.py b/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()

diff --git a/vllm/utils.py b/vllm/utils.py
@@ -30,7 +30,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
@@ -207,6 +206,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)