From 20ce7cc1eaf9350050a8ace0c7771bad8af60012 Mon Sep 17 00:00:00 2001
From: Michael Hudgins <michaelhudgins@google.com>
Date: Mon, 24 Jun 2024 16:19:46 -0400
Subject: [PATCH 01/37] Create cloud-tpu-presubmit

---
 .github/workflows/cloud-tpu-presubmit | 90 +++++++++++++++++++++++++++
 1 file changed, 90 insertions(+)
 create mode 100644 .github/workflows/cloud-tpu-presubmit

diff --git a/.github/workflows/cloud-tpu-presubmit b/.github/workflows/cloud-tpu-presubmit
new file mode 100644
index 000000000000..f15f51271666
--- /dev/null
+++ b/.github/workflows/cloud-tpu-presubmit
@@ -0,0 +1,90 @@
+# Cloud TPU CI
+name: Cloud TPU Presubmit
+# Run on pull_request that is labeled as "optional_ci_tpu" or workflow dispatch
+on:
+  pull_request:
+    branches:
+      - main
+    types: [labeled, synchronize]
+  workflow_dispatch:
+# Cancel any previous iterations if a new commit is pushed
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+jobs:
+  cloud-tpu-test:
+    # TODO: confirm final naming for optional label
+    if: contains(github.event.pull_request.labels.*.name, 'optional_ci_tpu')
+    name: "TPU v5e x 8 Presubmit"
+    env:
+      ENABLE_PJRT_COMPATIBILITY: 1
+    # TODO: Needs final runs-on value
+    runs-on: arc-linux-x86-ct5lp-224-8tpu
+    container:
+      # TODO: Needs newer, light weight image
+      image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
+    timeout-minutes: 45
+    defaults:
+      run:
+        shell: bash -ex {0}
+    steps:
+      - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
+      - name: Install JAX test requirements
+        run: |
+          pip install -U -r build/test-requirements.txt
+      - name: DEBUG HALT
+        run: |
+          echo "Halting"
+          sleep 180m
+      # TODO: build jax should be done on a step prior or we should just bazel test
+      - name: Build JAX
+        run: |
+          pip uninstall -y jaxlib
+          python3 build/build.py --use_clang 
+          pip install -e .
+          ls -la dist/*.whl
+          pip install dist/*.whl
+          # Note the version it installs! Should be today's date
+          pip install -U --no-index --pre libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html 
+          # Everything being built in this step downgrades numpy, reupgrade it    
+          pip install "numpy>=2.0.0"
+          python3 -c 'import sys; print("python version:", sys.version)'
+          python3 -c 'import jax; print("jax version:", jax.__version__)'
+          python3 -c 'import jaxlib; print("jaxlib version:", jaxlib.__version__)'
+          # strings $HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so | grep 'Built on'
+          python3 -c 'import jax; print("libtpu version:",
+            jax.lib.xla_bridge.get_backend().platform_version)'
+      - name: Run tests
+        env:
+          JAX_PLATFORMS: tpu,cpu
+          PY_COLORS: 1
+          NUM_TESTS: 8
+          JAX_NUM_GENERATED_CASES: 25
+        run: |
+          # Run single-accelerator tests in parallel
+          mkdir results
+          JAX_ENABLE_TPU_XDIST=true python3 -m pytest -n=$NUM_TESTS --tb=short \
+            --junitxml=results/singlejunit.xml --maxfail=20 -m "not multiaccelerator" tests examples
+          # Run multi-accelerator across all chips
+          python3 -m pytest --tb=short --junitxml=results/multijunit.xml \
+            --maxfail=20 -m "multiaccelerator" tests
+      - name: 'Upload Artifact'
+        if: success() || failure()
+        uses: actions/upload-artifact@65462800fd760344b1a7b4382951275a0abb4808 # ratchet:actions/upload-artifact@v4
+        with:
+          name: junit
+          path: |
+            results/singlejunit.xml
+            results/multijunit.xml
+          retention-days: 1
+
+
+
+
+
+
+
+
+
+
+

From 7c9506cca30e945976aade4a08d9e6788e2b42b0 Mon Sep 17 00:00:00 2001
From: Michael Hudgins <michaelhudgins@google.com>
Date: Mon, 24 Jun 2024 16:21:14 -0400
Subject: [PATCH 02/37] Rename cloud-tpu-presubmit to cloud-tpu-presubmit.yml

---
 .../workflows/{cloud-tpu-presubmit => cloud-tpu-presubmit.yml}    | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename .github/workflows/{cloud-tpu-presubmit => cloud-tpu-presubmit.yml} (100%)

diff --git a/.github/workflows/cloud-tpu-presubmit b/.github/workflows/cloud-tpu-presubmit.yml
similarity index 100%
rename from .github/workflows/cloud-tpu-presubmit
rename to .github/workflows/cloud-tpu-presubmit.yml

From 882973098e81f5a45c969ced8cafab76c73ca67c Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Wed, 26 Jun 2024 22:10:34 +0000
Subject: [PATCH 03/37] Add a tpu presubmit workflow

---
 .github/workflows/cloud-tpu-presubmit.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index f15f51271666..28832d466307 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -16,6 +16,10 @@ jobs:
     # TODO: confirm final naming for optional label
     if: contains(github.event.pull_request.labels.*.name, 'optional_ci_tpu')
     name: "TPU v5e x 8 Presubmit"
+    strategy:
+      fail-fast: false # don't cancel all jobs on failure
+      matrix:
+        instances: ["one", "two"]
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value

From c594236580a03bd0cc135a2a49c4c7b26e939c68 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Thu, 27 Jun 2024 16:25:12 -0700
Subject: [PATCH 04/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 28832d466307..4bc00edae716 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -18,8 +18,8 @@ jobs:
     name: "TPU v5e x 8 Presubmit"
     strategy:
       fail-fast: false # don't cancel all jobs on failure
-      matrix:
-        instances: ["one", "two"]
+#      matrix:
+#        instances: ["one", "two"]
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value

From 067e211bc226e40b341890772d6e4a7352d17262 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Thu, 27 Jun 2024 21:05:04 -0700
Subject: [PATCH 05/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 4bc00edae716..28832d466307 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -18,8 +18,8 @@ jobs:
     name: "TPU v5e x 8 Presubmit"
     strategy:
       fail-fast: false # don't cancel all jobs on failure
-#      matrix:
-#        instances: ["one", "two"]
+      matrix:
+        instances: ["one", "two"]
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value

From 9b4b45e4360fe79a780dcc2d1d49b002789b9697 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 8 Jul 2024 17:28:32 -0700
Subject: [PATCH 06/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 28832d466307..d8b0daf98a6f 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -19,7 +19,7 @@ jobs:
     strategy:
       fail-fast: false # don't cancel all jobs on failure
       matrix:
-        instances: ["one", "two"]
+        instances: ["one", "two", "three", "four"]
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value

From 6fa2f66f985044ff7ec95d5f73049a69624085bc Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Fri, 2 Aug 2024 09:19:37 -0700
Subject: [PATCH 07/37] Create gcs_access.yml

---
 .github/workflows/gcs_access.yml | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)
 create mode 100644 .github/workflows/gcs_access.yml

diff --git a/.github/workflows/gcs_access.yml b/.github/workflows/gcs_access.yml
new file mode 100644
index 000000000000..c5707f84db44
--- /dev/null
+++ b/.github/workflows/gcs_access.yml
@@ -0,0 +1,27 @@
+# WorkFlow dispatch to stress test the runner system
+name: GCS Access Test
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+# Cancel any previous iterations if a new commit is pushed
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+jobs:
+  cpu-stress:
+    strategy:
+      matrix:
+        runner: ["arc-linux-x86-n2-64"]
+        instances: ["1"]
+    runs-on: ${{ matrix.runner }}
+    timeout-minutes: 45
+    container:
+      # TODO: Needs newer, light weight image
+      image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
+    steps:
+      - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
+      - name: Gcloud
+        run: |
+          curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" "https://storage.googleapis.com/storage/v1/b/quoct-tf-test/o"

From 6181511ee781457264b402285f6acdd9b7b1e900 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Fri, 2 Aug 2024 09:48:03 -0700
Subject: [PATCH 08/37] Update gcs_access.yml

---
 .github/workflows/gcs_access.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.github/workflows/gcs_access.yml b/.github/workflows/gcs_access.yml
index c5707f84db44..1fd16c9ce172 100644
--- a/.github/workflows/gcs_access.yml
+++ b/.github/workflows/gcs_access.yml
@@ -25,3 +25,7 @@ jobs:
       - name: Gcloud
         run: |
           curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" "https://storage.googleapis.com/storage/v1/b/quoct-tf-test/o"
+      - name: Install JAX
+        run: |
+          pip install -U -r build/test-requirements.txt
+          pip install -U --pre jax==0.4.31.dev20240625 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

From b7d6a1a143af07ecca4291899b743decdf84085c Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 5 Aug 2024 11:29:05 -0700
Subject: [PATCH 09/37] Update gcs_access.yml

---
 .github/workflows/gcs_access.yml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.github/workflows/gcs_access.yml b/.github/workflows/gcs_access.yml
index 1fd16c9ce172..ff7775a7e530 100644
--- a/.github/workflows/gcs_access.yml
+++ b/.github/workflows/gcs_access.yml
@@ -22,6 +22,10 @@ jobs:
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
     steps:
       - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
+      - name: 'Set up Cloud SDK'
+        uses: 'google-github-actions/setup-gcloud@v2'
+        with:
+          version: '>= 363.0.0'
       - name: Gcloud
         run: |
           curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" "https://storage.googleapis.com/storage/v1/b/quoct-tf-test/o"

From b08fb4ae41fc08c1defd4519ae606a2293ae8532 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 13 Aug 2024 16:47:42 -0700
Subject: [PATCH 10/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 43 +----------------------
 1 file changed, 1 insertion(+), 42 deletions(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index d8b0daf98a6f..91de7d81075b 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: arc-linux-x86-ct5lp-224-8tpu
+    runs-on: arc-linux-x86-n2-32-dev
     container:
       # TODO: Needs newer, light weight image
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
@@ -40,47 +40,6 @@ jobs:
         run: |
           echo "Halting"
           sleep 180m
-      # TODO: build jax should be done on a step prior or we should just bazel test
-      - name: Build JAX
-        run: |
-          pip uninstall -y jaxlib
-          python3 build/build.py --use_clang 
-          pip install -e .
-          ls -la dist/*.whl
-          pip install dist/*.whl
-          # Note the version it installs! Should be today's date
-          pip install -U --no-index --pre libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html 
-          # Everything being built in this step downgrades numpy, reupgrade it    
-          pip install "numpy>=2.0.0"
-          python3 -c 'import sys; print("python version:", sys.version)'
-          python3 -c 'import jax; print("jax version:", jax.__version__)'
-          python3 -c 'import jaxlib; print("jaxlib version:", jaxlib.__version__)'
-          # strings $HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so | grep 'Built on'
-          python3 -c 'import jax; print("libtpu version:",
-            jax.lib.xla_bridge.get_backend().platform_version)'
-      - name: Run tests
-        env:
-          JAX_PLATFORMS: tpu,cpu
-          PY_COLORS: 1
-          NUM_TESTS: 8
-          JAX_NUM_GENERATED_CASES: 25
-        run: |
-          # Run single-accelerator tests in parallel
-          mkdir results
-          JAX_ENABLE_TPU_XDIST=true python3 -m pytest -n=$NUM_TESTS --tb=short \
-            --junitxml=results/singlejunit.xml --maxfail=20 -m "not multiaccelerator" tests examples
-          # Run multi-accelerator across all chips
-          python3 -m pytest --tb=short --junitxml=results/multijunit.xml \
-            --maxfail=20 -m "multiaccelerator" tests
-      - name: 'Upload Artifact'
-        if: success() || failure()
-        uses: actions/upload-artifact@65462800fd760344b1a7b4382951275a0abb4808 # ratchet:actions/upload-artifact@v4
-        with:
-          name: junit
-          path: |
-            results/singlejunit.xml
-            results/multijunit.xml
-          retention-days: 1
 
 
 

From ad8965dee05a60e84711345cc5850c7122596c99 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 13 Aug 2024 16:48:16 -0700
Subject: [PATCH 11/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 91de7d81075b..73b6b7dea8f3 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -19,7 +19,7 @@ jobs:
     strategy:
       fail-fast: false # don't cancel all jobs on failure
       matrix:
-        instances: ["one", "two", "three", "four"]
+        instances: ["one"]
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value

From f6d55dc9f952183245c2758c42d4044136a4d072 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 13 Aug 2024 16:48:27 -0700
Subject: [PATCH 12/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 73b6b7dea8f3..ad16a44314ad 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -19,7 +19,7 @@ jobs:
     strategy:
       fail-fast: false # don't cancel all jobs on failure
       matrix:
-        instances: ["one"]
+        instances: ["one", "two"]
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value

From 755a5416c83a2174608a3f98fe1ef23461216ada Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 13 Aug 2024 16:58:50 -0700
Subject: [PATCH 13/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index ad16a44314ad..ffd163c98014 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: arc-linux-x86-n2-32-dev
+    runs-on: arc-linux-x86-n2-64-dev
     container:
       # TODO: Needs newer, light weight image
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11

From ad5105bea250d003b0618897c120d83afad25109 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 13 Aug 2024 17:47:52 -0700
Subject: [PATCH 14/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index ffd163c98014..ce83d3792f43 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -39,7 +39,7 @@ jobs:
       - name: DEBUG HALT
         run: |
           echo "Halting"
-          sleep 180m
+          sleep 30m
 
 
 

From 7da0a6a2ea76835f7af914dc4dcf425bff5342bb Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Wed, 14 Aug 2024 15:27:42 -0700
Subject: [PATCH 15/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index ce83d3792f43..d43ddea5c31a 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: arc-linux-x86-n2-64-dev
+    runs-on: arc-linux-x86-n2-64-test
     container:
       # TODO: Needs newer, light weight image
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11

From 4388e50aab5ee3d5f3f1353d2fad4e068c071d5b Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 20 Aug 2024 11:34:08 -0700
Subject: [PATCH 16/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index d43ddea5c31a..0aaa5a9d6676 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: arc-linux-x86-n2-64-test
+    runs-on: t2a-standard-48-dev
     container:
       # TODO: Needs newer, light weight image
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11

From 4a000df885d512c42465629c6a139c411fd794fd Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 20 Aug 2024 12:52:58 -0700
Subject: [PATCH 17/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 0aaa5a9d6676..54fbfd332a87 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -26,7 +26,7 @@ jobs:
     runs-on: t2a-standard-48-dev
     container:
       # TODO: Needs newer, light weight image
-      image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
+      image: us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build-arm64:jax-latest-multi-python
     timeout-minutes: 45
     defaults:
       run:

From 22c801636c60b386de8024d3e7b97692d02e4e9a Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Fri, 23 Aug 2024 15:00:18 -0700
Subject: [PATCH 18/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 54fbfd332a87..8b2aa9e030f4 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: t2a-standard-48-dev
+    runs-on: arc-linux-x86-n2-64-test
     container:
       # TODO: Needs newer, light weight image
       image: us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build-arm64:jax-latest-multi-python

From d42fbee5e1bc2ed5f06be554a5ed561995a2bb85 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Fri, 23 Aug 2024 15:07:58 -0700
Subject: [PATCH 19/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 8b2aa9e030f4..d43ddea5c31a 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -26,7 +26,7 @@ jobs:
     runs-on: arc-linux-x86-n2-64-test
     container:
       # TODO: Needs newer, light weight image
-      image: us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build-arm64:jax-latest-multi-python
+      image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
     timeout-minutes: 45
     defaults:
       run:

From 7d0afdd1df9586ee2b7d204dd81c5543d8276b8f Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 8 Oct 2024 14:14:54 -0700
Subject: [PATCH 20/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index d43ddea5c31a..218b7e93cef4 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: arc-linux-x86-n2-64-test
+    runs-on: linux-x86-n2-16-dev
     container:
       # TODO: Needs newer, light weight image
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11

From 7cd8fe0e227d4d0f6e19453ccf4a38b5f5a1f772 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 8 Oct 2024 14:16:51 -0700
Subject: [PATCH 21/37] Delete .github/workflows/gcs_access.yml

---
 .github/workflows/gcs_access.yml | 35 --------------------------------
 1 file changed, 35 deletions(-)
 delete mode 100644 .github/workflows/gcs_access.yml

diff --git a/.github/workflows/gcs_access.yml b/.github/workflows/gcs_access.yml
deleted file mode 100644
index ff7775a7e530..000000000000
--- a/.github/workflows/gcs_access.yml
+++ /dev/null
@@ -1,35 +0,0 @@
-# WorkFlow dispatch to stress test the runner system
-name: GCS Access Test
-on:
-  pull_request:
-    branches:
-      - main
-  workflow_dispatch:
-# Cancel any previous iterations if a new commit is pushed
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
-jobs:
-  cpu-stress:
-    strategy:
-      matrix:
-        runner: ["arc-linux-x86-n2-64"]
-        instances: ["1"]
-    runs-on: ${{ matrix.runner }}
-    timeout-minutes: 45
-    container:
-      # TODO: Needs newer, light weight image
-      image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
-    steps:
-      - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
-      - name: 'Set up Cloud SDK'
-        uses: 'google-github-actions/setup-gcloud@v2'
-        with:
-          version: '>= 363.0.0'
-      - name: Gcloud
-        run: |
-          curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" "https://storage.googleapis.com/storage/v1/b/quoct-tf-test/o"
-      - name: Install JAX
-        run: |
-          pip install -U -r build/test-requirements.txt
-          pip install -U --pre jax==0.4.31.dev20240625 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

From aa6f6181ad14127794e80613613d6d25c47855a4 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Wed, 9 Oct 2024 08:55:31 -0700
Subject: [PATCH 22/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 218b7e93cef4..9abbd7c74534 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: linux-x86-n2-16-dev
+    runs-on: linux-x86-n2-16-quoct
     container:
       # TODO: Needs newer, light weight image
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11

From e6db8e201f2fe4eb8bf800d5d347bf1cf22ba988 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Thu, 10 Oct 2024 13:27:57 -0700
Subject: [PATCH 23/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 9abbd7c74534..4d5fd3df6c84 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -23,7 +23,7 @@ jobs:
     env:
       ENABLE_PJRT_COMPATIBILITY: 1
     # TODO: Needs final runs-on value
-    runs-on: linux-x86-n2-16-quoct
+    runs-on: linux-x86-n2-16
     container:
       # TODO: Needs newer, light weight image
       image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11

From b85462d8be89aac9b46349c2c73b33555f17defd Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 15:25:21 -0700
Subject: [PATCH 24/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 95 +++++++++++++----------
 1 file changed, 53 insertions(+), 42 deletions(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 4d5fd3df6c84..d87a7518bea2 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -1,53 +1,64 @@
-# Cloud TPU CI
-name: Cloud TPU Presubmit
-# Run on pull_request that is labeled as "optional_ci_tpu" or workflow dispatch
+name: Build JAX Artifacts
+
 on:
   pull_request:
     branches:
       - main
-    types: [labeled, synchronize]
-  workflow_dispatch:
-# Cancel any previous iterations if a new commit is pushed
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  workflow_call:
+
 jobs:
-  cloud-tpu-test:
-    # TODO: confirm final naming for optional label
-    if: contains(github.event.pull_request.labels.*.name, 'optional_ci_tpu')
-    name: "TPU v5e x 8 Presubmit"
-    strategy:
-      fail-fast: false # don't cancel all jobs on failure
-      matrix:
-        instances: ["one", "two"]
-    env:
-      ENABLE_PJRT_COMPATIBILITY: 1
-    # TODO: Needs final runs-on value
-    runs-on: linux-x86-n2-16
-    container:
-      # TODO: Needs newer, light weight image
-      image: index.docker.io/tensorflow/build@sha256:7fb38f0319bda36393cad7f40670aa22352b44421bb906f5cf34d543acd8e1d2 # ratchet:tensorflow/build:latest-python3.11
-    timeout-minutes: 45
+  build:
+    continue-on-error: true
     defaults:
       run:
-        shell: bash -ex {0}
-    steps:
-      - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
-      - name: Install JAX test requirements
-        run: |
-          pip install -U -r build/test-requirements.txt
-      - name: DEBUG HALT
-        run: |
-          echo "Halting"
-          sleep 30m
-
-
-
-
-
-
-
+        # Explicitly set the shell to bash to override the default Windows environment, i.e, cmd.
+        shell: bash
+    strategy:
+      matrix:
+        runner: ["windows-x86-n2-64-dev"]
+        artifact: ["jax", "jaxlib", "jax-cuda-pjrt", "jax-cuda-plugin"]
+        python: ["3.10", "3.11", "3.12"]
+        # jax-cuda-pjrt and jax are pure Python packages so they do not need to be built for each
+        # Python version.
+        exclude:
+          # Pure Python packages do not need to be built for each Python version.
+          - artifact: "jax-cuda-pjrt"
+            python: "3.10"
+          - artifact: "jax-cuda-pjrt"
+            python: "3.11"
+          - artifact: "jax"
+            python: "3.10"
+          - artifact: "jax"
+            python: "3.11"
+          # jax is a pure Python package so it does not need to be built on multiple platforms.
+          - artifact: "jax"
+            runner: "windows-x86-n2-64"
+          - artifact: "jax"
+            runner: "linux-arm64-t2a-48"
+          # jax-cuda-plugin and jax-cuda-pjrt are not supported on Windows.
+          - artifact: "jax-cuda-plugin"
+            runner: "windows-x86-n2-64"
+          - artifact: "jax-cuda-pjrt"
+            runner: "windows-x86-n2-64"
 
+    runs-on: ${{ matrix.runner }}
 
+    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
+                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') ||
+                   (contains(matrix.runner, 'windows-x86') && null) }}
 
+    env:
+      # Do not run Docker container for Linux runners. Linux runners already run in a Docker container.
+      JAXCI_RUN_DOCKER_CONTAINER: 0
+      # Use RBE to build the artifacts where possibl (Linux x86 and Windows).
+      JAXCI_BUILD_ARTIFACT_WITH_RBE: 1
 
+    steps:
+      - uses: actions/checkout@v3
+      # Halt for testing
+      - name: Wait For Connection
+        uses: ./actions/ci_connection/
+      - name: Build ${{ matrix.artifact }}
+        env:
+          JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
+        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/${{ matrix.artifact }}"

From d3d0729b8306490b9d134ed7cfa60546819f0b7a Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 15:37:12 -0700
Subject: [PATCH 25/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index d87a7518bea2..954a12d4b555 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,4 +61,5 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/${{ matrix.artifact }}"
+        run: choco install llvm --version=18.1.4 --yes --no-progress --allow-downgrade
+

From 5493f2dba8287ee81b878219e19e351c8d58e821 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 15:48:04 -0700
Subject: [PATCH 26/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 954a12d4b555..df492721bb50 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,5 +61,5 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: choco install llvm --version=18.1.4 --yes --no-progress --allow-downgrade
+        run: ls
 

From 2590d18c923380b55410b6cfdfde4a8d5134d136 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 15:56:04 -0700
Subject: [PATCH 27/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index df492721bb50..bf46d0b94c6e 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,5 +61,5 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: ls
+        run: ls && timeout 30 && ls
 

From ce2a59b15af60380316499fe1066724454d1bdc6 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 16:17:09 -0700
Subject: [PATCH 28/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index bf46d0b94c6e..82f0883343d8 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,5 +61,5 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: ls && timeout 30 && ls
+        run: ls; timeout 30; ls
 

From 0123a73e9bab4b71871bed49b685bed9054eb24d Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 16:19:13 -0700
Subject: [PATCH 29/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 82f0883343d8..9c87aecdd638 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,5 +61,5 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: ls; timeout 30; ls
+        run: timeout 100
 

From 81ffebe50d7d50d61ee177bfb909672b95135646 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 16:20:31 -0700
Subject: [PATCH 30/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 9c87aecdd638..90787cb318ec 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,5 +61,4 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: timeout 100
-
+        run: ls

From 12f0281c25f762df469449cb86c57bb7f9404b0e Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 16:22:19 -0700
Subject: [PATCH 31/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 90787cb318ec..16934708b8b4 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,4 +61,4 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: ls
+        run: sleep 100

From 5a226f232e813a4055dcef23aacaba681e801c9d Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 28 Oct 2024 16:39:33 -0700
Subject: [PATCH 32/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 16934708b8b4..36651ba97179 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -61,4 +61,4 @@ jobs:
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: sleep 100
+        run: sleep 400

From e2b9c92d6736f24d3dd31601352fda25d1a39163 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 3 Dec 2024 19:09:17 -0800
Subject: [PATCH 33/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 36651ba97179..d853d3de313f 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -15,9 +15,9 @@ jobs:
         shell: bash
     strategy:
       matrix:
-        runner: ["windows-x86-n2-64-dev"]
+        runner: ["linux-x86-n2-64"]
         artifact: ["jax", "jaxlib", "jax-cuda-pjrt", "jax-cuda-plugin"]
-        python: ["3.10", "3.11", "3.12"]
+        python: ["3.11"]
         # jax-cuda-pjrt and jax are pure Python packages so they do not need to be built for each
         # Python version.
         exclude:
@@ -43,9 +43,7 @@ jobs:
 
     runs-on: ${{ matrix.runner }}
 
-    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
-                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') ||
-                   (contains(matrix.runner, 'windows-x86') && null) }}
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:702526224"
 
     env:
       # Do not run Docker container for Linux runners. Linux runners already run in a Docker container.
@@ -55,10 +53,11 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
-      # Halt for testing
-      - name: Wait For Connection
-        uses: ./actions/ci_connection/
       - name: Build ${{ matrix.artifact }}
         env:
           JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: sleep 400
+        run: echo "Testing quoct"
+      - name: Another Test
+        run: ls
+      - name: Another Test
+        run: sleep 500

From f2a74c2cfce0c2b9fb4aae5d49227ab04dd98033 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Tue, 3 Dec 2024 19:45:05 -0800
Subject: [PATCH 34/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index d853d3de313f..0991d74fbd43 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -15,7 +15,7 @@ jobs:
         shell: bash
     strategy:
       matrix:
-        runner: ["linux-x86-n2-64"]
+        runner: ["linux-x86-n2-64-quoct"]
         artifact: ["jax", "jaxlib", "jax-cuda-pjrt", "jax-cuda-plugin"]
         python: ["3.11"]
         # jax-cuda-pjrt and jax are pure Python packages so they do not need to be built for each

From e70fd114b1ab3fc43ef36e610a2f9a7b5ff770f4 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Thu, 5 Dec 2024 14:32:35 -0800
Subject: [PATCH 35/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 0991d74fbd43..e94ca3535446 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -15,7 +15,7 @@ jobs:
         shell: bash
     strategy:
       matrix:
-        runner: ["linux-x86-n2-64-quoct"]
+        runner: ["linux-arm64-c4a-16"]
         artifact: ["jax", "jaxlib", "jax-cuda-pjrt", "jax-cuda-plugin"]
         python: ["3.11"]
         # jax-cuda-pjrt and jax are pure Python packages so they do not need to be built for each

From d721eac27b3897ef366a6e166adf607838c0fc41 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 16 Dec 2024 11:22:47 -0800
Subject: [PATCH 36/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index e94ca3535446..97f22038bd71 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -60,4 +60,4 @@ jobs:
       - name: Another Test
         run: ls
       - name: Another Test
-        run: sleep 500
+        run: sleep 10000

From af1467a3ab9c700289bbca8d3ba1bd344ec1a36c Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Mon, 16 Dec 2024 13:06:39 -0800
Subject: [PATCH 37/37] Update cloud-tpu-presubmit.yml

---
 .github/workflows/cloud-tpu-presubmit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/cloud-tpu-presubmit.yml b/.github/workflows/cloud-tpu-presubmit.yml
index 97f22038bd71..9365639d4916 100644
--- a/.github/workflows/cloud-tpu-presubmit.yml
+++ b/.github/workflows/cloud-tpu-presubmit.yml
@@ -43,7 +43,7 @@ jobs:
 
     runs-on: ${{ matrix.runner }}
 
-    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:702526224"
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
 
     env:
       # Do not run Docker container for Linux runners. Linux runners already run in a Docker container.