Add local compute concurrency limiter with tests
This commit is contained in:
parent
4497f610c0
commit
3aad8a51e4
3 changed files with 143 additions and 33 deletions
40
src/lib/server/compute/concurrency-limiter.ts
Normal file
40
src/lib/server/compute/concurrency-limiter.ts
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
export class ConcurrencyLimiter {
|
||||
private readonly maxInFlight: number;
|
||||
private inFlight = 0;
|
||||
private readonly queue: Array<() => void> = [];
|
||||
|
||||
constructor(limit: number) {
|
||||
this.maxInFlight = Number.isFinite(limit) && limit >= 1 ? Math.floor(limit) : 1;
|
||||
}
|
||||
|
||||
async run<T>(fn: () => Promise<T>): Promise<T> {
|
||||
await this.acquire();
|
||||
try {
|
||||
return await fn();
|
||||
} finally {
|
||||
this.release();
|
||||
}
|
||||
}
|
||||
|
||||
private acquire(): Promise<void> {
|
||||
if (this.inFlight < this.maxInFlight) {
|
||||
this.inFlight += 1;
|
||||
return Promise.resolve();
|
||||
}
|
||||
return new Promise((resolve) => {
|
||||
this.queue.push(() => {
|
||||
this.inFlight += 1;
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
private release(): void {
|
||||
this.inFlight = Math.max(0, this.inFlight - 1);
|
||||
const next = this.queue.shift();
|
||||
if (next) next();
|
||||
}
|
||||
}
|
||||
|
||||
export const LOCAL_PDF_LIMITER = new ConcurrencyLimiter(2);
|
||||
export const LOCAL_WHISPER_LIMITER = new ConcurrencyLimiter(1);
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import type { ComputeBackend, PdfLayoutInput, WhisperAlignInput, WhisperAlignResult } from '@/lib/server/compute/types';
|
||||
import { LOCAL_PDF_LIMITER, LOCAL_WHISPER_LIMITER } from '@/lib/server/compute/concurrency-limiter';
|
||||
import { getDocumentBlob } from '@/lib/server/documents/blobstore';
|
||||
import { getTtsSegmentAudioObject } from '@/lib/server/tts/segments-blobstore';
|
||||
import {
|
||||
|
|
@ -10,42 +11,46 @@ export class LocalComputeBackend implements ComputeBackend {
|
|||
readonly mode = 'local' as const;
|
||||
|
||||
async alignWords(input: WhisperAlignInput): Promise<WhisperAlignResult> {
|
||||
let audioBuffer = input.audioBuffer ?? null;
|
||||
if (!audioBuffer && input.audioObjectKey) {
|
||||
const bytes = new Uint8Array(await getTtsSegmentAudioObject(input.audioObjectKey));
|
||||
audioBuffer = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
|
||||
}
|
||||
if (!audioBuffer) {
|
||||
throw new Error('Local compute alignment requires audioBuffer or audioObjectKey');
|
||||
}
|
||||
return runWhisperAlignmentFromAudioBuffer({
|
||||
audioBuffer,
|
||||
text: input.text,
|
||||
cacheKey: input.cacheKey,
|
||||
lang: input.lang,
|
||||
return LOCAL_WHISPER_LIMITER.run(async () => {
|
||||
let audioBuffer = input.audioBuffer ?? null;
|
||||
if (!audioBuffer && input.audioObjectKey) {
|
||||
const bytes = new Uint8Array(await getTtsSegmentAudioObject(input.audioObjectKey));
|
||||
audioBuffer = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
|
||||
}
|
||||
if (!audioBuffer) {
|
||||
throw new Error('Local compute alignment requires audioBuffer or audioObjectKey');
|
||||
}
|
||||
return runWhisperAlignmentFromAudioBuffer({
|
||||
audioBuffer,
|
||||
text: input.text,
|
||||
cacheKey: input.cacheKey,
|
||||
lang: input.lang,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async parsePdfLayout(input: PdfLayoutInput) {
|
||||
let pdfBytes = input.pdfBytes ?? null;
|
||||
if (!pdfBytes && input.documentId && typeof input.namespace !== 'undefined') {
|
||||
const bytes = new Uint8Array(await getDocumentBlob(input.documentId, input.namespace));
|
||||
pdfBytes = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
|
||||
}
|
||||
if (!pdfBytes) {
|
||||
throw new Error('Local compute PDF layout requires pdfBytes or (documentId + namespace)');
|
||||
}
|
||||
return {
|
||||
parsed: (await runPdfLayoutFromPdfBuffer({
|
||||
documentId: input.documentId,
|
||||
pdfBytes,
|
||||
onPageParsed: (page) => input.onProgress?.({
|
||||
totalPages: page.totalPages,
|
||||
pagesParsed: page.pageNumber,
|
||||
currentPage: page.pageNumber,
|
||||
phase: 'infer',
|
||||
}),
|
||||
})).parsed,
|
||||
};
|
||||
return LOCAL_PDF_LIMITER.run(async () => {
|
||||
let pdfBytes = input.pdfBytes ?? null;
|
||||
if (!pdfBytes && input.documentId && typeof input.namespace !== 'undefined') {
|
||||
const bytes = new Uint8Array(await getDocumentBlob(input.documentId, input.namespace));
|
||||
pdfBytes = bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength);
|
||||
}
|
||||
if (!pdfBytes) {
|
||||
throw new Error('Local compute PDF layout requires pdfBytes or (documentId + namespace)');
|
||||
}
|
||||
return {
|
||||
parsed: (await runPdfLayoutFromPdfBuffer({
|
||||
documentId: input.documentId,
|
||||
pdfBytes,
|
||||
onPageParsed: (page) => input.onProgress?.({
|
||||
totalPages: page.totalPages,
|
||||
pagesParsed: page.pageNumber,
|
||||
currentPage: page.pageNumber,
|
||||
phase: 'infer',
|
||||
}),
|
||||
})).parsed,
|
||||
};
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
65
tests/unit/compute-concurrency-limiter.spec.ts
Normal file
65
tests/unit/compute-concurrency-limiter.spec.ts
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
import { expect, test } from '@playwright/test';
|
||||
import { ConcurrencyLimiter } from '../../src/lib/server/compute/concurrency-limiter';
|
||||
|
||||
function sleep(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
test.describe('compute-concurrency-limiter', () => {
|
||||
test('caps active jobs at configured limit', async () => {
|
||||
const limiter = new ConcurrencyLimiter(2);
|
||||
let inFlight = 0;
|
||||
let maxInFlightSeen = 0;
|
||||
|
||||
const jobs = Array.from({ length: 6 }, (_, i) => limiter.run(async () => {
|
||||
inFlight += 1;
|
||||
maxInFlightSeen = Math.max(maxInFlightSeen, inFlight);
|
||||
await sleep(20 + (i % 2) * 5);
|
||||
inFlight -= 1;
|
||||
return i;
|
||||
}));
|
||||
|
||||
const result = await Promise.all(jobs);
|
||||
expect(result).toEqual([0, 1, 2, 3, 4, 5]);
|
||||
expect(maxInFlightSeen).toBe(2);
|
||||
});
|
||||
|
||||
test('queues in FIFO order when saturated', async () => {
|
||||
const limiter = new ConcurrencyLimiter(1);
|
||||
const starts: number[] = [];
|
||||
|
||||
const one = limiter.run(async () => {
|
||||
starts.push(1);
|
||||
await sleep(25);
|
||||
});
|
||||
const two = limiter.run(async () => {
|
||||
starts.push(2);
|
||||
await sleep(5);
|
||||
});
|
||||
const three = limiter.run(async () => {
|
||||
starts.push(3);
|
||||
await sleep(1);
|
||||
});
|
||||
|
||||
await Promise.all([one, two, three]);
|
||||
expect(starts).toEqual([1, 2, 3]);
|
||||
});
|
||||
|
||||
test('releases slot after failure', async () => {
|
||||
const limiter = new ConcurrencyLimiter(1);
|
||||
let startedSecond = false;
|
||||
|
||||
const first = limiter.run(async () => {
|
||||
await sleep(10);
|
||||
throw new Error('boom');
|
||||
});
|
||||
const second = limiter.run(async () => {
|
||||
startedSecond = true;
|
||||
return 'ok';
|
||||
});
|
||||
|
||||
await expect(first).rejects.toThrow('boom');
|
||||
await expect(second).resolves.toBe('ok');
|
||||
expect(startedSecond).toBeTruthy();
|
||||
});
|
||||
});
|
||||
Loading…
Reference in a new issue