roboomp
941a931599
Extended ONNX Runtime CUDA sidecar repair to PI_TINY_DEVICE=auto and gpu on Linux x64, matching the CUDA-capable diagnostics path. Added focused regression coverage for both generic accelerated settings. Refs #4475
495 lines
19 KiB
TypeScript
495 lines
19 KiB
TypeScript
import * as fsp from "node:fs/promises";
|
|
import { createRequire } from "node:module";
|
|
import * as path from "node:path";
|
|
import type { ProgressInfo } from "@huggingface/transformers";
|
|
import {
|
|
ensureRuntimeInstalled,
|
|
getTinyModelsCacheDir,
|
|
installRuntimeModuleResolver,
|
|
isCompiledBinary,
|
|
resolveRuntimeModule,
|
|
} from "@oh-my-pi/pi-utils";
|
|
import packageJson from "../../package.json" with { type: "json" };
|
|
|
|
/**
|
|
* Child-side scaffolding shared by the ONNX inference worker bodies
|
|
* (`stt/asr-worker`, `tiny/worker`, `tts/tts-worker`). These are the helpers
|
|
* that run inside the spawned subprocess: error serialization, structured log
|
|
* and progress reporting over the worker's typed transport, side-runtime
|
|
* install (sharp stubbing + module-resolver patch), once-per-process runtime
|
|
* memoization, and the Transformers.js runtime loader. The parent/client-side
|
|
* complement lives in `worker-client.ts`.
|
|
*
|
|
* Each worker keeps its own strongly-typed transport / model-key / progress
|
|
* event; the structural {@link WorkerLogTransport} / {@link WorkerProgressTransport}
|
|
* interfaces below are the minimal shapes these helpers need, and every worker's
|
|
* concrete transport satisfies them.
|
|
*/
|
|
|
|
export const TRANSFORMERS_PACKAGE = "@huggingface/transformers";
|
|
const COMPILED_TRANSFORMERS_VERSION = process.env.PI_TINY_TRANSFORMERS_VERSION;
|
|
const ONNX_RUNTIME_NODE_PACKAGE = "onnxruntime-node";
|
|
const ONNX_RUNTIME_CUDA_INSTALL = "cuda12";
|
|
const ONNX_RUNTIME_CUDA_PROVIDER_FILES = [
|
|
"libonnxruntime_providers_cuda.so",
|
|
"libonnxruntime_providers_shared.so",
|
|
"libonnxruntime_providers_tensorrt.so",
|
|
] as const;
|
|
const LINUX_X64_ONNX_RUNTIME_CUDA_PROVIDER_DIR = path.join("bin", "napi-v6", "linux", "x64");
|
|
|
|
const sourceRequire = createRequire(import.meta.url);
|
|
|
|
// ── Error serialization ─────────────────────────────────────────────
|
|
|
|
export function errorText(error: unknown): string {
|
|
return error instanceof Error ? (error.stack ?? error.message) : String(error);
|
|
}
|
|
|
|
export function errorMessage(error: unknown): string {
|
|
return error instanceof Error ? error.message : String(error);
|
|
}
|
|
|
|
// ── Structured logging ──────────────────────────────────────────────
|
|
|
|
export type WorkerLogLevel = "debug" | "warn" | "error";
|
|
|
|
/** Minimal transport surface a worker exposes for forwarding log lines. */
|
|
export interface WorkerLogTransport {
|
|
send(message: { type: "log"; level: WorkerLogLevel; msg: string; meta?: Record<string, unknown> }): void;
|
|
}
|
|
|
|
export function sendLog(
|
|
transport: WorkerLogTransport,
|
|
level: WorkerLogLevel,
|
|
msg: string,
|
|
meta?: Record<string, unknown>,
|
|
): void {
|
|
transport.send({ type: "log", level, msg, meta });
|
|
}
|
|
|
|
// ── Progress reporting ──────────────────────────────────────────────
|
|
|
|
/**
|
|
* Generic worker progress event. Each worker's protocol declares an identical
|
|
* shape with its own `modelKey` type; this is the parameterized version the
|
|
* shared helpers emit, structurally assignable to each protocol's event.
|
|
*/
|
|
export interface WorkerProgressEvent<K> {
|
|
modelKey: K;
|
|
status: "initiate" | "download" | "progress" | "progress_total" | "done" | "ready" | "error";
|
|
name?: string;
|
|
file?: string;
|
|
progress?: number;
|
|
loaded?: number;
|
|
total?: number;
|
|
files?: Record<string, { loaded: number; total: number }>;
|
|
task?: string;
|
|
model?: string;
|
|
}
|
|
|
|
/** Minimal transport surface a worker exposes for emitting progress events. */
|
|
export interface WorkerProgressTransport<K> {
|
|
send(message: { type: "progress"; id: string; event: WorkerProgressEvent<K> }): void;
|
|
}
|
|
|
|
/** Map a Transformers.js {@link ProgressInfo} onto the worker progress event. */
|
|
function toProgressEvent<K>(modelKey: K, info: ProgressInfo): WorkerProgressEvent<K> {
|
|
if (info.status === "ready") {
|
|
return { modelKey, status: info.status, task: info.task, model: info.model };
|
|
}
|
|
if (info.status === "progress_total") {
|
|
return {
|
|
modelKey,
|
|
status: info.status,
|
|
name: info.name,
|
|
progress: info.progress,
|
|
loaded: info.loaded,
|
|
total: info.total,
|
|
files: info.files,
|
|
};
|
|
}
|
|
if (info.status === "progress") {
|
|
return {
|
|
modelKey,
|
|
status: info.status,
|
|
name: info.name,
|
|
file: info.file,
|
|
progress: info.progress,
|
|
loaded: info.loaded,
|
|
total: info.total,
|
|
};
|
|
}
|
|
return { modelKey, status: info.status, name: info.name, file: info.file };
|
|
}
|
|
|
|
export function sendProgress<K>(
|
|
transport: WorkerProgressTransport<K>,
|
|
id: string,
|
|
modelKey: K,
|
|
info: ProgressInfo,
|
|
): void {
|
|
transport.send({ type: "progress", id, event: toProgressEvent(modelKey, info) });
|
|
}
|
|
|
|
// ── Model cache ─────────────────────────────────────────────────────
|
|
|
|
/**
|
|
* If a model is already warming/warm in `cache`, replay a `ready` progress
|
|
* event for this request once it resolves and return the cached promise so the
|
|
* caller can short-circuit; otherwise return `undefined`.
|
|
*/
|
|
export function replayCachedReady<K, M>(
|
|
cache: Map<K, Promise<M>>,
|
|
modelKey: K,
|
|
transport: WorkerProgressTransport<K>,
|
|
requestId: string,
|
|
task: string,
|
|
model: string,
|
|
): Promise<M> | undefined {
|
|
const cached = cache.get(modelKey);
|
|
if (!cached) return undefined;
|
|
void cached
|
|
.then(() => {
|
|
transport.send({ type: "progress", id: requestId, event: { modelKey, status: "ready", task, model } });
|
|
})
|
|
.catch(() => undefined);
|
|
return cached;
|
|
}
|
|
|
|
// ── Side-runtime install scaffolding ────────────────────────────────
|
|
|
|
/**
|
|
* Stub `sharp` (the speech/text pipelines are not image codecs, so the native
|
|
* image dependency is dead weight) and patch the module resolver so a side
|
|
* runtime's bare requires resolve against its own `node_modules`. Returns the
|
|
* runtime's `node_modules` directory.
|
|
*/
|
|
export async function installSharpStubResolver(runtimeDir: string): Promise<string> {
|
|
const nodeModules = path.join(runtimeDir, "node_modules");
|
|
const sharpStub = path.join(runtimeDir, "omp-sharp-stub.cjs");
|
|
await Bun.write(sharpStub, "module.exports = {};\n");
|
|
installRuntimeModuleResolver({ runtimeNodeModules: nodeModules, stubs: { sharp: sharpStub } });
|
|
return nodeModules;
|
|
}
|
|
|
|
function shouldInstallOnnxRuntimeCudaProviders(device: string | undefined): boolean {
|
|
const normalized = device?.trim().toLowerCase();
|
|
return (
|
|
process.platform === "linux" &&
|
|
process.arch === "x64" &&
|
|
(normalized === "cuda" || normalized === "gpu" || normalized === "auto")
|
|
);
|
|
}
|
|
|
|
async function missingOnnxRuntimeCudaProviderFiles(binDir: string): Promise<string[]> {
|
|
const missing: string[] = [];
|
|
for (const file of ONNX_RUNTIME_CUDA_PROVIDER_FILES) {
|
|
try {
|
|
await fsp.access(path.join(binDir, file));
|
|
} catch {
|
|
missing.push(file);
|
|
}
|
|
}
|
|
return missing;
|
|
}
|
|
|
|
async function readPipe(stream: ReadableStream<Uint8Array> | null): Promise<string> {
|
|
if (!stream) return "";
|
|
return new Response(stream).text();
|
|
}
|
|
|
|
async function installOnnxRuntimeCudaProviders(packageDir: string, runtimeDir: string, binDir: string): Promise<void> {
|
|
const script = path.join(packageDir, "script", "install.js");
|
|
try {
|
|
await fsp.access(script);
|
|
} catch {
|
|
throw new Error(
|
|
`ONNX Runtime CUDA provider binaries are missing from ${binDir}, and ${script} is unavailable. Remove the tiny-model side runtime cache at ${runtimeDir} and retry.`,
|
|
);
|
|
}
|
|
|
|
const proc = Bun.spawn([process.execPath, script], {
|
|
cwd: runtimeDir,
|
|
env: { ...Bun.env, BUN_BE_BUN: "1", ONNXRUNTIME_NODE_INSTALL: ONNX_RUNTIME_CUDA_INSTALL },
|
|
stdout: "pipe",
|
|
stderr: "pipe",
|
|
});
|
|
const [stdout, stderr, exitCode] = await Promise.all([
|
|
readPipe(proc.stdout as ReadableStream<Uint8Array> | null),
|
|
readPipe(proc.stderr as ReadableStream<Uint8Array> | null),
|
|
proc.exited,
|
|
]);
|
|
if (exitCode !== 0) {
|
|
const output = `${stdout}\n${stderr}`.trim();
|
|
throw new Error(
|
|
`Failed to install ONNX Runtime CUDA provider binaries into ${binDir} with ${process.execPath} ${script} (exit ${exitCode}). Remove the tiny-model side runtime cache at ${runtimeDir} and retry with network access. ${output}`,
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Repairs the compiled Transformers side runtime when CUDA was requested and
|
|
* Bun skipped `onnxruntime-node`'s NuGet sidecar install.
|
|
*/
|
|
export async function ensureOnnxRuntimeCudaProviders(
|
|
runtimeDir: string,
|
|
device = process.env.PI_TINY_DEVICE,
|
|
): Promise<void> {
|
|
if (!shouldInstallOnnxRuntimeCudaProviders(device)) return;
|
|
const nodeModules = path.join(runtimeDir, "node_modules");
|
|
const manifest = resolveRuntimeModule(nodeModules, `${ONNX_RUNTIME_NODE_PACKAGE}/package.json`);
|
|
if (!manifest)
|
|
throw new Error(`Unable to resolve ${ONNX_RUNTIME_NODE_PACKAGE} in compiled runtime at ${nodeModules}`);
|
|
const packageDir = path.dirname(manifest);
|
|
const binDir = path.join(packageDir, LINUX_X64_ONNX_RUNTIME_CUDA_PROVIDER_DIR);
|
|
const missing = await missingOnnxRuntimeCudaProviderFiles(binDir);
|
|
if (missing.length === 0) return;
|
|
|
|
await installOnnxRuntimeCudaProviders(packageDir, runtimeDir, binDir);
|
|
const stillMissing = await missingOnnxRuntimeCudaProviderFiles(binDir);
|
|
if (stillMissing.length === 0) return;
|
|
throw new Error(
|
|
`ONNX Runtime CUDA provider install completed but ${stillMissing.join(", ")} are still missing from ${binDir}. Remove the tiny-model side runtime cache at ${runtimeDir} and retry.`,
|
|
);
|
|
}
|
|
|
|
/**
|
|
* Prepare a freshly-installed compiled runtime for loading and return the
|
|
* absolute entrypoint of `packageName` to `require`.
|
|
*/
|
|
async function prepareCompiledRuntime(runtimeDir: string, packageName: string): Promise<string> {
|
|
const nodeModules = await installSharpStubResolver(runtimeDir);
|
|
const entry = resolveRuntimeModule(nodeModules, packageName);
|
|
if (!entry) throw new Error(`Unable to resolve ${packageName} in compiled runtime at ${nodeModules}`);
|
|
return entry;
|
|
}
|
|
|
|
// ── Transformers version resolution ─────────────────────────────────
|
|
|
|
function resolveTransformersVersionSpec(): string {
|
|
const manifest = packageJson as {
|
|
optionalDependencies?: Record<string, string>;
|
|
dependencies?: Record<string, string>;
|
|
};
|
|
const versionSpec =
|
|
manifest.optionalDependencies?.[TRANSFORMERS_PACKAGE] ?? manifest.dependencies?.[TRANSFORMERS_PACKAGE];
|
|
if (!versionSpec) throw new Error(`${TRANSFORMERS_PACKAGE} is missing from package.json optionalDependencies`);
|
|
if (!versionSpec.startsWith("catalog:")) return versionSpec;
|
|
if (COMPILED_TRANSFORMERS_VERSION) return COMPILED_TRANSFORMERS_VERSION;
|
|
const installed = sourceRequire(`${TRANSFORMERS_PACKAGE}/package.json`) as { version: string };
|
|
return installed.version;
|
|
}
|
|
|
|
let cachedTransformersVersionSpec: string | undefined;
|
|
|
|
/**
|
|
* Lazily resolve (and memoize) the transformers version spec. In the `catalog:`
|
|
* case {@link resolveTransformersVersionSpec} `require`s the installed
|
|
* `@huggingface/transformers/package.json`, so it is only ever touched on the
|
|
* compiled-binary runtime-install path — loading a worker (smoke-test ping,
|
|
* online path) never triggers the transformers resolve/install dance.
|
|
*/
|
|
export function getTransformersVersionSpec(): string {
|
|
cachedTransformersVersionSpec ??= resolveTransformersVersionSpec();
|
|
return cachedTransformersVersionSpec;
|
|
}
|
|
|
|
// ── Transformers runtime loader ─────────────────────────────────────
|
|
|
|
/** The subset of the Transformers.js module surface {@link configureTransformers} touches. */
|
|
interface ConfigurableTransformers {
|
|
env: { cacheDir?: string; allowLocalModels?: boolean; logLevel?: unknown };
|
|
LogLevel: { ERROR: unknown };
|
|
}
|
|
|
|
export interface TransformersRuntimeMetadata {
|
|
__ompRuntimeNodeModules?: string;
|
|
__ompTransformersEntry?: string;
|
|
__ompCudaRepairError?: string;
|
|
}
|
|
|
|
function attachTransformersRuntimeMetadata<T extends ConfigurableTransformers>(
|
|
transformers: T,
|
|
metadata: TransformersRuntimeMetadata,
|
|
): T {
|
|
const runtime = transformers as T & TransformersRuntimeMetadata;
|
|
runtime.__ompRuntimeNodeModules = metadata.__ompRuntimeNodeModules;
|
|
runtime.__ompTransformersEntry = metadata.__ompTransformersEntry;
|
|
runtime.__ompCudaRepairError = metadata.__ompCudaRepairError;
|
|
return runtime;
|
|
}
|
|
|
|
const TRANSITIVE_CUDA_LIBRARY_RE =
|
|
/\b(lib(?:cu|nv)[A-Za-z0-9_.+-]*\.so(?:\.[0-9]+)*)\b[^:\n]*:\s*cannot open shared object file/iu;
|
|
const CUDA_DEVICE_UNAVAILABLE_RE = /\bCUDA failure 100\b|no CUDA-capable device is detected|cudaSetDevice|GPU=-1/iu;
|
|
|
|
function cudaDeviceUnavailable(error: unknown): boolean {
|
|
return CUDA_DEVICE_UNAVAILABLE_RE.test(errorText(error));
|
|
}
|
|
|
|
function missingCudaLibrary(error: unknown): string | undefined {
|
|
return TRANSITIVE_CUDA_LIBRARY_RE.exec(errorText(error))?.[1];
|
|
}
|
|
|
|
function cudaFailureCause(
|
|
metadata: TransformersRuntimeMetadata,
|
|
error: unknown,
|
|
missingFiles: readonly string[],
|
|
): string {
|
|
if (metadata.__ompCudaRepairError) {
|
|
return `ONNX Runtime CUDA provider install failed: ${metadata.__ompCudaRepairError}`;
|
|
}
|
|
if (missingFiles.length > 0) return `missing ONNX Runtime CUDA provider file(s): ${missingFiles.join(", ")}`;
|
|
const missingLibrary = missingCudaLibrary(error);
|
|
if (missingLibrary) return `${missingLibrary}: cannot open shared object file`;
|
|
if (cudaDeviceUnavailable(error)) {
|
|
return "CUDA provider files are present; CUDA runtime reports no CUDA-capable device";
|
|
}
|
|
return "CUDA provider files are present; inspect the original ONNX Runtime CUDA error";
|
|
}
|
|
|
|
function cudaFailureHint(
|
|
metadata: TransformersRuntimeMetadata,
|
|
error: unknown,
|
|
missingFiles: readonly string[],
|
|
): string {
|
|
if (metadata.__ompCudaRepairError) {
|
|
return "restore network access to nuget.org (or pre-populate the tiny side runtime) and rerun; CPU inference remained available";
|
|
}
|
|
if (missingFiles.length > 0) return "reinstall the tiny side runtime with ONNX Runtime postinstall enabled";
|
|
if (missingCudaLibrary(error)) {
|
|
return "install the matching CUDA/cuDNN shared libraries and expose them on the dynamic loader path";
|
|
}
|
|
if (cudaDeviceUnavailable(error)) {
|
|
return "make the NVIDIA GPU visible to this process/session, or use providers.tinyModelDevice=default/cpu";
|
|
}
|
|
return "check the host CUDA driver, device visibility, and ONNX Runtime CUDA compatibility";
|
|
}
|
|
|
|
function resolveOnnxRuntimePackageDir(metadata: TransformersRuntimeMetadata): string | null {
|
|
const entry = metadata.__ompTransformersEntry;
|
|
if (entry) {
|
|
try {
|
|
return path.dirname(createRequire(entry).resolve(`${ONNX_RUNTIME_NODE_PACKAGE}/package.json`));
|
|
} catch {
|
|
// Fall through to the side-runtime resolver below.
|
|
}
|
|
}
|
|
const nodeModules = metadata.__ompRuntimeNodeModules;
|
|
if (!nodeModules) return null;
|
|
const manifest = resolveRuntimeModule(nodeModules, `${ONNX_RUNTIME_NODE_PACKAGE}/package.json`);
|
|
return manifest ? path.dirname(manifest) : null;
|
|
}
|
|
|
|
export async function formatOnnxRuntimeCudaDiagnostics(
|
|
metadata: TransformersRuntimeMetadata,
|
|
requestedDevice: string,
|
|
error: unknown,
|
|
): Promise<string | null> {
|
|
const device = requestedDevice.trim().toLowerCase();
|
|
if (device !== "cuda" && device !== "gpu" && device !== "auto") return null;
|
|
if (process.platform !== "linux" || process.arch !== "x64") return null;
|
|
const packageDir = resolveOnnxRuntimePackageDir(metadata);
|
|
if (!packageDir) {
|
|
return [
|
|
"ONNX Runtime CUDA diagnostics:",
|
|
` PI_TINY_DEVICE=${requestedDevice} requested CUDAExecutionProvider`,
|
|
" cause: unable to resolve onnxruntime-node in the tiny-model runtime",
|
|
].join("\n");
|
|
}
|
|
const binDir = path.join(packageDir, LINUX_X64_ONNX_RUNTIME_CUDA_PROVIDER_DIR);
|
|
const missingFiles = await missingOnnxRuntimeCudaProviderFiles(binDir);
|
|
const sideRuntime = metadata.__ompRuntimeNodeModules;
|
|
const lines = [
|
|
"ONNX Runtime CUDA diagnostics:",
|
|
` PI_TINY_DEVICE=${requestedDevice} requested CUDAExecutionProvider`,
|
|
sideRuntime ? ` side runtime: ${sideRuntime}` : ` onnxruntime-node: ${packageDir}`,
|
|
` cause: ${cudaFailureCause(metadata, error, missingFiles)}`,
|
|
];
|
|
lines.push(` hint: ${cudaFailureHint(metadata, error, missingFiles)}`);
|
|
return lines.join("\n");
|
|
}
|
|
|
|
function configureTransformers<T extends ConfigurableTransformers>(transformers: T): T {
|
|
transformers.env.cacheDir = getTinyModelsCacheDir();
|
|
transformers.env.allowLocalModels = false;
|
|
transformers.env.logLevel = transformers.LogLevel.ERROR;
|
|
return transformers;
|
|
}
|
|
|
|
/**
|
|
* Memoize an async runtime load so it runs at most once per process, clearing
|
|
* the cache on failure so a later call can retry. Each worker holds one
|
|
* instance per runtime it loads.
|
|
*/
|
|
export class MemoizedRuntime<T> {
|
|
#promise: Promise<T> | null = null;
|
|
|
|
load(build: () => Promise<T>): Promise<T> {
|
|
if (this.#promise) return this.#promise;
|
|
const promise = build().catch(error => {
|
|
this.#promise = null;
|
|
throw error;
|
|
});
|
|
this.#promise = promise;
|
|
return promise;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Load the `@huggingface/transformers` runtime into `holder` (memoized): from
|
|
* the ambient install when running from source, or from a version-keyed side
|
|
* runtime (resolved lazily at `runtimeDir()`) when running as a compiled binary.
|
|
* The result is cast to the caller's concrete runtime type `T`.
|
|
*/
|
|
export function loadTransformersRuntime<T extends ConfigurableTransformers, K>(
|
|
holder: MemoizedRuntime<T>,
|
|
transport: WorkerProgressTransport<K>,
|
|
requestId: string,
|
|
modelKey: K,
|
|
runtimeDir: () => string,
|
|
): Promise<T> {
|
|
return holder.load(async () => {
|
|
if (!isCompiledBinary()) {
|
|
const entry = sourceRequire.resolve(TRANSFORMERS_PACKAGE);
|
|
return attachTransformersRuntimeMetadata(configureTransformers(sourceRequire(entry) as T), {
|
|
__ompTransformersEntry: entry,
|
|
});
|
|
}
|
|
const installedDir = await ensureRuntimeInstalled({
|
|
runtimeDir: runtimeDir(),
|
|
install: {
|
|
dependencies: { [TRANSFORMERS_PACKAGE]: getTransformersVersionSpec() },
|
|
trustedDependencies: ["onnxruntime-node"],
|
|
},
|
|
probePackage: TRANSFORMERS_PACKAGE,
|
|
onPhase: phase =>
|
|
transport.send({
|
|
type: "progress",
|
|
id: requestId,
|
|
event: {
|
|
modelKey,
|
|
status: phase,
|
|
name: `${TRANSFORMERS_PACKAGE}@${getTransformersVersionSpec()}`,
|
|
},
|
|
}),
|
|
});
|
|
let cudaRepairError: string | undefined;
|
|
try {
|
|
await ensureOnnxRuntimeCudaProviders(installedDir);
|
|
} catch (repairError) {
|
|
// Deferred failure: keep loading Transformers so `loadPipelineWithDeviceFallback`
|
|
// still gets its CUDA→CPU retry. The error is surfaced through the CUDA
|
|
// diagnostics attached to the runtime metadata.
|
|
cudaRepairError = errorMessage(repairError);
|
|
}
|
|
const entry = await prepareCompiledRuntime(installedDir, TRANSFORMERS_PACKAGE);
|
|
const require_ = createRequire(entry);
|
|
return attachTransformersRuntimeMetadata(configureTransformers(require_(entry) as T), {
|
|
__ompRuntimeNodeModules: path.join(installedDir, "node_modules"),
|
|
__ompTransformersEntry: entry,
|
|
__ompCudaRepairError: cudaRepairError,
|
|
});
|
|
});
|
|
}
|