Refactor scheduler to make it more readable

[liu-cong]

This refactor does the following:

1. Explicitly define ActiveModels and WaitingModels that map to the running and waiting adapters metric. Previously we combine both as ActiveModels. This change makes it clear.
2. Create a scheduling.Context object that holds the contextual info during a request scheduling cycle. This has 2 benefits: a) Any contextual info can be added to this struct instead of modifying the scheduler interface, making extending the scheduler easier (will soon need this for prefix caching implementation). b) Create a snapshot of the pod metrics for the scheduling cycle, to reduce concurrent access to the shared datastore, as well as provide a consistent view of the pods and metrics during scheduling. This makes debugging easier.
3. Created a simpleFilter which contains the user readable filter name plus filter function. Making the filter chain composition much cleaner.
4. This is the only change of logic. Previously we had 2 slightly different LoRA affinity filters, one is considered the "main" lora affinity filter, the other one is only used when all pods are significantly queued up (>128 waiting requests). As far as I can tell, there was no clear reason why we needed both. The reason was because we added the "main" one later to account for long queues. In this PR I consolidated to the "main" lora affinity filter. And benchmark results showed the refactored version had almost identical results with HEAD.

Benchmarks

I ran benchmarks using the LPG tool, with 10 lora adapters sharing the same traffic, and 8 vllm replicas running on H100 with max-lora=4.

• LPG reported data between baseline and refactor:

• EPP metrics and vLLM metrics between the baseline and refactor

  • baseline
    
    

  • Refactor
    
    

[liu-cong]

cc @kaushikmitr @kfswain

[netlify]

✅ Deploy Preview for gateway-api-inference-extension ready!

Name	Link
🔨 Latest commit	87162b4
🔍 Latest deploy log	https://app.netlify.com/sites/gateway-api-inference-extension/deploys/67f067432baed800080c0429
😎 Deploy Preview	https://deploy-preview-645--gateway-api-inference-extension.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

---

To edit notification comments on pull requests, go to your Netlify site configuration.

[liu-cong]

It's much cleaner this way to show how we handle critical vs. sheddable differently, than the previous filter.

[liu-cong]

This is where we consider both active and waiting as lora affinity. This refactor just makes it super clear.

[liu-cong]

This is the original lora affinity filter. This one considered a pod with an active lora adapter equally with a pod that has space to load a new lora.

Later on we introduced the loRASoftAffinityFilter below, which generally favors pods with active loras, but with a probability to choose a pod with space to load a new lora, to avoid hot spots.

The new affinity filter is considered better than the original one. There is no clear reason why we want to keep the original one.

[kfswain]

Ack! Will take a look later today, this looks awesome at a glance. Thanks!

[kaushikmitr]

If I read the first chart correctly the QPS goes all the way up to 3000? 😮

[liu-cong]

If I read the first chart correctly the QPS goes all the way up to 3000? 😮

Not the real QPS ... I set the QPS to 300, with 10 adapters, so that's how the 3000 is calculated. However, the LPG tool isn't able to send that high QPS though. If you look at the InferenceModel metrics on EPP, the real QPS was about 30 requests/s for each adapter.

[nirrozenbaum]

adding my two cents -
I think the current code is mixing filter and scorer terms.
Filter should be defined as a hard requirement, that if we cannot meet the request is dropped.
Scorer should be defined as a soft requirement, that we can use to score/rank pods.
with the above terminology - if a filter fails (ends up with 0 pods), the request should fail.

additionally, I think the current impl that uses "decision tree" is very hard to follow and makes it very hard to understand what is the pod selection logic. I would flatten this structure and have a slice of filters and a slice of scorers (could be defined differently based on the request and other properties, e.g., criticality).

the way I see it (using the above terminology) -
for each request a set of filters should run sequentially (no meaning for the order cause each is just a boolean conditional).
if we're left with pods that can serve the request, we need to rank them and run scorers (each scorer is independent, scorers can run in parallel). then we can calculate a weighted score for each pod based on scorers configuration (the weight of each).

another important point I would consider - current code makes it very hard for someone to fork the code and change filters. we need to make it a super easy thing to do by design.

@kfswain this is part of the discussion we started in the last community call and we probably need to continue.

[ahg-g]

I agree we need to move to a better structure that potentially mimics that of the pluggable framework of kube-scheduler, but I believe this is a good initial refactor, it is net positive. Leaving lgtm to others.

/approve

[k8s-ci-robot]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: ahg-g, liu-cong

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

• OWNERS [ahg-g]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[nirrozenbaum]

my comments were in general about the existing structuring and not about this PR.
this PR is an improvement comparing to the existing code.

agree with you and leaving the "better structuring" discussion out of the scope of this PR.

/lgtm

[kfswain]

additionally, I think the current impl that uses "decision tree" is very hard to follow and makes it very hard to understand what is the pod selection logic. I would flatten this structure and have a slice of filters and a slice of scorers (could be defined differently based on the request and other properties, e.g., criticality).

+1 current implementation is rough to mentally follow. I've been thinking about how we can separate this logic.

B/C our tooling has some gateway-like properties (like InferenceModel routing) and some LB-like properties (our scheduler) I think we can break this up.

We have a Routing layer, that handles any routing needs such as LoRA selection, prompt injection, maybe even checks that reject a request before it ever need be scheduled (and can be easily extended by a user) and then a Scheduling layer, that just scores endpoints (I don't know if we need to do any filtering, perhaps I'm wrong but i think we can indirectly filter via score) but we define that interface strongly.

We will also be introducing queuing at the EPP layer. So that makes the logic even more clean:

1. Routing layer: modifies request and validates that the request should even be sent, w/an interface to make extension simple
2. Queuing mechanism acts as flow control to scheduler
3. Scheduler: scoring mechanism to select an endpoint, also w/ a strong interface that allows for custom implementations to be easily plugged

I'll have more to present on the Th meeting. I'll get something written up/diagram/etc

[liu-cong]

Agree that we should simplify the filter tree, and a promising approach is to have a chain of filter and score plugins like the kube scheduler.

I did some initial experiment with changing the filters to "scorers" in https://github.com/liu-cong/llm-instance-gateway/tree/score

My initial benchmark results didn't performa as good as current implementation on LoRA affinity. Though with some tuning we can perhaps perform better. Will keep experimenting once I get some free cycles.