diff --git a/packages/node/meshnet_node/hardware.py b/packages/node/meshnet_node/hardware.py index 4f253fc..d929b4b 100644 --- a/packages/node/meshnet_node/hardware.py +++ b/packages/node/meshnet_node/hardware.py @@ -192,7 +192,7 @@ def detect_hardware() -> dict: } -def benchmark_throughput(device_str: str = "cpu") -> float: +def benchmark_throughput_checked(device_str: str = "cpu") -> tuple[float, bool, str | None]: """ Estimate compute throughput via a synthetic transformer GEMM benchmark. @@ -201,7 +201,8 @@ def benchmark_throughput(device_str: str = "cpu") -> float: The value is used as benchmark_tokens_per_sec in tracker registration for routing tiebreaks; it is not an absolute token rate. - Falls back to 1.0 if torch is unavailable. + Returns (score, ok, error). Score falls back to 1.0 when the requested + device cannot run the benchmark. """ try: import torch # type: ignore[import] @@ -233,6 +234,12 @@ def benchmark_throughput(device_str: str = "cpu") -> float: _sync() elapsed = time.perf_counter() - t0 - return round(n_iters / max(elapsed, 1e-9), 2) - except Exception: - return 1.0 + return round(n_iters / max(elapsed, 1e-9), 2), True, None + except Exception as exc: + return 1.0, False, f"{type(exc).__name__}: {exc}" + + +def benchmark_throughput(device_str: str = "cpu") -> float: + """Return only the numeric throughput index, preserving the legacy API.""" + score, _ok, _error = benchmark_throughput_checked(device_str) + return score diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 37c3427..fba42ee 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -14,7 +14,7 @@ from pathlib import Path from typing import Any from .downloader import compute_shard_checksum, download_shard -from .hardware import detect_hardware, benchmark_throughput +from .hardware import detect_hardware, benchmark_throughput_checked from .relay_bridge import RelayHttpBridge, peer_id_from_wallet from .server import StubNodeServer from .torch_server import TorchNodeServer @@ -386,7 +386,18 @@ def run_startup( print(f" Memory budget: {memory_budget_mb / 1024:.1f} GB {memory_budget_source}", flush=True) print("Benchmarking compute...", flush=True) - bench_tps = benchmark_throughput(device) + if device != "cuda" and gpu_name: + _cuda_score, cuda_ok, cuda_error = benchmark_throughput_checked("cuda") + hw["cuda_benchmark_ok"] = cuda_ok + if cuda_error: + hw["cuda_benchmark_error"] = cuda_error + if not cuda_ok: + print(f" CUDA benchmark unavailable: {cuda_error}; using CPU benchmark", flush=True) + bench_tps, bench_ok, bench_error = benchmark_throughput_checked(device) + hw["benchmark_device"] = device + hw["benchmark_ok"] = bench_ok + if bench_error: + hw["benchmark_error"] = bench_error device_label = "GPU" if device == "cuda" else "CPU" print(f" {device_label} throughput index: {bench_tps:,.0f}", flush=True) diff --git a/tests/conftest.py b/tests/conftest.py index ee9c069..891d10e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,6 @@ def _stub_benchmark_throughput(monkeypatch): """ try: import meshnet_node.startup as startup_mod - monkeypatch.setattr(startup_mod, "benchmark_throughput", lambda _device: 999.0) + monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", lambda _device: (999.0, True, None)) except ImportError: pass diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 4070a11..00a4b19 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -173,7 +173,7 @@ def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path): monkeypatch.setattr(startup_mod, "detect_hardware", lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384}) - monkeypatch.setattr(startup_mod, "benchmark_throughput", lambda _device: 42.5) + monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", lambda _device: (42.5, True, None)) monkeypatch.setattr(startup_mod, "TorchNodeServer", lambda **_kw: FakeNode()) monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24) monkeypatch.setattr(startup_mod, "RelayHttpBridge", None) @@ -193,6 +193,67 @@ def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path): node.stop() assert captured.get("benchmark_tokens_per_sec") == 42.5 + assert captured["hardware_profile"]["benchmark_device"] == "cpu" + assert captured["hardware_profile"]["benchmark_ok"] is True + + +def test_cuda_benchmark_failure_is_registered_for_inventory_only_gpu(monkeypatch, tmp_path, capsys): + import meshnet_node.startup as startup_mod + + captured: dict = {} + + class FakeNode: + backend = None + + def start(self): + return 7099 + + def stop(self): + pass + + def fake_benchmark(device): + if device == "cuda": + return 1.0, False, "AssertionError: Torch not compiled with CUDA enabled" + return 55.0, True, None + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: { + "device": "cpu", + "gpu_name": "NVIDIA GeForce RTX 4060 Laptop GPU", + "vram_mb": 8188, + "dedicated_vram_mb": 8188, + "shared_vram_mb": 40555, + "ram_mb": 81111, + "cuda_available": False, + }, + ) + monkeypatch.setattr(startup_mod, "benchmark_throughput_checked", fake_benchmark) + monkeypatch.setattr(startup_mod, "TorchNodeServer", lambda **_kw: FakeNode()) + monkeypatch.setattr(startup_mod, "_post_json", + lambda _url, payload, timeout=10.0: (captured.update(payload) or {"node_id": "x"})) + monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None) + + node = run_startup( + tracker_url="http://localhost:8080", + model_id="Qwen/Qwen2.5-0.5B-Instruct", + shard_start=0, + shard_end=23, + wallet_path=tmp_path / "wallet.json", + ) + node.stop() + + output = capsys.readouterr().out + assert "CUDA benchmark unavailable" in output + assert "Hardware: CPU (CUDA inactive)" in output + hw = captured["hardware_profile"] + assert hw["cuda_benchmark_ok"] is False + assert "Torch not compiled with CUDA enabled" in hw["cuda_benchmark_error"] + assert hw["benchmark_device"] == "cpu" + assert hw["benchmark_ok"] is True + assert captured["ram_bytes"] == 81111 * 1024 * 1024 + assert captured["vram_bytes"] == 0 def test_wallet_generates_new_keypair(tmp_path):