diff --git a/src/lib/adapters/docker/pull.test.ts b/src/lib/adapters/docker/pull.test.ts new file mode 100644 index 00000000000..c9f1f167840 --- /dev/null +++ b/src/lib/adapters/docker/pull.test.ts @@ -0,0 +1,220 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { EventEmitter } from "node:events"; + +import { afterEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../../runner", () => ({ + ROOT: "/repo/root", + run: vi.fn(), + runCapture: vi.fn(), +})); + +import { + type DockerPullChildProcess, + type DockerPullReadable, + dockerPullProgressSignature, + dockerPullWithProgressWatchdog, +} from "./pull"; + +class FakeReadable extends EventEmitter implements DockerPullReadable {} + +class FakeChild extends EventEmitter implements DockerPullChildProcess { + stdout = new FakeReadable(); + stderr = new FakeReadable(); + kill = vi.fn((signal?: NodeJS.Signals | number) => { + queueMicrotask(() => this.emit("close", null, signal ?? "SIGTERM")); + return true; + }); +} + +describe("docker pull progress watchdog", () => { + afterEach(() => { + vi.useRealTimers(); + }); + + it("parses Docker pull progress signatures from phase and byte-count lines", () => { + expect(dockerPullProgressSignature("latest: Pulling from nvidia/vllm")).toBe( + "source:latest: Pulling from nvidia/vllm", + ); + expect(dockerPullProgressSignature("abc123def: Downloading [==> ] 12.5MB/1.2GB")).toBe( + "layer:abc123def:Downloading:12.5MB/1.2GB", + ); + expect(dockerPullProgressSignature("abc123def: Extracting 250 MB/1 GB")).toBe( + "layer:abc123def:Extracting:250MB/1GB", + ); + expect(dockerPullProgressSignature("abc123def: Pull complete")).toBe( + "layer:abc123def:Pull complete", + ); + expect( + dockerPullProgressSignature( + "e20d54a357dc: Extracting [=====> ] 10.03MB/84.19MB", + ), + ).toBe("layer:e20d54a357dc:Extracting:10.03MB/84.19MB"); + expect( + dockerPullProgressSignature( + "e5eda78c7490: Downloading [============================> ] 17.83MB/30.76MB", + ), + ).toBe("layer:e5eda78c7490:Downloading:17.83MB/30.76MB"); + expect(dockerPullProgressSignature("b39b21d4717d: Download complete ")).toBe( + "layer:b39b21d4717d:Download complete", + ); + expect( + dockerPullProgressSignature("#12 sha256:abc123 25.4MB/100MB 4.0s 10.1MB/s"), + ).toBeNull(); + }); + + it("spawns plain docker pull without unsupported progress flags", async () => { + const child = new FakeChild(); + const spawnImpl = vi.fn(() => child); + const pull = dockerPullWithProgressWatchdog("example/image:latest", { + suppressOutput: true, + spawnImpl, + }); + + child.emit("close", 0, null); + + await expect(pull).resolves.toMatchObject({ status: 0 }); + expect(spawnImpl).toHaveBeenCalledWith( + ["pull", "example/image:latest"], + expect.objectContaining({ stdio: ["ignore", "pipe", "pipe"] }), + ); + }); + + it("allows slow pulls to continue while byte progress advances", async () => { + vi.useFakeTimers(); + const child = new FakeChild(); + const pull = dockerPullWithProgressWatchdog("example/image:latest", { + suppressOutput: true, + stallTimeoutMs: 1_000, + maxTimeoutMs: 10_000, + watchdogIntervalMs: 100, + spawnImpl: () => child, + }); + + child.stderr.emit("data", Buffer.from("abc123def: Downloading 1MB/10MB\r")); + await vi.advanceTimersByTimeAsync(900); + expect(child.kill).not.toHaveBeenCalled(); + + child.stderr.emit("data", Buffer.from("abc123def: Downloading 2MB/10MB\r")); + await vi.advanceTimersByTimeAsync(900); + expect(child.kill).not.toHaveBeenCalled(); + + child.emit("close", 0, null); + await expect(pull).resolves.toMatchObject({ + status: 0, + timedOut: false, + timeoutKind: null, + }); + }); + + it("kills a pull when output repeats without forward progress", async () => { + vi.useFakeTimers(); + const child = new FakeChild(); + const pull = dockerPullWithProgressWatchdog("example/image:latest", { + suppressOutput: true, + stallTimeoutMs: 1_000, + maxTimeoutMs: 10_000, + watchdogIntervalMs: 100, + spawnImpl: () => child, + }); + + child.stderr.emit("data", Buffer.from("abc123def: Downloading 1MB/10MB\r")); + await vi.advanceTimersByTimeAsync(500); + child.stderr.emit("data", Buffer.from("abc123def: Downloading 1MB/10MB\r")); + await vi.advanceTimersByTimeAsync(600); + + await expect(pull).resolves.toMatchObject({ + status: 124, + timedOut: true, + timeoutKind: "stall", + }); + expect(child.kill).toHaveBeenCalledWith("SIGTERM"); + }); + + it("keeps returned diagnostics to a bounded output tail", async () => { + const child = new FakeChild(); + const emitted: string[] = []; + const pull = dockerPullWithProgressWatchdog("example/image:latest", { + logLine: (line) => emitted.push(line), + spawnImpl: () => child, + }); + + for (let i = 0; i < 210; i += 1) { + child.stderr.emit("data", Buffer.from(`diagnostic line ${String(i)}\n`)); + } + child.emit("close", 1, null); + + const result = await pull; + expect(emitted).toHaveLength(210); + expect(result.output.split("\n")).toHaveLength(200); + expect(result.output).not.toContain("diagnostic line 0"); + expect(result.output).toContain("diagnostic line 209"); + }); + + it("clamps positive sub-millisecond watchdog intervals to 1ms", async () => { + vi.useFakeTimers(); + const child = new FakeChild(); + const pull = dockerPullWithProgressWatchdog("example/image:latest", { + suppressOutput: true, + stallTimeoutMs: 0.5, + maxTimeoutMs: 10_000, + watchdogIntervalMs: 0.5, + spawnImpl: () => child, + }); + + await vi.advanceTimersByTimeAsync(2); + + await expect(pull).resolves.toMatchObject({ + status: 124, + timedOut: true, + timeoutKind: "stall", + }); + expect(child.kill).toHaveBeenCalledWith("SIGTERM"); + }); + + it("enforces the maximum safety budget even when progress keeps advancing", async () => { + vi.useFakeTimers(); + const child = new FakeChild(); + const pull = dockerPullWithProgressWatchdog("example/image:latest", { + suppressOutput: true, + stallTimeoutMs: 1_000, + maxTimeoutMs: 2_000, + watchdogIntervalMs: 100, + spawnImpl: () => child, + }); + + child.stderr.emit("data", Buffer.from("abc123def: Downloading 1MB/10MB\r")); + await vi.advanceTimersByTimeAsync(900); + child.stderr.emit("data", Buffer.from("abc123def: Downloading 2MB/10MB\r")); + await vi.advanceTimersByTimeAsync(900); + child.stderr.emit("data", Buffer.from("abc123def: Downloading 3MB/10MB\r")); + await vi.advanceTimersByTimeAsync(300); + + await expect(pull).resolves.toMatchObject({ + status: 124, + timedOut: true, + timeoutKind: "max", + }); + expect(child.kill).toHaveBeenCalledWith("SIGTERM"); + }); + + it("returns a failed result when docker pull fails to spawn", async () => { + const child = new FakeChild(); + const pull = dockerPullWithProgressWatchdog("example/image:latest", { + suppressOutput: true, + spawnImpl: () => child, + }); + const error = Object.assign(new Error("ENOENT"), { code: "ENOENT" }); + + child.emit("error", error); + + await expect(pull).resolves.toMatchObject({ + status: 1, + timedOut: false, + timeoutKind: null, + error, + }); + }); +}); diff --git a/src/lib/adapters/docker/pull.ts b/src/lib/adapters/docker/pull.ts index 5c11c896001..ed156ea2ebd 100644 --- a/src/lib/adapters/docker/pull.ts +++ b/src/lib/adapters/docker/pull.ts @@ -1,8 +1,268 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +import type { SpawnOptions } from "node:child_process"; + +import { ROOT } from "../../runner"; +import { buildSubprocessEnv } from "../../subprocess-env"; +import { dockerSpawn } from "./exec"; import { dockerRun, type DockerRunOptions, type DockerRunResult } from "./run"; export function dockerPull(imageRef: string, opts: DockerRunOptions = {}): DockerRunResult { return dockerRun(["pull", imageRef], opts); } + +export const DEFAULT_DOCKER_PULL_STALL_TIMEOUT_MS = 120 * 1000; +export const DEFAULT_DOCKER_PULL_MAX_TIMEOUT_MS = 12 * 60 * 60 * 1000; +const DOCKER_PULL_OUTPUT_TAIL_LINES = 200; +const DOCKER_PULL_PROGRESS_STATE_LIMIT = 512; + +export interface DockerPullWatchdogOptions { + suppressOutput?: boolean; + stallTimeoutMs?: number; + maxTimeoutMs?: number; + watchdogIntervalMs?: number; + logLine?: (line: string) => void; + env?: NodeJS.ProcessEnv; + spawnImpl?: ( + args: readonly string[], + options: SpawnOptions, + ) => DockerPullChildProcess; +} + +export interface DockerPullWatchdogResult { + status: number; + signal: NodeJS.Signals | null; + output: string; + timedOut: boolean; + timeoutKind: "stall" | "max" | null; + error?: Error; +} + +export interface DockerPullReadable { + on(event: "data", listener: (chunk: Buffer | string) => void): this; +} + +export interface DockerPullChildProcess { + stdout: DockerPullReadable | null; + stderr: DockerPullReadable | null; + kill?(signal?: NodeJS.Signals | number): boolean; + on(event: "error", listener: (error: Error) => void): this; + on(event: "close", listener: (code: number | null, signal: NodeJS.Signals | null) => void): this; +} + +export function dockerPullProgressSignature(line: string): string | null { + const normalized = line.trim().replace(/\s+/g, " "); + if (!normalized) return null; + + const layerMatch = normalized.match( + /^([a-f0-9]{6,}):\s+([A-Za-z][A-Za-z ]*?)(?=\s+\[|\s+[\d.]+\s*[A-Za-z]+\s*\/|$)(?:\s+\[[^\]]*\])?(?:\s+([\d.]+\s*[A-Za-z]+)\s*\/\s*([\d.]+\s*[A-Za-z]+))?/, + ); + if (layerMatch) { + const layer = layerMatch[1]; + const phase = layerMatch[2].trim(); + const completed = layerMatch[3]?.replace(/\s+/g, ""); + const total = layerMatch[4]?.replace(/\s+/g, ""); + return completed && total + ? `layer:${layer}:${phase}:${completed}/${total}` + : `layer:${layer}:${phase}`; + } + + if (/^(?:[^:\s]+:\s+)?Pulling from \S+/.test(normalized)) { + return `source:${normalized}`; + } + if (/^Digest: sha256:[a-f0-9]{8,}/.test(normalized)) { + return `digest:${normalized}`; + } + if (/^Status: /.test(normalized)) { + return `status:${normalized}`; + } + + // This vLLM watchdog intentionally recognizes observed `docker pull` + // layer/status output only; BuildKit-style build output is not pull progress. + return null; +} + +function dockerPullProgressKey(signature: string): string { + const layerMatch = signature.match(/^layer:([^:]+):([^:]+)(?::.*)?$/); + if (layerMatch) return `layer:${layerMatch[1]}:${layerMatch[2]}`; + if (signature.startsWith("source:")) return "source"; + if (signature.startsWith("digest:")) return "digest"; + if (signature.startsWith("status:")) return "status"; + return signature; +} + +export async function dockerPullWithProgressWatchdog( + imageRef: string, + opts: DockerPullWatchdogOptions = {}, +): Promise { + const stallTimeoutMs = positiveMs( + opts.stallTimeoutMs, + DEFAULT_DOCKER_PULL_STALL_TIMEOUT_MS, + ); + const maxTimeoutMs = positiveMs(opts.maxTimeoutMs, DEFAULT_DOCKER_PULL_MAX_TIMEOUT_MS); + const watchdogIntervalMs = positiveMs( + opts.watchdogIntervalMs, + Math.min(1000, Math.max(100, Math.floor(stallTimeoutMs / 4))), + ); + const logLine = opts.logLine ?? ((line: string) => process.stderr.write(`${line}\n`)); + const spawnImpl = + opts.spawnImpl ?? + ((args: readonly string[], options: SpawnOptions) => + dockerSpawn(args, options) as DockerPullChildProcess); + const child = spawnImpl(["pull", imageRef], { + cwd: ROOT, + env: buildSubprocessEnv(normalizeExtraEnv(opts.env)), + stdio: ["ignore", "pipe", "pipe"], + }); + + return new Promise((resolve) => { + const lines: string[] = []; + const latestProgressByKey = new Map(); + const startedAt = Date.now(); + let lastProgressAt = startedAt; + let stdoutPending = ""; + let stderrPending = ""; + let settled = false; + let timeoutKind: "stall" | "max" | null = null; + let capturedError: Error | undefined; + let forceKillTimer: NodeJS.Timeout | null = null; + + function noteProgress(line: string) { + const signature = dockerPullProgressSignature(line); + if (!signature) return; + const key = dockerPullProgressKey(signature); + if (latestProgressByKey.get(key) === signature) return; + latestProgressByKey.delete(key); + latestProgressByKey.set(key, signature); + while (latestProgressByKey.size > DOCKER_PULL_PROGRESS_STATE_LIMIT) { + const oldestKey = latestProgressByKey.keys().next().value; + if (oldestKey === undefined) break; + latestProgressByKey.delete(oldestKey); + } + lastProgressAt = Date.now(); + } + + function rememberLine(line: string) { + lines.push(line); + if (lines.length > DOCKER_PULL_OUTPUT_TAIL_LINES) { + lines.splice(0, lines.length - DOCKER_PULL_OUTPUT_TAIL_LINES); + } + } + + function flushLine(line: string) { + const trimmed = line.trimEnd(); + if (!trimmed) return; + rememberLine(trimmed); + noteProgress(trimmed); + if (!opts.suppressOutput) logLine(trimmed); + } + + function consumeChunk(pending: string, chunk: Buffer | string): string { + const text = pending + chunk.toString(); + const parts = text.split(/[\r\n]+/); + const nextPending = parts.pop() ?? ""; + for (const part of parts) flushLine(part); + return nextPending; + } + + function flushPending() { + if (stdoutPending) { + flushLine(stdoutPending); + stdoutPending = ""; + } + if (stderrPending) { + flushLine(stderrPending); + stderrPending = ""; + } + } + + function requestKill(kind: "stall" | "max") { + if (settled || timeoutKind) return; + timeoutKind = kind; + const detail = + kind === "stall" + ? `docker pull stalled after ${formatSeconds(stallTimeoutMs)} without progress` + : `docker pull exceeded maximum safety budget ${formatSeconds(maxTimeoutMs)}`; + rememberLine(detail); + if (!opts.suppressOutput) logLine(` ${detail}`); + try { + child.kill?.("SIGTERM"); + } catch { + /* best effort */ + } + forceKillTimer = setTimeout(() => { + try { + child.kill?.("SIGKILL"); + } catch { + /* best effort */ + } + }, 5000); + forceKillTimer.unref?.(); + } + + const watchdog = setInterval(() => { + const now = Date.now(); + if (now - lastProgressAt >= stallTimeoutMs) { + requestKill("stall"); + return; + } + if (now - startedAt >= maxTimeoutMs) { + requestKill("max"); + } + }, watchdogIntervalMs); + watchdog.unref?.(); + + function finish(code: number | null, signal: NodeJS.Signals | null) { + if (settled) return; + settled = true; + flushPending(); + clearInterval(watchdog); + if (forceKillTimer) clearTimeout(forceKillTimer); + const status = code ?? (timeoutKind ? 124 : 1); + const result: DockerPullWatchdogResult = { + status, + signal, + output: lines.join("\n"), + timedOut: timeoutKind !== null, + timeoutKind, + ...(capturedError ? { error: capturedError } : {}), + }; + resolve(result); + } + + child.stdout?.on("data", (chunk) => { + stdoutPending = consumeChunk(stdoutPending, chunk); + }); + child.stderr?.on("data", (chunk) => { + stderrPending = consumeChunk(stderrPending, chunk); + }); + child.on("error", (error: Error) => { + capturedError = error; + rememberLine(`docker pull failed to start: ${error.message}`); + finish(1, null); + }); + child.on("close", (code: number | null, signal: NodeJS.Signals | null) => { + finish(code, signal); + }); + }); +} + +function normalizeExtraEnv(env: NodeJS.ProcessEnv | undefined): Record | undefined { + if (!env) return undefined; + const normalized: Record = {}; + for (const [key, value] of Object.entries(env)) { + if (value !== undefined) normalized[key] = value; + } + return normalized; +} + +function positiveMs(value: number | undefined, fallbackMs: number): number { + if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) return fallbackMs; + return value < 1 ? 1 : Math.floor(value); +} + +function formatSeconds(ms: number): string { + const seconds = Math.max(1, Math.round(ms / 1000)); + return `${seconds}s`; +} diff --git a/src/lib/inference/vllm.ts b/src/lib/inference/vllm.ts index 8f1401a4c0e..ed6f496ad81 100644 --- a/src/lib/inference/vllm.ts +++ b/src/lib/inference/vllm.ts @@ -5,10 +5,14 @@ // offer vLLM at all" lives in onboard.ts; this module owns picking the // right profile per platform and running the install. -const { runCapture, runShell } = require("../runner"); -const { dockerCapture, dockerSpawn } = require("../adapters/docker"); -const { VLLM_PORT } = require("../core/ports"); -const { getGpuIndicesByName } = require("./nim"); +import { + dockerCapture, + dockerPullWithProgressWatchdog, + dockerSpawn, +} from "../adapters/docker"; +import { VLLM_PORT } from "../core/ports"; +import { runCapture, runShell } from "../runner"; +import { getGpuIndicesByName } from "./nim"; import { DEFAULT_VLLM_MODEL, VLLM_MODELS, @@ -21,7 +25,7 @@ import { // Per-platform install recipe. Add new platforms by appending an entry to // the profile table at the bottom of this file. The menu key in onboard.ts // stays "install-vllm" regardless of platform. -interface VllmProfile { +export interface VllmProfile { name: string; // human label, e.g. "DGX Spark" image: string; // container image // Default model when NEMOCLAW_VLLM_MODEL is unset. Per-platform default @@ -39,7 +43,9 @@ interface VllmProfile { buildDockerRunFlags?: () => string[]; // Approximate first-run time shown in the confirmation prompt. estimatedMinutes: string; - // Image-pull deadline. First run on a slow link can be several minutes. + // Maximum wall-clock safety budget for image pulls. The Docker adapter uses + // a shorter progress watchdog for stalls, so slow-but-moving pulls can keep + // going until this last-ditch cap. pullTimeoutSec: number; // Marker emitter sees container output line by line. The patterns below // map a line to a user-visible "==> ..." progress marker. Order matters: @@ -118,7 +124,7 @@ const SPARK_PROFILE: VllmProfile = { "HF_HOME=/root/.cache/huggingface", ], estimatedMinutes: "10–30 minutes", - pullTimeoutSec: 900, + pullTimeoutSec: 12 * 60 * 60, loadTimeoutSec: 1800, progressMarkers: [ { @@ -219,19 +225,22 @@ function dockerPrereqsOk(): { ok: boolean; reason?: string } { return { ok: true }; } -function pullImage(profile: VllmProfile): { ok: boolean; reason?: string } { +export async function pullImage(profile: VllmProfile): Promise<{ ok: boolean; reason?: string }> { emit(`Pulling vLLM image: ${profile.image}`); - // GNU `timeout` enforces the pull deadline. macOS BSD coreutils omits it; - // fall back to a plain `docker pull` there. - const hasTimeout = !!runCapture(["sh", "-c", "command -v timeout"], { - ignoreError: true, - }).trim(); - const prefix = hasTimeout ? `timeout ${String(profile.pullTimeoutSec)} ` : ""; - const result = runShell(`${prefix}docker pull ${profile.image}`, { - ignoreError: true, - suppressOutput: true, + const result = await dockerPullWithProgressWatchdog(profile.image, { + maxTimeoutMs: profile.pullTimeoutSec * 1000, + logLine: emit, }); if (result.status !== 0) { + if (result.timeoutKind === "stall") { + return { ok: false, reason: "docker pull stalled with no progress" }; + } + if (result.timeoutKind === "max") { + return { + ok: false, + reason: `docker pull exceeded ${String(profile.pullTimeoutSec)}s safety budget`, + }; + } return { ok: false, reason: `docker pull failed (exit ${String(result.status)})` }; } return { ok: true }; @@ -291,8 +300,8 @@ function downloadModel( } } - proc.stdout.on("data", onChunk); - proc.stderr.on("data", onChunk); + proc.stdout?.on("data", onChunk); + proc.stderr?.on("data", onChunk); proc.on("error", (err: Error) => { resolve({ ok: false, reason: `spawn error: ${err.message}` }); @@ -409,10 +418,10 @@ function streamLogsUntilReady( let stdoutBuffer = ""; let stderrBuffer = ""; - proc.stdout.on("data", (raw: Buffer) => { + proc.stdout?.on("data", (raw: Buffer) => { stdoutBuffer = consumeChunk(stdoutBuffer, raw); }); - proc.stderr.on("data", (raw: Buffer) => { + proc.stderr?.on("data", (raw: Buffer) => { stderrBuffer = consumeChunk(stderrBuffer, raw); }); @@ -498,7 +507,7 @@ export async function installVllm( return { ok: false }; } - const pull = pullImage(profile); + const pull = await pullImage(profile); if (!pull.ok) { console.error(` vLLM install failed: ${String(pull.reason)}`); return { ok: false };