i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 17:30:23 +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

207
crates/aether-scheduler-core/src/candidate/mod.rs
View File

@@ -32,9 +32,9 @@ mod tests {
use super::{
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
candidate_supports_required_capability, collect_global_model_names_for_required_capability,
CandidateRuntimeSelectabilityInput, EnumerateMinimalCandidateSelectionInput,
SchedulerMinimalCandidateSelectionCandidate,
candidate_runtime_skip_reason_with_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability, CandidateRuntimeSelectabilityInput,
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
};
use crate::SchedulerAuthConstraints;
@@ -118,6 +118,15 @@ mod tests {
key
}
fn sample_key_with_concurrent_limit(
id: &str,
concurrent_limit: Option<i32>,
) -> StoredProviderCatalogKey {
let mut key = sample_key(id, 1.0);
key.concurrent_limit = concurrent_limit;
key
}
fn stored_candidate(
id: &str,
status: RequestCandidateStatus,
@@ -304,6 +313,198 @@ mod tests {
));
}
#[test]
fn provider_key_concurrency_limit_unset_or_zero_is_unlimited() {
let recent_candidates = vec![stored_candidate("one", RequestCandidateStatus::Pending, 95)];
for concurrent_limit in [None, Some(0)] {
let provider_key_rpm_states = BTreeMap::from([(
"key-1".to_string(),
sample_key_with_concurrent_limit("1", concurrent_limit),
)]);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
None
);
}
}
#[test]
fn provider_key_concurrency_limit_rejects_pending_active_with_exact_skip_reason() {
let recent_candidates = vec![stored_candidate("one", RequestCandidateStatus::Pending, 95)];
let provider_key_rpm_states = BTreeMap::from([(
"key-1".to_string(),
sample_key_with_concurrent_limit("1", Some(1)),
)]);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
Some("provider_key_concurrency_limit_reached")
);
}
#[test]
fn provider_key_concurrency_limit_rejects_streaming_active() {
let recent_candidates = vec![stored_candidate(
"streaming",
RequestCandidateStatus::Streaming,
95,
)];
let provider_key_rpm_states = BTreeMap::from([(
"key-1".to_string(),
sample_key_with_concurrent_limit("1", Some(1)),
)]);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
Some("provider_key_concurrency_limit_reached")
);
}
#[test]
fn provider_key_concurrency_limit_ignores_finished_and_stale_active_requests() {
let recent_candidates = vec![
stored_candidate("finished", RequestCandidateStatus::Success, 95),
stored_candidate("failed", RequestCandidateStatus::Failed, 96),
stored_candidate("cancelled", RequestCandidateStatus::Cancelled, 97),
stored_candidate("stale", RequestCandidateStatus::Pending, 699_000),
];
let provider_key_rpm_states = BTreeMap::from([(
"key-1".to_string(),
sample_key_with_concurrent_limit("1", Some(1)),
)]);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 1_000,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
None
);
}
#[test]
fn provider_key_concurrency_limit_missing_state_does_not_skip() {
let recent_candidates = vec![stored_candidate("one", RequestCandidateStatus::Pending, 95)];
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
None
);
}
#[test]
fn provider_key_concurrency_limit_preserves_key_circuit_and_rpm_checks() {
let mut circuit_open_key = sample_key_with_concurrent_limit("1", Some(2));
circuit_open_key.circuit_breaker_by_format = Some(serde_json::json!({
"openai:chat": {"open": true}
}));
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), circuit_open_key)]);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
Some("key_circuit_open")
);
let recent_candidates = vec![stored_candidate(
"one",
RequestCandidateStatus::Pending,
95_000,
)];
let provider_key_rpm_states = BTreeMap::from([(
"key-1".to_string(),
sample_key_with_concurrent_limit("1", Some(2)).with_rate_limit_fields(
Some(1),
Some(2),
None,
None,
None,
None,
None,
None,
None,
),
)]);
assert_eq!(
candidate_runtime_skip_reason_with_state(CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
Some("key_rpm_exhausted")
);
}
#[test]
fn candidate_selectability_rejects_quota_or_zero_health() {
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), sample_key("1", 0.0))]);

20
crates/aether-scheduler-core/src/candidate/selectability.rs
View File

@@ -86,9 +86,27 @@ pub fn candidate_runtime_skip_reason_with_state(
return Some("provider_concurrency_limit_reached");
}
let provider_key = provider_key_rpm_states.get(&candidate.key_id);
if let Some(provider_key) = provider_key {
if let Some(limit) = provider_key
.concurrent_limit
.filter(|limit| *limit > 0)
.and_then(|limit| usize::try_from(limit).ok())
{
if crate::count_recent_active_requests_for_provider_key(
recent_candidates,
candidate.key_id.as_str(),
now_unix_secs,
) >= limit
{
return Some("provider_key_concurrency_limit_reached");
}
}
}
let is_cached_user = cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(candidate, target));
if let Some(provider_key) = provider_key_rpm_states.get(&candidate.key_id) {
if let Some(provider_key) = provider_key {
if crate::is_provider_key_circuit_open(provider_key, candidate.endpoint_api_format.as_str())
{
return Some("key_circuit_open");