diff --git a/openpilot/starpilot/models/runtime_bridge/native.py b/openpilot/starpilot/models/runtime_bridge/native.py index 8ea101b544..1ac9cfd614 100644 --- a/openpilot/starpilot/models/runtime_bridge/native.py +++ b/openpilot/starpilot/models/runtime_bridge/native.py @@ -247,7 +247,6 @@ def _verified(path): def _load_driving(stream, unchanged): import io - import mmap import pickle import struct from tinygrad import Device @@ -262,21 +261,25 @@ def _load_driving(stream, unchanged): if stream.tell() == os.fstat(stream.fileno()).st_size: unchanged() return pickle.loads(opcodes) - # Map the already authenticated inode. Do not reopen a path or use a cached - # DISK:/proc/self/fd device whose descriptor number can be reused later. - mapping = mmap.mmap(stream.fileno(), 0, access=mmap.ACCESS_COPY) - view = memoryview(mapping)[stream.tell():] - buffers = Buffer(HCQ_RUNTIME_DEV.device, len(view), opaque=view) - arena, chunk_size = Buffer(Device.DEFAULT, buffers.nbytes, preallocate=True), 8 << 20 + # Copy from the authenticated descriptor with bounded host staging. Reserving + # a private mapping for every weight can fail even when GPU memory is free. + length = os.fstat(stream.fileno()).st_size - stream.tell() + arena = Buffer(Device.DEFAULT, length, preallocate=True) + chunk_size = min(length, 8 << 20) + staging = bytearray(chunk_size) + view = memoryview(staging) + buffers = Buffer(HCQ_RUNTIME_DEV.device, chunk_size, opaque=view) unchanged() - for off in range(0, buffers.nbytes, chunk_size): - size = min(chunk_size, buffers.nbytes - off) - arena.view(size, off).ensure_allocated().copy_from(buffers.view(size, off).ensure_allocated()) - Device.default.synchronize() + for off in range(0, length, chunk_size): + size = min(chunk_size, length - off) + if stream.readinto(view[:size]) != size: + raise ValueError('Truncated native OOB payload') + arena.view(size, off).ensure_allocated().copy_from(buffers.view(size, 0).ensure_allocated()) + # A staging window must remain unchanged until its transfer completes. + Device.default.synchronize() unchanged() del buffers view.release() - mapping.close() unpickler = pickle.Unpickler(io.BytesIO(opcodes)) cast("PersistentLoad", unpickler).persistent_load = lambda record: arena.view(*record) return unpickler.load() diff --git a/openpilot/starpilot/models/tests/test_native_driving_reader.py b/openpilot/starpilot/models/tests/test_native_driving_reader.py new file mode 100644 index 0000000000..67d0b08c92 --- /dev/null +++ b/openpilot/starpilot/models/tests/test_native_driving_reader.py @@ -0,0 +1,163 @@ +"""Exact native reader AST with inert GPU buffers and a real authenticated fd.""" +import ast +from contextlib import contextmanager +import hashlib +import os +from pathlib import Path +import pickle +import stat +import struct +import tempfile +from types import ModuleType, SimpleNamespace +from typing import Any, cast +import unittest +from unittest.mock import Mock, patch + +SOURCE = Path(os.environ.get('STARPILOT_NATIVE_READER_SOURCE', Path(__file__).resolve().parents[1] / 'runtime_bridge/native.py')) +CHUNK = 8 << 20 + + +class Stream: + def __init__(self, stream, events, pending, truncate=False): + self.stream, self.events, self.pending, self.truncate = stream, events, pending, truncate + def fileno(self): + return self.stream.fileno() + def tell(self): + return self.stream.tell() + def read(self, size): + self.events.append(('read', size)) + return self.stream.read(size) + def readinto(self, view): + if self.pending: + raise AssertionError('staging reused before synchronization') + self.events.append(('readinto', len(view))) + count = self.stream.readinto(view) + return count - 1 if self.truncate else count + + +class TestNativeDrivingReader(unittest.TestCase): + def setUp(self): + self.temp = self.enterContext(tempfile.TemporaryDirectory()) + self.path = Path(self.temp) / 'artifact.pkl' + self.events: list[Any] = [] + self.pending: list[Any] = [] + self.arenas: list[bytearray] = [] + self.stages: list[memoryview] = [] + self.allocations: list[Any] = [] + tree = ast.parse(SOURCE.read_text()) + nodes: list[ast.stmt] = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in ('_verified', '_load_driving')] + self.ns: dict[str, Any] = {"contextmanager": contextmanager, "os": os, "stat": stat, "hashlib": hashlib, + "Path": Path, "cast": cast, "_allowed": {}, "__name__": __name__} + exec(compile(ast.Module(body=nodes, type_ignores=[]), str(SOURCE), 'exec'), self.ns) + owner = self + + class Buffer: + def __init__(self, device, size, preallocate=False, opaque=None): + owner.allocations.append((device, size, preallocate, opaque is not None)) + self.device, self.offset, self.size = device, 0, size + self.data = opaque if opaque is not None else bytearray(size) + if opaque is not None: + owner.stages.append(opaque) + else: + owner.arenas.append(self.data) + def view(self, size, offset): + if size < 0 or offset < 0 or offset + size > self.size: + raise ValueError('view outside arena') + return View(self.data, self.offset + offset, size) + def __del__(self): + if self.device == 'HOST': + owner.events.append('staging-buffer-released') + + class View: + def __init__(self, data, offset, size): + self.data, self.offset, self.size = data, offset, size + def ensure_allocated(self): + return self + def copy_from(self, source): + owner.events.append(('copy', self.offset, self.size, source.offset)) + owner.pending.append((self, source)) + def synchronize(): + for target, source in self.pending: + target.data[target.offset:target.offset + target.size] = source.data[source.offset:source.offset + source.size] + self.pending.clear() + self.events.append('sync') + tg = ModuleType('tinygrad') + tg.Device = SimpleNamespace(DEFAULT='GPU', default=SimpleNamespace(synchronize=synchronize)) + device = ModuleType('tinygrad.device') + device.Buffer, device.HCQ_RUNTIME_DEV = Buffer, SimpleNamespace(device='HOST') + self.enterContext(patch.dict('sys.modules', {'tinygrad': tg, 'tinygrad.device': device})) + + def envelope(self, payload, opcodes=b'actual opaque opcodes'): + data = struct.pack('