tsan.yaml

name: CI - Free-threading and Thread Sanitizer (nightly)

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

on:
  schedule:
    - cron: "0 12 * * *" # Daily at 12:00 UTC
  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:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          path: jax
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: python/cpython
          path: cpython
          ref: v3.13.0
      - 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
      - name: Build CPython with TSAN enabled
        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
          # 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()"
          ${GITHUB_WORKSPACE}/cpython-tsan/bin/python3 -m venv ${GITHUB_WORKSPACE}/venv
      - name: Install JAX test requirements
        run: |
          source ${GITHUB_WORKSPACE}/venv/bin/activate
          cd jax
          python -m pip install -r build/test-requirements.txt
      - name: Build and install JAX
        run: |
          source ${GITHUB_WORKSPACE}/venv/bin/activate
          cd jax
          python build/build.py \
            --bazel_options=--color=yes \
            --bazel_options=--copt=-fsanitize=thread \
            --bazel_options=--linkopt="-fsanitize=thread" \
            --bazel_options=--@rules_python//python/config_settings:py_freethreaded="yes" \
            --bazel_options=--define=use_enabled_free_threading=true \
            --clang_path=/usr/bin/clang-18
          # We have to manually install nightly scipy, otherwise default scipy installation
          # is failing to build it here: ../meson.build:84:0: ERROR: Unknown compiler(s)
          python -m pip install -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple scipy
          python -m pip install dist/jaxlib-*.whl
          python -m pip install -e .
      - name: Run tests
        env:
          JAX_NUM_GENERATED_CASES: 1
          JAX_ENABLE_X64: true
          JAX_SKIP_SLOW_TESTS: true
          PY_COLORS: 1
        run: |
          source ${GITHUB_WORKSPACE}/venv/bin/activate
          cd jax
          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"
          # As we do not have yet free-threading support
          # there will be the following warning:
          # RuntimeWarning: The global interpreter lock (GIL) has been enabled to load module 'jaxlib.utils',
          # which has not declared that it can run safely without the GIL.
          # To avoid that we temporarily define PYTHON_GIL
          export PYTHON_GIL=0

          python -m pytest -s -vvv \
            tests/jaxpr_effects_test.py::EffectOrderingTest::test_different_threads_get_different_tokens \
            tests/api_test.py::CustomJVPTest::test_concurrent_initial_style \
            tests/api_test.py::APITest::test_concurrent_device_get_and_put \
            tests/api_test.py::JitTest::test_concurrent_jit