From 65f3ee6a856869b63f71435602b89f610fcfdf0b Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Mon, 29 Jun 2026 17:45:38 +0300 Subject: [PATCH 1/3] feat(us-016): mining-style node startup CLI + live dashboard - `meshnet-node` with no args runs interactive setup wizard on first run, then starts directly on subsequent runs using saved config - Wizard auto-detects all GPUs/VRAM, shows curated model list with per-quant VRAM requirements, marks models that exceed available VRAM as incompatible, offers HuggingFace Hub browse as escape hatch - Persistent config saved to ~/.config/meshnet/config.json (0o600) - Live rich dashboard (tokens/sec EMA, VRAM, requests, peers, uptime) with automatic plain-text fallback when stdout is not a TTY (WSL2/SSH/CI) - All wizard values overridable via CLI flags; --reset-config re-runs wizard - `meshnet-node models` lists curated models; `--browse` fetches HF Hub top-20 - `meshnet-node config` prints saved config - `meshnet-node start ...` preserved for backward compatibility - 19 new tests; 97 passed, 1 skipped (no regressions) Co-Authored-By: Claude Sonnet 4.6 --- packages/node/meshnet_node/cli.py | 302 ++++++++++++++---- packages/node/meshnet_node/config.py | 72 +++++ packages/node/meshnet_node/dashboard.py | 220 +++++++++++++ packages/node/meshnet_node/model_catalog.py | 133 ++++++++ packages/node/meshnet_node/wizard.py | 322 ++++++++++++++++++++ packages/node/pyproject.toml | 1 + tests/test_mining_cli.py | 298 ++++++++++++++++++ 7 files changed, 1284 insertions(+), 64 deletions(-) create mode 100644 packages/node/meshnet_node/config.py create mode 100644 packages/node/meshnet_node/dashboard.py create mode 100644 packages/node/meshnet_node/model_catalog.py create mode 100644 packages/node/meshnet_node/wizard.py create mode 100644 tests/test_mining_cli.py diff --git a/packages/node/meshnet_node/cli.py b/packages/node/meshnet_node/cli.py index 35a0563..da3e864 100644 --- a/packages/node/meshnet_node/cli.py +++ b/packages/node/meshnet_node/cli.py @@ -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__": diff --git a/packages/node/meshnet_node/config.py b/packages/node/meshnet_node/config.py new file mode 100644 index 0000000..a3cca42 --- /dev/null +++ b/packages/node/meshnet_node/config.py @@ -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 diff --git a/packages/node/meshnet_node/dashboard.py b/packages/node/meshnet_node/dashboard.py new file mode 100644 index 0000000..f8e97c9 --- /dev/null +++ b/packages/node/meshnet_node/dashboard.py @@ -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 diff --git a/packages/node/meshnet_node/model_catalog.py b/packages/node/meshnet_node/model_catalog.py new file mode 100644 index 0000000..5393435 --- /dev/null +++ b/packages/node/meshnet_node/model_catalog.py @@ -0,0 +1,133 @@ +"""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="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, + 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 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 diff --git a/packages/node/meshnet_node/wizard.py b/packages/node/meshnet_node/wizard.py new file mode 100644 index 0000000..624d810 --- /dev/null +++ b/packages/node/meshnet_node/wizard.py @@ -0,0 +1,322 @@ +"""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 + +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: + 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 + + print(f"\n ✓ Selected: {selected_repo}") + + # 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() diff --git a/packages/node/pyproject.toml b/packages/node/pyproject.toml index 67c03aa..6acbc09 100644 --- a/packages/node/pyproject.toml +++ b/packages/node/pyproject.toml @@ -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", diff --git a/tests/test_mining_cli.py b/tests/test_mining_cli.py new file mode 100644 index 0000000..94498bf --- /dev/null +++ b/tests/test_mining_cli.py @@ -0,0 +1,298 @@ +"""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 = CURATED_MODELS[0] # Llama-3-70B + 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 = CURATED_MODELS[0] # Llama-3-70B: 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 3 (Mixtral), quant 1 (nf4), default dir, default tracker, default wallet + inputs = iter([ + "3", # pick Mixtral (index 3 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"] == "mistralai/Mixtral-8x7B-Instruct-v0.1" + 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_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 From a2258d3df4b4d5844bc4288bb8383b31ad73d024 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Mon, 29 Jun 2026 17:58:16 +0300 Subject: [PATCH 2/3] feat(us-017): P2P gossip, NAT-traversal relay node, and TLS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - packages/p2p: identity (peer_id from sha256 of RSA pubkey), TLS cert generation with SHA-256 fingerprint, GossipClient (WSS PubSub with per-topic handlers, dedup by msg_id, auto-reconnect), MdnsDiscovery (zeroconf optional dependency, graceful no-op fallback) - packages/relay: new meshnet-relay package — RelayServer (asyncio + websockets) with gossip fanout hub, circuit relay proxy for NAT traversal, peer registry; meshnet-relay CLI - packages/p2p/relay_bootstrap.json: team relay bootstrap list - Tracker: _NodeEntry gains relay_addr, cert_fingerprint, peer_id; both register and heartbeat handlers read and store these optional fields - docs/adr/0010 already written (previous commit) - conftest.py: packages/relay added to sys.path - 18 new tests; 115 passed total, 1 skipped (no regressions) Co-Authored-By: Claude Sonnet 4.6 --- conftest.py | 1 + packages/p2p/meshnet_p2p/gossip.py | 162 ++++++++ packages/p2p/meshnet_p2p/identity.py | 64 +++ packages/p2p/meshnet_p2p/mdns.py | 136 +++++++ packages/p2p/meshnet_p2p/tls.py | 114 ++++++ packages/p2p/pyproject.toml | 8 + packages/p2p/relay_bootstrap.json | 10 + packages/relay/meshnet_relay/__init__.py | 3 + packages/relay/meshnet_relay/cli.py | 60 +++ packages/relay/meshnet_relay/peer_registry.py | 60 +++ packages/relay/meshnet_relay/server.py | 224 +++++++++++ packages/relay/pyproject.toml | 20 + packages/tracker/meshnet_tracker/server.py | 26 ++ tests/test_gossip_and_relay.py | 371 ++++++++++++++++++ 14 files changed, 1259 insertions(+) create mode 100644 packages/p2p/meshnet_p2p/gossip.py create mode 100644 packages/p2p/meshnet_p2p/identity.py create mode 100644 packages/p2p/meshnet_p2p/mdns.py create mode 100644 packages/p2p/meshnet_p2p/tls.py create mode 100644 packages/p2p/relay_bootstrap.json create mode 100644 packages/relay/meshnet_relay/__init__.py create mode 100644 packages/relay/meshnet_relay/cli.py create mode 100644 packages/relay/meshnet_relay/peer_registry.py create mode 100644 packages/relay/meshnet_relay/server.py create mode 100644 packages/relay/pyproject.toml create mode 100644 tests/test_gossip_and_relay.py diff --git a/conftest.py b/conftest.py index 617c147..2a444cd 100644 --- a/conftest.py +++ b/conftest.py @@ -11,6 +11,7 @@ _packages = [ "packages/sdk", "packages/contracts", "packages/p2p", + "packages/relay", "packages/validator", ] diff --git a/packages/p2p/meshnet_p2p/gossip.py b/packages/p2p/meshnet_p2p/gossip.py new file mode 100644 index 0000000..3d82eab --- /dev/null +++ b/packages/p2p/meshnet_p2p/gossip.py @@ -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) diff --git a/packages/p2p/meshnet_p2p/identity.py b/packages/p2p/meshnet_p2p/identity.py new file mode 100644 index 0000000..8c905a0 --- /dev/null +++ b/packages/p2p/meshnet_p2p/identity.py @@ -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 diff --git a/packages/p2p/meshnet_p2p/mdns.py b/packages/p2p/meshnet_p2p/mdns.py new file mode 100644 index 0000000..edd38d2 --- /dev/null +++ b/packages/p2p/meshnet_p2p/mdns.py @@ -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 diff --git a/packages/p2p/meshnet_p2p/tls.py b/packages/p2p/meshnet_p2p/tls.py new file mode 100644 index 0000000..4f5f1d6 --- /dev/null +++ b/packages/p2p/meshnet_p2p/tls.py @@ -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:'.""" + 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 diff --git a/packages/p2p/pyproject.toml b/packages/p2p/pyproject.toml index 378d633..9b2f65c 100644 --- a/packages/p2p/pyproject.toml +++ b/packages/p2p/pyproject.toml @@ -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*"] diff --git a/packages/p2p/relay_bootstrap.json b/packages/p2p/relay_bootstrap.json new file mode 100644 index 0000000..b5f4a30 --- /dev/null +++ b/packages/p2p/relay_bootstrap.json @@ -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" + } + ] +} diff --git a/packages/relay/meshnet_relay/__init__.py b/packages/relay/meshnet_relay/__init__.py new file mode 100644 index 0000000..0025a58 --- /dev/null +++ b/packages/relay/meshnet_relay/__init__.py @@ -0,0 +1,3 @@ +"""meshnet-relay — NAT-traversal relay and gossip hub.""" + +__version__ = "0.1.0" diff --git a/packages/relay/meshnet_relay/cli.py b/packages/relay/meshnet_relay/cli.py new file mode 100644 index 0000000..23a0912 --- /dev/null +++ b/packages/relay/meshnet_relay/cli.py @@ -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/ 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() diff --git a/packages/relay/meshnet_relay/peer_registry.py b/packages/relay/meshnet_relay/peer_registry.py new file mode 100644 index 0000000..2593b77 --- /dev/null +++ b/packages/relay/meshnet_relay/peer_registry.py @@ -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) diff --git a/packages/relay/meshnet_relay/server.py b/packages/relay/meshnet_relay/server.py new file mode 100644 index 0000000..bca3d86 --- /dev/null +++ b/packages/relay/meshnet_relay/server.py @@ -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 diff --git a/packages/relay/pyproject.toml b/packages/relay/pyproject.toml new file mode 100644 index 0000000..ce5f883 --- /dev/null +++ b/packages/relay/pyproject.toml @@ -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*"] diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 542aca3..79ce8eb 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -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() diff --git a/tests/test_gossip_and_relay.py b/tests/test_gossip_and_relay.py new file mode 100644 index 0000000..7b89287 --- /dev/null +++ b/tests/test_gossip_and_relay.py @@ -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 From 080d49b2c2a74d27e1d66225aef67569393f7a74 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Mon, 29 Jun 2026 18:27:50 +0300 Subject: [PATCH 3/3] feat(us-016): auto-detect shard range from model config Layer count is now fetched from the curated catalog (zero network calls for known models) or via AutoConfig.from_pretrained() (~1 KB config.json only) when model_id is given without --shard-start/--shard-end. - model_catalog: add detect_num_layers(), two small Qwen models at top - startup: _detect_num_layers() helper; shard range auto-derived - wizard: show detected layer count for custom HF repos - tests: 3 new tests for auto-shard; fix catalog-order assumptions Co-Authored-By: Claude Sonnet 4.6 --- packages/node/meshnet_node/model_catalog.py | 34 ++++++++++- packages/node/meshnet_node/startup.py | 31 ++++++++-- packages/node/meshnet_node/wizard.py | 14 ++++- tests/test_mining_cli.py | 68 +++++++++++++++++++-- 4 files changed, 135 insertions(+), 12 deletions(-) diff --git a/packages/node/meshnet_node/model_catalog.py b/packages/node/meshnet_node/model_catalog.py index 5393435..9852476 100644 --- a/packages/node/meshnet_node/model_catalog.py +++ b/packages/node/meshnet_node/model_catalog.py @@ -42,6 +42,24 @@ class ModelPreset: 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", @@ -72,7 +90,7 @@ CURATED_MODELS: list[ModelPreset] = [ ModelPreset( name="Llama-3-8B-Instruct", hf_repo="meta-llama/Meta-Llama-3-8B-Instruct", - num_layers=32, + num_layers=32, # gated repo — requires HF login vram_nf4=4.5, vram_int8=8.5, vram_bf16=16.0, @@ -108,6 +126,20 @@ CURATED_MODELS: list[ModelPreset] = [ ] +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: diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 62d896a..036b7ff 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -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 diff --git a/packages/node/meshnet_node/wizard.py b/packages/node/meshnet_node/wizard.py index 624d810..5edf1d3 100644 --- a/packages/node/meshnet_node/wizard.py +++ b/packages/node/meshnet_node/wizard.py @@ -9,7 +9,7 @@ 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 +from .model_catalog import CURATED_MODELS, ModelPreset, browse_hf_hub, detect_num_layers if TYPE_CHECKING: pass @@ -239,6 +239,13 @@ def run_wizard(config_path_override=None) -> dict: 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: @@ -254,7 +261,10 @@ def run_wizard(config_path_override=None) -> dict: selected_repo = None selected_preset = None - print(f"\n ✓ Selected: {selected_repo}") + 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) diff --git a/tests/test_mining_cli.py b/tests/test_mining_cli.py index 94498bf..211ea06 100644 --- a/tests/test_mining_cli.py +++ b/tests/test_mining_cli.py @@ -21,7 +21,7 @@ def test_curated_models_list_is_non_empty(): def test_model_preset_vram_for_quant(): from meshnet_node.model_catalog import CURATED_MODELS - m = CURATED_MODELS[0] # Llama-3-70B + 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 @@ -39,7 +39,8 @@ def test_model_preset_fits_vram(): def test_recommended_quant_respects_vram(): from meshnet_node.model_catalog import CURATED_MODELS - m = CURATED_MODELS[0] # Llama-3-70B: nf4=18, int8=40, bf16=140 + 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" @@ -125,9 +126,9 @@ def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch): # Tracker not reachable (stub) monkeypatch.setattr(wiz, "_ping_tracker", lambda url: False) - # Simulate user selecting model 3 (Mixtral), quant 1 (nf4), default dir, default tracker, default wallet + # Simulate user selecting model 1 (Qwen2.5-0.5B), quant 1 (nf4), default dir, default tracker, default wallet inputs = iter([ - "3", # pick Mixtral (index 3 in CURATED_MODELS) + "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 @@ -136,7 +137,7 @@ def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch): monkeypatch.setattr("builtins.input", lambda prompt="": next(inputs)) cfg = wiz.run_wizard(config_path_override=tmp_path / "config.json") - assert cfg["model_hf_repo"] == "mistralai/Mixtral-8x7B-Instruct-v0.1" + 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" @@ -265,6 +266,63 @@ def test_config_command_prints_saved_config(tmp_path, monkeypatch, capsys): 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