Merge pull request #255 from ROCm/ci-upstream-sync-135_1

CI: 03/04/25 upstream sync

Author: charleshofer

.bazelrc

@@ -1,6 +1,12 @@
 # #############################################################################
 # All default build options below. These apply to all build commands.
 # #############################################################################
+# TODO: Enable Bzlmod
+common --noenable_bzlmod
+
+# TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260
+common --noincompatible_enable_cc_toolchain_resolution
+
 # Make Bazel print out all options from rc files.
 common --announce_rc
 
@@ -41,6 +47,13 @@ build:linux --copt=-Wno-array-parameter
 build:macos --config=posix
 build:macos --apple_platform_type=macos
 
+# Bazel 7.0.0 no longer supports dynamic symbol lookup on macOS. To resolve
+# undefined symbol errors in macOS arm64 builds, explicitly add the necessary
+# linker flags until dependencies are well defined. See
+# https://github.com/bazelbuild/bazel/issues/19730.
+build:macos --linkopt=-Wl,-undefined,dynamic_lookup
+build:macos --host_linkopt=-Wl,-undefined,dynamic_lookup
+
 # Windows has a relatively short command line limit, which JAX has begun to hit.
 # See https://docs.bazel.build/versions/main/windows.html
 build:windows --features=compiler_param_file
@@ -267,6 +280,12 @@ build:resultstore --bes_instance_name="tensorflow-testing"
 build:resultstore --bes_results_url="https://source.cloud.google.com/results/invocations"
 build:resultstore --bes_timeout=600s
 
+# Configs for RBE cache. When using resultstore, we need to use these configs
+# as well to ensure that the logs that get uploaded to resultstore can be read
+# without any errors.
+build:rbe_cache --remote_cache=remotebuildexecution.googleapis.com
+build:rbe_cache --remote_instance_name=projects/tensorflow-testing/instances/default_instance
+
 build:rbe --config=resultstore
 build:rbe --repo_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1
 build:rbe --define=EXECUTOR=remote

.bazelversion

@@ -1 +1 @@
-6.5.0
+7.4.1

.github/workflows/asan.yaml

@@ -90,4 +90,4 @@ jobs:
           echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
           echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
           # The LD_PRELOAD works around https://github.com/google/sanitizers/issues/934#issuecomment-649516500
-          LD_PRELOAD=/lib/x86_64-linux-gnu/libstdc++.so.6 python -m pytest -n auto --tb=short --maxfail=20 tests
+          LD_PRELOAD=/lib/x86_64-linux-gnu/libstdc++.so.6 python -m pytest -n 32 --tb=short --maxfail=20 tests

.github/workflows/build_artifacts.yml

@@ -17,7 +17,7 @@ on:
         options:
         - "linux-x86-n2-16"
         - "linux-arm64-c4a-64"
-        - "windows-x86-n2-16"
+        - "windows-x86-n2-64"
       artifact:
         description: "Which JAX artifact to build?"
         type: choice
@@ -119,11 +119,11 @@ jobs:
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-      - name: Enable RBE if building on Linux x86 or Windows x86
-        if: contains(inputs.runner, 'linux-x86') || contains(inputs.runner, 'windows-x86')
+      - name: Enable RBE if building on Linux x86
+        if: contains(inputs.runner, 'linux-x86')
         run: echo "JAXCI_BUILD_ARTIFACT_WITH_RBE=1" >> $GITHUB_ENV
-      - name: Enable Bazel remote cache (with writes enabled) if building on Linux Aarch64
-        if: contains(inputs.runner, 'linux-arm64')
+      - name: Enable Bazel remote cache (with writes enabled) if building on Linux Aarch64 or Windows x86
+        if: contains(inputs.runner, 'linux-arm64') || contains(inputs.runner, 'windows-x86')
         run: echo "JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE=1" >> $GITHUB_ENV
       # Halt for testing
       - name: Wait For Connection

.github/workflows/ci-build.yaml

@@ -80,7 +80,7 @@ jobs:
         JAX_SKIP_SLOW_TESTS: true
         PY_COLORS: 1
       run: |
-        pip install -e .
+        uv pip install --system -e .
         echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
         echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
         echo "JAX_ENABLE_CUSTOM_PRNG=$JAX_ENABLE_CUSTOM_PRNG"
@@ -171,7 +171,7 @@ jobs:
         JAX_SKIP_SLOW_TESTS: true
         PY_COLORS: 1
       run: |
-        pip install -e .
+        uv pip install --system -e .
         echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
         echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
         echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"

.github/workflows/cloud-tpu-ci-nightly.yml

@@ -33,7 +33,7 @@ jobs:
         python-version: ["3.10"]
     name: "TPU test (jaxlib=${{ matrix.jaxlib-version }}, ${{ matrix.tpu.type }})"
     env:
-      LIBTPU_OLDEST_VERSION_DATE: 20241118
+      LIBTPU_OLDEST_VERSION_DATE: 20241205
       PYTHON: python${{ matrix.python-version }}
     runs-on: ${{ matrix.tpu.runner }}
     container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest"

.github/workflows/pytest_cpu.yml

@@ -79,11 +79,15 @@ jobs:
 
           # Get the major and minor version of Python.
           # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.10, then python_major_minor=310
-          python_major_minor=$(echo "$JAXCI_HERMETIC_PYTHON_VERSION" | tr -d '.')
+          # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.13-nogil, then python_major_minor=313t
+          python_major_minor=$(echo "${JAXCI_HERMETIC_PYTHON_VERSION//-nogil/t}" | tr -d '.')
 
           echo "OS=${os}" >> $GITHUB_ENV
           echo "ARCH=${arch}" >> $GITHUB_ENV
-          echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV
+          # Python wheels follow a naming convention: standard wheels use the pattern
+          # `*-cp<py_version>-cp<py_version>-*`, while free-threaded wheels use
+          # `*-cp<py_version>-cp<py_version>t-*`.
+          echo "PYTHON_MAJOR_MINOR=cp${python_major_minor%t}-cp${python_major_minor}-" >> $GITHUB_ENV
       - name: Download jaxlib wheel from GCS (non-Windows runs)
         id: download-wheel-artifacts-nw
         # Set continue-on-error to true to prevent actions from failing the workflow if this step
@@ -125,7 +129,19 @@ jobs:
           echo "Skipping the test run."
           exit 1
       - name: Install Python dependencies
-        run: $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+        run: |
+          # TODO(srnitin): Remove after uv is installed in the Windows Dockerfile
+          $JAXCI_PYTHON -m pip install uv~=0.5.30
+          # python 3.13t cannot compile zstandard 0.23.0 due to
+          # https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
+          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
+          # in requirements.in.
+          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
+            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
+            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
+          else
+            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+          fi
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/pytest_cuda.yml

@@ -77,11 +77,15 @@ jobs:
 
           # Get the major and minor version of Python.
           # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.10, then python_major_minor=310
-          python_major_minor=$(echo "$JAXCI_HERMETIC_PYTHON_VERSION" | tr -d '.')
+          # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.13-nogil, then python_major_minor=313t
+          python_major_minor=$(echo "${JAXCI_HERMETIC_PYTHON_VERSION//-nogil/t}" | tr -d '.')
 
           echo "OS=${os}" >> $GITHUB_ENV
           echo "ARCH=${arch}" >> $GITHUB_ENV
-          echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV
+          # Python wheels follow a naming convention: standard wheels use the pattern
+          # `*-cp<py_version>-cp<py_version>-*`, while free-threaded wheels use
+          # `*-cp<py_version>-cp<py_version>t-*`.
+          echo "PYTHON_MAJOR_MINOR=cp${python_major_minor%t}-cp${python_major_minor}-" >> $GITHUB_ENV
       - name: Download the wheel artifacts from GCS
         id: download-wheel-artifacts
         # Set continue-on-error to true to prevent actions from failing the workflow if this step
@@ -106,7 +110,17 @@ jobs:
           echo "Skipping the test run."
           exit 1
       - name: Install Python dependencies
-        run: $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+        run: |
+          # python 3.13t cannot compile zstandard 0.23.0 due to
+          # https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
+          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
+          # in requirements.in.
+          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
+            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
+            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
+          else
+            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+          fi
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/tsan-suppressions.txt

@@ -17,6 +17,7 @@ race:dump_traceback
 # https://github.com/python/cpython/issues/128137
 # Fixed in Python 3.14, but not backported to 3.13.
 race:immortalize_interned
+race:_PyUnicode_InternMortal
 
 # https://github.com/python/cpython/issues/128144
 # Fixed in Python 3.14, but not backported to 3.13.
@@ -31,6 +32,7 @@ race:split_keys_entry_added
 # https://github.com/python/cpython/issues/129748
 race:mi_block_set_nextx
 
+
 # Races because the LAPACK and BLAS in our scipy isn't TSAN instrumented.
 race:heevd_ffi
 race:gesdd_ffi
@@ -59,3 +61,6 @@ race:gemm_oncopy
 
 # https://github.com/python/cpython/issues/128133
 # race:bytes_hash
+
+# https://github.com/python/cpython/issues/130571
+# race:_PyObject_GetMethod

.github/workflows/tsan.yaml

@@ -199,7 +199,5 @@ jobs:
               --local_test_jobs=32 \
               --test_timeout=600 \
               --config=resultstore \
-              --spawn_strategy=local \
-              --remote_cache=remotebuildexecution.googleapis.com \
-              --remote_instance_name=projects/tensorflow-testing/instances/default_instance \
+              --config=rbe_cache \
               //tests:cpu_tests