iinit
This commit is contained in:
Dobromir Popov
693
scripts/run_distributed_llama.py
Normal file
693
scripts/run_distributed_llama.py
Normal file
@@ -0,0 +1,693 @@
|
||||
#!/usr/bin/env python3
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
import subprocess
|
||||
import time
|
||||
import signal
|
||||
import socket
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
# --- Configuration & Defaults ---
|
||||
CONFIG_FILE = Path.home() / ".config" / "distributed-llama.json"
|
||||
DEFAULT_MODELS_DIR = Path.home() / "models"
|
||||
|
||||
# Binary paths: override via env vars or persisted config
|
||||
DEFAULT_LLAMA_SERVER_BIN = os.getenv("LLAMA_SERVER_BIN", "llama-server")
|
||||
DEFAULT_RPC_SERVER_BIN = os.getenv("LLAMA_RPC_SERVER_BIN", "rpc-server")
|
||||
|
||||
REMOTE_SSH_PORT = os.getenv("REMOTE_PORT", "22")
|
||||
RPC_PORT = int(os.getenv("RPC_PORT", "50052"))
|
||||
LOCAL_HOST_PORT = os.getenv("LOCAL_HOST_PORT", "8080")
|
||||
|
||||
MODES = ["llama-server", "llama-cli", "llama-bench"]
|
||||
DEFAULT_MODE = "llama-server"
|
||||
|
||||
DEFAULT_HOSTS: list = []
|
||||
|
||||
KV_CACHE_QUANT_VALUES = ("q8_0", "q5_1", "q5_0", "q4_1", "q4_0", "iq4_nl")
|
||||
KV_CACHE_OPTIONS = {
|
||||
"off": "Disabled (full precision)",
|
||||
"q8_0": "Q8_0 (recommended)",
|
||||
"q5_1": "Q5_1",
|
||||
"q5_0": "Q5_0",
|
||||
"q4_1": "Q4_1",
|
||||
"q4_0": "Q4_0 (aggressive)",
|
||||
"iq4_nl": "IQ4_NL",
|
||||
}
|
||||
|
||||
HEALTH_CHECK_INTERVAL = 10 # seconds between liveness polls
|
||||
RESTART_COOLDOWN = 5 # seconds to wait before attempting restart
|
||||
PORT_WAIT_TIMEOUT = 30 # seconds to wait for a (re)started node's port
|
||||
|
||||
|
||||
# --- Helpers ---
|
||||
|
||||
def check_dependencies():
|
||||
missing = [t for t in ("dialog", "ssh") if not shutil.which(t)]
|
||||
if missing:
|
||||
print(f"Error: missing required tools: {', '.join(missing)}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def run_dialog(args):
|
||||
with tempfile.NamedTemporaryFile(mode="w+") as tf:
|
||||
try:
|
||||
res = subprocess.run(["dialog"] + args, stderr=tf, check=False)
|
||||
tf.seek(0)
|
||||
return tf.read().strip(), res.returncode
|
||||
except Exception:
|
||||
return None, 1
|
||||
|
||||
|
||||
def show_msg(title, msg):
|
||||
run_dialog(["--title", title, "--msgbox", msg, "10", "60"])
|
||||
|
||||
|
||||
def _port_open(ip, port, timeout=2):
|
||||
try:
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
s.settimeout(timeout)
|
||||
s.connect((ip, port))
|
||||
s.close()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _wait_for_port(ip, port, stop_event, timeout=PORT_WAIT_TIMEOUT):
|
||||
for _ in range(timeout):
|
||||
if stop_event.is_set():
|
||||
return False
|
||||
if _port_open(ip, port):
|
||||
return True
|
||||
time.sleep(1)
|
||||
return False
|
||||
|
||||
|
||||
# --- File Picker (GGUF only) ---
|
||||
|
||||
def _dir_contents(path):
|
||||
try:
|
||||
entries = os.listdir(path)
|
||||
except PermissionError:
|
||||
return [], []
|
||||
dirs = sorted(e for e in entries if os.path.isdir(os.path.join(path, e)))
|
||||
files = sorted(e for e in entries if e.endswith(".gguf"))
|
||||
return dirs, files
|
||||
|
||||
|
||||
def custom_file_picker(start_path):
|
||||
current = os.path.abspath(start_path if os.path.isdir(start_path) else os.path.dirname(start_path))
|
||||
while True:
|
||||
dirs, files = _dir_contents(current)
|
||||
items = []
|
||||
if current != "/":
|
||||
items += ["..", "Parent Directory"]
|
||||
for d in dirs:
|
||||
items += [d + "/", "<DIR>"]
|
||||
for f in files:
|
||||
items += [f, "<GGUF>"]
|
||||
if not items:
|
||||
items += [".", "Empty Directory"]
|
||||
|
||||
pretty = current if len(current) <= 50 else "..." + current[-47:]
|
||||
sel, code = run_dialog([
|
||||
"--title", "Select GGUF File",
|
||||
"--backtitle", f"Current: {pretty}",
|
||||
"--menu", "Navigate and select a .gguf file:", "20", "70", "12",
|
||||
*items,
|
||||
])
|
||||
if code != 0:
|
||||
return None
|
||||
sel = sel.strip()
|
||||
if sel == "..":
|
||||
current = os.path.dirname(current)
|
||||
elif sel.endswith("/"):
|
||||
current = os.path.join(current, sel[:-1])
|
||||
elif sel != ".":
|
||||
return os.path.join(current, sel)
|
||||
|
||||
|
||||
# --- Node Monitor (health check + auto-restart) ---
|
||||
|
||||
class NodeMonitor(threading.Thread):
|
||||
"""Watches one RPC node and restarts it via SSH if the port goes dark."""
|
||||
|
||||
def __init__(self, ip, rpc_port, ssh_port, remote_bin, stop_event):
|
||||
super().__init__(daemon=True, name=f"monitor-{ip}")
|
||||
self.ip = ip
|
||||
self.rpc_port = rpc_port
|
||||
self.ssh_port = ssh_port
|
||||
self.remote_bin = remote_bin # path to rpc-server on the remote machine
|
||||
self.stop_event = stop_event
|
||||
self.restart_count = 0
|
||||
self.status = "ok" # "ok" | "down" | "restarting" | "unreachable"
|
||||
|
||||
# --- internals ---
|
||||
|
||||
def _ssh_start(self):
|
||||
"""Kill any stale rpc-server, start a fresh one, return PID string or None."""
|
||||
script = (
|
||||
f"pkill -9 -f '{self.remote_bin}' 2>/dev/null || true\n"
|
||||
f"sleep 1\n"
|
||||
f"nohup {self.remote_bin} -H 0.0.0.0 -p {self.rpc_port} "
|
||||
f"> /tmp/rpc-server-{self.ip}.log 2>&1 < /dev/null &\n"
|
||||
f"echo $!\n"
|
||||
)
|
||||
res = subprocess.run(
|
||||
[
|
||||
"ssh", "-p", self.ssh_port,
|
||||
"-o", "StrictHostKeyChecking=no",
|
||||
"-o", "ConnectTimeout=10",
|
||||
"-o", "BatchMode=yes",
|
||||
self.ip, "bash -s",
|
||||
],
|
||||
input=script, text=True, capture_output=True,
|
||||
)
|
||||
if res.returncode != 0:
|
||||
return None
|
||||
for line in reversed(res.stdout.splitlines()):
|
||||
if line.strip().isdigit():
|
||||
return line.strip()
|
||||
return None
|
||||
|
||||
def _attempt_restart(self):
|
||||
self.status = "restarting"
|
||||
print(f"\n[MONITOR] {self.ip}: restarting RPC server...", flush=True)
|
||||
pid = self._ssh_start()
|
||||
if not pid:
|
||||
self.status = "unreachable"
|
||||
print(f"[MONITOR] {self.ip}: SSH failed — node unreachable.", flush=True)
|
||||
return
|
||||
print(f"[MONITOR] {self.ip}: started (PID {pid}), waiting for port...", flush=True)
|
||||
if _wait_for_port(self.ip, self.rpc_port, self.stop_event):
|
||||
self.restart_count += 1
|
||||
self.status = "ok"
|
||||
print(f"[MONITOR] {self.ip}: back online (restart #{self.restart_count})", flush=True)
|
||||
else:
|
||||
self.status = "unreachable"
|
||||
print(f"[MONITOR] {self.ip}: port did not open after restart.", flush=True)
|
||||
|
||||
# --- thread entry point ---
|
||||
|
||||
def run(self):
|
||||
while not self.stop_event.wait(HEALTH_CHECK_INTERVAL):
|
||||
if not _port_open(self.ip, self.rpc_port):
|
||||
self.status = "down"
|
||||
print(f"\n[MONITOR] {self.ip}:{self.rpc_port} is DOWN.", flush=True)
|
||||
time.sleep(RESTART_COOLDOWN)
|
||||
if not self.stop_event.is_set():
|
||||
self._attempt_restart()
|
||||
else:
|
||||
self.status = "ok"
|
||||
|
||||
|
||||
# --- AppState ---
|
||||
|
||||
class AppState:
|
||||
def __init__(self):
|
||||
self.model_path = ""
|
||||
self.mode = DEFAULT_MODE
|
||||
self.hosts: list = list(DEFAULT_HOSTS) # [[ip, enabled], ...]
|
||||
self.context_size = None
|
||||
self.bench_prefill = "512,8192,16384,32768,65536"
|
||||
self.bench_gen = "128"
|
||||
self.kv_cache_quant = None
|
||||
self.extra_args = ""
|
||||
self.bench_extra_args = ""
|
||||
self.llama_server_bin = DEFAULT_LLAMA_SERVER_BIN
|
||||
self.rpc_server_bin = DEFAULT_RPC_SERVER_BIN
|
||||
self.load_config()
|
||||
|
||||
@property
|
||||
def active_hosts(self):
|
||||
return [h[0] for h in self.hosts if h[1]]
|
||||
|
||||
def save_config(self):
|
||||
data = {
|
||||
"model_path": self.model_path,
|
||||
"mode": self.mode,
|
||||
"hosts": self.hosts,
|
||||
"context_size": self.context_size,
|
||||
"bench_prefill": self.bench_prefill,
|
||||
"bench_gen": self.bench_gen,
|
||||
"kv_cache_quant": self.kv_cache_quant,
|
||||
"extra_args": self.extra_args,
|
||||
"bench_extra_args": self.bench_extra_args,
|
||||
"llama_server_bin": self.llama_server_bin,
|
||||
"rpc_server_bin": self.rpc_server_bin,
|
||||
}
|
||||
try:
|
||||
CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
CONFIG_FILE.write_text(json.dumps(data, indent=2))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def load_config(self):
|
||||
if not CONFIG_FILE.is_file():
|
||||
return
|
||||
try:
|
||||
data = json.loads(CONFIG_FILE.read_text())
|
||||
except (json.JSONDecodeError, OSError):
|
||||
return
|
||||
if not isinstance(data, dict):
|
||||
return
|
||||
|
||||
mp = data.get("model_path", "")
|
||||
if isinstance(mp, str) and mp and os.path.isfile(mp):
|
||||
self.model_path = mp
|
||||
|
||||
md = data.get("mode")
|
||||
if isinstance(md, str) and md in MODES:
|
||||
self.mode = md
|
||||
|
||||
hs = data.get("hosts")
|
||||
if isinstance(hs, list) and hs:
|
||||
valid = [
|
||||
e for e in hs
|
||||
if isinstance(e, list) and len(e) == 2
|
||||
and isinstance(e[0], str) and isinstance(e[1], bool)
|
||||
]
|
||||
if valid:
|
||||
self.hosts = valid
|
||||
|
||||
cs = data.get("context_size")
|
||||
if cs is None or (isinstance(cs, int) and cs > 0):
|
||||
self.context_size = cs
|
||||
|
||||
kv = data.get("kv_cache_quant")
|
||||
if kv is None or (isinstance(kv, str) and kv in KV_CACHE_QUANT_VALUES):
|
||||
self.kv_cache_quant = kv
|
||||
|
||||
for attr in ("bench_prefill", "bench_gen", "extra_args", "bench_extra_args"):
|
||||
v = data.get(attr)
|
||||
if isinstance(v, str):
|
||||
setattr(self, attr, v)
|
||||
|
||||
for attr in ("llama_server_bin", "rpc_server_bin"):
|
||||
v = data.get(attr)
|
||||
if isinstance(v, str) and v:
|
||||
setattr(self, attr, v)
|
||||
|
||||
|
||||
# --- TUI actions ---
|
||||
|
||||
def select_model(state):
|
||||
start = state.model_path if state.model_path else str(DEFAULT_MODELS_DIR)
|
||||
if os.path.isfile(start):
|
||||
start = os.path.dirname(start)
|
||||
sel = custom_file_picker(start)
|
||||
if sel:
|
||||
state.model_path = sel
|
||||
|
||||
|
||||
def select_mode(state):
|
||||
items = []
|
||||
for m in MODES:
|
||||
items += [m, ""]
|
||||
sel, code = run_dialog([
|
||||
"--title", "Select Execution Mode",
|
||||
"--menu", "Choose how to run the model:", "12", "50", "5",
|
||||
*items,
|
||||
])
|
||||
if code == 0 and sel:
|
||||
state.mode = sel
|
||||
|
||||
|
||||
def select_context(state):
|
||||
if state.mode == "llama-bench":
|
||||
sel_p, code = run_dialog([
|
||||
"--title", "Bench Prefill Sizes (pp)",
|
||||
"--inputbox",
|
||||
"Comma-separated prompt sizes (e.g. 512,8192,16384).\nLeave empty to skip:",
|
||||
"10", "68", str(state.bench_prefill or ""),
|
||||
])
|
||||
if code == 0:
|
||||
state.bench_prefill = sel_p.strip() or None
|
||||
sel_n, code2 = run_dialog([
|
||||
"--title", "Bench Token Generation (tg)",
|
||||
"--inputbox",
|
||||
"Comma-separated generation lengths (e.g. 128).\nLeave empty to skip:",
|
||||
"10", "68", str(state.bench_gen or ""),
|
||||
])
|
||||
if code2 == 0:
|
||||
state.bench_gen = sel_n.strip() or None
|
||||
else:
|
||||
sel, code = run_dialog([
|
||||
"--title", "Context Size",
|
||||
"--inputbox",
|
||||
"Enter context size (e.g. 4096, 8192).\nLeave empty for model default:",
|
||||
"10", "60", str(state.context_size or ""),
|
||||
])
|
||||
if code == 0:
|
||||
state.context_size = int(sel.strip()) if sel.strip().isdigit() else None
|
||||
|
||||
|
||||
def select_kv_cache(state):
|
||||
items = []
|
||||
for k, v in KV_CACHE_OPTIONS.items():
|
||||
items += [k, v]
|
||||
sel, code = run_dialog([
|
||||
"--title", "KV Cache Quantization",
|
||||
"--menu", "Quantize the KV cache to reduce memory usage:", "14", "55", "5",
|
||||
*items,
|
||||
])
|
||||
if code == 0 and sel:
|
||||
state.kv_cache_quant = None if sel == "off" else sel
|
||||
|
||||
|
||||
def edit_extra_args(state):
|
||||
is_bench = state.mode == "llama-bench"
|
||||
current = state.bench_extra_args if is_bench else state.extra_args
|
||||
label = "(Bench)" if is_bench else "(Server/CLI)"
|
||||
sel, code = run_dialog([
|
||||
"--title", f"Extra Arguments {label}",
|
||||
"--inputbox",
|
||||
f"Extra CLI args for {state.mode} (appended as-is).\nLeave empty for none:",
|
||||
"12", "65", current,
|
||||
])
|
||||
if code == 0:
|
||||
if is_bench:
|
||||
state.bench_extra_args = sel.strip()
|
||||
else:
|
||||
state.extra_args = sel.strip()
|
||||
|
||||
|
||||
def edit_binaries(state):
|
||||
sel_s, code = run_dialog([
|
||||
"--title", "Local llama-server Binary",
|
||||
"--inputbox", "Path or name of the local llama-server binary:", "10", "65",
|
||||
state.llama_server_bin,
|
||||
])
|
||||
if code == 0 and sel_s.strip():
|
||||
state.llama_server_bin = sel_s.strip()
|
||||
|
||||
sel_r, code = run_dialog([
|
||||
"--title", "Remote rpc-server Binary",
|
||||
"--inputbox",
|
||||
"Path or name of rpc-server on the REMOTE machines.\n"
|
||||
"(e.g. /opt/llama/rpc-server or just rpc-server if in PATH):",
|
||||
"11", "65", state.rpc_server_bin,
|
||||
])
|
||||
if code == 0 and sel_r.strip():
|
||||
state.rpc_server_bin = sel_r.strip()
|
||||
|
||||
|
||||
def add_server(state):
|
||||
sel, code = run_dialog(["--title", "Add Server", "--inputbox", "Enter new server IP address:", "10", "50"])
|
||||
if code == 0 and sel.strip():
|
||||
state.hosts.append([sel.strip(), True])
|
||||
|
||||
|
||||
def remove_server(state):
|
||||
if not state.hosts:
|
||||
show_msg("Info", "No servers to remove.")
|
||||
return
|
||||
items = []
|
||||
for i, (ip, _) in enumerate(state.hosts):
|
||||
items += [str(i), ip]
|
||||
sel, code = run_dialog(["--title", "Remove Server", "--menu", "Select server to remove:", "15", "50", "5", *items])
|
||||
if code == 0 and sel:
|
||||
idx = int(sel)
|
||||
if 0 <= idx < len(state.hosts):
|
||||
del state.hosts[idx]
|
||||
|
||||
|
||||
def edit_server(state):
|
||||
if not state.hosts:
|
||||
show_msg("Info", "No servers to edit.")
|
||||
return
|
||||
items = []
|
||||
for i, (ip, _) in enumerate(state.hosts):
|
||||
items += [str(i), ip]
|
||||
sel, code = run_dialog(["--title", "Edit Server", "--menu", "Select server to edit:", "15", "50", "5", *items])
|
||||
if code == 0 and sel:
|
||||
idx = int(sel)
|
||||
if 0 <= idx < len(state.hosts):
|
||||
new_ip, code2 = run_dialog([
|
||||
"--title", "Edit Server IP",
|
||||
"--inputbox", "Enter new IP address:", "10", "50", state.hosts[idx][0],
|
||||
])
|
||||
if code2 == 0 and new_ip.strip():
|
||||
state.hosts[idx][0] = new_ip.strip()
|
||||
|
||||
|
||||
def toggle_servers(state):
|
||||
if not state.hosts:
|
||||
show_msg("Info", "No servers to configure. Add some first.")
|
||||
return
|
||||
items = []
|
||||
for i, (ip, enabled) in enumerate(state.hosts):
|
||||
items += [str(i), ip, "on" if enabled else "off"]
|
||||
sel_str, code = run_dialog([
|
||||
"--title", "Toggle Active Servers",
|
||||
"--checklist", "Space to toggle active nodes:", "15", "50", "5",
|
||||
*items,
|
||||
])
|
||||
if code == 0:
|
||||
for h in state.hosts:
|
||||
h[1] = False
|
||||
if sel_str:
|
||||
for x in sel_str.split():
|
||||
idx = int(x.strip('"'))
|
||||
if 0 <= idx < len(state.hosts):
|
||||
state.hosts[idx][1] = True
|
||||
|
||||
|
||||
def configure_servers(state):
|
||||
while True:
|
||||
menu = ["1", "Toggle Active Servers", "2", "Add Server", "3", "Remove Server", "4", "Edit Server", "5", "Back"]
|
||||
sel, code = run_dialog(["--title", "Manage Remote Servers", "--menu", "Choose an action:", "15", "50", "5", *menu])
|
||||
if code != 0 or sel == "5":
|
||||
break
|
||||
{"1": toggle_servers, "2": add_server, "3": remove_server, "4": edit_server}.get(sel, lambda s: None)(state)
|
||||
|
||||
|
||||
# --- Core: start one RPC node and return its PID ---
|
||||
|
||||
def _ssh_start_rpc(ip, rpc_port, ssh_port, remote_bin):
|
||||
script = (
|
||||
f"pkill -9 -f '{remote_bin}' 2>/dev/null || true\n"
|
||||
f"sleep 1\n"
|
||||
f"nohup {remote_bin} -H 0.0.0.0 -p {rpc_port} "
|
||||
f"> /tmp/rpc-server-{ip}.log 2>&1 < /dev/null &\n"
|
||||
f"echo $!\n"
|
||||
)
|
||||
res = subprocess.run(
|
||||
[
|
||||
"ssh", "-p", ssh_port,
|
||||
"-o", "StrictHostKeyChecking=no",
|
||||
"-o", "ConnectTimeout=10",
|
||||
"-o", "BatchMode=yes",
|
||||
ip, "bash -s",
|
||||
],
|
||||
input=script, text=True, capture_output=True,
|
||||
)
|
||||
if res.returncode != 0:
|
||||
return None, res.stderr.strip()
|
||||
for line in reversed(res.stdout.splitlines()):
|
||||
if line.strip().isdigit():
|
||||
return line.strip(), None
|
||||
return None, f"unexpected output: {res.stdout!r}"
|
||||
|
||||
|
||||
def _kill_rpc(ip, ssh_port, remote_bin, pid=None):
|
||||
parts = []
|
||||
if pid:
|
||||
parts.append(f"kill -9 {pid} 2>/dev/null || true")
|
||||
parts.append(f"pkill -9 -f '{remote_bin}' 2>/dev/null || true")
|
||||
subprocess.run(
|
||||
["ssh", "-p", ssh_port, "-o", "ConnectTimeout=5", "-o", "BatchMode=yes", ip, "; ".join(parts)],
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
|
||||
|
||||
# --- Run ---
|
||||
|
||||
def run_distributed(state):
|
||||
if not state.model_path or not os.path.exists(state.model_path):
|
||||
show_msg("Error", f"Model file not found:\n{state.model_path}")
|
||||
return
|
||||
if not state.active_hosts:
|
||||
show_msg("Error", "No remote servers selected.")
|
||||
return
|
||||
|
||||
subprocess.run(["clear"])
|
||||
print("=== Distributed llama.cpp ===")
|
||||
print(f"Model: {state.model_path}")
|
||||
print(f"Mode: {state.mode}")
|
||||
print(f"Local bin: {state.llama_server_bin}")
|
||||
print(f"Remote bin: {state.rpc_server_bin}")
|
||||
if state.mode == "llama-bench":
|
||||
print(f"Prefill: {state.bench_prefill or '-'}")
|
||||
print(f"Gen: {state.bench_gen or '-'}")
|
||||
else:
|
||||
print(f"Context: {state.context_size or 'default'}")
|
||||
print(f"KV Cache: {state.kv_cache_quant or 'off'}")
|
||||
print(f"Hosts: {state.active_hosts}")
|
||||
print("---")
|
||||
|
||||
active_ips = state.active_hosts
|
||||
remote_pids: dict[str, str | None] = {}
|
||||
stop_event = threading.Event()
|
||||
monitors: list[NodeMonitor] = []
|
||||
|
||||
def cleanup():
|
||||
print("\nStopping monitors and cleaning up remote RPC servers...")
|
||||
stop_event.set()
|
||||
for ip in active_ips:
|
||||
pid = remote_pids.get(ip)
|
||||
print(f" Killing RPC on {ip} (PID: {pid or '?'})...")
|
||||
_kill_rpc(ip, REMOTE_SSH_PORT, state.rpc_server_bin, pid)
|
||||
|
||||
def on_signal(sig, frame):
|
||||
cleanup()
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGINT, on_signal)
|
||||
signal.signal(signal.SIGTERM, on_signal)
|
||||
|
||||
try:
|
||||
rpc_endpoints = []
|
||||
|
||||
# 1. Start RPC servers on all remote nodes
|
||||
for ip in active_ips:
|
||||
print(f"-> Starting RPC server on {ip}...", end=" ", flush=True)
|
||||
pid, err = _ssh_start_rpc(ip, RPC_PORT, REMOTE_SSH_PORT, state.rpc_server_bin)
|
||||
if not pid:
|
||||
print(f"FAILED ({err})")
|
||||
cleanup()
|
||||
return
|
||||
remote_pids[ip] = pid
|
||||
print(f"PID {pid}", end=" ", flush=True)
|
||||
|
||||
print(f"| waiting for :{RPC_PORT}...", end=" ", flush=True)
|
||||
if not _wait_for_port(ip, RPC_PORT, stop_event):
|
||||
print("TIMEOUT")
|
||||
cleanup()
|
||||
return
|
||||
print("OK")
|
||||
rpc_endpoints.append(f"{ip}:{RPC_PORT}")
|
||||
|
||||
rpc_arg = ",".join(rpc_endpoints)
|
||||
print(f"All nodes ready. RPC: {rpc_arg}")
|
||||
|
||||
# 2. Start health monitor threads
|
||||
for ip in active_ips:
|
||||
m = NodeMonitor(ip, RPC_PORT, REMOTE_SSH_PORT, state.rpc_server_bin, stop_event)
|
||||
m.start()
|
||||
monitors.append(m)
|
||||
print(f"[MONITOR] Watching {len(monitors)} node(s) every {HEALTH_CHECK_INTERVAL}s.")
|
||||
print("---")
|
||||
|
||||
# 3. Build and run local command
|
||||
import shlex
|
||||
cmd = [state.llama_server_bin, "-m", state.model_path, "--rpc", rpc_arg]
|
||||
|
||||
if state.mode == "llama-server":
|
||||
cmd += ["--host", "0.0.0.0", "--port", LOCAL_HOST_PORT]
|
||||
if state.context_size:
|
||||
cmd += ["-c", str(state.context_size)]
|
||||
|
||||
elif state.mode == "llama-cli":
|
||||
cmd += ["-cnv", "-p", "You are a helpful assistant."]
|
||||
if state.context_size:
|
||||
cmd += ["-c", str(state.context_size)]
|
||||
|
||||
elif state.mode == "llama-bench":
|
||||
if state.bench_prefill:
|
||||
cmd += ["-p", state.bench_prefill.strip()]
|
||||
if state.bench_gen:
|
||||
cmd += ["-n", state.bench_gen.strip()]
|
||||
|
||||
if state.kv_cache_quant:
|
||||
cmd += ["--cache-type-k", state.kv_cache_quant, "--cache-type-v", state.kv_cache_quant]
|
||||
|
||||
extra = state.bench_extra_args if state.mode == "llama-bench" else state.extra_args
|
||||
if extra:
|
||||
cmd += shlex.split(extra)
|
||||
|
||||
print(f"CMD: {' '.join(cmd)}")
|
||||
print("---")
|
||||
proc = subprocess.Popen(cmd)
|
||||
proc.wait()
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n[ERROR] {e}")
|
||||
finally:
|
||||
cleanup()
|
||||
|
||||
input("\nDone. Press Enter to return to menu...")
|
||||
|
||||
|
||||
# --- Main Menu ---
|
||||
|
||||
def main_menu():
|
||||
state = AppState()
|
||||
|
||||
while True:
|
||||
model_display = Path(state.model_path).name if state.model_path else "(None)"
|
||||
servers_display = f"{len(state.active_hosts)} active"
|
||||
|
||||
if state.mode == "llama-bench":
|
||||
ctx_display = f"pp=[{state.bench_prefill or '-'}] tg={state.bench_gen or '-'}"
|
||||
ctx_label = "Bench: "
|
||||
run_label = "RUN BENCHMARK"
|
||||
else:
|
||||
ctx_display = str(state.context_size) if state.context_size else "default"
|
||||
ctx_label = "Context: "
|
||||
run_label = "RUN DISTRIBUTED SERVER"
|
||||
|
||||
kv_display = state.kv_cache_quant or "off"
|
||||
extra = (state.bench_extra_args if state.mode == "llama-bench" else state.extra_args) or "(none)"
|
||||
|
||||
menu_args = [
|
||||
"--clear", "--backtitle", "Distributed llama.cpp",
|
||||
"--title", "Main Menu",
|
||||
"--menu", "Configure and launch distributed inference:", "24", "70", "11",
|
||||
"1", f"Model: {model_display}",
|
||||
"2", f"Servers: {servers_display}",
|
||||
"3", f"Mode: {state.mode}",
|
||||
"4", f"{ctx_label}{ctx_display}",
|
||||
"5", f"KV Cache: {kv_display}",
|
||||
"6", f"Extra: {extra}",
|
||||
"7", f"Binaries: local={Path(state.llama_server_bin).name} remote={Path(state.rpc_server_bin).name}",
|
||||
"8", run_label,
|
||||
"9", "Exit",
|
||||
]
|
||||
|
||||
choice, code = run_dialog(menu_args)
|
||||
if code != 0:
|
||||
break
|
||||
|
||||
{
|
||||
"1": select_model,
|
||||
"2": configure_servers,
|
||||
"3": select_mode,
|
||||
"4": select_context,
|
||||
"5": select_kv_cache,
|
||||
"6": edit_extra_args,
|
||||
"7": edit_binaries,
|
||||
"8": run_distributed,
|
||||
}.get(choice, lambda s: None)(state)
|
||||
|
||||
if choice == "9":
|
||||
break
|
||||
|
||||
state.save_config()
|
||||
|
||||
state.save_config()
|
||||
subprocess.run(["clear"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
check_dependencies()
|
||||
try:
|
||||
main_menu()
|
||||
except KeyboardInterrupt:
|
||||
subprocess.run(["clear"])
|
||||
sys.exit(0)
|
||||
Reference in New Issue
Block a user