openai_account_scheduler.go 104 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021
  1. package service
  2. import (
  3. "container/heap"
  4. "context"
  5. "errors"
  6. "fmt"
  7. "hash/fnv"
  8. "log/slog"
  9. "math"
  10. "sort"
  11. "strconv"
  12. "strings"
  13. "sync"
  14. "sync/atomic"
  15. "time"
  16. "github.com/Wei-Shaw/sub2api/internal/config"
  17. "golang.org/x/sync/singleflight"
  18. )
  19. const (
  20. openAIAccountScheduleLayerPreviousResponse = "previous_response_id"
  21. openAIAccountScheduleLayerGuardianParent = "guardian_parent"
  22. openAIAccountScheduleLayerSessionSticky = "session_hash"
  23. openAIAccountScheduleLayerLoadBalance = "load_balance"
  24. openAIAdvancedSchedulerSettingKey = "openai_advanced_scheduler_enabled"
  25. )
  26. const (
  27. openAIAdvancedSchedulerSettingCacheTTL = 5 * time.Second
  28. openAIAdvancedSchedulerSettingDBTimeout = 2 * time.Second
  29. // ponytail: cap probes added when cost ordering expands configured Top-K;
  30. // use bulk acquisition if a measured workload needs a higher ceiling.
  31. openAIAccountSelectionProbeLimit = 64
  32. )
  33. const (
  34. openAIQuotaHeadroomNeutralFactor = 0.5
  35. openAIQuotaHeadroomSecondaryLowRemain = 0.10
  36. openAIQuotaHeadroomSnapshotStaleAfter = 8 * time.Hour
  37. openAIUpstreamCostNeutralFactor = 0.5
  38. defaultOpenAIOAuthSchedulingRateMultiplier = 1.0
  39. )
  40. type cachedOpenAIAdvancedSchedulerSetting struct {
  41. lowUpstreamRatePriorityEnabled bool
  42. oauthSchedulingRateMultiplier float64
  43. enabled bool
  44. stickyWeightedEnabled bool
  45. subscriptionPriorityEnabled bool
  46. lbTopKOverride int
  47. weightOverrides map[string]float64
  48. expiresAt int64
  49. }
  50. type openAIAdvancedSchedulerRuntimeSettings struct {
  51. lowUpstreamRatePriorityEnabled bool
  52. oauthSchedulingRateMultiplier float64
  53. enabled bool
  54. stickyWeightedEnabled bool
  55. subscriptionPriorityEnabled bool
  56. lbTopKOverride int
  57. weightOverrides map[string]float64
  58. }
  59. var openAIAdvancedSchedulerSettingCache atomic.Value // *cachedOpenAIAdvancedSchedulerSetting
  60. var openAIAdvancedSchedulerSettingSF singleflight.Group
  61. type OpenAIAccountScheduleRequest struct {
  62. GroupID *int64
  63. Platform string
  64. SessionHash string
  65. StickyAccountID int64
  66. GuardianParentAccountID int64
  67. StickyPreviousAccountID int64
  68. StickyWeighted bool
  69. SubscriptionPriority bool
  70. PreserveStickyBinding bool
  71. RequirePrivacySet bool
  72. PreviousResponseID string
  73. PreviousResponseCanMove bool
  74. UseUpstreamTokenCost bool
  75. RequestedModel string
  76. RequiredTransport OpenAIUpstreamTransport
  77. RequiredCapability OpenAIEndpointCapability
  78. RequiredImageCapability OpenAIImagesCapability
  79. // RequireCompact is only for legacy /responses/compact capability filtering
  80. // and compact_model_mapping; native remote compaction v2 leaves it false.
  81. RequireCompact bool
  82. ExcludedIDs map[int64]struct{}
  83. }
  84. type OpenAIAccountScheduleDecision struct {
  85. Layer string
  86. StickyPreviousHit bool
  87. StickySessionHit bool
  88. CandidateCount int
  89. TopK int
  90. LatencyMs int64
  91. LoadSkew float64
  92. SelectedAccountID int64
  93. SelectedAccountType string
  94. }
  95. type OpenAIAccountSchedulerMetricsSnapshot struct {
  96. SelectTotal int64
  97. StickyPreviousHitTotal int64
  98. StickySessionHitTotal int64
  99. LoadBalanceSelectTotal int64
  100. AccountSwitchTotal int64
  101. SchedulerLatencyMsTotal int64
  102. SchedulerLatencyMsAvg float64
  103. StickyHitRatio float64
  104. AccountSwitchRate float64
  105. LoadSkewAvg float64
  106. RuntimeStatsAccountCount int
  107. }
  108. type OpenAIAccountScheduler interface {
  109. Select(ctx context.Context, req OpenAIAccountScheduleRequest) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error)
  110. ReportResult(accountID int64, success bool, firstTokenMs *int)
  111. ReportSwitch()
  112. SnapshotMetrics() OpenAIAccountSchedulerMetricsSnapshot
  113. }
  114. type openAIAccountSchedulerMetrics struct {
  115. selectTotal atomic.Int64
  116. stickyPreviousHitTotal atomic.Int64
  117. stickySessionHitTotal atomic.Int64
  118. loadBalanceSelectTotal atomic.Int64
  119. accountSwitchTotal atomic.Int64
  120. latencyMsTotal atomic.Int64
  121. loadSkewMilliTotal atomic.Int64
  122. }
  123. type openAIAccountLoadPlan struct {
  124. allCandidates []openAIAccountCandidateScore
  125. candidates []openAIAccountCandidateScore
  126. staleSnapshotCompactRetry []openAIAccountCandidateScore
  127. selectionOrder []openAIAccountCandidateScore
  128. candidateCount int
  129. topK int
  130. loadSkew float64
  131. includeOverflowFallback bool
  132. }
  133. type openAIAccountLoadSelectionAttempt struct {
  134. result *AccountSelectionResult
  135. selectionOrder []openAIAccountCandidateScore
  136. candidateCount int
  137. topK int
  138. loadSkew float64
  139. compactBlocked bool
  140. noCompactCandidates bool
  141. err error
  142. }
  143. func (m *openAIAccountSchedulerMetrics) recordSelect(decision OpenAIAccountScheduleDecision) {
  144. if m == nil {
  145. return
  146. }
  147. m.selectTotal.Add(1)
  148. m.latencyMsTotal.Add(decision.LatencyMs)
  149. m.loadSkewMilliTotal.Add(int64(math.Round(decision.LoadSkew * 1000)))
  150. if decision.StickyPreviousHit {
  151. m.stickyPreviousHitTotal.Add(1)
  152. }
  153. if decision.StickySessionHit {
  154. m.stickySessionHitTotal.Add(1)
  155. }
  156. if decision.Layer == openAIAccountScheduleLayerLoadBalance {
  157. m.loadBalanceSelectTotal.Add(1)
  158. }
  159. }
  160. func (m *openAIAccountSchedulerMetrics) recordSwitch() {
  161. if m == nil {
  162. return
  163. }
  164. m.accountSwitchTotal.Add(1)
  165. }
  166. type openAIAccountRuntimeStats struct {
  167. accounts sync.Map
  168. accountCount atomic.Int64
  169. }
  170. type openAIAccountRuntimeStat struct {
  171. errorRateEWMABits atomic.Uint64
  172. ttftEWMABits atomic.Uint64
  173. }
  174. func newOpenAIAccountRuntimeStats() *openAIAccountRuntimeStats {
  175. return &openAIAccountRuntimeStats{}
  176. }
  177. func (s *openAIAccountRuntimeStats) loadOrCreate(accountID int64) *openAIAccountRuntimeStat {
  178. if value, ok := s.accounts.Load(accountID); ok {
  179. stat, _ := value.(*openAIAccountRuntimeStat)
  180. if stat != nil {
  181. return stat
  182. }
  183. }
  184. stat := &openAIAccountRuntimeStat{}
  185. stat.ttftEWMABits.Store(math.Float64bits(math.NaN()))
  186. actual, loaded := s.accounts.LoadOrStore(accountID, stat)
  187. if !loaded {
  188. s.accountCount.Add(1)
  189. return stat
  190. }
  191. existing, _ := actual.(*openAIAccountRuntimeStat)
  192. if existing != nil {
  193. return existing
  194. }
  195. return stat
  196. }
  197. func updateEWMAAtomic(target *atomic.Uint64, sample float64, alpha float64) {
  198. for {
  199. oldBits := target.Load()
  200. oldValue := math.Float64frombits(oldBits)
  201. newValue := alpha*sample + (1-alpha)*oldValue
  202. if target.CompareAndSwap(oldBits, math.Float64bits(newValue)) {
  203. return
  204. }
  205. }
  206. }
  207. func (s *openAIAccountRuntimeStats) report(accountID int64, success bool, firstTokenMs *int) {
  208. if s == nil || accountID <= 0 {
  209. return
  210. }
  211. const alpha = 0.2
  212. stat := s.loadOrCreate(accountID)
  213. errorSample := 1.0
  214. if success {
  215. errorSample = 0.0
  216. }
  217. updateEWMAAtomic(&stat.errorRateEWMABits, errorSample, alpha)
  218. if firstTokenMs != nil && *firstTokenMs > 0 {
  219. ttft := float64(*firstTokenMs)
  220. ttftBits := math.Float64bits(ttft)
  221. for {
  222. oldBits := stat.ttftEWMABits.Load()
  223. oldValue := math.Float64frombits(oldBits)
  224. if math.IsNaN(oldValue) {
  225. if stat.ttftEWMABits.CompareAndSwap(oldBits, ttftBits) {
  226. break
  227. }
  228. continue
  229. }
  230. newValue := alpha*ttft + (1-alpha)*oldValue
  231. if stat.ttftEWMABits.CompareAndSwap(oldBits, math.Float64bits(newValue)) {
  232. break
  233. }
  234. }
  235. }
  236. }
  237. func (s *openAIAccountRuntimeStats) snapshot(accountID int64) (errorRate float64, ttft float64, hasTTFT bool) {
  238. if s == nil || accountID <= 0 {
  239. return 0, 0, false
  240. }
  241. value, ok := s.accounts.Load(accountID)
  242. if !ok {
  243. return 0, 0, false
  244. }
  245. stat, _ := value.(*openAIAccountRuntimeStat)
  246. if stat == nil {
  247. return 0, 0, false
  248. }
  249. errorRate = clamp01(math.Float64frombits(stat.errorRateEWMABits.Load()))
  250. ttftValue := math.Float64frombits(stat.ttftEWMABits.Load())
  251. if math.IsNaN(ttftValue) {
  252. return errorRate, 0, false
  253. }
  254. return errorRate, ttftValue, true
  255. }
  256. func (s *openAIAccountRuntimeStats) size() int {
  257. if s == nil {
  258. return 0
  259. }
  260. return int(s.accountCount.Load())
  261. }
  262. type defaultOpenAIAccountScheduler struct {
  263. service *OpenAIGatewayService
  264. metrics openAIAccountSchedulerMetrics
  265. stats *openAIAccountRuntimeStats
  266. grokFreeQuotaGateCache sync.Map // key: int64(accountID), value: grokFreeQuotaGateCacheEntry
  267. }
  268. type openAISelectionProbeBudget struct {
  269. acquires int
  270. rechecks int
  271. attempted map[int64]struct{}
  272. limited bool
  273. }
  274. func newOpenAISelectionProbeBudget() *openAISelectionProbeBudget {
  275. return &openAISelectionProbeBudget{attempted: make(map[int64]struct{})}
  276. }
  277. func (b *openAISelectionProbeBudget) enableLimit() {
  278. if b != nil {
  279. b.limited = true
  280. }
  281. }
  282. func (b *openAISelectionProbeBudget) recordAcquire(accountID int64) bool {
  283. if b == nil {
  284. return false
  285. }
  286. if !b.limited {
  287. return true
  288. }
  289. if b.acquires >= openAIAccountSelectionProbeLimit {
  290. return false
  291. }
  292. if b.attempted == nil {
  293. b.attempted = make(map[int64]struct{})
  294. }
  295. b.acquires++
  296. b.attempted[accountID] = struct{}{}
  297. return true
  298. }
  299. func (b *openAISelectionProbeBudget) recordRecheck() bool {
  300. if b == nil {
  301. return false
  302. }
  303. if !b.limited {
  304. return true
  305. }
  306. if b.rechecks >= openAIAccountSelectionProbeLimit {
  307. return false
  308. }
  309. b.rechecks++
  310. return true
  311. }
  312. func (b *openAISelectionProbeBudget) acquireExhausted() bool {
  313. return b != nil && b.limited && b.acquires >= openAIAccountSelectionProbeLimit
  314. }
  315. func (b *openAISelectionProbeBudget) wasAttempted(accountID int64) bool {
  316. if b == nil {
  317. return false
  318. }
  319. _, ok := b.attempted[accountID]
  320. return ok
  321. }
  322. type openAIStickyEscapeConfig struct {
  323. enabled bool
  324. ttftMs float64
  325. errorRate float64
  326. }
  327. func newDefaultOpenAIAccountScheduler(service *OpenAIGatewayService, stats *openAIAccountRuntimeStats) OpenAIAccountScheduler {
  328. if stats == nil {
  329. stats = newOpenAIAccountRuntimeStats()
  330. }
  331. return &defaultOpenAIAccountScheduler{
  332. service: service,
  333. stats: stats,
  334. }
  335. }
  336. func (s *defaultOpenAIAccountScheduler) Select(
  337. ctx context.Context,
  338. req OpenAIAccountScheduleRequest,
  339. ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
  340. if s != nil && s.service != nil && s.service.openAIGroupRequiresPrivacySet(ctx, req.GroupID) {
  341. req.RequirePrivacySet = true
  342. }
  343. decision := OpenAIAccountScheduleDecision{}
  344. start := time.Now()
  345. defer func() {
  346. decision.LatencyMs = time.Since(start).Milliseconds()
  347. s.metrics.recordSelect(decision)
  348. }()
  349. previousResponseID := strings.TrimSpace(req.PreviousResponseID)
  350. if previousResponseID != "" && NormalizeOpenAICompatiblePlatform(req.Platform) == PlatformOpenAI &&
  351. (!req.StickyWeighted || !req.PreviousResponseCanMove) {
  352. selection, err := s.service.selectAccountByPreviousResponseIDForCapability(
  353. ctx,
  354. req.GroupID,
  355. previousResponseID,
  356. req.RequestedModel,
  357. req.ExcludedIDs,
  358. req.RequiredCapability,
  359. req.RequireCompact,
  360. )
  361. if err != nil {
  362. return nil, decision, err
  363. }
  364. if selection != nil && selection.Account != nil {
  365. compatible, _ := s.isAccountRequestCompatibleReason(ctx, selection.Account, req)
  366. hasGroupMetadata := len(selection.Account.GroupIDs) > 0 || len(selection.Account.AccountGroups) > 0
  367. groupCompatible := !hasGroupMetadata || openAIStickyAccountMatchesGroup(selection.Account, req.GroupID)
  368. if hasGroupMetadata && s.service != nil {
  369. groupCompatible = s.service.openAIAccountMatchesSchedulingGroup(selection.Account, req.GroupID)
  370. }
  371. if !groupCompatible ||
  372. !compatible || !s.isAccountTransportCompatible(selection.Account, req.RequiredTransport) {
  373. if selection.ReleaseFunc != nil {
  374. selection.ReleaseFunc()
  375. }
  376. selection = nil
  377. }
  378. }
  379. if selection != nil && selection.Account != nil {
  380. decision.Layer = openAIAccountScheduleLayerPreviousResponse
  381. decision.StickyPreviousHit = true
  382. decision.SelectedAccountID = selection.Account.ID
  383. decision.SelectedAccountType = selection.Account.Type
  384. if req.SessionHash != "" {
  385. _ = s.service.bindOpenAIStickySessionDuringSelection(ctx, req.GroupID, req.SessionHash, selection.Account.ID)
  386. }
  387. return selection, decision, nil
  388. }
  389. }
  390. if req.GuardianParentAccountID > 0 {
  391. parentReq := req
  392. parentReq.StickyAccountID = req.GuardianParentAccountID
  393. parentReq.PreserveStickyBinding = true
  394. selection, _, err := s.selectBySessionHash(ctx, parentReq)
  395. if err != nil {
  396. return nil, decision, err
  397. }
  398. if selection != nil && selection.Account != nil {
  399. decision.Layer = openAIAccountScheduleLayerGuardianParent
  400. decision.StickySessionHit = true
  401. decision.SelectedAccountID = selection.Account.ID
  402. decision.SelectedAccountType = selection.Account.Type
  403. return selection, decision, nil
  404. }
  405. }
  406. if !req.StickyWeighted {
  407. selection, escapedSticky, err := s.selectBySessionHash(ctx, req)
  408. if err != nil {
  409. return nil, decision, err
  410. }
  411. if selection != nil && selection.Account != nil {
  412. decision.Layer = openAIAccountScheduleLayerSessionSticky
  413. decision.StickySessionHit = true
  414. decision.SelectedAccountID = selection.Account.ID
  415. decision.SelectedAccountType = selection.Account.Type
  416. return selection, decision, nil
  417. }
  418. if escapedSticky {
  419. req.PreserveStickyBinding = true
  420. }
  421. }
  422. selection, candidateCount, topK, loadSkew, err := s.selectByLoadBalance(ctx, req)
  423. decision.Layer = openAIAccountScheduleLayerLoadBalance
  424. decision.CandidateCount = candidateCount
  425. decision.TopK = topK
  426. decision.LoadSkew = loadSkew
  427. if err != nil {
  428. return nil, decision, err
  429. }
  430. if selection != nil && selection.Account != nil {
  431. decision.SelectedAccountID = selection.Account.ID
  432. decision.SelectedAccountType = selection.Account.Type
  433. if req.StickyWeighted {
  434. if req.StickyPreviousAccountID > 0 && selection.Account.ID == req.StickyPreviousAccountID {
  435. decision.StickyPreviousHit = true
  436. }
  437. if req.StickyAccountID > 0 && selection.Account.ID == req.StickyAccountID {
  438. decision.StickySessionHit = true
  439. }
  440. }
  441. }
  442. return selection, decision, nil
  443. }
  444. func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
  445. ctx context.Context,
  446. req OpenAIAccountScheduleRequest,
  447. ) (*AccountSelectionResult, bool, error) {
  448. sessionHash := strings.TrimSpace(req.SessionHash)
  449. if sessionHash == "" || s == nil || s.service == nil || s.service.cache == nil {
  450. return nil, false, nil
  451. }
  452. accountID := req.StickyAccountID
  453. clearBinding := func() {
  454. if !req.PreserveStickyBinding {
  455. _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
  456. }
  457. }
  458. if accountID <= 0 {
  459. var err error
  460. accountID, err = s.service.getStickySessionAccountID(ctx, req.GroupID, sessionHash)
  461. if err != nil || accountID <= 0 {
  462. return nil, false, nil
  463. }
  464. }
  465. if accountID <= 0 {
  466. return nil, false, nil
  467. }
  468. if req.ExcludedIDs != nil {
  469. if _, excluded := req.ExcludedIDs[accountID]; excluded {
  470. return nil, false, nil
  471. }
  472. }
  473. account, err := s.service.getSchedulableAccount(ctx, accountID)
  474. if err != nil || account == nil {
  475. clearBinding()
  476. return nil, false, nil
  477. }
  478. if shouldClearStickySession(account, req.RequestedModel) || account.Platform != NormalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() || !account.IsSchedulable() {
  479. clearBinding()
  480. return nil, false, nil
  481. }
  482. if !s.isAccountRequestCompatible(ctx, account, req) {
  483. return nil, false, nil
  484. }
  485. if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
  486. clearBinding()
  487. return nil, false, nil
  488. }
  489. account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.GroupID, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability)
  490. if account == nil || !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) || !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
  491. clearBinding()
  492. return nil, false, nil
  493. }
  494. // Free-tier soft gate: sticky session must not pin an over-quota free OAuth account.
  495. // Admin QueryQuota / import probes do not use this path.
  496. if account != nil && len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 {
  497. clearBinding()
  498. return nil, false, nil
  499. }
  500. // Team+model cool: sticky must not pin a sibling under the same team 429 window.
  501. now := time.Now()
  502. upstreamModel := canonicalOpenAIAccountSchedulingModel(account, req.RequestedModel)
  503. if account != nil && isGrokTeamModelRateLimited(account, upstreamModel, now) {
  504. clearBinding()
  505. return nil, false, nil
  506. }
  507. if account != nil && isGrokModelQuotaBlocked(account.ID, upstreamModel, now) {
  508. clearBinding()
  509. return nil, false, nil
  510. }
  511. escapeCfg := s.service.openAIStickyEscapeConfig()
  512. if reason, errorRate, ttft, shouldEscape := s.shouldEscapeStickyAccount(accountID, escapeCfg); shouldEscape {
  513. slog.Info("sticky_escape_triggered",
  514. "account_id", accountID,
  515. "reason", reason,
  516. "error_rate", errorRate,
  517. "ttft", ttft,
  518. )
  519. return nil, true, nil
  520. }
  521. result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Mode1EffectiveConcurrency())
  522. if acquireErr == nil && result != nil && result.Acquired {
  523. if !req.PreserveStickyBinding {
  524. _ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL())
  525. }
  526. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  527. Account: account,
  528. Acquired: true,
  529. ReleaseFunc: result.ReleaseFunc,
  530. }), false, nil
  531. }
  532. cfg := s.service.schedulingConfig()
  533. // WaitPlan.MaxConcurrency 使用 Concurrency(非 EffectiveLoadFactor),因为 WaitPlan 控制的是 Redis 实际并发槽位等待。
  534. if s.service.concurrencyService != nil {
  535. if escapeCfg.enabled && acquireErr == nil && result != nil && !result.Acquired {
  536. errorRate, ttft, _ := s.stats.snapshot(accountID)
  537. slog.Info("sticky_escape_triggered",
  538. "account_id", accountID,
  539. "reason", "concurrency_full",
  540. "error_rate", errorRate,
  541. "ttft", ttft,
  542. )
  543. return nil, true, nil
  544. }
  545. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  546. Account: account,
  547. WaitPlan: &AccountWaitPlan{
  548. AccountID: accountID,
  549. MaxConcurrency: account.Mode1EffectiveConcurrency(),
  550. Timeout: cfg.StickySessionWaitTimeout,
  551. MaxWaiting: cfg.StickySessionMaxWaiting,
  552. },
  553. }), false, nil
  554. }
  555. return nil, false, nil
  556. }
  557. func openAIStickyAccountMatchesGroup(account *Account, groupID *int64) bool {
  558. if account == nil {
  559. return false
  560. }
  561. if groupID == nil {
  562. return len(account.AccountGroups) == 0 && len(account.GroupIDs) == 0
  563. }
  564. for _, accountGroupID := range account.GroupIDs {
  565. if accountGroupID == *groupID {
  566. return true
  567. }
  568. }
  569. for _, accountGroup := range account.AccountGroups {
  570. if accountGroup.GroupID == *groupID {
  571. return true
  572. }
  573. }
  574. return false
  575. }
  576. func openAIAccountSchedulingPriority(account *Account) int {
  577. if account == nil {
  578. return 0
  579. }
  580. return account.Priority
  581. }
  582. func (s *defaultOpenAIAccountScheduler) shouldEscapeStickyAccount(accountID int64, cfg openAIStickyEscapeConfig) (reason string, errorRate float64, ttft float64, shouldEscape bool) {
  583. if !cfg.enabled || s == nil || s.stats == nil || accountID <= 0 {
  584. return "", 0, 0, false
  585. }
  586. errorRate, ttft, hasTTFT := s.stats.snapshot(accountID)
  587. if hasTTFT && ttft > cfg.ttftMs {
  588. return "ttft", errorRate, ttft, true
  589. }
  590. if errorRate > cfg.errorRate {
  591. return "error_rate", errorRate, ttft, true
  592. }
  593. return "", errorRate, ttft, false
  594. }
  595. type openAIAccountCandidateScore struct {
  596. account *Account
  597. loadInfo *AccountLoadInfo
  598. loadKnown bool
  599. score float64
  600. priority int
  601. errorRate float64
  602. ttft float64
  603. hasTTFT bool
  604. }
  605. type openAIAccountCandidateHeap []openAIAccountCandidateScore
  606. func (h openAIAccountCandidateHeap) Len() int {
  607. return len(h)
  608. }
  609. func (h openAIAccountCandidateHeap) Less(i, j int) bool {
  610. // 最小堆根节点保存“最差”候选,便于 O(log k) 维护 topK。
  611. return isOpenAIAccountCandidateBetter(h[j], h[i])
  612. }
  613. func (h openAIAccountCandidateHeap) Swap(i, j int) {
  614. h[i], h[j] = h[j], h[i]
  615. }
  616. func (h *openAIAccountCandidateHeap) Push(x any) {
  617. candidate, ok := x.(openAIAccountCandidateScore)
  618. if !ok {
  619. panic("openAIAccountCandidateHeap: invalid element type")
  620. }
  621. *h = append(*h, candidate)
  622. }
  623. func (h *openAIAccountCandidateHeap) Pop() any {
  624. old := *h
  625. n := len(old)
  626. last := old[n-1]
  627. *h = old[:n-1]
  628. return last
  629. }
  630. func isOpenAIAccountCandidateBetter(left openAIAccountCandidateScore, right openAIAccountCandidateScore) bool {
  631. if left.score != right.score {
  632. return left.score > right.score
  633. }
  634. if left.account.Priority != right.account.Priority {
  635. return left.account.Priority < right.account.Priority
  636. }
  637. if left.loadInfo.LoadRate != right.loadInfo.LoadRate {
  638. return left.loadInfo.LoadRate < right.loadInfo.LoadRate
  639. }
  640. if left.loadInfo.WaitingCount != right.loadInfo.WaitingCount {
  641. return left.loadInfo.WaitingCount < right.loadInfo.WaitingCount
  642. }
  643. return left.account.ID < right.account.ID
  644. }
  645. func selectTopKOpenAICandidates(candidates []openAIAccountCandidateScore, topK int) []openAIAccountCandidateScore {
  646. if len(candidates) == 0 {
  647. return nil
  648. }
  649. if topK <= 0 {
  650. topK = 1
  651. }
  652. if topK >= len(candidates) {
  653. ranked := append([]openAIAccountCandidateScore(nil), candidates...)
  654. sort.Slice(ranked, func(i, j int) bool {
  655. return isOpenAIAccountCandidateBetter(ranked[i], ranked[j])
  656. })
  657. return ranked
  658. }
  659. best := make(openAIAccountCandidateHeap, 0, topK)
  660. for _, candidate := range candidates {
  661. if len(best) < topK {
  662. heap.Push(&best, candidate)
  663. continue
  664. }
  665. if isOpenAIAccountCandidateBetter(candidate, best[0]) {
  666. best[0] = candidate
  667. heap.Fix(&best, 0)
  668. }
  669. }
  670. ranked := make([]openAIAccountCandidateScore, len(best))
  671. copy(ranked, best)
  672. sort.Slice(ranked, func(i, j int) bool {
  673. return isOpenAIAccountCandidateBetter(ranked[i], ranked[j])
  674. })
  675. return ranked
  676. }
  677. type openAISelectionRNG struct {
  678. state uint64
  679. }
  680. func newOpenAISelectionRNG(seed uint64) openAISelectionRNG {
  681. if seed == 0 {
  682. seed = 0x9e3779b97f4a7c15
  683. }
  684. return openAISelectionRNG{state: seed}
  685. }
  686. func (r *openAISelectionRNG) nextUint64() uint64 {
  687. // xorshift64*
  688. x := r.state
  689. x ^= x >> 12
  690. x ^= x << 25
  691. x ^= x >> 27
  692. r.state = x
  693. return x * 2685821657736338717
  694. }
  695. func (r *openAISelectionRNG) nextFloat64() float64 {
  696. // [0,1)
  697. return float64(r.nextUint64()>>11) / (1 << 53)
  698. }
  699. func deriveOpenAISelectionSeed(req OpenAIAccountScheduleRequest) uint64 {
  700. hasher := fnv.New64a()
  701. writeValue := func(value string) {
  702. trimmed := strings.TrimSpace(value)
  703. if trimmed == "" {
  704. return
  705. }
  706. _, _ = hasher.Write([]byte(trimmed))
  707. _, _ = hasher.Write([]byte{0})
  708. }
  709. writeValue(req.SessionHash)
  710. writeValue(req.PreviousResponseID)
  711. writeValue(req.RequestedModel)
  712. if req.GroupID != nil {
  713. _, _ = hasher.Write([]byte(strconv.FormatInt(*req.GroupID, 10)))
  714. }
  715. seed := hasher.Sum64()
  716. // 对“无会话锚点”的纯负载均衡请求引入时间熵,避免固定命中同一账号。
  717. if strings.TrimSpace(req.SessionHash) == "" && strings.TrimSpace(req.PreviousResponseID) == "" {
  718. seed ^= uint64(time.Now().UnixNano())
  719. }
  720. if seed == 0 {
  721. seed = uint64(time.Now().UnixNano()) ^ 0x9e3779b97f4a7c15
  722. }
  723. return seed
  724. }
  725. func buildOpenAIWeightedSelectionOrder(
  726. candidates []openAIAccountCandidateScore,
  727. req OpenAIAccountScheduleRequest,
  728. ) []openAIAccountCandidateScore {
  729. if len(candidates) <= 1 {
  730. return append([]openAIAccountCandidateScore(nil), candidates...)
  731. }
  732. pool := append([]openAIAccountCandidateScore(nil), candidates...)
  733. weights := make([]float64, len(pool))
  734. minScore := pool[0].score
  735. for i := 1; i < len(pool); i++ {
  736. if pool[i].score < minScore {
  737. minScore = pool[i].score
  738. }
  739. }
  740. for i := range pool {
  741. // 将 top-K 分值平移到正区间,避免“单一最高分账号”长期垄断。
  742. weight := (pool[i].score - minScore) + 1.0
  743. if math.IsNaN(weight) || math.IsInf(weight, 0) || weight <= 0 {
  744. weight = 1.0
  745. }
  746. weights[i] = weight
  747. }
  748. order := make([]openAIAccountCandidateScore, 0, len(pool))
  749. rng := newOpenAISelectionRNG(deriveOpenAISelectionSeed(req))
  750. for len(pool) > 0 {
  751. total := 0.0
  752. for _, w := range weights {
  753. total += w
  754. }
  755. selectedIdx := 0
  756. if total > 0 {
  757. r := rng.nextFloat64() * total
  758. acc := 0.0
  759. for i, w := range weights {
  760. acc += w
  761. if r <= acc {
  762. selectedIdx = i
  763. break
  764. }
  765. }
  766. } else {
  767. selectedIdx = int(rng.nextUint64() % uint64(len(pool)))
  768. }
  769. order = append(order, pool[selectedIdx])
  770. pool = append(pool[:selectedIdx], pool[selectedIdx+1:]...)
  771. weights = append(weights[:selectedIdx], weights[selectedIdx+1:]...)
  772. }
  773. return order
  774. }
  775. func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan(
  776. ctx context.Context,
  777. req OpenAIAccountScheduleRequest,
  778. filtered []*Account,
  779. loadMap map[int64]*AccountLoadInfo,
  780. ) openAIAccountLoadPlan {
  781. allCandidates := make([]openAIAccountCandidateScore, 0, len(filtered))
  782. for _, account := range filtered {
  783. loadInfo, loadKnown := loadMap[account.ID]
  784. if !loadKnown || loadInfo == nil {
  785. loadInfo = &AccountLoadInfo{AccountID: account.ID}
  786. loadKnown = false
  787. }
  788. errorRate, ttft, hasTTFT := 0.0, 0.0, false
  789. if s.stats != nil {
  790. errorRate, ttft, hasTTFT = s.stats.snapshot(account.ID)
  791. }
  792. allCandidates = append(allCandidates, openAIAccountCandidateScore{
  793. account: account,
  794. loadInfo: loadInfo,
  795. loadKnown: loadKnown,
  796. errorRate: errorRate,
  797. ttft: ttft,
  798. hasTTFT: hasTTFT,
  799. })
  800. }
  801. candidates := allCandidates
  802. staleSnapshotCompactRetry := make([]openAIAccountCandidateScore, 0, len(allCandidates))
  803. if req.RequireCompact {
  804. candidates = make([]openAIAccountCandidateScore, 0, len(allCandidates))
  805. for _, candidate := range allCandidates {
  806. if openAICompactSupportTier(candidate.account) == 0 {
  807. staleSnapshotCompactRetry = append(staleSnapshotCompactRetry, candidate)
  808. continue
  809. }
  810. candidates = append(candidates, candidate)
  811. }
  812. }
  813. plan := openAIAccountLoadPlan{
  814. allCandidates: allCandidates,
  815. candidates: candidates,
  816. staleSnapshotCompactRetry: staleSnapshotCompactRetry,
  817. candidateCount: len(candidates),
  818. }
  819. if len(candidates) == 0 {
  820. plan.selectionOrder = s.buildOpenAISelectionOrder(req, plan)
  821. return plan
  822. }
  823. minPriority, maxPriority := openAIAccountSchedulingPriority(candidates[0].account), openAIAccountSchedulingPriority(candidates[0].account)
  824. maxWaiting := 1
  825. loadRateSum := 0.0
  826. loadRateSumSquares := 0.0
  827. minTTFT, maxTTFT := 0.0, 0.0
  828. hasTTFTSample := false
  829. for i := range candidates {
  830. candidate := &candidates[i]
  831. candidate.priority = openAIAccountSchedulingPriority(candidate.account)
  832. if candidate.priority < minPriority {
  833. minPriority = candidate.priority
  834. }
  835. if candidate.priority > maxPriority {
  836. maxPriority = candidate.priority
  837. }
  838. if candidate.loadInfo.WaitingCount > maxWaiting {
  839. maxWaiting = candidate.loadInfo.WaitingCount
  840. }
  841. if candidate.hasTTFT && candidate.ttft > 0 {
  842. if !hasTTFTSample {
  843. minTTFT, maxTTFT = candidate.ttft, candidate.ttft
  844. hasTTFTSample = true
  845. } else {
  846. if candidate.ttft < minTTFT {
  847. minTTFT = candidate.ttft
  848. }
  849. if candidate.ttft > maxTTFT {
  850. maxTTFT = candidate.ttft
  851. }
  852. }
  853. }
  854. loadRate := float64(candidate.loadInfo.LoadRate)
  855. loadRateSum += loadRate
  856. loadRateSumSquares += loadRate * loadRate
  857. }
  858. plan.loadSkew = calcLoadSkewByMoments(loadRateSum, loadRateSumSquares, len(candidates))
  859. weights := s.service.openAIWSSchedulerWeightsForRequest(ctx)
  860. now := time.Now()
  861. upstreamCostFactors := map[int64]float64(nil)
  862. if req.UseUpstreamTokenCost && weights.UpstreamCost > 0 {
  863. accounts := make([]*Account, 0, len(candidates))
  864. for _, candidate := range candidates {
  865. accounts = append(accounts, candidate.account)
  866. }
  867. upstreamCostFactors = openAIUpstreamCostFactors(accounts, now, s.service.openAIOAuthSchedulingRateMultiplier(ctx))
  868. for _, factor := range upstreamCostFactors {
  869. if factor != openAIUpstreamCostNeutralFactor {
  870. plan.includeOverflowFallback = true
  871. break
  872. }
  873. }
  874. }
  875. // Reset 因子(use-it-or-lose-it):在拥有「未来会话窗口结束时间」的账号中,
  876. // 剩余时间越短 → 因子越接近 1(越早重置越优先用尽)。无活跃窗口的账号因子为 0。
  877. // 仅在 weights.Reset > 0 时计算,默认关闭不影响原有行为。
  878. minResetRemaining, maxResetRemaining := 0.0, 0.0
  879. hasResetSample := false
  880. if weights.Reset > 0 {
  881. for _, candidate := range candidates {
  882. end := candidate.account.SessionWindowEnd
  883. if end == nil || !now.Before(*end) {
  884. continue
  885. }
  886. remaining := end.Sub(now).Seconds()
  887. if !hasResetSample {
  888. minResetRemaining, maxResetRemaining = remaining, remaining
  889. hasResetSample = true
  890. continue
  891. }
  892. if remaining < minResetRemaining {
  893. minResetRemaining = remaining
  894. }
  895. if remaining > maxResetRemaining {
  896. maxResetRemaining = remaining
  897. }
  898. }
  899. }
  900. for i := range candidates {
  901. item := &candidates[i]
  902. priorityFactor := 1.0
  903. if maxPriority > minPriority {
  904. priorityFactor = 1 - float64(item.priority-minPriority)/float64(maxPriority-minPriority)
  905. }
  906. loadFactor := 1 - clamp01(float64(item.loadInfo.LoadRate)/100.0)
  907. queueFactor := 1 - clamp01(float64(item.loadInfo.WaitingCount)/float64(maxWaiting))
  908. errorFactor := 1 - clamp01(item.errorRate)
  909. ttftFactor := 0.5
  910. if item.hasTTFT && hasTTFTSample && maxTTFT > minTTFT {
  911. ttftFactor = 1 - clamp01((item.ttft-minTTFT)/(maxTTFT-minTTFT))
  912. }
  913. resetFactor := 0.0
  914. if weights.Reset > 0 && hasResetSample {
  915. if end := item.account.SessionWindowEnd; end != nil && now.Before(*end) {
  916. if maxResetRemaining > minResetRemaining {
  917. resetFactor = 1 - clamp01((end.Sub(now).Seconds()-minResetRemaining)/(maxResetRemaining-minResetRemaining))
  918. } else {
  919. // 所有有窗口的账号剩余时间相同:一律给满分,让其优于无窗口账号。
  920. resetFactor = 1
  921. }
  922. }
  923. }
  924. quotaHeadroomFactor := 0.0
  925. if weights.QuotaHeadroom > 0 {
  926. quotaHeadroomFactor = openAIQuotaHeadroomFactor(item.account, now)
  927. }
  928. upstreamCostFactor := openAIUpstreamCostNeutralFactor
  929. if factor, ok := upstreamCostFactors[item.account.ID]; ok {
  930. upstreamCostFactor = factor
  931. }
  932. item.score = weights.Priority*priorityFactor +
  933. weights.Load*loadFactor +
  934. weights.Queue*queueFactor +
  935. weights.ErrorRate*errorFactor +
  936. weights.TTFT*ttftFactor +
  937. weights.Reset*resetFactor +
  938. weights.QuotaHeadroom*quotaHeadroomFactor +
  939. weights.UpstreamCost*(upstreamCostFactor-openAIUpstreamCostNeutralFactor)
  940. if req.StickyWeighted {
  941. if req.PreviousResponseCanMove && req.StickyPreviousAccountID > 0 && item.account.ID == req.StickyPreviousAccountID {
  942. item.score += weights.Previous
  943. }
  944. if req.StickyAccountID > 0 && item.account.ID == req.StickyAccountID {
  945. item.score += weights.SessionSticky
  946. }
  947. }
  948. }
  949. plan.candidates = candidates
  950. plan.topK = s.service.openAIWSLBTopKForRequest(ctx)
  951. if plan.topK > len(candidates) {
  952. plan.topK = len(candidates)
  953. }
  954. if plan.topK <= 0 {
  955. plan.topK = 1
  956. }
  957. plan.selectionOrder = s.buildOpenAISelectionOrder(req, plan)
  958. return plan
  959. }
  960. func (s *defaultOpenAIAccountScheduler) buildOpenAISelectionOrder(
  961. req OpenAIAccountScheduleRequest,
  962. plan openAIAccountLoadPlan,
  963. ) []openAIAccountCandidateScore {
  964. buildSelectionOrder := func(pool []openAIAccountCandidateScore) []openAIAccountCandidateScore {
  965. if len(pool) == 0 || plan.topK <= 0 {
  966. return nil
  967. }
  968. groupTopK := plan.topK
  969. if groupTopK > len(pool) {
  970. groupTopK = len(pool)
  971. }
  972. ranked := selectTopKOpenAICandidates(pool, groupTopK)
  973. var primary []openAIAccountCandidateScore
  974. if req.StickyWeighted {
  975. for _, stickyID := range []int64{req.StickyPreviousAccountID, req.StickyAccountID} {
  976. if stickyID <= 0 {
  977. continue
  978. }
  979. for i, candidate := range ranked {
  980. if candidate.account != nil && candidate.account.ID == stickyID {
  981. primary = append([]openAIAccountCandidateScore{candidate}, ranked[:i]...)
  982. primary = append(primary, ranked[i+1:]...)
  983. break
  984. }
  985. }
  986. if len(primary) > 0 {
  987. break
  988. }
  989. }
  990. }
  991. if len(primary) == 0 {
  992. primary = buildOpenAIWeightedSelectionOrder(ranked, req)
  993. }
  994. if !plan.includeOverflowFallback || groupTopK >= len(pool) {
  995. return primary
  996. }
  997. selected := make(map[int64]struct{}, len(primary))
  998. for _, candidate := range primary {
  999. selected[candidate.account.ID] = struct{}{}
  1000. }
  1001. overflow := make([]openAIAccountCandidateScore, 0, len(pool)-len(primary))
  1002. for _, candidate := range pool {
  1003. if _, ok := selected[candidate.account.ID]; !ok {
  1004. overflow = append(overflow, candidate)
  1005. }
  1006. }
  1007. sort.Slice(overflow, func(i, j int) bool {
  1008. return isOpenAIAccountCandidateBetter(overflow[i], overflow[j])
  1009. })
  1010. return append(primary, overflow...)
  1011. }
  1012. if req.RequireCompact {
  1013. supported := make([]openAIAccountCandidateScore, 0, len(plan.candidates))
  1014. unknown := make([]openAIAccountCandidateScore, 0, len(plan.candidates))
  1015. for _, candidate := range plan.candidates {
  1016. switch openAICompactSupportTier(candidate.account) {
  1017. case 2:
  1018. supported = append(supported, candidate)
  1019. case 1:
  1020. unknown = append(unknown, candidate)
  1021. }
  1022. }
  1023. selectionOrder := make([]openAIAccountCandidateScore, 0, len(plan.allCandidates))
  1024. selectionOrder = append(selectionOrder, buildSelectionOrder(supported)...)
  1025. selectionOrder = append(selectionOrder, buildSelectionOrder(unknown)...)
  1026. if len(plan.staleSnapshotCompactRetry) > 0 && s.service.schedulerSnapshot != nil {
  1027. selectionOrder = append(selectionOrder, sortOpenAICompactRetryCandidates(plan.staleSnapshotCompactRetry)...)
  1028. }
  1029. return selectionOrder
  1030. }
  1031. return buildSelectionOrder(plan.candidates)
  1032. }
  1033. func sortOpenAICompactRetryCandidates(pool []openAIAccountCandidateScore) []openAIAccountCandidateScore {
  1034. if len(pool) == 0 {
  1035. return nil
  1036. }
  1037. ordered := append([]openAIAccountCandidateScore(nil), pool...)
  1038. sort.SliceStable(ordered, func(i, j int) bool {
  1039. a, b := ordered[i], ordered[j]
  1040. if a.account.Priority != b.account.Priority {
  1041. return a.account.Priority < b.account.Priority
  1042. }
  1043. if a.loadInfo.LoadRate != b.loadInfo.LoadRate {
  1044. return a.loadInfo.LoadRate < b.loadInfo.LoadRate
  1045. }
  1046. if a.loadInfo.WaitingCount != b.loadInfo.WaitingCount {
  1047. return a.loadInfo.WaitingCount < b.loadInfo.WaitingCount
  1048. }
  1049. switch {
  1050. case a.account.LastUsedAt == nil && b.account.LastUsedAt != nil:
  1051. return true
  1052. case a.account.LastUsedAt != nil && b.account.LastUsedAt == nil:
  1053. return false
  1054. case a.account.LastUsedAt == nil && b.account.LastUsedAt == nil:
  1055. return false
  1056. default:
  1057. return a.account.LastUsedAt.Before(*b.account.LastUsedAt)
  1058. }
  1059. })
  1060. return ordered
  1061. }
  1062. func (s *defaultOpenAIAccountScheduler) tryAcquireOpenAISelectionOrder(
  1063. ctx context.Context,
  1064. req OpenAIAccountScheduleRequest,
  1065. selectionOrder []openAIAccountCandidateScore,
  1066. ) (*AccountSelectionResult, bool, error) {
  1067. budget := newOpenAISelectionProbeBudget()
  1068. budget.enableLimit()
  1069. return s.tryAcquireOpenAISelectionOrderWithBudget(ctx, req, selectionOrder, budget)
  1070. }
  1071. func (s *defaultOpenAIAccountScheduler) tryAcquireOpenAISelectionOrderWithBudget(
  1072. ctx context.Context,
  1073. req OpenAIAccountScheduleRequest,
  1074. selectionOrder []openAIAccountCandidateScore,
  1075. budget *openAISelectionProbeBudget,
  1076. ) (*AccountSelectionResult, bool, error) {
  1077. compactBlocked := false
  1078. release := func(result *AcquireResult) {
  1079. if result != nil && result.ReleaseFunc != nil {
  1080. result.ReleaseFunc()
  1081. }
  1082. }
  1083. for i := 0; i < len(selectionOrder); i++ {
  1084. candidate := selectionOrder[i]
  1085. if candidate.account == nil {
  1086. continue
  1087. }
  1088. if candidate.loadKnown && candidate.account.Mode1EffectiveConcurrency() > 0 &&
  1089. candidate.loadInfo.CurrentConcurrency >= candidate.account.Mode1EffectiveConcurrency() {
  1090. continue
  1091. }
  1092. result, attempted, acquireErr := s.tryAcquireOpenAIAccountSlot(ctx, candidate.account.ID, candidate.account.Mode1EffectiveConcurrency(), budget)
  1093. if !attempted {
  1094. break
  1095. }
  1096. if acquireErr != nil {
  1097. return nil, compactBlocked, acquireErr
  1098. }
  1099. if result == nil || !result.Acquired {
  1100. continue
  1101. }
  1102. fresh := s.service.resolveFreshSchedulableOpenAIAccount(ctx, candidate.account, req.Platform, req.RequestedModel, false, req.RequiredCapability)
  1103. if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
  1104. release(result)
  1105. continue
  1106. }
  1107. if !s.consumeOpenAISelectionDBRecheck(budget) {
  1108. release(result)
  1109. break
  1110. }
  1111. fresh = s.service.recheckSelectedOpenAIAccountFromDB(ctx, fresh, req.GroupID, req.Platform, req.RequestedModel, false, req.RequiredCapability)
  1112. if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
  1113. release(result)
  1114. continue
  1115. }
  1116. if req.RequireCompact && openAICompactSupportTier(fresh) == 0 {
  1117. compactBlocked = true
  1118. release(result)
  1119. continue
  1120. }
  1121. if fresh.Mode1EffectiveConcurrency() != candidate.account.Mode1EffectiveConcurrency() {
  1122. release(result)
  1123. result, attempted, acquireErr = s.tryAcquireOpenAIAccountSlot(ctx, fresh.ID, fresh.Mode1EffectiveConcurrency(), budget)
  1124. if !attempted {
  1125. continue
  1126. }
  1127. if acquireErr != nil {
  1128. return nil, compactBlocked, acquireErr
  1129. }
  1130. if result == nil || !result.Acquired {
  1131. continue
  1132. }
  1133. }
  1134. if req.SessionHash != "" && !req.PreserveStickyBinding {
  1135. _ = s.service.bindOpenAIStickySessionDuringSelection(ctx, req.GroupID, req.SessionHash, fresh.ID)
  1136. }
  1137. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  1138. Account: fresh,
  1139. Acquired: true,
  1140. ReleaseFunc: result.ReleaseFunc,
  1141. }), compactBlocked, nil
  1142. }
  1143. return nil, compactBlocked, nil
  1144. }
  1145. func (s *defaultOpenAIAccountScheduler) tryAcquireOpenAIAccountSlot(
  1146. ctx context.Context,
  1147. accountID int64,
  1148. maxConcurrency int,
  1149. budget *openAISelectionProbeBudget,
  1150. ) (*AcquireResult, bool, error) {
  1151. if s.service.concurrencyService != nil && maxConcurrency > 0 && !budget.recordAcquire(accountID) {
  1152. return nil, false, nil
  1153. }
  1154. result, err := s.service.tryAcquireAccountSlot(ctx, accountID, maxConcurrency)
  1155. return result, true, err
  1156. }
  1157. func (s *defaultOpenAIAccountScheduler) consumeOpenAISelectionDBRecheck(budget *openAISelectionProbeBudget) bool {
  1158. if s.service.schedulerSnapshot == nil || s.service.accountRepo == nil {
  1159. return true
  1160. }
  1161. return budget.recordRecheck()
  1162. }
  1163. func (s *defaultOpenAIAccountScheduler) tryFallbackToWeightedSticky(
  1164. ctx context.Context,
  1165. req OpenAIAccountScheduleRequest,
  1166. ) (*AccountSelectionResult, error) {
  1167. if !req.StickyWeighted {
  1168. return nil, nil
  1169. }
  1170. for _, accountID := range []int64{req.StickyPreviousAccountID, req.StickyAccountID} {
  1171. if accountID <= 0 {
  1172. continue
  1173. }
  1174. if req.ExcludedIDs != nil {
  1175. if _, excluded := req.ExcludedIDs[accountID]; excluded {
  1176. continue
  1177. }
  1178. }
  1179. account, err := s.service.getSchedulableAccount(ctx, accountID)
  1180. if err != nil || account == nil {
  1181. continue
  1182. }
  1183. if !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
  1184. continue
  1185. }
  1186. account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.GroupID, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability)
  1187. if account == nil {
  1188. if accountID == req.StickyAccountID && strings.TrimSpace(req.SessionHash) != "" {
  1189. _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, req.SessionHash)
  1190. }
  1191. continue
  1192. }
  1193. if !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) {
  1194. if accountID == req.StickyAccountID && strings.TrimSpace(req.SessionHash) != "" {
  1195. _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, req.SessionHash)
  1196. }
  1197. continue
  1198. }
  1199. if !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
  1200. continue
  1201. }
  1202. if req.RequireCompact && openAICompactSupportTier(account) == 0 {
  1203. continue
  1204. }
  1205. // Keep weighted sticky fallback subject to the same free-tier gate as the
  1206. // normal and sticky selection paths. Otherwise an over-quota free account
  1207. // could be reintroduced after the primary candidate pass.
  1208. if len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 {
  1209. continue
  1210. }
  1211. upstreamModel := canonicalOpenAIAccountSchedulingModel(account, req.RequestedModel)
  1212. now := time.Now()
  1213. if isGrokTeamModelRateLimited(account, upstreamModel, now) ||
  1214. isGrokModelQuotaBlocked(account.ID, upstreamModel, now) {
  1215. continue
  1216. }
  1217. result, acquireErr := s.service.tryAcquireAccountSlot(ctx, account.ID, account.Mode1EffectiveConcurrency())
  1218. if acquireErr != nil {
  1219. return nil, acquireErr
  1220. }
  1221. if result != nil && result.Acquired {
  1222. if req.SessionHash != "" && !req.PreserveStickyBinding {
  1223. _ = s.service.bindOpenAIStickySessionDuringSelection(ctx, req.GroupID, req.SessionHash, account.ID)
  1224. }
  1225. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  1226. Account: account,
  1227. Acquired: true,
  1228. ReleaseFunc: result.ReleaseFunc,
  1229. }), nil
  1230. }
  1231. if s.service.concurrencyService != nil {
  1232. cfg := s.service.schedulingConfig()
  1233. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  1234. Account: account,
  1235. WaitPlan: &AccountWaitPlan{
  1236. AccountID: account.ID,
  1237. MaxConcurrency: account.Mode1EffectiveConcurrency(),
  1238. Timeout: cfg.StickySessionWaitTimeout,
  1239. MaxWaiting: cfg.StickySessionMaxWaiting,
  1240. },
  1241. }), nil
  1242. }
  1243. }
  1244. return nil, nil
  1245. }
  1246. // openAISelectionFilterStats counts why candidates were dropped by the
  1247. // selectByLoadBalance initial filter. Historically these exclusions were
  1248. // silent (debug logs at best), so a "no available accounts" failure with
  1249. // excluded_account_count=0 was undiagnosable from the error alone (#4599).
  1250. // The reasons map is lazily allocated: on the happy path (nothing filtered
  1251. // out, or an account is eventually selected) no extra allocation happens.
  1252. type openAISelectionFilterStats struct {
  1253. pool int
  1254. reasons map[string]int
  1255. }
  1256. func (s *openAISelectionFilterStats) exclude(reason string) {
  1257. if s.reasons == nil {
  1258. s.reasons = make(map[string]int, 4)
  1259. }
  1260. s.reasons[reason]++
  1261. }
  1262. // summary renders deterministic exclusion statistics for scheduling error
  1263. // messages, e.g. "pool=3, filtered: model_not_supported=2 quota_auto_pause_7d=1".
  1264. // Reasons are sorted lexicographically so the output is stable for tests and
  1265. // log aggregation. extra, when non-empty, is appended as a trailing marker.
  1266. func (s openAISelectionFilterStats) summary(extra string) string {
  1267. var b strings.Builder
  1268. _, _ = b.WriteString("pool=")
  1269. _, _ = b.WriteString(strconv.Itoa(s.pool))
  1270. if len(s.reasons) > 0 {
  1271. reasons := make([]string, 0, len(s.reasons))
  1272. for reason := range s.reasons {
  1273. reasons = append(reasons, reason)
  1274. }
  1275. sort.Strings(reasons)
  1276. _, _ = b.WriteString(", filtered:")
  1277. for _, reason := range reasons {
  1278. _, _ = b.WriteString(" ")
  1279. _, _ = b.WriteString(reason)
  1280. _, _ = b.WriteString("=")
  1281. _, _ = b.WriteString(strconv.Itoa(s.reasons[reason]))
  1282. }
  1283. }
  1284. if extra != "" {
  1285. _, _ = b.WriteString(", ")
  1286. _, _ = b.WriteString(extra)
  1287. }
  1288. return b.String()
  1289. }
  1290. func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
  1291. ctx context.Context,
  1292. req OpenAIAccountScheduleRequest,
  1293. ) (*AccountSelectionResult, int, int, float64, error) {
  1294. budget := newOpenAISelectionProbeBudget()
  1295. accounts, err := s.service.listSchedulableAccounts(ctx, req.GroupID, req.Platform)
  1296. if err != nil {
  1297. return nil, 0, 0, 0, err
  1298. }
  1299. if len(accounts) == 0 {
  1300. return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary(""))
  1301. }
  1302. // Local free-tier soft gate on the Grok scheduling path only (not admin probe).
  1303. accounts = s.filterGrokFreeQuotaAccounts(ctx, accounts)
  1304. if len(accounts) == 0 {
  1305. return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_free_quota_soft_gate"))
  1306. }
  1307. // Team+model rate-limit cool: siblings of a 429'd team skip the hot model.
  1308. if req.Platform == PlatformGrok {
  1309. now := time.Now()
  1310. filtered := filterGrokTeamModelRateLimitedAccounts(accounts, req.RequestedModel, now)
  1311. if len(filtered) == 0 && len(accounts) > 0 {
  1312. return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_team_model_rate_limit"))
  1313. }
  1314. if filtered != nil {
  1315. accounts = filtered
  1316. }
  1317. // Per-account model free-usage soft-block (other models stay eligible).
  1318. modelFiltered := filterGrokModelQuotaBlockedAccounts(accounts, req.RequestedModel, now)
  1319. if len(modelFiltered) == 0 && len(accounts) > 0 {
  1320. return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_model_quota_block"))
  1321. }
  1322. accounts = modelFiltered
  1323. }
  1324. // require_privacy_set: 获取分组信息
  1325. var schedGroup *Group
  1326. if req.GroupID != nil && s.service.schedulerSnapshot != nil {
  1327. schedGroup, _ = s.service.schedulerSnapshot.GetGroupByID(ctx, *req.GroupID)
  1328. }
  1329. filterStats := openAISelectionFilterStats{pool: len(accounts)}
  1330. filtered := make([]*Account, 0, len(accounts))
  1331. loadReq := make([]AccountWithConcurrency, 0, len(accounts))
  1332. for i := range accounts {
  1333. account := &accounts[i]
  1334. if req.ExcludedIDs != nil {
  1335. if _, excluded := req.ExcludedIDs[account.ID]; excluded {
  1336. filterStats.exclude("excluded")
  1337. continue
  1338. }
  1339. }
  1340. if !account.IsSchedulable() {
  1341. filterStats.exclude("not_schedulable")
  1342. continue
  1343. }
  1344. if account.Platform != NormalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() {
  1345. filterStats.exclude("platform_mismatch")
  1346. continue
  1347. }
  1348. if s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
  1349. filterStats.exclude("runtime_blocked")
  1350. continue
  1351. }
  1352. // require_privacy_set is a group-scoped eligibility gate. Do not mutate the
  1353. // shared account: another group may intentionally allow accounts whose
  1354. // upstream privacy setting has not been confirmed.
  1355. if schedGroup != nil && schedGroup.RequirePrivacySet && !account.IsPrivacySet() {
  1356. filterStats.exclude("privacy_not_set")
  1357. continue
  1358. }
  1359. if compatible, reason := s.isAccountRequestCompatibleReason(ctx, account, req); !compatible {
  1360. filterStats.exclude(reason)
  1361. continue
  1362. }
  1363. if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
  1364. filterStats.exclude("transport_incompatible")
  1365. continue
  1366. }
  1367. filtered = append(filtered, account)
  1368. loadReq = append(loadReq, AccountWithConcurrency{
  1369. ID: account.ID,
  1370. MaxConcurrency: account.EffectiveLoadFactor(),
  1371. })
  1372. }
  1373. if len(filtered) == 0 {
  1374. return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, filterStats.summary(""))
  1375. }
  1376. loadMap := map[int64]*AccountLoadInfo{}
  1377. if s.service.concurrencyService != nil {
  1378. if batchLoad, loadErr := s.service.concurrencyService.GetAccountsLoadBatch(ctx, loadReq); loadErr == nil {
  1379. loadMap = batchLoad
  1380. }
  1381. }
  1382. if req.SubscriptionPriority {
  1383. subscriptionAccounts, regularAccounts := partitionOpenAIChatGPTSubscriptionAccounts(filtered)
  1384. if len(subscriptionAccounts) > 0 {
  1385. attempt := s.trySelectByLoadBalancePool(ctx, req, subscriptionAccounts, loadMap, budget)
  1386. if attempt.err != nil && (!attempt.noCompactCandidates || len(regularAccounts) <= 0) {
  1387. return nil, attempt.candidateCount, attempt.topK, attempt.loadSkew, attempt.err
  1388. }
  1389. if attempt.result != nil {
  1390. return attempt.result, attempt.candidateCount, attempt.topK, attempt.loadSkew, nil
  1391. }
  1392. if len(regularAccounts) > 0 {
  1393. regularAttempt := s.trySelectByLoadBalancePool(ctx, req, regularAccounts, loadMap, budget)
  1394. if regularAttempt.err != nil && !regularAttempt.noCompactCandidates {
  1395. return nil, regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew, regularAttempt.err
  1396. }
  1397. if regularAttempt.result != nil {
  1398. return regularAttempt.result, regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew, nil
  1399. }
  1400. var result *AccountSelectionResult
  1401. candidateCount, topK, loadSkew := regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew
  1402. fallbackErr := regularAttempt.err
  1403. if regularAttempt.err == nil {
  1404. result, candidateCount, topK, loadSkew, fallbackErr = s.finishLoadBalanceSelectionFallback(ctx, req, regularAttempt, budget, filterStats)
  1405. if fallbackErr == nil && result != nil {
  1406. return result, candidateCount, topK, loadSkew, nil
  1407. }
  1408. }
  1409. // 常规池既无法获取也无法排队(含仅剩不支持 compact 的候选)时,
  1410. // 回退到订阅池的等待计划:busy-but-waitable 的订阅账号不应因常规池存在
  1411. // 而被丢弃,否则开启订阅优先反而让本可排队成功的请求硬失败。
  1412. subResult, subCandidateCount, subTopK, subLoadSkew, subErr := s.finishLoadBalanceSelectionFallback(ctx, req, attempt, budget, filterStats)
  1413. if subErr == nil && subResult != nil {
  1414. return subResult, subCandidateCount, subTopK, subLoadSkew, nil
  1415. }
  1416. return result, candidateCount, topK, loadSkew, fallbackErr
  1417. }
  1418. return s.finishLoadBalanceSelectionFallback(ctx, req, attempt, budget, filterStats)
  1419. }
  1420. }
  1421. attempt := s.trySelectByLoadBalancePool(ctx, req, filtered, loadMap, budget)
  1422. if attempt.err != nil {
  1423. return nil, attempt.candidateCount, attempt.topK, attempt.loadSkew, attempt.err
  1424. }
  1425. if attempt.result != nil {
  1426. return attempt.result, attempt.candidateCount, attempt.topK, attempt.loadSkew, nil
  1427. }
  1428. return s.finishLoadBalanceSelectionFallback(ctx, req, attempt, budget, filterStats)
  1429. }
  1430. func partitionOpenAIChatGPTSubscriptionAccounts(accounts []*Account) ([]*Account, []*Account) {
  1431. subscriptionAccounts := make([]*Account, 0, len(accounts))
  1432. regularAccounts := make([]*Account, 0, len(accounts))
  1433. for _, account := range accounts {
  1434. if account != nil && account.IsOpenAIChatGPTSubscription() {
  1435. subscriptionAccounts = append(subscriptionAccounts, account)
  1436. continue
  1437. }
  1438. regularAccounts = append(regularAccounts, account)
  1439. }
  1440. return subscriptionAccounts, regularAccounts
  1441. }
  1442. func (s *defaultOpenAIAccountScheduler) trySelectByLoadBalancePool(
  1443. ctx context.Context,
  1444. req OpenAIAccountScheduleRequest,
  1445. filtered []*Account,
  1446. loadMap map[int64]*AccountLoadInfo,
  1447. budget *openAISelectionProbeBudget,
  1448. ) openAIAccountLoadSelectionAttempt {
  1449. plan := s.buildOpenAIAccountLoadPlan(ctx, req, filtered, loadMap)
  1450. if openAICostOverflowExpanded(req, plan) {
  1451. budget.enableLimit()
  1452. }
  1453. attempt := openAIAccountLoadSelectionAttempt{
  1454. selectionOrder: plan.selectionOrder,
  1455. candidateCount: plan.candidateCount,
  1456. topK: plan.topK,
  1457. loadSkew: plan.loadSkew,
  1458. }
  1459. if req.RequireCompact && len(plan.candidates) == 0 && len(plan.staleSnapshotCompactRetry) == 0 {
  1460. attempt.noCompactCandidates = true
  1461. attempt.err = ErrNoAvailableCompactAccounts
  1462. return attempt
  1463. }
  1464. if req.RequireCompact && len(attempt.selectionOrder) == 0 && s.service.schedulerSnapshot == nil {
  1465. attempt.noCompactCandidates = true
  1466. attempt.err = ErrNoAvailableCompactAccounts
  1467. return attempt
  1468. }
  1469. if len(attempt.selectionOrder) == 0 {
  1470. attempt.compactBlocked = req.RequireCompact && len(plan.allCandidates) > 0
  1471. return attempt
  1472. }
  1473. result, compactBlocked, acquireErr := s.tryAcquireOpenAISelectionOrderWithBudget(ctx, req, attempt.selectionOrder, budget)
  1474. attempt.compactBlocked = compactBlocked
  1475. if acquireErr != nil {
  1476. attempt.err = acquireErr
  1477. return attempt
  1478. }
  1479. if result != nil {
  1480. attempt.result = result
  1481. return attempt
  1482. }
  1483. if s.service.concurrencyService != nil && !budget.acquireExhausted() {
  1484. loadReq := buildOpenAIAccountLoadRequest(filtered)
  1485. if freshLoadMap, loadErr := s.service.concurrencyService.GetAccountsLoadBatchFresh(ctx, loadReq); loadErr == nil {
  1486. freshPlan := s.buildOpenAIAccountLoadPlan(ctx, req, filtered, freshLoadMap)
  1487. if openAICostOverflowExpanded(req, freshPlan) {
  1488. budget.enableLimit()
  1489. }
  1490. if len(freshPlan.selectionOrder) > 0 {
  1491. freshResult, freshCompactBlocked, freshAcquireErr := s.tryAcquireOpenAISelectionOrderWithBudget(ctx, req, freshPlan.selectionOrder, budget)
  1492. if freshAcquireErr != nil {
  1493. attempt.err = freshAcquireErr
  1494. return attempt
  1495. }
  1496. if freshResult != nil {
  1497. attempt.result = freshResult
  1498. attempt.selectionOrder = freshPlan.selectionOrder
  1499. attempt.candidateCount = freshPlan.candidateCount
  1500. attempt.topK = freshPlan.topK
  1501. attempt.loadSkew = freshPlan.loadSkew
  1502. return attempt
  1503. }
  1504. attempt.compactBlocked = attempt.compactBlocked || freshCompactBlocked
  1505. attempt.selectionOrder = freshPlan.selectionOrder
  1506. attempt.candidateCount = freshPlan.candidateCount
  1507. attempt.topK = freshPlan.topK
  1508. attempt.loadSkew = freshPlan.loadSkew
  1509. }
  1510. }
  1511. }
  1512. return attempt
  1513. }
  1514. func openAICostOverflowExpanded(req OpenAIAccountScheduleRequest, plan openAIAccountLoadPlan) bool {
  1515. if !plan.includeOverflowFallback || plan.topK <= 0 {
  1516. return false
  1517. }
  1518. if !req.RequireCompact {
  1519. return len(plan.candidates) > plan.topK
  1520. }
  1521. supported, unknown := 0, 0
  1522. for _, candidate := range plan.candidates {
  1523. switch openAICompactSupportTier(candidate.account) {
  1524. case 2:
  1525. supported++
  1526. case 1:
  1527. unknown++
  1528. }
  1529. }
  1530. return supported > plan.topK || unknown > plan.topK
  1531. }
  1532. func buildOpenAIAccountLoadRequest(accounts []*Account) []AccountWithConcurrency {
  1533. loadReq := make([]AccountWithConcurrency, 0, len(accounts))
  1534. for _, account := range accounts {
  1535. if account == nil {
  1536. continue
  1537. }
  1538. loadReq = append(loadReq, AccountWithConcurrency{
  1539. ID: account.ID,
  1540. MaxConcurrency: account.EffectiveLoadFactor(),
  1541. })
  1542. }
  1543. return loadReq
  1544. }
  1545. func (s *defaultOpenAIAccountScheduler) finishLoadBalanceSelectionFallback(
  1546. ctx context.Context,
  1547. req OpenAIAccountScheduleRequest,
  1548. attempt openAIAccountLoadSelectionAttempt,
  1549. budget *openAISelectionProbeBudget,
  1550. filterStats openAISelectionFilterStats,
  1551. ) (*AccountSelectionResult, int, int, float64, error) {
  1552. candidateCount := attempt.candidateCount
  1553. topK := attempt.topK
  1554. loadSkew := attempt.loadSkew
  1555. if len(attempt.selectionOrder) == 0 {
  1556. return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, attempt.compactBlocked, filterStats.summary("selection_order_empty"))
  1557. }
  1558. if stickyFallback, stickyErr := s.tryFallbackToWeightedSticky(ctx, req); stickyErr != nil {
  1559. return nil, candidateCount, topK, loadSkew, stickyErr
  1560. } else if stickyFallback != nil {
  1561. return stickyFallback, candidateCount, topK, loadSkew, nil
  1562. }
  1563. cfg := s.service.schedulingConfig()
  1564. compactBlocked := attempt.compactBlocked
  1565. // WaitPlan.MaxConcurrency 使用 Concurrency(非 EffectiveLoadFactor),因为 WaitPlan 控制的是 Redis 实际并发槽位等待。
  1566. passes := 1
  1567. if budget != nil && budget.limited {
  1568. passes = 4
  1569. }
  1570. for pass := 0; pass < passes; pass++ {
  1571. wantAttempted := pass == 1 || pass == 3
  1572. wantKnownFull := pass >= 2
  1573. for _, candidate := range attempt.selectionOrder {
  1574. if candidate.account == nil {
  1575. continue
  1576. }
  1577. if budget != nil && budget.limited {
  1578. knownFull := candidate.loadKnown && candidate.account.Mode1EffectiveConcurrency() > 0 &&
  1579. candidate.loadInfo.CurrentConcurrency >= candidate.account.Mode1EffectiveConcurrency()
  1580. if budget.wasAttempted(candidate.account.ID) != wantAttempted || knownFull != wantKnownFull {
  1581. continue
  1582. }
  1583. }
  1584. fresh := s.service.resolveFreshSchedulableOpenAIAccount(ctx, candidate.account, req.Platform, req.RequestedModel, false, req.RequiredCapability)
  1585. if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
  1586. continue
  1587. }
  1588. if !s.consumeOpenAISelectionDBRecheck(budget) {
  1589. return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, compactBlocked, filterStats.summary("selection_order_exhausted"))
  1590. }
  1591. fresh = s.service.recheckSelectedOpenAIAccountFromDB(ctx, fresh, req.GroupID, req.Platform, req.RequestedModel, false, req.RequiredCapability)
  1592. if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
  1593. continue
  1594. }
  1595. if req.RequireCompact && openAICompactSupportTier(fresh) == 0 {
  1596. compactBlocked = true
  1597. continue
  1598. }
  1599. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  1600. Account: fresh,
  1601. WaitPlan: &AccountWaitPlan{
  1602. AccountID: fresh.ID,
  1603. MaxConcurrency: fresh.Mode1EffectiveConcurrency(),
  1604. Timeout: cfg.FallbackWaitTimeout,
  1605. MaxWaiting: cfg.FallbackMaxWaiting,
  1606. },
  1607. }), candidateCount, topK, loadSkew, nil
  1608. }
  1609. }
  1610. return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, compactBlocked, filterStats.summary("selection_order_exhausted"))
  1611. }
  1612. func (s *defaultOpenAIAccountScheduler) isAccountTransportCompatible(account *Account, requiredTransport OpenAIUpstreamTransport) bool {
  1613. if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
  1614. return true
  1615. }
  1616. if s == nil || s.service == nil {
  1617. return false
  1618. }
  1619. return s.service.isOpenAIAccountTransportCompatible(account, requiredTransport)
  1620. }
  1621. func (s *defaultOpenAIAccountScheduler) lookupShadowParentAccount(ctx context.Context, id int64) *Account {
  1622. if s == nil || s.service == nil {
  1623. return nil
  1624. }
  1625. if s.service.schedulerSnapshot != nil {
  1626. if account, err := s.service.schedulerSnapshot.GetAccount(ctx, id); err == nil && account != nil {
  1627. return account
  1628. }
  1629. }
  1630. if s.service.accountRepo == nil {
  1631. return nil
  1632. }
  1633. account, _ := s.service.accountRepo.GetByID(ctx, id)
  1634. return account
  1635. }
  1636. func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(ctx context.Context, account *Account, req OpenAIAccountScheduleRequest) bool {
  1637. compatible, _ := s.isAccountRequestCompatibleReason(ctx, account, req)
  1638. return compatible
  1639. }
  1640. // isAccountRequestCompatibleReason reports whether the account can serve the
  1641. // request, and when it cannot, names the veto point. The reason feeds
  1642. // openAISelectionFilterStats so that "no available accounts" errors state why
  1643. // each candidate was dropped instead of failing silently (#4599).
  1644. func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx context.Context, account *Account, req OpenAIAccountScheduleRequest) (bool, string) {
  1645. if account == nil {
  1646. return false, "account_nil"
  1647. }
  1648. if req.RequirePrivacySet && !account.IsPrivacySet() {
  1649. return false, "privacy_not_set"
  1650. }
  1651. if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
  1652. return false, "runtime_blocked"
  1653. }
  1654. if s != nil && s.service != nil && s.service.isOpenAIProxyStreamQuarantined(ctx, account) {
  1655. return false, "proxy_stream_quarantined"
  1656. }
  1657. // Quota auto-pause must be evaluated during the initial filter too. Without it the
  1658. // TopK candidate pool can be filled with paused accounts and the later fresh/DB
  1659. // rechecks won't reach healthy accounts that fell outside TopK — manifesting as
  1660. // "no available accounts" even though healthy ones exist.
  1661. if paused, decision := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused {
  1662. if decision.reason != "" {
  1663. return false, decision.reason
  1664. }
  1665. reason := "quota_auto_pause"
  1666. if decision.window != "" {
  1667. reason += "_" + decision.window
  1668. }
  1669. return false, reason
  1670. }
  1671. // 母账号健康联动:影子账号的凭据来自母账号,母账号不可调度时影子也不应被选中。
  1672. // Parent-health gate: shadow borrows the parent's credentials; an unschedulable
  1673. // parent must block the shadow across all scheduler paths.
  1674. if !parentHealthyForShadow(account, func(id int64) *Account {
  1675. return s.lookupShadowParentAccount(ctx, id)
  1676. }) {
  1677. return false, "shadow_parent_unhealthy"
  1678. }
  1679. if req.RequestedModel != "" && !account.IsModelSupported(req.RequestedModel) {
  1680. return false, "model_not_supported"
  1681. }
  1682. if req.GroupID != nil && s != nil && s.service != nil &&
  1683. s.service.needsUpstreamChannelRestrictionCheck(ctx, req.GroupID) &&
  1684. s.service.isUpstreamModelRestrictedByChannel(ctx, *req.GroupID, account, req.RequestedModel, req.RequireCompact) {
  1685. return false, "channel_upstream_restricted"
  1686. }
  1687. if !accountSupportsOpenAICapabilities(account, req.RequiredCapability, req.RequiredImageCapability) {
  1688. return false, "capability_mismatch"
  1689. }
  1690. // 分组利润控制:不合格账号在候选过滤与抢槽后终检阶段即被排除,
  1691. // 排序/评分/粘性/熔断只在合格账号之间工作;named reason 进入 filter stats。
  1692. if vetoed, reason := openAIProfitControlVetoReason(ctx, account); vetoed {
  1693. return false, reason
  1694. }
  1695. return true, ""
  1696. }
  1697. func (s *defaultOpenAIAccountScheduler) ReportResult(accountID int64, success bool, firstTokenMs *int) {
  1698. if s == nil || s.stats == nil {
  1699. return
  1700. }
  1701. s.stats.report(accountID, success, firstTokenMs)
  1702. }
  1703. func (s *defaultOpenAIAccountScheduler) ReportSwitch() {
  1704. if s == nil {
  1705. return
  1706. }
  1707. s.metrics.recordSwitch()
  1708. }
  1709. func (s *defaultOpenAIAccountScheduler) SnapshotMetrics() OpenAIAccountSchedulerMetricsSnapshot {
  1710. if s == nil {
  1711. return OpenAIAccountSchedulerMetricsSnapshot{}
  1712. }
  1713. selectTotal := s.metrics.selectTotal.Load()
  1714. prevHit := s.metrics.stickyPreviousHitTotal.Load()
  1715. sessionHit := s.metrics.stickySessionHitTotal.Load()
  1716. switchTotal := s.metrics.accountSwitchTotal.Load()
  1717. latencyTotal := s.metrics.latencyMsTotal.Load()
  1718. loadSkewTotal := s.metrics.loadSkewMilliTotal.Load()
  1719. snapshot := OpenAIAccountSchedulerMetricsSnapshot{
  1720. SelectTotal: selectTotal,
  1721. StickyPreviousHitTotal: prevHit,
  1722. StickySessionHitTotal: sessionHit,
  1723. LoadBalanceSelectTotal: s.metrics.loadBalanceSelectTotal.Load(),
  1724. AccountSwitchTotal: switchTotal,
  1725. SchedulerLatencyMsTotal: latencyTotal,
  1726. RuntimeStatsAccountCount: s.stats.size(),
  1727. }
  1728. if selectTotal > 0 {
  1729. snapshot.SchedulerLatencyMsAvg = float64(latencyTotal) / float64(selectTotal)
  1730. snapshot.StickyHitRatio = float64(prevHit+sessionHit) / float64(selectTotal)
  1731. snapshot.AccountSwitchRate = float64(switchTotal) / float64(selectTotal)
  1732. snapshot.LoadSkewAvg = float64(loadSkewTotal) / 1000 / float64(selectTotal)
  1733. }
  1734. return snapshot
  1735. }
  1736. func (s *OpenAIGatewayService) openAIAdvancedSchedulerSettingRepo() SettingRepository {
  1737. if s == nil || s.rateLimitService == nil || s.rateLimitService.settingService == nil {
  1738. return nil
  1739. }
  1740. return s.rateLimitService.settingService.settingRepo
  1741. }
  1742. func (s *OpenAIGatewayService) openAIAdvancedSchedulerRuntimeSettings(ctx context.Context) openAIAdvancedSchedulerRuntimeSettings {
  1743. if cached, ok := openAIAdvancedSchedulerSettingCache.Load().(*cachedOpenAIAdvancedSchedulerSetting); ok && cached != nil {
  1744. if time.Now().UnixNano() < cached.expiresAt {
  1745. return openAIAdvancedSchedulerRuntimeSettings{
  1746. lowUpstreamRatePriorityEnabled: cached.lowUpstreamRatePriorityEnabled,
  1747. oauthSchedulingRateMultiplier: cached.oauthSchedulingRateMultiplier,
  1748. enabled: cached.enabled,
  1749. stickyWeightedEnabled: cached.stickyWeightedEnabled,
  1750. subscriptionPriorityEnabled: cached.subscriptionPriorityEnabled,
  1751. lbTopKOverride: cached.lbTopKOverride,
  1752. weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(cached.weightOverrides),
  1753. }
  1754. }
  1755. }
  1756. result, _, _ := openAIAdvancedSchedulerSettingSF.Do(openAIAdvancedSchedulerSettingKey, func() (any, error) {
  1757. if cached, ok := openAIAdvancedSchedulerSettingCache.Load().(*cachedOpenAIAdvancedSchedulerSetting); ok && cached != nil {
  1758. if time.Now().UnixNano() < cached.expiresAt {
  1759. return openAIAdvancedSchedulerRuntimeSettings{
  1760. lowUpstreamRatePriorityEnabled: cached.lowUpstreamRatePriorityEnabled,
  1761. oauthSchedulingRateMultiplier: cached.oauthSchedulingRateMultiplier,
  1762. enabled: cached.enabled,
  1763. stickyWeightedEnabled: cached.stickyWeightedEnabled,
  1764. subscriptionPriorityEnabled: cached.subscriptionPriorityEnabled,
  1765. lbTopKOverride: cached.lbTopKOverride,
  1766. weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(cached.weightOverrides),
  1767. }, nil
  1768. }
  1769. }
  1770. lowUpstreamRatePriorityEnabled := false
  1771. oauthSchedulingRateMultiplier := defaultOpenAIOAuthSchedulingRateMultiplier
  1772. enabled := false
  1773. stickyWeightedEnabled := false
  1774. subscriptionPriorityEnabled := false
  1775. lbTopKOverride := 0
  1776. weightOverrides := map[string]float64{}
  1777. if repo := s.openAIAdvancedSchedulerSettingRepo(); repo != nil {
  1778. dbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIAdvancedSchedulerSettingDBTimeout)
  1779. defer cancel()
  1780. if values, err := repo.GetMultiple(dbCtx, openAIAdvancedSchedulerRuntimeSettingKeys()); err == nil {
  1781. lowUpstreamRatePriorityEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAILowUpstreamRatePriorityEnabled]), "true")
  1782. oauthSchedulingRateMultiplier = parseOpenAIOAuthSchedulingRateMultiplier(values[SettingKeyOpenAIOAuthSchedulingRateMultiplier])
  1783. enabled = strings.EqualFold(strings.TrimSpace(values[openAIAdvancedSchedulerSettingKey]), "true")
  1784. stickyWeightedEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled]), "true")
  1785. subscriptionPriorityEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]), "true")
  1786. lbTopKOverride = parsePositiveIntOverride(values[SettingKeyOpenAIAdvancedSchedulerLBTopK])
  1787. weightOverrides = parseOpenAIAdvancedSchedulerWeightOverrides(values)
  1788. } else {
  1789. // 批量读取失败时逐键降级,覆盖全部键(含 TopK/权重),避免只加载布尔开关
  1790. // 而静默丢弃管理员配置的覆盖值;降级状态会被缓存一个 TTL,必须留痕。
  1791. slog.Warn("openai_advanced_scheduler_settings_batch_load_failed", "error", err)
  1792. fallbackValues := make(map[string]string)
  1793. for _, key := range openAIAdvancedSchedulerRuntimeSettingKeys() {
  1794. if value, valueErr := repo.GetValue(dbCtx, key); valueErr == nil {
  1795. fallbackValues[key] = value
  1796. }
  1797. }
  1798. lowUpstreamRatePriorityEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAILowUpstreamRatePriorityEnabled]), "true")
  1799. oauthSchedulingRateMultiplier = parseOpenAIOAuthSchedulingRateMultiplier(fallbackValues[SettingKeyOpenAIOAuthSchedulingRateMultiplier])
  1800. enabled = strings.EqualFold(strings.TrimSpace(fallbackValues[openAIAdvancedSchedulerSettingKey]), "true")
  1801. stickyWeightedEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled]), "true")
  1802. subscriptionPriorityEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]), "true")
  1803. lbTopKOverride = parsePositiveIntOverride(fallbackValues[SettingKeyOpenAIAdvancedSchedulerLBTopK])
  1804. weightOverrides = parseOpenAIAdvancedSchedulerWeightOverrides(fallbackValues)
  1805. }
  1806. }
  1807. openAIAdvancedSchedulerSettingCache.Store(&cachedOpenAIAdvancedSchedulerSetting{
  1808. lowUpstreamRatePriorityEnabled: lowUpstreamRatePriorityEnabled,
  1809. oauthSchedulingRateMultiplier: oauthSchedulingRateMultiplier,
  1810. enabled: enabled,
  1811. stickyWeightedEnabled: stickyWeightedEnabled,
  1812. subscriptionPriorityEnabled: subscriptionPriorityEnabled,
  1813. lbTopKOverride: lbTopKOverride,
  1814. weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(weightOverrides),
  1815. expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(),
  1816. })
  1817. return openAIAdvancedSchedulerRuntimeSettings{
  1818. lowUpstreamRatePriorityEnabled: lowUpstreamRatePriorityEnabled,
  1819. oauthSchedulingRateMultiplier: oauthSchedulingRateMultiplier,
  1820. enabled: enabled,
  1821. stickyWeightedEnabled: stickyWeightedEnabled,
  1822. subscriptionPriorityEnabled: subscriptionPriorityEnabled,
  1823. lbTopKOverride: lbTopKOverride,
  1824. weightOverrides: weightOverrides,
  1825. }, nil
  1826. })
  1827. settings, _ := result.(openAIAdvancedSchedulerRuntimeSettings)
  1828. return settings
  1829. }
  1830. func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerEnabled(ctx context.Context) bool {
  1831. return s.openAIAdvancedSchedulerRuntimeSettings(ctx).enabled
  1832. }
  1833. func (s *OpenAIGatewayService) isOpenAILowUpstreamRatePriorityEnabled(ctx context.Context) bool {
  1834. settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
  1835. return !settings.enabled && settings.lowUpstreamRatePriorityEnabled
  1836. }
  1837. func (s *OpenAIGatewayService) openAIOAuthSchedulingRateMultiplier(ctx context.Context) float64 {
  1838. return s.openAIAdvancedSchedulerRuntimeSettings(ctx).oauthSchedulingRateMultiplier
  1839. }
  1840. func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx context.Context) bool {
  1841. settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
  1842. return settings.enabled && settings.stickyWeightedEnabled
  1843. }
  1844. func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerSubscriptionPriorityEnabled(ctx context.Context) bool {
  1845. settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
  1846. return settings.enabled && settings.subscriptionPriorityEnabled
  1847. }
  1848. func openAIAdvancedSchedulerRuntimeSettingKeys() []string {
  1849. keys := []string{
  1850. SettingKeyOpenAILowUpstreamRatePriorityEnabled,
  1851. SettingKeyOpenAIOAuthSchedulingRateMultiplier,
  1852. openAIAdvancedSchedulerSettingKey,
  1853. SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled,
  1854. SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled,
  1855. SettingKeyOpenAIAdvancedSchedulerLBTopK,
  1856. }
  1857. for _, spec := range openAIAdvancedSchedulerWeightOverrideSpecs() {
  1858. keys = append(keys, spec.key)
  1859. }
  1860. return keys
  1861. }
  1862. type openAIAdvancedSchedulerWeightOverrideSpec struct {
  1863. key string
  1864. name string
  1865. }
  1866. func openAIAdvancedSchedulerWeightOverrideSpecs() []openAIAdvancedSchedulerWeightOverrideSpec {
  1867. return []openAIAdvancedSchedulerWeightOverrideSpec{
  1868. {key: SettingKeyOpenAIAdvancedSchedulerWeightPriority, name: "priority"},
  1869. {key: SettingKeyOpenAIAdvancedSchedulerWeightLoad, name: "load"},
  1870. {key: SettingKeyOpenAIAdvancedSchedulerWeightQueue, name: "queue"},
  1871. {key: SettingKeyOpenAIAdvancedSchedulerWeightErrorRate, name: "error_rate"},
  1872. {key: SettingKeyOpenAIAdvancedSchedulerWeightTTFT, name: "ttft"},
  1873. {key: SettingKeyOpenAIAdvancedSchedulerWeightReset, name: "reset"},
  1874. {key: SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom, name: "quota_headroom"},
  1875. {key: SettingKeyOpenAIAdvancedSchedulerWeightUpstreamCost, name: "upstream_cost"},
  1876. {key: SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse, name: "previous_response"},
  1877. {key: SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky, name: "session_sticky"},
  1878. }
  1879. }
  1880. func parsePositiveIntOverride(raw string) int {
  1881. raw = strings.TrimSpace(raw)
  1882. if raw == "" {
  1883. return 0
  1884. }
  1885. value, err := strconv.Atoi(raw)
  1886. if err != nil || value <= 0 {
  1887. return 0
  1888. }
  1889. return value
  1890. }
  1891. func parseOpenAIAdvancedSchedulerWeightOverrides(values map[string]string) map[string]float64 {
  1892. overrides := map[string]float64{}
  1893. for _, spec := range openAIAdvancedSchedulerWeightOverrideSpecs() {
  1894. raw := strings.TrimSpace(values[spec.key])
  1895. if raw == "" {
  1896. continue
  1897. }
  1898. value, err := strconv.ParseFloat(raw, 64)
  1899. if err != nil || value < 0 || math.IsNaN(value) || math.IsInf(value, 0) {
  1900. continue
  1901. }
  1902. overrides[spec.name] = value
  1903. }
  1904. return overrides
  1905. }
  1906. func cloneOpenAIAdvancedSchedulerWeightOverrides(in map[string]float64) map[string]float64 {
  1907. if len(in) == 0 {
  1908. return nil
  1909. }
  1910. out := make(map[string]float64, len(in))
  1911. for key, value := range in {
  1912. out[key] = value
  1913. }
  1914. return out
  1915. }
  1916. func (s *OpenAIGatewayService) getOpenAIAccountScheduler(ctx context.Context) OpenAIAccountScheduler {
  1917. if s == nil {
  1918. return nil
  1919. }
  1920. if !s.isOpenAIAdvancedSchedulerEnabled(ctx) {
  1921. return nil
  1922. }
  1923. s.openaiSchedulerOnce.Do(func() {
  1924. if s.openaiAccountStats == nil {
  1925. s.openaiAccountStats = newOpenAIAccountRuntimeStats()
  1926. }
  1927. if s.openaiScheduler == nil {
  1928. s.openaiScheduler = newDefaultOpenAIAccountScheduler(s, s.openaiAccountStats)
  1929. }
  1930. })
  1931. return s.openaiScheduler
  1932. }
  1933. func resetOpenAIAdvancedSchedulerSettingCacheForTest() {
  1934. openAIAdvancedSchedulerSettingCache = atomic.Value{}
  1935. openAIAdvancedSchedulerSettingSF = singleflight.Group{}
  1936. }
  1937. func (s *OpenAIGatewayService) SelectAccountWithScheduler(
  1938. ctx context.Context,
  1939. groupID *int64,
  1940. previousResponseID string,
  1941. sessionHash string,
  1942. requestedModel string,
  1943. excludedIDs map[int64]struct{},
  1944. requiredTransport OpenAIUpstreamTransport,
  1945. requireCompact bool,
  1946. ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
  1947. return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, "", "", requireCompact, PlatformOpenAI, false, true)
  1948. }
  1949. // SelectAccountWithSchedulerForCapability 按能力要求调度账号。
  1950. // previousResponseCanMove 表示首包 input 可自行重建工具续链,previous_response_id 允许跨账号迁移
  1951. // (粘性加权模式下改为加权偏好而非硬粘连)。
  1952. func (s *OpenAIGatewayService) SelectAccountWithSchedulerForCapability(
  1953. ctx context.Context,
  1954. groupID *int64,
  1955. previousResponseID string,
  1956. sessionHash string,
  1957. requestedModel string,
  1958. excludedIDs map[int64]struct{},
  1959. requiredTransport OpenAIUpstreamTransport,
  1960. requiredCapability OpenAIEndpointCapability,
  1961. requireCompact bool,
  1962. previousResponseCanMove bool,
  1963. useUpstreamTokenCost bool,
  1964. platformOverride ...string,
  1965. ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
  1966. platform := PlatformOpenAI
  1967. if len(platformOverride) > 0 {
  1968. platform = platformOverride[0]
  1969. }
  1970. return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, "", requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost)
  1971. }
  1972. func (s *OpenAIGatewayService) SelectAccountWithSchedulerForImages(
  1973. ctx context.Context,
  1974. groupID *int64,
  1975. sessionHash string,
  1976. requestedModel string,
  1977. excludedIDs map[int64]struct{},
  1978. requiredCapability OpenAIImagesCapability,
  1979. ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
  1980. selection, decision, err := s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", requiredCapability, false, PlatformOpenAI, false, false)
  1981. if err == nil && selection != nil && selection.Account != nil {
  1982. return selection, decision, nil
  1983. }
  1984. // 如果要求 native 能力(如指定了模型)但没有可用的 APIKey 账号,回退到 basic(OAuth 账号)
  1985. if requiredCapability == OpenAIImagesCapabilityNative {
  1986. return s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", OpenAIImagesCapabilityBasic, false, PlatformOpenAI, false, false)
  1987. }
  1988. return selection, decision, err
  1989. }
  1990. // selectAccountWithScheduler wraps selectAccountWithSchedulerOnce with a
  1991. // fail-open second pass for the proxy stream circuit (#5056): when the only
  1992. // reason no account is available is that every candidate sits behind a
  1993. // quarantined proxy, the quarantine must degrade to a preference instead of
  1994. // zeroing out capacity. The retry re-runs the exact same selection with the
  1995. // quarantine checks bypassed, so healthy proxies always win the first pass
  1996. // and quarantined ones only serve when nothing else can.
  1997. func (s *OpenAIGatewayService) selectAccountWithScheduler(
  1998. ctx context.Context,
  1999. groupID *int64,
  2000. previousResponseID string,
  2001. sessionHash string,
  2002. requestedModel string,
  2003. excludedIDs map[int64]struct{},
  2004. requiredTransport OpenAIUpstreamTransport,
  2005. requiredCapability OpenAIEndpointCapability,
  2006. requiredImageCapability OpenAIImagesCapability,
  2007. requireCompact bool,
  2008. platform string,
  2009. previousResponseCanMove bool,
  2010. useUpstreamTokenCost bool,
  2011. ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
  2012. selection, decision, err := s.selectAccountWithSchedulerOnce(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, requiredImageCapability, requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost)
  2013. if err == nil || openAIProxyStreamQuarantineBypassed(ctx) {
  2014. return selection, decision, err
  2015. }
  2016. if !errors.Is(err, ErrNoAvailableAccounts) && !errors.Is(err, ErrNoAvailableCompactAccounts) {
  2017. return selection, decision, err
  2018. }
  2019. // The circuit only ever quarantines PlatformOpenAI accounts.
  2020. if NormalizeOpenAICompatiblePlatform(platform) != PlatformOpenAI {
  2021. return selection, decision, err
  2022. }
  2023. blocked := s.getOpenAIProxyStreamCircuit().activeBlockCount(time.Now())
  2024. if blocked == 0 {
  2025. return selection, decision, err
  2026. }
  2027. s.logOpenAIProxyStreamQuarantineFailOpen(requestedModel, blocked)
  2028. return s.selectAccountWithSchedulerOnce(withOpenAIProxyStreamQuarantineBypass(ctx), groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, requiredImageCapability, requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost)
  2029. }
  2030. type openAIGroupPrivacyRequirementContextKey struct{}
  2031. type openAIGroupPrivacyRequirement struct {
  2032. groupID int64
  2033. required bool
  2034. }
  2035. func (s *OpenAIGatewayService) withOpenAIGroupPrivacyRequirement(ctx context.Context, groupID *int64) context.Context {
  2036. return context.WithValue(ctx, openAIGroupPrivacyRequirementContextKey{}, openAIGroupPrivacyRequirement{
  2037. groupID: derefGroupID(groupID),
  2038. required: s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID),
  2039. })
  2040. }
  2041. func (s *OpenAIGatewayService) openAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
  2042. if cached, ok := ctx.Value(openAIGroupPrivacyRequirementContextKey{}).(openAIGroupPrivacyRequirement); ok && cached.groupID == derefGroupID(groupID) {
  2043. return cached.required
  2044. }
  2045. return s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID)
  2046. }
  2047. func (s *OpenAIGatewayService) loadOpenAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
  2048. if s == nil || groupID == nil || s.schedulerSnapshot == nil {
  2049. return false
  2050. }
  2051. group, err := s.schedulerSnapshot.GetGroupByID(ctx, *groupID)
  2052. if err != nil {
  2053. return true
  2054. }
  2055. return group != nil && group.RequirePrivacySet
  2056. }
  2057. func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
  2058. ctx context.Context,
  2059. groupID *int64,
  2060. previousResponseID string,
  2061. sessionHash string,
  2062. requestedModel string,
  2063. excludedIDs map[int64]struct{},
  2064. requiredTransport OpenAIUpstreamTransport,
  2065. requiredCapability OpenAIEndpointCapability,
  2066. requiredImageCapability OpenAIImagesCapability,
  2067. requireCompact bool,
  2068. platform string,
  2069. previousResponseCanMove bool,
  2070. useUpstreamTokenCost bool,
  2071. ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
  2072. ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
  2073. ctx = s.withOpenAIGroupPrivacyRequirement(ctx, groupID)
  2074. // 分组利润控制:唯一文本调度入口的防御性装门。handler 文本
  2075. // 入口已在请求开始经 WithOpenAIRequestPricingContext 装门并固定 pricingAt,
  2076. // 此处对同分组门直接复用(failover 重入阈值稳定),仅为不经 handler 装配的
  2077. // 内部调用兜底。图片/视频调度不在利润门范围:requiredImageCapability 非空的
  2078. // Images 调度不装门;其他使用 Responses 能力的文本请求(包括原生远程压缩)
  2079. // 仍须装门。其余媒体路径通过 WithOpenAIProfitControlSuppressed 显式跳过。
  2080. if requiredImageCapability == "" {
  2081. ctx = s.withOpenAIProfitControlGate(ctx, groupID)
  2082. }
  2083. platform = NormalizeOpenAICompatiblePlatform(platform)
  2084. decision := OpenAIAccountScheduleDecision{}
  2085. preserveGuardianParentBinding := preserveOpenAIGuardianParentBinding(ctx, sessionHash)
  2086. guardianParentAccountID := int64(0)
  2087. if strings.TrimSpace(previousResponseID) == "" {
  2088. guardianParentAccountID = s.resolveOpenAIGuardianParentAccountID(ctx, groupID)
  2089. }
  2090. scheduler := s.getOpenAIAccountScheduler(ctx)
  2091. if scheduler == nil {
  2092. decision.Layer = openAIAccountScheduleLayerLoadBalance
  2093. if guardianParentAccountID > 0 {
  2094. if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
  2095. return nil, decision, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
  2096. }
  2097. fallbackScheduler := &defaultOpenAIAccountScheduler{service: s, stats: newOpenAIAccountRuntimeStats()}
  2098. selection, _, err := fallbackScheduler.selectBySessionHash(ctx, OpenAIAccountScheduleRequest{
  2099. GroupID: groupID,
  2100. Platform: platform,
  2101. SessionHash: sessionHash,
  2102. StickyAccountID: guardianParentAccountID,
  2103. PreserveStickyBinding: true,
  2104. RequestedModel: requestedModel,
  2105. RequiredTransport: requiredTransport,
  2106. RequiredCapability: requiredCapability,
  2107. RequiredImageCapability: requiredImageCapability,
  2108. RequireCompact: requireCompact,
  2109. ExcludedIDs: excludedIDs,
  2110. RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
  2111. })
  2112. if err != nil {
  2113. return nil, decision, err
  2114. }
  2115. if selection != nil && selection.Account != nil {
  2116. decision.Layer = openAIAccountScheduleLayerGuardianParent
  2117. decision.StickySessionHit = true
  2118. decision.SelectedAccountID = selection.Account.ID
  2119. decision.SelectedAccountType = selection.Account.Type
  2120. return selection, decision, nil
  2121. }
  2122. }
  2123. legacySessionHash := sessionHash
  2124. if preserveGuardianParentBinding {
  2125. legacySessionHash = ""
  2126. }
  2127. if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
  2128. effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
  2129. for {
  2130. selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, legacySessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
  2131. if err != nil {
  2132. return nil, decision, err
  2133. }
  2134. if selection == nil || selection.Account == nil {
  2135. return selection, decision, nil
  2136. }
  2137. if accountSupportsOpenAICapabilities(selection.Account, requiredCapability, requiredImageCapability) {
  2138. return selection, decision, nil
  2139. }
  2140. if selection.ReleaseFunc != nil {
  2141. selection.ReleaseFunc()
  2142. }
  2143. if effectiveExcludedIDs == nil {
  2144. effectiveExcludedIDs = make(map[int64]struct{})
  2145. }
  2146. if _, exists := effectiveExcludedIDs[selection.Account.ID]; exists {
  2147. return nil, decision, ErrNoAvailableAccounts
  2148. }
  2149. effectiveExcludedIDs[selection.Account.ID] = struct{}{}
  2150. }
  2151. }
  2152. effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
  2153. for {
  2154. selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, legacySessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
  2155. if err != nil {
  2156. return nil, decision, err
  2157. }
  2158. if selection == nil || selection.Account == nil {
  2159. return selection, decision, nil
  2160. }
  2161. if s.isOpenAIAccountTransportCompatible(selection.Account, requiredTransport) &&
  2162. accountSupportsOpenAICapabilities(selection.Account, requiredCapability, requiredImageCapability) {
  2163. return selection, decision, nil
  2164. }
  2165. if selection.ReleaseFunc != nil {
  2166. selection.ReleaseFunc()
  2167. }
  2168. if effectiveExcludedIDs == nil {
  2169. effectiveExcludedIDs = make(map[int64]struct{})
  2170. }
  2171. if _, exists := effectiveExcludedIDs[selection.Account.ID]; exists {
  2172. return nil, decision, ErrNoAvailableAccounts
  2173. }
  2174. effectiveExcludedIDs[selection.Account.ID] = struct{}{}
  2175. }
  2176. }
  2177. if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
  2178. slog.Warn("channel pricing restriction blocked request",
  2179. "group_id", derefGroupID(groupID),
  2180. "model", requestedModel)
  2181. return nil, decision, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
  2182. }
  2183. var stickyAccountID int64
  2184. if sessionHash != "" && s.cache != nil {
  2185. if accountID, err := s.getStickySessionAccountID(ctx, groupID, sessionHash); err == nil && accountID > 0 {
  2186. stickyAccountID = accountID
  2187. }
  2188. }
  2189. stickyWeighted := s.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx)
  2190. subscriptionPriority := s.isOpenAIAdvancedSchedulerSubscriptionPriorityEnabled(ctx)
  2191. stickyPreviousAccountID := int64(0)
  2192. if stickyWeighted && previousResponseCanMove && strings.TrimSpace(previousResponseID) != "" && platform == PlatformOpenAI {
  2193. stickyPreviousAccountID = s.ResolveAccountIDByPreviousResponseIDForScheduler(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact)
  2194. }
  2195. selection, decision, err := scheduler.Select(ctx, OpenAIAccountScheduleRequest{
  2196. GroupID: groupID,
  2197. Platform: platform,
  2198. SessionHash: sessionHash,
  2199. StickyAccountID: stickyAccountID,
  2200. GuardianParentAccountID: guardianParentAccountID,
  2201. StickyPreviousAccountID: stickyPreviousAccountID,
  2202. StickyWeighted: stickyWeighted,
  2203. SubscriptionPriority: subscriptionPriority,
  2204. PreserveStickyBinding: preserveGuardianParentBinding,
  2205. RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
  2206. PreviousResponseID: previousResponseID,
  2207. PreviousResponseCanMove: previousResponseCanMove,
  2208. UseUpstreamTokenCost: useUpstreamTokenCost,
  2209. RequestedModel: requestedModel,
  2210. RequiredTransport: requiredTransport,
  2211. RequiredCapability: requiredCapability,
  2212. RequiredImageCapability: requiredImageCapability,
  2213. RequireCompact: requireCompact,
  2214. ExcludedIDs: excludedIDs,
  2215. })
  2216. if err != nil || selection == nil || selection.Account == nil {
  2217. return selection, decision, err
  2218. }
  2219. // The advanced scheduler has several direct sticky/previous-response
  2220. // branches that construct AccountSelectionResult without going through
  2221. // newSelectionResult. Resolve only the request-time random proxy here: the
  2222. // other branches already re-check the account in the scheduler, and a full
  2223. // snapshot hydration for every non-random request would add an avoidable
  2224. // Redis/DB read on this hot path.
  2225. if selection.Account.IsRandomProxy() {
  2226. if hydrateErr := ResolveRandomProxyFromSource(ctx, selection.Account, s.accountRepo); hydrateErr != nil {
  2227. if selection.ReleaseFunc != nil {
  2228. selection.ReleaseFunc()
  2229. selection.ReleaseFunc = nil
  2230. }
  2231. return nil, decision, hydrateErr
  2232. }
  2233. }
  2234. return selection, decision, nil
  2235. }
  2236. func accountSupportsOpenAICapabilities(account *Account, requiredCapability OpenAIEndpointCapability, requiredImageCapability OpenAIImagesCapability) bool {
  2237. if account == nil {
  2238. return false
  2239. }
  2240. return account.SupportsOpenAIEndpointCapability(requiredCapability) &&
  2241. account.SupportsOpenAIImageCapability(requiredImageCapability)
  2242. }
  2243. func cloneExcludedAccountIDs(excludedIDs map[int64]struct{}) map[int64]struct{} {
  2244. if len(excludedIDs) == 0 {
  2245. return nil
  2246. }
  2247. cloned := make(map[int64]struct{}, len(excludedIDs))
  2248. for id := range excludedIDs {
  2249. cloned[id] = struct{}{}
  2250. }
  2251. return cloned
  2252. }
  2253. func (s *OpenAIGatewayService) isOpenAIAccountTransportCompatible(account *Account, requiredTransport OpenAIUpstreamTransport) bool {
  2254. if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
  2255. return true
  2256. }
  2257. if s == nil || account == nil {
  2258. return false
  2259. }
  2260. if requiredTransport == OpenAIUpstreamTransportResponsesWebsocketV2Ingress {
  2261. if s.cfg == nil || !s.cfg.Gateway.OpenAIWS.ModeRouterV2Enabled {
  2262. return s.getOpenAIWSProtocolResolver().Resolve(account).Transport == OpenAIUpstreamTransportResponsesWebsocketV2
  2263. }
  2264. mode := account.ResolveOpenAIResponsesWebSocketV2Mode(s.cfg.Gateway.OpenAIWS.IngressModeDefault)
  2265. switch mode {
  2266. case OpenAIWSIngressModeCtxPool, OpenAIWSIngressModePassthrough, OpenAIWSIngressModeHTTPBridge, OpenAIWSIngressModeShared, OpenAIWSIngressModeDedicated:
  2267. return true
  2268. default:
  2269. return false
  2270. }
  2271. }
  2272. return s.getOpenAIWSProtocolResolver().Resolve(account).Transport == requiredTransport
  2273. }
  2274. func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(account *Account, model string, success bool, firstTokenMs *int, observedErr ...error) bool {
  2275. if account == nil {
  2276. return false
  2277. }
  2278. accountID := account.ID
  2279. healthTripped := false
  2280. if s != nil && s.rateLimitService != nil {
  2281. if success {
  2282. s.rateLimitService.ObserveOpenAIAPIKeyHealthSuccess(context.Background(), account)
  2283. } else if len(observedErr) > 0 && observedErr[0] != nil {
  2284. healthTripped = s.rateLimitService.ObserveOpenAIAPIKeyHealthFailure(context.Background(), account, observedErr[0])
  2285. }
  2286. }
  2287. if success {
  2288. s.openaiOAuth429RetryStartedAt.Delete(accountID)
  2289. s.clearOpenAIAccountModelTransientState(accountID, normalizeOpenAIAccountModelTransientModel(model))
  2290. }
  2291. scheduler := s.getOpenAIAccountScheduler(context.Background())
  2292. if scheduler == nil {
  2293. return healthTripped
  2294. }
  2295. scheduler.ReportResult(accountID, success, firstTokenMs)
  2296. return healthTripped
  2297. }
  2298. // ObserveOpenAIAccountHealthFailure records failures that cannot reach the
  2299. // scheduler-result path, for example after semantic response bytes were sent.
  2300. func (s *OpenAIGatewayService) ObserveOpenAIAccountHealthFailure(ctx context.Context, account *Account, observedErr error) bool {
  2301. if s == nil || s.rateLimitService == nil || account == nil || observedErr == nil {
  2302. return false
  2303. }
  2304. return s.rateLimitService.ObserveOpenAIAPIKeyHealthFailure(ctx, account, observedErr)
  2305. }
  2306. func (s *OpenAIGatewayService) RecordOpenAIAccountSwitch() {
  2307. scheduler := s.getOpenAIAccountScheduler(context.Background())
  2308. if scheduler == nil {
  2309. return
  2310. }
  2311. scheduler.ReportSwitch()
  2312. }
  2313. func (s *OpenAIGatewayService) SnapshotOpenAIAccountSchedulerMetrics() OpenAIAccountSchedulerMetricsSnapshot {
  2314. scheduler := s.getOpenAIAccountScheduler(context.Background())
  2315. if scheduler == nil {
  2316. return OpenAIAccountSchedulerMetricsSnapshot{}
  2317. }
  2318. return scheduler.SnapshotMetrics()
  2319. }
  2320. func (s *OpenAIGatewayService) openAIWSSessionStickyTTL() time.Duration {
  2321. if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 {
  2322. return time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second
  2323. }
  2324. return openaiStickySessionTTL
  2325. }
  2326. func (s *OpenAIGatewayService) openAIWSLBTopK() int {
  2327. if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.LBTopK > 0 {
  2328. return s.cfg.Gateway.OpenAIWS.LBTopK
  2329. }
  2330. return 7
  2331. }
  2332. func (s *OpenAIGatewayService) openAIWSLBTopKForRequest(ctx context.Context) int {
  2333. base := s.openAIWSLBTopK()
  2334. settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
  2335. // DB 覆盖值与 stickyWeighted/subscriptionPriority 一样受总开关门控:
  2336. // 关闭高级调度器后所有调用方(含管理页分数快照)都应回到配置/默认行为。
  2337. if !settings.enabled {
  2338. return base
  2339. }
  2340. if settings.lbTopKOverride > 0 {
  2341. return settings.lbTopKOverride
  2342. }
  2343. return base
  2344. }
  2345. func (s *OpenAIGatewayService) openAIStickyEscapeConfig() openAIStickyEscapeConfig {
  2346. if s != nil && s.cfg != nil {
  2347. cfg := s.cfg.Gateway.OpenAIScheduler
  2348. enabled := cfg.StickyEscapeEnabled
  2349. if !enabled && cfg.StickyEscapeTTFTMs == 0 && cfg.StickyEscapeErrorRate == 0 {
  2350. enabled = true
  2351. }
  2352. ttftMs := float64(cfg.StickyEscapeTTFTMs)
  2353. if ttftMs <= 0 {
  2354. ttftMs = 15000
  2355. }
  2356. errorRate := cfg.StickyEscapeErrorRate
  2357. if errorRate < 0 || errorRate > 1 {
  2358. errorRate = 0.5
  2359. }
  2360. if errorRate == 0 && cfg.StickyEscapeTTFTMs == 0 && cfg.StickyEscapeErrorRate == 0 {
  2361. errorRate = 0.5
  2362. }
  2363. return openAIStickyEscapeConfig{
  2364. enabled: enabled,
  2365. ttftMs: ttftMs,
  2366. errorRate: errorRate,
  2367. }
  2368. }
  2369. return openAIStickyEscapeConfig{
  2370. enabled: true,
  2371. ttftMs: 15000,
  2372. errorRate: 0.5,
  2373. }
  2374. }
  2375. func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedulerScoreWeightsView {
  2376. if s != nil && s.cfg != nil {
  2377. return GatewayOpenAIWSSchedulerScoreWeightsView{
  2378. Priority: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority,
  2379. Load: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load,
  2380. Queue: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue,
  2381. ErrorRate: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate,
  2382. TTFT: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT,
  2383. Reset: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Reset,
  2384. QuotaHeadroom: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom,
  2385. UpstreamCost: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.UpstreamCost,
  2386. Previous: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.PreviousResponse,
  2387. SessionSticky: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky,
  2388. }
  2389. }
  2390. return GatewayOpenAIWSSchedulerScoreWeightsView{
  2391. Priority: 1.0,
  2392. Load: 1.0,
  2393. Queue: 0.7,
  2394. ErrorRate: 0.8,
  2395. TTFT: 0.5,
  2396. Reset: 0.0,
  2397. QuotaHeadroom: 0.0,
  2398. UpstreamCost: 0.0,
  2399. Previous: 5.0,
  2400. SessionSticky: 3.0,
  2401. }
  2402. }
  2403. func (s *OpenAIGatewayService) openAIWSSchedulerWeightsForRequest(ctx context.Context) GatewayOpenAIWSSchedulerScoreWeightsView {
  2404. weights := s.openAIWSSchedulerWeights()
  2405. settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
  2406. // 同 openAIWSLBTopKForRequest:总开关关闭时不应用 DB 覆盖值。
  2407. if !settings.enabled {
  2408. return weights
  2409. }
  2410. overridden := applyOpenAIAdvancedSchedulerWeightOverrides(weights, settings.weightOverrides)
  2411. if !overridden.configWeights().IsValid() {
  2412. return weights
  2413. }
  2414. return overridden
  2415. }
  2416. func applyOpenAIAdvancedSchedulerWeightOverrides(
  2417. weights GatewayOpenAIWSSchedulerScoreWeightsView,
  2418. overrides map[string]float64,
  2419. ) GatewayOpenAIWSSchedulerScoreWeightsView {
  2420. for key, value := range overrides {
  2421. switch key {
  2422. case "priority":
  2423. weights.Priority = value
  2424. case "load":
  2425. weights.Load = value
  2426. case "queue":
  2427. weights.Queue = value
  2428. case "error_rate":
  2429. weights.ErrorRate = value
  2430. case "ttft":
  2431. weights.TTFT = value
  2432. case "reset":
  2433. weights.Reset = value
  2434. case "quota_headroom":
  2435. weights.QuotaHeadroom = value
  2436. case "upstream_cost":
  2437. weights.UpstreamCost = value
  2438. case "previous_response":
  2439. weights.Previous = value
  2440. case "session_sticky":
  2441. weights.SessionSticky = value
  2442. }
  2443. }
  2444. return weights
  2445. }
  2446. type GatewayOpenAIWSSchedulerScoreWeightsView struct {
  2447. Priority float64
  2448. Load float64
  2449. Queue float64
  2450. ErrorRate float64
  2451. TTFT float64
  2452. // Reset 倾向「会话窗口最早重置」的账号;0 表示关闭(默认)。
  2453. Reset float64
  2454. QuotaHeadroom float64
  2455. UpstreamCost float64
  2456. Previous float64
  2457. SessionSticky float64
  2458. }
  2459. func (w GatewayOpenAIWSSchedulerScoreWeightsView) configWeights() config.GatewayOpenAIWSSchedulerScoreWeights {
  2460. return config.GatewayOpenAIWSSchedulerScoreWeights{
  2461. Priority: w.Priority,
  2462. Load: w.Load,
  2463. Queue: w.Queue,
  2464. ErrorRate: w.ErrorRate,
  2465. TTFT: w.TTFT,
  2466. Reset: w.Reset,
  2467. QuotaHeadroom: w.QuotaHeadroom,
  2468. UpstreamCost: w.UpstreamCost,
  2469. PreviousResponse: w.Previous,
  2470. SessionSticky: w.SessionSticky,
  2471. }
  2472. }
  2473. type OpenAIAccountSchedulerScoreSnapshot struct {
  2474. BaseScore float64
  2475. StickyScore float64
  2476. StickyScoreInfinity bool
  2477. StickyWeightedEnabled bool
  2478. }
  2479. func (s *RateLimitService) BuildOpenAIAccountSchedulerScoreSnapshot(
  2480. ctx context.Context,
  2481. accounts []*Account,
  2482. loadMap map[int64]*AccountLoadInfo,
  2483. ) map[int64]OpenAIAccountSchedulerScoreSnapshot {
  2484. gateway := &OpenAIGatewayService{cfg: nil, rateLimitService: s}
  2485. if s != nil {
  2486. gateway.cfg = s.cfg
  2487. }
  2488. return buildOpenAIAccountSchedulerScoreSnapshot(
  2489. accounts,
  2490. loadMap,
  2491. gateway.openAIWSSchedulerWeightsForRequest(ctx),
  2492. gateway.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx),
  2493. gateway.openAIOAuthSchedulingRateMultiplier(ctx),
  2494. )
  2495. }
  2496. func BuildOpenAIAccountSchedulerScoreSnapshot(
  2497. accounts []*Account,
  2498. loadMap map[int64]*AccountLoadInfo,
  2499. ) map[int64]OpenAIAccountSchedulerScoreSnapshot {
  2500. gateway := &OpenAIGatewayService{}
  2501. return buildOpenAIAccountSchedulerScoreSnapshot(accounts, loadMap, gateway.openAIWSSchedulerWeights(), false, defaultOpenAIOAuthSchedulingRateMultiplier)
  2502. }
  2503. func buildOpenAIAccountSchedulerScoreSnapshot(
  2504. accounts []*Account,
  2505. loadMap map[int64]*AccountLoadInfo,
  2506. weights GatewayOpenAIWSSchedulerScoreWeightsView,
  2507. stickyWeightedEnabled bool,
  2508. oauthSchedulingRateMultiplier float64,
  2509. ) map[int64]OpenAIAccountSchedulerScoreSnapshot {
  2510. if len(accounts) == 0 {
  2511. return nil
  2512. }
  2513. candidates := make([]openAIAccountCandidateScore, 0, len(accounts))
  2514. for _, account := range accounts {
  2515. if account == nil {
  2516. continue
  2517. }
  2518. loadInfo := loadMap[account.ID]
  2519. if loadInfo == nil {
  2520. loadInfo = &AccountLoadInfo{AccountID: account.ID}
  2521. }
  2522. candidates = append(candidates, openAIAccountCandidateScore{
  2523. account: account,
  2524. loadInfo: loadInfo,
  2525. errorRate: 0,
  2526. ttft: 0,
  2527. hasTTFT: false,
  2528. })
  2529. }
  2530. if len(candidates) == 0 {
  2531. return nil
  2532. }
  2533. minPriority, maxPriority := openAIAccountSchedulingPriority(candidates[0].account), openAIAccountSchedulingPriority(candidates[0].account)
  2534. maxWaiting := 1
  2535. for i := range candidates {
  2536. candidate := &candidates[i]
  2537. candidate.priority = openAIAccountSchedulingPriority(candidate.account)
  2538. if candidate.priority < minPriority {
  2539. minPriority = candidate.priority
  2540. }
  2541. if candidate.priority > maxPriority {
  2542. maxPriority = candidate.priority
  2543. }
  2544. if candidate.loadInfo.WaitingCount > maxWaiting {
  2545. maxWaiting = candidate.loadInfo.WaitingCount
  2546. }
  2547. }
  2548. minResetRemaining, maxResetRemaining := 0.0, 0.0
  2549. hasResetSample := false
  2550. now := time.Now()
  2551. upstreamCostFactors := map[int64]float64(nil)
  2552. if weights.UpstreamCost > 0 {
  2553. accounts := make([]*Account, 0, len(candidates))
  2554. for _, candidate := range candidates {
  2555. accounts = append(accounts, candidate.account)
  2556. }
  2557. upstreamCostFactors = openAIUpstreamCostFactors(accounts, now, oauthSchedulingRateMultiplier)
  2558. }
  2559. if weights.Reset > 0 {
  2560. for _, candidate := range candidates {
  2561. end := candidate.account.SessionWindowEnd
  2562. if end == nil || !now.Before(*end) {
  2563. continue
  2564. }
  2565. remaining := end.Sub(now).Seconds()
  2566. if !hasResetSample {
  2567. minResetRemaining, maxResetRemaining = remaining, remaining
  2568. hasResetSample = true
  2569. continue
  2570. }
  2571. if remaining < minResetRemaining {
  2572. minResetRemaining = remaining
  2573. }
  2574. if remaining > maxResetRemaining {
  2575. maxResetRemaining = remaining
  2576. }
  2577. }
  2578. }
  2579. result := make(map[int64]OpenAIAccountSchedulerScoreSnapshot, len(candidates))
  2580. for _, candidate := range candidates {
  2581. priorityFactor := 1.0
  2582. if maxPriority > minPriority {
  2583. priorityFactor = 1 - float64(candidate.priority-minPriority)/float64(maxPriority-minPriority)
  2584. }
  2585. loadFactor := 1 - clamp01(float64(candidate.loadInfo.LoadRate)/100.0)
  2586. queueFactor := 1 - clamp01(float64(candidate.loadInfo.WaitingCount)/float64(maxWaiting))
  2587. errorFactor := 1.0
  2588. ttftFactor := 0.5
  2589. resetFactor := 0.0
  2590. if weights.Reset > 0 && hasResetSample {
  2591. if end := candidate.account.SessionWindowEnd; end != nil && now.Before(*end) {
  2592. if maxResetRemaining > minResetRemaining {
  2593. resetFactor = 1 - clamp01((end.Sub(now).Seconds()-minResetRemaining)/(maxResetRemaining-minResetRemaining))
  2594. } else {
  2595. resetFactor = 1
  2596. }
  2597. }
  2598. }
  2599. quotaHeadroomFactor := 0.0
  2600. if weights.QuotaHeadroom > 0 {
  2601. quotaHeadroomFactor = openAIQuotaHeadroomFactor(candidate.account, now)
  2602. }
  2603. upstreamCostFactor := openAIUpstreamCostNeutralFactor
  2604. if factor, ok := upstreamCostFactors[candidate.account.ID]; ok {
  2605. upstreamCostFactor = factor
  2606. }
  2607. baseScore := weights.Priority*priorityFactor +
  2608. weights.Load*loadFactor +
  2609. weights.Queue*queueFactor +
  2610. weights.ErrorRate*errorFactor +
  2611. weights.TTFT*ttftFactor +
  2612. weights.Reset*resetFactor +
  2613. weights.QuotaHeadroom*quotaHeadroomFactor +
  2614. weights.UpstreamCost*(upstreamCostFactor-openAIUpstreamCostNeutralFactor)
  2615. score := OpenAIAccountSchedulerScoreSnapshot{
  2616. BaseScore: baseScore,
  2617. StickyWeightedEnabled: stickyWeightedEnabled,
  2618. StickyScoreInfinity: !stickyWeightedEnabled,
  2619. }
  2620. if stickyWeightedEnabled {
  2621. score.StickyScore = baseScore + weights.Previous + weights.SessionSticky
  2622. }
  2623. result[candidate.account.ID] = score
  2624. }
  2625. return result
  2626. }
  2627. func openAIUpstreamCostFactors(accounts []*Account, now time.Time, oauthSchedulingRateMultiplier float64) map[int64]float64 {
  2628. type rateSample struct {
  2629. accountID int64
  2630. rate float64
  2631. }
  2632. factors := make(map[int64]float64, len(accounts))
  2633. samples := make([]rateSample, 0, len(accounts))
  2634. eligibleCount := 0
  2635. for _, account := range accounts {
  2636. if account == nil {
  2637. continue
  2638. }
  2639. factors[account.ID] = openAIUpstreamCostNeutralFactor
  2640. if !account.IsOpenAIApiKey() && !account.IsOpenAIOAuthLike() {
  2641. continue
  2642. }
  2643. eligibleCount++
  2644. if rate, ok := openAISchedulingRate(account, now, oauthSchedulingRateMultiplier); ok {
  2645. samples = append(samples, rateSample{accountID: account.ID, rate: rate})
  2646. }
  2647. }
  2648. if len(samples) < 2 || eligibleCount == 0 {
  2649. return factors
  2650. }
  2651. allEqual := true
  2652. positiveLogs := make([]float64, 0, len(samples))
  2653. for i, sample := range samples {
  2654. if i > 0 && sample.rate != samples[0].rate {
  2655. allEqual = false
  2656. }
  2657. if sample.rate > 0 {
  2658. positiveLogs = append(positiveLogs, math.Log(sample.rate))
  2659. }
  2660. }
  2661. if allEqual || len(positiveLogs) == 0 {
  2662. return factors
  2663. }
  2664. sort.Float64s(positiveLogs)
  2665. middle := len(positiveLogs) / 2
  2666. medianLog := positiveLogs[middle]
  2667. if len(positiveLogs)%2 == 0 {
  2668. medianLog = (positiveLogs[middle-1] + positiveLogs[middle]) / 2
  2669. }
  2670. center := math.Exp(medianLog)
  2671. if center <= 0 || math.IsNaN(center) || math.IsInf(center, 0) {
  2672. return factors
  2673. }
  2674. coverage := float64(len(samples)) / float64(eligibleCount)
  2675. for _, sample := range samples {
  2676. rawFactor := 1.0
  2677. if sample.rate > 0 {
  2678. rawFactor = 1 / (1 + sample.rate/center)
  2679. }
  2680. factors[sample.accountID] = clamp01(openAIUpstreamCostNeutralFactor + coverage*(rawFactor-openAIUpstreamCostNeutralFactor))
  2681. }
  2682. return factors
  2683. }
  2684. type openAILegacyUpstreamRateOrder struct {
  2685. enabled bool
  2686. rates map[int64]float64
  2687. }
  2688. func newOpenAILegacyUpstreamRateOrder(accounts []*Account, now time.Time, oauthSchedulingRateMultiplier float64) openAILegacyUpstreamRateOrder {
  2689. rates := make(map[int64]float64, len(accounts))
  2690. var first float64
  2691. distinct := false
  2692. for _, account := range accounts {
  2693. if account == nil {
  2694. continue
  2695. }
  2696. // 与 openAIUpstreamCostFactors 使用同一道平台门控:只有 OpenAI 平台账号
  2697. // 的倍率参与 legacy 低倍率优先排序。上游自报倍率来自中转方,不能让它对
  2698. // 其他平台的调度产生影响——否则自报低价即可吸走流量,而实际结算走本地倍率。
  2699. if !account.IsOpenAIApiKey() && !account.IsOpenAIOAuthLike() {
  2700. continue
  2701. }
  2702. rate, ok := openAISchedulingRate(account, now, oauthSchedulingRateMultiplier)
  2703. if !ok {
  2704. continue
  2705. }
  2706. if len(rates) == 0 {
  2707. first = rate
  2708. } else if rate != first {
  2709. distinct = true
  2710. }
  2711. rates[account.ID] = rate
  2712. }
  2713. return openAILegacyUpstreamRateOrder{enabled: len(rates) >= 2 && distinct, rates: rates}
  2714. }
  2715. func openAISchedulingRate(account *Account, now time.Time, oauthSchedulingRateMultiplier float64) (float64, bool) {
  2716. if account != nil && account.IsOpenAIOAuthLike() {
  2717. return oauthSchedulingRateMultiplier, true
  2718. }
  2719. return openAIFreshUpstreamBillingRate(account, now)
  2720. }
  2721. // compare returns -1 when a should be selected before b, 1 when b should be
  2722. // selected first, and 0 when the rate signal does not distinguish them.
  2723. func (o openAILegacyUpstreamRateOrder) compare(a, b *Account) int {
  2724. if !o.enabled || a == nil || b == nil {
  2725. return 0
  2726. }
  2727. aRate, aKnown := o.rates[a.ID]
  2728. bRate, bKnown := o.rates[b.ID]
  2729. if aKnown != bKnown {
  2730. if aKnown {
  2731. return -1
  2732. }
  2733. return 1
  2734. }
  2735. if !aKnown || aRate == bRate {
  2736. return 0
  2737. }
  2738. if aRate < bRate {
  2739. return -1
  2740. }
  2741. return 1
  2742. }
  2743. func openAIFreshUpstreamBillingRate(account *Account, now time.Time) (float64, bool) {
  2744. if !isUpstreamBillingProbeAccount(account) {
  2745. return 0, false
  2746. }
  2747. snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra)
  2748. if snapshot == nil || (snapshot.Status != UpstreamBillingProbeStatusOK && snapshot.Status != UpstreamBillingProbeStatusFailed) ||
  2749. snapshot.ReceivedAt == nil || snapshot.ReceivedAt.IsZero() {
  2750. return 0, false
  2751. }
  2752. receivedAt := *snapshot.ReceivedAt
  2753. freshUntil := snapshot.FreshUntil
  2754. if freshUntil == nil && snapshot.Status == UpstreamBillingProbeStatusOK {
  2755. interval := snapshot.NextProbeAt.Sub(receivedAt)
  2756. if interval > 0 {
  2757. freshUntil = probeTimePtr(receivedAt.Add(2 * interval))
  2758. }
  2759. }
  2760. if freshUntil == nil || !freshUntil.After(receivedAt) || now.Before(receivedAt) || now.After(*freshUntil) {
  2761. return 0, false
  2762. }
  2763. return upstreamBillingRateAt(snapshot.Data, now)
  2764. }
  2765. func openAIQuotaHeadroomFactor(account *Account, now time.Time) float64 {
  2766. if account == nil || len(account.Extra) == 0 || openAIQuotaHeadroomSnapshotStale(account.Extra, now) {
  2767. return openAIQuotaHeadroomNeutralFactor
  2768. }
  2769. primaryUsedPercent, ok := resolveAccountExtraNumber(account.Extra, "codex_primary_used_percent", "codex_7d_used_percent")
  2770. if !ok || openAIQuotaWindowResetAny(account.Extra, now, "primary", "7d") {
  2771. return openAIQuotaHeadroomNeutralFactor
  2772. }
  2773. factor := 1 - clamp01(primaryUsedPercent/100)
  2774. if secondaryUsedPercent, ok := resolveAccountExtraNumber(account.Extra, "codex_secondary_used_percent", "codex_5h_used_percent"); ok &&
  2775. !openAIQuotaWindowResetAny(account.Extra, now, "secondary", "5h") {
  2776. secondaryRemaining := 1 - clamp01(secondaryUsedPercent/100)
  2777. if secondaryRemaining < openAIQuotaHeadroomSecondaryLowRemain {
  2778. factor *= openAIQuotaHeadroomNeutralFactor
  2779. }
  2780. }
  2781. return factor
  2782. }
  2783. func openAIQuotaHeadroomSnapshotStale(extra map[string]any, now time.Time) bool {
  2784. updatedRaw, ok := extra["codex_usage_updated_at"]
  2785. if !ok {
  2786. return true
  2787. }
  2788. updatedAt, err := parseTime(fmt.Sprint(updatedRaw))
  2789. if err != nil {
  2790. return true
  2791. }
  2792. return now.Sub(updatedAt) >= openAIQuotaHeadroomSnapshotStaleAfter
  2793. }
  2794. func openAIQuotaWindowResetAny(extra map[string]any, now time.Time, windows ...string) bool {
  2795. for _, window := range windows {
  2796. if openAIQuotaWindowReset(extra, window, now) {
  2797. return true
  2798. }
  2799. }
  2800. return false
  2801. }
  2802. func clamp01(value float64) float64 {
  2803. switch {
  2804. case value < 0:
  2805. return 0
  2806. case value > 1:
  2807. return 1
  2808. default:
  2809. return value
  2810. }
  2811. }
  2812. func calcLoadSkewByMoments(sum float64, sumSquares float64, count int) float64 {
  2813. if count <= 1 {
  2814. return 0
  2815. }
  2816. mean := sum / float64(count)
  2817. variance := sumSquares/float64(count) - mean*mean
  2818. if variance < 0 {
  2819. variance = 0
  2820. }
  2821. return math.Sqrt(variance)
  2822. }