Reuse run_torch_xla_tests from pytorch/xla (#59888)

Ailing Zhang · facebook-github-bot · commit a0e62c4da4ba · 2021-06-15T13:00:09.000-07:00
Summary: Pull Request resolved: #59888 Test Plan: Imported from OSS Reviewed By: samestep Differential Revision: D29114274 Pulled By: ailzhang fbshipit-source-id: d2845c7fc95d038cd68c10e22b68be8ad3cae736
diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh
@@ -37,8 +37,6 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
   # mainly used so that we're not spending extra cycles testing cpu
   # devices on expensive gpu machines
   export PYTORCH_TESTING_DEVICE_ONLY_FOR="cuda"
-elif [[ "$BUILD_ENVIRONMENT" == *xla* ]]; then
-  export PYTORCH_TESTING_DEVICE_ONLY_FOR="xla"
 fi
 
 if [[ "$BUILD_ENVIRONMENT" == *cuda11* ]]; then
@@ -339,23 +337,9 @@ test_torch_function_benchmark() {
 }
 
 test_xla() {
-  export XLA_USE_XRT=1 XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0"
-  # Issue #30717: randomize the port of XLA/gRPC workers is listening on to reduce flaky tests.
-  XLA_PORT=$(shuf -i 40701-40999 -n 1)
-  export XRT_WORKERS="localservice:0;grpc://localhost:$XLA_PORT"
-  pushd xla
-  echo "Running Python Tests"
-  ./test/run_tests.sh
-
-  # Disabled due to MNIST download issue.
-  # See https://github.com/pytorch/pytorch/issues/53267
-  # echo "Running MNIST Test"
-  # python test/test_train_mnist.py --tidy
-
-  echo "Running C++ Tests"
-  pushd test/cpp
-  CC=clang-9 CXX=clang++-9 ./run_tests.sh
-  popd
+  # shellcheck disable=SC1091
+  source "./xla/.circleci/common.sh"
+  run_torch_xla_tests "$(pwd)" "$(pwd)/xla"
   assert_git_not_dirty
 }
 
