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

Refine load balance candidate ranking

Browse Source
This commit is contained in:
fawney19
2026-05-11 02:32:11 +08:00
parent ef4cc55c9a
commit cc4512fbbb
6 changed files with 256 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
crates/aether-ai-serving/src/candidate_ranking.rs
View File

@@ -17,6 +17,7 @@ pub enum AiRankingSchedulingMode {
pub struct AiRankingContextConfig {
pub priority_mode: SchedulerPriorityMode,
pub scheduling_mode: AiRankingSchedulingMode,
pub load_balance_seed: u64,
}
#[derive(Debug, Clone, Copy, PartialEq)]
@@ -81,15 +82,7 @@ pub fn build_ai_rankable_candidate(
)
.unwrap_or((u8::MAX, u8::MAX));
let mut rankable =
SchedulerRankableCandidate::from_candidate(parts.candidate, parts.original_index);
// The scheduler order is the upstream tie-breaker; AI serving only adds transport facts.
rankable.provider_id.clear();
rankable.endpoint_id.clear();
rankable.key_id.clear();
rankable.selected_provider_model_name.clear();
rankable
SchedulerRankableCandidate::from_candidate(parts.candidate, parts.original_index)
.with_capability_priority(requested_capability_priority_for_candidate(
parts.required_capabilities,
parts.candidate,
@@ -107,7 +100,7 @@ pub fn ai_ranking_context(config: AiRankingContextConfig) -> SchedulerRankingCon
priority_mode: config.priority_mode,
ranking_mode: ai_ranking_mode(config.scheduling_mode),
include_health: false,
load_balance_seed: 0,
load_balance_seed: config.load_balance_seed,
}
}

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

@@ -4,7 +4,7 @@ mod priority;
mod reasons;
mod types;
use modes::{apply_load_balance_rotation, compare_rankable_candidates};
use modes::compare_rankable_candidates;
use priority::candidate_priority_slot;
use reasons::{demoted_by as ranking_demoted_by, promoted_by as ranking_promoted_by};
pub use reasons::{
@@ -37,7 +37,6 @@ fn scheduler_candidate_ranking_order(
order.sort_by(|left, right| {
compare_rankable_candidates(&candidates[*left], &candidates[*right], context)
});
apply_load_balance_rotation(&mut order, candidates, context);
order
}
@@ -139,6 +138,16 @@ mod tests {
.collect()
}
fn ranked_keys(
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<String> {
scheduler_candidate_ranking_order(candidates, context)
.into_iter()
.map(|index| candidates[index].key_id.clone())
.collect()
}
#[test]
fn provider_priority_mode_prefers_provider_priority_slot() {
let candidates = vec![
@@ -335,6 +344,40 @@ mod tests {
);
}
#[test]
fn fixed_order_randomizes_equal_priority_ties_without_crossing_priority_slots() {
let first = candidate("first", 0, 0, Some(0));
let second = candidate("second", 0, 0, Some(0));
let lower_priority = candidate("lower", 10, 0, Some(10));
let first_seed_order = ranked_ids(
&[first.clone(), second.clone(), lower_priority.clone()],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
);
let alternate_order = (1..128)
.map(|seed| {
ranked_ids(
&[first.clone(), second.clone(), lower_priority.clone()],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: seed,
},
)
})
.find(|order| order[0] != first_seed_order[0])
.expect("equal priority tie should vary by seed");
assert_eq!(first_seed_order[2], "provider-lower");
assert_eq!(alternate_order[2], "provider-lower");
}
#[test]
fn load_balance_does_not_rotate_across_cross_format_demotion_group() {
let same_format = candidate("same", 0, 0, Some(0));
@@ -356,22 +399,89 @@ mod tests {
}
#[test]
fn load_balance_rotates_only_within_same_priority_group() {
let first = candidate("first", 0, 0, Some(0));
let second = candidate("second", 0, 0, Some(0));
let third = candidate("third", 10, 0, Some(10));
fn load_balance_provider_mode_randomizes_providers_then_uses_internal_key_priority() {
let mut provider_a_primary = candidate("a-primary", 0, 0, Some(0));
provider_a_primary.provider_id = "provider-a".to_string();
provider_a_primary.key_id = "key-a-primary".to_string();
let mut provider_a_secondary = candidate("a-secondary", 0, 10, Some(10));
provider_a_secondary.provider_id = "provider-a".to_string();
provider_a_secondary.key_id = "key-a-secondary".to_string();
let mut provider_b = candidate("b", 100, 0, Some(100));
provider_b.provider_id = "provider-b".to_string();
provider_b.key_id = "key-b".to_string();
let candidates = vec![provider_a_secondary, provider_b, provider_a_primary];
let seed = (0..512)
.find(|seed| {
ranked_keys(
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::LoadBalance,
include_health: false,
load_balance_seed: *seed,
},
)
.first()
.is_some_and(|key| key == "key-b")
})
.expect("test seed should put provider-b first");
let order = ranked_keys(
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::LoadBalance,
include_health: false,
load_balance_seed: seed,
},
);
assert_eq!(order[0], "key-b");
let primary_index = order
.iter()
.position(|key| key == "key-a-primary")
.expect("primary key should be ranked");
let secondary_index = order
.iter()
.position(|key| key == "key-a-secondary")
.expect("secondary key should be ranked");
assert!(primary_index < secondary_index);
}
#[test]
fn load_balance_global_key_mode_randomizes_keys_ignoring_global_priority() {
let high_priority = candidate("high", 100, 0, Some(0));
let low_priority = candidate("low", 0, 0, Some(100));
let candidates = vec![high_priority, low_priority];
let seed = (0..512)
.find(|seed| {
ranked_ids(
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::GlobalKey,
ranking_mode: SchedulerRankingMode::LoadBalance,
include_health: false,
load_balance_seed: *seed,
},
)
.first()
.is_some_and(|provider| provider == "provider-low")
})
.expect("test seed should put lower global-priority key first");
assert_eq!(
ranked_ids(
&[first, second, third],
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
priority_mode: SchedulerPriorityMode::GlobalKey,
ranking_mode: SchedulerRankingMode::LoadBalance,
include_health: false,
load_balance_seed: 1,
load_balance_seed: seed,
},
),
vec!["provider-second", "provider-first", "provider-third"]
vec!["provider-low", "provider-high"]
);
}
}

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

@@ -1,10 +1,12 @@
use std::cmp::Ordering;
use sha2::{Digest, Sha256};
use super::compare_candidate_identity_for_ranking;
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::priority::compare_candidate_priority_slot;
use super::types::{SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode};
pub(super) fn compare_rankable_candidates(
@@ -30,6 +32,7 @@ fn compare_fixed_order(
.then_with(|| compare_demoted_format_preference(left, right))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then_with(|| compare_format_preference(left, right))
.then_with(|| compare_seeded_candidate_hash(left, right, context.load_balance_seed, "tie"))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
@@ -48,7 +51,7 @@ fn compare_cache_affinity(
.then(left.tunnel_bucket.cmp(&right.tunnel_bucket))
.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_affinity_or_seeded_hash(left, right, context.load_balance_seed))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
@@ -62,14 +65,37 @@ fn compare_load_balance_base(
.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_preference(left, right))
.then_with(|| compare_health(left, right, context.include_health))
.then(left.affinity_hash.cmp(&right.affinity_hash))
.then_with(|| compare_load_balance_distribution(left, right, context))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
fn compare_load_balance_distribution(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
match context.priority_mode {
crate::SchedulerPriorityMode::Provider => {
compare_seeded_provider_hash(left, right, context.load_balance_seed)
.then_with(|| {
if left.provider_id == right.provider_id {
left.key_internal_priority.cmp(&right.key_internal_priority)
} else {
Ordering::Equal
}
})
.then_with(|| {
compare_seeded_candidate_hash(left, right, context.load_balance_seed, "key")
})
}
crate::SchedulerPriorityMode::GlobalKey => {
compare_seeded_candidate_hash(left, right, context.load_balance_seed, "global-key")
}
}
}
fn compare_health(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
@@ -84,44 +110,78 @@ fn compare_health(
.then_with(|| right.health_score.total_cmp(&left.health_score))
}
pub(super) fn apply_load_balance_rotation(
sorted_indices: &mut [usize],
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) {
if context.ranking_mode != SchedulerRankingMode::LoadBalance || sorted_indices.len() < 2 {
return;
}
let mut start = 0usize;
while start < sorted_indices.len() {
let mut end = start + 1;
while end < sorted_indices.len()
&& candidates_share_load_balance_rotation_group(
&candidates[sorted_indices[start]],
&candidates[sorted_indices[end]],
context.priority_mode,
)
{
end += 1;
}
let group_len = end - start;
if group_len > 1 {
let offset = usize::try_from(context.load_balance_seed).unwrap_or(0) % group_len;
sorted_indices[start..end].rotate_left(offset);
}
start = end;
}
}
fn candidates_share_load_balance_rotation_group(
fn compare_affinity_or_seeded_hash(
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
seed: u64,
) -> Ordering {
match (left.affinity_hash, right.affinity_hash) {
(Some(left_hash), Some(right_hash)) => left_hash.cmp(&right_hash),
_ => compare_seeded_candidate_hash(left, right, seed, "tie"),
}
}
fn compare_seeded_provider_hash(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
seed: u64,
) -> Ordering {
seeded_rank_hash(seed, "provider", [left.provider_id.as_str()], 0).cmp(&seeded_rank_hash(
seed,
"provider",
[right.provider_id.as_str()],
0,
))
}
fn compare_seeded_candidate_hash(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
seed: u64,
salt: &str,
) -> Ordering {
seeded_rank_hash(
seed,
salt,
[
left.provider_id.as_str(),
left.endpoint_id.as_str(),
left.key_id.as_str(),
left.selected_provider_model_name.as_str(),
],
left.original_index,
)
.cmp(&seeded_rank_hash(
seed,
salt,
[
right.provider_id.as_str(),
right.endpoint_id.as_str(),
right.key_id.as_str(),
right.selected_provider_model_name.as_str(),
],
right.original_index,
))
}
fn seeded_rank_hash<'a>(
seed: u64,
salt: &str,
parts: impl IntoIterator<Item = &'a str>,
original_index: usize,
) -> u64 {
let mut hasher = Sha256::new();
hasher.update(seed.to_be_bytes());
hasher.update(b":");
hasher.update(salt.as_bytes());
for part in parts {
hasher.update(b":");
hasher.update(part.as_bytes());
}
hasher.update(b":");
hasher.update(original_index.to_be_bytes());
let digest = hasher.finalize();
u64::from_be_bytes([
digest[0], digest[1], digest[2], digest[3], digest[4], digest[5], digest[6], digest[7],
])
}

16
crates/aether-scheduler-core/src/ranking/priority.rs
View File

@@ -34,19 +34,3 @@ pub(super) fn compare_candidate_priority_slot(
.then(left.key_internal_priority.cmp(&right.key_internal_priority)),
}
}
pub(super) fn candidates_share_priority_group(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
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
}
}
}