Initial Scheduler Subsystem interface

[kfswain]

This PR extends the EPP Architecture proposal to provide a concrete definition for the Scheduler Subsystem.

Currently updating the readme, but cutting the PR to discuss the interface

[netlify]

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

Name	Link
🔨 Latest commit	0a2e3e0
🔍 Latest deploy log	https://app.netlify.com/projects/gateway-api-inference-extension/deploys/683a0e52ae4421000853590f
😎 Deploy Preview	https://deploy-preview-845--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 project configuration.

[nirrozenbaum]

there are two different points to address here:

1. extend the existing scheduler to work with multiple scheduling cycles and define the extension points. At this stage scheduler should be extensible via code.
2. After point 1 is settled, the next step is to configure it using a configuration file. there’s more than that though, like how do we load out of tree plugins (e.g., like kube scheduler using shared obj?). Anyway, I wouldn’t go into configuration file as this is a very advanced stage and current example seems like it might be missing fields (e.g., parameters for each plugin). we first need to define how it works, what are extension points and their names, under which layer each extension falls, etc.

bottom line - I would leave configuration file out of scope at this point.

[kfswain]

I actually disagree here. Your other comment, makes the point that a SchedulingProfile is more of a struct, and I think that actually emphasizes the need to have a config file. B/C if a profile is a struct, it needs a way for a user to express how that struct is populated.

If we don't provide config, we demand that anyone developing their algo rebuild their binary every single time they make a change, rather than update a config file and redeploy the EPP. I think the sooner much of this can be expressed outside of a new build, the better the experience, & the higher the rate of adoption.

parameters for each plugin

This is a great design question to bring up & something we should hash out. Should each profile have it's own separate set of configurable knobs (example: kv_cache_util). I could see that having merit.

[nirrozenbaum]

I think the sooner much of this can be expressed outside of a new build, the better the experience, & the higher the rate of adoption.

just to clarify - of course I agree that a configuration file is needed.

I was trying to make a point that same as a big PR is much more convenient to discuss when broken into smaller PRs, I think these are two completely separate topics and we could benefit from separating the discussions.

one topic is having a multi cycle scheduler and another topic is how to express the scheduler configuration in a file including the out of tree plugins handling. each of these topics is big on its own and can be solved separately.

[nirrozenbaum]

score is NOT a property of an endpoint.
it’s a result of scorers invocation in one cycle. this shouldn’t be kept as a field in endpoint (we’ve been in this discussion before in the latest scheduler changes)

[kfswain]

Is you prefer a different name, that is fine, but ultimately each endpoint in a Scheduling Cycle should have a score associated with it, however that is done. The scope of this field would be contained to the scheduler subsystem. So while I agree this is not a field on a ModelServer struct, with a lifecycle that is longer than a single scheduling cycle. An object with the same lifecycle as the scheduling cycle it is in seems reasonable to have a score field attached to it.

[nirrozenbaum]

today we have in a cycle a map[endpoint]float64 which lives in the scope of a cycle and is expressing the score of an endpoint during the cycle.

I don’t see a need for a struct/interface with SetScore/GetScore

[kfswain]

Making complex structs a key is a code smell to me: map[endpoint]float64, esp when the value is a simple float. You also mention further down with the desire to pick topX endpoints, which implies sorting, removing the value of a map.

[nirrozenbaum]

right. we had that issue also in today’s code. this is why we have ScoredPod struct which lives in the lifecycle of a cycle and is encapsulating endpoint with its score in the current cycle.
so in current code:
endpoint is an interface and lives as long as the endpoint is alive
scoredEndpoint lives within a cycle and holds ref to the endpoint + score

[kfswain]

So the usage of a newly defined Endpoint is intentional, you'll note that the only not standard lib import I make is to the current Scheduling library, and only for CycleState. The scheduling lib shouldn't have cross dependencies

[nirrozenbaum]

I think it worth doing some “playground” around this. there’s no way around the fact the we need access to data layer and specifically to endpoint state. unless the state is part of scheduling which doesn’t sound right to me

[nirrozenbaum]

The scheduling subsystem wouldn't have the distinction. I am treating this more like a library, in the context of a scheduling library all Endpoints would have scores, as that is one of the primary methods of interaction a scheduler has with its endpoints.

not sure I understand your intention here.
Endpoint lifecycle is different than Score lifecycle.
Endpoint lives as long as the model server exist. Score lives within the score of a cycle.

[kfswain]

The Endpoint object within the scheduling system != the endpoint object within the data layer

[kfswain]

The scheduling subsystem is self-contained & so the endpoint, from the perspective of the scheduling system, shares lifecycle, and strongly cares about score.

[nirrozenbaum]

I recommend we rename pre/post schedule to pre/post cycle.
current namings may become very confusing once the schedule becomes multi-cycle.
In my view these names represent the following:
PreSchedule == ProfileSelection
PostSchedule == results aggregation

[kfswain]

Seems reasonable enough

[nirrozenbaum]

This call should receive the request or some request properties in order to select the profiles (decision on which profiles should be based on the request props)

[nirrozenbaum]

why Scheduler is a plugin?

[nirrozenbaum]

This was removed from the diagram.
two questions regarding this:

• I do wonder if we actually need Pre/Post Cycle. what is the use case for using them? can we do without them?
• if we do keep PreCycle, I think we need also PostCycle in the diagram for symmetry.

I was looking on Prefix implementation (as it is the only example that uses PreCycle atm) and noticed that:

• PreCycle is used to call matchLongestPrefix which internally updates servers with their longest prefix match.
• this in theory may calculate longest prefix for servers that will be filtered by the filters. that means we might do redundant calculations as those servers will not be candidates for selection.
• it sounds to me like we could call matchLongestPrefix at the beginning of the Score function and calculate prefix only for the relevant servers.
• bottom line, the only use case where we currently use PreCycle is implemented in an insufficient way (calculating for all servers instead of just the ones that will remain relevant after filtering). intuitively, I would argue that we can remove Pre/PostCycle completely as I was not convinced (yet) that those are needed.

[kfswain]

So I'm glad you brought this up.

I thought Pre/Post schedule was a requirement you brought over. I've struggled to see a case for them.

The best I can see for Pre- is some sort of per-cycle digestion of data that will be used in multiple filters/scorers. But that could still technically happen in the ProfileSelection implementation.

Same thing with Post that could be some per-cycle` metrics emission or something. But that could likely happen in the aggregration step also.

I support removing them. @ahg-g any opinion here?

[ahg-g]

I recommend to not add any extensions that we don't have use cases for yet. We can always add new extensions, but we can't remove them once they are established.

[nirrozenbaum]

to summarize, I think we all agree pre/post should be removed.
having said that, prefix plugin currently uses both.
as a prerequisite step to remove both interfaces, we need to move the pre logic inside the score function, and the post (that stores the prefix for next requests) to PostResponse.
the first is fairly easy to do. the second is a bit more tricky since currently PostResponse doesn't get the prompt and this requires more changes.

[ahg-g]

I am thinking how we want to set the selected endpoint in the request response. One approach is to do that in a PostSchedule extension, the other is to leave that open and hardcode it in IGW director itself based on some convention we set on the SchedulingResult.

I am leaning towards the former to give us more flexibility in changing that behavior: e.g., for p/d, we will initially have to set one as the canonical endpoint that the proxy should forward the request to and the other as a header. Perhaps this can also help us implement fallbacks as well?

[kfswain]

Agreed, the former seems more appropriate to me as result aggregation is opinionated & system based. If we set a standard we inject opinions, reducing the extensability.

PostResponse doesn't get the prompt and this requires more changes.

We can update this fairly easily.

[nirrozenbaum]

+1 on PostSchedule

[nirrozenbaum]

just a thought -
having a Name function is enough to be a plugin. this looks a big strange (any struct with Name function is a plugin?).
plugin should have some kind of Run function.

[kfswain]

Using KubeScheduler as a frame of reference resonates well with this audience, so I stuck with those conventions: https://github.com/kubernetes/kubernetes/blob/fa6d252d845a415be53d8de0475404df19a838e6/pkg/scheduler/framework/interface.go#L443

[kfswain]

Additionally, The bigger the interface, the weaker the abstraction. And abstracting each specific interface to match a Run function would ruins the ability for a single plugin to implement multiple interfaces.

As is currently, I could write a plugin that implements Filter & Score, even though they have the same params/return type. Abstracting to Run eliminates that

[elevran]

IMO, Name() is not a useful interface and it seems to be used for code reuse (avoid writing Name() on the real abstraction?). I would like to suggest we extend the Plugin's interface to include:

• plugin type (can use Go reflection, but string names might be friendlier)
• plugin name (e.g., the prefix aware scorer in P and D cycles are of the same type but possibly different instances with different names)
• a Configure() method (or Init() or whatever else is a good name to for when it is configured from parameters in a file)
• the actual plugin function (Score, Filter, Pick etc.)

[nirrozenbaum]

I completely disagree on the Scheduler functions.
I think the Scheduler should have a single (public) function which is Schedule. all the others that are specified here are internal functions of the scheduler and not the public interface.

to break it down more, when calling to scheduler -
do you expect that the caller (e.g., director) will call ProfileSelection, then director goes iteratively on the profiles and for each it invokes the plugins in the correct order?
I don't think that was your intention.

what I'm expecting to see here is something along these lines:

Schedule(ctx context.Context, *schedulingtypes.Request) (result []*schedulingtypes.Result, err error)

and internally in the scheduler implementation I expect to see something like the following:

func (s *Scheduler) Schedule(ctx context.Context, req *types.Request) ([]*types.Result, error) {
   profiles := s.profilePicker.Pick(req, s.availableProfiles) // selecting profiles is a function of the request
   // Snapshot pod metrics from the datastore to:
   // 1. Reduce concurrent access to the datastore.
   // 2. Ensure consistent data during the scheduling operation of a request between all scheduling cycles.
   sCtx := types.NewSchedulingContext(ctx, req, types.ToSchedulerPodMetrics(s.datastore.PodGetAll()))
   
   result := []*types.Result{} // assuming the Result has a profileName field inside
   for name, profile := range profiles {
   	// run the selected profiles and collect results
   	// runProfileCycle internally runs the plugins according to their order
   	cycleResult, err := s.runSchedulerProfileCycle(sCtx, profile)
   	if err != nil {
   		return result, fmt.Errorf("failed to run all required scheduling profiles - %w", err)
   	}
   	result = append(result, cycleResult)
   }

   return result, nil
}

[ahg-g]

I agree that here we just need a Schedule function. I am also not sure we necessarily need to make the Scheduler an interface vs just a struct, a struct seems sufficient here, so it would be good to provide a reasoning why an interface is needed here.

I also recommend to define a type named EndpointPickingProfile in this PR, and show that the list of profiles are a parameter to a NewScheduler() instantiation call. This will help clarify the Schedule snippet above and the idea behind the profile picker extension point.

[ahg-g]

Actually one thing we discussed last week was having the ProfilePicker call return a single profile with a continue flag instead of returning a list of profiles. This allows conditional execution of profiles.

[nirrozenbaum]

@ahg-g right 👍🏻.
the above sample I wrote was the simplest form of implementing Schedule function just to emphasize that current interface functions are not needed.

[kfswain]

ProfileSelection & SchedulingResult are both pluggable functions. If we were to make a struct, we would need to make an interface that defines these functions anyway. The Director will call the scheduler framework entrypoint and THAT can be a struct. But this should absolutely be an interface or disaggregated pluggability (as well as other complex multi-profile implementations) is not possible

[kfswain]

Yes, MultiProfileHandler & Scheduler interface are basically the same signature. So sounds like we are on the same page (I personally, would still just call it a scheduler, as even the default implementation will still need to pick a single profile & aggregate the result, but tbd)

But yes, apologies, the config was indeed not sketched out here. I was adding it implicitly when I added the config file, agreed that may have been misleading. We are on the same page.

I may have shared this PR too early as some of the context wasn't filled in and I didnt get a chance to fill it in before discussion started. Updating this is a prio of mine today

[nirrozenbaum]

SG, sorry for the misunderstanding.
we're on the same page with regards to having a single plugin with both functionalities, yes.

[nirrozenbaum]

this discussion made me realize our name selection has to be great(!).
otherwise, it could become very confusing for a reader (even we were able to confuse).

now after resolving the misunderstanding, I would consider defining that plugin like the following (or something along these lines):

type MultiProfilePlugin interface {
    // SelectProfiles determines which profiles should run for the given request.
    SelectProfiles(request Request) []*SchedulerProfile

    // ProcessProfileResults handles the outcome of each selected profile.
    // It may aggregate results, log test profile outputs, or apply custom logic.
    ProcessProfileResults(request Request, results map[string]ProfileResult) (FinalResult, error)
}

I think the above names are much more descriptive than PreSchedule/PostSchedule which don't say a lot and are also not mentioning anything about profiles or their result processing.

MultiProfilePlugin suggests it's a plugin dealing with multiple profiles.
function names like SelectProfiles and ProcessProfileResults are also very descriptive.

[elevran]

Adding one more consideration into SelectProfile:
In the case of disaggregation we may want to invoke a cycle only after the previous cycle as made a choice (e.g., co-locate P/D in the same node, if possible).
I think this can be done if we mandate that profiles are run sequentially and in the order returned (i.e., SchedulerProfile[0] before SchedulerProfile[1], never in parallel) and their results are made available in the scheduling context to later profiles.
Further (and this might be a separate use case), we may want to decide if to run a profile only after knowing (in the MultiProfilePlugin) the result of the previous profile (e.g., if request prefix is 99% cache on the D, P is not needed).

Does it make sense to allow multiple return values from SelectProfile

• list of profiles to run (possibly empty)
• indication whether scheduling is done or another callback into ScheduleProfile is needed.

I think this would address both above use cases.
Thoughts?

[nirrozenbaum]

@elevran yeah, agreed on the comment.
my previous comments were more focused on name selection rather than on the function args/functionality.
in the last PR where I started initial support of multi cycle scheduler, the suggestion you made already exist.

current code functionality is -
request from profile picker iteratively the profiles to run based on the request properties and the results of previous profiles execution (if there are such). more than one profile can be returned (e.g., can return all).
the iteration ends when the profile picker returns no profiles to run.

see here:
https://github.com/kubernetes-sigs/gateway-api-inference-extension/blob/main/pkg/epp/scheduling/scheduler.go#L113-L130

[nirrozenbaum]

is this type trying to encapsulate the multi-cycle results and map profileName -> it's result?
I would define a SchedulingResult per cycle and not per the complete schedule function.
then the return value of Schedule call can be either []*SchedulingResult (e.g., if the result includes the profile name) or alternatively map[string]*SchedulingResult (map from profile name to the result of the profile cycle invocation).

this makes it clearer that each Profile cycle has it's own result.

additionally, a cycle result might have more than []Endpoint as a return value e.g., if we have equally scored endpoints and additional endpoints with very close score - we might want to have bestScoreEndpoints and backupEndpoints, and return them all to the PostSchedule to let the plugin decide which endpoint to use.

each cycle is making a local optimized endpoint selection. when receiving all cycles selection, we may want to do global optimization. for example in P/D (cycle for P and cycle for D) -
we may want to select not the best but the second best endpoint since it puts P and D on the same region (this is just a simple example to emphasize the idea).

[kfswain]

SchedulingResult is a simple data structure to communicate with the greater system.

I believe it is a strong answer to these questions:
What does a scheduler subsystem do? It selects endpoints (typically optimally, but thats up to user configuration).
How many? Up to user implementation.
What if i want to organize & sort my endpoints? Up to user implementation

is this type trying to encapsulate the multi-cycle results and map profileName -> it's result?

Doesn't have to map to profileName, just key(s) that a larger system understands

additionally, a cycle result might have more than []Endpoint as a return value e.g., if we have equally scored endpoints and additional endpoints with very close score - we might want to have bestScoreEndpoints and backupEndpoints

Absolutely right. Only SchedulingResult has the context to make those decisions. You may also want to run a scheduling cycle and do nothing but log its result.

Please note the Picker function returns []Endpoint

each cycle is making a local optimized endpoint selection. when receiving all cycles selection, we may want to do global optimization. for example in P/D (cycle for P and cycle for D) -
we may want to select not the best but the second best endpoint since it puts P and D on the same region (this is just a simple example to emphasize the idea).

Agreed, a custom ProfileSelection function would be your answer. Pack endpoints into the map[string][]Endpoint however you wish. Exactly how that data is consumed is intentionally out of scope of a Scheduler Subsystem.

[nirrozenbaum]

it seems like this definition is aiming to unify the Filter/Scorer/Picker interfaces (leaving Pre/Post out of discussions since we agreed they are not needed).
I think this is wrong, as every plugin has it's own unique functionality and the types should be different accordingly.
for example, this has caused to add Get/SetScore functions to Endpoint (although Score is not a property of Endpoint as wrote in other comment).

additionally, none of the above interfaces receive the request or request properties which is a requirement.
Some filters/scorers may depend on some request properties, like the target model or the prompt (e.g., scorer that scores endpoints if the endpoint ActiveModels has the request target model in it).

[kfswain]

the types should be different accordingly.

The types are different they are. The function signatures look similar yes. Because they all act on similar data.

additionally, none of the above interfaces receive the request or request properties which is a requirement.

CycleState is intentionally unstructured to allow for generic usage. Request data would go here.

[nirrozenbaum]

I think putting requests properties in CycleState misses the point.

Generally speaking, the scheduler has two main inputs when called: the request properties and the available endpoints along with their state (e.g., metrics).

decision of the output is based on those input fields.

CycleState by definition is a way to communicate between different plugins during the same cycle.

can we put request props in CycleState? sure, but that was not the intention.

we can also put the Endpoints in CycleState, can’t we?

[ahg-g]

We need to define request context type as a structured parameter, another parameter is the list of endpoints to select from. This is different from CycleState which indeed is mostly to enable plugins to store and communicate state across the extension points it implements.

[nirrozenbaum]

@ahg-g exactly.
my previous comment included a rhetorical question to try and make this point.

[nirrozenbaum]

function names are typically verbs or verb phrases, because they do something.
Pick/Select conveys an action: choosing the best endpoint.
Selection sounds like a data type or result, not an action. It implies an object, not behavior.
(e.g., Selection object could represents the result of a selection process).
so I suggest -
if we call the interface Picker - function should be called Pick.
if we call it Selector the function name should be called Select

[ahg-g]

I don't think we need this. A custom implementation of any part of the scheduler subsystem will need to be modeled as an extension point that a plugin implements. Interpreting the results can be done in a PostSchedule extension as proposed in https://github.com/kubernetes-sigs/gateway-api-inference-extension/pull/845/files#r2098163370

[nirrozenbaum]

++.
related also to other comment that none of these functions is needed, but only a Schedule function.
https://github.com/kubernetes-sigs/gateway-api-inference-extension/pull/845/files#r2096908076

[kfswain]

but only a Schedule function

Hard disagree here & we should reconcile this difference. PTAL at the other comments written.

But yes, I can rename the functions to pre/post schedule.

[nirrozenbaum]

as a side note comment on this discussion -
This PR aims to define the scheduler extension points, but I think a very related extension point that has to be addressed on the same time is the PostDispatch, since many plugins that needs to store some state that relates to the results will need to register to both.

this is not mentioned here anywhere, but is very relevant.
for example, until that is supported, we cannot get rid of the PostCycle since prefix plugin is using it (although it actually needs to use the PostDispatch)

[kfswain]

as a side note comment on this discussion - This PR aims to define the scheduler extension points, but I think a very related extension point that has to be addressed on the same time is the PostDispatch, since many plugins that needs to store some state that relates to the results will need to register to both.

this is not mentioned here anywhere, but is very relevant. for example, until that is supported, we cannot get rid of the PostCycle since prefix plugin is using it (although it actually needs to use the PostDispatch)

Yeah, this is where I expected the conversation to head.

I was originally considering anything post schedule to be out of scope of the scheduler. But if we want to allow each plugin to manage its own data, then I could see having different entrypoints for updating that data (postResponse is one you all feel strongly about, some might be fine with postDispatch.)

Part of me thinks that plugins shouldnt have their own data store. But that might couple the scheduler too much with another system.

[nirrozenbaum]

point (1) - I'd like to suggest using a more descriptive name. Additionally, this should be a plugin.
point (2) - I think it would be good to use a result struct rather than []Endpoint. we might find out that we need more than just an array of endpoints (e.g., maybe we need array of equally best endpoints and additional array of backup endpoints with very close lower scores?)

Suggested change

	// Scheduler is the implementation of a... scheduler.
	// The scheduler object is created at startup using the provided configuration.
	type Scheduler interface {
	// PreSchedule selects scheduling profiles through the implemented
	// logic, and returns:
	// - profiles - A subset of the registered scheduling profiles to be ran
	PreSchedule(request map[string]any, data scheduling.CycleState, results map[string][]Endpoint) map[string]SchedulingProfile
	// PostSchedule recieves the output of the result(s) of the scheduling cycle(s)
	// and makes sense of the data to be consumed by the calling system.
	// For example: suppose you have 2 profiles ShadowBoxing Profile & Production Profile.
	// PostSchedule would know to simply log the result of ShadowBoxing
	// profile, and do nothing else with it.
	PostSchedule(profileResults map[string][]Endpoint) SchedulingResult
	}
	type Result struct {
	Endpoints []Endpoint // set of selected endpoints
	}
	// MultiProfilePlugin defines the interface for handling multi profile scheduling.
	type MultiProfilePlugin interface {
	Plugin
	// selects the SchedulingProfiles to run while taking into consideration the request properties
	// and the previously executed SchedluderProfile cycles along with their results.
	SelectProfiles(request map[string]any, data scheduling.CycleState, executionResults map[string]*types.Result) map[string]*SchedulerProfile
	// ProcessProfileResults handles the outcome of each selected profile.
	// It may aggregate results, log test profile outputs, or apply custom logic.
	ProcessProfileResults(request map[string]any, results map[string]*types.Result) map[string]*types.Result
	}

[ahg-g]

Merging this so that we can continue to iterate over the proposal in @kfswain absence.

/lgtm
/approve

[k8s-ci-robot]

[APPROVALNOTIFIER] This PR is APPROVED

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

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,kfswain]

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

[nirrozenbaum]

/lgtm

[ahg-g]

/lgtm

[ahg-g]

Is there an interface that unifies SelectProfiles and ProcessProfileResults into a single interface? I think these (as @kfswain said, IIUC) are part of the same plugin and should always come as a single unit. They should not be implemented by different plugins.

Another way is to validate that a configuration always defines a SelectProfile and a ProcessResults (or whatever we call them) plugins. If we couple them into one, then the configuration API would need to deal with them as a single entity.

[nirrozenbaum]

Is there an interface that unifies SelectProfiles and ProcessProfileResults into a single interface? I think these (as @kfswain said, IIUC) are part of the same plugin and should always come as a single unit. They should not be implemented by different plugins.

Another way is to validate that a configuration always defines a SelectProfile and a ProcessResults (or whatever we call them) plugins. If we couple them into one, then the configuration API would need to deal with them as a single entity.

I'm fine with either of them.
I agree with @ahg-g it is sufficient to validate both exist, same as we should validate the picker plugin exists or fail on bootstrap.