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

Code before:
}

func (r *pdRouter) isPodSuitableForPromptLength(pod *v1.Pod, promptLength int) bool {
	minLength, maxLength := r.getPodPromptRange(pod)

	if minLength > maxLength {
		return false
	}
	// If no prompt length range is configured, the pod is assumed to be suitable for handling any length.
	if minLength == 0 && maxLength == math.MaxInt32 {
		return true
	}

	return promptLength >= minLength && promptLength <= maxLength
}

// getPodPromptRange retrieves the minimum and maximum prompt lengths from pod labels.
func (r *pdRouter) getPodPromptRange(pod *v1.Pod) (int, int) {
	minLength := 0
	maxLength := math.MaxInt32

	if val, ok := pod.Labels[PromptMinLength]; ok {
		if parsed, err := strconv.Atoi(val); err == nil {
			minLength = parsed
		} else {
			klog.Warningf("failed to parse pod label %s with value %s as integer: %v", PromptMinLength, val, err)
		}
	}

	if val, ok := pod.Labels[PromptMaxLength]; ok {
		if parsed, err := strconv.Atoi(val); err == nil {
			maxLength = parsed
		} else {
			klog.Warningf("failed to parse pod label %s with value %s as integer: %v", PromptMaxLength, val, err)
		}
	}

	return minLength, maxLength
}

// getPromptLength extracts the prompt length from the request body.
func (r *pdRouter) getPromptLength(routingCtx *types.RoutingContext) int {
	var requestBody map[string]interface{}
	if err := json.Unmarshal(routingCtx.ReqBody, &requestBody); err != nil {
		klog.ErrorS(err, "failed to unmarshal request body for prompt length", "request_id", routingCtx.RequestID)
		return 0
	}

	var prompt string
	if messages, ok := requestBody["messages"].([]interface{}); ok && len(messages) > 0 {
		var promptBuilder bytes.Buffer

Review comment (suggestion):
Please remove the warn log message from line 814 and 822. If a pod is inaccurately configured then it will pollute the logs.

For future reference, if you need to log something, please include log like this 

```
klog.InfoS("<message>", "request_id", routerCtx.RequestID, <variables>...)
```

Code after:

func (t *PrefillRequestTracker) GetPrefillRequestCountsForPods(pods []*v1.Pod) map[string]int32 {
	counts := make(map[string]int32)
	for _, pod := range pods {
		countInterface, exists := t.podRequestCounts.Load(pod.Name)
		if !exists {
			counts[pod.Name] = 0
		} else {
			counts[pod.Name] = countInterface.(*atomic.Int32).Load()
		}
	}
	return counts
}

func (r *pdRouter) isPodSuitableForPromptLength(pod *v1.Pod, promptLength int) bool {
	minLength, maxLength := r.getPodPromptRange(pod)

	if minLength > maxLength {
		return false
	}
	// If no prompt length range is configured, the pod is assumed to be suitable for handling any length.
	if minLength == 0 && maxLength == math.MaxInt32 {
		return true
	}

	return promptLength >= minLength && promptLength <= maxLength
}

// getPodPromptRange retrieves the minimum and maximum prompt lengths from pod labels.
func (r *pdRouter) getPodPromptRange(pod *v1.Pod) (int, int) {
	minLength := 0
	maxLength := math.MaxInt32

	if val, ok := pod.Labels[PromptMinLength]; ok {
		if parsed, err := strconv.Atoi(val); err == nil {
			minLength = parsed
		}
	}

	if val, ok := pod.Labels[PromptMaxLength]; ok {
		if parsed, err := strconv.Atoi(val); err == nil {
			maxLength = parsed
		}
	}

	return minLength, maxLength
}