129 lines
3.6 KiB
Python
129 lines
3.6 KiB
Python
"""
|
|
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
|