Files
llmgw/internal/bench/sampler.go
T
Ilker ManapandClaude Opus 5.5 9853936e0b Give every engine a key only the router holds
Engines listen on every address their node has, with no key. Checked on the
live cluster: from a neighbouring node, a request straight to an engine's
port was answered — past the gateway's authentication, quotas and
accounting, by anyone who could reach the node.

Each replica's engine is now started with a key and refuses requests without
it. The key is derived from the replica id with an HMAC over a secret kept
in the store, rather than stored per replica: the router has to send the
right key to replicas this process never started — adopted after a restart,
or placed by the other control plane — and a function of a shared secret
gives that without a table to keep in step. A key taken from one node opens
that replica only. The secret is made on first use and copied with the rest
of the store when moving to PostgreSQL; a new one would lock the router out
of every running engine at the moment of the move.

The key reaches llama.cpp and vLLM through their environment
(LLAMA_API_KEY, VLLM_API_KEY) rather than their command line, which every
user on the node can read; SGLang only takes it as an argument. The agent's
readiness check, the router's proxy and its metrics scrape send it, the
client's own Authorization still never reaches the engine, and `ctl bench`
fetches the key of the replica it is pointed at from the admin API.

Engines started before this keep running without a key until their replica
is next started. With no store there is nowhere to keep the secret, and
engines start without one, with a warning.

Also formats two files that went in unformatted.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-09-30 00:38:29 +02:00

113 lines
2.4 KiB
Go

package bench
import (
"context"
"net/http"
"sync"
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
)
// sampler reads an engine's metrics during a step.
//
// It is what turns a latency measurement into something admission can use: the
// queue depth and cache pressure observed while the SLO was still being met are
// the values target_waiting and kv_ceiling should have, rather than the round
// numbers somebody typed.
//
// Optional, because a run pointed at the gateway has no single engine to read.
type sampler struct {
url string
apiKey string
engine string
client *http.Client
cancel context.CancelFunc
mu sync.Mutex
wg sync.WaitGroup
waiting float64
running float64
kv float64
n int
}
func newSampler(metricsURL, engine string) *sampler {
return &sampler{
url: metricsURL,
engine: engine,
client: &http.Client{Timeout: 2 * time.Second},
}
}
func (s *sampler) start(ctx context.Context) {
if s.url == "" {
return
}
ctx, s.cancel = context.WithCancel(ctx)
s.wg.Add(1)
go func() {
defer s.wg.Done()
t := time.NewTicker(time.Second)
defer t.Stop()
for {
select {
case <-ctx.Done():
return
case <-t.C:
s.sample(ctx)
}
}
}()
}
func (s *sampler) sample(ctx context.Context) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.url, nil)
if err != nil {
return
}
if s.apiKey != "" {
req.Header.Set("Authorization", "Bearer "+s.apiKey)
}
resp, err := s.client.Do(req)
if err != nil {
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return
}
// The same mapping the router reads with, so a bench result cannot disagree
// with what the scheduler sees. The engine has to be named: reading a
// llama.cpp endpoint with vLLM's names finds nothing, and the run would
// report a queue depth of zero and suggest thresholds built on it.
stats, err := scrape.Parse(resp.Body, scrape.MappingFor(s.engine))
if err != nil {
return
}
s.mu.Lock()
s.waiting += float64(stats.Waiting)
s.running += float64(stats.Running)
s.kv += stats.KVCacheUsage
s.n++
s.mu.Unlock()
}
func (s *sampler) stop() {
if s.cancel != nil {
s.cancel()
}
s.wg.Wait()
}
func (s *sampler) averages() (waiting, running, kv float64) {
s.mu.Lock()
defer s.mu.Unlock()
if s.n == 0 {
return 0, 0, 0
}
n := float64(s.n)
return s.waiting / n, s.running / n, s.kv / n
}