modeld: fold metadata into jit pkl (#38042)

* modeld: fold metadata into jit pkl

* modeld

* no more metadata deps
This commit is contained in:
Armand du Parc Locmaria
2026-05-14 21:17:08 -07:00
committed by Shane Smiskol
parent 2d4ac33ed7
commit 74554a523f
4 changed files with 38 additions and 51 deletions
+9 -10
View File
@@ -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