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

fix: preserve pool scan budget for exhausted accounts

Browse Source
This commit is contained in:
fawney19
2026-05-26 23:49:26 +08:00
parent 4933ae9014
commit e21cd72181
1 changed files with 120 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+120 -5
View File
@@ -343,9 +343,11 @@ pub(crate) struct PoolKeyCursor<'a> {
pool_key_order: StoredPoolKeyCandidateOrder,
next_offset: u32,
scanned_keys: u32,
budget_scanned_keys: u32,
window_size: u32,
page_size: u32,
max_scanned_keys: u32,
absolute_max_scanned_keys: u32,
score_top_n: u32,
score_phase_loaded: bool,
skip_reason_counts: BTreeMap<&'static str, u32>,
@@ -393,13 +395,14 @@ impl<'a> PoolKeyCursor<'a> {
.map(|config| config.score_top_n)
.unwrap_or(u64::from(aether_dispatch_core::DEFAULT_POOL_PAGE_SIZE))
.clamp(1, u64::from(u32::MAX)) as u32;
let max_scanned_keys = pool_config
let configured_max_scanned_keys = pool_config
.as_ref()
.map(|config| config.score_fallback_scan_limit)
.unwrap_or(u64::from(aether_dispatch_core::DEFAULT_POOL_MAX_SCAN))
.clamp(1, u64::from(u32::MAX)) as u32;
let window_config = crate::dispatch::pool::default_pool_window_config().normalized();
let max_scanned_keys = max_scanned_keys.min(window_config.max_scan);
let max_scanned_keys = configured_max_scanned_keys.min(window_config.max_scan);
let absolute_max_scanned_keys = configured_max_scanned_keys.max(max_scanned_keys);
Self {
state,
group,
@@ -412,9 +415,11 @@ impl<'a> PoolKeyCursor<'a> {
pool_key_order,
next_offset: 0,
scanned_keys: 0,
budget_scanned_keys: 0,
window_size: window_config.window_size,
page_size: window_config.page_size,
max_scanned_keys: max_scanned_keys.max(window_config.window_size),
absolute_max_scanned_keys: absolute_max_scanned_keys.max(window_config.window_size),
score_top_n,
score_phase_loaded: false,
skip_reason_counts: BTreeMap::new(),
@@ -486,6 +491,7 @@ impl<'a> PoolKeyCursor<'a> {
extra_data: Some(serde_json::json!({
"pool_group_exhaustion": {
"scanned_keys": self.scanned_keys,
"budget_scanned_keys": self.budget_scanned_keys,
"skip_reason_counts": skip_reason_counts,
}
})),
@@ -504,6 +510,9 @@ impl<'a> PoolKeyCursor<'a> {
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
scanned_keys = self.scanned_keys,
budget_scanned_keys = self.budget_scanned_keys,
max_scanned_keys = self.max_scanned_keys,
absolute_max_scanned_keys = self.absolute_max_scanned_keys,
skip_reason_counts = ?self.skip_reason_counts,
"gateway pool scheduler exhausted pool group without a schedulable key"
);
@@ -548,13 +557,16 @@ impl<'a> PoolKeyCursor<'a> {
}
}
if self.scanned_keys >= self.max_scanned_keys {
if self.budget_scanned_keys >= self.max_scanned_keys
|| self.scanned_keys >= self.absolute_max_scanned_keys
{
return None;
}
let limit = self
.page_size
.min(self.max_scanned_keys - self.scanned_keys);
.min(self.max_scanned_keys - self.budget_scanned_keys)
.min(self.absolute_max_scanned_keys - self.scanned_keys);
let query = StoredPoolKeyCandidateRowsQuery {
api_format: self.group.candidate.endpoint_api_format.clone(),
provider_id: self.group.candidate.provider_id.clone(),
@@ -593,11 +605,21 @@ impl<'a> PoolKeyCursor<'a> {
}
self.scanned_keys += rows.len() as u32;
self.budget_scanned_keys += rows.len() as u32;
self.next_offset = self.next_offset.saturating_add(rows.len() as u32);
Some(self.build_page_eligible_candidates(rows).await)
}
async fn next_score_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
if self.scanned_keys >= self.absolute_max_scanned_keys {
return None;
}
let limit = self
.score_top_n
.min(self.absolute_max_scanned_keys - self.scanned_keys);
if limit == 0 {
return None;
}
let scope = provider_key_pool_score_scope();
let query = ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
@@ -608,7 +630,7 @@ impl<'a> PoolKeyCursor<'a> {
hard_states: vec![PoolMemberHardState::Available, PoolMemberHardState::Unknown],
probe_statuses: None,
offset: 0,
limit: self.score_top_n as usize,
limit: limit as usize,
};
let scores = match self.state.app().data.list_ranked_pool_members(&query).await {
Ok(scores) => scores,
@@ -667,6 +689,7 @@ impl<'a> PoolKeyCursor<'a> {
}
};
self.scanned_keys = self.scanned_keys.saturating_add(scores.len() as u32);
self.budget_scanned_keys = self.budget_scanned_keys.saturating_add(scores.len() as u32);
Some(self.build_page_eligible_candidates(rows).await)
}
@@ -1010,9 +1033,20 @@ impl<'a> PoolKeyCursor<'a> {
for skipped_candidate in skipped_candidates {
self.record_skip_reason(skipped_candidate.skip_reason);
}
let prefiltered_count = skipped_candidates
.iter()
.filter(|candidate| pool_skip_reason_releases_scan_budget(candidate.skip_reason))
.count();
self.budget_scanned_keys = self
.budget_scanned_keys
.saturating_sub(u32::try_from(prefiltered_count).unwrap_or(u32::MAX));
}
}
fn pool_skip_reason_releases_scan_budget(skip_reason: &str) -> bool {
skip_reason == POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
}
fn pool_candidate_transport_policy_facts(
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
) -> CandidateTransportPolicyFacts<'_> {
@@ -2967,6 +3001,87 @@ mod tests {
}));
}
#[tokio::test]
async fn pool_key_cursor_does_not_spend_effective_scan_budget_on_exhausted_accounts() {
let provider_config = Some(json!({
"pool_advanced": {
"skip_exhausted_accounts": true
}
}));
let (provider, endpoint, mut keys, rows) = large_pool_fixture(700, provider_config.clone());
for key in keys.iter_mut().take(600) {
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "openai",
"exhausted": true,
"usage_ratio": 1.0,
"windows": [
{
"code": "daily",
"used_ratio": 1.0,
"remaining_ratio": 0.0
}
]
}
}));
}
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
keys,
)),
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
)
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config,
);
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
assert_eq!(
cursor.max_scanned_keys,
aether_dispatch_core::DEFAULT_POOL_MAX_SCAN
);
assert!(
cursor.absolute_max_scanned_keys > cursor.max_scanned_keys,
"pool config scan limit should be retained as the absolute cap"
);
let candidate = cursor
.next_key()
.await
.expect("cursor should scan past exhausted accounts within the absolute cap");
let key_index = candidate
.candidate
.key_id
.strip_prefix("key-")
.and_then(|value| value.parse::<usize>().ok())
.expect("fixture key id should contain a numeric suffix");
assert!(
key_index >= 600,
"cursor should not return one of the exhausted leading keys"
);
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert_eq!(cursor.scanned_keys, 640);
assert_eq!(cursor.budget_scanned_keys, 40);
assert_eq!(
cursor
.skip_reason_counts
.get(aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON),
Some(&600)
);
}
#[tokio::test]
async fn pool_scheduler_skips_invalid_and_exhausted_high_priority_hot_pool_before_fallback_provider(
) {