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

Implement generic pool member scoring and probing

Browse Source
This commit is contained in:
fawney19
2026-05-12 01:46:22 +08:00
parent cb0ccb9cdb
commit b9e62d1667
76 changed files with 6276 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
apps/aether-gateway/src/ai_serving/mod.rs
View File

@@ -43,13 +43,13 @@ pub(crate) use self::planner::{
build_local_video_sync_plan_and_reports_for_kind,
build_openai_responses_stream_plan_from_decision,
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
build_standard_family_stream_attempt_source, build_standard_family_stream_plan_and_reports,
build_standard_family_sync_attempt_source, build_standard_family_sync_plan_and_reports,
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
set_local_openai_chat_execution_exhausted_diagnostic,
build_provider_key_pool_score_upsert, build_standard_family_stream_attempt_source,
build_standard_family_stream_plan_and_reports, build_standard_family_sync_attempt_source,
build_standard_family_sync_plan_and_reports, build_standard_stream_plan_from_decision,
build_standard_sync_plan_from_decision, extract_pool_sticky_session_token,
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
planner_is_matching_stream_request, set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
LocalExecutionAttemptSource, LocalResolvedOAuthRequestAuth, PlannerAppState,

2
apps/aether-gateway/src/ai_serving/planner/mod.rs
View File

@@ -16,6 +16,7 @@ mod materialization_policy;
mod passthrough;
mod plan_builders;
mod pool_scheduler;
pub(crate) mod pool_scores;
mod report_context;
mod route;
mod runtime_miss;
@@ -36,6 +37,7 @@ pub(crate) use self::plan_builders::{
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
pub(crate) use self::pool_scores::build_provider_key_pool_score_upsert;
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::specialized::{
build_local_gemini_files_stream_attempt_source_for_kind,

164
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
View File

@@ -9,7 +9,11 @@ use aether_ai_serving::{
};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
};
use aether_data_contracts::repository::pool_scores::{
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{Map, Value};
@@ -37,6 +41,8 @@ use crate::handlers::shared::{
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::provider_key_auth::provider_key_auth_semantics;
use super::pool_scores::provider_key_pool_score_scope;
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
const DEFAULT_POOL_KEY_PAGE_SIZE: u32 = 128;
const DEFAULT_POOL_MAX_SCANNED_KEYS: u32 = 1024;
@@ -176,6 +182,8 @@ pub(crate) struct PoolKeyCursor<'a> {
scanned_keys: u32,
page_size: u32,
max_scanned_keys: u32,
score_top_n: u32,
score_phase_loaded: bool,
skip_reason_counts: BTreeMap<&'static str, u32>,
next_pool_key_index: u32,
sticky_candidate_loaded: bool,
@@ -194,6 +202,17 @@ impl<'a> PoolKeyCursor<'a> {
request_auth_channel: Option<&str>,
) -> Self {
let pool_key_order = pool_key_candidate_order_for_group(&group);
let pool_config = pool_config_for_candidate(&group);
let score_top_n = pool_config
.as_ref()
.map(|config| config.score_top_n)
.unwrap_or(u64::from(DEFAULT_POOL_KEY_PAGE_SIZE))
.clamp(1, u64::from(u32::MAX)) as u32;
let max_scanned_keys = pool_config
.as_ref()
.map(|config| config.score_fallback_scan_limit)
.unwrap_or(u64::from(DEFAULT_POOL_MAX_SCANNED_KEYS))
.clamp(1, u64::from(u32::MAX)) as u32;
Self {
state,
group,
@@ -204,7 +223,9 @@ impl<'a> PoolKeyCursor<'a> {
next_offset: 0,
scanned_keys: 0,
page_size: DEFAULT_POOL_KEY_PAGE_SIZE,
max_scanned_keys: DEFAULT_POOL_MAX_SCANNED_KEYS,
max_scanned_keys,
score_top_n,
score_phase_loaded: false,
skip_reason_counts: BTreeMap::new(),
next_pool_key_index: 0,
sticky_candidate_loaded: false,
@@ -257,6 +278,13 @@ impl<'a> PoolKeyCursor<'a> {
}
async fn next_page_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
if !self.score_phase_loaded {
self.score_phase_loaded = true;
if let Some(score_candidates) = self.next_score_candidates().await {
return Some(score_candidates);
}
}
if self.scanned_keys >= self.max_scanned_keys {
return None;
}
@@ -306,6 +334,82 @@ impl<'a> PoolKeyCursor<'a> {
Some(self.build_page_eligible_candidates(rows).await)
}
async fn next_score_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
let scope = provider_key_pool_score_scope(
self.group.candidate.endpoint_api_format.as_str(),
Some(self.group.candidate.model_id.as_str()),
);
let query = ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: self.group.candidate.provider_id.clone(),
capability: scope.capability.clone(),
scope_kind: scope.scope_kind.clone(),
scope_id: scope.scope_id.clone(),
hard_states: vec![PoolMemberHardState::Available, PoolMemberHardState::Unknown],
probe_statuses: None,
offset: 0,
limit: self.score_top_n as usize,
};
let scores = match self.state.app().data.list_ranked_pool_members(&query).await {
Ok(scores) => scores,
Err(err) => {
warn!(
event_name = "pool_group_score_load_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
error = ?err,
"gateway pool scheduler failed to read ranked pool member scores"
);
return None;
}
};
if scores.is_empty() {
return None;
}
self.record_score_schedule_interest(&scores).await;
let key_ids = scores
.iter()
.map(|score| score.member_id.clone())
.collect::<Vec<_>>();
let rows_query = StoredPoolKeyCandidateRowsByKeyIdsQuery {
api_format: self.group.candidate.endpoint_api_format.clone(),
provider_id: self.group.candidate.provider_id.clone(),
endpoint_id: self.group.candidate.endpoint_id.clone(),
model_id: self.group.candidate.model_id.clone(),
selected_provider_model_name: self.group.candidate.selected_provider_model_name.clone(),
key_ids,
};
let rows = match self
.state
.app()
.list_pool_key_candidate_rows_for_group_key_ids(&rows_query)
.await
{
Ok(rows) => rows,
Err(err) => {
warn!(
event_name = "pool_group_score_key_load_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
score_count = scores.len(),
error = ?err,
"gateway pool scheduler failed to materialize ranked pool keys"
);
return None;
}
};
self.scanned_keys = self.scanned_keys.saturating_add(scores.len() as u32);
Some(self.build_page_eligible_candidates(rows).await)
}
async fn sticky_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
let pool_config = pool_config_for_candidate(&self.group)?;
if !admin_provider_pool_cache_affinity_enabled(&pool_config) {
@@ -465,6 +569,62 @@ impl<'a> PoolKeyCursor<'a> {
self.skipped_candidates.append(&mut skipped);
}
async fn record_score_schedule_interest(&self, scores: &[StoredPoolMemberScore]) {
if scores.is_empty() {
return;
}
let scheduled_at = current_unix_ms() / 1000;
let mut failed = 0usize;
for score in scores {
let identity = PoolMemberIdentity {
pool_kind: score.pool_kind.clone(),
pool_id: score.pool_id.clone(),
member_kind: score.member_kind.clone(),
member_id: score.member_id.clone(),
};
let scope = PoolScoreScope {
capability: score.capability.clone(),
scope_kind: score.scope_kind.clone(),
scope_id: score.scope_id.clone(),
};
let result = self
.state
.app()
.data
.record_pool_member_schedule_feedback(PoolMemberScheduleFeedback {
identity,
scope: Some(scope),
scheduled_at,
succeeded: None,
hard_state: None,
score_delta: None,
score_reason_patch: Some(serde_json::json!({
"last_schedule_interest": {
"provider_id": self.group.candidate.provider_id.as_str(),
"endpoint_id": self.group.candidate.endpoint_id.as_str(),
"model_id": self.group.candidate.model_id.as_str()
}
})),
})
.await;
if result.is_err() {
failed += 1;
}
}
if failed > 0 {
warn!(
event_name = "pool_group_score_interest_update_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
failed_count = failed,
score_count = scores.len(),
"gateway pool scheduler failed to record some pool score schedule interests"
);
}
}
async fn build_page_eligible_candidates(
&mut self,
rows: Vec<StoredMinimalCandidateSelectionRow>,

192
apps/aether-gateway/src/ai_serving/planner/pool_scores.rs Normal file
View File

@@ -0,0 +1,192 @@
use aether_ai_serving::{score_pool_member, PoolMemberScoreInput, POOL_SCORE_VERSION};
use aether_data_contracts::repository::pool_scores::{
PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore,
POOL_SCORE_SCOPE_KIND_MODEL,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::Value;
use crate::handlers::shared::{provider_key_health_summary, provider_key_status_snapshot_payload};
pub(crate) fn build_provider_key_pool_score_upsert(
key: &StoredProviderCatalogKey,
provider_type: &str,
api_format: &str,
model_id: Option<&str>,
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
now_unix_secs: u64,
) -> 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);
let input = provider_key_score_input(
key,
provider_type,
identity.clone(),
scope.clone(),
existing,
now_unix_secs,
);
let output = score_pool_member(&input);
UpsertPoolMemberScore {
id: provider_key_pool_score_id(&identity, &scope),
identity,
scope,
score: output.score,
hard_state: output.hard_state,
score_version: POOL_SCORE_VERSION,
score_reason: output.score_reason,
last_ranked_at: Some(now_unix_secs),
last_scheduled_at: existing.and_then(|score| score.last_scheduled_at),
last_success_at: existing.and_then(|score| score.last_success_at),
last_failure_at: existing.and_then(|score| score.last_failure_at),
failure_count: existing.map(|score| score.failure_count).unwrap_or(0),
last_probe_attempt_at: existing.and_then(|score| score.last_probe_attempt_at),
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
last_probe_failure_at: existing.and_then(|score| score.last_probe_failure_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),
updated_at: now_unix_secs,
}
}
pub(crate) fn provider_key_pool_score_scope(
api_format: &str,
model_id: Option<&str>,
) -> PoolScoreScope {
PoolScoreScope {
capability: api_format.trim().to_ascii_lowercase(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: model_id
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
}
}
pub(crate) fn provider_key_pool_score_id(
identity: &PoolMemberIdentity,
scope: &PoolScoreScope,
) -> String {
let raw = format!(
"{}:{}:{}:{}:{}:{}:{}",
identity.pool_kind,
identity.pool_id,
identity.member_kind,
identity.member_id,
scope.capability,
scope.scope_kind,
scope.scope_id.as_deref().unwrap_or("*")
);
format!(
"pms-{:016x}-{:016x}",
stable_hash(raw.as_bytes()),
stable_hash(identity.member_id.as_bytes())
)
}
fn provider_key_score_input(
key: &StoredProviderCatalogKey,
provider_type: &str,
identity: PoolMemberIdentity,
scope: PoolScoreScope,
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
now_unix_secs: u64,
) -> PoolMemberScoreInput {
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
let quota_snapshot = status_snapshot
.as_object()
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object);
let account_snapshot = status_snapshot
.as_object()
.and_then(|snapshot| snapshot.get("account"))
.and_then(Value::as_object);
let (health_score, _, _, any_circuit_open, _) = provider_key_health_summary(key);
let health_score = key
.health_by_format
.as_ref()
.and_then(Value::as_object)
.filter(|payload| !payload.is_empty())
.map(|_| health_score);
PoolMemberScoreInput {
identity,
scope: scope.clone(),
internal_priority: key.internal_priority,
is_active: key.is_active,
health_score,
quota_usage_ratio: quota_snapshot
.and_then(|quota| quota.get("usage_ratio"))
.and_then(json_f64)
.map(|value| value.clamp(0.0, 1.0)),
quota_exhausted: quota_snapshot
.and_then(|quota| quota.get("exhausted"))
.and_then(Value::as_bool)
.unwrap_or(false),
account_blocked: account_snapshot
.and_then(|account| account.get("blocked"))
.and_then(Value::as_bool)
.unwrap_or(false),
oauth_invalid_reason: key.oauth_invalid_reason.clone(),
circuit_open: any_circuit_open
|| circuit_open_for_scope(key.circuit_breaker_by_format.as_ref(), &scope),
success_count: key.success_count.unwrap_or(0).into(),
error_count: key.error_count.unwrap_or(0).into(),
total_response_time_ms: key.total_response_time_ms.unwrap_or(0).into(),
total_tokens: key.total_tokens,
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_status: existing
.map(|score| score.probe_status)
.unwrap_or(PoolMemberProbeStatus::Never),
now_unix_secs,
}
}
fn circuit_open_for_scope(circuit_by_format: Option<&Value>, scope: &PoolScoreScope) -> bool {
let Some(formats) = circuit_by_format.and_then(Value::as_object) else {
return false;
};
let keys = api_format_lookup_keys(&scope.capability);
keys.iter().any(|key| {
formats
.get(key)
.and_then(Value::as_object)
.and_then(|value| value.get("open"))
.and_then(Value::as_bool)
.unwrap_or(false)
})
}
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 api_format_lookup_keys(api_format: &str) -> Vec<String> {
let normalized = aether_ai_formats::normalize_api_format_alias(api_format);
let mut keys = aether_ai_formats::api_format_storage_aliases(&normalized);
if !keys.iter().any(|value| value == &normalized) {
keys.push(normalized);
}
keys.sort();
keys.dedup();
keys
}
fn stable_hash(bytes: &[u8]) -> u64 {
let mut hash = 0xcbf29ce484222325u64;
for byte in bytes {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}

12
apps/aether-gateway/src/control/route/admin/observability_families.rs
View File

@@ -240,6 +240,18 @@ pub(super) fn classify_admin_observability_family_route(
"admin:pool",
false,
))
} else if method == http::Method::GET
&& normalized_path_no_trailing.starts_with("/api/admin/pool/")
&& normalized_path_no_trailing.ends_with("/scores")
&& normalized_path_no_trailing.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"pool_manage",
"scores",
"admin:pool",
false,
))
} else if method == http::Method::POST
&& normalized_path_no_trailing.starts_with("/api/admin/pool/")
&& normalized_path_no_trailing.ends_with("/keys/batch-import")

8
apps/aether-gateway/src/control/tests/admin_pool.rs
View File

@@ -52,6 +52,14 @@ fn classifies_admin_pool_provider_key_routes_as_admin_proxy_route() {
assert_eq!(list.route_family.as_deref(), Some("pool_manage"));
assert_eq!(list.route_kind.as_deref(), Some("list_keys"));
let scores_uri: Uri = "/api/admin/pool/provider-1/scores?api_format=openai:responses"
.parse()
.expect("uri should parse");
let scores = classify_control_route(&http::Method::GET, &scores_uri, &headers)
.expect("route should classify");
assert_eq!(scores.route_family.as_deref(), Some("pool_manage"));
assert_eq!(scores.route_kind.as_deref(), Some("scores"));
let batch_import_uri: Uri = "/api/admin/pool/provider-1/keys/batch-import"
.parse()
.expect("uri should parse");

14
apps/aether-gateway/src/data/state/core.rs
View File

@@ -45,6 +45,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -86,6 +88,8 @@ impl GatewayDataState {
let gemini_file_mapping_writer = backends.write().gemini_file_mappings();
let provider_catalog_reader = backends.read().provider_catalog();
let provider_catalog_writer = backends.write().provider_catalog();
let pool_score_reader = backends.read().pool_scores();
let pool_score_writer = backends.write().pool_scores();
let provider_quota_reader = backends.read().provider_quotas();
let provider_quota_writer = backends.write().provider_quotas();
let usage_reader = backends.read().usage();
@@ -125,6 +129,8 @@ impl GatewayDataState {
request_candidate_writer,
provider_catalog_reader,
provider_catalog_writer,
pool_score_reader,
pool_score_writer,
provider_quota_reader,
provider_quota_writer,
usage_reader,
@@ -263,6 +269,14 @@ impl GatewayDataState {
self.provider_catalog_writer.is_some()
}
pub(crate) fn has_pool_score_reader(&self) -> bool {
self.pool_score_reader.is_some()
}
pub(crate) fn has_pool_score_writer(&self) -> bool {
self.pool_score_writer.is_some()
}
pub(crate) fn has_proxy_node_reader(&self) -> bool {
self.proxy_node_reader.is_some()
}

15
apps/aether-gateway/src/data/state/mod.rs
View File

@@ -93,7 +93,8 @@ use aether_data_contracts::repository::billing::{
};
use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
};
use aether_data_contracts::repository::candidates::{
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
@@ -107,6 +108,13 @@ use aether_data_contracts::repository::global_models::{
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use aether_data_contracts::repository::pool_scores::{
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
@@ -156,6 +164,8 @@ pub(crate) struct GatewayDataState {
request_candidate_writer: Option<Arc<dyn RequestCandidateWriteRepository>>,
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
provider_catalog_writer: Option<Arc<dyn ProviderCatalogWriteRepository>>,
pool_score_reader: Option<Arc<dyn PoolScoreReadRepository>>,
pool_score_writer: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
provider_quota_reader: Option<Arc<dyn ProviderQuotaReadRepository>>,
provider_quota_writer: Option<Arc<dyn ProviderQuotaWriteRepository>>,
usage_reader: Option<Arc<dyn UsageReadRepository>>,
@@ -257,6 +267,8 @@ impl fmt::Debug for GatewayDataState {
"has_provider_catalog_writer",
&self.provider_catalog_writer.is_some(),
)
.field("has_pool_score_reader", &self.pool_score_reader.is_some())
.field("has_pool_score_writer", &self.pool_score_writer.is_some())
.field(
"has_provider_quota_reader",
&self.provider_quota_reader.is_some(),
@@ -287,6 +299,7 @@ mod catalog;
mod core;
mod integrations;
mod models;
mod pool_scores;
mod runtime;
#[cfg(test)]
mod testing;

13
apps/aether-gateway/src/data/state/models.rs
View File

@@ -2,7 +2,8 @@ use super::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
DataLayerError, GatewayDataState, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsQuery,
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
StoredPublicGlobalModel, StoredPublicGlobalModelPage, StoredRequestedModelCandidateRowsQuery,
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
@@ -73,6 +74,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_pool_key_candidate_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
match &self.minimal_candidate_selection_reader {
Some(repository) => repository.list_pool_key_rows_for_group_key_ids(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_public_global_models(
&self,
query: &PublicGlobalModelQuery,

123
apps/aether-gateway/src/data/state/pool_scores.rs Normal file
View File

@@ -0,0 +1,123 @@
use super::{
DataLayerError, GatewayDataState, GetPoolMemberScoresByIdsQuery,
ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery, ListRankedPoolMembersQuery,
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt, PoolMemberProbeResult,
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore,
};
impl GatewayDataState {
pub(crate) async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.list_ranked_pool_members(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.list_pool_member_probe_candidates(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.list_pool_member_scores(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.get_pool_member_scores_by_ids(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<Option<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => repository.upsert_pool_member_score(score).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => repository.record_pool_member_probe_result(result).await,
None => Ok(0),
}
}
pub(crate) async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => repository.mark_pool_member_probe_in_progress(attempt).await,
None => Ok(0),
}
}
pub(crate) async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => {
repository
.record_pool_member_schedule_feedback(feedback)
.await
}
None => Ok(0),
}
}
pub(crate) async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => {
repository
.mark_pool_member_hard_state(identity, scope, hard_state, updated_at)
.await
}
None => Ok(0),
}
}
pub(crate) async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => {
repository
.delete_pool_member_scores_for_member(identity)
.await
}
None => Ok(0),
}
}
}

4
apps/aether-gateway/src/data/state/testing/announcements.rs
View File

@@ -34,6 +34,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -86,6 +88,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,

114
apps/aether-gateway/src/data/state/testing/mod.rs
View File

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};
use aether_data_contracts::repository::candidates::RequestCandidateRepository;
use aether_data_contracts::repository::pool_scores::PoolMemberScoreRepository;
use aether_data_contracts::repository::quota::ProviderQuotaRepository;
use aether_data_contracts::repository::usage::UsageRepository;
@@ -12,12 +13,13 @@ use super::{
GeminiFileMappingWriteRepository, GlobalModelReadRepository, GlobalModelWriteRepository,
ManagementTokenReadRepository, ManagementTokenWriteRepository,
MinimalCandidateSelectionReadRepository, OAuthProviderReadRepository,
OAuthProviderWriteRepository, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, ProxyNodeReadRepository,
ProxyNodeWriteRepository, RequestCandidateReadRepository, RequestCandidateWriteRepository,
SettlementWriteRepository, StoredSystemConfigEntry, StoredUserPreferenceRecord,
UsageReadRepository, UsageWriteRepository, UserReadRepository, VideoTaskReadRepository,
VideoTaskWriteRepository, WalletReadRepository, WalletWriteRepository,
OAuthProviderWriteRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, ProviderQuotaReadRepository,
ProviderQuotaWriteRepository, ProxyNodeReadRepository, ProxyNodeWriteRepository,
RequestCandidateReadRepository, RequestCandidateWriteRepository, SettlementWriteRepository,
StoredSystemConfigEntry, StoredUserPreferenceRecord, UsageReadRepository, UsageWriteRepository,
UserReadRepository, VideoTaskReadRepository, VideoTaskWriteRepository, WalletReadRepository,
WalletWriteRepository,
};
mod announcements;
@@ -67,6 +69,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -118,6 +122,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -165,6 +171,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -300,6 +308,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -333,6 +343,18 @@ impl GatewayDataState {
self
}
#[cfg(test)]
pub(crate) fn with_pool_score_repository_for_tests<T>(mut self, repository: Arc<T>) -> Self
where
T: PoolMemberScoreRepository + 'static,
{
let pool_score_reader: Arc<dyn PoolScoreReadRepository> = repository.clone();
let pool_score_writer: Arc<dyn PoolMemberScoreWriteRepository> = repository;
self.pool_score_reader = Some(pool_score_reader);
self.pool_score_writer = Some(pool_score_writer);
self
}
#[cfg(test)]
pub(crate) fn with_provider_catalog_and_request_candidate_reader_for_tests(
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
@@ -363,6 +385,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -419,6 +443,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -484,6 +510,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -531,6 +559,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -579,6 +609,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -638,6 +670,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -699,6 +733,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -744,6 +780,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(repository),
@@ -804,6 +842,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -857,6 +897,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -915,6 +957,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -974,6 +1018,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1032,6 +1078,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -1079,6 +1127,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1126,6 +1176,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1185,6 +1237,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1249,6 +1303,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1296,6 +1352,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1348,6 +1406,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1417,6 +1477,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1481,6 +1543,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1529,6 +1593,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1577,6 +1643,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1627,6 +1695,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_repository),
@@ -1675,6 +1745,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1723,6 +1795,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1779,6 +1853,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -1836,6 +1912,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -1896,6 +1974,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -1962,6 +2042,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2029,6 +2111,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2100,6 +2184,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2178,6 +2264,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2238,6 +2326,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -2289,6 +2379,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2336,6 +2428,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2389,6 +2483,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2446,6 +2542,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2504,6 +2602,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2555,6 +2655,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,

12
apps/aether-gateway/src/data/state/testing/video_tasks.rs
View File

@@ -38,6 +38,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -92,6 +94,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -143,6 +147,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -198,6 +204,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -257,6 +265,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -325,6 +335,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,

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

@@ -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);

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

@@ -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,

17
apps/aether-gateway/src/handlers/admin/provider/pool_admin/mod.rs
View File

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

196
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/scores.rs Normal file
View File

@@ -0,0 +1,196 @@
use super::{
admin_pool_provider_id_from_scores_path, build_admin_pool_error_response,
parse_admin_pool_page, parse_admin_pool_page_size,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::shared::query_param_value;
use crate::GatewayError;
use aether_data_contracts::repository::pool_scores::{
ListPoolMemberScoresQuery, PoolMemberHardState, PoolMemberProbeStatus,
POOL_KIND_PROVIDER_KEY_POOL, POOL_SCORE_SCOPE_KIND_MODEL,
};
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use std::collections::BTreeMap;
pub(super) async fn build_admin_pool_scores_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
let Some(provider_id) = admin_pool_provider_id_from_scores_path(request_context.path()) else {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
"provider_id 无效",
));
};
let query = request_context.query_string();
let page = match parse_admin_pool_page(query) {
Ok(value) => value,
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let page_size = match parse_admin_pool_page_size(query) {
Ok(value) => value.min(500),
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let offset = page.saturating_sub(1).saturating_mul(page_size);
let api_format = query_param_value(query, "api_format")
.map(|value| 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)
}

16
apps/aether-gateway/src/handlers/admin/provider/pool_admin/support.rs
View File

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

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

@@ -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,

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

@@ -11,7 +11,7 @@ pub(crate) use runtime::{
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_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_pool_quota_probe_worker, spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker,
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,

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

@@ -20,6 +20,8 @@ mod oauth_token_refresh;
mod pending_cleanup;
#[path = "runtime/pool_quota_probe.rs"]
mod pool_quota_probe;
#[path = "runtime/pool_score_rebuild.rs"]
mod pool_score_rebuild;
#[path = "runtime/provider_checkin.rs"]
mod provider_checkin;
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
@@ -67,6 +69,10 @@ pub(crate) use pool_quota_probe::{
select_pool_quota_probe_key_ids, spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary,
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,
};
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
use proxy_node_metrics_cleanup::*;
use proxy_node_staleness::*;

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

@@ -1,10 +1,16 @@
use std::collections::BTreeMap;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::pool_scores::{
ListPoolMemberProbeCandidatesQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
POOL_KIND_PROVIDER_KEY_POOL,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
use futures_util::{stream, StreamExt};
use serde_json::Value;
use tracing::{debug, info, warn};
@@ -19,6 +25,7 @@ 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;
const POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER: usize = 50;
const POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY: usize = 16;
const POOL_QUOTA_PROBE_PROVIDER_LOCK_TTL_MS: u64 = 30_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -50,6 +57,7 @@ impl PoolQuotaProbeRunSummary {
pub(crate) struct PoolQuotaProbeWorkerConfig {
pub(crate) scan_interval: Duration,
pub(crate) max_keys_per_provider: usize,
pub(crate) global_concurrency: usize,
}
impl PoolQuotaProbeWorkerConfig {
@@ -63,9 +71,15 @@ impl PoolQuotaProbeWorkerConfig {
"POOL_QUOTA_PROBE_MAX_KEYS_PER_PROVIDER",
POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER,
);
let global_concurrency = env_usize(
"POOL_QUOTA_PROBE_GLOBAL_CONCURRENCY",
POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY,
)
.clamp(1, 256);
Self {
scan_interval: Duration::from_secs(scan_interval_seconds),
max_keys_per_provider,
global_concurrency,
}
}
}
@@ -172,6 +186,49 @@ pub(crate) fn select_pool_quota_probe_key_ids(
stale.into_iter().map(|(_, key_id)| key_id).collect()
}
async fn select_score_probe_key_ids(
state: &AppState,
provider_id: &str,
now_ts: u64,
interval_seconds: u64,
limit: usize,
) -> Vec<String> {
if limit == 0 {
return Vec::new();
}
let stale_before_unix_secs = now_ts.saturating_sub(interval_seconds);
let query = ListPoolMemberProbeCandidatesQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: provider_id.to_string(),
capability: None,
stale_before_unix_secs,
limit: limit.saturating_mul(4).max(limit),
};
let scores = match state.data.list_pool_member_probe_candidates(&query).await {
Ok(scores) => scores,
Err(err) => {
debug!(
provider_id,
error = ?err,
"gateway pool quota probe: failed to read score probe candidates"
);
return Vec::new();
}
};
let mut selected = Vec::new();
let mut seen = std::collections::BTreeSet::new();
for score in scores {
if !seen.insert(score.member_id.clone()) {
continue;
}
selected.push(score.member_id);
if selected.len() >= limit {
break;
}
}
selected
}
fn probe_stamp_key(provider_id: &str, key_id: &str) -> String {
format!("{POOL_QUOTA_PROBE_REDIS_PREFIX}:{provider_id}:{key_id}")
}
@@ -295,7 +352,28 @@ async fn select_keys_for_provider(
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let probe_stamps = load_probe_timestamps(runtime, &provider.id, &key_ids).await;
let selected_ids = select_pool_quota_probe_key_ids(
let mut selected_ids = select_score_probe_key_ids(
state,
&provider.id,
now_ts,
interval_seconds,
max_keys_per_provider,
)
.await
.into_iter()
.filter(|key_id| key_ids.iter().any(|known_id| known_id == key_id))
.filter(|key_id| {
probe_stamps.get(key_id).is_none_or(|last_probe_ts| {
now_ts.saturating_sub(*last_probe_ts) >= interval_seconds
})
})
.collect::<Vec<_>>();
let mut selected_seen = selected_ids
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let remaining = max_keys_per_provider.saturating_sub(selected_ids.len());
let fallback_selected_ids = select_pool_quota_probe_key_ids(
&keys,
provider_type,
now_ts,
@@ -303,6 +381,14 @@ async fn select_keys_for_provider(
&probe_stamps,
max_keys_per_provider,
);
for key_id in fallback_selected_ids {
if selected_ids.len() >= max_keys_per_provider || remaining == 0 {
break;
}
if selected_seen.insert(key_id.clone()) {
selected_ids.push(key_id);
}
}
if selected_ids.is_empty() {
return Ok(Vec::new());
}
@@ -381,6 +467,166 @@ fn update_summary_from_payload(
.unwrap_or(0) as usize;
}
async fn record_score_probe_results_from_payload(
state: &AppState,
provider_id: &str,
selected_key_ids: &[String],
payload: Option<&Value>,
attempted_at: u64,
) {
let mut recorded = std::collections::BTreeSet::new();
if let Some(results) = payload
.and_then(|value| value.get("results"))
.and_then(Value::as_array)
{
for item in results {
let Some(key_id) = item
.get("key_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
recorded.insert(key_id.to_string());
record_score_probe_result_for_key(
state,
provider_id,
key_id,
attempted_at,
probe_result_succeeded(item),
probe_result_hard_state(item),
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": item.get("status").cloned().unwrap_or(Value::Null),
"status_code": item.get("status_code").cloned().unwrap_or(Value::Null),
"message": item.get("message").cloned().unwrap_or(Value::Null),
"auto_removed": item.get("auto_removed").cloned().unwrap_or(Value::Null)
}
}),
)
.await;
}
}
for key_id in selected_key_ids {
if recorded.contains(key_id) {
continue;
}
record_score_probe_result_for_key(
state,
provider_id,
key_id,
attempted_at,
false,
None,
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "missing_result"
}
}),
)
.await;
}
}
async fn record_score_probe_result_for_key(
state: &AppState,
provider_id: &str,
key_id: &str,
attempted_at: u64,
succeeded: bool,
hard_state: Option<PoolMemberHardState>,
score_reason_patch: Value,
) {
let result = PoolMemberProbeResult {
identity: PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.to_string()),
scope: None,
attempted_at,
succeeded,
hard_state,
probe_status: if succeeded {
PoolMemberProbeStatus::Ok
} else {
PoolMemberProbeStatus::Failed
},
score_reason_patch: Some(score_reason_patch),
};
if let Err(err) = state.data.record_pool_member_probe_result(result).await {
debug!(
provider_id,
key_id,
error = ?err,
"gateway pool quota probe: failed to record score probe result"
);
}
}
async fn record_score_probe_in_progress_for_key(
state: &AppState,
provider_id: &str,
key_id: &str,
attempted_at: u64,
) {
let attempt = PoolMemberProbeAttempt {
identity: PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.to_string()),
scope: None,
attempted_at,
score_reason_patch: Some(serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "in_progress"
}
})),
};
if let Err(err) = state.data.mark_pool_member_probe_in_progress(attempt).await {
debug!(
provider_id,
key_id,
error = ?err,
"gateway pool quota probe: failed to mark score probe in progress"
);
}
}
fn probe_result_succeeded(item: &Value) -> bool {
item.get("status")
.and_then(Value::as_str)
.is_some_and(|status| status == "success")
}
fn probe_result_hard_state(item: &Value) -> Option<PoolMemberHardState> {
if probe_result_succeeded(item) {
return Some(PoolMemberHardState::Available);
}
if item
.get("auto_removed")
.and_then(Value::as_bool)
.unwrap_or(false)
{
return Some(PoolMemberHardState::Banned);
}
let status = item
.get("status")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
match status.as_str() {
"auth_invalid" | "forbidden" => Some(PoolMemberHardState::AuthInvalid),
"workspace_deactivated" => Some(PoolMemberHardState::Banned),
"quota_exhausted" => Some(PoolMemberHardState::QuotaExhausted),
_ => match item.get("status_code").and_then(Value::as_u64) {
Some(401 | 403) => Some(PoolMemberHardState::AuthInvalid),
Some(402) => Some(PoolMemberHardState::QuotaExhausted),
Some(429 | 500..=599) => Some(PoolMemberHardState::Cooldown),
_ => None,
},
}
}
pub(crate) async fn perform_pool_quota_probe_once_with_config(
state: &AppState,
config: PoolQuotaProbeWorkerConfig,
@@ -408,6 +654,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
provider,
provider_type,
pool_config.probing_interval_minutes,
pool_config.probe_concurrency.clamp(1, 64) as usize,
))
} else {
None
@@ -421,7 +668,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
@@ -441,7 +688,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
..PoolQuotaProbeRunSummary::empty()
};
for (provider, provider_type, interval_minutes) in providers {
for (provider, provider_type, interval_minutes, probe_concurrency) in providers {
let endpoints = endpoints_by_provider
.remove(&provider.id)
.unwrap_or_default();
@@ -474,42 +721,98 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
summary.providers_probed += 1;
summary.selected_keys += selected_count;
let selected_key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
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>();
match refresh_provider_probe_keys(&admin_state, &provider, &endpoint, &provider_type, keys)
.await
{
Ok(payload) => {
update_summary_from_payload(&mut summary, selected_count, payload.as_ref());
let probe_success = payload
.as_ref()
.and_then(|value| value.get("success"))
.and_then(Value::as_u64)
.unwrap_or(0);
let probe_failed = payload
.as_ref()
.and_then(|value| value.get("failed"))
.and_then(Value::as_u64)
.unwrap_or(0);
info!(
provider_id = %provider_short_id,
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();
let admin_state = &admin_state;
let provider = &provider;
let endpoint = &endpoint;
let provider_type = provider_type.as_str();
async move {
let result = refresh_provider_probe_keys(
admin_state,
provider,
endpoint,
provider_type,
selected = selected_count,
success = probe_success,
failed = probe_failed,
"gateway pool quota probe completed"
);
vec![key],
)
.await;
(key_id, result)
}
Err(err) => {
summary.failed += selected_count;
warn!(
provider_id = %provider_short_id,
provider_type,
selected = selected_count,
error = ?err,
"gateway pool quota probe failed"
);
}))
.buffer_unordered(probe_concurrency)
.collect::<Vec<_>>()
.await;
let mut probe_success = 0usize;
let mut probe_failed = 0usize;
for (key_id, result) in probe_results {
match result {
Ok(payload) => {
update_summary_from_payload(&mut summary, 1, payload.as_ref());
probe_success += payload
.as_ref()
.and_then(|value| value.get("success"))
.and_then(Value::as_u64)
.unwrap_or(0) as usize;
probe_failed += payload
.as_ref()
.and_then(|value| value.get("failed"))
.and_then(Value::as_u64)
.unwrap_or(0) as usize;
record_score_probe_results_from_payload(
state,
&provider.id,
std::slice::from_ref(&key_id),
payload.as_ref(),
now_ts,
)
.await;
}
Err(err) => {
summary.failed += 1;
probe_failed += 1;
record_score_probe_result_for_key(
state,
&provider.id,
&key_id,
now_ts,
false,
Some(PoolMemberHardState::Cooldown),
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "worker_error",
"message": format!("{err:?}")
}
}),
)
.await;
warn!(
provider_id = %provider_short_id,
provider_type,
key_id,
error = ?err,
"gateway pool quota probe failed"
);
}
}
}
info!(
provider_id = %provider_short_id,
provider_type,
selected = selected_count,
success = probe_success,
failed = probe_failed,
concurrency = probe_concurrency,
"gateway pool quota probe completed"
);
}
Ok(summary)

405
apps/aether-gateway/src/maintenance/runtime/pool_score_rebuild.rs Normal file
View File

@@ -0,0 +1,405 @@
use std::collections::BTreeMap;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::global_models::AdminProviderModelListQuery;
use aether_data_contracts::repository::pool_scores::GetPoolMemberScoresByIdsQuery;
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::{AppState, GatewayError};
const POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS: u64 = 300;
const POOL_SCORE_REBUILD_MIN_INTERVAL_SECONDS: u64 = 30;
const POOL_SCORE_REBUILD_DEFAULT_MAX_UPSERTS_PER_TICK: usize = 20_000;
const POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY: &str = "ap:pool_score_rebuild:provider_cursor";
const POOL_SCORE_REBUILD_PROVIDER_OFFSET_PREFIX: &str = "ap:pool_score_rebuild:provider_offset";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct PoolScoreRebuildRunSummary {
pub(crate) providers_checked: usize,
pub(crate) providers_scored: usize,
pub(crate) keys_seen: usize,
pub(crate) scores_upserted: usize,
}
impl PoolScoreRebuildRunSummary {
const fn empty() -> Self {
Self {
providers_checked: 0,
providers_scored: 0,
keys_seen: 0,
scores_upserted: 0,
}
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct PoolScoreRebuildWorkerConfig {
pub(crate) interval: Duration,
pub(crate) max_upserts_per_tick: usize,
}
impl PoolScoreRebuildWorkerConfig {
fn from_env() -> Self {
let interval_seconds = env_u64(
"POOL_SCORE_REBUILD_INTERVAL_SECONDS",
POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS,
)
.max(POOL_SCORE_REBUILD_MIN_INTERVAL_SECONDS);
let max_upserts_per_tick = env_usize(
"POOL_SCORE_REBUILD_MAX_UPSERTS_PER_TICK",
POOL_SCORE_REBUILD_DEFAULT_MAX_UPSERTS_PER_TICK,
)
.max(1);
Self {
interval: Duration::from_secs(interval_seconds),
max_upserts_per_tick,
}
}
}
fn env_u64(name: &str, default_value: u64) -> u64 {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(default_value)
}
fn env_usize(name: &str, default_value: usize) -> usize {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(default_value)
}
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
async fn load_runtime_usize(state: &AppState, key: &str) -> usize {
state
.runtime_state
.kv_get(key)
.await
.ok()
.flatten()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(0)
}
async fn store_runtime_usize(state: &AppState, key: &str, value: usize) {
if let Err(err) = state
.runtime_state
.kv_set(key, value.to_string(), None)
.await
{
debug!(
key,
error = ?err,
"gateway pool score rebuild: failed to store cursor"
);
}
}
fn provider_offset_cursor_key(provider_id: &str) -> String {
format!("{POOL_SCORE_REBUILD_PROVIDER_OFFSET_PREFIX}:{provider_id}")
}
fn score_combo_indices(
flat_index: usize,
model_count: usize,
key_count: usize,
) -> (usize, usize, usize) {
let endpoint_stride = model_count.saturating_mul(key_count).max(1);
let endpoint_index = flat_index / endpoint_stride;
let remainder = flat_index % endpoint_stride;
let model_index = remainder / key_count.max(1);
let key_index = remainder % key_count.max(1);
(endpoint_index, model_index, key_index)
}
#[derive(Debug, Clone)]
struct ProviderScoreBuildItem {
endpoint_index: usize,
model_index: usize,
key_index: usize,
score_id: String,
}
pub(crate) async fn perform_pool_score_rebuild_once_with_config(
state: &AppState,
config: PoolScoreRebuildWorkerConfig,
) -> Result<PoolScoreRebuildRunSummary, GatewayError> {
if !state.has_provider_catalog_data_reader()
|| !state.data.has_pool_score_reader()
|| !state.data.has_pool_score_writer()
{
return Ok(PoolScoreRebuildRunSummary::empty());
}
let mut providers = state
.list_provider_catalog_providers(true)
.await?
.into_iter()
.filter(|provider| admin_provider_pool_config(provider).is_some())
.collect::<Vec<_>>();
providers.sort_by(|left, right| left.id.cmp(&right.id));
if providers.is_empty() {
return Ok(PoolScoreRebuildRunSummary::empty());
}
let provider_ids = providers
.iter()
.map(|provider| provider.id.clone())
.collect::<Vec<_>>();
let mut endpoints_by_provider = BTreeMap::new();
for endpoint in state
.list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
.await?
.into_iter()
.filter(|endpoint| endpoint.is_active)
{
endpoints_by_provider
.entry(endpoint.provider_id.clone())
.or_insert_with(Vec::new)
.push(endpoint);
}
let mut keys_by_provider = BTreeMap::new();
for key in state
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
.await?
{
keys_by_provider
.entry(key.provider_id.clone())
.or_insert_with(Vec::new)
.push(key);
}
let now = now_unix_secs();
let mut summary = PoolScoreRebuildRunSummary {
providers_checked: providers.len(),
..PoolScoreRebuildRunSummary::empty()
};
let start_provider_index =
load_runtime_usize(state, POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY).await % providers.len();
let mut last_provider_index = None;
for provider_index in
(0..providers.len()).map(|offset| (start_provider_index + offset) % providers.len())
{
if summary.scores_upserted >= config.max_upserts_per_tick {
break;
}
last_provider_index = Some(provider_index);
let provider = providers[provider_index].clone();
let endpoints = endpoints_by_provider
.remove(&provider.id)
.unwrap_or_default();
let keys = keys_by_provider.remove(&provider.id).unwrap_or_default();
if endpoints.is_empty() || keys.is_empty() {
continue;
}
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() {
continue;
}
let total_combinations = endpoints
.len()
.saturating_mul(models.len())
.saturating_mul(keys.len());
if total_combinations == 0 {
continue;
}
let provider_cursor_key = provider_offset_cursor_key(&provider.id);
let provider_cursor =
load_runtime_usize(state, &provider_cursor_key).await % total_combinations.max(1);
let remaining_budget = config
.max_upserts_per_tick
.saturating_sub(summary.scores_upserted);
let provider_budget = remaining_budget.min(total_combinations);
let mut build_items = Vec::with_capacity(provider_budget);
for offset in 0..provider_budget {
let flat_index = (provider_cursor + offset) % total_combinations;
let (endpoint_index, model_index, key_index) =
score_combo_indices(flat_index, models.len(), keys.len());
let endpoint = &endpoints[endpoint_index];
let api_format = endpoint.api_format.trim();
if api_format.is_empty() {
continue;
}
let model = &models[model_index];
let key = &keys[key_index];
let draft = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
api_format,
Some(model.id.as_str()),
None,
now,
);
build_items.push(ProviderScoreBuildItem {
endpoint_index,
model_index,
key_index,
score_id: draft.id,
});
}
if build_items.is_empty() {
store_runtime_usize(
state,
&provider_cursor_key,
(provider_cursor + provider_budget) % total_combinations,
)
.await;
continue;
}
let existing_scores = 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 rebuild: failed to read existing scores by id"
);
Vec::new()
})
.into_iter()
.map(|score| (score.id.clone(), score))
.collect::<BTreeMap<_, _>>();
let mut provider_upserts = 0usize;
summary.keys_seen = summary.keys_seen.saturating_add(keys.len());
for item in &build_items {
if summary.scores_upserted >= config.max_upserts_per_tick {
break;
}
let endpoint = &endpoints[item.endpoint_index];
let model = &models[item.model_index];
let key = &keys[item.key_index];
let existing = existing_scores.get(&item.score_id);
let upsert = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
endpoint.api_format.trim(),
Some(model.id.as_str()),
existing,
now,
);
if state
.data
.upsert_pool_member_score(upsert)
.await
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?
.is_some()
{
summary.scores_upserted = summary.scores_upserted.saturating_add(1);
provider_upserts = provider_upserts.saturating_add(1);
}
}
store_runtime_usize(
state,
&provider_cursor_key,
(provider_cursor + provider_budget) % total_combinations,
)
.await;
if provider_upserts > 0 {
summary.providers_scored = summary.providers_scored.saturating_add(1);
}
}
if let Some(last_provider_index) = last_provider_index {
store_runtime_usize(
state,
POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY,
(last_provider_index + 1) % providers.len(),
)
.await;
}
Ok(summary)
}
pub(crate) async fn perform_pool_score_rebuild_once(
state: &AppState,
) -> Result<PoolScoreRebuildRunSummary, GatewayError> {
perform_pool_score_rebuild_once_with_config(state, PoolScoreRebuildWorkerConfig::from_env())
.await
}
pub(crate) fn spawn_pool_score_rebuild_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {
if !state.has_provider_catalog_data_reader()
|| !state.data.has_pool_score_reader()
|| !state.data.has_pool_score_writer()
{
return None;
}
let config = PoolScoreRebuildWorkerConfig::from_env();
Some(tokio::spawn(async move {
if let Err(err) = perform_pool_score_rebuild_once_with_config(&state, config).await {
warn!(
error = ?err,
"gateway pool score rebuild initial tick failed"
);
}
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
interval.tick().await;
match perform_pool_score_rebuild_once_with_config(&state, config).await {
Ok(summary) if summary.scores_upserted > 0 => {
info!(
providers_checked = summary.providers_checked,
providers_scored = summary.providers_scored,
keys_seen = summary.keys_seen,
scores_upserted = summary.scores_upserted,
"gateway pool score rebuild completed"
);
}
Ok(_) => {}
Err(err) => {
warn!(
error = ?err,
"gateway pool score rebuild worker tick failed"
);
}
}
}
}))
}
#[cfg(test)]
mod tests {
use super::score_combo_indices;
#[test]
fn score_combo_indices_walks_endpoint_model_key_order() {
assert_eq!(score_combo_indices(0, 2, 3), (0, 0, 0));
assert_eq!(score_combo_indices(2, 2, 3), (0, 0, 2));
assert_eq!(score_combo_indices(3, 2, 3), (0, 1, 0));
assert_eq!(score_combo_indices(6, 2, 3), (1, 0, 0));
}
}

145
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -2,6 +2,9 @@ use std::collections::BTreeMap;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_data_contracts::repository::pool_scores::{
PoolMemberHardState, PoolMemberIdentity, PoolMemberScheduleFeedback,
};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
count_recent_rpm_requests_for_provider_key, ClientSessionAffinity, SchedulerAffinityTarget,
@@ -381,6 +384,19 @@ async fn record_sync_pool_success_effect(
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(true),
Some(PoolMemberHardState::Available),
Some(50),
serde_json::json!({
"last_request_feedback": {
"source": "sync_success"
}
}),
)
.await;
}
async fn record_adaptive_rate_limit_effect(
@@ -600,6 +616,19 @@ async fn record_stream_pool_success_effect(
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(true),
Some(PoolMemberHardState::Available),
Some(50),
serde_json::json!({
"last_request_feedback": {
"source": "stream_success"
}
}),
)
.await;
}
async fn record_pool_error_effect(
@@ -636,6 +665,21 @@ async fn record_pool_error_effect(
Some(effect.headers),
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(false),
pool_score_hard_state_for_status(effect.status_code, effect.error_body),
Some(pool_score_delta_for_status(effect.status_code)),
serde_json::json!({
"last_request_feedback": {
"source": "pool_error",
"status_code": effect.status_code,
"classification": format!("{:?}", effect.classification)
}
}),
)
.await;
}
async fn open_pool_key_circuit_breaker(
@@ -730,6 +774,21 @@ async fn record_oauth_invalidation_effect(
plan.provider_id, plan.endpoint_id, plan.key_id, err
);
}
record_pool_score_schedule_feedback(
state,
context,
Some(false),
Some(PoolMemberHardState::AuthInvalid),
Some(-2_000),
serde_json::json!({
"last_request_feedback": {
"source": "oauth_invalidation",
"status_code": effect.status_code,
"reason": invalid_reason
}
}),
)
.await;
}
fn resolve_local_oauth_invalid_reason(
@@ -792,6 +851,92 @@ async fn record_pool_stream_timeout_effect(
&pool_context.pool_config,
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(false),
Some(PoolMemberHardState::Cooldown),
Some(-250),
serde_json::json!({
"last_request_feedback": {
"source": "stream_timeout"
}
}),
)
.await;
}
async fn record_pool_score_schedule_feedback(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
succeeded: Option<bool>,
hard_state: Option<PoolMemberHardState>,
score_delta: Option<i32>,
score_reason_patch: Value,
) {
if context.plan.provider_id.trim().is_empty() || context.plan.key_id.trim().is_empty() {
return;
}
let feedback = PoolMemberScheduleFeedback {
identity: PoolMemberIdentity::provider_api_key(
context.plan.provider_id.clone(),
context.plan.key_id.clone(),
),
scope: None,
scheduled_at: current_unix_secs(),
succeeded,
hard_state,
score_delta,
score_reason_patch: Some(score_reason_patch),
};
if let Err(err) = state
.data
.record_pool_member_schedule_feedback(feedback)
.await
{
warn!(
provider_id = %context.plan.provider_id,
key_id = %context.plan.key_id,
error = ?err,
"gateway orchestration effects: failed to record pool score schedule feedback"
);
}
}
fn pool_score_hard_state_for_status(
status_code: u16,
error_body: Option<&str>,
) -> Option<PoolMemberHardState> {
match status_code {
401 | 403 => Some(PoolMemberHardState::AuthInvalid),
402 => Some(PoolMemberHardState::QuotaExhausted),
429 | 500..=599 => Some(PoolMemberHardState::Cooldown),
_ => {
let body = error_body.unwrap_or_default().to_ascii_lowercase();
if body.contains("quota") && body.contains("exceed") {
Some(PoolMemberHardState::QuotaExhausted)
} else if body.contains("invalid") && body.contains("token") {
Some(PoolMemberHardState::AuthInvalid)
} else if body.contains("banned")
|| body.contains("suspended")
|| body.contains("blocked")
{
Some(PoolMemberHardState::Banned)
} else {
None
}
}
}
}
fn pool_score_delta_for_status(status_code: u16) -> i32 {
match status_code {
401 | 403 => -2_000,
402 => -1_000,
429 => -500,
500..=599 => -300,
_ => -100,
}
}
#[cfg(test)]

48
apps/aether-gateway/src/state/catalog.rs
View File

@@ -1,7 +1,8 @@
use super::{AppState, GatewayError, LocalMutationOutcome, LocalProviderDeleteTaskState};
use crate::handlers::shared::sync_provider_key_oauth_status_snapshot;
use aether_data_contracts::repository::{candidates, global_models, provider_catalog};
use aether_data_contracts::repository::{candidates, global_models, pool_scores, provider_catalog};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
impl AppState {
pub fn has_provider_catalog_data_reader(&self) -> bool {
@@ -529,6 +530,25 @@ impl AppState {
.cleanup_deleted_provider_catalog_refs(provider_id, endpoint_ids, key_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
for key_id in key_ids {
if let Err(err) = self
.data
.delete_pool_member_scores_for_member(
&pool_scores::PoolMemberIdentity::provider_api_key(
provider_id.to_string(),
key_id.to_string(),
),
)
.await
{
warn!(
provider_id,
key_id,
error = ?err,
"gateway provider catalog cleanup: failed to delete pool member scores"
);
}
}
if !endpoint_ids.is_empty() || !key_ids.is_empty() {
self.clear_provider_transport_snapshot_cache();
}
@@ -620,12 +640,38 @@ impl AppState {
&self,
key_id: &str,
) -> Result<bool, GatewayError> {
let existing_key = self
.data
.list_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
.into_iter()
.next();
let deleted = self
.data
.delete_provider_catalog_key(key_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
if let Some(key) = existing_key {
if let Err(err) = self
.data
.delete_pool_member_scores_for_member(
&pool_scores::PoolMemberIdentity::provider_api_key(
key.provider_id.clone(),
key.id.clone(),
),
)
.await
{
warn!(
provider_id = %key.provider_id,
key_id = %key.id,
error = ?err,
"gateway provider catalog key delete: failed to delete pool member scores"
);
}
}
self.clear_provider_transport_snapshot_cache();
}
Ok(deleted)

5
apps/aether-gateway/src/state/core.rs
View File

@@ -47,6 +47,7 @@ use crate::maintenance::spawn_oauth_token_refresh_worker;
use crate::maintenance::spawn_pending_cleanup_worker;
use crate::maintenance::spawn_pool_monitor_worker;
use crate::maintenance::spawn_pool_quota_probe_worker;
use crate::maintenance::spawn_pool_score_rebuild_worker;
use crate::maintenance::spawn_provider_checkin_worker;
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
@@ -1113,6 +1114,10 @@ impl AppState {
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
spawn_pool_quota_probe_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_SCORE_REBUILD,
spawn_pool_score_rebuild_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
spawn_stats_hourly_aggregation_worker(self.data.clone()),

10
apps/aether-gateway/src/state/runtime/candidate_queries.rs
View File

@@ -57,6 +57,16 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_pool_key_candidate_rows_for_group_key_ids(
&self,
query: &candidate_selection::StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<candidate_selection::StoredMinimalCandidateSelectionRow>, GatewayError> {
self.data
.list_pool_key_candidate_rows_for_group_key_ids(query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn read_provider_quota_snapshot(
&self,
provider_id: &str,

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

@@ -22,6 +22,7 @@ pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
@@ -101,6 +102,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_SCORE_REBUILD,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_MONITOR,
TaskKind::Scheduled,

93
apps/aether-gateway/src/tests/control/admin/pool.rs
View File

@@ -1,8 +1,13 @@
use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::pool_scores::{
PoolMemberHardState, PoolMemberProbeStatus, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_SCOPE_KIND_MODEL,
};
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use axum::body::{to_bytes, Body, Bytes};
@@ -343,6 +348,94 @@ async fn gateway_handles_admin_pool_batch_import_locally_with_trusted_admin_prin
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_pool_scores_locally_with_trusted_admin_principal() {
let provider = sample_provider("provider-openai", "openai", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
let mut key = sample_key("key-openai-a", "provider-openai", "openai:chat", "sk-a");
key.name = "score key".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key.clone()],
));
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::seed(vec![
StoredPoolMemberScore {
id: "pms-provider-openai-key-openai-a-openai-chat-model-1".to_string(),
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-openai".to_string(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: "key-openai-a".to_string(),
capability: "openai:chat".to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
score: 0.875,
hard_state: PoolMemberHardState::Available,
score_version: 1,
score_reason: json!({ "weights": { "manual_priority": 0.3 } }),
last_ranked_at: Some(1_700_000_000),
last_scheduled_at: Some(1_700_000_010),
last_success_at: Some(1_700_000_020),
last_failure_at: None,
failure_count: 0,
last_probe_attempt_at: Some(1_700_000_030),
last_probe_success_at: Some(1_700_000_040),
last_probe_failure_at: None,
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Ok,
updated_at: 1_700_000_050,
},
]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
))
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository)),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/pool/provider-openai/scores?api_format=openai:chat&model_id=model-1"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(
payload["items"].as_array().map(|items| items.len()),
Some(1)
);
assert_eq!(payload["items"][0]["member_id"], json!("key-openai-a"));
assert_eq!(payload["items"][0]["key"]["name"], json!("score key"));
assert_eq!(payload["items"][0]["probe_status"], json!("ok"));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_pool_trailing_slash_routes_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));