Skip to content

Commit 169dcb0

Browse files
committed
fix: detect gpu for toolkit cdi preflight
1 parent 9ad9b5d commit 169dcb0

5 files changed

Lines changed: 152 additions & 36 deletions

File tree

src/lib/onboard.ts

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1577,8 +1577,9 @@ function waitForSandboxReady(sandboxName: string, attempts = 10, delaySeconds =
15771577

15781578
// ── Step 1: Preflight ────────────────────────────────────────────
15791579

1580-
// Keep the Docker CDI guard near preflight so resume hits the same early failure path.
1581-
// Jetson/Tegra uses Docker's NVIDIA runtime backend and is exempt from CDI.
1580+
/**
1581+
* Fails onboarding when Docker CDI injection is configured but the NVIDIA GPU spec is invalid.
1582+
*/
15821583
function assertCdiNvidiaGpuSpecPresent(
15831584
host: ReturnType<typeof assessHost>,
15841585
optedOutGpuPassthrough: boolean,
@@ -1601,11 +1602,9 @@ type PreflightOptions = Pick<
16011602
optedOutGpuPassthrough?: boolean;
16021603
};
16031604

1604-
// Reject unsupported container runtimes (currently only Podman with the
1605-
// Linux Docker-driver gateway) before any Docker-specific probes. Both
1606-
// the fresh preflight and `--resume` backstop call this — if `docker`
1607-
// resolves to Podman, surface the unsupported-runtime message instead of
1608-
// running bridge/DNS diagnostics that would be misleading.
1605+
/**
1606+
* Rejects unsupported runtimes before Docker-specific bridge, DNS, and CDI probes.
1607+
*/
16091608
function rejectUnsupportedContainerRuntime(host: ReturnType<typeof assessHost>): void {
16101609
if (isLinuxDockerDriverGatewayEnabled() && host.runtime === "podman") {
16111610
console.error(` ✗ ${cliDisplayName()} onboarding now uses OpenShell's Docker driver.`);
@@ -1615,13 +1614,14 @@ function rejectUnsupportedContainerRuntime(host: ReturnType<typeof assessHost>):
16151614
}
16161615
}
16171616

1617+
/**
1618+
* Runs host preflight and blocks early on Docker, GPU, CDI, and runtime problems.
1619+
*/
16181620
async function preflight(
16191621
preflightOpts: PreflightOptions = {},
16201622
): Promise<ReturnType<typeof nim.detectGpu>> {
16211623
step(1, 8, "Preflight checks");
1622-
16231624
const host = assessHost();
1624-
16251625
// Docker / runtime
16261626
if (!host.dockerReachable) {
16271627
console.error(" Docker is not reachable. Please fix Docker and try again.");
@@ -1640,12 +1640,11 @@ async function preflight(
16401640
device: preflightOpts.sandboxGpuDevice ?? null,
16411641
});
16421642
exitOnSandboxGpuConfigErrors(sandboxGpuConfig);
1643-
const optedOutGpuPassthrough =
1643+
const explicitlyOptedOutGpuPassthrough =
16441644
preflightOpts.optedOutGpuPassthrough === true ||
16451645
preflightOpts.noGpu === true ||
1646-
!sandboxGpuConfig.sandboxGpuEnabled;
1647-
assertCdiNvidiaGpuSpecPresent(host, optedOutGpuPassthrough, sandboxGpuConfig.hostGpuPlatform);
1648-
1646+
sandboxGpuConfig.mode === "0";
1647+
assertCdiNvidiaGpuSpecPresent(host, explicitlyOptedOutGpuPassthrough, sandboxGpuConfig.hostGpuPlatform);
16491648
assertDockerBridgeAndContainerDnsHealthy(host, isNonInteractive());
16501649

16511650
if (host.runtime !== "unknown") {

src/lib/onboard/machine/handlers/preflight.test.ts

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,13 +186,35 @@ describe("handlePreflightState", () => {
186186
expect(harness.deps.startRecordedStep).not.toHaveBeenCalled();
187187
expect(harness.deps.assertCdiNvidiaGpuSpecPresent).toHaveBeenCalledWith(
188188
{ cdiNvidiaGpuSpecMissing: false },
189-
true,
189+
false,
190190
undefined,
191191
);
192192
expect(harness.deps.validateSandboxGpuPreflight).toHaveBeenCalledOnce();
193193
expect(result.resumePreflight).toBe(true);
194194
});
195195

196+
it("keeps CDI guard active on resume when auto mode disables GPU after failed detection", async () => {
197+
const session = createSession();
198+
session.steps.preflight.status = "complete";
199+
session.gpuPassthrough = false;
200+
const assertCdiNvidiaGpuSpecPresent = vi.fn();
201+
const host = { cdiNvidiaGpuSpecMissing: true };
202+
const harness = createDeps({
203+
detectGpu: vi.fn(() => null),
204+
getResumeSandboxGpuOverrides: vi.fn(() => ({ flag: null, device: null })),
205+
resolveSandboxGpuConfig,
206+
assessHost: () => host,
207+
assertCdiNvidiaGpuSpecPresent,
208+
});
209+
210+
await handlePreflightState({
211+
...baseOptions(harness.deps, session),
212+
resume: true,
213+
});
214+
215+
expect(assertCdiNvidiaGpuSpecPresent).toHaveBeenCalledWith(host, false, null);
216+
});
217+
196218
it("passes host GPU platform into the resumed CDI guard", async () => {
197219
const session = createSession();
198220
session.steps.preflight.status = "complete";
@@ -220,7 +242,7 @@ describe("handlePreflightState", () => {
220242

221243
expect(assertCdiNvidiaGpuSpecPresent).toHaveBeenCalledWith(
222244
{ cdiNvidiaGpuSpecMissing: false },
223-
true,
245+
false,
224246
"jetson",
225247
);
226248
});

src/lib/onboard/machine/handlers/preflight.ts

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -93,10 +93,17 @@ export interface PreflightStateResult<Gpu, Config extends PreflightSandboxGpuCon
9393
stateResult: OnboardStateTransitionResult;
9494
}
9595

96+
/**
97+
* Checks whether the environment pins sandbox GPU mode or device selection.
98+
*/
9699
function envHasSandboxGpuOverride(env: NodeJS.ProcessEnv): boolean {
97100
return env.NEMOCLAW_SANDBOX_GPU !== undefined || env.NEMOCLAW_SANDBOX_GPU_DEVICE !== undefined;
98101
}
99102

103+
/**
104+
* Executes or revalidates the preflight state while preserving the user's
105+
* effective sandbox GPU intent across fresh onboarding and resume.
106+
*/
100107
export async function handlePreflightState<
101108
Gpu,
102109
SandboxEntry,
@@ -143,9 +150,7 @@ export async function handlePreflightState<
143150
});
144151
deps.validateSandboxGpuPreflight(resumeSandboxGpuConfig);
145152
const resumeOptedOutGpuPassthrough =
146-
noGpu ||
147-
(!gpuRequested && session?.gpuPassthrough === false) ||
148-
!resumeSandboxGpuConfig.sandboxGpuEnabled;
153+
noGpu || effectiveSandboxGpuFlag === "disable" || resumeSandboxGpuConfig.mode === "0";
149154
const resumeHost = deps.assessHost();
150155
// Reject unsupported runtimes (Podman) BEFORE the CDI GPU-spec
151156
// backstop and the Docker-specific bridge/DNS probes so Podman

src/lib/onboard/preflight-cdi.test.ts

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ import { assessHost, planHostRemediation } from "../../../dist/lib/onboard/prefl
88

99
type HostAssessment = Parameters<typeof planHostRemediation>[0];
1010

11+
/**
12+
* Creates a Linux Docker host assessment with NVIDIA CDI defaults for focused overrides.
13+
*/
1114
function baseAssessment(overrides: Partial<HostAssessment> = {}): HostAssessment {
1215
return {
1316
platform: "linux",
@@ -38,6 +41,9 @@ function baseAssessment(overrides: Partial<HostAssessment> = {}): HostAssessment
3841
};
3942
}
4043

44+
/**
45+
* Emulates the systemctl/stat probes needed by CDI staleness remediation tests.
46+
*/
4147
function healthySystemctlAndStat(command: readonly string[]) {
4248
if (command[0] === "systemctl" && command[1] === "is-enabled") return "enabled";
4349
if (command[0] === "systemctl" && command[1] === "is-active") return "active";
@@ -67,6 +73,51 @@ describe("assessHost — CDI", () => {
6773
expect(result.cdiNvidiaGpuSpecMissing).toBe(true);
6874
});
6975

76+
it("plans toolkit bootstrap when PCI detects NVIDIA hardware but nvidia-smi and nvidia-ctk are absent", () => {
77+
const result = assessHost({
78+
platform: "linux",
79+
env: {},
80+
release: "6.8.0-58-generic",
81+
readFileImpl: (filePath: string) =>
82+
filePath.endsWith("other.yaml")
83+
? "cdiVersion: 0.5.0\nkind: vendor.example/device\ndevices: []\n"
84+
: "Linux version 6.8.0-58-generic",
85+
readdirImpl: (dir: string) => (dir === "/etc/cdi" ? ["other.yaml"] : []),
86+
runCaptureImpl: (command: readonly string[]) =>
87+
new Map([
88+
[["sh", "-c", 'command -v "$1"', "--", "apt-get"].join("\0"), "/usr/bin/apt-get"],
89+
[
90+
["lspci", "-nn"].join("\0"),
91+
"01:00.0 VGA compatible controller: NVIDIA Corporation GA102 [GeForce RTX 3090]\n",
92+
],
93+
[["systemctl", "is-active", "docker"].join("\0"), "active"],
94+
[["systemctl", "is-enabled", "docker"].join("\0"), "enabled"],
95+
]).get(command.join("\0")) ?? "",
96+
dockerInfoOutput: JSON.stringify({
97+
ServerVersion: "27.0",
98+
OperatingSystem: "Ubuntu 24.04",
99+
CDISpecDirs: ["/etc/cdi", "/var/run/cdi"],
100+
}),
101+
commandExistsImpl: (name: string) =>
102+
name === "docker" || name === "lspci" || name === "systemctl",
103+
});
104+
105+
expect(result.hasNvidiaGpu).toBe(true);
106+
expect(result.nvidiaContainerToolkitInstalled).toBe(false);
107+
expect(result.cdiNvidiaGpuSpecMissing).toBe(true);
108+
109+
const action = planHostRemediation(result).find(
110+
(entry: { id: string }) => entry.id === "install_nvidia_container_toolkit",
111+
);
112+
expect(action).toBeTruthy();
113+
expect(action?.blocking).toBe(true);
114+
expect(action?.commands).toContain("sudo apt-get install -y nvidia-container-toolkit");
115+
expect(action?.commands.some((command) => command.includes("nvidia-ctk cdi generate"))).toBe(
116+
true,
117+
);
118+
expect(action?.commands.some((command) => command.includes("nvidia-ctk cdi list"))).toBe(true);
119+
});
120+
70121
it("does not flag the host when an nvidia.com/gpu YAML spec is present", () => {
71122
const result = assessHost({
72123
platform: "linux",

src/lib/onboard/preflight.ts

Lines changed: 57 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -385,13 +385,37 @@ function isHeadlessLikely(env: NodeJS.ProcessEnv): boolean {
385385
return !env.DISPLAY && !env.WAYLAND_DISPLAY && !env.TERM_PROGRAM;
386386
}
387387

388-
function detectNvidiaGpu(runCaptureImpl: RunCaptureFn): boolean {
389-
if (!commandExists("nvidia-smi", runCaptureImpl)) {
390-
return false;
391-
}
392-
return Boolean(String(runCaptureImpl(["nvidia-smi", "-L"], { ignoreError: true }) || "").trim());
388+
/**
389+
* Detects NVIDIA hardware via nvidia-smi first, then Linux PCI data as a toolkit-free fallback.
390+
*/
391+
function detectNvidiaGpu(opts: {
392+
platform: NodeJS.Platform | string;
393+
isWsl: boolean;
394+
runCaptureImpl: RunCaptureFn;
395+
commandExistsImpl?: (commandName: string) => boolean;
396+
}): boolean {
397+
const commandExistsImpl =
398+
opts.commandExistsImpl ??
399+
((commandName: string) => commandExists(commandName, opts.runCaptureImpl));
400+
if (commandExistsImpl("nvidia-smi")) {
401+
const smiOutput = opts.runCaptureImpl(["nvidia-smi", "-L"], { ignoreError: true });
402+
if (String(smiOutput || "").trim()) return true;
403+
}
404+
405+
if (opts.platform !== "linux" || opts.isWsl || !commandExistsImpl("lspci")) return false;
406+
const pciOutput = opts.runCaptureImpl(["lspci", "-nn"], { ignoreError: true });
407+
return String(pciOutput || "")
408+
.split(/\r?\n/)
409+
.some(
410+
(line) =>
411+
/nvidia/i.test(line) &&
412+
/(vga compatible controller|3d controller|display controller)/i.test(line),
413+
);
393414
}
394415

416+
/**
417+
* Detects the host package manager used for NVIDIA toolkit remediation commands.
418+
*/
395419
function detectPackageManager(runCaptureImpl: RunCaptureFn): PackageManager {
396420
if (commandExists("apt-get", runCaptureImpl)) return "apt";
397421
if (commandExists("dnf", runCaptureImpl)) return "dnf";
@@ -401,6 +425,9 @@ function detectPackageManager(runCaptureImpl: RunCaptureFn): PackageManager {
401425
return "unknown";
402426
}
403427

428+
/**
429+
* Normalizes systemctl active/enabled output into a tri-state service status.
430+
*/
404431
function parseSystemctlState(value = ""): boolean | null {
405432
const normalized = String(value || "")
406433
.trim()
@@ -447,6 +474,9 @@ export function buildContainerToolkitBootstrapCommands(
447474
];
448475
}
449476

477+
/**
478+
* Builds the host capability snapshot used to plan preflight remediation.
479+
*/
450480
export function assessHost(opts: AssessHostOpts = {}): HostAssessment {
451481
const platform = opts.platform ?? process.platform;
452482
const env = opts.env ?? process.env;
@@ -456,12 +486,33 @@ export function assessHost(opts: AssessHostOpts = {}): HostAssessment {
456486
runCapture(command, { ignoreError: options?.ignoreError ?? false }));
457487
const readFileImpl = opts.readFileImpl ?? fs.readFileSync;
458488
const readdirImpl = opts.readdirImpl ?? ((dir: string) => fs.readdirSync(dir));
489+
const shouldReadLinuxHostDetails = platform === "linux";
490+
const release = opts.release ?? (shouldReadLinuxHostDetails ? os.release() : "");
491+
const procVersion =
492+
opts.procVersion ??
493+
(shouldReadLinuxHostDetails
494+
? (() => {
495+
try {
496+
return readFileImpl("/proc/version", "utf-8");
497+
} catch {
498+
return "";
499+
}
500+
})()
501+
: "");
502+
const isWslHost = detectWsl({ platform, env, release, procVersion });
459503
const dockerInstalled =
460504
opts.commandExistsImpl?.("docker") ?? commandExists("docker", runCaptureImpl);
461505
const nodeInstalled = opts.commandExistsImpl?.("node") ?? commandExists("node", runCaptureImpl);
462506
const openshellInstalled =
463507
opts.commandExistsImpl?.("openshell") ?? commandExists("openshell", runCaptureImpl);
464-
const hasNvidiaGpu = opts.gpuProbeImpl?.() ?? detectNvidiaGpu(runCaptureImpl);
508+
const hasNvidiaGpu =
509+
opts.gpuProbeImpl?.() ??
510+
detectNvidiaGpu({
511+
platform,
512+
isWsl: isWslHost,
513+
runCaptureImpl,
514+
commandExistsImpl: opts.commandExistsImpl,
515+
});
465516
const nvidiaContainerToolkitInstalled =
466517
opts.commandExistsImpl?.("nvidia-ctk") ?? commandExists("nvidia-ctk", runCaptureImpl);
467518
const packageManager = detectPackageManager(runCaptureImpl);
@@ -480,22 +531,10 @@ export function assessHost(opts: AssessHostOpts = {}): HostAssessment {
480531
dockerReachable = true;
481532
dockerRunning = true;
482533
}
483-
484-
const release = opts.release ?? os.release();
485-
const procVersion =
486-
opts.procVersion ??
487-
(() => {
488-
try {
489-
return readFileImpl("/proc/version", "utf-8");
490-
} catch {
491-
return "";
492-
}
493-
})();
494534
let runtime = inferContainerRuntime(dockerInfoOutput);
495535
if (dockerReachable && runtime === "unknown" && platform === "linux") {
496536
runtime = "docker";
497537
}
498-
const isWslHost = detectWsl({ platform, env, release, procVersion });
499538
const dockerCgroupVersion = dockerReachable
500539
? parseDockerCgroupVersion(dockerInfoOutput)
501540
: "unknown";

0 commit comments

Comments
 (0)