1195 lines
37 KiB
TypeScript
1195 lines
37 KiB
TypeScript
import { createHash } from 'node:crypto';
|
|
import os from 'node:os';
|
|
import Fastify, { type FastifyRequest } from 'fastify';
|
|
import { z } from 'zod';
|
|
import {
|
|
connect,
|
|
nanos,
|
|
credsAuthenticator,
|
|
type NatsConnection,
|
|
} from '@nats-io/transport-node';
|
|
import {
|
|
AckPolicy,
|
|
DeliverPolicy,
|
|
ReplayPolicy,
|
|
RetentionPolicy,
|
|
StorageType,
|
|
jetstream,
|
|
jetstreamManager,
|
|
type Consumer,
|
|
type JetStreamClient,
|
|
type JetStreamManager,
|
|
type JsMsg,
|
|
} from '@nats-io/jetstream';
|
|
import { Kvm, type KV } from '@nats-io/kv';
|
|
import {
|
|
ensureComputeModels,
|
|
runPdfLayoutFromPdfBuffer,
|
|
runWhisperAlignmentFromAudioBuffer,
|
|
} from '@openreader/compute-core/local-runtime';
|
|
import {
|
|
getComputeTimeoutConfig,
|
|
withIdleTimeoutAndHardCap,
|
|
withTimeout,
|
|
} from '@openreader/compute-core/runtime/timeout-config';
|
|
import {
|
|
type PdfLayoutJobRequest,
|
|
type PdfLayoutJobResult,
|
|
type WhisperAlignJobRequest,
|
|
type WhisperAlignJobResult,
|
|
type WorkerJobErrorShape,
|
|
type WorkerJobState,
|
|
type WorkerJobTiming,
|
|
type WorkerOperationKind,
|
|
type WorkerOperationRequest,
|
|
type WorkerOperationState,
|
|
type PdfLayoutProgress,
|
|
} from '@openreader/compute-core/contracts';
|
|
import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3';
|
|
|
|
const JOBS_STREAM_NAME = 'compute_jobs';
|
|
const WHISPER_JOBS_SUBJECT = 'jobs.whisper';
|
|
const LAYOUT_JOBS_SUBJECT = 'jobs.layout';
|
|
const WHISPER_CONSUMER_NAME = 'compute_whisper';
|
|
const LAYOUT_CONSUMER_NAME = 'compute_layout';
|
|
const COMPUTE_STATE_BUCKET = 'compute_state';
|
|
const COMPUTE_STATE_TTL_MS = 24 * 60 * 60 * 1000;
|
|
const LOOP_ERROR_BACKOFF_MS = 500;
|
|
const SSE_POLL_INTERVAL_MS = 400;
|
|
const RUNNING_HEARTBEAT_MS = 5000;
|
|
const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i;
|
|
const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/;
|
|
const WHISPER_MAX_DELIVER = 1;
|
|
const NATS_API_TIMEOUT_MS = 60_000;
|
|
|
|
interface QueuedJob<TPayload> {
|
|
jobId: string;
|
|
opId: string;
|
|
opKey: string;
|
|
kind: WorkerOperationKind;
|
|
queuedAt: number;
|
|
payload: TPayload;
|
|
}
|
|
|
|
interface StoredJobState<Result> {
|
|
jobId: string;
|
|
opId: string;
|
|
opKey: string;
|
|
kind: WorkerOperationKind;
|
|
status: WorkerJobState;
|
|
timestamp: number;
|
|
startedAt?: number;
|
|
updatedAt: number;
|
|
result?: Result;
|
|
error?: WorkerJobErrorShape;
|
|
timing?: WorkerJobTiming;
|
|
progress?: PdfLayoutProgress;
|
|
}
|
|
|
|
interface OpIndexEntry {
|
|
opId: string;
|
|
}
|
|
|
|
type JsonCodec<T> = {
|
|
encode(value: T): Uint8Array;
|
|
decode(data: Uint8Array): T;
|
|
};
|
|
|
|
function requireEnv(name: string): string {
|
|
const value = process.env[name]?.trim();
|
|
if (!value) throw new Error(`${name} is required`);
|
|
return value;
|
|
}
|
|
|
|
function readIntEnv(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);
|
|
}
|
|
|
|
function parseBoolEnv(name: string, fallback: boolean): boolean {
|
|
const raw = process.env[name]?.trim();
|
|
if (!raw) return fallback;
|
|
const normalized = raw.toLowerCase();
|
|
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<string, unknown> {
|
|
const format = (process.env.COMPUTE_LOG_FORMAT?.trim().toLowerCase() || 'pretty');
|
|
if (format === 'json') return true;
|
|
return {
|
|
level: process.env.COMPUTE_LOG_LEVEL?.trim() || 'info',
|
|
transport: {
|
|
target: 'pino-pretty',
|
|
options: {
|
|
colorize: true,
|
|
translateTime: 'SYS:standard',
|
|
ignore: 'pid,hostname',
|
|
},
|
|
},
|
|
};
|
|
}
|
|
|
|
function normalizeS3Prefix(prefix: string | undefined): string {
|
|
const value = (prefix || 'openreader').trim();
|
|
return value ? value.replace(/^\/+|\/+$/g, '') : 'openreader';
|
|
}
|
|
|
|
function sanitizeNamespace(namespace: string | null): string | null {
|
|
if (!namespace) return null;
|
|
return SAFE_NAMESPACE_REGEX.test(namespace) ? namespace : null;
|
|
}
|
|
|
|
function documentParsedKey(id: string, namespace: string | null, prefix: string): string {
|
|
if (!DOCUMENT_ID_REGEX.test(id)) {
|
|
throw new Error(`Invalid document id: ${id}`);
|
|
}
|
|
const ns = sanitizeNamespace(namespace);
|
|
const nsSegment = ns ? `ns/${ns}/` : '';
|
|
return `${prefix}/documents_v1/parsed_v1/${nsSegment}${id}.json`;
|
|
}
|
|
|
|
function buildS3Client(): S3Client {
|
|
const bucket = requireEnv('S3_BUCKET');
|
|
const region = requireEnv('S3_REGION');
|
|
const accessKeyId = requireEnv('S3_ACCESS_KEY_ID');
|
|
const secretAccessKey = requireEnv('S3_SECRET_ACCESS_KEY');
|
|
const endpoint = process.env.S3_ENDPOINT?.trim() || undefined;
|
|
const forcePathStyle = parseBoolEnv('S3_FORCE_PATH_STYLE', false);
|
|
|
|
void bucket;
|
|
|
|
return new S3Client({
|
|
region,
|
|
endpoint,
|
|
forcePathStyle,
|
|
requestChecksumCalculation: 'WHEN_REQUIRED',
|
|
responseChecksumValidation: 'WHEN_REQUIRED',
|
|
credentials: {
|
|
accessKeyId,
|
|
secretAccessKey,
|
|
},
|
|
});
|
|
}
|
|
|
|
async function bodyToBuffer(body: unknown): Promise<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 Buffer.from(await maybe.transformToByteArray());
|
|
}
|
|
}
|
|
if (typeof body === 'object' && body !== null && 'on' in body) {
|
|
const stream = body as NodeJS.ReadableStream;
|
|
const chunks: Buffer[] = [];
|
|
for await (const chunk of stream) {
|
|
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 isAuthed(request: FastifyRequest, expectedToken: string): boolean {
|
|
const auth = request.headers.authorization;
|
|
if (!auth?.startsWith('Bearer ')) return false;
|
|
const token = auth.slice('Bearer '.length).trim();
|
|
return token === expectedToken;
|
|
}
|
|
|
|
function safeDurationMs(start: number | undefined, end: number | undefined): number | undefined {
|
|
if (!Number.isFinite(start) || !Number.isFinite(end)) return undefined;
|
|
return Math.max(0, Math.floor((end as number) - (start as number)));
|
|
}
|
|
|
|
function toErrorMessage(error: unknown): string {
|
|
if (error instanceof Error && error.message) return error.message;
|
|
return String(error);
|
|
}
|
|
|
|
function isAlreadyExistsError(error: unknown): boolean {
|
|
const message = toErrorMessage(error).toLowerCase();
|
|
return message.includes('already in use') || message.includes('already exists');
|
|
}
|
|
|
|
function isCasConflictError(error: unknown): boolean {
|
|
const message = toErrorMessage(error).toLowerCase();
|
|
return message.includes('wrong last sequence') || message.includes('key exists') || message.includes('wrong last');
|
|
}
|
|
|
|
function createJsonCodec<T>(): JsonCodec<T> {
|
|
const encoder = new TextEncoder();
|
|
const decoder = new TextDecoder();
|
|
return {
|
|
encode(value: T): Uint8Array {
|
|
return encoder.encode(JSON.stringify(value));
|
|
},
|
|
decode(data: Uint8Array): T {
|
|
return JSON.parse(decoder.decode(data)) as T;
|
|
},
|
|
};
|
|
}
|
|
|
|
function sleep(ms: number): Promise<void> {
|
|
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<void> {
|
|
if (this.inFlight < this.maxInFlight) {
|
|
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);
|
|
const next = this.queue.shift();
|
|
if (next) next();
|
|
}
|
|
}
|
|
|
|
function isTerminalStatus(status: WorkerJobState): boolean {
|
|
return status === 'succeeded' || status === 'failed';
|
|
}
|
|
|
|
function hashOpKey(opKey: string): string {
|
|
return createHash('sha256').update(opKey).digest('hex');
|
|
}
|
|
|
|
function opIndexKvKey(opKey: string): string {
|
|
return `op_index.${hashOpKey(opKey)}`;
|
|
}
|
|
|
|
function opStateKvKey(opId: string): string {
|
|
return `op_state.${opId}`;
|
|
}
|
|
|
|
function jobStateKvKey(jobId: string): string {
|
|
return `job_state.${jobId}`;
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
const alignSchema = 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 layoutSchema = 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),
|
|
});
|
|
|
|
const operationCreateSchema = z.discriminatedUnion('kind', [
|
|
z.object({
|
|
kind: z.literal('whisper_align'),
|
|
opKey: z.string().trim().min(1).max(1024),
|
|
payload: alignSchema,
|
|
}),
|
|
z.object({
|
|
kind: z.literal('pdf_layout'),
|
|
opKey: z.string().trim().min(1).max(1024),
|
|
payload: layoutSchema,
|
|
}),
|
|
]);
|
|
|
|
async function ensureJetStreamResources(
|
|
jsm: JetStreamManager,
|
|
whisperTimeoutMs: number,
|
|
pdfTimeoutMs: number,
|
|
pdfAttempts: number,
|
|
maxBytes: number,
|
|
): Promise<void> {
|
|
const streamConfig = {
|
|
name: JOBS_STREAM_NAME,
|
|
subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT],
|
|
retention: RetentionPolicy.Workqueue,
|
|
storage: StorageType.File,
|
|
max_bytes: maxBytes,
|
|
};
|
|
|
|
try {
|
|
await jsm.streams.add(streamConfig);
|
|
} catch (error) {
|
|
if (!isAlreadyExistsError(error)) throw error;
|
|
await jsm.streams.update(JOBS_STREAM_NAME, {
|
|
subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT],
|
|
max_bytes: maxBytes,
|
|
});
|
|
}
|
|
|
|
const ensureConsumer = async (
|
|
name: string,
|
|
subject: string,
|
|
ackWaitMs: number,
|
|
maxDeliver: number,
|
|
): Promise<void> => {
|
|
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 jsm.consumers.add(JOBS_STREAM_NAME, config);
|
|
} catch (error) {
|
|
if (!isAlreadyExistsError(error)) throw error;
|
|
await jsm.consumers.update(JOBS_STREAM_NAME, name, {
|
|
filter_subject: subject,
|
|
ack_wait: nanos(Math.max(ackWaitMs, 1_000)),
|
|
max_deliver: maxDeliver,
|
|
});
|
|
}
|
|
};
|
|
|
|
await Promise.all([
|
|
ensureConsumer(WHISPER_CONSUMER_NAME, WHISPER_JOBS_SUBJECT, whisperTimeoutMs + 15_000, WHISPER_MAX_DELIVER),
|
|
ensureConsumer(LAYOUT_CONSUMER_NAME, LAYOUT_JOBS_SUBJECT, pdfTimeoutMs + 15_000, pdfAttempts),
|
|
]);
|
|
}
|
|
|
|
async function main(): Promise<void> {
|
|
const port = readIntEnv('PORT', 8081);
|
|
const host = process.env.COMPUTE_WORKER_HOST?.trim() || '0.0.0.0';
|
|
const workerToken = requireEnv('COMPUTE_WORKER_TOKEN');
|
|
const natsUrl = requireEnv('NATS_URL');
|
|
const timeoutConfig = getComputeTimeoutConfig();
|
|
|
|
const jobConcurrency = readIntEnv('COMPUTE_JOB_CONCURRENCY', 1);
|
|
const whisperTimeoutMs = timeoutConfig.whisperTimeoutMs;
|
|
const pdfTimeoutMs = timeoutConfig.pdfTimeoutMs;
|
|
const pdfHardCapMs = timeoutConfig.pdfHardCapMs;
|
|
const pdfAttempts = readIntEnv('COMPUTE_PDF_JOB_ATTEMPTS', 1);
|
|
const prewarmModels = parseBoolEnv('COMPUTE_PREWARM_MODELS', true);
|
|
const jobsStreamMaxBytes = readIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024);
|
|
const jobStatesMaxBytes = readIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024);
|
|
const opStaleMs = readIntEnv(
|
|
'COMPUTE_OP_STALE_MS',
|
|
Math.max(30 * 60_000, Math.max(whisperTimeoutMs, pdfTimeoutMs) * 4),
|
|
);
|
|
|
|
const connectOpts: any = { servers: natsUrl };
|
|
const natsCreds = process.env.NATS_CREDS?.trim();
|
|
const natsCredsFile = process.env.NATS_CREDS_FILE?.trim();
|
|
|
|
if (natsCreds) {
|
|
console.log('[compute-worker] Connecting to NATS using credentials string from NATS_CREDS');
|
|
connectOpts.authenticator = credsAuthenticator(new TextEncoder().encode(natsCreds));
|
|
} else if (natsCredsFile) {
|
|
console.log(`[compute-worker] Connecting to NATS using credentials file: ${natsCredsFile}`);
|
|
const { readFileSync } = await import('node:fs');
|
|
const credsData = readFileSync(natsCredsFile);
|
|
connectOpts.authenticator = credsAuthenticator(credsData);
|
|
}
|
|
|
|
const nc: NatsConnection = await connect(connectOpts);
|
|
const js: JetStreamClient = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS });
|
|
const jsm: JetStreamManager = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS });
|
|
|
|
await ensureJetStreamResources(jsm, whisperTimeoutMs, pdfTimeoutMs, pdfAttempts, jobsStreamMaxBytes);
|
|
|
|
const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, {
|
|
history: 1,
|
|
ttl: COMPUTE_STATE_TTL_MS,
|
|
max_bytes: jobStatesMaxBytes,
|
|
});
|
|
|
|
const s3 = buildS3Client();
|
|
const s3Bucket = requireEnv('S3_BUCKET');
|
|
const s3Prefix = normalizeS3Prefix(process.env.S3_PREFIX);
|
|
|
|
const ensureSafeKey = (key: string): string => {
|
|
const trimmed = key.trim();
|
|
if (!trimmed.startsWith(`${s3Prefix}/`)) {
|
|
throw new Error('Object key prefix mismatch');
|
|
}
|
|
return trimmed;
|
|
};
|
|
|
|
const readObjectByKey = async (key: string): Promise<ArrayBuffer> => {
|
|
const safeKey = ensureSafeKey(key);
|
|
const response = await s3.send(new GetObjectCommand({
|
|
Bucket: s3Bucket,
|
|
Key: safeKey,
|
|
}));
|
|
const bytes = await bodyToBuffer(response.Body);
|
|
return toArrayBuffer(new Uint8Array(bytes));
|
|
};
|
|
|
|
const putParsedObject = async (documentId: string, namespace: string | null, parsed: unknown): Promise<string> => {
|
|
const key = documentParsedKey(documentId, namespace, s3Prefix);
|
|
const body = Buffer.from(JSON.stringify(parsed));
|
|
await s3.send(new PutObjectCommand({
|
|
Bucket: s3Bucket,
|
|
Key: key,
|
|
Body: body,
|
|
ContentType: 'application/json',
|
|
ServerSideEncryption: 'AES256',
|
|
}));
|
|
return key;
|
|
};
|
|
|
|
if (prewarmModels) {
|
|
await ensureComputeModels();
|
|
}
|
|
|
|
const app = Fastify({ logger: buildLoggerConfig() });
|
|
app.log.info({
|
|
jobConcurrency,
|
|
whisperTimeoutMs,
|
|
pdfTimeoutMs,
|
|
pdfAttempts,
|
|
opStaleMs,
|
|
availableCpuCores: getAvailableCpuCores(),
|
|
onnxThreadsPerJob: getOnnxThreadsPerJob(jobConcurrency),
|
|
natsApiTimeoutMs: NATS_API_TIMEOUT_MS,
|
|
pdfLayoutHardCapMs: pdfHardCapMs,
|
|
}, 'compute runtime config');
|
|
|
|
const opIndexCodec = createJsonCodec<OpIndexEntry>();
|
|
const opStateCodec = createJsonCodec<WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>>();
|
|
const whisperJobCodec = createJsonCodec<QueuedJob<WhisperAlignJobRequest>>();
|
|
const layoutJobCodec = createJsonCodec<QueuedJob<PdfLayoutJobRequest>>();
|
|
const jobStateCodec = createJsonCodec<StoredJobState<WhisperAlignJobResult | PdfLayoutJobResult>>();
|
|
|
|
const putOpState = async (state: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>): Promise<void> => {
|
|
await kv.put(opStateKvKey(state.opId), opStateCodec.encode(state));
|
|
};
|
|
|
|
const getOpState = async (opId: string): Promise<WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> | null> => {
|
|
const entry = await kv.get(opStateKvKey(opId));
|
|
if (!entry || entry.operation !== 'PUT') return null;
|
|
return opStateCodec.decode(entry.value);
|
|
};
|
|
|
|
const putJobState = async (state: StoredJobState<WhisperAlignJobResult | PdfLayoutJobResult>): Promise<void> => {
|
|
await kv.put(jobStateKvKey(state.jobId), jobStateCodec.encode(state));
|
|
};
|
|
|
|
const publishQueuedJob = async (
|
|
op: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>,
|
|
payload: WhisperAlignJobRequest | PdfLayoutJobRequest,
|
|
): Promise<void> => {
|
|
if (op.kind === 'whisper_align') {
|
|
await js.publish(WHISPER_JOBS_SUBJECT, whisperJobCodec.encode({
|
|
jobId: op.jobId,
|
|
opId: op.opId,
|
|
opKey: op.opKey,
|
|
kind: 'whisper_align',
|
|
queuedAt: op.queuedAt,
|
|
payload: payload as WhisperAlignJobRequest,
|
|
}));
|
|
return;
|
|
}
|
|
|
|
await js.publish(LAYOUT_JOBS_SUBJECT, layoutJobCodec.encode({
|
|
jobId: op.jobId,
|
|
opId: op.opId,
|
|
opKey: op.opKey,
|
|
kind: 'pdf_layout',
|
|
queuedAt: op.queuedAt,
|
|
payload: payload as PdfLayoutJobRequest,
|
|
}));
|
|
};
|
|
|
|
const enqueueOrReuseOperation = async (
|
|
req: WorkerOperationRequest,
|
|
): Promise<WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>> => {
|
|
const opKey = req.opKey.trim();
|
|
const indexKey = opIndexKvKey(opKey);
|
|
|
|
for (let attemptNo = 0; attemptNo < 10; attemptNo += 1) {
|
|
const indexEntry = await kv.get(indexKey);
|
|
if (indexEntry && indexEntry.operation === 'PUT') {
|
|
const pointer = opIndexCodec.decode(indexEntry.value);
|
|
const current = await getOpState(pointer.opId);
|
|
|
|
if (!current) {
|
|
await sleep(25);
|
|
continue;
|
|
}
|
|
|
|
if (current && current.kind === req.kind) {
|
|
const ageMs = Date.now() - current.updatedAt;
|
|
if (current.status === 'succeeded') return current;
|
|
if ((current.status === 'queued' || current.status === 'running') && ageMs <= opStaleMs) {
|
|
return current;
|
|
}
|
|
}
|
|
|
|
const now = Date.now();
|
|
const replacement: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
opId: crypto.randomUUID(),
|
|
opKey,
|
|
kind: req.kind,
|
|
jobId: crypto.randomUUID(),
|
|
status: 'queued',
|
|
queuedAt: now,
|
|
updatedAt: now,
|
|
};
|
|
|
|
try {
|
|
await kv.update(indexKey, opIndexCodec.encode({ opId: replacement.opId }), indexEntry.revision);
|
|
} catch (error) {
|
|
if (isCasConflictError(error)) continue;
|
|
throw error;
|
|
}
|
|
|
|
await putOpState(replacement);
|
|
await putJobState({
|
|
jobId: replacement.jobId,
|
|
opId: replacement.opId,
|
|
opKey: replacement.opKey,
|
|
kind: replacement.kind,
|
|
status: 'queued',
|
|
timestamp: now,
|
|
updatedAt: now,
|
|
});
|
|
|
|
try {
|
|
await publishQueuedJob(replacement, req.payload);
|
|
return replacement;
|
|
} catch (error) {
|
|
const failed: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
...replacement,
|
|
status: 'failed',
|
|
updatedAt: Date.now(),
|
|
error: { message: toErrorMessage(error) },
|
|
};
|
|
await putOpState(failed);
|
|
await putJobState({
|
|
jobId: replacement.jobId,
|
|
opId: replacement.opId,
|
|
opKey: replacement.opKey,
|
|
kind: replacement.kind,
|
|
status: 'failed',
|
|
timestamp: replacement.queuedAt,
|
|
updatedAt: failed.updatedAt,
|
|
error: failed.error,
|
|
});
|
|
return failed;
|
|
}
|
|
}
|
|
|
|
const now = Date.now();
|
|
const created: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
opId: crypto.randomUUID(),
|
|
opKey,
|
|
kind: req.kind,
|
|
jobId: crypto.randomUUID(),
|
|
status: 'queued',
|
|
queuedAt: now,
|
|
updatedAt: now,
|
|
};
|
|
|
|
try {
|
|
await kv.create(indexKey, opIndexCodec.encode({ opId: created.opId }));
|
|
} catch (error) {
|
|
if (isCasConflictError(error)) continue;
|
|
throw error;
|
|
}
|
|
|
|
await putOpState(created);
|
|
await putJobState({
|
|
jobId: created.jobId,
|
|
opId: created.opId,
|
|
opKey: created.opKey,
|
|
kind: created.kind,
|
|
status: 'queued',
|
|
timestamp: now,
|
|
updatedAt: now,
|
|
});
|
|
|
|
try {
|
|
await publishQueuedJob(created, req.payload);
|
|
return created;
|
|
} catch (error) {
|
|
const failed: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
...created,
|
|
status: 'failed',
|
|
updatedAt: Date.now(),
|
|
error: { message: toErrorMessage(error) },
|
|
};
|
|
await putOpState(failed);
|
|
await putJobState({
|
|
jobId: created.jobId,
|
|
opId: created.opId,
|
|
opKey: created.opKey,
|
|
kind: created.kind,
|
|
status: 'failed',
|
|
timestamp: created.queuedAt,
|
|
updatedAt: failed.updatedAt,
|
|
error: failed.error,
|
|
});
|
|
return failed;
|
|
}
|
|
}
|
|
|
|
throw new Error('Unable to reserve operation after repeated CAS conflicts');
|
|
};
|
|
|
|
app.addHook('onRequest', async (request, reply) => {
|
|
const path = request.url.split('?')[0] ?? request.url;
|
|
if (path === '/health/live' || path === '/health/ready') return;
|
|
if (!isAuthed(request, workerToken)) {
|
|
return reply.code(401).send({ error: 'Unauthorized' });
|
|
}
|
|
return;
|
|
});
|
|
|
|
app.get('/health/live', async () => ({ ok: true }));
|
|
|
|
app.get('/health/ready', async (_request, reply) => {
|
|
try {
|
|
await nc.flush();
|
|
return { ok: true };
|
|
} catch (error) {
|
|
reply.code(503);
|
|
return {
|
|
ok: false,
|
|
error: toErrorMessage(error),
|
|
};
|
|
}
|
|
});
|
|
|
|
app.post('/ops', async (request, reply) => {
|
|
const parsed = operationCreateSchema.safeParse(request.body);
|
|
if (!parsed.success) {
|
|
reply.code(400);
|
|
return {
|
|
error: 'Invalid request body',
|
|
issues: parsed.error.issues,
|
|
};
|
|
}
|
|
|
|
const op = await enqueueOrReuseOperation(parsed.data as WorkerOperationRequest);
|
|
reply.code(202);
|
|
return op;
|
|
});
|
|
|
|
app.get('/ops/:opId', async (request, reply) => {
|
|
const params = z.object({ opId: z.string().trim().min(1) }).safeParse(request.params);
|
|
if (!params.success) {
|
|
reply.code(400);
|
|
return { error: 'Invalid op id' };
|
|
}
|
|
|
|
const state = await getOpState(params.data.opId);
|
|
if (!state) {
|
|
reply.code(404);
|
|
return { error: 'Operation not found' };
|
|
}
|
|
|
|
return state;
|
|
});
|
|
|
|
app.get('/ops/:opId/events', async (request, reply) => {
|
|
const params = z.object({ opId: z.string().trim().min(1) }).safeParse(request.params);
|
|
if (!params.success) {
|
|
reply.code(400);
|
|
return { error: 'Invalid op id' };
|
|
}
|
|
|
|
const initial = await getOpState(params.data.opId);
|
|
if (!initial) {
|
|
reply.code(404);
|
|
return { error: 'Operation not found' };
|
|
}
|
|
|
|
reply.hijack();
|
|
reply.raw.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
|
|
reply.raw.setHeader('Cache-Control', 'no-cache, no-transform');
|
|
reply.raw.setHeader('Connection', 'keep-alive');
|
|
reply.raw.setHeader('X-Accel-Buffering', 'no');
|
|
|
|
const writeSnapshot = (snapshot: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>): void => {
|
|
if (closed || reply.raw.writableEnded) return;
|
|
reply.raw.write(`event: snapshot\ndata: ${JSON.stringify(snapshot)}\n\n`);
|
|
};
|
|
|
|
let closed = false;
|
|
request.raw.on('close', () => {
|
|
closed = true;
|
|
});
|
|
|
|
try {
|
|
let current = initial;
|
|
writeSnapshot(current);
|
|
let signature = JSON.stringify(current);
|
|
|
|
while (!closed && !isTerminalStatus(current.status)) {
|
|
await sleep(SSE_POLL_INTERVAL_MS);
|
|
const next = await getOpState(params.data.opId);
|
|
if (!next) break;
|
|
const nextSignature = JSON.stringify(next);
|
|
if (nextSignature !== signature) {
|
|
current = next;
|
|
signature = nextSignature;
|
|
writeSnapshot(current);
|
|
}
|
|
}
|
|
} catch (error) {
|
|
app.log.warn({
|
|
opId: params.data.opId,
|
|
error: toErrorMessage(error),
|
|
}, 'op events stream loop error');
|
|
}
|
|
|
|
if (!reply.raw.writableEnded) {
|
|
reply.raw.end();
|
|
}
|
|
|
|
return reply;
|
|
});
|
|
|
|
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;
|
|
|
|
const runWhisper = async (
|
|
payload: WhisperAlignJobRequest,
|
|
queueWaitMs: number,
|
|
): Promise<WhisperAlignJobResult> => {
|
|
const parsed = alignSchema.parse(payload);
|
|
|
|
const s3FetchStartedAt = Date.now();
|
|
const audioBuffer = await readObjectByKey(parsed.audioObjectKey);
|
|
const s3FetchMs = Date.now() - s3FetchStartedAt;
|
|
|
|
const computeStartedAt = Date.now();
|
|
const result = await withTimeout(
|
|
runWhisperAlignmentFromAudioBuffer({
|
|
audioBuffer,
|
|
text: parsed.text,
|
|
cacheKey: parsed.cacheKey,
|
|
lang: parsed.lang,
|
|
}),
|
|
whisperTimeoutMs,
|
|
'whisper alignment job',
|
|
);
|
|
|
|
const computeMs = Date.now() - computeStartedAt;
|
|
return {
|
|
...result,
|
|
timing: {
|
|
queueWaitMs,
|
|
s3FetchMs,
|
|
computeMs,
|
|
},
|
|
};
|
|
};
|
|
|
|
const runLayout = async (
|
|
payload: PdfLayoutJobRequest,
|
|
queueWaitMs: number,
|
|
hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise<void> },
|
|
): Promise<PdfLayoutJobResult> => {
|
|
const parsed = layoutSchema.parse(payload);
|
|
|
|
const s3FetchStartedAt = Date.now();
|
|
const pdfBytes = await readObjectByKey(parsed.documentObjectKey);
|
|
const s3FetchMs = Date.now() - s3FetchStartedAt;
|
|
|
|
let lastTotalPages = 0;
|
|
let lastPagesParsed = 0;
|
|
const computeStartedAt = Date.now();
|
|
const result = await withIdleTimeoutAndHardCap({
|
|
idleTimeoutMs: Math.max(pdfTimeoutMs, 1_000),
|
|
hardCapMs: pdfHardCapMs,
|
|
label: 'pdf layout job',
|
|
run: async (touchProgress) => runPdfLayoutFromPdfBuffer({
|
|
documentId: parsed.documentId,
|
|
pdfBytes,
|
|
onPageParsed: async ({ pageNumber, totalPages }) => {
|
|
touchProgress();
|
|
lastTotalPages = totalPages;
|
|
lastPagesParsed = pageNumber;
|
|
if (!hooks?.onProgress) return;
|
|
await hooks.onProgress({
|
|
totalPages,
|
|
pagesParsed: pageNumber,
|
|
currentPage: pageNumber,
|
|
phase: 'infer',
|
|
});
|
|
},
|
|
}),
|
|
});
|
|
|
|
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 putParsedObject(parsed.documentId, parsed.namespace, result.parsed);
|
|
return {
|
|
parsedObjectKey,
|
|
timing: {
|
|
queueWaitMs,
|
|
s3FetchMs,
|
|
computeMs,
|
|
},
|
|
};
|
|
};
|
|
|
|
async function processMessage<TPayload, TResult>(input: {
|
|
msg: JsMsg;
|
|
codec: JsonCodec<QueuedJob<TPayload>>;
|
|
run: (
|
|
payload: TPayload,
|
|
queueWaitMs: number,
|
|
hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise<void> },
|
|
) => Promise<TResult>;
|
|
workerLabel: string;
|
|
}): Promise<void> {
|
|
let decoded: QueuedJob<TPayload> | null = null;
|
|
let heartbeat: NodeJS.Timeout | null = null;
|
|
let latestProgress: PdfLayoutProgress | undefined;
|
|
try {
|
|
decoded = input.codec.decode(input.msg.data);
|
|
const startedAt = Date.now();
|
|
const queueWaitMs = safeDurationMs(decoded.queuedAt, startedAt);
|
|
|
|
const runningState: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
opId: decoded.opId,
|
|
opKey: decoded.opKey,
|
|
kind: decoded.kind,
|
|
jobId: decoded.jobId,
|
|
status: 'running',
|
|
queuedAt: decoded.queuedAt,
|
|
startedAt,
|
|
updatedAt: startedAt,
|
|
...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
};
|
|
|
|
await putOpState(runningState);
|
|
await putJobState({
|
|
jobId: decoded.jobId,
|
|
opId: decoded.opId,
|
|
opKey: decoded.opKey,
|
|
kind: decoded.kind,
|
|
status: 'running',
|
|
timestamp: decoded.queuedAt,
|
|
startedAt,
|
|
updatedAt: startedAt,
|
|
...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
});
|
|
|
|
const persistRunningState = async (updatedAt: number): Promise<void> => {
|
|
const runningOpState: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
opId: decoded!.opId,
|
|
opKey: decoded!.opKey,
|
|
kind: decoded!.kind,
|
|
jobId: decoded!.jobId,
|
|
status: 'running',
|
|
queuedAt: decoded!.queuedAt,
|
|
startedAt,
|
|
updatedAt,
|
|
...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
};
|
|
|
|
await putOpState(runningOpState);
|
|
await putJobState({
|
|
jobId: decoded!.jobId,
|
|
opId: decoded!.opId,
|
|
opKey: decoded!.opKey,
|
|
kind: decoded!.kind,
|
|
status: 'running',
|
|
timestamp: decoded!.queuedAt,
|
|
startedAt,
|
|
updatedAt,
|
|
...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
});
|
|
};
|
|
|
|
heartbeat = setInterval(() => {
|
|
const now = Date.now();
|
|
void persistRunningState(now).catch((stateError) => {
|
|
app.log.error({
|
|
worker: input.workerLabel,
|
|
opId: decoded?.opId,
|
|
jobId: decoded?.jobId,
|
|
error: toErrorMessage(stateError),
|
|
}, 'failed to persist running heartbeat state');
|
|
});
|
|
}, RUNNING_HEARTBEAT_MS);
|
|
|
|
const result = await input.run(decoded.payload, queueWaitMs ?? 0, {
|
|
onProgress: async (progress) => {
|
|
latestProgress = progress;
|
|
await persistRunningState(Date.now());
|
|
},
|
|
});
|
|
const resultTiming = result && typeof result === 'object' && 'timing' in result
|
|
? (result as { timing?: WorkerJobTiming }).timing
|
|
: undefined;
|
|
const now = Date.now();
|
|
|
|
const succeededState: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
opId: decoded.opId,
|
|
opKey: decoded.opKey,
|
|
kind: decoded.kind,
|
|
jobId: decoded.jobId,
|
|
status: 'succeeded',
|
|
queuedAt: decoded.queuedAt,
|
|
startedAt,
|
|
updatedAt: now,
|
|
result: result as WhisperAlignJobResult | PdfLayoutJobResult,
|
|
...(resultTiming ? { timing: resultTiming } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
};
|
|
|
|
await putOpState(succeededState);
|
|
await putJobState({
|
|
jobId: decoded.jobId,
|
|
opId: decoded.opId,
|
|
opKey: decoded.opKey,
|
|
kind: decoded.kind,
|
|
status: 'succeeded',
|
|
timestamp: decoded.queuedAt,
|
|
startedAt,
|
|
updatedAt: now,
|
|
result: result as WhisperAlignJobResult | PdfLayoutJobResult,
|
|
...(resultTiming ? { timing: resultTiming } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
});
|
|
|
|
input.msg.ack();
|
|
app.log.info({
|
|
worker: input.workerLabel,
|
|
opId: decoded.opId,
|
|
jobId: decoded.jobId,
|
|
resultRef: extractResultRef(decoded.kind, result),
|
|
timing: resultTiming,
|
|
}, 'job succeeded');
|
|
} catch (error) {
|
|
const message = toErrorMessage(error);
|
|
const deliveryCount = input.msg.info.deliveryCount;
|
|
const isWhisperAlign = decoded?.kind === 'whisper_align';
|
|
const maxAttempts = isWhisperAlign ? WHISPER_MAX_DELIVER : pdfAttempts;
|
|
const hasRetriesLeft = !isWhisperAlign && deliveryCount < maxAttempts;
|
|
|
|
if (decoded) {
|
|
const now = Date.now();
|
|
const queueWaitMs = safeDurationMs(decoded.queuedAt, now);
|
|
|
|
const status: WorkerJobState = hasRetriesLeft ? 'running' : 'failed';
|
|
const opState: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
|
opId: decoded.opId,
|
|
opKey: decoded.opKey,
|
|
kind: decoded.kind,
|
|
jobId: decoded.jobId,
|
|
status,
|
|
queuedAt: decoded.queuedAt,
|
|
updatedAt: now,
|
|
...(status === 'failed' ? { error: { message } } : {}),
|
|
...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
};
|
|
|
|
await putOpState(opState).catch((stateError) => {
|
|
app.log.error({
|
|
worker: input.workerLabel,
|
|
opId: decoded?.opId,
|
|
jobId: decoded?.jobId,
|
|
error: toErrorMessage(stateError),
|
|
}, 'failed to persist operation state');
|
|
});
|
|
|
|
await putJobState({
|
|
jobId: decoded.jobId,
|
|
opId: decoded.opId,
|
|
opKey: decoded.opKey,
|
|
kind: decoded.kind,
|
|
status,
|
|
timestamp: decoded.queuedAt,
|
|
updatedAt: now,
|
|
...(status === 'failed' ? { error: { message } } : {}),
|
|
...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}),
|
|
...(latestProgress ? { progress: latestProgress } : {}),
|
|
}).catch((stateError) => {
|
|
app.log.error({
|
|
worker: input.workerLabel,
|
|
opId: decoded?.opId,
|
|
jobId: decoded?.jobId,
|
|
error: toErrorMessage(stateError),
|
|
}, 'failed to persist job state');
|
|
});
|
|
}
|
|
|
|
if (hasRetriesLeft) {
|
|
input.msg.nak();
|
|
app.log.error({
|
|
worker: input.workerLabel,
|
|
opId: decoded?.opId,
|
|
jobId: decoded?.jobId,
|
|
error: message,
|
|
deliveryCount,
|
|
maxAttempts,
|
|
}, 'job failed, nacked for retry');
|
|
} else {
|
|
input.msg.term(message);
|
|
app.log.error({
|
|
worker: input.workerLabel,
|
|
opId: decoded?.opId,
|
|
jobId: decoded?.jobId,
|
|
error: message,
|
|
deliveryCount,
|
|
maxAttempts,
|
|
retrySuppressed: isWhisperAlign ? 'whisper_align' : undefined,
|
|
}, 'job failed, max attempts reached');
|
|
}
|
|
} finally {
|
|
if (heartbeat) clearInterval(heartbeat);
|
|
}
|
|
}
|
|
|
|
async function createWorkerLoop<TPayload, TResult>(input: {
|
|
consumer: Consumer;
|
|
codec: JsonCodec<QueuedJob<TPayload>>;
|
|
run: (
|
|
payload: TPayload,
|
|
queueWaitMs: number,
|
|
hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise<void> },
|
|
) => Promise<TResult>;
|
|
workerLabel: string;
|
|
}): Promise<void> {
|
|
while (!stopping) {
|
|
await jobGate.acquire();
|
|
if (stopping) {
|
|
jobGate.release();
|
|
return;
|
|
}
|
|
|
|
let msg: JsMsg | null = null;
|
|
try {
|
|
try {
|
|
msg = await input.consumer.next();
|
|
} 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<void>[] = [];
|
|
|
|
for (let i = 0; i < jobConcurrency; i += 1) {
|
|
workerLoops.push(createWorkerLoop({
|
|
consumer: whisperConsumer,
|
|
codec: whisperJobCodec,
|
|
run: runWhisper,
|
|
workerLabel: `whisper-${i + 1}`,
|
|
}));
|
|
}
|
|
|
|
for (let i = 0; i < jobConcurrency; i += 1) {
|
|
workerLoops.push(createWorkerLoop({
|
|
consumer: layoutConsumer,
|
|
codec: layoutJobCodec,
|
|
run: runLayout,
|
|
workerLabel: `layout-${i + 1}`,
|
|
}));
|
|
}
|
|
|
|
const close = async (): Promise<void> => {
|
|
if (stopping) return;
|
|
stopping = true;
|
|
await app.close();
|
|
await Promise.allSettled(workerLoops);
|
|
await nc.drain();
|
|
};
|
|
|
|
process.once('SIGINT', () => {
|
|
void close().finally(() => process.exit(0));
|
|
});
|
|
|
|
process.once('SIGTERM', () => {
|
|
void close().finally(() => process.exit(0));
|
|
});
|
|
|
|
await app.listen({ host, port });
|
|
app.log.info({ host, port }, 'compute worker listening');
|
|
}
|
|
|
|
void main().catch((error) => {
|
|
console.error('[compute-worker] fatal startup error', error);
|
|
process.exit(1);
|
|
});
|