Repository: google/jax
Pull request #22580 (2024-07-22 22:05:09)

Add a beacon metric to report tasks disabled Jax compilation cache.
Add a beacon metric to report tasks disabled Jax compilation cache.

- Create metric '/jax/compilation_cache/task_disabled_cache' as a beacon metric to monitor tasks which have disabled compilation cache.
- Modified existing logic for reporting the '/jax/compilation_cache/tasks_using_cache' metric and make it easier to find the two adoption related metrics in the code.



Patch:
diff --git a/jax/_src/compilation_cache.py b/jax/_src/compilation_cache.py
index 1a21eb143..7fcbe860f 100644
--- a/jax/_src/compilation_cache.py
+++ b/jax/_src/compilation_cache.py
@@ -42,26 +42,47 @@ _cache: CacheInterface | None = None
 
 _cache_initialized: bool = False
 
+_cache_checked: bool = False
+
 _cache_used: bool = False
 
-# Mutex to protect _cache_initialized and _cache_used.
+# Mutex to protect _cache_initialized, _cache_checked and _cache_used.
 _cache_initialized_mutex = threading.Lock()
 
 _UNSUPPORTED_RUNTIMES: set[str] = set()
 
-def set_once_cache_used(f) -> None:
-  """One-time setting of _cache_used.
-
-  If _cache_used is False, set it to True and execute the provided function
-  f. No action if _cache_used is True. This provides a mechanism to execute f
-  once per task. Note that reset_cache() will reset _cache_used also.
+def is_cache_used(backend: xla_client.Client) -> bool:
+  """Check if cache is used and report adoption metrics one-time per task.
+  The cache may be initialized during the first call to this function.
   """
-  global _cache_used
+  # Return _cache_used directly if _cache_checked is True. If _cache_checked is
+  # False, set it to True, report metrics and return if cache is used. This
+  # provides a mechanism to report the metrics once per task. Note that
+  # reset_cache() will reset _cache_checked and _cache_used also.
+  global _cache_checked, _cache_used
+  with _cache_initialized_mutex:
+    if _cache_checked:
+      return _cache_used
+
   with _cache_initialized_mutex:
-    if not _cache_used:
-      _cache_used = True
-      if f is not None:
-        f()
+    if not _cache_checked:
+      _cache_checked = True
+
+      # Persistent compilation cache only implemented on TPU and GPU and the
+      # backend that supports serialization of executables.
+      # TODO(skye): add warning when initializing cache on unsupported default
+      # platform
+      supported_platforms = ["tpu", "gpu", "cpu"]
+
+      if not _is_cache_enabled():
+        monitoring.record_event('/jax/compilation_cache/task_disabled_cache')
+      elif (
+          backend.platform in supported_platforms
+          and getattr(backend, "supports_executable_serialization", True)
+      ):
+        monitoring.record_event('/jax/compilation_cache/tasks_using_cache')
+        _cache_used = True
+      return _cache_used
 
 
 def get_file_cache(path: str) -> tuple[CacheInterface, str] | None:
@@ -267,12 +288,14 @@ def reset_cache() -> None:
   """Get back to pristine, uninitialized state."""
   global _cache
   global _cache_initialized
+  global _cache_checked
   global _cache_used
   logger.info("Resetting cache at %s.",
                _cache._path if _cache is not None else "<empty>")
   _cache = None
   with _cache_initialized_mutex:
     _cache_initialized = False
+    _cache_checked = False
     _cache_used = False
 
 
diff --git a/jax/_src/compiler.py b/jax/_src/compiler.py
index f4a3a3190..d778abaf3 100644
--- a/jax/_src/compiler.py
+++ b/jax/_src/compiler.py
@@ -268,23 +268,12 @@ def compile_or_get_cached(
   if dumped_to := mlir.dump_module_to_file(computation, "compile"):
     logging.info("Dumped the module to %s.", dumped_to)
 
-  # Persistent compilation cache only implemented on TPU and GPU and the backend
-  # that supports serialization of executables.
-  # TODO(skye): add warning when initializing cache on unsupported default platform
-  supported_platforms = ["tpu", "gpu", "cpu"]
-  use_compilation_cache = (
-      config.enable_compilation_cache.value
-      and getattr(backend, "supports_executable_serialization", True)
-      and backend.platform in supported_platforms
-  )
+  use_compilation_cache = compilation_cache.is_cache_used(backend)
 
   if not use_compilation_cache:
     return backend_compile(backend, computation, compile_options,
                            host_callbacks)
 
-  compilation_cache.set_once_cache_used(
-      lambda: monitoring.record_event(
-          "/jax/compilation_cache/tasks_using_cache"))
   monitoring.record_event('/jax/compilation_cache/compile_requests_use_cache')
 
   try:


Test patch:
diff --git a/tests/compilation_cache_test.py b/tests/compilation_cache_test.py
index caf354cd5..441a93236 100644
--- a/tests/compilation_cache_test.py
+++ b/tests/compilation_cache_test.py
@@ -574,6 +574,20 @@ class CompilationCacheDisabledTest(CompilationCacheTestCase):
       f(1)
       self.assertEqual(count_cache_items(), 0)
 
+  def test_tasks_disable_cache_metric(self):
+    with config.enable_compilation_cache(False):
+      count_before_first_use = _counts[
+          "/jax/compilation_cache/task_disabled_cache"]
+      jit(lambda x: x + 1)(1)
+      count_after_first_use = _counts[
+          "/jax/compilation_cache/task_disabled_cache"]
+      self.assertEqual(count_after_first_use - count_before_first_use, 1)
+
+      # Verify that the count is incremented only once per task.
+      jit(lambda x: x + 3)(3)
+      count_after_second_use = _counts[
+          "/jax/compilation_cache/task_disabled_cache"]
+      self.assertEqual(count_after_second_use, count_after_first_use)
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())