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, subslice_index: int | None = None) → 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 (str) – Desired subslice topology (e.g. "2x4").

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

  • chips_per_vm (int | None) – 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 (Dict[str, float] | None) – Per-bundle resources. Defaults to {"CPU": 1, "TPU": chips_per_vm}.

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

  • name (str) – Optional placement group name.

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

  • head_reservation_timeout_s (float | None) – Maximum seconds to wait for TPU head placement groups. Defaults to DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S.

  • subslice_index (int | None) – Optional index of the subslice to select. If specified, only the subslice at this index within the physical slice will be considered. If that subslice is busy, the request will fail even if other subslices are idle.

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.

Return type:

SubslicePlacementGroup

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.