Connor

selfdrive/modeld/modeld.py

@@ -99,6 +99,7 @@ def _should_publish_model_output(model_output, vipc_dropped_frames: int, externa
 MIN_LAT_CONTROL_SPEED = 0.3
 BIG_MODEL_LOAD_WAIT_TIMEOUT_MS = 30000
 BIG_MODEL_RUN_WAIT_TIMEOUT_MS = 3000
+EXTERNAL_GPU_TEST_POWER_LIMIT_W = 60
 EXTERNAL_GPU_POWER_READY_MV = 10000
 EXTERNAL_GPU_EGMP_READY_MV = 12500
 EXTERNAL_GPU_POWER_STABLE_SECONDS = 3.0
@@ -115,6 +116,12 @@ def _set_hcq_wait_timeout(timeout_ms: int) -> None:
   getenv.cache_clear()
 
 
+def _set_external_gpu_power_limit(watts: int) -> None:
+  os.environ["AM_POWER_LIMIT"] = str(watts)
+  from tinygrad.helpers import getenv
+  getenv.cache_clear()
+
+
 def _external_gpu_power_voltage(device_type: str, panda_states, peripheral_state) -> int | None:
   if device_type == "tici":
     voltage = int(peripheral_state.voltage)
@@ -1053,6 +1060,9 @@ def main(demo=False):
   external_artifact = MODELS_PATH / f"{big_model_id}_driving_tinygrad.pkl"
   external_artifact_ready = external_model_selected and file_chunked_exists(external_artifact)
   external_gpu_requested = usbgpu_present_now and (bool(big_model_id) or model_lab_ready)
+  if external_gpu_requested:
+    _set_external_gpu_power_limit(EXTERNAL_GPU_TEST_POWER_LIMIT_W)
+    cloudlog.warning(f"external GPU test power limit set to {EXTERNAL_GPU_TEST_POWER_LIMIT_W} W")
   params.put_bool("UsbGpuPresent", usbgpu_present_now)
   params.put_bool("UsbGpuCompiled", external_artifact_ready or model_lab_ready)
   params.put_bool("UsbGpuActive", False)

selfdrive/modeld/tests/test_usbgpu_helpers.py

@@ -1,4 +1,5 @@
 import io
+import os
 import struct
 from types import MethodType
 from types import SimpleNamespace
@@ -39,6 +40,17 @@ def test_external_gpu_uses_a_longer_load_watchdog():
   assert modeld.BIG_MODEL_RUN_WAIT_TIMEOUT_MS == 3000
 
 
+def test_external_gpu_test_power_limit(monkeypatch):
+  from tinygrad.helpers import getenv
+
+  monkeypatch.delenv("AM_POWER_LIMIT", raising=False)
+  getenv.cache_clear()
+  modeld._set_external_gpu_power_limit(modeld.EXTERNAL_GPU_TEST_POWER_LIMIT_W)
+
+  assert os.environ["AM_POWER_LIMIT"] == "60"
+  assert getenv("AM_POWER_LIMIT", 0.0) == 60.0
+
+
 def test_external_gpu_voltage_uses_hardware_specific_source():
   panda_type = modeld.log.PandaState.PandaType
   panda_states = [SimpleNamespace(pandaType=panda_type.dos, voltage=230)]