mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
Route both Python exit paths through pythonExitError so the reason survives the scrubber; OOM/segfault stay operational and keep "out of memory" for the lighter-model fallback.
845 lines
28 KiB
TypeScript
845 lines
28 KiB
TypeScript
import { type ChildProcess, spawn } from "node:child_process";
|
|
import { randomUUID } from "node:crypto";
|
|
import { dirname, resolve } from "node:path";
|
|
import { fileURLToPath } from "node:url";
|
|
import { context, propagation, SpanStatusCode, trace } from "@opentelemetry/api";
|
|
import { SafeError } from "@snapotter/shared";
|
|
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");
|
|
|
|
function appendEnvPath(base: string, suffix: string): string {
|
|
const normalizedBase = base.replace(/\/+$/, "");
|
|
return `${normalizedBase || "/"}${normalizedBase === "" ? "" : "/"}${suffix}`;
|
|
}
|
|
|
|
/**
|
|
* Build a minimal environment for spawned Python processes.
|
|
* Only passes through variables needed for venv, CUDA, model cache,
|
|
* package resolution, and locale -- avoids leaking secrets or
|
|
* application config from the parent process.
|
|
*/
|
|
function buildMinimalEnv(): Record<string, string> {
|
|
const env: Record<string, string> = {
|
|
PYTHONUNBUFFERED: "1",
|
|
LANG: process.env.LANG || "C.UTF-8",
|
|
};
|
|
const passthrough = [
|
|
"PATH",
|
|
"HOME",
|
|
"VIRTUAL_ENV",
|
|
"PYTHONPATH",
|
|
"CUDA_VISIBLE_DEVICES",
|
|
"LD_LIBRARY_PATH",
|
|
// Application-specific vars the sidecar scripts depend on
|
|
"DATA_DIR",
|
|
"MODELS_PATH",
|
|
"U2NET_HOME",
|
|
"PROCESSING_TIMEOUT_S",
|
|
"DISPATCHER_MAX_REQUESTS",
|
|
"PYTHON_VENV_PATH",
|
|
"SNAPOTTER_GPU",
|
|
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
|
"OTEL_EXPORTER_OTLP_PROTOCOL",
|
|
"OTEL_EXPORTER_OTLP_HEADERS",
|
|
];
|
|
for (const key of passthrough) {
|
|
if (process.env[key] !== undefined) {
|
|
env[key] = process.env[key] as string;
|
|
}
|
|
}
|
|
|
|
env.DATA_DIR ??= "./data";
|
|
env.MODELS_PATH ??= appendEnvPath(env.DATA_DIR, "ai/models");
|
|
|
|
// Runtime model downloads are allowed by default (public model weights
|
|
// only, never user data). SNAPOTTER_ALLOW_MODEL_DOWNLOAD=0 enables strict
|
|
// offline mode for airgapped deployments: the sidecar then gets the
|
|
// Hugging Face offline flags and every download fallback raises an
|
|
// actionable error instead of fetching. Bundle installs stay exempt
|
|
// because install_feature.py lifts the flags in its own process.
|
|
const allowModelDownload = process.env.SNAPOTTER_ALLOW_MODEL_DOWNLOAD;
|
|
if (allowModelDownload !== undefined) {
|
|
env.SNAPOTTER_ALLOW_MODEL_DOWNLOAD = allowModelDownload;
|
|
}
|
|
if (allowModelDownload === "0" || allowModelDownload?.toLowerCase() === "false") {
|
|
env.HF_HUB_OFFLINE = "1";
|
|
env.TRANSFORMERS_OFFLINE = "1";
|
|
}
|
|
return env;
|
|
}
|
|
|
|
/** Try venv first, then system python. */
|
|
function getPythonPath(): string {
|
|
const venvPath = process.env.PYTHON_VENV_PATH || resolve(__dirname, "../../../.venv");
|
|
return `${venvPath}/bin/python3`;
|
|
}
|
|
|
|
/**
|
|
* Extract a user-friendly error from a Python process error.
|
|
*/
|
|
const OOM_EXIT_TEXT =
|
|
/out of memory|failed to allocate|cudaerrormemoryallocation|cublas_status_alloc_failed|bad_alloc/i;
|
|
|
|
/**
|
|
* Build the error to reject a non-zero Python exit with. The extracted reason is
|
|
* kept as the message (it is already surfaced to callers) but wrapped in a
|
|
* SafeError so it survives the API's Sentry scrubber, which otherwise reduces a
|
|
* plain Error to "Error: Error" (NODE-24). Classification is deliberate: a kill
|
|
* signal / OS OOM (SIGKILL, exit 137) and a segfault (SIGSEGV, exit 139) are
|
|
* operational, and their constant messages keep the "out of memory" text so
|
|
* memory-aware callers still retry on a lighter model. Any other non-zero exit
|
|
* is a bug, with the extracted reason (or an OOM reported in-band) preserved.
|
|
*/
|
|
function pythonExitError(code: number | null, signal: string | null, extracted: string): SafeError {
|
|
if (signal === "SIGSEGV" || code === 139) {
|
|
return new SafeError("Process crashed (segmentation fault)", {
|
|
kind: "operational",
|
|
code: `exit-${code ?? signal ?? "unknown"}`,
|
|
});
|
|
}
|
|
if (signal === "SIGKILL" || code === 137) {
|
|
return new SafeError("Process killed (out of memory) -- try a lighter model or smaller image", {
|
|
kind: "operational",
|
|
code: `exit-${code ?? signal ?? "unknown"}`,
|
|
});
|
|
}
|
|
const message = extracted || `Python script exited with code ${code}`;
|
|
return new SafeError(message, {
|
|
kind: OOM_EXIT_TEXT.test(message) ? "operational" : "bug",
|
|
code: `exit-${code ?? "unknown"}`,
|
|
});
|
|
}
|
|
|
|
function extractPythonError(error: unknown): string {
|
|
if (error && typeof error === "object") {
|
|
const pErr = error as {
|
|
stderr?: string;
|
|
stdout?: string;
|
|
message?: string;
|
|
};
|
|
for (const output of [pErr.stdout, pErr.stderr]) {
|
|
if (output) {
|
|
try {
|
|
const parsed = JSON.parse(output.trim());
|
|
if (parsed.error) return parsed.error;
|
|
} catch {
|
|
const trimmed = output.trim();
|
|
if (trimmed && !trimmed.startsWith("Traceback")) {
|
|
return trimmed;
|
|
}
|
|
// Extract the last meaningful line from a Python Traceback
|
|
if (trimmed) {
|
|
const lines = trimmed
|
|
.split("\n")
|
|
.map((l) => l.trim())
|
|
.filter(Boolean);
|
|
const lastLine = lines[lines.length - 1];
|
|
if (lastLine && lastLine !== "Traceback (most recent call last):") {
|
|
return lastLine;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if (pErr.message) return pErr.message;
|
|
// Return empty string for {stdout, stderr} objects with no useful content
|
|
// so the caller's fallback message (e.g. exit code) kicks in.
|
|
return "";
|
|
}
|
|
if (error instanceof Error) return error.message;
|
|
return String(error);
|
|
}
|
|
|
|
export type ProgressCallback = (percent: number, stage: string) => void;
|
|
|
|
// ── PythonDispatcher class ─────────────────────────────────────────
|
|
|
|
interface PendingRequest {
|
|
resolve: (result: { stdout: string; stderr: string }) => void;
|
|
reject: (err: Error) => void;
|
|
onProgress?: ProgressCallback;
|
|
stderrLines: string[];
|
|
/** Which child generation this request was written to (see startChild). */
|
|
generation: number;
|
|
}
|
|
|
|
// Crash recovery constants
|
|
const CRASH_WINDOW_MS = 60_000;
|
|
const MAX_CONSECUTIVE_CRASHES = 5;
|
|
const BASE_BACKOFF_MS = 1_000;
|
|
|
|
export interface DispatcherStatus {
|
|
running: boolean;
|
|
ready: boolean;
|
|
failed: boolean;
|
|
gpu: boolean;
|
|
pid: number | null;
|
|
consecutiveCrashes: number;
|
|
}
|
|
|
|
/**
|
|
* A single Python dispatcher instance parameterized by profile.
|
|
* The "ai" profile loads heavy ML libraries at startup; the "docs"
|
|
* profile starts lean with only document-processing scripts allowed.
|
|
*/
|
|
export class PythonDispatcher {
|
|
private profile: "ai" | "docs";
|
|
private child: ChildProcess | null = null;
|
|
private childReady = false;
|
|
private childFailed = false;
|
|
private gpuAvail = false;
|
|
private pending = new Map<string, PendingRequest>();
|
|
private crashes = 0;
|
|
private lastCrashTs = 0;
|
|
private backoffEnd = 0;
|
|
/**
|
|
* Monotonic child counter. Each spawned child and every request written to
|
|
* it carry the generation current at spawn time, so a stale child's late
|
|
* close/error events (SIGTERM delivery can lag a replacement spawn) only
|
|
* ever touch their own generation's pending requests.
|
|
*/
|
|
private generation = 0;
|
|
/**
|
|
* Children we SIGTERMed on purpose (shutdown/reload). Tracked per child
|
|
* rather than as an instance-wide flag: an instance flag reset by the next
|
|
* spawn would let the stale child's close event record a phantom crash and
|
|
* null out the fresh child. The request-timeout kill path deliberately does
|
|
* NOT add to this set, so a genuinely hung script still counts as a crash.
|
|
*/
|
|
private stoppedChildren = new WeakSet<ChildProcess>();
|
|
|
|
constructor(opts: { profile: "ai" | "docs" }) {
|
|
this.profile = opts.profile;
|
|
}
|
|
|
|
private buildEnv(): Record<string, string> {
|
|
const env = buildMinimalEnv();
|
|
env.DISPATCHER_PROFILE = this.profile;
|
|
return env;
|
|
}
|
|
|
|
private recordCrash(): void {
|
|
const now = Date.now();
|
|
if (now - this.lastCrashTs > CRASH_WINDOW_MS) {
|
|
this.crashes = 1;
|
|
} else {
|
|
this.crashes++;
|
|
}
|
|
this.lastCrashTs = now;
|
|
|
|
if (this.crashes >= MAX_CONSECUTIVE_CRASHES) {
|
|
console.error(
|
|
`[bridge] Dispatcher crashed ${this.crashes} times in ${CRASH_WINDOW_MS / 1000}s, disabling permanently`,
|
|
);
|
|
this.childFailed = true;
|
|
return;
|
|
}
|
|
|
|
const delay = BASE_BACKOFF_MS * 2 ** (this.crashes - 1);
|
|
this.backoffEnd = now + delay;
|
|
console.warn(
|
|
`[bridge] Dispatcher crash #${this.crashes}, backing off ${delay}ms before restart`,
|
|
);
|
|
}
|
|
|
|
/** Reject and drop the pending requests written to one child generation. */
|
|
private rejectPendingForGeneration(generation: number, message: string): void {
|
|
for (const [id, req] of this.pending.entries()) {
|
|
if (req.generation !== generation) continue;
|
|
req.reject(new Error(message));
|
|
this.pending.delete(id);
|
|
}
|
|
}
|
|
|
|
private startChild(): ChildProcess | null {
|
|
if (this.childFailed) return null;
|
|
this.generation++;
|
|
const gen = this.generation;
|
|
|
|
try {
|
|
const proc = spawn(getPythonPath(), [resolve(PYTHON_DIR, "dispatcher.py")], {
|
|
stdio: ["pipe", "pipe", "pipe"],
|
|
env: this.buildEnv(),
|
|
});
|
|
|
|
proc.stdin?.on("error", (err: NodeJS.ErrnoException) => {
|
|
if (err.code === "EPIPE" || err.code === "ERR_STREAM_DESTROYED") {
|
|
console.error(
|
|
`[bridge] Dispatcher stdin pipe broken (${err.code}), rejecting pending requests`,
|
|
);
|
|
this.rejectPendingForGeneration(gen, "Python dispatcher stdin closed unexpectedly");
|
|
// An intentional shutdown() ends stdin then SIGTERMs the child,
|
|
// which can surface here as an EPIPE/ERR_STREAM_DESTROYED. That is
|
|
// not a crash: counting it would let repeated legitimate restarts
|
|
// (shutdownDispatcher() runs after every AI bundle install) trip
|
|
// the crash limit and permanently disable the dispatcher. Guard
|
|
// mirrors the "close" handler below.
|
|
if (!this.stoppedChildren.has(proc)) this.recordCrash();
|
|
if (this.child === proc) {
|
|
this.child = null;
|
|
this.childReady = false;
|
|
}
|
|
}
|
|
});
|
|
|
|
let stderrBuf = "";
|
|
let stdoutBuf = "";
|
|
|
|
proc.stderr?.on("data", (chunk: Buffer) => {
|
|
stderrBuf += chunk.toString();
|
|
const lines = stderrBuf.split("\n");
|
|
stderrBuf = lines.pop() ?? "";
|
|
|
|
for (const line of lines) {
|
|
const trimmed = line.trim();
|
|
if (!trimmed) continue;
|
|
|
|
try {
|
|
const parsed = JSON.parse(trimmed);
|
|
|
|
// Readiness signal. Ignore it from a superseded child so a stale
|
|
// process can't mark a not-yet-ready replacement as ready.
|
|
if (parsed.ready === true) {
|
|
if (this.child === proc) {
|
|
this.childReady = true;
|
|
this.gpuAvail = parsed.gpu === true;
|
|
this.crashes = 0;
|
|
console.log(`[bridge] Python dispatcher ready (GPU: ${parsed.gpu === true})`);
|
|
}
|
|
continue;
|
|
}
|
|
|
|
// Progress event - route to the currently active request
|
|
if (typeof parsed.progress === "number" && typeof parsed.stage === "string") {
|
|
// Progress goes to this child's pending requests (only one
|
|
// should be active at a time since Python processes synchronously)
|
|
for (const req of this.pending.values()) {
|
|
if (req.generation === gen) req.onProgress?.(parsed.progress, parsed.stage);
|
|
}
|
|
continue;
|
|
}
|
|
|
|
// Diagnostic notices are forwarded so they reach docker logs
|
|
// instead of being silently dropped for matching neither shape.
|
|
if (typeof parsed.info === "string") {
|
|
console.log(`[python] ${parsed.info}`);
|
|
continue;
|
|
}
|
|
if (typeof parsed.warning === "string") {
|
|
console.warn(`[python] ${parsed.warning}`);
|
|
}
|
|
} catch {
|
|
// Not JSON - forward diagnostic messages to Node.js logger,
|
|
// collect the rest as error output for pending requests.
|
|
if (trimmed.startsWith("[")) {
|
|
console.log(`[python] ${trimmed}`);
|
|
}
|
|
for (const req of this.pending.values()) {
|
|
if (req.generation === gen) req.stderrLines.push(trimmed);
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
proc.stdout?.on("data", (chunk: Buffer) => {
|
|
stdoutBuf += chunk.toString();
|
|
const lines = stdoutBuf.split("\n");
|
|
stdoutBuf = lines.pop() ?? "";
|
|
|
|
for (const line of lines) {
|
|
const trimmed = line.trim();
|
|
if (!trimmed) continue;
|
|
|
|
try {
|
|
const response = JSON.parse(trimmed);
|
|
const reqId = response.id;
|
|
const req = this.pending.get(reqId);
|
|
if (req) {
|
|
this.pending.delete(reqId);
|
|
if (response.exitCode !== 0) {
|
|
const extracted = extractPythonError({
|
|
stdout: response.stdout,
|
|
stderr: req.stderrLines.join("\n"),
|
|
});
|
|
req.reject(pythonExitError(response.exitCode, null, extracted));
|
|
} else {
|
|
req.resolve({
|
|
stdout: response.stdout || "",
|
|
stderr: req.stderrLines.join("\n"),
|
|
});
|
|
}
|
|
}
|
|
} catch {
|
|
// Not a valid response line
|
|
}
|
|
}
|
|
});
|
|
|
|
proc.on("error", (err: NodeJS.ErrnoException) => {
|
|
console.error(`[bridge] Dispatcher error: ${err.message} (code: ${err.code})`);
|
|
if (err.code === "ENOENT") {
|
|
this.childFailed = true;
|
|
} else if (!this.stoppedChildren.has(proc)) {
|
|
// Skip crash accounting when we initiated the teardown (shutdown()
|
|
// marks the child stopped before killing it); mirrors "close".
|
|
this.recordCrash();
|
|
}
|
|
this.rejectPendingForGeneration(gen, extractPythonError(err));
|
|
if (this.child === proc) {
|
|
this.child = null;
|
|
this.childReady = false;
|
|
}
|
|
});
|
|
|
|
proc.on("close", (code) => {
|
|
this.rejectPendingForGeneration(gen, "Python dispatcher exited unexpectedly");
|
|
if (code !== 0 && !this.stoppedChildren.has(proc)) {
|
|
this.recordCrash();
|
|
}
|
|
if (this.child === proc) {
|
|
this.child = null;
|
|
this.childReady = false;
|
|
}
|
|
});
|
|
|
|
return proc;
|
|
} catch {
|
|
this.childFailed = true;
|
|
return null;
|
|
}
|
|
}
|
|
|
|
private getChild(): ChildProcess | null {
|
|
if (this.childFailed) return null;
|
|
if (!this.child || this.child.killed) {
|
|
if (Date.now() < this.backoffEnd) return null;
|
|
this.child = this.startChild();
|
|
}
|
|
return this.child;
|
|
}
|
|
|
|
/**
|
|
* Send a request to the persistent Python dispatcher.
|
|
* Returns null if the dispatcher is unavailable (caller should fall back).
|
|
*/
|
|
private dispatcherRun(
|
|
scriptName: string,
|
|
args: string[],
|
|
options: { onProgress?: ProgressCallback; timeout?: number } = {},
|
|
): Promise<{ stdout: string; stderr: string }> | null {
|
|
const proc = this.getChild();
|
|
if (!proc || !proc.stdin || !this.childReady) return null;
|
|
|
|
const id = randomUUID();
|
|
const timeout =
|
|
options.timeout ??
|
|
(process.env.PROCESSING_TIMEOUT_S && parseInt(process.env.PROCESSING_TIMEOUT_S, 10) > 0
|
|
? parseInt(process.env.PROCESSING_TIMEOUT_S, 10) * 1000
|
|
: 600000);
|
|
|
|
return new Promise((resolvePromise, rejectPromise) => {
|
|
const timer = setTimeout(() => {
|
|
this.pending.delete(id);
|
|
// Kill the stuck dispatcher so it restarts on the next request instead of
|
|
// blocking all subsequent operations behind the timed-out script.
|
|
if (this.child && !this.child.killed) {
|
|
this.child.kill("SIGTERM");
|
|
}
|
|
rejectPromise(
|
|
new SafeError("Python script timed out", { kind: "operational", code: "timeout" }),
|
|
);
|
|
}, timeout);
|
|
|
|
const wrappedResolve = (result: { stdout: string; stderr: string }) => {
|
|
clearTimeout(timer);
|
|
resolvePromise(result);
|
|
};
|
|
|
|
const wrappedReject = (err: Error) => {
|
|
clearTimeout(timer);
|
|
rejectPromise(err);
|
|
};
|
|
|
|
this.pending.set(id, {
|
|
resolve: wrappedResolve,
|
|
reject: wrappedReject,
|
|
onProgress: options.onProgress,
|
|
stderrLines: [],
|
|
// getChild() above either reused or just spawned the child this
|
|
// request is written to, so the current generation is its generation.
|
|
generation: this.generation,
|
|
});
|
|
|
|
const msg: Record<string, unknown> = { id, script: scriptName.replace(".py", ""), args };
|
|
const otelCarrier: Record<string, string> = {};
|
|
propagation.inject(context.active(), otelCarrier);
|
|
if (otelCarrier.traceparent) {
|
|
msg._otel = {
|
|
traceparent: otelCarrier.traceparent,
|
|
tracestate: otelCarrier.tracestate,
|
|
};
|
|
}
|
|
const request = JSON.stringify(msg);
|
|
try {
|
|
proc.stdin!.write(request + "\n");
|
|
} catch {
|
|
this.pending.delete(id);
|
|
clearTimeout(timer);
|
|
rejectPromise(new Error("Python dispatcher stdin closed unexpectedly"));
|
|
}
|
|
});
|
|
}
|
|
|
|
// ── Per-request fallback (original implementation) ──────────────
|
|
|
|
private runPerRequest(
|
|
scriptName: string,
|
|
args: string[],
|
|
options: {
|
|
onProgress?: ProgressCallback;
|
|
timeout?: number;
|
|
} = {},
|
|
): Promise<{ stdout: string; stderr: string }> {
|
|
// Mirror the dispatcher's feature gate. The persistent dispatcher rejects
|
|
// scripts whose bundle is not installed (in Python); this fallback spawns
|
|
// scripts directly, so without the same check it would run them ungated
|
|
// when the dispatcher is down (e.g. right after a model repair). Reject
|
|
// with the same message the dispatcher path surfaces.
|
|
if (missingBundleForScript(scriptName)) {
|
|
return Promise.reject(new Error("feature_not_installed"));
|
|
}
|
|
const scriptPath = resolve(PYTHON_DIR, scriptName);
|
|
const timeout =
|
|
options.timeout ??
|
|
(process.env.PROCESSING_TIMEOUT_S && parseInt(process.env.PROCESSING_TIMEOUT_S, 10) > 0
|
|
? parseInt(process.env.PROCESSING_TIMEOUT_S, 10) * 1000
|
|
: 600000);
|
|
|
|
return new Promise((resolvePromise, rejectPromise) => {
|
|
const trySpawn = (pythonBin: string, isFallback: boolean) => {
|
|
const proc = spawn(pythonBin, [scriptPath, ...args], {
|
|
stdio: ["ignore", "pipe", "pipe"],
|
|
env: this.buildEnv(),
|
|
});
|
|
|
|
let stdout = "";
|
|
const stderrLines: string[] = [];
|
|
let stderrBuffer = "";
|
|
let timedOut = false;
|
|
|
|
const timer = setTimeout(() => {
|
|
timedOut = true;
|
|
proc.kill("SIGTERM");
|
|
}, timeout);
|
|
|
|
proc.stdout.on("data", (chunk: Buffer) => {
|
|
stdout += chunk.toString();
|
|
});
|
|
|
|
proc.stderr.on("data", (chunk: Buffer) => {
|
|
stderrBuffer += chunk.toString();
|
|
const lines = stderrBuffer.split("\n");
|
|
stderrBuffer = lines.pop() ?? "";
|
|
|
|
for (const line of lines) {
|
|
const trimmed = line.trim();
|
|
if (!trimmed) continue;
|
|
|
|
try {
|
|
const parsed = JSON.parse(trimmed);
|
|
if (typeof parsed.progress === "number" && typeof parsed.stage === "string") {
|
|
options.onProgress?.(parsed.progress, parsed.stage);
|
|
continue;
|
|
}
|
|
} catch {
|
|
// Not JSON - collect as regular stderr
|
|
}
|
|
stderrLines.push(trimmed);
|
|
}
|
|
});
|
|
|
|
proc.on("error", (err: NodeJS.ErrnoException) => {
|
|
clearTimeout(timer);
|
|
if (err.code === "ENOENT" && !isFallback) {
|
|
trySpawn("python3", true);
|
|
} else {
|
|
rejectPromise(new Error(extractPythonError(err)));
|
|
}
|
|
});
|
|
|
|
proc.on("close", (code, signal) => {
|
|
clearTimeout(timer);
|
|
|
|
if (stderrBuffer.trim()) {
|
|
stderrLines.push(stderrBuffer.trim());
|
|
}
|
|
|
|
if (timedOut) {
|
|
rejectPromise(
|
|
new SafeError("Python script timed out", { kind: "operational", code: "timeout" }),
|
|
);
|
|
return;
|
|
}
|
|
|
|
const stderr = stderrLines.join("\n");
|
|
|
|
if (code !== 0) {
|
|
rejectPromise(
|
|
pythonExitError(code, signal, extractPythonError({ stdout: stdout.trim(), stderr })),
|
|
);
|
|
return;
|
|
}
|
|
|
|
resolvePromise({ stdout: stdout.trim(), stderr });
|
|
});
|
|
};
|
|
|
|
trySpawn(getPythonPath(), false);
|
|
});
|
|
}
|
|
|
|
// ── Public API ──────────────────────────────────────────────────
|
|
|
|
/**
|
|
* Whether the Python dispatcher detected a CUDA GPU at startup.
|
|
*/
|
|
isGpuAvailable(): boolean {
|
|
return this.gpuAvail;
|
|
}
|
|
|
|
getStatus(): DispatcherStatus {
|
|
return {
|
|
running: this.child !== null && !this.child.killed,
|
|
ready: this.childReady,
|
|
failed: this.childFailed,
|
|
gpu: this.gpuAvail,
|
|
pid: this.child?.pid ?? null,
|
|
consecutiveCrashes: this.crashes,
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Shut down the persistent dispatcher process.
|
|
*/
|
|
shutdown(): void {
|
|
if (this.child && !this.child.killed) {
|
|
this.stoppedChildren.add(this.child);
|
|
this.child.stdin?.end();
|
|
this.child.kill("SIGTERM");
|
|
this.child = null;
|
|
this.childReady = false;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Eagerly start the Python dispatcher and wait for its readiness signal.
|
|
* Returns the GPU status once ready, or {ready: false} on timeout/failure.
|
|
* Safe to call multiple times -- idempotent if the dispatcher is already running.
|
|
*/
|
|
init(timeoutMs = 30_000): Promise<{ ready: boolean; gpu: boolean }> {
|
|
if (this.childReady) {
|
|
return Promise.resolve({ ready: true, gpu: this.gpuAvail });
|
|
}
|
|
if (this.childFailed) {
|
|
return Promise.resolve({ ready: false, gpu: false });
|
|
}
|
|
|
|
const proc = this.getChild();
|
|
if (!proc) {
|
|
return Promise.resolve({ ready: false, gpu: false });
|
|
}
|
|
|
|
return new Promise((resolvePromise) => {
|
|
const timer = setTimeout(() => {
|
|
clearInterval(poll);
|
|
resolvePromise({ ready: false, gpu: false });
|
|
}, timeoutMs);
|
|
|
|
const poll = setInterval(() => {
|
|
if (this.childReady) {
|
|
clearTimeout(timer);
|
|
clearInterval(poll);
|
|
resolvePromise({ ready: true, gpu: this.gpuAvail });
|
|
} else if (this.childFailed) {
|
|
clearTimeout(timer);
|
|
clearInterval(poll);
|
|
resolvePromise({ ready: false, gpu: false });
|
|
}
|
|
}, 50);
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Run a Python script with real-time progress streaming via stderr.
|
|
*
|
|
* Tries the persistent dispatcher first for warm-start performance.
|
|
* Falls back to per-request spawning if the dispatcher is unavailable.
|
|
*/
|
|
run(
|
|
scriptName: string,
|
|
args: string[],
|
|
options: {
|
|
onProgress?: ProgressCallback;
|
|
timeout?: number;
|
|
} = {},
|
|
): Promise<{ stdout: string; stderr: string }> {
|
|
const tracer = trace.getTracer("snapotter-sidecar");
|
|
const span = trace.getActiveSpan()
|
|
? tracer.startSpan("sidecar.execute", {
|
|
attributes: {
|
|
"sidecar.script": scriptName.replace(".py", ""),
|
|
"sidecar.profile": this.profile,
|
|
},
|
|
})
|
|
: null;
|
|
|
|
const doRun = (): Promise<{ stdout: string; stderr: string }> => {
|
|
// Try persistent dispatcher first
|
|
const dispatcherPromise = this.dispatcherRun(scriptName, args, options);
|
|
if (dispatcherPromise) {
|
|
return dispatcherPromise.catch((err: Error) => {
|
|
if (
|
|
err.message === "Python dispatcher exited unexpectedly" ||
|
|
err.message === "Python dispatcher stdin closed unexpectedly"
|
|
) {
|
|
console.warn(
|
|
`[bridge] Dispatcher crashed during ${scriptName}, retrying with per-request process`,
|
|
);
|
|
return this.runPerRequest(scriptName, args, options).then((result) => ({
|
|
...result,
|
|
stderr: `${result.stderr}\n[bridge] retried after dispatcher crash`,
|
|
}));
|
|
}
|
|
throw err;
|
|
});
|
|
}
|
|
|
|
// Fall back to per-request spawning
|
|
return this.runPerRequest(scriptName, args, options);
|
|
};
|
|
|
|
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;
|
|
},
|
|
);
|
|
};
|
|
|
|
// 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));
|
|
}
|
|
}
|
|
|
|
// ── Singleton AI dispatcher (preserves existing module-level API) ──
|
|
|
|
let aiDispatcher: PythonDispatcher | null = null;
|
|
|
|
function getAiDispatcher(): PythonDispatcher {
|
|
if (!aiDispatcher) aiDispatcher = new PythonDispatcher({ profile: "ai" });
|
|
return aiDispatcher;
|
|
}
|
|
|
|
/**
|
|
* Whether the Python dispatcher detected a CUDA GPU at startup.
|
|
*/
|
|
export function isGpuAvailable(): boolean {
|
|
return getAiDispatcher().isGpuAvailable();
|
|
}
|
|
|
|
export function getDispatcherStatus(): DispatcherStatus {
|
|
return getAiDispatcher().getStatus();
|
|
}
|
|
|
|
/**
|
|
* Shut down the persistent dispatcher process.
|
|
*/
|
|
export function shutdownDispatcher(): void {
|
|
getAiDispatcher().shutdown();
|
|
}
|
|
|
|
/**
|
|
* Eagerly start the Python dispatcher and wait for its readiness signal.
|
|
* Returns the GPU status once ready, or {ready: false} on timeout/failure.
|
|
* Safe to call multiple times -- idempotent if the dispatcher is already running.
|
|
*/
|
|
export function initDispatcher(timeoutMs = 30_000): Promise<{ ready: boolean; gpu: boolean }> {
|
|
return getAiDispatcher().init(timeoutMs);
|
|
}
|
|
|
|
/**
|
|
* Run a Python script with real-time progress streaming via stderr.
|
|
*
|
|
* Tries the persistent dispatcher first for warm-start performance.
|
|
* Falls back to per-request spawning if the dispatcher is unavailable.
|
|
*/
|
|
export function runPythonWithProgress(
|
|
scriptName: string,
|
|
args: string[],
|
|
options: {
|
|
onProgress?: ProgressCallback;
|
|
timeout?: number;
|
|
} = {},
|
|
): Promise<{ stdout: string; stderr: string }> {
|
|
return getAiDispatcher().run(scriptName, args, options);
|
|
}
|
|
|
|
// biome-ignore lint/suspicious/noExplicitAny: matches JSON.parse return type
|
|
export function parseStdoutJson(stdout: string): any {
|
|
const matched = stdout.match(/\{[\s\S]*\}$/);
|
|
if (!matched) throw new Error("No JSON response from Python script");
|
|
return JSON.parse(matched[0]);
|
|
}
|
|
|
|
// ── Docs-profile dispatcher ──────────────────────────────────────
|
|
|
|
let docsDispatcher: PythonDispatcher | null = null;
|
|
|
|
/** Lazy docs-profile dispatcher instance (same protocol, no new IPC). */
|
|
export function getDocsDispatcher(): PythonDispatcher {
|
|
if (!docsDispatcher) docsDispatcher = new PythonDispatcher({ profile: "docs" });
|
|
return docsDispatcher;
|
|
}
|
|
|
|
export async function runDocsScript(
|
|
script: string,
|
|
args: Record<string, unknown>,
|
|
opts?: { timeoutMs?: number },
|
|
): Promise<string> {
|
|
// Ensure .py extension: the dispatcher protocol strips it for the
|
|
// in-process path, while the per-request fallback needs it for the
|
|
// file path on disk.
|
|
const scriptFile = script.endsWith(".py") ? script : `${script}.py`;
|
|
const result = await getDocsDispatcher().run(
|
|
scriptFile,
|
|
[JSON.stringify(args)],
|
|
opts?.timeoutMs ? { timeout: opts.timeoutMs } : undefined,
|
|
);
|
|
return result.stdout;
|
|
}
|
|
|
|
export async function shutdownDocsDispatcher(): Promise<void> {
|
|
docsDispatcher?.shutdown();
|
|
docsDispatcher = null;
|
|
}
|