Added CI job with TSAN and free-threading

.github/workflows/tsan.yaml

@@ -0,0 +1,85 @@
+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
+          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"
+          python -m pytest -n auto --tb=short --maxfail=20 tests

WORKSPACE

@@ -3,8 +3,19 @@ load("//third_party/xla:workspace.bzl", jax_xla_workspace = "repo")
 jax_xla_workspace()
 
 # Initialize hermetic Python
-load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
-python_init_rules()
+# load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
+# python_init_rules()
+
+# THIS CAN BE REMOVED WHEN MERGED: https://github.com/openxla/xla/pull/19346
+# Use latest rules_python
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+http_archive(
+    name = "rules_python",
+    sha256 = "62ddebb766b4d6ddf1712f753dac5740bea072646f630eb9982caa09ad8a7687",
+    strip_prefix = "rules_python-0.39.0",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.39.0/rules_python-0.39.0.tar.gz",
+)
+
 
 load("@xla//third_party/py:python_init_repositories.bzl", "python_init_repositories")
 python_init_repositories(