mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Isolate core candidate selection helpers
This commit is contained in:
@@ -453,6 +453,52 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
|
|||||||
"candidate/mod.rs should depend directly on core model helper namespace"
|
"candidate/mod.rs should depend directly on core model helper namespace"
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let core_candidate_enumeration =
|
||||||
|
read_workspace_file("crates/aether-scheduler-core/src/candidate/enumeration.rs");
|
||||||
|
for forbidden in [
|
||||||
|
"apply_scheduler_candidate_ranking",
|
||||||
|
"SchedulerRankableCandidate",
|
||||||
|
"candidate_affinity_hash",
|
||||||
|
"requested_capability_priority_for_candidate_descriptors",
|
||||||
|
] {
|
||||||
|
assert!(
|
||||||
|
!core_candidate_enumeration.contains(forbidden),
|
||||||
|
"core candidate/enumeration.rs should only enumerate theoretical candidates, not rank with {forbidden}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let core_candidate_selectability =
|
||||||
|
read_workspace_file("crates/aether-scheduler-core/src/candidate/selectability.rs");
|
||||||
|
for forbidden in [
|
||||||
|
"apply_scheduler_candidate_ranking",
|
||||||
|
"with_affinity_hash",
|
||||||
|
"collect_selectable_candidates_from_keys",
|
||||||
|
"reorder_candidates_by_scheduler_health",
|
||||||
|
] {
|
||||||
|
assert!(
|
||||||
|
!core_candidate_selectability.contains(forbidden),
|
||||||
|
"core candidate/selectability.rs should only decide selectability, not rank with {forbidden}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
workspace_file_exists("crates/aether-scheduler-core/src/candidate/selection.rs"),
|
||||||
|
"core candidate/selection.rs should host legacy selection compatibility helpers"
|
||||||
|
);
|
||||||
|
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!(
|
||||||
|
core_candidate_selection.contains(expected),
|
||||||
|
"core candidate/selection.rs should host compatibility helper {expected}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
let affinity_cache = read_workspace_file("apps/aether-gateway/src/cache/scheduler_affinity.rs");
|
let affinity_cache = read_workspace_file("apps/aether-gateway/src/cache/scheduler_affinity.rs");
|
||||||
assert!(
|
assert!(
|
||||||
affinity_cache.contains("aether_scheduler_core::SchedulerAffinityTarget"),
|
affinity_cache.contains("aether_scheduler_core::SchedulerAffinityTarget"),
|
||||||
|
|||||||
@@ -3,49 +3,10 @@ use std::collections::BTreeSet;
|
|||||||
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
|
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
|
||||||
use aether_data_contracts::DataLayerError;
|
use aether_data_contracts::DataLayerError;
|
||||||
|
|
||||||
use super::capability::{
|
|
||||||
enabled_required_capabilities, requested_capability_priority_for_candidate_descriptors,
|
|
||||||
};
|
|
||||||
use super::types::{
|
use super::types::{
|
||||||
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
|
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(
|
pub fn enumerate_minimal_candidate_selection(
|
||||||
input: BuildMinimalCandidateSelectionInput<'_>,
|
input: BuildMinimalCandidateSelectionInput<'_>,
|
||||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
|
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
|
||||||
|
|||||||
@@ -2,20 +2,23 @@ pub mod capability;
|
|||||||
pub mod enumeration;
|
pub mod enumeration;
|
||||||
pub mod identity;
|
pub mod identity;
|
||||||
pub mod selectability;
|
pub mod selectability;
|
||||||
|
pub mod selection;
|
||||||
pub mod types;
|
pub mod types;
|
||||||
|
|
||||||
pub use capability::{
|
pub use capability::{
|
||||||
candidate_supports_required_capability, requested_capability_priority_for_candidate,
|
candidate_supports_required_capability, requested_capability_priority_for_candidate,
|
||||||
};
|
};
|
||||||
pub use enumeration::{
|
pub use enumeration::{
|
||||||
build_minimal_candidate_selection, collect_global_model_names_for_required_capability,
|
collect_global_model_names_for_required_capability, enumerate_minimal_candidate_selection,
|
||||||
enumerate_minimal_candidate_selection,
|
|
||||||
};
|
};
|
||||||
pub use identity::compare_candidates_by_priority_mode;
|
pub use identity::compare_candidates_by_priority_mode;
|
||||||
pub use selectability::{
|
pub use selectability::{
|
||||||
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
|
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
|
||||||
candidate_runtime_skip_reason_with_state, collect_selectable_candidates_from_keys,
|
candidate_runtime_skip_reason_with_state, CandidateRuntimeSelectabilityInput,
|
||||||
reorder_candidates_by_scheduler_health, CandidateRuntimeSelectabilityInput,
|
};
|
||||||
|
pub use selection::{
|
||||||
|
build_minimal_candidate_selection, collect_selectable_candidates_from_keys,
|
||||||
|
reorder_candidates_by_scheduler_health,
|
||||||
};
|
};
|
||||||
pub use types::{
|
pub use types::{
|
||||||
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
|
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
|
||||||
|
|||||||
@@ -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::candidates::StoredRequestCandidate;
|
||||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
|
|
||||||
use super::capability::{
|
use super::types::SchedulerMinimalCandidateSelectionCandidate;
|
||||||
enabled_required_capabilities, requested_capability_priority_for_candidate_descriptors,
|
|
||||||
};
|
|
||||||
use super::types::{SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode};
|
|
||||||
|
|
||||||
pub fn auth_api_key_concurrency_limit_reached(
|
pub fn auth_api_key_concurrency_limit_reached(
|
||||||
recent_candidates: &[StoredRequestCandidate],
|
recent_candidates: &[StoredRequestCandidate],
|
||||||
@@ -22,82 +19,6 @@ pub fn auth_api_key_concurrency_limit_reached(
|
|||||||
>= concurrent_limit
|
>= 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)]
|
#[derive(Clone, Copy, Debug)]
|
||||||
pub struct CandidateRuntimeSelectabilityInput<'a> {
|
pub struct CandidateRuntimeSelectabilityInput<'a> {
|
||||||
pub candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
|
pub candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
|
||||||
@@ -190,15 +111,3 @@ pub fn candidate_runtime_skip_reason_with_state(
|
|||||||
|
|
||||||
None
|
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)
|
|
||||||
}
|
|
||||||
|
|||||||
137
crates/aether-scheduler-core/src/candidate/selection.rs
Normal file
137
crates/aether-scheduler-core/src/candidate/selection.rs
Normal file
@@ -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)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user