fix: reject invalid token tracker weights#1030
Conversation
Signed-off-by: pm-ju <pmdevops29@gmail.com>
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enhances token weight validation by ensuring values are finite and non-negative, specifically handling NaN and Infinite values. It introduces a helper function isValidTokenWeight, updates the configuration logic to use it, and adds comprehensive unit tests. Review feedback suggests improving log observability by including default values in warning messages and recommends deduplicating the validation logic by merging the new helper function with an existing identical one in the same package.
| inputWeight = w | ||
| } else { | ||
| klog.Warningf("Invalid FAIRNESS_INPUT_TOKEN_WEIGHT: %v, using default", err) | ||
| klog.Warningf("Invalid FAIRNESS_INPUT_TOKEN_WEIGHT: %q, using default", inputWeightStr) |
There was a problem hiding this comment.
For better observability and consistency with other configuration parsing logic in this file (such as createFairnessQueueConfig), the warning message should include the default value being used.
| klog.Warningf("Invalid FAIRNESS_INPUT_TOKEN_WEIGHT: %q, using default", inputWeightStr) | |
| klog.Warningf("Invalid FAIRNESS_INPUT_TOKEN_WEIGHT: %q, using default %v", inputWeightStr, defaultInputTokenWeight) |
| outputWeight = w | ||
| } else { | ||
| klog.Warningf("Invalid FAIRNESS_OUTPUT_TOKEN_WEIGHT: %v, using default", err) | ||
| klog.Warningf("Invalid FAIRNESS_OUTPUT_TOKEN_WEIGHT: %q, using default", outputWeightStr) |
There was a problem hiding this comment.
For better observability and consistency with other configuration parsing logic in this file, the warning message should include the default value being used.
| klog.Warningf("Invalid FAIRNESS_OUTPUT_TOKEN_WEIGHT: %q, using default", outputWeightStr) | |
| klog.Warningf("Invalid FAIRNESS_OUTPUT_TOKEN_WEIGHT: %q, using default %v", outputWeightStr, defaultOutputTokenWeight) |
| func isValidTokenWeight(value float64) bool { | ||
| return !math.IsNaN(value) && !math.IsInf(value, 0) && value >= 0 | ||
| } |
There was a problem hiding this comment.
The function isValidTokenWeight is identical to isValidFairnessWeight already defined in pkg/kthena-router/datastore/store.go. Since both functions perform the same validation (checking for finite, non-negative values) and reside in the same package, they should be unified into a single helper function to improve maintainability and ensure consistency across the package.
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Tightens validation for token-tracker weight configuration so non-finite values (NaN/Inf) no longer poison fairness/token accounting.
Changes:
- Add
isValidTokenWeightto require weights be finite and non-negative. - Validate env-provided weights (
FAIRNESS_INPUT_TOKEN_WEIGHT,FAIRNESS_OUTPUT_TOKEN_WEIGHT) before applying; otherwise fall back to defaults. - Extend tests to cover NaN/Inf and mixed-validity scenarios.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| pkg/kthena-router/datastore/token_tracker.go | Adds shared finite/non-negative validation and applies it in WithTokenWeights. |
| pkg/kthena-router/datastore/store.go | Validates env-derived weights with the new helper and adjusts warning logs. |
| pkg/kthena-router/datastore/token_tracker_test.go | Adds unit tests covering NaN/Inf weight handling in options/config. |
| pkg/kthena-router/datastore/store_test.go | Adds integration-style test ensuring env invalid weights fall back to defaults. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| inputWeight = w | ||
| } else { | ||
| klog.Warningf("Invalid FAIRNESS_INPUT_TOKEN_WEIGHT: %v, using default", err) | ||
| klog.Warningf("Invalid FAIRNESS_INPUT_TOKEN_WEIGHT: %q, using default", inputWeightStr) |
| outputWeight = w | ||
| } else { | ||
| klog.Warningf("Invalid FAIRNESS_OUTPUT_TOKEN_WEIGHT: %v, using default", err) | ||
| klog.Warningf("Invalid FAIRNESS_OUTPUT_TOKEN_WEIGHT: %q, using default", outputWeightStr) |
| t.Setenv("FAIRNESS_INPUT_TOKEN_WEIGHT", "NaN") | ||
| t.Setenv("FAIRNESS_OUTPUT_TOKEN_WEIGHT", strconv.FormatFloat(math.Inf(1), 'f', -1, 64)) | ||
|
|
||
| tracker := createTokenTracker().(*InMemorySlidingWindowTokenTracker) |
| t.Setenv("FAIRNESS_INPUT_TOKEN_WEIGHT", "1.5") | ||
| t.Setenv("FAIRNESS_OUTPUT_TOKEN_WEIGHT", "-2") | ||
|
|
||
| tracker = createTokenTracker().(*InMemorySlidingWindowTokenTracker) |
/kind bug
Validate token tracker weights before applying them.
Previously
FAIRNESS_INPUT_TOKEN_WEIGHT,FAIRNESS_OUTPUT_TOKEN_WEIGHT, andWithTokenWeightsrejected negative values but still accepted non-finite values likeNaNandInf. That could poison token accounting and fairness/rate-limit calculations with non-finite totals.This change requires token weights to be finite and non-negative. Invalid env values fall back to defaults, while valid values are still applied.
Verification:
go test ./pkg/kthena-router/datastore