fix(pool): honor quota exhaustion toggle and default ignore to off

This commit is contained in:
elky
2026-10-10 10:04:41 +08:00
parent 29ba6c8cb1
commit 6bafef672f
8 changed files with 99 additions and 37 deletions
@@ -326,7 +326,9 @@ fn active_probe_member_is_unschedulable_for_request(
return true;
}
key_context.is_some_and(|context| {
context.account_blocked || context.quota_exhausted || context.quota_hard_blocked
context.account_blocked
|| (pool_config.skip_exhausted_accounts && context.quota_exhausted)
|| context.quota_hard_blocked
})
}
@@ -1529,12 +1531,14 @@ async fn read_pool_catalog_key_contexts_by_id(
provider_type,
provider_model_name,
);
context.quota_exhausted |= reserve_minimum_quota_key_ids.contains(&key.id)
let reserve_reached = reserve_minimum_quota_key_ids.contains(&key.id)
&& admin_provider_pool_pure::admin_pool_key_minimum_quota_reached(
&key,
provider_type,
provider_model_name,
);
context.quota_exhausted |= reserve_reached;
context.quota_hard_blocked |= reserve_reached;
(key.id.clone(), context)
})
.collect::<BTreeMap<_, _>>();
@@ -2303,6 +2307,31 @@ mod tests {
);
}
#[test]
fn pool_scheduler_allows_quota_exhausted_key_when_ignored() {
let exhausted = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-exhausted",
10,
Some(json!({ "pool_advanced": { "ignore_exhausted_accounts": true } })),
);
let contexts = BTreeMap::from([(
"key-exhausted".to_string(),
PoolCatalogKeyContext {
quota_exhausted: true,
..PoolCatalogKeyContext::default()
},
)]);
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![exhausted],
&BTreeMap::new(),
&contexts,
);
assert_eq!(scheduled.len(), 1);
assert!(skipped.is_empty());
}
#[test]
fn pool_scheduler_skips_quota_exhausted_key_when_flag_is_false() {
let ready = sample_eligible_candidate(
@@ -410,7 +410,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
}],
unschedulable_rules: Vec::new(),
lru_enabled: false,
skip_exhausted_accounts: false,
skip_exhausted_accounts: true,
reserve_minimum_quota: false,
sticky_session_ttl_seconds: 3600,
latency_window_seconds: 3600,
@@ -444,9 +444,10 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
scheduling_presets,
unschedulable_rules,
skip_exhausted_accounts: pool_advanced
.get("skip_exhausted_accounts")
.get("ignore_exhausted_accounts")
.and_then(Value::as_bool)
.unwrap_or(false),
.map(|ignore| !ignore)
.unwrap_or(true),
reserve_minimum_quota: pool_advanced
.get("reserve_minimum_quota")
.and_then(Value::as_bool)
@@ -574,11 +575,25 @@ mod tests {
}
#[test]
fn defaults_skip_exhausted_accounts_to_false() {
fn parses_ignore_exhausted_accounts_with_default_off() {
for ignore in [false, true] {
let provider = sample_provider(json!({
"pool_advanced": {
"ignore_exhausted_accounts": ignore,
"skip_exhausted_accounts": true
}
}));
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
assert_eq!(config.skip_exhausted_accounts, !ignore);
}
}
#[test]
fn defaults_skip_exhausted_accounts_to_true() {
let provider = sample_provider(json!({ "pool_advanced": {} }));
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
assert!(!config.skip_exhausted_accounts);
assert!(config.skip_exhausted_accounts);
assert!(!config.reserve_minimum_quota);
}
@@ -593,7 +608,7 @@ mod tests {
}));
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
assert_eq!(config.reserve_minimum_quota, enabled);
assert!(!config.skip_exhausted_accounts);
assert!(config.skip_exhausted_accounts);
}
}
@@ -385,9 +385,10 @@ async fn read_provider_pool_state_map(
.and_then(|value| value.get("pool_advanced"));
let skip_exhausted_accounts = pool_advanced
.and_then(serde_json::Value::as_object)
.and_then(|value| value.get("skip_exhausted_accounts"))
.and_then(|value| value.get("ignore_exhausted_accounts"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
.map(|ignore| !ignore)
.unwrap_or(true);
let reserve_minimum_quota = pool_advanced
.and_then(serde_json::Value::as_object)
.and_then(|value| value.get("reserve_minimum_quota"))