WIP

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

@@ -2,9 +2,7 @@ name: Optimum TPU / Test TGI on TPU
 
 on:
   push:
-    branches: [ main ]
-    paths:
-      - "text-generation-inference/**"
+    branches: [ quick-ci-test ]
   pull_request:
     branches: [ main ]
     paths:
@@ -31,13 +29,3 @@ jobs:
       - name: Build and test TGI server
         run: |
           HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_TPU_CI }} make tgi_test
-
-      # Use a different step to test the Jetstream Pytorch version, to avoid conflicts with torch-xla[tpu]
-      - name: Install and test TGI server (Jetstream Pytorch)
-        run: |
-          pip install -U .[jetstream-pt] \
-            -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-            -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html \
-            -f https://storage.googleapis.com/libtpu-releases/index.html
-          JETSTREAM_PT=1 HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_TPU_CI }} python -m \
-            pytest -sv text-generation-inference/tests -k jetstream

Makefile

@@ -90,7 +90,7 @@ tgi_server:
 tgi_test: test_installs tgi_server
 	find text-generation-inference -name "text_generation_server-$(VERSION)-py3-none-any.whl" \
 	                               -exec python -m pip install --force-reinstall {} \;
-	python -m pytest -sv text-generation-inference/tests
+	python -m pytest -sv text-generation-inference/tests -k gemma-2b
 
 tgi_docker_test: tpu-tgi
 	python -m pip install -r text-generation-inference/integration-tests/requirements.txt