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

Fix active probe pool fallback tracing

Browse Source
This commit is contained in:
fawney19
2026-05-17 20:34:06 +08:00
parent a2f91b4108
commit 7d569127ae
4 changed files with 1181 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -7,7 +7,13 @@ use aether_ai_serving::{
AiCandidatePreselectionOutcome, AiSkippedCandidatePersistencePort,
};
use aether_dispatch_core::{DispatchSequence, DispatchSequenceItem};
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
use aether_routing_core::{
rank_vector_for_candidate, CandidateKind, ResolvedRoutingPolicy, RoutingCandidateFacts,
RoutingCandidateTrace, RoutingDecisionTrace,
};
use aether_scheduler_core::{
ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome,
};
use async_trait::async_trait;
use serde_json::Value;
use std::collections::VecDeque;
@@ -19,6 +25,7 @@ use tracing::warn;
use uuid::Uuid;
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate_at_epoch;
use crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy;
use crate::ai_serving::planner::candidate_resolution::{
resolve_and_rank_logical_local_execution_candidates, EligibleLocalExecutionCandidate,
LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
@@ -36,7 +43,7 @@ use crate::dispatch::refs::dispatch_ref_for_local_candidate;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
use crate::scheduler::candidate::API_KEY_CONCURRENCY_LIMIT_SKIP_REASON;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::{AppState, GatewayError};
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
@@ -74,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>>,
@@ -110,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();
@@ -179,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> {
@@ -189,6 +239,7 @@ struct GatewayLocalCandidateMaterializationPort<'a, F, G> {
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&'a ClientSessionAffinity>,
required_capabilities: Option<&'a Value>,
routing_policy: Option<&'a ResolvedRoutingPolicy>,
sticky_session_token: Option<&'a str>,
request_auth_channel: Option<&'a str>,
persistence_policy: LocalCandidatePersistencePolicy<'a>,
@@ -243,6 +294,7 @@ where
self.auth_snapshot,
self.client_session_affinity,
self.required_capabilities,
self.routing_policy,
self.sticky_session_token,
self.request_auth_channel,
self.resolution_mode,
@@ -281,6 +333,8 @@ where
.skipped
.record_runtime_miss_diagnostic,
candidates,
self.routing_policy,
self.client_api_format,
self.sticky_session_token,
self.requested_model,
self.request_auth_channel,
@@ -294,6 +348,12 @@ where
starting_candidate_index: u32,
skipped_candidates: Vec<Self::Skipped>,
) -> Result<(), Self::Error> {
let skipped_candidates = attach_routing_trace_to_skipped_candidates(
self.routing_policy,
self.client_api_format,
starting_candidate_index,
skipped_candidates,
);
persist_skipped_local_execution_candidates_with_context(
self.state.app(),
self.trace_id,
@@ -432,6 +492,7 @@ pub(crate) async fn materialize_local_execution_candidates_with_serving<F, G>(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&Value>,
routing_policy: Option<&ResolvedRoutingPolicy>,
sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
persistence_policy: LocalCandidatePersistencePolicy<'_>,
@@ -445,7 +506,8 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
{
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
let scheduler_cache_affinity_enabled =
scheduler_cache_affinity_enabled(state, routing_policy).await;
let port = GatewayLocalCandidateMaterializationPort {
state,
trace_id,
@@ -454,6 +516,7 @@ where
auth_snapshot,
client_session_affinity,
required_capabilities,
routing_policy,
sticky_session_token,
request_auth_channel,
persistence_policy,
@@ -478,6 +541,7 @@ pub(crate) async fn build_local_execution_candidate_attempt_source_with_serving<
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&Value>,
routing_policy: Option<&ResolvedRoutingPolicy>,
sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
persistence_policy: LocalCandidatePersistencePolicy<'_>,
@@ -491,7 +555,8 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
{
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
let scheduler_cache_affinity_enabled =
scheduler_cache_affinity_enabled(state, routing_policy).await;
let _ = build_available_extra_data;
let (candidates, resolved_skipped) = resolve_and_rank_logical_local_execution_candidates(
state,
@@ -501,6 +566,7 @@ where
auth_snapshot,
client_session_affinity,
required_capabilities,
routing_policy,
sticky_session_token,
request_auth_channel,
resolution_mode,
@@ -529,7 +595,12 @@ where
trace_id,
persistence_policy.skipped,
u32::try_from(candidates.len()).unwrap_or(u32::MAX),
skipped_candidates,
attach_routing_trace_to_skipped_candidates(
routing_policy,
client_api_format,
u32::try_from(candidates.len()).unwrap_or(u32::MAX),
skipped_candidates,
),
)
.await;
@@ -542,6 +613,14 @@ where
sticky_session_token,
requested_model,
request_auth_channel,
routing_policy,
Some(PoolGroupExhaustionPersistenceContext::new(
state.app().clone(),
trace_id,
persistence_policy.skipped,
client_api_format,
routing_policy,
)),
);
(
@@ -559,6 +638,8 @@ fn build_logical_candidate_items<'a>(
sticky_session_token: Option<&str>,
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;
@@ -578,12 +659,13 @@ fn build_logical_candidate_items<'a>(
}
}
LocalExecutionCandidateKind::PoolGroup => {
let cursor = PoolKeyCursor::new(
let cursor = PoolKeyCursor::new_with_routing_policy(
state,
candidate,
sticky_session_token,
requested_model,
request_auth_channel,
routing_policy,
);
let cursor = if let Some(trace_id) = trace_id {
cursor.with_runtime_miss_diagnostic(trace_id, record_runtime_miss_diagnostic)
@@ -594,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(),
});
}
}
@@ -615,6 +698,7 @@ pub(crate) async fn build_lazy_requested_model_execution_candidate_attempt_sourc
auth_snapshot: &GatewayAuthApiKeySnapshot,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&Value>,
routing_policy: Option<&ResolvedRoutingPolicy>,
sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
persistence_policy: LocalCandidatePersistencePolicy<'_>,
@@ -628,7 +712,8 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync + 'a,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync + 'a,
{
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
let scheduler_cache_affinity_enabled =
scheduler_cache_affinity_enabled(state, routing_policy).await;
let _ = build_available_extra_data;
let decorate_skipped_candidate = Arc::new(decorate_skipped_candidate);
let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic;
@@ -639,6 +724,7 @@ where
require_streaming,
required_capabilities,
auth_snapshot,
routing_policy,
client_session_affinity,
use_api_format_alias_match,
key_mode,
@@ -652,6 +738,7 @@ where
auth_snapshot: auth_snapshot.clone(),
client_session_affinity: client_session_affinity.cloned(),
required_capabilities: required_capabilities.cloned(),
routing_policy: routing_policy.cloned(),
sticky_session_token: sticky_session_token.map(str::to_string),
request_auth_channel: request_auth_channel.map(str::to_string),
skipped_user_id: persistence_policy.skipped.user_id.to_string(),
@@ -693,6 +780,7 @@ struct RequestedModelAttemptPageCursor<'a> {
auth_snapshot: GatewayAuthApiKeySnapshot,
client_session_affinity: Option<ClientSessionAffinity>,
required_capabilities: Option<Value>,
routing_policy: Option<ResolvedRoutingPolicy>,
sticky_session_token: Option<String>,
request_auth_channel: Option<String>,
skipped_user_id: String,
@@ -756,6 +844,7 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
Some(&self.auth_snapshot),
self.client_session_affinity.as_ref(),
self.required_capabilities.as_ref(),
self.routing_policy.as_ref(),
self.sticky_session_token.as_deref(),
self.request_auth_channel.as_deref(),
self.resolution_mode,
@@ -794,6 +883,17 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
self.sticky_session_token.as_deref(),
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));
@@ -814,7 +914,12 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
&self.trace_id,
skipped_persistence,
skipped_starting_candidate_index,
skipped_candidates,
attach_routing_trace_to_skipped_candidates(
self.routing_policy.as_ref(),
&self.client_api_format,
skipped_starting_candidate_index,
skipped_candidates,
),
)
.await;
}
@@ -858,7 +963,12 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
&self.trace_id,
skipped_persistence,
self.next_candidate_index,
skipped_candidates,
attach_routing_trace_to_skipped_candidates(
self.routing_policy.as_ref(),
&self.client_api_format,
self.next_candidate_index,
skipped_candidates,
),
)
.await;
self.next_candidate_index = self
@@ -900,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();
@@ -925,19 +1044,14 @@ async fn pop_attempt_from_items(
}
}
async fn scheduler_cache_affinity_enabled(state: PlannerAppState<'_>) -> bool {
match read_scheduler_ordering_config(state.app()).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
Err(error) => {
warn!(
event_name = "planner_scheduler_affinity_config_load_failed",
log_type = "event",
error = ?error,
"failed to load scheduler config while checking cache affinity mode"
);
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
}
}
async fn scheduler_cache_affinity_enabled(
state: PlannerAppState<'_>,
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> bool {
scheduler_ordering_config_for_routing_policy(state, routing_policy)
.await
.scheduling_mode
== SchedulerSchedulingMode::CacheAffinity
}
pub(crate) fn remember_first_local_candidate_affinity(
@@ -1038,6 +1152,8 @@ async fn materialize_logical_local_execution_candidate_attempts<F>(
context: LocalAvailableCandidatePersistenceContext<'_>,
record_runtime_miss_diagnostic: bool,
candidates: Vec<EligibleLocalExecutionCandidate>,
routing_policy: Option<&ResolvedRoutingPolicy>,
client_api_format: &str,
sticky_session_token: Option<&str>,
requested_model: Option<&str>,
request_auth_channel: Option<&str>,
@@ -1059,18 +1175,21 @@ where
context,
candidate,
candidate_index,
routing_policy,
client_api_format,
build_extra_data,
)
.await,
);
}
LocalExecutionCandidateKind::PoolGroup => {
let mut cursor = PoolKeyCursor::new(
let mut cursor = PoolKeyCursor::new_with_routing_policy(
state,
candidate,
sticky_session_token,
requested_model,
request_auth_channel,
routing_policy,
)
.with_runtime_miss_diagnostic(trace_id, record_runtime_miss_diagnostic);
let attempt_count_before_pool = attempts.len();
@@ -1082,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;
}
}
}
}
@@ -1097,6 +1238,8 @@ async fn persist_available_local_execution_candidate_at_index<F>(
context: LocalAvailableCandidatePersistenceContext<'_>,
candidate: EligibleLocalExecutionCandidate,
candidate_index: u32,
routing_policy: Option<&ResolvedRoutingPolicy>,
client_api_format: &str,
build_extra_data: &F,
) -> Vec<LocalExecutionCandidateAttempt>
where
@@ -1107,6 +1250,16 @@ where
available_candidate_base_extra_data_with_dispatch_ref(&candidate, build_extra_data),
candidate.ranking.as_ref(),
);
let extra_data = attach_routing_trace_to_extra_data(
routing_policy,
client_api_format,
&candidate.candidate,
candidate.kind,
candidate.ranking.as_ref(),
None,
Some(candidate_index),
extra_data,
);
let should_persist = should_persist_available_local_candidate(&candidate);
let mut attempts = Vec::with_capacity(attempt_slots as usize);
let mut owned_candidate = Some(candidate);
@@ -1190,6 +1343,160 @@ where
Some(Value::Object(object))
}
fn attach_routing_trace_to_skipped_candidates(
routing_policy: Option<&ResolvedRoutingPolicy>,
client_api_format: &str,
starting_candidate_index: u32,
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
) -> Vec<SkippedLocalExecutionCandidate> {
skipped_candidates
.into_iter()
.enumerate()
.map(|(offset, skipped)| {
let selected_order =
starting_candidate_index.saturating_add(u32::try_from(offset).unwrap_or(u32::MAX));
attach_routing_trace_to_skipped_candidate(
routing_policy,
client_api_format,
selected_order,
skipped,
)
})
.collect()
}
fn attach_routing_trace_to_skipped_candidate(
routing_policy: Option<&ResolvedRoutingPolicy>,
client_api_format: &str,
selected_order: u32,
mut skipped_candidate: SkippedLocalExecutionCandidate,
) -> SkippedLocalExecutionCandidate {
let kind = if skipped_candidate
.transport
.as_ref()
.is_some_and(|transport| {
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref())
.is_some()
}) {
LocalExecutionCandidateKind::PoolGroup
} else {
LocalExecutionCandidateKind::SingleKey
};
skipped_candidate.extra_data = attach_routing_trace_to_extra_data(
routing_policy,
client_api_format,
&skipped_candidate.candidate,
kind,
skipped_candidate.ranking.as_ref(),
Some(skipped_candidate.skip_reason),
Some(selected_order),
skipped_candidate.extra_data,
);
skipped_candidate
}
#[allow(clippy::too_many_arguments)]
fn attach_routing_trace_to_extra_data(
routing_policy: Option<&ResolvedRoutingPolicy>,
client_api_format: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
kind: LocalExecutionCandidateKind,
ranking: Option<&SchedulerRankingOutcome>,
skip_reason: Option<&'static str>,
selected_order: Option<u32>,
extra_data: Option<Value>,
) -> Option<Value> {
let Some(policy) = routing_policy else {
return extra_data;
};
let routing_trace = routing_trace_for_candidate(
policy,
client_api_format,
candidate,
kind,
ranking,
skip_reason,
selected_order,
);
Some(merge_routing_trace_into_extra_data(
extra_data,
routing_trace,
))
}
fn merge_routing_trace_into_extra_data(
extra_data: Option<Value>,
routing_trace: RoutingDecisionTrace,
) -> Value {
let mut object = match extra_data {
Some(Value::Object(object)) => object,
Some(value) => {
let mut object = serde_json::Map::new();
object.insert("extra".to_string(), value);
object
}
None => serde_json::Map::new(),
};
object.insert(
"routing_trace".to_string(),
serde_json::json!(routing_trace),
);
Value::Object(object)
}
fn routing_trace_for_candidate(
policy: &ResolvedRoutingPolicy,
client_api_format: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
kind: LocalExecutionCandidateKind,
ranking: Option<&SchedulerRankingOutcome>,
skip_reason: Option<&'static str>,
selected_order: Option<u32>,
) -> RoutingDecisionTrace {
let candidate_kind = routing_candidate_kind(kind);
let mut trace = crate::routing::build_routing_trace_seed(policy, client_api_format);
trace.global_candidates.push(RoutingCandidateTrace {
candidate_kind,
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
model_id: candidate.model_id.clone(),
key_id: match candidate_kind {
CandidateKind::Provider => Some(candidate.key_id.clone()),
CandidateKind::PoolGroup => None,
},
ranking_vector: rank_vector_for_candidate(
&policy.ranking_overlay,
&RoutingCandidateFacts {
candidate_kind,
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
model_id: candidate.model_id.clone(),
key_id: match candidate_kind {
CandidateKind::Provider => Some(candidate.key_id.clone()),
CandidateKind::PoolGroup => None,
},
provider_priority: candidate.provider_priority,
key_priority: candidate
.key_global_priority_for_format
.unwrap_or(candidate.key_internal_priority),
},
),
skip_reason: skip_reason.map(str::to_string),
selected_order,
});
if let Some(ranking) = ranking {
trace.runtime_facts.cache_affinity_hit = ranking.promoted_by == Some("cached_affinity");
}
trace
}
fn routing_candidate_kind(kind: LocalExecutionCandidateKind) -> CandidateKind {
match kind {
LocalExecutionCandidateKind::SingleKey => CandidateKind::Provider,
LocalExecutionCandidateKind::PoolGroup => CandidateKind::PoolGroup,
}
}
fn dispatch_sequence_from_attempts(
attempts: Vec<LocalExecutionCandidateAttempt>,
) -> DispatchSequence<LocalExecutionCandidateAttempt> {
@@ -1252,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,
@@ -1426,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,
@@ -1644,6 +1986,7 @@ mod tests {
auth_snapshot: Some(&auth_snapshot),
client_session_affinity: None,
required_capabilities: None,
routing_policy: None,
sticky_session_token: None,
request_auth_channel: None,
persistence_policy: LocalCandidatePersistencePolicy {
@@ -1718,6 +2061,8 @@ mod tests {
false,
vec![pool_group, sample_eligible("normal-key", None)],
None,
"openai:chat",
None,
Some("gpt-5"),
None,
&|_| None,
@@ -1732,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"))
@@ -1861,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());

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

@@ -14,8 +14,11 @@ 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};
use tracing::warn;
use crate::ai_serving::{
@@ -31,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,
};
@@ -40,6 +44,7 @@ use crate::orchestration::LocalExecutionCandidateMetadata;
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
const POOL_ACTIVE_PROBE_SEALED_SKIP_REASON: &str = "pool_active_probe_sealed";
const ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON: &str = "routing_profile_disallowed_key";
type PoolCatalogKeyContext = PoolMemberSignals;
@@ -110,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<_>>();
@@ -125,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 {
@@ -150,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,
@@ -187,6 +328,7 @@ pub(crate) struct PoolKeyCursor<'a> {
sticky_session_token: Option<String>,
requested_model: Option<String>,
request_auth_channel: Option<String>,
routing_overlay: Option<RankingOverlay>,
runtime_miss_trace_id: Option<String>,
record_runtime_miss_diagnostic: bool,
pool_key_order: StoredPoolKeyCandidateOrder,
@@ -216,7 +358,26 @@ impl<'a> PoolKeyCursor<'a> {
requested_model: Option<&str>,
request_auth_channel: Option<&str>,
) -> Self {
let pool_key_order = pool_key_candidate_order_for_group(&group);
Self::new_with_routing_policy(
state,
group,
sticky_session_token,
requested_model,
request_auth_channel,
None,
)
}
pub(crate) fn new_with_routing_policy(
state: PlannerAppState<'a>,
group: EligibleLocalExecutionCandidate,
sticky_session_token: Option<&str>,
requested_model: Option<&str>,
request_auth_channel: Option<&str>,
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> Self {
let pool_key_order = pool_key_candidate_order_for_group(&group, routing_policy);
let routing_overlay = routing_policy.map(|policy| policy.ranking_overlay.clone());
let pool_config = pool_config_for_candidate(&group);
let score_top_n = pool_config
.as_ref()
@@ -236,6 +397,7 @@ impl<'a> PoolKeyCursor<'a> {
sticky_session_token: sticky_session_token.map(str::to_string),
requested_model: requested_model.map(str::to_string),
request_auth_channel: request_auth_channel.map(str::to_string),
routing_overlay,
runtime_miss_trace_id: None,
record_runtime_miss_diagnostic: false,
pool_key_order,
@@ -295,6 +457,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;
@@ -551,6 +739,9 @@ impl<'a> PoolKeyCursor<'a> {
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
while let Some(candidate) = self.queued_candidates.pop_front() {
let mut candidate = candidate;
if self.skip_candidate_if_routing_profile_disallowed(&candidate) {
continue;
}
if self.skip_candidate_if_runtime_cooldown(&candidate).await {
continue;
}
@@ -562,6 +753,28 @@ impl<'a> PoolKeyCursor<'a> {
None
}
fn skip_candidate_if_routing_profile_disallowed(
&mut self,
candidate: &EligibleLocalExecutionCandidate,
) -> bool {
let Some(overlay) = self.routing_overlay.as_ref() else {
return false;
};
if overlay.key_allowed(candidate.candidate.key_id.as_str()) {
return false;
}
self.record_skip_reason(ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON);
self.skipped_candidates
.push(SkippedLocalExecutionCandidate {
candidate: candidate.candidate.clone(),
skip_reason: ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON,
transport: Some(candidate.transport.clone()),
ranking: candidate.ranking.clone(),
extra_data: None,
});
true
}
async fn skip_candidate_if_runtime_cooldown(
&mut self,
candidate: &EligibleLocalExecutionCandidate,
@@ -583,6 +796,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,
@@ -602,6 +816,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;
@@ -864,6 +1103,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()
@@ -879,13 +1237,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,
@@ -958,19 +1317,38 @@ fn should_trigger_active_probe_burst_for_request(
fn pool_key_candidate_order_for_group(
group: &EligibleLocalExecutionCandidate,
routing_policy: Option<&ResolvedRoutingPolicy>,
) -> StoredPoolKeyCandidateOrder {
let Some(pool_config) = pool_config_for_candidate(group) else {
return StoredPoolKeyCandidateOrder::InternalPriority;
};
let presets = pool_config
.scheduling_presets
.iter()
.map(|preset| PoolSchedulingPreset {
preset: preset.preset.clone(),
enabled: preset.enabled,
mode: preset.mode.clone(),
let override_presets = routing_policy
.and_then(|policy| {
policy
.pool_policy_overrides
.get(group.candidate.provider_id.as_str())
})
.collect::<Vec<_>>();
.filter(|override_policy| !override_policy.scheduling_presets.is_empty());
let presets = match override_presets {
Some(override_policy) => override_policy
.scheduling_presets
.iter()
.map(|preset| PoolSchedulingPreset {
preset: preset.preset.clone(),
enabled: preset.enabled,
mode: preset.mode.clone(),
})
.collect::<Vec<_>>(),
None => pool_config
.scheduling_presets
.iter()
.map(|preset| PoolSchedulingPreset {
preset: preset.preset.clone(),
enabled: preset.enabled,
mode: preset.mode.clone(),
})
.collect::<Vec<_>>(),
};
let active_presets = ProviderPoolService::with_builtin_adapters()
.normalize_scheduling_presets(group.transport.provider.provider_type.as_str(), &presets)
.into_iter()
@@ -1068,9 +1446,14 @@ 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,
};
use crate::ai_serving::{
apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate,
@@ -1096,6 +1479,9 @@ mod tests {
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
};
use aether_routing_core::{
RankingOverlay, ResolvedRoutingPolicy, RoutingSchedulingMode, RoutingSetPriorityMode,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
@@ -1498,7 +1884,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
@@ -1525,16 +2046,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()
);
}
@@ -2103,6 +2627,59 @@ mod tests {
);
}
#[tokio::test]
async fn pool_key_cursor_filters_expanded_keys_by_routing_profile_allowed_keys() {
let app = AppState::new().expect("state should build");
let provider_config = Some(json!({ "pool_advanced": { "lru_enabled": true } }));
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config.clone(),
);
let routing_policy = routing_policy_with_allowed_keys(["key-b"]);
let mut cursor = PoolKeyCursor::new_with_routing_policy(
PlannerAppState::new(&app),
group,
None,
None,
None,
Some(&routing_policy),
);
cursor.queued_candidates = VecDeque::from([
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
provider_config.clone(),
),
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config),
]);
let candidate = cursor
.next_key()
.await
.expect("cursor should skip disallowed pool key and return allowed key");
assert_eq!(candidate.candidate.key_id, "key-b");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert_eq!(
cursor
.skip_reason_counts
.get(ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON),
Some(&1)
);
let skipped = cursor.take_skipped_candidates();
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![("key-a", ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON)]
);
}
#[tokio::test]
async fn pool_key_cursor_allows_parallel_requests_to_use_same_healthy_key() {
let app = AppState::new().expect("state should build");
@@ -2668,6 +3245,28 @@ mod tests {
(provider, endpoint, keys, rows)
}
fn routing_policy_with_allowed_keys<const N: usize>(
key_ids: [&str; N],
) -> ResolvedRoutingPolicy {
ResolvedRoutingPolicy {
group_id: Some("routing-group-1".to_string()),
group_version: Some(1),
selection_source: "test".to_string(),
requested_model: "gpt-5".to_string(),
resolved_model: "gpt-5".to_string(),
priority_mode: RoutingSetPriorityMode::Provider,
scheduling_mode: RoutingSchedulingMode::CacheAffinity,
keep_priority_on_conversion: false,
ranking_overlay: RankingOverlay {
allowed_keys: key_ids.into_iter().map(str::to_string).collect(),
..RankingOverlay::default()
},
mutation_plan: Default::default(),
pool_policy_overrides: BTreeMap::new(),
matched_rules: Vec::new(),
}
}
fn sample_eligible_candidate(
provider_id: &str,
endpoint_id: &str,

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" => "池化提供商没有可调度账号",