|
| 1 | +import { spawn, ChildProcess } from "node:child_process"; |
| 2 | +import { logger } from "./logger"; |
| 3 | + |
| 4 | +// Inference backends node-llama-cpp can't serve (it only ships CUDA/Vulkan/ |
| 5 | +// Metal prebuilds) are run as external server processes that expose the |
| 6 | +// OpenAI-compatible chat-completions API, and chat traffic goes through the |
| 7 | +// same openai-compatible client the cloud providers use: |
| 8 | +// |
| 9 | +// - "mlx": Apple-Silicon inference via Python's mlx_lm package |
| 10 | +// (`python3 -m mlx_lm.server`). Models are Hugging Face repo ids or local |
| 11 | +// paths; the server downloads/loads them itself. |
| 12 | +// - "rocm": AMD-GPU inference via a ROCm/HIP build of llama.cpp's |
| 13 | +// `llama-server` binary (the official llama.cpp releases ship one), run |
| 14 | +// against the same GGUF files the built-in llama.cpp backend uses. |
| 15 | +export type LocalBackendId = "mlx" | "rocm"; |
| 16 | + |
| 17 | +export interface LocalBackendConfig { |
| 18 | + // Path to the ROCm llama-server binary. No sensible default beyond PATH |
| 19 | + // lookup — the user downloads a HIP build themselves. |
| 20 | + rocmServerPath?: string; |
| 21 | + // Python interpreter used to launch mlx_lm.server (needs `pip install mlx-lm`). |
| 22 | + mlxPythonPath?: string; |
| 23 | +} |
| 24 | + |
| 25 | +interface RunningServer { |
| 26 | + process: ChildProcess; |
| 27 | + model: string; |
| 28 | + baseUrl: string; |
| 29 | + exited: boolean; |
| 30 | +} |
| 31 | + |
| 32 | +// Fixed per-backend ports so a restarted app reconnects rather than leaking |
| 33 | +// orphan servers across random ports. |
| 34 | +const PORTS: Record<LocalBackendId, number> = { mlx: 8790, rocm: 8791 }; |
| 35 | +// First startup can include downloading/loading a multi-GB model. |
| 36 | +const STARTUP_TIMEOUT_MS = 180_000; |
| 37 | +const HEALTH_POLL_MS = 750; |
| 38 | + |
| 39 | +const servers = new Map<LocalBackendId, RunningServer>(); |
| 40 | + |
| 41 | +export function buildServerCommand( |
| 42 | + backend: LocalBackendId, |
| 43 | + model: string, |
| 44 | + config: LocalBackendConfig |
| 45 | +): { command: string; args: string[] } { |
| 46 | + const port = PORTS[backend]; |
| 47 | + if (backend === "mlx") { |
| 48 | + return { |
| 49 | + command: config.mlxPythonPath?.trim() || "python3", |
| 50 | + args: ["-m", "mlx_lm.server", "--model", model, "--port", String(port), "--host", "127.0.0.1"], |
| 51 | + }; |
| 52 | + } |
| 53 | + return { |
| 54 | + command: config.rocmServerPath?.trim() || "llama-server", |
| 55 | + args: [ |
| 56 | + "-m", model, |
| 57 | + "--port", String(port), |
| 58 | + "--host", "127.0.0.1", |
| 59 | + // Offload everything; llama-server clamps to what actually fits. |
| 60 | + "--n-gpu-layers", "999", |
| 61 | + ], |
| 62 | + }; |
| 63 | +} |
| 64 | + |
| 65 | +export function describeSpawnFailure(backend: LocalBackendId): string { |
| 66 | + return backend === "mlx" |
| 67 | + ? "Couldn't launch the MLX server — it needs Python with the mlx-lm package (pip install mlx-lm), available on Apple Silicon Macs." |
| 68 | + : "Couldn't launch llama-server — set the path to a ROCm (HIP) build of llama.cpp's llama-server binary in Settings."; |
| 69 | +} |
| 70 | + |
| 71 | +// Any HTTP response means the server socket is up (a 404 from a route probe |
| 72 | +// is still proof of life); only a network-level failure counts as down. |
| 73 | +async function isReachable(baseUrl: string): Promise<boolean> { |
| 74 | + try { |
| 75 | + await fetch(`${baseUrl}/v1/models`, { signal: AbortSignal.timeout(2000) }); |
| 76 | + return true; |
| 77 | + } catch { |
| 78 | + return false; |
| 79 | + } |
| 80 | +} |
| 81 | + |
| 82 | +function sleep(ms: number): Promise<void> { |
| 83 | + return new Promise((resolve) => setTimeout(resolve, ms)); |
| 84 | +} |
| 85 | + |
| 86 | +export async function ensureServer( |
| 87 | + backend: LocalBackendId, |
| 88 | + model: string, |
| 89 | + config: LocalBackendConfig |
| 90 | +): Promise<string> { |
| 91 | + const existing = servers.get(backend); |
| 92 | + if (existing && !existing.exited && existing.model === model) { |
| 93 | + if (await isReachable(existing.baseUrl)) return existing.baseUrl; |
| 94 | + // Process alive but unresponsive — restart it below. |
| 95 | + } |
| 96 | + if (existing) { |
| 97 | + existing.process.kill(); |
| 98 | + servers.delete(backend); |
| 99 | + } |
| 100 | + |
| 101 | + const baseUrl = `http://127.0.0.1:${PORTS[backend]}`; |
| 102 | + const { command, args } = buildServerCommand(backend, model, config); |
| 103 | + logger.info(`Starting ${backend} server: ${command} ${args.join(" ")}`); |
| 104 | + |
| 105 | + let child: ChildProcess; |
| 106 | + try { |
| 107 | + child = spawn(command, args, { stdio: "ignore" }); |
| 108 | + } catch { |
| 109 | + throw new Error(describeSpawnFailure(backend)); |
| 110 | + } |
| 111 | + |
| 112 | + const entry: RunningServer = { process: child, model, baseUrl, exited: false }; |
| 113 | + servers.set(backend, entry); |
| 114 | + |
| 115 | + let spawnError: string | null = null; |
| 116 | + child.on("error", (err) => { |
| 117 | + spawnError = (err as NodeJS.ErrnoException).code === "ENOENT" ? describeSpawnFailure(backend) : err.message; |
| 118 | + entry.exited = true; |
| 119 | + }); |
| 120 | + child.on("exit", (code) => { |
| 121 | + entry.exited = true; |
| 122 | + if (code !== 0 && code !== null) logger.warn(`${backend} server exited with code ${code}`); |
| 123 | + }); |
| 124 | + |
| 125 | + const deadline = Date.now() + STARTUP_TIMEOUT_MS; |
| 126 | + while (Date.now() < deadline) { |
| 127 | + if (spawnError) throw new Error(spawnError); |
| 128 | + if (entry.exited) { |
| 129 | + servers.delete(backend); |
| 130 | + throw new Error( |
| 131 | + backend === "mlx" |
| 132 | + ? "The MLX server exited during startup — check that mlx-lm is installed and the model id is valid." |
| 133 | + : "llama-server exited during startup — check that the binary is a working ROCm build and the model file is a valid GGUF." |
| 134 | + ); |
| 135 | + } |
| 136 | + if (await isReachable(baseUrl)) return baseUrl; |
| 137 | + await sleep(HEALTH_POLL_MS); |
| 138 | + } |
| 139 | + child.kill(); |
| 140 | + servers.delete(backend); |
| 141 | + throw new Error(`The ${backend} server didn't become reachable within ${STARTUP_TIMEOUT_MS / 1000}s.`); |
| 142 | +} |
| 143 | + |
| 144 | +export function stopServer(backend: LocalBackendId): void { |
| 145 | + const entry = servers.get(backend); |
| 146 | + if (entry) { |
| 147 | + entry.process.kill(); |
| 148 | + servers.delete(backend); |
| 149 | + } |
| 150 | +} |
| 151 | + |
| 152 | +export function stopAll(): void { |
| 153 | + for (const backend of [...servers.keys()]) stopServer(backend); |
| 154 | +} |
| 155 | + |
| 156 | +export function getRunningBackends(): { backend: LocalBackendId; model: string }[] { |
| 157 | + return [...servers.entries()] |
| 158 | + .filter(([, s]) => !s.exited) |
| 159 | + .map(([backend, s]) => ({ backend, model: s.model })); |
| 160 | +} |
0 commit comments