Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d71c88498b | ||
|
|
c5a91cf647 | ||
|
|
aac68386d8 | ||
|
|
28753d11ed | ||
|
|
3068c68730 | ||
|
|
da361f2a21 | ||
|
|
1e17706367 | ||
|
|
49c9071a59 | ||
|
|
16ddd6d4ad | ||
|
|
fadc658ac0 | ||
|
|
bb7a490f5c |
@@ -4,9 +4,13 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"math"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
|
||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
|
||||
@@ -21,6 +25,10 @@ import (
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Size the Go scheduler to the container's CPU quota before any server
|
||||
// goroutine starts.
|
||||
configureGOMAXPROCS()
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Initialize OpenTelemetry tracing
|
||||
@@ -177,3 +185,53 @@ func main() {
|
||||
log.Printf("gateway shutdown complete")
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
// configureGOMAXPROCS sets GOMAXPROCS from the cgroup CPU quota so the runtime
|
||||
// sizes its scheduler to the container limit rather than the host CPU count.
|
||||
// The default is left alone when no quota is configured.
|
||||
func configureGOMAXPROCS() {
|
||||
quota, period := readCPUQuota()
|
||||
if quota <= 0 || period <= 0 {
|
||||
return
|
||||
}
|
||||
n := int(math.Ceil(float64(quota) / float64(period)))
|
||||
if n < 1 {
|
||||
n = 1
|
||||
}
|
||||
runtime.GOMAXPROCS(n)
|
||||
log.Printf("GOMAXPROCS set to %d (cgroup cpu quota %d/%d)", n, quota, period)
|
||||
}
|
||||
|
||||
// readCPUQuota returns the CPU quota and period from cgroup v2, falling back to
|
||||
// cgroup v1. A quota of 0 means no limit.
|
||||
func readCPUQuota() (int64, int64) {
|
||||
// cgroup v2: "<quota> <period>" or "max <period>".
|
||||
if data, err := os.ReadFile("/sys/fs/cgroup/cpu.max"); err == nil {
|
||||
fields := strings.Fields(string(data))
|
||||
if len(fields) >= 2 {
|
||||
period, perr := strconv.ParseInt(fields[1], 10, 64)
|
||||
if perr == nil && period > 0 {
|
||||
if fields[0] == "max" {
|
||||
return 0, period
|
||||
}
|
||||
if quota, qerr := strconv.ParseInt(fields[0], 10, 64); qerr == nil {
|
||||
return quota, period
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cgroup v1: cpu.cfs_quota_us is -1 when unlimited.
|
||||
period := int64(100000)
|
||||
if data, err := os.ReadFile("/sys/fs/cgroup/cpu/cpu.cfs_period_us"); err == nil {
|
||||
if v, perr := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64); perr == nil && v > 0 {
|
||||
period = v
|
||||
}
|
||||
}
|
||||
if data, err := os.ReadFile("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"); err == nil {
|
||||
if v, qerr := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64); qerr == nil {
|
||||
return v, period
|
||||
}
|
||||
}
|
||||
return 0, 0
|
||||
}
|
||||
|
||||
@@ -15,7 +15,7 @@ type Metrics struct {
|
||||
|
||||
// Request latencies: request_duration_seconds (histogram)
|
||||
// Stored as cumulative buckets for Prometheus text format
|
||||
requestDuration map[string]int64 // stores duration samples in milliseconds
|
||||
requestDuration map[string]int64 // stores duration samples in milliseconds
|
||||
requestDurationBuckets map[string]map[float64]int64 // histogram buckets
|
||||
|
||||
// Bytes counters: gateway_bytes{direction, route, upstream}
|
||||
@@ -41,22 +41,33 @@ type Metrics struct {
|
||||
warmupRequestsDelayed int64
|
||||
prefixCacheHitTokens int64
|
||||
queueFullRequests int64
|
||||
|
||||
// Admission control metrics (#64)
|
||||
admissionParked int64
|
||||
admissionRejected map[string]int64
|
||||
engineStateStale int64
|
||||
|
||||
// Engine gauge readings mirrored from vLLM (#64)
|
||||
vllmRequestsRunning int64
|
||||
vllmKVUsagePerc float64
|
||||
observerDegraded int
|
||||
}
|
||||
|
||||
// NewMetrics creates a new Metrics instance.
|
||||
func NewMetrics() *Metrics {
|
||||
return &Metrics{
|
||||
requestTotal: make(map[string]int64),
|
||||
requestDuration: make(map[string]int64),
|
||||
requestDurationBuckets: make(map[string]map[float64]int64),
|
||||
bytesIn: make(map[string]int64),
|
||||
bytesOut: make(map[string]int64),
|
||||
upstreamHealth: make(map[string]int),
|
||||
requestTotal: make(map[string]int64),
|
||||
requestDuration: make(map[string]int64),
|
||||
requestDurationBuckets: make(map[string]map[float64]int64),
|
||||
bytesIn: make(map[string]int64),
|
||||
bytesOut: make(map[string]int64),
|
||||
upstreamHealth: make(map[string]int),
|
||||
streamingResponsesTotal: make(map[string]int64),
|
||||
streamingByteCount: make(map[string]int64),
|
||||
ttftMs: make(map[string][]int64),
|
||||
itlMs: make(map[string][]int64),
|
||||
tokenCount: make(map[string]int64),
|
||||
streamingByteCount: make(map[string]int64),
|
||||
ttftMs: make(map[string][]int64),
|
||||
itlMs: make(map[string][]int64),
|
||||
tokenCount: make(map[string]int64),
|
||||
admissionRejected: make(map[string]int64),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -133,7 +144,7 @@ func (m *Metrics) RecordStreamingResponse(route, upstream string, totalBytes int
|
||||
|
||||
func (m *Metrics) recordDuration(key string, duration time.Duration) {
|
||||
m.requestDuration[key] += int64(duration.Milliseconds())
|
||||
|
||||
|
||||
if _, ok := m.requestDurationBuckets[key]; !ok {
|
||||
m.requestDurationBuckets[key] = make(map[float64]int64)
|
||||
}
|
||||
@@ -288,6 +299,12 @@ func (m *Metrics) Reset() {
|
||||
m.warmupRequestsDelayed = 0
|
||||
m.prefixCacheHitTokens = 0
|
||||
m.queueFullRequests = 0
|
||||
m.admissionParked = 0
|
||||
m.admissionRejected = make(map[string]int64)
|
||||
m.engineStateStale = 0
|
||||
m.vllmRequestsRunning = 0
|
||||
m.vllmKVUsagePerc = 0
|
||||
m.observerDegraded = 0
|
||||
}
|
||||
|
||||
// IncrQueueFull increments queue_full_requests_total.
|
||||
@@ -332,6 +349,98 @@ func (m *Metrics) GetPrefixCacheHitTokens() int64 {
|
||||
return m.prefixCacheHitTokens
|
||||
}
|
||||
|
||||
// IncrAdmissionParked increments admission_parked_total.
|
||||
func (m *Metrics) IncrAdmissionParked() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.admissionParked++
|
||||
}
|
||||
|
||||
// GetAdmissionParked returns admission_parked_total.
|
||||
func (m *Metrics) GetAdmissionParked() int64 {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.admissionParked
|
||||
}
|
||||
|
||||
// IncrAdmissionRejected increments admission_rejected_total{reason=reason}.
|
||||
func (m *Metrics) IncrAdmissionRejected(reason string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.admissionRejected[reason]++
|
||||
}
|
||||
|
||||
// GetAdmissionRejected returns a copy of admission_rejected_total by reason.
|
||||
func (m *Metrics) GetAdmissionRejected() map[string]int64 {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
out := make(map[string]int64, len(m.admissionRejected))
|
||||
for reason, count := range m.admissionRejected {
|
||||
out[reason] = count
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// IncrEngineStateStale increments engine_state_stale_total.
|
||||
func (m *Metrics) IncrEngineStateStale() {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.engineStateStale++
|
||||
}
|
||||
|
||||
// GetEngineStateStale returns engine_state_stale_total.
|
||||
func (m *Metrics) GetEngineStateStale() int64 {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.engineStateStale
|
||||
}
|
||||
|
||||
// SetVLLMRequestsRunning sets vllm_requests_running from the gateway's own reading.
|
||||
func (m *Metrics) SetVLLMRequestsRunning(n int64) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.vllmRequestsRunning = n
|
||||
}
|
||||
|
||||
// GetVLLMRequestsRunning returns vllm_requests_running.
|
||||
func (m *Metrics) GetVLLMRequestsRunning() int64 {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.vllmRequestsRunning
|
||||
}
|
||||
|
||||
// SetVLLMKVUsagePerc sets vllm_kv_usage_perc from the gateway's own reading.
|
||||
func (m *Metrics) SetVLLMKVUsagePerc(v float64) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.vllmKVUsagePerc = v
|
||||
}
|
||||
|
||||
// GetVLLMKVUsagePerc returns vllm_kv_usage_perc.
|
||||
func (m *Metrics) GetVLLMKVUsagePerc() float64 {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.vllmKVUsagePerc
|
||||
}
|
||||
|
||||
// SetObserverDegraded sets observer_degraded (1 = startup self-check failed).
|
||||
func (m *Metrics) SetObserverDegraded(degraded bool) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if degraded {
|
||||
m.observerDegraded = 1
|
||||
} else {
|
||||
m.observerDegraded = 0
|
||||
}
|
||||
}
|
||||
|
||||
// GetObserverDegraded returns observer_degraded.
|
||||
func (m *Metrics) GetObserverDegraded() int {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.observerDegraded
|
||||
}
|
||||
|
||||
// globalMetrics is the singleton Metrics instance for gateway-wide telemetry.
|
||||
var (
|
||||
globalMetrics *Metrics
|
||||
|
||||
@@ -48,6 +48,38 @@ func (p *PrometheusExporter) Export() string {
|
||||
lines = append(lines, "# TYPE vllm_prefix_cache_hit_tokens_total counter")
|
||||
lines = append(lines, fmt.Sprintf("vllm_prefix_cache_hit_tokens_total %d", p.metrics.GetPrefixCacheHitTokens()))
|
||||
|
||||
lines = append(lines, "# HELP admission_parked_total Requests parked waiting for engine admission")
|
||||
lines = append(lines, "# TYPE admission_parked_total counter")
|
||||
lines = append(lines, fmt.Sprintf("admission_parked_total %d", p.metrics.GetAdmissionParked()))
|
||||
|
||||
lines = append(lines, "# HELP admission_rejected_total Requests rejected by engine admission control")
|
||||
lines = append(lines, "# TYPE admission_rejected_total counter")
|
||||
rejected := p.metrics.GetAdmissionRejected()
|
||||
reasons := make([]string, 0, len(rejected))
|
||||
for reason := range rejected {
|
||||
reasons = append(reasons, reason)
|
||||
}
|
||||
sort.Strings(reasons)
|
||||
for _, reason := range reasons {
|
||||
lines = append(lines, fmt.Sprintf(`admission_rejected_total{reason="%s"} %d`, reason, rejected[reason]))
|
||||
}
|
||||
|
||||
lines = append(lines, "# HELP engine_state_stale_total Times the engine metrics feed went stale")
|
||||
lines = append(lines, "# TYPE engine_state_stale_total counter")
|
||||
lines = append(lines, fmt.Sprintf("engine_state_stale_total %d", p.metrics.GetEngineStateStale()))
|
||||
|
||||
lines = append(lines, "# HELP vllm_requests_running vLLM requests running as read by the gateway")
|
||||
lines = append(lines, "# TYPE vllm_requests_running gauge")
|
||||
lines = append(lines, fmt.Sprintf("vllm_requests_running %d", p.metrics.GetVLLMRequestsRunning()))
|
||||
|
||||
lines = append(lines, "# HELP vllm_kv_usage_perc vLLM KV cache usage fraction as read by the gateway")
|
||||
lines = append(lines, "# TYPE vllm_kv_usage_perc gauge")
|
||||
lines = append(lines, fmt.Sprintf("vllm_kv_usage_perc %g", p.metrics.GetVLLMKVUsagePerc()))
|
||||
|
||||
lines = append(lines, "# HELP observer_degraded Engine metrics observer startup self-check failed (1) or healthy (0)")
|
||||
lines = append(lines, "# TYPE observer_degraded gauge")
|
||||
lines = append(lines, fmt.Sprintf("observer_degraded %d", p.metrics.GetObserverDegraded()))
|
||||
|
||||
return strings.Join(lines, "\n") + "\n"
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/logging"
|
||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||
)
|
||||
|
||||
// Metric names published by the vLLM engine. Admission is gated on these
|
||||
// rather than on a per-pod semaphore: every gateway replica reads the same
|
||||
// engine, so the limit is global by construction (#64).
|
||||
const (
|
||||
metricRequestsRunning = "vllm:num_requests_running"
|
||||
metricRequestsWaiting = "vllm:num_requests_waiting"
|
||||
metricKVCacheUsage = "vllm:kv_cache_usage_perc"
|
||||
|
||||
// enginePollInterval is how often each backend's /metrics is scraped.
|
||||
enginePollInterval = 2 * time.Second
|
||||
|
||||
// engineStaleThreshold is the number of consecutive failed polls after
|
||||
// which the state is marked unhealthy and admission fails open.
|
||||
engineStaleThreshold = 3
|
||||
|
||||
// engineMetricsTimeout bounds a single scrape.
|
||||
engineMetricsTimeout = 2 * time.Second
|
||||
)
|
||||
|
||||
// EngineState is the gateway's view of one vLLM engine, mirrored from its
|
||||
// own metrics endpoint.
|
||||
type EngineState struct {
|
||||
Running int64
|
||||
Waiting int64
|
||||
KVUsage float64
|
||||
UpdatedAt time.Time
|
||||
Healthy bool
|
||||
}
|
||||
|
||||
// InFlight returns the number of requests the engine currently holds:
|
||||
// running plus waiting.
|
||||
func (s EngineState) InFlight() int64 {
|
||||
return s.Running + s.Waiting
|
||||
}
|
||||
|
||||
// EngineObserver polls one vLLM backend's /metrics and publishes the gauges
|
||||
// admission is gated on. Healthy is false until the first successful poll and
|
||||
// after engineStaleThreshold consecutive failures.
|
||||
type EngineObserver struct {
|
||||
addr string
|
||||
interval time.Duration
|
||||
client *http.Client
|
||||
|
||||
mu sync.Mutex
|
||||
state EngineState
|
||||
failures int
|
||||
degraded bool
|
||||
tick chan struct{}
|
||||
}
|
||||
|
||||
// NewEngineObserver creates an observer for the backend at addr.
|
||||
func NewEngineObserver(addr string, interval time.Duration) *EngineObserver {
|
||||
return &EngineObserver{
|
||||
addr: addr,
|
||||
interval: interval,
|
||||
// Own client: the proxy transports carry socket options and connect
|
||||
// timeouts tuned for inference, not for a short metrics scrape.
|
||||
client: &http.Client{Timeout: engineMetricsTimeout},
|
||||
tick: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Start launches the polling goroutine. It returns immediately; the goroutine
|
||||
// exits when ctx is cancelled.
|
||||
func (o *EngineObserver) Start(ctx context.Context) {
|
||||
go func() {
|
||||
ticker := time.NewTicker(o.interval)
|
||||
defer ticker.Stop()
|
||||
// The first scrape waits one interval rather than firing at Start().
|
||||
// Scraping immediately makes the gateway contact every observed backend
|
||||
// before any request needs one, which is wasted work and a surprising
|
||||
// side effect for anything asserting on upstream traffic. State is
|
||||
// unknown until the first poll, and unknown fails open, so the deferral
|
||||
// costs at most one interval of ungated admission. The startup
|
||||
// self-check (decision Q21) still runs, on that first tick.
|
||||
startup := true
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
o.poll(ctx, startup)
|
||||
startup = false
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Get returns the most recent engine state.
|
||||
func (o *EngineObserver) Get() EngineState {
|
||||
o.mu.Lock()
|
||||
defer o.mu.Unlock()
|
||||
return o.state
|
||||
}
|
||||
|
||||
// Tick returns a channel that is closed and replaced after every successful
|
||||
// poll, waking requests parked on engine admission.
|
||||
func (o *EngineObserver) Tick() <-chan struct{} {
|
||||
o.mu.Lock()
|
||||
defer o.mu.Unlock()
|
||||
return o.tick
|
||||
}
|
||||
|
||||
// Degraded reports whether the startup self-check found the metrics feed
|
||||
// unusable (parse failure or missing series).
|
||||
func (o *EngineObserver) Degraded() bool {
|
||||
o.mu.Lock()
|
||||
defer o.mu.Unlock()
|
||||
return o.degraded
|
||||
}
|
||||
|
||||
// Snapshot returns the latest state together with the tick channel that the
|
||||
// next successful poll will close. Both are read under one lock on purpose:
|
||||
// calling Get and Tick separately can miss the wakeup that lands between the
|
||||
// two calls, parking a request for an extra poll interval.
|
||||
func (o *EngineObserver) Snapshot() (EngineState, <-chan struct{}) {
|
||||
o.mu.Lock()
|
||||
defer o.mu.Unlock()
|
||||
return o.state, o.tick
|
||||
}
|
||||
|
||||
// markDegraded records whether the metrics feed is currently unusable and
|
||||
// mirrors that into observer_degraded. It is cleared on the next parseable
|
||||
// scrape, so one truncated response cannot pin the gauge for the process
|
||||
// lifetime while metric-name drift after the first poll is still reported.
|
||||
func (o *EngineObserver) markDegraded(v bool) {
|
||||
o.mu.Lock()
|
||||
changed := o.degraded != v
|
||||
o.degraded = v
|
||||
o.mu.Unlock()
|
||||
if changed {
|
||||
observability.Global().SetObserverDegraded(v)
|
||||
}
|
||||
}
|
||||
|
||||
// poll scrapes /metrics once. startup marks the first poll, whose failure is a
|
||||
// self-diagnosis signal rather than a transient blip.
|
||||
func (o *EngineObserver) poll(ctx context.Context, startup bool) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://"+o.addr+"/metrics", nil)
|
||||
if err != nil {
|
||||
o.recordFailure()
|
||||
logging.Warnf("engine metrics request build failed", map[string]string{
|
||||
"address": o.addr,
|
||||
"err": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
resp, err := o.client.Do(req)
|
||||
if err != nil {
|
||||
o.recordFailure()
|
||||
logging.Warnf("engine metrics scrape failed", map[string]string{
|
||||
"address": o.addr,
|
||||
"err": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
o.recordFailure()
|
||||
logging.Warnf("engine metrics scrape returned non-200", map[string]string{
|
||||
"address": o.addr,
|
||||
"status": strconv.Itoa(resp.StatusCode),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
running, waiting, kvUsage, err := ParseEngineMetrics(resp.Body)
|
||||
if err != nil {
|
||||
// Metric-name drift must be loud, never a silent fall-back to fail-open
|
||||
// admission (decision Q21). The gauge tracks current state, so it clears
|
||||
// on the next parseable scrape instead of pinning for the pod's life.
|
||||
o.markDegraded(true)
|
||||
if startup {
|
||||
logging.Errorf("engine metrics unusable at startup", err, map[string]string{
|
||||
"address": o.addr,
|
||||
})
|
||||
} else {
|
||||
logging.Warnf("engine metrics parse failed", map[string]string{
|
||||
"address": o.addr,
|
||||
"err": err.Error(),
|
||||
})
|
||||
}
|
||||
o.recordFailure()
|
||||
return
|
||||
}
|
||||
|
||||
o.markDegraded(false)
|
||||
o.recordSuccess(running, waiting, kvUsage)
|
||||
}
|
||||
|
||||
// recordSuccess stores a fresh reading and wakes parked waiters.
|
||||
func (o *EngineObserver) recordSuccess(running, waiting int64, kvUsage float64) {
|
||||
o.mu.Lock()
|
||||
o.state = EngineState{
|
||||
Running: running,
|
||||
Waiting: waiting,
|
||||
KVUsage: kvUsage,
|
||||
UpdatedAt: time.Now(),
|
||||
Healthy: true,
|
||||
}
|
||||
o.failures = 0
|
||||
close(o.tick)
|
||||
o.tick = make(chan struct{})
|
||||
o.mu.Unlock()
|
||||
|
||||
m := observability.Global()
|
||||
m.SetVLLMRequestsRunning(running)
|
||||
m.SetVLLMKVUsagePerc(kvUsage)
|
||||
}
|
||||
|
||||
// recordFailure counts a failed poll and, on the third consecutive failure,
|
||||
// marks the state unhealthy and counts the staleness transition once — not
|
||||
// once per failed poll. Last known values are kept.
|
||||
func (o *EngineObserver) recordFailure() {
|
||||
o.mu.Lock()
|
||||
o.failures++
|
||||
stale := o.failures == engineStaleThreshold
|
||||
if o.failures >= engineStaleThreshold {
|
||||
o.state.Healthy = false
|
||||
}
|
||||
failures := o.failures
|
||||
o.mu.Unlock()
|
||||
|
||||
if stale {
|
||||
observability.Global().IncrEngineStateStale()
|
||||
logging.Warnf("engine metrics stale; admission is failing open", map[string]string{
|
||||
"address": o.addr,
|
||||
"failures": strconv.Itoa(failures),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ParseEngineMetrics extracts the running, waiting and KV-cache-usage gauges
|
||||
// from a vLLM Prometheus exposition body. Comment lines, unrelated series and
|
||||
// non-finite samples are ignored; values are summed across repeated label sets.
|
||||
// It returns an error naming any of the three series that is absent.
|
||||
func ParseEngineMetrics(r io.Reader) (running int64, waiting int64, kvUsage float64, err error) {
|
||||
scanner := bufio.NewScanner(r)
|
||||
// The real exposition body is large; give the scanner room per line.
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
||||
|
||||
var hasRunning, hasWaiting, hasKV bool
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" || strings.HasPrefix(line, "#") {
|
||||
continue
|
||||
}
|
||||
name, value, ok := splitMetricSample(line)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
switch name {
|
||||
case metricRequestsRunning, metricRequestsWaiting, metricKVCacheUsage:
|
||||
default:
|
||||
continue
|
||||
}
|
||||
f, err := strconv.ParseFloat(value, 64)
|
||||
if err != nil || math.IsNaN(f) || math.IsInf(f, 0) {
|
||||
continue
|
||||
}
|
||||
switch name {
|
||||
case metricRequestsRunning:
|
||||
running += int64(f)
|
||||
hasRunning = true
|
||||
case metricRequestsWaiting:
|
||||
waiting += int64(f)
|
||||
hasWaiting = true
|
||||
case metricKVCacheUsage:
|
||||
kvUsage += f
|
||||
hasKV = true
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("read engine metrics: %w", err)
|
||||
}
|
||||
|
||||
var missing []string
|
||||
if !hasRunning {
|
||||
missing = append(missing, metricRequestsRunning)
|
||||
}
|
||||
if !hasWaiting {
|
||||
missing = append(missing, metricRequestsWaiting)
|
||||
}
|
||||
if !hasKV {
|
||||
missing = append(missing, metricKVCacheUsage)
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
return 0, 0, 0, fmt.Errorf("engine metrics missing series: %s", strings.Join(missing, ", "))
|
||||
}
|
||||
return running, waiting, kvUsage, nil
|
||||
}
|
||||
|
||||
// splitMetricSample splits one Prometheus exposition line into its metric name
|
||||
// and its first sample value. Lines without a value are reported as not ok.
|
||||
func splitMetricSample(line string) (name, value string, ok bool) {
|
||||
end := strings.IndexAny(line, "{ \t")
|
||||
if end <= 0 {
|
||||
return "", "", false
|
||||
}
|
||||
name = line[:end]
|
||||
rest := line[end:]
|
||||
if rest[0] == '{' {
|
||||
closeBrace := strings.LastIndex(rest, "}")
|
||||
if closeBrace < 0 {
|
||||
return "", "", false
|
||||
}
|
||||
rest = rest[closeBrace+1:]
|
||||
}
|
||||
fields := strings.Fields(rest)
|
||||
if len(fields) == 0 {
|
||||
return "", "", false
|
||||
}
|
||||
return name, fields[0], true
|
||||
}
|
||||
@@ -3,7 +3,6 @@ package proxy
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
@@ -12,87 +11,37 @@ import (
|
||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||
)
|
||||
|
||||
// LLMMetricsCapture wraps a response writer to capture TTFT and ITL metrics
|
||||
type LLMMetricsCapture struct {
|
||||
writer io.WriteCloser
|
||||
model string
|
||||
metrics *observability.Metrics
|
||||
firstTokenTime time.Time
|
||||
lastTokenTime time.Time
|
||||
requestStartTime time.Time
|
||||
ttftRecorded bool
|
||||
tokenCount int64
|
||||
responseStartTime time.Time
|
||||
}
|
||||
|
||||
// NewLLMMetricsCapture creates a new metrics capture wrapper
|
||||
func NewLLMMetricsCapture(writer io.WriteCloser, model string, metrics *observability.Metrics, startTime time.Time) *LLMMetricsCapture {
|
||||
return &LLMMetricsCapture{
|
||||
writer: writer,
|
||||
model: model,
|
||||
metrics: metrics,
|
||||
requestStartTime: startTime,
|
||||
responseStartTime: time.Now(),
|
||||
}
|
||||
}
|
||||
|
||||
// Write intercepts writes to detect tokens and record metrics
|
||||
func (c *LLMMetricsCapture) Write(p []byte) (int, error) {
|
||||
// Record first token time
|
||||
if !c.ttftRecorded && len(p) > 0 {
|
||||
now := time.Now()
|
||||
ttft := now.Sub(c.requestStartTime).Milliseconds()
|
||||
c.metrics.RecordTTFT(c.model, ttft)
|
||||
c.ttftRecorded = true
|
||||
c.firstTokenTime = now
|
||||
c.lastTokenTime = now
|
||||
}
|
||||
|
||||
// Count tokens in SSE stream (simple: count "data: " lines)
|
||||
if c.ttftRecorded {
|
||||
tokenCount := strings.Count(string(p), "data: ")
|
||||
if tokenCount > 0 {
|
||||
now := time.Now()
|
||||
if !c.firstTokenTime.IsZero() && c.lastTokenTime != now {
|
||||
itl := now.Sub(c.lastTokenTime).Milliseconds()
|
||||
c.metrics.RecordITL(c.model, itl)
|
||||
}
|
||||
c.lastTokenTime = now
|
||||
c.tokenCount += int64(tokenCount)
|
||||
}
|
||||
}
|
||||
|
||||
return c.writer.Write(p)
|
||||
}
|
||||
|
||||
// Close records final metrics and closes writer
|
||||
func (c *LLMMetricsCapture) Close() error {
|
||||
if c.tokenCount > 0 {
|
||||
c.metrics.RecordTokenCount(c.model, c.tokenCount)
|
||||
}
|
||||
return c.writer.Close()
|
||||
}
|
||||
|
||||
// ResponseWriterWrapper wraps http.ResponseWriter to capture metrics
|
||||
// ResponseWriterWrapper wraps an http.ResponseWriter to capture per-request LLM
|
||||
// latency metrics (TTFT, ITL, generated tokens) while a streaming response is
|
||||
// copied from the upstream by httputil.ReverseProxy.
|
||||
//
|
||||
// Only server-sent-event responses are measured. A non-streaming completion is
|
||||
// written in a single Write, so recording it would emit one "TTFT" sample equal
|
||||
// to the whole request latency and no ITL samples at all — numbers that look
|
||||
// plausible but mean nothing.
|
||||
type ResponseWriterWrapper struct {
|
||||
writer http.ResponseWriter
|
||||
statusCode int
|
||||
metrics *observability.Metrics
|
||||
model string
|
||||
startTime time.Time
|
||||
firstByteTime time.Time
|
||||
lastWriteTime time.Time
|
||||
ttftRecorded bool
|
||||
writer http.ResponseWriter
|
||||
statusCode int
|
||||
metrics *observability.Metrics
|
||||
model string
|
||||
startTime time.Time
|
||||
firstByteTime time.Time
|
||||
lastWriteTime time.Time
|
||||
ttftRecorded bool
|
||||
streamChecked bool
|
||||
streaming bool
|
||||
}
|
||||
|
||||
// NewResponseWriterWrapper creates a wrapper for response writer
|
||||
// NewResponseWriterWrapper creates a wrapper for response writer.
|
||||
// startTime must be the time the request entered the gateway, not the time this
|
||||
// wrapper was constructed, otherwise TTFT measures the wrong interval.
|
||||
func NewResponseWriterWrapper(w http.ResponseWriter, model string, metrics *observability.Metrics, startTime time.Time) *ResponseWriterWrapper {
|
||||
return &ResponseWriterWrapper{
|
||||
writer: w,
|
||||
model: model,
|
||||
metrics: metrics,
|
||||
startTime: startTime,
|
||||
statusCode: 200,
|
||||
writer: w,
|
||||
model: model,
|
||||
metrics: metrics,
|
||||
startTime: startTime,
|
||||
statusCode: 200,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -101,29 +50,48 @@ func (w *ResponseWriterWrapper) Header() http.Header {
|
||||
return w.writer.Header()
|
||||
}
|
||||
|
||||
// isStreaming reports whether this is an SSE response. Resolved once from the
|
||||
// Content-Type header, which ReverseProxy has already set before the first write.
|
||||
func (w *ResponseWriterWrapper) isStreaming() bool {
|
||||
if !w.streamChecked {
|
||||
w.streamChecked = true
|
||||
w.streaming = strings.Contains(w.writer.Header().Get("Content-Type"), "text/event-stream")
|
||||
}
|
||||
return w.streaming
|
||||
}
|
||||
|
||||
// countSSETokens counts generated tokens in a chunk of an SSE stream. These
|
||||
// upstreams emit one `data:` event per token, so the event count is the token
|
||||
// count; the `[DONE]` terminator is not a token.
|
||||
func countSSETokens(b []byte) int64 {
|
||||
var n int64
|
||||
for _, line := range strings.Split(string(b), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if strings.HasPrefix(line, "data:") && !strings.HasPrefix(line, "data: [DONE]") {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// Write implements http.ResponseWriter
|
||||
func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
|
||||
// Record TTFT on first write
|
||||
if !w.ttftRecorded && len(b) > 0 {
|
||||
if len(b) > 0 && w.isStreaming() {
|
||||
now := time.Now()
|
||||
ttft := now.Sub(w.startTime).Milliseconds()
|
||||
w.metrics.RecordTTFT(w.model, ttft)
|
||||
w.ttftRecorded = true
|
||||
w.firstByteTime = now
|
||||
w.lastWriteTime = now
|
||||
}
|
||||
|
||||
// Record ITL for subsequent writes (for streaming)
|
||||
if w.ttftRecorded && len(b) > 0 {
|
||||
now := time.Now()
|
||||
if !w.firstByteTime.IsZero() && w.lastWriteTime != now {
|
||||
itl := now.Sub(w.lastWriteTime).Milliseconds()
|
||||
// Only record if ITL > 0 (avoid recording same millisecond twice)
|
||||
if itl > 0 {
|
||||
w.metrics.RecordITL(w.model, itl)
|
||||
}
|
||||
if !w.ttftRecorded {
|
||||
w.metrics.RecordTTFT(w.model, now.Sub(w.startTime).Milliseconds())
|
||||
w.ttftRecorded = true
|
||||
w.firstByteTime = now
|
||||
} else if itl := now.Sub(w.lastWriteTime).Milliseconds(); itl > 0 {
|
||||
// Only record positive ITL: a zero means two chunks landed in the
|
||||
// same millisecond, which carries no timing information.
|
||||
w.metrics.RecordITL(w.model, itl)
|
||||
}
|
||||
w.lastWriteTime = now
|
||||
|
||||
if n := countSSETokens(b); n > 0 {
|
||||
w.metrics.RecordTokenCount(w.model, n)
|
||||
}
|
||||
}
|
||||
|
||||
return w.writer.Write(b)
|
||||
@@ -150,4 +118,8 @@ func (w *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||||
return nil, nil, fmt.Errorf("response writer does not implement Hijacker")
|
||||
}
|
||||
|
||||
|
||||
// Unwrap exposes the underlying writer to http.ResponseController so optional
|
||||
// capabilities (Flush, Hijack, deadlines) keep working through the wrapper.
|
||||
func (w *ResponseWriterWrapper) Unwrap() http.ResponseWriter {
|
||||
return w.writer
|
||||
}
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||
)
|
||||
|
||||
// exportMetrics renders the Prometheus text exposition for m.
|
||||
func exportMetrics(m *observability.Metrics) string {
|
||||
return observability.NewPrometheusExporter(m).Export()
|
||||
}
|
||||
|
||||
// requireSeries asserts the exporter emitted a sample line with the exact value.
|
||||
func requireSeries(t *testing.T, out, want string) {
|
||||
t.Helper()
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
if strings.TrimSpace(line) == want {
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Errorf("missing metrics series %q\n--- got ---\n%s", want, out)
|
||||
}
|
||||
|
||||
// TestResponseWriterWrapperRecordsSSEMetrics guards the regression where the
|
||||
// LLM latency wrapper existed but was never constructed, leaving
|
||||
// llm_ttft_seconds / llm_itl_seconds / llm_tokens_total as HELP-only stubs.
|
||||
func TestResponseWriterWrapperRecordsSSEMetrics(t *testing.T) {
|
||||
m := observability.NewMetrics()
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
rec.Header().Set("Content-Type", "text/event-stream")
|
||||
|
||||
w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now().Add(-50*time.Millisecond))
|
||||
|
||||
if _, err := w.Write([]byte("data: {\"token\":\"a\"}\n\n")); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
if _, err := w.Write([]byte("data: {\"token\":\"b\"}\n\n")); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
// The terminator is not a generated token.
|
||||
if _, err := w.Write([]byte("data: [DONE]\n\n")); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
|
||||
out := exportMetrics(m)
|
||||
requireSeries(t, out, `llm_ttft_seconds_count{model="reasoning"} 1`)
|
||||
requireSeries(t, out, `llm_itl_seconds_count{model="reasoning"} 1`)
|
||||
requireSeries(t, out, `llm_tokens_total{model="reasoning"} 2`)
|
||||
}
|
||||
|
||||
// TestResponseWriterWrapperIgnoresNonStreaming verifies a single-shot JSON
|
||||
// response does not produce a bogus TTFT sample equal to full request latency.
|
||||
func TestResponseWriterWrapperIgnoresNonStreaming(t *testing.T) {
|
||||
m := observability.NewMetrics()
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
rec.Header().Set("Content-Type", "application/json")
|
||||
|
||||
w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now().Add(-50*time.Millisecond))
|
||||
if _, err := w.Write([]byte(`{"choices":[]}`)); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
|
||||
out := exportMetrics(m)
|
||||
if strings.Contains(out, `llm_ttft_seconds_count{model="reasoning"}`) {
|
||||
t.Errorf("non-streaming response recorded TTFT\n--- got ---\n%s", out)
|
||||
}
|
||||
if strings.Contains(out, `llm_tokens_total{model="reasoning"}`) {
|
||||
t.Errorf("non-streaming response recorded tokens\n--- got ---\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponseWriterWrapperPreservesStreaming verifies the wrapper stays
|
||||
// transparent: bytes reach the client and Flush still works, so wrapping the
|
||||
// writer cannot silently break SSE (issue #33).
|
||||
func TestResponseWriterWrapperPreservesStreaming(t *testing.T) {
|
||||
m := observability.NewMetrics()
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
rec.Header().Set("Content-Type", "text/event-stream")
|
||||
|
||||
w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now())
|
||||
payload := "data: {\"token\":\"a\"}\n\n"
|
||||
if _, err := w.Write([]byte(payload)); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
w.Flush()
|
||||
|
||||
if got := rec.Body.String(); got != payload {
|
||||
t.Errorf("body = %q, want %q", got, payload)
|
||||
}
|
||||
if !rec.Flushed {
|
||||
t.Error("Flush did not reach the underlying writer")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandlerRecordsLLMMetricsForStreaming is the end-to-end guard: it drives a
|
||||
// real streaming request through the Handler, so it fails if the wrapper is ever
|
||||
// constructed nowhere (the original defect) rather than merely being broken.
|
||||
func TestHandlerRecordsLLMMetricsForStreaming(t *testing.T) {
|
||||
// Unique model name keeps the observations isolated in the global registry.
|
||||
model := fmt.Sprintf("metrics-test-%d", time.Now().UnixNano())
|
||||
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
rc := http.NewResponseController(w)
|
||||
for _, tok := range []string{"a", "b", "c"} {
|
||||
time.Sleep(2 * time.Millisecond) // guarantee distinct ITL samples
|
||||
fmt.Fprintf(w, "data: {\"token\":%q}\n\n", tok)
|
||||
if err := rc.Flush(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
_ = rc.Flush()
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
cfg := &config.Config{
|
||||
Models: map[string]*config.ModelUpstream{
|
||||
model: {
|
||||
Name: model,
|
||||
Address: strings.TrimPrefix(upstream.URL, "http://"),
|
||||
Path: "/v1/chat/completions",
|
||||
},
|
||||
},
|
||||
Routes: make(map[string]*config.Route),
|
||||
}
|
||||
|
||||
handler := New(cfg)
|
||||
defer handler.Close()
|
||||
srv := httptest.NewServer(handler)
|
||||
defer srv.Close()
|
||||
|
||||
body := fmt.Sprintf(`{"model":%q,"stream":true,"messages":[{"role":"user","content":"hi"}]}`, model)
|
||||
resp, err := http.Post(srv.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
|
||||
if err != nil {
|
||||
t.Fatalf("request failed: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", resp.StatusCode)
|
||||
}
|
||||
if _, err := io.ReadAll(resp.Body); err != nil {
|
||||
t.Fatalf("read body: %v", err)
|
||||
}
|
||||
|
||||
out := exportMetrics(observability.Global())
|
||||
requireSeries(t, out, fmt.Sprintf(`llm_ttft_seconds_count{model=%q} 1`, model))
|
||||
requireSeries(t, out, fmt.Sprintf(`llm_tokens_total{model=%q} 3`, model))
|
||||
|
||||
itlCount := regexp.MustCompile(fmt.Sprintf(`llm_itl_seconds_count\{model=%q\} (\d+)`, model))
|
||||
match := itlCount.FindStringSubmatch(out)
|
||||
if match == nil {
|
||||
t.Errorf("no llm_itl_seconds samples recorded for %s\n--- got ---\n%s", model, out)
|
||||
} else if match[1] == "0" {
|
||||
t.Errorf("llm_itl_seconds_count = 0, want > 0")
|
||||
}
|
||||
}
|
||||
+236
-21
@@ -2,6 +2,7 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -47,6 +48,10 @@ type Handler struct {
|
||||
schedulers map[string]*Scheduler
|
||||
// CAG: delay applied to warmup requests before queue entry (#45)
|
||||
warmupDelay time.Duration
|
||||
// engineObservers polls each model backend's /metrics for admission (#64)
|
||||
engineObservers map[string]*EngineObserver
|
||||
// maxBodyBytes caps synthesized model-route request bodies (MAX_BODY_BYTES)
|
||||
maxBodyBytes int64
|
||||
// transportsMu guards h.transports for goroutine-safe reads/writes.
|
||||
transportsMu sync.RWMutex
|
||||
}
|
||||
@@ -122,18 +127,80 @@ func New(cfg *config.Config) *Handler {
|
||||
maxSlots = n
|
||||
}
|
||||
}
|
||||
maxQueue := 20
|
||||
maxQueue := 8
|
||||
if v := os.Getenv("VLLM_MAX_QUEUE_DEPTH"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
maxQueue = n
|
||||
}
|
||||
}
|
||||
|
||||
// Issue #64: engine-wide admission thresholds. MaxInFlight defaults to the
|
||||
// same value as VLLM_MAX_NUM_SEQS; the observers below read the engine's
|
||||
// own gauges, which every replica shares.
|
||||
admissionMaxInFlight := maxSlots
|
||||
if v := os.Getenv("ADMISSION_MAX_INFLIGHT"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
admissionMaxInFlight = n
|
||||
}
|
||||
}
|
||||
admissionMaxKVUsage := 0.90
|
||||
if v := os.Getenv("ADMISSION_MAX_KV_USAGE"); v != "" {
|
||||
if f, err := strconv.ParseFloat(v, 64); err == nil && f > 0 {
|
||||
admissionMaxKVUsage = f
|
||||
}
|
||||
}
|
||||
maxParkSeconds := 120
|
||||
if v := os.Getenv("ADMISSION_MAX_PARK_SECONDS"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||
maxParkSeconds = n
|
||||
}
|
||||
}
|
||||
h.maxBodyBytes = 1048576
|
||||
if v := os.Getenv("MAX_BODY_BYTES"); v != "" {
|
||||
if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {
|
||||
h.maxBodyBytes = n
|
||||
}
|
||||
}
|
||||
|
||||
limits := AdmissionLimits{
|
||||
MaxInFlight: admissionMaxInFlight,
|
||||
MaxKVUsage: admissionMaxKVUsage,
|
||||
MaxPark: time.Duration(maxParkSeconds) * time.Second,
|
||||
}
|
||||
|
||||
// Issue #64: only backends that actually publish the vLLM scheduler gauges
|
||||
// may be observed. `reasoning` is the only vLLM engine here; ornith and
|
||||
// qwen-cpu are Ollama and embeddings/rerank are TEI, none of which export
|
||||
// vllm:* series. Observing them would trip the startup self-check and hold
|
||||
// observer_degraded at 1 forever. Unset (the default) observes nothing and
|
||||
// leaves admission exactly as it was before this change.
|
||||
observeModels := make(map[string]bool)
|
||||
for _, name := range strings.Split(os.Getenv("ADMISSION_OBSERVE_MODELS"), ",") {
|
||||
if name = strings.TrimSpace(name); name != "" {
|
||||
observeModels[name] = true
|
||||
}
|
||||
}
|
||||
|
||||
h.engineObservers = make(map[string]*EngineObserver)
|
||||
h.schedulers = make(map[string]*Scheduler)
|
||||
for modelName, model := range cfg.Models {
|
||||
addrs := model.AllAddresses()
|
||||
if len(addrs) > 0 {
|
||||
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue)
|
||||
if len(addrs) == 0 {
|
||||
continue
|
||||
}
|
||||
if observeModels[modelName] {
|
||||
for _, addr := range addrs {
|
||||
if _, ok := h.engineObservers[addr]; !ok {
|
||||
h.engineObservers[addr] = NewEngineObserver(addr, enginePollInterval)
|
||||
}
|
||||
}
|
||||
}
|
||||
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue, h.engineObservers, limits)
|
||||
}
|
||||
// The handler has no request-path lifecycle to hang cancellation on, so the
|
||||
// observers poll for the lifetime of the process.
|
||||
for _, o := range h.engineObservers {
|
||||
o.Start(context.Background())
|
||||
}
|
||||
// Pre-populate transports for all scheduler backends at startup (single-threaded).
|
||||
// getOrCreateTransport is not goroutine-safe; pre-populating avoids concurrent writes.
|
||||
@@ -151,7 +218,7 @@ func New(cfg *config.Config) *Handler {
|
||||
// Create a transport per unique upstream address for connection reuse
|
||||
transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)
|
||||
|
||||
upstreamURL, _ := url.Parse("http://" + route.Upstream.Address)
|
||||
upstreamURL := &url.URL{Scheme: "http", Host: route.Upstream.Address}
|
||||
|
||||
r := &Route{
|
||||
Name: name,
|
||||
@@ -314,6 +381,10 @@ func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, d
|
||||
|
||||
// ServeHTTP implements http.Handler.
|
||||
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
// Captured before any work so TTFT measures from the moment the gateway
|
||||
// accepted the request, not from when the upstream response is copied.
|
||||
requestStart := time.Now()
|
||||
|
||||
if h.handleAuthEndpoint(w, r) {
|
||||
return
|
||||
}
|
||||
@@ -358,7 +429,14 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
// For streaming responses (SSE, chunked), disable buffering. Issue #33.
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
|
||||
proxy.ServeHTTP(w, r)
|
||||
// Capture per-request LLM latency metrics (TTFT/ITL/tokens) while the
|
||||
// response is copied from the upstream. Inert for non-SSE responses.
|
||||
var respWriter http.ResponseWriter = w
|
||||
if route.ModelName != "" {
|
||||
respWriter = NewResponseWriterWrapper(w, route.ModelName, observability.Global(), requestStart)
|
||||
}
|
||||
|
||||
proxy.ServeHTTP(respWriter, r)
|
||||
}
|
||||
|
||||
|
||||
@@ -384,27 +462,104 @@ func (h *Handler) applyWarmup(ctx context.Context, r *http.Request) bool {
|
||||
// acquireScheduledBackend selects a backend from sched, applies warmup
|
||||
// deprioritisation, acquires a concurrency slot, and patches route when the
|
||||
// selected backend differs from the config default. It returns the backend
|
||||
// to Release (caller must defer), the (possibly patched) route, and ok=true.
|
||||
// If the slot cannot be acquired it writes a 503 and returns (nil, nil, false).
|
||||
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool) {
|
||||
// to Release (caller must defer), the (possibly patched) route, whether SSE
|
||||
// headers were already written while parking, and ok=true.
|
||||
// If the slot cannot be acquired it writes a 503 and returns ok=false.
|
||||
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool, bool) {
|
||||
backend := sched.Select(cacheKey)
|
||||
if !h.applyWarmup(ctx, r) {
|
||||
return nil, nil, false
|
||||
return nil, nil, false, false
|
||||
}
|
||||
if backend == nil {
|
||||
return nil, route, true
|
||||
return nil, route, false, true
|
||||
}
|
||||
if err := backend.Acquire(ctx); err != nil {
|
||||
if errors.Is(err, ErrQueueFull) {
|
||||
|
||||
// Decision Q22: a request that actually parks on the engine gate and wants
|
||||
// an SSE stream gets its headers pre-written plus a 15 s keep-alive, so no
|
||||
// intermediary times the connection out before the engine admits it.
|
||||
var (
|
||||
prewritten bool
|
||||
hbStop chan struct{}
|
||||
hbDone chan struct{}
|
||||
)
|
||||
onPark := func() {
|
||||
if !requestWantsStream(r) {
|
||||
return
|
||||
}
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
prewritten = true
|
||||
hbStop = make(chan struct{})
|
||||
hbDone = make(chan struct{})
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
fmt.Fprint(w, ": keep-alive\n\n")
|
||||
flusher.Flush()
|
||||
go func() {
|
||||
defer close(hbDone)
|
||||
ticker := time.NewTicker(15 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-hbStop:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
fmt.Fprint(w, ": keep-alive\n\n")
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
err := backend.acquire(ctx, onPark)
|
||||
if prewritten {
|
||||
close(hbStop)
|
||||
<-hbDone
|
||||
}
|
||||
if err != nil {
|
||||
// Once the SSE status line is spent, report the rejection as an event
|
||||
// instead of a second WriteHeader.
|
||||
reject := func(body string) {
|
||||
if prewritten {
|
||||
fmt.Fprintf(w, "event: error\ndata: {\"error\":%q}\n\n", body)
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, "%s", body)
|
||||
}
|
||||
if errors.Is(err, ErrAdmissionTimeout) {
|
||||
observability.Global().IncrAdmissionRejected("park_timeout")
|
||||
logging.Warnf("admission timeout", map[string]string{
|
||||
"backend": backend.Address,
|
||||
"path": r.URL.Path,
|
||||
"model": route.ModelName,
|
||||
})
|
||||
if !prewritten {
|
||||
w.Header().Set("Retry-After", "5")
|
||||
}
|
||||
reject("admission timeout")
|
||||
} else if errors.Is(err, ErrQueueFull) {
|
||||
observability.Global().IncrQueueFull()
|
||||
observability.Global().IncrAdmissionRejected("queue_full")
|
||||
logging.Warnf("backend queue full", map[string]string{
|
||||
"backend": backend.Address,
|
||||
"path": r.URL.Path,
|
||||
"model": route.ModelName,
|
||||
})
|
||||
w.Header().Set("Retry-After", "5")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, "queue full")
|
||||
if !prewritten {
|
||||
w.Header().Set("Retry-After", "5")
|
||||
}
|
||||
reject("queue full")
|
||||
} else {
|
||||
logging.Warnf("backend acquire failed", map[string]string{
|
||||
"backend": backend.Address,
|
||||
@@ -412,15 +567,14 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
|
||||
"model": route.ModelName,
|
||||
"err": err.Error(),
|
||||
})
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, "upstream busy")
|
||||
reject("upstream busy")
|
||||
}
|
||||
return nil, nil, false
|
||||
return nil, nil, prewritten, false
|
||||
}
|
||||
if backend.Address != route.Upstream.Address {
|
||||
patchedCfg := *route.Upstream
|
||||
patchedCfg.Address = backend.Address
|
||||
patchedURL, _ := url.Parse("http://" + backend.Address)
|
||||
patchedURL := &url.URL{Scheme: "http", Host: backend.Address}
|
||||
route = &Route{
|
||||
Name: route.Name,
|
||||
ModelName: route.ModelName,
|
||||
@@ -431,7 +585,27 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
|
||||
},
|
||||
}
|
||||
}
|
||||
return backend, route, true
|
||||
return backend, route, prewritten, true
|
||||
}
|
||||
|
||||
// requestWantsStream reports whether the JSON request body asks for an SSE
|
||||
// stream ("stream": true). The body is restored for the upstream.
|
||||
func requestWantsStream(r *http.Request) bool {
|
||||
if r.Body == nil {
|
||||
return false
|
||||
}
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
r.Body = io.NopCloser(bytes.NewReader(body))
|
||||
var probe struct {
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &probe); err != nil {
|
||||
return false
|
||||
}
|
||||
return probe.Stream
|
||||
}
|
||||
|
||||
// handleAuthEndpoint routes /auth/* requests and returns true if handled.
|
||||
@@ -594,9 +768,10 @@ func (h *Handler) enforceBodyLimit(w http.ResponseWriter, r *http.Request, route
|
||||
// Returns (nil, nil, false) when the backend slot cannot be acquired (response written).
|
||||
func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *http.Request, route *Route, sched *Scheduler, cacheKey string) (*httputil.ReverseProxy, *Backend, bool) {
|
||||
var selectedBackend *Backend
|
||||
var headersPrewritten bool
|
||||
if sched != nil {
|
||||
var ok bool
|
||||
selectedBackend, route, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
|
||||
selectedBackend, route, headersPrewritten, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
|
||||
if !ok {
|
||||
return nil, nil, false
|
||||
}
|
||||
@@ -611,12 +786,52 @@ func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *h
|
||||
|
||||
capturedBackend := selectedBackend
|
||||
capturedRoute := route
|
||||
if headersPrewritten {
|
||||
// The 200 + text/event-stream status is already on the wire, so an
|
||||
// upstream status can no longer be forwarded: httputil.ReverseProxy
|
||||
// calls WriteHeader again and net/http keeps the first one. Frame a
|
||||
// failure as an SSE event rather than letting a raw error body land
|
||||
// inside the stream, which a streaming client reads as an empty
|
||||
// completion and reports as success.
|
||||
proxy.ModifyResponse = func(resp *http.Response) error {
|
||||
if resp.StatusCode < http.StatusBadRequest {
|
||||
return nil
|
||||
}
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
resp.Body.Close()
|
||||
logging.Warnf("upstream error after parked stream started", map[string]string{
|
||||
"upstream": capturedRoute.Upstream.Address,
|
||||
"status": strconv.Itoa(resp.StatusCode),
|
||||
})
|
||||
resp.Body = io.NopCloser(strings.NewReader(fmt.Sprintf(
|
||||
"event: error\ndata: {\"upstream_status\":%d,\"error\":%q}\n\n",
|
||||
resp.StatusCode, strings.TrimSpace(string(body)))))
|
||||
resp.ContentLength = -1
|
||||
resp.Header.Del("Content-Length")
|
||||
resp.Header.Set("Content-Type", "text/event-stream")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
|
||||
if capturedBackend != nil {
|
||||
if s := h.schedulers[capturedRoute.ModelName]; s != nil {
|
||||
s.MarkFailed(capturedBackend.Address)
|
||||
}
|
||||
}
|
||||
if headersPrewritten {
|
||||
// SSE headers and a 200 are already on the wire: a 502 here would be
|
||||
// an ignored WriteHeader and would tear the event stream. Report the
|
||||
// failure as an event instead.
|
||||
fmt.Fprint(w, "event: error\ndata: {\"error\":\"upstream error\"}\n\n")
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
logging.Errorf("upstream error", err, map[string]string{
|
||||
"upstream": capturedRoute.Upstream.Address,
|
||||
"path": r.URL.Path,
|
||||
})
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
fmt.Fprintf(w, "upstream error")
|
||||
logging.Errorf("upstream error", err, map[string]string{
|
||||
|
||||
+112
-12
@@ -69,9 +69,10 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
// Restore the body so it can be read again by the upstream
|
||||
r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
|
||||
// Parse the JSON to find the model field
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(bodyBytes, &payload); err != nil {
|
||||
// Extract only the top-level "model" value with a token walk: the common
|
||||
// path must not materialise the whole body (messages and all) as a map.
|
||||
modelRaw, hasModel, err := extractTopLevelModel(bodyBytes)
|
||||
if err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_json",
|
||||
Message: "request body is not valid JSON",
|
||||
@@ -79,7 +80,6 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
|
||||
// Extract the model name
|
||||
modelVal, hasModel := payload["model"]
|
||||
if !hasModel {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "missing_model",
|
||||
@@ -88,7 +88,7 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
|
||||
// Handle null model
|
||||
if modelVal == nil {
|
||||
if bytes.Equal(bytes.TrimSpace(modelRaw), []byte("null")) {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "null_model",
|
||||
Message: "'model' field is null",
|
||||
@@ -96,8 +96,8 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
|
||||
// Extract as string
|
||||
modelName, ok := modelVal.(string)
|
||||
if !ok {
|
||||
var modelName string
|
||||
if err := json.Unmarshal(modelRaw, &modelName); err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "missing_model",
|
||||
Message: "'model' field must be a string",
|
||||
@@ -122,10 +122,26 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// If upstream expects a different model name, rewrite the body
|
||||
// If upstream expects a different model name, rewrite the body. This path is
|
||||
// uncommon, so it may decode the body as raw fields; the common path above
|
||||
// does not build a full map.
|
||||
if modelUpstream.UpstreamModel != "" && modelUpstream.UpstreamModel != modelName {
|
||||
payload["model"] = modelUpstream.UpstreamModel
|
||||
newBody, err := json.Marshal(payload)
|
||||
var fields map[string]json.RawMessage
|
||||
if err := json.Unmarshal(bodyBytes, &fields); err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_request",
|
||||
Message: fmt.Sprintf("failed to rewrite model name: %v", err),
|
||||
}
|
||||
}
|
||||
newModel, err := json.Marshal(modelUpstream.UpstreamModel)
|
||||
if err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_request",
|
||||
Message: fmt.Sprintf("failed to rewrite model name: %v", err),
|
||||
}
|
||||
}
|
||||
fields["model"] = newModel
|
||||
newBody, err := json.Marshal(fields)
|
||||
if err != nil {
|
||||
return nil, &modelValidationError{
|
||||
Kind: "invalid_request",
|
||||
@@ -151,11 +167,11 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
ConnectTimeout: h.defaultConnectTimeout,
|
||||
ReadTimeout: h.defaultReadTimeout,
|
||||
WriteTimeout: h.defaultWriteTimeout,
|
||||
MaxBodySize: h.defaultMaxBodySize,
|
||||
MaxBodySize: h.maxBodyBytes,
|
||||
AuthRequired: false,
|
||||
}
|
||||
|
||||
targetURL, _ := url.Parse("http://" + modelUpstream.Address)
|
||||
targetURL := &url.URL{Scheme: "http", Host: modelUpstream.Address}
|
||||
|
||||
route := &Route{
|
||||
Name: "v1-chat-" + modelName,
|
||||
@@ -169,3 +185,87 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
||||
|
||||
return route, nil
|
||||
}
|
||||
|
||||
// extractTopLevelModel walks a JSON object body with a streaming decoder and
|
||||
// returns the raw value of the top-level "model" field. It never builds the
|
||||
// full document as a map, so a large "messages" array is skipped by token.
|
||||
// hasModel is false when the key is absent or the document is a bare null.
|
||||
func extractTopLevelModel(body []byte) (model json.RawMessage, hasModel bool, err error) {
|
||||
dec := json.NewDecoder(bytes.NewReader(body))
|
||||
tok, err := dec.Token()
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if tok == nil {
|
||||
// A bare `null` body decoded to a nil map before, i.e. no model field.
|
||||
if _, err := dec.Token(); err != io.EOF {
|
||||
if err == nil {
|
||||
return nil, false, fmt.Errorf("unexpected data after JSON null")
|
||||
}
|
||||
return nil, false, err
|
||||
}
|
||||
return nil, false, nil
|
||||
}
|
||||
if delim, ok := tok.(json.Delim); !ok || delim != '{' {
|
||||
return nil, false, fmt.Errorf("request body is not a JSON object")
|
||||
}
|
||||
for dec.More() {
|
||||
keyTok, err := dec.Token()
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
key, _ := keyTok.(string)
|
||||
if key == "model" {
|
||||
if err := dec.Decode(&model); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
hasModel = true
|
||||
continue
|
||||
}
|
||||
if err := skipJSONValue(dec); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
}
|
||||
// Consume the closing brace and reject trailing data, preserving the
|
||||
// invalid_json answer the old whole-body decode produced.
|
||||
if _, err := dec.Token(); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if _, err := dec.Token(); err != io.EOF {
|
||||
if err == nil {
|
||||
return nil, false, fmt.Errorf("unexpected data after JSON object")
|
||||
}
|
||||
return nil, false, err
|
||||
}
|
||||
return model, hasModel, nil
|
||||
}
|
||||
|
||||
// skipJSONValue advances dec past the next JSON value without materialising it.
|
||||
func skipJSONValue(dec *json.Decoder) error {
|
||||
tok, err := dec.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
delim, ok := tok.(json.Delim)
|
||||
if !ok {
|
||||
return nil // scalar
|
||||
}
|
||||
if delim != '{' && delim != '[' {
|
||||
return fmt.Errorf("unexpected delimiter %v", delim)
|
||||
}
|
||||
depth := 1
|
||||
for depth > 0 {
|
||||
tok, err := dec.Token()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if d, ok := tok.(json.Delim); ok {
|
||||
if d == '{' || d == '[' {
|
||||
depth++
|
||||
} else {
|
||||
depth--
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -8,20 +8,41 @@ import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||
)
|
||||
|
||||
// ErrQueueFull is returned when the in-memory queue is at capacity.
|
||||
var ErrQueueFull = errors.New("backend queue full")
|
||||
|
||||
// ErrAdmissionTimeout is returned when a request parked on engine admission
|
||||
// exceeds the configured maximum park duration.
|
||||
var ErrAdmissionTimeout = errors.New("admission timeout")
|
||||
|
||||
// AdmissionLimits are the engine-wide admission thresholds applied to every
|
||||
// backend of a scheduler (#64).
|
||||
type AdmissionLimits struct {
|
||||
MaxInFlight int
|
||||
MaxKVUsage float64
|
||||
MaxPark time.Duration
|
||||
}
|
||||
|
||||
// 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)
|
||||
//
|
||||
// The engine gate is applied before the queue: admission reads the engine's
|
||||
// own published gauges via observer, which is global across replicas.
|
||||
type Backend struct {
|
||||
Address string
|
||||
queue chan struct{}
|
||||
slots chan struct{}
|
||||
Address string
|
||||
queue chan struct{}
|
||||
slots chan struct{}
|
||||
observer *EngineObserver
|
||||
maxInFlight int
|
||||
maxKVUsage float64
|
||||
maxPark time.Duration
|
||||
}
|
||||
|
||||
// newBackend creates a Backend with given queue depth and concurrency limit.
|
||||
@@ -40,7 +61,48 @@ func newBackend(addr string, maxSlots, maxQueue int) *Backend {
|
||||
// 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.
|
||||
// Returns ErrAdmissionTimeout if the engine gate parks the request past maxPark.
|
||||
func (b *Backend) Acquire(ctx context.Context) error {
|
||||
return b.acquire(ctx, nil)
|
||||
}
|
||||
|
||||
// acquire is Acquire with an optional onPark callback. onPark is invoked once,
|
||||
// on the calling goroutine, the first time the request has to wait for engine
|
||||
// admission; the parked streaming path uses it to start an SSE heartbeat.
|
||||
func (b *Backend) acquire(ctx context.Context, onPark func()) error {
|
||||
// Engine gate (#64): admit on the engine's own gauges, not on this pod's
|
||||
// counters. Failing OPEN on !Healthy is deliberate — a metrics outage must
|
||||
// not become an inference outage; the three-replica gate is only a fairness
|
||||
// improvement on top of the engine's own queueing.
|
||||
if b.observer != nil {
|
||||
deadline := time.Now().Add(b.maxPark)
|
||||
timer := time.NewTimer(time.Until(deadline))
|
||||
defer timer.Stop()
|
||||
parked := false
|
||||
for {
|
||||
// State and tick channel are read together under one lock: taking
|
||||
// them separately can miss a poll that lands between the two calls.
|
||||
st, tick := b.observer.Snapshot()
|
||||
if !st.Healthy || (st.InFlight() < int64(b.maxInFlight) && st.KVUsage < b.maxKVUsage) {
|
||||
break
|
||||
}
|
||||
if !parked {
|
||||
observability.Global().IncrAdmissionParked()
|
||||
parked = true
|
||||
if onPark != nil {
|
||||
onPark()
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-tick:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return ErrAdmissionTimeout
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Non-blocking queue entry: reject immediately when queue is full.
|
||||
select {
|
||||
case b.queue <- struct{}{}:
|
||||
@@ -76,10 +138,15 @@ type Scheduler struct {
|
||||
|
||||
// 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 {
|
||||
// observers holds one engine observer per backend address (may be nil).
|
||||
func NewScheduler(addresses []string, maxSlots, maxQueue int, observers map[string]*EngineObserver, limits AdmissionLimits) *Scheduler {
|
||||
backends := make([]*Backend, len(addresses))
|
||||
for i, addr := range addresses {
|
||||
backends[i] = newBackend(addr, maxSlots, maxQueue)
|
||||
backends[i].observer = observers[addr]
|
||||
backends[i].maxInFlight = limits.MaxInFlight
|
||||
backends[i].maxKVUsage = limits.MaxKVUsage
|
||||
backends[i].maxPark = limits.MaxPark
|
||||
}
|
||||
return &Scheduler{
|
||||
backends: backends,
|
||||
|
||||
@@ -67,7 +67,7 @@ func TestBackendAcquireReturnsErrQueueFullWhenQueueFull(t *testing.T) {
|
||||
// always selects the same backend (rendezvous / highest-random-weight hashing).
|
||||
func TestSchedulerSelectRendezvousHashDeterministic(t *testing.T) {
|
||||
addrs := []string{"a:80", "b:80", "c:80"}
|
||||
s := NewScheduler(addrs, 10, 10)
|
||||
s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})
|
||||
|
||||
first := s.Select("same-cache-key")
|
||||
if first == nil {
|
||||
@@ -85,7 +85,7 @@ func TestSchedulerSelectRendezvousHashDeterministic(t *testing.T) {
|
||||
// requests evenly across all backends when no cache key is given.
|
||||
func TestSchedulerSelectNoKeyRoundRobin(t *testing.T) {
|
||||
addrs := []string{"a:80", "b:80", "c:80"}
|
||||
s := NewScheduler(addrs, 10, 10)
|
||||
s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})
|
||||
|
||||
counts := make(map[string]int)
|
||||
for i := 0; i < 9; i++ {
|
||||
@@ -160,7 +160,7 @@ func TestBackendAcquireCancelledContext(t *testing.T) {
|
||||
// selected within the fail window (30 s default).
|
||||
func TestSchedulerMarkFailedExcludesBackend(t *testing.T) {
|
||||
addrs := []string{"a:80", "b:80"}
|
||||
s := NewScheduler(addrs, 10, 10)
|
||||
s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})
|
||||
|
||||
s.MarkFailed("a:80")
|
||||
|
||||
|
||||
@@ -102,6 +102,22 @@ func GetSessionSpec() *Spec {
|
||||
},
|
||||
},
|
||||
},
|
||||
// sandbox — per-session k8s Pod+PVC+Service lifecycle (§4.15.4)
|
||||
{
|
||||
Name: "sandbox",
|
||||
Methods: []Method{
|
||||
{
|
||||
Verb: "POST",
|
||||
UpstreamPath: "/v1/sessions/sandbox",
|
||||
Auth: &Auth{Required: true, Capability: "sessions:write"},
|
||||
},
|
||||
{
|
||||
Verb: "DELETE",
|
||||
UpstreamPath: "/v1/sessions/sandbox",
|
||||
Auth: &Auth{Required: true, Capability: "sessions:write"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
// Package sessions — FleetClient calls the SandboxFleet controller REST API
|
||||
// to assign sessions to fleet pods. This is the primary sandbox path;
|
||||
// the raw-k8s SandboxManager is the fallback when the fleet controller
|
||||
// is unavailable (e.g. during fleet rollout).
|
||||
package sessions
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultFleetControllerURL = "http://tool-definition-controller.poimen.svc.cluster.local:8788"
|
||||
)
|
||||
|
||||
// FleetClient calls the SandboxFleet controller REST API.
|
||||
type FleetClient struct {
|
||||
baseURL string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// FleetAssignResult is returned by the controller assign endpoint.
|
||||
type FleetAssignResult struct {
|
||||
PodSvcDns string `json:"podSvcDns"`
|
||||
PodName string `json:"podName"`
|
||||
}
|
||||
|
||||
// NewFleetClient creates a FleetClient from env or default URL.
|
||||
// FLEET_CONTROLLER_URL — default: defaultFleetControllerURL
|
||||
func NewFleetClient() *FleetClient {
|
||||
url := os.Getenv("FLEET_CONTROLLER_URL")
|
||||
if url == "" {
|
||||
url = defaultFleetControllerURL
|
||||
}
|
||||
return &FleetClient{
|
||||
baseURL: url,
|
||||
client: &http.Client{Timeout: 35 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// Assign asks the fleet controller to assign sessionID to a fleet pod.
|
||||
// Returns the pod's ClusterIP service DNS and pod name.
|
||||
func (f *FleetClient) Assign(ctx context.Context, sessionID string) (*FleetAssignResult, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
|
||||
f.baseURL+"/sessions/"+sessionID+"/assign", nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("fleet controller unreachable: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 300 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return nil, fmt.Errorf("fleet controller returned %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
var result FleetAssignResult
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return nil, fmt.Errorf("decode: %w", err)
|
||||
}
|
||||
if result.PodSvcDns == "" {
|
||||
return nil, fmt.Errorf("fleet controller returned empty podSvcDns")
|
||||
}
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// Release notifies the fleet controller that a session has ended.
|
||||
// Cleans up the per-session ServiceAccount (audit trail SA).
|
||||
// Non-fatal: fleet controller handles this async.
|
||||
func (f *FleetClient) Release(ctx context.Context, sessionID string) error {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete,
|
||||
f.baseURL+"/sessions/"+sessionID, bytes.NewReader(nil))
|
||||
if err != nil {
|
||||
return fmt.Errorf("build request: %w", err)
|
||||
}
|
||||
resp, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("fleet controller unreachable: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode >= 300 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return fmt.Errorf("fleet controller returned %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -12,9 +12,11 @@ package sessions
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
@@ -30,9 +32,13 @@ const (
|
||||
|
||||
// Handler handles session REST endpoints via internal service calls.
|
||||
// All traffic (reads and writes) goes to graphqlURL (poimen-graphql).
|
||||
// Sandbox creates/deletes pi-session pods via the SandboxFleet controller
|
||||
// (primary) with SandboxManager as direct-k8s fallback.
|
||||
type Handler struct {
|
||||
graphqlURL string
|
||||
client *http.Client
|
||||
sandbox *SandboxManager
|
||||
fleet *FleetClient
|
||||
}
|
||||
|
||||
// NewHandler creates a Handler reading upstream URL from env var.
|
||||
@@ -42,10 +48,34 @@ func NewHandler() *Handler {
|
||||
if gql == "" {
|
||||
gql = defaultGraphQLURL
|
||||
}
|
||||
return &Handler{
|
||||
h := &Handler{
|
||||
graphqlURL: gql,
|
||||
client: &http.Client{Timeout: 30 * time.Second, Transport: tracing.NewTransport(nil)},
|
||||
sandbox: NewSandboxManager(),
|
||||
fleet: NewFleetClient(),
|
||||
}
|
||||
h.startIdleReaper()
|
||||
return h
|
||||
}
|
||||
|
||||
// startIdleReaper launches a background goroutine that deletes exited pi-session
|
||||
// pods (bridge.js idle-timeout triggered) while preserving their PVCs.
|
||||
func (h *Handler) startIdleReaper() {
|
||||
intervalStr := os.Getenv("SANDBOX_REAPER_INTERVAL_MINUTES")
|
||||
interval := defaultReaperInterval
|
||||
if intervalStr != "" {
|
||||
var min int
|
||||
if n, _ := fmt.Sscanf(intervalStr, "%d", &min); n == 1 && min > 0 {
|
||||
interval = time.Duration(min) * time.Minute
|
||||
}
|
||||
}
|
||||
go func() {
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
h.sandbox.ReapIdlePods(context.Background())
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// ServeHTTP routes by X-Upstream-Path set by the dispatcher.
|
||||
@@ -66,6 +96,8 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
h.handleListEpisodes(w, r)
|
||||
case "/v1/sessions/episodes/get":
|
||||
h.handleGetEpisode(w, r)
|
||||
case "/v1/sessions/sandbox":
|
||||
h.handleSandbox(w, r)
|
||||
default:
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "sessions endpoint not found"})
|
||||
}
|
||||
@@ -180,9 +212,12 @@ func (h *Handler) handleList(w http.ResponseWriter, r *http.Request) {
|
||||
searchQuery = fmt.Sprintf("%s cwd:%s", searchQuery, cwd)
|
||||
}
|
||||
|
||||
// entities query (not hybridSearch) — project arg is optional, so it queries
|
||||
// across ALL projects. Required because each session lives in its own
|
||||
// session-<uuid> project; there is no single project to hybridSearch against.
|
||||
gqlQuery := `query($q:String!,$limit:Int!){
|
||||
hybridSearch(query:$q,limit:$limit,entityType:SESSION){
|
||||
entities{id name summary tCreated}
|
||||
entities(query:$q,limit:$limit,entityType:SESSION){
|
||||
entities{id name summary projectId tCreated}
|
||||
searchMode
|
||||
}
|
||||
}`
|
||||
@@ -495,6 +530,94 @@ func (h *Handler) gql(r *http.Request, query string, variables map[string]any) (
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// handleSandbox dispatches POST (create/resume) and DELETE (teardown).
|
||||
// Both take ?session_id= as a query param.
|
||||
func (h *Handler) handleSandbox(w http.ResponseWriter, r *http.Request) {
|
||||
sessionID := r.URL.Query().Get("session_id")
|
||||
if sessionID == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "session_id query param required"})
|
||||
return
|
||||
}
|
||||
|
||||
switch r.Method {
|
||||
case http.MethodPost:
|
||||
h.handleSandboxCreate(w, r, sessionID)
|
||||
case http.MethodDelete:
|
||||
h.handleSandboxDelete(w, r, sessionID)
|
||||
default:
|
||||
writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
|
||||
}
|
||||
}
|
||||
|
||||
// POST ?session_id= — assign session to SandboxFleet pod (primary) or
|
||||
// create a dedicated pod via SandboxManager (fallback).
|
||||
func (h *Handler) handleSandboxCreate(w http.ResponseWriter, r *http.Request, sessionID string) {
|
||||
// Primary: fleet controller assigns session to a pre-warmed pod.
|
||||
if assign, err := h.fleet.Assign(r.Context(), sessionID); err == nil {
|
||||
svcDns := assign.PodSvcDns + fmt.Sprintf(":%d", sandboxPort)
|
||||
log.Printf("sandbox: fleet assigned session=%s pod=%s dns=%s", sessionID, assign.PodName, svcDns)
|
||||
go h.updateSandboxSummary(r, sessionID, assign.PodName, "", svcDns)
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"session_id": sessionID,
|
||||
"pod_name": assign.PodName,
|
||||
"svc_dns": svcDns,
|
||||
"resumed": false,
|
||||
"fleet": true,
|
||||
})
|
||||
return
|
||||
} else {
|
||||
log.Printf("sandbox: fleet unavailable for session=%s (%v), falling back to direct k8s", sessionID, err)
|
||||
}
|
||||
|
||||
// Fallback: create a dedicated pod via SandboxManager.
|
||||
result, err := h.sandbox.CreateOrResume(r.Context(), sessionID)
|
||||
if err != nil {
|
||||
log.Printf("sandbox create/resume session=%s: %v", sessionID, err)
|
||||
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "sandbox error: " + err.Error()})
|
||||
return
|
||||
}
|
||||
go h.updateSandboxSummary(r, sessionID, result.PodName, result.PVCName, result.SvcDNS)
|
||||
log.Printf("sandbox: direct pod session=%s pod=%s resumed=%v dns=%s", sessionID, result.PodName, result.Resumed, result.SvcDNS)
|
||||
writeJSON(w, http.StatusOK, result)
|
||||
}
|
||||
|
||||
// updateSandboxSummary merges pod/svc metadata into the session memory_entity (fire-and-forget).
|
||||
// Uses context.Background() because r.Context() is cancelled once the HTTP handler returns.
|
||||
func (h *Handler) updateSandboxSummary(r *http.Request, sessionID, podName, pvcName, svcDns string) {
|
||||
// Build a detached request carrying just the auth header.
|
||||
detached, _ := http.NewRequestWithContext(context.Background(), r.Method, r.URL.String(), nil)
|
||||
if auth := r.Header.Get("Authorization"); auth != "" {
|
||||
detached.Header.Set("Authorization", auth)
|
||||
}
|
||||
current, _ := h.fetchEntitySummary(detached, "session-"+sessionID, sessionID, "SESSION")
|
||||
if current == nil {
|
||||
current = map[string]any{}
|
||||
}
|
||||
current["pod_name"] = podName
|
||||
current["svc_dns"] = svcDns
|
||||
current["last_active_at"] = time.Now().UTC().Format(time.RFC3339)
|
||||
if pvcName != "" {
|
||||
current["pvc_name"] = pvcName
|
||||
}
|
||||
if b, err := json.Marshal(current); err == nil {
|
||||
h.upsertEntity(r, "session-"+sessionID, sessionID, "SESSION", string(b)) //nolint:errcheck
|
||||
}
|
||||
}
|
||||
|
||||
// DELETE ?session_id= — release from fleet + delete dedicated pod+svc+pvc if any.
|
||||
func (h *Handler) handleSandboxDelete(w http.ResponseWriter, r *http.Request, sessionID string) {
|
||||
// Release from fleet (cleans up per-session SA; non-fatal if session was fleet-based).
|
||||
if err := h.fleet.Release(r.Context(), sessionID); err != nil {
|
||||
log.Printf("sandbox: fleet release session=%s: %v (continuing)", sessionID, err)
|
||||
}
|
||||
// Also attempt direct k8s cleanup (no-ops if resources don't exist).
|
||||
if err := h.sandbox.Delete(r.Context(), sessionID); err != nil {
|
||||
log.Printf("sandbox delete session=%s: %v", sessionID, err)
|
||||
}
|
||||
log.Printf("sandbox: released session=%s", sessionID)
|
||||
writeJSON(w, http.StatusOK, map[string]string{"session_id": sessionID, "status": "deleted"})
|
||||
}
|
||||
|
||||
// EpisodeID computes the deterministic episode ID per §4.14 Q6:
|
||||
// sha256(session_id + milestone_id + attempt).
|
||||
// Exported so callers can verify the value the server would produce.
|
||||
|
||||
@@ -18,7 +18,8 @@ import (
|
||||
// "query" field and records captured request bodies.
|
||||
//
|
||||
// - "upsertEntity" → mutation response
|
||||
// - "hybridSearch" → search response
|
||||
// - "hybridSearch" → search response (get, patch, episodes handlers)
|
||||
// - "entities" → list response (sessions list handler)
|
||||
// - "entity" → single entity response
|
||||
func newGQLServer(t *testing.T, captured *[]byte) *httptest.Server {
|
||||
t.Helper()
|
||||
@@ -33,7 +34,11 @@ func newGQLServer(t *testing.T, captured *[]byte) *httptest.Server {
|
||||
case strings.Contains(bs, "upsertEntity"):
|
||||
fmt.Fprint(w, `{"data":{"upsertEntity":{"id":"test-entity-id","created":true}}}`)
|
||||
case strings.Contains(bs, "hybridSearch"):
|
||||
// get, patch, episodes handlers still use hybridSearch
|
||||
fmt.Fprint(w, `{"data":{"hybridSearch":{"entities":[]}}}`)
|
||||
case strings.Contains(bs, "entities"):
|
||||
// sessions list handler uses entities query
|
||||
fmt.Fprint(w, `{"data":{"entities":{"entities":[]}}}`)
|
||||
case strings.Contains(bs, "entity"):
|
||||
fmt.Fprint(w, `{"data":{"entity":null}}`)
|
||||
default:
|
||||
|
||||
@@ -0,0 +1,492 @@
|
||||
// Package sessions — SandboxManager provisions per-session k8s resources:
|
||||
//
|
||||
// PVC workspace-session-<uuid> Longhorn 1Gi
|
||||
// Pod pi-session-<uuid> single container, /workspace <- PVC
|
||||
// Svc pi-session-<uuid> ClusterIP :8787
|
||||
// NetPol pi-session-sandbox shared for all pi-session pods (created once)
|
||||
//
|
||||
// All k8s API calls use raw HTTP + in-cluster SA token (no client-go).
|
||||
// Same pattern as internal/webhook/promoter.go.
|
||||
package sessions
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
sandboxNamespace = "poimen"
|
||||
sandboxImage = "forgejo.riotpiao.com/riotpiao-poimen/poimen-pi-session-runtime:latest"
|
||||
sandboxPort = 8787
|
||||
defaultPVCSize = "1Gi"
|
||||
defaultStorageClass = "longhorn"
|
||||
defaultReadyTimeout = 30 * time.Second
|
||||
defaultReaperInterval = 5 * time.Minute
|
||||
sandboxSATokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token"
|
||||
sandboxCACertPath = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
|
||||
)
|
||||
|
||||
// SandboxManager manages the k8s lifecycle of pi-session Pods, PVCs, Services.
|
||||
type SandboxManager struct {
|
||||
k8sHost string
|
||||
k8sClient *http.Client
|
||||
pvcSize string
|
||||
readyTimeout time.Duration
|
||||
}
|
||||
|
||||
// SandboxResult is returned by CreateOrResume.
|
||||
type SandboxResult struct {
|
||||
SessionID string `json:"session_id"`
|
||||
PodName string `json:"pod_name"`
|
||||
PVCName string `json:"pvc_name"`
|
||||
SvcDNS string `json:"svc_dns"`
|
||||
Resumed bool `json:"resumed"` // true = pod already existed
|
||||
}
|
||||
|
||||
// NewSandboxManager creates a SandboxManager from env + in-cluster SA files.
|
||||
func NewSandboxManager() *SandboxManager {
|
||||
host := os.Getenv("K8S_API_HOST")
|
||||
if host == "" {
|
||||
host = "https://kubernetes.default.svc"
|
||||
}
|
||||
pvcSize := os.Getenv("SANDBOX_PVC_SIZE")
|
||||
if pvcSize == "" {
|
||||
pvcSize = defaultPVCSize
|
||||
}
|
||||
caCert, _ := os.ReadFile(sandboxCACertPath)
|
||||
return &SandboxManager{
|
||||
k8sHost: host,
|
||||
k8sClient: buildSandboxK8sClient(caCert),
|
||||
pvcSize: pvcSize,
|
||||
readyTimeout: defaultReadyTimeout,
|
||||
}
|
||||
}
|
||||
|
||||
func buildSandboxK8sClient(caCert []byte) *http.Client {
|
||||
transport := &http.Transport{}
|
||||
if len(caCert) > 0 {
|
||||
pool := x509.NewCertPool()
|
||||
pool.AppendCertsFromPEM(caCert)
|
||||
transport.TLSClientConfig = &tls.Config{RootCAs: pool}
|
||||
} else {
|
||||
log.Printf("sandbox: WARNING no SA CA cert at %s — using InsecureSkipVerify", sandboxCACertPath)
|
||||
transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} //nolint:gosec
|
||||
}
|
||||
return &http.Client{Timeout: 30 * time.Second, Transport: transport}
|
||||
}
|
||||
|
||||
func sandboxSAToken() string {
|
||||
b, err := os.ReadFile(sandboxSATokenPath)
|
||||
if err != nil {
|
||||
log.Printf("sandbox: WARNING failed to read SA token %s: %v", sandboxSATokenPath, err)
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
|
||||
func pvcName(sessionID string) string { return "workspace-session-" + sessionID }
|
||||
func podName(sessionID string) string { return "pi-session-" + sessionID }
|
||||
func svcName(sessionID string) string { return "pi-session-" + sessionID }
|
||||
func svcDNS(sessionID string) string {
|
||||
return fmt.Sprintf("%s.%s.svc.cluster.local:%d", svcName(sessionID), sandboxNamespace, sandboxPort)
|
||||
}
|
||||
|
||||
// ── public API ────────────────────────────────────────────────────────────
|
||||
|
||||
// CreateOrResume provisions or resumes a sandbox for sessionID.
|
||||
// Idempotent: if the pod is already Running, returns its DNS immediately.
|
||||
func (s *SandboxManager) CreateOrResume(ctx context.Context, sessionID string) (*SandboxResult, error) {
|
||||
res := &SandboxResult{
|
||||
SessionID: sessionID,
|
||||
PodName: podName(sessionID),
|
||||
PVCName: pvcName(sessionID),
|
||||
SvcDNS: svcDNS(sessionID),
|
||||
}
|
||||
|
||||
// 1. Ensure PVC exists
|
||||
if err := s.ensurePVC(ctx, sessionID); err != nil {
|
||||
return nil, fmt.Errorf("ensure PVC: %w", err)
|
||||
}
|
||||
|
||||
// 2. Check if pod already running (resume case)
|
||||
phase, err := s.getPodPhase(ctx, sessionID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("check pod: %w", err)
|
||||
}
|
||||
if phase == "Running" {
|
||||
res.Resumed = true
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// 3. Ensure NetworkPolicy for all pi-session pods exists
|
||||
if err := s.ensureNetworkPolicy(ctx); err != nil {
|
||||
log.Printf("sandbox: WARNING failed to ensure network policy: %v", err)
|
||||
}
|
||||
|
||||
// 4. Create Pod + Service
|
||||
if err := s.createPod(ctx, sessionID); err != nil {
|
||||
return nil, fmt.Errorf("create pod: %w", err)
|
||||
}
|
||||
if err := s.ensureService(ctx, sessionID); err != nil {
|
||||
return nil, fmt.Errorf("ensure service: %w", err)
|
||||
}
|
||||
|
||||
// 5. Wait for Ready
|
||||
if err := s.waitPodReady(ctx, sessionID); err != nil {
|
||||
return nil, fmt.Errorf("wait ready: %w", err)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// Delete removes the Pod, Service, and PVC for sessionID.
|
||||
func (s *SandboxManager) Delete(ctx context.Context, sessionID string) error {
|
||||
var errs []string
|
||||
for _, fn := range []func(context.Context, string) error{
|
||||
s.deletePod, s.deleteService, s.deletePVC,
|
||||
} {
|
||||
if err := fn(ctx, sessionID); err != nil && !isNotFound(err) {
|
||||
errs = append(errs, err.Error())
|
||||
}
|
||||
}
|
||||
if len(errs) > 0 {
|
||||
return fmt.Errorf("delete sandbox: %s", strings.Join(errs, "; "))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReapIdlePods lists all pi-session-* pods and deletes those that have exited
|
||||
// (Succeeded/Failed) — bridge.js self-exits after idle timeout.
|
||||
// The PVC is preserved so the next CreateOrResume reattaches the same workspace.
|
||||
func (s *SandboxManager) ReapIdlePods(ctx context.Context) {
|
||||
pods, err := s.listSandboxPods(ctx)
|
||||
if err != nil {
|
||||
log.Printf("sandbox reaper: list pods: %v", err)
|
||||
return
|
||||
}
|
||||
for _, podAny := range pods {
|
||||
pod, _ := podAny.(map[string]any)
|
||||
phase, _ := pod["status"].(map[string]any)["phase"].(string)
|
||||
name, _ := pod["metadata"].(map[string]any)["name"].(string)
|
||||
if phase == "Succeeded" || phase == "Failed" {
|
||||
log.Printf("sandbox reaper: deleting exited pod %s (phase=%s)", name, phase)
|
||||
if err := s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+name); err != nil && !isNotFound(err) {
|
||||
log.Printf("sandbox reaper: delete pod %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── k8s resource helpers ─────────────────────────────────────────────────
|
||||
|
||||
func (s *SandboxManager) ensurePVC(ctx context.Context, sessionID string) error {
|
||||
name := pvcName(sessionID)
|
||||
_, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/persistentvolumeclaims/"+name)
|
||||
if err == nil {
|
||||
return nil // already exists
|
||||
}
|
||||
if !isNotFound(err) {
|
||||
return err
|
||||
}
|
||||
// Create
|
||||
body := map[string]any{
|
||||
"apiVersion": "v1",
|
||||
"kind": "PersistentVolumeClaim",
|
||||
"metadata": map[string]any{
|
||||
"name": name, "namespace": sandboxNamespace,
|
||||
"labels": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||
},
|
||||
"spec": map[string]any{
|
||||
"accessModes": []string{"ReadWriteOnce"},
|
||||
"storageClassName": defaultStorageClass,
|
||||
"resources": map[string]any{
|
||||
"requests": map[string]string{"storage": s.pvcSize},
|
||||
},
|
||||
},
|
||||
}
|
||||
return s.k8sCreate(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/persistentvolumeclaims", body)
|
||||
}
|
||||
|
||||
func (s *SandboxManager) getPodPhase(ctx context.Context, sessionID string) (string, error) {
|
||||
pod, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+podName(sessionID))
|
||||
if isNotFound(err) {
|
||||
return "", nil
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
phase, _ := pod["status"].(map[string]any)["phase"].(string)
|
||||
return phase, nil
|
||||
}
|
||||
|
||||
func (s *SandboxManager) createPod(ctx context.Context, sessionID string) error {
|
||||
body := map[string]any{
|
||||
"apiVersion": "v1",
|
||||
"kind": "Pod",
|
||||
"metadata": map[string]any{
|
||||
"name": podName(sessionID), "namespace": sandboxNamespace,
|
||||
"labels": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||
},
|
||||
"spec": map[string]any{
|
||||
"restartPolicy": "Never",
|
||||
"automountServiceAccountToken": false,
|
||||
"securityContext": map[string]any{
|
||||
"runAsNonRoot": true,
|
||||
"runAsUser": int64(1000),
|
||||
"seccompProfile": map[string]string{"type": "RuntimeDefault"},
|
||||
},
|
||||
"containers": []map[string]any{{
|
||||
"name": "pi-session",
|
||||
"image": sandboxImage,
|
||||
"ports": []map[string]any{{"containerPort": sandboxPort}},
|
||||
"securityContext": map[string]any{
|
||||
"readOnlyRootFilesystem": true,
|
||||
"allowPrivilegeEscalation": false,
|
||||
"capabilities": map[string]any{"drop": []string{"ALL"}},
|
||||
},
|
||||
"resources": map[string]any{
|
||||
"requests": map[string]string{"cpu": "500m", "memory": "512Mi"},
|
||||
"limits": map[string]string{"cpu": "1", "memory": "1Gi"},
|
||||
},
|
||||
"readinessProbe": map[string]any{
|
||||
"tcpSocket": map[string]any{"port": sandboxPort},
|
||||
"initialDelaySeconds": 3,
|
||||
"periodSeconds": 2,
|
||||
"failureThreshold": 15,
|
||||
},
|
||||
"volumeMounts": []map[string]any{
|
||||
{"name": "workspace", "mountPath": "/workspace"},
|
||||
{"name": "tmp", "mountPath": "/tmp"},
|
||||
},
|
||||
}},
|
||||
"volumes": []map[string]any{
|
||||
{"name": "workspace", "persistentVolumeClaim": map[string]string{"claimName": pvcName(sessionID)}},
|
||||
{"name": "tmp", "emptyDir": map[string]any{}},
|
||||
},
|
||||
},
|
||||
}
|
||||
return s.k8sCreate(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods", body)
|
||||
}
|
||||
|
||||
func (s *SandboxManager) ensureService(ctx context.Context, sessionID string) error {
|
||||
name := svcName(sessionID)
|
||||
_, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/services/"+name)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if !isNotFound(err) {
|
||||
return err
|
||||
}
|
||||
body := map[string]any{
|
||||
"apiVersion": "v1",
|
||||
"kind": "Service",
|
||||
"metadata": map[string]any{
|
||||
"name": name, "namespace": sandboxNamespace,
|
||||
"labels": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||
},
|
||||
"spec": map[string]any{
|
||||
"type": "ClusterIP",
|
||||
"selector": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||
"ports": []map[string]any{{
|
||||
"port": sandboxPort, "targetPort": sandboxPort, "protocol": "TCP",
|
||||
}},
|
||||
},
|
||||
}
|
||||
return s.k8sCreate(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/services", body)
|
||||
}
|
||||
|
||||
func (s *SandboxManager) ensureNetworkPolicy(ctx context.Context) error {
|
||||
const npName = "pi-session-sandbox"
|
||||
_, err := s.k8sGet(ctx, "/apis/networking.k8s.io/v1/namespaces/"+sandboxNamespace+"/networkpolicies/"+npName)
|
||||
if err == nil {
|
||||
return nil // already exists
|
||||
}
|
||||
if !isNotFound(err) {
|
||||
return err
|
||||
}
|
||||
body := map[string]any{
|
||||
"apiVersion": "networking.k8s.io/v1",
|
||||
"kind": "NetworkPolicy",
|
||||
"metadata": map[string]any{
|
||||
"name": npName, "namespace": sandboxNamespace,
|
||||
"labels": map[string]string{"app": "pi-session"},
|
||||
},
|
||||
"spec": map[string]any{
|
||||
"podSelector": map[string]any{"matchLabels": map[string]string{"app": "pi-session"}},
|
||||
"policyTypes": []string{"Ingress", "Egress"},
|
||||
"ingress": []map[string]any{{
|
||||
"from": []map[string]any{{
|
||||
"podSelector": map[string]any{"matchLabels": map[string]string{"app": "poimen-frontend"}},
|
||||
}},
|
||||
"ports": []map[string]any{{"protocol": "TCP", "port": sandboxPort}},
|
||||
}},
|
||||
"egress": []map[string]any{
|
||||
// DNS
|
||||
{
|
||||
"to": []map[string]any{{"namespaceSelector": map[string]any{"matchLabels": map[string]string{"kubernetes.io/metadata.name": "kube-system"}}}},
|
||||
"ports": []map[string]any{{"protocol": "UDP", "port": 53}, {"protocol": "TCP", "port": 53}},
|
||||
},
|
||||
// api-gateway (LLM inference)
|
||||
{
|
||||
"to": []map[string]any{{"namespaceSelector": map[string]any{"matchLabels": map[string]string{"kubernetes.io/metadata.name": "api"}}}},
|
||||
"ports": []map[string]any{{"protocol": "TCP", "port": 8080}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
return s.k8sCreate(ctx, "/apis/networking.k8s.io/v1/namespaces/"+sandboxNamespace+"/networkpolicies", body)
|
||||
}
|
||||
|
||||
func (s *SandboxManager) waitPodReady(ctx context.Context, sessionID string) error {
|
||||
deadline := time.Now().Add(s.readyTimeout)
|
||||
for time.Now().Before(deadline) {
|
||||
pod, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+podName(sessionID))
|
||||
if err != nil {
|
||||
return fmt.Errorf("poll pod: %w", err)
|
||||
}
|
||||
status, _ := pod["status"].(map[string]any)
|
||||
phase, _ := status["phase"].(string)
|
||||
if phase == "Failed" || phase == "Succeeded" {
|
||||
return fmt.Errorf("pod entered terminal phase %s before ready", phase)
|
||||
}
|
||||
if phase == "Running" {
|
||||
for _, cond := range toSlice(status["conditions"]) {
|
||||
if c, ok := cond.(map[string]any); ok {
|
||||
if c["type"] == "Ready" && c["status"] == "True" {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-time.After(2 * time.Second):
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("pod not ready after %v", s.readyTimeout)
|
||||
}
|
||||
|
||||
func (s *SandboxManager) deletePod(ctx context.Context, sessionID string) error {
|
||||
return s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+podName(sessionID))
|
||||
}
|
||||
func (s *SandboxManager) deleteService(ctx context.Context, sessionID string) error {
|
||||
return s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/services/"+svcName(sessionID))
|
||||
}
|
||||
func (s *SandboxManager) deletePVC(ctx context.Context, sessionID string) error {
|
||||
return s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/persistentvolumeclaims/"+pvcName(sessionID))
|
||||
}
|
||||
|
||||
func (s *SandboxManager) listSandboxPods(ctx context.Context) ([]any, error) {
|
||||
body, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods?labelSelector=app%3Dpi-session")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items, _ := body["items"].([]any)
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// ── raw k8s HTTP ─────────────────────────────────────────────────────────
|
||||
|
||||
func (s *SandboxManager) k8sGet(ctx context.Context, path string) (map[string]any, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.k8sHost+path, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+sandboxSAToken())
|
||||
resp, err := s.k8sClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("k8s GET %s: %w", path, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode == 404 {
|
||||
return nil, &k8sNotFound{path: path}
|
||||
}
|
||||
if resp.StatusCode >= 400 {
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return nil, fmt.Errorf("k8s GET %s returned %d: %s", path, resp.StatusCode, b)
|
||||
}
|
||||
var result map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return nil, fmt.Errorf("decode: %w", err)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *SandboxManager) k8sCreate(ctx context.Context, path string, body any) error {
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal: %w", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.k8sHost+path, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+sandboxSAToken())
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := s.k8sClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("k8s POST %s: %w", path, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode == 409 { // AlreadyExists — idempotent
|
||||
return nil
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return fmt.Errorf("k8s POST %s returned %d: %s", path, resp.StatusCode, b)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SandboxManager) k8sDelete(ctx context.Context, path string) error {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, s.k8sHost+path, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+sandboxSAToken())
|
||||
resp, err := s.k8sClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("k8s DELETE %s: %w", path, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode == 404 {
|
||||
return &k8sNotFound{path: path}
|
||||
}
|
||||
if resp.StatusCode >= 400 {
|
||||
b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return fmt.Errorf("k8s DELETE %s returned %d: %s", path, resp.StatusCode, b)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ── error helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
type k8sNotFound struct{ path string }
|
||||
|
||||
func (e *k8sNotFound) Error() string { return "not found: " + e.path }
|
||||
|
||||
func isNotFound(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
_, ok := err.(*k8sNotFound)
|
||||
return ok
|
||||
}
|
||||
|
||||
func toSlice(v any) []any {
|
||||
if s, ok := v.([]any); ok {
|
||||
return s
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
package sessions
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ── helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
// runningReadyPod returns minimal k8s pod JSON: phase Running + Ready condition.
|
||||
func runningReadyPod(name string) map[string]any {
|
||||
return map[string]any{
|
||||
"metadata": map[string]any{"name": name},
|
||||
"status": map[string]any{
|
||||
"phase": "Running",
|
||||
"conditions": []any{
|
||||
map[string]any{"type": "Ready", "status": "True"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// respondK8sJSON writes a JSON body with the given status code.
|
||||
func respondK8sJSON(w http.ResponseWriter, status int, body any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
json.NewEncoder(w).Encode(body) //nolint:errcheck
|
||||
}
|
||||
|
||||
// newTestSM builds a SandboxManager pointed at the given test server URL.
|
||||
func newTestSM(url string) *SandboxManager {
|
||||
return &SandboxManager{
|
||||
k8sHost: url,
|
||||
k8sClient: &http.Client{Timeout: 5 * time.Second},
|
||||
pvcSize: "1Gi",
|
||||
readyTimeout: 3 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestSandboxCreateOrResume_CreateNew ───────────────────────────────────────
|
||||
|
||||
// No existing resources → creates PVC, NetworkPolicy, Pod, Service; Resumed=false.
|
||||
func TestSandboxCreateOrResume_CreateNew(t *testing.T) {
|
||||
id := "newid"
|
||||
var podGetCount int32
|
||||
|
||||
mux := http.NewServeMux()
|
||||
|
||||
// PVC: GET → 404, POST → 201
|
||||
mux.HandleFunc("GET /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
})
|
||||
mux.HandleFunc("POST /api/v1/namespaces/poimen/persistentvolumeclaims", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
})
|
||||
|
||||
// NetworkPolicy (fixed name "pi-session-sandbox"): GET → 404, POST → 201
|
||||
mux.HandleFunc("GET /apis/networking.k8s.io/v1/namespaces/poimen/networkpolicies/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
})
|
||||
mux.HandleFunc("POST /apis/networking.k8s.io/v1/namespaces/poimen/networkpolicies", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
})
|
||||
|
||||
// Pod: first GET (getPodPhase) → 404; second GET (waitPodReady) → Running+Ready; POST → 201
|
||||
mux.HandleFunc("GET /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
if atomic.AddInt32(&podGetCount, 1) == 1 {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
respondK8sJSON(w, http.StatusOK, runningReadyPod(podName(id)))
|
||||
})
|
||||
mux.HandleFunc("POST /api/v1/namespaces/poimen/pods", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
})
|
||||
|
||||
// Service: GET → 404, POST → 201
|
||||
mux.HandleFunc("GET /api/v1/namespaces/poimen/services/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
})
|
||||
mux.HandleFunc("POST /api/v1/namespaces/poimen/services", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
result, err := newTestSM(srv.URL).CreateOrResume(context.Background(), id)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateOrResume: %v", err)
|
||||
}
|
||||
if want := svcDNS(id); result.SvcDNS != want {
|
||||
t.Errorf("SvcDNS: want %q, got %q", want, result.SvcDNS)
|
||||
}
|
||||
if result.Resumed {
|
||||
t.Error("want Resumed=false, got true")
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestSandboxCreateOrResume_Resume ─────────────────────────────────────────
|
||||
|
||||
// PVC exists + pod Running → Resumed=true, no POST calls.
|
||||
func TestSandboxCreateOrResume_Resume(t *testing.T) {
|
||||
id := "resumeid"
|
||||
|
||||
mux := http.NewServeMux()
|
||||
|
||||
// PVC exists
|
||||
mux.HandleFunc("GET /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
respondK8sJSON(w, http.StatusOK, map[string]any{"metadata": map[string]any{"name": pvcName(id)}})
|
||||
})
|
||||
|
||||
// Pod exists and Running (getPodPhase uses GET /{name})
|
||||
mux.HandleFunc("GET /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
respondK8sJSON(w, http.StatusOK, map[string]any{
|
||||
"metadata": map[string]any{"name": podName(id)},
|
||||
"status": map[string]any{"phase": "Running"},
|
||||
})
|
||||
})
|
||||
|
||||
// Any POST → 500 to assert no creates are issued
|
||||
mux.HandleFunc("POST /", func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Errorf("unexpected POST to %s", r.URL.Path)
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
result, err := newTestSM(srv.URL).CreateOrResume(context.Background(), id)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateOrResume: %v", err)
|
||||
}
|
||||
if !result.Resumed {
|
||||
t.Error("want Resumed=true, got false")
|
||||
}
|
||||
if want := svcDNS(id); result.SvcDNS != want {
|
||||
t.Errorf("SvcDNS: want %q, got %q", want, result.SvcDNS)
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestSandboxDelete ─────────────────────────────────────────────────────────
|
||||
|
||||
// All DELETEs succeed → no error returned.
|
||||
func TestSandboxDelete(t *testing.T) {
|
||||
id := "deleteid"
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/services/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
if err := newTestSM(srv.URL).Delete(context.Background(), id); err != nil {
|
||||
t.Errorf("Delete: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestSandboxDelete_NotFound ────────────────────────────────────────────────
|
||||
|
||||
// All DELETEs return 404 → isNotFound errors swallowed → no error returned.
|
||||
func TestSandboxDelete_NotFound(t *testing.T) {
|
||||
id := "ghostid"
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
})
|
||||
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/services/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
})
|
||||
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
if err := newTestSM(srv.URL).Delete(context.Background(), id); err != nil {
|
||||
t.Errorf("Delete_NotFound: want nil, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestSandboxReapIdlePods ───────────────────────────────────────────────────
|
||||
|
||||
// pod-a=Succeeded → deleted; pod-b=Running → skipped.
|
||||
func TestSandboxReapIdlePods(t *testing.T) {
|
||||
var deletedA int32
|
||||
|
||||
mux := http.NewServeMux()
|
||||
|
||||
// List all pi-session pods (query string not part of mux path matching)
|
||||
mux.HandleFunc("GET /api/v1/namespaces/poimen/pods", func(w http.ResponseWriter, _ *http.Request) {
|
||||
respondK8sJSON(w, http.StatusOK, map[string]any{
|
||||
"items": []any{
|
||||
map[string]any{
|
||||
"metadata": map[string]any{"name": "pod-a"},
|
||||
"status": map[string]any{"phase": "Succeeded"},
|
||||
},
|
||||
map[string]any{
|
||||
"metadata": map[string]any{"name": "pod-b"},
|
||||
"status": map[string]any{"phase": "Running"},
|
||||
},
|
||||
},
|
||||
})
|
||||
})
|
||||
|
||||
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, r *http.Request) {
|
||||
name := r.PathValue("name")
|
||||
if name == "pod-b" {
|
||||
t.Errorf("pod-b (Running) must NOT be deleted")
|
||||
}
|
||||
if name == "pod-a" {
|
||||
atomic.AddInt32(&deletedA, 1)
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
|
||||
srv := httptest.NewServer(mux)
|
||||
defer srv.Close()
|
||||
|
||||
newTestSM(srv.URL).ReapIdlePods(context.Background())
|
||||
|
||||
if got := atomic.LoadInt32(&deletedA); got != 1 {
|
||||
t.Errorf("pod-a delete count: want 1, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestSandboxNameHelpers ────────────────────────────────────────────────────
|
||||
|
||||
func TestSandboxNameHelpers(t *testing.T) {
|
||||
cases := []struct{ fn, want, got string }{
|
||||
{"pvcName", "workspace-session-abc", pvcName("abc")},
|
||||
{"podName", "pi-session-abc", podName("abc")},
|
||||
{"svcName", "pi-session-abc", svcName("abc")},
|
||||
{"svcDNS", "pi-session-abc.poimen.svc.cluster.local:8787", svcDNS("abc")},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if c.got != c.want {
|
||||
t.Errorf("%s: want %q, got %q", c.fn, c.want, c.got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── TestSandboxIsNotFound ─────────────────────────────────────────────────────
|
||||
|
||||
func TestSandboxIsNotFound(t *testing.T) {
|
||||
if !isNotFound(&k8sNotFound{}) {
|
||||
t.Error("isNotFound(&k8sNotFound{}) must be true")
|
||||
}
|
||||
if isNotFound(fmt.Errorf("other")) {
|
||||
t.Error("isNotFound(fmt.Errorf) must be false")
|
||||
}
|
||||
if isNotFound(nil) {
|
||||
t.Error("isNotFound(nil) must be false")
|
||||
}
|
||||
}
|
||||
@@ -36,6 +36,21 @@ func (rw *responseWriter) Write(b []byte) (int, error) {
|
||||
return n, err
|
||||
}
|
||||
|
||||
// Flush forwards to the wrapped writer so that streaming handlers can still
|
||||
// detect http.Flusher through this wrapper. Without it a bare
|
||||
// `w.(http.Flusher)` assertion fails for every request that passes through
|
||||
// this middleware, which silently disables SSE flushing and keep-alives.
|
||||
func (rw *responseWriter) Flush() {
|
||||
if f, ok := rw.ResponseWriter.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
// Unwrap exposes the wrapped writer to http.ResponseController.
|
||||
func (rw *responseWriter) Unwrap() http.ResponseWriter {
|
||||
return rw.ResponseWriter
|
||||
}
|
||||
|
||||
// Middleware returns an HTTP middleware that adds tracing to requests.
|
||||
func Middleware(next http.Handler) http.Handler {
|
||||
tracer := otel.Tracer(tracerName)
|
||||
|
||||
File diff suppressed because one or more lines are too long
+22
-3
@@ -21,9 +21,10 @@ spec:
|
||||
labels:
|
||||
app: api-gateway
|
||||
component: gateway
|
||||
# Required by the llm-serving-default-deny NetworkPolicy, which admits
|
||||
# only pods labelled llm-client=true (from any namespace) on port 8080.
|
||||
# Without it every upstream dial times out and dispatch returns 502.
|
||||
# The llm-serving-ingress NetworkPolicy in the llm-serving namespace now
|
||||
# selects by namespace (kubernetes.io/metadata.name: api), not by this
|
||||
# label. Kept for older NetworkPolicy revisions that still admit on
|
||||
# llm-client=true.
|
||||
llm-client: "true"
|
||||
annotations:
|
||||
reloader.stakater.com/auto: "true"
|
||||
@@ -67,6 +68,24 @@ spec:
|
||||
value: "5m"
|
||||
- name: LOG_LEVEL
|
||||
value: "info"
|
||||
- name: ADMISSION_MAX_PARK_SECONDS
|
||||
value: "120"
|
||||
# Must track the engine's --max-num-seqs (homelab: reasoning.yaml).
|
||||
# The gauge this compares against is engine-wide, so this number is a
|
||||
# ceiling on concurrent requests across all gateway replicas, not a
|
||||
# per-pod slot count. Drift here is silent; startup logs the value.
|
||||
- name: ADMISSION_MAX_INFLIGHT
|
||||
value: "4"
|
||||
# Only backends that publish the vLLM scheduler gauges may be observed.
|
||||
# `reasoning` is the only vLLM engine; ornith/qwen-cpu are Ollama and
|
||||
# embeddings/rerank are TEI, so observing them would hold
|
||||
# observer_degraded at 1 forever.
|
||||
- name: ADMISSION_OBSERVE_MODELS
|
||||
value: "reasoning"
|
||||
- name: VLLM_MAX_QUEUE_DEPTH
|
||||
value: "8"
|
||||
- name: MAX_BODY_BYTES
|
||||
value: "1048576"
|
||||
# OpenTelemetry tracing configuration
|
||||
- name: OTEL_EXPORTER_OTLP_ENDPOINT
|
||||
value: "otel-collector.tracing.svc.cluster.local:4317"
|
||||
|
||||
+11
-3
@@ -21,11 +21,17 @@ spec:
|
||||
ports:
|
||||
- protocol: TCP
|
||||
port: 8080
|
||||
# Allow from Prometheus scraping (if in monitoring namespace)
|
||||
# Allow from Prometheus scraping (monitoring namespace)
|
||||
# MUST select on kubernetes.io/metadata.name, not a plain `name:` label.
|
||||
# Only ingress-nginx carries a plain `name` label in this cluster; every
|
||||
# other namespace has just kubernetes.io/metadata.name, so `name: monitoring`
|
||||
# matched nothing and silently denied every Prometheus scrape of :8080
|
||||
# (surfaced as context-deadline-exceeded on all 3 gateway replicas while the
|
||||
# ServiceMonitor itself was discovered and healthy).
|
||||
- from:
|
||||
- namespaceSelector:
|
||||
matchLabels:
|
||||
name: monitoring
|
||||
kubernetes.io/metadata.name: monitoring
|
||||
ports:
|
||||
- protocol: TCP
|
||||
port: 8080
|
||||
@@ -122,7 +128,7 @@ spec:
|
||||
ports:
|
||||
- protocol: TCP
|
||||
port: 7233
|
||||
# Allow to Poimen (memory:8080 + graphql:8090 + redis:6379)
|
||||
# Allow to Poimen (memory:8080 + graphql:8090 + redis:6379 + fleet-controller:8788)
|
||||
- to:
|
||||
- namespaceSelector:
|
||||
matchLabels:
|
||||
@@ -134,6 +140,8 @@ spec:
|
||||
port: 8090
|
||||
- protocol: TCP
|
||||
port: 6379
|
||||
- protocol: TCP
|
||||
port: 8788
|
||||
# Allow to Gotify (push notifications)
|
||||
- to:
|
||||
- namespaceSelector:
|
||||
|
||||
@@ -26,3 +26,37 @@ roleRef:
|
||||
kind: ClusterRole
|
||||
name: image-promoter
|
||||
apiGroup: rbac.authorization.k8s.io
|
||||
---
|
||||
# Role + RoleBinding for pi-session sandbox manager (§4.15.4).
|
||||
# api-gateway SA needs get/create/delete for pods, pvcs, services,
|
||||
# and networkpolicies in the poimen namespace.
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: Role
|
||||
metadata:
|
||||
name: session-sandbox-manager
|
||||
namespace: poimen
|
||||
labels:
|
||||
app: api-gateway
|
||||
rules:
|
||||
- apiGroups: [""]
|
||||
resources: ["pods", "persistentvolumeclaims", "services"]
|
||||
verbs: ["get", "list", "create", "delete"]
|
||||
- apiGroups: ["networking.k8s.io"]
|
||||
resources: ["networkpolicies"]
|
||||
verbs: ["get", "create", "delete"]
|
||||
---
|
||||
apiVersion: rbac.authorization.k8s.io/v1
|
||||
kind: RoleBinding
|
||||
metadata:
|
||||
name: session-sandbox-manager-api-gateway
|
||||
namespace: poimen
|
||||
labels:
|
||||
app: api-gateway
|
||||
subjects:
|
||||
- kind: ServiceAccount
|
||||
name: api-gateway
|
||||
namespace: api
|
||||
roleRef:
|
||||
kind: Role
|
||||
name: session-sandbox-manager
|
||||
apiGroup: rbac.authorization.k8s.io
|
||||
|
||||
Reference in New Issue
Block a user