[Ray TPU Webhook] Pod Informer inconsistency for large RayCluster sizes

[ryanaoleary]

The Ray TPU webhook sometimes fails to assign unique TPU_WORKER_IDs to all workers in a slice, resulting in the Jax initialization timing out. In manual testing this only occurs with large slice sizes (i.e. a v6e-256) around ~20% of the time (2 out of 10 trials) and is not reproduced in unit tests. I've previously tested the webhook with up to 4x4x8 v4 slices and did not observe this issue. The error is consistently fixed by re-creating the RayCluster. Initial debugging yielded the following results:

1. I initially thought the bug was due to a race condition with multiple concurrent mutation requests, but adding a sync.Mutex lock for the Pod mutation yielded no change. Additionally, mutating admission webhooks appear to execute requests sequentially, although I'm unclear whether this behavior is specific to Gatekeeper.
2. Log statements show that the error occurs when an admission request calculates the next TPU_WORKER_ID with a stale Informer cache (the PodLister does not show the updated TPU_WORKER_ID from the previous request) and re-assigns a worker ID.
3. Adding an additional WaitForCacheSync has no effect and checking cache.WaitForCacheSync(stopCh, podInformer.HasSynced) returns true even when the error occurs.
4. Adding a sufficiently long time.Sleep() after the Mutate function solves the issue but adds significant latency.

It appears the bug is a result of a slow update time to the PodInformer cache for a large number of KubeRay worker Pods. I'm currently investigating a better method to ensure the PodInformer retains an up-to-date list of Pods without resorting to calling time.Sleep().

[ryanaoleary]

cc: @andrewsykim