Isolate core candidate selection helpers

crates/aether-scheduler-core/src/candidate/enumeration.rs

@@ -3,49 +3,10 @@ use std::collections::BTreeSet;
 use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
 use aether_data_contracts::DataLayerError;
 
-use super::capability::{
-    enabled_required_capabilities, requested_capability_priority_for_candidate_descriptors,
-};
 use super::types::{
     BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
 };
 
-pub fn build_minimal_candidate_selection(
-    input: BuildMinimalCandidateSelectionInput<'_>,
-) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
-    let priority_mode = input.priority_mode;
-    let affinity_key = input.affinity_key.map(str::to_string);
-    let required_capabilities = enabled_required_capabilities(input.required_capabilities);
-    let mut candidates = enumerate_minimal_candidate_selection(input)?;
-    let rankables = candidates
-        .iter()
-        .enumerate()
-        .map(|(index, candidate)| {
-            crate::SchedulerRankableCandidate::from_candidate(candidate, index)
-                .with_capability_priority(requested_capability_priority_for_candidate_descriptors(
-                    required_capabilities.iter().copied(),
-                    candidate,
-                ))
-                .with_affinity_hash(
-                    affinity_key
-                        .as_deref()
-                        .map(|key| crate::candidate_affinity_hash(key, candidate)),
-                )
-        })
-        .collect::<Vec<_>>();
-    crate::apply_scheduler_candidate_ranking(
-        &mut candidates,
-        &rankables,
-        crate::SchedulerRankingContext {
-            priority_mode,
-            ranking_mode: crate::SchedulerRankingMode::CacheAffinity,
-            include_health: false,
-            load_balance_seed: 0,
-        },
-    );
-    Ok(candidates)
-}
-
 pub fn enumerate_minimal_candidate_selection(
     input: BuildMinimalCandidateSelectionInput<'_>,
 ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {

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

@@ -2,20 +2,23 @@ pub mod capability;
 pub mod enumeration;
 pub mod identity;
 pub mod selectability;
+pub mod selection;
 pub mod types;
 
 pub use capability::{
     candidate_supports_required_capability, requested_capability_priority_for_candidate,
 };
 pub use enumeration::{
-    build_minimal_candidate_selection, collect_global_model_names_for_required_capability,
-    enumerate_minimal_candidate_selection,
+    collect_global_model_names_for_required_capability, enumerate_minimal_candidate_selection,
 };
 pub use identity::compare_candidates_by_priority_mode;
 pub use selectability::{
     auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
-    candidate_runtime_skip_reason_with_state, collect_selectable_candidates_from_keys,
-    reorder_candidates_by_scheduler_health, CandidateRuntimeSelectabilityInput,
+    candidate_runtime_skip_reason_with_state, CandidateRuntimeSelectabilityInput,
+};
+pub use selection::{
+    build_minimal_candidate_selection, collect_selectable_candidates_from_keys,
+    reorder_candidates_by_scheduler_health,
 };
 pub use types::{
     BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,

crates/aether-scheduler-core/src/candidate/selectability.rs

@@ -1,12 +1,9 @@
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::BTreeMap;
 
 use aether_data_contracts::repository::candidates::StoredRequestCandidate;
 use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
 
-use super::capability::{
-    enabled_required_capabilities, requested_capability_priority_for_candidate_descriptors,
-};
-use super::types::{SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode};
+use super::types::SchedulerMinimalCandidateSelectionCandidate;
 
 pub fn auth_api_key_concurrency_limit_reached(
     recent_candidates: &[StoredRequestCandidate],
@@ -22,82 +19,6 @@ pub fn auth_api_key_concurrency_limit_reached(
         >= concurrent_limit
 }
 
-pub fn collect_selectable_candidates_from_keys(
-    candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
-    selectable_keys: &BTreeSet<(String, String, String)>,
-    cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
-) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
-    let mut promoted = None;
-    let mut selected = Vec::with_capacity(candidates.len());
-    let mut emitted_keys = BTreeSet::new();
-
-    for candidate in candidates {
-        let key = crate::candidate_key(&candidate);
-        if !selectable_keys.contains(&key) || !emitted_keys.insert(key) {
-            continue;
-        }
-        if promoted.is_none()
-            && cached_affinity_target
-                .is_some_and(|target| crate::matches_affinity_target(&candidate, target))
-        {
-            promoted = Some(candidate);
-        } else {
-            selected.push(candidate);
-        }
-    }
-
-    if let Some(candidate) = promoted {
-        selected.insert(0, candidate);
-    }
-
-    selected
-}
-
-pub fn reorder_candidates_by_scheduler_health(
-    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
-    provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
-    required_capabilities: Option<&serde_json::Value>,
-    affinity_key: Option<&str>,
-    priority_mode: SchedulerPriorityMode,
-) {
-    let required_capabilities = enabled_required_capabilities(required_capabilities);
-    let rankables = candidates
-        .iter()
-        .enumerate()
-        .map(|(index, candidate)| {
-            crate::SchedulerRankableCandidate::from_candidate(candidate, index)
-                .with_capability_priority(requested_capability_priority_for_candidate_descriptors(
-                    required_capabilities.iter().copied(),
-                    candidate,
-                ))
-                .with_affinity_hash(
-                    affinity_key.map(|key| crate::candidate_affinity_hash(key, candidate)),
-                )
-                .with_health(
-                    provider_key_rpm_states
-                        .get(&candidate.key_id)
-                        .and_then(|key| {
-                            crate::provider_key_health_bucket(
-                                key,
-                                candidate.endpoint_api_format.as_str(),
-                            )
-                        }),
-                    candidate_provider_key_health_score(candidate, Some(provider_key_rpm_states)),
-                )
-        })
-        .collect::<Vec<_>>();
-    crate::apply_scheduler_candidate_ranking(
-        candidates,
-        &rankables,
-        crate::SchedulerRankingContext {
-            priority_mode,
-            ranking_mode: crate::SchedulerRankingMode::CacheAffinity,
-            include_health: true,
-            load_balance_seed: 0,
-        },
-    );
-}
-
 #[derive(Clone, Copy, Debug)]
 pub struct CandidateRuntimeSelectabilityInput<'a> {
     pub candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
@@ -190,15 +111,3 @@ pub fn candidate_runtime_skip_reason_with_state(
 
     None
 }
-
-fn candidate_provider_key_health_score(
-    candidate: &SchedulerMinimalCandidateSelectionCandidate,
-    provider_key_rpm_states: Option<&BTreeMap<String, StoredProviderCatalogKey>>,
-) -> f64 {
-    provider_key_rpm_states
-        .and_then(|states| states.get(&candidate.key_id))
-        .and_then(|key| {
-            crate::effective_provider_key_health_score(key, candidate.endpoint_api_format.as_str())
-        })
-        .unwrap_or(1.0)
-}

crates/aether-scheduler-core/src/candidate/selection.rs

@@ -0,0 +1,137 @@
+use std::collections::{BTreeMap, BTreeSet};
+
+use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
+use aether_data_contracts::DataLayerError;
+
+use super::capability::{
+    enabled_required_capabilities, requested_capability_priority_for_candidate_descriptors,
+};
+use super::enumeration::enumerate_minimal_candidate_selection;
+use super::types::{
+    BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
+    SchedulerPriorityMode,
+};
+
+pub fn build_minimal_candidate_selection(
+    input: BuildMinimalCandidateSelectionInput<'_>,
+) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
+    let priority_mode = input.priority_mode;
+    let affinity_key = input.affinity_key.map(str::to_string);
+    let required_capabilities = enabled_required_capabilities(input.required_capabilities);
+    let mut candidates = enumerate_minimal_candidate_selection(input)?;
+    let rankables = candidates
+        .iter()
+        .enumerate()
+        .map(|(index, candidate)| {
+            crate::SchedulerRankableCandidate::from_candidate(candidate, index)
+                .with_capability_priority(requested_capability_priority_for_candidate_descriptors(
+                    required_capabilities.iter().copied(),
+                    candidate,
+                ))
+                .with_affinity_hash(
+                    affinity_key
+                        .as_deref()
+                        .map(|key| crate::candidate_affinity_hash(key, candidate)),
+                )
+        })
+        .collect::<Vec<_>>();
+    crate::apply_scheduler_candidate_ranking(
+        &mut candidates,
+        &rankables,
+        crate::SchedulerRankingContext {
+            priority_mode,
+            ranking_mode: crate::SchedulerRankingMode::CacheAffinity,
+            include_health: false,
+            load_balance_seed: 0,
+        },
+    );
+    Ok(candidates)
+}
+
+pub fn collect_selectable_candidates_from_keys(
+    candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
+    selectable_keys: &BTreeSet<(String, String, String)>,
+    cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
+) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
+    let mut promoted = None;
+    let mut selected = Vec::with_capacity(candidates.len());
+    let mut emitted_keys = BTreeSet::new();
+
+    for candidate in candidates {
+        let key = crate::candidate_key(&candidate);
+        if !selectable_keys.contains(&key) || !emitted_keys.insert(key) {
+            continue;
+        }
+        if promoted.is_none()
+            && cached_affinity_target
+                .is_some_and(|target| crate::matches_affinity_target(&candidate, target))
+        {
+            promoted = Some(candidate);
+        } else {
+            selected.push(candidate);
+        }
+    }
+
+    if let Some(candidate) = promoted {
+        selected.insert(0, candidate);
+    }
+
+    selected
+}
+
+pub fn reorder_candidates_by_scheduler_health(
+    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
+    provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
+    required_capabilities: Option<&serde_json::Value>,
+    affinity_key: Option<&str>,
+    priority_mode: SchedulerPriorityMode,
+) {
+    let required_capabilities = enabled_required_capabilities(required_capabilities);
+    let rankables = candidates
+        .iter()
+        .enumerate()
+        .map(|(index, candidate)| {
+            crate::SchedulerRankableCandidate::from_candidate(candidate, index)
+                .with_capability_priority(requested_capability_priority_for_candidate_descriptors(
+                    required_capabilities.iter().copied(),
+                    candidate,
+                ))
+                .with_affinity_hash(
+                    affinity_key.map(|key| crate::candidate_affinity_hash(key, candidate)),
+                )
+                .with_health(
+                    provider_key_rpm_states
+                        .get(&candidate.key_id)
+                        .and_then(|key| {
+                            crate::provider_key_health_bucket(
+                                key,
+                                candidate.endpoint_api_format.as_str(),
+                            )
+                        }),
+                    candidate_provider_key_health_score(candidate, Some(provider_key_rpm_states)),
+                )
+        })
+        .collect::<Vec<_>>();
+    crate::apply_scheduler_candidate_ranking(
+        candidates,
+        &rankables,
+        crate::SchedulerRankingContext {
+            priority_mode,
+            ranking_mode: crate::SchedulerRankingMode::CacheAffinity,
+            include_health: true,
+            load_balance_seed: 0,
+        },
+    );
+}
+
+fn candidate_provider_key_health_score(
+    candidate: &SchedulerMinimalCandidateSelectionCandidate,
+    provider_key_rpm_states: Option<&BTreeMap<String, StoredProviderCatalogKey>>,
+) -> f64 {
+    provider_key_rpm_states
+        .and_then(|states| states.get(&candidate.key_id))
+        .and_then(|key| {
+            crate::effective_provider_key_health_score(key, candidate.endpoint_api_format.as_str())
+        })
+        .unwrap_or(1.0)
+}