ray.util.tpu.get_tpu_num_slices_for_workers#

ray.util.tpu.get_tpu_num_slices_for_workers(topology: str, accelerator_type: str, num_workers: int, resources_per_worker: Dict[str, float] | None = None) int[source]#

Calculates the number of slices needed to accommodate the specified number of workers.

Parameters:
  • topology – The TPU topology string.

  • accelerator_type – The accelerator type string.

  • num_workers – The desired number of workers.

  • resources_per_worker – Optional dict of resources per worker.

Returns:

The number of slices required. Returns 1 if inputs are invalid or incomplete.

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