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

Merge branch 'codex/pool-hot-trace-fix'

Browse Source 
# Conflicts:
#	apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
This commit is contained in:
fawney19
2026-05-19 14:47:37 +08:00
parent 7e9ca88e00 7d569127ae
commit 63202a63ef
4 changed files with 786 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

226
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
View File

@@ -81,6 +81,7 @@ enum LocalExecutionCandidateAttemptSourceItem<'a> {
cursor: PoolKeyCursor<'a>,
candidate_index: u32,
pending_attempts: DispatchSequence<LocalExecutionCandidateAttempt>,
pool_exhaustion_persistence: Option<PoolGroupExhaustionPersistenceContext>,
},
RequestedModelPage {
cursor: Box<RequestedModelAttemptPageCursor<'a>>,
@@ -117,11 +118,20 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
cursor,
candidate_index,
pending_attempts,
pool_exhaustion_persistence,
} => {
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
let Some(candidate) = cursor.next_key().await else {
if let Some(skipped) = cursor.exhausted_group_skipped_candidate() {
persist_pool_group_exhaustion_skipped_candidate(
pool_exhaustion_persistence.as_ref(),
*candidate_index,
skipped,
)
.await;
}
cursor.log_exhausted();
let _ = cursor.take_skipped_candidates();
self.items.pop_front();
@@ -186,6 +196,39 @@ pub(crate) struct LocalSkippedCandidatePersistenceContext<'a> {
pub(crate) record_runtime_miss_diagnostic: bool,
}
#[derive(Clone)]
struct PoolGroupExhaustionPersistenceContext {
app: AppState,
trace_id: String,
user_id: String,
api_key_id: String,
required_capabilities: Option<Value>,
error_context: &'static str,
client_api_format: String,
routing_policy: Option<ResolvedRoutingPolicy>,
}
impl PoolGroupExhaustionPersistenceContext {
fn new(
app: AppState,
trace_id: &str,
context: LocalSkippedCandidatePersistenceContext<'_>,
client_api_format: &str,
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> Self {
Self {
app,
trace_id: trace_id.to_string(),
user_id: context.user_id.to_string(),
api_key_id: context.api_key_id.to_string(),
required_capabilities: context.required_capabilities.cloned(),
error_context: context.error_context,
client_api_format: client_api_format.to_string(),
routing_policy: routing_policy.cloned(),
}
}
}
pub(crate) use aether_ai_serving::AiCandidateResolutionMode as LocalCandidateResolutionMode;
struct GatewayLocalCandidateMaterializationPort<'a, F, G> {
@@ -571,6 +614,13 @@ where
requested_model,
request_auth_channel,
routing_policy,
Some(PoolGroupExhaustionPersistenceContext::new(
state.app().clone(),
trace_id,
persistence_policy.skipped,
client_api_format,
routing_policy,
)),
);
(
@@ -589,6 +639,7 @@ fn build_logical_candidate_items<'a>(
requested_model: Option<&str>,
request_auth_channel: Option<&str>,
routing_policy: Option<&ResolvedRoutingPolicy>,
pool_exhaustion_persistence: Option<PoolGroupExhaustionPersistenceContext>,
) -> (VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>, u32) {
let mut items = VecDeque::new();
let mut next_candidate_index = starting_candidate_index;
@@ -625,6 +676,7 @@ fn build_logical_candidate_items<'a>(
cursor,
candidate_index,
pending_attempts: DispatchSequence::new(Vec::new()),
pool_exhaustion_persistence: pool_exhaustion_persistence.clone(),
});
}
}
@@ -832,6 +884,16 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
Some(&self.requested_model),
self.request_auth_channel.as_deref(),
self.routing_policy.as_ref(),
Some(PoolGroupExhaustionPersistenceContext {
app: self.state.app().clone(),
trace_id: self.trace_id.clone(),
user_id: self.skipped_user_id.clone(),
api_key_id: self.skipped_api_key_id.clone(),
required_capabilities: self.skipped_required_capabilities.clone(),
error_context: self.skipped_error_context,
client_api_format: self.client_api_format.clone(),
routing_policy: self.routing_policy.clone(),
}),
);
self.next_candidate_index = next_candidate_index
.saturating_add(u32::try_from(skipped_candidate_count).unwrap_or(u32::MAX));
@@ -948,11 +1010,20 @@ async fn pop_attempt_from_items(
cursor,
candidate_index,
pending_attempts,
pool_exhaustion_persistence,
} => {
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
let Some(candidate) = cursor.next_key().await else {
if let Some(skipped) = cursor.exhausted_group_skipped_candidate() {
persist_pool_group_exhaustion_skipped_candidate(
pool_exhaustion_persistence.as_ref(),
*candidate_index,
skipped,
)
.await;
}
cursor.log_exhausted();
let _ = cursor.take_skipped_candidates();
items.pop_front();
@@ -1130,7 +1201,29 @@ where
}
let _ = cursor.take_skipped_candidates();
if attempts.len() == attempt_count_before_pool {
let skipped = cursor.exhausted_group_skipped_candidate();
cursor.log_exhausted();
if let Some(skipped) = skipped {
let pool_exhaustion_context = PoolGroupExhaustionPersistenceContext::new(
state.app().clone(),
trace_id,
LocalSkippedCandidatePersistenceContext {
user_id: context.user_id,
api_key_id: context.api_key_id,
required_capabilities: context.required_capabilities,
error_context: context.error_context,
record_runtime_miss_diagnostic: false,
},
client_api_format,
routing_policy,
);
persist_pool_group_exhaustion_skipped_candidate(
Some(&pool_exhaustion_context),
candidate_index,
skipped,
)
.await;
}
}
}
}
@@ -1466,6 +1559,40 @@ fn build_unpersisted_local_execution_candidate_attempts(
attempts
}
async fn persist_pool_group_exhaustion_skipped_candidate(
context: Option<&PoolGroupExhaustionPersistenceContext>,
candidate_index: u32,
skipped: SkippedLocalExecutionCandidate,
) {
let Some(context) = context else {
return;
};
let skipped = attach_routing_trace_to_skipped_candidate(
context.routing_policy.as_ref(),
&context.client_api_format,
candidate_index,
skipped,
);
let extra_data =
ai_candidate_extra_data_with_ranking(skipped.extra_data.clone(), skipped.ranking.as_ref());
let candidate_id = Uuid::new_v4().to_string();
persist_skipped_local_execution_candidate(
&context.app,
&context.trace_id,
&context.user_id,
&context.api_key_id,
&skipped.candidate,
candidate_index,
candidate_id.as_str(),
context.required_capabilities.as_ref(),
skipped.skip_reason,
extra_data,
context.error_context,
false,
)
.await;
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn persist_skipped_local_execution_candidate(
state: &AppState,
@@ -1640,6 +1767,7 @@ mod tests {
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
@@ -1949,11 +2077,26 @@ mod tests {
.read_request_candidates_by_request_id("trace-logical-pool")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 1);
assert_eq!(stored.len(), 2);
assert_eq!(stored[0].key_id.as_deref(), Some("pool-group"));
assert_eq!(stored[0].status, RequestCandidateStatus::Skipped);
assert_eq!(
stored[0].skip_reason.as_deref(),
Some("pool_group_exhausted")
);
assert_eq!(stored[0].candidate_index, 0);
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("pool_group_exhaustion"))
.and_then(|value| value.get("scanned_keys")),
Some(&json!(0))
);
assert_eq!(stored[1].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[1].candidate_index, 1);
assert_eq!(
stored[1]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
@@ -2078,6 +2221,83 @@ mod tests {
assert!(source.next_attempt().await.is_none());
}
#[tokio::test]
async fn dynamic_pool_exhaustion_persists_group_skip_summary() {
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
Arc::new(InMemoryAuthApiKeySnapshotRepository::default()),
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::default()),
Arc::new(InMemoryProviderCatalogReadRepository::seed(
Vec::new(),
Vec::new(),
Vec::new(),
)),
Arc::clone(&request_candidate_repository),
"test-encryption-key",
),
);
let mut pool_group = sample_eligible("pool-group", None);
pool_group.kind = LocalExecutionCandidateKind::PoolGroup;
pool_group.transport = sample_transport(
"pool-group",
Some(json!({ "pool_advanced": { "scheduling_presets": [] } })),
);
let cursor = PoolKeyCursor::new(
PlannerAppState::new(&app),
pool_group,
None,
Some("gpt-5"),
None,
);
let pool_exhaustion_persistence = PoolGroupExhaustionPersistenceContext::new(
app.clone(),
"trace-dynamic-pool",
LocalSkippedCandidatePersistenceContext {
user_id: "user-1",
api_key_id: "api-key-1",
required_capabilities: None,
error_context: "persist should not fail",
record_runtime_miss_diagnostic: false,
},
"openai:chat",
None,
);
let mut source = LocalExecutionCandidateAttemptSource {
items: VecDeque::from([LocalExecutionCandidateAttemptSourceItem::Pool {
cursor,
candidate_index: 0,
pending_attempts: DispatchSequence::new(Vec::new()),
pool_exhaustion_persistence: Some(pool_exhaustion_persistence),
}]),
};
assert!(source.next_attempt().await.is_none());
let stored = app
.read_request_candidates_by_request_id("trace-dynamic-pool")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].status, RequestCandidateStatus::Skipped);
assert_eq!(
stored[0].skip_reason.as_deref(),
Some("pool_group_exhausted")
);
assert_eq!(stored[0].candidate_index, 0);
assert_eq!(stored[0].key_id.as_deref(), Some("pool-group"));
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("pool_group_exhaustion"))
.and_then(|value| value.get("scanned_keys")),
Some(&json!(0))
);
}
#[tokio::test]
async fn pool_internal_skipped_candidates_are_not_persisted() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());

485
apps/aether-gateway/src/dispatch/pool_scheduler.rs
View File

@@ -14,6 +14,8 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
use aether_pool_core::{
run_pool_scheduler, PoolCandidateFacts, PoolCandidateInput, PoolCandidateOrchestration,
PoolMemberSignals, PoolRuntimeState, PoolSchedulingConfig, PoolSchedulingPreset,
POOL_ACCOUNT_BLOCKED_SKIP_REASON, POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
POOL_COOLDOWN_SKIP_REASON, POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
use aether_provider_pool::ProviderPoolService;
use aether_routing_core::{RankingOverlay, ResolvedRoutingPolicy};
@@ -32,6 +34,7 @@ use crate::handlers::shared::provider_pool::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
};
use crate::handlers::shared::provider_pool::{
admin_provider_pool_quota_probe_active_members_key,
read_admin_provider_pool_key_cooldown_reason, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState,
};
@@ -112,6 +115,7 @@ async fn schedule_pool_page_candidates(
let key_context_by_id = read_pool_catalog_key_contexts_by_id(state, &candidates).await;
let mut runtime_by_provider = BTreeMap::new();
let mut pool_config_by_provider = BTreeMap::new();
let mut burst_provider_ids = BTreeSet::<String>::new();
for (provider_id, (pool_config, key_ids)) in provider_runtime_requirements {
let key_ids = key_ids.into_iter().collect::<Vec<_>>();
@@ -127,17 +131,45 @@ async fn schedule_pool_page_candidates(
)
.await
};
if should_trigger_active_probe_burst_for_request(&pool_config, &runtime) {
burst_provider_ids.insert(provider_id.clone());
}
pool_config_by_provider.insert(provider_id.clone(), pool_config);
runtime_by_provider.insert(provider_id, runtime);
}
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
let preflight_evictions = prune_unschedulable_active_probe_members_for_request(
&mut runtime_by_provider,
&candidates,
&key_context_by_id,
);
spawn_active_probe_member_evictions_for_request(state, &preflight_evictions);
burst_provider_ids.extend(preflight_evictions.keys().cloned());
for (provider_id, pool_config) in &pool_config_by_provider {
let Some(runtime) = runtime_by_provider.get(provider_id) else {
continue;
};
if should_trigger_active_probe_burst_for_request(pool_config, runtime) {
burst_provider_ids.insert(provider_id.clone());
}
}
let outcome = apply_local_execution_pool_scheduler_with_runtime_map_outcome(
candidates,
&runtime_by_provider,
&key_context_by_id,
);
let scheduled = outcome.candidates;
let skipped = outcome.skipped;
burst_provider_ids.extend(outcome.active_probe_seal_fallback_provider_ids);
spawn_active_probe_member_evictions_for_request(
state,
&outcome.active_probe_evicted_members_by_provider,
);
burst_provider_ids.extend(
outcome
.active_probe_evicted_members_by_provider
.keys()
.cloned(),
);
for skipped_candidate in &skipped {
if skipped_candidate.skip_reason == POOL_ACTIVE_PROBE_SEALED_SKIP_REASON {
@@ -152,6 +184,113 @@ async fn schedule_pool_page_candidates(
(scheduled, skipped)
}
async fn remove_active_probe_members_for_request(
state: PlannerAppState<'_>,
evicted_members_by_provider: &BTreeMap<String, BTreeSet<String>>,
) {
remove_active_probe_members(state.app().clone(), evicted_members_by_provider).await;
}
fn spawn_active_probe_member_evictions_for_request(
state: PlannerAppState<'_>,
evicted_members_by_provider: &BTreeMap<String, BTreeSet<String>>,
) {
if evicted_members_by_provider.is_empty() {
return;
}
let app = state.app().clone();
let evicted_members_by_provider = evicted_members_by_provider.clone();
tokio::spawn(async move {
remove_active_probe_members(app, &evicted_members_by_provider).await;
});
}
async fn remove_active_probe_members(
app: crate::AppState,
evicted_members_by_provider: &BTreeMap<String, BTreeSet<String>>,
) {
for (provider_id, key_ids) in evicted_members_by_provider {
let set_key = admin_provider_pool_quota_probe_active_members_key(provider_id);
for key_id in key_ids {
if let Err(err) = app
.runtime_state
.as_ref()
.set_remove(&set_key, key_id)
.await
{
warn!(
event_name = "pool_active_probe_member_evict_failed",
log_type = "event",
provider_id,
key_id,
error = ?err,
"gateway pool scheduler failed to evict unschedulable active probe member"
);
}
}
}
}
fn prune_unschedulable_active_probe_members_for_request(
runtime_by_provider: &mut BTreeMap<String, AdminProviderPoolRuntimeState>,
candidates: &[EligibleLocalExecutionCandidate],
key_context_by_id: &BTreeMap<String, PoolCatalogKeyContext>,
) -> BTreeMap<String, BTreeSet<String>> {
let mut evicted = BTreeMap::<String, BTreeSet<String>>::new();
for candidate in candidates {
let Some(pool_config) = pool_config_for_candidate(candidate) else {
continue;
};
if !should_enforce_active_probe_sealed_pool(&pool_config) {
continue;
}
let provider_id = candidate.candidate.provider_id.as_str();
let key_id = candidate.candidate.key_id.as_str();
let Some(runtime) = runtime_by_provider.get_mut(provider_id) else {
continue;
};
if !runtime.active_probe_member_ids.contains(key_id) {
continue;
}
if !active_probe_member_is_unschedulable_for_request(
&pool_config,
runtime,
key_id,
key_context_by_id.get(key_id),
) {
continue;
}
runtime.active_probe_member_ids.remove(key_id);
evicted
.entry(provider_id.to_string())
.or_default()
.insert(key_id.to_string());
}
evicted
}
fn active_probe_member_is_unschedulable_for_request(
pool_config: &AdminProviderPoolConfig,
runtime: &AdminProviderPoolRuntimeState,
key_id: &str,
key_context: Option<&PoolCatalogKeyContext>,
) -> bool {
if runtime.cooldown_reason_by_key.contains_key(key_id) {
return true;
}
if pool_config.cost_limit_per_key_tokens.is_some_and(|limit| {
runtime
.cost_window_usage_by_key
.get(key_id)
.copied()
.unwrap_or(0)
>= limit
}) {
return true;
}
key_context.is_some_and(|context| context.account_blocked || context.quota_exhausted)
}
async fn expand_pool_group_candidate(
state: PlannerAppState<'_>,
group: EligibleLocalExecutionCandidate,
@@ -332,6 +471,32 @@ impl<'a> PoolKeyCursor<'a> {
std::mem::take(&mut self.skipped_candidates)
}
pub(crate) fn exhausted_group_skipped_candidate(
&self,
) -> Option<SkippedLocalExecutionCandidate> {
if self.returned_key_count > 0 {
return None;
}
let skip_reason_counts = self
.skip_reason_counts
.iter()
.map(|(reason, count)| ((*reason).to_string(), serde_json::json!(count)))
.collect::<serde_json::Map<String, serde_json::Value>>();
Some(SkippedLocalExecutionCandidate {
candidate: self.group.candidate.clone(),
skip_reason: self.runtime_miss_pool_exhaustion_skip_reason(),
transport: Some(self.group.transport.clone()),
ranking: self.group.ranking.clone(),
extra_data: Some(serde_json::json!({
"pool_group_exhaustion": {
"scanned_keys": self.scanned_keys,
"skip_reason_counts": skip_reason_counts,
}
})),
})
}
pub(crate) fn log_exhausted(&mut self) {
if self.exhausted_logged {
return;
@@ -645,6 +810,7 @@ impl<'a> PoolKeyCursor<'a> {
ranking: candidate.ranking.clone(),
extra_data: None,
});
self.spawn_active_probe_member_eviction_and_replenish(candidate);
true
}
Ok(None) => false,
@@ -664,6 +830,31 @@ impl<'a> PoolKeyCursor<'a> {
}
}
fn spawn_active_probe_member_eviction_and_replenish(
&self,
candidate: &EligibleLocalExecutionCandidate,
) {
let Some(config) = pool_config_for_candidate(candidate) else {
return;
};
if !should_enforce_active_probe_sealed_pool(&config) {
return;
}
let provider_id = candidate.candidate.provider_id.as_str();
let key_id = candidate.candidate.key_id.as_str();
spawn_active_probe_member_evictions_for_request(
self.state,
&BTreeMap::from([(
provider_id.to_string(),
BTreeSet::from([key_id.to_string()]),
)]),
);
let _ = spawn_pool_quota_probe_replenish_for_request(
self.state.app().clone(),
provider_id.to_string(),
);
}
async fn record_score_schedule_interest(&self, scores: &[StoredPoolMemberScore]) {
if scores.is_empty() {
return;
@@ -926,6 +1117,125 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let outcome = apply_local_execution_pool_scheduler_with_runtime_map_outcome(
candidates,
runtime_by_provider,
key_context_by_id,
);
(outcome.candidates, outcome.skipped)
}
struct PoolSchedulerApplyOutcome {
candidates: Vec<EligibleLocalExecutionCandidate>,
skipped: Vec<SkippedLocalExecutionCandidate>,
active_probe_seal_fallback_provider_ids: BTreeSet<String>,
active_probe_evicted_members_by_provider: BTreeMap<String, BTreeSet<String>>,
}
fn apply_local_execution_pool_scheduler_with_runtime_map_outcome(
candidates: Vec<EligibleLocalExecutionCandidate>,
runtime_by_provider: &BTreeMap<String, AdminProviderPoolRuntimeState>,
key_context_by_id: &BTreeMap<String, PoolCatalogKeyContext>,
) -> PoolSchedulerApplyOutcome {
let (scheduled, skipped) = run_local_execution_pool_scheduler_with_runtime_map(
candidates.clone(),
runtime_by_provider,
key_context_by_id,
true,
);
let mut active_probe_evicted_members_by_provider =
active_probe_evicted_members_from_skipped(&skipped, runtime_by_provider);
let active_probe_seal_fallback_provider_ids = if scheduled.is_empty() {
skipped
.iter()
.filter(|skipped| skipped.skip_reason == POOL_ACTIVE_PROBE_SEALED_SKIP_REASON)
.map(|skipped| skipped.candidate.provider_id.clone())
.collect::<BTreeSet<_>>()
} else {
BTreeSet::new()
};
if active_probe_seal_fallback_provider_ids.is_empty() {
return PoolSchedulerApplyOutcome {
candidates: scheduled,
skipped,
active_probe_seal_fallback_provider_ids,
active_probe_evicted_members_by_provider,
};
}
let (scheduled, skipped) = run_local_execution_pool_scheduler_with_runtime_map(
candidates,
runtime_by_provider,
key_context_by_id,
false,
);
merge_active_probe_evictions(
&mut active_probe_evicted_members_by_provider,
active_probe_evicted_members_from_skipped(&skipped, runtime_by_provider),
);
PoolSchedulerApplyOutcome {
candidates: scheduled,
skipped,
active_probe_seal_fallback_provider_ids,
active_probe_evicted_members_by_provider,
}
}
fn merge_active_probe_evictions(
target: &mut BTreeMap<String, BTreeSet<String>>,
source: BTreeMap<String, BTreeSet<String>>,
) {
for (provider_id, key_ids) in source {
target.entry(provider_id).or_default().extend(key_ids);
}
}
fn active_probe_evicted_members_from_skipped(
skipped: &[SkippedLocalExecutionCandidate],
runtime_by_provider: &BTreeMap<String, AdminProviderPoolRuntimeState>,
) -> BTreeMap<String, BTreeSet<String>> {
if skipped.is_empty() {
return BTreeMap::new();
}
let mut evicted = BTreeMap::<String, BTreeSet<String>>::new();
for skipped_candidate in skipped {
if !matches!(
skipped_candidate.skip_reason,
POOL_ACCOUNT_BLOCKED_SKIP_REASON
| POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
| POOL_COOLDOWN_SKIP_REASON
| POOL_COST_LIMIT_REACHED_SKIP_REASON
) {
continue;
}
let Some(runtime) = runtime_by_provider.get(&skipped_candidate.candidate.provider_id)
else {
continue;
};
if runtime
.active_probe_member_ids
.contains(&skipped_candidate.candidate.key_id)
{
evicted
.entry(skipped_candidate.candidate.provider_id.clone())
.or_default()
.insert(skipped_candidate.candidate.key_id.clone());
}
}
evicted
}
fn run_local_execution_pool_scheduler_with_runtime_map(
candidates: Vec<EligibleLocalExecutionCandidate>,
runtime_by_provider: &BTreeMap<String, AdminProviderPoolRuntimeState>,
key_context_by_id: &BTreeMap<String, PoolCatalogKeyContext>,
enforce_active_probe_seal: bool,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let scheduler_runtime_by_provider = runtime_by_provider
.iter()
@@ -941,13 +1251,14 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
let admin_pool_config = pool_config_for_candidate(&candidate);
if let Some(config) = admin_pool_config.as_ref() {
if should_enforce_active_probe_sealed_pool(config) {
if enforce_active_probe_seal && should_enforce_active_probe_sealed_pool(config) {
let active_member_ids = runtime_by_provider
.get(&candidate.candidate.provider_id)
.map(|runtime| &runtime.active_probe_member_ids);
if !active_member_ids
.is_some_and(|members| members.contains(&candidate.candidate.key_id))
{
let should_seal_cold_member = active_member_ids.is_some_and(|members| {
!members.is_empty() && !members.contains(&candidate.candidate.key_id)
});
if should_seal_cold_member {
skipped_candidates.push(SkippedLocalExecutionCandidate {
candidate: candidate.candidate.clone(),
skip_reason: POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
@@ -1149,8 +1460,12 @@ fn apply_pool_orchestration(
#[cfg(test)]
mod tests {
use super::{
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
pool_config_for_candidate, should_trigger_active_probe_burst_for_request,
admin_provider_pool_quota_probe_active_members_key,
apply_local_execution_pool_scheduler_with_runtime_map,
apply_local_execution_pool_scheduler_with_runtime_map_outcome,
build_pool_catalog_key_context, pool_config_for_candidate,
prune_unschedulable_active_probe_members_for_request,
remove_active_probe_members_for_request, should_trigger_active_probe_burst_for_request,
PoolCatalogKeyContext, PoolKeyCursor, POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON,
};
@@ -1583,7 +1898,142 @@ mod tests {
}
#[test]
fn pool_scheduler_keeps_pool_out_keys_sealed_when_active_probe_pool_is_empty() {
fn pool_scheduler_falls_back_when_active_probe_members_are_unschedulable() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let key_hot = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-hot",
10,
provider_config.clone(),
);
let key_cold = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-cold",
10,
provider_config,
);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-hot".to_string()]),
cooldown_reason_by_key: BTreeMap::from([(
"key-hot".to_string(),
"429".to_string(),
)]),
provider_desired_hot: 1,
..AdminProviderPoolRuntimeState::default()
},
)]);
let outcome = apply_local_execution_pool_scheduler_with_runtime_map_outcome(
vec![key_hot, key_cold],
&runtime_by_provider,
&BTreeMap::new(),
);
let scheduled = outcome.candidates;
let skipped = outcome.skipped;
assert_eq!(
scheduled
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-cold"]
);
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![("key-hot", "pool_cooldown")]
);
assert_eq!(
outcome
.active_probe_evicted_members_by_provider
.get("provider-pool"),
Some(&BTreeSet::from(["key-hot".to_string()]))
);
}
#[test]
fn pool_scheduler_prunes_cold_active_probe_members_before_scheduling() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let key_hot = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-hot",
10,
provider_config,
);
let mut runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-hot".to_string()]),
cooldown_reason_by_key: BTreeMap::from([(
"key-hot".to_string(),
"429".to_string(),
)]),
provider_desired_hot: 1,
..AdminProviderPoolRuntimeState::default()
},
)]);
let evicted = prune_unschedulable_active_probe_members_for_request(
&mut runtime_by_provider,
&[key_hot],
&BTreeMap::new(),
);
assert_eq!(
evicted.get("provider-pool"),
Some(&BTreeSet::from(["key-hot".to_string()]))
);
assert!(runtime_by_provider
.get("provider-pool")
.expect("runtime should exist")
.active_probe_member_ids
.is_empty());
}
#[tokio::test]
async fn pool_scheduler_removes_unschedulable_member_from_active_probe_set() {
let app = AppState::new().expect("state should build");
let set_key = admin_provider_pool_quota_probe_active_members_key("provider-pool");
app.runtime_state
.set_add(&set_key, "key-hot")
.await
.expect("active member should insert");
remove_active_probe_members_for_request(
PlannerAppState::new(&app),
&BTreeMap::from([(
"provider-pool".to_string(),
BTreeSet::from(["key-hot".to_string()]),
)]),
)
.await;
let members = app
.runtime_state
.set_members(&set_key)
.await
.expect("active members should read");
assert!(members.is_empty());
}
#[test]
fn pool_scheduler_allows_cold_start_when_active_probe_pool_is_empty() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
@@ -1610,16 +2060,19 @@ mod tests {
&BTreeMap::new(),
);
assert!(scheduled.is_empty());
assert_eq!(
scheduled
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-a", "key-b"]
);
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![
("key-a", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON),
("key-b", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON),
]
Vec::<(&str, &str)>::new()
);
}

93
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
View File

@@ -4,7 +4,8 @@ use super::keys::{
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolUnschedulableRule,
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolUnschedulableRule,
};
use aether_runtime_state::RuntimeState;
use regex::Regex;
@@ -311,6 +312,27 @@ async fn set_pool_cooldown(
std::time::Duration::from_secs(ttl_seconds.saturating_add(60)),
)
.await;
spawn_remove_pool_active_probe_member(runtime, provider_id, key_id);
}
fn spawn_remove_pool_active_probe_member(runtime: &RuntimeState, provider_id: &str, key_id: &str) {
let runtime = runtime.clone();
let provider_id = provider_id.to_string();
let key_id = key_id.to_string();
tokio::spawn(async move {
if let Err(err) = runtime
.set_remove(
&admin_provider_pool_quota_probe_active_members_key(&provider_id),
&key_id,
)
.await
{
warn!(
"gateway admin provider pool: failed to remove active probe member for provider {provider_id} key {key_id}: {:?}",
err
);
}
});
}
async fn invalidate_pool_oauth_cache(runtime: &RuntimeState, key_id: &str) {
@@ -423,10 +445,12 @@ pub(crate) async fn record_admin_provider_pool_error(
if status_code == 401 {
invalidate_pool_oauth_cache(runtime, key_id).await;
spawn_remove_pool_active_probe_member(runtime, provider_id, key_id);
return;
}
if status_code == 402 {
spawn_remove_pool_active_probe_member(runtime, provider_id, key_id);
return;
}
@@ -435,6 +459,7 @@ pub(crate) async fn record_admin_provider_pool_error(
.iter()
.any(|pattern| error_message.contains(pattern))
{
spawn_remove_pool_active_probe_member(runtime, provider_id, key_id);
return;
}
set_pool_cooldown(
@@ -569,8 +594,8 @@ mod tests {
};
use crate::handlers::admin::provider::pool::runtime::reads::read_admin_provider_pool_runtime_state;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule,
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
};
use crate::AppState;
use aether_runtime_state::{RedisClientConfig, RuntimeState, RuntimeStateConfig};
@@ -642,6 +667,24 @@ mod tests {
.with_runtime_state(std::sync::Arc::new(runtime_state))
}
async fn wait_for_active_probe_members_empty(runtime: &RuntimeState, set_key: &str) {
for _ in 0..20 {
let members = runtime
.set_members(set_key)
.await
.expect("active members should read");
if members.is_empty() {
return;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let members = runtime
.set_members(set_key)
.await
.expect("active members should read");
assert!(members.is_empty());
}
#[test]
fn parses_google_quota_cooldown_from_reset_timestamp() {
let now_unix_secs = chrono::DateTime::parse_from_rfc3339("2026-04-17T10:00:00Z")
@@ -910,6 +953,50 @@ mod tests {
.is_some_and(|ttl| *ttl <= 120 && *ttl >= 100));
}
#[tokio::test]
async fn error_feedback_removes_active_probe_member_when_key_becomes_unschedulable() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_evict_active_probe").await;
let runtime = app.runtime_state.as_ref();
let pool_config = sample_pool_config();
let set_key = admin_provider_pool_quota_probe_active_members_key("provider-1");
runtime
.set_add(&set_key, "key-2")
.await
.expect("active member should insert");
record_admin_provider_pool_error(
runtime,
"provider-1",
"key-2",
&pool_config,
429,
Some(r#"{"error":{"message":"rate limited"}}"#),
None,
)
.await;
wait_for_active_probe_members_empty(runtime, &set_key).await;
runtime
.set_add(&set_key, "key-402")
.await
.expect("active member should insert");
record_admin_provider_pool_error(
runtime,
"provider-1",
"key-402",
&pool_config,
402,
Some(r#"{"error":{"message":"quota exhausted"}}"#),
None,
)
.await;
wait_for_active_probe_members_empty(runtime, &set_key).await;
}
#[tokio::test]
async fn error_feedback_uses_google_quota_cooldown_when_retry_after_missing() {
let Some(redis) = start_managed_redis_or_skip().await else {

7
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -1741,9 +1741,9 @@ fn local_execution_runtime_miss_all_candidates_skipped_detail(
(_, Some(summary), Some(model)) => format!(
"支持模型 {model} 的候选提供商全部不可用：{summary}（原因代码: all_candidates_skipped）"
),
(_, Some(summary), None) => format!(
"候选提供商全部不可用：{summary}（原因代码: all_candidates_skipped）"
),
(_, Some(summary), None) => {
format!("候选提供商全部不可用：{summary}（原因代码: all_candidates_skipped）")
}
(count, None, Some(model)) if count > 0 => format!(
"找到 {count} 个支持模型 {model} 的候选提供商，但都不满足本次{request_mode}请求要求（原因代码: all_candidates_skipped）"
),
@@ -1801,6 +1801,7 @@ fn local_execution_runtime_miss_skip_reason_label(reason: &str) -> &str {
"key_inactive" => "API Key 未启用",
"key_model_disabled" => "API Key 未允许该模型",
"mapped_model_missing" => "模型映射缺失",
"pool_active_probe_sealed" => "池内账号未进入主动探测热池",
"pool_cooldown" => "池内账号处于冷却中",
"pool_cost_limit_reached" => "池内账号成本额度已用尽",
"pool_group_exhausted" => "池化提供商没有可调度账号",