From 8ba2657b325dccdbcabbaf792ad34ab835e5338d Mon Sep 17 00:00:00 2001
From: Ilya Panfilov <Ilya.Panfilov@amd.com>
Date: Fri, 16 Jan 2026 23:49:18 -0500
Subject: [PATCH] Do not fail CI on known failed JAX test

---
 ci/_utils.sh | 10 +++++++++-
 ci/jax.sh    |  8 +++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/ci/_utils.sh b/ci/_utils.sh
index d3d867d3f..966efbcad 100644
--- a/ci/_utils.sh
+++ b/ci/_utils.sh
@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 
@@ -25,6 +25,8 @@ export CI=1
 
 _script_error_count=0
 _run_error_count=0
+_ignored_error_count=0
+TEST_ERROR_IGNORE=""
 
 script_error() {
     _script_error_count=$((_script_error_count+1))
@@ -32,6 +34,11 @@ script_error() {
 }
 
 test_run_error() {
+    if [ -n "$TEST_ERROR_IGNORE" ]; then
+        _ignored_error_count=$((_ignored_error_count+1))
+        test -n "$@" && echo "Ignore error in test $@" >&2
+        return
+    fi
     _run_error_count=$((_run_error_count+1))
     test -n "$@" && echo "Error in test $@" >&2
 }
@@ -39,6 +46,7 @@ test_run_error() {
 return_run_results() {
     test $_script_error_count -ne 0 && echo Detected $_script_error_count script errors during tests run at level $TEST_LEVEL >&2
     test $_run_error_count -ne 0 && echo Got $_run_error_count test errors during run at level $TEST_LEVEL >&2
+    test $_ignored_error_count -ne 0 && echo Ignored $_ignored_error_count test errors during run at level $TEST_LEVEL >&2
     test $_run_error_count -eq 0 -a $_script_error_count -eq 0
 }
 
diff --git a/ci/jax.sh b/ci/jax.sh
index 6229a7de3..77ca1fba6 100755
--- a/ci/jax.sh
+++ b/ci/jax.sh
@@ -1,5 +1,5 @@
 #!/bin/sh
-# Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (c) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
 
@@ -79,7 +79,13 @@ run_test_config_mgpu() {
     else
         _dfa_level=3
     fi
+    # Do not fail automated CI if test_distributed_fused_attn is hung
+    # If the sctipt run w/o TEST_LEVEL the test error will be honored
+    if [ "$TEST_LEVEL" -le 3 ]; then
+        TEST_ERROR_IGNORE="1"
+    fi
     run $_dfa_level test_distributed_fused_attn.py $_timeout_args
+    TEST_ERROR_IGNORE=""
     run_default_fa 3 test_distributed_layernorm.py
     run_default_fa 2 test_distributed_layernorm_mlp.py $_timeout_args
     run_default_fa 3 test_distributed_softmax.py