Unify candidate ranking pipeline

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

@@ -0,0 +1,12 @@
+use std::cmp::Ordering;
+
+use super::types::SchedulerRankableCandidate;
+
+pub fn compare_format_state(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+) -> Ordering {
+    left.demote_cross_format
+        .cmp(&right.demote_cross_format)
+        .then(left.format_preference.cmp(&right.format_preference))
+}

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

@@ -0,0 +1,294 @@
+mod format;
+mod modes;
+mod priority;
+mod reasons;
+mod types;
+
+pub use format::compare_format_state;
+pub use modes::{apply_load_balance_rotation, compare_rankable_candidates};
+pub use priority::{
+    candidate_priority_slot, candidates_share_priority_group, compare_candidate_priority_slot,
+};
+pub use reasons::{
+    demoted_by as ranking_demoted_by, promoted_by as ranking_promoted_by,
+    RANKING_REASON_CACHED_AFFINITY, RANKING_REASON_CROSS_FORMAT, RANKING_REASON_LOCAL_TUNNEL,
+};
+pub use types::{
+    SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
+    SchedulerRankingOutcome, SchedulerTunnelAffinityBucket,
+};
+
+pub fn compare_candidate_identity_for_ranking(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+) -> std::cmp::Ordering {
+    left.provider_id
+        .cmp(&right.provider_id)
+        .then(left.endpoint_id.cmp(&right.endpoint_id))
+        .then(left.key_id.cmp(&right.key_id))
+        .then(
+            left.selected_provider_model_name
+                .cmp(&right.selected_provider_model_name),
+        )
+}
+
+pub fn scheduler_candidate_ranking_order(
+    candidates: &[SchedulerRankableCandidate],
+    context: SchedulerRankingContext,
+) -> Vec<usize> {
+    let mut order = (0..candidates.len()).collect::<Vec<_>>();
+    order.sort_by(|left, right| {
+        compare_rankable_candidates(&candidates[*left], &candidates[*right], context)
+    });
+    apply_load_balance_rotation(&mut order, candidates, context);
+    order
+}
+
+pub fn scheduler_ranking_outcomes(
+    candidates: &[SchedulerRankableCandidate],
+    context: SchedulerRankingContext,
+) -> Vec<SchedulerRankingOutcome> {
+    scheduler_candidate_ranking_order(candidates, context)
+        .into_iter()
+        .enumerate()
+        .map(|(ranking_index, original_index)| {
+            let candidate = &candidates[original_index];
+            SchedulerRankingOutcome {
+                original_index,
+                ranking_index,
+                priority_mode: context.priority_mode,
+                ranking_mode: context.ranking_mode,
+                priority_slot: candidate_priority_slot(candidate, context.priority_mode),
+                promoted_by: ranking_promoted_by(candidate, context.ranking_mode),
+                demoted_by: ranking_demoted_by(candidate),
+            }
+        })
+        .collect()
+}
+
+pub fn apply_scheduler_candidate_ranking<T>(
+    items: &mut [T],
+    candidates: &[SchedulerRankableCandidate],
+    context: SchedulerRankingContext,
+) -> Vec<SchedulerRankingOutcome> {
+    let outcomes = scheduler_ranking_outcomes(candidates, context);
+    apply_order(
+        items,
+        outcomes
+            .iter()
+            .map(|outcome| outcome.original_index)
+            .collect(),
+    );
+    outcomes
+}
+
+fn apply_order<T>(items: &mut [T], sorted_old_indices: Vec<usize>) {
+    if items.len() < 2 {
+        return;
+    }
+
+    let mut target_positions = vec![0usize; sorted_old_indices.len()];
+    for (new_position, old_position) in sorted_old_indices.into_iter().enumerate() {
+        target_positions[old_position] = new_position;
+    }
+
+    for index in 0..items.len() {
+        while target_positions[index] != index {
+            let target = target_positions[index];
+            items.swap(index, target);
+            target_positions.swap(index, target);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{SchedulerPriorityMode, SchedulerTunnelAffinityBucket};
+
+    fn candidate(
+        id: &str,
+        provider_priority: i32,
+        key_priority: i32,
+        global_key_priority: Option<i32>,
+    ) -> SchedulerRankableCandidate {
+        SchedulerRankableCandidate {
+            provider_id: format!("provider-{id}"),
+            endpoint_id: format!("endpoint-{id}"),
+            key_id: format!("key-{id}"),
+            selected_provider_model_name: "gpt-5".to_string(),
+            provider_priority,
+            key_internal_priority: key_priority,
+            key_global_priority_for_format: global_key_priority,
+            capability_priority: (0, 0),
+            cached_affinity_match: false,
+            affinity_hash: None,
+            tunnel_bucket: SchedulerTunnelAffinityBucket::Neutral,
+            demote_cross_format: false,
+            format_preference: (0, 0),
+            health_bucket: None,
+            health_score: 1.0,
+            original_index: 0,
+        }
+    }
+
+    fn ranked_ids(
+        candidates: &[SchedulerRankableCandidate],
+        context: SchedulerRankingContext,
+    ) -> Vec<String> {
+        scheduler_candidate_ranking_order(candidates, context)
+            .into_iter()
+            .map(|index| candidates[index].provider_id.clone())
+            .collect()
+    }
+
+    #[test]
+    fn provider_priority_mode_prefers_provider_priority_slot() {
+        let candidates = vec![
+            candidate("global", 10, 0, Some(0)),
+            candidate("provider", 0, 10, Some(10)),
+        ];
+
+        assert_eq!(
+            ranked_ids(
+                &candidates,
+                SchedulerRankingContext {
+                    priority_mode: SchedulerPriorityMode::Provider,
+                    ranking_mode: SchedulerRankingMode::FixedOrder,
+                    include_health: false,
+                    load_balance_seed: 0,
+                },
+            ),
+            vec!["provider-provider", "provider-global"]
+        );
+    }
+
+    #[test]
+    fn global_key_priority_mode_prefers_global_key_priority_slot() {
+        let candidates = vec![
+            candidate("provider", 0, 10, Some(10)),
+            candidate("global", 10, 0, Some(0)),
+        ];
+
+        assert_eq!(
+            ranked_ids(
+                &candidates,
+                SchedulerRankingContext {
+                    priority_mode: SchedulerPriorityMode::GlobalKey,
+                    ranking_mode: SchedulerRankingMode::FixedOrder,
+                    include_health: false,
+                    load_balance_seed: 0,
+                },
+            ),
+            vec!["provider-global", "provider-provider"]
+        );
+    }
+
+    #[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);
+
+        let mut higher_priority = candidate("higher", 0, 0, Some(0));
+        higher_priority.demote_cross_format = true;
+        higher_priority.format_preference = (9, 9);
+
+        assert_eq!(
+            ranked_ids(
+                &[lower_priority, higher_priority],
+                SchedulerRankingContext {
+                    priority_mode: SchedulerPriorityMode::Provider,
+                    ranking_mode: SchedulerRankingMode::FixedOrder,
+                    include_health: false,
+                    load_balance_seed: 0,
+                },
+            ),
+            vec!["provider-higher", "provider-lower"]
+        );
+    }
+
+    #[test]
+    fn cache_affinity_can_promote_cached_candidate_and_reports_reason() {
+        let high_priority = candidate("high", 0, 0, Some(0));
+        let mut cached = candidate("cached", 10, 0, Some(10));
+        cached.cached_affinity_match = true;
+        let candidates = vec![high_priority, cached];
+        let context = SchedulerRankingContext {
+            priority_mode: SchedulerPriorityMode::Provider,
+            ranking_mode: SchedulerRankingMode::CacheAffinity,
+            include_health: false,
+            load_balance_seed: 0,
+        };
+
+        let outcomes = scheduler_ranking_outcomes(&candidates, context);
+        assert_eq!(outcomes[0].original_index, 1);
+        assert_eq!(
+            outcomes[0].promoted_by,
+            Some(RANKING_REASON_CACHED_AFFINITY)
+        );
+    }
+
+    #[test]
+    fn cache_affinity_without_cache_hit_keeps_priority_before_tunnel() {
+        let mut higher_priority = candidate("higher", 0, 0, Some(0));
+        higher_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::RemoteTunnel;
+
+        let mut lower_priority = candidate("lower", 10, 0, Some(10));
+        lower_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::LocalTunnel;
+
+        assert_eq!(
+            ranked_ids(
+                &[lower_priority, higher_priority],
+                SchedulerRankingContext {
+                    priority_mode: SchedulerPriorityMode::Provider,
+                    ranking_mode: SchedulerRankingMode::CacheAffinity,
+                    include_health: false,
+                    load_balance_seed: 0,
+                },
+            ),
+            vec!["provider-higher", "provider-lower"]
+        );
+    }
+
+    #[test]
+    fn cache_affinity_keeps_cross_format_demotion_before_priority() {
+        let same_format_low_priority = candidate("same", 10, 0, Some(10));
+        let mut cross_format_high_priority = candidate("cross", 0, 0, Some(0));
+        cross_format_high_priority.demote_cross_format = true;
+
+        assert_eq!(
+            ranked_ids(
+                &[cross_format_high_priority, same_format_low_priority],
+                SchedulerRankingContext {
+                    priority_mode: SchedulerPriorityMode::Provider,
+                    ranking_mode: SchedulerRankingMode::CacheAffinity,
+                    include_health: false,
+                    load_balance_seed: 0,
+                },
+            ),
+            vec!["provider-same", "provider-cross"]
+        );
+    }
+
+    #[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));
+
+        assert_eq!(
+            ranked_ids(
+                &[first, second, third],
+                SchedulerRankingContext {
+                    priority_mode: SchedulerPriorityMode::Provider,
+                    ranking_mode: SchedulerRankingMode::LoadBalance,
+                    include_health: false,
+                    load_balance_seed: 1,
+                },
+            ),
+            vec!["provider-second", "provider-first", "provider-third"]
+        );
+    }
+}

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

@@ -0,0 +1,111 @@
+use std::cmp::Ordering;
+
+use super::types::{SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode};
+use super::{
+    candidates_share_priority_group, compare_candidate_identity_for_ranking,
+    compare_candidate_priority_slot, compare_format_state,
+};
+
+pub fn compare_rankable_candidates(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+    context: SchedulerRankingContext,
+) -> Ordering {
+    match context.ranking_mode {
+        SchedulerRankingMode::FixedOrder => compare_fixed_order(left, right, context),
+        SchedulerRankingMode::CacheAffinity => compare_cache_affinity(left, right, context),
+        SchedulerRankingMode::LoadBalance => compare_load_balance_base(left, right, context),
+    }
+}
+
+fn compare_fixed_order(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+    context: SchedulerRankingContext,
+) -> Ordering {
+    left.capability_priority
+        .cmp(&right.capability_priority)
+        .then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
+        .then_with(|| compare_format_state(left, right))
+        .then_with(|| compare_candidate_identity_for_ranking(left, right))
+        .then(left.original_index.cmp(&right.original_index))
+}
+
+fn compare_cache_affinity(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+    context: SchedulerRankingContext,
+) -> Ordering {
+    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_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_health(left, right, context.include_health))
+        .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_base(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+    context: SchedulerRankingContext,
+) -> Ordering {
+    left.capability_priority
+        .cmp(&right.capability_priority)
+        .then(left.demote_cross_format.cmp(&right.demote_cross_format))
+        .then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
+        .then(left.format_preference.cmp(&right.format_preference))
+        .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))
+        .then(left.original_index.cmp(&right.original_index))
+}
+
+fn compare_health(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+    include_health: bool,
+) -> Ordering {
+    if !include_health {
+        return Ordering::Equal;
+    }
+    right
+        .health_bucket
+        .cmp(&left.health_bucket)
+        .then_with(|| right.health_score.total_cmp(&left.health_score))
+}
+
+pub 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_priority_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;
+    }
+}

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

@@ -0,0 +1,52 @@
+use std::cmp::Ordering;
+
+use crate::SchedulerPriorityMode;
+
+use super::types::SchedulerRankableCandidate;
+
+pub fn candidate_priority_slot(
+    candidate: &SchedulerRankableCandidate,
+    priority_mode: SchedulerPriorityMode,
+) -> i32 {
+    match priority_mode {
+        SchedulerPriorityMode::Provider => candidate.provider_priority,
+        SchedulerPriorityMode::GlobalKey => {
+            candidate.key_global_priority_for_format.unwrap_or(i32::MAX)
+        }
+    }
+}
+
+pub fn compare_candidate_priority_slot(
+    left: &SchedulerRankableCandidate,
+    right: &SchedulerRankableCandidate,
+    priority_mode: SchedulerPriorityMode,
+) -> Ordering {
+    match priority_mode {
+        SchedulerPriorityMode::Provider => left
+            .provider_priority
+            .cmp(&right.provider_priority)
+            .then(left.key_internal_priority.cmp(&right.key_internal_priority)),
+        SchedulerPriorityMode::GlobalKey => left
+            .key_global_priority_for_format
+            .unwrap_or(i32::MAX)
+            .cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
+            .then(left.provider_priority.cmp(&right.provider_priority))
+            .then(left.key_internal_priority.cmp(&right.key_internal_priority)),
+    }
+}
+
+pub 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
+        }
+    }
+}

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

@@ -0,0 +1,28 @@
+use super::types::{
+    SchedulerRankableCandidate, SchedulerRankingMode, SchedulerTunnelAffinityBucket,
+};
+
+pub const RANKING_REASON_CACHED_AFFINITY: &str = "cached_affinity";
+pub const RANKING_REASON_LOCAL_TUNNEL: &str = "local_tunnel";
+pub const RANKING_REASON_CROSS_FORMAT: &str = "cross_format";
+
+pub fn promoted_by(
+    candidate: &SchedulerRankableCandidate,
+    ranking_mode: SchedulerRankingMode,
+) -> Option<&'static str> {
+    if ranking_mode == SchedulerRankingMode::CacheAffinity && candidate.cached_affinity_match {
+        return Some(RANKING_REASON_CACHED_AFFINITY);
+    }
+    if ranking_mode == SchedulerRankingMode::CacheAffinity
+        && candidate.tunnel_bucket == SchedulerTunnelAffinityBucket::LocalTunnel
+    {
+        return Some(RANKING_REASON_LOCAL_TUNNEL);
+    }
+    None
+}
+
+pub fn demoted_by(candidate: &SchedulerRankableCandidate) -> Option<&'static str> {
+    candidate
+        .demote_cross_format
+        .then_some(RANKING_REASON_CROSS_FORMAT)
+}

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

@@ -0,0 +1,131 @@
+use crate::{
+    ProviderKeyHealthBucket, SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
+};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
+pub enum SchedulerRankingMode {
+    FixedOrder,
+    #[default]
+    CacheAffinity,
+    LoadBalance,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
+pub enum SchedulerTunnelAffinityBucket {
+    LocalTunnel = 0,
+    #[default]
+    Neutral = 1,
+    RemoteTunnel = 2,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct SchedulerRankableCandidate {
+    pub provider_id: String,
+    pub endpoint_id: String,
+    pub key_id: String,
+    pub selected_provider_model_name: String,
+    pub provider_priority: i32,
+    pub key_internal_priority: i32,
+    pub key_global_priority_for_format: Option<i32>,
+    pub capability_priority: (u32, u32),
+    pub cached_affinity_match: bool,
+    pub affinity_hash: Option<u64>,
+    pub tunnel_bucket: SchedulerTunnelAffinityBucket,
+    pub demote_cross_format: bool,
+    pub format_preference: (u8, u8),
+    pub health_bucket: Option<ProviderKeyHealthBucket>,
+    pub health_score: f64,
+    pub original_index: usize,
+}
+
+impl SchedulerRankableCandidate {
+    pub fn from_candidate(
+        candidate: &SchedulerMinimalCandidateSelectionCandidate,
+        original_index: usize,
+    ) -> Self {
+        Self {
+            provider_id: candidate.provider_id.clone(),
+            endpoint_id: candidate.endpoint_id.clone(),
+            key_id: candidate.key_id.clone(),
+            selected_provider_model_name: candidate.selected_provider_model_name.clone(),
+            provider_priority: candidate.provider_priority,
+            key_internal_priority: candidate.key_internal_priority,
+            key_global_priority_for_format: candidate.key_global_priority_for_format,
+            capability_priority: (0, 0),
+            cached_affinity_match: false,
+            affinity_hash: None,
+            tunnel_bucket: SchedulerTunnelAffinityBucket::Neutral,
+            demote_cross_format: false,
+            format_preference: (0, 0),
+            health_bucket: None,
+            health_score: 1.0,
+            original_index,
+        }
+    }
+
+    pub fn with_capability_priority(mut self, value: (u32, u32)) -> Self {
+        self.capability_priority = value;
+        self
+    }
+
+    pub fn with_cached_affinity_match(mut self, value: bool) -> Self {
+        self.cached_affinity_match = value;
+        self
+    }
+
+    pub fn with_affinity_hash(mut self, value: Option<u64>) -> Self {
+        self.affinity_hash = value;
+        self
+    }
+
+    pub fn with_tunnel_bucket(mut self, value: SchedulerTunnelAffinityBucket) -> Self {
+        self.tunnel_bucket = value;
+        self
+    }
+
+    pub fn with_format_state(
+        mut self,
+        demote_cross_format: bool,
+        format_preference: (u8, u8),
+    ) -> Self {
+        self.demote_cross_format = demote_cross_format;
+        self.format_preference = format_preference;
+        self
+    }
+
+    pub fn with_health(mut self, bucket: Option<ProviderKeyHealthBucket>, score: f64) -> Self {
+        self.health_bucket = bucket;
+        self.health_score = score;
+        self
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct SchedulerRankingContext {
+    pub priority_mode: SchedulerPriorityMode,
+    pub ranking_mode: SchedulerRankingMode,
+    pub include_health: bool,
+    pub load_balance_seed: u64,
+}
+
+impl Default for SchedulerRankingContext {
+    fn default() -> Self {
+        Self {
+            priority_mode: SchedulerPriorityMode::Provider,
+            ranking_mode: SchedulerRankingMode::CacheAffinity,
+            include_health: false,
+            load_balance_seed: 0,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
+pub struct SchedulerRankingOutcome {
+    pub original_index: usize,
+    pub ranking_index: usize,
+    pub priority_mode: SchedulerPriorityMode,
+    pub ranking_mode: SchedulerRankingMode,
+    pub priority_slot: i32,
+    pub promoted_by: Option<&'static str>,
+    pub demoted_by: Option<&'static str>,
+}