mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-12 04:09:48 +08:00
fix(pool): honor quota exhaustion toggle and default ignore to off
This commit is contained in:
@@ -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"))
|
||||
|
||||
Reference in New Issue
Block a user