Fix compile_cache_test flakiness caused by the persistent cache compile time threshold

[bzantium]

Description

test_train_step_cache_hit relies 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 showed up on a CI run where the two compilations of jit_train_step logged identical shapes and argument mappings, but the AOT one finished in 0.941s and the runtime one in 1.790s. Only the second cleared the threshold, so the single file in the cache directory was written by the runtime compilation rather than read by it, and it was the only one of the 20 compilations in that subprocess above 1 second.

This fix sets the threshold to 0 in the subprocess environment so that the AOT compilation is cached however fast it is, leaving the test to verify only what it is meant to verify: that the AOT and runtime signatures produce the same cache key.

Tests

Checked against jax 0.10.0 that JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0 reaches jax.config.jax_persistent_cache_min_compile_time_secs, and that a sub-second compilation writes a cache entry with the threshold at 0 while writing none at the default. Lowering the threshold also caches the other functions the subprocess compiles, but only one entry carries the jit_train_step prefix the test filters on, so the existing assertions still hold.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!