fix: pipe error logs in inspect gpu command

giladgd · giladgd · commit 2527413e3e46 · 2025-01-06T01:40:41.000+02:00
diff --git a/src/bindings/Llama.ts b/src/bindings/Llama.ts
@@ -14,7 +14,7 @@ import {BindingModule} from "./AddonTypes.js";
 import {BuildGpu, BuildMetadataFile, LlamaGpuType, LlamaLocks, LlamaLogLevel, LlamaLogLevelGreaterThanOrEqual} from "./types.js";
 import {MemoryOrchestrator, MemoryReservation} from "./utils/MemoryOrchestrator.js";
 
-const LlamaLogLevelToAddonLogLevel: ReadonlyMap<LlamaLogLevel, number> = new Map([
+export const LlamaLogLevelToAddonLogLevel: ReadonlyMap<LlamaLogLevel, number> = new Map([
     [LlamaLogLevel.disabled, 0],
     [LlamaLogLevel.fatal, 1],
     [LlamaLogLevel.error, 2],
diff --git a/src/bindings/getLlama.ts b/src/bindings/getLlama.ts
@@ -342,10 +342,12 @@ export async function getLlamaForOptions({
     debug = defaultLlamaCppDebugMode
 }: LlamaOptions, {
     updateLastBuildInfoOnCompile = false,
-    skipLlamaInit = false
+    skipLlamaInit = false,
+    pipeBinaryTestErrorLogs = false
 }: {
     updateLastBuildInfoOnCompile?: boolean,
-    skipLlamaInit?: boolean
+    skipLlamaInit?: boolean,
+    pipeBinaryTestErrorLogs?: boolean
 } = {}): Promise<Llama> {
     const platform = getPlatform();
     const arch = process.arch;
@@ -463,7 +465,8 @@ export async function getLlamaForOptions({
                             ? "falling back to building from source"
                             : null
                     ),
-                debug
+                debug,
+                pipeBinaryTestErrorLogs
             });
 
             if (llama != null)
@@ -556,7 +559,8 @@ async function loadExistingLlamaBinary({
     vramPadding,
     ramPadding,
     fallbackMessage,
-    debug
+    debug,
+    pipeBinaryTestErrorLogs
 }: {
     buildOptions: BuildOptions,
     canUsePrebuiltBinaries: boolean,
@@ -571,7 +575,8 @@ async function loadExistingLlamaBinary({
     vramPadding: Required<LlamaOptions>["vramPadding"],
     ramPadding: Required<LlamaOptions>["ramPadding"],
     fallbackMessage: string | null,
-    debug: boolean
+    debug: boolean,
+    pipeBinaryTestErrorLogs: boolean
 }) {
     const buildFolderName = await getBuildFolderNameForBuildOptions(buildOptions);
 
@@ -590,7 +595,7 @@ async function loadExistingLlamaBinary({
             });
             const resolvedBindingPath = await resolveActualBindingBinaryPath(localBuildBinPath);
             const binaryCompatible = shouldTestBinaryBeforeLoading
-                ? await testBindingBinary(resolvedBindingPath, buildOptions.gpu)
+                ? await testBindingBinary(resolvedBindingPath, buildOptions.gpu, undefined, pipeBinaryTestErrorLogs)
                 : true;
 
             if (binaryCompatible) {
@@ -649,7 +654,7 @@ async function loadExistingLlamaBinary({
                 });
                 const resolvedBindingPath = await resolveActualBindingBinaryPath(prebuiltBinDetails.binaryPath);
                 const binaryCompatible = shouldTestBinaryBeforeLoading
-                    ? await testBindingBinary(resolvedBindingPath, buildOptions.gpu)
+                    ? await testBindingBinary(resolvedBindingPath, buildOptions.gpu, undefined, pipeBinaryTestErrorLogs)
                     : true;
 
                 if (binaryCompatible) {
diff --git a/src/bindings/utils/testBindingBinary.ts b/src/bindings/utils/testBindingBinary.ts
@@ -4,15 +4,21 @@ import {createRequire} from "module";
 import path from "path";
 import {getConsoleLogPrefix} from "../../utils/getConsoleLogPrefix.js";
 import {runningInElectron} from "../../utils/runtime.js";
-import {BuildGpu} from "../types.js";
+import {BuildGpu, LlamaLogLevel} from "../types.js";
+import {LlamaLogLevelToAddonLogLevel} from "../Llama.js";
 import type {BindingModule} from "../AddonTypes.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: string, gpu: BuildGpu, testTimeout: number = 1000 * 60 * 5): Promise<boolean> {
+export async function testBindingBinary(
+    bindingBinaryPath: string,
+    gpu: BuildGpu,
+    testTimeout: number = 1000 * 60 * 5,
+    pipeOutputOnNode: boolean = false
+): Promise<boolean> {
     if (!detectedFileName.startsWith(expectedFileName)) {
         console.warn(
             getConsoleLogPrefix() +
@@ -57,7 +63,8 @@ export async function testBindingBinary(bindingBinaryPath: string, gpu: BuildGpu
         onExit(code: number): void
     }): {
         sendMessage(message: ParentToChildMessage): void,
-        killProcess(): void
+        killProcess(): void,
+        pipeMessages(): void
     } {
         if (forkFunction.type === "electron") {
             let exited = false;
@@ -88,20 +95,28 @@ export async function testBindingBinary(bindingBinaryPath: string, gpu: BuildGpu
 
             return {
                 sendMessage: (message: ParentToChildMessage) => subProcess.postMessage(message),
-                killProcess: cleanupElectronFork
+                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");
 
@@ -121,9 +136,36 @@ export async function testBindingBinary(bindingBinaryPath: string, gpu: BuildGpu
             onExit(subProcess.exitCode ?? -1);
         }
 
+        function onStdout(data: string) {
+            if (!pipeSet)
+                return;
+
+            process.stdout.write(data);
+        }
+
+        function onStderr(data: string) {
+            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: ParentToChildMessage) => subProcess.send(message),
-            killProcess: cleanupNodeFork
+            killProcess: cleanupNodeFork,
+            pipeMessages
         };
     }
 
@@ -169,6 +211,13 @@ export async function testBindingBinary(bindingBinaryPath: string, gpu: BuildGpu
                             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"});
@@ -189,13 +238,28 @@ export async function testBindingBinary(bindingBinaryPath: string, gpu: BuildGpu
 }
 
 if (process.env.TEST_BINDING_CP === "true" && (process.parentPort != null || process.send != null)) {
+    let binding: BindingModule;
     const sendMessage = process.parentPort != null
         ? (message: ChildToParentMessage) => process.parentPort.postMessage(message)
         : (message: ChildToParentMessage) => process.send!(message);
     const onMessage = async (message: ParentToChildMessage) => {
         if (message.type === "start") {
             try {
-                const binding: BindingModule = require(message.bindingBinaryPath);
+                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();
@@ -235,9 +299,13 @@ type ParentToChildMessage = {
     type: "start",
     bindingBinaryPath: string,
     gpu: BuildGpu
+} | {
+    type: "test",
+    bindingBinaryPath: string,
+    gpu: BuildGpu
 } | {
     type: "exit"
 };
 type ChildToParentMessage = {
-    type: "ready" | "done"
+    type: "ready" | "loaded" | "done"
 };
diff --git a/src/cli/commands/inspect/commands/InspectGpuCommand.ts b/src/cli/commands/inspect/commands/InspectGpuCommand.ts
@@ -172,7 +172,8 @@ async function getLlamaForGpu(gpu: BuildGpu) {
             logLevel: LlamaLogLevel.warn,
             vramPadding: 0
         }, {
-            skipLlamaInit: true
+            skipLlamaInit: true,
+            pipeBinaryTestErrorLogs: true
         });
     } catch (err) {
         return undefined;
