diff --git a/openpilot/common/params_keys.h b/openpilot/common/params_keys.h index 111099fce9..4d8ffb64eb 100644 --- a/openpilot/common/params_keys.h +++ b/openpilot/common/params_keys.h @@ -195,11 +195,10 @@ inline static std::unordered_map keys = { // Model Manager params {"ModelManager_ActiveBundle", {PERSISTENT, JSON}}, - {"ModelManager_ActiveJson", {CLEAR_ON_MANAGER_START, STRING}}, - {"ModelManager_PrevBundle", {PERSISTENT, JSON}}, - {"ModelManager_PrevBundle_USBGPU", {PERSISTENT, JSON}}, + {"ModelManager_ActiveBundleUSBGPU", {PERSISTENT, JSON}}, + {"ModelManager_ActiveJson", {CLEAR_ON_MANAGER_START, JSON}}, {"ModelManager_ClearCache", {CLEAR_ON_MANAGER_START, BOOL}}, - {"ModelManager_DownloadIndex", {CLEAR_ON_MANAGER_START | CLEAR_ON_ONROAD_TRANSITION, INT}}, + {"ModelManager_DownloadRef", {CLEAR_ON_MANAGER_START | CLEAR_ON_ONROAD_TRANSITION, STRING}}, {"ModelManager_Favs", {PERSISTENT | BACKUP, STRING}}, {"ModelManager_LastSyncTime", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, INT, "0"}}, {"ModelManager_LastSyncTime_USBGPU", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, INT, "0"}}, diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index 9f3d709537..d180012279 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -91,7 +91,7 @@ class ModelState(ModelStateBase): if env_pkl and os.path.exists(env_pkl): model_bundle = None else: - model_bundle = get_active_bundle() + model_bundle = get_active_bundle(usbgpu=usbgpu) self.generation = model_bundle.generation if model_bundle is not None else None overrides = {override.key: override.value for override in model_bundle.overrides} if model_bundle else {} diff --git a/openpilot/sunnypilot/models/fetcher.py b/openpilot/sunnypilot/models/fetcher.py index c9e86edd0c..f7a4813161 100644 --- a/openpilot/sunnypilot/models/fetcher.py +++ b/openpilot/sunnypilot/models/fetcher.py @@ -141,41 +141,57 @@ class ModelFetcher: MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v21.json" MODEL_URL_USBGPU = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_usbgpu_v22.json" + MODEL_SOURCES = { + "qcom": (MODEL_URL, ""), + "usbgpu": (MODEL_URL_USBGPU, "_USBGPU"), + } + def __init__(self, params: Params): self.params = params self.model_parser = ModelParser() - self._is_usbgpu: bool | None = None - self.model_cache = ModelCache(params) + self._active_json_published = False + self.model_caches = { + source: ModelCache(params, suffix=suffix) + for source, (_, suffix) in self.MODEL_SOURCES.items() + } self.model_url = self.MODEL_URL + self._update_model_source() - def _update_model_source(self, chestnut_present: bool) -> None: - """Updates what json to use based on chestnut hardware presence via deviceState""" - is_usbgpu = chestnut_present - if is_usbgpu != self._is_usbgpu: - self._is_usbgpu = is_usbgpu - self.model_cache = ModelCache(self.params, suffix="_USBGPU" if is_usbgpu else "") - self.model_url = self.MODEL_URL_USBGPU if is_usbgpu else self.MODEL_URL - self.params.put("ModelManager_ActiveJson", self.model_url, block=True) + @staticmethod + def active_source(chestnut_present: bool) -> str: + return "usbgpu" if chestnut_present else "qcom" - def _fetch_and_cache_models(self) -> list[custom.ModelManagerSP.ModelBundle] | None: + def _update_model_source(self) -> None: + """Publishes the manifest URLs for both sources""" + if not self._active_json_published: + self._active_json_published = True + self.params.put("ModelManager_ActiveJson", { + "qcom": self.MODEL_URL, + "usbgpu": self.MODEL_URL_USBGPU, + }, block=True) + + def _fetch_and_cache_models(self, source: str) -> list[custom.ModelManagerSP.ModelBundle] | None: """Fetches fresh model data from remote and updates cache. Returns None on transport errors. Raises on 404 and other fatal HTTP errors. """ + model_url, _ = self.MODEL_SOURCES[source] try: - response = requests.get(self.model_url, timeout=10) + response = requests.get(model_url, timeout=10) # Explicitly handle 404 differently if response.status_code == 404: - cloudlog.error(f"Models URL returned 404 Not Found: {self.model_url}") - raise HTTPError(f"404 Not Found: {self.model_url}", response=response) + cloudlog.error(f"Models URL returned 404 Not Found: {model_url}") + raise HTTPError(f"404 Not Found: {model_url}", response=response) # Raise for any other 4xx/5xx response.raise_for_status() json_data = response.json() - self.model_cache.set(json_data) - cloudlog.debug("Successfully updated models cache") - return self.model_parser.parse_models(json_data) + parsed = self.model_parser.parse_models(json_data) + if parsed: + self.model_caches[source].set(json_data) + cloudlog.debug(f"Successfully updated models cache for {source}") + return parsed except ConnectionError as e: cloudlog.warning(f"DNS/connection error while fetching models: {e}") @@ -188,16 +204,34 @@ class ModelFetcher: return None - def get_available_bundles(self, chestnut_present: bool = False) -> list[custom.ModelManagerSP.ModelBundle]: - """Gets the list of available models, with smart cache handling""" - self._update_model_source(chestnut_present) - cached_data, is_expired = self.model_cache.get() + @staticmethod + def _cache_matches_source(source: str, cached_data: dict) -> bool: + """Confirms a cached manifest contains requested source's models.""" + bundles = cached_data.get("bundles", []) + if source == "usbgpu": + return any(bundle.get("is_big") is True for bundle in bundles) + return not any(bundle.get("is_big") is True for bundle in bundles) + + def _get_source_bundles(self, source: str) -> list[custom.ModelManagerSP.ModelBundle]: + cached_data, is_expired = self.model_caches[source].get() if cached_data and not is_expired: - cloudlog.debug("Using valid cached models data") - return self.model_parser.parse_models(cached_data) + if self._cache_matches_source(source, cached_data): + try: + parsed = self.model_parser.parse_models(cached_data) + except Exception: + cloudlog.warning(f"Failed to parse cached models for {source}; refetching", exc_info=True) + else: + if parsed: + cloudlog.debug(f"Using valid cached models data for source {source}") + return parsed + # a source-matching cache that yields no valid bundles is stale (e.g. an old + # manifest version) - do not trust it, refetch so the source is repopulated + cloudlog.warning(f"Cached models for {source} have no valid bundles; refetching") + else: + cloudlog.warning(f"Cached models for {source} not valid; refetching") - fetched_bundles = self._fetch_and_cache_models() + fetched_bundles = self._fetch_and_cache_models(source) if fetched_bundles is not None: return fetched_bundles @@ -205,14 +239,41 @@ class ModelFetcher: cloudlog.warning("Failed to fetch fresh data and no cache available") cloudlog.warning("Failed to fetch fresh data. Using expired cache as fallback") - return self.model_parser.parse_models(cached_data) + try: + return self.model_parser.parse_models(cached_data) + except Exception: + return [] + + def get_bundles_for_source(self, source: str) -> list[custom.ModelManagerSP.ModelBundle]: + """Gets the list of available models for a specific source, with smart cache handling.""" + if source not in self.MODEL_SOURCES: + cloudlog.warning(f"Unknown model source: {source}") + return [] + return self._get_source_bundles(source) + + +def get_cached_bundles(params: Params, source: str) -> list[custom.ModelManagerSP.ModelBundle]: + """Reads a source's cached manifest from params and parses it into bundles.""" + + if source not in ModelFetcher.MODEL_SOURCES: + cloudlog.warning(f"Unknown model source: {source}") + return [] + _, suffix = ModelFetcher.MODEL_SOURCES[source] + cached_data = params.get(f"ModelManager_ModelsCache{suffix}") + if not cached_data: + return [] + try: + return ModelParser.parse_models(cached_data) + except Exception as e: + cloudlog.warning(f"Failed to parse cached models for source {source}: {e}") + return [] if __name__ == "__main__": from openpilot.selfdrive.modeld.helpers import usbgpu_present params = Params() model_fetcher = ModelFetcher(params) - bundles = model_fetcher.get_available_bundles(chestnut_present=usbgpu_present()) + bundles = model_fetcher.get_bundles_for_source(ModelFetcher.active_source(usbgpu_present())) for bundle in bundles: for model in bundle.models: model_overrides = {override.key: override.value for override in bundle.overrides} diff --git a/openpilot/sunnypilot/models/helpers.py b/openpilot/sunnypilot/models/helpers.py index e33cc445d1..584b97ba8b 100644 --- a/openpilot/sunnypilot/models/helpers.py +++ b/openpilot/sunnypilot/models/helpers.py @@ -16,6 +16,7 @@ from openpilot.common.params import Params from openpilot.common.swaglog import cloudlog from openpilot.sunnypilot.models.constants import Meta, MetaSimPose, MetaTombRaider from openpilot.common.hardware.hw import Paths +from openpilot.selfdrive.modeld.helpers import usbgpu_present # SET ME TO THE EXACT JSON VERSION WE SET IN SUNNYPILOT_MODELS REPO REQUIRED_JSON_VERSION = 18 @@ -24,6 +25,12 @@ CUSTOM_MODEL_PATH = Paths.model_root() METADATA_PATH = Path(__file__).parent / '../models/supercombo_metadata.pkl' ModelManager = custom.ModelManagerSP +ACTIVE_BUNDLE_KEYS = { + "qcom": "ModelManager_ActiveBundle", + "usbgpu": "ModelManager_ActiveBundleUSBGPU", +} +_LAST_VALIDATED_RAW: dict[str, bytes | None] = {} + def _compute_hash(file_path: str) -> str | None: from openpilot.common.file_chunker import open_file_chunked @@ -85,11 +92,11 @@ def _bundle_needs_reset(active_bundle: custom.ModelManagerSP.ModelBundle, availa if available_bundles is not None: matching_bundle = None for bundle in available_bundles: - if active_bundle.ref and bundle.ref: + if getattr(active_bundle, 'ref', None) and getattr(bundle, 'ref', None): if active_bundle.ref == bundle.ref: matching_bundle = bundle break - elif active_bundle.internalName == bundle.internalName: + elif getattr(active_bundle, 'internalName', None) == getattr(bundle, 'internalName', None): matching_bundle = bundle break @@ -97,55 +104,77 @@ def _bundle_needs_reset(active_bundle: custom.ModelManagerSP.ModelBundle, availa return True if active_bundle.minimumSelectorVersion != matching_bundle.minimumSelectorVersion: return True - if active_bundle.runner.raw != matching_bundle.runner.raw: - return True + + active_runner = getattr(active_bundle, 'runner', None) + matching_runner = getattr(matching_bundle, 'runner', None) + if active_runner is not None and matching_runner is not None: + if getattr(active_runner, 'raw', active_runner) != getattr(matching_runner, 'raw', matching_runner): + return True if set(_bundle_artifacts(active_bundle)) != set(_bundle_artifacts(matching_bundle)): return True - # missing files trigger re-download, not selection reset - return False + return not _bundle_is_valid_locally(active_bundle) -def _prev_bundle_key(is_usbgpu: bool) -> str: - return "ModelManager_PrevBundle_USBGPU" if is_usbgpu else "ModelManager_PrevBundle" - - -def validate_active_bundle(params: Params, available_bundles: list[custom.ModelManagerSP.ModelBundle] | None = None, - is_usbgpu: bool = False) -> None: - raw_bundle = params.get("ModelManager_ActiveBundle") - if not raw_bundle: - prev = params.get(_prev_bundle_key(is_usbgpu)) - if prev and (prev_bundle := get_active_bundle(params, raw_bundle_dict=prev)) is not None: - if not _bundle_needs_reset(prev_bundle, available_bundles): - params.put("ModelManager_ActiveBundle", prev, block=True) - return - - active_bundle = get_active_bundle(params, raw_bundle_dict=raw_bundle) - if active_bundle is None or _bundle_needs_reset(active_bundle, available_bundles): - cloudlog.warning("Active model bundle invalid; resetting to default") - params.put(_prev_bundle_key(not is_usbgpu), raw_bundle, block=True) - - prev = params.get(_prev_bundle_key(is_usbgpu)) - if prev and (prev_bundle := get_active_bundle(params, raw_bundle_dict=prev)) is not None: - if not _bundle_needs_reset(prev_bundle, available_bundles): - params.put("ModelManager_ActiveBundle", prev, block=True) - return - - params.remove("ModelManager_ActiveBundle") - params.put("ModelRunnerTypeCache", int(custom.ModelManagerSP.Runner.stock), block=True) - - -def get_active_bundle(params: Params | None = None, raw_bundle_dict: dict | bytes | None = None) -> "custom.ModelManagerSP.ModelBundle | None": - params = params or Params() +def _parse_active_bundle(raw_bundle) -> "custom.ModelManagerSP.ModelBundle | None": try: - active_bundle_dict = raw_bundle_dict if raw_bundle_dict is not None else (params.get("ModelManager_ActiveBundle") or {}) - if isinstance(active_bundle_dict, dict) and active_bundle_dict and is_bundle_version_compatible(active_bundle_dict): - return custom.ModelManagerSP.ModelBundle(**active_bundle_dict) + if isinstance(raw_bundle, dict) and raw_bundle and is_bundle_version_compatible(raw_bundle): + return custom.ModelManagerSP.ModelBundle(**raw_bundle) except Exception: pass return None +def get_selected_bundle(params: Params | None = None, source: str = "qcom") -> "custom.ModelManagerSP.ModelBundle | None": + params = params or Params() + return _parse_active_bundle(params.get(ACTIVE_BUNDLE_KEYS.get(source, "ModelManager_ActiveBundle"))) + + +def get_active_bundle(params: Params | None = None, *, usbgpu: bool | None = None) -> "custom.ModelManagerSP.ModelBundle | None": + params = params or Params() + if usbgpu: + return get_selected_bundle(params, "usbgpu") + else: + return get_selected_bundle(params, "qcom") + + +def resolve_bundle_by_ref( + ref: str, source_bundles: dict[str, list[custom.ModelManagerSP.ModelBundle]], +) -> "tuple[custom.ModelManagerSP.ModelBundle, str] | None": + """Finds the bundle matching a ref across all sources.""" + for source, bundles in source_bundles.items(): + for bundle in bundles: + if bundle.ref == ref: + return bundle, source + return None + + +def _validate_active_bundle(params: Params, source: str, available_bundles: list[custom.ModelManagerSP.ModelBundle] | None = None) -> None: + global _LAST_VALIDATED_RAW + + key = ACTIVE_BUNDLE_KEYS[source] + raw_bundle = params.get(key) + if not raw_bundle: + return + + if _LAST_VALIDATED_RAW.get(key) == raw_bundle: + return + + active_bundle = _parse_active_bundle(raw_bundle) + if active_bundle is None or _bundle_needs_reset(active_bundle, available_bundles): + cloudlog.warning(f"Active model bundle invalid for {source}; resetting to default") + params.remove(key) + params.put("ModelRunnerTypeCache", int(custom.ModelManagerSP.Runner.stock), block=True) + _LAST_VALIDATED_RAW[key] = None + else: + _LAST_VALIDATED_RAW[key] = raw_bundle + + +def validate_active_bundles(params: Params, source_bundles: dict[str, list[custom.ModelManagerSP.ModelBundle]]) -> None: + for source, bundles in source_bundles.items(): + _validate_active_bundle(params, source, bundles) + + def get_active_model_runner(params: Params | None = None, force_check: bool = False) -> int: params = params or Params() cached_runner_type = params.get("ModelRunnerTypeCache") diff --git a/openpilot/sunnypilot/models/manager.py b/openpilot/sunnypilot/models/manager.py index e47cf7536c..bdc6c88741 100644 --- a/openpilot/sunnypilot/models/manager.py +++ b/openpilot/sunnypilot/models/manager.py @@ -17,7 +17,8 @@ from openpilot.common.hardware.hw import Paths from openpilot.cereal import messaging, custom from openpilot.sunnypilot.models.fetcher import ModelFetcher -from openpilot.sunnypilot.models.helpers import get_active_bundle, validate_active_bundle, verify_file +from openpilot.sunnypilot.models.helpers import (ACTIVE_BUNDLE_KEYS, get_active_bundle, get_selected_bundle, + resolve_bundle_by_ref, validate_active_bundles, verify_file) # (connect, read) seconds. read is per-request inactivity, not a total cap DOWNLOAD_TIMEOUT = (30, 30) @@ -31,9 +32,11 @@ class ModelManagerSP: self.model_fetcher = ModelFetcher(self.params) self.pm = messaging.PubMaster(["modelManagerSP"]) self.sm = messaging.SubMaster(["deviceState"]) + self.chestnut_present = False self.available_models: list[custom.ModelManagerSP.ModelBundle] = [] + self.source_models: dict[str, list[custom.ModelManagerSP.ModelBundle]] = {} self.selected_bundle: custom.ModelManagerSP.ModelBundle = None - self.active_bundle: custom.ModelManagerSP.ModelBundle = get_active_bundle(self.params) + self.active_bundle: custom.ModelManagerSP.ModelBundle = get_active_bundle(self.params, usbgpu=self.chestnut_present) self._chunk_size = 128 * 1000 # 128 KB chunks self._download_start_times: dict[str, float] = {} # Track start time per model @@ -77,7 +80,7 @@ class ModelManagerSP: f.write(chunk) bytes_downloaded += len(chunk) - if self.params.get("ModelManager_DownloadIndex") is None: + if self.params.get("ModelManager_DownloadRef") is None: raise Exception("Download cancelled") if total_size > 0: @@ -115,7 +118,7 @@ class ModelManagerSP: for data in response.iter_content(chunk_size=self._chunk_size): f.write(data) chunk_downloaded += len(data) - if self.params.get("ModelManager_DownloadIndex") is None: + if self.params.get("ModelManager_DownloadRef") is None: raise Exception("Download cancelled") intra = chunk_downloaded / max(chunk_size, 1) progress = min(99.0, ((i + intra) / num_chunks) * 100) @@ -217,8 +220,8 @@ class ModelManagerSP: model_manager_state.availableBundles = self.available_models self.pm.send('modelManagerSP', msg) - async def _download_bundle(self, model_bundle: custom.ModelManagerSP.ModelBundle, destination_path: str) -> None: - """Downloads all models in a bundle""" + async def _download_bundle(self, model_bundle: custom.ModelManagerSP.ModelBundle, destination_path: str, source: str) -> None: + """Downloads a bundle and sets it as the active bundle for its source""" self.selected_bundle = model_bundle self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.downloading for model in self.selected_bundle.models: @@ -240,10 +243,9 @@ class ModelManagerSP: seen_artifacts.add(artifact.fileName) await self._process_artifact(artifact, destination_path) - self.active_bundle = self.selected_bundle - self.active_bundle.status = custom.ModelManagerSP.DownloadStatus.downloaded - self.params.put("ModelManager_ActiveBundle", self.active_bundle.to_dict(), block=True) - self.selected_bundle = None + self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.downloaded + self.params.put(ACTIVE_BUNDLE_KEYS[source], model_bundle.to_dict(), block=True) + self.active_bundle = get_active_bundle(self.params, usbgpu=self.chestnut_present) except Exception: if self.selected_bundle is not None: @@ -253,37 +255,32 @@ class ModelManagerSP: finally: self._report_status() - def download(self, model_bundle: custom.ModelManagerSP.ModelBundle, destination_path: str) -> None: + def download(self, model_bundle: custom.ModelManagerSP.ModelBundle, destination_path: str, source: str) -> None: """Main entry point for downloading a model bundle""" - asyncio.run(self._download_bundle(model_bundle, destination_path)) - - BOOT_SETTLE_TICKS = 10 # seconds at 1 Hz before validating active bundle + asyncio.run(self._download_bundle(model_bundle, destination_path, source)) def main_thread(self) -> None: """Main thread for model management""" rk = Ratekeeper(1, print_delay_threshold=None) - boot_ticks = 0 while True: try: self.sm.update(0) - chestnut_present = self.sm['deviceState'].chestnutPresent - self.available_models = self.model_fetcher.get_available_bundles(chestnut_present) - if boot_ticks >= self.BOOT_SETTLE_TICKS: - validate_active_bundle(self.params, self.available_models, is_usbgpu=chestnut_present) - boot_ticks = min(boot_ticks + 1, self.BOOT_SETTLE_TICKS) - self.active_bundle = get_active_bundle(self.params) + self.chestnut_present = self.sm['deviceState'].chestnutPresent + self.source_models = {source: self.model_fetcher.get_bundles_for_source(source) for source in ModelFetcher.MODEL_SOURCES} + self.available_models = self.source_models[ModelFetcher.active_source(self.chestnut_present)] + validate_active_bundles(self.params, self.source_models) + self.active_bundle = get_active_bundle(self.params, usbgpu=self.chestnut_present) - if (index_to_download := self.params.get("ModelManager_DownloadIndex")) is not None: - if self.active_bundle and self.active_bundle.index == index_to_download: - self.params.remove("ModelManager_DownloadIndex") - elif model_to_download := next((model for model in self.available_models if model.index == index_to_download), None): + if (ref_to_download := self.params.get("ModelManager_DownloadRef")) is not None: + if resolved := resolve_bundle_by_ref(ref_to_download, self.source_models): + model_to_download, source = resolved try: - self.download(model_to_download, Paths.model_root()) + self.download(model_to_download, Paths.model_root(), source) except Exception as e: cloudlog.exception(e) finally: - self.params.remove("ModelManager_DownloadIndex") + self.params.remove("ModelManager_DownloadRef") self.selected_bundle = None if self.params.get("ModelManager_ClearCache"): @@ -302,12 +299,14 @@ class ModelManagerSP: Clears the model cache directory of all files except those in the active model bundle. """ - # Get list of files used by active model bundle + # Get list of files used by both slots' selected bundles (either may become + # the truly active bundle depending on hardware availability) active_files = [] - if self.active_bundle is not None: # When the default model is active - for model in self.active_bundle.models: - if hasattr(model, 'artifact') and model.artifact.fileName: - active_files.append(model.artifact.fileName) + for source in ACTIVE_BUNDLE_KEYS: + if selected_bundle := get_selected_bundle(self.params, source): + for model in selected_bundle.models: + if hasattr(model, 'artifact') and model.artifact.fileName: + active_files.append(model.artifact.fileName) # Remove all files except active ones (including their chunk files) model_dir = Paths.model_root() diff --git a/openpilot/sunnypilot/models/tests/test_manager_download.py b/openpilot/sunnypilot/models/tests/test_manager_download.py index 67fb9023af..d98f1fccd1 100644 --- a/openpilot/sunnypilot/models/tests/test_manager_download.py +++ b/openpilot/sunnypilot/models/tests/test_manager_download.py @@ -11,6 +11,7 @@ import http.server import os import tempfile import threading +import time import unittest from typing import Any from unittest import mock @@ -23,6 +24,8 @@ from openpilot.common.test import OpenpilotTestCase from openpilot.common.file_chunker import get_chunk_name, get_manifest_path from openpilot.selfdrive.test.helpers import http_server_context from openpilot.sunnypilot.models import manager as manager_module +from openpilot.sunnypilot.models.fetcher import ModelFetcher, get_cached_bundles +from openpilot.sunnypilot.models.helpers import get_active_bundle, get_selected_bundle, resolve_bundle_by_ref from openpilot.sunnypilot.models.manager import ModelManagerSP CHUNK_BODIES = [b'A' * 5000, b'B' * 5000, b'C' * 3000] @@ -103,6 +106,7 @@ class ManagerDownloadTestBase(OpenpilotTestCase): self.manager.selected_bundle = None self.manager.active_bundle = None self.manager.available_models = [] + self.manager.chestnut_present = False self.manager._chunk_size = 1024 self.manager._download_start_times = {} @@ -249,6 +253,85 @@ class TestManagerDownload(ManagerDownloadTestBase): assert self.manager._download_start_times == {} self.run_with_server(body) + def test_download_ref_present_keeps_download_alive(self): + """A pending download request (DownloadRef set) must not be cancelled mid-transfer.""" + def body(): + artifact = self.make_artifact(chunked=True) + base_path = os.path.join(self.dest, artifact.fileName) + self.manager.params.get.side_effect = lambda key: b"ref" if key == "ModelManager_DownloadRef" else None + asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + assert os.path.isfile(get_manifest_path(base_path)) + self.run_with_server(body) + + def test_cancellation_via_download_ref(self): + """Removing DownloadRef mid-transfer cancels the download.""" + def body(): + artifact = self.make_artifact(chunked=True) + base_path = os.path.join(self.dest, artifact.fileName) + checks = {"n": 0} + + def get(key): + if key == "ModelManager_DownloadRef": + checks["n"] += 1 + return b"ref" if checks["n"] <= 2 else None + return b"0" + + self.manager.params.get.side_effect = get + with self.assertRaises(Exception) as ctx: + asyncio.run(self.manager._download_chunked(artifact.downloadUri.uri, base_path, artifact)) + assert 'cancelled' in str(ctx.exception).lower() + assert not os.path.isfile(get_manifest_path(base_path)) + self.run_with_server(body) + + def _make_params_with_store(self): + params = mock.MagicMock() + store = {} + + def get(key, *args, **kwargs): + return store.get(key, b"0") # b"0" -> download not cancelled + + def put(key, value, *args, **kwargs): + store[key] = value + + params.get.side_effect = get + params.put.side_effect = put + return params, store + + def test_download_writes_qcom_slot(self): + """A download resolved to the qcom source writes the qcom active bundle slot only.""" + def body(): + artifact = self.make_artifact(chunked=True) + self._bundle.ref = "test-ref" + self._bundle.minimumSelectorVersion = 18 + params, store = self._make_params_with_store() + self.manager.params = params + asyncio.run(self.manager._download_bundle(self._bundle, self.dest, "qcom")) + + assert "ModelManager_ActiveBundle" in store, "qcom download must write the qcom slot" + assert "ModelManager_ActiveBundleUSBGPU" not in store, "qcom download must not touch the usbgpu slot" + assert self.manager.selected_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded + assert self.manager.active_bundle is not None and self.manager.active_bundle.ref == "test-ref" + assert self.manager.active_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded + chunk_names = [get_chunk_name(artifact.fileName, i, len(artifact.chunks)) for i in range(len(artifact.chunks))] + missing = [c for c in chunk_names if not os.path.isfile(os.path.join(self.dest, c))] + assert missing == [], f"chunks missing from the cache: {missing}" + self.run_with_server(body) + + def test_download_writes_usbgpu_slot(self): + """A download resolved to the usbgpu source writes the usbgpu active bundle slot only.""" + def body(): + self.make_artifact(chunked=True) + self._bundle.ref = "big-ref" + self._bundle.minimumSelectorVersion = 18 + params, store = self._make_params_with_store() + self.manager.params = params + asyncio.run(self.manager._download_bundle(self._bundle, self.dest, "usbgpu")) + + assert "ModelManager_ActiveBundleUSBGPU" in store, "usbgpu download must write the usbgpu slot" + assert "ModelManager_ActiveBundle" not in store, "usbgpu download must not touch the qcom slot" + assert self.manager.selected_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded + self.run_with_server(body) + class TestManagerImports(OpenpilotTestCase): """Catches undeclared dependencies. aiohttp lived only in the AGNOS venv; 19.6 dropped @@ -267,6 +350,253 @@ class TestManagerImports(OpenpilotTestCase): assert connect > 0 and read > 0, "requests defaults to no timeout; downloads would hang forever" +class TestResolveBundleByRef(OpenpilotTestCase): + """A ref resolves to (bundle, source) across both hardware manifests. Refs are + unique per manifest and never overlap across sources, so a ref maps to exactly + one slot. Shared by the manager's download flow and the settings UI.""" + + @staticmethod + def _bundle(ref: str): + bundle = custom.ModelManagerSP.ModelBundle.new_message() + bundle.ref = ref + return bundle + + def test_qcom_ref_resolves_to_qcom_slot(self): + small = self._bundle("small") + assert resolve_bundle_by_ref("small", {"qcom": [small], "usbgpu": []}) == (small, "qcom") + + def test_usbgpu_ref_resolves_to_usbgpu_slot(self): + big = self._bundle("big") + assert resolve_bundle_by_ref("big", {"qcom": [], "usbgpu": [big]}) == (big, "usbgpu") + + def test_unknown_ref_returns_none(self): + source_bundles = {"qcom": [self._bundle("small")], "usbgpu": []} + assert resolve_bundle_by_ref("nope", source_bundles) is None + + +def manifest_bundle(short_name: str, ref: str, index: int = 0, is_big: bool = False) -> dict: + """Minimal manifest bundle dict, version-compatible (no chunks to avoid disk side effects). + Big (usbgpu) bundles carry `is_big: true` in the manifest JSON.""" + return { + "index": index, + "short_name": short_name, + "display_name": short_name.upper(), + "generation": 1, + "environment": "release", + "runner": "tinygrad", + "is_big": is_big, + "minimum_selector_version": "18", + "ref": ref, + "models": [{ + "type": "supercombo", + "artifact": { + "file_name": f"{short_name}.pkl", + "download_uri": {"url": f"https://example.com/{short_name}.pkl", "sha256": "s"}, + }, + }], + } + + +def fresh_sync_time() -> int: + return int(time.monotonic() * 1e9) + + +class TestModelFetcherSources(OpenpilotTestCase): + """Both manifests are always maintained: get_bundles_for_source exposes either + source by name, and active_source picks which one matches the attached hardware.""" + + def _make_params(self, qcom_manifest, usbgpu_manifest): + params = mock.MagicMock() + + def get(key): + if key == "ModelManager_ModelsCache": + return qcom_manifest + if key == "ModelManager_ModelsCache_USBGPU": + return usbgpu_manifest + if key in ("ModelManager_LastSyncTime", "ModelManager_LastSyncTime_USBGPU"): + return fresh_sync_time() + return None + + params.get.side_effect = get + return params + + def test_active_source_follows_chestnut_presence(self): + assert ModelFetcher.active_source(False) == "qcom" + assert ModelFetcher.active_source(True) == "usbgpu" + + def test_get_bundles_for_source_returns_each_source(self): + params = self._make_params({"bundles": [manifest_bundle("small", "aaa")]}, + {"bundles": [manifest_bundle("big", "bbb", is_big=True)]}) + fetcher = ModelFetcher(params) + assert [bundle.ref for bundle in fetcher.get_bundles_for_source("qcom")] == ["aaa"] + assert [bundle.ref for bundle in fetcher.get_bundles_for_source("usbgpu")] == ["bbb"] + + def test_get_bundles_for_source_unknown(self): + assert ModelFetcher(mock.MagicMock()).get_bundles_for_source("bogus") == [] + + def test_get_cached_bundles_parses_source(self): + params = self._make_params({"bundles": [manifest_bundle("small", "aaa")]}, + {"bundles": [manifest_bundle("big", "bbb", is_big=True)]}) + qcom_bundles = get_cached_bundles(params, "qcom") + usbgpu_bundles = get_cached_bundles(params, "usbgpu") + assert [b.ref for b in qcom_bundles] == ["aaa"] + assert [b.ref for b in usbgpu_bundles] == ["bbb"] + assert qcom_bundles[0].displayName == "SMALL" + + def test_get_cached_bundles_empty_when_missing(self): + params = mock.MagicMock() + params.get.return_value = None + assert get_cached_bundles(params, "qcom") == [] + assert get_cached_bundles(params, "usbgpu") == [] + + def test_get_cached_bundles_unknown_source(self): + assert get_cached_bundles(mock.MagicMock(), "bogus") == [] + + def test_active_json_has_both_urls(self): + params = mock.MagicMock() + ModelFetcher(params) + active_json_calls = [call for call in params.put.call_args_list if call.args[0] == "ModelManager_ActiveJson"] + assert active_json_calls, "expected ModelManager_ActiveJson to be written" + assert active_json_calls[-1].args[1] == { + "qcom": ModelFetcher.MODEL_URL, + "usbgpu": ModelFetcher.MODEL_URL_USBGPU, + } + + + +class TestSourceCacheIntegrity(OpenpilotTestCase): + """Each source's cached manifest must contain only that source's models; the + `is_big` flag in the JSON marks the big (usbgpu) models. A mismatched cache is + legacy data from before the per-source split (the active manifest was cached + under the unsuffixed key regardless of hardware) and is refetched. This + replaces the old one-time bundle migration.""" + + def _make_params(self, qcom_manifest, usbgpu_manifest): + params = mock.MagicMock() + + def get(key): + if key == "ModelManager_ModelsCache": + return qcom_manifest + if key == "ModelManager_ModelsCache_USBGPU": + return usbgpu_manifest + if key in ("ModelManager_LastSyncTime", "ModelManager_LastSyncTime_USBGPU"): + return fresh_sync_time() + return None + + params.get.side_effect = get + return params + + def _fetched(self, *bundles): + return ModelFetcher(mock.MagicMock()).model_parser.parse_models({"bundles": list(bundles)}) + + def test_qcom_cache_with_big_models_is_refetched(self): + """Legacy: the unsuffixed cache holds the big manifest. is_big confirms it is + the wrong set for qcom, so a fresh fetch replaces it.""" + params = self._make_params({"bundles": [manifest_bundle("big", "bbb", is_big=True)]}, + {"bundles": [manifest_bundle("big2", "ccc", is_big=True)]}) + fetcher = ModelFetcher(params) + fetched = self._fetched(manifest_bundle("small", "aaa")) + with mock.patch.object(fetcher, "_fetch_and_cache_models", return_value=fetched): + bundles = fetcher.get_bundles_for_source("qcom") + assert [bundle.ref for bundle in bundles] == ["aaa"] + + def test_usbgpu_cache_without_big_models_is_refetched(self): + params = self._make_params({"bundles": [manifest_bundle("small", "aaa")]}, + {"bundles": [manifest_bundle("big2", "ccc")]}) + fetcher = ModelFetcher(params) + fetched = self._fetched(manifest_bundle("big", "bbb", is_big=True)) + with mock.patch.object(fetcher, "_fetch_and_cache_models", return_value=fetched): + bundles = fetcher.get_bundles_for_source("usbgpu") + assert [bundle.ref for bundle in bundles] == ["bbb"] + + def test_matching_caches_are_used_without_fetch(self): + params = self._make_params({"bundles": [manifest_bundle("small", "aaa")]}, + {"bundles": [manifest_bundle("big", "bbb", is_big=True)]}) + fetcher = ModelFetcher(params) + with mock.patch.object(fetcher, "_fetch_and_cache_models", side_effect=AssertionError("cache should be used")): + assert [bundle.ref for bundle in fetcher.get_bundles_for_source("qcom")] == ["aaa"] + assert [bundle.ref for bundle in fetcher.get_bundles_for_source("usbgpu")] == ["bbb"] + + def test_stale_version_cache_is_refetched(self): + """A source-matching cache whose bundles are all filtered by the selector version + check parses to zero valid bundles; it is stale (e.g. an old manifest) and must be + refetched instead of silently returning an empty list forever.""" + stale = manifest_bundle("small", "aaa") + stale["minimum_selector_version"] = "16" + params = self._make_params({"bundles": [stale]}, + {"bundles": [manifest_bundle("big", "bbb", is_big=True)]}) + fetcher = ModelFetcher(params) + fetched = self._fetched(manifest_bundle("small2", "ddd")) + with mock.patch.object(fetcher, "_fetch_and_cache_models", return_value=fetched) as fetch: + bundles = fetcher.get_bundles_for_source("qcom") + fetch.assert_called_once_with("qcom") + assert [bundle.ref for bundle in bundles] == ["ddd"] + + def test_corrupt_cache_is_refetched(self): + """A cache that fails to parse (e.g. truncated/foreign JSON) must trigger a + refetch instead of raising every loop and never recovering.""" + corrupt = {"bundles": [{"short_name": "broken"}]} # missing required fields + params = self._make_params(corrupt, {"bundles": [manifest_bundle("big", "bbb", is_big=True)]}) + fetcher = ModelFetcher(params) + fetched = self._fetched(manifest_bundle("small", "aaa")) + with mock.patch.object(fetcher, "_fetch_and_cache_models", return_value=fetched) as fetch: + bundles = fetcher.get_bundles_for_source("qcom") + fetch.assert_called_once_with("qcom") + assert [bundle.ref for bundle in bundles] == ["aaa"] + + +class TestActiveBundleSelection(OpenpilotTestCase): + """The effective active bundle follows the hardware: the usbgpu slot wins when a GPU + is present and compiled, otherwise the qcom slot. Each slot keeps its own selection.""" + + @staticmethod + def _raw_bundle(ref: str) -> dict: + bundle = custom.ModelManagerSP.ModelBundle.new_message() + bundle.ref = ref + bundle.minimumSelectorVersion = 18 + return bundle.to_dict() + + def _params(self, qcom=None, usbgpu=None): + params = mock.MagicMock() + + def get(key, *args, **kwargs): + if key == "ModelManager_ActiveBundle": + return qcom + if key == "ModelManager_ActiveBundleUSBGPU": + return usbgpu + return None + + params.get.side_effect = get + return params + + def test_selected_bundle_is_per_slot(self): + params = self._params(qcom=self._raw_bundle("small"), usbgpu=self._raw_bundle("big")) + assert get_selected_bundle(params, "qcom").ref == "small" + assert get_selected_bundle(params, "usbgpu").ref == "big" + + +class TestEffectiveSource(OpenpilotTestCase): + """One gate decides the active source. With no flags it is runtime truth (GPU + attached); display callers (mici) pass the ui_state flags, which additionally + require the big model to be loading, active, or the device offroad. The active + bundle is simply the selected bundle of that source.""" + + @staticmethod + def _raw_bundle(ref: str) -> dict: + bundle = custom.ModelManagerSP.ModelBundle.new_message() + bundle.ref = ref + bundle.minimumSelectorVersion = 18 + return bundle.to_dict() + + + def test_active_bundle_follows_source(self): + params = mock.MagicMock() + params.get.side_effect = lambda key: {"ModelManager_ActiveBundle": self._raw_bundle("small"), + "ModelManager_ActiveBundleUSBGPU": self._raw_bundle("big")}.get(key) + with mock.patch("openpilot.sunnypilot.models.helpers.usbgpu_present", return_value=False): + assert get_active_bundle(params).ref == "small" + + @unittest.skipUnless(os.environ.get('RUN_INTEGRATION_TESTS'), 'requires external network') class TestLiveModelManifest(OpenpilotTestCase): """Every artifact and chunk URL in the published manifest must resolve.""" diff --git a/openpilot/sunnypilot/sunnylink/statsd.py b/openpilot/sunnypilot/sunnylink/statsd.py index 7e8faf6327..a221fc084f 100755 --- a/openpilot/sunnypilot/sunnylink/statsd.py +++ b/openpilot/sunnypilot/sunnylink/statsd.py @@ -65,6 +65,7 @@ def sp_stats(end_event): 'MadsSteeringMode', 'MadsUnifiedEngagementMode', 'ModelManager_ActiveBundle', + 'ModelManager_ActiveBundleUSBGPU', 'ModelManager_Favs', 'EnableSunnylinkUploader', 'SunnylinkEnabled',