Make the resuable pytest cpu workflow as default

Author: nitins17

.github/workflows/pytest_cpu.yml

@@ -1,9 +1,9 @@
 name: Run Pytest CPU tests
 
 on:
-  # pull_request:
-  #   branches:
-  #     - main
+  pull_request:
+    branches:
+      - main
   workflow_dispatch:
     inputs:
       halt-for-connection:
@@ -13,27 +13,37 @@ on:
         default: 'no'
         options:
         - 'yes'
-        - 'no'
 
 jobs:
-  build:
+  build_jaxlib_artifact:
+    name: "Build the jaxlib aritfact using latest XLA"
+    uses: ./.github/workflows/build_artifacts.yml
+    with:
+      wheel_list: "jaxlib"
+      python_list: "3.10"
+      platform_list: "linux_x86,linux_arm64"
+      clone_main_xla: 1
+      upload_artifacts: true
+      upload_destination: '${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
+
+  run_pytest:
+    name: "Run CPU tests with Pytest"
+    needs: build_jaxlib_artifact
     continue-on-error: true
     defaults:
       run:
         # Explicitly set the shell to bash to override the default Windows environment, i.e, cmd.
         shell: bash
     strategy:
       matrix:
-        runner: ["windows-x86-n2-64", "linux-x86-n2-64", "linux-arm64-t2a-48"]
+        runner: ["linux-x86-n2-64", "linux-arm64-t2a-48"]
         python: ["3.10"]
 
     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/ml-build-arm64:latest') ||
-                (contains(matrix.runner, 'windows-x86') && null) }}
+                (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') }}
 
     env:
-      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
 
     steps:
@@ -43,8 +53,15 @@ jobs:
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Build jaxlib
-        run: ./ci/build_artifacts.sh "jaxlib"
+      - name: Set Platform
+        run: |
+          os=$(uname -s | awk '{print tolower($0)}')
+          arch=$(uname -m)
+          echo "PLATFORM=${os}_${arch}" >> $GITHUB_ENV
+      - name: Download the artifacts built in the "build_artifacts" job
+        run: >-
+          mkdir -p $(pwd)/dist &&
+          gsutil -m cp -r gs://general-ml-ci-transient/jax-github-actions/"${{ github.workflow }}"/${{ github.run_number }}/${{ github.run_attempt }}/$PLATFORM/python${JAXCI_HERMETIC_PYTHON_VERSION} $(pwd)/dist/
       - name: Install pytest
         env:
           JAXCI_PYTHON: python${{ matrix.python }}