333 lines
12 KiB
Python
333 lines
12 KiB
Python
"""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, detect_num_layers
|
|
|
|
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:
|
|
# 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:
|
|
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
|
|
|
|
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)
|
|
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()
|