Test python3.14 nightly CI workflows

.github/workflows/pytest_cpu.yml

@@ -50,7 +50,7 @@ jobs:
         # Explicitly set the shell to bash to override Windows's default (cmd)
         shell: bash
     runs-on: ${{ inputs.runner }}
-    container: ${{ (contains(inputs.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
+    container: ${{ (contains(inputs.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-rbe:752380653') ||
                    (contains(inputs.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') ||
                    (contains(inputs.runner, 'windows-x86') && null) }}
     name: "Pytest CPU (${{ inputs.runner }}, Python ${{ inputs.python }}, x64=${{ inputs.enable-x64 }})"
@@ -128,10 +128,28 @@ jobs:
       - name: Install Python dependencies
         run: |
           $JAXCI_PYTHON -m pip install uv~=0.5.30
+
+          # For py 3.14.0a7, some pre-built dependency wheels aren't available,
+          # so we need to download their deps or build them from source.
+          if [[ $JAXCI_PYTHON == "python3.14" ]]; then
+            echo "Building some dependencies from source for Python 3.14.0a7"
+            # Build numpy 2.3.0.dev0 from source at the current github repo head. This is
+            # needed to include fixes for https://github.com/numpy/numpy/issues/28681.
+            $JAXCI_PYTHON -m uv pip install "git+https://github.com/numpy/numpy@2ddc86163e6b194edb975629bdfc60bd5af40c0f"
+
+            # SciPy build requirements
+            apt update && apt-get install -y gfortran libopenblas-dev liblapack-dev pkg-config --no-install-recommends
+          else
+            echo "JAXCI_HERMETIC_PYTHON_VERSION is "$JAXCI_HERMETIC_PYTHON_VERSION
+          fi
+
+          echo "Check glibc version"
+          ldd --version
+
           $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
 
           # CPU Pytests crash with NumPy 2.2+ on Linux Aarch64; b/399168632
-          if [[ $OS == "linux" && $ARCH == "aarch64" ]]; then
+          if [[ $OS == "linux" && $ARCH == "aarch64" && $JAXCI_PYTHON != "python3.14" ]]; then
             $JAXCI_PYTHON -m uv pip install numpy~=2.1.0
           fi
           # python 3.13t cannot compile zstandard 0.23.0 due to

.github/workflows/pytest_cuda.yml

@@ -114,7 +114,22 @@ jobs:
           echo "Skipping the test run."
           exit 1
       - name: Install Python dependencies
-        run: $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
+        run: |
+          # For py 3.14.0a7, some pre-built dependency wheels aren't available,
+          # so we need to download their deps or build them from source.
+          if [[ $JAXCI_HERMETIC_PYTHON_VERSION == "python3.14" ]]; then
+            echo "Building some dependencies from source for Python 3.14.0a7"
+            # Build numpy 2.3.0.dev0 from source at the current github repo head. This is
+            # needed to include fixes for https://github.com/numpy/numpy/issues/28681.
+            $JAXCI_PYTHON -m uv pip install "git+https://github.com/numpy/numpy@2ddc86163e6b194edb975629bdfc60bd5af40c0f"
+
+            # SciPy build requirements
+            apt update && apt-get install -y gfortran libopenblas-dev liblapack-dev pkg-config --no-install-recommends
+          else
+            echo "JAXCI_HERMETIC_PYTHON_VERSION is "$JAXCI_HERMETIC_PYTHON_VERSION
+          fi
+
+          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/wheel_tests_nightly_release.yml

@@ -38,7 +38,7 @@ jobs:
           # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
           # that build the wheels.
           runner: ["linux-x86-n2-64", "linux-arm64-t2a-48", "windows-x86-n2-64"]
-          python: ["3.10","3.11", "3.12", "3.13", "3.13-nogil"]
+          python: ["3.10","3.11", "3.12", "3.13", "3.13-nogil", "3.14"]
           enable-x64: [0]
           exclude:
             - runner: "windows-x86-n2-64"
@@ -59,7 +59,7 @@ jobs:
           # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
           # that build the wheels.
           runner: ["linux-x86-g2-48-l4-4gpu"]
-          python: ["3.10","3.11", "3.12", "3.13", "3.13-nogil"]
+          python: ["3.10","3.11", "3.12", "3.13", "3.13-nogil", "3.14"]
           cuda: ["12.1", "12.8"]
           enable-x64: [0]
     name: "Pytest CUDA (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"

build/test-requirements.txt

@@ -6,7 +6,7 @@ filelock
 flatbuffers
 hypothesis
 mpmath>=1.3
-pillow>=10.4.0
+pillow>=10.4.0; python_version<="3.13"
 # TODO(kanglan): Remove once psutil from portpicker supports python 3.13t
 portpicker; python_version<"3.13"
 pytest-xdist
@@ -16,6 +16,6 @@ setuptools
 # matplotlib 3.9.0 pins NumPy 1.23, which is incompatible with the requirement
 # below.
 matplotlib~=3.8.4; python_version=="3.10"
-matplotlib; python_version>="3.11"
+matplotlib; python_version>="3.11" and python_version<="3.13"
 opt-einsum
 auditwheel