From a0f6f00806c87bb51ca47dcab0f38a2ccc188db4 Mon Sep 17 00:00:00 2001
From: hiro <vuonghoainam.work@gmail.com>
Date: Tue, 23 Jan 2024 16:12:03 +0700
Subject: [PATCH] fix: Add fix for physical cpu core count

---
 .../src/node/index.ts                         |  9 +--
 .../src/node/utils.ts                         | 56 +++++++++++++++++++
 2 files changed, 61 insertions(+), 4 deletions(-)
 create mode 100644 extensions/inference-nitro-extension/src/node/utils.ts

diff --git a/extensions/inference-nitro-extension/src/node/index.ts b/extensions/inference-nitro-extension/src/node/index.ts
index 3aef77afd..3a342eec5 100644
--- a/extensions/inference-nitro-extension/src/node/index.ts
+++ b/extensions/inference-nitro-extension/src/node/index.ts
@@ -3,11 +3,12 @@ import path from "path";
 import { ChildProcessWithoutNullStreams, spawn } from "child_process";
 import tcpPortUsed from "tcp-port-used";
 import fetchRT from "fetch-retry";
-import osUtils from "os-utils";
 import { log, getJanDataFolderPath } from "@janhq/core/node";
 import { getNitroProcessInfo, updateNvidiaInfo } from "./nvidia";
 import { Model, InferenceEngine, ModelSettingParams } from "@janhq/core";
 import { executableNitroFile } from "./execute";
+import { physicalCpuCount } from "./utils";
+
 // Polyfill fetch with retry
 const fetchRetry = fetchRT(fetch);
 
@@ -121,8 +122,8 @@ async function runModel(
     currentSettings = {
       llama_model_path: currentModelFile,
       ...wrapper.model.settings,
-      // This is critical and requires real system information
-      cpu_threads: Math.max(1, Math.round(nitroResourceProbe.numCpuPhysicalCore / 2)),
+      // This is critical and requires real CPU physical core count (or performance core)
+      cpu_threads: Math.max(1, nitroResourceProbe.numCpuPhysicalCore),
     };
     return runNitroAndLoadModel();
   }
@@ -348,7 +349,7 @@ function spawnNitroProcess(): Promise<any> {
  */
 function getResourcesInfo(): Promise<ResourcesInfo> {
   return new Promise(async (resolve) => {
-    const cpu = await osUtils.cpuCount();
+    const cpu = await physicalCpuCount();
     log(`[NITRO]::CPU informations - ${cpu}`);
     const response: ResourcesInfo = {
       numCpuPhysicalCore: cpu,
diff --git a/extensions/inference-nitro-extension/src/node/utils.ts b/extensions/inference-nitro-extension/src/node/utils.ts
new file mode 100644
index 000000000..a2ef70df4
--- /dev/null
+++ b/extensions/inference-nitro-extension/src/node/utils.ts
@@ -0,0 +1,56 @@
+import os from "os";
+import childProcess from "child_process";
+
+function exec(command: string): Promise<string> {
+  return new Promise((resolve, reject) => {
+    childProcess.exec(command, { encoding: "utf8" }, (error, stdout) => {
+      if (error) {
+        reject(error);
+      } else {
+        resolve(stdout);
+      }
+    });
+  });
+}
+
+let amount: number;
+const platform = os.platform();
+
+export async function physicalCpuCount(): Promise<number> {
+  return new Promise((resolve, reject) => {
+    if (platform === "linux") {
+      exec('lscpu -p | egrep -v "^#" | sort -u -t, -k 2,4 | wc -l')
+        .then((output) => {
+          amount = parseInt(output.trim(), 10);
+          resolve(amount);
+        })
+        .catch(reject);
+    } else if (platform === "darwin") {
+      exec("sysctl -n hw.physicalcpu_max")
+        .then((output) => {
+          amount = parseInt(output.trim(), 10);
+          resolve(amount);
+        })
+        .catch(reject);
+    } else if (platform === "win32") {
+      exec("WMIC CPU Get NumberOfCores")
+        .then((output) => {
+          amount = output
+            .split(os.EOL)
+            .map((line: string) => parseInt(line))
+            .filter((value: number) => !isNaN(value))
+            .reduce((sum: number, number: number) => sum + number, 0);
+          resolve(amount);
+        })
+        .catch(reject);
+    } else {
+      const cores = os.cpus().filter((cpu: any, index: number) => {
+        const hasHyperthreading = cpu.model.includes("Intel");
+        const isOdd = index % 2 === 1;
+        return !hasHyperthreading || isOdd;
+      });
+      amount = cores.length;
+      resolve(amount);
+    }
+  });
+}