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

fix scheduler affinity candidate selection

Browse Source
This commit is contained in:
fawney19
2026-04-30 16:27:24 +08:00
parent 558abfcfa3
commit 33aa70c22b
24 changed files with 1715 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
crates/aether-ai-formats/src/matrix.rs
View File

@@ -31,7 +31,12 @@ const NON_COMPACT_STANDARD_CANDIDATE_API_FORMATS: &[&str] = &[
"claude:messages",
"gemini:generate_content",
];
const STANDARD_API_FAMILY_ORDER: &[&str] = &["openai", "claude", "gemini"];
const STANDARD_API_FORMAT_ORDER: &[&str] = &[
"openai:chat",
"openai:responses",
"claude:messages",
"gemini:generate_content",
];
pub fn request_candidate_api_format_preference(
client_api_format: &str,
@@ -60,7 +65,7 @@ pub fn request_candidate_api_format_preference(
Some((
preference_bucket,
standard_api_family_priority(provider_family),
standard_api_format_priority(provider_api_format.as_str()),
))
}
@@ -213,11 +218,12 @@ pub fn normalized_same_standard_api_format(left: &str, right: &str) -> bool {
api_format_alias_matches(left, right)
}
fn standard_api_family_priority(family: &str) -> u8 {
STANDARD_API_FAMILY_ORDER
fn standard_api_format_priority(api_format: &str) -> u8 {
let api_format = normalize_api_format_alias(api_format);
STANDARD_API_FORMAT_ORDER
.iter()
.position(|candidate| *candidate == family)
.unwrap_or(STANDARD_API_FAMILY_ORDER.len()) as u8
.position(|candidate| *candidate == api_format)
.unwrap_or(STANDARD_API_FORMAT_ORDER.len()) as u8
}
#[cfg(test)]
@@ -369,6 +375,19 @@ mod tests {
"gemini:generate_content"
]
);
assert_eq!(
request_candidate_api_formats("claude:messages", false),
vec![
"claude:messages",
"openai:chat",
"openai:responses",
"gemini:generate_content"
]
);
assert!(
request_candidate_api_format_preference("claude:messages", "openai:chat")
< request_candidate_api_format_preference("claude:messages", "openai:responses")
);
assert_eq!(
request_candidate_api_formats("openai:cli", false),
Vec::<&'static str>::new()

9
crates/aether-ai-pipeline/src/conversion/registry.rs
View File

@@ -460,6 +460,15 @@ mod tests {
"gemini:generate_content",
]
);
assert_eq!(
request_candidate_api_formats("claude:messages", false),
vec![
"claude:messages",
"openai:chat",
"openai:responses",
"gemini:generate_content",
]
);
assert_eq!(
request_candidate_api_formats("openai:cli", false),
Vec::<&'static str>::new()

24
crates/aether-scheduler-core/src/ranking/format.rs
View File

@@ -2,11 +2,27 @@ use std::cmp::Ordering;
use super::types::SchedulerRankableCandidate;
pub(super) fn compare_format_state(
pub(super) fn compare_cross_format_demotion(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
) -> Ordering {
left.demote_cross_format
.cmp(&right.demote_cross_format)
.then(left.format_preference.cmp(&right.format_preference))
left.demote_cross_format.cmp(&right.demote_cross_format)
}
pub(super) fn compare_format_preference(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
) -> Ordering {
left.format_preference.cmp(&right.format_preference)
}
pub(super) fn compare_demoted_format_preference(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
) -> Ordering {
if left.demote_cross_format && right.demote_cross_format {
compare_format_preference(left, right)
} else {
Ordering::Equal
}
}

86
crates/aether-scheduler-core/src/ranking/mod.rs
View File

@@ -182,19 +182,15 @@ mod tests {
}
#[test]
fn fixed_order_keeps_priority_before_affinity_tunnel_and_format_preference() {
let mut lower_priority = candidate("lower", 10, 0, Some(10));
lower_priority.cached_affinity_match = true;
lower_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::LocalTunnel;
lower_priority.format_preference = (0, 0);
fn fixed_order_demotes_cross_format_before_priority() {
let lower_priority_same_format = candidate("same", 10, 0, Some(10));
let mut higher_priority = candidate("higher", 0, 0, Some(0));
higher_priority.demote_cross_format = true;
higher_priority.format_preference = (9, 9);
let mut higher_priority_cross_format = candidate("cross", 0, 0, Some(0));
higher_priority_cross_format.demote_cross_format = true;
assert_eq!(
ranked_ids(
&[lower_priority, higher_priority],
&[higher_priority_cross_format, lower_priority_same_format],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
@@ -202,7 +198,7 @@ mod tests {
load_balance_seed: 0,
},
),
vec!["provider-higher", "provider-lower"]
vec!["provider-same", "provider-cross"]
);
}
@@ -289,6 +285,76 @@ mod tests {
);
}
#[test]
fn cache_affinity_promotes_cached_candidate_before_cross_format_demotion() {
let same_format = candidate("same", 10, 0, Some(10));
let mut cached_cross_format = candidate("cross", 0, 0, Some(0));
cached_cross_format.cached_affinity_match = true;
cached_cross_format.demote_cross_format = true;
let outcomes = scheduler_ranking_outcomes(
&[cached_cross_format, same_format],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
},
);
assert_eq!(outcomes[0].original_index, 0);
assert_eq!(
outcomes[0].promoted_by,
Some(RANKING_REASON_CACHED_AFFINITY)
);
assert_eq!(outcomes[0].demoted_by, Some(RANKING_REASON_CROSS_FORMAT));
assert_eq!(outcomes[1].original_index, 1);
}
#[test]
fn demoted_cross_format_candidates_follow_format_preference_before_priority() {
let mut openai_responses_high_priority = candidate("responses", 0, 0, Some(0));
openai_responses_high_priority.demote_cross_format = true;
openai_responses_high_priority.format_preference = (3, 1);
let mut openai_chat_low_priority = candidate("chat", 10, 0, Some(10));
openai_chat_low_priority.demote_cross_format = true;
openai_chat_low_priority.format_preference = (3, 0);
assert_eq!(
ranked_ids(
&[openai_responses_high_priority, openai_chat_low_priority],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-chat", "provider-responses"]
);
}
#[test]
fn load_balance_does_not_rotate_across_cross_format_demotion_group() {
let same_format = candidate("same", 0, 0, Some(0));
let mut cross_format = candidate("cross", 0, 0, Some(0));
cross_format.demote_cross_format = true;
assert_eq!(
ranked_ids(
&[same_format, cross_format],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::LoadBalance,
include_health: false,
load_balance_seed: 1,
},
),
vec!["provider-same", "provider-cross"]
);
}
#[test]
fn load_balance_rotates_only_within_same_priority_group() {
let first = candidate("first", 0, 0, Some(0));

31
crates/aether-scheduler-core/src/ranking/modes.rs
View File

@@ -1,7 +1,9 @@
use std::cmp::Ordering;
use super::compare_candidate_identity_for_ranking;
use super::format::compare_format_state;
use super::format::{
compare_cross_format_demotion, compare_demoted_format_preference, compare_format_preference,
};
use super::priority::{candidates_share_priority_group, compare_candidate_priority_slot};
use super::types::{SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode};
@@ -24,8 +26,10 @@ fn compare_fixed_order(
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then_with(|| compare_cross_format_demotion(left, right))
.then_with(|| compare_demoted_format_preference(left, right))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then_with(|| compare_format_state(left, right))
.then_with(|| compare_format_preference(left, right))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
@@ -38,10 +42,11 @@ fn compare_cache_affinity(
left.capability_priority
.cmp(&right.capability_priority)
.then_with(|| right.cached_affinity_match.cmp(&left.cached_affinity_match))
.then(left.demote_cross_format.cmp(&right.demote_cross_format))
.then_with(|| compare_cross_format_demotion(left, right))
.then_with(|| compare_demoted_format_preference(left, right))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then(left.tunnel_bucket.cmp(&right.tunnel_bucket))
.then(left.format_preference.cmp(&right.format_preference))
.then_with(|| compare_format_preference(left, right))
.then_with(|| compare_health(left, right, context.include_health))
.then(left.affinity_hash.cmp(&right.affinity_hash))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
@@ -55,9 +60,10 @@ fn compare_load_balance_base(
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then(left.demote_cross_format.cmp(&right.demote_cross_format))
.then_with(|| compare_cross_format_demotion(left, right))
.then_with(|| compare_demoted_format_preference(left, right))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then(left.format_preference.cmp(&right.format_preference))
.then_with(|| compare_format_preference(left, right))
.then_with(|| compare_health(left, right, context.include_health))
.then(left.affinity_hash.cmp(&right.affinity_hash))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
@@ -91,7 +97,7 @@ pub(super) fn apply_load_balance_rotation(
while start < sorted_indices.len() {
let mut end = start + 1;
while end < sorted_indices.len()
&& candidates_share_priority_group(
&& candidates_share_load_balance_rotation_group(
&candidates[sorted_indices[start]],
&candidates[sorted_indices[end]],
context.priority_mode,
@@ -108,3 +114,14 @@ pub(super) fn apply_load_balance_rotation(
start = end;
}
}
fn candidates_share_load_balance_rotation_group(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
priority_mode: crate::SchedulerPriorityMode,
) -> bool {
candidates_share_priority_group(left, right, priority_mode)
&& left.capability_priority == right.capability_priority
&& left.demote_cross_format == right.demote_cross_format
&& left.format_preference == right.format_preference
}