Lazy sorting added for replay buffer

[chirayuharyan]

• I've read the .github/CONTRIBUTING.md file
• My code follows the typing guidelines
• I've added appropriate tests
• I've run pre-commit hooks locally

Lazy sorting for replay buffer

Summary

Adds an opt-in lazy_sort mode to ReplayBuffer that defers concatenation and sorting until the buffer reaches 2x capacity, significantly reducing the cost of _local_add on buffer manager ranks.

Problem

On buffer manager ranks, every incoming batch triggered an extend() (which copies the entire buffer via torch.cat) followed by an argsort + truncation. With 32 training agents each submitting ~100 batches, this resulted in ~3,200 extend-sort-truncate cycles — dominated by the quadratic copy cost of repeatedly concatenating onto a growing tensor.

Changes

replay_buffer.py

• lazy_sort parameter — new bool flag on ReplayBuffer.__init__ (default False, backward compatible). When enabled, incoming batches are accumulated in _pending_batches list instead of being immediately concatenated.
• _flush_pending() — merges all pending batches into a single combined batch first, then extends training_container once — avoids the quadratic cost of extending one-by-one.
• _sort_and_truncate() — extracted from _local_add into a helper; argsorts by log_rewards (if prioritized) and truncates to capacity. Both paths (lazy and eager) share this method.
• _local_add() — branches on self.lazy_sort: lazy path accumulates and flushes at 2x capacity; eager path preserves the original extend-sort-truncate-every-call behavior.
• __len__ — includes _pending_len so length queries reflect uncommitted items.
• sample() — samples from committed container only (does not flush pending). Guard updated to handle empty training_container.
• save() — calls _flush_pending() before serialization.
• NormBasedDiversePrioritizedReplayBuffer — lazy_sort parameter threaded through to super().__init__. _local_add flushes pending batches when the buffer is full (diversity filtering requires a consistent view), but delegates to the base lazy path while filling up.

replay_buffer_manager.py

• Passes lazy_sort=True and timing=self.timing to both ReplayBuffer and NormBasedDiversePrioritizedReplayBuffer constructors.

Performance

Benchmarked on 32 training agents, 1 buffer manager (Intel EMR (64 cores) x 2, single node), 100 iterations:

Phase	Before (eager)	After (lazy sort)	Reduction
replay_add (total)	173.6s (3200 calls)	97.6s (3200 calls)	44%
local_add/extend	53.2s (3200 calls)	61.7s (399 calls)	per-call larger, but fewer calls
local_add/sort	111.6s (3200 calls)	33.7s (399 calls)	70%
local_add/truncate	0.3s (3200 calls)	0.06s (399 calls)	—

Backward Compatibility

lazy_sort defaults to False — existing code is unaffected. The eager path is unchanged aside from being wrapped in timers and delegating to the shared _sort_and_truncate helper.

[codecov]

Codecov Report

❌ Patch coverage is 55.55556% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.14%. Comparing base (0408921) to head (c759a80).
⚠️ Report is 7 commits behind head on master.

Files with missing lines	Patch %	Lines
src/gfn/containers/replay_buffer.py	55.55%	18 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #523      +/-   ##
==========================================
- Coverage   73.25%   73.14%   -0.12%     
==========================================
  Files          59       59              
  Lines       10082    10115      +33     
  Branches     1402     1408       +6     
==========================================
+ Hits         7386     7399      +13     
- Misses       2237     2255      +18     
- Partials      459      461       +2     

Files with missing lines	Coverage Δ
src/gfn/containers/replay_buffer_manager.py	0.00% <ø> (ø)
src/gfn/containers/replay_buffer.py	53.63% <55.55%> (-1.84%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[josephdviviano]

lgtm, great work, thank you