diff --git a/openpilot/common/file_chunker.py b/openpilot/common/file_chunker.py index 57dfc35531..74f3013f71 100755 --- a/openpilot/common/file_chunker.py +++ b/openpilot/common/file_chunker.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +import io import sys import math import os @@ -21,13 +22,12 @@ def get_chunk_targets(path, file_size): 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)) + actual_num_chunks = max(1, math.ceil(os.path.getsize(path) / 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]) + with open(path, 'rb') as f: + for chunk_path in chunk_paths: + with open(chunk_path, 'wb') as out: + out.write(f.read(CHUNK_SIZE)) Path(manifest_path).write_text(str(len(chunk_paths))) os.remove(path) @@ -39,14 +39,40 @@ def get_existing_chunks(path): return _chunk_paths(path, num_chunks) raise FileNotFoundError(path) -def read_file_chunked(path): +class ChunkStream(io.RawIOBase): + def __init__(self, paths): + self._paths = iter(paths) + self._buf = memoryview(b'') + + def readable(self): + return True + + def readinto(self, b): + n = 0 + while n < len(b): + if not self._buf: + p = next(self._paths, None) + if p is None: + break + with open(p, 'rb') as f: + self._buf = memoryview(f.read()) + continue + take = min(len(b) - n, len(self._buf)) + b[n:n + take] = self._buf[:take] + self._buf = self._buf[take:] + n += take + return n + +def open_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) + paths = [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)] + elif os.path.isfile(path): + paths = [path] + else: + raise FileNotFoundError(path) + return io.BufferedReader(ChunkStream(paths)) if __name__ == "__main__": diff --git a/openpilot/selfdrive/modeld/compile_modeld.py b/openpilot/selfdrive/modeld/compile_modeld.py index 20e2eeeded..769fb69eb3 100755 --- a/openpilot/selfdrive/modeld/compile_modeld.py +++ b/openpilot/selfdrive/modeld/compile_modeld.py @@ -6,11 +6,14 @@ import os import pickle import tempfile import time +import shutil from functools import partial from collections import namedtuple import numpy as np +from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob + def _patch_tinygrad_fetch_fw(): import hashlib import pathlib @@ -27,6 +30,23 @@ def _patch_tinygrad_fetch_fw(): helpers.fetch_fw = fetch_fw _patch_tinygrad_fetch_fw() +def _patch_tinygrad_buffer_reduce(): + from tinygrad.device import Buffer + def __reduce_ex__(self, protocol): + buf = None + if self._base is not None: + return self.__class__, (self.device, self.size, self.dtype, None, None, None, 0, self.base, self.offset, self.is_allocated()) + if self.device == "NPY": + return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None, self.uop_refcount) + if self.is_allocated(): + buf = bytearray(self.nbytes) + self.copyout(memoryview(buf)) + if protocol >= 5: + buf = pickle.PickleBuffer(buf) + return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount) + Buffer.__reduce_ex__ = __reduce_ex__ +_patch_tinygrad_buffer_reduce() + from tinygrad.tensor import Tensor from tinygrad.helpers import Context from tinygrad.device import Device @@ -255,7 +275,10 @@ def compile_jit(jit, make_random_inputs, input_keys, make_queues): print('capture + replay') test_val, test_buffers = random_inputs_run(jit, SEED) print('pickle round trip') - jit = pickle.loads(pickle.dumps(jit)) + with tempfile.TemporaryFile(dir=".") as f: + dump_oob(jit, f) + f.seek(0) + jit = load_oob(f) random_inputs_run(jit, SEED, test_val, test_buffers, expect_match=True) random_inputs_run(jit, SEED+1, test_val, test_buffers, expect_match=False) return jit @@ -266,12 +289,11 @@ def _parse_size(s): return int(w), int(h) -def read_file_chunked_to_shm(path): - from openpilot.common.file_chunker import read_file_chunked - from openpilot.common.hardware.hw import Paths - with tempfile.NamedTemporaryFile(prefix='compile_modeld_', dir=Paths.shm_path(), delete=False) as f: - f.write(read_file_chunked(path)) - tmp_path = f.name +def read_file_chunked_to_disk(path): + from openpilot.common.file_chunker import open_file_chunked + tmp_path = f'{path}.unchunked' + with open(tmp_path, 'wb') as f, open_file_chunked(path) as src: + shutil.copyfileobj(src, f) atexit.register(lambda: os.path.exists(tmp_path) and os.remove(tmp_path)) return tmp_path @@ -289,7 +311,7 @@ if __name__ == "__main__": p.add_argument('--frame-skip', type=int, required=True) args = p.parse_args() - model_path = read_file_chunked_to_shm(args.onnx) + model_path = read_file_chunked_to_disk(args.onnx) model_w, model_h = args.model_size model_runner = OnnxRunner(model_path) @@ -310,5 +332,5 @@ if __name__ == "__main__": out[(cam_w,cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues) with open(args.output, "wb") as f: - pickle.dump(out, f) + dump_oob(out, f) print(f"Saved JITs to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)") diff --git a/openpilot/selfdrive/modeld/dmonitoringmodeld.py b/openpilot/selfdrive/modeld/dmonitoringmodeld.py index 2c659d9ac6..554407a223 100755 --- a/openpilot/selfdrive/modeld/dmonitoringmodeld.py +++ b/openpilot/selfdrive/modeld/dmonitoringmodeld.py @@ -14,7 +14,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.common.file_chunker import open_file_chunked from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp PROCESS_NAME = "openpilot.selfdrive.modeld.dmonitoringmodeld" @@ -43,7 +43,7 @@ class ModelState: self.frame_buf_params = get_nv12_info(cam_w, cam_h) self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} self._blob_cache : dict[int, Tensor] = {} - self.model_run = pickle.loads(read_file_chunked(str(MODEL_PKL_PATH))) + self.model_run = pickle.load(open_file_chunked(str(MODEL_PKL_PATH))) with open(MODELS_DIR / f'dm_warp_{cam_w}x{cam_h}_tinygrad.pkl', "rb") as f: self.image_warp = pickle.load(f) diff --git a/openpilot/selfdrive/modeld/helpers.py b/openpilot/selfdrive/modeld/helpers.py index 64bf28873f..608bc6aa64 100644 --- a/openpilot/selfdrive/modeld/helpers.py +++ b/openpilot/selfdrive/modeld/helpers.py @@ -1,4 +1,9 @@ +import io import json +import pickle +import shutil +import struct +import tempfile from pathlib import Path MODELS_DIR = Path(__file__).resolve().parent / 'models' @@ -15,6 +20,34 @@ def modeld_pkl_path(usbgpu: bool): prefix = 'big_' if usbgpu else '' return MODELS_DIR / f'{prefix}driving_tinygrad.pkl' +def dump_oob(obj, f): + with tempfile.TemporaryFile(dir=".") as tmp: + def buffer_callback(pb: pickle.PickleBuffer): + m = pb.raw() + tmp.write(struct.pack(' bool: for d in Path("/sys/bus/usb/devices").glob("*"): try: diff --git a/openpilot/selfdrive/modeld/modeld.py b/openpilot/selfdrive/modeld/modeld.py index 3d5e04873d..115873158f 100755 --- a/openpilot/selfdrive/modeld/modeld.py +++ b/openpilot/selfdrive/modeld/modeld.py @@ -3,7 +3,6 @@ import os os.environ['GMMU'] = '0' # for usbgpu fast loading, noop for qcom from tinygrad.tensor import Tensor import time -import pickle import numpy as np import openpilot.cereal.messaging as messaging from openpilot.cereal import log @@ -23,9 +22,9 @@ 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.compile_modeld import make_input_queues, WARP_INPUTS, POLICY_INPUTS from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_driving_model_data, fill_pose_msg, PublishState -from openpilot.common.file_chunker import read_file_chunked, get_manifest_path +from openpilot.common.file_chunker import open_file_chunked, get_manifest_path from openpilot.selfdrive.modeld.constants import ModelConstants, Plan -from openpilot.selfdrive.modeld.helpers import usbgpu_present, modeld_pkl_path, get_tg_input_devices +from openpilot.selfdrive.modeld.helpers import usbgpu_present, modeld_pkl_path, get_tg_input_devices, load_oob PROCESS_NAME = "openpilot.selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') @@ -79,7 +78,7 @@ class ModelState: def __init__(self, cam_w: int, cam_h: int, usbgpu: bool): input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu) self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV'] - jits = pickle.loads(read_file_chunked(modeld_pkl_path(usbgpu))) + jits = load_oob(open_file_chunked(modeld_pkl_path(usbgpu))) metadata = jits['metadata'] self.input_shapes = metadata['input_shapes'] self.vision_input_names = [k for k in self.input_shapes if 'img' in k]