Refine load balance candidate ranking

apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs

@@ -9,6 +9,7 @@ use async_trait::async_trait;
 use tracing::warn;
 
 use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
+use crate::clock::current_unix_ms;
 use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
 use crate::scheduler::config::{
     read_scheduler_ordering_config, SchedulerOrderingConfig, SchedulerSchedulingMode,
@@ -176,6 +177,7 @@ fn ai_ranking_context_config(ordering_config: SchedulerOrderingConfig) -> AiRank
     AiRankingContextConfig {
         priority_mode: ordering_config.priority_mode,
         scheduling_mode: ai_ranking_scheduling_mode(ordering_config.scheduling_mode),
+        load_balance_seed: current_unix_ms(),
     }
 }
 

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

@@ -606,7 +606,7 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
 }
 
 #[tokio::test]
-async fn load_balance_rotates_same_priority_group_and_ignores_cached_affinity() {
+async fn load_balance_ignores_provider_priority_and_cached_affinity() {
     let mut first = sample_row();
     first.provider_id = "provider-a".to_string();
     first.provider_name = "provider-a".to_string();
@@ -623,9 +623,9 @@ async fn load_balance_rotates_same_priority_group_and_ignores_cached_affinity()
     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 = 100;
     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": 100}));
 
     let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
         first, second,
@@ -664,6 +664,29 @@ async fn load_balance_rotates_same_priority_group_and_ignores_cached_affinity()
     )
     .await
     .expect("first pass should succeed");
+    let mut provider_b_first_seed = None;
+    for seed in 101..600 {
+        let pass = collect_selectable_candidates(
+            state.data.as_ref(),
+            &state,
+            "openai:chat",
+            "gpt-4.1",
+            false,
+            Some(&auth_snapshot),
+            seed,
+        )
+        .await
+        .expect("seeded pass should succeed");
+        if pass
+            .first()
+            .is_some_and(|candidate| candidate.provider_id == "provider-b")
+        {
+            provider_b_first_seed = Some(seed);
+            break;
+        }
+    }
+    let provider_b_first_seed = provider_b_first_seed
+        .expect("test seed should allow provider-b to win despite lower priority");
     let second_pass = collect_selectable_candidates(
         state.data.as_ref(),
         &state,
@@ -671,27 +694,14 @@ async fn load_balance_rotates_same_priority_group_and_ignores_cached_affinity()
         "gpt-4.1",
         false,
         Some(&auth_snapshot),
-        101,
+        provider_b_first_seed,
     )
     .await
     .expect("second pass should succeed");
 
     assert_eq!(first_pass.len(), 2);
     assert_eq!(second_pass.len(), 2);
-    assert_ne!(first_pass[0].provider_id, second_pass[0].provider_id);
+    assert_eq!(second_pass[0].provider_id, "provider-b");
-    assert!(
-        first_pass[0].provider_id != "provider-b" || second_pass[0].provider_id != "provider-b"
-    );
-    assert_ne!(
-        first_pass
-            .iter()
-            .map(|candidate| candidate.provider_id.as_str())
-            .collect::<Vec<_>>(),
-        second_pass
-            .iter()
-            .map(|candidate| candidate.provider_id.as_str())
-            .collect::<Vec<_>>()
-    );
 }
 
 #[tokio::test]

crates/aether-ai-serving/src/candidate_ranking.rs

@@ -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)
-        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
         .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,
     }
 }
 

crates/aether-scheduler-core/src/ranking/mod.rs

@@ -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() {
+    fn load_balance_provider_mode_randomizes_providers_then_uses_internal_key_priority() {
-        let first = candidate("first", 0, 0, Some(0));
+        let mut provider_a_primary = candidate("a-primary", 0, 0, Some(0));
-        let second = candidate("second", 0, 0, Some(0));
+        provider_a_primary.provider_id = "provider-a".to_string();
-        let third = candidate("third", 10, 0, Some(10));
+        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"]
         );
     }
 }

crates/aether-scheduler-core/src/ranking/modes.rs

@@ -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_with(|| compare_load_balance_distribution(left, right, context))
-        .then(left.affinity_hash.cmp(&right.affinity_hash))
         .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(
+fn compare_affinity_or_seeded_hash(
-    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(
     left: &SchedulerRankableCandidate,
     right: &SchedulerRankableCandidate,
-    priority_mode: crate::SchedulerPriorityMode,
+    seed: u64,
-) -> bool {
+) -> Ordering {
-    candidates_share_priority_group(left, right, priority_mode)
+    match (left.affinity_hash, right.affinity_hash) {
-        && left.capability_priority == right.capability_priority
+        (Some(left_hash), Some(right_hash)) => left_hash.cmp(&right_hash),
-        && left.demote_cross_format == right.demote_cross_format
+        _ => compare_seeded_candidate_hash(left, right, seed, "tie"),
-        && left.format_preference == right.format_preference
+    }
+}
+
+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],
+    ])
 }

crates/aether-scheduler-core/src/ranking/priority.rs

@@ -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
-        }
-    }
-}