From 2e659a5bed16b1138d627d7326451289a705a41a Mon Sep 17 00:00:00 2001
From: srinivamd <52507740+srinivamd@users.noreply.github.com>
Date: Tue, 9 Jun 2026 08:27:12 -0700
Subject: [PATCH] [ROCm] Deselect crashing FusedAttention bwd8 + add
 pytest-timeout

test_fused_attention_bwd8 causes a SIGSEGV in the XLA Triton compiler
(KernelReuseCache::GetWithStatus via ThunkEmitter::EmitTritonCustomCall)
on gfx942. pytest-xdist replaces the crashed worker, but the replacement
deadlocks during GPU runtime re-initialization, blocking the entire test
suite for 9+ hours until the job-level timeout fires.

Changes:
1. Deselect test_fused_attention_bwd8 in ci/run_pytest_rocm.sh (sgpu only).
   This variant is NOT covered by the existing in-test skipTest guards,
   which skip bwd{1,2,7,9} but not bwd8.
2. Add pytest-timeout to build/test-requirements.txt and pass
   --timeout=1800 (30 min) to both sgpu and mgpu pytest invocations.
   This is a safety net: if any future test hangs, it will be killed
   after 30 minutes instead of wasting the full 10-hour job budget.

Tracked in: ROCM-25626
Upstream: openxla/xla#38972, jax-ml/jax#34722
---
 build/test-requirements.txt | 1 +
 ci/run_pytest_rocm.sh       | 3 +++
 2 files changed, 4 insertions(+)
 mode change 100755 => 100644 ci/run_pytest_rocm.sh

diff --git a/build/test-requirements.txt b/build/test-requirements.txt
index c6975ec45acf..8987c394b0ae 100644
--- a/build/test-requirements.txt
+++ b/build/test-requirements.txt
@@ -7,6 +7,7 @@ mpmath==1.3.0  # TODO(dsuo): Remove pin sinc(inf) is resolved as either nan or 0
 pillow>=11.3
 portpicker
 pytest<9.0  # Works around https://github.com/pytest-dev/pytest/issues/13895
+pytest-timeout
 pytest-xdist
 rich
 matplotlib
diff --git a/ci/run_pytest_rocm.sh b/ci/run_pytest_rocm.sh
old mode 100755
new mode 100644
index e1abdb52e801..db76396125c7
--- a/ci/run_pytest_rocm.sh
+++ b/ci/run_pytest_rocm.sh
@@ -125,6 +125,7 @@ set +e
 
 # Run single-accelerator tests in parallel
 "$JAXCI_PYTHON" -m pytest -n $num_processes --tb=short \
+--timeout=1800 \
 --json-report --json-report-file=${LOGS_DIR}/pytest_results_single.json \
 --junitxml=test-artifacts/junit-single.xml \
 --dist=loadfile \
@@ -132,6 +133,7 @@ set +e
 --deselect=tests/multi_device_test.py::MultiDeviceTest::test_computation_follows_data \
 --deselect=tests/multiprocess_gpu_test.py::MultiProcessGpuTest::test_distributed_jax_visible_devices \
 --deselect=tests/compilation_cache_test.py::CompilationCacheTest::test_task_using_cache_metric \
+--deselect=tests/pallas/gpu_ops_test.py::FusedAttentionTest::test_fused_attention_bwd8 \
 tests
 
 first_cmd_retval=$?
@@ -141,6 +143,7 @@ if [[ $gpu_count -gt 1 ]]; then
   unset JAX_ENABLE_ROCM_XDIST
 
   "$JAXCI_PYTHON" -m pytest --tb=short \
+    --timeout=1800 \
     --json-report --json-report-file=${LOGS_DIR}/pytest_results_multi.json \
     --junitxml=test-artifacts/junit-multi.xml \
     -m "multiaccelerator" \