models: bind a download to its ref so cancel and reselect work everywhere

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"):

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 = {}