diff --git a/src/lib/inference/vllm.test.ts b/src/lib/inference/vllm.test.ts index 6f2e7db2d95..e8c6e37b5dd 100644 --- a/src/lib/inference/vllm.test.ts +++ b/src/lib/inference/vllm.test.ts @@ -1,33 +1,54 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +import { EventEmitter } from "node:events"; +import os from "node:os"; +import path from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; const mocks = vi.hoisted(() => ({ dockerCapture: vi.fn(), + dockerForceRm: vi.fn(), dockerPullWithProgressWatchdog: vi.fn(), + dockerRunDetached: vi.fn(), dockerSpawn: vi.fn(), - getGpuIndicesByName: vi.fn(() => []), + dockerStop: vi.fn(), + getGpuIndicesByName: vi.fn<(_pattern: RegExp) => number[]>(() => []), runCapture: vi.fn(), - runShell: vi.fn(), })); vi.mock("../runner", () => ({ runCapture: mocks.runCapture, - runShell: mocks.runShell, })); vi.mock("../adapters/docker", () => ({ dockerCapture: mocks.dockerCapture, + dockerForceRm: mocks.dockerForceRm, dockerPullWithProgressWatchdog: mocks.dockerPullWithProgressWatchdog, + dockerRunDetached: mocks.dockerRunDetached, dockerSpawn: mocks.dockerSpawn, + dockerStop: mocks.dockerStop, })); vi.mock("./nim", () => ({ getGpuIndicesByName: mocks.getGpuIndicesByName, })); -import { buildVllmRunCommand, detectVllmProfile, installVllm, pullImage } from "./vllm"; +import { buildVllmRunArgs, detectVllmProfile, installVllm, pullImage } from "./vllm"; + +function mockDockerSpawnSuccess(): EventEmitter & { + stdout: EventEmitter; + stderr: EventEmitter; +} { + const proc = new EventEmitter() as EventEmitter & { + stdout: EventEmitter; + stderr: EventEmitter; + }; + proc.stdout = new EventEmitter(); + proc.stderr = new EventEmitter(); + process.nextTick(() => proc.emit("exit", 0)); + return proc; +} describe("vLLM profile detection", () => { beforeEach(() => { @@ -128,23 +149,58 @@ describe("vLLM run command", () => { it("adds --restart unless-stopped so the container survives a host reboot (#4886)", () => { const profile = detectVllmProfile({ platform: "spark", type: "nvidia" }); expect(profile).not.toBeNull(); - const cmd = buildVllmRunCommand( - profile!, - profile!.defaultModel, - profile!.dockerRunFlags.join(" "), + const args = buildVllmRunArgs(profile!, profile!.defaultModel, profile!.dockerRunFlags); + expect(args.slice(0, 2)).toEqual(["--restart", "unless-stopped"]); + expect(args).toContain("--name"); + expect(args[args.indexOf("--name") + 1]).toBe(profile!.containerName); + expect(args).toContain("8000:8000"); + }); + + it("preserves profile run flags and image as argv tokens", () => { + const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); + expect(profile).not.toBeNull(); + const args = buildVllmRunArgs(profile!, profile!.defaultModel, [ + "--gpus", + "device=0,1", + "--ipc=host", + ]); + expect(args).toEqual(expect.arrayContaining(["--gpus", "device=0,1", "--ipc=host"])); + expect(args).toContain(profile!.image); + expect(args).toEqual(expect.arrayContaining(["--entrypoint", "/bin/bash"])); + expect(args.join(" ")).not.toContain("docker run"); + }); + + it("keeps shell metacharacters in Docker argv tokens instead of shell composing them", () => { + const profile = detectVllmProfile({ platform: "spark", type: "nvidia" }); + expect(profile).not.toBeNull(); + const labelValue = "profile=$(touch /tmp/nemoclaw-vllm-pwn)"; + const args = buildVllmRunArgs(profile!, profile!.defaultModel, ["--label", labelValue], { + HF_TOKEN: "hf_test", + } as NodeJS.ProcessEnv); + + expect(args).toEqual(expect.arrayContaining(["--label", labelValue, "-e", "HF_TOKEN"])); + expect(args).not.toContain(`--label ${labelValue}`); + expect(args).not.toContain("-e HF_TOKEN"); + }); + + it("uses os.homedir for the Hugging Face cache mount without shell quoting", () => { + const profile = detectVllmProfile({ platform: "spark", type: "nvidia" }); + expect(profile).not.toBeNull(); + const mount = profile!.dockerRunFlags[profile!.dockerRunFlags.indexOf("-v") + 1]; + + expect(mount).toBe( + `${path.join(os.homedir(), ".cache", "huggingface")}:/root/.cache/huggingface`, ); - expect(cmd).toContain("docker run -d --restart unless-stopped"); - expect(cmd).toContain(`--name ${profile!.containerName}`); - expect(cmd).toContain(":8000"); }); - it("preserves the profile run flags and image", () => { + it("builds the Station multi-GPU flag without shell-only quotes", () => { + mocks.getGpuIndicesByName.mockReturnValue([0, 1]); const profile = detectVllmProfile({ platform: "station", type: "nvidia" }); expect(profile).not.toBeNull(); - const cmd = buildVllmRunCommand(profile!, profile!.defaultModel, "--gpus device=0 --ipc=host"); - expect(cmd).toContain("--restart unless-stopped --gpus device=0 --ipc=host"); - expect(cmd).toContain(profile!.image); - expect(cmd).toContain("--entrypoint /bin/bash"); + const flags = profile!.buildDockerRunFlags!(); + + expect(flags).toEqual(expect.arrayContaining(["--gpus", "device=0,1"])); + expect(flags).not.toContain(`'"device=0,1"'`); }); }); @@ -244,4 +300,47 @@ describe("installVllm model resolution", () => { const errors = errSpy.mock.calls.map((c: unknown[]) => String(c[0])).join("\n"); expect(errors).toMatch(/gated on Hugging Face/); }); + + it("starts the long-lived vLLM container through Docker argv, not a shell command", async () => { + process.env.HF_TOKEN = "hf_test"; + const profile = detectVllmProfile({ platform: "spark", type: "nvidia" })!; + mocks.runCapture.mockImplementation((cmd: readonly string[]) => { + if (cmd[0] === "sh") return "/usr/bin/tool\n"; + if (cmd[0] === "curl") return '{"data":[]}'; + return ""; + }); + mocks.dockerPullWithProgressWatchdog.mockResolvedValue({ + status: 0, + signal: null, + output: "", + timedOut: false, + timeoutKind: null, + }); + mocks.dockerSpawn.mockReturnValue(mockDockerSpawnSuccess()); + mocks.dockerRunDetached.mockReturnValue({ status: 0, stdout: "", stderr: "", error: null }); + mocks.dockerCapture.mockReturnValue(`${profile.containerName}\n`); + + const result = await installVllm(profile, { + hasImage: true, + nonInteractive: true, + promptFn: vi.fn(), + }); + + expect(result).toEqual({ ok: true }); + expect(mocks.dockerForceRm).toHaveBeenCalledWith( + profile.containerName, + expect.objectContaining({ ignoreError: true, suppressOutput: true }), + ); + expect(mocks.dockerRunDetached).toHaveBeenCalledTimes(1); + const [args, opts] = mocks.dockerRunDetached.mock.calls[0] as [ + string[], + { env?: Record }, + ]; + expect(args).toEqual( + expect.arrayContaining(["--restart", "unless-stopped", "-e", "HF_TOKEN", profile.image]), + ); + expect(args.some((arg) => arg.includes("docker run"))).toBe(false); + expect(args[args.indexOf("-lc") + 1]).toContain("vllm serve"); + expect(opts).toEqual(expect.objectContaining({ env: { HF_TOKEN: "hf_test" } })); + }); }); diff --git a/src/lib/inference/vllm.ts b/src/lib/inference/vllm.ts index 303ef4c2629..030855ba4bc 100644 --- a/src/lib/inference/vllm.ts +++ b/src/lib/inference/vllm.ts @@ -5,10 +5,19 @@ // offer vLLM at all" lives in onboard.ts; this module owns picking the // right profile per platform and running the install. -import { dockerCapture, dockerPullWithProgressWatchdog, dockerSpawn } from "../adapters/docker"; +import os from "node:os"; +import path from "node:path"; +import { + dockerCapture, + dockerForceRm, + dockerPullWithProgressWatchdog, + dockerRunDetached, + dockerSpawn, + dockerStop, +} from "../adapters/docker"; import { buildValidatedCurlCommandArgs } from "../adapters/http/curl-args"; import { VLLM_PORT } from "../core/ports"; -import { runCapture, runShell } from "../runner"; +import { runCapture } from "../runner"; import { getGpuIndicesByName } from "./nim"; import { VLLM_EXTRA_ARGS_ENV, @@ -77,6 +86,27 @@ function qwen35bNvfp4Model(): VllmModelDef { const HF_TOKEN_ENV_KEYS = ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"] as const; const MODEL_DOWNLOAD_HEARTBEAT_MS = 30_000; const VLLM_LAUNCH_HEARTBEAT_MS = 30_000; +const HF_CACHE_CONTAINER_DIR = "/root/.cache/huggingface"; + +function hostHfCacheDir(): string { + return path.join(os.homedir(), ".cache", "huggingface"); +} + +function hfCacheMount(): string { + return `${hostHfCacheDir()}:${HF_CACHE_CONTAINER_DIR}`; +} + +function vllmDockerRunFlags(gpuFlag = "all"): string[] { + return [ + "--gpus", + gpuFlag, + "--ipc=host", + "-v", + hfCacheMount(), + "-e", + `HF_HOME=${HF_CACHE_CONTAINER_DIR}`, + ]; +} function pickHfTokenEntry( env: NodeJS.ProcessEnv = process.env, @@ -108,10 +138,9 @@ export function buildHfTokenDockerArgs(env: NodeJS.ProcessEnv = process.env): st /** * Companion to `buildHfTokenDockerArgs`: returns the `{ KEY: value }` map * that has to be merged into the subprocess env so docker can see the - * token when `-e KEY` (key-only) tells it to forward by name. The CLI's - * `runShell` strips non-allowlisted env names by default (see - * subprocess-env.ts), so callers that go through that path must pass - * this map via the runner's `env` option. + * token when `-e KEY` (key-only) tells it to forward by name. The CLI runner + * strips non-allowlisted env names by default (see subprocess-env.ts), so + * Docker callers must pass this map via the runner's `env` option. */ export function buildHfTokenForwardEnv( env: NodeJS.ProcessEnv = process.env, @@ -126,15 +155,7 @@ const SPARK_PROFILE: VllmProfile = { image: VLLM_IMAGES.ngc2605Post1, defaultModel: qwen35bNvfp4Model(), containerName: "nemoclaw-vllm", - dockerRunFlags: [ - "--gpus", - "all", - "--ipc=host", - "-v", - `${process.env.HOME}/.cache/huggingface:/root/.cache/huggingface`, - "-e", - "HF_HOME=/root/.cache/huggingface", - ], + dockerRunFlags: vllmDockerRunFlags(), pullTimeoutSec: 12 * 60 * 60, loadTimeoutSec: 1800, }; @@ -154,16 +175,8 @@ const STATION_PROFILE: VllmProfile = { ? "all" : indices.length === 1 ? `device=${indices[0]}` - : `'"device=${indices.join(",")}"'`; - return [ - "--gpus", - gpuFlag, - "--ipc=host", - "-v", - `${process.env.HOME}/.cache/huggingface:/root/.cache/huggingface`, - "-e", - "HF_HOME=/root/.cache/huggingface", - ]; + : `device=${indices.join(",")}`; + return vllmDockerRunFlags(gpuFlag); }, pullTimeoutSec: SPARK_PROFILE.pullTimeoutSec, loadTimeoutSec: SPARK_PROFILE.loadTimeoutSec, @@ -260,9 +273,9 @@ function downloadModel( "--entrypoint", "hf", "-v", - `${process.env.HOME}/.cache/huggingface:/root/.cache/huggingface`, + hfCacheMount(), "-e", - "HF_HOME=/root/.cache/huggingface", + `HF_HOME=${HF_CACHE_CONTAINER_DIR}`, ...buildHfTokenDockerArgs(), profile.image, "download", @@ -336,21 +349,51 @@ function downloadModel( }); } -// Build the `docker run` command for the long-lived vLLM inference container. +function validateDockerArg(value: string, label: string): string { + if (value.length === 0) { + throw new Error(`${label} must not be empty`); + } + if (value.includes("\0")) { + throw new Error(`${label} must not contain NUL bytes`); + } + return value; +} + +function validateDockerArgs(args: readonly string[], label: string): string[] { + return args.map((arg, index) => validateDockerArg(String(arg), `${label}[${String(index)}]`)); +} + +// Build the `docker run` argv for the long-lived vLLM inference container. // Exported for testing. `--restart unless-stopped` makes the container come // back after a host reboot or Docker daemon restart (#4886); without a restart // policy the container stays down after a reboot and `nemoclaw inference get` // fails until a full `nemoclaw onboard --fresh --gpu` recreates it. -export function buildVllmRunCommand( +export function buildVllmRunArgs( profile: VllmProfile, model: VllmModelDef, - runFlags: string, -): string { - const extra = runFlags ? ` ${runFlags}` : ""; - return ( - `docker run -d --restart unless-stopped${extra} -p ${String(VLLM_PORT)}:8000 ` + - `--name ${profile.containerName} --entrypoint /bin/bash ${profile.image} -lc ${JSON.stringify(buildVllmServeCommand(model))}` + runFlags: readonly string[], + env: NodeJS.ProcessEnv = process.env, +): string[] { + const image = validateDockerArg(profile.image, "vLLM image"); + const containerName = validateDockerArg(profile.containerName, "vLLM container name"); + const safeRunFlags = validateDockerArgs( + [...runFlags, ...buildHfTokenDockerArgs(env)], + "vLLM docker run flags", ); + return [ + "--restart", + "unless-stopped", + ...safeRunFlags, + "-p", + `${String(VLLM_PORT)}:8000`, + "--name", + containerName, + "--entrypoint", + "/bin/bash", + image, + "-lc", + buildVllmServeCommand(model, env), + ]; } function startContainer( @@ -359,7 +402,7 @@ function startContainer( ): { ok: boolean; reason?: string } { emit(`Starting vLLM container (${profile.containerName})`); // Idempotent: tear down any prior container by the same name first. - runShell(`docker rm -f ${profile.containerName}`, { + dockerForceRm(profile.containerName, { ignoreError: true, suppressOutput: true, }); @@ -372,10 +415,13 @@ function startContainer( // token from the docker subprocess env by default, so we have to put it // back via the `env:` option; the docker argv only carries `-e KEY` so // the value stays out of /proc//cmdline. - const hfTokenFlags = buildHfTokenDockerArgs().join(" "); - const flags = [resolvedFlags.join(" "), hfTokenFlags].filter(Boolean).join(" "); - const cmd = buildVllmRunCommand(profile, model, flags); - const result = runShell(cmd, { + let runArgs: string[]; + try { + runArgs = buildVllmRunArgs(profile, model, resolvedFlags); + } catch (err) { + return { ok: false, reason: (err as Error).message }; + } + const result = dockerRunDetached(runArgs, { ignoreError: true, suppressOutput: true, env: buildHfTokenForwardEnv(), @@ -572,7 +618,7 @@ export async function installVllm( const ready = await waitForVllmReady(profile); if (!ready.ok) { printContainerLogTail(profile); - runShell(`docker stop ${profile.containerName}`, { + dockerStop(profile.containerName, { ignoreError: true, suppressOutput: true, });