From 0e8acf5d59fcb249931b8138cde91a97c947bbdd Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Mon, 6 Jul 2026 23:57:59 +0300 Subject: [PATCH] DL tracker fixed!!!! now make it nice! --- packages/node/meshnet_node/downloader.py | 128 +++++++++++++++++------ 1 file changed, 94 insertions(+), 34 deletions(-) diff --git a/packages/node/meshnet_node/downloader.py b/packages/node/meshnet_node/downloader.py index 2577e33..6c14d98 100644 --- a/packages/node/meshnet_node/downloader.py +++ b/packages/node/meshnet_node/downloader.py @@ -122,6 +122,7 @@ def _fetch_source_file( file_url: str, dest: Path, timeout: float, + on_chunk=None, ) -> int: """Download one file; skip if *dest* already matches the remote size. @@ -135,6 +136,8 @@ def _fetch_source_file( length = resp.getheader("Content-Length") if hasattr(resp, "getheader") else None expected = int(length) if isinstance(length, str) and length.isdigit() else None if expected is not None and dest.exists() and dest.stat().st_size == expected: + if on_chunk is not None: + on_chunk(expected) return expected # complete from an earlier attempt/run received = 0 with dest.open("wb") as out: @@ -144,11 +147,63 @@ def _fetch_source_file( break out.write(chunk) received += len(chunk) + if on_chunk is not None: + on_chunk(len(chunk)) if expected is not None and received != expected: raise OSError(f"short read: got {received} of {expected} bytes") return received +class _SourceProgress: + """tqdm bar over the whole per-file download (total / speed / ETA), + falling back to plain per-file prints when tqdm is unavailable.""" + + def __init__(self, enabled: bool, label: str, total_bytes: int | None): + self._label = label + self._bar = None + self._enabled = enabled + if enabled and total_bytes: + try: + from tqdm import tqdm # type: ignore[import] + + self._bar = tqdm( + total=total_bytes, + unit="B", + unit_scale=True, + unit_divisor=1024, + desc=f"Downloading ({label})", + ) + except Exception: + self._bar = None + + def add_bytes(self, n: int) -> None: + if self._bar is not None: + self._bar.update(n) + + def rewind(self, n: int) -> None: + """Roll the bar back after a failed attempt so retries don't double-count.""" + if self._bar is not None and n: + self._bar.update(-n) + + def file_done(self, index: int, count: int, rel: str, size: int, reused: bool) -> None: + note = "already complete, " if reused else "" + line = f" {self._label}: [{index}/{count}] {rel} ({note}{size / 1e9:.2f} GB)" + if self._bar is not None: + self._bar.set_postfix_str(f"{index}/{count} files", refresh=False) + elif self._enabled: + print(line, flush=True) + + def message(self, text: str) -> None: + if self._bar is not None: + self._bar.write(text) + elif self._enabled: + print(text, flush=True) + + def close(self) -> None: + if self._bar is not None: + self._bar.close() + + def _reuse_local_file(expected: int, dest: Path, final: Path) -> bool: """Reuse an already-complete copy of a file instead of re-downloading.""" if dest.exists() and dest.stat().st_size == expected: @@ -188,41 +243,46 @@ def _download_source_files( partial_dir = shard_dir.parent / f"{shard_dir.name}.partial" partial_dir.mkdir(parents=True, exist_ok=True) sep = "&" if "?" in url else "?" - for index, rel in enumerate(rel_files, start=1): - dest = partial_dir / rel - dest.parent.mkdir(parents=True, exist_ok=True) - known_size = sizes.get(rel) - if isinstance(known_size, int) and _reuse_local_file(known_size, dest, shard_dir / rel): - if progress: - print( - f" {label}: [{index}/{len(rel_files)}] {rel} " - f"(already complete, {known_size / 1e9:.2f} GB)", - flush=True, - ) - continue - file_url = f"{url}{sep}{urllib.parse.urlencode({'file': rel})}" - for attempt in range(1, _FILE_RETRY_ATTEMPTS + 1): - try: - size = _fetch_source_file(file_url, dest, timeout) - if progress: - print( - f" {label}: [{index}/{len(rel_files)}] {rel} ({size / 1e9:.2f} GB)", - flush=True, - ) - break - except _TarOnlySource: - if progress: - print(f" {label}: no single-file support — using tar stream", flush=True) - return None - except Exception as exc: - if progress: - print( - f" {label}: {rel} attempt {attempt}/{_FILE_RETRY_ATTEMPTS} failed: {exc!r}", - flush=True, - ) - if attempt == _FILE_RETRY_ATTEMPTS: + known_sizes = [sizes.get(rel) for rel in rel_files] + total_bytes = sum(s for s in known_sizes if isinstance(s, int)) if all( + isinstance(s, int) for s in known_sizes + ) else None + tracker_bar = _SourceProgress(progress, label, total_bytes) + try: + for index, rel in enumerate(rel_files, start=1): + dest = partial_dir / rel + dest.parent.mkdir(parents=True, exist_ok=True) + known_size = sizes.get(rel) + if isinstance(known_size, int) and _reuse_local_file(known_size, dest, shard_dir / rel): + tracker_bar.add_bytes(known_size) + tracker_bar.file_done(index, len(rel_files), rel, known_size, reused=True) + continue + file_url = f"{url}{sep}{urllib.parse.urlencode({'file': rel})}" + for attempt in range(1, _FILE_RETRY_ATTEMPTS + 1): + counted = 0 + + def _on_chunk(n: int) -> None: + nonlocal counted + counted += n + tracker_bar.add_bytes(n) + + try: + size = _fetch_source_file(file_url, dest, timeout, on_chunk=_on_chunk) + tracker_bar.file_done(index, len(rel_files), rel, size, reused=False) + break + except _TarOnlySource: + tracker_bar.message(f" {label}: no single-file support — using tar stream") return None - time.sleep(1.0 * attempt) + except Exception as exc: + tracker_bar.rewind(counted) + tracker_bar.message( + f" {label}: {rel} attempt {attempt}/{_FILE_RETRY_ATTEMPTS} failed: {exc!r}" + ) + if attempt == _FILE_RETRY_ATTEMPTS: + return None + time.sleep(1.0 * attempt) + finally: + tracker_bar.close() if shard_dir.exists(): shutil.rmtree(shard_dir) shutil.move(str(partial_dir), str(shard_dir))