|
| 1 | +/* |
| 2 | +Copyright 2026 The llm-d Authors. |
| 3 | +
|
| 4 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +you may not use this file except in compliance with the License. |
| 6 | +You may obtain a copy of the License at |
| 7 | +
|
| 8 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +
|
| 10 | +Unless required by applicable law or agreed to in writing, software |
| 11 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +See the License for the specific language governing permissions and |
| 14 | +limitations under the License. |
| 15 | +*/ |
| 16 | + |
| 17 | +package headerprofile |
| 18 | + |
| 19 | +import ( |
| 20 | + "context" |
| 21 | + "encoding/json" |
| 22 | + "fmt" |
| 23 | + "maps" |
| 24 | + "slices" |
| 25 | + "strings" |
| 26 | + |
| 27 | + "sigs.k8s.io/controller-runtime/pkg/log" |
| 28 | + |
| 29 | + fwkplugin "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" |
| 30 | + fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" |
| 31 | +) |
| 32 | + |
| 33 | +const ( |
| 34 | + // HeaderProfileHandlerType is the type of the HeaderProfileHandler. |
| 35 | + HeaderProfileHandlerType = "header-profile-handler" |
| 36 | + |
| 37 | + // defaultHeaderName is the request header read when parameters.HeaderName is empty. |
| 38 | + // Kept mixed-case for the README and for tests that use it as a mixed-case |
| 39 | + // constructor input; defaultHeaderNameLower is the form actually used as a header |
| 40 | + // key. |
| 41 | + defaultHeaderName = "EPP-Profile" |
| 42 | + |
| 43 | + // defaultProfileName is the scheduling profile run when parameters.DefaultProfile is |
| 44 | + // empty. |
| 45 | + defaultProfileName = "decode" |
| 46 | +) |
| 47 | + |
| 48 | +// defaultHeaderNameLower is defaultHeaderName normalized once at init, the same way |
| 49 | +// NewHeaderProfileHandler normalizes any configured headerName. |
| 50 | +var defaultHeaderNameLower = strings.ToLower(defaultHeaderName) |
| 51 | + |
| 52 | +// compile-time type assertion |
| 53 | +var _ fwksched.ProfileHandler = &HeaderProfileHandler{} |
| 54 | + |
| 55 | +// parameters configures the HeaderProfileHandler. |
| 56 | +type parameters struct { |
| 57 | + // HeaderName is the request header whose value names the scheduling profile to run. |
| 58 | + // Defaults to defaultHeaderName when empty. |
| 59 | + HeaderName string `json:"headerName"` |
| 60 | + // DefaultProfile is the scheduling profile to run when the header is missing or |
| 61 | + // blank. Defaults to defaultProfileName when empty. Useful for requests that never |
| 62 | + // carry the header at all - pass-through calls (e.g. /models) or a deployment that |
| 63 | + // doesn't disaggregate. |
| 64 | + DefaultProfile string `json:"defaultProfile"` |
| 65 | +} |
| 66 | + |
| 67 | +// HeaderProfileHandlerFactory defines the factory function for HeaderProfileHandler. |
| 68 | +func HeaderProfileHandlerFactory(name string, rawParameters *json.Decoder, _ fwkplugin.Handle) (fwkplugin.Plugin, error) { |
| 69 | + params := parameters{} |
| 70 | + if rawParameters != nil { |
| 71 | + if err := rawParameters.Decode(¶ms); err != nil { |
| 72 | + return nil, fmt.Errorf("failed to parse the parameters of the '%s' profile handler - %w", HeaderProfileHandlerType, err) |
| 73 | + } |
| 74 | + } |
| 75 | + |
| 76 | + return NewHeaderProfileHandler(params.HeaderName, params.DefaultProfile).WithName(name), nil |
| 77 | +} |
| 78 | + |
| 79 | +// NewHeaderProfileHandler trims and defaults both arguments, additionally |
| 80 | +// lowercasing headerName to match how the EPP stores ingested header keys |
| 81 | +// (pkg/epp/handlers/request.go); defaultProfile stays case-sensitive since it names a |
| 82 | +// schedulingProfiles entry. |
| 83 | +func NewHeaderProfileHandler(headerName, defaultProfile string) *HeaderProfileHandler { |
| 84 | + headerName = strings.ToLower(strings.TrimSpace(headerName)) |
| 85 | + if headerName == "" { |
| 86 | + headerName = defaultHeaderNameLower |
| 87 | + } |
| 88 | + |
| 89 | + defaultProfile = strings.TrimSpace(defaultProfile) |
| 90 | + if defaultProfile == "" { |
| 91 | + defaultProfile = defaultProfileName |
| 92 | + } |
| 93 | + |
| 94 | + return &HeaderProfileHandler{ |
| 95 | + typedName: fwkplugin.TypedName{Type: HeaderProfileHandlerType, Name: HeaderProfileHandlerType}, |
| 96 | + headerName: headerName, |
| 97 | + defaultProfile: defaultProfile, |
| 98 | + } |
| 99 | +} |
| 100 | + |
| 101 | +// HeaderProfileHandler runs exactly one scheduling profile per request, named by a |
| 102 | +// request header, falling back to the sole configured profile or to defaultProfile when |
| 103 | +// there's nothing to choose or the header is missing/blank; an unrecognized header value |
| 104 | +// is still an error. |
| 105 | +type HeaderProfileHandler struct { |
| 106 | + typedName fwkplugin.TypedName |
| 107 | + headerName string |
| 108 | + defaultProfile string |
| 109 | +} |
| 110 | + |
| 111 | +// TypedName returns the type and name tuple of this plugin instance. |
| 112 | +func (h *HeaderProfileHandler) TypedName() fwkplugin.TypedName { |
| 113 | + return h.typedName |
| 114 | +} |
| 115 | + |
| 116 | +// WithName sets the name of the profile handler. |
| 117 | +func (h *HeaderProfileHandler) WithName(name string) *HeaderProfileHandler { |
| 118 | + h.typedName.Name = name |
| 119 | + return h |
| 120 | +} |
| 121 | + |
| 122 | +// profileHeader returns the trimmed value of the profile-selecting header, or "" when |
| 123 | +// request is nil or the header is absent or blank. Trimming avoids surprising lookup |
| 124 | +// failures when the header carries incidental leading/trailing whitespace. |
| 125 | +func (h *HeaderProfileHandler) profileHeader(request *fwksched.InferenceRequest) string { |
| 126 | + if request == nil { |
| 127 | + return "" |
| 128 | + } |
| 129 | + return strings.TrimSpace(request.Headers[h.headerName]) |
| 130 | +} |
| 131 | + |
| 132 | +// noMatchError explains why no configured scheduling profile matches profileName, the |
| 133 | +// already-trimmed value of the profile-selecting header. |
| 134 | +func (h *HeaderProfileHandler) noMatchError(profileName string) error { |
| 135 | + if profileName == "" { |
| 136 | + return fmt.Errorf("header profile handler: missing %q header", h.headerName) |
| 137 | + } |
| 138 | + return fmt.Errorf("header profile handler: no scheduling profile configured for %q header value %q", h.headerName, profileName) |
| 139 | +} |
| 140 | + |
| 141 | +// defaultProfileNotConfiguredError reports defaultProfile itself missing from |
| 142 | +// schedulingProfiles - a config bug, distinct from noMatchError's per-request issue. |
| 143 | +func (h *HeaderProfileHandler) defaultProfileNotConfiguredError() error { |
| 144 | + return fmt.Errorf("header profile handler: defaultProfile %q is not a configured scheduling profile", h.defaultProfile) |
| 145 | +} |
| 146 | + |
| 147 | +// Pick implements fwksched.ProfileHandler.Pick; see README.md for behavior. |
| 148 | +func (h *HeaderProfileHandler) Pick(ctx context.Context, request *fwksched.InferenceRequest, profiles map[string]fwksched.SchedulerProfile, |
| 149 | + profileResults map[string]*fwksched.ProfileRunResult) map[string]fwksched.SchedulerProfile { |
| 150 | + if len(profileResults) > 0 { // the selected profile has already run |
| 151 | + return map[string]fwksched.SchedulerProfile{} |
| 152 | + } |
| 153 | + |
| 154 | + // With exactly one configured profile there is nothing to choose: always run |
| 155 | + // it, so a deployment scaled down to a single stage works without swapping profile |
| 156 | + // handlers or requiring every caller to send the header. |
| 157 | + if len(profiles) == 1 { |
| 158 | + for name, profile := range profiles { |
| 159 | + return map[string]fwksched.SchedulerProfile{name: profile} |
| 160 | + } |
| 161 | + } |
| 162 | + |
| 163 | + profileName := h.profileHeader(request) |
| 164 | + resolvedProfileName := profileName |
| 165 | + if resolvedProfileName == "" { |
| 166 | + resolvedProfileName = h.defaultProfile |
| 167 | + } |
| 168 | + |
| 169 | + profile, ok := profiles[resolvedProfileName] |
| 170 | + if !ok { |
| 171 | + err := h.noMatchError(profileName) |
| 172 | + if profileName == "" { |
| 173 | + err = h.defaultProfileNotConfiguredError() |
| 174 | + } |
| 175 | + registeredProfiles := slices.Sorted(maps.Keys(profiles)) |
| 176 | + // Logged here, not returned from ProcessResults (skipped when Pick returns empty), |
| 177 | + // at Info, not Error, since a bad or missing header is a client issue an operator |
| 178 | + // still needs to see by default. registeredProfiles is included to help spot typos |
| 179 | + // or config/header misalignment. |
| 180 | + // See README.md for why the client only gets a generic 429 instead of this reason. |
| 181 | + log.FromContext(ctx).Info("no scheduling profile selected for request", |
| 182 | + "error", err, "resolvedProfileName", resolvedProfileName, "registeredProfiles", registeredProfiles) |
| 183 | + return map[string]fwksched.SchedulerProfile{} |
| 184 | + } |
| 185 | + |
| 186 | + return map[string]fwksched.SchedulerProfile{resolvedProfileName: profile} |
| 187 | +} |
| 188 | + |
| 189 | +// ProcessResults handles the outcome of the single profile run selected by Pick. |
| 190 | +// It specifies in the SchedulingResult the key of the primary profile that should be |
| 191 | +// used to get the request's selected destination. |
| 192 | +func (h *HeaderProfileHandler) ProcessResults(_ context.Context, request *fwksched.InferenceRequest, |
| 193 | + profileResults map[string]*fwksched.ProfileRunResult) (*fwksched.SchedulingResult, error) { |
| 194 | + switch len(profileResults) { |
| 195 | + case 0: |
| 196 | + return nil, h.noMatchError(h.profileHeader(request)) |
| 197 | + case 1: |
| 198 | + // exactly one profile ran, handled below |
| 199 | + default: |
| 200 | + return nil, fmt.Errorf("header profile handler is intended to run a single profile per request, got %d", len(profileResults)) |
| 201 | + } |
| 202 | + |
| 203 | + var profileName string |
| 204 | + for name := range profileResults { |
| 205 | + profileName = name |
| 206 | + } |
| 207 | + |
| 208 | + if profileResults[profileName] == nil { // there was an error while running the profile |
| 209 | + return nil, fmt.Errorf("failed to run scheduler profile '%s'", profileName) |
| 210 | + } |
| 211 | + |
| 212 | + return &fwksched.SchedulingResult{ |
| 213 | + ProfileResults: profileResults, |
| 214 | + PrimaryProfileName: profileName, |
| 215 | + }, nil |
| 216 | +} |
0 commit comments