diff --git a/compute/core/src/pdf-layout/parsePdf.ts b/compute/core/src/pdf-layout/parsePdf.ts index 0207b8f..8321b2e 100644 --- a/compute/core/src/pdf-layout/parsePdf.ts +++ b/compute/core/src/pdf-layout/parsePdf.ts @@ -1,4 +1,5 @@ import path from 'path'; +import { createRequire } from 'node:module'; import type { TextItem } from 'pdfjs-dist/types/src/display/api'; import type { ParsedPdfDocument, ParsedPdfPage } from '../types/parsed-pdf'; import type { PdfTextItem } from './types'; @@ -19,6 +20,13 @@ interface ParsePdfInput { } const LAYOUT_RENDER_SCALE = 1.5; +const require = createRequire(import.meta.url); + +function resolvePdfjsStandardFontDataUrl(): string { + const pdfjsPackageJson = require.resolve('pdfjs-dist/package.json'); + const standardFontDir = path.join(path.dirname(pdfjsPackageJson), 'standard_fonts'); + return `${standardFontDir.replace(/\/?$/, '/')}`; +} export function normalizeTextItemsForLayout(items: TextItem[], pageHeight: number): PdfTextItem[] { return items @@ -68,8 +76,7 @@ export async function parsePdf(input: ParsePdfInput): Promise pdfjs.GlobalWorkerOptions.workerSrc = 'pdfjs-dist/legacy/build/pdf.worker.mjs'; pdfjs.GlobalWorkerOptions.workerPort = null; } - const standardFontDir = path.join(process.cwd(), 'node_modules', 'pdfjs-dist', 'standard_fonts'); - const standardFontDataUrl = `${standardFontDir.replace(/\/?$/, '/')}`; + const standardFontDataUrl = resolvePdfjsStandardFontDataUrl(); const loadingTask = pdfjs.getDocument({ data: pdfBytesForText, diff --git a/compute/core/src/pdf-layout/renderPage.ts b/compute/core/src/pdf-layout/renderPage.ts index e68fb66..cb86d70 100644 --- a/compute/core/src/pdf-layout/renderPage.ts +++ b/compute/core/src/pdf-layout/renderPage.ts @@ -1,4 +1,5 @@ import path from 'path'; +import { createRequire } from 'node:module'; import type { Canvas } from '@napi-rs/canvas'; type CanvasRuntime = { @@ -8,6 +9,13 @@ type CanvasRuntime = { }; let canvasRuntimePromise: Promise | null = null; +const require = createRequire(import.meta.url); + +function resolvePdfjsStandardFontDataUrl(): string { + const pdfjsPackageJson = require.resolve('pdfjs-dist/package.json'); + const standardFontDir = path.join(path.dirname(pdfjsPackageJson), 'standard_fonts'); + return `${standardFontDir.replace(/\/?$/, '/')}`; +} async function loadCanvasRuntime(): Promise { if (!canvasRuntimePromise) { @@ -112,8 +120,7 @@ export async function renderPage({ pdfjs.GlobalWorkerOptions.workerPort = null; } - const standardFontDir = path.join(process.cwd(), 'node_modules', 'pdfjs-dist', 'standard_fonts'); - const standardFontDataUrl = `${standardFontDir.replace(/\/?$/, '/')}`; + const standardFontDataUrl = resolvePdfjsStandardFontDataUrl(); const loadingTask = pdfjs.getDocument({ data: isolatedBytes, diff --git a/compute/core/src/pdf-layout/runLayoutModel.ts b/compute/core/src/pdf-layout/runLayoutModel.ts index f7b80eb..24a7ac8 100644 --- a/compute/core/src/pdf-layout/runLayoutModel.ts +++ b/compute/core/src/pdf-layout/runLayoutModel.ts @@ -2,6 +2,7 @@ import * as ort from 'onnxruntime-node'; import { readFile } from 'fs/promises'; import type { LayoutRegion, PdfTextItem } from './types'; import { ensureModel, MODEL_CONFIG_PATH, MODEL_PREPROCESSOR_PATH } from './ensureModel'; +import { getOnnxThreadsPerJob } from '../runtime/cpu-budget'; interface RunLayoutInput { pageWidth: number; @@ -115,9 +116,16 @@ async function getSession(): Promise { if (!sessionPromise) { sessionPromise = (async () => { const modelPath = await ensureModel(); - return ort.InferenceSession.create(modelPath, { + const onnxThreadsPerJob = getOnnxThreadsPerJob(); + const stableSessionOptions: ort.InferenceSession.SessionOptions = { executionProviders: ['cpu'], graphOptimizationLevel: 'all', + intraOpNumThreads: onnxThreadsPerJob, + interOpNumThreads: 1, + executionMode: 'sequential', + }; + return ort.InferenceSession.create(modelPath, { + ...stableSessionOptions, }); })(); } diff --git a/compute/core/src/runtime/cpu-budget.ts b/compute/core/src/runtime/cpu-budget.ts new file mode 100644 index 0000000..50d161b --- /dev/null +++ b/compute/core/src/runtime/cpu-budget.ts @@ -0,0 +1,28 @@ +import os from 'node:os'; + +function readPositiveInt(name: string, fallback: number): number { + const raw = process.env[name]?.trim(); + if (!raw) return fallback; + const parsed = Number(raw); + if (!Number.isFinite(parsed) || parsed <= 0) return fallback; + return Math.floor(parsed); +} + +export function getComputeJobConcurrency(): number { + return readPositiveInt('COMPUTE_JOB_CONCURRENCY', 1); +} + +export function getAvailableCpuCores(): number { + if (typeof os.availableParallelism === 'function') { + const value = os.availableParallelism(); + if (Number.isFinite(value) && value >= 1) return Math.floor(value); + } + const fallback = os.cpus().length; + return Number.isFinite(fallback) && fallback >= 1 ? Math.floor(fallback) : 1; +} + +export function getOnnxThreadsPerJob(): number { + const concurrency = getComputeJobConcurrency(); + const usableCores = Math.max(1, getAvailableCpuCores() - 1); + return Math.max(1, Math.floor(usableCores / concurrency)); +} diff --git a/compute/core/src/whisper/alignment.ts b/compute/core/src/whisper/alignment.ts index 7f6607f..e3b3085 100644 --- a/compute/core/src/whisper/alignment.ts +++ b/compute/core/src/whisper/alignment.ts @@ -9,6 +9,7 @@ import { Tokenizer } from '@huggingface/tokenizers'; import JSZip from 'jszip'; import type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment } from '../types/tts'; import { getFFmpegPath } from '../runtime/ffmpeg'; +import { getOnnxThreadsPerJob } from '../runtime/cpu-budget'; import { mapWordsToSentenceOffsets, type WhisperWord, @@ -575,10 +576,11 @@ async function getRuntime(): Promise { const defaultLanguageToken = Number(defaultLanguageFromForced ?? tokenizer.token_to_id('<|en|>') ?? 50259); const transcribeToken = Number(transcribeFromForced ?? tokenizer.token_to_id('<|transcribe|>') ?? 50359); + const onnxThreadsPerJob = getOnnxThreadsPerJob(); const stableSessionOptions: ort.InferenceSession.SessionOptions = { executionProviders: ['cpu'], graphOptimizationLevel: 'disabled', - intraOpNumThreads: 1, + intraOpNumThreads: onnxThreadsPerJob, interOpNumThreads: 1, executionMode: 'sequential', enableCpuMemArena: false, diff --git a/compute/worker/.env.example b/compute/worker/.env.example index c8b889a..7f9d345 100644 --- a/compute/worker/.env.example +++ b/compute/worker/.env.example @@ -30,10 +30,9 @@ S3_FORCE_PATH_STYLE=true # Queue + execution tuning # COMPUTE_PREWARM_MODELS=true -# COMPUTE_WHISPER_CONCURRENCY=1 -# COMPUTE_PDF_CONCURRENCY=2 +# COMPUTE_JOB_CONCURRENCY=1 # COMPUTE_WHISPER_TIMEOUT_MS=30000 -# COMPUTE_PDF_TIMEOUT_MS=90000 +# COMPUTE_PDF_TIMEOUT_MS=300000 # COMPUTE_PDF_JOB_ATTEMPTS=2 # COMPUTE_JOBS_STREAM_MAX_BYTES=268435456 # COMPUTE_JOB_STATES_MAX_BYTES=67108864 diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index 6800068..7c2821e 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -1,4 +1,5 @@ import { createHash } from 'node:crypto'; +import os from 'node:os'; import Fastify, { type FastifyRequest } from 'fastify'; import { z } from 'zod'; import { @@ -112,6 +113,20 @@ function parseBoolEnv(name: string, fallback: boolean): boolean { return normalized === '1' || normalized === 'true' || normalized === 'yes' || normalized === 'on'; } +function getAvailableCpuCores(): number { + if (typeof os.availableParallelism === 'function') { + const value = os.availableParallelism(); + if (Number.isFinite(value) && value >= 1) return Math.floor(value); + } + const fallback = os.cpus().length; + return Number.isFinite(fallback) && fallback >= 1 ? Math.floor(fallback) : 1; +} + +function getOnnxThreadsPerJob(jobConcurrency: number): number { + const usableCores = Math.max(1, getAvailableCpuCores() - 1); + return Math.max(1, Math.floor(usableCores / Math.max(1, jobConcurrency))); +} + function buildLoggerConfig(): boolean | Record { const format = (process.env.COMPUTE_LOG_FORMAT?.trim().toLowerCase() || 'pretty'); if (format === 'json') return true; @@ -311,6 +326,36 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } +class ConcurrencyGate { + 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 acquire(): Promise { + if (this.inFlight < this.maxInFlight) { + this.inFlight += 1; + return; + } + + await new Promise((resolve) => { + this.queue.push(() => { + this.inFlight += 1; + resolve(); + }); + }); + } + + release(): void { + this.inFlight = Math.max(0, this.inFlight - 1); + const next = this.queue.shift(); + if (next) next(); + } +} + function isTerminalStatus(status: WorkerJobState): boolean { return status === 'succeeded' || status === 'failed'; } @@ -428,10 +473,9 @@ async function main(): Promise { const workerToken = requireEnv('COMPUTE_WORKER_TOKEN'); const natsUrl = requireEnv('NATS_URL'); - const whisperConcurrency = readIntEnv('COMPUTE_WHISPER_CONCURRENCY', 1); - const pdfConcurrency = readIntEnv('COMPUTE_PDF_CONCURRENCY', 2); + const jobConcurrency = readIntEnv('COMPUTE_JOB_CONCURRENCY', 1); const whisperTimeoutMs = readIntEnv('COMPUTE_WHISPER_TIMEOUT_MS', 30_000); - const pdfTimeoutMs = readIntEnv('COMPUTE_PDF_TIMEOUT_MS', 90_000); + const pdfTimeoutMs = readIntEnv('COMPUTE_PDF_TIMEOUT_MS', 5 * 60_000); const pdfAttempts = readIntEnv('COMPUTE_PDF_JOB_ATTEMPTS', 2); const prewarmModels = parseBoolEnv('COMPUTE_PREWARM_MODELS', true); const jobsStreamMaxBytes = readIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); @@ -507,6 +551,17 @@ async function main(): Promise { } const app = Fastify({ logger: buildLoggerConfig() }); + app.log.info({ + jobConcurrency, + whisperTimeoutMs, + pdfTimeoutMs, + pdfAttempts, + opStaleMs, + availableCpuCores: getAvailableCpuCores(), + onnxThreadsPerJob: getOnnxThreadsPerJob(jobConcurrency), + natsApiTimeoutMs: NATS_API_TIMEOUT_MS, + pdfLayoutHardCapMs: PDF_LAYOUT_HARD_CAP_MS, + }, 'compute runtime config'); const opIndexCodec = createJsonCodec(); const opStateCodec = createJsonCodec>(); @@ -806,6 +861,7 @@ async function main(): Promise { const whisperConsumer = await js.consumers.get(JOBS_STREAM_NAME, WHISPER_CONSUMER_NAME); const layoutConsumer = await js.consumers.get(JOBS_STREAM_NAME, LAYOUT_CONSUMER_NAME); + const jobGate = new ConcurrencyGate(jobConcurrency); let stopping = false; @@ -1125,29 +1181,39 @@ async function main(): Promise { workerLabel: string; }): Promise { while (!stopping) { - let msg: JsMsg | null = null; - try { - msg = await input.consumer.next({ expires: PULL_EXPIRES_MS }); - } catch (error) { - if (stopping) return; - app.log.error({ error: toErrorMessage(error), worker: input.workerLabel }, 'worker pull failed'); - await sleep(LOOP_ERROR_BACKOFF_MS); - continue; + await jobGate.acquire(); + if (stopping) { + jobGate.release(); + return; } - if (!msg) continue; - await processMessage({ - msg, - codec: input.codec, - run: input.run, - workerLabel: input.workerLabel, - }); + let msg: JsMsg | null = null; + try { + try { + msg = await input.consumer.next({ expires: PULL_EXPIRES_MS }); + } catch (error) { + if (stopping) return; + app.log.error({ error: toErrorMessage(error), worker: input.workerLabel }, 'worker pull failed'); + await sleep(LOOP_ERROR_BACKOFF_MS); + continue; + } + + if (!msg) continue; + await processMessage({ + msg, + codec: input.codec, + run: input.run, + workerLabel: input.workerLabel, + }); + } finally { + jobGate.release(); + } } } const workerLoops: Promise[] = []; - for (let i = 0; i < whisperConcurrency; i += 1) { + for (let i = 0; i < jobConcurrency; i += 1) { workerLoops.push(createWorkerLoop({ consumer: whisperConsumer, codec: whisperJobCodec, @@ -1156,7 +1222,7 @@ async function main(): Promise { })); } - for (let i = 0; i < pdfConcurrency; i += 1) { + for (let i = 0; i < jobConcurrency; i += 1) { workerLoops.push(createWorkerLoop({ consumer: layoutConsumer, codec: layoutJobCodec, diff --git a/docs-site/docs/deploy/compute-worker.md b/docs-site/docs/deploy/compute-worker.md index fef4995..302ce93 100644 --- a/docs-site/docs/deploy/compute-worker.md +++ b/docs-site/docs/deploy/compute-worker.md @@ -48,10 +48,9 @@ Common optional: Advanced tuning (usually leave unset unless you need overrides): - `COMPUTE_PREWARM_MODELS=true` -- `COMPUTE_WHISPER_CONCURRENCY=1` -- `COMPUTE_PDF_CONCURRENCY=2` +- `COMPUTE_JOB_CONCURRENCY=1` (shared total compute jobs across whisper + PDF) - `COMPUTE_WHISPER_TIMEOUT_MS=30000` -- `COMPUTE_PDF_TIMEOUT_MS=90000` +- `COMPUTE_PDF_TIMEOUT_MS=300000` - `COMPUTE_PDF_JOB_ATTEMPTS=2` (PDF layout retry attempts) - `COMPUTE_JOBS_STREAM_MAX_BYTES=268435456` (256MB JetStream jobs stream cap) - `COMPUTE_JOB_STATES_MAX_BYTES=67108864` (64MB JetStream KV bucket cap) @@ -121,10 +120,9 @@ COMPUTE_WORKER_HOST=0.0.0.0 COMPUTE_WORKER_TOKEN= # Optional advanced tuning overrides (defaults shown): # COMPUTE_PREWARM_MODELS=true -# COMPUTE_WHISPER_CONCURRENCY=1 -# COMPUTE_PDF_CONCURRENCY=2 +# COMPUTE_JOB_CONCURRENCY=1 # COMPUTE_WHISPER_TIMEOUT_MS=30000 -# COMPUTE_PDF_TIMEOUT_MS=90000 +# COMPUTE_PDF_TIMEOUT_MS=300000 # COMPUTE_PDF_JOB_ATTEMPTS=2 # COMPUTE_JOBS_STREAM_MAX_BYTES=268435456 # COMPUTE_JOB_STATES_MAX_BYTES=67108864 diff --git a/src/lib/server/compute/concurrency-limiter.ts b/src/lib/server/compute/concurrency-limiter.ts index 08102db..e733cba 100644 --- a/src/lib/server/compute/concurrency-limiter.ts +++ b/src/lib/server/compute/concurrency-limiter.ts @@ -1,3 +1,5 @@ +import { getComputeJobConcurrency } from '@/lib/server/compute/cpu-budget'; + export class ConcurrencyLimiter { private readonly maxInFlight: number; private inFlight = 0; @@ -36,5 +38,4 @@ export class ConcurrencyLimiter { } } -export const LOCAL_PDF_LIMITER = new ConcurrencyLimiter(2); -export const LOCAL_WHISPER_LIMITER = new ConcurrencyLimiter(1); +export const LOCAL_COMPUTE_LIMITER = new ConcurrencyLimiter(getComputeJobConcurrency()); diff --git a/src/lib/server/compute/cpu-budget.ts b/src/lib/server/compute/cpu-budget.ts new file mode 100644 index 0000000..50d161b --- /dev/null +++ b/src/lib/server/compute/cpu-budget.ts @@ -0,0 +1,28 @@ +import os from 'node:os'; + +function readPositiveInt(name: string, fallback: number): number { + const raw = process.env[name]?.trim(); + if (!raw) return fallback; + const parsed = Number(raw); + if (!Number.isFinite(parsed) || parsed <= 0) return fallback; + return Math.floor(parsed); +} + +export function getComputeJobConcurrency(): number { + return readPositiveInt('COMPUTE_JOB_CONCURRENCY', 1); +} + +export function getAvailableCpuCores(): number { + if (typeof os.availableParallelism === 'function') { + const value = os.availableParallelism(); + if (Number.isFinite(value) && value >= 1) return Math.floor(value); + } + const fallback = os.cpus().length; + return Number.isFinite(fallback) && fallback >= 1 ? Math.floor(fallback) : 1; +} + +export function getOnnxThreadsPerJob(): number { + const concurrency = getComputeJobConcurrency(); + const usableCores = Math.max(1, getAvailableCpuCores() - 1); + return Math.max(1, Math.floor(usableCores / concurrency)); +} diff --git a/src/lib/server/compute/local.ts b/src/lib/server/compute/local.ts index 4d3ee0c..c5f88e1 100644 --- a/src/lib/server/compute/local.ts +++ b/src/lib/server/compute/local.ts @@ -1,5 +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 { LOCAL_COMPUTE_LIMITER } from '@/lib/server/compute/concurrency-limiter'; import { getDocumentBlob } from '@/lib/server/documents/blobstore'; import { getTtsSegmentAudioObject } from '@/lib/server/tts/segments-blobstore'; import { @@ -11,7 +11,7 @@ export class LocalComputeBackend implements ComputeBackend { readonly mode = 'local' as const; async alignWords(input: WhisperAlignInput): Promise { - return LOCAL_WHISPER_LIMITER.run(async () => { + return LOCAL_COMPUTE_LIMITER.run(async () => { let audioBuffer = input.audioBuffer ?? null; if (!audioBuffer && input.audioObjectKey) { const bytes = new Uint8Array(await getTtsSegmentAudioObject(input.audioObjectKey)); @@ -30,7 +30,7 @@ export class LocalComputeBackend implements ComputeBackend { } async parsePdfLayout(input: PdfLayoutInput) { - return LOCAL_PDF_LIMITER.run(async () => { + return LOCAL_COMPUTE_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)); diff --git a/src/lib/server/pdf-layout/parsePdf.ts b/src/lib/server/pdf-layout/parsePdf.ts index cbb8f29..172d9bd 100644 --- a/src/lib/server/pdf-layout/parsePdf.ts +++ b/src/lib/server/pdf-layout/parsePdf.ts @@ -1,4 +1,5 @@ import path from 'path'; +import { createRequire } from 'node:module'; import type { TextItem } from 'pdfjs-dist/types/src/display/api'; import type { ParsedPdfDocument, ParsedPdfPage } from '@/types/parsed-pdf'; import type { PdfTextItem } from '@/lib/server/pdf-layout/types'; @@ -14,6 +15,13 @@ interface ParsePdfInput { } const LAYOUT_RENDER_SCALE = 1.5; +const require = createRequire(import.meta.url); + +function resolvePdfjsStandardFontDataUrl(): string { + const pdfjsPackageJson = require.resolve('pdfjs-dist/package.json'); + const standardFontDir = path.join(path.dirname(pdfjsPackageJson), 'standard_fonts'); + return `${standardFontDir.replace(/\/?$/, '/')}`; +} export function normalizeTextItemsForLayout(items: TextItem[], pageHeight: number): PdfTextItem[] { return items @@ -63,8 +71,7 @@ export async function parsePdf(input: ParsePdfInput): Promise pdfjs.GlobalWorkerOptions.workerSrc = 'pdfjs-dist/legacy/build/pdf.worker.mjs'; pdfjs.GlobalWorkerOptions.workerPort = null; } - const standardFontDir = path.join(process.cwd(), 'node_modules', 'pdfjs-dist', 'standard_fonts'); - const standardFontDataUrl = `${standardFontDir.replace(/\/?$/, '/')}`; + const standardFontDataUrl = resolvePdfjsStandardFontDataUrl(); const loadingTask = pdfjs.getDocument({ data: pdfBytesForText, diff --git a/src/lib/server/pdf-layout/renderPage.ts b/src/lib/server/pdf-layout/renderPage.ts index e68fb66..cb86d70 100644 --- a/src/lib/server/pdf-layout/renderPage.ts +++ b/src/lib/server/pdf-layout/renderPage.ts @@ -1,4 +1,5 @@ import path from 'path'; +import { createRequire } from 'node:module'; import type { Canvas } from '@napi-rs/canvas'; type CanvasRuntime = { @@ -8,6 +9,13 @@ type CanvasRuntime = { }; let canvasRuntimePromise: Promise | null = null; +const require = createRequire(import.meta.url); + +function resolvePdfjsStandardFontDataUrl(): string { + const pdfjsPackageJson = require.resolve('pdfjs-dist/package.json'); + const standardFontDir = path.join(path.dirname(pdfjsPackageJson), 'standard_fonts'); + return `${standardFontDir.replace(/\/?$/, '/')}`; +} async function loadCanvasRuntime(): Promise { if (!canvasRuntimePromise) { @@ -112,8 +120,7 @@ export async function renderPage({ pdfjs.GlobalWorkerOptions.workerPort = null; } - const standardFontDir = path.join(process.cwd(), 'node_modules', 'pdfjs-dist', 'standard_fonts'); - const standardFontDataUrl = `${standardFontDir.replace(/\/?$/, '/')}`; + const standardFontDataUrl = resolvePdfjsStandardFontDataUrl(); const loadingTask = pdfjs.getDocument({ data: isolatedBytes, diff --git a/src/lib/server/pdf-layout/runLayoutModel.ts b/src/lib/server/pdf-layout/runLayoutModel.ts index a03f21c..3463921 100644 --- a/src/lib/server/pdf-layout/runLayoutModel.ts +++ b/src/lib/server/pdf-layout/runLayoutModel.ts @@ -2,6 +2,7 @@ import * as ort from 'onnxruntime-node'; import { readFile } from 'fs/promises'; import type { LayoutRegion, PdfTextItem } from '@/lib/server/pdf-layout/types'; import { ensureModel, MODEL_CONFIG_PATH, MODEL_PREPROCESSOR_PATH } from '@/lib/server/pdf-layout/ensureModel'; +import { getOnnxThreadsPerJob } from '@/lib/server/compute/cpu-budget'; interface RunLayoutInput { pageWidth: number; @@ -115,9 +116,16 @@ async function getSession(): Promise { if (!sessionPromise) { sessionPromise = (async () => { const modelPath = await ensureModel(); - return ort.InferenceSession.create(modelPath, { + const onnxThreadsPerJob = getOnnxThreadsPerJob(); + const stableSessionOptions: ort.InferenceSession.SessionOptions = { executionProviders: ['cpu'], graphOptimizationLevel: 'all', + intraOpNumThreads: onnxThreadsPerJob, + interOpNumThreads: 1, + executionMode: 'sequential', + }; + return ort.InferenceSession.create(modelPath, { + ...stableSessionOptions, }); })(); } diff --git a/src/lib/server/whisper/alignment.ts b/src/lib/server/whisper/alignment.ts index d0d8f14..f225ea2 100644 --- a/src/lib/server/whisper/alignment.ts +++ b/src/lib/server/whisper/alignment.ts @@ -8,6 +8,7 @@ import { Tokenizer } from '@huggingface/tokenizers'; import JSZip from 'jszip'; import type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment } from '@/types/tts'; import { getFFmpegPath } from '@/lib/server/audiobooks/ffmpeg-bin'; +import { getOnnxThreadsPerJob } from '@/lib/server/compute/cpu-budget'; import { mapWordsToSentenceOffsets, type WhisperWord, @@ -573,10 +574,11 @@ async function getRuntime(): Promise { const defaultLanguageToken = Number(defaultLanguageFromForced ?? tokenizer.token_to_id('<|en|>') ?? 50259); const transcribeToken = Number(transcribeFromForced ?? tokenizer.token_to_id('<|transcribe|>') ?? 50359); + const onnxThreadsPerJob = getOnnxThreadsPerJob(); const stableSessionOptions: ort.InferenceSession.SessionOptions = { executionProviders: ['cpu'], graphOptimizationLevel: 'disabled', - intraOpNumThreads: 1, + intraOpNumThreads: onnxThreadsPerJob, interOpNumThreads: 1, executionMode: 'sequential', enableCpuMemArena: false,