From 943e82a7cb634e4c86360bd3590c6a9e2786ccef Mon Sep 17 00:00:00 2001
From: Tianqi Chen <tqchen@users.noreply.github.com>
Date: Mon, 15 Jun 2020 10:28:59 -0700
Subject: [PATCH] [CI] Limit number of threads in all jobs (#5815)

---
 tests/scripts/task_cpp_unittest.sh       | 4 ++++
 tests/scripts/task_golang.sh             | 4 ++++
 tests/scripts/task_java_unittest.sh      | 4 ++++
 tests/scripts/task_python_docs.sh        | 4 ++++
 tests/scripts/task_python_integration.sh | 2 ++
 tests/scripts/task_python_topi.sh        | 5 +++++
 tests/scripts/task_python_vta_fsim.sh    | 4 ++++
 tests/scripts/task_python_vta_tsim.sh    | 4 ++++
 tests/scripts/task_rust.sh               | 4 ++++
 9 files changed, 35 insertions(+)

diff --git a/tests/scripts/task_cpp_unittest.sh b/tests/scripts/task_cpp_unittest.sh
index 751e98e9abdc..5ac1843253d4 100755
--- a/tests/scripts/task_cpp_unittest.sh
+++ b/tests/scripts/task_cpp_unittest.sh
@@ -23,6 +23,10 @@ export LD_LIBRARY_PATH="lib:${LD_LIBRARY_PATH:-}"
 # NOTE: important to use abspath, when VTA is enabled.
 export VTA_HW_PATH=`pwd`/3rdparty/vta-hw
 
+# to avoid CI thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=1
+
 # Remove existing testcases
 rm -f build/*_test
 
diff --git a/tests/scripts/task_golang.sh b/tests/scripts/task_golang.sh
index 49965793f6b3..0ff6c39d602c 100755
--- a/tests/scripts/task_golang.sh
+++ b/tests/scripts/task_golang.sh
@@ -24,5 +24,9 @@ export LD_LIBRARY_PATH="lib:${LD_LIBRARY_PATH:-}"
 tvm_root="$(git rev-parse --show-toplevel)"
 export PYTHONPATH="$tvm_root/python":"$tvm_root/topi/python"
 
+# to avoid CI CPU thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=1
+
 # Golang tests
 make -C golang tests
diff --git a/tests/scripts/task_java_unittest.sh b/tests/scripts/task_java_unittest.sh
index 63f16fd755f6..7ab4afae3c2e 100755
--- a/tests/scripts/task_java_unittest.sh
+++ b/tests/scripts/task_java_unittest.sh
@@ -22,6 +22,10 @@ set -u
 export PYTHONPATH=python
 export LD_LIBRARY_PATH="lib:${LD_LIBRARY_PATH:-}"
 
+# to avoid CI CPU thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=1
+
 CURR_DIR=$(cd `dirname $0`; pwd)
 SCRIPT_DIR=$CURR_DIR/../../jvm/core/src/test/scripts
 TEMP_DIR=$(mktemp -d)
diff --git a/tests/scripts/task_python_docs.sh b/tests/scripts/task_python_docs.sh
index 4c52e1e4ca5e..c239abb7bc9d 100755
--- a/tests/scripts/task_python_docs.sh
+++ b/tests/scripts/task_python_docs.sh
@@ -21,6 +21,10 @@ set -u
 
 source tests/scripts/setup-pytest-env.sh
 
+# to avoid CI CPU thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=4
+
 cleanup()
 {
     rm -rf /tmp/$$.log.txt
diff --git a/tests/scripts/task_python_integration.sh b/tests/scripts/task_python_integration.sh
index 296f9338bfb6..f7539d6a55fd 100755
--- a/tests/scripts/task_python_integration.sh
+++ b/tests/scripts/task_python_integration.sh
@@ -22,6 +22,8 @@ set -u
 source tests/scripts/setup-pytest-env.sh
 export PYTHONPATH=${PYTHONPATH}:${TVM_PATH}/apps/extension/python
 export LD_LIBRARY_PATH="build:${LD_LIBRARY_PATH:-}"
+
+# to avoid CI CPU thread throttling.
 export TVM_BIND_THREADS=0
 export TVM_NUM_THREADS=2
 
diff --git a/tests/scripts/task_python_topi.sh b/tests/scripts/task_python_topi.sh
index 66585d56d1db..e483d5f7f4a6 100755
--- a/tests/scripts/task_python_topi.sh
+++ b/tests/scripts/task_python_topi.sh
@@ -20,6 +20,11 @@ set -e
 set -u
 
 source tests/scripts/setup-pytest-env.sh
+
+# to avoid CI thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=1
+
 # Rebuild cython
 make cython3
 
diff --git a/tests/scripts/task_python_vta_fsim.sh b/tests/scripts/task_python_vta_fsim.sh
index f48b7d6d1594..8080bbe756c7 100755
--- a/tests/scripts/task_python_vta_fsim.sh
+++ b/tests/scripts/task_python_vta_fsim.sh
@@ -20,6 +20,10 @@ set -e
 set -u
 
 source tests/scripts/setup-pytest-env.sh
+# to avoid CI thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=1
+
 export PYTHONPATH=${PYTHONPATH}:${TVM_PATH}/vta/python
 export VTA_HW_PATH=`pwd`/3rdparty/vta-hw
 
diff --git a/tests/scripts/task_python_vta_tsim.sh b/tests/scripts/task_python_vta_tsim.sh
index 8e9a5eafbd9a..c87d5483b8a5 100755
--- a/tests/scripts/task_python_vta_tsim.sh
+++ b/tests/scripts/task_python_vta_tsim.sh
@@ -23,6 +23,10 @@ source tests/scripts/setup-pytest-env.sh
 export PYTHONPATH=${PYTHONPATH}:${TVM_PATH}/vta/python
 export VTA_HW_PATH=`pwd`/3rdparty/vta-hw
 
+# to avoid CI CPU thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=1
+
 # cleanup pycache
 find . -type f -path "*.pyc" | xargs rm -f
 
diff --git a/tests/scripts/task_rust.sh b/tests/scripts/task_rust.sh
index 1cbf4fa2cd07..17bad38fa71b 100755
--- a/tests/scripts/task_rust.sh
+++ b/tests/scripts/task_rust.sh
@@ -27,6 +27,10 @@ export RUST_DIR="$TVM_HOME/rust"
 export LLVM_CONFIG_PATH=`which llvm-config-10`
 echo "Using $LLVM_CONFIG_PATH"
 
+# to avoid CI CPU thread throttling.
+export TVM_BIND_THREADS=0
+export OMP_NUM_THREADS=1
+
 cd $RUST_DIR
 cargo fmt -- --check