dispatch#

ray.util.tpu.dispatch(fn: Any, *args: Any, topology: str | None = None, accelerator_version: str | None = None, tpu_slice: SlicePlacementGroup | None = None, num_slices: int = 1, chips_per_vm: int | None = None, head_reservation_timeout_s: float | None = 100.0, pg_ready_timeout_s: float | None = None, **kwargs: Any) → List[ray._raylet.ObjectRef][source]#

Run a remote function on every host in a TPU slice.

Dispatches one task per host in the slice, pinning each task to its corresponding placement-group bundle via PlacementGroupSchedulingStrategy. The function blocks until the underlying placement group is scheduled, then returns a list of object references — one per host — that can be passed directly to ray.get.

Resource options (num_cpus=0, resources={"TPU": N}, and scheduling_strategy) are applied automatically via .options() and override any values set in the @ray.remote decorator.

Parameters:

  • fn – A @ray.remote-decorated function to run on every host.

  • *args – Positional arguments broadcast to every task invocation.

  • topology – The TPU topology string (e.g. "4x4", "2x2x2"). Required when tpu_slice is None; ignored otherwise.

  • accelerator_version – The TPU accelerator generation (e.g. "v4", "v6e"). Required when tpu_slice is None; ignored otherwise.

  • tpu_slice – An existing SlicePlacementGroup to schedule onto. When provided, the slice is used directly and dispatch does not create, modify, or tear down any placement groups. When None (default), a new slice is reserved internally and its head placement groups are released once the worker placement group becomes ready.

  • num_slices – Number of TPU slices to reserve. Ignored when tpu_slice is provided. Defaults to 1.

  • chips_per_vm – Optional override for the number of chips per VM. Ignored when tpu_slice is provided.

  • head_reservation_timeout_s – Seconds to wait for each head placement group during slice reservation. Ignored when tpu_slice is provided. Defaults to DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S.

  • pg_ready_timeout_s – Seconds to wait for the worker placement group to become ready after reservation. Pass None to wait indefinitely (default).

  • **kwargs – Keyword arguments broadcast to every task invocation.

Returns:

One object reference per host in the slice.

Pass the list to ray.get to retrieve results.

Return type:

List[ray.ObjectRef]

Raises:

  • TypeError – If fn is not a @ray.remote-decorated function (i.e. it has no .options() method).

  • ValueError – If tpu_slice is None and either topology or accelerator_version is not provided.

  • TimeoutError – If the placement group does not become ready within pg_ready_timeout_s seconds. When the slice was created internally, it is shut down before the error is raised to avoid leaking resources.

Examples:

import ray
from ray.util.tpu import dispatch, slice_placement_group

@ray.remote
def my_tpu_task():
    import jax
    return jax.device_count()

# One-shot: reserve a v6e 4x4 slice, run on every host, then
# release automatically when the driver exits.
results = ray.get(
    dispatch(my_tpu_task, topology="4x4", accelerator_version="v6e")
)

# Reuse an existing slice across multiple calls.
slice_handle = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(slice_handle.placement_group.ready())

results1 = ray.get(dispatch(my_tpu_task, tpu_slice=slice_handle))
results2 = ray.get(dispatch(my_tpu_task, tpu_slice=slice_handle))
slice_handle.shutdown()

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

On this page