[Core][TPU] Add per_slice_pgs paramater to SlicePlacementGroup

[ryanaoleary]

Description

Currently SlicePlacementGroup accepts a num_slices parameter - atomically reserving that many slices- and returns a single PG handle spanning all the Ray nodes across those num_slices slices. This works well for multi-slice training or inference where the MEGASCALE backend it utilized to perform collectives across all the TPU hosts.

However, for other use-cases like RL or disaggregated serving it'd be useful for users to be able to reserve multiple TPU slices, but utilize the individual slices at their discretion for different workloads. This PR introduced a per_slice_pgs that when set exposes a per_slice_placement_groups property. per_slice_placement_groups is a list of placement group handles where each PG reserves 1 TPU slice.

Related issues

#64071

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[ryanaoleary]

cc: @spencer-p

[gemini-code-assist]

Code Review

This pull request introduces the per_slice_pgs option to slice_placement_group, allowing users to create individual placement groups per TPU slice instead of a single placement group across all slices. It also updates the cleanup logic and adds unit tests for this new feature. The review comments point out two critical issues: first, a potential resource leak if an exception occurs during slice reservation because partially created placement groups are not yet tracked in self._managed_pgs; second, a potential IndexError when accessing placement_group after shutdown() has cleared self._managed_pgs.

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.

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[ryanaoleary]

not stale, just needs review

[andrewsykim]

Very minor nits on naming, feel free to ignore

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit b2fd409. Configure here.}

[ryanaoleary]

cc: @andrewsykim some changes made in 660450d after approval to fix a cursor comment, otherwise I think this is good to merge