i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Unify candidate ranking pipeline

Browse Source
This commit is contained in:
fawney19
2026-04-27 12:34:03 +08:00
parent 9b866a6d17
commit 3b542434a2
46 changed files with 3635 additions and 1914 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1074
crates/aether-scheduler-core/src/candidate.rs
View File

File diff suppressed because it is too large Load Diff

124
crates/aether-scheduler-core/src/candidate/capability.rs Normal file
View File

@@ -0,0 +1,124 @@
use super::types::SchedulerMinimalCandidateSelectionCandidate;
#[derive(Debug, Clone, Copy)]
pub(crate) struct RequiredCapabilityDescriptor<'a> {
pub(crate) name: &'a str,
pub(crate) compatible: bool,
}
pub fn candidate_supports_required_capability(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
required_capability: &str,
) -> bool {
let required_capability = required_capability.trim();
if required_capability.is_empty() {
return true;
}
let Some(capabilities) = candidate.key_capabilities.as_ref() else {
return false;
};
if let Some(object) = capabilities.as_object() {
return object.iter().any(|(key, value)| {
key.eq_ignore_ascii_case(required_capability)
&& match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => {
value.as_i64().is_some_and(|value| value > 0)
}
_ => false,
}
});
}
if let Some(items) = capabilities.as_array() {
return items.iter().any(|value| {
value
.as_str()
.is_some_and(|value| value.eq_ignore_ascii_case(required_capability))
});
}
false
}
pub fn requested_capability_priority_for_candidate(
required_capabilities: Option<&serde_json::Value>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (u32, u32) {
let Some(required_capabilities) = required_capabilities.and_then(serde_json::Value::as_object)
else {
return (0, 0);
};
requested_capability_priority_for_candidate_descriptors(
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),
})
}),
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>(
required_capabilities: I,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (u32, u32)
where
I: IntoIterator<Item = RequiredCapabilityDescriptor<'a>>,
{
let mut exclusive_misses = 0u32;
let mut compatible_misses = 0u32;
for capability in required_capabilities {
if candidate_supports_required_capability(candidate, capability.name) {
continue;
}
if capability.compatible {
compatible_misses += 1;
} else {
exclusive_misses += 1;
}
}
(exclusive_misses, compatible_misses)
}
fn requested_capability_is_enabled(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => value.as_i64().is_some_and(|value| value > 0),
_ => false,
}
}
fn requested_capability_is_compatible(capability: &str) -> bool {
matches!(
capability.trim().to_ascii_lowercase().as_str(),
"cache_1h" | "context_1m"
)
}

163
crates/aether-scheduler-core/src/candidate/enumeration.rs Normal file
View File

@@ -0,0 +1,163 @@
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> {
let BuildMinimalCandidateSelectionInput {
rows,
normalized_api_format,
requested_model_name,
resolved_global_model_name,
require_streaming,
auth_constraints,
..
} = input;
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !crate::auth_constraints_allow_api_format(auth_constraints, normalized_api_format) {
return Ok(Vec::new());
}
if !crate::auth_constraints_allow_model(
auth_constraints,
requested_model_name,
resolved_global_model_name,
) {
return Ok(Vec::new());
}
let mut candidates = Vec::with_capacity(rows.len());
for row in rows {
if !crate::auth_constraints_allow_provider(
auth_constraints,
&row.provider_id,
&row.provider_name,
&row.provider_type,
) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
let Some((selected_provider_model_name, mapping_matched_model)) =
crate::resolve_provider_model_name(&row, requested_model_name, normalized_api_format)
else {
continue;
};
candidates.push(SchedulerMinimalCandidateSelectionCandidate {
provider_id: row.provider_id,
provider_name: row.provider_name,
provider_type: row.provider_type,
provider_priority: row.provider_priority,
endpoint_id: row.endpoint_id,
endpoint_api_format: row.endpoint_api_format,
key_id: row.key_id,
key_name: row.key_name,
key_auth_type: row.key_auth_type,
key_internal_priority: row.key_internal_priority,
key_global_priority_for_format: crate::extract_global_priority_for_format(
row.key_global_priority_by_format.as_ref(),
normalized_api_format,
)?,
key_capabilities: row.key_capabilities,
model_id: row.model_id,
global_model_id: row.global_model_id,
global_model_name: row.global_model_name,
selected_provider_model_name,
mapping_matched_model,
});
}
Ok(candidates)
}
pub fn collect_global_model_names_for_required_capability(
rows: Vec<StoredMinimalCandidateSelectionRow>,
normalized_api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_constraints: Option<&crate::SchedulerAuthConstraints>,
) -> Vec<String> {
if normalized_api_format.is_empty() || required_capability.trim().is_empty() {
return Vec::new();
}
if !crate::auth_constraints_allow_api_format(auth_constraints, normalized_api_format) {
return Vec::new();
}
let mut model_names = BTreeSet::new();
for row in rows {
if !crate::auth_constraints_allow_provider(
auth_constraints,
&row.provider_id,
&row.provider_name,
&row.provider_type,
) {
continue;
}
if !crate::row_supports_required_capability(&row, required_capability) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
if !crate::auth_constraints_allow_model(
auth_constraints,
&row.global_model_name,
&row.global_model_name,
) {
continue;
}
model_names.insert(row.global_model_name);
}
model_names.into_iter().collect()
}

39
crates/aether-scheduler-core/src/candidate/identity.rs Normal file
View File

@@ -0,0 +1,39 @@
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),
)
}

482
crates/aether-scheduler-core/src/candidate/mod.rs Normal file
View File

@@ -0,0 +1,482 @@
pub mod capability;
pub mod enumeration;
pub mod identity;
pub mod selectability;
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,
};
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,
};
pub use types::{
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
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,
};
use crate::SchedulerAuthConstraints;
fn sample_row(id: &str) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{id}"),
provider_name: format!("Provider {id}"),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{id}"),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: format!("key-{id}"),
key_name: format!("prod-{id}"),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: Some(serde_json::json!({"cache_1h": true})),
key_internal_priority: 50,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 2})),
model_id: format!("model-{id}"),
global_model_id: format!("global-model-{id}"),
global_model_name: "gpt-5".to_string(),
global_model_mappings: Some(vec!["gpt-5(?:\\.\\d+)?".to_string()]),
global_model_supports_streaming: Some(true),
model_provider_model_name: format!("gpt-5-upstream-{id}"),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: format!("gpt-5-canary-{id}"),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]),
model_supports_streaming: None,
model_is_active: true,
model_is_available: true,
}
}
fn sample_candidate(
id: &str,
capabilities: Option<serde_json::Value>,
) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
provider_id: format!("provider-{id}"),
provider_name: format!("Provider {id}"),
provider_type: "openai".to_string(),
provider_priority: 0,
endpoint_id: format!("endpoint-{id}"),
endpoint_api_format: "openai:chat".to_string(),
key_id: format!("key-{id}"),
key_name: format!("key-{id}"),
key_auth_type: "bearer".to_string(),
key_internal_priority: 0,
key_global_priority_for_format: None,
key_capabilities: capabilities,
model_id: format!("model-{id}"),
global_model_id: format!("global-model-{id}"),
global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
mapping_matched_model: None,
}
}
fn sample_key(id: &str, health_score: f64) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
format!("key-{id}"),
format!("provider-{id}"),
format!("key-{id}"),
"api_key".to_string(),
None,
true,
)
.expect("provider key should build");
key.health_by_format = Some(serde_json::json!({
"openai:chat": {
"health_score": health_score
}
}));
key
}
fn stored_candidate(
id: &str,
status: RequestCandidateStatus,
created_at_unix_ms: i64,
) -> StoredRequestCandidate {
let finished_at_unix_ms = match status {
RequestCandidateStatus::Pending | RequestCandidateStatus::Streaming => None,
_ => Some(created_at_unix_ms),
};
StoredRequestCandidate::new(
id.to_string(),
format!("req-{id}"),
None,
None,
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
status,
None,
false,
None,
None,
None,
None,
None,
None,
None,
created_at_unix_ms,
Some(created_at_unix_ms),
finished_at_unix_ms,
)
.expect("candidate should build")
}
#[test]
fn reads_required_capability_from_object_and_array_forms() {
assert!(candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!({"vision": true}))),
"vision"
));
assert!(candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!(["vision", "tools"]))),
"tools"
));
assert!(!candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!({"vision": false}))),
"vision"
));
}
#[test]
fn builds_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();
let constraints = SchedulerAuthConstraints {
allowed_providers: Some(vec!["provider-1".to_string()]),
allowed_api_formats: Some(vec!["OPENAI:CHAT".to_string()]),
allowed_models: Some(vec!["gpt-5".to_string()]),
};
let candidates = build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows: vec![sample_row("1"), disallowed],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
resolved_global_model_name: "gpt-5",
require_streaming: false,
required_capabilities: None,
auth_constraints: Some(&constraints),
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate selection should build");
assert_eq!(candidates.len(), 1);
assert_eq!(candidates[0].provider_id, "provider-1");
assert_eq!(candidates[0].selected_provider_model_name, "gpt-5-canary-1");
}
#[test]
fn enumeration_preserves_theoretical_candidate_order_without_final_sorting() {
let mut later_priority = sample_row("1");
later_priority.provider_priority = 10;
let mut earlier_priority = sample_row("2");
earlier_priority.provider_priority = 0;
let candidates =
super::enumerate_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows: vec![later_priority, earlier_priority],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
resolved_global_model_name: "gpt-5",
require_streaming: false,
required_capabilities: None,
auth_constraints: None,
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate enumeration should build");
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].provider_id, "provider-1");
assert_eq!(candidates[1].provider_id, "provider-2");
}
#[test]
fn collects_global_model_names_for_required_capability_with_auth_constraints() {
let mut disallowed = sample_row("2");
disallowed.global_model_name = "gpt-4.1".to_string();
disallowed.provider_id = "provider-blocked".to_string();
disallowed.provider_name = "Blocked".to_string();
let constraints = SchedulerAuthConstraints {
allowed_providers: Some(vec!["provider-1".to_string()]),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_models: Some(vec!["gpt-5".to_string()]),
};
let model_names = collect_global_model_names_for_required_capability(
vec![sample_row("1"), disallowed],
"openai:chat",
"cache_1h",
false,
Some(&constraints),
);
assert_eq!(model_names, vec!["gpt-5".to_string()]);
}
#[test]
fn minimal_candidate_selection_prefers_matching_requested_capabilities_before_priority() {
let mut missing_capability = sample_row("1");
missing_capability.key_capabilities = Some(serde_json::json!({"cache_1h": false}));
missing_capability.provider_priority = 0;
let mut matching_capability = sample_row("2");
matching_capability.key_capabilities = Some(serde_json::json!({"cache_1h": true}));
matching_capability.provider_priority = 10;
let required_capabilities = serde_json::json!({"cache_1h": true});
let candidates = build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows: vec![missing_capability, matching_capability],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
resolved_global_model_name: "gpt-5",
require_streaming: false,
required_capabilities: Some(&required_capabilities),
auth_constraints: None,
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate selection should build");
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].key_id, "key-2");
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)];
let provider_concurrent_limits = BTreeMap::from([("provider-1".to_string(), 1)]);
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &provider_concurrent_limits,
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
},
));
}
#[test]
fn candidate_selectability_rejects_quota_or_zero_health() {
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), sample_key("1", 0.0))]);
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
},
));
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: true,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
},
));
}
#[test]
fn candidate_selectability_rejects_exhausted_account_quota() {
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: true,
oauth_invalid: false,
rpm_reset_at: None,
},
));
}
#[test]
fn candidate_selectability_rejects_oauth_invalid_keys() {
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: true,
rpm_reset_at: None,
},
));
}
#[test]
fn detects_auth_api_key_concurrency_limit_from_recent_active_requests() {
let recent_candidates = vec![StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95,
Some(95),
None,
)
.expect("candidate should build")];
assert!(auth_api_key_concurrency_limit_reached(
&recent_candidates,
100,
"api-key-1",
1,
));
assert!(!auth_api_key_concurrency_limit_reached(
&recent_candidates,
100,
"api-key-1",
2,
));
}
}

204
crates/aether-scheduler-core/src/candidate/selectability.rs Normal file
View File

@@ -0,0 +1,204 @@
use std::collections::{BTreeMap, BTreeSet};
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};
pub fn auth_api_key_concurrency_limit_reached(
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
api_key_id: &str,
concurrent_limit: usize,
) -> bool {
if api_key_id.trim().is_empty() || concurrent_limit == 0 {
return false;
}
crate::count_recent_active_requests_for_api_key(recent_candidates, api_key_id, now_unix_secs)
>= 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,
pub recent_candidates: &'a [StoredRequestCandidate],
pub provider_concurrent_limits: &'a BTreeMap<String, usize>,
pub provider_key_rpm_states: &'a BTreeMap<String, StoredProviderCatalogKey>,
pub now_unix_secs: u64,
pub cached_affinity_target: Option<&'a crate::SchedulerAffinityTarget>,
pub provider_quota_blocks_requests: bool,
pub account_quota_exhausted: bool,
pub oauth_invalid: bool,
pub rpm_reset_at: Option<u64>,
}
pub fn candidate_is_selectable_with_runtime_state(
input: CandidateRuntimeSelectabilityInput<'_>,
) -> bool {
candidate_runtime_skip_reason_with_state(input).is_none()
}
pub fn candidate_runtime_skip_reason_with_state(
input: CandidateRuntimeSelectabilityInput<'_>,
) -> Option<&'static str> {
let CandidateRuntimeSelectabilityInput {
candidate,
recent_candidates,
provider_concurrent_limits,
provider_key_rpm_states,
now_unix_secs,
cached_affinity_target,
provider_quota_blocks_requests,
account_quota_exhausted,
oauth_invalid,
rpm_reset_at,
} = input;
if provider_quota_blocks_requests {
return Some("provider_quota_blocked");
}
if account_quota_exhausted {
return Some("account_quota_exhausted");
}
if oauth_invalid {
return Some("oauth_invalid");
}
if crate::is_candidate_in_recent_failure_cooldown(
recent_candidates,
candidate.provider_id.as_str(),
candidate.endpoint_id.as_str(),
candidate.key_id.as_str(),
now_unix_secs,
) {
return Some("recent_failure_cooldown");
}
if provider_concurrent_limits
.get(&candidate.provider_id)
.is_some_and(|limit| {
crate::count_recent_active_requests_for_provider(
recent_candidates,
candidate.provider_id.as_str(),
now_unix_secs,
) >= *limit
})
{
return Some("provider_concurrency_limit_reached");
}
let is_cached_user = cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(candidate, target));
if let Some(provider_key) = provider_key_rpm_states.get(&candidate.key_id) {
if crate::is_provider_key_circuit_open(provider_key, candidate.endpoint_api_format.as_str())
{
return Some("key_circuit_open");
}
if crate::provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())
.is_some_and(|score| score <= 0.0)
{
return Some("key_health_score_zero");
}
if !crate::provider_key_rpm_allows_request_since(
provider_key,
recent_candidates,
now_unix_secs,
is_cached_user,
rpm_reset_at,
) {
return Some("key_rpm_exhausted");
}
}
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)
}

41
crates/aether-scheduler-core/src/candidate/types.rs Normal file
View File

@@ -0,0 +1,41 @@
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub enum SchedulerPriorityMode {
#[default]
Provider,
GlobalKey,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct SchedulerMinimalCandidateSelectionCandidate {
pub provider_id: String,
pub provider_name: String,
pub provider_type: String,
pub provider_priority: i32,
pub endpoint_id: String,
pub endpoint_api_format: String,
pub key_id: String,
pub key_name: String,
pub key_auth_type: String,
pub key_internal_priority: i32,
pub key_global_priority_for_format: Option<i32>,
pub key_capabilities: Option<serde_json::Value>,
pub model_id: String,
pub global_model_id: String,
pub global_model_name: String,
pub selected_provider_model_name: String,
pub mapping_matched_model: Option<String>,
}
pub struct BuildMinimalCandidateSelectionInput<'a> {
pub rows: Vec<StoredMinimalCandidateSelectionRow>,
pub normalized_api_format: &'a str,
pub requested_model_name: &'a str,
pub resolved_global_model_name: &'a str,
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,
}

18
crates/aether-scheduler-core/src/lib.rs
View File

@@ -4,6 +4,7 @@ mod candidate;
mod health;
mod model;
mod provider;
mod ranking;
mod request_candidate;
pub use affinity::{
@@ -19,9 +20,10 @@ pub use candidate::{
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,
reorder_candidates_by_scheduler_health, requested_capability_priority_for_candidate,
BuildMinimalCandidateSelectionInput, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
enumerate_minimal_candidate_selection, reorder_candidates_by_scheduler_health,
requested_capability_priority_for_candidate, BuildMinimalCandidateSelectionInput,
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
@@ -35,9 +37,17 @@ pub use health::{
pub use model::{
candidate_model_names, extract_global_priority_for_format, matches_model_mapping,
normalize_api_format, resolve_provider_model_name, resolve_requested_global_model_name,
row_supports_required_capability, select_provider_model_name,
row_supports_requested_model, row_supports_required_capability, select_provider_model_name,
};
pub use provider::{build_provider_concurrent_limit_map, should_skip_provider_quota};
pub use ranking::{
apply_scheduler_candidate_ranking, candidate_priority_slot, candidates_share_priority_group,
compare_candidate_identity_for_ranking, compare_candidate_priority_slot,
scheduler_candidate_ranking_order, scheduler_ranking_outcomes, SchedulerRankableCandidate,
SchedulerRankingContext, SchedulerRankingMode, SchedulerRankingOutcome,
SchedulerTunnelAffinityBucket, RANKING_REASON_CACHED_AFFINITY, RANKING_REASON_CROSS_FORMAT,
RANKING_REASON_LOCAL_TUNNEL,
};
pub use request_candidate::{
build_execution_request_candidate_seed, build_local_request_candidate_status_record,
build_report_request_candidate_status_record, execution_error_details,

70
crates/aether-scheduler-core/src/model.rs
View File

@@ -11,30 +11,56 @@ pub fn resolve_requested_global_model_name(
requested_model_name: &str,
api_format: &str,
) -> Option<String> {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_name == requested_model_name
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
resolve_global_model_name_by(rows, |row| row.global_model_name == requested_model_name)
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_name == requested_model_name
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
})
})
}
pub fn row_supports_requested_model(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
) -> bool {
row.global_model_name == requested_model_name
|| row.model_provider_model_name == requested_model_name
|| row
.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
|| row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
}
fn resolve_global_model_name_by<F>(

12
crates/aether-scheduler-core/src/ranking/format.rs Normal file
View File

@@ -0,0 +1,12 @@
use std::cmp::Ordering;
use super::types::SchedulerRankableCandidate;
pub fn compare_format_state(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
) -> Ordering {
left.demote_cross_format
.cmp(&right.demote_cross_format)
.then(left.format_preference.cmp(&right.format_preference))
}

294
crates/aether-scheduler-core/src/ranking/mod.rs Normal file
View File

@@ -0,0 +1,294 @@
mod format;
mod modes;
mod priority;
mod reasons;
mod types;
pub use format::compare_format_state;
pub use modes::{apply_load_balance_rotation, compare_rankable_candidates};
pub use priority::{
candidate_priority_slot, candidates_share_priority_group, compare_candidate_priority_slot,
};
pub use reasons::{
demoted_by as ranking_demoted_by, promoted_by as ranking_promoted_by,
RANKING_REASON_CACHED_AFFINITY, RANKING_REASON_CROSS_FORMAT, RANKING_REASON_LOCAL_TUNNEL,
};
pub use types::{
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
SchedulerRankingOutcome, SchedulerTunnelAffinityBucket,
};
pub fn compare_candidate_identity_for_ranking(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
) -> 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),
)
}
pub fn scheduler_candidate_ranking_order(
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<usize> {
let mut order = (0..candidates.len()).collect::<Vec<_>>();
order.sort_by(|left, right| {
compare_rankable_candidates(&candidates[*left], &candidates[*right], context)
});
apply_load_balance_rotation(&mut order, candidates, context);
order
}
pub fn scheduler_ranking_outcomes(
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<SchedulerRankingOutcome> {
scheduler_candidate_ranking_order(candidates, context)
.into_iter()
.enumerate()
.map(|(ranking_index, original_index)| {
let candidate = &candidates[original_index];
SchedulerRankingOutcome {
original_index,
ranking_index,
priority_mode: context.priority_mode,
ranking_mode: context.ranking_mode,
priority_slot: candidate_priority_slot(candidate, context.priority_mode),
promoted_by: ranking_promoted_by(candidate, context.ranking_mode),
demoted_by: ranking_demoted_by(candidate),
}
})
.collect()
}
pub fn apply_scheduler_candidate_ranking<T>(
items: &mut [T],
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<SchedulerRankingOutcome> {
let outcomes = scheduler_ranking_outcomes(candidates, context);
apply_order(
items,
outcomes
.iter()
.map(|outcome| outcome.original_index)
.collect(),
);
outcomes
}
fn apply_order<T>(items: &mut [T], sorted_old_indices: Vec<usize>) {
if items.len() < 2 {
return;
}
let mut target_positions = vec![0usize; sorted_old_indices.len()];
for (new_position, old_position) in sorted_old_indices.into_iter().enumerate() {
target_positions[old_position] = new_position;
}
for index in 0..items.len() {
while target_positions[index] != index {
let target = target_positions[index];
items.swap(index, target);
target_positions.swap(index, target);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{SchedulerPriorityMode, SchedulerTunnelAffinityBucket};
fn candidate(
id: &str,
provider_priority: i32,
key_priority: i32,
global_key_priority: Option<i32>,
) -> SchedulerRankableCandidate {
SchedulerRankableCandidate {
provider_id: format!("provider-{id}"),
endpoint_id: format!("endpoint-{id}"),
key_id: format!("key-{id}"),
selected_provider_model_name: "gpt-5".to_string(),
provider_priority,
key_internal_priority: key_priority,
key_global_priority_for_format: global_key_priority,
capability_priority: (0, 0),
cached_affinity_match: false,
affinity_hash: None,
tunnel_bucket: SchedulerTunnelAffinityBucket::Neutral,
demote_cross_format: false,
format_preference: (0, 0),
health_bucket: None,
health_score: 1.0,
original_index: 0,
}
}
fn ranked_ids(
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<String> {
scheduler_candidate_ranking_order(candidates, context)
.into_iter()
.map(|index| candidates[index].provider_id.clone())
.collect()
}
#[test]
fn provider_priority_mode_prefers_provider_priority_slot() {
let candidates = vec![
candidate("global", 10, 0, Some(0)),
candidate("provider", 0, 10, Some(10)),
];
assert_eq!(
ranked_ids(
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-provider", "provider-global"]
);
}
#[test]
fn global_key_priority_mode_prefers_global_key_priority_slot() {
let candidates = vec![
candidate("provider", 0, 10, Some(10)),
candidate("global", 10, 0, Some(0)),
];
assert_eq!(
ranked_ids(
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::GlobalKey,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-global", "provider-provider"]
);
}
#[test]
fn fixed_order_keeps_priority_before_affinity_tunnel_and_format_preference() {
let mut lower_priority = candidate("lower", 10, 0, Some(10));
lower_priority.cached_affinity_match = true;
lower_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::LocalTunnel;
lower_priority.format_preference = (0, 0);
let mut higher_priority = candidate("higher", 0, 0, Some(0));
higher_priority.demote_cross_format = true;
higher_priority.format_preference = (9, 9);
assert_eq!(
ranked_ids(
&[lower_priority, higher_priority],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-higher", "provider-lower"]
);
}
#[test]
fn cache_affinity_can_promote_cached_candidate_and_reports_reason() {
let high_priority = candidate("high", 0, 0, Some(0));
let mut cached = candidate("cached", 10, 0, Some(10));
cached.cached_affinity_match = true;
let candidates = vec![high_priority, cached];
let context = SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
};
let outcomes = scheduler_ranking_outcomes(&candidates, context);
assert_eq!(outcomes[0].original_index, 1);
assert_eq!(
outcomes[0].promoted_by,
Some(RANKING_REASON_CACHED_AFFINITY)
);
}
#[test]
fn cache_affinity_without_cache_hit_keeps_priority_before_tunnel() {
let mut higher_priority = candidate("higher", 0, 0, Some(0));
higher_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::RemoteTunnel;
let mut lower_priority = candidate("lower", 10, 0, Some(10));
lower_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::LocalTunnel;
assert_eq!(
ranked_ids(
&[lower_priority, higher_priority],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-higher", "provider-lower"]
);
}
#[test]
fn cache_affinity_keeps_cross_format_demotion_before_priority() {
let same_format_low_priority = candidate("same", 10, 0, Some(10));
let mut cross_format_high_priority = candidate("cross", 0, 0, Some(0));
cross_format_high_priority.demote_cross_format = true;
assert_eq!(
ranked_ids(
&[cross_format_high_priority, same_format_low_priority],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-same", "provider-cross"]
);
}
#[test]
fn load_balance_rotates_only_within_same_priority_group() {
let first = candidate("first", 0, 0, Some(0));
let second = candidate("second", 0, 0, Some(0));
let third = candidate("third", 10, 0, Some(10));
assert_eq!(
ranked_ids(
&[first, second, third],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::LoadBalance,
include_health: false,
load_balance_seed: 1,
},
),
vec!["provider-second", "provider-first", "provider-third"]
);
}
}

111
crates/aether-scheduler-core/src/ranking/modes.rs Normal file
View File

@@ -0,0 +1,111 @@
use std::cmp::Ordering;
use super::types::{SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode};
use super::{
candidates_share_priority_group, compare_candidate_identity_for_ranking,
compare_candidate_priority_slot, compare_format_state,
};
pub fn compare_rankable_candidates(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
match context.ranking_mode {
SchedulerRankingMode::FixedOrder => compare_fixed_order(left, right, context),
SchedulerRankingMode::CacheAffinity => compare_cache_affinity(left, right, context),
SchedulerRankingMode::LoadBalance => compare_load_balance_base(left, right, context),
}
}
fn compare_fixed_order(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then_with(|| compare_format_state(left, right))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
fn compare_cache_affinity(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then_with(|| right.cached_affinity_match.cmp(&left.cached_affinity_match))
.then(left.demote_cross_format.cmp(&right.demote_cross_format))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then(left.tunnel_bucket.cmp(&right.tunnel_bucket))
.then(left.format_preference.cmp(&right.format_preference))
.then_with(|| compare_health(left, right, context.include_health))
.then(left.affinity_hash.cmp(&right.affinity_hash))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
fn compare_load_balance_base(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then(left.demote_cross_format.cmp(&right.demote_cross_format))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then(left.format_preference.cmp(&right.format_preference))
.then_with(|| compare_health(left, right, context.include_health))
.then(left.affinity_hash.cmp(&right.affinity_hash))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
fn compare_health(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
include_health: bool,
) -> Ordering {
if !include_health {
return Ordering::Equal;
}
right
.health_bucket
.cmp(&left.health_bucket)
.then_with(|| right.health_score.total_cmp(&left.health_score))
}
pub fn apply_load_balance_rotation(
sorted_indices: &mut [usize],
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) {
if context.ranking_mode != SchedulerRankingMode::LoadBalance || sorted_indices.len() < 2 {
return;
}
let mut start = 0usize;
while start < sorted_indices.len() {
let mut end = start + 1;
while end < sorted_indices.len()
&& candidates_share_priority_group(
&candidates[sorted_indices[start]],
&candidates[sorted_indices[end]],
context.priority_mode,
)
{
end += 1;
}
let group_len = end - start;
if group_len > 1 {
let offset = usize::try_from(context.load_balance_seed).unwrap_or(0) % group_len;
sorted_indices[start..end].rotate_left(offset);
}
start = end;
}
}

52
crates/aether-scheduler-core/src/ranking/priority.rs Normal file
View File

@@ -0,0 +1,52 @@
use std::cmp::Ordering;
use crate::SchedulerPriorityMode;
use super::types::SchedulerRankableCandidate;
pub fn candidate_priority_slot(
candidate: &SchedulerRankableCandidate,
priority_mode: SchedulerPriorityMode,
) -> i32 {
match priority_mode {
SchedulerPriorityMode::Provider => candidate.provider_priority,
SchedulerPriorityMode::GlobalKey => {
candidate.key_global_priority_for_format.unwrap_or(i32::MAX)
}
}
}
pub fn compare_candidate_priority_slot(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
priority_mode: SchedulerPriorityMode,
) -> Ordering {
match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority)),
}
}
pub fn candidates_share_priority_group(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
priority_mode: SchedulerPriorityMode,
) -> bool {
match priority_mode {
SchedulerPriorityMode::Provider => {
left.provider_priority == right.provider_priority
&& left.key_internal_priority == right.key_internal_priority
}
SchedulerPriorityMode::GlobalKey => {
left.key_global_priority_for_format == right.key_global_priority_for_format
}
}
}

28
crates/aether-scheduler-core/src/ranking/reasons.rs Normal file
View File

@@ -0,0 +1,28 @@
use super::types::{
SchedulerRankableCandidate, SchedulerRankingMode, SchedulerTunnelAffinityBucket,
};
pub const RANKING_REASON_CACHED_AFFINITY: &str = "cached_affinity";
pub const RANKING_REASON_LOCAL_TUNNEL: &str = "local_tunnel";
pub const RANKING_REASON_CROSS_FORMAT: &str = "cross_format";
pub fn promoted_by(
candidate: &SchedulerRankableCandidate,
ranking_mode: SchedulerRankingMode,
) -> Option<&'static str> {
if ranking_mode == SchedulerRankingMode::CacheAffinity && candidate.cached_affinity_match {
return Some(RANKING_REASON_CACHED_AFFINITY);
}
if ranking_mode == SchedulerRankingMode::CacheAffinity
&& candidate.tunnel_bucket == SchedulerTunnelAffinityBucket::LocalTunnel
{
return Some(RANKING_REASON_LOCAL_TUNNEL);
}
None
}
pub fn demoted_by(candidate: &SchedulerRankableCandidate) -> Option<&'static str> {
candidate
.demote_cross_format
.then_some(RANKING_REASON_CROSS_FORMAT)
}

131
crates/aether-scheduler-core/src/ranking/types.rs Normal file
View File

@@ -0,0 +1,131 @@
use crate::{
ProviderKeyHealthBucket, SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub enum SchedulerRankingMode {
FixedOrder,
#[default]
CacheAffinity,
LoadBalance,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
pub enum SchedulerTunnelAffinityBucket {
LocalTunnel = 0,
#[default]
Neutral = 1,
RemoteTunnel = 2,
}
#[derive(Debug, Clone, PartialEq)]
pub struct SchedulerRankableCandidate {
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
pub selected_provider_model_name: String,
pub provider_priority: i32,
pub key_internal_priority: i32,
pub key_global_priority_for_format: Option<i32>,
pub capability_priority: (u32, u32),
pub cached_affinity_match: bool,
pub affinity_hash: Option<u64>,
pub tunnel_bucket: SchedulerTunnelAffinityBucket,
pub demote_cross_format: bool,
pub format_preference: (u8, u8),
pub health_bucket: Option<ProviderKeyHealthBucket>,
pub health_score: f64,
pub original_index: usize,
}
impl SchedulerRankableCandidate {
pub fn from_candidate(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
original_index: usize,
) -> Self {
Self {
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
selected_provider_model_name: candidate.selected_provider_model_name.clone(),
provider_priority: candidate.provider_priority,
key_internal_priority: candidate.key_internal_priority,
key_global_priority_for_format: candidate.key_global_priority_for_format,
capability_priority: (0, 0),
cached_affinity_match: false,
affinity_hash: None,
tunnel_bucket: SchedulerTunnelAffinityBucket::Neutral,
demote_cross_format: false,
format_preference: (0, 0),
health_bucket: None,
health_score: 1.0,
original_index,
}
}
pub fn with_capability_priority(mut self, value: (u32, u32)) -> Self {
self.capability_priority = value;
self
}
pub fn with_cached_affinity_match(mut self, value: bool) -> Self {
self.cached_affinity_match = value;
self
}
pub fn with_affinity_hash(mut self, value: Option<u64>) -> Self {
self.affinity_hash = value;
self
}
pub fn with_tunnel_bucket(mut self, value: SchedulerTunnelAffinityBucket) -> Self {
self.tunnel_bucket = value;
self
}
pub fn with_format_state(
mut self,
demote_cross_format: bool,
format_preference: (u8, u8),
) -> Self {
self.demote_cross_format = demote_cross_format;
self.format_preference = format_preference;
self
}
pub fn with_health(mut self, bucket: Option<ProviderKeyHealthBucket>, score: f64) -> Self {
self.health_bucket = bucket;
self.health_score = score;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SchedulerRankingContext {
pub priority_mode: SchedulerPriorityMode,
pub ranking_mode: SchedulerRankingMode,
pub include_health: bool,
pub load_balance_seed: u64,
}
impl Default for SchedulerRankingContext {
fn default() -> Self {
Self {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct SchedulerRankingOutcome {
pub original_index: usize,
pub ranking_index: usize,
pub priority_mode: SchedulerPriorityMode,
pub ranking_mode: SchedulerRankingMode,
pub priority_slot: i32,
pub promoted_by: Option<&'static str>,
pub demoted_by: Option<&'static str>,
}