fix scheduler affinity candidate selection

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

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

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

@@ -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));

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

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