IQ.Pilot Release Commit @ b6534c0

iqpilot/selfdrive/iqmodeld/tests/test_fused_runner_guards.py

@@ -0,0 +1,128 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+from __future__ import annotations
+
+import pickle
+from dataclasses import dataclass
+
+import numpy as np
+import pytest
+
+from iqpilot.selfdrive.iqmodeld.models.runners import model_runner as model_runner_mod
+from iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelType
+from iqpilot.selfdrive.iqmodeld.models.runners.tinygrad import fused_runner as fused_mod
+
+
+class _View:
+  def __init__(self, shape):
+    self.shape = shape
+
+
+class _Captured:
+  def __init__(self, expected_names, expected_input_info):
+    self.expected_names = expected_names
+    self.expected_input_info = expected_input_info
+
+
+class _FakeJit:
+  def __init__(self, expected_names, expected_input_info):
+    self.captured = _Captured(expected_names, expected_input_info)
+
+  def __call__(self, **kwargs):
+    raise AssertionError("policy jit should not run in this test")
+
+
+class _FakeTensor:
+  def __init__(self, arr, device=None):
+    self.shape = tuple(np.asarray(arr).shape)
+
+  def contiguous(self):
+    return self
+
+  def realize(self):
+    return self
+
+
+class _FakeDevice:
+  DEFAULT = "FAKE"
+
+
+@dataclass
+class _Type:
+  raw: int
+
+
+@dataclass
+class _Artifact:
+  fileName: str
+
+
+class _Model:
+  def __init__(self, file_name):
+    self.type = _Type(ModelType.vision)
+    self.artifact = _Artifact(file_name)
+    self.metadata = None
+
+
+class _Bundle:
+  def __init__(self, file_name):
+    self.models = [_Model(file_name)]
+    self.is20hz = True
+
+
+POLICY_INPUTS = ["action_t", "big_img", "desire", "desire_q", "feat_q", "img", "traffic_convention"]
+POLICY_SHAPES = {
+  "action_t": (1, 2), "big_img": (1, 12, 128, 256), "desire": (1, 8), "desire_q": (1, 100, 8),
+  "feat_q": (1, 99, 512), "img": (1, 12, 128, 256), "traffic_convention": (1, 2),
+}
+
+
+def _write_fused_pkl(path, policy_inputs):
+  info = [(_View(POLICY_SHAPES[n]), (), None, "NPY") for n in policy_inputs]
+  role_meta = {
+    "input_shapes": {"desire_pulse": (1, 100, 8), "traffic_convention": (1, 2), "features_buffer": (1, 99, 512)},
+    "output_slices": {},
+  }
+  blob = {
+    "metadata": {
+      "vision": {"input_shapes": {"img": (1, 12, 128, 256), "big_img": (1, 12, 128, 256)}, "output_slices": {}},
+      "on_policy": role_meta,
+      "off_policy": role_meta,
+    },
+    "run_policy": _FakeJit(policy_inputs, info),
+    "frame_skip": 4,
+    (1928, 1208): _FakeJit(["frame"], [(_View((1,)), (), None, "NPY")]),
+  }
+  with open(path, "wb") as f:
+    pickle.dump(blob, f)
+
+
+@pytest.fixture
+def fused_runner(tmp_path, monkeypatch):
+  def _build(policy_inputs):
+    name = "driving_fused_test.pkl"
+    _write_fused_pkl(tmp_path / name, policy_inputs)
+    monkeypatch.setattr(model_runner_mod, "_fetch_bundle", lambda params=None: _Bundle(name))
+    monkeypatch.setattr(fused_mod, "CUSTOM_MODEL_PATH", str(tmp_path))
+    monkeypatch.setattr(fused_mod, "_tinygrad_imports", lambda: (_FakeTensor, _FakeDevice))
+    return fused_mod.TinygradFusedRunner()
+  return _build
+
+
+def test_action_t_allocated_when_only_the_jit_declares_it(fused_runner):
+  runner = fused_runner(POLICY_INPUTS)
+  assert "action_t" not in runner._on_meta["input_shapes"]
+
+  runner._ensure_queues(1928, 1208)
+
+  assert runner._npy_buffers["action_t"].shape == POLICY_SHAPES["action_t"]
+  assert runner._npy_buffers["traffic_convention"].shape == POLICY_SHAPES["traffic_convention"]
+
+
+def test_action_t_absent_when_the_jit_does_not_take_it(fused_runner):
+  runner = fused_runner([n for n in POLICY_INPUTS if n != "action_t"])
+
+  runner._ensure_queues(1928, 1208)
+
+  assert "action_t" not in runner._npy_buffers