revert the repro change

pytorch · Feb 6, 2024 · cba8e5d · cba8e5d
1 parent a2c4167
commit cba8e5d
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 21 deletions.
diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml
@@ -50,34 +50,34 @@ jobs:
       matrix:
         include:
           # Use readable strings as they define the workflow titles.
-          # - run_benchmark_tests: 'benchmark_tests'
-          # - run_cpp_tests1: 'cpp_tests1'
-          # - run_cpp_tests2: 'cpp_tests2'
+          - run_benchmark_tests: 'benchmark_tests'
+          - run_cpp_tests1: 'cpp_tests1'
+          - run_cpp_tests2: 'cpp_tests2'
           - run_python_tests: 'python_tests'
             run_xla_op_tests1: 'xla_op1'
-          # - run_python_tests: 'python_tests'
-          #   run_xla_op_tests2: 'xla_op2'
-          # - run_python_tests: 'python_tests'
-          #   run_xla_op_tests3: 'xla_op3'
-          # - run_python_tests: 'python'
-          #   run_torch_mp_op_tests: 'torch_mp_op'
+          - run_python_tests: 'python_tests'
+            run_xla_op_tests2: 'xla_op2'
+          - run_python_tests: 'python_tests'
+            run_xla_op_tests3: 'xla_op3'
+          - run_python_tests: 'python'
+            run_torch_mp_op_tests: 'torch_mp_op'
     timeout-minutes: ${{ inputs.timeout-minutes }}
     env:
       DOCKER_IMAGE: ${{ inputs.docker-image }}
       WORKDIR: /var/lib/jenkins/workspace
       GCLOUD_SERVICE_KEY: ${{ secrets.gcloud-service-key }}
       USE_COVERAGE: ${{ inputs.collect-coverage && '1' || '0' }}
-      # XLA_SKIP_XRT_TESTS: ${{ inputs.disable-xrt }}
-      # XLA_SKIP_TORCH_OP_TESTS: ${{ inputs.disable-pjrt }}
-      # XLA_SKIP_MP_OP_TESTS: ${{ inputs.disable-pjrt }}
-      # RUN_BENCHMARK_TESTS: ${{ matrix.run_benchmark_tests }}
-      # RUN_CPP_TESTS1: ${{ matrix.run_cpp_tests1 }}
-      # RUN_CPP_TESTS2: ${{ matrix.run_cpp_tests2 }}
+      XLA_SKIP_XRT_TESTS: ${{ inputs.disable-xrt }}
+      XLA_SKIP_TORCH_OP_TESTS: ${{ inputs.disable-pjrt }}
+      XLA_SKIP_MP_OP_TESTS: ${{ inputs.disable-pjrt }}
+      RUN_BENCHMARK_TESTS: ${{ matrix.run_benchmark_tests }}
+      RUN_CPP_TESTS1: ${{ matrix.run_cpp_tests1 }}
+      RUN_CPP_TESTS2: ${{ matrix.run_cpp_tests2 }}
       RUN_PYTHON_TESTS: ${{ matrix.run_python_tests }}
       RUN_XLA_OP_TESTS1: ${{ matrix.run_xla_op_tests1 }}
-      # RUN_XLA_OP_TESTS2: ${{ matrix.run_xla_op_tests2 }}
-      # RUN_XLA_OP_TESTS3: ${{ matrix.run_xla_op_tests3 }}
-      # RUN_TORCH_MP_OP_TESTS: ${{ matrix.run_torch_mp_op_tests }}
+      RUN_XLA_OP_TESTS2: ${{ matrix.run_xla_op_tests2 }}
+      RUN_XLA_OP_TESTS3: ${{ matrix.run_xla_op_tests3 }}
+      RUN_TORCH_MP_OP_TESTS: ${{ matrix.run_torch_mp_op_tests }}
     steps:
       - name: Setup Linux
         uses: pytorch/test-infra/.github/actions/setup-linux@main
@@ -112,7 +112,7 @@ jobs:
         run: |
           echo "DOCKER_IMAGE: ${DOCKER_IMAGE}"
           docker pull "${DOCKER_IMAGE}"
-          pid=$(docker run --shm-size=16g ${GPU_FLAG:-} -e USE_COVERAGE -e RUN_PYTHON_TESTS -e RUN_XLA_OP_TESTS1 -t -d -w "$WORKDIR" "${DOCKER_IMAGE}")
+          pid=$(docker run --shm-size=16g ${GPU_FLAG:-} -e USE_COVERAGE -e XLA_SKIP_XRT_TESTS -e XLA_SKIP_TORCH_OP_TESTS -e XLA_SKIP_MP_OP_TESTS -e RUN_BENCHMARK_TESTS -e RUN_CPP_TESTS1 -e RUN_CPP_TESTS2 -e RUN_PYTHON_TESTS -e RUN_XLA_OP_TESTS1 -e RUN_XLA_OP_TESTS2 -e RUN_XLA_OP_TESTS3 -e RUN_TORCH_MP_OP_TESTS -t -d -w "$WORKDIR" "${DOCKER_IMAGE}")
           echo "${GCLOUD_SERVICE_KEY}" | docker exec -i "${pid}" sh -c "cat >> /tmp/pytorch/xla/default_credentials.json"
           echo "pid=${pid}" >> "${GITHUB_ENV}"
       - name: Test

diff --git a/test/run_tests.sh b/test/run_tests.sh
@@ -154,10 +154,38 @@ function run_torch_op_tests {
 
 # DO NOT MODIFY
 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_test "$CDIR/test_grad_checkpoint.py"
+  run_test "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
+  run_test_without_functionalization "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
+  run_pt_xla_debug "$CDIR/debug_tool/test_pt_xla_debug.py"
+  run_test "$CDIR/test_async_closures.py"
+  run_test "$CDIR/test_hlo_metadata.py"
+  run_test "$CDIR/test_profiler.py"
   run_test "$CDIR/pjrt/test_runtime.py"
   run_test "$CDIR/pjrt/test_runtime_single_proc_gpu.py"
   run_test "$CDIR/pjrt/test_runtime_multi_gpu.py"
   run_test "$CDIR/pjrt/test_runtime_multi_cpu.py"
+  run_test "$CDIR/pjrt/test_internal_tpu.py"
+  run_test "$CDIR/pjrt/test_ddp.py"
+  run_test "$CDIR/pjrt/test_mesh_service.py"
+  run_test "$CDIR/test_ops.py"
+  run_test "$CDIR/test_metrics.py"
+  run_test "$CDIR/test_zero1.py"
+  run_test "$CDIR/dynamo/test_dynamo_integrations_util.py"
+  run_test "$CDIR/dynamo/test_dynamo.py"
+  run_test "$CDIR/dynamo/test_bridge.py"
+  run_test "$CDIR/dynamo/test_num_output.py"
+  run_save_tensor_ir "$CDIR/dynamo/test_dynamo_graph_dump.py"
+  run_use_bf16 "$CDIR/test_data_type.py"
+  run_xla_ir_debug "$CDIR/test_env_var_mapper.py"
+  run_xla_hlo_debug "$CDIR/test_env_var_mapper.py"
+  run_xla_hlo_debug "$CDIR/stablehlo/test_stablehlo_save_load.py"
+  run_save_tensor_ir "$CDIR/spmd/test_spmd_graph_dump.py"
+  run_save_tensor_hlo "$CDIR/spmd/test_spmd_graph_dump.py"
 }
 
 # DO NOT MODIFY

diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py
@@ -148,8 +148,6 @@ def _setup_tpu_vm_library_path() -> bool:
 
 def _prepare_to_exit():
   try:
-    device = _XLAC._xla_get_default_device()
-    _XLAC._set_all_reduce_token(device, None)
     _XLAC._prepare_to_exit()
     if int(os.environ.get('PT_XLA_DEBUG', '0')):
       _summarize_fn_tracker()