Enable pytest-xdist for TVM CI

Jenkinsfile

@@ -149,6 +149,9 @@ stage('Prepare') {
 }
 
 stage("Sanity Check") {
+  environment {
+    CI = 'true'
+  }
   timeout(time: max_time, unit: 'MINUTES') {
     node('CPU') {
       ws(per_exec_ws("tvm/sanity")) {
@@ -162,7 +165,7 @@ stage("Sanity Check") {
 // Run make. First try to do an incremental make from a previous workspace in hope to
 // accelerate the compilation. If something wrong, clean the workspace and then
 // build from scratch.
-def make(docker_type, path, make_flag) {
+def make(docker_type, path, make_flag='') {
   timeout(time: max_time, unit: 'MINUTES') {
     try {
       sh "${docker_run} ${docker_type} ./tests/scripts/task_build.sh ${path} ${make_flag}"
@@ -201,16 +204,19 @@ def unpack_lib(name, libs) {
 }
 
 stage('Build') {
+  environment {
+    CI = 'true'
+  }
   parallel 'BUILD: GPU': {
     node('GPUBUILD') {
       ws(per_exec_ws("tvm/build-gpu")) {
         init_git()
         sh "${docker_run} ${ci_gpu} ./tests/scripts/task_config_build_gpu.sh"
-        make(ci_gpu, 'build', '-j2')
+        make(ci_gpu, 'build')
         pack_lib('gpu', tvm_multilib)
         // compiler test
         sh "${docker_run} ${ci_gpu} ./tests/scripts/task_config_build_gpu_vulkan.sh"
-        make(ci_gpu, 'build2', '-j2')
+        make(ci_gpu, 'build2')
       }
     }
   },
@@ -219,7 +225,7 @@ stage('Build') {
       ws(per_exec_ws("tvm/build-cpu")) {
         init_git()
         sh "${docker_run} ${ci_cpu} ./tests/scripts/task_config_build_cpu.sh"
-        make(ci_cpu, 'build', '-j2')
+        make(ci_cpu, 'build')
         pack_lib('cpu', tvm_multilib_tsim)
         timeout(time: max_time, unit: 'MINUTES') {
           sh "${docker_run} ${ci_cpu} ./tests/scripts/task_ci_setup.sh"
@@ -240,7 +246,7 @@ stage('Build') {
       ws(per_exec_ws("tvm/build-wasm")) {
         init_git()
         sh "${docker_run} ${ci_wasm} ./tests/scripts/task_config_build_wasm.sh"
-        make(ci_wasm, 'build', '-j2')
+        make(ci_wasm, 'build')
         timeout(time: max_time, unit: 'MINUTES') {
           sh "${docker_run} ${ci_wasm} ./tests/scripts/task_ci_setup.sh"
           sh "${docker_run} ${ci_wasm} ./tests/scripts/task_web_wasm.sh"
@@ -253,7 +259,7 @@ stage('Build') {
       ws(per_exec_ws("tvm/build-i386")) {
         init_git()
         sh "${docker_run} ${ci_i386} ./tests/scripts/task_config_build_i386.sh"
-        make(ci_i386, 'build', '-j2')
+        make(ci_i386, 'build')
         pack_lib('i386', tvm_multilib_tsim)
       }
     }
@@ -263,7 +269,7 @@ stage('Build') {
       ws(per_exec_ws("tvm/build-arm")) {
         init_git()
         sh "${docker_run} ${ci_arm} ./tests/scripts/task_config_build_arm.sh"
-        make(ci_arm, 'build', '-j4')
+        make(ci_arm, 'build')
         pack_lib('arm', tvm_multilib)
       }
     }
@@ -273,7 +279,7 @@ stage('Build') {
       ws(per_exec_ws("tvm/build-qemu")) {
         init_git()
         sh "${docker_run} ${ci_qemu} ./tests/scripts/task_config_build_qemu.sh"
-        make(ci_qemu, 'build', '-j2')
+        make(ci_qemu, 'build')
         timeout(time: max_time, unit: 'MINUTES') {
           sh "${docker_run} ${ci_qemu} ./tests/scripts/task_ci_setup.sh"
           sh "${docker_run} ${ci_qemu} ./tests/scripts/task_python_microtvm.sh"
@@ -285,6 +291,9 @@ stage('Build') {
 }
 
 stage('Unit Test') {
+  environment {
+    CI = 'true'
+  }
   parallel 'python3: GPU': {
     node('TensorCore') {
       ws(per_exec_ws("tvm/ut-python-gpu")) {
@@ -345,6 +354,9 @@ stage('Unit Test') {
 }
 
 stage('Integration Test') {
+  environment {
+    CI = 'true'
+  }
   parallel 'topi: GPU': {
     node('GPU') {
       ws(per_exec_ws("tvm/topi-python-gpu")) {
@@ -401,6 +413,9 @@ stage('Integration Test') {
 
 /*
 stage('Build packages') {
+  environment {
+    CI = 'true'
+  }
   parallel 'conda CPU': {
     node('CPU') {
       sh "${docker_run} tlcpack/conda-cpu ./conda/build_cpu.sh
@@ -418,6 +433,9 @@ stage('Build packages') {
 */
 
 stage('Deploy') {
+    environment {
+      CI = 'true'
+    }
     node('doc') {
       ws(per_exec_ws("tvm/deploy-docs")) {
         if (env.BRANCH_NAME == "main") {

docker/bash.sh

@@ -33,16 +33,24 @@
 #     With -i, execute interactively.
 #
 
-set -euo pipefail
+set -xeuo pipefail
 
 
 function show_usage() {
     cat <<EOF
-Usage: docker/bash.sh [-i|--interactive] [--net=host] [-t|--tty]
+Usage: docker/bash.sh [-i|--interactive] [--net=host] [--cpuset-cpus=<cpus>]
          [--mount MOUNT_DIR] [--repo-mount-point REPO_MOUNT_POINT]
          [--dry-run]
          <DOCKER_IMAGE_NAME> [--] [COMMAND]
 
+--cpuset-cpus=<cpus>
+
+    Restrict docker container to use specific CPUs. See
+    docker run --help for further documentation of this parameter.
+    When launched from the CI (the "CI" environment variable is set),
+    this parameter is inferred from the "NODE_NAME" and "EXECUTOR_NUMBER"
+    environment variables.
+
 -h, --help
 
     Display this help message.
@@ -55,20 +63,20 @@ Usage: docker/bash.sh [-i|--interactive] [--net=host] [-t|--tty]
 
     Start the docker session with a pseudo terminal (tty).
 
---net=host
-
-    Expose servers run into the container to the host, passing the
-    "--net=host" argument through to docker.  On MacOS, this is
-    instead passed as "-p 8888:8888" since the host networking driver
-    isn't supported.
-
 --mount MOUNT_DIR
 
     Expose MOUNT_DIR as an additional mount point inside the docker
     container.  The mount point inside the container is the same as
     the folder location outside the container.  This option can be
     specified multiple times.
 
+--net=host
+
+    Expose servers run into the container to the host, passing the
+    "--net=host" argument through to docker.  On MacOS, this is
+    instead passed as "-p 8888:8888" since the host networking driver
+    isn't supported.
+
 --repo-mount-point REPO_MOUNT_POINT
 
     The directory inside the docker container at which the TVM
@@ -111,6 +119,7 @@ EOF
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
 REPO_DIR="$(dirname "${SCRIPT_DIR}")"
 
+CPUSET_CPUS=
 DRY_RUN=false
 INTERACTIVE=false
 TTY=false
@@ -146,6 +155,16 @@ break_joined_flag='if (( ${#1} == 2 )); then shift; else set -- -"${1#-i}" "${@:
 
 while (( $# )); do
     case "$1" in
+        --cpuset-cpus=?*)
+            CPUSET_CPUS="${1#*=}"
+            shift
+            ;;
+
+        --dry-run)
+            DRY_RUN=true
+            shift
+            ;;
+
         -h|--help)
             show_usage
             exit 0
@@ -180,8 +199,8 @@ while (( $# )); do
             shift
             ;;
 
-        --dry-run)
-            DRY_RUN=true
+        --net=host)
+            USE_NET_HOST=true
             shift
             ;;
 
@@ -276,6 +295,46 @@ DOCKER_ENV+=( --env CI_BUILD_HOME="${REPO_MOUNT_POINT}"
               --env CI_IMAGE_NAME="${DOCKER_IMAGE_NAME}"
             )
 
+# Choose CPUs on which this container will execute.
+if [ -n "${CI+x}" -a -z "${CPUSET_CPUS}" ]; then
+    if [ -n "${CI_NUM_EXECUTORS-}" ]; then
+        if [ -n "${CI_CPUSET_LOWER_BOUND-}" -a -n "${CI_CPUSET_UPPER_BOUND-}" ]; then
+            TOTAL_CPUS=$(expr "${CI_CPUSET_UPPER_BOUND}" - "${CI_CPUSET_LOWER_BOUND}" + 1) || /bin/true
+            if [ "${TOTAL_CPUS}" -lt 1 ]; then
+                echo "ERROR: computed TOTAL_CPUS=${TOTAL_CPUS} based on CI_CPUSET_{UPPER,LOWER}_BOUND!"
+                exit 2
+            fi
+        else
+            TOTAL_CPUS=$(nproc)
+            CI_CPUSET_LOWER_BOUND=0
+        fi
+        CPUS_PER_EXECUTOR=$(expr "${TOTAL_CPUS}" / "${CI_NUM_EXECUTORS}")
+        # NOTE: Expr exit status varies by the computed value (good and bad!).
+        CPUSET_CPUS_LOWER_BOUND=$(expr "${CI_CPUSET_LOWER_BOUND}" + \( "${CPUS_PER_EXECUTOR}" '*' "${EXECUTOR_NUMBER}" \) ) || /bin/true
+        CPUSET_CPUS_UPPER_BOUND=$(expr "${CPUSET_CPUS_LOWER_BOUND}" + "${CPUS_PER_EXECUTOR}" - 1) || /bin/true
+        CPUSET_CPUS="${CPUSET_CPUS_LOWER_BOUND}-${CPUSET_CPUS_UPPER_BOUND}"
+        echo "COMPUTE TOTAL_CPUS=${TOTAL_CPUS} CPUS_PER_EXECUTOR=${CPUS_PER_EXECUTOR} CPUSET_CPUS_LOWER_BOUND=${CPUSET_CPUS_LOWER_BOUND} CPUSET_CPUS_UPPER_BOUND=${CPUSET_CPUS_UPPER_BOUND}"
+    else
+        echo "WARNING: CI_NUM_EXECUTORS environment variable not set."
+        echo "No CPU parallism will be used in this CI build, so it may be quite slow."
+    fi
+fi
+
+if [ -n "${CPUSET_CPUS}" ]; then
+    if [ -z "$(echo ${CPUSET_CPUS} | sed -E '/^[0-9]+-[0-9]+$/ p; /.*/ d')" ]; then
+        echo "error: --cpuset-cpus: must specify in the form <m>-<n>; got ${CPUSET_CPUS}"
+        exit 2
+    fi
+    CPUSET_CPUS_LOWER_BOUND=$(echo "${CPUSET_CPUS}" | sed -E 's/^([0-9]+)-.*$/\1/g')
+    CPUSET_CPUS_UPPER_BOUND=$(echo "${CPUSET_CPUS}" | sed -E 's/^.*-([0-9]+)$/\1/g')
+    CPUSET_NUM_CPUS=$(expr "${CPUSET_CPUS_UPPER_BOUND}" - "${CPUSET_CPUS_LOWER_BOUND}" + 1) || /bin/true
+    DOCKER_FLAGS+=(
+        "--cpuset-cpus=${CPUSET_CPUS}"
+        "--env" "CI_CPUSET_CPUS=${CPUSET_CPUS}"
+        "--env" "CI_CPUSET_NUM_CPUS=${CPUSET_NUM_CPUS}"
+    )
+    echo "USING CPUSET_CPUS ${CPUSET_CPUS}"
+fi
 
 # Pass tvm test data folder through to the docker container, to avoid
 # repeated downloads.  Check if we have permissions to write to the

python/tvm/contrib/nvcc.py

@@ -97,7 +97,10 @@ def compile_cuda(code, target="ptx", arch=None, options=None, path_target=None):
     # if cxx_compiler_path != "":
     #    cmd += ["-ccbin", cxx_compiler_path]
 
-    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    # NOTE(areusch): Per https://github.com/lpereira/lwan/issues/106, stdin must be left open.
+    proc = subprocess.Popen(
+        cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+    )
 
     (out, _) = proc.communicate()
 

python/tvm/testing/plugin.py

@@ -36,6 +36,7 @@
 
 import tvm
 from tvm.testing import utils
+from xdist.scheduler.loadscope import LoadScopeScheduling
 
 
 MARKERS = {
@@ -288,3 +289,16 @@ def _parametrize_correlated_parameters(metafunc):
             names = ",".join(name for name, values in params)
             value_sets = zip(*[values for name, values in params])
             metafunc.parametrize(names, value_sets, indirect=True, ids=ids)
+
+
+class TvmTestScheduler(LoadScopeScheduling):
+    def _split_scope(self, nodeid):
+        # NOTE: test_tvm_testing_features contains parametrization-related tests, and must be
+        # serialized on a single host.
+        if "test_tvm_testing_features" in nodeid:
+            return "functional-tests"
+        return nodeid
+
+
+def pytest_xdist_make_scheduler(config, log):
+    return TvmTestScheduler(config, log)

tests/python/unittest/test_auto_scheduler_search_policy.py

@@ -210,11 +210,4 @@ def apply_func(search_policy, state, stage_id):
 
 
 if __name__ == "__main__":
-    test_workload_registry_empty_policy()
-    test_sketch_search_policy_basic()
-    test_sketch_search_policy_basic_spawn()
-    test_sketch_search_policy_xgbmodel()
-    test_sketch_search_policy_cuda_rpc_runner()
-    test_sketch_search_policy_cuda_xgbmodel_rpc_runner()
-    test_sketch_search_policy_zero_rank()
-    test_sketch_search_policy_custom_sketch()
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))

tests/scripts/setup-pytest-env.sh

@@ -32,7 +32,35 @@ export PYTHONPATH="${TVM_PATH}/python"
 export TVM_PYTEST_RESULT_DIR="${TVM_PATH}/build/pytest-results"
 mkdir -p "${TVM_PYTEST_RESULT_DIR}"
 
+if [ -n "${CI_CPUSET_NUM_CPUS-}" ]; then
+    # When the # of CPUs has been restricted (e.g. when --cpuset-cpus has been passed to docker by
+    # docker/bash.sh), explicitly use all available CPUs. This environment variable is set by
+    # docker/bash.sh when it sets --cpuset-cpus.
+    PYTEST_NUM_CPUS="${CI_CPUSET_NUM_CPUS}"
+else
+    # Else attempt to use $(nproc) - 1.
+    PYTEST_NUM_CPUS=$(nproc)
+    if [ -z "${PYTEST_NUM_CPUS}" ]; then
+        echo "WARNING: nproc failed; running pytest with only 1 CPU"
+        PYTEST_NUM_CPUS=1
+    elif [ ${PYTEST_NUM_CPUS} -gt 1 ]; then
+        PYTEST_NUM_CPUS=$(expr ${PYTEST_NUM_CPUS} - 1)  # Don't nuke interactive work.
+    fi
+
+    # Don't use >4 CPUs--in general, we only use 4 CPUs in testing, so we want to retain this
+    # maximum for the purposes of reproducing the CI. You can still override this by setting
+    # --cpuset-cpus in docker/bash.sh.
+    if [ ${PYTEST_NUM_CPUS} -gt 4 ]; then
+        PYTEST_NUM_CPUS=4
+    fi
+fi
+
 function run_pytest() {
+    local extra_args=( )
+    if [ "$1" == "--parallel" ]; then
+        extra_args=( -n "${PYTEST_NUM_CPUS}" )
+        shift
+    fi
     local ffi_type="$1"
     shift
     local test_suite_name="$1"
@@ -43,8 +71,10 @@ function run_pytest() {
         exit 2
     fi
     TVM_FFI=${ffi_type} python3 -m pytest \
+           --timeout=480 \
            -o "junit_suite_name=${test_suite_name}-${ffi_type}" \
            "--junit-xml=${TVM_PYTEST_RESULT_DIR}/${test_suite_name}-${ffi_type}.xml" \
            "--junit-prefix=${ffi_type}" \
+           "${extra_args[@]}" \
            "$@"
 }

tests/scripts/task_build.sh

@@ -15,5 +15,18 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+set -eux
+
 export VTA_HW_PATH=`pwd`/3rdparty/vta-hw
-cd $1 && cmake .. -DCMAKE_BUILD_TYPE=RelWithDebInfo && make $2 && cd ..
+MAKE_ARG=( )
+if [ -n "${2+x}" ]; then
+    MAKE_ARG=( "${2}" )
+fi
+
+if [ -n "${CI_CPUSET_NUM_CPUS+x}" -a "a${MAKE_ARG[@]:+b}" == "a" ]; then
+    MAKE_ARG=( "-j${CI_CPUSET_NUM_CPUS}" )
+fi
+
+cd "$1"
+cmake .. -DCMAKE_BUILD_TYPE=RelWithDebInfo
+make ${MAKE_ARG[@]+"${MAKE_ARG[@]}"}

tests/scripts/task_ci_setup.sh

@@ -30,7 +30,7 @@ set -o pipefail
 #
 echo "Addtiional setup in" ${CI_IMAGE_NAME}
 
-python3 -m pip install --user tlcpack-sphinx-addon==0.2.1 synr==0.4.0
+python3 -m pip install --user tlcpack-sphinx-addon==0.2.1 synr==0.4.0 pytest-timeout
 
 # Rebuild standalone_crt in build/ tree. This file is not currently archived by pack_lib() in
 # Jenkinsfile. We expect config.cmake to be present from pack_lib().