From f8672b073e39c41162fafea5f69cf9a939a87ac0 Mon Sep 17 00:00:00 2001 From: Richard R Date: Thu, 21 May 2026 15:58:29 -0600 Subject: [PATCH] refactor(compute): unify job concurrency and CPU thread allocation logic Standardize compute job concurrency configuration by introducing a shared COMPUTE_JOB_CONCURRENCY setting, replacing separate PDF and Whisper concurrency controls. Add cross-platform CPU core detection and thread budgeting utilities, and update both ONNX model execution and concurrency limiters to use dynamic thread allocation per job. Refactor environment variable docs and examples to reflect unified concurrency management. Streamline PDF.js font path resolution for improved reliability across environments. --- compute/core/src/pdf-layout/parsePdf.ts | 11 +- compute/core/src/pdf-layout/renderPage.ts | 11 +- compute/core/src/pdf-layout/runLayoutModel.ts | 10 +- compute/core/src/runtime/cpu-budget.ts | 28 +++++ compute/core/src/whisper/alignment.ts | 4 +- compute/worker/.env.example | 5 +- compute/worker/src/server.ts | 106 ++++++++++++++---- docs-site/docs/deploy/compute-worker.md | 10 +- src/lib/server/compute/concurrency-limiter.ts | 5 +- src/lib/server/compute/cpu-budget.ts | 28 +++++ src/lib/server/compute/local.ts | 6 +- src/lib/server/pdf-layout/parsePdf.ts | 11 +- src/lib/server/pdf-layout/renderPage.ts | 11 +- src/lib/server/pdf-layout/runLayoutModel.ts | 10 +- src/lib/server/whisper/alignment.ts | 4 +- 15 files changed, 214 insertions(+), 46 deletions(-) create mode 100644 compute/core/src/runtime/cpu-budget.ts create mode 100644 src/lib/server/compute/cpu-budget.ts 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,