Skip to content

passing headers to scheduler plugins #775

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 2, 2025
Merged

passing headers to scheduler plugins #775

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
581a6a8
passing headers to scheduler plugins
nirrozenbaum May 2, 2025
7b8a83e
addressed code review comments
nirrozenbaum May 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 4 additions & 6 deletions pkg/epp/handlers/request.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,11 @@ import (
)

// HandleRequestBody always returns the requestContext even in the error case, as the request context is used in error handling.
func (s *StreamingServer) HandleRequestBody(
ctx context.Context,
reqCtx *RequestContext,
) (*RequestContext, error) {
var requestBodyBytes []byte
func (s *StreamingServer) HandleRequestBody(ctx context.Context, reqCtx *RequestContext) (*RequestContext, error) {
logger := log.FromContext(ctx)
requestBodyMap := reqCtx.Request.Body

var requestBodyBytes []byte
requestBodyMap := reqCtx.Request.Body
// Resolve target models.
model, ok := requestBodyMap["model"].(string)
if !ok {
Expand Down Expand Up @@ -70,6 +67,7 @@ func (s *StreamingServer) HandleRequestBody(
ResolvedTargetModel: modelName,
Critical: modelObj.Spec.Criticality != nil && *modelObj.Spec.Criticality == v1alpha2.Critical,
Prompt: prompt,
Headers: reqCtx.Request.Headers,
}
logger.V(logutil.DEBUG).Info("LLM request assembled", "request", llmReq)

Expand Down
2 changes: 1 addition & 1 deletion pkg/epp/scheduling/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ package scheduling

import "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"

// SchedulerConfig creates a new SchedulerConfig object with the given plugins.
// NewSchedulerConfig creates a new SchedulerConfig object with the given plugins.
func NewSchedulerConfig(preSchedulePlugins []plugins.PreSchedule, filters []plugins.Filter, scorers map[plugins.Scorer]int,
picker plugins.Picker, postSchedulePlugins []plugins.PostSchedule) *SchedulerConfig {
return &SchedulerConfig{
Expand Down
16 changes: 10 additions & 6 deletions pkg/epp/scheduling/types/types.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,17 +28,21 @@ import (

// LLMRequest is a structured representation of the fields we parse out of the LLMRequest body.
type LLMRequest struct {
// Model is the name of the model that the user specified in the request body.
Model string
// Target models is a map of target model name to weight.
TargetModels map[string]int
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed TargetModels from LLMRequest, which is never set and is always empty.
LLM request holds the fields that are passed to scheduler. the scheduler doesn't care about the target models, but only about the ResolvedTargetModel (traffic splitting is done in request.go. when LLMRequest is passed to scheduler the target model has been resolved already).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not this PR, but I was wondering if we should even keep the LLMRequest object. I worry that reqCtx is this monolithic super-object. But also, its been very useful to have, and its not really expensive to pass around since you'd just be passing a pointer.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in terms of best practices, it would probably be best to expose to each layer only the things it needs.
one mono super-object may be very convenient, but it opens the door for updating fields that the layer shouldn't.
we might be able to keep the mono object but work with layer interfaces that expose to each layer only the fields it needs.
anyway, as you said, this is not in scope of this PR

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, I was noodling on how best to deal with this since I commented, I think I have some ideas. WIP tho

Prompt string
// Resolved target model is the final target model after traffic split.
// ResolvedTargetModel is the final target model after traffic split.
ResolvedTargetModel string
Critical bool
// Critical is a boolean that specifies if a request is critical or not.
Critical bool
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should make this an int since we already have 3 criticality levels, and likely will be extended to have more.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In current code scheduler changes its algorithm based on the question of whether the request is critical or not.
I agree that this field should be changed, but this requires adaptation in more places other than this field.
should be done in a separate PR dedicated for this point (not in this PR scope).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

++

// Prompt is the prompt that was sent in the request body.
Prompt string
// Headers is a map of the request headers.
Headers map[string]string
}

func (r *LLMRequest) String() string {
return fmt.Sprintf("Model: %s, TargetModels: %v, ResolvedTargetModel: %s, Critical: %t, PromptLength: %v", r.Model, r.TargetModels, r.ResolvedTargetModel, r.Critical, len(r.Prompt))
return fmt.Sprintf("Model: %s, ResolvedTargetModel: %s, Critical: %t, PromptLength: %d, Headers: %v",
r.Model, r.ResolvedTargetModel, r.Critical, len(r.Prompt), r.Headers)
}

type Pod interface {
Expand Down