feat: --gpu flag to steer offer search toward specific hardware (#2)

* feat: --gpu flag to steer offer search toward specific hardware

Add a repeatable, case-insensitive substring filter on GPU name to
`aiod spin` and `aiod estimate`, e.g. `--gpu rtx6000` to target RTX
6000Ada / A6000 cards.

- vast: fetch a wider offer pool, filter by normalised name substring,
  trim to limit (vast's gpu_name API filter is exact-match only).
- runpod: filter gpu types by displayName substring.
- spin: "no offer found matching --gpu ..." when the filter empties results.

* fix: show real rented GPU in estimate table + smarter --gpu token matching

The estimate table showed the sizing tier (e.g. 'A100 80GB') instead of
the actual matched offer, so --gpu looked broken even when it matched.
Show offer.gpu_name/num_gpus instead.

Matching is now token-AND within a query (whitespace-split): --gpu 'rtx 6000'
matches 'RTX PRO 6000 WS'. Shared gpu_name_matches() used by vast + runpod.

Author: nathanmarlor

aiod/cli.py

@@ -233,6 +233,9 @@ def estimate(
     context: int = typer.Option(None, "--context", help="Context length for the KV estimate"),
     concurrency: int = typer.Option(4, "--concurrency", help="Concurrent sequences for KV estimate"),
     max_price: float = typer.Option(None, "--max-price", help="Cap $/hr in the offer search"),
+    gpu: list[str] = typer.Option(
+        None, "--gpu", help="Only consider GPUs whose name contains this (repeatable), e.g. --gpu 'rtx 6000'"
+    ),
 ):
     """Size a model from its HuggingFace link and show live provider cost ($0)."""
     provider = provider.lower()
@@ -274,10 +277,11 @@ def estimate(
             with console.status(f"Querying live {provider} offers..."):
                 for p in sizing.plans:
                     disk = recommend_disk_gb(p.weights_gb)
-                    priced = client.price_plan(p, disk, max_price=max_p)
+                    priced = client.price_plan(p, disk, max_price=max_p, gpu_match=gpu or None)
                     best = _pick_cheapest(priced)
                     if best:
-                        fit = f"{best.option.num_gpus}x {best.option.tier.name}"
+                        # Show the real rented card, not the sizing tier proxy.
+                        fit = f"{best.offer.num_gpus}x {best.offer.gpu_name}"
                         hr = f"${best.offer.dph_total:.2f}"
                         four = f"${best.offer.dph_total * 4:.2f}"
                     else:
@@ -306,6 +310,9 @@ def spin(
         None, "--quant", "-q", help="vLLM: bf16/fp8/awq-int4 · GGUF: a repo quant tag"
     ),
     max_price: float = typer.Option(None, "--max-price", help="Hard cap $/hr (default from .env)"),
+    gpu: list[str] = typer.Option(
+        None, "--gpu", help="Only rent GPUs whose name contains this (repeatable), e.g. --gpu 'rtx 6000'"
+    ),
     ttl: float = typer.Option(None, "--ttl", help="Auto-destroy reminder window, hours"),
     idle: int = typer.Option(
         None, "--idle", help="Auto-shutdown after N idle minutes (starts a local watcher)"
@@ -388,12 +395,13 @@ def spin(
 
     with client_cm as client:
         with console.status("Finding the cheapest GPU that fits..."):
-            priced = client.price_plan(plan, disk, max_price=max_p)
+            priced = client.price_plan(plan, disk, max_price=max_p, gpu_match=gpu or None)
         best = _pick_cheapest(priced)
         if not best:
+            gpu_note = f" matching --gpu {' '.join(gpu)}" if gpu else ""
             console.print(
                 f"[red]No {provider} offer found[/] under ${max_p:.2f}/hr for {m.repo_id} "
-                f"({quant}). Try a higher --max-price or a smaller quant."
+                f"({quant}){gpu_note}. Try a higher --max-price or a smaller quant."
             )
             raise typer.Exit(1)
 

aiod/runpod.py

@@ -25,7 +25,7 @@
 
 from .bootstrap import ServerConfig
 from .sizing import QuantPlan
-from .vast import PricedOption
+from .vast import PricedOption, gpu_name_matches
 
 REST = "https://rest.runpod.io/v1"
 GRAPHQL = "https://api.runpod.io/graphql"
@@ -94,13 +94,20 @@ def _fetch_gpu_types(self) -> list[dict]:
         self._gpu_types = data.get("data", {}).get("gpuTypes", []) or []
         return self._gpu_types
 
-    def _cheapest_gpu_for(self, vram_gb: float) -> dict | None:
+    def _cheapest_gpu_for(
+        self, vram_gb: float, gpu_match: list[str] | None = None
+    ) -> dict | None:
         """Cheapest SECURE gpu type in the same VRAM class as `vram_gb`."""
+        queries = [g for g in (gpu_match or []) if g.strip()]
         cands: list[tuple[float, dict]] = []
         for g in self._fetch_gpu_types():
             mem, price = g.get("memoryInGb"), g.get("securePrice")
             if not mem or not price or not g.get("secureCloud"):
                 continue
+            if queries:
+                name = str(g.get("displayName", g.get("id", "")))
+                if not gpu_name_matches(name, queries):
+                    continue
             if vram_gb * 0.95 <= mem <= vram_gb * 1.6:
                 cands.append((price, g))
         if not cands:
@@ -109,14 +116,14 @@ def _cheapest_gpu_for(self, vram_gb: float) -> dict | None:
 
     def price_plan(
         self, plan: QuantPlan, min_disk_gb: float, max_price: float | None = None,
-        max_candidates: int = 3,
+        max_candidates: int = 3, gpu_match: list[str] | None = None,
     ) -> list[PricedOption]:
         """Price the few least-wasteful configs that fit (mirrors the vast backend;
         all pricing comes from one cached GraphQL call, so this is cheap)."""
         opts = sorted((o for o in plan.options if o.fits), key=lambda o: o.total_vram_gb)
         priced: list[PricedOption] = []
         for opt in opts[:max_candidates]:
-            g = self._cheapest_gpu_for(opt.tier.vram_gb)
+            g = self._cheapest_gpu_for(opt.tier.vram_gb, gpu_match=gpu_match)
             offer = None
             if g:
                 total = g["securePrice"] * opt.num_gpus

aiod/vast.py

@@ -36,6 +36,19 @@ class VastError(Exception):
     pass
 
 
+def gpu_name_matches(name: str, queries: list[str]) -> bool:
+    """True if `name` matches any --gpu query. A query matches when ALL its
+    whitespace-separated tokens appear (case-insensitively) in the GPU name, so
+    `rtx 6000` matches "RTX PRO 6000 WS" and `a6000` matches "RTX A6000". A single
+    glued token like `rtx6000` only matches a contiguous run, so prefer spaces."""
+    n = name.lower()
+    for q in queries:
+        toks = q.lower().split()
+        if toks and all(t in n for t in toks):
+            return True
+    return False
+
+
 @dataclass
 class Offer:
     id: int
@@ -115,6 +128,7 @@ def search_offers(
         min_compute_cap: int = 800,
         min_inet_mbps: int = 400,
         gpu_names: list[str] | None = None,
+        gpu_match: list[str] | None = None,
         limit: int = 30,
     ) -> list[Offer]:
         query: dict = {
@@ -137,16 +151,22 @@ def search_offers(
             "inet_down": {"gte": min_inet_mbps},
             "type": "ondemand",
             "order": [["dph_total", "asc"]],
-            "limit": limit,
+            # When substring-filtering by GPU name we fetch a wider pool and trim
+            # client-side, since vast's gpu_name filter only does exact matches.
+            "limit": max(limit, 50) if gpu_match else limit,
         }
         if max_price is not None:
             query["dph_total"] = {"lte": float(max_price)}
         if gpu_names:
             query["gpu_name"] = {"in": gpu_names}
 
         data = self._post("/api/v0/bundles/", query)
-        offers = data.get("offers", []) or []
-        return [self._parse_offer(o) for o in offers]
+        offers = [self._parse_offer(o) for o in (data.get("offers", []) or [])]
+        if gpu_match:
+            queries = [g for g in gpu_match if g.strip()]
+            offers = [o for o in offers if gpu_name_matches(o.gpu_name, queries)]
+            offers = offers[:limit]
+        return offers
 
     @staticmethod
     def _parse_offer(o: dict) -> Offer:
@@ -271,6 +291,7 @@ def price_plan(
         min_disk_gb: float,
         max_price: float | None = None,
         max_candidates: int = 3,
+        gpu_match: list[str] | None = None,
     ) -> list[PricedOption]:
         """Price the few least-wasteful GPU configs that fit and return them all
         (the caller picks the cheapest by actual price). We sort fitting configs by
@@ -298,6 +319,7 @@ def price_plan(
                 max_price=max_price,
                 min_compute_cap=min_cc,
                 min_inet_mbps=min_inet,
+                gpu_match=gpu_match,
                 limit=10,
             )
             priced.append(PricedOption(option=opt, offer=offers[0] if offers else None))