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

Code before:
		h.timestamps = newTimestamps
	}
}

func (p *prefixCacheAndLoadRouter) Route(ctx *types.RoutingContext, pods types.PodList) (string, error) {
	readyPods := utils.FilterRoutablePods(pods.All())

	if len(readyPods) == 0 {
		return "", fmt.Errorf("no pods to forward request")
	}
	if len(readyPods) == 1 {
		for _, pod := range readyPods {
			ctx.SetTargetPod(pod)
			return ctx.TargetAddress(), nil
		}
	}

	var podUpdateNeeded bool
	func() {
		p.podsMu.RLock()
		defer p.podsMu.RUnlock()
		podUpdateNeeded = len(readyPods) != p.numPods
	}()

	if podUpdateNeeded {
		klog.Infof("requestID: %s, num pods in data structure: %d", ctx.RequestID, p.numPods)
		klog.Infof("requestID: %s, current actual ready pods: %d", ctx.RequestID, len(readyPods))
		p.podsMu.Lock()
		p.updatePodSet(readyPods) // Move update pod logic to separate function
		p.podsMu.Unlock()
		klog.Infof("requestID: %s, num pods in data structure after updatePodSet: %d", ctx.RequestID, p.numPods)
	}

	tokens, err := utils.TokenizeInputText(ctx.Message)
	if err != nil {
		klog.Errorf("requestID: %s, Tokenization failed: %v", ctx.RequestID, err)
		return "", err
	}

	node, matchedTokens, _ := p.cache.AddPrefix(tokens, ctx.Model, "")
	var matchedPods []*v1.Pod
	var matchedPodsNames []string
	klog.V(5).Infof("requestID: %s, node.GetModelToPods()[%s]: %v", ctx.RequestID, ctx.Model, node.GetModelToPods()[ctx.Model])
	if modelPods, ok := node.GetModelToPods()[ctx.Model]; ok {
		readyPodsMap := make(map[string]*v1.Pod)
		for _, pod := range readyPods {
			readyPodsMap[pod.Name] = pod
		}
		for podName := range modelPods {
			if pod, exists := readyPodsMap[podName]; exists {
				klog.V(5).Infof("requestID: %s, Matched ready pod: %s", ctx.RequestID, podName)

Review comment (bug):
To keep logging same, use klog.InfoS.

Better to merge multiple log statements in one and also try to minimize the number of log messages for better debugability.

Code after:
		defer p.podsMu.RUnlock()
		podUpdateNeeded = len(readyPods) != p.numPods
	}()

	if podUpdateNeeded {
		p.podsMu.Lock()
		p.updatePodSet(readyPods) // Move update pod logic to separate function
		p.podsMu.Unlock()
		klog.InfoS("Request processing", "requestID", ctx.RequestID, "updatePodSet", p.numPods)
	}

	tokens, err := utils.TokenizeInputText(ctx.Message)
	if err != nil {
		klog.Errorf("requestID: %s, Tokenization failed: %v", ctx.RequestID, err)
		return "", err
	}

	node, matchedTokens, _ := p.cache.AddPrefix(tokens, ctx.Model, "")
	var matchedPods []*v1.Pod
	var matchedPodsNames []string
	if modelPods, ok := node.GetModelToPods()[ctx.Model]; ok {
		readyPodsMap := make(map[string]*v1.Pod)
		for _, pod := range readyPods {
			readyPodsMap[pod.Name] = pod
		}
		for podName := range modelPods {
			if pod, exists := readyPodsMap[podName]; exists {
				matchedPods = append(matchedPods, pod)
				matchedPodsNames = append(matchedPodsNames, podName)
			}
		}
	}

	var targetPod *v1.Pod
	matchRatio := float64(len(matchedTokens)) / float64(len(tokens))
	prefixRoutingThreshold := 0.5
	klog.InfoS("requestID: %s, Matched tokens/Total tokens: %d/%d, Matching ratio: %.0f%%, len(matchedPodsNames): %d, matchedPodsNames: %v", "requestID", ctx.RequestID, "matchedTokens", len(matchedTokens), "totalTokens", len(tokens), "matchingRatio", matchRatio*100, "matchedPodsNamesCount", len(matchedPods), "matchedPodsNames", matchedPodsNames)

	if matchRatio > prefixRoutingThreshold {
		klog.InfoS("requestID: %s, Do prefix-aware routing! (matching ratio: %.2f > %.2f)", "requestID", ctx.RequestID, "matchRatio", matchRatio, "threshold", prefixRoutingThreshold)
		var prefixMatches []prefixMatch

		currentNode := node
		for currentNode != nil {
			if modelPods, ok := currentNode.GetModelToPods()[ctx.Model]; ok {
				var nodePods []*v1.Pod
				for podName := range modelPods {
					for _, pod := range readyPods {
						if pod.Name == podName {
							nodePods = append(nodePods, pod)
						}