filter cached benchmarks by template server-side

lookup_cached_benchmark pulled every perf row for a gpu_name/num_gpus combo and
matched the template client-side, discarding the rest: 3275 rows for a common
combo like 1x RTX 5090, once per GPU in the sweep.

The benchmarks table is indexed on (template_hash, last_update desc) and
(template_id, last_update desc), so neither index could be used. Filter on the
template columns instead, newest-first and capped, which is index-backed and
returns ~2 rows. A workergroup is created with either template_id or
template_hash, so a row can carry either; query once per column and dedupe.

search_benchmarks grows order/limit params to pass these through.

Author: vastzuby

tests/cli/test_benchmarks_commands.py

@@ -244,7 +244,21 @@ def _run_cli(parse_argv, argv, *, create_resp=None, workers_seq=None,
     vast.client = MagicMock(api_key="k")
     vast.search_templates.return_value = [template]
     vast.search_offers.return_value = fake_offer_list
-    vast.search_benchmarks.return_value = list(benchmark_rows or [])
+    # Template matching happens server-side now, so the fake has to apply
+    # select_filters; returning every row regardless of query would let a
+    # different template's rows count as a cache hit.
+    def _fake_search_benchmarks(query=None, order=None, limit=None):
+        def matches(row):
+            for col, cond in (query or {}).items():
+                if "eq" in cond and row.get(col) != cond["eq"]:
+                    return False
+                if "gte" in cond and (row.get(col) or 0) < cond["gte"]:
+                    return False
+            return True
+        hits = [r for r in (benchmark_rows or []) if matches(r)]
+        return hits[:limit] if limit else hits
+
+    vast.search_benchmarks.side_effect = _fake_search_benchmarks
     vast.show_instance.return_value = instance_resp
     vast.create_endpoint.return_value = {"success": True, "result": 11}
     vast.delete_endpoint.return_value = {}
@@ -491,3 +505,26 @@ def test_cache_query_shape(self, parse_argv):
         assert abs((time.time() - cutoff)
                    - DEFAULT_CACHE_MAX_AGE_DAYS * 86400) < 60
 
+    def test_cache_query_is_template_filtered_and_bounded(self, parse_argv):
+        """The template is filtered server-side, one query per template column.
+
+        Both columns are indexed with last_update, so filtering server-side and
+        ordering newest-first keeps the query index-backed and bounded instead
+        of pulling every perf row for the GPU and discarding most of them.
+        """
+        from vastai.cli.commands.benchmarks import MAX_CACHE_ROWS
+        _, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "-y", "--raw"],
+            benchmark_rows=[_bench_row(20.0)],
+        )
+        calls = vast.search_benchmarks.call_args_list
+        assert [c.kwargs["query"].get("template_hash") for c in calls] == [
+            {"eq": "x"}, None]
+        assert [c.kwargs["query"].get("template_id") for c in calls] == [
+            None, {"eq": 99999}]
+        for c in calls:
+            assert c.kwargs["limit"] == MAX_CACHE_ROWS
+            assert c.kwargs["order"] == [{"col": "last_update", "dir": "desc"}]
+

vastai/api/offers.py

@@ -113,17 +113,24 @@ def search_templates(client: VastClient, query: dict = None) -> list:
     return r.json().get("templates", [])
 
 
-def search_benchmarks(client: VastClient, query: dict = None) -> list:
+def search_benchmarks(client: VastClient, query: dict = None, order: list = None,
+                      limit: int = None) -> list:
     """Search for benchmarks using a query dict.
 
     Args:
         client: VastClient instance.
         query: Pre-parsed query dict of select_filters.
+        order: List of {"col": ..., "dir": ...} dicts, e.g. [{"col": "last_update", "dir": "desc"}].
+        limit: Max number of results. Omit for an unbounded result set.
 
     Returns:
         List of benchmark dicts.
     """
     query_args = {"select_cols": ["*"], "select_filters": query or {}}
+    if order is not None:
+        query_args["order_by"] = order
+    if limit is not None:
+        query_args["limit"] = int(limit)
     r = client.get("/benchmarks", query_args=query_args)
     r.raise_for_status()
     return r.json()

vastai/cli/commands/benchmarks.py

@@ -303,30 +303,49 @@ def pick_num_gpus(gpu_name, extra_filters):
 
 
 DEFAULT_CACHE_MAX_AGE_DAYS = 30
+MAX_CACHE_ROWS = 100  # newest rows per template column; caps popular template+GPU combos
 
 
 def lookup_cached_benchmark(vast, *, gpu_name, num_gpus, template_hash,
                              template_id, max_age_days):
     """Recent reported benchmarks for this exact spec (median + spread), or None.
 
-    Template is matched client-side: rows carry template_hash or template_id
-    depending on how the benchmarked workergroup was created.
+    Queried once per template column rather than filtered client-side: a
+    workergroup is created with either template_id or template_hash, so a row
+    can carry either one. Both columns are indexed with last_update on the
+    benchmarks table, so each query is index-backed instead of pulling every
+    perf row for the GPU and discarding the non-matching ones.
     """
-    query = {
+    base = {
         "type": {"eq": "perf"},
         "gpu_name": {"eq": gpu_name},
         "num_gpus": {"eq": num_gpus},
         "last_update": {"gte": time.time() - max_age_days * 86400},
     }
-    rows = vast.search_benchmarks(query=query)
-    if not isinstance(rows, list):
-        return None
+    rows = []
+    seen_ids = set()
+    for col, value in (("template_hash", template_hash), ("template_id", template_id)):
+        if not value:
+            continue
+        found = vast.search_benchmarks(
+            query=dict(base, **{col: {"eq": value}}),
+            order=[{"col": "last_update", "dir": "desc"}],
+            limit=MAX_CACHE_ROWS,
+        )
+        for r in found if isinstance(found, list) else []:
+            if not isinstance(r, dict):
+                continue
+            # A row carrying both columns comes back from both queries; dedupe on
+            # id, but keep rows without one rather than collapsing them together.
+            row_id = r.get("id")
+            if row_id is not None:
+                if row_id in seen_ids:
+                    continue
+                seen_ids.add(row_id)
+            rows.append(r)
     matched = [
         r for r in rows
-        if isinstance(r, dict)
-        and ((template_hash and r.get("template_hash") == template_hash)
-             or (template_id and r.get("template_id") == template_id))
-        and isinstance(r.get("value"), (int, float)) and r["value"] > 0
+        if isinstance(r.get("value"), (int, float)) and r["value"] > 0
     ]
     if not matched:
         return None

vastai/sdk.py

@@ -276,9 +276,12 @@ def search_templates(self, query: Optional[str] = None) -> list[dict]:
         """Search for templates."""
         return offers.search_templates(self.client, query=query)
 
-    def search_benchmarks(self, query: Optional[Union[str, dict]] = None) -> list[dict]:
+    def search_benchmarks(self, query: Optional[Union[str, dict]] = None,
+                          order: Optional[list] = None,
+                          limit: Optional[int] = None) -> list[dict]:
         """Search for benchmarks."""
-        return offers.search_benchmarks(self.client, query=query)
+        return offers.search_benchmarks(self.client, query=query, order=order,
+                                        limit=limit)
 
     def search_volumes(self, query: Optional[str] = None, **kwargs) -> list[dict]:
         """Search for volume offers."""