Disable JAX's compilation cache since Pathways manages it's own compi…

…lation cache. PiperOrigin-RevId: 694159092
AI-Hypercomputer · Nov 12, 2024 · 80e27a8 · 80e27a8
1 parent ea39d05
commit 80e27a8
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/pathwaysutils/__init__.py b/pathwaysutils/__init__.py
@@ -60,6 +60,11 @@ def _is_persistence_enabled():
     pathways_orbax_handler.register_pathways_handlers(
         datetime.timedelta(minutes=10)
     )
+
+  # Turn off JAX compilation cache because Pathways handles its own compilation
+  # cache.
+  jax.config.update("jax_enable_compilation_cache", False)
+
   try:
     cloud_logging.setup()
   except OSError as e: