modeld: fold metadata into jit pkl (#38042)

* modeld: fold metadata into jit pkl

* modeld

* no more metadata deps

selfdrive/modeld/SConscript

@@ -60,18 +60,9 @@ compiled_flags_node = lenv.Command(
 # tinygrad calls brew which needs a $HOME in the env
 mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''
 
-# Get model metadata
-for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
-  fn = File(f"models/{model_name}").abspath
-  script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)]
-  cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
-  lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [compiled_flags_node], cmd)
-
 modeld_dir = Dir("#selfdrive/modeld").abspath
 compile_modeld_script = [File(f"{modeld_dir}/compile_modeld.py")]
-compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")]
 driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']]
-driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']]
 
 model_w, model_h = MEDMODEL_INPUT_SIZE
 frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
@@ -83,14 +74,21 @@ cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
        f'--vision-onnx {File("models/driving_vision.onnx").abspath} '
        f'--policy-onnx {File("models/driving_policy.onnx").abspath} '
        f'--output {pkl_path} --frame-skip {frame_skip}')
-node = lenv.Command(pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [Value(camera_res_args), chunker_file, compiled_flags_node], cmd)
+node = lenv.Command(pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(camera_res_args), chunker_file, compiled_flags_node], cmd)
 onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
 chunk_targets = get_chunk_paths(pkl_path, estimate_pickle_max_size(onnx_sizes_sum)*2) # TODO make weight dedupe work on QCOM
 def do_chunk(target, source, env, pkl=pkl_path, chunks=chunk_targets):
   chunk_file(pkl, chunks)
 lenv.Command(chunk_targets, node, do_chunk)
 
+# get model metadata
+fn = File(f"models/dmonitoring_model").abspath
+script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)]
+cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
+lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [compiled_flags_node], cmd)
+
 dm_w, dm_h = DM_INPUT_SIZE
+compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")]
 for cam_w, cam_h in CAMERA_CONFIGS:
   dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
   cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_dm_warp.py '
@@ -98,6 +96,7 @@ for cam_w, cam_h in CAMERA_CONFIGS:
          f'--output {dm_pkl_path}')
   lenv.Command(dm_pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd)
 
+driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']]
 def tg_compile(flags, model_name):
   pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
   fn = File(f"models/{model_name}").abspath

selfdrive/modeld/compile_modeld.py

@@ -4,7 +4,7 @@ import os
 import pickle
 import time
 from functools import partial
-from collections import namedtuple
+from collections import namedtuple, defaultdict
 
 import numpy as np
 from tinygrad.tensor import Tensor
@@ -158,10 +158,13 @@ def make_run_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, mode
 
 
 def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
-                   vision_runner, policy_runner, vision_features_slice,
-                   vision_input_shapes, policy_input_shapes):
+                   vision_runner, policy_runner, vision_metadata, policy_metadata):
   print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
 
+  vision_features_slice = vision_metadata['output_slices']['hidden_state']
+  vision_input_shapes = vision_metadata['input_shapes']
+  policy_input_shapes = policy_metadata['input_shapes']
+
   _run = make_run_policy(vision_runner, policy_runner, nv12, model_w, model_h,
                          vision_features_slice, frame_skip, prepare_only)
   run_policy_jit = TinyJit(_run, prune=True)
@@ -229,26 +232,20 @@ if __name__ == "__main__":
   p.add_argument('--frame-skip', type=int, required=True)
   args = p.parse_args()
 
-  model_w, model_h = args.model_size
-
+  out = defaultdict(dict)
   # init runners once so weights are shared
-  from get_model_metadata import metadata_path_for
+  from get_model_metadata import make_metadata_dict
   vision_runner = OnnxRunner(args.vision_onnx)
   policy_runner = OnnxRunner(args.policy_onnx)
-  with open(metadata_path_for(args.vision_onnx), 'rb') as f:
-    vision_metadata = pickle.load(f)
-    vision_features_slice = vision_metadata['output_slices']['hidden_state']
-    vision_input_shapes = vision_metadata['input_shapes']
-  with open(metadata_path_for(args.policy_onnx), 'rb') as f:
-    policy_input_shapes = pickle.load(f)['input_shapes']
+  out['metadata']['vision'] = make_metadata_dict(args.vision_onnx)
+  out['metadata']['policy'] = make_metadata_dict(args.policy_onnx)
 
-  out = {}
   for cam_w, cam_h in args.camera_resolutions:
     nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+    model_w, model_h = args.model_size
     out[(cam_w,cam_h)] = {
       name: compile_modeld(nv12, model_w, model_h, prepare_only, args.frame_skip,
-                           vision_runner, policy_runner, vision_features_slice,
-                           vision_input_shapes, policy_input_shapes)
+                           vision_runner, policy_runner, out['metadata']['vision'], out['metadata']['policy'])
       for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
     }
 

selfdrive/modeld/get_model_metadata.py

@@ -7,10 +7,6 @@ from typing import Any
 
 from tinygrad.nn.onnx import OnnxPBParser
 
-def metadata_path_for(onnx_path) -> pathlib.Path:
-  p = pathlib.Path(onnx_path)
-  return p.parent / (p.stem + '_metadata.pkl')
-
 
 class MetadataOnnxPBParser(OnnxPBParser):
   def _parse_ModelProto(self) -> dict:
@@ -39,21 +35,21 @@ def get_metadata_value_by_name(model: dict[str, Any], name: str) -> str | Any:
   return None
 
 
-if __name__ == "__main__":
-  model_path = pathlib.Path(sys.argv[1])
+def make_metadata_dict(model_path):
   model = MetadataOnnxPBParser(model_path).parse()
   output_slices = get_metadata_value_by_name(model, 'output_slices')
   assert output_slices is not None, 'output_slices not found in metadata'
-
-  metadata = {
+  return {
     'model_checkpoint': get_metadata_value_by_name(model, 'model_checkpoint'),
     'output_slices': pickle.loads(codecs.decode(output_slices.encode(), "base64")),
     'input_shapes': dict(get_name_and_shape(x) for x in model["graph"]["input"]),
     'output_shapes': dict(get_name_and_shape(x) for x in model["graph"]["output"]),
   }
 
-  metadata_path = metadata_path_for(model_path)
-  with open(metadata_path, 'wb') as f:
-    pickle.dump(metadata, f)
 
+if __name__ == "__main__":
+  model_path = pathlib.Path(sys.argv[1])
+  metadata_path = model_path.parent / (model_path.stem + '_metadata.pkl')
+  with open(metadata_path, 'wb') as f:
+    pickle.dump(make_metadata_dict(model_path), f)
   print(f'saved metadata to {metadata_path}')

selfdrive/modeld/modeld.py

@@ -35,9 +35,6 @@ from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
 PROCESS_NAME = "selfdrive.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 
-VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl'
-POLICY_METADATA_PATH = MODELS_DIR / 'driving_policy_metadata.pkl'
-
 LAT_SMOOTH_SECONDS = 0.0
 LONG_SMOOTH_SECONDS = 0.3
 MIN_LAT_CONTROL_SPEED = 0.3
@@ -81,16 +78,15 @@ class ModelState:
   prev_desire: np.ndarray  # for tracking the rising edge of the pulse
 
   def __init__(self, cam_w: int, cam_h: int):
-    with open(VISION_METADATA_PATH, 'rb') as f:
-      vision_metadata = pickle.load(f)
-      self.vision_input_shapes =  vision_metadata['input_shapes']
-      self.vision_input_names = list(self.vision_input_shapes.keys())
-      self.vision_output_slices = vision_metadata['output_slices']
+    jits = pickle.loads(read_file_chunked(MODELS_DIR / 'driving_tinygrad.pkl'))
+    vision_metadata = jits['metadata']['vision']
+    self.vision_input_shapes =  vision_metadata['input_shapes']
+    self.vision_input_names = list(self.vision_input_shapes.keys())
+    self.vision_output_slices = vision_metadata['output_slices']
 
-    with open(POLICY_METADATA_PATH, 'rb') as f:
-      policy_metadata = pickle.load(f)
-      self.policy_input_shapes =  policy_metadata['input_shapes']
-      self.policy_output_slices = policy_metadata['output_slices']
+    policy_metadata = jits['metadata']['policy']
+    self.policy_input_shapes =  policy_metadata['input_shapes']
+    self.policy_output_slices = policy_metadata['output_slices']
 
     self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
 
@@ -100,9 +96,8 @@ class ModelState:
     self._blob_cache : dict[int, Tensor] = {}
     self.parser = Parser()
     self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')}
-    jits = pickle.loads(read_file_chunked(MODELS_DIR / 'driving_tinygrad.pkl'))[(cam_w,cam_h)]
-    self.run_policy = jits['run_policy']
-    self.warp_enqueue = jits['warp_enqueue']
+    self.run_policy = jits[(cam_w,cam_h)]['run_policy']
+    self.warp_enqueue = jits[(cam_w,cam_h)]['warp_enqueue']
     self.warp_enqueue(
       **self.input_queues,
       frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(),