From b0a63f198b4ba343e9d5cfcd779293653f2a43fc Mon Sep 17 00:00:00 2001 From: Jason Wen Date: Wed, 26 Aug 2026 03:08:04 -0400 Subject: [PATCH] models: bind a download to its ref so cancel and reselect work everywhere --- openpilot/sunnypilot/models/manager.py | 23 +++++++- .../models/tests/test_manager_download.py | 56 +++++++++++++++++++ 2 files changed, 76 insertions(+), 3 deletions(-) diff --git a/openpilot/sunnypilot/models/manager.py b/openpilot/sunnypilot/models/manager.py index 2405566d55..930eddef6e 100644 --- a/openpilot/sunnypilot/models/manager.py +++ b/openpilot/sunnypilot/models/manager.py @@ -39,6 +39,17 @@ class ModelManagerSP: 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 + self._download_ref: bytes | str | None = None + + def _download_interrupted(self) -> bool: + # DownloadRef removed means cancel; a different ref means the user picked + # another model mid-download. Either way this download must stop. + return self.params.get("ModelManager_DownloadRef") != self._download_ref + + def _release_download_ref(self) -> None: + if not self._download_interrupted(): + self.params.remove("ModelManager_DownloadRef") + 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.""" @@ -80,7 +91,7 @@ class ModelManagerSP: f.write(chunk) bytes_downloaded += len(chunk) - if self.params.get("ModelManager_DownloadRef") is None: + if self._download_interrupted(): raise Exception("Download cancelled") if total_size > 0: @@ -118,7 +129,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_DownloadRef") is None: + if self._download_interrupted(): raise Exception("Download cancelled") intra = chunk_downloaded / max(chunk_size, 1) progress = min(99.0, ((i + intra) / num_chunks) * 100) @@ -137,6 +148,9 @@ class ModelManagerSP: async def _process_artifact(self, artifact, destination_path: str) -> None: if not artifact.downloadUri.uri: return None + if self._download_interrupted(): + # raised before the try so a cancel never deletes files already on disk + raise Exception("Download cancelled") url = artifact.downloadUri.uri expected_hash = artifact.downloadUri.sha256 @@ -242,6 +256,8 @@ class ModelManagerSP: seen_artifacts.add(artifact.fileName) await self._process_artifact(artifact, destination_path) + if self._download_interrupted(): + raise Exception("Download cancelled") 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) @@ -274,12 +290,13 @@ class ModelManagerSP: 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 + self._download_ref = ref_to_download try: self.download(model_to_download, Paths.model_root(), source) except Exception as e: cloudlog.exception(e) finally: - self.params.remove("ModelManager_DownloadRef") + self._release_download_ref() self.selected_bundle = None if self.params.get("ModelManager_ClearCache"): diff --git a/openpilot/sunnypilot/models/tests/test_manager_download.py b/openpilot/sunnypilot/models/tests/test_manager_download.py index d74deb03e6..2d990afa52 100644 --- a/openpilot/sunnypilot/models/tests/test_manager_download.py +++ b/openpilot/sunnypilot/models/tests/test_manager_download.py @@ -103,6 +103,7 @@ class ManagerDownloadTestBase(OpenpilotTestCase): self.manager = ModelManagerSP.__new__(ModelManagerSP) self.manager.params = mock.MagicMock() self.manager.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.manager.selected_bundle = None @@ -261,6 +262,7 @@ class TestManagerDownload(ManagerDownloadTestBase): 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 + self.manager._download_ref = b"ref" 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) @@ -279,12 +281,66 @@ class TestManagerDownload(ManagerDownloadTestBase): 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_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 test_replaced_download_ref_cancels_transfer(self): + """Selecting another model mid-transfer cancels the running 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 b"other-ref" + 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_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 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): + """Cancelling must stop an already-on-disk bundle before it is applied to the slot.""" + def body(): + artifact = self.make_artifact(chunked=True) + base_path = os.path.join(self.dest, artifact.fileName) + for i, data in enumerate(CHUNK_BODIES): + with open(get_chunk_name(base_path, i, len(CHUNK_BODIES)), 'wb') as f: + f.write(data) + self._bundle.ref = "test-ref" + params, store = self._make_params_with_store() + self.manager.params = params + self.manager._download_ref = b"ref" # store has no DownloadRef -> cancelled + 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 all(os.path.isfile(p) for p in self.chunk_paths(base_path)), "cancel must not delete cached chunks" + self.run_with_server(body) + def _make_params_with_store(self): params = mock.MagicMock() store = {}