Clean up validation dataloader workers - #26014
Conversation
|
All Contributors have signed the CLA. ✅ |
|
👋 Hello @Marchematics, thank you for submitting a
For more guidance, please refer to our Contributing Guide. Don't hesitate to leave a comment if you have any questions. Thank you for contributing to Ultralytics! 🚀 |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Pull request overview
This PR aims to prevent persistent validation DataLoader worker processes from lingering across validation passes by explicitly shutting them down after each validation loop, while ensuring the custom InfiniteDataLoader can be safely iterated again after shutdown.
Changes:
- Wrap validation iteration in a
try/finallyand calldataloader.close()to release worker processes after each pass. - Make
InfiniteDataLoader.__iter__()recreate its internal iterator when it has been cleared byclose(). - Add a test asserting
close()shuts down workers and that the loader remains reusable afterward.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| ultralytics/engine/validator.py | Adds try/finally cleanup to close the validation dataloader after each pass. |
| ultralytics/data/build.py | Makes InfiniteDataLoader restartable by recreating the iterator on-demand and nulling it on close(). |
| tests/test_python.py | Adds coverage to validate worker shutdown and restart behavior for the dataloader. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
UltralyticsAssistant
left a comment
There was a problem hiding this comment.
🔍 PR Review
Made with ❤️ by Ultralytics Actions
Reviewed the dataloader lifecycle, validator reuse, training validation, Comet integration, and depth calibration consumers. The iterator restart logic is sound, but cleanup currently happens too early for post-validation dataloader consumers and does not cover setup failures, leaving worker-lifecycle regressions in those paths.
💬 Posted 2 inline comments
- 💡 MEDIUM
ultralytics/engine/validator.py:275This closes the loader before later consumers run. For example, when Comet image prediction logging is enabled,on_fit_epoch_endsubsequently iteratestrainer.validator.dataloader;InfiniteDataLoader.__iter__then creates a new worker pool becauseiteratorisNone, leaving those workers alive until the next validation or final teardown and causing worker shutdown/startup churn every epoch. Cleanup should occur after these post-validation consumers, or each consumer that performs an ex… - 📝 LOW
ultralytics/engine/validator.py:244Thetrystarts only afteron_val_start,TQDM(...), andinit_metrics(...), butInfiniteDataLoadercreates its worker iterator during construction. If any setup step raises—for example, a useron_val_startcallback—thefinallyis skipped and the validation workers remain alive. Put thetry/finallyaround the whole post-dataloader validation setup and pass so failed validations are cleaned up as well.
|
I have read the CLA Document and I sign the CLA |
|
Benchmarked and cold-reviewed exact head COCO8 training benchmarkEnvironment: NVIDIA RTX PRO 6000 Blackwell Server Edition, Python 3.12.3, PyTorch 2.11.0+cu128. I ran three measured repetitions per commit after a warm-up, alternating commit order. The matched training configuration was:
The PR increased end-to-end training time by 25.3% and reduced throughput by 20.2%. Instrumentation also confirmed the mechanism: base retained the same two-worker iterator across all epochs, while the PR set it to The high-resolution unified-memory pressure reported in #25827 is a real workload constraint, but unconditional teardown moves that cost to every multi-batch validation on every epoch and regresses the small-dataset workloads where worker reuse matters most. Review findingEven aside from the performance regression, cleanup is not complete on the training path. Validation returns at RecommendationDo not merge this PR; close it as implemented and keep persistent validation workers as the default. A future solution to the reported memory-constrained workload should solve that constraint without unconditionally discarding reusable workers and prefetched batches for all training jobs, and it should include matched end-to-end speed and peak-memory evidence across both the original high-resolution case and small datasets such as COCO8. |
|
Follow-up: I benchmarked the resource-salvaging middle ground on exact base The prototype left the training COCO8 speed, 20 epochs,
|
| Validation loader | Run 1 | Run 2 | Run 3 | Mean ± SD | Throughput | Time delta |
|---|---|---|---|---|---|---|
| Current infinite, prefetch 4 | 10.246 s | 10.474 s | 10.165 s | 10.295 ± 0.160 s | 1.943 epochs/s | — |
| Finite persistent, prefetch 2 | 10.842 s | 10.397 s | 10.524 s | 10.588 ± 0.229 s | 1.889 epochs/s | +2.8% |
| Finite persistent, prefetch 1 | 11.217 s | 10.579 s | 10.677 s | 10.824 ± 0.344 s | 1.848 epochs/s | +5.1% |
For comparison, the PR's close/recreate approach was +25.3% in the preceding matched benchmark. Excluding the timing column, every per-epoch training/validation metric was identical across these three variants.
Settled resources after validation, imgsz=1920
I used an isolated 2 GiB container shared-memory allocation and measured container-cgroup RAM after each epoch validation and the final validation. COCO8 produces two validation workers here:
| Validation loader | Outstanding val tasks | Queued val batches | Mean charged RAM | Shared memory | CUDA reserved |
|---|---|---|---|---|---|
| Current infinite, prefetch 4 | 8 | 8 | 2890 MiB | 533 MiB | 3478–3482 MiB |
| Finite persistent, prefetch 2 | 0 | 0 | 2681 MiB | 309 MiB | 3478–3482 MiB |
| Finite persistent, prefetch 1 | 0 | 0 | 2679 MiB | 309 MiB | 3478–3482 MiB |
Both finite variants released 224 MiB of shared memory and about 210 MiB of total charged RAM between validations in this two-worker COCO8 case, while retaining the workers. They did not change CUDA-reserved memory, confirming that the salvageable resource here is host/shared batch storage rather than model VRAM. The original 20-worker, large-batch workload should have a much larger absolute difference, but needs direct measurement on that unified-memory host rather than extrapolation.
Updated recommendation
The current PR should still not merge. If we pursue #25827, finite persistent validation workers with prefetch_factor=2 are the strongest replacement direction:
- retains almost all small-dataset speed (2.8% cost versus 25.3% for reconstruction),
- empties validation queues between passes,
- halves the current in-validation prefetch depth from 4 to 2,
- and saves the same between-validation RAM as prefetch 1 in this test.
Prefetch 1 adds another speed penalty without additional settled RAM savings, so it only merits consideration if direct measurement shows peak memory during validation itself remains the limiting problem.
|
Added the requested finite-persistent / COCO8 speedSame RTX PRO 6000, Python 3.12.3, PyTorch 2.11.0+cu128,
This is effectively tied with finite-persistent/prefetch-2's previous +2.8% result and remains far better than close/recreate's +25.3%. COCO8 has only two finite validation batches under this configuration, so prefetch 2 already fills the available finite work; prefetch 4 cannot improve this particular speed case. Excluding the timing column, all per-epoch training and validation metrics matched the control exactly. Settled resources after validationMatched
Finite-persistent/prefetch-4 reclaimed about 195 MiB total charged RAM and 192–224 MiB shared memory between validations while retaining both workers. CUDA-reserved memory was unchanged. RecommendationFinite persistent with The current close/recreate implementation should still not merge. A replacement should keep training infinite, make validation finite with persistent workers, retain prefetch 4 initially, and verify peak/settled memory on the reported unified-memory host before merge. |
|
Thanks for the detailed benchmarks and the finite-persistent direction. Implemented finite persistent validation workers with |
The endless _RepeatSampler restarted the sampler while the previous epoch was still prefetching, so DistributedSampler.set_epoch landed one epoch late: epochs 0 and 1 trained on the identical shuffle and every later epoch used the previous epoch's seed. It also kept prefetching the next validation pass between epochs, retaining worker memory after validation completed. A plain DataLoader with persistent workers reuses the worker pool across epochs, only prefetches an epoch once __iter__ runs after set_epoch, and drains between passes. close_dataloader() shuts the pool down at the end of training, after standalone validation, and at close_mosaic so fresh workers pick up the new transforms.
…shuffle in _RepeatSampler Per the benchmarks on this PR: the infinite training loader stays because cross-epoch prefetch keeps small-dataset training fast, and validation uses a finite DataLoader with persistent workers and prefetch 4 so it drains between passes without rebuilding workers. _RepeatSampler now advances a DistributedSampler's epoch when it restarts, because the next pass is drawn before the trainer's set_epoch runs; without this DDP epochs 0 and 1 train on the identical order and every later epoch lags one behind. One close_dataloader() replaces InfiniteDataLoader.close/reset for both loader kinds.
Standalone validation drops its validator after the pass, and the persistent iterator's own __del__ shuts the workers down, so no explicit close is needed there.
⚡ Actions TriggerMade with ❤️ by Ultralytics Actions GitHub Actions below triggered via workflow dispatch for this PR at 2026-09-03 11:48:37 UTC with
|
Every caller already passes shuffle=True only for training loaders, so the infinite loader keys off it; classify now passes shuffle explicitly so its validation loaders are sequential and finite like the other tasks.
⚡ Actions TriggerMade with ❤️ by Ultralytics Actions GitHub Actions below triggered via workflow dispatch for this PR at 2026-09-03 15:34:53 UTC with
|
Signed-off-by: Glenn Jocher <[email protected]>
GPU validation: 100 epochs, main vs this PRNVIDIA GH200 480GB, Python 3.14.7, torch 2.14.0+cu130, fresh uv venv.
mAP is identical to every digit in all eight runs: on a single GPU the training loader's ordering is unchanged and the validation loader cannot affect metrics. The DDP reshuffle fix is not observable on a one-GPU host. Speed: coco8 at batch 1 is 5.7% faster; coco128 at batch 16 is 3.6% slower. Isolating the coco128 cost with 100 back-to-back validation passes on coco128 (4 batches, 4 workers, no training):
The finite loader pays about 22 ms per pass to drain and re-prime its workers (2.2 s over 100 epochs) plus about 1.1 s once for the lazy worker spawn at first iteration, which accounts for most of the coco128 gap. This is the trade-off measured earlier on this PR for releasing queued validation batches between epochs. The coco8 speedup is consistent with the drained validation queue no longer competing for CPU with a batch-1 training loop. Side observation: at interpreter exit main's infinite validation loader emitted a worker |
# Conflicts: # ultralytics/data/build.py
|
Merged Where this leaves the trade-off, using the profile from #26057 and the diagnostic above on coco128 (4 validation workers, 100 epochs):
What this PR still adds on top of main: zero validation batches retained during training instead of |
# Conflicts: # ultralytics/models/yolo/depth/calibrate.py
|
Fresh comparison on ultra5: this PR drains the validation queue, but was slower on both small training workloads tested. Compared exact PR head Real
Every per-epoch validation metric matched between variants in both repetitions, including across mosaic closure. All runs finished with zero remaining child processes. Validation worker PIDs remained persistent between epochs in both implementations. Main retained up to 4 outstanding validation tasks for COCO8 and 12 for COCO128; the PR retained zero. A separate 1920px COCO8 validation-loader check used the actual detection trainer's loader, batch 2, two workers, real model inference, and five complete passes. At the settled fifth pass:
That is approximately 64 MiB less retained shared memory with the PR. Prediction hashes were identical across every pass and both versions. This proves the queue-retention mechanism on this small workload; it does not reproduce the original GB10/1920px/batch16/workers15 memory-pressure case or justify extrapolating these timings to a large dataset. Recommendation: no default loader change on these results alone, and no claim of a speed improvement. The memory benefit is real, but this host had no memory pressure and the measured training workloads became slower. Keep the existing persistent-worker behavior; do not revive unconditional teardown/recreation. The smallest remaining decision input is the reported memory-constrained configuration, comparing peak/settled process and container memory, epoch wall time, and accuracy against today's prefetch-2 main. This single-GPU run does not validate the PR's separate DDP reshuffle behavior. |
Release validation dataloader workers' queued batches between passes, following the benchmark direction settled on this PR: training keeps the infinite loader, validation becomes a finite loader with persistent workers and
prefetch_factor=4.Fixes #25827.
What changes
build_dataloaderreturnsInfiniteDataLoaderfor shuffled (training) loaders and a plaintorch.utils.data.DataLoader(persistent_workers=True)for unshuffled (validation) ones. Every caller already distinguishes the two throughshuffle; classify now passes it explicitly, which also makes its validation loaders sequential like the other tasks. The finite loader reuses its workers but exhausts its sampler each pass, so no next-validation batches sit queued while training resumes (the retention mechanism in Validation retains InfiniteDataLoader workers and prefetched batches between epochs #25827). The infinite loader is untouched for training, where cross-epoch prefetch is what keeps small datasets fast (+2.9% for finite validation alone versus +25% for close-and-recreate, per the benchmarks above)._RepeatSampler. The endless sampler restarts and draws the next epoch while the current one is still prefetching, before the trainer'sDistributedSampler.set_epoch. Result on DDP: epochs 0 and 1 train on the identical order and every later epoch lags one epoch behind its seed._RepeatSamplernow advances the wrapped sampler's epoch when it restarts. Verified with a 64-sample dataset,DistributedSampler(num_replicas=2, rank=0)and two workers: all four epochs are distinct and match a plainDataLoaderepoch for epoch. Single-GPU runs useRandomSamplerwith a generator and were unaffected.close_dataloader()replacesInfiniteDataLoader.close()andreset()and works on torch's own persistent iterator for both loader kinds: the trainer calls it at the end of training (Fix leaked dataloader workers at end of training (atexitkilled by signal: Terminatedcrash) #25024) and atclose_mosaicso fresh workers pick up the new transforms. The infinite loader creates its iterator lazily on first iteration instead of in__init__, which is what makes the restart a plain re-iteration. Standalonemodel.val()needs nothing: the validator is dropped after the pass and the iterator's own__del__shuts its workers down (verified: no live children aftermodel.val()with two workers). The torch<2.0prefetch_factorgate goes because the kwarg is only passed when there are workers._rewindgoes because its validation loader now restarts on everyfor.Validation
DistributedSampler(num_replicas=2, rank=0)and two workers, four epochs of the infinite loader are all distinct and match a plainDataLoaderepoch for epoch; main gives epochs 0 and 1 the same order.close_dataloader().tests/test_engine.pyand the train/val/dataloader subset oftests/test_python.pypass (45 passed).@ultralytics/run-all: slow matrix torch 1.8 to 2.11, GPU, macOS, Windows, ARM, Raspberry Pi and Docker builds green; the Jetson jobs only failed while sharing the runner's weights directory with a concurrent run and are being re-run alone.Before merge: verify peak and settled memory on the reported GB10 unified-memory host, as required above.