diff --git a/tools/lib/cache.py b/tools/lib/cache.py index d142955e59..fd214f6bb5 100644 --- a/tools/lib/cache.py +++ b/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): - dir_ = os.path.join(DEFAULT_CACHE_DIR, "local") +def cache_path_for_file_path(fn, cache_dir=DEFAULT_CACHE_DIR): + dir_ = os.path.join(cache_dir, "local") mkdirs_exists_ok(dir_) fn_parsed = urllib.parse.urlparse(fn) if fn_parsed.scheme == '': diff --git a/tools/lib/framereader.py b/tools/lib/framereader.py index cbc9310790..06fee6857d 100644 --- a/tools/lib/framereader.py +++ b/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_path = cache_path_for_file_path(fn, cache_prefix) + cache_dir = kwargs.pop('cache_dir', DEFAULT_CACHE_DIR) + 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): - cache_path = cache_path_for_file_path(fn, cache_prefix) +def index_video(fn, frame_type=None, cache_dir=DEFAULT_CACHE_DIR): + 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): - cache_path = cache_path_for_file_path(fn, cache_prefix) +def get_video_index(fn, frame_type, cache_dir=DEFAULT_CACHE_DIR): + 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)