feat(cag): memory standing-context adapter, warmup semaphore, cache-aware routing
Issue #44: code-register memory adapter with standing-context resource; fix dispatcher method-level Auth.Capability override (was always using verb-derived capability, ignoring per-method Auth field). Issue #45: per-backend bounded in-memory queue (VLLM_MAX_QUEUE_DEPTH=20) + concurrency semaphore (VLLM_MAX_NUM_SEQS=4); warmup requests (X-Poimen-Warmup: true) sleep WARMUP_DELAY_MS before queue entry; /metrics endpoint wired. Issue #46: rendezvous-hash backend selection on X-Poimen-Cache-Key; round-robin fallback; MarkFailed excludes unhealthy backends 30s; ModifyResponse captures vllm-num-cached-tokens + emits structured log + increments counter. Closes #44 Closes #45 Closes #46
This commit is contained in:
@@ -0,0 +1,155 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"hash/fnv"
|
||||
"sort"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrQueueFull is returned when the in-memory queue is at capacity.
|
||||
var ErrQueueFull = errors.New("backend queue full")
|
||||
|
||||
// Backend represents a single vLLM upstream with a bounded in-memory queue.
|
||||
// Two channels model the queue:
|
||||
//
|
||||
// queue — waiting requests (bounded by VLLM_MAX_QUEUE_DEPTH)
|
||||
// slots — concurrent in-service requests (bounded by VLLM_MAX_NUM_SEQS)
|
||||
type Backend struct {
|
||||
Address string
|
||||
queue chan struct{}
|
||||
slots chan struct{}
|
||||
}
|
||||
|
||||
// newBackend creates a Backend with given queue depth and concurrency limit.
|
||||
func newBackend(addr string, maxSlots, maxQueue int) *Backend {
|
||||
slots := make(chan struct{}, maxSlots)
|
||||
for i := 0; i < maxSlots; i++ {
|
||||
slots <- struct{}{}
|
||||
}
|
||||
return &Backend{
|
||||
Address: addr,
|
||||
queue: make(chan struct{}, maxQueue),
|
||||
slots: slots,
|
||||
}
|
||||
}
|
||||
|
||||
// Acquire enters the in-memory queue then waits for a vLLM concurrency slot.
|
||||
// Returns ErrQueueFull immediately if the queue channel is at capacity (no goroutine queued).
|
||||
// Returns ctx.Err() if the context is cancelled while waiting for a slot.
|
||||
func (b *Backend) Acquire(ctx context.Context) error {
|
||||
// Non-blocking queue entry: reject immediately when queue is full.
|
||||
select {
|
||||
case b.queue <- struct{}{}:
|
||||
// entered queue
|
||||
default:
|
||||
return ErrQueueFull
|
||||
}
|
||||
// Wait for a concurrency slot, then leave the queue.
|
||||
select {
|
||||
case <-b.slots:
|
||||
<-b.queue // leave queue; now in service
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
<-b.queue // leave queue; cancelled
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// Release returns a concurrency slot to the pool after the upstream request completes.
|
||||
func (b *Backend) Release() {
|
||||
b.slots <- struct{}{}
|
||||
}
|
||||
|
||||
// Scheduler routes LLM requests to backends using rendezvous hashing (cache-aware)
|
||||
// or round-robin (no cache key). Failed backends are excluded for a configurable window.
|
||||
type Scheduler struct {
|
||||
mu sync.RWMutex
|
||||
backends []*Backend
|
||||
failedUntil map[string]time.Time
|
||||
failWindow time.Duration
|
||||
rrIdx atomic.Uint64
|
||||
}
|
||||
|
||||
// NewScheduler creates a Scheduler for the given addresses.
|
||||
// maxSlots = VLLM_MAX_NUM_SEQS, maxQueue = VLLM_MAX_QUEUE_DEPTH.
|
||||
func NewScheduler(addresses []string, maxSlots, maxQueue int) *Scheduler {
|
||||
backends := make([]*Backend, len(addresses))
|
||||
for i, addr := range addresses {
|
||||
backends[i] = newBackend(addr, maxSlots, maxQueue)
|
||||
}
|
||||
return &Scheduler{
|
||||
backends: backends,
|
||||
failedUntil: make(map[string]time.Time),
|
||||
failWindow: 30 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// Select returns a backend. Uses rendezvous hashing when cacheKey is non-empty
|
||||
// (same key always routes to the same backend, skipping unhealthy ones).
|
||||
// Uses round-robin when cacheKey is empty.
|
||||
func (s *Scheduler) Select(cacheKey string) *Backend {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
if len(s.backends) == 0 {
|
||||
return nil
|
||||
}
|
||||
if len(s.backends) == 1 {
|
||||
return s.backends[0]
|
||||
}
|
||||
if cacheKey != "" {
|
||||
return s.selectRendezvous(cacheKey)
|
||||
}
|
||||
return s.selectRoundRobin()
|
||||
}
|
||||
|
||||
// selectRendezvous picks the healthy backend with the highest FNV32a score
|
||||
// for the key (highest random weight / rendezvous hashing).
|
||||
// Falls back to next-highest score when primary is excluded by MarkFailed.
|
||||
// Must be called with mu held.
|
||||
func (s *Scheduler) selectRendezvous(cacheKey string) *Backend {
|
||||
type scored struct {
|
||||
b *Backend
|
||||
score uint32
|
||||
}
|
||||
now := time.Now()
|
||||
scores := make([]scored, len(s.backends))
|
||||
h := fnv.New32a()
|
||||
for i, b := range s.backends {
|
||||
h.Reset()
|
||||
h.Write([]byte(cacheKey + ":" + b.Address))
|
||||
scores[i] = scored{b: b, score: h.Sum32()}
|
||||
}
|
||||
sort.Slice(scores, func(i, j int) bool {
|
||||
return scores[i].score > scores[j].score
|
||||
})
|
||||
for _, sc := range scores {
|
||||
if t, failed := s.failedUntil[sc.b.Address]; !failed || now.After(t) {
|
||||
if now.After(t) {
|
||||
delete(s.failedUntil, sc.b.Address)
|
||||
}
|
||||
return sc.b
|
||||
}
|
||||
}
|
||||
// All backends unhealthy: return highest-scored (fail-open).
|
||||
return scores[0].b
|
||||
}
|
||||
|
||||
// selectRoundRobin distributes requests evenly when no cache key is present.
|
||||
// Must be called with mu held.
|
||||
func (s *Scheduler) selectRoundRobin() *Backend {
|
||||
idx := s.rrIdx.Add(1) - 1
|
||||
return s.backends[idx%uint64(len(s.backends))]
|
||||
}
|
||||
|
||||
// MarkFailed excludes addr from routing for the fail window (default 30 s).
|
||||
// Called by the proxy ErrorHandler on upstream connection failure.
|
||||
func (s *Scheduler) MarkFailed(addr string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.failedUntil[addr] = time.Now().Add(s.failWindow)
|
||||
}
|
||||
Reference in New Issue
Block a user