From 8196b743afa5d321983bb0728e3e206c3af10e5d Mon Sep 17 00:00:00 2001 From: Armand du Parc Locmaria Date: Tue, 30 Jun 2026 13:13:23 -0700 Subject: [PATCH] Revert "usbgpu: never materialize model weights in RAM" (#38271) Revert "usbgpu: never materialize model weights in RAM (#38265)" This reverts commit 724971afc381b864edd6d3ab4ff57cb2710a4299. --- openpilot/common/file_chunker.py | 41 +++++-------------- openpilot/selfdrive/modeld/compile_modeld.py | 40 ++++-------------- .../selfdrive/modeld/dmonitoringmodeld.py | 4 +- openpilot/selfdrive/modeld/helpers.py | 29 ------------- openpilot/selfdrive/modeld/modeld.py | 7 ++-- 5 files changed, 26 insertions(+), 95 deletions(-) diff --git a/openpilot/common/file_chunker.py b/openpilot/common/file_chunker.py index 6af17495b..57dfc3553 100755 --- a/openpilot/common/file_chunker.py +++ b/openpilot/common/file_chunker.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 -import io import sys import math import os @@ -22,12 +21,13 @@ def get_chunk_targets(path, file_size): def chunk_file(path, targets): manifest_path, *chunk_paths = targets - 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}" 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)) + 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) @@ -39,33 +39,14 @@ def get_existing_chunks(path): return _chunk_paths(path, num_chunks) raise FileNotFoundError(path) -class ChunkStream(io.RawIOBase): - def __init__(self, paths): - self._files = (open(p, 'rb') for p in paths) - self._cur = next(self._files, None) - - def readable(self): - return True - - def readinto(self, b): - while self._cur is not None: - n = self._cur.readinto(b) - if n: - return n - self._cur.close() - self._cur = next(self._files, None) - return 0 - -def open_file_chunked(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()) - 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)) + 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) if __name__ == "__main__": diff --git a/openpilot/selfdrive/modeld/compile_modeld.py b/openpilot/selfdrive/modeld/compile_modeld.py index 769fb69eb..20e2eeede 100755 --- a/openpilot/selfdrive/modeld/compile_modeld.py +++ b/openpilot/selfdrive/modeld/compile_modeld.py @@ -6,14 +6,11 @@ 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 @@ -30,23 +27,6 @@ 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 @@ -275,10 +255,7 @@ 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') - with tempfile.TemporaryFile(dir=".") as f: - dump_oob(jit, f) - f.seek(0) - jit = load_oob(f) + jit = pickle.loads(pickle.dumps(jit)) 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 @@ -289,11 +266,12 @@ def _parse_size(s): return int(w), int(h) -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) +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 atexit.register(lambda: os.path.exists(tmp_path) and os.remove(tmp_path)) return tmp_path @@ -311,7 +289,7 @@ if __name__ == "__main__": p.add_argument('--frame-skip', type=int, required=True) args = p.parse_args() - model_path = read_file_chunked_to_disk(args.onnx) + model_path = read_file_chunked_to_shm(args.onnx) model_w, model_h = args.model_size model_runner = OnnxRunner(model_path) @@ -332,5 +310,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: - dump_oob(out, f) + pickle.dump(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 554407a22..2c659d9ac 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 open_file_chunked +from openpilot.common.file_chunker import read_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.load(open_file_chunked(str(MODEL_PKL_PATH))) + self.model_run = pickle.loads(read_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 4a85248d9..64bf28873 100644 --- a/openpilot/selfdrive/modeld/helpers.py +++ b/openpilot/selfdrive/modeld/helpers.py @@ -1,9 +1,4 @@ -import io import json -import pickle -import shutil -import struct -import tempfile from pathlib import Path MODELS_DIR = Path(__file__).resolve().parent / 'models' @@ -20,30 +15,6 @@ 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 115873158..3d5e04873 100755 --- a/openpilot/selfdrive/modeld/modeld.py +++ b/openpilot/selfdrive/modeld/modeld.py @@ -3,6 +3,7 @@ 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 @@ -22,9 +23,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 open_file_chunked, get_manifest_path +from openpilot.common.file_chunker import read_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, load_oob +from openpilot.selfdrive.modeld.helpers import usbgpu_present, modeld_pkl_path, get_tg_input_devices PROCESS_NAME = "openpilot.selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') @@ -78,7 +79,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 = load_oob(open_file_chunked(modeld_pkl_path(usbgpu))) + jits = pickle.loads(read_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]