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

feat: add adaptive pool metrics and self-check

Browse Source
This commit is contained in:
fawney19
2026-05-15 01:46:24 +08:00
parent bf511f9f8c
commit 54a8312e46
34 changed files with 4521 additions and 300 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
apps/aether-gateway/src/handlers/admin/mod.rs
View File

@@ -26,8 +26,10 @@ pub(crate) use self::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, update_existing_provider_oauth_catalog_key,
};
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::provider_quota_refresh_endpoint_for_provider;
pub(crate) use self::provider::oauth::quota::shared::provider_type_supports_quota_refresh;
pub(crate) use self::provider::oauth::quota::shared::{
persist_provider_quota_refresh_state, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh,
};
pub(crate) use self::provider::oauth::runtime::{
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
refresh_provider_oauth_account_state_after_update,
@@ -35,6 +37,9 @@ pub(crate) use self::provider::oauth::runtime::{
pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response;
pub(crate) use self::provider::pool::config::admin_provider_pool_config;
pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response;
pub(crate) use self::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
};
pub(crate) use self::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
pub(crate) use self::provider::{
maybe_build_local_admin_provider_oauth_response, maybe_build_local_admin_providers_response,

91
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -37,6 +37,38 @@ fn json_f64(value: &Value) -> Option<f64> {
})
}
fn parse_pool_probe_target_percent(pool_advanced: &Map<String, Value>) -> Option<f64> {
pool_advanced
.get("probing_target_percent")
.or_else(|| pool_advanced.get("probing_active_target_percent"))
.or_else(|| pool_advanced.get("active_probe_target_percent"))
.and_then(json_f64)
.filter(|value| value.is_finite() && *value > 0.0)
.map(|value| value.clamp(0.0, 100.0))
}
fn parse_pool_probe_target_count(pool_advanced: &Map<String, Value>) -> Option<u64> {
pool_advanced
.get("probing_target_count")
.or_else(|| pool_advanced.get("probing_active_target_count"))
.or_else(|| pool_advanced.get("active_probe_target_count"))
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(100_000))
}
fn parse_pool_account_self_check_method(pool_advanced: &Map<String, Value>) -> String {
pool_advanced
.get("account_self_check_method")
.or_else(|| pool_advanced.get("self_check_method"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_ascii_lowercase)
.filter(|value| matches!(value.as_str(), "quota_refresh" | "custom_request"))
.unwrap_or_else(|| "quota_refresh".to_string())
}
fn pool_score_weight(object: &Map<String, Value>, names: &[&str], current: f64) -> f64 {
names
.iter()
@@ -388,7 +420,14 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probing_target_percent: None,
probing_target_count: None,
probe_concurrency: 4,
account_self_check_enabled: false,
account_self_check_interval_minutes: 60,
account_self_check_concurrency: 4,
account_self_check_method: "quota_refresh".to_string(),
account_self_check_request: None,
score_top_n: 128,
score_fallback_scan_limit: 1024,
score_rules: PoolMemberScoreRules::default(),
@@ -456,12 +495,39 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.filter(|value| *value > 0)
.map(|value| value.min(1440))
.unwrap_or(10),
probing_target_percent: parse_pool_probe_target_percent(pool_advanced),
probing_target_count: parse_pool_probe_target_count(pool_advanced),
probe_concurrency: pool_advanced
.get("probe_concurrency")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(64))
.unwrap_or(4),
account_self_check_enabled: pool_advanced
.get("account_self_check_enabled")
.or_else(|| pool_advanced.get("self_check_enabled"))
.and_then(Value::as_bool)
.unwrap_or(false),
account_self_check_interval_minutes: pool_advanced
.get("account_self_check_interval_minutes")
.or_else(|| pool_advanced.get("self_check_interval_minutes"))
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(1440))
.unwrap_or(60),
account_self_check_concurrency: pool_advanced
.get("account_self_check_concurrency")
.or_else(|| pool_advanced.get("self_check_concurrency"))
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(64))
.unwrap_or(4),
account_self_check_method: parse_pool_account_self_check_method(pool_advanced),
account_self_check_request: pool_advanced
.get("account_self_check_request")
.or_else(|| pool_advanced.get("self_check_request"))
.filter(|value| value.is_object())
.cloned(),
score_top_n: pool_advanced
.get("score_top_n")
.and_then(json_u64)
@@ -547,7 +613,18 @@ mod tests {
"health_policy_enabled": false,
"probing_enabled": true,
"probing_interval_minutes": 20,
"probing_target_percent": 25,
"probing_target_count": 3,
"probe_concurrency": 6,
"account_self_check_enabled": true,
"account_self_check_interval_minutes": 90,
"account_self_check_concurrency": 5,
"account_self_check_method": "custom_request",
"account_self_check_request": {
"method": "GET",
"path": "/v1/me",
"success_status_codes": [200]
},
"score_top_n": 256,
"score_fallback_scan_limit": 2048,
"score_rules": {
@@ -584,7 +661,21 @@ mod tests {
assert!(!config.health_policy_enabled);
assert!(config.probing_enabled);
assert_eq!(config.probing_interval_minutes, 20);
assert_eq!(config.probing_target_percent, Some(25.0));
assert_eq!(config.probing_target_count, Some(3));
assert_eq!(config.probe_concurrency, 6);
assert!(config.account_self_check_enabled);
assert_eq!(config.account_self_check_interval_minutes, 90);
assert_eq!(config.account_self_check_concurrency, 5);
assert_eq!(config.account_self_check_method, "custom_request");
assert_eq!(
config
.account_self_check_request
.as_ref()
.and_then(|value| value.get("path"))
.and_then(serde_json::Value::as_str),
Some("/v1/me")
);
assert_eq!(config.score_top_n, 256);
assert_eq!(config.score_fallback_scan_limit, 2048);
assert_eq!(config.score_rules.weights.manual_priority, 0.4);

47
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -5,10 +5,15 @@ use super::keys::{
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState,
};
use crate::maintenance::PoolQuotaProbeWorkerConfig;
use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
};
use aether_runtime_state::{DataLayerError, RuntimeState};
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
@@ -19,6 +24,10 @@ fn current_unix_secs() -> u64 {
.as_secs()
}
fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bool {
pool_config.probing_enabled
}
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState,
provider_ids: &[String],
@@ -102,6 +111,40 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
}
if should_load_active_probe_members(pool_config) {
state.active_probe_member_ids = runtime
.set_members(&admin_provider_pool_quota_probe_active_members_key(
provider_id,
))
.await
.map(|values| {
values
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
}
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let demand_snapshot = read_provider_pool_demand_snapshot(
runtime,
provider_id,
key_ids.len(),
probe_config.max_keys_per_provider,
)
.await;
state.provider_in_flight = demand_snapshot.in_flight;
state.provider_ema_in_flight = demand_snapshot.ema_in_flight;
state.provider_desired_hot = if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
};
state.provider_burst_pending =
pool_config.probing_enabled && provider_pool_burst_pending(runtime, provider_id).await;
if !cooldown_keys.is_empty() {
let cooldown_reasons = runtime
.kv_get_many(&cooldown_keys)

10
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/status.rs
View File

@@ -23,6 +23,11 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"pool_enabled": false,
"total_keys": 0,
"total_sticky_sessions": 0,
"provider_hot_count": 0,
"provider_desired_hot": 0,
"provider_in_flight": 0,
"provider_ema_in_flight": 0.0,
"provider_burst_pending": false,
"keys": [],
}));
};
@@ -68,6 +73,11 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"pool_enabled": true,
"total_keys": key_payloads.len(),
"total_sticky_sessions": runtime.total_sticky_sessions,
"provider_hot_count": runtime.active_probe_member_ids.len(),
"provider_desired_hot": runtime.provider_desired_hot,
"provider_in_flight": runtime.provider_in_flight,
"provider_ema_in_flight": runtime.provider_ema_in_flight,
"provider_burst_pending": runtime.provider_burst_pending,
"keys": key_payloads,
}))
}

7
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
View File

@@ -615,7 +615,14 @@ mod tests {
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probing_target_percent: None,
probing_target_count: None,
probe_concurrency: 4,
account_self_check_enabled: false,
account_self_check_interval_minutes: 60,
account_self_check_concurrency: 4,
account_self_check_method: "quota_refresh".to_string(),
account_self_check_request: None,
score_top_n: 128,
score_fallback_scan_limit: 1024,
score_rules: aether_pool_core::PoolMemberScoreRules::default(),

82
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/overview.rs
View File

@@ -5,7 +5,12 @@ use super::{
read_admin_provider_pool_cooldown_counts,
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
};
use crate::handlers::admin::provider::shared::support::admin_provider_pool_quota_probe_active_members_key;
use crate::handlers::admin::request::AdminAppState;
use crate::maintenance::PoolQuotaProbeWorkerConfig;
use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
};
use crate::GatewayError;
use aether_admin::provider::pool as admin_provider_pool_pure;
use axum::{
@@ -14,6 +19,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use serde_json::{json, Value};
pub(super) async fn build_admin_pool_overview_response(
state: &AdminAppState<'_>,
@@ -60,17 +66,77 @@ pub(super) async fn build_admin_pool_overview_response(
.map(|item| (item.provider_id.clone(), item))
.collect::<BTreeMap<_, _>>();
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let mut runtime_metrics_by_provider = BTreeMap::new();
for (provider, pool_config) in &pool_enabled_providers {
let active_keys = key_stats_by_provider
.get(&provider.id)
.map(|item| item.active_keys as usize)
.unwrap_or(0);
let hot_count = if pool_config.probing_enabled {
state
.runtime_state()
.set_len(&admin_provider_pool_quota_probe_active_members_key(
&provider.id,
))
.await
.unwrap_or(0)
} else {
0
};
let demand_snapshot = read_provider_pool_demand_snapshot(
state.runtime_state(),
&provider.id,
active_keys,
probe_config.max_keys_per_provider,
)
.await;
let burst_pending = pool_config.probing_enabled
&& provider_pool_burst_pending(state.runtime_state(), &provider.id).await;
runtime_metrics_by_provider.insert(
provider.id.clone(),
json!({
"provider_hot_count": hot_count,
"provider_desired_hot": if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
},
"provider_in_flight": demand_snapshot.in_flight,
"provider_ema_in_flight": demand_snapshot.ema_in_flight,
"provider_burst_pending": burst_pending,
}),
);
}
let providers = pool_enabled_providers
.into_iter()
.map(|(provider, _)| provider)
.collect::<Vec<_>>();
Ok(
Json(admin_provider_pool_pure::build_admin_pool_overview_payload(
&providers,
&key_stats_by_provider,
&cooldown_counts_by_provider,
))
.into_response(),
)
let mut payload = admin_provider_pool_pure::build_admin_pool_overview_payload(
&providers,
&key_stats_by_provider,
&cooldown_counts_by_provider,
);
if let Some(items) = payload.get_mut("items").and_then(Value::as_array_mut) {
for item in items {
let Some(provider_id) = item.get("provider_id").and_then(Value::as_str) else {
continue;
};
let Some(metrics) = runtime_metrics_by_provider.get(provider_id) else {
continue;
};
let Some(item_object) = item.as_object_mut() else {
continue;
};
if let Some(metrics_object) = metrics.as_object() {
for (key, value) in metrics_object {
item_object.insert(key.clone(), value.clone());
}
}
}
}
Ok(Json(payload).into_response())
}

20
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
View File

@@ -2,15 +2,21 @@ use crate::handlers::admin::request::AdminAppState;
use crate::LocalProviderDeleteTaskState;
use aether_pool_core::PoolMemberScoreRules;
use serde_json::json;
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_KEYS: usize = 200;
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_MODELS: usize = 500;
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_FETCH_LIMIT: usize = 10_000;
pub(crate) const ADMIN_PROVIDER_POOL_SCAN_BATCH: u64 = 200;
pub(crate) const ADMIN_PROVIDER_POOL_QUOTA_PROBE_ACTIVE_SET_PREFIX: &str =
"ap:quota_probe:active_members";
pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str =
"Admin provider OAuth data unavailable";
pub(crate) fn admin_provider_pool_quota_probe_active_members_key(provider_id: &str) -> String {
format!("{ADMIN_PROVIDER_POOL_QUOTA_PROBE_ACTIVE_SET_PREFIX}:{provider_id}")
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AdminProviderPoolSchedulingPreset {
pub(crate) preset: String,
@@ -40,7 +46,14 @@ pub(crate) struct AdminProviderPoolConfig {
pub(crate) health_policy_enabled: bool,
pub(crate) probing_enabled: bool,
pub(crate) probing_interval_minutes: u64,
pub(crate) probing_target_percent: Option<f64>,
pub(crate) probing_target_count: Option<u64>,
pub(crate) probe_concurrency: u64,
pub(crate) account_self_check_enabled: bool,
pub(crate) account_self_check_interval_minutes: u64,
pub(crate) account_self_check_concurrency: u64,
pub(crate) account_self_check_method: String,
pub(crate) account_self_check_request: Option<serde_json::Value>,
pub(crate) score_top_n: u64,
pub(crate) score_fallback_scan_limit: u64,
pub(crate) score_rules: PoolMemberScoreRules,
@@ -54,6 +67,11 @@ pub(crate) struct AdminProviderPoolRuntimeState {
pub(crate) total_sticky_sessions: usize,
pub(crate) sticky_sessions_by_key: BTreeMap<String, usize>,
pub(crate) sticky_bound_key_id: Option<String>,
pub(crate) active_probe_member_ids: BTreeSet<String>,
pub(crate) provider_in_flight: usize,
pub(crate) provider_ema_in_flight: f64,
pub(crate) provider_desired_hot: usize,
pub(crate) provider_burst_pending: bool,
pub(crate) cooldown_reason_by_key: BTreeMap<String, String>,
pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>,
pub(crate) cost_window_usage_by_key: BTreeMap<String, u64>,

3
apps/aether-gateway/src/handlers/shared/provider_pool.rs
View File

@@ -8,6 +8,7 @@ pub(crate) use super::super::admin::provider::pool::runtime::{
release_admin_provider_pool_key_lease,
};
pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule, ADMIN_PROVIDER_POOL_SCAN_BATCH,
};