[Misc][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

…urrent safety and stateful routing Signed-off-by: Jingyuan Zhang <[email protected]>

Signed-off-by: Jingyuan Zhang <[email protected]>

cmd/metadata/main.go

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

pkg/cache/cache_init.go

+	return InitForGateway(config, stopCh, redisClient)
+}
+
+func InitForGateway(config *rest.Config, stopCh <-chan struct{}, redisClient *redis.Client) *Store {


We should expose a featureFlag list to such initialization logic. we can do this later.

Yes, since the cache has limited use cases, I only expose initialization interfaces with a fixed feature set to avoid misconfiguration. More flexible initialization function can be added later.

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.

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.

Jeffwan · 2025-03-24T22:36:43Z

pkg/utils/pod.go

@@ -140,33 +140,59 @@ func CountReadyPods(podList *v1.PodList) (int64, error) {
 	return int64(readyPodCount), nil
 }

-// FilterReadyPods filters and returns a list of pods that have a valid PodIP.
-func FilterReadyPods(pods map[string]*v1.Pod) []*v1.Pod {


Seems FilterReadyPods has been deleted? do we have any other references?

It has been renamed to FilterRoutablePods to avoid confusion.

Ok. If we do not have other caller, that's ok. If we do, we may need to keep a copy. but seems we do not have it

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

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

Jingyuan Zhang added 2 commits March 24, 2025 14:31

Rebase: Cache and Router refactoring for concurrent performance, conc…

2c047f5

…urrent safety and stateful routing Signed-off-by: Jingyuan Zhang <[email protected]>

bug ifx

cea9419

Signed-off-by: Jingyuan Zhang <[email protected]>

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

zhangjyr commented Mar 19, 2025

Jeffwan commented Mar 19, 2025

Jeffwan commented Mar 21, 2025

Jeffwan Mar 24, 2025

zhangjyr Mar 24, 2025

Jeffwan Mar 24, 2025

zhangjyr Mar 24, 2025

Jeffwan Mar 24, 2025 •

edited

Loading

zhangjyr Mar 24, 2025

Jeffwan Mar 24, 2025

zhangjyr Mar 24, 2025

Jeffwan Mar 24, 2025

zhangjyr Mar 25, 2025

Jeffwan Mar 24, 2025

zhangjyr Mar 25, 2025

Jeffwan Mar 24, 2025

Jeffwan Mar 24, 2025

zhangjyr Mar 25, 2025

Jeffwan Mar 24, 2025

zhangjyr Mar 25, 2025

Jeffwan Mar 25, 2025

Jeffwan commented Mar 24, 2025 •

edited

Loading

Xunzhuo commented Mar 25, 2025

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

Are you sure you want to change the base?

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

Conversation

zhangjyr commented Mar 19, 2025

Pull Request Description

Related Issues

Pull Request Title Format

Submission Checklist

Jeffwan commented Mar 19, 2025

Jeffwan commented Mar 21, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jeffwan Mar 24, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jeffwan commented Mar 24, 2025 • edited Loading

Xunzhuo commented Mar 25, 2025

Jeffwan Mar 24, 2025 •

edited

Loading

Jeffwan commented Mar 24, 2025 •

edited

Loading