chunk tinygrad pkl below GitHub max size - NoCache and AlwaysBuild (#37194)

* nocache

* +

* fixes

* lint

* not split

* use pathlib

* cleanup

* better

* even better

.gitignore

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

release/build_release.sh

@@ -72,8 +72,7 @@ find . -name '*.pyc' -delete
 find . -name 'moc_*' -delete
 find . -name '__pycache__' -delete
 rm -rf .sconsign.dblite Jenkinsfile release/
-rm selfdrive/modeld/models/driving_vision.onnx
-rm selfdrive/modeld/models/driving_policy.onnx
+rm -f selfdrive/modeld/models/*.onnx
 
 find third_party/ -name '*x86*' -exec rm -r {} +
 find third_party/ -name '*Darwin*' -exec rm -r {} +

selfdrive/modeld/SConscript

@@ -3,6 +3,7 @@ import glob
 
 Import('env', 'arch')
 lenv = env.Clone()
+CHUNK_BYTES = int(os.environ.get("TG_CHUNK_BYTES", str(45 * 1024 * 1024)))
 
 tinygrad_root = env.Dir("#").abspath
 tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
@@ -36,12 +37,28 @@ lenv.Command(warp_targets, tinygrad_files + script_files, cmd)
 def tg_compile(flags, model_name):
   pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
   fn = File(f"models/{model_name}").abspath
-  return lenv.Command(
-    fn + "_tinygrad.pkl",
+
+  out = fn + "_tinygrad.pkl"
+  full = out + ".full"
+  parts = out + ".parts"
+
+  full_node = lenv.Command(
+    full,
     [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 {full}'
   )
 
+  split_script = File(Dir("#selfdrive/modeld").File("external_pickle.py").abspath)
+  parts_node = lenv.Command(
+    parts,
+    [full_node, split_script, Value(str(CHUNK_BYTES))],
+    [f'python3 {split_script.abspath} {full} {out} {CHUNK_BYTES}', Delete(full)],
+  )
+
+  lenv.NoCache(parts_node)
+  lenv.AlwaysBuild(parts_node)
+  return parts_node
+
 # Compile small models
 for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
   tg_compile(tg_flags, model_name)

selfdrive/modeld/dmonitoringmodeld.py

@@ -17,6 +17,7 @@ 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.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp
+from openpilot.selfdrive.modeld.external_pickle import load_external_pickle
 
 PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
@@ -44,8 +45,7 @@ class ModelState:
     self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
     self._blob_cache : dict[int, Tensor] = {}
     self.image_warp = None
-    with open(MODEL_PKL_PATH, "rb") as f:
-      self.model_run = pickle.load(f)
+    self.model_run = load_external_pickle(MODEL_PKL_PATH)
 
   def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]:
     self.numpy_inputs['calib'][0,:] = calib

selfdrive/modeld/external_pickle.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+import hashlib
+import pickle
+import sys
+from pathlib import Path
+
+def split_pickle(full_path: Path, out_prefix: Path, chunk_bytes: int) -> None:
+  data = full_path.read_bytes()
+  out_dir = out_prefix.parent
+
+  for p in out_dir.glob(f"{out_prefix.name}.data-*"):
+    p.unlink()
+
+  total = (len(data) + chunk_bytes - 1) // chunk_bytes
+  names = []
+  for i in range(0, len(data), chunk_bytes):
+    name = f"{out_prefix.name}.data-{(i // chunk_bytes) + 1:04d}-of-{total:04d}"
+    (out_dir / name).write_bytes(data[i:i + chunk_bytes])
+    names.append(name)
+
+  manifest = hashlib.sha256(data).hexdigest() + "\n" + "\n".join(names) + "\n"
+  (out_dir / (out_prefix.name + ".parts")).write_text(manifest)
+
+def load_external_pickle(prefix: Path):
+  parts = prefix.parent / (prefix.name + ".parts")
+  lines = parts.read_text().splitlines()
+  expected_hash, chunk_names = lines[0], lines[1:]
+
+  data = bytearray()
+  for name in chunk_names:
+    data += (prefix.parent / name).read_bytes()
+
+  if hashlib.sha256(data).hexdigest() != expected_hash:
+    raise RuntimeError(f"hash mismatch loading {prefix}")
+  return pickle.loads(data)
+
+if __name__ == "__main__":
+  split_pickle(Path(sys.argv[1]), Path(sys.argv[2]), int(sys.argv[3]))

selfdrive/modeld/modeld.py

@@ -28,6 +28,7 @@ from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_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.selfdrive.modeld.constants import ModelConstants, Plan
+from openpilot.selfdrive.modeld.external_pickle import load_external_pickle
 
 
 PROCESS_NAME = "selfdrive.modeld.modeld"
@@ -177,11 +178,8 @@ class ModelState:
     self.parser = Parser()
     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 = load_external_pickle(VISION_PKL_PATH)
+    self.policy_run = load_external_pickle(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()}