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

feat: add provider-key concurrent limit (#352)

Browse Source 
* feat: add provider-key concurrent limit

* Fix provider key concurrent limit checks

---------

Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
Kayphoon
2026-05-03 01:55:23 +08:00
committed by GitHub
parent 11c5884d4f
commit fe27fb17fb
23 changed files with 1827 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
apps/aether-gateway/src/scheduler/candidate/affinity.rs
View File

@@ -1,5 +1,6 @@
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
matches_affinity_target, SchedulerAffinityTarget,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
@@ -21,6 +22,20 @@ pub(super) fn build_scheduler_affinity_cache_key(
build_scheduler_affinity_cache_key_for_api_key_id(api_key_id, api_format, global_model_name)
}
pub(super) fn scheduler_candidate_affinity_hash(
affinity_key: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> u64 {
candidate_affinity_hash(affinity_key, candidate)
}
pub(super) fn scheduler_candidate_matches_affinity_target(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
target: &SchedulerAffinityTarget,
) -> bool {
matches_affinity_target(candidate, target)
}
#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn remember_scheduler_affinity(
affinity_cache_key: Option<&str>,
@@ -30,13 +45,14 @@ pub(super) fn remember_scheduler_affinity(
let Some(cache_key) = affinity_cache_key else {
return;
};
let (provider_id, endpoint_id, key_id) = candidate_key(candidate);
state.remember_scheduler_affinity_target(
cache_key,
SchedulerAffinityTarget {
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
provider_id,
endpoint_id,
key_id,
},
SCHEDULER_AFFINITY_TTL,
SCHEDULER_AFFINITY_MAX_ENTRIES,

21
apps/aether-gateway/src/scheduler/candidate/ranking.rs
View File

@@ -1,12 +1,15 @@
use aether_scheduler_core::{
apply_scheduler_candidate_ranking, candidate_affinity_hash,
effective_provider_key_health_score, matches_affinity_target, provider_key_health_bucket,
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
apply_scheduler_candidate_ranking, effective_provider_key_health_score,
provider_key_health_bucket, requested_capability_priority_for_candidate,
SchedulerAffinityTarget, SchedulerRankableCandidate, SchedulerRankingContext,
SchedulerRankingMode,
};
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use super::affinity::{
scheduler_candidate_affinity_hash, scheduler_candidate_matches_affinity_target,
};
use super::runtime::CandidateRuntimeSelectionSnapshot;
use super::SchedulerMinimalCandidateSelectionCandidate;
@@ -31,12 +34,12 @@ pub(super) fn rank_scheduler_candidates(
required_capabilities,
candidate,
))
.with_cached_affinity_match(
cached_affinity_target
.is_some_and(|target| matches_affinity_target(candidate, target)),
)
.with_cached_affinity_match(cached_affinity_target.is_some_and(|target| {
scheduler_candidate_matches_affinity_target(candidate, target)
}))
.with_affinity_hash(
priority_affinity_key.map(|key| candidate_affinity_hash(key, candidate)),
priority_affinity_key
.map(|key| scheduler_candidate_affinity_hash(key, candidate)),
)
.with_health(
provider_key.and_then(|key| {

389
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -5,10 +5,13 @@ use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelect
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
use aether_data_contracts::repository::candidate_selection::StoredProviderModelMapping;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
@@ -99,6 +102,102 @@ async fn collect_selectable_candidates_with_skip_reasons(
.await
}
fn provider_key_concurrency_row(
provider_id: &str,
endpoint_id: &str,
key_id: &str,
key_name: &str,
provider_priority: i32,
key_priority: i32,
) -> StoredMinimalCandidateSelectionRow {
let mut row = sample_row();
row.provider_id = provider_id.to_string();
row.provider_name = provider_id.to_string();
row.endpoint_id = endpoint_id.to_string();
row.key_id = key_id.to_string();
row.key_name = key_name.to_string();
row.provider_priority = provider_priority;
row.key_internal_priority = key_priority;
row.key_global_priority_by_format = Some(serde_json::json!({"openai:chat": key_priority}));
row
}
fn provider_key_with_concurrent_limit(
key_id: &str,
provider_id: &str,
concurrent_limit: Option<i32>,
) -> StoredProviderCatalogKey {
let mut key = sample_key(key_id, provider_id, Some(10));
key.concurrent_limit = concurrent_limit;
key
}
fn active_provider_key_candidate(
candidate_id: &str,
request_id: &str,
provider_id: &str,
endpoint_id: &str,
key_id: &str,
status: RequestCandidateStatus,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
candidate_id.to_string(),
request_id.to_string(),
None,
None,
None,
None,
0,
0,
Some(provider_id.to_string()),
Some(endpoint_id.to_string()),
Some(key_id.to_string()),
status,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95_000,
Some(95_000),
None,
)
.expect("candidate should build")
}
fn provider_key_concurrency_state(
rows: Vec<StoredMinimalCandidateSelectionRow>,
keys: Vec<StoredProviderCatalogKey>,
request_candidates: Vec<StoredRequestCandidate>,
) -> AppState {
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("test-provider-a", None),
sample_provider("test-provider-b", None),
],
Vec::new(),
keys,
));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(request_candidates));
AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_provider_catalog_quota_and_request_candidates_for_tests(
candidates,
provider_catalog,
quotas,
request_candidates,
),
)
}
#[test]
fn skips_only_exhausted_monthly_quota_provider() {
let inactive = StoredProviderQuotaSnapshot::new(
@@ -860,6 +959,294 @@ async fn selects_next_candidate_when_first_provider_concurrent_limit_is_reached(
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn provider_key_concurrency_selects_next_key_when_first_provider_key_concurrent_limit_is_reached(
) {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
],
vec![active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Streaming,
)],
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "test-provider-a");
assert_eq!(selected.key_id, "provider-key-b");
}
#[tokio::test]
async fn provider_key_concurrency_selects_next_provider_when_all_provider_keys_concurrent_limit_reached(
) {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
provider_key_concurrency_row(
"test-provider-b",
"endpoint-b",
"provider-key-c",
"gamma",
1,
0,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-c", "test-provider-b", Some(1)),
],
vec![
active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Pending,
),
active_provider_key_candidate(
"cand-provider-key-b",
"req-provider-key-b",
"test-provider-a",
"endpoint-a",
"provider-key-b",
RequestCandidateStatus::Streaming,
),
],
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.provider_id, "test-provider-b");
assert_eq!(selected.key_id, "provider-key-c");
}
#[tokio::test]
async fn provider_key_concurrency_returns_none_when_all_provider_keys_concurrent_limit_reached() {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
provider_key_concurrency_row(
"test-provider-b",
"endpoint-b",
"provider-key-c",
"gamma",
1,
0,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-c", "test-provider-b", Some(1)),
],
vec![
active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Pending,
),
active_provider_key_candidate(
"cand-provider-key-b",
"req-provider-key-b",
"test-provider-a",
"endpoint-a",
"provider-key-b",
RequestCandidateStatus::Streaming,
),
active_provider_key_candidate(
"cand-provider-key-c",
"req-provider-key-c",
"test-provider-b",
"endpoint-b",
"provider-key-c",
RequestCandidateStatus::Streaming,
),
],
);
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed");
assert!(selected.is_none());
let (selected_candidates, skipped_candidates) =
collect_selectable_candidates_with_skip_reasons(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed");
assert!(selected_candidates.is_empty());
assert_eq!(skipped_candidates.len(), 3);
assert!(skipped_candidates
.iter()
.all(|skipped| { skipped.skip_reason == "provider_key_concurrency_limit_reached" }));
assert!(!is_exact_all_skipped_by_auth_limit(
&selected_candidates,
&skipped_candidates,
));
}
#[tokio::test]
async fn provider_key_concurrency_collects_exact_skip_reason_for_saturated_provider_keys() {
let state = provider_key_concurrency_state(
vec![
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-a",
"alpha",
0,
0,
),
provider_key_concurrency_row(
"test-provider-a",
"endpoint-a",
"provider-key-b",
"beta",
0,
1,
),
],
vec![
provider_key_with_concurrent_limit("provider-key-a", "test-provider-a", Some(1)),
provider_key_with_concurrent_limit("provider-key-b", "test-provider-a", Some(1)),
],
vec![active_provider_key_candidate(
"cand-provider-key-a",
"req-provider-key-a",
"test-provider-a",
"endpoint-a",
"provider-key-a",
RequestCandidateStatus::Pending,
)],
);
let (selected_candidates, skipped_candidates) =
collect_selectable_candidates_with_skip_reasons(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
None,
100,
)
.await
.expect("selection should succeed");
assert_eq!(selected_candidates.len(), 1);
assert_eq!(selected_candidates[0].provider_id, "test-provider-a");
assert_eq!(selected_candidates[0].key_id, "provider-key-b");
assert_eq!(skipped_candidates.len(), 1);
assert_eq!(
skipped_candidates[0].candidate.provider_id,
"test-provider-a"
);
assert_eq!(skipped_candidates[0].candidate.key_id, "provider-key-a");
assert_eq!(
skipped_candidates[0].skip_reason,
"provider_key_concurrency_limit_reached",
);
}
#[tokio::test]
async fn returns_none_when_auth_api_key_concurrent_limit_is_reached() {
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![

12
apps/aether-gateway/src/scheduler/candidate/tests/support.rs
View File

@@ -90,7 +90,17 @@ pub(super) fn sample_key(
true,
)
.expect("key should build")
.with_rate_limit_fields(rpm_limit, None, None, None, None, None, Some(20), Some(20))
.with_rate_limit_fields(
rpm_limit,
None,
None,
None,
None,
None,
None,
Some(20),
Some(20),
)
}
pub(super) fn sample_auth_snapshot(api_key_id: &str) -> GatewayAuthApiKeySnapshot {