in ray-on-gke/tpu/kuberay-tpu-webhook/main.go [470:494]
func getNextWorkerID(sliceToWorkerIDs map[slice][]int, podSlice slice, namespace string, replicaIndex int) (int, error) {
tpuWorkerID := 0 // defaults to 0 (first Pod in slice)
if len(sliceToWorkerIDs) == 0 || len(sliceToWorkerIDs[podSlice]) == 0 {
return tpuWorkerID, nil
}
sort.Ints(sliceToWorkerIDs[podSlice])
// iterate through existing workers and get the next lowest, unused ID
lastID := 0
for index, workerID := range sliceToWorkerIDs[podSlice] {
// check for incorrect assignment of IDs
if index == 0 {
lastID = workerID
} else if workerID == lastID {
return 0, errors.New("Identical TPU_WORKER_ID assigned to multiple TPU workers in slice")
}
// get the next lowest, valid TPU_WORKER_ID
if workerID != tpuWorkerID {
break
}
lastID = workerID
tpuWorkerID++
}
klog.V(1).InfoS("getNextWorkerID", "RayCluster", namespace+"/"+podSlice.clusterName, "Worker Group", podSlice.groupName, "replicaIndex", replicaIndex, "TPU_WORKER_ID", tpuWorkerID)
return tpuWorkerID, nil
}