ray.util.tpu.get_tpu_coordinator_env_vars#

ray.util.tpu.get_tpu_coordinator_env_vars(coordinator_address: str, num_slices: int, slice_id: int, coordinator_port: str = '8081') → Dict[str, str][source]#

Returns the environment variables required for JAX multi-slice coordination.

Parameters:

  • coordinator_address – The IP address or hostname of the coordinator.

  • num_slices – The total number of slices in the cluster.

  • slice_id – The index of the current slice.

  • coordinator_port – The port the coordinator is listening on.

Returns:

A dictionary mapping environment variable names to their values.

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