From 3d5dfd2a88a76aaa96fe3b34e1dad80d249b1aa0 Mon Sep 17 00:00:00 2001 From: Richard R Date: Fri, 12 Jun 2026 14:07:15 -0600 Subject: [PATCH] refactor(compute-worker): migrate compute logic to modular architecture - Remove legacy compute-worker monolith under `src/compute/`, including PDF and Whisper inference, control-plane, and runtime orchestration - Move PDF and Whisper inference logic to new `src/inference/` module - Move control-plane, orchestrator, and state machine logic to `src/operations/` - Move NATS/JetStream adapters to `src/infrastructure/` - Move job orchestration, progress, and artifact persistence to `src/jobs/` - Update imports throughout tests and main app to reference new module structure - Update Dockerfile and scripts to use new asset paths - Add new entrypoints: `src/api/app.ts` and `src/api/contracts.ts` - Remove obsolete files and update `.gitignore` for new dev artifacts This refactor modularizes the compute-worker, improving maintainability and separation of concerns. No functional changes to inference or orchestration logic. BREAKING CHANGE: compute-worker internal APIs, directory structure, and imports have changed; downstream code must update imports and integration points. --- .gitignore | 4 +- Dockerfile | 4 +- compute-worker/scripts/generate-openapi.ts | 2 +- compute-worker/src/{runtime.ts => api/app.ts} | 881 ++---------------- .../index.ts => api/contracts.ts} | 12 +- compute-worker/src/api/operation-keys.ts | 2 +- compute-worker/src/api/public-operation.ts | 2 +- .../config/cpu-budget.ts | 0 .../{compute => inference}/config/timeout.ts | 0 .../src/{compute => inference}/index.ts | 4 +- .../{compute => inference}/local-runtime.ts | 0 .../pdf/assets/LICENSE.txt | 0 .../pdf/assets/manifest.json | 0 .../src/{compute => inference}/pdf/merge.ts | 0 .../src/{compute => inference}/pdf/model.ts | 0 .../pdf/normalize-text.ts | 0 .../src/{compute => inference}/pdf/parse.ts | 0 .../pdf/parser-version-key.ts | 0 .../pdf/parser-version.ts | 0 .../pdf/pdfjs-runtime.ts | 0 .../src/{compute => inference}/pdf/render.ts | 0 .../pdf/runLayoutModel.ts | 0 .../src/{compute => inference}/pdf/stitch.ts | 0 .../src/{compute => inference}/pdf/types.ts | 0 .../platform/docstore.ts | 0 .../{compute => inference}/platform/ffmpeg.ts | 0 .../src/{compute => inference}/types/index.ts | 0 .../types/parsed-pdf.ts | 0 .../src/{compute => inference}/types/tts.ts | 0 .../{compute => inference}/whisper/align.ts | 0 .../whisper/alignment-map.ts | 0 .../whisper/assets/LICENSE.txt | 0 .../whisper/assets/manifest.json | 0 .../whisper/assets/mel_filters.npz | Bin .../{compute => inference}/whisper/model.ts | 0 .../whisper/spectral.ts | 0 .../whisper/token-timestamps.ts | 0 .../json-codec.ts | 0 .../nats-adapters.ts} | 4 +- compute-worker/src/infrastructure/nats.ts | 104 +++ compute-worker/src/infrastructure/storage.ts | 127 +++ compute-worker/src/jobs/handlers.ts | 126 +++ .../{ => jobs}/pdf-artifact-persistence.ts | 0 compute-worker/src/{ => jobs}/pdf-progress.ts | 2 +- .../src/{ => jobs}/worker-loop-policy.ts | 2 +- compute-worker/src/jobs/worker-loop.ts | 312 +++++++ .../control-plane => operations}/index.ts | 2 +- .../recovery.ts} | 2 +- .../orchestrator.ts => operations/service.ts} | 2 +- .../control-plane => operations}/sse.ts | 0 .../state-machine.ts | 2 +- .../control-plane => operations}/types.ts | 2 +- compute-worker/src/server.ts | 2 +- .../src/storage/artifact-addressing.ts | 4 +- compute-worker/tests/api/routes.test.ts | 2 +- .../pdf-merge-text-with-regions.test.ts | 2 +- .../pdf-parse-normalize-text-items.test.ts | 2 +- .../pdf-stitch-cross-page-blocks.test.ts | 4 +- .../algorithms/support/document-fixtures.ts | 4 +- .../whisper-alignment-mapping.test.ts | 2 +- .../algorithms/whisper-spectral.test.ts | 2 +- .../whisper-token-timestamps.test.ts | 2 +- .../control-plane/orchestrator.test.ts | 4 +- .../control-plane/run-layout-model.test.ts | 8 +- .../tests/compute/control-plane/sse.test.ts | 2 +- .../control-plane/state-machine.test.ts | 4 +- .../helpers/in-memory-control-plane.ts | 4 +- .../tests/fixtures/fake-control-plane.ts | 4 +- .../tests/unit/jetstream-adapters.test.ts | 6 +- .../tests/unit/key-and-progress.test.ts | 4 +- .../tests/unit/orphan-recovery.test.ts | 2 +- .../unit/pdf-artifact-persistence.test.ts | 2 +- .../tests/unit/worker-loop-policy.test.ts | 2 +- scripts/check-next-server-bundle.mjs | 2 +- ...pute-worker-client-contract.vitest.spec.ts | 2 +- 75 files changed, 810 insertions(+), 860 deletions(-) rename compute-worker/src/{runtime.ts => api/app.ts} (56%) rename compute-worker/src/{compute/api-contracts/index.ts => api/contracts.ts} (87%) rename compute-worker/src/{compute => inference}/config/cpu-budget.ts (100%) rename compute-worker/src/{compute => inference}/config/timeout.ts (100%) rename compute-worker/src/{compute => inference}/index.ts (93%) rename compute-worker/src/{compute => inference}/local-runtime.ts (100%) rename compute-worker/src/{compute => inference}/pdf/assets/LICENSE.txt (100%) rename compute-worker/src/{compute => inference}/pdf/assets/manifest.json (100%) rename compute-worker/src/{compute => inference}/pdf/merge.ts (100%) rename compute-worker/src/{compute => inference}/pdf/model.ts (100%) rename compute-worker/src/{compute => inference}/pdf/normalize-text.ts (100%) rename compute-worker/src/{compute => inference}/pdf/parse.ts (100%) rename compute-worker/src/{compute => inference}/pdf/parser-version-key.ts (100%) rename compute-worker/src/{compute => inference}/pdf/parser-version.ts (100%) rename compute-worker/src/{compute => inference}/pdf/pdfjs-runtime.ts (100%) rename compute-worker/src/{compute => inference}/pdf/render.ts (100%) rename compute-worker/src/{compute => inference}/pdf/runLayoutModel.ts (100%) rename compute-worker/src/{compute => inference}/pdf/stitch.ts (100%) rename compute-worker/src/{compute => inference}/pdf/types.ts (100%) rename compute-worker/src/{compute => inference}/platform/docstore.ts (100%) rename compute-worker/src/{compute => inference}/platform/ffmpeg.ts (100%) rename compute-worker/src/{compute => inference}/types/index.ts (100%) rename compute-worker/src/{compute => inference}/types/parsed-pdf.ts (100%) rename compute-worker/src/{compute => inference}/types/tts.ts (100%) rename compute-worker/src/{compute => inference}/whisper/align.ts (100%) rename compute-worker/src/{compute => inference}/whisper/alignment-map.ts (100%) rename compute-worker/src/{compute => inference}/whisper/assets/LICENSE.txt (100%) rename compute-worker/src/{compute => inference}/whisper/assets/manifest.json (100%) rename compute-worker/src/{compute => inference}/whisper/assets/mel_filters.npz (100%) rename compute-worker/src/{compute => inference}/whisper/model.ts (100%) rename compute-worker/src/{compute => inference}/whisper/spectral.ts (100%) rename compute-worker/src/{compute => inference}/whisper/token-timestamps.ts (100%) rename compute-worker/src/{control-plane => infrastructure}/json-codec.ts (100%) rename compute-worker/src/{control-plane/jetstream.ts => infrastructure/nats-adapters.ts} (99%) create mode 100644 compute-worker/src/infrastructure/nats.ts create mode 100644 compute-worker/src/infrastructure/storage.ts create mode 100644 compute-worker/src/jobs/handlers.ts rename compute-worker/src/{ => jobs}/pdf-artifact-persistence.ts (100%) rename compute-worker/src/{ => jobs}/pdf-progress.ts (89%) rename compute-worker/src/{ => jobs}/worker-loop-policy.ts (91%) create mode 100644 compute-worker/src/jobs/worker-loop.ts rename compute-worker/src/{compute/control-plane => operations}/index.ts (71%) rename compute-worker/src/{orphan-recovery.ts => operations/recovery.ts} (98%) rename compute-worker/src/{compute/control-plane/orchestrator.ts => operations/service.ts} (99%) rename compute-worker/src/{compute/control-plane => operations}/sse.ts (100%) rename compute-worker/src/{compute/control-plane => operations}/state-machine.ts (98%) rename compute-worker/src/{compute/control-plane => operations}/types.ts (98%) diff --git a/.gitignore b/.gitignore index 55b2755..466d577 100644 --- a/.gitignore +++ b/.gitignore @@ -56,5 +56,7 @@ node_modules/ # vscode .vscode -# .agents +# Agents .agents +.codex +.claude diff --git a/Dockerfile b/Dockerfile index 93e32ab..f4ebc3e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -87,7 +87,7 @@ COPY --from=app-builder /app/THIRD_PARTY_LICENSES /licenses # Include SeaweedFS license text for the copied weed binary. COPY --from=seaweedfs-builder /tmp/SeaweedFS-LICENSE.txt /licenses/SeaweedFS-LICENSE.txt # Include static model notices for runtime-downloaded assets. -COPY --from=app-builder /app/compute-worker/src/compute/pdf/assets/LICENSE.txt /licenses/pp-doclayoutv3-LICENSE.txt +COPY --from=app-builder /app/compute-worker/src/inference/pdf/assets/LICENSE.txt /licenses/pp-doclayoutv3-LICENSE.txt # Copy seaweedfs weed binary for optional embedded local S3. COPY --from=seaweedfs-builder /tmp/weed /usr/local/bin/weed @@ -97,7 +97,7 @@ COPY --from=nats-builder /tmp/nats-server /usr/local/bin/nats-server RUN chmod +x /usr/local/bin/nats-server # Include OpenAI Whisper license text for runtime-downloaded ONNX artifacts. -COPY --from=app-builder /app/compute-worker/src/compute/whisper/assets/LICENSE.txt /licenses/openai-whisper-LICENSE.txt +COPY --from=app-builder /app/compute-worker/src/inference/whisper/assets/LICENSE.txt /licenses/openai-whisper-LICENSE.txt # Match the app's historical container port now that standalone server.js # is started directly instead of `next start -p 3003`. diff --git a/compute-worker/scripts/generate-openapi.ts b/compute-worker/scripts/generate-openapi.ts index 7af98d5..0cfb141 100644 --- a/compute-worker/scripts/generate-openapi.ts +++ b/compute-worker/scripts/generate-openapi.ts @@ -1,6 +1,6 @@ import { mkdir, writeFile } from 'node:fs/promises'; import { dirname, resolve } from 'node:path'; -import { createComputeWorkerApp } from '../src/runtime'; +import { createComputeWorkerApp } from '../src/api/app'; process.env.COMPUTE_WORKER_TOKEN ||= 'openapi-generation-token'; process.env.NATS_URL ||= 'nats://127.0.0.1:4222'; diff --git a/compute-worker/src/runtime.ts b/compute-worker/src/api/app.ts similarity index 56% rename from compute-worker/src/runtime.ts rename to compute-worker/src/api/app.ts index 5cafa7c..1775f4e 100644 --- a/compute-worker/src/runtime.ts +++ b/compute-worker/src/api/app.ts @@ -1,78 +1,72 @@ import Fastify, { type FastifyInstance, type FastifyRequest } from 'fastify'; import swagger from '@fastify/swagger'; -import { z } from 'zod'; import { connect, - nanos, credsAuthenticator, type NatsConnection, } from '@nats-io/transport-node'; import { - AckPolicy, - DeliverPolicy, - ReplayPolicy, - RetentionPolicy, - StorageType, jetstream, jetstreamManager, type Consumer, type JetStreamClient, type JetStreamManager, - type JsMsg, } from '@nats-io/jetstream'; import { Kvm } from '@nats-io/kv'; import { ensureComputeModels, - runPdfLayoutFromPdfBuffer, - runWhisperAlignmentFromAudioBuffer, -} from './compute/local-runtime'; +} from '../inference/local-runtime'; import { getComputeTimeoutConfig, getComputeOpStaleMs, getAvailableCpuCores, getOnnxThreadsPerJob, - withIdleTimeoutAndHardCap, - withTimeout, -} from './compute'; -import { encodeSseFrame, OperationOrchestrator } from './compute/control-plane'; +} from '../inference'; +import { encodeSseFrame, OperationOrchestrator } from '../operations'; import type { PdfLayoutJobRequest, PdfLayoutJobResult, WorkerOperationEvent, WhisperAlignJobRequest, WhisperAlignJobResult, - WorkerJobState, WorkerJobTiming, - WorkerOperationKind, WorkerOperationRequest, WorkerOperationState, - PdfLayoutProgress, -} from './compute/api-contracts'; -import { - DeleteObjectCommand, - GetObjectCommand, - HeadObjectCommand, - PutObjectCommand, - S3Client, -} from '@aws-sdk/client-s3'; +} from '../api/contracts'; import { JetStreamOperationEventStream, - OP_EVENTS_SUBJECT_WILDCARD, JetStreamOperationQueue, JetStreamOperationStateStore, hashOpKey, -} from './control-plane/jetstream'; -import { type JsonCodec, createJsonCodec } from './control-plane/json-codec'; +} from '../infrastructure/nats-adapters'; +import { createJsonCodec } from '../infrastructure/json-codec'; import { recoverOrphanedOperations, type StreamedOperationState, -} from './orphan-recovery'; -import { buildInferProgressForPageParsed, buildInferProgressForPageStart } from './pdf-progress'; -import { buildQueueWaitTiming, decideRetryAction } from './worker-loop-policy'; -import { persistParsedPdfWhileSourceExists } from './pdf-artifact-persistence'; -import { parsedPdfArtifactKey } from './storage/artifact-addressing'; -import { buildPdfOperationKey, buildWhisperOperationKey } from './api/operation-keys'; -import { toPublicOperation, type PublicOperationEvent } from './api/public-operation'; +} from '../operations/recovery'; +import { parsedPdfArtifactKey } from '../storage/artifact-addressing'; +import { + createArtifactStorage, + createS3ClientFromEnv, + normalizeS3Prefix, + type ArtifactStorage, +} from '../infrastructure/storage'; +import { createJobHandlers } from '../jobs/handlers'; +import { createWorkerLoopController, type QueuedJob } from '../jobs/worker-loop'; +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 { apiErrorResponseSchema, artifactReferenceSchema, @@ -89,26 +83,14 @@ import { publicOperationSchema, ttsSentenceAlignmentSchema, whisperOperationCreateSchema, -} from './api/schemas'; +} from './schemas'; -const JOBS_STREAM_NAME = 'compute_jobs'; -const WHISPER_JOBS_SUBJECT = 'jobs.whisper'; -const LAYOUT_JOBS_SUBJECT = 'jobs.layout'; -const WHISPER_CONSUMER_NAME = 'compute_whisper'; -const LAYOUT_CONSUMER_NAME = 'compute_layout'; -const EVENTS_STREAM_NAME = 'compute_events'; -const COMPUTE_STATE_BUCKET = 'compute_state'; -const COMPUTE_STATE_TTL_MS = 24 * 60 * 60 * 1000; -const LOOP_ERROR_BACKOFF_MS = 500; -const RUNNING_HEARTBEAT_MS = 5000; 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; -const WHISPER_MAX_DELIVER = 1; -const NATS_API_TIMEOUT_MS = 60_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. @@ -118,22 +100,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 PULL_EXPIRES_MS = 5_000; const REQUEST_STARTED_AT_MS_KEY = Symbol('request-started-at-ms'); const REQUEST_COUNTED_KEY = Symbol('request-activity-counted'); -const SLOW_JOB_LOG_THRESHOLD_MS_BY_KIND: Record = { - whisper_align: 15_000, - pdf_layout: 120_000, -}; - -interface QueuedJob { - jobId: string; - opId: string; - opKey: string; - kind: WorkerOperationKind; - queuedAt: number; - payload: TPayload; -} +const WHISPER_MAX_DELIVER = 1; interface NatsSession { nc: NatsConnection; @@ -267,64 +236,6 @@ function buildLoggerConfig(): boolean | Record { }; } -function normalizeS3Prefix(prefix: string | undefined): string { - const value = (prefix || 'openreader').trim(); - return value ? value.replace(/^\/+|\/+$/g, '') : 'openreader'; -} - -function buildS3Client(): S3Client { - const bucket = requireEnv('S3_BUCKET'); - const region = requireEnv('S3_REGION'); - const accessKeyId = requireEnv('S3_ACCESS_KEY_ID'); - const secretAccessKey = requireEnv('S3_SECRET_ACCESS_KEY'); - const endpoint = process.env.S3_ENDPOINT?.trim() || undefined; - const forcePathStyle = parseBoolEnv('S3_FORCE_PATH_STYLE', false); - - void bucket; - - return new S3Client({ - region, - endpoint, - forcePathStyle, - requestChecksumCalculation: 'WHEN_REQUIRED', - responseChecksumValidation: 'WHEN_REQUIRED', - credentials: { - accessKeyId, - secretAccessKey, - }, - }); -} - -async function bodyToBuffer(body: unknown): Promise { - if (!body) return Buffer.alloc(0); - if (body instanceof Uint8Array) return Buffer.from(body); - if (ArrayBuffer.isView(body)) return Buffer.from(body.buffer, body.byteOffset, body.byteLength); - if (body instanceof ArrayBuffer) return Buffer.from(body); - if (typeof body === 'object' && body !== null && 'transformToByteArray' in body) { - const maybe = body as { transformToByteArray?: () => Promise }; - if (typeof maybe.transformToByteArray === 'function') { - return Buffer.from(await maybe.transformToByteArray()); - } - } - if (typeof body === 'object' && body !== null && 'on' in body) { - const stream = body as NodeJS.ReadableStream; - const chunks: Buffer[] = []; - for await (const chunk of stream) { - if (Buffer.isBuffer(chunk)) chunks.push(chunk); - else if (typeof chunk === 'string') chunks.push(Buffer.from(chunk)); - else chunks.push(Buffer.from(chunk as Uint8Array)); - } - return Buffer.concat(chunks); - } - throw new Error('Unsupported S3 response body type'); -} - -function toArrayBuffer(bytes: Uint8Array): ArrayBuffer { - const copy = new Uint8Array(bytes.byteLength); - copy.set(bytes); - return copy.buffer; -} - function isAuthed(request: FastifyRequest, expectedToken: string): boolean { const auth = request.headers.authorization; if (!auth?.startsWith('Bearer ')) return false; @@ -332,11 +243,6 @@ function isAuthed(request: FastifyRequest, expectedToken: string): boolean { return token === expectedToken; } -function safeDurationMs(start: number | undefined, end: number | undefined): number | undefined { - if (!Number.isFinite(start) || !Number.isFinite(end)) return undefined; - return Math.max(0, Math.floor((end as number) - (start as number))); -} - function toErrorMessage(error: unknown): string { if (error instanceof Error && error.message) return error.message; return String(error); @@ -370,155 +276,12 @@ function extractOpId(request: FastifyRequest, path: string): string | null { } } -function isAlreadyExistsError(error: unknown): boolean { - const message = toErrorMessage(error).toLowerCase(); - return message.includes('already in use') || message.includes('already exists'); -} - -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - -class ConcurrencyGate { - private readonly maxInFlight: number; - private inFlight = 0; - private readonly queue: Array<() => void> = []; - - constructor(limit: number) { - this.maxInFlight = Number.isFinite(limit) && limit >= 1 ? Math.floor(limit) : 1; - } - - async acquire(): Promise { - if (this.inFlight < this.maxInFlight) { - this.inFlight += 1; - return; - } - - await new Promise((resolve) => { - this.queue.push(() => { - this.inFlight += 1; - resolve(); - }); - }); - } - - release(): void { - this.inFlight = Math.max(0, this.inFlight - 1); - const next = this.queue.shift(); - if (next) next(); - } -} - -function isTerminalStatus(status: WorkerJobState): boolean { +function isTerminalStatus(status: import('../api/contracts').WorkerJobState): boolean { return status === 'succeeded' || status === 'failed'; } -function extractResultRef(kind: WorkerOperationKind, result: unknown): string | undefined { - if (kind !== 'pdf_layout' || !result || typeof result !== 'object') return undefined; - const maybe = result as { parsedObjectKey?: unknown }; - return typeof maybe.parsedObjectKey === 'string' ? maybe.parsedObjectKey : undefined; -} - -const alignSchema = z.object({ - text: z.string().trim().min(1), - lang: z.string().trim().min(1).max(16).optional(), - cacheKey: z.string().trim().min(1).max(256).optional(), - audioObjectKey: z.string().trim().min(1).max(2048), -}); - -const layoutSchema = z.object({ - documentId: z.string().trim().min(1), - namespace: z.string().trim().min(1).max(128).nullable(), - documentObjectKey: z.string().trim().min(1).max(2048), -}); - const errorResponseSchema = jsonSchema(apiErrorResponseSchema); -async function ensureJetStreamResources( - jsm: JetStreamManager, - whisperTimeoutMs: number, - pdfTimeoutMs: number, - pdfAttempts: number, - jobsMaxBytes: number, - eventsMaxBytes: number, - natsReplicas: number, -): Promise { - const streamConfig = { - name: JOBS_STREAM_NAME, - subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT], - retention: RetentionPolicy.Workqueue, - storage: StorageType.File, - max_bytes: jobsMaxBytes, - num_replicas: natsReplicas, - }; - - try { - await jsm.streams.add(streamConfig); - } catch (error) { - if (!isAlreadyExistsError(error)) throw error; - await jsm.streams.update(JOBS_STREAM_NAME, { - subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT], - max_bytes: jobsMaxBytes, - num_replicas: natsReplicas, - }); - } - - const eventsStreamConfig = { - name: EVENTS_STREAM_NAME, - subjects: [OP_EVENTS_SUBJECT_WILDCARD], - retention: RetentionPolicy.Limits, - storage: StorageType.File, - max_bytes: eventsMaxBytes, - max_age: nanos(COMPUTE_STATE_TTL_MS), - num_replicas: natsReplicas, - }; - - try { - await jsm.streams.add(eventsStreamConfig); - } catch (error) { - if (!isAlreadyExistsError(error)) throw error; - await jsm.streams.update(EVENTS_STREAM_NAME, { - subjects: [OP_EVENTS_SUBJECT_WILDCARD], - max_bytes: eventsMaxBytes, - max_age: nanos(COMPUTE_STATE_TTL_MS), - num_replicas: natsReplicas, - }); - } - - const ensureConsumer = async ( - name: string, - subject: string, - ackWaitMs: number, - maxDeliver: number, - ): Promise => { - const config = { - durable_name: name, - ack_policy: AckPolicy.Explicit, - deliver_policy: DeliverPolicy.All, - replay_policy: ReplayPolicy.Instant, - filter_subject: subject, - ack_wait: nanos(Math.max(ackWaitMs, 1_000)), - max_deliver: maxDeliver, - }; - - try { - await jsm.consumers.add(JOBS_STREAM_NAME, config); - } catch (error) { - if (!isAlreadyExistsError(error)) throw error; - await jsm.consumers.update(JOBS_STREAM_NAME, name, { - filter_subject: subject, - ack_wait: nanos(Math.max(ackWaitMs, 1_000)), - max_deliver: maxDeliver, - }); - } - }; - - await Promise.all([ - ensureConsumer(WHISPER_CONSUMER_NAME, WHISPER_JOBS_SUBJECT, whisperTimeoutMs + 15_000, WHISPER_MAX_DELIVER), - ensureConsumer(LAYOUT_CONSUMER_NAME, LAYOUT_JOBS_SUBJECT, pdfTimeoutMs + 15_000, pdfAttempts), - ]); -} - export async function createComputeWorkerApp(options: CreateComputeWorkerAppOptions = {}): Promise { const port = options.port ?? readIntEnv('PORT', 8081); const host = options.host ?? (process.env.COMPUTE_WORKER_HOST?.trim() || '0.0.0.0'); @@ -560,11 +323,9 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti // the session via ensureConnected(). let session: NatsSession | null = null; let connecting: Promise | null = null; - let workerLoops: Promise[] = []; let idleTimer: NodeJS.Timeout | null = null; let orphanSweepTimer: NodeJS.Timeout | null = null; let stopping = false; - let loopStopRequested = 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 @@ -575,7 +336,6 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti let lastActivityAt = Date.now(); let lastActivityReason = 'startup'; let lastIdleStatusLogAt = 0; - const jobGate = new ConcurrencyGate(jobConcurrency); const markActivity = (reason: string): void => { lastActivityAt = Date.now(); @@ -641,7 +401,6 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti // Clear synchronously (before any await) so concurrent requests reconnect a // fresh session instead of using the connection we're about to close. session = null; - loopStopRequested = true; if (idleTimer) { clearInterval(idleTimer); idleTimer = null; @@ -655,9 +414,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti } catch { // ignore close errors } - await Promise.allSettled(workerLoops); - workerLoops = []; - loopStopRequested = false; + await workerLoops.stop(); app.log.info({ reason }, 'nats disconnected'); } @@ -668,15 +425,15 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti 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( + await ensureJetStreamResources({ jsm, whisperTimeoutMs, pdfTimeoutMs, pdfAttempts, - jobsStreamMaxBytes, - eventsStreamMaxBytes, + jobsMaxBytes: jobsStreamMaxBytes, + eventsMaxBytes: eventsStreamMaxBytes, natsReplicas, - ); + }); const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { replicas: natsReplicas, history: 1, @@ -690,7 +447,12 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti sessionGeneration += 1; orphanRecoveryDoneForGeneration = -1; markActivity('nats_connected'); - startWorkerLoops(next); + 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 @@ -709,82 +471,22 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti } } - const s3 = disableWorkers ? null : buildS3Client(); - const s3Bucket = disableWorkers ? '' : requireEnv('S3_BUCKET'); const s3Prefix = normalizeS3Prefix(process.env.S3_PREFIX); - - const ensureSafeKey = (key: string): string => { - const trimmed = key.trim(); - if (!trimmed.startsWith(`${s3Prefix}/`)) { - throw new Error('Object key prefix mismatch'); - } - return trimmed; + const storageDisabled = async (): Promise => { + throw new Error('S3 access is disabled for this worker app instance'); }; - - const readObjectByKey = async (key: string): Promise => { - if (!s3) { - throw new Error('S3 access is disabled for this worker app instance'); + const storage: ArtifactStorage = disableWorkers + ? { + readObject: storageDisabled, + objectExists: storageDisabled, + deleteObject: storageDisabled, + putParsedPdf: storageDisabled, } - const safeKey = ensureSafeKey(key); - const response = await s3.send(new GetObjectCommand({ - Bucket: s3Bucket, - Key: safeKey, - })); - const bytes = await bodyToBuffer(response.Body); - return toArrayBuffer(new Uint8Array(bytes)); - }; - - const objectExists = async (key: string): Promise => { - if (!s3) { - throw new Error('S3 access is disabled for this worker app instance'); - } - const safeKey = ensureSafeKey(key); - try { - await s3.send(new HeadObjectCommand({ - Bucket: s3Bucket, - Key: safeKey, - })); - return true; - } catch (error) { - const maybe = error as { name?: string; Code?: string; $metadata?: { httpStatusCode?: number } }; - if ( - maybe.$metadata?.httpStatusCode === 404 - || maybe.name === 'NotFound' - || maybe.name === 'NoSuchKey' - || maybe.Code === 'NotFound' - || maybe.Code === 'NoSuchKey' - ) { - return false; - } - throw error; - } - }; - - const deleteObjectByKey = async (key: string): Promise => { - if (!s3) { - throw new Error('S3 access is disabled for this worker app instance'); - } - await s3.send(new DeleteObjectCommand({ - Bucket: s3Bucket, - Key: ensureSafeKey(key), - })); - }; - - const putParsedObject = async (documentId: string, namespace: string | null, parsed: unknown): Promise => { - if (!s3) { - throw new Error('S3 access is disabled for this worker app instance'); - } - const key = parsedPdfArtifactKey({ documentId, namespace, prefix: s3Prefix }); - const body = Buffer.from(JSON.stringify(parsed)); - await s3.send(new PutObjectCommand({ - Bucket: s3Bucket, - Key: key, - Body: body, - ContentType: 'application/json', - ServerSideEncryption: 'AES256', - })); - return key; - }; + : createArtifactStorage({ + bucket: requireEnv('S3_BUCKET'), + prefix: s3Prefix, + client: createS3ClientFromEnv(requireEnv), + }); if (prewarmModels && !disableWorkers) { await ensureComputeModels(); @@ -1081,7 +783,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti namespace: parsed.data.namespace, prefix: s3Prefix, }); - const hasArtifact = await (options.routeDeps?.artifactExists ?? objectExists)(artifactKey); + 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; @@ -1257,451 +959,28 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti return reply; }); - const runWhisper = async ( - payload: WhisperAlignJobRequest, - queueWaitMs: number, - ): Promise => { - const parsed = alignSchema.parse(payload); + const jobHandlers = createJobHandlers({ + storage, + whisperTimeoutMs, + pdfTimeoutMs, + pdfHardCapMs, + }); - const s3FetchStartedAt = Date.now(); - const audioBuffer = await withTimeout( - readObjectByKey(parsed.audioObjectKey), - whisperTimeoutMs, - 'whisper s3 fetch', - ); - const s3FetchMs = Date.now() - s3FetchStartedAt; - - const computeStartedAt = Date.now(); - const result = await withTimeout( - runWhisperAlignmentFromAudioBuffer({ - audioBuffer, - text: parsed.text, - cacheKey: parsed.cacheKey, - lang: parsed.lang, - }), - whisperTimeoutMs, - 'whisper alignment job', - ); - - const computeMs = Date.now() - computeStartedAt; - return { - ...result, - timing: { - queueWaitMs, - s3FetchMs, - computeMs, - }, - }; - }; - - const runLayout = async ( - payload: PdfLayoutJobRequest, - queueWaitMs: number, - hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise }, - ): Promise => { - const parsed = layoutSchema.parse(payload); - - const s3FetchStartedAt = Date.now(); - const pdfBytes = await withTimeout( - readObjectByKey(parsed.documentObjectKey), - Math.max(pdfTimeoutMs, 1_000), - 'pdf s3 fetch', - ); - const s3FetchMs = Date.now() - s3FetchStartedAt; - - let lastTotalPages = 0; - let lastPagesParsed = 0; - const computeStartedAt = Date.now(); - const result = await withIdleTimeoutAndHardCap({ - idleTimeoutMs: Math.max(pdfTimeoutMs, 1_000), - hardCapMs: pdfHardCapMs, - label: 'pdf layout job', - run: async (touchProgress) => runPdfLayoutFromPdfBuffer({ - documentId: parsed.documentId, - pdfBytes, - onPageStarted: async ({ pageNumber, totalPages }) => { - touchProgress(); - lastTotalPages = totalPages; - if (!hooks?.onProgress) return; - await hooks.onProgress(buildInferProgressForPageStart({ - pageNumber, - totalPages, - })); - }, - onPageParsed: async ({ pageNumber, totalPages }) => { - touchProgress(); - lastTotalPages = totalPages; - lastPagesParsed = pageNumber; - if (!hooks?.onProgress) return; - await hooks.onProgress(buildInferProgressForPageParsed({ - pageNumber, - totalPages, - })); - }, - }), - }); - - const computeMs = Date.now() - computeStartedAt; - if (hooks?.onProgress && lastTotalPages > 0) { - await hooks.onProgress({ - totalPages: lastTotalPages, - pagesParsed: lastPagesParsed, - currentPage: lastPagesParsed || undefined, - phase: 'merge', - }); - } - const parsedObjectKey = await persistParsedPdfWhileSourceExists({ - sourceObjectKey: parsed.documentObjectKey, - sourceExists: objectExists, - putParsedObject: () => putParsedObject(parsed.documentId, parsed.namespace, result.parsed), - deleteParsedObject: deleteObjectByKey, - }); - return { - parsedObjectKey, - timing: { - queueWaitMs, - s3FetchMs, - computeMs, - }, - }; - }; - - type ProcessMessageInput = { - msg: JsMsg; - codec: JsonCodec>; - run: ( - payload: TPayload, - queueWaitMs: number, - hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise }, - ) => Promise; - workerLabel: string; - }; - - type ProcessMessageContext = { - decoded: QueuedJob; - workerLabel: string; - startedAt: number; - queueWaitTiming?: { queueWaitMs: number }; - latestProgress?: PdfLayoutProgress; - }; - - function extractTiming(result: unknown): WorkerJobTiming | undefined { - if (!result || typeof result !== 'object' || !('timing' in result)) return undefined; - return (result as { timing?: WorkerJobTiming }).timing; - } - - async function markRunning( - context: ProcessMessageContext, - updatedAt: number, - options?: { includeStartedAt?: boolean }, - ): Promise { - if (context.latestProgress) { - await orchestrator.markProgress({ - opId: context.decoded.opId, - progress: context.latestProgress, - updatedAt, - ...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}), - }); - return; - } - - await orchestrator.markRunning({ - opId: context.decoded.opId, - ...(options?.includeStartedAt === false ? {} : { startedAt: context.startedAt }), - updatedAt, - ...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}), - }); - } - - async function markProgress( - context: ProcessMessageContext, - progress: PdfLayoutProgress, - updatedAt: number, - ): Promise { - context.latestProgress = progress; - await markRunning(context, updatedAt); - } - - async function markTerminal(input: { - context: ProcessMessageContext; - status: 'succeeded' | 'failed'; - result?: TResult; - errorMessage?: string; - timing?: WorkerJobTiming; - updatedAt: number; - }): Promise { - if (input.status === 'succeeded') { - await orchestrator.markSucceeded({ - opId: input.context.decoded.opId, - result: input.result as WhisperAlignJobResult | PdfLayoutJobResult, - updatedAt: input.updatedAt, - ...(input.timing ? { timing: input.timing } : {}), - }); - return; - } - - await orchestrator.markFailed({ - opId: input.context.decoded.opId, - error: { message: input.errorMessage ?? 'unknown worker failure' }, - updatedAt: input.updatedAt, - ...(input.timing ? { timing: input.timing } : {}), - }); - } - - async function handleRetry(input: { - context: ProcessMessageContext | null; - msg: JsMsg; - errorMessage: string; - }): Promise { - const deliveryCount = input.msg.info.deliveryCount; - const kind = input.context?.decoded.kind ?? 'pdf_layout'; - const retryAction = decideRetryAction({ - kind, - deliveryCount, - pdfAttempts, - whisperMaxDeliver: WHISPER_MAX_DELIVER, - }); - const hasRetriesLeft = retryAction === 'nak_retry'; - const maxAttempts = kind === 'whisper_align' ? WHISPER_MAX_DELIVER : pdfAttempts; - - if (input.context) { - const now = Date.now(); - const retryTiming = buildQueueWaitTiming(input.context.decoded.queuedAt, now); - const persistOpUpdate = hasRetriesLeft - ? (input.context.latestProgress - ? orchestrator.markProgress({ - opId: input.context.decoded.opId, - progress: input.context.latestProgress, - updatedAt: now, - ...(retryTiming ? { timing: retryTiming } : {}), - }) - : orchestrator.markRunning({ - opId: input.context.decoded.opId, - updatedAt: now, - ...(retryTiming ? { timing: retryTiming } : {}), - })) - : markTerminal({ - context: input.context, - status: 'failed', - errorMessage: input.errorMessage, - updatedAt: now, - ...(retryTiming ? { timing: retryTiming } : {}), - }); - - await persistOpUpdate.catch((stateError) => { - app.log.error({ - worker: input.context?.workerLabel, - opId: input.context?.decoded.opId, - jobId: input.context?.decoded.jobId, - error: toErrorMessage(stateError), - }, 'failed to persist operation state'); - }); - } - - if (hasRetriesLeft) { - input.msg.nak(); - app.log.error({ - worker: input.context?.workerLabel, - kind: input.context?.decoded.kind, - opId: input.context?.decoded.opId, - jobId: input.context?.decoded.jobId, - status: 'running', - error: input.errorMessage, - deliveryCount, - maxAttempts, - retryAction: 'nack_retry', - }, 'job.terminal'); - return; - } - - input.msg.term(input.errorMessage); - app.log.error({ - worker: input.context?.workerLabel, - kind: input.context?.decoded.kind, - opId: input.context?.decoded.opId, - jobId: input.context?.decoded.jobId, - status: 'failed', - error: input.errorMessage, - deliveryCount, - maxAttempts, - retrySuppressed: kind === 'whisper_align' ? 'whisper_align' : undefined, - retryAction: 'term', - }, 'job.terminal'); - } - - async function processMessage(input: ProcessMessageInput): Promise { - let context: ProcessMessageContext | null = null; - let heartbeat: NodeJS.Timeout | null = null; - try { - const decoded = input.codec.decode(input.msg.data); - const startedAt = Date.now(); - context = { - decoded, - workerLabel: input.workerLabel, - startedAt, - queueWaitTiming: buildQueueWaitTiming(decoded.queuedAt, startedAt), - }; - - await markRunning(context, startedAt); - app.log.info({ - worker: input.workerLabel, - kind: decoded.kind, - opId: decoded.opId, - jobId: decoded.jobId, - queueWaitMs: context.queueWaitTiming?.queueWaitMs ?? null, - deliveryCount: input.msg.info.deliveryCount, - }, 'job.started'); - - heartbeat = setInterval(() => { - const now = Date.now(); - void markRunning(context!, now).catch((stateError) => { - app.log.error({ - worker: input.workerLabel, - opId: context?.decoded.opId, - jobId: context?.decoded.jobId, - error: toErrorMessage(stateError), - }, 'failed to persist operation heartbeat state'); - }); - }, RUNNING_HEARTBEAT_MS); - - const result = await input.run(decoded.payload, context.queueWaitTiming?.queueWaitMs ?? 0, { - onProgress: async (progress) => { - try { - input.msg.working(); - } catch (ackError) { - app.log.warn({ - worker: input.workerLabel, - kind: context?.decoded.kind, - opId: context?.decoded.opId, - jobId: context?.decoded.jobId, - error: toErrorMessage(ackError), - }, 'failed to extend JetStream ack wait on progress'); - } - await markProgress(context!, progress, Date.now()); - }, - }); - const resultTiming = extractTiming(result); - const now = Date.now(); - - await markTerminal({ - context, - status: 'succeeded', - result, - timing: resultTiming, - updatedAt: now, - }); - - input.msg.ack(); - const terminalDurationMs = safeDurationMs(context.startedAt, now); - const slowJobLogThresholdMs = SLOW_JOB_LOG_THRESHOLD_MS_BY_KIND[decoded.kind]; - if ((terminalDurationMs ?? 0) >= slowJobLogThresholdMs) { - app.log.info({ - worker: input.workerLabel, - kind: decoded.kind, - opId: decoded.opId, - jobId: decoded.jobId, - durationMs: terminalDurationMs ?? null, - timing: resultTiming ?? null, - }, 'job.stage'); - } - app.log.info({ - worker: input.workerLabel, - kind: decoded.kind, - opId: decoded.opId, - jobId: decoded.jobId, - status: 'succeeded', - durationMs: terminalDurationMs ?? null, - resultRef: extractResultRef(decoded.kind, result), - timing: resultTiming ?? null, - }, 'job.terminal'); - } catch (error) { - await handleRetry({ - context, - msg: input.msg, - errorMessage: toErrorMessage(error), - }); - } finally { - if (heartbeat) clearInterval(heartbeat); - } - } - - async function createWorkerLoop(input: { - owner: NatsSession; - consumer: Consumer; - codec: JsonCodec>; - run: ( - payload: TPayload, - queueWaitMs: number, - hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise }, - ) => Promise; - workerLabel: string; - }): Promise { - // Exit when the loop's connection is no longer the active session (idle - // disconnect, unexpected close, or replaced by a reconnect). - const detached = (): boolean => stopping || loopStopRequested || session !== input.owner; - while (!detached()) { - let msg: JsMsg | null = null; - try { - try { - // Bounded pull so the loop yields periodically and exits promptly when - // the session is torn down for idle (nc.close() rejects the pending pull). - msg = await input.consumer.next({ expires: PULL_EXPIRES_MS }); - } catch (error) { - if (detached()) return; - app.log.error({ error: toErrorMessage(error), worker: input.workerLabel }, 'worker pull failed'); - await sleep(LOOP_ERROR_BACKOFF_MS); - continue; - } - - // An empty pull is not activity; let the idle window advance. - if (!msg) continue; - markActivity(`job_received:${input.workerLabel}`); - inFlightJobs += 1; - await jobGate.acquire(); - if (detached()) { - return; - } - await processMessage({ - msg, - codec: input.codec, - run: input.run, - workerLabel: input.workerLabel, - }); - } finally { - if (msg) { - jobGate.release(); - inFlightJobs = Math.max(0, inFlightJobs - 1); - markActivity(`job_completed:${input.workerLabel}`); - } - } - } - } - - function startWorkerLoops(active: NatsSession): void { - if (disableWorkers) return; - // Always starts a fresh set bound to the new session. Any loops from a prior - // session self-terminate once they observe session !== their owner. - workerLoops = []; - loopStopRequested = false; - for (let i = 0; i < jobConcurrency; i += 1) { - workerLoops.push(createWorkerLoop({ - owner: active, - consumer: active.whisperConsumer, - codec: whisperJobCodec, - run: runWhisper, - workerLabel: `whisper-${i + 1}`, - })); - } - for (let i = 0; i < jobConcurrency; i += 1) { - workerLoops.push(createWorkerLoop({ - owner: active, - consumer: active.layoutConsumer, - codec: layoutJobCodec, - run: runLayout, - workerLabel: `layout-${i + 1}`, - })); - } - } + const workerLoops = createWorkerLoopController({ + orchestrator, + handlers: jobHandlers, + logger: app.log, + jobConcurrency, + pdfAttempts, + whisperCodec: whisperJobCodec, + pdfCodec: layoutJobCodec, + isOwnerActive: (owner) => session === owner, + isStopping: () => stopping, + markActivity, + onInFlightJobsChanged: (delta) => { + inFlightJobs = Math.max(0, inFlightJobs + delta); + }, + }); const close = async (): Promise => { if (stopping) return; @@ -1715,7 +994,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti orphanSweepTimer = null; } await app.close(); - await Promise.allSettled(workerLoops); + await workerLoops.stop(); const current = session; session = null; if (current) { diff --git a/compute-worker/src/compute/api-contracts/index.ts b/compute-worker/src/api/contracts.ts similarity index 87% rename from compute-worker/src/compute/api-contracts/index.ts rename to compute-worker/src/api/contracts.ts index 487a618..d6e87f3 100644 --- a/compute-worker/src/compute/api-contracts/index.ts +++ b/compute-worker/src/api/contracts.ts @@ -1,24 +1,24 @@ -import type { TTSSentenceAlignment } from '../types/tts'; -import type { ParsedPdfDocument } from '../types/parsed-pdf'; +import type { TTSSentenceAlignment } from '../inference/types/tts'; +import type { ParsedPdfDocument } from '../inference/types/parsed-pdf'; export type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment, TTSSentenceWord, -} from '../types/tts'; +} from '../inference/types/tts'; export type { ParsedPdfBlockKind, ParsedPdfBlockFragment, ParsedPdfBlock, ParsedPdfPage, ParsedPdfDocument, -} from '../types/parsed-pdf'; +} from '../inference/types/parsed-pdf'; export const ALIGN_QUEUE_NAME = 'whisper-align'; export const PDF_LAYOUT_QUEUE_NAME = 'pdf-layout'; -export { PDF_PARSER_VERSION } from '../pdf/parser-version'; -export { encodeParserVersion } from '../pdf/parser-version-key'; +export { PDF_PARSER_VERSION } from '../inference/pdf/parser-version'; +export { encodeParserVersion } from '../inference/pdf/parser-version-key'; export interface WhisperAlignJobBase { text: string; diff --git a/compute-worker/src/api/operation-keys.ts b/compute-worker/src/api/operation-keys.ts index 9c990b5..a6763fe 100644 --- a/compute-worker/src/api/operation-keys.ts +++ b/compute-worker/src/api/operation-keys.ts @@ -1,5 +1,5 @@ import { createHash } from 'node:crypto'; -import { PDF_PARSER_VERSION } from '../compute/pdf/parser-version'; +import { PDF_PARSER_VERSION } from '../inference/pdf/parser-version'; function sha256Hex(input: string): string { return createHash('sha256').update(input).digest('hex'); diff --git a/compute-worker/src/api/public-operation.ts b/compute-worker/src/api/public-operation.ts index 9ab98a4..ac9b43b 100644 --- a/compute-worker/src/api/public-operation.ts +++ b/compute-worker/src/api/public-operation.ts @@ -1,4 +1,4 @@ -import type { WorkerOperationState } from '../compute/api-contracts'; +import type { WorkerOperationState } from '../api/contracts'; import { pdfSubjectFromOperationKey } from './operation-keys'; export type PublicOperationSubject = diff --git a/compute-worker/src/compute/config/cpu-budget.ts b/compute-worker/src/inference/config/cpu-budget.ts similarity index 100% rename from compute-worker/src/compute/config/cpu-budget.ts rename to compute-worker/src/inference/config/cpu-budget.ts diff --git a/compute-worker/src/compute/config/timeout.ts b/compute-worker/src/inference/config/timeout.ts similarity index 100% rename from compute-worker/src/compute/config/timeout.ts rename to compute-worker/src/inference/config/timeout.ts diff --git a/compute-worker/src/compute/index.ts b/compute-worker/src/inference/index.ts similarity index 93% rename from compute-worker/src/compute/index.ts rename to compute-worker/src/inference/index.ts index ea7b118..255fdce 100644 --- a/compute-worker/src/compute/index.ts +++ b/compute-worker/src/inference/index.ts @@ -1,4 +1,4 @@ -export * from './api-contracts'; +export * from '../api/contracts'; export { getComputeJobConcurrency, getAvailableCpuCores, @@ -23,4 +23,4 @@ export { normalizeTextItemsForLayout } from './pdf/normalize-text'; export { mapWordsToSentenceOffsets, type WhisperWord } from './whisper/alignment-map'; export { buildGoertzelCoefficients, goertzelPower } from './whisper/spectral'; export { buildWordsFromTimestampedTokens, extractTokenStartTimestamps } from './whisper/token-timestamps'; -export * from './control-plane'; +export * from '../operations'; diff --git a/compute-worker/src/compute/local-runtime.ts b/compute-worker/src/inference/local-runtime.ts similarity index 100% rename from compute-worker/src/compute/local-runtime.ts rename to compute-worker/src/inference/local-runtime.ts diff --git a/compute-worker/src/compute/pdf/assets/LICENSE.txt b/compute-worker/src/inference/pdf/assets/LICENSE.txt similarity index 100% rename from compute-worker/src/compute/pdf/assets/LICENSE.txt rename to compute-worker/src/inference/pdf/assets/LICENSE.txt diff --git a/compute-worker/src/compute/pdf/assets/manifest.json b/compute-worker/src/inference/pdf/assets/manifest.json similarity index 100% rename from compute-worker/src/compute/pdf/assets/manifest.json rename to compute-worker/src/inference/pdf/assets/manifest.json diff --git a/compute-worker/src/compute/pdf/merge.ts b/compute-worker/src/inference/pdf/merge.ts similarity index 100% rename from compute-worker/src/compute/pdf/merge.ts rename to compute-worker/src/inference/pdf/merge.ts diff --git a/compute-worker/src/compute/pdf/model.ts b/compute-worker/src/inference/pdf/model.ts similarity index 100% rename from compute-worker/src/compute/pdf/model.ts rename to compute-worker/src/inference/pdf/model.ts diff --git a/compute-worker/src/compute/pdf/normalize-text.ts b/compute-worker/src/inference/pdf/normalize-text.ts similarity index 100% rename from compute-worker/src/compute/pdf/normalize-text.ts rename to compute-worker/src/inference/pdf/normalize-text.ts diff --git a/compute-worker/src/compute/pdf/parse.ts b/compute-worker/src/inference/pdf/parse.ts similarity index 100% rename from compute-worker/src/compute/pdf/parse.ts rename to compute-worker/src/inference/pdf/parse.ts diff --git a/compute-worker/src/compute/pdf/parser-version-key.ts b/compute-worker/src/inference/pdf/parser-version-key.ts similarity index 100% rename from compute-worker/src/compute/pdf/parser-version-key.ts rename to compute-worker/src/inference/pdf/parser-version-key.ts diff --git a/compute-worker/src/compute/pdf/parser-version.ts b/compute-worker/src/inference/pdf/parser-version.ts similarity index 100% rename from compute-worker/src/compute/pdf/parser-version.ts rename to compute-worker/src/inference/pdf/parser-version.ts diff --git a/compute-worker/src/compute/pdf/pdfjs-runtime.ts b/compute-worker/src/inference/pdf/pdfjs-runtime.ts similarity index 100% rename from compute-worker/src/compute/pdf/pdfjs-runtime.ts rename to compute-worker/src/inference/pdf/pdfjs-runtime.ts diff --git a/compute-worker/src/compute/pdf/render.ts b/compute-worker/src/inference/pdf/render.ts similarity index 100% rename from compute-worker/src/compute/pdf/render.ts rename to compute-worker/src/inference/pdf/render.ts diff --git a/compute-worker/src/compute/pdf/runLayoutModel.ts b/compute-worker/src/inference/pdf/runLayoutModel.ts similarity index 100% rename from compute-worker/src/compute/pdf/runLayoutModel.ts rename to compute-worker/src/inference/pdf/runLayoutModel.ts diff --git a/compute-worker/src/compute/pdf/stitch.ts b/compute-worker/src/inference/pdf/stitch.ts similarity index 100% rename from compute-worker/src/compute/pdf/stitch.ts rename to compute-worker/src/inference/pdf/stitch.ts diff --git a/compute-worker/src/compute/pdf/types.ts b/compute-worker/src/inference/pdf/types.ts similarity index 100% rename from compute-worker/src/compute/pdf/types.ts rename to compute-worker/src/inference/pdf/types.ts diff --git a/compute-worker/src/compute/platform/docstore.ts b/compute-worker/src/inference/platform/docstore.ts similarity index 100% rename from compute-worker/src/compute/platform/docstore.ts rename to compute-worker/src/inference/platform/docstore.ts diff --git a/compute-worker/src/compute/platform/ffmpeg.ts b/compute-worker/src/inference/platform/ffmpeg.ts similarity index 100% rename from compute-worker/src/compute/platform/ffmpeg.ts rename to compute-worker/src/inference/platform/ffmpeg.ts diff --git a/compute-worker/src/compute/types/index.ts b/compute-worker/src/inference/types/index.ts similarity index 100% rename from compute-worker/src/compute/types/index.ts rename to compute-worker/src/inference/types/index.ts diff --git a/compute-worker/src/compute/types/parsed-pdf.ts b/compute-worker/src/inference/types/parsed-pdf.ts similarity index 100% rename from compute-worker/src/compute/types/parsed-pdf.ts rename to compute-worker/src/inference/types/parsed-pdf.ts diff --git a/compute-worker/src/compute/types/tts.ts b/compute-worker/src/inference/types/tts.ts similarity index 100% rename from compute-worker/src/compute/types/tts.ts rename to compute-worker/src/inference/types/tts.ts diff --git a/compute-worker/src/compute/whisper/align.ts b/compute-worker/src/inference/whisper/align.ts similarity index 100% rename from compute-worker/src/compute/whisper/align.ts rename to compute-worker/src/inference/whisper/align.ts diff --git a/compute-worker/src/compute/whisper/alignment-map.ts b/compute-worker/src/inference/whisper/alignment-map.ts similarity index 100% rename from compute-worker/src/compute/whisper/alignment-map.ts rename to compute-worker/src/inference/whisper/alignment-map.ts diff --git a/compute-worker/src/compute/whisper/assets/LICENSE.txt b/compute-worker/src/inference/whisper/assets/LICENSE.txt similarity index 100% rename from compute-worker/src/compute/whisper/assets/LICENSE.txt rename to compute-worker/src/inference/whisper/assets/LICENSE.txt diff --git a/compute-worker/src/compute/whisper/assets/manifest.json b/compute-worker/src/inference/whisper/assets/manifest.json similarity index 100% rename from compute-worker/src/compute/whisper/assets/manifest.json rename to compute-worker/src/inference/whisper/assets/manifest.json diff --git a/compute-worker/src/compute/whisper/assets/mel_filters.npz b/compute-worker/src/inference/whisper/assets/mel_filters.npz similarity index 100% rename from compute-worker/src/compute/whisper/assets/mel_filters.npz rename to compute-worker/src/inference/whisper/assets/mel_filters.npz diff --git a/compute-worker/src/compute/whisper/model.ts b/compute-worker/src/inference/whisper/model.ts similarity index 100% rename from compute-worker/src/compute/whisper/model.ts rename to compute-worker/src/inference/whisper/model.ts diff --git a/compute-worker/src/compute/whisper/spectral.ts b/compute-worker/src/inference/whisper/spectral.ts similarity index 100% rename from compute-worker/src/compute/whisper/spectral.ts rename to compute-worker/src/inference/whisper/spectral.ts diff --git a/compute-worker/src/compute/whisper/token-timestamps.ts b/compute-worker/src/inference/whisper/token-timestamps.ts similarity index 100% rename from compute-worker/src/compute/whisper/token-timestamps.ts rename to compute-worker/src/inference/whisper/token-timestamps.ts diff --git a/compute-worker/src/control-plane/json-codec.ts b/compute-worker/src/infrastructure/json-codec.ts similarity index 100% rename from compute-worker/src/control-plane/json-codec.ts rename to compute-worker/src/infrastructure/json-codec.ts diff --git a/compute-worker/src/control-plane/jetstream.ts b/compute-worker/src/infrastructure/nats-adapters.ts similarity index 99% rename from compute-worker/src/control-plane/jetstream.ts rename to compute-worker/src/infrastructure/nats-adapters.ts index e5367eb..00d2526 100644 --- a/compute-worker/src/control-plane/jetstream.ts +++ b/compute-worker/src/infrastructure/nats-adapters.ts @@ -8,12 +8,12 @@ import type { OperationState, OperationStateStore, QueuedOperation, -} from '../compute/control-plane'; +} from '../operations'; import type { PdfLayoutJobRequest, WhisperAlignJobRequest, WorkerOperationKind, -} from '../compute/api-contracts'; +} from '../api/contracts'; import { createJsonCodec } from './json-codec'; export interface KvEntryLike { diff --git a/compute-worker/src/infrastructure/nats.ts b/compute-worker/src/infrastructure/nats.ts new file mode 100644 index 0000000..429a03d --- /dev/null +++ b/compute-worker/src/infrastructure/nats.ts @@ -0,0 +1,104 @@ +import { + AckPolicy, + DeliverPolicy, + ReplayPolicy, + RetentionPolicy, + StorageType, + type JetStreamManager, +} from '@nats-io/jetstream'; +import { nanos } from '@nats-io/transport-node'; +import { OP_EVENTS_SUBJECT_WILDCARD } from './nats-adapters'; + +export const JOBS_STREAM_NAME = 'compute_jobs'; +export const WHISPER_JOBS_SUBJECT = 'jobs.whisper'; +export const LAYOUT_JOBS_SUBJECT = 'jobs.layout'; +export const WHISPER_CONSUMER_NAME = 'compute_whisper'; +export const LAYOUT_CONSUMER_NAME = 'compute_layout'; +export const EVENTS_STREAM_NAME = 'compute_events'; +export const COMPUTE_STATE_BUCKET = 'compute_state'; +export const COMPUTE_STATE_TTL_MS = 24 * 60 * 60 * 1000; +export const NATS_API_TIMEOUT_MS = 60_000; + +const WHISPER_MAX_DELIVER = 1; + +function isAlreadyExistsError(error: unknown): boolean { + const message = (error instanceof Error ? error.message : String(error)).toLowerCase(); + return message.includes('already in use') || message.includes('already exists'); +} + +export async function ensureJetStreamResources(input: { + jsm: JetStreamManager; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + pdfAttempts: number; + jobsMaxBytes: number; + eventsMaxBytes: number; + natsReplicas: number; +}): Promise { + const streamConfig = { + name: JOBS_STREAM_NAME, + subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT], + retention: RetentionPolicy.Workqueue, + storage: StorageType.File, + max_bytes: input.jobsMaxBytes, + num_replicas: input.natsReplicas, + }; + try { + await input.jsm.streams.add(streamConfig); + } catch (error) { + if (!isAlreadyExistsError(error)) throw error; + await input.jsm.streams.update(JOBS_STREAM_NAME, { + subjects: streamConfig.subjects, + max_bytes: input.jobsMaxBytes, + num_replicas: input.natsReplicas, + }); + } + + const eventsStreamConfig = { + name: EVENTS_STREAM_NAME, + subjects: [OP_EVENTS_SUBJECT_WILDCARD], + retention: RetentionPolicy.Limits, + storage: StorageType.File, + max_bytes: input.eventsMaxBytes, + max_age: nanos(COMPUTE_STATE_TTL_MS), + num_replicas: input.natsReplicas, + }; + try { + await input.jsm.streams.add(eventsStreamConfig); + } catch (error) { + if (!isAlreadyExistsError(error)) throw error; + await input.jsm.streams.update(EVENTS_STREAM_NAME, { + subjects: eventsStreamConfig.subjects, + max_bytes: input.eventsMaxBytes, + max_age: eventsStreamConfig.max_age, + num_replicas: input.natsReplicas, + }); + } + + const ensureConsumer = async (name: string, subject: string, ackWaitMs: number, maxDeliver: number) => { + const config = { + durable_name: name, + ack_policy: AckPolicy.Explicit, + deliver_policy: DeliverPolicy.All, + replay_policy: ReplayPolicy.Instant, + filter_subject: subject, + ack_wait: nanos(Math.max(ackWaitMs, 1_000)), + max_deliver: maxDeliver, + }; + try { + await input.jsm.consumers.add(JOBS_STREAM_NAME, config); + } catch (error) { + if (!isAlreadyExistsError(error)) throw error; + await input.jsm.consumers.update(JOBS_STREAM_NAME, name, { + filter_subject: subject, + ack_wait: config.ack_wait, + max_deliver: maxDeliver, + }); + } + }; + + await Promise.all([ + ensureConsumer(WHISPER_CONSUMER_NAME, WHISPER_JOBS_SUBJECT, input.whisperTimeoutMs + 15_000, WHISPER_MAX_DELIVER), + ensureConsumer(LAYOUT_CONSUMER_NAME, LAYOUT_JOBS_SUBJECT, input.pdfTimeoutMs + 15_000, input.pdfAttempts), + ]); +} diff --git a/compute-worker/src/infrastructure/storage.ts b/compute-worker/src/infrastructure/storage.ts new file mode 100644 index 0000000..1190da3 --- /dev/null +++ b/compute-worker/src/infrastructure/storage.ts @@ -0,0 +1,127 @@ +import { + DeleteObjectCommand, + GetObjectCommand, + HeadObjectCommand, + PutObjectCommand, + S3Client, +} from '@aws-sdk/client-s3'; +import { parsedPdfArtifactKey } from '../storage/artifact-addressing'; + +export interface ArtifactStorage { + readObject(key: string): Promise; + objectExists(key: string): Promise; + deleteObject(key: string): Promise; + putParsedPdf(documentId: string, namespace: string | null, parsed: unknown): Promise; +} + +export interface ArtifactStorageConfig { + bucket: string; + prefix: string; + client: S3Client; +} + +function bodyToBuffer(body: unknown): Promise | Buffer { + if (!body) return Buffer.alloc(0); + if (body instanceof Uint8Array) return Buffer.from(body); + if (ArrayBuffer.isView(body)) return Buffer.from(body.buffer, body.byteOffset, body.byteLength); + if (body instanceof ArrayBuffer) return Buffer.from(body); + if (typeof body === 'object' && body !== null && 'transformToByteArray' in body) { + const maybe = body as { transformToByteArray?: () => Promise }; + if (typeof maybe.transformToByteArray === 'function') { + return maybe.transformToByteArray().then((bytes) => Buffer.from(bytes)); + } + } + if (typeof body === 'object' && body !== null && 'on' in body) { + return (async () => { + const chunks: Buffer[] = []; + for await (const chunk of body as NodeJS.ReadableStream) { + if (Buffer.isBuffer(chunk)) chunks.push(chunk); + else if (typeof chunk === 'string') chunks.push(Buffer.from(chunk)); + else chunks.push(Buffer.from(chunk as Uint8Array)); + } + return Buffer.concat(chunks); + })(); + } + throw new Error('Unsupported S3 response body type'); +} + +function toArrayBuffer(bytes: Uint8Array): ArrayBuffer { + const copy = new Uint8Array(bytes.byteLength); + copy.set(bytes); + return copy.buffer; +} + +function isNotFound(error: unknown): boolean { + const maybe = error as { name?: string; Code?: string; $metadata?: { httpStatusCode?: number } }; + return maybe.$metadata?.httpStatusCode === 404 + || maybe.name === 'NotFound' + || maybe.name === 'NoSuchKey' + || maybe.Code === 'NotFound' + || maybe.Code === 'NoSuchKey'; +} + +export function normalizeS3Prefix(prefix: string | undefined): string { + const value = (prefix || 'openreader').trim(); + return value ? value.replace(/^\/+|\/+$/g, '') : 'openreader'; +} + +export function createS3ClientFromEnv(requireEnv: (name: string) => string): S3Client { + return new S3Client({ + region: requireEnv('S3_REGION'), + endpoint: process.env.S3_ENDPOINT?.trim() || undefined, + forcePathStyle: ['1', 'true', 'yes', 'on'].includes(process.env.S3_FORCE_PATH_STYLE?.trim().toLowerCase() ?? ''), + requestChecksumCalculation: 'WHEN_REQUIRED', + responseChecksumValidation: 'WHEN_REQUIRED', + credentials: { + accessKeyId: requireEnv('S3_ACCESS_KEY_ID'), + secretAccessKey: requireEnv('S3_SECRET_ACCESS_KEY'), + }, + }); +} + +export function createArtifactStorage(config: ArtifactStorageConfig): ArtifactStorage { + const safeKey = (key: string): string => { + const trimmed = key.trim(); + if (!trimmed.startsWith(`${config.prefix}/`)) throw new Error('Object key prefix mismatch'); + return trimmed; + }; + + return { + async readObject(key) { + const response = await config.client.send(new GetObjectCommand({ + Bucket: config.bucket, + Key: safeKey(key), + })); + return toArrayBuffer(new Uint8Array(await bodyToBuffer(response.Body))); + }, + async objectExists(key) { + try { + await config.client.send(new HeadObjectCommand({ + Bucket: config.bucket, + Key: safeKey(key), + })); + return true; + } catch (error) { + if (isNotFound(error)) return false; + throw error; + } + }, + async deleteObject(key) { + await config.client.send(new DeleteObjectCommand({ + Bucket: config.bucket, + Key: safeKey(key), + })); + }, + async putParsedPdf(documentId, namespace, parsed) { + const key = parsedPdfArtifactKey({ documentId, namespace, prefix: config.prefix }); + await config.client.send(new PutObjectCommand({ + Bucket: config.bucket, + Key: key, + Body: Buffer.from(JSON.stringify(parsed)), + ContentType: 'application/json', + ServerSideEncryption: 'AES256', + })); + return key; + }, + }; +} diff --git a/compute-worker/src/jobs/handlers.ts b/compute-worker/src/jobs/handlers.ts new file mode 100644 index 0000000..90a4068 --- /dev/null +++ b/compute-worker/src/jobs/handlers.ts @@ -0,0 +1,126 @@ +import { z } from 'zod'; +import { + runPdfLayoutFromPdfBuffer, + runWhisperAlignmentFromAudioBuffer, +} from '../inference/local-runtime'; +import { withIdleTimeoutAndHardCap, withTimeout } from '../inference'; +import type { + PdfLayoutJobRequest, + PdfLayoutJobResult, + PdfLayoutProgress, + WhisperAlignJobRequest, + WhisperAlignJobResult, +} from '../api/contracts'; +import type { ArtifactStorage } from '../infrastructure/storage'; +import { persistParsedPdfWhileSourceExists } from './pdf-artifact-persistence'; +import { buildInferProgressForPageParsed, buildInferProgressForPageStart } from './pdf-progress'; + +const whisperRequestSchema = z.object({ + text: z.string().trim().min(1), + lang: z.string().trim().min(1).max(16).optional(), + cacheKey: z.string().trim().min(1).max(256).optional(), + audioObjectKey: z.string().trim().min(1).max(2048), +}); + +const pdfRequestSchema = z.object({ + documentId: z.string().trim().min(1), + namespace: z.string().trim().min(1).max(128).nullable(), + documentObjectKey: z.string().trim().min(1).max(2048), +}); + +export interface JobHandlers { + runWhisper(payload: WhisperAlignJobRequest, queueWaitMs: number): Promise; + runPdfLayout( + payload: PdfLayoutJobRequest, + queueWaitMs: number, + hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise }, + ): Promise; +} + +export function createJobHandlers(input: { + storage: ArtifactStorage; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + pdfHardCapMs: number; +}): JobHandlers { + return { + async runWhisper(payload, queueWaitMs) { + const parsed = whisperRequestSchema.parse(payload); + const s3FetchStartedAt = Date.now(); + const audioBuffer = await withTimeout( + input.storage.readObject(parsed.audioObjectKey), + input.whisperTimeoutMs, + 'whisper s3 fetch', + ); + const s3FetchMs = Date.now() - s3FetchStartedAt; + const computeStartedAt = Date.now(); + const result = await withTimeout( + runWhisperAlignmentFromAudioBuffer({ + audioBuffer, + text: parsed.text, + cacheKey: parsed.cacheKey, + lang: parsed.lang, + }), + input.whisperTimeoutMs, + 'whisper alignment job', + ); + return { + ...result, + timing: { queueWaitMs, s3FetchMs, computeMs: Date.now() - computeStartedAt }, + }; + }, + + async runPdfLayout(payload, queueWaitMs, hooks) { + const parsed = pdfRequestSchema.parse(payload); + const s3FetchStartedAt = Date.now(); + const pdfBytes = await withTimeout( + input.storage.readObject(parsed.documentObjectKey), + Math.max(input.pdfTimeoutMs, 1_000), + 'pdf s3 fetch', + ); + const s3FetchMs = Date.now() - s3FetchStartedAt; + let lastTotalPages = 0; + let lastPagesParsed = 0; + const computeStartedAt = Date.now(); + const result = await withIdleTimeoutAndHardCap({ + idleTimeoutMs: Math.max(input.pdfTimeoutMs, 1_000), + hardCapMs: input.pdfHardCapMs, + label: 'pdf layout job', + run: async (touchProgress) => runPdfLayoutFromPdfBuffer({ + documentId: parsed.documentId, + pdfBytes, + onPageStarted: async ({ pageNumber, totalPages }) => { + touchProgress(); + lastTotalPages = totalPages; + await hooks?.onProgress?.(buildInferProgressForPageStart({ pageNumber, totalPages })); + }, + onPageParsed: async ({ pageNumber, totalPages }) => { + touchProgress(); + lastTotalPages = totalPages; + lastPagesParsed = pageNumber; + await hooks?.onProgress?.(buildInferProgressForPageParsed({ pageNumber, totalPages })); + }, + }), + }); + const computeMs = Date.now() - computeStartedAt; + if (hooks?.onProgress && lastTotalPages > 0) { + await hooks.onProgress({ + totalPages: lastTotalPages, + pagesParsed: lastPagesParsed, + currentPage: lastPagesParsed || undefined, + phase: 'merge', + }); + } + const parsedObjectKey = await persistParsedPdfWhileSourceExists({ + sourceObjectKey: parsed.documentObjectKey, + sourceExists: input.storage.objectExists, + putParsedObject: () => input.storage.putParsedPdf(parsed.documentId, parsed.namespace, result.parsed), + deleteParsedObject: input.storage.deleteObject, + }); + return { + parsedObjectKey, + timing: { queueWaitMs, s3FetchMs, computeMs }, + }; + }, + }; +} diff --git a/compute-worker/src/pdf-artifact-persistence.ts b/compute-worker/src/jobs/pdf-artifact-persistence.ts similarity index 100% rename from compute-worker/src/pdf-artifact-persistence.ts rename to compute-worker/src/jobs/pdf-artifact-persistence.ts diff --git a/compute-worker/src/pdf-progress.ts b/compute-worker/src/jobs/pdf-progress.ts similarity index 89% rename from compute-worker/src/pdf-progress.ts rename to compute-worker/src/jobs/pdf-progress.ts index cc5c025..4aefe08 100644 --- a/compute-worker/src/pdf-progress.ts +++ b/compute-worker/src/jobs/pdf-progress.ts @@ -1,4 +1,4 @@ -import type { PdfLayoutProgress } from './compute/api-contracts'; +import type { PdfLayoutProgress } from '../api/contracts'; export function buildInferProgressForPageStart(input: { pageNumber: number; diff --git a/compute-worker/src/worker-loop-policy.ts b/compute-worker/src/jobs/worker-loop-policy.ts similarity index 91% rename from compute-worker/src/worker-loop-policy.ts rename to compute-worker/src/jobs/worker-loop-policy.ts index bce8653..ce3547a 100644 --- a/compute-worker/src/worker-loop-policy.ts +++ b/compute-worker/src/jobs/worker-loop-policy.ts @@ -1,4 +1,4 @@ -import type { WorkerOperationKind } from './compute/api-contracts'; +import type { WorkerOperationKind } from '../api/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 new file mode 100644 index 0000000..194bf62 --- /dev/null +++ b/compute-worker/src/jobs/worker-loop.ts @@ -0,0 +1,312 @@ +import type { Consumer, JsMsg } from '@nats-io/jetstream'; +import type { + PdfLayoutJobRequest, + PdfLayoutJobResult, + PdfLayoutProgress, + WhisperAlignJobRequest, + WhisperAlignJobResult, + WorkerJobTiming, + WorkerOperationKind, +} from '../api/contracts'; +import type { JsonCodec } from '../infrastructure/json-codec'; +import type { JobHandlers } from './handlers'; +import { buildQueueWaitTiming, decideRetryAction } from './worker-loop-policy'; + +const LOOP_ERROR_BACKOFF_MS = 500; +const RUNNING_HEARTBEAT_MS = 5000; +const PULL_EXPIRES_MS = 5_000; +const WHISPER_MAX_DELIVER = 1; +const SLOW_JOB_LOG_THRESHOLD_MS_BY_KIND: Record = { + whisper_align: 15_000, + pdf_layout: 120_000, +}; + +export interface QueuedJob { + jobId: string; + opId: string; + opKey: string; + kind: WorkerOperationKind; + queuedAt: number; + payload: TPayload; +} + +interface WorkerLoopOrchestrator { + markRunning(input: { opId: string; startedAt?: number; updatedAt?: number; timing?: WorkerJobTiming }): Promise; + markProgress(input: { + opId: string; + progress: PdfLayoutProgress; + 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; +} + +interface WorkerLogger { + info(data: unknown, message?: string): void; + warn(data: unknown, message?: string): void; + error(data: unknown, message?: string): void; +} + +class ConcurrencyGate { + private inFlight = 0; + private readonly queue: Array<() => void> = []; + + constructor(private readonly limit: number) {} + + async acquire(): Promise { + if (this.inFlight < this.limit) { + this.inFlight += 1; + return; + } + await new Promise((resolve) => { + this.queue.push(() => { + this.inFlight += 1; + resolve(); + }); + }); + } + + release(): void { + this.inFlight = Math.max(0, this.inFlight - 1); + this.queue.shift()?.(); + } +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +function toErrorMessage(error: unknown): string { + return error instanceof Error && error.message ? error.message : String(error); +} + +function safeDurationMs(start: number, end: number): number { + return Math.max(0, Math.floor(end - start)); +} + +function extractTiming(result: unknown): WorkerJobTiming | undefined { + if (!result || typeof result !== 'object' || !('timing' in result)) return undefined; + return (result as { timing?: WorkerJobTiming }).timing; +} + +function extractResultRef(kind: WorkerOperationKind, result: unknown): string | undefined { + if (kind !== 'pdf_layout' || !result || typeof result !== 'object') return undefined; + const maybe = result as { parsedObjectKey?: unknown }; + return typeof maybe.parsedObjectKey === 'string' ? maybe.parsedObjectKey : undefined; +} + +export function createWorkerLoopController(input: { + orchestrator: WorkerLoopOrchestrator; + handlers: JobHandlers; + logger: WorkerLogger; + jobConcurrency: number; + pdfAttempts: number; + whisperCodec: JsonCodec>; + pdfCodec: JsonCodec>; + isOwnerActive: (owner: object) => boolean; + isStopping: () => boolean; + markActivity: (reason: string) => void; + onInFlightJobsChanged: (delta: number) => void; +}) { + const gate = new ConcurrencyGate(Math.max(1, Math.floor(input.jobConcurrency))); + let loops: Promise[] = []; + let stopRequested = false; + + type Context = { + decoded: QueuedJob; + workerLabel: string; + startedAt: number; + queueWaitTiming?: { queueWaitMs: number }; + latestProgress?: PdfLayoutProgress; + }; + + const markRunning = async (context: Context, updatedAt: number): Promise => { + if (context.latestProgress) { + await input.orchestrator.markProgress({ + opId: context.decoded.opId, + progress: context.latestProgress, + updatedAt, + ...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}), + }); + return; + } + await input.orchestrator.markRunning({ + opId: context.decoded.opId, + startedAt: context.startedAt, + updatedAt, + ...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}), + }); + }; + + const processMessage = async (work: { + msg: JsMsg; + codec: JsonCodec>; + run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout']; + workerLabel: string; + }): Promise => { + let context: Context | null = null; + let heartbeat: NodeJS.Timeout | null = null; + try { + const decoded = work.codec.decode(work.msg.data); + const startedAt = Date.now(); + context = { + decoded, + workerLabel: work.workerLabel, + startedAt, + queueWaitTiming: buildQueueWaitTiming(decoded.queuedAt, startedAt), + }; + await markRunning(context, startedAt); + input.logger.info({ + worker: work.workerLabel, + kind: decoded.kind, + opId: decoded.opId, + jobId: decoded.jobId, + queueWaitMs: context.queueWaitTiming?.queueWaitMs ?? null, + deliveryCount: work.msg.info.deliveryCount, + }, 'job.started'); + heartbeat = setInterval(() => { + void markRunning(context!, Date.now()).catch((error) => { + input.logger.error({ + worker: work.workerLabel, + opId: context?.decoded.opId, + jobId: context?.decoded.jobId, + error: toErrorMessage(error), + }, 'failed to persist operation heartbeat state'); + }); + }, RUNNING_HEARTBEAT_MS); + const result = await work.run(decoded.payload as never, context.queueWaitTiming?.queueWaitMs ?? 0, { + onProgress: async (progress) => { + try { + work.msg.working(); + } catch (error) { + input.logger.warn({ + worker: work.workerLabel, + kind: context?.decoded.kind, + opId: context?.decoded.opId, + jobId: context?.decoded.jobId, + error: toErrorMessage(error), + }, 'failed to extend JetStream ack wait on progress'); + } + context!.latestProgress = progress; + await markRunning(context!, Date.now()); + }, + }); + const timing = extractTiming(result); + const now = Date.now(); + await input.orchestrator.markSucceeded({ + opId: decoded.opId, + result, + updatedAt: now, + ...(timing ? { timing } : {}), + }); + work.msg.ack(); + const durationMs = safeDurationMs(startedAt, now); + if (durationMs >= SLOW_JOB_LOG_THRESHOLD_MS_BY_KIND[decoded.kind]) { + input.logger.info({ worker: work.workerLabel, kind: decoded.kind, opId: decoded.opId, jobId: decoded.jobId, durationMs, timing: timing ?? null }, 'job.stage'); + } + input.logger.info({ + worker: work.workerLabel, + kind: decoded.kind, + opId: decoded.opId, + jobId: decoded.jobId, + status: 'succeeded', + durationMs, + resultRef: extractResultRef(decoded.kind, result), + timing: timing ?? null, + }, 'job.terminal'); + } catch (error) { + const errorMessage = toErrorMessage(error); + const deliveryCount = work.msg.info.deliveryCount; + const kind = context?.decoded.kind ?? 'pdf_layout'; + const action = decideRetryAction({ kind, deliveryCount, pdfAttempts: input.pdfAttempts, whisperMaxDeliver: WHISPER_MAX_DELIVER }); + const timing = context ? buildQueueWaitTiming(context.decoded.queuedAt, Date.now()) : undefined; + if (context) { + const update = action === 'nak_retry' + ? markRunning(context, Date.now()) + : input.orchestrator.markFailed({ + opId: context.decoded.opId, + error: { message: errorMessage }, + updatedAt: Date.now(), + ...(timing ? { timing } : {}), + }); + await update.catch((stateError) => input.logger.error({ + worker: context?.workerLabel, + opId: context?.decoded.opId, + jobId: context?.decoded.jobId, + error: toErrorMessage(stateError), + }, 'failed to persist operation state')); + } + if (action === 'nak_retry') work.msg.nak(); + else work.msg.term(errorMessage); + input.logger.error({ + worker: context?.workerLabel, + kind: context?.decoded.kind, + opId: context?.decoded.opId, + jobId: context?.decoded.jobId, + status: action === 'nak_retry' ? 'running' : 'failed', + error: errorMessage, + deliveryCount, + retryAction: action === 'nak_retry' ? 'nack_retry' : 'term', + }, 'job.terminal'); + } finally { + if (heartbeat) clearInterval(heartbeat); + } + }; + + const runLoop = async (work: { + owner: object; + consumer: Consumer; + codec: JsonCodec>; + run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout']; + workerLabel: string; + }): Promise => { + const detached = () => input.isStopping() || stopRequested || !input.isOwnerActive(work.owner); + while (!detached()) { + let msg: JsMsg | null = null; + try { + try { + msg = await work.consumer.next({ expires: PULL_EXPIRES_MS }); + } catch (error) { + if (detached()) return; + input.logger.error({ error: toErrorMessage(error), worker: work.workerLabel }, 'worker pull failed'); + await sleep(LOOP_ERROR_BACKOFF_MS); + continue; + } + if (!msg) continue; + input.markActivity(`job_received:${work.workerLabel}`); + input.onInFlightJobsChanged(1); + await gate.acquire(); + if (detached()) return; + await processMessage({ ...work, msg }); + } finally { + if (msg) { + gate.release(); + input.onInFlightJobsChanged(-1); + input.markActivity(`job_completed:${work.workerLabel}`); + } + } + } + }; + + return { + start(owner: object, consumers: { whisper: Consumer; pdfLayout: Consumer }): void { + stopRequested = false; + loops = []; + for (let i = 0; i < input.jobConcurrency; i += 1) { + loops.push(runLoop({ owner, consumer: consumers.whisper, codec: input.whisperCodec, run: input.handlers.runWhisper, workerLabel: `whisper-${i + 1}` })); + loops.push(runLoop({ owner, consumer: consumers.pdfLayout, codec: input.pdfCodec, run: input.handlers.runPdfLayout, workerLabel: `layout-${i + 1}` })); + } + }, + async stop(): Promise { + stopRequested = true; + await Promise.allSettled(loops); + loops = []; + }, + }; +} diff --git a/compute-worker/src/compute/control-plane/index.ts b/compute-worker/src/operations/index.ts similarity index 71% rename from compute-worker/src/compute/control-plane/index.ts rename to compute-worker/src/operations/index.ts index 5b8255d..35699ba 100644 --- a/compute-worker/src/compute/control-plane/index.ts +++ b/compute-worker/src/operations/index.ts @@ -1,4 +1,4 @@ export * from './types'; export * from './state-machine'; -export * from './orchestrator'; +export * from './service'; export * from './sse'; diff --git a/compute-worker/src/orphan-recovery.ts b/compute-worker/src/operations/recovery.ts similarity index 98% rename from compute-worker/src/orphan-recovery.ts rename to compute-worker/src/operations/recovery.ts index 8257b4d..5aaaa10 100644 --- a/compute-worker/src/orphan-recovery.ts +++ b/compute-worker/src/operations/recovery.ts @@ -4,7 +4,7 @@ import type { WorkerJobTiming, WorkerJobState, WorkerOperationState, -} from './compute/api-contracts'; +} from '../api/contracts'; export type StreamedOperationState = WorkerOperationState; diff --git a/compute-worker/src/compute/control-plane/orchestrator.ts b/compute-worker/src/operations/service.ts similarity index 99% rename from compute-worker/src/compute/control-plane/orchestrator.ts rename to compute-worker/src/operations/service.ts index 7f621a4..1469277 100644 --- a/compute-worker/src/compute/control-plane/orchestrator.ts +++ b/compute-worker/src/operations/service.ts @@ -5,7 +5,7 @@ import type { WorkerOperationKind, WorkerOperationRequest, WorkerOperationState, -} from '../api-contracts'; +} from '../api/contracts'; import { buildQueuedState, createErrorShape, diff --git a/compute-worker/src/compute/control-plane/sse.ts b/compute-worker/src/operations/sse.ts similarity index 100% rename from compute-worker/src/compute/control-plane/sse.ts rename to compute-worker/src/operations/sse.ts diff --git a/compute-worker/src/compute/control-plane/state-machine.ts b/compute-worker/src/operations/state-machine.ts similarity index 98% rename from compute-worker/src/compute/control-plane/state-machine.ts rename to compute-worker/src/operations/state-machine.ts index d280215..3144e8c 100644 --- a/compute-worker/src/compute/control-plane/state-machine.ts +++ b/compute-worker/src/operations/state-machine.ts @@ -4,7 +4,7 @@ import type { WorkerOperationKind, WorkerOperationRequest, WorkerOperationState, -} from '../api-contracts'; +} from '../api/contracts'; export function isTerminalStatus(status: WorkerJobState): boolean { return status === 'succeeded' || status === 'failed'; diff --git a/compute-worker/src/compute/control-plane/types.ts b/compute-worker/src/operations/types.ts similarity index 98% rename from compute-worker/src/compute/control-plane/types.ts rename to compute-worker/src/operations/types.ts index 9fc66ff..1d60200 100644 --- a/compute-worker/src/compute/control-plane/types.ts +++ b/compute-worker/src/operations/types.ts @@ -2,7 +2,7 @@ import type { WorkerOperationEvent, WorkerOperationKind, WorkerOperationState, -} from '../api-contracts'; +} from '../api/contracts'; export type OperationState = WorkerOperationState; export type OperationEvent = WorkerOperationEvent; diff --git a/compute-worker/src/server.ts b/compute-worker/src/server.ts index 1bf176c..3a110cc 100644 --- a/compute-worker/src/server.ts +++ b/compute-worker/src/server.ts @@ -1,4 +1,4 @@ -import { startComputeWorkerFromEnv } from './runtime'; +import { startComputeWorkerFromEnv } from './api/app'; void startComputeWorkerFromEnv().catch((error) => { console.error('[compute-worker] fatal startup error', error); diff --git a/compute-worker/src/storage/artifact-addressing.ts b/compute-worker/src/storage/artifact-addressing.ts index a4d1ee6..e341f23 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 '../compute/pdf/parser-version'; -import { encodeParserVersion } from '../compute/pdf/parser-version-key'; +import { PDF_PARSER_VERSION } from '../inference/pdf/parser-version'; +import { encodeParserVersion } from '../inference/pdf/parser-version-key'; const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; diff --git a/compute-worker/tests/api/routes.test.ts b/compute-worker/tests/api/routes.test.ts index e7336ee..4dea817 100644 --- a/compute-worker/tests/api/routes.test.ts +++ b/compute-worker/tests/api/routes.test.ts @@ -1,5 +1,5 @@ import { afterEach, beforeEach, describe, expect, test } from 'vitest'; -import { createComputeWorkerApp } from '../../src/runtime'; +import { createComputeWorkerApp } from '../../src/api/app'; import { FakeControlPlane } from '../fixtures/fake-control-plane'; const AUTH = { authorization: 'Bearer test-token' }; diff --git a/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts b/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts index 71436d0..d035613 100644 --- a/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts +++ b/compute-worker/tests/compute/algorithms/pdf-merge-text-with-regions.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import { mergeTextWithRegions } from '../../../src/compute'; +import { mergeTextWithRegions } from '../../../src/inference'; describe('mergeTextWithRegions', () => { test('assigns text items to containing regions by centroid and joins in reading order', () => { diff --git a/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts b/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts index f11272d..167dc46 100644 --- a/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts +++ b/compute-worker/tests/compute/algorithms/pdf-parse-normalize-text-items.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'vitest'; -import { normalizeTextItemsForLayout } from '../../../src/compute'; +import { normalizeTextItemsForLayout } from '../../../src/inference'; import type { TextItem } from 'pdfjs-dist/types/src/display/api'; function makeTextItem( diff --git a/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts b/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts index fb7b602..97d4b19 100644 --- a/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts +++ b/compute-worker/tests/compute/algorithms/pdf-stitch-cross-page-blocks.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'vitest'; -import { stitchCrossPageBlocks } from '../../../src/compute'; -import type { ParsedPdfBlock, ParsedPdfDocument, ParsedPdfBlockKind } from '../../../src/compute/types'; +import { stitchCrossPageBlocks } from '../../../src/inference'; +import type { ParsedPdfBlock, ParsedPdfDocument, ParsedPdfBlockKind } from '../../../src/inference/types'; import { makeParsedPdfDocument, makeParsedPdfPage } from './support/document-fixtures'; function makeBlock( diff --git a/compute-worker/tests/compute/algorithms/support/document-fixtures.ts b/compute-worker/tests/compute/algorithms/support/document-fixtures.ts index ffcde2d..f8db418 100644 --- a/compute-worker/tests/compute/algorithms/support/document-fixtures.ts +++ b/compute-worker/tests/compute/algorithms/support/document-fixtures.ts @@ -1,5 +1,5 @@ -import type { BaseDocument } from '../../../src/types/documents'; -import type { ParsedPdfBlock, ParsedPdfBlockKind, ParsedPdfDocument, ParsedPdfPage } from '../../../../src/compute/types'; +import type { BaseDocument } from '../../../../../src/types/documents'; +import type { ParsedPdfBlock, ParsedPdfBlockKind, ParsedPdfDocument, ParsedPdfPage } from '../../../../src/inference/types'; export function makeBaseDocument(overrides: Partial = {}): BaseDocument { return { diff --git a/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts b/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts index 2eea68f..fedea8b 100644 --- a/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts +++ b/compute-worker/tests/compute/algorithms/whisper-alignment-mapping.test.ts @@ -1,7 +1,7 @@ import { describe, expect, test } from 'vitest'; import { mapWordsToSentenceOffsets, -} from '../../../src/compute'; +} from '../../../src/inference'; describe('whisper alignment mapping', () => { test('maps words to sentence offsets with punctuation and repeated spaces', () => { diff --git a/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts b/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts index ac4a3d8..6b51d62 100644 --- a/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts +++ b/compute-worker/tests/compute/algorithms/whisper-spectral.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import { buildGoertzelCoefficients, goertzelPower } from '../../../src/compute'; +import { buildGoertzelCoefficients, goertzelPower } from '../../../src/inference'; function dftPower(samples: Float32Array, k: number): number { const n = samples.length; diff --git a/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts b/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts index a6ff75d..6f8bced 100644 --- a/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts +++ b/compute-worker/tests/compute/algorithms/whisper-token-timestamps.test.ts @@ -3,7 +3,7 @@ import * as ort from 'onnxruntime-node'; import { buildWordsFromTimestampedTokens, extractTokenStartTimestamps, -} from '../../../src/compute'; +} from '../../../src/inference'; describe('whisper token timestamp alignment', () => { test('extracts monotonic token timestamps from cross-attention maps', () => { diff --git a/compute-worker/tests/compute/control-plane/orchestrator.test.ts b/compute-worker/tests/compute/control-plane/orchestrator.test.ts index 6c9bbab..66964d3 100644 --- a/compute-worker/tests/compute/control-plane/orchestrator.test.ts +++ b/compute-worker/tests/compute/control-plane/orchestrator.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'vitest'; -import type { WorkerOperationRequest } from '../../../src/compute/api-contracts'; -import { OperationOrchestrator } from '../../../src/compute/control-plane'; +import type { WorkerOperationRequest } from '../../../src/api/contracts'; +import { OperationOrchestrator } from '../../../src/operations'; import { InMemoryOperationEventStream, InMemoryOperationQueue, diff --git a/compute-worker/tests/compute/control-plane/run-layout-model.test.ts b/compute-worker/tests/compute/control-plane/run-layout-model.test.ts index 27ef706..c8278c9 100644 --- a/compute-worker/tests/compute/control-plane/run-layout-model.test.ts +++ b/compute-worker/tests/compute/control-plane/run-layout-model.test.ts @@ -70,13 +70,13 @@ vi.mock('@napi-rs/canvas', () => { }; }); -vi.mock('../../../src/compute/pdf/model', () => ({ +vi.mock('../../../src/inference/pdf/model', () => ({ ensureModel: vi.fn(async () => '/tmp/model.onnx'), MODEL_CONFIG_PATH: '/tmp/model-config.json', MODEL_PREPROCESSOR_PATH: '/tmp/model-preprocessor.json', })); -vi.mock('../../../src/compute/config/cpu-budget', () => ({ +vi.mock('../../../src/inference/config/cpu-budget', () => ({ getOnnxThreadsPerJob: vi.fn(() => 1), })); @@ -107,7 +107,7 @@ describe('runLayoutModel', () => { }, }; - const { runLayoutModel } = await import('../../../src/compute/pdf/runLayoutModel'); + const { runLayoutModel } = await import('../../../src/inference/pdf/runLayoutModel'); const regions = await runLayoutModel({ pageWidth: 100, pageHeight: 100, @@ -154,7 +154,7 @@ describe('runLayoutModel', () => { }, }; - const { runLayoutModel } = await import('../../../src/compute/pdf/runLayoutModel'); + const { runLayoutModel } = await import('../../../src/inference/pdf/runLayoutModel'); const regions = await runLayoutModel({ pageWidth: 100, pageHeight: 100, diff --git a/compute-worker/tests/compute/control-plane/sse.test.ts b/compute-worker/tests/compute/control-plane/sse.test.ts index 3461229..2697d6d 100644 --- a/compute-worker/tests/compute/control-plane/sse.test.ts +++ b/compute-worker/tests/compute/control-plane/sse.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import { encodeSseFrame, parseSseEventId, parseSsePayload } from '../../../src/compute/control-plane/sse'; +import { encodeSseFrame, parseSseEventId, parseSsePayload } from '../../../src/operations/sse'; describe('sse codec', () => { test('encodes event id and payload and decodes both reliably', () => { 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 a9d5390..8fba44a 100644 --- a/compute-worker/tests/compute/control-plane/state-machine.test.ts +++ b/compute-worker/tests/compute/control-plane/state-machine.test.ts @@ -1,11 +1,11 @@ import { describe, expect, test } from 'vitest'; -import type { WorkerOperationState } from '../../../src/compute/api-contracts'; +import type { WorkerOperationState } from '../../../src/api/contracts'; import { explainReplacementReason, isInflightStatus, isTerminalStatus, shouldReuseExistingOperation, -} from '../../../src/compute/control-plane/state-machine'; +} from '../../../src/operations/state-machine'; function runningState(overrides: Partial = {}): WorkerOperationState { return { 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 7395b2b..811521b 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/compute/api-contracts'; +import type { WorkerOperationKind } from '../../../src/api/contracts'; import type { OperationEvent, OperationEventStream, @@ -8,7 +8,7 @@ import type { OperationState, OperationStateStore, QueuedOperation, -} from '../../../src/compute/control-plane/types'; +} from '../../../src/operations/types'; function topicFor(opId: string): string { return `op.${opId}`; diff --git a/compute-worker/tests/fixtures/fake-control-plane.ts b/compute-worker/tests/fixtures/fake-control-plane.ts index 40ab638..23f2760 100644 --- a/compute-worker/tests/fixtures/fake-control-plane.ts +++ b/compute-worker/tests/fixtures/fake-control-plane.ts @@ -4,8 +4,8 @@ import type { WorkerOperationEvent, WorkerOperationRequest, WorkerOperationState, -} from '../../src/compute/api-contracts'; -import type { ComputeWorkerRouteDeps } from '../../src/runtime'; +} from '../../src/api/contracts'; +import type { ComputeWorkerRouteDeps } from '../../src/api/app'; type ComputeResult = WhisperAlignJobResult | PdfLayoutJobResult; type ComputeState = WorkerOperationState; diff --git a/compute-worker/tests/unit/jetstream-adapters.test.ts b/compute-worker/tests/unit/jetstream-adapters.test.ts index ae86c06..a238996 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/compute/control-plane'; -import type { WorkerOperationRequest } from '../../src/compute/api-contracts'; +import { OperationOrchestrator } from '../../src/operations'; +import type { WorkerOperationRequest } from '../../src/api/contracts'; import { JetStreamOperationEventStream, JetStreamOperationQueue, @@ -11,7 +11,7 @@ import { opStateKvKey, type KvEntryLike, type KvStoreLike, -} from '../../src/control-plane/jetstream'; +} from '../../src/infrastructure/nats-adapters'; class FakeKvStore implements KvStoreLike { private readonly data = new Map(); diff --git a/compute-worker/tests/unit/key-and-progress.test.ts b/compute-worker/tests/unit/key-and-progress.test.ts index cbaf5ac..d6274db 100644 --- a/compute-worker/tests/unit/key-and-progress.test.ts +++ b/compute-worker/tests/unit/key-and-progress.test.ts @@ -1,9 +1,9 @@ import { describe, expect, test } from 'vitest'; -import { hashOpKey, opIndexKvKey, opStateKvKey } from '../../src/control-plane/jetstream'; +import { hashOpKey, opIndexKvKey, opStateKvKey } from '../../src/infrastructure/nats-adapters'; import { buildInferProgressForPageParsed, buildInferProgressForPageStart, -} from '../../src/pdf-progress'; +} from '../../src/jobs/pdf-progress'; import { buildPdfOperationKey } from '../../src/api/operation-keys'; import { parsedPdfArtifactKey } from '../../src/storage/artifact-addressing'; diff --git a/compute-worker/tests/unit/orphan-recovery.test.ts b/compute-worker/tests/unit/orphan-recovery.test.ts index 712fdf6..7ec771d 100644 --- a/compute-worker/tests/unit/orphan-recovery.test.ts +++ b/compute-worker/tests/unit/orphan-recovery.test.ts @@ -1,5 +1,5 @@ import { afterEach, describe, expect, test, vi } from 'vitest'; -import { recoverOrphanedOperations } from '../../src/orphan-recovery'; +import { recoverOrphanedOperations } from '../../src/operations/recovery'; import { FakeControlPlane } from '../fixtures/fake-control-plane'; describe('orphan recovery', () => { diff --git a/compute-worker/tests/unit/pdf-artifact-persistence.test.ts b/compute-worker/tests/unit/pdf-artifact-persistence.test.ts index 9378435..68f3de5 100644 --- a/compute-worker/tests/unit/pdf-artifact-persistence.test.ts +++ b/compute-worker/tests/unit/pdf-artifact-persistence.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test, vi } from 'vitest'; -import { persistParsedPdfWhileSourceExists } from '../../src/pdf-artifact-persistence'; +import { persistParsedPdfWhileSourceExists } from '../../src/jobs/pdf-artifact-persistence'; describe('PDF artifact persistence', () => { test('does not write parsed output after the source was deleted', async () => { diff --git a/compute-worker/tests/unit/worker-loop-policy.test.ts b/compute-worker/tests/unit/worker-loop-policy.test.ts index 818e459..c200e70 100644 --- a/compute-worker/tests/unit/worker-loop-policy.test.ts +++ b/compute-worker/tests/unit/worker-loop-policy.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import { buildQueueWaitTiming, decideRetryAction } from '../../src/worker-loop-policy'; +import { buildQueueWaitTiming, decideRetryAction } from '../../src/jobs/worker-loop-policy'; describe('worker loop policy', () => { test('returns queue wait timing with non-negative clamped duration', () => { diff --git a/scripts/check-next-server-bundle.mjs b/scripts/check-next-server-bundle.mjs index 2b3cbcb..8174d40 100644 --- a/scripts/check-next-server-bundle.mjs +++ b/scripts/check-next-server-bundle.mjs @@ -13,7 +13,7 @@ const forbidden = [ 'onnxruntime-node', '@huggingface/tokenizers', '@openreader/compute-worker', - '/compute-worker/src/compute/', + '/compute-worker/src/inference/', ]; const includeExt = new Set(['.js', '.mjs', '.cjs']); diff --git a/tests/unit/compute-worker-client-contract.vitest.spec.ts b/tests/unit/compute-worker-client-contract.vitest.spec.ts index e37f234..39adac6 100644 --- a/tests/unit/compute-worker-client-contract.vitest.spec.ts +++ b/tests/unit/compute-worker-client-contract.vitest.spec.ts @@ -1,5 +1,5 @@ import { afterEach, describe, expect, test } from 'vitest'; -import { createComputeWorkerApp, type ComputeWorkerApp } from '../../compute-worker/src/runtime'; +import { createComputeWorkerApp, type ComputeWorkerApp } from '../../compute-worker/src/api/app'; import { FakeControlPlane } from '../../compute-worker/tests/fixtures/fake-control-plane'; import { ComputeWorkerClient } from '../../src/lib/server/compute-worker/client';