From a2844e5dd8fae7203953c9db74cd3cb92a27150f Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Wed, 28 Feb 2024 12:24:28 +0000
Subject: [PATCH] WIP

---
 .github/actions/pytorch-xla-tpu/action.yml    |  2 +
 .../workflows/test-pytorch-xla-tpu-tgi.yml    | 46 +++++++++++--------
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/.github/actions/pytorch-xla-tpu/action.yml b/.github/actions/pytorch-xla-tpu/action.yml
index 8653756e..856721ef 100644
--- a/.github/actions/pytorch-xla-tpu/action.yml
+++ b/.github/actions/pytorch-xla-tpu/action.yml
@@ -8,5 +8,7 @@ inputs:
 runs:
   using: 'docker'
   image: 'Dockerfile'
+  container:
+    options: --shm-size "16gb" --ipc host --privileged
   args:
     - ${{ inputs.make-target }}
diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi.yml b/.github/workflows/test-pytorch-xla-tpu-tgi.yml
index f089bf35..a23773d6 100644
--- a/.github/workflows/test-pytorch-xla-tpu-tgi.yml
+++ b/.github/workflows/test-pytorch-xla-tpu-tgi.yml
@@ -27,23 +27,29 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
-      - name: Create venv and install Pytorch/XLA
-        run: |
-          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 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 venv-tpu-pytorch/bin/activate
-          make tgi_test
+
+      - name: Containerized Build and Test TGI
+        uses: ./.github/actions/pytorch-xla-tpu
+        with:
+          make-target: tgi_test
+
+      # - name: Create venv and install Pytorch/XLA
+      #   run: |
+      #     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 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 venv-tpu-pytorch/bin/activate
+      #     make tgi_test