Add TSAN numpy step

jax-ml · Feb 3, 2025 · 024f4f6 · 024f4f6
1 parent cb188a0
commit 024f4f6
Showing 1 changed file with 68 additions and 1 deletion.
diff --git a/.github/workflows/tsan.yaml b/.github/workflows/tsan.yaml
@@ -67,7 +67,7 @@ jobs:
           # Create archive to be used with bazel as hermetic python:
           cd ${GITHUB_WORKSPACE} && tar -czpf python-tsan.tgz cpython-tsan
 
-      - name: Save CPython with TSAN
+      - name: Save TSAN CPython
         id: cache-cpython-tsan-save
         if: steps.cache-cpython-tsan-restore.outputs.cache-hit != 'true'
         uses: actions/cache/save@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0
@@ -76,6 +76,73 @@ 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"
+            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 -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: