modeld_v2: big to small model fallback

openpilot/cereal/custom.capnp

@@ -192,6 +192,7 @@ struct ModelManagerSP @0xaedffd8f31e7b55d {
     ref @9 :Text;
     minimumSelectorVersion @10 :UInt32;
     overrides @11 :List(Override);
+    isBig @12 :Bool;
   }
 }
 

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -8,22 +8,22 @@ See the LICENSE.md file in the root directory for more details.
 
 import os
 os.environ['GMMU'] = '0'
+import numpy as np
+import threading
+import time
+from setproctitle import setproctitle
+from tinygrad.tensor import Tensor
+
+import openpilot.cereal.messaging as messaging
 from openpilot.common.hardware import COMMA_HARDWARE
 from openpilot.selfdrive.modeld.helpers import chestnut_present, load_oob
-import time
-import numpy as np
-import openpilot.cereal.messaging as messaging
 from openpilot.cereal import log
 from opendbc.car.structs import car
 from openpilot.cereal.services import SERVICE_LIST
-from setproctitle import setproctitle
 from openpilot.cereal.messaging import PubMaster, SubMaster
 from openpilot.cereal.visionipc import VisionStreamType
 from msgq.visionipc import VisionIpcClient, VisionBuf
 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.swaglog import cloudlog
 from openpilot.common.params import Params
@@ -42,13 +42,13 @@ from openpilot.sunnypilot.modeld_v2.constants import Plan
 from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants
 from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper
 from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues, make_supercombo_input_queues, WARP_INPUTS, POLICY_INPUTS
-
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
 from openpilot.sunnypilot.models.helpers import get_active_bundle
 from openpilot.sunnypilot.selfdrive.controls.lib.relc import RoadEdgeLaneChangeController
 
 PROCESS_NAME = "openpilot.selfdrive.modeld.modeld_tinygrad"
+BIG_MODEL_TIMEOUT = 60
 
 
 def _pkl_exists(path):
@@ -59,15 +59,13 @@ def _pkl_exists(path):
 def _find_driving_pkl(bundle):
   if (override := os.environ.get('COMBINED_MODEL_PKL')) and _pkl_exists(override):
     return override
-  if bundle is None or not bundle.models:
-    return None
-  from openpilot.common.hardware.hw import Paths
-  model_root = Paths.model_root()
-
-  pkl_name = bundle.models[0].artifact.fileName
-  pkl_path = os.path.join(model_root, pkl_name)
-  if _pkl_exists(pkl_path):
-    return pkl_path
+  if bundle and bundle.models:
+    from openpilot.common.hardware.hw import Paths
+    for model in bundle.models:
+      if (file_name := getattr(model.artifact, 'fileName', None)):
+        if _pkl_exists(path := os.path.join(Paths.model_root(), file_name)):
+          return path
+  return None
 
 
 class FrameMeta:
@@ -92,6 +90,8 @@ class ModelState(ModelStateBase):
       model_bundle = None
     else:
       model_bundle = get_active_bundle(chestnut=chestnut)
+    if model_bundle and chestnut != model_bundle.isBig:
+      model_bundle = None
     self.generation = model_bundle.generation if model_bundle is not None else None
     overrides = {override.key: override.value for override in model_bundle.overrides} if model_bundle else {}
 
@@ -102,7 +102,7 @@ class ModelState(ModelStateBase):
     self.chestnut = chestnut
 
     pkl_path = _find_driving_pkl(model_bundle)
-    assert pkl_path is not None, "No driving pkl found — all models must be compiled with compile_modeld.py"
+    assert pkl_path is not None, f"No driving pkl found for {'chestnut' if chestnut else 'small model'} — all models must be compiled with compile_modeld.py"
     self._init_combined(pkl_path, cam_w, cam_h, model_bundle)
 
   def _init_combined(self, pkl_path, cam_w, cam_h, bundle):
@@ -185,9 +185,6 @@ class ModelState(ModelStateBase):
     else:
       self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=frame_tensor, big_frame=big_frame_tensor)
 
-    if self.chestnut:
-      self.warmup()
-
   def warmup(self) -> None:
     dummy_frames = {k: np.zeros(self.frame_buf_params[k][3], dtype=np.uint8) for k in self._vision_input_names}
     transforms = {k: np.eye(3, dtype=np.float32) for k in [self._road_key, self._wide_key] if k}
@@ -363,21 +360,26 @@ def main(demo=False):
 
   model = None
   if CHESTNUT:
-    import threading
-    def load():
-      nonlocal model
-      model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=True)
-    t = threading.Thread(target=load, daemon=True)
-    t.start()
-    t.join(60)
-    if model is None:
-      params.put_bool("ChestnutActive", False)
-      raise RuntimeError("chestnut model load failed or timed out (60s)")
-    params.put_bool("ChestnutActive", True)
-  else:
-    model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False)
+    big_model = None
+    def load_big():
+      nonlocal big_model
+      try:
+        m = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=True)
+        m.warmup()
+        big_model = m
+      except Exception:
+        cloudlog.exception("chestnut load failed")
+    loader = threading.Thread(target=load_big, daemon=True)
+    loader.start()
+    loader.join(BIG_MODEL_TIMEOUT)
+    model = big_model
+    params.put_bool("ChestnutActive", model is not None)
 
+  small_model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False) if model is None or CHESTNUT else None
+  if model is None:
+    model = small_model
   params.put_bool("ChestnutLoading", False)
+  assert model is not None
   cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
 
   # messaging
@@ -386,7 +388,7 @@ def main(demo=False):
   sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState", "carControl", "lateralDelay"])
 
   publish_state = PublishState()
-  chestnut_state = ChestnutState(pm, CHESTNUT) if CHESTNUT else None
+  chestnut_state = ChestnutState(pm, model.chestnut) if CHESTNUT else None
 
   # setup filter to track dropped frames
   frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ)
@@ -509,7 +511,19 @@ def main(demo=False):
       inputs['action_t'] = np.array([lat_action_t, long_action_t], dtype=np.float32)
 
     mt1 = time.perf_counter()
-    model_output = model.run(bufs, transforms, inputs, prepare_only)
+    try:
+      model_output = model.run(bufs, transforms, inputs, prepare_only)
+    except Exception:
+      if not params.get_bool("ChestnutActive"):
+        raise
+      cloudlog.exception("chestnut failed, falling back to small")
+      params.put_bool("ChestnutActive", False)
+      assert small_model is not None
+      model = small_model
+      if chestnut_state is not None:
+        chestnut_state.big = False
+      run_count = 0
+      model_output = None
     mt2 = time.perf_counter()
     model_execution_time = mt2 - mt1
 

openpilot/sunnypilot/modeld_v2/tests/helpers.py

@@ -37,10 +37,11 @@ class DummyModel:
 
 
 class DummyBundle:
-  def __init__(self, is_20hz=False, models=None, generation=10):
+  def __init__(self, is_20hz=False, models=None, generation=10, is_big=False):
     self.overrides = [DummyOverride('lat', '.1'), DummyOverride('long', '.3')]
     self.generation = generation
     self.is20hz = is_20hz
+    self.isBig = is_big
     self.models = models or []
 
 

openpilot/sunnypilot/modeld_v2/tests/test_fallback.py

@@ -0,0 +1,127 @@
+import io
+import requests
+from unittest import mock
+
+from openpilot.common.file_chunker import get_chunk_name
+from openpilot.common.hardware import hw
+from openpilot.common.test import OpenpilotTestCase
+from openpilot.selfdrive.modeld.helpers import dump_oob
+import openpilot.sunnypilot.modeld_v2.modeld as modeld_module
+from openpilot.sunnypilot.modeld_v2.tests import helpers as tests_helpers
+from openpilot.sunnypilot.modeld_v2.tests.helpers import DummyModel, DummyBundle, CAM_W, CAM_H
+from openpilot.sunnypilot.models.fetcher import ModelParser, ModelFetcher
+
+tmp_path = tests_helpers.tmp_path
+
+
+class TestFallback(OpenpilotTestCase):
+  def test_find_dual_model_in_bundle(self, tmp_path, monkeypatch):
+    lebowski_file = 'driving_lebowski.pkl'
+    tsfdo_file = 'driving_tsfdo.pkl'
+    (tmp_path / lebowski_file).write_bytes(b'fkasdjfkljf')
+    (tmp_path / tsfdo_file).write_bytes(b'dskfajklsdjlsfka')
+
+    monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
+    big_bundle = DummyBundle(models=[DummyModel('supercombo', lebowski_file)], is_big=True)
+    small_bundle = DummyBundle(models=[DummyModel('supercombo', tsfdo_file)], is_big=False)
+    big_pkl = modeld_module._find_driving_pkl(big_bundle)
+    small_pkl = modeld_module._find_driving_pkl(small_bundle)
+
+    assert big_pkl is not None and lebowski_file in big_pkl
+    assert small_pkl is not None and tsfdo_file in small_pkl
+
+  def test_download_models_and_init_modelstate_fallback(self, tmp_path, monkeypatch):
+    monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
+    big_json = requests.get(ModelFetcher.MODEL_URL_CHESTNUT).json()
+    big_bundle = ModelParser.parse_models(big_json)[-1]
+    small_json = requests.get(ModelFetcher.MODEL_URL).json()
+    small_bundle = ModelParser.parse_models(small_json)[-1]
+
+    buf = io.BytesIO()
+    dump_oob(tests_helpers.make_pkl_data(tests_helpers.ARCHETYPES['supercombo_non20hz']), buf)
+    oob_bytes = buf.getvalue()
+
+    for bundle in (big_bundle, small_bundle):
+      artifact = bundle.models[0].artifact
+      for i in range(len(artifact.chunks)):
+        (tmp_path / get_chunk_name(artifact.fileName, i, len(artifact.chunks))).write_bytes(oob_bytes if i == 0 else b"")
+
+    monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None, *, chestnut=None: small_bundle)
+    assert modeld_module.ModelState(CAM_W, CAM_H, chestnut=False).chestnut is False
+
+    monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None, *, chestnut=None: big_bundle)
+    try:
+      assert modeld_module.ModelState(CAM_W, CAM_H, chestnut=True).chestnut is True
+    except Exception as e:
+      assert "AMD" in str(e) or "device" in str(e).lower()
+
+  def test_runtime_fallback_from_big_to_small(self):
+    params = mock.MagicMock()
+    params.get_bool.return_value = True
+
+    big = mock.MagicMock(chestnut=True)
+    big.run.side_effect = RuntimeError("eGPU error")
+
+    small = mock.MagicMock(chestnut=False)
+    small.run.return_value = {"plan": []}
+
+    chestnut = mock.MagicMock(big=True)
+    model, run_count = big, 50
+
+    try:
+      model.run()
+    except Exception:
+      if not params.get_bool("chestnutActive"):
+        raise
+      params.put_bool("chestnutActive", False)
+      model = small
+      chestnut.big = False
+      run_count = 0
+
+    assert model is small
+    assert not chestnut.big
+    assert run_count == 0
+    params.put_bool.assert_called_with("chestnutActive", False)
+
+    model.run()
+    small.run.assert_called_once()
+
+  def test_runtime_stays_on_big_model_if_no_errors(self):
+    params = mock.MagicMock()
+    params.get_bool.return_value = True
+
+    big = mock.MagicMock(chestnut=True)
+    big.run.return_value = {"plan": [1, 2, 3]}
+
+    small = mock.MagicMock(chestnut=False)
+    chestnut = mock.MagicMock(big=True)
+    model, run_count = big, 50
+    try:
+      model.run()
+    except Exception:
+      if not params.get_bool("ChestnutActive"):
+        raise
+      params.put_bool("ChestnutActive", False)
+      model = small
+      chestnut.big = False
+      run_count = 0
+
+    assert model is big
+    assert chestnut.big is True
+    assert run_count == 50
+    params.put_bool.assert_not_called()
+    small.run.assert_not_called()
+
+  def test_runtime_exception_on_small_model_raises(self):
+    params = mock.MagicMock()
+    params.get_bool.return_value = False
+
+    model = mock.MagicMock(chestnut=False)
+    model.run.side_effect = RuntimeError("CPU error")
+
+    with self.assertRaises(RuntimeError):
+      try:
+        model.run()
+      except Exception:
+        if not params.get_bool("ChestnutActive"):
+          raise

openpilot/sunnypilot/models/fetcher.py

@@ -87,6 +87,7 @@ class ModelParser:
     model_bundle.environment = bundle["environment"]
     model_bundle.runner = bundle.get("runner", custom.ModelManagerSP.Runner.snpe)
     model_bundle.is20hz = bundle.get("is_20hz", False)
+    model_bundle.isBig = bundle.get("is_big", False)
     model_bundle.minimumSelectorVersion = int(bundle["minimum_selector_version"])
     model_bundle.overrides = ModelParser._parse_overrides(bundle.get("overrides", {}))
     model_bundle.ref = bundle.get("ref")