func getNextWorkerID()

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
}