neuron-tai/packages/tracker/meshnet_tracker/cli.py

"""meshnet-tracker CLI entry point."""

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

from .accounts import DEFAULT_ACCOUNTS_DB_PATH
from .billing import DEFAULT_BILLING_DB_PATH
from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH
from .logging_setup import (
    DEFAULT_LOG_BACKUP_COUNT,
    DEFAULT_LOG_DIR,
    DEFAULT_LOG_MAX_BYTES,
    configure_tracker_file_logging,
)
from .routing_stats import RoutingConfig
from .server import (
    DEFAULT_CALLER_CREDIT_USDT,
    DEFAULT_DEVNET_TOPUP_USDT,
    TrackerServer,
    derive_relay_url_from_public_tracker_url,
)

DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3"


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 tracker env defaults before parsing arguments."""
    _load_env_file(Path.cwd() / ".env")
    _load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env")


def _routing_config_from_args(args: argparse.Namespace) -> RoutingConfig | None:
    """Build a RoutingConfig from CLI flags; None keeps env-var/server defaults."""
    overrides = {
        "explore_share": args.route_explore_share,
        "weight_alpha": args.route_weight_alpha,
        "stats_half_life_seconds": args.route_stats_half_life,
    }
    set_values = {key: value for key, value in overrides.items() if value is not None}
    if not set_values:
        return None
    return RoutingConfig(**set_values)


def main() -> None:
    _load_env_defaults()
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on")
    common.add_argument("--port", type=int, default=8080, help="Port to listen on")
    common.add_argument(
        "--heartbeat-timeout",
        type=float,
        default=30.0,
        help="Seconds before a node is removed from the registry after missed heartbeat",
    )
    common.add_argument(
        "--cluster-peers",
        default="",
        help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)",
    )
    common.add_argument(
        "--self-url",
        default=None,
        help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)",
    )
    common.add_argument(
        "--stats-db",
        default=None,
        metavar="PATH",
        help="SQLite database path for persistent model usage statistics",
    )
    common.add_argument(
        "--relay-url",
        default=None,
        help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws",
    )
    common.add_argument(
        "--billing-db",
        default=DEFAULT_BILLING_DB_PATH,
        metavar="PATH",
        help=(
            "SQLite database path for the USDT billing ledger "
            f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)"
        ),
    )
    common.add_argument(
        "--no-billing",
        action="store_true",
        help="Disable the USDT billing ledger",
    )
    common.add_argument(
        "--max-charge-per-request",
        type=float,
        default=None,
        help=(
            "Reject chat completion requests whose prompt plus requested completion "
            "token bound would cost more than this many USDT"
        ),
    )
    common.add_argument(
        "--starting-credit",
        type=float,
        default=DEFAULT_CALLER_CREDIT_USDT,
        metavar="USDT",
        help=(
            "One-time Caller Credit granted when an account creates its first "
            f"API key (default: {DEFAULT_CALLER_CREDIT_USDT}; set 0 to require "
            "deposits before inference)"
        ),
    )
    common.add_argument(
        "--devnet-topup",
        type=float,
        default=DEFAULT_DEVNET_TOPUP_USDT,
        metavar="USDT",
        help=(
            "Dashboard devnet top-up faucet: each click credits this many USDT "
            f"to one of the account's keys (default: {DEFAULT_DEVNET_TOPUP_USDT}; "
            "MUST be 0 on mainnet deployments)"
        ),
    )
    common.add_argument(
        "--registry-db",
        default=DEFAULT_REGISTRY_DB_PATH,
        metavar="PATH",
        help=(
            "SQLite database path for persisted strike/ban/reputation registry "
            f"state (default: {DEFAULT_REGISTRY_DB_PATH})"
        ),
    )
    common.add_argument(
        "--no-registry-contracts",
        action="store_true",
        help="Disable the local contract registry used for strike/ban/reputation enforcement",
    )
    common.add_argument(
        "--accounts-db",
        default=DEFAULT_ACCOUNTS_DB_PATH,
        metavar="PATH",
        help=(
            "SQLite database path for dashboard user accounts "
            f"(default: {DEFAULT_ACCOUNTS_DB_PATH})"
        ),
    )
    common.add_argument(
        "--no-accounts",
        action="store_true",
        help="Disable dashboard user accounts (registration/login)",
    )
    common.add_argument(
        "--solana-rpc-url",
        default=None,
        help="Solana RPC URL (e.g. https://api.devnet.solana.com); enables the on-chain treasury",
    )
    common.add_argument(
        "--usdt-mint",
        default=None,
        help="SPL mint address of (mock) USDT — see scripts/devnet_setup.py",
    )
    common.add_argument(
        "--treasury-keypair",
        default=None,
        metavar="PATH",
        help="Treasury keypair JSON path (only on settlement-capable trackers)",
    )
    common.add_argument(
        "--settle-period",
        type=float,
        default=86400.0,
        help="Max seconds between payouts to a node (dev: 60, prod: 86400)",
    )
    common.add_argument(
        "--payout-threshold",
        type=float,
        default=5.0,
        help="Pending USDT that triggers an immediate payout (dev: 0)",
    )
    common.add_argument(
        "--payout-dust-floor",
        type=float,
        default=0.01,
        help="Never pay out less than this many USDT",
    )
    common.add_argument(
        "--validator-service-token",
        default=None,
        help=(
            "Service token the validator uses on POST /v1/billing/forfeit "
            "(default: MESHNET_VALIDATOR_SERVICE_TOKEN env; ADR-0017)"
        ),
    )
    common.add_argument(
        "--hive-secret",
        default=None,
        help=(
            "Shared secret authenticating gossip between tracker peers "
            "(default: MESHNET_HIVE_SECRET env; required for multi-tracker replication)"
        ),
    )
    common.add_argument(
        "--toploc-calibration-db",
        default=None,
        metavar="PATH",
        help=(
            "SQLite path for the AH-021 honest-noise TOPLOC calibration corpus "
            "(enables POST /v1/calibration/toploc/run + GET /v1/calibration/toploc/results)"
        ),
    )
    common.add_argument(
        "--toploc-reference-node-url",
        default=None,
        help="Reference node the calibration job teacher-forces claimed tokens against (see validator README)",
    )
    common.add_argument(
        "--toploc-calibration-gate-min-hardware-profiles",
        type=int,
        default=1,
        help=(
            "Distinct (GPU model, dtype) profiles the corpus must cover before "
            "gate_status.ready is true (alpha exception: fleet size is acceptable)"
        ),
    )
    common.add_argument(
        "--enable-hf-pricing",
        action="store_true",
        help=(
            "Enable the daily dynamic pricing refresh (issue 23): for presets with a "
            "curated hf_aliases list, sets the client price to 80%% of the cheapest "
            "matching HuggingFace inference-marketplace rate. Presets without "
            "hf_aliases are unaffected and keep their static price."
        ),
    )
    common.add_argument(
        "--hf-pricing-log-db",
        default=None,
        metavar="PATH",
        help=(
            "SQLite database path for the dynamic pricing change log "
            f"(default when --enable-hf-pricing is set: {DEFAULT_HF_PRICING_LOG_DB_PATH}; "
            "enables GET /v1/pricing/hf/history)"
        ),
    )
    common.add_argument(
        "--hf-pricing-refresh-interval",
        type=float,
        default=86400.0,
        help="Seconds between dynamic pricing refresh passes (default: daily)",
    )
    common.add_argument(
        "--models-dir",
        default=None,
        metavar="PATH",
        help="Local HuggingFace snapshot root advertised as tracker model-file source (default: MESHNET_MODELS_DIR)",
    )
    common.add_argument(
        "--route-explore-share",
        type=float,
        default=None,
        metavar="FRACTION",
        help=(
            "Fraction of requests routed down unproven/stale routes to gather "
            "throughput statistics (ADR-0021; default 0.3, lower once traffic grows)"
        ),
    )
    common.add_argument(
        "--route-weight-alpha",
        type=float,
        default=None,
        metavar="ALPHA",
        help=(
            "Traffic weight exponent among proven routes: share ∝ tps^alpha "
            "(default 1.0 — a 1.5x-faster route gets 1.5x the traffic)"
        ),
    )
    common.add_argument(
        "--route-stats-half-life",
        type=float,
        default=None,
        metavar="SECONDS",
        help="Half-life for decaying route throughput observations (default 600)",
    )
    common.add_argument(
        "--log-dir",
        default=DEFAULT_LOG_DIR,
        metavar="PATH",
        help=(
            "Directory for rotating tracker logs "
            f"(default: {DEFAULT_LOG_DIR}; files: info.log, warning.log, error.log)"
        ),
    )
    common.add_argument(
        "--log-max-bytes",
        type=int,
        default=DEFAULT_LOG_MAX_BYTES,
        metavar="BYTES",
        help=f"Rotate each tracker log file after this many bytes (default: {DEFAULT_LOG_MAX_BYTES})",
    )
    common.add_argument(
        "--log-backup-count",
        type=int,
        default=DEFAULT_LOG_BACKUP_COUNT,
        metavar="N",
        help=f"Number of rotated tracker log files to keep per level (default: {DEFAULT_LOG_BACKUP_COUNT})",
    )
    common.add_argument(
        "--no-file-logs",
        action="store_true",
        help="Disable rotating tracker log files and only write to the terminal",
    )

    parser = argparse.ArgumentParser(
        prog="meshnet-tracker",
        description="Distributed Inference Network node registry and route selection",
        parents=[common],
    )
    subparsers = parser.add_subparsers(dest="command")

    subparsers.add_parser("start", help="Start the tracker server", parents=[common])

    args = parser.parse_args()

    if args.command in {None, "start"}:
        if not args.no_file_logs:
            log_dir = configure_tracker_file_logging(
                args.log_dir,
                max_bytes=args.log_max_bytes,
                backup_count=args.log_backup_count,
            )
            print(f"meshnet-tracker logs: {log_dir}", flush=True)
        cluster_peers = [u.strip() for u in args.cluster_peers.split(",") if u.strip()]
        relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url)
        treasury = None
        if args.solana_rpc_url and args.usdt_mint and args.treasury_keypair:
            from meshnet_contracts.solana_adapter import SolanaCustodialTreasury

            treasury = SolanaCustodialTreasury(
                args.solana_rpc_url, args.usdt_mint, args.treasury_keypair,
            )
        contracts = None
        if not args.no_registry_contracts:
            from meshnet_contracts import LocalSolanaContracts  # type: ignore[import-not-found]

            contracts = LocalSolanaContracts(registry_db=args.registry_db)
        server = TrackerServer(
            host=args.host,
            port=args.port,
            heartbeat_timeout=args.heartbeat_timeout,
            cluster_peers=cluster_peers or None,
            cluster_self_url=args.self_url,
            stats_db=getattr(args, "stats_db", None),
            relay_url=relay_url,
            enable_billing=not args.no_billing,
            billing_db=None if args.no_billing else args.billing_db,
            max_charge_per_request=args.max_charge_per_request,
            starting_credit=args.starting_credit,
            devnet_topup_amount=args.devnet_topup,
            contracts=contracts,
            accounts_db=None if args.no_accounts else args.accounts_db,
            treasury=treasury,
            settle_period=args.settle_period,
            payout_threshold=args.payout_threshold,
            payout_dust_floor=args.payout_dust_floor,
            validator_service_token=args.validator_service_token,
            hive_secret=args.hive_secret,
            toploc_calibration_db=args.toploc_calibration_db,
            toploc_reference_node_url=args.toploc_reference_node_url,
            toploc_calibration_gate_min_hardware_profiles=args.toploc_calibration_gate_min_hardware_profiles,
            enable_hf_pricing=args.enable_hf_pricing,
            hf_pricing_log_db=(
                args.hf_pricing_log_db
                or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None)
            ),
            hf_pricing_refresh_interval=args.hf_pricing_refresh_interval,
            models_dir=args.models_dir,
            routing_config=_routing_config_from_args(args),
        )
        port = server.start()
        print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True)
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            server.stop()
            sys.exit(0)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()