framereader: add cache_dir argument (#29904)

* framereader: add cache_dir argument

* make it an env var

tools/lib/cache.py

@@ -2,10 +2,10 @@ import os
 import urllib.parse
 from openpilot.common.file_helpers import mkdirs_exists_ok
 
-DEFAULT_CACHE_DIR = os.path.expanduser("~/.commacache")
+DEFAULT_CACHE_DIR = os.getenv("CACHE_ROOT", os.path.expanduser("~/.commacache"))
 
-def cache_path_for_file_path(fn, cache_prefix=None):
+def cache_path_for_file_path(fn, cache_dir=DEFAULT_CACHE_DIR):
-  dir_ = os.path.join(DEFAULT_CACHE_DIR, "local")
+  dir_ = os.path.join(cache_dir, "local")
   mkdirs_exists_ok(dir_)
   fn_parsed = urllib.parse.urlparse(fn)
   if fn_parsed.scheme == '':

tools/lib/framereader.py

@@ -12,7 +12,7 @@ import numpy as np
 from lru import LRU
 
 import _io
-from openpilot.tools.lib.cache import cache_path_for_file_path
+from openpilot.tools.lib.cache import cache_path_for_file_path, DEFAULT_CACHE_DIR
 from openpilot.tools.lib.exceptions import DataUnreadableError
 from openpilot.common.file_helpers import atomic_write_in_dir
 
@@ -106,8 +106,8 @@ def cache_fn(func):
     if kwargs.pop('no_cache', None):
       cache_path = None
     else:
-      cache_prefix = kwargs.pop('cache_prefix', None)
+      cache_dir = kwargs.pop('cache_dir', DEFAULT_CACHE_DIR)
-      cache_path = cache_path_for_file_path(fn, cache_prefix)
+      cache_path = cache_path_for_file_path(fn, cache_dir)
 
     if cache_path and os.path.exists(cache_path):
       with open(cache_path, "rb") as cache_file:
@@ -140,18 +140,18 @@ def index_stream(fn, typ):
   }
 
 
-def index_videos(camera_paths, cache_prefix=None):
+def index_videos(camera_paths, cache_dir=DEFAULT_CACHE_DIR):
   """Requires that paths in camera_paths are contiguous and of the same type."""
   if len(camera_paths) < 1:
     raise ValueError("must provide at least one video to index")
 
   frame_type = fingerprint_video(camera_paths[0])
   for fn in camera_paths:
-    index_video(fn, frame_type, cache_prefix)
+    index_video(fn, frame_type, cache_dir)
 
 
-def index_video(fn, frame_type=None, cache_prefix=None):
+def index_video(fn, frame_type=None, cache_dir=DEFAULT_CACHE_DIR):
-  cache_path = cache_path_for_file_path(fn, cache_prefix)
+  cache_path = cache_path_for_file_path(fn, cache_dir)
 
   if os.path.exists(cache_path):
     return
@@ -160,16 +160,16 @@ def index_video(fn, frame_type=None, cache_prefix=None):
     frame_type = fingerprint_video(fn[0])
 
   if frame_type == FrameType.h265_stream:
-    index_stream(fn, "hevc", cache_prefix=cache_prefix)
+    index_stream(fn, "hevc", cache_dir=cache_dir)
   else:
     raise NotImplementedError("Only h265 supported")
 
 
-def get_video_index(fn, frame_type, cache_prefix=None):
+def get_video_index(fn, frame_type, cache_dir=DEFAULT_CACHE_DIR):
-  cache_path = cache_path_for_file_path(fn, cache_prefix)
+  cache_path = cache_path_for_file_path(fn, cache_dir)
 
   if not os.path.exists(cache_path):
-    index_video(fn, frame_type, cache_prefix)
+    index_video(fn, frame_type, cache_dir)
 
   if not os.path.exists(cache_path):
     return None
@@ -284,13 +284,13 @@ class BaseFrameReader:
     raise NotImplementedError
 
 
-def FrameReader(fn, cache_prefix=None, readahead=False, readbehind=False, index_data=None):
+def FrameReader(fn, cache_dir=DEFAULT_CACHE_DIR, readahead=False, readbehind=False, index_data=None):
   frame_type = fingerprint_video(fn)
   if frame_type == FrameType.raw:
     return RawFrameReader(fn)
   elif frame_type in (FrameType.h265_stream,):
     if not index_data:
-      index_data = get_video_index(fn, frame_type, cache_prefix)
+      index_data = get_video_index(fn, frame_type, cache_dir)
     return StreamFrameReader(fn, frame_type, index_data, readahead=readahead, readbehind=readbehind)
   else:
     raise NotImplementedError(frame_type)