Make 2 more run_tests.sh support test selection (#9192)

Co-authored-by: Zhanyong Wan <[email protected]>

Author: zhanyong-wan

benchmarks/run_benchmark.sh

@@ -5,7 +5,7 @@ LOGFILE=/tmp/benchmark_test.log
 
 # Note [Keep Going]
 #
-# Set the `CONTINUE_ON_ERROR` flag to `true` to make the CI tests continue on error.
+# Set the `CONTINUE_ON_ERROR` flag to `1` to make the CI tests continue on error.
 # This will allow you to see all the failures on your PR, not stopping with the first
 # test failure like the default behavior.
 CONTINUE_ON_ERROR="${CONTINUE_ON_ERROR:-0}"

test/benchmarks/run_tests.sh

@@ -9,7 +9,7 @@ export PYTHONPATH=$PYTHONPATH:$CDIR/../../benchmarks/
 
 # Note [Keep Going]
 #
-# Set the `CONTINUE_ON_ERROR` flag to `true` to make the CI tests continue on error.
+# Set the `CONTINUE_ON_ERROR` flag to `1` to make the CI tests continue on error.
 # This will allow you to see all the failures on your PR, not stopping with the first
 # test failure like the default behavior.
 CONTINUE_ON_ERROR="${CONTINUE_ON_ERROR:-0}"

test/benchmarks/run_torchbench_tests.sh

@@ -14,7 +14,7 @@ TORCHBENCH_DIR=$PYTORCH_DIR/benchmark
 
 # Note [Keep Going]
 #
-# Set the `CONTINUE_ON_ERROR` flag to `true` to make the CI tests continue on error.
+# Set the `CONTINUE_ON_ERROR` flag to `1` to make the CI tests continue on error.
 # This will allow you to see all the failures on your PR, not stopping with the first
 # test failure like the default behavior.
 CONTINUE_ON_ERROR="${CONTINUE_ON_ERROR:-0}"

test/neuron/run_tests.sh

@@ -1,43 +1,25 @@
 #!/bin/bash
 set -exo pipefail
 CDIR="$(cd "$(dirname "$0")"/../ ; pwd -P)"
-LOGFILE=/tmp/pytorch_py_test.log
-MAX_GRAPH_SIZE=500
-GRAPH_CHECK_FREQUENCY=100
-VERBOSITY=2
 
 # Utils file
 source "${CDIR}/utils/run_tests_utils.sh"
 
+parse_options_to_vars $@
+
+# Consume the parsed commandline arguments.
+shift $(($OPTIND - 1))
+
 # Note [Keep Going]
 #
-# Set the `CONTINUE_ON_ERROR` flag to `true` to make the CI tests continue on error.
+# Set the `CONTINUE_ON_ERROR` flag to `1` to make the CI tests continue on error.
 # This will allow you to see all the failures on your PR, not stopping with the first
 # test failure like the default behavior.
 CONTINUE_ON_ERROR="${CONTINUE_ON_ERROR:-0}"
 if [[ "$CONTINUE_ON_ERROR" == "1" ]]; then
   set +e
 fi
 
-while getopts 'LM:C:V:' OPTION
-do
-  case $OPTION in
-    L)
-      LOGFILE=
-      ;;
-    M)
-      MAX_GRAPH_SIZE=$OPTARG
-      ;;
-    C)
-      GRAPH_CHECK_FREQUENCY=$OPTARG
-      ;;
-    V)
-      VERBOSITY=$OPTARG
-      ;;
-  esac
-done
-shift $(($OPTIND - 1))
-
 export TRIM_GRAPH_SIZE=$MAX_GRAPH_SIZE
 export TRIM_GRAPH_CHECK_FREQUENCY=$GRAPH_CHECK_FREQUENCY
 export TORCH_TEST_DEVICES="$CDIR/pytorch_test_base.py"
@@ -49,6 +31,9 @@ TORCH_XLA_DIR=$(cd ~; dirname "$(python -c 'import torch_xla; print(torch_xla.__
 COVERAGE_FILE="$CDIR/../.coverage"
 
 function run_coverage {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   if [ "${USE_COVERAGE:-0}" != "0" ]; then
     coverage run --source="$TORCH_XLA_DIR" -p "$@"
   else
@@ -57,61 +42,97 @@ function run_coverage {
 }
 
 function run_test {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running in PjRt runtime: $@"
   PJRT_DEVICE=NEURON NEURON_NUM_DEVICES=1 run_coverage "$@"
 }
 
 function run_test_without_functionalization {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running with XLA_DISABLE_FUNCTIONALIZATION: $@"
   XLA_DISABLE_FUNCTIONALIZATION=1 run_test "$@"
 }
 
 function run_xla_ir_debug {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running with XLA_IR_DEBUG: $@"
   XLA_IR_DEBUG=1 run_test "$@"
 }
 
 function run_use_bf16 {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running with XLA_USE_BF16: $@"
   XLA_USE_BF16=1 run_test "$@"
 }
 
 function run_downcast_bf16 {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running with XLA_DOWNCAST_BF16: $@"
   XLA_DOWNCAST_BF16=1 run_test "$@"
 }
 
 function run_xla_hlo_debug {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running with XLA_IR_DEBUG: $@"
   XLA_HLO_DEBUG=1 run_test "$@"
 }
 
 function run_dynamic {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running in DynamicShape mode: $@"
   XLA_EXPERIMENTAL="nonzero:masked_select:masked_scatter:nms" run_test "$@"
 }
 
 function run_eager_debug {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running in Eager Debug mode: $@"
   XLA_USE_EAGER_DEBUG_MODE=1 run_test "$@"
 }
 
 function run_pt_xla_debug {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running in save tensor file mode: $@"
   PT_XLA_DEBUG=1 PT_XLA_DEBUG_FILE="/tmp/pt_xla_debug.txt" run_test "$@"
 }
 
 function run_pt_xla_debug_level1 {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running in save tensor file mode: $@"
   PT_XLA_DEBUG_LEVEL=1 PT_XLA_DEBUG_FILE="/tmp/pt_xla_debug.txt" run_test "$@"
 }
 
 function run_pt_xla_debug_level2 {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   echo "Running in save tensor file mode: $@"
   PT_XLA_DEBUG_LEVEL=2 PT_XLA_DEBUG_FILE="/tmp/pt_xla_debug.txt" run_test "$@"
 }
 
 function run_torchrun {
+  if ! test_is_selected "$1"; then
+    return
+  fi
   PJRT_DEVICE=NEURON torchrun --nnodes 1 --nproc-per-node 2 $@
 }
 
@@ -325,6 +346,8 @@ function run_tests {
   fi
 }
 
+set_test_filter $@
+
 if [ "$LOGFILE" != "" ]; then
   run_tests 2>&1 | tee $LOGFILE
 else

test/run_tests.sh

@@ -2,49 +2,17 @@
 set -exo pipefail
 
 # Absolute path to the directory of this script.
-CDIR="$(cd "$(dirname "$0")" ; pwd -P)"
+CDIR="$(
+  cd "$(dirname "$0")"
+  pwd -P
+)"
 
 # Import utilities.
 source "${CDIR}/utils/run_tests_utils.sh"
 
-# Default option values. Can be overridden via commandline flags.
-LOGFILE=/tmp/pytorch_py_test.log
-MAX_GRAPH_SIZE=500
-GRAPH_CHECK_FREQUENCY=100
-VERBOSITY=2
-
-# Parse commandline flags:
-#   -L
-#      disable writing to the log file at $LOGFILE.
-#   -M max_graph_size
-#   -C graph_check_frequency
-#   -V verbosity
-#   -h
-#      print the help string
-while getopts 'LM:C:V:h' OPTION
-do
-  case $OPTION in
-    L)
-      LOGFILE=
-      ;;
-    M)
-      MAX_GRAPH_SIZE=$OPTARG
-      ;;
-    C)
-      GRAPH_CHECK_FREQUENCY=$OPTARG
-      ;;
-    V)
-      VERBOSITY=$OPTARG
-      ;;
-    h)
-      echo -e "Usage: $0 TEST_FILTER...\nwhere TEST_FILTERs are globs match .py test files. If no test filter is provided, runs all tests."
-      exit 0
-      ;;
-    \?)  # This catches all invalid options.
-      echo "ERROR: Invalid commandline flag."
-      exit 1
-  esac
-done
+parse_options_to_vars $@
+
+# Consume the parsed commandline arguments.
 shift $(($OPTIND - 1))
 
 # Set the `CONTINUE_ON_ERROR` flag to `1` to make the CI tests continue on error.
@@ -62,35 +30,12 @@ export PYTORCH_TEST_WITH_SLOW=1
 export XLA_DUMP_FATAL_STACK=1
 export CPU_NUM_DEVICES=4
 
-TORCH_XLA_DIR=$(cd ~; dirname "$(python -c 'import torch_xla; print(torch_xla.__file__)')")
+TORCH_XLA_DIR=$(
+  cd ~
+  dirname "$(python -c 'import torch_xla; print(torch_xla.__file__)')"
+)
 COVERAGE_FILE="$CDIR/../.coverage"
 
-# Given $1 as a (possibly not normalized) test filepath, returns successfully
-# if it matches any of the space-separated globs $_TEST_FILTER. If
-# $_TEST_FILTER is empty, returns successfully. 
-function test_is_selected {
-  if [[ -z "$_TEST_FILTER" ]]; then
-    return 0  # success
-  fi
-
-  # _TEST_FILTER is a space-separate list of globs. Loop through the
-  # list elements.
-  for _FILTER in $_TEST_FILTER; do
-    # realpath normalizes the paths (e.g. resolving `..` and relative paths)
-    # so that they can be compared.
-    case `realpath $1` in
-      `realpath $_FILTER`)
-        return 0  # success
-        ;;
-      *)
-        # No match
-        ;;
-    esac
-  done
-
-  return 1  # failure
-}
-
 function run_coverage {
   if ! test_is_selected "$1"; then
     return
@@ -230,7 +175,7 @@ function run_xla_op_tests1 {
   run_dynamic "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
   run_dynamic "$CDIR/ds/test_dynamic_shapes.py"
   run_dynamic "$CDIR/ds/test_dynamic_shape_models.py" "$@" --verbosity=$VERBOSITY
-  run_eager_debug  "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
+  run_eager_debug "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
   run_test "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
   run_test "$CDIR/test_xla_graph_execution.py" "$@" --verbosity=$VERBOSITY
   run_pt_xla_debug_level2 "$CDIR/test_xla_graph_execution.py" "$@" --verbosity=$VERBOSITY
@@ -263,7 +208,7 @@ function run_xla_op_tests1 {
   run_test "$CDIR/test_python_ops.py"
   run_test "$CDIR/test_ops.py"
   run_test "$CDIR/test_metrics.py"
-  if [ -f "/tmp/metrics.txt" ] ; then
+  if [ -f "/tmp/metrics.txt" ]; then
     rm /tmp/metrics.txt
   fi
   XLA_METRICS_FILE=/tmp/metrics.txt run_test "$CDIR/test_metrics.py"
@@ -371,7 +316,7 @@ function run_xla_op_tests3 {
     # Please keep PJRT_DEVICE and GPU_NUM_DEVICES explicit in the following test commands.
     echo "single-host-single-process"
     PJRT_DEVICE=CUDA GPU_NUM_DEVICES=1 python3 test/test_train_mp_imagenet.py --fake_data --batch_size=16 --num_epochs=1 --num_cores=1 --num_steps=25 --model=resnet18
-    PJRT_DEVICE=CUDA torchrun --nnodes=1 --node_rank=0 --nproc_per_node=1 test/test_train_mp_imagenet.py --fake_data --pjrt_distributed --batch_size=16 --num_epochs=1  --num_steps=25 --model=resnet18
+    PJRT_DEVICE=CUDA torchrun --nnodes=1 --node_rank=0 --nproc_per_node=1 test/test_train_mp_imagenet.py --fake_data --pjrt_distributed --batch_size=16 --num_epochs=1 --num_steps=25 --model=resnet18
 
     echo "single-host-multi-process"
     num_devices=$(nvidia-smi --list-gpus | wc -l)
@@ -488,16 +433,7 @@ function run_tests {
   fi
 }
 
-if [[ $# -ge 1 ]]; then
-  # There are positional arguments - set $_TEST_FILTER to them.
-  _TEST_FILTER=$@
-  # Sometimes a test may fail even if it doesn't match _TEST_FILTER. Therefore,
-  # we need to set this to be able to get to the test(s) we want to run.
-  CONTINUE_ON_ERROR=1
-else
-  # No positional argument - run all tests.
-  _TEST_FILTER=""
-fi
+set_test_filter $@
 
 if [ "$LOGFILE" != "" ]; then
   run_tests 2>&1 | tee $LOGFILE