Test CI scripts and workflows

.bazelrc

@@ -210,6 +210,9 @@ build:macos_cache --remote_cache="https://storage.googleapis.com/tensorflow-maco
 # Cache pushes are limited to JAX's CI system.
 build:macos_cache_push --config=macos_cache --remote_upload_local_results=true --google_default_credentials
 
+build:public_cache --remote_cache="https://storage.googleapis.com/jax-bazel-cache/" --remote_upload_local_results=false
+build:public_cache_push --config=public_cache --remote_upload_local_results=true --google_default_credentials
+
 # #############################################################################
 # CI Build config options below.
 # JAX uses these configs in CI builds for building artifacts and when running

.github/workflows/build_artifacts.yml

@@ -109,7 +109,6 @@ jobs:
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
-      JAXCI_CLONE_MAIN_XLA: "${{ inputs.clone_main_xla }}"
 
     name: Build ${{ inputs.artifact }} (${{ inputs.runner }}, Python ${{ inputs.python }}, clone main XLA=${{ inputs.clone_main_xla }})
 
@@ -118,7 +117,14 @@ jobs:
       gcs_upload_uri: ${{ steps.store-gcs-upload-uri.outputs.gcs_upload_uri }}
 
     steps:
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - name: Checkout JAX repository
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2      
+      - name: Checkout XLA repository
+        if: contains("${{ inputs.clone_main_xla }}", '1')
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          repository: openxla/xla
+          path: jax/xla
       - name: Enable RBE if building on Linux x86 or Windows x86
         if: contains(inputs.runner, 'linux-x86') || contains(inputs.runner, 'windows-x86')
         run: echo "JAXCI_BUILD_ARTIFACT_WITH_RBE=1" >> $GITHUB_ENV

.github/workflows/build_artifacts_single_invocation.yml

@@ -0,0 +1,37 @@
+name: Build multiple artifacts with build CLI
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+        - 'yes'
+        - 'no'
+
+jobs:
+  build_artifacts:
+
+    runs-on: "linux-x86-n2-64"
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
+
+    env:
+      JAXCI_HERMETIC_PYTHON_VERSION: 3.12
+
+    steps:
+      - uses: actions/checkout@v3
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Mark GitHub repo as safe
+        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"
+      - name: Build all artifacts with a single invocation of build CLI
+        run: python build/build.py build --editable --use_clang --output_path=$(pwd)/dist --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose

.github/workflows/ci_duplicate.yml

@@ -0,0 +1,41 @@
+name: CI (duplicate)
+
+on:
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read  # to fetch code
+  actions: write  # to cancel previous workflows
+
+jobs:
+  lint_and_typecheck:
+    runs-on: linux-x86-n2-16
+    container:
+      image: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest" # ratchet:ubuntu:20.04
+    timeout-minutes: 5
+    steps:
+      - name: Install dependencies for setting up Python
+        env:
+            DEBIAN_FRONTEND: noninteractive
+        run: |
+            apt update && apt install libssl-dev libsqlite3-dev git -y
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - run: python -m pip install pre-commit
+    #   - name: Add to PATH
+    #     run: export PATH=$PATH:/__w/_tool/Python/3.11.10/x64/bin
+    #   - name: Set python3 alias
+    #     run: echo "alias python3=/__w/_tool/Python/3.11.10/x64/bin/python3" >> ~/.bashrc
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+      - uses: actions/cache@6849a6489940f00c2f30c0fb92c6274307ccb58a # v4.1.2
+        with:
+          path: ~/.cache/pre-commit
+          key: pre-commit-${{ env.pythonLocation }}-${{ hashFiles('.pre-commit-config.yaml', 'setup.py') }}
+      - run: pre-commit run --show-diff-on-failure --color=always --all-files

.github/workflows/env_var_debug.yml

@@ -0,0 +1,46 @@
+name: CI - Debug
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  run-tests:
+
+    strategy:
+        fail-fast: false # don't cancel all jobs on failure
+        matrix:
+          # Python values need to match the matrix stategy in the build artifacts job above
+          runner: ["linux-x86-g2-48-l4-4gpu",]
+          python: ["3.11",]
+
+    runs-on: ${{ matrix.runner }}
+    container:  "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
+
+    name: "Debug env var setting"
+
+    env:
+      JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
+      JAXCI_PYTHON: "python${{ matrix.python }}"
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set env vars for use in artifact download URL
+        run: |
+            os=$(uname -s | awk '{print tolower($0)}')
+            arch=$(uname -m)
+
+            # Get the major and minor version of Python.
+            # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.10, then python_major_minor=310
+            echo $os $arch $JAXCI_HERMETIC_PYTHON_VERSION
+
+            python_major_minor_1=$(echo "$JAXCI_HERMETIC_PYTHON_VERSION" | tr -d '.')
+            echo $os $arch $JAXCI_HERMETIC_PYTHON_VERSION $python_major_minor_1
+
+            python_major_minor_2="${JAXCI_HERMETIC_PYTHON_VERSION//./}"
+            echo $os $arch $JAXCI_HERMETIC_PYTHON_VERSION $python_major_minor_1 $python_major_minor_2
+
+            echo "OS=${os}" >> $GITHUB_ENV
+            echo "ARCH=${arch}" >> $GITHUB_ENV
+            echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV

.github/workflows/pytest_cpu_crash_debug.yml

@@ -0,0 +1,57 @@
+name: CI - Pytest CPU (crash debug)
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  run-tests:
+    defaults:
+      run:
+        # Explicitly set the shell to bash to override the default Windows environment, i.e, cmd.
+        shell: bash
+
+    strategy:
+        fail-fast: false # don't cancel all jobs on failure
+        matrix:
+          runner: ["linux-x86-t2a-48-dev",]
+          python: ["3.10",]
+          enable-x64: [1, 0]
+
+    runs-on: ${{ matrix.runner }}
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest"
+
+    name: "Pytest CPU (${{ matrix.runner }}, Python ${{ matrix.python }}, x64=${{ matrix.enable-x64 }})"
+
+    env:
+      JAXCI_HERMETIC_PYTHON_VERSION: "${{ matrix.python }}"
+      JAXCI_PYTHON: "python${{ matrix.python }}"
+      JAXCI_ENABLE_X64: ${{ matrix.enable-x64 }}
+
+
+    steps:
+      - uses: actions/checkout@v3
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+      # Checkout XLA at head, if we're building jaxlib at head.
+      - name: Checkout XLA at head
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          repository: openxla/xla
+          path: xla
+      # We need to mark the GitHub workspace as safe as otherwise git commands will fail.
+      - name: Mark GitHub workspace as safe
+        run: |
+          git config --global --add safe.directory "$GITHUB_WORKSPACE"
+      - name: Build jaxlib
+        run: |
+            $JAXCI_PYTHON build/build.py build --wheels=jaxlib \
+                        --bazel_options=--config=ci_linux_aarch64 \
+                        --local_xla_path="$(pwd)/xla" \
+                        --verbose
+      - name: Install dependencies
+        run: $JAXCI_PYTHON -m pip install -r build/requirements.in
+      - name: Run Pytest CPU tests
+        run: ./ci/run_pytest_cpu.sh

.github/workflows/pytest_cuda.yml

@@ -58,7 +58,7 @@ jobs:
     runs-on: ${{ inputs.runner }}
     # Test the oldest and newest supported CUDA versions.
     container:  ${{ (contains(inputs.cuda, '12.1') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.1-cudnn9.8:latest') ||
-                (contains(inputs.cuda, '12.8') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.8-cudnn9.8:latest') }}
+                (contains(inputs.cuda, '12.8') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.8-cudnn9.8:latest')}}
     name: "Pytest CUDA (${{ inputs.runner }}, CUDA ${{ inputs.cuda }}, Python ${{ inputs.python }}, x64=${{ inputs.enable-x64 }})"
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"

.github/workflows/wheel_tests_continuous.yml

@@ -20,6 +20,9 @@ name: CI - Wheel Tests (Continuous)
 on:
   schedule:
     - cron: "0 0 * * *" # Run once per day
+  pull_request:
+    branches:
+      - main 
   workflow_dispatch: # allows triggering the workflow run manually
 
 concurrency:
@@ -55,7 +58,7 @@ jobs:
         runner: ${{ matrix.runner }}
         artifact: ${{ matrix.artifact }}
         python: ${{ matrix.python }}
-        clone_main_xla: 1
+        clone_main_xla: 0
         upload_artifacts_to_gcs: true
         gcs_upload_uri: 'gs://general-ml-ci-transient/jax-github-actions/jax-fork/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
 
@@ -73,7 +76,7 @@ jobs:
       runner: ${{ matrix.runner }}
       artifact: ${{ matrix.artifact }}
       python: ${{ matrix.python }}
-      clone_main_xla: 1
+      clone_main_xla: 0
       upload_artifacts_to_gcs: true
       gcs_upload_uri: 'gs://general-ml-ci-transient/jax-github-actions/jax-fork/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'