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

@@ -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

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<void> {
   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<void> {
   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<typeof connect>[0] = { servers: natsUrl };
@@ -533,8 +542,9 @@ async function main(): Promise<void> {
       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<void> {
     availableCpuCores: getAvailableCpuCores(),
     onnxThreadsPerJob: getOnnxThreadsPerJob(),
     natsApiTimeoutMs: NATS_API_TIMEOUT_MS,
+    natsReplicas,
     pdfLayoutHardCapMs: pdfHardCapMs,
   }, 'compute runtime config');
 

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=<long-random-shared-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)