experiment with reusable workflows

Author: nitins17

.github/workflows/build_artifacts.yml

@@ -15,6 +15,42 @@ on:
         - 'yes'
         - 'no'
   workflow_call:
+    inputs:
+      build_jax:
+        description: "Should the jax artifact be built?"
+        required: true
+        default: true
+        type: boolean
+      build_jaxlib:
+        description: "Should the jaxlib artifact be built?"
+        required: true
+        default: true
+        type: boolean
+      build_jax_cuda_plugin:
+        description: "Should the jax-cuda-plugin artifact be built?"
+        required: true
+        default: true
+        type: boolean
+      build_jax_cuda_pjrt:
+        description: "Should the jax-cuda-pjrt artifact be built?"
+        required: true
+        default: true
+        type: boolean
+      clone_main_xla:
+        description: "Should latest XLA be used? (1 to enable, 0 to disable)"
+        type: string
+        required: true
+        default: "0"
+      upload_artifacts:
+        description: "Should the artifacts be uploaded to a GCS bucket?"
+        required: true
+        default: false
+        type: boolean
+      upload_destination:
+        description: "GCS location to where the artifacts should be uploaded"
+        required: true
+        default: '${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
+        type: string
 
 jobs:
   build:
@@ -27,7 +63,7 @@ jobs:
       matrix:
         runner: ["windows-x86-n2-64", "linux-x86-n2-16", "linux-arm64-t2a-16"]
         artifact: ["jax", "jaxlib", "jax-cuda-pjrt", "jax-cuda-plugin"]
-        python: ["3.10", "3.11", "3.12"]
+        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:
@@ -58,8 +94,8 @@ jobs:
                    (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
+      JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
+      JAXCI_CLONE_MAIN_XLA: "${{ inputs.clone_main_xla }}"
 
     steps:
       - uses: actions/checkout@v3
@@ -68,7 +104,19 @@ jobs:
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: Build ${{ matrix.artifact }}
-        env:
-          JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
-        run: ./ci/build_artifacts.sh "${{ matrix.artifact }}"
+      - name: Build jax
+        if: inputs.build_jax && matrix.artifact == 'jax'
+        run: ./ci/build_artifacts.sh "jax"
+      - name: Build jaxlib
+        if: inputs.build_jaxlib && matrix.artifact == 'jaxlib'
+        run: ./ci/build_artifacts.sh "jaxlib"
+      - name: Build jax-cuda-plugin
+        if: inputs.build_jax_cuda_plugin && matrix.artifact == 'jax-cuda-plugin'
+        run: ./ci/build_artifacts.sh "jax-cuda-plugin"
+      - name: Build jax-cuda-pjrt
+        if: inputs.build_jax_cuda_pjrt && matrix.artifact == 'jax-cuda-pjrt'
+        run: ./ci/build_artifacts.sh "jax-cuda-pjrt"
+      - name: Upload artifacts to GCS bucket
+        if: inputs.upload_artifacts
+        run:  ~/usr/local/bin/google-cloud-sdk/bin/gsutil -m cp -r $(pwd)/dist gs://general-ml-ci-transient/jax-github-actions/"${{ inputs.upload_destination }}"
+      

.github/workflows/pytest_cpu_reuse.yml

@@ -1,9 +1,9 @@
 name: Run Pytest CPU tests (resuable workflow)
 
 on:
-  # pull_request:
-  #   branches:
-  #     - main
+  pull_request:
+    branches:
+      - main
   workflow_dispatch:
     inputs:
       halt-for-connection:
@@ -15,28 +15,36 @@ on:
         - 'yes'
 
 jobs:
-  build_jaxlib_artifacts:
+  build_jaxlib_artifact:
+    name: "Build the jaxlib aritfact using latest XLA"
     uses: ./.github/workflows/build_artifacts.yml
+    with:
+      build_jax: false
+      build_jaxlib: true
+      build_jax_cuda_plugin: false
+      build_jax_cuda_pjrt: false
+      clone_main_xla: 1
+      upload_artifacts: true
+      upload_destination: '${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
 
   run_pytest:
-    needs: build_jaxlib_artifacts
+    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/linux-arm64-arc-container: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:
@@ -46,6 +54,8 @@ jobs:
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Download the artifacts built in the "build_artifacts" job
+        run: mkdir -p $(pwd)/dist && ~/usr/local/bin/google-cloud-sdk/bin/gsutil -m cp -r gs://general-ml-ci-transient/jax-github-actions/"${{ github.workflow }}"/${{ github.run_number }}/${{ github.run_attempt }} $(pwd)/dist
       - name: Install pytest
         env:
           JAXCI_PYTHON: python${{ matrix.python }}