Debug GPU failures

Author: nitins17

.github/workflows/bazel_gpu_local.yml

@@ -37,3 +37,5 @@ jobs:
         uses: ./actions/ci_connection/
       - name: Run Bazel GPU tests locally
         run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_gpu_local"
+      - name: Test if step runs
+        run: echo "This step ran!"

ci/run_bazel_test.sh

@@ -59,14 +59,16 @@ if [[ $JAXCI_RUN_BAZEL_TEST_GPU_LOCAL == 1 ]]; then
             --repo_env=HERMETIC_PYTHON_VERSION="$JAXCI_HERMETIC_PYTHON_VERSION" \
             --override_repository=xla="${JAXCI_XLA_GIT_DIR}" \
             --run_under "${JAXCI_JAX_GIT_DIR}/build/parallel_accelerator_execute.sh" \
-            //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests
+            //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests || true
+      echo "Finished running non-multiaccelerator tests..."
 
       # Runs multiaccelerator tests with all GPUs.
       check_if_to_run_in_docker bazel --bazelrc=ci/.bazelrc test --config=ci_linux_x86_64_cuda \
             --config=multiaccelerator_local \
             --repo_env=HERMETIC_PYTHON_VERSION="$JAXCI_HERMETIC_PYTHON_VERSION" \
             --override_repository=xla="${JAXCI_XLA_GIT_DIR}" \
-            //tests:gpu_tests //tests/pallas:gpu_tests
+            //tests:gpu_tests //tests/pallas:gpu_tests || true
+      echo "Finished running multiaccelerator tests..."
 fi
 
 # Run Bazel GPU tests with RBE.