Skip to content

CI - Free-threading and Thread Sanitizer (nightly) #68

CI - Free-threading and Thread Sanitizer (nightly)

CI - Free-threading and Thread Sanitizer (nightly) #68

Workflow file for this run

name: CI - Free-threading and Thread Sanitizer (nightly)
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
on:
schedule:
- 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:
- main
paths:
- '**/workflows/tsan.yaml'
jobs:
tsan:
runs-on: linux-x86-n2-64
container:
image: index.docker.io/library/ubuntu@sha256:b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b # ratchet:ubuntu:24.04
strategy:
fail-fast: false
defaults:
run:
shell: bash -l {0}
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.
- name: Install clang 18
env:
DEBIAN_FRONTEND: noninteractive
run: |
apt update
apt install -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
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
path: jax
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
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 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 enabled TSAN
if: steps.cache-cpython-tsan-restore.outputs.cache-hit != 'true'
run: |
cd cpython
mkdir ${GITHUB_WORKSPACE}/cpython-tsan
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
PYTHON_GIL=0 ${GITHUB_WORKSPACE}/cpython-tsan/bin/python3 -c "import sys; assert not sys._is_gil_enabled()"
# Create archive to be used with bazel as hermetic python:
cd ${GITHUB_WORKSPACE} && tar -czpf python-tsan.tgz cpython-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
with:
path: |
./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:
JAX_NUM_GENERATED_CASES: 1
JAX_ENABLE_X64: true
JAX_SKIP_SLOW_TESTS: true
PY_COLORS: 1
run: |
cd jax
export PYTHON_SHA256=($(sha256sum ${GITHUB_WORKSPACE}/python-tsan.tgz))
echo "Python sha256: ${PYTHON_SHA256}"
python3 -VV
python3 build/build.py build --configure_only \
--python_version=3.13-ft \
--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 \
--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
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
./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=$PWD/.github/workflows/tsan-suppressions.txt \
--test_env=JAX_TEST_NUM_THREADS=8 \
--test_output=errors \
--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 \
//tests:cpu_tests