313 lines
11 KiB
Go
313 lines
11 KiB
Go
package ollama
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/local/glpi-neural-brain/internal/workqueue"
|
|
)
|
|
|
|
func poolTestServer(t *testing.T, calls *atomic.Int64, fail bool) *httptest.Server {
|
|
t.Helper()
|
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
switch r.URL.Path {
|
|
case "/api/tags":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []map[string]any{{"name": "qwen3:8b", "digest": "chat"}, {"name": "embeddinggemma:latest", "digest": "embed"}}})
|
|
case "/api/chat":
|
|
calls.Add(1)
|
|
if fail {
|
|
http.Error(w, "busy", http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"message": map[string]any{"content": `{"ok":true}`}})
|
|
case "/api/embed":
|
|
calls.Add(1)
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"embeddings": [][]float64{{0.1, 0.2}}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
}
|
|
|
|
func TestPoolLeastInflightBalancesSerialRequests(t *testing.T) {
|
|
var aCalls, bCalls atomic.Int64
|
|
a := poolTestServer(t, &aCalls, false)
|
|
defer a.Close()
|
|
b := poolTestServer(t, &bCalls, false)
|
|
defer b.Close()
|
|
c := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "a", URL: a.URL}, {Name: "b", URL: b.URL}}, RoutingMode: "least_inflight", NodeMaxInflight: 1, HealthInterval: time.Minute, FailureCooldown: time.Second, RequestTimeout: time.Second, FailoverEnabled: true, RequireSameModelDigest: true, RequireEmbeddingModel: true}, "qwen3:8b", "embeddinggemma")
|
|
if err := c.Ping(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for i := 0; i < 4; i++ {
|
|
var out struct {
|
|
OK bool `json:"ok"`
|
|
}
|
|
if err := c.ChatJSON(context.Background(), "system", "user", map[string]any{"type": "object"}, &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if aCalls.Load() != 2 || bCalls.Load() != 2 {
|
|
t.Fatalf("distribution a=%d b=%d", aCalls.Load(), bCalls.Load())
|
|
}
|
|
}
|
|
|
|
func TestPoolFailsOverOnRetryableError(t *testing.T) {
|
|
var badCalls, goodCalls atomic.Int64
|
|
bad := poolTestServer(t, &badCalls, true)
|
|
defer bad.Close()
|
|
good := poolTestServer(t, &goodCalls, false)
|
|
defer good.Close()
|
|
c := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "a-bad", URL: bad.URL}, {Name: "b-good", URL: good.URL}}, RoutingMode: "least_inflight", NodeMaxInflight: 1, HealthInterval: time.Minute, FailureCooldown: time.Second, RequestTimeout: time.Second, FailoverEnabled: true, FailoverAttempts: 2, RequireSameModelDigest: true, RequireEmbeddingModel: true}, "qwen3:8b", "embeddinggemma")
|
|
if err := c.Ping(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var out struct {
|
|
OK bool `json:"ok"`
|
|
}
|
|
if err := c.ChatJSON(context.Background(), "system", "user", map[string]any{"type": "object"}, &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if badCalls.Load() != 1 || goodCalls.Load() != 1 {
|
|
t.Fatalf("failover bad=%d good=%d", badCalls.Load(), goodCalls.Load())
|
|
}
|
|
}
|
|
|
|
func TestLowPriorityWaiterYieldsToNormalWaiter(t *testing.T) {
|
|
client := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "only", URL: "http://unused"}}, RoutingMode: "least_inflight", NodeMaxInflight: 1}, "qwen3:8b", "embeddinggemma")
|
|
client.mu.Lock()
|
|
node := client.nodes[0]
|
|
node.healthy = true
|
|
node.compatible = true
|
|
node.chatModel = true
|
|
node.embeddingModel = true
|
|
node.inflight = 1 // hold capacity until both waiters are registered
|
|
client.healthReady = true
|
|
client.mu.Unlock()
|
|
|
|
type result struct {
|
|
name string
|
|
node *nodeState
|
|
err error
|
|
}
|
|
results := make(chan result, 2)
|
|
lowCtx, lowCancel := context.WithTimeout(WithLowPriority(context.Background()), 2*time.Second)
|
|
defer lowCancel()
|
|
go func() {
|
|
n, err := client.acquireNode(lowCtx, "chat", map[*nodeState]bool{})
|
|
results <- result{name: "low", node: n, err: err}
|
|
}()
|
|
waitForPoolWaiters(t, client, 0, 1)
|
|
|
|
normalCtx, normalCancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer normalCancel()
|
|
go func() {
|
|
n, err := client.acquireNode(normalCtx, "chat", map[*nodeState]bool{})
|
|
results <- result{name: "normal", node: n, err: err}
|
|
}()
|
|
waitForPoolWaiters(t, client, 1, 1)
|
|
|
|
client.mu.Lock()
|
|
node.inflight = 0
|
|
client.mu.Unlock()
|
|
|
|
first := <-results
|
|
if first.err != nil {
|
|
t.Fatal(first.err)
|
|
}
|
|
if first.name != "normal" {
|
|
t.Fatalf("low-priority acquisition overtook a normal waiter: first=%s", first.name)
|
|
}
|
|
client.releaseNode(first.node, time.Millisecond, nil)
|
|
|
|
second := <-results
|
|
if second.err != nil {
|
|
t.Fatal(second.err)
|
|
}
|
|
if second.name != "low" {
|
|
t.Fatalf("expected low-priority waiter second, got %s", second.name)
|
|
}
|
|
client.releaseNode(second.node, time.Millisecond, nil)
|
|
}
|
|
|
|
func waitForPoolWaiters(t *testing.T, client *Client, normal, low int) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(time.Second)
|
|
for time.Now().Before(deadline) {
|
|
client.mu.Lock()
|
|
gotNormal, gotLow := client.normalWaiters, client.lowWaiters
|
|
client.mu.Unlock()
|
|
if gotNormal == normal && gotLow == low {
|
|
return
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
client.mu.Lock()
|
|
gotNormal, gotLow := client.normalWaiters, client.lowWaiters
|
|
client.mu.Unlock()
|
|
t.Fatalf("waiter counters did not reach normal=%d low=%d; got normal=%d low=%d", normal, low, gotNormal, gotLow)
|
|
}
|
|
|
|
func TestChatJSONModelUsesRequestedModelOnSamePool(t *testing.T) {
|
|
var requested atomic.Value
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
switch r.URL.Path {
|
|
case "/api/tags":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []map[string]any{
|
|
{"name": "qwen3:8b", "digest": "chat"},
|
|
{"name": "gemma3:12b", "digest": "author"},
|
|
{"name": "embeddinggemma:latest", "digest": "embed"},
|
|
}})
|
|
case "/api/chat":
|
|
var body map[string]any
|
|
_ = json.NewDecoder(r.Body).Decode(&body)
|
|
requested.Store(body["model"])
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"message": map[string]any{"content": `{"ok":true}`}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "only", URL: srv.URL}}, NodeMaxInflight: 1, HealthInterval: time.Minute, RequestTimeout: time.Second, RequireSameModelDigest: true, RequireEmbeddingModel: true}, "qwen3:8b", "embeddinggemma")
|
|
if err := client.Ping(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var out struct {
|
|
OK bool `json:"ok"`
|
|
}
|
|
if err := client.ChatJSONModel(context.Background(), "gemma3:12b", "system", "user", map[string]any{"type": "object"}, &out); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got, _ := requested.Load().(string); got != "gemma3:12b" {
|
|
t.Fatalf("requested model=%q", got)
|
|
}
|
|
if !out.OK {
|
|
t.Fatal("expected structured response")
|
|
}
|
|
status := client.ModelStatus("gemma3:12b")
|
|
if status["healthy_nodes_with_model"] != 1 {
|
|
t.Fatalf("unexpected model status: %#v", status)
|
|
}
|
|
}
|
|
|
|
func TestBusyOllamaWaiterDoesNotOccupySharedResearchSlot(t *testing.T) {
|
|
var chatCalls atomic.Int64
|
|
firstStarted := make(chan struct{})
|
|
releaseFirst := make(chan struct{})
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
switch r.URL.Path {
|
|
case "/api/tags":
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []map[string]any{{"name": "qwen3:8b", "digest": "chat"}, {"name": "embeddinggemma:latest", "digest": "embed"}}})
|
|
case "/api/chat":
|
|
call := chatCalls.Add(1)
|
|
if call == 1 {
|
|
close(firstStarted)
|
|
<-releaseFirst
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"message": map[string]any{"content": `{"ok":true}`}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
|
|
client := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "only", URL: srv.URL}}, NodeMaxInflight: 1, HealthInterval: time.Minute, RequestTimeout: 2 * time.Second, FailoverEnabled: true, FailoverAttempts: 1, RequireSameModelDigest: true, RequireEmbeddingModel: true}, "qwen3:8b", "embeddinggemma")
|
|
shared := workqueue.New(2, 8)
|
|
client.SetSharedLimiter(shared)
|
|
if err := client.Ping(context.Background()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
type response struct{ err error }
|
|
firstDone := make(chan response, 1)
|
|
secondDone := make(chan response, 1)
|
|
go func() {
|
|
var out struct {
|
|
OK bool `json:"ok"`
|
|
}
|
|
firstDone <- response{err: client.ChatJSON(context.Background(), "system", "one", map[string]any{"type": "object"}, &out)}
|
|
}()
|
|
select {
|
|
case <-firstStarted:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("first Ollama request did not start")
|
|
}
|
|
go func() {
|
|
var out struct {
|
|
OK bool `json:"ok"`
|
|
}
|
|
secondDone <- response{err: client.ChatJSON(context.Background(), "system", "two", map[string]any{"type": "object"}, &out)}
|
|
}()
|
|
waitForPoolWaiters(t, client, 1, 0)
|
|
|
|
status := shared.Status()
|
|
if status.Active != 1 || status.Kinds["ollama.chat"].Active != 1 {
|
|
t.Fatalf("busy Ollama waiter consumed shared slot: %+v", status)
|
|
}
|
|
releaseSearch, err := shared.AcquireKind(context.Background(), "searxng.search")
|
|
if err != nil {
|
|
t.Fatalf("SearXNG should still get the second shared slot: %v; status=%+v", err, shared.Status())
|
|
}
|
|
releaseSearch()
|
|
|
|
close(releaseFirst)
|
|
if r := <-firstDone; r.err != nil {
|
|
t.Fatal(r.err)
|
|
}
|
|
if r := <-secondDone; r.err != nil {
|
|
t.Fatal(r.err)
|
|
}
|
|
if got := shared.Status().Active; got != 0 {
|
|
t.Fatalf("shared queue leaked active slots: %+v", shared.Status())
|
|
}
|
|
}
|
|
|
|
func TestAcquireNodeWaitsForHealthyNodeCooldown(t *testing.T) {
|
|
client := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "only", URL: "http://unused"}}, RoutingMode: "least_inflight", NodeMaxInflight: 1}, "qwen3:8b", "embeddinggemma")
|
|
client.mu.Lock()
|
|
node := client.nodes[0]
|
|
node.healthy = true
|
|
node.compatible = true
|
|
node.chatModel = true
|
|
node.embeddingModel = true
|
|
node.cooldownUntil = time.Now().Add(80 * time.Millisecond)
|
|
client.healthReady = true
|
|
client.mu.Unlock()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
started := time.Now()
|
|
got, err := client.acquireNode(ctx, "chat", map[*nodeState]bool{})
|
|
if err != nil {
|
|
t.Fatalf("healthy cooling node should be queued, not rejected: %v", err)
|
|
}
|
|
if elapsed := time.Since(started); elapsed < 50*time.Millisecond {
|
|
t.Fatalf("acquisition did not wait for cooldown: %s", elapsed)
|
|
}
|
|
client.releaseNode(got, time.Millisecond, nil)
|
|
}
|
|
|
|
func TestNodeMaxInflightCanChangeAtRuntime(t *testing.T) {
|
|
client := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "only", URL: "http://unused"}}, NodeMaxInflight: 1}, "qwen3:8b", "embeddinggemma")
|
|
if got := client.NodeMaxInflight(); got != 1 {
|
|
t.Fatalf("initial node max inflight=%d", got)
|
|
}
|
|
client.SetNodeMaxInflight(7)
|
|
if got := client.NodeMaxInflight(); got != 7 {
|
|
t.Fatalf("runtime node max inflight=%d", got)
|
|
}
|
|
client.SetNodeMaxInflight(0)
|
|
if got := client.NodeMaxInflight(); got != 1 {
|
|
t.Fatalf("invalid runtime limit should clamp to 1, got %d", got)
|
|
}
|
|
}
|