From f856c26fcb158ff4beb15c78d592fac69e0d6b55 Mon Sep 17 00:00:00 2001 From: SnapOtter Date: Tue, 5 May 2026 21:14:56 +0800 Subject: [PATCH] fix: upscale tool times out on CPU-only systems (NAS/low-power hardware) The upscale function called runPythonWithProgress without a timeout parameter, defaulting to the bridge's 10-minute hard limit. On CPU-only systems like Synology NAS devices, Real-ESRGAN 4x upscaling easily exceeds this for modest images. Additionally, when the timeout fired on the dispatcher path, the Python process was left running and blocked all subsequent AI operations. This fix adds an adaptive timeout based on input megapixels, scale factor, and GPU availability (180s/effective-MP on CPU, 30s/effective-MP on GPU, floor of 10 minutes). It also kills the dispatcher on timeout so subsequent requests can proceed via a fresh restart. Closes #119 --- packages/ai/src/bridge.ts | 5 ++ packages/ai/src/upscaling.ts | 18 ++++- tests/unit/ai/bridge.test.ts | 19 +++++ tests/unit/ai/upscaling.test.ts | 135 +++++++++++++++++++++++++++++++- 4 files changed, 174 insertions(+), 3 deletions(-) diff --git a/packages/ai/src/bridge.ts b/packages/ai/src/bridge.ts index 51721907..c82d603d 100644 --- a/packages/ai/src/bridge.ts +++ b/packages/ai/src/bridge.ts @@ -266,6 +266,11 @@ function dispatcherRun( return new Promise((resolvePromise, rejectPromise) => { const timer = setTimeout(() => { pendingRequests.delete(id); + // Kill the stuck dispatcher so it restarts on the next request instead of + // blocking all subsequent AI operations behind the timed-out script. + if (dispatcher && !dispatcher.killed) { + dispatcher.kill("SIGTERM"); + } rejectPromise(new Error("Python script timed out")); }, timeout); diff --git a/packages/ai/src/upscaling.ts b/packages/ai/src/upscaling.ts index d7ffeac9..b46f79d8 100644 --- a/packages/ai/src/upscaling.ts +++ b/packages/ai/src/upscaling.ts @@ -1,7 +1,12 @@ import { readFile, writeFile } from "node:fs/promises"; import { join } from "node:path"; import sharp from "sharp"; -import { type ProgressCallback, parseStdoutJson, runPythonWithProgress } from "./bridge.js"; +import { + isGpuAvailable, + type ProgressCallback, + parseStdoutJson, + runPythonWithProgress, +} from "./bridge.js"; export interface UpscaleOptions { scale?: number; @@ -31,10 +36,19 @@ export async function upscale( const pngBuffer = await sharp(inputBuffer).png().toBuffer(); await writeFile(inputPath, pngBuffer); + + const meta = await sharp(pngBuffer).metadata(); + const megapixels = ((meta.width ?? 0) * (meta.height ?? 0)) / 1_000_000; + const scale = options.scale ?? 2; + const effectiveMp = megapixels * scale ** 2; + // CPU inference is ~50-100x slower than GPU; be generous for self-hosted NAS hardware + const rateMs = isGpuAvailable() ? 30_000 : 180_000; + const timeout = Math.max(600_000, effectiveMp * rateMs); + const { stdout } = await runPythonWithProgress( "upscale.py", [inputPath, outputPath, JSON.stringify(options)], - { onProgress }, + { onProgress, timeout }, ); const result = parseStdoutJson(stdout); diff --git a/tests/unit/ai/bridge.test.ts b/tests/unit/ai/bridge.test.ts index 3dc4702d..9a46342c 100644 --- a/tests/unit/ai/bridge.test.ts +++ b/tests/unit/ai/bridge.test.ts @@ -1557,6 +1557,25 @@ describe("bridge - dispatcher request timeout", () => { await expect(promise).rejects.toThrow("Python script timed out"); vi.useRealTimers(); }); + + it("kills the dispatcher on timeout so subsequent requests can proceed", async () => { + vi.useFakeTimers(); + const mock = createMockProcess(); + vi.mocked(spawn).mockReturnValue(mock.process); + + const initPromise = initDispatcher(); + mock.stderr.emit("data", Buffer.from('{"ready": true, "gpu": false}\n')); + vi.advanceTimersByTime(100); + await initPromise; + + const promise = runPythonWithProgress("stuck.py", [], { timeout: 2000 }); + + vi.advanceTimersByTime(3000); + + await expect(promise).rejects.toThrow("Python script timed out"); + expect(mock.process.kill).toHaveBeenCalledWith("SIGTERM"); + vi.useRealTimers(); + }); }); // ── Max consecutive crash threshold ───────────────────────────────── diff --git a/tests/unit/ai/upscaling.test.ts b/tests/unit/ai/upscaling.test.ts index a3e0bb7c..b1886a71 100644 --- a/tests/unit/ai/upscaling.test.ts +++ b/tests/unit/ai/upscaling.test.ts @@ -5,6 +5,7 @@ vi.mock("sharp", () => { const mockSharp = vi.fn(() => ({ png: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }), })); return { default: mockSharp }; }); @@ -17,10 +18,15 @@ vi.mock("node:fs/promises", () => ({ vi.mock("../../../packages/ai/src/bridge.js", () => ({ runPythonWithProgress: vi.fn(), parseStdoutJson: vi.fn(), + isGpuAvailable: vi.fn().mockReturnValue(false), })); import sharp from "sharp"; -import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js"; +import { + isGpuAvailable, + parseStdoutJson, + runPythonWithProgress, +} from "../../../packages/ai/src/bridge.js"; import { upscale } from "../../../packages/ai/src/upscaling.js"; const FAKE_INPUT = Buffer.from("fake-small-image"); @@ -46,8 +52,10 @@ beforeEach(() => { ({ png: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }), }) as unknown as ReturnType, ); + vi.mocked(isGpuAvailable).mockReturnValue(false); }); afterEach(() => { @@ -257,4 +265,129 @@ describe("upscale", () => { expect(options.onProgress).toBeUndefined(); }); }); + + describe("timeout calculation", () => { + it("passes a timeout to runPythonWithProgress", async () => { + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR); + + const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; + expect(options.timeout).toBeTypeOf("number"); + expect(options.timeout).toBeGreaterThan(0); + }); + + it("uses minimum timeout of 600s for small images", async () => { + vi.mocked(sharp).mockImplementation( + () => + ({ + png: vi.fn().mockReturnThis(), + toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 100, height: 100 }), + }) as unknown as ReturnType, + ); + + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 2 }); + + const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; + expect(options.timeout).toBe(600_000); + }); + + it("scales timeout with image megapixels", async () => { + vi.mocked(sharp).mockImplementation( + () => + ({ + png: vi.fn().mockReturnThis(), + toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 2000, height: 1500 }), + }) as unknown as ReturnType, + ); + + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 2 }); + const timeout2x = vi.mocked(runPythonWithProgress).mock.calls[0][2].timeout!; + + vi.mocked(runPythonWithProgress).mockClear(); + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 4 }); + const timeout4x = vi.mocked(runPythonWithProgress).mock.calls[0][2].timeout!; + + expect(timeout4x).toBeGreaterThan(timeout2x); + }); + + it("scales timeout with scale factor squared", async () => { + vi.mocked(sharp).mockImplementation( + () => + ({ + png: vi.fn().mockReturnThis(), + toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 1000, height: 1000 }), + }) as unknown as ReturnType, + ); + + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 2 }); + const timeout2x = vi.mocked(runPythonWithProgress).mock.calls[0][2].timeout!; + + vi.mocked(runPythonWithProgress).mockClear(); + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 4 }); + const timeout4x = vi.mocked(runPythonWithProgress).mock.calls[0][2].timeout!; + + // 4x scale has 4x the effective megapixels vs 2x scale (16/4 = 4) + expect(timeout4x / timeout2x).toBe(4); + }); + + it("uses shorter timeout when GPU is available", async () => { + vi.mocked(sharp).mockImplementation( + () => + ({ + png: vi.fn().mockReturnThis(), + toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 2000, height: 1500 }), + }) as unknown as ReturnType, + ); + + vi.mocked(isGpuAvailable).mockReturnValue(false); + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 4 }); + const cpuTimeout = vi.mocked(runPythonWithProgress).mock.calls[0][2].timeout!; + + vi.mocked(runPythonWithProgress).mockClear(); + vi.mocked(isGpuAvailable).mockReturnValue(true); + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 4 }); + const gpuTimeout = vi.mocked(runPythonWithProgress).mock.calls[0][2].timeout!; + + expect(gpuTimeout).toBeLessThan(cpuTimeout); + }); + + it("provides generous timeout for CPU upscale at high scale", async () => { + // Simulates user scenario: ~2MP image at 4x on CPU (Synology NAS) + vi.mocked(sharp).mockImplementation( + () => + ({ + png: vi.fn().mockReturnThis(), + toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 1600, height: 1200 }), + }) as unknown as ReturnType, + ); + vi.mocked(isGpuAvailable).mockReturnValue(false); + + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR, { scale: 4 }); + const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; + + // ~1.92MP * 16 (4^2) * 180_000ms = ~55 minutes. Must be well above 10 min. + expect(options.timeout).toBeGreaterThan(10 * 60 * 1000); + }); + + it("defaults scale to 2 when not provided", async () => { + vi.mocked(sharp).mockImplementation( + () => + ({ + png: vi.fn().mockReturnThis(), + toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), + metadata: vi.fn().mockResolvedValue({ width: 2000, height: 2000 }), + }) as unknown as ReturnType, + ); + + await upscale(FAKE_INPUT, FAKE_OUTPUT_DIR); + const options = vi.mocked(runPythonWithProgress).mock.calls[0][2]; + + // 4MP * 4 (2^2) * 180_000 = 2_880_000ms + expect(options.timeout).toBe(2_880_000); + }); + }); });