81 lines
2.2 KiB
Go
81 lines
2.2 KiB
Go
package autotune
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/example/ollama-fair-gateway/internal/config"
|
|
"github.com/example/ollama-fair-gateway/internal/worker"
|
|
)
|
|
|
|
func waitProfile(t *testing.T, m *Manager, id string) Profile {
|
|
t.Helper()
|
|
deadline := time.Now().Add(5 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
p, ok := m.Get(id)
|
|
if ok && p.Status != "queued" && p.Status != "running" {
|
|
return p
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
t.Fatalf("profile %s did not finish", id)
|
|
return Profile{}
|
|
}
|
|
|
|
func TestBenchmarkRecommendationApplyAndReload(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/api/generate" {
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/x-ndjson")
|
|
f, _ := w.(http.Flusher)
|
|
fmt.Fprintln(w, `{"response":"x","done":false}`)
|
|
if f != nil {
|
|
f.Flush()
|
|
}
|
|
time.Sleep(20 * time.Millisecond)
|
|
fmt.Fprintln(w, `{"done":true,"prompt_eval_count":10,"prompt_eval_duration":10000000,"eval_count":20,"eval_duration":20000000}`)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
wc := config.WorkerConfig{Name: "gpu", URL: ts.URL, MaxConcurrent: 2}
|
|
pool := worker.New([]config.WorkerConfig{wc}, "gpu")
|
|
path := filepath.Join(t.TempDir(), "auto-tune.json")
|
|
cfg := config.AutoTuningConfig{Enabled: true, MaxConcurrency: 2, SamplesPerLevel: 1, MaxTokens: 32, Timeout: config.Duration(5 * time.Second), Prompt: "test", TTFTWeight: .05, ThroughputWeight: 1}
|
|
m, err := New(cfg, pool, path)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
p, err := m.Start(context.Background(), StartRequest{Worker: "gpu", Model: "qwen:test"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
p = waitProfile(t, m, p.ID)
|
|
if p.Status != "completed" {
|
|
t.Fatalf("status=%s err=%s", p.Status, p.Error)
|
|
}
|
|
if len(p.Levels) != 2 {
|
|
t.Fatalf("levels=%d", len(p.Levels))
|
|
}
|
|
if p.RecommendedConcurrency != 2 {
|
|
t.Fatalf("recommendation=%d want 2; levels=%+v", p.RecommendedConcurrency, p.Levels)
|
|
}
|
|
if _, err := m.Apply(p.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
m2, err := New(cfg, pool, path)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got := m2.Applied()["gpu"]["qwen:test"]; got != 2 {
|
|
t.Fatalf("reloaded applied=%d want 2", got)
|
|
}
|
|
}
|