121 lines
3.9 KiB
Python
121 lines
3.9 KiB
Python
"""Download progress and cancellation state operations."""
|
|
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable, Dict, Optional
|
|
|
|
# Shared mutable download state has a single owner in this module.
|
|
download_progress: Dict[str, Dict[str, Any]] = {}
|
|
download_lock = threading.Lock()
|
|
cancelled_downloads: set = set()
|
|
aria2_lock = threading.RLock()
|
|
aria2_transfers: Dict[str, Dict[str, Any]] = {}
|
|
aria2_action_locks: Dict[str, Any] = {}
|
|
aria2_desired_states: Dict[str, Dict[str, Any]] = {}
|
|
xet_transfers: Dict[str, Dict[str, Any]] = {}
|
|
xet_transfers_lock = threading.Lock()
|
|
|
|
|
|
def create_initial_progress(
|
|
*,
|
|
url: str,
|
|
path: str,
|
|
filename: str,
|
|
directory: str,
|
|
download_backend: str,
|
|
start_time: float,
|
|
total_size: int = 0,
|
|
) -> Dict[str, Any]:
|
|
"""Build the shared initial progress record for a download."""
|
|
return {
|
|
"status": "starting",
|
|
"progress": 0,
|
|
"total_size": total_size,
|
|
"downloaded": 0,
|
|
"filename": filename,
|
|
"path": path,
|
|
"directory": directory,
|
|
"url": url,
|
|
"error": None,
|
|
"speed": 0,
|
|
"start_time": start_time,
|
|
"download_backend": download_backend,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DownloadStateDependencies:
|
|
"""State and side-effect dependencies supplied by the composition layer."""
|
|
|
|
download_progress: Dict[str, Dict[str, Any]]
|
|
download_lock: Any
|
|
cancelled_downloads: set
|
|
aria2_lock: Any
|
|
aria2_transfers: Dict[str, Dict[str, Any]]
|
|
aria2_desired_states: Dict[str, Dict[str, Any]]
|
|
xet_transfers: Dict[str, Dict[str, Any]]
|
|
xet_transfers_lock: Any
|
|
set_download_progress_status: Callable[..., Any]
|
|
force_remove_aria2_transfer: Callable[..., Any]
|
|
thread_factory: Callable[..., Any]
|
|
log: Any
|
|
|
|
|
|
def get_progress(
|
|
download_id: str,
|
|
state: DownloadStateDependencies,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Get progress for a specific download."""
|
|
with state.download_lock:
|
|
return state.download_progress.get(download_id, {}).copy()
|
|
|
|
|
|
def get_all_progress(state: DownloadStateDependencies) -> Dict[str, Dict[str, Any]]:
|
|
"""Get progress for all downloads."""
|
|
with state.download_lock:
|
|
return {key: value.copy() for key, value in state.download_progress.items()}
|
|
|
|
|
|
def cancel_download(download_id: str, state: DownloadStateDependencies) -> bool:
|
|
"""Cancel a download in progress."""
|
|
with state.download_lock:
|
|
progress = state.download_progress.get(download_id)
|
|
if not progress or progress.get("status") in {"completed", "error", "cancelled"}:
|
|
return True
|
|
state.cancelled_downloads.add(download_id)
|
|
progress.update(status="cancelling", speed=0)
|
|
with state.aria2_lock:
|
|
state.aria2_desired_states.pop(download_id, None)
|
|
|
|
transfer = state.aria2_transfers.get(download_id)
|
|
if transfer and transfer.get("gid"):
|
|
state.thread_factory(
|
|
target=state.force_remove_aria2_transfer,
|
|
args=(download_id, transfer["gid"]),
|
|
daemon=True,
|
|
).start()
|
|
|
|
with state.xet_transfers_lock:
|
|
xet_transfer = dict(state.xet_transfers.get(download_id) or {})
|
|
xet_handle = xet_transfer.get("handle")
|
|
xet_cancel = getattr(xet_handle, "cancel", None)
|
|
if callable(xet_cancel):
|
|
try:
|
|
xet_cancel()
|
|
except Exception as exc:
|
|
state.log.warning(f"Could not cancel Hugging Face Xet transfer {download_id}: {exc}")
|
|
return True
|
|
|
|
|
|
def clear_completed_downloads(state: DownloadStateDependencies):
|
|
"""Clear completed/failed downloads from progress tracking."""
|
|
with state.download_lock:
|
|
to_remove = [
|
|
download_id
|
|
for download_id, info in state.download_progress.items()
|
|
if info.get("status") in ("completed", "error", "cancelled")
|
|
]
|
|
for download_id in to_remove:
|
|
del state.download_progress[download_id]
|
|
state.cancelled_downloads.discard(download_id)
|