From 4186d0bcf1afd15f86bfba1c9c893649dd6885df Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 07:16:04 +0300 Subject: [PATCH] Add download progress tracking and byte formatting utilities - Introduced `fmt_bytes` function for human-readable byte size formatting. - Added `DownloadProgress` class to track and display download progress with speed and completion percentage. - Updated `download` functions in `civitai_fetch.py` and `llamacpp_fetch.py` to utilize progress tracking. - Enhanced `ollama_pull.py` to support streaming progress updates during model pulls. - Updated tests to validate new formatting and progress tracking functionalities, ensuring accurate output and user feedback. --- src/gpu_rent/remote/civitai_fetch.py | 81 ++++++++++++++++++++-- src/gpu_rent/remote/llamacpp_fetch.py | 99 +++++++++++++++++++++++---- src/gpu_rent/remote/ollama_pull.py | 81 ++++++++++++++++++++-- tests/test_civitai_fetch.py | 23 ++++++- 4 files changed, 259 insertions(+), 25 deletions(-) diff --git a/src/gpu_rent/remote/civitai_fetch.py b/src/gpu_rent/remote/civitai_fetch.py index 4907bfe..e221243 100644 --- a/src/gpu_rent/remote/civitai_fetch.py +++ b/src/gpu_rent/remote/civitai_fetch.py @@ -5,6 +5,7 @@ from __future__ import annotations import hashlib import json import sys +import time import urllib.parse import urllib.request from pathlib import Path @@ -14,6 +15,68 @@ JOBS_PATH = Path("/tmp/gpu-rent-civitai-jobs.json") MARKER = Path("/mnt/swarm_data/.gpu-rent-models-seeded") +def fmt_bytes(n: float) -> str: + n = float(n) + for unit, div in (("GB", 1024**3), ("MB", 1024**2), ("KB", 1024), ("B", 1)): + if n >= div or unit == "B": + if unit == "B": + return f"{int(n)}B" + return f"{n / div:.1f}{unit}" + return f"{n:.0f}B" + + +def progress_line( + label: str, + done: int, + total: int | None, + speed: float, + *, + width: int = 22, +) -> str: + """One SSH-safe progress line (newline, not \\r).""" + if total and total > 0: + pct = min(100.0, 100.0 * done / total) + filled = int(width * done / total) + filled = min(width, max(0, filled)) + bar = "#" * filled + "-" * (width - filled) + return ( + f"{label} [{bar}] {pct:5.1f}% " + f"{fmt_bytes(done)}/{fmt_bytes(total)} {fmt_bytes(speed)}/s" + ) + return f"{label} {fmt_bytes(done)} {fmt_bytes(speed)}/s" + + +class DownloadProgress: + """Print size + speed about once per second (SSH readline-friendly).""" + + def __init__(self, label: str, total: int | None) -> None: + self.label = label + self.total = total if total and total > 0 else None + self.done = 0 + self.t0 = time.monotonic() + self.last_print = 0.0 + + def add(self, n: int) -> None: + self.done += n + now = time.monotonic() + if now - self.last_print < 1.0 and not ( + self.total is not None and self.done >= self.total + ): + return + self.last_print = now + self._emit() + + def finish(self) -> None: + self._emit() + + def _emit(self) -> None: + elapsed = max(time.monotonic() - self.t0, 0.001) + print( + progress_line(self.label, self.done, self.total, self.done / elapsed), + flush=True, + ) + + def sha256_path(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as fh: @@ -38,7 +101,7 @@ def should_skip(dest: Path, expect_sha: str) -> tuple[bool, str]: return False, "sha не совпал — перекачиваю" -def download(url: str, dest: Path, token: str) -> None: +def download(url: str, dest: Path, token: str, *, label: str) -> None: dest.parent.mkdir(parents=True, exist_ok=True) partial = dest.with_suffix(dest.suffix + ".partial") @@ -60,11 +123,19 @@ def download(url: str, dest: Path, token: str) -> None: url, headers={"Authorization": f"Bearer {token}", "User-Agent": "gpu-rent/0.1"} ) with opener.open(req, timeout=600) as response, partial.open("wb") as out: + total = response.headers.get("Content-Length") + try: + total_n = int(total) if total else None + except ValueError: + total_n = None + prog = DownloadProgress(label, total_n) while True: chunk = response.read(1024 * 1024) if not chunk: break out.write(chunk) + prog.add(len(chunk)) + prog.finish() partial.replace(dest) @@ -85,7 +156,8 @@ def main() -> int: skip, reason = should_skip(dest, expect) if skip: skipped += 1 - print(f"{prefix} {reason}: {dest.name}") + size = dest.stat().st_size if dest.is_file() else 0 + print(f"{prefix} {reason}: {dest.name} ({fmt_bytes(size)})") # refresh sidecars even on skip for extra_name, extra_text in (job.get("sidecars") or {}).items(): extra = dest.parent / extra_name @@ -95,8 +167,8 @@ def main() -> int: if reason: print(f"{prefix} {reason}: {dest.name}") try: - print(f"{prefix} качаю: {dest.name}") - download(job["url"], dest, token) + print(f"{prefix} качаю: {dest.name}", flush=True) + download(job["url"], dest, token, label=f"{prefix} {dest.name}") downloaded += 1 if expect: got = sha256_path(dest).lower() @@ -106,6 +178,7 @@ def main() -> int: for extra_name, extra_text in (job.get("sidecars") or {}).items(): extra = dest.parent / extra_name extra.write_text(extra_text, encoding="utf-8") + print(f"{prefix} ok {dest.name} ({fmt_bytes(dest.stat().st_size)})") except Exception as exc: failed += 1 print(f"{prefix} FAIL {dest.name}: {exc}", file=sys.stderr) diff --git a/src/gpu_rent/remote/llamacpp_fetch.py b/src/gpu_rent/remote/llamacpp_fetch.py index 0e7ebca..d2dc4ac 100644 --- a/src/gpu_rent/remote/llamacpp_fetch.py +++ b/src/gpu_rent/remote/llamacpp_fetch.py @@ -5,6 +5,7 @@ from __future__ import annotations import json import os import sys +import time import urllib.error import urllib.request from pathlib import Path @@ -14,6 +15,85 @@ MODELS_DIR = Path("/mnt/swarm_data/llamacpp/models") TOKEN_FILE = Path("/tmp/gpu-rent-hf.token") +def fmt_bytes(n: float) -> str: + n = float(n) + for unit, div in (("GB", 1024**3), ("MB", 1024**2), ("KB", 1024), ("B", 1)): + if n >= div or unit == "B": + if unit == "B": + return f"{int(n)}B" + return f"{n / div:.1f}{unit}" + return f"{n:.0f}B" + + +def progress_line( + label: str, + done: int, + total: int | None, + speed: float, + *, + width: int = 22, +) -> str: + if total and total > 0: + pct = min(100.0, 100.0 * done / total) + filled = int(width * done / total) + filled = min(width, max(0, filled)) + bar = "#" * filled + "-" * (width - filled) + return ( + f"{label} [{bar}] {pct:5.1f}% " + f"{fmt_bytes(done)}/{fmt_bytes(total)} {fmt_bytes(speed)}/s" + ) + return f"{label} {fmt_bytes(done)} {fmt_bytes(speed)}/s" + + +class DownloadProgress: + def __init__(self, label: str, total: int | None) -> None: + self.label = label + self.total = total if total and total > 0 else None + self.done = 0 + self.t0 = time.monotonic() + self.last_print = 0.0 + + def add(self, n: int) -> None: + self.done += n + now = time.monotonic() + if now - self.last_print < 1.0 and not ( + self.total is not None and self.done >= self.total + ): + return + self.last_print = now + self._emit() + + def finish(self) -> None: + self._emit() + + def _emit(self) -> None: + elapsed = max(time.monotonic() - self.t0, 0.001) + print( + progress_line(self.label, self.done, self.total, self.done / elapsed), + flush=True, + ) + + +def download(url: str, dest: Path, headers: dict[str, str], *, label: str) -> None: + partial = dest.with_suffix(dest.suffix + ".partial") + req = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(req, timeout=600) as resp, partial.open("wb") as out: + cl = resp.headers.get("Content-Length") + try: + total_n = int(cl) if cl else None + except ValueError: + total_n = None + prog = DownloadProgress(label, total_n) + while True: + chunk = resp.read(1024 * 1024) + if not chunk: + break + out.write(chunk) + prog.add(len(chunk)) + prog.finish() + partial.replace(dest) + + def main() -> int: if TOKEN_FILE.is_file(): try: @@ -43,29 +123,22 @@ def main() -> int: if not name: name = url.rstrip("/").rsplit("/", 1)[-1] or "model.gguf" dest = MODELS_DIR / name + prefix = f"[{i}/{len(jobs)}]" if dest.is_file() and dest.stat().st_size > 1_000_000: - print(f"[{i}/{len(jobs)}] уже есть {name} ({dest.stat().st_size} bytes)") + print(f"{prefix} уже есть {name} ({fmt_bytes(dest.stat().st_size)})") continue - print(f"[{i}/{len(jobs)}] download {name}") - partial = dest.with_suffix(dest.suffix + ".partial") + print(f"{prefix} качаю {name}", flush=True) headers = {"User-Agent": "gpu-rent/1"} if token: headers["Authorization"] = f"Bearer {token}" - req = urllib.request.Request(url, headers=headers) try: - with urllib.request.urlopen(req, timeout=600) as resp, partial.open("wb") as out: - while True: - chunk = resp.read(1024 * 1024) - if not chunk: - break - out.write(chunk) - partial.replace(dest) - print(f"ok {name} ({dest.stat().st_size} bytes)") + download(url, dest, headers, label=f"{prefix} {name}") + print(f"{prefix} ok {name} ({fmt_bytes(dest.stat().st_size)})") except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError) as exc: failed += 1 print(f"FAIL {name}: {exc}", file=sys.stderr) try: - partial.unlink(missing_ok=True) + dest.with_suffix(dest.suffix + ".partial").unlink(missing_ok=True) except OSError: pass if failed: diff --git a/src/gpu_rent/remote/ollama_pull.py b/src/gpu_rent/remote/ollama_pull.py index 79e0cd5..7624490 100644 --- a/src/gpu_rent/remote/ollama_pull.py +++ b/src/gpu_rent/remote/ollama_pull.py @@ -1,15 +1,43 @@ #!/usr/bin/env python3 -"""Pull Ollama models from a JSON list. Stdlib only. Runs on the VM.""" +"""Pull Ollama models from a JSON list. Stdlib only. Runs on the VM. + +Uses POST /api/pull with stream JSON for completed/total + speed lines. +""" from __future__ import annotations import json import subprocess import sys import time +import urllib.error +import urllib.request from pathlib import Path JOBS = Path("/tmp/gpu-rent-ollama-models.json") MARKER = Path("/mnt/swarm_data/.gpu-rent-ollama-pulling") +OLLAMA = "http://127.0.0.1:11434" + + +def fmt_bytes(n: float) -> str: + n = float(n) + for unit, div in (("GB", 1024**3), ("MB", 1024**2), ("KB", 1024), ("B", 1)): + if n >= div or unit == "B": + if unit == "B": + return f"{int(n)}B" + return f"{n / div:.1f}{unit}" + return f"{n:.0f}B" + + +def progress_line(label: str, done: int, total: int | None, speed: float, *, width: int = 22) -> str: + if total and total > 0: + pct = min(100.0, 100.0 * done / total) + filled = min(width, max(0, int(width * done / total))) + bar = "#" * filled + "-" * (width - filled) + return ( + f"{label} [{bar}] {pct:5.1f}% " + f"{fmt_bytes(done)}/{fmt_bytes(total)} {fmt_bytes(speed)}/s" + ) + return f"{label} {fmt_bytes(done)} {fmt_bytes(speed)}/s" def listed() -> set[str]: @@ -32,7 +60,6 @@ def already_have(have: set[str], wanted: str) -> bool: """Exact tag match only — qwen2.5:3b must not satisfy qwen2.5:7b.""" if wanted in have: return True - # ollama list sometimes omits :latest if ":" not in wanted and f"{wanted}:latest" in have: return True if wanted.endswith(":latest") and wanted.rsplit(":", 1)[0] in have: @@ -40,6 +67,45 @@ def already_have(have: set[str], wanted: str) -> bool: return False +def pull_stream(name: str, label: str) -> None: + body = json.dumps({"name": name, "stream": True}).encode("utf-8") + req = urllib.request.Request( + f"{OLLAMA}/api/pull", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + t0 = time.monotonic() + last_print = 0.0 + with urllib.request.urlopen(req, timeout=7200) as resp: + while True: + raw = resp.readline() + if not raw: + break + try: + ev = json.loads(raw.decode("utf-8")) + except json.JSONDecodeError: + continue + if ev.get("error"): + raise RuntimeError(str(ev["error"])) + status = str(ev.get("status") or "") + total = int(ev.get("total") or 0) or None + done = int(ev.get("completed") or 0) + now = time.monotonic() + if total and done: + if now - last_print >= 1.0 or done >= total: + elapsed = max(now - t0, 0.001) + speed = done / elapsed + print(progress_line(label, done, total, speed), flush=True) + last_print = now + elif status and status not in {"success"} and now - last_print >= 2.0: + print(f"{label} {status}", flush=True) + last_print = now + if status == "success": + break + print(f"{label} ok", flush=True) + + def main() -> int: if not JOBS.is_file(): print("no jobs file", file=sys.stderr) @@ -57,14 +123,15 @@ def main() -> int: name = str(name).strip() if not name: continue + prefix = f"[{i}/{len(models)}]" if already_have(have, name): - print(f"[{i}/{len(models)}] уже есть {name}") + print(f"{prefix} уже есть {name}") continue - print(f"[{i}/{len(models)}] ollama pull {name}") + print(f"{prefix} ollama pull {name}", flush=True) try: - subprocess.check_call(["ollama", "pull", name]) + pull_stream(name, label=f"{prefix} {name}") have.add(name) - except subprocess.CalledProcessError as exc: + except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError, RuntimeError) as exc: failed += 1 print(f"FAIL pull {name}: {exc}", file=sys.stderr) finally: @@ -76,4 +143,4 @@ def main() -> int: if __name__ == "__main__": - sys.exit(main()) + raise SystemExit(main()) diff --git a/tests/test_civitai_fetch.py b/tests/test_civitai_fetch.py index 7ef7a78..7e91439 100644 --- a/tests/test_civitai_fetch.py +++ b/tests/test_civitai_fetch.py @@ -1,6 +1,6 @@ from pathlib import Path -from gpu_rent.remote.civitai_fetch import should_skip +from gpu_rent.remote.civitai_fetch import fmt_bytes, progress_line, should_skip def test_should_skip_missing(tmp_path: Path): @@ -33,3 +33,24 @@ def test_should_not_skip_bad_sha(tmp_path: Path): skip, reason = should_skip(dest, "deadbeef" * 8) assert skip is False assert "перекачиваю" in reason + + +def test_fmt_bytes(): + assert fmt_bytes(500) == "500B" + assert "KB" in fmt_bytes(2048) + assert "MB" in fmt_bytes(5 * 1024 * 1024) + assert "GB" in fmt_bytes(2 * 1024**3) + + +def test_progress_line_with_total(): + line = progress_line("file", 512 * 1024 * 1024, 1024 * 1024 * 1024, 10 * 1024 * 1024, width=10) + assert "[" in line and "]" in line + assert "50.0%" in line + assert "/s" in line + assert "file" in line + + +def test_progress_line_unknown_total(): + line = progress_line("file", 1024 * 1024, None, 100_000) + assert "%" not in line + assert "/s" in line