[Feature] Add replay readiness waits and counters

[vmoens]

No description provided.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4408

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures

As of commit b874303 with merge base 2d258fe ():

NEW FAILURES - The following jobs have failed:

• Unit-tests on Linux / tests-cpu (3.10, mp) / linux-job (gh)
  test/test_loggers.py::TestTrackioLogger::test_log_histogram[steps1]
• Unit-tests on Linux / tests-cpu (3.13, mp) / linux-job (gh)
  test/test_loggers.py::TestTrackioLogger::test_log_histogram[steps1]
• Unit-tests on Linux / tests-gpu (3.12, 13.0, 2) / linux-job (gh)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]
• Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 2) / linux-job (gh)
  test/test_inference_server.py::TestProcessSlotTransport::test_server_batched_pass_on_cuda[True]

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

Review (COMMENT) — own PR, comment-only.

Read the whole thing carefully, including the concurrency choreography. The core design holds up:

• Notification placement is right: _notify_replay_state_change() is called after releasing _replay_lock/_write_lock in add/extend/empty, so there's no lock-order inversion with the condition variable, and the cancel_event polling caps condition.wait at 0.1s only when cancellation was requested.
• __init__ delegates to share(), so both shared and non-shared buffers always have a usable _readiness_condition (I went looking for the "plain buffer never gets one" bug — it isn't there).
• Pickle discipline is thorough: spawn-inheritance keeps Value/Condition handles, checkpointing snapshots them to plain ints, __setstate__ reconstructs per shared, and load_state_dict/dumps/loads round-trip is pinned by tests.
• The done-vs-terminated-style split between storage_size and sampleable_size plus 64-bit write_count migration (with old-pickle compat via plain-int payload) is exactly what long-running async training needs.
• RemoteTensorDictReplayBuffer waits run on the owning process where RPC-issued extends also land, so server-side waits do get notified — good.

Findings, in descending order of importance:

1. Service-backend / distributed-transport clients silently hang with timeout=None. The docs say Ray and "the fixed-layout distributed transport" don't support blocking readiness, but only RayReplayBuffer/_RayReplayBufferClient actually raise. A _DistributedReplayService client is a ReplayBuffer subclass whose local storage never fills (writes go over the wire), so sample(wait=True)/wait_until_sampleable() on it blocks forever with the default timeout=None — no error, no wakeup, since no local notification ever fires. Either raise NotImplementedError on those clients like Ray does, or implement a control-channel polling loop (the stats channel you just extended has everything needed). Silent infinite hang is the worst of the three options.

2. size semantics diverge between base and ensemble. Base now reports "size" = sampleable_size (docstring: "mirrors len(buffer)"), but EnsembleReplayBuffer.stats()["size"] remains the physical storage count while also exposing a different "sampleable_size". Consumers diffing stats()["size"] against stats()["sampleable_size"] to detect consumption pressure get something meaningful on ReplayBuffer and a no-op on the ensemble. Either make ensemble's "size" mirror len(self) too, or document the divergence on the ensemble side.

3. sample(wait=True) is advisory, not a reservation. With a consuming sampler and two concurrent learners, both can pass wait_until_sampleable and the second still fails inside _sample. That's inherent to check-then-act across processes, but the docstring should say so explicitly — async learners will otherwise treat wait=True as a guarantee.

4. Minor: test_stats_counters_checkpoint_roundtrip's dumps leg extends restored with a zero row before loads — if that's just to trigger storage init, a comment saying so saves the next reader a head-scratch.

[vmoens]

CI triage for b874303ff:

• The Python 3.10/3.13 mp jobs fail only in Trackio logger setup because the installed trackio imports the removed huggingface_hub.hf_api.CommitOperationAdd.
• Both GPU shard-2 jobs fail in TestProcessSlotTransport::test_server_batched_pass_on_cuda[True] with PyTorch's Invalid stream capture status during CUDA graph capture.

These exact failures recur on the independent #4407 head and on the stacked #4418/#4419 heads. Neither failure exercises this PR's replay-readiness/counter changes, so I found no branch-caused CI failure to patch.