inference working
This commit is contained in:
Dobromir Popov
332
packages/node/meshnet_node/wizard.py
Normal file
332
packages/node/meshnet_node/wizard.py
Normal file
@@ -0,0 +1,332 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user