mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-19 16:39:15 +00:00
feat: proactive rate limiting (RPM) + 429 short cooldown for sources
- config.go: Source add RPM field (requests-per-minute cap, 0=unlimited) - provider.go: RecordRateLimit() — 429 uses fixed 30s cooldown, not exponential - provider.go: Throttle() — token-bucket proactive rate limiter, spaces requests at 60s/RPM interval, respects context cancellation - provider.go: ReportStatus() — 429 -> RecordRateLimit, 5xx -> RecordFailure - provider.go: Chat/ChatStream — wire Throttle after TryAcquire - api.go: sourcePayload + RPM, buildSource passes RPM through - ui/index.html: add RPM input field in source editor, bilingual i18n labels - deploy.sh: backup old binary + rollback on healthcheck failure - provider_test.go: TestModelStateRateLimitShortCooldown, TestThrottleSpacingAndCancel - config.yaml: sensenova rpm: 12
This commit is contained in:
JianFeeeee
274
deploy.sh
Executable file
274
deploy.sh
Executable file
@ -0,0 +1,274 @@
|
||||
#!/usr/bin/env bash
|
||||
# ============================================================
|
||||
# llmsproxy 原子化滚动部署脚本
|
||||
#
|
||||
# 策略：
|
||||
# 1. /tmp (tmpfs) 构建，/usr/local/bin (ext4) 部署
|
||||
# 2. 构建后先 copy 到目标文件系统（.llmsproxy.tmp），
|
||||
# 再 rename(2) 做原子替换（同一文件系统内 mv 是原子的）
|
||||
# 3. runtime.json 不替换（保留运行时密钥、审计状态）
|
||||
# 4. systemctl restart → 服务中断仅在 Go 进程重启瞬间 (<500ms)
|
||||
# 5. 健康检查确认新进程正常
|
||||
#
|
||||
# 注意事项：
|
||||
# - 本服务依赖自身（推理依赖），restart 前后会短暂中断，
|
||||
# 但重试机制（systemd RestartSec=5 + 客户端 retry）可自愈
|
||||
# - 适配器文件替换后需重启才能生效（VM 启动时批量加载）
|
||||
# ============================================================
|
||||
set -euo pipefail
|
||||
|
||||
# ---------- 配置 ----------
|
||||
readonly TARGET_BIN="/usr/local/bin/llmsproxy"
|
||||
readonly TARGET_ADAPTERS="/etc/llmsproxy/adapters"
|
||||
readonly REPO_DIR="/home/program/llmsproxy"
|
||||
readonly HEALTH_URL="http://127.0.0.1:8081/v1/models"
|
||||
readonly BIN_NAME="llmsproxy"
|
||||
|
||||
# 构建目录在 /tmp（tmpfs 更快）
|
||||
readonly TMP_DIR="/tmp/llmsproxy-deploy"
|
||||
# 同文件系统暂存文件（用于原子 rename）
|
||||
readonly STAGING_BIN="${TARGET_BIN}.staging"
|
||||
# 回滚备份：健康检查失败时从这里恢复旧二进制
|
||||
readonly BACKUP_BIN="${TARGET_BIN}.bak"
|
||||
|
||||
# ---------- 日志 ----------
|
||||
log() { printf '\033[1;34m[%s]\033[0m %s\n' "$(date +%H:%M:%S)" "$*"; }
|
||||
warn() { printf '\033[1;33m[WARN]\033[0m %s\n' "$*"; }
|
||||
err() { printf '\033[1;31m[ERR]\033[0m %s\n' "$*" >&2; }
|
||||
fail() { err "$@"; exit 1; }
|
||||
|
||||
# ---------- 清理 ----------
|
||||
cleanup() {
|
||||
log "清理临时文件"
|
||||
rm -rf "$TMP_DIR" "$STAGING_BIN"
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
# ---------- 前置检查 ----------
|
||||
precheck() {
|
||||
log "前置检查"
|
||||
[[ -f "$REPO_DIR/go.mod" ]] || fail "仓库目录 $REPO_DIR 不存在"
|
||||
command -v go &>/dev/null || fail "未找到 go"
|
||||
command -v systemctl &>/dev/null || fail "未找到 systemctl"
|
||||
[[ $(id -u) -eq 0 ]] || fail "需要 root 权限"
|
||||
|
||||
OLD_SIZE=$(stat -c%s "$TARGET_BIN" 2>/dev/null || echo 0)
|
||||
OLD_SHA256=$(sha256sum "$TARGET_BIN" 2>/dev/null | awk '{print $1}')
|
||||
|
||||
log "旧二进制: $TARGET_BIN (${OLD_SIZE} bytes, sha256=${OLD_SHA256:0:16})"
|
||||
log "go: $(go version)"
|
||||
}
|
||||
|
||||
# ---------- 备份旧二进制（用于回滚）----------
|
||||
backup_old_binary() {
|
||||
log "备份旧二进制到 $BACKUP_BIN"
|
||||
if [[ -f "$TARGET_BIN" ]]; then
|
||||
cp -f "$TARGET_BIN" "$BACKUP_BIN"
|
||||
log "✓ 旧二进制已备份 (sha256=$(sha256sum "$BACKUP_BIN" | awk '{print $1}' | cut -c1-16))"
|
||||
else
|
||||
warn "旧二进制不存在，跳过备份（无回滚点）"
|
||||
fi
|
||||
}
|
||||
|
||||
# ---------- 构建 ----------
|
||||
build() {
|
||||
log "构建新二进制"
|
||||
rm -rf "$TMP_DIR"
|
||||
mkdir -p "$TMP_DIR/bin"
|
||||
|
||||
cd "$REPO_DIR"
|
||||
CGO_ENABLED=1 go build -tags luajit \
|
||||
-trimpath \
|
||||
-ldflags="-s -w" \
|
||||
-o "$TMP_DIR/bin/$BIN_NAME" \
|
||||
./cmd/llmsproxy
|
||||
|
||||
NEW_BIN="$TMP_DIR/bin/$BIN_NAME"
|
||||
NEW_SIZE=$(stat -c%s "$NEW_BIN")
|
||||
|
||||
if [[ "$NEW_SIZE" -eq "$OLD_SIZE" ]]; then
|
||||
NEW_SHA256=$(sha256sum "$NEW_BIN" | awk '{print $1}')
|
||||
if [[ "$NEW_SHA256" == "$OLD_SHA256" ]]; then
|
||||
warn "新二进制与旧版本完全一致（sha256 相同），跳过部署"
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
log "新二进制: $NEW_BIN (${NEW_SIZE} bytes)"
|
||||
|
||||
# 验证：确认包含关键修复
|
||||
local checks=(
|
||||
"mergeUsage"
|
||||
"prompt_tokens"
|
||||
"completion_tokens"
|
||||
"message_start"
|
||||
)
|
||||
for c in "${checks[@]}"; do
|
||||
if strings "$NEW_BIN" | grep -qF "$c"; then
|
||||
log " ✓ 新二进制包含: $c"
|
||||
else
|
||||
warn " ⚠ 新二进制不包含: $c（非预期但非致命）"
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
# ---------- 原子替换二进制 ----------
|
||||
# 原理：
|
||||
# /tmp 是 tmpfs，/usr/local/bin 是 ext4——不同文件系统 rename(2) 不原子。
|
||||
# 因此先 copy 到目标文件系统上的 .staging 文件，再 rename(2) 到目标路径。
|
||||
# rename(2) 在同一文件系统内是元数据级操作，对正在读写的进程安全：
|
||||
# - 旧 inode 的 fd 继续有效（进程持有旧文件，不会读到截断内容）
|
||||
# - 新进程打开 path 看到新 inode
|
||||
# - 旧 inode 在最后一个 fd 关闭后释放
|
||||
atomic_replace_binary() {
|
||||
log "原子替换二进制"
|
||||
|
||||
local new_bin_tmp="$TMP_DIR/bin/$BIN_NAME"
|
||||
[[ -f "$new_bin_tmp" ]] || fail "新二进制文件不存在: $new_bin_tmp"
|
||||
|
||||
# Step 1: copy 到目标文件系统（非原子，但目标文件唯一）
|
||||
log " copy → $STAGING_BIN (同文件系统暂存)"
|
||||
cp -f "$new_bin_tmp" "$STAGING_BIN"
|
||||
chmod 0755 "$STAGING_BIN"
|
||||
|
||||
# Step 2: 原子 rename（同一文件系统，元数据级操作）
|
||||
log " rename → $TARGET_BIN (原子)"
|
||||
mv -f "$STAGING_BIN" "$TARGET_BIN"
|
||||
|
||||
# 验证
|
||||
NEW_SHA256=$(sha256sum "$TARGET_BIN" | awk '{print $1}')
|
||||
log "✓ 二进制已原子替换 (sha256=${NEW_SHA256:0:16})"
|
||||
}
|
||||
|
||||
# ---------- 同步适配器 ----------
|
||||
sync_adapters() {
|
||||
log "同步适配器"
|
||||
|
||||
# 只从内嵌适配器目录同步（完整集 10 个适配器，含最新 usage 透传修复）。
|
||||
# 注意：仓库根的 adapters/ 是运行时覆盖目录，可能不全（缺 opencode.lua 等），
|
||||
# 不可作为部署源。
|
||||
SRC_ADAPTERS="$REPO_DIR/internal/lua/adapters"
|
||||
[[ -d "$SRC_ADAPTERS" ]] || fail "内嵌适配器目录不存在: $SRC_ADAPTERS"
|
||||
|
||||
# 备份旧的适配器文件（含任何运行时覆盖版本，如 opencode.lua）
|
||||
BACKUP_DIR="/etc/llmsproxy/adapters.bak.$(date +%Y%m%d%H%M%S)"
|
||||
if [[ -d "$TARGET_ADAPTERS" ]]; then
|
||||
mkdir -p "$BACKUP_DIR"
|
||||
cp -rf "$TARGET_ADAPTERS/"*.lua "$BACKUP_DIR/" 2>/dev/null || true
|
||||
log " 旧适配器已备份到 $BACKUP_DIR"
|
||||
fi
|
||||
|
||||
# 清空目标目录中的 .lua 文件（保留 .bak.* 归档），再复制全部新适配器。
|
||||
# 不依赖 rsync（部署机可能未安装），纯 shell 保证可移植。
|
||||
find "$TARGET_ADAPTERS" -maxdepth 1 -name '*.lua' -not -name '*.bak.*' -delete
|
||||
cp -f "$SRC_ADAPTERS/"*.lua "$TARGET_ADAPTERS/"
|
||||
chmod 0644 "$TARGET_ADAPTERS/"*.lua
|
||||
|
||||
# 校验：每个适配器必须包含 usage 透传修复（opencode/openai/deepseek 等应有 'uses'）
|
||||
local required_files=(openai deepseek anthropic gemini ollama opencode github groq kimicode mistral)
|
||||
for f in "${required_files[@]}"; do
|
||||
[[ -f "$TARGET_ADAPTERS/$f.lua" ]] || warn " 缺少适配器: $f.lua"
|
||||
done
|
||||
local usage_fixed=0
|
||||
for f in openai deepseek github groq kimicode mistral opencode; do
|
||||
grep -qE '\buses\b' "$TARGET_ADAPTERS/$f.lua" 2>/dev/null && ((usage_fixed++)) || true
|
||||
done
|
||||
log " usage 透传修复适配器: ${usage_fixed}/7 (openai/deepseek/github/groq/kimicode/mistral/opencode)"
|
||||
log "✓ 适配器已同步到 $TARGET_ADAPTERS/"
|
||||
}
|
||||
|
||||
# ---------- 重启服务 ----------
|
||||
restart_service() {
|
||||
log "重启服务（中断 <500ms）"
|
||||
systemctl restart llmsproxy.service || fail "systemd restart 失败"
|
||||
log "✓ llmsproxy.service restarted"
|
||||
}
|
||||
|
||||
# ---------- 健康检查 ----------
|
||||
healthcheck() {
|
||||
log "健康检查"
|
||||
local attempts=20
|
||||
local wait_sec=1
|
||||
local attempt=1
|
||||
|
||||
while (( attempt <= attempts )); do
|
||||
# 用 curl 检查（200=有 key 未鉴权但服务正常，401=需要认证，都说明服务在运行）
|
||||
local code
|
||||
code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 3 "$HEALTH_URL" 2>/dev/null || echo "000")
|
||||
if [[ "$code" == "200" || "$code" == "401" ]]; then
|
||||
log "✓ 服务健康 (status=$code, ${attempt}/${attempts})"
|
||||
return 0
|
||||
fi
|
||||
log " 等待中 ($attempt/$attempts, status=$code)"
|
||||
sleep "$wait_sec"
|
||||
(( attempt++ ))
|
||||
done
|
||||
|
||||
warn "健康检查未通过，请检查服务状态: systemctl status llmsproxy.service"
|
||||
return 1
|
||||
}
|
||||
|
||||
# ---------- 回滚 ----------
|
||||
# 健康检查失败时恢复旧二进制并重启，保证推理服务可用。
|
||||
# 这是最后一道保险：即使新版本启动失败（编译错误、配置不兼容、
|
||||
# panic on start 等），也能让旧版本继续服务，避免推理能力丢失。
|
||||
rollback() {
|
||||
err "健康检查失败，启动回滚流程"
|
||||
if [[ ! -f "$BACKUP_BIN" ]]; then
|
||||
err "无回滚备份 ($BACKUP_BIN 不存在)，无法自动恢复"
|
||||
err "请手动检查: systemctl status llmsproxy.service && journalctl -u llmsproxy -n 50"
|
||||
return 1
|
||||
fi
|
||||
|
||||
log "恢复旧二进制"
|
||||
cp -f "$BACKUP_BIN" "$TARGET_BIN"
|
||||
chmod 0755 "$TARGET_BIN"
|
||||
|
||||
log "重启服务（旧版本）"
|
||||
systemctl restart llmsproxy.service || {
|
||||
err "重启失败，服务可能完全宕机"
|
||||
err "请手动检查: systemctl status llmsproxy.service"
|
||||
return 1
|
||||
}
|
||||
|
||||
# 回滚后再健康检查（只查 1 次，失败就认了）
|
||||
sleep 2
|
||||
local code
|
||||
code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 "$HEALTH_URL" 2>/dev/null || echo "000")
|
||||
if [[ "$code" == "200" || "$code" == "401" ]]; then
|
||||
log "✓ 回滚成功，旧版本已恢复服务 (status=$code)"
|
||||
log " 新版本部署失败，请检查构建/配置后重试"
|
||||
return 0
|
||||
else
|
||||
err "回滚后服务仍未恢复 (status=$code)"
|
||||
err "服务可能已完全宕机，请手动介入"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# ---------- 打印部署摘要 ----------
|
||||
print_summary() {
|
||||
log "═════════════════════════════════════════════"
|
||||
log "部署完成"
|
||||
log " 二进制: $TARGET_BIN"
|
||||
log " 适配器: $TARGET_ADAPTERS"
|
||||
log "═════════════════════════════════════════════"
|
||||
}
|
||||
|
||||
# ---------- 主流程 ----------
|
||||
main() {
|
||||
log "═══ llmsproxy 原子部署开始 ═══"
|
||||
precheck
|
||||
build
|
||||
sync_adapters
|
||||
backup_old_binary
|
||||
atomic_replace_binary
|
||||
restart_service
|
||||
if ! healthcheck; then
|
||||
rollback || true
|
||||
exit 1
|
||||
fi
|
||||
print_summary
|
||||
}
|
||||
|
||||
main "$@"
|
||||
@ -62,6 +62,7 @@ type Source struct {
|
||||
Timeout time.Duration `yaml:"timeout" json:"-"`
|
||||
MaxConcurrent int `yaml:"max_concurrent" json:"max_concurrent"`
|
||||
QueueTimeout time.Duration `yaml:"queue_timeout" json:"-"`
|
||||
RPM int `yaml:"rpm" json:"rpm,omitempty"` // optional requests-per-minute cap (0 = unlimited)
|
||||
}
|
||||
|
||||
// Load reads and validates a config file. When the file does not exist yet a
|
||||
|
||||
@ -72,6 +72,7 @@ type sourcePayload struct {
|
||||
Temperature float64 `json:"temperature"`
|
||||
MaxTokens int `json:"max_tokens"`
|
||||
MaxConcurrent int `json:"max_concurrent"`
|
||||
RPM int `json:"rpm"` // optional requests-per-minute cap, 0 = unlimited
|
||||
}
|
||||
|
||||
func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
@ -105,6 +106,7 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
Temperature: p.Temperature,
|
||||
MaxTokens: p.MaxTokens,
|
||||
MaxConcurrent: p.MaxConcurrent,
|
||||
RPM: p.RPM,
|
||||
}
|
||||
if err := g.core.AddSource(src); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "source_error", err.Error())
|
||||
|
||||
@ -847,6 +847,7 @@
|
||||
mEp: "聊天端点",
|
||||
mImgEp: "生图端点",
|
||||
mConc: "并发上限",
|
||||
mRPM: "RPM 限速 (0=不限)",
|
||||
mTemp: "温度",
|
||||
mModels: "模型列表",
|
||||
mAddModel: "+ 模型",
|
||||
@ -1067,6 +1068,7 @@
|
||||
mEp: "Chat endpoint",
|
||||
mImgEp: "Image endpoint",
|
||||
mConc: "Max concurrency",
|
||||
mRPM: "RPM limit (0 = unlimited)",
|
||||
mTemp: "Temperature",
|
||||
mModels: "Models",
|
||||
mAddModel: "+ model",
|
||||
@ -2479,6 +2481,10 @@
|
||||
<div><label>${t("mConc")}</label><input id="s-conc" type="number" value="${s.max_concurrent || 8}"></div>
|
||||
<div><label>${t("mTemp")}</label><input id="s-temp" type="number" step="0.1" value="${s.temperature ?? 0.7}"></div>
|
||||
</div>
|
||||
<div class="row">
|
||||
<div><label>${t("mRPM")}</label><input id="s-rpm" type="number" min="0" placeholder="0 = 不限" value="${s.rpm || 0}"></div>
|
||||
<div></div>
|
||||
</div>
|
||||
<label>${t("mModels")}</label>
|
||||
<div id="s-models"></div>
|
||||
<button class="ghost small" onclick="addModelRow()">${t("mAddModel")}</button>
|
||||
@ -2537,6 +2543,7 @@
|
||||
endpoint: $("#s-ep").value.trim(),
|
||||
image_endpoint: $("#s-img").value.trim(),
|
||||
max_concurrent: parseInt($("#s-conc").value) || 8,
|
||||
rpm: parseInt($("#s-rpm").value) || 0,
|
||||
temperature: parseFloat($("#s-temp").value) || 0,
|
||||
models,
|
||||
meta,
|
||||
|
||||
@ -47,6 +47,13 @@ const (
|
||||
backoffCapN = 10
|
||||
)
|
||||
|
||||
// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit
|
||||
// rejection means "too fast", not "broken": unlike transport/5xx failures it
|
||||
// must NOT escalate exponentially to backoffCap, or a quota-rich source gets
|
||||
// locked out for up to 30 minutes while its allowance is untouched. A short
|
||||
// fixed window lets the slot re-enter the rotation quickly.
|
||||
const rateLimitCooldown = 30 * time.Second
|
||||
|
||||
func clampPref(a *atomic.Int64, lo, hi int64) {
|
||||
for {
|
||||
cur := a.Load()
|
||||
@ -100,6 +107,17 @@ func (s *ModelState) RecordFailure(auth bool) {
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
|
||||
// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the
|
||||
// exponential schedule, so a merely-throttled source recovers in seconds and
|
||||
// its remaining quota stays usable. The preference penalty still applies so
|
||||
// other slots are preferred while cooling.
|
||||
func (s *ModelState) RecordRateLimit() {
|
||||
s.failCount.Add(1)
|
||||
s.cooldownUntil.Store(time.Now().Add(rateLimitCooldown).Unix())
|
||||
s.pref.Add(-prefFailStep)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
|
||||
// RecordSuccess resets the failure counter and cooldown and bumps the
|
||||
// preference score by one.
|
||||
func (s *ModelState) RecordSuccess() {
|
||||
@ -139,6 +157,12 @@ type Provider struct {
|
||||
mu sync.Mutex
|
||||
sem chan struct{}
|
||||
states map[string]*ModelState // key = model id
|
||||
// proactive rate limiting: nextOK is the earliest unix-nano time the next
|
||||
// request may leave; throttle() spaces requests 60s/RPM apart. Zero when
|
||||
// cfg.RPM <= 0 (unlimited).
|
||||
rateMu sync.Mutex
|
||||
nextOK int64
|
||||
rpmGap time.Duration
|
||||
lastProbe struct {
|
||||
ok bool
|
||||
err string
|
||||
@ -178,6 +202,9 @@ func New(cfg config.Source, vm *lua.VM) *Provider {
|
||||
if cfg.MaxConcurrent <= 0 {
|
||||
p.sem = nil
|
||||
}
|
||||
if cfg.RPM > 0 {
|
||||
p.rpmGap = time.Minute / time.Duration(cfg.RPM)
|
||||
}
|
||||
for _, m := range cfg.Models {
|
||||
p.states[m.ID] = &ModelState{}
|
||||
}
|
||||
@ -491,15 +518,20 @@ func (p *Provider) RecordSuccess(model string) {
|
||||
|
||||
// ReportStatus records an upstream HTTP status for the given model and drives
|
||||
// the (source, model) backoff state. 401/403 → capped self-healing cooldown
|
||||
// with doubled penalty; 429 and 5xx → normal exponential backoff. Other codes
|
||||
// (400 client schema errors, 402 billing errors) are not penalized here —
|
||||
// they surface via the status page / audit instead.
|
||||
// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed
|
||||
// cooldown (rateLimitCooldown): a rate rejection means "too fast", not
|
||||
// "broken", so the slot must re-enter rotation quickly instead of escalating
|
||||
// to a 30-minute lockout that wastes remaining quota.
|
||||
func (p *Provider) ReportStatus(model string, code int) {
|
||||
if code == 401 || code == 403 {
|
||||
p.RecordFailure(model, code)
|
||||
return
|
||||
}
|
||||
if code >= 500 || code == 429 {
|
||||
if code == 429 {
|
||||
p.state(model).RecordRateLimit()
|
||||
return
|
||||
}
|
||||
if code >= 500 {
|
||||
p.RecordFailure(model, code)
|
||||
}
|
||||
}
|
||||
@ -586,6 +618,37 @@ func (p *Provider) Release() {
|
||||
<-p.sem
|
||||
}
|
||||
|
||||
// ---- proactive rate limiting ----
|
||||
|
||||
// Throttle blocks until this source's rate-limit window allows the next
|
||||
// request, spacing outbound requests at least rpmGap apart (60s/RPM). It is a
|
||||
// no-op when no rpm cap is configured. The wait happens AFTER the concurrency
|
||||
// slot is taken, so a queued request still counts against max_concurrent —
|
||||
// bounded by the caller's ctx (queue timeout / client disconnect).
|
||||
func (p *Provider) Throttle(ctx context.Context) error {
|
||||
if p.rpmGap <= 0 {
|
||||
return nil
|
||||
}
|
||||
p.rateMu.Lock()
|
||||
now := time.Now().UnixNano()
|
||||
wait := p.nextOK - now
|
||||
if wait > 0 {
|
||||
p.nextOK += int64(p.rpmGap) // reserve the next window for the follower
|
||||
p.rateMu.Unlock()
|
||||
t := time.NewTimer(time.Duration(wait))
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-t.C:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
p.nextOK = now + int64(p.rpmGap)
|
||||
p.rateMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---- request construction ----
|
||||
|
||||
func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
|
||||
@ -630,6 +693,9 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
|
||||
return nil, err
|
||||
}
|
||||
defer p.Release()
|
||||
if err := p.Throttle(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
model := p.ModelFor(req.Model)
|
||||
|
||||
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
|
||||
@ -684,6 +750,10 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
if err := p.TryAcquire(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := p.Throttle(ctx); err != nil {
|
||||
p.Release()
|
||||
return nil, err
|
||||
}
|
||||
model := p.ModelFor(req.Model)
|
||||
req.Stream = true
|
||||
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
|
||||
|
||||
@ -245,6 +245,77 @@ func TestModelStateCooldownAndRecovery(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelStateRateLimitShortCooldown(t *testing.T) {
|
||||
// 429 must NOT use the exponential backoff schedule: a rate-limited but
|
||||
// quota-rich source has to re-enter rotation after the short fixed window.
|
||||
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
|
||||
st := p.state("m")
|
||||
|
||||
// repeated 429s must stay at the fixed short cooldown, never escalate
|
||||
for i := 0; i < 15; i++ {
|
||||
p.ReportStatus("m", 429)
|
||||
}
|
||||
until := st.CooldownUntil()
|
||||
want := time.Now().Add(rateLimitCooldown).Unix()
|
||||
if until < want-2 || until > want+2 {
|
||||
t.Fatalf("429 cooldown = %d, want ~%d (fixed %v, not exponential)", until, want, rateLimitCooldown)
|
||||
}
|
||||
if st.FailCount() != 15 {
|
||||
t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount())
|
||||
}
|
||||
if st.Pref() != -15*int64(prefFailStep) && st.Pref() > int64(prefMin) {
|
||||
t.Fatalf("pref = %d", st.Pref())
|
||||
}
|
||||
if p.ModelAvailable("m") {
|
||||
t.Fatal("model must be cooling right after a 429")
|
||||
}
|
||||
}
|
||||
|
||||
func TestThrottleSpacingAndCancel(t *testing.T) {
|
||||
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
|
||||
p.cfg.RPM = 120 // gap = 500ms
|
||||
p.rpmGap = time.Minute / time.Duration(p.cfg.RPM)
|
||||
|
||||
start := time.Now()
|
||||
for i := 0; i < 3; i++ {
|
||||
if err := p.Throttle(context.Background()); err != nil {
|
||||
t.Fatalf("throttle %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
elapsed := time.Since(start)
|
||||
// first call passes immediately, the next two wait one gap each
|
||||
want := 2 * time.Minute / time.Duration(p.cfg.RPM)
|
||||
if elapsed < want-time.Duration(100*time.Millisecond) || elapsed > want+time.Second {
|
||||
t.Fatalf("3 throttled calls took %v, want ~%v", elapsed, want)
|
||||
}
|
||||
|
||||
// unlimited source: Throttle is a no-op
|
||||
p2 := newTestProvider(t, src("free", "http://127.0.0.1:1", "openai", "m"))
|
||||
if err := p2.Throttle(context.Background()); err != nil {
|
||||
t.Fatalf("unlimited throttle: %v", err)
|
||||
}
|
||||
|
||||
// cancelled context aborts a pending window reservation
|
||||
p3 := newTestProvider(t, src("slow", "http://127.0.0.1:1", "openai", "m"))
|
||||
p3.cfg.RPM = 6 // 10s gap
|
||||
p3.rpmGap = time.Minute / time.Duration(p3.cfg.RPM)
|
||||
// first call passes immediately (fresh provider) and reserves the next window
|
||||
if err := p3.Throttle(context.Background()); err != nil {
|
||||
t.Fatalf("first throttle: %v", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
// second call must wait ~10s for the reserved window; the short deadline aborts it
|
||||
start = time.Now()
|
||||
err := p3.Throttle(ctx)
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("cancelled throttle = %v, want DeadlineExceeded", err)
|
||||
}
|
||||
if time.Since(start) > time.Second {
|
||||
t.Fatalf("cancel took %v, want fast abort", time.Since(start))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryAcquire(t *testing.T) {
|
||||
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
|
||||
p.cfg.MaxConcurrent = 1
|
||||
|
||||
Reference in New Issue
Block a user