modify scripts to run tpu tests

Author: nitins17

.github/workflows/pytest_tpu.yml

@@ -0,0 +1,48 @@
+name: Run Pytest TPU tests
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        runner: ["linux-x86-ct5lp-224-8tpu"]
+        python: ["3.10"]
+
+    runs-on: ${{ matrix.runner }}
+    container:
+      image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
+
+    env:
+      # GitHub actions run in Docker by defaut. Disable running the `setup_docker.sh` script.
+      JAXCI_RUN_DOCKER_CONTAINER: 0
+      # Use RBE to build the artifacts.
+      JAXCI_BUILD_ARTIFACT_WITH_RBE: 1
+      # Setup the test environment (disable x64 mode and clone XLA at HEAD)
+      JAXCI_SETUP_TEST_ENVIRONMENT: 1
+      JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
+
+    steps:
+      - uses: actions/checkout@v3
+      # Halt for testing
+      - name: Wait For Connection
+        uses: ./actions/ci_connection/
+      - name: Build jaxlib
+        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jaxlib"
+      - name: Install pytest
+        env:
+          JAXCI_PYTHON: python${{ matrix.python }}
+        run: $JAXCI_PYTHON -m pip install pytest
+      - name: Install Test requirements
+        env:
+          JAXCI_PYTHON: python${{ matrix.python }}
+        run: | 
+          $JAXCI_PYTHON -m pip install -r build/test-requirements.txt
+          $JAXCI_PYTHON -m pip install -r build/collect-profile-requirements.txt
+      - name: Install Libtpu
+        run: $JAXCI_PYTHON -m pip install --pre libtpu-nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
+      - name: Run Pytest TPU tests
+        run: ./ci/run_pytest.sh "ci/envs/run_tests/pytest_tpu"

ci/run_pytest.sh

@@ -41,6 +41,8 @@ if [[ $JAXCI_RUN_PYTEST_TPU == 1 ]]; then
   echo "Running TPU tests..."
   # Run single-accelerator tests in parallel
   export JAX_ENABLE_TPU_XDIST=true
+  
+  check_if_to_run_in_docker "$JAXCI_PYTHON" -c 'import jax; print("libtpu version:",jax.lib.xla_bridge.get_backend().platform_version)'
   check_if_to_run_in_docker "$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
@@ -51,4 +53,4 @@ if [[ $JAXCI_RUN_PYTEST_TPU == 1 ]]; then
 
   # Run multi-accelerator across all chips
   check_if_to_run_in_docker "$JAXCI_PYTHON" -m pytest --tb=short --maxfail=20 -m "multiaccelerator" tests
-fi
+fi