Added file system cache interface

PiperOrigin-RevId: 388473011

jax/experimental/compilation_cache/cache_interface.py

@@ -0,0 +1,24 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+
+class CacheInterface(ABC):
+    @abstractmethod
+    def get(self, key: str):
+        pass
+
+    @abstractmethod
+    def put(self, key: str, value: bytes):
+        pass

jax/experimental/compilation_cache/compilation_cache.py

@@ -27,7 +27,7 @@ def initialize_cache(path, max_cache_size_bytes=32 * 2**30):
        max_cache_sixe defaults to 32GiB.
     """
     global _cache
-    assert _cache == None, f"The cache path has already been initialized to {_cache}"
+    assert _cache == None, f"The cache path has already been initialized to {_cache._path}"
     _cache = FileSystemCache(path, max_cache_size_bytes)
     logging.warning(f"Initialized persistent compilation cache at {path}")
 

jax/experimental/compilation_cache/file_system_cache.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 
 import os
+from jax.experimental.compilation_cache.cache_interface import CacheInterface
 import tempfile
 from typing import Optional
 import warnings
 
-class FileSystemCache:
+class FileSystemCache(CacheInterface):
 
   def __init__(self, path: str, max_cache_size_bytes=32 * 2**30):
     """Sets up a cache at 'path'. Cached values may already be present."""
@@ -63,7 +64,6 @@ class FileSystemCache:
     if new_file_size >= self._max_cache_size_bytes:
       return False
 
-    #TODO(colemanliyah): Refactor this section so the whole directory doesn't need to be checked
     while new_file_size + self._get_cache_directory_size() > self._max_cache_size_bytes:
       last_time = float('inf')
       file_to_delete = None

jax/interpreters/xla.py

@@ -70,6 +70,7 @@ def compile_or_get_cached(backend, computation, compile_options):
     if cc.is_initialized():
         cached_executable = cc.get_executable(computation, compile_options)
         if cached_executable is not None:
+            logging.info('Persistent compilation cache hit')
             return cached_executable
         else:
             compiled = backend_compile(backend, computation, compile_options)