From 653dbedaee19aa43b4d498bfe238710008bc3a36 Mon Sep 17 00:00:00 2001 From: Richard R Date: Fri, 22 May 2026 09:32:25 -0600 Subject: [PATCH] refactor(worker): add configurable NATS JetStream replica count Introduce COMPUTE_NATS_REPLICAS environment variable to control the number of replicas for JetStream streams and key-value buckets. Only values 1, 3, or 5 are accepted, defaulting to 1 for other inputs. Update worker server logic and deployment documentation to describe and support this configuration. --- compute/worker/.env.example | 1 + compute/worker/src/server.ts | 13 ++++++++++++- docs-site/docs/deploy/compute-worker.md | 3 +++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/compute/worker/.env.example b/compute/worker/.env.example index 941c118..800dd93 100644 --- a/compute/worker/.env.example +++ b/compute/worker/.env.example @@ -40,6 +40,7 @@ S3_FORCE_PATH_STYLE=true # COMPUTE_PDF_JOB_ATTEMPTS=1 # COMPUTE_JOBS_STREAM_MAX_BYTES=268435456 # COMPUTE_JOB_STATES_MAX_BYTES=67108864 +# COMPUTE_NATS_REPLICAS=1 # Optional stale window for reusing in-flight opKey entries before forcing a new attempt # Default is max(30m, 4x max compute timeout); running jobs also refresh heartbeat every 5s # COMPUTE_OP_STALE_MS=1800000 diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index e545b83..a1ce849 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -128,6 +128,11 @@ function readIntEnv(name: string, fallback: number): number { return Math.floor(parsed); } +function normalizeNatsReplicas(value: number): number { + if (value === 3 || value === 5) return value; + return 1; +} + function parseBoolEnv(name: string, fallback: boolean): boolean { const raw = process.env[name]?.trim(); if (!raw) return fallback; @@ -383,6 +388,7 @@ async function ensureJetStreamResources( pdfTimeoutMs: number, pdfAttempts: number, maxBytes: number, + natsReplicas: number, ): Promise { const streamConfig = { name: JOBS_STREAM_NAME, @@ -390,6 +396,7 @@ async function ensureJetStreamResources( retention: RetentionPolicy.Workqueue, storage: StorageType.File, max_bytes: maxBytes, + num_replicas: natsReplicas, }; try { @@ -399,6 +406,7 @@ async function ensureJetStreamResources( await jsm.streams.update(JOBS_STREAM_NAME, { subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT], max_bytes: maxBytes, + num_replicas: natsReplicas, }); } @@ -451,6 +459,7 @@ async function main(): Promise { const prewarmModels = parseBoolEnv('COMPUTE_PREWARM_MODELS', true); const jobsStreamMaxBytes = readIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); const jobStatesMaxBytes = readIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024); + const natsReplicas = normalizeNatsReplicas(readIntEnv('COMPUTE_NATS_REPLICAS', 1)); const opStaleMs = getComputeOpStaleMs(); const connectOpts: Parameters[0] = { servers: natsUrl }; @@ -533,8 +542,9 @@ async function main(): Promise { const nc: NatsConnection = await connect(connectOpts); const js: JetStreamClient = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS }); const jsm: JetStreamManager = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS }); - await ensureJetStreamResources(jsm, whisperTimeoutMs, pdfTimeoutMs, pdfAttempts, jobsStreamMaxBytes); + await ensureJetStreamResources(jsm, whisperTimeoutMs, pdfTimeoutMs, pdfAttempts, jobsStreamMaxBytes, natsReplicas); const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { + replicas: natsReplicas, history: 1, ttl: COMPUTE_STATE_TTL_MS, max_bytes: jobStatesMaxBytes, @@ -614,6 +624,7 @@ async function main(): Promise { availableCpuCores: getAvailableCpuCores(), onnxThreadsPerJob: getOnnxThreadsPerJob(), natsApiTimeoutMs: NATS_API_TIMEOUT_MS, + natsReplicas, pdfLayoutHardCapMs: pdfHardCapMs, }, 'compute runtime config'); diff --git a/docs-site/docs/deploy/compute-worker.md b/docs-site/docs/deploy/compute-worker.md index 13296d4..21171c5 100644 --- a/docs-site/docs/deploy/compute-worker.md +++ b/docs-site/docs/deploy/compute-worker.md @@ -56,6 +56,7 @@ Advanced tuning (usually leave unset unless you need overrides): - `COMPUTE_PDF_JOB_ATTEMPTS=1` (PDF layout retry attempts) - `COMPUTE_JOBS_STREAM_MAX_BYTES=268435456` (256MB JetStream jobs stream cap) - `COMPUTE_JOB_STATES_MAX_BYTES=67108864` (64MB JetStream KV bucket cap) +- `COMPUTE_NATS_REPLICAS=1` (JetStream stream + KV replicas; valid: `1`, `3`, `5`) - `COMPUTE_OP_STALE_MS=1800000` (stale op replacement window) ## App server environment variables (worker mode) @@ -158,6 +159,7 @@ COMPUTE_WORKER_TOKEN= # COMPUTE_PDF_JOB_ATTEMPTS=1 # COMPUTE_JOBS_STREAM_MAX_BYTES=268435456 # COMPUTE_JOB_STATES_MAX_BYTES=67108864 +# COMPUTE_NATS_REPLICAS=1 NATS_URL=tls://connect.ngs.global:4222 NATS_CREDS="-----BEGIN NATS USER JWT----- @@ -180,6 +182,7 @@ Notes: - On Railway, leave `PORT` managed by the platform. - If your platform supports mounted files, you can use `NATS_CREDS_FILE` instead of `NATS_CREDS`. - `COMPUTE_JOBS_STREAM_MAX_BYTES` and `COMPUTE_JOB_STATES_MAX_BYTES` are optional; defaults are `268435456` (256MiB) and `67108864` (64MiB). +- `COMPUTE_NATS_REPLICAS` is optional; default is `1`. Valid values are `1`, `3`, `5`. ### 4. Configure the OpenReader app server (worker mode)