rework tpu job and scripts to match upstream

Author: nitins17

.github/workflows/pytest_tpu.yml

@@ -1,9 +1,9 @@
 name: Run Pytest TPU tests
 
 on:
-  # pull_request:
-  #   branches:
-  #     - main
+  pull_request:
+    branches:
+      - main
   workflow_dispatch:
     inputs:
       halt-for-connection:
@@ -16,19 +16,26 @@ on:
         - 'no'
 
 jobs:
-  run_tests:
+  run_tpu_tests:
     strategy:
+      fail-fast: false
       matrix:
-        runner: ["linux-x86-ct5lp-224-8tpu"]
-        tpu_cores: ["8"]
+        jaxlib-version: ["pypi_latest", "nightly", "nightly+oldest_supported_libtpu"]
+        tpu: [
+          # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
+          {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
+          {type: "v5e-8", cores: "8", 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"
+    runs-on: ${{ matrix.tpu.runner }}
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
+
+    name: "TPU test (jaxlib=${{ matrix.jaxlib-version }}, ${{ matrix.tpu.type }})"
 
     env:
-      JAXCI_CLONE_MAIN_XLA: 1
+      LIBTPU_OLDEST_VERSION_DATE: 20240722
+      ENABLE_PJRT_COMPATIBILITY: ${{ matrix.jaxlib-version == 'nightly+oldest_supported_libtpu' }}
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
 
     steps:
@@ -38,23 +45,33 @@ jobs:
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Build jaxlib
-        run: ./ci/build_artifacts.sh "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
+      - name: Install JAX test requirements
         env:
           JAXCI_PYTHON: python${{ matrix.python }}
-        run: $JAXCI_PYTHON -m pip install --pre libtpu-nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
+        run: |
+          pip install -U -r build/test-requirements.txt
+          pip install -U -r build/collect-profile-requirements.txt
+      - name: Install JAX
+        run: |
+          if [ "${{ matrix.jaxlib-version }}" == "pypi_latest" ]; then
+            pip install .[tpu] \
+              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+          elif [ "${{ matrix.jaxlib-version }}" == "nightly" ]; then
+            pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+            pip install --pre libtpu \
+              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+            pip install requests
+          elif [ "${{ matrix.jaxlib-version }}" == "nightly+oldest_supported_libtpu" ]; then
+            pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+            # TODO(phawkins): switch to libtpu, when the oldest release we support is a libtpu release.
+            pip install --pre libtpu-nightly==0.1.dev${{ env.LIBTPU_OLDEST_VERSION_DATE }} \
+              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+            pip install requests
+          else
+            echo "Unknown jaxlib-version: ${{ matrix.jaxlib-version }}"
+            exit 1
+          fi
       - name: Run Pytest TPU tests
         env:
-          JAXCI_TPU_CORES: ${{ matrix.tpu_cores }}
+          JAXCI_TPU_CORES: ${{ matrix.tpu.cores }}
         run: ./ci/run_pytest_tpu.sh

ci/run_pytest_tpu.sh

@@ -26,37 +26,30 @@ set -exu -o history -o allexport
 # Inherit default JAXCI environment variables.
 source ci/envs/default.env
 
-# Install jaxlib wheel on the system. Requires a jaxlib wheel to be present
-# inside $JAXCI_OUTPUT_DIR (../dist)
-echo "Installing wheels locally..."
-source ./ci/utilities/install_wheels_locally.sh
-
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
-export PY_COLORS=1
-export JAX_SKIP_SLOW_TESTS=true
-
 "$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))"
-
 "$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 /usr/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..."
+# Set up common test environment variables
+export PY_COLORS=1
+export JAX_SKIP_SLOW_TESTS=true
 export JAX_PLATFORMS=tpu,cpu
-# Run single-accelerator tests in parallel
-export JAX_ENABLE_TPU_XDIST=true
+# End of common test environment variable setup
 
-"$JAXCI_PYTHON" -m pytest -n="$JAXCI_TPU_CORES" --tb=short \
+echo "Running TPU tests..."
+# Run single-accelerator tests in parallel
+JAX_ENABLE_TPU_XDIST=true "$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
 
 # Run Pallas printing tests, which need to run with I/O capturing disabled.
-export TPU_STDERR_LOG_LEVEL=0
-"$JAXCI_PYTHON" -m pytest -s tests/pallas/tpu_pallas_test.py::PallasCallPrintTest
+TPU_STDERR_LOG_LEVEL=0 "$JAXCI_PYTHON" -m pytest -s tests/pallas/tpu_pallas_test.py::PallasCallPrintTest
 
 # Run multi-accelerator across all chips
 "$JAXCI_PYTHON" -m pytest --tb=short --maxfail=20 -m "multiaccelerator" tests