mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Drop legacy candidate ordering helpers
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
}
|
||||
@@ -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)];
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user