mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
Implement generic pool member scoring and probing
This commit is contained in:
@@ -270,6 +270,9 @@ 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,
|
||||
stream_timeout_threshold: 3,
|
||||
stream_timeout_window_seconds: 1800,
|
||||
stream_timeout_cooldown_seconds: 300,
|
||||
@@ -333,6 +336,24 @@ 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),
|
||||
stream_timeout_threshold: pool_advanced
|
||||
.get("stream_timeout_threshold")
|
||||
.and_then(json_u64)
|
||||
@@ -405,6 +426,9 @@ 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,
|
||||
"stream_timeout_threshold": 4,
|
||||
"stream_timeout_window_seconds": 900,
|
||||
"stream_timeout_cooldown_seconds": 180
|
||||
@@ -424,6 +448,9 @@ 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.stream_timeout_threshold, 4);
|
||||
assert_eq!(config.stream_timeout_window_seconds, 900);
|
||||
assert_eq!(config.stream_timeout_cooldown_seconds, 180);
|
||||
|
||||
@@ -615,6 +615,9 @@ 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,
|
||||
stream_timeout_threshold: 3,
|
||||
stream_timeout_window_seconds: 1800,
|
||||
stream_timeout_cooldown_seconds: 300,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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| aether_ai_formats::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)
|
||||
}
|
||||
@@ -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")
|
||||
|
||||
@@ -39,6 +39,9 @@ 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) stream_timeout_threshold: u64,
|
||||
pub(crate) stream_timeout_window_seconds: u64,
|
||||
pub(crate) stream_timeout_cooldown_seconds: u64,
|
||||
|
||||
Reference in New Issue
Block a user