From c28bb981e1bcc9a9783cc197ea2634b78d694f1e Mon Sep 17 00:00:00 2001
From: Sing_chan <51314274+betterpig@users.noreply.github.com>
Date: Mon, 1 Aug 2022 17:19:08 +0800
Subject: [PATCH] set parallel_job according to CUDA memory in Windows CI
 unittest (#44695)

* set parallel_job according to CUDA memory

* fix bug: add whitespace between conten and [] or condition wont work
---
 tools/windows/run_unittests.sh | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tools/windows/run_unittests.sh b/tools/windows/run_unittests.sh
index 3c1a95f3691..1b8da6aa51e 100644
--- a/tools/windows/run_unittests.sh
+++ b/tools/windows/run_unittests.sh
@@ -336,7 +336,13 @@ function run_unittest_gpu() {
     if [ "$2" == "" ]; then
         parallel_job=$parallel_level_base
     else
-        parallel_job=`expr $2 \* $parallel_level_base`
+        # set parallel_job according to CUDA memory and suggested parallel num,
+        # the latter is derived in linux server with 16G CUDA memory.
+        cuda_memory=$(nvidia-smi --query-gpu=memory.total --format=csv | tail -1 | awk -F ' ' '{print $1}')
+        parallel_job=$(($2 * $cuda_memory / 16000))
+        if [ $parallel_job -lt 1 ]; then
+            parallel_job=1
+        fi
     fi
     echo "************************************************************************"
     echo "********These unittests run $parallel_job job each time with 1 GPU**********"
-- 
GitLab