neuron-tai/packages/node/meshnet_node/cli.py

"""meshnet-node CLI entry point — mining-style UX."""

from __future__ import annotations

import argparse
import os
import socket
import sys
import time
from pathlib import Path


def _load_env_file(path: Path) -> None:
    """Load simple KEY=VALUE pairs from an env file without overriding env vars."""
    if not path.exists():
        return
    try:
        lines = path.read_text().splitlines()
    except OSError:
        return
    for line in lines:
        text = line.strip()
        if not text or text.startswith("#"):
            continue
        if text.startswith("export "):
            text = text[len("export "):].strip()
        if "=" not in text:
            continue
        key, value = text.split("=", 1)
        key = key.strip()
        if not key or key in os.environ:
            continue
        value = value.strip()
        if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
            value = value[1:-1]
        os.environ[key] = value


def _load_env_defaults() -> None:
    """Load local and user-level node env defaults before config defaults are imported."""
    _load_env_file(Path.cwd() / ".env")
    _load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env")


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 None,
            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", "auto").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"),
            advertise_host=cfg.get("advertise_host"),
            route_timeout=float(cfg.get("route_timeout", 30.0)),
            vram_mb_override=cfg.get("vram_mb_override"),
            max_loaded_shards=int(cfg.get("max_loaded_shards", 1)),
            debug=bool(cfg.get("debug", False)),
            tracker_source_disabled=bool(cfg.get("tracker_source_disabled", False)),
            torch_threads=cfg.get("torch_threads"),
            torch_interop_threads=cfg.get("torch_interop_threads"),
        )
    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 _resolve_model_flags(
    model: str | None,
    model_id: str | None,
) -> tuple[str | None, str | None]:
    """Return (model_name, hf_repo_or_none) from --model / --model-id flags."""
    explicit = model_id or model
    if not explicit:
        return None, None
    if "/" in explicit:
        return explicit.split("/")[-1], explicit
    return explicit, None


def _first_available_port(host: str, start: int = 7000, attempts: int = 100) -> int:
    """Return the first TCP port bindable on host, starting at start."""
    bind_host = "" if host == "0.0.0.0" else host
    for port in range(start, start + attempts):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            try:
                sock.bind((bind_host, port))
            except OSError:
                continue
            return port
    raise OSError(f"no available TCP port in range {start}-{start + attempts - 1}")


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 = {}
    model_name, hf_repo = _resolve_model_flags(args.model, getattr(args, "model_id", None))
    if model_name is not None:
        overrides["model_name"] = model_name
        overrides["model_hf_repo"] = hf_repo or ""
    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 args.advertise_host:
        overrides["advertise_host"] = args.advertise_host
    if args.route_timeout != 30.0:
        overrides["route_timeout"] = args.route_timeout
    if getattr(args, "memory", None) is not None:
        overrides["vram_mb_override"] = args.memory
    if getattr(args, "max_shards", None) is not None:
        overrides["max_loaded_shards"] = args.max_shards
    if args.debug:
        overrides["debug"] = True
    if getattr(args, "tracker_source_disabled", False):
        overrides["tracker_source_disabled"] = True
    if getattr(args, "torch_threads", None) is not None:
        overrides["torch_threads"] = args.torch_threads
    if getattr(args, "torch_interop_threads", None) is not None:
        overrides["torch_interop_threads"] = args.torch_interop_threads

    if 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 DEFAULTS

    # Build a transient config from flags (don't write to disk)
    cfg = dict(DEFAULTS)
    if args.tracker:
        cfg["tracker_url"] = args.tracker
    cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
    model_name, hf_repo = _resolve_model_flags(
        args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or None,
        args.model_id,
    )
    if model_name is not None:
        cfg["model_name"] = model_name
        cfg["model_hf_repo"] = hf_repo or ""
    cfg["quantization"] = args.quantization
    cfg["host"] = args.host
    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") or None,
            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),
            route_timeout=getattr(args, "route_timeout", 30.0),
            vram_mb_override=getattr(args, "memory", None),
            max_loaded_shards=getattr(args, "max_shards", 1),
            debug=getattr(args, "debug", False),
            tracker_source_disabled=getattr(args, "tracker_source_disabled", False),
            torch_threads=getattr(args, "torch_threads", None),
            torch_interop_threads=getattr(args, "torch_interop_threads", 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:
    _load_env_defaults()

    parser = argparse.ArgumentParser(
        prog="meshnet-node",
        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="MODEL", help="Model name or HuggingFace repo ID to serve")
    parser.add_argument("--model-id", metavar="MODEL", help="Alias for --model (catalog name or HuggingFace repo)")
    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("--tracker-source-disabled", action="store_true",
                        help="Skip tracker/peer model-file sources and download from HuggingFace directly")
    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("--advertise-host", metavar="ADDR", help="Host/IP advertised to the tracker")
    parser.add_argument("--route-timeout", type=float, metavar="SEC", default=30.0,
                        help="Seconds to wait for tracker route lookup (default 30)")
    parser.add_argument("--memory", type=int, metavar="MB", default=None,
                        help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
    parser.add_argument("--max-shards", type=int, metavar="N", default=None,
                        help="Maximum shard slots this node advertises to the tracker (default 1)")
    parser.add_argument("--torch-threads", type=int, metavar="N",
                        help="Set PyTorch intra-op CPU worker threads")
    parser.add_argument("--torch-interop-threads", type=int, metavar="N",
                        help="Set PyTorch inter-op CPU worker threads")
    parser.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
    parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
    parser.add_argument("--compact", action="store_true", help="Single-line status output")
    parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists")

    subparsers = parser.add_subparsers(dest="command")

    # 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")
    start_cmd.add_argument("--port", type=int)
    start_cmd.add_argument("--model", help="Model name or HuggingFace repo ID")
    start_cmd.add_argument("--model-id", help="Alias for --model (catalog name or HuggingFace repo)")
    start_cmd.add_argument("--shard-start", type=int)
    start_cmd.add_argument("--shard-end", type=int)
    start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto")
    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")
    start_cmd.add_argument("--route-timeout", type=float, default=30.0,
                           help="Seconds to wait for tracker route lookup (default 30)")
    start_cmd.add_argument("--memory", type=int, default=None, metavar="MB",
                           help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
    start_cmd.add_argument("--max-shards", type=int, default=1, metavar="N",
                           help="Maximum shard slots this node advertises to the tracker (default 1)")
    start_cmd.add_argument("--torch-threads", type=int, metavar="N",
                           help="Set PyTorch intra-op CPU worker threads")
    start_cmd.add_argument("--torch-interop-threads", type=int, metavar="N",
                           help="Set PyTorch inter-op CPU worker threads")
    start_cmd.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
    start_cmd.add_argument("--tracker-source-disabled", action="store_true",
                           help="Skip tracker/peer model-file sources and download from HuggingFace directly")

    args = parser.parse_args()

    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:
        # Default: wizard or start with saved config
        sys.exit(_cmd_default(args))


if __name__ == "__main__":
    main()