[misc][ci] fix cpu test with plugins (#7489)

This commit is contained in:
youkaichao
2024-08-13 19:27:46 -07:00
committed by GitHub
parent 97992802f3
commit ea49e6a3c8
4 changed files with 15 additions and 8 deletions
+10 -3
View File
@@ -10,7 +10,6 @@ from pathlib import Path
from typing import Any, Callable, Dict, List, Optional
import openai
import ray
import requests
from transformers import AutoTokenizer
from typing_extensions import ParamSpec
@@ -18,9 +17,10 @@ from typing_extensions import ParamSpec
from vllm.distributed import (ensure_model_parallel_initialized,
init_distributed_environment)
from vllm.entrypoints.openai.cli_args import make_arg_parser
from vllm.platforms import current_platform
from vllm.utils import FlexibleArgumentParser, get_open_port, is_hip
if is_hip():
if current_platform.is_rocm():
from amdsmi import (amdsmi_get_gpu_vram_usage,
amdsmi_get_processor_handles, amdsmi_init,
amdsmi_shut_down)
@@ -32,7 +32,7 @@ if is_hip():
yield
finally:
amdsmi_shut_down()
else:
elif current_platform.is_cuda():
from pynvml import (nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo,
nvmlInit, nvmlShutdown)
@@ -43,6 +43,11 @@ else:
yield
finally:
nvmlShutdown()
else:
@contextmanager
def _nvml():
yield
VLLM_PATH = Path(__file__).parent.parent
@@ -293,6 +298,8 @@ def multi_process_parallel(
pp_size: int,
test_target: Any,
) -> None:
import ray
# Using ray helps debugging the error when it failed
# as compared to multiprocessing.
# NOTE: We need to set working_dir for distributed tests,