121 lines
3.2 KiB
Go
121 lines
3.2 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net"
|
|
"strings"
|
|
"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 TestSubmitRunModelQueueUsesRoutePolicyBeforeProviderInstancePolicy(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)
|
|
nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})
|
|
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "prov-node-route-policy", Client: edgeClient})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "prov-node-route-policy",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{
|
|
Name: "ollama-local",
|
|
Enabled: true,
|
|
Capacity: 1,
|
|
MaxQueue: 10,
|
|
},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{Concurrency: 1},
|
|
})
|
|
|
|
bus := edgeevents.NewBus()
|
|
svc := New(reg, bus)
|
|
svc.SetNodeStore(store)
|
|
|
|
// run 1: fills capacity (capacity is 1)
|
|
res1, err := svc.SubmitRun(context.Background(), SubmitRunRequest{
|
|
ModelGroupKey: "route-policy-group",
|
|
Adapter: "ollama-local",
|
|
Background: true,
|
|
MaxQueue: 1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("run1 error: %v", err)
|
|
}
|
|
defer res1.Close()
|
|
|
|
// run 2: enters queue in a separate goroutine
|
|
ctx2, cancel2 := context.WithCancel(context.Background())
|
|
defer cancel2()
|
|
errCh2 := make(chan error, 1)
|
|
go func() {
|
|
res2, err := svc.SubmitRun(ctx2, SubmitRunRequest{
|
|
ModelGroupKey: "route-policy-group",
|
|
Adapter: "ollama-local",
|
|
Background: true,
|
|
MaxQueue: 1,
|
|
})
|
|
if err == nil {
|
|
res2.Close()
|
|
}
|
|
errCh2 <- err
|
|
}()
|
|
|
|
// Wait explicitly until run 2 enters the queue
|
|
waitForQueueLen(t, svc.queue, "route-policy-group", 1)
|
|
|
|
// run 3: should be rejected immediately because queue is full (max queue = 1, and run2 is in the queue)
|
|
// We use a short timeout context to prevent blocking indefinitely if something fails.
|
|
ctx3, cancel3 := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel3()
|
|
|
|
_, err3 := svc.SubmitRun(ctx3, SubmitRunRequest{
|
|
ModelGroupKey: "route-policy-group",
|
|
Adapter: "ollama-local",
|
|
Background: true,
|
|
MaxQueue: 1,
|
|
})
|
|
if err3 == nil {
|
|
t.Fatal("expected third run to fail immediately due to queue full")
|
|
}
|
|
if !strings.Contains(err3.Error(), "queue is full") {
|
|
t.Errorf("expected queue is full error, got: %v", err3)
|
|
}
|
|
|
|
// clean up run2 and verify it exits
|
|
cancel2()
|
|
select {
|
|
case err2 := <-errCh2:
|
|
if err2 == nil {
|
|
t.Error("expected run2 to fail with context cancelled, got nil")
|
|
} else if !errors.Is(err2, context.Canceled) {
|
|
t.Errorf("expected run2 to fail with context cancelled, got: %v", err2)
|
|
}
|
|
case <-time.After(1 * time.Second):
|
|
t.Error("timeout waiting for run2 goroutine to exit")
|
|
}
|
|
}
|