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); + }); + }); });