Warm up external GPU model before activation

selfdrive/modeld/modeld.py

@@ -4,6 +4,7 @@ from openpilot.system.hardware import TICI
 os.environ['GMMU'] = '0'
 os.environ['DEV'] = 'QCOM' if TICI else 'LLVM'
 from tinygrad.tensor import Tensor
+import threading
 import time
 import pickle
 import numpy as np
@@ -63,6 +64,7 @@ def _model_smooth_seconds(params, key, default):
   value = params.get_float(key, return_default=True, default=default)
   return round(min(max(value, SMOOTH_SECONDS_STEP), 2.0) / SMOOTH_SECONDS_STEP) * SMOOTH_SECONDS_STEP
 MIN_LAT_CONTROL_SPEED = 0.3
+BIG_MODEL_TIMEOUT = 60
 
 
 def _get_param_str(params: Params, key: str, default: str = "") -> str:
@@ -294,6 +296,7 @@ class ModelState:
     self.off_policy_enabled = "off_policy" in self.policy_order
     self.off_policy_numpy_inputs = dict(self.numpy_inputs) if self.off_policy_enabled else {}
     self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
+    self.nonfinite_count = 0
     self.parser = Parser()
     self.aux_parser = Parser(ignore_missing=True)
     self.frame_buf_size = get_nv12_info(cam_w, cam_h)[3]
@@ -377,6 +380,46 @@ class ModelState:
     parsed.update(policy_results[primary_key])
     return parsed
 
+  def _reset_state(self) -> None:
+    if self.model_type == "supercombo":
+      self.input_queues, self.npy = make_supercombo_input_queues(
+        self.policy_input_shapes, self.frame_skip, self.QUEUE_DEV,
+      )
+    else:
+      vision_shapes = self.metadata["vision"]["input_shapes"]
+      self.input_queues, self.npy = make_split_input_queues(
+        vision_shapes, self.policy_input_shapes, self.frame_skip, self.QUEUE_DEV,
+      )
+
+    for value in self.numpy_inputs.values():
+      value.fill(0)
+    self.prev_desire.fill(0)
+    if self.prev_desired_curv_key is not None:
+      self.full_prev_desired_curv.fill(0)
+    self._blob_cache.clear()
+
+  def warmup(self) -> None:
+    dummy_frames = {
+      key: np.zeros(self.frame_buf_size, dtype=np.uint8)
+      for key in self.vision_input_names
+    }
+    # A host pointer is not a valid camera buffer for every warp backend. Match
+    # upstream and substitute realized device buffers for warmup only.
+    self._blob_cache.update({
+      (key, value.ctypes.data): Tensor.zeros(value.shape, dtype="uint8", device=self.WARP_DEV).realize()
+      for key, value in dummy_frames.items()
+    })
+    eye = np.eye(3, dtype=np.float32)
+    inputs = {self.desire_key: np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)}
+    for name, value in self.numpy_inputs.items():
+      if name in (self.desire_key, self.prev_desired_curv_key):
+        continue
+      shape = value.shape[1:] if value.ndim > 1 and value.shape[0] == 1 else value.shape
+      inputs[name] = np.zeros(shape, dtype=value.dtype)
+
+    self.run(dummy_frames, dict.fromkeys(self.vision_input_names, eye), inputs, False)
+    self._reset_state()
+
   def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
           inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None:
     frames: dict[str, Tensor] = {}
@@ -426,12 +469,16 @@ class ModelState:
       )
     outputs = [output.numpy().flatten() for output in output_tensors]
 
-    # USB GPU failures can produce NaNs/Infs instead of raising. Never publish
-    # one of those frames; the next frame can recover without affecting the
-    # native GPU/CPU model paths.
+    # A corrupted inference poisons recurrent state. Reset and retry a few
+    # times, then raise so modeld can fall back to the already-loaded CPU model.
     if self.uses_external_gpu and any(not np.isfinite(output).all() for output in outputs):
-      cloudlog.error("external GPU produced non-finite model output, dropping frame")
+      self.nonfinite_count += 1
+      cloudlog.error(f"external GPU produced non-finite model output, resetting state ({self.nonfinite_count})")
+      self._reset_state()
+      if self.nonfinite_count >= 5:
+        raise RuntimeError("external GPU produced non-finite output after state reset")
       return None
+    self.nonfinite_count = 0
 
     if self.model_type == "supercombo":
       model_output = outputs[0]
@@ -523,9 +570,40 @@ def main(demo=False):
 
   start_time = time.monotonic()
   cloudlog.warning("loading model")
+  model = None
+  small_model = None
   if external_gpu_requested:
-    wait_usbgpu_link()
-  model = _load_model_state(vipc_client_main.width, vipc_client_main.height, selected_model, external_gpu_requested, params)
+    big_model = None
+
+    def load_big_model() -> None:
+      nonlocal big_model
+      try:
+        wait_usbgpu_link()
+        candidate = ModelState(vipc_client_main.width, vipc_client_main.height, True)
+        if not candidate.uses_external_gpu:
+          raise RuntimeError("external GPU model resolved to the builtin model")
+        candidate.warmup()
+        big_model = candidate
+      except Exception:
+        cloudlog.exception("external GPU model load or warmup failed")
+
+    loader = threading.Thread(target=load_big_model, name="big_model_loader", daemon=True)
+    loader.start()
+    loader.join(BIG_MODEL_TIMEOUT)
+    if loader.is_alive():
+      cloudlog.error(f"external GPU model load timed out after {BIG_MODEL_TIMEOUT}s")
+    model = big_model
+
+    # Keep the native model ready so a GPU error never takes modeld down.
+    small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False)
+    if model is None:
+      model = small_model
+    else:
+      params.put("ModelVersion", model.policy_generation)
+      params.put("DrivingModelVersion", model.policy_generation)
+  else:
+    model = _load_model_state(vipc_client_main.width, vipc_client_main.height, selected_model, False, params)
+
   external_gpu_active = model.uses_external_gpu
   params.put_bool("UsbGpuCompiled", external_model_selected and file_chunked_exists(external_artifact))
   params.put_bool("UsbGpuActive", external_gpu_active)
@@ -680,7 +758,17 @@ def main(demo=False):
       inputs['lateral_control_params'] = lateral_control_params
 
     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 external_gpu_active or small_model is None:
+        raise
+      cloudlog.exception("external GPU model failed, falling back to builtin model")
+      params.put_bool("UsbGpuActive", False)
+      model = small_model
+      external_gpu_active = False
+      run_count = 0
+      model_output = None
     mt2 = time.perf_counter()
     model_execution_time = mt2 - mt1
 

selfdrive/modeld/tests/test_usbgpu_helpers.py

@@ -1,7 +1,9 @@
 import io
+from types import MethodType
 
 import numpy as np
 
+from openpilot.selfdrive.modeld import modeld
 from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob, tinygrad_dev_config
 from scripts import model_compiler
 
@@ -34,3 +36,58 @@ def test_external_gpu_probe_matches_upstream_retry_loop(monkeypatch):
   model_compiler.wait_for_external_gpu()
 
   assert calls == ["probe", ("sleep", 1), "probe", ("sleep", 1), "probe"]
+
+
+def test_external_gpu_warmup_runs_a_complete_frame_and_resets(monkeypatch):
+  class FakeTensor:
+    @staticmethod
+    def zeros(shape, **kwargs):
+      calls.append(("tensor", shape, kwargs))
+      return FakeTensor()
+
+    def realize(self):
+      return self
+
+  calls = []
+  state = modeld.ModelState.__new__(modeld.ModelState)
+  state.frame_buf_size = 32
+  state.vision_input_names = ["img", "big_img"]
+  state._blob_cache = {}
+  state._warp_dev = "QCOM"
+  state.desire_key = "desire"
+  state.prev_desired_curv_key = "prev_desired_curv"
+  state.numpy_inputs = {
+    "desire": np.zeros((1, 8), dtype=np.float32),
+    "traffic_convention": np.zeros((1, 2), dtype=np.float32),
+    "action_t": np.zeros((1, 2), dtype=np.float32),
+    "prev_desired_curv": np.zeros((1, 5, 1), dtype=np.float32),
+  }
+
+  def fake_run(self, bufs, transforms, inputs, prepare_only):
+    calls.append((
+      "run",
+      {key: value.shape for key, value in bufs.items()},
+      {key: value.shape for key, value in transforms.items()},
+      {key: value.shape for key, value in inputs.items()},
+      prepare_only,
+    ))
+    return {}
+
+  state.run = MethodType(fake_run, state)
+  state._reset_state = MethodType(lambda self: calls.append(("reset",)), state)
+  monkeypatch.setattr(modeld, "Tensor", FakeTensor)
+
+  state.warmup()
+
+  assert calls == [
+    ("tensor", (32,), {"dtype": "uint8", "device": "QCOM"}),
+    ("tensor", (32,), {"dtype": "uint8", "device": "QCOM"}),
+    (
+      "run",
+      {"img": (32,), "big_img": (32,)},
+      {"img": (3, 3), "big_img": (3, 3)},
+      {"desire": (8,), "traffic_convention": (2,), "action_t": (2,)},
+      False,
+    ),
+    ("reset",),
+  ]