101 lines
2.7 KiB
Python
101 lines
2.7 KiB
Python
"""
|
|
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 iqpilot.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
|