From 7c0b45f0c5d607f4c213b6ee06f616f0041ca124 Mon Sep 17 00:00:00 2001
From: hiento09 <136591877+hiento09@users.noreply.github.com>
Date: Tue, 26 Dec 2023 17:33:20 +0700
Subject: [PATCH] Fix switch thread crash nitro windows linux (#1214)

* Add fix bug #1204 on windows

* nitro gpu exit on kill

* correct bat script syntax

* Remove wait 5 second before start nitro

---------

Co-authored-by: Hien To <tominhhien97@gmail.com>
---
 .../inference-nitro-extension/bin/linux-start.sh  | 15 ++++++++-------
 .../inference-nitro-extension/bin/win-start.bat   |  7 ++++---
 .../inference-nitro-extension/src/module.ts       |  1 +
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/extensions/inference-nitro-extension/bin/linux-start.sh b/extensions/inference-nitro-extension/bin/linux-start.sh
index 3cf9c8013..199589ad2 100644
--- 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
+./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 $?
-else
-    echo "nitro_linux_amd64_cuda encountered an error, attempting to run nitro_linux_amd64..."
-    cd ../linux-cpu
-    ./nitro "$@"
-    exit $?
-fi
\ No newline at end of file
+)
diff --git a/extensions/inference-nitro-extension/bin/win-start.bat b/extensions/inference-nitro-extension/bin/win-start.bat
index e18c97a8e..250e0d218 100644
--- 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
index bca0b6fcc..a990ef81b 100644
--- 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)