popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix Windows memory budget detection

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-01 10:49:06 +02:00
parent 278be49539
commit d778b23e1e
3 changed files with 219 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
packages/node/meshnet_node/hardware.py
View File

@@ -1,5 +1,6 @@
"""GPU hardware detection with graceful CPU fallback."""
import json
import os
import subprocess
import time
@@ -12,7 +13,96 @@ def _detect_ram_mb() -> int:
page_size = os.sysconf("SC_PAGE_SIZE")
return int((pages * page_size) // (1024 * 1024))
except (AttributeError, OSError, ValueError):
return 0
pass
return _detect_windows_ram_mb()
def _detect_windows_ram_mb() -> int:
"""Return Windows physical RAM in MB, or 0."""
try:
import ctypes
class _MemoryStatusEx(ctypes.Structure):
_fields_ = [
("dwLength", ctypes.c_ulong),
("dwMemoryLoad", ctypes.c_ulong),
("ullTotalPhys", ctypes.c_ulonglong),
("ullAvailPhys", ctypes.c_ulonglong),
("ullTotalPageFile", ctypes.c_ulonglong),
("ullAvailPageFile", ctypes.c_ulonglong),
("ullTotalVirtual", ctypes.c_ulonglong),
("ullAvailVirtual", ctypes.c_ulonglong),
("ullAvailExtendedVirtual", ctypes.c_ulonglong),
]
status = _MemoryStatusEx()
status.dwLength = ctypes.sizeof(_MemoryStatusEx)
if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
return int(status.ullTotalPhys // (1024 * 1024))
except (AttributeError, OSError, ValueError):
pass
try:
result = subprocess.run(
[
"powershell",
"-NoProfile",
"-Command",
"(Get-CimInstance Win32_ComputerSystem).TotalPhysicalMemory",
],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode == 0 and result.stdout.strip():
return int(result.stdout.strip()) // (1024 * 1024)
except (FileNotFoundError, subprocess.TimeoutExpired, ValueError):
pass
return 0
def _detect_windows_gpu_memory() -> dict | None:
"""Return Windows GPU memory metadata from Win32_VideoController, if available."""
try:
result = subprocess.run(
[
"powershell",
"-NoProfile",
"-Command",
(
"Get-CimInstance Win32_VideoController | "
"Select-Object Name,AdapterRAM | ConvertTo-Json -Compress"
),
],
capture_output=True,
text=True,
timeout=5,
)
except (FileNotFoundError, subprocess.TimeoutExpired):
return None
if result.returncode != 0 or not result.stdout.strip():
return None
try:
raw = json.loads(result.stdout)
except json.JSONDecodeError:
return None
entries = raw if isinstance(raw, list) else [raw]
best: dict | None = None
for entry in entries:
if not isinstance(entry, dict):
continue
name = str(entry.get("Name") or "").strip()
if not name:
continue
try:
adapter_ram = int(entry.get("AdapterRAM") or 0)
except (TypeError, ValueError):
adapter_ram = 0
vram_mb = max(0, adapter_ram // (1024 * 1024))
if best is None or vram_mb > best["vram_mb"]:
best = {"gpu_name": name, "vram_mb": vram_mb}
return best
def detect_hardware() -> dict:
@@ -25,7 +115,15 @@ def detect_hardware() -> dict:
name = torch.cuda.get_device_name(idx)
props = torch.cuda.get_device_properties(idx)
vram_mb = props.total_memory // (1024 * 1024)
return {"device": "cuda", "gpu_name": name, "vram_mb": vram_mb, "ram_mb": ram_mb}
shared_vram_mb = max(0, ram_mb // 2)
return {
"device": "cuda",
"gpu_name": name,
"vram_mb": vram_mb,
"dedicated_vram_mb": vram_mb,
"shared_vram_mb": shared_vram_mb,
"ram_mb": ram_mb,
}
except ImportError:
pass
@@ -39,11 +137,37 @@ def detect_hardware() -> dict:
parts = line.split(",", 1)
gpu_name = parts[0].strip()
vram_mb = int(parts[1].strip()) if len(parts) > 1 else 0
return {"device": "cuda", "gpu_name": gpu_name, "vram_mb": vram_mb, "ram_mb": ram_mb}
shared_vram_mb = max(0, ram_mb // 2)
return {
"device": "cuda",
"gpu_name": gpu_name,
"vram_mb": vram_mb,
"dedicated_vram_mb": vram_mb,
"shared_vram_mb": shared_vram_mb,
"ram_mb": ram_mb,
}
except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError):
pass
return {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": ram_mb}
windows_gpu = _detect_windows_gpu_memory()
if windows_gpu is not None:
return {
"device": "cpu",
"gpu_name": windows_gpu["gpu_name"],
"vram_mb": windows_gpu["vram_mb"],
"dedicated_vram_mb": windows_gpu["vram_mb"],
"shared_vram_mb": max(0, ram_mb // 2),
"ram_mb": ram_mb,
}
return {
"device": "cpu",
"gpu_name": None,
"vram_mb": 0,
"dedicated_vram_mb": 0,
"shared_vram_mb": 0,
"ram_mb": ram_mb,
}
def benchmark_throughput(device_str: str = "cpu") -> float:

29
packages/node/meshnet_node/startup.py
View File

@@ -24,9 +24,11 @@ from .wallet import load_or_create_wallet
_DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024
def _memory_budget(vram_mb: int, ram_mb: int) -> tuple[int, str]:
def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int = 0) -> tuple[int, str]:
"""Return the capacity budget in MB and whether it came from VRAM or RAM."""
if vram_mb > 0:
if device == "cuda" and vram_mb > 0:
if shared_vram_mb > 0:
return vram_mb + shared_vram_mb, "VRAM + shared RAM"
return vram_mb, "VRAM"
return max(0, ram_mb), "RAM"
@@ -348,17 +350,28 @@ def run_startup(
device: str = hw["device"]
gpu_name: str | None = hw.get("gpu_name")
vram_mb: int = hw.get("vram_mb", 0)
shared_vram_mb: int = hw.get("shared_vram_mb", 0)
ram_mb: int = hw.get("ram_mb", 16 * 1024)
if vram_mb_override is not None:
vram_mb = vram_mb_override
shared_vram_mb = 0
print(f" Memory budget overridden to {vram_mb / 1024:.1f} GB via --memory", flush=True)
elif device == "cpu":
print(f" WARNING: No CUDA GPU detected — running in CPU mode ({ram_mb / 1024:.1f} GB RAM)", flush=True)
gpu_suffix = ""
if gpu_name and vram_mb > 0:
gpu_suffix = f"; detected {gpu_name} ({vram_mb / 1024:.1f} GB dedicated VRAM, {shared_vram_mb / 1024:.1f} GB shared)"
print(f" WARNING: No CUDA GPU detected — running in CPU mode ({ram_mb / 1024:.1f} GB RAM{gpu_suffix})", flush=True)
else:
print(f" GPU: {gpu_name} ({vram_mb / 1024:.1f} GB VRAM, {ram_mb / 1024:.1f} GB RAM)", flush=True)
shared_suffix = f", {shared_vram_mb / 1024:.1f} GB shared" if shared_vram_mb > 0 else ""
print(f" GPU: {gpu_name} ({vram_mb / 1024:.1f} GB dedicated VRAM{shared_suffix}, {ram_mb / 1024:.1f} GB RAM)", flush=True)
memory_budget_mb, memory_budget_source = _memory_budget(vram_mb, ram_mb)
if vram_mb_override is not None:
memory_budget_mb = vram_mb
memory_budget_source = "memory override"
else:
memory_budget_mb, memory_budget_source = _memory_budget(device, vram_mb, ram_mb, shared_vram_mb)
assignment_vram_mb = memory_budget_mb if device == "cuda" or vram_mb_override is not None else 0
print(f" Memory budget: {memory_budget_mb / 1024:.1f} GB {memory_budget_source}", flush=True)
print("Benchmarking compute...", flush=True)
@@ -367,7 +380,7 @@ def run_startup(
print(f" {device_label} throughput index: {bench_tps:,.0f}", flush=True)
registration_capabilities = {
"vram_bytes": max(0, int(vram_mb)) * 1024 * 1024,
"vram_bytes": max(0, int(assignment_vram_mb)) * 1024 * 1024,
"ram_bytes": max(0, int(ram_mb)) * 1024 * 1024,
"max_loaded_shards": max_loaded_shards,
"benchmark_tokens_per_sec": bench_tps,
@@ -397,7 +410,7 @@ def run_startup(
if shard_start is None and shard_end is None:
try:
qs = urllib.parse.urlencode({
"device": device, "vram_mb": vram_mb, "ram_mb": ram_mb, "hf_repo": model_id,
"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb, "hf_repo": model_id,
})
net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
@@ -495,7 +508,7 @@ def run_startup(
# 3a. Auto-join: query tracker for network-wide HF model assignment.
print("Querying tracker for network assignment...", flush=True)
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb, "ram_mb": ram_mb})
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
net_assignment: dict = {}
try:
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")

72
tests/test_node_startup.py
View File

@@ -13,6 +13,7 @@ from meshnet_node.downloader import download_shard, write_shard_archive
from meshnet_node.hardware import detect_hardware, benchmark_throughput
from meshnet_node.startup import (
_infer_relay_url_from_tracker,
_memory_budget,
_probationary_status_line,
run_startup,
)
@@ -34,13 +35,80 @@ def test_detect_hardware_returns_valid_profile():
assert isinstance(hw.get("ram_mb"), int)
assert hw["ram_mb"] > 0
if hw["device"] == "cpu":
assert hw["gpu_name"] is None
assert hw["vram_mb"] == 0
assert hw["gpu_name"] is None or isinstance(hw["gpu_name"], str)
assert hw["vram_mb"] >= 0
else:
assert isinstance(hw["gpu_name"], str) and hw["gpu_name"]
assert hw["vram_mb"] > 0
def test_windows_ram_fallback_is_used_when_sysconf_is_unavailable(monkeypatch):
"""Windows hosts do not have os.sysconf; RAM must not collapse to 0 MB."""
import meshnet_node.hardware as hardware_mod
monkeypatch.setattr(
hardware_mod.os,
"sysconf",
lambda _name: (_ for _ in ()).throw(AttributeError()),
raising=False,
)
monkeypatch.setattr(hardware_mod, "_detect_windows_ram_mb", lambda: 64 * 1024)
assert hardware_mod._detect_ram_mb() == 64 * 1024
def test_windows_gpu_memory_fallback_preserves_cpu_execution(monkeypatch):
"""A Windows-visible GPU is reported, but CUDA execution is not claimed without CUDA."""
import meshnet_node.hardware as hardware_mod
calls = []
class FakeResult:
def __init__(self, stdout):
self.returncode = 0
self.stdout = stdout
def fake_run(command, *args, **kwargs):
calls.append(command)
joined = " ".join(command)
if "nvidia-smi" in joined:
raise FileNotFoundError
if "Win32_ComputerSystem" in joined:
return FakeResult(str(80 * 1024 * 1024 * 1024))
if "Win32_VideoController" in joined:
return FakeResult('{"Name":"NVIDIA GeForce RTX Laptop GPU","AdapterRAM":8589934592}')
raise AssertionError(command)
monkeypatch.setattr(
hardware_mod.os,
"sysconf",
lambda _name: (_ for _ in ()).throw(AttributeError()),
raising=False,
)
monkeypatch.setattr(hardware_mod.subprocess, "run", fake_run)
monkeypatch.setattr(hardware_mod, "_detect_windows_ram_mb", lambda: 80 * 1024)
monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(cuda=types.SimpleNamespace(is_available=lambda: False)))
hw = hardware_mod.detect_hardware()
assert hw["device"] == "cpu"
assert hw["gpu_name"] == "NVIDIA GeForce RTX Laptop GPU"
assert hw["vram_mb"] == 8192
assert hw["shared_vram_mb"] == 40 * 1024
assert hw["ram_mb"] == 80 * 1024
def test_memory_budget_uses_ram_for_cpu_and_shared_memory_for_cuda():
assert _memory_budget("cpu", vram_mb=8192, ram_mb=80 * 1024, shared_vram_mb=40 * 1024) == (
80 * 1024,
"RAM",
)
assert _memory_budget("cuda", vram_mb=8192, ram_mb=80 * 1024, shared_vram_mb=40 * 1024) == (
48 * 1024,
"VRAM + shared RAM",
)
def test_benchmark_throughput_cpu_returns_positive():
"""CPU benchmark returns a positive float greater than the 1.0 error fallback."""
result = benchmark_throughput("cpu")