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

Unify candidate ranking pipeline

Browse Source
This commit is contained in:
fawney19
2026-04-27 12:34:03 +08:00
parent 9b866a6d17
commit 3b542434a2
46 changed files with 3635 additions and 1914 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
apps/aether-gateway/src/scheduler/candidate/enumeration.rs Normal file
View File

@@ -0,0 +1,28 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
enumerate_minimal_candidate_selection_with_required_capabilities,
MinimalCandidateSelectionRowSource,
};
use crate::GatewayError;
pub(super) async fn enumerate_scheduler_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
enumerate_minimal_candidate_selection_with_required_capabilities(
selection_row_source,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
required_capabilities,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}

4
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -1,10 +1,12 @@
use self::affinity::candidate_affinity_hash;
use self::selection::{
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
};
use super::state::SchedulerRuntimeState;
mod affinity;
mod enumeration;
mod ranking;
mod resolution;
mod runtime;
mod selection;

76
apps/aether-gateway/src/scheduler/candidate/ranking.rs Normal file
View File

@@ -0,0 +1,76 @@
use aether_scheduler_core::{
apply_scheduler_candidate_ranking, effective_provider_key_health_score,
matches_affinity_target, provider_key_health_bucket,
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
};
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use super::affinity::candidate_affinity_hash;
use super::runtime::CandidateRuntimeSelectionSnapshot;
use super::SchedulerMinimalCandidateSelectionCandidate;
pub(super) fn rank_scheduler_candidates(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
ordering_config: SchedulerOrderingConfig,
required_capabilities: Option<&serde_json::Value>,
priority_affinity_key: Option<&str>,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
now_unix_secs: u64,
) {
let rankables = candidates
.iter()
.enumerate()
.map(|(index, candidate)| {
let provider_key = runtime_snapshot
.provider_key_rpm_states
.get(&candidate.key_id);
SchedulerRankableCandidate::from_candidate(candidate, index)
.with_capability_priority(requested_capability_priority_for_candidate(
required_capabilities,
candidate,
))
.with_cached_affinity_match(
cached_affinity_target
.is_some_and(|target| matches_affinity_target(candidate, target)),
)
.with_affinity_hash(
priority_affinity_key.map(|key| candidate_affinity_hash(key, candidate)),
)
.with_health(
provider_key.and_then(|key| {
provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
}),
provider_key
.and_then(|key| {
effective_provider_key_health_score(
key,
candidate.endpoint_api_format.as_str(),
)
})
.unwrap_or(1.0),
)
})
.collect::<Vec<_>>();
apply_scheduler_candidate_ranking(
candidates,
&rankables,
SchedulerRankingContext {
priority_mode: ordering_config.priority_mode,
ranking_mode: scheduler_ranking_mode(ordering_config.scheduling_mode),
include_health: true,
load_balance_seed: now_unix_secs,
},
);
}
fn scheduler_ranking_mode(mode: SchedulerSchedulingMode) -> SchedulerRankingMode {
match mode {
SchedulerSchedulingMode::FixedOrder => SchedulerRankingMode::FixedOrder,
SchedulerSchedulingMode::CacheAffinity => SchedulerRankingMode::CacheAffinity,
SchedulerSchedulingMode::LoadBalance => SchedulerRankingMode::LoadBalance,
}
}

45
apps/aether-gateway/src/scheduler/candidate/resolution.rs Normal file
View File

@@ -0,0 +1,45 @@
use std::collections::BTreeSet;
use aether_scheduler_core::SchedulerAffinityTarget;
use super::affinity::candidate_key;
use super::runtime::{current_candidate_runtime_skip_reason, CandidateRuntimeSelectionSnapshot};
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerSkippedCandidate};
pub(super) fn resolve_scheduler_candidate_selectability(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
) -> (
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
) {
let mut selected = Vec::with_capacity(candidates.len());
let mut skipped = Vec::new();
let mut emitted_selected_keys = BTreeSet::new();
let mut emitted_skipped_keys = BTreeSet::new();
for candidate in candidates {
let key = candidate_key(&candidate);
if let Some(skip_reason) = current_candidate_runtime_skip_reason(
&candidate,
runtime_snapshot,
now_unix_secs,
cached_affinity_target,
) {
if emitted_skipped_keys.insert(key) {
skipped.push(SchedulerSkippedCandidate {
candidate,
skip_reason,
});
}
continue;
}
if emitted_selected_keys.insert(key) {
selected.push(candidate);
}
}
(selected, skipped)
}

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

@@ -1,27 +1,15 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::{
collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health as reorder_candidates_by_scheduler_health_in_core,
SchedulerPriorityMode,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities,
MinimalCandidateSelectionRowSource,
};
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::GatewayError;
use super::affinity::{
build_scheduler_affinity_cache_key, candidate_key, remember_scheduler_affinity,
};
use super::affinity::{build_scheduler_affinity_cache_key, remember_scheduler_affinity};
use super::enumeration::enumerate_scheduler_candidates;
use super::ranking::rank_scheduler_candidates;
use super::resolution::resolve_scheduler_candidate_selectability;
use super::runtime::{
auth_snapshot_concurrency_limit_reached, current_candidate_runtime_skip_reason,
read_candidate_runtime_selection_snapshot,
auth_snapshot_concurrency_limit_reached, read_candidate_runtime_selection_snapshot,
};
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
@@ -44,69 +32,6 @@ pub(super) fn is_exact_all_skipped_by_auth_limit(
.all(|candidate| candidate.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
}
pub(super) fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
required_capabilities: Option<&serde_json::Value>,
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) {
reorder_candidates_by_scheduler_health_in_core(
candidates,
provider_key_rpm_states,
required_capabilities,
affinity_key,
priority_mode,
);
}
fn apply_load_balance_rotation(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
priority_mode: SchedulerPriorityMode,
now_unix_secs: u64,
) {
if candidates.len() < 2 {
return;
}
let mut start = 0usize;
while start < candidates.len() {
let mut end = start + 1;
while end < candidates.len()
&& candidates_share_load_balance_group(
&candidates[start],
&candidates[end],
priority_mode,
)
{
end += 1;
}
let group_len = end - start;
if group_len > 1 {
let offset = usize::try_from(now_unix_secs).unwrap_or(0) % group_len;
candidates[start..end].rotate_left(offset);
}
start = end;
}
}
fn candidates_share_load_balance_group(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
) -> bool {
match priority_mode {
SchedulerPriorityMode::Provider => {
left.provider_priority == right.provider_priority
&& left.key_internal_priority == right.key_internal_priority
}
SchedulerPriorityMode::GlobalKey => {
left.key_global_priority_for_format == right.key_global_priority_for_format
}
}
}
#[cfg_attr(not(test), allow(dead_code))]
pub(super) async fn select_minimal_candidate(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
@@ -182,36 +107,18 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let priority_affinity_key =
scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
let mut candidates =
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
let mut candidates = enumerate_scheduler_candidates(
selection_row_source,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
ordering_config.priority_mode,
priority_affinity_key,
required_capabilities,
auth_snapshot,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
.await?;
let runtime_snapshot =
read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
.await?;
reorder_candidates_by_scheduler_health(
&mut candidates,
&runtime_snapshot.provider_key_rpm_states,
required_capabilities,
priority_affinity_key,
ordering_config.priority_mode,
);
if ordering_config.scheduling_mode == SchedulerSchedulingMode::LoadBalance {
apply_load_balance_rotation(
&mut candidates,
ordering_config.priority_mode,
now_unix_secs,
);
}
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let cached_affinity_target = if ordering_config.scheduling_mode
@@ -225,6 +132,15 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
};
if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) {
rank_scheduler_candidates(
&mut candidates,
&runtime_snapshot,
ordering_config,
required_capabilities,
priority_affinity_key,
cached_affinity_target.as_ref(),
now_unix_secs,
);
return Ok((
Vec::new(),
candidates
@@ -237,37 +153,23 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
));
}
let mut selected_keys = BTreeSet::new();
let mut skipped = Vec::new();
let mut emitted_skipped_keys = BTreeSet::new();
let (mut selected, skipped) = resolve_scheduler_candidate_selectability(
candidates,
&runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
);
rank_scheduler_candidates(
&mut selected,
&runtime_snapshot,
ordering_config,
required_capabilities,
priority_affinity_key,
cached_affinity_target.as_ref(),
now_unix_secs,
);
for candidate in &candidates {
let key = candidate_key(candidate);
if let Some(skip_reason) = current_candidate_runtime_skip_reason(
candidate,
&runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
) {
if emitted_skipped_keys.insert(key) {
skipped.push(SchedulerSkippedCandidate {
candidate: candidate.clone(),
skip_reason,
});
}
continue;
}
selected_keys.insert(key);
}
Ok((
collect_selectable_candidates_from_keys(
candidates,
&selected_keys,
cached_affinity_target.as_ref(),
),
skipped,
))
Ok((selected, skipped))
}
fn scheduling_priority_affinity_key<'a>(

48
apps/aether-gateway/src/scheduler/candidate/tests/model.rs
View File

@@ -144,6 +144,54 @@ async fn read_minimal_candidate_selection_resolves_provider_model_alias() {
assert_eq!(selection[0].selected_provider_model_name, "gpt-5.2");
}
#[tokio::test]
async fn read_minimal_candidate_selection_keeps_all_rows_supporting_requested_model() {
let mut exact = sample_row();
exact.provider_id = "provider-exact".to_string();
exact.endpoint_id = "endpoint-exact".to_string();
exact.key_id = "key-exact".to_string();
exact.model_id = "model-exact".to_string();
exact.global_model_id = "global-exact".to_string();
exact.global_model_name = "gpt-5".to_string();
exact.model_provider_model_name = "gpt-5".to_string();
exact.model_provider_model_mappings = None;
let mut mapped = sample_row();
mapped.provider_id = "provider-mapped".to_string();
mapped.endpoint_id = "endpoint-mapped".to_string();
mapped.key_id = "key-mapped".to_string();
mapped.model_id = "model-mapped".to_string();
mapped.global_model_id = "global-mapped".to_string();
mapped.global_model_name = "claude-sonnet".to_string();
mapped.global_model_mappings = Some(vec!["gpt-5".to_string()]);
mapped.model_provider_model_name = "claude-sonnet-upstream".to_string();
mapped.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "claude-sonnet-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
exact, mapped,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let selection = read_minimal_candidate_selection(&state, "openai:chat", "gpt-5", false, None)
.await
.expect("selection should succeed");
let provider_ids = selection
.iter()
.map(|candidate| candidate.provider_id.as_str())
.collect::<Vec<_>>();
assert_eq!(provider_ids, vec!["provider-exact", "provider-mapped"]);
assert_eq!(
selection[1].selected_provider_model_name,
"claude-sonnet-upstream"
);
}
#[tokio::test]
async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_snapshot() {
let mut row = sample_row();

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

@@ -384,9 +384,9 @@ async fn fixed_order_disables_same_priority_affinity_hash_tiebreaker() {
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.provider_priority = 0;
second.provider_priority = 1;
second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
@@ -438,9 +438,9 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.provider_priority = 1;
second.provider_priority = 0;
second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,