ray.util.tpu.get_tpu_worker_resources#

ray.util.tpu.get_tpu_worker_resources(topology: str, accelerator_type: str, resources_per_unit: Dict[str, float] | None = None, num_slices: int = 1) → Tuple[int, Dict[str, float]][source]#

Calculates the number of workers and the resources required for each worker to run based on a TPU topology.

Parameters:

  • topology – The TPU topology string.

  • accelerator_type – The accelerator string.

  • resources_per_unit – Optional manual override for resources per unit. If unspecified, the number of TPU chips in a host is assumed.

  • num_slices – The number of TPU slices.

Returns:

  • num_workers: Total workers required.

  • unit_resources: The resource dictionary for a single worker.

Return type:

A tuple containing

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