Update cloud-tpu-presubmit.yml

Author: quoctruong

.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 }}"