From c177811670f211ca4cedc3874ffaca7d6cb9a80d Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Thu, 7 Nov 2024 10:08:54 -0800 Subject: [PATCH] Disable JAX's compilation cache since Pathways manages it's own compilation cache. PiperOrigin-RevId: 694159092 --- pathwaysutils/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pathwaysutils/__init__.py b/pathwaysutils/__init__.py index 3af7ad9..84e6e19 100644 --- 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: