starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(core): normalize account scheduling load

Browse Source
This commit is contained in:
Zyy
2026-08-14 23:27:06 +08:00
parent 4117840f75
commit b8ea1cd076
2 changed files with 101 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-core/src/engine/credential.rs
+62 -22
View File
@@ -1588,14 +1588,26 @@ impl AccountSelector {
let candidate = match context.policy.strategy() {
RotationStrategy::QuotaResetPriority => {
eligible.sort_by_key(|candidate| {
let default_concurrency = context.policy.max_concurrent_per_account();
eligible.sort_by(|left, right| {
(
candidate.signals.quota_reset_at.is_none(),
candidate.signals.quota_reset_at,
candidate.signals.in_flight,
candidate.signals.last_started_at,
candidate.account.id().clone(),
left.signals.quota_reset_at.is_none(),
left.signals.quota_reset_at,
)
.cmp(&(
right.signals.quota_reset_at.is_none(),
right.signals.quota_reset_at,
))
.then_with(|| {
capacity_utilization(left, default_concurrency)
.total_cmp(&capacity_utilization(right, default_concurrency))
})
.then_with(|| {
left.signals
.last_started_at
.cmp(&right.signals.last_started_at)
})
.then_with(|| left.account.id().cmp(right.account.id()))
});
eligible.first().copied()?
}
@@ -1604,9 +1616,11 @@ impl AccountSelector {
let index = context.round_robin_cursor as usize % eligible.len();
eligible.get(index).copied()?
}
RotationStrategy::Smart => {
select_smart_candidate(&eligible, context.round_robin_cursor)?
}
RotationStrategy::Smart => select_smart_candidate(
&eligible,
context.policy.max_concurrent_per_account(),
context.round_robin_cursor,
)?,
RotationStrategy::Sticky => {
eligible.sort_by_key(|candidate| {
(
@@ -1676,8 +1690,18 @@ const SMART_QUOTA_WEIGHT: f64 = 0.8;
const SMART_FAILURE_WEIGHT: f64 = 1.0;
const SMART_LATENCY_WEIGHT: f64 = 0.5;
fn capacity_utilization(candidate: &AccountCandidate, default_concurrency: NonZeroU32) -> f64 {
f64::from(candidate.signals.in_flight)
/ f64::from(
candidate
.account
.effective_concurrency(default_concurrency)
.get(),
)
}
struct SmartNormalization {
max_in_flight: u32,
max_capacity_utilization: f64,
min_quota: Option<u64>,
max_quota: Option<u64>,
min_latency_ms: Option<u64>,
@@ -1685,11 +1709,11 @@ struct SmartNormalization {
}
impl SmartNormalization {
fn from_candidates(candidates: &[&AccountCandidate]) -> Self {
let max_in_flight = candidates
fn from_candidates(candidates: &[&AccountCandidate], default_concurrency: NonZeroU32) -> Self {
let max_capacity_utilization = candidates
.iter()
.map(|candidate| candidate.signals.in_flight)
.max()
.map(|candidate| capacity_utilization(candidate, default_concurrency))
.reduce(f64::max)
.unwrap_or_default();
let min_quota = candidates
.iter()
@@ -1710,7 +1734,7 @@ impl SmartNormalization {
.filter(|latency| *latency > 0)
.max();
Self {
max_in_flight,
max_capacity_utilization,
min_quota,
max_quota,
min_latency_ms,
@@ -1721,12 +1745,18 @@ impl SmartNormalization {
fn select_smart_candidate<'a>(
candidates: &[&'a AccountCandidate],
default_concurrency: NonZeroU32,
cursor: u64,
) -> Option<&'a AccountCandidate> {
let normalization = SmartNormalization::from_candidates(candidates);
let normalization = SmartNormalization::from_candidates(candidates, default_concurrency);
let mut ranked = candidates
.iter()
.map(|candidate| (*candidate, smart_score(candidate, &normalization)))
.map(|candidate| {
(
*candidate,
smart_score(candidate, default_concurrency, &normalization),
)
})
.collect::<Vec<_>>();
ranked.sort_by(|(left, left_score), (right, right_score)| {
right_score
@@ -1742,11 +1772,14 @@ fn select_smart_candidate<'a>(
Some(ranked[index].0)
}
fn smart_score(candidate: &AccountCandidate, normalization: &SmartNormalization) -> f64 {
let load = lower_is_better(
u64::from(candidate.signals.in_flight),
0,
u64::from(normalization.max_in_flight),
fn smart_score(
candidate: &AccountCandidate,
default_concurrency: NonZeroU32,
normalization: &SmartNormalization,
) -> f64 {
let load = lower_ratio_is_better(
capacity_utilization(candidate, default_concurrency),
normalization.max_capacity_utilization,
);
let quota = candidate.signals.quota_remaining_rank.map_or(0.5, |quota| {
higher_is_better(
@@ -1781,6 +1814,13 @@ fn smart_score(candidate: &AccountCandidate, normalization: &SmartNormalization)
+ SMART_LATENCY_WEIGHT * latency
}
fn lower_ratio_is_better(value: f64, maximum: f64) -> f64 {
if maximum <= 0.0 {
return 1.0;
}
1.0 - (value / maximum).clamp(0.0, 1.0)
}
fn lower_is_better(value: u64, minimum: u64, maximum: u64) -> f64 {
if maximum <= minimum {
return 1.0;
backend/crates/gateway-core/tests/engine/credential.rs
+39
View File
@@ -51,6 +51,15 @@ fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandida
}
}
fn candidate_with_concurrency(id: &str, in_flight: u32, concurrency: u32) -> AccountCandidate {
let mut candidate = candidate(id, in_flight, Some(100));
candidate.account = candidate.account.with_scheduling(
Some(AccountConcurrencyLimit::new(concurrency).expect("concurrency override")),
AccountWeight::DEFAULT,
);
candidate
}
fn weighted_candidate(id: &str, weight: u16, in_flight: u32) -> AccountCandidate {
let mut candidate = candidate(id, in_flight, None);
candidate.account = candidate.account.with_scheduling(
@@ -348,6 +357,20 @@ fn smart_selector_should_prefer_lower_inflight_count() {
assert_eq!(selected.candidate().account.id().as_str(), "acct_idle");
}
#[test]
fn smart_selector_should_prefer_lower_capacity_utilization_over_lower_inflight_count() {
let candidates = [
candidate_with_concurrency("acct_smaller", 1, 2),
candidate_with_concurrency("acct_larger", 2, 10),
];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::Smart))
.expect("candidate available");
assert_eq!(selected.candidate().account.id().as_str(), "acct_larger");
}
#[test]
fn selector_should_prioritize_the_highest_weight_for_every_rotation_strategy() {
let candidates = vec![
@@ -622,6 +645,22 @@ fn quota_reset_selector_should_prefer_known_earliest_window() {
assert_eq!(selected.candidate().account.id().as_str(), "acct_earlier");
}
#[test]
fn quota_reset_selector_should_use_capacity_utilization_as_load_tiebreaker() {
let reset_at = SystemTime::now() + Duration::from_secs(60);
let mut smaller = candidate_with_concurrency("acct_smaller", 1, 2);
smaller.signals.quota_reset_at = Some(reset_at);
let mut larger = candidate_with_concurrency("acct_larger", 2, 10);
larger.signals.quota_reset_at = Some(reset_at);
let candidates = [smaller, larger];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::QuotaResetPriority))
.expect("candidate available");
assert_eq!(selected.candidate().account.id().as_str(), "acct_larger");
}
#[test]
fn smart_selector_should_prefer_known_quota_over_unknown_after_load_is_equal() {
let candidates = vec![