subslice_placement_group#

ray.util.tpu.subslice_placement_group(subslice_topology: str, accelerator_version: str, chips_per_vm: int | None = None, resources_per_bundle: Dict[str, float] | None = None, strategy: str = 'STRICT_SPREAD', name: str = '', lifetime: str | None = None, head_reservation_timeout_s: float | None = 100.0) → SubslicePlacementGroup[source]#

Asynchronously creates a PlacementGroup for a TPU subslice.

A subslice placement group reserves a contiguous subset of workers within a larger TPU slice, enabling multiple workloads to share a physical slice while maintaining ICI topology alignment.

On the first call for a given topology this function temporarily reserves a full parent slice to discover the physical chip layout, computes subslice labels, and releases unused workers. Subsequent calls reuse the cached data.

Parameters:

  • subslice_topology – Desired subslice topology (e.g. "2x4").

  • accelerator_version – TPU accelerator generation (e.g. "v6e").

  • chips_per_vm – Optional override for chips per VM. Useful for ambiguous topologies like v6e 2x4 which can be 1 VM (8 chips) or 2 VMs (4 chips each).

  • resources_per_bundle – Per-bundle resources. Defaults to {"CPU": 1, "TPU": chips_per_vm}.

  • strategy – Placement group strategy (default "STRICT_SPREAD").

  • name – Optional placement group name.

  • lifetime – Placement group lifetime (None or "detached").

  • head_reservation_timeout_s – Maximum seconds to wait for TPU head placement groups. Defaults to DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S.

Returns:

A SubslicePlacementGroup handle.

Raises:

  • ValueError – If the subslice topology is invalid for the accelerator, or if no suitable parent topology is found in the cluster.

  • RuntimeError – If all slices are occupied, or if libtpu is missing.

Examples:

import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.tpu import subslice_placement_group

sg = subslice_placement_group(
    subslice_topology="2x4",
    accelerator_version="v6e",
)

@ray.remote(num_cpus=0, resources={"TPU": 4})
def train(world, rank):
    ...

tasks = [
    train.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=sg.placement_group,
        )
    ).remote(world=sg.num_hosts, rank=i)
    for i in range(sg.num_hosts)
]

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