tidy test scripts

pytorch · May 8, 2024 · 6173fed · 6173fed
1 parent 98d24de
commit 6173fed
Show file tree

Hide file tree

Showing 4 changed files with 153 additions and 90 deletions.
diff --git a/tests/common-test-functionality.sh b/tests/common-test-functionality.sh
@@ -0,0 +1,90 @@
+#!/bin/bash
+
+# Will catch exit code 5 when tests are deselected from previous passing run
+# (relevent for --last-failed-no-failures none)
+last_failed_no_failures_code=5
+
+#  functions shared across test files
+run_tests() {
+    # Set defaults
+    local core_args="-vvv tests"
+    local cache_dir=".unknown-cache"
+    local skip_distrib_tests=1
+    local match_tests_expression=""
+    local trap_deselected_exit_code=1
+    local use_last_failed=0
+    local use_coverage=0
+    # Always clean up pytest.ini
+    trap 'rm -f pytest.ini' RETURN
+    # Parse arguments
+    while [[ $# -gt 0 ]]
+    do
+        key="$1"
+        case $key in
+            --core_args)
+            core_args="$2"
+            shift
+            shift
+            ;;
+            --cache_dir)
+            cache_dir="$2"
+            shift
+            shift
+            ;;
+            --skip_distrib_tests)
+            skip_distrib_tests="$2"
+            shift
+            shift
+            ;;
+            --match_tests_expression)
+            match_tests_expression="$2"
+            shift
+            shift
+            ;;
+            --trap_deselected_exit_code)
+            trap_deselected_exit_code="$2"
+            shift
+            shift
+            ;;
+            --use_last_failed)
+            use_last_failed="$2"
+            shift
+            shift
+            ;;
+            --use_coverage)
+            use_coverage="$2"
+            shift
+            shift
+            ;;
+            *)
+            shift
+            ;;
+        esac
+    done
+
+    if [ "${skip_distrib_tests}" -eq "1" ]; then
+        # can be overwritten by core_args
+        skip_distrib_opt="-m 'not distributed and not tpu and not multinode_distributed'"
+    else
+        skip_distrib_opt=""
+    fi
+
+
+    echo [pytest] > pytest.ini ; echo "cache_dir=${cache_dir}" >> pytest.ini
+
+    # Assemble options for the pytest command
+    pytest_args="${skip_distrib_opt} ${core_args} -k '${match_tests_expression}' tests"
+    if [ "${use_last_failed:-0}" -eq "1" ] && [ -d "${cache_dir}" ]; then
+        pytest_args="--last-failed --last-failed-no-failures none ${pytest_args}"
+    fi
+    if [ "${use_coverage}" -eq "1" ]; then
+        pytest_args="--cov ignite --cov-append --cov-report term-missing --cov-report xml ${pytest_args}"
+    fi
+
+    # Run the command
+    if [ "$trap_deselected_exit_code" -eq "1" ]; then
+        CUDA_VISIBLE_DEVICES="" eval "pytest ${pytest_args}" || { exit_code=$?; if [ "$exit_code" -eq ${last_failed_no_failures_code} ]; then echo "All tests deselected"; else exit $exit_code; fi; }
+    else
+        CUDA_VISIBLE_DEVICES="" eval "pytest ${pytest_args}"
+    fi
+}
diff --git a/tests/run_cpu_tests.sh b/tests/run_cpu_tests.sh
@@ -1,39 +1,30 @@
 #!/bin/bash
-
+source "$(dirname "$0")/common-test-functionality.sh"
 set -xeu
 
-if [ "${SKIP_DISTRIB_TESTS:-0}" -eq "1" ]; then
-    skip_distrib_opt="not distributed and not tpu and not multinode_distributed"
-else
-    skip_distrib_opt=""
-fi
-
-MATCH_TESTS_EXPRESSION=${1:-""}
+skip_distrib_tests=${SKIP_DISTRIB_TESTS:-0}
+use_last_failed=${USE_LAST_FAILED:-0}
+match_tests_expression=${1:-""}
 
-# Will catch exit code 5 when tests are deselected from previous passing run
-EXIT_CODE_ALL_TESTS_DESELECTED=5
 
-CACHE_DIR=.cpu-not-distrib
-echo [pytest] > pytest.ini ; echo "cache_dir=${CACHE_DIR}" >> pytest.ini
-PYTEST_ARGS="--tx 4*popen//python=python --cov ignite --cov-report term-missing --cov-report xml -vvv tests -m '${skip_distrib_opt}' -k '${MATCH_TESTS_EXPRESSION}'"
-if [ "${USE_LAST_FAILED:-0}" -eq "1" ] && [ -d "${CACHE_DIR}" ]; then
-    PYTEST_ARGS="--last-failed --last-failed-no-failures none ${PYTEST_ARGS}"
-fi
-CUDA_VISIBLE_DEVICES="" eval "pytest ${PYTEST_ARGS}"  || { exit_code=$?; if [ "$exit_code" -eq 5 ]; then echo "All tests deselected"; else exit $exit_code; fi;}
+run_tests \
+    --core-args "--tx 4*popen//python=python -vvv tests" \
+    --cache_dir ".cpu-not-distrib" \
+    --skip_distrib_tests "${skip_distrib_tests}" \
+    --use_coverage 1 \
+    --match_tests_expression "${match_tests_expression}" \
+    --use_last_failed ${use_last_failed}
 
 # https://pubs.opengroup.org/onlinepubs/009695399/utilities/xcu_chap02.html#tag_02_06_02
-if [ "${SKIP_DISTRIB_TESTS:-0}" -eq "1" ]; then
+if [ "${skip_distrib_tests}" -eq "1" ]; then
     exit 0
 fi
 
-export WORLD_SIZE=2
-CACHE_DIR=.cpu-distrib
-echo [pytest] > pytest.ini ; echo "cache_dir=${CACHE_DIR}" >> pytest.ini
-PYTEST_ARGS="--cov ignite --cov-append --cov-report term-missing --cov-report xml --dist=each --tx ${WORLD_SIZE}*popen//python=python tests -m distributed -vvv -k '${MATCH_TESTS_EXPRESSION}'"
-if [ "${USE_LAST_FAILED:-0}" -eq "1" ] && [ -d "${CACHE_DIR}" ]; then
-    PYTEST_ARGS="--last-failed --last-failed-no-failures none ${PYTEST_ARGS}"
-fi
-CUDA_VISIBLE_DEVICES="" eval "pytest ${PYTEST_ARGS}"
-unset WORLD_SIZE
-
-rm -f pytest.ini
+# Run 2 processes with --dist=each
+run_tests \
+    --core-args "--dist=each --tx 2*popen//python=python -m distributed -vvv tests" \
+    --cache_dir ".cpu-distrib" \
+    --skip_distrib_tests 0 \
+    --use_coverage 1 \
+    --match_tests_expression "${match_tests_expression}" \
+    --use_last_failed ${use_last_failed}
diff --git a/tests/run_gpu_tests.sh b/tests/run_gpu_tests.sh
@@ -1,58 +1,46 @@
 #!/bin/bash
+source "$(dirname "$0")/common-test-functionality.sh"
+set -xeu
 
-if [ -z "$1" ]; then
-    ngpus=1
-else
-    ngpus=$1
-fi
-
-MATCH_TESTS_EXPRESSION=${2:-""}
+skip_distrib_tests=${SKIP_DISTRIB_TESTS:-1}
+use_last_failed=${USE_LAST_FAILED:-0}
+ngpus=${1:-1}
 
-if [ -z "$MATCH_TESTS_EXPRESSION" ]; then
+match_tests_expression=${2:-""}
+if [ -z "$match_tests_expression" ]; then
     cuda_pattern="cuda"
 else
-    cuda_pattern="cuda and $MATCH_TESTS_EXPRESSION"
-fi
-
-# Will catch exit code 5 when tests are deselected from previous passing run
-EXIT_CODE_ALL_TESTS_DESELECTED=5
-
-set -xeu
-
-CACHE_DIR=.gpu-cuda
-echo [pytest] > pytest.ini ; echo "cache_dir=${CACHE_DIR}" >> pytest.ini
-PYTEST_ARGS="--cov ignite --cov-report term-missing --cov-report xml -vvv tests/ -k '${cuda_pattern}'"
-if [ "${USE_LAST_FAILED:-0}" -eq "1" ] && [ -d "${CACHE_DIR}" ]; then
-    PYTEST_ARGS="--last-failed --last-failed-no-failures none ${PYTEST_ARGS}"
+    cuda_pattern="cuda and $match_tests_expression"
 fi
-CUDA_VISIBLE_DEVICES="" eval "pytest ${PYTEST_ARGS}" || { exit_code=$?; if [ "$exit_code" -eq 5 ]; then echo "All tests deselected"; else exit $exit_code; fi;}
-
 
+run_tests \
+    --core-args "-vvv tests" \
+    --cache_dir ".gpu-cuda" \
+    --skip_distrib_tests "${skip_distrib_tests}" \
+    --use_coverage 1 \
+    --match_tests_expression "${cuda_pattern}" \
+    --use_last_failed ${use_last_failed}
 
 # https://pubs.opengroup.org/onlinepubs/009695399/utilities/xcu_chap02.html#tag_02_06_02
-if [ "${SKIP_DISTRIB_TESTS:-0}" -eq "1" ]; then
+if [ "${skip_distrib_tests}" -eq "1" ]; then
     exit 0
 fi
 
-CACHE_DIR=.gpu-distrib
-echo [pytest] > pytest.ini ; echo "cache_dir=${CACHE_DIR}" >> pytest.ini
-PYTEST_ARGS="--cov ignite --cov-append --cov-report term-missing --cov-report xml -vvv tests/ -m distributed -k '${MATCH_TESTS_EXPRESSION}'"
-if [ "${USE_LAST_FAILED:-0}" -eq "1" ] && [ -d "${CACHE_DIR}" ]; then
-    PYTEST_ARGS="--last-failed --last-failed-no-failures none ${PYTEST_ARGS}"
-fi
-CUDA_VISIBLE_DEVICES="" eval "pytest ${PYTEST_ARGS}" || { exit_code=$?; if [ "$exit_code" -eq 5 ]; then echo "All tests deselected"; else exit $exit_code; fi;}
+run_tests \
+    --core-args "-vvv -m distributed tests" \
+    --cache_dir ".gpu-distrib" \
+    --skip_distrib_tests 0 \
+    --use_coverage 1 \
+    --match_tests_expression "${match_tests_expression}" \
+    --use_last_failed ${use_last_failed}
 
-if [ ${ngpus} -gt 1 ]; then
-
-    export WORLD_SIZE=${ngpus}
-    CACHE_DIR=.gpu-distrib-multi
-    echo [pytest] > pytest.ini ; echo "cache_dir=${CACHE_DIR}" >> pytest.ini
-    PYTEST_ARGS="--cov ignite --cov-append --cov-report term-missing --cov-report xml --dist=each --tx ${WORLD_SIZE}*popen//python=python tests -m distributed -vvv -k '${MATCH_TESTS_EXPRESSION}'"
-    if [ "${USE_LAST_FAILED:-0}" -eq "1" ] && [ -d "${CACHE_DIR}" ]; then
-        PYTEST_ARGS="--last-failed --last-failed-no-failures none ${PYTEST_ARGS}"
-    fi
-    CUDA_VISIBLE_DEVICES="" eval "pytest ${PYTEST_ARGS}"
-    unset WORLD_SIZE
 
+if [ ${ngpus} -gt 1 ]; then
+    run_tests \
+        --core-args "--tx ${ngpus}*popen//python=python --dist=each -vvv -m distributed tests" \
+        --cache_dir ".gpu-distrib-multi" \
+        --skip_distrib_tests 0 \
+        --use_coverage 1 \
+        --match_tests_expression "${match_tests_expression}" \
+        --use_last_failed ${use_last_failed}
 fi
-rm -f pytest.ini
diff --git a/tests/run_tpu_tests.sh b/tests/run_tpu_tests.sh
@@ -1,26 +1,20 @@
 #!/bin/bash
-# Will catch exit code 5 when tests are deselected from previous passing run
-EXIT_CODE_ALL_TESTS_DESELECTED=5
-
+source "$(dirname "$0")/common-test-functionality.sh"
 set -xeu
+use_last_failed=${USE_LAST_FAILED:-0}
 
-CACHE_DIR=.tpu
-echo [pytest] > pytest.ini ; echo "cache_dir=${CACHE_DIR}" >> pytest.ini
-PYTEST_ARGS="--cov ignite --cov-report term-missing --cov-report xml tests/ -vvv -m tpu"
-if [ "${USE_LAST_FAILED:-0}" -eq "1" ] && [ -d "${CACHE_DIR}" ]; then
-    PYTEST_ARGS="--last-failed --last-failed-no-failures none ${PYTEST_ARGS}"
-fi
-CUDA_VISIBLE_DEVICES="" eval "pytest ${PYTEST_ARGS}" || { exit_code=$?; if [ "$exit_code" -eq 5 ]; then echo "All tests deselected"; else exit $exit_code; fi;}
+run_tests \
+    --core-args "-vvv -m tpu tests" \
+    --cache_dir ".tpu" \
+    --use_coverage 1 \
+    --use_last_failed ${use_last_failed}
 
 
 if [ -z ${NUM_TPU_WORKERS+x} ]; then
     export NUM_TPU_WORKERS=1
-    CACHE_DIR=.tpu-multi
-    echo [pytest] > pytest.ini ; echo "cache_dir=${CACHE_DIR}" >> pytest.ini
-    PYTEST_ARGS="--cov ignite --cov-append --cov-report term-missing --cov-report xml tests/ -vvv -m tpu"
-    if [ "${USE_LAST_FAILED:-0}" -eq "1" ] && [ -d "${CACHE_DIR}" ]; then
-        PYTEST_ARGS="--last-failed --last-failed-no-failures none ${PYTEST_ARGS}"
-    fi
-    CUDA_VISIBLE_DEVICES="" eval "pytest ${PYTEST_ARGS}" || { exit_code=$?; if [ "$exit_code" -eq 5 ]; then echo "All tests deselected"; else exit $exit_code; fi;}
+    run_tests \
+        --core-args "-vvv -m tpu tests" \
+        --cache_dir ".tpu-multi" \
+        --use_coverage 1 \
+        --use_last_failed ${use_last_failed}
 fi
-rm -f pytest.ini