From 74554a523f57c5da17cd40bf98d42d62a2dcfcef Mon Sep 17 00:00:00 2001 From: Armand du Parc Locmaria Date: Thu, 14 May 2026 21:17:08 -0700 Subject: [PATCH] modeld: fold metadata into jit pkl (#38042) * modeld: fold metadata into jit pkl * modeld * no more metadata deps --- selfdrive/modeld/SConscript | 19 +++++++++--------- selfdrive/modeld/compile_modeld.py | 27 ++++++++++++-------------- selfdrive/modeld/get_model_metadata.py | 18 +++++++---------- selfdrive/modeld/modeld.py | 25 ++++++++++-------------- 4 files changed, 38 insertions(+), 51 deletions(-) diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 29ac9a47c..73dcb4a6f 100644 --- a/selfdrive/modeld/SConscript +++ b/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 diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index f6ad39892..f5d92edce 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/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)] } diff --git a/selfdrive/modeld/get_model_metadata.py b/selfdrive/modeld/get_model_metadata.py index ee08e9fb1..e4c173957 100755 --- a/selfdrive/modeld/get_model_metadata.py +++ b/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}') diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 5784099a3..527cc5d7a 100755 --- a/selfdrive/modeld/modeld.py +++ b/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(),