diff --git a/docs/kthena/docs/user-guide/session-boost.md b/docs/kthena/docs/user-guide/session-boost.md index d12b2dc44..00d170992 100644 --- a/docs/kthena/docs/user-guide/session-boost.md +++ b/docs/kthena/docs/user-guide/session-boost.md @@ -112,6 +112,7 @@ env: | `ENABLE_SESSION_BOOST` | Enable the session-boost scheduling strategy | `false` | Mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING` | | `SESSION_BOOST_HEADER` | HTTP header used to identify conversation sessions | `X-Session-ID` | Must match what your clients send | | `SESSION_BOOST_MAX_SESSIONS` | Max number of recently-completed sessions kept "warm" for boosting (LRU-bounded) | `4096` | When the cache is full, the least-recently-used session is evicted automatically. Size it by the number of concurrent conversations you want to keep boosted—no time-based tuning required | +| `SESSION_BOOST_MAX_OVERTAKES` | Max times a queued boosted request may be overtaken by newer completions | `16` | Bounds starvation: once a boosted request has been overtaken this many times it is "forced" and ranks ahead of newer completions, guaranteeing it runs before hitting `SESSION_BOOST_TIMEOUT`. Lower it to bound tail latency; raise it to favor prefix-cache reuse | | `SESSION_BOOST_INFLIGHT_PER_POD` | Inflight requests admitted per backend pod | `16` | The total inflight limit is this value multiplied by the number of backend pods. Size it from the per-pod concurrency (e.g., vLLM's `--max-num-seqs`) | | `SESSION_BOOST_TIMEOUT` | Maximum time a request may wait in the queue before it is rejected with 504 | `30s` | Enabled by default. A request queued longer than this is rejected with `504 Gateway Timeout`. Set a non-positive duration (e.g. `0s`) to disable it. This is the only server-side queue-wait bound in session-boost mode; `FAIRNESS_QUEUE_TIMEOUT` does not apply | @@ -169,7 +170,10 @@ This design is intentional: inference engines such as vLLM evict their prefix (K The session boost queue uses a simple two-level priority: 1. **Boosted requests** (session tracked in the LRU cache) are always dequeued before non-boosted requests. -2. **Within the same boost level**, requests are served in FIFO order (earliest arrival first). +2. **Within the boosted group**, requests are ordered by their session's **completion time**: the session whose previous turn completed most recently is served first. Its prefix (KV) cache is the most likely to still be warm on the backend, so serving it first maximizes cache reuse. This ordering does not depend on when the follow-up request arrived, only on how recently the session's last turn finished; ties are broken FIFO by arrival time. +3. **Among non-boosted requests**, ordering is FIFO (earliest arrival first). + +Ordering by completion time (rather than a plain LIFO on arrival) keeps the queue fair: a boosted request cannot be starved by newer arrivals, because a session's completion time is fixed once its turn finishes and does not keep getting pushed back. Every boosted request also still leapfrogs the entire non-boosted tier, so a follow-up turn is always prioritized over unrelated traffic. ### Backpressure Control @@ -197,13 +201,13 @@ If a request waits in the queue longer than `SESSION_BOOST_TIMEOUT`, the router Session boost and user fairness are two mutually exclusive scheduling strategies for the per-model request queue; they configure that queue for different goals: -| Aspect | Session Boost | User Fairness | -| ---------------- | ---------------------------------- | --------------------------------- | -| Goal | Maximize prefix cache hits | Equitable resource allocation | -| Activation | `ENABLE_SESSION_BOOST=true` | `ENABLE_FAIRNESS_SCHEDULING=true` | -| Requires user ID | No | Yes | -| Priority logic | Boosted > non-boosted, FIFO within | Lower recent usage wins | -| Best for | Multi-turn latency optimization | Multi-tenant contention | +| Aspect | Session Boost | User Fairness | +| ---------------- | ------------------------------------------------------------ | --------------------------------- | +| Goal | Maximize prefix cache hits | Equitable resource allocation | +| Activation | `ENABLE_SESSION_BOOST=true` | `ENABLE_FAIRNESS_SCHEDULING=true` | +| Requires user ID | No | Yes | +| Priority logic | Boosted > non-boosted; most-recently-completed session first | Lower recent usage wins | +| Best for | Multi-turn latency optimization | Multi-tenant contention | The two strategies are **mutually exclusive**; enabling both is a configuration error. Enable one or the other based on your primary scheduling concern. @@ -286,9 +290,9 @@ The grace period is an **advanced, scenario-specific** tuning knob. It is **disa By default, when a request completes the queue immediately dequeues the next request. The grace period (`SESSION_BOOST_GRACE_PERIOD`) instead makes the queue **briefly hold the freed dequeue slot** after a completion, waiting for a follow-up from the *same* session to arrive so it can ride the still-warm prefix cache: -- If a boosted (same-session) request arrives during the grace window, it is admitted first when the window ends, because boosted requests outrank others in the queue. +- If a boosted (same-session) request arrives during the grace window, it is admitted first when the window ends, because boosted requests outrank others in the queue—and among boosted requests the one whose session completed most recently is at the head. - If no boosted request arrives before the window expires, the next non-boosted request proceeds as normal. -- If the head of the queue is already a boosted request, the grace period is skipped entirely. +- The queue always waits for the full grace window even when a boosted request is already at the head, because a still-newer same-session follow-up may yet arrive and should be the one to ride the warm prefix cache. ### When it might help diff --git a/docs/proposal/session-boost-strategy.md b/docs/proposal/session-boost-strategy.md index 05b9b12eb..737ac7f46 100644 --- a/docs/proposal/session-boost-strategy.md +++ b/docs/proposal/session-boost-strategy.md @@ -79,11 +79,12 @@ Session boost is a priority strategy of the shared per-model request priority qu │ ┌──────────────────┐ ┌─────────────────────────┐ │ │ │ SessionTracker │<────│ MarkSessionRequest │ │ │ │ (bounded LRU) │ │ Completed() │ │ -│ │ keys: sessionID │ │ (after response sent) │ │ -│ │ cap: 4096 default│ └─────────────────────────┘ │ +│ │ key: sessionID │ │ (after response sent) │ │ +│ │ value: completeAt│ └─────────────────────────┘ │ +│ │ cap: 4096 default│ │ │ └────────┬─────────┘ │ │ │ │ -│ │ HasRecentCompletion(sessionID)? │ +│ │ CompletionTime(sessionID)? │ │ ▼ │ │ ┌──────────────────┐ │ │ │ PushRequest() │ │ @@ -98,9 +99,10 @@ Session boost is a priority strategy of the shared per-model request priority qu │ │ │ │ │ │ Ordering: │ │ │ │ 1. SessionBoost=true > false │ │ -│ │ 2. Within same boost: FIFO │ │ +│ │ 2. Within boosted: most recently │ │ +│ │ completed session first │ │ │ │ │ │ -│ │ [Boosted-1] [Boosted-2] [Normal-1] │ │ +│ │ [Boosted-warm][Boosted-cooler][Norm]│ │ │ └──────────────────────────────────────┘ │ │ │ │ │ ▼ │ @@ -159,14 +161,14 @@ Timeline: ▼ ▼ ▼ ─────┼──┼──────────────────────────────────┼───── │ │ │ - │ │ Case A: Head already boosted → │ - │ │ skip grace, dequeue immediately │ - │ │ │ - │ │ Case B: Otherwise hold the slot │ - │ │ for the full grace period, then │ - │ │ dequeue the head (boosted ranks │ - │ │ first, so a follow-up that │ - │ │ arrived wins automatically) │ + │ │ Hold the slot for the full │ + │ │ grace period, then dequeue the │ + │ │ head. Boosted requests rank │ + │ │ first and, among boosted, the │ + │ │ newest-arrived wins, so a │ + │ │ same-session follow-up that │ + │ │ arrived during the window is │ + │ │ dequeued first automatically. │ │ └────────────────────────────────┐ ``` @@ -178,16 +180,17 @@ It is important that grace is **triggered only by release events**, because a re 1. A request finishes → `Release()` runs → `inflightCount` is decremented → a signal is sent on `releaseCh`. 2. The freed slot would normally be claimed immediately by the current heap head (which may be an unrelated, non-boosted request). The grace wait instead **holds that just-freed slot for up to `SessionBoostGracePeriod`**, betting that a same-session follow-up is about to arrive and reuse the warm prefix cache. -3. When the wait resolves, `tryBackpressureDequeue` runs and re-checks **both** capacity gates before admitting anyone. A boosted follow-up that arrived during grace is admitted first because boosted requests outrank others in the heap, and only when `inflight < MaxConcurrent` **and** at least one backend pod reports capacity. +3. When the wait resolves, `tryBackpressureDequeue` runs and re-checks **both** capacity gates before admitting anyone. A boosted follow-up that arrived during grace is admitted first because boosted requests outrank others in the heap and, among boosted requests, the one whose session completed most recently is at the head — and only when `inflight < MaxConcurrent` **and** at least one backend pod reports capacity. + +The grace wait always runs for the full window (subject to shutdown), and in every case the two capacity gates are the final arbiter: -The grace wait can resolve in three ways, and in every case the two capacity gates are the final arbiter: +| Outcome | Trigger | What happens next | +| ----------------- | ---------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Grace expires** | The timer fires (`timer.C`). | Stop holding the slot and fall through to the capacity gates, which admit the heap head. If a same-session follow-up arrived during the wait it is now boosted and sits at the head (session with the most recent completion first), so it wins automatically (subject to inflight + backend capacity). | -| Outcome | Trigger | What happens next | -| -------------------------- | --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| **Skip grace (fast path)** | The heap head is *already* session-boosted when the release fires (`isHeadSessionBoosted()`). | No wait at all — go straight to the capacity gates and admit if they pass. There is nothing to wait for. | -| **Grace expires** | The timer fires (`timer.C`). | Stop holding the slot and fall through to the capacity gates, which admit the heap head. If a same-session follow-up arrived during the wait it is now boosted and sits at the head, so it wins automatically (subject to inflight + backend capacity). | +> There is no "head already boosted" fast path. The queue holds the freed slot for the full grace window even when a boosted request is already at the head, because a still-newer same-session follow-up may yet arrive during the window and should be the one to ride the warm prefix cache. -This ordering matters because the three checks answer different questions and run in sequence: +This ordering matters because the timing layer and the two gates answer different questions and run in sequence: ``` Release frees a slot @@ -215,7 +218,7 @@ Two interactions are worth calling out explicitly: - **Grace never overrides backpressure.** If the grace window ends but the inflight limit is already reached or every backend pod is busy, the request is **not** admitted — `tryBackpressureDequeue` simply holds (and drains any cancelled requests from the heap) until the next release or new arrival reopens the gates. Grace only chooses *who* tries next; the capacity gates decide *whether* anyone runs. - **Fresh arrivals on an idle queue bypass grace.** Grace is tied to `releaseCh` (a freed slot), not to `notifyCh` (a new arrival). When a request lands on an otherwise idle queue with no pending release, it goes straight to the capacity gates with no grace delay, so enabling grace adds no admission latency to first turns. The only place a new arrival waits is *inside* an already-running grace window, where it is precisely the boosted follow-up the window exists to catch. When both a release and a new arrival are pending at once, the release is preferred so the freed slot is the one held for the grace window. -The net effect is a strict precedence: **grace timing → inflight gate → backend gate**. The grace layer is purely additive and optional (`SessionBoostGracePeriod = 0` removes it entirely, taking the no-grace fast path), and it can only ever *delay* an admission to favor a session-boosted follow-up — it can never admit a request that the inflight or backend gates would otherwise reject. +The net effect is a strict precedence: **grace timing → inflight gate → backend gate**. The grace layer is purely additive and optional (`SessionBoostGracePeriod = 0` removes it entirely, dequeuing immediately on each release), and it can only ever *delay* an admission to favor a session-boosted follow-up — it can never admit a request that the inflight or backend gates would otherwise reject. #### Configuration @@ -224,6 +227,7 @@ The net effect is a strict precedence: **grace timing → inflight gate → back | `ENABLE_SESSION_BOOST` | `false` | Enable the session-boost scheduling strategy (mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING`) | | `SESSION_BOOST_HEADER` | `X-Session-ID` | HTTP header used to identify conversation sessions | | `SESSION_BOOST_MAX_SESSIONS` | `4096` | Maximum number of recently-completed sessions kept warm for boosting. Bounds an LRU cache; the least-recently-used session is evicted when exceeded. Sized by session count, not time | +| `SESSION_BOOST_MAX_OVERTAKES` | `16` | Bounds how many times a queued boosted request may be overtaken by newer completions before it is "forced" and ranks ahead of the completion-time lane. Prevents sustained boosted traffic from starving an older request into hitting `SESSION_BOOST_TIMEOUT` | | `SESSION_BOOST_GRACE_PERIOD` | `0` | Wait time after release for same-session follow-up. Disabled by default; enable only when you understand the latency trade-off | | `SESSION_BOOST_INFLIGHT_PER_POD` | `16` | Inflight requests admitted per backend pod; total inflight = perPod x backend pod count. Size it from the estimated per-pod concurrency (e.g. vLLM's --max-num-seqs) | | `SESSION_BOOST_TIMEOUT` | `30s` | Maximum time a request may wait in the queue before it is rejected with HTTP 504. Enabled by default; set a non-positive duration (e.g. `0s`) to disable it. It is the only server-side queue-wait bound in session-boost mode (`FAIRNESS_QUEUE_TIMEOUT` does not apply) | @@ -266,9 +270,29 @@ type RequestPriorityQueue struct { // Session-boost ordering (RequestPriorityQueue.Less when sessionBoost == true): // 1. SessionBoost=true comes before SessionBoost=false -// 2. Within same boost status: earlier RequestTime comes first (FIFO) +// 2. Within boosted requests: a "forced" request (overtaken past the bound) comes +// before non-forced ones; among forced, earlier RequestTime first (FIFO) +// 3. Within non-forced boosted requests: the session whose previous turn completed +// most recently comes first (warmest prefix cache); ties broken FIFO by +// RequestTime +// 4. Within non-boosted requests: earlier RequestTime comes first (FIFO) ``` +> **Bounded overtaking (anti-starvation).** Ordering boosted requests by completion +> time means an older queued request keeps the timestamp captured when it enqueued, +> while every later-completing session gets a newer one and jumps ahead. Under +> sustained boosted traffic that older request could be overtaken indefinitely and +> hit `SESSION_BOOST_TIMEOUT` instead of ever running. To bound this without slowing +> enqueue, the queue keeps a single monotonic counter of boosted dequeues and each +> request snapshots it on arrival; the difference at any later moment is exactly how +> many boosted requests have been served ahead of it. At dequeue the loop checks +> only the longest-waiting boosted request (tracked in arrival order), and once its +> overtake count reaches `SESSION_BOOST_MAX_OVERTAKES` it is marked *forced* and +> moved into place with `heap.Fix`, so it ranks ahead of the completion-time lane and +> can no longer be pushed back. Enqueue stays O(log n) (a counter snapshot plus a +> heap push) and the check adds O(log n) only when a request actually needs +> promoting, while still preferring cache-warm sessions in the common case. + #### Session Identification Sessions are identified by a configurable HTTP header (default: `X-Session-ID`, controlled by `SESSION_BOOST_HEADER` environment variable). This is a client-provided identifier that groups related requests in a multi-turn conversation: diff --git a/pkg/kthena-router/datastore/fairness_queue.go b/pkg/kthena-router/datastore/fairness_queue.go index 099887fc2..dfd526c25 100644 --- a/pkg/kthena-router/datastore/fairness_queue.go +++ b/pkg/kthena-router/datastore/fairness_queue.go @@ -73,6 +73,15 @@ type FairnessQueueConfig struct { // default of defaultSessionBoostMaxSessions is used. SessionBoostMaxSessions int + // SessionBoostMaxOvertakes bounds how many times a queued boosted request may be + // overtaken by newer same-priority boosted requests (sessions that completed + // more recently) before it is "aged": once a boosted request has been overtaken + // this many times it ranks ahead of the completion-time lane and can no longer + // be pushed back, so sustained boosted traffic cannot starve it into hitting + // SESSION_BOOST_TIMEOUT. When <= 0, a default of defaultSessionBoostMaxOvertakes + // is used. Only meaningful when SessionBoostEnabled is true. + SessionBoostMaxOvertakes int + // SessionBoostGracePeriod is the duration to wait after a release before // dequeuing the next request in backpressure mode. This gives the same session // time to submit a follow-up request that benefits from prefix cache. @@ -96,6 +105,11 @@ const defaultSessionBoostInflightPerPod = 16 // (<= 0). Each entry is tiny (a session ID), so the default is generous. const defaultSessionBoostMaxSessions = 4096 +// defaultSessionBoostMaxOvertakes is the bound on how many times a queued boosted +// request may be overtaken by newer completions before it is aged and forced ahead +// of the completion-time lane. Used when SessionBoostMaxOvertakes is not set (<= 0). +const defaultSessionBoostMaxOvertakes = 16 + // defaultSessionBoostHeader is the HTTP header used to identify conversation // sessions when SESSION_BOOST_HEADER is not set. const defaultSessionBoostHeader = "X-Session-ID" @@ -112,6 +126,7 @@ func DefaultFairnessQueueConfig() FairnessQueueConfig { SessionBoostEnabled: false, SessionIDHeader: defaultSessionBoostHeader, SessionBoostMaxSessions: defaultSessionBoostMaxSessions, + SessionBoostMaxOvertakes: defaultSessionBoostMaxOvertakes, SessionBoostGracePeriod: 0, InflightPerPod: defaultSessionBoostInflightPerPod, } @@ -148,10 +163,14 @@ type Request struct { SessionID string // Session identifier for multi-turn conversations Priority float64 // Priority (lower value means higher priority) SessionBoost bool // Whether this request has session priority boost (recently completed session) - RequestTime time.Time - NotifyChan chan struct{} - CancelCh <-chan struct{} // Request-scoped cancellation signal - Release func() // Set by the queue when a permit is acquired + // LastTurnCompletedAt is the time the session's previous turn completed, captured + // when the request is boosted. Boosted requests are ordered by this timestamp + // (most recent first) so the session with the warmest prefix cache runs first. + LastTurnCompletedAt time.Time + RequestTime time.Time + NotifyChan chan struct{} + CancelCh <-chan struct{} // Request-scoped cancellation signal + Release func() // Set by the queue when a permit is acquired // admitMu serializes admission (by the dequeue loop) against abandonment (by // the waiting caller on timeout/cancel), closing the race where the loop has @@ -160,6 +179,19 @@ type Request struct { admitMu sync.Mutex admitted bool // admission committed: Release is set and about to be signalled abandoned bool // caller gave up before admission; the loop must not admit + + // Bounded-overtaking bookkeeping for session-boost ordering. All three fields + // are guarded by the queue lock (pq.mu). heapIndex is the request's current + // position in the heap slice (maintained by Swap/Push/Pop) so it can be + // promoted in place with heap.Fix; it is -1 when the request is not in the + // heap. boostEnqueueAdmits snapshots pq.boostAdmitCount at enqueue, so + // pq.boostAdmitCount-boostEnqueueAdmits is how many boosted requests have been + // served ahead of this one (i.e. how many times it has been overtaken). + // boostForced is set once that count reaches the bound; it makes the request + // rank ahead of the completion-time lane and is never cleared. + heapIndex int + boostEnqueueAdmits int64 + boostForced bool } // commitAdmission runs fn under the request lock, but only if the caller has not @@ -227,6 +259,15 @@ 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 + + // Bounded-overtaking state (session-boost mode), guarded by mu. boostAdmitCount + // is a monotonic count of boosted requests dequeued so far; each request + // snapshots it at enqueue so the difference gives how many times it has been + // overtaken. boostArrivals holds queued boosted requests in arrival order so the + // oldest still-waiting one can be found in amortized O(1) (stale entries are + // skipped lazily), avoiding any per-enqueue scan of the heap. + boostAdmitCount int64 + boostArrivals []*Request } var _ heap.Interface = &RequestPriorityQueue{} @@ -273,11 +314,32 @@ func NewRequestPriorityQueueWithConfig(metricsInstance *metrics.Metrics, cfg Fai func (pq *RequestPriorityQueue) Len() int { return len(pq.heap) } func (pq *RequestPriorityQueue) Less(i, j int) bool { - // Session-boost mode: boosted requests outrank others; ties broken FIFO. + // Session-boost mode: boosted requests outrank others. Among boosted requests, + // the session whose previous turn completed most recently wins, because its + // prefix cache is the most likely to still be warm on the backend; ties are + // broken FIFO by arrival time. To keep this ordering from starving an older + // boosted request under a steady stream of newer completions, a request that + // has been overtaken SessionBoostMaxOvertakes times is marked "forced" and ranks + // ahead of the completion-time lane. Non-boosted requests keep FIFO ordering. if pq.sessionBoost { if pq.heap[i].SessionBoost != pq.heap[j].SessionBoost { return pq.heap[i].SessionBoost } + if pq.heap[i].SessionBoost { + // Forced requests (overtaken past the bound) rank ahead of the + // completion-time lane so newer traffic can no longer push them back. + if pq.heap[i].boostForced != pq.heap[j].boostForced { + return pq.heap[i].boostForced + } + if pq.heap[i].boostForced { + // Both forced: oldest arrival first so the longest waiter runs next. + return pq.heap[i].RequestTime.Before(pq.heap[j].RequestTime) + } + if !pq.heap[i].LastTurnCompletedAt.Equal(pq.heap[j].LastTurnCompletedAt) { + return pq.heap[i].LastTurnCompletedAt.After(pq.heap[j].LastTurnCompletedAt) + } + return pq.heap[i].RequestTime.Before(pq.heap[j].RequestTime) + } return pq.heap[i].RequestTime.Before(pq.heap[j].RequestTime) } // same user, FIFO @@ -294,10 +356,13 @@ func (pq *RequestPriorityQueue) Less(i, j int) bool { func (pq *RequestPriorityQueue) Swap(i, j int) { pq.heap[i], pq.heap[j] = pq.heap[j], pq.heap[i] + pq.heap[i].heapIndex = i + pq.heap[j].heapIndex = j } func (pq *RequestPriorityQueue) Push(x interface{}) { item := x.(*Request) + item.heapIndex = len(pq.heap) pq.heap = append(pq.heap, item) } @@ -308,6 +373,7 @@ func (pq *RequestPriorityQueue) Pop() interface{} { } item := pq.heap[n-1] pq.heap[n-1] = nil + item.heapIndex = -1 pq.heap = pq.heap[0 : n-1] return item } @@ -317,11 +383,26 @@ func (pq *RequestPriorityQueue) PushRequest(r *Request) error { defer pq.mu.Unlock() // In session-boost mode, promote requests whose session recently completed. - if pq.sessionBoost && pq.sessionTracker != nil && r.SessionID != "" && - pq.sessionTracker.HasRecentCompletion(r.SessionID) { - r.SessionBoost = true + // Capture the session's completion time so boosted requests can be ordered by + // prefix-cache warmth (most recently completed first). + if pq.sessionBoost && pq.sessionTracker != nil && r.SessionID != "" { + if completedAt, ok := pq.sessionTracker.CompletionTime(r.SessionID); ok { + r.SessionBoost = true + r.LastTurnCompletedAt = completedAt + } } + if r.SessionBoost { + // Bounded overtaking: snapshot how many boosted requests have already been + // admitted. Later, pq.boostAdmitCount-boostEnqueueAdmits is exactly how many + // boosted requests were served ahead of this one; when that reaches the + // bound the dequeue loop forces this request (see promoteStarvedBoostLocked). + // This is O(1) here, so enqueue stays O(log n) via heap.Push instead of the + // previous per-enqueue heap scan. boostArrivals keeps arrival order so the + // oldest waiter can be found cheaply. + r.boostEnqueueAdmits = pq.boostAdmitCount + pq.boostArrivals = append(pq.boostArrivals, r) + } heap.Push(pq, r) // Update queue size metrics @@ -340,6 +421,44 @@ func (pq *RequestPriorityQueue) PushRequest(r *Request) error { return nil } +// oldestWaitingBoostLocked returns the boosted request that has been waiting in the +// queue the longest, or nil if none. It lazily discards entries from the front of +// boostArrivals that have already left the heap (admitted, forced out, or drained +// as cancelled), so amortized cost is O(1). Caller must hold pq.mu. +func (pq *RequestPriorityQueue) oldestWaitingBoostLocked() *Request { + for len(pq.boostArrivals) > 0 { + r := pq.boostArrivals[0] + if r.heapIndex >= 0 && !r.isCancelled() { + return r + } + pq.boostArrivals[0] = nil + pq.boostArrivals = pq.boostArrivals[1:] + } + return nil +} + +// promoteStarvedBoostLocked enforces the bounded-overtaking guarantee. If the +// longest-waiting boosted request has been overtaken by at least +// SessionBoostMaxOvertakes newer boosted admissions, it is marked forced and moved +// into place with heap.Fix (O(log n)) so the next pop returns it. Only the oldest +// waiter needs checking: boostEnqueueAdmits is monotonic with arrival order, so the +// oldest has been overtaken the most and crosses the bound first. Caller must hold +// pq.mu. +func (pq *RequestPriorityQueue) promoteStarvedBoostLocked() { + oldest := pq.oldestWaitingBoostLocked() + if oldest == nil || oldest.boostForced { + return + } + maxOvertakes := pq.config.SessionBoostMaxOvertakes + if maxOvertakes <= 0 { + maxOvertakes = defaultSessionBoostMaxOvertakes + } + if pq.boostAdmitCount-oldest.boostEnqueueAdmits >= int64(maxOvertakes) { + oldest.boostForced = true + heap.Fix(pq, oldest.heapIndex) + } +} + // popWhenAvailable blocks until an item is available or the context is done, then pops one item. // Cancelled/timed-out requests are skipped automatically. func (pq *RequestPriorityQueue) popWhenAvailable(ctx context.Context) (*Request, error) { @@ -347,6 +466,12 @@ func (pq *RequestPriorityQueue) popWhenAvailable(ctx context.Context) (*Request, for { pq.mu.Lock() if len(pq.heap) > 0 { + // Anti-starvation: before picking the heap head, promote the + // longest-waiting boosted request if it has been overtaken too many + // times, so completion-time ordering cannot starve it. + if pq.sessionBoost { + pq.promoteStarvedBoostLocked() + } req := heap.Pop(pq).(*Request) // Skip cancelled/timed-out requests @@ -400,6 +525,12 @@ func (pq *RequestPriorityQueue) popWhenAvailable(ctx context.Context) (*Request, pq.metricRecordDuration(req.ModelName, req.UserID, time.Since(req.RequestTime)) pq.metricIncDequeue(req.ModelName, req.UserID) + // Count this boosted dequeue so the overtake bound advances for every + // still-waiting boosted request (see promoteStarvedBoostLocked). + if pq.sessionBoost && req.SessionBoost { + pq.boostAdmitCount++ + } + pq.mu.Unlock() return req, nil } diff --git a/pkg/kthena-router/datastore/session_boost_queue.go b/pkg/kthena-router/datastore/session_boost_queue.go index 95ee4d5dd..eea0c0dc3 100644 --- a/pkg/kthena-router/datastore/session_boost_queue.go +++ b/pkg/kthena-router/datastore/session_boost_queue.go @@ -56,10 +56,18 @@ type PodCounter func() int // keeping a few extra entries is harmless because boosting only matters when the // queue is contended. type SessionTracker struct { - // cache maps a session ID to a presence marker, ordered by recency. The - // underlying hashicorp/golang-lru cache is safe for concurrent use and evicts - // the least-recently-used session once capacity is exceeded. - cache *lru.Cache[string, struct{}] + // cache maps a session ID to the time its most recent request completed, + // ordered by recency. The underlying hashicorp/golang-lru cache is safe for + // concurrent use and evicts the least-recently-used session once capacity is + // exceeded. The stored timestamp is used to order boosted requests so that the + // session whose previous turn completed most recently—and is therefore most + // likely to still have a warm prefix cache—is served first. + cache *lru.Cache[string, time.Time] + + // now returns the current time when a completion is recorded. It defaults to + // time.Now and is only overridden in tests to make completion timestamps + // deterministic. + now func() time.Time } // NewSessionTracker creates a new session tracker that remembers up to capacity @@ -69,24 +77,24 @@ func NewSessionTracker(capacity int) *SessionTracker { if capacity <= 0 { capacity = defaultSessionBoostMaxSessions } - cache, err := lru.NewWithEvict(capacity, func(sessionID string, _ struct{}) { + cache, err := lru.NewWithEvict(capacity, func(sessionID string, _ time.Time) { klog.V(4).Infof("[SessionTracker] evicted LRU session %q", sessionID) }) if err != nil { // capacity is guaranteed positive above, so this is unreachable in practice. klog.Errorf("[SessionTracker] failed to create LRU cache (capacity=%d): %v", capacity, err) } - return &SessionTracker{cache: cache} + return &SessionTracker{cache: cache, now: time.Now} } // MarkRequestCompleted records that a request from the given session has completed, -// promoting it to the most-recently-used position. When the cache exceeds its -// capacity, the least-recently-used session is evicted. +// storing the completion time and promoting it to the most-recently-used position. +// When the cache exceeds its capacity, the least-recently-used session is evicted. func (st *SessionTracker) MarkRequestCompleted(sessionID string) { if sessionID == "" { return } - st.cache.Add(sessionID, struct{}{}) + st.cache.Add(sessionID, st.now()) } // HasRecentCompletion reports whether the given session ID is currently tracked @@ -99,6 +107,16 @@ func (st *SessionTracker) HasRecentCompletion(sessionID string) bool { return st.cache.Contains(sessionID) } +// CompletionTime returns the time the given session's most recent request completed +// and whether the session is currently tracked. It is a pure read (via Peek) and +// does not change recency ordering. +func (st *SessionTracker) CompletionTime(sessionID string) (time.Time, bool) { + if sessionID == "" { + return time.Time{}, false + } + return st.cache.Peek(sessionID) +} + // ActiveSessions returns the number of sessions currently tracked. func (st *SessionTracker) ActiveSessions() int { return st.cache.Len() @@ -233,30 +251,16 @@ func (pq *RequestPriorityQueue) drainPendingRelease() bool { } } -// isHeadSessionBoosted checks if the highest-priority request in the queue has a session boost. -func (pq *RequestPriorityQueue) isHeadSessionBoosted() bool { - pq.mu.RLock() - defer pq.mu.RUnlock() - if len(pq.heap) == 0 { - return false - } - return pq.heap[0].SessionBoost -} - // waitGraceAndDequeue holds a just-freed slot for up to SessionBoostGracePeriod // before dispatching, giving a same-session follow-up time to arrive. It does not // need to watch for the follow-up itself: boosted requests outrank others in the -// heap, so once the timer fires tryBackpressureDequeue admits the boosted request -// first if one showed up. The wait stays responsive to shutdown via ctx/stopCh. +// heap and, among boosted requests, the one whose session completed most recently +// wins, so once the timer fires tryBackpressureDequeue admits the warmest boosted +// follow-up first if one showed up. The wait always runs for the full window even +// when the head is already boosted, because a follow-up from an even-more-recently +// completed session may yet arrive and should be the one to ride the warm prefix +// cache. The wait stays responsive to shutdown via ctx/stopCh. func (pq *RequestPriorityQueue) waitGraceAndDequeue(ctx context.Context) { - // Fast path: a boosted follow-up is already at the head, so there is nothing - // to wait for. - if pq.isHeadSessionBoosted() { - klog.V(4).Info("[SessionBoost] grace: head already boosted, skipping wait") - pq.tryBackpressureDequeue(ctx) - return - } - klog.V(4).Infof("[SessionBoost] grace: holding freed slot for %v", pq.config.SessionBoostGracePeriod) timer := time.NewTimer(pq.config.SessionBoostGracePeriod) defer timer.Stop() @@ -290,6 +294,7 @@ func (pq *RequestPriorityQueue) drainCancelledLocked() int { pq.metricDecSize(req.ModelName, req.UserID) pq.metricRecordDuration(req.ModelName, req.UserID, time.Since(req.RequestTime)) pq.metricIncCancelled(req.ModelName, req.UserID) + req.heapIndex = -1 // no longer in the heap; lets boostArrivals skip it continue } kept = append(kept, req) @@ -303,6 +308,12 @@ func (pq *RequestPriorityQueue) drainCancelledLocked() int { pq.heap[i] = nil } pq.heap = kept + // Reset positions before heap.Init: compaction shifted elements, and heap.Init + // only updates indices for elements it moves, so stale heapIndex values must be + // corrected here to keep heap.Fix targeting the right slots. + for idx, req := range pq.heap { + req.heapIndex = idx + } heap.Init(pq) return drained } diff --git a/pkg/kthena-router/datastore/session_boost_queue_test.go b/pkg/kthena-router/datastore/session_boost_queue_test.go index 30859b3e2..d63ff391c 100644 --- a/pkg/kthena-router/datastore/session_boost_queue_test.go +++ b/pkg/kthena-router/datastore/session_boost_queue_test.go @@ -18,6 +18,7 @@ package datastore import ( "context" + "fmt" "sync" "testing" "time" @@ -187,8 +188,16 @@ func TestSessionBoostQueue_MultipleSessions(t *testing.T) { q := newSessionBoostQueue(cfg, nil) defer q.Close() - q.MarkSessionRequestCompleted("conv-A") - q.MarkSessionRequestCompleted("conv-B") + // Use a controllable clock so conv-B's completion is strictly after conv-A's, + // keeping the completion-time ordering below deterministic. + st := q.GetSessionTracker() + base := time.Now() + markCompletedAt := func(sessionID string, completedAt time.Time) { + st.now = func() time.Time { return completedAt } + q.MarkSessionRequestCompleted(sessionID) + } + markCompletedAt("conv-A", base) + markCompletedAt("conv-B", base.Add(time.Second)) now := time.Now() requests := []*Request{ @@ -217,12 +226,13 @@ func TestSessionBoostQueue_MultipleSessions(t *testing.T) { t.Errorf("Last two should not be boosted: third=%v fourth=%v", third.SessionBoost, fourth.SessionBoost) } - // Among boosted: FIFO order (boost-A arrived before boost-B) - if first.UserID != "u2" { - t.Errorf("Expected boost-A first (earlier arrival), got %s", first.UserID) + // Among boosted: the session that completed most recently wins. conv-B was + // marked completed after conv-A, so boost-B (u3) is dequeued before boost-A (u2). + if first.UserID != "u3" { + t.Errorf("Expected boost-B first (most recently completed session), got %s", first.UserID) } - if second.UserID != "u3" { - t.Errorf("Expected boost-B second, got %s", second.UserID) + if second.UserID != "u2" { + t.Errorf("Expected boost-A second, got %s", second.UserID) } // Among normal: FIFO order @@ -234,6 +244,129 @@ func TestSessionBoostQueue_MultipleSessions(t *testing.T) { } } +// TestSessionBoostQueue_OrderedByCompletionTime verifies that boosted requests are +// dequeued by their session's completion time (most recently completed first), +// independent of the order in which the requests arrived in the queue. +func TestSessionBoostQueue_OrderedByCompletionTime(t *testing.T) { + cfg := sessionBoostConfig() + q := newSessionBoostQueue(cfg, nil) + defer q.Close() + + // Drive the tracker with a controllable clock so each completion gets a + // strictly increasing, distinct timestamp regardless of the wall clock's + // resolution or how fast the test executes. This keeps the completion-time + // ordering deterministic instead of falling back to the RequestTime tiebreak. + st := q.GetSessionTracker() + base := time.Now() + markCompletedAt := func(sessionID string, completedAt time.Time) { + st.now = func() time.Time { return completedAt } + q.MarkSessionRequestCompleted(sessionID) + } + + // Completion order: conv-A, then conv-B, then conv-C (conv-C most recent). + markCompletedAt("conv-A", base) + markCompletedAt("conv-B", base.Add(time.Second)) + markCompletedAt("conv-C", base.Add(2*time.Second)) + + now := time.Now() + // Arrival order deliberately differs from completion order: B arrives first, + // then C, then A. This distinguishes completion-time ordering from both FIFO + // (which would yield B, C, A) and LIFO (which would yield A, C, B). + requests := []*Request{ + {UserID: "u-B", ModelName: "m", SessionID: "conv-B", RequestTime: now}, + {UserID: "u-C", ModelName: "m", SessionID: "conv-C", RequestTime: now.Add(time.Millisecond)}, + {UserID: "u-A", ModelName: "m", SessionID: "conv-A", RequestTime: now.Add(2 * time.Millisecond)}, + } + for _, r := range requests { + if err := q.PushRequest(r); err != nil { + t.Fatalf("PushRequest failed: %v", err) + } + } + + // Expect dequeue order by most-recent completion: C, then B, then A. + for i, expected := range []string{"u-C", "u-B", "u-A"} { + got, err := q.popWhenAvailable(context.Background()) + if err != nil { + t.Fatalf("Pop %d failed: %v", i, err) + } + if !got.SessionBoost { + t.Errorf("Position %d: expected boosted request, got non-boosted %s", i, got.UserID) + } + if got.UserID != expected { + t.Errorf("Position %d: expected %s (by completion time), got %s", i, expected, got.UserID) + } + } +} + +// TestSessionBoostQueue_BoundedOvertaking verifies that completion-time ordering +// cannot starve an older boosted request. A steady stream of newer completions is +// injected—one before every dequeue—so a naive newest-first ordering would push +// the original request back forever until it timed out. With bounded overtaking +// the victim must be admitted no later than SessionBoostMaxOvertakes dequeues. +func TestSessionBoostQueue_BoundedOvertaking(t *testing.T) { + cfg := sessionBoostConfig() + cfg.SessionBoostMaxOvertakes = 3 + q := newSessionBoostQueue(cfg, nil) + defer q.Close() + + // Controllable clock so each completion gets a strictly newer timestamp, + // making every injected request rank ahead of the victim under plain + // completion-time ordering. + st := q.GetSessionTracker() + base := time.Now() + tick := 0 + markCompletedNewer := func(sessionID string) { + tick++ + ts := base.Add(time.Duration(tick) * time.Second) + st.now = func() time.Time { return ts } + q.MarkSessionRequestCompleted(sessionID) + } + + // The victim: an older boosted request whose session completed first. + markCompletedNewer("conv-victim") + victim := &Request{UserID: "victim", ModelName: "m", SessionID: "conv-victim", RequestTime: base} + if err := q.PushRequest(victim); err != nil { + t.Fatalf("PushRequest failed: %v", err) + } + + // Before each dequeue, a newer session completes and enqueues a boosted + // follow-up that would outrank the victim by completion time. Without a bound + // the victim would be overtaken indefinitely. + dequeuedAt := -1 + const rounds = 10 + for i := 0; i < rounds; i++ { + sid := fmt.Sprintf("conv-new-%d", i) + markCompletedNewer(sid) + newer := &Request{ + UserID: sid, + ModelName: "m", + SessionID: sid, + RequestTime: base.Add(time.Duration(i+1) * time.Millisecond), + } + if err := q.PushRequest(newer); err != nil { + t.Fatalf("PushRequest %d failed: %v", i, err) + } + + got, err := q.popWhenAvailable(context.Background()) + if err != nil { + t.Fatalf("Pop %d failed: %v", i, err) + } + if got.UserID == "victim" { + dequeuedAt = i + break + } + } + + if dequeuedAt < 0 { + t.Fatalf("victim was starved: not dequeued within %d rounds of newer completions", rounds) + } + // The victim can be overtaken at most SessionBoostMaxOvertakes times, so it must + // be admitted no later than that many dequeues (index <= max). + if dequeuedAt > cfg.SessionBoostMaxOvertakes { + t.Errorf("victim dequeued after %d overtakes, want <= %d", dequeuedAt, cfg.SessionBoostMaxOvertakes) + } +} + func TestSessionBoostQueue_BackpressureDrainsCancelledWhenBusy(t *testing.T) { // Backends are always busy, so tryBackpressureDequeue never pops. Cancelled // requests must still be drained from the heap so the queue size does not stay diff --git a/pkg/kthena-router/datastore/store.go b/pkg/kthena-router/datastore/store.go index 131e68cb8..4d8051c34 100644 --- a/pkg/kthena-router/datastore/store.go +++ b/pkg/kthena-router/datastore/store.go @@ -511,6 +511,14 @@ func applySessionBoostConfigFromEnv(cfg *FairnessQueueConfig) { } } + if v := os.Getenv("SESSION_BOOST_MAX_OVERTAKES"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + cfg.SessionBoostMaxOvertakes = n + } else { + klog.Warningf("Invalid SESSION_BOOST_MAX_OVERTAKES: %q, using default %d", v, cfg.SessionBoostMaxOvertakes) + } + } + if v := os.Getenv("SESSION_BOOST_INFLIGHT_PER_POD"); v != "" { if n, err := strconv.Atoi(v); err == nil && n > 0 { cfg.InflightPerPod = n