From 5403aa00cc592751d69fce0351adf4e04d5fbbf8 Mon Sep 17 00:00:00 2001 From: Richard R Date: Sat, 13 Jun 2026 05:04:24 -0600 Subject: [PATCH] refactor(compute-worker): modularize operation contracts and API, update protocol naming - Move operation contracts from api/contracts to operations/contracts for clearer separation - Delete legacy API modules: contracts.ts, operation-keys.ts, public-operation.ts - Add new modular files: compute-operation.ts, http-hooks.ts, routes.ts, operations/keys.ts, operations/reconciliation.ts - Refactor all imports to use new operations/contracts and keys modules - Rename protocol types from PublicOperation to ComputeOperation for clarity - Update OpenAPI schema and all references to use ComputeOperation/Event naming - Extract whisper decoder helpers to inference/whisper/decoder.ts - Add new infra: nats-session.ts for NATS session management - Update tests and storage logic to use new contracts and keys modules - Refactor worker loop and orchestrator to use modular job definitions - Update client and protocol in src/lib/server/compute-worker to use new ComputeOperation types This refactor modularizes compute-worker operation logic, clarifies type naming, and improves maintainability by separating contracts, keys, and reconciliation logic. No breaking changes to external API, but all internal references and tests are updated for the new structure. --- compute-worker/openapi.json | 8 +- compute-worker/src/api/app.ts | 865 ++---------------- ...blic-operation.ts => compute-operation.ts} | 18 +- compute-worker/src/api/http-hooks.ts | 80 ++ compute-worker/src/api/routes.ts | 356 +++++++ compute-worker/src/api/schemas.ts | 8 +- compute-worker/src/inference/pdf/parse.ts | 2 +- compute-worker/src/inference/whisper/align.ts | 117 +-- .../src/inference/whisper/decoder.ts | 111 +++ compute-worker/src/infrastructure/config.ts | 36 +- .../src/infrastructure/nats-adapters.ts | 2 +- .../src/infrastructure/nats-session.ts | 204 +++++ compute-worker/src/jobs/handlers.ts | 2 +- compute-worker/src/jobs/pdf-progress.ts | 2 +- compute-worker/src/jobs/worker-loop-policy.ts | 2 +- compute-worker/src/jobs/worker-loop.ts | 39 +- .../src/{api => operations}/contracts.ts | 6 +- .../operation-keys.ts => operations/keys.ts} | 2 +- .../src/operations/reconciliation.ts | 95 ++ compute-worker/src/operations/recovery.ts | 2 +- compute-worker/src/operations/service.ts | 2 +- .../src/operations/state-machine.ts | 2 +- compute-worker/src/operations/types.ts | 2 +- .../src/storage/artifact-addressing.ts | 4 +- .../algorithms/whisper-decoder.test.ts | 61 ++ .../control-plane/orchestrator.test.ts | 2 +- .../control-plane/state-machine.test.ts | 2 +- .../helpers/in-memory-control-plane.ts | 2 +- .../tests/fixtures/fake-control-plane.ts | 2 +- .../tests/unit/jetstream-adapters.test.ts | 2 +- .../tests/unit/key-and-progress.test.ts | 2 +- compute-worker/tests/unit/worker-loop.test.ts | 170 ++++ src/lib/server/compute-worker/client.ts | 20 +- src/lib/server/compute-worker/generated.ts | 6 +- src/lib/server/compute-worker/protocol.ts | 22 +- src/lib/server/pdf-parse/operation.ts | 10 +- src/lib/server/pdf-parse/snapshot.ts | 8 +- 37 files changed, 1313 insertions(+), 963 deletions(-) rename compute-worker/src/api/{public-operation.ts => compute-operation.ts} (72%) create mode 100644 compute-worker/src/api/http-hooks.ts create mode 100644 compute-worker/src/api/routes.ts create mode 100644 compute-worker/src/inference/whisper/decoder.ts create mode 100644 compute-worker/src/infrastructure/nats-session.ts rename compute-worker/src/{api => operations}/contracts.ts (95%) rename compute-worker/src/{api/operation-keys.ts => operations/keys.ts} (96%) create mode 100644 compute-worker/src/operations/reconciliation.ts create mode 100644 compute-worker/tests/compute/algorithms/whisper-decoder.test.ts create mode 100644 compute-worker/tests/unit/worker-loop.test.ts diff --git a/compute-worker/openapi.json b/compute-worker/openapi.json index c0973f1..5b1f63a 100644 --- a/compute-worker/openapi.json +++ b/compute-worker/openapi.json @@ -272,7 +272,7 @@ ], "additionalProperties": false }, - "PublicOperation": { + "ComputeOperation": { "type": "object", "properties": { "opId": { @@ -413,7 +413,7 @@ ], "additionalProperties": false }, - "PublicOperationEvent": { + "ComputeOperationEvent": { "type": "object", "properties": { "eventId": { @@ -1706,12 +1706,12 @@ ], "responses": { "200": { - "description": "Server-sent PublicOperationEvent stream", + "description": "Server-sent ComputeOperationEvent stream", "content": { "application/json": { "schema": { "type": "string", - "description": "Server-sent PublicOperationEvent stream" + "description": "Server-sent ComputeOperationEvent stream" } } } diff --git a/compute-worker/src/api/app.ts b/compute-worker/src/api/app.ts index 77b5a55..7ea327c 100644 --- a/compute-worker/src/api/app.ts +++ b/compute-worker/src/api/app.ts @@ -1,18 +1,9 @@ -import Fastify, { type FastifyInstance, type FastifyRequest } from 'fastify'; +import Fastify, { type FastifyInstance } from 'fastify'; import swagger from '@fastify/swagger'; import { - connect, credsAuthenticator, - type NatsConnection, + type ConnectionOptions, } from '@nats-io/transport-node'; -import { - jetstream, - jetstreamManager, - type Consumer, - type JetStreamClient, - type JetStreamManager, -} from '@nats-io/jetstream'; -import { Kvm } from '@nats-io/kv'; import { ensureComputeModels, } from '../inference/runtime'; @@ -21,30 +12,26 @@ import { getComputeOpStaleMs, getAvailableCpuCores, getOnnxThreadsPerJob, + buildLoggerConfig, + normalizeNatsReplicas, + readBoolEnv, + readPositiveIntEnv, + requireEnv, } from '../infrastructure/config'; -import { encodeSseFrame, OperationOrchestrator } from '../operations'; +import { OperationOrchestrator } from '../operations'; import type { PdfLayoutJobRequest, PdfLayoutJobResult, - WorkerOperationEvent, WhisperAlignJobRequest, WhisperAlignJobResult, - WorkerJobTiming, - WorkerOperationRequest, - WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; import { JetStreamOperationEventStream, JetStreamOperationQueue, JetStreamOperationStateStore, - hashOpKey, } from '../infrastructure/nats-adapters'; import { createJsonCodec } from '../infrastructure/json-codec'; -import { - recoverOrphanedOperations, - type StreamedOperationState, -} from '../operations/recovery'; -import { parsedPdfArtifactKey } from '../storage/artifact-addressing'; +import { createOperationReconciler } from '../operations/reconciliation'; import { createArtifactStorage, createS3ClientFromEnv, @@ -53,44 +40,31 @@ import { } from '../infrastructure/storage'; import { createJobHandlers } from '../jobs/handlers'; import { createWorkerLoopController, type QueuedJob } from '../jobs/worker-loop'; +import { createNatsSessionManager } from '../infrastructure/nats-session'; import { - COMPUTE_STATE_BUCKET, - COMPUTE_STATE_TTL_MS, EVENTS_STREAM_NAME, - JOBS_STREAM_NAME, - LAYOUT_CONSUMER_NAME, LAYOUT_JOBS_SUBJECT, NATS_API_TIMEOUT_MS, - WHISPER_CONSUMER_NAME, WHISPER_JOBS_SUBJECT, - ensureJetStreamResources, } from '../infrastructure/nats'; -import { buildPdfOperationKey, buildWhisperOperationKey } from './operation-keys'; -import { toPublicOperation, type PublicOperationEvent } from './public-operation'; +import { registerHttpHooks } from './http-hooks'; +import { + registerComputeWorkerRoutes, + type ComputeWorkerRouteDeps, +} from './routes'; import { apiErrorResponseSchema, artifactReferenceSchema, jsonSchema, - operationEventsQuerySchema, - operationParamsSchema, operationErrorSchema, parsedPdfDocumentSchema, pdfLayoutProgressSchema, pdfLayoutResolutionSchema, - pdfOperationCreateSchema, - pdfResolveSchema, - publicOperationEventSchema, - publicOperationSchema, + computeOperationEventSchema, + computeOperationSchema, ttsSentenceAlignmentSchema, - whisperOperationCreateSchema, } from './schemas'; -const OP_EVENTS_KEEPALIVE_MS = 15_000; -// Reconnection delay handed to the browser EventSource via the SSE `retry:` -// directive. When a silent stream is torn down for idle sleep, this keeps the -// client from immediately reconnecting and re-waking the worker; instead it -// reconnects on a slow cadence so the container stays asleep most of the time. -const OP_EVENTS_RECONNECT_HINT_MS = 120_000; // Disconnect from NATS after this much continuous idle so the worker stops // generating outbound traffic (pull polling + keepalive PINGs) and Railway can // put it to sleep. Reconnect happens lazily on the next inbound request. @@ -100,76 +74,9 @@ const IDLE_STATUS_LOG_INTERVAL_MS = 60_000; const ORPHAN_SWEEP_INTERVAL_MS = 15_000; // Bounded pull window so consumer loops yield periodically and can be stopped // cleanly when going idle, instead of blocking on a long-lived pull. -const REQUEST_STARTED_AT_MS_KEY = Symbol('request-started-at-ms'); -const REQUEST_COUNTED_KEY = Symbol('request-activity-counted'); const WHISPER_MAX_DELIVER = 1; -interface NatsSession { - nc: NatsConnection; - js: JetStreamClient; - jsm: JetStreamManager; - kv: Awaited>; - whisperConsumer: Consumer; - layoutConsumer: Consumer; -} - -interface OperationEventStreamLike { - subscribe(input: { - opId: string; - sinceEventId?: number; - onEvent: (event: WorkerOperationEvent) => void | Promise; - onError?: (error: unknown) => void; - }): Promise<() => void>; -} - -interface OperationStateStoreLike { - getOpState(opId: string): Promise; - getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; - getOpIndex?(opKey: string): Promise<{ opId: string } | null>; - listOpStates?(): Promise; -} - -interface OrchestratorLike { - enqueueOrReuse(request: WorkerOperationRequest): Promise; - markRunning(input: { - opId: string; - startedAt?: number; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markProgress(input: { - opId: string; - progress: WorkerOperationState['progress']; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markSucceeded(input: { - opId: string; - result: unknown; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markFailed(input: { - opId: string; - error: { message: string; code?: string } | string; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markFailedIfUnchanged?(input: { - current: StreamedOperationState; - expectedRevision: number; - error: { message: string; code?: string } | string; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; -} - -export interface ComputeWorkerRouteDeps { - orchestrator: OrchestratorLike; - operationStateStore: OperationStateStoreLike; - operationEventStream: OperationEventStreamLike; - artifactExists?: (key: string) => Promise; -} +export type { ComputeWorkerRouteDeps } from './routes'; export interface CreateComputeWorkerAppOptions { host?: string; @@ -187,122 +94,27 @@ export interface ComputeWorkerApp { close(): Promise; } -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 normalizeNatsReplicas(value: number): number { - if (value === 3 || value === 5) return value; - return 1; -} - -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 buildLoggerConfig(): boolean | Record { - const format = (process.env.LOG_FORMAT?.trim().toLowerCase() || 'pretty'); - const level = process.env.COMPUTE_LOG_LEVEL?.trim() || 'info'; - if (format === 'json') { - return { - level, - base: null, - }; - } - return { - level, - base: null, - transport: { - target: 'pino-pretty', - options: { - colorize: true, - translateTime: 'SYS:standard', - ignore: 'pid,hostname', - }, - }, - }; -} - -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 toErrorMessage(error: unknown): string { - if (error instanceof Error && error.message) return error.message; - return String(error); -} - -function requestPath(request: FastifyRequest): string { - return request.url.split('?')[0] ?? request.url; -} - -function isHealthPath(path: string): boolean { - return path === '/health/live' || path === '/health/ready'; -} - -function extractTraceId(request: FastifyRequest): string | null { - const header = request.headers['x-openreader-trace-id']; - if (Array.isArray(header)) return header[0] ?? null; - return typeof header === 'string' ? header : null; -} - -function extractOpId(request: FastifyRequest, path: string): string | null { - const params = request.params as { opId?: unknown } | undefined; - if (params && typeof params.opId === 'string' && params.opId.trim()) { - return params.opId.trim(); - } - const match = path.match(/^\/v1\/operations\/([^/]+)/); - if (!match?.[1]) return null; - try { - return decodeURIComponent(match[1]); - } catch { - return match[1]; - } -} - -function isTerminalStatus(status: import('../api/contracts').WorkerJobState): boolean { - return status === 'succeeded' || status === 'failed'; -} - -const errorResponseSchema = jsonSchema(apiErrorResponseSchema); - export async function createComputeWorkerApp(options: CreateComputeWorkerAppOptions = {}): Promise { - const port = options.port ?? readIntEnv('PORT', 8081); + const port = options.port ?? readPositiveIntEnv('PORT', 8081); const host = options.host ?? (process.env.COMPUTE_WORKER_HOST?.trim() || '0.0.0.0'); const workerToken = options.workerToken ?? requireEnv('COMPUTE_WORKER_TOKEN'); const disableWorkers = options.disableWorkers ?? false; const natsUrl = requireEnv('NATS_URL'); const timeoutConfig = getComputeTimeoutConfig(); - const jobConcurrency = readIntEnv('COMPUTE_JOB_CONCURRENCY', 1); + const jobConcurrency = readPositiveIntEnv('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', false); - const jobsStreamMaxBytes = readIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); - const eventsStreamMaxBytes = readIntEnv('COMPUTE_EVENTS_STREAM_MAX_BYTES', 128 * 1024 * 1024); - const jobStatesMaxBytes = readIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024); - const natsReplicas = normalizeNatsReplicas(readIntEnv('COMPUTE_NATS_REPLICAS', 1)); + const pdfAttempts = readPositiveIntEnv('COMPUTE_PDF_JOB_ATTEMPTS', 1); + const prewarmModels = readBoolEnv('COMPUTE_PREWARM_MODELS', false); + const jobsStreamMaxBytes = readPositiveIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); + const eventsStreamMaxBytes = readPositiveIntEnv('COMPUTE_EVENTS_STREAM_MAX_BYTES', 128 * 1024 * 1024); + const jobStatesMaxBytes = readPositiveIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024); + const natsReplicas = normalizeNatsReplicas(readPositiveIntEnv('COMPUTE_NATS_REPLICAS', 1)); const opStaleMs = getComputeOpStaleMs(); - const connectOpts: Parameters[0] = { servers: natsUrl }; + const connectOpts: ConnectionOptions = { servers: natsUrl }; const natsCreds = process.env.NATS_CREDS?.trim(); const natsCredsFile = process.env.NATS_CREDS_FILE?.trim(); @@ -316,160 +128,18 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti connectOpts.authenticator = credsAuthenticator(credsData); } - // Lazy NATS connection lifecycle. The worker connects on demand (first request - // that needs the queue/KV) and disconnects after IDLE_DISCONNECT_MS of full idle - // so it stops emitting outbound traffic and Railway can sleep it. Reconnect is - // transparent: any inbound operation request both wakes the container and re-establishes - // the session via ensureConnected(). - let session: NatsSession | null = null; - let connecting: Promise | null = null; - let idleTimer: NodeJS.Timeout | null = null; - let orphanSweepTimer: NodeJS.Timeout | null = null; let stopping = false; - - // Activity accounting feeding the idle detector. The worker is considered idle - // only when no HTTP request is in flight, no SSE stream is open, no job is - // processing, and nothing has happened for IDLE_DISCONNECT_MS. let inFlightHttp = 0; let activeSse = 0; let inFlightJobs = 0; let lastActivityAt = Date.now(); - let lastActivityReason = 'startup'; - let lastIdleStatusLogAt = 0; - + let lastActivityReason = "startup"; const markActivity = (reason: string): void => { lastActivityAt = Date.now(); lastActivityReason = reason; }; - - function startIdleTimer(): void { - if (idleTimer) return; - idleTimer = setInterval(() => { - if (!session || stopping) return; - const now = Date.now(); - const idleForMs = now - lastActivityAt; - if (now - lastIdleStatusLogAt >= IDLE_STATUS_LOG_INTERVAL_MS) { - lastIdleStatusLogAt = now; - app.log.info({ - activeSse, - idleForMs, - inFlightHttp, - inFlightJobs, - lastActivityReason, - disconnectEligible: inFlightHttp === 0 - && inFlightJobs === 0 - && idleForMs >= IDLE_DISCONNECT_MS, - }, 'nats idle status'); - } - // Hard work in flight always blocks idle. An open SSE no longer blocks just - // by existing — only by delivering events, which refresh lastActivityAt via - // markActivity() in the stream's onEvent. So a silent/stuck-op stream lets - // the idle window elapse and is torn down by disconnect() below. - if (inFlightHttp > 0 || inFlightJobs > 0) return; - if (idleForMs < IDLE_DISCONNECT_MS) return; - void disconnect('idle'); - }, IDLE_CHECK_INTERVAL_MS); - // Don't let the idle checker keep the process alive on its own. - idleTimer.unref?.(); - } - - function startOrphanSweepTimer(): void { - if (orphanSweepTimer) return; - orphanSweepTimer = setInterval(() => { - if (!session || stopping) return; - void runOrphanedOpRecovery({ force: true }).catch((error) => { - app.log.error({ - error: toErrorMessage(error), - }, 'periodic orphaned operation recovery failed'); - }); - }, ORPHAN_SWEEP_INTERVAL_MS); - orphanSweepTimer.unref?.(); - } - - async function disconnect(reason: string): Promise { - const current = session; - if (!current) return; - // Snapshot what's still attached so a dropped connection is visible in logs - // (e.g. SSE streams torn down because we went idle while they were open). - app.log.info({ - reason, - activeSse, - inFlightHttp, - inFlightJobs, - idleForMs: Date.now() - lastActivityAt, - }, 'nats dropping connection'); - // Clear synchronously (before any await) so concurrent requests reconnect a - // fresh session instead of using the connection we're about to close. - session = null; - if (idleTimer) { - clearInterval(idleTimer); - idleTimer = null; - } - if (orphanSweepTimer) { - clearInterval(orphanSweepTimer); - orphanSweepTimer = null; - } - try { - await current.nc.close(); - } catch { - // ignore close errors - } - await workerLoops.stop(); - app.log.info({ reason }, 'nats disconnected'); - } - - async function ensureConnected(): Promise { - if (session) return session; - if (connecting) return connecting; - connecting = (async () => { - 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, - jobsMaxBytes: jobsStreamMaxBytes, - eventsMaxBytes: eventsStreamMaxBytes, - natsReplicas, - }); - const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { - replicas: natsReplicas, - history: 1, - ttl: COMPUTE_STATE_TTL_MS, - max_bytes: jobStatesMaxBytes, - }); - 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 next: NatsSession = { nc, js, jsm, kv, whisperConsumer, layoutConsumer }; - session = next; - sessionGeneration += 1; - orphanRecoveryDoneForGeneration = -1; - markActivity('nats_connected'); - if (!disableWorkers) { - workerLoops.start(next, { - whisper: next.whisperConsumer, - pdfLayout: next.layoutConsumer, - }); - } - startIdleTimer(); - startOrphanSweepTimer(); - // Safety net: if the connection closes for any reason (network drop after - // exhausting reconnects, or our own disconnect), drop the stale session so - // the next request reconnects cleanly. - void nc.closed().then(() => { - if (session?.nc === nc) session = null; - }); - app.log.info('nats connected'); - return next; - })(); - try { - return await connecting; - } finally { - connecting = null; - } - } + let sessionManager!: ReturnType; + const ensureConnected = () => sessionManager.ensureConnected(); const s3Prefix = normalizeS3Prefix(process.env.S3_PREFIX); const storageDisabled = async (): Promise => { @@ -517,8 +187,8 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti ErrorResponse: jsonSchema(apiErrorResponseSchema), OperationError: jsonSchema(operationErrorSchema), PdfLayoutProgress: jsonSchema(pdfLayoutProgressSchema), - PublicOperation: jsonSchema(publicOperationSchema), - PublicOperationEvent: jsonSchema(publicOperationEventSchema), + ComputeOperation: jsonSchema(computeOperationSchema), + ComputeOperationEvent: jsonSchema(computeOperationEventSchema), PdfLayoutResolution: jsonSchema(pdfLayoutResolutionSchema), }, }, @@ -571,392 +241,44 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti const operationStateStore = options.routeDeps?.operationStateStore ?? defaultOperationStateStore; const operationEventStream = options.routeDeps?.operationEventStream ?? defaultOperationEventStream; const orchestrator = options.routeDeps?.orchestrator ?? defaultOrchestrator; - let orphanRecoveryPromise: Promise | null = null; - let orphanRecoveryDoneForGeneration = -1; - let sessionGeneration = options.routeDeps ? 0 : -1; + const reconciler = createOperationReconciler({ + stateStore: operationStateStore, + orchestrator, + whisperTimeoutMs, + pdfTimeoutMs, + opStaleMs, + getGeneration: () => options.routeDeps ? 0 : sessionManager.getGeneration(), + logger: app.log, + }); + const ensureOrphanedOpRecovery = () => reconciler.run(); + const getOpState = (opId: string) => reconciler.getOpState(opId); - const runOrphanedOpRecovery = async (options?: { force?: boolean }): Promise => { - if (typeof operationStateStore.listOpStates !== 'function') return; - if (typeof operationStateStore.getOpStateRecord !== 'function') return; - if (typeof orchestrator.markFailedIfUnchanged !== 'function') return; - if (!options?.force && orphanRecoveryDoneForGeneration === sessionGeneration) return; - if (orphanRecoveryPromise) { - await orphanRecoveryPromise; - return; - } - - orphanRecoveryPromise = (async () => { - const recoveredStates = await recoverOrphanedOperations({ - operationStateStore: operationStateStore as typeof operationStateStore & { - getOpStateRecord(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; - listOpStates(): Promise; - }, - orchestrator: orchestrator as typeof orchestrator & { - markFailedIfUnchanged(input: { - current: StreamedOperationState; - expectedRevision: number; - error: { message: string; code?: string } | string; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - }, - whisperTimeoutMs, - pdfTimeoutMs, - opStaleMs, - }); - - if (recoveredStates.length > 0) { - app.log.warn({ - recoveredCount: recoveredStates.length, - ops: recoveredStates.map((state) => ({ - opId: state.opId, - kind: state.kind, - status: state.status, - })), - }, 'recovered stale in-flight operations during reconciliation'); - } - - orphanRecoveryDoneForGeneration = sessionGeneration; - })().finally(() => { - orphanRecoveryPromise = null; - }); - - await orphanRecoveryPromise; - }; - - const ensureOrphanedOpRecovery = async (): Promise => { - await runOrphanedOpRecovery(); - }; - - const getOpState = async (opId: string): Promise => { - await ensureOrphanedOpRecovery(); - return await operationStateStore.getOpState(opId); - }; - - const releaseHttp = (request: FastifyRequest): void => { - const counted = request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean }; - if (!counted[REQUEST_COUNTED_KEY]) return; - counted[REQUEST_COUNTED_KEY] = false; - inFlightHttp = Math.max(0, inFlightHttp - 1); - markActivity('http_completed'); - }; - - app.addHook('onRequest', async (request, reply) => { - const path = requestPath(request); - (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY] = Date.now(); - // Count every request as in-flight activity so the idle detector never - // disconnects mid-request. Released in onResponse, or manually after hijack - // for SSE streams (where onResponse does not fire). - (request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean })[REQUEST_COUNTED_KEY] = true; - inFlightHttp += 1; - markActivity(`http_started:${path}`); - if (isHealthPath(path)) return; - if (!isAuthed(request, workerToken)) { - return reply.code(401).send({ error: 'Unauthorized' }); - } - return; + const { releaseHttp } = registerHttpHooks({ + app, + workerToken, + markActivity, + onInFlightHttpChanged: (delta) => { + inFlightHttp = Math.max(0, inFlightHttp + delta); + }, }); - app.addHook('onResponse', async (request, reply) => { - releaseHttp(request); - const path = requestPath(request); - if (isHealthPath(path)) return; - if (reply.statusCode >= 500) { - const startedAt = (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY]; - const durationMs = Number.isFinite(startedAt) ? Math.max(0, Date.now() - (startedAt as number)) : -1; - app.log.error({ - reqId: request.id, - method: request.method, - path, - statusCode: reply.statusCode, - durationMs, - traceId: extractTraceId(request) ?? null, - opId: extractOpId(request, path), - }, 'http.error'); - } - }); - - app.get('/health/live', { - schema: { - security: [], - response: { 200: { type: 'object', properties: { ok: { type: 'boolean' } }, required: ['ok'] } }, + registerComputeWorkerRoutes({ + app, + deps: { + orchestrator, + operationStateStore, + operationEventStream, + artifactExists: options.routeDeps?.artifactExists ?? storage.objectExists, }, - }, async () => ({ ok: true })); - - // Reports readiness without forcing a NATS round-trip. Probing NATS here would - // reconnect (and keep) the connection open, defeating idle sleep, so we only - // report the current connection state. The worker reconnects lazily on the next - // operation request regardless of what this returns. - app.get('/health/ready', { - schema: { - security: [], - response: { - 200: { - type: 'object', - properties: { ok: { type: 'boolean' }, natsConnected: { type: 'boolean' } }, - required: ['ok', 'natsConnected'], - }, - }, + s3Prefix, + ensureOrphanedOpRecovery, + getOpState, + getNatsConnected: () => sessionManager.isConnected(), + releaseHttp, + markActivity, + onActiveSseChanged: (delta) => { + activeSse = Math.max(0, activeSse + delta); }, - }, async () => ({ ok: true, natsConnected: session !== null })); - - app.post('/v1/whisper-align/operations', { - schema: { - body: jsonSchema(whisperOperationCreateSchema), - response: { 202: jsonSchema(publicOperationSchema), 400: errorResponseSchema }, - }, - }, async (request, reply) => { - const parsed = whisperOperationCreateSchema.safeParse(request.body); - if (!parsed.success) { - reply.code(400); - return { - error: 'Invalid request body', - issues: parsed.error.issues, - }; - } - - const requestOp: WorkerOperationRequest = { - kind: 'whisper_align', - opKey: buildWhisperOperationKey(parsed.data), - payload: parsed.data, - }; - await ensureOrphanedOpRecovery(); - const op = await orchestrator.enqueueOrReuse(requestOp); - app.log.info({ - kind: requestOp.kind, - opId: op.opId, - jobId: op.jobId, - status: op.status, - opKeyHash: hashOpKey(requestOp.opKey.trim()).slice(0, 16), - }, 'op.accepted'); - reply.code(202); - return toPublicOperation(op); - }); - - app.post('/v1/pdf-layout/operations', { - schema: { - body: jsonSchema(pdfOperationCreateSchema), - response: { 202: jsonSchema(publicOperationSchema), 400: errorResponseSchema }, - }, - }, async (request, reply) => { - const parsed = pdfOperationCreateSchema.safeParse(request.body); - if (!parsed.success) { - reply.code(400); - return { - error: 'Invalid request body', - issues: parsed.error.issues, - }; - } - - const requestOp: WorkerOperationRequest = { - kind: 'pdf_layout', - opKey: buildPdfOperationKey(parsed.data), - payload: { - documentId: parsed.data.documentId, - namespace: parsed.data.namespace, - documentObjectKey: parsed.data.documentObjectKey, - }, - }; - await ensureOrphanedOpRecovery(); - const op = await orchestrator.enqueueOrReuse(requestOp); - reply.code(202); - return toPublicOperation(op); - }); - - app.post('/v1/pdf-layout/resolve', { - schema: { - body: jsonSchema(pdfResolveSchema), - response: { - 200: jsonSchema(pdfLayoutResolutionSchema), - 400: errorResponseSchema, - }, - }, - }, async (request, reply) => { - const parsed = pdfResolveSchema.safeParse(request.body); - if (!parsed.success) { - reply.code(400); - return { error: 'Invalid request body', issues: parsed.error.issues }; - } - await ensureOrphanedOpRecovery(); - const artifactKey = parsedPdfArtifactKey({ - documentId: parsed.data.documentId, - namespace: parsed.data.namespace, - prefix: s3Prefix, - }); - const hasArtifact = await (options.routeDeps?.artifactExists ?? storage.objectExists)(artifactKey); - const opKey = buildPdfOperationKey(parsed.data); - const index = await operationStateStore.getOpIndex?.(opKey); - const operation = index?.opId ? await operationStateStore.getOpState(index.opId) : null; - return { - artifact: hasArtifact ? { objectKey: artifactKey } : null, - operation: operation ? toPublicOperation(operation) : null, - }; - }); - - app.get('/v1/operations/:opId', { - schema: { - params: jsonSchema(operationParamsSchema), - response: { 200: jsonSchema(publicOperationSchema), 400: errorResponseSchema, 404: errorResponseSchema }, - }, - }, async (request, reply) => { - const params = operationParamsSchema.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 toPublicOperation(state); - }); - - app.get('/v1/operations/:opId/events', { - schema: { - params: jsonSchema(operationParamsSchema), - querystring: jsonSchema(operationEventsQuerySchema), - response: { - 200: { type: 'string', description: 'Server-sent PublicOperationEvent stream' }, - 400: errorResponseSchema, - 404: errorResponseSchema, - }, - }, - }, async (request, reply) => { - const params = operationParamsSchema.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' }; - } - - const cursorQueryRaw = request.query as { sinceEventId?: string | number | null } | undefined; - const cursorFromQuery = Number(cursorQueryRaw?.sinceEventId ?? 0); - const lastEventIdHeader = request.headers['last-event-id']; - const cursorFromHeader = Number( - Array.isArray(lastEventIdHeader) ? (lastEventIdHeader[0] ?? 0) : (lastEventIdHeader ?? 0), - ); - const sinceEventId = Math.max( - 0, - Number.isFinite(cursorFromQuery) ? Math.floor(cursorFromQuery) : 0, - Number.isFinite(cursorFromHeader) ? Math.floor(cursorFromHeader) : 0, - ); - - reply.hijack(); - // onResponse will not fire for a hijacked reply, so release the HTTP in-flight - // count here and track the long-lived stream via activeSse instead. - releaseHttp(request); - activeSse += 1; - markActivity('sse_started'); - 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'); - // Tell the browser EventSource to back off before reconnecting. If this stream - // is torn down because the worker went idle (NATS dropped), we don't want the - // client to reconnect immediately and re-wake the container. - reply.raw.write(encodeSseFrame({ retry: OP_EVENTS_RECONNECT_HINT_MS })); - - let closed = false; - let unsubscribe: (() => void) | null = null; - let keepalive: NodeJS.Timeout | null = null; - - const writeSnapshot = (snapshot: StreamedOperationState, eventId: number): void => { - if (closed || reply.raw.writableEnded) return; - const frameEvent: PublicOperationEvent = { - eventId, - snapshot: toPublicOperation(snapshot), - }; - reply.raw.write(encodeSseFrame({ - id: eventId, - event: 'snapshot', - data: frameEvent, - })); - }; - - const closeStream = (): void => { - if (closed) return; - closed = true; - if (unsubscribe) { - unsubscribe(); - unsubscribe = null; - } - if (keepalive) { - clearInterval(keepalive); - keepalive = null; - } - activeSse = Math.max(0, activeSse - 1); - markActivity('sse_closed'); - if (!reply.raw.writableEnded) { - reply.raw.end(); - } - }; - - request.raw.on('close', () => { - closeStream(); - }); - - try { - let current = initial; - let signature = JSON.stringify(current); - writeSnapshot(current, sinceEventId > 0 ? sinceEventId : 0); - if (isTerminalStatus(current.status)) { - return reply; - } - - keepalive = setInterval(() => { - if (closed || reply.raw.writableEnded) return; - reply.raw.write(': keepalive\n\n'); - }, OP_EVENTS_KEEPALIVE_MS); - - unsubscribe = await operationEventStream.subscribe({ - opId: params.data.opId, - sinceEventId, - onEvent: (event) => { - if (closed) return; - if (event.snapshot.opId !== params.data.opId) return; - const nextSignature = JSON.stringify(event.snapshot); - if (nextSignature !== signature) { - current = event.snapshot; - signature = nextSignature; - // A real event is progress: refresh the idle window so an actively - // streaming op keeps the worker awake. A silent stream does not. - markActivity('sse_event'); - writeSnapshot(current, event.eventId); - } - if (isTerminalStatus(event.snapshot.status)) { - closeStream(); - } - }, - onError: (error) => { - app.log.warn({ - opId: params.data.opId, - error: toErrorMessage(error), - }, 'op events stream loop error'); - closeStream(); - }, - }); - - await new Promise((resolve) => { - request.raw.once('close', () => resolve()); - }); - } catch (error) { - app.log.warn({ - opId: params.data.opId, - error: toErrorMessage(error), - }, 'op events stream loop error'); - } finally { - closeStream(); - } - - return reply; }); const jobHandlers = createJobHandlers({ @@ -974,7 +296,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti pdfAttempts, whisperCodec: whisperJobCodec, pdfCodec: layoutJobCodec, - isOwnerActive: (owner) => session === owner, + isOwnerActive: (owner) => sessionManager.isOwnerActive(owner), isStopping: () => stopping, markActivity, onInFlightJobsChanged: (delta) => { @@ -982,32 +304,41 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti }, }); + sessionManager = createNatsSessionManager({ + connectOptions: connectOpts, + logger: app.log, + whisperTimeoutMs, + pdfTimeoutMs, + pdfAttempts, + jobsStreamMaxBytes, + eventsStreamMaxBytes, + jobStatesMaxBytes, + natsReplicas, + isStopping: () => stopping, + getActivity: () => ({ + activeSse, + inFlightHttp, + inFlightJobs, + lastActivityAt, + lastActivityReason, + }), + markActivity, + startWorkers: (session) => { + if (disableWorkers) return; + workerLoops.start(session, { + whisper: session.whisperConsumer, + pdfLayout: session.layoutConsumer, + }); + }, + stopWorkers: () => workerLoops.stop(), + runReconciliation: () => reconciler.run({ force: true }), + }); + const close = async (): Promise => { if (stopping) return; stopping = true; - if (idleTimer) { - clearInterval(idleTimer); - idleTimer = null; - } - if (orphanSweepTimer) { - clearInterval(orphanSweepTimer); - orphanSweepTimer = null; - } await app.close(); - await workerLoops.stop(); - const current = session; - session = null; - if (current) { - try { - await current.nc.drain(); - } catch { - try { - await current.nc.close(); - } catch { - // ignore close errors - } - } - } + await sessionManager.close(); }; const registerSignalHandlers = (): void => { diff --git a/compute-worker/src/api/public-operation.ts b/compute-worker/src/api/compute-operation.ts similarity index 72% rename from compute-worker/src/api/public-operation.ts rename to compute-worker/src/api/compute-operation.ts index ac9b43b..03c6360 100644 --- a/compute-worker/src/api/public-operation.ts +++ b/compute-worker/src/api/compute-operation.ts @@ -1,13 +1,13 @@ -import type { WorkerOperationState } from '../api/contracts'; -import { pdfSubjectFromOperationKey } from './operation-keys'; +import type { WorkerOperationState } from '../operations/contracts'; +import { pdfSubjectFromOperationKey } from '../operations/keys'; -export type PublicOperationSubject = +export type ComputeOperationSubject = | { kind: 'whisper_align' } | { kind: 'pdf_layout'; documentId: string; namespace: string | null }; -export interface PublicOperation { +export interface ComputeOperation { opId: string; - subject: PublicOperationSubject; + subject: ComputeOperationSubject; status: WorkerOperationState['status']; queuedAt: number; updatedAt: number; @@ -18,14 +18,14 @@ export interface PublicOperation { progress?: WorkerOperationState['progress']; } -export interface PublicOperationEvent { +export interface ComputeOperationEvent { eventId: number; - snapshot: PublicOperation; + snapshot: ComputeOperation; } -export function toPublicOperation( +export function toComputeOperation( state: WorkerOperationState, -): PublicOperation { +): ComputeOperation { const subject = state.kind === 'pdf_layout' ? (pdfSubjectFromOperationKey(state.opKey) ?? { kind: 'pdf_layout', documentId: '', namespace: null }) : { kind: 'whisper_align' as const }; diff --git a/compute-worker/src/api/http-hooks.ts b/compute-worker/src/api/http-hooks.ts new file mode 100644 index 0000000..d633da3 --- /dev/null +++ b/compute-worker/src/api/http-hooks.ts @@ -0,0 +1,80 @@ +import type { FastifyInstance, FastifyRequest } from 'fastify'; + +const REQUEST_STARTED_AT_MS_KEY = Symbol('request-started-at-ms'); +const REQUEST_COUNTED_KEY = Symbol('request-activity-counted'); + +function requestPath(request: FastifyRequest): string { + return request.url.split('?')[0] ?? request.url; +} + +function isHealthPath(path: string): boolean { + return path === '/health/live' || path === '/health/ready'; +} + +function isAuthed(request: FastifyRequest, expectedToken: string): boolean { + const auth = request.headers.authorization; + return auth?.startsWith('Bearer ') === true + && auth.slice('Bearer '.length).trim() === expectedToken; +} + +function extractTraceId(request: FastifyRequest): string | null { + const header = request.headers['x-openreader-trace-id']; + if (Array.isArray(header)) return header[0] ?? null; + return typeof header === 'string' ? header : null; +} + +function extractOpId(request: FastifyRequest, path: string): string | null { + const params = request.params as { opId?: unknown } | undefined; + if (typeof params?.opId === 'string' && params.opId.trim()) return params.opId.trim(); + const match = path.match(/^\/v1\/operations\/([^/]+)/); + if (!match?.[1]) return null; + try { + return decodeURIComponent(match[1]); + } catch { + return match[1]; + } +} + +export function registerHttpHooks(input: { + app: FastifyInstance; + workerToken: string; + markActivity: (reason: string) => void; + onInFlightHttpChanged: (delta: number) => void; +}) { + const releaseHttp = (request: FastifyRequest): void => { + const counted = request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean }; + if (!counted[REQUEST_COUNTED_KEY]) return; + counted[REQUEST_COUNTED_KEY] = false; + input.onInFlightHttpChanged(-1); + input.markActivity('http_completed'); + }; + + input.app.addHook('onRequest', async (request, reply) => { + const path = requestPath(request); + (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY] = Date.now(); + (request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean })[REQUEST_COUNTED_KEY] = true; + input.onInFlightHttpChanged(1); + input.markActivity(`http_started:${path}`); + if (!isHealthPath(path) && !isAuthed(request, input.workerToken)) { + return reply.code(401).send({ error: 'Unauthorized' }); + } + }); + + input.app.addHook('onResponse', async (request, reply) => { + releaseHttp(request); + const path = requestPath(request); + if (isHealthPath(path) || reply.statusCode < 500) return; + const startedAt = (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY]; + input.app.log.error({ + reqId: request.id, + method: request.method, + path, + statusCode: reply.statusCode, + durationMs: Number.isFinite(startedAt) ? Math.max(0, Date.now() - (startedAt as number)) : -1, + traceId: extractTraceId(request), + opId: extractOpId(request, path), + }, 'http.error'); + }); + + return { releaseHttp }; +} diff --git a/compute-worker/src/api/routes.ts b/compute-worker/src/api/routes.ts new file mode 100644 index 0000000..481e965 --- /dev/null +++ b/compute-worker/src/api/routes.ts @@ -0,0 +1,356 @@ +import type { FastifyInstance, FastifyRequest } from 'fastify'; +import { encodeSseFrame } from '../operations'; +import type { + PdfLayoutJobResult, + WorkerJobTiming, + WorkerOperationEvent, + WorkerOperationRequest, + WorkerOperationState, + WhisperAlignJobResult, +} from '../operations/contracts'; +import { hashOpKey } from '../infrastructure/nats-adapters'; +import type { StreamedOperationState } from '../operations/recovery'; +import type { ReconciliationStateStore } from '../operations/reconciliation'; +import { parsedPdfArtifactKey } from '../storage/artifact-addressing'; +import { buildPdfOperationKey, buildWhisperOperationKey } from '../operations/keys'; +import { toComputeOperation, type ComputeOperationEvent } from './compute-operation'; +import { + apiErrorResponseSchema, + jsonSchema, + operationEventsQuerySchema, + operationParamsSchema, + pdfLayoutResolutionSchema, + pdfOperationCreateSchema, + pdfResolveSchema, + computeOperationSchema, + whisperOperationCreateSchema, +} from './schemas'; + +const OP_EVENTS_KEEPALIVE_MS = 15_000; +const OP_EVENTS_RECONNECT_HINT_MS = 120_000; +const errorResponseSchema = jsonSchema(apiErrorResponseSchema); + +interface OperationEventStreamLike { + subscribe(input: { + opId: string; + sinceEventId?: number; + onEvent: (event: WorkerOperationEvent) => void | Promise; + onError?: (error: unknown) => void; + }): Promise<() => void>; +} + +interface OperationStateStoreLike extends ReconciliationStateStore {} + +interface OrchestratorLike { + enqueueOrReuse(request: WorkerOperationRequest): Promise; + markRunning(input: { + opId: string; + startedAt?: number; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markProgress(input: { + opId: string; + progress: WorkerOperationState['progress']; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markSucceeded(input: { + opId: string; + result: unknown; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markFailed(input: { + opId: string; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markFailedIfUnchanged?(input: { + current: StreamedOperationState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; +} + +export interface ComputeWorkerRouteDeps { + orchestrator: OrchestratorLike; + operationStateStore: OperationStateStoreLike; + operationEventStream: OperationEventStreamLike; + artifactExists?: (key: string) => Promise; +} + +function toErrorMessage(error: unknown): string { + if (error instanceof Error && error.message) return error.message; + return String(error); +} + +function isTerminalStatus(status: import('../operations/contracts').WorkerJobState): boolean { + return status === 'succeeded' || status === 'failed'; +} + +export function registerComputeWorkerRoutes(input: { + app: FastifyInstance; + deps: ComputeWorkerRouteDeps; + s3Prefix: string; + ensureOrphanedOpRecovery: () => Promise; + getOpState: (opId: string) => Promise; + getNatsConnected: () => boolean; + releaseHttp: (request: FastifyRequest) => void; + markActivity: (reason: string) => void; + onActiveSseChanged: (delta: number) => void; +}) { + const { + app, + deps, + s3Prefix, + ensureOrphanedOpRecovery, + getOpState, + getNatsConnected, + releaseHttp, + markActivity, + onActiveSseChanged, + } = input; + + app.get('/health/live', { + schema: { + security: [], + response: { 200: { type: 'object', properties: { ok: { type: 'boolean' } }, required: ['ok'] } }, + }, + }, async () => ({ ok: true })); + + app.get('/health/ready', { + schema: { + security: [], + response: { + 200: { + type: 'object', + properties: { ok: { type: 'boolean' }, natsConnected: { type: 'boolean' } }, + required: ['ok', 'natsConnected'], + }, + }, + }, + }, async () => ({ ok: true, natsConnected: getNatsConnected() })); + + app.post('/v1/whisper-align/operations', { + schema: { + body: jsonSchema(whisperOperationCreateSchema), + response: { 202: jsonSchema(computeOperationSchema), 400: errorResponseSchema }, + }, + }, async (request, reply) => { + const parsed = whisperOperationCreateSchema.safeParse(request.body); + if (!parsed.success) { + reply.code(400); + return { error: 'Invalid request body', issues: parsed.error.issues }; + } + + const requestOp: WorkerOperationRequest = { + kind: 'whisper_align', + opKey: buildWhisperOperationKey(parsed.data), + payload: parsed.data, + }; + await ensureOrphanedOpRecovery(); + const op = await deps.orchestrator.enqueueOrReuse(requestOp); + app.log.info({ + kind: requestOp.kind, + opId: op.opId, + jobId: op.jobId, + status: op.status, + opKeyHash: hashOpKey(requestOp.opKey.trim()).slice(0, 16), + }, 'op.accepted'); + reply.code(202); + return toComputeOperation(op); + }); + + app.post('/v1/pdf-layout/operations', { + schema: { + body: jsonSchema(pdfOperationCreateSchema), + response: { 202: jsonSchema(computeOperationSchema), 400: errorResponseSchema }, + }, + }, async (request, reply) => { + const parsed = pdfOperationCreateSchema.safeParse(request.body); + if (!parsed.success) { + reply.code(400); + return { error: 'Invalid request body', issues: parsed.error.issues }; + } + + const requestOp: WorkerOperationRequest = { + kind: 'pdf_layout', + opKey: buildPdfOperationKey(parsed.data), + payload: { + documentId: parsed.data.documentId, + namespace: parsed.data.namespace, + documentObjectKey: parsed.data.documentObjectKey, + }, + }; + await ensureOrphanedOpRecovery(); + const op = await deps.orchestrator.enqueueOrReuse(requestOp); + reply.code(202); + return toComputeOperation(op); + }); + + app.post('/v1/pdf-layout/resolve', { + schema: { + body: jsonSchema(pdfResolveSchema), + response: { 200: jsonSchema(pdfLayoutResolutionSchema), 400: errorResponseSchema }, + }, + }, async (request, reply) => { + const parsed = pdfResolveSchema.safeParse(request.body); + if (!parsed.success) { + reply.code(400); + return { error: 'Invalid request body', issues: parsed.error.issues }; + } + await ensureOrphanedOpRecovery(); + const artifactKey = parsedPdfArtifactKey({ + documentId: parsed.data.documentId, + namespace: parsed.data.namespace, + prefix: s3Prefix, + }); + const hasArtifact = await deps.artifactExists?.(artifactKey) ?? false; + const opKey = buildPdfOperationKey(parsed.data); + const index = await deps.operationStateStore.getOpIndex?.(opKey); + const operation = index?.opId ? await deps.operationStateStore.getOpState(index.opId) : null; + return { + artifact: hasArtifact ? { objectKey: artifactKey } : null, + operation: operation ? toComputeOperation(operation) : null, + }; + }); + + app.get('/v1/operations/:opId', { + schema: { + params: jsonSchema(operationParamsSchema), + response: { 200: jsonSchema(computeOperationSchema), 400: errorResponseSchema, 404: errorResponseSchema }, + }, + }, async (request, reply) => { + const params = operationParamsSchema.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 toComputeOperation(state); + }); + + app.get('/v1/operations/:opId/events', { + schema: { + params: jsonSchema(operationParamsSchema), + querystring: jsonSchema(operationEventsQuerySchema), + response: { + 200: { type: 'string', description: 'Server-sent ComputeOperationEvent stream' }, + 400: errorResponseSchema, + 404: errorResponseSchema, + }, + }, + }, async (request, reply) => { + const params = operationParamsSchema.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' }; + } + + const cursorQueryRaw = request.query as { sinceEventId?: string | number | null } | undefined; + const cursorFromQuery = Number(cursorQueryRaw?.sinceEventId ?? 0); + const lastEventIdHeader = request.headers['last-event-id']; + const cursorFromHeader = Number( + Array.isArray(lastEventIdHeader) ? (lastEventIdHeader[0] ?? 0) : (lastEventIdHeader ?? 0), + ); + const sinceEventId = Math.max( + 0, + Number.isFinite(cursorFromQuery) ? Math.floor(cursorFromQuery) : 0, + Number.isFinite(cursorFromHeader) ? Math.floor(cursorFromHeader) : 0, + ); + + reply.hijack(); + releaseHttp(request); + onActiveSseChanged(1); + markActivity('sse_started'); + 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'); + reply.raw.write(encodeSseFrame({ retry: OP_EVENTS_RECONNECT_HINT_MS })); + + let closed = false; + let unsubscribe: (() => void) | null = null; + let keepalive: NodeJS.Timeout | null = null; + + const writeSnapshot = (snapshot: StreamedOperationState, eventId: number): void => { + if (closed || reply.raw.writableEnded) return; + const frameEvent: ComputeOperationEvent = { + eventId, + snapshot: toComputeOperation(snapshot), + }; + reply.raw.write(encodeSseFrame({ id: eventId, event: 'snapshot', data: frameEvent })); + }; + + const closeStream = (): void => { + if (closed) return; + closed = true; + if (unsubscribe) { + unsubscribe(); + unsubscribe = null; + } + if (keepalive) { + clearInterval(keepalive); + keepalive = null; + } + onActiveSseChanged(-1); + markActivity('sse_closed'); + if (!reply.raw.writableEnded) reply.raw.end(); + }; + + request.raw.on('close', closeStream); + + try { + let current = initial; + let signature = JSON.stringify(current); + writeSnapshot(current, sinceEventId > 0 ? sinceEventId : 0); + if (isTerminalStatus(current.status)) return reply; + + keepalive = setInterval(() => { + if (!closed && !reply.raw.writableEnded) reply.raw.write(': keepalive\n\n'); + }, OP_EVENTS_KEEPALIVE_MS); + + unsubscribe = await deps.operationEventStream.subscribe({ + opId: params.data.opId, + sinceEventId, + onEvent: (event) => { + if (closed || event.snapshot.opId !== params.data.opId) return; + const nextSignature = JSON.stringify(event.snapshot); + if (nextSignature !== signature) { + current = event.snapshot; + signature = nextSignature; + markActivity('sse_event'); + writeSnapshot(current, event.eventId); + } + if (isTerminalStatus(event.snapshot.status)) closeStream(); + }, + onError: (error) => { + app.log.warn({ opId: params.data.opId, error: toErrorMessage(error) }, 'op events stream loop error'); + closeStream(); + }, + }); + + await new Promise((resolve) => request.raw.once('close', resolve)); + } catch (error) { + app.log.warn({ opId: params.data.opId, error: toErrorMessage(error) }, 'op events stream loop error'); + } finally { + closeStream(); + } + return reply; + }); +} diff --git a/compute-worker/src/api/schemas.ts b/compute-worker/src/api/schemas.ts index 66df5c5..0175d86 100644 --- a/compute-worker/src/api/schemas.ts +++ b/compute-worker/src/api/schemas.ts @@ -95,7 +95,7 @@ export const pdfLayoutProgressSchema = z.object({ phase: z.enum(['infer', 'merge']), }); -export const publicOperationSchema = z.object({ +export const computeOperationSchema = z.object({ opId: z.string(), subject: z.discriminatedUnion('kind', [ z.object({ kind: z.literal('whisper_align') }), @@ -119,14 +119,14 @@ export const publicOperationSchema = z.object({ progress: pdfLayoutProgressSchema.optional(), }); -export const publicOperationEventSchema = z.object({ +export const computeOperationEventSchema = z.object({ eventId: z.number(), - snapshot: publicOperationSchema, + snapshot: computeOperationSchema, }); export const pdfLayoutResolutionSchema = z.object({ artifact: artifactReferenceSchema.nullable(), - operation: publicOperationSchema.nullable(), + operation: computeOperationSchema.nullable(), }); export function jsonSchema(schema: z.ZodType): Record { diff --git a/compute-worker/src/inference/pdf/parse.ts b/compute-worker/src/inference/pdf/parse.ts index b4491b2..4f70f53 100644 --- a/compute-worker/src/inference/pdf/parse.ts +++ b/compute-worker/src/inference/pdf/parse.ts @@ -4,7 +4,7 @@ import { ensureModel } from './model'; import { runLayoutModel } from './layout-model'; import { mergeTextWithRegions } from './document-layout'; import { configurePdfjsNodeRuntime, resolvePdfjsStandardFontDataUrl } from './pdfjs'; -import { PDF_PARSER_VERSION } from '../../api/contracts'; +import { PDF_PARSER_VERSION } from '../../operations/contracts'; import { stitchCrossPageBlocks } from './stitch'; import { renderPage } from './render'; import { normalizeTextItemsForLayout } from './normalize-text'; diff --git a/compute-worker/src/inference/whisper/align.ts b/compute-worker/src/inference/whisper/align.ts index 073afc4..a271623 100644 --- a/compute-worker/src/inference/whisper/align.ts +++ b/compute-worker/src/inference/whisper/align.ts @@ -29,6 +29,11 @@ import { WHISPER_DECODER_MERGED_MODEL_PATH, WHISPER_DECODER_WITH_PAST_MODEL_PATH, } from './model'; +import { + applyTokenSuppression, + applyWhisperTimestampLogitsRules, + argmax, +} from './decoder'; export function buildGoertzelCoefficients(freqBins: number, fftSize: number): Float64Array { const coeffs = new Float64Array(freqBins); @@ -448,118 +453,6 @@ function buildEmptyPastFeeds() { return state.emptyPastFeedsTemplate; } -function argmax(values: Float32Array): number | null { - let bestIdx = 0; - let bestScore = Number.NEGATIVE_INFINITY; - - for (let i = 0; i < values.length; i += 1) { - const score = values[i]; - if (score > bestScore) { - bestScore = score; - bestIdx = i; - } - } - - return Number.isFinite(bestScore) ? bestIdx : null; -} - -function applyTokenSuppression(logits: Float32Array, tokens: Set) { - for (const tokenId of tokens) { - if (tokenId >= 0 && tokenId < logits.length) { - logits[tokenId] = Number.NEGATIVE_INFINITY; - } - } -} - -function logSoftmax(input: Float32Array): Float32Array { - let max = Number.NEGATIVE_INFINITY; - for (let i = 0; i < input.length; i += 1) { - if (input[i] > max) max = input[i]; - } - if (!Number.isFinite(max)) { - return new Float32Array(input.length).fill(Number.NEGATIVE_INFINITY); - } - - let sum = 0; - for (let i = 0; i < input.length; i += 1) { - sum += Math.exp(input[i] - max); - } - const logSum = Math.log(sum); - - const out = new Float32Array(input.length); - for (let i = 0; i < input.length; i += 1) { - out[i] = input[i] - max - logSum; - } - return out; -} - -function applyWhisperTimestampLogitsRules(input: { - logits: Float32Array; - generated: number[]; - beginIndex: number; - eosTokenId: number; - noTimestampsTokenId: number; - timestampBeginTokenId: number; - maxInitialTimestampIndex: number; -}) { - const { - logits, - generated, - beginIndex, - eosTokenId, - noTimestampsTokenId, - timestampBeginTokenId, - maxInitialTimestampIndex, - } = input; - - if (noTimestampsTokenId >= 0 && noTimestampsTokenId < logits.length) { - logits[noTimestampsTokenId] = Number.NEGATIVE_INFINITY; - } - - if (generated.length === beginIndex) { - const upper = Math.min(timestampBeginTokenId, logits.length); - for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - - const seq = generated.slice(beginIndex); - const lastWasTimestamp = seq.length >= 1 && seq[seq.length - 1] >= timestampBeginTokenId; - const penultimateWasTimestamp = seq.length < 2 || seq[seq.length - 2] >= timestampBeginTokenId; - - if (lastWasTimestamp) { - if (penultimateWasTimestamp) { - for (let i = timestampBeginTokenId; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } else { - const upper = Math.min(eosTokenId, logits.length); - for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - } - - if (generated.length === beginIndex && Number.isFinite(maxInitialTimestampIndex)) { - const lastAllowed = timestampBeginTokenId + maxInitialTimestampIndex; - for (let i = lastAllowed + 1; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - - const textUpper = Math.min(timestampBeginTokenId, logits.length); - if (textUpper <= 0 || textUpper >= logits.length) return; - - const logprobs = logSoftmax(logits); - - let maxTextTokenLogprob = Number.NEGATIVE_INFINITY; - for (let i = 0; i < textUpper; i += 1) { - if (logprobs[i] > maxTextTokenLogprob) maxTextTokenLogprob = logprobs[i]; - } - - let timestampProbMass = 0; - for (let i = textUpper; i < logprobs.length; i += 1) { - timestampProbMass += Math.exp(logprobs[i]); - } - const timestampLogprob = timestampProbMass > 0 ? Math.log(timestampProbMass) : Number.NEGATIVE_INFINITY; - - if (timestampLogprob > maxTextTokenLogprob) { - for (let i = 0; i < textUpper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } -} - async function getRuntime(): Promise { if (state.runtimePromise) return state.runtimePromise; diff --git a/compute-worker/src/inference/whisper/decoder.ts b/compute-worker/src/inference/whisper/decoder.ts new file mode 100644 index 0000000..777047e --- /dev/null +++ b/compute-worker/src/inference/whisper/decoder.ts @@ -0,0 +1,111 @@ +export function argmax(values: Float32Array): number | null { + let bestIdx = 0; + let bestScore = Number.NEGATIVE_INFINITY; + + for (let i = 0; i < values.length; i += 1) { + const score = values[i]; + if (score > bestScore) { + bestScore = score; + bestIdx = i; + } + } + + return Number.isFinite(bestScore) ? bestIdx : null; +} + +export function applyTokenSuppression(logits: Float32Array, tokens: Set) { + for (const tokenId of tokens) { + if (tokenId >= 0 && tokenId < logits.length) { + logits[tokenId] = Number.NEGATIVE_INFINITY; + } + } +} + +function logSoftmax(input: Float32Array): Float32Array { + let max = Number.NEGATIVE_INFINITY; + for (let i = 0; i < input.length; i += 1) { + if (input[i] > max) max = input[i]; + } + if (!Number.isFinite(max)) { + return new Float32Array(input.length).fill(Number.NEGATIVE_INFINITY); + } + + let sum = 0; + for (let i = 0; i < input.length; i += 1) { + sum += Math.exp(input[i] - max); + } + const logSum = Math.log(sum); + + const out = new Float32Array(input.length); + for (let i = 0; i < input.length; i += 1) { + out[i] = input[i] - max - logSum; + } + return out; +} + +export function applyWhisperTimestampLogitsRules(input: { + logits: Float32Array; + generated: number[]; + beginIndex: number; + eosTokenId: number; + noTimestampsTokenId: number; + timestampBeginTokenId: number; + maxInitialTimestampIndex: number; +}) { + const { + logits, + generated, + beginIndex, + eosTokenId, + noTimestampsTokenId, + timestampBeginTokenId, + maxInitialTimestampIndex, + } = input; + + if (noTimestampsTokenId >= 0 && noTimestampsTokenId < logits.length) { + logits[noTimestampsTokenId] = Number.NEGATIVE_INFINITY; + } + + if (generated.length === beginIndex) { + const upper = Math.min(timestampBeginTokenId, logits.length); + for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + + const seq = generated.slice(beginIndex); + const lastWasTimestamp = seq.length >= 1 && seq[seq.length - 1] >= timestampBeginTokenId; + const penultimateWasTimestamp = seq.length < 2 || seq[seq.length - 2] >= timestampBeginTokenId; + + if (lastWasTimestamp) { + if (penultimateWasTimestamp) { + for (let i = timestampBeginTokenId; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } else { + const upper = Math.min(eosTokenId, logits.length); + for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + } + + if (generated.length === beginIndex && Number.isFinite(maxInitialTimestampIndex)) { + const lastAllowed = timestampBeginTokenId + maxInitialTimestampIndex; + for (let i = lastAllowed + 1; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + + const textUpper = Math.min(timestampBeginTokenId, logits.length); + if (textUpper <= 0 || textUpper >= logits.length) return; + + const logprobs = logSoftmax(logits); + + let maxTextTokenLogprob = Number.NEGATIVE_INFINITY; + for (let i = 0; i < textUpper; i += 1) { + if (logprobs[i] > maxTextTokenLogprob) maxTextTokenLogprob = logprobs[i]; + } + + let timestampProbMass = 0; + for (let i = textUpper; i < logprobs.length; i += 1) { + timestampProbMass += Math.exp(logprobs[i]); + } + const timestampLogprob = timestampProbMass > 0 ? Math.log(timestampProbMass) : Number.NEGATIVE_INFINITY; + + if (timestampLogprob > maxTextTokenLogprob) { + for (let i = 0; i < textUpper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } +} diff --git a/compute-worker/src/infrastructure/config.ts b/compute-worker/src/infrastructure/config.ts index 2547f3b..791ac4e 100644 --- a/compute-worker/src/infrastructure/config.ts +++ b/compute-worker/src/infrastructure/config.ts @@ -19,7 +19,13 @@ export type IdleTimeoutAndHardCapInput = { label: string; }; -function readPositiveIntEnv(name: string, fallback: number): number { +export function requireEnv(name: string): string { + const value = process.env[name]?.trim(); + if (!value) throw new Error(`${name} is required`); + return value; +} + +export function readPositiveIntEnv(name: string, fallback: number): number { const raw = process.env[name]?.trim(); if (!raw) return fallback; const parsed = Number(raw); @@ -27,6 +33,34 @@ function readPositiveIntEnv(name: string, fallback: number): number { return Math.floor(parsed); } +export function readBoolEnv(name: string, fallback: boolean): boolean { + const raw = process.env[name]?.trim(); + if (!raw) return fallback; + return ['1', 'true', 'yes', 'on'].includes(raw.toLowerCase()); +} + +export function normalizeNatsReplicas(value: number): number { + return value === 3 || value === 5 ? value : 1; +} + +export function buildLoggerConfig(): boolean | Record { + const format = process.env.LOG_FORMAT?.trim().toLowerCase() || 'pretty'; + const level = process.env.COMPUTE_LOG_LEVEL?.trim() || 'info'; + if (format === 'json') return { level, base: null }; + return { + level, + base: null, + transport: { + target: 'pino-pretty', + options: { + colorize: true, + translateTime: 'SYS:standard', + ignore: 'pid,hostname', + }, + }, + }; +} + let timeoutConfigCache: ComputeTimeoutConfig | null = null; let opStaleMsCache: number | null = null; diff --git a/compute-worker/src/infrastructure/nats-adapters.ts b/compute-worker/src/infrastructure/nats-adapters.ts index 00d2526..d6e4ddb 100644 --- a/compute-worker/src/infrastructure/nats-adapters.ts +++ b/compute-worker/src/infrastructure/nats-adapters.ts @@ -13,7 +13,7 @@ import type { PdfLayoutJobRequest, WhisperAlignJobRequest, WorkerOperationKind, -} from '../api/contracts'; +} from '../operations/contracts'; import { createJsonCodec } from './json-codec'; export interface KvEntryLike { diff --git a/compute-worker/src/infrastructure/nats-session.ts b/compute-worker/src/infrastructure/nats-session.ts new file mode 100644 index 0000000..356a0de --- /dev/null +++ b/compute-worker/src/infrastructure/nats-session.ts @@ -0,0 +1,204 @@ +import type { Consumer, JetStreamClient, JetStreamManager } from '@nats-io/jetstream'; +import { jetstream, jetstreamManager } from '@nats-io/jetstream'; +import { Kvm } from '@nats-io/kv'; +import { connect, type ConnectionOptions, type NatsConnection } from '@nats-io/transport-node'; +import { + COMPUTE_STATE_BUCKET, + COMPUTE_STATE_TTL_MS, + JOBS_STREAM_NAME, + LAYOUT_CONSUMER_NAME, + NATS_API_TIMEOUT_MS, + WHISPER_CONSUMER_NAME, + ensureJetStreamResources, +} from './nats'; + +const IDLE_DISCONNECT_MS = 120_000; +const IDLE_CHECK_INTERVAL_MS = 5_000; +const IDLE_STATUS_LOG_INTERVAL_MS = 60_000; +const ORPHAN_SWEEP_INTERVAL_MS = 15_000; + +export interface NatsSession { + nc: NatsConnection; + js: JetStreamClient; + jsm: JetStreamManager; + kv: Awaited>; + whisperConsumer: Consumer; + layoutConsumer: Consumer; +} + +interface NatsSessionLogger { + info(data: unknown, message?: string): void; + info(message: string): void; + error(data: unknown, message?: string): void; +} + +export interface NatsActivitySnapshot { + activeSse: number; + inFlightHttp: number; + inFlightJobs: number; + lastActivityAt: number; + lastActivityReason: string; +} + +function toErrorMessage(error: unknown): string { + return error instanceof Error && error.message ? error.message : String(error); +} + +export function createNatsSessionManager(input: { + connectOptions: ConnectionOptions; + logger: NatsSessionLogger; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + pdfAttempts: number; + jobsStreamMaxBytes: number; + eventsStreamMaxBytes: number; + jobStatesMaxBytes: number; + natsReplicas: number; + isStopping: () => boolean; + getActivity: () => NatsActivitySnapshot; + markActivity: (reason: string) => void; + startWorkers: (session: NatsSession) => void; + stopWorkers: () => Promise; + runReconciliation: () => Promise; +}) { + let session: NatsSession | null = null; + let connecting: Promise | null = null; + let idleTimer: NodeJS.Timeout | null = null; + let orphanSweepTimer: NodeJS.Timeout | null = null; + let lastIdleStatusLogAt = 0; + let generation = -1; + + const clearTimers = () => { + if (idleTimer) clearInterval(idleTimer); + if (orphanSweepTimer) clearInterval(orphanSweepTimer); + idleTimer = null; + orphanSweepTimer = null; + }; + + const disconnect = async (reason: string): Promise => { + const current = session; + if (!current) return; + const activity = input.getActivity(); + input.logger.info({ + reason, + activeSse: activity.activeSse, + inFlightHttp: activity.inFlightHttp, + inFlightJobs: activity.inFlightJobs, + idleForMs: Date.now() - activity.lastActivityAt, + }, 'nats dropping connection'); + session = null; + clearTimers(); + try { + await current.nc.close(); + } catch { + // Closing an already-closed session is harmless. + } + await input.stopWorkers(); + input.logger.info({ reason }, 'nats disconnected'); + }; + + const startTimers = () => { + if (!idleTimer) { + idleTimer = setInterval(() => { + if (!session || input.isStopping()) return; + const activity = input.getActivity(); + const now = Date.now(); + const idleForMs = now - activity.lastActivityAt; + if (now - lastIdleStatusLogAt >= IDLE_STATUS_LOG_INTERVAL_MS) { + lastIdleStatusLogAt = now; + input.logger.info({ + ...activity, + idleForMs, + disconnectEligible: activity.inFlightHttp === 0 + && activity.inFlightJobs === 0 + && idleForMs >= IDLE_DISCONNECT_MS, + }, 'nats idle status'); + } + if (activity.inFlightHttp > 0 || activity.inFlightJobs > 0 || idleForMs < IDLE_DISCONNECT_MS) return; + void disconnect('idle'); + }, IDLE_CHECK_INTERVAL_MS); + idleTimer.unref?.(); + } + if (!orphanSweepTimer) { + orphanSweepTimer = setInterval(() => { + if (!session || input.isStopping()) return; + void input.runReconciliation().catch((error) => { + input.logger.error({ error: toErrorMessage(error) }, 'periodic orphaned operation recovery failed'); + }); + }, ORPHAN_SWEEP_INTERVAL_MS); + orphanSweepTimer.unref?.(); + } + }; + + const ensureConnected = async (): Promise => { + if (session) return session; + if (connecting) return connecting; + connecting = (async () => { + const nc = await connect(input.connectOptions); + const js = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS }); + const jsm = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS }); + await ensureJetStreamResources({ + jsm, + whisperTimeoutMs: input.whisperTimeoutMs, + pdfTimeoutMs: input.pdfTimeoutMs, + pdfAttempts: input.pdfAttempts, + jobsMaxBytes: input.jobsStreamMaxBytes, + eventsMaxBytes: input.eventsStreamMaxBytes, + natsReplicas: input.natsReplicas, + }); + const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { + replicas: input.natsReplicas, + history: 1, + ttl: COMPUTE_STATE_TTL_MS, + max_bytes: input.jobStatesMaxBytes, + }); + const next: NatsSession = { + nc, + js, + jsm, + kv, + whisperConsumer: await js.consumers.get(JOBS_STREAM_NAME, WHISPER_CONSUMER_NAME), + layoutConsumer: await js.consumers.get(JOBS_STREAM_NAME, LAYOUT_CONSUMER_NAME), + }; + session = next; + generation += 1; + input.markActivity('nats_connected'); + input.startWorkers(next); + startTimers(); + void nc.closed().then(() => { + if (session?.nc === nc) session = null; + }); + input.logger.info('nats connected'); + return next; + })(); + try { + return await connecting; + } finally { + connecting = null; + } + }; + + return { + ensureConnected, + disconnect, + isConnected: () => session !== null, + isOwnerActive: (owner: object) => session === owner, + getGeneration: () => generation, + async close(): Promise { + clearTimers(); + await input.stopWorkers(); + const current = session; + session = null; + if (!current) return; + try { + await current.nc.drain(); + } catch { + try { + await current.nc.close(); + } catch { + // Closing an already-closed session is harmless. + } + } + }, + }; +} diff --git a/compute-worker/src/jobs/handlers.ts b/compute-worker/src/jobs/handlers.ts index a013602..0f64c24 100644 --- a/compute-worker/src/jobs/handlers.ts +++ b/compute-worker/src/jobs/handlers.ts @@ -10,7 +10,7 @@ import type { PdfLayoutProgress, WhisperAlignJobRequest, WhisperAlignJobResult, -} from '../api/contracts'; +} from '../operations/contracts'; import type { ArtifactStorage } from '../infrastructure/storage'; import { persistParsedPdfWhileSourceExists } from './pdf-artifact-persistence'; import { buildInferProgressForPageParsed, buildInferProgressForPageStart } from './pdf-progress'; diff --git a/compute-worker/src/jobs/pdf-progress.ts b/compute-worker/src/jobs/pdf-progress.ts index 4aefe08..acc3867 100644 --- a/compute-worker/src/jobs/pdf-progress.ts +++ b/compute-worker/src/jobs/pdf-progress.ts @@ -1,4 +1,4 @@ -import type { PdfLayoutProgress } from '../api/contracts'; +import type { PdfLayoutProgress } from '../operations/contracts'; export function buildInferProgressForPageStart(input: { pageNumber: number; diff --git a/compute-worker/src/jobs/worker-loop-policy.ts b/compute-worker/src/jobs/worker-loop-policy.ts index ce3547a..c5f9c4e 100644 --- a/compute-worker/src/jobs/worker-loop-policy.ts +++ b/compute-worker/src/jobs/worker-loop-policy.ts @@ -1,4 +1,4 @@ -import type { WorkerOperationKind } from '../api/contracts'; +import type { WorkerOperationKind } from '../operations/contracts'; export type RetryAction = 'nak_retry' | 'term_fail'; diff --git a/compute-worker/src/jobs/worker-loop.ts b/compute-worker/src/jobs/worker-loop.ts index 194bf62..27c5afe 100644 --- a/compute-worker/src/jobs/worker-loop.ts +++ b/compute-worker/src/jobs/worker-loop.ts @@ -7,7 +7,7 @@ import type { WhisperAlignJobResult, WorkerJobTiming, WorkerOperationKind, -} from '../api/contracts'; +} from '../operations/contracts'; import type { JsonCodec } from '../infrastructure/json-codec'; import type { JobHandlers } from './handlers'; import { buildQueueWaitTiming, decideRetryAction } from './worker-loop-policy'; @@ -30,7 +30,7 @@ export interface QueuedJob { payload: TPayload; } -interface WorkerLoopOrchestrator { +export interface WorkerLoopOrchestrator { markRunning(input: { opId: string; startedAt?: number; updatedAt?: number; timing?: WorkerJobTiming }): Promise; markProgress(input: { opId: string; @@ -47,7 +47,7 @@ interface WorkerLoopOrchestrator { }): Promise; } -interface WorkerLogger { +export interface WorkerLogger { info(data: unknown, message?: string): void; warn(data: unknown, message?: string): void; error(data: unknown, message?: string): void; @@ -126,6 +126,17 @@ export function createWorkerLoopController(input: { latestProgress?: PdfLayoutProgress; }; + type JobRunner = ( + payload: TPayload, + queueWaitMs: number, + hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise }, + ) => Promise; + + type WorkDefinition = { + codec: JsonCodec>; + run: JobRunner; + }; + const markRunning = async (context: Context, updatedAt: number): Promise => { if (context.latestProgress) { await input.orchestrator.markProgress({ @@ -144,10 +155,8 @@ export function createWorkerLoopController(input: { }); }; - const processMessage = async (work: { + const processMessage = async (work: WorkDefinition & { msg: JsMsg; - codec: JsonCodec>; - run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout']; workerLabel: string; }): Promise => { let context: Context | null = null; @@ -180,7 +189,7 @@ export function createWorkerLoopController(input: { }, 'failed to persist operation heartbeat state'); }); }, RUNNING_HEARTBEAT_MS); - const result = await work.run(decoded.payload as never, context.queueWaitTiming?.queueWaitMs ?? 0, { + const result = await work.run(decoded.payload, context.queueWaitTiming?.queueWaitMs ?? 0, { onProgress: async (progress) => { try { work.msg.working(); @@ -259,11 +268,9 @@ export function createWorkerLoopController(input: { } }; - const runLoop = async (work: { + const runLoop = async (work: WorkDefinition & { owner: object; consumer: Consumer; - codec: JsonCodec>; - run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout']; workerLabel: string; }): Promise => { const detached = () => input.isStopping() || stopRequested || !input.isOwnerActive(work.owner); @@ -298,9 +305,17 @@ export function createWorkerLoopController(input: { start(owner: object, consumers: { whisper: Consumer; pdfLayout: Consumer }): void { stopRequested = false; loops = []; + const whisperWork: WorkDefinition = { + codec: input.whisperCodec, + run: input.handlers.runWhisper, + }; + const pdfWork: WorkDefinition = { + codec: input.pdfCodec, + run: input.handlers.runPdfLayout, + }; 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}` })); + loops.push(runLoop({ owner, consumer: consumers.whisper, ...whisperWork, workerLabel: `whisper-${i + 1}` })); + loops.push(runLoop({ owner, consumer: consumers.pdfLayout, ...pdfWork, workerLabel: `layout-${i + 1}` })); } }, async stop(): Promise { diff --git a/compute-worker/src/api/contracts.ts b/compute-worker/src/operations/contracts.ts similarity index 95% rename from compute-worker/src/api/contracts.ts rename to compute-worker/src/operations/contracts.ts index 4f78782..5180972 100644 --- a/compute-worker/src/api/contracts.ts +++ b/compute-worker/src/operations/contracts.ts @@ -1,18 +1,18 @@ -import type { TTSSentenceAlignment, ParsedPdfDocument } from './types'; +import type { TTSSentenceAlignment, ParsedPdfDocument } from '../api/types'; export type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment, TTSSentenceWord, -} from './types'; +} from '../api/types'; export type { ParsedPdfBlockKind, ParsedPdfBlockFragment, ParsedPdfBlock, ParsedPdfPage, ParsedPdfDocument, -} from './types'; +} from '../api/types'; export const ALIGN_QUEUE_NAME = 'whisper-align'; export const PDF_LAYOUT_QUEUE_NAME = 'pdf-layout'; diff --git a/compute-worker/src/api/operation-keys.ts b/compute-worker/src/operations/keys.ts similarity index 96% rename from compute-worker/src/api/operation-keys.ts rename to compute-worker/src/operations/keys.ts index d814fb2..02fc2c8 100644 --- a/compute-worker/src/api/operation-keys.ts +++ b/compute-worker/src/operations/keys.ts @@ -1,5 +1,5 @@ import { createHash } from 'node:crypto'; -import { PDF_PARSER_VERSION } from '../api/contracts'; +import { PDF_PARSER_VERSION } from './contracts'; function sha256Hex(input: string): string { return createHash('sha256').update(input).digest('hex'); diff --git a/compute-worker/src/operations/reconciliation.ts b/compute-worker/src/operations/reconciliation.ts new file mode 100644 index 0000000..df8dc9e --- /dev/null +++ b/compute-worker/src/operations/reconciliation.ts @@ -0,0 +1,95 @@ +import type { WorkerJobTiming } from './contracts'; +import { + recoverOrphanedOperations, + type StreamedOperationState, +} from './recovery'; + +export interface ReconciliationStateStore { + getOpState(opId: string): Promise; + getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; + getOpIndex?(opKey: string): Promise<{ opId: string } | null>; + listOpStates?(): Promise; +} + +export interface ReconciliationOrchestrator { + markFailedIfUnchanged?(input: { + current: StreamedOperationState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; +} + +interface ReconciliationLogger { + warn(data: unknown, message?: string): void; +} + +function supportsRecovery( + stateStore: ReconciliationStateStore, + orchestrator: ReconciliationOrchestrator, +): boolean { + return typeof stateStore.listOpStates === 'function' + && typeof stateStore.getOpStateRecord === 'function' + && typeof orchestrator.markFailedIfUnchanged === 'function'; +} + +export function createOperationReconciler(input: { + stateStore: ReconciliationStateStore; + orchestrator: ReconciliationOrchestrator; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + opStaleMs: number; + getGeneration: () => number; + logger: ReconciliationLogger; +}) { + let recoveryPromise: Promise | null = null; + let recoveredGeneration = -1; + + const run = async (options?: { force?: boolean }): Promise => { + if (!supportsRecovery(input.stateStore, input.orchestrator)) return; + const generation = input.getGeneration(); + if (!options?.force && recoveredGeneration === generation) return; + if (recoveryPromise) return await recoveryPromise; + + recoveryPromise = (async () => { + const recoveredStates = await recoverOrphanedOperations({ + operationStateStore: { + getOpStateRecord: input.stateStore.getOpStateRecord!, + listOpStates: input.stateStore.listOpStates!, + }, + orchestrator: { + markFailedIfUnchanged: async (request) => { + const result = await input.orchestrator.markFailedIfUnchanged!(request); + return result as StreamedOperationState | null; + }, + }, + whisperTimeoutMs: input.whisperTimeoutMs, + pdfTimeoutMs: input.pdfTimeoutMs, + opStaleMs: input.opStaleMs, + }); + if (recoveredStates.length > 0) { + input.logger.warn({ + recoveredCount: recoveredStates.length, + ops: recoveredStates.map((state) => ({ + opId: state.opId, + kind: state.kind, + status: state.status, + })), + }, 'recovered stale in-flight operations during reconciliation'); + } + recoveredGeneration = generation; + })().finally(() => { + recoveryPromise = null; + }); + await recoveryPromise; + }; + + return { + run, + async getOpState(opId: string): Promise { + await run(); + return await input.stateStore.getOpState(opId); + }, + }; +} diff --git a/compute-worker/src/operations/recovery.ts b/compute-worker/src/operations/recovery.ts index 5aaaa10..b453a9a 100644 --- a/compute-worker/src/operations/recovery.ts +++ b/compute-worker/src/operations/recovery.ts @@ -4,7 +4,7 @@ import type { WorkerJobTiming, WorkerJobState, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; export type StreamedOperationState = WorkerOperationState; diff --git a/compute-worker/src/operations/service.ts b/compute-worker/src/operations/service.ts index 1469277..67430e9 100644 --- a/compute-worker/src/operations/service.ts +++ b/compute-worker/src/operations/service.ts @@ -5,7 +5,7 @@ import type { WorkerOperationKind, WorkerOperationRequest, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; import { buildQueuedState, createErrorShape, diff --git a/compute-worker/src/operations/state-machine.ts b/compute-worker/src/operations/state-machine.ts index 3144e8c..cd4076f 100644 --- a/compute-worker/src/operations/state-machine.ts +++ b/compute-worker/src/operations/state-machine.ts @@ -4,7 +4,7 @@ import type { WorkerOperationKind, WorkerOperationRequest, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; export function isTerminalStatus(status: WorkerJobState): boolean { return status === 'succeeded' || status === 'failed'; diff --git a/compute-worker/src/operations/types.ts b/compute-worker/src/operations/types.ts index 1d60200..d3f24c2 100644 --- a/compute-worker/src/operations/types.ts +++ b/compute-worker/src/operations/types.ts @@ -2,7 +2,7 @@ import type { WorkerOperationEvent, WorkerOperationKind, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; export type OperationState = WorkerOperationState; export type OperationEvent = WorkerOperationEvent; diff --git a/compute-worker/src/storage/artifact-addressing.ts b/compute-worker/src/storage/artifact-addressing.ts index bf4e312..10bb820 100644 --- a/compute-worker/src/storage/artifact-addressing.ts +++ b/compute-worker/src/storage/artifact-addressing.ts @@ -1,5 +1,5 @@ -import { PDF_PARSER_VERSION } from '../api/contracts'; -import { encodeParserVersion } from '../api/contracts'; +import { PDF_PARSER_VERSION } from '../operations/contracts'; +import { encodeParserVersion } from '../operations/contracts'; const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; diff --git a/compute-worker/tests/compute/algorithms/whisper-decoder.test.ts b/compute-worker/tests/compute/algorithms/whisper-decoder.test.ts new file mode 100644 index 0000000..07f405e --- /dev/null +++ b/compute-worker/tests/compute/algorithms/whisper-decoder.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, test } from 'vitest'; +import { + applyTokenSuppression, + applyWhisperTimestampLogitsRules, + argmax, +} from '../../../src/inference/whisper/decoder'; + +describe('whisper decoder logits', () => { + test('selects finite maximums and suppresses valid token ids', () => { + const logits = new Float32Array([1, 5, 3]); + + applyTokenSuppression(logits, new Set([-1, 1, 10])); + + expect(argmax(logits)).toBe(2); + expect(logits[1]).toBe(Number.NEGATIVE_INFINITY); + expect(argmax(new Float32Array([Number.NEGATIVE_INFINITY]))).toBeNull(); + }); + + test('forces an initial timestamp within the configured range', () => { + const logits = new Float32Array([4, 3, 2, 1, 8, 7, 6, 5]); + + applyWhisperTimestampLogitsRules({ + logits, + generated: [10, 11], + beginIndex: 2, + eosTokenId: 3, + noTimestampsTokenId: 2, + timestampBeginTokenId: 4, + maxInitialTimestampIndex: 1, + }); + + expect(Array.from(logits.slice(0, 4))).toEqual([ + Number.NEGATIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.NEGATIVE_INFINITY, + ]); + expect(logits[4]).toBe(8); + expect(logits[5]).toBe(7); + expect(logits[6]).toBe(Number.NEGATIVE_INFINITY); + expect(logits[7]).toBe(Number.NEGATIVE_INFINITY); + }); + + test('requires text after a timestamp pair', () => { + const logits = new Float32Array([4, 3, 2, 1, 8, 7]); + + applyWhisperTimestampLogitsRules({ + logits, + generated: [10, 4, 5], + beginIndex: 1, + eosTokenId: 3, + noTimestampsTokenId: -1, + timestampBeginTokenId: 4, + maxInitialTimestampIndex: Number.POSITIVE_INFINITY, + }); + + expect(logits[4]).toBe(Number.NEGATIVE_INFINITY); + expect(logits[5]).toBe(Number.NEGATIVE_INFINITY); + expect(logits[0]).toBe(4); + }); +}); diff --git a/compute-worker/tests/compute/control-plane/orchestrator.test.ts b/compute-worker/tests/compute/control-plane/orchestrator.test.ts index 66964d3..092cfdb 100644 --- a/compute-worker/tests/compute/control-plane/orchestrator.test.ts +++ b/compute-worker/tests/compute/control-plane/orchestrator.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import type { WorkerOperationRequest } from '../../../src/api/contracts'; +import type { WorkerOperationRequest } from '../../../src/operations/contracts'; import { OperationOrchestrator } from '../../../src/operations'; import { InMemoryOperationEventStream, diff --git a/compute-worker/tests/compute/control-plane/state-machine.test.ts b/compute-worker/tests/compute/control-plane/state-machine.test.ts index 8fba44a..667084f 100644 --- a/compute-worker/tests/compute/control-plane/state-machine.test.ts +++ b/compute-worker/tests/compute/control-plane/state-machine.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import type { WorkerOperationState } from '../../../src/api/contracts'; +import type { WorkerOperationState } from '../../../src/operations/contracts'; import { explainReplacementReason, isInflightStatus, diff --git a/compute-worker/tests/compute/helpers/in-memory-control-plane.ts b/compute-worker/tests/compute/helpers/in-memory-control-plane.ts index 811521b..0348bff 100644 --- a/compute-worker/tests/compute/helpers/in-memory-control-plane.ts +++ b/compute-worker/tests/compute/helpers/in-memory-control-plane.ts @@ -1,5 +1,5 @@ import { EventEmitter } from 'node:events'; -import type { WorkerOperationKind } from '../../../src/api/contracts'; +import type { WorkerOperationKind } from '../../../src/operations/contracts'; import type { OperationEvent, OperationEventStream, diff --git a/compute-worker/tests/fixtures/fake-control-plane.ts b/compute-worker/tests/fixtures/fake-control-plane.ts index 23f2760..0cd1c86 100644 --- a/compute-worker/tests/fixtures/fake-control-plane.ts +++ b/compute-worker/tests/fixtures/fake-control-plane.ts @@ -4,7 +4,7 @@ import type { WorkerOperationEvent, WorkerOperationRequest, WorkerOperationState, -} from '../../src/api/contracts'; +} from '../../src/operations/contracts'; import type { ComputeWorkerRouteDeps } from '../../src/api/app'; type ComputeResult = WhisperAlignJobResult | PdfLayoutJobResult; diff --git a/compute-worker/tests/unit/jetstream-adapters.test.ts b/compute-worker/tests/unit/jetstream-adapters.test.ts index a238996..8425061 100644 --- a/compute-worker/tests/unit/jetstream-adapters.test.ts +++ b/compute-worker/tests/unit/jetstream-adapters.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'vitest'; import { OperationOrchestrator } from '../../src/operations'; -import type { WorkerOperationRequest } from '../../src/api/contracts'; +import type { WorkerOperationRequest } from '../../src/operations/contracts'; import { JetStreamOperationEventStream, JetStreamOperationQueue, diff --git a/compute-worker/tests/unit/key-and-progress.test.ts b/compute-worker/tests/unit/key-and-progress.test.ts index d6274db..dc1656a 100644 --- a/compute-worker/tests/unit/key-and-progress.test.ts +++ b/compute-worker/tests/unit/key-and-progress.test.ts @@ -4,7 +4,7 @@ import { buildInferProgressForPageParsed, buildInferProgressForPageStart, } from '../../src/jobs/pdf-progress'; -import { buildPdfOperationKey } from '../../src/api/operation-keys'; +import { buildPdfOperationKey } from '../../src/operations/keys'; import { parsedPdfArtifactKey } from '../../src/storage/artifact-addressing'; describe('compute worker helpers', () => { diff --git a/compute-worker/tests/unit/worker-loop.test.ts b/compute-worker/tests/unit/worker-loop.test.ts new file mode 100644 index 0000000..3b9ba07 --- /dev/null +++ b/compute-worker/tests/unit/worker-loop.test.ts @@ -0,0 +1,170 @@ +import type { Consumer, JsMsg } from '@nats-io/jetstream'; +import { describe, expect, test, vi } from 'vitest'; +import type { PdfLayoutProgress } from '../../src/operations/contracts'; +import { createJsonCodec } from '../../src/infrastructure/json-codec'; +import { + createWorkerLoopController, + type QueuedJob, + type WorkerLoopOrchestrator, +} from '../../src/jobs/worker-loop'; + +function createMessage(job: QueuedJob, deliveryCount = 1) { + const codec = createJsonCodec>(); + const ack = vi.fn(); + const nak = vi.fn(); + const term = vi.fn(); + const working = vi.fn(); + return { + codec, + msg: { + data: codec.encode(job), + info: { deliveryCount }, + ack, + nak, + term, + working, + } as unknown as JsMsg, + ack, + nak, + term, + working, + }; +} + +function createConsumer(message?: JsMsg): Consumer { + let delivered = false; + return { + next: async () => { + if (!delivered && message) { + delivered = true; + return message; + } + await new Promise((resolve) => setTimeout(resolve, 1)); + return null; + }, + } as unknown as Consumer; +} + +function createOrchestrator() { + const calls: Array<{ method: string; input: unknown }> = []; + const record = (method: string) => async (input: unknown) => { + calls.push({ method, input }); + return input; + }; + const orchestrator: WorkerLoopOrchestrator = { + markRunning: record('running'), + markProgress: record('progress'), + markSucceeded: record('succeeded'), + markFailed: record('failed'), + }; + return { orchestrator, calls }; +} + +describe('worker loop controller', () => { + test('processes PDF progress, persists success, ACKs, and releases in-flight activity', async () => { + const owner = {}; + let active = true; + let inFlight = 0; + const { orchestrator, calls } = createOrchestrator(); + let complete!: () => void; + const completed = new Promise((resolve) => { complete = resolve; }); + const progress: PdfLayoutProgress = { totalPages: 2, pagesParsed: 1, currentPage: 1, phase: 'infer' }; + const pdf = createMessage({ + jobId: 'job-pdf', + opId: 'op-pdf', + opKey: 'pdf-key', + kind: 'pdf_layout', + queuedAt: Date.now() - 10, + payload: { + documentId: 'a'.repeat(64), + namespace: null, + documentObjectKey: 'openreader/doc.pdf', + }, + }); + const whisperCodec = createJsonCodec>(); + const controller = createWorkerLoopController({ + orchestrator, + handlers: { + runWhisper: async () => ({ alignments: [] }), + runPdfLayout: async (_payload, _queueWaitMs, hooks) => { + await hooks?.onProgress?.(progress); + active = false; + complete(); + return { parsedObjectKey: 'openreader/parsed.json' }; + }, + }, + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, + jobConcurrency: 1, + pdfAttempts: 2, + whisperCodec, + pdfCodec: pdf.codec, + isOwnerActive: () => active, + isStopping: () => false, + markActivity: vi.fn(), + onInFlightJobsChanged: (delta) => { inFlight += delta; }, + }); + + controller.start(owner, { whisper: createConsumer(), pdfLayout: createConsumer(pdf.msg) }); + await completed; + await controller.stop(); + + expect(pdf.working).toHaveBeenCalledOnce(); + expect(pdf.ack).toHaveBeenCalledOnce(); + expect(pdf.nak).not.toHaveBeenCalled(); + expect(inFlight).toBe(0); + expect(calls.map((call) => call.method)).toContain('progress'); + expect(calls.map((call) => call.method)).toContain('succeeded'); + }); + + test('NAKs a retryable PDF failure and does not mark it terminal', async () => { + const owner = {}; + let active = true; + const { orchestrator, calls } = createOrchestrator(); + let attempted!: () => void; + const attemptCompleted = new Promise((resolve) => { attempted = resolve; }); + const pdf = createMessage({ + jobId: 'job-pdf', + opId: 'op-pdf', + opKey: 'pdf-key', + kind: 'pdf_layout', + queuedAt: Date.now(), + payload: { + documentId: 'a'.repeat(64), + namespace: null, + documentObjectKey: 'openreader/doc.pdf', + }, + }); + const controller = createWorkerLoopController({ + orchestrator, + handlers: { + runWhisper: async () => ({ alignments: [] }), + runPdfLayout: async () => { + active = false; + attempted(); + throw new Error('retry me'); + }, + }, + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, + jobConcurrency: 1, + pdfAttempts: 2, + whisperCodec: createJsonCodec(), + pdfCodec: pdf.codec, + isOwnerActive: () => active, + isStopping: () => false, + markActivity: vi.fn(), + onInFlightJobsChanged: vi.fn(), + }); + + controller.start(owner, { whisper: createConsumer(), pdfLayout: createConsumer(pdf.msg) }); + await attemptCompleted; + await new Promise((resolve) => setTimeout(resolve, 0)); + await controller.stop(); + + expect(pdf.nak).toHaveBeenCalledOnce(); + expect(pdf.term).not.toHaveBeenCalled(); + expect(calls.map((call) => call.method)).not.toContain('failed'); + }); +}); diff --git a/src/lib/server/compute-worker/client.ts b/src/lib/server/compute-worker/client.ts index 2f8c0b6..a8f35bf 100644 --- a/src/lib/server/compute-worker/client.ts +++ b/src/lib/server/compute-worker/client.ts @@ -4,10 +4,10 @@ import type { PdfLayoutResult, WhisperAlignRequest, WhisperAlignResult, - WorkerOperation, - WorkerOperationEvent, + ComputeOperation, + ComputeOperationEvent, } from './protocol'; -import { isWorkerOperation } from './protocol'; +import { isComputeOperation } from './protocol'; class WorkerHttpError extends Error { constructor( @@ -88,7 +88,7 @@ function parseSseFrame(frame: string): { id: number | null; data: string | null return { id, data: data.length > 0 ? data.join('\n') : null }; } -function isTerminal(operation: WorkerOperation): boolean { +function isTerminal(operation: ComputeOperation): boolean { return operation.status === 'succeeded' || operation.status === 'failed'; } @@ -105,7 +105,7 @@ export class ComputeWorkerClient { let lastError: unknown; for (let attempt = 1; attempt <= 2; attempt += 1) { try { - const operation = await this.requestJson>( + const operation = await this.requestJson>( 'POST', '/v1/whisper-align/operations', input, @@ -129,7 +129,7 @@ export class ComputeWorkerClient { throw lastError instanceof Error ? lastError : new Error('Unknown compute worker failure'); } - createPdfLayoutOperation(input: PdfLayoutRequest): Promise> { + createPdfLayoutOperation(input: PdfLayoutRequest): Promise> { return this.requestJson('POST', '/v1/pdf-layout/operations', input); } @@ -141,7 +141,7 @@ export class ComputeWorkerClient { }); } - async getOperation(opId: string): Promise | null> { + async getOperation(opId: string): Promise | null> { try { return await this.requestJson('GET', `/v1/operations/${encodeURIComponent(opId)}`); } catch (error) { @@ -199,7 +199,7 @@ export class ComputeWorkerClient { return error instanceof Error && /network|timeout|fetch failed/i.test(error.message); } - private async waitForOperation(opId: string, timeoutMs: number): Promise> { + private async waitForOperation(opId: string, timeoutMs: number): Promise> { const controller = new AbortController(); const timeout = setTimeout(() => controller.abort(), timeoutMs); let lastEventId: number | null = null; @@ -227,8 +227,8 @@ export class ComputeWorkerClient { buffer = buffer.slice(end + 2); if (frame.id !== null) lastEventId = frame.id; if (!frame.data) continue; - const event = JSON.parse(frame.data) as WorkerOperationEvent; - if (!event?.snapshot || !isWorkerOperation(event.snapshot) || event.snapshot.opId !== opId) continue; + const event = JSON.parse(frame.data) as ComputeOperationEvent; + if (!event?.snapshot || !isComputeOperation(event.snapshot) || event.snapshot.opId !== opId) continue; if (isTerminal(event.snapshot)) return event.snapshot; } } diff --git a/src/lib/server/compute-worker/generated.ts b/src/lib/server/compute-worker/generated.ts index feb2060..cdac950 100644 --- a/src/lib/server/compute-worker/generated.ts +++ b/src/lib/server/compute-worker/generated.ts @@ -469,7 +469,7 @@ export interface paths { }; requestBody?: never; responses: { - /** @description Server-sent PublicOperationEvent stream */ + /** @description Server-sent ComputeOperationEvent stream */ 200: { headers: { [name: string]: unknown; @@ -579,7 +579,7 @@ export interface components { /** @enum {string} */ phase: "infer" | "merge"; }; - PublicOperation: { + ComputeOperation: { opId: string; subject: { /** @enum {string} */ @@ -613,7 +613,7 @@ export interface components { phase: "infer" | "merge"; }; }; - PublicOperationEvent: { + ComputeOperationEvent: { eventId: number; snapshot: { opId: string; diff --git a/src/lib/server/compute-worker/protocol.ts b/src/lib/server/compute-worker/protocol.ts index 73b196c..2fa0c83 100644 --- a/src/lib/server/compute-worker/protocol.ts +++ b/src/lib/server/compute-worker/protocol.ts @@ -9,17 +9,17 @@ export type ParsedPdfBlockFragment = ParsedPdfBlock['fragments'][number]; export type TTSSentenceAlignment = components['schemas']['TTSSentenceAlignment']; export type TTSSentenceWord = TTSSentenceAlignment['words'][number]; -export type PdfLayoutProgress = NonNullable; -export type WorkerOperationStatus = components['schemas']['PublicOperation']['status']; -export type WorkerOperationSubject = components['schemas']['PublicOperation']['subject']; +export type PdfLayoutProgress = NonNullable; +export type ComputeOperationStatus = components['schemas']['ComputeOperation']['status']; +export type ComputeOperationSubject = components['schemas']['ComputeOperation']['subject']; -export type WorkerOperation = - Omit +export type ComputeOperation = + Omit & { result?: Result }; -export type WorkerOperationEvent = { +export type ComputeOperationEvent = { eventId: number; - snapshot: WorkerOperation; + snapshot: ComputeOperation; }; export type WhisperAlignRequest = @@ -31,20 +31,20 @@ export type PdfLayoutResolveRequest = export type WhisperAlignResult = { alignments: TTSSentenceAlignment[]; - timing?: components['schemas']['PublicOperation']['timing']; + timing?: components['schemas']['ComputeOperation']['timing']; }; export type PdfLayoutResult = { parsedObjectKey: string; - timing?: components['schemas']['PublicOperation']['timing']; + timing?: components['schemas']['ComputeOperation']['timing']; }; export type PdfLayoutResolution = { artifact: { objectKey: string } | null; - operation: WorkerOperation | null; + operation: ComputeOperation | null; }; -export function isWorkerOperation(value: unknown): value is WorkerOperation { +export function isComputeOperation(value: unknown): value is ComputeOperation { if (!value || typeof value !== 'object') return false; const record = value as Record; return typeof record.opId === 'string' diff --git a/src/lib/server/pdf-parse/operation.ts b/src/lib/server/pdf-parse/operation.ts index 60a924a..a10136b 100644 --- a/src/lib/server/pdf-parse/operation.ts +++ b/src/lib/server/pdf-parse/operation.ts @@ -3,7 +3,7 @@ import { documentKey } from '@/lib/server/documents/blobstore'; import type { PdfLayoutResolution, PdfLayoutResult, - WorkerOperation, + ComputeOperation, } from '@/lib/server/compute-worker/protocol'; function currentPdfOperationInput(documentId: string, namespace: string | null, forceToken?: string): { @@ -33,7 +33,7 @@ export async function resolveCurrentPdfParse(input: { export async function lookupCurrentPdfParseOperation(input: { documentId: string; namespace: string | null; -}): Promise | null> { +}): Promise | null> { return (await resolveCurrentPdfParse(input)).operation; } @@ -41,7 +41,7 @@ export async function createOrReuseCurrentPdfParseOperation(input: { documentId: string; namespace: string | null; forceToken?: string; -}): Promise> { +}): Promise> { return getComputeWorkerClient().createPdfLayoutOperation(currentPdfOperationInput( input.documentId, input.namespace, @@ -49,12 +49,12 @@ export async function createOrReuseCurrentPdfParseOperation(input: { )); } -export async function fetchPdfParseOperation(opId: string): Promise | null> { +export async function fetchPdfParseOperation(opId: string): Promise | null> { return getComputeWorkerClient().getOperation(opId); } export function isPdfParseOperationForDocument( - state: WorkerOperation, + state: ComputeOperation, input: { documentId: string; namespace: string | null; diff --git a/src/lib/server/pdf-parse/snapshot.ts b/src/lib/server/pdf-parse/snapshot.ts index 2fba7f9..b4e8924 100644 --- a/src/lib/server/pdf-parse/snapshot.ts +++ b/src/lib/server/pdf-parse/snapshot.ts @@ -1,8 +1,8 @@ -import type { PdfLayoutResult, WorkerOperation } from '@/lib/server/compute-worker/protocol'; +import type { PdfLayoutResult, ComputeOperation } from '@/lib/server/compute-worker/protocol'; import type { PdfParseStatus } from '@/types/parsed-pdf'; import type { PdfParseSnapshot } from '@/lib/server/pdf-parse/types'; -function mapWorkerStatusToParseStatus(status: WorkerOperation['status']): PdfParseStatus { +function mapWorkerStatusToParseStatus(status: ComputeOperation['status']): PdfParseStatus { switch (status) { case 'queued': return 'pending'; @@ -19,7 +19,7 @@ function mapWorkerStatusToParseStatus(status: WorkerOperation['status']): PdfPar } export function parsedObjectKeyFromWorkerState( - state: WorkerOperation, + state: ComputeOperation, ): string | null { const result = state.result; if (!result || typeof result !== 'object' || !('parsedObjectKey' in result)) return null; @@ -30,7 +30,7 @@ export function parsedObjectKeyFromWorkerState( } export function pdfParseSnapshotFromWorkerState( - state: WorkerOperation, + state: ComputeOperation, ): PdfParseSnapshot { const parseStatus = mapWorkerStatusToParseStatus(state.status); return {