|
| 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 costguard implements a cost-minimising scorer for the ModelSelector |
| 18 | +// framework. It routes inference to the model with the lowest observed cost |
| 19 | +// via an epsilon-Greedy Multi-Arm Bandit, using a risk-aware rank that |
| 20 | +// combines each model's body cost (trimmed mean) and Conditional Tail |
| 21 | +// Expectation of the cost distribution. See the package README for the full |
| 22 | +// algorithm and configuration reference. |
| 23 | +package costguard |
| 24 | + |
| 25 | +import ( |
| 26 | + "context" |
| 27 | + "encoding/json" |
| 28 | + "fmt" |
| 29 | + "math" |
| 30 | + "time" |
| 31 | + |
| 32 | + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/datalayer" |
| 33 | + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/modelselector" |
| 34 | + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin" |
| 35 | + "github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling" |
| 36 | +) |
| 37 | + |
| 38 | +const ( |
| 39 | + // PluginType is the identifier used when registering this scorer. |
| 40 | + PluginType = "costguard" |
| 41 | + |
| 42 | + // z95 is the standard normal quantile for a 95% confidence interval, used |
| 43 | + // in sampleThreshold to derive the per-model minimum sample size from the |
| 44 | + // configured percentile margin of error. |
| 45 | + z95 = 1.96 |
| 46 | + |
| 47 | + // neutralScore is returned for any model the scorer has no strong opinion |
| 48 | + // on (missing / empty / under-explored cost data, degenerate sigmoid |
| 49 | + // inputs). It composes with other scorers in the pipeline without pushing |
| 50 | + // selection either way. |
| 51 | + neutralScore = 0.5 |
| 52 | + |
| 53 | + defaultEpsilon = 0.1 |
| 54 | + defaultAlpha = 0.95 |
| 55 | + defaultLambda = 1.0 |
| 56 | + defaultWindowDuration = "2h" |
| 57 | + defaultPercentileMarginError = 0.03 |
| 58 | +) |
| 59 | + |
| 60 | +// compile-time interface assertion |
| 61 | +var _ modelselector.Scorer = &CostGuardScorer{} |
| 62 | + |
| 63 | +// Config defines the JSON configuration for the plugin. All fields have |
| 64 | +// sensible defaults, so an empty config is valid. |
| 65 | +type Config struct { |
| 66 | + // Epsilon is the probability of random exploration on each request. Must |
| 67 | + // be in [0, 1]. Defaults to 0.1. Setting Epsilon to 1.0 forces full random |
| 68 | + // selection (with replacement) on every request. |
| 69 | + Epsilon float64 `json:"epsilon"` |
| 70 | + |
| 71 | + // Alpha is the quantile that separates the body of the cost distribution |
| 72 | + // from the tail. Must be in (0, 1). Defaults to 0.95. |
| 73 | + Alpha float64 `json:"alpha"` |
| 74 | + |
| 75 | + // Lambda is the penalty weight applied to the tail cost contribution |
| 76 | + // (Conditional Tail Expectation). Must be >= 0. Defaults to 1.0. |
| 77 | + Lambda float64 `json:"lambda"` |
| 78 | + |
| 79 | + // WindowDuration is the epoch length. Parsed as a Go duration string |
| 80 | + // (e.g. "2h", "30m"). Must be > 0. Defaults to "2h". |
| 81 | + // |
| 82 | + // TODO: costguard-epoch: consume in the final PR for costguard epoch. |
| 83 | + // Parsed and validated today, but not consulted at scoring time. |
| 84 | + // The requestcostmetadata extractor plugin will own the window |
| 85 | + // in the concluding PR for CostGuard. |
| 86 | + WindowDuration string `json:"windowDuration"` |
| 87 | + |
| 88 | + // PercentileMarginError is the error around the alpha-quantile |
| 89 | + // of the observed cost distribution. Must be in (0, 1). |
| 90 | + // Defaults to 0.03. Smaller values require quadratically more |
| 91 | + // samples to consider a model sufficiently explored. |
| 92 | + PercentileMarginError float64 `json:"percentileMarginError"` |
| 93 | +} |
| 94 | + |
| 95 | +// CostGuardScorer will consume per-model cost samples via the datalayer |
| 96 | +// (published by the requestcostmetadata extractor) and score candidate models |
| 97 | +// to minimize typical cost and risk of the high cost in the tail. |
| 98 | +// See README for how it works. |
| 99 | +// |
| 100 | +// This PR creates only the configuration surface, plugin registration, and a stub |
| 101 | +// Score that returns a neutral score for every model. The exploit and explore |
| 102 | +// branches follow in subsequent PRs; the requestcostmetadata extractor takes |
| 103 | +// ownership of the scorer epoch window in the final PR. |
| 104 | +type CostGuardScorer struct { |
| 105 | + typedName plugin.TypedName |
| 106 | + epsilon float64 |
| 107 | + alpha float64 |
| 108 | + lambda float64 |
| 109 | + sampleThreshold uint64 |
| 110 | + windowDuration time.Duration |
| 111 | +} |
| 112 | + |
| 113 | +// ScorerFactory creates a CostGuardScorer from the given raw parameters. |
| 114 | +// It validates every field of Config and precomputes the sample threshold so |
| 115 | +// Score can rank candidate models with high statistical fidelity. |
| 116 | +func ScorerFactory(name string, rawParameters json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) { |
| 117 | + config := Config{ |
| 118 | + Epsilon: defaultEpsilon, |
| 119 | + Alpha: defaultAlpha, |
| 120 | + Lambda: defaultLambda, |
| 121 | + WindowDuration: defaultWindowDuration, |
| 122 | + PercentileMarginError: defaultPercentileMarginError, |
| 123 | + } |
| 124 | + if len(rawParameters) > 0 { |
| 125 | + if err := json.Unmarshal(rawParameters, &config); err != nil { |
| 126 | + return nil, fmt.Errorf("costguard %q: failed to parse parameters: %w", name, err) |
| 127 | + } |
| 128 | + } |
| 129 | + |
| 130 | + if config.Epsilon < 0 || config.Epsilon > 1 { |
| 131 | + return nil, fmt.Errorf("costguard %q: epsilon must be in [0, 1], got %f", name, config.Epsilon) |
| 132 | + } |
| 133 | + if config.Alpha <= 0 || config.Alpha >= 1 { |
| 134 | + return nil, fmt.Errorf("costguard %q: alpha must be in (0, 1), got %f", name, config.Alpha) |
| 135 | + } |
| 136 | + if config.Lambda < 0 { |
| 137 | + return nil, fmt.Errorf("costguard %q: lambda must be >= 0, got %f", name, config.Lambda) |
| 138 | + } |
| 139 | + if config.PercentileMarginError <= 0 || config.PercentileMarginError >= 1 { |
| 140 | + return nil, fmt.Errorf("costguard %q: percentileMarginError must be in (0, 1), got %f", name, config.PercentileMarginError) |
| 141 | + } |
| 142 | + windowDuration, err := time.ParseDuration(config.WindowDuration) |
| 143 | + if err != nil { |
| 144 | + return nil, fmt.Errorf("costguard %q: invalid windowDuration %q: %w", name, config.WindowDuration, err) |
| 145 | + } |
| 146 | + if windowDuration <= 0 { |
| 147 | + return nil, fmt.Errorf("costguard %q: windowDuration must be > 0, got %s", name, windowDuration) |
| 148 | + } |
| 149 | + |
| 150 | + return NewCostGuardScorer( |
| 151 | + config.Epsilon, |
| 152 | + config.Alpha, |
| 153 | + config.Lambda, |
| 154 | + windowDuration, |
| 155 | + config.PercentileMarginError, |
| 156 | + ).WithName(name), nil |
| 157 | +} |
| 158 | + |
| 159 | +// NewCostGuardScorer constructs a scorer with the given parameters. |
| 160 | +// It performs no range checks; callers are responsible for passing values in |
| 161 | +// the ranges documented on Config. |
| 162 | +// ScorerFactory is the intended path for validated construction from raw JSON parameters. |
| 163 | +func NewCostGuardScorer(epsilon, alpha, lambda float64, windowDuration time.Duration, percentileMarginError float64) *CostGuardScorer { |
| 164 | + return &CostGuardScorer{ |
| 165 | + typedName: plugin.TypedName{Type: PluginType, Name: PluginType}, |
| 166 | + epsilon: epsilon, |
| 167 | + alpha: alpha, |
| 168 | + lambda: lambda, |
| 169 | + sampleThreshold: sampleThreshold(alpha, percentileMarginError), |
| 170 | + windowDuration: windowDuration, |
| 171 | + } |
| 172 | +} |
| 173 | + |
| 174 | +// TypedName returns the typed name of the plugin instance. |
| 175 | +func (s *CostGuardScorer) TypedName() plugin.TypedName { |
| 176 | + return s.typedName |
| 177 | +} |
| 178 | + |
| 179 | +// WithName sets the instance name. |
| 180 | +func (s *CostGuardScorer) WithName(name string) *CostGuardScorer { |
| 181 | + s.typedName.Name = name |
| 182 | + return s |
| 183 | +} |
| 184 | + |
| 185 | +// Score returns the neutral score for every model. The exploit and explore |
| 186 | +// branches follow in subsequent PRs. |
| 187 | +func (s *CostGuardScorer) Score(_ context.Context, _ *plugin.CycleState, _ *requesthandling.InferenceRequest, models []datalayer.Model) map[datalayer.Model]float64 { |
| 188 | + scores := make(map[datalayer.Model]float64, len(models)) |
| 189 | + for _, m := range models { |
| 190 | + scores[m] = neutralScore |
| 191 | + } |
| 192 | + return scores |
| 193 | +} |
| 194 | + |
| 195 | +// sampleThreshold returns the minimum sample count for a model's alpha-quantile |
| 196 | +// estimate to have a two-sided 95% CI of half-width w (Wald formula). |
| 197 | +func sampleThreshold(alpha, w float64) uint64 { |
| 198 | + return uint64(math.Ceil(z95 * z95 * alpha * (1 - alpha) / (w * w))) |
| 199 | +} |
0 commit comments