diff --git a/openpilot/common/file_chunker.py b/openpilot/common/file_chunker.py index 3baa618fb1..2d080c3fff 100755 --- a/openpilot/common/file_chunker.py +++ b/openpilot/common/file_chunker.py @@ -40,13 +40,9 @@ def get_existing_chunks(path): raise FileNotFoundError(path) class ChunkStream(io.RawIOBase): - def __init__(self, paths, total=0, progress_cb=None): + def __init__(self, paths): self._paths = iter(paths) self._f = None - self._total = total - self._progress_cb = progress_cb - self._read = 0 - self._last_pct = -1 def readable(self): return True @@ -66,15 +62,9 @@ class ChunkStream(io.RawIOBase): self._f = None continue n += count - self._read += n - if self._progress_cb and self._total: - pct = min(100, self._read * 100 // self._total) - if pct != self._last_pct: - self._last_pct = pct - self._progress_cb(pct) return n -def open_file_chunked(path, progress_cb=None): +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()) @@ -83,8 +73,7 @@ def open_file_chunked(path, progress_cb=None): paths = [path] else: raise FileNotFoundError(path) - total = sum(os.path.getsize(p) for p in paths) - return io.BufferedReader(ChunkStream(paths, total, progress_cb)) + return io.BufferedReader(ChunkStream(paths)) if __name__ == "__main__": diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index f31a25771d..a8eeb924da 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -24,7 +24,7 @@ from opendbc.car.car_helpers import get_demo_car_params from tinygrad.tensor import Tensor -from openpilot.common.file_chunker import open_file_chunked +from openpilot.common.file_chunker import open_file_chunked, get_existing_chunks from openpilot.common.swaglog import cloudlog from openpilot.common.params import Params from openpilot.common.filter_simple import FirstOrderFilter @@ -56,6 +56,32 @@ def _pkl_exists(path): return os.path.exists(path) or os.path.exists(get_manifest_path(path)) +class _ProgressReader: + # wraps a chunked stream, reports byte-read % via cb (throttled to whole percent) + def __init__(self, inner, total, cb): + self._inner, self._total, self._cb = inner, total, cb + self._read = 0 + self._pct = -1 + + def _bump(self, n): + self._read += n + if self._total: + pct = min(100, self._read * 100 // self._total) + if pct != self._pct: + self._pct = pct + self._cb(pct) + + def read(self, size=-1): + data = self._inner.read(size) + self._bump(len(data)) + return data + + def readinto(self, b): + n = self._inner.readinto(b) + self._bump(n) + return n + + def _find_driving_pkl(bundle): if (override := os.environ.get('COMBINED_MODEL_PKL')) and _pkl_exists(override): return override @@ -107,8 +133,11 @@ class ModelState(ModelStateBase): def _init_combined(self, pkl_path, cam_w, cam_h, bundle): cloudlog.warning(f"loading combined pkl: {pkl_path}") - progress_cb = (lambda pct: Params().put("UsbGpuLoadProgress", pct)) if self.usbgpu else None - jits = load_oob(open_file_chunked(pkl_path, progress_cb)) + stream = open_file_chunked(pkl_path) + if self.usbgpu: + total = sum(os.path.getsize(p) for p in get_existing_chunks(pkl_path)) + stream = _ProgressReader(stream, total, lambda pct: Params().put("UsbGpuLoadProgress", pct)) + jits = load_oob(stream) self.WARP_DEV = 'QCOM' if COMMA_HARDWARE else 'CPU' self.DEV = 'AMD' if self.usbgpu else self.WARP_DEV