Merge pull request #4222 from AI-Hypercomputer:igorts/b517327969-nightly-build-failure

Google-ML-Automation · Google-ML-Automation · commit fd5bf7cddf4a · 2026-06-23T10:05:50.000-07:00
PiperOrigin-RevId: 936737860
diff --git a/tests/unit/compile_cache_test.py b/tests/unit/compile_cache_test.py
@@ -106,8 +106,11 @@ def test_train_step_cache_hit():
         "cache was not writeable or the JAX cache configuration was ignored."
     )
 
-    assert len(cache_files) == 1, (
-        f"Expected exactly 1 JAX compilation cache file, but found {len(cache_files)}: {cache_files}. "
+    train_step_cache_files = [f for f in cache_files if f.startswith("jit_train_step")]
+    assert len(train_step_cache_files) == 1, (
+        f"Expected exactly 1 JAX compilation cache file for 'jit_train_step', "
+        f"but found {len(train_step_cache_files)}: {train_step_cache_files} "
+        f"(all cache files: {cache_files}). "
         "This indicates a cache miss where AOT compilation and runtime execution generated different keys, "
         "causing train_step to be compiled twice (double-compilation regression)."
     )

Original file line number	Diff line number	Diff line change
`@@ -106,8 +106,11 @@ def test_train_step_cache_hit():`
`106`	`106`	`"cache was not writeable or the JAX cache configuration was ignored."`
`107`	`107`	`)`
`108`	`108`
`109`		`- assert len(cache_files) == 1, (`
`110`		`- f"Expected exactly 1 JAX compilation cache file, but found {len(cache_files)}: {cache_files}. "`
	`109`	`+ train_step_cache_files = [f for f in cache_files if f.startswith("jit_train_step")]`
	`110`	`+ assert len(train_step_cache_files) == 1, (`
	`111`	`+ f"Expected exactly 1 JAX compilation cache file for 'jit_train_step', "`
	`112`	`+ f"but found {len(train_step_cache_files)}: {train_step_cache_files} "`
	`113`	`+ f"(all cache files: {cache_files}). "`
`111`	`114`	`"This indicates a cache miss where AOT compilation and runtime execution generated different keys, "`
`112`	`115`	`"causing train_step to be compiled twice (double-compilation regression)."`
`113`	`116`	`)`