mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 01:17:46 +08:00
feat: improve routing policy diagnostics
This commit is contained in:
fawney19
@@ -688,8 +688,18 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
return None;
|
||||
}
|
||||
};
|
||||
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);
|
||||
let materialized_row_count = rows.len() as u32;
|
||||
let missing_score_count = scores.len().saturating_sub(rows.len());
|
||||
if missing_score_count > 0 {
|
||||
*self
|
||||
.skip_reason_counts
|
||||
.entry("pool_score_member_missing")
|
||||
.or_insert(0) += u32::try_from(missing_score_count).unwrap_or(u32::MAX);
|
||||
}
|
||||
self.scanned_keys = self.scanned_keys.saturating_add(materialized_row_count);
|
||||
self.budget_scanned_keys = self
|
||||
.budget_scanned_keys
|
||||
.saturating_add(materialized_row_count);
|
||||
Some(self.build_page_eligible_candidates(rows).await)
|
||||
}
|
||||
|
||||
@@ -1044,7 +1054,10 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
}
|
||||
|
||||
fn pool_skip_reason_releases_scan_budget(skip_reason: &str) -> bool {
|
||||
skip_reason == POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
|
||||
matches!(
|
||||
skip_reason,
|
||||
POOL_ACCOUNT_EXHAUSTED_SKIP_REASON | POOL_ACCOUNT_BLOCKED_SKIP_REASON
|
||||
)
|
||||
}
|
||||
|
||||
fn pool_candidate_transport_policy_facts(
|
||||
@@ -1587,7 +1600,8 @@ mod tests {
|
||||
ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate,
|
||||
apply_local_runtime_candidate_terminal_reason, provider_key_pool_score_id,
|
||||
provider_key_pool_score_scope, EligibleLocalExecutionCandidate,
|
||||
LocalExecutionCandidateKind, PlannerAppState,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
@@ -1597,10 +1611,14 @@ mod tests {
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
|
||||
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
||||
};
|
||||
use aether_data_contracts::repository::pool_scores::{
|
||||
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeStatus, StoredPoolMemberScore,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
@@ -3082,6 +3100,131 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_does_not_spend_effective_scan_budget_on_blocked_accounts() {
|
||||
const BLOCKED_COUNT: usize = 1_600;
|
||||
let provider_config = Some(json!({ "pool_advanced": {} }));
|
||||
let (provider, endpoint, mut keys, rows) =
|
||||
large_pool_fixture(BLOCKED_COUNT + 100, provider_config.clone());
|
||||
for key in keys.iter_mut().take(BLOCKED_COUNT) {
|
||||
key.oauth_invalid_reason = Some("blocked account".to_string());
|
||||
}
|
||||
|
||||
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 >= u32::try_from(BLOCKED_COUNT + 1).unwrap(),
|
||||
"default absolute scan cap should allow scanning past a large blocked prefix"
|
||||
);
|
||||
|
||||
let candidate = cursor
|
||||
.next_key()
|
||||
.await
|
||||
.expect("cursor should scan past blocked 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 >= BLOCKED_COUNT,
|
||||
"cursor should not return one of the blocked leading keys"
|
||||
);
|
||||
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
|
||||
assert!(
|
||||
cursor.budget_scanned_keys <= aether_dispatch_core::DEFAULT_POOL_PAGE_SIZE,
|
||||
"blocked accounts should not consume effective scan budget"
|
||||
);
|
||||
assert_eq!(
|
||||
cursor
|
||||
.skip_reason_counts
|
||||
.get(aether_pool_core::POOL_ACCOUNT_BLOCKED_SKIP_REASON),
|
||||
Some(&(BLOCKED_COUNT as u32))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_does_not_spend_scan_budget_on_missing_score_rows() {
|
||||
let provider_config = Some(json!({
|
||||
"pool_advanced": {
|
||||
"score_top_n": 128
|
||||
}
|
||||
}));
|
||||
let (provider, endpoint, keys, rows) = large_pool_fixture(1, provider_config.clone());
|
||||
let scores = (0..128)
|
||||
.map(|index| {
|
||||
sample_provider_key_pool_score(
|
||||
"provider-pool",
|
||||
&format!("missing-key-{index:03}"),
|
||||
1_000.0 - index as f64,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
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_pool_score_repository_for_tests(Arc::new(
|
||||
InMemoryPoolMemberScoreRepository::seed(scores),
|
||||
))
|
||||
.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);
|
||||
let candidate = cursor
|
||||
.next_key()
|
||||
.await
|
||||
.expect("cursor should fall back to catalog rows after stale scores");
|
||||
|
||||
assert_eq!(candidate.candidate.key_id, "key-00000");
|
||||
assert_eq!(cursor.scanned_keys, 1);
|
||||
assert_eq!(cursor.budget_scanned_keys, 1);
|
||||
assert_eq!(
|
||||
cursor.skip_reason_counts.get("pool_score_member_missing"),
|
||||
Some(&128)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_scheduler_skips_invalid_and_exhausted_high_priority_hot_pool_before_fallback_provider(
|
||||
) {
|
||||
@@ -3766,6 +3909,40 @@ mod tests {
|
||||
(provider, endpoint, keys, rows)
|
||||
}
|
||||
|
||||
fn sample_provider_key_pool_score(
|
||||
provider_id: &str,
|
||||
key_id: &str,
|
||||
score: f64,
|
||||
) -> StoredPoolMemberScore {
|
||||
let identity = PoolMemberIdentity::provider_api_key(provider_id, key_id);
|
||||
let scope = provider_key_pool_score_scope();
|
||||
StoredPoolMemberScore {
|
||||
id: provider_key_pool_score_id(&identity, &scope),
|
||||
pool_kind: identity.pool_kind,
|
||||
pool_id: identity.pool_id,
|
||||
member_kind: identity.member_kind,
|
||||
member_id: identity.member_id,
|
||||
capability: scope.capability,
|
||||
scope_kind: scope.scope_kind,
|
||||
scope_id: scope.scope_id,
|
||||
score,
|
||||
hard_state: PoolMemberHardState::Available,
|
||||
score_version: 1,
|
||||
score_reason: json!({}),
|
||||
last_ranked_at: Some(1_000),
|
||||
last_scheduled_at: None,
|
||||
last_success_at: None,
|
||||
last_failure_at: None,
|
||||
failure_count: 0,
|
||||
last_probe_attempt_at: None,
|
||||
last_probe_success_at: None,
|
||||
last_probe_failure_at: None,
|
||||
probe_failure_count: 0,
|
||||
probe_status: PoolMemberProbeStatus::Ok,
|
||||
updated_at: 1_000,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_codex_pool_provider(
|
||||
provider_id: &str,
|
||||
provider_priority: i32,
|
||||
|
||||
Reference in New Issue
Block a user