test-pytorch-xla-tpu-tgi.yml

name: Optimum TPU / Test TGI on TPU

on:
  push:
    branches: [ main, enable-ci-cd ]
    paths:
      - "text-generation-inference/**"
      - ".github/**"
  pull_request:
    branches: [ main ]
    paths:
      - "text-generation-inference/**"

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
  cancel-in-progress: true

jobs:
  do-the-job:
    name: Run TGI tests
    runs-on: optimum-tpu
    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
      # image: ubuntu:22.04
      options: --shm-size "16gb" --ipc host --privileged
    env:
      PJRT_DEVICE: TPU
    steps:
      - name: Checkout
        uses: actions/checkout@v4

      - name: Test env
        run: |
          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: 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