37 lines
853 B
Python
37 lines
853 B
Python
import os
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clean_gpu_memory_between_tests():
|
|
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
|
|
yield
|
|
return
|
|
|
|
# Wait for GPU memory to be cleared before starting the test
|
|
import gc
|
|
|
|
from tests.utils import wait_for_gpu_memory_to_clear
|
|
|
|
num_gpus = torch.cuda.device_count()
|
|
if num_gpus > 0:
|
|
try:
|
|
wait_for_gpu_memory_to_clear(
|
|
devices=list(range(num_gpus)),
|
|
threshold_ratio=0.1,
|
|
)
|
|
except ValueError as e:
|
|
logger.info("Failed to clean GPU memory: %s", e)
|
|
|
|
yield
|
|
|
|
# Clean up GPU memory after the test
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|