mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix: GPU deployment robustness (6 fixes from end-to-end testing on an RTX 4070) (#334)
* fix(docker): pin CUDA base to 12.6 so the GPU image starts on R560+ drivers The amd64 base nvidia/cuda:12.9.2-cudnn-runtime bakes a cuda>=12.9 driver gate enforced by nvidia-container-toolkit at container start, so the image fails to launch on common production drivers (e.g. 570.x / CUDA 12.8). The AI bundles are all cu126 wheels and the image installs libcublas-12-6, so 12.9 was misaligned with the workload. Pin to nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04 to match the wheels and lower the driver floor to R560+. * fix(ai): broaden OOM detection so the rembg lighter-model fallback fires onnxruntime/CUDA allocation failures surface as 'Failed to allocate memory for requested buffer', CUBLAS_STATUS_ALLOC_FAILED, or bad_alloc, not just 'out of memory'. The background-removal and transparency-fixer fallback-to-lighter-model paths only matched the literal 'out of memory', so the fallback was dead code and transparency-fixer (default birefnet-hr-matting) always failed with an allocation error. Add isMemoryAllocError() and use it in both checks. * fix(ai): use bundled PaddleOCR models so OCR runs offline ocr.py passed no model dirs to PaddleOCR, so PaddleX resolved models from ~/.paddlex and downloaded them from HuggingFace at runtime (slow first use, broken air-gapped), ignoring the models the OCR bundle ships in MODELS_PATH; it also pulled doc-orientation/unwarping models that are not bundled. Pin detection, recognition and textline models to the bundled dirs in MODELS_PATH (per language) and disable use_doc_orientation_classify / use_doc_unwarping, with per-component fallback when a model is absent. Verified: OCR runs with zero HuggingFace requests. * fix(docker): add CAP_KILL so container shutdown is graceful cap_drop: ALL without re-adding KILL meant tini (PID 1, root) could not forward SIGTERM to the gosu-dropped snapotter process (root minus CAP_KILL cannot signal a different UID). docker stop logged '[FATAL tini] forwarding signal: Operation not permitted', never delivered the signal, and fell back to SIGKILL after the 10s timeout. Add KILL to cap_add in both compose files. Verified: docker stop completes in 0s with SIGTERM delivered (exit 143) and no FATAL tini. * fix(ai): serialize bundle installs against AI jobs to prevent sidecar segfault A feature bundle install rewrites the shared Python venv (pip + copytree of site-packages/*.so) as a background subprocess, with no coordination against AI tool jobs that dlopen native libs (torch / onnxruntime CUDA) from the same venv; a job loading a shared object while it is overwritten segfaults the sidecar. Add a process-wide async mutex (venv-lock.ts): bridge.run() acquires it before every AI script and the install route holds it across the installer subprocess. Both run in the same Node process so a module-level lock suffices. Verified: concurrent install + AI job produces zero segfaults and the job serializes behind the install. * fix(ai): make the venv lock read/write so concurrent AI jobs are not serialized The first cut used an exclusive mutex, which (a) deferred the dispatcher spawn by a microtask and broke unit tests that synchronously drive the mocked spawn, and (b) serialized AI jobs against each other, removing the dispatcher's by-id request multiplexing. Make it a writer-preferring read/write lock: AI jobs are shared readers (with a synchronous fast path so spawn still happens in-tick) and a bundle install is the exclusive writer. Verified: all 764 AI unit tests pass. * fix(ai): degrade OCR to Tesseract on CPU-only hosts instead of segfaulting The amd64 AI bundle ships paddlepaddle-gpu, whose native libs dlopen libcuda.so.1 at import and segfault on a host without a GPU (libcuda is the driver lib, injected only by nvidia-container-toolkit on GPU hosts). The segfault crashed the shared long-lived AI dispatcher and, after a few attempts, tripped the bridge crash-recovery permanent-disable, wedging all AI until a container restart. The standalone ocr tool defaults to quality=balanced (PaddleOCR), so it hit this on every CPU-only deployment; ocr-pdf already hardcoded Tesseract and was unaffected. ocr.py now gates the PaddleOCR tiers on gpu_available(): balanced/best transparently fall back to fast (Tesseract, CPU-capable) when no usable GPU is present, and run_paddleocr_v5/run_paddleocr_vl refuse before importing paddle so the GPU build is never dlopen'd on CPU. GPU hosts are unchanged. Verified on a CPU-only Windows/WSL2 box: ocr returns Tesseract text across repeated runs with the dispatcher staying healthy (no wedge).
This commit is contained in:
@@ -15,6 +15,19 @@ export interface RemoveBackgroundOptions {
|
||||
const MAX_REMBG_PX = Number(process.env.MAX_REMBG_PX) || 2048;
|
||||
const OOM_FALLBACK_MODEL = "u2net";
|
||||
|
||||
/**
|
||||
* onnxruntime / CUDA surface allocation failures with several different
|
||||
* messages ("out of memory", "Failed to allocate memory for requested buffer",
|
||||
* CUBLAS_STATUS_ALLOC_FAILED, bad_alloc). Match them all so the lighter-model
|
||||
* fallback actually triggers instead of failing the job.
|
||||
*/
|
||||
export function isMemoryAllocError(err: unknown): boolean {
|
||||
if (!(err instanceof Error)) return false;
|
||||
return /out of memory|failed to allocate|cudaerrormemoryallocation|cublas_status_alloc_failed|bad_alloc/i.test(
|
||||
err.message,
|
||||
);
|
||||
}
|
||||
|
||||
export async function removeBackground(
|
||||
inputBuffer: Buffer,
|
||||
outputDir: string,
|
||||
@@ -80,7 +93,7 @@ async function runAndParse(
|
||||
}
|
||||
return readFile(outputPath);
|
||||
} catch (err) {
|
||||
const isOom = err instanceof Error && err.message.includes("out of memory");
|
||||
const isOom = isMemoryAllocError(err);
|
||||
const canFallback = isOom && options.model !== OOM_FALLBACK_MODEL;
|
||||
|
||||
if (!canFallback) throw err;
|
||||
|
||||
+27
-16
@@ -4,6 +4,7 @@ import { dirname, resolve } from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { context, propagation, SpanStatusCode, trace } from "@opentelemetry/api";
|
||||
import { missingBundleForScript } from "./feature-gate.js";
|
||||
import { acquireVenvRead, tryAcquireVenvRead } from "./venv-lock.js";
|
||||
|
||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
const PYTHON_DIR = resolve(__dirname, "../python");
|
||||
@@ -624,23 +625,33 @@ export class PythonDispatcher {
|
||||
return this.runPerRequest(scriptName, args, options);
|
||||
};
|
||||
|
||||
if (!span) return doRun();
|
||||
const runWithSpan = (): Promise<{ stdout: string; stderr: string }> => {
|
||||
if (!span) return doRun();
|
||||
return doRun().then(
|
||||
(result) => {
|
||||
span.end();
|
||||
return result;
|
||||
},
|
||||
(err) => {
|
||||
span.setStatus({
|
||||
code: SpanStatusCode.ERROR,
|
||||
message: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
span.recordException(err instanceof Error ? err : new Error(String(err)));
|
||||
span.end();
|
||||
throw err;
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
return doRun().then(
|
||||
(result) => {
|
||||
span.end();
|
||||
return result;
|
||||
},
|
||||
(err) => {
|
||||
span.setStatus({
|
||||
code: SpanStatusCode.ERROR,
|
||||
message: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
span.recordException(err instanceof Error ? err : new Error(String(err)));
|
||||
span.end();
|
||||
throw err;
|
||||
},
|
||||
);
|
||||
// Serialize against bundle installs that mutate the shared venv: a job that
|
||||
// dlopens native libs (torch / onnxruntime CUDA) while pip rewrites them
|
||||
// segfaults the sidecar. In the common case the lock is free, so acquire it
|
||||
// synchronously and run inline (no microtask deferral, preserving the prior
|
||||
// call timing); only when an install holds it do we await.
|
||||
const release = tryAcquireVenvRead();
|
||||
if (release) return runWithSpan().finally(release);
|
||||
return acquireVenvRead().then((r) => runWithSpan().finally(r));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
export { removeBackground } from "./background-removal.js";
|
||||
export { isMemoryAllocError, removeBackground } from "./background-removal.js";
|
||||
export type { DispatcherStatus } from "./bridge.js";
|
||||
export {
|
||||
getDispatcherStatus,
|
||||
@@ -29,3 +29,4 @@ export { seamCarve } from "./seam-carving.js";
|
||||
export type { TranscribeOptions, TranscriptionResult, TranscriptSegment } from "./transcription.js";
|
||||
export { transcribeAudio } from "./transcription.js";
|
||||
export { upscale } from "./upscaling.js";
|
||||
export { acquireVenvLock } from "./venv-lock.js";
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* Read/write lock between AI tool jobs (readers) and bundle installs (writers).
|
||||
*
|
||||
* A bundle install rewrites the shared venv's site-packages (pip + copytree of
|
||||
* *.so), while AI tool jobs dlopen native libraries (torch / onnxruntime CUDA)
|
||||
* from that same venv. Loading a shared object while it is being overwritten
|
||||
* segfaults the sidecar, so an install must not overlap with any AI job.
|
||||
*
|
||||
* AI jobs may run concurrently with each other -- the persistent dispatcher
|
||||
* multiplexes requests by id -- so they are shared readers; only an install
|
||||
* needs exclusivity, so it is the writer. Writer-preferring, so an install is
|
||||
* not starved by a steady stream of jobs.
|
||||
*
|
||||
* Both sides live in the same Node process (the Fastify route spawns the
|
||||
* installer; in-process BullMQ workers run the jobs), so a module-level lock
|
||||
* shared via Node's module cache is sufficient.
|
||||
*
|
||||
* Readers have a synchronous fast path (tryAcquireVenvRead) so a job's work
|
||||
* begins in the same tick when no install is active or pending -- exactly as
|
||||
* before this lock existed. Every acquire returns a release function; always
|
||||
* call it (it is idempotent).
|
||||
*/
|
||||
type Release = () => void;
|
||||
|
||||
let readers = 0;
|
||||
let writerActive = false;
|
||||
const writeWaiters: Array<() => void> = [];
|
||||
const readWaiters: Array<() => void> = [];
|
||||
|
||||
function readRelease(): Release {
|
||||
let done = false;
|
||||
return () => {
|
||||
if (done) return;
|
||||
done = true;
|
||||
readers--;
|
||||
if (readers === 0 && writeWaiters.length > 0) {
|
||||
writerActive = true;
|
||||
writeWaiters.shift()?.();
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
function writeRelease(): Release {
|
||||
let done = false;
|
||||
return () => {
|
||||
if (done) return;
|
||||
done = true;
|
||||
writerActive = false;
|
||||
if (writeWaiters.length > 0) {
|
||||
writerActive = true;
|
||||
writeWaiters.shift()?.();
|
||||
} else {
|
||||
const granted = readWaiters.splice(0);
|
||||
for (const grant of granted) grant();
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/** AI-job (reader) sync fast path; null if an install is active or waiting. */
|
||||
export function tryAcquireVenvRead(): Release | null {
|
||||
if (writerActive || writeWaiters.length > 0) return null;
|
||||
readers++;
|
||||
return readRelease();
|
||||
}
|
||||
|
||||
/** AI-job (reader) acquire; waits while an install holds or is awaiting the lock. */
|
||||
export function acquireVenvRead(): Promise<Release> {
|
||||
const r = tryAcquireVenvRead();
|
||||
if (r) return Promise.resolve(r);
|
||||
return new Promise<Release>((resolve) => {
|
||||
readWaiters.push(() => {
|
||||
readers++;
|
||||
resolve(readRelease());
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** Bundle-install (writer) exclusive acquire; waits for in-flight AI jobs. */
|
||||
export function acquireVenvLock(): Promise<Release> {
|
||||
if (!writerActive && readers === 0 && writeWaiters.length === 0) {
|
||||
writerActive = true;
|
||||
return Promise.resolve(writeRelease());
|
||||
}
|
||||
return new Promise<Release>((resolve) => {
|
||||
writeWaiters.push(() => resolve(writeRelease()));
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user