diff --git a/docs/claude/planning/gpu-oom-mitigation.md b/docs/claude/planning/gpu-oom-mitigation.md new file mode 100644 index 00000000..b81159c0 --- /dev/null +++ b/docs/claude/planning/gpu-oom-mitigation.md @@ -0,0 +1,223 @@ +# GPU Out-of-Memory Mitigation for the Antenna Worker + +Status: implemented on branch `fix/gpu-oom-mitigation` (August 2026). The +operational half of the mitigation (fewer worker processes per card, smaller +fetch batches) has been confirmed in production; see "Production results" +below. The allocator default originally proposed here (item 4) was tried in +production, failed on the deployed hardware, and has been removed from the +branch — see "Correction: the allocator default". + +## Problem + +In a production deployment (August 2026), ML jobs failed with hundreds of +`CUDA out of memory` errors plus a few `CUBLAS_STATUS_ALLOC_FAILED` errors. +The deployment shape: one 24 GiB GPU shared by four worker processes serving +multiple environments. Measured from the error logs: + +- Failed allocations were consistently 426 MiB (504 occurrences) or 736 MiB + (8 occurrences) — the same allocation retried batch after batch, not random + sizes. Each failed batch is caught in `_process_batch` and reported to + Antenna as per-image task errors, so the job keeps running and keeps + failing on every subsequent batch. +- Free memory at failure ranged 61–457 MiB; the failing process held + ~14.5 GiB total, of which 13.49 GiB was allocated by PyTorch. +- Reserved-but-unallocated memory was consistently 757–899 MiB + (fragmentation, since the allocator could not satisfy a 426 MiB request + despite ~800 MiB reserved). +- A neighbouring process held 6.12 GiB at the same time; two others held + ~125 MiB each (CUDA context only, effectively idle). + +## Diagnosis: peak working set, not a leak + +Everything below is established by code reading unless labelled measured. + +### The GPU batch size is governed by the wrong setting + +On the Antenna worker path, the size of every GPU forward pass is +`antenna_api_batch_size` (default 24), not the GPU batch-size settings: + +- `get_rest_dataloader` (`trapdata/antenna/datasets.py:427`) sizes the + `RESTDataset` fetch with `settings.antenna_api_batch_size`, and each + yielded batch is the entire API fetch, pre-collated + (`datasets.py:317-320`). +- `_process_batch` runs the whole batch of full-resolution (~4K) images + through FasterRCNN in a single forward call + (`trapdata/antenna/worker.py:254`). Neither + `InferenceBaseClass.predict_batch` (`trapdata/ml/models/base.py:276`) nor + `ObjectDetector.predict_batch` (`trapdata/ml/models/localization.py:163`) + chunks its input. +- All crops from all detections in the batch are stacked into one tensor and + classified in a single forward call (`worker.py:321-322` for the terminal + classifier, `worker.py:176-177` for the binary filter). The detector allows + up to 500 detections per image (`ml/models/localization.py:255`), so this + stack is unbounded in practice. +- `localization_batch_size` (default 8, described in settings as "reduce this + if you run out of memory") and `classification_batch_size` (default 20) are + never referenced under `trapdata/antenna/`. The synchronous FastAPI path + does pass them (`trapdata/api/api.py:223,239,279`); the worker path was + built without them. Even `benchmark.py:73` sets + `settings.localization_batch_size` from its `--gpu-batch-size` flag, but + nothing on the worker path ever reads it — the knob is a silent no-op. + +### What the peak is made of (rough estimates from code reading) + +- Current batch of 24 full-resolution images as float32 tensors: + ~2.4 GiB (24 × 3 × 2160 × 4096 × 4 B). +- `CUDAPrefetcher` (`datasets.py:445`) holds the *next* full batch on the GPU + while the current one is processed — roughly another 2.4 GiB. +- `image_tensors` (`worker.py:275`) keeps all full-resolution GPU tensors + alive through detection *and* classification (needed for crop slicing). +- FasterRCNN activations for a 24-image forward pass (internally resized to + ≤1333 px): several GiB. +- Model weights for detector + binary filter + terminal classifier: ~1–2 GiB. + +These estimates are consistent with the measured 13.49 GiB PyTorch-allocated +peak. With four processes sharing 24 GiB, the card is exhausted whenever more +than one process is mid-batch at the same time — which also explains why the +same configuration often works: a single busy process fits. + +### Why this is not a leak + +- Per-batch intermediates are local to `_process_batch` (`worker.py:203`, the + docstring states this deliberately) and `torch.cuda.empty_cache()` runs + after every batch (`worker.py:395`). +- Models are constructed per job, lazily on the first batch + (`worker.py:464-478`), and are function locals released by reference + counting when `_process_job` returns. There is no cross-job model cache. +- An existing regression test (`trapdata/antenna/tests/test_memory_leak.py`) + pins host-RSS stability across batches. A prior analysis of host-RAM + blowup (DataLoader `pin_memory` × `prefetch_factor`, `datasets.py:439-441`) + concerns host memory in the DataLoader subprocesses, which never touch + CUDA — a separate problem from this one. + +One cross-job gap does exist: nothing releases cached allocator blocks when a +job *ends*. `empty_cache()` runs per batch and at the start of the *next* +claimed job (`worker.py:436`), so a process idling between jobs retains +reserved VRAM (roughly its freed model weights plus remnants) that co-tenant +processes on the shared card cannot use. + +### Allocator configuration + +No `PYTORCH_ALLOC_CONF` / `PYTORCH_CUDA_ALLOC_CONF` is set anywhere in the +repo, and `torch.cuda.set_per_process_memory_fraction` is not used. The +measured 757–899 MiB reserved-but-unallocated at failure is the fragmentation +signature that `expandable_segments:True` targets (the OOM message itself +recommends it). However, that option is not usable on the deployed vGPU +hardware — see "Correction: the allocator default" below before considering +it. + +## Fix + +Smallest changes that address the dominant cause, in order of impact: + +1. **Honor the GPU batch-size settings on the worker path.** Construct the + worker's models with `batch_size` from settings (mirroring the FastAPI + path): detector gets `localization_batch_size`, binary filter and terminal + classifier get `classification_batch_size`. Add a chunked-inference helper + in `worker.py` that runs `predict_batch` + `post_process_batch` over the + input in chunks of `model.batch_size`, so the forward-pass peak is capped + regardless of `antenna_api_batch_size`. Results are unchanged: detection + and per-crop classification are independent per item, and softmax is + per-row. (For mixed-size image batches, sub-batching can change FasterRCNN's + internal padding, which can perturb detections near image borders + negligibly.) +2. **Bounded adaptive backoff.** If a chunk still hits CUDA OOM (a co-tenant + process spiking), the helper halves the chunk size and retries, down to a + chunk of 1, calling `empty_cache()` between attempts. The reduction lasts + for the remainder of that call; the next batch starts fresh from the + configured size, so transient neighbour pressure shrinks chunks only while + it persists. +3. **Release GPU memory at the end of every job.** In `_process_job`'s + `finally` block, drop references to the models, prefetcher, and loader, + then call `empty_cache()`, so an idle process returns cached VRAM to the + shared card instead of holding it until its next job. +4. ~~**Default the allocator to `expandable_segments:True`.**~~ **Withdrawn — + see "Correction: the allocator default" below.** The worker sets no + allocator default; it only logs the effective `PYTORCH_ALLOC_CONF` / + `PYTORCH_CUDA_ALLOC_CONF` values at startup. Operators can opt in via the + environment on hardware that supports it. +5. **Update stale documentation** in `datasets.py` (which currently says the + async worker uses `antenna_api_batch_size` for the GPU batch) and the + worker-tuning notes. + +After this change `antenna_api_batch_size` controls only fetch granularity +and how many downloaded images are resident per batch — no longer the +forward-pass size. Operators can additionally lower it to shrink the resident +image tensors and the prefetcher's double-buffer. + +## Deliberately not changed + +- **`CUDAPrefetcher` double-buffering** — a throughput feature with a bounded + cost (one extra batch of images); lowering `antenna_api_batch_size` shrinks + it without a code change. +- **Per-batch `empty_cache()`** (`worker.py:395`) — slightly hurts throughput + but is polite on a shared card; not the target of this fix. +- **`set_per_process_memory_fraction`** — would relabel the failure (OOM at + the cap) without reducing demand, and caps legitimate bursts when the card + is otherwise idle. +- **Cross-job model caching** — would avoid per-job weight reloads but pins + VRAM in idle processes, the opposite of what a shared card needs. +- **Default values of the batch-size settings** — 8 (localization) and + 20 (classification) are the long-standing defaults of the synchronous path. + +## Correction: the allocator default (tried in production, removed) + +An earlier revision of this plan (item 4) proposed defaulting the CUDA +allocator to `expandable_segments:True` at worker startup, and this document +claimed that "PyTorch falls back with a warning where unsupported". **Both +claims were wrong.** When the default was deployed to production hosts +running on NVIDIA vGPU (H100 24 GB vGPU profile), every CUDA initialization +hard-failed with: + +``` +CUDA driver error: operation not supported +``` + +Both hosts failed identically and one job run was lost. `expandable_segments` +requires the CUDA virtual-memory-management driver APIs, which the vGPU +profile does not expose, and PyTorch raises rather than falling back. Do not +re-attempt an allocator default here: the worker now sets nothing and only +logs the effective `PYTORCH_ALLOC_CONF` / `PYTORCH_CUDA_ALLOC_CONF` values at +startup, so an operator-set value is visible in the logs. Opting in remains +possible via the environment on hardware that supports it. + +## Production results (measured, August 2026) + +The operational half of this mitigation was applied to the affected +deployment before the code changes shipped: worker processes per 24 GiB card +reduced from two to one (per environment pool), and the API fetch batch size +set to 8. Measured outcome: + +- A production job with ~3,200 large images completed 100% at 88.3 + images/min with zero CUDA out-of-memory errors. The previous run of the + same workload ran at 3.1 images/min, produced 512 `CUDA out of memory` + errors, and was killed by the stale-job reaper at 7% progress. +- A second job (~400 images) completed at 103.4 images/min with zero errors. + +Throughput went *up* with half the processes: the second process on the card +was causing allocator thrash and failed-batch retries, not adding capacity. +This confirms the peak-working-set diagnosis above (13–17.4 GiB per busy +process; two busy processes cannot coexist on a 24 GiB card). + +## What still needs verification (no GPU available in this environment) + +- That the chunked detector/classifier forward passes produce identical + results on a real GPU end-to-end run (CPU-path tests pass; the CUDA + prefetcher path is not exercised without a GPU). +- The actual post-fix peak VRAM per busy process (estimated, not measured: + roughly 5–9 GiB with defaults) — observe `nvidia-smi` during a real job. +- Whether the backoff path ever triggers in steady state (its warning log is + the signal that the card is still over-committed and process count or batch + sizes need ops-side tuning). + +## Validation protocol for production + +1. Deploy to one host; leave the rest unchanged as control. +2. During a real job, watch per-process GPU memory (`nvidia-smi`) — expect + busy-process peaks well below the previous ~14.5 GiB, and idle processes + dropping to near context-only after a job completes. +3. Grep worker logs for `CUDA out of memory` recurrence and for the new + chunk-size reduction warnings. +4. Compare seconds-per-image from the per-batch log line before and after — + chunking the detector forward is expected to be roughly throughput-neutral + since the GPU already serializes the work internally. diff --git a/trapdata/antenna/datasets.py b/trapdata/antenna/datasets.py index 7ecc7bdf..98d9ba18 100644 --- a/trapdata/antenna/datasets.py +++ b/trapdata/antenna/datasets.py @@ -34,9 +34,16 @@ Settings quick-reference (prefix with AMI_ as env vars): localization_batch_size (default 8) - How many images the GPU processes at once (detection). Larger = - more GPU memory. These are full-resolution images (~4K). - Async worker use antennna_api_batch_size for this. + How many full-resolution (~4K) images go through the detector in + one forward pass. The worker slices each fetched batch into chunks + of this size before inference (``_predict_in_chunks`` in worker.py), + so this — not antenna_api_batch_size — bounds the detector's GPU + memory peak. + + classification_batch_size (default 20) + How many detection crops go through the binary and species + classifiers in one forward pass, using the same chunking mechanism + as above. num_workers (default 4) DataLoader subprocesses. Each independently fetches tasks and @@ -44,12 +51,14 @@ GPU, at the cost of CPU/RAM. With 0 workers, fetching and inference are sequential (useful for debugging). - antenna_api_batch_size (default 16) + antenna_api_batch_size (default 24) How many task URLs to request from Antenna per API call. Determines how many images are downloaded concurrently per - thread pool invocation. Should be >= localization_batch_size - so one API call can fill at least one GPU batch without an - extra round trip. + thread pool invocation, and therefore how many decoded image + tensors are resident per batch (doubled while the CUDA prefetcher + holds the next batch). Should be >= localization_batch_size so one + API call can fill at least one GPU inference chunk without an + extra round trip; lowering it reduces GPU memory residency. prefetch_factor (PyTorch default: 2 when num_workers > 0) Batches prefetched per worker. Not overridden here — the @@ -421,7 +430,8 @@ def get_rest_dataloader( job_id: Job ID to fetch tasks for settings: Settings object. Relevant fields: - antenna_api_base_url / antenna_api_auth_token - - antenna_api_batch_size (tasks per API call and GPU batch size) + - antenna_api_batch_size (tasks per API call; GPU inference is + chunked separately by the batch-size settings, see worker.py) - num_workers (DataLoader subprocesses) """ dataset = RESTDataset( diff --git a/trapdata/antenna/gpu_memory_bench.py b/trapdata/antenna/gpu_memory_bench.py new file mode 100644 index 00000000..63728e92 --- /dev/null +++ b/trapdata/antenna/gpu_memory_bench.py @@ -0,0 +1,408 @@ +"""Multi-process GPU memory pressure benchmark for the Antenna worker. + +Reproduces the failure mode where several worker processes share one GPU and +collectively exhaust it: each process runs the real worker code path +(``_process_job`` with the real detector, binary filter, and species +classifier) against an in-process mock Antenna API serving the bundled +full-resolution (4096x2160) test images. This is the same infrastructure the +integration tests use, so no network or external services are involved. + +The point is to demonstrate *concurrent* pressure — a single process +allocating until it dies only proves the GPU is finite. Run it with the +process count and batch sizes that match a deployment to reproduce the +failure signature (a few-hundred-MiB allocation failing with little free +memory on the card), then run it again with the fix or different parameters +to demonstrate the difference. + +Usage (single 24 GiB GPU, shapes similar to a shared production card):: + + python -m trapdata.antenna.gpu_memory_bench \ + --processes 4 --jobs 2 --tasks-per-job 48 \ + --api-batch-size 24 --device 0 + +To compare against an older revision of the worker, copy this file into a +git worktree of that revision and run it from there with the same arguments. + +What it reports, per process: + +- ``torch.cuda.memory_allocated`` / ``memory_reserved`` / + ``max_memory_allocated`` and device-wide free memory at: start, after a + reference model load, after each batch, and between jobs. +- The count of task results whose error carries a CUDA allocation-failure + signature ("out of memory" / ``ALLOC_FAILED``), plus the first such + message (which includes free-at-failure details reported by the + allocator). + +How to read the between-jobs numbers: if memory held after a job completes +keeps growing job over job, something is not being released (a leak). If it +returns to a stable baseline and failures only appear with several +concurrent processes, the card is oversubscribed and the levers are process +count and batch sizes. +""" + +import argparse +import multiprocessing +import os +import queue +import sys +import time +from pathlib import Path + +GIB = 1024**3 + +# The allocator env vars torch recognizes (newer and older spelling). +_ALLOC_ENV_VARS = ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF") + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Multi-process GPU memory pressure benchmark for the worker" + ) + parser.add_argument( + "--processes", + type=int, + default=4, + help="Worker processes sharing the GPU (default: 4)", + ) + parser.add_argument( + "--jobs", + type=int, + default=2, + help="Sequential jobs per process (default: 2)", + ) + parser.add_argument( + "--tasks-per-job", + type=int, + default=48, + help="Tasks (images) per job (default: 48)", + ) + parser.add_argument( + "--api-batch-size", + type=int, + default=24, + help="antenna_api_batch_size: images fetched per API call (default: 24)", + ) + parser.add_argument( + "--localization-batch-size", + type=int, + default=8, + help="Detector GPU batch size setting (default: 8)", + ) + parser.add_argument( + "--classification-batch-size", + type=int, + default=20, + help="Classifier GPU batch size setting (default: 20)", + ) + parser.add_argument( + "--pipeline", + default="quebec_vermont_moths_2023", + help="Pipeline slug to run (default: quebec_vermont_moths_2023)", + ) + parser.add_argument( + "--alloc-conf", + default=None, + help=( + "Value for the CUDA allocator env vars, e.g. " + "'expandable_segments:True'. Default: leave the environment " + "as-is. Applied in each child before any CUDA allocation. " + "Note: expandable_segments raises 'CUDA driver error: operation " + "not supported' on hardware without virtual-memory-management " + "driver APIs (e.g. NVIDIA vGPU)." + ), + ) + parser.add_argument( + "--device", + type=int, + default=0, + help="CUDA device index to load (default: 0)", + ) + return parser.parse_args(argv) + + +def _sample(rank: int, device: int, phase: str) -> dict: + """Record and print the CUDA memory counters for this process.""" + import torch + + free_b, total_b = torch.cuda.mem_get_info(device) + stats = { + "phase": phase, + "alloc_gib": torch.cuda.memory_allocated(device) / GIB, + "reserved_gib": torch.cuda.memory_reserved(device) / GIB, + "max_alloc_gib": torch.cuda.max_memory_allocated(device) / GIB, + "free_gib": free_b / GIB, + "total_gib": total_b / GIB, + } + print( + f"[P{rank}] {phase}: " + f"alloc={stats['alloc_gib']:.2f} GiB, " + f"reserved={stats['reserved_gib']:.2f} GiB, " + f"max_alloc={stats['max_alloc_gib']:.2f} GiB, " + f"device free={stats['free_gib']:.2f}/{stats['total_gib']:.2f} GiB", + flush=True, + ) + return stats + + +def _is_memory_error_message(message: str) -> bool: + return "out of memory" in message or "ALLOC_FAILED" in message + + +def _load_model_reference(pipeline: str, batch_sizes: dict, on_loaded) -> None: + """Load the pipeline's model stack once, sample via ``on_loaded``, release. + + Gives an "after model load" baseline to compare the between-jobs numbers + against, and measures the model footprint itself. The sample callback runs + while the models are alive; they are released when this function returns. + """ + from trapdata.api.api import CLASSIFIER_CHOICES, should_filter_detections + from trapdata.api.models.classification import MothClassifierBinary + from trapdata.api.models.localization import APIMothDetector + + classifier_class = CLASSIFIER_CHOICES[pipeline] + models = [ + classifier_class( + source_images=[], + detections=[], + batch_size=batch_sizes["classification"], + ), + APIMothDetector([], batch_size=batch_sizes["localization"]), + ] + if should_filter_detections(classifier_class): + models.append( + MothClassifierBinary( + source_images=[], + detections=[], + terminal=False, + batch_size=batch_sizes["classification"], + ) + ) + on_loaded() + del models + + +def _child_main( + rank: int, + args: argparse.Namespace, + barrier, + results_queue, +) -> None: + """One simulated worker process: run jobs and report memory samples.""" + # Allocator config must be in the environment before the first CUDA + # allocation of this process to take effect. + if args.alloc_conf is not None: + for var in _ALLOC_ENV_VARS: + os.environ[var] = args.alloc_conf + + import torch + from fastapi.testclient import TestClient + + from trapdata.antenna.schemas import AntennaPipelineProcessingTask + from trapdata.antenna.tests import antenna_api_server + from trapdata.antenna.tests.antenna_api_server import app as antenna_app + from trapdata.antenna.worker import _process_job + from trapdata.api.tests.image_server import StaticFileTestServer + from trapdata.api.tests.utils import patch_antenna_api_requests + from trapdata.settings import Settings + from trapdata.tests import TEST_IMAGES_BASE_PATH + + device = args.device + torch.cuda.set_device(device) + effective_conf = {var: os.environ.get(var, "(unset)") for var in _ALLOC_ENV_VARS} + print(f"[P{rank}] allocator env: {effective_conf}", flush=True) + + samples: list[dict] = [] + oom_messages: list[str] = [] + job_exceptions: list[str] = [] + + samples.append(_sample(rank, device, "start")) + + batch_sizes = { + "localization": args.localization_batch_size, + "classification": args.classification_batch_size, + } + _load_model_reference( + args.pipeline, + batch_sizes, + on_loaded=lambda: samples.append(_sample(rank, device, "model-load reference")), + ) + torch.cuda.empty_cache() + samples.append(_sample(rank, device, "post-release baseline")) + + # Each process runs its own in-process mock Antenna API and image server. + images_dir = Path(TEST_IMAGES_BASE_PATH) + file_server = StaticFileTestServer(images_dir) + file_server.start() + client = TestClient(antenna_app, follow_redirects=False) + + image_paths = sorted((images_dir / "vermont").glob("*.jpg")) + image_urls = [file_server.get_url(p.relative_to(images_dir)) for p in image_paths] + + settings = Settings() + settings.antenna_api_base_url = "http://testserver/api/v2" + settings.antenna_api_auth_token = "benchmark-token" + settings.antenna_api_batch_size = args.api_batch_size + settings.num_workers = 0 + settings.localization_batch_size = args.localization_batch_size + settings.classification_batch_size = args.classification_batch_size + + # Line up all processes so the jobs actually overlap. + barrier.wait() + start_time = time.monotonic() + + between_jobs_alloc: list[float] = [] + try: + for j in range(args.jobs): + job_id = 9000 + rank * 100 + j + antenna_api_server.reset() + tasks = [ + AntennaPipelineProcessingTask( + id=f"task_{rank}_{j}_{i}", + image_id=f"img_{rank}_{j}_{i}", + image_url=image_urls[i % len(image_urls)], + reply_subject=f"reply_{rank}_{j}_{i}", + ) + for i in range(args.tasks_per_job) + ] + antenna_api_server.setup_job(job_id=job_id, tasks=tasks) + + torch.cuda.reset_peak_memory_stats(device) + + def on_batch(batch_num: int, items: int, job_index: int = j): + label = f"job {job_index} batch {batch_num + 1}" + if batch_num == 0: + label += " (models loaded)" + samples.append(_sample(rank, device, label)) + + try: + with patch_antenna_api_requests(client): + _process_job( + args.pipeline, + job_id, + settings, + device=torch.device("cuda", device), + on_batch_complete=on_batch, + ) + except Exception as e: + # A job-level failure (e.g. OOM during prefetch) is itself a + # data point; record it and keep the process running. + job_exceptions.append(f"job {j}: {type(e).__name__}: {e}") + print(f"[P{rank}] job {j} raised: {e}", flush=True) + + stats = _sample(rank, device, f"between jobs (after job {j})") + between_jobs_alloc.append(stats["alloc_gib"]) + + # Collect allocation-failure signatures from the posted results. + for result in antenna_api_server.get_posted_results(job_id): + error = getattr(result.result, "error", None) + if error and _is_memory_error_message(error): + oom_messages.append(error) + finally: + file_server.stop() + + elapsed = time.monotonic() - start_time + if oom_messages: + print( + f"[P{rank}] first allocation-failure message: {oom_messages[0][:400]}", + flush=True, + ) + + results_queue.put( + { + "rank": rank, + "elapsed_s": elapsed, + "samples": samples, + "between_jobs_alloc_gib": between_jobs_alloc, + "oom_error_count": len(oom_messages), + "first_oom_message": oom_messages[0] if oom_messages else None, + "job_exceptions": job_exceptions, + "peak_alloc_gib": max(s["max_alloc_gib"] for s in samples), + } + ) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + + import torch + + if not torch.cuda.is_available(): + print("CUDA is not available; this benchmark needs a GPU.", file=sys.stderr) + return 2 + + # The parent deliberately avoids creating a CUDA context (it would hold + # a few hundred MiB on the device and distort the measurements); the + # children report device totals themselves. + print( + f"Device index {args.device}: " + f"{args.processes} processes x {args.jobs} jobs x " + f"{args.tasks_per_job} tasks, api_batch_size={args.api_batch_size}, " + f"localization_batch_size={args.localization_batch_size}, " + f"classification_batch_size={args.classification_batch_size}, " + f"alloc_conf={args.alloc_conf!r}", + flush=True, + ) + + ctx = multiprocessing.get_context("spawn") + barrier = ctx.Barrier(args.processes) + results_queue = ctx.Queue() + procs = [ + ctx.Process( + target=_child_main, + args=(rank, args, barrier, results_queue), + name=f"gpu-bench-{rank}", + ) + for rank in range(args.processes) + ] + for p in procs: + p.start() + + results = [] + for _ in procs: + try: + # Generous ceiling; a normal run finishes in minutes. A missing + # result means a child died hard (e.g. CUDA abort) — report it + # rather than hanging forever. + results.append(results_queue.get(timeout=1800)) + except queue.Empty: + print("Timed out waiting for a worker process result.", flush=True) + break + for p in procs: + p.join(timeout=30) + if p.is_alive(): + print(f"Terminating unresponsive process {p.name}", flush=True) + p.terminate() + + results.sort(key=lambda r: r["rank"]) + total_oom = sum(r["oom_error_count"] for r in results) + + print("\n=== Summary ===", flush=True) + for r in results: + between = ", ".join(f"{a:.2f}" for a in r["between_jobs_alloc_gib"]) + print( + f"[P{r['rank']}] peak_alloc={r['peak_alloc_gib']:.2f} GiB, " + f"between-jobs alloc per job=[{between}] GiB, " + f"oom_errors={r['oom_error_count']}, " + f"job_exceptions={len(r['job_exceptions'])}, " + f"elapsed={r['elapsed_s']:.0f}s", + flush=True, + ) + print( + f"\nAllocation-failure signature reproduced: " + f"{'YES' if total_oom else 'NO'} " + f"({total_oom} allocation-failure task errors across " + f"{args.processes} processes)", + flush=True, + ) + print( + "Interpretation: growing between-jobs alloc across jobs suggests a " + "leak; a stable between-jobs baseline with failures only under " + "concurrency indicates an oversubscribed card (tune process count / " + "batch sizes).", + flush=True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/trapdata/antenna/tests/test_gpu_memory.py b/trapdata/antenna/tests/test_gpu_memory.py new file mode 100644 index 00000000..64a1230a --- /dev/null +++ b/trapdata/antenna/tests/test_gpu_memory.py @@ -0,0 +1,249 @@ +"""Unit tests for the worker's GPU memory management. + +Covers the chunked-inference helpers (``_predict_in_chunks`` and +``_classify_crops_in_chunks``) and the per-job model factory that applies +the GPU batch-size settings. All tests use fake models and run on CPU; the +out-of-memory scenarios are simulated by raising the same exception types +the CUDA allocator raises. +""" + +from types import SimpleNamespace +from unittest import TestCase +from unittest.mock import MagicMock, patch + +import torch + +from trapdata.antenna.worker import ( + _classify_crops_in_chunks, + _init_job_models, + _is_cuda_memory_error, + _predict_in_chunks, +) + + +class FakeModel: + """Inference stub that records the chunk sizes ``predict_batch`` receives. + + ``post_process_batch`` negates each item so tests can verify that outputs + pass through post-processing and stay in input order. ``fail_first_n`` + makes the first N ``predict_batch`` calls raise a simulated CUDA + out-of-memory error before any work is recorded. + """ + + def __init__(self, batch_size: int, fail_first_n: int = 0): + self.batch_size = batch_size + self.seen_chunk_sizes: list[int] = [] + self._failures_remaining = fail_first_n + + def predict_batch(self, chunk): + if self._failures_remaining > 0: + self._failures_remaining -= 1 + raise torch.OutOfMemoryError("CUDA out of memory. (simulated)") + self.seen_chunk_sizes.append(len(chunk)) + return list(chunk) + + def post_process_batch(self, output): + return [-item for item in output] + + +class TestPredictInChunks(TestCase): + """The detector's forward-pass peak must be bounded by its batch_size.""" + + def _items(self, n: int) -> list[torch.Tensor]: + return [torch.tensor(float(i)) for i in range(n)] + + def test_list_input_is_chunked_and_order_preserved(self): + model = FakeModel(batch_size=4) + results = _predict_in_chunks(model, self._items(10)) + + assert model.seen_chunk_sizes == [4, 4, 2] + assert [float(r) for r in results] == [-float(i) for i in range(10)] + + def test_tensor_input_is_sliced(self): + model = FakeModel(batch_size=4) + results = _predict_in_chunks(model, torch.arange(10.0)) + + assert model.seen_chunk_sizes == [4, 4, 2] + assert [float(r) for r in results] == [-float(i) for i in range(10)] + + def test_single_chunk_when_batch_size_exceeds_items(self): + model = FakeModel(batch_size=8) + results = _predict_in_chunks(model, self._items(3)) + + assert model.seen_chunk_sizes == [3] + assert len(results) == 3 + + def test_oom_halves_chunk_size_and_retries(self): + """A simulated allocation failure must shrink the chunk, not fail the batch.""" + model = FakeModel(batch_size=8, fail_first_n=1) + results = _predict_in_chunks(model, self._items(8)) + + # The failed 8-item attempt is not recorded; the retry runs at 4. + assert model.seen_chunk_sizes == [4, 4] + assert [float(r) for r in results] == [-float(i) for i in range(8)] + + def test_oom_at_chunk_size_one_reraises(self): + model = FakeModel(batch_size=1, fail_first_n=100) + + with self.assertRaises(torch.OutOfMemoryError): + _predict_in_chunks(model, self._items(2)) + + def test_unrelated_runtime_error_propagates_without_retry(self): + model = FakeModel(batch_size=4) + model.predict_batch = MagicMock(side_effect=RuntimeError("size mismatch")) + + with self.assertRaises(RuntimeError): + _predict_in_chunks(model, self._items(4)) + # No retry: a non-memory error must fail on the first call. + assert model.predict_batch.call_count == 1 + + +def _make_detection(image_id: str, x1: int, y1: int, x2: int, y2: int): + """A minimal stand-in for DetectionResponse: just bbox and source image id.""" + return SimpleNamespace( + source_image_id=image_id, + bbox=SimpleNamespace(x1=x1, y1=y1, x2=x2, y2=y2), + ) + + +class FakeCropClassifier(FakeModel): + """FakeModel plus the transform hook ``_classify_crops_in_chunks`` needs. + + The transform trims each crop to its top-left pixel, whose value encodes + the crop's x position in the test image — so predictions can be traced + back to detections, and crops stack regardless of bbox size. + """ + + def get_transforms(self): + return lambda crop: crop[:, :1, :1] + + def post_process_batch(self, output): + # Identify each crop by the value at its top-left corner. + return [float(item[0, 0, 0]) for item in output] + + +class TestClassifyCropsInChunks(TestCase): + """Crop construction itself must be chunked, not just the forward pass. + + If all crops were built before inference, a dense batch (hundreds of + detections per image) would allocate every crop tensor up front and the + chunked forward pass would cap nothing. + """ + + def _image_tensors(self) -> dict[str, torch.Tensor]: + # One 3x100x100 "image" whose pixel values encode the x coordinate, + # so each crop is identifiable by its top-left corner value. + gradient = torch.arange(100.0).repeat(3, 100, 1) + return {"img_a": gradient} + + def test_chunked_and_mapped_back_to_detections(self): + detections = [ + _make_detection("img_a", x1=x, y1=0, x2=x + 10, y2=10) for x in range(5) + ] + model = FakeCropClassifier(batch_size=2) + + predictions, valid_indices = _classify_crops_in_chunks( + model, detections, self._image_tensors() + ) + + assert model.seen_chunk_sizes == [2, 2, 1] + assert valid_indices == [0, 1, 2, 3, 4] + # Each prediction carries its crop's top-left value = detection's x1. + assert predictions == [0.0, 1.0, 2.0, 3.0, 4.0] + + def test_invalid_bboxes_are_skipped(self): + detections = [ + _make_detection("img_a", x1=0, y1=0, x2=10, y2=10), + _make_detection("img_a", x1=5, y1=5, x2=5, y2=10), # zero width + _make_detection("img_a", x1=20, y1=0, x2=30, y2=10), + ] + model = FakeCropClassifier(batch_size=4) + + predictions, valid_indices = _classify_crops_in_chunks( + model, detections, self._image_tensors() + ) + + assert valid_indices == [0, 2] + assert predictions == [0.0, 20.0] + + def test_empty_detections(self): + model = FakeCropClassifier(batch_size=4) + predictions, valid_indices = _classify_crops_in_chunks( + model, [], self._image_tensors() + ) + + assert predictions == [] + assert valid_indices == [] + + def test_oom_halves_chunk_size_and_retries(self): + detections = [ + _make_detection("img_a", x1=x, y1=0, x2=x + 10, y2=10) for x in range(4) + ] + model = FakeCropClassifier(batch_size=4, fail_first_n=1) + + predictions, valid_indices = _classify_crops_in_chunks( + model, detections, self._image_tensors() + ) + + assert model.seen_chunk_sizes == [2, 2] + assert predictions == [0.0, 1.0, 2.0, 3.0] + + +class TestIsCudaMemoryError(TestCase): + def test_out_of_memory_error_is_detected(self): + assert _is_cuda_memory_error(torch.OutOfMemoryError("CUDA out of memory.")) + + def test_cublas_alloc_failure_is_detected(self): + exc = RuntimeError( + "CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle)" + ) + assert _is_cuda_memory_error(exc) + + def test_other_runtime_errors_are_not_detected(self): + assert not _is_cuda_memory_error(RuntimeError("size mismatch")) + + +class TestInitJobModels(TestCase): + """The worker must apply the GPU batch-size settings to its models. + + Guards against the regression where the worker constructed models without + a batch size, so the whole API fetch batch went through one forward pass + and the AMI_LOCALIZATION_BATCH_SIZE / AMI_CLASSIFICATION_BATCH_SIZE + settings were silently ignored. + """ + + def _settings(self) -> MagicMock: + settings = MagicMock() + settings.localization_batch_size = 8 + settings.classification_batch_size = 20 + return settings + + @patch("trapdata.antenna.worker.MothClassifierBinary") + @patch("trapdata.antenna.worker.APIMothDetector") + def test_batch_sizes_from_settings(self, mock_detector, mock_binary): + classifier_class = MagicMock() + + classifier, detector, binary_filter = _init_job_models( + classifier_class, use_binary_filter=True, settings=self._settings() + ) + + classifier_class.assert_called_once_with( + source_images=[], detections=[], batch_size=20 + ) + mock_detector.assert_called_once_with([], batch_size=8) + mock_binary.assert_called_once_with( + source_images=[], detections=[], terminal=False, batch_size=20 + ) + assert classifier is classifier_class.return_value + assert detector is mock_detector.return_value + assert binary_filter is mock_binary.return_value + + @patch("trapdata.antenna.worker.MothClassifierBinary") + @patch("trapdata.antenna.worker.APIMothDetector") + def test_no_binary_filter_when_not_used(self, mock_detector, mock_binary): + _, _, binary_filter = _init_job_models( + MagicMock(), use_binary_filter=False, settings=self._settings() + ) + + assert binary_filter is None + mock_binary.assert_not_called() diff --git a/trapdata/antenna/tests/test_memory_leak.py b/trapdata/antenna/tests/test_memory_leak.py index a09c14c1..1e07a2f0 100644 --- a/trapdata/antenna/tests/test_memory_leak.py +++ b/trapdata/antenna/tests/test_memory_leak.py @@ -55,6 +55,9 @@ def _make_settings(self): settings.antenna_api_batch_size = 2 settings.num_workers = 0 settings.localization_batch_size = 2 + settings.classification_batch_size = 2 + # Real integer; ResultPoster compares payload sizes against it + settings.antenna_result_post_max_bytes = 25 * 1024 * 1024 return settings @pytest.mark.slow diff --git a/trapdata/antenna/tests/test_worker.py b/trapdata/antenna/tests/test_worker.py index f6b90796..4a073719 100644 --- a/trapdata/antenna/tests/test_worker.py +++ b/trapdata/antenna/tests/test_worker.py @@ -260,6 +260,9 @@ def _make_settings(self): settings.antenna_api_batch_size = 2 settings.num_workers = 0 # Disable multiprocessing for tests settings.localization_batch_size = 2 # Real integer for batch processing + settings.classification_batch_size = 2 # Real integer for chunked inference + # Real integer; ResultPoster compares payload sizes against it + settings.antenna_result_post_max_bytes = 25 * 1024 * 1024 return settings def test_empty_queue(self): @@ -421,6 +424,9 @@ def _make_settings(self): settings.antenna_api_batch_size = 2 settings.num_workers = 0 settings.localization_batch_size = 2 # Real integer for batch processing + settings.classification_batch_size = 2 # Real integer for chunked inference + # Real integer; ResultPoster compares payload sizes against it + settings.antenna_result_post_max_bytes = 25 * 1024 * 1024 return settings def test_full_workflow_with_real_inference(self): diff --git a/trapdata/antenna/worker.py b/trapdata/antenna/worker.py index 832deb93..cb9e49f9 100644 --- a/trapdata/antenna/worker.py +++ b/trapdata/antenna/worker.py @@ -3,6 +3,7 @@ from __future__ import annotations import datetime +import os import time from collections.abc import Callable @@ -15,7 +16,7 @@ from trapdata.antenna.result_posting import ResultPoster from trapdata.antenna.schemas import AntennaTaskResult, AntennaTaskResultError from trapdata.api.api import CLASSIFIER_CHOICES, should_filter_detections -from trapdata.api.models.classification import MothClassifierBinary +from trapdata.api.models.classification import APIMothClassifier, MothClassifierBinary from trapdata.api.models.localization import APIMothDetector from trapdata.api.schemas import ( DetectionResponse, @@ -86,6 +87,19 @@ def _worker_loop(gpu_id: int, pipelines: list[str]): f"AMI worker instance {gpu_id} pinned to GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}" ) + # Log the allocator config once per process; PyTorch reads these env vars + # at the first CUDA allocation, so this line records what took effect. + # The worker deliberately sets no default: ``expandable_segments:True`` + # raises ``CUDA driver error: operation not supported`` on NVIDIA vGPU, + # which does not expose the virtual-memory-management driver APIs it + # needs. Operators can still opt in via the environment on hardware that + # supports it. + logger.info( + "CUDA allocator config: " + f"PYTORCH_ALLOC_CONF={os.environ.get('PYTORCH_ALLOC_CONF', '(unset)')}, " + f"PYTORCH_CUDA_ALLOC_CONF={os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '(unset)')}" + ) + # Build full service name with hostname full_service_name = get_full_service_name(settings.antenna_service_name) logger.info(f"Running worker as: {full_service_name}") @@ -129,6 +143,162 @@ def _worker_loop(gpu_id: int, pipelines: list[str]): time.sleep(SLEEP_TIME_SECONDS) +def _is_cuda_memory_error(exc: RuntimeError) -> bool: + """Whether an exception is a CUDA memory-allocation failure. + + Covers the caching allocator's ``torch.OutOfMemoryError`` as well as the + plain ``RuntimeError`` that CUDA libraries raise when their own workspace + allocation fails under memory pressure (for example + ``CUBLAS_STATUS_ALLOC_FAILED`` from a linear layer). + """ + if isinstance(exc, torch.OutOfMemoryError): + return True + message = str(exc) + return "out of memory" in message or "ALLOC_FAILED" in message + + +def _predict_in_chunks( + model, + items: torch.Tensor | list[torch.Tensor], +) -> list: + """Run inference over ``items`` in chunks of ``model.batch_size``. + + Caps the GPU memory peak of a forward pass. The DataLoader batch is sized + by ``antenna_api_batch_size``, which is a fetch/transfer concern; the + model's ``batch_size`` (``localization_batch_size`` for the detector) + reflects what fits on the GPU. Without chunking, one forward pass over an + entire fetch batch of full-resolution images can exhaust a GPU that is + shared with other worker processes. + + If a chunk itself hits a CUDA memory-allocation failure (for example + because another process on the same GPU is spiking), the chunk size is + halved and the chunk retried, down to single items. The reduction lasts + only for the current call; the next call starts fresh from + ``model.batch_size``. + + Args: + model: Inference model providing ``batch_size``, ``predict_batch``, + and ``post_process_batch``. + items: A stacked tensor, or a list of tensors. Chunks keep the same + form: tensor slices, or sub-lists (the detector accepts + variable-size image lists). + + Returns: + Concatenated post-processed outputs, one entry per input item, in + input order. + """ + chunk_size = max(1, int(model.batch_size)) + results: list = [] + i = 0 + while i < len(items): + chunk = items[i : i + chunk_size] + try: + output = model.predict_batch(chunk) + except RuntimeError as e: + if not _is_cuda_memory_error(e): + raise + failed_size = len(chunk) + del chunk # Release the chunk itself before clearing the cache + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if chunk_size == 1: + raise + chunk_size = max(1, chunk_size // 2) + logger.warning( + f"CUDA memory allocation failed on a chunk of {failed_size} items; " + f"retrying with chunk size {chunk_size}" + ) + continue + results.extend(model.post_process_batch(output)) + i += len(chunk) + return results + + +def _classify_crops_in_chunks( + model, + detections: list[DetectionResponse], + image_tensors: dict[str, torch.Tensor], +) -> tuple[list, list[int]]: + """Crop, transform, and classify detections in chunks of ``model.batch_size``. + + Owning the whole crop lifecycle here — slice, transform, stack, forward — + keeps at most one chunk of crop tensors alive at a time. Building every + crop up front and only chunking the forward pass would allocate all crops + (potentially hundreds per image) before inference starts, which defeats + the purpose of chunking on a memory-constrained GPU. + + Detections with an invalid (empty) bounding box are skipped with a + warning; the returned index list maps each prediction back to its + detection. + + On a CUDA memory-allocation failure the chunk size is halved and the + chunk retried, down to single crops, matching ``_predict_in_chunks``. + The reduction lasts only for the current call. + + Args: + model: Classifier providing ``batch_size``, ``get_transforms``, + ``predict_batch``, and ``post_process_batch``. + detections: Detections whose bounding boxes select the crops. + image_tensors: Source-image tensors on the inference device, keyed by + source image id. + + Returns: + ``(predictions, valid_indices)`` where ``predictions[i]`` is the + post-processed output for ``detections[valid_indices[i]]``. + """ + transforms = model.get_transforms() + valid_indices: list[int] = [] + for idx, dresp in enumerate(detections): + bbox = dresp.bbox + y1, y2 = int(bbox.y1), int(bbox.y2) + x1, x2 = int(bbox.x1), int(bbox.x2) + if y1 >= y2 or x1 >= x2: + logger.warning( + f"Skipping detection {idx} with invalid bbox: " + f"({x1},{y1})->({x2},{y2})" + ) + continue + valid_indices.append(idx) + + chunk_size = max(1, int(model.batch_size)) + predictions: list = [] + pos = 0 + while pos < len(valid_indices): + chunk_indices = valid_indices[pos : pos + chunk_size] + try: + crops = [] + for idx in chunk_indices: + dresp = detections[idx] + image_tensor = image_tensors[dresp.source_image_id] + bbox = dresp.bbox + crop = image_tensor[ + :, int(bbox.y1) : int(bbox.y2), int(bbox.x1) : int(bbox.x2) + ] + crops.append(transforms(crop)) + batch = torch.stack(crops) + del crops + output = model.predict_batch(batch) + del batch + except RuntimeError as e: + if not _is_cuda_memory_error(e): + raise + # Drop any partially built chunk tensors before clearing the cache + crops = batch = None + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if chunk_size == 1: + raise + chunk_size = max(1, chunk_size // 2) + logger.warning( + f"CUDA memory allocation failed on a chunk of {len(chunk_indices)} " + f"crops; retrying with chunk size {chunk_size}" + ) + continue + predictions.extend(model.post_process_batch(output)) + pos += len(chunk_indices) + return predictions, valid_indices + + def _apply_binary_classification( binary_filter: "MothClassifierBinary", detector_results: list[DetectionResponse], @@ -148,54 +318,32 @@ def _apply_binary_classification( """ binary_filter.reset(detector_results) - # Process binary classification crops - binary_crops = [] - binary_valid_indices = [] - binary_transforms = binary_filter.get_transforms() - - for idx, dresp in enumerate(detector_results): - image_tensor = image_tensors[dresp.source_image_id] - bbox = dresp.bbox - y1, y2 = int(bbox.y1), int(bbox.y2) - x1, x2 = int(bbox.x1), int(bbox.x2) - if y1 >= y2 or x1 >= x2: - logger.warning( - f"Skipping binary classification {idx} with invalid bbox: " - f"({x1},{y1})->({x2},{y2})" - ) - continue - crop = image_tensor[:, y1:y2, x1:x2] - crop_transformed = binary_transforms(crop) - binary_crops.append(crop_transformed) - binary_valid_indices.append(idx) + # Crop and classify chunk by chunk so only one chunk of crop tensors is + # on the GPU at a time. + binary_out, binary_valid_indices = _classify_crops_in_chunks( + binary_filter, detector_results, image_tensors + ) moth_detections = [] non_moth_detections = [] - if binary_crops: - batched_binary_crops = torch.stack(binary_crops) - binary_out = binary_filter.predict_batch(batched_binary_crops) - binary_out = binary_filter.post_process_batch(binary_out) - - for crop_i, idx in enumerate(binary_valid_indices): - dresp = detector_results[idx] - detection = binary_filter.update_detection_classification( - seconds_per_item=0, - image_id=dresp.source_image_id, - detection_idx=idx, - predictions=binary_out[crop_i], - ) + for crop_i, idx in enumerate(binary_valid_indices): + dresp = detector_results[idx] + detection = binary_filter.update_detection_classification( + seconds_per_item=0, + image_id=dresp.source_image_id, + detection_idx=idx, + predictions=binary_out[crop_i], + ) - # Separate moth from non-moth detections - for classification in detection.classifications: - if classification.classification == binary_filter.positive_binary_label: - moth_detections.append(detection) - elif ( - classification.classification == binary_filter.negative_binary_label - ): - non_moth_detections.append(detection) - image_detections[detection.source_image_id].append(detection) - break + # Separate moth from non-moth detections + for classification in detection.classifications: + if classification.classification == binary_filter.positive_binary_label: + moth_detections.append(detection) + elif classification.classification == binary_filter.negative_binary_label: + non_moth_detections.append(detection) + image_detections[detection.source_image_id].append(detection) + break return moth_detections, non_moth_detections @@ -211,9 +359,9 @@ def _process_batch( ) -> tuple[int, int, list[AntennaTaskResult], float, float]: """Process a single batch of images through detection and classification. - All large intermediates (image_tensors, crops, batched_crops, image_detections) - are local to this function and freed by Python's reference counting when it - returns, preventing memory accumulation across batches. + All large intermediates (image_tensors, crops, image_detections) are local + to this function and freed by Python's reference counting when it returns, + preventing memory accumulation across batches. Args: batch: Dictionary with images, image_ids, reply_subjects, image_urls, failed_items @@ -248,13 +396,14 @@ def _process_batch( batch_start_time = datetime.datetime.now() - # output is dict of "boxes", "labels", "scores" + # Each output item is the list of bounding boxes kept for one image. + # Chunked so the detector's forward-pass memory peak is bounded by + # localization_batch_size, not by the (larger) API fetch batch. batch_output = [] if len(images) > 0: - batch_output = detector.predict_batch(images) + batch_output = _predict_in_chunks(detector, images) n_items = len(batch_output) - batch_output = list(detector.post_process_batch(batch_output)) # Convert image_ids to list if needed if isinstance(image_ids, (np.ndarray, torch.Tensor)): @@ -292,46 +441,26 @@ def _process_batch( detections_for_terminal_classifier = detector_results detections_to_return = [] - # Run terminal classifier on filtered detections + # Run terminal classifier on filtered detections. Cropping and + # inference are interleaved chunk by chunk so only one chunk of crop + # tensors is on the GPU at a time. classifier.reset(detections_for_terminal_classifier) - classify_transforms = classifier.get_transforms() - # Collect and transform all crops for batched classification - crops = [] - valid_indices = [] n_detections = 0 - for idx, dresp in enumerate(detections_for_terminal_classifier): - image_tensor = image_tensors[dresp.source_image_id] - bbox = dresp.bbox - y1, y2 = int(bbox.y1), int(bbox.y2) - x1, x2 = int(bbox.x1), int(bbox.x2) - if y1 >= y2 or x1 >= x2: - logger.warning( - f"Skipping detection {idx} with invalid bbox: " - f"({x1},{y1})->({x2},{y2})" - ) - continue - crop = image_tensor[:, y1:y2, x1:x2] - crop_transformed = classify_transforms(crop) - crops.append(crop_transformed) - valid_indices.append(idx) - classify_start = datetime.datetime.now() - if crops: - batched_crops = torch.stack(crops) - classifier_out = classifier.predict_batch(batched_crops) - classifier_out = classifier.post_process_batch(classifier_out) - - for crop_i, idx in enumerate(valid_indices): - dresp = detections_for_terminal_classifier[idx] - detection = classifier.update_detection_classification( - seconds_per_item=0, - image_id=dresp.source_image_id, - detection_idx=idx, - predictions=classifier_out[crop_i], - ) - image_detections[dresp.source_image_id].append(detection) - n_detections += 1 + classifier_out, valid_indices = _classify_crops_in_chunks( + classifier, detections_for_terminal_classifier, image_tensors + ) + for crop_i, idx in enumerate(valid_indices): + dresp = detections_for_terminal_classifier[idx] + detection = classifier.update_detection_classification( + seconds_per_item=0, + image_id=dresp.source_image_id, + detection_idx=idx, + predictions=classifier_out[crop_i], + ) + image_detections[dresp.source_image_id].append(detection) + n_detections += 1 classify_time = (datetime.datetime.now() - classify_start).total_seconds() # Count non-moth detections returned from binary filter @@ -396,6 +525,40 @@ def _process_batch( return n_items, n_detections, batch_results, detect_time, classify_time +def _init_job_models( + classifier_class: type, + use_binary_filter: bool, + settings: Settings, +) -> tuple[APIMothClassifier, APIMothDetector, MothClassifierBinary | None]: + """Instantiate the models for one job, sized for chunked GPU inference. + + The GPU batch-size settings are applied here, mirroring the synchronous + API path in ``trapdata/api/api.py``: the detector processes + full-resolution images in chunks of ``localization_batch_size`` and the + classifiers process crops in chunks of ``classification_batch_size`` + (see ``_predict_in_chunks``). + + Returns: + ``(classifier, detector, binary_filter)``; ``binary_filter`` is None + when the pipeline does not use the moth/non-moth filter. + """ + classifier = classifier_class( + source_images=[], + detections=[], + batch_size=settings.classification_batch_size, + ) + detector = APIMothDetector([], batch_size=settings.localization_batch_size) + binary_filter = None + if use_binary_filter: + binary_filter = MothClassifierBinary( + source_images=[], + detections=[], + terminal=False, + batch_size=settings.classification_batch_size, + ) + return classifier, detector, binary_filter + + @torch.no_grad() def _process_job( pipeline: str, @@ -463,20 +626,14 @@ def _process_job( # Defer instantiation of poster, detector and classifiers until we have data if not classifier: - classifier = classifier_class(source_images=[], detections=[]) - detector = APIMothDetector([]) + classifier, detector, binary_filter = _init_job_models( + classifier_class, use_binary_filter, settings + ) result_poster = ResultPoster( max_pending=MAX_PENDING_POSTS, max_post_bytes=settings.antenna_result_post_max_bytes, ) - if use_binary_filter: - binary_filter = MothClassifierBinary( - source_images=[], - detections=[], - terminal=False, - ) - assert detector is not None, "Detector not initialized" assert classifier is not None, "Classifier not initialized" assert result_poster is not None, "ResultPoster not initialized" @@ -509,7 +666,7 @@ def _process_job( ) batch_total, t_total = t_total() logger.info( - f"Batch {i + 1}: {batch_total/max(n_items, 1):.2f}s/image, " + f"Batch {i + 1}: {batch_total / max(n_items, 1):.2f}s/image, " f"Classification time: {cls_time:.2f}s, Detection time: {det_time:.2f}s, " f"Load time: {load_time:.2f}s" ) @@ -542,3 +699,15 @@ def _process_job( finally: if result_poster: result_poster.shutdown() + # Release per-job GPU state: drop the references that keep model + # weights and prefetched batches alive, then return the freed blocks + # to the driver. Without the empty_cache() call, a process idling + # between jobs keeps its peak reserved memory cached, and other + # worker processes sharing the same GPU cannot use it. + if isinstance(batch_source, CUDAPrefetcher): + # The prefetcher's buffered next batch is a full set of images on + # the GPU; clear it explicitly rather than relying on the del. + batch_source.next_batch = None + del classifier, detector, binary_filter, batch_source, loader + if torch.cuda.is_available(): + torch.cuda.empty_cache()