diff --git a/extensions/inference-nitro-extension/package.json b/extensions/inference-nitro-extension/package.json index ef74fff08..ecbbf17a8 100644 --- a/extensions/inference-nitro-extension/package.json +++ b/extensions/inference-nitro-extension/package.json @@ -36,6 +36,7 @@ "kill-port": "^2.0.1", "path-browserify": "^1.0.1", "rxjs": "^7.8.1", + "systeminformation": "^5.21.20", "tcp-port-used": "^1.0.2", "ts-loader": "^9.5.0", "ulid": "^2.3.0" @@ -52,6 +53,7 @@ "tcp-port-used", "kill-port", "fetch-retry", - "electron-log" + "electron-log", + "systeminformation" ] } diff --git a/extensions/inference-nitro-extension/src/@types/global.d.ts b/extensions/inference-nitro-extension/src/@types/global.d.ts index 642f10909..f93a3e4c9 100644 --- a/extensions/inference-nitro-extension/src/@types/global.d.ts +++ b/extensions/inference-nitro-extension/src/@types/global.d.ts @@ -24,3 +24,8 @@ interface ModelOperationResponse { error?: any; modelFile?: string; } + +interface ResourcesInfo { + numCpuPhysicalCore: number; + memAvailable: number; +} \ No newline at end of file diff --git a/extensions/inference-nitro-extension/src/module.ts b/extensions/inference-nitro-extension/src/module.ts index d36553f40..64a7393fc 100644 --- a/extensions/inference-nitro-extension/src/module.ts +++ b/extensions/inference-nitro-extension/src/module.ts @@ -4,6 +4,7 @@ const path = require("path"); const { spawn } = require("child_process"); const tcpPortUsed = require("tcp-port-used"); const fetchRetry = require("fetch-retry")(global.fetch); +const si = require("systeminformation"); const log = require("electron-log"); @@ -167,7 +168,7 @@ async function checkAndUnloadNitro() { * Should run exactly platform specified Nitro binary version */ async function spawnNitroProcess(): Promise { - return new Promise((resolve, reject) => { + return new Promise(async (resolve, reject) => { let binaryFolder = path.join(__dirname, "bin"); // Current directory by default let binaryName; @@ -190,10 +191,20 @@ async function spawnNitroProcess(): Promise { const binaryPath = path.join(binaryFolder, binaryName); + // Gather system information for CPU physical cores and memory + const nitroResourceProbe = await getResourcesInfo(); + console.log( + "Nitro with physical core: " + nitroResourceProbe.numCpuPhysicalCore + ); + // Execute the binary - subprocess = spawn(binaryPath, [1, "127.0.0.1", PORT], { - cwd: binaryFolder, - }); + subprocess = spawn( + binaryPath, + [nitroResourceProbe.numCpuPhysicalCore, "127.0.0.1", PORT], + { + cwd: binaryFolder, + } + ); // Handle subprocess output subprocess.stdout.on("data", (data) => { @@ -263,15 +274,28 @@ function validateModelVersion(): Promise { }); } -/** - * Cleans up any registered resources. - * Its module specific function, should be called when application is closed - */ + function dispose() { // clean other registered resources here killSubprocess(); } +/** + * Get the system resources information + */ +async function getResourcesInfo(): Promise { + return new Promise(async (resolve) => { + const cpu = await si.cpu(); + const mem = await si.mem(); + + const response = { + numCpuPhysicalCore: cpu.physicalCores, + memAvailable: mem.available, + }; + resolve(response); + }); +} + module.exports = { initModel, killSubprocess,