feat: show model download progress in aiod status (#3)

Track the model's download size (weights_gb) on the instance at spin
time, then in `aiod status` pull live disk_usage + gpu_util from the
provider and render a Download row, e.g.
  Download   239 / 343 GB  (70%) — downloading weights
Phase is inferred from telemetry (downloading -> loading into VRAM once
the GPU is active) and the row hides once the model is serving. Older
instances backfill weights_gb on first status via a sizing lookup.

state.load() now drops unknown fields so an older binary won't lose a
newer instance. Adds tests/test_state.py.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: nathanmarlor

aiod/cli.py

@@ -458,6 +458,7 @@ def spin(
             status="creating",
             provider=provider,
             idle_minutes=idle_m,
+            weights_gb=plan.weights_gb,
         )
         state.save(inst)
         console.print(f"[green]✓[/] Instance [bold]{instance_id}[/] created. Waiting for boot...")
@@ -647,18 +648,33 @@ def status():
 
     s = Settings.load()
     live_status = "?"
+    disk_usage = gpu_util = None
     if providers.api_key_for(inst.provider, s):
         try:
             with providers.get_client(inst.provider, s) as client:
                 vi = client.get_instance(inst.instance_id)
                 live_status = client.status_of(vi)
+                disk_usage = vi.get("disk_usage")
+                gpu_util = vi.get("gpu_util")
                 ep = client.endpoint_of(vi, CONTAINER_PORT)
                 if ep and not inst.base_url:
                     inst.host, inst.port = ep
                     state.save(inst)
         except providers.PROVIDER_ERRORS as e:
             live_status = f"error: {e}"
 
+    # Backfill the download size for instances spun before it was tracked, so the
+    # progress % still works. Best-effort: a sizing call hits HF, so only do it
+    # while still loading and cache the result back to state.
+    if inst.weights_gb is None and inst.status != "running" and disk_usage is not None:
+        try:
+            plan = size_any(inst.repo_id, hf_token=s.hf_token).plan(inst.quant)
+            if plan:
+                inst.weights_gb = plan.weights_gb
+                state.save(inst)
+        except Exception:  # noqa: BLE001 - cosmetic; never break `status`
+            pass
+
     over = inst.expires_in_hours < 0
     table = Table(show_header=False)
     table.add_row("Instance", str(inst.instance_id))
@@ -674,6 +690,10 @@ def status():
         if over
         else f"{inst.expires_in_hours:.2f} h left",
     )
+    if inst.status != "running":
+        dl = state.download_progress(disk_usage, inst.weights_gb, gpu_util)
+        if dl:
+            table.add_row("Download", dl)
     console.print(table)
     if evs:
         _print_events(evs)

aiod/engine.py

@@ -96,6 +96,7 @@ def emit(phase: str, msg: str = "") -> None:
                 status="creating",
                 provider=provider,
                 idle_minutes=idle_minutes,
+                weights_gb=plan.weights_gb,
             )
             state.save(inst)
 

aiod/state.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 import json
 import time
 from dataclasses import asdict, dataclass
@@ -30,6 +31,7 @@ class Instance:
     status: str = "creating"  # creating | loading | running | error
     provider: str = "vast"  # which backend rented this (vast | runpod | ...)
     idle_minutes: int | None = None  # auto-shutdown threshold, if a watcher is running
+    weights_gb: float | None = None  # model download size, for download-progress %
 
     @property
     def base_url(self) -> str | None:
@@ -50,6 +52,27 @@ def est_cost_so_far(self) -> float:
         return self.age_hours * self.price_per_hr
 
 
+def download_progress(
+    disk_usage_gb: float | None, weights_gb: float | None, gpu_util: float | None = None
+) -> str | None:
+    """A human 'Download' line from live telemetry, e.g.
+    '135 / 343 GB  (39%) — downloading weights'. Returns None when there's no
+    usable figure. `disk_usage_gb` counts the container image too, so the % runs a
+    little optimistic; GPU activity is what tells us the load phase has started."""
+    if disk_usage_gb is None:
+        return None
+    if not weights_gb:
+        return f"{disk_usage_gb:.0f} GB on disk"
+    pct = min(100.0, disk_usage_gb / weights_gb * 100)
+    if gpu_util and gpu_util > 0:
+        phase = "loading into VRAM"
+    elif pct >= 99:
+        phase = "download complete — loading"
+    else:
+        phase = "downloading weights"
+    return f"{disk_usage_gb:.0f} / {weights_gb:.0f} GB  ({pct:.0f}%) — {phase}"
+
+
 def save(inst: Instance) -> None:
     STATE_DIR.mkdir(parents=True, exist_ok=True)
     STATE_FILE.write_text(json.dumps(asdict(inst), indent=2))
@@ -60,7 +83,10 @@ def load() -> Instance | None:
         return None
     try:
         data = json.loads(STATE_FILE.read_text())
-        return Instance(**data)
+        # Tolerate fields written by a newer version (drop unknowns) so an older
+        # binary doesn't lose track of a live instance.
+        fields = {f.name for f in dataclasses.fields(Instance)}
+        return Instance(**{k: v for k, v in data.items() if k in fields})
     except (json.JSONDecodeError, TypeError):
         return None
 

aiod/tui.py

@@ -374,6 +374,7 @@ def _launch_worker(self, quant: str, best: PricedOption, provider: str, idle_m:
                     status="creating",
                     provider=provider,
                     idle_minutes=idle_m,
+                    weights_gb=(p.weights_gb if (p := self.sizing.plan(quant)) else None),
                 )
                 state.save(inst)
                 self.call_from_thread(self._log, f"[green]✓[/] Instance {instance_id} created.")

tests/test_state.py

@@ -0,0 +1,49 @@
+from aiod import state
+
+
+def test_download_progress_percent_and_phase():
+    line = state.download_progress(135, 343, gpu_util=0)
+    assert "135 / 343 GB" in line
+    assert "(39%)" in line
+    assert "downloading weights" in line
+
+
+def test_download_progress_gpu_active_means_loading():
+    line = state.download_progress(343, 343, gpu_util=42)
+    assert "loading into VRAM" in line
+
+
+def test_download_progress_complete_but_idle_gpu():
+    line = state.download_progress(343, 343, gpu_util=0)
+    assert "(100%)" in line
+    assert "download complete" in line
+
+
+def test_download_progress_caps_at_100():
+    # disk_usage counts the image too, so it can exceed the weights size.
+    line = state.download_progress(360, 343, gpu_util=0)
+    assert "(100%)" in line
+
+
+def test_download_progress_no_target_size():
+    assert state.download_progress(80, None) == "80 GB on disk"
+
+
+def test_download_progress_no_telemetry():
+    assert state.download_progress(None, 343) is None
+
+
+def test_load_tolerates_unknown_fields(tmp_path, monkeypatch):
+    monkeypatch.setattr(state, "STATE_DIR", tmp_path)
+    monkeypatch.setattr(state, "STATE_FILE", tmp_path / "state.json")
+    inst = state.Instance(
+        instance_id=1, repo_id="org/m", quant="bf16", gpu_desc="1x X",
+        price_per_hr=1.0, created_at=0.0, ttl_hours=4.0, weights_gb=343.0,
+    )
+    state.save(inst)
+    # Simulate a field written by a newer version.
+    raw = (tmp_path / "state.json").read_text().rstrip().rstrip("}")
+    (tmp_path / "state.json").write_text(raw + ',  "future_field": 99\n}')
+    loaded = state.load()
+    assert loaded is not None
+    assert loaded.weights_gb == 343.0