diff --git a/benchmarks/profiling/benchmark_latency.py b/benchmarks/profiling/benchmark_latency.py index 96ebbf760f25b..8a5834ff7b044 100644 --- 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 index 35d8fa1224661..46b587c16a8ca 100644 --- 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 index 465e7a1ac4ec3..a8d58a97dcec4 100644 --- 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)