Fix switch thread crash nitro windows linux (janhq#1214)

* Add fix bug janhq#1204 on windows * nitro gpu exit on kill * correct bat script syntax * Remove wait 5 second before start nitro --------- Co-authored-by: Hien To <[email protected]>
VPaulV · Dec 26, 2023 · 7c0b45f · 7c0b45f
1 parent 4cd7466
commit 7c0b45f
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 10 deletions.
diff --git a/extensions/inference-nitro-extension/bin/linux-start.sh b/extensions/inference-nitro-extension/bin/linux-start.sh
@@ -28,11 +28,12 @@ export CUDA_VISIBLE_DEVICES=$selectedGpuId
 
 # Attempt to run nitro_linux_amd64_cuda
 cd linux-cuda
-if ./nitro "$@"; then
-    exit $?
-else
-    echo "nitro_linux_amd64_cuda encountered an error, attempting to run nitro_linux_amd64..."
-    cd ../linux-cpu
-    ./nitro "$@"
+./nitro "$@" > output.log 2>&1 || (
+    echo "Check output log" && 
+    if grep -q "CUDA error" output.log; then
+        echo "CUDA error detected, attempting to run nitro_linux_amd64..."
+        cd ../linux-cpu && ./nitro "$@"
+        exit $?
+    fi
     exit $?
-fi
+)
diff --git a/extensions/inference-nitro-extension/bin/win-start.bat b/extensions/inference-nitro-extension/bin/win-start.bat
@@ -31,9 +31,10 @@ set CUDA_VISIBLE_DEVICES=!gpuId!
 
 rem Attempt to run nitro_windows_amd64_cuda.exe
 cd win-cuda
-nitro.exe %*
-if %errorlevel% neq 0 goto RunCpuVersion
-goto End
+
+nitro.exe %* > output.log
+type output.log | findstr /C:"CUDA error" >nul
+if %errorlevel% equ 0 ( goto :RunCpuVersion ) else ( goto :End )
 
 :RunCpuVersion
 rem Run nitro_windows_amd64.exe...

diff --git a/extensions/inference-nitro-extension/src/module.ts b/extensions/inference-nitro-extension/src/module.ts
@@ -67,6 +67,7 @@ async function loadModel(nitroResourceProbe: any | undefined) {
   // Gather system information for CPU physical cores and memory
   if (!nitroResourceProbe) nitroResourceProbe = await getResourcesInfo();
   return killSubprocess()
+    .then(() => tcpPortUsed.waitUntilFree(PORT, 300, 5000))
     .then(() => spawnNitroProcess(nitroResourceProbe))
     .then(() => loadLLMModel(currentSettings))
     .then(validateModelStatus)