llmgw/internal/control/routing_test.go

package control

import (
	"log/slog"
	"testing"
	"time"

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

type fixedFeed struct {
	reps []RoutedReplica
	ok   bool
}

func (f fixedFeed) Ready() ([]RoutedReplica, bool) { return f.reps, f.ok }

func routingController(t *testing.T, sink *recordSink) *Controller {
	t.Helper()
	cat, err := catalog.New([]catalog.Model{{Name: "llama", Source: "/models/llama", Aliases: []string{"ll"},
		Scaling:  catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "one", GPUClass: "h100-80g", Engine: "vllm", Priority: 3}}}})
	if err != nil {
		t.Fatal(err)
	}
	cfg := DefaultConfig()
	cfg.Reconciler.EngineKey = EngineKeys([]byte("secret"))
	return New(cfg, catalog.NewLive(cat), noDispatch{}, sink, slog.New(slog.DiscardHandler))
}

// standby is a leadership that does not lead: one without a store always
// would.
func standby() *Leadership {
	l := NewLeadership(nil, 0, nil)
	l.leader.Store(false)
	return l
}

// A standby following the leader routes to the leader's replicas, each with
// the key it derives itself and the engine and priority its own catalog
// gives the profile; a replica of a model its catalog lacks is left out.
func TestAStandbyRoutesToTheLeadersReplicas(t *testing.T) {
	sink := &recordSink{}
	c := routingController(t, sink)
	c.SetLeadership(standby())
	c.SetReplicaFeed(fixedFeed{ok: true, reps: []RoutedReplica{
		{ID: "llama-r0", Model: "llama", Profile: "one", Pool: "a", Endpoint: "http://10.0.0.5:20001/", Weight: 2},
		{ID: "gone-r0", Model: "gone", Profile: "x", Endpoint: "http://10.0.0.6:20001"},
	}})
	c.Publish(t.Context())

	if len(sink.infos) != 1 {
		t.Fatalf("published %+v", sink.infos)
	}
	e := sink.infos[0]
	if e.ID != "llama-r0" || e.URL != "http://10.0.0.5:20001" || e.Pool != "a" || e.Weight != 2 ||
		e.Engine != "vllm" || e.Priority != 3 || len(e.Models) != 2 || e.Models[1] != "ll" {
		t.Errorf("endpoint %+v", e)
	}
	if e.APIKey == "" || e.APIKey != EngineKeys([]byte("secret"))("llama-r0") {
		t.Errorf("the key is not the one the leader gave the engine: %q", e.APIKey)
	}
}

// One that has not heard from the leader lately routes to nothing, and the
// leader never reads a feed: its replicas are its reconciler's.
func TestAStandbyCutOffRoutesNothingAndTheLeaderIgnoresTheFeed(t *testing.T) {
	sink := &recordSink{}
	c := routingController(t, sink)
	l := standby()
	c.SetLeadership(l)
	c.SetReplicaFeed(fixedFeed{ok: false, reps: []RoutedReplica{{ID: "llama-r0", Model: "llama", Profile: "one"}}})
	if got := c.endpoints(); len(got) != 0 {
		t.Errorf("a standby that has not heard from the leader routes to %+v", got)
	}

	l.leader.Store(true)
	c.SetReplicaFeed(fixedFeed{ok: true, reps: []RoutedReplica{{ID: "llama-r0", Model: "llama", Profile: "one"}}})
	if got := c.endpoints(); len(got) != 0 {
		t.Errorf("the leader routed to a feed rather than its own replicas: %+v", got)
	}
}

// The catalog's revision moves with every catalog put in force, so a
// standby knows to read it again.
func TestTheCatalogRevisionMovesWithTheCatalog(t *testing.T) {
	c := routingController(t, &recordSink{})
	before := c.Routing().Catalog
	c.ReplaceCatalog(t.Context(), c.cat.Load())
	if after := c.Routing().Catalog; after == before {
		t.Errorf("the revision stayed %q", after)
	}
}

// rateSink is a router whose counts can be shared: its own, and what it was
// last told the others counted.
type rateSink struct {
	recordSink
	own    accounting.Rates
	others *accounting.Rates
}

func (s *rateSink) Rates() accounting.Rates { return s.own }
func (s *rateSink) SetOtherRates(r accounting.Rates) {
	s.others = &r
}

func counts(k string, n int64) accounting.Rates {
	r := accounting.NewRates(time.Now())
	r.Requests.Current = map[string]int64{k: n}
	return r
}

// Each standby is given what the leader and the other standbys counted, never
// its own; the leader puts the standbys' counts in force on its own router.
func TestLimitsAreCountedAcrossTheControlPlanes(t *testing.T) {
	sink := &rateSink{own: counts("key-1", 5)}
	cat, _ := catalog.New(nil)
	c := New(DefaultConfig(), catalog.NewLive(cat), noDispatch{}, sink, slog.New(slog.DiscardHandler))

	b, d := counts("key-1", 3), counts("key-1", 7)
	c.RecordActivity(Activity{From: "ctl-b", Rates: &b})
	c.RecordActivity(Activity{From: "ctl-d", Rates: &d})

	if got := c.RoutingFor("ctl-b").Rates.Requests.Current["key-1"]; got != 12 {
		t.Errorf("ctl-b was told %d, want the leader's 5 and ctl-d's 7", got)
	}
	c.shareRates()
	if sink.others == nil || sink.others.Requests.Current["key-1"] != 10 {
		t.Errorf("the leader's router was told %+v, want the standbys' 10", sink.others)
	}

	// A standby that stopped reporting stops counting.
	c.rates.mu.Lock()
	r := c.rates.by["ctl-d"]
	r.at = r.at.Add(-ratesFresh - time.Second)
	c.rates.by["ctl-d"] = r
	c.rates.mu.Unlock()
	if got := c.RoutingFor("ctl-b").Rates.Requests.Current["key-1"]; got != 5 {
		t.Errorf("a quiet standby still counts: %d", got)
	}
}