process_replay: helpers (#28367)

* process_replay helper

* Refactor

* refactor check_enabled uses

* fix __init__

* noqa in __init__ imports

* Move compare_logs outside of process_replay

* Move save_logs to tools.lib.helpers

* Remove save_log from compare_logs

* fix lint issues

selfdrive/debug/run_process_on_route.py

@@ -2,10 +2,10 @@
 
 import argparse
 
-from selfdrive.test.process_replay.compare_logs import save_log
 from selfdrive.test.process_replay.process_replay import CONFIGS, replay_process
 from tools.lib.logreader import MultiLogIterator
 from tools.lib.route import Route
+from tools.lib.helpers import save_log
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser(description="Run process on route and create new logs",

selfdrive/test/process_replay/__init__.py

@@ -0,0 +1 @@
+from selfdrive.test.process_replay.process_replay import CONFIGS, get_process_config, replay_process, replay_process_with_name  # noqa: F401

selfdrive/test/process_replay/compare_logs.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python3
-import bz2
 import sys
 import math
 import capnp
@@ -12,16 +11,6 @@ from tools.lib.logreader import LogReader
 EPSILON = sys.float_info.epsilon
 
 
-def save_log(dest, log_msgs, compress=True):
-  dat = b"".join(msg.as_builder().to_bytes() for msg in log_msgs)
-
-  if compress:
-    dat = bz2.compress(dat)
-
-  with open(dest, "wb") as f:
-    f.write(dat)
-
-
 def remove_ignored_fields(msg, ignore):
   msg = msg.as_builder()
   for key in ignore:

selfdrive/test/process_replay/model_replay.py

@@ -14,11 +14,12 @@ from common.transformations.camera import tici_f_frame_size, tici_d_frame_size
 from system.hardware import PC
 from selfdrive.manager.process_config import managed_processes
 from selfdrive.test.openpilotci import BASE_URL, get_url
-from selfdrive.test.process_replay.compare_logs import compare_logs, save_log
+from selfdrive.test.process_replay.compare_logs import compare_logs
 from selfdrive.test.process_replay.test_processes import format_diff
 from system.version import get_commit
 from tools.lib.framereader import FrameReader
 from tools.lib.logreader import LogReader
+from tools.lib.helpers import save_log
 
 TEST_ROUTE = "2f4452b03ccb98f0|2022-12-03--13-45-30"
 SEGMENT = 6

selfdrive/test/process_replay/process_replay.py

@@ -333,7 +333,28 @@ def get_process_config(name):
     raise Exception(f"Cannot find process config with name: {name}") from ex
 
 
-def replay_process(cfg, lr, fingerprint=None):
+def replay_process_with_name(name, lr, *args, **kwargs):
+  cfg = get_process_config(name)
+  return replay_process(cfg, lr, *args, **kwargs)
+
+
+def replay_process(cfg, lr, fingerprint=None, return_all_logs=False):
+  all_msgs = list(lr)
+  process_logs = _replay_single_process(cfg, all_msgs, fingerprint)
+
+  if return_all_logs:
+    keys = set(cfg.subs)
+    modified_logs = [m for m in all_msgs if m.which() not in keys]
+    modified_logs.extend(process_logs)
+    modified_logs.sort(key=lambda m: m.logMonoTime)
+    log_msgs = modified_logs
+  else:
+    log_msgs = process_logs
+
+  return log_msgs
+
+
+def _replay_single_process(cfg, lr, fingerprint):
   with OpenpilotPrefix():
     controlsState = None
     initialized = False
@@ -484,7 +505,7 @@ def setup_env(CP=None, cfg=None, controlsState=None, lr=None, fingerprint=None,
       params.put_bool("ExperimentalLongitudinalEnabled", True)
 
 
-def check_enabled(msgs):
+def check_openpilot_enabled(msgs):
   cur_enabled_count = 0
   max_enabled_count = 0
   for msg in msgs:

selfdrive/test/process_replay/regen.py

@@ -19,7 +19,7 @@ from panda.python import Panda
 from selfdrive.car.toyota.values import EPS_SCALE
 from selfdrive.manager.process import ensure_running
 from selfdrive.manager.process_config import managed_processes
-from selfdrive.test.process_replay.process_replay import CONFIGS, FAKEDATA, setup_env, check_enabled
+from selfdrive.test.process_replay.process_replay import CONFIGS, FAKEDATA, setup_env, check_openpilot_enabled
 from selfdrive.test.update_ci_routes import upload_route
 from tools.lib.route import Route
 from tools.lib.framereader import FrameReader
@@ -343,7 +343,7 @@ def regen_segment(lr, frs=None, daemons="all", outdir=FAKEDATA, disable_tqdm=Fal
   segment = params.get("CurrentRoute", encoding='utf-8') + "--0"
   seg_path = os.path.join(outdir, segment)
   # check to make sure openpilot is engaged in the route
-  if not check_enabled(LogReader(os.path.join(seg_path, "rlog"))):
+  if not check_openpilot_enabled(LogReader(os.path.join(seg_path, "rlog"))):
     raise Exception(f"Route did not engage for long enough: {segment}")
 
   return seg_path

selfdrive/test/process_replay/test_processes.py

@@ -9,11 +9,12 @@ from typing import Any, DefaultDict, Dict
 
 from selfdrive.car.car_helpers import interface_names
 from selfdrive.test.openpilotci import get_url, upload_file
-from selfdrive.test.process_replay.compare_logs import compare_logs, save_log
-from selfdrive.test.process_replay.process_replay import CONFIGS, PROC_REPLAY_DIR, FAKEDATA, check_enabled, replay_process
+from selfdrive.test.process_replay.compare_logs import compare_logs
+from selfdrive.test.process_replay.process_replay import CONFIGS, PROC_REPLAY_DIR, FAKEDATA, check_openpilot_enabled, replay_process
 from system.version import get_commit
 from tools.lib.filereader import FileReader
 from tools.lib.logreader import LogReader
+from tools.lib.helpers import save_log
 
 source_segments = [
   ("BODY", "937ccb7243511b65|2022-05-24--16-03-09--1"),        # COMMA.BODY
@@ -104,7 +105,7 @@ def test_process(cfg, lr, segment, ref_log_path, new_log_path, ignore_fields=Non
 
   # check to make sure openpilot is engaged in the route
   if cfg.proc_name == "controlsd":
-    if not check_enabled(log_msgs):
+    if not check_openpilot_enabled(log_msgs):
       return f"Route did not enable at all or for long enough: {new_log_path}", log_msgs
 
   try:

tools/lib/helpers.py

@@ -1,3 +1,4 @@
+import bz2
 import datetime
 
 TIME_FMT = "%Y-%m-%d--%H-%M-%S"
@@ -13,8 +14,19 @@ class RE:
   EXPLORER_FILE = r'^(?P<segment_name>{})--(?P<file_name>[a-z]+\.[a-z0-9]+)$'.format(SEGMENT_NAME)
   OP_SEGMENT_DIR = r'^(?P<segment_name>{})$'.format(SEGMENT_NAME)
 
+
 def timestamp_to_datetime(t: str) -> datetime.datetime:
   """
     Convert an openpilot route timestamp to a python datetime
   """
   return datetime.datetime.strptime(t, TIME_FMT)
+
+
+def save_log(dest, log_msgs, compress=True):
+  dat = b"".join(msg.as_builder().to_bytes() for msg in log_msgs)
+
+  if compress:
+    dat = bz2.compress(dat)
+
+  with open(dest, "wb") as f:
+    f.write(dat)

tools/plotjuggler/juggle.py

@@ -11,10 +11,10 @@ import requests
 import argparse
 
 from common.basedir import BASEDIR
-from selfdrive.test.process_replay.compare_logs import save_log
 from selfdrive.test.openpilotci import get_url
 from tools.lib.logreader import LogReader
 from tools.lib.route import Route, SegmentName
+from tools.lib.helpers import save_log
 from urllib.parse import urlparse, parse_qs
 
 juggle_dir = os.path.dirname(os.path.realpath(__file__))