Repository: vllm-project/aibrix
File: pkg/plugins/gateway/algorithms/pd_disaggregation.go (Go)

Code before:
// Pods without PodGroupIndex label are also included for backward compatibility.
func (r *pdRouter) filterPrefillDecodePods(routingCtx *types.RoutingContext, readyPods []*v1.Pod) (*v1.Pod, *v1.Pod, error) {
	prefillPods, decodePods := []*v1.Pod{}, []*v1.Pod{}

	var promptLength int
	if aibrixPromptLengthBucketing {
		promptLength = r.getPromptLength(routingCtx)
		klog.V(4).InfoS("prompt length based filtering enabled", "request_id", routingCtx.RequestID, "prompt_length", promptLength)
	}

	for _, pod := range readyPods {
		if _, ok := pod.Labels[PDRoleSetIdentifier]; !ok {
			continue
		}
		if _, ok := pod.Labels[PDRoleIdentifier]; !ok {
			continue
		}

		// For multi-node scenarios, only select pods from node_rank=0 (PodGroupIndex=0)
		// which have the HTTP server running
		if !isPodWithHTTPServer(pod) {
			continue
		}

		if aibrixPromptLengthBucketing {
			if !r.isPodSuitableForPromptLength(pod, promptLength) {
				continue
			}
		}

		switch pod.Labels[PDRoleIdentifier] {
		case "prefill":
			prefillPods = append(prefillPods, pod)
		case "decode":
			decodePods = append(decodePods, pod)
		}
	}

	if len(prefillPods) == 0 || len(decodePods) == 0 {
		return nil, nil, fmt.Errorf("prefill or decode pods are not ready: prefill=%d, decode=%d", len(prefillPods), len(decodePods))
	}
	// check for prefill and decode imbalance
	targetPod, isImbalanced := r.loadImbalanceSelectPrefillPod(prefillPods, r.prefillRequestTracker.GetPrefillRequestCountsForPods(prefillPods))
	if isImbalanced {
		klog.InfoS("load imbalance detected, selecting least-loaded prefill pod", "request_id", routingCtx.RequestID, "selected_prefill_pod", targetPod.Name)
		prefillPods = []*v1.Pod{targetPod}
		decodePods = utils.FilterPodsByLabel(decodePods, PDRoleSetIdentifier, targetPod.Labels[PDRoleSetIdentifier])
	}

	targetPod, maxRequestCount, maxThroughput, maxFreeGPUUsage, podRequestCounts, podThroughputs, podFreeGpuUsage := r.loadImbalanceSelectDecodePod(routingCtx, decodePods)
	if targetPod != nil {

Review comment (suggestion):
Assuming, if no pod was suitable for prompt length then final prefillPods array will be empty.

Better alternative will be

```
var prefillPodsByPromptLength []*v1.Pod

for _, pod := range readyPods {
	role := pod.Labels[PDRoleIdentifier]

	switch role {
	case "prefill":
		prefillPods = append(prefillPods, pod)

		if aibrixPromptLengthBucketing && r.isPodSuitableForPromptLength(pod) {
			prefillPodsByPromptLength = append(prefillPodsByPromptLength, pod)
		}

	case "decode":
		decodePods = append(decodePods, pod)
	}
}

// Override prefill pods only if bucketing produced results
if aibrixPromptLengthBucketing && len(prefillPodsByPromptLength) > 0 {
	prefillPods = prefillPodsByPromptLength
}
```



Code after:
// Pods without PodGroupIndex label are also included for backward compatibility.
func (r *pdRouter) filterPrefillDecodePods(routingCtx *types.RoutingContext, readyPods []*v1.Pod) (*v1.Pod, *v1.Pod, error) {
	prefillPods, decodePods := []*v1.Pod{}, []*v1.Pod{}
	promptLengthBucketingPrefillPods, promptLengthBucketingDecodePods := []*v1.Pod{}, []*v1.Pod{}

	var promptLength int
	if aibrixPromptLengthBucketing {
		promptLength, _ = routingCtx.PromptLength()
		klog.V(4).InfoS("prompt length based filtering enabled", "request_id", routingCtx.RequestID, "prompt_length", promptLength)
	}

	for _, pod := range readyPods {
		if _, ok := pod.Labels[PDRoleSetIdentifier]; !ok {
			continue
		}
		if _, ok := pod.Labels[PDRoleIdentifier]; !ok {
			continue
		}

		// For multi-node scenarios, only select pods from node_rank=0 (PodGroupIndex=0)
		// which have the HTTP server running
		if !isPodWithHTTPServer(pod) {
			continue
		}

		switch pod.Labels[PDRoleIdentifier] {
		case "prefill":
			prefillPods = append(prefillPods, pod)

			if aibrixPromptLengthBucketing && r.isPodSuitableForPromptLength(pod, promptLength) {
				promptLengthBucketingPrefillPods = append(promptLengthBucketingPrefillPods, pod)
			}
		case "decode":
			decodePods = append(decodePods, pod)

			if aibrixPromptLengthBucketing && r.isPodSuitableForPromptLength(pod, promptLength) {
				promptLengthBucketingDecodePods = append(promptLengthBucketingDecodePods, pod)
			}
		}
	}

	// Override prefill pods only if bucketing produced results
	if aibrixPromptLengthBucketing && len(promptLengthBucketingPrefillPods) > 0 {
		prefillPods = promptLengthBucketingPrefillPods
	}

	if aibrixPromptLengthBucketing && len(promptLengthBucketingDecodePods) > 0 {
		decodePods = promptLengthBucketingDecodePods
	}

	if len(prefillPods) == 0 || len(decodePods) == 0 {