cache benchmark runs

.github/workflows/vast-sdk-testing.yml

@@ -99,7 +99,7 @@ jobs:
   # detect-changes filters this job out.
   # --------------------------------------------------------------------------
   unit-and-integration-check:
-    name: unit-and-integration-check
+    name: unit-and-integration
     needs: unit-and-integration
     if: always()
     runs-on: ubuntu-latest

tests/cli/test_benchmarks_commands.py

@@ -7,6 +7,7 @@
      measured_perf, and returns rows.
 """
 
+from contextlib import ExitStack
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -28,7 +29,7 @@ def _stub_gpu_catalog():
 
 
 # ---------------------------------------------------------------------------
-# _benchmark_gpu — cleanup invariant
+# benchmark_gpu — cleanup invariant
 # ---------------------------------------------------------------------------
 
 
@@ -80,7 +81,7 @@ def test_happy_path_returns_ok(self):
                           side_effect=[0, 0, 1, 2, 3, 4, 5, 6]):
             active_wgs = set()
             active_eps = set()
-            gpu, num_gpus, status, perf, err, price = bench._benchmark_gpu(
+            gpu, num_gpus, status, perf, err, price = bench.benchmark_gpu(
                 vast,
                 gpu_name="RTX 4080", num_gpus=1, timeout=60,
                 workergroups=active_wgs, endpoints=active_eps,
@@ -105,7 +106,7 @@ def test_template_id_wins_over_hash(self):
         with patch.object(bench.time, "sleep", return_value=None), \
              patch.object(bench.time, "monotonic",
                           side_effect=[0, 0, 1, 2, 3, 4, 5, 6]):
-            bench._benchmark_gpu(
+            bench.benchmark_gpu(
                 vast,
                 gpu_name="RTX 3060", num_gpus=1, timeout=60,
                 workergroups=set(), endpoints=set(),
@@ -123,7 +124,7 @@ def test_template_hash_used_when_no_id(self):
         with patch.object(bench.time, "sleep", return_value=None), \
              patch.object(bench.time, "monotonic",
                           side_effect=[0, 0, 1, 2, 3, 4, 5, 6]):
-            bench._benchmark_gpu(
+            bench.benchmark_gpu(
                 vast,
                 gpu_name="RTX 3060", num_gpus=1, timeout=60,
                 workergroups=set(), endpoints=set(),
@@ -141,7 +142,7 @@ def test_timeout_still_tears_down(self):
                           side_effect=[0, 0, 2, 3, 4, 5]):
             active_wgs = set()
             active_eps = set()
-            gpu, num_gpus, status, perf, err, price = bench._benchmark_gpu(
+            gpu, num_gpus, status, perf, err, price = bench.benchmark_gpu(
                 vast,
                 gpu_name="RTX 3060", num_gpus=1, timeout=1,
                 workergroups=active_wgs,
@@ -162,7 +163,7 @@ def test_create_failure_records_error_and_no_teardown(self):
             active_wgs = set()
             active_eps = set()
             with pytest.raises(RuntimeError):
-                bench._benchmark_gpu(
+                bench.benchmark_gpu(
                     vast,
                     gpu_name="RTX 3060", num_gpus=1, timeout=10,
                     workergroups=active_wgs,
@@ -176,7 +177,7 @@ def test_create_failure_records_error_and_no_teardown(self):
     def test_create_returns_no_id_reports_error(self):
         vast = _mk_vast(create_workergroup={"success": False})
         with patch.object(bench.time, "monotonic", return_value=0):
-            gpu, num_gpus, status, perf, err, price = bench._benchmark_gpu(
+            gpu, num_gpus, status, perf, err, price = bench.benchmark_gpu(
                 vast,
                 gpu_name="RTX 3060", num_gpus=1, timeout=10,
                 workergroups=set(), endpoints=set(),
@@ -196,7 +197,7 @@ def test_all_workers_terminal_bails_fast(self):
         with patch.object(bench.time, "sleep"), \
              patch.object(bench.time, "monotonic",
                           side_effect=[0, 0, 0, 5, 5, 50, 50]):
-            gpu, num_gpus, status, perf, err, price = bench._benchmark_gpu(
+            gpu, num_gpus, status, perf, err, price = bench.benchmark_gpu(
                 vast,
                 gpu_name="RTX 3060", num_gpus=1, timeout=600,
                 workergroups=set(), endpoints=set(),
@@ -214,7 +215,7 @@ def test_delete_failure_does_not_raise(self):
                         show_instance={"dph_total": 1.0})
         with patch.object(bench.time, "sleep"), \
              patch.object(bench.time, "monotonic", side_effect=[0, 0, 1]):
-            gpu, num_gpus, status, perf, err, price = bench._benchmark_gpu(
+            gpu, num_gpus, status, perf, err, price = bench.benchmark_gpu(
                 vast,
                 gpu_name="RTX 3060", num_gpus=1, timeout=60,
                 workergroups=set(), endpoints=set(),
@@ -232,7 +233,8 @@ def test_delete_failure_does_not_raise(self):
 
 def _run_cli(parse_argv, argv, *, create_resp=None, workers_seq=None,
              rental_dph=None, template=None, preflight_offers=1,
-             create_workergroup_raises=None):
+             create_workergroup_raises=None, benchmark_rows=None,
+             input_side_effect=None):
     """Parse argv and invoke the command with a mocked VastAI."""
     template = template if template is not None else _FAKE_TEMPLATE
     fake_offer_list = [{"id": i} for i in range(preflight_offers)]
@@ -242,6 +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
+    # 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 = {}
@@ -254,8 +271,13 @@ def _run_cli(parse_argv, argv, *, create_resp=None, workers_seq=None,
     vast.get_endpoint_workers.side_effect = list(workers_seq or [[]])
     vast.delete_workergroup.return_value = {}
 
-    with patch.object(bench, "VastAI", return_value=vast), \
-         patch.object(bench.time, "sleep"):
+    ctx = [patch.object(bench, "VastAI", return_value=vast),
+           patch.object(bench.time, "sleep")]
+    if input_side_effect is not None:
+        ctx.append(patch("builtins.input", side_effect=input_side_effect))
+    with ExitStack() as stack:
+        for c in ctx:
+            stack.enter_context(c)
         args = parse_argv(argv)
         return args.func(args), vast
 
@@ -297,8 +319,10 @@ def test_endpoint_name_includes_gpu_spec(self, parse_argv):
             rental_dph=0.5,
         )
         name = vast.create_endpoint.call_args.kwargs["endpoint_name"]
-        assert name.startswith("benchmark 1x RTX 3060 (")
-        assert name.endswith(")")
+        assert name.startswith("benchmark 1x RTX 3060 ")
+        # Backend rejects shell chars in endpoint_name (e.g. parens), so the
+        # uniqueness suffix must avoid them.
+        assert not set(name) & set(";&|(){}$`<>*?[]")
 
     def test_endpoint_deleted_even_on_exception(self, parse_argv):
         rows, vast = _run_cli(
@@ -344,3 +368,163 @@ def test_template_not_found_errors(self, parse_argv, capsys):
             rc = args.func(args)
             assert rc == 1
             assert "not found" in capsys.readouterr().err
+
+
+# ---------------------------------------------------------------------------
+# Cached-benchmark pre-flight
+# ---------------------------------------------------------------------------
+
+
+def _bench_row(value, *, template_hash="x", template_id=None, age_days=1):
+    import time
+    return {"type": "perf", "gpu_name": "RTX 3060", "num_gpus": 1,
+            "template_hash": template_hash, "template_id": template_id,
+            "value": value, "last_update": time.time() - age_days * 86400}
+
+
+class TestBenchmarkCache:
+    def test_cache_hit_skips_rental_and_uses_median(self, parse_argv):
+        rows, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "-y", "--raw"],
+            benchmark_rows=[_bench_row(10.0), _bench_row(30.0),
+                            _bench_row(20.0)],
+        )
+        assert rows[0]["status"] == "cached"
+        assert rows[0]["measured_perf"] == 20.0
+        vast.create_endpoint.assert_not_called()
+        vast.create_workergroup.assert_not_called()
+
+    def test_cached_rows_have_no_price(self, parse_argv):
+        # Cached rows report perf only; $/hr would come from a different
+        # machine than the one benchmarked, so it is omitted.
+        rows, _ = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "-y", "--raw"],
+            benchmark_rows=[_bench_row(20.0)],
+        )
+        assert rows[0]["rental_dph"] is None
+        assert rows[0]["perf_per_dollar"] is None
+
+    def test_cached_unrentable_is_flagged(self, parse_argv, capsys):
+        rows, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "-y", "--raw"],
+            benchmark_rows=[_bench_row(10.0)],
+            preflight_offers=0,
+        )
+        assert rows[0]["status"] == "cached"
+        assert rows[0]["measured_perf"] == 10.0
+        vast.create_endpoint.assert_not_called()
+        err = " ".join(capsys.readouterr().err.split())  # rich may wrap lines
+        assert "no offers available to rent" in err.lower()
+
+    def test_cache_hit_interactive_decline_uses_cache(self, parse_argv, capsys):
+        # No -y/--raw: user is prompted and declines (empty input) -> keep cached, no rental.
+        _, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "--timeout", "60"],
+            benchmark_rows=[_bench_row(20.0)],
+            input_side_effect=[""],
+        )
+        vast.create_endpoint.assert_not_called()
+        out = capsys.readouterr()
+        assert "cached" in (out.out + out.err)
+
+    def test_cache_hit_interactive_accept_runs_fresh(self, parse_argv):
+        # User opts into a fresh run ('y'), then confirms the rental ('y').
+        _, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "--timeout", "60"],
+            benchmark_rows=[_bench_row(20.0)],
+            workers_seq=[[{"id": 1, "measured_perf": 5.0, "status": "idle"}]],
+            rental_dph=0.5,
+            input_side_effect=["y", "y"],
+        )
+        vast.create_endpoint.assert_called()
+        vast.create_workergroup.assert_called_once()
+
+    def test_no_cache_bypasses_cache(self, parse_argv):
+        rows, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "--timeout", "60", "-y", "--raw",
+             "--no-cache"],
+            benchmark_rows=[_bench_row(10.0)],
+            workers_seq=[[{"id": 1, "measured_perf": 5.0, "status": "idle"}]],
+            rental_dph=0.5,
+        )
+        vast.search_benchmarks.assert_not_called()
+        assert rows[0]["status"] == "ok"
+        assert rows[0]["measured_perf"] == 5.0
+
+    def test_other_template_rows_are_a_miss(self, parse_argv):
+        rows, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "--timeout", "60", "-y", "--raw"],
+            benchmark_rows=[_bench_row(10.0, template_hash="other")],
+            workers_seq=[[{"id": 1, "measured_perf": 5.0, "status": "idle"}]],
+            rental_dph=0.5,
+        )
+        assert rows[0]["status"] == "ok"
+        vast.create_workergroup.assert_called_once()
+
+    def test_row_matching_by_template_id_only(self, parse_argv):
+        # Autoscaler rows may carry only template_id (no hash) depending on
+        # how the workergroup was created.
+        rows, _ = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "-y", "--raw"],
+            benchmark_rows=[_bench_row(7.0, template_hash=None,
+                                       template_id=99999)],
+        )
+        assert rows[0]["status"] == "cached"
+        assert rows[0]["measured_perf"] == 7.0
+
+    def test_cache_query_shape(self, parse_argv):
+        import time
+        from vastai.cli.commands.benchmarks import DEFAULT_CACHE_MAX_AGE_DAYS
+        _, vast = _run_cli(
+            parse_argv,
+            ["run", "benchmarks", "--template_id", "99999",
+             "--gpus", "RTX_3060", "-y", "--raw"],
+            benchmark_rows=[_bench_row(20.0)],
+        )
+        query = vast.search_benchmarks.call_args.kwargs["query"]
+        assert query["type"] == {"eq": "perf"}
+        assert query["gpu_name"] == {"eq": "RTX 3060"}
+        assert query["num_gpus"] == {"eq": 1}
+        cutoff = query["last_update"]["gte"]
+        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"}]
+

tests/cli/test_benchmarks_gpu_catalog.py

@@ -50,7 +50,7 @@ def test_parse_gpu_spec_resolves_catalog_only_name(monkeypatch):
     monkeypatch.setattr(bench, "_get_gpu_types", lambda: [
         {"canonical_name": "RTX 6090", "gpu_ram_mb": 49152},
     ])
-    assert bench._parse_gpu_spec("2x rtx_6090", 1) == ("RTX 6090", 2)
+    assert bench.parse_gpu_spec("2x rtx_6090", 1) == ("RTX 6090", 2)
 
 
 # ---------------------------------------------------------------------------

tests/cli/test_update_command.py

@@ -157,7 +157,7 @@ def test_prune_failure_does_not_fail_the_update(self, install_root, fake_uv):
         fake_uv.write_text(FAKE_UV.replace('touch "$(dirname "$0")/cache-pruned"', "exit 1"))
         _seed_env(install_root, "1.2.3")
         perform_update("1.3.0", MANIFEST)
-        assert "1.3.0" in (install_root / "env" / "bin" / "vastai").read_text()
+        assert "1.3.0" in (install_root / "current" / "bin" / "vastai").read_text()
 
     def test_build_failure_surfaces_real_error_not_a_fabricated_one(self, install_root):
         # A uv that "succeeds" but never builds the env: the verify step then

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