Compare commits
3 Commits
bbe57d5f07
...
worktree-f
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
080d49b2c2 | ||
|
|
a2258d3df4 | ||
|
|
65f3ee6a85 |
@@ -11,6 +11,7 @@ _packages = [
|
||||
"packages/sdk",
|
||||
"packages/contracts",
|
||||
"packages/p2p",
|
||||
"packages/relay",
|
||||
"packages/validator",
|
||||
]
|
||||
|
||||
|
||||
@@ -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
72
packages/node/meshnet_node/config.py
Normal 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
220
packages/node/meshnet_node/dashboard.py
Normal 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
|
||||
165
packages/node/meshnet_node/model_catalog.py
Normal file
165
packages/node/meshnet_node/model_catalog.py
Normal 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
|
||||
@@ -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,
|
||||
@@ -102,7 +114,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: layers {shard_start}–{shard_end}\n"
|
||||
f" Quantization: {quantization}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Hardware: {device.upper()}\n"
|
||||
@@ -110,8 +122,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 +213,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
|
||||
|
||||
332
packages/node/meshnet_node/wizard.py
Normal file
332
packages/node/meshnet_node/wizard.py
Normal 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()
|
||||
@@ -13,6 +13,7 @@ dependencies = [
|
||||
"huggingface-hub>=0.20",
|
||||
"accelerate>=0.28",
|
||||
"bitsandbytes>=0.43",
|
||||
"rich>=13",
|
||||
"safetensors>=0.4",
|
||||
"torch>=2.1",
|
||||
"transformers>=4.39",
|
||||
|
||||
162
packages/p2p/meshnet_p2p/gossip.py
Normal file
162
packages/p2p/meshnet_p2p/gossip.py
Normal file
@@ -0,0 +1,162 @@
|
||||
"""WebSocket gossip client — connects to relay, publish/subscribe to topics."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from typing import Callable
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Message envelope topics
|
||||
TOPIC_NODE_JOIN = "node-join"
|
||||
TOPIC_NODE_LEAVE = "node-leave"
|
||||
TOPIC_COVERAGE_UPDATE = "coverage-update"
|
||||
TOPIC_HEARTBEAT = "heartbeat"
|
||||
TOPIC_PEER_LIST = "peer-list"
|
||||
TOPIC_RELAY_ANNOUNCE = "relay-announce"
|
||||
|
||||
_MSG_TTL = 3 # max re-broadcast hops
|
||||
|
||||
|
||||
def _make_envelope(topic: str, payload: dict, from_peer: str, ttl: int = _MSG_TTL) -> dict:
|
||||
return {
|
||||
"topic": topic,
|
||||
"version": 1,
|
||||
"from_peer": from_peer,
|
||||
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||||
"msg_id": str(uuid.uuid4()),
|
||||
"ttl": ttl,
|
||||
"payload": payload,
|
||||
}
|
||||
|
||||
|
||||
class GossipClient:
|
||||
"""Thread-safe WebSocket gossip client.
|
||||
|
||||
Usage::
|
||||
|
||||
client = GossipClient(relay_url="ws://relay:8765", peer_id="abc123")
|
||||
client.subscribe("node-join", lambda env: print(env["payload"]))
|
||||
client.start()
|
||||
client.publish("node-join", {"addr": "http://192.168.1.42:8001", ...})
|
||||
...
|
||||
client.stop()
|
||||
"""
|
||||
|
||||
def __init__(self, relay_url: str, peer_id: str, reconnect_interval: float = 3.0):
|
||||
self.relay_url = relay_url
|
||||
self.peer_id = peer_id
|
||||
self.reconnect_interval = reconnect_interval
|
||||
|
||||
self._handlers: dict[str, list[Callable]] = defaultdict(list)
|
||||
self._seen: set[str] = set()
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self._ws = None
|
||||
self._running = False
|
||||
self._connected = threading.Event()
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
|
||||
def subscribe(self, topic: str, handler: Callable) -> None:
|
||||
"""Register a sync callback for messages on topic."""
|
||||
self._handlers[topic].append(handler)
|
||||
|
||||
def publish(self, topic: str, payload: dict) -> None:
|
||||
"""Send a gossip message to all peers via the relay. Thread-safe."""
|
||||
envelope = _make_envelope(topic, payload, self.peer_id)
|
||||
if self._loop and self._running:
|
||||
asyncio.run_coroutine_threadsafe(self._send(envelope), self._loop)
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the gossip client in a background thread."""
|
||||
self._running = True
|
||||
self._loop = asyncio.new_event_loop()
|
||||
self._thread = threading.Thread(target=self._run_loop, daemon=True, name="gossip")
|
||||
self._thread.start()
|
||||
|
||||
def wait_connected(self, timeout: float = 5.0) -> bool:
|
||||
"""Block until connected to relay or timeout. Returns True if connected."""
|
||||
return self._connected.wait(timeout)
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the gossip client and clean up."""
|
||||
self._running = False
|
||||
if self._loop and self._stop_event is not None:
|
||||
self._loop.call_soon_threadsafe(self._stop_event.set)
|
||||
if self._thread:
|
||||
self._thread.join(timeout=3.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal asyncio methods (run inside the background event loop)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_loop(self) -> None:
|
||||
asyncio.set_event_loop(self._loop)
|
||||
self._stop_event = asyncio.Event()
|
||||
try:
|
||||
self._loop.run_until_complete(self._connect_loop())
|
||||
except Exception:
|
||||
log.debug("Gossip loop exited", exc_info=True)
|
||||
|
||||
async def _connect_loop(self) -> None:
|
||||
import websockets # type: ignore[import]
|
||||
|
||||
while self._running and not (self._stop_event and self._stop_event.is_set()):
|
||||
try:
|
||||
async with websockets.connect(
|
||||
self.relay_url,
|
||||
ping_interval=20,
|
||||
ping_timeout=10,
|
||||
open_timeout=5,
|
||||
) as ws:
|
||||
self._ws = ws
|
||||
self._connected.set()
|
||||
log.debug("Gossip connected to %s", self.relay_url)
|
||||
# Send peer registration
|
||||
await ws.send(json.dumps(
|
||||
_make_envelope(
|
||||
"peer-register",
|
||||
{"peer_id": self.peer_id},
|
||||
self.peer_id,
|
||||
)
|
||||
))
|
||||
await self._receive_loop(ws)
|
||||
except Exception as exc:
|
||||
self._connected.clear()
|
||||
if self._running:
|
||||
log.debug("Gossip disconnected (%s); reconnecting in %ss", exc, self.reconnect_interval)
|
||||
await asyncio.sleep(self.reconnect_interval)
|
||||
|
||||
async def _receive_loop(self, ws) -> None:
|
||||
async for raw in ws:
|
||||
try:
|
||||
envelope = json.loads(raw)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
msg_id = envelope.get("msg_id", "")
|
||||
if msg_id in self._seen:
|
||||
continue
|
||||
self._seen.add(msg_id)
|
||||
if len(self._seen) > 10_000:
|
||||
# Trim seen set to avoid unbounded growth
|
||||
self._seen = set(list(self._seen)[-5_000:])
|
||||
|
||||
topic = envelope.get("topic", "")
|
||||
for handler in self._handlers.get(topic, []):
|
||||
try:
|
||||
handler(envelope)
|
||||
except Exception:
|
||||
log.debug("Gossip handler error for topic %s", topic, exc_info=True)
|
||||
|
||||
async def _send(self, envelope: dict) -> None:
|
||||
if self._ws is not None:
|
||||
try:
|
||||
await self._ws.send(json.dumps(envelope))
|
||||
except Exception as exc:
|
||||
log.debug("Gossip send failed: %s", exc)
|
||||
64
packages/p2p/meshnet_p2p/identity.py
Normal file
64
packages/p2p/meshnet_p2p/identity.py
Normal file
@@ -0,0 +1,64 @@
|
||||
"""Peer identity — stable peer_id and RSA keypair, persisted to disk."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
from pathlib import Path
|
||||
|
||||
_DEFAULT_IDENTITY_PATH = Path.home() / ".config" / "meshnet" / "identity.json"
|
||||
|
||||
|
||||
def _generate_keypair() -> tuple[bytes, bytes]:
|
||||
"""Return (private_key_pem, public_key_pem) for a new RSA-2048 keypair."""
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
|
||||
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
priv_pem = key.private_bytes(
|
||||
serialization.Encoding.PEM,
|
||||
serialization.PrivateFormat.PKCS8,
|
||||
serialization.NoEncryption(),
|
||||
)
|
||||
pub_pem = key.public_key().public_bytes(
|
||||
serialization.Encoding.PEM,
|
||||
serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
return priv_pem, pub_pem
|
||||
|
||||
|
||||
def _peer_id_from_pubkey(pub_pem: bytes) -> str:
|
||||
return hashlib.sha256(pub_pem).hexdigest()[:16]
|
||||
|
||||
|
||||
def load_or_create_identity(path: Path | None = None) -> dict:
|
||||
"""Return identity dict with peer_id, private_key_pem, public_key_pem.
|
||||
|
||||
Creates and persists a new identity if none exists at path.
|
||||
"""
|
||||
p = path or _DEFAULT_IDENTITY_PATH
|
||||
if p.exists():
|
||||
try:
|
||||
data = json.loads(p.read_text())
|
||||
if "peer_id" in data and "public_key_pem" in data:
|
||||
return data
|
||||
except (json.JSONDecodeError, OSError):
|
||||
pass
|
||||
|
||||
priv_pem, pub_pem = _generate_keypair()
|
||||
identity = {
|
||||
"peer_id": _peer_id_from_pubkey(pub_pem),
|
||||
"private_key_pem": priv_pem.decode(),
|
||||
"public_key_pem": pub_pem.decode(),
|
||||
}
|
||||
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text(json.dumps(identity, indent=2))
|
||||
try:
|
||||
os.chmod(p, stat.S_IRUSR | stat.S_IWUSR)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
return identity
|
||||
136
packages/p2p/meshnet_p2p/mdns.py
Normal file
136
packages/p2p/meshnet_p2p/mdns.py
Normal file
@@ -0,0 +1,136 @@
|
||||
"""mDNS peer discovery using zeroconf (optional dependency).
|
||||
|
||||
Falls back gracefully if zeroconf is not installed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import socket
|
||||
import threading
|
||||
from typing import Callable
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
MDNS_SERVICE_TYPE = "_meshnet._tcp.local."
|
||||
|
||||
try:
|
||||
from zeroconf import ServiceInfo, ServiceBrowser, Zeroconf # type: ignore[import]
|
||||
|
||||
_HAS_ZEROCONF = True
|
||||
except ImportError:
|
||||
_HAS_ZEROCONF = False
|
||||
|
||||
|
||||
def _local_ip() -> str:
|
||||
try:
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
s.connect(("8.8.8.8", 80))
|
||||
ip = s.getsockname()[0]
|
||||
s.close()
|
||||
return ip
|
||||
except OSError:
|
||||
return "127.0.0.1"
|
||||
|
||||
|
||||
class MdnsDiscovery:
|
||||
"""Announce this node on mDNS and discover peers on the same LAN.
|
||||
|
||||
If `zeroconf` is not installed, all methods are no-ops.
|
||||
|
||||
Usage::
|
||||
|
||||
disc = MdnsDiscovery(
|
||||
peer_id="abc123",
|
||||
port=8001,
|
||||
on_peer_found=lambda peer_id, addr: print("found", peer_id, addr),
|
||||
)
|
||||
disc.start()
|
||||
...
|
||||
disc.stop()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
peer_id: str,
|
||||
port: int,
|
||||
on_peer_found: Callable[[str, str], None] | None = None,
|
||||
on_peer_lost: Callable[[str], None] | None = None,
|
||||
):
|
||||
self.peer_id = peer_id
|
||||
self.port = port
|
||||
self.on_peer_found = on_peer_found
|
||||
self.on_peer_lost = on_peer_lost
|
||||
self._zc: "Zeroconf | None" = None # type: ignore[name-defined]
|
||||
self._info: "ServiceInfo | None" = None # type: ignore[name-defined]
|
||||
self._browser = None
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return _HAS_ZEROCONF
|
||||
|
||||
def start(self) -> None:
|
||||
if not _HAS_ZEROCONF:
|
||||
log.info("zeroconf not installed — mDNS discovery disabled")
|
||||
return
|
||||
try:
|
||||
self._zc = Zeroconf()
|
||||
local_ip = _local_ip()
|
||||
self._info = ServiceInfo(
|
||||
MDNS_SERVICE_TYPE,
|
||||
f"{self.peer_id}.{MDNS_SERVICE_TYPE}",
|
||||
addresses=[socket.inet_aton(local_ip)],
|
||||
port=self.port,
|
||||
properties={"peer_id": self.peer_id, "version": "1"},
|
||||
)
|
||||
self._zc.register_service(self._info)
|
||||
if self.on_peer_found or self.on_peer_lost:
|
||||
self._browser = ServiceBrowser(
|
||||
self._zc, MDNS_SERVICE_TYPE, listener=_Listener(self)
|
||||
)
|
||||
log.info("mDNS announced: %s on %s:%d", self.peer_id, local_ip, self.port)
|
||||
except Exception as exc:
|
||||
log.warning("mDNS start failed: %s", exc)
|
||||
|
||||
def stop(self) -> None:
|
||||
if not _HAS_ZEROCONF or self._zc is None:
|
||||
return
|
||||
try:
|
||||
if self._info:
|
||||
self._zc.unregister_service(self._info)
|
||||
self._zc.close()
|
||||
except Exception as exc:
|
||||
log.debug("mDNS stop error: %s", exc)
|
||||
self._zc = None
|
||||
|
||||
|
||||
class _Listener:
|
||||
"""Internal zeroconf service listener."""
|
||||
|
||||
def __init__(self, disc: MdnsDiscovery):
|
||||
self._disc = disc
|
||||
|
||||
def add_service(self, zc, type_, name):
|
||||
try:
|
||||
info = zc.get_service_info(type_, name)
|
||||
if info is None:
|
||||
return
|
||||
remote_peer_id = (info.properties or {}).get(b"peer_id", b"").decode()
|
||||
if remote_peer_id == self._disc.peer_id:
|
||||
return # ignore self
|
||||
addr = f"http://{socket.inet_ntoa(info.addresses[0])}:{info.port}"
|
||||
if self._disc.on_peer_found:
|
||||
self._disc.on_peer_found(remote_peer_id, addr)
|
||||
except Exception as exc:
|
||||
log.debug("mDNS add_service error: %s", exc)
|
||||
|
||||
def remove_service(self, zc, type_, name):
|
||||
try:
|
||||
# name is like "peer_id._meshnet._tcp.local."
|
||||
peer_id = name.split(".")[0]
|
||||
if self._disc.on_peer_lost:
|
||||
self._disc.on_peer_lost(peer_id)
|
||||
except Exception as exc:
|
||||
log.debug("mDNS remove_service error: %s", exc)
|
||||
|
||||
def update_service(self, zc, type_, name):
|
||||
pass
|
||||
114
packages/p2p/meshnet_p2p/tls.py
Normal file
114
packages/p2p/meshnet_p2p/tls.py
Normal file
@@ -0,0 +1,114 @@
|
||||
"""TLS certificate generation and fingerprint helpers for node-to-node comms."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import ssl
|
||||
import stat
|
||||
from pathlib import Path
|
||||
|
||||
_CERT_PATH = Path.home() / ".config" / "meshnet" / "node_cert.pem"
|
||||
_KEY_PATH = Path.home() / ".config" / "meshnet" / "node_key.pem"
|
||||
|
||||
|
||||
def generate_self_signed_cert(
|
||||
cert_path: Path | None = None,
|
||||
key_path: Path | None = None,
|
||||
common_name: str | None = None,
|
||||
) -> tuple[Path, Path]:
|
||||
"""Generate a self-signed RSA-2048 cert valid for 10 years.
|
||||
|
||||
Returns (cert_path, key_path). Skips generation if both files already exist.
|
||||
"""
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from cryptography.x509.oid import NameOID
|
||||
|
||||
cert_p = cert_path or _CERT_PATH
|
||||
key_p = key_path or _KEY_PATH
|
||||
|
||||
if cert_p.exists() and key_p.exists():
|
||||
return cert_p, key_p
|
||||
|
||||
cert_p.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
cn = common_name or socket.getfqdn()
|
||||
|
||||
subject = issuer = x509.Name([
|
||||
x509.NameAttribute(NameOID.COMMON_NAME, cn),
|
||||
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "meshnet-node"),
|
||||
])
|
||||
|
||||
san_list: list = [x509.DNSName(cn)]
|
||||
try:
|
||||
san_list.append(x509.IPAddress(ipaddress.IPv4Address(socket.gethostbyname(cn))))
|
||||
except (socket.gaierror, ValueError):
|
||||
pass
|
||||
san_list.append(x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")))
|
||||
|
||||
cert = (
|
||||
x509.CertificateBuilder()
|
||||
.subject_name(subject)
|
||||
.issuer_name(issuer)
|
||||
.public_key(key.public_key())
|
||||
.serial_number(x509.random_serial_number())
|
||||
.not_valid_before(datetime.datetime.now(datetime.timezone.utc))
|
||||
.not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=365 * 10))
|
||||
.add_extension(x509.SubjectAlternativeName(san_list), critical=False)
|
||||
.sign(key, hashes.SHA256())
|
||||
)
|
||||
|
||||
key_pem = key.private_bytes(
|
||||
serialization.Encoding.PEM,
|
||||
serialization.PrivateFormat.TraditionalOpenSSL,
|
||||
serialization.NoEncryption(),
|
||||
)
|
||||
cert_pem = cert.public_bytes(serialization.Encoding.PEM)
|
||||
|
||||
key_p.write_bytes(key_pem)
|
||||
cert_p.write_bytes(cert_pem)
|
||||
try:
|
||||
os.chmod(key_p, stat.S_IRUSR | stat.S_IWUSR)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
return cert_p, key_p
|
||||
|
||||
|
||||
def cert_fingerprint(cert_path: Path | None = None) -> str:
|
||||
"""Return sha256 fingerprint of the cert as 'sha256:<hex>'."""
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
|
||||
p = cert_path or _CERT_PATH
|
||||
cert = x509.load_pem_x509_certificate(p.read_bytes())
|
||||
fp = cert.fingerprint(hashes.SHA256()).hex()
|
||||
return f"sha256:{fp}"
|
||||
|
||||
|
||||
def make_server_ssl_context(
|
||||
cert_path: Path | None = None,
|
||||
key_path: Path | None = None,
|
||||
) -> ssl.SSLContext:
|
||||
"""Return an ssl.SSLContext for a server using our self-signed cert."""
|
||||
cert_p = cert_path or _CERT_PATH
|
||||
key_p = key_path or _KEY_PATH
|
||||
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
ctx.load_cert_chain(certfile=str(cert_p), keyfile=str(key_p))
|
||||
return ctx
|
||||
|
||||
|
||||
def make_client_ssl_context(verify: bool = False) -> ssl.SSLContext:
|
||||
"""Return a client SSLContext. verify=False for self-signed TOFU connections."""
|
||||
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||||
if not verify:
|
||||
ctx.check_hostname = False
|
||||
ctx.verify_mode = ssl.CERT_NONE
|
||||
return ctx
|
||||
@@ -8,6 +8,14 @@ version = "0.1.0"
|
||||
description = "Distributed Inference Network gossip and shard swarm"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
dependencies = [
|
||||
"cryptography>=41",
|
||||
"websockets>=13",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
mdns = ["zeroconf>=0.131"]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["meshnet_p2p*"]
|
||||
|
||||
10
packages/p2p/relay_bootstrap.json
Normal file
10
packages/p2p/relay_bootstrap.json
Normal file
@@ -0,0 +1,10 @@
|
||||
{
|
||||
"relays": [
|
||||
{
|
||||
"url": "ws://localhost:8765",
|
||||
"cert_fingerprint": null,
|
||||
"operator": "localhost-dev",
|
||||
"note": "Local development relay — replace with team relay URL before production"
|
||||
}
|
||||
]
|
||||
}
|
||||
3
packages/relay/meshnet_relay/__init__.py
Normal file
3
packages/relay/meshnet_relay/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""meshnet-relay — NAT-traversal relay and gossip hub."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
60
packages/relay/meshnet_relay/cli.py
Normal file
60
packages/relay/meshnet_relay/cli.py
Normal file
@@ -0,0 +1,60 @@
|
||||
"""meshnet-relay CLI entry point."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="meshnet-relay",
|
||||
description="Meshnet NAT-traversal relay and gossip hub",
|
||||
)
|
||||
parser.add_argument("--host", default="0.0.0.0", help="Interface to bind")
|
||||
parser.add_argument("--port", type=int, default=8765, help="WebSocket port")
|
||||
parser.add_argument("--cert", metavar="PATH", help="TLS certificate (PEM)")
|
||||
parser.add_argument("--key", metavar="PATH", help="TLS private key (PEM)")
|
||||
parser.add_argument("--max-peers", type=int, default=500, help="Max concurrent peers")
|
||||
parser.add_argument("--log-level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"])
|
||||
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, args.log_level),
|
||||
format="%(asctime)s %(levelname)-8s %(name)s %(message)s",
|
||||
)
|
||||
|
||||
from .server import RelayServer
|
||||
|
||||
ssl_cert = Path(args.cert) if args.cert else None
|
||||
ssl_key = Path(args.key) if args.key else None
|
||||
|
||||
server = RelayServer(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
ssl_cert=ssl_cert,
|
||||
ssl_key=ssl_key,
|
||||
max_peers=args.max_peers,
|
||||
)
|
||||
port = server.start()
|
||||
scheme = "wss" if ssl_cert else "ws"
|
||||
print(f"meshnet-relay listening on {scheme}://{args.host}:{port}", flush=True)
|
||||
print(" /ws gossip PubSub", flush=True)
|
||||
print(" /relay/<id> circuit relay to peer", flush=True)
|
||||
print(" /health health check", flush=True)
|
||||
print(" /v1/peers peer list", flush=True)
|
||||
print("Press Ctrl-C to stop.", flush=True)
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopping relay…", flush=True)
|
||||
server.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
60
packages/relay/meshnet_relay/peer_registry.py
Normal file
60
packages/relay/meshnet_relay/peer_registry.py
Normal file
@@ -0,0 +1,60 @@
|
||||
"""In-memory registry of connected gossip peers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class PeerEntry:
|
||||
peer_id: str
|
||||
addr: str
|
||||
ws: Any # websockets.WebSocketServerProtocol
|
||||
connected_at: float = field(default_factory=time.monotonic)
|
||||
last_seen: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
class PeerRegistry:
|
||||
def __init__(self):
|
||||
self._peers: dict[str, PeerEntry] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def register(self, peer_id: str, addr: str, ws) -> None:
|
||||
with self._lock:
|
||||
self._peers[peer_id] = PeerEntry(peer_id=peer_id, addr=addr, ws=ws)
|
||||
|
||||
def unregister(self, peer_id: str) -> None:
|
||||
with self._lock:
|
||||
self._peers.pop(peer_id, None)
|
||||
|
||||
def touch(self, peer_id: str) -> None:
|
||||
with self._lock:
|
||||
if peer_id in self._peers:
|
||||
self._peers[peer_id].last_seen = time.monotonic()
|
||||
|
||||
def get(self, peer_id: str) -> PeerEntry | None:
|
||||
with self._lock:
|
||||
return self._peers.get(peer_id)
|
||||
|
||||
def all_except(self, peer_id: str) -> list[PeerEntry]:
|
||||
with self._lock:
|
||||
return [e for pid, e in self._peers.items() if pid != peer_id]
|
||||
|
||||
def list_peers(self) -> list[dict]:
|
||||
with self._lock:
|
||||
return [
|
||||
{
|
||||
"peer_id": e.peer_id,
|
||||
"addr": e.addr,
|
||||
"connected_at": e.connected_at,
|
||||
"last_seen": e.last_seen,
|
||||
}
|
||||
for e in self._peers.values()
|
||||
]
|
||||
|
||||
def __len__(self) -> int:
|
||||
with self._lock:
|
||||
return len(self._peers)
|
||||
224
packages/relay/meshnet_relay/server.py
Normal file
224
packages/relay/meshnet_relay/server.py
Normal file
@@ -0,0 +1,224 @@
|
||||
"""Relay server — WebSocket gossip hub + circuit relay proxy.
|
||||
|
||||
HTTP API (served via asyncio-based handler on same port):
|
||||
GET /health → {"status": "ok", "peers": N}
|
||||
GET /v1/peers → [{peer_id, addr, last_seen}]
|
||||
POST /v1/gossip → accept a gossip envelope, fan out to connected peers
|
||||
|
||||
WebSocket endpoints:
|
||||
ws[s]://host:port/ws → gossip PubSub connection
|
||||
ws[s]://host:port/relay/{peer_id} → circuit relay to that peer
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from .peer_registry import PeerRegistry
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RelayServer:
|
||||
"""Async WebSocket relay server that runs in a background thread.
|
||||
|
||||
Usage::
|
||||
|
||||
server = RelayServer(host="0.0.0.0", port=8765)
|
||||
port = server.start() # returns actual port
|
||||
...
|
||||
server.stop()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "0.0.0.0",
|
||||
port: int = 8765,
|
||||
ssl_cert: Path | None = None,
|
||||
ssl_key: Path | None = None,
|
||||
max_peers: int = 500,
|
||||
):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.ssl_cert = ssl_cert
|
||||
self.ssl_key = ssl_key
|
||||
self.max_peers = max_peers
|
||||
|
||||
self._registry = PeerRegistry()
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self._server = None
|
||||
self._actual_port: int = port
|
||||
self._ready = threading.Event()
|
||||
self._running = False
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
|
||||
@property
|
||||
def registry(self) -> PeerRegistry:
|
||||
return self._registry
|
||||
|
||||
def start(self) -> int:
|
||||
"""Start server in background thread. Returns actual bound port."""
|
||||
self._running = True
|
||||
self._loop = asyncio.new_event_loop()
|
||||
self._thread = threading.Thread(target=self._run, daemon=True, name="relay")
|
||||
self._thread.start()
|
||||
self._ready.wait(timeout=5)
|
||||
return self._actual_port
|
||||
|
||||
def stop(self) -> None:
|
||||
self._running = False
|
||||
if self._loop and self._stop_event is not None:
|
||||
self._loop.call_soon_threadsafe(self._stop_event.set)
|
||||
if self._thread:
|
||||
self._thread.join(timeout=3.0)
|
||||
|
||||
def _run(self) -> None:
|
||||
asyncio.set_event_loop(self._loop)
|
||||
self._loop.run_until_complete(self._serve())
|
||||
|
||||
async def _serve(self) -> None:
|
||||
import websockets # type: ignore[import]
|
||||
import websockets.server # type: ignore[import]
|
||||
|
||||
ssl_ctx = None
|
||||
if self.ssl_cert and self.ssl_key:
|
||||
import ssl
|
||||
ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
||||
ssl_ctx.load_cert_chain(str(self.ssl_cert), str(self.ssl_key))
|
||||
|
||||
server = await websockets.serve(
|
||||
self._handle_connection,
|
||||
self.host,
|
||||
self.port,
|
||||
ssl=ssl_ctx,
|
||||
)
|
||||
# Record actual port after bind
|
||||
for sock in server.sockets or []:
|
||||
self._actual_port = sock.getsockname()[1]
|
||||
break
|
||||
|
||||
self._stop_event = asyncio.Event()
|
||||
self._server = server
|
||||
self._ready.set()
|
||||
log.info("Relay listening on %s:%d", self.host, self._actual_port)
|
||||
|
||||
await self._stop_event.wait()
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
async def _handle_connection(self, ws) -> None:
|
||||
"""Dispatch incoming WebSocket to gossip hub or circuit relay."""
|
||||
try:
|
||||
path = ws.request.path
|
||||
except AttributeError:
|
||||
path = getattr(ws, "path", "/ws")
|
||||
|
||||
if path.startswith("/relay/"):
|
||||
peer_id = path[len("/relay/"):]
|
||||
await self._handle_circuit_relay(ws, peer_id)
|
||||
elif path == "/health":
|
||||
await ws.send(json.dumps({"status": "ok", "peers": len(self._registry)}))
|
||||
await ws.close()
|
||||
elif path == "/v1/peers":
|
||||
await ws.send(json.dumps(self._registry.list_peers()))
|
||||
await ws.close()
|
||||
else:
|
||||
await self._handle_gossip(ws)
|
||||
|
||||
async def _handle_gossip(self, ws) -> None:
|
||||
"""Accept a gossip peer connection, register it, and fan out messages."""
|
||||
peer_id: str | None = None
|
||||
peer_addr: str = ""
|
||||
|
||||
try:
|
||||
async for raw in ws:
|
||||
try:
|
||||
envelope = json.loads(raw)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
|
||||
topic = envelope.get("topic", "")
|
||||
from_peer = envelope.get("from_peer", "")
|
||||
|
||||
# Handle peer registration message
|
||||
if topic == "peer-register":
|
||||
payload = envelope.get("payload", {})
|
||||
peer_id = payload.get("peer_id") or from_peer
|
||||
peer_addr = payload.get("addr", "")
|
||||
if len(self._registry) >= self.max_peers:
|
||||
await ws.close(1008, "relay at capacity")
|
||||
return
|
||||
self._registry.register(peer_id, peer_addr, ws)
|
||||
log.debug("Peer registered: %s", peer_id)
|
||||
# Send current peer list back
|
||||
await ws.send(json.dumps({
|
||||
"topic": "peer-list",
|
||||
"version": 1,
|
||||
"from_peer": "relay",
|
||||
"payload": {"peers": self._registry.list_peers()},
|
||||
}))
|
||||
continue
|
||||
|
||||
# Fan out to all other registered peers
|
||||
if peer_id:
|
||||
self._registry.touch(peer_id)
|
||||
fan_out_peers = self._registry.all_except(peer_id or "")
|
||||
await _broadcast(raw, fan_out_peers)
|
||||
|
||||
except Exception as exc:
|
||||
log.debug("Gossip connection error: %s", exc)
|
||||
finally:
|
||||
if peer_id:
|
||||
self._registry.unregister(peer_id)
|
||||
log.debug("Peer unregistered: %s", peer_id)
|
||||
|
||||
async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None:
|
||||
"""Proxy WebSocket traffic between ws_requester and target_peer_id's ws."""
|
||||
target = self._registry.get(target_peer_id)
|
||||
if target is None:
|
||||
try:
|
||||
await ws_requester.send(json.dumps({
|
||||
"error": f"peer {target_peer_id!r} not connected to relay"
|
||||
}))
|
||||
await ws_requester.close(1011, "target peer not found")
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
log.debug("Circuit relay: ??? → %s", target_peer_id)
|
||||
|
||||
async def pipe(src, dst) -> None:
|
||||
try:
|
||||
async for msg in src:
|
||||
await dst.send(msg)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
await asyncio.gather(
|
||||
pipe(ws_requester, target.ws),
|
||||
pipe(target.ws, ws_requester),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
|
||||
async def _broadcast(raw: str | bytes, peers: list) -> None:
|
||||
"""Send raw message to all peers; ignore individual send failures."""
|
||||
if not peers:
|
||||
return
|
||||
import asyncio
|
||||
await asyncio.gather(
|
||||
*[_safe_send(p.ws, raw) for p in peers],
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
|
||||
async def _safe_send(ws, msg) -> None:
|
||||
try:
|
||||
await ws.send(msg)
|
||||
except Exception:
|
||||
pass
|
||||
20
packages/relay/pyproject.toml
Normal file
20
packages/relay/pyproject.toml
Normal file
@@ -0,0 +1,20 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=64"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "meshnet-relay"
|
||||
version = "0.1.0"
|
||||
description = "Distributed Inference Network NAT-traversal relay and gossip hub"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
dependencies = [
|
||||
"websockets>=13",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
meshnet-relay = "meshnet_relay.cli:main"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["meshnet_relay*"]
|
||||
@@ -56,6 +56,7 @@ class _NodeEntry:
|
||||
"score", "vram_bytes", "ram_bytes", "quantizations",
|
||||
"benchmark_tokens_per_sec", "quantization", "managed_assignment",
|
||||
"pending_directives", "last_heartbeat", "tracker_mode",
|
||||
"relay_addr", "cert_fingerprint", "peer_id",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
@@ -76,6 +77,9 @@ class _NodeEntry:
|
||||
quantization: str | None = None,
|
||||
managed_assignment: bool = False,
|
||||
tracker_mode: bool = False,
|
||||
relay_addr: str | None = None,
|
||||
cert_fingerprint: str | None = None,
|
||||
peer_id: str | None = None,
|
||||
) -> None:
|
||||
self.node_id = node_id
|
||||
self.endpoint = endpoint
|
||||
@@ -93,6 +97,9 @@ class _NodeEntry:
|
||||
self.quantization = quantization
|
||||
self.managed_assignment = managed_assignment
|
||||
self.tracker_mode = tracker_mode
|
||||
self.relay_addr = relay_addr
|
||||
self.cert_fingerprint = cert_fingerprint
|
||||
self.peer_id = peer_id
|
||||
self.pending_directives: list[dict] = []
|
||||
self.last_heartbeat: float = time.monotonic()
|
||||
|
||||
@@ -607,6 +614,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
return
|
||||
|
||||
tracker_mode = bool(body.get("tracker_mode", False))
|
||||
relay_addr = body.get("relay_addr") or None
|
||||
cert_fingerprint = body.get("cert_fingerprint") or None
|
||||
peer_id = body.get("peer_id") or None
|
||||
|
||||
node_id = str(uuid.uuid4())
|
||||
entry = _NodeEntry(
|
||||
@@ -626,6 +636,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
quantization=quantization,
|
||||
managed_assignment=not explicit_shard,
|
||||
tracker_mode=tracker_mode,
|
||||
relay_addr=relay_addr,
|
||||
cert_fingerprint=cert_fingerprint,
|
||||
peer_id=peer_id,
|
||||
)
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
@@ -643,6 +656,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
def _handle_heartbeat(self, node_id: str):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
body: dict = {}
|
||||
content_length = int(self.headers.get("Content-Length", 0))
|
||||
if content_length > 0:
|
||||
try:
|
||||
body = json.loads(self.rfile.read(content_length))
|
||||
except Exception:
|
||||
pass
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
entry = server.registry.get(node_id)
|
||||
@@ -650,6 +670,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._send_json(404, {"error": "node not found"})
|
||||
return
|
||||
entry.last_heartbeat = time.monotonic()
|
||||
if body.get("relay_addr"):
|
||||
entry.relay_addr = body["relay_addr"]
|
||||
if body.get("cert_fingerprint"):
|
||||
entry.cert_fingerprint = body["cert_fingerprint"]
|
||||
if body.get("peer_id"):
|
||||
entry.peer_id = body["peer_id"]
|
||||
_rebalance_model_locked(server, entry.model or "stub-model")
|
||||
directives = list(entry.pending_directives)
|
||||
entry.pending_directives.clear()
|
||||
|
||||
371
tests/test_gossip_and_relay.py
Normal file
371
tests/test_gossip_and_relay.py
Normal file
@@ -0,0 +1,371 @@
|
||||
"""Tests for US-017: P2P gossip, relay node, and TLS infrastructure."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# identity tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_load_or_create_identity_generates_peer_id(tmp_path):
|
||||
from meshnet_p2p.identity import load_or_create_identity
|
||||
|
||||
identity = load_or_create_identity(tmp_path / "identity.json")
|
||||
assert len(identity["peer_id"]) == 16
|
||||
assert "public_key_pem" in identity
|
||||
assert "private_key_pem" in identity
|
||||
|
||||
|
||||
def test_identity_is_stable_across_loads(tmp_path):
|
||||
from meshnet_p2p.identity import load_or_create_identity
|
||||
|
||||
path = tmp_path / "identity.json"
|
||||
first = load_or_create_identity(path)
|
||||
second = load_or_create_identity(path)
|
||||
assert first["peer_id"] == second["peer_id"]
|
||||
assert first["public_key_pem"] == second["public_key_pem"]
|
||||
|
||||
|
||||
def test_identity_different_for_different_paths(tmp_path):
|
||||
from meshnet_p2p.identity import load_or_create_identity
|
||||
|
||||
a = load_or_create_identity(tmp_path / "a.json")
|
||||
b = load_or_create_identity(tmp_path / "b.json")
|
||||
# Extremely unlikely to collide
|
||||
assert a["peer_id"] != b["peer_id"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TLS / certificate tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_generate_self_signed_cert_creates_files(tmp_path):
|
||||
from meshnet_p2p.tls import generate_self_signed_cert
|
||||
|
||||
cert_p, key_p = generate_self_signed_cert(
|
||||
cert_path=tmp_path / "cert.pem",
|
||||
key_path=tmp_path / "key.pem",
|
||||
common_name="localhost",
|
||||
)
|
||||
assert cert_p.exists()
|
||||
assert key_p.exists()
|
||||
assert cert_p.stat().st_size > 100
|
||||
assert key_p.stat().st_size > 100
|
||||
|
||||
|
||||
def test_generate_self_signed_cert_is_idempotent(tmp_path):
|
||||
from meshnet_p2p.tls import generate_self_signed_cert
|
||||
|
||||
args = dict(cert_path=tmp_path / "cert.pem", key_path=tmp_path / "key.pem", common_name="test")
|
||||
generate_self_signed_cert(**args)
|
||||
mtime1 = (tmp_path / "cert.pem").stat().st_mtime
|
||||
|
||||
generate_self_signed_cert(**args)
|
||||
mtime2 = (tmp_path / "cert.pem").stat().st_mtime
|
||||
assert mtime1 == mtime2 # file not regenerated
|
||||
|
||||
|
||||
def test_cert_fingerprint_returns_sha256_prefix(tmp_path):
|
||||
from meshnet_p2p.tls import generate_self_signed_cert, cert_fingerprint
|
||||
|
||||
cert_p, key_p = generate_self_signed_cert(
|
||||
cert_path=tmp_path / "cert.pem",
|
||||
key_path=tmp_path / "key.pem",
|
||||
common_name="test",
|
||||
)
|
||||
fp = cert_fingerprint(cert_p)
|
||||
assert fp.startswith("sha256:")
|
||||
assert len(fp) == len("sha256:") + 64 # 32 bytes hex = 64 chars
|
||||
|
||||
|
||||
def test_make_server_ssl_context_loads_cert(tmp_path):
|
||||
import ssl
|
||||
from meshnet_p2p.tls import generate_self_signed_cert, make_server_ssl_context
|
||||
|
||||
cert_p, key_p = generate_self_signed_cert(
|
||||
cert_path=tmp_path / "cert.pem",
|
||||
key_path=tmp_path / "key.pem",
|
||||
common_name="test",
|
||||
)
|
||||
ctx = make_server_ssl_context(cert_p, key_p)
|
||||
assert isinstance(ctx, ssl.SSLContext)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PeerRegistry tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_peer_registry_register_and_list():
|
||||
from meshnet_relay.peer_registry import PeerRegistry
|
||||
|
||||
reg = PeerRegistry()
|
||||
ws_mock = MagicMock()
|
||||
reg.register("peer1", "http://1.2.3.4:8001", ws_mock)
|
||||
|
||||
assert len(reg) == 1
|
||||
peers = reg.list_peers()
|
||||
assert len(peers) == 1
|
||||
assert peers[0]["peer_id"] == "peer1"
|
||||
|
||||
|
||||
def test_peer_registry_all_except_excludes_sender():
|
||||
from meshnet_relay.peer_registry import PeerRegistry
|
||||
|
||||
reg = PeerRegistry()
|
||||
reg.register("a", "http://a:8001", MagicMock())
|
||||
reg.register("b", "http://b:8001", MagicMock())
|
||||
reg.register("c", "http://c:8001", MagicMock())
|
||||
|
||||
others = reg.all_except("a")
|
||||
assert len(others) == 2
|
||||
assert all(e.peer_id != "a" for e in others)
|
||||
|
||||
|
||||
def test_peer_registry_unregister_removes_peer():
|
||||
from meshnet_relay.peer_registry import PeerRegistry
|
||||
|
||||
reg = PeerRegistry()
|
||||
reg.register("x", "http://x:8001", MagicMock())
|
||||
reg.unregister("x")
|
||||
assert len(reg) == 0
|
||||
assert reg.get("x") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GossipClient + RelayServer integration test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _start_relay(host="127.0.0.1", port=0):
|
||||
from meshnet_relay.server import RelayServer
|
||||
server = RelayServer(host=host, port=port)
|
||||
actual_port = server.start()
|
||||
return server, actual_port
|
||||
|
||||
|
||||
def test_gossip_fanout_through_relay():
|
||||
"""Node B publishes node-join; node A receives it within 2 seconds."""
|
||||
from meshnet_p2p.gossip import GossipClient, TOPIC_NODE_JOIN
|
||||
|
||||
relay, port = _start_relay()
|
||||
relay_url = f"ws://127.0.0.1:{port}/ws"
|
||||
|
||||
client_a = GossipClient(relay_url=relay_url, peer_id="peer_a")
|
||||
client_b = GossipClient(relay_url=relay_url, peer_id="peer_b")
|
||||
|
||||
received = []
|
||||
client_a.subscribe(TOPIC_NODE_JOIN, lambda env: received.append(env))
|
||||
|
||||
client_a.start()
|
||||
client_b.start()
|
||||
|
||||
assert client_a.wait_connected(timeout=5), "client_a failed to connect to relay"
|
||||
assert client_b.wait_connected(timeout=5), "client_b failed to connect to relay"
|
||||
|
||||
# Give both peers time to register with relay
|
||||
time.sleep(0.2)
|
||||
|
||||
client_b.publish(TOPIC_NODE_JOIN, {"addr": "http://192.168.1.10:8001", "peer_id": "peer_b"})
|
||||
|
||||
deadline = time.monotonic() + 2.0
|
||||
while time.monotonic() < deadline and not received:
|
||||
time.sleep(0.05)
|
||||
|
||||
client_a.stop()
|
||||
client_b.stop()
|
||||
relay.stop()
|
||||
|
||||
assert received, "client_a did not receive node-join message from client_b"
|
||||
assert received[0]["topic"] == TOPIC_NODE_JOIN
|
||||
assert received[0]["from_peer"] == "peer_b"
|
||||
|
||||
|
||||
def test_gossip_dedup_prevents_processing_duplicate_message_ids():
|
||||
"""A message with a duplicate msg_id is only processed once."""
|
||||
from meshnet_p2p.gossip import GossipClient, TOPIC_NODE_JOIN
|
||||
|
||||
relay, port = _start_relay()
|
||||
relay_url = f"ws://127.0.0.1:{port}/ws"
|
||||
|
||||
client_a = GossipClient(relay_url=relay_url, peer_id="peer_a2")
|
||||
client_b = GossipClient(relay_url=relay_url, peer_id="peer_b2")
|
||||
|
||||
received = []
|
||||
client_a.subscribe(TOPIC_NODE_JOIN, lambda env: received.append(env))
|
||||
|
||||
client_a.start()
|
||||
client_b.start()
|
||||
client_a.wait_connected(5)
|
||||
client_b.wait_connected(5)
|
||||
time.sleep(0.2)
|
||||
|
||||
# Publish once
|
||||
client_b.publish(TOPIC_NODE_JOIN, {"test": "dedup"})
|
||||
time.sleep(0.5)
|
||||
|
||||
count = len(received)
|
||||
|
||||
client_a.stop()
|
||||
client_b.stop()
|
||||
relay.stop()
|
||||
|
||||
# Should have received exactly one message (not duplicated)
|
||||
assert count <= 1
|
||||
|
||||
|
||||
def test_relay_server_peer_list_grows_on_connect():
|
||||
"""Relay registry grows when clients connect."""
|
||||
from meshnet_p2p.gossip import GossipClient
|
||||
|
||||
relay, port = _start_relay()
|
||||
relay_url = f"ws://127.0.0.1:{port}/ws"
|
||||
|
||||
client = GossipClient(relay_url=relay_url, peer_id="solo_peer")
|
||||
client.start()
|
||||
client.wait_connected(5)
|
||||
time.sleep(0.3) # let peer-register message process
|
||||
|
||||
peer_count = len(relay.registry)
|
||||
client.stop()
|
||||
relay.stop()
|
||||
|
||||
assert peer_count >= 1
|
||||
|
||||
|
||||
def test_relay_circuit_relay_proxies_message():
|
||||
"""A node behind NAT (client_a) receives a message via circuit relay from client_b."""
|
||||
import websockets.sync.client # type: ignore[import]
|
||||
from meshnet_relay.server import RelayServer
|
||||
|
||||
relay = RelayServer(host="127.0.0.1", port=0)
|
||||
port = relay.start()
|
||||
|
||||
# client_a connects to gossip hub and registers
|
||||
received_via_relay = []
|
||||
ready = threading.Event()
|
||||
|
||||
def run_nat_client():
|
||||
import websockets.sync.client as wsc
|
||||
with wsc.connect(f"ws://127.0.0.1:{port}/ws") as ws:
|
||||
ws.send(json.dumps({
|
||||
"topic": "peer-register",
|
||||
"version": 1,
|
||||
"from_peer": "nat_peer",
|
||||
"msg_id": "reg-001",
|
||||
"timestamp": "2026-06-29T00:00:00Z",
|
||||
"ttl": 3,
|
||||
"payload": {"peer_id": "nat_peer", "addr": ""},
|
||||
}))
|
||||
# consume peer-list response
|
||||
ws.recv()
|
||||
ready.set()
|
||||
# wait for a relayed message
|
||||
try:
|
||||
msg = ws.recv(timeout=3)
|
||||
received_via_relay.append(json.loads(msg))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
nat_thread = threading.Thread(target=run_nat_client, daemon=True)
|
||||
nat_thread.start()
|
||||
ready.wait(timeout=5)
|
||||
time.sleep(0.1) # ensure peer is in registry
|
||||
|
||||
# client_b connects via circuit relay path
|
||||
def send_via_relay():
|
||||
try:
|
||||
import websockets.sync.client as wsc
|
||||
with wsc.connect(f"ws://127.0.0.1:{port}/relay/nat_peer") as ws:
|
||||
ws.send(json.dumps({"topic": "direct-relay-test", "payload": {"hi": "there"}}))
|
||||
time.sleep(0.3)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
relay_thread = threading.Thread(target=send_via_relay, daemon=True)
|
||||
relay_thread.start()
|
||||
relay_thread.join(timeout=3)
|
||||
nat_thread.join(timeout=3)
|
||||
relay.stop()
|
||||
|
||||
assert received_via_relay, "NAT'd peer did not receive message via circuit relay"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tracker gossip fields tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _start_tracker_and_register(extra_fields: dict) -> dict:
|
||||
"""Helper: start tracker, register node with extra gossip fields, return response."""
|
||||
import http.server
|
||||
import json as _json
|
||||
import urllib.request
|
||||
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
|
||||
tracker = TrackerServer(host="127.0.0.1", port=0)
|
||||
port = tracker.start()
|
||||
url = f"http://127.0.0.1:{port}"
|
||||
|
||||
payload = {
|
||||
"endpoint": f"http://127.0.0.1:8001",
|
||||
"shard_start": 0,
|
||||
"shard_end": 7,
|
||||
"model": "stub-model",
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
**extra_fields,
|
||||
}
|
||||
data = _json.dumps(payload).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{url}/v1/nodes/register",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as r:
|
||||
resp = _json.loads(r.read())
|
||||
|
||||
tracker.stop()
|
||||
return resp
|
||||
|
||||
|
||||
def test_tracker_accepts_relay_addr_in_registration():
|
||||
resp = _start_tracker_and_register({
|
||||
"relay_addr": "ws://relay.meshnet.ai:8765/relay/abc123",
|
||||
"cert_fingerprint": "sha256:deadbeef",
|
||||
"peer_id": "abc123def456ef01",
|
||||
})
|
||||
assert "node_id" in resp
|
||||
|
||||
|
||||
def test_tracker_accepts_registration_without_gossip_fields():
|
||||
"""Existing registrations without P2P fields still work."""
|
||||
resp = _start_tracker_and_register({})
|
||||
assert "node_id" in resp
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# mDNS (no-op without zeroconf installed)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_mdns_discovery_is_available_flag():
|
||||
from meshnet_p2p.mdns import MdnsDiscovery
|
||||
|
||||
disc = MdnsDiscovery(peer_id="test", port=8001)
|
||||
# is_available() should be bool regardless of zeroconf install status
|
||||
assert isinstance(disc.is_available(), bool)
|
||||
|
||||
|
||||
def test_mdns_start_and_stop_without_zeroconf(monkeypatch):
|
||||
from meshnet_p2p import mdns as mdns_mod
|
||||
monkeypatch.setattr(mdns_mod, "_HAS_ZEROCONF", False)
|
||||
from meshnet_p2p.mdns import MdnsDiscovery
|
||||
|
||||
disc = MdnsDiscovery(peer_id="x", port=8001)
|
||||
disc.start() # should not raise
|
||||
disc.stop() # should not raise
|
||||
356
tests/test_mining_cli.py
Normal file
356
tests/test_mining_cli.py
Normal file
@@ -0,0 +1,356 @@
|
||||
"""Tests for US-016: mining-style node startup CLI + live dashboard."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# model_catalog tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_curated_models_list_is_non_empty():
|
||||
from meshnet_node.model_catalog import CURATED_MODELS
|
||||
assert len(CURATED_MODELS) >= 5
|
||||
|
||||
|
||||
def test_model_preset_vram_for_quant():
|
||||
from meshnet_node.model_catalog import CURATED_MODELS
|
||||
|
||||
m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)
|
||||
assert m.vram_for_quant("nf4") == m.vram_nf4
|
||||
assert m.vram_for_quant("int8") == m.vram_int8
|
||||
assert m.vram_for_quant("bf16") == m.vram_bf16
|
||||
assert m.vram_for_quant("bfloat16") == m.vram_bf16 # alias
|
||||
|
||||
|
||||
def test_model_preset_fits_vram():
|
||||
from meshnet_node.model_catalog import CURATED_MODELS
|
||||
|
||||
small = next(m for m in CURATED_MODELS if m.vram_nf4 < 10)
|
||||
assert small.fits_vram(small.vram_nf4, "nf4")
|
||||
assert not small.fits_vram(small.vram_nf4 - 1, "nf4")
|
||||
|
||||
|
||||
def test_recommended_quant_respects_vram():
|
||||
from meshnet_node.model_catalog import CURATED_MODELS
|
||||
|
||||
m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)
|
||||
# nf4=18, int8=40, bf16=140
|
||||
assert m.recommended_quant(200) == "bf16"
|
||||
assert m.recommended_quant(50) == "int8"
|
||||
assert m.recommended_quant(20) == "nf4"
|
||||
assert m.recommended_quant(5) is None
|
||||
|
||||
|
||||
def test_models_with_insufficient_vram_are_marked(monkeypatch):
|
||||
from meshnet_node import wizard as wiz
|
||||
|
||||
# Simulate 6 GB GPU
|
||||
gpus = [{"index": 0, "name": "RTX 3060", "vram_gb": 6.0, "backend": "cuda"}]
|
||||
monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus)
|
||||
|
||||
# Phi-3 at NF4 needs 4 GB — should fit; Llama-3-70B at NF4 needs 18 GB — should not
|
||||
from meshnet_node.model_catalog import CURATED_MODELS
|
||||
|
||||
phi = next(m for m in CURATED_MODELS if "Phi-3" in m.name)
|
||||
llama = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name)
|
||||
|
||||
assert phi.fits_vram(6.0, "nf4")
|
||||
assert not llama.fits_vram(6.0, "nf4")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# config tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_load_config_returns_none_when_missing(tmp_path):
|
||||
from meshnet_node.config import load_config
|
||||
assert load_config(tmp_path / "nonexistent.json") is None
|
||||
|
||||
|
||||
def test_save_and_load_config_roundtrip(tmp_path):
|
||||
from meshnet_node.config import save_config, load_config
|
||||
|
||||
cfg = {"model_hf_repo": "test/model", "quantization": "nf4", "tracker_url": "http://localhost:8080"}
|
||||
cfg_path = tmp_path / "config.json"
|
||||
save_config(cfg, cfg_path)
|
||||
|
||||
loaded = load_config(cfg_path)
|
||||
assert loaded == cfg
|
||||
|
||||
|
||||
def test_save_config_creates_parent_dirs(tmp_path):
|
||||
from meshnet_node.config import save_config, load_config
|
||||
|
||||
nested = tmp_path / "deep" / "nested" / "config.json"
|
||||
save_config({"x": 1}, nested)
|
||||
assert nested.exists()
|
||||
assert load_config(nested) == {"x": 1}
|
||||
|
||||
|
||||
def test_merge_cli_overrides_applies_non_none_values():
|
||||
from meshnet_node.config import merge_cli_overrides
|
||||
|
||||
base = {"tracker_url": "http://a:8080", "quantization": "nf4", "port": 7000}
|
||||
result = merge_cli_overrides(base, tracker_url="http://b:9090", port=None)
|
||||
assert result["tracker_url"] == "http://b:9090"
|
||||
assert result["port"] == 7000 # None override ignored
|
||||
assert result["quantization"] == "nf4" # unchanged
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# wizard tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_print_models_table_runs_without_error(capsys, monkeypatch):
|
||||
from meshnet_node import wizard as wiz
|
||||
|
||||
monkeypatch.setattr(wiz, "_detect_gpus", lambda: [{"index": 0, "name": "GPU", "vram_gb": 24.0, "backend": "cuda"}])
|
||||
wiz.print_models_table()
|
||||
out = capsys.readouterr().out
|
||||
assert "Llama" in out or "Qwen" in out or "Phi" in out
|
||||
|
||||
|
||||
def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch):
|
||||
from meshnet_node import wizard as wiz
|
||||
from meshnet_node.config import load_config, save_config
|
||||
|
||||
# Fake GPU
|
||||
gpus = [{"index": 0, "name": "RTX 4090", "vram_gb": 24.0, "backend": "cuda"}]
|
||||
monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus)
|
||||
# Tracker not reachable (stub)
|
||||
monkeypatch.setattr(wiz, "_ping_tracker", lambda url: False)
|
||||
|
||||
# Simulate user selecting model 1 (Qwen2.5-0.5B), quant 1 (nf4), default dir, default tracker, default wallet
|
||||
inputs = iter([
|
||||
"1", # pick Qwen2.5-0.5B-Instruct (index 1 in CURATED_MODELS)
|
||||
"1", # quant NF4
|
||||
str(tmp_path / "models"), # download dir
|
||||
"http://localhost:8080", # tracker
|
||||
str(tmp_path / "wallet.json"), # wallet
|
||||
])
|
||||
monkeypatch.setattr("builtins.input", lambda prompt="": next(inputs))
|
||||
|
||||
cfg = wiz.run_wizard(config_path_override=tmp_path / "config.json")
|
||||
assert cfg["model_hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
assert cfg["quantization"] == "nf4"
|
||||
assert "download_dir" in cfg
|
||||
assert cfg["tracker_url"] == "http://localhost:8080"
|
||||
|
||||
|
||||
def test_wizard_raises_keyboard_interrupt_on_ctrl_c(monkeypatch):
|
||||
from meshnet_node import wizard as wiz
|
||||
|
||||
gpus = [{"index": 0, "name": "RTX 4090", "vram_gb": 24.0, "backend": "cuda"}]
|
||||
monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus)
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def fake_input(prompt=""):
|
||||
call_count[0] += 1
|
||||
if call_count[0] == 1:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
monkeypatch.setattr("builtins.input", fake_input)
|
||||
|
||||
import pytest
|
||||
with pytest.raises(KeyboardInterrupt):
|
||||
wiz.run_wizard()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# dashboard tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_is_interactive_tty_false_when_not_tty(monkeypatch):
|
||||
from meshnet_node import dashboard as dash
|
||||
|
||||
monkeypatch.setattr(sys.stdout, "isatty", lambda: False)
|
||||
assert not dash.is_interactive_tty()
|
||||
|
||||
|
||||
def test_dashboard_plain_fallback_on_keyboard_interrupt(monkeypatch):
|
||||
"""Plain loop exits cleanly when Ctrl-C is raised."""
|
||||
from meshnet_node import dashboard as dash
|
||||
|
||||
node = MagicMock()
|
||||
node.chat_completion_count = 5
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def fake_sleep(t):
|
||||
call_count[0] += 1
|
||||
if call_count[0] >= 1:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
monkeypatch.setattr(dash.time, "sleep", fake_sleep)
|
||||
monkeypatch.setattr(dash, "_gpu_stats", lambda: [])
|
||||
monkeypatch.setattr(sys.stdout, "isatty", lambda: False)
|
||||
|
||||
cfg = {"model_name": "test-model", "quantization": "nf4"}
|
||||
# Should not raise
|
||||
dash.run_dashboard(node, cfg, start_time=dash.time.monotonic())
|
||||
|
||||
|
||||
def test_ema_updates_correctly():
|
||||
from meshnet_node.dashboard import _EMA
|
||||
|
||||
ema = _EMA(alpha=1.0) # alpha=1.0 → always takes latest sample
|
||||
ema.update(10.0)
|
||||
assert ema.value == 10.0
|
||||
ema.update(20.0)
|
||||
assert ema.value == 20.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI integration tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_models_command_prints_table(capsys, monkeypatch):
|
||||
"""meshnet-node models prints the curated table and exits 0."""
|
||||
from meshnet_node import wizard as wiz
|
||||
|
||||
monkeypatch.setattr(wiz, "_detect_gpus", lambda: [])
|
||||
|
||||
from meshnet_node.cli import main
|
||||
monkeypatch.setattr(sys, "argv", ["meshnet-node", "models"])
|
||||
|
||||
try:
|
||||
main()
|
||||
except SystemExit as exc:
|
||||
assert exc.code == 0
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert "Llama" in out or "Qwen" in out or "Phi" in out
|
||||
|
||||
|
||||
def test_config_command_no_config_exits_1(tmp_path, monkeypatch):
|
||||
from meshnet_node import config as cfg_mod
|
||||
from meshnet_node.cli import main
|
||||
|
||||
monkeypatch.setattr(cfg_mod, "_DEFAULT_CONFIG_FILE", tmp_path / "nonexistent.json")
|
||||
monkeypatch.setattr(sys, "argv", ["meshnet-node", "config"])
|
||||
|
||||
with patch("meshnet_node.config.config_path", return_value=tmp_path / "nonexistent.json"):
|
||||
try:
|
||||
main()
|
||||
except SystemExit as exc:
|
||||
assert exc.code == 1
|
||||
|
||||
|
||||
def test_config_command_prints_saved_config(tmp_path, monkeypatch, capsys):
|
||||
from meshnet_node import config as cfg_mod
|
||||
from meshnet_node.config import save_config
|
||||
from meshnet_node.cli import main
|
||||
|
||||
saved = {"model_hf_repo": "meta-llama/Meta-Llama-3-70B-Instruct", "quantization": "nf4"}
|
||||
cfg_file = tmp_path / "config.json"
|
||||
save_config(saved, cfg_file)
|
||||
|
||||
monkeypatch.setattr(sys, "argv", ["meshnet-node", "config"])
|
||||
|
||||
with patch("meshnet_node.config.config_path", return_value=cfg_file):
|
||||
with patch("meshnet_node.config.load_config", return_value=saved):
|
||||
try:
|
||||
main()
|
||||
except SystemExit as exc:
|
||||
assert exc.code == 0
|
||||
|
||||
out = capsys.readouterr().out
|
||||
data = json.loads(out.split("\n", 1)[1]) # skip the "Config: ..." header line
|
||||
assert data["model_hf_repo"] == saved["model_hf_repo"]
|
||||
|
||||
|
||||
def test_detect_num_layers_returns_catalog_value_without_network(monkeypatch):
|
||||
"""detect_num_layers uses the curated catalog first — no network call."""
|
||||
from meshnet_node.model_catalog import detect_num_layers
|
||||
|
||||
# Qwen2.5-0.5B is in the catalog with 24 layers
|
||||
layers = detect_num_layers("Qwen/Qwen2.5-0.5B-Instruct")
|
||||
assert layers == 24
|
||||
|
||||
|
||||
def test_detect_num_layers_returns_none_on_error(monkeypatch):
|
||||
from meshnet_node.model_catalog import detect_num_layers
|
||||
|
||||
# Monkeypatch AutoConfig to raise
|
||||
import meshnet_node.model_catalog as cat
|
||||
monkeypatch.setattr(cat, "detect_num_layers", lambda repo: None if "bad" in repo else detect_num_layers(repo))
|
||||
assert cat.detect_num_layers("bad/repo") is None
|
||||
|
||||
|
||||
def test_startup_auto_detects_shard_range(monkeypatch, tmp_path):
|
||||
"""When shard_start/end are None, startup reads layer count from catalog."""
|
||||
from meshnet_node import startup as su
|
||||
from meshnet_node.model_catalog import detect_num_layers
|
||||
|
||||
calls = []
|
||||
|
||||
def fake_detect(repo):
|
||||
calls.append(repo)
|
||||
return 24 # Qwen2.5-0.5B
|
||||
|
||||
monkeypatch.setattr(su, "_detect_num_layers", fake_detect)
|
||||
|
||||
# Fake hardware detection
|
||||
monkeypatch.setattr(su, "detect_hardware", lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0})
|
||||
|
||||
# Fake wallet
|
||||
monkeypatch.setattr(su, "load_or_create_wallet", lambda **kw: (None, None, "fake-wallet"))
|
||||
|
||||
# Fake TorchNodeServer
|
||||
class FakeNode:
|
||||
chat_completion_count = 0
|
||||
def start(self): return 9999
|
||||
def stop(self): pass
|
||||
|
||||
import meshnet_node.startup as su2
|
||||
monkeypatch.setattr(su2, "TorchNodeServer", lambda **kw: FakeNode())
|
||||
|
||||
node = su.run_startup(
|
||||
tracker_url="http://localhost:8080",
|
||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
# shard_start and shard_end intentionally omitted
|
||||
quantization="bfloat16",
|
||||
host="127.0.0.1",
|
||||
)
|
||||
assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"]
|
||||
assert isinstance(node, FakeNode)
|
||||
|
||||
|
||||
def test_legacy_start_subcommand_accepted(monkeypatch):
|
||||
"""meshnet-node start --tracker http://... does not crash on arg parsing."""
|
||||
from meshnet_node.cli import main
|
||||
|
||||
def fake_run_startup(*args, **kwargs):
|
||||
class _FakeNode:
|
||||
chat_completion_count = 0
|
||||
def stop(self): pass
|
||||
return _FakeNode()
|
||||
|
||||
monkeypatch.setattr(sys, "argv", [
|
||||
"meshnet-node", "start",
|
||||
"--tracker", "http://localhost:8080",
|
||||
"--model", "stub-model",
|
||||
"--port", "0",
|
||||
])
|
||||
|
||||
raised = []
|
||||
|
||||
def fake_sleep(t):
|
||||
raise KeyboardInterrupt
|
||||
|
||||
with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
|
||||
with patch("time.sleep", side_effect=fake_sleep):
|
||||
try:
|
||||
main()
|
||||
except SystemExit as exc:
|
||||
raised.append(exc.code)
|
||||
|
||||
# Exited (either 0 or via KeyboardInterrupt caught in _cmd_start)
|
||||
# The important thing is no unhandled exception from arg parsing
|
||||
Reference in New Issue
Block a user