Remove ranked minimal selection compatibility helper

This commit is contained in:
fawney19
2026-04-27 15:52:26 +08:00
parent e5e09b49f4
commit e9f03d8d29
11 changed files with 103 additions and 161 deletions

View File

@@ -65,26 +65,7 @@ pub fn requested_capability_priority_for_candidate(
)
}
pub(crate) fn enabled_required_capabilities(
required_capabilities: Option<&serde_json::Value>,
) -> Vec<RequiredCapabilityDescriptor<'_>> {
let Some(required_capabilities) = required_capabilities.and_then(serde_json::Value::as_object)
else {
return Vec::new();
};
required_capabilities
.iter()
.filter_map(|(capability, value)| {
requested_capability_is_enabled(value).then_some(RequiredCapabilityDescriptor {
name: capability.as_str(),
compatible: requested_capability_is_compatible(capability),
})
})
.collect()
}
pub(crate) fn requested_capability_priority_for_candidate_descriptors<'a, I>(
fn requested_capability_priority_for_candidate_descriptors<'a, I>(
required_capabilities: I,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (u32, u32)

View File

@@ -4,13 +4,13 @@ use aether_data_contracts::repository::candidate_selection::StoredMinimalCandida
use aether_data_contracts::DataLayerError;
use super::types::{
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
};
pub fn enumerate_minimal_candidate_selection(
input: BuildMinimalCandidateSelectionInput<'_>,
input: EnumerateMinimalCandidateSelectionInput<'_>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let BuildMinimalCandidateSelectionInput {
let EnumerateMinimalCandidateSelectionInput {
rows,
normalized_api_format,
requested_model_name,

View File

@@ -1,7 +1,6 @@
pub mod capability;
pub mod enumeration;
pub mod selectability;
pub mod selection;
pub mod types;
pub use capability::{
@@ -14,9 +13,8 @@ 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_ranked_minimal_candidate_selection;
pub use types::{
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
@@ -33,11 +31,10 @@ mod tests {
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use super::{
auth_api_key_concurrency_limit_reached, build_ranked_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability, BuildMinimalCandidateSelectionInput,
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
candidate_supports_required_capability, collect_global_model_names_for_required_capability,
CandidateRuntimeSelectabilityInput, EnumerateMinimalCandidateSelectionInput,
SchedulerMinimalCandidateSelectionCandidate,
};
use crate::SchedulerAuthConstraints;
@@ -176,7 +173,7 @@ mod tests {
}
#[test]
fn builds_ranked_minimal_candidate_selection_with_auth_constraints() {
fn enumerates_minimal_candidate_selection_with_auth_constraints() {
let mut disallowed = sample_row("2");
disallowed.provider_id = "provider-blocked".to_string();
disallowed.provider_name = "Blocked".to_string();
@@ -187,7 +184,7 @@ mod tests {
allowed_models: Some(vec!["gpt-5".to_string()]),
};
let candidates =
build_ranked_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
super::enumerate_minimal_candidate_selection(EnumerateMinimalCandidateSelectionInput {
rows: vec![sample_row("1"), disallowed],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
@@ -195,8 +192,6 @@ mod tests {
require_streaming: false,
required_capabilities: None,
auth_constraints: Some(&constraints),
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate selection should build");
@@ -213,7 +208,7 @@ mod tests {
earlier_priority.provider_priority = 0;
let candidates =
super::enumerate_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
super::enumerate_minimal_candidate_selection(EnumerateMinimalCandidateSelectionInput {
rows: vec![later_priority, earlier_priority],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
@@ -221,8 +216,6 @@ mod tests {
require_streaming: false,
required_capabilities: None,
auth_constraints: None,
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate enumeration should build");
@@ -255,8 +248,7 @@ mod tests {
}
#[test]
fn ranked_minimal_candidate_selection_prefers_matching_requested_capabilities_before_priority()
{
fn requested_capability_priority_counts_missing_compatible_capabilities() {
let mut missing_capability = sample_row("1");
missing_capability.key_capabilities = Some(serde_json::json!({"cache_1h": false}));
missing_capability.provider_priority = 0;
@@ -267,7 +259,7 @@ mod tests {
let required_capabilities = serde_json::json!({"cache_1h": true});
let candidates =
build_ranked_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
super::enumerate_minimal_candidate_selection(EnumerateMinimalCandidateSelectionInput {
rows: vec![missing_capability, matching_capability],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
@@ -275,14 +267,20 @@ mod tests {
require_streaming: false,
required_capabilities: Some(&required_capabilities),
auth_constraints: None,
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate selection should build");
let priority = candidates
.iter()
.map(|candidate| {
super::requested_capability_priority_for_candidate(
Some(&required_capabilities),
candidate,
)
})
.collect::<Vec<_>>();
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].key_id, "key-2");
assert_eq!(candidates[1].key_id, "key-1");
assert_eq!(priority, vec![(0, 1), (0, 0)]);
}
#[test]

View File

@@ -1,45 +0,0 @@
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,
};
pub fn build_ranked_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)
}

View File

@@ -28,7 +28,7 @@ pub struct SchedulerMinimalCandidateSelectionCandidate {
pub mapping_matched_model: Option<String>,
}
pub struct BuildMinimalCandidateSelectionInput<'a> {
pub struct EnumerateMinimalCandidateSelectionInput<'a> {
pub rows: Vec<StoredMinimalCandidateSelectionRow>,
pub normalized_api_format: &'a str,
pub requested_model_name: &'a str,
@@ -36,6 +36,4 @@ pub struct BuildMinimalCandidateSelectionInput<'a> {
pub require_streaming: bool,
pub required_capabilities: Option<&'a serde_json::Value>,
pub auth_constraints: Option<&'a crate::SchedulerAuthConstraints>,
pub affinity_key: Option<&'a str>,
pub priority_mode: SchedulerPriorityMode,
}

View File

@@ -16,12 +16,12 @@ pub use auth::{
auth_constraints_allow_provider, provider_matches_allowed_value, SchedulerAuthConstraints,
};
pub use candidate::{
auth_api_key_concurrency_limit_reached, build_ranked_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,
enumerate_minimal_candidate_selection, requested_capability_priority_for_candidate,
BuildMinimalCandidateSelectionInput, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
candidate_runtime_skip_reason_with_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability, enumerate_minimal_candidate_selection,
requested_capability_priority_for_candidate, CandidateRuntimeSelectabilityInput,
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,

View File

@@ -209,6 +209,26 @@ mod tests {
);
}
#[test]
fn capability_priority_precedes_provider_priority() {
let matching_capability = candidate("matching", 10, 0, None);
let mut missing_compatible_capability = candidate("missing", 0, 0, None);
missing_compatible_capability.capability_priority = (0, 1);
assert_eq!(
ranked_ids(
&[missing_compatible_capability, matching_capability],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-matching", "provider-missing"]
);
}
#[test]
fn cache_affinity_can_promote_cached_candidate_and_reports_reason() {
let high_priority = candidate("high", 0, 0, Some(0));