mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-19 05:47:48 +00:00
feat(config): advertise a concurrency limit to clients
Self-hosted upstreams serve few requests at once. A client that opens many requests overloads them. The new maxConcurrency option sends X-Concurrency-Limit on every response. Clients that honour the header cap their own in-flight requests. The default of 0 sends no header, so behaviour does not change.
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func newConcurrencyTestGateway(t *testing.T, maxConcurrency int) (http.Handler, error) {
|
||||
t.Helper()
|
||||
cfg := CreateConfig()
|
||||
cfg.MaxConcurrency = maxConcurrency
|
||||
cfg.Providers = map[string]*ProviderConfig{
|
||||
"openai": {Type: "openai", APIKey: "sk-test", Models: []string{"gpt-test"}},
|
||||
}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "tester", Group: "default", APIKey: "sk-user"}}}
|
||||
return New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
}
|
||||
|
||||
func TestMaxConcurrency_AdvertisedOnResponses(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
header string
|
||||
max int
|
||||
}{
|
||||
{"configured limit is advertised", "3", 3},
|
||||
{"zero sends no header", "", 0},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
h, err := newConcurrencyTestGateway(t, tc.max)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
req.Header.Set("Authorization", "Bearer sk-user")
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if got := rec.Header().Get("X-Concurrency-Limit"); got != tc.header {
|
||||
t.Fatalf("X-Concurrency-Limit = %q, want %q", got, tc.header)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaxConcurrency_NegativeRejected(t *testing.T) {
|
||||
t.Parallel()
|
||||
if _, err := newConcurrencyTestGateway(t, -1); err == nil {
|
||||
t.Fatal("New accepted a negative maxConcurrency")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user