Skip to content

Fix/jetson gpu detection #1106

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Junior00619 wants to merge 1 commit into from
+95 −16
Closed

Fix/jetson gpu detection #1106

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 16 additions & 13 deletions bin/lib/nim.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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 }
);
Expand All @@ -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 {
Expand All @@ -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 }
);
Expand All @@ -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 */ }
}
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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 }
);
Expand All @@ -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 }
);
Expand Down
82 changes: 79 additions & 3 deletions test/nim.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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");

Expand Down Expand Up @@ -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();
}
Comment on lines +138 to +150
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Tighten the negative-case assertion to avoid false positives.

The if (gpu) guard allows this test to pass even when it validates nothing. Assert that unified-memory fallback was not taken by checking command behavior (no free -m call), and keep the spark assertion.

Suggested test hardening 
       const gpu = nim.detectGpu();
-      // Should NOT match as a unified-memory device
-      if (gpu) {
-        expect(gpu.spark).toBeUndefined();
-      }
+      const commands = runner.runCapture.mock.calls.map(([cmd]) => cmd);
+      expect(commands.some((cmd) => cmd.includes("--query-gpu=name"))).toBe(true);
+      expect(commands.some((cmd) => cmd.includes("free -m"))).toBe(false);
+      expect(gpu?.spark).toBeUndefined();
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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("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();
const commands = runner.runCapture.mock.calls.map(([cmd]) => cmd);
expect(commands.some((cmd) => cmd.includes("--query-gpu=name"))).toBe(true);
expect(commands.some((cmd) => cmd.includes("free -m"))).toBe(false);
expect(gpu?.spark).toBeUndefined();
});
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@test/nim.test.js` around lines 138 - 150, The test currently uses an if (gpu)
guard which can hide failures; instead assert the GPU was detected and that
unified-memory fallback was not taken: first assert gpu is defined (or truthy)
after calling nim.detectGpu(), then assert runner.runCapture was never called
with a command containing "free -m" (to ensure no fallback probe occurred), and
finally assert gpu.spark is undefined to confirm no unified-memory tagging;
reference nim.detectGpu and runner.runCapture in your changes.

});

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");
Expand Down