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

Improve pool score probing rules

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

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

23
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/create.rs
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);

23
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs
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);

164
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
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" => {
@@ -273,6 +391,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
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,
@@ -281,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),
@@ -354,6 +474,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.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)
@@ -429,6 +550,21 @@ mod tests {
"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
@@ -451,6 +587,13 @@ mod tests {
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);
@@ -489,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!({

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

@@ -618,6 +618,7 @@ mod tests {
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,

2
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
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;
@@ -42,6 +43,7 @@ pub(crate) struct AdminProviderPoolConfig {
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,

7
apps/aether-gateway/src/maintenance/mod.rs
View File

@@ -4,9 +4,10 @@ mod tests;
pub(crate) use runtime::{
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
inspect_proxy_upgrade_rollout, list_admin_cleanup_run_records,
perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once,
rebuild_admin_stats_once, record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
ensure_provider_key_pool_scores_for_keys, inspect_proxy_upgrade_rollout,
list_admin_cleanup_run_records, perform_oauth_token_refresh_once,
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,
run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,

5
apps/aether-gateway/src/maintenance/runtime.rs
View File

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

53
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
View File

@@ -21,6 +21,8 @@ use crate::admin_api::{
};
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_DEFAULT_SCAN_INTERVAL_SECONDS: u64 = 60;
const POOL_QUOTA_PROBE_MIN_SCAN_INTERVAL_SECONDS: u64 = 15;
@@ -495,7 +497,9 @@ async fn record_score_probe_results_from_payload(
key_id,
attempted_at,
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!({
"last_probe": {
"source": "pool_quota_probe",
@@ -520,7 +524,7 @@ async fn record_score_probe_results_from_payload(
key_id,
attempted_at,
false,
None,
Some(PoolMemberHardState::Cooldown),
serde_json::json!({
"last_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)| {
let pool_config = admin_provider_pool_config(&provider)?;
if pool_config.probing_enabled {
Some((
provider,
provider_type,
pool_config.probing_interval_minutes,
pool_config.probe_concurrency.clamp(1, 64) as usize,
))
Some((provider, provider_type, pool_config))
} else {
None
}
@@ -668,7 +667,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
let provider_ids = providers
.iter()
.map(|(provider, _, _, _)| provider.id.clone())
.map(|(provider, _, _)| provider.id.clone())
.collect::<Vec<_>>();
let mut endpoints_by_provider = BTreeMap::<String, Vec<StoredProviderCatalogEndpoint>>::new();
for endpoint in state
@@ -688,7 +687,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
..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
.remove(&provider.id)
.unwrap_or_default();
@@ -702,6 +701,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
continue;
};
let interval_minutes = pool_config.probing_interval_minutes;
let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60);
let keys = select_keys_for_provider(
state,
@@ -722,11 +722,44 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
summary.selected_keys += selected_count;
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 {
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 probe_concurrency = pool_config.probe_concurrency.clamp(1, 64) as usize;
let probe_concurrency = probe_concurrency.min(config.global_concurrency).max(1);
let probe_results = stream::iter(keys.into_iter().map(|key| {
let key_id = key.id.clone();

146
apps/aether-gateway/src/maintenance/runtime/pool_score_rebuild.rs
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::pool_scores::GetPoolMemberScoresByIdsQuery;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use tracing::{debug, info, warn};
use crate::admin_api::admin_provider_pool_config;
use crate::ai_serving::build_provider_key_pool_score_upsert;
use crate::handlers::shared::provider_pool::AdminProviderPoolConfig;
use crate::{AppState, GatewayError};
const POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS: u64 = 300;
@@ -130,6 +134,136 @@ struct ProviderScoreBuildItem {
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(
state: &AppState,
config: PoolScoreRebuildWorkerConfig,
@@ -145,16 +279,18 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.list_provider_catalog_providers(true)
.await?
.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<_>>();
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() {
return Ok(PoolScoreRebuildRunSummary::empty());
}
let provider_ids = providers
.iter()
.map(|provider| provider.id.clone())
.map(|(provider, _)| provider.id.clone())
.collect::<Vec<_>>();
let mut endpoints_by_provider = BTreeMap::new();
for endpoint in state
@@ -196,7 +332,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
break;
}
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
.remove(&provider.id)
.unwrap_or_default();
@@ -252,6 +388,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
Some(model.id.as_str()),
None,
now,
pool_config.score_rules,
);
build_items.push(ProviderScoreBuildItem {
endpoint_index,
@@ -306,6 +443,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
Some(model.id.as_str()),
existing,
now,
pool_config.score_rules,
);
if state
.data