diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py index 4f1b30a7..669dd3c6 100644 --- a/py/routes/handlers/model_handlers.py +++ b/py/routes/handlers/model_handlers.py @@ -1904,8 +1904,18 @@ class ModelDownloadHandler: try: status_filter = request.query.get("status") or None service = await DownloadQueueService.get_instance() - cleared = await service.clear_queue(status_filter=status_filter) - return web.json_response({"success": True, "cleared": cleared}) + cleared_ids = await service.clear_queue(status_filter=status_filter) + # Clearing the queue rows alone would orphan any in-memory tasks + # and persisted aria2 state for those downloads, leaving them + # polling the daemon invisibly. Tear that tracking down too. + try: + await self._download_coordinator.discard_cleared_downloads(cleared_ids) + except Exception: + self._logger.warning( + "Failed to discard in-memory state for cleared downloads", + exc_info=True, + ) + return web.json_response({"success": True, "cleared": len(cleared_ids)}) except Exception as exc: self._logger.error( "Error clearing download queue: %s", exc, exc_info=True diff --git a/py/services/aria2_downloader.py b/py/services/aria2_downloader.py index e5541a96..9acb7aab 100644 --- a/py/services/aria2_downloader.py +++ b/py/services/aria2_downloader.py @@ -217,8 +217,9 @@ class Aria2Downloader: """Call get_status with retry for transient RPC failures. Only retries on :exc:`Aria2Error` (RPC-level failure). Returns - ``None`` immediately when the download_id is not tracked (a missing - transfer is not a transient condition, so retrying is pointless). + ``None`` immediately when the transfer is not tracked or its GID is + gone from the daemon (a missing transfer is not a transient + condition, so retrying is pointless). A single failed RPC call should not immediately fail the download, because aria2 may be temporarily busy (e.g. finalizing multiple @@ -332,7 +333,13 @@ class Aria2Downloader: return transfer async def get_status(self, download_id: str) -> Optional[Dict[str, Any]]: - """Return the raw aria2 status payload for a known download.""" + """Return the raw aria2 status payload for a known download. + + Returns ``None`` when the download_id is not tracked or the daemon no + longer knows the transfer's GID (daemon restart / forceRemove). A + forgotten GID is permanent, not transient, so the caller's recovery + path handles it instead of burning retry attempts on a dead GID. + """ transfer = self._transfers.get(download_id) if transfer is None: @@ -348,8 +355,17 @@ class Aria2Downloader: "files", ] try: - status = await self._rpc_call("aria2.tellStatus", [transfer.gid, keys]) + status = await self._rpc_call( + "aria2.tellStatus", [transfer.gid, keys], log_errors=False + ) except Exception as exc: + if "not found" in str(exc).lower(): + logger.debug( + "aria2 GID %s for download %s is gone; treating as lost transfer", + transfer.gid, + download_id, + ) + return None raise Aria2Error(f"Failed to query aria2 download status: {exc}") from exc if isinstance(status, dict): @@ -367,7 +383,9 @@ class Aria2Downloader: "files", ] try: - status = await self._rpc_call("aria2.tellStatus", [gid, keys]) + status = await self._rpc_call( + "aria2.tellStatus", [gid, keys], log_errors=False + ) except Exception as exc: message = str(exc) if "cannot be found" in message.lower() or "not found" in message.lower(): @@ -434,8 +452,19 @@ class Aria2Downloader: try: await self._rpc_call("aria2.forceRemove", [transfer.gid]) except Exception as exc: - return {"success": False, "error": str(exc)} + if "not found" not in str(exc).lower(): + return {"success": False, "error": str(exc)} + # The daemon already forgot this GID (restart / prior removal), + # so the transfer is effectively cancelled. + logger.debug( + "aria2 GID %s for download %s already gone during cancel", + transfer.gid, + download_id, + ) + # Drop the in-memory entry as well so a concurrent poll loop does + # not mistake the removal for a lost transfer and re-register it. + self._transfers.pop(download_id, None) await self._state_store.remove(download_id) return {"success": True, "message": "Download cancelled successfully"} @@ -725,7 +754,9 @@ class Aria2Downloader: return isinstance(result, dict) - async def _rpc_call(self, method: str, params: list[Any]) -> Any: + async def _rpc_call( + self, method: str, params: list[Any], *, log_errors: bool = True + ) -> Any: if not self._rpc_url: raise Aria2Error("aria2 RPC endpoint is not initialized") @@ -756,7 +787,10 @@ class Aria2Downloader: error = body["error"] or {} code = error.get("code") if isinstance(error, dict) else None message = error.get("message") if isinstance(error, dict) else str(error) - logger.error( + # Probing calls (e.g. tellStatus for a GID the daemon may have + # forgotten) pass log_errors=False: an expected "not found" must + # not spam the log at ERROR level. + (logger.error if log_errors else logger.debug)( "aria2 RPC %s failed with HTTP %s, code=%s, message=%s", method, response.status, @@ -771,7 +805,7 @@ class Aria2Downloader: raise Aria2Error(status_message or "Unknown aria2 RPC error") if response.status != 200: - logger.error( + (logger.error if log_errors else logger.debug)( "aria2 RPC %s returned unexpected HTTP status %s without error payload: %s", method, response.status, diff --git a/py/services/download_coordinator.py b/py/services/download_coordinator.py index 4577ba57..7bc65db6 100644 --- a/py/services/download_coordinator.py +++ b/py/services/download_coordinator.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging -from typing import Any, Awaitable, Callable, Dict, Optional +from typing import Any, Awaitable, Callable, Dict, Iterable, Optional from .downloader import DownloadProgress @@ -186,6 +186,14 @@ class DownloadCoordinator: download_manager = await self._download_manager_factory() return await download_manager.get_active_downloads() + async def discard_cleared_downloads(self, download_ids: Iterable[str]) -> int: + """Tear down in-memory/aria2 tracking for queue-cleared downloads.""" + + if not download_ids: + return 0 + download_manager = await self._download_manager_factory() + return await download_manager.discard_cleared_downloads(download_ids) + def _parse_optional_int(self, value: Any, field: str) -> Optional[int]: """Parse an optional integer from user input.""" diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 959e62df..5948bf21 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -13,7 +13,7 @@ import zipfile from concurrent.futures import ThreadPoolExecutor from collections import OrderedDict import uuid -from typing import Any, Dict, List, Optional, Set, Tuple, cast +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, cast from urllib.parse import urlparse from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata from ..utils.constants import ( @@ -1100,6 +1100,11 @@ class DownloadManager: save_path = self._resolve_save_path_from_persisted_record(record) if save_path is None: + # No resolvable target path (e.g. a queued download whose + # paths were never resolved before shutdown): the record + # can never be restored, so drop it instead of letting it + # accumulate in the state store forever. + await self._aria2_state_store.remove(download_id) continue if ( @@ -2897,6 +2902,64 @@ class DownloadManager: # Preserve aria2 state store entry so the partial download # info survives restarts and can be resumed later + async def discard_cleared_downloads(self, download_ids: Iterable[str]) -> int: + """Stop in-memory tracking for downloads cleared from the queue. + + Cancels asyncio tasks, removes live aria2 transfers and drops the + persisted aria2 state so cleared downloads cannot keep polling the + daemon or be resurrected as ghost entries on the next restart. + Partial files on disk are preserved; unlike ``cancel_download`` no + files are deleted. + + Returns the number of downloads that had any in-memory or persisted + tracking removed. + """ + discarded = 0 + aria2_downloader = None + + for download_id in download_ids: + task = self._download_tasks.get(download_id) + info = self._active_downloads.get(download_id) + persisted = await self._aria2_state_store.get(download_id) + if task is None and info is None and persisted is None: + continue + + discarded += 1 + + if task is not None: + task.cancel() + + pause_control = self._pause_events.pop(download_id, None) + if pause_control is not None: + pause_control.resume() + + if task is not None: + try: + await asyncio.wait_for(asyncio.shield(task), timeout=2.0) + except (asyncio.CancelledError, asyncio.TimeoutError): + pass + + self._download_tasks.pop(download_id, None) + self._active_downloads.pop(download_id, None) + + backend = (info or persisted or {}).get("transfer_backend") or "python" + if backend == "aria2": + if aria2_downloader is None: + aria2_downloader = await get_aria2_downloader() + if await aria2_downloader.has_transfer(download_id): + try: + await aria2_downloader.cancel_download(download_id) + except Exception as exc: + logger.warning( + "Failed to remove aria2 transfer for cleared download %s: %s", + download_id, + exc, + ) + + await self._aria2_state_store.remove(download_id) + + return discarded + async def pause_download(self, download_id: str) -> Dict[str, Any]: """Pause an active download without losing progress.""" diff --git a/py/services/download_queue_service.py b/py/services/download_queue_service.py index 7eae08e5..2f2fed89 100644 --- a/py/services/download_queue_service.py +++ b/py/services/download_queue_service.py @@ -6,7 +6,7 @@ import logging import os import sqlite3 import time -from typing import Any, Optional +from typing import Any, List, Optional from ..utils.cache_paths import get_cache_base_dir @@ -390,23 +390,31 @@ class DownloadQueueService: conn.commit() return True - async def clear_queue(self, status_filter: Optional[str] = None) -> int: + async def clear_queue(self, status_filter: Optional[str] = None) -> List[str]: """Remove items from the queue. When *status_filter* is provided only items with that status are - deleted. Returns the number of deleted rows. + deleted. Returns the ``download_id`` values of the deleted rows so + callers can also tear down any in-memory tracking for them. """ async with self._lock: conn = self._get_conn() if status_filter is not None: - cursor = conn.execute( + rows = conn.execute( + "SELECT download_id FROM download_queue WHERE status = ?", + (status_filter,), + ).fetchall() + conn.execute( "DELETE FROM download_queue WHERE status = ?", (status_filter,), ) else: - cursor = conn.execute("DELETE FROM download_queue") + rows = conn.execute( + "SELECT download_id FROM download_queue" + ).fetchall() + conn.execute("DELETE FROM download_queue") conn.commit() - return cursor.rowcount + return [row["download_id"] for row in rows] async def complete_download( self, diff --git a/tests/services/test_aria2_downloader.py b/tests/services/test_aria2_downloader.py index 3feead5f..67665a55 100644 --- a/tests/services/test_aria2_downloader.py +++ b/tests/services/test_aria2_downloader.py @@ -57,7 +57,7 @@ async def test_download_file_polls_until_complete(tmp_path, monkeypatch): ] ) - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): rpc_calls.append((method, params)) if method == "aria2.addUri": return "gid-1" @@ -139,7 +139,7 @@ async def test_download_file_keeps_auth_headers_when_civitai_does_not_redirect( ] ) - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): rpc_calls.append((method, params)) if method == "aria2.addUri": return "gid-1" @@ -178,7 +178,7 @@ async def test_pause_resume_cancel_forward_to_rpc(monkeypatch): calls = [] - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): calls.append((method, params)) return "gid-1" @@ -232,7 +232,7 @@ async def test_download_file_reuses_existing_transfer_without_add_uri( ] ) - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): rpc_calls.append((method, params)) if method == "aria2.tellStatus": return next(statuses) @@ -265,7 +265,7 @@ async def test_download_file_recovers_when_transfer_lost_mid_poll( add_uri_count = {"n": 0} poll_count = {"n": 0} - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): if method == "aria2.addUri": add_uri_count["n"] += 1 return "gid-1" if add_uri_count["n"] == 1 else "gid-2" @@ -317,7 +317,7 @@ async def test_download_file_recovers_when_rpc_fails_mid_poll(tmp_path, monkeypa add_uri_count = {"n": 0} poll_count = {"n": 0} - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): if method == "aria2.addUri": add_uri_count["n"] += 1 return "gid-1" if add_uri_count["n"] == 1 else "gid-2" @@ -366,7 +366,7 @@ async def test_download_file_fails_after_recovery_attempts_exhausted( save_path = tmp_path / "downloads" / "model.safetensors" add_uri_count = {"n": 0} - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): if method == "aria2.addUri": add_uri_count["n"] += 1 return f"gid-{add_uri_count['n']}" @@ -402,7 +402,7 @@ async def test_download_file_concurrent_same_id_schedules_once(tmp_path, monkeyp add_uri_count = {"n": 0} poll_count = {"n": 0} - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): if method == "aria2.addUri": add_uri_count["n"] += 1 return "gid-1" @@ -458,7 +458,7 @@ async def test_download_file_cleanup_preserves_newer_registration(tmp_path, monk save_path = tmp_path / "downloads" / "model.safetensors" poll_count = {"n": 0} - async def fake_rpc_call(method, params): + async def fake_rpc_call(method, params, **_kwargs): if method == "aria2.addUri": return "gid-1" if method == "aria2.tellStatus": @@ -808,3 +808,121 @@ def test_stderr_error_report_prunes_expired_entries(): assert old_line not in downloader._stderr_error_report assert new_line in downloader._stderr_error_report + + +@pytest.mark.asyncio +async def test_get_status_returns_none_without_retry_when_gid_not_found(monkeypatch): + """A forgotten GID is permanent: no retry attempts on a dead GID.""" + downloader = Aria2Downloader() + downloader._transfers["download-1"] = Aria2Transfer( + gid="gone-gid", save_path="/tmp/model.safetensors" + ) + + calls = [] + + async def fake_rpc_call(method, params, **_kwargs): + calls.append(method) + raise Aria2Error("GID gone-gid is not found") + + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock()) + + assert await downloader._get_status_with_retry("download-1") is None + assert calls == ["aria2.tellStatus"] + + +@pytest.mark.asyncio +async def test_get_status_still_raises_on_transient_rpc_error(monkeypatch): + downloader = Aria2Downloader() + downloader._transfers["download-1"] = Aria2Transfer( + gid="gid-1", save_path="/tmp/model.safetensors" + ) + + async def fake_rpc_call(method, params, **_kwargs): + raise Aria2Error("connection reset") + + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock()) + + with pytest.raises(Aria2Error, match="Failed to query aria2 download status"): + await downloader._get_status_with_retry("download-1") + + +@pytest.mark.asyncio +async def test_cancel_download_pops_transfer_on_success(monkeypatch): + downloader = Aria2Downloader() + downloader._transfers["download-1"] = Aria2Transfer( + gid="gid-1", save_path="/tmp/model.safetensors" + ) + + async def fake_rpc_call(method, params, **_kwargs): + return "gid-1" + + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + + result = await downloader.cancel_download("download-1") + + assert result["success"] is True + assert "download-1" not in downloader._transfers + + +@pytest.mark.asyncio +async def test_cancel_download_tolerates_missing_gid(monkeypatch): + """Cancelling a transfer the daemon already forgot still succeeds.""" + downloader = Aria2Downloader() + downloader._transfers["download-1"] = Aria2Transfer( + gid="gone-gid", save_path="/tmp/model.safetensors" + ) + await downloader._state_store.upsert( + "download-1", {"gid": "gone-gid", "status": "downloading"} + ) + + async def fake_rpc_call(method, params, **_kwargs): + raise Aria2Error("GID gone-gid is not found") + + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + + result = await downloader.cancel_download("download-1") + + assert result["success"] is True + assert "download-1" not in downloader._transfers + assert await downloader._state_store.get("download-1") is None + + +@pytest.mark.asyncio +async def test_rpc_call_suppresses_error_log_when_log_errors_false( + monkeypatch, caplog +): + """Probing calls must not spam ERROR for an expected failure.""" + + class FakeResponse: + status = 400 + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def text(self): + return '{"jsonrpc": "2.0", "error": {"code": 1, "message": "GID x is not found"}}' + + class FakeSession: + closed = False + + def post(self, *args, **kwargs): + return FakeResponse() + + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + monkeypatch.setattr(downloader, "_get_rpc_session", AsyncMock(return_value=FakeSession())) + + with caplog.at_level(logging.DEBUG, logger="py.services.aria2_downloader"): + with pytest.raises(Aria2Error, match="not found"): + await downloader._rpc_call("aria2.tellStatus", ["x"], log_errors=False) + + error_records = [r for r in caplog.records if r.levelno == logging.ERROR] + debug_records = [r for r in caplog.records if r.levelno == logging.DEBUG] + assert error_records == [] + assert any("GID x is not found" in r.message for r in debug_records) diff --git a/tests/services/test_download_manager_basic.py b/tests/services/test_download_manager_basic.py index ae807463..c58503f4 100644 --- a/tests/services/test_download_manager_basic.py +++ b/tests/services/test_download_manager_basic.py @@ -2003,3 +2003,98 @@ def test_resolve_target_file_returns_none_for_no_match(): assert DownloadManager._resolve_target_file(files, {"id": 9999}) is None assert DownloadManager._resolve_target_file(files, None) is None assert DownloadManager._resolve_target_file(files, {}) is None + + +@pytest.mark.asyncio +async def test_restore_drops_unrestorable_persisted_records(monkeypatch, tmp_path): + """Records without any resolvable target path can never be restored; + the restore sweep must delete them instead of skipping them forever.""" + manager = DownloadManager() + + await manager._aria2_state_store.upsert( + "download-orphan", + { + "download_id": "download-orphan", + "transfer_backend": "aria2", + "status": "failed", + # no save_path / file_path / resume_context + }, + ) + + class DummyAria2Downloader: + async def get_status_by_gid(self, gid): + return None + + monkeypatch.setattr( + download_manager, + "get_aria2_downloader", + AsyncMock(return_value=DummyAria2Downloader()), + ) + + downloads = await manager.get_active_downloads() + + assert downloads["downloads"] == [] + assert await manager._aria2_state_store.get("download-orphan") is None + + +@pytest.mark.asyncio +async def test_discard_cleared_downloads_stops_tracking_and_preserves_files( + monkeypatch, tmp_path +): + manager = DownloadManager() + + save_path = tmp_path / "file.safetensors" + save_path.write_text("partial") + control_path = tmp_path / "file.safetensors.aria2" + control_path.write_text("control") + + async def _pending(): + await asyncio.sleep(3600) + + task = asyncio.create_task(_pending()) + manager._download_tasks["download-1"] = task + manager._pause_events["download-1"] = download_manager.DownloadStreamControl() + manager._active_downloads["download-1"] = { + "status": "downloading", + "transfer_backend": "aria2", + "file_path": str(save_path), + } + await manager._aria2_state_store.upsert( + "download-1", + { + "download_id": "download-1", + "transfer_backend": "aria2", + "status": "downloading", + "save_path": str(save_path), + "gid": "gid-1", + }, + ) + + cancelled = [] + + class DummyAria2Downloader: + async def has_transfer(self, download_id): + return True + + async def cancel_download(self, download_id): + cancelled.append(download_id) + return {"success": True} + + monkeypatch.setattr( + download_manager, + "get_aria2_downloader", + AsyncMock(return_value=DummyAria2Downloader()), + ) + + discarded = await manager.discard_cleared_downloads(["download-1", "unknown-id"]) + + assert discarded == 1 + assert cancelled == ["download-1"] + assert task.cancelled() + assert "download-1" not in manager._download_tasks + assert "download-1" not in manager._active_downloads + assert "download-1" not in manager._pause_events + assert await manager._aria2_state_store.get("download-1") is None + # Partial files are preserved for a future resume from disk. + assert save_path.exists() + assert control_path.exists() diff --git a/tests/services/test_download_queue_service.py b/tests/services/test_download_queue_service.py index 48b2f711..e63c4147 100644 --- a/tests/services/test_download_queue_service.py +++ b/tests/services/test_download_queue_service.py @@ -495,3 +495,29 @@ async def test_dedup_history_collapses_same_file_and_legacy_rows(tmp_path: Path) history = await svc.get_history() remaining = {item["download_id"] for item in history["items"]} assert remaining == {"dl-x2", "dl-y2"} + + +# --------------------------------------------------------------------------- +# clear_queue +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_clear_queue_returns_deleted_ids(tmp_path: Path) -> None: + """clear_queue reports the download_ids it removed so callers can tear + down any in-memory tracking for them.""" + svc = _make_service(tmp_path) + await svc.add_to_queue(download_id="dl-1", model_id=1) + await svc.add_to_queue(download_id="dl-2", model_id=2) + await svc.add_to_queue(download_id="dl-3", model_id=3) + await svc.update_status("dl-3", "downloading") + + cleared = await svc.clear_queue(status_filter="queued") + assert sorted(cleared) == ["dl-1", "dl-2"] + + remaining = await svc.get_queue() + assert [row["download_id"] for row in remaining] == ["dl-3"] + + cleared_all = await svc.clear_queue() + assert cleared_all == ["dl-3"] + assert await svc.get_queue() == []