Clean up GPU memory after killing sglang processes (#2457)

Signed-off-by: Ata Fatahi <immrata@gmail.com>
This commit is contained in:
Ata Fatahi
2024-12-17 06:42:40 -05:00
committed by GitHub
parent e21026690d
commit ce094a5d79
2 changed files with 6 additions and 1 deletions

View File

@@ -60,7 +60,7 @@ jobs:
pip install --force-reinstall dist/*.whl
- name: Run e2e test
run: |
bash scripts/killall_sglang.sh
bash scripts/killall_sglang.sh "nuk_gpus"
cd sgl-router/py_test
python3 run_suite.py

View File

@@ -11,4 +11,9 @@ kill -9 $(ps aux | grep 'sglang.bench' | grep -v 'grep' | awk '{print $2}') 2>/d
# Clean all GPU processes if any argument is provided
if [ $# -gt 0 ]; then
kill -9 $(nvidia-smi | sed -n '/Processes:/,$p' | grep " [0-9]" | awk '{print $5}') 2>/dev/null
lsof /dev/nvidia* | awk '{print $2}' | xargs kill -9 2>/dev/null
fi
# Show GPU status after clean up
nvidia-smi