Stop ML jobs failing with CUDA out-of-memory when workers share one GPU

[mihow]

Summary

When several worker processes share one GPU, ML jobs can fail with waves of CUDA out of memory errors. The root cause is that the worker ran every GPU forward pass at the size of the API fetch batch (AMI_ANTENNA_API_BATCH_SIZE, default 24 full-resolution images) and stacked every detection crop from the batch into a single classifier call. The AMI_LOCALIZATION_BATCH_SIZE and AMI_CLASSIFICATION_BATCH_SIZE settings — which exist precisely to cap GPU memory — were silently ignored on the worker path. In a production deployment (August 2026), this appeared as a job crawling at 3.1 images/min with 512 out-of-memory errors before being killed as stale at 7% progress: the peak working set of a busy process was 13–17.4 GiB, on a 24 GiB card shared by four worker processes.

This PR makes the batch-size settings actually govern the forward-pass size, chunks crop classification so dense images can't blow the peak, retries with smaller chunks under transient pressure from a co-tenant process, and releases GPU memory at the end of every job so idle processes stop holding VRAM their neighbours need.

Important framing: this is a peak working set problem, not a leak. The failing allocation was the same 426 MiB size, 504 times, at a stable 13.49 GiB allocated — a fixed peak colliding with a co-tenant on a shared card, not upward drift. A separate, unrelated defect (host RAM growing with work done in the DataLoader subprocesses, which never touch CUDA) is addressed independently — see PR #163 (worker drain-and-recycle, implementing #147) and PR #148. The two were conflated in earlier triage; they have different memory, different processes, different mechanisms, and different fixes.

List of Changes

1. ML jobs no longer exceed GPU memory when the fetch batch is larger than what fits in one forward pass. The worker's detector and classifiers are now constructed with AMI_LOCALIZATION_BATCH_SIZE / AMI_CLASSIFICATION_BATCH_SIZE (mirroring the synchronous API path), and a chunked-inference helper (_predict_in_chunks) slices every forward pass to that size. AMI_ANTENNA_API_BATCH_SIZE now controls only fetch/transfer granularity.
2. Images dense with insects no longer spike the memory peak during classification. Crops are sliced, transformed, stacked, and classified one chunk at a time (_classify_crops_in_chunks) instead of every crop being built up front — the detector permits up to 500 detections per image, so the up-front stack was unbounded in practice.
3. A transient memory spike from a neighbouring process degrades throughput instead of failing the batch. On a CUDA allocation failure the chunk size is halved and retried, down to single items; the next batch starts fresh at the configured size.
4. Idle workers return GPU memory to the shared card. At the end of every job the models, the CUDA prefetcher's buffered batch, and the cached allocator blocks are released, so a process idling between jobs drops to near context-only.
5. The worker sets no CUDA allocator default. An earlier commit on this branch defaulted the allocator to expandable_segments:True; deployed on NVIDIA vGPU (H100 24 GB vGPU profile) it raises CUDA driver error: operation not supported at CUDA initialization — PyTorch does not fall back. That default is removed; the worker now only logs the effective PYTORCH_ALLOC_CONF / PYTORCH_CUDA_ALLOC_CONF values at startup so an operator opt-in stays visible.
6. A reproducible pressure benchmark (trapdata/antenna/gpu_memory_bench.py) runs N worker processes against an in-process mock API on one GPU to reproduce the shared-card exhaustion signature and compare configurations or revisions.
7. Unit tests cover both chunk helpers (chunking, order preservation, invalid-bbox skipping, OOM halve-and-retry, non-memory errors propagating) and pin that the batch-size settings reach the models (trapdata/antenna/tests/test_gpu_memory.py).

Measured results from production

The operational half of this mitigation shipped ahead of the code: one worker process per 24 GiB card instead of two, and fetch batch size 8. Measured outcome: a production job with ~3,200 large images completed at 88.3 images/min with zero CUDA errors (its previous run: 3.1 images/min, 512 OOM errors, killed at 7%); a ~400-image job completed at 103.4 images/min with zero errors. Throughput went up with half the processes — the second process was causing thrash, not adding capacity. The code changes in this PR bound the per-process peak so that process count and batch sizes have predictable headroom instead of relying on process count alone.

What still needs verification

• There is no GPU in the environment this was developed in. The chunked forward passes are covered by CPU-path unit tests only; an end-to-end GPU run should confirm results are unchanged (sub-batching can change FasterRCNN's internal padding for mixed-size batches, which can perturb detections near image borders negligibly).
• Post-fix peak VRAM per busy process is an estimate from code reading (roughly 5–9 GiB with default settings), not a measurement — observe nvidia-smi during a real job.
• Whether the halve-and-retry backoff ever fires in steady state; its warning log is the signal that the card is still over-committed and batch sizes or process count need tuning.
• A before/after run of gpu_memory_bench.py on real hardware.

Test and lint results

• trapdata/antenna/tests/test_gpu_memory.py: 15 passed.
• Full suite (uv run pytest --import-mode=importlib): 60 passed, 1 skipped, 1 failed — the failure is trapdata/api/tests/test_models.py::TestSourceImageSchema::test_url, which live-fetches an image from Wikimedia and received an HTTP 400 from the external server; pre-existing environmental flake, unrelated to this branch.
• pre-commit (black, isort, flake8 + bugbear, autoflake): all hooks pass.

Base branch

Opened against fix/result-post-chunking (PR #149) — this branch is stacked on it, so the diff here shows only this PR's work. Merge #149 first, then this PR can retarget main.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HcXFHJRXMrsHPX7xz9ZifF

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 3a759b36-3460-48ff-9e60-31a75eeb31e0

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}