diff --git a/pkg/kthena-router/datastore/fairness_queue.go b/pkg/kthena-router/datastore/fairness_queue.go index 8719118d10..4d2d638b9b 100644 --- a/pkg/kthena-router/datastore/fairness_queue.go +++ b/pkg/kthena-router/datastore/fairness_queue.go @@ -234,6 +234,7 @@ type RequestPriorityQueue struct { podCounter PodCounter // Optional; counts backend pods for inflight scaling inflightCount atomic.Int64 // In-flight requests in session-boost mode releaseCh chan struct{} // Signals a permit release in session-boost mode + releaseMu sync.Mutex // Makes a release's freed slot and its signal visible together recheckCh chan struct{} // Signals refreshed backend state in session-boost mode } diff --git a/pkg/kthena-router/datastore/session_boost_queue.go b/pkg/kthena-router/datastore/session_boost_queue.go index dc8cc17acc..38c7e19ab9 100644 --- a/pkg/kthena-router/datastore/session_boost_queue.go +++ b/pkg/kthena-router/datastore/session_boost_queue.go @@ -161,11 +161,15 @@ func (pq *RequestPriorityQueue) admitSessionBoost(req *Request) bool { releaseOnce := sync.Once{} req.Release = func() { releaseOnce.Do(func() { + // Free the slot and signal under releaseMu, so + // dequeueOrHoldPendingRelease sees both or neither. + pq.releaseMu.Lock() pq.inflightCount.Add(-1) select { case pq.releaseCh <- struct{}{}: default: } + pq.releaseMu.Unlock() pq.metricDecInflight(req.ModelName) }) } @@ -244,12 +248,24 @@ func (pq *RequestPriorityQueue) runSessionBoostMode(ctx context.Context) { // an arrival or a backend metrics refresh. With grace enabled, if a release is // also pending, it routes through the grace path so the freed slot is held for a // same-session follow-up; otherwise it re-checks admission directly. +// +// The release check and the inflight count are read together under releaseMu. +// A slot freed after that read is left out of this pass, so it is still held for +// a follow-up through the grace path. func (pq *RequestPriorityQueue) dequeueOrHoldPendingRelease(ctx context.Context, grace bool) { - if grace && pq.drainPendingRelease() { + if !grace { + pq.tryBackpressureDequeue(ctx) + return + } + pq.releaseMu.Lock() + released := pq.drainPendingRelease() + inflight := pq.inflightCount.Load() + pq.releaseMu.Unlock() + if released { pq.waitGraceAndDequeue(ctx) return } - pq.tryBackpressureDequeue(ctx) + pq.tryBackpressureDequeueAbove(ctx, inflight) } // RecheckBackpressure asks the session-boost dequeue loop to re-evaluate held @@ -350,6 +366,13 @@ func (pq *RequestPriorityQueue) drainCancelledLocked() int { // the queue is empty. This avoids the one-request-per-tick bottleneck during // initial ramp-up and whenever spare capacity exists. func (pq *RequestPriorityQueue) tryBackpressureDequeue(ctx context.Context) { + pq.tryBackpressureDequeueAbove(ctx, 0) +} + +// tryBackpressureDequeueAbove is tryBackpressureDequeue with a lower bound for the +// inflight count: the count the caller read, plus the requests admitted in this +// pass. Slots freed after the caller's read are not used here. +func (pq *RequestPriorityQueue) tryBackpressureDequeueAbove(ctx context.Context, minInflight int64) { // In session-boost mode, the total inflight limit is InflightPerPod scaled by // the number of backend pods serving the model. perPod := pq.config.InflightPerPod @@ -366,7 +389,7 @@ func (pq *RequestPriorityQueue) tryBackpressureDequeue(ctx context.Context) { } for { - currentInflight := pq.inflightCount.Load() + currentInflight := max(pq.inflightCount.Load(), minInflight) if currentInflight >= maxInflight { pq.mu.Lock() @@ -405,6 +428,7 @@ func (pq *RequestPriorityQueue) tryBackpressureDequeue(ctx context.Context) { // continue admitting the next queued request. continue } + minInflight++ klog.V(4).Infof("[SessionBoost] backpressure dequeue: user=%s model=%s sessionBoost=%v inflight=%d/%d", req.UserID, req.ModelName, req.SessionBoost, pq.inflightCount.Load(), maxInflight) diff --git a/pkg/kthena-router/datastore/session_boost_queue_test.go b/pkg/kthena-router/datastore/session_boost_queue_test.go index a2b159142b..3cbafd571e 100644 --- a/pkg/kthena-router/datastore/session_boost_queue_test.go +++ b/pkg/kthena-router/datastore/session_boost_queue_test.go @@ -574,22 +574,25 @@ func TestSessionBoostQueue_GracePeriod_BoostedArrives(t *testing.T) { } // TestSessionBoostQueue_DequeueOrHoldPendingRelease covers the handling of an -// arrival or a recheck. When a release is pending at the same time, select may -// pick either signal first, so this handling must not hand the freed slot to the -// head of the queue without the grace wait. The context is cancelled up front: -// a grace wait then returns without dispatching, while a direct dequeue still -// admits, so whether the held request is admitted shows which path was taken. +// arrival or a recheck. When a release is pending at the same time (select may +// pick either signal first), or lands while the handling runs, this handling must +// not hand the freed slot to the head of the queue without the grace wait. The +// context is cancelled up front: a grace wait then returns without dispatching, +// while a direct dequeue still admits, so whether the held request is admitted +// shows which path was taken. func TestSessionBoostQueue_DequeueOrHoldPendingRelease(t *testing.T) { tests := []struct { - name string - gracePeriod time.Duration - releasePending bool - wantAdmitted bool - wantReleasePending bool + name string + gracePeriod time.Duration + releasePending bool + releaseDuringHandling bool + wantAdmitted bool + wantReleasePending bool }{ {name: "pending release holds the freed slot for grace", gracePeriod: time.Hour, releasePending: true, wantAdmitted: false, wantReleasePending: false}, {name: "no pending release admits directly", gracePeriod: time.Hour, releasePending: false, wantAdmitted: true, wantReleasePending: false}, {name: "grace disabled admits directly", gracePeriod: 0, releasePending: true, wantAdmitted: true, wantReleasePending: true}, + {name: "release during handling holds the freed slot for grace", gracePeriod: time.Hour, releaseDuringHandling: true, wantAdmitted: false, wantReleasePending: true}, } for _, tt := range tests { @@ -619,11 +622,20 @@ func TestSessionBoostQueue_DequeueOrHoldPendingRelease(t *testing.T) { } <-q.notifyCh - // Free the slot. Without a pending release, its signal has already - // been handled. - req1.Release() - if !tt.releasePending { - <-q.releaseCh + if tt.releaseDuringHandling { + // podCounter runs inside the handling, after the release check, + // so freeing the slot from it lands the release in between. + q.podCounter = func() int { + req1.Release() + return 1 + } + } else { + // Free the slot. Without a pending release, its signal has + // already been handled. + req1.Release() + if !tt.releasePending { + <-q.releaseCh + } } ctx, cancel := context.WithCancel(context.Background())