Improve pool score probing rules

This commit is contained in:
fawney19
2026-05-12 09:08:10 +08:00
parent c578689356
commit fad28eee2c
18 changed files with 959 additions and 43 deletions

View File

@@ -1,4 +1,6 @@
use aether_ai_serving::{score_pool_member, PoolMemberScoreInput, POOL_SCORE_VERSION}; use aether_ai_serving::{
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreRules, POOL_SCORE_VERSION,
};
use aether_data_contracts::repository::pool_scores::{ use aether_data_contracts::repository::pool_scores::{
PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore, PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore,
POOL_SCORE_SCOPE_KIND_MODEL, POOL_SCORE_SCOPE_KIND_MODEL,
@@ -15,6 +17,7 @@ pub(crate) fn build_provider_key_pool_score_upsert(
model_id: Option<&str>, model_id: Option<&str>,
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>, existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
now_unix_secs: u64, now_unix_secs: u64,
score_rules: PoolMemberScoreRules,
) -> UpsertPoolMemberScore { ) -> UpsertPoolMemberScore {
let identity = PoolMemberIdentity::provider_api_key(key.provider_id.clone(), key.id.clone()); 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 scope = provider_key_pool_score_scope(api_format, model_id);
@@ -26,7 +29,7 @@ pub(crate) fn build_provider_key_pool_score_upsert(
existing, existing,
now_unix_secs, now_unix_secs,
); );
let output = score_pool_member(&input); let output = score_pool_member_with_rules(&input, score_rules);
UpsertPoolMemberScore { UpsertPoolMemberScore {
id: provider_key_pool_score_id(&identity, &scope), id: provider_key_pool_score_id(&identity, &scope),
identity, identity,
@@ -139,6 +142,7 @@ fn provider_key_score_input(
total_cost_usd: key.total_cost_usd, total_cost_usd: key.total_cost_usd,
last_used_at: key.last_used_at_unix_secs, last_used_at: key.last_used_at_unix_secs,
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at), last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
probe_failure_count: existing.map(|score| score.probe_failure_count).unwrap_or(0),
probe_status: existing probe_status: existing
.map(|score| score.probe_status) .map(|score| score.probe_status)
.unwrap_or(PoolMemberProbeStatus::Never), .unwrap_or(PoolMemberProbeStatus::Never),

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::paths::admin_provider_id_for_keys;
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest; use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; 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::provider_key_auth::provider_key_effective_api_formats;
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError}; use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
use axum::{ use axum::{
@@ -98,6 +100,27 @@ pub(super) async fn maybe_handle(
let endpoints = state let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id)) .list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?; .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 = let api_formats =
provider_key_effective_api_formats(&created, &provider.provider_type, &endpoints); 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::paths::admin_update_key_id;
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch; use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; 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::provider_key_auth::provider_key_effective_api_formats;
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError}; use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
use axum::{ use axum::{
@@ -121,6 +123,27 @@ pub(super) async fn maybe_handle(
let endpoints = state let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id)) .list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?; .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 = let api_formats =
provider_key_effective_api_formats(&updated, &provider.provider_type, &endpoints); provider_key_effective_api_formats(&updated, &provider.provider_type, &endpoints);

View File

@@ -1,6 +1,7 @@
use crate::handlers::admin::provider::shared::support::{ use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule, AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
}; };
use aether_ai_serving::{PoolMemberScoreRules, PoolMemberScoreWeights};
use serde_json::{Map, Value}; use serde_json::{Map, Value};
const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[ 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())) .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> { fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
match preset { match preset {
"free_first" | "team_first" | "plus_first" | "pro_first" => { "free_first" | "team_first" | "plus_first" | "pro_first" => {
@@ -273,6 +391,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
probe_concurrency: 4, probe_concurrency: 4,
score_top_n: 128, score_top_n: 128,
score_fallback_scan_limit: 1024, score_fallback_scan_limit: 1024,
score_rules: PoolMemberScoreRules::default(),
stream_timeout_threshold: 3, stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800, stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300, stream_timeout_cooldown_seconds: 300,
@@ -281,6 +400,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
let scheduling_presets = parse_pool_scheduling_presets(pool_advanced); let scheduling_presets = parse_pool_scheduling_presets(pool_advanced);
let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced); let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced);
let score_rules = parse_pool_score_rules(pool_advanced);
Some(AdminProviderPoolConfig { Some(AdminProviderPoolConfig {
lru_enabled: admin_provider_pool_lru_enabled(&scheduling_presets), lru_enabled: admin_provider_pool_lru_enabled(&scheduling_presets),
@@ -354,6 +474,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.filter(|value| *value > 0) .filter(|value| *value > 0)
.map(|value| value.min(50_000)) .map(|value| value.min(50_000))
.unwrap_or(1024), .unwrap_or(1024),
score_rules,
stream_timeout_threshold: pool_advanced stream_timeout_threshold: pool_advanced
.get("stream_timeout_threshold") .get("stream_timeout_threshold")
.and_then(json_u64) .and_then(json_u64)
@@ -429,6 +550,21 @@ mod tests {
"probe_concurrency": 6, "probe_concurrency": 6,
"score_top_n": 256, "score_top_n": 256,
"score_fallback_scan_limit": 2048, "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_threshold": 4,
"stream_timeout_window_seconds": 900, "stream_timeout_window_seconds": 900,
"stream_timeout_cooldown_seconds": 180 "stream_timeout_cooldown_seconds": 180
@@ -451,6 +587,13 @@ mod tests {
assert_eq!(config.probe_concurrency, 6); assert_eq!(config.probe_concurrency, 6);
assert_eq!(config.score_top_n, 256); assert_eq!(config.score_top_n, 256);
assert_eq!(config.score_fallback_scan_limit, 2048); 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_threshold, 4);
assert_eq!(config.stream_timeout_window_seconds, 900); assert_eq!(config.stream_timeout_window_seconds, 900);
assert_eq!(config.stream_timeout_cooldown_seconds, 180); assert_eq!(config.stream_timeout_cooldown_seconds, 180);
@@ -489,6 +632,27 @@ mod tests {
assert_eq!(config.sticky_session_ttl_seconds, 0); 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] #[test]
fn parses_pool_config_from_generic_config_value() { fn parses_pool_config_from_generic_config_value() {
let config = admin_provider_pool_config_from_config_value(Some(&json!({ let config = admin_provider_pool_config_from_config_value(Some(&json!({

View File

@@ -618,6 +618,7 @@ mod tests {
probe_concurrency: 4, probe_concurrency: 4,
score_top_n: 128, score_top_n: 128,
score_fallback_scan_limit: 1024, score_fallback_scan_limit: 1024,
score_rules: aether_ai_serving::PoolMemberScoreRules::default(),
stream_timeout_threshold: 3, stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800, stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300, stream_timeout_cooldown_seconds: 300,

View File

@@ -1,5 +1,6 @@
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use crate::LocalProviderDeleteTaskState; use crate::LocalProviderDeleteTaskState;
use aether_ai_serving::PoolMemberScoreRules;
use serde_json::json; use serde_json::json;
use std::collections::BTreeMap; use std::collections::BTreeMap;
@@ -42,6 +43,7 @@ pub(crate) struct AdminProviderPoolConfig {
pub(crate) probe_concurrency: u64, pub(crate) probe_concurrency: u64,
pub(crate) score_top_n: u64, pub(crate) score_top_n: u64,
pub(crate) score_fallback_scan_limit: u64, pub(crate) score_fallback_scan_limit: u64,
pub(crate) score_rules: PoolMemberScoreRules,
pub(crate) stream_timeout_threshold: u64, pub(crate) stream_timeout_threshold: u64,
pub(crate) stream_timeout_window_seconds: u64, pub(crate) stream_timeout_window_seconds: u64,
pub(crate) stream_timeout_cooldown_seconds: u64, pub(crate) stream_timeout_cooldown_seconds: u64,

View File

@@ -4,9 +4,10 @@ mod tests;
pub(crate) use runtime::{ pub(crate) use runtime::{
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts, cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
inspect_proxy_upgrade_rollout, list_admin_cleanup_run_records, ensure_provider_key_pool_scores_for_keys, inspect_proxy_upgrade_rollout,
perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once, list_admin_cleanup_run_records, perform_oauth_token_refresh_once,
rebuild_admin_stats_once, record_completed_cleanup_run, record_proxy_upgrade_traffic_success, perform_pool_quota_probe_once, perform_provider_checkin_once, rebuild_admin_stats_once,
record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node, restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker, run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker, spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,

View File

@@ -70,8 +70,9 @@ pub(crate) use pool_quota_probe::{
PoolQuotaProbeWorkerConfig, PoolQuotaProbeWorkerConfig,
}; };
pub(crate) use pool_score_rebuild::{ pub(crate) use pool_score_rebuild::{
perform_pool_score_rebuild_once, perform_pool_score_rebuild_once_with_config, ensure_provider_key_pool_scores_for_keys, perform_pool_score_rebuild_once,
spawn_pool_score_rebuild_worker, PoolScoreRebuildRunSummary, PoolScoreRebuildWorkerConfig, perform_pool_score_rebuild_once_with_config, spawn_pool_score_rebuild_worker,
PoolScoreRebuildRunSummary, PoolScoreRebuildWorkerConfig,
}; };
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary}; pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
use proxy_node_metrics_cleanup::*; use proxy_node_metrics_cleanup::*;

View File

@@ -21,6 +21,8 @@ use crate::admin_api::{
}; };
use crate::{AppState, GatewayError}; use crate::{AppState, GatewayError};
use super::pool_score_rebuild::ensure_provider_key_pool_scores_for_keys;
const POOL_QUOTA_PROBE_REDIS_PREFIX: &str = "ap:quota_probe:last"; const POOL_QUOTA_PROBE_REDIS_PREFIX: &str = "ap:quota_probe:last";
const POOL_QUOTA_PROBE_DEFAULT_SCAN_INTERVAL_SECONDS: u64 = 60; const POOL_QUOTA_PROBE_DEFAULT_SCAN_INTERVAL_SECONDS: u64 = 60;
const POOL_QUOTA_PROBE_MIN_SCAN_INTERVAL_SECONDS: u64 = 15; const POOL_QUOTA_PROBE_MIN_SCAN_INTERVAL_SECONDS: u64 = 15;
@@ -495,7 +497,9 @@ async fn record_score_probe_results_from_payload(
key_id, key_id,
attempted_at, attempted_at,
probe_result_succeeded(item), probe_result_succeeded(item),
probe_result_hard_state(item), probe_result_hard_state(item).or_else(|| {
(!probe_result_succeeded(item)).then_some(PoolMemberHardState::Cooldown)
}),
serde_json::json!({ serde_json::json!({
"last_probe": { "last_probe": {
"source": "pool_quota_probe", "source": "pool_quota_probe",
@@ -520,7 +524,7 @@ async fn record_score_probe_results_from_payload(
key_id, key_id,
attempted_at, attempted_at,
false, false,
None, Some(PoolMemberHardState::Cooldown),
serde_json::json!({ serde_json::json!({
"last_probe": { "last_probe": {
"source": "pool_quota_probe", "source": "pool_quota_probe",
@@ -650,12 +654,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
.filter_map(|(provider, provider_type)| { .filter_map(|(provider, provider_type)| {
let pool_config = admin_provider_pool_config(&provider)?; let pool_config = admin_provider_pool_config(&provider)?;
if pool_config.probing_enabled { if pool_config.probing_enabled {
Some(( Some((provider, provider_type, pool_config))
provider,
provider_type,
pool_config.probing_interval_minutes,
pool_config.probe_concurrency.clamp(1, 64) as usize,
))
} else { } else {
None None
} }
@@ -668,7 +667,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
let provider_ids = providers let provider_ids = providers
.iter() .iter()
.map(|(provider, _, _, _)| provider.id.clone()) .map(|(provider, _, _)| provider.id.clone())
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let mut endpoints_by_provider = BTreeMap::<String, Vec<StoredProviderCatalogEndpoint>>::new(); let mut endpoints_by_provider = BTreeMap::<String, Vec<StoredProviderCatalogEndpoint>>::new();
for endpoint in state for endpoint in state
@@ -688,7 +687,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
..PoolQuotaProbeRunSummary::empty() ..PoolQuotaProbeRunSummary::empty()
}; };
for (provider, provider_type, interval_minutes, probe_concurrency) in providers { for (provider, provider_type, pool_config) in providers {
let endpoints = endpoints_by_provider let endpoints = endpoints_by_provider
.remove(&provider.id) .remove(&provider.id)
.unwrap_or_default(); .unwrap_or_default();
@@ -702,6 +701,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
continue; continue;
}; };
let interval_minutes = pool_config.probing_interval_minutes;
let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60); let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60);
let keys = select_keys_for_provider( let keys = select_keys_for_provider(
state, state,
@@ -722,11 +722,44 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
summary.selected_keys += selected_count; summary.selected_keys += selected_count;
let selected_key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>(); let selected_key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize)
.min(50_000)
.max(selected_count.min(50_000));
match ensure_provider_key_pool_scores_for_keys(
state,
&provider,
&pool_config,
&endpoints,
&keys,
now_ts,
score_ensure_budget,
)
.await
{
Ok(upserted) if upserted > 0 => {
debug!(
provider_id = %provider.id,
key_count = selected_count,
scores_upserted = upserted,
"gateway pool quota probe: ensured score rows for selected probe keys"
);
}
Ok(_) => {}
Err(err) => {
warn!(
provider_id = %provider.id,
key_count = selected_count,
error = ?err,
"gateway pool quota probe: failed to ensure score rows for selected probe keys"
);
}
}
for key_id in &selected_key_ids { for key_id in &selected_key_ids {
record_score_probe_in_progress_for_key(state, &provider.id, key_id, now_ts).await; record_score_probe_in_progress_for_key(state, &provider.id, key_id, now_ts).await;
} }
let provider_short_id = provider.id.chars().take(8).collect::<String>(); let provider_short_id = provider.id.chars().take(8).collect::<String>();
let probe_concurrency = pool_config.probe_concurrency.clamp(1, 64) as usize;
let probe_concurrency = probe_concurrency.min(config.global_concurrency).max(1); let probe_concurrency = probe_concurrency.min(config.global_concurrency).max(1);
let probe_results = stream::iter(keys.into_iter().map(|key| { let probe_results = stream::iter(keys.into_iter().map(|key| {
let key_id = key.id.clone(); let key_id = key.id.clone();

View File

@@ -3,10 +3,14 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::global_models::AdminProviderModelListQuery; use aether_data_contracts::repository::global_models::AdminProviderModelListQuery;
use aether_data_contracts::repository::pool_scores::GetPoolMemberScoresByIdsQuery; use aether_data_contracts::repository::pool_scores::GetPoolMemberScoresByIdsQuery;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
use crate::admin_api::admin_provider_pool_config; use crate::admin_api::admin_provider_pool_config;
use crate::ai_serving::build_provider_key_pool_score_upsert; use crate::ai_serving::build_provider_key_pool_score_upsert;
use crate::handlers::shared::provider_pool::AdminProviderPoolConfig;
use crate::{AppState, GatewayError}; use crate::{AppState, GatewayError};
const POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS: u64 = 300; const POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS: u64 = 300;
@@ -130,6 +134,136 @@ struct ProviderScoreBuildItem {
score_id: String, score_id: String,
} }
pub(crate) async fn ensure_provider_key_pool_scores_for_keys(
state: &AppState,
provider: &StoredProviderCatalogProvider,
pool_config: &AdminProviderPoolConfig,
endpoints: &[StoredProviderCatalogEndpoint],
keys: &[StoredProviderCatalogKey],
now_unix_secs: u64,
max_upserts: usize,
) -> Result<usize, GatewayError> {
if max_upserts == 0
|| keys.is_empty()
|| !state.data.has_pool_score_reader()
|| !state.data.has_pool_score_writer()
{
return Ok(0);
}
let endpoints = endpoints
.iter()
.filter(|endpoint| endpoint.is_active && !endpoint.api_format.trim().is_empty())
.collect::<Vec<_>>();
let keys = keys
.iter()
.filter(|key| key.is_active && key.provider_id == provider.id)
.collect::<Vec<_>>();
if endpoints.is_empty() || keys.is_empty() {
return Ok(0);
}
let models = state
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: provider.id.clone(),
is_active: Some(true),
offset: 0,
limit: 10_000,
})
.await?
.into_iter()
.filter(|model| model.is_available)
.collect::<Vec<_>>();
if models.is_empty() {
return Ok(0);
}
let max_items = endpoints
.len()
.saturating_mul(models.len())
.saturating_mul(keys.len())
.min(max_upserts);
let mut build_items = Vec::with_capacity(max_items);
'outer: for (endpoint_index, endpoint) in endpoints.iter().enumerate() {
for (model_index, model) in models.iter().enumerate() {
for (key_index, key) in keys.iter().enumerate() {
let draft = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
endpoint.api_format.trim(),
Some(model.id.as_str()),
None,
now_unix_secs,
pool_config.score_rules,
);
build_items.push(ProviderScoreBuildItem {
endpoint_index,
model_index,
key_index,
score_id: draft.id,
});
if build_items.len() >= max_upserts {
break 'outer;
}
}
}
}
if build_items.is_empty() {
return Ok(0);
}
let existing_score_ids = state
.data
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: build_items
.iter()
.map(|item| item.score_id.clone())
.collect(),
})
.await
.unwrap_or_else(|err| {
debug!(
provider_id = %provider.id,
error = ?err,
"gateway pool score ensure: failed to read existing scores by id"
);
Vec::new()
})
.into_iter()
.map(|score| score.id)
.collect::<std::collections::BTreeSet<_>>();
let mut upserted = 0usize;
for item in &build_items {
if existing_score_ids.contains(&item.score_id) {
continue;
}
let endpoint = endpoints[item.endpoint_index];
let model = &models[item.model_index];
let key = keys[item.key_index];
let upsert = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
endpoint.api_format.trim(),
Some(model.id.as_str()),
None,
now_unix_secs,
pool_config.score_rules,
);
if state
.data
.upsert_pool_member_score(upsert)
.await
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?
.is_some()
{
upserted = upserted.saturating_add(1);
}
}
Ok(upserted)
}
pub(crate) async fn perform_pool_score_rebuild_once_with_config( pub(crate) async fn perform_pool_score_rebuild_once_with_config(
state: &AppState, state: &AppState,
config: PoolScoreRebuildWorkerConfig, config: PoolScoreRebuildWorkerConfig,
@@ -145,16 +279,18 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.list_provider_catalog_providers(true) .list_provider_catalog_providers(true)
.await? .await?
.into_iter() .into_iter()
.filter(|provider| admin_provider_pool_config(provider).is_some()) .filter_map(|provider| {
admin_provider_pool_config(&provider).map(|config| (provider, config))
})
.collect::<Vec<_>>(); .collect::<Vec<_>>();
providers.sort_by(|left, right| left.id.cmp(&right.id)); providers.sort_by(|left, right| left.0.id.cmp(&right.0.id));
if providers.is_empty() { if providers.is_empty() {
return Ok(PoolScoreRebuildRunSummary::empty()); return Ok(PoolScoreRebuildRunSummary::empty());
} }
let provider_ids = providers let provider_ids = providers
.iter() .iter()
.map(|provider| provider.id.clone()) .map(|(provider, _)| provider.id.clone())
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let mut endpoints_by_provider = BTreeMap::new(); let mut endpoints_by_provider = BTreeMap::new();
for endpoint in state for endpoint in state
@@ -196,7 +332,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
break; break;
} }
last_provider_index = Some(provider_index); last_provider_index = Some(provider_index);
let provider = providers[provider_index].clone(); let (provider, pool_config) = providers[provider_index].clone();
let endpoints = endpoints_by_provider let endpoints = endpoints_by_provider
.remove(&provider.id) .remove(&provider.id)
.unwrap_or_default(); .unwrap_or_default();
@@ -252,6 +388,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
Some(model.id.as_str()), Some(model.id.as_str()),
None, None,
now, now,
pool_config.score_rules,
); );
build_items.push(ProviderScoreBuildItem { build_items.push(ProviderScoreBuildItem {
endpoint_index, endpoint_index,
@@ -306,6 +443,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
Some(model.id.as_str()), Some(model.id.as_str()),
existing, existing,
now, now,
pool_config.score_rules,
); );
if state if state
.data .data

View File

@@ -101,8 +101,11 @@ pub use pool_scheduler::{
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON, AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
}; };
pub use pool_scores::{ pub use pool_scores::{
probe_freshness_score, score_pool_member, PoolMemberScoreInput, PoolMemberScoreOutput, probe_freshness_score, probe_freshness_score_with_ttl, score_pool_member,
POOL_SCORE_VERSION, PROBE_FRESHNESS_TTL_SECONDS, score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreOutput,
PoolMemberScoreRules, PoolMemberScoreWeights, POOL_SCORE_VERSION,
PROBE_FAILURE_COOLDOWN_THRESHOLD, PROBE_FAILURE_PENALTY, PROBE_FRESHNESS_TTL_SECONDS,
REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP,
}; };
pub use ranking_metadata::append_ai_ranking_metadata_to_object; pub use ranking_metadata::append_ai_ranking_metadata_to_object;
pub use report_context::{ pub use report_context::{

View File

@@ -5,6 +5,127 @@ use serde_json::{json, Value};
pub const POOL_SCORE_VERSION: u64 = 1; pub const POOL_SCORE_VERSION: u64 = 1;
pub const PROBE_FRESHNESS_TTL_SECONDS: u64 = 30 * 60; pub const PROBE_FRESHNESS_TTL_SECONDS: u64 = 30 * 60;
pub const UNSCHEDULABLE_SCORE_CAP: f64 = 0.05;
pub const PROBE_FAILURE_PENALTY: f64 = 0.05;
pub const REQUEST_FAILURE_PENALTY: f64 = 0.005;
pub const PROBE_FAILURE_COOLDOWN_THRESHOLD: u64 = 3;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PoolMemberScoreWeights {
pub manual_priority: f64,
pub health: f64,
pub probe_freshness: f64,
pub quota_remaining: f64,
pub latency: f64,
pub cost_lru: f64,
}
impl Default for PoolMemberScoreWeights {
fn default() -> Self {
Self {
manual_priority: 0.30,
health: 0.20,
probe_freshness: 0.15,
quota_remaining: 0.15,
latency: 0.10,
cost_lru: 0.10,
}
}
}
impl PoolMemberScoreWeights {
pub fn normalized(self) -> Self {
let sanitized = Self {
manual_priority: finite_non_negative(self.manual_priority),
health: finite_non_negative(self.health),
probe_freshness: finite_non_negative(self.probe_freshness),
quota_remaining: finite_non_negative(self.quota_remaining),
latency: finite_non_negative(self.latency),
cost_lru: finite_non_negative(self.cost_lru),
};
let total = sanitized.manual_priority
+ sanitized.health
+ sanitized.probe_freshness
+ sanitized.quota_remaining
+ sanitized.latency
+ sanitized.cost_lru;
if total <= f64::EPSILON {
return sanitized;
}
Self {
manual_priority: sanitized.manual_priority / total,
health: sanitized.health / total,
probe_freshness: sanitized.probe_freshness / total,
quota_remaining: sanitized.quota_remaining / total,
latency: sanitized.latency / total,
cost_lru: sanitized.cost_lru / total,
}
}
fn as_reason_json(self) -> Value {
json!({
"manual_priority": self.manual_priority,
"health": self.health,
"probe_freshness": self.probe_freshness,
"quota_remaining": self.quota_remaining,
"latency": self.latency,
"cost_lru": self.cost_lru
})
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PoolMemberScoreRules {
pub weights: PoolMemberScoreWeights,
pub probe_freshness_ttl_seconds: u64,
pub unschedulable_score_cap: f64,
pub probe_failure_penalty: f64,
pub request_failure_penalty: f64,
pub probe_failure_cooldown_threshold: u64,
}
impl Default for PoolMemberScoreRules {
fn default() -> Self {
Self {
weights: PoolMemberScoreWeights::default(),
probe_freshness_ttl_seconds: PROBE_FRESHNESS_TTL_SECONDS,
unschedulable_score_cap: UNSCHEDULABLE_SCORE_CAP,
probe_failure_penalty: PROBE_FAILURE_PENALTY,
request_failure_penalty: REQUEST_FAILURE_PENALTY,
probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
}
}
}
impl PoolMemberScoreRules {
pub fn effective(self) -> Self {
let defaults = Self::default();
Self {
weights: self.weights.normalized(),
probe_freshness_ttl_seconds: if self.probe_freshness_ttl_seconds == 0 {
defaults.probe_freshness_ttl_seconds
} else {
self.probe_freshness_ttl_seconds
},
unschedulable_score_cap: if self.unschedulable_score_cap.is_finite() {
self.unschedulable_score_cap.clamp(0.0, 1.0)
} else {
defaults.unschedulable_score_cap
},
probe_failure_penalty: if self.probe_failure_penalty.is_finite() {
self.probe_failure_penalty.clamp(0.0, 1.0)
} else {
defaults.probe_failure_penalty
},
request_failure_penalty: if self.request_failure_penalty.is_finite() {
self.request_failure_penalty.clamp(0.0, 1.0)
} else {
defaults.request_failure_penalty
},
probe_failure_cooldown_threshold: self.probe_failure_cooldown_threshold,
}
}
}
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct PoolMemberScoreInput { pub struct PoolMemberScoreInput {
@@ -25,6 +146,7 @@ pub struct PoolMemberScoreInput {
pub total_cost_usd: f64, pub total_cost_usd: f64,
pub last_used_at: Option<u64>, pub last_used_at: Option<u64>,
pub last_probe_success_at: Option<u64>, pub last_probe_success_at: Option<u64>,
pub probe_failure_count: u64,
pub probe_status: PoolMemberProbeStatus, pub probe_status: PoolMemberProbeStatus,
pub now_unix_secs: u64, pub now_unix_secs: u64,
} }
@@ -37,13 +159,23 @@ pub struct PoolMemberScoreOutput {
} }
pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput { pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput {
let hard_state = derive_hard_state(input); score_pool_member_with_rules(input, PoolMemberScoreRules::default())
}
pub fn score_pool_member_with_rules(
input: &PoolMemberScoreInput,
rules: PoolMemberScoreRules,
) -> PoolMemberScoreOutput {
let rules = rules.effective();
let weights = rules.weights;
let hard_state = derive_hard_state(input, &rules);
let manual_priority = manual_priority_score(input.internal_priority); let manual_priority = manual_priority_score(input.internal_priority);
let health = input.health_score.unwrap_or(0.5).clamp(0.0, 1.0); let health = input.health_score.unwrap_or(0.5).clamp(0.0, 1.0);
let probe_freshness = probe_freshness_score( let probe_freshness = probe_freshness_score_with_ttl(
input.last_probe_success_at, input.last_probe_success_at,
input.probe_status, input.probe_status,
input.now_unix_secs, input.now_unix_secs,
rules.probe_freshness_ttl_seconds,
); );
let quota_remaining = input let quota_remaining = input
.quota_usage_ratio .quota_usage_ratio
@@ -52,14 +184,20 @@ pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput
let latency = latency_score(input.success_count, input.total_response_time_ms); let latency = latency_score(input.success_count, input.total_response_time_ms);
let cost_lru = cost_lru_score(input.total_cost_usd, input.total_tokens, input.last_used_at); let cost_lru = cost_lru_score(input.total_cost_usd, input.total_tokens, input.last_used_at);
let mut score = manual_priority * 0.30 let weighted_score = manual_priority * weights.manual_priority
+ health * 0.20 + health * weights.health
+ probe_freshness * 0.15 + probe_freshness * weights.probe_freshness
+ quota_remaining * 0.15 + quota_remaining * weights.quota_remaining
+ latency * 0.10 + latency * weights.latency
+ cost_lru * 0.10; + cost_lru * weights.cost_lru;
let probe_failure_penalty =
(input.probe_failure_count.min(10) as f64 * rules.probe_failure_penalty).min(0.5);
let request_failure_penalty =
(input.error_count.min(20) as f64 * rules.request_failure_penalty).min(0.5);
let total_penalty = (probe_failure_penalty + request_failure_penalty).min(1.0);
let mut score = weighted_score - total_penalty;
if !hard_state.schedulable() { if !hard_state.schedulable() {
score = score.min(0.05); score = score.min(rules.unschedulable_score_cap);
} }
score = score.clamp(0.0, 1.0); score = score.clamp(0.0, 1.0);
@@ -67,14 +205,7 @@ pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput
score, score,
hard_state, hard_state,
score_reason: json!({ score_reason: json!({
"weights": { "weights": weights.as_reason_json(),
"manual_priority": 0.30,
"health": 0.20,
"probe_freshness": 0.15,
"quota_remaining": 0.15,
"latency": 0.10,
"cost_lru": 0.10
},
"factors": { "factors": {
"manual_priority": manual_priority, "manual_priority": manual_priority,
"health": health, "health": health,
@@ -83,13 +214,36 @@ pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput
"latency": latency, "latency": latency,
"cost_lru": cost_lru "cost_lru": cost_lru
}, },
"rules": {
"probe_freshness_ttl_seconds": rules.probe_freshness_ttl_seconds,
"unschedulable_score_cap": rules.unschedulable_score_cap,
"probe_failure_penalty": rules.probe_failure_penalty,
"request_failure_penalty": rules.request_failure_penalty,
"probe_failure_cooldown_threshold": rules.probe_failure_cooldown_threshold
},
"penalties": {
"probe_failure": probe_failure_penalty,
"request_failure": request_failure_penalty,
"total": total_penalty
},
"hard_state": hard_state.as_database(), "hard_state": hard_state.as_database(),
"score_version": POOL_SCORE_VERSION "score_version": POOL_SCORE_VERSION
}), }),
} }
} }
fn derive_hard_state(input: &PoolMemberScoreInput) -> PoolMemberHardState { fn finite_non_negative(value: f64) -> f64 {
if value.is_finite() {
value.max(0.0)
} else {
0.0
}
}
fn derive_hard_state(
input: &PoolMemberScoreInput,
rules: &PoolMemberScoreRules,
) -> PoolMemberHardState {
if !input.is_active { if !input.is_active {
return PoolMemberHardState::Inactive; return PoolMemberHardState::Inactive;
} }
@@ -109,6 +263,12 @@ fn derive_hard_state(input: &PoolMemberScoreInput) -> PoolMemberHardState {
if input.circuit_open { if input.circuit_open {
return PoolMemberHardState::Cooldown; return PoolMemberHardState::Cooldown;
} }
if input.probe_status == PoolMemberProbeStatus::Failed
&& rules.probe_failure_cooldown_threshold > 0
&& input.probe_failure_count >= rules.probe_failure_cooldown_threshold
{
return PoolMemberHardState::Cooldown;
}
if input.health_score.is_some() || input.probe_status == PoolMemberProbeStatus::Ok { if input.health_score.is_some() || input.probe_status == PoolMemberProbeStatus::Ok {
PoolMemberHardState::Available PoolMemberHardState::Available
} else { } else {
@@ -124,6 +284,20 @@ pub fn probe_freshness_score(
last_probe_success_at: Option<u64>, last_probe_success_at: Option<u64>,
probe_status: PoolMemberProbeStatus, probe_status: PoolMemberProbeStatus,
now_unix_secs: u64, now_unix_secs: u64,
) -> f64 {
probe_freshness_score_with_ttl(
last_probe_success_at,
probe_status,
now_unix_secs,
PROBE_FRESHNESS_TTL_SECONDS,
)
}
pub fn probe_freshness_score_with_ttl(
last_probe_success_at: Option<u64>,
probe_status: PoolMemberProbeStatus,
now_unix_secs: u64,
ttl_seconds: u64,
) -> f64 { ) -> f64 {
if probe_status != PoolMemberProbeStatus::Ok { if probe_status != PoolMemberProbeStatus::Ok {
return 0.0; return 0.0;
@@ -131,11 +305,12 @@ pub fn probe_freshness_score(
let Some(success_at) = last_probe_success_at else { let Some(success_at) = last_probe_success_at else {
return 0.0; return 0.0;
}; };
let ttl_seconds = ttl_seconds.max(1);
let age = now_unix_secs.saturating_sub(success_at); let age = now_unix_secs.saturating_sub(success_at);
if age >= PROBE_FRESHNESS_TTL_SECONDS { if age >= ttl_seconds {
0.0 0.0
} else { } else {
1.0 - (age as f64 / PROBE_FRESHNESS_TTL_SECONDS as f64) 1.0 - (age as f64 / ttl_seconds as f64)
} }
} }
@@ -203,6 +378,7 @@ mod tests {
total_cost_usd: 0.01, total_cost_usd: 0.01,
last_used_at: None, last_used_at: None,
last_probe_success_at: Some(1_000), last_probe_success_at: Some(1_000),
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Ok, probe_status: PoolMemberProbeStatus::Ok,
now_unix_secs: 1_000, now_unix_secs: 1_000,
} }
@@ -234,4 +410,131 @@ mod tests {
0.0 0.0
); );
} }
#[test]
fn custom_rules_change_weights_and_probe_ttl() {
let mut input = input();
input.last_probe_success_at = Some(1_000);
input.now_unix_secs = 1_900;
let rules = PoolMemberScoreRules {
weights: PoolMemberScoreWeights {
manual_priority: 0.0,
health: 0.0,
probe_freshness: 1.0,
quota_remaining: 0.0,
latency: 0.0,
cost_lru: 0.0,
},
probe_freshness_ttl_seconds: 1_000,
unschedulable_score_cap: 0.05,
probe_failure_penalty: 0.0,
request_failure_penalty: 0.0,
probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
};
let output = score_pool_member_with_rules(&input, rules);
assert!((output.score - 0.1).abs() < 0.000_001);
assert_eq!(
output.score_reason["rules"]["probe_freshness_ttl_seconds"],
1_000
);
assert_eq!(output.score_reason["weights"]["probe_freshness"], 1.0);
}
#[test]
fn custom_rules_normalize_weights() {
let rules = PoolMemberScoreRules {
weights: PoolMemberScoreWeights {
manual_priority: 2.0,
health: 2.0,
probe_freshness: 0.0,
quota_remaining: 0.0,
latency: -1.0,
cost_lru: f64::NAN,
},
probe_freshness_ttl_seconds: 0,
unschedulable_score_cap: f64::INFINITY,
probe_failure_penalty: f64::NAN,
request_failure_penalty: f64::INFINITY,
probe_failure_cooldown_threshold: 2,
}
.effective();
assert_eq!(rules.weights.manual_priority, 0.5);
assert_eq!(rules.weights.health, 0.5);
assert_eq!(
rules.probe_freshness_ttl_seconds,
PROBE_FRESHNESS_TTL_SECONDS
);
assert_eq!(rules.unschedulable_score_cap, UNSCHEDULABLE_SCORE_CAP);
assert_eq!(rules.probe_failure_penalty, PROBE_FAILURE_PENALTY);
assert_eq!(rules.request_failure_penalty, REQUEST_FAILURE_PENALTY);
assert_eq!(rules.probe_failure_cooldown_threshold, 2);
}
#[test]
fn custom_rules_preserve_zero_weight_total() {
let rules = PoolMemberScoreRules {
weights: PoolMemberScoreWeights {
manual_priority: 0.0,
health: 0.0,
probe_freshness: 0.0,
quota_remaining: 0.0,
latency: 0.0,
cost_lru: 0.0,
},
..PoolMemberScoreRules::default()
}
.effective();
assert_eq!(
rules.weights,
PoolMemberScoreWeights {
manual_priority: 0.0,
health: 0.0,
probe_freshness: 0.0,
quota_remaining: 0.0,
latency: 0.0,
cost_lru: 0.0,
}
);
}
#[test]
fn repeated_probe_failures_move_member_to_cooldown() {
let mut input = input();
input.probe_status = PoolMemberProbeStatus::Failed;
input.probe_failure_count = 3;
let output = score_pool_member(&input);
assert_eq!(output.hard_state, PoolMemberHardState::Cooldown);
assert!(output.score <= UNSCHEDULABLE_SCORE_CAP);
assert_eq!(
output.score_reason["rules"]["probe_failure_cooldown_threshold"],
PROBE_FAILURE_COOLDOWN_THRESHOLD
);
}
#[test]
fn probe_and_request_failures_penalize_schedulable_score() {
let mut input = input();
input.last_probe_success_at = None;
input.probe_status = PoolMemberProbeStatus::Never;
input.probe_failure_count = 1;
input.error_count = 2;
let rules = PoolMemberScoreRules {
probe_failure_penalty: 0.1,
request_failure_penalty: 0.01,
probe_failure_cooldown_threshold: 3,
..PoolMemberScoreRules::default()
};
let output = score_pool_member_with_rules(&input, rules);
assert_eq!(output.hard_state, PoolMemberHardState::Available);
assert_eq!(output.score_reason["penalties"]["probe_failure"], 0.1);
assert_eq!(output.score_reason["penalties"]["request_failure"], 0.02);
}
} }

View File

@@ -70,6 +70,7 @@ impl InMemoryPoolMemberScoreRepository {
scores.sort_by(|left, right| { scores.sort_by(|left, right| {
probe_priority(left) probe_priority(left)
.cmp(&probe_priority(right)) .cmp(&probe_priority(right))
.then_with(|| right.probe_failure_count.cmp(&left.probe_failure_count))
.then_with(|| { .then_with(|| {
left.last_probe_success_at left.last_probe_success_at
.unwrap_or(0) .unwrap_or(0)

View File

@@ -223,6 +223,7 @@ impl PoolScoreReadRepository for MysqlPoolMemberScoreRepository {
WHEN probe_status = 'stale' THEN 3 WHEN probe_status = 'stale' THEN 3
ELSE 4 ELSE 4
END ASC, END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC, COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC, COALESCE(last_scheduled_at, 0) DESC,
member_id ASC member_id ASC

View File

@@ -235,6 +235,7 @@ impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository {
WHEN probe_status = 'stale' THEN 3 WHEN probe_status = 'stale' THEN 3
ELSE 4 ELSE 4
END ASC, END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC, COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC, COALESCE(last_scheduled_at, 0) DESC,
member_id ASC member_id ASC

View File

@@ -223,6 +223,7 @@ impl PoolScoreReadRepository for SqlitePoolMemberScoreRepository {
WHEN probe_status = 'stale' THEN 3 WHEN probe_status = 'stale' THEN 3
ELSE 4 ELSE 4
END ASC, END ASC,
probe_failure_count DESC,
COALESCE(last_probe_success_at, 0) ASC, COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC, COALESCE(last_scheduled_at, 0) DESC,
member_id ASC member_id ASC

View File

@@ -521,6 +521,24 @@ export interface SchedulingPresetItem {
mode?: string | null mode?: string | null
} }
export interface PoolScoreWeights {
manual_priority?: number | null
health?: number | null
probe_freshness?: number | null
quota_remaining?: number | null
latency?: number | null
cost_lru?: number | null
}
export interface PoolScoreRules {
weights?: PoolScoreWeights | null
probe_freshness_ttl_seconds?: number | null
unschedulable_score_cap?: number | null
probe_failure_penalty?: number | null
request_failure_penalty?: number | null
probe_failure_cooldown_threshold?: number | null
}
export interface PoolAdvancedConfig { export interface PoolAdvancedConfig {
global_priority?: number | null global_priority?: number | null
sticky_session_ttl_seconds?: number | null sticky_session_ttl_seconds?: number | null
@@ -551,6 +569,7 @@ export interface PoolAdvancedConfig {
probe_concurrency?: number | null probe_concurrency?: number | null
score_top_n?: number | null score_top_n?: number | null
score_fallback_scan_limit?: number | null score_fallback_scan_limit?: number | null
score_rules?: PoolScoreRules | null
probing_enabled?: boolean probing_enabled?: boolean
probing_interval_minutes?: number | null probing_interval_minutes?: number | null
auto_remove_banned_keys?: boolean auto_remove_banned_keys?: boolean

View File

@@ -290,6 +290,161 @@
</section> </section>
</div> </div>
<section class="space-y-4 rounded-2xl border border-border/60 bg-card/70 p-4 sm:p-5">
<div class="space-y-1">
<div class="flex flex-wrap items-center gap-2">
<h3 class="text-sm font-semibold">
分数规则
</h3>
<span class="rounded-full bg-muted px-2 py-0.5 text-[11px] text-muted-foreground">
候选排序
</span>
</div>
<p class="text-xs leading-5 text-muted-foreground">
调整主动探测、健康、额度、延迟和使用成本进入号池候选排序时的权重。
</p>
</div>
<div class="grid gap-3 lg:grid-cols-3">
<div class="space-y-1.5">
<Label>优先级权重</Label>
<Input
:model-value="form.score_weight_manual_priority ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.30"
@update:model-value="(v) => form.score_weight_manual_priority = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>健康权重</Label>
<Input
:model-value="form.score_weight_health ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.20"
@update:model-value="(v) => form.score_weight_health = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>探测新鲜度权重</Label>
<Input
:model-value="form.score_weight_probe_freshness ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.15"
@update:model-value="(v) => form.score_weight_probe_freshness = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>额度剩余权重</Label>
<Input
:model-value="form.score_weight_quota_remaining ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.15"
@update:model-value="(v) => form.score_weight_quota_remaining = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>延迟权重</Label>
<Input
:model-value="form.score_weight_latency ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.10"
@update:model-value="(v) => form.score_weight_latency = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>成本/LRU 权重</Label>
<Input
:model-value="form.score_weight_cost_lru ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.10"
@update:model-value="(v) => form.score_weight_cost_lru = parseNum(v)"
/>
</div>
</div>
<div class="grid gap-3 sm:grid-cols-2 lg:grid-cols-3">
<div class="space-y-1.5">
<Label>
探测新鲜度 TTL
<span class="text-xs text-muted-foreground">(秒)</span>
</Label>
<Input
:model-value="form.probe_freshness_ttl_seconds ?? ''"
type="number"
min="1"
max="604800"
placeholder="1800"
@update:model-value="(v) => form.probe_freshness_ttl_seconds = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>探测失败惩罚</Label>
<Input
:model-value="form.probe_failure_penalty ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.05"
@update:model-value="(v) => form.probe_failure_penalty = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>请求失败惩罚</Label>
<Input
:model-value="form.request_failure_penalty ?? ''"
type="number"
min="0"
max="1"
step="0.001"
placeholder="0.005"
@update:model-value="(v) => form.request_failure_penalty = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>探测失败冷却阈值</Label>
<Input
:model-value="form.probe_failure_cooldown_threshold ?? ''"
type="number"
min="0"
max="100"
placeholder="3"
@update:model-value="(v) => form.probe_failure_cooldown_threshold = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>不可调度分数上限</Label>
<Input
:model-value="form.unschedulable_score_cap ?? ''"
type="number"
min="0"
max="1"
step="0.01"
placeholder="0.05"
@update:model-value="(v) => form.unschedulable_score_cap = parseNum(v)"
/>
</div>
</div>
</section>
<section <section
v-if="isClaudeCode" v-if="isClaudeCode"
class="space-y-4 rounded-2xl border border-border/60 bg-card/70 p-4 sm:p-5" class="space-y-4 rounded-2xl border border-border/60 bg-card/70 p-4 sm:p-5"
@@ -506,6 +661,17 @@ const form = ref({
probe_concurrency: null as number | null | undefined, probe_concurrency: null as number | null | undefined,
score_top_n: null as number | null | undefined, score_top_n: null as number | null | undefined,
score_fallback_scan_limit: null as number | null | undefined, score_fallback_scan_limit: null as number | null | undefined,
score_weight_manual_priority: null as number | null | undefined,
score_weight_health: null as number | null | undefined,
score_weight_probe_freshness: null as number | null | undefined,
score_weight_quota_remaining: null as number | null | undefined,
score_weight_latency: null as number | null | undefined,
score_weight_cost_lru: null as number | null | undefined,
probe_freshness_ttl_seconds: null as number | null | undefined,
unschedulable_score_cap: null as number | null | undefined,
probe_failure_penalty: null as number | null | undefined,
request_failure_penalty: null as number | null | undefined,
probe_failure_cooldown_threshold: null as number | null | undefined,
probing_enabled: false, probing_enabled: false,
probing_interval_minutes: null as number | null | undefined, probing_interval_minutes: null as number | null | undefined,
auto_remove_banned_keys: false, auto_remove_banned_keys: false,
@@ -571,6 +737,8 @@ watch(() => props.modelValue, (open) => {
if (!open) return if (!open) return
const cfg = props.currentConfig const cfg = props.currentConfig
const scoreRules = cfg?.score_rules
const scoreWeights = scoreRules?.weights
form.value = { form.value = {
global_priority: cfg?.global_priority ?? null, global_priority: cfg?.global_priority ?? null,
sticky_session_ttl_seconds: cfg?.sticky_session_ttl_seconds ?? null, sticky_session_ttl_seconds: cfg?.sticky_session_ttl_seconds ?? null,
@@ -584,6 +752,17 @@ watch(() => props.modelValue, (open) => {
probe_concurrency: cfg?.probe_concurrency ?? null, probe_concurrency: cfg?.probe_concurrency ?? null,
score_top_n: cfg?.score_top_n ?? null, score_top_n: cfg?.score_top_n ?? null,
score_fallback_scan_limit: cfg?.score_fallback_scan_limit ?? null, score_fallback_scan_limit: cfg?.score_fallback_scan_limit ?? null,
score_weight_manual_priority: scoreWeights?.manual_priority ?? null,
score_weight_health: scoreWeights?.health ?? null,
score_weight_probe_freshness: scoreWeights?.probe_freshness ?? null,
score_weight_quota_remaining: scoreWeights?.quota_remaining ?? null,
score_weight_latency: scoreWeights?.latency ?? null,
score_weight_cost_lru: scoreWeights?.cost_lru ?? null,
probe_freshness_ttl_seconds: scoreRules?.probe_freshness_ttl_seconds ?? null,
unschedulable_score_cap: scoreRules?.unschedulable_score_cap ?? null,
probe_failure_penalty: scoreRules?.probe_failure_penalty ?? null,
request_failure_penalty: scoreRules?.request_failure_penalty ?? null,
probe_failure_cooldown_threshold: scoreRules?.probe_failure_cooldown_threshold ?? null,
probing_enabled: cfg?.probing_enabled ?? false, probing_enabled: cfg?.probing_enabled ?? false,
probing_interval_minutes: cfg?.probing_interval_minutes ?? null, probing_interval_minutes: cfg?.probing_interval_minutes ?? null,
auto_remove_banned_keys: cfg?.auto_remove_banned_keys ?? false, auto_remove_banned_keys: cfg?.auto_remove_banned_keys ?? false,
@@ -605,6 +784,23 @@ watch(() => props.modelValue, (open) => {
async function handleSave() { async function handleSave() {
loading.value = true loading.value = true
try { try {
const scoreRules = {
...(props.currentConfig?.score_rules ?? {}),
weights: {
...(props.currentConfig?.score_rules?.weights ?? {}),
manual_priority: form.value.score_weight_manual_priority ?? undefined,
health: form.value.score_weight_health ?? undefined,
probe_freshness: form.value.score_weight_probe_freshness ?? undefined,
quota_remaining: form.value.score_weight_quota_remaining ?? undefined,
latency: form.value.score_weight_latency ?? undefined,
cost_lru: form.value.score_weight_cost_lru ?? undefined,
},
probe_freshness_ttl_seconds: form.value.probe_freshness_ttl_seconds ?? undefined,
unschedulable_score_cap: form.value.unschedulable_score_cap ?? undefined,
probe_failure_penalty: form.value.probe_failure_penalty ?? undefined,
request_failure_penalty: form.value.request_failure_penalty ?? undefined,
probe_failure_cooldown_threshold: form.value.probe_failure_cooldown_threshold ?? undefined,
}
// 合并已有配置(保留 scheduling_presets 等不在此对话框编辑的字段) // 合并已有配置(保留 scheduling_presets 等不在此对话框编辑的字段)
const poolAdvanced: Record<string, unknown> = { const poolAdvanced: Record<string, unknown> = {
...(props.currentConfig ?? {}), ...(props.currentConfig ?? {}),
@@ -620,6 +816,7 @@ async function handleSave() {
probe_concurrency: form.value.probe_concurrency ?? undefined, probe_concurrency: form.value.probe_concurrency ?? undefined,
score_top_n: form.value.score_top_n ?? undefined, score_top_n: form.value.score_top_n ?? undefined,
score_fallback_scan_limit: form.value.score_fallback_scan_limit ?? undefined, score_fallback_scan_limit: form.value.score_fallback_scan_limit ?? undefined,
score_rules: scoreRules,
probing_enabled: form.value.probing_enabled, probing_enabled: form.value.probing_enabled,
probing_interval_minutes: form.value.probing_enabled probing_interval_minutes: form.value.probing_enabled
? (form.value.probing_interval_minutes ?? undefined) ? (form.value.probing_interval_minutes ?? undefined)