232 lines
8.4 KiB
Python
232 lines
8.4 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
|
|
"""
|
|
import os
|
|
import pickle as _pk
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import numpy as np
|
|
from iqpilot.cereal import custom
|
|
from iqpilot.common.swaglog import cloudlog
|
|
from iqpilot.system.hardware import TICI
|
|
from iqpilot.system.hardware.hw import Paths as _hw_paths
|
|
from iqpilot.selfdrive.iqmodeld.models.helpers import get_active_bundle as _fetch_bundle
|
|
from iqpilot.selfdrive.iqmodeld.models.combined_artifact import has_combined_split_artifact
|
|
|
|
# ---- runtime type surface (native OpenCL/frame handles resolve to Any off-device) ----
|
|
if TYPE_CHECKING:
|
|
from iqpilot.selfdrive.iqmodeld.native.iqmodel_pyx import GpuMemorySlot, RoadProjector
|
|
else:
|
|
def _resolve_native_types() -> tuple[Any, Any]:
|
|
try:
|
|
from iqpilot.selfdrive.iqmodeld.native.iqmodel_pyx import GpuMemorySlot as iq_clmem
|
|
from iqpilot.selfdrive.iqmodeld.native.iqmodel_pyx import RoadProjector as iq_frame
|
|
return iq_clmem, iq_frame
|
|
except (ModuleNotFoundError, ImportError):
|
|
return Any, Any
|
|
|
|
GpuMemorySlot, RoadProjector = _resolve_native_types()
|
|
|
|
NumpyDict = dict[str, np.ndarray]
|
|
ShapeDict = dict[str, tuple[int, ...]]
|
|
SliceDict = dict[str, slice]
|
|
CLMemDict = dict[str, GpuMemorySlot]
|
|
FrameDict = dict[str, RoadProjector]
|
|
|
|
ModelType = custom.IQModelManager.Model.Type
|
|
Model = custom.IQModelManager.Model
|
|
|
|
SEND_RAW_PRED = os.getenv("SEND_RAW_PRED")
|
|
CUSTOM_MODEL_PATH = _hw_paths.model_root()
|
|
|
|
_META_FIELDS = ("input_shapes", "output_slices")
|
|
|
|
USBGPU = "USBGPU" in os.environ
|
|
|
|
|
|
def _configure_accelerator():
|
|
"""Point tinygrad at the right backend. Must run before tinygrad is imported,
|
|
which is why it fires at module import."""
|
|
backend, extra = ("QCOM" if TICI else "CPU"), {}
|
|
if USBGPU:
|
|
backend, extra = "AMD", {"AMD_IFACE": "USB"}
|
|
elif TICI:
|
|
extra = {"QCOM_PRIORITY": "8"}
|
|
os.environ["DEV"] = backend
|
|
os.environ.update(extra)
|
|
|
|
|
|
_configure_accelerator()
|
|
|
|
|
|
# real metadata pkls are a few KB; anything bigger is a model artifact wrongly
|
|
# referenced as metadata (pre-fix manifests self-referenced the artifact), and
|
|
# unpickling it here double-loads the model onto the GPU
|
|
_META_MAX_BYTES = 1 << 20
|
|
|
|
|
|
def load_artifact_metadata(metadata_filename):
|
|
"""Read one artifact's metadata pkl: (input shapes, output slices)."""
|
|
try:
|
|
path = os.path.join(CUSTOM_MODEL_PATH, metadata_filename)
|
|
if os.path.getsize(path) > _META_MAX_BYTES:
|
|
cloudlog.error(f"metadata pkl {metadata_filename} is artifact-sized, refusing to unpickle it")
|
|
return tuple({} for _ in _META_FIELDS)
|
|
with open(path, 'rb') as fh:
|
|
blob = _pk.load(fh)
|
|
return tuple(blob.get(field, {}) for field in _META_FIELDS)
|
|
except Exception:
|
|
cloudlog.exception(f"unreadable metadata pkl {metadata_filename}, continuing without it")
|
|
return tuple({} for _ in _META_FIELDS)
|
|
|
|
|
|
@dataclass
|
|
class ArtifactSpec:
|
|
"""One model of the active bundle plus its unpacked metadata."""
|
|
model: Any
|
|
metadata: Any = None
|
|
input_shapes: ShapeDict = field(default_factory=dict)
|
|
output_slices: SliceDict = field(default_factory=dict)
|
|
|
|
def __post_init__(self):
|
|
self.metadata = self.model.metadata
|
|
if self.metadata:
|
|
self.input_shapes, self.output_slices = load_artifact_metadata(self.metadata.fileName)
|
|
|
|
|
|
# kept name: some runners annotate against the old alias
|
|
ModelData = ArtifactSpec
|
|
|
|
|
|
class RunnerRoot:
|
|
"""Shared root of the runner hierarchy.
|
|
|
|
Both ModelRunner and the per-model parser mixins (model_types.py) inherit
|
|
this, so the concrete `TinygradRunner(ModelRunner, *Tinygrad)` diamond keeps
|
|
one consistent parser registry + slice implementation.
|
|
"""
|
|
|
|
parser_method_dict: dict
|
|
_model_data: "ArtifactSpec | None"
|
|
|
|
def _slice_outputs(self, model_outputs):
|
|
raise NotImplementedError
|
|
|
|
|
|
class ModelRunner(RunnerRoot):
|
|
"""Base for the tinygrad/ONNX runners.
|
|
|
|
Owns the active bundle's ArtifactSpecs and the shared slice/parse plumbing;
|
|
subclasses provide input staging (prepare_inputs) and execution (_run_model).
|
|
"""
|
|
|
|
# False for fused runners, which warp + manage temporal buffers inside the JIT
|
|
uses_opencl_warp = True
|
|
|
|
def __init__(self):
|
|
active = _fetch_bundle()
|
|
if not active:
|
|
raise ValueError("runner started without an active model bundle")
|
|
|
|
self.models = {spec.type.raw: ArtifactSpec(spec) for spec in _qcom_models(active)}
|
|
self.is_20hz_3d = False
|
|
self.is_20hz = active.is20hz
|
|
self.inputs = {}
|
|
self.parser_method_dict = {}
|
|
self._model_data = None # active spec for the current operation
|
|
self._parser = self._constants = None
|
|
|
|
def _active_spec(self):
|
|
spec = self._model_data
|
|
if spec is None:
|
|
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
|
return spec
|
|
|
|
# views proxied straight off the active artifact spec; kept out of the class
|
|
# body (served via __getattr__) so the read surface stays data-driven
|
|
_SPEC_VIEW = frozenset(("input_shapes", "output_slices"))
|
|
|
|
def __getattr__(self, name):
|
|
if name == "constants":
|
|
return self._constants
|
|
if name == "vision_input_names":
|
|
return list(self._active_spec().input_shapes)
|
|
if name in ModelRunner._SPEC_VIEW:
|
|
return getattr(self._active_spec(), name)
|
|
raise AttributeError(name)
|
|
|
|
def prepare_inputs(self, imgs_cl, numpy_inputs, frames):
|
|
"""Stage image + numpy inputs for inference; implemented per backend."""
|
|
raise NotImplementedError
|
|
|
|
def _run_model(self):
|
|
"""Execute inference over the staged inputs; implemented per backend."""
|
|
raise NotImplementedError
|
|
|
|
def run_model(self):
|
|
# parsing happens inside each backend's _run_model
|
|
return self._run_model()
|
|
|
|
def _slice_outputs(self, model_outputs):
|
|
"""Split the flat output vector into named views per the artifact's slice table."""
|
|
sliced = {}
|
|
for tag, span in self._active_spec().output_slices.items():
|
|
sliced[tag] = model_outputs[np.newaxis, span]
|
|
if SEND_RAW_PRED:
|
|
sliced["raw_pred"] = model_outputs.copy()
|
|
return sliced
|
|
|
|
|
|
# ---- runner selection (which backend to build for the active bundle) ----------
|
|
|
|
def _qcom_models(bundle) -> list:
|
|
# usbeMac artifacts ride along in a bundle for the eGPU host; they are never
|
|
# loaded on QCOM and must not affect runner classification
|
|
return [m for m in bundle.models if m.type.raw != ModelType.usbeMac]
|
|
|
|
|
|
def _single_artifact_prefix(bundle, prefix: str) -> bool:
|
|
models = _qcom_models(bundle)
|
|
return len(models) == 1 and models[0].artifact.fileName.startswith(prefix)
|
|
|
|
|
|
def _is_fused_bundle(bundle) -> bool:
|
|
return _single_artifact_prefix(bundle, "driving_fused_")
|
|
|
|
|
|
def _is_supercombo_bundle(bundle) -> bool:
|
|
return _single_artifact_prefix(bundle, "driving_supercombo_")
|
|
|
|
|
|
def _is_split_bundle(bundle) -> bool:
|
|
present = {m.type.raw for m in _qcom_models(bundle)}
|
|
split_kinds = {ModelType.vision, ModelType.policy, ModelType.offPolicy, ModelType.onPolicy}
|
|
return not present.isdisjoint(split_kinds)
|
|
|
|
|
|
def get_model_runner() -> "ModelRunner":
|
|
"""Build the runner backend that fits the active bundle (supercombo / fused /
|
|
combined-split / split / single). Concrete runners are imported lazily so one
|
|
backend failing to load can't take down the others at import time."""
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.tinygrad_runner import (TinygradRunner,
|
|
TinygradSplitRunner)
|
|
bundle = _fetch_bundle()
|
|
# an eMac-only bundle (no QCOM-loadable models) runs the stock default on
|
|
# device; the big host serves the bundle's precompiled artifact
|
|
if not (bundle and bundle.models and _qcom_models(bundle)):
|
|
return TinygradRunner(ModelType.supercombo)
|
|
|
|
if _is_supercombo_bundle(bundle):
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.supercombo_runner import TinygradSupercomboRunner
|
|
return TinygradSupercomboRunner()
|
|
if _is_fused_bundle(bundle):
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.fused_runner import TinygradFusedRunner
|
|
return TinygradFusedRunner()
|
|
if _is_split_bundle(bundle) and has_combined_split_artifact(bundle):
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.combined_split_runner import TinygradCombinedSplitRunner
|
|
return TinygradCombinedSplitRunner()
|
|
if _is_split_bundle(bundle):
|
|
return TinygradSplitRunner()
|
|
return TinygradRunner(_qcom_models(bundle)[0].type.raw)
|