move out

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

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