diff --git a/openpilot/sunnypilot/models/helpers.py b/openpilot/sunnypilot/models/helpers.py index 707b86f722..8e045e11e8 100644 --- a/openpilot/sunnypilot/models/helpers.py +++ b/openpilot/sunnypilot/models/helpers.py @@ -41,12 +41,7 @@ def _compute_hash(file_path: str) -> str | None: return None -async def verify_file(file_path: str, expected_hash: str) -> bool: - file_hash = _compute_hash(file_path) - return file_hash == expected_hash.lower() if file_hash else False - - -def _verify_file(file_path: str, expected_hash: str) -> bool: +def verify_file(file_path: str, expected_hash: str) -> bool: file_hash = _compute_hash(file_path) return file_hash == expected_hash.lower() if file_hash else False @@ -81,7 +76,7 @@ def _bundle_artifacts(bundle: custom.ModelManagerSP.ModelBundle) -> list[tuple[s def _bundle_is_valid_locally(bundle: custom.ModelManagerSP.ModelBundle) -> bool: model_root = Paths.model_root() - return all(_verify_file(os.path.join(model_root, file_name), expected_hash) + return all(verify_file(os.path.join(model_root, file_name), expected_hash) for file_name, expected_hash in _bundle_artifacts(bundle)) diff --git a/openpilot/sunnypilot/models/manager.py b/openpilot/sunnypilot/models/manager.py index 178d6c04e5..c78f7d95b9 100644 --- a/openpilot/sunnypilot/models/manager.py +++ b/openpilot/sunnypilot/models/manager.py @@ -10,6 +10,7 @@ import os import time import requests +from openpilot.common.file_chunker import get_chunk_name from openpilot.common.params import Params from openpilot.common.realtime import Ratekeeper from openpilot.common.swaglog import cloudlog @@ -22,6 +23,8 @@ 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) +# how many download+verify rounds before giving up on a chunk that won't verify +MAX_CHUNK_VERIFY_ATTEMPTS = 3 class DownloadCancelled(Exception): @@ -153,6 +156,35 @@ class ModelManagerSP: os.remove(base_path) del self._download_start_times[artifact.fileName] + async def _verify_chunks_parallel(self, artifact, full_path: str) -> set[int]: + num_chunks = len(artifact.chunks) + start = time.monotonic() + futures = [] + chunk_index: dict[asyncio.Future, int] = {} + for i, chunk in enumerate(artifact.chunks): + fut = asyncio.ensure_future(asyncio.to_thread(verify_file, get_chunk_name(full_path, i, num_chunks), chunk.sha256)) + futures.append(fut) + chunk_index[fut] = i + + valid_chunks: set[int] = set() + pending = set(futures) + while pending: + if self._download_interrupted(): + for t in pending: + t.cancel() + raise DownloadCancelled("Download cancelled") + done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED) + for fut in done: + if fut.result(): + valid_chunks.add(chunk_index[fut]) + artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.verifying + artifact.downloadProgress.progress = (len(valid_chunks) / num_chunks) * 100 + self._sync_artifact_progress(artifact) + self._report_status() + elapsed = time.monotonic() - start + cloudlog.info(f"Verified {len(valid_chunks)}/{num_chunks} chunks of {artifact.fileName} in {elapsed:.2f}s") + return valid_chunks + async def _process_artifact(self, artifact, destination_path: str) -> None: if not artifact.downloadUri.uri: return None @@ -170,20 +202,10 @@ class ModelManagerSP: is_cached = False valid_chunks: set[int] = set() if len(artifact.chunks) > 0: - from openpilot.common.file_chunker import get_chunk_name - num_chunks = len(artifact.chunks) - for i, chunk in enumerate(artifact.chunks): - if self._download_interrupted(): - raise DownloadCancelled("Download cancelled") - if await verify_file(get_chunk_name(full_path, i, num_chunks), chunk.sha256): - valid_chunks.add(i) - artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.verifying - artifact.downloadProgress.progress = (len(valid_chunks) / num_chunks) * 100 - self._sync_artifact_progress(artifact) - self._report_status() - is_cached = len(valid_chunks) == num_chunks + valid_chunks = await self._verify_chunks_parallel(artifact, full_path) + is_cached = len(valid_chunks) == len(artifact.chunks) else: - if await verify_file(full_path, expected_hash): + if verify_file(full_path, expected_hash): is_cached = True if is_cached: @@ -195,15 +217,18 @@ class ModelManagerSP: return if len(artifact.chunks) > 0: - await self._download_chunked(url, full_path, artifact, skip=valid_chunks) - from openpilot.common.file_chunker import get_chunk_name - for i, chunk in enumerate(artifact.chunks): - chunk_path = get_chunk_name(full_path, i, len(artifact.chunks)) - if not await verify_file(chunk_path, chunk.sha256): - raise ValueError(f"Hash validation failed for chunk {i+1} of {filename}") + attempts = 0 + while len(valid_chunks) < len(artifact.chunks): + attempts += 1 + cloudlog.warning(f"Re-downloading {len(artifact.chunks) - len(valid_chunks)} invalid chunk(s) of {filename} (attempt {attempts})") + await self._download_chunked(url, full_path, artifact, skip=valid_chunks) + valid_chunks = await self._verify_chunks_parallel(artifact, full_path) + if len(valid_chunks) < len(artifact.chunks) and attempts >= MAX_CHUNK_VERIFY_ATTEMPTS: + missing = next(i for i in range(len(artifact.chunks)) if i not in valid_chunks) + raise ValueError(f"Hash validation failed for chunk {missing+1} of {filename} after {attempts} attempts") else: await self._download_file(url, full_path, artifact) - if not await verify_file(full_path, expected_hash): + if not verify_file(full_path, expected_hash): raise ValueError(f"Hash validation failed for {filename}") artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloaded