diff --git a/src/lib/inference/model-acquisition/hugging-face.test.ts b/src/lib/inference/model-acquisition/hugging-face.test.ts new file mode 100644 index 00000000000..838b50b208b --- /dev/null +++ b/src/lib/inference/model-acquisition/hugging-face.test.ts @@ -0,0 +1,373 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { EventEmitter } from "node:events"; + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { + acquireHuggingFaceModel, + buildHfTokenDockerArgs, + buildHfTokenForwardEnv, + buildHuggingFaceModelDownloadArgv, + type HuggingFaceModelAcquisitionObserver, + type HuggingFaceModelAcquisitionRequest, + hfDownloadAuthentication, +} from "./hugging-face"; + +const dockerSpawn = vi.fn(); + +interface MockProcess extends EventEmitter { + readonly stderr: EventEmitter; + readonly stdout: EventEmitter; +} + +function mockProcess(): MockProcess { + const proc = new EventEmitter() as MockProcess; + Object.defineProperties(proc, { + stderr: { value: new EventEmitter() }, + stdout: { value: new EventEmitter() }, + }); + return proc; +} + +function request( + overrides: Partial = {}, +): HuggingFaceModelAcquisitionRequest { + return { + credentialEnv: { HF_TOKEN: "hf_test_token" }, + dockerEnv: { DOCKER_HOST: "ssh://spark.example.test" }, + downloaderImage: `nvcr.io/nvidia/vllm@sha256:${"a".repeat(64)}`, + hostCacheDir: "/home/nvidia/.cache/huggingface", + repository: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-GGUF", + revision: "0123456789abcdef0123456789abcdef01234567", + spawnDocker: dockerSpawn, + userIdentity: "1001:1001", + ...overrides, + }; +} + +function observer(): HuggingFaceModelAcquisitionObserver { + return { logLine: vi.fn(), onRateLimit: vi.fn() }; +} + +describe("Hugging Face model acquisition", () => { + let stdoutWrite: ReturnType; + let stderrWrite: ReturnType; + + beforeEach(() => { + vi.clearAllMocks(); + stdoutWrite = vi.spyOn(process.stdout, "write").mockImplementation(() => true); + stderrWrite = vi.spyOn(process.stderr, "write").mockImplementation(() => true); + }); + + afterEach(() => { + stdoutWrite.mockRestore(); + stderrWrite.mockRestore(); + vi.useRealTimers(); + }); + + it("builds the existing cache, identity, token, and revision contract with one exact file (#8279)", () => { + const input = request({ filename: "model/Nemotron.Q4_K_M.gguf" }); + + expect(buildHuggingFaceModelDownloadArgv(input)).toEqual([ + "run", + "-t", + "--rm", + "--pull=never", + "--user", + "1001:1001", + "--entrypoint", + "hf", + "-v", + "/home/nvidia/.cache/huggingface:/tmp/nemoclaw-huggingface", + "-e", + "HF_HOME=/tmp/nemoclaw-huggingface", + "-e", + "HF_TOKEN", + input.downloaderImage, + "download", + input.repository, + "model/Nemotron.Q4_K_M.gguf", + "--revision", + input.revision, + ]); + }); + + it("preserves whole-repository anonymous downloads without revision or host identity (#8279)", () => { + const input = request({ + credentialEnv: {}, + revision: undefined, + userIdentity: null, + }); + + expect(buildHuggingFaceModelDownloadArgv(input)).toEqual([ + "run", + "-t", + "--rm", + "--pull=never", + "--entrypoint", + "hf", + "-v", + "/home/nvidia/.cache/huggingface:/tmp/nemoclaw-huggingface", + "-e", + "HF_HOME=/tmp/nemoclaw-huggingface", + input.downloaderImage, + "download", + input.repository, + ]); + }); + + it.each([ + "", + "../model.gguf", + "/model.gguf", + "--revision", + "model/../other.gguf", + ])("rejects an exact filename that is not a normalized repository-relative path %j (#8279)", (filename) => { + expect(() => buildHuggingFaceModelDownloadArgv(request({ filename }))).toThrow( + "Hugging Face filename must be one normalized repository-relative path", + ); + }); + + it.each([ + "", + "--revision", + " nvidia/model", + "nvidia/model ", + "nvidia/bad model", + "nvidia/model/extra", + "nvidia/bad..model", + "nvidia/model.git", + `nvidia/${"a".repeat(90)}`, + "nvidia/\u001b[31mmodel", + ])("rejects an invalid Hugging Face repository ID %j (#8279)", (repository) => { + expect(() => buildHuggingFaceModelDownloadArgv(request({ repository }))).toThrow( + "Hugging Face repository must be one valid repository ID", + ); + }); + + it("rejects a cache path that is not normalized before starting Docker (#8279)", () => { + expect(() => + buildHuggingFaceModelDownloadArgv(request({ hostCacheDir: "/cache/../foreign" })), + ).toThrow("Hugging Face cache must be a normalized absolute path"); + }); + + it("rejects a root download identity before starting Docker (#8279)", () => { + expect(() => buildHuggingFaceModelDownloadArgv(request({ userIdentity: "0:0" }))).toThrow( + "Hugging Face model download user must be one non-root numeric uid and numeric gid", + ); + }); + + it("forwards the token by key outside argv and reports completion (#8279)", async () => { + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const input = request(); + const events = observer(); + const resultPromise = acquireHuggingFaceModel(input, events); + proc.emit("exit", 0); + + await expect(resultPromise).resolves.toEqual({ ok: true }); + const [argv, options] = dockerSpawn.mock.calls[0] as [ + string[], + { env: Record; stdio: string[] }, + ]; + expect(argv).toContain("HF_TOKEN"); + expect(argv.join("\n")).not.toContain("hf_test_token"); + expect(options).toEqual({ + env: { + DOCKER_HOST: "ssh://spark.example.test", + HF_TOKEN: "hf_test_token", + }, + stdio: ["ignore", "pipe", "pipe"], + }); + expect(events.logLine).toHaveBeenNthCalledWith( + 1, + `Pre-downloading model with hf: ${input.repository}`, + ); + expect(events.logLine).toHaveBeenNthCalledWith(2, "Model download complete"); + }); + + it("redacts a token split across UTF-8 output chunks and detects rate limiting (#8279)", async () => { + const token = `hf_${"r".repeat(32)}`; + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const events = observer(); + const resultPromise = acquireHuggingFaceModel( + request({ credentialEnv: { HF_TOKEN: token } }), + events, + ); + const unicode = Buffer.from("Downloading café\n"); + const unicodeSplit = unicode.indexOf(0xc3) + 1; + const tokenSplit = 17; + proc.stdout.emit("data", unicode.subarray(0, unicodeSplit)); + proc.stdout.emit("data", unicode.subarray(unicodeSplit)); + proc.stdout.emit("data", Buffer.from(`value=${token.slice(0, tokenSplit)}`)); + proc.stderr.emit( + "data", + Buffer.from(`${token.slice(tokenSplit)} HTTP 429 Too Many Requests\n`), + ); + proc.stdout.emit("data", Buffer.from("\n")); + proc.emit("exit", 1); + + await expect(resultPromise).resolves.toEqual({ + ok: false, + reason: "hf download failed (exit 1)", + }); + const stdout = stdoutWrite.mock.calls.map((call: unknown[]) => String(call[0])).join(""); + const stderr = stderrWrite.mock.calls.map((call: unknown[]) => String(call[0])).join(""); + expect(`${stdout}\n${stderr}`).not.toContain(token); + expect(`${stdout}\n${stderr}`).not.toContain(token.slice(0, tokenSplit)); + expect(`${stdout}\n${stderr}`).not.toContain(token.slice(tokenSplit)); + expect(stdout).toContain("Downloading café"); + expect(stdout).toContain("value="); + expect(stderr).toContain("HTTP 429 Too Many Requests"); + expect(`${stdout}\n${stderr}`).not.toContain("�"); + expect(events.onRateLimit).toHaveBeenCalledOnce(); + }); + + it("redacts a contextual bearer credential split across same-stream chunks (#8279)", async () => { + const secret = "opaque-bearer-value"; + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const resultPromise = acquireHuggingFaceModel(request({ credentialEnv: {} }), observer()); + proc.stdout.emit("data", Buffer.from("Authorization: Bearer ")); + proc.stdout.emit("data", Buffer.from(`${secret} request failed\n`)); + proc.emit("exit", 1); + + await expect(resultPromise).resolves.toEqual({ + ok: false, + reason: "hf download failed (exit 1)", + }); + const output = stdoutWrite.mock.calls.map((call: unknown[]) => String(call[0])).join(""); + expect(output).toContain("Authorization: Bearer request failed"); + expect(output).not.toContain(secret); + }); + + it("redacts a contextual bearer credential across interleaved streams (#8279)", async () => { + const secret = "opaque-bearer-value"; + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const resultPromise = acquireHuggingFaceModel(request({ credentialEnv: {} }), observer()); + proc.stdout.emit("data", Buffer.from("Authorization: Bearer ")); + proc.stderr.emit("data", Buffer.from("progress\n")); + proc.stdout.emit("data", Buffer.from(`${secret} request failed\n`)); + proc.emit("exit", 1); + + await expect(resultPromise).resolves.toEqual({ + ok: false, + reason: "hf download failed (exit 1)", + }); + const output = [...stdoutWrite.mock.calls, ...stderrWrite.mock.calls] + .map((call: unknown[]) => String(call[0])) + .join(""); + expect(output).toContain("Authorization: Bearer request failed"); + expect(output).toContain("progress"); + expect(output).not.toContain(secret); + }); + + it("redacts a folded authorization value across a split CRLF boundary (#8279)", async () => { + const secret = "opaque-folded-value"; + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const resultPromise = acquireHuggingFaceModel(request({ credentialEnv: {} }), observer()); + proc.stdout.emit("data", Buffer.from("error: Authorization:\r")); + proc.stdout.emit("data", Buffer.from("\n")); + proc.stdout.emit("data", Buffer.from(`\t${secret}\r`)); + const progressFrameOutput = stdoutWrite.mock.calls + .map((call: unknown[]) => String(call[0])) + .join(""); + expect(progressFrameOutput).not.toContain(secret); + proc.stdout.emit("data", Buffer.from("\nnext diagnostic\n")); + proc.emit("exit", 1); + + await expect(resultPromise).resolves.toEqual({ + ok: false, + reason: "hf download failed (exit 1)", + }); + const output = [...stdoutWrite.mock.calls, ...stderrWrite.mock.calls] + .map((call: unknown[]) => String(call[0])) + .join(""); + expect(output).toContain("Authorization: "); + expect(output).toContain("next diagnostic"); + expect(output).not.toContain(secret); + }); + + it("flushes carriage-return progress frames promptly (#8279)", async () => { + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const resultPromise = acquireHuggingFaceModel(request({ credentialEnv: {} }), observer()); + proc.stdout.emit("data", Buffer.from("Downloading 1%\r")); + const progressOutput = stdoutWrite.mock.calls + .map((call: unknown[]) => String(call[0])) + .join(""); + expect(progressOutput).toContain("Downloading 1%\r"); + proc.emit("exit", 0); + + await expect(resultPromise).resolves.toEqual({ ok: true }); + }); + + it("bounds unterminated output by suppressing the affected stream (#8279)", async () => { + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const events = observer(); + const resultPromise = acquireHuggingFaceModel(request({ credentialEnv: {} }), events); + proc.stdout.emit("data", Buffer.from("x".repeat(70_000))); + proc.stdout.emit("data", Buffer.from("must-not-be-emitted\n")); + proc.emit("exit", 0); + + await expect(resultPromise).resolves.toEqual({ ok: true }); + const output = stdoutWrite.mock.calls.map((call: unknown[]) => String(call[0])).join(""); + expect(output).not.toContain("must-not-be-emitted"); + expect(output).not.toContain("x".repeat(100)); + expect(events.logLine).toHaveBeenCalledWith( + "Hugging Face output suppressed after exceeding the redaction buffer limit", + ); + }); + + it("returns validation failures without spawning Docker (#8279)", async () => { + await expect( + acquireHuggingFaceModel(request({ hostCacheDir: "/cache/../foreign" }), observer()), + ).resolves.toEqual({ + ok: false, + reason: "Hugging Face cache must be a normalized absolute path without ':'", + }); + expect(dockerSpawn).not.toHaveBeenCalled(); + }); + + it("returns one spawn error and clears the download heartbeat (#8279)", async () => { + vi.useFakeTimers(); + const proc = mockProcess(); + dockerSpawn.mockReturnValue(proc); + const events = observer(); + const resultPromise = acquireHuggingFaceModel(request(), events); + proc.emit("error", new Error("docker unavailable")); + proc.emit("exit", 125); + + await expect(resultPromise).resolves.toEqual({ + ok: false, + reason: "spawn error: docker unavailable", + }); + await vi.advanceTimersByTimeAsync(60_000); + expect(events.logLine).toHaveBeenCalledTimes(1); + expect(events.onRateLimit).not.toHaveBeenCalled(); + expect(stderrWrite.mock.calls.flat().join("\n")).not.toContain("hf output lines"); + }); + + it("keeps token selection and metadata behavior independent of the serving provider (#8279)", () => { + const env = { + HF_TOKEN: "hf_primary", + HUGGING_FACE_HUB_TOKEN: "hf_secondary", + } as NodeJS.ProcessEnv; + expect(buildHfTokenDockerArgs(env)).toEqual(["-e", "HF_TOKEN"]); + expect(buildHfTokenForwardEnv(env)).toEqual({ HF_TOKEN: "hf_primary" }); + expect(hfDownloadAuthentication(env)).toEqual({ + authenticated: true, + source: "HF_TOKEN", + }); + expect( + hfDownloadAuthentication({ HF_TOKEN: " ", HUGGING_FACE_HUB_TOKEN: "hf_fallback" }), + ).toEqual({ authenticated: true, source: "HUGGING_FACE_HUB_TOKEN" }); + expect(hfDownloadAuthentication({})).toEqual({ authenticated: false }); + }); +}); diff --git a/src/lib/inference/model-acquisition/hugging-face.ts b/src/lib/inference/model-acquisition/hugging-face.ts new file mode 100644 index 00000000000..639034e1e56 --- /dev/null +++ b/src/lib/inference/model-acquisition/hugging-face.ts @@ -0,0 +1,387 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import type { spawn } from "node:child_process"; +import path from "node:path"; +import { StringDecoder } from "node:string_decoder"; + +import { redactFull } from "../../security/redact"; + +const HF_TOKEN_ENV_KEYS = ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"] as const; +const HF_RATE_LIMIT_PATTERN = /\b429\b|too many requests|rate[\s_-]*limit/i; +const MODEL_DOWNLOAD_HEARTBEAT_MS = 30_000; +const HF_DOWNLOAD_CACHE_CONTAINER_DIR = "/tmp/nemoclaw-huggingface"; +const HF_REPOSITORY_ID_MAX_LENGTH = 96; +const HF_PENDING_OUTPUT_MAX_CHARS = 64 * 1024; +const HF_REPOSITORY_COMPONENT_PATTERN = /^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$/; +const HF_SENSITIVE_HEADER_AT_STREAM_BOUNDARY = + /\b(?:authorization|proxy-authorization|cookie|set-cookie)[ \t]*[:=][^\r\n]*(?:\r\n|[\r\n])$/i; +const TAIL_MAX = 50; + +export interface HuggingFaceModelAcquisitionRequest { + readonly credentialEnv?: NodeJS.ProcessEnv; + readonly dockerEnv: Readonly>; + readonly downloaderImage: string; + /** When set, download exactly one repository-relative file. */ + readonly filename?: string; + readonly hostCacheDir: string; + readonly repository: string; + readonly revision?: string; + readonly spawnDocker: ( + args: readonly string[], + options?: Parameters[2], + ) => ReturnType; + readonly userIdentity: string | null; +} + +export interface HuggingFaceModelAcquisitionObserver { + readonly logLine: (line: string) => void; + readonly onRateLimit: () => void; +} + +export type HuggingFaceModelAcquisitionResult = + | { readonly ok: true } + | { readonly ok: false; readonly reason: string }; + +export type HfDownloadAuthentication = + | { readonly authenticated: false } + | { readonly authenticated: true; readonly source: (typeof HF_TOKEN_ENV_KEYS)[number] }; + +function pickHfTokenEntry( + env: NodeJS.ProcessEnv = process.env, +): { key: (typeof HF_TOKEN_ENV_KEYS)[number]; value: string } | null { + for (const key of HF_TOKEN_ENV_KEYS) { + const value = String(env[key] ?? "").trim(); + if (value) return { key, value }; + } + return null; +} + +/** Return only the presence and source of Hugging Face authentication, never its value. */ +export function hfDownloadAuthentication( + env: NodeJS.ProcessEnv = process.env, +): HfDownloadAuthentication { + const entry = pickHfTokenEntry(env); + return entry ? { authenticated: true, source: entry.key } : { authenticated: false }; +} + +/** + * Return the key-only Docker environment argument so a token never enters the + * host process list. + */ +export function buildHfTokenDockerArgs(env: NodeJS.ProcessEnv = process.env): string[] { + const entry = pickHfTokenEntry(env); + return entry ? ["-e", entry.key] : []; +} + +/** Supply the matching token value to Docker's environment outside argv. */ +export function buildHfTokenForwardEnv( + env: NodeJS.ProcessEnv = process.env, +): Record { + const entry = pickHfTokenEntry(env); + return entry ? { [entry.key]: entry.value } : {}; +} + +function hfDownloadCacheMount(hostCacheDir: string): string { + const normalized = path.posix.normalize(hostCacheDir); + if ( + !path.posix.isAbsolute(hostCacheDir) || + normalized !== hostCacheDir || + hostCacheDir.includes(":") + ) { + throw new Error("Hugging Face cache must be a normalized absolute path without ':'"); + } + return `${normalized}:${HF_DOWNLOAD_CACHE_CONTAINER_DIR}`; +} + +function hostUserDockerArgs(identity: string | null): string[] { + if (identity !== null && !/^[1-9][0-9]*:(?:0|[1-9][0-9]*)$/.test(identity)) { + throw new Error( + "Hugging Face model download user must be one non-root numeric uid and numeric gid", + ); + } + return identity ? ["--user", identity] : []; +} + +function exactFilenameArgs(filename: string | undefined): string[] { + if (filename === undefined) return []; + const normalized = path.posix.normalize(filename); + if ( + filename.length === 0 || + filename.includes("\0") || + filename.startsWith("-") || + path.posix.isAbsolute(filename) || + normalized !== filename || + normalized === "." || + normalized.startsWith("../") + ) { + throw new Error("Hugging Face filename must be one normalized repository-relative path"); + } + return [filename]; +} + +function repositoryArg(repository: string): string { + const components = repository.split("/"); + if ( + repository.length === 0 || + repository.length > HF_REPOSITORY_ID_MAX_LENGTH || + components.length > 2 || + components.some((component) => !HF_REPOSITORY_COMPONENT_PATTERN.test(component)) || + repository.includes("--") || + repository.includes("..") || + repository.endsWith(".git") + ) { + throw new Error("Hugging Face repository must be one valid repository ID"); + } + return repository; +} + +export function buildHuggingFaceModelDownloadArgv( + request: HuggingFaceModelAcquisitionRequest, +): string[] { + const credentialEnv = request.credentialEnv ?? process.env; + return [ + "run", + "-t", + "--rm", + "--pull=never", + ...hostUserDockerArgs(request.userIdentity), + "--entrypoint", + "hf", + "-v", + hfDownloadCacheMount(request.hostCacheDir), + "-e", + `HF_HOME=${HF_DOWNLOAD_CACHE_CONTAINER_DIR}`, + ...buildHfTokenDockerArgs(credentialEnv), + request.downloaderImage, + "download", + repositoryArg(request.repository), + ...exactFilenameArgs(request.filename), + ...(request.revision ? ["--revision", request.revision] : []), + ]; +} + +function formatElapsed(ms: number): string { + const totalSeconds = Math.max(0, Math.floor(ms / 1000)); + const minutes = Math.floor(totalSeconds / 60); + const seconds = totalSeconds % 60; + if (minutes === 0) return `${String(seconds)}s`; + return `${String(minutes)}m ${String(seconds)}s`; +} + +function redactHfDownloadOutput(text: string, tokenValue: string | null): string { + const withoutKnownToken = tokenValue ? text.split(tokenValue).join("") : text; + return redactFull(withoutKnownToken); +} + +function redactHfDownloadOutputChunks( + chunks: readonly { text: string; stream: NodeJS.WriteStream }[], + tokenValue: string | null, +): string[] { + const joined = chunks.map((chunk) => chunk.text).join(""); + const tokenSpans: { start: number; end: number }[] = []; + if (tokenValue) { + let searchFrom = 0; + while (searchFrom < joined.length) { + const start = joined.indexOf(tokenValue, searchFrom); + if (start < 0) break; + tokenSpans.push({ start, end: start + tokenValue.length }); + searchFrom = start + tokenValue.length; + } + } + + let chunkStart = 0; + const tokenRedactedChunks = chunks.map((chunk) => { + const chunkEnd = chunkStart + chunk.text.length; + let cursor = chunkStart; + let safeText = ""; + for (const span of tokenSpans) { + if (span.end <= chunkStart || span.start >= chunkEnd) continue; + safeText += joined.slice(cursor, Math.max(cursor, span.start)); + if (span.start >= chunkStart) safeText += ""; + cursor = Math.max(cursor, Math.min(chunkEnd, span.end)); + } + safeText += joined.slice(cursor, chunkEnd); + chunkStart = chunkEnd; + return { text: safeText, stream: chunk.stream }; + }); + + const redactedChunks = tokenRedactedChunks.map(() => ""); + const streamGroups = new Map(); + for (const [index, chunk] of tokenRedactedChunks.entries()) { + const group = streamGroups.get(chunk.stream); + streamGroups.set(chunk.stream, { + text: `${group?.text ?? ""}${chunk.text}`, + lastIndex: index, + }); + } + for (const group of streamGroups.values()) { + redactedChunks[group.lastIndex] = redactHfDownloadOutput(group.text, null); + } + return redactedChunks; +} + +/** + * Run `hf download` in a caller-supplied image that Docker already has. + * Reuse the configured Hugging Face cache. + */ +export function acquireHuggingFaceModel( + request: HuggingFaceModelAcquisitionRequest, + observer: HuggingFaceModelAcquisitionObserver, +): Promise { + return new Promise((resolve) => { + const credentialEnv = request.credentialEnv ?? process.env; + const tokenValue = pickHfTokenEntry(credentialEnv)?.value ?? null; + let argv: string[]; + try { + argv = buildHuggingFaceModelDownloadArgv(request); + } catch (err) { + resolve({ + ok: false, + reason: err instanceof Error ? err.message : String(err), + }); + return; + } + observer.logLine(`Pre-downloading model with hf: ${request.repository}`); + + let proc: ReturnType; + try { + proc = request.spawnDocker(argv, { + env: { ...request.dockerEnv, ...buildHfTokenForwardEnv(credentialEnv) }, + stdio: ["ignore", "pipe", "pipe"], + }); + } catch (err) { + resolve({ + ok: false, + reason: `spawn error: ${err instanceof Error ? err.message : String(err)}`, + }); + return; + } + + const tail: string[] = []; + const outputDecoders = [ + { decoder: new StringDecoder("utf8"), stream: process.stdout }, + { decoder: new StringDecoder("utf8"), stream: process.stderr }, + ]; + let pendingOutput: { text: string; stream: NodeJS.WriteStream }[] = []; + let pendingOutputChars = 0; + const suppressedOutputStreams = new Set(); + let reportedOutputSuppression = false; + let resolved = false; + let decodersFinalized = false; + const start = Date.now(); + let lastOutputAt = start; + let lastOutputEndedCleanly = true; + const heartbeat = setInterval(() => { + const now = Date.now(); + if (now - lastOutputAt >= MODEL_DOWNLOAD_HEARTBEAT_MS) { + if (!lastOutputEndedCleanly) process.stdout.write("\n"); + observer.logLine( + `Model download still running (${formatElapsed(now - start)} elapsed; no new output)`, + ); + lastOutputAt = now; + lastOutputEndedCleanly = true; + } + }, MODEL_DOWNLOAD_HEARTBEAT_MS); + heartbeat.unref?.(); + + function done(result: HuggingFaceModelAcquisitionResult): void { + if (resolved) return; + resolved = true; + clearInterval(heartbeat); + resolve(result); + } + + function rememberTail(text: string): void { + for (const segment of text.split(/[\r\n]+/)) { + if (!segment) continue; + tail.push(segment); + if (tail.length > TAIL_MAX) tail.shift(); + } + } + + function pendingStreamsEndCleanly(): boolean { + const streamText = new Map(); + for (const chunk of pendingOutput) { + streamText.set(chunk.stream, `${streamText.get(chunk.stream) ?? ""}${chunk.text}`); + } + return [...streamText.values()].every( + (text) => /[\r\n]$/.test(text) && !HF_SENSITIVE_HEADER_AT_STREAM_BOUNDARY.test(text), + ); + } + + function flushOutput(flushAll = false): void { + if (pendingOutput.length === 0 || (!flushAll && !pendingStreamsEndCleanly())) return; + const selected = pendingOutput; + pendingOutput = []; + pendingOutputChars = 0; + const safeChunks = redactHfDownloadOutputChunks(selected, tokenValue); + for (const [index, safeText] of safeChunks.entries()) { + if (!safeText) continue; + selected[index].stream.write(safeText); + lastOutputEndedCleanly = /[\r\n]$/.test(safeText); + rememberTail(safeText); + } + } + + function queueOutput(text: string, stream: NodeJS.WriteStream): void { + if (!text || suppressedOutputStreams.has(stream)) return; + pendingOutput.push({ text, stream }); + pendingOutputChars += text.length; + flushOutput(); + if (pendingOutputChars <= HF_PENDING_OUTPUT_MAX_CHARS) return; + + for (const chunk of pendingOutput) suppressedOutputStreams.add(chunk.stream); + pendingOutput = []; + pendingOutputChars = 0; + if (!reportedOutputSuppression) { + reportedOutputSuppression = true; + observer.logLine( + "Hugging Face output suppressed after exceeding the redaction buffer limit", + ); + } + } + + function finalizeOutputDecoders(): void { + if (decodersFinalized) return; + decodersFinalized = true; + for (const state of outputDecoders) { + const text = state.decoder.end(); + queueOutput(text, state.stream); + } + flushOutput(true); + } + + function onChunk(buf: Buffer, state: (typeof outputDecoders)[number]): void { + lastOutputAt = Date.now(); + const text = state.decoder.write(buf); + queueOutput(text, state.stream); + } + + proc.stdout?.on("data", (buf: Buffer) => onChunk(buf, outputDecoders[0])); + proc.stderr?.on("data", (buf: Buffer) => onChunk(buf, outputDecoders[1])); + + proc.on("error", (err: Error) => { + finalizeOutputDecoders(); + done({ ok: false, reason: `spawn error: ${err.message}` }); + }); + + proc.on("exit", (code: number | null) => { + if (resolved) return; + finalizeOutputDecoders(); + if (code === 0) { + if (!lastOutputEndedCleanly) process.stdout.write("\n"); + observer.logLine("Model download complete"); + done({ ok: true }); + return; + } + if (tail.length > 0) { + process.stderr.write(` --- Last ${String(tail.length)} hf output lines: ---\n`); + for (const line of tail) process.stderr.write(` ${line}\n`); + process.stderr.write(" ---\n"); + } + if (HF_RATE_LIMIT_PATTERN.test(tail.join("\n"))) observer.onRateLimit(); + done({ ok: false, reason: `hf download failed (exit ${String(code)})` }); + }); + }); +} diff --git a/src/lib/inference/vllm-download-model.test.ts b/src/lib/inference/vllm-download-model.test.ts new file mode 100644 index 00000000000..88bc5faada9 --- /dev/null +++ b/src/lib/inference/vllm-download-model.test.ts @@ -0,0 +1,59 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + acquireHuggingFaceModel: vi.fn(), + dockerSpawn: vi.fn(), +})); + +vi.mock("../adapters/docker", async (importOriginal) => ({ + ...(await importOriginal()), + dockerSpawn: mocks.dockerSpawn, +})); + +vi.mock("./model-acquisition/hugging-face", async (importOriginal) => ({ + ...(await importOriginal()), + acquireHuggingFaceModel: mocks.acquireHuggingFaceModel, +})); + +import { detectVllmProfile, downloadModel } from "./vllm"; + +describe("vLLM model acquisition adapter", () => { + it("maps the existing vLLM contract to shared acquisition and preserves failures (#8279)", async () => { + const failure = { + ok: false as const, + reason: "hf download failed (exit 1)", + }; + mocks.acquireHuggingFaceModel.mockResolvedValue(failure); + const profile = detectVllmProfile({ platform: "spark", type: "nvidia" })!; + const model = profile.defaultModel; + const dockerEnv = { DOCKER_HOST: "ssh://spark.example.test" }; + + await expect( + downloadModel(profile, model, dockerEnv, { + hostCacheDir: "/home/nvidia/.cache/huggingface", + userIdentity: "1001:1001", + }), + ).resolves.toBe(failure); + + expect(mocks.acquireHuggingFaceModel).toHaveBeenCalledOnce(); + expect(mocks.acquireHuggingFaceModel).toHaveBeenCalledWith( + { + dockerEnv, + downloaderImage: profile.image, + hostCacheDir: "/home/nvidia/.cache/huggingface", + repository: model.id, + revision: model.revision, + spawnDocker: mocks.dockerSpawn, + userIdentity: "1001:1001", + }, + { + logLine: expect.any(Function), + onRateLimit: expect.any(Function), + }, + ); + expect(mocks.dockerSpawn).not.toHaveBeenCalled(); + }); +}); diff --git a/src/lib/inference/vllm.test.ts b/src/lib/inference/vllm.test.ts index 42d7cc3b5ba..9b1b8c4ca84 100644 --- a/src/lib/inference/vllm.test.ts +++ b/src/lib/inference/vllm.test.ts @@ -54,11 +54,11 @@ vi.mock("./vllm-storage", async (importOriginal) => { }); import { currentPhaseActivityLabel } from "../core/phase-activity"; +import { hfDownloadAuthentication } from "./model-acquisition/hugging-face"; import { assertVllmRegistryDigestRef, buildVllmRunArgs, detectVllmProfile, - hfDownloadAuthentication, installVllm, isNemoClawManagedVllmRunning, NEMOCLAW_VLLM_CONTAINER_NAME, diff --git a/src/lib/inference/vllm.ts b/src/lib/inference/vllm.ts index a58cae26b5e..e43e0d69a52 100644 --- a/src/lib/inference/vllm.ts +++ b/src/lib/inference/vllm.ts @@ -8,7 +8,6 @@ import fs from "node:fs"; import os from "node:os"; import path from "node:path"; -import { StringDecoder } from "node:string_decoder"; import { isDeepStrictEqual } from "node:util"; import { dockerCapture, @@ -29,8 +28,11 @@ import { VLLM_PORT } from "../core/ports"; import { shellQuote } from "../core/shell-quote"; import { isAffirmativeAnswer } from "../onboard/prompt-helpers"; import { runCapture } from "../runner"; -import { redactFull } from "../security/redact"; import { isSafeModelId } from "../validation"; +import { + acquireHuggingFaceModel, + hfDownloadAuthentication, +} from "./model-acquisition/hugging-face"; import { getGpuIndicesByName } from "./nim"; import { buildLocalDualStationDockerEnv, @@ -182,14 +184,10 @@ function qwen35bNvfp4Model(): VllmModelDef { return match; } -const HF_TOKEN_ENV_KEYS = ["HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"] as const; const HF_TOKEN_SETTINGS_URL = "https://huggingface.co/settings/tokens"; -const HF_RATE_LIMIT_PATTERN = /\b429\b|too many requests|rate[\s_-]*limit/i; -const MODEL_DOWNLOAD_HEARTBEAT_MS = 30_000; const VLLM_LAUNCH_HEARTBEAT_MS = 30_000; const VLLM_MAX_STARTUP_RESTARTS = 3; const HF_CACHE_CONTAINER_DIR = "/root/.cache/huggingface"; -const HF_DOWNLOAD_CACHE_CONTAINER_DIR = "/tmp/nemoclaw-huggingface"; const HF_CACHE_COMPONENT_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]*$/; export const NEMOCLAW_VLLM_CONTAINER_NAME = "nemoclaw-vllm"; export const NEMOCLAW_VLLM_MANAGED_LABEL = "com.nvidia.nemoclaw.managed-vllm"; @@ -203,18 +201,6 @@ function hfCacheMount(): string { return `${hostHfCacheDir()}:${HF_CACHE_CONTAINER_DIR}`; } -function hfDownloadCacheMount(hostCacheDir = hostHfCacheDir()): string { - const normalized = path.posix.normalize(hostCacheDir); - if ( - !path.posix.isAbsolute(hostCacheDir) || - normalized !== hostCacheDir || - hostCacheDir.includes(":") - ) { - throw new Error("vLLM model cache must be a normalized absolute path without ':'"); - } - return `${normalized}:${HF_DOWNLOAD_CACHE_CONTAINER_DIR}`; -} - function hfModelCacheKey(model: VllmModelDef): string | null { const modelParts = model.id.split("/"); if (modelParts.some((part) => !HF_CACHE_COMPONENT_PATTERN.test(part))) return null; @@ -243,13 +229,6 @@ function hostUserIdentity(): string | null { return `${String(process.getuid())}:${String(process.getgid())}`; } -function hostUserDockerArgs(identity = hostUserIdentity()): string[] { - if (identity !== null && !/^[1-9][0-9]*:(?:0|[1-9][0-9]*)$/.test(identity)) { - throw new Error("vLLM model download user must be one non-root numeric uid and numeric gid"); - } - return identity ? ["--user", identity] : []; -} - function vllmDockerRunFlags(gpuFlag = "all"): string[] { return [ "--gpus", @@ -262,28 +241,6 @@ function vllmDockerRunFlags(gpuFlag = "all"): string[] { ]; } -function pickHfTokenEntry( - env: NodeJS.ProcessEnv = process.env, -): { key: (typeof HF_TOKEN_ENV_KEYS)[number]; value: string } | null { - for (const key of HF_TOKEN_ENV_KEYS) { - const value = String(env[key] ?? "").trim(); - if (value) return { key, value }; - } - return null; -} - -export type HfDownloadAuthentication = - | { authenticated: false } - | { authenticated: true; source: (typeof HF_TOKEN_ENV_KEYS)[number] }; - -/** Return only the presence and source of Hugging Face authentication, never its value. */ -export function hfDownloadAuthentication( - env: NodeJS.ProcessEnv = process.env, -): HfDownloadAuthentication { - const entry = pickHfTokenEntry(env); - return entry ? { authenticated: true, source: entry.key } : { authenticated: false }; -} - function printHfDownloadAuthentication(nonInteractive: boolean): void { const authentication = hfDownloadAuthentication(); if (authentication.authenticated) { @@ -311,44 +268,6 @@ function printHfDownloadAuthentication(nonInteractive: boolean): void { console.log(" The token is passed only to the temporary model downloader."); } -function redactHfDownloadOutput(text: string, tokenValue: string | null): string { - const withoutKnownToken = tokenValue ? text.split(tokenValue).join("") : text; - return redactFull(withoutKnownToken); -} - -function redactHfDownloadOutputChunks( - chunks: readonly { text: string; stream: NodeJS.WriteStream }[], - tokenValue: string | null, -): string[] { - const joined = chunks.map((chunk) => chunk.text).join(""); - const tokenSpans: { start: number; end: number }[] = []; - if (tokenValue) { - let searchFrom = 0; - while (searchFrom < joined.length) { - const start = joined.indexOf(tokenValue, searchFrom); - if (start < 0) break; - tokenSpans.push({ start, end: start + tokenValue.length }); - searchFrom = start + tokenValue.length; - } - } - - let chunkStart = 0; - return chunks.map((chunk) => { - const chunkEnd = chunkStart + chunk.text.length; - let cursor = chunkStart; - let safeText = ""; - for (const span of tokenSpans) { - if (span.end <= chunkStart || span.start >= chunkEnd) continue; - safeText += joined.slice(cursor, Math.max(cursor, span.start)); - if (span.start >= chunkStart) safeText += ""; - cursor = Math.max(cursor, Math.min(chunkEnd, span.end)); - } - safeText += joined.slice(cursor, chunkEnd); - chunkStart = chunkEnd; - return redactHfDownloadOutput(safeText, null); - }); -} - function printHfRateLimitRecovery(): void { process.stderr.write(" Hugging Face rate limiting was detected.\n"); process.stderr.write(` Create a read token at ${HF_TOKEN_SETTINGS_URL}.\n`); @@ -359,36 +278,6 @@ function printHfRateLimitRecovery(): void { ); } -/** - * Forward a Hugging Face token from the host into the one-shot `hf download` - * container so gated model weights can be fetched. - * - * Returns the bare `-e KEY` form (no `=value`) so the token never lands in - * the host process list. Docker reads the actual value from its own - * environment, which the caller is responsible for populating via - * `buildHfTokenForwardEnv` when spawning through the runner allowlist. - * The download container can live for several minutes during a cold pull; - * argv-embedded secrets would be visible via `ps` for that whole window. - */ -export function buildHfTokenDockerArgs(env: NodeJS.ProcessEnv = process.env): string[] { - const entry = pickHfTokenEntry(env); - return entry ? ["-e", entry.key] : []; -} - -/** - * 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 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, -): Record { - const entry = pickHfTokenEntry(env); - return entry ? { [entry.key]: entry.value } : {}; -} - const SPARK_PROFILE: VllmProfile = { name: "DGX Spark", platform: "spark", @@ -566,156 +455,25 @@ export async function pullImage( return { ok: true }; } -// Run `hf download ` inside a one-shot container of the same image. +// Preserve the vLLM downloadModel API while acquireHuggingFaceModel runs `hf download`. export function downloadModel( profile: VllmProfile, model: VllmModelDef, dockerEnv: Record = buildVllmDockerEnv(), target: { hostCacheDir?: string; userIdentity?: string } = {}, ): Promise<{ ok: boolean; reason?: string }> { - emit(`Pre-downloading model with hf: ${model.id}`); - return new Promise((resolve) => { - const tokenValue = pickHfTokenEntry()?.value ?? null; - const proc = dockerSpawn( - [ - "run", - "-t", - "--rm", - "--pull=never", - ...hostUserDockerArgs(target.userIdentity), - "--entrypoint", - "hf", - "-v", - hfDownloadCacheMount(target.hostCacheDir), - "-e", - `HF_HOME=${HF_DOWNLOAD_CACHE_CONTAINER_DIR}`, - ...buildHfTokenDockerArgs(), - profile.image, - "download", - model.id, - ...(model.revision ? ["--revision", model.revision] : []), - ], - { - env: { ...dockerEnv, ...buildHfTokenForwardEnv() }, - stdio: ["ignore", "pipe", "pipe"], - }, - ); - - const tail: string[] = []; - const outputDecoders = [ - { decoder: new StringDecoder("utf8"), stream: process.stdout }, - { decoder: new StringDecoder("utf8"), stream: process.stderr }, - ]; - let pendingOutput: { text: string; stream: NodeJS.WriteStream }[] = []; - const TAIL_MAX = 50; - let resolved = false; - let decodersFinalized = false; - const start = Date.now(); - let lastOutputAt = start; - let lastOutputEndedCleanly = true; - const heartbeat = setInterval(() => { - const now = Date.now(); - if (now - lastOutputAt >= MODEL_DOWNLOAD_HEARTBEAT_MS) { - if (!lastOutputEndedCleanly) process.stdout.write("\n"); - emit(`Model download still running (${formatElapsed(now - start)} elapsed; no new output)`); - lastOutputAt = now; - lastOutputEndedCleanly = true; - } - }, MODEL_DOWNLOAD_HEARTBEAT_MS); - heartbeat.unref?.(); - - function done(result: { ok: boolean; reason?: string }): void { - if (resolved) return; - resolved = true; - clearInterval(heartbeat); - resolve(result); - } - - function rememberTail(text: string): void { - for (const segment of text.split(/[\r\n]+/)) { - if (!segment) continue; - tail.push(segment); - if (tail.length > TAIL_MAX) tail.shift(); - } - } - - function takePendingOutput(end: number): { text: string; stream: NodeJS.WriteStream }[] { - const selected: { text: string; stream: NodeJS.WriteStream }[] = []; - let remaining = end; - while (remaining > 0 && pendingOutput.length > 0) { - const chunk = pendingOutput[0]; - if (chunk.text.length <= remaining) { - selected.push(chunk); - pendingOutput.shift(); - remaining -= chunk.text.length; - continue; - } - selected.push({ text: chunk.text.slice(0, remaining), stream: chunk.stream }); - pendingOutput[0] = { text: chunk.text.slice(remaining), stream: chunk.stream }; - remaining = 0; - } - return selected; - } - - function flushOutput(flushAll = false): void { - const pendingText = pendingOutput.map((chunk) => chunk.text).join(""); - const end = flushAll - ? pendingText.length - : Math.max(pendingText.lastIndexOf("\n"), pendingText.lastIndexOf("\r")) + 1; - if (end <= 0) return; - const selected = takePendingOutput(end); - const safeChunks = redactHfDownloadOutputChunks(selected, tokenValue); - for (const [index, safeText] of safeChunks.entries()) { - if (!safeText) continue; - selected[index].stream.write(safeText); - lastOutputEndedCleanly = /[\r\n]$/.test(safeText); - rememberTail(safeText); - } - } - - function finalizeOutputDecoders(): void { - if (decodersFinalized) return; - decodersFinalized = true; - for (const state of outputDecoders) { - const text = state.decoder.end(); - if (text) pendingOutput.push({ text, stream: state.stream }); - } - flushOutput(true); - } - - function onChunk(buf: Buffer, state: (typeof outputDecoders)[number]): void { - lastOutputAt = Date.now(); - const text = state.decoder.write(buf); - if (text) pendingOutput.push({ text, stream: state.stream }); - flushOutput(); - } - - proc.stdout?.on("data", (buf: Buffer) => onChunk(buf, outputDecoders[0])); - proc.stderr?.on("data", (buf: Buffer) => onChunk(buf, outputDecoders[1])); - - proc.on("error", (err: Error) => { - finalizeOutputDecoders(); - done({ ok: false, reason: `spawn error: ${err.message}` }); - }); - - proc.on("exit", (code: number | null) => { - finalizeOutputDecoders(); - if (code === 0) { - if (!lastOutputEndedCleanly) process.stdout.write("\n"); - emit("Model download complete"); - done({ ok: true }); - return; - } - // Surface the last few sanitized lines so a failure has actionable context. - if (tail.length > 0) { - process.stderr.write(` --- Last ${String(tail.length)} hf output lines: ---\n`); - for (const line of tail) process.stderr.write(` ${line}\n`); - process.stderr.write(" ---\n"); - } - if (HF_RATE_LIMIT_PATTERN.test(tail.join("\n"))) printHfRateLimitRecovery(); - done({ ok: false, reason: `hf download failed (exit ${String(code)})` }); - }); - }); + return acquireHuggingFaceModel( + { + dockerEnv, + downloaderImage: profile.image, + hostCacheDir: target.hostCacheDir ?? hostHfCacheDir(), + repository: model.id, + revision: model.revision, + spawnDocker: dockerSpawn, + userIdentity: target.userIdentity ?? hostUserIdentity(), + }, + { logLine: emit, onRateLimit: printHfRateLimitRecovery }, + ); } function validateDockerArg(value: string, label: string): string { diff --git a/test/detect-vllm-profile.test.ts b/test/detect-vllm-profile.test.ts index 70d45b06805..3336e38dfe2 100644 --- a/test/detect-vllm-profile.test.ts +++ b/test/detect-vllm-profile.test.ts @@ -6,8 +6,8 @@ import { describe, expect, it, vi } from "vitest"; import { buildHfTokenDockerArgs, buildHfTokenForwardEnv, - detectVllmProfile, -} from "../src/lib/inference/vllm.js"; +} from "../src/lib/inference/model-acquisition/hugging-face.js"; +import { detectVllmProfile } from "../src/lib/inference/vllm.js"; describe("detectVllmProfile", () => { it("returns the Spark profile when gpu.platform === 'spark'", () => {