i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor pool candidate scheduling

Browse Source
This commit is contained in:
fawney19
2026-05-03 20:14:29 +08:00
parent 8ebee9922c
commit a24e4a793d
55 changed files with 4825 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
apps/aether-gateway/src/scheduler/candidate/ranking.rs
View File

@@ -26,9 +26,16 @@ pub(super) fn rank_scheduler_candidates(
.iter()
.enumerate()
.map(|(index, candidate)| {
let provider_key = runtime_snapshot
.provider_key_rpm_states
.get(&candidate.key_id);
let pool_group = runtime_snapshot
.pool_provider_ids
.contains(candidate.provider_id.as_str());
let provider_key = (!pool_group)
.then(|| {
runtime_snapshot
.provider_key_rpm_states
.get(&candidate.key_id)
})
.flatten();
SchedulerRankableCandidate::from_candidate(candidate, index)
.with_capability_priority(requested_capability_priority_for_candidate(
required_capabilities,

110
apps/aether-gateway/src/scheduler/candidate/runtime.rs
View File

@@ -22,6 +22,7 @@ pub(super) struct CandidateRuntimeSelectionSnapshot {
pub(super) recent_candidates: Vec<StoredRequestCandidate>,
pub(super) provider_concurrent_limits: BTreeMap<String, usize>,
pub(super) provider_key_rpm_states: BTreeMap<String, StoredProviderCatalogKey>,
pub(super) pool_provider_ids: BTreeSet<String>,
provider_quota_blocks_requests: BTreeMap<String, bool>,
key_account_quota_exhausted: BTreeMap<String, bool>,
key_oauth_invalid: BTreeMap<String, bool>,
@@ -35,8 +36,15 @@ pub(super) async fn read_candidate_runtime_selection_snapshot(
) -> Result<CandidateRuntimeSelectionSnapshot, GatewayError> {
let recent_candidates = state.read_recent_request_candidates(128).await?;
let provider_concurrent_limits = read_provider_concurrent_limits(state, candidates).await?;
let provider_skip_exhausted_accounts =
read_provider_skip_exhausted_account_map(state, candidates).await?;
let provider_pool_state = read_provider_pool_state_map(state, candidates).await?;
let provider_skip_exhausted_accounts = provider_pool_state
.iter()
.map(|(provider_id, state)| (provider_id.clone(), state.skip_exhausted_accounts))
.collect::<BTreeMap<_, _>>();
let pool_provider_ids = provider_pool_state
.iter()
.filter_map(|(provider_id, state)| state.pool_enabled.then_some(provider_id.clone()))
.collect::<BTreeSet<_>>();
let provider_key_rpm_states = read_provider_key_rpm_states(state, candidates).await?;
let key_account_quota_exhausted = read_key_account_quota_exhaustion_map(
candidates,
@@ -54,6 +62,7 @@ pub(super) async fn read_candidate_runtime_selection_snapshot(
recent_candidates,
provider_concurrent_limits,
provider_key_rpm_states,
pool_provider_ids,
provider_quota_blocks_requests,
key_account_quota_exhausted,
key_oauth_invalid,
@@ -93,6 +102,9 @@ pub(super) fn is_candidate_selectable(
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
) -> bool {
let pool_group = snapshot
.pool_provider_ids
.contains(candidate.provider_id.as_str());
candidate_is_selectable_with_runtime_state(CandidateRuntimeSelectabilityInput {
candidate,
recent_candidates: &snapshot.recent_candidates,
@@ -105,20 +117,26 @@ pub(super) fn is_candidate_selectable(
.get(candidate.provider_id.as_str())
.copied()
.unwrap_or(false),
account_quota_exhausted: snapshot
.key_account_quota_exhausted
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
oauth_invalid: snapshot
.key_oauth_invalid
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
rpm_reset_at: snapshot
.provider_key_rpm_reset_ats
.get(candidate.key_id.as_str())
.copied()
account_quota_exhausted: !pool_group
&& snapshot
.key_account_quota_exhausted
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
oauth_invalid: !pool_group
&& snapshot
.key_oauth_invalid
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
rpm_reset_at: (!pool_group)
.then(|| {
snapshot
.provider_key_rpm_reset_ats
.get(candidate.key_id.as_str())
.copied()
.flatten()
})
.flatten(),
})
}
@@ -129,15 +147,22 @@ pub(super) fn current_candidate_runtime_skip_reason(
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
) -> Option<&'static str> {
let pool_group = snapshot
.pool_provider_ids
.contains(candidate.provider_id.as_str());
let provider_quota_blocks_requests = snapshot
.provider_quota_blocks_requests
.get(candidate.provider_id.as_str())
.copied()
.unwrap_or(false);
let rpm_reset_at = snapshot
.provider_key_rpm_reset_ats
.get(candidate.key_id.as_str())
.copied()
let rpm_reset_at = (!pool_group)
.then(|| {
snapshot
.provider_key_rpm_reset_ats
.get(candidate.key_id.as_str())
.copied()
.flatten()
})
.flatten();
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
@@ -148,16 +173,18 @@ pub(super) fn current_candidate_runtime_skip_reason(
now_unix_secs,
cached_affinity_target,
provider_quota_blocks_requests,
account_quota_exhausted: snapshot
.key_account_quota_exhausted
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
oauth_invalid: snapshot
.key_oauth_invalid
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
account_quota_exhausted: !pool_group
&& snapshot
.key_account_quota_exhausted
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
oauth_invalid: !pool_group
&& snapshot
.key_oauth_invalid
.get(candidate.key_id.as_str())
.copied()
.unwrap_or(false),
rpm_reset_at,
})
}
@@ -228,10 +255,16 @@ async fn read_provider_quota_block_map(
Ok(quota_blocks)
}
async fn read_provider_skip_exhausted_account_map(
#[derive(Debug, Clone, Copy, Default)]
struct ProviderPoolState {
pool_enabled: bool,
skip_exhausted_accounts: bool,
}
async fn read_provider_pool_state_map(
state: &(impl SchedulerRuntimeState + ?Sized),
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
) -> Result<BTreeMap<String, bool>, GatewayError> {
) -> Result<BTreeMap<String, ProviderPoolState>, GatewayError> {
let provider_ids = candidates
.iter()
.map(|candidate| candidate.provider_id.clone())
@@ -248,15 +281,22 @@ async fn read_provider_skip_exhausted_account_map(
Ok(providers
.into_iter()
.map(|provider| {
let skip_exhausted_accounts = provider
let pool_advanced = provider
.config
.as_ref()
.and_then(|value| value.get("pool_advanced"))
.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(serde_json::Value::as_bool)
.unwrap_or(false);
(provider.id, skip_exhausted_accounts)
(
provider.id,
ProviderPoolState {
pool_enabled: pool_advanced.is_some(),
skip_exhausted_accounts,
},
)
})
.collect())
}

26
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -1638,11 +1638,14 @@ async fn skips_codex_candidate_when_account_quota_is_exhausted_and_pool_flag_ena
.await
.expect("selection should succeed");
assert_eq!(selected.len(), 1);
assert_eq!(selected[0].provider_id, "provider-openai");
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.provider_id, "provider-codex");
assert_eq!(skipped[0].skip_reason, "account_quota_exhausted");
assert_eq!(selected.len(), 2);
assert!(selected
.iter()
.any(|item| item.provider_id == "provider-codex"));
assert!(selected
.iter()
.any(|item| item.provider_id == "provider-openai"));
assert!(skipped.is_empty());
}
#[tokio::test]
@@ -2240,11 +2243,14 @@ async fn skips_kiro_candidate_when_account_quota_is_exhausted_and_pool_flag_enab
.await
.expect("selection should succeed");
assert_eq!(selected.len(), 1);
assert_eq!(selected[0].provider_id, "provider-openai");
assert_eq!(skipped.len(), 1);
assert_eq!(skipped[0].candidate.provider_id, "provider-kiro");
assert_eq!(skipped[0].skip_reason, "account_quota_exhausted");
assert_eq!(selected.len(), 2);
assert!(selected
.iter()
.any(|item| item.provider_id == "provider-kiro"));
assert!(selected
.iter()
.any(|item| item.provider_id == "provider-openai"));
assert!(skipped.is_empty());
}
#[tokio::test]