diff --git a/ci/L0_multi_gpu_vllm/vllm_backend/vllm_multi_gpu_test.py b/ci/L0_multi_gpu_vllm/vllm_backend/vllm_multi_gpu_test.py index 21591845..f809d9a3 100644 --- a/ci/L0_multi_gpu_vllm/vllm_backend/vllm_multi_gpu_test.py +++ b/ci/L0_multi_gpu_vllm/vllm_backend/vllm_multi_gpu_test.py @@ -29,13 +29,15 @@ import unittest from functools import partial -import pynvml import tritonclient.grpc as grpcclient from tritonclient.utils import * +from vllm.utils import import_pynvml sys.path.append("../../common") from test_util import TestResultCollector, UserData, callback, create_vllm_request +pynvml = import_pynvml() + class VLLMMultiGPUTest(TestResultCollector): def setUp(self):