Update test-pytorch-xla-tpu-tgi.yml

.github/workflows/test-pytorch-xla-tpu-tgi.yml

@@ -18,21 +18,25 @@ jobs:
   do-the-job:
     name: Run TGI tests 
     runs-on: tpu-runners-dind
-    container:
+    #container:
       # Use an image that works with TPU with Pytorch 2.3.0 (release was not working)
-      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:8f1dcd5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526
-      options: --shm-size "16gb" --ipc host --privileged
+      #image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla@sha256:8f1dcd5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526
+      #options: --shm-size "16gb" --ipc host --privileged
     env:
       PJRT_DEVICE: TPU
     steps:
-      - name: Checkout
-        uses: actions/checkout@v4
+      #- name: Checkout
+        #uses: actions/checkout@v4
       - name: slepp
-        run: sleep 10m
+        run: |
+          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+          pip install --upgrade pip
+          python -c 'import jax; print("TPU cores:", jax.device_count())'
+          python -c "import torch_xla.core.xla_model as xm; assert xm.xla_device().type == 'xla', 'XLA device not available'"
       #- name: Install dependencies
       #  run: pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
-      - name: Checking Pytorch/XLA installation
-        run: python -c "import torch_xla.core.xla_model as xm; assert xm.xla_device().type == 'xla', 'XLA device not available'"
+      #- name: Checking Pytorch/XLA installation
+      #  run: python -c "import torch_xla.core.xla_model as xm; assert xm.xla_device().type == 'xla', 'XLA device not available'"
 
       #- name: Build and test TGI server
       #  run: |