[Feature Enhancement] Set ordered replica index label to support multi-slice

[ryanaoleary]

Why are these changes needed?

This is a feature enhancement to the RayMultiHostIndexing feature that adds replica id and host-index labels to Pods created by KubeRay that was implemented in this PR: #3998.

To support the multi-host use-case for TPUs and GPUs, we already set a replica id label (a unique string based on the group name) and a host-index (a unique int from 0-NumOfHosts-1) when this feature is enabled and NumOfHosts > 1. These labels add value to users through observability when running SPMD workloads, and through atomic creation and deletion/re-creation of multi-host groups.

However, these labels do not support the use-case for multi-slice workloads where it is important to know the ordered index of the replica within the multi-slice set. Frameworks like JAX require the slice ID, which is an int between 0 and the # of slices, to be set to configure multi-slice workloads (source).

To solve this issue, this PR adds a label for the ordered replica index, an int value between 0 and replicas-1 for each worker group when this feature is enabled. This label will greatly simplify the process of setting environment variables like MEGASCALE_SLICE_ID for multi-slice workloads that use JAX. We can then check these KubeRay labels from the KubeRay TPU webhook when injecting environment vars. Before the TPU webhook change, these labels are still useful because users can pass them to the Pod environment using downward API.

Related issue number

#3902

Checks

• I've made sure the tests are passing.
• Testing Strategy
  • Unit tests
  • Manual tests
  • This PR is not tested :(

[ryanaoleary]

cc: @andrewsykim this is the change we discussed offline

[andrewsykim]

@ryanaoleary if we want to include this in v1.5, we need to ensure that the code path is unchanged when the feature gate is disabled, it seems we are modifying some code paths even when the feature gate is disabled, could we guard all the changes with the feature gate by end of day tomorrow?

[ryanaoleary]

[d7c48a2](/ray-project/kuberay/pull/4163/commits/d7c48a2412a20e2870e2040d28b4e64109832bb7)

Should be fixed with d7c48a2 and 784993b, all the new logic is behind a feature gate and separated into a different code path than the regular reconcilePods. The only change not explicitly behind a feature gate is the change to the function signature of DefaultWorkerPodTemplate, but the replica and host index args are ignored unless the feature is enabled.

[andrewsykim]

Reviewed and confirmed the behavior when feature gate is disabled is unchanged. Given this has a feature gate I think it's OK to include in v1.5. I will also do some manual testing after the next v1.5 release candidate