TPU simple tests

Signed-off-by: Vladimir Suvorov <suvorovv@google.com>

Author: mydatascience

.github/workflows/tpu-tests.yml

@@ -34,34 +34,17 @@ jobs:
   tpu_unit_tests:
     runs-on: [linux-x86-ct5lp-224-8tpu]
     container:
-      image: python:3.12-slim
-      options: --privileged --cpus=2 --memory=4Gi
+      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
+      options: --privileged
       env:
-        TPU_ACCELERATOR_TYPE: ""
-        JAX_PLATFORMS: "cpu"
-        CUDA_VISIBLE_DEVICES: ""
-        TF_CPP_MIN_LOG_LEVEL: "3"
+        CLOUD_TPU_ACCELERATOR: v5e-8
+        JAX_PLATFORMS: tpu
     steps:
     - name: Checkout code
       uses: actions/checkout@v4
       with:
         fetch-depth: 0
 
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.12'
-    
-    - name: Install system dependencies
-      run: |
-        sudo apt-get update
-        sudo apt-get install -y git curl
-    
-    - name: Set up JAX for TPU
-      run: |
-        pip install --upgrade pip
-        pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-    
     - name: Install tunix dependencies
       run: |
         pip install -e .
@@ -115,9 +98,10 @@ jobs:
         print('Llama3 params loaded successfully')
         "
 
+
   notify_failure:
     name: Notify failed build
-    needs: [tpu_unit_tests]
+    needs: [tpu_unit_tests, tpu_integration_tests]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:
@@ -144,7 +128,7 @@ jobs:
   notify_success_and_close:
     name: Close issue after 3 successful builds
     if: ${{ success() && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }}
-    needs: [tpu_unit_tests]
+    needs: [tpu_unit_tests, tpu_integration_tests]
     runs-on: ubuntu-latest
     permissions:
       issues: write