mirror of
https://github.com/infiniteCable2/openpilot.git
synced 2026-08-04 15:56:05 +08:00
Armand du Parc Locmaria
a2c554805b
* 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
60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
import io
|
|
import json
|
|
import pickle
|
|
import shutil
|
|
import struct
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
MODELS_DIR = Path(__file__).resolve().parent / 'models'
|
|
TG_INPUT_DEVICES_PATH = MODELS_DIR / 'tg_input_devices.json'
|
|
USBGPU_VID = 0xADD1
|
|
USBGPU_PID = 0x0001
|
|
|
|
|
|
def get_tg_input_devices(process_name: str, usbgpu: bool):
|
|
with open(TG_INPUT_DEVICES_PATH) as f:
|
|
return json.load(f)[process_name]['default' if not usbgpu else 'usbgpu']
|
|
|
|
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:
|
|
if int((d / "idVendor").read_text(), 16) == USBGPU_VID and \
|
|
int((d / "idProduct").read_text(), 16) == USBGPU_PID:
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return False
|