Reapply "usbgpu: never materialize model weights in RAM" (#38273)

* Reapply "usbgpu: never materialize model weights in RAM" (#38271)

This reverts commit 8196b743af.

* release oob buffers as we load them

* don't leak fds

* lint

* 10s faster loading time for big model

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__":

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)")

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)
 

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('<q', m.nbytes))
+      tmp.write(m)
+      pb.release() # keep peak ram at ~1 buffer
+    stream = io.BytesIO()
+    pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
+    opcodes = stream.getvalue()
+    f.write(struct.pack('<q', len(opcodes)))
+    f.write(opcodes)
+    tmp.seek(0)
+    shutil.copyfileobj(tmp, f)
+
+def load_oob(f):
+  opcodes = f.read(struct.unpack('<q', f.read(8))[0])
+  def buffers():
+    prev = None
+    while (h := f.read(8)):
+      if prev is not None:
+        prev.release()
+      buf = bytearray(struct.unpack('<q', h)[0])
+      f.readinto(buf)
+      prev = pickle.PickleBuffer(buf)
+      yield prev
+  return pickle.load(io.BytesIO(opcodes), buffers=buffers())
+
 def usbgpu_present() -> bool:
   for d in Path("/sys/bus/usb/devices").glob("*"):
     try:

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]