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}")