From 80596da7733cc7d1a2adbb751e6086e0ded3cf20 Mon Sep 17 00:00:00 2001 From: Richard R Date: Wed, 3 Jun 2026 01:44:18 -0600 Subject: [PATCH] feat(worker): recover and fail stale in-flight pdf ops on startup Add orphaned operation recovery logic to detect and mark stale in-flight pdf_layout jobs as failed during worker startup. Extend OperationStateStore with listOpStates for state enumeration. Update tests and documentation to cover recovery behavior and new environment variable COMPUTE_PDF_JOB_ATTEMPTS. --- .env.example | 1 + compute/worker/src/control-plane/jetstream.ts | 13 +++++ compute/worker/src/runtime.ts | 55 +++++++++++++++++++ compute/worker/tests/api/routes.test.ts | 33 +++++++++++ .../tests/fixtures/fake-control-plane.ts | 5 ++ docs-site/docs/deploy/compute-worker.md | 3 +- docs-site/docs/deploy/local-development.md | 4 +- .../docs/reference/environment-variables.md | 8 +++ 8 files changed, 120 insertions(+), 2 deletions(-) diff --git a/.env.example b/.env.example index 8173d6d..02c3e57 100644 --- a/.env.example +++ b/.env.example @@ -65,6 +65,7 @@ S3_BUCKET= # COMPUTE_JOB_CONCURRENCY=1 # COMPUTE_WHISPER_TIMEOUT_MS=30000 # COMPUTE_PDF_TIMEOUT_MS=300000 +# COMPUTE_PDF_JOB_ATTEMPTS=1 # COMPUTE_OP_STALE_MS=1800000 # WHISPER_MODEL_BASE_URL=https://huggingface.co/onnx-community/whisper-base_timestamped/resolve/main # PDF_LAYOUT_MODEL_BASE_URL=https://huggingface.co/Bei0001/PP-DocLayoutV3-ONNX/resolve/main diff --git a/compute/worker/src/control-plane/jetstream.ts b/compute/worker/src/control-plane/jetstream.ts index 642906d..8844673 100644 --- a/compute/worker/src/control-plane/jetstream.ts +++ b/compute/worker/src/control-plane/jetstream.ts @@ -27,6 +27,7 @@ export interface KvStoreLike { put(key: string, data: Uint8Array): Promise; create(key: string, data: Uint8Array): Promise; update(key: string, data: Uint8Array, version: number): Promise; + keys(filter?: string | string[]): Promise>; } function toErrorMessage(error: unknown): string { @@ -91,6 +92,18 @@ export class JetStreamOperationStateStore implements Operation await kv.put(opStateKvKey(state.opId), this.opStateCodec.encode(state)); } + async listOpStates(): Promise[]> { + const kv = await this.getKv(); + const keys = await kv.keys('op_state.*'); + const states: OperationState[] = []; + for await (const key of keys) { + const entry = await kv.get(key); + if (!isPut(entry)) continue; + states.push(this.opStateCodec.decode(entry.value)); + } + return states; + } + async getOpIndex(opKey: string): Promise<{ opId: string } | null> { const kv = await this.getKv(); const entry = await kv.get(opIndexKvKey(opKey)); diff --git a/compute/worker/src/runtime.ts b/compute/worker/src/runtime.ts index 61b18ca..efd9208 100644 --- a/compute/worker/src/runtime.ts +++ b/compute/worker/src/runtime.ts @@ -120,6 +120,7 @@ interface OperationEventStreamLike { interface OperationStateStoreLike { getOpState(opId: string): Promise; + listOpStates?(): Promise; } interface OrchestratorLike { @@ -178,6 +179,10 @@ function requireEnv(name: string): string { return value; } +function isInflightStatus(status: WorkerJobState): boolean { + return status === 'queued' || status === 'running'; +} + function readIntEnv(name: string, fallback: number): number { const raw = process.env[name]?.trim(); if (!raw) return fallback; @@ -617,6 +622,8 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti const layoutConsumer = await js.consumers.get(JOBS_STREAM_NAME, LAYOUT_CONSUMER_NAME); const next: NatsSession = { nc, js, jsm, kv, whisperConsumer, layoutConsumer }; session = next; + sessionGeneration += 1; + orphanRecoveryDoneForGeneration = -1; markActivity(); startWorkerLoops(next); startIdleTimer(); @@ -731,8 +738,55 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti const operationStateStore = options.routeDeps?.operationStateStore ?? defaultOperationStateStore; const operationEventStream = options.routeDeps?.operationEventStream ?? defaultOperationEventStream; const orchestrator = options.routeDeps?.orchestrator ?? defaultOrchestrator; + let orphanRecoveryPromise: Promise | null = null; + let orphanRecoveryDoneForGeneration = -1; + let sessionGeneration = options.routeDeps ? 0 : -1; + + const ensureOrphanedOpRecovery = async (): Promise => { + if (typeof operationStateStore.listOpStates !== 'function') return; + if (orphanRecoveryDoneForGeneration === sessionGeneration) return; + if (orphanRecoveryPromise) { + await orphanRecoveryPromise; + return; + } + + orphanRecoveryPromise = (async () => { + const now = Date.now(); + const states = await operationStateStore.listOpStates!(); + const stalePdfStates = states.filter((state) => ( + state.kind === 'pdf_layout' + && isInflightStatus(state.status) + && (now - state.updatedAt) > opStaleMs + )); + + for (const state of stalePdfStates) { + await orchestrator.markFailed({ + opId: state.opId, + error: { + code: 'WORKER_ORPHANED_OP', + message: `Worker stopped before completion; stale operation recovered on startup after ${opStaleMs}ms`, + }, + updatedAt: now, + }); + } + + if (stalePdfStates.length > 0) { + app.log.warn({ + recoveredCount: stalePdfStates.length, + opIds: stalePdfStates.map((state) => state.opId), + }, 'recovered stale in-flight pdf operations on startup'); + } + + orphanRecoveryDoneForGeneration = sessionGeneration; + })().finally(() => { + orphanRecoveryPromise = null; + }); + + await orphanRecoveryPromise; + }; const getOpState = async (opId: string): Promise => { + await ensureOrphanedOpRecovery(); return await operationStateStore.getOpState(opId); }; @@ -798,6 +852,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti } const requestOp = parsed.data as WorkerOperationRequest; + await ensureOrphanedOpRecovery(); const op = await orchestrator.enqueueOrReuse(requestOp); app.log.info({ kind: requestOp.kind, diff --git a/compute/worker/tests/api/routes.test.ts b/compute/worker/tests/api/routes.test.ts index 77c4d05..5be0a59 100644 --- a/compute/worker/tests/api/routes.test.ts +++ b/compute/worker/tests/api/routes.test.ts @@ -119,4 +119,37 @@ describe('compute worker API routes', () => { expect(stream.body).toContain('id: 7'); expect(stream.body).toContain('"status":"succeeded"'); }); + + test('marks stale in-flight pdf ops failed during startup reconciliation', async () => { + fake.seedState({ + opId: 'op-stale', + opKey: 'k-stale', + kind: 'pdf_layout', + jobId: 'job-op-stale', + status: 'running', + queuedAt: 1, + updatedAt: 1, + }); + + const fetch = await runtime.app.inject({ + method: 'GET', + url: '/ops/op-stale', + headers: AUTH, + }); + + expect(fetch.statusCode).toBe(200); + expect(fetch.json()).toMatchObject({ + opId: 'op-stale', + status: 'failed', + error: { + code: 'WORKER_ORPHANED_OP', + }, + }); + expect(fake.getState('op-stale')).toMatchObject({ + status: 'failed', + error: { + code: 'WORKER_ORPHANED_OP', + }, + }); + }); }); diff --git a/compute/worker/tests/fixtures/fake-control-plane.ts b/compute/worker/tests/fixtures/fake-control-plane.ts index 2d18f2b..305304e 100644 --- a/compute/worker/tests/fixtures/fake-control-plane.ts +++ b/compute/worker/tests/fixtures/fake-control-plane.ts @@ -47,6 +47,7 @@ export class FakeControlPlane { }, operationStateStore: { getOpState: async (opId) => this.stateByOpId.get(opId) ?? null, + listOpStates: async () => Array.from(this.stateByOpId.values()), }, operationEventStream: { subscribe: async ({ opId, sinceEventId, onEvent }) => { @@ -71,6 +72,10 @@ export class FakeControlPlane { this.eventsByOpId.set(opId, list); } + getState(opId: string): ComputeState | null { + return this.stateByOpId.get(opId) ?? null; + } + private async enqueueOrReuse(request: WorkerOperationRequest): Promise { const existingId = this.opIdByOpKey.get(request.opKey); if (existingId) { diff --git a/docs-site/docs/deploy/compute-worker.md b/docs-site/docs/deploy/compute-worker.md index b5ec152..cd279db 100644 --- a/docs-site/docs/deploy/compute-worker.md +++ b/docs-site/docs/deploy/compute-worker.md @@ -41,7 +41,7 @@ S3_SECRET_ACCESS_KEY=... - Embedded/local mode: configure the root `.env` only. - External worker mode: set `COMPUTE_WORKER_URL` and `COMPUTE_WORKER_TOKEN` on the app, and worker runtime values on the worker service. -- Keep shared values aligned across app and worker: `COMPUTE_WORKER_TOKEN`, `S3_*`, `COMPUTE_WHISPER_TIMEOUT_MS`, `COMPUTE_PDF_TIMEOUT_MS`, and `COMPUTE_OP_STALE_MS`. +- Keep shared values aligned across app and worker: `COMPUTE_WORKER_TOKEN`, `S3_*`, `COMPUTE_WHISPER_TIMEOUT_MS`, `COMPUTE_PDF_TIMEOUT_MS`, `COMPUTE_PDF_JOB_ATTEMPTS`, and `COMPUTE_OP_STALE_MS`. ::: Common optional variables: @@ -76,6 +76,7 @@ COMPUTE_WORKER_TOKEN= # Optional shared overrides: # COMPUTE_WHISPER_TIMEOUT_MS=30000 # COMPUTE_PDF_TIMEOUT_MS=300000 +# COMPUTE_PDF_JOB_ATTEMPTS=1 # COMPUTE_OP_STALE_MS=1800000 ``` diff --git a/docs-site/docs/deploy/local-development.md b/docs-site/docs/deploy/local-development.md index 338aeb0..7454b7b 100644 --- a/docs-site/docs/deploy/local-development.md +++ b/docs-site/docs/deploy/local-development.md @@ -189,9 +189,11 @@ COMPUTE_WORKER_TOKEN= ``` Ownership in external worker mode: -- root `.env`: app routing/auth (`COMPUTE_WORKER_URL`, `COMPUTE_WORKER_TOKEN`) plus optional shared timeout/stale overrides +- root `.env`: app routing/auth (`COMPUTE_WORKER_URL`, `COMPUTE_WORKER_TOKEN`) plus optional shared timeout/stale/retry overrides such as `COMPUTE_PDF_JOB_ATTEMPTS` - `compute/worker/.env*` (or worker platform env): worker runtime variables (`NATS_*`, `S3_*`, model base URLs, worker tuning) +For embedded worker startup (`COMPUTE_WORKER_URL` unset), worker tuning values such as `COMPUTE_PDF_JOB_ATTEMPTS` must be set in the root `.env` because `compute/worker/.env*` is ignored in that mode. + Worker mode requires worker-reachable shared object storage (S3-compatible endpoint). For external worker mode, object storage must be shared/reachable by both app and worker services. diff --git a/docs-site/docs/reference/environment-variables.md b/docs-site/docs/reference/environment-variables.md index 45786fc..d80bd8e 100644 --- a/docs-site/docs/reference/environment-variables.md +++ b/docs-site/docs/reference/environment-variables.md @@ -48,6 +48,7 @@ Runtime site features are seeded with `RUNTIME_SEED_JSON` / `RUNTIME_SEED_JSON_P | `COMPUTE_JOB_CONCURRENCY` | Compute | `1` | Shared compute concurrency cap | | `COMPUTE_WHISPER_TIMEOUT_MS` | Compute | `30000` | Whisper alignment timeout budget | | `COMPUTE_PDF_TIMEOUT_MS` | Compute | `300000` | PDF parse timeout budget | +| `COMPUTE_PDF_JOB_ATTEMPTS` | Compute | `1` | Max JetStream deliveries for PDF layout jobs | | `COMPUTE_OP_STALE_MS` | Compute | `max(30m, 4x max compute timeout)` | Shared stale window for compute op replacement | | `WHISPER_MODEL_BASE_URL` | Compute model source | onnx-community default | Override Whisper ONNX model base URL | | `PDF_LAYOUT_MODEL_BASE_URL` | Compute model source | PP-DocLayoutV3 default | Override PDF layout ONNX model base URL | @@ -307,6 +308,13 @@ PDF parse timeout budget. - Default: `300000` +### COMPUTE_PDF_JOB_ATTEMPTS + +Max JetStream deliveries for PDF layout jobs. + +- Default: `1` +- In embedded worker mode, set this in the root `.env` + ### COMPUTE_OP_STALE_MS Stale operation window before worker/app cleanup logic can replace an op.