From 931ade395d449b5a81b0b7d77b407e1d22b4fff8 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Tue, 27 Feb 2024 15:00:35 +0000
Subject: [PATCH] WIP test

---
 .github/workflows/test_tpu_tgi.yml | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/test_tpu_tgi.yml b/.github/workflows/test_tpu_tgi.yml
index 1ee83a15..ab06082b 100644
--- a/.github/workflows/test_tpu_tgi.yml
+++ b/.github/workflows/test_tpu_tgi.yml
@@ -2,9 +2,10 @@ name: Optimum TPU / Test TGI on TPU
 
 on:
   push:
-    branches: [ main ]
+    branches: [ main, enable-ci-cd ]
     paths:
       - "text-generation-inference/**"
+      - ".github/**"
   pull_request:
     branches: [ main ]
     paths:
@@ -17,17 +18,32 @@ concurrency:
 jobs:
   do-the-job:
     name: Run TGI tests
-    runs-on: [optimum-tpu] # run the configured Google TPU
+    runs-on: optimum-tpu
+    container:
+      image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.2.0_3.10_tpuvm
+      options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ --privileged
+    env:
+      PJRT_DEVICE: TPU
     steps:
       - name: Checkout
         uses: actions/checkout@v2
       - name: Create venv and install Pytorch/XLA
         run: |
-          python3 -m venv optimum-tpu-pytorch
-          source optimum-tpu-pytorch/bin/activate
+          echo "PJRT_DEVICE set to $PJRT_DEVICE"
+          pwd
+          ls
+          cat /etc/issue
+          echo "-----"
+          # sudo apt install python3.10-venv -y
+          # echo "Creating virtual environment..."
+          # python3 -m venv venv-tpu-pytorch
+          # source venv-tpu-pytorch/bin/activate
+          echo "Installing basic packages"
           python -m pip install -U pip
-          python -m pip install "torch~=2.2.0" "torch_xla[tpu]~=2.2.0" -f https://storage.googleapis.com/libtpu-releases/index.html
+          # python -m pip install "torch~=2.2.0" "torch_xla[tpu]~=2.2.0" -f https://storage.googleapis.com/libtpu-releases/index.html numpy
+          echo "Testing XLA installation..."
+          python -c "import torch_xla.core.xla_model as xm; assert xm.xla_device().type == 'xla', 'XLA device not available'"
       - name: Run TGI server python tests
         run: |
-          source optimum-tpu-pytorch/bin/activate
+          # source venv-tpu-pytorch/bin/activate
           make tgi_test