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

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