Merge pull request #31328 from olupton:test-dump-on-deserialize

PiperOrigin-RevId: 840268767

Author: Google-ML-Automation

tests/compilation_cache_test.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from collections import Counter
+import glob
 import logging
 import math
 import os
@@ -648,6 +649,35 @@ def test_persistent_cache_enable_xla_caches(self):
         self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_per_fusion_autotune_cache_dir, f"jax-cache{s}xla_gpu_per_fusion_autotune_cache_dir")
         self.assertEqual(compile_options.executable_build_options.debug_options.xla_gpu_experimental_autotune_cache_mode, xc.AutotuneCacheMode.UPDATE)
 
+  @jtu.skip_on_devices("tpu") # TPU backend does not dump on deserialize
+  def test_dump_on_cache_hit(self):
+    previous_counts = Counter(_counts)
+    with (
+      config.persistent_cache_min_compile_time_secs(0),
+      config.persistent_cache_min_entry_size_bytes(0),
+      tempfile.TemporaryDirectory() as dump_dir1,
+      tempfile.TemporaryDirectory() as dump_dir2
+    ):
+      jit(lambda x: x + 1, compiler_options={"xla_dump_to": dump_dir1})(1)
+      self.assertEqual(
+        _counts["/jax/compilation_cache/cache_hits"],
+        previous_counts["/jax/compilation_cache/cache_hits"],
+      )
+      jit(lambda x: x + 1, compiler_options={"xla_dump_to": dump_dir2, "xla_dump_hlo_as_proto": True, "xla_dump_hlo_as_text": True})(1)
+      self.assertEqual(
+          _counts["/jax/compilation_cache/cache_hits"],
+          previous_counts["/jax/compilation_cache/cache_hits"] + 1,
+      1)
+      dump1_files = glob.glob(os.path.join(dump_dir1, "*after_optimizations.txt"))
+      dump2_files = glob.glob(os.path.join(dump_dir2, "*after_optimizations.txt"))
+      self.assertEqual(len(dump1_files), 1)
+      self.assertEqual(len(dump2_files), 1)
+      with (open(dump1_files[0]) as file1, open(dump2_files[0]) as file2):
+        self.assertEqual(file1.read(), file2.read())
+      dump2_pbs = glob.glob(os.path.join(dump_dir2, "*after_optimizations.hlo.pb"))
+      self.assertEqual(len(dump2_pbs), 1)
+
+
 @jtu.with_config(
     jax_enable_compilation_cache=False,
     jax_persistent_cache_min_compile_time_secs=0,