From ffd6a1bf092be6f6ace9cdfff115e88ea33a60fa Mon Sep 17 00:00:00 2001
From: George Guanheng Zhang <zhangguanheng@fb.com>
Date: Thu, 4 Mar 2021 08:35:03 -0800
Subject: [PATCH] Run CI unittests in parallel (#3445)

Summary:
* enable parallel tests

* disable parallelism for GPU tests

* [test] limit maximum processes on linux

* [debug] limit max processes even further

* [test] use subprocesses over threads

* [test] limit intra-op threads

* only limit intra op threads for CPU tests

* [poc] use low timeout for showcasing

* [poc] fix syntax

* set timeout to 5 minutes

* fix timeout on windows

Reviewed By: fmassa

Differential Revision: D26756257

fbshipit-source-id: f2fc4753a67a1505f01116119926eec365693ab9

Co-authored-by: Francisco Massa <fvsmassa@gmail.com>
---
 .../unittest/linux/scripts/environment.yml      |  2 ++
 .circleci/unittest/linux/scripts/run_test.sh    | 17 ++++++++++++++++-
 .../unittest/windows/scripts/environment.yml    |  2 ++
 .circleci/unittest/windows/scripts/run_test.sh  | 17 ++++++++++++++++-
 4 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/.circleci/unittest/linux/scripts/environment.yml b/.circleci/unittest/linux/scripts/environment.yml
index dcad1abfa31..04aa9007e88 100644
--- a/.circleci/unittest/linux/scripts/environment.yml
+++ b/.circleci/unittest/linux/scripts/environment.yml
@@ -5,6 +5,8 @@ channels:
   - conda-forge
 dependencies:
   - pytest
+  - pytest-xdist
+  - pytest-timeout
   - pytest-cov
   - codecov
   - pip
diff --git a/.circleci/unittest/linux/scripts/run_test.sh b/.circleci/unittest/linux/scripts/run_test.sh
index 419b9eb562c..a0a2e0dae3f 100755
--- a/.circleci/unittest/linux/scripts/run_test.sh
+++ b/.circleci/unittest/linux/scripts/run_test.sh
@@ -6,5 +6,20 @@ eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
 export PYTORCH_TEST_WITH_SLOW='1'
+if [ "${CU_VERSION:-}" == cpu ] ; then
+    NUMPROCESSES="auto"
+    export OMP_NUM_THREADS="1"
+else
+    NUMPROCESSES="1"
+fi
+
 python -m torch.utils.collect_env
-pytest --cov=torchvision --junitxml=test-results/junit.xml -v --durations 20 test --ignore=test/test_datasets_download.py
+pytest \
+    --numprocesses=$NUMPROCESSES \
+    --timeout=300 \
+    --cov=torchvision \
+    --junitxml=test-results/junit.xml \
+    --verbose \
+    --durations 20 \
+    --ignore=test/test_datasets_download.py \
+    test
diff --git a/.circleci/unittest/windows/scripts/environment.yml b/.circleci/unittest/windows/scripts/environment.yml
index b4f32cb3cad..069d57e1e03 100644
--- a/.circleci/unittest/windows/scripts/environment.yml
+++ b/.circleci/unittest/windows/scripts/environment.yml
@@ -5,6 +5,8 @@ channels:
   - conda-forge
 dependencies:
   - pytest
+  - pytest-xdist
+  - pytest-timeout
   - pytest-cov
   - codecov
   - pip
diff --git a/.circleci/unittest/windows/scripts/run_test.sh b/.circleci/unittest/windows/scripts/run_test.sh
index 96d9cbd6b2d..a70956347ac 100644
--- a/.circleci/unittest/windows/scripts/run_test.sh
+++ b/.circleci/unittest/windows/scripts/run_test.sh
@@ -6,5 +6,20 @@ eval "$(./conda/Scripts/conda.exe 'shell.bash' 'hook')"
 conda activate ./env
 
 export PYTORCH_TEST_WITH_SLOW='1'
+if [ "${CU_VERSION:-}" == cpu ] ; then
+    NUMPROCESSES="auto"
+    export OMP_NUM_THREADS="1"
+else
+    NUMPROCESSES="1"
+fi
+
 python -m torch.utils.collect_env
-pytest --cov=torchvision --junitxml=test-results/junit.xml -v --durations 20 test --ignore=test/test_datasets_download.py
+pytest \
+    --numprocesses=$NUMPROCESSES \
+    --timeout=300 \
+    --cov=torchvision \
+    --junitxml=test-results/junit.xml \
+    --verbose \
+    --durations 20 \
+    --ignore=test/test_datasets_download.py \
+    test