From 0c946642b0c395a2459701176808774b201b0603 Mon Sep 17 00:00:00 2001
From: Mason Chang <mason@masonchang.com>
Date: Thu, 1 Feb 2024 10:38:07 -0800
Subject: [PATCH] Revert "Revert: #6368 to fix nightly CI coverage run.
 (#6369)"

This reverts commit 48a1ae3dbb87b44c84503f75f105ff1aa0656070.
---
 .circleci/common.sh                  | 2 +-
 .github/workflows/_test.yml          | 8 +-------
 .github/workflows/build_and_test.yml | 2 +-
 test/pytorch_test_base.py            | 2 +-
 4 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/.circleci/common.sh b/.circleci/common.sh
index 4ffc01c69dd..ad0a1378e8e 100755
--- a/.circleci/common.sh
+++ b/.circleci/common.sh
@@ -221,7 +221,7 @@ function run_torch_xla_tests() {
   export CXX_ABI=$(python -c "import torch;print(int(torch._C._GLIBCXX_USE_CXX11_ABI))")
 
   # TODO(yeounoh) test coverage workflow is not parallelized.
-  if [[ -z "$RUN_BENCHMARK_TESTS" && -z "$RUN_CPP_TESTS1" && -z "$RUN_CPP_TESTS2" && -z "$RUN_PYTHON_TESTS" ]]; then
+  if [[ -z "$RUN_BENCHMARK_TESTS" && -z "$RUN_CPP_TESTS1" && -z "$RUN_CPP_TESTS2" && -z "$RUN_PYTHON_TESTS" || "$USE_COVERAGE" != "0" ]]; then
     run_torch_xla_python_tests $PYTORCH_DIR $XLA_DIR $USE_COVERAGE
     run_torch_xla_cpp_tests $PYTORCH_DIR $XLA_DIR $USE_COVERAGE
     run_torch_xla_benchmark_tests $XLA_DIR
diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml
index b0119ecf1d6..9f7e6d468d6 100644
--- a/.github/workflows/_test.yml
+++ b/.github/workflows/_test.yml
@@ -125,15 +125,9 @@ jobs:
         env:
           CIRCLE_WORKFLOW_ID: ${{ github.run_id }}
           CIRCLE_BUILD_NUM: ${{ github.run_number }}
-          BENCHMARK_TEST_NAME: ${{ env.RUN_BENCHMARK_TESTS }}
-          PYTHON_TEST_NAME: ${{ env.RUN_PYTHON_TESTS }}${{ env.RUN_XLA_OP_TESTS1 }}${{ env.RUN_XLA_OP_TESTS2 }}${{ env.RUN_XLA_OP_TESTS3 }}${{ env.RUN_TORCH_MP_OP_TESTS }}
+          PYTHON_TEST_NAME: ${{ env.RUN_BENCHMARK_TESTS }}${{ env.RUN_PYTHON_TESTS }}${{ env.RUN_XLA_OP_TESTS1 }}${{ env.RUN_XLA_OP_TESTS2 }}${{ env.RUN_XLA_OP_TESTS3 }}${{ env.RUN_TORCH_MP_OP_TESTS }}
           CPP_TEST_NAME: ${{ env.RUN_CPP_TESTS1 }}${{ env.RUN_CPP_TESTS2 }}
         run: |
-            # TODO(yeounoh) collect coverage report as needed.
-            if [ -n "${BENCHMARK_TEST_NAME}" ] || [ "${PYTHON_TEST_NAME}" == "pythontorch_mp_op" ]; then
-                exit 0
-            fi
-
             docker cp "${pid}":/home/jenkins/htmlcov "${GITHUB_WORKSPACE}"
             if [ -n "${GPU_FLAG:-}" ]; then
               if [ -n "${PYTHON_TEST_NAME}" ]; then
diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
index bd1f9be5b0e..a3496479268 100644
--- a/.github/workflows/build_and_test.yml
+++ b/.github/workflows/build_and_test.yml
@@ -44,7 +44,7 @@ jobs:
       docker-image: ${{ needs.build.outputs.docker-image }}
       runner: linux.8xlarge.nvidia.gpu
       timeout-minutes: 300
-      collect-coverage: ${{ github.event_name == 'push' && github.event.ref == 'refs/heads/master' }}
+      collect-coverage: false #${{ github.event_name == 'push' && github.event.ref == 'refs/heads/master' }}
       disable-xrt: 1
     secrets:
       gcloud-service-key: ${{ secrets.GCLOUD_SERVICE_KEY }}
diff --git a/test/pytorch_test_base.py b/test/pytorch_test_base.py
index 4101d703d0d..64c8e41e4f4 100644
--- a/test/pytorch_test_base.py
+++ b/test/pytorch_test_base.py
@@ -19,7 +19,6 @@
     'test_sum_xla_bfloat16': 0.1,
     'test_put_xla_bfloat16': 0.05,
     'test_take_xla_bfloat16': 0.05,
-    'test_EmbeddingBag_per_sample_weights_and_new_offsets_xla': 0.01,
 }
 
 DISABLED_TORCH_TESTS_ANY = {
@@ -316,6 +315,7 @@
         'test_embedding_bag_device',  # FIXME! Unsupported device type for sparse layout: xla
         'test_embedding_scalar_weight_error_xla',  # tsl::CurrentStackTrace[abi:cxx11]
         'test_EmbeddingBag_per_sample_weights_and_no_offsets',  # FIXME! Unsupported device type for sparse layout: xla
+        'test_EmbeddingBag_per_sample_weights_and_new_offsets',  # precision
     },
 
     # test/nn/test_convolution.py