wehub-resource-sync
82 lines
2.6 KiB
Go
82 lines
2.6 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/labstack/echo/v4"
|
|
"github.com/mudler/LocalAI/core/config"
|
|
"github.com/mudler/LocalAI/core/services/routing/admission"
|
|
"github.com/mudler/LocalAI/core/services/routing/pii"
|
|
)
|
|
|
|
// AdmissionControl runs after RouteModel so the limit applies to the
|
|
// SERVED model — a router fanout that lands on a saturated downstream
|
|
// model gets rejected even though the requested router-model has slack.
|
|
//
|
|
// On reject: HTTP 503, Retry-After header, error JSON. An audit row
|
|
// goes into the shared event store under KindAdmission so admins see
|
|
// rejection rates alongside PII and proxy events.
|
|
//
|
|
// Models without limits.max_concurrent (the common case) hit a fast
|
|
// no-op path — Acquire returns immediately for max <= 0.
|
|
func AdmissionControl(limiter *admission.Limiter, events pii.EventStore) echo.MiddlewareFunc {
|
|
return func(next echo.HandlerFunc) echo.HandlerFunc {
|
|
return func(c echo.Context) error {
|
|
cfg, ok := c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig)
|
|
if !ok || cfg == nil {
|
|
return next(c)
|
|
}
|
|
max := cfg.Limits.MaxConcurrent
|
|
release, ok := limiter.Acquire(cfg.Name, max)
|
|
if !ok {
|
|
retryAfter := admission.RetryAfter(cfg.Limits.RetryAfterSeconds)
|
|
recordAdmissionRejection(events, cfg.Name, retryAfter)
|
|
c.Response().Header().Set("Retry-After", strconv.Itoa(int(retryAfter.Seconds())))
|
|
return c.JSON(http.StatusServiceUnavailable, map[string]any{
|
|
"error": map[string]any{
|
|
"type": "admission_rejected",
|
|
"message": fmt.Sprintf("model %q is at capacity (max_concurrent=%d); retry after %s", cfg.Name, max, retryAfter),
|
|
},
|
|
})
|
|
}
|
|
defer release()
|
|
return next(c)
|
|
}
|
|
}
|
|
}
|
|
|
|
// admissionEventSeq scopes IDs across the process so rapid
|
|
// rejections under load get unique row IDs without coordinating
|
|
// with the rest of the event-store ID schemes.
|
|
var admissionEventSeq atomic.Uint64
|
|
|
|
func recordAdmissionRejection(events pii.EventStore, modelName string, retryAfter time.Duration) {
|
|
if events == nil {
|
|
return
|
|
}
|
|
statusCode := http.StatusServiceUnavailable
|
|
durMS := retryAfter.Milliseconds()
|
|
id := fmt.Sprintf("adm_%d_%s", admissionEventSeq.Add(1), randHex(4))
|
|
_ = events.Record(context.Background(), pii.PIIEvent{
|
|
ID: id,
|
|
Kind: pii.KindAdmission,
|
|
Host: modelName,
|
|
StatusCode: statusCode,
|
|
DurationMS: durMS,
|
|
CreatedAt: time.Now().UTC(),
|
|
})
|
|
}
|
|
|
|
func randHex(n int) string {
|
|
b := make([]byte, n)
|
|
_, _ = rand.Read(b)
|
|
return hex.EncodeToString(b)
|
|
}
|