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

Implement generic pool member scoring and probing

Browse Source
This commit is contained in:
fawney19
2026-05-12 01:46:22 +08:00
parent cb0ccb9cdb
commit b9e62d1667
76 changed files with 6276 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/ai_serving/planner/mod.rs
View File

@@ -16,6 +16,7 @@ mod materialization_policy;
mod passthrough;
mod plan_builders;
mod pool_scheduler;
pub(crate) mod pool_scores;
mod report_context;
mod route;
mod runtime_miss;
@@ -36,6 +37,7 @@ pub(crate) use self::plan_builders::{
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
pub(crate) use self::pool_scores::build_provider_key_pool_score_upsert;
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::specialized::{
build_local_gemini_files_stream_attempt_source_for_kind,

164
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
View File

@@ -9,7 +9,11 @@ use aether_ai_serving::{
};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
};
use aether_data_contracts::repository::pool_scores::{
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{Map, Value};
@@ -37,6 +41,8 @@ use crate::handlers::shared::{
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::provider_key_auth::provider_key_auth_semantics;
use super::pool_scores::provider_key_pool_score_scope;
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
const DEFAULT_POOL_KEY_PAGE_SIZE: u32 = 128;
const DEFAULT_POOL_MAX_SCANNED_KEYS: u32 = 1024;
@@ -176,6 +182,8 @@ pub(crate) struct PoolKeyCursor<'a> {
scanned_keys: u32,
page_size: u32,
max_scanned_keys: u32,
score_top_n: u32,
score_phase_loaded: bool,
skip_reason_counts: BTreeMap<&'static str, u32>,
next_pool_key_index: u32,
sticky_candidate_loaded: bool,
@@ -194,6 +202,17 @@ impl<'a> PoolKeyCursor<'a> {
request_auth_channel: Option<&str>,
) -> Self {
let pool_key_order = pool_key_candidate_order_for_group(&group);
let pool_config = pool_config_for_candidate(&group);
let score_top_n = pool_config
.as_ref()
.map(|config| config.score_top_n)
.unwrap_or(u64::from(DEFAULT_POOL_KEY_PAGE_SIZE))
.clamp(1, u64::from(u32::MAX)) as u32;
let max_scanned_keys = pool_config
.as_ref()
.map(|config| config.score_fallback_scan_limit)
.unwrap_or(u64::from(DEFAULT_POOL_MAX_SCANNED_KEYS))
.clamp(1, u64::from(u32::MAX)) as u32;
Self {
state,
group,
@@ -204,7 +223,9 @@ impl<'a> PoolKeyCursor<'a> {
next_offset: 0,
scanned_keys: 0,
page_size: DEFAULT_POOL_KEY_PAGE_SIZE,
max_scanned_keys: DEFAULT_POOL_MAX_SCANNED_KEYS,
max_scanned_keys,
score_top_n,
score_phase_loaded: false,
skip_reason_counts: BTreeMap::new(),
next_pool_key_index: 0,
sticky_candidate_loaded: false,
@@ -257,6 +278,13 @@ impl<'a> PoolKeyCursor<'a> {
}
async fn next_page_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
if !self.score_phase_loaded {
self.score_phase_loaded = true;
if let Some(score_candidates) = self.next_score_candidates().await {
return Some(score_candidates);
}
}
if self.scanned_keys >= self.max_scanned_keys {
return None;
}
@@ -306,6 +334,82 @@ impl<'a> PoolKeyCursor<'a> {
Some(self.build_page_eligible_candidates(rows).await)
}
async fn next_score_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
let scope = provider_key_pool_score_scope(
self.group.candidate.endpoint_api_format.as_str(),
Some(self.group.candidate.model_id.as_str()),
);
let query = ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: self.group.candidate.provider_id.clone(),
capability: scope.capability.clone(),
scope_kind: scope.scope_kind.clone(),
scope_id: scope.scope_id.clone(),
hard_states: vec![PoolMemberHardState::Available, PoolMemberHardState::Unknown],
probe_statuses: None,
offset: 0,
limit: self.score_top_n as usize,
};
let scores = match self.state.app().data.list_ranked_pool_members(&query).await {
Ok(scores) => scores,
Err(err) => {
warn!(
event_name = "pool_group_score_load_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
error = ?err,
"gateway pool scheduler failed to read ranked pool member scores"
);
return None;
}
};
if scores.is_empty() {
return None;
}
self.record_score_schedule_interest(&scores).await;
let key_ids = scores
.iter()
.map(|score| score.member_id.clone())
.collect::<Vec<_>>();
let rows_query = StoredPoolKeyCandidateRowsByKeyIdsQuery {
api_format: self.group.candidate.endpoint_api_format.clone(),
provider_id: self.group.candidate.provider_id.clone(),
endpoint_id: self.group.candidate.endpoint_id.clone(),
model_id: self.group.candidate.model_id.clone(),
selected_provider_model_name: self.group.candidate.selected_provider_model_name.clone(),
key_ids,
};
let rows = match self
.state
.app()
.list_pool_key_candidate_rows_for_group_key_ids(&rows_query)
.await
{
Ok(rows) => rows,
Err(err) => {
warn!(
event_name = "pool_group_score_key_load_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
score_count = scores.len(),
error = ?err,
"gateway pool scheduler failed to materialize ranked pool keys"
);
return None;
}
};
self.scanned_keys = self.scanned_keys.saturating_add(scores.len() as u32);
Some(self.build_page_eligible_candidates(rows).await)
}
async fn sticky_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
let pool_config = pool_config_for_candidate(&self.group)?;
if !admin_provider_pool_cache_affinity_enabled(&pool_config) {
@@ -465,6 +569,62 @@ impl<'a> PoolKeyCursor<'a> {
self.skipped_candidates.append(&mut skipped);
}
async fn record_score_schedule_interest(&self, scores: &[StoredPoolMemberScore]) {
if scores.is_empty() {
return;
}
let scheduled_at = current_unix_ms() / 1000;
let mut failed = 0usize;
for score in scores {
let identity = PoolMemberIdentity {
pool_kind: score.pool_kind.clone(),
pool_id: score.pool_id.clone(),
member_kind: score.member_kind.clone(),
member_id: score.member_id.clone(),
};
let scope = PoolScoreScope {
capability: score.capability.clone(),
scope_kind: score.scope_kind.clone(),
scope_id: score.scope_id.clone(),
};
let result = self
.state
.app()
.data
.record_pool_member_schedule_feedback(PoolMemberScheduleFeedback {
identity,
scope: Some(scope),
scheduled_at,
succeeded: None,
hard_state: None,
score_delta: None,
score_reason_patch: Some(serde_json::json!({
"last_schedule_interest": {
"provider_id": self.group.candidate.provider_id.as_str(),
"endpoint_id": self.group.candidate.endpoint_id.as_str(),
"model_id": self.group.candidate.model_id.as_str()
}
})),
})
.await;
if result.is_err() {
failed += 1;
}
}
if failed > 0 {
warn!(
event_name = "pool_group_score_interest_update_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
failed_count = failed,
score_count = scores.len(),
"gateway pool scheduler failed to record some pool score schedule interests"
);
}
}
async fn build_page_eligible_candidates(
&mut self,
rows: Vec<StoredMinimalCandidateSelectionRow>,

192
apps/aether-gateway/src/ai_serving/planner/pool_scores.rs Normal file
View File

@@ -0,0 +1,192 @@
use aether_ai_serving::{score_pool_member, PoolMemberScoreInput, POOL_SCORE_VERSION};
use aether_data_contracts::repository::pool_scores::{
PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore,
POOL_SCORE_SCOPE_KIND_MODEL,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::Value;
use crate::handlers::shared::{provider_key_health_summary, provider_key_status_snapshot_payload};
pub(crate) fn build_provider_key_pool_score_upsert(
key: &StoredProviderCatalogKey,
provider_type: &str,
api_format: &str,
model_id: Option<&str>,
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
now_unix_secs: u64,
) -> UpsertPoolMemberScore {
let identity = PoolMemberIdentity::provider_api_key(key.provider_id.clone(), key.id.clone());
let scope = provider_key_pool_score_scope(api_format, model_id);
let input = provider_key_score_input(
key,
provider_type,
identity.clone(),
scope.clone(),
existing,
now_unix_secs,
);
let output = score_pool_member(&input);
UpsertPoolMemberScore {
id: provider_key_pool_score_id(&identity, &scope),
identity,
scope,
score: output.score,
hard_state: output.hard_state,
score_version: POOL_SCORE_VERSION,
score_reason: output.score_reason,
last_ranked_at: Some(now_unix_secs),
last_scheduled_at: existing.and_then(|score| score.last_scheduled_at),
last_success_at: existing.and_then(|score| score.last_success_at),
last_failure_at: existing.and_then(|score| score.last_failure_at),
failure_count: existing.map(|score| score.failure_count).unwrap_or(0),
last_probe_attempt_at: existing.and_then(|score| score.last_probe_attempt_at),
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
last_probe_failure_at: existing.and_then(|score| score.last_probe_failure_at),
probe_failure_count: existing.map(|score| score.probe_failure_count).unwrap_or(0),
probe_status: existing
.map(|score| score.probe_status)
.unwrap_or(PoolMemberProbeStatus::Never),
updated_at: now_unix_secs,
}
}
pub(crate) fn provider_key_pool_score_scope(
api_format: &str,
model_id: Option<&str>,
) -> PoolScoreScope {
PoolScoreScope {
capability: api_format.trim().to_ascii_lowercase(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: model_id
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
}
}
pub(crate) fn provider_key_pool_score_id(
identity: &PoolMemberIdentity,
scope: &PoolScoreScope,
) -> String {
let raw = format!(
"{}:{}:{}:{}:{}:{}:{}",
identity.pool_kind,
identity.pool_id,
identity.member_kind,
identity.member_id,
scope.capability,
scope.scope_kind,
scope.scope_id.as_deref().unwrap_or("*")
);
format!(
"pms-{:016x}-{:016x}",
stable_hash(raw.as_bytes()),
stable_hash(identity.member_id.as_bytes())
)
}
fn provider_key_score_input(
key: &StoredProviderCatalogKey,
provider_type: &str,
identity: PoolMemberIdentity,
scope: PoolScoreScope,
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
now_unix_secs: u64,
) -> PoolMemberScoreInput {
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
let quota_snapshot = status_snapshot
.as_object()
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object);
let account_snapshot = status_snapshot
.as_object()
.and_then(|snapshot| snapshot.get("account"))
.and_then(Value::as_object);
let (health_score, _, _, any_circuit_open, _) = provider_key_health_summary(key);
let health_score = key
.health_by_format
.as_ref()
.and_then(Value::as_object)
.filter(|payload| !payload.is_empty())
.map(|_| health_score);
PoolMemberScoreInput {
identity,
scope: scope.clone(),
internal_priority: key.internal_priority,
is_active: key.is_active,
health_score,
quota_usage_ratio: quota_snapshot
.and_then(|quota| quota.get("usage_ratio"))
.and_then(json_f64)
.map(|value| value.clamp(0.0, 1.0)),
quota_exhausted: quota_snapshot
.and_then(|quota| quota.get("exhausted"))
.and_then(Value::as_bool)
.unwrap_or(false),
account_blocked: account_snapshot
.and_then(|account| account.get("blocked"))
.and_then(Value::as_bool)
.unwrap_or(false),
oauth_invalid_reason: key.oauth_invalid_reason.clone(),
circuit_open: any_circuit_open
|| circuit_open_for_scope(key.circuit_breaker_by_format.as_ref(), &scope),
success_count: key.success_count.unwrap_or(0).into(),
error_count: key.error_count.unwrap_or(0).into(),
total_response_time_ms: key.total_response_time_ms.unwrap_or(0).into(),
total_tokens: key.total_tokens,
total_cost_usd: key.total_cost_usd,
last_used_at: key.last_used_at_unix_secs,
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
probe_status: existing
.map(|score| score.probe_status)
.unwrap_or(PoolMemberProbeStatus::Never),
now_unix_secs,
}
}
fn circuit_open_for_scope(circuit_by_format: Option<&Value>, scope: &PoolScoreScope) -> bool {
let Some(formats) = circuit_by_format.and_then(Value::as_object) else {
return false;
};
let keys = api_format_lookup_keys(&scope.capability);
keys.iter().any(|key| {
formats
.get(key)
.and_then(Value::as_object)
.and_then(|value| value.get("open"))
.and_then(Value::as_bool)
.unwrap_or(false)
})
}
fn json_f64(value: &Value) -> Option<f64> {
value.as_f64().or_else(|| {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| value.parse::<f64>().ok())
})
}
fn api_format_lookup_keys(api_format: &str) -> Vec<String> {
let normalized = aether_ai_formats::normalize_api_format_alias(api_format);
let mut keys = aether_ai_formats::api_format_storage_aliases(&normalized);
if !keys.iter().any(|value| value == &normalized) {
keys.push(normalized);
}
keys.sort();
keys.dedup();
keys
}
fn stable_hash(bytes: &[u8]) -> u64 {
let mut hash = 0xcbf29ce484222325u64;
for byte in bytes {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}