xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-23 23:11:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: Improve GPU cleanup in tests

Browse Source 
Centralizes the GPU memory cleanup logic into a single static method to prevent flaky test failures from OOM errors.

Signed-off-by: AzizCode92 <azizbenothman76@gmail.com>
This commit is contained in:
AzizCode92 2025-09-06 22:59:30 +02:00
parent 276ec32ba9
commit 2c9ed5b738
1 changed files with 8 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
tests/utils.py
View File

@ -170,14 +170,18 @@ class RemoteOpenAIServer:
except subprocess.TimeoutExpired: except subprocess.TimeoutExpired:
# force kill if needed # force kill if needed
self.proc.kill() self.proc.kill()
# GPU memory cleanup self.__cleanup_gpu_memory()
@staticmethod
def __cleanup_gpu_memory():
try: try:
if current_platform.is_cuda() or current_platform.is_rocm(): if current_platform.is_cuda_alike():
num_devices = cuda_device_count_stateless() num_devices = cuda_device_count_stateless()
if num_devices > 0: if num_devices > 0:
wait_for_gpu_memory_to_clear(devices=list( wait_for_gpu_memory_to_clear(devices=list(
range(num_devices)), range(num_devices)),
threshold_ratio=0.05) threshold_ratio=0.05,
timeout_s=60)
except Exception as e: except Exception as e:
print(f"GPU cleanup warning: {e}") print(f"GPU cleanup warning: {e}")
@ -276,16 +280,7 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):
# force kill if needed # force kill if needed
self.proc.kill() self.proc.kill()
# GPU memory cleaning self.__cleanup_gpu_memory()
try:
if current_platform.is_cuda() or current_platform.is_rocm():
num_devices = cuda_device_count_stateless()
if num_devices > 0:
wait_for_gpu_memory_to_clear(devices=list(
range(num_devices)),
threshold_ratio=0.05)
except Exception as e:
print(f"GPU cleanup warning: {e}")
def _test_completion( def _test_completion(