diff --git a/pkg/kthena-router/filters/ratelimit/ratelimit.go b/pkg/kthena-router/filters/ratelimit/ratelimit.go index 1d40d1e7a..155a68f51 100644 --- a/pkg/kthena-router/filters/ratelimit/ratelimit.go +++ b/pkg/kthena-router/filters/ratelimit/ratelimit.go @@ -66,7 +66,8 @@ type TokenRateLimiter struct { outputLimiter map[string]Limiter // Redis client for global rate limiting - redisClient *redis.Client + redisClient *redis.Client + appliedLimits map[string]*networkingv1alpha1.RateLimit tokenizer tokenizer.Tokenizer } @@ -94,6 +95,7 @@ func NewTokenRateLimiter() *TokenRateLimiter { inputLimiter: make(map[string]Limiter), outputLimiter: make(map[string]Limiter), tokenizer: tokenizer.NewSimpleEstimateTokenizer(), + appliedLimits: make(map[string]*networkingv1alpha1.RateLimit), } } @@ -150,13 +152,15 @@ func (r *TokenRateLimiter) AddOrUpdateLimiter(model string, ratelimit *networkin r.redisClient = redis.NewClient(&redis.Options{ Addr: ratelimit.Global.Redis.Address, }) + } - // Test connection - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if err := r.redisClient.Ping(ctx).Err(); err != nil { - return fmt.Errorf("failed to connect to redis: %w", err) - } + // Test connection + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := r.redisClient.Ping(ctx).Err(); err != nil { + _ = r.redisClient.Close() + r.redisClient = nil + return fmt.Errorf("failed to connect to redis: %w", err) } // Create global rate limiters @@ -200,6 +204,7 @@ func (r *TokenRateLimiter) AddOrUpdateLimiter(model string, ratelimit *networkin } } + r.appliedLimits[model] = ratelimit return nil } @@ -210,6 +215,19 @@ func (r *TokenRateLimiter) DeleteLimiter(model string) { delete(r.inputLimiter, model) delete(r.outputLimiter, model) + if r.appliedLimits != nil { + delete(r.appliedLimits, model) + } +} + +// GetAppliedLimiter returns the applied rate limit configuration for a model +func (r *TokenRateLimiter) GetAppliedLimiter(model string) *networkingv1alpha1.RateLimit { + r.mutex.RLock() + defer r.mutex.RUnlock() + if r.appliedLimits == nil { + return nil + } + return r.appliedLimits[model] } func getTimeUnitDuration(unit networkingv1alpha1.RateLimitUnit) time.Duration { diff --git a/pkg/kthena-router/router/router.go b/pkg/kthena-router/router/router.go index 14604d750..71561c893 100644 --- a/pkg/kthena-router/router/router.go +++ b/pkg/kthena-router/router/router.go @@ -28,6 +28,7 @@ import ( "net" "net/http" "os" + "reflect" "strconv" "strings" "time" @@ -143,8 +144,15 @@ func NewRouter(store datastore.Store, routerConfigPath string) *Router { switch data.EventType { case datastore.EventAdd, datastore.EventUpdate: if data.ModelRoute == nil || data.ModelRoute.Spec.RateLimit == nil { + loadRateLimiter.DeleteLimiter(data.ModelName) return } + + appliedLimit := loadRateLimiter.GetAppliedLimiter(data.ModelName) + if reflect.DeepEqual(appliedLimit, data.ModelRoute.Spec.RateLimit) { + return // RateLimit already successfully applied, nothing to do + } + klog.Infof("add or update rate limit for model %s", data.ModelName) // Configure the unified rate limiter for this model diff --git a/pkg/kthena-router/router/router_test.go b/pkg/kthena-router/router/router_test.go index 9d71343b5..4c9b97466 100644 --- a/pkg/kthena-router/router/router_test.go +++ b/pkg/kthena-router/router/router_test.go @@ -2696,3 +2696,90 @@ type failingEnqueueStore struct { func (s *failingEnqueueStore) Enqueue(req *datastore.Request) error { return fmt.Errorf("injected enqueue failure") } + +func TestRouter_RateLimit_ReconciliationIdempotency(t *testing.T) { + backendHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, `{"id":"resp"}`) + }) + router, store, backend := setupTestRouter(t, backendHandler) + defer backend.Close() + + modelName := "rate-limited-model" + tokens := uint32(10) + unit := aiv1alpha1.Second + + modelRoute := &aiv1alpha1.ModelRoute{ + ObjectMeta: v1.ObjectMeta{Name: "mr-ratelimit", Namespace: "default"}, + Spec: aiv1alpha1.ModelRouteSpec{ + ModelName: modelName, + RateLimit: &aiv1alpha1.RateLimit{ + InputTokensPerUnit: &tokens, + Unit: unit, + }, + Rules: []*aiv1alpha1.Rule{ + { + TargetModels: []*aiv1alpha1.TargetModel{ + {ModelServerName: "ms-1"}, + }, + }, + }, + }, + } + + err := store.AddOrUpdateModelRoute(modelRoute) + assert.NoError(t, err) + + assert.Eventually(t, func() bool { + applied := router.loadRateLimiter.GetAppliedLimiter(modelName) + return applied != nil && *applied.InputTokensPerUnit == 10 + }, 1*time.Second, 10*time.Millisecond) + + prompt := "hello world" + for i := 0; i < 3; i++ { + err := router.loadRateLimiter.RateLimit(modelName, prompt) + assert.NoError(t, err) + } + + err = router.loadRateLimiter.RateLimit(modelName, prompt) + assert.Error(t, err) + + sameRoute := modelRoute.DeepCopy() + sameRoute.Labels = map[string]string{"reconciled": "true"} + + err = store.AddOrUpdateModelRoute(sameRoute) + assert.NoError(t, err) + + time.Sleep(50 * time.Millisecond) + + err = router.loadRateLimiter.RateLimit(modelName, prompt) + assert.Error(t, err, "rate limit state should be preserved across non-config updates") + + newTokens := uint32(50) + updatedRoute := modelRoute.DeepCopy() + updatedRoute.Spec.RateLimit = &aiv1alpha1.RateLimit{ + InputTokensPerUnit: &newTokens, + Unit: unit, + } + + err = store.AddOrUpdateModelRoute(updatedRoute) + assert.NoError(t, err) + + assert.Eventually(t, func() bool { + applied := router.loadRateLimiter.GetAppliedLimiter(modelName) + return applied != nil && *applied.InputTokensPerUnit == 50 + }, 1*time.Second, 10*time.Millisecond) + + err = router.loadRateLimiter.RateLimit(modelName, prompt) + assert.NoError(t, err, "limiter should be updated when rate limit config changes") + + noLimitRoute := modelRoute.DeepCopy() + noLimitRoute.Spec.RateLimit = nil + + err = store.AddOrUpdateModelRoute(noLimitRoute) + assert.NoError(t, err) + + assert.Eventually(t, func() bool { + return router.loadRateLimiter.GetAppliedLimiter(modelName) == nil + }, 1*time.Second, 10*time.Millisecond) +}