diff --git a/.gitignore b/.gitignore index af18f0662..062801d78 100644 --- a/.gitignore +++ b/.gitignore @@ -64,8 +64,7 @@ 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 new file mode 100644 index 000000000..ac9ddbb38 --- /dev/null +++ b/common/file_chunker.py @@ -0,0 +1,37 @@ +import math +import os +from pathlib import Path + +CHUNK_SIZE = 45 * 1024 * 1024 # 45MB, 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_manifest_path(name): + return f"{name}.chunkmanifest" + +def get_chunk_paths(path, file_size): + num_chunks = math.ceil(file_size / CHUNK_SIZE) + return [get_manifest_path(path)] + [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)] + +def chunk_file(path, targets): + manifest_path, *chunk_paths = targets + with open(path, 'rb') as f: + data = f.read() + actual_num_chunks = max(1, math.ceil(len(data) / CHUNK_SIZE)) + assert len(chunk_paths) >= actual_num_chunks, f"Allowed {len(chunk_paths)} chunks but needs at least {actual_num_chunks}, for path {path}" + for i, chunk_path in enumerate(chunk_paths): + with open(chunk_path, 'wb') as f: + f.write(data[i * CHUNK_SIZE:(i + 1) * CHUNK_SIZE]) + Path(manifest_path).write_text(str(len(chunk_paths))) + os.remove(path) + + +def read_file_chunked(path): + manifest_path = get_manifest_path(path) + if os.path.isfile(manifest_path): + num_chunks = int(Path(manifest_path).read_text().strip()) + return b''.join(Path(get_chunk_name(path, i, num_chunks)).read_bytes() for i in range(num_chunks)) + if os.path.isfile(path): + return Path(path).read_bytes() + raise FileNotFoundError(path) diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 1808cfec2..ed5c50beb 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,14 +1,18 @@ 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 @@ -26,39 +30,29 @@ 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(warp_targets, tinygrad_files + script_files, compile_warp_cmd) def tg_compile(flags, model_name): pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' fn = File(f"models/{model_name}").abspath - - 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}' + 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, 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] ) - 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 956ea8a6a..28190db3e 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 = load_external_pickle(MODEL_PKL_PATH) + 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 diff --git a/selfdrive/modeld/external_pickle.py b/selfdrive/modeld/external_pickle.py deleted file mode 100755 index d60a9632a..000000000 --- a/selfdrive/modeld/external_pickle.py +++ /dev/null @@ -1,38 +0,0 @@ -#!/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 3fe3e0e6d..bff59366d 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 = load_external_pickle(VISION_PKL_PATH) - self.policy_run = load_external_pickle(POLICY_PKL_PATH) + 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()}