From f709c7723fc8e4d3f32be68faeda4d8eaf784c2a Mon Sep 17 00:00:00 2001 From: Louis Date: Thu, 21 Nov 2024 23:05:08 +0700 Subject: [PATCH] fix: factory reset hang on wiping data --- .../inference-cortex-extension/src/index.ts | 4 +- .../src/node/cpuInfo.ts | 27 ++++++++++++ .../src/node/execute.ts | 43 +++++++++++++------ web/hooks/useFactoryReset.test.ts | 8 ++++ web/hooks/useFactoryReset.ts | 11 ++++- 5 files changed, 78 insertions(+), 15 deletions(-) create mode 100644 extensions/inference-cortex-extension/src/node/cpuInfo.ts diff --git a/extensions/inference-cortex-extension/src/index.ts b/extensions/inference-cortex-extension/src/index.ts index 3ff5505042..531b407f2c 100644 --- a/extensions/inference-cortex-extension/src/index.ts +++ b/extensions/inference-cortex-extension/src/index.ts @@ -83,11 +83,11 @@ export default class JanInferenceCortexExtension extends LocalOAIEngine { }) } - onUnload(): void { + async onUnload() { console.log('Clean up cortex.cpp services') this.shouldReconnect = false this.clean() - executeOnMain(NODE, 'dispose') + await executeOnMain(NODE, 'dispose') super.onUnload() } diff --git a/extensions/inference-cortex-extension/src/node/cpuInfo.ts b/extensions/inference-cortex-extension/src/node/cpuInfo.ts new file mode 100644 index 0000000000..4366a995b6 --- /dev/null +++ b/extensions/inference-cortex-extension/src/node/cpuInfo.ts @@ -0,0 +1,27 @@ +import { cpuInfo } from 'cpu-instructions' + +// Check the CPU info and determine the supported instruction set +const info = cpuInfo.cpuInfo().some((e) => e.toUpperCase() === 'AVX512') + ? 'avx512' + : cpuInfo.cpuInfo().some((e) => e.toUpperCase() === 'AVX2') + ? 'avx2' + : cpuInfo.cpuInfo().some((e) => e.toUpperCase() === 'AVX') + ? 'avx' + : 'noavx' + +// Send the result and wait for confirmation before exiting +new Promise((resolve, reject) => { + // @ts-ignore + process.send(info, (error: Error | null) => { + if (error) { + reject(error) + } else { + resolve() + } + }) +}) + .then(() => process.exit(0)) + .catch((error) => { + console.error('Failed to send info:', error) + process.exit(1) + }) diff --git a/extensions/inference-cortex-extension/src/node/execute.ts b/extensions/inference-cortex-extension/src/node/execute.ts index 44b85d515b..0b091d464d 100644 --- a/extensions/inference-cortex-extension/src/node/execute.ts +++ b/extensions/inference-cortex-extension/src/node/execute.ts @@ -1,6 +1,6 @@ import * as path from 'path' -import { cpuInfo } from 'cpu-instructions' import { GpuSetting, appResourcePath, log } from '@janhq/core/node' +import { fork } from 'child_process' export interface CortexExecutableOptions { enginePath: string @@ -52,7 +52,9 @@ const extension = (): '.exe' | '' => { */ const cudaVersion = (settings?: GpuSetting): '11-7' | '12-0' | undefined => { const isUsingCuda = - settings?.vulkan !== true && settings?.run_mode === 'gpu' && !os().includes('mac') + settings?.vulkan !== true && + settings?.run_mode === 'gpu' && + !os().includes('mac') if (!isUsingCuda) return undefined return settings?.cuda?.version === '11' ? '11-7' : '12-0' @@ -62,15 +64,29 @@ const cudaVersion = (settings?: GpuSetting): '11-7' | '12-0' | undefined => { * The CPU instructions that will be set - either 'avx512', 'avx2', 'avx', or 'noavx'. * @returns */ -const cpuInstructions = (): string => { +const cpuInstructions = async (): Promise => { if (process.platform === 'darwin') return '' - return cpuInfo.cpuInfo().some((e) => e.toUpperCase() === 'AVX512') - ? 'avx512' - : cpuInfo.cpuInfo().some((e) => e.toUpperCase() === 'AVX2') - ? 'avx2' - : cpuInfo.cpuInfo().some((e) => e.toUpperCase() === 'AVX') - ? 'avx' - : 'noavx' + + const child = fork(path.join(__dirname, './cpuInfo.js')) // Path to the child process file + + return new Promise((resolve, reject) => { + child.on('message', (cpuInfo?: string) => { + resolve(cpuInfo ?? 'noavx') + child.kill() // Kill the child process after receiving the result + }) + + child.on('error', (err) => { + resolve('noavx') + child.kill() + }) + + child.on('exit', (code) => { + if (code !== 0) { + resolve('noavx') + child.kill() + } + }) + }) } /** @@ -94,8 +110,11 @@ export const executableCortexFile = ( /** * Find which variant to run based on the current platform. */ -export const engineVariant = (gpuSetting?: GpuSetting): string => { - const cpuInstruction = cpuInstructions() +export const engineVariant = async ( + gpuSetting?: GpuSetting +): Promise => { + const cpuInstruction = await cpuInstructions() + log(`[CORTEX]: CPU instruction: ${cpuInstruction}`) let engineVariant = [ os(), gpuSetting?.vulkan diff --git a/web/hooks/useFactoryReset.test.ts b/web/hooks/useFactoryReset.test.ts index b9ec10d6b4..7655da24b3 100644 --- a/web/hooks/useFactoryReset.test.ts +++ b/web/hooks/useFactoryReset.test.ts @@ -17,6 +17,14 @@ jest.mock('@janhq/core', () => ({ fs: { rm: jest.fn(), }, + EngineManager: { + instance: jest.fn().mockReturnValue({ + get: jest.fn(), + engines: { + values: jest.fn().mockReturnValue([]) + } + }), + }, })) describe('useFactoryReset', () => { diff --git a/web/hooks/useFactoryReset.ts b/web/hooks/useFactoryReset.ts index a8e3efb9ab..f68a6fd8c1 100644 --- a/web/hooks/useFactoryReset.ts +++ b/web/hooks/useFactoryReset.ts @@ -1,6 +1,6 @@ import { useCallback } from 'react' -import { fs, AppConfiguration } from '@janhq/core' +import { fs, AppConfiguration, EngineManager } from '@janhq/core' import { atom, useAtomValue, useSetAtom } from 'jotai' import { useActiveModel } from './useActiveModel' @@ -37,6 +37,15 @@ export default function useFactoryReset() { // 1: Stop running model setFactoryResetState(FactoryResetState.StoppingModel) await stopModel() + + await Promise.all( + EngineManager.instance() + .engines.values() + .map(async (engine) => { + await engine.onUnload() + }) + ) + await new Promise((resolve) => setTimeout(resolve, 4000)) // 2: Delete the old jan data folder