diff --git a/openpilot/cereal/custom.capnp b/openpilot/cereal/custom.capnp index b8c833ed8..d0e9dd463 100644 --- a/openpilot/cereal/custom.capnp +++ b/openpilot/cereal/custom.capnp @@ -138,6 +138,9 @@ struct ModelManagerSP @0xaedffd8f31e7b55d { status @0 :DownloadStatus; progress @1 :Float32; eta @2 :UInt32; + speed @3 :Float32; + canceling @4 :Bool; + resuming @5 :Bool; } struct Chunk { diff --git a/openpilot/selfdrive/ui/sunnypilot/layouts/settings/models.py b/openpilot/selfdrive/ui/sunnypilot/layouts/settings/models.py index d036d8d28..1d12f9f5b 100644 --- a/openpilot/selfdrive/ui/sunnypilot/layouts/settings/models.py +++ b/openpilot/selfdrive/ui/sunnypilot/layouts/settings/models.py @@ -235,6 +235,8 @@ class ModelsLayout(Widget): progress = sum(p.progress for p in progresses) / len(progresses) ds = custom.ModelManagerSP.DownloadStatus + if any(p.canceling for p in progresses): + return {"name": name, "downloading": True, "progress": progress, "status_text": tr("canceling...")} if ds.failed in statuses: # close.png is authored black and a tint cannot lift it, hence close2 return {"name": name, "status_text": tr("download failed"), "text_color": rl.RED, "icon": "icons/close2.png"} diff --git a/openpilot/selfdrive/ui/sunnypilot/mici/layouts/models.py b/openpilot/selfdrive/ui/sunnypilot/mici/layouts/models.py index 404ff2e30..99e0b06e8 100644 --- a/openpilot/selfdrive/ui/sunnypilot/mici/layouts/models.py +++ b/openpilot/selfdrive/ui/sunnypilot/mici/layouts/models.py @@ -251,11 +251,12 @@ class ModelsLayoutMici(NavScroller): self.current_model_info.info_text.set_text(tr("download failed")) elif manager.selectedBundle and manager.selectedBundle.status == custom.ModelManagerSP.DownloadStatus.downloading: - self.cancel_download_btn.set_visible(True) + self.cancel_download_btn.set_visible(ui_state.params.get("ModelManager_DownloadRef") is not None) device.set_override_interactive_timeout(5) progress = 0.0 count = 0 verifying = False + canceling = False for model in manager.selectedBundle.models: count += 1 p = model.artifact.downloadProgress @@ -263,11 +264,13 @@ class ModelsLayoutMici(NavScroller): custom.ModelManagerSP.DownloadStatus.verifying): progress += p.progress verifying = verifying or p.status == custom.ModelManagerSP.DownloadStatus.verifying + canceling = canceling or p.canceling elif p.status in (custom.ModelManagerSP.DownloadStatus.downloaded, custom.ModelManagerSP.DownloadStatus.cached): progress += 100.0 - self.current_model_info.current_model_header.set_text(tr("verifying") if verifying else tr("downloading")) + header = tr("canceling...") if canceling else tr("verifying") if verifying else tr("downloading") + self.current_model_info.current_model_header.set_text(header) self.cancel_download_btn.set_text(tr("cancel verification") if verifying else tr("cancel download")) self.current_model_info.current_model_header._shimmer = True name_text = manager.selectedBundle.internalName.lower() diff --git a/openpilot/sunnypilot/models/manager.py b/openpilot/sunnypilot/models/manager.py index 66b0fe28a..bdee8ac94 100644 --- a/openpilot/sunnypilot/models/manager.py +++ b/openpilot/sunnypilot/models/manager.py @@ -7,7 +7,10 @@ See the LICENSE.md file in the root directory for more details. import asyncio import os +import re +import threading import time +from concurrent.futures import ThreadPoolExecutor import requests from openpilot.common.params import Params @@ -22,12 +25,99 @@ from openpilot.sunnypilot.models.helpers import (ACTIVE_BUNDLE_KEYS, get_active_ # (connect, read) seconds. read is per-request inactivity, not a total cap DOWNLOAD_TIMEOUT = (30, 30) +# HuggingFace throttles each connection to ~1-2 MB/s but not the connection count; 12 saturates the device link +MAX_CONCURRENT_PIECES = 12 +PIECE_SIZE = 8 * 1024 * 1024 +PIECE_ATTEMPTS = 3 +PIECE_BACKOFF = 1.0 # seconds, growing linearly: waits 1 s then 2 s between a piece's three attempts +TRANSIENT_HTTP_STATUS = frozenset({408, 429, 500, 502, 503, 504}) +REPORT_INTERVAL = 0.5 # seconds between progress publications +# byte ranges count encoded bytes but iter_content yields decoded ones, so ask for an unencoded body +NO_COMPRESSION = {"Accept-Encoding": "identity"} class DownloadCancelled(Exception): pass +def _is_transient(e: BaseException) -> bool: + if isinstance(e, requests.HTTPError): + return e.response is not None and e.response.status_code in TRANSIENT_HTTP_STATUS + return isinstance(e, (requests.ConnectionError, requests.Timeout, requests.exceptions.ChunkedEncodingError)) + + +class _Progress: + def __init__(self): + self._bytes = 0 + self._lock = threading.Lock() + + def add_bytes(self, n: int) -> None: + with self._lock: + self._bytes += n + + def snapshot(self) -> int: + with self._lock: + return self._bytes + + +def _piece_ranges(total: int, piece_size: int) -> list[tuple[int, int]]: + return [(start, min(start + piece_size, total)) for start in range(0, total, piece_size)] + + +def _prepare_target(path: str, total: int) -> None: + with open(path, "wb") as f: + f.truncate(total) + + +def _content_range_total(response: requests.Response) -> int | None: + if response.status_code != 206: + return None + match = re.fullmatch(r"bytes \d+-\d+/(\d+)", response.headers.get("Content-Range", "")) + return int(match.group(1)) if match else None + + +def _probe_size(url: str) -> tuple[int, bool]: + with requests.get(url, headers={**NO_COMPRESSION, "Range": "bytes=0-0"}, stream=True, timeout=DOWNLOAD_TIMEOUT) as response: + response.raise_for_status() + total = _content_range_total(response) + if total is not None: + return total, True + return int(response.headers.get("Content-Length", 0)), False + + +def _fetch_piece(url: str, path: str, index: int, start: int, end: int | None, progress: _Progress, + cancel: threading.Event, block_size: int) -> None: + headers = {**NO_COMPRESSION, "Range": f"bytes={start}-{end - 1}"} if end is not None else NO_COMPRESSION + label = f"{os.path.basename(path)} piece {index}" + for attempt in range(PIECE_ATTEMPTS): + written = 0 + try: + with requests.get(url, headers=headers, stream=True, timeout=DOWNLOAD_TIMEOUT) as response: + response.raise_for_status() + if end is not None and response.status_code != 206: + raise ValueError(f"server ignored byte range for {label}") + with open(path, "r+b") as f: + f.seek(start) + for block in response.iter_content(chunk_size=block_size): + if cancel.is_set(): + raise DownloadCancelled("Download cancelled") + f.write(block) + written += len(block) + progress.add_bytes(len(block)) + if end is None: + f.truncate() # unknown length: the body defines the file size + if end is not None and written != end - start: + raise requests.exceptions.ConnectionError(f"short read: {written} of {end - start} bytes") + return + except requests.RequestException as e: + progress.add_bytes(-written) # the piece restarts from scratch + if not _is_transient(e) or attempt == PIECE_ATTEMPTS - 1: + raise + cloudlog.warning(f"retrying {label} after {type(e).__name__}: {e}") + if cancel.wait(PIECE_BACKOFF * (1 + attempt)): + raise DownloadCancelled("Download cancelled") from None + + class ModelManagerSP: """Manages model downloads and status reporting""" @@ -41,7 +131,7 @@ class ModelManagerSP: 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, chestnut=self.chestnut_present) - self._chunk_size = 128 * 1000 # 128 KB chunks + self._block_size = 128 * 1000 # 128 KB network read blocks self._download_start_times: dict[str, float] = {} # Track start time per model self._download_ref: bytes | str | None = None @@ -56,7 +146,6 @@ class ModelManagerSP: self._download_ref = None def _sync_artifact_progress(self, source_artifact) -> None: - """Mirror download progress to all artifacts sharing the same filename in the selected bundle.""" if not self.selected_bundle: return for model in self.selected_bundle.models: @@ -67,7 +156,6 @@ class ModelManagerSP: artifact.downloadProgress.eta = source_artifact.downloadProgress.eta def _calculate_eta(self, filename: str, progress: float) -> int: - """Calculate ETA based on elapsed time and current progress""" if filename not in self._download_start_times or progress <= 0: return 60 # Default ETA for new downloads @@ -75,39 +163,58 @@ class ModelManagerSP: if elapsed_time <= 0: return 60 - # If we're at X% after Y seconds, we can estimate total time as (Y / X) * 100 total_estimated_time = (elapsed_time / progress) * 100 eta = total_estimated_time - elapsed_time - return max(1, int(eta)) # Return at least 1 second if download is ongoing - async def _download_file(self, url: str, path: str, model) -> None: - """Downloads a file with progress tracking""" - self._download_start_times[model.fileName] = time.monotonic() + def _set_progress(self, artifact, status, progress: float, eta: int = 0) -> None: + artifact.downloadProgress.status = status + artifact.downloadProgress.progress = progress + artifact.downloadProgress.eta = eta + self._sync_artifact_progress(artifact) + self._report_status() - with requests.get(url, stream=True, timeout=DOWNLOAD_TIMEOUT) as response: # noqa: ASYNC210 - response.raise_for_status() - total_size = int(response.headers.get("content-length", 0)) - bytes_downloaded = 0 + def _publish_progress(self, artifact, done_bytes: int, total: int) -> None: + # held at 99 until the file passes its hash check + progress = min(99.0, done_bytes / total * 100) if total > 0 else 0.0 + eta = self._calculate_eta(artifact.fileName, progress) + self._set_progress(artifact, custom.ModelManagerSP.DownloadStatus.downloading, progress, eta) - with open(path, 'wb') as f: # noqa: ASYNC230 - for chunk in response.iter_content(chunk_size=self._chunk_size): # type: bytes - f.write(chunk) - bytes_downloaded += len(chunk) + async def _report_until_done(self, tasks: list[asyncio.Future], artifact, progress: _Progress, total: int) -> None: + pending = set(tasks) + while pending: + done, pending = await asyncio.wait(pending, timeout=REPORT_INTERVAL) + for task in done: + task.result() # re-raises a worker failure + if self._download_interrupted(): + raise DownloadCancelled("Download cancelled") + self._publish_progress(artifact, progress.snapshot(), total) - if self._download_interrupted(): - raise DownloadCancelled("Download cancelled") + async def _download_file(self, url: str, path: str, artifact) -> None: + self._download_start_times[artifact.fileName] = time.monotonic() + loop = asyncio.get_running_loop() - if total_size > 0: - progress = (bytes_downloaded / total_size) * 100 - model.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading - model.downloadProgress.progress = progress - model.downloadProgress.eta = self._calculate_eta(model.fileName, progress) - self._sync_artifact_progress(model) - self._report_status() + total, ranged = await loop.run_in_executor(None, _probe_size, url) + pieces: list[tuple[int, int | None]] = list(_piece_ranges(total, PIECE_SIZE)) if ranged else [(0, None)] + await loop.run_in_executor(None, _prepare_target, path, total) + progress = _Progress() + cancel = threading.Event() - # Clean up start time after download completes - del self._download_start_times[model.fileName] + with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_PIECES) as pool: + tasks = [loop.run_in_executor(pool, _fetch_piece, url, path, i, start, end, progress, cancel, self._block_size) + for i, (start, end) in enumerate(pieces)] + try: + await self._report_until_done(tasks, artifact, progress, total) + except BaseException as e: + cancel.set() + if isinstance(e, DownloadCancelled): + artifact.downloadProgress.canceling = True + self._set_progress(artifact, custom.ModelManagerSP.DownloadStatus.downloading, artifact.downloadProgress.progress) + for task in tasks: + task.cancel() # drops queued pieces; running ones see the event at their next block + await asyncio.gather(*tasks, return_exceptions=True) + raise + del self._download_start_times[artifact.fileName] async def _process_artifact(self, artifact, destination_path: str) -> None: if not artifact.downloadUri.uri: @@ -119,33 +226,30 @@ class ModelManagerSP: expected_hash = artifact.downloadUri.sha256 filename = artifact.fileName full_path = os.path.join(destination_path, filename) + status = custom.ModelManagerSP.DownloadStatus + artifact.downloadProgress.canceling = False try: + self._set_progress(artifact, status.verifying, 0) if await verify_file(full_path, expected_hash): - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached - artifact.downloadProgress.progress = 100 - artifact.downloadProgress.eta = 0 - self._sync_artifact_progress(artifact) - self._report_status() + self._set_progress(artifact, status.cached, 100) return + self._set_progress(artifact, status.downloading, 0) await self._download_file(url, full_path, artifact) + self._set_progress(artifact, status.verifying, 99) if not await verify_file(full_path, expected_hash): raise ValueError(f"Hash validation failed for {filename}") - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloaded - artifact.downloadProgress.progress = 100 - artifact.downloadProgress.eta = 0 - self._sync_artifact_progress(artifact) - self._report_status() + self._set_progress(artifact, status.downloaded, 100) except DownloadCancelled: self._download_start_times.pop(artifact.fileName, None) - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed + artifact.downloadProgress.status = status.failed artifact.downloadProgress.eta = 0 self._sync_artifact_progress(artifact) if self.selected_bundle: - self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed + self.selected_bundle.status = status.failed self._report_status() raise @@ -153,11 +257,11 @@ class ModelManagerSP: cloudlog.error(f"Error downloading {filename}: {str(e)}") if os.path.isfile(full_path): # noqa: ASYNC240 os.remove(full_path) - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed + artifact.downloadProgress.status = status.failed artifact.downloadProgress.eta = 0 self._sync_artifact_progress(artifact) if self.selected_bundle: - self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed + self.selected_bundle.status = status.failed self._report_status() self._download_start_times.pop(artifact.fileName, None) raise @@ -272,12 +376,7 @@ class ModelManagerSP: rk.keep_time() def clear_model_cache(self) -> None: - """ - Clears the model cache directory of all files except those in the active model bundle. - """ - - # Get list of files used by both slots' selected bundles (either may become - # the truly active bundle depending on hardware availability) + """ Clears the model cache directory of all files except those in the active model bundle.""" active_files = [] for source in ACTIVE_BUNDLE_KEYS: if selected_bundle := get_selected_bundle(self.params, source): @@ -285,7 +384,6 @@ class ModelManagerSP: if model.artifact.fileName: active_files.append(model.artifact.fileName) - # Remove all files except active ones model_dir = Paths.model_root() try: for filename in os.listdir(model_dir): diff --git a/openpilot/sunnypilot/models/tests/test_manager_download.py b/openpilot/sunnypilot/models/tests/test_manager_download.py index 0ae0273ec..662a776e9 100644 --- a/openpilot/sunnypilot/models/tests/test_manager_download.py +++ b/openpilot/sunnypilot/models/tests/test_manager_download.py @@ -6,9 +6,15 @@ See the LICENSE.md file in the root directory for more details. """ import asyncio +import contextlib +import dataclasses import hashlib import http.server +import math import os +import random +import re +import requests import tempfile import threading import time @@ -16,251 +22,188 @@ import unittest from typing import Any from unittest import mock -import requests - from openpilot.cereal import custom from openpilot.common.test import OpenpilotTestCase -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 import helpers -from openpilot.sunnypilot.models.helpers import (get_active_bundle, get_active_source, get_selected_bundle, +from openpilot.sunnypilot.models.helpers import (ACTIVE_BUNDLE_KEYS, get_active_bundle, get_active_source, get_selected_bundle, resolve_bundle_by_ref, validate_active_bundles) -from openpilot.sunnypilot.models.manager import ModelManagerSP +from openpilot.sunnypilot.models.manager import DownloadCancelled, ModelManagerSP -WHOLE_BODY = b'Z' * 9000 +FILE_NAME = 'driving_test_tinygrad.pkl' +# fixed-seed random bytes: every piece differs, so a misplaced piece changes the file +WHOLE_BODY = random.Random(0).randbytes(10_240) +PIECE_SIZE = 1000 # bytes per test piece -> 11 pieces, the last one short +NUM_PIECES = math.ceil(len(WHOLE_BODY) / PIECE_SIZE) def sha256(data: bytes) -> str: return hashlib.sha256(data).hexdigest() -class DownloadHandler(http.server.BaseHTTPRequestHandler): - """Serves the fixture bodies. Class attributes are reset per test.""" - request_paths: list[str] = [] - fail_paths: dict[str, int] = {} - stall_paths: set[str] = set() +def piece_range(index: int) -> tuple[int, int]: + start = index * PIECE_SIZE + return start, min(start + PIECE_SIZE, len(WHOLE_BODY)) - 1 + + +@dataclasses.dataclass +class ServerState: + request_ranges: list[tuple[int, int] | None] = dataclasses.field(default_factory=list) # None: no (honored) Range header + request_encodings: list[str | None] = dataclasses.field(default_factory=list) # Accept-Encoding of every request + fail_ranges: dict[tuple[int, int], int] = dataclasses.field(default_factory=dict) # range -> HTTP status, every time + fail_times: dict[tuple[int, int], int] = dataclasses.field(default_factory=dict) # range -> how many more requests 503 + short_once: set[tuple[int, int]] = dataclasses.field(default_factory=set) # ranges whose first answer is cut short + stall_ranges: set[tuple[int, int]] = dataclasses.field(default_factory=set) stall_event: threading.Event | None = None + stall_released_by_event: bool | None = None + support_ranges: bool = True + corrupt: bool = False + + +class DownloadHandler(http.server.BaseHTTPRequestHandler): + state = ServerState() def log_message(self, format: str, *args: Any) -> None: # noqa: A002 pass - def _body_for(self, path): - return WHOLE_BODY if path.endswith('.pkl') else None - def do_GET(self): - type(self).request_paths.append(self.path) + state = self.state + body = WHOLE_BODY + if state.corrupt: + # flip every bit of one mid-file byte: same size, so only the hash check can catch it + body = body[:5000] + bytes([body[5000] ^ 0xFF]) + body[5001:] - status = type(self).fail_paths.get(self.path) - if status: - self.send_response(status) - self.end_headers() - return + rng = None + match = re.fullmatch(r'bytes=(\d+)-(\d+)', self.headers.get('Range', '')) + if match and state.support_ranges: + rng = (int(match.group(1)), int(match.group(2))) + state.request_ranges.append(rng) + state.request_encodings.append(self.headers.get('Accept-Encoding')) - body = self._body_for(self.path) - if body is None: - self.send_response(404) - self.end_headers() - return + if rng is not None: + status = state.fail_ranges.get(rng) + if status is None and state.fail_times.get(rng, 0) > 0: + state.fail_times[rng] -= 1 + status = 503 + if status is not None: + self.send_response(status) + self.end_headers() + return - self.send_response(200) - self.send_header('Content-Length', str(len(body))) + if rng is None: + payload = body + self.send_response(200) + else: + start, end = rng + payload = body[start:end + 1] + self.send_response(206) + self.send_header('Content-Range', f'bytes {start}-{end}/{len(body)}') + if rng in state.short_once: + state.short_once.discard(rng) + payload = payload[:100] + self.send_header('Content-Length', str(len(payload))) self.end_headers() - if self.path in type(self).stall_paths: - # write a little, then wait so the test can cancel mid-transfer - self.wfile.write(body[:100]) + if rng in state.stall_ranges: + # write a little, then hold the connection until the test releases it + self.wfile.write(payload[:100]) self.wfile.flush() - if type(self).stall_event is not None: - type(self).stall_event.wait(timeout=5) - self.wfile.write(body[100:]) + if state.stall_event is not None: + state.stall_released_by_event = state.stall_event.wait(timeout=5) + self.wfile.write(payload[100:]) else: - self.wfile.write(body) + self.wfile.write(payload) + + +@contextlib.contextmanager +def threaded_server(handler): + server = http.server.ThreadingHTTPServer(('127.0.0.1', 0), handler) + thread = threading.Thread(target=server.serve_forever) + thread.start() + try: + yield f'http://127.0.0.1:{server.server_port}' + finally: + server.shutdown() + server.server_close() + thread.join() class ManagerDownloadTestBase(OpenpilotTestCase): def setUp(self): super().setUp() - DownloadHandler.request_paths = [] - DownloadHandler.fail_paths = {} - DownloadHandler.stall_paths = set() - DownloadHandler.stall_event = None + self.server = DownloadHandler.state = ServerState() + self.base_url = self.enterContext(threaded_server(DownloadHandler)) self._tmp = tempfile.TemporaryDirectory() self.addCleanup(self._tmp.cleanup) self.dest = self._tmp.name + for name, value in (('PIECE_SIZE', PIECE_SIZE), ('REPORT_INTERVAL', 0.05), ('PIECE_BACKOFF', 0.01)): + self.enterContext(mock.patch.object(manager_module, name, value)) + self.reported: list[float] = [] + self.reported_statuses: list[int] = [] self.manager = ModelManagerSP.__new__(ModelManagerSP) - self.manager.params = mock.MagicMock() - self.manager.params.get.return_value = b'0' # not cancelled + self.params = self.manager.params = mock.MagicMock() + self.params.get.return_value = b'0' # not cancelled self.manager._download_ref = b'0' - self.manager.pm = mock.MagicMock() - self.manager.pm.send.side_effect = self._record_progress + self.pm = self.manager.pm = mock.MagicMock() + self.pm.send.side_effect = self._record_progress 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._block_size = 256 self.manager._download_start_times = {} def _record_progress(self, *args) -> None: - """Runs on every real _report_status send.""" artifact = getattr(self, 'artifact', None) if artifact is not None: self.reported.append(float(artifact.downloadProgress.progress)) + status = artifact.downloadProgress.status + self.reported_statuses.append(getattr(status, 'raw', status)) # .raw: _DynamicEnum equals its int but does not hash like it def make_artifact(self): bundle = custom.ModelManagerSP.ModelBundle.new_message() bundle.init('models', 1) artifact = bundle.models[0].artifact - artifact.fileName = 'driving_test_tinygrad.pkl' - artifact.downloadUri.uri = self.base_url + '/driving_test_tinygrad.pkl' + artifact.fileName = FILE_NAME + artifact.downloadUri.uri = self.base_url + '/' + FILE_NAME artifact.downloadUri.sha256 = sha256(WHOLE_BODY) self._bundle = bundle self.artifact = artifact return artifact -class TestManagerDownload(ManagerDownloadTestBase): - """Exercises the real _download_file against a local server.""" + @property + def path(self) -> str: + return os.path.join(self.dest, FILE_NAME) - def run_with_server(self, fn): - with http_server_context(handler=DownloadHandler) as (host, port): - self.base_url = f'http://{host}:{port}' - return fn() + def piece_requests(self) -> list[tuple[int, int]]: + return [r for r in self.server.request_ranges if r is not None and r != (0, 0)] - def test_download_file_writes_exact_bytes(self): - def body(): - artifact = self.make_artifact() - path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, path, artifact)) - with open(path, 'rb') as f: - written = f.read() - assert written == WHOLE_BODY - assert sha256(written) == artifact.downloadUri.sha256 - assert artifact.fileName not in self.manager._download_start_times - self.run_with_server(body) + def stall(self, index: int) -> threading.Event: + self.server.stall_ranges = {piece_range(index)} + self.server.stall_event = threading.Event() + return self.server.stall_event - def test_progress_is_monotonic_and_bounded(self): - def body(): - artifact = self.make_artifact() - base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) + def download_file(self): + artifact = self.make_artifact() + asyncio.run(self.manager._download_file(artifact.downloadUri.uri, self.path, artifact)) + return artifact - assert self.reported, "expected progress reports" - for a, b in zip(self.reported, self.reported[1:], strict=False): - assert b >= a, f"progress went backwards: {a} -> {b}" - self.run_with_server(body) + def process_artifact(self): + artifact = self.make_artifact() + asyncio.run(self.manager._process_artifact(artifact, self.dest)) + return artifact - def test_http_error_propagates(self): - def body(): - artifact = self.make_artifact() - base_path = os.path.join(self.dest, artifact.fileName) - DownloadHandler.fail_paths = {'/' + artifact.fileName: 404} + def download_bundle(self, source: str): + asyncio.run(self.manager._download_bundle(self._bundle, self.dest, source)) - with self.assertRaises(requests.exceptions.HTTPError): - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) - assert not os.path.isfile(base_path) - self.run_with_server(body) - - def test_cancellation_mid_transfer(self): - """Cancellation is checked inside the byte loop; it must still fire after the port.""" - def body(): - artifact = self.make_artifact() - base_path = os.path.join(self.dest, artifact.fileName) - self.manager.params.get.return_value = None # cancelled - - with self.assertRaises(Exception) as ctx: - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) - assert 'cancelled' in str(ctx.exception).lower() - self.run_with_server(body) - - def test_repeat_downloads_are_stable(self): - """Back-to-back runs must produce identical bytes and leak no start-time state.""" - def body(): - for _ in range(2): - artifact = self.make_artifact() - base_path = os.path.join(self.dest, artifact.fileName) - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) - with open(base_path, 'rb') as f: - assert f.read() == WHOLE_BODY - 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() - 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 - self.manager._download_ref = b"ref" - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) - assert os.path.isfile(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() - 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 - self.manager._download_ref = b"ref" - with self.assertRaises(Exception) as ctx: - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) - assert 'cancelled' in str(ctx.exception).lower() - self.run_with_server(body) - - def test_replaced_download_ref_queues_instead_of_cancelling(self): - """Selecting another model mid-transfer lets the running download finish.""" - def body(): - artifact = self.make_artifact() - base_path = os.path.join(self.dest, artifact.fileName) - self.manager.params.get.side_effect = lambda key: b"other-ref" if key == "ModelManager_DownloadRef" else None - self.manager._download_ref = b"ref" - asyncio.run(self.manager._download_file(artifact.downloadUri.uri, base_path, artifact)) - assert os.path.isfile(base_path) - self.run_with_server(body) - - def test_replaced_download_ref_is_kept(self): - """A selection made during a download must survive that download's cleanup.""" - self.manager.params.get.return_value = b"new-ref" - self.manager._download_ref = b"old-ref" - self.manager._release_download_ref() - self.manager.params.remove.assert_not_called() - - def test_own_download_ref_is_released(self): - self.manager.params.get.return_value = b"ref" - self.manager._download_ref = b"ref" - self.manager._release_download_ref() - self.manager.params.remove.assert_called_once_with("ModelManager_DownloadRef") - - def test_cached_bundle_cancel_skips_slot_write(self): - """A cancel must stop an already-on-disk bundle before it is applied to the slot.""" - def body(): - artifact = self.make_artifact() - base_path = os.path.join(self.dest, artifact.fileName) - with open(base_path, 'wb') as f: - f.write(WHOLE_BODY) - self._bundle.ref = "test-ref" - params, store = self._make_params_with_store() - store["ModelManager_DownloadRef"] = None # removed -> cancelled - self.manager.params = params - self.manager._download_ref = b"ref" - with self.assertRaises(Exception) as ctx: - asyncio.run(self.manager._download_bundle(self._bundle, self.dest, "qcom")) - assert 'cancelled' in str(ctx.exception).lower() - assert "ModelManager_ActiveBundle" not in store - assert os.path.isfile(base_path), "cancel must not delete the cached model" - self.run_with_server(body) + def read_path(self) -> bytes: + with open(self.path, 'rb') as f: + return f.read() def _make_params_with_store(self): params = mock.MagicMock() @@ -276,84 +219,156 @@ class TestManagerDownload(ManagerDownloadTestBase): 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() - self._bundle.ref = "test-ref" - self._bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION - 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_ActiveBundleChestnut" not in store, "qcom download must not touch the chestnut 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 - assert os.path.isfile(os.path.join(self.dest, artifact.fileName)) - self.run_with_server(body) +class TestManagerDownload(ManagerDownloadTestBase): + def test_download(self): + artifact = self.download_file() + assert self.read_path() == WHOLE_BODY + assert sorted(self.piece_requests()) == [piece_range(i) for i in range(NUM_PIECES)], "pieces must tile the file once" + assert self.server.request_encodings == ['identity'] * (NUM_PIECES + 1), "ranges count encoded bytes, so ask for none" + assert self.reported and self.reported == sorted(self.reported), "progress only rises" + assert max(self.reported) <= 99.0, "progress holds below 100 until the hash check" + assert artifact.fileName not in self.manager._download_start_times - def test_download_writes_chestnut_slot(self): - """A download resolved to the chestnut source writes the chestnut active bundle slot only.""" - def body(): - self.make_artifact() - self._bundle.ref = "big-ref" - self._bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION - params, store = self._make_params_with_store() - self.manager.params = params - asyncio.run(self.manager._download_bundle(self._bundle, self.dest, "chestnut")) + def test_pieces_download_in_parallel(self): + stall_event = self.stall(0) + others = {piece_range(i) for i in range(1, NUM_PIECES)} - assert "ModelManager_ActiveBundleChestnut" in store, "chestnut download must write the chestnut slot" - assert "ModelManager_ActiveBundle" not in store, "chestnut download must not touch the qcom slot" - assert self.manager.selected_bundle.status == custom.ModelManagerSP.DownloadStatus.downloaded - self.run_with_server(body) + def release_when_others_requested(): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + if others <= set(self.piece_requests()): + stall_event.set() + return + time.sleep(0.01) + + threading.Thread(target=release_when_others_requested).start() + self.download_file() + assert self.server.stall_released_by_event is True, "other pieces did not download while one was stalled" + + def test_fallback_when_server_ignores_ranges(self): + self.server.support_ranges = False + self.download_file() + assert self.read_path() == WHOLE_BODY + assert self.server.request_ranges == [None, None], "expected the probe plus one whole-body request" + assert self.server.request_encodings == ['identity', 'identity'] + + def test_transient_failures_are_retried(self): + self.server.fail_times = {piece_range(2): 1} + self.server.short_once = {piece_range(4)} + self.download_file() + assert self.read_path() == WHOLE_BODY + assert self.piece_requests().count(piece_range(2)) == 2 + assert self.piece_requests().count(piece_range(4)) == 2 + + def test_permanent_failure_is_not_retried(self): + self.server.fail_ranges = {piece_range(1): 404} + with self.assertRaises(requests.exceptions.HTTPError): + self.download_file() + assert self.piece_requests().count(piece_range(1)) == 1 + + def test_removing_download_ref_cancels(self): + stall_event = self.stall(3) + refs = iter([b"ref"]) + self.params.get.side_effect = lambda key: next(refs, None) # removed after the first poll + self.manager._download_ref = b"ref" + + def release_once_canceling(*args): + # the "canceling" report comes after the cancel event is set, so the piece can't finish first + if self.artifact.downloadProgress.canceling: + stall_event.set() + self.pm.send.side_effect = release_once_canceling + with self.assertRaises(DownloadCancelled): + self.download_file() + assert self.server.stall_released_by_event is True, "the cancel must have caught the piece mid-transfer" + end = piece_range(3)[1] + 1 + assert self.read_path()[end - 100:end] == bytes(100), "the cancelled piece must stop before its last bytes land" + + def test_replaced_download_ref_queues_instead_of_canceling(self): + self.params.get.return_value = b"other-ref" + self.manager._download_ref = b"ref" + self.download_file() + assert self.read_path() == WHOLE_BODY -class TestManagerImports(OpenpilotTestCase): - """Catches undeclared dependencies. aiohttp lived only in the AGNOS venv; 19.6 dropped - it and models_manager died on device while CI stayed green.""" +class TestProcessArtifact(ManagerDownloadTestBase): + def test_download_writes_only_its_slot(self): + ds = custom.ModelManagerSP.DownloadStatus + for source, other in (("qcom", "chestnut"), ("chestnut", "qcom")): # bernoulli + with self.subTest(source=source): + self.make_artifact() + self._bundle.ref = f"{source}-ref" + self._bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION + self.manager.params, store = self._make_params_with_store() + self.download_bundle(source) + assert ACTIVE_BUNDLE_KEYS[source] in store and ACTIVE_BUNDLE_KEYS[other] not in store + assert self.manager.selected_bundle.status == ds.downloaded + assert self.artifact.downloadProgress.progress == 100 and self.artifact.downloadProgress.eta == 0 + assert self.read_path() == WHOLE_BODY + os.remove(self.path) - def test_manager_imports(self): - assert manager_module.ModelManagerSP is not None + def test_hash_mismatch_discards_file(self): + self.server.corrupt = True + with self.assertRaises(ValueError): + self.process_artifact() + assert not os.path.isfile(self.path) + assert self.artifact.downloadProgress.status == custom.ModelManagerSP.DownloadStatus.failed - def test_no_undeclared_http_client(self): - with open(manager_module.__file__) as f: - src = f.read() - assert 'import aiohttp' not in src, "aiohttp is not available on AGNOS 19.6; use requests" + def test_probe_failure_moves_past_verifying(self): + self.server.fail_ranges = {(0, 0): 404} # the size probe + with self.assertRaises(requests.exceptions.HTTPError): + self.process_artifact() + ds = custom.ModelManagerSP.DownloadStatus + assert self.reported_statuses == [ds.verifying, ds.downloading, ds.failed] - def test_download_timeout_is_explicit(self): - connect, read = manager_module.DOWNLOAD_TIMEOUT - assert connect > 0 and read > 0, "requests defaults to no timeout; downloads would hang forever" + def test_cached_file_skips_network(self): + with open(self.path, 'wb') as f: + f.write(WHOLE_BODY) + self.process_artifact() + assert self.server.request_ranges == [], "cached file must not hit the network" + ds = custom.ModelManagerSP.DownloadStatus + assert self.reported_statuses == [ds.verifying, ds.cached] + + def test_cached_bundle_cancel_skips_slot_write(self): + with open(self.path, 'wb') as f: + f.write(WHOLE_BODY) + self.make_artifact() + self.manager.params, store = self._make_params_with_store() + store["ModelManager_DownloadRef"] = None # removed -> cancelled + self.manager._download_ref = b"ref" + with self.assertRaises(DownloadCancelled): + self.download_bundle("qcom") + assert ACTIVE_BUNDLE_KEYS["qcom"] not in store + assert os.path.isfile(self.path), "cancel must not delete a cached model" + + def test_download_ref_is_released_only_if_unchanged(self): + for current, removed in ((b"new-ref", False), (b"ref", True)): + self.params = self.manager.params = mock.MagicMock() + self.params.get.return_value = current + self.manager._download_ref = b"ref" + self.manager._release_download_ref() + assert self.params.remove.called == removed + + +def raw_bundle(ref: str, runner: int | None = None) -> dict: + bundle = custom.ModelManagerSP.ModelBundle.new_message() + bundle.ref = ref + bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION + if runner is not None: + bundle.runner = runner + return bundle.to_dict() 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], "chestnut": []}) == (small, "qcom") - - def test_chestnut_ref_resolves_to_chestnut_slot(self): - big = self._bundle("big") - assert resolve_bundle_by_ref("big", {"qcom": [], "chestnut": [big]}) == (big, "chestnut") - - def test_unknown_ref_returns_none(self): - source_bundles = {"qcom": [self._bundle("small")], "chestnut": []} - assert resolve_bundle_by_ref("nope", source_bundles) is None + def test_ref_resolves_to_its_slot(self): + small, big = (custom.ModelManagerSP.ModelBundle(**raw_bundle(ref)) for ref in ("small", "big")) + catalog = {"qcom": [small], "chestnut": [big]} + assert resolve_bundle_by_ref("small", catalog) == (small, "qcom") + assert resolve_bundle_by_ref("big", catalog) == (big, "chestnut") + assert resolve_bundle_by_ref("nope", catalog) 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 (chestnut) bundles carry `is_big: true` in the manifest JSON.""" return { "index": index, "short_name": short_name, @@ -374,309 +389,127 @@ def manifest_bundle(short_name: str, ref: str, index: int = 0, is_big: bool = Fa } -def fresh_sync_time() -> int: - return int(time.monotonic() * 1e9) +def manifest_params(qcom_manifest, chestnut_manifest=None): + values = { + "ModelManager_ModelsCache": qcom_manifest, + "ModelManager_ModelsCache_Chestnut": chestnut_manifest or {"bundles": [manifest_bundle("big", "bbb", is_big=True)]}, + "ModelManager_LastSyncTime": int(time.monotonic() * 1e9), + "ModelManager_LastSyncTime_Chestnut": int(time.monotonic() * 1e9), + } + params = mock.MagicMock() + params.get.side_effect = values.get + return params -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, chestnut_manifest): - params = mock.MagicMock() - - def get(key): - if key == "ModelManager_ModelsCache": - return qcom_manifest - if key == "ModelManager_ModelsCache_Chestnut": - return chestnut_manifest - if key in ("ModelManager_LastSyncTime", "ModelManager_LastSyncTime_Chestnut"): - return fresh_sync_time() - return None - - params.get.side_effect = get - return params +class TestModelFetcher(OpenpilotTestCase): + def refetch(self, fetcher, *bundles): + fetched = fetcher.model_parser.parse_models({"bundles": list(bundles)}) + return mock.patch.object(fetcher, "_fetch_and_cache_models", return_value=fetched) def test_active_source_follows_chestnut_presence(self): assert ModelFetcher.active_source(False) == "qcom" assert ModelFetcher.active_source(True) == "chestnut" - 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("chestnut")] == ["bbb"] - - def test_get_bundles_for_source_unknown(self): + def test_unknown_source_has_no_bundles(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") - chestnut_bundles = get_cached_bundles(params, "chestnut") - assert [b.ref for b in qcom_bundles] == ["aaa"] - assert [b.ref for b in chestnut_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, "chestnut") == [] - - def test_get_cached_bundles_unknown_source(self): assert get_cached_bundles(mock.MagicMock(), "bogus") == [] + def test_get_cached_bundles_reads_its_source(self): + params = manifest_params({"bundles": [manifest_bundle("small", "aaa")]}) + assert [b.ref for b in get_cached_bundles(params, "qcom")] == ["aaa"] + assert [b.ref for b in get_cached_bundles(params, "chestnut")] == ["bbb"] + 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, - "chestnut": ModelFetcher.MODEL_URL_CHESTNUT, - } - - - -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 (chestnut) 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, chestnut_manifest): - params = mock.MagicMock() - - def get(key): - if key == "ModelManager_ModelsCache": - return qcom_manifest - if key == "ModelManager_ModelsCache_Chestnut": - return chestnut_manifest - if key in ("ModelManager_LastSyncTime", "ModelManager_LastSyncTime_Chestnut"): - 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_chestnut_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("chestnut") - assert [bundle.ref for bundle in bundles] == ["bbb"] + params.put.assert_any_call("ModelManager_ActiveJson", {"qcom": ModelFetcher.MODEL_URL, "chestnut": ModelFetcher.MODEL_URL_CHESTNUT}, block=True) 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) + fetcher = ModelFetcher(manifest_params({"bundles": [manifest_bundle("small", "aaa")]})) 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("chestnut")] == ["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_unusable_cache_is_refetched(self): + stale = {**manifest_bundle("small", "aaa"), "minimum_selector_version": "16"} # parses to zero valid bundles + corrupt = {"short_name": "broken"} # missing required fields + for cached in (stale, corrupt): + with self.subTest(cached=cached): + fetcher = ModelFetcher(manifest_params({"bundles": [cached]})) + with self.refetch(fetcher, manifest_bundle("small2", "ddd")) as fetch: + assert [bundle.ref for bundle in fetcher.get_bundles_for_source("qcom")] == ["ddd"] + fetch.assert_called_once_with("qcom") - def test_mismatched_refetch_happens_once(self): - """If the fresh manifest still fails the source check, the URL is authoritative: - trust it instead of refetching at 1 Hz forever.""" - 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("big", "bbb", is_big=True)) - with mock.patch.object(fetcher, "_fetch_and_cache_models", return_value=fetched) as fetch: + def test_mismatched_cache_is_refetched_once(self): + fetcher = ModelFetcher(manifest_params({"bundles": [manifest_bundle("big", "bbb", is_big=True)]})) + with self.refetch(fetcher, manifest_bundle("big", "bbb", is_big=True)) as fetch: first = fetcher.get_bundles_for_source("qcom") second = fetcher.get_bundles_for_source("qcom") fetch.assert_called_once_with("qcom") - assert [bundle.ref for bundle in first] == ["bbb"] - assert [bundle.ref for bundle in second] == ["bbb"] - - 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"] + assert [bundle.ref for bundle in first] == [bundle.ref for bundle in second] == ["bbb"] -class TestActiveBundleValidation(OpenpilotTestCase): - """Validation is per-slot: a failed fetch (empty bundle list) must not reset a slot, - and resetting one slot must not stomp the runner cache derived from the other.""" +def ref_of(bundle) -> str | None: + return None if bundle is None else bundle.ref + +def slot_params(qcom=None, chestnut=None): + params = mock.MagicMock() + params.get.side_effect = {"ModelManager_ActiveBundle": qcom, "ModelManager_ActiveBundleChestnut": chestnut}.get + return params + + +class TestActiveBundle(OpenpilotTestCase): def setUp(self): super().setUp() helpers._LAST_VALIDATED_RAW.clear() - @staticmethod - def _raw_bundle(ref: str, runner: int | None = None) -> dict: - bundle = custom.ModelManagerSP.ModelBundle.new_message() - bundle.ref = ref - bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION - if runner is not None: - bundle.runner = runner - return bundle.to_dict() + def test_active_source(self): + for gpu, source in ((False, "qcom"), (True, "chestnut")): + with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=gpu): + assert get_active_source() == source + for state, source in ( + ((True, False, False, True), "chestnut"), # offroad + ((True, False, True, False), "chestnut"), # onroad, big model loading + ((True, True, False, False), "chestnut"), # onroad, big model running + ((True, False, False, False), "qcom"), # onroad, big model idle + ((True, None, False, False), "qcom"), # onroad, big model state not known yet + ((False, False, False, True), "qcom"), # offroad, no GPU + ): + chestnut, active, loading, offroad = state + assert get_active_source(chestnut=chestnut, chestnut_active=active, chestnut_loading=loading, offroad=offroad) == source, state - def _params(self, qcom=None, chestnut=None): - params = mock.MagicMock() - - def get(key, *args, **kwargs): - return {"ModelManager_ActiveBundle": qcom, "ModelManager_ActiveBundleChestnut": chestnut}.get(key) - - params.get.side_effect = get - return params + def test_active_bundle_is_the_active_slot(self): + params = slot_params(qcom=raw_bundle("small"), chestnut=raw_bundle("big")) + assert ref_of(get_selected_bundle(params, "qcom")) == "small" + assert ref_of(get_selected_bundle(params, "chestnut")) == "big" + for gpu, ref in ((False, "small"), (True, "big")): + with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=gpu): + assert ref_of(get_active_bundle(params)) == ref + with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=True): + assert get_active_bundle(slot_params(qcom=raw_bundle("small"))) is None, \ + "an empty chestnut slot must not fall back to the qcom pick" def test_empty_catalog_does_not_reset_slot(self): - params = self._params(qcom=self._raw_bundle("small")) + params = slot_params(qcom=raw_bundle("small")) with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=False): validate_active_bundles(params, {"qcom": [], "chestnut": []}) params.remove.assert_not_called() def test_reset_recomputes_runner_from_surviving_slot(self): tinygrad = int(custom.ModelManagerSP.Runner.tinygrad) - big_raw = self._raw_bundle("big", runner=tinygrad) - params = self._params(qcom=self._raw_bundle("gone"), chestnut=big_raw) - catalog = {"qcom": [custom.ModelManagerSP.ModelBundle(**self._raw_bundle("other"))], + big_raw = raw_bundle("big", runner=tinygrad) + params = slot_params(qcom=raw_bundle("gone"), chestnut=big_raw) + catalog = {"qcom": [custom.ModelManagerSP.ModelBundle(**raw_bundle("other"))], "chestnut": [custom.ModelManagerSP.ModelBundle(**big_raw)]} with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=True): validate_active_bundles(params, catalog) params.remove.assert_called_once_with("ModelManager_ActiveBundle") - runner_puts = [call for call in params.put.call_args_list if call.args[0] == "ModelRunnerTypeCache"] - assert [call.args[1] for call in runner_puts] == [tinygrad] - - -class TestActiveBundleSelection(OpenpilotTestCase): - """The effective active bundle is the active source's slot: chestnut when a GPU is - present, qcom otherwise. An empty active slot means the hardware default (stock - runner), never the other slot's pick - modeld_v2 requires a real bundle.""" - - @staticmethod - def _raw_bundle(ref: str) -> dict: - bundle = custom.ModelManagerSP.ModelBundle.new_message() - bundle.ref = ref - bundle.minimumSelectorVersion = helpers.REQUIRED_JSON_VERSION - return bundle.to_dict() - - def _params(self, qcom=None, chestnut=None): - params = mock.MagicMock() - - def get(key, *args, **kwargs): - if key == "ModelManager_ActiveBundle": - return qcom - if key == "ModelManager_ActiveBundleChestnut": - return chestnut - 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"), chestnut=self._raw_bundle("big")) - assert get_selected_bundle(params, "qcom").ref == "small" - assert get_selected_bundle(params, "chestnut").ref == "big" - - def test_no_gpu_uses_qcom_slot(self): - params = self._params(qcom=self._raw_bundle("small"), chestnut=self._raw_bundle("big")) - with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=False): - assert get_active_bundle(params).ref == "small" - - def test_gpu_uses_chestnut_slot(self): - params = self._params(qcom=self._raw_bundle("small"), chestnut=self._raw_bundle("big")) - with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=True): - assert get_active_bundle(params).ref == "big" - - def test_gpu_without_big_selection_is_hardware_default(self): - params = self._params(qcom=self._raw_bundle("small"), chestnut=None) - with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=True): - assert get_active_bundle(params) is None - - -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 = helpers.REQUIRED_JSON_VERSION - return bundle.to_dict() - - def test_runtime_no_gpu(self): - with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=False): - assert get_active_source() == "qcom" - - def test_runtime_gpu_present(self): - with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=True): - assert get_active_source() == "chestnut" - - def test_display_offroad_gpu_present_shows_big(self): - assert get_active_source(chestnut=True, chestnut_active=False, chestnut_loading=False, offroad=True) == "chestnut" - - def test_display_onroad_gpu_loading_shows_big(self): - assert get_active_source(chestnut=True, chestnut_active=False, chestnut_loading=True, offroad=False) == "chestnut" - - def test_display_onroad_gpu_active_shows_big(self): - assert get_active_source(chestnut=True, chestnut_active=True, chestnut_loading=False, offroad=False) == "chestnut" - - def test_display_onroad_gpu_idle_shows_small(self): - assert get_active_source(chestnut=True, chestnut_active=False, chestnut_loading=False, offroad=False) == "qcom" - - def test_display_active_none_is_idle(self): - assert get_active_source(chestnut=True, chestnut_active=None, chestnut_loading=False, offroad=False) == "qcom" - - def test_active_bundle_follows_source(self): - params = mock.MagicMock() - params.get.side_effect = lambda key: {"ModelManager_ActiveBundle": self._raw_bundle("small"), - "ModelManager_ActiveBundleChestnut": self._raw_bundle("big")}.get(key) - with mock.patch("openpilot.sunnypilot.models.helpers.chestnut_present", return_value=False): - assert get_active_bundle(params).ref == "small" - assert get_selected_bundle(params, get_active_source(chestnut=True, chestnut_active=False, - chestnut_loading=False, offroad=True)).ref == "big" + params.put.assert_called_once_with("ModelRunnerTypeCache", tinygrad, block=True) @unittest.skipUnless(os.environ.get('RUN_INTEGRATION_TESTS'), 'requires external network') class TestLiveModelManifest(OpenpilotTestCase): - """Every artifact URL in the published manifest must resolve.""" - def test_all_manifest_urls_available(self): from openpilot.sunnypilot.models.fetcher import ModelFetcher @@ -687,15 +520,13 @@ class TestLiveModelManifest(OpenpilotTestCase): for bundle in manifest.get('bundles', []): for model in bundle.get('models', []): artifact = model['artifact'] - urls = [artifact['download_uri']['url']] - for u in urls: - try: - r = session.head(u, timeout=15, allow_redirects=True) - if r.status_code != 200: - dead.append(f"{bundle.get('short_name')}: HTTP {r.status_code} {u}") - except requests.RequestException as e: - dead.append(f"{bundle.get('short_name')}: {type(e).__name__} {u}") - + url = artifact['download_uri']['url'] + try: + r = session.head(url, timeout=15, allow_redirects=True) + if r.status_code != 200: + dead.append(f"{bundle.get('short_name')}: HTTP {r.status_code} {url}") + except requests.RequestException as e: + dead.append(f"{bundle.get('short_name')}: {type(e).__name__} {url}") assert not dead, "unreachable model URLs:\n" + "\n".join(dead) diff --git a/openpilot/system/ui/sunnypilot/widgets/download_status.py b/openpilot/system/ui/sunnypilot/widgets/download_status.py index c85814d0e..e0f45186b 100644 --- a/openpilot/system/ui/sunnypilot/widgets/download_status.py +++ b/openpilot/system/ui/sunnypilot/widgets/download_status.py @@ -57,7 +57,7 @@ class DownloadStatusAction(ItemAction): self.icon: str | None = None self.icon_color: rl.Color | None = None self._font = gui_app.font(FontWeight.NORMAL) - # raw progress arrives in steps, one per 128KB chunk the manager publishes + # raw progress arrives in steps, one per progress report from the manager self._progress = FirstOrderFilter(0.0, 0.5, 1 / gui_app.target_fps) # integrated per frame; (t * speed) % span jumps whenever the fill width changes self._sweep = 0.0