mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 10:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #27094 from vfdev-5:fix-tsan-numpy-install-patch

Browse Source 
PiperOrigin-RevId: 738080051
This commit is contained in:
jax authors 2025-03-18 11:24:15 -07:00
commit ee0073e605
2 changed files with 22 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
.github/workflows/requirements_lock_3_13_ft.patch vendored
View File

@ -1,8 +1,8 @@
diff --git a/build/requirements_lock_3_13_ft.txt b/build/requirements_lock_3_13_ft.txt
index dfefaf042..2700e140e 100644
index e7a2968e9..d37e11ee3 100644
--- a/build/requirements_lock_3_13_ft.txt
+++ b/build/requirements_lock_3_13_ft.txt
@@ -4,6 +4,12 @@
@@ -4,6 +4,11 @@
#
# pip-compile --allow-unsafe --generate-hashes --output-file=build/requirements_lock_3_13_ft.txt build/requirements.in
#
@ -10,12 +10,11 @@ index dfefaf042..2700e140e 100644
+--pre
+--extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple
+numpy
+
+
absl-py==2.1.0 \
--hash=sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308 \
--hash=sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff
@@ -328,68 +334,6 @@ mpmath==1.3.0 \
@@ -328,68 +333,6 @@ mpmath==1.3.0 \
--hash=sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f \
--hash=sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c
# via -r build/test-requirements.txt
@ -81,6 +80,6 @@ index dfefaf042..2700e140e 100644
- # matplotlib
- # ml-dtypes
- # scipy
opt-einsum==3.4.0 \
--hash=sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd \
--hash=sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac
nvidia-cublas-cu12==12.8.3.14 ; sys_platform == "linux" \
--hash=sha256:3f0e05e7293598cf61933258b73e66a160c27d59c4422670bf0b79348c04be44 \
--hash=sha256:93a4e0e386cc7f6e56c822531396de8170ed17068a1e18f987574895044cd8c3 \

19
.github/workflows/tsan.yaml vendored
View File

@ -173,12 +173,18 @@ jobs:
--bazel_options=--copt=-g \
--clang_path=/usr/bin/clang-18
# Update the patch to use TSAN instrumented numpy
# Patch build/requirements_lock_3_13_ft.txt 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
git apply .github/workflows/requirements_lock_3_13_ft.patch || exit 1
# 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
# Display the content for debugging in logs
cat build/requirements_lock_3_13_ft.txt | head -15
# Check the patch
cat build/requirements_lock_3_13_ft.txt | head -15 | grep -E "(--pre|.*${GITHUB_WORKSPACE}/wheelhouse/|numpy)"
if [ "$?" == "1" ]; then echo "Could not find the patch in the requirements_lock_3_13_ft.txt"; exit 1; fi
cat build/requirements_lock_3_13_ft.txt | grep -E "(numpy==)"
if [ "$?" == "0" ]; then "Found original numpy dependency in the requirements_lock_3_13_ft.txt"; exit 1; fi
echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
@ -188,6 +194,13 @@ jobs:
bazel_exec=($(ls bazel-*))
ln -s ${bazel_exec} bazel
# Check python version
./bazel run --@rules_python//python/config_settings:py_freethreaded="yes" @python//:python3 -- -VV
# Check numpy version
./bazel cquery @pypi_numpy//:* | grep whl
# Build JAX and run tests
./bazel test \
--test_env=JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES \
--test_env=JAX_ENABLE_X64=$JAX_ENABLE_X64 \