From 2ba6df2506a3b3e664131f67513b5bcd0c08075d Mon Sep 17 00:00:00 2001 From: YassineYousfi Date: Fri, 13 Feb 2026 10:14:24 -0800 Subject: [PATCH] chunk tinygrad pkl below GitHub max size - NoCache and AlwaysBuild (#37194) * nocache * + * fixes * lint * not split * use pathlib * cleanup * better * even better --- .gitignore | 1 + release/build_release.sh | 3 +-- selfdrive/modeld/SConscript | 23 +++++++++++++--- selfdrive/modeld/dmonitoringmodeld.py | 4 +-- selfdrive/modeld/external_pickle.py | 38 +++++++++++++++++++++++++++ selfdrive/modeld/modeld.py | 8 +++--- 6 files changed, 65 insertions(+), 12 deletions(-) create mode 100755 selfdrive/modeld/external_pickle.py diff --git a/.gitignore b/.gitignore index 1fef0a125..af18f0662 100644 --- a/.gitignore +++ b/.gitignore @@ -65,6 +65,7 @@ cppcheck_report.txt comma*.sh selfdrive/modeld/models/*.pkl +selfdrive/modeld/models/*.pkl.* # openpilot log files *.bz2 diff --git a/release/build_release.sh b/release/build_release.sh index 220da05c1..7bc6732c6 100755 --- a/release/build_release.sh +++ b/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 {} + diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index b13a84f70..1808cfec2 100644 --- a/selfdrive/modeld/SConscript +++ b/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) diff --git a/selfdrive/modeld/dmonitoringmodeld.py b/selfdrive/modeld/dmonitoringmodeld.py index 94871c199..956ea8a6a 100755 --- a/selfdrive/modeld/dmonitoringmodeld.py +++ b/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 diff --git a/selfdrive/modeld/external_pickle.py b/selfdrive/modeld/external_pickle.py new file mode 100755 index 000000000..d60a9632a --- /dev/null +++ b/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])) diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index df000979e..a8b633bf0 100755 --- a/selfdrive/modeld/modeld.py +++ b/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()}