Update GPU/TPU test scripts

Author: nitins17

.github/workflows/pytest_tpu.yml

@@ -6,7 +6,7 @@ on:
       - main
 
 jobs:
-  build:
+  run_tests:
     strategy:
       matrix:
         runner: ["linux-x86-ct5lp-224-8tpu"]
@@ -43,4 +43,7 @@ jobs:
           JAXCI_PYTHON: python${{ matrix.python }}
         run: $JAXCI_PYTHON -m pip install --pre libtpu-nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
       - name: Run Pytest TPU tests
+        env:
+          JAX_PLATFORMS: tpu,cpu
+          PY_COLORS: 1
         run: ./ci/run_pytest.sh "ci/envs/run_tests/pytest_tpu.env"

ci/run_bazel_test.sh

@@ -68,13 +68,13 @@ if [[ $JAXCI_RUN_BAZEL_TEST_GPU_RBE == 1 ]]; then
             //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests
 fi
 
-# Run Bazel GPU tests locally.
+# Run Non-RBE Bazel GPU tests (single accelerator and multiaccelerator tetsts).
 if [[ $JAXCI_RUN_BAZEL_TEST_GPU_NON_RBE == 1 ]]; then
       export NCCL_DEBUG=WARN
       nvidia-smi
-      echo "Running local GPU tests..."
+      echo "Running (non-RBE) GPU tests..."
 
-      # Runs non-multiaccelerator tests with one GPU apiece.
+      # Runs single accelerator tests with one GPU apiece.
       # It appears --run_under needs an absolute path.
       # The product of the `JAX_ACCELERATOR_COUNT`` and `JAX_TESTS_PER_ACCELERATOR`
       # should match the VM's CPU core count (set in `--local_test_jobs`).

ci/run_pytest.sh

@@ -29,6 +29,10 @@ if [[ $JAXCI_RUN_PYTEST_CPU == 1 ]]; then
 fi
 
 if [[ $JAXCI_RUN_PYTEST_GPU == 1 ]]; then
+  nvidia-smi
+  export NCCL_DEBUG=WARN
+  export TF_CPP_MIN_LOG_LEVEL=0
+
   echo "Running GPU tests..."
   export XLA_PYTHON_CLIENT_ALLOCATOR=platform
   export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1
@@ -41,11 +45,17 @@ if [[ $JAXCI_RUN_PYTEST_GPU == 1 ]]; then
 fi
 
 if [[ $JAXCI_RUN_PYTEST_TPU == 1 ]]; then
+
+  "$JAXCI_PYTHON" -c 'import sys; print("python version:", sys.version)'
+  "$JAXCI_PYTHON" -c 'import jax; print("jax version:", jax.__version__)'
+  "$JAXCI_PYTHON" -c 'import jaxlib; print("jaxlib version:", jaxlib.__version__)'
+  strings $HOME/.local/lib/"$JAXCI_PYTHON"/site-packages/libtpu/libtpu.so | grep 'Built on'
+  "$JAXCI_PYTHON" -c 'import jax; print("libtpu version:",jax.lib.xla_bridge.get_backend().platform_version)'
+
   echo "Running TPU tests..."
   # Run single-accelerator tests in parallel
   export JAX_ENABLE_TPU_XDIST=true
 
-  "$JAXCI_PYTHON" -c 'import jax; print("libtpu version:",jax.lib.xla_bridge.get_backend().platform_version)'
   "$JAXCI_PYTHON" -m pytest -n="$JAXCI_TPU_CORES" --tb=short \
     --deselect=tests/pallas/tpu_pallas_test.py::PallasCallPrintTest \
     --maxfail=20 -m "not multiaccelerator" tests examples