monitor chestnut USB in hardwared (#38741)

hardwared: monitor chestnut USB independently
(cherry picked from commit 7d5596d5c3527f274a3615181f52095909ffb8e9)

openpilot/cereal/log.capnp

@@ -2593,6 +2593,7 @@ struct Event {
     clocks @35 :Clocks;
     deviceState @6 :DeviceState;
     chestnutState @152 :ChestnutState;
+    chestnutGpuState @153 :ChestnutState;
     logMessage @18 :Text;
     errorLogMessage @85 :Text;
 

openpilot/cereal/services.py

@@ -26,6 +26,7 @@ _services: dict[str, tuple] = {
   "temperatureSensor": (True, 2., 200),
   "deviceState": (True, 2., 1),
   "chestnutState": (True, 10., 10),
+  "chestnutGpuState": (False, 10.),
   "touch": (True, 20., 1),
   "can": (True, 100., 2053, QueueSize.BIG),  # decimation gives ~3 msgs in a full segment
   "controlsState": (True, 100., 10, QueueSize.MEDIUM),

openpilot/selfdrive/modeld/modeld.py

@@ -4,7 +4,6 @@ import ctypes
 from functools import cached_property
 import os
 os.environ['GMMU'] = '0' # for chestnut fast loading, noop for qcom
-from tinygrad.tensor import Tensor
 from tinygrad.device import Device
 import usb1
 import struct
@@ -29,18 +28,17 @@ from openpilot.common.transformations.model import get_warp_matrix
 from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
 from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, should_stop, smooth_value, get_curvature_from_plan
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
-from openpilot.selfdrive.modeld.compile_modeld import make_input_queues, WARP_INPUTS, POLICY_INPUTS
+from openpilot.selfdrive.modeld.compile_modeld import make_input_queues, nv12_copy_size, MODELD_INPUTS
 from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_driving_model_data, fill_pose_msg, PublishState
 from openpilot.common.file_chunker import open_file_chunked
 from openpilot.common.hardware.usb import CHESTNUT_USB_IDS
 from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
-from openpilot.selfdrive.modeld.helpers import chestnut_present, chestnut_compiled, modeld_pkl_path, get_tg_input_devices, load_oob
+from openpilot.selfdrive.modeld.helpers import chestnut_present, chestnut_compiled, chestnut_ready, modeld_pkl_path, load_oob
 
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
 from openpilot.sunnypilot.selfdrive.controls.lib.relc import RoadEdgeLaneChangeController
 
-PROCESS_NAME = "openpilot.selfdrive.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 
 LAT_SMOOTH_SECONDS = 0.0
@@ -182,9 +180,9 @@ class ModelState(ModelStateBase):
 
   def __init__(self, cam_w: int, cam_h: int, chestnut: bool):
     ModelStateBase.__init__(self)
-    input_devices = get_tg_input_devices(PROCESS_NAME, chestnut)
-    self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV']
     jits = load_oob(open_file_chunked(modeld_pkl_path(chestnut)))
+    input_devices = jits['input_devices']
+    self.model_device = input_devices['model']
     metadata = jits['metadata']
     self.input_shapes = metadata['input_shapes']
     self.vision_input_names = [k for k in self.input_shapes if 'img' in k]
@@ -194,13 +192,11 @@ class ModelState(ModelStateBase):
     self.chestnut = chestnut
 
     self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-    self.input_queues, self.npy = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
-    self.full_frames: dict[str, Tensor] = {}
-    self._blob_cache: dict[tuple[str, int], Tensor] = {}
+    self.frame_copy_size = nv12_copy_size(*get_nv12_info(cam_w, cam_h)[:3])
+    self.input_queues, self.npy, self.frame_views = make_input_queues(
+      self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size)
     self.parser = Parser()
-    self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')}
-    self.run_policy = jits['run_policy']
-    self.warp = jits[(cam_w,cam_h)]
+    self.run_model = jits['run_model'][(cam_w,cam_h)]
 
   def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
     parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
@@ -208,14 +204,8 @@ class ModelState(ModelStateBase):
 
   def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
           inputs: dict[str, np.ndarray], after_enqueue: Callable[[], None] | None = None) -> dict[str, np.ndarray]:
-    for key in bufs.keys():
-      ptr = np.frombuffer(bufs[key].data, dtype=np.uint8).ctypes.data
-      yuv_size = self.frame_buf_params[key][3]
-      # There is a ringbuffer of imgs, just cache tensors pointing to all of them
-      cache_key = (key, ptr)
-      if cache_key not in self._blob_cache:
-        self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.WARP_DEV)
-      self.full_frames[key] = self._blob_cache[cache_key]
+    for key, buf in bufs.items():
+      np.copyto(self.frame_views[key], np.frombuffer(buf.data, dtype=np.uint8, count=self.frame_copy_size))
 
     # Model decides when action is completed, so desire input is just a pulse triggered on rising edge
     inputs['desire_pulse'][0] = 0
@@ -226,11 +216,7 @@ class ModelState(ModelStateBase):
     self.npy['tfm'][:,:] = transforms['img'][:,:]
     self.npy['big_tfm'][:,:] = transforms['big_img'][:,:]
 
-    warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'])
-
-    outs, = self.run_policy(
-      **{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped
-    )
+    outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS})
     if after_enqueue is not None:
       after_enqueue()
     model_output = outs.numpy()[0]
@@ -244,25 +230,37 @@ class ModelState(ModelStateBase):
     return outputs_dict
 
   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}
+    dummy_frames = {k: np.zeros(self.frame_copy_size, dtype=np.uint8) for k in self.vision_input_names}
     eye = np.eye(3, dtype=np.float32)
     dims = {'desire_pulse': ModelConstants.DESIRE_LEN, 'traffic_convention': 2, 'action_t': 2}
     self.run(dummy_frames, dict.fromkeys(self.vision_input_names, eye), {k: np.zeros(v, dtype=np.float32) for k, v in dims.items()})
-    self.input_queues, self.npy = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
+    self.input_queues, self.npy, self.frame_views = make_input_queues(
+      self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size)
     self.prev_desire[:] = 0
-    self.full_frames.clear()
-    self._blob_cache.clear()
 
 
 def main(demo=False):
   cloudlog.warning("modeld init")
 
-  CHESTNUT = chestnut_present() and chestnut_compiled()
+  chestnut_available = chestnut_present() and chestnut_compiled()
+  CHESTNUT = False
+  if chestnut_available:
+    poller = messaging.Poller()
+    sock = messaging.sub_sock("chestnutState", poller=poller, conflate=True)
+    deadline = time.monotonic() + 4. / SERVICE_LIST['deviceState'].frequency
+    while not CHESTNUT and (remaining := deadline - time.monotonic()) > 0.:
+      if not poller.poll(round(remaining * 1000)):
+        break
+      msg = messaging.recv_one_or_none(sock)
+      CHESTNUT = msg is not None and msg.valid and chestnut_ready(msg.chestnutState)
   if CHESTNUT:
     os.environ['HCQDEV_WAIT_TIMEOUT_MS'] = '3000'
   params = Params()
   params.put_bool("ChestnutLoading", CHESTNUT)
-  params.remove("ChestnutActive")
+  if chestnut_available and not CHESTNUT:
+    params.put_bool("ChestnutActive", False)
+  else:
+    params.remove("ChestnutActive")
 
   config_realtime_process(7, 54)
 
@@ -306,7 +304,11 @@ def main(demo=False):
     loader.start()
     loader.join(BIG_MODEL_TIMEOUT)
     model = big_model
+    if model is None:
+      params.put_bool("ChestnutModelError", True)
     params.put_bool("ChestnutActive", model is not None)
+    if model is not None:
+      params.remove("ChestnutModelError")
 
   small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False) if model is None or CHESTNUT else None
   if model is None:
@@ -440,6 +442,7 @@ def main(demo=False):
         raise
       # fallback to small model
       cloudlog.exception("big model failed, fall back to small")
+      params.put_bool("ChestnutModelError", True)
       params.put_bool("ChestnutActive", False)
       assert small_model is not None
       model = small_model
@@ -466,12 +469,11 @@ def main(demo=False):
       l_lane_change_prob = desire_state[log.Desire.laneChangeLeft]
       r_lane_change_prob = desire_state[log.Desire.laneChangeRight]
       lane_change_prob = l_lane_change_prob + r_lane_change_prob
-      DH.update(sm['carState'], sm['carControl'].latActive, lane_change_prob)
-      modelv2_send.modelV2.meta.laneChangeState = DH.lane_change_state
-      modelv2_send.modelV2.meta.laneChangeDirection = DH.lane_change_direction
-
       mdv2sp_send = messaging.new_message('modelDataV2SP')
       left_edge, right_edge = RELC.update_and_fill(modelv2_send.modelV2, mdv2sp_send.modelDataV2SP, v_ego)
+      DH.update(sm['carState'], sm['carControl'].latActive, lane_change_prob, left_edge, right_edge)
+      modelv2_send.modelV2.meta.laneChangeState = DH.lane_change_state
+      modelv2_send.modelV2.meta.laneChangeDirection = DH.lane_change_direction
       mdv2sp_send.modelDataV2SP.laneTurnDirection = DH.lane_turn_direction
 
       fill_driving_model_data(drivingdata_send, modelv2_send)

openpilot/system/hardware/chestnut/monitoring.py

@@ -0,0 +1,124 @@
+import struct
+from contextlib import suppress
+
+import usb1
+
+import openpilot.cereal.messaging as messaging
+from openpilot.cereal.services import SERVICE_LIST
+from openpilot.common.hardware.usb import CHESTNUT_USB_IDS
+
+
+USB_TIMEOUT_MS = 100
+PCIE_LTSSM_ADDRESS = 0xB450
+
+
+class ChestnutUsb:
+  def __init__(self):
+    self.context: usb1.USBContext | None = None
+    self.handle = None
+
+  def close(self) -> None:
+    handle, context = self.handle, self.context
+    self.handle = None
+    self.context = None
+    with suppress(Exception):
+      if handle is not None:
+        handle.close()
+    with suppress(Exception):
+      if context is not None:
+        context.close()
+
+  def connect(self) -> bool:
+    if self.handle is not None:
+      return True
+
+    context = usb1.USBContext()
+    for vendor_id, product_id in CHESTNUT_USB_IDS:
+      handle = context.openByVendorIDAndProductID(vendor_id, product_id, skip_on_error=True)
+      if handle is not None:
+        self.context = context
+        self.handle = handle
+        return True
+    context.close()
+    return False
+
+  def _read(self, request: int, value: int, length: int) -> bytes:
+    if self.handle is None:
+      raise usb1.USBErrorNoDevice
+    raw = bytes(self.handle.controlRead(0xC0, request, value, 0, length, timeout=USB_TIMEOUT_MS))
+    if len(raw) != length:
+      raise ValueError(f"short chestnut USB response: {len(raw)}/{length}")
+    return raw
+
+  def read_ina(self) -> tuple[int, int, bool]:
+    return struct.unpack('<Hh?', self._read(0xC0, 0, 5))
+
+  def read_pcie_ltssm(self) -> int:
+    return self._read(0xE4, PCIE_LTSSM_ADDRESS, 1)[0]
+
+
+class ChestnutMonitoring:
+  def __init__(self, usb: ChestnutUsb | None = None):
+    self.usb = usb or ChestnutUsb()
+    self.gpu_state = None
+    self.seen = False
+    self.enabled = False
+    self.usb_failed = False
+
+  def set_enabled(self, enabled: bool) -> None:
+    if self.enabled == enabled:
+      return
+    self.enabled = enabled
+    self.usb.close()
+    self.usb_failed = False
+
+  def retry(self) -> None:
+    self.usb_failed = False
+
+  def model_alive(self, sm: messaging.SubMaster, now: float) -> bool:
+    modeld = next((p for p in sm['managerState'].processes if p.name == 'modeld'), None)
+    if modeld is not None and modeld.shouldBeRunning and not modeld.running:
+      return False
+    recv_time = sm.recv_time['chestnutGpuState']
+    return recv_time > 0. and now - recv_time < 10. / SERVICE_LIST['chestnutGpuState'].frequency
+
+  def update_gpu_state(self, sm: messaging.SubMaster, now: float) -> None:
+    if sm.updated['chestnutGpuState']:
+      self.gpu_state = sm['chestnutGpuState'] if sm.valid['chestnutGpuState'] else None
+    elif not self.model_alive(sm, now):
+      self.gpu_state = None
+
+  def update(self, sm: messaging.SubMaster, now: float, model_loading: bool = False):
+    self.update_gpu_state(sm, now)
+    return self.build_message(model_loading)
+
+  def build_message(self, model_loading: bool = False):
+    if not self.enabled:
+      return None
+
+    msg = messaging.new_message('chestnutState')
+    if self.gpu_state is not None:
+      msg.chestnutState = self.gpu_state
+    state = msg.chestnutState
+
+    if self.usb_failed:
+      return msg if self.seen else None
+
+    try:
+      if not self.usb.connect():
+        self.usb_failed = True
+        return msg if self.seen else None
+
+      self.seen = True
+      voltage, current, fault = self.usb.read_ina()
+      pcie_ltssm = self.usb.read_pcie_ltssm()
+      state.supplyVoltage = voltage
+      state.supplyCurrent = current
+      state.supplyFault = fault
+      state.pcieLtssm = pcie_ltssm
+      msg.valid = True
+    except Exception as e:
+      if not model_loading or not isinstance(e, usb1.USBErrorTimeout):
+        self.usb.close()
+        self.usb_failed = True
+    return msg

openpilot/system/hardware/chestnut/status.py

@@ -21,14 +21,13 @@ class ChestnutStatus:
     self.power_lost = False
     self.power_restored = False
     self.link_failures = 0
-    self.model_loading_seen = False
     self.model_attempted = False
     self.overheated = False
     self.usb_seen = False
     self.usb_failed = False
 
   def update(self, offroad: bool, branch: str, usb_state: list[dict], firmware_failed: bool,
-             model_loading: bool, model_active: bool | None, state, set_alert) -> None:
+             model_active: bool | None, state, usb_failed: bool, set_alert) -> None:
     detected = [d for d in usb_state if is_chestnut_usb_id(d["vendorId"], d["productId"], include_bootloader=True)]
     devices = [d for d in detected if is_chestnut_usb_id(d["vendorId"], d["productId"])]
     firmware_ok = len(devices) == 1 and devices[0]["product"] == CHESTNUT_USB_PRODUCT
@@ -40,16 +39,15 @@ class ChestnutStatus:
       self.power_lost = False
       self.power_restored = False
       self.link_failures = 0
-      self.model_loading_seen = False
       self.model_attempted = False
       self.usb_seen = firmware_ok
       self.usb_failed = False
 
-    self.model_loading_seen |= model_loading
-    self.model_attempted |= self.model_loading_seen and not model_loading and model_active is not None
+    self.model_attempted |= model_active is not None
 
-    if not offroad and self.usb_seen and not firmware_ok:
-      self.usb_failed = True
+    if not offroad:
+      self.usb_seen |= firmware_ok
+    self.usb_failed = not offroad and self.usb_seen and (not firmware_ok or usb_failed)
 
     if not offroad and state is not None:
       powered = state.supplyVoltage >= CHESTNUT_POWERED_VOLTAGE
@@ -80,17 +78,28 @@ class ChestnutStatus:
     release = branch in CHESTNUT_RELEASE_BRANCHES
     missing = self.usb_failed or (offroad and release and time.monotonic() - self.started > 10. and len(detected) != 1)
     slow_usb = offroad and len(devices) == 1 and devices[0]["speedMbps"] < 5000
-    set_alert("Offroad_ChestnutBranch", not release and len(devices) == 1)
-    set_alert("Offroad_ChestnutNotDetected", missing)
-    set_alert("Offroad_ChestnutOverheated", self.overheated, f"{state.tempC:.0f} °C" if state is not None else None)
-    set_alert("Offroad_ChestnutUsbSlow", slow_usb, f"{devices[0]['speedMbps']} Mbps" if slow_usb else None)
+    update_failed = offroad and firmware_failed
+    compiled = firmware_ok and chestnut_compiled()
+    uncompiled = offroad and firmware_ok and not compiled
+
     if self.power_lost:
       pcie_alert = ("Chestnut power restored. 12V is stable again, cycle ignition." if self.power_restored else
                     "Chestnut power disconnected. Check 12V connection, then cycle ignition." if self.power_unavailable else
                     "Chestnut power lost. Possibly caused by an engine-crank voltage drop. Check 12V connection, then cycle ignition.")
     else:
       pcie_alert = "Chestnut GPU unavailable. PCIe link is not up. Check the GPU is securely seated."
-    set_alert("Offroad_ChestnutPcieUnavailable", self.pcie_failed, pcie_alert)
-    set_alert("Offroad_ChestnutUncompiled", offroad and firmware_ok and not chestnut_compiled())
-    set_alert("Offroad_ChestnutUpdateFailed", offroad and firmware_failed)
+
+    alerts = (
+      ("Offroad_ChestnutNotDetected", missing, None),
+      ("Offroad_ChestnutUpdateFailed", update_failed, None),
+      ("Offroad_ChestnutUncompiled", uncompiled, None),
+      ("Offroad_ChestnutPcieUnavailable", self.pcie_failed, pcie_alert),
+      ("Offroad_ChestnutOverheated", self.overheated, f"{state.tempC:.0f} °C" if state is not None else None),
+      ("Offroad_ChestnutUsbSlow", slow_usb, f"{devices[0]['speedMbps']} Mbps" if slow_usb else None),
+    )
+    active_alert = next((name for name, active, _ in alerts if active), None)
+
+    set_alert("Offroad_ChestnutBranch", not release and len(devices) == 1 and not missing)
+    for name, _, extra_text in alerts:
+      set_alert(name, name == active_alert, extra_text)
     self.offroad = offroad