Skip to content

Commit

Permalink
fix: Add fix for physical cpu core count
Browse files Browse the repository at this point in the history
  • Loading branch information
hiro-v committed Jan 23, 2024
1 parent 17479e7 commit a0f6f00
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 4 deletions.
9 changes: 5 additions & 4 deletions extensions/inference-nitro-extension/src/node/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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,
Expand Down
56 changes: 56 additions & 0 deletions extensions/inference-nitro-extension/src/node/utils.ts
Original file line number Diff line number Diff line change
@@ -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);
}
});
}

0 comments on commit a0f6f00

Please sign in to comment.