From 280ca2e6deda5ac28d0eef65bcf1524192b4b5c7 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Tue, 14 Jul 2026 11:02:57 +0200 Subject: [PATCH] Improve HuggingFace Xet downloads with native progress snapshots, transfer metrics, cancellation, and finalizing UI state --- core/downloader.py | 279 +++++++++++++++--- web/resolver/actions/queue_methods.js | 16 +- .../actions/resolve_download_methods.js | 66 +++-- web/resolver/utils/render_format_methods.js | 44 ++- 4 files changed, 346 insertions(+), 59 deletions(-) diff --git a/core/downloader.py b/core/downloader.py index 5319592..07d2051 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -53,6 +53,8 @@ aria2_action_locks: Dict[str, threading.Lock] = {} aria2_desired_states: Dict[str, Dict[str, Any]] = {} aria2_idle_timer: Optional[threading.Timer] = None aria2_process_started_by_resolver = False +xet_transfers: Dict[str, Dict[str, Any]] = {} +xet_transfers_lock = threading.Lock() # Speed calculation settings SPEED_HISTORY_SIZE = 5 # Number of samples for smoothing @@ -1270,6 +1272,25 @@ def _delete_partial_download_files(dest_path: str) -> None: log.warning(f"Could not delete incomplete download file {path}: {exc}") +def _delete_xet_partial_file(partial_path: str, attempts: int = 5) -> bool: + """Delete a stopped Xet partial file, retrying while Windows releases it.""" + attempts = max(1, int(attempts or 1)) + last_error: Optional[Exception] = None + for attempt in range(attempts): + try: + if not os.path.exists(partial_path): + return True + os.remove(partial_path) + return True + except Exception as exc: + last_error = exc + if attempt + 1 < attempts: + time.sleep(0.25) + + log.warning(f"Could not delete incomplete Xet file {partial_path}: {last_error}") + return False + + class _HuggingFaceXetDownloadCancelled(Exception): """Raised by the Xet progress adapter when a download is cancelled.""" @@ -1281,7 +1302,65 @@ class _HuggingFaceXetProgressAdapter: self.download_id = download_id self.total_size = max(0, int(total_size or 0)) self.downloaded = 0 + self.transfer_downloaded = 0 + self.transfer_total_size = 0 self.last_update = start_time + self.speed = 0 + + def _publish( + self, + downloaded: int, + speed: int, + transfer_downloaded: int = 0, + transfer_total_size: int = 0, + ) -> None: + downloaded = max(0, int(downloaded or 0)) + if self.total_size > 0: + downloaded = min(downloaded, self.total_size) + self.downloaded = max(self.downloaded, downloaded) + progress = ( + min(100.0, round((self.downloaded / self.total_size) * 100, 1)) + if self.total_size > 0 + else 0 + ) + + self.transfer_downloaded = max( + self.transfer_downloaded, + max(0, int(transfer_downloaded or 0)), + ) + self.transfer_total_size = max( + self.transfer_total_size, + max(0, int(transfer_total_size or 0)), + ) + transfer_progress = ( + min( + 100.0, + round( + (self.transfer_downloaded / self.transfer_total_size) * 100, + 1, + ), + ) + if self.transfer_total_size > 0 + else 0 + ) + + self.speed = max(0, int(speed or 0)) + with download_lock: + state = download_progress.get(self.download_id) + if state is not None: + state.update( + { + "status": "downloading", + "progress": progress, + "downloaded": self.downloaded, + "total_size": self.total_size, + "speed": self.speed, + "transfer_downloaded": self.transfer_downloaded, + "transfer_total_size": self.transfer_total_size, + "transfer_progress": transfer_progress, + "download_backend": "huggingface_xet", + } + ) def update(self, byte_delta: Any) -> None: if self.download_id in cancelled_downloads: @@ -1292,30 +1371,157 @@ class _HuggingFaceXetProgressAdapter: except (TypeError, ValueError): delta = 0 - now = time.time() - elapsed = max(now - self.last_update, 0.001) - self.last_update = now - self.downloaded += delta - speed = int(delta / elapsed) if delta else 0 - progress = ( - min(100, int((self.downloaded / self.total_size) * 100)) - if self.total_size > 0 - else 0 + downloaded = self.downloaded + delta + # The legacy one-argument callback exposes logical byte increments only, + # not network-transfer speed. Reporting a derived rate here would confuse + # fast cache/file reconstruction with actual download throughput. + self._publish(downloaded, 0) + + def __call__(self, total_update: Any, item_updates: Any) -> None: + """Receive hf_xet's detailed 200 ms progress snapshots when available.""" + if self.download_id in cancelled_downloads: + raise _HuggingFaceXetDownloadCancelled("Download cancelled") + + downloaded = int(getattr(total_update, "total_bytes_completed", 0) or 0) + if downloaded <= 0: + items = item_updates.values() if isinstance(item_updates, dict) else item_updates + for item in items or []: + downloaded = max( + downloaded, + int(getattr(item, "bytes_completed", 0) or 0), + ) + + transfer_downloaded = int( + getattr(total_update, "total_transfer_bytes_completed", 0) or 0 + ) + transfer_total_size = int( + getattr(total_update, "total_transfer_bytes", 0) or 0 + ) + transfer_rate = getattr( + total_update, + "total_transfer_bytes_completion_rate", + None, + ) + try: + speed = max(0, int(float(transfer_rate or 0))) + except (TypeError, ValueError): + speed = 0 + self.last_update = time.time() + self._publish( + downloaded, + speed, + transfer_downloaded, + transfer_total_size, ) - with download_lock: - state = download_progress.get(self.download_id) - if state is not None: - state.update( - { - "status": "downloading", - "progress": progress, - "downloaded": self.downloaded, - "total_size": self.total_size, - "speed": speed, - "download_backend": "huggingface_xet", - } + +def _run_huggingface_xet_transfer( + incomplete_path: Path, + xet_file_data: Any, + request_headers: Dict[str, str], + expected_size: int, + filename: str, + progress_adapter: _HuggingFaceXetProgressAdapter, +) -> None: + """Use detailed native Xet progress when supported, with legacy fallback.""" + import hf_xet + from huggingface_hub.file_download import xet_get + from huggingface_hub.utils import refresh_xet_connection_info + + supports_session_progress = all( + hasattr(hf_xet, name) + for name in ("XetFileInfo", "XetSession") + ) + if supports_session_progress: + try: + from huggingface_hub.utils._xet import ( + get_xet_session, + xet_headers_without_auth, + ) + except ImportError: + supports_session_progress = False + + if supports_session_progress: + session = get_xet_session() + xet_headers = xet_headers_without_auth(request_headers) + group = session.new_file_download_group( + token_refresh_url=xet_file_data.refresh_route, + token_refresh_headers=request_headers, + custom_headers=xet_headers, + progress_callback=progress_adapter, + progress_interval_ms=200, + ) + try: + with group: + handle = group.start_download_file( + hf_xet.XetFileInfo(xet_file_data.file_hash, expected_size or None), + str(incomplete_path.absolute()), ) + with xet_transfers_lock: + xet_transfers[progress_adapter.download_id] = { + "handle": handle, + "partial_path": str(incomplete_path), + } + if progress_adapter.download_id in cancelled_downloads: + cancel = getattr(handle, "cancel", None) + if callable(cancel): + cancel() + finally: + with xet_transfers_lock: + xet_transfers.pop(progress_adapter.download_id, None) + return + + supports_detailed_progress = all( + hasattr(hf_xet, name) + for name in ( + "PyItemProgressUpdate", + "PyTotalProgressUpdate", + "PyXetDownloadInfo", + "download_files", + ) + ) + if not supports_detailed_progress: + xet_get( + incomplete_path=incomplete_path, + xet_file_data=xet_file_data, + headers=request_headers, + expected_size=expected_size or None, + displayed_filename=filename, + _tqdm_bar=progress_adapter, + ) + return + + connection_info = refresh_xet_connection_info( + file_data=xet_file_data, + headers=request_headers, + ) + if connection_info is None: + raise ValueError("Failed to refresh Hugging Face Xet connection info") + + def token_refresher() -> tuple[str, int]: + refreshed = refresh_xet_connection_info( + file_data=xet_file_data, + headers=request_headers, + ) + if refreshed is None: + raise ValueError("Failed to refresh Hugging Face Xet access token") + return refreshed.access_token, refreshed.expiration_unix_epoch + + download_info = hf_xet.PyXetDownloadInfo( + destination_path=str(incomplete_path.absolute()), + hash=xet_file_data.file_hash, + file_size=expected_size or None, + ) + hf_xet.download_files( + [download_info], + endpoint=connection_info.endpoint, + token_info=( + connection_info.access_token, + connection_info.expiration_unix_epoch, + ), + token_refresher=token_refresher, + progress_updater=[progress_adapter], + ) def _download_huggingface_xet( @@ -1337,7 +1543,7 @@ def _download_huggingface_xet( try: __import__("hf_xet") - from huggingface_hub.file_download import get_hf_file_metadata, xet_get + from huggingface_hub.file_download import get_hf_file_metadata except ImportError: return None @@ -1403,13 +1609,13 @@ def _download_huggingface_xet( os.remove(partial_path) log.info(f"Starting Hugging Face Xet download: {filename}") - xet_get( - incomplete_path=Path(partial_path), - xet_file_data=xet_file_data, - headers=request_headers, - expected_size=expected_size or None, - displayed_filename=filename, - _tqdm_bar=progress_adapter, + _run_huggingface_xet_transfer( + Path(partial_path), + xet_file_data, + request_headers, + expected_size, + filename, + progress_adapter, ) if download_id in cancelled_downloads: @@ -1465,11 +1671,7 @@ def _download_huggingface_xet( isinstance(exc, _HuggingFaceXetDownloadCancelled) or download_id in cancelled_downloads ) - try: - if os.path.exists(partial_path): - os.remove(partial_path) - except Exception as cleanup_exc: - log.warning(f"Could not delete incomplete Xet file: {cleanup_exc}") + _delete_xet_partial_file(partial_path) error_msg = ( "Download cancelled" if was_cancelled else _sanitize_download_error(exc) @@ -2354,6 +2556,15 @@ def cancel_download(download_id: str) -> bool: args=(download_id, transfer["gid"]), daemon=True, ).start() + with xet_transfers_lock: + xet_transfer = dict(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: + log.warning(f"Could not cancel Hugging Face Xet transfer {download_id}: {exc}") return True diff --git a/web/resolver/actions/queue_methods.js b/web/resolver/actions/queue_methods.js index 97a56bc..70d5410 100644 --- a/web/resolver/actions/queue_methods.js +++ b/web/resolver/actions/queue_methods.js @@ -1253,7 +1253,8 @@ export const queueMethods = { const progress = typeof this.applyPendingDownloadStatus === 'function' ? this.applyPendingDownloadStatus(info, info.lastProgress || {}) : (info.lastProgress || {}); - const percent = Math.max(0, Math.min(100, Number(progress.progress) || 0)); + const displayProgress = this.getDownloadDisplayProgress(progress); + const percent = displayProgress.percent; const filename = progress.filename || info.filename || info.missing?.download_source?.filename @@ -1261,17 +1262,22 @@ export const queueMethods = { || 'model'; const category = this.getCategoryDisplayName?.(info.category || info.missing?.category || '') || info.category || ''; const nodeLabel = info.missing?.subgraph_name || info.missing?.node_type || (info.missing?.subgraph_id ? 'Subgraph' : 'Node'); - const downloaded = this.formatBytes(progress.downloaded || 0); - const total = progress.total_size ? this.formatBytes(progress.total_size) : ''; + const downloaded = this.formatBytes(displayProgress.downloaded); + const total = displayProgress.totalSize ? this.formatBytes(displayProgress.totalSize) : ''; const progressMeta = this.formatDownloadProgressMeta?.(progress) || ''; const status = progress.status || info.lastStatus || 'starting'; const backend = String(progress.download_backend || info.downloadBackend || '').toLowerCase(); const isAria2 = backend === 'aria2'; const isPaused = status === 'paused'; + const percentLabel = this.formatDownloadPercent?.(percent) ?? String(Math.round(percent)); const statusLabel = status === 'downloading' - ? `${Math.round(percent)}%` + ? (displayProgress.isFinalizing ? 'Finalizing' : `${percentLabel}%`) : (status === 'starting' ? 'Starting' : (isPaused ? 'Paused' : status.replace(/_/g, ' '))); - const sizeText = total ? `${downloaded} / ${total}` : downloaded; + const logicalDownloaded = this.formatBytes(progress.downloaded || 0); + const logicalTotal = progress.total_size ? this.formatBytes(progress.total_size) : ''; + const sizeText = displayProgress.isFinalizing && logicalTotal + ? `Finalizing file: ${logicalDownloaded} / ${logicalTotal}` + : (total ? `${downloaded} / ${total}` : downloaded); const targetPath = progress.directory || info.downloadDirectory || info.downloadPath || ''; const targetLabel = targetPath ? targetPath.split(/[\/\\]/).filter(Boolean).pop() || targetPath : ''; const workflowLabel = this.getDownloadWorkflowLabel?.(info) || 'Unknown workflow'; diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index 4c3643f..373f310 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/web/resolver/actions/resolve_download_methods.js @@ -463,7 +463,7 @@ export const resolveDownloadMethods = { name: progress?.filename || info?.filename || 'Download', path: targetPath, resolved_path: targetPath, - open_path: filePath || directory, + open_path: directory || filePath, folder_path: directory || targetPath, download_directory: directory || '', download_path: filePath || '', @@ -1282,8 +1282,12 @@ export const resolveDownloadMethods = { }); }, - getDownloadProgressStatusLabel(status = '', percent = 0) { - if (status === 'downloading') return `${Math.round(percent)}%`; + getDownloadProgressStatusLabel(status = '', percent = 0, isFinalizing = false) { + if (status === 'downloading') { + if (isFinalizing) return 'Finalizing'; + const percentLabel = this.formatDownloadPercent?.(percent) ?? String(Math.round(percent)); + return `${percentLabel}%`; + } if (status === 'starting') return 'Starting'; if (status === 'paused') return 'Paused'; if (status === 'completed_checking') return 'Checking'; @@ -1309,13 +1313,19 @@ export const resolveDownloadMethods = { const backend = String(progress.download_backend || snapshot.downloadBackend || snapshot.download_backend || '').toLowerCase(); const isAria2 = backend === 'aria2'; const isPaused = status === 'paused'; - const percent = Math.max(0, Math.min(100, Number(progress.progress) || 0)); - const downloaded = this.formatBytes(progress.downloaded || 0); - const total = this.formatBytes(progress.total_size || 0); + const displayProgress = this.getDownloadDisplayProgress(progress); + const percent = displayProgress.percent; + const downloaded = this.formatBytes(displayProgress.downloaded); + const total = this.formatBytes(displayProgress.totalSize); const progressMeta = this.formatDownloadProgressMeta(progress); - const leftText = progress.total_size - ? `${downloaded} / ${total} (${percent}%)${isPaused ? ' - Paused' : ''}` - : 'Connecting...'; + const percentLabel = this.formatDownloadPercent?.(percent) ?? String(Math.round(percent)); + const logicalDownloaded = this.formatBytes(progress.downloaded || 0); + const logicalTotal = progress.total_size ? this.formatBytes(progress.total_size) : 'Unknown'; + const leftText = displayProgress.isFinalizing + ? `Finalizing file: ${logicalDownloaded} / ${logicalTotal}` + : (displayProgress.totalSize + ? `${downloaded} / ${total} (${percentLabel}%)${isPaused ? ' - Paused' : ''}` + : 'Connecting...'); let actionClass = canCancel ? 'cancel-download-btn mr-btn mr-btn-danger mr-btn-sm' : ''; let actionText = canCancel ? 'Cancel' : ''; let actionsHtml = ''; @@ -1390,9 +1400,10 @@ export const resolveDownloadMethods = { renderDownloadProgressGroupItem(downloadId, snapshot) { const progress = snapshot.progress || {}; const status = snapshot.status || progress.status || ''; - const percent = Math.max(0, Math.min(100, Number(progress.progress) || 0)); + const displayProgress = this.getDownloadDisplayProgress(progress); + const percent = displayProgress.percent; const filename = progress.filename || snapshot.filename || 'Download'; - const statusLabel = this.getDownloadProgressStatusLabel(status, percent); + const statusLabel = this.getDownloadProgressStatusLabel(status, percent, displayProgress.isFinalizing); const itemIdAttr = downloadId ? ` data-download-id="${this.escapeHtml(String(downloadId))}"` : ''; return `
@@ -1432,20 +1443,28 @@ export const resolveDownloadMethods = { if (!downloadBtn) return; if (shouldRenderProgress) { - const percent = Number(progress.progress); + const displayProgress = this.getDownloadDisplayProgress(progress); + const percent = displayProgress.percent; downloadBtn.disabled = true; downloadBtn.classList.remove('mr-is-success-action', 'mr-btn-primary'); - const label = status === 'paused' - ? `Paused ${Number.isFinite(percent) && percent > 0 ? `${Math.round(percent)}%` : ''}`.trim() - : (Number.isFinite(percent) && percent > 0 ? `Downloading ${Math.round(percent)}%` : 'Starting download...'); + const percentLabel = Number.isFinite(percent) && percent > 0 + ? `${this.formatDownloadPercent?.(percent) ?? Math.round(percent)}%` + : ''; + const label = displayProgress.isFinalizing + ? 'Finalizing download...' + : (status === 'paused' + ? `Paused ${percentLabel}`.trim() + : (percentLabel ? `Downloading ${percentLabel}` : 'Starting download...')); if (downloadBtn.classList.contains('search-download-btn')) { downloadBtn.innerHTML = getSvgIcon('download'); downloadBtn.setAttribute('data-tooltip', label); downloadBtn.setAttribute('aria-label', label); } else { - downloadBtn.textContent = status === 'paused' - ? 'Paused' - : (Number.isFinite(percent) && percent > 0 ? `${Math.round(percent)}%` : 'Starting...'); + downloadBtn.textContent = displayProgress.isFinalizing + ? 'Finalizing...' + : (status === 'paused' + ? 'Paused' + : (Number.isFinite(percent) && percent > 0 ? `${Math.round(percent)}%` : 'Starting...')); } } else if (status === 'cancelling') { downloadBtn.disabled = true; @@ -1992,6 +2011,14 @@ export const resolveDownloadMethods = { /** * Poll download progress */ + getDownloadProgressPollDelay(progress = {}, info = {}) { + if (progress.status === 'paused') return 1500; + const backend = String( + progress.download_backend || info.downloadBackend || info.download_backend || '' + ).toLowerCase(); + return backend === 'huggingface_xet' ? 200 : 1000; + }, + async pollDownloadProgress(downloadId) { const info = this.activeDownloads[downloadId]; if (!info) return; @@ -2017,7 +2044,8 @@ export const resolveDownloadMethods = { } // Continue polling - setTimeout(() => this.pollDownloadProgress(downloadId), progress.status === 'paused' ? 1500 : 1000); + const pollDelay = this.getDownloadProgressPollDelay(progress, info); + setTimeout(() => this.pollDownloadProgress(downloadId), pollDelay); } else if (progress.status === 'cancelling') { this.renderDownloadSnapshot(downloadId, snapshot, { progressDiv, downloadBtn }); diff --git a/web/resolver/utils/render_format_methods.js b/web/resolver/utils/render_format_methods.js index 979feae..84af6ab 100644 --- a/web/resolver/utils/render_format_methods.js +++ b/web/resolver/utils/render_format_methods.js @@ -350,9 +350,51 @@ export const renderFormatMethods = { return remainingSeconds > 0 ? `${minutes}m ${remainingSeconds}s` : `${minutes}m`; }, + formatDownloadPercent(value) { + const percent = Math.max(0, Math.min(100, Number(value) || 0)); + if (percent > 0 && percent < 0.1) return '<0.1'; + if (percent > 0 && percent < 10) { + return percent.toFixed(1).replace(/\.0$/, ''); + } + return String(Math.round(percent)); + }, + + getDownloadDisplayProgress(progress = {}) { + const logicalDownloaded = Math.max(0, Number(progress.downloaded) || 0); + const logicalTotalSize = Math.max(0, Number(progress.total_size) || 0); + const logicalPercent = Math.max(0, Math.min(100, Number(progress.progress) || 0)); + const backend = String(progress.download_backend || '').toLowerCase(); + const transferDownloaded = Math.max(0, Number(progress.transfer_downloaded) || 0); + const transferTotalSize = Math.max(0, Number(progress.transfer_total_size) || 0); + + if (backend === 'huggingface_xet' && logicalTotalSize > 0) { + return { + percent: Math.max(0, Math.min(100, (transferDownloaded / logicalTotalSize) * 100)), + downloaded: transferDownloaded, + totalSize: logicalTotalSize, + isTransfer: true, + isFinalizing: transferTotalSize > 0 + && transferDownloaded >= transferTotalSize + && logicalTotalSize > 0 + && logicalDownloaded < logicalTotalSize + }; + } + + return { + percent: logicalPercent, + downloaded: logicalDownloaded, + totalSize: logicalTotalSize, + isTransfer: false, + isFinalizing: false + }; + }, + getDownloadEtaText(progress = {}) { + const backend = String(progress.download_backend || '').toLowerCase(); const totalSize = Number(progress.total_size) || 0; - const downloaded = Number(progress.downloaded) || 0; + const downloaded = backend === 'huggingface_xet' + ? (Number(progress.transfer_downloaded) || 0) + : (Number(progress.downloaded) || 0); const speed = Number(progress.speed) || 0; if (totalSize <= 0 || downloaded <= 0 || speed <= 0 || downloaded >= totalSize) return '';