Add a presubmit check to test against oldest supported numpy

PiperOrigin-RevId: 738525650

Author: nitins17

.github/workflows/oldest_supported_numpy.yml

@@ -0,0 +1,60 @@
+# CI - Oldest Supported NumPy (presubmit)
+# This workflow tests the oldest supported NumPy and jaxlib versions.
+
+name: CI - Oldest Supported NumPy (presubmit)
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+      - 'release/**'
+
+# This should also be set to read-only in the project settings, but it's nice to
+# document and enforce the permissions here.
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  # Don't cancel in-progress jobs for main/release branches.
+  cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
+
+jobs:
+  test-oldest-supported-numpy:
+    if: github.event.repository.fork == false
+    defaults:
+      run:
+        shell: bash
+    runs-on: "linux-x86-n2-64"
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"
+# Begin Presubmit Naming Check - name modification requires internal check to be updated
+    name: "CI - Oldest Supported NumPy (Python 3.10, x64=0)"
+# End Presubmit Naming Check github-oldest-supported-numpy-presubmit
+
+    env:
+      JAXCI_PYTHON: "python3.10"
+      JAXCI_ENABLE_X64: 0
+      JAX_NUM_GENERATED_CASES: 5
+
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - name: Install Python dependencies
+        run: |
+          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
+
+          # Install NumPy and SciPy with the oldest supported versions
+          $JAXCI_PYTHON -m uv pip install numpy==1.25.2 scipy==1.11.1
+
+          # Install JAX using the changes in the PR
+          $JAXCI_PYTHON -m uv pip install -e .[minimum-jaxlib]
+      # Halt for testing
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Run Pytest CPU tests
+        timeout-minutes: 30
+        run: ./ci/run_pytest_cpu.sh

ci/run_pytest_cpu.sh

@@ -26,13 +26,13 @@ set -exu -o history -o allexport
 # Source default JAXCI environment variables.
 source ci/envs/default.env
 
+# Set up the build environment.
+source "ci/utilities/setup_build_environment.sh"
+
 # Install jaxlib wheel inside the $JAXCI_OUTPUT_DIR directory on the system.
 echo "Installing wheels locally..."
 source ./ci/utilities/install_wheels_locally.sh
 
-# Set up the build environment.
-source "ci/utilities/setup_build_environment.sh"
-
 # Print all the installed packages
 echo "Installed packages:"
 "$JAXCI_PYTHON" -m uv pip list

ci/utilities/install_wheels_locally.sh

@@ -30,23 +30,25 @@ for i in "${!WHEELS[@]}"; do
   fi
 done
 
-if [[ -z "${WHEELS[@]}" ]]; then
-  echo "ERROR: No wheels found under $JAXCI_OUTPUT_DIR"
-  exit 1
-fi
+if [[ -n "${WHEELS[@]}" ]]; then
+  echo "Installing the following wheels:"
+  echo "${WHEELS[@]}"
 
-echo "Installing the following wheels:"
-echo "${WHEELS[@]}"
-
-# Install `uv` if it's not already installed. `uv` is much faster than pip for
-# installing Python packages.
-if ! command -v uv >/dev/null 2>&1; then
-  pip install uv~=0.5.30
-fi
+  # Install `uv` if it's not already installed. `uv` is much faster than pip for
+  # installing Python packages.
+  if ! command -v uv >/dev/null 2>&1; then
+    pip install uv~=0.5.30
+  fi
 
-# On Windows, convert MSYS Linux-like paths to Windows paths.
-if [[ $(uname -s) =~ "MSYS_NT" ]]; then
-  "$JAXCI_PYTHON" -m uv pip install $(cygpath -w "${WHEELS[@]}")
+  # On Windows, convert MSYS Linux-like paths to Windows paths.
+  if [[ $(uname -s) =~ "MSYS_NT" ]]; then
+    "$JAXCI_PYTHON" -m uv pip install $(cygpath -w "${WHEELS[@]}")
+  else
+    "$JAXCI_PYTHON" -m uv pip install "${WHEELS[@]}"
+  fi
 else
-  "$JAXCI_PYTHON" -m uv pip install "${WHEELS[@]}"
+  # Note that we don't exit here because the wheels may have been installed
+  # earlier in a different step in the CI job.
+  echo "INFO: No wheels found under $JAXCI_OUTPUT_DIR"
+  echo "INFO: Skipping local wheel installation."
 fi