Merge pull request #4079 from janhq/fix/factory-reset-hang-on-wiping-…

…data
janhq · Nov 21, 2024 · e9de9e7 · e9de9e7
2 parents 3c3b5bf + f709c77
commit e9de9e7
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 15 deletions.
diff --git 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
@@ -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<void>((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
@@ -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<string> => {
   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<string> => {
+  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
@@ -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
@@ -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