Fix compile_cache_test to assert single jit_train_step cache file

[Liauuu]

Description

Fixes a failure in the nightly cpu-unit job (Build Nightly Docker Images, commit 85690dab).

compile_cache_test.py assumed the JAX persistent compilation cache directory would contain exactly one file total. In nightly CI, two entries are written:

• jit_train_step-...
• jit_initialize_state-...

initialize_state is a separate JIT used during model state initialization and is unrelated to the train_step double-compilation regression this test is meant to catch. Requiring a total cache file count of 1 therefore fails even when train_step caching behaves correctly.

This change updates the test to:

• Assert exactly one cache file whose name starts with jit_train_step
• Keep the existing log assertion that runtime execution hits the persistent cache for jit_train_step

Other cache entries (e.g. jit_initialize_state) are allowed and no longer cause a false failure.

Context

Related to #4000 (partial — MoE and GPU integration failures in the same nightly run are separate issues).

Error log (before fix)

FAILED tests/unit/compile_cache_test.py::test_train_step_cache_hit - AssertionError: Expected exactly 1 JAX compilation cache file, but found 2: ['jit_train_step-425341b98f4b4ea0715158828c3b68ee2deab5f2df3c74ad475a08c09652d19f-cache', 'jit_initialize_state-68b1de74b0881dd7ff1205f613bc9988724719034daf764e4c2f6893b61098ca-cache']. This indicates a cache miss where AOT compilation and runtime execution generated different keys, causing train_step to be compiled twice (double-compilation regression). assert 2 == 1

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.