Drop legacy candidate ordering helpers

This commit is contained in:
fawney19
2026-04-27 14:54:56 +08:00
parent 8d2cbf32df
commit 2101a4ecc7
7 changed files with 23 additions and 235 deletions

View File

@@ -1,6 +1,6 @@
pub(super) use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
compare_affinity_order, matches_affinity_target, SchedulerAffinityTarget,
matches_affinity_target, SchedulerAffinityTarget,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;

View File

@@ -381,7 +381,6 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
);
for pattern in [
"candidate_affinity_hash",
"compare_affinity_order",
"matches_affinity_target",
"candidate_key",
] {
@@ -483,14 +482,12 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
assert!(
workspace_file_exists("crates/aether-scheduler-core/src/candidate/selection.rs"),
"core candidate/selection.rs should host legacy selection compatibility helpers"
"core candidate/selection.rs should host legacy minimal selection compatibility helper"
);
let core_candidate_selection =
read_workspace_file("crates/aether-scheduler-core/src/candidate/selection.rs");
for expected in [
"build_minimal_candidate_selection",
"collect_selectable_candidates_from_keys",
"reorder_candidates_by_scheduler_health",
"apply_scheduler_candidate_ranking",
] {
assert!(
@@ -498,6 +495,16 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"core candidate/selection.rs should host compatibility helper {expected}"
);
}
for forbidden in [
"collect_selectable_candidates_from_keys",
"reorder_candidates_by_scheduler_health",
"compare_candidates_by_priority_mode",
] {
assert!(
!core_candidate_selection.contains(forbidden),
"core candidate/selection.rs should not keep removed legacy helper {forbidden}"
);
}
let affinity_cache = read_workspace_file("apps/aether-gateway/src/cache/scheduler_affinity.rs");
assert!(

View File

@@ -29,18 +29,6 @@ pub fn build_scheduler_affinity_cache_key_for_api_key_id(
))
}
pub fn compare_affinity_order(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
affinity_key: Option<&str>,
) -> std::cmp::Ordering {
let Some(affinity_key) = affinity_key else {
return std::cmp::Ordering::Equal;
};
candidate_affinity_hash(affinity_key, left).cmp(&candidate_affinity_hash(affinity_key, right))
}
pub fn candidate_affinity_hash(
affinity_key: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
@@ -82,7 +70,7 @@ pub fn candidate_key(
mod tests {
use super::{
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
compare_affinity_order, matches_affinity_target, SchedulerAffinityTarget,
matches_affinity_target, SchedulerAffinityTarget,
};
use crate::SchedulerMinimalCandidateSelectionCandidate;
@@ -133,7 +121,7 @@ mod tests {
}
#[test]
fn affinity_hash_and_order_are_candidate_specific() {
fn affinity_hash_is_candidate_specific() {
let left = sample_candidate("1");
let right = sample_candidate("2");
@@ -141,14 +129,6 @@ mod tests {
candidate_affinity_hash("api-key-1", &left),
candidate_affinity_hash("api-key-1", &right)
);
assert_ne!(
compare_affinity_order(&left, &right, Some("api-key-1")),
std::cmp::Ordering::Equal
);
assert_eq!(
compare_affinity_order(&left, &right, None),
std::cmp::Ordering::Equal
);
}
#[test]

View File

@@ -1,39 +0,0 @@
use super::types::{SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode};
pub fn compare_candidates_by_priority_mode(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
affinity_key: Option<&str>,
) -> std::cmp::Ordering {
match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then_with(|| compare_candidate_identity(left, right)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| compare_candidate_identity(left, right)),
}
}
pub(crate) fn compare_candidate_identity(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
) -> 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),
)
}

View File

@@ -1,6 +1,5 @@
pub mod capability;
pub mod enumeration;
pub mod identity;
pub mod selectability;
pub mod selection;
pub mod types;
@@ -11,15 +10,11 @@ pub use capability::{
pub use enumeration::{
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, CandidateRuntimeSelectabilityInput,
};
pub use selection::{
build_minimal_candidate_selection, collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health,
};
pub use selection::build_minimal_candidate_selection;
pub use types::{
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
@@ -27,7 +22,7 @@ pub use types::{
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use std::collections::BTreeMap;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
@@ -40,10 +35,9 @@ mod tests {
use super::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability,
collect_selectable_candidates_from_keys, reorder_candidates_by_scheduler_health,
BuildMinimalCandidateSelectionInput, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
collect_global_model_names_for_required_capability, BuildMinimalCandidateSelectionInput,
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
use crate::SchedulerAuthConstraints;
@@ -288,66 +282,6 @@ mod tests {
assert_eq!(candidates[1].key_id, "key-1");
}
#[test]
fn reorders_candidates_by_health_before_affinity_tiebreak() {
let mut candidates = vec![
sample_candidate("1", None),
sample_candidate("2", None),
sample_candidate("3", None),
];
let provider_key_rpm_states = BTreeMap::from([
("key-1".to_string(), sample_key("1", 0.95)),
("key-2".to_string(), sample_key("2", 0.40)),
("key-3".to_string(), sample_key("3", 0.95)),
]);
reorder_candidates_by_scheduler_health(
&mut candidates,
&provider_key_rpm_states,
None,
Some("api-key-1"),
SchedulerPriorityMode::GlobalKey,
);
assert_ne!(candidates[0].key_id, "key-2");
assert_ne!(candidates[1].key_id, "key-2");
assert_eq!(candidates[2].key_id, "key-2");
}
#[test]
fn collects_selectable_candidates_with_affinity_priority_and_dedup() {
let candidates = vec![
sample_candidate("1", None),
sample_candidate("2", None),
sample_candidate("1", None),
];
let selectable_keys = BTreeSet::from([
(
"provider-1".to_string(),
"endpoint-1".to_string(),
"key-1".to_string(),
),
(
"provider-2".to_string(),
"endpoint-2".to_string(),
"key-2".to_string(),
),
]);
let selected = collect_selectable_candidates_from_keys(
candidates,
&selectable_keys,
Some(&crate::SchedulerAffinityTarget {
provider_id: "provider-2".to_string(),
endpoint_id: "endpoint-2".to_string(),
key_id: "key-2".to_string(),
}),
);
assert_eq!(selected.len(), 2);
assert_eq!(selected[0].key_id, "key-2");
assert_eq!(selected[1].key_id, "key-1");
}
#[test]
fn candidate_selectability_respects_provider_concurrency_limit() {
let recent_candidates = vec![stored_candidate("one", RequestCandidateStatus::Pending, 95)];

View File

@@ -1,6 +1,3 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data_contracts::DataLayerError;
use super::capability::{
@@ -9,7 +6,6 @@ use super::capability::{
use super::enumeration::enumerate_minimal_candidate_selection;
use super::types::{
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
pub fn build_minimal_candidate_selection(
@@ -47,91 +43,3 @@ pub fn build_minimal_candidate_selection(
);
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)
}

View File

@@ -9,7 +9,7 @@ mod request_candidate;
pub use affinity::{
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
compare_affinity_order, matches_affinity_target, SchedulerAffinityTarget,
matches_affinity_target, SchedulerAffinityTarget,
};
pub use auth::{
auth_constraints_allow_api_format, auth_constraints_allow_model,
@@ -19,11 +19,9 @@ pub use candidate::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_runtime_skip_reason_with_state,
candidate_supports_required_capability, collect_global_model_names_for_required_capability,
collect_selectable_candidates_from_keys, compare_candidates_by_priority_mode,
enumerate_minimal_candidate_selection, reorder_candidates_by_scheduler_health,
requested_capability_priority_for_candidate, BuildMinimalCandidateSelectionInput,
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
enumerate_minimal_candidate_selection, requested_capability_priority_for_candidate,
BuildMinimalCandidateSelectionInput, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,