Use a local clang config for TSAN builds

yuriivcs · Google-ML-Automation · commit 2f0f001b4bbc · 2026-02-03T18:13:33.000-08:00
Update the TSAN GitHub workflow to use `rbe_linux_x86_64_clang_local`, which ensures clang-based builds run locally, as required for ThreadSanitizer.

PiperOrigin-RevId: 864992841
diff --git a/.bazelrc b/.bazelrc
@@ -292,7 +292,7 @@ common:ci_linux_x86_64 --color=yes
 # toolchain for both CPU and GPU builds.
 common:ci_linux_x86_64_clang_local --config=ci_linux_x86_64
 common:ci_linux_x86_64_clang_local --config=clang_local
-common:ci_linux_x86_64_clang_local --repo_env=TF_SYSROOT="/dt9"
+#common:ci_linux_x86_64_clang_local --repo_env=TF_SYSROOT="/dt9"
 
 # Clang path needs to be set for remote toolchain to be configured correctly.
 common:ci_linux_x86_64_clang_local --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-18/bin/clang"
@@ -429,6 +429,8 @@ common:rbe_linux_x86_64_base --platforms="@ml_build_config_platform//:platform"
 common:rbe_linux_x86_64 --config=rbe_linux_x86_64_base
 common:rbe_linux_x86_64 --config=ci_linux_x86_64
 
+common:rbe_linux_x86_64_clang_local --config=ci_linux_x86_64_clang_local
+
 common:rbe_linux_x86_64_cuda_common --config=rbe_linux_x86_64_base
 common:rbe_linux_x86_64_cuda_common --config=rbe_gpu_pool
 # Update UMD version when RBE CUDA driver is updated.
diff --git a/.github/workflows/tsan.yaml b/.github/workflows/tsan.yaml
@@ -5,7 +5,7 @@ concurrency:
   cancel-in-progress: true
 on:
   schedule:
-    - cron: "0 5 * * *"  # Daily at 05:00 UTC == 00:00 EST == 21:00 PST
+    - cron: "0 5 * * *" # Daily at 05:00 UTC == 00:00 EST == 21:00 PST
   workflow_dispatch: # allows triggering the workflow run manually
   pull_request: # Automatically trigger on pull requests affecting this file
     branches:
@@ -14,20 +14,19 @@ on:
       - '**/workflows/tsan.yaml'
       - '**/workflows/tsan-suppressions*.txt'
 permissions: {}
-
-env:
-  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
-  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
-
 jobs:
   tsan:
     runs-on: linux-x86-n4-64
     container:
-      image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build@sha256:ea67e8453d8b09c2ba48853da5e79efef4b65804b4a48dfae4b4da89ffd38405 # ml-build container (based on Ubuntu 22.04)
+      image: index.docker.io/library/ubuntu@sha256:b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b # ratchet:ubuntu:24.04
     strategy:
       fail-fast: false
       matrix:
         include:
+          - name-prefix: "with 3.13"
+            python-version: "3.13"
+            github_branch: "3.13"
+            requirements_lock_name: "requirements_lock_3_13_ft"
           - name-prefix: "with 3.14"
             python-version: "3.14"
             github_branch: "3.14"
@@ -38,56 +37,32 @@ jobs:
     steps:
       # Install git before actions/checkout as otherwise it will download the code with the GitHub
       # REST API and therefore any subsequent git commands will fail.
-      # Also install dependencies for Google Cloud SDK (curl, python3, etc)
-      - name: Install dependencies
+      - name: Install clang 18
         env:
           DEBIAN_FRONTEND: noninteractive
         run: |
-          apt-get update
-          apt-get install -q -y \
-            clang-18 \
-            libclang-common-18-dev \
-            libclang-rt-18-dev \
-            libc++abi-18-dev \
-            lld-18 \
-            libstdc++-12-dev \
-            libc++-18-dev \
-            build-essential \
-            libssl-dev \
-            zlib1g-dev \
-            libbz2-dev \
-            libreadline-dev \
-            libsqlite3-dev \
-            curl \
-            git \
-            libncursesw5-dev \
-            xz-utils \
-            tk-dev \
-            libxml2-dev \
-            libxmlsec1-dev \
-            libffi-dev \
-            liblzma-dev \
-            file \
-            vim \
-            wget \
-            zip \
-            zstd
-      - name: Install Google Cloud SDK
-        run: |
-          echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
-          curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -
-          apt update && apt install -y google-cloud-cli
-
-      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+          apt update
+          apt install -q -y clang-18 libstdc++-14-dev build-essential libssl-dev \
+            zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev curl git \
+            libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev \
+            libffi-dev liblzma-dev file zip
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           path: jax
           persist-credentials: false
-      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           repository: numpy/numpy
           path: numpy
           submodules: true
           persist-credentials: false
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        if: ${{ matrix.python-version == '3.14' }}
+        with:
+          repository: scipy/scipy
+          path: scipy
+          submodules: true
+          persist-credentials: false
 
       - name: Get year & week number
         id: get-date
@@ -96,13 +71,13 @@ jobs:
 
       - name: Restore cached TSAN CPython ${{ matrix.python-version }}
         id: cache-cpython-tsan-restore
-        uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+        uses: actions/cache/restore@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
         with:
           path: |
             ./python-tsan.tgz
           key: ${{ runner.os }}-cpython-tsan-${{ matrix.python-version }}-${{ steps.get-date.outputs.date }}
 
-      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         if: steps.cache-cpython-tsan-restore.outputs.cache-hit != 'true'
         with:
           repository: python/cpython
@@ -116,7 +91,7 @@ jobs:
         run: |
           cd cpython
           mkdir ${GITHUB_WORKSPACE}/cpython-tsan
-          CC=clang-18 CXX=clang++-18 CFLAGS=" -O0 -g" CXXFLAGS="-stdlib=libc++" ./configure --prefix ${GITHUB_WORKSPACE}/cpython-tsan --disable-gil --with-thread-sanitizer --with-mimalloc
+          CC=clang-18 CXX=clang++-18 ./configure --prefix ${GITHUB_WORKSPACE}/cpython-tsan --disable-gil --with-thread-sanitizer
           make -j64
           make install -j64
           # Check whether free-threading mode is enabled
@@ -128,38 +103,22 @@ jobs:
       - name: Save TSAN CPython ${{ matrix.python-version }}
         id: cache-cpython-tsan-save
         if: steps.cache-cpython-tsan-restore.outputs.cache-hit != 'true'
-        uses: actions/cache/save@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+        uses: actions/cache/save@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
         with:
           path: |
             ./python-tsan.tgz
           key: ${{ runner.os }}-cpython-tsan-${{ matrix.python-version }}-${{ steps.get-date.outputs.date }}
 
-      # Upload the Python tarball to GCS so RBE can access it.
-      - name: Upload TSAN CPython to GCS
-        run: |
-          GCS_DEST="gs://general-ml-ci-transient/jax-github-actions/jax/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}"
-          echo "Uploading python-tsan.tgz to $GCS_DEST"
-          gcloud storage cp python-tsan.tgz "$GCS_DEST/"
-
-          # Output the HTTP URL for Bazel
-          BASE_URL="https://storage.googleapis.com/general-ml-ci-transient/jax-github-actions/jax/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}/python-tsan.tgz"
-          # URL-encode the path (handles spaces in workflow name)
-          PUBLIC_URL=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${BASE_URL}', safe=':/'))")
-          echo "HERMETIC_PYTHON_URL=$PUBLIC_URL" >> $GITHUB_ENV
-
-      # --- Condtional NumPy steps for pre-release Python ---
-
       - name: Restore cached TSAN Numpy
         id: cache-numpy-tsan-restore
-        if: matrix.python-version == '3.15'
-        uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+        uses: actions/cache/restore@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
         with:
           path: |
             ./wheelhouse
           key: ${{ runner.os }}-numpy-tsan-${{ matrix.python-version }}-${{ hashFiles('numpy/pyproject.toml') }}-${{ steps.get-date.outputs.date }}
 
       - name: Build TSAN Numpy wheel
-        if: matrix.python-version == '3.15' && steps.cache-numpy-tsan-restore.outputs.cache-hit != 'true'
+        if: steps.cache-numpy-tsan-restore.outputs.cache-hit != 'true'
         run: |
           set -eux
           cd numpy
@@ -177,6 +136,7 @@ jobs:
           export PATH=${GITHUB_WORKSPACE}/cpython-tsan/bin/:$PATH
 
           python3 -m pip install uv~=0.5.30
+          python3 -m pip install --upgrade pip
           python3 -m uv pip install -r requirements/build_requirements.txt
 
           CC=clang-18 CXX=clang++-18 python3 -m pip wheel --wheel-dir dist -v . --no-build-isolation -Csetup-args=-Db_sanitize=thread -Csetup-args=-Dbuildtype=debugoptimized
@@ -205,17 +165,91 @@ jobs:
 
       - name: Save TSAN Numpy wheel
         id: cache-numpy-tsan-save
-        if: matrix.python-version == '3.15' && steps.cache-numpy-tsan-restore.outputs.cache-hit != 'true'
-        uses: actions/cache/save@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+        if: steps.cache-numpy-tsan-restore.outputs.cache-hit != 'true'
+        uses: actions/cache/save@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
         with:
           path: |
             ./wheelhouse
           key: ${{ runner.os }}-numpy-tsan-${{ matrix.python-version }}-${{ hashFiles('numpy/pyproject.toml') }}-${{ steps.get-date.outputs.date }}
 
-      # --- End Conditional NumPy Steps ---
+      - name: Restore cached Scipy
+        if: ${{ matrix.python-version == '3.14' }}
+        id: cache-scipy-restore
+        uses: actions/cache/restore@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
+        with:
+          path: |
+            ./wheelhouse
+          key: ${{ runner.os }}-scipy-${{ matrix.python-version }}-${{ hashFiles('scipy/pyproject.toml') }}-${{ steps.get-date.outputs.date }}
+
+      - name: Build Scipy wheel
+        if: ${{ steps.cache-scipy-restore.outputs.cache-hit != 'true' && matrix.python-version == '3.14' }}
+        env:
+          DEBIAN_FRONTEND: noninteractive
+        run: |
+          # Install scipy dependencies:
+          apt install -q -y gfortran libopenblas-dev liblapack-dev pkg-config --no-install-recommends
+
+          cd scipy
+
+          # If we restored cpython from cache, we need to get python interpreter from python-tsan.tgz
+          if [ ! -d ${GITHUB_WORKSPACE}/cpython-tsan/bin/ ]; then
+            echo "Extract cpython from python-tsan.tgz"
+            pushd .
+            ls ${GITHUB_WORKSPACE}/python-tsan.tgz
+            cd ${GITHUB_WORKSPACE} && tar -xzf python-tsan.tgz
+            ls ${GITHUB_WORKSPACE}/cpython-tsan/bin/
+            popd
+          fi
+
+          export PATH=${GITHUB_WORKSPACE}/cpython-tsan/bin/:$PATH
+
+          python3 -m pip install uv~=0.5.30
+          python3 -m pip install --upgrade pip
+
+          python3 -m uv pip install -U --pre numpy --extra-index-url file://${GITHUB_WORKSPACE}/wheelhouse/
+          python3 -m uv pip install cython pythran pybind11 meson-python ninja
+
+          python3 -m uv pip list | grep -E "(numpy|pythran|cython|pybind11)"
+
+          export CC=clang-18
+          export CXX=clang++-18
+          python3 -m pip wheel --wheel-dir dist -vvv . --no-build-isolation --no-deps -Csetup-args=-Dbuildtype=debugoptimized
+
+          # Create simple index and copy the wheel
+          mkdir -p ${GITHUB_WORKSPACE}/wheelhouse/scipy
+
+          scipy_whl_name=($(cd dist && ls scipy*.whl))
+          if [ -z "${scipy_whl_name}" ]; then exit 1; fi
+
+          echo "Built TSAN Scipy wheel: ${scipy_whl_name}"
+
+          cp dist/${scipy_whl_name} ${GITHUB_WORKSPACE}/wheelhouse/scipy
+
+          # Recreate wheelhouse index with Numpy and Scipy
+          cat << EOF > ${GITHUB_WORKSPACE}/wheelhouse/index.html
+          <!DOCTYPE html><html><body>
+          <a href="numpy">numpy></a></br>
+          <a href="scipy">scipy></a></br>
+          </body></html>
+          EOF
+
+          cat << EOF > ${GITHUB_WORKSPACE}/wheelhouse/scipy/index.html
+          <!DOCTYPE html><html><body>
+          <a href="${scipy_whl_name}">${scipy_whl_name}</a></br>
+          </body></html>
+          EOF
+
+      - name: Save Scipy wheel
+        id: cache-scipy-save
+        if: ${{ steps.cache-scipy-restore.outputs.cache-hit != 'true' && matrix.python-version == '3.14' }}
+        uses: actions/cache/save@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
+        with:
+          path: |
+            ./wheelhouse
+          key: ${{ runner.os }}-scipy-${{ matrix.python-version }}-${{ hashFiles('scipy/pyproject.toml') }}-${{ steps.get-date.outputs.date }}
 
       - name: Build Jax and run tests
-        timeout-minutes: 180
+        timeout-minutes: 120
         env:
           JAX_NUM_GENERATED_CASES: 1
           JAX_ENABLE_X64: true
@@ -226,57 +260,70 @@ jobs:
           set -x
           cd jax
 
-          # Calculate SHA256 of the Python tarball
-          export PYTHON_SHA256=($(sha256sum ${GITHUB_WORKSPACE}/python-tsan.tgz | awk '{ print $1 }'))
+          export PYTHON_SHA256=($(sha256sum ${GITHUB_WORKSPACE}/python-tsan.tgz))
           echo "Python sha256: ${PYTHON_SHA256}"
-          echo "Python URL: $HERMETIC_PYTHON_URL"
 
-          # Configure Bazel
           python3 build/build.py build --configure_only \
             --python_version=${{ matrix.python-version }}-ft \
-            --bazel_options=--repo_env=HERMETIC_PYTHON_URL="${HERMETIC_PYTHON_URL}" \
+            --bazel_options=--repo_env=HERMETIC_PYTHON_URL="file://${GITHUB_WORKSPACE}/python-tsan.tgz" \
             --bazel_options=--repo_env=HERMETIC_PYTHON_SHA256=${PYTHON_SHA256} \
             --bazel_options=--repo_env=HERMETIC_PYTHON_PREFIX="cpython-tsan/" \
             --bazel_options=--color=yes \
             --bazel_options=--copt=-fsanitize=thread \
             --bazel_options=--linkopt="-fsanitize=thread" \
-            --bazel_options=--copt=-g
+            --bazel_options=--copt=-g \
+            --clang_path=/usr/bin/clang-18
 
           mkdir -p dist
+          # Check whether we have numpy wheel or exit with error
+          ls ${GITHUB_WORKSPACE}/wheelhouse/numpy/*.whl || exit 1
+          cp -v ${GITHUB_WORKSPACE}/wheelhouse/numpy/*.whl dist/
+          if [ "${{ matrix.python-version }}" == "3.14" ]; then
+            # Check whether we have scipy wheel or exit with error
+            ls ${GITHUB_WORKSPACE}/wheelhouse/scipy/*.whl || exit 1
+            cp -v ${GITHUB_WORKSPACE}/wheelhouse/scipy/*.whl dist/
+
+            # Patch build/requirements_lock_3_14_ft.txt to use TSAN instrumented NumPy and Scipy
+            sed -i "s|--extra-index-url.*|--extra-index-url file://${GITHUB_WORKSPACE}/wheelhouse/|" build/${{ matrix.requirements_lock_name }}.txt
 
-          # Copy custom Numpy wheel only if using 3.15 (and if it exists)
-          if [ "${{ matrix.python-version }}" == "3.15" ]; then
-            ls ${GITHUB_WORKSPACE}/wheelhouse/numpy/*.whl || exit 1
-            cp -v ${GITHUB_WORKSPACE}/wheelhouse/numpy/*.whl dist/
+            # We should install jpeg dev package to be able to build Pillow from source:
+            apt install -q -y libjpeg-dev --no-install-recommends
 
-            # Patch requirements lock to use TSAN-instrumented NumPy
-            sed -i "s|--extra-index-url.*|--extra-index-url file://${GITHUB_WORKSPACE}/wheelhouse/|" build/${{ matrix.requirements_lock_name }}.txt
+            # Install scipy runtime dependencies (in case we restore scipy wheel from cache):
+            apt install -q -y libopenblas-dev liblapack-dev --no-install-recommends
           fi
 
           echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
           echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
           echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
 
+          # Set symlink to the bazel executable
+          bazel_exec=($(ls bazel-*))
+          ln -s ${bazel_exec} bazel
 
           # Check python version
-          python_version="${{ matrix.python-version }}"
-          bazel run --config=rbe_linux_x86_64 --@rules_python//python/config_settings:py_freethreaded="yes" @python_${python_version//./_}_host//:python -- -VV
+          ./bazel run --@rules_python//python/config_settings:py_freethreaded="yes" @python//:python3 -- -VV
 
           # Check numpy version
-          bazel cquery @pypi_numpy//:* | grep whl
+          ./bazel cquery @pypi_numpy//:* | grep whl
+
+          if [ "${{ matrix.python-version }}" == "3.14" ]; then
+            # Check scipy version
+            ./bazel cquery @pypi_scipy//:* | grep whl
+          fi
 
           # Build JAX and run tests
-          bazel test \
+          ./bazel test \
               --test_env=JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES \
               --test_env=JAX_ENABLE_X64=$JAX_ENABLE_X64 \
               --test_env=JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS \
               --test_env=PYTHON_GIL=0 \
-              --test_env=TSAN_OPTIONS="halt_on_error=1,suppressions=tests/config/tsan-suppressions_${{ matrix.python-version }}.txt" \
+              --test_env=TSAN_OPTIONS=halt_on_error=1,suppressions=$PWD/.github/workflows/tsan-suppressions_${{ matrix.python-version }}.txt \
               --test_env=JAX_TEST_NUM_THREADS=8 \
-              --@rules_python//python/config_settings:py_freethreaded="yes" \
               --test_output=errors \
+              --local_test_jobs=32 \
               --test_timeout=1800 \
-              --config=rbe_linux_x86_64 \
-              --test_tag_filters=-notsan \
-              --run_under=//tests/config:oss_tsan_wrapper_sh \
-              //tests:cpu_tests
+              --config=resultstore \
+              --config=ci_rbe_cache \
+              --config=clang_local \
+              //tests:cpu_tests
