toki
a3a3451c6a
- Add queue capability contract and FIFO admission queue planning - Update edge node mapper and runtime types for queue support - Extend ollama and vllm adapters with availability/capacity features - Add config updates for queue and provider settings - Update proto definitions for runtime queue messages
443 lines
12 KiB
Go
443 lines
12 KiB
Go
package node_test
|
|
|
|
import (
|
|
"testing"
|
|
|
|
edgenode "iop/apps/edge/internal/node"
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
func TestBuildConfigPayload_OllamaEnabled(t *testing.T) {
|
|
// Simulates the post-normalisation state produced by LoadEdge for a legacy
|
|
// single-instance ollama config: OllamaInstances is populated by LoadEdge,
|
|
// so callers that bypass LoadEdge must populate it manually.
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{
|
|
Name: "ollama",
|
|
Enabled: true,
|
|
BaseURL: "http://localhost:11434",
|
|
ContextSize: 262144,
|
|
Capacity: 3,
|
|
MaxQueue: 7,
|
|
QueueTimeoutMS: 1500,
|
|
RequestTimeoutMS: 30000,
|
|
},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{
|
|
Concurrency: 4,
|
|
WorkspaceRoot: "/tmp/ws",
|
|
},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
if payload.Runtime == nil {
|
|
t.Fatal("expected runtime config, got nil")
|
|
}
|
|
if payload.Runtime.Concurrency != 4 {
|
|
t.Errorf("expected concurrency 4, got %d", payload.Runtime.Concurrency)
|
|
}
|
|
|
|
if len(payload.Adapters) < 2 {
|
|
t.Fatalf("expected at least 2 adapters, got %d", len(payload.Adapters))
|
|
}
|
|
|
|
mockFound := false
|
|
ollamaFound := false
|
|
for _, a := range payload.Adapters {
|
|
if a.Type == "mock" {
|
|
mockFound = true
|
|
}
|
|
if a.Type == "ollama" {
|
|
ollamaFound = true
|
|
ollamaCfg := a.GetOllama()
|
|
if ollamaCfg == nil {
|
|
t.Fatal("expected ollama adapter config")
|
|
}
|
|
if ollamaCfg.BaseUrl != "http://localhost:11434" {
|
|
t.Errorf("expected ollama base url %q, got %q", "http://localhost:11434", ollamaCfg.BaseUrl)
|
|
}
|
|
if ollamaCfg.ContextSize != 262144 {
|
|
t.Errorf("expected ollama context size %d, got %d", 262144, ollamaCfg.ContextSize)
|
|
}
|
|
if ollamaCfg.Capacity != 3 || ollamaCfg.MaxQueue != 7 || ollamaCfg.QueueTimeoutMs != 1500 || ollamaCfg.RequestTimeoutMs != 30000 {
|
|
t.Errorf("unexpected ollama queue config: %+v", ollamaCfg)
|
|
}
|
|
}
|
|
}
|
|
if !mockFound {
|
|
t.Fatal("expected mock adapter")
|
|
}
|
|
if !ollamaFound {
|
|
t.Fatal("expected ollama adapter when enabled")
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_CLIProfiles(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{
|
|
Enabled: true,
|
|
Profiles: map[string]config.CLIProfileConf{
|
|
"default": {
|
|
Command: "echo",
|
|
Args: []string{"hello"},
|
|
Env: []string{"FOO=bar"},
|
|
Persistent: false,
|
|
Terminal: true,
|
|
ResponseIdleTimeoutMS: 5000,
|
|
StartupIdleTimeoutMS: 10000,
|
|
OutputFormat: "plain",
|
|
Mode: "implicit",
|
|
ResumeArgs: []string{"--resume"},
|
|
CompletionMarker: config.CompletionMarkerConf{
|
|
Line: "DONE",
|
|
Regex: "^DONE$",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{
|
|
WorkspaceRoot: "/tmp/ws",
|
|
},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
cliFound := false
|
|
for _, a := range payload.Adapters {
|
|
if a.Type == "cli" {
|
|
cliFound = true
|
|
cliCfg := a.GetCli()
|
|
if cliCfg == nil {
|
|
t.Fatal("expected cli adapter config")
|
|
}
|
|
if len(cliCfg.Profiles) != 1 {
|
|
t.Fatalf("expected 1 cli profile, got %d", len(cliCfg.Profiles))
|
|
}
|
|
p, ok := cliCfg.Profiles["default"]
|
|
if !ok {
|
|
t.Fatal("expected default profile")
|
|
}
|
|
if p.Command != "echo" {
|
|
t.Errorf("expected command %q, got %q", "echo", p.Command)
|
|
}
|
|
if len(p.Args) != 1 || p.Args[0] != "hello" {
|
|
t.Errorf("expected args [hello], got %v", p.Args)
|
|
}
|
|
if p.OutputFormat != "plain" {
|
|
t.Errorf("expected format %q, got %q", "plain", p.OutputFormat)
|
|
}
|
|
if cm := p.CompletionMarker; cm == nil || cm.Line != "DONE" {
|
|
t.Errorf("expected completion marker line %q, got %v", "DONE", cm)
|
|
}
|
|
}
|
|
}
|
|
if !cliFound {
|
|
t.Fatal("expected cli adapter when enabled")
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_CLIProfileOpencodeSSEMode(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{
|
|
Enabled: true,
|
|
Profiles: map[string]config.CLIProfileConf{
|
|
"opencode": {
|
|
Command: "/usr/local/bin/opencode",
|
|
Args: []string{"--model", "m"},
|
|
Mode: "opencode-sse",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{WorkspaceRoot: "/tmp/ws"},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
var found bool
|
|
for _, a := range payload.Adapters {
|
|
if a.Type != "cli" {
|
|
continue
|
|
}
|
|
cli := a.GetCli()
|
|
if cli == nil {
|
|
t.Fatal("expected cli config")
|
|
}
|
|
p, ok := cli.Profiles["opencode"]
|
|
if !ok {
|
|
t.Fatal("expected opencode profile in payload")
|
|
}
|
|
if p.Mode != "opencode-sse" {
|
|
t.Errorf("mode: got %q, want %q", p.Mode, "opencode-sse")
|
|
}
|
|
if len(p.Args) != 2 || p.Args[0] != "--model" || p.Args[1] != "m" {
|
|
t.Errorf("args: got %v", p.Args)
|
|
}
|
|
found = true
|
|
}
|
|
if !found {
|
|
t.Fatal("expected cli adapter in payload")
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_GenericRecordWithoutAdapters(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "generic",
|
|
Token: "token",
|
|
AgentKind: config.AgentKindGenericNode,
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed for generic record: %v", err)
|
|
}
|
|
if payload.GetRuntime() == nil {
|
|
t.Fatal("expected runtime config, got nil")
|
|
}
|
|
if len(payload.GetAdapters()) < 1 || payload.GetAdapters()[0].GetType() != "mock" {
|
|
t.Fatalf("expected mock adapter for generic record, got %+v", payload.GetAdapters())
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_MockAlwaysPresent(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
Ollama: config.OllamaConf{Enabled: false},
|
|
Vllm: config.VllmConf{Enabled: false},
|
|
CLI: config.CLIConf{Enabled: false},
|
|
},
|
|
Runtime: config.RuntimeConf{
|
|
WorkspaceRoot: "/tmp/ws",
|
|
},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
if len(payload.Adapters) < 1 {
|
|
t.Fatalf("expected at least 1 adapter, got %d", len(payload.Adapters))
|
|
}
|
|
|
|
first := payload.Adapters[0]
|
|
if first.Type != "mock" {
|
|
t.Errorf("expected first adapter type %q, got %q", "mock", first.Type)
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_MultiOllamaInstances(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "local", Enabled: true, BaseURL: "http://127.0.0.1:11434", ContextSize: 131072, Capacity: 2, MaxQueue: 4, QueueTimeoutMS: 1000, RequestTimeoutMS: 20000},
|
|
{Name: "dgx", Enabled: true, BaseURL: "http://192.168.0.91:11434", ContextSize: 262144, Capacity: 6, MaxQueue: 12, QueueTimeoutMS: 2000, RequestTimeoutMS: 60000},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{WorkspaceRoot: "/tmp/ws"},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
var ollamaAdapters []*iop.AdapterConfig
|
|
for _, a := range payload.Adapters {
|
|
if a.Type == "ollama" {
|
|
ollamaAdapters = append(ollamaAdapters, a)
|
|
}
|
|
}
|
|
if len(ollamaAdapters) != 2 {
|
|
t.Fatalf("expected 2 ollama adapters, got %d", len(ollamaAdapters))
|
|
}
|
|
names := map[string]bool{}
|
|
for _, a := range ollamaAdapters {
|
|
names[a.Name] = true
|
|
switch a.Name {
|
|
case "local":
|
|
if a.GetOllama().GetCapacity() != 2 || a.GetOllama().GetMaxQueue() != 4 || a.GetOllama().GetQueueTimeoutMs() != 1000 || a.GetOllama().GetRequestTimeoutMs() != 20000 {
|
|
t.Errorf("local queue config mismatch: %+v", a.GetOllama())
|
|
}
|
|
case "dgx":
|
|
if a.GetOllama().GetCapacity() != 6 || a.GetOllama().GetMaxQueue() != 12 || a.GetOllama().GetQueueTimeoutMs() != 2000 || a.GetOllama().GetRequestTimeoutMs() != 60000 {
|
|
t.Errorf("dgx queue config mismatch: %+v", a.GetOllama())
|
|
}
|
|
}
|
|
}
|
|
if !names["local"] || !names["dgx"] {
|
|
t.Errorf("expected instance names local and dgx, got %v", names)
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_MultiVllmInstances(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "a100", Enabled: true, Endpoint: "http://10.0.0.5:8000", Capacity: 4, MaxQueue: 10, QueueTimeoutMS: 1500, RequestTimeoutMS: 45000},
|
|
{Name: "h100", Enabled: true, Endpoint: "http://10.0.0.6:8000", Capacity: 8, MaxQueue: 16, QueueTimeoutMS: 2500, RequestTimeoutMS: 90000},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{WorkspaceRoot: "/tmp/ws"},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
var vllmAdapters []*iop.AdapterConfig
|
|
for _, a := range payload.Adapters {
|
|
if a.Type == "vllm" {
|
|
vllmAdapters = append(vllmAdapters, a)
|
|
}
|
|
}
|
|
if len(vllmAdapters) != 2 {
|
|
t.Fatalf("expected 2 vllm adapters, got %d", len(vllmAdapters))
|
|
}
|
|
names := map[string]bool{}
|
|
for _, a := range vllmAdapters {
|
|
names[a.Name] = true
|
|
if a.GetVllm() == nil {
|
|
t.Errorf("instance %q: expected vllm typed config", a.Name)
|
|
}
|
|
switch a.Name {
|
|
case "a100":
|
|
if a.GetVllm().GetCapacity() != 4 || a.GetVllm().GetMaxQueue() != 10 || a.GetVllm().GetQueueTimeoutMs() != 1500 || a.GetVllm().GetRequestTimeoutMs() != 45000 {
|
|
t.Errorf("a100 queue config mismatch: %+v", a.GetVllm())
|
|
}
|
|
case "h100":
|
|
if a.GetVllm().GetCapacity() != 8 || a.GetVllm().GetMaxQueue() != 16 || a.GetVllm().GetQueueTimeoutMs() != 2500 || a.GetVllm().GetRequestTimeoutMs() != 90000 {
|
|
t.Errorf("h100 queue config mismatch: %+v", a.GetVllm())
|
|
}
|
|
}
|
|
}
|
|
if !names["a100"] || !names["h100"] {
|
|
t.Errorf("expected instance names a100 and h100, got %v", names)
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_LegacyOllamaViaInstances(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
// Simulate what LoadEdge produces after legacy normalisation.
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "ollama", Enabled: true, BaseURL: "http://localhost:11434", ContextSize: 4096, Capacity: 5, MaxQueue: 9, QueueTimeoutMS: 1700, RequestTimeoutMS: 55000},
|
|
},
|
|
},
|
|
Runtime: config.RuntimeConf{WorkspaceRoot: "/tmp/ws"},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
var found *iop.AdapterConfig
|
|
for _, a := range payload.Adapters {
|
|
if a.Type == "ollama" {
|
|
found = a
|
|
break
|
|
}
|
|
}
|
|
if found == nil {
|
|
t.Fatal("expected ollama adapter in payload")
|
|
}
|
|
if found.Name != "ollama" {
|
|
t.Errorf("expected name %q, got %q", "ollama", found.Name)
|
|
}
|
|
if found.GetOllama().GetBaseUrl() != "http://localhost:11434" {
|
|
t.Errorf("unexpected base_url: %q", found.GetOllama().GetBaseUrl())
|
|
}
|
|
if found.GetOllama().GetCapacity() != 5 || found.GetOllama().GetMaxQueue() != 9 || found.GetOllama().GetQueueTimeoutMs() != 1700 || found.GetOllama().GetRequestTimeoutMs() != 55000 {
|
|
t.Errorf("unexpected queue config: %+v", found.GetOllama())
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_DuplicateOllamaNameError(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "dup", Enabled: true, BaseURL: "http://127.0.0.1:11434"},
|
|
{Name: "dup", Enabled: true, BaseURL: "http://127.0.0.2:11434"},
|
|
},
|
|
},
|
|
}
|
|
|
|
_, err := edgenode.BuildConfigPayload(rec)
|
|
if err == nil {
|
|
t.Fatal("expected error for duplicate ollama instance name")
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_MockUsesTypedConfig(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "test",
|
|
Token: "token",
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("BuildConfigPayload failed: %v", err)
|
|
}
|
|
|
|
var mockAdapter *iop.AdapterConfig
|
|
for _, a := range payload.Adapters {
|
|
if a.GetType() == "mock" {
|
|
mockAdapter = a
|
|
break
|
|
}
|
|
}
|
|
if mockAdapter == nil {
|
|
t.Fatal("expected mock adapter in payload")
|
|
}
|
|
if mockAdapter.GetSettings() != nil {
|
|
t.Fatalf("mock adapter must not populate legacy Settings, got %v", mockAdapter.GetSettings())
|
|
}
|
|
if mockAdapter.GetMock() == nil {
|
|
t.Fatal("mock adapter must use typed MockAdapterConfig oneof")
|
|
}
|
|
}
|