Merge remote-tracking branch 'origin/codex/pool-member-scores' into aether-rust-pioneer

This commit is contained in:
fawney19
2026-05-12 09:17:25 +08:00
81 changed files with 7206 additions and 77 deletions

View File

@@ -1,6 +1,8 @@
use crate::handlers::admin::admin_provider_pool_config;
use crate::handlers::admin::provider::shared::paths::admin_provider_id_for_keys;
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
use crate::provider_key_auth::provider_key_effective_api_formats;
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
use axum::{
@@ -98,6 +100,27 @@ pub(super) async fn maybe_handle(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
if let Some(pool_config) = admin_provider_pool_config(&provider) {
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
state.as_ref(),
&provider,
&pool_config,
&endpoints,
std::slice::from_ref(&created),
now_unix_secs,
score_ensure_budget,
)
.await
{
tracing::debug!(
provider_id = %provider.id,
key_id = %created.id,
error = ?err,
"gateway admin provider key create: failed to seed pool score rows"
);
}
}
let api_formats =
provider_key_effective_api_formats(&created, &provider.provider_type, &endpoints);

View File

@@ -1,6 +1,8 @@
use crate::handlers::admin::admin_provider_pool_config;
use crate::handlers::admin::provider::shared::paths::admin_update_key_id;
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
use crate::provider_key_auth::provider_key_effective_api_formats;
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
use axum::{
@@ -121,6 +123,27 @@ pub(super) async fn maybe_handle(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
if let Some(pool_config) = admin_provider_pool_config(&provider) {
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
state.as_ref(),
&provider,
&pool_config,
&endpoints,
std::slice::from_ref(&updated),
now_unix_secs,
score_ensure_budget,
)
.await
{
tracing::debug!(
provider_id = %provider.id,
key_id = %updated.id,
error = ?err,
"gateway admin provider key update: failed to seed pool score rows"
);
}
}
let api_formats =
provider_key_effective_api_formats(&updated, &provider.provider_type, &endpoints);

View File

@@ -1,6 +1,7 @@
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
};
use aether_ai_serving::{PoolMemberScoreRules, PoolMemberScoreWeights};
use serde_json::{Map, Value};
const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[
@@ -26,6 +27,123 @@ fn json_u64(value: &Value) -> Option<u64> {
.or_else(|| value.as_i64().and_then(|raw| u64::try_from(raw).ok()))
}
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 pool_score_weight(object: &Map<String, Value>, names: &[&str], current: f64) -> f64 {
names
.iter()
.find_map(|name| {
object
.get(*name)
.and_then(json_f64)
.filter(|value| value.is_finite() && *value >= 0.0)
})
.unwrap_or(current)
}
fn parse_pool_score_weights(
raw_weights: Option<&Map<String, Value>>,
current: PoolMemberScoreWeights,
) -> PoolMemberScoreWeights {
let Some(raw_weights) = raw_weights else {
return current;
};
PoolMemberScoreWeights {
manual_priority: pool_score_weight(
raw_weights,
&["manual_priority", "priority", "internal_priority"],
current.manual_priority,
),
health: pool_score_weight(raw_weights, &["health"], current.health),
probe_freshness: pool_score_weight(
raw_weights,
&["probe_freshness", "freshness", "probe"],
current.probe_freshness,
),
quota_remaining: pool_score_weight(
raw_weights,
&["quota_remaining", "quota", "quota_available"],
current.quota_remaining,
),
latency: pool_score_weight(raw_weights, &["latency"], current.latency),
cost_lru: pool_score_weight(
raw_weights,
&["cost_lru", "cost_remaining", "cost", "lru"],
current.cost_lru,
),
}
}
fn parse_pool_score_rules(pool_advanced: &Map<String, Value>) -> PoolMemberScoreRules {
let mut rules = PoolMemberScoreRules::default();
for key in ["score_weights", "pool_score_weights", "scoring_weights"] {
rules.weights = parse_pool_score_weights(
pool_advanced.get(key).and_then(Value::as_object),
rules.weights,
);
}
if let Some(score_rules) = pool_advanced
.get("score_rules")
.or_else(|| pool_advanced.get("pool_score_rules"))
.and_then(Value::as_object)
{
rules.weights = parse_pool_score_weights(
score_rules.get("weights").and_then(Value::as_object),
rules.weights,
);
if let Some(ttl_seconds) = score_rules
.get("probe_freshness_ttl_seconds")
.or_else(|| score_rules.get("score_probe_freshness_ttl_seconds"))
.and_then(json_u64)
.filter(|value| *value > 0)
{
rules.probe_freshness_ttl_seconds = ttl_seconds.min(7 * 24 * 3600);
}
if let Some(cap) = score_rules
.get("unschedulable_score_cap")
.or_else(|| score_rules.get("hard_state_score_cap"))
.and_then(json_f64)
.filter(|value| value.is_finite())
{
rules.unschedulable_score_cap = cap.clamp(0.0, 1.0);
}
if let Some(penalty) = score_rules
.get("probe_failure_penalty")
.and_then(json_f64)
.filter(|value| value.is_finite())
{
rules.probe_failure_penalty = penalty.clamp(0.0, 1.0);
}
if let Some(penalty) = score_rules
.get("request_failure_penalty")
.or_else(|| score_rules.get("runtime_failure_penalty"))
.and_then(json_f64)
.filter(|value| value.is_finite())
{
rules.request_failure_penalty = penalty.clamp(0.0, 1.0);
}
if let Some(threshold) = score_rules
.get("probe_failure_cooldown_threshold")
.or_else(|| score_rules.get("probe_failure_hard_state_threshold"))
.and_then(json_u64)
{
rules.probe_failure_cooldown_threshold = threshold.min(100);
}
}
rules.effective()
}
fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
match preset {
"free_first" | "team_first" | "plus_first" | "pro_first" => {
@@ -270,6 +388,10 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probe_concurrency: 4,
score_top_n: 128,
score_fallback_scan_limit: 1024,
score_rules: PoolMemberScoreRules::default(),
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,
@@ -278,6 +400,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
let scheduling_presets = parse_pool_scheduling_presets(pool_advanced);
let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced);
let score_rules = parse_pool_score_rules(pool_advanced);
Some(AdminProviderPoolConfig {
lru_enabled: admin_provider_pool_lru_enabled(&scheduling_presets),
@@ -333,6 +456,25 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.filter(|value| *value > 0)
.map(|value| value.min(1440))
.unwrap_or(10),
probe_concurrency: pool_advanced
.get("probe_concurrency")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(64))
.unwrap_or(4),
score_top_n: pool_advanced
.get("score_top_n")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(4096))
.unwrap_or(128),
score_fallback_scan_limit: pool_advanced
.get("score_fallback_scan_limit")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(50_000))
.unwrap_or(1024),
score_rules,
stream_timeout_threshold: pool_advanced
.get("stream_timeout_threshold")
.and_then(json_u64)
@@ -405,6 +547,24 @@ mod tests {
"health_policy_enabled": false,
"probing_enabled": true,
"probing_interval_minutes": 20,
"probe_concurrency": 6,
"score_top_n": 256,
"score_fallback_scan_limit": 2048,
"score_rules": {
"weights": {
"manual_priority": 0.4,
"health": 0.2,
"probe_freshness": 0.2,
"quota_remaining": 0.1,
"latency": 0.05,
"cost_lru": 0.05
},
"probe_freshness_ttl_seconds": 1200,
"unschedulable_score_cap": 0.03,
"probe_failure_penalty": 0.08,
"request_failure_penalty": 0.01,
"probe_failure_cooldown_threshold": 2
},
"stream_timeout_threshold": 4,
"stream_timeout_window_seconds": 900,
"stream_timeout_cooldown_seconds": 180
@@ -424,6 +584,16 @@ mod tests {
assert!(!config.health_policy_enabled);
assert!(config.probing_enabled);
assert_eq!(config.probing_interval_minutes, 20);
assert_eq!(config.probe_concurrency, 6);
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);
assert_eq!(config.score_rules.weights.health, 0.2);
assert_eq!(config.score_rules.probe_freshness_ttl_seconds, 1200);
assert_eq!(config.score_rules.unschedulable_score_cap, 0.03);
assert_eq!(config.score_rules.probe_failure_penalty, 0.08);
assert_eq!(config.score_rules.request_failure_penalty, 0.01);
assert_eq!(config.score_rules.probe_failure_cooldown_threshold, 2);
assert_eq!(config.stream_timeout_threshold, 4);
assert_eq!(config.stream_timeout_window_seconds, 900);
assert_eq!(config.stream_timeout_cooldown_seconds, 180);
@@ -462,6 +632,27 @@ mod tests {
assert_eq!(config.sticky_session_ttl_seconds, 0);
}
#[test]
fn parses_legacy_pool_score_weights_from_pool_advanced() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scoring_weights": {
"manual_priority": 0,
"health": 2,
"probe": 1,
"quota_remaining": 0,
"latency": 0,
"cost_remaining": 1
}
}
})))
.expect("pool config should parse");
assert_eq!(config.score_rules.weights.health, 0.5);
assert_eq!(config.score_rules.weights.probe_freshness, 0.25);
assert_eq!(config.score_rules.weights.cost_lru, 0.25);
}
#[test]
fn parses_pool_config_from_generic_config_value() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({

View File

@@ -615,6 +615,10 @@ mod tests {
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probe_concurrency: 4,
score_top_n: 128,
score_fallback_scan_limit: 1024,
score_rules: aether_ai_serving::PoolMemberScoreRules::default(),
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,

View File

@@ -24,6 +24,8 @@ mod read_overview;
mod read_presets;
#[path = "read_routes/resolve_selection.rs"]
mod read_resolve_selection;
#[path = "read_routes/scores.rs"]
mod read_scores;
pub(crate) mod selection;
mod support;
@@ -34,11 +36,11 @@ pub(crate) use self::batch_shared::{
AdminPoolBatchImportRequest,
};
pub(crate) use self::support::{
admin_pool_provider_id_from_path, parse_admin_pool_key_sort, parse_admin_pool_page,
parse_admin_pool_page_size, parse_admin_pool_quick_selectors, parse_admin_pool_search,
parse_admin_pool_status_filter, AdminPoolKeySort, AdminPoolKeySortDirection,
AdminPoolKeySortField, AdminPoolResolveSelectionRequest,
ADMIN_POOL_BANNED_KEY_CLEANUP_EMPTY_MESSAGE,
admin_pool_provider_id_from_path, admin_pool_provider_id_from_scores_path,
parse_admin_pool_key_sort, parse_admin_pool_page, parse_admin_pool_page_size,
parse_admin_pool_quick_selectors, parse_admin_pool_search, parse_admin_pool_status_filter,
AdminPoolKeySort, AdminPoolKeySortDirection, AdminPoolKeySortField,
AdminPoolResolveSelectionRequest, ADMIN_POOL_BANNED_KEY_CLEANUP_EMPTY_MESSAGE,
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
ADMIN_POOL_PROVIDER_CATALOG_WRITER_UNAVAILABLE_DETAIL,
};
@@ -103,6 +105,11 @@ pub(crate) async fn maybe_build_local_admin_pool_response(
read_keys::build_admin_pool_list_keys_response(state, request_context).await?,
));
}
Some("scores") => {
return Ok(Some(
read_scores::build_admin_pool_scores_response(state, request_context).await?,
));
}
Some("resolve_selection") => {
return Ok(Some(
read_resolve_selection::build_admin_pool_resolve_selection_response(

View File

@@ -0,0 +1,196 @@
use super::{
admin_pool_provider_id_from_scores_path, build_admin_pool_error_response,
parse_admin_pool_page, parse_admin_pool_page_size,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::shared::query_param_value;
use crate::GatewayError;
use aether_data_contracts::repository::pool_scores::{
ListPoolMemberScoresQuery, PoolMemberHardState, PoolMemberProbeStatus,
POOL_KIND_PROVIDER_KEY_POOL, POOL_SCORE_SCOPE_KIND_MODEL,
};
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use std::collections::BTreeMap;
pub(super) async fn build_admin_pool_scores_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
let Some(provider_id) = admin_pool_provider_id_from_scores_path(request_context.path()) else {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
"provider_id 无效",
));
};
let query = request_context.query_string();
let page = match parse_admin_pool_page(query) {
Ok(value) => value,
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let page_size = match parse_admin_pool_page_size(query) {
Ok(value) => value.min(500),
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let offset = page.saturating_sub(1).saturating_mul(page_size);
let api_format = query_param_value(query, "api_format")
.map(|value| crate::ai_serving::normalize_api_format_alias(value.as_str()));
let model_id = query_param_value(query, "model_id");
let hard_states = match parse_hard_state_filter(query) {
Ok(value) => value,
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let probe_statuses = match parse_probe_status_filter(query) {
Ok(value) => value,
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let scores = state
.app()
.data
.list_pool_member_scores(&ListPoolMemberScoresQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: provider_id.clone(),
capability: api_format.clone(),
scope_kind: Some(POOL_SCORE_SCOPE_KIND_MODEL.to_string()),
scope_id: model_id.clone(),
hard_states,
probe_statuses,
offset,
limit: page_size,
})
.await
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?;
let key_ids = scores
.iter()
.map(|score| score.member_id.clone())
.collect::<Vec<_>>();
let keys = state
.app()
.read_provider_catalog_keys_by_ids(&key_ids)
.await
.unwrap_or_default()
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>();
let items = scores
.into_iter()
.map(|score| {
let key = keys.get(&score.member_id);
json!({
"id": score.id,
"pool_kind": score.pool_kind,
"pool_id": score.pool_id,
"member_kind": score.member_kind,
"member_id": score.member_id,
"capability": score.capability,
"scope_kind": score.scope_kind,
"scope_id": score.scope_id,
"score": score.score,
"hard_state": score.hard_state.as_database(),
"score_version": score.score_version,
"score_reason": score.score_reason,
"last_ranked_at": score.last_ranked_at,
"last_scheduled_at": score.last_scheduled_at,
"last_success_at": score.last_success_at,
"last_failure_at": score.last_failure_at,
"failure_count": score.failure_count,
"last_probe_attempt_at": score.last_probe_attempt_at,
"last_probe_success_at": score.last_probe_success_at,
"last_probe_failure_at": score.last_probe_failure_at,
"probe_failure_count": score.probe_failure_count,
"probe_status": score.probe_status.as_database(),
"updated_at": score.updated_at,
"key": key.map(|key| json!({
"id": key.id,
"name": key.name,
"auth_type": key.auth_type,
"is_active": key.is_active,
"internal_priority": key.internal_priority,
"last_used_at": key.last_used_at_unix_secs,
}))
})
})
.collect::<Vec<_>>();
Ok(Json(json!({
"provider_id": provider_id,
"page": page,
"page_size": page_size,
"filters": {
"api_format": api_format,
"model_id": model_id,
"hard_state": query_param_value(query, "hard_state"),
"probe_status": query_param_value(query, "probe_status")
},
"items": items
}))
.into_response())
}
fn parse_hard_state_filter(query: Option<&str>) -> Result<Vec<PoolMemberHardState>, String> {
let Some(raw) = query_param_value(query, "hard_state") else {
return Ok(Vec::new());
};
raw.split(',')
.map(|value| match value.trim() {
"available" => Ok(PoolMemberHardState::Available),
"unknown" => Ok(PoolMemberHardState::Unknown),
"cooldown" => Ok(PoolMemberHardState::Cooldown),
"quota_exhausted" => Ok(PoolMemberHardState::QuotaExhausted),
"auth_invalid" => Ok(PoolMemberHardState::AuthInvalid),
"banned" => Ok(PoolMemberHardState::Banned),
"inactive" => Ok(PoolMemberHardState::Inactive),
_ => Err("hard_state must be one of: available, unknown, cooldown, quota_exhausted, auth_invalid, banned, inactive".to_string()),
})
.collect()
}
fn parse_probe_status_filter(
query: Option<&str>,
) -> Result<Option<Vec<PoolMemberProbeStatus>>, String> {
let Some(raw) = query_param_value(query, "probe_status") else {
return Ok(None);
};
raw.split(',')
.map(|value| match value.trim() {
"never" => Ok(PoolMemberProbeStatus::Never),
"ok" => Ok(PoolMemberProbeStatus::Ok),
"failed" => Ok(PoolMemberProbeStatus::Failed),
"stale" => Ok(PoolMemberProbeStatus::Stale),
"in_progress" => Ok(PoolMemberProbeStatus::InProgress),
_ => Err(
"probe_status must be one of: never, ok, failed, stale, in_progress".to_string(),
),
})
.collect::<Result<Vec<_>, _>>()
.map(Some)
}

View File

@@ -149,6 +149,18 @@ pub(crate) fn admin_pool_provider_id_from_path(request_path: &str) -> Option<Str
}
}
pub(crate) fn admin_pool_provider_id_from_scores_path(request_path: &str) -> Option<String> {
let raw = request_path.strip_prefix("/api/admin/pool/")?;
let mut segments = raw.split('/');
let provider_id = segments.next()?.trim();
let scores_segment = segments.next()?.trim_end_matches('/').trim();
if provider_id.is_empty() || scores_segment != "scores" {
None
} else {
Some(provider_id.to_string())
}
}
pub(crate) fn is_admin_pool_route(request_context: &AdminRequestContext<'_>) -> bool {
let normalized_path = request_context.path().trim_end_matches('/');
let path = if normalized_path.is_empty() {
@@ -164,6 +176,10 @@ pub(crate) fn is_admin_pool_route(request_context: &AdminRequestContext<'_>) ->
&& path.starts_with("/api/admin/pool/")
&& path.ends_with("/keys")
&& path.matches('/').count() == 5)
|| (request_context.method() == http::Method::GET
&& path.starts_with("/api/admin/pool/")
&& path.ends_with("/scores")
&& path.matches('/').count() == 5)
|| (request_context.method() == http::Method::POST
&& path.starts_with("/api/admin/pool/")
&& path.ends_with("/keys/batch-import")

View File

@@ -1,5 +1,6 @@
use crate::handlers::admin::request::AdminAppState;
use crate::LocalProviderDeleteTaskState;
use aether_ai_serving::PoolMemberScoreRules;
use serde_json::json;
use std::collections::BTreeMap;
@@ -39,6 +40,10 @@ pub(crate) struct AdminProviderPoolConfig {
pub(crate) health_policy_enabled: bool,
pub(crate) probing_enabled: bool,
pub(crate) probing_interval_minutes: u64,
pub(crate) probe_concurrency: u64,
pub(crate) score_top_n: u64,
pub(crate) score_fallback_scan_limit: u64,
pub(crate) score_rules: PoolMemberScoreRules,
pub(crate) stream_timeout_threshold: u64,
pub(crate) stream_timeout_window_seconds: u64,
pub(crate) stream_timeout_cooldown_seconds: u64,