forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ 0798119
This commit is contained in:
3
iqpilot/selfdrive/iqmodeld/models/__init__.py
Normal file
3
iqpilot/selfdrive/iqmodeld/models/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
IQ model selection and runner support that is actively used by iqmodeld.
|
||||
"""
|
||||
100
iqpilot/selfdrive/iqmodeld/models/combined_artifact.py
Normal file
100
iqpilot/selfdrive/iqmodeld/models/combined_artifact.py
Normal file
@@ -0,0 +1,100 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from openpilot.system.hardware.hw import Paths
|
||||
|
||||
|
||||
_MODEL_ROOT = Path(Paths.model_root())
|
||||
_OVERRIDE_KEYS = (
|
||||
"combinedRuntimeArtifact",
|
||||
"combinedSplitArtifact",
|
||||
"iqCombinedArtifact",
|
||||
)
|
||||
_SPLIT_ROLE_PATTERN = re.compile(r"^driving_(vision|policy|off_policy|on_policy)_(.+)_tinygrad\.pkl$")
|
||||
|
||||
|
||||
def _bundle_models(bundle) -> list:
|
||||
models = getattr(bundle, "models", None)
|
||||
return list(models) if models is not None else []
|
||||
|
||||
|
||||
def _bundle_override_map(bundle) -> dict[str, str]:
|
||||
result: dict[str, str] = {}
|
||||
for override in getattr(bundle, "overrides", None) or []:
|
||||
key = getattr(override, "key", None)
|
||||
value = getattr(override, "value", None)
|
||||
if key and value:
|
||||
result[str(key)] = str(value)
|
||||
return result
|
||||
|
||||
|
||||
def _artifact_name(model) -> str:
|
||||
return getattr(getattr(model, "artifact", None), "fileName", "") or ""
|
||||
|
||||
|
||||
def _split_suffixes(bundle) -> list[str]:
|
||||
suffixes: list[str] = []
|
||||
for model in _bundle_models(bundle):
|
||||
match = _SPLIT_ROLE_PATTERN.match(_artifact_name(model))
|
||||
if match:
|
||||
suffixes.append(match.group(2))
|
||||
return suffixes
|
||||
|
||||
|
||||
def _derived_candidates(bundle) -> list[str]:
|
||||
seen: set[str] = set()
|
||||
candidates: list[str] = []
|
||||
|
||||
for suffix in _split_suffixes(bundle):
|
||||
for candidate in (
|
||||
f"driving_combined_{suffix}.pkl",
|
||||
f"iqmodeld_combined_{suffix}.pkl",
|
||||
):
|
||||
if candidate not in seen:
|
||||
seen.add(candidate)
|
||||
candidates.append(candidate)
|
||||
|
||||
ref = getattr(bundle, "ref", None)
|
||||
if ref:
|
||||
short_ref = str(ref)[:8]
|
||||
for candidate in (
|
||||
f"driving_combined_{short_ref}.pkl",
|
||||
f"iqmodeld_combined_{short_ref}.pkl",
|
||||
):
|
||||
if candidate not in seen:
|
||||
seen.add(candidate)
|
||||
candidates.append(candidate)
|
||||
|
||||
return candidates
|
||||
|
||||
|
||||
def combined_split_artifact_candidates(bundle) -> list[Path]:
|
||||
explicit_env = os.getenv("IQMODEL_COMBINED_PKL")
|
||||
if explicit_env:
|
||||
explicit_path = Path(explicit_env)
|
||||
return [explicit_path if explicit_path.is_absolute() else _MODEL_ROOT / explicit_path]
|
||||
|
||||
overrides = _bundle_override_map(bundle)
|
||||
explicit_names = [overrides[key] for key in _OVERRIDE_KEYS if key in overrides]
|
||||
if explicit_names:
|
||||
return [_MODEL_ROOT / name for name in explicit_names]
|
||||
|
||||
return [_MODEL_ROOT / name for name in _derived_candidates(bundle)]
|
||||
|
||||
|
||||
def resolve_combined_split_artifact(bundle) -> Path | None:
|
||||
for candidate in combined_split_artifact_candidates(bundle):
|
||||
if candidate.is_file():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def has_combined_split_artifact(bundle) -> bool:
|
||||
return resolve_combined_split_artifact(bundle) is not None
|
||||
14
iqpilot/selfdrive/iqmodeld/models/fetcher.py
Normal file
14
iqpilot/selfdrive/iqmodeld/models/fetcher.py
Normal file
@@ -0,0 +1,14 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright (c) IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
|
||||
|
||||
Public entry point for the model-manifest fetcher: prefers the compiled private
|
||||
bundle, falling back to the in-tree source. The default-runner fallback lives in
|
||||
ManifestDecoder now, so no post-import patching is needed.
|
||||
"""
|
||||
from openpilot.iqpilot._proprietary_loader import ProprietaryModuleMissing, load_private_module
|
||||
|
||||
try:
|
||||
load_private_module(__name__, "iqpilot_private.models.fetcher")
|
||||
except ProprietaryModuleMissing:
|
||||
from iqpilot.models_private_src.fetcher import * # noqa: F403
|
||||
10
iqpilot/selfdrive/iqmodeld/models/git_auth.py
Normal file
10
iqpilot/selfdrive/iqmodeld/models/git_auth.py
Normal file
@@ -0,0 +1,10 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright (c) IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
|
||||
"""
|
||||
from openpilot.iqpilot._proprietary_loader import ProprietaryModuleMissing, load_private_module
|
||||
|
||||
try:
|
||||
load_private_module(__name__, "iqpilot_private.models.git_auth")
|
||||
except ProprietaryModuleMissing:
|
||||
from iqpilot.models_private_src.git_auth import * # noqa: F403
|
||||
309
iqpilot/selfdrive/iqmodeld/models/helpers.py
Normal file
309
iqpilot/selfdrive/iqmodeld/models/helpers.py
Normal file
@@ -0,0 +1,309 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright (c) IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from cereal import custom
|
||||
from openpilot.common.params import Params
|
||||
from openpilot.common.swaglog import cloudlog
|
||||
from openpilot.iqpilot._proprietary_loader import ProprietaryModuleMissing, load_private_module
|
||||
from openpilot.system.hardware.hw import Paths
|
||||
|
||||
try:
|
||||
load_private_module(__name__, "iqpilot_private.models.helpers")
|
||||
except ProprietaryModuleMissing:
|
||||
try:
|
||||
from iqpilot.models_private_src.helpers import * # noqa: F403
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
ModelBundle = custom.IQModelManager.ModelBundle
|
||||
Runner = custom.IQModelManager.Runner
|
||||
_MODEL_ROOT = Path(Paths.model_root())
|
||||
_ACTIVE_BUNDLE_KEY = "ModelManager_ActiveBundle"
|
||||
_MODELS_CACHE_KEY = "ModelManager_ModelsCache"
|
||||
_RUNNER_CACHE_KEY = "ModelRunnerTypeCache"
|
||||
_DOWNLOAD_INDEX_KEY = "ModelManager_DownloadIndex"
|
||||
_PENDING_MODEL_RESTORE_FILE = "/data/k3_pending_model_restore"
|
||||
_STOCK_RUNNER = int(Runner.stock)
|
||||
_TINYGRAD_RUNNER = int(Runner.tinygrad)
|
||||
_SNPE_RUNNER = int(Runner.snpe)
|
||||
|
||||
_DEFAULT_MODEL_DIR = Path(__file__).resolve().parents[1] / "default_model"
|
||||
_DEFAULT_BUNDLE_JSON = _DEFAULT_MODEL_DIR / "bundle.json"
|
||||
_DEFAULT_BUNDLE_REF = "default"
|
||||
|
||||
|
||||
def get_default_model_bundle(_bundles):
|
||||
"""Legacy compatibility hook: stock default is preinstalled, not a manifest bundle."""
|
||||
return None
|
||||
|
||||
|
||||
def _coerce_runner_value(value) -> int | None:
|
||||
raw = getattr(value, "raw", value)
|
||||
try:
|
||||
return int(raw)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _bundle_models(bundle) -> list:
|
||||
models = getattr(bundle, "models", None)
|
||||
return list(models) if models is not None else []
|
||||
|
||||
|
||||
def _bundle_needs_runtime_upgrade(bundle) -> bool:
|
||||
if bundle is None:
|
||||
return False
|
||||
|
||||
if _coerce_runner_value(getattr(bundle, "runner", None)) == _SNPE_RUNNER:
|
||||
return True
|
||||
|
||||
for model in _bundle_models(bundle):
|
||||
file_name = getattr(getattr(model, "artifact", None), "fileName", "") or ""
|
||||
if file_name.endswith(".thneed"):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _load_cached_manifest_bundles(params: Params):
|
||||
cached = params.get(_MODELS_CACHE_KEY) or {}
|
||||
bundles = []
|
||||
for raw_bundle in cached.get("bundles", []):
|
||||
try:
|
||||
min_selector_version = int(raw_bundle.get("minimumSelectorVersion", raw_bundle.get("minimum_selector_version", 0)))
|
||||
compatibility_view = dict(raw_bundle)
|
||||
compatibility_view["minimumSelectorVersion"] = min_selector_version
|
||||
is_compatible = globals().get("is_bundle_version_compatible")
|
||||
if is_compatible is not None and not is_compatible(compatibility_view):
|
||||
continue
|
||||
|
||||
if "short_name" in raw_bundle:
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.fetcher import ManifestDecoder
|
||||
bundles.append(ManifestDecoder._decode_bundle(raw_bundle))
|
||||
continue
|
||||
|
||||
if "internalName" in raw_bundle:
|
||||
bundles.append(ModelBundle(**raw_bundle))
|
||||
continue
|
||||
|
||||
bundle = ModelBundle()
|
||||
bundle.index = int(raw_bundle["index"])
|
||||
bundle.internalName = raw_bundle.get("short_name")
|
||||
bundle.displayName = raw_bundle.get("display_name")
|
||||
bundle.status = 0
|
||||
bundle.generation = int(raw_bundle["generation"])
|
||||
bundle.environment = raw_bundle["environment"]
|
||||
bundle.runner = raw_bundle.get("runner", Runner.tinygrad)
|
||||
bundle.is20hz = raw_bundle.get("is_20hz", False)
|
||||
bundle.minimumSelectorVersion = int(min_selector_version)
|
||||
bundle.ref = raw_bundle.get("ref")
|
||||
bundle.overrides = []
|
||||
for key, value in raw_bundle.get("overrides", {}).items():
|
||||
override = custom.IQModelManager.Override()
|
||||
override.key = key
|
||||
override.value = value
|
||||
bundle.overrides.append(override)
|
||||
|
||||
bundle.models = []
|
||||
for raw_model in raw_bundle.get("models", []):
|
||||
model = custom.IQModelManager.Model()
|
||||
model.type = raw_model.get("type")
|
||||
for attr_name in ("artifact", "metadata"):
|
||||
raw_artifact = raw_model.get(attr_name)
|
||||
if not raw_artifact:
|
||||
continue
|
||||
artifact = custom.IQModelManager.Artifact()
|
||||
artifact.fileName = raw_artifact.get("file_name")
|
||||
download_uri = custom.IQModelManager.DownloadUri()
|
||||
download_uri.uri = raw_artifact.get("download_uri", {}).get("url")
|
||||
download_uri.sha256 = raw_artifact.get("download_uri", {}).get("sha256")
|
||||
artifact.downloadUri = download_uri
|
||||
setattr(model, attr_name, artifact)
|
||||
bundle.models.append(model)
|
||||
|
||||
bundles.append(bundle)
|
||||
except Exception:
|
||||
continue
|
||||
return bundles
|
||||
|
||||
|
||||
def _bundle_match_key(bundle) -> tuple[str | None, str | None, str | None]:
|
||||
return (
|
||||
getattr(bundle, "ref", None),
|
||||
getattr(bundle, "internalName", None),
|
||||
getattr(bundle, "displayName", None),
|
||||
)
|
||||
|
||||
|
||||
def _find_runtime_upgrade(bundle, params: Params, available_bundles=None):
|
||||
if not _bundle_needs_runtime_upgrade(bundle):
|
||||
return bundle
|
||||
|
||||
candidate_bundles = available_bundles if available_bundles is not None else _load_cached_manifest_bundles(params)
|
||||
ref, internal_name, display_name = _bundle_match_key(bundle)
|
||||
|
||||
for candidate in candidate_bundles:
|
||||
if getattr(candidate, "ref", None) and getattr(candidate, "ref", None) == ref:
|
||||
return candidate
|
||||
|
||||
for candidate in candidate_bundles:
|
||||
if getattr(candidate, "internalName", None) == internal_name:
|
||||
return candidate
|
||||
|
||||
for candidate in candidate_bundles:
|
||||
if getattr(candidate, "displayName", None) == display_name:
|
||||
return candidate
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def bundle_files_ready(bundle) -> bool:
|
||||
if bundle is None:
|
||||
return False
|
||||
|
||||
for model in _bundle_models(bundle):
|
||||
artifact = getattr(model, "artifact", None)
|
||||
metadata = getattr(model, "metadata", None)
|
||||
for file_name in (getattr(metadata, "fileName", None), getattr(artifact, "fileName", None)):
|
||||
if file_name and not (_MODEL_ROOT / file_name).is_file():
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def persist_active_bundle(params: Params, bundle) -> None:
|
||||
params.put(_ACTIVE_BUNDLE_KEY, bundle.to_dict())
|
||||
params.remove(_RUNNER_CACHE_KEY)
|
||||
|
||||
|
||||
def _load_default_bundle_dict() -> dict:
|
||||
return json.loads(_DEFAULT_BUNDLE_JSON.read_text())
|
||||
|
||||
|
||||
def _default_bundle_filenames(bundle_dict: dict) -> list[str]:
|
||||
names = []
|
||||
for model in bundle_dict.get("models", []):
|
||||
for artifact in (model.get("metadata"), model.get("artifact")):
|
||||
file_name = artifact.get("fileName", "") if isinstance(artifact, dict) else ""
|
||||
if file_name:
|
||||
names.append(file_name)
|
||||
return names
|
||||
|
||||
|
||||
def is_default_bundle(bundle) -> bool:
|
||||
return bool(bundle is not None and getattr(bundle, "ref", None) == _DEFAULT_BUNDLE_REF)
|
||||
|
||||
|
||||
def ensure_default_model_files(bundle_dict: dict = None) -> None:
|
||||
bundle_dict = bundle_dict if bundle_dict is not None else _load_default_bundle_dict()
|
||||
try:
|
||||
_MODEL_ROOT.mkdir(parents=True, exist_ok=True)
|
||||
except OSError as e:
|
||||
cloudlog.exception(f"default_model: cannot create model root: {e}")
|
||||
return
|
||||
for file_name in _default_bundle_filenames(bundle_dict):
|
||||
src = _DEFAULT_MODEL_DIR / file_name
|
||||
dst = _MODEL_ROOT / file_name
|
||||
if not src.is_file():
|
||||
cloudlog.error(f"default_model: shipped asset missing {src}")
|
||||
continue
|
||||
if dst.is_file() and dst.stat().st_size == src.stat().st_size:
|
||||
continue
|
||||
try:
|
||||
shutil.copy2(src, dst)
|
||||
cloudlog.warning(f"default_model: staged {file_name} into model root")
|
||||
except OSError as e:
|
||||
cloudlog.exception(f"default_model: failed staging {file_name}: {e}")
|
||||
|
||||
|
||||
def select_default_model(params: Params = None) -> None:
|
||||
params = Params() if params is None else params
|
||||
bundle_dict = _load_default_bundle_dict()
|
||||
ensure_default_model_files(bundle_dict)
|
||||
params.remove(_DOWNLOAD_INDEX_KEY)
|
||||
params.put(_ACTIVE_BUNDLE_KEY, bundle_dict)
|
||||
params.remove(_RUNNER_CACHE_KEY)
|
||||
params.put(_RUNNER_CACHE_KEY, _TINYGRAD_RUNNER)
|
||||
try:
|
||||
if os.path.isfile(_PENDING_MODEL_RESTORE_FILE):
|
||||
os.remove(_PENDING_MODEL_RESTORE_FILE)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def seed_default_bundle_if_unset(params: Params = None) -> None:
|
||||
params = Params() if params is None else params
|
||||
if params.get(_ACTIVE_BUNDLE_KEY) or params.get(_DOWNLOAD_INDEX_KEY) is not None:
|
||||
return
|
||||
try:
|
||||
select_default_model(params)
|
||||
cloudlog.warning("default_model: seeded Default (CD210) as active bundle")
|
||||
except Exception as e:
|
||||
cloudlog.exception(f"default_model: failed to seed default bundle: {e}")
|
||||
|
||||
|
||||
def get_runtime_bundle_upgrade(bundle, params: Params = None, available_bundles=None):
|
||||
params = Params() if params is None else params
|
||||
return _find_runtime_upgrade(bundle, params, available_bundles)
|
||||
|
||||
|
||||
def get_active_bundle(params: Params = None):
|
||||
params = Params() if params is None else params
|
||||
|
||||
try:
|
||||
active_bundle = params.get(_ACTIVE_BUNDLE_KEY) or {}
|
||||
if not active_bundle:
|
||||
return None
|
||||
is_compatible = globals().get("is_bundle_version_compatible")
|
||||
if is_compatible is not None and not is_compatible(active_bundle):
|
||||
return None
|
||||
bundle = ModelBundle(**active_bundle)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
replacement = _find_runtime_upgrade(bundle, params)
|
||||
if replacement is not None and replacement is not bundle and bundle_files_ready(replacement):
|
||||
persist_active_bundle(params, replacement)
|
||||
return replacement
|
||||
|
||||
return bundle
|
||||
|
||||
|
||||
def get_active_model_runner(params: Params = None, force_check=False):
|
||||
params = Params() if params is None else params
|
||||
|
||||
active_bundle = get_active_bundle(params)
|
||||
if not active_bundle:
|
||||
seed_default_bundle_if_unset(params)
|
||||
active_bundle = get_active_bundle(params)
|
||||
if not active_bundle:
|
||||
if params.get(_RUNNER_CACHE_KEY) != str(_TINYGRAD_RUNNER):
|
||||
params.put(_RUNNER_CACHE_KEY, _TINYGRAD_RUNNER)
|
||||
return _TINYGRAD_RUNNER
|
||||
|
||||
cached_runner_type = params.get(_RUNNER_CACHE_KEY)
|
||||
if cached_runner_type and not force_check and isinstance(cached_runner_type, str) and cached_runner_type.isdigit():
|
||||
return int(cached_runner_type)
|
||||
|
||||
runner_type = _coerce_runner_value(active_bundle.runner)
|
||||
if runner_type == _SNPE_RUNNER:
|
||||
replacement = _find_runtime_upgrade(active_bundle, params)
|
||||
if replacement is not None and replacement is not active_bundle and bundle_files_ready(replacement):
|
||||
persist_active_bundle(params, replacement)
|
||||
runner_type = _coerce_runner_value(replacement.runner)
|
||||
else:
|
||||
if replacement is not None and getattr(replacement, "index", None) is not None and params.get(_DOWNLOAD_INDEX_KEY) is None:
|
||||
params.put(_DOWNLOAD_INDEX_KEY, int(replacement.index))
|
||||
cloudlog.warning(f"Queued tinygrad migration for retired bundle {getattr(active_bundle, 'internalName', '<unknown>')}")
|
||||
runner_type = _TINYGRAD_RUNNER
|
||||
|
||||
if cached_runner_type != runner_type:
|
||||
params.put(_RUNNER_CACHE_KEY, int(runner_type))
|
||||
|
||||
return runner_type
|
||||
13
iqpilot/selfdrive/iqmodeld/models/inference_state.py
Normal file
13
iqpilot/selfdrive/iqmodeld/models/inference_state.py
Normal file
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
|
||||
|
||||
Common base for the per-process inference/runtime states. It seeds the lateral
|
||||
steer delay from the cached learned value so every subclass starts with a usable
|
||||
number before its first liveDelay message arrives.
|
||||
"""
|
||||
from openpilot.iqpilot.common.steer_delay import cached_steer_delay
|
||||
|
||||
|
||||
class InferenceStateBase:
|
||||
def __init__(self):
|
||||
self.lat_delay = cached_steer_delay()
|
||||
391
iqpilot/selfdrive/iqmodeld/models/manager.py
Normal file
391
iqpilot/selfdrive/iqmodeld/models/manager.py
Normal file
@@ -0,0 +1,391 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright (c) IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
|
||||
"""
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
from cereal import custom
|
||||
from openpilot.common.realtime import Ratekeeper
|
||||
from openpilot.common.time_helpers import system_time_valid
|
||||
from openpilot.iqpilot._proprietary_loader import ProprietaryModuleMissing, load_private_module
|
||||
from openpilot.common.swaglog import cloudlog
|
||||
from openpilot.system.hardware.hw import Paths
|
||||
|
||||
_TIME_SYNC_WAIT_TIMEOUT_S = 30.0
|
||||
_TIME_SYNC_POLL_S = 0.5
|
||||
|
||||
|
||||
def _wait_for_valid_clock(timeout: float = _TIME_SYNC_WAIT_TIMEOUT_S) -> None:
|
||||
if system_time_valid():
|
||||
return
|
||||
cloudlog.warning("models_manager: system clock not yet valid, waiting for NTP before fetching")
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if system_time_valid():
|
||||
cloudlog.warning("models_manager: system clock is now valid, resuming")
|
||||
return
|
||||
time.sleep(_TIME_SYNC_POLL_S)
|
||||
cloudlog.warning("models_manager: gave up waiting for a valid clock, proceeding anyway")
|
||||
|
||||
try:
|
||||
load_private_module(__name__, "iqpilot_private.models.manager")
|
||||
_BaseIQModelManager = IQModelManager # noqa: F821
|
||||
except ProprietaryModuleMissing:
|
||||
from iqpilot.models_private_src.manager import IQModelManager as _BaseIQModelManager
|
||||
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.git_auth import get_aiohttp_auth
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.helpers import (
|
||||
bundle_files_ready,
|
||||
get_active_bundle,
|
||||
get_runtime_bundle_upgrade,
|
||||
is_default_bundle,
|
||||
persist_active_bundle,
|
||||
)
|
||||
|
||||
|
||||
_ACTIVE_BUNDLE_KEY = "ModelManager_ActiveBundle"
|
||||
_DOWNLOAD_INDEX_KEY = "ModelManager_DownloadIndex"
|
||||
_RUNNER_CACHE_KEY = "ModelRunnerTypeCache"
|
||||
|
||||
|
||||
class IQModelManager(_BaseIQModelManager):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._validated_active_key: tuple[tuple[str, str], ...] | None = None
|
||||
self._manifest_refresh_key: tuple[tuple[str, str], ...] | None = None
|
||||
|
||||
@staticmethod
|
||||
def _bundle_index(bundle) -> int | None:
|
||||
try:
|
||||
return int(getattr(bundle, "index", -1))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _bundle_files(bundle) -> list[tuple[str, str]]:
|
||||
files = []
|
||||
for model in getattr(bundle, "models", []) or []:
|
||||
for artifact in (getattr(model, "metadata", None), getattr(model, "artifact", None)):
|
||||
filename = getattr(artifact, "fileName", "") if artifact is not None else ""
|
||||
if not filename:
|
||||
continue
|
||||
download_uri = getattr(artifact, "downloadUri", None)
|
||||
sha256 = getattr(download_uri, "sha256", "") if download_uri is not None else ""
|
||||
files.append((filename, sha256 or ""))
|
||||
return files
|
||||
|
||||
@staticmethod
|
||||
def _safe_model_path(filename: str) -> Path | None:
|
||||
if not filename or os.path.basename(filename) != filename:
|
||||
cloudlog.warning(f"Ignoring unsafe model filename {filename!r}")
|
||||
return None
|
||||
|
||||
root = Path(Paths.model_root()).resolve()
|
||||
path = (root / filename).resolve()
|
||||
try:
|
||||
path.relative_to(root)
|
||||
except ValueError:
|
||||
cloudlog.warning(f"Ignoring model path outside model root {path}")
|
||||
return None
|
||||
return path
|
||||
|
||||
@staticmethod
|
||||
def _verify_file_sync(path: Path, expected_hash: str) -> bool:
|
||||
if not path.is_file():
|
||||
return False
|
||||
if not expected_hash:
|
||||
return True
|
||||
|
||||
sha256_hash = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
sha256_hash.update(chunk)
|
||||
return sha256_hash.hexdigest().lower() == expected_hash.lower()
|
||||
|
||||
def _bundle_validation_key(self, bundle) -> tuple[tuple[str, str], ...]:
|
||||
return tuple(self._bundle_files(bundle))
|
||||
|
||||
def _bundle_files_valid(self, bundle) -> bool:
|
||||
for filename, expected_hash in self._bundle_files(bundle):
|
||||
path = self._safe_model_path(filename)
|
||||
if path is None or not self._verify_file_sync(path, expected_hash):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _remove_bundle_files(self, bundle) -> None:
|
||||
for filename, _expected_hash in self._bundle_files(bundle):
|
||||
path = self._safe_model_path(filename)
|
||||
if path is None:
|
||||
continue
|
||||
for candidate in (path, Path(f"{path}.download")):
|
||||
try:
|
||||
if candidate.is_file():
|
||||
candidate.unlink()
|
||||
except OSError as e:
|
||||
cloudlog.exception(f"Failed to remove model artifact {candidate}: {e}")
|
||||
|
||||
def _find_available_bundle(self, target):
|
||||
target_index = self._bundle_index(target)
|
||||
target_ref = getattr(target, "ref", None)
|
||||
target_internal = getattr(target, "internalName", None)
|
||||
target_display = getattr(target, "displayName", None)
|
||||
|
||||
for bundle in self.available_models:
|
||||
if target_index is not None and self._bundle_index(bundle) == target_index:
|
||||
return bundle
|
||||
if target_ref and getattr(bundle, "ref", None) == target_ref:
|
||||
return bundle
|
||||
if target_internal and getattr(bundle, "internalName", None) == target_internal:
|
||||
return bundle
|
||||
if target_display and getattr(bundle, "displayName", None) == target_display:
|
||||
return bundle
|
||||
return None
|
||||
|
||||
def _bundle_matches(self, left, right) -> bool:
|
||||
if left is None or right is None:
|
||||
return False
|
||||
|
||||
left_index = self._bundle_index(left)
|
||||
right_index = self._bundle_index(right)
|
||||
if left_index is not None and right_index is not None and left_index == right_index:
|
||||
return True
|
||||
|
||||
for attr in ("ref", "internalName", "displayName"):
|
||||
left_value = getattr(left, attr, None)
|
||||
if left_value and left_value == getattr(right, attr, None):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _clear_active_bundle(self) -> None:
|
||||
self.params.remove(_ACTIVE_BUNDLE_KEY)
|
||||
self.params.remove(_RUNNER_CACHE_KEY)
|
||||
self.active_bundle = None
|
||||
self._validated_active_key = None
|
||||
|
||||
def _download_request_matches(self, bundle) -> bool:
|
||||
bundle_index = self._bundle_index(bundle)
|
||||
return bundle_index is not None and self._download_index() == bundle_index
|
||||
|
||||
def _queue_active_redownload_if_invalid(self) -> None:
|
||||
if self.active_bundle is None:
|
||||
self._validated_active_key = None
|
||||
return
|
||||
|
||||
validation_key = self._bundle_validation_key(self.active_bundle)
|
||||
if validation_key == self._validated_active_key:
|
||||
return
|
||||
|
||||
if self._bundle_files_valid(self.active_bundle):
|
||||
self._validated_active_key = validation_key
|
||||
return
|
||||
|
||||
bundle = self._find_available_bundle(self.active_bundle) or self.active_bundle
|
||||
bundle_index = self._bundle_index(bundle)
|
||||
cloudlog.warning(f"Active model {_display_bundle_name(self.active_bundle)} is missing or corrupt; queueing redownload")
|
||||
self._remove_bundle_files(bundle)
|
||||
self._clear_active_bundle()
|
||||
if bundle_index is not None and self._download_index() is None:
|
||||
self.params.put(_DOWNLOAD_INDEX_KEY, bundle_index)
|
||||
|
||||
def _find_manifest_counterpart(self, target):
|
||||
# never match by index: indexes shift between manifest generations, and a
|
||||
# positional match could redownload a different model than the user selected
|
||||
for attr in ("ref", "internalName", "displayName"):
|
||||
value = getattr(target, attr, None)
|
||||
if not value:
|
||||
continue
|
||||
for bundle in self.available_models:
|
||||
if getattr(bundle, attr, None) == value:
|
||||
return bundle
|
||||
return None
|
||||
|
||||
def _queue_active_manifest_refresh(self) -> None:
|
||||
active = self.active_bundle
|
||||
if active is None or is_default_bundle(active):
|
||||
return
|
||||
if self._download_index() is not None:
|
||||
return
|
||||
|
||||
counterpart = self._find_manifest_counterpart(active)
|
||||
if counterpart is None:
|
||||
return
|
||||
counterpart_index = self._bundle_index(counterpart)
|
||||
if counterpart_index is None:
|
||||
return
|
||||
|
||||
active_files = dict(self._bundle_files(active))
|
||||
stale = False
|
||||
for filename, sha in self._bundle_files(counterpart):
|
||||
if not sha:
|
||||
continue
|
||||
active_sha = active_files.get(filename)
|
||||
# an empty recorded hash can't prove a mismatch, so it never triggers a redownload
|
||||
if active_sha is None or (active_sha and active_sha.lower() != sha.lower()):
|
||||
stale = True
|
||||
break
|
||||
if not stale:
|
||||
self._manifest_refresh_key = None
|
||||
return
|
||||
|
||||
# the manifest may be an expired offline cache, so keep the active bundle and its
|
||||
# files in place: the download flow replaces artifacts atomically and only persists
|
||||
# the counterpart as active once everything landed. One attempt per bundle per run
|
||||
# so a dead network doesn't turn the 1Hz loop into a download-retry storm.
|
||||
key = self._bundle_validation_key(active)
|
||||
if key == self._manifest_refresh_key:
|
||||
return
|
||||
self._manifest_refresh_key = key
|
||||
|
||||
cloudlog.warning(f"Active model {_display_bundle_name(active)} artifacts are stale vs current manifest; queueing redownload")
|
||||
self.params.put(_DOWNLOAD_INDEX_KEY, counterpart_index)
|
||||
|
||||
async def _download_file(self, url: str, path: str, model) -> None:
|
||||
temp_path = f"{path}.download"
|
||||
self._download_start_times[model.fileName] = time.monotonic()
|
||||
|
||||
try:
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
|
||||
async with aiohttp.ClientSession(auth=get_aiohttp_auth()) as session:
|
||||
async with session.get(url) as response:
|
||||
response.raise_for_status()
|
||||
total_size = int(response.headers.get("content-length", 0))
|
||||
bytes_downloaded = 0
|
||||
|
||||
with open(temp_path, "wb") as f:
|
||||
async for chunk in response.content.iter_chunked(self._chunk_size):
|
||||
f.write(chunk)
|
||||
bytes_downloaded += len(chunk)
|
||||
|
||||
if self._download_index() is None:
|
||||
raise Exception("Download cancelled")
|
||||
|
||||
if total_size > 0:
|
||||
progress = (bytes_downloaded / total_size) * 100
|
||||
model.downloadProgress.status = custom.IQModelManager.DownloadStatus.downloading
|
||||
model.downloadProgress.progress = progress
|
||||
model.downloadProgress.eta = self._calculate_eta(model.fileName, progress)
|
||||
self._report_status()
|
||||
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
|
||||
os.replace(temp_path, path)
|
||||
|
||||
except Exception:
|
||||
if os.path.exists(temp_path):
|
||||
os.remove(temp_path)
|
||||
raise
|
||||
|
||||
finally:
|
||||
self._download_start_times.pop(model.fileName, None)
|
||||
|
||||
async def _download_bundle(self, model_bundle: custom.IQModelManager.ModelBundle, destination_path: str) -> None:
|
||||
self.selected_bundle = model_bundle
|
||||
self.selected_bundle.status = custom.IQModelManager.DownloadStatus.downloading
|
||||
os.makedirs(destination_path, exist_ok=True)
|
||||
|
||||
try:
|
||||
if not self._download_request_matches(model_bundle):
|
||||
raise RuntimeError("Download cancelled")
|
||||
|
||||
tasks = [self._process_model(model, destination_path) for model in self.selected_bundle.models]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
if not self._download_request_matches(model_bundle):
|
||||
raise RuntimeError("Download cancelled")
|
||||
|
||||
self.active_bundle = self.selected_bundle
|
||||
self.active_bundle.status = custom.IQModelManager.DownloadStatus.downloaded
|
||||
self.params.put(_ACTIVE_BUNDLE_KEY, self.active_bundle.to_dict())
|
||||
self.params.remove(_RUNNER_CACHE_KEY)
|
||||
self.selected_bundle = None
|
||||
|
||||
except Exception:
|
||||
if self._download_request_matches(model_bundle) and self.selected_bundle is not None:
|
||||
self.selected_bundle.status = custom.IQModelManager.DownloadStatus.failed
|
||||
else:
|
||||
self.selected_bundle = None
|
||||
raise
|
||||
|
||||
finally:
|
||||
self._report_status()
|
||||
|
||||
def download(self, model_bundle: custom.IQModelManager.ModelBundle, destination_path: str) -> None:
|
||||
asyncio.run(self._download_bundle(model_bundle, destination_path))
|
||||
|
||||
def _queue_tinygrad_upgrade(self) -> None:
|
||||
if self.active_bundle is None:
|
||||
return
|
||||
|
||||
replacement = get_runtime_bundle_upgrade(self.active_bundle, self.params, self.available_models)
|
||||
if replacement is None or replacement is self.active_bundle:
|
||||
return
|
||||
|
||||
if bundle_files_ready(replacement):
|
||||
persist_active_bundle(self.params, replacement)
|
||||
self.active_bundle = replacement
|
||||
return
|
||||
|
||||
if self._download_index() is None and getattr(replacement, "index", None) is not None:
|
||||
self.params.put("ModelManager_DownloadIndex", int(replacement.index))
|
||||
cloudlog.warning(f"Queued tinygrad upgrade for retired bundle {getattr(self.active_bundle, 'internalName', '<unknown>')}")
|
||||
|
||||
def main_thread(self) -> None:
|
||||
_wait_for_valid_clock()
|
||||
rk = Ratekeeper(1, print_delay_threshold=None)
|
||||
|
||||
while True:
|
||||
try:
|
||||
# before NTP the TLS cert reads "not yet valid" and every fetch SSL-fails; one line, not spam
|
||||
if not system_time_valid():
|
||||
if not getattr(self, "_ntp_wait_logged", False):
|
||||
cloudlog.warning("models_manager: waiting for NTP before fetching (system clock not valid)")
|
||||
self._ntp_wait_logged = True
|
||||
rk.keep_time()
|
||||
continue
|
||||
self._ntp_wait_logged = False
|
||||
|
||||
self.available_models = self.model_fetcher.get_available_bundles()
|
||||
self.active_bundle = get_active_bundle(self.params)
|
||||
self._queue_active_redownload_if_invalid()
|
||||
self._queue_tinygrad_upgrade()
|
||||
self._queue_active_manifest_refresh()
|
||||
|
||||
if (index_to_download := self._download_index()) is not None:
|
||||
if model_to_download := next((model for model in self.available_models if model.index == index_to_download), None):
|
||||
try:
|
||||
self.download(model_to_download, Paths.model_root())
|
||||
except Exception as e:
|
||||
cloudlog.exception(e)
|
||||
finally:
|
||||
self.params.remove("ModelManager_DownloadIndex")
|
||||
self.selected_bundle = None
|
||||
|
||||
if self.params.get("ModelManager_ClearCache"):
|
||||
self.clear_model_cache()
|
||||
self.params.remove("ModelManager_ClearCache")
|
||||
|
||||
self._report_status()
|
||||
rk.keep_time()
|
||||
|
||||
except Exception as e:
|
||||
cloudlog.exception(f"Error in main thread: {str(e)}")
|
||||
rk.keep_time()
|
||||
|
||||
|
||||
def _display_bundle_name(bundle) -> str:
|
||||
return getattr(bundle, "internalName", None) or getattr(bundle, "displayName", None) or "<unknown>"
|
||||
|
||||
|
||||
def main():
|
||||
IQModelManager().main_thread()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
3
iqpilot/selfdrive/iqmodeld/models/runners/__init__.py
Normal file
3
iqpilot/selfdrive/iqmodeld/models/runners/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
Runner interfaces used by iqmodeld model execution.
|
||||
"""
|
||||
207
iqpilot/selfdrive/iqmodeld/models/runners/model_runner.py
Normal file
207
iqpilot/selfdrive/iqmodeld/models/runners/model_runner.py
Normal file
@@ -0,0 +1,207 @@
|
||||
"""
|
||||
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 cereal import custom
|
||||
from openpilot.system.hardware import TICI
|
||||
from openpilot.system.hardware.hw import Paths as _hw_paths
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.helpers import get_active_bundle as _fetch_bundle
|
||||
from openpilot.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 openpilot.iqpilot.selfdrive.iqmodeld.native.iqmodel_pyx import GpuMemorySlot, RoadProjector
|
||||
else:
|
||||
def _resolve_native_types() -> tuple[Any, Any]:
|
||||
try:
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.native.iqmodel_pyx import GpuMemorySlot as iq_clmem
|
||||
from openpilot.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()
|
||||
|
||||
|
||||
def load_artifact_metadata(metadata_filename):
|
||||
"""Read one artifact's metadata pkl: (input shapes, output slices)."""
|
||||
with open(os.path.join(CUSTOM_MODEL_PATH, metadata_filename), 'rb') as fh:
|
||||
blob = _pk.load(fh)
|
||||
return tuple(blob.get(field, {}) for field 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 active.models}
|
||||
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 _single_artifact_prefix(bundle, prefix: str) -> bool:
|
||||
return len(bundle.models) == 1 and bundle.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 bundle.models}
|
||||
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 openpilot.iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.tinygrad_runner import (TinygradRunner,
|
||||
TinygradSplitRunner)
|
||||
bundle = _fetch_bundle()
|
||||
if not (bundle and bundle.models):
|
||||
return TinygradRunner(ModelType.supercombo)
|
||||
|
||||
if _is_supercombo_bundle(bundle):
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.supercombo_runner import TinygradSupercomboRunner
|
||||
return TinygradSupercomboRunner()
|
||||
if _is_fused_bundle(bundle):
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.fused_runner import TinygradFusedRunner
|
||||
return TinygradFusedRunner()
|
||||
if _is_split_bundle(bundle) and has_combined_split_artifact(bundle):
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.combined_split_runner import TinygradCombinedSplitRunner
|
||||
return TinygradCombinedSplitRunner()
|
||||
if _is_split_bundle(bundle):
|
||||
return TinygradSplitRunner()
|
||||
return TinygradRunner(bundle.models[0].type.raw)
|
||||
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
ONNX runner support for iqmodeld.
|
||||
"""
|
||||
@@ -0,0 +1,57 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import CLMemDict, FrameDict, ModelType, NumpyDict, ShapeDict
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld import MODEL_PATH
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.config import ModelConstants
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.parser import ArchiveParser
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.runtime.ort import ORT_TYPES_TO_NP_TYPES, make_onnx_cpu_runner
|
||||
|
||||
|
||||
def _onnx_dtype_table(session) -> dict[str, np.dtype]:
|
||||
return {
|
||||
tensor_info.name: ORT_TYPES_TO_NP_TYPES[tensor_info.type]
|
||||
for tensor_info in session.get_inputs()
|
||||
}
|
||||
|
||||
|
||||
class ONNXRunner(ModelRunner):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.runner = make_onnx_cpu_runner(MODEL_PATH)
|
||||
self._constants = ModelConstants
|
||||
self._model_data = self.models.get(ModelType.supercombo)
|
||||
self._input_dtypes = _onnx_dtype_table(self.runner)
|
||||
self._parser = ArchiveParser()
|
||||
self.parser_method_dict[ModelType.supercombo] = self._parser.parse_outputs
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
return {tensor_info.name: tensor_info.shape for tensor_info in self.runner.get_inputs()}
|
||||
|
||||
def _frame_as_numpy(self, stream_name: str, imgs_cl: CLMemDict, frames: FrameDict) -> np.ndarray:
|
||||
flattened = frames[stream_name].as_numpy(imgs_cl[stream_name])
|
||||
shaped = flattened.reshape(self.input_shapes[stream_name])
|
||||
return shaped.astype(self._input_dtypes[stream_name])
|
||||
|
||||
def prepare_inputs(self, imgs_cl: CLMemDict, numpy_inputs: NumpyDict, frames: FrameDict) -> dict:
|
||||
staged_inputs = dict(numpy_inputs)
|
||||
for stream_name in imgs_cl:
|
||||
staged_inputs[stream_name] = self._frame_as_numpy(stream_name, imgs_cl, frames)
|
||||
self.inputs = staged_inputs
|
||||
return staged_inputs
|
||||
|
||||
def _parse_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
if self._model_data is None:
|
||||
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
||||
return self.parser_method_dict[self._model_data.model.type.raw](self._slice_outputs(model_outputs))
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
combined = self.runner.run(None, self.inputs)[0].reshape(-1)
|
||||
return self._parse_outputs(combined)
|
||||
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
Tinygrad runner support for iqmodeld.
|
||||
"""
|
||||
@@ -0,0 +1,245 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.combined_artifact import resolve_combined_split_artifact
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.helpers import get_active_bundle
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import NumpyDict, ShapeDict, SliceDict
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
|
||||
|
||||
def _tinygrad_imports():
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.tensor import Tensor
|
||||
return Tensor, Device
|
||||
|
||||
|
||||
def _phase_roles(meta_by_role: dict[str, dict]) -> list[str]:
|
||||
return [name for name in meta_by_role if name != "vision"]
|
||||
|
||||
|
||||
def _phase_desire_key(policy_shapes: dict[str, tuple[int, ...]]) -> str:
|
||||
for key in policy_shapes:
|
||||
if key.startswith("desire"):
|
||||
return key
|
||||
raise KeyError("No desire-like key found in policy inputs")
|
||||
|
||||
|
||||
def _phase_image_keys(vision_shapes: dict[str, tuple[int, ...]]) -> tuple[str, str]:
|
||||
names = sorted(name for name in vision_shapes if "img" in name)
|
||||
road_key = next((name for name in names if "big" not in name), None)
|
||||
wide_key = next((name for name in names if "big" in name), None)
|
||||
if road_key is None or wide_key is None:
|
||||
raise ValueError(f"Unable to resolve road/wide image keys from {list(vision_shapes)}")
|
||||
return road_key, wide_key
|
||||
|
||||
|
||||
def _base_policy_keys(policy_shapes: dict[str, tuple[int, ...]]) -> set[str]:
|
||||
desired_key = _phase_desire_key(policy_shapes)
|
||||
return {desired_key, "features_buffer", "traffic_convention", "action_t"}
|
||||
|
||||
|
||||
def _slice_map(raw_blob: np.ndarray, slices: dict[str, slice]) -> NumpyDict:
|
||||
return {name: raw_blob[np.newaxis, section] for name, section in slices.items() if name != "pad"}
|
||||
|
||||
|
||||
class TinygradCombinedSplitRunner(ModelRunner):
|
||||
uses_opencl_warp: bool = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._constants = SplitModelConstants
|
||||
self._parser = PhaseParser()
|
||||
self._bundle = get_active_bundle()
|
||||
self._artifact_path = resolve_combined_split_artifact(self._bundle)
|
||||
if self._artifact_path is None:
|
||||
raise FileNotFoundError("No IQ combined split artifact is available for the active bundle")
|
||||
|
||||
with open(self._artifact_path, "rb") as artifact:
|
||||
runtime_package: dict[Any, Any] = pickle.load(artifact)
|
||||
|
||||
self._meta_by_role = runtime_package.get("meta_by_role", runtime_package.get("metadata", {}))
|
||||
self._policy_roles = runtime_package.get("roles", _phase_roles(self._meta_by_role))
|
||||
self._camera_programs = {
|
||||
camera_key: spec
|
||||
for camera_key, spec in runtime_package.items()
|
||||
if isinstance(camera_key, tuple) and isinstance(spec, dict)
|
||||
}
|
||||
self._execute_bundle = runtime_package.get("execute_bundle", runtime_package.get("run_policy"))
|
||||
self._frame_stride = int(runtime_package.get("frame_stride", runtime_package.get("frame_skip", 1)))
|
||||
|
||||
if "vision" not in self._meta_by_role:
|
||||
raise ValueError("Combined split artifact is missing vision metadata")
|
||||
if not self._policy_roles:
|
||||
raise ValueError("Combined split artifact is missing policy roles")
|
||||
if self._execute_bundle is None:
|
||||
raise ValueError("Combined split artifact is missing execute_bundle")
|
||||
|
||||
self._vision_meta = self._meta_by_role["vision"]
|
||||
self._primary_policy_meta = self._meta_by_role[self._policy_roles[0]]
|
||||
self._desired_key = _phase_desire_key(self._primary_policy_meta["input_shapes"])
|
||||
self._road_key, self._wide_key = _phase_image_keys(self._vision_meta["input_shapes"])
|
||||
self._extra_policy_keys = [
|
||||
key for key in self._primary_policy_meta["input_shapes"]
|
||||
if key not in _base_policy_keys(self._primary_policy_meta["input_shapes"])
|
||||
]
|
||||
|
||||
self._queue_tensors: dict[str, Any] | None = None
|
||||
self._numpy_state: dict[str, np.ndarray] | None = None
|
||||
self._camera_shape: tuple[int, int] | None = None
|
||||
self._blob_cache: dict[tuple[str, int], Any] = {}
|
||||
self._last_desire = np.zeros(self._primary_policy_meta["input_shapes"][self._desired_key][2], dtype=np.float32)
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
return [self._road_key, self._wide_key]
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
merged: ShapeDict = dict(self._vision_meta["input_shapes"])
|
||||
for role in self._policy_roles:
|
||||
merged.update(self._meta_by_role[role]["input_shapes"])
|
||||
return merged
|
||||
|
||||
@property
|
||||
def output_slices(self) -> SliceDict:
|
||||
merged: SliceDict = dict(self._vision_meta["output_slices"])
|
||||
for role in self._policy_roles:
|
||||
merged.update(self._meta_by_role[role]["output_slices"])
|
||||
return merged
|
||||
|
||||
def prepare_inputs(self, imgs_cl, numpy_inputs, frames):
|
||||
raise RuntimeError("Combined split runner manages its own warp + queue state; use run_fused()")
|
||||
|
||||
def _frame_blob(self, stream_name: str, buf):
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
raw_frame = np.frombuffer(buf.data, dtype=np.uint8)
|
||||
cache_key = (stream_name, raw_frame.ctypes.data)
|
||||
tensor = self._blob_cache.get(cache_key)
|
||||
if tensor is None:
|
||||
tensor = Tensor.from_blob(raw_frame.ctypes.data, (raw_frame.size,), dtype="uint8", device=Device.DEFAULT)
|
||||
self._blob_cache[cache_key] = tensor
|
||||
return tensor
|
||||
|
||||
def _allocate_runtime_state(self, camera_width: int, camera_height: int) -> None:
|
||||
if self._queue_tensors is not None and self._camera_shape == (camera_width, camera_height):
|
||||
return
|
||||
if (camera_width, camera_height) not in self._camera_programs:
|
||||
raise RuntimeError(f"No combined split kernels available for {camera_width}x{camera_height}")
|
||||
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
vision_shapes = self._vision_meta["input_shapes"]
|
||||
policy_shapes = self._primary_policy_meta["input_shapes"]
|
||||
|
||||
image_shape = vision_shapes[self._road_key]
|
||||
frame_history = image_shape[1] // 6
|
||||
queue_depth = self._frame_stride * (frame_history - 1) + 1
|
||||
frame_queue_shape = (queue_depth, 6, image_shape[2], image_shape[3])
|
||||
|
||||
feature_shape = policy_shapes["features_buffer"]
|
||||
desired_shape = policy_shapes[self._desired_key]
|
||||
traffic_shape = policy_shapes["traffic_convention"]
|
||||
action_shape = policy_shapes.get("action_t", traffic_shape)
|
||||
|
||||
numpy_state = {
|
||||
"tfm": np.zeros((3, 3), dtype=np.float32),
|
||||
"big_tfm": np.zeros((3, 3), dtype=np.float32),
|
||||
"desire": np.zeros(desired_shape[2], dtype=np.float32),
|
||||
"traffic_convention": np.zeros(traffic_shape, dtype=np.float32),
|
||||
"action_t": np.zeros(action_shape, dtype=np.float32),
|
||||
}
|
||||
for key in self._extra_policy_keys:
|
||||
numpy_state[key] = np.zeros(policy_shapes[key], dtype=np.float32)
|
||||
|
||||
queue_tensors = {
|
||||
"img_q": Tensor(np.zeros(frame_queue_shape, dtype=np.uint8), device=Device.DEFAULT).contiguous().realize(),
|
||||
"big_img_q": Tensor(np.zeros(frame_queue_shape, dtype=np.uint8), device=Device.DEFAULT).contiguous().realize(),
|
||||
"feat_q": Tensor(
|
||||
np.zeros((self._frame_stride * (feature_shape[1] - 1) + 1, feature_shape[0], feature_shape[2]), dtype=np.float32),
|
||||
device=Device.DEFAULT,
|
||||
).contiguous().realize(),
|
||||
"desire_q": Tensor(
|
||||
np.zeros((self._frame_stride * desired_shape[1], desired_shape[0], desired_shape[2]), dtype=np.float32),
|
||||
device=Device.DEFAULT,
|
||||
).contiguous().realize(),
|
||||
**{name: Tensor(value, device="NPY").realize() for name, value in numpy_state.items()},
|
||||
}
|
||||
|
||||
self._queue_tensors = queue_tensors
|
||||
self._numpy_state = numpy_state
|
||||
self._camera_shape = (camera_width, camera_height)
|
||||
|
||||
def _policy_inputs(self) -> dict[str, Any]:
|
||||
assert self._queue_tensors is not None
|
||||
tensor_names = ["feat_q", "desire_q", "desire", "traffic_convention", "action_t", *self._extra_policy_keys]
|
||||
return {name: self._queue_tensors[name] for name in tensor_names if name in self._queue_tensors}
|
||||
|
||||
def _merge_policy_outputs(self, raw_outputs: tuple[Any, ...]) -> NumpyDict:
|
||||
outputs = self._parser.parse_vision_outputs(
|
||||
_slice_map(raw_outputs[0].numpy().flatten(), self._vision_meta["output_slices"])
|
||||
)
|
||||
|
||||
has_on_policy = any(role == "on_policy" for role in self._policy_roles)
|
||||
for role_name, tensor_out in zip(self._policy_roles, raw_outputs[1:], strict=True):
|
||||
parsed = self._parser.parse_policy_outputs(
|
||||
_slice_map(tensor_out.numpy().flatten(), self._meta_by_role[role_name]["output_slices"])
|
||||
)
|
||||
if role_name == "off_policy" and has_on_policy:
|
||||
parsed.pop("plan", None)
|
||||
outputs.update(parsed)
|
||||
|
||||
if "planplus" in outputs and "plan" in outputs:
|
||||
outputs["plan"] = outputs["plan"] + outputs["planplus"]
|
||||
return outputs
|
||||
|
||||
def run_fused(self, bufs: dict, transforms: dict[str, np.ndarray], numpy_inputs: NumpyDict) -> NumpyDict:
|
||||
main_buf = bufs[self._road_key]
|
||||
self._allocate_runtime_state(main_buf.width, main_buf.height)
|
||||
assert self._queue_tensors is not None and self._numpy_state is not None and self._camera_shape is not None
|
||||
|
||||
self._numpy_state["tfm"][:] = transforms[self._road_key]
|
||||
self._numpy_state["big_tfm"][:] = transforms[self._wide_key]
|
||||
|
||||
current_desire = numpy_inputs[self._desired_key].copy()
|
||||
current_desire[0] = 0
|
||||
self._numpy_state["desire"][:] = np.where(current_desire - self._last_desire > 0.99, current_desire, 0)
|
||||
self._last_desire[:] = current_desire
|
||||
|
||||
if "traffic_convention" in numpy_inputs:
|
||||
self._numpy_state["traffic_convention"][:] = numpy_inputs["traffic_convention"]
|
||||
if "action_t" in numpy_inputs:
|
||||
self._numpy_state["action_t"][:] = numpy_inputs["action_t"]
|
||||
for key in self._extra_policy_keys:
|
||||
if key in numpy_inputs:
|
||||
self._numpy_state[key][:] = numpy_inputs[key]
|
||||
|
||||
stage_inputs = self._camera_programs[self._camera_shape].get("stage_inputs", self._camera_programs[self._camera_shape].get("warp_enqueue"))
|
||||
if stage_inputs is None:
|
||||
raise RuntimeError("Combined split artifact camera entry is missing stage_inputs")
|
||||
|
||||
staged_main, staged_wide = stage_inputs(
|
||||
img_q=self._queue_tensors["img_q"],
|
||||
big_img_q=self._queue_tensors["big_img_q"],
|
||||
tfm=self._queue_tensors["tfm"],
|
||||
big_tfm=self._queue_tensors["big_tfm"],
|
||||
frame=self._frame_blob(self._road_key, bufs[self._road_key]),
|
||||
big_frame=self._frame_blob(self._wide_key, bufs[self._wide_key]),
|
||||
)
|
||||
raw_outputs = self._execute_bundle(img=staged_main, big_img=staged_wide, **self._policy_inputs())
|
||||
if not isinstance(raw_outputs, tuple):
|
||||
raw_outputs = (raw_outputs,)
|
||||
return self._merge_policy_outputs(raw_outputs)
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
raise RuntimeError("Combined split runner executes through run_fused()")
|
||||
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import (
|
||||
CUSTOM_MODEL_PATH, NumpyDict, ShapeDict, SliceDict,
|
||||
)
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
|
||||
|
||||
def _tinygrad_imports():
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.device import Device
|
||||
return Tensor, Device
|
||||
|
||||
|
||||
WARP_DEV = os.getenv('WARP_DEV')
|
||||
|
||||
|
||||
class TinygradFusedRunner(ModelRunner):
|
||||
"""Runs a fused warp+vision+policy pkl. Bundle ships one `driving_fused_*` artifact."""
|
||||
|
||||
uses_opencl_warp: bool = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self._constants = SplitModelConstants
|
||||
self._parser = PhaseParser()
|
||||
|
||||
if len(self.models) != 1:
|
||||
raise ValueError(f"fused bundle must have exactly one artifact, got {list(self.models)}")
|
||||
self._model_data = next(iter(self.models.values()))
|
||||
|
||||
pkl_path = os.path.join(CUSTOM_MODEL_PATH, self._model_data.model.artifact.fileName)
|
||||
with open(pkl_path, 'rb') as f:
|
||||
self._fused: dict[Any, Any] = pickle.load(f)
|
||||
|
||||
self._vision_meta = self._fused['metadata']['vision']
|
||||
self._on_meta = self._fused['metadata']['on_policy']
|
||||
self._off_meta = self._fused['metadata']['off_policy']
|
||||
self._run_policy = self._fused['run_policy']
|
||||
self._warp_jits: dict[tuple[int, int], Any] = {k: v for k, v in self._fused.items() if isinstance(k, tuple)}
|
||||
if not self._warp_jits:
|
||||
raise ValueError("fused pkl has no warp JITs")
|
||||
|
||||
self._frame_skip: int = int(self._fused.get('frame_skip', 4))
|
||||
|
||||
self._queues: dict[str, Any] | None = None
|
||||
self._npy_buffers: dict[str, np.ndarray] | None = None
|
||||
self._cam_resolution: tuple[int, int] | None = None
|
||||
self._blob_cache: dict[tuple[str, int], Any] = {}
|
||||
|
||||
def _frame_tensor(self, key, buf):
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
arr = np.frombuffer(buf.data, dtype=np.uint8)
|
||||
ck = (key, arr.ctypes.data)
|
||||
t = self._blob_cache.get(ck)
|
||||
if t is None:
|
||||
t = Tensor.from_blob(arr.ctypes.data, (arr.size,), dtype='uint8', device=Device.DEFAULT)
|
||||
self._blob_cache[ck] = t
|
||||
return t
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
return ['img', 'big_img']
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
return {**self._vision_meta['input_shapes'], **self._on_meta['input_shapes']}
|
||||
|
||||
@property
|
||||
def output_slices(self) -> SliceDict:
|
||||
merged: SliceDict = {}
|
||||
for src in (self._vision_meta['output_slices'], self._on_meta['output_slices'], self._off_meta['output_slices']):
|
||||
merged.update({k: v for k, v in src.items() if k != 'pad'})
|
||||
return merged
|
||||
|
||||
def prepare_inputs(self, imgs_cl, numpy_inputs, frames):
|
||||
raise RuntimeError("fused runner has no OpenCL path; use run_fused()")
|
||||
|
||||
def _ensure_queues(self, cam_w: int, cam_h: int) -> None:
|
||||
if self._queues is not None and self._cam_resolution == (cam_w, cam_h):
|
||||
return
|
||||
if (cam_w, cam_h) not in self._warp_jits:
|
||||
raise RuntimeError(f"no warp JIT for {cam_w}x{cam_h}; have {sorted(self._warp_jits)}")
|
||||
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
img_shape = self._vision_meta['input_shapes']['img']
|
||||
fb = self._on_meta['input_shapes']['features_buffer']
|
||||
dp = self._on_meta['input_shapes']['desire_pulse']
|
||||
n_frames = img_shape[1] // 6
|
||||
img_buf_shape = (self._frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3])
|
||||
|
||||
zeros_u8 = lambda shp: Tensor(np.zeros(shp, dtype=np.uint8), device=Device.DEFAULT).contiguous().realize()
|
||||
zeros_f32 = lambda shp: Tensor(np.zeros(shp, dtype=np.float32), device=Device.DEFAULT).contiguous().realize()
|
||||
|
||||
self._queues = {
|
||||
'img_q': zeros_u8(img_buf_shape),
|
||||
'big_img_q': zeros_u8(img_buf_shape),
|
||||
'feat_q': zeros_f32((self._frame_skip * (fb[1] - 1) + 1, fb[0], fb[2])),
|
||||
'desire_q': zeros_f32((self._frame_skip * dp[1], dp[0], dp[2])),
|
||||
}
|
||||
# shapes must match the captured run_policy JIT inputs
|
||||
on_shapes = self._on_meta['input_shapes']
|
||||
self._npy_buffers = {
|
||||
'desire': np.zeros(dp[2], dtype=np.float32),
|
||||
'traffic_convention': np.zeros(on_shapes['traffic_convention'], dtype=np.float32),
|
||||
'action_t': np.zeros(on_shapes['action_t'], dtype=np.float32),
|
||||
'tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
'big_tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
}
|
||||
self._cam_resolution = (cam_w, cam_h)
|
||||
|
||||
def run_fused(self, bufs: dict, transforms: dict[str, np.ndarray], numpy_inputs: NumpyDict) -> NumpyDict:
|
||||
"""warp + vision + policy in one pass from raw NV12 bufs + transform matrices."""
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
|
||||
main_buf = bufs['img']
|
||||
self._ensure_queues(main_buf.width, main_buf.height)
|
||||
assert self._queues is not None and self._npy_buffers is not None
|
||||
|
||||
desire_key = next((k for k in numpy_inputs if k.startswith('desire')), None)
|
||||
if desire_key is not None:
|
||||
self._npy_buffers['desire'][:] = numpy_inputs[desire_key]
|
||||
if 'traffic_convention' in numpy_inputs:
|
||||
self._npy_buffers['traffic_convention'][:] = numpy_inputs['traffic_convention']
|
||||
if 'action_t' in numpy_inputs:
|
||||
self._npy_buffers['action_t'][:] = numpy_inputs['action_t']
|
||||
self._npy_buffers['tfm'][:] = transforms['img']
|
||||
self._npy_buffers['big_tfm'][:] = transforms['big_img']
|
||||
|
||||
npy = lambda key: Tensor(self._npy_buffers[key], device='NPY')
|
||||
|
||||
# frames go on the compute device to match the captured warp JIT
|
||||
frame = self._frame_tensor('img', bufs['img'])
|
||||
big_frame = self._frame_tensor('big_img', bufs['big_img'])
|
||||
|
||||
warp_jit = self._warp_jits[self._cam_resolution]
|
||||
img, big_img = warp_jit(img_q=self._queues['img_q'], big_img_q=self._queues['big_img_q'],
|
||||
tfm=npy('tfm'), big_tfm=npy('big_tfm'), frame=frame, big_frame=big_frame)
|
||||
|
||||
vision_out_t, on_out_t, off_out_t = self._run_policy(
|
||||
img=img, big_img=big_img, feat_q=self._queues['feat_q'], desire_q=self._queues['desire_q'],
|
||||
desire=npy('desire'), traffic_convention=npy('traffic_convention'), action_t=npy('action_t'))
|
||||
|
||||
# parse each model's output on its own sliced dict; parsing a merged dict
|
||||
# would run parse_dynamic_outputs twice and double-parse plan/lead
|
||||
def _slice(tensor_out, meta) -> NumpyDict:
|
||||
flat = tensor_out.numpy().flatten()
|
||||
return {k: flat[np.newaxis, sl] for k, sl in meta['output_slices'].items() if k != 'pad'}
|
||||
|
||||
parsed: NumpyDict = {}
|
||||
parsed.update(self._parser.parse_vision_outputs(_slice(vision_out_t, self._vision_meta)))
|
||||
parsed.update(self._parser.parse_policy_outputs(_slice(off_out_t, self._off_meta)))
|
||||
parsed.update(self._parser.parse_policy_outputs(_slice(on_out_t, self._on_meta)))
|
||||
return parsed
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
raise RuntimeError("fused path goes through run_fused(), not _run_model()")
|
||||
@@ -0,0 +1,60 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC
|
||||
from collections.abc import Callable
|
||||
|
||||
import numpy as np
|
||||
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelType, NumpyDict
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import RunnerRoot
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.parser import ArchiveParser, PhaseParser
|
||||
|
||||
|
||||
class _ParserRole(RunnerRoot, ABC):
|
||||
def _bind_parser_role(self,
|
||||
selector: int,
|
||||
parser_builder: Callable[[], object],
|
||||
projector: Callable[[object, NumpyDict], NumpyDict]) -> None:
|
||||
parser = parser_builder()
|
||||
self.parser_method_dict[selector] = lambda model_blob: projector(parser, self._slice_outputs(model_blob))
|
||||
|
||||
|
||||
def _phase_policy(parser: PhaseParser, sliced_outputs: NumpyDict) -> NumpyDict:
|
||||
return parser.parse_policy_outputs(sliced_outputs)
|
||||
|
||||
|
||||
def _phase_vision(parser: PhaseParser, sliced_outputs: NumpyDict) -> NumpyDict:
|
||||
return parser.parse_vision_outputs(sliced_outputs)
|
||||
|
||||
|
||||
def _archive_combined(parser: ArchiveParser, sliced_outputs: NumpyDict) -> NumpyDict:
|
||||
return parser.parse_outputs(sliced_outputs)
|
||||
|
||||
|
||||
class OffPolicyTinygrad(_ParserRole, ABC):
|
||||
def __init__(self):
|
||||
self._bind_parser_role(ModelType.offPolicy, PhaseParser, _phase_policy)
|
||||
|
||||
|
||||
class OnPolicyTinygrad(_ParserRole, ABC):
|
||||
def __init__(self):
|
||||
self._bind_parser_role(ModelType.onPolicy, PhaseParser, _phase_policy)
|
||||
|
||||
|
||||
class PolicyTinygrad(_ParserRole, ABC):
|
||||
def __init__(self):
|
||||
self._bind_parser_role(ModelType.policy, PhaseParser, _phase_policy)
|
||||
|
||||
|
||||
class VisionTinygrad(_ParserRole, ABC):
|
||||
def __init__(self):
|
||||
self._bind_parser_role(ModelType.vision, PhaseParser, _phase_vision)
|
||||
|
||||
|
||||
class SupercomboTinygrad(_ParserRole, ABC):
|
||||
def __init__(self):
|
||||
self._bind_parser_role(ModelType.supercombo, ArchiveParser, _archive_combined)
|
||||
@@ -0,0 +1,344 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import math
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from openpilot.common.params import Params
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import CUSTOM_MODEL_PATH, NumpyDict, ShapeDict, SliceDict
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
|
||||
|
||||
def _tinygrad_imports():
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.device import Device
|
||||
return Tensor, Device
|
||||
|
||||
|
||||
def _captured_queue_depth(warp_jit: Any) -> int | None:
|
||||
captured = getattr(warp_jit, "captured", None)
|
||||
infos = getattr(captured, "expected_input_info", None)
|
||||
if not infos or len(infos) < 2:
|
||||
return None
|
||||
|
||||
view_repr = repr(infos[1][0])
|
||||
dims = [int(val) for val in re.findall(r"arg=(\d+)", view_repr)]
|
||||
return dims[0] if len(dims) >= 4 else None
|
||||
|
||||
|
||||
def _captured_devices(warp_jit: Any) -> set[str]:
|
||||
captured = getattr(warp_jit, "captured", None)
|
||||
infos = getattr(captured, "expected_input_info", None)
|
||||
if not infos:
|
||||
return set()
|
||||
|
||||
devices: set[str] = set()
|
||||
for info in infos:
|
||||
if isinstance(info, tuple) and len(info) >= 4 and isinstance(info[3], str):
|
||||
devices.add(info[3])
|
||||
return devices
|
||||
|
||||
|
||||
def _captured_expected_names(jit_obj: Any) -> list[str]:
|
||||
captured = getattr(jit_obj, "captured", None)
|
||||
names = getattr(captured, "expected_names", None)
|
||||
return list(names) if names else []
|
||||
|
||||
|
||||
def _file_sha256(path: str) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _is_jit_arg_mismatch(err: BaseException) -> bool:
|
||||
return "args mismatch in JIT" in str(err)
|
||||
|
||||
|
||||
class TinygradSupercomboRunner(ModelRunner):
|
||||
"""Runs a single combined supercombo pkl. Bundle ships one `driving_supercombo_*` artifact."""
|
||||
|
||||
uses_opencl_warp: bool = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._constants = SplitModelConstants
|
||||
self._parser = PhaseParser()
|
||||
|
||||
if len(self.models) != 1:
|
||||
raise ValueError(f"supercombo bundle must have exactly one artifact, got {list(self.models)}")
|
||||
self._model_data = next(iter(self.models.values()))
|
||||
|
||||
pkl_path = os.path.join(CUSTOM_MODEL_PATH, self._model_data.model.artifact.fileName)
|
||||
self._pkl_path = pkl_path
|
||||
self._expected_sha256 = getattr(getattr(self._model_data.model.artifact, "downloadUri", None), "sha256", "") or ""
|
||||
self._verify_artifact_file()
|
||||
with open(pkl_path, 'rb') as f:
|
||||
self._m: dict[Any, Any] = pickle.load(f)
|
||||
|
||||
self._meta = self._m['metadata']
|
||||
self._ish = self._meta['input_shapes']
|
||||
self._slices = {k: v for k, v in self._meta['output_slices'].items() if k != 'pad'}
|
||||
self._hidden_slice = self._meta['output_slices']['hidden_state']
|
||||
self._run_policy = self._m['run_policy']
|
||||
self._warp_jits: dict[tuple[int, int], Any] = {k: v for k, v in self._m.items() if isinstance(k, tuple)}
|
||||
if not self._warp_jits:
|
||||
raise ValueError("supercombo pkl has no warp JITs")
|
||||
self._frame_skip = int(self._m.get('frame_skip', 4))
|
||||
self._validate_warp_jits(pkl_path)
|
||||
self._validate_jit_names()
|
||||
|
||||
self._queues: dict[str, Any] | None = None
|
||||
self._npy: dict[str, np.ndarray] | None = None
|
||||
self._cam: tuple[int, int] | None = None
|
||||
self._prev_desire = np.zeros(self._ish['desire_pulse'][2], dtype=np.float32)
|
||||
self._blob_cache: dict[tuple[str, int], Any] = {}
|
||||
|
||||
def _verify_artifact_file(self) -> None:
|
||||
if not self._expected_sha256:
|
||||
return
|
||||
|
||||
actual_sha256 = _file_sha256(self._pkl_path)
|
||||
if actual_sha256 == self._expected_sha256:
|
||||
return
|
||||
|
||||
try:
|
||||
os.remove(self._pkl_path)
|
||||
except OSError:
|
||||
pass
|
||||
redownload_msg = self._schedule_active_bundle_redownload()
|
||||
|
||||
raise RuntimeError(
|
||||
"supercombo artifact SHA mismatch: "
|
||||
f"expected {self._expected_sha256}, got {actual_sha256} for {self._pkl_path}. "
|
||||
f"Deleted the stale cached file{redownload_msg}."
|
||||
)
|
||||
|
||||
def _validate_warp_jits(self, pkl_path: str) -> None:
|
||||
img = self._ish['img']
|
||||
n_frames = img[1] // 6
|
||||
expected_depth = self._frame_skip * (n_frames - 1) + 1
|
||||
expected_device = os.getenv('DEV')
|
||||
|
||||
mismatches: list[str] = []
|
||||
for cam, warp_jit in sorted(self._warp_jits.items()):
|
||||
captured_depth = _captured_queue_depth(warp_jit)
|
||||
captured_devices = _captured_devices(warp_jit)
|
||||
if captured_depth is not None and captured_depth != expected_depth:
|
||||
mismatches.append(
|
||||
f"{cam[0]}x{cam[1]} queue-depth captured={captured_depth} expected={expected_depth}"
|
||||
)
|
||||
if expected_device and captured_devices and expected_device not in captured_devices:
|
||||
mismatches.append(
|
||||
f"{cam[0]}x{cam[1]} device captured={sorted(captured_devices)} expected={expected_device}"
|
||||
)
|
||||
|
||||
if mismatches:
|
||||
details = "; ".join(mismatches)
|
||||
raise RuntimeError(
|
||||
"supercombo warp JIT compatibility mismatch: "
|
||||
f"{details}. Bundle {pkl_path} was compiled with the wrong backend, frame_skip, or queue shape; "
|
||||
"re-download or rebuild this model artifact."
|
||||
)
|
||||
|
||||
def _validate_jit_names(self) -> None:
|
||||
expected_warp_names = ['big_frame', 'big_tfm', 'frame', 'tfm']
|
||||
expected_policy_names = ['big_img_q', 'desire_q', 'feat_q', 'img_q', 'packed_npy_inputs', 'warped']
|
||||
|
||||
mismatches: list[str] = []
|
||||
|
||||
policy_names = sorted(_captured_expected_names(self._run_policy))
|
||||
if policy_names and policy_names != expected_policy_names:
|
||||
mismatches.append(f"run_policy captured={policy_names} expected={expected_policy_names}")
|
||||
|
||||
for cam, warp_jit in sorted(self._warp_jits.items()):
|
||||
warp_names = sorted(_captured_expected_names(warp_jit))
|
||||
if warp_names and warp_names != expected_warp_names:
|
||||
mismatches.append(f"{cam[0]}x{cam[1]} warp captured={warp_names} expected={expected_warp_names}")
|
||||
|
||||
if mismatches:
|
||||
details = "; ".join(mismatches)
|
||||
actual_sha = None
|
||||
try:
|
||||
actual_sha = _file_sha256(self._pkl_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if actual_sha and self._expected_sha256 and actual_sha != self._expected_sha256:
|
||||
try:
|
||||
os.remove(self._pkl_path)
|
||||
except OSError:
|
||||
pass
|
||||
redownload_msg = self._schedule_active_bundle_redownload()
|
||||
raise RuntimeError(
|
||||
"supercombo artifact contract mismatch with stale cached SHA: "
|
||||
f"{details}. Expected SHA {self._expected_sha256}, got {actual_sha}. "
|
||||
f"Deleted the stale cached file{redownload_msg}."
|
||||
)
|
||||
|
||||
raise RuntimeError(
|
||||
"supercombo artifact JIT argument mismatch: "
|
||||
f"{details}. This model file does not match the current IQPilot runtime contract. "
|
||||
"Re-download or rebuild this model artifact."
|
||||
)
|
||||
|
||||
def _handle_runtime_jit_mismatch(self, err: BaseException) -> None:
|
||||
if not _is_jit_arg_mismatch(err):
|
||||
raise err
|
||||
|
||||
actual_sha = None
|
||||
try:
|
||||
actual_sha = _file_sha256(self._pkl_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if actual_sha and self._expected_sha256 and actual_sha != self._expected_sha256:
|
||||
try:
|
||||
os.remove(self._pkl_path)
|
||||
except OSError:
|
||||
pass
|
||||
redownload_msg = self._schedule_active_bundle_redownload()
|
||||
raise RuntimeError(
|
||||
"supercombo artifact runtime JIT mismatch with stale cached SHA: "
|
||||
f"expected {self._expected_sha256}, got {actual_sha} for {self._pkl_path}. "
|
||||
f"Deleted the stale cached file{redownload_msg}."
|
||||
) from err
|
||||
|
||||
raise RuntimeError(
|
||||
"supercombo artifact runtime JIT mismatch: "
|
||||
f"{err}. This model file does not match the current IQPilot runtime contract. "
|
||||
"Re-download or rebuild this model artifact."
|
||||
) from err
|
||||
|
||||
def _schedule_active_bundle_redownload(self) -> str:
|
||||
try:
|
||||
params = Params()
|
||||
active_bundle = params.get("ModelManager_ActiveBundle") or {}
|
||||
index = active_bundle.get("index") if isinstance(active_bundle, dict) else None
|
||||
if isinstance(index, str) and index.isdigit():
|
||||
index = int(index)
|
||||
if isinstance(index, int) and index >= 0:
|
||||
params.put("ModelManager_DownloadIndex", str(index))
|
||||
params.remove("ModelRunnerTypeCache")
|
||||
return "; scheduled automatic re-download of the active model"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return "; unable to schedule automatic re-download"
|
||||
|
||||
def _frame_tensor(self, key: str, buf):
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
arr = np.frombuffer(buf.data, dtype=np.uint8)
|
||||
ck = (key, arr.ctypes.data)
|
||||
t = self._blob_cache.get(ck)
|
||||
if t is None:
|
||||
t = Tensor.from_blob(arr.ctypes.data, (arr.size,), dtype='uint8', device=Device.DEFAULT)
|
||||
self._blob_cache[ck] = t
|
||||
return t
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
return ['img', 'big_img']
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
return dict(self._ish)
|
||||
|
||||
@property
|
||||
def output_slices(self) -> SliceDict:
|
||||
return dict(self._slices)
|
||||
|
||||
def prepare_inputs(self, imgs_cl, numpy_inputs, frames):
|
||||
raise RuntimeError("supercombo runner has no OpenCL path; use run_fused()")
|
||||
|
||||
def _ensure_queues(self, cam_w: int, cam_h: int) -> None:
|
||||
if self._queues is not None and self._cam == (cam_w, cam_h):
|
||||
return
|
||||
if (cam_w, cam_h) not in self._warp_jits:
|
||||
raise RuntimeError(f"no warp JIT for {cam_w}x{cam_h}; have {sorted(self._warp_jits)}")
|
||||
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
fs = self._frame_skip
|
||||
img = self._ish['img']
|
||||
n_frames = img[1] // 6
|
||||
img_buf = (fs * (n_frames - 1) + 1, 6, img[2], img[3])
|
||||
fb = self._ish['features_buffer']
|
||||
dp = self._ish['desire_pulse']
|
||||
tc = self._ish['traffic_convention']
|
||||
at = self._ish['action_t']
|
||||
|
||||
zeros_u8 = lambda s: Tensor(np.zeros(s, dtype=np.uint8), device=Device.DEFAULT).contiguous().realize()
|
||||
zeros_f32 = lambda s: Tensor(np.zeros(s, dtype=np.float32), device=Device.DEFAULT).contiguous().realize()
|
||||
|
||||
# packed npy block (single NPY tensor, mutated in place via views): order matches run_policy.split
|
||||
shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at), 'prev_feat': (fb[0], fb[2])}
|
||||
sizes = [math.prod(s) for s in shapes.values()]
|
||||
packed = np.zeros(sum(sizes), dtype=np.float32)
|
||||
views = {k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed, np.cumsum(sizes[:-1])), strict=True)}
|
||||
|
||||
self._npy = {'tfm': np.zeros((3, 3), dtype=np.float32), 'big_tfm': np.zeros((3, 3), dtype=np.float32), **views}
|
||||
self._queues = {
|
||||
'img_q': zeros_u8(img_buf),
|
||||
'big_img_q': zeros_u8(img_buf),
|
||||
'feat_q': zeros_f32((fs * fb[1], fb[0], fb[2])),
|
||||
'desire_q': zeros_f32((fs * dp[1], dp[0], dp[2])),
|
||||
'tfm': Tensor(self._npy['tfm'], device='NPY'),
|
||||
'big_tfm': Tensor(self._npy['big_tfm'], device='NPY'),
|
||||
'packed_npy_inputs': Tensor(packed, device='NPY'),
|
||||
}
|
||||
self._cam = (cam_w, cam_h)
|
||||
|
||||
def run_fused(self, bufs: dict, transforms: dict[str, np.ndarray], numpy_inputs: NumpyDict) -> NumpyDict:
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
main_buf = bufs['img']
|
||||
self._ensure_queues(main_buf.width, main_buf.height)
|
||||
assert self._queues is not None and self._npy is not None
|
||||
|
||||
self._npy['tfm'][:] = transforms['img']
|
||||
self._npy['big_tfm'][:] = transforms['big_img']
|
||||
|
||||
desire_key = next((k for k in numpy_inputs if k.startswith('desire')), None)
|
||||
cur = numpy_inputs[desire_key].copy() if desire_key is not None else np.zeros_like(self._prev_desire)
|
||||
cur[0] = 0
|
||||
self._npy['desire'][:] = np.where(cur - self._prev_desire > .99, cur, 0)
|
||||
self._prev_desire[:] = cur
|
||||
if 'traffic_convention' in numpy_inputs:
|
||||
self._npy['traffic_convention'][:] = numpy_inputs['traffic_convention']
|
||||
if 'action_t' in numpy_inputs:
|
||||
self._npy['action_t'][:] = numpy_inputs['action_t']
|
||||
# self._npy['prev_feat'] holds last frame's hidden_state (zeros on the first frame)
|
||||
|
||||
frame = self._frame_tensor('img', bufs['img'])
|
||||
big_frame = self._frame_tensor('big_img', bufs['big_img'])
|
||||
|
||||
warp = self._warp_jits[self._cam]
|
||||
try:
|
||||
warped = warp(tfm=self._queues['tfm'], big_tfm=self._queues['big_tfm'], frame=frame, big_frame=big_frame)
|
||||
out, = self._run_policy(warped=warped, img_q=self._queues['img_q'], big_img_q=self._queues['big_img_q'],
|
||||
feat_q=self._queues['feat_q'], desire_q=self._queues['desire_q'],
|
||||
packed_npy_inputs=self._queues['packed_npy_inputs'])
|
||||
except Exception as err:
|
||||
self._handle_runtime_jit_mismatch(err)
|
||||
raise
|
||||
flat = out.numpy().flatten()
|
||||
|
||||
# feed hidden_state back as prev_feat for the next frame
|
||||
self._npy['prev_feat'][:] = flat[self._hidden_slice].reshape(self._npy['prev_feat'].shape)
|
||||
|
||||
sliced = {k: flat[np.newaxis, sl] for k, sl in self._slices.items()}
|
||||
return self._parser.parse_vision_outputs(sliced) # single-pass; parse_outputs double-parses a combined dict
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
raise RuntimeError("supercombo path goes through run_fused(), not _run_model()")
|
||||
@@ -0,0 +1,186 @@
|
||||
"""
|
||||
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
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import (
|
||||
CLMemDict,
|
||||
CUSTOM_MODEL_PATH,
|
||||
FrameDict,
|
||||
ModelType,
|
||||
NumpyDict,
|
||||
ShapeDict,
|
||||
SliceDict,
|
||||
)
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.model_types import (
|
||||
OffPolicyTinygrad,
|
||||
OnPolicyTinygrad,
|
||||
PolicyTinygrad,
|
||||
SupercomboTinygrad,
|
||||
VisionTinygrad,
|
||||
)
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.config import ModelConstants
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.runtime.tinygrad import qcom_tensor_from_opencl_address
|
||||
from openpilot.system.hardware import TICI
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TensorShapePlan:
|
||||
dtype: object
|
||||
device: str
|
||||
|
||||
|
||||
def _artifact_path(filename: str) -> str:
|
||||
return f"{CUSTOM_MODEL_PATH}/{filename}"
|
||||
|
||||
|
||||
def _load_program_blob(filename: str):
|
||||
with open(_artifact_path(filename), "rb") as artifact:
|
||||
try:
|
||||
return pickle.load(artifact)
|
||||
except FileNotFoundError as exc:
|
||||
assert "/dev/kgsl-3d0" not in str(exc), "Model was built on C3 or C3X, but is being loaded on PC"
|
||||
raise
|
||||
|
||||
|
||||
def _compile_input_plan(captured) -> dict[str, _TensorShapePlan]:
|
||||
plan: dict[str, _TensorShapePlan] = {}
|
||||
for name, info in zip(captured.expected_names, captured.expected_input_info, strict=True):
|
||||
plan[name] = _TensorShapePlan(dtype=info[2], device=info[3])
|
||||
return plan
|
||||
|
||||
|
||||
def _merge_step_outputs(output_groups: list[NumpyDict]) -> NumpyDict:
|
||||
stitched: NumpyDict = {}
|
||||
for payload in output_groups:
|
||||
stitched.update(payload)
|
||||
if "planplus" in stitched and "plan" in stitched:
|
||||
stitched["plan"] = stitched["plan"] + stitched["planplus"]
|
||||
return stitched
|
||||
|
||||
|
||||
class TinygradRunner(ModelRunner, SupercomboTinygrad, PolicyTinygrad, VisionTinygrad, OffPolicyTinygrad, OnPolicyTinygrad):
|
||||
def __init__(self, model_type: int = ModelType.supercombo):
|
||||
ModelRunner.__init__(self)
|
||||
for initializer in (SupercomboTinygrad, PolicyTinygrad, VisionTinygrad, OffPolicyTinygrad, OnPolicyTinygrad):
|
||||
initializer.__init__(self)
|
||||
|
||||
self._constants = ModelConstants
|
||||
self._model_data = self.models.get(model_type)
|
||||
if self._model_data is None or self._model_data.model is None:
|
||||
raise ValueError(f"Model data for type {model_type} not available.")
|
||||
|
||||
asset_name = self._model_data.model.artifact.fileName
|
||||
assert asset_name.endswith("_tinygrad.pkl"), f"Invalid model file {asset_name} for TinygradRunner"
|
||||
|
||||
self.model_run = _load_program_blob(asset_name)
|
||||
self._input_plan = _compile_input_plan(self.model_run.captured)
|
||||
self.input_to_dtype = {name: spec.dtype for name, spec in self._input_plan.items()}
|
||||
self.input_to_device = {name: spec.device for name, spec in self._input_plan.items()}
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
return [stream_name for stream_name in self.input_shapes if "img" in stream_name]
|
||||
|
||||
def _attach_vision_tensor(self, stream_name: str, frame_buffers: CLMemDict, frame_views: FrameDict) -> None:
|
||||
spec = self._input_plan[stream_name]
|
||||
frame_buffer = frame_buffers[stream_name]
|
||||
if TICI:
|
||||
self.inputs[stream_name] = qcom_tensor_from_opencl_address(frame_buffer.mem_address,
|
||||
self.input_shapes[stream_name],
|
||||
dtype=spec.dtype)
|
||||
return
|
||||
|
||||
mirrored = frame_views[stream_name].as_numpy(frame_buffer).reshape(self.input_shapes[stream_name])
|
||||
self.inputs[stream_name] = Tensor(mirrored, device=spec.device, dtype=spec.dtype).realize()
|
||||
|
||||
def _attach_state_tensor(self, tensor_name: str, tensor_value: np.ndarray) -> None:
|
||||
spec = self._input_plan[tensor_name]
|
||||
self.inputs[tensor_name] = Tensor(tensor_value, device=spec.device, dtype=spec.dtype).realize()
|
||||
|
||||
def prepare_vision_inputs(self, imgs_cl: CLMemDict, frames: FrameDict):
|
||||
for stream_name in imgs_cl:
|
||||
if stream_name not in self.inputs or not TICI:
|
||||
self._attach_vision_tensor(stream_name, imgs_cl, frames)
|
||||
|
||||
def prepare_policy_inputs(self, numpy_inputs: NumpyDict):
|
||||
for tensor_name, tensor_value in numpy_inputs.items():
|
||||
self._attach_state_tensor(tensor_name, tensor_value)
|
||||
|
||||
def prepare_inputs(self, imgs_cl: CLMemDict, numpy_inputs: NumpyDict, frames: FrameDict) -> dict:
|
||||
self.prepare_vision_inputs(imgs_cl, frames)
|
||||
self.prepare_policy_inputs(numpy_inputs)
|
||||
return self.inputs
|
||||
|
||||
def _parse_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
||||
if self._model_data is None:
|
||||
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
||||
return self.parser_method_dict[self._model_data.model.type.raw](model_outputs)
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
raw_output = self.model_run(**self.inputs).numpy().reshape(-1)
|
||||
return self._parse_outputs(raw_output)
|
||||
|
||||
|
||||
class TinygradSplitRunner(ModelRunner):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.is_20hz_3d = True
|
||||
self._constants = SplitModelConstants
|
||||
self.vision_runner = TinygradRunner(ModelType.vision)
|
||||
self.policy_runner = TinygradRunner(ModelType.policy) if self.models.get(ModelType.policy) else None
|
||||
self.off_policy_runner = TinygradRunner(ModelType.offPolicy) if self.models.get(ModelType.offPolicy) else None
|
||||
self.on_policy_runner = TinygradRunner(ModelType.onPolicy) if self.models.get(ModelType.onPolicy) else None
|
||||
|
||||
def _policy_units(self) -> list[TinygradRunner]:
|
||||
return [runner for runner in (self.policy_runner, self.off_policy_runner, self.on_policy_runner) if runner is not None]
|
||||
|
||||
def run_vision(self) -> NumpyDict:
|
||||
return self.vision_runner.run_model()
|
||||
|
||||
def run_policy(self) -> NumpyDict:
|
||||
return _merge_step_outputs([runner.run_model() for runner in self._policy_units()])
|
||||
|
||||
def refresh_policy_features(self, features_buffer: np.ndarray) -> None:
|
||||
for runner in self._policy_units():
|
||||
if "features_buffer" in runner._input_plan:
|
||||
runner._attach_state_tensor("features_buffer", features_buffer)
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
return _merge_step_outputs([self.run_vision(), self.run_policy()])
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
return list(self.vision_runner.vision_input_names)
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
composite: ShapeDict = dict(self.vision_runner.input_shapes)
|
||||
for runner in self._policy_units():
|
||||
composite.update(runner.input_shapes)
|
||||
return composite
|
||||
|
||||
@property
|
||||
def output_slices(self) -> SliceDict:
|
||||
composite: SliceDict = dict(self.vision_runner.output_slices)
|
||||
for runner in self._policy_units():
|
||||
composite.update(runner.output_slices)
|
||||
return composite
|
||||
|
||||
def prepare_inputs(self, imgs_cl: CLMemDict, numpy_inputs: NumpyDict, frames: FrameDict) -> dict:
|
||||
self.vision_runner.prepare_vision_inputs(imgs_cl, frames)
|
||||
assembled_inputs = dict(self.vision_runner.inputs)
|
||||
for runner in self._policy_units():
|
||||
runner.prepare_policy_inputs(numpy_inputs)
|
||||
assembled_inputs.update(runner.inputs)
|
||||
self.inputs = assembled_inputs
|
||||
return assembled_inputs
|
||||
96
iqpilot/selfdrive/iqmodeld/models/split_model_constants.py
Normal file
96
iqpilot/selfdrive/iqmodeld/models/split_model_constants.py
Normal file
@@ -0,0 +1,96 @@
|
||||
# openpilot model I/O constants (comma.ai, MIT — see LICENSE)
|
||||
import numpy as np
|
||||
|
||||
|
||||
def index_function(idx, max_val=192, max_idx=32):
|
||||
return max_val * ((idx/max_idx)**2)
|
||||
|
||||
|
||||
class SplitModelConstants:
|
||||
# time and distance indices
|
||||
IDX_N = 33
|
||||
T_IDXS = [index_function(idx, max_val=10.0) for idx in range(IDX_N)]
|
||||
X_IDXS = [index_function(idx, max_val=192.0) for idx in range(IDX_N)]
|
||||
LEAD_T_IDXS = [0., 2., 4., 6., 8., 10.]
|
||||
LEAD_T_OFFSETS = [0., 2., 4.]
|
||||
META_T_IDXS = [2., 4., 6., 8., 10.]
|
||||
|
||||
# split-model temporal / history run parameters
|
||||
MODEL_FREQ = 20
|
||||
HISTORY_FREQ = 5
|
||||
HISTORY_LEN_SECONDS = 5
|
||||
TEMPORAL_SKIP = MODEL_FREQ // HISTORY_FREQ
|
||||
FULL_HISTORY_BUFFER_LEN = MODEL_FREQ * HISTORY_LEN_SECONDS
|
||||
INPUT_HISTORY_BUFFER_LEN = HISTORY_FREQ * HISTORY_LEN_SECONDS
|
||||
|
||||
FEATURE_LEN = 512
|
||||
|
||||
DESIRE_LEN = 8
|
||||
TRAFFIC_CONVENTION_LEN = 2
|
||||
LAT_PLANNER_STATE_LEN = 4
|
||||
LATERAL_CONTROL_PARAMS_LEN = 2
|
||||
PREV_DESIRED_CURV_LEN = 1
|
||||
|
||||
# model outputs constants
|
||||
FCW_THRESHOLDS_5MS2 = np.array([.05, .05, .15, .15, .15], dtype=np.float32)
|
||||
FCW_THRESHOLDS_3MS2 = np.array([.7, .7], dtype=np.float32)
|
||||
FCW_5MS2_PROBS_WIDTH = 5
|
||||
FCW_3MS2_PROBS_WIDTH = 2
|
||||
|
||||
DISENGAGE_WIDTH = 5
|
||||
POSE_WIDTH = 6
|
||||
WIDE_FROM_DEVICE_WIDTH = 3
|
||||
LEAD_WIDTH = 4
|
||||
LANE_LINES_WIDTH = 2
|
||||
ROAD_EDGES_WIDTH = 2
|
||||
PLAN_WIDTH = 15
|
||||
DESIRE_PRED_WIDTH = 8
|
||||
LAT_PLANNER_SOLUTION_WIDTH = 4
|
||||
DESIRED_CURV_WIDTH = 1
|
||||
ACTION_WIDTH = 2
|
||||
|
||||
NUM_LANE_LINES = 4
|
||||
NUM_ROAD_EDGES = 2
|
||||
|
||||
LEAD_TRAJ_LEN = 6
|
||||
DESIRE_PRED_LEN = 4
|
||||
|
||||
PLAN_MHP_N = 5
|
||||
LEAD_MHP_N = 2
|
||||
PLAN_MHP_SELECTION = 1
|
||||
LEAD_MHP_SELECTION = 3
|
||||
|
||||
FCW_THRESHOLD_5MS2_HIGH = 0.15
|
||||
FCW_THRESHOLD_5MS2_LOW = 0.05
|
||||
FCW_THRESHOLD_3MS2 = 0.7
|
||||
|
||||
CONFIDENCE_BUFFER_LEN = 5
|
||||
RYG_GREEN = 0.01165
|
||||
RYG_YELLOW = 0.06157
|
||||
|
||||
POLY_PATH_DEGREE = 4
|
||||
|
||||
|
||||
# model outputs slices
|
||||
class Plan:
|
||||
POSITION = slice(0, 3)
|
||||
VELOCITY = slice(3, 6)
|
||||
ACCELERATION = slice(6, 9)
|
||||
T_FROM_CURRENT_EULER = slice(9, 12)
|
||||
ORIENTATION_RATE = slice(12, 15)
|
||||
|
||||
|
||||
class Meta:
|
||||
ENGAGED = slice(0, 1)
|
||||
# next 2, 4, 6, 8, 10 seconds
|
||||
GAS_DISENGAGE = slice(1, 31, 6)
|
||||
BRAKE_DISENGAGE = slice(2, 31, 6)
|
||||
STEER_OVERRIDE = slice(3, 31, 6)
|
||||
HARD_BRAKE_3 = slice(4, 31, 6)
|
||||
HARD_BRAKE_4 = slice(5, 31, 6)
|
||||
HARD_BRAKE_5 = slice(6, 31, 6)
|
||||
# next 0, 2, 4, 6, 8, 10 seconds
|
||||
GAS_PRESS = slice(31, 55, 4)
|
||||
BRAKE_PRESS = slice(32, 55, 4)
|
||||
LEFT_BLINKER = slice(33, 55, 4)
|
||||
RIGHT_BLINKER = slice(34, 55, 4)
|
||||
Reference in New Issue
Block a user