mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
fix(openai-ws): 对齐 OpenAI 新版上下文池容量计算
This commit is contained in:
@@ -1274,7 +1274,9 @@ type GatewayOpenAIWSConfig struct {
|
||||
MaxConnsPerAccount int `mapstructure:"max_conns_per_account"`
|
||||
MinIdlePerAccount int `mapstructure:"min_idle_per_account"`
|
||||
MaxIdlePerAccount int `mapstructure:"max_idle_per_account"`
|
||||
// DynamicMaxConnsByAccountConcurrencyEnabled: 是否按账号并发动态计算连接池上限
|
||||
// DynamicMaxConnsByAccountConcurrencyEnabled: 是否按账号并发动态计算连接池上限。
|
||||
// 旧版及 mode_router_v2 的 ctx_pool 共用此开关和类型系数;关闭后使用 max_conns_per_account。
|
||||
// mode_router_v2 下并发数 <= 0 的账号仍不可调度。
|
||||
DynamicMaxConnsByAccountConcurrencyEnabled bool `mapstructure:"dynamic_max_conns_by_account_concurrency_enabled"`
|
||||
// OAuthMaxConnsFactor: OAuth 账号连接池系数(effective=ceil(concurrency*factor))
|
||||
OAuthMaxConnsFactor float64 `mapstructure:"oauth_max_conns_factor"`
|
||||
|
||||
@@ -2156,14 +2156,8 @@ func (p *openAIWSConnPool) effectiveMaxConnsByAccount(account *Account) int {
|
||||
if hardCap <= 0 {
|
||||
return 0
|
||||
}
|
||||
if p.modeRouterV2Enabled() {
|
||||
if account == nil {
|
||||
return hardCap
|
||||
}
|
||||
if account.Concurrency <= 0 {
|
||||
return 0
|
||||
}
|
||||
return min(account.Concurrency, hardCap)
|
||||
if p.modeRouterV2Enabled() && account != nil && account.Concurrency <= 0 {
|
||||
return 0
|
||||
}
|
||||
if account == nil || !p.dynamicMaxConnsEnabled() {
|
||||
return hardCap
|
||||
|
||||
@@ -1164,21 +1164,95 @@ func TestOpenAIWSConnPool_EffectiveMaxConnsDisabledFallbackHardCap(t *testing.T)
|
||||
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(account), "关闭动态模式后应保持旧行为")
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2RespectsHardCap(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
|
||||
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 0.3
|
||||
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 0.6
|
||||
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
accountType string
|
||||
concurrency int
|
||||
factor float64
|
||||
dynamic bool
|
||||
want int
|
||||
}{
|
||||
{name: "oauth expansion", accountType: AccountTypeOAuth, concurrency: 1, factor: 5, dynamic: true, want: 5},
|
||||
{name: "apikey expansion", accountType: AccountTypeAPIKey, concurrency: 1, factor: 5, dynamic: true, want: 5},
|
||||
{name: "oauth fraction", accountType: AccountTypeOAuth, concurrency: 20, factor: 0.3, dynamic: true, want: 6},
|
||||
{name: "apikey rounding", accountType: AccountTypeAPIKey, concurrency: 3, factor: 0.6, dynamic: true, want: 2},
|
||||
{name: "minimum one", accountType: AccountTypeOAuth, concurrency: 1, factor: 0.3, dynamic: true, want: 1},
|
||||
{name: "hard cap", accountType: AccountTypeOAuth, concurrency: 3, factor: 5, dynamic: true, want: 8},
|
||||
{name: "dynamic disabled", accountType: AccountTypeAPIKey, concurrency: 2, factor: 0.6, want: 8},
|
||||
{name: "zero concurrency", accountType: AccountTypeOAuth, factor: 5, dynamic: true, want: 0},
|
||||
{name: "negative concurrency", accountType: AccountTypeAPIKey, concurrency: -1, factor: 5, dynamic: true, want: 0},
|
||||
{name: "zero concurrency dynamic disabled", accountType: AccountTypeOAuth, factor: 5, want: 0},
|
||||
{name: "negative concurrency dynamic disabled", accountType: AccountTypeAPIKey, concurrency: -1, factor: 5, want: 0},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
|
||||
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = tt.dynamic
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 1
|
||||
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 1
|
||||
if tt.accountType == AccountTypeOAuth {
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = tt.factor
|
||||
} else {
|
||||
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = tt.factor
|
||||
}
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
t.Cleanup(pool.Close)
|
||||
account := &Account{Platform: PlatformOpenAI, Type: tt.accountType, Concurrency: tt.concurrency}
|
||||
require.Equal(t, tt.want, pool.effectiveMaxConnsByAccount(account))
|
||||
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(nil))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
func TestOpenAIWSConnPool_AcquireRetainedSessionsUsesScaledCapacity(t *testing.T) {
|
||||
for _, mode := range []struct {
|
||||
name string
|
||||
v2 bool
|
||||
}{
|
||||
{name: "legacy"},
|
||||
{name: "mode router v2", v2: true},
|
||||
} {
|
||||
t.Run(mode.name, func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = mode.v2
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 3
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 3
|
||||
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 1
|
||||
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
|
||||
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 2
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
pool.setClientDialerForTest(&openAIWSFakeDialer{})
|
||||
t.Cleanup(pool.Close)
|
||||
req := openAIWSAcquireRequest{
|
||||
Account: &Account{ID: 902, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 2},
|
||||
WSURL: "wss://example.com/v1/responses",
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
leases := make([]*openAIWSConnLease, 0, 3)
|
||||
for range 3 {
|
||||
lease, err := pool.Acquire(ctx, req)
|
||||
require.NoError(t, err, "轮次之间仍持有连接的会话应可使用系数扩出的容量")
|
||||
t.Cleanup(lease.Release)
|
||||
leases = append(leases, lease)
|
||||
}
|
||||
|
||||
high := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 20}
|
||||
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(high), "v2 路径也必须受连接池硬上限约束")
|
||||
fullCtx, fullCancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer fullCancel()
|
||||
_, err := pool.Acquire(fullCtx, req)
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded, "扩容后仍必须受全局硬上限约束")
|
||||
|
||||
nonPositive := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 0}
|
||||
require.Equal(t, 0, pool.effectiveMaxConnsByAccount(nonPositive), "并发数<=0 时应不可调度")
|
||||
leases[0].Release()
|
||||
reused, err := pool.Acquire(ctx, req)
|
||||
require.NoError(t, err)
|
||||
defer reused.Release()
|
||||
require.Equal(t, leases[0].ConnID(), reused.ConnID())
|
||||
require.True(t, reused.Reused())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_AcquireRejectsWhenEffectiveMaxConnsIsZero(t *testing.T) {
|
||||
|
||||
@@ -357,6 +357,8 @@ gateway:
|
||||
max_idle_per_account: 12
|
||||
# 是否按账号并发动态计算连接池上限:
|
||||
# effective_max_conns = min(max_conns_per_account, ceil(account.concurrency * factor))
|
||||
# 旧版及 mode_router_v2 的 ctx_pool 共用此开关和下方类型系数;关闭后使用 max_conns_per_account。
|
||||
# mode_router_v2 下并发数 <= 0 的账号仍不可调度。
|
||||
dynamic_max_conns_by_account_concurrency_enabled: true
|
||||
# 按账号类型分别设置系数(OAuth / API Key)
|
||||
oauth_max_conns_factor: 1.0
|
||||
|
||||
@@ -613,7 +613,7 @@ export default {
|
||||
wsModeShared: 'Shared (shared)',
|
||||
wsModeDedicated: 'Dedicated (dedicated)',
|
||||
wsModeConcurrencyHint:
|
||||
'When WS mode is enabled, account concurrency becomes the WS connection pool limit for this account.',
|
||||
'Account concurrency limits requests in flight; the context pool connection limit is controlled by the gateway connection pool settings.',
|
||||
wsModePassthroughHint: 'Passthrough mode does not use the WS connection pool.',
|
||||
oauthResponsesWebsocketsV2: 'OAuth WebSocket Mode',
|
||||
oauthResponsesWebsocketsV2Desc:
|
||||
|
||||
@@ -698,7 +698,7 @@ export default {
|
||||
wsModeHttpBridge: 'HTTP 桥接(http_bridge)',
|
||||
wsModeShared: '共享(shared)',
|
||||
wsModeDedicated: '独享(dedicated)',
|
||||
wsModeConcurrencyHint: '启用 WS mode 后,该账号并发数将作为该账号 WS 连接池上限。',
|
||||
wsModeConcurrencyHint: '账号并发数限制同时执行的请求数;上下文池的连接上限由网关连接池配置决定。',
|
||||
wsModePassthroughHint: 'passthrough 模式不使用 WS 连接池。',
|
||||
oauthResponsesWebsocketsV2: 'OAuth WebSocket Mode',
|
||||
oauthResponsesWebsocketsV2Desc:
|
||||
|
||||
Reference in New Issue
Block a user