Skip to content

fix(ratelimit): isolate rate limiters per ModelRoute #1171

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

Closed
nXtCyberNet wants to merge 12 commits into from
Closed

fix(ratelimit): isolate rate limiters per ModelRoute #1171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
30f4c0b
updated the route namestruct from model to namespace/modelroutename
nXtCyberNet May 29, 2026
e101323
changed the key in main flow
nXtCyberNet Jun 3, 2026
3eeb9ec
added err log
nXtCyberNet Jun 3, 2026
f908438
solved the issue
nXtCyberNet Jun 3, 2026
5a2b611
.
nXtCyberNet Jun 3, 2026
8a28c36
changed ratelimit to doloadbalancer
nXtCyberNet Jun 5, 2026
b0a6445
ssolved ci
nXtCyberNet Jun 5, 2026
2031650
resolve conflict
nXtCyberNet Jun 6, 2026
e5a0077
remove the test
nXtCyberNet Jun 12, 2026
63fe482
.
nXtCyberNet Jun 18, 2026
ac11592
saving ratelimiterkey to gin context
nXtCyberNet Jun 22, 2026
d78422f
converted set to struct
nXtCyberNet Jun 23, 2026
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
30 changes: 15 additions & 15 deletions pkg/kthena-router/filters/ratelimit/ratelimit.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ func NewTokenRateLimiter() *TokenRateLimiter {
}

// RateLimit checks if the request is within rate limits for both input and output tokens
func (r *TokenRateLimiter) RateLimit(model, prompt string) error {
func (r *TokenRateLimiter) RateLimit(limiterKey, prompt string) error {
// Estimate input tokens
tokens, err := r.tokenizer.CalculateTokenNum(prompt)
if err != nil {
Expand All @@ -107,8 +107,8 @@ func (r *TokenRateLimiter) RateLimit(model, prompt string) error {
}

r.mutex.RLock()
inputLimiter, hasInputLimit := r.inputLimiter[model]
outputLimiter, hasOutputLimit := r.outputLimiter[model]
inputLimiter, hasInputLimit := r.inputLimiter[limiterKey]
outputLimiter, hasOutputLimit := r.outputLimiter[limiterKey]
r.mutex.RUnlock()

// Check input token rate limit
Expand All @@ -126,9 +126,9 @@ func (r *TokenRateLimiter) RateLimit(model, prompt string) error {
}

// RecordOutputTokens records the actual output tokens consumed after response generation
func (r *TokenRateLimiter) RecordOutputTokens(model string, tokenCount int) {
func (r *TokenRateLimiter) RecordOutputTokens(limiterKey string, tokenCount int) {
r.mutex.RLock()
outputLimiter, exists := r.outputLimiter[model]
outputLimiter, exists := r.outputLimiter[limiterKey]
r.mutex.RUnlock()

if exists {
Expand All @@ -137,7 +137,7 @@ func (r *TokenRateLimiter) RecordOutputTokens(model string, tokenCount int) {
}

// AddOrUpdateLimiter adds or updates rate limiter for a model
func (r *TokenRateLimiter) AddOrUpdateLimiter(model string, ratelimit *networkingv1alpha1.RateLimit) error {
func (r *TokenRateLimiter) AddOrUpdateLimiter(limiterKey string, ratelimit *networkingv1alpha1.RateLimit) error {

@LiZhenCheng9527 LiZhenCheng9527 Jun 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

To be honest, I don’t really understand what ‘limiterKey’ refers to. Changing the name of this formal parameter isn’t a good idea.

@nXtCyberNet nXtCyberNet Jun 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Hi @LiZhenCheng9527 ,
So currently the ratelimit was defined by only the model which cause there is no namespace level isolation that cause the follow ratelimit bug :

ModelRoute-1 (namespace-a, model: llama-70b, tokenLimit: 100 tokens/min) 
  → Redis key: "llama-70b"

ModelRoute-2 (namespace-b, model: llama-70b, tokenLimit: 50 tokens/min)
  → Redis key: "llama-70b"  ← SAME KEY ❌

Result: ModelRoute-2's update OVERWRITES ModelRoute-1's token bucket
        Both namespaces now share the SAME token bucket with 50 tokens/min
        namespace-a requests get rate-limited incorrectly ❌

Which is inconsistent and breaks ratelimit if multiple teams were using the same model with different configs ,

So after changing the model name key to limiterkey - namespace/routename , since the routename is crd , so it's itself a unique name , so this will also allow per route rate limit isolation too
Like this

ModelRoute-1 (namespace-a, model: llama-70b, tokenLimit: 100 tokens/min)
  → Redis key: "namespace-a/modelroute-1"
  → Independent token bucket with 100 tokens/min

ModelRoute-2 (namespace-b, model: llama-70b, tokenLimit: 50 tokens/min)
  → Redis key: "namespace-b/modelroute-2"
  → Independent token bucket with 50 tokens/min ✅

Result: Each ModelRoute has its own isolated token bucket
namespace-a can use 100 tokens/min, namespace-b uses 50 tokens/min

,

Also it act same even if multiple routes were created in a same namespace

And the limitkey is a code level refactoring so also there is no userside level change , WDYT?

Comment thread
nXtCyberNet marked this conversation as resolved.
Comment on lines 139 to +140
Comment on lines 139 to +140
Comment on lines 139 to +140
r.mutex.Lock()
defer r.mutex.Unlock()
Comment thread
hzxuzhonghu marked this conversation as resolved.

Expand All @@ -161,21 +161,21 @@ func (r *TokenRateLimiter) AddOrUpdateLimiter(model string, ratelimit *networkin

// Create global rate limiters
if ratelimit.InputTokensPerUnit != nil {
r.inputLimiter[model] = NewGlobalRateLimiter(
r.inputLimiter[limiterKey] = NewGlobalRateLimiter(
r.redisClient,
"kthena:ratelimit",
model,
limiterKey,
"input",
*ratelimit.InputTokensPerUnit,
ratelimit.Unit,
)
Comment thread
hzxuzhonghu marked this conversation as resolved.
}

if ratelimit.OutputTokensPerUnit != nil {
r.outputLimiter[model] = NewGlobalRateLimiter(
r.outputLimiter[limiterKey] = NewGlobalRateLimiter(
r.redisClient,
"kthena:ratelimit",
model,
limiterKey,
"output",
*ratelimit.OutputTokensPerUnit,
ratelimit.Unit,
Expand All @@ -186,14 +186,14 @@ func (r *TokenRateLimiter) AddOrUpdateLimiter(model string, ratelimit *networkin
duration := getTimeUnitDuration(ratelimit.Unit)

if ratelimit.InputTokensPerUnit != nil {
r.inputLimiter[model] = NewLocalLimiter(
r.inputLimiter[limiterKey] = NewLocalLimiter(
rate.Limit(float64(*ratelimit.InputTokensPerUnit)/duration.Seconds()),
int(*ratelimit.InputTokensPerUnit),
)
}

if ratelimit.OutputTokensPerUnit != nil {
r.outputLimiter[model] = NewLocalLimiter(
r.outputLimiter[limiterKey] = NewLocalLimiter(
rate.Limit(float64(*ratelimit.OutputTokensPerUnit)/duration.Seconds()),
int(*ratelimit.OutputTokensPerUnit),
)
Expand All @@ -204,12 +204,12 @@ func (r *TokenRateLimiter) AddOrUpdateLimiter(model string, ratelimit *networkin
}

// DeleteLimiter deletes rate limiter for a model
func (r *TokenRateLimiter) DeleteLimiter(model string) {
func (r *TokenRateLimiter) DeleteLimiter(limiterKey string) {
Comment thread
nXtCyberNet marked this conversation as resolved.
Comment thread
nXtCyberNet marked this conversation as resolved.
Comment on lines 206 to +207
Comment on lines 206 to +207
Comment on lines 206 to +207
r.mutex.Lock()
defer r.mutex.Unlock()

delete(r.inputLimiter, model)
delete(r.outputLimiter, model)
delete(r.inputLimiter, limiterKey)
delete(r.outputLimiter, limiterKey)
}

func getTimeUnitDuration(unit networkingv1alpha1.RateLimitUnit) time.Duration {
Expand Down
88 changes: 47 additions & 41 deletions pkg/kthena-router/filters/ratelimit/ratelimit_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ limitations under the License.
package ratelimit

import (
"fmt"
"testing"
"time"

Expand All @@ -25,26 +26,26 @@ import (

func TestTokenRateLimiter_Basic(t *testing.T) {
rl := NewTokenRateLimiter()
model := "test-model"
limiterKey := testLimiterKey("default", "test-route")
prompt := "hello world" // 3 tokens
tokens := uint32(10)
unit := networkingv1alpha1.Second

rl.AddOrUpdateLimiter(model, &networkingv1alpha1.RateLimit{
rl.AddOrUpdateLimiter(limiterKey, &networkingv1alpha1.RateLimit{
InputTokensPerUnit: &tokens,
Unit: unit,
})

// Should allow up to 10 tokens immediately
for i := 0; i < 3; i++ {
err := rl.RateLimit(model, prompt)
err := rl.RateLimit(limiterKey, prompt)
if err != nil {
t.Fatalf("unexpected error on allowed request: %v, %d", err, i)
}
}

// 4th request should be rate limited
err := rl.RateLimit(model, prompt)
err := rl.RateLimit(limiterKey, prompt)
if err == nil {
t.Fatalf("expected rate limit error, got nil")
}
Expand All @@ -56,33 +57,33 @@ func TestTokenRateLimiter_Basic(t *testing.T) {
func TestTokenRateLimiter_NoLimiter(t *testing.T) {
rl := NewTokenRateLimiter()
// No limiter added, should always allow
err := rl.RateLimit("unknown-model", "test")
err := rl.RateLimit(testLimiterKey("unknown", "route"), "test")
if err != nil {
t.Fatalf("expected nil error for unknown model, got %v", err)
}
}

func TestTokenRateLimiter_ResetAfterTime(t *testing.T) {
rl := NewTokenRateLimiter()
model := "test-model"
limiterKey := testLimiterKey("default", "test-route")
prompt := "hello world"
tokens := uint32(10)
unit := networkingv1alpha1.Second

rl.AddOrUpdateLimiter(model, &networkingv1alpha1.RateLimit{
rl.AddOrUpdateLimiter(limiterKey, &networkingv1alpha1.RateLimit{
InputTokensPerUnit: &tokens,
Unit: unit,
})

// Use up tokens
for i := 0; i < 3; i++ {
err := rl.RateLimit(model, prompt)
err := rl.RateLimit(limiterKey, prompt)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
// Should be rate limited now
err := rl.RateLimit(model, prompt)
err := rl.RateLimit(limiterKey, prompt)
if err == nil {
t.Fatalf("expected rate limit error, got nil")
}
Expand All @@ -92,56 +93,56 @@ func TestTokenRateLimiter_ResetAfterTime(t *testing.T) {

// Wait for refill
time.Sleep(1100 * time.Millisecond)
err = rl.RateLimit(model, prompt)
err = rl.RateLimit(limiterKey, prompt)
if err != nil {
t.Fatalf("expected nil after refill, got %v", err)
}
}

func TestTokenRateLimiter_OutputTokenRecording(t *testing.T) {
rl := NewTokenRateLimiter()
model := "test-model"
limiterKey := testLimiterKey("default", "test-route")
tokens := uint32(10)
unit := networkingv1alpha1.Second

rl.AddOrUpdateLimiter(model, &networkingv1alpha1.RateLimit{
rl.AddOrUpdateLimiter(limiterKey, &networkingv1alpha1.RateLimit{
OutputTokensPerUnit: &tokens,
Unit: unit,
})

// Record output tokens - this should not block/error
rl.RecordOutputTokens(model, 5)
rl.RecordOutputTokens(model, 3)
rl.RecordOutputTokens(model, 2) // Total: 10 tokens consumed
rl.RecordOutputTokens(limiterKey, 5)
rl.RecordOutputTokens(limiterKey, 3)
rl.RecordOutputTokens(limiterKey, 2) // Total: 10 tokens consumed

// Recording more tokens should still work (just consumes from the bucket)
rl.RecordOutputTokens(model, 1)
rl.RecordOutputTokens(limiterKey, 1)
}

func TestTokenRateLimiter_CombinedInputOutput(t *testing.T) {
rl := NewTokenRateLimiter()
model := "test-model"
limiterKey := testLimiterKey("default", "test-route")
prompt := "hello world hello world" // Should be ~6 tokens
inputTokens := uint32(8) // Allow only one request (6 tokens < 8, but two requests = 12 > 8)
outputTokens := uint32(10) // Allow output recording
unit := networkingv1alpha1.Second

rl.AddOrUpdateLimiter(model, &networkingv1alpha1.RateLimit{
rl.AddOrUpdateLimiter(limiterKey, &networkingv1alpha1.RateLimit{
InputTokensPerUnit: &inputTokens,
OutputTokensPerUnit: &outputTokens,
Unit: unit,
})

// First request should be allowed
err := rl.RateLimit(model, prompt)
err := rl.RateLimit(limiterKey, prompt)
if err != nil {
t.Fatalf("unexpected error on first request: %v", err)
}
// Record output tokens used
rl.RecordOutputTokens(model, 2)
rl.RecordOutputTokens(limiterKey, 2)

// Second request should be rate limited due to input token exhaustion
err = rl.RateLimit(model, prompt)
err = rl.RateLimit(limiterKey, prompt)
if err == nil {
t.Fatalf("expected rate limit error after exhausting input tokens")
}
Expand All @@ -153,72 +154,72 @@ func TestTokenRateLimiter_CombinedInputOutput(t *testing.T) {
func TestTokenRateLimiter_OutputNoLimiter(t *testing.T) {
rl := NewTokenRateLimiter()
// No limiter added, should not error when recording output tokens
rl.RecordOutputTokens("unknown-model", 100)
rl.RecordOutputTokens(testLimiterKey("unknown", "route"), 100)
// RecordOutputTokens doesn't return error, just silently does nothing
}

func TestTokenRateLimiter_DeleteLimiter(t *testing.T) {
rl := NewTokenRateLimiter()
model := "test-model"
limiterKey := testLimiterKey("default", "test-route")
inputTokens := uint32(3) // Very restrictive
outputTokens := uint32(5)
unit := networkingv1alpha1.Second

rl.AddOrUpdateLimiter(model, &networkingv1alpha1.RateLimit{
rl.AddOrUpdateLimiter(limiterKey, &networkingv1alpha1.RateLimit{
InputTokensPerUnit: &inputTokens,
OutputTokensPerUnit: &outputTokens,
Unit: unit,
})

// Verify limiter exists and restricts
err := rl.RateLimit(model, "hello world") // ~3 tokens
err := rl.RateLimit(limiterKey, "hello world") // ~3 tokens
if err != nil {
t.Fatalf("first request should be allowed: %v", err)
}

err = rl.RateLimit(model, "hello world") // Should be rate limited
err = rl.RateLimit(limiterKey, "hello world") // Should be rate limited
if err == nil {
t.Fatalf("expected rate limit error")
}

// Delete limiters
rl.DeleteLimiter(model)
rl.DeleteLimiter(limiterKey)

// Should now be unrestricted
for i := 0; i < 10; i++ {
err = rl.RateLimit(model, "hello world")
err = rl.RateLimit(limiterKey, "hello world")
if err != nil {
t.Fatalf("expected nil after deletion, got %v", err)
}
}

// Recording output tokens should work without error
rl.RecordOutputTokens(model, 100)
rl.RecordOutputTokens(limiterKey, 100)
}

func TestTokenRateLimiter_OutputRateLimit(t *testing.T) {
rl := NewTokenRateLimiter()
model := "test-model"
limiterKey := testLimiterKey("default", "test-route")
prompt := "hello world"
outputTokens := uint32(5) // Very low limit
unit := networkingv1alpha1.Second

rl.AddOrUpdateLimiter(model, &networkingv1alpha1.RateLimit{
rl.AddOrUpdateLimiter(limiterKey, &networkingv1alpha1.RateLimit{
OutputTokensPerUnit: &outputTokens,
Unit: unit,
})

// First request should be allowed (has 5 tokens available)
err := rl.RateLimit(model, prompt)
err := rl.RateLimit(limiterKey, prompt)
if err != nil {
t.Fatalf("first request should be allowed: %v", err)
}

// Consume most tokens
rl.RecordOutputTokens(model, 5)
rl.RecordOutputTokens(limiterKey, 5)

// Next request should be blocked due to insufficient output tokens
err = rl.RateLimit(model, prompt)
err = rl.RateLimit(limiterKey, prompt)
if err == nil {
t.Fatalf("expected output rate limit error")
}
Expand All @@ -229,19 +230,19 @@ func TestTokenRateLimiter_OutputRateLimit(t *testing.T) {

func TestTokenRateLimiter_InputAndOutputErrors(t *testing.T) {
rl := NewTokenRateLimiter()
model := "test-model"
longPrompt := "hello world hello world hello world" // Should be ~9 tokens
inputTokens := uint32(5) // Very low input limit
outputTokens := uint32(10) // Higher output limit
unit := networkingv1alpha1.Second

// Test input rate limit error
rl.AddOrUpdateLimiter(model+"-input", &networkingv1alpha1.RateLimit{
inputLimiterKey := testLimiterKey("default", "test-route-input")
rl.AddOrUpdateLimiter(inputLimiterKey, &networkingv1alpha1.RateLimit{
InputTokensPerUnit: &inputTokens,
Unit: unit,
})

err := rl.RateLimit(model+"-input", longPrompt)
err := rl.RateLimit(inputLimiterKey, longPrompt)
if err == nil {
t.Fatalf("expected input rate limit error")
}
Expand All @@ -250,29 +251,34 @@ func TestTokenRateLimiter_InputAndOutputErrors(t *testing.T) {
}

// Test output rate limit error
rl.AddOrUpdateLimiter(model+"-output", &networkingv1alpha1.RateLimit{
outputLimiterKey := testLimiterKey("default", "test-route-output")
rl.AddOrUpdateLimiter(outputLimiterKey, &networkingv1alpha1.RateLimit{
OutputTokensPerUnit: &outputTokens,
Unit: unit,
})

// First make a successful request to establish the limiter
err = rl.RateLimit(model+"-output", "short")
err = rl.RateLimit(outputLimiterKey, "short")
if err != nil {
t.Fatalf("first request should succeed: %v", err)
}

// Consume all available output tokens
rl.RecordOutputTokens(model+"-output", 10) // Consume all 10 tokens
rl.RecordOutputTokens(outputLimiterKey, 10) // Consume all 10 tokens

// Wait a bit for the tokens to be recorded
time.Sleep(10 * time.Millisecond)

// Next request should be blocked due to insufficient output tokens (< 1 token available)
err = rl.RateLimit(model+"-output", "short") // Short prompt to avoid input limit
err = rl.RateLimit(outputLimiterKey, "short") // Short prompt to avoid input limit
if err == nil {
t.Fatalf("expected output rate limit error")
}
if _, ok := err.(*OutputRateLimitExceededError); !ok {
t.Fatalf("expected OutputRateLimitExceededError, got %T: %v", err, err)
}
}

func testLimiterKey(namespace, modelRouteName string) string {
return fmt.Sprintf("%s/%s", namespace, modelRouteName)
}
Loading
Loading