ci: faster model_replay (#34036)

* cache draft

* fix

* fix

* fix

* better

* zst

* more

* try

* pool

* fix

* fix

* revert :C

* better

* cleanup

* no cache

* this too

Jenkinsfile

@@ -243,7 +243,7 @@ node {
       'replay': {
         deviceStage("model-replay", "tici-replay", ["UNSAFE=1"], [
           step("build", "cd system/manager && ./build.py", [diffPaths: ["selfdrive/modeld/", "tinygrad_repo", "selfdrive/test/process_replay/model_replay.py"]]),
-          step("model replay", "selfdrive/test/process_replay/model_replay.py", [diffPaths: ["selfdrive/modeld/"]]),
+          step("model replay", "selfdrive/test/process_replay/model_replay.py", [diffPaths: ["selfdrive/modeld/", "tinygrad_repo", "selfdrive/test/process_replay/model_replay.py"]]),
         ])
       },
       'tizi': {

selfdrive/test/process_replay/model_replay.py

@@ -7,19 +7,20 @@ import tempfile
 from itertools import zip_longest
 
 import matplotlib.pyplot as plt
+import numpy as np
 
 from openpilot.common.git import get_commit
 from openpilot.system.hardware import PC
 from openpilot.tools.lib.openpilotci import get_url
 from openpilot.selfdrive.test.process_replay.compare_logs import compare_logs, format_diff
 from openpilot.selfdrive.test.process_replay.process_replay import get_process_config, replay_process
-from openpilot.tools.lib.framereader import FrameReader
+from openpilot.tools.lib.framereader import FrameReader, NumpyFrameReader
 from openpilot.tools.lib.logreader import LogReader, save_log
 from openpilot.tools.lib.github_utils import GithubUtils
 
 TEST_ROUTE = "2f4452b03ccb98f0|2022-12-03--13-45-30"
 SEGMENT = 6
-MAX_FRAMES = 100 if PC else 600
+MAX_FRAMES = 100 if PC else 400
 
 NO_MODEL = "NO_MODEL" in os.environ
 SEND_EXTRA_INPUTS = bool(int(os.getenv("SEND_EXTRA_INPUTS", "0")))
@@ -31,7 +32,7 @@ GITHUB = GithubUtils(API_TOKEN, DATA_TOKEN)
 
 
 def get_log_fn(test_route, ref="master"):
-  return f"{test_route}_model_tici_{ref}.bz2"
+  return f"{test_route}_model_tici_{ref}.zst"
 
 def plot(proposed, master, title, tmp):
   proposed = list(proposed)
@@ -151,21 +152,44 @@ def model_replay(lr, frs):
   dmonitoringmodeld = get_process_config("dmonitoringmodeld")
 
   modeld_msgs = replay_process(modeld, modeld_logs, frs)
+  if isinstance(frs['roadCameraState'], NumpyFrameReader):
+    del frs['roadCameraState'].frames
+    del frs['wideRoadCameraState'].frames
   dmonitoringmodeld_msgs = replay_process(dmonitoringmodeld, dmodeld_logs, frs)
   return modeld_msgs + dmonitoringmodeld_msgs
 
 
+def get_frames():
+  regen_cache = "--regen-cache" in sys.argv
+  cache = "--cache" in sys.argv or not PC or regen_cache
+  videos = ('fcamera.hevc', 'dcamera.hevc', 'ecamera.hevc')
+  cams = ('roadCameraState', 'driverCameraState', 'wideRoadCameraState')
+
+  if cache:
+    frames_cache = '/tmp/model_replay_cache' if PC else '/data/model_replay_cache'
+    os.makedirs(frames_cache, exist_ok=True)
+
+    cache_size = 200
+    for v in videos:
+      if not all(os.path.isfile(f'{frames_cache}/{TEST_ROUTE}_{v}_{i}.npy') for i in range(MAX_FRAMES//cache_size)) or regen_cache:
+        f = FrameReader(get_url(TEST_ROUTE, SEGMENT, v)).get(0, MAX_FRAMES + 1, pix_fmt="nv12")
+        print(f'Caching {v}...')
+        for i in range(MAX_FRAMES//cache_size):
+          np.save(f'{frames_cache}/{TEST_ROUTE}_{v}_{i}', f[(i * cache_size) + 1:((i + 1) * cache_size) + 1])
+        del f
+
+    return {c : NumpyFrameReader(f"{frames_cache}/{TEST_ROUTE}_{v}", 1928, 1208, cache_size) for c,v in zip(cams, videos, strict=True)}
+  else:
+    return {c : FrameReader(get_url(TEST_ROUTE, SEGMENT, v), readahead=True) for c,v in zip(cams, videos, strict=True)}
+
+
 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.bz2")))
-  frs = {
-    'roadCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "fcamera.hevc"), readahead=True),
-    'driverCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "dcamera.hevc"), readahead=True),
-    'wideRoadCameraState': FrameReader(get_url(TEST_ROUTE, SEGMENT, "ecamera.hevc"), readahead=True)
-  }
+  lr = list(LogReader(get_url(TEST_ROUTE, SEGMENT, "rlog.zst")))
+  frs = get_frames()
 
   log_msgs = []
   # run replays

tools/lib/framereader.py

@@ -535,3 +535,25 @@ def FrameIterator(fn, pix_fmt, **kwargs):
   else:
     for i in range(fr.frame_count):
       yield fr.get(i, pix_fmt=pix_fmt)[0]
+
+
+class NumpyFrameReader:
+  def __init__(self, name, w, h, cache_size):
+    self.name = name
+    self.pos = -1
+    self.frames = None
+    self.w = w
+    self.h = h
+    self.cache_size = cache_size
+
+  def close(self):
+    pass
+
+  def get(self, num, count=1, pix_fmt="nv12"):
+    num -= 1
+    q = num // self.cache_size
+    if q != self.pos:
+      del self.frames
+      self.pos = q
+      self.frames = np.load(f'{self.name}_{self.pos}.npy')
+    return [self.frames[num % self.cache_size]]