update SConscript

Author: Armandpl

selfdrive/modeld/SConscript

@@ -31,13 +31,25 @@ image_flag = {
      'larch64': 'IMAGE=2',
 }.get(arch, 'IMAGE=0')
 script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)]
-compile_warp_cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
+compile_warp_cmd = f'{tg_flags} {image_flag} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
+driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_on_policy', 'driving_off_policy']]
+driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_on_policy']]
 warp_targets = []
+policy_pkls = []
 for cam in [_ar_ox_fisheye, _os_fisheye]:
   w, h = cam.width, cam.height
+  policy_pkls.append(File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath)
   warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath]
-lenv.Command(warp_targets, tinygrad_files + script_files, compile_warp_cmd)
+compile_node = lenv.Command(warp_targets, tinygrad_files + script_files + driving_onnx_deps + driving_metadata_deps + [chunker_file], compile_warp_cmd)
+
+# chunk the combined policy pkls (they contain model weights)
+for policy_pkl in policy_pkls:
+  onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
+  chunk_targets = get_chunk_paths(policy_pkl, estimate_pickle_max_size(onnx_sizes_sum))
+  def do_chunk(target, source, env, pkl=policy_pkl, chunks=chunk_targets):
+    chunk_file(pkl, chunks)
+  lenv.Command(chunk_targets, compile_node, do_chunk)
 
 def tg_compile(flags, model_name):
   pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
@@ -58,6 +70,5 @@ def tg_compile(flags, model_name):
     do_chunk,
   )
 
-# Compile small models
-for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']:
-  tg_compile(tg_flags, model_name)
+# Compile dmonitoring model (driving models are in the combined policy JIT)
+tg_compile(tg_flags, 'dmonitoring_model')