diff --git a/bin/lib/nim.js b/bin/lib/nim.js index 25c050097..8da775701 100644 --- a/bin/lib/nim.js +++ b/bin/lib/nim.js @@ -1,259 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 // -// NIM container management — pull, start, stop, health-check NIM images. +// Thin re-export shim — the implementation lives in src/lib/nim.ts, +// compiled to dist/lib/nim.js. -const { run, runCapture, shellQuote } = require("./runner"); -const nimImages = require("./nim-images.json"); -const UNIFIED_MEMORY_GPU_TAGS = ["GB10", "Thor", "Orin", "Xavier"]; - -function containerName(sandboxName) { - return `nemoclaw-nim-${sandboxName}`; -} - -function getImageForModel(modelName) { - const entry = nimImages.models.find((m) => m.name === modelName); - return entry ? entry.image : null; -} - -function listModels() { - return nimImages.models.map((m) => ({ - name: m.name, - image: m.image, - minGpuMemoryMB: m.minGpuMemoryMB, - })); -} - -function canRunNimWithMemory(totalMemoryMB) { - return nimImages.models.some((m) => m.minGpuMemoryMB <= totalMemoryMB); -} - -function detectGpu() { - // Try NVIDIA first — query VRAM - try { - const output = runCapture("nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits", { - ignoreError: true, - }); - if (output) { - const lines = output.split("\n").filter((l) => l.trim()); - 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); - return { - type: "nvidia", - count: perGpuMB.length, - totalMemoryMB, - perGpuMB: perGpuMB[0], - nimCapable: canRunNimWithMemory(totalMemoryMB), - }; - } - } - } catch { - /* ignored */ - } - - // Fallback: unified-memory NVIDIA devices where discrete VRAM is not queryable. - try { - const nameOutput = runCapture("nvidia-smi --query-gpu=name --format=csv,noheader,nounits", { - ignoreError: true, - }); - const gpuNames = nameOutput - .split("\n") - .map((line) => line.trim()) - .filter(Boolean); - const unifiedGpuNames = gpuNames.filter((name) => - UNIFIED_MEMORY_GPU_TAGS.some((tag) => new RegExp(tag, "i").test(name)), - ); - if (unifiedGpuNames.length > 0) { - let totalMemoryMB = 0; - try { - const memLine = runCapture("free -m | awk '/Mem:/ {print $2}'", { ignoreError: true }); - if (memLine) totalMemoryMB = parseInt(memLine.trim(), 10) || 0; - } catch { - /* ignored */ - } - const count = unifiedGpuNames.length; - const perGpuMB = count > 0 ? Math.floor(totalMemoryMB / count) : totalMemoryMB; - const isSpark = unifiedGpuNames.some((name) => /GB10/i.test(name)); - return { - type: "nvidia", - name: unifiedGpuNames[0], - count, - totalMemoryMB, - perGpuMB: perGpuMB || totalMemoryMB, - nimCapable: canRunNimWithMemory(totalMemoryMB), - unifiedMemory: true, - spark: isSpark, - }; - } - } catch { - /* ignored */ - } - - // macOS: detect Apple Silicon or discrete GPU - if (process.platform === "darwin") { - try { - const spOutput = runCapture("system_profiler SPDisplaysDataType 2>/dev/null", { - ignoreError: true, - }); - if (spOutput) { - const chipMatch = spOutput.match(/Chipset Model:\s*(.+)/); - const vramMatch = spOutput.match(/VRAM.*?:\s*(\d+)\s*(MB|GB)/i); - const coresMatch = spOutput.match(/Total Number of Cores:\s*(\d+)/); - - if (chipMatch) { - const name = chipMatch[1].trim(); - let memoryMB = 0; - - if (vramMatch) { - memoryMB = parseInt(vramMatch[1], 10); - if (vramMatch[2].toUpperCase() === "GB") memoryMB *= 1024; - } else { - // Apple Silicon shares system RAM — read total memory - try { - const memBytes = runCapture("sysctl -n hw.memsize", { ignoreError: true }); - if (memBytes) memoryMB = Math.floor(parseInt(memBytes, 10) / 1024 / 1024); - } catch { - /* ignored */ - } - } - - return { - type: "apple", - name, - count: 1, - cores: coresMatch ? parseInt(coresMatch[1], 10) : null, - totalMemoryMB: memoryMB, - perGpuMB: memoryMB, - nimCapable: false, - }; - } - } - } catch { - /* ignored */ - } - } - - return null; -} - -function pullNimImage(model) { - const image = getImageForModel(model); - if (!image) { - console.error(` Unknown model: ${model}`); - process.exit(1); - } - console.log(` Pulling NIM image: ${image}`); - run(`docker pull ${shellQuote(image)}`); - return image; -} - -function startNimContainer(sandboxName, model, port = 8000) { - const name = containerName(sandboxName); - return startNimContainerByName(name, model, port); -} - -function startNimContainerByName(name, model, port = 8000) { - const image = getImageForModel(model); - if (!image) { - console.error(` Unknown model: ${model}`); - process.exit(1); - } - - // Stop any existing container with same name - const qn = shellQuote(name); - run(`docker rm -f ${qn} 2>/dev/null || true`, { ignoreError: true }); - - console.log(` Starting NIM container: ${name}`); - run( - `docker run -d --gpus all -p ${Number(port)}:8000 --name ${qn} --shm-size 16g ${shellQuote(image)}`, - ); - return name; -} - -function waitForNimHealth(port = 8000, timeout = 300) { - const start = Date.now(); - const intervalSec = 5; - const hostPort = Number(port); - console.log(` Waiting for NIM health on port ${hostPort} (timeout: ${timeout}s)...`); - - while ((Date.now() - start) / 1000 < timeout) { - try { - const result = runCapture(`curl -sf http://localhost:${hostPort}/v1/models`, { - ignoreError: true, - }); - if (result) { - console.log(" NIM is healthy."); - return true; - } - } catch { - /* ignored */ - } - require("child_process").spawnSync("sleep", [String(intervalSec)]); - } - console.error(` NIM did not become healthy within ${timeout}s.`); - return false; -} - -function stopNimContainer(sandboxName) { - const name = containerName(sandboxName); - stopNimContainerByName(name); -} - -function stopNimContainerByName(name) { - const qn = shellQuote(name); - console.log(` Stopping NIM container: ${name}`); - run(`docker stop ${qn} 2>/dev/null || true`, { ignoreError: true }); - run(`docker rm ${qn} 2>/dev/null || true`, { ignoreError: true }); -} - -function nimStatus(sandboxName, port) { - const name = containerName(sandboxName); - return nimStatusByName(name, port); -} - -function nimStatusByName(name, port) { - try { - const qn = shellQuote(name); - const state = runCapture(`docker inspect --format '{{.State.Status}}' ${qn} 2>/dev/null`, { - ignoreError: true, - }); - if (!state) return { running: false, container: name }; - - let healthy = false; - if (state === "running") { - let resolvedHostPort = port != null ? Number(port) : 0; - if (!resolvedHostPort) { - const mapping = runCapture(`docker port ${qn} 8000 2>/dev/null`, { - ignoreError: true, - }); - const m = mapping && mapping.match(/:(\d+)\s*$/); - resolvedHostPort = m ? Number(m[1]) : 8000; - } - const health = runCapture( - `curl -sf http://localhost:${resolvedHostPort}/v1/models 2>/dev/null`, - { ignoreError: true }, - ); - healthy = !!health; - } - return { running: state === "running", healthy, container: name, state }; - } catch { - return { running: false, container: name }; - } -} - -module.exports = { - containerName, - getImageForModel, - listModels, - canRunNimWithMemory, - detectGpu, - pullNimImage, - startNimContainer, - startNimContainerByName, - waitForNimHealth, - stopNimContainer, - stopNimContainerByName, - nimStatus, - nimStatusByName, -}; +module.exports = require("../../dist/lib/nim"); diff --git a/test/nim.test.js b/src/lib/nim.test.ts similarity index 76% rename from test/nim.test.js rename to src/lib/nim.test.ts index 44f613a55..ab3cbe5f9 100644 --- a/test/nim.test.js +++ b/src/lib/nim.test.ts @@ -3,26 +3,29 @@ import { createRequire } from "module"; import { describe, it, expect, vi } from "vitest"; -import nim from "../bin/lib/nim"; +import type { Mock } from "vitest"; + +// Import from compiled dist/ for coverage attribution. +import nim from "../../dist/lib/nim"; const require = createRequire(import.meta.url); -const NIM_PATH = require.resolve("../bin/lib/nim"); -const RUNNER_PATH = require.resolve("../bin/lib/runner"); +const NIM_DIST_PATH = require.resolve("../../dist/lib/nim"); +const RUNNER_PATH = require.resolve("../../bin/lib/runner"); -function loadNimWithMockedRunner(runCapture) { +function loadNimWithMockedRunner(runCapture: Mock) { const runner = require(RUNNER_PATH); const originalRun = runner.run; const originalRunCapture = runner.runCapture; - delete require.cache[NIM_PATH]; + delete require.cache[NIM_DIST_PATH]; runner.run = vi.fn(); runner.runCapture = runCapture; - const nimModule = require(NIM_PATH); + const nimModule = require(NIM_DIST_PATH); return { nimModule, restore() { - delete require.cache[NIM_PATH]; + delete require.cache[NIM_DIST_PATH]; runner.run = originalRun; runner.runCapture = originalRunCapture; }, @@ -90,7 +93,7 @@ describe("nim", () => { }); it("detects GB10 unified-memory GPUs as Spark-capable NVIDIA devices", () => { - const runCapture = vi.fn((cmd) => { + const runCapture = vi.fn((cmd: string) => { if (cmd.includes("memory.total")) return ""; if (cmd.includes("query-gpu=name")) return "NVIDIA GB10"; if (cmd.includes("free -m")) return "131072"; @@ -115,7 +118,7 @@ describe("nim", () => { }); it("detects Orin unified-memory GPUs without marking them as Spark", () => { - const runCapture = vi.fn((cmd) => { + const runCapture = vi.fn((cmd: string) => { if (cmd.includes("memory.total")) return ""; if (cmd.includes("query-gpu=name")) return "NVIDIA Jetson AGX Orin"; if (cmd.includes("free -m")) return "32768"; @@ -140,7 +143,7 @@ describe("nim", () => { }); it("marks low-memory unified-memory NVIDIA devices as not NIM-capable", () => { - const runCapture = vi.fn((cmd) => { + const runCapture = vi.fn((cmd: string) => { if (cmd.includes("memory.total")) return ""; if (cmd.includes("query-gpu=name")) return "NVIDIA Xavier"; if (cmd.includes("free -m")) return "4096"; @@ -172,7 +175,7 @@ describe("nim", () => { describe("nimStatusByName", () => { it("uses provided port directly", () => { - const runCapture = vi.fn((cmd) => { + const runCapture = vi.fn((cmd: string) => { if (cmd.includes("docker inspect")) return "running"; if (cmd.includes("http://localhost:9000/v1/models")) return '{"data":[]}'; return ""; @@ -181,7 +184,7 @@ describe("nim", () => { try { const st = nimModule.nimStatusByName("foo", 9000); - const commands = runCapture.mock.calls.map(([cmd]) => cmd); + const commands = runCapture.mock.calls.map(([cmd]: [string]) => cmd); expect(st).toMatchObject({ running: true, @@ -189,8 +192,10 @@ describe("nim", () => { container: "foo", state: "running", }); - expect(commands.some((cmd) => cmd.includes("docker port"))).toBe(false); - expect(commands.some((cmd) => cmd.includes("http://localhost:9000/v1/models"))).toBe(true); + expect(commands.some((cmd: string) => cmd.includes("docker port"))).toBe(false); + expect(commands.some((cmd: string) => cmd.includes("http://localhost:9000/v1/models"))).toBe( + true, + ); } finally { restore(); } @@ -198,7 +203,7 @@ describe("nim", () => { it("uses published docker port when no port is provided", () => { for (const mapping of ["0.0.0.0:9000", "127.0.0.1:9000", "[::]:9000", ":::9000"]) { - const runCapture = vi.fn((cmd) => { + const runCapture = vi.fn((cmd: string) => { if (cmd.includes("docker inspect")) return "running"; if (cmd.includes("docker port")) return mapping; if (cmd.includes("http://localhost:9000/v1/models")) return '{"data":[]}'; @@ -208,18 +213,10 @@ describe("nim", () => { try { const st = nimModule.nimStatusByName("foo"); - const commands = runCapture.mock.calls.map(([cmd]) => cmd); + const commands = runCapture.mock.calls.map(([cmd]: [string]) => cmd); - expect(st).toMatchObject({ - running: true, - healthy: true, - container: "foo", - state: "running", - }); - expect(commands.some((cmd) => cmd.includes("docker port"))).toBe(true); - expect(commands.some((cmd) => cmd.includes("http://localhost:9000/v1/models"))).toBe( - true, - ); + expect(st).toMatchObject({ running: true, healthy: true, container: "foo", state: "running" }); + expect(commands.some((cmd: string) => cmd.includes("docker port"))).toBe(true); } finally { restore(); } @@ -227,7 +224,7 @@ describe("nim", () => { }); it("falls back to 8000 when docker port lookup fails", () => { - const runCapture = vi.fn((cmd) => { + const runCapture = vi.fn((cmd: string) => { if (cmd.includes("docker inspect")) return "running"; if (cmd.includes("docker port")) return ""; if (cmd.includes("http://localhost:8000/v1/models")) return '{"data":[]}'; @@ -237,23 +234,14 @@ describe("nim", () => { try { const st = nimModule.nimStatusByName("foo"); - const commands = runCapture.mock.calls.map(([cmd]) => cmd); - - expect(st).toMatchObject({ - running: true, - healthy: true, - container: "foo", - state: "running", - }); - expect(commands.some((cmd) => cmd.includes("docker port"))).toBe(true); - expect(commands.some((cmd) => cmd.includes("http://localhost:8000/v1/models"))).toBe(true); + expect(st).toMatchObject({ running: true, healthy: true, container: "foo", state: "running" }); } finally { restore(); } }); it("does not run health check when container is not running", () => { - const runCapture = vi.fn((cmd) => { + const runCapture = vi.fn((cmd: string) => { if (cmd.includes("docker inspect")) return "exited"; return ""; }); @@ -261,17 +249,8 @@ describe("nim", () => { try { const st = nimModule.nimStatusByName("foo"); - const commands = runCapture.mock.calls.map(([cmd]) => cmd); - - expect(st).toMatchObject({ - running: false, - healthy: false, - container: "foo", - state: "exited", - }); - expect(commands).toHaveLength(1); - expect(commands.some((cmd) => cmd.includes("docker port"))).toBe(false); - expect(commands.some((cmd) => cmd.includes("http://localhost:"))).toBe(false); + expect(st).toMatchObject({ running: false, healthy: false, container: "foo", state: "exited" }); + expect(runCapture.mock.calls).toHaveLength(1); } finally { restore(); } diff --git a/src/lib/nim.ts b/src/lib/nim.ts new file mode 100644 index 000000000..2e10d1e60 --- /dev/null +++ b/src/lib/nim.ts @@ -0,0 +1,275 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// NIM container management — pull, start, stop, health-check NIM images. + +// eslint-disable-next-line @typescript-eslint/no-require-imports +const { run, runCapture, shellQuote } = require("../../bin/lib/runner"); +// eslint-disable-next-line @typescript-eslint/no-require-imports +const nimImages = require("../../bin/lib/nim-images.json"); + +const UNIFIED_MEMORY_GPU_TAGS = ["GB10", "Thor", "Orin", "Xavier"]; + +export interface NimModel { + name: string; + image: string; + minGpuMemoryMB: number; +} + +export interface GpuDetection { + type: string; + name?: string; + count: number; + totalMemoryMB: number; + perGpuMB: number; + cores?: number | null; + nimCapable: boolean; + unifiedMemory?: boolean; + spark?: boolean; +} + +export interface NimStatus { + running: boolean; + healthy?: boolean; + container: string; + state?: string; +} + +export function containerName(sandboxName: string): string { + return `nemoclaw-nim-${sandboxName}`; +} + +export function getImageForModel(modelName: string): string | null { + const entry = nimImages.models.find((m: NimModel) => m.name === modelName); + return entry ? entry.image : null; +} + +export function listModels(): NimModel[] { + return nimImages.models.map((m: NimModel) => ({ + name: m.name, + image: m.image, + minGpuMemoryMB: m.minGpuMemoryMB, + })); +} + +export function canRunNimWithMemory(totalMemoryMB: number): boolean { + return nimImages.models.some((m: NimModel) => m.minGpuMemoryMB <= totalMemoryMB); +} + +export function detectGpu(): GpuDetection | null { + // Try NVIDIA first — query VRAM + try { + const output = runCapture( + "nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits", + { ignoreError: true }, + ); + if (output) { + const lines = output.split("\n").filter((l: string) => l.trim()); + const perGpuMB = lines + .map((l: string) => parseInt(l.trim(), 10)) + .filter((n: number) => !isNaN(n)); + if (perGpuMB.length > 0) { + const totalMemoryMB = perGpuMB.reduce((a: number, b: number) => a + b, 0); + return { + type: "nvidia", + count: perGpuMB.length, + totalMemoryMB, + perGpuMB: perGpuMB[0], + nimCapable: canRunNimWithMemory(totalMemoryMB), + }; + } + } + } catch { + /* ignored */ + } + + // Fallback: unified-memory NVIDIA devices + try { + const nameOutput = runCapture( + "nvidia-smi --query-gpu=name --format=csv,noheader,nounits", + { ignoreError: true }, + ); + const gpuNames = nameOutput + .split("\n") + .map((line: string) => line.trim()) + .filter(Boolean); + const unifiedGpuNames = gpuNames.filter((name: string) => + UNIFIED_MEMORY_GPU_TAGS.some((tag) => new RegExp(tag, "i").test(name)), + ); + if (unifiedGpuNames.length > 0) { + let totalMemoryMB = 0; + try { + const memLine = runCapture("free -m | awk '/Mem:/ {print $2}'", { ignoreError: true }); + if (memLine) totalMemoryMB = parseInt(memLine.trim(), 10) || 0; + } catch { + /* ignored */ + } + const count = unifiedGpuNames.length; + const perGpuMB = count > 0 ? Math.floor(totalMemoryMB / count) : totalMemoryMB; + const isSpark = unifiedGpuNames.some((name: string) => /GB10/i.test(name)); + return { + type: "nvidia", + name: unifiedGpuNames[0], + count, + totalMemoryMB, + perGpuMB: perGpuMB || totalMemoryMB, + nimCapable: canRunNimWithMemory(totalMemoryMB), + unifiedMemory: true, + spark: isSpark, + }; + } + } catch { + /* ignored */ + } + + // macOS: detect Apple Silicon or discrete GPU + if (process.platform === "darwin") { + try { + const spOutput = runCapture("system_profiler SPDisplaysDataType 2>/dev/null", { + ignoreError: true, + }); + if (spOutput) { + const chipMatch = spOutput.match(/Chipset Model:\s*(.+)/); + const vramMatch = spOutput.match(/VRAM.*?:\s*(\d+)\s*(MB|GB)/i); + const coresMatch = spOutput.match(/Total Number of Cores:\s*(\d+)/); + + if (chipMatch) { + const name = chipMatch[1].trim(); + let memoryMB = 0; + + if (vramMatch) { + memoryMB = parseInt(vramMatch[1], 10); + if (vramMatch[2].toUpperCase() === "GB") memoryMB *= 1024; + } else { + try { + const memBytes = runCapture("sysctl -n hw.memsize", { ignoreError: true }); + if (memBytes) memoryMB = Math.floor(parseInt(memBytes, 10) / 1024 / 1024); + } catch { + /* ignored */ + } + } + + return { + type: "apple", + name, + count: 1, + cores: coresMatch ? parseInt(coresMatch[1], 10) : null, + totalMemoryMB: memoryMB, + perGpuMB: memoryMB, + nimCapable: false, + }; + } + } + } catch { + /* ignored */ + } + } + + return null; +} + +export function pullNimImage(model: string): string { + const image = getImageForModel(model); + if (!image) { + console.error(` Unknown model: ${model}`); + process.exit(1); + } + console.log(` Pulling NIM image: ${image}`); + run(`docker pull ${shellQuote(image)}`); + return image; +} + +export function startNimContainer(sandboxName: string, model: string, port = 8000): string { + const name = containerName(sandboxName); + return startNimContainerByName(name, model, port); +} + +export function startNimContainerByName(name: string, model: string, port = 8000): string { + const image = getImageForModel(model); + if (!image) { + console.error(` Unknown model: ${model}`); + process.exit(1); + } + + const qn = shellQuote(name); + run(`docker rm -f ${qn} 2>/dev/null || true`, { ignoreError: true }); + + console.log(` Starting NIM container: ${name}`); + run( + `docker run -d --gpus all -p ${Number(port)}:8000 --name ${qn} --shm-size 16g ${shellQuote(image)}`, + ); + return name; +} + +export function waitForNimHealth(port = 8000, timeout = 300): boolean { + const start = Date.now(); + const intervalSec = 5; + const hostPort = Number(port); + console.log(` Waiting for NIM health on port ${hostPort} (timeout: ${timeout}s)...`); + + while ((Date.now() - start) / 1000 < timeout) { + try { + const result = runCapture(`curl -sf http://localhost:${hostPort}/v1/models`, { + ignoreError: true, + }); + if (result) { + console.log(" NIM is healthy."); + return true; + } + } catch { + /* ignored */ + } + // eslint-disable-next-line @typescript-eslint/no-require-imports + require("child_process").spawnSync("sleep", [String(intervalSec)]); + } + console.error(` NIM did not become healthy within ${timeout}s.`); + return false; +} + +export function stopNimContainer(sandboxName: string): void { + const name = containerName(sandboxName); + stopNimContainerByName(name); +} + +export function stopNimContainerByName(name: string): void { + const qn = shellQuote(name); + console.log(` Stopping NIM container: ${name}`); + run(`docker stop ${qn} 2>/dev/null || true`, { ignoreError: true }); + run(`docker rm ${qn} 2>/dev/null || true`, { ignoreError: true }); +} + +export function nimStatus(sandboxName: string, port?: number): NimStatus { + const name = containerName(sandboxName); + return nimStatusByName(name, port); +} + +export function nimStatusByName(name: string, port?: number): NimStatus { + try { + const qn = shellQuote(name); + const state = runCapture( + `docker inspect --format '{{.State.Status}}' ${qn} 2>/dev/null`, + { ignoreError: true }, + ); + if (!state) return { running: false, container: name }; + + let healthy = false; + if (state === "running") { + let resolvedHostPort = port != null ? Number(port) : 0; + if (!resolvedHostPort) { + const mapping = runCapture(`docker port ${qn} 8000 2>/dev/null`, { + ignoreError: true, + }); + const m = mapping && mapping.match(/:(\d+)\s*$/); + resolvedHostPort = m ? Number(m[1]) : 8000; + } + const health = runCapture( + `curl -sf http://localhost:${resolvedHostPort}/v1/models 2>/dev/null`, + { ignoreError: true }, + ); + healthy = !!health; + } + return { running: state === "running", healthy, container: name, state }; + } catch { + return { running: false, container: name }; + } +}