ray.util.tpu.slice_placement_group#

ray.util.tpu.slice_placement_group(topology: str, accelerator_version: str, resources_per_bundle: Dict[str, float] | None = None, num_slices: int = 1, chips_per_vm: int | None = None, **kwargs) → SlicePlacementGroup[source]#

Asynchronously creates a PlacementGroup for a TPU slice.

A slice placement group reserves num_slices TPU slice(s) and creates a placement group for scheduling tasks or actors.

Parameters:

  • topology – The desired TPU pod topology (e.g. “4x4”, “2x8”).

  • accelerator_version – The TPU accelerator generation, (e.g. “v4”, “v5p”, “v6e”).

  • resources_per_bundle – Specify the number of resources to reserve per bundle. When unspecified, SlicePlacementGroup defaults to reserving 1 bundle per TPU host in a topology, with the bundle resources set to the number of TPU in a host. Ex: Specifying {“TPU”: 1} for a 4x4 topology would result in 16 bundles, each with 1 TPU. If resources_per_bundle=None for the same topology, there would be 4 bundles with 4 TPU each.

  • num_slices – The number of tpu slices within the placement group.

  • chips_per_vm – An optional override for the number of chips per TPU VM. Useful for ambiguous topologies like v6e 2x4 which have 1 host, but can be provisioned as either 1 VM (8 chips) or 2 VMs (4 chips each).

  • **kwargs – Additional arguments for the placement group, such as ‘name’, ‘lifetime’, or ‘strategy’.

Returns:

The handle for the created SlicePlacementGroup.

PublicAPI (alpha): This API is in alpha and may change before becoming stable.