fix(metrics): move pass@k from trainer to rollout manager

[HJSang]

Summary

• Move pass@k computation from the Megatron trainer (log_passrate in log_utils.py) to the rollout manager (log_rollout_data in metrics.py)
• The trainer-side log_passrate assumed raw_reward had contiguous groups of n_samples_per_prompt, but custom converters can filter individual samples from within groups, breaking alignment and crashing the reshape
• The rollout side runs before convert_samples_to_train_data, so all samples are present and group boundaries are intact

Motivation

Custom convert_samples_to_train_data implementations (e.g. for agent RL) filter out individual failed/zero-reward samples per-prompt to save forward-pass compute. This shrinks the batch below rollout_batch_size * n_samples_per_prompt and can leave partial groups (e.g., 78 samples with group_size=8). The trainer-side log_passrate then crashes on:

assert len(flat_rewards) == num_groups * group_size  # 78 != 16*8

Moving pass@k to the rollout side avoids the issue entirely — the data is complete and aligned at that point. The eval path (log_eval_rollout_data) already computed pass@k on the rollout side and is unchanged.

Changes

File	Change
miles/ray/rollout/metrics.py	Add _compute_passrate_from_samples, call it from log_rollout_data
miles/backends/training_utils/log_utils.py	Remove log_passrate function and call site
tests/fast/ray/rollout/test_metrics.py	Add TestComputePassrateFromSamples

Metric namespace change

Before: passrate/pass@k (logged via gather_log_data("passrate", ...) on trainer)
After: rollout/pass@k (logged inline with other rollout metrics)

The log_passrate config flag still controls whether the metric is computed.

Test plan

• Unit tests for _compute_passrate_from_samples (perfect/zero/mixed/non-divisible)
• Verify W&B metrics appear under rollout/pass@k on a real run
• Confirm no crash with custom converters that filter samples

[gemini-code-assist]

Code Review

This pull request moves the pass@k metrics calculation from the trainer side (in log_utils.py) to the rollout side (in metrics.py) and adds corresponding unit tests. The feedback highlights a discrepancy between the docstring and the actual implementation of _compute_passrate_from_samples: the docstring claims the function groups samples by group_index to robustly handle downstream filtering, but the code actually uses a flat list and skips the entire calculation if the sample count is not divisible by the group size. The reviewer suggests implementing the grouping logic as described and updating the unit tests to match.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The docstring states that this function groups samples by their group_index field and is robust to downstream filtering by excluding incomplete groups. However, the current implementation simply takes a flat list of rewards and skips the entire pass@k calculation if any sample is missing (due to the divisibility check).

Grouping by group_index and only processing completed groups of size group_size makes the calculation robust to missing/filtered samples without skipping the entire batch.

Suggested change

	flat_rewards = [sample.get_reward_value(args) for sample in all_samples]
	if len(flat_rewards) % group_size != 0:
	logger.warning(
	f"Skipping rollout pass@k: {len(flat_rewards)} samples is not divisible by "
	f"n_samples_per_prompt={group_size}."
	)
	return {}
	return compute_pass_rate(
	flat_rewards=flat_rewards,
	group_size=group_size,
	)
	groups = group_by(all_samples, lambda s: s.group_index)
	completed_groups = [g for g in groups.values() if len(g) == group_size]
	if not completed_groups:
	return {}
	flat_rewards = [
	sample.get_reward_value(args)
	for group in completed_groups
	for sample in group
	]
	return compute_pass_rate(
	flat_rewards=flat_rewards,
	group_size=group_size,
	)

[gemini-code-assist]

Update the test to assert that pass@k is computed for the completed groups instead of skipping the entire batch when the total number of samples is not divisible by group_size.

Suggested change

	def test_skips_when_not_divisible(self):
	args = make_args(n_samples_per_prompt=4, reward_key=None)
	samples = make_samples_grouped(2, 4, rewards=[1.0] * 8)
	samples = samples[:7]
	out = _compute_passrate_from_samples(args, samples)
	assert out == {}
	def test_handles_incomplete_groups_by_skipping_them(self):
	args = make_args(n_samples_per_prompt=4, reward_key=None)
	samples = make_samples_grouped(2, 4, rewards=[1.0] * 8)
	samples = samples[:7]
	out = _compute_passrate_from_samples(args, samples)
	assert out["pass@1"] == pytest.approx(1.0)
	assert out["pass@2"] == pytest.approx(1.0)
	assert out["pass@4"] == pytest.approx(1.0)