Share RAM image caches across DataLoader workers - #26086

👋 Hello @glenn-jocher, thank you for submitting a ultralytics/ultralytics 🚀 PR! This automated message confirms your contribution was received, and an Ultralytics engineer will assist with the review. To ensure a seamless integration of your work, please review the following checklist:

• ✅ Define a Purpose: Clearly explain the purpose of your fix or feature in your PR description, and link to any relevant issues. Ensure your commit messages are clear, concise, and adhere to the project's conventions.
• ✅ Synchronize with Source: Confirm your PR is synchronized with the ultralytics/ultralytics main branch. If it's behind, update it by clicking the 'Update branch' button or by running git pull and git merge main locally.
• ✅ Ensure CI Checks Pass: Verify all Ultralytics Continuous Integration (CI) checks are passing. If any checks fail, please address the issues.
• ✅ Update Documentation: Update the relevant documentation for any new or modified features.
• ✅ Add Tests: If applicable, include or update tests to cover your changes, and confirm that all tests are passing.
• ✅ Sign the CLA: Please ensure you have signed our Contributor License Agreement if this is your first Ultralytics PR by writing "I have read the CLA Document and I sign the CLA" in a new message.
• ✅ Minimize Changes: Limit your changes to the minimum necessary for your bug fix or feature addition. "It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is." — Bruce Lee

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

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Follow-up design review and handoff

Reviewed head: 0e0b968da243953238d6fd48ef743f6c6d372522. Keep this PR draft; the current RawArray implementation is not ready for LGTM. No implementation changes are made by this comment.

Better candidate to investigate

Replace the existing packed NumPy pixel buffer with a packed CPU Torch uint8 tensor, and use NumPy views when retrieving images. Investigate relying on PyTorch multiprocessing storage reducers to establish sharing only when a dataset is serialized for spawn/forkserver workers, rather than allocating shared backing unconditionally.

The intended benefit is to preserve ordinary RAM/copy-on-write behavior for zero-worker and fork loaders while avoiding a full pixel-cache copy per spawned worker. This is a proposal, not a validated drop-in fix. Verify the reducer behavior and lifetime on supported platforms and the Python/PyTorch support floor before relying on it. Independent distributed-training ranks are outside this sharing scope.

Blockers and design constraints

• RawArray introduces a backing-storage requirement even for paths that previously needed only RAM. A real decoded repository image (2,624,400 bytes) successfully allocated in NumPy under an OS file-size limit of 65,536 bytes, while the proposed RawArray cache raised OSError 27. Shared backing can also fail through mechanisms that a Python OSError handler does not cover.
• Torch shared storage still has shared-memory/platform capacity limits, including constrained container configurations, and may temporarily require both the original and shared buffer. Reusing a standard reducer does not solve capacity by itself.
• Do not treat catching OSError and silently retaining an unshared image list as a complete fix: it restores per-worker duplication and does not cover SIGBUS-style failures.
• The new per-image copy needs justification and measurement. Standard inspected augmentation paths already allocate before in-place changes, so cache corruption in those paths has not been demonstrated. However, true shared storage makes accidental writes visible across workers. Establish the intended mutation contract and validate it before retaining or removing the copy.

Next steps / acceptance evidence

1. Prototype at the existing _ImageCache owner, keeping the diff minimal and leaving worker start-method defaults unchanged. Avoid parallel cache implementations or speculative platform scaffolding.
2. Validate real detection and classification loaders on Linux, macOS, and Windows, with fork where supported, spawn/forkserver, zero workers, and eight workers. Include the supported version floor, repeated loader resets, and worker/storage cleanup.
3. Prove that fork/zero-worker paths do not acquire a new shared-backing requirement, and that spawned workers share one pixel buffer. Measure backing identity/PSS or platform equivalents and parent peak memory; RSS alone cannot distinguish shared mappings from duplication.
4. Exercise constrained shared memory/backing storage in disposable processes. Resolve the capacity behavior before calling this a compatibility-preserving replacement; a hard crash or unexplained memory multiplication is not acceptable.
5. Verify cache contents through real augmentation paths and measure full-loader/training throughput and peak memory against the current implementation. Use a representative cache, not only a tiny smoke dataset.
6. Resume the same reviewer on the implementation delta, then obtain a cold full-diff review of the final live head. If the candidate cannot meet these constraints, retaining the current cache and documenting spawn replication is preferable to claiming an incomplete universal fix.

Related context: PyTorch #13246 describes longstanding worker memory replication; Ultralytics #9824 and merged #24670/#24673 cover earlier cache improvements, while #24364 explored shared tensor storage. Recheck their exact scope rather than assuming they prove this proposal. These are memory/reliability concerns; no verified security advisory is established here.

Claude Code Fable 5.1 (high) recommends investigating the Torch-storage approach, but it has not been implemented or validated and is not yet LGTM.