Fix for linux-x86-ct5lp-224-8tpu

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

Author: mydatascience

.github/workflows/tpu-tests.yml

@@ -31,16 +31,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  prelim:
-    runs-on: ["linux-x86-n2-32"]
-    steps:
-    - name: Test gsutil installation
-      run: which gsutil >/dev/null 2>&1 || { echo >&2 "gsutil is required but not installed. Aborting"; exit 24;}
-    - name: Cleanup old docker images
-      run: docker system prune --all --force
-
   tpu_unit_tests:
-    needs: prelim
     runs-on: [linux-x86-ct5lp-224-8tpu]
     container:
       image: python:3.12-slim
@@ -124,48 +115,6 @@ jobs:
         print('Llama3 params loaded successfully')
         "
 
-  tpu_integration_tests:
-    needs: prelim
-    runs-on: [self-hosted, linux-x86-ct5lp-224-8tpu]
-    container:
-      image: python:3.12-slim
-      options: --privileged --cpus=2 --memory=4Gi
-      env:
-        TPU_ACCELERATOR_TYPE: ""
-        JAX_PLATFORMS: "cpu"
-        CUDA_VISIBLE_DEVICES: ""
-        TF_CPP_MIN_LOG_LEVEL: "3"
-    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 .
-        pip install pytest pytest-xdist
-    
-    - name: Run integration tests
-      run: |
-        # Run more comprehensive tests that might take longer
-        python -m pytest tests/ -v --tb=short -m "integration_test" --timeout=300
-
   notify_failure:
     name: Notify failed build
     needs: [tpu_unit_tests, tpu_integration_tests]