node-llama-cpp
Version:
Run AI models locally on your machine with node.js bindings for llama.cpp. Enforce a JSON schema on the model output on the generation level
238 lines • 9.15 kB
JavaScript
import { fork } from "node:child_process";
import { fileURLToPath } from "url";
import { createRequire } from "module";
import path from "path";
import { getConsoleLogPrefix } from "../../utils/getConsoleLogPrefix.js";
import { runningInElectron } from "../../utils/runtime.js";
import { LlamaLogLevel } from "../types.js";
import { LlamaLogLevelToAddonLogLevel } from "../Llama.js";
const require = createRequire(import.meta.url);
const __filename = fileURLToPath(import.meta.url);
const detectedFileName = path.basename(__filename);
const expectedFileName = "testBindingBinary";
export async function testBindingBinary(bindingBinaryPath, gpu, testTimeout = 1000 * 60 * 5, pipeOutputOnNode = false) {
if (!detectedFileName.startsWith(expectedFileName)) {
console.warn(getConsoleLogPrefix() +
`"${expectedFileName}.js" file is not independent, so testing a binding binary with the current system` +
"prior to importing it cannot be done.\n" +
getConsoleLogPrefix() +
"Assuming the test passed with the risk that the process may crash due to an incompatible binary.\n" +
getConsoleLogPrefix() +
'To resolve this issue, make sure that "node-llama-cpp" is not bundled together with other code and is imported as an external module with its original file structure.');
return true;
}
async function getForkFunction() {
if (runningInElectron) {
try {
const { utilityProcess } = await import("electron");
return {
type: "electron",
fork: utilityProcess.fork.bind(utilityProcess)
};
}
catch (err) {
// do nothing
}
}
return {
type: "node",
fork
};
}
const forkFunction = await getForkFunction();
function createTestProcess({ onMessage, onExit }) {
if (forkFunction.type === "electron") {
let exited = false;
const subProcess = forkFunction.fork(__filename, [], {
env: {
...process.env,
TEST_BINDING_CP: "true"
}
});
function cleanupElectronFork() {
if (subProcess.pid != null || !exited) {
subProcess.kill();
exited = true;
}
process.off("exit", cleanupElectronFork);
}
process.on("exit", cleanupElectronFork);
subProcess.on("message", onMessage);
subProcess.on("exit", (code) => {
exited = true;
cleanupElectronFork();
onExit(code);
});
return {
sendMessage: (message) => subProcess.postMessage(message),
killProcess: cleanupElectronFork,
pipeMessages: () => void 0
};
}
let pipeSet = false;
const subProcess = forkFunction.fork(__filename, [], {
detached: false,
silent: true,
stdio: pipeOutputOnNode
? ["ignore", "pipe", "pipe", "ipc"]
: ["ignore", "ignore", "ignore", "ipc"],
env: {
...process.env,
TEST_BINDING_CP: "true"
}
});
function cleanupNodeFork() {
subProcess.stdout?.off("data", onStdout);
subProcess.stderr?.off("data", onStderr);
if (subProcess.exitCode == null)
subProcess.kill("SIGKILL");
process.off("exit", cleanupNodeFork);
}
process.on("exit", cleanupNodeFork);
subProcess.on("message", onMessage);
subProcess.on("exit", (code) => {
cleanupNodeFork();
onExit(code ?? -1);
});
if (subProcess.killed || subProcess.exitCode != null) {
cleanupNodeFork();
onExit(subProcess.exitCode ?? -1);
}
function onStdout(data) {
if (!pipeSet)
return;
process.stdout.write(data);
}
function onStderr(data) {
if (!pipeSet)
return;
process.stderr.write(data);
}
if (pipeOutputOnNode) {
subProcess.stdout?.on("data", onStdout);
subProcess.stderr?.on("data", onStderr);
}
function pipeMessages() {
if (!pipeOutputOnNode || pipeSet)
return;
pipeSet = true;
}
return {
sendMessage: (message) => subProcess.send(message),
killProcess: cleanupNodeFork,
pipeMessages
};
}
let testPassed = false;
let forkSucceeded = false;
let timeoutHandle = null;
let subProcess = undefined;
let testFinished = false;
function cleanup() {
testFinished = true;
if (timeoutHandle != null)
clearTimeout(timeoutHandle);
subProcess?.killProcess();
}
return Promise.race([
new Promise((_, reject) => {
timeoutHandle = setTimeout(() => {
reject(new Error("Binding binary load test timed out"));
cleanup();
}, testTimeout);
}),
new Promise((resolve, reject) => {
function done() {
if (!forkSucceeded)
reject(new Error(`Binding binary test failed to run a test process via file "${__filename}"`));
else
resolve(testPassed);
cleanup();
}
subProcess = createTestProcess({
onMessage(message) {
if (message.type === "ready") {
forkSucceeded = true;
subProcess.sendMessage({
type: "start",
bindingBinaryPath,
gpu
});
}
else if (message.type === "loaded") {
subProcess.pipeMessages(); // only start piping error logs if the binary loaded successfully
subProcess.sendMessage({
type: "test",
bindingBinaryPath,
gpu
});
}
else if (message.type === "done") {
testPassed = true;
subProcess.sendMessage({ type: "exit" });
}
},
onExit(code) {
if (code !== 0)
testPassed = false;
done();
}
});
if (testFinished)
subProcess.killProcess();
})
]);
}
if (process.env.TEST_BINDING_CP === "true" && (process.parentPort != null || process.send != null)) {
let binding;
const sendMessage = process.parentPort != null
? (message) => process.parentPort.postMessage(message)
: (message) => process.send(message);
const onMessage = async (message) => {
if (message.type === "start") {
try {
binding = require(message.bindingBinaryPath);
const errorLogLevel = LlamaLogLevelToAddonLogLevel.get(LlamaLogLevel.error);
if (errorLogLevel != null)
binding.setLoggerLogLevel(errorLogLevel);
sendMessage({ type: "loaded" });
}
catch (err) {
console.error(err);
process.exit(1);
}
}
else if (message.type === "test") {
try {
if (binding == null)
throw new Error("Binding binary is not loaded");
binding.loadBackends();
const loadedGpu = binding.getGpuType();
if (loadedGpu == null || (loadedGpu === false && message.gpu !== false))
binding.loadBackends(path.dirname(path.resolve(message.bindingBinaryPath)));
await binding.init();
binding.getGpuVramInfo();
binding.getGpuDeviceInfo();
const gpuType = binding.getGpuType();
void gpuType;
if (gpuType !== message.gpu)
throw new Error(`Binary GPU type mismatch. Expected: ${message.gpu}, got: ${gpuType}`);
binding.ensureGpuDeviceIsSupported();
sendMessage({ type: "done" });
}
catch (err) {
console.error(err);
process.exit(1);
}
}
else if (message.type === "exit") {
process.exit(0);
}
};
if (process.parentPort != null)
process.parentPort.on("message", (message) => onMessage(message.data));
else
process.on("message", onMessage);
sendMessage({ type: "ready" });
}
//# sourceMappingURL=testBindingBinary.js.map