converge the Group type alias next to Sample

[yueming-yuan]

Follow-up to #2241, which corrected generate_and_rm_group's return annotation to the shape it actually produces. That fix spelled the type out at the call site, making it the third place the same union is written; this converges them and applies the same correction to the legacy producer #2241 did not cover.

What was wrong

Group = list[Sample | list[Sample]] was defined independently in two modules (fully_async_rollout.py, multi_lora/async_rollout.py), and spelled out inline in two more after #2241. Consumers rediscovered the shape on their own — hence the hand-rolled isinstance(group[0], list) unwrapping scattered across first_sample / iter_samples / rollout_data_conversion._first_sample.

sglang_rollout.generate_and_rm_group still declared -> list[Sample] while doing the same asyncio.gather over a generate_and_rm that returns Sample | list[Sample] — the exact bug #2241 fixed on the new path.

Change

One definition in miles/utils/types.py, next to Sample:

# One prompt's finished trajectories. A generate function may expand one trajectory
# into several samples (e.g. multi-agent), so an element is a Sample or a list of them.
Group = list[Sample | list[Sample]]

types.py is a leaf module (stdlib + numpy/torch only) and is already imported by every consumer, so this adds no import edges and cannot cycle. Placing it in inference_rollout_common.py next to the producer was the other candidate, but multi_lora consumes the legacy sglang_rollout producer and would then depend on the new inference path for a type alone.

Both producers now declare -> Group; the four consumers import it instead of redefining or respelling it.

Annotations only — no runtime statement changes.

Notes

• Stacked on fix: correct generated group type annotations #2241 (branch jiajun/fix-generated-group-type, plus a merge of current main). Until that merges, the diff here shows its commit too; it rebases away cleanly afterwards.
• Not executed: the devbox this repo's fast tests run in expired. Verified with py_compile and pre-commit run --all-files only. Please run tests/fast/rollout/ before merging.

🤖 Generated with Claude Code