Use a precompiled eGPU driving model (#38930)

* Ship precompiled eGPU model and camera warps

Compile f78ed37d-afad-4dbc-8050-40ea885eedde/12864 through xx/ml_tools/openpilot_compile using the pinned tinygrad version.

* Precompile the existing master driving model

Use the unchanged master ONNX (SHA-256 6fee5937923c74848df4a63f6239eb6331c6274dd4bdb7a5d6ec0388a8b543d5) instead of updating the trained model.

* Compile camera warps on device

* Remove obsolete ONNX chunking and big model build check

* Chunk model artifacts only during release packaging

* Require model and camera warps for Chestnut readiness

* Recompile precompiled CPU helpers for the runtime host

* Ship the eGPU model with an ARM submission helper

* Exempt model pickles from the build product size limit

.gitattributes

@@ -3,6 +3,7 @@
 # to move existing files into LFS:
 # git add --renormalize .
 *.onnx filter=lfs diff=lfs merge=lfs -text
+openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl filter=lfs diff=lfs merge=lfs -text
 *.svg filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 *.gif filter=lfs diff=lfs merge=lfs -text

.github/workflows/docs.yaml

@@ -18,7 +18,7 @@ concurrency:
 env:
   GIT_CONFIG_COUNT: 1
   GIT_CONFIG_KEY_0: lfs.fetchexclude
-  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx
+  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
 
 jobs:
   docs:

.github/workflows/release.yaml

@@ -7,7 +7,7 @@ on:
 env:
   GIT_CONFIG_COUNT: 1
   GIT_CONFIG_KEY_0: lfs.fetchexclude
-  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx
+  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
 
 jobs:
   build_masterci:

.github/workflows/repo-maintenance.yaml

@@ -11,7 +11,7 @@ env:
   PYTHONPATH: ${{ github.workspace }}
   GIT_CONFIG_COUNT: 1
   GIT_CONFIG_KEY_0: lfs.fetchexclude
-  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx
+  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
 
 jobs:
   package_updates:

.github/workflows/tests.yaml

@@ -22,7 +22,7 @@ env:
   PYTHONPATH: ${{ github.workspace }}
   GIT_CONFIG_COUNT: 1
   GIT_CONFIG_KEY_0: lfs.fetchexclude
-  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx
+  GIT_CONFIG_VALUE_0: openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
 
 jobs:
   build_release:

.gitignore

@@ -50,6 +50,7 @@ st[0-9A-Za-z][0-9A-Za-z][0-9A-Za-z][0-9A-Za-z][0-9A-Za-z][0-9A-Za-z]
 *.stats
 *.pkl
 *.pkl*
+!openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
 config.json
 compile_commands.json
 compare_runtime*.html

SConstruct

@@ -342,6 +342,8 @@ AddPostAction(BUILD_TARGETS or [Dir('.')], prune_cache_dir)
 def check_build_product_size(target, source, env):
   limit = 50 * 1024 * 1024  # GitHub max size
   for t in target:
+    if str(t).endswith('.pkl'):  # chunked during release packaging
+      continue
     if hasattr(t, 'isfile') and t.isfile() and (size := os.path.getsize(t.abspath)) > limit:
       raise SCons.Errors.UserError(f"{t} is {size / (1024 * 1024):.1f} MiB, exceeding the {limit / (1024 * 1024):.1f} MiB limit")
 if not GetOption('extras'):

openpilot/common/file_chunker.py

@@ -32,14 +32,6 @@ def chunk_file(path, targets):
   Path(manifest_path).write_text(str(len(chunk_paths)))
   os.remove(path)
 
-def get_existing_chunks(path):
-  if os.path.isfile(path):
-    return [path]
-  if os.path.isfile(manifest := get_manifest_path(path)):
-    num_chunks = int(Path(manifest).read_text().strip())
-    return _chunk_paths(path, num_chunks)
-  raise FileNotFoundError(path)
-
 class ChunkStream(io.RawIOBase):
   def __init__(self, paths):
     self._paths = iter(paths)
@@ -67,11 +59,11 @@ class ChunkStream(io.RawIOBase):
 
 def open_file_chunked(path):
   manifest_path = get_manifest_path(path)
-  if os.path.isfile(manifest_path):
+  if os.path.isfile(path):
+    paths = [path]
+  elif os.path.isfile(manifest_path):
     num_chunks = int(Path(manifest_path).read_text().strip())
     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))

openpilot/selfdrive/modeld/SConscript

@@ -1,18 +1,14 @@
 import glob
 import os
-import shutil
-import tempfile
 import time
 from SCons.Script import Action, Value
-from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks, open_file_chunked
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
-from openpilot.selfdrive.modeld.helpers import chestnut_present, modeld_pkl_path
+from openpilot.selfdrive.modeld.helpers import chestnut_present
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 
 
 Import('env', 'arch')
-chunker_file = File("#openpilot/common/file_chunker.py")
 lenv = env.Clone()
 lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath)
 
@@ -20,11 +16,6 @@ tinygrad_root = env.Dir("#").abspath
 tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
                   if 'pycache' not in x and os.path.isfile(os.path.join(tinygrad_root, x))]
 
-def estimate_pickle_max_size(onnx_size):
-  # QCOM programs for models with spatial recurrent features can approach 2x
-  # the ONNX size. Overestimating only adds an empty trailing chunk.
-  return 2.0 * onnx_size + 10 * 1024 * 1024
-
 camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]
 
 if arch == 'comma_arm64':
@@ -47,7 +38,7 @@ compiler = Dir('#tinygrad_repo/examples/openpilot').abspath
 # CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
 taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''
 
-def chestnut_action(command, pkl=None, chunks=()):
+def chestnut_action(command):
   def do_compile(target, source, env):
     from openpilot.system.hardware.chestnut.flash import link_up
     # chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars
@@ -56,47 +47,27 @@ def chestnut_action(command, pkl=None, chunks=()):
         break
       time.sleep(1)
     else:
-      print("Chestnut not ready, skipping big model build")
+      print("Chestnut not ready, skipping warp build")
       return
-    if ret := env.Execute(command):
-      return ret
-    if chunks:
-      chunk_file(pkl, chunks)
+    return env.Execute(command)
   return Action(do_compile, "  [CHESTNUT] $TARGET")
 
-def compile_model(onnx_path, pkl_path, flags, chestnut=False):
+def compile_model(onnx_path, pkl_path):
   onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
-  onnx_deps = get_existing_chunks(onnx_path)
-  cmd = (f'{flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
-         f'"{{onnx}}" "{target_pkl_path}" --device-input "*" --out-of-band --benchmark-runs 1')
-  def do_compile(target, source, env):
-    if os.path.isfile(onnx_path):
-      return env.Execute(cmd.format(onnx=onnx_path))
-    # TODO: Remove ONNX chunk reassembly once models are precompiled.
-    with tempfile.NamedTemporaryFile(dir=os.path.dirname(onnx_path), suffix='.onnx') as tmp, open_file_chunked(onnx_path) as src:
-      shutil.copyfileobj(src, tmp)
-      tmp.flush()
-      return env.Execute(cmd.format(onnx=tmp.name))
-  compile_action = Action(do_compile, "  [ONNX] $TARGET")
-  onnx_sizes_sum = sum(os.path.getsize(f) for f in onnx_deps)
-  chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
-  def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets):
-    chunk_file(pkl, chunks)
-  actions = chestnut_action(compile_action, target_pkl_path, chunk_targets) if chestnut else [compile_action, Action(do_chunk, "  [CHUNK] $TARGET")]
-  node = lenv.Command(
-    chunk_targets,
-    tinygrad_files + onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
-    actions,
+  cmd = (f'{tg_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
+         f'"{onnx_path}" "{target_pkl_path}" --device-input "*" --out-of-band --benchmark-runs 1')
+  lenv.Command(
+    target_pkl_path,
+    tinygrad_files + [onnx_path, Value(cmd)],
+    Action(cmd, "  [ONNX] $TARGET"),
   )
-  if chestnut:
-    lenv.SideEffect(chestnut_lock, node)
 
-compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl', tg_flags)
+compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl')
+compile_model('models/driving_supercombo.onnx', 'models/driving_tinygrad.pkl')
 
 model_w, model_h = MEDMODEL_INPUT_SIZE
 for chestnut in [False, True] if CHESTNUT else [False]:
   file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
-  compile_model(f'models/{file_prefix}driving_supercombo.onnx', modeld_pkl_path(chestnut), cmd_flags, chestnut)
   for cam_w, cam_h in camera_configs:
     warp_pkl_path = File(f"models/{file_prefix}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
     stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)

openpilot/selfdrive/modeld/helpers.py

@@ -35,4 +35,6 @@ def chestnut_present() -> bool:
   return False
 
 def chestnut_compiled() -> bool:
-  return Path(get_manifest_path(modeld_pkl_path(chestnut=True))).is_file()
+  path = modeld_pkl_path(chestnut=True)
+  return (path.is_file() or Path(get_manifest_path(path)).is_file()) and all(
+    (MODELS_DIR / f'big_driving_warp_{size}_tinygrad.pkl').is_file() for size in ('1344x760', '1928x1208'))

openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6fee5937923c74848df4a63f6239eb6331c6274dd4bdb7a5d6ec0388a8b543d5
-size 766018462

openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76cc0a9bc3af7a318889483dcbe126337f8d338f5abcbe664a8988c9b18b6639
+size 776634338

openpilot/selfdrive/test/chestnut.sh

@@ -3,7 +3,4 @@ set -e
 
 sudo python3 openpilot/system/hardware/chestnut/flash.py
 
-TARGET=openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl.chunkmanifest
-rm -f "$TARGET"
 SCONSFLAGS="-j4" ./openpilot/system/manager/build.py
-test -s "$TARGET"

openpilot/selfdrive/test/setup_device_ci.sh

@@ -64,9 +64,8 @@ pull_lfs() {
     return
   fi
 
-  # The big driving model is not used on these devices yet. Keep its pointer in
-  # the worktree, but don't download or copy the 1.8 GB LFS object.
-  LFS_EXCLUDE="openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx"
+  # Keep the precompiled big model as a pointer on devices without Chestnut.
+  LFS_EXCLUDE="openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl"
 
   git config --local lfs.fetchexclude "$LFS_EXCLUDE"
   git lfs pull --exclude="$LFS_EXCLUDE"

tools/release/build_release.sh

@@ -49,9 +49,6 @@ for policy in /sys/devices/system/cpu/cpufreq/policy*; do
 done
 
 scons
-if [ -n "$INCLUDE_BIG_MODEL" ]; then
-  test -f openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl.chunkmanifest
-fi
 
 if [ -z "$PANDA_DEBUG_BUILD" ]; then
   # release panda fw
@@ -61,6 +58,11 @@ else
   scons panda/
 fi
 
+find openpilot/selfdrive/modeld/models -name '*.pkl' -size +95M -exec ./openpilot/common/file_chunker.py {} \;
+if [ -n "$INCLUDE_BIG_MODEL" ]; then
+  test -f openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl.chunkmanifest
+fi
+
 # Ensure no submodules in release
 if test "$(git submodule--helper list | wc -l)" -gt "0"; then
   echo "submodules found:"

tools/release/build_stripped.sh

@@ -39,7 +39,7 @@ cd "$SOURCE_DIR"
 cd "$TARGET_DIR"
 rm -rf .git/modules/
 
-find openpilot/selfdrive/modeld/models -name '*.onnx' -size +95M -exec ./openpilot/common/file_chunker.py {} \;
+find openpilot/selfdrive/modeld/models -name '*.pkl' -size +95M -exec ./openpilot/common/file_chunker.py {} \;
 
 # include source commit hash and build date in commit
 GIT_HASH=$(git --git-dir="$SOURCE_DIR/.git" rev-parse HEAD)

tools/release/release_files.py

@@ -32,7 +32,7 @@ if __name__ == "__main__":
       continue
 
     rf = os.fsdecode(tracked_file)
-    if not os.getenv("INCLUDE_BIG_MODEL") and rf.startswith("openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx"):
+    if not os.getenv("INCLUDE_BIG_MODEL") and rf.startswith("openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl"):
       continue
     blacklisted = any(re.search(p, rf) for p in blacklist)
     whitelisted = any(re.search(p, rf) for p in whitelist)