diff --git a/bin/lib/nim.js b/bin/lib/nim.js index 64c074dad..ceeeb9e46 100644 --- a/bin/lib/nim.js +++ b/bin/lib/nim.js @@ -3,7 +3,8 @@ // // NIM container management — pull, start, stop, health-check NIM images. -const { run, runCapture, shellQuote } = require("./runner"); +const runner = require("./runner"); +const { run, shellQuote } = runner; const nimImages = require("./nim-images.json"); function containerName(sandboxName) { @@ -26,7 +27,7 @@ function listModels() { function detectGpu() { // Try NVIDIA first — query VRAM try { - const output = runCapture( + const output = runner.runCapture( "nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits", { ignoreError: true } ); @@ -48,17 +49,19 @@ function detectGpu() { } } catch { /* ignored */ } - // Fallback: DGX Spark (GB10) — VRAM not queryable due to unified memory architecture + // Fallback: unified-memory NVIDIA devices (VRAM not separately queryable) + // Covers DGX Spark (GB10), Jetson AGX Thor, and Jetson AGX Orin. + const UNIFIED_MEMORY_GPU_TAGS = ["GB10", "Thor", "Orin"]; try { - const nameOutput = runCapture( + const nameOutput = runner.runCapture( "nvidia-smi --query-gpu=name --format=csv,noheader,nounits", { ignoreError: true } ); - if (nameOutput && nameOutput.includes("GB10")) { - // GB10 has 128GB unified memory shared with Grace CPU — use system RAM + if (nameOutput && UNIFIED_MEMORY_GPU_TAGS.some((tag) => nameOutput.includes(tag))) { + // Unified memory shared with CPU — use system RAM as VRAM estimate let totalMemoryMB = 0; try { - const memLine = runCapture("free -m | awk '/Mem:/ {print $2}'", { ignoreError: true }); + const memLine = runner.runCapture("free -m | awk '/Mem:/ {print $2}'", { ignoreError: true }); if (memLine) totalMemoryMB = parseInt(memLine.trim(), 10) || 0; } catch { /* ignored */ } return { @@ -75,7 +78,7 @@ function detectGpu() { // macOS: detect Apple Silicon or discrete GPU if (process.platform === "darwin") { try { - const spOutput = runCapture( + const spOutput = runner.runCapture( "system_profiler SPDisplaysDataType 2>/dev/null", { ignoreError: true } ); @@ -94,7 +97,7 @@ function detectGpu() { } else { // Apple Silicon shares system RAM — read total memory try { - const memBytes = runCapture("sysctl -n hw.memsize", { ignoreError: true }); + const memBytes = runner.runCapture("sysctl -n hw.memsize", { ignoreError: true }); if (memBytes) memoryMB = Math.floor(parseInt(memBytes, 10) / 1024 / 1024); } catch { /* ignored */ } } @@ -158,7 +161,7 @@ function waitForNimHealth(port = 8000, timeout = 300) { while ((Date.now() - start) / 1000 < timeout) { try { - const result = runCapture(`curl -sf http://localhost:${hostPort}/v1/models`, { + const result = runner.runCapture(`curl -sf http://localhost:${hostPort}/v1/models`, { ignoreError: true, }); if (result) { @@ -192,7 +195,7 @@ function nimStatus(sandboxName, port) { function nimStatusByName(name, port) { try { const qn = shellQuote(name); - const state = runCapture( + const state = runner.runCapture( `docker inspect --format '{{.State.Status}}' ${qn} 2>/dev/null`, { ignoreError: true } ); @@ -202,13 +205,13 @@ function nimStatusByName(name, port) { if (state === "running") { let resolvedHostPort = port != null ? Number(port) : 0; if (!resolvedHostPort) { - const mapping = runCapture(`docker port ${qn} 8000 2>/dev/null`, { + const mapping = runner.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( + const health = runner.runCapture( `curl -sf http://localhost:${resolvedHostPort}/v1/models 2>/dev/null`, { ignoreError: true } ); diff --git a/test/nim.test.js b/test/nim.test.js index 468a55c94..7af856432 100644 --- a/test/nim.test.js +++ b/test/nim.test.js @@ -1,11 +1,19 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -import { createRequire } from "module"; -import { describe, it, expect, vi } from "vitest"; -import nim from "../bin/lib/nim"; +import { describe, it, expect, afterEach, vi } from "vitest"; +import { createRequire } from "node:module"; const require = createRequire(import.meta.url); +const runner = require("../bin/lib/runner"); +const nim = require("../bin/lib/nim"); + +const _originalRunCapture = runner.runCapture; +afterEach(() => { + runner.runCapture = _originalRunCapture; + vi.restoreAllMocks(); +}); + const NIM_PATH = require.resolve("../bin/lib/nim"); const RUNNER_PATH = require.resolve("../bin/lib/runner"); @@ -88,6 +96,74 @@ describe("nim", () => { }); }); + describe("detectGpu unified-memory fallback", () => { + /** Build a runCapture mock where VRAM query returns [N/A] and GPU name returns `name`. */ + function mockUnifiedMemoryGpu(name, systemMemMB = "65536") { + runner.runCapture = vi.fn().mockImplementation((cmd) => { + if (cmd.includes("--query-gpu=memory.total")) return "[N/A]"; + if (cmd.includes("--query-gpu=name")) return name; + if (cmd.includes("free -m")) return systemMemMB; + return ""; + }); + } + + it("detects DGX Spark (GB10) via unified-memory fallback", () => { + mockUnifiedMemoryGpu("NVIDIA Graphics Device GB10"); + const gpu = nim.detectGpu(); + expect(gpu).not.toBeNull(); + expect(gpu.type).toBe("nvidia"); + expect(gpu.nimCapable).toBe(true); + expect(gpu.spark).toBe(true); + expect(gpu.totalMemoryMB).toBe(65536); + }); + + it("detects Jetson AGX Thor via unified-memory fallback", () => { + mockUnifiedMemoryGpu("Jetson AGX Thor"); + const gpu = nim.detectGpu(); + expect(gpu).not.toBeNull(); + expect(gpu.type).toBe("nvidia"); + expect(gpu.nimCapable).toBe(true); + expect(gpu.spark).toBe(true); + }); + + it("detects Jetson AGX Orin via unified-memory fallback", () => { + mockUnifiedMemoryGpu("Jetson AGX Orin"); + const gpu = nim.detectGpu(); + expect(gpu).not.toBeNull(); + expect(gpu.type).toBe("nvidia"); + expect(gpu.nimCapable).toBe(true); + expect(gpu.spark).toBe(true); + }); + + it("does not trigger fallback for desktop GPU names", () => { + // Desktop GPUs return valid VRAM, but even if VRAM were [N/A], + // the name must not match any unified-memory tag. + runner.runCapture = vi.fn().mockImplementation((cmd) => { + if (cmd.includes("--query-gpu=memory.total")) return "[N/A]"; + if (cmd.includes("--query-gpu=name")) return "NVIDIA GeForce RTX 4090"; + return ""; + }); + const gpu = nim.detectGpu(); + // Should NOT match as a unified-memory device + if (gpu) { + expect(gpu.spark).toBeUndefined(); + } + }); + + it("skips fallback when VRAM is queryable", () => { + runner.runCapture = vi.fn().mockImplementation((cmd) => { + if (cmd.includes("--query-gpu=memory.total")) return "24564"; + if (cmd.includes("--query-gpu=name")) return "NVIDIA GeForce RTX 4090"; + return ""; + }); + const gpu = nim.detectGpu(); + expect(gpu).not.toBeNull(); + expect(gpu.type).toBe("nvidia"); + expect(gpu.totalMemoryMB).toBe(24564); + expect(gpu.spark).toBeUndefined(); + }); + }); + describe("nimStatus", () => { it("returns not running for nonexistent container", () => { const st = nim.nimStatus("nonexistent-test-xyz");