Chunk big model files (#37134)

* file chunking

* try this

* more cleanup

* cleaner

.gitignore

@@ -64,7 +64,7 @@ flycheck_*
 cppcheck_report.txt
 comma*.sh
 
-selfdrive/modeld/models/*.pkl
+selfdrive/modeld/models/*.pkl*
 
 # openpilot log files
 *.bz2

common/file_chunker.py

@@ -0,0 +1,27 @@
+import glob
+import os
+import sys
+from pathlib import Path
+
+CHUNK_SIZE = 49 * 1024 * 1024  # 49MB, under GitHub's 50MB limit
+
+
+def chunk_file(path):
+  with open(path, 'rb') as f:
+    data = f.read()
+  for i in range(0, len(data), CHUNK_SIZE):
+    with open(f"{path}.chunk{i // CHUNK_SIZE:02d}", 'wb') as f:
+      f.write(data[i:i + CHUNK_SIZE])
+  os.remove(path)
+
+
+def read_file_chunked(path):
+  files = sorted(glob.glob(f"{path}.chunk*")) or ([path] if os.path.isfile(path) else [])
+  if not files:
+    raise FileNotFoundError(path)
+  return b''.join(Path(f).read_bytes() for f in files)
+
+
+if __name__ == "__main__":
+  for path in sys.argv[1:]:
+    chunk_file(path)

selfdrive/modeld/SConscript

@@ -8,6 +8,8 @@ tinygrad_root = env.Dir("#").abspath
 tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
                   if 'pycache' not in x and os.path.isfile(os.path.join(tinygrad_root, x))]
 
+chunk_cmd = f'python3 {Dir("#common").abspath}/file_chunker.py'
+
 # Get model metadata
 for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
   fn = File(f"models/{model_name}").abspath
@@ -24,21 +26,23 @@ image_flag = {
      'larch64': 'IMAGE=2',
 }.get(arch, 'IMAGE=0')
 script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)]
-cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
+compile_warp_cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 warp_targets = []
 for cam in [_ar_ox_fisheye, _os_fisheye]:
   w, h = cam.width, cam.height
   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, cmd)
+lenv.Command([t + ".chunk00" for t in warp_targets], tinygrad_files + script_files,
+             f'{compile_warp_cmd} && {chunk_cmd} {" ".join(warp_targets)}')
 
 def tg_compile(flags, model_name):
   pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
   fn = File(f"models/{model_name}").abspath
+  pkl = fn + "_tinygrad.pkl"
   return lenv.Command(
-    fn + "_tinygrad.pkl",
+    pkl + ".chunk00",
     [fn + ".onnx"] + tinygrad_files,
-    f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {fn}_tinygrad.pkl'
+    f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl} && {chunk_cmd} {pkl}'
   )
 
 # Compile small models

selfdrive/modeld/dmonitoringmodeld.py

@@ -16,6 +16,7 @@ from openpilot.common.realtime import config_realtime_process
 from openpilot.common.transformations.model import dmonitoringmodel_intrinsics
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+from openpilot.common.file_chunker import read_file_chunked
 from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp
 
 PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld"
@@ -43,8 +44,7 @@ class ModelState:
     self.frame_buf_params = None
     self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
     self.image_warp = None
-    with open(MODEL_PKL_PATH, "rb") as f:
-      self.model_run = pickle.load(f)
+    self.model_run = pickle.loads(read_file_chunked(str(MODEL_PKL_PATH)))
 
   def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]:
     self.numpy_inputs['calib'][0,:] = calib
@@ -54,8 +54,7 @@ class ModelState:
     if self.image_warp is None:
       self.frame_buf_params = get_nv12_info(buf.width, buf.height)
       warp_path = MODELS_DIR / f'dm_warp_{buf.width}x{buf.height}_tinygrad.pkl'
-      with open(warp_path, "rb") as f:
-        self.image_warp = pickle.load(f)
+      self.image_warp = pickle.loads(read_file_chunked(str(warp_path)))
     self.warp_inputs['frame'] = Tensor.from_blob(buf.data.ctypes.data, (self.frame_buf_params[3],), dtype='uint8').realize()
 
     self.warp_inputs_np['transform'][:] = transform[:]

selfdrive/modeld/modeld.py

@@ -27,6 +27,7 @@ from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
 from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value, get_curvature_from_plan
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
 from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState
+from openpilot.common.file_chunker import read_file_chunked
 from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
 
 
@@ -177,11 +178,8 @@ class ModelState:
     self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {}
     self.update_imgs = None
 
-    with open(VISION_PKL_PATH, "rb") as f:
-      self.vision_run = pickle.load(f)
-
-    with open(POLICY_PKL_PATH, "rb") as f:
-      self.policy_run = pickle.load(f)
+    self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH)))
+    self.policy_run = pickle.loads(read_file_chunked(str(POLICY_PKL_PATH)))
 
   def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
     parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
@@ -198,8 +196,7 @@ class ModelState:
         w, h = bufs[key].width, bufs[key].height
         self.frame_buf_params[key] = get_nv12_info(w, h)
       warp_path = MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl'
-      with open(warp_path, "rb") as f:
-        self.update_imgs = pickle.load(f)
+      self.update_imgs = pickle.loads(read_file_chunked(str(warp_path)))
 
 
     for key in bufs.keys():