feat(cag): CAG integration — memory adapter, warmup semaphore, cache-aware routing (#47)
CI / CI (push) Successful in 3m20s
CI / CI (push) Successful in 3m20s
## Summary Three issues in one PR. Implements Cache-Augmented Generation (CAG) support in the api-gw. --- ### #44 — Memory ServiceAdapter: standing-context resource - Code-registers `memory` adapter with `standing-context` resource (`POST /memory/standing-context`) - Fixes dispatcher bug: method-level `Auth.Capability` override now respected (was always using verb-derived capability) - Auth: `memory:read` required via method-level override; `memory:write` for ingest - `MEMORY_UPSTREAM_URL` env configures upstream ### #45 — Continuous batching: warmup deprioritisation + per-backend semaphore - In-memory bounded queue per backend: `VLLM_MAX_QUEUE_DEPTH` slots (default 20) - Concurrency limiter: `VLLM_MAX_NUM_SEQS` slots (default 4) - `X-Poimen-Warmup: true` → sleep `WARMUP_DELAY_MS` (default 500ms) before queue entry, strip header, add `X-Priority: low` - Queue full → 503 immediately, no goroutine pile-up - Metrics: `warmup_requests_delayed_total`, `queue_full_requests_total` at `/metrics` ### #46 — Cache-aware routing: rendezvous hash + prefix cache hit logging - `addresses: [...]` YAML field for multi-backend models (backward-compat with `address`) - `X-Poimen-Cache-Key` header → rendezvous-hash backend selection (same key = same GPU) - No cache key → round-robin - Failed backend excluded 30s; next-highest-score fallback - `vllm-num-cached-tokens` response header captured → `vllm_prefix_cache_hit_tokens_total` counter + structured log --- ## Verification ``` go build ./... ✅ go vet ./... ✅ go test -race ./... ✅ (all packages, 0 races) ``` ## Reviewer verdicts - reviewer-sre: SIGN-OFF - reviewer-quality: SIGN-OFF ## Non-blocking findings (🟡🔵 only) - `getOrCreateTransport` not goroutine-safe by design; pre-population invariant documented - `time.Sleep(warmupDelay)` holds goroutine; context-cancel alternative left as TODO - OTel trace_id not propagated into CAG log events (existing gap) Closes #44 Closes #45 Closes #46 --------- Co-authored-by: Admin Bot <[email protected]> Reviewed-on: #47
This commit was merged in pull request #47.
This commit is contained in:
@@ -0,0 +1,157 @@
|
||||
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 {
|
||||
_, failed := s.failedUntil[sc.b.Address]
|
||||
if !failed {
|
||||
return sc.b
|
||||
}
|
||||
expiry := s.failedUntil[sc.b.Address]
|
||||
if now.After(expiry) {
|
||||
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