[API] Cache and Router refactoring for concurrent performance, concurrent safety and stateful routing. #884

zhangjyr · 2025-03-19T21:04:06Z

Pull Request Description

Refactoring for cache:

Merge multiple pod, model, and metric mapping by adding Pod metadata and Model metadata and using two main thread-safe registries for metadata.
Eliminate the global cache mutex lock on reading and replace it with multiple layers of locks in the metadata.
Cache serves in different contexts: control manager, metadata, and gateway. Added cache initialization helpers for each context for later customization.
Added more unit test cases to cover the pod and model adapter changes.

Refactoring for router

Eliminate thread-unsafe map access in the Router interface.
Merge two contexts, context.Context and routing.RoutingContext, as RoutingContext
Abstract away the Router interface and the RoutingContext for shared access from both the routing and cache package.
RoutingContext now supports both sync and async routing resolution by following the Promise paradigm.

Related Issues

Resolves: #868

Important: Before submitting, please complete the description above and review the checklist below.

Contribution Guidelines (Expand for Details)

We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:

Pull Request Title Format

Your PR title should start with one of these prefixes to indicate the nature of the change:

[Bug]: Corrections to existing functionality
[CI]: Changes to build process or CI pipeline
[Docs]: Updates or additions to documentation
[API]: Modifications to aibrix's API or interface
[CLI]: Changes or additions to the Command Line Interface
[Misc]: For changes not covered above (use sparingly)

Note: For changes spanning multiple categories, use multiple prefixes in order of importance.

Submission Checklist

PR title includes appropriate prefix(es)
Changes are clearly explained in the PR description
New and existing tests pass successfully
Code adheres to project style and best practices
Documentation updated to reflect changes (if applicable)
Thorough testing completed, no regressions introduced

By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.

Jeffwan · 2025-03-19T23:04:46Z

Some code changes might be conflict with #878. that's not a problem now. We can review the core logic first and then rebase the changes later.

Jeffwan · 2025-03-21T16:42:21Z

I will have a check today

cmd/metadata/main.go

pkg/cache/cache_init.go

Jeffwan · 2025-03-24T22:16:15Z

pkg/cache/cache_api.go

 	// Parameters:
 	//   modelName: Name of the model
 	// Returns:
 	//   bool: True if model exists, false otherwise
-	GetModel(modelName string) bool


Why do you delete ListPods and GetModel in the interface?

There is no use case for ListPods except in unit tests. For this reason, I only provide a non-efficient implementation to return all pods in the cache and delete the function from the API.

The name of GetModel implies it will return some model object. However, the original GetModel API only returns true or false. HasModel would be a more proper name for this API.

Jeffwan · 2025-03-24T22:21:15Z

pkg/cache/cache_api.go

 	// ListPodsByModel gets pods associated with a model
 	// Parameters:
 	//   modelName: Name of the model
 	// Returns:
 	//   map[string]*v1.Pod: Pod objects matching the criteria
 	//   error: Error information if operation fails
-	ListPodsByModel(modelName string) (map[string]*v1.Pod, error)


please help me recap the motivation here. the core part is to refactor the cache store registry with better and efficiency data structure? could you remind me the painpoint here? is it just for code refactor or there's some challenges extend it or performance issue?

The original API exposes a thread-unsafe map object while the implementation returns what is stored in the cache without making a copy. This will cause panic on map read accessing while the map is written in the Lock() context.

Because there is no need for map keys, the new implementation returns a slice with read-through-cache acceleration.

Jeffwan · 2025-03-24T22:21:44Z

pkg/cache/cache_api.go

@@ -73,7 +70,7 @@ type ModelCache interface {
 	// Returns:


Will we have other informations in future? in that case, we may need complex structure as well?

That is possible. As *v1.Pod suggests, if our model(CRD) definition contains more options, such as labels, we may need to update returns here with complex structure. In this case, I would suggest adding the model definition to github.com/vllm-project/aibrix/api/model/v1alpha1

Jeffwan · 2025-03-24T22:23:53Z

pkg/cache/cache_api.go

 	//   requestID: Unique request identifier
 	//   modelName: Name of the model
 	// Returns:
 	//   int64: Trace term identifier
-	AddRequestCount(requestID string, modelName string) (traceTerm int64)
+	AddRequestCount(ctx *types.RoutingContext, requestID string, modelName string) (traceTerm int64)


here's an interface change . AddRequestCount and DoneRequestCount is moved under TraceCache

In fact, AddRequestCount is paired with DoneRequestCount or DoneRequestTrace. For example
AddRequestCount + DoneRequestCount or AddRequestCount + DoneRequestTrace. While these functions were designed for trace only, they are also used as an entry point for real-time statistics. In contrast, the metric cache offers an interface for accessing cached metrics/real-time statistics.
In summary:

MetricCache is used for reading metrics

TraceCache is used to generate real-time statistics besides trace. I think MetricCollector may be a better name for this interface. I welcome other options.
Related changes:
enableGPUOptimizerTracing is moved from gateway to cache to offer:

While enableGPUOptimizerTracing = false, the trace serving routine (for writing to redis) will not start altogether.

The trace collection operations in AddRequestCount, DoneRequestCount, and DoneRequestTrace can be disabled individually.

I renamed TraceCache to RequestTracker

Jeffwan · 2025-03-24T22:24:52Z

pkg/cache/cache_impl.go

-	defer c.mu.RUnlock()
-
-	pod, ok := c.Pods[podName]
+	metaPod, ok := c.metaPods.Load(podName)


I want to get more context here. (Sorry I didn't clearly review the issue yet). original way is very common in golang, what's the problem?

R lock is ok. Since we use sync.map to replace thread-unsafe maps, Rlock is not needed anymore.
Using sync.map lower the lock granularity from cache-wide to map-wide.

pkg/utils/sync_map.go

Jeffwan · 2025-03-24T22:34:27Z

pkg/utils/pod_array.go

+)
+
+const (
+	DeploymentIdentifier string = "app.kubernetes.io/name"


some scenarios like distributed inference use other orchestrators and we may not have this label.
better to have a group identifier could be configured

I am thinking maybe client-go indexer is another option if we think map[]xxx is not efficient

Let me define a new interface to replace PodArray. We can replace the PodArray with Indexer if that's more efficient later.
For DeploymentIdentifier, I can also add something like PodGrouper for alternative implementations.

A types.PodList is added to decouple PodArray implementation. I hardly think the Indexer in the informer may be helpful in our context because we have retrieved pods' information and cached out of informer caches.

pkg/utils/pod.go

Jeffwan · 2025-03-24T22:51:00Z

/cc @Xunzhuo @varungup90 Please help review it.

@zhangjyr this is still a huge PR, if we have some agreement on the data structure and helper utilities. I suggest to add those add only code first to make this code more simpler

Xunzhuo · 2025-03-25T08:04:02Z

I think we need to let e2e test passed before reviews

pkg/types/router_context.go

varungup90

Nice PR, overall LGTM. Left some minor comments.

pkg/utils/pod.go

varungup90 · 2025-03-26T02:29:57Z

pkg/plugins/gateway/gateway.go


 		case *extProcPb.ProcessingRequest_RequestBody:
-			resp, model, targetPodIP, stream, traceTerm = s.HandleRequestBody(ctx, requestID, req, user, routingStrategy)
+			resp, model, routerCtx, stream, traceTerm = s.HandleRequestBody(ctx, requestID, req, user, routingAlgorithm)


technically model and stream can be added to routingCtx as well. Does not need to be part of this PR, it can be done in follow up PR.

routerCtx may not be available for requests that have not specified a routing strategy. For now, the "model" and "stream" are still needed here. However, since the routing context also included the RoutingAlgorithm field, it is possible to eliminate these fields later.

pkg/plugins/gateway/gateway_req_body.go

varungup90 · 2025-03-26T02:41:58Z

pkg/plugins/gateway/gateway_req_headers.go

 		klog.ErrorS(nil, "incorrect routing strategy", "routing-strategy", routingStrategy)
 		return generateErrorResponse(
 			envoyTypePb.StatusCode_BadRequest,
 			[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
 				Key: HeaderErrorInvalidRouting, RawValue: []byte(routingStrategy),
-			}}}, "incorrect routing strategy"), utils.User{}, rpm, routingStrategy
+			}}}, "incorrect routing strategy"), utils.User{}, rpm, routingAlgorithm


In text "incorrect routing algorithm". To a global replace for strategy -> algorithm

Let's keep strategy as user-faced string variable. I don't plan to change all header staff. And algorithms be valid strategies that verifed.

pkg/cache/registry.go

pkg/cache/informers.go

…urrent safety and stateful routing Signed-off-by: Jingyuan Zhang <jingyuan.zhang0929@bytedance.com>

Jeffwan · 2025-03-27T18:50:52Z

@Xunzhuo jingyuan fixed the integration test issues, do you get a chance to help take a look?

Jeffwan · 2025-03-28T18:49:03Z

@Xunzhuo Just wanted to loop you in—this PR is currently blocking some performance-related changes (prefix-cache aware routing) that community users are waiting for. it's quite large and deserves more attention. We plan to merge it to unblock some following changes this time as an exception.

If you have some bandwidth, please help review this change and @zhangjyr will address those comments in future PRs
We do want to avoid such large PRs in the future. Big PRs like this are really tough to review properly and can slow everyone down. Contributors should split their changes into smaller, reviewable chunks to make the process smoother for everyone.

…rent safety and stateful routing. (vllm-project#884) Cache and Router refactoring for concurrent performance, concurrent safety and stateful routing Signed-off-by: Jingyuan Zhang <jingyuan.zhang0929@bytedance.com> Co-authored-by: Jingyuan Zhang <jingyuan.zhang0929@bytedance.com>

zhangjyr requested review from Jeffwan, varungup90, gangmuk and nwangfw March 19, 2025 21:04

zhangjyr changed the title ~~[Misc] Cache and Router refactoring for concurrent performance, concurrent safety and stateful routing.~~ [Misc][API] Cache and Router refactoring for concurrent performance, concurrent safety and stateful routing. Mar 19, 2025

zhangjyr mentioned this pull request Mar 20, 2025

[API] Refactor: core cache design and impl #878

Merged

zhangjyr force-pushed the feature/cache_router_refactor branch from dd7d52d to cea9419 Compare March 24, 2025 21:36

Jeffwan reviewed Mar 24, 2025

View reviewed changes

cmd/metadata/main.go Show resolved Hide resolved

Jeffwan reviewed Mar 24, 2025

View reviewed changes

zhangjyr requested a review from Jeffwan March 26, 2025 01:22