diff --git a/.github/workflows/vast-sdk-testing.yml b/.github/workflows/vast-sdk-testing.yml index a38ec598..c7aacde8 100644 --- a/.github/workflows/vast-sdk-testing.yml +++ b/.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 diff --git a/tests/cli/test_benchmarks_commands.py b/tests/cli/test_benchmarks_commands.py index 9af34438..4601a87c 100644 --- a/tests/cli/test_benchmarks_commands.py +++ b/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"}] + diff --git a/tests/cli/test_benchmarks_gpu_catalog.py b/tests/cli/test_benchmarks_gpu_catalog.py index 465c40de..be9c8508 100644 --- a/tests/cli/test_benchmarks_gpu_catalog.py +++ b/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) # --------------------------------------------------------------------------- diff --git a/tests/cli/test_update_command.py b/tests/cli/test_update_command.py index 930ca41b..6a4036c9 100644 --- a/tests/cli/test_update_command.py +++ b/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 diff --git a/vastai/api/offers.py b/vastai/api/offers.py index da503ce9..c00bdd2e 100644 --- a/vastai/api/offers.py +++ b/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() diff --git a/vastai/cli/commands/benchmarks.py b/vastai/cli/commands/benchmarks.py index dde89e32..bd9af8f5 100644 --- a/vastai/cli/commands/benchmarks.py +++ b/vastai/cli/commands/benchmarks.py @@ -1,6 +1,8 @@ -"""vastai run benchmarks: benchmark a template against one or more GPUs. -Each GPU is rented in parallel; the per-GPU measured perf comes from the -template's pyworker benchmark. +""" +A CLI command to benchmark a template against multiple GPUs in parallel. +The template's pyworker benchmark determines the measured perf (ex. tokens/$). + +vastai run benchmarks """ import argparse @@ -11,6 +13,7 @@ import random import re import signal +import statistics import sys import threading import time @@ -35,34 +38,29 @@ # Default GPU sweep set + per-card VRAM fallback (MB). Catalog gpu_ram_mb # overrides these values when populated; the keys stay the editorial sweep set. -_DEFAULT_GPUS = { +DEFAULT_GPUS = { "RTX 5090": 32607, "RTX 4090": 24564, "RTX 3090": 24576, "RTX A6000": 49140, } -_ENDPOINT_CONFIG = {"cold_workers": 1, "max_workers": 1, "min_load": 1.0} +# One worker per GPU, kept alive while idle so we can read its measured perf. +ENDPOINT_CONFIG = {"cold_workers": 1, "max_workers": 1, "min_load": 1.0} -_WORKER_POLL_INTERVAL = 10.0 -_DEFAULT_BENCHMARK_TIMEOUT = 60 * 60 # seconds to wait for a benchmark to complete before giving up -_RENTAL_TIMEOUT = 120 # fail fast if the autoscaler hasn't rented anything by then -_MAX_ABANDONS = 3 # bail after this many workers get abandoned without one reaching idle -_PARALLEL_SUBMIT_DELAY = 4.0 # spacing between parallel submits +RENTAL_TIMEOUT = 120 +DEFAULT_BENCHMARK_TIMEOUT = 30 * 60 +WORKER_POLL_INTERVAL = 10.0 +PARALLEL_SUBMIT_DELAY = 4.0 # spacing between parallel endpoint submits +MAX_ABANDONS = 3 # bail after this many workers vanish without one reaching idle -# Worker states the autoscaler does NOT recover from. -# these will timeout if they persist for >_TERMINAL_GRACE_SECONDS without any worker reaching idle -# `error` is excluded since the autoscaler retries with error -> rebooting -> model_loading. -_TERMINAL_STATES = {"stopped", "destroying", "unavail"} -_TERMINAL_GRACE_SECONDS = 30 # how long a worker can sit in a terminal state before we give up on it -# statuses to use for the live table -_TERMINAL_STATUSES = {"done", "skipped", "timeout", "failed", "no_worker"} +# states in which the autoscaler shouldnt retry the worker +TERMINAL_STATES = {"stopped", "destroying", "unavail"} +TERMINAL_GRACE_SECONDS = 30 +TERMINAL_STATUSES = {"done", "skipped", "timeout", "failed", "no_worker"} -# key in worker states that says "no worker rented yet" -_WAITING_KEY = "__waiting__" - -def _format_time_elapsed(seconds): +def format_time_elapsed(seconds): m, s = divmod(int(seconds), 60) return f"{m}:{s:02d}" @@ -101,24 +99,21 @@ def _catalog_vram_map(): def _per_card_vram_mb(gpu_name): """Per-card VRAM in MB: catalog gpu_ram_mb if populated, else the hardcoded - _DEFAULT_GPUS value. None when neither source knows the GPU. + DEFAULT_GPUS value. None when neither source knows the GPU. """ - return _catalog_vram_map().get(gpu_name) or _DEFAULT_GPUS.get(gpu_name) - + return _catalog_vram_map().get(gpu_name) or DEFAULT_GPUS.get(gpu_name) -def _parse_gpu_spec(token, default_num_gpus): - """parses the gpu arg into (gpu_name, num_gpus) - ex. "2x RTX_4090" -> ("RTX 4090", 2) - if both --num_gpus and inline Nx (1x, 2x, etc) prefix are given, Nx wins - if no count per gpu isprovided, uses default num_gpus +def parse_gpu_spec(token, default_num_gpus): + """Parse a GPU token into (canonical_gpu_name, num_gpus). + Example: 2x RTX_4090" -> ("RTX 4090", 2). """ - # so users can input "rtx_4090" or "RTX 4090" or "Rtx_4090" and all resolves to "RTX 4090" + # so "rtx_4090", "RTX 4090", "Rtx_4090" all map to "RTX 4090" canonical = _canonical_gpu_names() token = token.strip() - # converts "2x RTX_4090" to ("RTX 4090", 2) + # string to tuple. "2x RTX_4090" to ("RTX 4090", 2) m = re.match(r"^(\d+)[\s_]*x[\s_]*(.+)$", token, re.IGNORECASE) if m: raw_name, n = m.group(2).strip().replace("_", " "), int(m.group(1)) @@ -127,7 +122,7 @@ def _parse_gpu_spec(token, default_num_gpus): return (canonical.get(raw_name.lower(), raw_name), n) -def _format_filter_query(filters): +def format_filter_query(filters): """Render extra_filters dict into comparison string. e.g. {"gpu_ram": {"gte": 16000}} -> "gpu_ram>=16000" """ @@ -142,14 +137,13 @@ def _format_filter_query(filters): # ------------------------------------------------------------------------------------ -# the next few functions are pre flight helpers -# they check whether each (template, gpu_name, and num_gpus) is rentable -# so we dont run the benchmark only for it to fail due to 0 matching offers -# if no mathcing offers, returns user friendly msg -# ----------------------------------------------------------------------------------- +# Pre-flight helpers: check whether each (template, gpu_name, num_gpus) is rentable +# before we benchmark it, so a run can't fail late due to 0 matching offers. On a +# miss they build a user-friendly explanation of which filter blocked it. +# ------------------------------------------------------------------------------------ -def _has_matching_offer(vast, *, gpu_name, num_gpus, extra_filters): +def has_matching_offer(vast, *, gpu_name, num_gpus, extra_filters): """True if any verified+rentable offer matches the template's filters""" query = dict(extra_filters or {}) query["gpu_name"] = {"eq": gpu_name} @@ -157,14 +151,14 @@ def _has_matching_offer(vast, *, gpu_name, num_gpus, extra_filters): return bool(vast.search_offers(query=query, limit=1)) -def _format_skip_message(vast, *, gpu_name, num_gpus, extra_filters): - """Build a user friendlyexplanation for why a GPU has 0 matching offers. - Calls search offers to find the blocking filter +def format_skip_message(vast, *, gpu_name, num_gpus, extra_filters): + """Build a user-friendly explanation for why a GPU has 0 matching offers. + Calls search_offers to find the blocking filter. """ if not extra_filters: return (f"no offers for {gpu_name} (num_gpus={num_gpus})") - diag = _find_blockers( + diag = find_blockers( vast, gpu_name=gpu_name, num_gpus=num_gpus, extra_filters=extra_filters, ) @@ -183,7 +177,7 @@ def _format_skip_message(vast, *, gpu_name, num_gpus, extra_filters): continue gpu_value = sample.get(key) for op, threshold in ops.items(): - if _check_template_filter(gpu_value, op, threshold) is True: + if check_template_filter(gpu_value, op, threshold) is True: continue # search-engine quirk, skip sym = OP_TO_STR.get(op, op) if gpu_value is None: @@ -200,10 +194,10 @@ def _format_skip_message(vast, *, gpu_name, num_gpus, extra_filters): ops = extra_filters.get("gpu_total_ram") or {} if ops: op, value = next(iter(ops.items())) - raw = _min_gpus_for_ram( + raw = min_gpus_for_ram( value, op, diag.get("per_card_gpu_ram")) if raw and raw > num_gpus: - viable = _get_gpu_chunk_size(raw) + viable = get_gpu_chunk_size(raw) lines.append( f" hint: try {viable}x {gpu_name} " f"(host total then satisfies {value})") @@ -212,15 +206,13 @@ def _format_skip_message(vast, *, gpu_name, num_gpus, extra_filters): return ("0 offers match. No single filter is the culprit; the combination " f"of template filters excludes {gpu_name} " - f"({_format_filter_query(extra_filters)})") + f"({format_filter_query(extra_filters)})") -def _find_blockers(vast, *, gpu_name, num_gpus, extra_filters): - """Identify why ``_has_matching_offer`` returned False. +def find_blockers(vast, *, gpu_name, num_gpus, extra_filters): + """Find which template filter blocks this GPU from having offers. - Skips filters the GPU mathematically satisfies (avoids API call + dodges - Vast search-engine quirks where a passing filter still returns 0). Bails - on the first real blocker (user fixes one and re-runs). Returns a dict + Bails on the first real blocker. Returns a dict {base_count, single_blockers, per_card_gpu_ram, sample_offer}. """ base_query = {"gpu_name": {"eq": gpu_name}, @@ -236,11 +228,11 @@ def _find_blockers(vast, *, gpu_name, num_gpus, extra_filters): if isinstance(ops, dict) and sample is not None: gpu_value = sample.get(key) if gpu_value is not None and all( - _check_template_filter(gpu_value, op, threshold) is True + check_template_filter(gpu_value, op, threshold) is True for op, threshold in ops.items() ): continue - if not _has_matching_offer( + if not has_matching_offer( vast, gpu_name=gpu_name, num_gpus=num_gpus, extra_filters={key: ops}, ): @@ -250,7 +242,7 @@ def _find_blockers(vast, *, gpu_name, num_gpus, extra_filters): "per_card_gpu_ram": per_card, "sample_offer": sample} -def _check_template_filter(gpu_value, op, threshold): +def check_template_filter(gpu_value, op, threshold): """Checks if gpu_value matches the template filters (e.g. gpu_total_ram>=16000)""" if gpu_value is None or threshold is None: return None @@ -267,7 +259,7 @@ def _check_template_filter(gpu_value, op, threshold): return None -def _min_gpus_for_ram(threshold, op, per_card_gpu_ram): +def min_gpus_for_ram(threshold, op, per_card_gpu_ram): """Smallest ``num_gpus`` such that ``num_gpus * per_card_gpu_ram`` satisfies the template's ``gpu_total_ram`` filter. Returns None if any input is missing or the operator isn't a comparison we can solve. @@ -281,7 +273,7 @@ def _min_gpus_for_ram(threshold, op, per_card_gpu_ram): return None -def _get_gpu_chunk_size(n): +def get_gpu_chunk_size(n): """Round n up to the next common GPU chunk size (1, 2, 4, 8). Returns n unchanged for >8 since that's rare enough we'd rather be honest than guess. """ @@ -291,7 +283,7 @@ def _get_gpu_chunk_size(n): return n -def _pick_num_gpus(gpu_name, extra_filters): +def pick_num_gpus(gpu_name, extra_filters): """Smallest num_gpus where ``num_gpus * per_card_ram`` satisfies the template's ``gpu_total_ram`` filter, rounded up to a chassis size. Falls back to 1 when the template has no such filter or one card @@ -304,13 +296,71 @@ def _pick_num_gpus(gpu_name, extra_filters): if not ops: return 1 op, threshold = next(iter(ops.items())) - raw_min = _min_gpus_for_ram(threshold, op, per_card) + raw_min = min_gpus_for_ram(threshold, op, per_card) if raw_min is None or raw_min <= 1: return 1 - return _get_gpu_chunk_size(raw_min) + return get_gpu_chunk_size(raw_min) + + +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. + + 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. + """ + base = { + "type": {"eq": "perf"}, + "gpu_name": {"eq": gpu_name}, + "num_gpus": {"eq": num_gpus}, + "last_update": {"gte": time.time() - max_age_days * 86400}, + } + 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.get("value"), (int, float)) and r["value"] > 0 + ] + if not matched: + return None + values = [r["value"] for r in matched] + newest = max((r.get("last_update") or 0) for r in matched) + return { + "median": statistics.median(values), + "low": min(values), + "high": max(values), + "n": len(matched), + "age_days": max(0.0, (time.time() - newest) / 86400), + } -def _update_worker_states(worker_states, current_workers, gpu_name): +def update_worker_states(worker_states, current_workers, gpu_name): """Update worker_states with current poll, and print only on worker rotation (the live rich table already shows current status / elapsed). Tracks state_started per worker for the terminal-state grace check. @@ -318,12 +368,13 @@ def _update_worker_states(worker_states, current_workers, gpu_name): """ now = time.monotonic() abandoned = 0 + waiting_key = "__waiting__" # non-int sentinel: "no worker rented yet" if not current_workers: - worker_states.setdefault(_WAITING_KEY, {"started": now}) + worker_states.setdefault(waiting_key, {"started": now}) return abandoned - worker_states.pop(_WAITING_KEY, None) + worker_states.pop(waiting_key, None) seen = set() for w in current_workers: @@ -342,11 +393,11 @@ def _update_worker_states(worker_states, current_workers, gpu_name): for wid in list(worker_states): if not isinstance(wid, int): - continue # sentinel keys (e.g. _WAITING_KEY) + continue # skip the non-int waiting sentinel if wid not in seen: # Worker abandoned; live table only shows current worker, so rotation is otherwise invisible. prev = worker_states[wid] - elapsed = _format_time_elapsed(now - prev["state_started"]) + elapsed = format_time_elapsed(now - prev["state_started"]) print(f"[{gpu_name}] worker {wid} abandoned " f"(last state={prev['status']} for {elapsed})", file=sys.stderr) @@ -355,7 +406,7 @@ def _update_worker_states(worker_states, current_workers, gpu_name): return abandoned -def _update_row(class_states, row_id, **fields): +def update_row(class_states, row_id, **fields): """Update one GPU's row in the live-table state dict. Sets ``run_started`` on first non-``queued`` status and ``run_ended`` on terminal status (so elapsed freezes); each thread only writes its own @@ -367,16 +418,16 @@ def _update_row(class_states, row_id, **fields): new_status = fields.get("status") now_ts = time.monotonic() if new_status and new_status != cur.get("status"): - if "run_started" not in cur and new_status != "queued": + if "run_started" not in cur and new_status not in ("queued", "cached"): cur["run_started"] = now_ts - if new_status in _TERMINAL_STATUSES: + if new_status in TERMINAL_STATUSES: cur["run_ended"] = now_ts else: cur.pop("run_ended", None) # recovered from a transient terminal status; unfreeze elapsed cur.update(fields) -def _render_table(class_states): +def render_table(class_states): """Render the per-GPU progress as a rich.Table for the live display.""" table = Table(show_header=True, header_style="bold cyan") table.add_column("GPU", no_wrap=True) @@ -400,9 +451,9 @@ def _render_table(class_states): if run_started is None: elapsed_str = "-" elif run_ended is not None: - elapsed_str = _format_time_elapsed(run_ended - run_started) + elapsed_str = format_time_elapsed(run_ended - run_started) else: - elapsed_str = _format_time_elapsed(now - run_started) + elapsed_str = format_time_elapsed(now - run_started) perf = s.get("perf") dph = s.get("dph") pps = (perf / dph) if (perf and dph) else None @@ -419,7 +470,7 @@ def _render_table(class_states): return table -def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, +def benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, workergroups, endpoints, template_hash=None, template_id=None, auto_instance=None, autoscaler_url=None, @@ -430,11 +481,12 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, """ endpoint_id = None workergroup_id = None - # need the uuid so the endpoint name is unique if we benchmark the same GPU in parallel - endpoint_name = f"benchmark {num_gpus}x {gpu_name} ({uuid.uuid4().hex[:8]})" + # uuid suffix keeps the name unique when benchmarking the same GPU in parallel; + # no parens/shell chars, the backend rejects them in endpoint_name. + endpoint_name = f"benchmark {num_gpus}x {gpu_name} {uuid.uuid4().hex[:8]}" row_id = f"{num_gpus}x {gpu_name}" # id for the live table so 1x/2x/4x of the same GPU gets its own row start = time.monotonic() - _update_row(class_states, row_id, status="provisioning") + update_row(class_states, row_id, status="provisioning") search = f"gpu_name={gpu_name.replace(' ', '_')} num_gpus={num_gpus}" if autoscaler_url: status_vast = VastAI(api_key=vast.client.api_key, server_url=autoscaler_url, retry=vast.client.retry) @@ -442,7 +494,7 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, status_vast = vast try: - endpoint_kwargs = dict(_ENDPOINT_CONFIG, endpoint_name=endpoint_name) + endpoint_kwargs = dict(ENDPOINT_CONFIG, endpoint_name=endpoint_name) if auto_instance is not None: endpoint_kwargs["auto_instance"] = auto_instance ep_resp = vast.create_endpoint(**endpoint_kwargs) @@ -452,9 +504,9 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, return (gpu_name, num_gpus, "error", None, f"create_endpoint returned no id: {ep_resp!r}", None) endpoints.add(endpoint_id) - _update_row(class_states, row_id, endpoint_id=endpoint_id) + update_row(class_states, row_id, endpoint_id=endpoint_id) - # cold_workers and min_load are set on _ENDPOINT_CONFIG, not here; the autoscaler reads + # cold_workers and min_load are set on ENDPOINT_CONFIG, not here; the autoscaler reads # both from the endpoint group, not the workergroup (verified via autoscaler.cpp). workergroup_kwargs = dict( endpoint_id=endpoint_id, @@ -473,16 +525,16 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, return (gpu_name, num_gpus, "error", None, f"create_workergroup returned no id: {resp!r}", None) workergroups.add(workergroup_id) - _update_row(class_states, row_id, status="waiting_for_worker") + update_row(class_states, row_id, status="waiting_for_worker") worker_states = {} dph_by_worker = {} # cache so we don't re-fetch dph on every poll - abandoned_count = 0 # bail when this hits _MAX_ABANDONS; autoscaler is renting + giving up repeatedly, no progress + abandoned_count = 0 # bail when this hits MAX_ABANDONS; autoscaler is renting + giving up repeatedly, no progress while time.monotonic() - start < timeout: if stop_event is not None and stop_event.is_set(): - _update_row(class_states, row_id, status="aborted") + update_row(class_states, row_id, status="aborted") return (gpu_name, num_gpus, "aborted", None, "user aborted", None) - # calls get endpoint workers which includes measured perf(get workergroup workers doesnt) + # calls get endpoint workers which includes measured perf(get workergroup workers doesnt) resp = status_vast.get_endpoint_workers(endpoint_id) if isinstance(resp, list): workers = resp @@ -490,9 +542,9 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, workers = resp["workers"] else: workers = [] - abandoned_count += _update_worker_states(worker_states, workers, gpu_name) - if abandoned_count >= _MAX_ABANDONS: - _update_row(class_states, row_id, status="failed") + abandoned_count += update_worker_states(worker_states, workers, gpu_name) + if abandoned_count >= MAX_ABANDONS: + update_row(class_states, row_id, status="failed") return (gpu_name, num_gpus, "failed", None, f"autoscaler abandoned {abandoned_count} workers without " f"any reaching idle (likely template + GPU config not loadable " @@ -516,7 +568,7 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, file=sys.stderr) if primary_id and dph_by_worker.get(primary_id) is not None: fields["dph"] = dph_by_worker[primary_id] - _update_row(class_states, row_id, **fields) + update_row(class_states, row_id, **fields) # measured_perf is only real once status==idle; before that it's a dlperf placeholder. ready = [w for w in workers if str(w.get("status", "")).lower() == "idle" @@ -525,46 +577,46 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, # dph was already fetched and cached when this worker_id first appeared in polling. worker_id = ready[0].get("id") dph = dph_by_worker.get(worker_id) if worker_id else None - _update_row(class_states, row_id, status="done", + update_row(class_states, row_id, status="done", perf=ready[0]["measured_perf"], dph=dph) return (gpu_name, num_gpus, "ok", ready[0]["measured_perf"], None, dph) # Fail fast: every worker stuck terminal past the grace period, none reached idle. now_ts = time.monotonic() terminal_workers_long = workers and all( - str(w.get("status", "")).lower() in _TERMINAL_STATES + str(w.get("status", "")).lower() in TERMINAL_STATES and isinstance(w.get("id"), int) and w["id"] in worker_states and now_ts - worker_states[w["id"]]["state_started"] - > _TERMINAL_GRACE_SECONDS + > TERMINAL_GRACE_SECONDS for w in workers ) if terminal_workers_long: states = sorted({str(w.get("status", "")).lower() for w in workers}) - _update_row(class_states, row_id, status="failed") + update_row(class_states, row_id, status="failed") return (gpu_name, num_gpus, "failed", None, f"all workers terminal ({', '.join(states)}) for " - f">{_TERMINAL_GRACE_SECONDS}s without reaching idle; " + f">{TERMINAL_GRACE_SECONDS}s without reaching idle; " f"autoscaler not rotating", None) # Fail fast if no real worker (int id) has ever appeared. if (not any(isinstance(k, int) for k in worker_states) - and time.monotonic() - start > _RENTAL_TIMEOUT): - _update_row(class_states, row_id, status="no_worker") + and time.monotonic() - start > RENTAL_TIMEOUT): + update_row(class_states, row_id, status="no_worker") return (gpu_name, num_gpus, "no_worker", None, - f"autoscaler did not rent in {_RENTAL_TIMEOUT}s " + f"autoscaler did not rent in {RENTAL_TIMEOUT}s " f"(possible causes: insufficient credit, scoring issue, " f"all candidates failed silently, or template+GPU mismatch " f"missed by pre-flight)", None) # Jitter so N parallel threads don't all hit the autoscaler at the same instant. - delay = _WORKER_POLL_INTERVAL + random.uniform(0, 2.0) + delay = WORKER_POLL_INTERVAL + random.uniform(0, 2.0) if stop_event is not None: if stop_event.wait(timeout=delay): - _update_row(class_states, row_id, status="aborted") + update_row(class_states, row_id, status="aborted") return (gpu_name, num_gpus, "aborted", None, "user aborted", None) else: time.sleep(delay) - _update_row(class_states, row_id, status="timeout") + update_row(class_states, row_id, status="timeout") return (gpu_name, num_gpus, "timeout", None, f"no measured_perf in {timeout}s", None) finally: @@ -594,15 +646,17 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, @parser.command( argument("--template_hash", type=str, default=None, - help="(required, one of --template_hash or --template_id) template hash to benchmark"), + help="the template to benchmark; provide either --template_hash or --template_id"), argument("--template_id", type=int, default=None, - help="(required, one of --template_hash or --template_id) template id"), + help="the template to benchmark; provide either --template_hash or --template_id"), argument("--gpus", type=str, - help="comma-separated GPU names (e.g. RTX_4090,RTX_3090); optional Nx prefix takes precedence over --num_gpus (e.g. \"2x RTX_4090\")"), + help="comma-separated GPU names to benchmark (e.g. RTX_4090,RTX_3090). Prefix a name with a count to set GPUs per instance, e.g. \"2x RTX_4090\" (overrides --num_gpus)"), argument("--num_gpus", type=int, default=None, help="GPUs per instance for tokens without an Nx prefix (default 1); overridden by inline Nx in --gpus"), - argument("--timeout", type=int, default=_DEFAULT_BENCHMARK_TIMEOUT, - help=f"max seconds to wait for a benchmark before giving up (default {_DEFAULT_BENCHMARK_TIMEOUT})"), + argument("--timeout", type=int, default=DEFAULT_BENCHMARK_TIMEOUT, + help=f"max seconds to wait for a benchmark before giving up (default {DEFAULT_BENCHMARK_TIMEOUT})"), + argument("--no-cache", dest="no_cache", action="store_true", + help="re-measure instead of reusing recent cached benchmark results"), argument("-y", "--yes", action="store_true", help="Skip confirmation prompt"), argument("--auto_instance", type=str, default=None, help=argparse.SUPPRESS), @@ -613,6 +667,11 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, Rents one instance per GPU in parallel, measures perf, tears down. Each rental runs for up to --timeout seconds and costs real money. + Specs benchmarked recently (same template, GPU, and count, by any + user) are served from the benchmarks table as an estimated perf, and + you're asked whether to reuse that or run a fresh benchmark. Pass + --no-cache to always re-measure, or -y to always reuse the cache. + Examples: # auto-sweep the default GPUs against TGI vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5 @@ -626,8 +685,11 @@ def _benchmark_gpu(vast, *, gpu_name, num_gpus, timeout, # default count for tokens without an Nx prefix vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5 --gpus RTX_4090,RTX_3090 --num_gpus 2 - # shorter timeout (30 min), skipping the cost prompt - vastai run benchmarks --template_hash 393fa8572e6c73c927c8275fe4dffd53 --timeout 1800 -y + # shorter timeout (10 min), skipping the cost prompt + vastai run benchmarks --template_hash 393fa8572e6c73c927c8275fe4dffd53 --timeout 600 -y + + # ignore cached results and re-measure everything + vastai run benchmarks --template_hash 79ebdd2ebfb9d42cedf7a221c42d37a5 --no-cache # raw JSON output for piping into another tool vastai run benchmarks --template_hash 40ef49becc953aa910ee05bd4653b9b3 --raw @@ -639,18 +701,19 @@ def run__benchmarks(args): "(run `vastai run benchmarks --help` for usage)", file=sys.stderr) return 1 - # bump retry budget so parallel polls don't exhaust VastClient's tiny default (3 attempts, ~0.7s) under autoscaler rate limits. - retry = max(args.retry, 8) - vast = VastAI(api_key=args.api_key, server_url=args.url, retry=retry, - explain=getattr(args, 'explain', False), - curl=getattr(args, 'curl', False)) - if args.template_id is not None: query = {"id": {"eq": args.template_id}} ident = f"id={args.template_id}" else: query = {"hash_id": {"eq": args.template_hash}} ident = f"hash={args.template_hash}" + + # bump retry budget so parallel polls don't exhaust VastClient's tiny default (3 attempts, ~0.7s) under autoscaler rate limits. + retry = max(args.retry, 8) + vast = VastAI(api_key=args.api_key, server_url=args.url, retry=retry, + explain=getattr(args, 'explain', False), + curl=getattr(args, 'curl', False)) + templates = vast.search_templates(query=query) if not templates: print(f"error: template not found ({ident})", file=sys.stderr) @@ -662,24 +725,24 @@ def run__benchmarks(args): extra_filters = {} template_name = template.get("name") or template.get("title") or "?" - filter_summary = _format_filter_query(extra_filters) or "none" + filter_summary = format_filter_query(extra_filters) or "none" print(f"\nTemplate id={template.get('id')} {template_name}, filters: {filter_summary}") # if user provides specific GPU count like "2x RTX 4090", that takes precedence over --num_gpus if args.gpus: - gpu_specs = [_parse_gpu_spec(t, args.num_gpus or 1) + gpu_specs = [parse_gpu_spec(t, args.num_gpus or 1) for t in args.gpus.split(",") if t.strip()] elif args.num_gpus is not None: - gpu_specs = [(g, args.num_gpus) for g in _DEFAULT_GPUS] + gpu_specs = [(g, args.num_gpus) for g in DEFAULT_GPUS] else: gpu_specs = [ - (g, _pick_num_gpus(g, extra_filters)) - for g in _DEFAULT_GPUS + (g, pick_num_gpus(g, extra_filters)) + for g in DEFAULT_GPUS ] console = Console(stderr=True) - # deduplicates (gpu, num_gpus) so we dont benchmark the same GPU config twice + # deduplicates (gpu, num_gpus) so we dont benchmark the same GPU config twice seen = set() deduped = [] for g, n in gpu_specs: @@ -691,16 +754,54 @@ def run__benchmarks(args): deduped.append((g, n)) gpu_specs = deduped - # Pre-flight: skip GPU specs that have 0 matching offers before prompting, - # so the user sees skip reasons before approving the rentals. + # Pre-flight: for each spec, serve a recent cached benchmark if one exists + # (unless --no-cache), else skip specs with 0 matching offers, so the user + # sees cache hits and skip reasons before approving any rentals. compatible_specs = [] skipped_results = [] + cached_results = [] + interactive = not (args.yes or args.raw) for g, n in gpu_specs: - if not _has_matching_offer( + if not args.no_cache: + hit = lookup_cached_benchmark( + vast, gpu_name=g, num_gpus=n, + template_hash=template.get("hash_id"), + template_id=template.get("id"), + max_age_days=DEFAULT_CACHE_MAX_AGE_DAYS, + ) + if hit: + # Perf only: a cached row's $/hr would come from a different machine than the one benchmarked. + rentable = has_matching_offer( + vast, gpu_name=g, num_gpus=n, extra_filters=extra_filters) + age = (f"{hit['age_days']:.0f}d" if hit["age_days"] >= 1 + else "<1d") + note = None if rentable else "no offers available to rent right now" + msg = (f"[green][{n}x {g}] cached:[/green] median perf " + f"{hit['median']:.1f} " + f"(n={hit['n']}, range {hit['low']:.1f}-{hit['high']:.1f}, " + f"newest {age} ago)") + if note: + msg += f"; [yellow]{note}[/yellow]" + console.print(msg, highlight=False) + # Default is to reuse the cached result; only an interactive user + # who opts in pays to re-measure (and only when offers exist). + run_fresh = ( + interactive and rentable + and input( + f" [{n}x {g}] run a fresh benchmark anyway? this rents a " + f"real GPU and charges your account [y/N] " + ).strip().lower() in ("y", "yes") + ) + if run_fresh: + compatible_specs.append((g, n)) # offers already confirmed + else: + cached_results.append((g, n, "cached", hit["median"], note, None)) + continue + if not has_matching_offer( vast, gpu_name=g, num_gpus=n, extra_filters=extra_filters, ): - msg = _format_skip_message( + msg = format_skip_message( vast, gpu_name=g, num_gpus=n, extra_filters=extra_filters, ) @@ -710,13 +811,27 @@ def run__benchmarks(args): else: compatible_specs.append((g, n)) + # Live-table state. Pre-populate every GPU so the table is complete. + class_states = {} + for g, n in compatible_specs: + update_row(class_states, f"{n}x {g}", status="queued") + for sr in skipped_results: + update_row(class_states, f"{sr[1]}x {sr[0]}", status="skipped") + for cr in cached_results: + update_row(class_states, f"{cr[1]}x {cr[0]}", status="cached", + perf=cr[3], dph=cr[5]) + timeout_minutes = args.timeout / 60.0 n = len(compatible_specs) if n == 0: - console.print( - "\nNo compatible GPUs to benchmark for this template.", - style="bold red") - return _print_results(args, skipped_results) + if cached_results: + console.print() + console.print(render_table(class_states)) + else: + console.print( + "\nNo compatible GPUs to benchmark for this template.", + style="bold red") + return print_results(args, skipped_results + cached_results) spec_strs = [f"{cnt}x {gpu}" for gpu, cnt in compatible_specs] spec_summary = ", ".join(spec_strs) @@ -766,16 +881,9 @@ def _cleanup(): pass atexit.register(_cleanup) - # Live-table state. Pre-populate every GPU so the table is complete. - class_states = {} - for g, n in compatible_specs: - _update_row(class_states, f"{n}x {g}", status="queued") - for sr in skipped_results: - _update_row(class_states, f"{sr[1]}x {sr[0]}", status="skipped") - def _run_one_gpu(g, n): try: - return _benchmark_gpu( + return benchmark_gpu( vast, gpu_name=g, num_gpus=n, @@ -790,7 +898,7 @@ def _run_one_gpu(g, n): stop_event=stop_event, ) except Exception as e: - _update_row(class_states, f"{n}x {g}", status="error") + update_row(class_states, f"{n}x {g}", status="error") # Surface response body for HTTPError so backend rate-limit / validation # messages aren't swallowed (default __str__ on HTTPError just shows the URL). body = "" @@ -810,9 +918,9 @@ def _run_one_gpu(g, n): for i, (g, n_for_class) in enumerate(compatible_specs): futures.append(executor.submit(_run_one_gpu, g, n_for_class)) if i < len(compatible_specs) - 1: - time.sleep(_PARALLEL_SUBMIT_DELAY) + time.sleep(PARALLEL_SUBMIT_DELAY) try: - with Live(_render_table(class_states), console=console, + with Live(render_table(class_states), console=console, refresh_per_second=2, transient=False) as live: pending = set(futures) while pending: @@ -827,7 +935,7 @@ def _run_one_gpu(g, n): console.print( f"[red][{n}x {g}] {status}:[/red] {err}", highlight=False) - live.update(_render_table(class_states)) + live.update(render_table(class_states)) except KeyboardInterrupt: # Set stop_event so threads exit their poll/sleep immediately, then cleanup synchronously. console.print( @@ -842,11 +950,11 @@ def _run_one_gpu(g, n): executor.shutdown(wait=True) _cleanup() - results = skipped_results + run_results - return _print_results(args, results) + results = skipped_results + cached_results + run_results + return print_results(args, results) -def _print_results(args, results): +def print_results(args, results): rows = [] for gpu, num_gpus, status, perf, err, price in results: pps = (perf / price) if (perf and price) else None @@ -864,4 +972,6 @@ def _print_results(args, results): return rows n_ok = sum(1 for r in rows if r["status"] == "ok") - print(f"\nBenchmark complete: {n_ok}/{len(rows)} GPUs measured.") + n_cached = sum(1 for r in rows if r["status"] == "cached") + line = f"\nBenchmark complete: {n_ok + n_cached}/{len(rows)} GPUs measured" + print(line + (f" ({n_cached} from cache)." if n_cached else ".")) diff --git a/vastai/sdk.py b/vastai/sdk.py index 1f72579e..abf35be5 100644 --- a/vastai/sdk.py +++ b/vastai/sdk.py @@ -3,7 +3,7 @@ from __future__ import annotations import warnings -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Union from vastai._base import _resolve_api_key, _APIKEY_SENTINEL from vastai.api.client import VastClient @@ -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[str] = 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."""