From d6af0e6eb50960c9bb7c4fffd44dfe87cb6888c1 Mon Sep 17 00:00:00 2001 From: Bruce Wayne Date: Fri, 20 Feb 2026 16:43:43 -0800 Subject: [PATCH] Revert "Simpler file chunker (#37276)" This reverts commit b27fa58444766c47f1ef980af79d324a669012c4. --- .gitignore | 3 +- common/file_chunker.py | 31 ------------------ selfdrive/modeld/SConscript | 45 ++++++++++++++------------- selfdrive/modeld/dmonitoringmodeld.py | 7 +++-- selfdrive/modeld/external_pickle.py | 38 ++++++++++++++++++++++ selfdrive/modeld/modeld.py | 9 +++--- 6 files changed, 72 insertions(+), 61 deletions(-) delete mode 100644 common/file_chunker.py create mode 100755 selfdrive/modeld/external_pickle.py diff --git a/.gitignore b/.gitignore index 062801d78..af18f0662 100644 --- a/.gitignore +++ b/.gitignore @@ -64,7 +64,8 @@ flycheck_* cppcheck_report.txt comma*.sh -selfdrive/modeld/models/*.pkl* +selfdrive/modeld/models/*.pkl +selfdrive/modeld/models/*.pkl.* # openpilot log files *.bz2 diff --git a/common/file_chunker.py b/common/file_chunker.py deleted file mode 100644 index f03d04a38..000000000 --- a/common/file_chunker.py +++ /dev/null @@ -1,31 +0,0 @@ -import math -import os -from pathlib import Path - -CHUNK_SIZE = 49 * 1024 * 1024 # 49MB, under GitHub's 50MB limit - -def get_chunk_name(name, idx, num_chunks): - return f"{name}.chunk{idx+1:02d}of{num_chunks:02d}" - -def get_chunk_paths(path, file_size): - num_chunks = math.ceil(file_size / CHUNK_SIZE) - return [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)] - -def chunk_file(path, num_chunks): - with open(path, 'rb') as f: - data = f.read() - actual_num_chunks = max(1, math.ceil(len(data) / CHUNK_SIZE)) - assert num_chunks >= actual_num_chunks, f"Allowed {num_chunks} chunks but needs at least {actual_num_chunks}, for path {path}" - for i in range(num_chunks): - with open(get_chunk_name(path, i, num_chunks), 'wb') as f: - f.write(data[i * CHUNK_SIZE:(i + 1) * CHUNK_SIZE]) - - -def read_file_chunked(path): - for num_chunks in range(1, 100): - if os.path.isfile(get_chunk_name(path, 0, num_chunks)): - files = [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)] - return b''.join(Path(f).read_bytes() for f in files) - if os.path.isfile(path): - return Path(path).read_bytes() - raise FileNotFoundError(path) diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 35be2acc0..1808cfec2 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,18 +1,14 @@ import os import glob -from openpilot.common.file_chunker import chunk_file, get_chunk_paths Import('env', 'arch') -chunker_file = File("#common/file_chunker.py") 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) if 'pycache' not in x and os.path.isfile(os.path.join(tinygrad_root, x))] -def estimate_pickle_max_size(onnx_size): - return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty - # Get model metadata for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: fn = File(f"models/{model_name}").abspath @@ -30,34 +26,39 @@ 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 ' +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] -def chunk_warps(target, source, env): - for t in warp_targets: - chunk_file(t, 1) -chunk_targets = sum([get_chunk_paths(t, estimate_pickle_max_size(0)) for t in warp_targets], []) -lenv.Command(chunk_targets, tinygrad_files + script_files + [chunker_file], - [compile_warp_cmd, chunk_warps]) +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 - pkl = fn + "_tinygrad.pkl" - onnx_path = fn + ".onnx" - chunk_targets = get_chunk_paths(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path))) - def do_chunk(target, source, env): - chunk_file(pkl, len(chunk_targets)) - return lenv.Command( - chunk_targets, - [onnx_path] + tinygrad_files + [chunker_file], - [f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', - do_chunk] + + 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 {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 6befe210a..956ea8a6a 100755 --- a/selfdrive/modeld/dmonitoringmodeld.py +++ b/selfdrive/modeld/dmonitoringmodeld.py @@ -16,8 +16,8 @@ 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 +from openpilot.selfdrive.modeld.external_pickle import load_external_pickle PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') @@ -45,7 +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 - self.model_run = pickle.loads(read_file_chunked(str(MODEL_PKL_PATH))) + 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 @@ -55,7 +55,8 @@ 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' - self.image_warp = pickle.loads(read_file_chunked(str(warp_path))) + with open(warp_path, "rb") as f: + self.image_warp = pickle.load(f) ptr = buf.data.ctypes.data # There is a ringbuffer of imgs, just cache tensors pointing to all of them if ptr not in self._blob_cache: 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 8cfbea02c..3fe3e0e6d 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -27,8 +27,8 @@ 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 +from openpilot.selfdrive.modeld.external_pickle import load_external_pickle PROCESS_NAME = "selfdrive.modeld.modeld" @@ -178,8 +178,8 @@ class ModelState: self.parser = Parser() self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {} self.update_imgs = None - self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH))) - self.policy_run = pickle.loads(read_file_chunked(str(POLICY_PKL_PATH))) + 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()} @@ -196,7 +196,8 @@ 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' - self.update_imgs = pickle.loads(read_file_chunked(str(warp_path))) + with open(warp_path, "rb") as f: + self.update_imgs = pickle.load(f) for key in bufs.keys(): ptr = bufs[key].data.ctypes.data