model replay: framereader cache (#35576)

* Simpler cache version

* cachetools

* different LRU

* lint

* smaller

* just write LRU

* mypy

* same length

pyproject.toml

@@ -106,7 +106,6 @@ dev = [
   "azure-storage-blob",
   "dbus-next",
   "dictdiffer",
-  "lru-dict",
   "matplotlib",
   "parameterized >=0.8, <0.9",
   "pyautogui",

selfdrive/test/process_replay/model_replay.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import os
+import pickle
 import sys
 from collections import defaultdict
 from typing import Any
@@ -194,17 +195,39 @@ def model_replay(lr, frs):
   return msgs
 
 
+def get_frames():
+  regen_cache = "--regen-cache" in sys.argv
+  frames_cache = '/tmp/model_replay_cache' if PC else '/data/model_replay_cache'
+  os.makedirs(frames_cache, exist_ok=True)
+
+  cache_name = f'{frames_cache}/{TEST_ROUTE}_{SEGMENT}_{START_FRAME}_{END_FRAME}.pkl'
+  if os.path.isfile(cache_name) and not regen_cache:
+    try:
+      print(f"Loading frames from cache {cache_name}")
+      return pickle.load(open(cache_name, "rb"))
+    except Exception as e:
+      print(f"Failed to load frames from cache {cache_name}: {e}")
+
+  frs = {
+    'roadCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "fcamera.hevc"), pix_fmt='nv12', cache_size=END_FRAME - START_FRAME),
+    'driverCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "dcamera.hevc"), pix_fmt='nv12', cache_size=END_FRAME - START_FRAME),
+    'wideRoadCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "ecamera.hevc"), pix_fmt='nv12', cache_size=END_FRAME - START_FRAME),
+  }
+  for fr in frs.values():
+    for fidx in range(START_FRAME, END_FRAME):
+      fr.get(fidx)
+    fr.it = None
+  print(f"Dumping frame cache {cache_name}")
+  pickle.dump(frs, open(cache_name, "wb"))
+  return frs
+
 if __name__ == "__main__":
   update = "--update" in sys.argv or (os.getenv("GIT_BRANCH", "") == 'master')
   replay_dir = os.path.dirname(os.path.abspath(__file__))
 
   # load logs
   lr = list(LogReader(get_url(TEST_ROUTE, SEGMENT, "rlog.zst")))
-  frs = {
-    'roadCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "fcamera.hevc"), pix_fmt='nv12'),
-    'driverCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "dcamera.hevc"), pix_fmt='nv12'),
-    'wideRoadCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "ecamera.hevc"), pix_fmt='nv12')
-  }
+  frs = get_frames()
 
   log_msgs = []
   # run replays

tools/lib/framereader.py

@@ -2,10 +2,9 @@ import os
 import subprocess
 import json
 from collections.abc import Iterator
+from collections import OrderedDict
 
 import numpy as np
-from lru import LRU
-
 from openpilot.tools.lib.filereader import FileReader, resolve_name
 from openpilot.tools.lib.exceptions import DataUnreadableError
 from openpilot.tools.lib.vidindex import hevc_index
@@ -16,6 +15,24 @@ HEVC_SLICE_P = 1
 HEVC_SLICE_I = 2
 
 
+class LRUCache:
+    def __init__(self, capacity: int):
+      self._cache: OrderedDict = OrderedDict()
+      self.capacity = capacity
+
+    def __getitem__(self, key):
+      self._cache.move_to_end(key)
+      return self._cache[key]
+
+    def __setitem__(self, key, value):
+      self._cache[key] = value
+      if len(self._cache) > self.capacity:
+          self._cache.popitem(last=False)
+
+    def __contains__(self, key):
+      return key in self._cache
+
+
 def assert_hvec(fn: str) -> None:
   with FileReader(fn) as f:
     header = f.read(4)
@@ -135,7 +152,7 @@ class FrameReader:
                cache_size: int = 30, pix_fmt: str = "rgb24"):
     self.decoder = FfmpegDecoder(fn, index_data, pix_fmt)
     self.iframes = self.decoder.iframes
-    self._cache: LRU[int, np.ndarray] = LRU(cache_size)
+    self._cache: LRUCache = LRUCache(cache_size)
     self.w, self.h, self.frame_count, = self.decoder.w, self.decoder.h, self.decoder.frame_count
     self.pix_fmt = pix_fmt