Update cache.py

cache.py

@@ -1,7 +1,7 @@
 # Copyright 2025 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 not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
@@ -12,56 +12,60 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from diskcache import Cache
 import os
 import shutil
 import tempfile
 import zipfile
 import logging
+from diskcache import Cache
 
-cache = Cache(os.environ.get("CACHE_DIR", "/cache"))
-# Print cache statistics after loading
+# Resolve cache directory from environment or default to HF_HOME or /tmp
+cache_directory = os.environ.get("CACHE_DIR", os.getenv("HF_HOME", "/tmp"))
+cache = Cache(cache_directory)
+
+# Print cache statistics
 try:
     item_count = len(cache)
     size_bytes = cache.volume()
-    print(f"Cache loaded: {item_count} items, approx {size_bytes} bytes")
+    print(f"[Cache] Loaded: {item_count} items, ~{size_bytes} bytes")
 except Exception as e:
-    print(f"Could not retrieve cache statistics: {e}")
+    print(f"[Cache] Unable to retrieve statistics: {e}")
 
 def create_cache_zip():
+    """
+    Archives the cache directory into a ZIP file, ensuring it is safe and consistent.
+    Returns the archive path and an error message (if any).
+    """
     temp_dir = tempfile.gettempdir()
-    base_name = os.path.join(temp_dir, "cache_archive") # A more descriptive name
-    archive_path = base_name + ".zip"
-    cache_directory = os.environ.get("CACHE_DIR", "/cache")
-    
+    archive_path = os.path.join(temp_dir, "cache_archive.zip")
+
     if not os.path.isdir(cache_directory):
-        logging.error(f"Cache directory not found at {cache_directory}")
-        return None, f"Cache directory not found on server: {cache_directory}"
-    
-    logging.info("Forcing a cache checkpoint for safe backup...")
+        logging.error(f"[Cache] Directory does not exist: {cache_directory}")
+        return None, f"Cache directory not found: {cache_directory}"
+
+    logging.info("[Cache] Initiating checkpoint before archiving...")
     try:
-        # Open and immediately close a connection.
-        # This forces SQLite to perform a checkpoint, merging the .wal file
-        # into the main .db file, ensuring the on-disk files are consistent.
+        # Force SQLite checkpoint to flush WAL
         with Cache(cache_directory) as temp_cache:
             temp_cache.close()
-        
-        # Clean up temporary files before archiving.
-        tmp_path = os.path.join(cache_directory, 'tmp')
-        if os.path.isdir(tmp_path):
-            logging.info(f"Removing temporary cache directory: {tmp_path}")
-            shutil.rmtree(tmp_path)
 
-        logging.info(f"Checkpoint complete. Creating zip archive of {cache_directory} to {archive_path}")
+        # Remove unnecessary temp cache files
+        tmp_subdir = os.path.join(cache_directory, 'tmp')
+        if os.path.isdir(tmp_subdir):
+            logging.info(f"[Cache] Removing temporary subdirectory: {tmp_subdir}")
+            shutil.rmtree(tmp_subdir)
+
+        logging.info(f"[Cache] Creating archive at: {archive_path}")
         with zipfile.ZipFile(archive_path, 'w', zipfile.ZIP_DEFLATED, compresslevel=9) as zipf:
             for root, _, files in os.walk(cache_directory):
                 for file in files:
                     file_path = os.path.join(root, file)
                     arcname = os.path.relpath(file_path, cache_directory)
                     zipf.write(file_path, arcname)
-        logging.info("Zip archive created successfully.")
+
+        logging.info("[Cache] Archive created successfully.")
         return archive_path, None
-        
+
     except Exception as e:
-        logging.error(f"Error creating zip archive of cache directory: {e}", exc_info=True)
-        return None, f"Error creating zip archive: {e}"
+        logging.exception("[Cache] Failed to create archive")
+        return None, f"Error creating cache archive: {e}"