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

[fix]: ensure the cleaning of the GPU memory is hardware-agnostic

Browse Source 
Signed-off-by: AzizCode92 <azizbenothman76@gmail.com>
This commit is contained in:
AzizCode92 2025-09-05 09:34:10 +02:00
parent 4c8625fdb1
commit 963fe4b087
1 changed files with 10 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
tests/utils.py
View File

@ -172,10 +172,11 @@ class RemoteOpenAIServer:
self.proc.kill() self.proc.kill()
# GPU memory cleanup # GPU memory cleanup
try: try:
if torch.cuda.is_available(): if current_platform.is_cuda() or current_platform.is_rocm():
devices_to_clear = list(range(torch.cuda.device_count())) num_devices = cuda_device_count_stateless()
if devices_to_clear: if num_devices > 0:
wait_for_gpu_memory_to_clear(devices=devices_to_clear, wait_for_gpu_memory_to_clear(devices=list(
range(num_devices)),
threshold_ratio=0.05) threshold_ratio=0.05)
except Exception as e: except Exception as e:
print(f"GPU cleanup warning: {e}") print(f"GPU cleanup warning: {e}")
@ -277,10 +278,11 @@ class RemoteOpenAIServerCustom(RemoteOpenAIServer):
# GPU memory cleaning # GPU memory cleaning
try: try:
if torch.cuda.is_available(): if current_platform.is_cuda() or current_platform.is_rocm():
devices_to_clear = list(range(torch.cuda.device_count())) num_devices = cuda_device_count_stateless()
if devices_to_clear: if num_devices > 0:
wait_for_gpu_memory_to_clear(devices=devices_to_clear, wait_for_gpu_memory_to_clear(devices=list(
range(num_devices)),
threshold_ratio=0.05) threshold_ratio=0.05)
except Exception as e: except Exception as e:
print(f"GPU cleanup warning: {e}") print(f"GPU cleanup warning: {e}")