166 lines
6.5 KiB
Go
166 lines
6.5 KiB
Go
package server
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/example/ollama-fair-gateway/internal/auth"
|
|
"github.com/example/ollama-fair-gateway/internal/batch"
|
|
"github.com/example/ollama-fair-gateway/internal/config"
|
|
"github.com/example/ollama-fair-gateway/internal/cost"
|
|
"github.com/example/ollama-fair-gateway/internal/metrics"
|
|
px "github.com/example/ollama-fair-gateway/internal/proxy"
|
|
"github.com/example/ollama-fair-gateway/internal/quota"
|
|
"github.com/example/ollama-fair-gateway/internal/scheduler"
|
|
"github.com/example/ollama-fair-gateway/internal/usage"
|
|
"github.com/example/ollama-fair-gateway/internal/worker"
|
|
)
|
|
|
|
func TestDurableBatchEndToEndThroughGatewayPipeline(t *testing.T) {
|
|
var upstreamBody string
|
|
backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/api/ps":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"models":[{"name":"qwen3:8b","model":"qwen3:8b"}]}`)
|
|
case "/api/tags":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"models":[{"name":"qwen3:8b"}]}`)
|
|
case "/api/chat":
|
|
b, _ := io.ReadAll(r.Body)
|
|
upstreamBody = string(b)
|
|
w.Header().Set("Content-Type", "application/x-ndjson")
|
|
_, _ = io.WriteString(w, "{\"message\":{\"content\":\"batch-ok\"},\"done\":true,\"prompt_eval_count\":3,\"eval_count\":2}\n")
|
|
default:
|
|
w.WriteHeader(http.StatusNotFound)
|
|
}
|
|
}))
|
|
defer backend.Close()
|
|
|
|
cfg := &config.Config{
|
|
Server: config.ServerConfig{MaxBodyBytes: 1 << 20, MaxRequestDuration: config.Duration(time.Minute)},
|
|
Auth: config.AuthConfig{APIKeys: []config.APIKeyConfig{{Name: "client", Key: "batch-key", Tenant: "team-a", Subject: "alice"}}},
|
|
Scheduler: config.SchedulerConfig{GlobalConcurrency: 2, MaxQueue: 32, MaxQueuePerActor: 8, QueueTimeout: config.Duration(time.Second), DefaultTenantWeight: 1, DefaultActorWeight: 1, ComputePaths: []string{"/api/chat"}},
|
|
Cost: config.CostConfig{Default: config.ModelRate{InputCreditsPer1K: 1, OutputCreditsPer1K: 3}, DefaultMaxOutputTokens: 16},
|
|
Workers: []config.WorkerConfig{{Name: "w", URL: backend.URL, MaxConcurrent: 2, HealthInterval: config.Duration(time.Hour)}},
|
|
ServiceClasses: config.ServiceClassesConfig{Default: "interactive", Classes: map[string]config.ServiceClassConfig{
|
|
"interactive": {Weight: 1, MaxQueueWait: config.Duration(time.Second), MaxConcurrent: 2},
|
|
"batch": {Weight: 0.25, MaxQueueWait: config.Duration(time.Second), MaxConcurrent: 1},
|
|
}},
|
|
BatchJobs: config.BatchJobsConfig{Enabled: true, Retention: config.Duration(time.Hour), MaxJobs: 100, MaxConcurrent: 1, MaxInputBytes: 1 << 20},
|
|
}
|
|
a, err := auth.New(context.Background(), cfg.Auth)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
wp := worker.New(cfg.Workers, "w")
|
|
root, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
wp.Start(root)
|
|
rec, err := usage.New("", 100, time.Second, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer rec.Close()
|
|
dir := t.TempDir()
|
|
bm, err := batch.New(cfg.BatchJobs, filepath.Join(dir, "batch-jobs.json"), filepath.Join(dir, "batch"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
sv := New(cfg, Dependencies{Auth: a, Scheduler: scheduler.NewLocal(2, 32, 8), Quota: quota.Disabled{}, Estimator: cost.New(cfg.Cost), Workers: wp, Proxy: px.New(), Usage: rec, Metrics: metrics.New(), Logger: slog.Default(), BatchJobs: bm})
|
|
bm.Start(root, sv.ExecuteBatch)
|
|
front := httptest.NewServer(sv.Handler())
|
|
defer front.Close()
|
|
|
|
payload := []byte(`{"path":"/api/chat","body":{"model":"qwen3:8b","messages":[{"role":"user","content":"run batch"}]}}`)
|
|
req, _ := http.NewRequest(http.MethodPost, front.URL+"/gateway/v1/batches", bytes.NewReader(payload))
|
|
req.Header.Set("Authorization", "Bearer batch-key")
|
|
req.Header.Set("Content-Type", "application/json")
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body, _ := io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
if resp.StatusCode != http.StatusAccepted {
|
|
t.Fatalf("create status=%d body=%s", resp.StatusCode, body)
|
|
}
|
|
var created batch.Job
|
|
if err := json.Unmarshal(body, &created); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if created.ID == "" || created.ServiceClass != "batch" || created.Identity.Tenant != "team-a" || created.Identity.Actor != "alice" {
|
|
t.Fatalf("created=%#v", created)
|
|
}
|
|
if !strings.HasPrefix(resp.Header.Get("Location"), "/gateway/v1/batches/") {
|
|
t.Fatalf("location=%q", resp.Header.Get("Location"))
|
|
}
|
|
|
|
var completed batch.Job
|
|
deadline := time.Now().Add(3 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
req, _ = http.NewRequest(http.MethodGet, front.URL+"/gateway/v1/batches/"+created.ID, nil)
|
|
req.Header.Set("Authorization", "Bearer batch-key")
|
|
resp, err = http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body, _ = io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Fatalf("get status=%d body=%s", resp.StatusCode, body)
|
|
}
|
|
if err := json.Unmarshal(body, &completed); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if completed.State == batch.StateCompleted {
|
|
break
|
|
}
|
|
if completed.State == batch.StateFailed || completed.State == batch.StateCancelled {
|
|
t.Fatalf("unexpected terminal state: %#v", completed)
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
if completed.State != batch.StateCompleted || completed.HTTPStatus != http.StatusOK || completed.ExecutionRequestID == "" || completed.OutputRef == "" {
|
|
t.Fatalf("completed=%#v", completed)
|
|
}
|
|
if !strings.Contains(upstreamBody, `"model":"qwen3:8b"`) || !strings.Contains(upstreamBody, `"run batch"`) {
|
|
t.Fatalf("upstream body=%s", upstreamBody)
|
|
}
|
|
|
|
req, _ = http.NewRequest(http.MethodGet, front.URL+"/gateway/v1/batches/"+created.ID+"/output", nil)
|
|
req.Header.Set("Authorization", "Bearer batch-key")
|
|
resp, err = http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
body, _ = io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), `"batch-ok"`) {
|
|
t.Fatalf("output status=%d body=%s", resp.StatusCode, body)
|
|
}
|
|
|
|
events := rec.Recent(10)
|
|
found := false
|
|
for _, e := range events {
|
|
if e.ID == completed.ExecutionRequestID {
|
|
found = true
|
|
if e.ServiceClass != "batch" || e.Tenant != "team-a" || e.Actor != "alice" || e.Usage.PromptTokens != 3 || e.Usage.CompletionTokens != 2 {
|
|
t.Fatalf("usage event=%#v", e)
|
|
}
|
|
}
|
|
}
|
|
if !found {
|
|
t.Fatalf("execution request %s missing from usage: %#v", completed.ExecutionRequestID, events)
|
|
}
|
|
}
|