mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-21 05:57:53 +00:00
feat(config): advertise concurrency via RateLimit-Policy with autodetect
The gateway now sends the concurrency limit as a RateLimit-Policy field from draft-ietf-httpapi-ratelimit-headers. The value is "concurrency";q=<n>;qu="concurrent-requests". The invented X-Concurrency-Limit header is gone. The new concurrencyAutodetect option is off by default. When on, the gateway learns a limit for each configured provider. A 429, 503, 504 or upstream timeout sets the limit to half the in-flight requests, with a minimum of 1. Four consecutive successes add 1. At 64 the limit counts as unknown again. Responses served via a provider advertise the smaller of maxConcurrency and the learned value. An explicit maxConcurrency always wins. Routing, retry and failover do not change.
This commit is contained in:
+165
@@ -0,0 +1,165 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// rateLimitPolicyHeader carries the concurrency limit as an RFC 8941
|
||||
// structured field, per draft-ietf-httpapi-ratelimit-headers.
|
||||
const rateLimitPolicyHeader = "RateLimit-Policy"
|
||||
|
||||
const (
|
||||
concurrencyPolicyName = "concurrency"
|
||||
concurrencyUnit = "concurrent-requests"
|
||||
// successesPerRegrow is how many consecutive upstream successes raise
|
||||
// the learned limit by one.
|
||||
successesPerRegrow = 4
|
||||
// learnedCeiling is the learned value at which the limit counts as
|
||||
// unknown again.
|
||||
learnedCeiling = 64
|
||||
)
|
||||
|
||||
// concurrencyPolicyValue renders n as a RateLimit-Policy list member.
|
||||
func concurrencyPolicyValue(n int) string {
|
||||
return `"` + concurrencyPolicyName + `";q=` + strconv.Itoa(n) + `;qu="` + concurrencyUnit + `"`
|
||||
}
|
||||
|
||||
// providerConcurrency is one provider's learned state.
|
||||
type providerConcurrency struct {
|
||||
inFlight int
|
||||
learned int // 0 means unknown
|
||||
successes int
|
||||
}
|
||||
|
||||
// concurrencyLearner learns, per configured provider, how many requests the
|
||||
// upstream serves at once. Keys are configured provider names only, so
|
||||
// memory stays bounded.
|
||||
type concurrencyLearner struct {
|
||||
byProvider map[string]*providerConcurrency
|
||||
mu sync.Mutex
|
||||
explicit int
|
||||
}
|
||||
|
||||
func newConcurrencyLearner(explicit int) *concurrencyLearner {
|
||||
return &concurrencyLearner{byProvider: map[string]*providerConcurrency{}, explicit: explicit}
|
||||
}
|
||||
|
||||
func (c *concurrencyLearner) state(provider string) *providerConcurrency {
|
||||
s := c.byProvider[provider]
|
||||
if s == nil {
|
||||
s = &providerConcurrency{}
|
||||
c.byProvider[provider] = s
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// begin counts one in-flight upstream request; the returned func ends it.
|
||||
func (c *concurrencyLearner) begin(provider string) func() {
|
||||
c.mu.Lock()
|
||||
c.state(provider).inFlight++
|
||||
c.mu.Unlock()
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
c.mu.Lock()
|
||||
c.state(provider).inFlight--
|
||||
c.mu.Unlock()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// observe feeds one upstream attempt outcome into the provider's state.
|
||||
func (c *concurrencyLearner) observe(provider string, resp *http.Response, err error) {
|
||||
overload := false
|
||||
success := false
|
||||
switch {
|
||||
case resp != nil:
|
||||
code := resp.StatusCode
|
||||
overload = code == http.StatusTooManyRequests || code == http.StatusServiceUnavailable || code == http.StatusGatewayTimeout
|
||||
success = code < http.StatusBadRequest
|
||||
case err != nil:
|
||||
overload = isDeadlineExceeded(err)
|
||||
}
|
||||
if !overload && !success {
|
||||
return
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
s := c.state(provider)
|
||||
if overload {
|
||||
s.successes = 0
|
||||
next := s.inFlight / 2
|
||||
if next < 1 {
|
||||
next = 1
|
||||
}
|
||||
if s.learned == 0 || next < s.learned {
|
||||
s.learned = next
|
||||
}
|
||||
return
|
||||
}
|
||||
if s.learned == 0 {
|
||||
return
|
||||
}
|
||||
s.successes++
|
||||
if s.successes < successesPerRegrow {
|
||||
return
|
||||
}
|
||||
s.successes = 0
|
||||
s.learned++
|
||||
if s.learned >= learnedCeiling {
|
||||
s.learned = 0
|
||||
}
|
||||
}
|
||||
|
||||
// limit returns the effective advertised limit for provider: the smaller of
|
||||
// the explicit and learned values, or 0 when neither is known.
|
||||
func (c *concurrencyLearner) limit(provider string) int {
|
||||
c.mu.Lock()
|
||||
learned := 0
|
||||
if s := c.byProvider[provider]; s != nil {
|
||||
learned = s.learned
|
||||
}
|
||||
c.mu.Unlock()
|
||||
switch {
|
||||
case c.explicit > 0 && (learned == 0 || c.explicit < learned):
|
||||
return c.explicit
|
||||
default:
|
||||
return learned
|
||||
}
|
||||
}
|
||||
|
||||
// setConcurrencyHeader sets or clears the RateLimit-Policy header on h.
|
||||
func setConcurrencyHeader(h http.Header, n int) {
|
||||
if n > 0 {
|
||||
h.Set(rateLimitPolicyHeader, concurrencyPolicyValue(n))
|
||||
return
|
||||
}
|
||||
h.Del(rateLimitPolicyHeader)
|
||||
}
|
||||
|
||||
// trackUpstream wires per-provider accounting into ctx: the limiter's
|
||||
// attempt record and, when concurrencyAutodetect is on, in-flight tracking,
|
||||
// learning and the RateLimit-Policy header on w. The returned func must run
|
||||
// when the upstream exchange ends.
|
||||
func (g *Gateway) trackUpstream(ctx context.Context, w http.ResponseWriter, provider, model string) (context.Context, func()) {
|
||||
c := g.concurrency
|
||||
if c == nil {
|
||||
return withAttemptRecorder(ctx, func(resp *http.Response, attemptErr error) {
|
||||
g.limiter.recordProviderAttempt(provider, model, resp, attemptErr)
|
||||
}), func() {}
|
||||
}
|
||||
end := c.begin(provider)
|
||||
setConcurrencyHeader(w.Header(), c.limit(provider))
|
||||
return withAttemptRecorder(ctx, func(resp *http.Response, attemptErr error) {
|
||||
g.limiter.recordProviderAttempt(provider, model, resp, attemptErr)
|
||||
c.observe(provider, resp, attemptErr)
|
||||
// The watchdog also calls this from a timer goroutine with a nil
|
||||
// resp; only a real response means the request goroutine still owns w.
|
||||
if resp != nil {
|
||||
setConcurrencyHeader(w.Header(), c.limit(provider))
|
||||
}
|
||||
}), end
|
||||
}
|
||||
Reference in New Issue
Block a user