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

"""Full node startup sequence — self-configuring, non-interactive."""

from __future__ import annotations

import json
import socket
import sys
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any

from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware
from .server import StubNodeServer
from .torch_server import TorchNodeServer
from .wallet import load_or_create_wallet


def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict:
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        url, data=data, headers={"Content-Type": "application/json"}, method="POST"
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def _get_json(url: str, timeout: float = 10.0) -> dict:
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.loads(r.read())


def run_startup(
    tracker_url: str,
    port: int = 0,
    model: str = "stub-model",
    model_id: str | None = None,
    shard_start: int | None = None,
    shard_end: int | None = None,
    quantization: str = "bfloat16",
    wallet_path: Path | None = None,
    cache_dir: Path | None = None,
    host: str = "127.0.0.1",
    advertise_host: str | None = None,
    contracts: Any | None = None,
) -> StubNodeServer | TorchNodeServer:
    """Execute the full startup sequence and return a running node server.

    Steps (all non-interactive):
      1. Detect GPU / hardware profile
      2. Load or generate Solana wallet keypair
      3. Query tracker for optimal shard assignment
      4. Download (or stub) the assigned shard from peers, then HuggingFace
      5. Start local HTTP server
      6. Register with tracker

    Prints a compact status summary on completion.
    """

    tracker_url = tracker_url.rstrip("/")

    # 1. Hardware detection
    print("Detecting hardware...", flush=True)
    hw = detect_hardware()
    device: str = hw["device"]
    gpu_name: str | None = hw.get("gpu_name")
    vram_mb: int = hw.get("vram_mb", 0)

    if device == "cpu":
        print("  WARNING: No CUDA GPU detected — running in CPU mode", flush=True)
    else:
        print(f"  GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True)

    # 2. Wallet
    print("Loading wallet...", flush=True)
    wallet_kwargs: dict = {}
    if wallet_path is not None:
        wallet_kwargs["path"] = wallet_path
    _, _, address = load_or_create_wallet(**wallet_kwargs)
    print(f"  Wallet: {address}", flush=True)
    probationary_line = _probationary_status_line(contracts, address)
    if probationary_line is not None:
        print(f"  {probationary_line}", flush=True)

    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,
            port=port,
            model_id=model_id,
            shard_start=shard_start,
            shard_end=shard_end,
            quantization=quantization,
        )
        actual_port = node.start()
        public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
        endpoint = f"http://{public_host}:{actual_port}"
        print(
            f"\n{'=' * 32}\n"
            f"meshnet-node ready\n"
            f"  Wallet:       {address}\n"
            f"  Model ID:     {model_id}\n"
            f"  Shard:        layers {shard_start}–{shard_end}\n"
            f"  Quantization: {quantization}\n"
            f"  Endpoint:     {endpoint}\n"
            f"  Hardware:     {device.upper()}\n"
            f"{'=' * 32}",
            flush=True,
        )
        return node
    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)
    assign_qs = urllib.parse.urlencode({
        "model": model,
        "device": device,
        "vram_mb": vram_mb,
    })
    try:
        assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}")
    except urllib.error.URLError as exc:
        print(f"  ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True)
        raise

    shard_start: int = assignment["shard_start"]
    shard_end: int = assignment["shard_end"]
    assigned_model: str = assignment.get("model", model)
    hf_repo: str | None = assignment.get("hf_repo")
    peers: list[dict] = assignment.get("peers", [])
    print(f"  Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)

    # 4. Download shard
    print("Downloading shard...", flush=True)
    dl_kwargs: dict = {}
    if cache_dir is not None:
        dl_kwargs["cache_dir"] = cache_dir
    if hf_repo is not None:
        dl_kwargs["hf_repo"] = hf_repo
    if peers:
        dl_kwargs["peers"] = peers
    shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
    shard_checksum = compute_shard_checksum(shard_path)
    print(f"  Cached at: {shard_path}", flush=True)

    # 5. Start HTTP server
    is_last = shard_end >= assignment.get("model_layers_end", shard_end)
    node = StubNodeServer(
        host=host,
        port=port,
        shard_start=shard_start,
        shard_end=shard_end,
        is_last_shard=is_last,
        model=assigned_model,
        shard_path=shard_path,
    )
    actual_port = node.start()
    public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
    endpoint = f"http://{public_host}:{actual_port}"

    # 6. Register with tracker
    print("Registering with tracker...", flush=True)
    try:
        reg_resp = _post_json(
            f"{tracker_url}/v1/nodes/register",
            {
                "endpoint": endpoint,
                "model": assigned_model,
                "shard_start": shard_start,
                "shard_end": shard_end,
                "shard_checksum": shard_checksum,
                "hardware_profile": hw,
                "wallet_address": address,
                "score": 1.0,
            },
        )
        node_id: str = reg_resp["node_id"]
    except Exception:
        node.stop()
        raise

    # Status summary
    hw_str = device.upper()
    if gpu_name:
        hw_str += f" ({gpu_name}, {vram_mb} MB)"
    print(
        f"\n{'=' * 32}\n"
        f"meshnet-node ready\n"
        f"  Wallet:   {address}\n"
        f"  Shard:    layers {shard_start}-{shard_end} ({assigned_model})\n"
        f"  Endpoint: {endpoint}\n"
        f"  Node ID:  {node_id}\n"
        f"  Hardware: {hw_str}\n"
        f"{'=' * 32}",
        flush=True,
    )

    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
    remaining = contracts.registry.probationary_jobs_remaining(wallet_address)
    if remaining <= 0:
        return "Probationary period complete: earning enabled"
    suffix = "job" if remaining == 1 else "jobs"
    return f"Probationary period: {remaining} {suffix} remaining before earning"