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

-cpu flag

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-09 08:19:15 +03:00
parent 4ed585bf54
commit 9ec4ca9ce1
8 changed files with 121 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
packages/node/meshnet_node/cli.py
View File

@@ -69,6 +69,7 @@ def _run_node(cfg: dict) -> None:
torch_threads=cfg.get("torch_threads"),
torch_interop_threads=cfg.get("torch_interop_threads"),
node_name=cfg.get("node_name"),
force_cpu=bool(cfg.get("force_cpu", False)),
)
except Exception as exc:
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
@@ -174,6 +175,8 @@ def _cmd_default(args) -> int:
overrides["torch_threads"] = args.torch_threads
if getattr(args, "torch_interop_threads", None) is not None:
overrides["torch_interop_threads"] = args.torch_interop_threads
if getattr(args, "cpu", False):
overrides["force_cpu"] = True
if overrides:
cfg = merge_cli_overrides(cfg, **overrides)
@@ -276,6 +279,7 @@ def _cmd_start(args) -> int:
torch_threads=getattr(args, "torch_threads", None),
torch_interop_threads=getattr(args, "torch_interop_threads", None),
node_name=cfg.get("node_name"),
force_cpu=getattr(args, "cpu", False),
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
@@ -332,6 +336,8 @@ def main() -> None:
help="Set PyTorch intra-op CPU worker threads")
parser.add_argument("--torch-interop-threads", type=int, metavar="N",
help="Set PyTorch inter-op CPU worker threads")
parser.add_argument("--cpu", action="store_true",
help="Force CPU inference even when a GPU is available")
parser.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
parser.add_argument("--compact", action="store_true", help="Single-line status output")
@@ -372,6 +378,8 @@ def main() -> None:
help="Set PyTorch intra-op CPU worker threads")
start_cmd.add_argument("--torch-interop-threads", type=int, metavar="N",
help="Set PyTorch inter-op CPU worker threads")
start_cmd.add_argument("--cpu", action="store_true",
help="Force CPU inference even when a GPU is available")
start_cmd.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
start_cmd.add_argument("--tracker-source-disabled", action="store_true",
help="Skip tracker/peer model-file sources and download from HuggingFace directly")

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

@@ -171,6 +171,18 @@ def _gpu_inventory_profile(gpu: dict | None, ram_mb: int) -> dict | None:
return profile
def with_forced_cpu(hw: dict) -> dict:
"""Return a hardware profile forced to CPU execution.
Keeps detected GPU metadata for diagnostics and tracker registration context,
but clears CUDA availability so startup and the model backend stay on CPU.
"""
forced = dict(hw)
forced["device"] = "cpu"
forced["cuda_available"] = False
return forced
def detect_hardware() -> dict:
"""Detect GPU model and available VRAM. Returns hardware profile dict."""
ram_mb = _detect_ram_mb()

11
packages/node/meshnet_node/model_backend.py
View File

@@ -212,6 +212,7 @@ class TorchModelShard:
shard_end: int,
quantization: Quantization = "auto",
cache_dir: Path | None = None,
force_cpu: bool = False,
) -> None:
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
raise ValueError("shard_start must be <= shard_end and non-negative")
@@ -229,7 +230,10 @@ class TorchModelShard:
) from exc
self.torch = torch
self.device = torch.device("cuda" if _torch_cuda_is_executable(torch) else "cpu")
if force_cpu:
self.device = torch.device("cpu")
else:
self.device = torch.device("cuda" if _torch_cuda_is_executable(torch) else "cpu")
load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
quant_config, dtype, uses_quantized_weights = _model_load_plan(
AutoConfig,
@@ -682,8 +686,11 @@ def load_torch_shard(
shard_end: int,
quantization: Quantization = "auto",
cache_dir: Path | None = None,
force_cpu: bool = False,
) -> TorchModelShard:
return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
return TorchModelShard(
model_id, shard_start, shard_end, quantization, cache_dir, force_cpu=force_cpu
)
def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None:

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

@@ -15,7 +15,7 @@ from pathlib import Path
from typing import Any
from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware, benchmark_throughput_checked
from .hardware import detect_hardware, benchmark_throughput_checked, with_forced_cpu
from .model_catalog import model_metadata_for
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
from .server import StubNodeServer
@@ -640,6 +640,7 @@ def run_startup(
torch_threads: int | None = None,
torch_interop_threads: int | None = None,
node_name: str | None = None,
force_cpu: bool = False,
) -> StubNodeServer | TorchNodeServer:
"""Execute the full startup sequence and return a running node server.
@@ -681,6 +682,8 @@ def run_startup(
print("Detecting hardware...", flush=True)
hw = detect_hardware()
if force_cpu:
hw = with_forced_cpu(hw)
torch_thread_config = _configure_torch_threads(torch_threads, torch_interop_threads)
if torch_thread_config:
hw.update(torch_thread_config)
@@ -697,6 +700,16 @@ def run_startup(
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 force_cpu:
if gpu_name and vram_mb > 0:
print(
f" --cpu: ignoring {gpu_name} "
f"({vram_mb / 1024:.1f} GB VRAM); running in CPU mode "
f"({ram_mb / 1024:.1f} GB RAM)",
flush=True,
)
else:
print(f" --cpu: running in CPU mode ({ram_mb / 1024:.1f} GB RAM)", flush=True)
elif device == "cpu":
gpu_suffix = ""
if gpu_name and vram_mb > 0:
@@ -718,7 +731,7 @@ def run_startup(
print(f" Memory budget: {memory_budget_mb / 1024:.1f} GB {memory_budget_source}", flush=True)
print("Benchmarking compute...", flush=True)
if device != "cuda" and gpu_name:
if device != "cuda" and gpu_name and not force_cpu:
_cuda_score, cuda_ok, cuda_error = benchmark_throughput_checked("cuda")
hw["cuda_benchmark_ok"] = cuda_ok
if cuda_error:
@@ -833,6 +846,7 @@ def run_startup(
cache_dir=cache_dir,
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
)
_node_start_time = time.monotonic()
actual_port = node.start()
@@ -986,6 +1000,7 @@ def run_startup(
cache_dir=cache_dir,
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
)
_node_start_time = time.monotonic()
actual_port = node.start()
@@ -1163,6 +1178,7 @@ def run_startup(
cache_dir=shard_path,
debug=debug,
max_loaded_shards=max_loaded_shards,
force_cpu=force_cpu,
)
actual_port = node.start()
total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers

7
packages/node/meshnet_node/torch_server.py
View File

@@ -1087,6 +1087,7 @@ class TorchNodeServer:
cache_dir: Path | None = None,
debug: bool = False,
max_loaded_shards: int = 1,
force_cpu: bool = False,
) -> None:
self._host = host
self._requested_port = port
@@ -1097,6 +1098,7 @@ class TorchNodeServer:
shard_end,
quantization,
cache_dir,
force_cpu=force_cpu,
)
self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
@@ -1245,12 +1247,15 @@ def _load_backend(
shard_end: int,
quantization: str,
cache_dir: Path | None = None,
force_cpu: bool = False,
) -> TorchModelShard:
from .model_backend import load_torch_shard
quant = validate_quantization(quantization)
try:
return load_torch_shard(model_id, shard_start, shard_end, quant, cache_dir)
return load_torch_shard(
model_id, shard_start, shard_end, quant, cache_dir, force_cpu=force_cpu
)
except MissingModelDependencyError:
raise
except InsufficientVRAMError as exc: