1468 lines
47 KiB
Go
1468 lines
47 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
toki "git.toki-labs.com/toki/proto-socket/go"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
edgeevents "iop/apps/edge/internal/events"
|
|
edgenode "iop/apps/edge/internal/node"
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
// waitForQueueLen polls until the group queue reaches wantLen or times out.
|
|
func waitForQueueLen(t *testing.T, m *modelQueueManager, groupKey string, wantLen int) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(200 * time.Millisecond)
|
|
for time.Now().Before(deadline) {
|
|
m.mu.Lock()
|
|
g, ok := m.groups[groupKey]
|
|
got := 0
|
|
if ok {
|
|
got = len(g.queue)
|
|
}
|
|
m.mu.Unlock()
|
|
if got >= wantLen {
|
|
return
|
|
}
|
|
time.Sleep(1 * time.Millisecond)
|
|
}
|
|
t.Fatalf("timeout: queue for %q did not reach length %d", groupKey, wantLen)
|
|
}
|
|
|
|
// TestModelQueueFIFOOrdering verifies that queued items are dispatched in FIFO
|
|
// order when a slot becomes available.
|
|
func TestModelQueueFIFOOrdering(t *testing.T) {
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-q1",
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
})
|
|
|
|
entry := &edgenode.NodeEntry{NodeID: "node-q1"}
|
|
cands := []candidateNode{{entry: entry, capacity: 1}}
|
|
defPolicy := groupPolicy{}
|
|
|
|
m := newModelQueueManager(store)
|
|
|
|
// Fill the only capacity slot.
|
|
first, err := m.admit(context.Background(), "g-fifo", "", "", cands, defPolicy)
|
|
if err != nil || first == nil {
|
|
t.Fatalf("initial admit: %v", err)
|
|
}
|
|
|
|
// Manually insert two items in known order so we can verify FIFO.
|
|
item1 := &queueItem{
|
|
candidates: cands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(2 * time.Second),
|
|
}
|
|
item2 := &queueItem{
|
|
candidates: cands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(2 * time.Second),
|
|
}
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-fifo", groupPolicy{})
|
|
g.queue = append(g.queue, item1, item2)
|
|
m.mu.Unlock()
|
|
|
|
// Release one slot — item1 (head) must be dispatched first.
|
|
m.mu.Lock()
|
|
m.releaseSlotLocked("g-fifo", "node-q1", "")
|
|
m.mu.Unlock()
|
|
|
|
select {
|
|
case res := <-item1.waitCh:
|
|
if res.err != nil {
|
|
t.Fatalf("item1 expected dispatch, got error: %v", res.err)
|
|
}
|
|
if res.candidate == nil || res.candidate.entry.NodeID != "node-q1" {
|
|
t.Fatalf("item1: unexpected candidate %v", res.candidate)
|
|
}
|
|
case <-time.After(100 * time.Millisecond):
|
|
t.Fatal("timeout: item1 was not dispatched after slot release")
|
|
}
|
|
|
|
// item2 must still be waiting (no slot available yet).
|
|
select {
|
|
case <-item2.waitCh:
|
|
t.Fatal("item2 should not have been dispatched yet")
|
|
default:
|
|
}
|
|
|
|
// Release the slot item1 holds so item2 gets dispatched.
|
|
m.mu.Lock()
|
|
m.releaseSlotLocked("g-fifo", "node-q1", "")
|
|
m.mu.Unlock()
|
|
|
|
select {
|
|
case res := <-item2.waitCh:
|
|
if res.err != nil {
|
|
t.Fatalf("item2 expected dispatch, got error: %v", res.err)
|
|
}
|
|
case <-time.After(100 * time.Millisecond):
|
|
t.Fatal("timeout: item2 was not dispatched after second slot release")
|
|
}
|
|
}
|
|
|
|
// TestModelQueueOverflow verifies that admit rejects requests when the queue
|
|
// is at max_queue capacity.
|
|
func TestModelQueueOverflow(t *testing.T) {
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-ov1",
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
})
|
|
|
|
entry := &edgenode.NodeEntry{NodeID: "node-ov1"}
|
|
cands := []candidateNode{{entry: entry, capacity: 1}}
|
|
|
|
m := newModelQueueManager(store)
|
|
|
|
// Fill the node's capacity and set max_queue=1.
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-overflow", groupPolicy{})
|
|
g.policy.maxQueue = 1
|
|
g.inflight["node-ov1"] = 1
|
|
m.mu.Unlock()
|
|
|
|
// Queue one item — should succeed.
|
|
item := &queueItem{
|
|
candidates: cands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(2 * time.Second),
|
|
}
|
|
m.mu.Lock()
|
|
m.groups["g-overflow"].queue = append(m.groups["g-overflow"].queue, item)
|
|
m.mu.Unlock()
|
|
|
|
// Second admit should fail immediately with errQueueFull.
|
|
_, err := m.admit(context.Background(), "g-overflow", "", "", cands, groupPolicy{})
|
|
if !errors.Is(err, errQueueFull) {
|
|
t.Fatalf("expected errQueueFull, got: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestModelQueueTimeout verifies that queued items receive a timeout error
|
|
// when no slot becomes available before the deadline.
|
|
func TestModelQueueTimeout(t *testing.T) {
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-to1",
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
})
|
|
|
|
entry := &edgenode.NodeEntry{NodeID: "node-to1"}
|
|
cands := []candidateNode{{entry: entry, capacity: 1}}
|
|
|
|
m := newModelQueueManager(store)
|
|
|
|
// Fill capacity and set a very short queue timeout.
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-timeout", groupPolicy{})
|
|
g.policy.queueTimeout = 20 * time.Millisecond
|
|
g.inflight["node-to1"] = 1
|
|
m.mu.Unlock()
|
|
|
|
start := time.Now()
|
|
_, err := m.admit(context.Background(), "g-timeout", "", "", cands, groupPolicy{})
|
|
elapsed := time.Since(start)
|
|
|
|
if !errors.Is(err, errQueueTimeout) {
|
|
t.Fatalf("expected errQueueTimeout, got: %v", err)
|
|
}
|
|
if elapsed < 15*time.Millisecond {
|
|
t.Fatalf("timed out too fast: %v", elapsed)
|
|
}
|
|
}
|
|
|
|
// TestModelQueueTerminalReleaseDispatchesNext verifies that a terminal run event
|
|
// releases the in-flight slot and dispatches the next queued item.
|
|
func TestModelQueueTerminalReleaseDispatchesNext(t *testing.T) {
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-tr1",
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
})
|
|
|
|
entry := &edgenode.NodeEntry{NodeID: "node-tr1"}
|
|
cands := []candidateNode{{entry: entry, capacity: 1}}
|
|
defPolicy := groupPolicy{}
|
|
|
|
bus := edgeevents.NewBus()
|
|
m := newModelQueueManager(store)
|
|
stop := m.startEventWatcher(bus)
|
|
defer stop()
|
|
|
|
// Fill capacity and record inflight.
|
|
selected, err := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy)
|
|
if err != nil {
|
|
t.Fatalf("admit: %v", err)
|
|
}
|
|
m.trackInflight("g-tr", "run-tr-001", selected.entry.NodeID, selected.providerID)
|
|
|
|
// Queue a second item in a goroutine.
|
|
resultCh := make(chan admitResult, 1)
|
|
go func() {
|
|
n, e := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy)
|
|
resultCh <- admitResult{candidate: n, err: e}
|
|
}()
|
|
|
|
waitForQueueLen(t, m, "g-tr", 1)
|
|
|
|
// Terminal event for the first run — should unblock the queued admit.
|
|
for _, termType := range []string{"complete", "error", "cancelled"} {
|
|
t.Run(termType, func(t *testing.T) {
|
|
// Reset state for each sub-test.
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-tr-"+termType, groupPolicy{})
|
|
g.inflight["node-tr1"] = 1
|
|
item := &queueItem{
|
|
candidates: cands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(2 * time.Second),
|
|
}
|
|
g.queue = []*queueItem{item}
|
|
runID := "run-tr-" + termType
|
|
m.inflightByRun[runID] = inflightRec{groupKey: "g-tr-" + termType, nodeID: "node-tr1"}
|
|
m.mu.Unlock()
|
|
|
|
bus.PublishRun(&iop.RunEvent{RunId: runID, Type: termType})
|
|
|
|
select {
|
|
case res := <-item.waitCh:
|
|
if res.err != nil {
|
|
t.Fatalf("expected dispatch, got error: %v", res.err)
|
|
}
|
|
if res.candidate == nil {
|
|
t.Fatal("expected non-nil candidate")
|
|
}
|
|
case <-time.After(500 * time.Millisecond):
|
|
t.Fatalf("timeout: queued item not dispatched after %q terminal event", termType)
|
|
}
|
|
})
|
|
}
|
|
|
|
// Cleanup the goroutine from the outer admit.
|
|
m.mu.Lock()
|
|
if g, ok := m.groups["g-tr"]; ok && len(g.queue) > 0 {
|
|
g.inflight["node-tr1"] = 1
|
|
m.tryDispatchLocked(g)
|
|
}
|
|
m.mu.Unlock()
|
|
select {
|
|
case <-resultCh:
|
|
case <-time.After(100 * time.Millisecond):
|
|
}
|
|
}
|
|
|
|
// TestModelQueueNodeDisconnectReleasesInflight verifies that a node disconnect
|
|
// event removes the node from queued item candidate lists, and that subsequent
|
|
// dispatch goes to a remaining live candidate (not the disconnected node).
|
|
func TestModelQueueNodeDisconnectReleasesInflight(t *testing.T) {
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-nd1",
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
})
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-nd2",
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
})
|
|
|
|
entry1 := &edgenode.NodeEntry{NodeID: "node-nd1"}
|
|
entry2 := &edgenode.NodeEntry{NodeID: "node-nd2"}
|
|
// Queued item that can use either node (both at capacity 1).
|
|
bothCands := []candidateNode{
|
|
{entry: entry1, capacity: 1},
|
|
{entry: entry2, capacity: 1},
|
|
}
|
|
|
|
bus := edgeevents.NewBus()
|
|
m := newModelQueueManager(store)
|
|
stop := m.startEventWatcher(bus)
|
|
defer stop()
|
|
|
|
// Fill node-nd1 and node-nd2 capacities and record inflight.
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-nd", groupPolicy{})
|
|
g.inflight["node-nd1"] = 1
|
|
g.inflight["node-nd2"] = 1
|
|
m.inflightByRun["run-nd-x"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd1"}
|
|
m.inflightByRun["run-nd-y"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd2"}
|
|
// Queue an item that can use either node.
|
|
item := &queueItem{
|
|
candidates: bothCands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(2 * time.Second),
|
|
}
|
|
g.queue = []*queueItem{item}
|
|
m.mu.Unlock()
|
|
|
|
// node-nd1 disconnects: its inflight is freed and it is removed from candidates.
|
|
// node-nd2 is still full, so no dispatch yet.
|
|
bus.PublishNode(&iop.EdgeNodeEvent{
|
|
NodeId: "node-nd1",
|
|
Type: "node.disconnected",
|
|
})
|
|
|
|
// node-nd2's run terminates: now nd2 has capacity, dispatch goes to nd2.
|
|
bus.PublishRun(&iop.RunEvent{RunId: "run-nd-y", Type: "complete"})
|
|
|
|
select {
|
|
case res := <-item.waitCh:
|
|
if res.err != nil {
|
|
t.Fatalf("expected dispatch after disconnect+terminal, got error: %v", res.err)
|
|
}
|
|
if res.candidate == nil {
|
|
t.Fatal("expected non-nil candidate after disconnect release")
|
|
}
|
|
// Must be nd2 — nd1 was removed from candidates on disconnect.
|
|
if res.candidate.entry.NodeID != "node-nd2" {
|
|
t.Errorf("expected dispatch to node-nd2, got %q (nd1 was disconnected)", res.candidate.entry.NodeID)
|
|
}
|
|
case <-time.After(500 * time.Millisecond):
|
|
t.Fatal("timeout: queued item not dispatched after node disconnect and terminal event")
|
|
}
|
|
|
|
// The inflight record for run-nd-x should be gone (removed at disconnect).
|
|
m.mu.Lock()
|
|
_, stillInFlight := m.inflightByRun["run-nd-x"]
|
|
m.mu.Unlock()
|
|
if stillInFlight {
|
|
t.Error("run-nd-x should be removed from inflightByRun after node disconnect")
|
|
}
|
|
}
|
|
|
|
// TestModelQueueUsesProviderCapacity verifies that the capacity supplied in the
|
|
// candidateNode (derived from adapter config) overrides the default of 1.
|
|
func TestModelQueueUsesProviderCapacity(t *testing.T) {
|
|
entry := &edgenode.NodeEntry{NodeID: "node-pc1"}
|
|
// Provider capacity = 2 (two concurrent slots on this node).
|
|
cands := []candidateNode{{entry: entry, capacity: 2}}
|
|
defPolicy := groupPolicy{}
|
|
|
|
m := newModelQueueManager(nil)
|
|
|
|
// First admit: inflight=0 < cap=2 → dispatched immediately.
|
|
n1, err := m.admit(context.Background(), "g-pc", "", "", cands, defPolicy)
|
|
if err != nil || n1 == nil {
|
|
t.Fatalf("first admit: %v", err)
|
|
}
|
|
|
|
// Second admit: inflight=1 < cap=2 → still dispatched (not queued).
|
|
n2, err := m.admit(context.Background(), "g-pc", "", "", cands, defPolicy)
|
|
if err != nil || n2 == nil {
|
|
t.Fatalf("second admit (capacity=2 should allow): %v", err)
|
|
}
|
|
|
|
// Third request must queue because inflight=2 == cap=2.
|
|
item := &queueItem{
|
|
candidates: cands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(200 * time.Millisecond),
|
|
}
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-pc", groupPolicy{})
|
|
g.queue = append(g.queue, item)
|
|
m.mu.Unlock()
|
|
|
|
// Releasing one slot should dispatch the queued item.
|
|
m.releaseSlot("g-pc", "node-pc1")
|
|
|
|
select {
|
|
case res := <-item.waitCh:
|
|
if res.err != nil {
|
|
t.Fatalf("expected dispatch after slot release, got: %v", res.err)
|
|
}
|
|
if res.candidate == nil || res.candidate.entry.NodeID != "node-pc1" {
|
|
t.Fatalf("unexpected candidate: %v", res.candidate)
|
|
}
|
|
case <-time.After(200 * time.Millisecond):
|
|
t.Fatal("timeout: item not dispatched after slot release")
|
|
}
|
|
}
|
|
|
|
// TestRefreshProviderQueuePolicyUpdatesExistingGroup verifies that a tightened
|
|
// queue policy (e.g. provider max_queue lowered via config refresh) propagates
|
|
// to an already-created group on the next admission, while in-flight slots are
|
|
// preserved.
|
|
func TestRefreshProviderQueuePolicyUpdatesExistingGroup(t *testing.T) {
|
|
entry := &edgenode.NodeEntry{NodeID: "node-rp"}
|
|
cands := []candidateNode{{entry: entry, capacity: 2}}
|
|
m := newModelQueueManager(nil)
|
|
|
|
// Create the group with a generous policy and take one slot.
|
|
initial := groupPolicy{maxQueue: 8, queueTimeout: 5 * time.Second}
|
|
if _, err := m.admit(context.Background(), "g-rp", "", "", cands, initial); err != nil {
|
|
t.Fatalf("first admit: %v", err)
|
|
}
|
|
m.mu.Lock()
|
|
if got := m.groups["g-rp"].policy.maxQueue; got != 8 {
|
|
m.mu.Unlock()
|
|
t.Fatalf("group created with maxQueue=%d, want 8", got)
|
|
}
|
|
m.mu.Unlock()
|
|
|
|
// Config refresh tightens the policy to maxQueue=1; the recomputed admit
|
|
// carries it and must update the existing group's policy. capacity=2 means
|
|
// this second admit is dispatched (not queued).
|
|
refreshed := groupPolicy{maxQueue: 1, queueTimeout: 5 * time.Second}
|
|
if _, err := m.admit(context.Background(), "g-rp", "", "", cands, refreshed); err != nil {
|
|
t.Fatalf("second admit: %v", err)
|
|
}
|
|
|
|
m.mu.Lock()
|
|
g := m.groups["g-rp"]
|
|
if g.policy.maxQueue != 1 {
|
|
m.mu.Unlock()
|
|
t.Fatalf("existing group policy not refreshed: maxQueue=%d, want 1", g.policy.maxQueue)
|
|
}
|
|
// Both slots now in-flight (cap=2). Queue one item to reach the refreshed
|
|
// maxQueue=1 so the next admission must overflow.
|
|
g.queue = append(g.queue, &queueItem{
|
|
candidates: cands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(5 * time.Second),
|
|
})
|
|
m.mu.Unlock()
|
|
|
|
if _, err := m.admit(context.Background(), "g-rp", "", "", cands, refreshed); !errors.Is(err, errQueueFull) {
|
|
t.Fatalf("expected errQueueFull under refreshed maxQueue=1, got: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestModelQueueUsesProviderQueuePolicy verifies that max_queue and
|
|
// queue_timeout from the policy parameter are respected.
|
|
func TestModelQueueUsesProviderQueuePolicy(t *testing.T) {
|
|
t.Run("maxQueue enforced", func(t *testing.T) {
|
|
entry := &edgenode.NodeEntry{NodeID: "node-pq1"}
|
|
cands := []candidateNode{{entry: entry, capacity: 1}}
|
|
// Provider policy: maxQueue=1.
|
|
policy := groupPolicy{maxQueue: 1, queueTimeout: 5 * time.Second}
|
|
|
|
m := newModelQueueManager(nil)
|
|
|
|
// Fill capacity and queue one item manually.
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-pq-max", policy)
|
|
g.inflight["node-pq1"] = 1
|
|
g.queue = []*queueItem{{
|
|
candidates: cands,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: time.Now().Add(5 * time.Second),
|
|
}}
|
|
m.mu.Unlock()
|
|
|
|
// Second admit should fail: queue already at maxQueue=1.
|
|
_, err := m.admit(context.Background(), "g-pq-max", "", "", cands, policy)
|
|
if !errors.Is(err, errQueueFull) {
|
|
t.Fatalf("expected errQueueFull, got: %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("queueTimeout enforced", func(t *testing.T) {
|
|
entry := &edgenode.NodeEntry{NodeID: "node-pq2"}
|
|
cands := []candidateNode{{entry: entry, capacity: 1}}
|
|
// Provider policy: very short timeout.
|
|
policy := groupPolicy{maxQueue: 16, queueTimeout: 20 * time.Millisecond}
|
|
|
|
m := newModelQueueManager(nil)
|
|
|
|
// Fill capacity.
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-pq-to", policy)
|
|
g.inflight["node-pq2"] = 1
|
|
m.mu.Unlock()
|
|
|
|
start := time.Now()
|
|
_, err := m.admit(context.Background(), "g-pq-to", "", "", cands, policy)
|
|
elapsed := time.Since(start)
|
|
|
|
if !errors.Is(err, errQueueTimeout) {
|
|
t.Fatalf("expected errQueueTimeout, got: %v", err)
|
|
}
|
|
if elapsed < 15*time.Millisecond {
|
|
t.Fatalf("timed out too fast: %v", elapsed)
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestModelQueueProviderLoadRatioSelection verifies that admit selects the
|
|
// candidate with the lowest in_flight/capacity ratio, not simply the first one.
|
|
func TestModelQueueProviderLoadRatioSelection(t *testing.T) {
|
|
entryA := &edgenode.NodeEntry{NodeID: "node-lr-a"}
|
|
entryB := &edgenode.NodeEntry{NodeID: "node-lr-b"}
|
|
// A: capacity=4, B: capacity=2.
|
|
cands := []candidateNode{
|
|
{entry: entryA, capacity: 4, providerID: "prov-a"},
|
|
{entry: entryB, capacity: 2, providerID: "prov-b"},
|
|
}
|
|
m := newModelQueueManager(nil)
|
|
|
|
// Put A at inflight=2 (ratio=0.5) and B at inflight=0 (ratio=0.0).
|
|
// Uses provider-aware slot keys per REVIEW_API-2.
|
|
m.mu.Lock()
|
|
g := m.getOrCreateGroupLocked("g-lr", groupPolicy{})
|
|
g.inflight["node-lr-a:prov-a"] = 2
|
|
m.mu.Unlock()
|
|
|
|
// Admit should pick B (ratio=0.0 < 0.5).
|
|
sel, err := m.admit(context.Background(), "g-lr", "", "", cands, groupPolicy{})
|
|
if err != nil {
|
|
t.Fatalf("admit: %v", err)
|
|
}
|
|
if sel == nil || sel.entry.NodeID != "node-lr-b" {
|
|
t.Fatalf("expected node-lr-b (lower ratio), got %v", sel)
|
|
}
|
|
|
|
// B is now at inflight=1 (ratio=0.5). A is still at inflight=2 (ratio=0.5).
|
|
// Equal-ratio rotation advances after the last selected slot (B), so A wins.
|
|
sel2, err := m.admit(context.Background(), "g-lr", "", "", cands, groupPolicy{})
|
|
if err != nil {
|
|
t.Fatalf("admit2: %v", err)
|
|
}
|
|
if sel2 == nil || sel2.entry.NodeID != "node-lr-a" {
|
|
t.Fatalf("expected node-lr-a (tie-break by providerID), got %v", sel2)
|
|
}
|
|
}
|
|
|
|
func TestModelQueueProviderLoadRatioFillsProportionally(t *testing.T) {
|
|
gx10 := &edgenode.NodeEntry{NodeID: "node-gx10"}
|
|
onex := &edgenode.NodeEntry{NodeID: "node-onex"}
|
|
cands := []candidateNode{
|
|
{entry: gx10, capacity: 4, providerID: "gx10-vllm"},
|
|
{entry: onex, capacity: 3, providerID: "onexplayer-lemonade"},
|
|
}
|
|
m := newModelQueueManager(nil)
|
|
|
|
var got []string
|
|
for i := 0; i < 7; i++ {
|
|
sel, err := m.admit(context.Background(), "g-proportional", "", "", cands, groupPolicy{})
|
|
if err != nil {
|
|
t.Fatalf("admit %d: %v", i+1, err)
|
|
}
|
|
got = append(got, sel.providerID)
|
|
}
|
|
|
|
want := []string{
|
|
"gx10-vllm",
|
|
"onexplayer-lemonade",
|
|
"gx10-vllm",
|
|
"onexplayer-lemonade",
|
|
"gx10-vllm",
|
|
"onexplayer-lemonade",
|
|
"gx10-vllm",
|
|
}
|
|
if len(got) != len(want) {
|
|
t.Fatalf("dispatch count: got %d want %d", len(got), len(want))
|
|
}
|
|
for i := range want {
|
|
if got[i] != want[i] {
|
|
t.Fatalf("dispatch[%d]: got %q want %q; full sequence=%v", i, got[i], want[i], got)
|
|
}
|
|
}
|
|
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
g := m.groups["g-proportional"]
|
|
if g.inflight["node-gx10:gx10-vllm"] != 4 || g.inflight["node-onex:onexplayer-lemonade"] != 3 {
|
|
t.Fatalf("inflight: got gx10=%d onex=%d, want 4/3", g.inflight["node-gx10:gx10-vllm"], g.inflight["node-onex:onexplayer-lemonade"])
|
|
}
|
|
}
|
|
|
|
func TestModelQueueProviderEqualIdleTieRotates(t *testing.T) {
|
|
gx10 := &edgenode.NodeEntry{NodeID: "node-gx10-idle"}
|
|
onex := &edgenode.NodeEntry{NodeID: "node-onex-idle"}
|
|
cands := []candidateNode{
|
|
{entry: gx10, capacity: 4, providerID: "gx10-vllm"},
|
|
{entry: onex, capacity: 3, providerID: "onexplayer-lemonade"},
|
|
}
|
|
m := newModelQueueManager(nil)
|
|
|
|
sel1, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
|
|
if err != nil {
|
|
t.Fatalf("first admit: %v", err)
|
|
}
|
|
if sel1.providerID != "gx10-vllm" {
|
|
t.Fatalf("first admit provider: got %q want gx10-vllm", sel1.providerID)
|
|
}
|
|
m.releaseSlot("g-idle-rotate", sel1.entry.NodeID, sel1.providerID)
|
|
|
|
sel2, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
|
|
if err != nil {
|
|
t.Fatalf("second admit: %v", err)
|
|
}
|
|
if sel2.providerID != "onexplayer-lemonade" {
|
|
t.Fatalf("second admit provider: got %q want onexplayer-lemonade", sel2.providerID)
|
|
}
|
|
m.releaseSlot("g-idle-rotate", sel2.entry.NodeID, sel2.providerID)
|
|
|
|
sel3, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
|
|
if err != nil {
|
|
t.Fatalf("third admit: %v", err)
|
|
}
|
|
if sel3.providerID != "gx10-vllm" {
|
|
t.Fatalf("third admit provider: got %q want gx10-vllm", sel3.providerID)
|
|
}
|
|
}
|
|
|
|
// TestModelQueueProviderServedTargetRewrite verifies that the selected
|
|
// candidateNode carries its servedTarget so callers can rewrite req.Target.
|
|
func TestModelQueueProviderServedTargetRewrite(t *testing.T) {
|
|
entry := &edgenode.NodeEntry{NodeID: "node-tr-rewrite"}
|
|
cands := []candidateNode{{
|
|
entry: entry,
|
|
capacity: 2,
|
|
providerID: "prov-vllm",
|
|
servedTarget: "qwen3-72b-instruct",
|
|
}}
|
|
m := newModelQueueManager(nil)
|
|
|
|
sel, err := m.admit(context.Background(), "g-tr-rewrite", "", "", cands, groupPolicy{})
|
|
if err != nil {
|
|
t.Fatalf("admit: %v", err)
|
|
}
|
|
if sel == nil {
|
|
t.Fatal("expected non-nil candidate")
|
|
}
|
|
if sel.servedTarget != "qwen3-72b-instruct" {
|
|
t.Errorf("servedTarget: got %q, want %q", sel.servedTarget, "qwen3-72b-instruct")
|
|
}
|
|
if sel.providerID != "prov-vllm" {
|
|
t.Errorf("providerID: got %q, want %q", sel.providerID, "prov-vllm")
|
|
}
|
|
}
|
|
|
|
// TestProviderStatusInflightTracking verifies that getStatsForProviderLocked
|
|
// correctly counts in-flight runs keyed by (nodeID, providerID).
|
|
func TestProviderStatusInflightTracking(t *testing.T) {
|
|
m := newModelQueueManager(nil)
|
|
|
|
m.mu.Lock()
|
|
m.inflightByRun["run-p1"] = inflightRec{groupKey: "g-alias", nodeID: "node-x", providerID: "prov-1"}
|
|
m.inflightByRun["run-p2"] = inflightRec{groupKey: "g-alias", nodeID: "node-x", providerID: "prov-1"}
|
|
m.inflightByRun["run-p3"] = inflightRec{groupKey: "g-alias2", nodeID: "node-x", providerID: "prov-2"}
|
|
m.mu.Unlock()
|
|
|
|
m.mu.Lock()
|
|
inf1, q1 := m.getStatsForProviderLocked("node-x", "prov-1")
|
|
inf2, q2 := m.getStatsForProviderLocked("node-x", "prov-2")
|
|
inf3, q3 := m.getStatsForProviderLocked("node-y", "prov-1")
|
|
m.mu.Unlock()
|
|
|
|
if inf1 != 2 || q1 != 0 {
|
|
t.Errorf("prov-1 on node-x: inflight=%d queued=%d, want 2/0", inf1, q1)
|
|
}
|
|
if inf2 != 1 || q2 != 0 {
|
|
t.Errorf("prov-2 on node-x: inflight=%d queued=%d, want 1/0", inf2, q2)
|
|
}
|
|
if inf3 != 0 || q3 != 0 {
|
|
t.Errorf("prov-1 on node-y: inflight=%d queued=%d, want 0/0", inf3, q3)
|
|
}
|
|
}
|
|
|
|
// TestModelQueueContextCancelRemovesQueuedItem verifies that cancelling the
|
|
// context of a queued admit removes the item and returns context.Canceled.
|
|
func TestModelQueueContextCancelRemovesQueuedItem(t *testing.T) {
|
|
entry := &edgenode.NodeEntry{NodeID: "node-cc1"}
|
|
cands := []candidateNode{{entry: entry, capacity: 1}}
|
|
defPolicy := groupPolicy{}
|
|
|
|
m := newModelQueueManager(nil)
|
|
|
|
// Fill the only slot.
|
|
n, err := m.admit(context.Background(), "g-cc", "", "", cands, defPolicy)
|
|
if err != nil || n == nil {
|
|
t.Fatalf("initial admit: %v", err)
|
|
}
|
|
|
|
// Queue a second admit with a cancellable context.
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
resultCh := make(chan error, 1)
|
|
go func() {
|
|
_, err := m.admit(ctx, "g-cc", "", "", cands, defPolicy)
|
|
resultCh <- err
|
|
}()
|
|
|
|
waitForQueueLen(t, m, "g-cc", 1)
|
|
cancel()
|
|
|
|
select {
|
|
case err := <-resultCh:
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Errorf("expected context.Canceled, got: %v", err)
|
|
}
|
|
case <-time.After(200 * time.Millisecond):
|
|
t.Fatal("timeout: context cancellation not handled")
|
|
}
|
|
|
|
// Queue must be empty after cancellation.
|
|
m.mu.Lock()
|
|
qLen := 0
|
|
if g, ok := m.groups["g-cc"]; ok {
|
|
qLen = len(g.queue)
|
|
}
|
|
m.mu.Unlock()
|
|
if qLen != 0 {
|
|
t.Errorf("queue should be empty after cancel, got %d items", qLen)
|
|
}
|
|
}
|
|
|
|
// TestModelQueueProviderCapacityIsPerProviderSlot verifies that same-node
|
|
// multiple provider candidates each have independent capacity/in-flight
|
|
// accounting per REVIEW_API-2.
|
|
func TestModelQueueProviderCapacityIsPerProviderSlot(t *testing.T) {
|
|
entry := &edgenode.NodeEntry{NodeID: "node-same"}
|
|
// Two provider candidates on the same node, each capacity=1.
|
|
cands := []candidateNode{
|
|
{entry: entry, capacity: 1, providerID: "prov-a", adapter: "vllm"},
|
|
{entry: entry, capacity: 1, providerID: "prov-b", adapter: "vllm"},
|
|
}
|
|
m := newModelQueueManager(nil)
|
|
|
|
// First admit: pick prov-a (deterministic tie-break by providerID).
|
|
sel1, err := m.admit(context.Background(), "g-same", "", "", cands, groupPolicy{})
|
|
if err != nil || sel1 == nil {
|
|
t.Fatalf("first admit: %v", err)
|
|
}
|
|
if sel1.providerID != "prov-a" {
|
|
t.Fatalf("expected prov-a, got %s", sel1.providerID)
|
|
}
|
|
|
|
// Second admit: should pick prov-b independently (prov-a is full, prov-b has capacity).
|
|
sel2, err := m.admit(context.Background(), "g-same", "", "", cands, groupPolicy{})
|
|
if err != nil || sel2 == nil {
|
|
t.Fatalf("second admit (prov-b should be available): %v", err)
|
|
}
|
|
if sel2.providerID != "prov-b" {
|
|
t.Fatalf("expected prov-b, got %s", sel2.providerID)
|
|
}
|
|
|
|
// Third admit: both providers at capacity, should queue.
|
|
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
|
defer cancel()
|
|
_, err = m.admit(ctx, "g-same", "", "", cands, groupPolicy{})
|
|
if !errors.Is(err, errQueueTimeout) && !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Errorf("expected queue timeout or deadline exceeded, got: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestResolveProviderPoolCandidatesFiltersInvalidProviders verifies that
|
|
// unavailable providers, served model mismatch, empty adapter providers,
|
|
// and capacity zero/unknown providers are excluded from dispatch candidates.
|
|
// This test calls the actual Service.resolveProviderPoolCandidates method.
|
|
func TestResolveProviderPoolCandidatesFiltersInvalidProviders(t *testing.T) {
|
|
// Build model catalog entry for "qwen3.6:35b".
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "qwen3.6:35b",
|
|
Providers: map[string]string{
|
|
"prov-available": "served-qwen",
|
|
"prov-unavailable": "served-qwen",
|
|
"prov-mismatch": "served-qwen",
|
|
"prov-no-adapter": "served-qwen",
|
|
"prov-cap-zero": "served-qwen",
|
|
"prov-cap-unknown": "served-qwen",
|
|
},
|
|
},
|
|
}
|
|
|
|
// Build NodeStore with multiple providers.
|
|
store := edgenode.NewNodeStore()
|
|
|
|
// Valid provider: available, has adapter, has capacity, served model matches provider's models.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-valid",
|
|
Runtime: config.RuntimeConf{Concurrency: 4},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-available",
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-qwen", "served-llama"},
|
|
Health: "available",
|
|
Capacity: 2,
|
|
},
|
|
},
|
|
})
|
|
|
|
// Invalid: health = "unavailable".
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-bad-health",
|
|
Runtime: config.RuntimeConf{Concurrency: 4},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-unavailable",
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-qwen"},
|
|
Health: "unavailable",
|
|
Capacity: 2,
|
|
},
|
|
},
|
|
})
|
|
|
|
// Invalid: served model not in provider's own models list.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-mismatch",
|
|
Runtime: config.RuntimeConf{Concurrency: 4},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-mismatch",
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-llama"}, // does NOT include "served-qwen"
|
|
Health: "available",
|
|
Capacity: 2,
|
|
},
|
|
},
|
|
})
|
|
|
|
// Provider-first (empty adapter, health available, capacity > 0): now a valid candidate.
|
|
// adapter key = provider ID "prov-no-adapter".
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-no-adapter",
|
|
Runtime: config.RuntimeConf{Concurrency: 4},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-no-adapter",
|
|
Type: "vllm",
|
|
Adapter: "", // provider-first: adapter key derived from ID
|
|
Models: []string{"served-qwen"},
|
|
Health: "available",
|
|
Capacity: 2,
|
|
},
|
|
},
|
|
})
|
|
|
|
// Invalid: capacity = 0.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-cap-zero",
|
|
Runtime: config.RuntimeConf{Concurrency: 4},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-cap-zero",
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-qwen"},
|
|
Health: "available",
|
|
Capacity: 0,
|
|
},
|
|
},
|
|
})
|
|
|
|
// Invalid: capacity < 0 (negative/unknown).
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-cap-unknown",
|
|
Runtime: config.RuntimeConf{Concurrency: 4},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-cap-unknown",
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-qwen"},
|
|
Health: "available",
|
|
Capacity: -1,
|
|
},
|
|
},
|
|
})
|
|
|
|
// Build a fake registry with all nodes.
|
|
reg := edgenode.NewRegistry()
|
|
allRecs := store.All()
|
|
for _, rec := range allRecs {
|
|
entry := &edgenode.NodeEntry{
|
|
NodeID: rec.ID,
|
|
LifecycleState: edgenode.LifecycleConnected,
|
|
}
|
|
reg.Register(entry)
|
|
}
|
|
|
|
// Create Service with catalog and node store.
|
|
svc := New(reg, nil)
|
|
svc.SetNodeStore(store)
|
|
svc.SetModelCatalog(catalog)
|
|
|
|
// Call the actual resolveProviderPoolCandidates.
|
|
req := SubmitRunRequest{
|
|
ModelGroupKey: "qwen3.6:35b",
|
|
ProviderPool: true,
|
|
}
|
|
storeSnapshot, catalogSnapshot := svc.runtimeConfigSnapshot()
|
|
candidates, policy, err := svc.resolveProviderPoolCandidates(req, storeSnapshot, catalogSnapshot)
|
|
if err != nil {
|
|
t.Fatalf("resolveProviderPoolCandidates: %v", err)
|
|
}
|
|
|
|
// prov-available (legacy/compat) and prov-no-adapter (provider-first) both pass.
|
|
// All others are filtered: unavailable health, model mismatch, cap=0, cap<0.
|
|
if len(candidates) != 2 {
|
|
ids := make([]string, len(candidates))
|
|
for i, c := range candidates {
|
|
ids[i] = c.providerID
|
|
}
|
|
t.Fatalf("expected 2 candidates, got %d: %v", len(candidates), ids)
|
|
}
|
|
|
|
byID := map[string]*candidateNode{}
|
|
for i := range candidates {
|
|
byID[candidates[i].providerID] = &candidates[i]
|
|
}
|
|
|
|
// Legacy/compat candidate uses explicit adapter key.
|
|
avail := byID["prov-available"]
|
|
if avail == nil {
|
|
t.Fatal("prov-available not in candidates")
|
|
}
|
|
if avail.adapter != "vllm-gpu" {
|
|
t.Errorf("prov-available adapter: got %q, want %q", avail.adapter, "vllm-gpu")
|
|
}
|
|
if avail.servedTarget != "served-qwen" {
|
|
t.Errorf("prov-available servedTarget: got %q", avail.servedTarget)
|
|
}
|
|
|
|
// Provider-first candidate uses provider ID as adapter key.
|
|
noAdp := byID["prov-no-adapter"]
|
|
if noAdp == nil {
|
|
t.Fatal("prov-no-adapter not in candidates (provider-first dispatch must use provider ID as adapter key)")
|
|
}
|
|
if noAdp.adapter != "prov-no-adapter" {
|
|
t.Errorf("prov-no-adapter adapter: got %q, want provider ID %q", noAdp.adapter, "prov-no-adapter")
|
|
}
|
|
if noAdp.servedTarget != "served-qwen" {
|
|
t.Errorf("prov-no-adapter servedTarget: got %q", noAdp.servedTarget)
|
|
}
|
|
|
|
// Policy should use defaults since none of the providers set policy.
|
|
if policy.maxQueue <= 0 {
|
|
t.Errorf("policy.maxQueue: got %d, expected > 0", policy.maxQueue)
|
|
}
|
|
if policy.queueTimeout <= 0 {
|
|
t.Errorf("policy.queueTimeout: got %v, expected > 0", policy.queueTimeout)
|
|
}
|
|
}
|
|
|
|
// TestSubmitRunProviderPoolRewritesAdapterAndTarget verifies that provider-pool
|
|
// SubmitRun rewrites both req.Adapter and req.Target from the selected candidate.
|
|
// This test uses a fake TCP client via net.Pipe to capture the actual RunRequest
|
|
// sent by SubmitRun(ProviderPool=true), confirming that the winning candidate's
|
|
// adapter and servedTarget propagate correctly through the full service path.
|
|
func TestSubmitRunProviderPoolRewritesAdapterAndTarget(t *testing.T) {
|
|
// Use net.Pipe to create a fake node connection that captures the RunRequest.
|
|
edgeConn, nodeConn := net.Pipe()
|
|
defer edgeConn.Close()
|
|
defer nodeConn.Close()
|
|
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.RunRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
|
|
// Capture the RunRequest received by the fake node.
|
|
var capturedReq *iop.RunRequest
|
|
var capturedMu sync.Mutex
|
|
toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
|
|
capturedMu.Lock()
|
|
capturedReq = req
|
|
capturedMu.Unlock()
|
|
})
|
|
|
|
// Build the model catalog with provider references.
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "qwen3.6:35b",
|
|
Providers: map[string]string{
|
|
"prov-vllm-01": "served-qwen",
|
|
},
|
|
},
|
|
}
|
|
|
|
// Build NodeStore with a provider-pool provider.
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-pool",
|
|
Runtime: config.RuntimeConf{Concurrency: 4},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-vllm-01",
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-qwen"},
|
|
Health: "available",
|
|
Capacity: 2,
|
|
},
|
|
},
|
|
})
|
|
|
|
// Build registry with the fake node.
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-pool",
|
|
LifecycleState: edgenode.LifecycleConnected,
|
|
Client: edgeClient,
|
|
})
|
|
|
|
// Create Service with queue and catalog.
|
|
// events bus must be non-nil to activate the queue path for provider-pool.
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
svc.SetModelCatalog(catalog)
|
|
|
|
// SubmitRun with ProviderPool=true.
|
|
result, err := svc.SubmitRun(context.Background(), SubmitRunRequest{
|
|
RunID: "run-pool-test-001",
|
|
ModelGroupKey: "qwen3.6:35b",
|
|
ProviderPool: true,
|
|
Background: true,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("SubmitRun: %v", err)
|
|
}
|
|
if result == nil {
|
|
t.Fatal("expected non-nil RunResult")
|
|
}
|
|
|
|
// Wait for the fake node to receive the request.
|
|
time.Sleep(50 * time.Millisecond)
|
|
|
|
capturedMu.Lock()
|
|
defer capturedMu.Unlock()
|
|
|
|
if capturedReq == nil {
|
|
t.Fatal("no RunRequest captured from fake node; SubmitRun did not send")
|
|
}
|
|
|
|
// Verify that the adapter and target were rewritten from the provider-pool candidate.
|
|
if capturedReq.GetAdapter() != "vllm-gpu" {
|
|
t.Errorf("adapter: got %q, want %q", capturedReq.GetAdapter(), "vllm-gpu")
|
|
}
|
|
if capturedReq.GetTarget() != "served-qwen" {
|
|
t.Errorf("target: got %q, want %q", capturedReq.GetTarget(), "served-qwen")
|
|
}
|
|
if capturedReq.GetRunId() != "run-pool-test-001" {
|
|
t.Errorf("runID: got %q, want %q", capturedReq.GetRunId(), "run-pool-test-001")
|
|
}
|
|
}
|
|
|
|
// TestResolveProviderPoolCandidatesAdapterInstanceValidation verifies that the
|
|
// defensive isProviderAdapterInstanceValid check in resolveProviderPoolCandidates
|
|
// excludes providers whose adapter name resolves to a disabled exact instance,
|
|
// an ambiguous type route, a zero-instance type route, or an unknown adapter key,
|
|
// while valid named instances and single enabled type routes remain as candidates.
|
|
func TestResolveProviderPoolCandidatesAdapterInstanceValidation(t *testing.T) {
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "ai-group",
|
|
Providers: map[string]string{
|
|
"prov-disabled": "ai-model",
|
|
"prov-ambiguous": "ai-model",
|
|
"prov-valid": "ai-model",
|
|
"prov-zero-type": "ai-model",
|
|
"prov-missing": "ai-model",
|
|
},
|
|
},
|
|
}
|
|
|
|
store := edgenode.NewNodeStore()
|
|
|
|
// prov-disabled: adapter="vllm-gpu" but VllmInstance "vllm-gpu" is disabled.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-disabled",
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: false, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-disabled", Adapter: "vllm-gpu", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
|
|
},
|
|
})
|
|
|
|
// prov-ambiguous: adapter="openai_compat" type route with 2 enabled instances → ambiguous.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-ambiguous",
|
|
Adapters: config.AdaptersConf{
|
|
OpenAICompatInstances: []config.OpenAICompatInstanceConf{
|
|
{Name: "compat-a", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
{Name: "compat-b", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-ambiguous", Adapter: "openai_compat", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
|
|
},
|
|
})
|
|
|
|
// prov-valid: adapter="vllm-gpu" and VllmInstance "vllm-gpu" is enabled.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-valid",
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-valid", Adapter: "vllm-gpu", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
|
|
},
|
|
})
|
|
|
|
// prov-zero-type: adapter="ollama" type route but no enabled ollama instances → zero-instance type route excluded.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-zero-type",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "ollama-disabled", Enabled: false, BaseURL: "http://127.0.0.1:11434"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-zero-type", Adapter: "ollama", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
|
|
},
|
|
})
|
|
|
|
// prov-missing: adapter="unknown-missing" not in any instance list and not a known type key → unknown adapter excluded.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-missing",
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:9000/v1"},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-missing", Adapter: "unknown-missing", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
|
|
},
|
|
})
|
|
|
|
reg := edgenode.NewRegistry()
|
|
allRecs := store.All()
|
|
for _, rec := range allRecs {
|
|
reg.Register(&edgenode.NodeEntry{NodeID: rec.ID, LifecycleState: edgenode.LifecycleConnected})
|
|
}
|
|
|
|
svc := New(reg, nil)
|
|
svc.SetNodeStore(store)
|
|
svc.SetModelCatalog(catalog)
|
|
|
|
req := SubmitRunRequest{ModelGroupKey: "ai-group", ProviderPool: true}
|
|
storeSnap, catalogSnap := svc.runtimeConfigSnapshot()
|
|
candidates, _, err := svc.resolveProviderPoolCandidates(req, storeSnap, catalogSnap)
|
|
if err != nil {
|
|
t.Fatalf("resolveProviderPoolCandidates: %v", err)
|
|
}
|
|
|
|
// Only prov-valid should be a candidate.
|
|
got := make(map[string]bool)
|
|
for _, c := range candidates {
|
|
got[c.providerID] = true
|
|
}
|
|
if len(candidates) != 1 {
|
|
t.Fatalf("expected 1 candidate (prov-valid), got %d: %v", len(candidates), got)
|
|
}
|
|
if !got["prov-valid"] {
|
|
t.Error("prov-valid (enabled exact instance) must be a candidate")
|
|
}
|
|
if got["prov-disabled"] {
|
|
t.Error("prov-disabled (disabled exact instance) must be excluded")
|
|
}
|
|
if got["prov-ambiguous"] {
|
|
t.Error("prov-ambiguous (ambiguous type route) must be excluded")
|
|
}
|
|
if got["prov-zero-type"] {
|
|
t.Error("prov-zero-type (zero-instance type route) must be excluded")
|
|
}
|
|
if got["prov-missing"] {
|
|
t.Error("prov-missing (unknown adapter) must be excluded")
|
|
}
|
|
}
|
|
|
|
// TestGetSnapshotForNodeCatalogFirstNoDuplicates verifies that when a node has
|
|
// both adapter instances and a providers[] catalog, getSnapshotForNode returns
|
|
// only catalog snapshots (no adapter duplicates).
|
|
func TestGetSnapshotForNodeCatalogFirstNoDuplicates(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-catalog",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{Enabled: true},
|
|
OpenAICompatInstances: []config.OpenAICompatInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-a", Adapter: "vllm-gpu", Models: []string{"model-a"}, Health: "available", Capacity: 2},
|
|
{ID: "prov-b", Adapter: "vllm-gpu", Models: []string{"model-b"}, Health: "available", Capacity: 3},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(rec)
|
|
|
|
m := newModelQueueManager(store)
|
|
|
|
snaps := m.getSnapshotForNode("node-catalog", rec)
|
|
|
|
// Must return exactly the catalog providers — no CLI or openai_compat adapter duplicates.
|
|
if len(snaps) != 2 {
|
|
t.Fatalf("expected 2 catalog snapshots, got %d: %+v", len(snaps), snaps)
|
|
}
|
|
ids := map[string]bool{}
|
|
for _, s := range snaps {
|
|
ids[s.Id] = true
|
|
}
|
|
if !ids["prov-a"] || !ids["prov-b"] {
|
|
t.Errorf("expected catalog provider IDs prov-a and prov-b, got: %v", ids)
|
|
}
|
|
// No adapter-only snapshots (no empty Id entries that correspond to adapter sections).
|
|
for _, s := range snaps {
|
|
if s.Id == "" {
|
|
t.Errorf("unexpected adapter-only snapshot (empty Id) in catalog-first result: %+v", s)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestResolveProviderPoolCandidatesSkipsDisabledProvider verifies that a
|
|
// provider with enabled=false is excluded from dispatch candidates.
|
|
func TestResolveProviderPoolCandidatesSkipsDisabledProvider(t *testing.T) {
|
|
disabled := false
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "model-x",
|
|
Providers: map[string]string{
|
|
"prov-active": "served-x",
|
|
"prov-disabled": "served-x",
|
|
},
|
|
},
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-d",
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-active", Type: "vllm", Adapter: "vllm-gpu", Models: []string{"served-x"}, Health: "available", Capacity: 2},
|
|
{ID: "prov-disabled", Type: "vllm", Adapter: "vllm-gpu", Models: []string{"served-x"}, Health: "available", Capacity: 2, Enabled: &disabled},
|
|
},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"}},
|
|
},
|
|
})
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-d", LifecycleState: edgenode.LifecycleConnected})
|
|
|
|
svc := New(reg, nil)
|
|
svc.SetNodeStore(store)
|
|
svc.SetModelCatalog(catalog)
|
|
|
|
storeSnap, catalogSnap := svc.runtimeConfigSnapshot()
|
|
candidates, _, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{ModelGroupKey: "model-x", ProviderPool: true}, storeSnap, catalogSnap)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if len(candidates) != 1 {
|
|
t.Fatalf("expected 1 candidate (disabled filtered), got %d", len(candidates))
|
|
}
|
|
if candidates[0].providerID != "prov-active" {
|
|
t.Errorf("expected prov-active, got %q", candidates[0].providerID)
|
|
}
|
|
}
|
|
|
|
// TestResolveProviderPoolCandidatesUsesProviderIDWhenAdapterOmitted verifies
|
|
// that a provider-first provider (empty Adapter field) uses its ID as the
|
|
// dispatch adapter key.
|
|
func TestResolveProviderPoolCandidatesUsesProviderIDWhenAdapterOmitted(t *testing.T) {
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "model-y",
|
|
Providers: map[string]string{
|
|
"prov-first": "served-y",
|
|
},
|
|
},
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-pf",
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-first", Type: "vllm", Adapter: "", Models: []string{"served-y"}, Health: "available", Capacity: 3},
|
|
},
|
|
})
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-pf", LifecycleState: edgenode.LifecycleConnected})
|
|
|
|
svc := New(reg, nil)
|
|
svc.SetNodeStore(store)
|
|
svc.SetModelCatalog(catalog)
|
|
|
|
storeSnap, catalogSnap := svc.runtimeConfigSnapshot()
|
|
candidates, _, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{ModelGroupKey: "model-y", ProviderPool: true}, storeSnap, catalogSnap)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if len(candidates) != 1 {
|
|
t.Fatalf("expected 1 candidate, got %d", len(candidates))
|
|
}
|
|
if candidates[0].adapter != "prov-first" {
|
|
t.Errorf("provider-first adapter key: got %q, want provider ID %q", candidates[0].adapter, "prov-first")
|
|
}
|
|
if candidates[0].providerID != "prov-first" {
|
|
t.Errorf("providerID: got %q", candidates[0].providerID)
|
|
}
|
|
}
|
|
|
|
// TestProviderSnapshotsReportDisabledProvider verifies that a disabled
|
|
// provider appears in snapshots with status=disabled and effective capacity 0.
|
|
func TestProviderSnapshotsReportDisabledProvider(t *testing.T) {
|
|
disabled := false
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-snap",
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-on", Type: "vllm", Health: "available", Capacity: 4},
|
|
{ID: "prov-off", Type: "vllm", Health: "available", Capacity: 4, Enabled: &disabled},
|
|
},
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(rec)
|
|
m := newModelQueueManager(store)
|
|
|
|
snaps := m.getSnapshotForNode("node-snap", rec)
|
|
if len(snaps) != 2 {
|
|
t.Fatalf("expected 2 snapshots, got %d", len(snaps))
|
|
}
|
|
|
|
byID := map[string]*iop.ProviderSnapshot{}
|
|
for _, s := range snaps {
|
|
byID[s.Id] = s
|
|
}
|
|
|
|
on := byID["prov-on"]
|
|
if on == nil {
|
|
t.Fatal("prov-on snapshot missing")
|
|
}
|
|
if on.Status != "available" {
|
|
t.Errorf("prov-on status: got %q, want available", on.Status)
|
|
}
|
|
if on.Capacity != 4 {
|
|
t.Errorf("prov-on capacity: got %d, want 4", on.Capacity)
|
|
}
|
|
|
|
off := byID["prov-off"]
|
|
if off == nil {
|
|
t.Fatal("prov-off snapshot missing")
|
|
}
|
|
if off.Status != "disabled" {
|
|
t.Errorf("prov-off status: got %q, want disabled", off.Status)
|
|
}
|
|
if off.Health != "disabled" {
|
|
t.Errorf("prov-off health: got %q, want disabled", off.Health)
|
|
}
|
|
if off.Capacity != 0 {
|
|
t.Errorf("prov-off effective capacity: got %d, want 0", off.Capacity)
|
|
}
|
|
}
|
|
|
|
// TestStatusProviderProviderFirstNoAdapterDuplicates verifies that provider-first
|
|
// providers (empty Adapter) appear once in snapshots with no adapter duplicate.
|
|
func TestStatusProviderProviderFirstNoAdapterDuplicates(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-pf-snap",
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-first-a", Type: "vllm", Adapter: "", Models: []string{"model-a"}, Health: "available", Capacity: 2},
|
|
{ID: "prov-first-b", Type: "vllm", Adapter: "", Models: []string{"model-b"}, Health: "available", Capacity: 2},
|
|
},
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(rec)
|
|
m := newModelQueueManager(store)
|
|
|
|
snaps := m.getSnapshotForNode("node-pf-snap", rec)
|
|
// catalog-first: exactly 2 provider snapshots, no adapter duplicate.
|
|
if len(snaps) != 2 {
|
|
t.Fatalf("expected 2 provider-first snapshots, got %d: %+v", len(snaps), snaps)
|
|
}
|
|
for _, s := range snaps {
|
|
if s.Id == "" {
|
|
t.Errorf("unexpected adapter-only snapshot (empty Id): %+v", s)
|
|
}
|
|
if s.Adapter != "" {
|
|
t.Errorf("provider-first snapshot should have empty Adapter, got %q for id=%q", s.Adapter, s.Id)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestGetSnapshotForNodeLegacyAdapterFallback verifies that a node with no
|
|
// providers[] catalog returns adapter snapshots (CLI, OllamaInstances, etc.)
|
|
// via the legacy adapter fallback path.
|
|
func TestGetSnapshotForNodeLegacyAdapterFallback(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-legacy",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{Enabled: true},
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "ollama-local", Enabled: true, Capacity: 3},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 2},
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(rec)
|
|
|
|
m := newModelQueueManager(store)
|
|
|
|
snaps := m.getSnapshotForNode("node-legacy", rec)
|
|
|
|
// Must return adapter snapshots: CLI + ollama-local = 2 entries.
|
|
if len(snaps) != 2 {
|
|
t.Fatalf("expected 2 adapter snapshots (cli + ollama-local), got %d: %+v", len(snaps), snaps)
|
|
}
|
|
adapters := map[string]bool{}
|
|
for _, s := range snaps {
|
|
adapters[s.Adapter] = true
|
|
}
|
|
if !adapters["cli"] {
|
|
t.Error("expected cli adapter snapshot in legacy fallback")
|
|
}
|
|
if !adapters["ollama-local"] {
|
|
t.Error("expected ollama-local adapter snapshot in legacy fallback")
|
|
}
|
|
}
|