rocm_jax/.github/workflows/asan.yaml

name: CI - Address 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/asan.yaml'

jobs:
  asan:
    # Don't execute in fork due to runner type
    if: github.repository == 'jax-ml/jax'
    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: v3.13.0
      - name: Build CPython with ASAN enabled
        env:
          ASAN_OPTIONS: detect_leaks=0
        run: |
          cd cpython
          mkdir ${GITHUB_WORKSPACE}/cpythonasan
          CC=clang-18 CXX=clang++-18 ./configure --prefix ${GITHUB_WORKSPACE}/cpythonasan --with-address-sanitizer --without-pymalloc
          make -j64
          make install
          ${GITHUB_WORKSPACE}/cpythonasan/bin/python3 -m venv ${GITHUB_WORKSPACE}/venv
      - name: Install JAX test requirements
        env:
          ASAN_OPTIONS: detect_leaks=0
        run: |
          source ${GITHUB_WORKSPACE}/venv/bin/activate
          cd jax
          pip install uv~=0.5.30
          uv pip install -r build/test-requirements.txt
      - name: Build and install JAX
        env:
          ASAN_OPTIONS: detect_leaks=0
        run: |
          source ${GITHUB_WORKSPACE}/venv/bin/activate
          cd jax
          python build/build.py build --wheels=jaxlib --verbose \
            --bazel_options=--color=yes \
            --bazel_options=--copt=-fsanitize=address \
            --clang_path=/usr/bin/clang-18
          uv pip install dist/jaxlib-*.whl \
            -e .
      - name: Run tests
        env:
          ASAN_OPTIONS: detect_leaks=0
          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"
          # The LD_PRELOAD works around https://github.com/google/sanitizers/issues/934#issuecomment-649516500
          LD_PRELOAD=/lib/x86_64-linux-gnu/libstdc++.so.6 python -m pytest -n 32 --tb=short --maxfail=20 tests