refactor(compute-worker): migrate compute logic to modular architecture

- Remove legacy compute-worker monolith under `src/compute/`, including PDF and Whisper inference, control-plane, and runtime orchestration
- Move PDF and Whisper inference logic to new `src/inference/` module
- Move control-plane, orchestrator, and state machine logic to `src/operations/`
- Move NATS/JetStream adapters to `src/infrastructure/`
- Move job orchestration, progress, and artifact persistence to `src/jobs/`
- Update imports throughout tests and main app to reference new module structure
- Update Dockerfile and scripts to use new asset paths
- Add new entrypoints: `src/api/app.ts` and `src/api/contracts.ts`
- Remove obsolete files and update `.gitignore` for new dev artifacts

This refactor modularizes the compute-worker, improving maintainability and separation of concerns. No functional changes to inference or orchestration logic.

BREAKING CHANGE: compute-worker internal APIs, directory structure, and imports have changed; downstream code must update imports and integration points.
This commit is contained in:
Richard R 2026-06-12 14:07:15 -06:00
parent a56aaa2be9
commit 3d5dfd2a88
75 changed files with 810 additions and 860 deletions

4
.gitignore vendored
View file

@ -56,5 +56,7 @@ node_modules/
# vscode
.vscode
# .agents
# Agents
.agents
.codex
.claude

View file

@ -87,7 +87,7 @@ COPY --from=app-builder /app/THIRD_PARTY_LICENSES /licenses
# Include SeaweedFS license text for the copied weed binary.
COPY --from=seaweedfs-builder /tmp/SeaweedFS-LICENSE.txt /licenses/SeaweedFS-LICENSE.txt
# Include static model notices for runtime-downloaded assets.
COPY --from=app-builder /app/compute-worker/src/compute/pdf/assets/LICENSE.txt /licenses/pp-doclayoutv3-LICENSE.txt
COPY --from=app-builder /app/compute-worker/src/inference/pdf/assets/LICENSE.txt /licenses/pp-doclayoutv3-LICENSE.txt
# Copy seaweedfs weed binary for optional embedded local S3.
COPY --from=seaweedfs-builder /tmp/weed /usr/local/bin/weed
@ -97,7 +97,7 @@ COPY --from=nats-builder /tmp/nats-server /usr/local/bin/nats-server
RUN chmod +x /usr/local/bin/nats-server
# Include OpenAI Whisper license text for runtime-downloaded ONNX artifacts.
COPY --from=app-builder /app/compute-worker/src/compute/whisper/assets/LICENSE.txt /licenses/openai-whisper-LICENSE.txt
COPY --from=app-builder /app/compute-worker/src/inference/whisper/assets/LICENSE.txt /licenses/openai-whisper-LICENSE.txt
# Match the app's historical container port now that standalone server.js
# is started directly instead of `next start -p 3003`.

View file

@ -1,6 +1,6 @@
import { mkdir, writeFile } from 'node:fs/promises';
import { dirname, resolve } from 'node:path';
import { createComputeWorkerApp } from '../src/runtime';
import { createComputeWorkerApp } from '../src/api/app';
process.env.COMPUTE_WORKER_TOKEN ||= 'openapi-generation-token';
process.env.NATS_URL ||= 'nats://127.0.0.1:4222';

View file

@ -1,24 +1,24 @@
import type { TTSSentenceAlignment } from '../types/tts';
import type { ParsedPdfDocument } from '../types/parsed-pdf';
import type { TTSSentenceAlignment } from '../inference/types/tts';
import type { ParsedPdfDocument } from '../inference/types/parsed-pdf';
export type {
TTSAudioBuffer,
TTSAudioBytes,
TTSSentenceAlignment,
TTSSentenceWord,
} from '../types/tts';
} from '../inference/types/tts';
export type {
ParsedPdfBlockKind,
ParsedPdfBlockFragment,
ParsedPdfBlock,
ParsedPdfPage,
ParsedPdfDocument,
} from '../types/parsed-pdf';
} from '../inference/types/parsed-pdf';
export const ALIGN_QUEUE_NAME = 'whisper-align';
export const PDF_LAYOUT_QUEUE_NAME = 'pdf-layout';
export { PDF_PARSER_VERSION } from '../pdf/parser-version';
export { encodeParserVersion } from '../pdf/parser-version-key';
export { PDF_PARSER_VERSION } from '../inference/pdf/parser-version';
export { encodeParserVersion } from '../inference/pdf/parser-version-key';
export interface WhisperAlignJobBase {
text: string;

View file

@ -1,5 +1,5 @@
import { createHash } from 'node:crypto';
import { PDF_PARSER_VERSION } from '../compute/pdf/parser-version';
import { PDF_PARSER_VERSION } from '../inference/pdf/parser-version';
function sha256Hex(input: string): string {
return createHash('sha256').update(input).digest('hex');

View file

@ -1,4 +1,4 @@
import type { WorkerOperationState } from '../compute/api-contracts';
import type { WorkerOperationState } from '../api/contracts';
import { pdfSubjectFromOperationKey } from './operation-keys';
export type PublicOperationSubject =

View file

@ -1,4 +1,4 @@
export * from './api-contracts';
export * from '../api/contracts';
export {
getComputeJobConcurrency,
getAvailableCpuCores,
@ -23,4 +23,4 @@ export { normalizeTextItemsForLayout } from './pdf/normalize-text';
export { mapWordsToSentenceOffsets, type WhisperWord } from './whisper/alignment-map';
export { buildGoertzelCoefficients, goertzelPower } from './whisper/spectral';
export { buildWordsFromTimestampedTokens, extractTokenStartTimestamps } from './whisper/token-timestamps';
export * from './control-plane';
export * from '../operations';

View file

@ -8,12 +8,12 @@ import type {
OperationState,
OperationStateStore,
QueuedOperation,
} from '../compute/control-plane';
} from '../operations';
import type {
PdfLayoutJobRequest,
WhisperAlignJobRequest,
WorkerOperationKind,
} from '../compute/api-contracts';
} from '../api/contracts';
import { createJsonCodec } from './json-codec';
export interface KvEntryLike {

View file

@ -0,0 +1,104 @@
import {
AckPolicy,
DeliverPolicy,
ReplayPolicy,
RetentionPolicy,
StorageType,
type JetStreamManager,
} from '@nats-io/jetstream';
import { nanos } from '@nats-io/transport-node';
import { OP_EVENTS_SUBJECT_WILDCARD } from './nats-adapters';
export const JOBS_STREAM_NAME = 'compute_jobs';
export const WHISPER_JOBS_SUBJECT = 'jobs.whisper';
export const LAYOUT_JOBS_SUBJECT = 'jobs.layout';
export const WHISPER_CONSUMER_NAME = 'compute_whisper';
export const LAYOUT_CONSUMER_NAME = 'compute_layout';
export const EVENTS_STREAM_NAME = 'compute_events';
export const COMPUTE_STATE_BUCKET = 'compute_state';
export const COMPUTE_STATE_TTL_MS = 24 * 60 * 60 * 1000;
export const NATS_API_TIMEOUT_MS = 60_000;
const WHISPER_MAX_DELIVER = 1;
function isAlreadyExistsError(error: unknown): boolean {
const message = (error instanceof Error ? error.message : String(error)).toLowerCase();
return message.includes('already in use') || message.includes('already exists');
}
export async function ensureJetStreamResources(input: {
jsm: JetStreamManager;
whisperTimeoutMs: number;
pdfTimeoutMs: number;
pdfAttempts: number;
jobsMaxBytes: number;
eventsMaxBytes: number;
natsReplicas: number;
}): Promise<void> {
const streamConfig = {
name: JOBS_STREAM_NAME,
subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT],
retention: RetentionPolicy.Workqueue,
storage: StorageType.File,
max_bytes: input.jobsMaxBytes,
num_replicas: input.natsReplicas,
};
try {
await input.jsm.streams.add(streamConfig);
} catch (error) {
if (!isAlreadyExistsError(error)) throw error;
await input.jsm.streams.update(JOBS_STREAM_NAME, {
subjects: streamConfig.subjects,
max_bytes: input.jobsMaxBytes,
num_replicas: input.natsReplicas,
});
}
const eventsStreamConfig = {
name: EVENTS_STREAM_NAME,
subjects: [OP_EVENTS_SUBJECT_WILDCARD],
retention: RetentionPolicy.Limits,
storage: StorageType.File,
max_bytes: input.eventsMaxBytes,
max_age: nanos(COMPUTE_STATE_TTL_MS),
num_replicas: input.natsReplicas,
};
try {
await input.jsm.streams.add(eventsStreamConfig);
} catch (error) {
if (!isAlreadyExistsError(error)) throw error;
await input.jsm.streams.update(EVENTS_STREAM_NAME, {
subjects: eventsStreamConfig.subjects,
max_bytes: input.eventsMaxBytes,
max_age: eventsStreamConfig.max_age,
num_replicas: input.natsReplicas,
});
}
const ensureConsumer = async (name: string, subject: string, ackWaitMs: number, maxDeliver: number) => {
const config = {
durable_name: name,
ack_policy: AckPolicy.Explicit,
deliver_policy: DeliverPolicy.All,
replay_policy: ReplayPolicy.Instant,
filter_subject: subject,
ack_wait: nanos(Math.max(ackWaitMs, 1_000)),
max_deliver: maxDeliver,
};
try {
await input.jsm.consumers.add(JOBS_STREAM_NAME, config);
} catch (error) {
if (!isAlreadyExistsError(error)) throw error;
await input.jsm.consumers.update(JOBS_STREAM_NAME, name, {
filter_subject: subject,
ack_wait: config.ack_wait,
max_deliver: maxDeliver,
});
}
};
await Promise.all([
ensureConsumer(WHISPER_CONSUMER_NAME, WHISPER_JOBS_SUBJECT, input.whisperTimeoutMs + 15_000, WHISPER_MAX_DELIVER),
ensureConsumer(LAYOUT_CONSUMER_NAME, LAYOUT_JOBS_SUBJECT, input.pdfTimeoutMs + 15_000, input.pdfAttempts),
]);
}

View file

@ -0,0 +1,127 @@
import {
DeleteObjectCommand,
GetObjectCommand,
HeadObjectCommand,
PutObjectCommand,
S3Client,
} from '@aws-sdk/client-s3';
import { parsedPdfArtifactKey } from '../storage/artifact-addressing';
export interface ArtifactStorage {
readObject(key: string): Promise<ArrayBuffer>;
objectExists(key: string): Promise<boolean>;
deleteObject(key: string): Promise<void>;
putParsedPdf(documentId: string, namespace: string | null, parsed: unknown): Promise<string>;
}
export interface ArtifactStorageConfig {
bucket: string;
prefix: string;
client: S3Client;
}
function bodyToBuffer(body: unknown): Promise<Buffer> | Buffer {
if (!body) return Buffer.alloc(0);
if (body instanceof Uint8Array) return Buffer.from(body);
if (ArrayBuffer.isView(body)) return Buffer.from(body.buffer, body.byteOffset, body.byteLength);
if (body instanceof ArrayBuffer) return Buffer.from(body);
if (typeof body === 'object' && body !== null && 'transformToByteArray' in body) {
const maybe = body as { transformToByteArray?: () => Promise<Uint8Array> };
if (typeof maybe.transformToByteArray === 'function') {
return maybe.transformToByteArray().then((bytes) => Buffer.from(bytes));
}
}
if (typeof body === 'object' && body !== null && 'on' in body) {
return (async () => {
const chunks: Buffer[] = [];
for await (const chunk of body as NodeJS.ReadableStream) {
if (Buffer.isBuffer(chunk)) chunks.push(chunk);
else if (typeof chunk === 'string') chunks.push(Buffer.from(chunk));
else chunks.push(Buffer.from(chunk as Uint8Array));
}
return Buffer.concat(chunks);
})();
}
throw new Error('Unsupported S3 response body type');
}
function toArrayBuffer(bytes: Uint8Array): ArrayBuffer {
const copy = new Uint8Array(bytes.byteLength);
copy.set(bytes);
return copy.buffer;
}
function isNotFound(error: unknown): boolean {
const maybe = error as { name?: string; Code?: string; $metadata?: { httpStatusCode?: number } };
return maybe.$metadata?.httpStatusCode === 404
|| maybe.name === 'NotFound'
|| maybe.name === 'NoSuchKey'
|| maybe.Code === 'NotFound'
|| maybe.Code === 'NoSuchKey';
}
export function normalizeS3Prefix(prefix: string | undefined): string {
const value = (prefix || 'openreader').trim();
return value ? value.replace(/^\/+|\/+$/g, '') : 'openreader';
}
export function createS3ClientFromEnv(requireEnv: (name: string) => string): S3Client {
return new S3Client({
region: requireEnv('S3_REGION'),
endpoint: process.env.S3_ENDPOINT?.trim() || undefined,
forcePathStyle: ['1', 'true', 'yes', 'on'].includes(process.env.S3_FORCE_PATH_STYLE?.trim().toLowerCase() ?? ''),
requestChecksumCalculation: 'WHEN_REQUIRED',
responseChecksumValidation: 'WHEN_REQUIRED',
credentials: {
accessKeyId: requireEnv('S3_ACCESS_KEY_ID'),
secretAccessKey: requireEnv('S3_SECRET_ACCESS_KEY'),
},
});
}
export function createArtifactStorage(config: ArtifactStorageConfig): ArtifactStorage {
const safeKey = (key: string): string => {
const trimmed = key.trim();
if (!trimmed.startsWith(`${config.prefix}/`)) throw new Error('Object key prefix mismatch');
return trimmed;
};
return {
async readObject(key) {
const response = await config.client.send(new GetObjectCommand({
Bucket: config.bucket,
Key: safeKey(key),
}));
return toArrayBuffer(new Uint8Array(await bodyToBuffer(response.Body)));
},
async objectExists(key) {
try {
await config.client.send(new HeadObjectCommand({
Bucket: config.bucket,
Key: safeKey(key),
}));
return true;
} catch (error) {
if (isNotFound(error)) return false;
throw error;
}
},
async deleteObject(key) {
await config.client.send(new DeleteObjectCommand({
Bucket: config.bucket,
Key: safeKey(key),
}));
},
async putParsedPdf(documentId, namespace, parsed) {
const key = parsedPdfArtifactKey({ documentId, namespace, prefix: config.prefix });
await config.client.send(new PutObjectCommand({
Bucket: config.bucket,
Key: key,
Body: Buffer.from(JSON.stringify(parsed)),
ContentType: 'application/json',
ServerSideEncryption: 'AES256',
}));
return key;
},
};
}

View file

@ -0,0 +1,126 @@
import { z } from 'zod';
import {
runPdfLayoutFromPdfBuffer,
runWhisperAlignmentFromAudioBuffer,
} from '../inference/local-runtime';
import { withIdleTimeoutAndHardCap, withTimeout } from '../inference';
import type {
PdfLayoutJobRequest,
PdfLayoutJobResult,
PdfLayoutProgress,
WhisperAlignJobRequest,
WhisperAlignJobResult,
} from '../api/contracts';
import type { ArtifactStorage } from '../infrastructure/storage';
import { persistParsedPdfWhileSourceExists } from './pdf-artifact-persistence';
import { buildInferProgressForPageParsed, buildInferProgressForPageStart } from './pdf-progress';
const whisperRequestSchema = z.object({
text: z.string().trim().min(1),
lang: z.string().trim().min(1).max(16).optional(),
cacheKey: z.string().trim().min(1).max(256).optional(),
audioObjectKey: z.string().trim().min(1).max(2048),
});
const pdfRequestSchema = z.object({
documentId: z.string().trim().min(1),
namespace: z.string().trim().min(1).max(128).nullable(),
documentObjectKey: z.string().trim().min(1).max(2048),
});
export interface JobHandlers {
runWhisper(payload: WhisperAlignJobRequest, queueWaitMs: number): Promise<WhisperAlignJobResult>;
runPdfLayout(
payload: PdfLayoutJobRequest,
queueWaitMs: number,
hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise<void> },
): Promise<PdfLayoutJobResult>;
}
export function createJobHandlers(input: {
storage: ArtifactStorage;
whisperTimeoutMs: number;
pdfTimeoutMs: number;
pdfHardCapMs: number;
}): JobHandlers {
return {
async runWhisper(payload, queueWaitMs) {
const parsed = whisperRequestSchema.parse(payload);
const s3FetchStartedAt = Date.now();
const audioBuffer = await withTimeout(
input.storage.readObject(parsed.audioObjectKey),
input.whisperTimeoutMs,
'whisper s3 fetch',
);
const s3FetchMs = Date.now() - s3FetchStartedAt;
const computeStartedAt = Date.now();
const result = await withTimeout(
runWhisperAlignmentFromAudioBuffer({
audioBuffer,
text: parsed.text,
cacheKey: parsed.cacheKey,
lang: parsed.lang,
}),
input.whisperTimeoutMs,
'whisper alignment job',
);
return {
...result,
timing: { queueWaitMs, s3FetchMs, computeMs: Date.now() - computeStartedAt },
};
},
async runPdfLayout(payload, queueWaitMs, hooks) {
const parsed = pdfRequestSchema.parse(payload);
const s3FetchStartedAt = Date.now();
const pdfBytes = await withTimeout(
input.storage.readObject(parsed.documentObjectKey),
Math.max(input.pdfTimeoutMs, 1_000),
'pdf s3 fetch',
);
const s3FetchMs = Date.now() - s3FetchStartedAt;
let lastTotalPages = 0;
let lastPagesParsed = 0;
const computeStartedAt = Date.now();
const result = await withIdleTimeoutAndHardCap({
idleTimeoutMs: Math.max(input.pdfTimeoutMs, 1_000),
hardCapMs: input.pdfHardCapMs,
label: 'pdf layout job',
run: async (touchProgress) => runPdfLayoutFromPdfBuffer({
documentId: parsed.documentId,
pdfBytes,
onPageStarted: async ({ pageNumber, totalPages }) => {
touchProgress();
lastTotalPages = totalPages;
await hooks?.onProgress?.(buildInferProgressForPageStart({ pageNumber, totalPages }));
},
onPageParsed: async ({ pageNumber, totalPages }) => {
touchProgress();
lastTotalPages = totalPages;
lastPagesParsed = pageNumber;
await hooks?.onProgress?.(buildInferProgressForPageParsed({ pageNumber, totalPages }));
},
}),
});
const computeMs = Date.now() - computeStartedAt;
if (hooks?.onProgress && lastTotalPages > 0) {
await hooks.onProgress({
totalPages: lastTotalPages,
pagesParsed: lastPagesParsed,
currentPage: lastPagesParsed || undefined,
phase: 'merge',
});
}
const parsedObjectKey = await persistParsedPdfWhileSourceExists({
sourceObjectKey: parsed.documentObjectKey,
sourceExists: input.storage.objectExists,
putParsedObject: () => input.storage.putParsedPdf(parsed.documentId, parsed.namespace, result.parsed),
deleteParsedObject: input.storage.deleteObject,
});
return {
parsedObjectKey,
timing: { queueWaitMs, s3FetchMs, computeMs },
};
},
};
}

View file

@ -1,4 +1,4 @@
import type { PdfLayoutProgress } from './compute/api-contracts';
import type { PdfLayoutProgress } from '../api/contracts';
export function buildInferProgressForPageStart(input: {
pageNumber: number;

View file

@ -1,4 +1,4 @@
import type { WorkerOperationKind } from './compute/api-contracts';
import type { WorkerOperationKind } from '../api/contracts';
export type RetryAction = 'nak_retry' | 'term_fail';

View file

@ -0,0 +1,312 @@
import type { Consumer, JsMsg } from '@nats-io/jetstream';
import type {
PdfLayoutJobRequest,
PdfLayoutJobResult,
PdfLayoutProgress,
WhisperAlignJobRequest,
WhisperAlignJobResult,
WorkerJobTiming,
WorkerOperationKind,
} from '../api/contracts';
import type { JsonCodec } from '../infrastructure/json-codec';
import type { JobHandlers } from './handlers';
import { buildQueueWaitTiming, decideRetryAction } from './worker-loop-policy';
const LOOP_ERROR_BACKOFF_MS = 500;
const RUNNING_HEARTBEAT_MS = 5000;
const PULL_EXPIRES_MS = 5_000;
const WHISPER_MAX_DELIVER = 1;
const SLOW_JOB_LOG_THRESHOLD_MS_BY_KIND: Record<WorkerOperationKind, number> = {
whisper_align: 15_000,
pdf_layout: 120_000,
};
export interface QueuedJob<TPayload> {
jobId: string;
opId: string;
opKey: string;
kind: WorkerOperationKind;
queuedAt: number;
payload: TPayload;
}
interface WorkerLoopOrchestrator {
markRunning(input: { opId: string; startedAt?: number; updatedAt?: number; timing?: WorkerJobTiming }): Promise<unknown>;
markProgress(input: {
opId: string;
progress: PdfLayoutProgress;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
markSucceeded(input: { opId: string; result: unknown; updatedAt?: number; timing?: WorkerJobTiming }): Promise<unknown>;
markFailed(input: {
opId: string;
error: { message: string; code?: string } | string;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
}
interface WorkerLogger {
info(data: unknown, message?: string): void;
warn(data: unknown, message?: string): void;
error(data: unknown, message?: string): void;
}
class ConcurrencyGate {
private inFlight = 0;
private readonly queue: Array<() => void> = [];
constructor(private readonly limit: number) {}
async acquire(): Promise<void> {
if (this.inFlight < this.limit) {
this.inFlight += 1;
return;
}
await new Promise<void>((resolve) => {
this.queue.push(() => {
this.inFlight += 1;
resolve();
});
});
}
release(): void {
this.inFlight = Math.max(0, this.inFlight - 1);
this.queue.shift()?.();
}
}
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
function toErrorMessage(error: unknown): string {
return error instanceof Error && error.message ? error.message : String(error);
}
function safeDurationMs(start: number, end: number): number {
return Math.max(0, Math.floor(end - start));
}
function extractTiming(result: unknown): WorkerJobTiming | undefined {
if (!result || typeof result !== 'object' || !('timing' in result)) return undefined;
return (result as { timing?: WorkerJobTiming }).timing;
}
function extractResultRef(kind: WorkerOperationKind, result: unknown): string | undefined {
if (kind !== 'pdf_layout' || !result || typeof result !== 'object') return undefined;
const maybe = result as { parsedObjectKey?: unknown };
return typeof maybe.parsedObjectKey === 'string' ? maybe.parsedObjectKey : undefined;
}
export function createWorkerLoopController(input: {
orchestrator: WorkerLoopOrchestrator;
handlers: JobHandlers;
logger: WorkerLogger;
jobConcurrency: number;
pdfAttempts: number;
whisperCodec: JsonCodec<QueuedJob<WhisperAlignJobRequest>>;
pdfCodec: JsonCodec<QueuedJob<PdfLayoutJobRequest>>;
isOwnerActive: (owner: object) => boolean;
isStopping: () => boolean;
markActivity: (reason: string) => void;
onInFlightJobsChanged: (delta: number) => void;
}) {
const gate = new ConcurrencyGate(Math.max(1, Math.floor(input.jobConcurrency)));
let loops: Promise<void>[] = [];
let stopRequested = false;
type Context<TPayload> = {
decoded: QueuedJob<TPayload>;
workerLabel: string;
startedAt: number;
queueWaitTiming?: { queueWaitMs: number };
latestProgress?: PdfLayoutProgress;
};
const markRunning = async <TPayload>(context: Context<TPayload>, updatedAt: number): Promise<void> => {
if (context.latestProgress) {
await input.orchestrator.markProgress({
opId: context.decoded.opId,
progress: context.latestProgress,
updatedAt,
...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}),
});
return;
}
await input.orchestrator.markRunning({
opId: context.decoded.opId,
startedAt: context.startedAt,
updatedAt,
...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}),
});
};
const processMessage = async <TPayload, TResult>(work: {
msg: JsMsg;
codec: JsonCodec<QueuedJob<TPayload>>;
run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout'];
workerLabel: string;
}): Promise<void> => {
let context: Context<TPayload> | null = null;
let heartbeat: NodeJS.Timeout | null = null;
try {
const decoded = work.codec.decode(work.msg.data);
const startedAt = Date.now();
context = {
decoded,
workerLabel: work.workerLabel,
startedAt,
queueWaitTiming: buildQueueWaitTiming(decoded.queuedAt, startedAt),
};
await markRunning(context, startedAt);
input.logger.info({
worker: work.workerLabel,
kind: decoded.kind,
opId: decoded.opId,
jobId: decoded.jobId,
queueWaitMs: context.queueWaitTiming?.queueWaitMs ?? null,
deliveryCount: work.msg.info.deliveryCount,
}, 'job.started');
heartbeat = setInterval(() => {
void markRunning(context!, Date.now()).catch((error) => {
input.logger.error({
worker: work.workerLabel,
opId: context?.decoded.opId,
jobId: context?.decoded.jobId,
error: toErrorMessage(error),
}, 'failed to persist operation heartbeat state');
});
}, RUNNING_HEARTBEAT_MS);
const result = await work.run(decoded.payload as never, context.queueWaitTiming?.queueWaitMs ?? 0, {
onProgress: async (progress) => {
try {
work.msg.working();
} catch (error) {
input.logger.warn({
worker: work.workerLabel,
kind: context?.decoded.kind,
opId: context?.decoded.opId,
jobId: context?.decoded.jobId,
error: toErrorMessage(error),
}, 'failed to extend JetStream ack wait on progress');
}
context!.latestProgress = progress;
await markRunning(context!, Date.now());
},
});
const timing = extractTiming(result);
const now = Date.now();
await input.orchestrator.markSucceeded({
opId: decoded.opId,
result,
updatedAt: now,
...(timing ? { timing } : {}),
});
work.msg.ack();
const durationMs = safeDurationMs(startedAt, now);
if (durationMs >= SLOW_JOB_LOG_THRESHOLD_MS_BY_KIND[decoded.kind]) {
input.logger.info({ worker: work.workerLabel, kind: decoded.kind, opId: decoded.opId, jobId: decoded.jobId, durationMs, timing: timing ?? null }, 'job.stage');
}
input.logger.info({
worker: work.workerLabel,
kind: decoded.kind,
opId: decoded.opId,
jobId: decoded.jobId,
status: 'succeeded',
durationMs,
resultRef: extractResultRef(decoded.kind, result),
timing: timing ?? null,
}, 'job.terminal');
} catch (error) {
const errorMessage = toErrorMessage(error);
const deliveryCount = work.msg.info.deliveryCount;
const kind = context?.decoded.kind ?? 'pdf_layout';
const action = decideRetryAction({ kind, deliveryCount, pdfAttempts: input.pdfAttempts, whisperMaxDeliver: WHISPER_MAX_DELIVER });
const timing = context ? buildQueueWaitTiming(context.decoded.queuedAt, Date.now()) : undefined;
if (context) {
const update = action === 'nak_retry'
? markRunning(context, Date.now())
: input.orchestrator.markFailed({
opId: context.decoded.opId,
error: { message: errorMessage },
updatedAt: Date.now(),
...(timing ? { timing } : {}),
});
await update.catch((stateError) => input.logger.error({
worker: context?.workerLabel,
opId: context?.decoded.opId,
jobId: context?.decoded.jobId,
error: toErrorMessage(stateError),
}, 'failed to persist operation state'));
}
if (action === 'nak_retry') work.msg.nak();
else work.msg.term(errorMessage);
input.logger.error({
worker: context?.workerLabel,
kind: context?.decoded.kind,
opId: context?.decoded.opId,
jobId: context?.decoded.jobId,
status: action === 'nak_retry' ? 'running' : 'failed',
error: errorMessage,
deliveryCount,
retryAction: action === 'nak_retry' ? 'nack_retry' : 'term',
}, 'job.terminal');
} finally {
if (heartbeat) clearInterval(heartbeat);
}
};
const runLoop = async <TPayload>(work: {
owner: object;
consumer: Consumer;
codec: JsonCodec<QueuedJob<TPayload>>;
run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout'];
workerLabel: string;
}): Promise<void> => {
const detached = () => input.isStopping() || stopRequested || !input.isOwnerActive(work.owner);
while (!detached()) {
let msg: JsMsg | null = null;
try {
try {
msg = await work.consumer.next({ expires: PULL_EXPIRES_MS });
} catch (error) {
if (detached()) return;
input.logger.error({ error: toErrorMessage(error), worker: work.workerLabel }, 'worker pull failed');
await sleep(LOOP_ERROR_BACKOFF_MS);
continue;
}
if (!msg) continue;
input.markActivity(`job_received:${work.workerLabel}`);
input.onInFlightJobsChanged(1);
await gate.acquire();
if (detached()) return;
await processMessage({ ...work, msg });
} finally {
if (msg) {
gate.release();
input.onInFlightJobsChanged(-1);
input.markActivity(`job_completed:${work.workerLabel}`);
}
}
}
};
return {
start(owner: object, consumers: { whisper: Consumer; pdfLayout: Consumer }): void {
stopRequested = false;
loops = [];
for (let i = 0; i < input.jobConcurrency; i += 1) {
loops.push(runLoop({ owner, consumer: consumers.whisper, codec: input.whisperCodec, run: input.handlers.runWhisper, workerLabel: `whisper-${i + 1}` }));
loops.push(runLoop({ owner, consumer: consumers.pdfLayout, codec: input.pdfCodec, run: input.handlers.runPdfLayout, workerLabel: `layout-${i + 1}` }));
}
},
async stop(): Promise<void> {
stopRequested = true;
await Promise.allSettled(loops);
loops = [];
},
};
}

View file

@ -1,4 +1,4 @@
export * from './types';
export * from './state-machine';
export * from './orchestrator';
export * from './service';
export * from './sse';

View file

@ -4,7 +4,7 @@ import type {
WorkerJobTiming,
WorkerJobState,
WorkerOperationState,
} from './compute/api-contracts';
} from '../api/contracts';
export type StreamedOperationState = WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>;

View file

@ -5,7 +5,7 @@ import type {
WorkerOperationKind,
WorkerOperationRequest,
WorkerOperationState,
} from '../api-contracts';
} from '../api/contracts';
import {
buildQueuedState,
createErrorShape,

View file

@ -4,7 +4,7 @@ import type {
WorkerOperationKind,
WorkerOperationRequest,
WorkerOperationState,
} from '../api-contracts';
} from '../api/contracts';
export function isTerminalStatus(status: WorkerJobState): boolean {
return status === 'succeeded' || status === 'failed';

View file

@ -2,7 +2,7 @@ import type {
WorkerOperationEvent,
WorkerOperationKind,
WorkerOperationState,
} from '../api-contracts';
} from '../api/contracts';
export type OperationState<Result = unknown> = WorkerOperationState<Result>;
export type OperationEvent<Result = unknown> = WorkerOperationEvent<Result>;

View file

@ -1,4 +1,4 @@
import { startComputeWorkerFromEnv } from './runtime';
import { startComputeWorkerFromEnv } from './api/app';
void startComputeWorkerFromEnv().catch((error) => {
console.error('[compute-worker] fatal startup error', error);

View file

@ -1,5 +1,5 @@
import { PDF_PARSER_VERSION } from '../compute/pdf/parser-version';
import { encodeParserVersion } from '../compute/pdf/parser-version-key';
import { PDF_PARSER_VERSION } from '../inference/pdf/parser-version';
import { encodeParserVersion } from '../inference/pdf/parser-version-key';
const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i;
const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/;

View file

@ -1,5 +1,5 @@
import { afterEach, beforeEach, describe, expect, test } from 'vitest';
import { createComputeWorkerApp } from '../../src/runtime';
import { createComputeWorkerApp } from '../../src/api/app';
import { FakeControlPlane } from '../fixtures/fake-control-plane';
const AUTH = { authorization: 'Bearer test-token' };

View file

@ -1,5 +1,5 @@
import { describe, expect, test } from 'vitest';
import { mergeTextWithRegions } from '../../../src/compute';
import { mergeTextWithRegions } from '../../../src/inference';
describe('mergeTextWithRegions', () => {
test('assigns text items to containing regions by centroid and joins in reading order', () => {

View file

@ -1,6 +1,6 @@
import { describe, expect, test } from 'vitest';
import { normalizeTextItemsForLayout } from '../../../src/compute';
import { normalizeTextItemsForLayout } from '../../../src/inference';
import type { TextItem } from 'pdfjs-dist/types/src/display/api';
function makeTextItem(

View file

@ -1,6 +1,6 @@
import { describe, expect, test } from 'vitest';
import { stitchCrossPageBlocks } from '../../../src/compute';
import type { ParsedPdfBlock, ParsedPdfDocument, ParsedPdfBlockKind } from '../../../src/compute/types';
import { stitchCrossPageBlocks } from '../../../src/inference';
import type { ParsedPdfBlock, ParsedPdfDocument, ParsedPdfBlockKind } from '../../../src/inference/types';
import { makeParsedPdfDocument, makeParsedPdfPage } from './support/document-fixtures';
function makeBlock(

View file

@ -1,5 +1,5 @@
import type { BaseDocument } from '../../../src/types/documents';
import type { ParsedPdfBlock, ParsedPdfBlockKind, ParsedPdfDocument, ParsedPdfPage } from '../../../../src/compute/types';
import type { BaseDocument } from '../../../../../src/types/documents';
import type { ParsedPdfBlock, ParsedPdfBlockKind, ParsedPdfDocument, ParsedPdfPage } from '../../../../src/inference/types';
export function makeBaseDocument(overrides: Partial<BaseDocument> = {}): BaseDocument {
return {

View file

@ -1,7 +1,7 @@
import { describe, expect, test } from 'vitest';
import {
mapWordsToSentenceOffsets,
} from '../../../src/compute';
} from '../../../src/inference';
describe('whisper alignment mapping', () => {
test('maps words to sentence offsets with punctuation and repeated spaces', () => {

View file

@ -1,5 +1,5 @@
import { describe, expect, test } from 'vitest';
import { buildGoertzelCoefficients, goertzelPower } from '../../../src/compute';
import { buildGoertzelCoefficients, goertzelPower } from '../../../src/inference';
function dftPower(samples: Float32Array, k: number): number {
const n = samples.length;

View file

@ -3,7 +3,7 @@ import * as ort from 'onnxruntime-node';
import {
buildWordsFromTimestampedTokens,
extractTokenStartTimestamps,
} from '../../../src/compute';
} from '../../../src/inference';
describe('whisper token timestamp alignment', () => {
test('extracts monotonic token timestamps from cross-attention maps', () => {

View file

@ -1,6 +1,6 @@
import { describe, expect, test } from 'vitest';
import type { WorkerOperationRequest } from '../../../src/compute/api-contracts';
import { OperationOrchestrator } from '../../../src/compute/control-plane';
import type { WorkerOperationRequest } from '../../../src/api/contracts';
import { OperationOrchestrator } from '../../../src/operations';
import {
InMemoryOperationEventStream,
InMemoryOperationQueue,

View file

@ -70,13 +70,13 @@ vi.mock('@napi-rs/canvas', () => {
};
});
vi.mock('../../../src/compute/pdf/model', () => ({
vi.mock('../../../src/inference/pdf/model', () => ({
ensureModel: vi.fn(async () => '/tmp/model.onnx'),
MODEL_CONFIG_PATH: '/tmp/model-config.json',
MODEL_PREPROCESSOR_PATH: '/tmp/model-preprocessor.json',
}));
vi.mock('../../../src/compute/config/cpu-budget', () => ({
vi.mock('../../../src/inference/config/cpu-budget', () => ({
getOnnxThreadsPerJob: vi.fn(() => 1),
}));
@ -107,7 +107,7 @@ describe('runLayoutModel', () => {
},
};
const { runLayoutModel } = await import('../../../src/compute/pdf/runLayoutModel');
const { runLayoutModel } = await import('../../../src/inference/pdf/runLayoutModel');
const regions = await runLayoutModel({
pageWidth: 100,
pageHeight: 100,
@ -154,7 +154,7 @@ describe('runLayoutModel', () => {
},
};
const { runLayoutModel } = await import('../../../src/compute/pdf/runLayoutModel');
const { runLayoutModel } = await import('../../../src/inference/pdf/runLayoutModel');
const regions = await runLayoutModel({
pageWidth: 100,
pageHeight: 100,

View file

@ -1,5 +1,5 @@
import { describe, expect, test } from 'vitest';
import { encodeSseFrame, parseSseEventId, parseSsePayload } from '../../../src/compute/control-plane/sse';
import { encodeSseFrame, parseSseEventId, parseSsePayload } from '../../../src/operations/sse';
describe('sse codec', () => {
test('encodes event id and payload and decodes both reliably', () => {

View file

@ -1,11 +1,11 @@
import { describe, expect, test } from 'vitest';
import type { WorkerOperationState } from '../../../src/compute/api-contracts';
import type { WorkerOperationState } from '../../../src/api/contracts';
import {
explainReplacementReason,
isInflightStatus,
isTerminalStatus,
shouldReuseExistingOperation,
} from '../../../src/compute/control-plane/state-machine';
} from '../../../src/operations/state-machine';
function runningState(overrides: Partial<WorkerOperationState> = {}): WorkerOperationState {
return {

View file

@ -1,5 +1,5 @@
import { EventEmitter } from 'node:events';
import type { WorkerOperationKind } from '../../../src/compute/api-contracts';
import type { WorkerOperationKind } from '../../../src/api/contracts';
import type {
OperationEvent,
OperationEventStream,
@ -8,7 +8,7 @@ import type {
OperationState,
OperationStateStore,
QueuedOperation,
} from '../../../src/compute/control-plane/types';
} from '../../../src/operations/types';
function topicFor(opId: string): string {
return `op.${opId}`;

View file

@ -4,8 +4,8 @@ import type {
WorkerOperationEvent,
WorkerOperationRequest,
WorkerOperationState,
} from '../../src/compute/api-contracts';
import type { ComputeWorkerRouteDeps } from '../../src/runtime';
} from '../../src/api/contracts';
import type { ComputeWorkerRouteDeps } from '../../src/api/app';
type ComputeResult = WhisperAlignJobResult | PdfLayoutJobResult;
type ComputeState = WorkerOperationState<ComputeResult>;

View file

@ -1,6 +1,6 @@
import { describe, expect, test } from 'vitest';
import { OperationOrchestrator } from '../../src/compute/control-plane';
import type { WorkerOperationRequest } from '../../src/compute/api-contracts';
import { OperationOrchestrator } from '../../src/operations';
import type { WorkerOperationRequest } from '../../src/api/contracts';
import {
JetStreamOperationEventStream,
JetStreamOperationQueue,
@ -11,7 +11,7 @@ import {
opStateKvKey,
type KvEntryLike,
type KvStoreLike,
} from '../../src/control-plane/jetstream';
} from '../../src/infrastructure/nats-adapters';
class FakeKvStore implements KvStoreLike {
private readonly data = new Map<string, KvEntryLike>();

View file

@ -1,9 +1,9 @@
import { describe, expect, test } from 'vitest';
import { hashOpKey, opIndexKvKey, opStateKvKey } from '../../src/control-plane/jetstream';
import { hashOpKey, opIndexKvKey, opStateKvKey } from '../../src/infrastructure/nats-adapters';
import {
buildInferProgressForPageParsed,
buildInferProgressForPageStart,
} from '../../src/pdf-progress';
} from '../../src/jobs/pdf-progress';
import { buildPdfOperationKey } from '../../src/api/operation-keys';
import { parsedPdfArtifactKey } from '../../src/storage/artifact-addressing';

View file

@ -1,5 +1,5 @@
import { afterEach, describe, expect, test, vi } from 'vitest';
import { recoverOrphanedOperations } from '../../src/orphan-recovery';
import { recoverOrphanedOperations } from '../../src/operations/recovery';
import { FakeControlPlane } from '../fixtures/fake-control-plane';
describe('orphan recovery', () => {

View file

@ -1,5 +1,5 @@
import { describe, expect, test, vi } from 'vitest';
import { persistParsedPdfWhileSourceExists } from '../../src/pdf-artifact-persistence';
import { persistParsedPdfWhileSourceExists } from '../../src/jobs/pdf-artifact-persistence';
describe('PDF artifact persistence', () => {
test('does not write parsed output after the source was deleted', async () => {

View file

@ -1,5 +1,5 @@
import { describe, expect, test } from 'vitest';
import { buildQueueWaitTiming, decideRetryAction } from '../../src/worker-loop-policy';
import { buildQueueWaitTiming, decideRetryAction } from '../../src/jobs/worker-loop-policy';
describe('worker loop policy', () => {
test('returns queue wait timing with non-negative clamped duration', () => {

View file

@ -13,7 +13,7 @@ const forbidden = [
'onnxruntime-node',
'@huggingface/tokenizers',
'@openreader/compute-worker',
'/compute-worker/src/compute/',
'/compute-worker/src/inference/',
];
const includeExt = new Set(['.js', '.mjs', '.cjs']);

View file

@ -1,5 +1,5 @@
import { afterEach, describe, expect, test } from 'vitest';
import { createComputeWorkerApp, type ComputeWorkerApp } from '../../compute-worker/src/runtime';
import { createComputeWorkerApp, type ComputeWorkerApp } from '../../compute-worker/src/api/app';
import { FakeControlPlane } from '../../compute-worker/tests/fixtures/fake-control-plane';
import { ComputeWorkerClient } from '../../src/lib/server/compute-worker/client';