Allow disabling compilation cache for particular runtime_types.

PiperOrigin-RevId: 640264856

Author: pschuh

jax/_src/compilation_cache.py

@@ -47,6 +47,7 @@
 # Mutex to protect _cache_initialized 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.
@@ -134,10 +135,13 @@ def _initialize_cache() -> None:
       logger.debug("Initialized persistent compilation cache at %s", path)
 
 
-def _get_cache() -> CacheInterface | None:
+def _get_cache(backend) -> CacheInterface | None:
   # TODO(b/289098047): consider making this an API and changing the callers of
   # get_executable_and_time() and put_executable_and_time() to call get_cache()
   # and passing the result to them.
+  if backend.runtime_type in _UNSUPPORTED_RUNTIMES:
+    logger.debug("_get_cache: Unsupported runtime: %s", backend.runtime_type)
+    return None
   if _cache is None:
     _initialize_cache()  # initialization is done at most once; see above
   return _cache
@@ -158,9 +162,9 @@ def decompress_executable(executable):
     return zlib.decompress(executable)
 
 
-def is_executable_in_cache(cache_key: str) -> bool:
+def is_executable_in_cache(backend, cache_key: str) -> bool:
   """Checks if the executable is in the cache."""
-  cache = _get_cache()
+  cache = _get_cache(backend)
   if cache is None:
     return False
 
@@ -175,7 +179,7 @@ def get_executable_and_time(
   """Returns the cached executable and its compilation time if present, or None
   otherwise.
   """
-  cache = _get_cache()
+  cache = _get_cache(backend)
   if cache is None:
     logger.debug("get_executable_and_time: cache is disabled/not initialized")
     return None, None
@@ -201,7 +205,7 @@ def put_executable_and_time(
   """Adds the 'executable' and its compilation time to the cache, possibly
   evicting older entries.
   """
-  cache = _get_cache()
+  cache = _get_cache(backend)
   if cache is None:
     logger.debug("put_executable_and_time: cache is disabled/not initialized")
     return

jax/_src/compiler.py

@@ -308,7 +308,7 @@ def compile_or_get_cached(
         computation, devices, compile_options, backend)
     compile_options.executable_build_options.fdo_profile = fdo_profile
 
-    if _is_executable_in_cache(pgle_profiled_module_key):
+    if _is_executable_in_cache(backend, pgle_profiled_module_key):
       # Load PGLE profiled module from the persistent cache.
       cache_key = pgle_profiled_module_key
       if pgle_profiler is not None:
@@ -614,11 +614,11 @@ def _compile_and_write_cache(
   )
   return executable
 
-def _is_executable_in_cache(cache_key) -> bool:
+def _is_executable_in_cache(backend, cache_key) -> bool:
   """Checks if executable is presented in cache on a given key
   """
   try:
-    return compilation_cache.is_executable_in_cache(cache_key)
+    return compilation_cache.is_executable_in_cache(backend, cache_key)
   except Exception as ex:
     if config.raise_persistent_cache_errors.value:
       raise

tests/compilation_cache_test.py

@@ -231,9 +231,10 @@ def f(x):
   def test_cache_write_warning(self):
     f = jit(lambda x: x * x)
 
+    backend = xla_bridge.get_backend()
     with (
       config.raise_persistent_cache_errors(False),
-      mock.patch.object(cc._get_cache().__class__, "put") as mock_put,
+      mock.patch.object(cc._get_cache(backend).__class__, "put") as mock_put,
       warnings.catch_warnings(record=True) as w,
     ):
       mock_put.side_effect = RuntimeError("test error")
@@ -252,9 +253,10 @@ def test_cache_write_warning(self):
   def test_cache_read_warning(self):
     f = jit(lambda x: x * x)
 
+    backend = xla_bridge.get_backend()
     with (
       config.raise_persistent_cache_errors(False),
-      mock.patch.object(cc._get_cache().__class__, "get") as mock_get,
+      mock.patch.object(cc._get_cache(backend).__class__, "get") as mock_get,
       warnings.catch_warnings(record=True) as w,
     ):
       mock_get.side_effect = RuntimeError("test error")