llmgw/internal/router/canary.go

package router

import (
	"hash/fnv"
	"net/http"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/canary"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)

// A canary: a share of a model's requests sent to a candidate, usually the
// same model at newer weights, and the two compared as they are served.

// variantHeader pins a request to one side, for testing either on purpose:
// "canary" or "stable". Either is the model the key asked for, so pinning
// grants nothing the key did not have.
const variantHeader = "x-llmgw-variant"

// chooseVariant says whether a request goes to the canary. A session — the
// x-session-id header or OpenAI's user field — always lands on the same
// side, so one conversation does not hop between weights; a request without
// one is placed by its own id.
func chooseVariant(req *parsedRequest, cn *catalog.Canary) bool {
	switch strings.ToLower(req.Header.Get(variantHeader)) {
	case "canary":
		return true
	case "stable":
		return false
	}
	if cn.Percent <= 0 {
		return false
	}
	if cn.Percent >= 100 {
		return true
	}
	key := req.SessionID
	if key == "" {
		key = req.RequestID
	}
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32()%100) < cn.Percent
}

type canaryStats struct {
	mu sync.Mutex
	by map[string]*canary.Stat // by the model the canary is of
}

func (c *canaryStats) observe(parent, cand string, onCanary bool, status int, streamed bool, ttft, took time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.by == nil {
		c.by = map[string]*canary.Stat{}
	}
	s := c.by[parent]
	if s == nil || s.Canary != cand {
		// A new candidate starts from nothing: the last one's numbers are
		// not this one's.
		s = &canary.Stat{Canary: cand, Since: time.Now()}
		c.by[parent] = s
	}
	if onCanary {
		s.Candidate.Add(status, streamed, ttft, took)
	} else {
		s.Stable.Add(status, streamed, ttft, took)
	}
}

// CanaryStats is how each model's canary and its stable side have served
// since the canary began, kept by this router.
func (r *Router) CanaryStats() map[string]canary.Stat {
	r.canaries.mu.Lock()
	defer r.canaries.mu.Unlock()
	out := map[string]canary.Stat{}
	for k, v := range r.canaries.by {
		out[k] = *v
	}
	return out
}

// ResetCanary forgets what a model's canary served, for a canary that has
// been taken out or promoted.
func (r *Router) ResetCanary(model string) {
	r.canaries.mu.Lock()
	defer r.canaries.mu.Unlock()
	delete(r.canaries.by, model)
}

// observeCanary counts a finished request against its side of a canary.
func (r *Router) observeCanary(req *parsedRequest, res dispatchResult, status int, took time.Duration) {
	if res.canaryOf == "" || status == 499 {
		return
	}
	r.canaries.observe(res.canaryOf, res.canary, res.onCanary, status, req.Stream, res.outcome.TTFT, took)
}

// isCanaryTarget reports whether a name is some model's canary: reached
// through that model, never by its own name.
func (r *Router) isCanaryTarget(name string) bool {
	canonical, _ := r.resolve(name)
	if p := r.canaryTargets.Load(); p != nil {
		_, ok := (*p)[canonical]
		return ok
	}
	return false
}

// hiddenModel is a model clients cannot name: a judge, or a canary.
func (r *Router) hiddenModel(name string) bool {
	return r.isGuardModel(name) || r.isCanaryTarget(name)
}

// setVariantHeader tells the client which side served it.
func setVariantHeader(h http.Header, onCanary bool) {
	if onCanary {
		h.Set(variantHeader, "canary")
	} else {
		h.Set(variantHeader, "stable")
	}
}