694 lines
22 KiB
Python
694 lines
22 KiB
Python
#!/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)
|