perf(batch): run the speaker-embedding trunk once per chunk (+30% per GPU); opt-in multi-process worker

Merged
GI
3d ago
+797-68

What

Two batch GPU worker changes ship in one worker roll.

  1. The speaker-embedding trunk runs once per chunk (on by default, +30% throughput per GPU). Stock pyannote
    get_embeddings runs the WeSpeaker ResNet once per (10 s chunk, local speaker) pair, which means three forwards per chunk for
    community-1. The worker now installs a per-unit shared-trunk get_embeddings: forward_frames once per chunk, then
    forward_embedding per speaker mask. It reuses the row and trunk code that already existed for flight batching, but it
    never packs rows across units, so a unit's embeddings don't depend on its flight. That cross-unit dependence is why flight
    batching stays dark.
    • FP32 and the pinned embedding batch size are unchanged.
    • Version guard: pyannote 4.0.x with the expected layout. Anything else keeps the stock path.
    • A CUDA OOM or a non-fatal error falls back to the stock forward for that call. Other fatal CUDA errors still quarantine.
    • Kill switch: worker.sharedEmbeddingTrunkEnabled: false (env BATCH_PIPELINE_SHARED_EMBEDDING_TRUNK_ENABLED).
  2. Multi-process supervisor, opt-in (default 1, so no behaviour change). worker.processes: 2
    (env BATCH_PIPELINE_WORKER_PROCESSES) runs two independent pump processes on the pod's GPU under a torch-free supervisor.
    • Each child is the unchanged runtime with its own models, scheduler id, leader holder and decoded bound.
    • Pod /ready, /health and /metrics aggregate the children. Only child 0 samples the queue.
    • SIGTERM fans out to every child. Any child exit drains its siblings and restarts the pod.
    • Each child gets its own shadow slot owner, because a shared one would let a sibling drop the other's shadow job (found by GLM review).
    • Two-process pods request 24Gi / 6 CPU, with limits of 26Gi / 8 CPU.

Why (measured)

The bench runs the real worker chain (load_worker, all models) on a dev T4 over 248 public-audio units (AMI, FLEURS in 7
languages, synthetic, generated silence), in flights of 8 with 2 lanes and simulated storage latency. Env comes from this
chart's prod values. The gate compares results against a frozen single-process reference.

treeprocessesaudio-h per GPU-hembedding s / diarize sgate
main137.47 / 36.83 (two runs)69.4 / 78.6reference
this PR (shared trunk)148.85 / 47.83 (+30%)24.1 / 33.1identical words and speakers
main245.4 (bench_mp, vs 33.8 for one)—identical
this PR245.1 (bench_mp, 3-CPU pod)—identical

On the 3-CPU T4 pod, two processes stopped helping once the shared trunk cut the GPU work, because the CPU became the limit. That's
why the default is 1. Prod L4 nodes have 8 vCPU; a dev L4 run was blocked by stockout (GCE out of resources). A later
canary with worker.processes=2 can decide. Prod fit for 2 (read-only check): node allocatable 27.6 GiB / 7.91 CPU, about 1 GiB / 0.6 CPU
used by other pods; the 7-day GPU framebuffer peak was 10.0 GiB per process.

Rollout and readback (lead)

Worker roll (gcp_diarizer, Recreate; re-apply repair flags after). In the first hour:

  • event=batch_shared_trunk outcome=enabled once per pod; event=batch_shared_trunk_fallback close to 0;
  • embedding and diarize stage seconds per diarized unit down about 60% (omi_batch_pipeline_stage_seconds);
    worker_flight seconds per completed unit down;
  • GPU framebuffer below the previous 10 GiB peak plus a small margin (the embedding forward now batches up to 8 chunks);
  • hard-3 clean.

Rollback: worker.sharedEmbeddingTrunkEnabled=false, or Helm rollback.

Tests

  • All 69 tests/unit/test_batch_pipeline*.py files plus the flag, env-ratchet, chart and runtime suites pass (72 files on the combined
    branch, then re-run after each fix).
  • Real-model CPU test (torch 2.8, pyannote.audio 4.0.3): shared-trunk embeddings match stock within 1e-4 for both
    exclude_overlap modes, and PyannoteDiarizer output is identical on 5 clips.
  • Supervisor tests with real child subprocesses.

Reviews: Luna, during the lane, found one supervisor issue, fixed. GLM on the supervisor found HIGH: shared shadow id, fixed with a test.
GLM on the shared trunk said SHIP; minor findings fixed (VRAM docstring, a chart kill-switch test).

Not proven: no L4 parity or throughput run (L4 stocked out in dev). Equivalence is proven on CPU, plus identical bench output
on T4. The prod readback above closes this.

🤖 Generated with Claude Code