toki
4dcf6f0cf9
- 노드 연결 2단계 핸드셰이크 추가 (Register→NodeReady→DispatchReady) - ConnectionGeneration 기반 연결 세대 관리로 stale 연결 차폐 - configured 노드 catalog 기반 snapshot rebuild (Connected 상태 분리) - provider 리스 소유권 일원화: edge가 소유권 승인·반환 전까지 대기 - 모델 대기열 승인/해제/스냅샷 서비스 구현 - 재연결 준비도 통합 테스트, 아카이브된 하위태스크 8건 포함
2302 lines
77 KiB
Go
2302 lines
77 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
"runtime"
|
|
"sync"
|
|
"sync/atomic"
|
|
"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"
|
|
)
|
|
|
|
func TestListNodeSnapshotsUsesEdgeQueueState(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-q-1", Alias: "node-q-1"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-q-1",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{
|
|
Name: "ollama-inst",
|
|
Enabled: true,
|
|
Capacity: 3,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
svc.queue.mu.Lock()
|
|
svc.queue.groups["test-group"] = &modelQueueGroup{
|
|
key: "test-group",
|
|
adapter: "ollama-inst",
|
|
inflight: map[string]int{"node-q-1": 2},
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: reg.All()[0], capacity: 3},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
if len(snap.ProviderSnapshots) != 1 {
|
|
t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
|
|
}
|
|
p := snap.ProviderSnapshots[0]
|
|
if p.Adapter != "ollama-inst" {
|
|
t.Errorf("expected adapter ollama-inst, got %s", p.Adapter)
|
|
}
|
|
if p.Capacity != 3 {
|
|
t.Errorf("expected capacity 3, got %d", p.Capacity)
|
|
}
|
|
if p.InFlight != 2 {
|
|
t.Errorf("expected inflight 2, got %d", p.InFlight)
|
|
}
|
|
if p.Queued != 1 {
|
|
t.Errorf("expected queued 1, got %d", p.Queued)
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsTypeRouteSingleNamedInstance(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-q-2", Alias: "node-q-2"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-q-2",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{
|
|
Name: "ollama-local",
|
|
Enabled: true,
|
|
Capacity: 5,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
svc.queue.mu.Lock()
|
|
svc.queue.groups["test-group"] = &modelQueueGroup{
|
|
key: "test-group",
|
|
adapter: "ollama",
|
|
inflight: map[string]int{"node-q-2": 3},
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: reg.All()[0], capacity: 5},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
if len(snap.ProviderSnapshots) != 1 {
|
|
t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
|
|
}
|
|
p := snap.ProviderSnapshots[0]
|
|
if p.Adapter != "ollama-local" {
|
|
t.Errorf("expected adapter ollama-local, got %s", p.Adapter)
|
|
}
|
|
if p.Capacity != 5 {
|
|
t.Errorf("expected capacity 5, got %d", p.Capacity)
|
|
}
|
|
if p.InFlight != 3 {
|
|
t.Errorf("expected inflight 3, got %d", p.InFlight)
|
|
}
|
|
if p.Queued != 1 {
|
|
t.Errorf("expected queued 1, got %d", p.Queued)
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsProviderEmptyAdapterFallbackToProviderID(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-e-1", Alias: "node-e-1"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-e-1",
|
|
Adapters: config.AdaptersConf{},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "standalone-provider",
|
|
Type: "cli",
|
|
Category: config.CategoryCLI,
|
|
Adapter: "",
|
|
Models: []string{"test-model"},
|
|
Capacity: 2,
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 2},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
svc.queue.mu.Lock()
|
|
key := providerResourceKey{nodeID: "node-e-1", providerID: "standalone-provider"}
|
|
res := svc.queue.resources[key]
|
|
if res == nil {
|
|
res = &providerResourceState{
|
|
nodeID: "node-e-1",
|
|
providerID: "standalone-provider",
|
|
capacity: 2,
|
|
enabled: true,
|
|
}
|
|
svc.queue.resources[key] = res
|
|
}
|
|
res.inFlight = 1
|
|
|
|
svc.queue.groups["test-group"] = &modelQueueGroup{
|
|
key: "test-group",
|
|
adapter: "standalone-provider",
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: reg.All()[0], capacity: 2, providerID: "standalone-provider"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
|
|
found := false
|
|
for i, ps := range snap.ProviderSnapshots {
|
|
if ps.Id == "standalone-provider" {
|
|
found = true
|
|
if ps.InFlight != 1 {
|
|
t.Errorf("expected InFlight=1 for standalone-provider (adapter empty), got %d", ps.InFlight)
|
|
}
|
|
if ps.Queued != 1 {
|
|
t.Errorf("expected Queued=1 for standalone-provider (adapter empty), got %d", ps.Queued)
|
|
}
|
|
if ps.Capacity != 2 {
|
|
t.Errorf("expected Capacity=2 for standalone-provider, got %d", ps.Capacity)
|
|
}
|
|
expectedLoadRatio := float32(1) / float32(2)
|
|
if ps.LoadRatio != expectedLoadRatio {
|
|
t.Errorf("expected LoadRatio=%.2f for standalone-provider, got %.2f", expectedLoadRatio, ps.LoadRatio)
|
|
}
|
|
t.Logf("provider snapshot[%d]: id=%s, adapter=%s, InFlight=%d, Queued=%d, Capacity=%d, LoadRatio=%.2f",
|
|
i, ps.Id, ps.Adapter, ps.InFlight, ps.Queued, ps.Capacity, ps.LoadRatio)
|
|
break
|
|
}
|
|
}
|
|
if !found {
|
|
t.Fatalf("expected provider snapshot with id=standalone-provider, got %d provider snapshots", len(snap.ProviderSnapshots))
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsProviderCatalogSupersedesAdapterSnapshots(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-cs-1", Alias: "node-cs-1"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-cs-1",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{Enabled: true},
|
|
OpenAICompatInstances: []config.OpenAICompatInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-cat-1", Adapter: "vllm-gpu", Models: []string{"model-x"}, Health: "available", Capacity: 4},
|
|
{ID: "prov-cat-2", Adapter: "vllm-gpu", Models: []string{"model-y"}, Health: "available", Capacity: 2},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
|
|
if len(snap.ProviderSnapshots) != 2 {
|
|
t.Fatalf("expected 2 catalog snapshots, got %d: %+v",
|
|
len(snap.ProviderSnapshots), snap.ProviderSnapshots)
|
|
}
|
|
ids := map[string]bool{}
|
|
for _, p := range snap.ProviderSnapshots {
|
|
if p.Id == "" {
|
|
t.Errorf("unexpected adapter-only snapshot (empty Id): %+v", p)
|
|
}
|
|
ids[p.Id] = true
|
|
}
|
|
if !ids["prov-cat-1"] || !ids["prov-cat-2"] {
|
|
t.Errorf("expected catalog IDs prov-cat-1 and prov-cat-2, got: %v", ids)
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsCatalogProviderZeroCapacityDoesNotUseRuntimeFallback(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-zero-cap", Alias: "node-zero-cap"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-zero-cap",
|
|
Adapters: config.AdaptersConf{
|
|
OpenAICompatInstances: []config.OpenAICompatInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Capacity: 7},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-zero", Adapter: "vllm-gpu", Models: []string{"model-zero"}, Health: "available", Capacity: 0},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 7},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
if len(snaps[0].ProviderSnapshots) != 1 {
|
|
t.Fatalf("expected 1 provider snapshot, got %d", len(snaps[0].ProviderSnapshots))
|
|
}
|
|
p := snaps[0].ProviderSnapshots[0]
|
|
if p.Id != "prov-zero" {
|
|
t.Fatalf("expected provider id prov-zero, got %q", p.Id)
|
|
}
|
|
if p.Capacity != 0 {
|
|
t.Fatalf("provider catalog capacity must not fall back to runtime concurrency: got %d", p.Capacity)
|
|
}
|
|
if p.LoadRatio != 0 {
|
|
t.Fatalf("expected zero load ratio for zero capacity provider, got %.2f", p.LoadRatio)
|
|
}
|
|
}
|
|
|
|
func TestStaticProviderCatalogSnapshotsZeroCapacityDoesNotUseRuntimeFallback(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-static-zero-cap",
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-zero", Adapter: "vllm-gpu", Models: []string{"model-zero"}, Health: "available", Capacity: 0},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 7},
|
|
}
|
|
|
|
snaps := staticProviderCatalogSnapshots(rec, true)
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 provider snapshot, got %d", len(snaps))
|
|
}
|
|
if snaps[0].Capacity != 0 {
|
|
t.Fatalf("static provider catalog capacity must not fall back to runtime concurrency: got %d", snaps[0].Capacity)
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsLegacyAdapterFallbackWhenNoProviders(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-legacy-snap", Alias: "node-legacy-snap"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-legacy-snap",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{Enabled: true},
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "ollama-local", Enabled: true, Capacity: 3},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 2},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
|
|
if len(snap.ProviderSnapshots) != 2 {
|
|
t.Fatalf("expected 2 adapter snapshots (cli + ollama-local), got %d: %+v",
|
|
len(snap.ProviderSnapshots), snap.ProviderSnapshots)
|
|
}
|
|
adapters := map[string]bool{}
|
|
for _, p := range snap.ProviderSnapshots {
|
|
adapters[p.Adapter] = true
|
|
}
|
|
if !adapters["cli"] || !adapters["ollama-local"] {
|
|
t.Errorf("expected cli and ollama-local adapter snapshots, got: %v", adapters)
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsProviderIdDiffersFromAdapterKey(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-p-1", Alias: "node-p-1"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-p-1",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true, Capacity: 4},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "provider-vllm-primary",
|
|
Type: "vllm",
|
|
Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"nvidia/Qwen3.6-35B"},
|
|
Capacity: 4,
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 2},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
svc.queue.mu.Lock()
|
|
key := providerResourceKey{nodeID: "node-p-1", providerID: "provider-vllm-primary"}
|
|
res := svc.queue.resources[key]
|
|
if res == nil {
|
|
res = &providerResourceState{
|
|
nodeID: "node-p-1",
|
|
providerID: "provider-vllm-primary",
|
|
capacity: 4,
|
|
enabled: true,
|
|
}
|
|
svc.queue.resources[key] = res
|
|
}
|
|
res.inFlight = 2
|
|
|
|
svc.queue.groups["test-group"] = &modelQueueGroup{
|
|
key: "test-group",
|
|
adapter: "vllm-gpu",
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: reg.All()[0], capacity: 4, providerID: "provider-vllm-primary"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
|
|
var p *config.NodeProviderConf
|
|
for i := range snap.ProviderSnapshots {
|
|
ps := snap.ProviderSnapshots[i]
|
|
if ps.Id == "provider-vllm-primary" {
|
|
p = &config.NodeProviderConf{
|
|
Capacity: int(ps.Capacity),
|
|
}
|
|
if ps.InFlight != 2 {
|
|
t.Errorf("expected InFlight=2 for provider-vllm-primary, got %d", ps.InFlight)
|
|
}
|
|
if ps.Queued != 1 {
|
|
t.Errorf("expected Queued=1 for provider-vllm-primary, got %d", ps.Queued)
|
|
}
|
|
if ps.Adapter != "vllm-gpu" {
|
|
t.Errorf("expected Adapter=vllm-gpu for provider-vllm-primary, got %s", ps.Adapter)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
if p == nil {
|
|
t.Fatalf("expected provider snapshot with id=provider-vllm-primary, got %d provider snapshots", len(snap.ProviderSnapshots))
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsLongContextFields(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-long-1", Alias: "node-long-1"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-long-1",
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "long-provider",
|
|
Type: "vllm",
|
|
Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"long-model"},
|
|
Capacity: 4,
|
|
LongContextCapacity: 2,
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 2},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
svc.queue.mu.Lock()
|
|
key := providerResourceKey{nodeID: "node-long-1", providerID: "long-provider"}
|
|
res := svc.queue.resources[key]
|
|
if res == nil {
|
|
res = &providerResourceState{
|
|
nodeID: "node-long-1",
|
|
providerID: "long-provider",
|
|
capacity: 4,
|
|
longCapacity: 2,
|
|
enabled: true,
|
|
}
|
|
svc.queue.resources[key] = res
|
|
}
|
|
res.inFlight = 1
|
|
res.longInFlight = 1
|
|
|
|
svc.queue.groups["test-group"] = &modelQueueGroup{
|
|
key: "test-group",
|
|
adapter: "vllm-gpu",
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: reg.All()[0], capacity: 4, providerID: "long-provider"},
|
|
},
|
|
long: true,
|
|
},
|
|
},
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
if len(snap.ProviderSnapshots) != 1 {
|
|
t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
|
|
}
|
|
p := snap.ProviderSnapshots[0]
|
|
if p.Id != "long-provider" {
|
|
t.Errorf("expected long-provider, got %s", p.Id)
|
|
}
|
|
if p.LongContextCapacity != 2 {
|
|
t.Errorf("expected LongContextCapacity 2, got %d", p.LongContextCapacity)
|
|
}
|
|
if p.LongInFlight != 1 {
|
|
t.Errorf("expected LongInFlight 1, got %d", p.LongInFlight)
|
|
}
|
|
if p.LongQueued != 1 {
|
|
t.Errorf("expected LongQueued 1, got %d", p.LongQueued)
|
|
}
|
|
}
|
|
|
|
func TestProviderSnapshotQueuedCountsCandidatePressureAcrossModelGroups(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
entry := &edgenode.NodeEntry{NodeID: "node-pressure", Alias: "node-pressure"}
|
|
reg.Register(entry)
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-pressure",
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "provider-shared",
|
|
Type: "vllm",
|
|
Category: config.CategoryAPI,
|
|
Models: []string{"model-a", "model-b"},
|
|
Health: "available",
|
|
Capacity: 1,
|
|
LongContextCapacity: 1,
|
|
},
|
|
{
|
|
ID: "provider-fallback",
|
|
Type: "vllm",
|
|
Category: config.CategoryAPI,
|
|
Models: []string{"model-a"},
|
|
Health: "available",
|
|
Capacity: 1,
|
|
LongContextCapacity: 1,
|
|
},
|
|
},
|
|
}
|
|
store.Add(rec)
|
|
|
|
svc := New(reg, edgeevents.NewBus())
|
|
svc.SetNodeStore(store)
|
|
|
|
svc.queue.mu.Lock()
|
|
svc.queue.groups["model-a"] = &modelQueueGroup{
|
|
key: "model-a",
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: entry, providerID: "provider-shared"},
|
|
{entry: entry, providerID: "provider-fallback"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
svc.queue.groups["model-b"] = &modelQueueGroup{
|
|
key: "model-b",
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: entry, providerID: "provider-shared"},
|
|
},
|
|
long: true,
|
|
},
|
|
},
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
if len(snaps[0].ProviderSnapshots) != 2 {
|
|
t.Fatalf("expected 2 provider snapshots, got %d", len(snaps[0].ProviderSnapshots))
|
|
}
|
|
|
|
byID := make(map[string]*iop.ProviderSnapshot, len(snaps[0].ProviderSnapshots))
|
|
for _, snapshot := range snaps[0].ProviderSnapshots {
|
|
byID[snapshot.GetId()] = snapshot
|
|
}
|
|
|
|
shared := byID["provider-shared"]
|
|
if shared == nil {
|
|
t.Fatal("provider-shared snapshot not found")
|
|
}
|
|
if shared.GetQueued() != 2 || shared.GetLongQueued() != 1 {
|
|
t.Fatalf("shared candidate pressure mismatch: queued=%d long_queued=%d", shared.GetQueued(), shared.GetLongQueued())
|
|
}
|
|
|
|
fallback := byID["provider-fallback"]
|
|
if fallback == nil {
|
|
t.Fatal("provider-fallback snapshot not found")
|
|
}
|
|
if fallback.GetQueued() != 1 || fallback.GetLongQueued() != 0 {
|
|
t.Fatalf("fallback candidate pressure mismatch: queued=%d long_queued=%d", fallback.GetQueued(), fallback.GetLongQueued())
|
|
}
|
|
|
|
if total := shared.GetQueued() + fallback.GetQueued(); total != 3 {
|
|
t.Fatalf("provider candidate pressure should be non-additive across 2 Edge pending requests: got total=%d", total)
|
|
}
|
|
}
|
|
|
|
func TestProviderSnapshotReadsResourceState(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-snap-test", Alias: "node-snap-test"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-snap-test",
|
|
Runtime: config.RuntimeConf{Concurrency: 5},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "snap-prov",
|
|
Adapter: "vllm",
|
|
Capacity: 3,
|
|
LongContextCapacity: 2,
|
|
Enabled: nil,
|
|
},
|
|
},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
svc.queue.mu.Lock()
|
|
key := providerResourceKey{nodeID: "node-snap-test", providerID: "snap-prov"}
|
|
res := svc.queue.resources[key]
|
|
if res == nil {
|
|
svc.queue.mu.Unlock()
|
|
t.Fatal("resource state not found after SetNodeStore")
|
|
}
|
|
res.inFlight = 2
|
|
res.longInFlight = 1
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot, got %d", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
if len(snap.ProviderSnapshots) != 1 {
|
|
t.Fatalf("expected 1 provider snapshot, got %d", len(snap.ProviderSnapshots))
|
|
}
|
|
p := snap.ProviderSnapshots[0]
|
|
if p.InFlight != 2 {
|
|
t.Errorf("expected InFlight 2, got %d", p.InFlight)
|
|
}
|
|
if p.LongInFlight != 1 {
|
|
t.Errorf("expected LongInFlight 1, got %d", p.LongInFlight)
|
|
}
|
|
}
|
|
|
|
func TestProviderSnapshotIgnoresLegacyFallbackState(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-fallback-test", Alias: "node-fallback-test"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-fallback-test",
|
|
Runtime: config.RuntimeConf{Concurrency: 5},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "fallback-prov",
|
|
Adapter: "vllm",
|
|
Capacity: 3,
|
|
LongContextCapacity: 2,
|
|
Enabled: nil,
|
|
},
|
|
},
|
|
}
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
trackTestLease(svc.queue, "test-group", "run-legacy", "node-fallback-test", "fallback-prov", true)
|
|
|
|
svc.queue.mu.Lock()
|
|
svc.queue.groups["test-group"] = &modelQueueGroup{
|
|
key: "test-group",
|
|
adapter: "vllm",
|
|
inflight: map[string]int{"node-fallback-test:fallback-prov": 1},
|
|
longInflight: map[string]int{"node-fallback-test:fallback-prov": 1},
|
|
}
|
|
key := providerResourceKey{nodeID: "node-fallback-test", providerID: "fallback-prov"}
|
|
res := svc.queue.resources[key]
|
|
if res == nil {
|
|
svc.queue.mu.Unlock()
|
|
t.Fatal("resource state not found")
|
|
}
|
|
res.inFlight = 0
|
|
res.longInFlight = 0
|
|
svc.queue.mu.Unlock()
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("expected 1 node snapshot")
|
|
}
|
|
p := snaps[0].ProviderSnapshots[0]
|
|
if p.InFlight != 0 {
|
|
t.Errorf("expected InFlight 0, got %d", p.InFlight)
|
|
}
|
|
if p.LongInFlight != 0 {
|
|
t.Errorf("expected LongInFlight 0, got %d", p.LongInFlight)
|
|
}
|
|
}
|
|
|
|
// findProviderSnapshotOnNode finds a provider snapshot by nodeID and providerID
|
|
// from the production ListNodeSnapshots output. It is used by lifecycle tests
|
|
// that assert provider state through the actual snapshot surface rather than
|
|
// reading internal queue/resource maps directly.
|
|
func findProviderSnapshotOnNode(snaps []NodeSnapshot, nodeID, providerID string) *iop.ProviderSnapshot {
|
|
for _, snap := range snaps {
|
|
if snap.NodeID == nodeID {
|
|
for _, ps := range snap.ProviderSnapshots {
|
|
if ps.Id == providerID {
|
|
return ps
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// assertProviderSnapshotState asserts individual provider snapshot fields
|
|
// read from production ListNodeSnapshots. All four counters are independent so
|
|
// that mixed-field regressions (e.g., inFlight set but queued not cleared) are
|
|
// caught. The production surface is the only assertion channel: the internal
|
|
// queue/resource maps are implementation details.
|
|
func assertProviderSnapshotState(t *testing.T, snaps []NodeSnapshot, nodeID, providerID, label string,
|
|
inFlight, queued, longInFlight, longQueued int32) {
|
|
t.Helper()
|
|
ps := findProviderSnapshotOnNode(snaps, nodeID, providerID)
|
|
if ps == nil {
|
|
t.Fatalf("%s: provider %s snapshot not found on %s in ListNodeSnapshots", label, providerID, nodeID)
|
|
}
|
|
if ps.GetId() != providerID {
|
|
t.Errorf("%s %s id mismatch: want %s, got %s", label, providerID, providerID, ps.GetId())
|
|
}
|
|
if ps.GetInFlight() != inFlight {
|
|
t.Errorf("%s %s inFlight: want %d, got %d", label, providerID, inFlight, ps.GetInFlight())
|
|
}
|
|
if ps.GetQueued() != queued {
|
|
t.Errorf("%s %s queued: want %d, got %d", label, providerID, queued, ps.GetQueued())
|
|
}
|
|
if ps.GetLongInFlight() != longInFlight {
|
|
t.Errorf("%s %s longInFlight: want %d, got %d", label, providerID, longInFlight, ps.GetLongInFlight())
|
|
}
|
|
if ps.GetLongQueued() != longQueued {
|
|
t.Errorf("%s %s longQueued: want %d, got %d", label, providerID, longQueued, ps.GetLongQueued())
|
|
}
|
|
}
|
|
|
|
// snapshotTuple describes the exact provider snapshot tuple (status, configured
|
|
// capacity, in-flight, queued, long-context in-flight, long-context queued).
|
|
// Used by the pure validateSnapshotTuple so every concurrent reader compares
|
|
// the same canonical field set without touching *testing.T.
|
|
type snapshotTuple struct {
|
|
status string
|
|
capacity int32
|
|
inFlight int32
|
|
queued int32
|
|
longInFlight int32
|
|
longQueued int32
|
|
}
|
|
|
|
// validateSnapshotTuple compares a snapshotTuple against an expected tuple and
|
|
// returns nil on exact match or an error listing the first mismatched field.
|
|
// Pure function: does not touch *testing.T so storm readers can collect
|
|
// failures through a bounded channel without blocking on t.Fatal.
|
|
func validateSnapshotTuple(got, want snapshotTuple) error {
|
|
if got.status != want.status {
|
|
return fmt.Errorf("status: want %q, got %q", want.status, got.status)
|
|
}
|
|
if got.capacity != want.capacity {
|
|
return fmt.Errorf("capacity: want %d, got %d", want.capacity, got.capacity)
|
|
}
|
|
if got.inFlight != want.inFlight {
|
|
return fmt.Errorf("inFlight: want %d, got %d", want.inFlight, got.inFlight)
|
|
}
|
|
if got.queued != want.queued {
|
|
return fmt.Errorf("queued: want %d, got %d", want.queued, got.queued)
|
|
}
|
|
if got.longInFlight != want.longInFlight {
|
|
return fmt.Errorf("longInFlight: want %d, got %d", want.longInFlight, got.longInFlight)
|
|
}
|
|
if got.longQueued != want.longQueued {
|
|
return fmt.Errorf("longQueued: want %d, got %d", want.longQueued, got.longQueued)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// snapshotFailure carries one reader's first validation error for bounded
|
|
// delivery through a snapReaders-sized channel during the concurrent storm.
|
|
// The main goroutine owns all t.Fatal/Error calls.
|
|
type snapshotFailure struct {
|
|
reader int
|
|
err error
|
|
}
|
|
|
|
// exactProviderTuples builds a provider-ID → snapshotTuple map from a node's
|
|
// snapshot, rejecting missing or duplicate providers with an error. Used by the
|
|
// concurrent storm reader so each expected provider is proven present and only
|
|
// once before any tuple comparison.
|
|
func exactProviderTuples(snaps []NodeSnapshot, nodeID string, expectedIDs ...string) (map[string]snapshotTuple, error) {
|
|
var nodeSnap *NodeSnapshot
|
|
for i := range snaps {
|
|
if snaps[i].NodeID == nodeID {
|
|
nodeSnap = &snaps[i]
|
|
break
|
|
}
|
|
}
|
|
if nodeSnap == nil {
|
|
return nil, fmt.Errorf("node %s not found in ListNodeSnapshots", nodeID)
|
|
}
|
|
seen := make(map[string]snapshotTuple, len(expectedIDs))
|
|
for _, ps := range nodeSnap.ProviderSnapshots {
|
|
pid := ps.GetId()
|
|
if pid == "" {
|
|
continue
|
|
}
|
|
t := snapshotTuple{
|
|
status: ps.GetStatus(),
|
|
capacity: ps.GetCapacity(),
|
|
inFlight: ps.GetInFlight(),
|
|
queued: ps.GetQueued(),
|
|
longInFlight: ps.GetLongInFlight(),
|
|
longQueued: ps.GetLongQueued(),
|
|
}
|
|
if _, dup := seen[pid]; dup {
|
|
return nil, fmt.Errorf("duplicate provider %s on %s", pid, nodeID)
|
|
}
|
|
seen[pid] = t
|
|
}
|
|
for _, wantID := range expectedIDs {
|
|
if _, ok := seen[wantID]; !ok {
|
|
return nil, fmt.Errorf("expected provider %s missing on %s", wantID, nodeID)
|
|
}
|
|
}
|
|
return seen, nil
|
|
}
|
|
|
|
// reportFirstSnapshotFailure sends err through errCh as a reader's first
|
|
// validation failure. Non-blocking: if the channel is full, the reader exits
|
|
// silently so a stuck channel never deadlocks the storm.
|
|
func reportFirstSnapshotFailure(errCh chan<- snapshotFailure, readerID int, err error) {
|
|
select {
|
|
case errCh <- snapshotFailure{reader: readerID, err: err}:
|
|
default:
|
|
}
|
|
}
|
|
|
|
// TestProviderSnapshotLiveCandidatePressureFollowsRefreshAcrossModelGroups
|
|
// verifies that a provider snapshot reflects live candidate pressure from actual
|
|
// SubmitRun waiter dispatches across model groups, covering both the normal and
|
|
// long-context lifecycle (enqueue, dispatch, terminal zero). The holder and
|
|
// waiter go through real admission → pump → terminal lifecycle, not direct
|
|
// queue injection.
|
|
func TestProviderSnapshotLiveCandidatePressureFollowsRefreshAcrossModelGroups(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
contextClass string
|
|
wantLong int
|
|
}{
|
|
{name: "normal"},
|
|
{name: "long", contextClass: contextClassLong, wantLong: 1},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.RunRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
|
|
edgeConn, nodeConn := net.Pipe()
|
|
defer edgeConn.Close()
|
|
defer nodeConn.Close()
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
|
|
const nodeID = "node-lifecycle"
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Client: edgeClient})
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "model-a",
|
|
Providers: map[string]string{"prov-shared": "served-a"},
|
|
},
|
|
{
|
|
ID: "model-b",
|
|
Providers: map[string]string{"prov-fallback": "served-b"},
|
|
},
|
|
{
|
|
ID: "model-wait",
|
|
Providers: map[string]string{
|
|
"prov-shared": "served-shared",
|
|
"prov-fallback": "served-fallback",
|
|
},
|
|
},
|
|
}
|
|
|
|
makeStore := func(longCap int) *edgenode.NodeStore {
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: nodeID,
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-shared", Type: "vllm", Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu", Models: []string{"served-a", "served-shared"},
|
|
Health: "available", Capacity: 1, LongContextCapacity: longCap,
|
|
},
|
|
{
|
|
ID: "prov-fallback", Type: "vllm", Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu", Models: []string{"served-b", "served-fallback"},
|
|
Health: "available", Capacity: 1, LongContextCapacity: longCap,
|
|
},
|
|
},
|
|
})
|
|
return store
|
|
}
|
|
|
|
svc.SetRuntimeConfig(makeStore(tc.wantLong), catalog, groupPolicy{})
|
|
|
|
// 1. Holder A for model-a → dispatched to prov-shared immediately.
|
|
holderA, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-a", ProviderPool: true, Background: true,
|
|
ContextClass: tc.contextClass,
|
|
})
|
|
if err != nil || holderA == nil {
|
|
t.Fatalf("holder A submit: %v", err)
|
|
}
|
|
hA := holderA.Dispatch()
|
|
if got := hA.ProviderID; got != "prov-shared" {
|
|
t.Fatalf("holder A expected prov-shared, got %s", got)
|
|
}
|
|
|
|
// 2. Holder B for model-b → dispatched to prov-fallback immediately.
|
|
holderB, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-b", ProviderPool: true, Background: true,
|
|
ContextClass: tc.contextClass,
|
|
})
|
|
if err != nil || holderB == nil {
|
|
t.Fatalf("holder B submit: %v", err)
|
|
}
|
|
hB := holderB.Dispatch()
|
|
if got := hB.ProviderID; got != "prov-fallback" {
|
|
t.Fatalf("holder B expected prov-fallback, got %s", got)
|
|
}
|
|
|
|
// Snapshot while both holders are in-flight and no waiter: queued must be 0.
|
|
snaps := svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step1-holders",
|
|
1, 0, int32(tc.wantLong), 0)
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step1-holders",
|
|
1, 0, int32(tc.wantLong), 0)
|
|
|
|
// 3. Waiter for model-wait → queues because both providers are full.
|
|
type submitResult struct {
|
|
run RunResult
|
|
err error
|
|
}
|
|
resultCh := make(chan submitResult, 1)
|
|
go func() {
|
|
run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-wait", ProviderPool: true, Background: true,
|
|
ContextClass: tc.contextClass,
|
|
})
|
|
resultCh <- submitResult{run: run, err: err}
|
|
}()
|
|
|
|
// Poll until the waiter is enqueued.
|
|
deadline := time.After(2 * time.Second)
|
|
waiterQueued := false
|
|
for {
|
|
svc.queue.mu.Lock()
|
|
group := svc.queue.groups["model-wait"]
|
|
count := 0
|
|
if group != nil {
|
|
count = len(group.queue)
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
if count == 1 {
|
|
waiterQueued = true
|
|
break
|
|
}
|
|
select {
|
|
case <-deadline:
|
|
t.Fatal("waiter was not enqueued within timeout")
|
|
default:
|
|
runtime.Gosched()
|
|
}
|
|
}
|
|
if !waiterQueued {
|
|
t.Fatal("waiter not enqueued")
|
|
}
|
|
|
|
// Snapshot while waiter is queued: each provider has in_flight=1
|
|
// from its holder AND queued=1 from the waiter's candidate.
|
|
snaps = svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step2-queued",
|
|
1, 1, int32(tc.wantLong), int32(tc.wantLong))
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step2-queued",
|
|
1, 1, int32(tc.wantLong), int32(tc.wantLong))
|
|
|
|
// 4. Release holder B (fallback) → waiter dispatches to fallback.
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: hB.RunID, Type: "complete"})
|
|
|
|
var waiterRes submitResult
|
|
select {
|
|
case waiterRes = <-resultCh:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("waiter did not dispatch after fallback release")
|
|
}
|
|
if waiterRes.err != nil {
|
|
t.Fatalf("waiter submit error: %v", waiterRes.err)
|
|
}
|
|
if waiterRes.run == nil {
|
|
t.Fatal("waiter expected non-nil result")
|
|
}
|
|
defer waiterRes.run.Close()
|
|
wDispatch := waiterRes.run.Dispatch()
|
|
if got := wDispatch.ProviderID; got != "prov-fallback" {
|
|
t.Fatalf("waiter expected dispatch to prov-fallback, got %s", got)
|
|
}
|
|
|
|
// Snapshot after dispatch.
|
|
snaps = svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step3-dispatched",
|
|
1, 0, int32(tc.wantLong), 0)
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step3-dispatched",
|
|
1, 0, int32(tc.wantLong), 0)
|
|
|
|
// 5. Terminal for holder A (shared).
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: hA.RunID, Type: "complete"})
|
|
snaps = svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step4-shared-terminal",
|
|
0, 0, 0, 0)
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step4-shared-terminal",
|
|
1, 0, int32(tc.wantLong), 0)
|
|
|
|
// 6. Terminal for waiter → full drain.
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: wDispatch.RunID, Type: "complete"})
|
|
snaps = svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "step5-all-terminal",
|
|
0, 0, 0, 0)
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "step5-all-terminal",
|
|
0, 0, 0, 0)
|
|
|
|
// Final counters and lease must be fully zero.
|
|
aInFlight, aLong := providerResourceCounts(svc.queue, nodeID, "prov-shared")
|
|
bInFlight, bLong := providerResourceCounts(svc.queue, nodeID, "prov-fallback")
|
|
if aInFlight != 0 || aLong != 0 {
|
|
t.Fatalf("expected prov-shared drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
|
|
}
|
|
if bInFlight != 0 || bLong != 0 {
|
|
t.Fatalf("expected prov-fallback drained to 0, got inFlight=%d longInFlight=%d", bInFlight, bLong)
|
|
}
|
|
if lc := leaseCount(svc.queue); lc != 0 {
|
|
t.Fatalf("expected leaseCount=0 after all terminals, got %d", lc)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestProviderSnapshotConcurrentRefreshPressureUsesLiveCandidates drives
|
|
// concurrent runtime refresh and snapshot reads against a queued waiter.
|
|
// Every snapshot tuple is asserted to be a fully-consistent old-or-new state:
|
|
// enabled providers must have (status=available, configured capacity, in-flight
|
|
// from current lease, queued from live candidates); disabled providers must have
|
|
// every effective counter and capacity zero. Each reader records at most one
|
|
// invalid-tuple failure so the storm cannot block on the bounded channel.
|
|
func TestProviderSnapshotConcurrentRefreshPressureUsesLiveCandidates(t *testing.T) {
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.RunRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
|
|
edgeConn, nodeConn := net.Pipe()
|
|
defer edgeConn.Close()
|
|
defer nodeConn.Close()
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
|
|
const nodeID = "node-concurrent"
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Client: edgeClient})
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "model-a",
|
|
Providers: map[string]string{"prov-shared": "served-a"},
|
|
},
|
|
{
|
|
ID: "model-b",
|
|
Providers: map[string]string{"prov-fallback": "served-b"},
|
|
},
|
|
{
|
|
ID: "model-conc",
|
|
Providers: map[string]string{
|
|
"prov-shared": "served-shared",
|
|
"prov-fallback": "served-fallback",
|
|
},
|
|
},
|
|
}
|
|
|
|
makeStore := func(fallbackEnabled bool) *edgenode.NodeStore {
|
|
enabledFlag := fallbackEnabled
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: nodeID,
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-shared", Type: "vllm", Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu", Models: []string{"served-a", "served-shared"},
|
|
Health: "available", Capacity: 1,
|
|
},
|
|
{
|
|
ID: "prov-fallback", Type: "vllm", Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu", Models: []string{"served-b", "served-fallback"},
|
|
Health: "available", Capacity: 1, Enabled: &enabledFlag,
|
|
},
|
|
},
|
|
})
|
|
return store
|
|
}
|
|
|
|
enabledStore := makeStore(true)
|
|
concDisabledStore := makeStore(false)
|
|
|
|
svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
|
|
|
|
holderA, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-a", ProviderPool: true, Background: true,
|
|
})
|
|
if err != nil || holderA == nil {
|
|
t.Fatalf("holder A submit: %v", err)
|
|
}
|
|
defer holderA.Close()
|
|
|
|
holderB, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-b", ProviderPool: true, Background: true,
|
|
})
|
|
if err != nil || holderB == nil {
|
|
t.Fatalf("holder B submit: %v", err)
|
|
}
|
|
defer holderB.Close()
|
|
|
|
type submitResult struct {
|
|
run RunResult
|
|
err error
|
|
}
|
|
resultCh := make(chan submitResult, 1)
|
|
go func() {
|
|
run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-conc", ProviderPool: true, Background: true,
|
|
})
|
|
resultCh <- submitResult{run: run, err: err}
|
|
}()
|
|
|
|
waitForQueueLen(t, svc.queue, "model-conc", 1)
|
|
|
|
// Snapshot while queued with both enabled — using production ListNodeSnapshots
|
|
// so the assertion exercises the actual snapshot surface.
|
|
snaps := svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "conc-pre-storm",
|
|
1, 1, 0, 0)
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "conc-pre-storm",
|
|
1, 1, 0, 0)
|
|
|
|
// Expected tuples for exact concurrent storm validation.
|
|
// prov-shared is always enabled: capacity=1, inFlight=1 (holder A), queued=1 (waiter candidate).
|
|
expectedShared := snapshotTuple{
|
|
status: "available", capacity: 1, inFlight: 1, queued: 1,
|
|
longInFlight: 0, longQueued: 0,
|
|
}
|
|
// prov-fallback alternates: enabled (same tuple as shared) or disabled (all zeros).
|
|
expectedFallbackEnabled := snapshotTuple{
|
|
status: "available", capacity: 1, inFlight: 1, queued: 1,
|
|
longInFlight: 0, longQueued: 0,
|
|
}
|
|
expectedFallbackDisabled := snapshotTuple{
|
|
status: "disabled", capacity: 0, inFlight: 0, queued: 0,
|
|
longInFlight: 0, longQueued: 0,
|
|
}
|
|
|
|
// --- Concurrent storm ---
|
|
// Bounded failure channel sized to reader count: each reader delivers at
|
|
// most one first-error, so the channel never fills and the storm cannot
|
|
// deadlock on send. Main goroutine owns all t.Fatal/Error.
|
|
const snapReaders = 4
|
|
errCh := make(chan snapshotFailure, snapReaders)
|
|
|
|
start := make(chan struct{})
|
|
stop := make(chan struct{})
|
|
var wg sync.WaitGroup
|
|
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
<-start
|
|
for {
|
|
select {
|
|
case <-stop:
|
|
return
|
|
default:
|
|
svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
|
|
svc.SetRuntimeConfig(concDisabledStore, catalog, groupPolicy{})
|
|
}
|
|
}
|
|
}()
|
|
|
|
var (
|
|
fallbackSeenAvailable atomic.Bool
|
|
fallbackSeenDisabled atomic.Bool
|
|
)
|
|
|
|
for r := 0; r < snapReaders; r++ {
|
|
wg.Add(1)
|
|
go func(readerID int) {
|
|
defer wg.Done()
|
|
<-start
|
|
for {
|
|
select {
|
|
case <-stop:
|
|
return
|
|
default:
|
|
}
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) == 0 {
|
|
runtime.Gosched()
|
|
continue
|
|
}
|
|
|
|
byID, err := exactProviderTuples(snaps, nodeID, "prov-shared", "prov-fallback")
|
|
if err != nil {
|
|
reportFirstSnapshotFailure(errCh, readerID, err)
|
|
return
|
|
}
|
|
if err := validateSnapshotTuple(byID["prov-shared"], expectedShared); err != nil {
|
|
reportFirstSnapshotFailure(errCh, readerID, err)
|
|
return
|
|
}
|
|
switch fb := byID["prov-fallback"]; {
|
|
case validateSnapshotTuple(fb, expectedFallbackEnabled) == nil:
|
|
fallbackSeenAvailable.Store(true)
|
|
case validateSnapshotTuple(fb, expectedFallbackDisabled) == nil:
|
|
fallbackSeenDisabled.Store(true)
|
|
default:
|
|
reportFirstSnapshotFailure(errCh, readerID, fmt.Errorf(
|
|
"fallback tuple mismatch: got status=%q capacity=%d inflight=%d queued=%d",
|
|
fb.status, fb.capacity, fb.inFlight, fb.queued))
|
|
return
|
|
}
|
|
}
|
|
}(r)
|
|
}
|
|
|
|
close(start)
|
|
time.Sleep(300 * time.Millisecond)
|
|
close(stop)
|
|
wg.Wait()
|
|
close(errCh)
|
|
|
|
// Main goroutine owns all t.Fatal/Error: each reader delivered at most one
|
|
// first-error through the bounded channel, so this loop never blocks.
|
|
for f := range errCh {
|
|
t.Errorf("reader %d: %v", f.reader, f.err)
|
|
}
|
|
|
|
// Fallback must have been observed in both enabled and disabled states during
|
|
// the storm — a torn-tuple regression that silently passed every fallback read
|
|
// as one state would leave one of the flags unset.
|
|
if !fallbackSeenAvailable.Load() {
|
|
t.Fatal("concurrent storm never sampled fallback in available state")
|
|
}
|
|
if !fallbackSeenDisabled.Load() {
|
|
t.Fatal("concurrent storm never sampled fallback in disabled state")
|
|
}
|
|
|
|
// --- Post-storm ---
|
|
svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
|
|
snaps = svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "conc-post-storm",
|
|
1, 1, 0, 0)
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "conc-post-storm",
|
|
1, 1, 0, 0)
|
|
|
|
// --- Dispatch the waiter via holder release ---
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: holderB.Dispatch().RunID, Type: "complete"})
|
|
|
|
var waiterRes submitResult
|
|
select {
|
|
case waiterRes = <-resultCh:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("waiter did not dispatch after holder release")
|
|
}
|
|
if waiterRes.err != nil {
|
|
t.Fatalf("waiter submit error: %v", waiterRes.err)
|
|
}
|
|
if waiterRes.run == nil {
|
|
t.Fatal("waiter expected non-nil result")
|
|
}
|
|
defer waiterRes.run.Close()
|
|
wDispatch := waiterRes.run.Dispatch()
|
|
if got := wDispatch.ProviderID; got != "prov-fallback" {
|
|
t.Fatalf("waiter expected dispatch to prov-fallback, got %s", got)
|
|
}
|
|
|
|
snaps = svc.ListNodeSnapshots()
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-shared", "conc-after-dispatch",
|
|
1, 0, 0, 0)
|
|
assertProviderSnapshotState(t, snaps, nodeID, "prov-fallback", "conc-after-dispatch",
|
|
1, 0, 0, 0)
|
|
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: holderA.Dispatch().RunID, Type: "complete"})
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: wDispatch.RunID, Type: "complete"})
|
|
|
|
aInFlight, aLong := providerResourceCounts(svc.queue, nodeID, "prov-shared")
|
|
bInFlight, bLong := providerResourceCounts(svc.queue, nodeID, "prov-fallback")
|
|
if aInFlight != 0 || aLong != 0 {
|
|
t.Fatalf("expected prov-shared drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
|
|
}
|
|
if bInFlight != 0 || bLong != 0 {
|
|
t.Fatalf("expected prov-fallback drained to 0, got inFlight=%d longInFlight=%d", bInFlight, bLong)
|
|
}
|
|
if lc := leaseCount(svc.queue); lc != 0 {
|
|
t.Fatalf("expected leaseCount=0 after all terminals, got %d", lc)
|
|
}
|
|
}
|
|
|
|
// TestProviderSnapshotResolverErrorDoesNotUseStaleCandidate verifies that when
|
|
// a queued item's live resolver returns an error or an empty slice, the
|
|
// production ListNodeSnapshots surface does not fall back to the enqueue-time
|
|
// candidate list. The resolver path is exercised by setting resolveCandidates
|
|
// on an actual queued item, then asserting the snapshot reports zero candidate
|
|
// pressure (queued=0, longQueued=0) for the provider the item targets.
|
|
//
|
|
// This regression test exists because a previous loop removed it: with a
|
|
// regular expression test list, a zero-match run exits successfully, so the
|
|
// missing test was silently passing. Both error and empty cases are covered
|
|
// here so the resolver contract (no stale enqueue-time snapshot fallback) is
|
|
// enforced.
|
|
func TestProviderSnapshotResolverErrorDoesNotUseStaleCandidate(t *testing.T) {
|
|
const nodeID = "node-resolver-stale"
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Alias: nodeID})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: nodeID,
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-shared",
|
|
Type: "vllm",
|
|
Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-a"},
|
|
Health: "available",
|
|
Capacity: 1,
|
|
},
|
|
},
|
|
})
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
// Pre-fill provider resource state to simulate an in-flight holder already
|
|
// consuming the slot. The queued waiter's pressure must not show up in the
|
|
// snapshot if its resolver is broken.
|
|
svc.queue.mu.Lock()
|
|
key := providerResourceKey{nodeID: nodeID, providerID: "prov-shared"}
|
|
res := svc.queue.resources[key]
|
|
if res == nil {
|
|
res = &providerResourceState{
|
|
nodeID: nodeID,
|
|
providerID: "prov-shared",
|
|
capacity: 1,
|
|
enabled: true,
|
|
}
|
|
svc.queue.resources[key] = res
|
|
}
|
|
res.inFlight = 1
|
|
|
|
entry := reg.All()[0]
|
|
svc.queue.groups["model-a"] = &modelQueueGroup{
|
|
key: "model-a",
|
|
queue: []*queueItem{
|
|
{
|
|
candidates: []candidateNode{
|
|
{entry: entry, providerID: "prov-shared", capacity: 1},
|
|
},
|
|
waitCh: make(chan admitResult, 1),
|
|
},
|
|
},
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
|
|
assertZeroPressure := func(t *testing.T, label string) {
|
|
t.Helper()
|
|
snaps := svc.ListNodeSnapshots()
|
|
ps := findProviderSnapshotOnNode(snaps, nodeID, "prov-shared")
|
|
if ps == nil {
|
|
t.Fatalf("%s: provider snapshot not found in ListNodeSnapshots", label)
|
|
}
|
|
if got := ps.GetQueued(); got != 0 {
|
|
t.Errorf("%s: expected queued=0, got %d", label, got)
|
|
}
|
|
if got := ps.GetLongQueued(); got != 0 {
|
|
t.Errorf("%s: expected longQueued=0, got %d", label, got)
|
|
}
|
|
}
|
|
|
|
// Case 1: resolver returns error.
|
|
t.Run("resolverError", func(t *testing.T) {
|
|
svc.queue.mu.Lock()
|
|
wg := svc.queue.groups["model-a"]
|
|
if len(wg.queue) == 1 {
|
|
wg.queue[0].resolveCandidates = func() ([]candidateNode, error) {
|
|
return nil, fmt.Errorf("resolver error")
|
|
}
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
assertZeroPressure(t, "resolverError")
|
|
})
|
|
|
|
// Case 2: resolver returns empty slice.
|
|
t.Run("resolverEmpty", func(t *testing.T) {
|
|
svc.queue.mu.Lock()
|
|
wg := svc.queue.groups["model-a"]
|
|
if len(wg.queue) == 1 {
|
|
wg.queue[0].resolveCandidates = func() ([]candidateNode, error) {
|
|
return []candidateNode{}, nil
|
|
}
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
assertZeroPressure(t, "resolverEmpty")
|
|
})
|
|
|
|
// Case 3: resolver works normally — pressure should be reported.
|
|
t.Run("resolverOkHasPressure", func(t *testing.T) {
|
|
svc.queue.mu.Lock()
|
|
wg := svc.queue.groups["model-a"]
|
|
if len(wg.queue) == 1 {
|
|
wg.queue[0].resolveCandidates = func() ([]candidateNode, error) {
|
|
return []candidateNode{{entry: entry, providerID: "prov-shared", capacity: 1}}, nil
|
|
}
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
snaps := svc.ListNodeSnapshots()
|
|
ps := findProviderSnapshotOnNode(snaps, nodeID, "prov-shared")
|
|
if ps == nil {
|
|
t.Fatal("resolverOk: provider snapshot not found")
|
|
}
|
|
if got := ps.GetQueued(); got != 1 {
|
|
t.Errorf("resolverOk: expected queued=1, got %d", got)
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestProviderSnapshotRuntimeRefreshIsOldOrNew forces a queue-held writer to
|
|
// update the runtime config while a reader holds the queue lock, then asserts
|
|
// that the returned snapshot is either the complete old tuple or the complete
|
|
// new tuple — never a torn mix such as available/1/0 between an old available
|
|
// (capacity=1, queued=1) and a new disabled (capacity=0, queued=0) state. The
|
|
// ordering is enforced by having the writer block until the reader holds the
|
|
// queue lock, then releasing; the reader must capture a single consistent
|
|
// snapshot under the queue critical section established in
|
|
// SetRuntimeConfig → ListNodeSnapshots.
|
|
func TestProviderSnapshotRuntimeRefreshIsOldOrNew(t *testing.T) {
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.RunRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
edgeConn, nodeConn := net.Pipe()
|
|
defer edgeConn.Close()
|
|
defer nodeConn.Close()
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
_ = toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
|
|
const nodeID = "node-old-or-new"
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: nodeID, Client: edgeClient})
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
|
|
catalog := []config.ModelCatalogEntry{
|
|
{
|
|
ID: "model-x",
|
|
Providers: map[string]string{"prov-x": "served-x"},
|
|
},
|
|
}
|
|
|
|
enabledFlag := true
|
|
disabledFlag := false
|
|
enabledStore := edgenode.NewNodeStore()
|
|
enabledStore.Add(&edgenode.NodeRecord{
|
|
ID: nodeID,
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm-gpu", Enabled: true},
|
|
},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-x", Type: "vllm", Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{"served-x"},
|
|
Health: "available",
|
|
Capacity: 1,
|
|
Enabled: &enabledFlag,
|
|
},
|
|
},
|
|
})
|
|
|
|
disabledStore := edgenode.NewNodeStore()
|
|
disabledStore.Add(&edgenode.NodeRecord{
|
|
ID: nodeID,
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: "prov-x", Type: "vllm", Category: config.CategoryAPI,
|
|
Health: "available", Capacity: 1,
|
|
Enabled: &disabledFlag,
|
|
},
|
|
},
|
|
})
|
|
|
|
svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
|
|
|
|
// Submit a holder so the provider has inflight+queued pressure under the
|
|
// enabled config. This makes the old tuple (available/capacity=1/inflight>0)
|
|
// distinguishable from the new tuple (disabled/capacity=0/inflight=0).
|
|
holder, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-x", ProviderPool: true, Background: true,
|
|
})
|
|
if err != nil || holder == nil {
|
|
t.Fatalf("holder submit: %v", err)
|
|
}
|
|
defer holder.Close()
|
|
|
|
// Force the queue to hold the lease so inflight > 0 is observable.
|
|
waitForInflight(t, svc.queue, nodeID, "prov-x", 1)
|
|
|
|
// Submit a waiter so the model-x group has one queued item targeting
|
|
// prov-x. This creates live candidate pressure (queued=1) that, combined
|
|
// with the holder's in-flight lease (inFlight=1), makes the old tuple
|
|
// (available/capacity=1/inFlight=1/queued=1) distinguishable from the new
|
|
// disabled tuple (disabled/0/0/0/0/0). Without the waiter the reader would
|
|
// always see queued=0 and the regression would be silent.
|
|
type submitResult struct {
|
|
run RunResult
|
|
err error
|
|
}
|
|
waiterResult := make(chan submitResult, 1)
|
|
go func() {
|
|
run, err := svc.SubmitRun(t.Context(), SubmitRunRequest{
|
|
ModelGroupKey: "model-x", ProviderPool: true, Background: true,
|
|
})
|
|
waiterResult <- submitResult{run: run, err: err}
|
|
}()
|
|
waitForQueueLen(t, svc.queue, "model-x", 1)
|
|
|
|
// --- Public reader barrier: enforce queue→service lock ordering ---
|
|
// Hold the service write lock so the reader goroutine (calling production
|
|
// ListNodeSnapshots) blocks on the service read lock inside the queue-held
|
|
// critical section. This proves the reader actually acquires and holds the
|
|
// queue lock, and the writer must wait for it.
|
|
svc.mu.Lock()
|
|
serviceLocked := true
|
|
defer func() {
|
|
if serviceLocked {
|
|
svc.mu.Unlock()
|
|
}
|
|
}()
|
|
|
|
readerResult := make(chan []NodeSnapshot, 1)
|
|
go func() {
|
|
readerResult <- svc.ListNodeSnapshots()
|
|
}()
|
|
if err := waitForQueueLockHeld(svc.queue, time.Second); err != nil {
|
|
svc.mu.Unlock()
|
|
serviceLocked = false
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// Start writer while reader still holds queue lock and waits on service lock.
|
|
writerDone := make(chan struct{})
|
|
go func() {
|
|
svc.SetRuntimeConfig(disabledStore, catalog, groupPolicy{})
|
|
close(writerDone)
|
|
}()
|
|
svc.mu.Unlock()
|
|
serviceLocked = false
|
|
|
|
// Bounded waits for reader, writer completions.
|
|
var oldSnaps []NodeSnapshot
|
|
select {
|
|
case oldSnaps = <-readerResult:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("public snapshot reader did not complete")
|
|
}
|
|
select {
|
|
case <-writerDone:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("runtime config writer did not complete")
|
|
}
|
|
|
|
expectedOld := snapshotTuple{
|
|
status: "available", capacity: 1, inFlight: 1, queued: 1,
|
|
longInFlight: 0, longQueued: 0,
|
|
}
|
|
assertExactSnapshotTuple(t, oldSnaps, nodeID, "prov-x", expectedOld)
|
|
|
|
// Public surface now returns the new tuple (disabled/capacity=0).
|
|
newSnaps := svc.ListNodeSnapshots()
|
|
expectedNew := snapshotTuple{
|
|
status: "disabled", capacity: 0, inFlight: 0, queued: 0,
|
|
longInFlight: 0, longQueued: 0,
|
|
}
|
|
assertExactSnapshotTuple(t, newSnaps, nodeID, "prov-x", expectedNew)
|
|
|
|
svc.SetRuntimeConfig(enabledStore, catalog, groupPolicy{})
|
|
// Drain holder
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: holder.Dispatch().RunID, Type: "complete"})
|
|
// Drain waiter (holder release unblocks it)
|
|
var waiterRes submitResult
|
|
select {
|
|
case waiterRes = <-waiterResult:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("waiter did not dispatch after holder release")
|
|
}
|
|
if waiterRes.err != nil || waiterRes.run == nil {
|
|
t.Fatalf("waiter result: run=%v err=%v", waiterRes.run, waiterRes.err)
|
|
}
|
|
// Dispatch and complete the waiter to fully drain the provider.
|
|
wDispatch := waiterRes.run.Dispatch()
|
|
svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: wDispatch.RunID, Type: "complete"})
|
|
|
|
// Terminal drain verification: provider resource counts and lease must be zero.
|
|
aInFlight, aLong := providerResourceCounts(svc.queue, nodeID, "prov-x")
|
|
if aInFlight != 0 || aLong != 0 {
|
|
t.Fatalf("expected prov-x drained to 0, got inFlight=%d longInFlight=%d", aInFlight, aLong)
|
|
}
|
|
if lc := leaseCount(svc.queue); lc != 0 {
|
|
t.Fatalf("expected leaseCount=0 after all terminals, got %d", lc)
|
|
}
|
|
}
|
|
|
|
// assertExactSnapshotTuple asserts that the production ListNodeSnapshots
|
|
// surface returns an exact tuple for (nodeID, providerID). Used by the
|
|
// tuple regression test to confirm the reader captures a single consistent
|
|
// snapshot under the queue lock via the public surface.
|
|
func assertExactSnapshotTuple(t *testing.T, snaps []NodeSnapshot, nodeID, providerID string, want snapshotTuple) {
|
|
t.Helper()
|
|
ps := findProviderSnapshotOnNode(snaps, nodeID, providerID)
|
|
if ps == nil {
|
|
t.Fatalf("provider %s snapshot not found on %s in ListNodeSnapshots", providerID, nodeID)
|
|
}
|
|
got := snapshotTuple{
|
|
status: ps.GetStatus(), capacity: ps.GetCapacity(),
|
|
inFlight: ps.GetInFlight(), queued: ps.GetQueued(),
|
|
longInFlight: ps.GetLongInFlight(), longQueued: ps.GetLongQueued(),
|
|
}
|
|
if err := validateSnapshotTuple(got, want); err != nil {
|
|
t.Fatalf("tuple mismatch for %s on %s: %v", providerID, nodeID, err)
|
|
}
|
|
}
|
|
|
|
// waitForInflight blocks until the queue for (nodeID, providerID) observes at
|
|
// least minInflight in-flight leases. Returns immediately if already satisfied.
|
|
func waitForInflight(t *testing.T, q *modelQueueManager, nodeID, providerID string, minInflight int) {
|
|
t.Helper()
|
|
for i := 0; i < 100; i++ {
|
|
inFlight, _ := providerResourceCounts(q, nodeID, providerID)
|
|
if inFlight >= minInflight {
|
|
return
|
|
}
|
|
runtime.Gosched()
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
}
|
|
|
|
// waitForQueueLockHeld proves that another goroutine holds q.mu by attempting
|
|
// TryLock from the caller's goroutine. It returns nil if TryLock fails within
|
|
// the deadline (proving the lock is held exclusively) or a formatted error if
|
|
// TryLock succeeds before the deadline (lock was not held).
|
|
func waitForQueueLockHeld(q *modelQueueManager, deadline time.Duration) error {
|
|
deadlineTimer := time.NewTimer(deadline)
|
|
defer deadlineTimer.Stop()
|
|
ticker := time.NewTicker(time.Millisecond)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-deadlineTimer.C:
|
|
return fmt.Errorf("queue lock not held within %v (TryLock succeeded: lock was not held exclusively)", deadline)
|
|
case <-ticker.C:
|
|
if !q.mu.TryLock() {
|
|
return nil
|
|
}
|
|
q.mu.Unlock()
|
|
runtime.Gosched()
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestListNodeSnapshotsIncludesConfiguredOfflineNodes verifies that configured
|
|
// nodes that are not currently connected still appear in the snapshot with
|
|
// Connected=false and preserve their configured identity and provider catalog.
|
|
// It also verifies that reconnecting with the same NodeID recovers Connected=true
|
|
// while keeping the same identity and stable position, and that the snapshot
|
|
// order is deterministic across no-queue and queue-backed paths.
|
|
func TestListNodeSnapshotsIncludesConfiguredOfflineNodes(t *testing.T) {
|
|
t.Run("noQueue", func(t *testing.T) {
|
|
t.Parallel()
|
|
store := edgenode.NewNodeStore()
|
|
// Add records in reverse index order to prove All() sorts by Index.
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-c",
|
|
Alias: "gamma",
|
|
Index: 2,
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{Enabled: true},
|
|
},
|
|
})
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-a",
|
|
Alias: "alpha",
|
|
Index: 0,
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{Enabled: true},
|
|
},
|
|
})
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-b",
|
|
Alias: "beta",
|
|
Index: 1,
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{Enabled: true},
|
|
},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-b", Type: "cli", Category: config.CategoryCLI, Adapter: "cli", Models: []string{"m"}, Capacity: 2},
|
|
},
|
|
})
|
|
|
|
reg := edgenode.NewRegistry()
|
|
// Only node-a is connected.
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-a", Alias: "alpha"})
|
|
|
|
svc := New(reg, nil)
|
|
if svc.queue != nil {
|
|
t.Fatal("expected no-queue service")
|
|
}
|
|
svc.SetNodeStore(store)
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 3 {
|
|
t.Fatalf("expected 3 node snapshots, got %d", len(snaps))
|
|
}
|
|
|
|
// Stable order: index 0, 1, 2.
|
|
expectedIDs := []string{"node-a", "node-b", "node-c"}
|
|
expectedLabels := []string{"node0", "node1", "node2"}
|
|
expectedConnected := []bool{true, false, false}
|
|
for i, s := range snaps {
|
|
if s.NodeID != expectedIDs[i] {
|
|
t.Errorf("snap[%d].NodeID = %q, want %q", i, s.NodeID, expectedIDs[i])
|
|
}
|
|
if s.Label != expectedLabels[i] {
|
|
t.Errorf("snap[%d].Label = %q, want %q", i, s.Label, expectedLabels[i])
|
|
}
|
|
if s.Connected != expectedConnected[i] {
|
|
t.Errorf("snap[%d].Connected = %v, want %v", i, s.Connected, expectedConnected[i])
|
|
}
|
|
}
|
|
|
|
// The no-queue path preserves the configured provider catalog. Legacy
|
|
// adapter snapshot parity is outside this catalog regression.
|
|
for _, s := range snaps {
|
|
if s.NodeID == "node-b" {
|
|
if len(s.ProviderSnapshots) != 1 {
|
|
t.Errorf("node-b: expected 1 catalog snapshot (prov-b), got %d", len(s.ProviderSnapshots))
|
|
}
|
|
if len(s.ProviderSnapshots) > 0 {
|
|
provider := s.ProviderSnapshots[0]
|
|
// Offline node: catalog identity preserved (id/adapter/type/category/models),
|
|
// but effective status/health/capacity drop to unavailable/offline/0.
|
|
if provider.Id != "prov-b" || provider.Adapter != "cli" || provider.Status != "unavailable" || provider.Health != "offline" || provider.Capacity != 0 || provider.InFlight != 0 || provider.Queued != 0 || len(provider.ServedModels) != 1 || provider.ServedModels[0] != "m" {
|
|
t.Errorf("node-b: unexpected configured provider snapshot (expected unavailable/offline/zero tuple with catalog identity): %+v", provider)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("withQueue", func(t *testing.T) {
|
|
t.Parallel()
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{ID: "node-c", Alias: "gamma", Index: 2, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})
|
|
store.Add(&edgenode.NodeRecord{ID: "node-a", Alias: "alpha", Index: 0, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-b",
|
|
Alias: "beta",
|
|
Index: 1,
|
|
Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}},
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-b", Type: "cli", Category: config.CategoryCLI, Adapter: "cli", Models: []string{"m"}, Capacity: 2},
|
|
},
|
|
})
|
|
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-a", Alias: "alpha"})
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
if svc.queue == nil {
|
|
t.Fatal("expected queue-backed service")
|
|
}
|
|
svc.SetNodeStore(store)
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 3 {
|
|
t.Fatalf("expected 3 node snapshots, got %d", len(snaps))
|
|
}
|
|
|
|
expectedIDs := []string{"node-a", "node-b", "node-c"}
|
|
for i, s := range snaps {
|
|
if s.NodeID != expectedIDs[i] {
|
|
t.Errorf("snap[%d].NodeID = %q, want %q", i, s.NodeID, expectedIDs[i])
|
|
}
|
|
if s.Connected != (s.NodeID == "node-a") {
|
|
t.Errorf("snap[%d] (%s) Connected = %v, want %v", i, s.NodeID, s.Connected, s.NodeID == "node-a")
|
|
}
|
|
if s.NodeID == "node-b" {
|
|
if len(s.ProviderSnapshots) != 1 {
|
|
t.Errorf("node-b: expected 1 queue-backed catalog snapshot, got %d", len(s.ProviderSnapshots))
|
|
} else if provider := s.ProviderSnapshots[0]; provider.Id != "prov-b" || provider.Adapter != "cli" || provider.Status != "unavailable" || provider.Health != "offline" || provider.Capacity != 0 || provider.InFlight != 0 || provider.Queued != 0 || provider.LongInFlight != 0 || provider.LongQueued != 0 {
|
|
t.Errorf("node-b: unexpected queue-backed provider snapshot (expected unavailable/offline/zero tuple): %+v", provider)
|
|
}
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("reconnectPreservesIdentity", func(t *testing.T) {
|
|
t.Parallel()
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{ID: "node-a", Alias: "alpha", Index: 0, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})
|
|
store.Add(&edgenode.NodeRecord{ID: "node-b", Alias: "beta", Index: 1, Adapters: config.AdaptersConf{CLI: config.CLIConf{Enabled: true}}})
|
|
|
|
// Initial: both configured nodes are connected, with node-b on its first
|
|
// accepted registry generation.
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-a", Alias: "alpha"})
|
|
firstNodeB := &edgenode.NodeEntry{NodeID: "node-b", Alias: "beta"}
|
|
reg.Register(firstNodeB)
|
|
if firstNodeB.ConnectionGeneration == 0 {
|
|
t.Fatal("first node-b registration did not receive a generation")
|
|
}
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
assertSnapshot := func(stage string, wantNodeBConnected bool) {
|
|
t.Helper()
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 2 {
|
|
t.Fatalf("%s: expected 2 snapshots, got %d", stage, len(snaps))
|
|
}
|
|
if snaps[0].NodeID != "node-a" || snaps[1].NodeID != "node-b" {
|
|
t.Fatalf("%s: stable identity/order changed: got [%s %s]", stage, snaps[0].NodeID, snaps[1].NodeID)
|
|
}
|
|
if snaps[1].Alias != "beta" || snaps[1].Label != "node1" {
|
|
t.Errorf("%s: node-b identity changed: alias=%q label=%q", stage, snaps[1].Alias, snaps[1].Label)
|
|
}
|
|
if snaps[1].Connected != wantNodeBConnected {
|
|
t.Errorf("%s: node-b Connected=%v, want %v", stage, snaps[1].Connected, wantNodeBConnected)
|
|
}
|
|
}
|
|
assertSnapshot("connected", true)
|
|
|
|
// Disconnect the current owner. The configured identity remains in the
|
|
// same position but connectivity becomes false.
|
|
reg.Unregister("node-b")
|
|
if _, ok := reg.CurrentGeneration("node-b"); ok {
|
|
t.Fatal("node-b generation remained live after unregister")
|
|
}
|
|
assertSnapshot("disconnected", false)
|
|
|
|
// Reconnect the same Node ID through a new registry entry. The registry
|
|
// must mint a later generation while the configured identity/order and
|
|
// snapshot position stay stable.
|
|
secondNodeB := &edgenode.NodeEntry{NodeID: "node-b", Alias: "beta"}
|
|
reg.Register(secondNodeB)
|
|
if secondNodeB.ConnectionGeneration <= firstNodeB.ConnectionGeneration {
|
|
t.Fatalf("reconnect generation=%d, want > first generation=%d", secondNodeB.ConnectionGeneration, firstNodeB.ConnectionGeneration)
|
|
}
|
|
assertSnapshot("reconnected", true)
|
|
})
|
|
}
|
|
|
|
// TestProviderSnapshotOfflineVsDisabled verifies that connectivity flipping
|
|
// only affects enabled providers: disabled providers keep status=disabled and
|
|
// health=disabled regardless of connectivity, while enabled disconnected
|
|
// providers show status=unavailable and health=offline. Catalog identity (id,
|
|
// adapter, type, category, models) is preserved in both cases.
|
|
func TestProviderSnapshotOfflineVsDisabled(t *testing.T) {
|
|
disabled := false
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-off-disabled",
|
|
Providers: []config.NodeProviderConf{
|
|
{ID: "prov-enabled", Type: "vllm", Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu", Models: []string{"served-a"},
|
|
Health: "available", Capacity: 3},
|
|
{ID: "prov-disabled", Type: "vllm", Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu", Models: []string{"served-b"},
|
|
Health: "available", Capacity: 3, Enabled: &disabled},
|
|
},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
|
|
},
|
|
}
|
|
|
|
// wrapForFind converts []*iop.ProviderSnapshot into a single-element
|
|
// []NodeSnapshot so we can reuse findProviderSnapshotOnNode.
|
|
wrapForFind := func(ps []*iop.ProviderSnapshot) []NodeSnapshot {
|
|
return []NodeSnapshot{{NodeID: "node-off-disabled", ProviderSnapshots: ps}}
|
|
}
|
|
|
|
assertDisabled := func(stage string, snaps []NodeSnapshot) {
|
|
t.Helper()
|
|
p := findProviderSnapshotOnNode(snaps, "node-off-disabled", "prov-disabled")
|
|
if p == nil {
|
|
t.Fatalf("%s: prov-disabled not found", stage)
|
|
}
|
|
if p.GetStatus() != "disabled" {
|
|
t.Errorf("%s prov-disabled: status=%q want disabled", stage, p.GetStatus())
|
|
}
|
|
if p.GetHealth() != "disabled" {
|
|
t.Errorf("%s prov-disabled: health=%q want disabled", stage, p.GetHealth())
|
|
}
|
|
if p.GetCapacity() != 0 || p.GetInFlight() != 0 || p.GetQueued() != 0 || p.GetLongInFlight() != 0 || p.GetLongQueued() != 0 {
|
|
t.Errorf("%s prov-disabled: counters must be zero", stage)
|
|
}
|
|
}
|
|
assertEnabled := func(stage string, snaps []NodeSnapshot, wantCap int32, wantStatus, wantHealth string) {
|
|
t.Helper()
|
|
p := findProviderSnapshotOnNode(snaps, "node-off-disabled", "prov-enabled")
|
|
if p == nil {
|
|
t.Fatalf("%s: prov-enabled not found", stage)
|
|
}
|
|
if p.GetStatus() != wantStatus {
|
|
t.Errorf("%s prov-enabled: status=%q want %q", stage, p.GetStatus(), wantStatus)
|
|
}
|
|
if p.GetHealth() != wantHealth {
|
|
t.Errorf("%s prov-enabled: health=%q want %q", stage, p.GetHealth(), wantHealth)
|
|
}
|
|
if p.GetCapacity() != wantCap {
|
|
t.Errorf("%s prov-enabled: capacity=%d want %d", stage, p.GetCapacity(), wantCap)
|
|
}
|
|
if p.GetInFlight() != 0 || p.GetQueued() != 0 || p.GetLongInFlight() != 0 || p.GetLongQueued() != 0 {
|
|
t.Errorf("%s prov-enabled: effective counters must be zero", stage)
|
|
}
|
|
}
|
|
|
|
// --- Disabled: connected ---
|
|
snaps := staticProviderCatalogSnapshots(rec, true)
|
|
assertDisabled("connected", wrapForFind(snaps))
|
|
assertEnabled("connected", wrapForFind(snaps), 3, "available", "available")
|
|
|
|
// --- Disabled: disconnected (must still be disabled, not unavailable) ---
|
|
snaps = staticProviderCatalogSnapshots(rec, false)
|
|
assertDisabled("disconnected", wrapForFind(snaps))
|
|
assertEnabled("disconnected", wrapForFind(snaps), 0, "unavailable", "offline")
|
|
}
|
|
|
|
// TestProviderSnapshotReconnectRestoresCapacityAndAdmission verifies the full
|
|
// generation-aware lease lifecycle across disconnect→offline→reconnect. The
|
|
// first connection acquires a real long-context lease, authoritative disconnect
|
|
// settlement releases it and orphans that generation, and a later connection
|
|
// restores configured capacity while exposing only its own newly admitted lease.
|
|
func TestProviderSnapshotReconnectRestoresCapacityAndAdmission(t *testing.T) {
|
|
const (
|
|
nodeID = "node-reconnect-cycle"
|
|
providerID = "prov-rc"
|
|
modelGroup = "model-rc"
|
|
servedModel = "served-x"
|
|
)
|
|
|
|
rec := &edgenode.NodeRecord{
|
|
ID: nodeID,
|
|
Providers: []config.NodeProviderConf{
|
|
{
|
|
ID: providerID,
|
|
Type: "vllm",
|
|
Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{servedModel},
|
|
Health: "available",
|
|
Capacity: 5,
|
|
LongContextCapacity: 2,
|
|
},
|
|
},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
|
|
},
|
|
}
|
|
|
|
reg := edgenode.NewRegistry()
|
|
firstEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "rc"}
|
|
reg.Register(firstEntry)
|
|
if firstEntry.ConnectionGeneration == 0 {
|
|
t.Fatal("first connection generation must be non-zero")
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(rec)
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
policy := NewGroupPolicy(16, time.Second)
|
|
svc.SetRuntimeConfig(store, []config.ModelCatalogEntry{
|
|
{ID: modelGroup, Providers: map[string]string{providerID: servedModel}},
|
|
}, policy)
|
|
|
|
req := SubmitRunRequest{
|
|
ModelGroupKey: modelGroup,
|
|
ContextClass: contextClassLong,
|
|
ProviderPool: true,
|
|
}
|
|
admitForGeneration := func(stage string, generation uint64) *candidateNode {
|
|
t.Helper()
|
|
candidates, _, err := svc.resolveQueueCandidates(req)
|
|
if err != nil {
|
|
t.Fatalf("%s: resolve provider-pool candidates: %v", stage, err)
|
|
}
|
|
selected, err := svc.queue.admit(
|
|
t.Context(), modelGroup, "", "", candidates, policy,
|
|
svc.resolveQueueCandidatesClosure(req), true, true,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("%s: admit long-context lease: %v", stage, err)
|
|
}
|
|
if selected.providerID != providerID {
|
|
t.Fatalf("%s: admitted provider=%q, want %q", stage, selected.providerID, providerID)
|
|
}
|
|
if selected.generation != generation {
|
|
t.Fatalf("%s: admitted generation=%d, want %d", stage, selected.generation, generation)
|
|
}
|
|
if selected.leaseID == 0 {
|
|
t.Fatalf("%s: admitted lease id must be non-zero", stage)
|
|
}
|
|
svc.queue.trackLease(selected.leaseID, "run-"+stage)
|
|
return selected
|
|
}
|
|
|
|
assertSnapshot := func(stage, wantStatus, wantHealth string, wantCapacity, wantInFlight, wantLongCapacity, wantLongInFlight int32) {
|
|
t.Helper()
|
|
ps := findProviderSnapshotOnNode(svc.ListNodeSnapshots(), nodeID, providerID)
|
|
if ps == nil {
|
|
t.Fatalf("%s: %s snapshot not found", stage, providerID)
|
|
}
|
|
if ps.GetId() != providerID || ps.GetAdapter() != "vllm-gpu" {
|
|
t.Errorf("%s: catalog identity=(%q,%q), want (%q,%q)", stage, ps.GetId(), ps.GetAdapter(), providerID, "vllm-gpu")
|
|
}
|
|
if len(ps.GetServedModels()) != 1 || ps.GetServedModels()[0] != servedModel {
|
|
t.Errorf("%s: served models=%v, want [%s]", stage, ps.GetServedModels(), servedModel)
|
|
}
|
|
if ps.GetStatus() != wantStatus || ps.GetHealth() != wantHealth {
|
|
t.Errorf("%s: status/health=(%q,%q), want (%q,%q)", stage, ps.GetStatus(), ps.GetHealth(), wantStatus, wantHealth)
|
|
}
|
|
if ps.GetCapacity() != wantCapacity || ps.GetInFlight() != wantInFlight {
|
|
t.Errorf("%s: capacity/in-flight=(%d,%d), want (%d,%d)", stage, ps.GetCapacity(), ps.GetInFlight(), wantCapacity, wantInFlight)
|
|
}
|
|
if ps.GetLongContextCapacity() != wantLongCapacity || ps.GetLongInFlight() != wantLongInFlight {
|
|
t.Errorf("%s: long capacity/in-flight=(%d,%d), want (%d,%d)", stage, ps.GetLongContextCapacity(), ps.GetLongInFlight(), wantLongCapacity, wantLongInFlight)
|
|
}
|
|
if ps.GetQueued() != 0 || ps.GetLongQueued() != 0 {
|
|
t.Errorf("%s: queued/long-queued=(%d,%d), want (0,0)", stage, ps.GetQueued(), ps.GetLongQueued())
|
|
}
|
|
t.Logf("%s: generation status=%s health=%s capacity=%d in_flight=%d long_capacity=%d long_in_flight=%d",
|
|
stage, ps.GetStatus(), ps.GetHealth(), ps.GetCapacity(), ps.GetInFlight(), ps.GetLongContextCapacity(), ps.GetLongInFlight())
|
|
}
|
|
|
|
// Stage 1: the first accepted generation owns one real tracked long lease.
|
|
firstLease := admitForGeneration("first-generation", firstEntry.ConnectionGeneration)
|
|
assertSnapshot("first-generation-admitted", "available", "available", 5, 1, 2, 1)
|
|
|
|
// Stage 2: remove the current registry owner, then settle that exact owner
|
|
// generation through the same authoritative service hook used by transport.
|
|
firstGeneration := firstEntry.ConnectionGeneration
|
|
reg.Unregister(nodeID)
|
|
if _, ok := reg.CurrentGeneration(nodeID); ok {
|
|
t.Fatal("disconnected node retained a current registry generation")
|
|
}
|
|
svc.HandleNodeDisconnect(nodeID, firstGeneration, "test authoritative disconnect")
|
|
|
|
svc.queue.mu.Lock()
|
|
_, firstLeaseStillLive := svc.queue.leases[firstLease.leaseID]
|
|
resource := svc.queue.resources[providerResourceKey{nodeID: nodeID, providerID: providerID}]
|
|
resourceOrphan := resource != nil && resource.orphan
|
|
resourceGeneration := uint64(0)
|
|
resourceInFlight, resourceLongInFlight := 0, 0
|
|
if resource != nil {
|
|
resourceGeneration = resource.generation
|
|
resourceInFlight, resourceLongInFlight = resource.snapshotCounts()
|
|
}
|
|
svc.queue.mu.Unlock()
|
|
if firstLeaseStillLive {
|
|
t.Fatal("authoritative disconnect did not settle the first-generation lease")
|
|
}
|
|
if resource == nil || !resourceOrphan || resourceGeneration != firstGeneration {
|
|
t.Fatalf("settled resource=(present=%t orphan=%t generation=%d), want (true,true,%d)", resource != nil, resourceOrphan, resourceGeneration, firstGeneration)
|
|
}
|
|
if resourceInFlight != 0 || resourceLongInFlight != 0 {
|
|
t.Fatalf("settled resource counters=(%d,%d), want (0,0)", resourceInFlight, resourceLongInFlight)
|
|
}
|
|
assertSnapshot("offline-after-settlement", "unavailable", "offline", 0, 0, 0, 0)
|
|
|
|
// Stage 3: reconnect the same identity. Configured capacity is visible, but
|
|
// the old generation contributes no counter before a new admission occurs.
|
|
secondEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "rc"}
|
|
reg.Register(secondEntry)
|
|
if secondEntry.ConnectionGeneration <= firstGeneration {
|
|
t.Fatalf("reconnect generation=%d, want > disconnected generation=%d", secondEntry.ConnectionGeneration, firstGeneration)
|
|
}
|
|
assertSnapshot("reconnected-before-admission", "available", "available", 5, 0, 2, 0)
|
|
|
|
// Stage 4: a real lease on the higher generation is admitted and is the only
|
|
// counter exposed. Releasing that lease returns both counters to zero.
|
|
secondLease := admitForGeneration("second-generation", secondEntry.ConnectionGeneration)
|
|
assertSnapshot("reconnected-after-admission", "available", "available", 5, 1, 2, 1)
|
|
svc.queue.releaseLease(secondLease.leaseID, "test release")
|
|
assertSnapshot("reconnected-after-release", "available", "available", 5, 0, 2, 0)
|
|
}
|
|
|
|
// TestReconnectConnectHookActivatesResourceGenerationInSnapshot verifies the
|
|
// auxiliary S15 convergence: the accepted-connect hook (HandleNodeConnect) alone,
|
|
// with no new admission, clears the disconnect orphan marker, advances the
|
|
// provider resource generation to the reconnect's, and returns the status
|
|
// snapshot to available with the configured capacity and zero in-flight counters.
|
|
func TestReconnectConnectHookActivatesResourceGenerationInSnapshot(t *testing.T) {
|
|
const (
|
|
nodeID = "node-hook-activate"
|
|
providerID = "prov-hook"
|
|
modelGroup = "model-hook"
|
|
servedModel = "served-hook"
|
|
)
|
|
|
|
rec := &edgenode.NodeRecord{
|
|
ID: nodeID,
|
|
Providers: []config.NodeProviderConf{{
|
|
ID: providerID,
|
|
Type: "vllm",
|
|
Category: config.CategoryAPI,
|
|
Adapter: "vllm-gpu",
|
|
Models: []string{servedModel},
|
|
Health: "available",
|
|
Capacity: 4,
|
|
LongContextCapacity: 2,
|
|
}},
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
|
|
},
|
|
}
|
|
|
|
reg := edgenode.NewRegistry()
|
|
firstEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "hook"}
|
|
reg.Register(firstEntry)
|
|
genOne := firstEntry.ConnectionGeneration
|
|
if genOne == 0 {
|
|
t.Fatal("first connection generation must be non-zero")
|
|
}
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(rec)
|
|
|
|
svc := New(reg, edgeevents.NewBus())
|
|
policy := NewGroupPolicy(16, time.Second)
|
|
svc.SetRuntimeConfig(store, []config.ModelCatalogEntry{
|
|
{ID: modelGroup, Providers: map[string]string{providerID: servedModel}},
|
|
}, policy)
|
|
|
|
// Admit a real lease on the first generation so the resource generation is
|
|
// genOne (not zero) before the disconnect.
|
|
req := SubmitRunRequest{ModelGroupKey: modelGroup, ProviderPool: true}
|
|
cands, _, err := svc.resolveQueueCandidates(req)
|
|
if err != nil {
|
|
t.Fatalf("resolve provider-pool candidates: %v", err)
|
|
}
|
|
lease, err := svc.queue.admit(
|
|
t.Context(), modelGroup, "", "", cands, policy,
|
|
svc.resolveQueueCandidatesClosure(req), false, true,
|
|
)
|
|
if err != nil || lease == nil {
|
|
t.Fatalf("admit generation-one lease: %v", err)
|
|
}
|
|
if lease.generation != genOne {
|
|
t.Fatalf("admitted lease generation=%d, want %d", lease.generation, genOne)
|
|
}
|
|
svc.queue.trackLease(lease.leaseID, "run-hook")
|
|
|
|
if present, orphan, gen, inFlight := resourceStateForTest(svc.queue, nodeID, providerID); !present || orphan || gen != genOne || inFlight != 1 {
|
|
t.Fatalf("generation-one resource=(present=%t orphan=%t generation=%d in-flight=%d), want (true,false,%d,1)", present, orphan, gen, inFlight, genOne)
|
|
}
|
|
|
|
// Disconnect the current owner: the lease settles and the resource is orphaned
|
|
// at generation one.
|
|
reg.Unregister(nodeID)
|
|
svc.HandleNodeDisconnect(nodeID, genOne, "disconnected")
|
|
if present, orphan, gen, inFlight := resourceStateForTest(svc.queue, nodeID, providerID); !present || !orphan || gen != genOne || inFlight != 0 {
|
|
t.Fatalf("offline resource=(present=%t orphan=%t generation=%d in-flight=%d), want (true,true,%d,0)", present, orphan, gen, inFlight, genOne)
|
|
}
|
|
if ps := findProviderSnapshotOnNode(svc.ListNodeSnapshots(), nodeID, providerID); ps == nil || ps.GetStatus() != "unavailable" || ps.GetHealth() != "offline" {
|
|
t.Fatalf("offline snapshot = %+v, want status=unavailable health=offline", ps)
|
|
}
|
|
|
|
// Reconnect and drive the accepted-connect hook alone — no new admission,
|
|
// refresh, or lease release.
|
|
secondEntry := &edgenode.NodeEntry{NodeID: nodeID, Alias: "hook"}
|
|
reg.Register(secondEntry)
|
|
genTwo := secondEntry.ConnectionGeneration
|
|
if genTwo <= genOne {
|
|
t.Fatalf("reconnect generation=%d, want > disconnected generation=%d", genTwo, genOne)
|
|
}
|
|
svc.HandleNodeConnect(nodeID, genTwo)
|
|
|
|
present, orphan, gen, inFlight := resourceStateForTest(svc.queue, nodeID, providerID)
|
|
if !present || orphan || gen != genTwo || inFlight != 0 {
|
|
t.Fatalf("activated resource=(present=%t orphan=%t generation=%d in-flight=%d), want (true,false,%d,0)", present, orphan, gen, inFlight, genTwo)
|
|
}
|
|
ps := findProviderSnapshotOnNode(svc.ListNodeSnapshots(), nodeID, providerID)
|
|
if ps == nil {
|
|
t.Fatal("provider snapshot missing after activation")
|
|
}
|
|
if ps.GetStatus() != "available" || ps.GetHealth() != "available" {
|
|
t.Errorf("activated snapshot status/health=(%q,%q), want (available,available)", ps.GetStatus(), ps.GetHealth())
|
|
}
|
|
if ps.GetCapacity() != 4 || ps.GetInFlight() != 0 {
|
|
t.Errorf("activated snapshot capacity/in-flight=(%d,%d), want (4,0)", ps.GetCapacity(), ps.GetInFlight())
|
|
}
|
|
if ps.GetLongContextCapacity() != 2 || ps.GetLongInFlight() != 0 {
|
|
t.Errorf("activated snapshot long capacity/in-flight=(%d,%d), want (2,0)", ps.GetLongContextCapacity(), ps.GetLongInFlight())
|
|
}
|
|
}
|