From 63e57cf253bf74978c32d9f2bd93dd5b045a5a47 Mon Sep 17 00:00:00 2001 From: bzantium Date: Fri, 31 Jul 2026 23:59:46 +0900 Subject: [PATCH] Fix compile_cache_test flakiness by removing the persistent cache compile time threshold. Previously, test_train_step_cache_hit relied on the AOT compilation of train_step taking longer than jax_persistent_cache_min_compile_time_secs, which defaults to 1 second. JAX writes an entry to the persistent cache only when a compilation is slower than this threshold, so on a fast runner the AOT compilation is never cached, the runtime execution recompiles instead of hitting the cache, and the test fails even though the AOT and runtime cache keys agree. This fix sets the threshold to 0 in the subprocess environment so that the AOT compilation is cached regardless of how long it takes, leaving the test to verify only what it is meant to verify: that the AOT and runtime signatures produce the same cache key. --- tests/unit/compile_cache_test.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/compile_cache_test.py b/tests/unit/compile_cache_test.py index a1f35b98a1..d752aa6328 100644 --- a/tests/unit/compile_cache_test.py +++ b/tests/unit/compile_cache_test.py @@ -83,6 +83,11 @@ def test_train_step_cache_hit(): env["JAX_ENABLE_COMPILATION_CACHE"] = "true" env["JAX_COMPILATION_CACHE_DIR"] = temp_dir env["JAX_LOG_COMPILES"] = "1" + # JAX only caches a computation whose compilation is slower than this + # threshold, which defaults to 1 second. The model here is small enough that + # the AOT compilation can finish below it, so the cache is left empty and the + # runtime execution has nothing to hit. + env["JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS"] = "0" print("Running CPU training subprocess:", " ".join(cmd)) result = subprocess.run(cmd, env=env, capture_output=True, text=True, check=True)