diff --git a/src/lib/onboard.ts b/src/lib/onboard.ts index 461eb3045ff..2b985b6d498 100644 --- a/src/lib/onboard.ts +++ b/src/lib/onboard.ts @@ -294,6 +294,7 @@ const { resolveSandboxImageTagFromCreateOutput } = const nim: typeof import("./inference/nim") = require("./inference/nim"); const onboardSession: typeof import("./state/onboard-session") = require("./state/onboard-session"); const { OnboardRuntimeBoundary }: typeof import("./onboard/runtime-boundary") = require("./onboard/runtime-boundary"); +const { handlePreflightState }: typeof import("./onboard/machine/handlers/preflight") = require("./onboard/machine/handlers/preflight"); const policies: typeof import("./policy") = require("./policy"); const tiers: typeof import("./policy/tiers") = require("./policy/tiers"); const { ensureUsageNoticeConsent } = require("./onboard/usage-notice"); @@ -9283,7 +9284,6 @@ async function onboard(opts: OnboardOptions = {}): Promise { const recordedSandboxName = session?.steps?.sandbox?.status === "complete" ? session?.sandboxName || null : null; - const resumeSandboxNameForGpu = recordedSandboxName || requestedSandboxName || null; console.log(""); console.log(` ${cliDisplayName()} Onboarding`); @@ -9292,59 +9292,46 @@ async function onboard(opts: OnboardOptions = {}): Promise { console.log(" ==================="); const explicitSandboxGpuFlag = resolveSandboxGpuFlagFromOptions(opts); - const resumePreflight = resume && session?.steps?.preflight?.status === "complete"; - const resumeHasResolvedGpuIntent = - resumePreflight && - explicitSandboxGpuFlag === null && - opts.sandboxGpuDevice == null && - process.env.NEMOCLAW_SANDBOX_GPU === undefined && - process.env.NEMOCLAW_SANDBOX_GPU_DEVICE === undefined; - const resumedSandboxGpuOverrides = resumeHasResolvedGpuIntent - ? getResumeSandboxGpuOverrides( - resumeSandboxNameForGpu ? registry.getSandbox(resumeSandboxNameForGpu) : null, - session?.gpuPassthrough, - ) - : { flag: null, device: null }; - const effectiveSandboxGpuFlag = explicitSandboxGpuFlag ?? resumedSandboxGpuOverrides.flag; - const effectiveSandboxGpuDevice = opts.sandboxGpuDevice ?? resumedSandboxGpuOverrides.device; - let gpu; - if (resumePreflight) { - skippedStepMessage("preflight", "cached"); - gpu = nim.detectGpu(); - const resumeSandboxGpuConfig = resolveSandboxGpuConfig(gpu, { - flag: effectiveSandboxGpuFlag, - device: effectiveSandboxGpuDevice, - }); - exitOnSandboxGpuConfigErrors(resumeSandboxGpuConfig); - // Re-check the CDI spec gap on resume (#3152). The cached preflight - // result does not capture host CDI state, and the original onboard - // attempt that wrote the cache likely aborted at gateway-start with - // exactly this CDI failure — so resuming without re-checking would - // walk into the same wall. Honour persisted `gpuPassthrough: false` - // from the prior session as an opt-out, since the resume invocation - // does not need to re-pass `--no-gpu` to keep that intent (the same - // resolution is replayed a few lines below for `gpuPassthrough`). - const resumeOptedOutGpuPassthrough = - opts.noGpu === true || - (opts.gpu !== true && session?.gpuPassthrough === false) || - !resumeSandboxGpuConfig.sandboxGpuEnabled; - assertCdiNvidiaGpuSpecPresent(assessHost(), resumeOptedOutGpuPassthrough); - validateSandboxGpuPreflight(resumeSandboxGpuConfig); - } else { - await startRecordedStep("preflight"); - gpu = await preflight({ ...opts, optedOutGpuPassthrough: opts.noGpu === true }); - await recordStepComplete("preflight"); - } - const sandboxGpuConfig = resolveSandboxGpuConfig(gpu, { - flag: effectiveSandboxGpuFlag, - device: effectiveSandboxGpuDevice, + const recordedGpuPassthroughBeforePreflight = session?.gpuPassthrough === true; + const preflightResult = await handlePreflightState({ + resume, + session, + recordedSandboxName, + requestedSandboxName, + explicitSandboxGpuFlag, + sandboxGpuDevice: opts.sandboxGpuDevice ?? null, + gpuRequested: opts.gpu === true, + noGpu: opts.noGpu === true, + env: process.env, + deps: { + getSandbox: registry.getSandbox.bind(registry), + getResumeSandboxGpuOverrides, + detectGpu: nim.detectGpu, + runPreflight: (preflightOptions) => preflight({ ...opts, ...preflightOptions }), + assessHost, + assertCdiNvidiaGpuSpecPresent, + resolveSandboxGpuConfig, + validateSandboxGpuPreflight: (config) => { + exitOnSandboxGpuConfigErrors(config); + validateSandboxGpuPreflight(config); + }, + skippedStepMessage, + startRecordedStep, + recordStepComplete, + updateSession: onboardSession.updateSession, + }, }); - - const requestedGpuPassthrough = opts.gpu === true; - const gpuPassthrough = sandboxGpuConfig.sandboxGpuEnabled; + session = preflightResult.session; + const { + sandboxGpuConfig, + resumeHasResolvedGpuIntent, + requestedGpuPassthrough, + gpuPassthrough, + } = preflightResult; + const gpu = preflightResult.gpu ?? null; if (gpuPassthrough) { note( - resumeHasResolvedGpuIntent && session?.gpuPassthrough === true + resumeHasResolvedGpuIntent && recordedGpuPassthroughBeforePreflight ? " [resume] Continuing GPU passthrough from the saved onboarding session." : requestedGpuPassthrough || sandboxGpuConfig.mode === "1" ? " GPU passthrough requested; passing --gpu to OpenShell gateway and sandbox creation." @@ -9363,13 +9350,6 @@ async function onboard(opts: OnboardOptions = {}): Promise { /* lspci not available — skip hint */ } } - // Persist GPU intent in the session so resume can restore it. - if (session && session.gpuPassthrough !== gpuPassthrough) { - session = onboardSession.updateSession((current: Session) => { - current.gpuPassthrough = gpuPassthrough; - return current; - }); - } dockerGpuLocalInference.configureLocalInferenceForDockerGpuHostNetwork(sandboxGpuConfig, { dockerDriverGateway: isLinuxDockerDriverGatewayEnabled(), note, diff --git a/src/lib/onboard/machine/handlers/preflight.test.ts b/src/lib/onboard/machine/handlers/preflight.test.ts new file mode 100644 index 00000000000..fa4b859915a --- /dev/null +++ b/src/lib/onboard/machine/handlers/preflight.test.ts @@ -0,0 +1,183 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it, vi } from "vitest"; + +import { createSession, type Session } from "../../../state/onboard-session"; +import { handlePreflightState, type PreflightStateOptions } from "./preflight"; + +type Gpu = { type: string } | null; +type SandboxEntry = { sandboxGpuEnabled?: boolean }; +type Host = { cdiNvidiaGpuSpecMissing?: boolean }; + +function createDeps(overrides: Partial["deps"]> = {}) { + let session = createSession(); + return { + calls: { + start: vi.fn(), + complete: vi.fn(), + skipped: vi.fn(), + detectGpu: vi.fn(() => ({ type: "nvidia" }) as Gpu), + runPreflight: vi.fn(async () => ({ type: "nvidia" }) as Gpu), + validate: vi.fn(), + cdi: vi.fn(), + updateSession: vi.fn(), + getSandbox: vi.fn(() => ({ sandboxGpuEnabled: true })), + getOverrides: vi.fn(() => ({ flag: "enable" as const, device: "0" })), + }, + deps: { + getSandbox: (name: string) => { + const value = ({ sandboxGpuEnabled: true } satisfies SandboxEntry); + return overrides.getSandbox ? overrides.getSandbox(name) : value; + }, + getResumeSandboxGpuOverrides: ( + sandbox: SandboxEntry | null, + sessionGpuPassthrough: boolean | null | undefined, + ) => { + if (overrides.getResumeSandboxGpuOverrides) { + return overrides.getResumeSandboxGpuOverrides(sandbox, sessionGpuPassthrough); + } + return { flag: "enable" as const, device: "0" }; + }, + detectGpu: () => ({ type: "nvidia" }) as Gpu, + runPreflight: async () => ({ type: "nvidia" }) as Gpu, + assessHost: () => ({ cdiNvidiaGpuSpecMissing: false }), + assertCdiNvidiaGpuSpecPresent: vi.fn(), + resolveSandboxGpuConfig: (_gpu: Gpu, opts: { flag: "enable" | "disable" | null; device: string | null | undefined }) => ({ + sandboxGpuEnabled: opts.flag === "enable", + mode: opts.flag === "enable" ? "1" : "0", + sandboxGpuDevice: opts.device, + }), + validateSandboxGpuPreflight: vi.fn(), + skippedStepMessage: vi.fn(), + startRecordedStep: vi.fn(async () => undefined), + recordStepComplete: vi.fn(async () => session), + updateSession: vi.fn((mutator: (value: Session) => Session | void) => { + session = mutator(session) ?? session; + return session; + }), + ...overrides, + }, + getSession: () => session, + }; +} + +function baseOptions( + deps: PreflightStateOptions["deps"], + session: Session | null = createSession(), +): PreflightStateOptions { + return { + resume: false, + session, + recordedSandboxName: null, + requestedSandboxName: "my-assistant", + explicitSandboxGpuFlag: null, + sandboxGpuDevice: null, + gpuRequested: false, + noGpu: false, + env: {}, + deps, + }; +} + +describe("handlePreflightState", () => { + it("runs full preflight through recorded step boundaries", async () => { + const harness = createDeps({ + startRecordedStep: vi.fn(async () => undefined), + runPreflight: vi.fn(async () => ({ type: "nvidia" }) as Gpu), + recordStepComplete: vi.fn(async () => createSession()), + }); + + const result = await handlePreflightState({ + ...baseOptions(harness.deps), + explicitSandboxGpuFlag: "enable", + sandboxGpuDevice: "GPU-0", + }); + + expect(harness.deps.startRecordedStep).toHaveBeenCalledWith("preflight"); + expect(harness.deps.runPreflight).toHaveBeenCalledWith({ optedOutGpuPassthrough: false }); + expect(harness.deps.recordStepComplete).toHaveBeenCalledWith("preflight"); + expect(result.sandboxGpuConfig).toMatchObject({ + sandboxGpuEnabled: true, + mode: "1", + sandboxGpuDevice: "GPU-0", + }); + expect(result.gpuPassthrough).toBe(true); + }); + + it("skips full preflight on resume but re-detects GPU and revalidates CDI/sandbox GPU", async () => { + const session = createSession(); + session.steps.preflight.status = "complete"; + session.gpuPassthrough = false; + const harness = createDeps({ + detectGpu: vi.fn(() => ({ type: "nvidia" }) as Gpu), + assertCdiNvidiaGpuSpecPresent: vi.fn(), + validateSandboxGpuPreflight: vi.fn(), + skippedStepMessage: vi.fn(), + startRecordedStep: vi.fn(async () => undefined), + runPreflight: vi.fn(async () => ({ type: "should-not-run" }) as Gpu), + }); + + const result = await handlePreflightState({ + ...baseOptions(harness.deps, session), + resume: true, + gpuRequested: false, + }); + + expect(harness.deps.skippedStepMessage).toHaveBeenCalledWith("preflight", "cached"); + expect(harness.deps.detectGpu).toHaveBeenCalledOnce(); + expect(harness.deps.runPreflight).not.toHaveBeenCalled(); + expect(harness.deps.startRecordedStep).not.toHaveBeenCalled(); + expect(harness.deps.assertCdiNvidiaGpuSpecPresent).toHaveBeenCalledWith( + { cdiNvidiaGpuSpecMissing: false }, + true, + ); + expect(harness.deps.validateSandboxGpuPreflight).toHaveBeenCalledOnce(); + expect(result.resumePreflight).toBe(true); + }); + + it("restores saved sandbox GPU intent only when resume has no explicit override", async () => { + const session = createSession(); + session.steps.preflight.status = "complete"; + session.gpuPassthrough = true; + const getResumeSandboxGpuOverrides = vi.fn(() => ({ flag: "enable" as const, device: "1" })); + const getSandbox = vi.fn(() => ({ sandboxGpuEnabled: true })); + const harness = createDeps({ getResumeSandboxGpuOverrides, getSandbox }); + + const result = await handlePreflightState({ + ...baseOptions(harness.deps, session), + resume: true, + recordedSandboxName: "saved", + }); + + expect(getSandbox).toHaveBeenCalledWith("saved"); + expect(getResumeSandboxGpuOverrides).toHaveBeenCalledWith( + { sandboxGpuEnabled: true }, + true, + ); + expect(result.resumeHasResolvedGpuIntent).toBe(true); + expect(result.effectiveSandboxGpuFlag).toBe("enable"); + expect(result.effectiveSandboxGpuDevice).toBe("1"); + + await handlePreflightState({ + ...baseOptions(harness.deps, session), + resume: true, + explicitSandboxGpuFlag: "disable", + }); + expect(getResumeSandboxGpuOverrides).toHaveBeenCalledTimes(1); + }); + + it("persists effective GPU passthrough intent for later resume", async () => { + const session = createSession(); + session.gpuPassthrough = false; + const harness = createDeps(); + + const result = await handlePreflightState({ + ...baseOptions(harness.deps, session), + explicitSandboxGpuFlag: "enable", + }); + + expect(result.session?.gpuPassthrough).toBe(true); + expect(harness.deps.updateSession).toHaveBeenCalledOnce(); + }); +}); diff --git a/src/lib/onboard/machine/handlers/preflight.ts b/src/lib/onboard/machine/handlers/preflight.ts new file mode 100644 index 00000000000..ccecf7860db --- /dev/null +++ b/src/lib/onboard/machine/handlers/preflight.ts @@ -0,0 +1,147 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import type { Session } from "../../../state/onboard-session"; + +export type PreflightSandboxGpuFlag = "enable" | "disable" | null; + +export interface PreflightSandboxGpuOverrides { + flag: PreflightSandboxGpuFlag; + device: string | null; +} + +export interface PreflightSandboxGpuConfig { + sandboxGpuEnabled: boolean; + mode: string; + sandboxGpuDevice?: string | null; + errors?: readonly string[]; +} + +export interface PreflightStateOptions< + Gpu, + SandboxEntry, + Host, + Config extends PreflightSandboxGpuConfig, +> { + resume: boolean; + session: Session | null; + recordedSandboxName: string | null; + requestedSandboxName: string | null; + explicitSandboxGpuFlag: PreflightSandboxGpuFlag; + sandboxGpuDevice?: string | null; + gpuRequested: boolean; + noGpu: boolean; + env: NodeJS.ProcessEnv; + deps: { + getSandbox(name: string): SandboxEntry | null; + getResumeSandboxGpuOverrides( + sandbox: SandboxEntry | null, + sessionGpuPassthrough: boolean | null | undefined, + ): PreflightSandboxGpuOverrides; + detectGpu(): Gpu; + runPreflight(options: { optedOutGpuPassthrough?: boolean }): Promise; + assessHost(): Host; + assertCdiNvidiaGpuSpecPresent(host: Host, optedOutGpuPassthrough: boolean): void; + resolveSandboxGpuConfig( + gpu: Gpu, + options: { flag: PreflightSandboxGpuFlag; device: string | null | undefined }, + ): Config; + validateSandboxGpuPreflight(config: Config): void; + skippedStepMessage(stepName: string, detail?: string | null): void; + startRecordedStep(stepName: string): Promise; + recordStepComplete(stepName: string): Promise; + updateSession(mutator: (session: Session) => Session | void): Session; + }; +} + +export interface PreflightStateResult { + gpu: Gpu; + sandboxGpuConfig: Config; + resumePreflight: boolean; + resumeHasResolvedGpuIntent: boolean; + requestedGpuPassthrough: boolean; + gpuPassthrough: boolean; + effectiveSandboxGpuFlag: PreflightSandboxGpuFlag; + effectiveSandboxGpuDevice: string | null | undefined; + session: Session | null; +} + +function envHasSandboxGpuOverride(env: NodeJS.ProcessEnv): boolean { + return env.NEMOCLAW_SANDBOX_GPU !== undefined || env.NEMOCLAW_SANDBOX_GPU_DEVICE !== undefined; +} + +export async function handlePreflightState< + Gpu, + SandboxEntry, + Host, + Config extends PreflightSandboxGpuConfig, +>({ + resume, + session, + recordedSandboxName, + requestedSandboxName, + explicitSandboxGpuFlag, + sandboxGpuDevice, + gpuRequested, + noGpu, + env, + deps, +}: PreflightStateOptions): Promise> { + const resumeSandboxNameForGpu = recordedSandboxName || requestedSandboxName || null; + const resumePreflight = resume && session?.steps?.preflight?.status === "complete"; + const resumeHasResolvedGpuIntent = + resumePreflight && + explicitSandboxGpuFlag === null && + sandboxGpuDevice == null && + !envHasSandboxGpuOverride(env); + const resumedSandboxGpuOverrides = resumeHasResolvedGpuIntent + ? deps.getResumeSandboxGpuOverrides( + resumeSandboxNameForGpu ? deps.getSandbox(resumeSandboxNameForGpu) : null, + session?.gpuPassthrough, + ) + : { flag: null, device: null }; + const effectiveSandboxGpuFlag = explicitSandboxGpuFlag ?? resumedSandboxGpuOverrides.flag; + const effectiveSandboxGpuDevice = sandboxGpuDevice ?? resumedSandboxGpuOverrides.device; + + let gpu: Gpu; + if (resumePreflight) { + deps.skippedStepMessage("preflight", "cached"); + gpu = deps.detectGpu(); + const resumeSandboxGpuConfig = deps.resolveSandboxGpuConfig(gpu, { + flag: effectiveSandboxGpuFlag, + device: effectiveSandboxGpuDevice, + }); + deps.validateSandboxGpuPreflight(resumeSandboxGpuConfig); + const resumeOptedOutGpuPassthrough = + noGpu || (!gpuRequested && session?.gpuPassthrough === false) || !resumeSandboxGpuConfig.sandboxGpuEnabled; + deps.assertCdiNvidiaGpuSpecPresent(deps.assessHost(), resumeOptedOutGpuPassthrough); + } else { + await deps.startRecordedStep("preflight"); + gpu = await deps.runPreflight({ optedOutGpuPassthrough: noGpu }); + session = await deps.recordStepComplete("preflight"); + } + + const sandboxGpuConfig = deps.resolveSandboxGpuConfig(gpu, { + flag: effectiveSandboxGpuFlag, + device: effectiveSandboxGpuDevice, + }); + const gpuPassthrough = sandboxGpuConfig.sandboxGpuEnabled; + if (session && session.gpuPassthrough !== gpuPassthrough) { + session = deps.updateSession((current) => { + current.gpuPassthrough = gpuPassthrough; + return current; + }); + } + + return { + gpu, + sandboxGpuConfig, + resumePreflight, + resumeHasResolvedGpuIntent, + requestedGpuPassthrough: gpuRequested, + gpuPassthrough, + effectiveSandboxGpuFlag, + effectiveSandboxGpuDevice, + session, + }; +}