diff --git a/bin/lib/nim.js b/bin/lib/nim.js index 4f2233e435f..e880a967923 100644 --- a/bin/lib/nim.js +++ b/bin/lib/nim.js @@ -6,15 +6,21 @@ const { run, runCapture } = require("./runner"); const nimImages = require("./nim-images.json"); +/** @param {string} sandboxName @returns {string} Docker container name. */ function containerName(sandboxName) { + if (!/^[a-zA-Z0-9][a-zA-Z0-9_.-]*$/.test(sandboxName)) { + throw new Error(`Invalid sandbox name: ${sandboxName}`); + } return `nemoclaw-nim-${sandboxName}`; } +/** @param {string} modelName @returns {string|null} NIM container image or null. */ function getImageForModel(modelName) { const entry = nimImages.models.find((m) => m.name === modelName); return entry ? entry.image : null; } +/** @returns {Array<{name: string, image: string, minGpuMemoryMB: number}>} */ function listModels() { return nimImages.models.map((m) => ({ name: m.name, @@ -23,10 +29,21 @@ function listModels() { })); } -function detectGpu() { +/** + * Detect GPU hardware. Returns an object describing the GPU (type, count, + * memory, capabilities) or null if no GPU is found. + * @param {object} [opts] - Optional overrides for dependency injection. + * @param {Function} [opts.runCapture] - Command runner (default: runner.runCapture). + * @param {string} [opts.platform] - OS platform (default: process.platform). + * @returns {{ type: string, count: number, totalMemoryMB: number, perGpuMB: number, nimCapable: boolean, spark?: boolean, name?: string, cores?: number } | null} + */ +function detectGpu(opts) { + const runCmd = (opts && opts.runCapture) || runCapture; + const platform = (opts && opts.platform) || process.platform; + // Try NVIDIA first — query VRAM try { - const output = runCapture( + const output = runCmd( "nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits", { ignoreError: true } ); @@ -35,8 +52,18 @@ function detectGpu() { const perGpuMB = lines.map((l) => parseInt(l.trim(), 10)).filter((n) => !isNaN(n)); if (perGpuMB.length > 0) { const totalMemoryMB = perGpuMB.reduce((a, b) => a + b, 0); + // Query GPU name for display + let name; + try { + name = runCmd( + "nvidia-smi --query-gpu=name --format=csv,noheader,nounits", + { ignoreError: true } + ); + if (name) name = name.split("\n")[0].trim(); + } catch {} return { type: "nvidia", + name, count: perGpuMB.length, totalMemoryMB, perGpuMB: perGpuMB[0], @@ -48,7 +75,7 @@ function detectGpu() { // Fallback: DGX Spark (GB10) — VRAM not queryable due to unified memory architecture try { - const nameOutput = runCapture( + const nameOutput = runCmd( "nvidia-smi --query-gpu=name --format=csv,noheader,nounits", { ignoreError: true } ); @@ -56,7 +83,7 @@ function detectGpu() { // GB10 has 128GB unified memory shared with Grace CPU — use system RAM let totalMemoryMB = 0; try { - const memLine = runCapture("free -m | awk '/Mem:/ {print $2}'", { ignoreError: true }); + const memLine = runCmd("free -m | awk '/Mem:/ {print $2}'", { ignoreError: true }); if (memLine) totalMemoryMB = parseInt(memLine.trim(), 10) || 0; } catch {} return { @@ -71,9 +98,9 @@ function detectGpu() { } catch {} // macOS: detect Apple Silicon or discrete GPU - if (process.platform === "darwin") { + if (platform === "darwin") { try { - const spOutput = runCapture( + const spOutput = runCmd( "system_profiler SPDisplaysDataType 2>/dev/null", { ignoreError: true } ); @@ -92,7 +119,7 @@ function detectGpu() { } else { // Apple Silicon shares system RAM — read total memory try { - const memBytes = runCapture("sysctl -n hw.memsize", { ignoreError: true }); + const memBytes = runCmd("sysctl -n hw.memsize", { ignoreError: true }); if (memBytes) memoryMB = Math.floor(parseInt(memBytes, 10) / 1024 / 1024); } catch {} } @@ -101,7 +128,7 @@ function detectGpu() { type: "apple", name, count: 1, - cores: coresMatch ? parseInt(coresMatch[1], 10) : null, + ...(coresMatch ? { cores: parseInt(coresMatch[1], 10) } : {}), totalMemoryMB: memoryMB, perGpuMB: memoryMB, nimCapable: false, @@ -114,23 +141,47 @@ function detectGpu() { return null; } +/** + * Suggest NIM models ranked by fit for a given GPU. + * Returns models sorted by VRAM requirement (descending), with the largest + * model that uses ≤90% of available VRAM marked as recommended. + * @param {{ totalMemoryMB: number, nimCapable: boolean } | null} gpu + * @returns {Array<{ name: string, image: string, minGpuMemoryMB: number, recommended: boolean }>} + */ +function suggestModelsForGpu(gpu) { + if (!gpu || !gpu.nimCapable) return []; + const vram = gpu.totalMemoryMB; + const fits = listModels() + .filter((m) => m.minGpuMemoryMB <= vram) + .sort((a, b) => b.minGpuMemoryMB - a.minGpuMemoryMB); + + // Mark the largest model that fits within 90% VRAM as recommended + const threshold = vram * 0.9; + let recommended = false; + return fits.map((m) => { + const rec = !recommended && m.minGpuMemoryMB <= threshold; + if (rec) recommended = true; + return { ...m, recommended: rec }; + }); +} + +/** @param {string} model - Model name to pull. @returns {string} Image tag. */ function pullNimImage(model) { const image = getImageForModel(model); if (!image) { - console.error(` Unknown model: ${model}`); - process.exit(1); + throw new Error(`Unknown model: ${model}`); } console.log(` Pulling NIM image: ${image}`); run(`docker pull ${image}`); return image; } +/** @param {string} sandboxName @param {string} model @param {number} [port=8000] @returns {string} Container name. */ function startNimContainer(sandboxName, model, port = 8000) { const name = containerName(sandboxName); const image = getImageForModel(model); if (!image) { - console.error(` Unknown model: ${model}`); - process.exit(1); + throw new Error(`Unknown model: ${model}`); } // Stop any existing container with same name @@ -143,6 +194,7 @@ function startNimContainer(sandboxName, model, port = 8000) { return name; } +/** @param {number} [port=8000] @param {number} [timeout=300] @returns {boolean} True if healthy. */ function waitForNimHealth(port = 8000, timeout = 300) { const start = Date.now(); const interval = 5000; @@ -150,8 +202,9 @@ function waitForNimHealth(port = 8000, timeout = 300) { while ((Date.now() - start) / 1000 < timeout) { try { - const result = runCapture(`curl -sf http://localhost:${port}/v1/models`, { + const result = runCapture(`curl -sf --max-time 10 http://localhost:${port}/v1/models`, { ignoreError: true, + timeout: 15000, }); if (result) { console.log(" NIM is healthy."); @@ -165,6 +218,7 @@ function waitForNimHealth(port = 8000, timeout = 300) { return false; } +/** @param {string} sandboxName - Stop and remove the NIM container. */ function stopNimContainer(sandboxName) { const name = containerName(sandboxName); console.log(` Stopping NIM container: ${name}`); @@ -172,7 +226,8 @@ function stopNimContainer(sandboxName) { run(`docker rm ${name} 2>/dev/null || true`, { ignoreError: true }); } -function nimStatus(sandboxName) { +/** @param {string} sandboxName @param {number} [port=8000] @returns {{running: boolean, healthy?: boolean, container: string, state?: string}} */ +function nimStatus(sandboxName, port = 8000) { const name = containerName(sandboxName); try { const state = runCapture( @@ -183,7 +238,7 @@ function nimStatus(sandboxName) { let healthy = false; if (state === "running") { - const health = runCapture(`curl -sf http://localhost:8000/v1/models 2>/dev/null`, { + const health = runCapture(`curl -sf http://localhost:${port}/v1/models 2>/dev/null`, { ignoreError: true, }); healthy = !!health; @@ -199,6 +254,7 @@ module.exports = { getImageForModel, listModels, detectGpu, + suggestModelsForGpu, pullNimImage, startNimContainer, waitForNimHealth, diff --git a/bin/lib/onboard.js b/bin/lib/onboard.js index 23f19b01a89..14797bef2a9 100644 --- a/bin/lib/onboard.js +++ b/bin/lib/onboard.js @@ -314,7 +314,8 @@ async function preflight() { // GPU const gpu = nim.detectGpu(); if (gpu && gpu.type === "nvidia") { - console.log(` ✓ NVIDIA GPU detected: ${gpu.count} GPU(s), ${gpu.totalMemoryMB} MB VRAM`); + const label = gpu.name ? `${gpu.name}, ` : ""; + console.log(` ✓ NVIDIA GPU detected: ${label}${gpu.count} GPU(s), ${gpu.totalMemoryMB} MB VRAM`); } else if (gpu && gpu.type === "apple") { console.log(` ✓ Apple GPU detected: ${gpu.name}${gpu.cores ? ` (${gpu.cores} cores)` : ""}, ${gpu.totalMemoryMB} MB unified memory`); console.log(" ⓘ NIM requires NVIDIA GPU — will use cloud inference"); @@ -539,8 +540,8 @@ async function setupNim(sandboxName, gpu) { } if (selected.key === "nim") { - // List models that fit GPU VRAM - const models = nim.listModels().filter((m) => m.minGpuMemoryMB <= gpu.totalMemoryMB); + // List models that fit GPU VRAM, ranked with recommendation + const models = nim.suggestModelsForGpu(gpu); if (models.length === 0) { console.log(" No NIM models fit your GPU VRAM. Falling back to cloud API."); } else { @@ -560,7 +561,8 @@ async function setupNim(sandboxName, gpu) { console.log(""); console.log(" Models that fit your GPU:"); models.forEach((m, i) => { - console.log(` ${i + 1}) ${m.name} (min ${m.minGpuMemoryMB} MB)`); + const tag = m.recommended ? " (recommended)" : ""; + console.log(` ${i + 1}) ${m.name} (min ${m.minGpuMemoryMB} MB)${tag}`); }); console.log(""); @@ -570,14 +572,22 @@ async function setupNim(sandboxName, gpu) { } model = sel.name; + const nimPort = 8000; + const nimPortCheck = await checkPortAvailable(nimPort); + if (!nimPortCheck.ok) { + console.error(` Port ${nimPort} is already in use (${nimPortCheck.process || "unknown"}).`); + console.error(" Stop the existing service or choose a different provider."); + process.exit(1); + } + console.log(` Pulling NIM image for ${model}...`); nim.pullNimImage(model); console.log(" Starting NIM container..."); - nimContainer = nim.startNimContainer(sandboxName, model); + nimContainer = nim.startNimContainer(sandboxName, model, nimPort); console.log(" Waiting for NIM to become healthy..."); - if (!nim.waitForNimHealth()) { + if (!nim.waitForNimHealth(nimPort)) { console.error(" NIM failed to start. Falling back to cloud API."); model = null; nimContainer = null; diff --git a/test/nim.test.js b/test/nim.test.js index 8166cf6c430..6e1bd1618ae 100644 --- a/test/nim.test.js +++ b/test/nim.test.js @@ -39,6 +39,18 @@ describe("nim", () => { it("prefixes with nemoclaw-nim-", () => { assert.equal(nim.containerName("my-sandbox"), "nemoclaw-nim-my-sandbox"); }); + + it("rejects names with shell metacharacters", () => { + assert.throws(() => nim.containerName("foo;rm -rf /"), /Invalid sandbox name/); + }); + + it("rejects empty string", () => { + assert.throws(() => nim.containerName(""), /Invalid sandbox name/); + }); + + it("accepts alphanumeric with dots and underscores", () => { + assert.equal(nim.containerName("my_sandbox.1"), "nemoclaw-nim-my_sandbox.1"); + }); }); describe("detectGpu", () => { @@ -68,10 +80,187 @@ describe("nim", () => { }); }); + describe("suggestModelsForGpu", () => { + it("returns empty for null GPU", () => { + assert.deepEqual(nim.suggestModelsForGpu(null), []); + }); + + it("returns empty for non-nimCapable GPU", () => { + assert.deepEqual(nim.suggestModelsForGpu({ totalMemoryMB: 16384, nimCapable: false }), []); + }); + + it("filters models that exceed VRAM", () => { + const models = nim.suggestModelsForGpu({ totalMemoryMB: 8000, nimCapable: true }); + for (const m of models) { + assert.ok(m.minGpuMemoryMB <= 8000, `${m.name} requires ${m.minGpuMemoryMB} MB`); + } + }); + + it("sorts by VRAM descending", () => { + const models = nim.suggestModelsForGpu({ totalMemoryMB: 200000, nimCapable: true }); + for (let i = 1; i < models.length; i++) { + assert.ok(models[i - 1].minGpuMemoryMB >= models[i].minGpuMemoryMB, + "models should be sorted by VRAM descending"); + } + }); + + it("marks exactly one model as recommended", () => { + const models = nim.suggestModelsForGpu({ totalMemoryMB: 200000, nimCapable: true }); + const recommended = models.filter((m) => m.recommended); + assert.equal(recommended.length, 1, "exactly one model should be recommended"); + }); + + it("recommended model fits within 90% VRAM", () => { + const vram = 32000; + const models = nim.suggestModelsForGpu({ totalMemoryMB: vram, nimCapable: true }); + const rec = models.find((m) => m.recommended); + if (rec) { + assert.ok(rec.minGpuMemoryMB <= vram * 0.9, + `recommended model (${rec.minGpuMemoryMB} MB) should fit within 90% of ${vram} MB`); + } + }); + + it("each entry has recommended boolean", () => { + const models = nim.suggestModelsForGpu({ totalMemoryMB: 200000, nimCapable: true }); + for (const m of models) { + assert.equal(typeof m.recommended, "boolean"); + } + }); + }); + + describe("pullNimImage", () => { + it("throws for unknown model instead of process.exit", () => { + assert.throws(() => nim.pullNimImage("bogus/nonexistent"), /Unknown model/); + }); + }); + describe("nimStatus", () => { it("returns not running for nonexistent container", () => { const st = nim.nimStatus("nonexistent-test-xyz"); assert.equal(st.running, false); }); }); + + describe("detectGpu (injected)", () => { + function mockRunCapture(responses) { + return function (cmd) { + for (const [pattern, response] of responses) { + if (cmd.includes(pattern)) { + if (response instanceof Error) throw response; + return response; + } + } + throw new Error("mock: no match for " + cmd); + }; + } + + it("detects standard NVIDIA GPU", () => { + const gpu = nim.detectGpu({ + runCapture: mockRunCapture([ + ["memory.total", "8192"], + ["query-gpu=name", "NVIDIA GeForce RTX 4090"], + ]), + }); + assert.equal(gpu.type, "nvidia"); + assert.equal(gpu.name, "NVIDIA GeForce RTX 4090"); + assert.equal(gpu.count, 1); + assert.equal(gpu.totalMemoryMB, 8192); + assert.equal(gpu.perGpuMB, 8192); + assert.equal(gpu.nimCapable, true); + assert.equal(gpu.spark, undefined); + }); + + it("detects multiple NVIDIA GPUs", () => { + const gpu = nim.detectGpu({ + runCapture: mockRunCapture([ + ["memory.total", "8192\n8192"], + ]), + }); + assert.equal(gpu.type, "nvidia"); + assert.equal(gpu.count, 2); + assert.equal(gpu.totalMemoryMB, 16384); + assert.equal(gpu.perGpuMB, 8192); + }); + + it("detects DGX Spark GB10", () => { + const gpu = nim.detectGpu({ + runCapture: mockRunCapture([ + ["memory.total", ""], + ["name", "NVIDIA GB10"], + ["free -m", "122880"], + ]), + }); + assert.equal(gpu.type, "nvidia"); + assert.equal(gpu.spark, true); + assert.equal(gpu.count, 1); + assert.equal(gpu.totalMemoryMB, 122880); + }); + + it("handles Spark with free -m failure", () => { + const gpu = nim.detectGpu({ + runCapture: mockRunCapture([ + ["memory.total", ""], + ["name", "NVIDIA GB10"], + ["free -m", new Error("command failed")], + ]), + }); + assert.equal(gpu.type, "nvidia"); + assert.equal(gpu.spark, true); + assert.equal(gpu.totalMemoryMB, 0); + }); + + it("detects macOS discrete GPU via VRAM", () => { + const gpu = nim.detectGpu({ + platform: "darwin", + runCapture: mockRunCapture([ + ["memory.total", new Error("no nvidia-smi")], + ["name", new Error("no nvidia-smi")], + ["system_profiler", "Chipset Model: Apple M2 Pro\n VRAM (Total): 16 GB\n Total Number of Cores: 19"], + ]), + }); + assert.equal(gpu.type, "apple"); + assert.equal(gpu.name, "Apple M2 Pro"); + assert.equal(gpu.nimCapable, false); + assert.equal(gpu.totalMemoryMB, 16384); + assert.equal(gpu.cores, 19); + }); + + it("detects Apple Silicon with unified memory", () => { + const gpu = nim.detectGpu({ + platform: "darwin", + runCapture: mockRunCapture([ + ["memory.total", new Error("no nvidia-smi")], + ["query-gpu=name", new Error("no nvidia-smi")], + ["system_profiler", "Chipset Model: Apple M4\n Total Number of Cores: 10"], + ["hw.memsize", "17179869184"], + ]), + }); + assert.equal(gpu.type, "apple"); + assert.equal(gpu.name, "Apple M4"); + assert.equal(gpu.nimCapable, false); + assert.equal(gpu.totalMemoryMB, 16384); + assert.equal(gpu.cores, 10); + }); + + it("returns null when no GPU detected", () => { + const gpu = nim.detectGpu({ + platform: "linux", + runCapture: mockRunCapture([ + ["memory.total", new Error("no nvidia-smi")], + ["name", new Error("no nvidia-smi")], + ]), + }); + assert.equal(gpu, null); + }); + + it("non-GB10 NVIDIA has no spark property", () => { + const gpu = nim.detectGpu({ + runCapture: mockRunCapture([ + ["memory.total", "24576"], + ]), + }); + assert.equal(gpu.type, "nvidia"); + assert.equal(gpu.spark, undefined); + }); + }); });