Batch Processing Transcripts with Celery

Within the Async Transcription Queue Management layer of the Transcription & Speaker Diarization pipeline, this page solves one exact production scenario: running GPU-bound Whisper Large V3 inference and Pyannote diarization as Celery batch tasks that survive worker restarts, OOM bursts, and transient broker drops without losing or double-charging a single audio job.

When you process podcast back-catalogs, multi-track interviews, or archival audio, the naive approach — calling inference inline from a request handler — collapses the moment two large files land at once. Decoupling heavy inference from synchronous gateways through deterministic Celery task routing gives you explicit retry semantics, bounded concurrency, and horizontal worker scaling. The configuration below is the difference between a queue that drains predictably and one that wedges on a CUDA out-of-memory crash at 03:00.

Problem Framing

The failure is visible in the worker log within minutes of a burst. Two transcription tasks land on the same GPU worker, the second allocates its Whisper activations before the first has freed them, and the runtime aborts mid-generation:

[ERROR] celery.worker: Task tasks.transcription.process_transcription_batch[a1f3]
  raised unexpected: OutOfMemoryError('CUDA out of memory. Tried to allocate
  2.04 GiB. GPU 0 has 23.69 GiB total, 22.91 GiB reserved by PyTorch')
[WARNING] celery.worker: Worker exited prematurely: signal 9 (SIGKILL)

Two things compound the damage. First, the default prefetch pulls extra tasks onto a worker that is already saturated, so the broker believes work is in flight while the process is being killed. Second, with early acknowledgment (Celery’s default), the task was acked the instant it was received — so when the worker dies, the job is gone, not requeued. The queue silently loses transcripts, and the only symptom downstream is a gap in your archival metadata. The fix is a combination of acknowledgment policy, prefetch flattening, and per-child task ceilings.

Solution Architecture

The remedy is to make each GPU worker process exactly one task at a time, acknowledge only after success, and recycle the worker process periodically to flush CUDA context fragmentation. Audio is split on voice-activity (VAD) intervals before inference so no single chunk can exceed the memory envelope, and diarization runs once over the full file and is aligned back onto the per-chunk transcript.

Resilient Celery batch-transcription flow on a single GPU worker A job leaves the Redis Celery queue and is reserved one at a time by a single GPU worker running pool=solo with concurrency=1, prefetch_multiplier=1, and acks_late=True. Inside the worker the audio passes through a VAD chunker that emits windows of at most thirty seconds, into a per-chunk Whisper Large V3 inference loop, then through Pyannote 3.1 diarization, and finally a label-alignment stage that maps speakers onto each segment. On success the worker acknowledges the message and writes results to the Redis result backend. If the worker is killed mid-task by a CUDA out-of-memory SIGKILL, task_reject_on_worker_lost returns the still-unacknowledged job to the broker to be re-reserved. GPU worker · pool=solo concurrency = 1 · prefetch_multiplier = 1 · acks_late = True reserve 1 for each chunk ack worker lost (CUDA OOM · SIGKILL) → task_reject_on_worker_lost requeues the job Redis broker celery queue VAD chunker ≤ 30 s windows Whisper Large V3 fp16 inference per-chunk loop Pyannote 3.1 diarization over full file Align labels speaker → each segment Result backend Redis db 1

The key invariants: worker_prefetch_multiplier = 1 plus worker_concurrency = 1 means a worker reserves a job only when idle; task_acks_late = True with task_reject_on_worker_lost = True means a SIGKILL returns the job to the queue instead of dropping it; and worker_max_tasks_per_child forces a clean process restart that reclaims leaked VRAM.

Implementation

Step 1 — Broker and GPU worker configuration

Persist the broker, track results for state polling, and isolate inference to one task per worker. For broker-specific tuning beyond this, validate parameters against the official Celery configuration documentation.

# celery_config.py
# celery==5.3.6, redis==5.0.3
import os

# Broker and backend configuration
broker_url = os.getenv("CELERY_BROKER_URL", "redis://redis-broker:6379/0")
result_backend = os.getenv("CELERY_RESULT_BACKEND", "redis://redis-broker:6379/1")

# Prevent task loss during worker restarts: ack only after success,
# and requeue if the worker is killed mid-task (e.g. CUDA OOM SIGKILL).
task_acks_late = True
task_reject_on_worker_lost = True
task_ignore_result = False

# GPU worker isolation: one task at a time to prevent CUDA OOM,
# and a flat prefetch so a saturated worker never hoards jobs.
worker_concurrency = 1
worker_prefetch_multiplier = 1
worker_max_tasks_per_child = 50  # Recycle the process to flush leaked VRAM

# Retry and timeout defaults (seconds)
task_default_retry_delay = 15
task_default_max_retries = 3
task_soft_time_limit = 1800   # raise SoftTimeLimitExceeded for graceful cleanup
task_time_limit = 2100        # hard SIGKILL ceiling

# Serialization
task_serializer = "json"
result_serializer = "json"
accept_content = ["json"]

# Logging
worker_hijack_root_logger = False
worker_log_format = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
worker_task_log_format = "%(asctime)s [%(levelname)s] %(name)s.%(task_name)s: %(message)s"

worker_max_tasks_per_child is the single most important line for long-running Whisper loops: PyTorch’s caching allocator fragments over hundreds of variable-length generations, and a periodic child restart is the only reliable way to reclaim it without restarting the whole fleet.

Step 2 — The chunked transcription task

Models are loaded once at import so they stay resident across tasks. Audio is split on energy-based VAD and capped at 30 seconds per chunk — Whisper’s native receptive field — then speaker diarization with Pyannote labels are aligned onto each segment.

# tasks/transcription.py
# torch==2.2.1, transformers==4.39.3, pyannote.audio==3.1.1, librosa==0.10.1
import logging
import os
import traceback
import numpy as np
import torch
import librosa
from celery import Celery, shared_task
from pyannote.audio import Pipeline
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

app = Celery()
app.config_from_object("celery_config")

logger = logging.getLogger(__name__)

WHISPER_MODEL_ID = "openai/whisper-large-v3"
PYANNOTE_MODEL_ID = "pyannote/speaker-diarization-3.1"

# Load once per worker process so weights stay resident across tasks.
try:
    whisper_processor = AutoProcessor.from_pretrained(WHISPER_MODEL_ID)
    whisper_model = AutoModelForSpeechSeq2Seq.from_pretrained(
        WHISPER_MODEL_ID,
        torch_dtype=torch.float16,     # fp16 halves VRAM vs fp32
        low_cpu_mem_usage=True,
        use_safetensors=True,
    ).to("cuda")
    whisper_model.eval()

    diarization_pipeline = Pipeline.from_pretrained(
        PYANNOTE_MODEL_ID,
        use_auth_token=os.getenv("HF_TOKEN"),
    )
    logger.info("Loaded Whisper Large V3 and Pyannote 3.1.")
except Exception as e:
    # Fail loudly: a worker with no models must never accept tasks.
    logger.critical("Model initialization failed: %s", e)
    raise RuntimeError("GPU model preload failed. Worker cannot start.") from e


def chunk_audio_by_vad(
    audio_path: str,
    sample_rate: int = 16000,   # Whisper's required input rate
    max_chunk_sec: int = 30,    # cap each window to bound activation memory
    top_db: int = 30,           # silence threshold for VAD splitting
) -> list[tuple[np.ndarray, float, float]]:
    """Split audio on energy-based VAD intervals, capped to prevent OOM."""
    y, sr = librosa.load(audio_path, sr=sample_rate)
    intervals = librosa.effects.split(y, top_db=top_db)
    max_samples = int(max_chunk_sec * sr)
    chunks: list[tuple[np.ndarray, float, float]] = []
    for start_idx, end_idx in intervals:
        cursor = start_idx
        while cursor < end_idx:
            window_end = min(cursor + max_samples, end_idx)
            chunks.append((y[cursor:window_end], cursor / sr, window_end / sr))
            cursor = window_end
    return chunks


@shared_task(bind=True, max_retries=3, default_retry_delay=15)
def process_transcription_batch(self, audio_url: str, task_id: str) -> dict:
    """Batch transcription with chunking, diarization, and explicit diagnostics."""
    temp_path = None
    try:
        logger.info("Task %s: preparing audio: %s", self.request.id, audio_url)
        temp_path = f"/tmp/{task_id}.wav"
        # In production, replace with a streaming downloader writing to temp_path.
        # download_audio(audio_url, temp_path)

        chunks = chunk_audio_by_vad(temp_path)
        logger.info("Task %s: split into %d chunks.", self.request.id, len(chunks))

        transcript_segments = []
        for i, (chunk_audio, start_sec, end_sec) in enumerate(chunks):
            # Match the model's fp16 dtype to avoid a CUDA dtype mismatch at forward.
            inputs = whisper_processor(
                chunk_audio, sampling_rate=16000, return_tensors="pt"
            ).to("cuda", dtype=torch.float16)
            with torch.no_grad():
                generated_ids = whisper_model.generate(inputs["input_features"])
            transcription = whisper_processor.batch_decode(
                generated_ids, skip_special_tokens=True
            )[0].strip()
            transcript_segments.append(
                {"chunk_index": i, "start": start_sec, "end": end_sec, "text": transcription}
            )

            if torch.cuda.is_available():
                mem_mb = torch.cuda.memory_allocated() / (1024 ** 2)
                logger.debug("Chunk %d done. VRAM allocated: %.2f MB", i, mem_mb)

        logger.info("Task %s: running Pyannote diarization.", self.request.id)
        diarization_result = diarization_pipeline(temp_path)

        aligned_transcript = []
        for seg in transcript_segments:
            speaker = "SPEAKER_00"
            for turn, _, speaker_label in diarization_result.itertracks(yield_label=True):
                if turn.start <= seg["start"] <= turn.end:
                    speaker = speaker_label
                    break
            aligned_transcript.append(
                {"start": seg["start"], "end": seg["end"],
                 "speaker": speaker, "text": seg["text"]}
            )

        logger.info("Task %s: aligned %d segments.", self.request.id, len(aligned_transcript))
        return {"status": "success", "segments": aligned_transcript}

    except Exception as exc:
        logger.error("Task %s failed: %s", self.request.id, traceback.format_exc())
        # Exponential backoff: 15s, 30s, 60s — absorbs transient broker/GPU stalls.
        raise self.retry(exc=exc, countdown=15 * (2 ** self.request.retries))
    finally:
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            logger.debug("Task %s: CUDA cache cleared.", self.request.id)

The finally block runs on both success and failure, so the temp file and the allocator cache are released even when a chunk raises mid-loop. The exponential countdown (15s → 30s → 60s) gives a contended GPU or a flapping broker time to recover before the next attempt instead of hammering it.

Step 3 — Launching workers

Start a dedicated GPU worker with a single-process pool so CUDA context initialization can never race:

# celery==5.3.6
celery -A tasks.transcription worker \
  --loglevel=info \
  --pool=solo \
  --hostname=gpu_worker@%h

--pool=solo enforces single-threaded execution per worker, which prevents two threads from touching the same CUDA context during model warmup. Run one worker per physical GPU and pin it with CUDA_VISIBLE_DEVICES.

Verification

Confirm the three invariants directly. First, watch steady-state VRAM — it should plateau, not climb monotonically:

# Should hold roughly flat between tasks; a steady climb means leaked context.
nvidia-smi --query-gpu=memory.used --format=csv -l 5

Second, check that the broker reflects late acknowledgment — an in-flight job stays unacked until it finishes:

# Redis broker: queued (waiting) tasks for the default queue.
redis-cli -h redis-broker -n 0 LLEN celery

Third, prove requeue-on-loss: kill a worker mid-task and confirm the job reappears rather than vanishing.

# In another shell, SIGKILL the worker child during a task, then:
celery -A tasks.transcription inspect reserved   # job should be re-reserved

For live queue depth, retry counts, and per-task runtime, run celery flower and watch the task succeeded/retried/failed counters move as a batch drains.

Failure Modes & Edge Cases

Edge case Symptom Remediation
Two tasks on one GPU CUDA out of memory + SIGKILL mid-generation Keep worker_concurrency = 1 and worker_prefetch_multiplier = 1; one task per GPU
Early-ack job loss Transcript silently missing after a worker crash task_acks_late = True and task_reject_on_worker_lost = True
VRAM creep over a long batch nvidia-smi used-memory climbs each task Lower worker_max_tasks_per_child to force earlier process recycling
Chunk split mid-word Truncated words at 30s boundaries Tune top_db; align cuts to VAD silences, not fixed windows
fp16/fp32 dtype mismatch RuntimeError: expected scalar type Half at forward Cast inputs with .to("cuda", dtype=torch.float16) to match the model
Job exceeds time limit SoftTimeLimitExceeded, then hard SIGKILL Split the source earlier or raise task_time_limit; route stragglers to a fallback
Low-quality source audio Whisper hallucinated phantom phrases Pre-trim with librosa.effects.trim() before chunking

For boundary cases where chunk edges drift relative to the picture, aligning speaker labels with video cuts covers the downstream correction. If hallucinated text survives normalization, reducing hallucinations in transcription outputs addresses the decoding side.

FAQ

Why one task per GPU worker instead of higher concurrency?

Whisper Large V3 in fp16 reserves several gigabytes of activations per in-flight generation, and Pyannote adds its own footprint. With more than one concurrent task on a single GPU, peak allocations overlap and trigger CUDA out of memory. Setting worker_concurrency = 1 and worker_prefetch_multiplier = 1 makes each worker reserve exactly one job when idle, so you scale by adding workers (one per GPU), not threads.

Do I lose a transcript if a worker is killed by OOM?

Not if you use late acknowledgment. With task_acks_late = True and task_reject_on_worker_lost = True, a job is only acked after it returns successfully; a SIGKILL returns it to the broker for another worker to pick up. This requires idempotent tasks — keyed on task_id so a re-run overwrites rather than duplicates results.

What does worker_max_tasks_per_child actually fix?

PyTorch’s caching allocator fragments over many variable-length generations, so nvidia-smi shows used memory creeping up even though each task calls torch.cuda.empty_cache(). Recycling the worker child after a fixed number of tasks restarts the process cleanly and reclaims that fragmented memory without restarting the whole fleet. Tune it down if you still see creep across a long batch.