diff --git a/apps/api/src/index.ts b/apps/api/src/index.ts index 7fe143d8..2c3fdf11 100644 --- a/apps/api/src/index.ts +++ b/apps/api/src/index.ts @@ -1,7 +1,7 @@ import { randomUUID } from "node:crypto"; import cors from "@fastify/cors"; import rateLimit from "@fastify/rate-limit"; -import { getDispatcherStatus, isGpuAvailable } from "@snapotter/ai"; +import { getDispatcherStatus, initDispatcher, isGpuAvailable } from "@snapotter/ai"; import { APP_VERSION } from "@snapotter/shared"; import { eq } from "drizzle-orm"; import Fastify from "fastify"; @@ -242,12 +242,13 @@ const cleanupCron = startCleanupCron(); // Start try { await app.listen({ port: env.PORT, host: "0.0.0.0" }); - const dispatcherStatus = getDispatcherStatus(); - const gpuLine = !dispatcherStatus.ready - ? "[INFO] GPU status: waiting for AI sidecar startup..." - : dispatcherStatus.gpu - ? "[INFO] GPU detected — AI tools will use CUDA acceleration" - : "[WARN] No GPU detected — AI tools will use CPU (slower)"; + + const dispatcherResult = await initDispatcher(); + const gpuLine = dispatcherResult.ready + ? dispatcherResult.gpu + ? "[INFO] GPU detected -- AI tools will use CUDA acceleration" + : "[WARN] No GPU detected -- AI tools will use CPU (slower)" + : "[WARN] AI sidecar did not start -- AI tools will use per-request Python (slower)"; console.log( [ `SnapOtter v${APP_VERSION} running on port ${env.PORT}`, diff --git a/packages/ai/python/gpu.py b/packages/ai/python/gpu.py index 8f915e0e..77bbe1d3 100644 --- a/packages/ai/python/gpu.py +++ b/packages/ai/python/gpu.py @@ -60,10 +60,18 @@ def onnx_providers(): """Return (providers, device) tuple. providers: ONNX Runtime execution providers in priority order. - device: "cuda" or "cpu" — reflects which hardware will actually be used. + device: "cuda" or "cpu" -- reflects which hardware will actually be used. """ if gpu_available(): - return (["CUDAExecutionProvider", "CPUExecutionProvider"], "cuda") + try: + import onnxruntime as _ort + available = _ort.get_available_providers() + if "CUDAExecutionProvider" in available: + return (["CUDAExecutionProvider", "CPUExecutionProvider"], "cuda") + emit_info("GPU detected by torch but CUDAExecutionProvider not available in onnxruntime " + "-- install onnxruntime-gpu for GPU acceleration") + except ImportError: + emit_info("onnxruntime not installed, cannot check CUDA provider") emit_info("No GPU detected, processing on CPU") return (["CPUExecutionProvider"], "cpu") diff --git a/packages/ai/src/bridge.ts b/packages/ai/src/bridge.ts index 398ef3e7..51721907 100644 --- a/packages/ai/src/bridge.ts +++ b/packages/ai/src/bridge.ts @@ -216,12 +216,14 @@ function startDispatcher(): ChildProcess | null { dispatcherReady = false; }); - child.on("close", () => { + child.on("close", (code) => { for (const [id, req] of pendingRequests.entries()) { req.reject(new Error("Python dispatcher exited unexpectedly")); pendingRequests.delete(id); } - recordCrash(); + if (code !== 0) { + recordCrash(); + } dispatcher = null; dispatcherReady = false; }); @@ -328,6 +330,44 @@ export function shutdownDispatcher(): void { } } +/** + * 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 }> { + if (dispatcherReady) { + return Promise.resolve({ ready: true, gpu: dispatcherGpuAvailable }); + } + if (dispatcherFailed) { + return Promise.resolve({ ready: false, gpu: false }); + } + + const proc = getDispatcher(); + if (!proc) { + return Promise.resolve({ ready: false, gpu: false }); + } + + return new Promise((resolve) => { + const timer = setTimeout(() => { + clearInterval(poll); + resolve({ ready: false, gpu: false }); + }, timeoutMs); + + const poll = setInterval(() => { + if (dispatcherReady) { + clearTimeout(timer); + clearInterval(poll); + resolve({ ready: true, gpu: dispatcherGpuAvailable }); + } else if (dispatcherFailed) { + clearTimeout(timer); + clearInterval(poll); + resolve({ ready: false, gpu: false }); + } + }, 50); + }); +} + // ── Per-request fallback (original implementation) ────────────────── function runPythonPerRequest( diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 2e3c4c4a..a750aa3a 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -1,6 +1,11 @@ export { removeBackground } from "./background-removal.js"; export type { DispatcherStatus } from "./bridge.js"; -export { getDispatcherStatus, isGpuAvailable, shutdownDispatcher } from "./bridge.js"; +export { + getDispatcherStatus, + initDispatcher, + isGpuAvailable, + shutdownDispatcher, +} from "./bridge.js"; export { colorize } from "./colorization.js"; export type { DetectFacesResult, FaceRegion } from "./face-detection.js"; export { blurFaces, detectFaces } from "./face-detection.js"; diff --git a/tests/unit/ai/bridge.test.ts b/tests/unit/ai/bridge.test.ts index 72b84240..d4bb652a 100644 --- a/tests/unit/ai/bridge.test.ts +++ b/tests/unit/ai/bridge.test.ts @@ -1035,6 +1035,59 @@ describe("bridge - dispatcher lifecycle via runPythonWithProgress", () => { // The important thing is no crash -- the line is collected for pending requests }); + it("does not count exit code 0 as a crash (normal MAX_REQUESTS restart)", async () => { + const mockDispatcher = createMockProcess(); + const mockPerReq = createMockProcess(); + let callCount = 0; + + vi.mocked(spawn).mockImplementation(() => { + callCount++; + if (callCount === 1) return mockDispatcher.process; + return mockPerReq.process; + }); + + const promise = runPythonWithProgress("test.py", []); + + // Dispatcher exits with code 0 (normal MAX_REQUESTS shutdown) + mockDispatcher.emitEvent("close", 0, null); + + await new Promise((r) => setTimeout(r, 10)); + + mockPerReq.stdout.emit("data", Buffer.from('{"ok": true}\n')); + mockPerReq.emitEvent("close", 0, null); + await promise; + + const status = getDispatcherStatus(); + expect(status.consecutiveCrashes).toBe(0); + expect(status.failed).toBe(false); + }); + + it("still counts non-zero exit codes as crashes", async () => { + const mockDispatcher = createMockProcess(); + const mockPerReq = createMockProcess(); + let callCount = 0; + + vi.mocked(spawn).mockImplementation(() => { + callCount++; + if (callCount === 1) return mockDispatcher.process; + return mockPerReq.process; + }); + + const promise = runPythonWithProgress("test.py", []); + + // Dispatcher exits with code 1 (real crash) + mockDispatcher.emitEvent("close", 1, null); + + await new Promise((r) => setTimeout(r, 10)); + + mockPerReq.stdout.emit("data", Buffer.from('{"ok": true}\n')); + mockPerReq.emitEvent("close", 0, null); + await promise; + + const status = getDispatcherStatus(); + expect(status.consecutiveCrashes).toBeGreaterThanOrEqual(1); + }); + it("per-request fallback retries with python3 when venv python fails with ENOENT", async () => { const mockDispatcher = createMockProcess(); const mockVenvPython = createMockProcess(); @@ -1074,3 +1127,94 @@ describe("bridge - dispatcher lifecycle via runPythonWithProgress", () => { expect(callCount).toBe(3); }); }); + +describe("bridge - initDispatcher", () => { + let initDispatcher: typeof import("../../../packages/ai/src/bridge.js").initDispatcher; + let getDispatcherStatus: typeof import("../../../packages/ai/src/bridge.js").getDispatcherStatus; + let shutdownDispatcher: typeof import("../../../packages/ai/src/bridge.js").shutdownDispatcher; + + beforeEach(async () => { + vi.resetModules(); + vi.mocked(spawn).mockReset(); + + const mod = await import("../../../packages/ai/src/bridge.js"); + initDispatcher = mod.initDispatcher; + getDispatcherStatus = mod.getDispatcherStatus; + shutdownDispatcher = mod.shutdownDispatcher; + }); + + afterEach(() => { + shutdownDispatcher(); + vi.restoreAllMocks(); + }); + + it("resolves with ready=true and gpu status after dispatcher emits readiness", async () => { + const mock = createMockProcess(); + vi.mocked(spawn).mockReturnValue(mock.process); + + const promise = initDispatcher(); + + // Dispatcher emits readiness signal + mock.stderr.emit("data", Buffer.from('{"ready": true, "gpu": true}\n')); + + const result = await promise; + expect(result).toEqual({ ready: true, gpu: true }); + expect(getDispatcherStatus().ready).toBe(true); + expect(getDispatcherStatus().gpu).toBe(true); + }); + + it("resolves with ready=false when dispatcher fails with ENOENT", async () => { + const mock = createMockProcess(); + vi.mocked(spawn).mockReturnValue(mock.process); + + const promise = initDispatcher(); + + const err = new Error("spawn ENOENT") as NodeJS.ErrnoException; + err.code = "ENOENT"; + mock.emitEvent("error", err); + + const result = await promise; + expect(result).toEqual({ ready: false, gpu: false }); + }); + + it("resolves with ready=false after timeout when dispatcher never signals ready", async () => { + vi.useFakeTimers(); + const mock = createMockProcess(); + vi.mocked(spawn).mockReturnValue(mock.process); + + const promise = initDispatcher(500); + + vi.advanceTimersByTime(600); + + const result = await promise; + expect(result).toEqual({ ready: false, gpu: false }); + + vi.useRealTimers(); + }); + + it("resolves with gpu=false when dispatcher reports no GPU", async () => { + const mock = createMockProcess(); + vi.mocked(spawn).mockReturnValue(mock.process); + + const promise = initDispatcher(); + + mock.stderr.emit("data", Buffer.from('{"ready": true, "gpu": false}\n')); + + const result = await promise; + expect(result).toEqual({ ready: true, gpu: false }); + }); + + it("is idempotent -- second call returns same result without respawning", async () => { + const mock = createMockProcess(); + vi.mocked(spawn).mockReturnValue(mock.process); + + const promise1 = initDispatcher(); + mock.stderr.emit("data", Buffer.from('{"ready": true, "gpu": true}\n')); + await promise1; + + const result2 = await initDispatcher(); + expect(result2).toEqual({ ready: true, gpu: true }); + // spawn should only have been called once + expect(spawn).toHaveBeenCalledTimes(1); + }); +});