From 93f11330d131541e196cc82623e09dd78e04e8aa Mon Sep 17 00:00:00 2001
From: Kamil Tokarski <ktokarski@nvidia.com>
Date: Tue, 16 Apr 2024 11:00:57 +0200
Subject: [PATCH] Disable some JAX iterator tests in sanitizer run (#5427)

Signed-off-by: Kamil Tokarski <ktokarski@nvidia.com>
---
 qa/TL0_FW_iterators/test_jax.sh | 17 ++++++++++++++---
 qa/setup_packages.py            |  2 +-
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/qa/TL0_FW_iterators/test_jax.sh b/qa/TL0_FW_iterators/test_jax.sh
index 6bcf4db5674..8872e37324d 100755
--- a/qa/TL0_FW_iterators/test_jax.sh
+++ b/qa/TL0_FW_iterators/test_jax.sh
@@ -10,12 +10,23 @@ do_once() {
 }
 
 test_body() {
-    # General tests for iterators
-    ${python_new_invoke_test} -A 'jax' test_fw_iterators
+    # Updating JAX 0.4.13 -> 0.4.16 (or the most recent 0.4.26) causes
+    # asan to fail with suposed stack-overflow
+    # TODO(ktokarski) Investigate what may be the underlying cause
+    if [ -z "$DALI_ENABLE_SANITIZERS" ]; then
+        # General tests for iterators
+        ${python_new_invoke_test} -A 'jax' test_fw_iterators
+    fi
 
     # More specific JAX tests
     ${python_new_invoke_test} -s jax_plugin/ test_integration test_iterator test_peekable_iterator
-    ${python_new_invoke_test} checkpointing.test_dali_checkpointing_fw_iterators.TestJax
+
+    # Updating JAX 0.4.13 -> 0.4.16 (or the most recent 0.4.26) causes
+    # asan to fail with suposed stack-overflow
+    # TODO(ktokarski) Investigate what may be the underlying cause
+    if [ -z "$DALI_ENABLE_SANITIZERS" ]; then
+        ${python_new_invoke_test} checkpointing.test_dali_checkpointing_fw_iterators.TestJax
+    fi
 }
 
 pushd ../..
diff --git a/qa/setup_packages.py b/qa/setup_packages.py
index 7d95208d80e..60d7c562234 100755
--- a/qa/setup_packages.py
+++ b/qa/setup_packages.py
@@ -565,7 +565,7 @@ def get_pyvers_name(self, url, cuda_version):
                 PckgVer(
                     "0.4.13", python_min_ver="3.8", python_max_ver="3.8", dependencies=["jaxlib"]
                 ),
-                # dax.fn.jax_function requires 0.4.16 which is the first one supporting
+                # dax.fn.jax_function requires at least 0.4.16 which is the first one supporting
                 # `__dlpack__` method, while 0.4.13 is the last one supported with Python3.8
                 PckgVer("0.4.16", python_min_ver="3.9", dependencies=["jaxlib"]),
             ]