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