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

inference working

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-29 23:54:35 +03:00
parent 607d49f5b0
commit 1bdfce657d
16 changed files with 1899 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,83 +1,257 @@
"""meshnet-node CLI entry point."""
"""meshnet-node CLI entry point — mining-style UX."""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
def _run_node(cfg: dict) -> None:
"""Start the node and hand off to the live dashboard. Blocks until Ctrl-C."""
from .startup import run_startup
from .dashboard import run_dashboard
start_time = time.monotonic()
try:
node = run_startup(
tracker_url=cfg["tracker_url"],
port=cfg.get("port", 7000),
model=cfg.get("model_name") or "stub-model",
model_id=cfg.get("model_hf_repo") or None,
shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"),
quantization=cfg.get("quantization", "bfloat16").replace("bf16", "bfloat16"),
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg.get("host", "0.0.0.0"),
)
except Exception as exc:
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
sys.exit(1)
try:
run_dashboard(node, cfg, start_time)
except KeyboardInterrupt:
pass
finally:
node.stop()
req = getattr(node, "chat_completion_count", 0)
elapsed = time.monotonic() - start_time
h, rem = divmod(int(elapsed), 3600)
m, s = divmod(rem, 60)
print(
f"\nmeshnet-node stopped. "
f"Served {req} requests in {h:02d}:{m:02d}:{s:02d}.",
flush=True,
)
def _cmd_default(args) -> int:
"""No subcommand: wizard if no config, else start with saved config."""
from .config import load_config, save_config, merge_cli_overrides
from .wizard import run_wizard
cfg = load_config()
if cfg is None or args.reset_config:
if args.reset_config and cfg is not None:
print("Resetting config — re-running setup wizard.\n")
try:
cfg = run_wizard()
except KeyboardInterrupt:
print("\nSetup cancelled.")
return 1
save_config(cfg)
print(f"\nConfig saved to ~/.config/meshnet/config.json\n")
# Apply CLI overrides on top of saved config
overrides: dict = {}
if args.model:
overrides["model_hf_repo"] = args.model
overrides["model_name"] = args.model.split("/")[-1]
if args.quantization:
overrides["quantization"] = args.quantization
if args.download_dir:
overrides["download_dir"] = args.download_dir
if args.tracker:
overrides["tracker_url"] = args.tracker
if args.wallet:
overrides["wallet_path"] = args.wallet
if args.shard_start is not None:
overrides["shard_start"] = args.shard_start
if args.shard_end is not None:
overrides["shard_end"] = args.shard_end
if args.port is not None:
overrides["port"] = args.port
if args.host:
overrides["host"] = args.host
if overrides:
cfg = merge_cli_overrides(cfg, **overrides)
_run_node(cfg)
return 0
def _cmd_models(args) -> int:
"""List curated models (with optional HF Hub browse)."""
from .wizard import print_models_table, _browse_hf_interactive
if args.browse:
from .model_catalog import browse_hf_hub
print("Fetching HuggingFace Hub top models...\n")
try:
models = browse_hf_hub(top_n=20)
print(f"{'#':<4} {'Repo':<60} {'Downloads':>12}")
print(f"{'─'*4} {'─'*60} {'─'*12}")
for i, m in enumerate(models, 1):
dl = m["downloads"]
dl_str = (
f"{dl/1e6:.1f}M" if dl >= 1_000_000
else f"{dl/1e3:.0f}k" if dl >= 1000
else str(dl)
)
print(f"{i:<4} {m['repo']:<60} {dl_str:>12}")
except RuntimeError as exc:
print(f"Error: {exc}", file=sys.stderr)
return 1
else:
print_models_table()
return 0
def _cmd_config(args) -> int:
"""Print current config."""
import json
from .config import load_config, config_path
cfg = load_config()
if cfg is None:
print("No config file found. Run `meshnet-node` to start setup.")
return 1
print(f"Config: {config_path()}")
print(json.dumps(cfg, indent=2))
return 0
def _cmd_start(args) -> int:
"""Legacy `start` subcommand — preserves backward compatibility with existing tests."""
from .config import load_config, DEFAULTS
# Build a transient config from flags (don't write to disk)
cfg = dict(DEFAULTS)
cfg["tracker_url"] = args.tracker
cfg["port"] = args.port
cfg["model_name"] = args.model
cfg["quantization"] = args.quantization
cfg["host"] = args.host
if args.model_id:
cfg["model_hf_repo"] = args.model_id
if args.shard_start is not None:
cfg["shard_start"] = args.shard_start
if args.shard_end is not None:
cfg["shard_end"] = args.shard_end
if args.wallet:
cfg["wallet_path"] = args.wallet
if args.download_dir:
cfg["download_dir"] = args.download_dir
# Legacy start: just run without the dashboard (keep original blocking loop)
from .startup import run_startup
try:
node = run_startup(
tracker_url=cfg["tracker_url"],
port=cfg["port"],
model=cfg["model_name"],
model_id=cfg.get("model_hf_repo"),
shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"),
quantization=cfg["quantization"].replace("bf16", "bfloat16"),
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg["host"],
advertise_host=getattr(args, "advertise_host", None),
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
sys.exit(1)
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
node.stop()
return 0
def main() -> None:
parser = argparse.ArgumentParser(
prog="meshnet-node",
description="Distributed Inference Network node client",
description="Distributed AI Inference — Node Client",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=(
"Run with no arguments to start the setup wizard.\n"
"After first setup, `meshnet-node` starts using your saved config.\n\n"
"Subcommands:\n"
" models List supported models\n"
" models --browse Browse HuggingFace Hub\n"
" config Show current config\n"
),
)
# Flags that apply to the no-subcommand (default) path
parser.add_argument("--model", metavar="HF_REPO", help="HuggingFace repo ID to serve")
parser.add_argument("--quantization", "-q", choices=["bf16", "int8", "nf4", "bfloat16"],
help="Quantization level")
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
parser.add_argument("--tracker", metavar="URL", help="Tracker URL")
parser.add_argument("--wallet", metavar="PATH", help="Wallet file path")
parser.add_argument("--shard-start", type=int, metavar="N", help="Pin shard start layer")
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer")
parser.add_argument("--port", type=int, metavar="N", help="Port to listen on")
parser.add_argument("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)")
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")
parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists")
subparsers = parser.add_subparsers(dest="command")
start_cmd = subparsers.add_parser("start", help="Start the node server")
start_cmd.add_argument(
"--tracker", default="http://localhost:8080", help="Tracker URL"
)
start_cmd.add_argument("--port", type=int, default=7000, help="Port to listen on")
start_cmd.add_argument(
"--model", default="stub-model", help="Model preset to request from tracker"
)
start_cmd.add_argument(
"--model-id",
help="HuggingFace model id for the real PyTorch backend",
)
start_cmd.add_argument("--shard-start", type=int, help="First layer index for an explicit shard")
start_cmd.add_argument("--shard-end", type=int, help="Exclusive layer end index for an explicit shard")
start_cmd.add_argument(
"--quantization",
choices=["bfloat16", "int8", "nf4"],
default="bfloat16",
help="Weight quantization for the real PyTorch backend",
)
start_cmd.add_argument(
"--host", default="0.0.0.0", help="Interface to bind to"
)
start_cmd.add_argument(
"--advertise-host",
help="Reachable host/IP to advertise to the tracker (defaults to FQDN when binding 0.0.0.0)",
)
start_cmd.add_argument(
"--tracker-mode",
action="store_true",
help="Enable client-facing /v1/chat/completions (auto-enabled when shard-start=0)",
)
start_cmd.add_argument(
"--tracker-url",
default=None,
help="Tracker URL for route selection (used in tracker mode)",
)
# models subcommand
models_cmd = subparsers.add_parser("models", help="List supported models")
models_cmd.add_argument("--browse", action="store_true", help="Browse HuggingFace Hub top-20")
# config subcommand
subparsers.add_parser("config", help="Show current saved config")
# start subcommand (legacy / backward-compat)
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
start_cmd.add_argument("--tracker", default="http://localhost:8080")
start_cmd.add_argument("--port", type=int, default=7000)
start_cmd.add_argument("--model", default="stub-model")
start_cmd.add_argument("--model-id", help="HuggingFace repo ID")
start_cmd.add_argument("--shard-start", type=int)
start_cmd.add_argument("--shard-end", type=int)
start_cmd.add_argument("--quantization", choices=["bfloat16", "int8", "nf4", "bf16"], default="bfloat16")
start_cmd.add_argument("--host", default="0.0.0.0")
start_cmd.add_argument("--advertise-host")
start_cmd.add_argument("--tracker-mode", action="store_true")
start_cmd.add_argument("--tracker-url", default=None)
start_cmd.add_argument("--wallet")
start_cmd.add_argument("--download-dir")
args = parser.parse_args()
if args.command == "start":
from meshnet_node.startup import run_startup
try:
node = run_startup(
tracker_url=args.tracker,
port=args.port,
model=args.model,
model_id=args.model_id,
shard_start=args.shard_start,
shard_end=args.shard_end,
quantization=args.quantization,
host=args.host,
advertise_host=args.advertise_host,
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
sys.exit(1)
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
node.stop()
sys.exit(0)
if args.command == "models":
sys.exit(_cmd_models(args))
elif args.command == "config":
sys.exit(_cmd_config(args))
elif args.command == "start":
sys.exit(_cmd_start(args))
else:
parser.print_help()
# Default: wizard or start with saved config
sys.exit(_cmd_default(args))
if __name__ == "__main__":

72
packages/node/meshnet_node/config.py Normal file
View File

@@ -0,0 +1,72 @@
"""Persistent node configuration — stored in ~/.config/meshnet/config.json."""
from __future__ import annotations
import json
import os
import stat
from pathlib import Path
_DEFAULT_CONFIG_DIR = Path.home() / ".config" / "meshnet"
_DEFAULT_CONFIG_FILE = _DEFAULT_CONFIG_DIR / "config.json"
_DEFAULT_DOWNLOAD_DIR = Path.home() / ".meshnet" / "models"
_DEFAULT_TRACKER_URL = "http://localhost:8080"
_DEFAULT_WALLET_PATH = str(Path.home() / ".config" / "meshnet" / "wallet.json")
_DEFAULT_QUANTIZATION = "nf4"
DEFAULTS = {
"model_hf_repo": "",
"model_name": "",
"quantization": _DEFAULT_QUANTIZATION,
"download_dir": str(_DEFAULT_DOWNLOAD_DIR),
"tracker_url": _DEFAULT_TRACKER_URL,
"wallet_path": _DEFAULT_WALLET_PATH,
"shard_start": None,
"shard_end": None,
"port": 7000,
"host": "0.0.0.0",
}
def config_path(override: Path | None = None) -> Path:
return override or _DEFAULT_CONFIG_FILE
def load_config(path: Path | None = None) -> dict | None:
"""Return parsed config dict, or None if no config file exists."""
p = config_path(path)
if not p.exists():
return None
try:
cfg = json.loads(p.read_text())
if not isinstance(cfg, dict):
return None
return cfg
except (json.JSONDecodeError, OSError):
return None
def save_config(cfg: dict, path: Path | None = None) -> None:
"""Write config to disk with restricted permissions (0o600)."""
p = config_path(path)
p.parent.mkdir(parents=True, exist_ok=True)
p.write_text(json.dumps(cfg, indent=2))
try:
os.chmod(p, stat.S_IRUSR | stat.S_IWUSR)
except OSError:
pass # Windows / some filesystems don't support chmod
def delete_config(path: Path | None = None) -> None:
p = config_path(path)
if p.exists():
p.unlink()
def merge_cli_overrides(cfg: dict, **cli_kwargs) -> dict:
"""Return a copy of cfg with any non-None CLI values applied on top."""
result = dict(cfg)
for key, val in cli_kwargs.items():
if val is not None:
result[key] = val
return result

220
packages/node/meshnet_node/dashboard.py Normal file
View File

@@ -0,0 +1,220 @@
"""Live node status dashboard — rich TUI with plain-text fallback."""
from __future__ import annotations
import os
import sys
import time
from collections import deque
from typing import TYPE_CHECKING
if TYPE_CHECKING:
pass
def is_interactive_tty() -> bool:
"""Return True when stdout is a real terminal (not CI / redirected / WSL2 dumb)."""
if not sys.stdout.isatty():
return False
term = os.environ.get("TERM", "")
if term in ("dumb", ""):
return False
return True
def _format_uptime(seconds: float) -> str:
s = int(seconds)
h, rem = divmod(s, 3600)
m, sec = divmod(rem, 60)
return f"{h:02d}:{m:02d}:{sec:02d}"
def _gpu_stats() -> list[dict]:
"""Return per-GPU utilization and VRAM stats, or empty list on CPU."""
try:
import torch # type: ignore[import]
if not torch.cuda.is_available():
return []
stats = []
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
used = torch.cuda.memory_allocated(i)
total = props.total_memory
# Utilization requires pynvml; skip gracefully if not available
util = _nvml_gpu_util(i)
stats.append(
{
"index": i,
"name": props.name,
"used_gb": used / 1e9,
"total_gb": total / 1e9,
"util_pct": util,
}
)
return stats
except ImportError:
return []
def _nvml_gpu_util(index: int) -> int | None:
"""Return GPU utilization % via pynvml, or None if unavailable."""
try:
import pynvml # type: ignore[import]
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
rates = pynvml.nvmlDeviceGetUtilizationRates(handle)
return rates.gpu
except Exception:
return None
class _EMA:
"""Exponential moving average for tokens/sec."""
def __init__(self, alpha: float = 0.1):
self._alpha = alpha
self._value: float | None = None
def update(self, sample: float) -> float:
if self._value is None:
self._value = sample
else:
self._value = self._alpha * sample + (1 - self._alpha) * self._value
return self._value
@property
def value(self) -> float:
return self._value or 0.0
def _make_bar(pct: float, width: int = 10) -> str:
filled = round(pct / 100 * width)
return "█" * filled + "░" * (width - filled)
def run_dashboard(node, config: dict, start_time: float) -> None:
"""Start the live dashboard. Blocks until Ctrl-C. Returns cleanly."""
if not is_interactive_tty():
_run_plain_loop(node, config, start_time)
return
try:
from rich.live import Live # type: ignore[import]
_run_rich_dashboard(node, config, start_time)
except ImportError:
_run_plain_loop(node, config, start_time)
def _build_rich_renderable(
node, config: dict, start_time: float, tps_ema: _EMA, prev_req: list[int]
):
from rich.table import Table # type: ignore[import]
from rich.panel import Panel # type: ignore[import]
from rich.columns import Columns # type: ignore[import]
from rich.text import Text # type: ignore[import]
uptime = time.monotonic() - start_time
req_count = getattr(node, "chat_completion_count", 0)
# Tokens/sec EMA (approximate: 20 tokens per request heuristic when no real counter)
delta_req = req_count - prev_req[0]
prev_req[0] = req_count
if delta_req > 0:
approx_tokens = delta_req * 20
tps_ema.update(approx_tokens / 2.0) # 2s interval
gpu_stats = _gpu_stats()
model_name = config.get("model_name") or config.get("model_hf_repo", "unknown").split("/")[-1]
shard = ""
if config.get("shard_start") is not None:
shard = f" shard {config['shard_start']}–{config['shard_end']}"
# Header line
header = Text(
f"meshnet-node {model_name} [{config.get('quantization', 'bf16')}]{shard}"
f" up {_format_uptime(uptime)}",
style="bold white",
)
# GPU table
gpu_table = Table(show_header=False, box=None, padding=(0, 1))
gpu_table.add_column("label", style="dim", no_wrap=True)
gpu_table.add_column("bar", no_wrap=True)
gpu_table.add_column("vram", no_wrap=True, style="cyan")
if gpu_stats:
for g in gpu_stats:
util = g["util_pct"]
util_str = f"{_make_bar(util)} {util:3d}%" if util is not None else " n/a"
vram_str = f"VRAM {g['used_gb']:.1f}/{g['total_gb']:.1f} GB"
gpu_table.add_row(f"GPU {g['index']} {g['name'][:20]}", util_str, vram_str)
else:
gpu_table.add_row("CPU mode", "", "no GPU detected")
# Stats panel
tps = tps_ema.value
bar_len = min(8, max(0, int(tps / 10)))
tps_bar = "▁▂▃▄▅▆▇█"[:bar_len].ljust(8)
stats_lines = [
f"Tokens/sec {tps_bar} {tps:.1f} t/s (EMA)",
f"Requests {req_count:,} served",
f"Peers 0 connected (gossip: US-017)",
f"TAI earned 0.00 TAI (payments: US-006)",
f"Uptime {_format_uptime(uptime)}",
"",
"[q] quit [c] compact view",
]
from rich.console import Group # type: ignore[import]
return Panel(
Group(header, gpu_table, Text("\n".join(stats_lines))),
title="[bold green]meshnet-node[/bold green]",
border_style="green",
)
def _run_rich_dashboard(node, config: dict, start_time: float) -> None:
from rich.live import Live # type: ignore[import]
tps_ema = _EMA()
prev_req = [0]
try:
with Live(
_build_rich_renderable(node, config, start_time, tps_ema, prev_req),
refresh_per_second=0.5,
screen=False,
) as live:
while True:
time.sleep(2)
live.update(
_build_rich_renderable(node, config, start_time, tps_ema, prev_req)
)
except KeyboardInterrupt:
pass
def _run_plain_loop(node, config: dict, start_time: float) -> None:
model_name = config.get("model_name") or config.get("model_hf_repo", "unknown").split("/")[-1]
try:
while True:
uptime = time.monotonic() - start_time
req = getattr(node, "chat_completion_count", 0)
gpu_stats = _gpu_stats()
vram_str = ""
if gpu_stats:
g = gpu_stats[0]
vram_str = f" VRAM{g['used_gb']:.1f}GB"
print(
f"[{model_name} req{req}{vram_str} up{_format_uptime(uptime)}]",
flush=True,
)
time.sleep(2)
except KeyboardInterrupt:
pass

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

@@ -241,6 +241,24 @@ class TorchModelShard:
yield token_text
t.join()
def count_prompt_tokens(self, messages: list[dict]) -> int:
"""Return tokenizer-backed prompt token count for OpenAI usage metadata."""
encoded = self._encode_messages(messages)
input_ids = encoded["input_ids"]
return int(input_ids.shape[-1])
def count_text_tokens(self, text: str) -> int:
"""Return tokenizer-backed completion token count for OpenAI usage metadata."""
try:
encoded = self.tokenizer(
text,
return_tensors="pt",
add_special_tokens=False,
)
except TypeError:
encoded = self.tokenizer(text, return_tensors="pt")
return int(encoded["input_ids"].shape[-1])
def _encode_messages(self, messages: list[dict]) -> dict:
"""Format messages with chat template (if available) and tokenize."""
if hasattr(self.tokenizer, "apply_chat_template"):

165
packages/node/meshnet_node/model_catalog.py Normal file
View File

@@ -0,0 +1,165 @@
"""Curated list of models supported by the network with VRAM requirements."""
from __future__ import annotations
from dataclasses import dataclass
@dataclass
class ModelPreset:
name: str
hf_repo: str
num_layers: int
# VRAM in GB at each quantization level (None = too large to quantize this way)
vram_nf4: float
vram_int8: float
vram_bf16: float
description: str
def vram_for_quant(self, quant: str) -> float:
"""Return VRAM requirement in GB for the given quantization."""
q = quant.lower().replace("bfloat16", "bf16")
if q == "nf4":
return self.vram_nf4
if q in ("int8", "int8"):
return self.vram_int8
if q in ("bf16", "bfloat16"):
return self.vram_bf16
raise ValueError(f"unknown quantization: {quant!r}")
def fits_vram(self, available_gb: float, quant: str) -> bool:
return self.vram_for_quant(quant) <= available_gb
def recommended_quant(self, available_gb: float) -> str | None:
"""Return the highest-quality quantization that fits available VRAM, or None."""
if self.vram_bf16 <= available_gb:
return "bf16"
if self.vram_int8 <= available_gb:
return "int8"
if self.vram_nf4 <= available_gb:
return "nf4"
return None
CURATED_MODELS: list[ModelPreset] = [
ModelPreset(
name="Qwen2.5-0.5B-Instruct",
hf_repo="Qwen/Qwen2.5-0.5B-Instruct",
num_layers=24,
vram_nf4=0.4,
vram_int8=0.6,
vram_bf16=1.0,
description="Smallest no-gating model — great for testing, ~1 GB",
),
ModelPreset(
name="Qwen2.5-1.5B-Instruct",
hf_repo="Qwen/Qwen2.5-1.5B-Instruct",
num_layers=28,
vram_nf4=1.0,
vram_int8=1.8,
vram_bf16=3.2,
description="Fast no-gating model — good quality, ~3 GB",
),
ModelPreset(
name="Llama-3-70B-Instruct",
hf_repo="meta-llama/Meta-Llama-3-70B-Instruct",
num_layers=80,
vram_nf4=18.0,
vram_int8=40.0,
vram_bf16=140.0,
description="Meta's flagship 70B instruction model",
),
ModelPreset(
name="Qwen2.5-72B-Instruct",
hf_repo="Qwen/Qwen2.5-72B-Instruct",
num_layers=80,
vram_nf4=19.0,
vram_int8=41.0,
vram_bf16=145.0,
description="Alibaba's 72B multilingual instruction model",
),
ModelPreset(
name="Mixtral-8x7B-Instruct",
hf_repo="mistralai/Mixtral-8x7B-Instruct-v0.1",
num_layers=32,
vram_nf4=7.0,
vram_int8=14.0,
vram_bf16=27.0,
description="Mistral's sparse MoE — fast and efficient",
),
ModelPreset(
name="Llama-3-8B-Instruct",
hf_repo="meta-llama/Meta-Llama-3-8B-Instruct",
num_layers=32, # gated repo — requires HF login
vram_nf4=4.5,
vram_int8=8.5,
vram_bf16=16.0,
description="Meta's compact 8B model — good for low-VRAM nodes",
),
ModelPreset(
name="Phi-3-medium-128k",
hf_repo="microsoft/Phi-3-medium-128k-instruct",
num_layers=40,
vram_nf4=4.0,
vram_int8=8.0,
vram_bf16=15.0,
description="Microsoft's efficient 14B model with 128k context",
),
ModelPreset(
name="Gemma-2-27B-IT",
hf_repo="google/gemma-2-27b-it",
num_layers=46,
vram_nf4=10.0,
vram_int8=20.0,
vram_bf16=54.0,
description="Google's 27B instruction-tuned model",
),
ModelPreset(
name="DeepSeek-V2-Lite-Chat",
hf_repo="deepseek-ai/DeepSeek-V2-Lite-Chat",
num_layers=27,
vram_nf4=5.0,
vram_int8=9.0,
vram_bf16=16.0,
description="DeepSeek's efficient MoE — strong coding + reasoning",
),
]
def detect_num_layers(hf_repo: str) -> int | None:
"""Return num_hidden_layers from HuggingFace config.json (downloads ~1 KB only)."""
# Check curated list first (no network call)
for m in CURATED_MODELS:
if m.hf_repo == hf_repo:
return m.num_layers
try:
from transformers import AutoConfig # type: ignore[import]
cfg = AutoConfig.from_pretrained(hf_repo)
return int(cfg.num_hidden_layers)
except Exception:
return None
def browse_hf_hub(top_n: int = 20) -> list[dict]:
"""Fetch top downloaded text-generation models from HuggingFace Hub."""
try:
from huggingface_hub import list_models # type: ignore[import]
models = list(
list_models(
pipeline_tag="text-generation",
library="transformers",
sort="downloads",
direction=-1,
limit=top_n,
)
)
return [
{
"repo": m.id,
"downloads": getattr(m, "downloads", 0) or 0,
}
for m in models
]
except Exception as exc:
raise RuntimeError(f"HuggingFace Hub lookup failed: {exc}") from exc

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

@@ -84,7 +84,19 @@ def run_startup(
if probationary_line is not None:
print(f" {probationary_line}", flush=True)
if model_id is not None and shard_start is not None and shard_end is not None:
if model_id is not None:
# Auto-detect shard range from model config if not explicitly provided
if shard_start is None or shard_end is None:
detected = _detect_num_layers(model_id)
if detected is None:
raise ValueError(
f"Could not read num_hidden_layers from {model_id} config. "
"Pass --shard-start and --shard-end explicitly."
)
shard_start = shard_start if shard_start is not None else 0
shard_end = shard_end if shard_end is not None else detected - 1
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
print("Loading real PyTorch model shard...", flush=True)
node = TorchNodeServer(
host=host,
@@ -93,8 +105,15 @@ def run_startup(
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
tracker_url=tracker_url,
)
actual_port = node.start()
total_layers = getattr(node.backend, "total_layers", None)
if isinstance(total_layers, int) and total_layers > 0:
layer_count = shard_end - shard_start + 1
shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}"
else:
shard_label = f"layers {shard_start}–{shard_end}"
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
print(
@@ -102,7 +121,7 @@ def run_startup(
f"meshnet-node ready\n"
f" Wallet: {address}\n"
f" Model ID: {model_id}\n"
f" Shard: layers {shard_start}-{shard_end}\n"
f" Shard: {shard_label}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Hardware: {device.upper()}\n"
@@ -110,8 +129,8 @@ def run_startup(
flush=True,
)
return node
if model_id is not None or shard_start is not None or shard_end is not None:
raise ValueError("--model-id, --shard-start, and --shard-end must be provided together")
if shard_start is not None or shard_end is not None:
raise ValueError("--shard-start / --shard-end require --model-id")
# 3. Shard assignment from tracker
print("Querying tracker for shard assignment...", flush=True)
@@ -201,6 +220,17 @@ def run_startup(
return node
def _detect_num_layers(model_id: str) -> int | None:
"""Fetch num_hidden_layers from HuggingFace model config (downloads ~1 KB config.json only)."""
try:
from transformers import AutoConfig # type: ignore[import]
cfg = AutoConfig.from_pretrained(model_id)
return int(cfg.num_hidden_layers)
except Exception as exc:
print(f" Warning: could not read model config from HF: {exc}", flush=True)
return None
def _probationary_status_line(contracts: Any | None, wallet_address: str) -> str | None:
if contracts is None:
return None

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

@@ -11,6 +11,7 @@ import urllib.error
import urllib.parse
import urllib.request
import uuid
from typing import Any
from .model_backend import (
InsufficientVRAMError,
@@ -232,7 +233,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
)
else:
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
self._send_openai_response(text, model_name, False)
self._send_openai_response(text, model_name, False, messages)
except Exception as exc:
self._send_json(500, {"error": f"generation failed: {exc}"})
return
@@ -250,7 +251,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
return
remaining_route = self._get_remaining_route(model_name)
result_text = self._run_downstream_pipeline(payload, remaining_route)
self._send_openai_response(result_text, model_name, stream)
self._send_openai_response(result_text, model_name, stream, messages)
def _get_remaining_route(self, model: str) -> list[str]:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
@@ -367,10 +368,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.wfile.write(b"data: [DONE]\n\n")
self.wfile.flush()
def _send_openai_response(self, text: str, model: str, stream: bool) -> None:
def _send_openai_response(
self,
text: str,
model: str,
stream: bool,
messages: list[dict] | None = None,
) -> None:
chunk_id = "chatcmpl-node"
created = int(time.time())
if not stream:
usage = _usage_for_response(self.server.backend, messages or [], text) # type: ignore[attr-defined]
self._send_json(200, {
"id": chunk_id,
"object": "chat.completion",
@@ -381,7 +389,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
"message": {"role": "assistant", "content": text},
"finish_reason": "stop",
}],
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
"usage": usage,
})
return
self.send_response(200)
@@ -412,6 +420,52 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.wfile.flush()
def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
prompt_tokens = _backend_token_count(
backend,
"count_prompt_tokens",
messages,
fallback=_fallback_message_token_count(messages),
)
completion_tokens = _backend_token_count(
backend,
"count_text_tokens",
completion_text,
fallback=_fallback_text_token_count(completion_text),
)
return {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int:
method: Any = getattr(backend, method_name, None)
if callable(method):
try:
return max(0, int(method(value)))
except Exception:
pass
return max(0, int(fallback))
def _fallback_message_token_count(messages: list[dict]) -> int:
text = " ".join(
str(message.get("content", ""))
for message in messages
if isinstance(message, dict)
)
return _fallback_text_token_count(text)
def _fallback_text_token_count(text: str) -> int:
parts = text.split()
if parts:
return len(parts)
return 1 if text else 0
class TorchNodeServer:
"""HTTP server backed by a HuggingFace causal language model shard."""

332
packages/node/meshnet_node/wizard.py Normal file
View File

@@ -0,0 +1,332 @@
"""Interactive first-run setup wizard — mining-client style."""
from __future__ import annotations
import sys
import urllib.error
import urllib.request
from pathlib import Path
from typing import TYPE_CHECKING
from .config import DEFAULTS, _DEFAULT_DOWNLOAD_DIR, _DEFAULT_TRACKER_URL, _DEFAULT_WALLET_PATH
from .model_catalog import CURATED_MODELS, ModelPreset, browse_hf_hub, detect_num_layers
if TYPE_CHECKING:
pass
_HEADER = """\
╔══════════════════════════════════════════════════════════════════╗
║ meshnet-node v0.1.0 ║
║ Distributed AI Inference — Node Setup ║
╚══════════════════════════════════════════════════════════════════╝
"""
_QUANT_LABELS = {"nf4": "NF4 (4-bit)", "int8": "INT8 (8-bit)", "bf16": "BF16 (full)"}
def _ask(prompt: str, default: str = "", validator=None) -> str:
"""Prompt user and return answer. Returns default on empty input or EOF."""
display = f"{prompt} [{default}]: " if default else f"{prompt}: "
while True:
try:
raw = input(display).strip()
except (EOFError, KeyboardInterrupt):
print()
raise KeyboardInterrupt
value = raw or default
if validator is None or validator(value):
return value
# validator returned error string
print(f" ✗ {validator(value)}")
def _ask_int(prompt: str, default: int, lo: int, hi: int) -> int:
def validate(s: str) -> bool | str:
try:
v = int(s)
except ValueError:
return "Please enter a number."
if not (lo <= v <= hi):
return f"Please enter a number between {lo} and {hi}."
return True
while True:
raw = _ask(prompt, str(default))
try:
v = int(raw)
if lo <= v <= hi:
return v
except ValueError:
pass
print(f" ✗ Enter a number between {lo} and {hi}.")
def _ask_yn(prompt: str, default: bool = True) -> bool:
hint = "Y/n" if default else "y/N"
raw = _ask(f"{prompt} [{hint}]").lower()
if not raw:
return default
return raw.startswith("y")
def _detect_gpus() -> list[dict]:
"""Return list of detected GPU dicts with name and vram_gb."""
gpus: list[dict] = []
try:
import torch # type: ignore[import]
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
gpus.append(
{
"index": i,
"name": props.name,
"vram_gb": props.total_memory / 1e9,
"backend": "cuda",
}
)
except ImportError:
pass
return gpus
def _total_vram_gb(gpus: list[dict]) -> float:
return sum(g["vram_gb"] for g in gpus)
def _print_gpus(gpus: list[dict]) -> None:
if not gpus:
print(" ⚠ No CUDA GPU detected — running in CPU mode")
print(" CPU inference is very slow. Consider a machine with an NVIDIA GPU.")
return
for g in gpus:
vram = g["vram_gb"]
print(f" GPU {g['index']}: {g['name']} {vram:.0f} GB VRAM ✓")
def _print_model_table(gpus: list[dict], quant: str = "nf4") -> None:
available_gb = _total_vram_gb(gpus)
print()
print(f" # {'Model':<30} {'Layers':>6} {'NF4':>6} {'INT8':>6} {'BF16':>6}")
print(f" {'─'*4} {'─'*30} {'─'*6} {'─'*6} {'─'*6} {'─'*6}")
for i, m in enumerate(CURATED_MODELS, 1):
fits_nf4 = "✓" if m.vram_nf4 <= available_gb else "✗"
fits_int8 = "✓" if m.vram_int8 <= available_gb else "✗"
fits_bf16 = "✓" if m.vram_bf16 <= available_gb else "✗"
nf4_str = f"{fits_nf4}{m.vram_nf4:.0f}GB"
int8_str = f"{fits_int8}{m.vram_int8:.0f}GB"
bf16_str = f"{fits_bf16}{m.vram_bf16:.0f}GB"
print(f" {i:<3} {m.name:<30} {m.num_layers:>6} {nf4_str:>6} {int8_str:>6} {bf16_str:>6}")
print(f" {m.description}")
idx = len(CURATED_MODELS) + 1
print(f" {idx:<3} {'[Browse HuggingFace Hub...]':<30}")
print()
def _browse_hf_interactive() -> str | None:
"""Show HF Hub top-20 and let user enter a repo ID. Returns repo ID or None to go back."""
print("\nFetching top models from HuggingFace Hub...")
try:
models = browse_hf_hub(top_n=20)
except RuntimeError as exc:
print(f" ✗ {exc}")
return None
print(f"\n {'#':<4} {'HuggingFace Repo':<50} Downloads")
print(f" {'─'*4} {'─'*50} {'─'*10}")
for i, m in enumerate(models, 1):
dl = m["downloads"]
dl_str = f"{dl/1e6:.1f}M" if dl >= 1_000_000 else f"{dl/1e3:.0f}k" if dl >= 1000 else str(dl)
print(f" {i:<4} {m['repo']:<50} {dl_str}")
print()
raw = _ask(
"Enter a number to select, or paste any HuggingFace repo ID (or press Enter to go back)",
default="",
)
if not raw:
return None
try:
idx = int(raw) - 1
if 0 <= idx < len(models):
return models[idx]["repo"]
except ValueError:
pass
# Treat raw input as a repo ID
if "/" in raw:
return raw
print(" ✗ Invalid input — please enter a number or a full repo ID like 'org/model-name'")
return None
def _ask_quant(gpus: list[dict], model: ModelPreset | None) -> str:
available_gb = _total_vram_gb(gpus)
print("\nQuantization level:")
options: list[tuple[str, str]] = []
for quant, label in [("nf4", "NF4 4-bit"), ("int8", "INT8 8-bit"), ("bf16", "BF16 full precision")]:
if model is not None:
vram = model.vram_for_quant(quant)
fits = "✓" if vram <= available_gb else "✗ insufficient VRAM"
suffix = f" ({vram:.0f} GB needed — {fits})"
else:
suffix = ""
options.append((quant, f"{label}{suffix}"))
for i, (_, label) in enumerate(options, 1):
print(f" {i}) {label}")
# Recommend the best fitting quant
if model is not None:
rec = model.recommended_quant(available_gb)
rec_idx = next((i for i, (q, _) in enumerate(options, 1) if q == rec), 1) if rec else 1
default_idx = rec_idx
print(f" (Recommended: {rec.upper() if rec else 'NF4'} for your GPU)")
else:
default_idx = 1
choice = _ask_int("Enter number", default_idx, 1, 3)
return options[choice - 1][0]
def _validate_dir(path_str: str) -> bool | str:
p = Path(path_str).expanduser()
try:
p.mkdir(parents=True, exist_ok=True)
return True
except OSError as exc:
return f"Cannot create directory: {exc}"
def _validate_tracker(url: str) -> bool | str:
if not url.startswith(("http://", "https://")):
return "URL must start with http:// or https://"
return True
def _ping_tracker(url: str) -> bool:
"""Return True if tracker responds to /health."""
try:
with urllib.request.urlopen(f"{url.rstrip('/')}/health", timeout=3):
return True
except Exception:
return False
def run_wizard(config_path_override=None) -> dict:
"""Run the interactive setup wizard and return a config dict.
Raises KeyboardInterrupt if user presses Ctrl-C.
"""
print(_HEADER)
# Step 1: GPU detection
print("Detecting hardware...")
gpus = _detect_gpus()
_print_gpus(gpus)
available_gb = _total_vram_gb(gpus)
if available_gb == 0:
available_gb = 9999 # CPU — don't filter models by VRAM
# Step 2 & 3: Model selection
print("\nSelect a model to serve:\n")
selected_repo: str | None = None
selected_preset: ModelPreset | None = None
while selected_repo is None:
_print_model_table(gpus)
lo, hi = 1, len(CURATED_MODELS) + 1
choice = _ask_int("Enter number", 1, lo, hi)
if choice == len(CURATED_MODELS) + 1:
repo = _browse_hf_interactive()
if repo:
# Look up layer count for custom repo
print(f" Checking {repo} config...", end=" ", flush=True)
layers = detect_num_layers(repo)
if layers:
print(f"{layers} layers")
else:
print("(layer count unknown — will detect on start)")
selected_repo = repo
selected_preset = None
else:
selected_preset = CURATED_MODELS[choice - 1]
selected_repo = selected_preset.hf_repo
if selected_preset.recommended_quant(available_gb) is None:
print(
f"\n ⚠ Warning: {selected_preset.name} requires at least "
f"{selected_preset.vram_nf4:.0f} GB VRAM at NF4 — even the smallest "
f"quantization may be too large for your GPU."
)
if not _ask_yn("Continue anyway?", default=False):
selected_repo = None
selected_preset = None
num_layers = (selected_preset.num_layers if selected_preset
else detect_num_layers(selected_repo or ""))
layers_str = f" {num_layers} layers" if num_layers else ""
print(f"\n ✓ Selected: {selected_repo}{layers_str}")
# Step 3b: Quantization
quant = _ask_quant(gpus, selected_preset)
print(f" ✓ Quantization: {quant.upper()}")
# Step 4: Download directory
print()
dl_dir = _ask(
"Download directory",
default=str(_DEFAULT_DOWNLOAD_DIR),
validator=lambda v: _validate_dir(v) if v else "Directory is required.",
)
print(f" ✓ Download dir: {dl_dir}")
# Step 5: Tracker URL
print()
tracker_url = _DEFAULT_TRACKER_URL
raw_tracker = _ask("Tracker URL", default=_DEFAULT_TRACKER_URL, validator=_validate_tracker)
tracker_url = raw_tracker
if _ping_tracker(tracker_url):
print(f" ✓ Tracker reachable: {tracker_url}")
else:
print(f" ⚠ Tracker not reachable at {tracker_url} (will retry on start)")
# Step 6: Wallet path
print()
wallet_path = _ask("Wallet path", default=_DEFAULT_WALLET_PATH)
print(f" ✓ Wallet: {wallet_path}")
cfg = {
"model_hf_repo": selected_repo,
"model_name": selected_preset.name if selected_preset else selected_repo.split("/")[-1],
"quantization": quant,
"download_dir": dl_dir,
"tracker_url": tracker_url,
"wallet_path": wallet_path,
"shard_start": None,
"shard_end": None,
"port": DEFAULTS["port"],
"host": DEFAULTS["host"],
}
return cfg
def print_models_table(available_gb: float | None = None) -> None:
"""Print curated model table for `meshnet-node models`."""
gpus: list[dict] = []
if available_gb is None:
gpus = _detect_gpus()
available_gb = _total_vram_gb(gpus) or 9999
else:
gpus = [{"index": 0, "name": "GPU", "vram_gb": available_gb, "backend": "cuda"}]
print(f"\n{'#':<4} {'Model':<32} {'HuggingFace Repo':<45} {'Layers':>6} {'NF4':>8} {'INT8':>8} {'BF16':>8}")
print(f"{'─'*4} {'─'*32} {'─'*45} {'─'*6} {'─'*8} {'─'*8} {'─'*8}")
for i, m in enumerate(CURATED_MODELS, 1):
def _cell(vram: float) -> str:
fits = "✓" if vram <= available_gb else "✗"
return f"{fits}{vram:.0f}GB"
print(
f"{i:<4} {m.name:<32} {m.hf_repo:<45} {m.num_layers:>6} "
f"{_cell(m.vram_nf4):>8} {_cell(m.vram_int8):>8} {_cell(m.vram_bf16):>8}"
)
print()