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

refactor: lazy pool key scheduling

Browse Source
This commit is contained in:
fawney19
2026-05-11 00:12:05 +08:00
parent 1a0f1a7b72
commit bacb14e5f0
26 changed files with 1856 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

364
crates/aether-ai-serving/src/pool_scheduler.rs
View File

@@ -206,6 +206,10 @@ fn schedule_pool_group<Candidate>(
.unwrap_or_default();
let active_presets =
normalize_enabled_pool_presets(&pool_config.scheduling_presets, provider_type.as_str());
let lru_distribution_enabled = pool_config.lru_enabled
&& !active_presets
.iter()
.any(|preset| pool_preset_mutex_group(&preset.preset).is_some());
let mut available = Vec::new();
let mut skipped = Vec::new();
@@ -285,7 +289,7 @@ fn schedule_pool_group<Candidate>(
let sort_vectors = build_pool_sort_vectors(
&available,
&active_presets,
pool_config.lru_enabled,
lru_distribution_enabled,
group_sort_seed(
provider_type.as_str(),
available.first().map(|item| &item.item.facts),
@@ -300,7 +304,7 @@ fn schedule_pool_group<Candidate>(
.cmp(&sort_vectors.get(&right.item.facts.key_id))
.then(left.original_index.cmp(&right.original_index))
});
} else if pool_config.lru_enabled {
} else if lru_distribution_enabled {
let lru_ranks = lru_rank_indices(&available, false);
available.sort_by(|left, right| {
lru_ranks
@@ -359,6 +363,16 @@ fn build_pool_sort_vectors<Candidate>(
let lru_ranks = lru_rank_indices(items, false);
let cache_affinity_ranks = lru_rank_indices(items, true);
if lru_enabled {
for item in items {
let key_id = item.item.facts.key_id.clone();
vectors
.entry(key_id.clone())
.or_default()
.push(*lru_ranks.get(&key_id).unwrap_or(&0));
}
}
for preset in presets {
let ranks = match preset.preset.as_str() {
"cache_affinity" => cache_affinity_ranks.clone(),
@@ -385,16 +399,6 @@ fn build_pool_sort_vectors<Candidate>(
}
}
if lru_enabled {
for item in items {
let key_id = item.item.facts.key_id.clone();
vectors
.entry(key_id.clone())
.or_default()
.push(*lru_ranks.get(&key_id).unwrap_or(&0));
}
}
vectors
}
@@ -408,13 +412,13 @@ fn lru_rank_indices<Candidate>(
fn priority_first_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
lru_ranks: &BTreeMap<String, usize>,
_lru_ranks: &BTreeMap<String, usize>,
) -> BTreeMap<String, usize> {
let scores = collect_metric_scores(items, |item| {
Some(f64::from(item.item.facts.key_internal_priority))
});
if !score_map_has_variation(&scores) {
return lru_ranks.clone();
return neutral_rank_indices(items);
}
rank_indices_from_score_map(items, &scores, false)
}
@@ -422,27 +426,35 @@ fn priority_first_ranks<Candidate>(
fn single_account_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
) -> BTreeMap<String, usize> {
let n = items.len().saturating_sub(1).max(1) as f64;
let priority_scores = collect_metric_scores(items, |item| {
Some(f64::from(item.item.facts.key_internal_priority))
});
let priority_ranks = rank_indices_from_score_map(items, &priority_scores, false);
let lru_desc_ranks = lru_rank_indices(items, true);
let combined_scores = items
let mut decorated = items
.iter()
.map(|item| {
let key_id = item.item.facts.key_id.clone();
let priority_rank = *priority_ranks.get(&key_id).unwrap_or(&0) as f64 / n;
let lru_rank = *lru_desc_ranks.get(&key_id).unwrap_or(&0) as f64 / n;
(key_id, Some((priority_rank * 0.75) + (lru_rank * 0.25)))
(
item.item.facts.key_internal_priority,
*lru_desc_ranks.get(&key_id).unwrap_or(&0),
item.original_index,
key_id,
)
})
.collect::<BTreeMap<_, _>>();
rank_indices_from_score_map(items, &combined_scores, false)
.collect::<Vec<_>>();
decorated.sort_by(|left, right| {
left.0
.cmp(&right.0)
.then(left.1.cmp(&right.1))
.then(left.2.cmp(&right.2))
});
decorated
.into_iter()
.enumerate()
.map(|(rank, (_, _, _, key_id))| (key_id, rank))
.collect()
}
fn plan_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
lru_ranks: &BTreeMap<String, usize>,
_lru_ranks: &BTreeMap<String, usize>,
mode: Option<&str>,
) -> BTreeMap<String, usize> {
let scores = items
@@ -458,14 +470,14 @@ fn plan_ranks<Candidate>(
})
.collect::<BTreeMap<_, _>>();
if !score_map_has_variation(&scores) {
return lru_ranks.clone();
return neutral_rank_indices(items);
}
rank_indices_from_score_map(items, &scores, false)
}
fn health_first_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
lru_ranks: &BTreeMap<String, usize>,
_lru_ranks: &BTreeMap<String, usize>,
) -> BTreeMap<String, usize> {
let scores = collect_metric_scores(items, |item| {
item.item
@@ -474,39 +486,39 @@ fn health_first_ranks<Candidate>(
.map(|score| 1.0 - score.clamp(0.0, 1.0))
});
if !score_map_has_signal(&scores) {
return lru_ranks.clone();
return neutral_rank_indices(items);
}
rank_indices_from_score_map(items, &scores, false)
}
fn latency_first_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
lru_ranks: &BTreeMap<String, usize>,
_lru_ranks: &BTreeMap<String, usize>,
) -> BTreeMap<String, usize> {
let scores = collect_metric_scores(items, |item| item.item.key_context.latency_avg_ms);
if !score_map_has_signal(&scores) {
return lru_ranks.clone();
return neutral_rank_indices(items);
}
rank_indices_from_score_map(items, &scores, false)
}
fn cost_first_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
lru_ranks: &BTreeMap<String, usize>,
_lru_ranks: &BTreeMap<String, usize>,
cost_limit_per_key_tokens: Option<u64>,
) -> BTreeMap<String, usize> {
let scores = collect_metric_scores(items, |item| {
cost_penalty(item, cost_limit_per_key_tokens).or(item.item.key_context.quota_usage_ratio)
});
if !score_map_has_signal(&scores) {
return lru_ranks.clone();
return neutral_rank_indices(items);
}
rank_indices_from_score_map(items, &scores, false)
}
fn quota_balanced_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
lru_ranks: &BTreeMap<String, usize>,
_lru_ranks: &BTreeMap<String, usize>,
cost_limit_per_key_tokens: Option<u64>,
) -> BTreeMap<String, usize> {
let scores = collect_metric_scores(items, |item| {
@@ -516,18 +528,18 @@ fn quota_balanced_ranks<Candidate>(
.or_else(|| cost_penalty(item, cost_limit_per_key_tokens))
});
if !score_map_has_signal(&scores) {
return lru_ranks.clone();
return neutral_rank_indices(items);
}
rank_indices_from_score_map(items, &scores, false)
}
fn recent_refresh_ranks<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
lru_ranks: &BTreeMap<String, usize>,
_lru_ranks: &BTreeMap<String, usize>,
) -> BTreeMap<String, usize> {
let scores = collect_metric_scores(items, |item| item.item.key_context.quota_reset_seconds);
if !score_map_has_signal(&scores) {
return lru_ranks.clone();
return neutral_rank_indices(items);
}
rank_indices_from_score_map(items, &scores, false)
}
@@ -646,6 +658,15 @@ fn rank_indices_from_score_map<Candidate>(
.collect()
}
fn neutral_rank_indices<Candidate>(
items: &[PoolGroupCandidateOrdering<Candidate>],
) -> BTreeMap<String, usize> {
items
.iter()
.map(|item| (item.item.facts.key_id.clone(), 0))
.collect()
}
fn cost_penalty<Candidate>(
item: &PoolGroupCandidateOrdering<Candidate>,
cost_limit_per_key_tokens: Option<u64>,
@@ -740,47 +761,42 @@ fn normalize_enabled_pool_presets(
entries.push((entries.len(), "recent_refresh".to_string(), true, None));
}
let mut group_anchor_index = BTreeMap::<String, usize>::new();
for (index, preset, _, _) in &entries {
let Some(mutex_group) = pool_preset_mutex_group(preset) else {
continue;
};
group_anchor_index
.entry(mutex_group.to_string())
.or_insert(*index);
}
let mut ordered_enabled = Vec::<(usize, usize, String, Option<String>)>::new();
let mut group_enabled = BTreeMap::<String, (usize, usize, String, Option<String>)>::new();
let mut distribution_mode = None::<(usize, String, Option<String>)>;
let mut strategy_presets = Vec::<(usize, String, Option<String>)>::new();
for (index, preset, enabled, mode) in entries {
if !enabled
|| preset == "lru"
|| !pool_preset_supported_for_provider(&preset, &provider_type)
{
if !enabled || !pool_preset_supported_for_provider(&preset, &provider_type) {
continue;
}
let Some(mutex_group) = pool_preset_mutex_group(&preset) else {
ordered_enabled.push((index, index, preset, mode));
strategy_presets.push((index, preset, mode));
continue;
};
let anchor = group_anchor_index
.get(mutex_group)
.copied()
.unwrap_or(index);
let existing = group_enabled.get(mutex_group);
if existing.is_none_or(|current| index < current.1) {
group_enabled.insert(mutex_group.to_string(), (anchor, index, preset, mode));
if mutex_group == "distribution_mode"
&& distribution_mode
.as_ref()
.is_none_or(|current| index < current.0)
{
distribution_mode = Some((index, preset, mode));
}
}
ordered_enabled.extend(group_enabled.into_values());
ordered_enabled.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
ordered_enabled
.into_iter()
.map(|(_, _, preset, mode)| NormalizedPoolPreset { preset, mode })
.collect()
let mut normalized = Vec::new();
if let Some((_, preset, mode)) = distribution_mode.filter(|(_, preset, _)| preset != "lru") {
normalized.push(NormalizedPoolPreset { preset, mode });
}
strategy_presets.sort_by_key(|left| left.0);
normalized.extend(
strategy_presets
.into_iter()
.map(|(_, preset, mode)| NormalizedPoolPreset { preset, mode }),
);
normalized
}
fn pool_preset_supported_for_provider(preset: &str, provider_type: &str) -> bool {
@@ -959,7 +975,193 @@ mod tests {
}
#[test]
fn normalizes_distribution_mutex_group_to_first_enabled_member() {
fn pool_scheduler_applies_distribution_mode_before_strategy_presets() {
let key_cache_hit =
sample_candidate("provider-pool", "endpoint-1", "key-cache-hit", 50, true)
.with_presets(vec![
AiPoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
},
]);
let key_high_priority =
sample_candidate("provider-pool", "endpoint-1", "key-high-priority", 10, true)
.with_presets(vec![
AiPoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
},
]);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
lru_score_by_key: BTreeMap::from([
("key-cache-hit".to_string(), 200.0),
("key-high-priority".to_string(), 10.0),
]),
..AiPoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(
vec![key_cache_hit, key_high_priority],
&runtime_by_provider,
"seed",
);
assert!(outcome.skipped_candidates.is_empty());
assert_eq!(
outcome
.candidates
.iter()
.map(|item| item.candidate.as_str())
.collect::<Vec<_>>(),
vec!["key-cache-hit", "key-high-priority"]
);
}
#[test]
fn load_balance_distribution_is_not_overridden_by_priority_strategy() {
let key_random_first =
sample_candidate("provider-pool", "endpoint-1", "key-random-first", 50, true)
.with_presets(vec![
AiPoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
},
]);
let key_high_priority =
sample_candidate("provider-pool", "endpoint-1", "key-high-priority", 10, true)
.with_presets(vec![
AiPoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
},
]);
let nonce = (0..1000)
.map(|index| format!("seed-{index}"))
.find(|nonce| {
let group_seed = format!("codex:provider-pool:endpoint-1:model-1:gpt-5:{nonce}");
stable_hash_score(format!("{group_seed}:key-random-first").as_str())
< stable_hash_score(format!("{group_seed}:key-high-priority").as_str())
})
.expect("test seed should exist");
let outcome = run_ai_pool_scheduler(
vec![key_random_first, key_high_priority],
&BTreeMap::new(),
nonce.as_str(),
);
assert!(outcome.skipped_candidates.is_empty());
assert_eq!(
outcome
.candidates
.iter()
.map(|item| item.candidate.as_str())
.collect::<Vec<_>>(),
vec!["key-random-first", "key-high-priority"]
);
}
#[test]
fn single_account_distribution_orders_by_priority_then_reverse_lru() {
let key_priority_old =
sample_candidate("provider-pool", "endpoint-1", "key-priority-old", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
preset: "single_account".to_string(),
enabled: true,
mode: None,
}]);
let key_priority_recent = sample_candidate(
"provider-pool",
"endpoint-1",
"key-priority-recent",
10,
true,
)
.with_presets(vec![AiPoolSchedulingPreset {
preset: "single_account".to_string(),
enabled: true,
mode: None,
}]);
let key_lower_priority_recent = sample_candidate(
"provider-pool",
"endpoint-1",
"key-lower-priority-recent",
50,
true,
)
.with_presets(vec![AiPoolSchedulingPreset {
preset: "single_account".to_string(),
enabled: true,
mode: None,
}]);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
lru_score_by_key: BTreeMap::from([
("key-priority-old".to_string(), 10.0),
("key-priority-recent".to_string(), 200.0),
("key-lower-priority-recent".to_string(), 500.0),
]),
..AiPoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(
vec![
key_priority_old,
key_lower_priority_recent,
key_priority_recent,
],
&runtime_by_provider,
"seed",
);
assert!(outcome.skipped_candidates.is_empty());
assert_eq!(
outcome
.candidates
.iter()
.map(|item| item.candidate.as_str())
.collect::<Vec<_>>(),
vec![
"key-priority-recent",
"key-priority-old",
"key-lower-priority-recent"
]
);
}
#[test]
fn normalizes_distribution_mode_before_strategy_presets() {
let presets = normalize_enabled_ai_pool_presets(
&[
AiPoolSchedulingPreset {
@@ -989,6 +1191,32 @@ mod tests {
assert_eq!(presets, ["single_account", "priority_first"]);
}
#[test]
fn normalizes_lru_as_mutually_exclusive_distribution_mode() {
let presets = normalize_enabled_ai_pool_presets(
&[
AiPoolSchedulingPreset {
preset: "lru".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
},
],
"openai",
);
assert_eq!(presets, ["priority_first"]);
}
fn sample_candidate(
provider_id: &str,
endpoint_id: &str,

5
crates/aether-data-contracts/src/repository/candidate_selection/mod.rs
View File

@@ -2,6 +2,7 @@ mod types;
pub use types::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsQuery,
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};

14
crates/aether-data-contracts/src/repository/candidate_selection/types.rs
View File

@@ -42,6 +42,18 @@ pub struct StoredMinimalCandidateSelectionRow {
pub model_is_available: bool,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum StoredPoolKeyCandidateOrder {
#[default]
InternalPriority,
Lru,
CacheAffinity,
SingleAccount,
LoadBalance {
seed: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredPoolKeyCandidateRowsQuery {
pub api_format: String,
@@ -49,6 +61,8 @@ pub struct StoredPoolKeyCandidateRowsQuery {
pub endpoint_id: String,
pub model_id: String,
pub selected_provider_model_name: String,
#[serde(default)]
pub order: StoredPoolKeyCandidateOrder,
pub offset: u32,
pub limit: u32,
}

45
crates/aether-data/src/repository/candidate_selection/memory.rs
View File

@@ -4,7 +4,8 @@ use async_trait::async_trait;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
};
use crate::DataLayerError;
@@ -130,11 +131,7 @@ impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelecti
&& row.model_id == query.model_id
})
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
left.key_internal_priority
.cmp(&right.key_internal_priority)
.then(left.key_id.cmp(&right.key_id))
});
sort_pool_key_rows(&mut rows, &query.order);
Ok(rows
.into_iter()
.skip(query.offset as usize)
@@ -143,6 +140,38 @@ impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelecti
}
}
fn sort_pool_key_rows(
rows: &mut [StoredMinimalCandidateSelectionRow],
order: &StoredPoolKeyCandidateOrder,
) {
rows.sort_by(|left, right| match order {
StoredPoolKeyCandidateOrder::LoadBalance { seed } => {
stable_pool_key_hash(seed.as_str(), left.key_id.as_str())
.cmp(&stable_pool_key_hash(seed.as_str(), right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
}
_ => left
.key_internal_priority
.cmp(&right.key_internal_priority)
.then(left.key_id.cmp(&right.key_id)),
});
}
fn stable_pool_key_hash(seed: &str, key_id: &str) -> u64 {
let mut hash = 0xcbf29ce484222325u64;
for byte in seed
.as_bytes()
.iter()
.copied()
.chain(std::iter::once(b':'))
.chain(key_id.as_bytes().iter().copied())
{
hash ^= u64::from(byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}
fn normalize_api_format(value: &str) -> String {
aether_ai_formats::normalize_api_format_alias(value)
}
@@ -215,7 +244,8 @@ mod tests {
use super::InMemoryMinimalCandidateSelectionReadRepository;
use crate::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
};
fn sample_row(
@@ -402,6 +432,7 @@ mod tests {
endpoint_id: "endpoint-pool".to_string(),
model_id: "model-pool".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
order: StoredPoolKeyCandidateOrder::InternalPriority,
offset: 2,
limit: 2,
})

5
crates/aether-data/src/repository/candidate_selection/mod.rs
View File

@@ -6,8 +6,9 @@ mod sqlite;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsQuery,
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
pub use mysql::MysqlMinimalCandidateSelectionReadRepository;

80
crates/aether-data/src/repository/candidate_selection/mysql.rs
View File

@@ -5,7 +5,7 @@ use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::mysql::MysqlPool;
@@ -35,6 +35,7 @@ SELECT
pak.capabilities AS key_capabilities,
pak.internal_priority AS key_internal_priority,
pak.global_priority_by_format AS key_global_priority_by_format,
pak.last_used_at AS key_last_used_at_unix_secs,
m.id AS model_id,
m.global_model_id AS global_model_id,
gm.name AS global_model_name,
@@ -67,6 +68,7 @@ struct CandidateSelectionRow {
row: StoredMinimalCandidateSelectionRow,
provider_pool_enabled: bool,
key_auth_config: Option<String>,
key_last_used_at_unix_secs: Option<u64>,
}
impl MysqlMinimalCandidateSelectionReadRepository {
@@ -180,18 +182,18 @@ impl MinimalCandidateSelectionReadRepository for MysqlMinimalCandidateSelectionR
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.map(|item| item.row)
.filter(|row| {
row.provider_id == query.provider_id
&& row.endpoint_id == query.endpoint_id
&& row.model_id == query.model_id
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
})
.collect::<Vec<_>>();
let mut rows = sort_pool_key_rows(rows);
let mut rows = sort_pool_key_rows(rows, &query.order);
Ok(rows
.drain(..)
.skip(query.offset as usize)
.take(query.limit as usize)
.map(|item| item.row)
.collect())
}
}
@@ -246,16 +248,66 @@ fn sort_rows(
}
fn sort_pool_key_rows(
mut rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> {
rows.sort_by(|left, right| {
left.key_internal_priority
.cmp(&right.key_internal_priority)
.then(left.key_id.cmp(&right.key_id))
mut rows: Vec<CandidateSelectionRow>,
order: &StoredPoolKeyCandidateOrder,
) -> Vec<CandidateSelectionRow> {
rows.sort_by(|left, right| match order {
StoredPoolKeyCandidateOrder::InternalPriority => compare_pool_key_internal(left, right),
StoredPoolKeyCandidateOrder::Lru => left
.key_last_used_at_unix_secs
.cmp(&right.key_last_used_at_unix_secs)
.then_with(|| compare_pool_key_internal(left, right)),
StoredPoolKeyCandidateOrder::CacheAffinity => right
.key_last_used_at_unix_secs
.cmp(&left.key_last_used_at_unix_secs)
.then_with(|| compare_pool_key_internal(left, right)),
StoredPoolKeyCandidateOrder::SingleAccount => left
.row
.key_internal_priority
.cmp(&right.row.key_internal_priority)
.then_with(|| {
right
.key_last_used_at_unix_secs
.cmp(&left.key_last_used_at_unix_secs)
})
.then(left.row.key_id.cmp(&right.row.key_id)),
StoredPoolKeyCandidateOrder::LoadBalance { seed } => {
stable_pool_key_hash(seed.as_str(), left.row.key_id.as_str())
.cmp(&stable_pool_key_hash(
seed.as_str(),
right.row.key_id.as_str(),
))
.then(left.row.key_id.cmp(&right.row.key_id))
}
});
rows
}
fn compare_pool_key_internal(
left: &CandidateSelectionRow,
right: &CandidateSelectionRow,
) -> std::cmp::Ordering {
left.row
.key_internal_priority
.cmp(&right.row.key_internal_priority)
.then(left.row.key_id.cmp(&right.row.key_id))
}
fn stable_pool_key_hash(seed: &str, key_id: &str) -> u64 {
let mut hash = 0xcbf29ce484222325u64;
for byte in seed
.as_bytes()
.iter()
.copied()
.chain(std::iter::once(b':'))
.chain(key_id.as_bytes().iter().copied())
{
hash ^= u64::from(byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}
fn row_matches_requested_model(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
@@ -394,6 +446,10 @@ fn map_candidate_selection_row(row: &MySqlRow) -> Result<CandidateSelectionRow,
},
provider_pool_enabled,
key_auth_config: row.try_get("key_auth_config").map_sql_err()?,
key_last_used_at_unix_secs: row
.try_get::<Option<i64>, _>("key_last_used_at_unix_secs")
.map_sql_err()?
.and_then(|value| u64::try_from(value).ok()),
})
}

82
crates/aether-data/src/repository/candidate_selection/postgres.rs
View File

@@ -5,7 +5,7 @@ use std::collections::BTreeSet;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::{error::SqlxResultExt, DataLayerError};
@@ -505,6 +505,36 @@ LIMIT $7
OFFSET $8
"#;
fn pool_key_candidate_order_by_sql(order: &StoredPoolKeyCandidateOrder) -> &'static str {
match order {
StoredPoolKeyCandidateOrder::InternalPriority => {
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::Lru => {
"ORDER BY\n pak.last_used_at ASC NULLS FIRST,\n pak.internal_priority ASC,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::CacheAffinity => {
"ORDER BY\n pak.last_used_at DESC NULLS LAST,\n pak.internal_priority ASC,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::SingleAccount => {
"ORDER BY\n pak.internal_priority ASC,\n pak.last_used_at DESC NULLS LAST,\n pak.id ASC"
}
StoredPoolKeyCandidateOrder::LoadBalance { .. } => {
"ORDER BY\n md5($9 || ':' || pak.id) ASC,\n pak.id ASC"
}
}
}
fn pool_key_candidate_selection_sql(order: &StoredPoolKeyCandidateOrder) -> String {
let default_order =
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC\nLIMIT $7\nOFFSET $8\n";
let replacement = format!(
"{}\nLIMIT $7\nOFFSET $8\n",
pool_key_candidate_order_by_sql(order)
);
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, &replacement)
}
#[derive(Debug, Clone)]
pub struct SqlxMinimalCandidateSelectionReadRepository {
pool: PgPool,
@@ -650,19 +680,23 @@ impl SqlxMinimalCandidateSelectionReadRepository {
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let limit = i64::from(query.limit.max(1));
let offset = i64::from(query.offset);
let sql = pool_key_candidate_selection_sql(&query.order);
for api_format in storage_aliases {
let mut query_builder = sqlx::query(sql.as_str())
.bind(api_format)
.bind(query.provider_id.as_str())
.bind(query.endpoint_id.as_str())
.bind(query.model_id.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(limit)
.bind(offset);
if let StoredPoolKeyCandidateOrder::LoadBalance { seed } = &query.order {
query_builder = query_builder.bind(seed.as_str());
}
rows.extend(
Self::collect_query_rows(
sqlx::query(LIST_POOL_KEYS_FOR_GROUP_SQL)
.bind(api_format)
.bind(query.provider_id.as_str())
.bind(query.endpoint_id.as_str())
.bind(query.model_id.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(limit)
.bind(offset)
.fetch(&self.pool),
query_builder.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
@@ -1039,13 +1073,16 @@ mod tests {
use serde_json::json;
use super::{
parse_provider_model_mappings, parse_string_list, requested_model_selection_page_sql,
requested_model_selection_sql, SqlxMinimalCandidateSelectionReadRepository,
parse_provider_model_mappings, parse_string_list, pool_key_candidate_selection_sql,
requested_model_selection_page_sql, requested_model_selection_sql,
SqlxMinimalCandidateSelectionReadRepository,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL, LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
};
use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory};
use crate::repository::candidate_selection::StoredProviderModelMapping;
use crate::repository::candidate_selection::{
StoredPoolKeyCandidateOrder, StoredProviderModelMapping,
};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
@@ -1100,6 +1137,23 @@ mod tests {
assert!(sql.ends_with("LIMIT $5\nOFFSET $6"));
}
#[test]
fn pool_key_selection_sql_applies_query_order() {
let load_balance_sql =
pool_key_candidate_selection_sql(&StoredPoolKeyCandidateOrder::LoadBalance {
seed: "seed".to_string(),
});
assert!(load_balance_sql.contains("md5($9 || ':' || pak.id) ASC"));
assert!(load_balance_sql.ends_with("LIMIT $7\nOFFSET $8\n"));
let lru_sql = pool_key_candidate_selection_sql(&StoredPoolKeyCandidateOrder::Lru);
assert!(lru_sql.contains("pak.last_used_at ASC NULLS FIRST"));
let cache_affinity_sql =
pool_key_candidate_selection_sql(&StoredPoolKeyCandidateOrder::CacheAffinity);
assert!(cache_affinity_sql.contains("pak.last_used_at DESC NULLS LAST"));
}
#[test]
fn parse_string_list_accepts_stringified_array() {
let parsed = parse_string_list(

85
crates/aether-data/src/repository/candidate_selection/sqlite.rs
View File

@@ -5,7 +5,7 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::sqlite::SqlitePool;
@@ -35,6 +35,7 @@ SELECT
pak.capabilities AS key_capabilities,
pak.internal_priority AS key_internal_priority,
pak.global_priority_by_format AS key_global_priority_by_format,
pak.last_used_at AS key_last_used_at_unix_secs,
m.id AS model_id,
m.global_model_id AS global_model_id,
gm.name AS global_model_name,
@@ -67,6 +68,7 @@ struct CandidateSelectionRow {
row: StoredMinimalCandidateSelectionRow,
provider_pool_enabled: bool,
key_auth_config: Option<String>,
key_last_used_at_unix_secs: Option<u64>,
}
impl SqliteMinimalCandidateSelectionReadRepository {
@@ -180,18 +182,18 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.map(|item| item.row)
.filter(|row| {
row.provider_id == query.provider_id
&& row.endpoint_id == query.endpoint_id
&& row.model_id == query.model_id
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
})
.collect::<Vec<_>>();
let mut rows = sort_pool_key_rows(rows);
let mut rows = sort_pool_key_rows(rows, &query.order);
Ok(rows
.drain(..)
.skip(query.offset as usize)
.take(query.limit as usize)
.map(|item| item.row)
.collect())
}
}
@@ -246,16 +248,66 @@ fn sort_rows(
}
fn sort_pool_key_rows(
mut rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> {
rows.sort_by(|left, right| {
left.key_internal_priority
.cmp(&right.key_internal_priority)
.then(left.key_id.cmp(&right.key_id))
mut rows: Vec<CandidateSelectionRow>,
order: &StoredPoolKeyCandidateOrder,
) -> Vec<CandidateSelectionRow> {
rows.sort_by(|left, right| match order {
StoredPoolKeyCandidateOrder::InternalPriority => compare_pool_key_internal(left, right),
StoredPoolKeyCandidateOrder::Lru => left
.key_last_used_at_unix_secs
.cmp(&right.key_last_used_at_unix_secs)
.then_with(|| compare_pool_key_internal(left, right)),
StoredPoolKeyCandidateOrder::CacheAffinity => right
.key_last_used_at_unix_secs
.cmp(&left.key_last_used_at_unix_secs)
.then_with(|| compare_pool_key_internal(left, right)),
StoredPoolKeyCandidateOrder::SingleAccount => left
.row
.key_internal_priority
.cmp(&right.row.key_internal_priority)
.then_with(|| {
right
.key_last_used_at_unix_secs
.cmp(&left.key_last_used_at_unix_secs)
})
.then(left.row.key_id.cmp(&right.row.key_id)),
StoredPoolKeyCandidateOrder::LoadBalance { seed } => {
stable_pool_key_hash(seed.as_str(), left.row.key_id.as_str())
.cmp(&stable_pool_key_hash(
seed.as_str(),
right.row.key_id.as_str(),
))
.then(left.row.key_id.cmp(&right.row.key_id))
}
});
rows
}
fn compare_pool_key_internal(
left: &CandidateSelectionRow,
right: &CandidateSelectionRow,
) -> std::cmp::Ordering {
left.row
.key_internal_priority
.cmp(&right.row.key_internal_priority)
.then(left.row.key_id.cmp(&right.row.key_id))
}
fn stable_pool_key_hash(seed: &str, key_id: &str) -> u64 {
let mut hash = 0xcbf29ce484222325u64;
for byte in seed
.as_bytes()
.iter()
.copied()
.chain(std::iter::once(b':'))
.chain(key_id.as_bytes().iter().copied())
{
hash ^= u64::from(byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}
fn row_matches_requested_model(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
@@ -394,6 +446,10 @@ fn map_candidate_selection_row(row: &SqliteRow) -> Result<CandidateSelectionRow,
},
provider_pool_enabled,
key_auth_config: row.try_get("key_auth_config").map_sql_err()?,
key_last_used_at_unix_secs: row
.try_get::<Option<i64>, _>("key_last_used_at_unix_secs")
.map_sql_err()?
.and_then(|value| u64::try_from(value).ok()),
})
}
@@ -620,8 +676,8 @@ mod tests {
use super::SqliteMinimalCandidateSelectionReadRepository;
use crate::lifecycle::migrate::run_sqlite_migrations;
use crate::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
MinimalCandidateSelectionReadRepository, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
};
#[tokio::test]
@@ -673,6 +729,7 @@ mod tests {
endpoint_id: "endpoint-1".to_string(),
model_id: "model-1".to_string(),
selected_provider_model_name: "provider-model".to_string(),
order: StoredPoolKeyCandidateOrder::InternalPriority,
offset: 1,
limit: 1,
})