Merge pull request #26195 from vfdev-5:use-tsan-numpy-in-tsan-ci

PiperOrigin-RevId: 721899249
jax-ml · Jan 31, 2025 · 3f73758 · 3f73758
2 parents 0bf8462 + f4199fc
commit 3f73758
Showing 1 changed file with 77 additions and 2 deletions.
diff --git a/.github/workflows/tsan.yaml b/.github/workflows/tsan.yaml
@@ -44,16 +44,21 @@ jobs:
           repository: python/cpython
           path: cpython
           ref: "3.13"
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          repository: numpy/numpy
+          path: numpy
+          submodules: true
 
-      - name: Restore cached CPython with TSAN
+      - name: Restore cached TSAN CPython
         id: cache-cpython-tsan-restore
         uses: actions/cache/restore@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
         with:
           path: |
             ./python-tsan.tgz
           key: ${{ runner.os }}-cpython-tsan-${{ hashFiles('cpython/configure.ac') }}
 
-      - name: Build CPython with TSAN enabled
+      - name: Build CPython with enabled TSAN
         if: steps.cache-cpython-tsan-restore.outputs.cache-hit != 'true'
         run: |
           cd cpython
@@ -76,6 +81,71 @@ jobs:
             ./python-tsan.tgz
           key: ${{ runner.os }}-cpython-tsan-${{ hashFiles('cpython/configure.ac') }}
 
+      - name: Get year & week number
+        id: get-date
+        run: echo "date=$(/bin/date "+%Y-%U")" >> $GITHUB_OUTPUT
+        shell: bash -l {0}
+
+      - name: Restore cached TSAN Numpy
+        id: cache-numpy-tsan-restore
+        uses: actions/cache/restore@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
+        with:
+          path: |
+            ./wheelhouse
+          key: ${{ runner.os }}-numpy-tsan-${{ hashFiles('numpy/pyproject.toml') }}-${{ steps.get-date.outputs.date }}
+
+      - name: Build TSAN Numpy wheel
+        if: steps.cache-numpy-tsan-restore.outputs.cache-hit != 'true'
+        run: |
+          cd numpy
+
+          # 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"
+            ls ${GITHUB_WORKSPACE}/python-tsan.tgz
+            cd ${GITHUB_WORKSPACE} && tar -xvzf python-tsan.tgz
+            ls ${GITHUB_WORKSPACE}/cpython-tsan/bin/
+          fi
+
+          export PATH=${GITHUB_WORKSPACE}/cpython-tsan/bin/:$PATH
+
+          python3 -m pip install -r requirements/build_requirements.txt
+          # Make sure to install a compatible Cython version (master branch is best for now)
+          python3 -m pip install -U git+https://github.com/cython/cython
+
+          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
+
+          # Create simple index and copy the wheel
+          mkdir -p ${GITHUB_WORKSPACE}/wheelhouse/numpy
+
+          numpy_whl_name=($(cd dist && ls numpy*.whl))
+          if [ -z "${numpy_whl_name}" ]; then exit 1; fi
+
+          echo "Built TSAN Numpy wheel: ${numpy_whl_name}"
+
+          cp dist/${numpy_whl_name} ${GITHUB_WORKSPACE}/wheelhouse/numpy
+
+          cat << EOF > ${GITHUB_WORKSPACE}/wheelhouse/index.html
+          <!DOCTYPE html><html><body>
+          <a href="numpy">numpy></a></br>
+          </body></html>
+          EOF
+
+          cat << EOF > ${GITHUB_WORKSPACE}/wheelhouse/numpy/index.html
+          <!DOCTYPE html><html><body>
+          <a href="${numpy_whl_name}">${numpy_whl_name}</a></br>
+          </body></html>
+          EOF
+
+      - name: Save TSAN Numpy wheel
+        id: cache-numpy-tsan-save
+        if: steps.cache-numpy-tsan-restore.outputs.cache-hit != 'true'
+        uses: actions/cache/save@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
+        with:
+          path: |
+            ./wheelhouse
+          key: ${{ runner.os }}-numpy-tsan-${{ hashFiles('numpy/pyproject.toml') }}-${{ steps.get-date.outputs.date }}
+
       - name: Build Jax and run tests
         timeout-minutes: 120
         env:
@@ -101,6 +171,11 @@ jobs:
             --bazel_options=--copt=-g \
             --clang_path=/usr/bin/clang-18
 
+          # Update the patch to use TSAN instrumented numpy
+
+          sed -i "s|+--extra-index-url.*|+--extra-index-url file://${GITHUB_WORKSPACE}/wheelhouse/|" .github/workflows/requirements_lock_3_13_ft.patch
+          cat .github/workflows/requirements_lock_3_13_ft.patch
+
           # Apply a patch to numpy in requirements lock 3.13 ft to use the nightly version
           git apply .github/workflows/requirements_lock_3_13_ft.patch