From 0f863b20926904273f6bc262c10a556559c1529b Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Thu, 3 Oct 2024 10:23:39 -0400
Subject: [PATCH] Change asan build to use a self-hosted runner and Ubuntu
 24.04.

---
 .github/workflows/asan.yaml | 40 ++++++++++++++++++++++---------------
 tests/logging_test.py       |  4 ++++
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/.github/workflows/asan.yaml b/.github/workflows/asan.yaml
index fd3054aabbef..ca2842edcad9 100644
--- a/.github/workflows/asan.yaml
+++ b/.github/workflows/asan.yaml
@@ -8,42 +8,48 @@ 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
+  pull_request: # Automatically trigger on pull requests affecting this file
     branches:
       - main
     paths:
-      - '**workflows/asan.yml'
+      - '**/workflows/asan.yml'
 
 jobs:
-  upstream-dev:
-    runs-on: ubuntu-20.04-16core
+  asan:
+    runs-on: linux-x86-n2-64
+    container:
+      image: index.docker.io/library/ubuntu@sha256:b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b # ratchet:ubuntu:24.04
     strategy:
       fail-fast: false
-      matrix:
-        python-version: ["3.12"]
+    defaults:
+      run:
+        shell: bash -l {0}
     steps:
-      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938  # v4.2.0
+      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938 # v4.2.0
         with:
           path: jax
-      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938  # v4.2.0
+      - uses: actions/checkout@d632683dd7b4114ad314bca15554477dd762a938 # v4.2.0
         with:
           repository: python/cpython
           path: cpython
           ref: v3.12.6
       - name: Install clang 18
+        env:
+          DEBIAN_FRONTEND: noninteractive
         run: |
-          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
-          echo deb http://apt.llvm.org/focal/ llvm-toolchain-focal-18 main | sudo tee -a /etc/apt/sources.list
-          sudo apt update
-          sudo apt install clang-18 libstdc++-10-dev
+          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 ASAN enabled
         env:
           ASAN_OPTIONS: detect_leaks=0
         run: |
           cd cpython
           mkdir ${GITHUB_WORKSPACE}/cpythonasan
-          CC=clang-18 ./configure --prefix ${GITHUB_WORKSPACE}/cpythonasan --with-address-sanitizer --without-pymalloc
-          make -j16
+          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
@@ -59,7 +65,10 @@ jobs:
         run: |
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
-          python build/build.py --bazel_options=--copt=-fsanitize=address --clang_path=/usr/bin/clang-18
+          python build/build.py \
+            --bazel_options=--color=yes
+            --bazel_options=--copt=-fsanitize=address \
+            --clang_path=/usr/bin/clang-18 \
           pip install dist/jaxlib-*.whl
           pip install -e .
       - name: Run tests
@@ -77,4 +86,3 @@ jobs:
           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 auto --tb=short --maxfail=20 tests
- 
\ No newline at end of file
diff --git a/tests/logging_test.py b/tests/logging_test.py
index 70f619de5ee6..a1d6695a1e37 100644
--- a/tests/logging_test.py
+++ b/tests/logging_test.py
@@ -92,10 +92,14 @@ def test_no_log_spam(self):
       python = sys.executable
       assert "python" in python
       env_variables = {"TF_CPP_MIN_LOG_LEVEL": "1"}
+      if os.getenv("ASAN_OPTIONS"):
+        env_variables["ASAN_OPTIONS"] = os.getenv("ASAN_OPTIONS")
       if os.getenv("PYTHONPATH"):
         env_variables["PYTHONPATH"] = os.getenv("PYTHONPATH")
       if os.getenv("LD_LIBRARY_PATH"):
         env_variables["LD_LIBRARY_PATH"] = os.getenv("LD_LIBRARY_PATH")
+      if os.getenv("LD_PRELOAD"):
+        env_variables["LD_PRELOAD"] = os.getenv("LD_PRELOAD")
       # Make sure C++ logging is at default level for the test process.
       proc = subprocess.run(
           [python, f.name],