[release-test] Fix GPU memory usage detection (#2576)

pytorch · Jan 28, 2025 · 1231062 · 1231062
1 parent 72006c0
commit 1231062
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/.github/workflows/userbenchmark-a100.yml b/.github/workflows/userbenchmark-a100.yml
@@ -27,9 +27,9 @@ jobs:
       - name: Install Conda
         run: |
           bash ./.ci/torchbench/install-conda.sh
-      - name: Install TorchBench
-        run: |
-          bash ./.ci/torchbench/install.sh
+#      - name: Install TorchBench
+#        run: |
+#          bash ./.ci/torchbench/install.sh
       - name: Run user benchmark
         run: |
           set -x

diff --git a/userbenchmark/release-test/monitor_proc.sh b/userbenchmark/release-test/monitor_proc.sh
@@ -16,12 +16,12 @@ get_gpu_max_memory_usage_cuda() {
     local my_pid=$1
     local max=$2
     local curr
-    # Some processes might not use the GPU
-    if ! nvidia-smi pmon -s m -c 1 -o T | grep "${my_pid}" >/dev/null 2>/dev/null; then
+    curr=$(nvidia-smi dmon -s m -c 1 -o T -i 0 | tail -n +3 | awk '{print $3}' | sort -n | tail -1 | grep -o "[0-9.]*")
+    # Some processes might not use the GPU, then memory usage should be 0
+    if [ "${curr}" -eq 0 ] ; then
         echo "${max}"
         return
     fi
-    curr=$(nvidia-smi pmon -s m -c 1 -o T | grep "${my_pid}" | awk '{print $5}' | sort | tail -1 | grep -o "[0-9.]*")
     max "${curr}" "${max}"
 }