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

Code before:
		promptLength = r.getPromptLength(routingCtx)
		klog.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)
		case "decode":
			decodePods = append(decodePods, pod)
		}

		if aibrixPromptLengthBucketing {
			if !r.isPodSuitableForPromptLength(pod, promptLength) {
				continue
			}
			switch pod.Labels[PDRoleIdentifier] {
			case "prefill":
				promptLengthBucketingPrefillPods = append(promptLengthBucketingPrefillPods, pod)
			case "decode":
				promptLengthBucketingDecodePods = append(promptLengthBucketingDecodePods, pod)
			}
		}
	}

	if len(promptLengthBucketingPrefillPods) != 0 {
		// If pods that meet the given length exist, use them; otherwise, use all pods as a fallback.
		prefillPods = promptLengthBucketingPrefillPods
	}

	if len(promptLengthBucketingDecodePods) != 0 {
		decodePods = promptLengthBucketingDecodePods
	}

	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

Review comment (refactor):
Seems the codes from 197-201 is duplicated. 

is the code a filter? If so, can we write in similar way like this?

```
		// 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
			}
}
```

Code after:
	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 {
		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 {