[Fix] discard failed and filtered rollout_state

[YanhuiDua]

修复点

1. 将failed / filtered的样本直接丢弃，也就是rollout_state中全部字段清空，并且不再放到replay buffer中去
2. 由于修复点1，那就不能再从replaybuffer中统计leftover的failed和filtered的样本，于是改成在ProduceProgress中进行统计

修复后内存曲线

• 机0:

• 机1:

[YanhuiDua]

@claude review

[claude]

Claude: ## Summary

Discards FAILED and FILTERED rollout groups in ReplayBuffer.put() to prevent unbounded accumulation, and adds free_rollout_state_refs / discard_rollout_state helpers to release Ray object refs. Also fixes reset_rollout_response to handle list[ObjectRef] routed_experts.

ProduceBatchResult impact: leftover_failed and leftover_filtered will now always be 0 (groups discarded before storage). This is the intended fix. However, the in-place mutation bug below can cause FAILED/FILTERED groups to be miscounted as COMPLETED by the producer.

RoutedExperts impact: Positive — fixes a gap where list[ObjectRef] routed_experts were not freed in reset_rollout_response. New free_rollout_state_refs properly handles all ref types.

Issues

Critical

• replay_buffer.py:466-471 — discard_rollout_state mutates items in-place, resetting status to INIT. The caller (put_generated_group at produce_utils.py:187) re-reads get_group_status(group) after put() returns and sees COMPLETED for what was actually a FAILED/FILTERED group. This inflates the producer's completed_sample_count and can cause premature termination.

Warning

• rl_data.py:167-178 — Setting the required field message to None creates an invalid RolloutState that would crash any code reading it.
• Missing regression test — per project standards, bug fix PRs must include a test that reproduces the original issue.

Nit

• ray_utils.py:97-105 — id()-based dedup operates on Python object identity, not Ray ObjectRef identity. Two wrappers for the same Ray object won't be deduplicated.
• free_rollout_state_refs is a public function missing a docstring.

Verdict

REQUEST_CHANGES

[YanhuiDua]

@claude review

[claude]

Claude: Nit — progress: Any bypasses mypy in a strict-checked path.

Both ProduceProgress and DisaggProduceProgress implement consume_discarded(task_name) -> tuple[int, int], but using Any here skips static verification. A lightweight Protocol or explicit union type would catch future signature drift:

Suggested change

	def _merge_discarded_counts(result: ProduceBatchResult, progress: Any, task_name: str) -> None:
	def _merge_discarded_counts(result: ProduceBatchResult, progress: ProduceProgress | DisaggProduceProgress, task_name: str) -> None:

(Requires importing DisaggProduceProgress — alternatively a Protocol with the consume_discarded method.)

[claude]

Claude: Nit — fallback keys "failed_groups" / "discarded_failed_groups" don't appear in any existing checkpoint format.

If these were intermediate names during development that never shipped, the nested .get() chains are dead code. Simplifying to just state.get("failed_samples", {}) makes the contract clearer and avoids future confusion about which key names are "supported."

[claude]

Claude: ## Updated Review (all 3 commits)

The previous review flagged critical issues in commit 1 (in-place mutation corrupting group status, None on required fields, id()-based dedup). All three issues were correctly fixed in commits 2–3:

• In-place mutation: discard now happens in put_generated_group with initial_status captured before mutation; return False immediately.
• Required fields: discard_rollout_state now skips required fields (continue), preserving message intact.
• Dedup: uses ref.hex() (Ray ObjectRef identity) instead of id() (Python wrapper identity).

ProduceBatchResult impact: leftover_failed / leftover_filtered renamed to failed_samples / filtered_samples with updated semantics — now includes buffer leftovers + discarded-before-insertion counts. Buffer leftover portion will drain to 0 after deployment since FAILED/FILTERED groups no longer enter the buffer.

RoutedExperts impact: Positive — free_rollout_state_refs correctly handles all ref ownership shapes (single ObjectRef, list, and nested refs in extra_fields). reset_rollout_response now handles list[ObjectRef] routed_experts.

Remaining Issues

Nit

• produce_utils.py:341 — progress: Any bypasses mypy strict checking.
• disagg_producer.py:185-192 — Fallback keys in load_state_dict reference names that don't exist in any shipped checkpoint format.

Verdict

APPROVE — the memory fix is correct and well-tested. Only cosmetic nits remain.