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

Signed-off-by: Ata Fatahi <[email protected]>
sgl-project · Dec 17, 2024 · ce094a5 · ce094a5
1 parent e210266
commit ce094a5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml
@@ -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
 

diff --git a/scripts/killall_sglang.sh b/scripts/killall_sglang.sh
@@ -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