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

@@ -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}")
+            print(f"{prefix} качаю: {dest.name}", flush=True)
-            download(job["url"], dest, token)
+            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)

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}")
+        print(f"{prefix} качаю {name}", flush=True)
-        partial = dest.with_suffix(dest.suffix + ".partial")
         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:
+            download(url, dest, headers, label=f"{prefix} {name}")
-                while True:
+            print(f"{prefix} ok {name} ({fmt_bytes(dest.stat().st_size)})")
-                    chunk = resp.read(1024 * 1024)
-                    if not chunk:
-                        break
-                    out.write(chunk)
-            partial.replace(dest)
-            print(f"ok {name} ({dest.stat().st_size} bytes)")
         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:

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())

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