diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi-jetstream.yml b/.github/workflows/test-pytorch-xla-tpu-tgi-jetstream.yml index 8640ea8c..e669baf5 100644 --- a/.github/workflows/test-pytorch-xla-tpu-tgi-jetstream.yml +++ b/.github/workflows/test-pytorch-xla-tpu-tgi-jetstream.yml @@ -1,14 +1,12 @@ name: Optimum TPU / Test TGI on TPU / Jetstream Pytorch on: - push: - branches: [ main ] - paths: - - "text-generation-inference/**" pull_request: branches: [ main ] paths: - "text-generation-inference/**" + # This can be used to trigger workflow from the web interface + workflow_dispatch: concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} @@ -17,10 +15,11 @@ concurrency: jobs: do-the-job: name: Run TGI tests - Jetstream Pytorch - runs-on: optimum-tpu + runs-on: + group: gcp-ct5lp-hightpu-8t container: image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm - options: --shm-size "16gb" --ipc host --privileged + options: --shm-size "16gb" --ipc host --privileged ${{ vars.V5_LITEPOD_8_ENV}} env: PJRT_DEVICE: TPU steps: diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi-nightly-jetstream.yml b/.github/workflows/test-pytorch-xla-tpu-tgi-nightly-jetstream.yml index da476e76..f85dc1c0 100644 --- a/.github/workflows/test-pytorch-xla-tpu-tgi-nightly-jetstream.yml +++ b/.github/workflows/test-pytorch-xla-tpu-tgi-nightly-jetstream.yml @@ -16,11 +16,11 @@ concurrency: jobs: do-the-job: name: Build and Run slow tests - runs-on: optimum-tpu + runs-on: + group: gcp-ct5lp-hightpu-8t container: - # Use a nightly image that works with TPU (release was not working) image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm - options: --shm-size "16gb" --ipc host --privileged + options: --shm-size "16gb" --ipc host --privileged ${{ vars.V5_LITEPOD_8_ENV}} env: PJRT_DEVICE: TPU HF_TOKEN: ${{ secrets.HF_TOKEN_OPTIMUM_TPU_CI }} diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi-nightly.yml b/.github/workflows/test-pytorch-xla-tpu-tgi-nightly.yml index 8a704d41..b294442f 100644 --- a/.github/workflows/test-pytorch-xla-tpu-tgi-nightly.yml +++ b/.github/workflows/test-pytorch-xla-tpu-tgi-nightly.yml @@ -14,11 +14,11 @@ concurrency: jobs: do-the-job: name: Build and Run slow tests - runs-on: optimum-tpu + runs-on: + group: gcp-ct5lp-hightpu-8t container: - # Use a nightly image that works with TPU (release was not working) image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm - options: --shm-size "16gb" --ipc host --privileged + options: --shm-size "16gb" --ipc host --privileged ${{ vars.V5_LITEPOD_8_ENV}} env: PJRT_DEVICE: TPU HF_TOKEN: ${{ secrets.HF_TOKEN_OPTIMUM_TPU_CI }} diff --git a/.github/workflows/test-pytorch-xla-tpu-tgi.yml b/.github/workflows/test-pytorch-xla-tpu-tgi.yml index 86873328..d26b5016 100644 --- a/.github/workflows/test-pytorch-xla-tpu-tgi.yml +++ b/.github/workflows/test-pytorch-xla-tpu-tgi.yml @@ -17,10 +17,11 @@ concurrency: jobs: do-the-job: name: Run TGI tests - runs-on: optimum-tpu + runs-on: + group: gcp-ct5lp-hightpu-8t container: image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm - options: --shm-size "16gb" --ipc host --privileged + options: --shm-size "16gb" --ipc host --privileged ${{ vars.V5_LITEPOD_8_ENV}} env: PJRT_DEVICE: TPU steps: diff --git a/.github/workflows/test-pytorch-xla-tpu.yml b/.github/workflows/test-pytorch-xla-tpu.yml index 6a389efe..5585679d 100644 --- a/.github/workflows/test-pytorch-xla-tpu.yml +++ b/.github/workflows/test-pytorch-xla-tpu.yml @@ -17,11 +17,11 @@ concurrency: jobs: do-the-job: name: Run optimum tpu tests - runs-on: optimum-tpu + runs-on: + group: gcp-ct5lp-hightpu-8t container: - # Use a nightly image that works with TPU (release was not working) image: us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.4.0_3.10_tpuvm - options: --shm-size "16gb" --ipc host --privileged + options: --shm-size "16gb" --ipc host --privileged ${{ vars.V5_LITEPOD_8_ENV}} env: PJRT_DEVICE: TPU steps: