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

51
crates/aether-data-contracts/src/repository/provider_catalog/types.rs
View File

@@ -263,6 +263,7 @@ pub struct StoredProviderCatalogKey {
pub proxy: Option<serde_json::Value>,
pub fingerprint: Option<serde_json::Value>,
pub rpm_limit: Option<u32>,
pub concurrent_limit: Option<i32>,
pub learned_rpm_limit: Option<u32>,
pub concurrent_429_count: Option<u32>,
pub rpm_429_count: Option<u32>,
@@ -338,6 +339,7 @@ impl StoredProviderCatalogKey {
proxy: None,
fingerprint: None,
rpm_limit: None,
concurrent_limit: None,
learned_rpm_limit: None,
concurrent_429_count: None,
rpm_429_count: None,
@@ -410,6 +412,7 @@ impl StoredProviderCatalogKey {
pub fn with_rate_limit_fields(
mut self,
rpm_limit: Option<u32>,
concurrent_limit: Option<i32>,
learned_rpm_limit: Option<u32>,
concurrent_429_count: Option<u32>,
rpm_429_count: Option<u32>,
@@ -419,6 +422,7 @@ impl StoredProviderCatalogKey {
success_count: Option<u32>,
) -> Self {
self.rpm_limit = rpm_limit;
self.concurrent_limit = concurrent_limit;
self.learned_rpm_limit = learned_rpm_limit;
self.concurrent_429_count = concurrent_429_count;
self.rpm_429_count = rpm_429_count;
@@ -460,6 +464,53 @@ impl StoredProviderCatalogKey {
}
}
#[cfg(test)]
mod transport_tests {
use super::StoredProviderCatalogKey;
#[test]
fn provider_catalog_key_defaults_concurrent_limit_to_none() {
let key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"default".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build");
assert_eq!(key.concurrent_limit, None);
}
#[test]
fn provider_catalog_key_rate_limit_builder_sets_concurrent_limit() {
let key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"default".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_rate_limit_fields(
Some(120),
Some(3),
None,
None,
None,
None,
None,
None,
None,
);
assert_eq!(key.rpm_limit, Some(120));
assert_eq!(key.concurrent_limit, Some(3));
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub enum ProviderCatalogKeyListOrder {
#[default]

30
crates/aether-data/src/repository/provider_catalog/memory.rs
View File

@@ -845,6 +845,36 @@ mod tests {
assert_eq!(stored[0].provider_id, "provider-1");
}
#[tokio::test]
async fn provider_api_keys_concurrent_limit_defaults_and_round_trips_in_memory() {
let repository = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("test-provider-a")],
vec![],
vec![],
);
let mut key = sample_key("provider-key-a", "test-provider-a");
assert_eq!(key.concurrent_limit, None);
key.concurrent_limit = Some(1);
let created = repository
.create_key(&key)
.await
.expect("key should create");
assert_eq!(created.concurrent_limit, Some(1));
let mut updated = created.clone();
updated.concurrent_limit = None;
repository
.update_key(&updated)
.await
.expect("key should update");
let reloaded = repository
.list_keys_by_ids(&["provider-key-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded[0].concurrent_limit, None);
}
#[tokio::test]
async fn creates_endpoint() {
let repository = InMemoryProviderCatalogReadRepository::seed(

113
crates/aether-data/src/repository/provider_catalog/sql.rs
View File

@@ -155,6 +155,7 @@ SELECT
proxy,
fingerprint,
rpm_limit,
concurrent_limit,
learned_rpm_limit,
concurrent_429_count,
rpm_429_count,
@@ -213,6 +214,7 @@ SELECT
proxy,
fingerprint,
rpm_limit,
concurrent_limit,
learned_rpm_limit,
concurrent_429_count,
rpm_429_count,
@@ -273,6 +275,7 @@ SELECT
NULL::jsonb AS proxy,
NULL::jsonb AS fingerprint,
NULL::integer AS rpm_limit,
NULL::integer AS concurrent_limit,
NULL::integer AS learned_rpm_limit,
NULL::integer AS concurrent_429_count,
NULL::integer AS rpm_429_count,
@@ -611,6 +614,7 @@ SELECT
proxy,
fingerprint,
rpm_limit,
concurrent_limit,
learned_rpm_limit,
concurrent_429_count,
rpm_429_count,
@@ -1199,6 +1203,7 @@ INSERT INTO provider_api_keys (
internal_priority,
global_priority_by_format,
rpm_limit,
concurrent_limit,
learned_rpm_limit,
allowed_models,
capabilities,
@@ -1331,6 +1336,7 @@ INSERT INTO provider_api_keys (
.bind(key.internal_priority)
.bind(&key.global_priority_by_format)
.bind(key.rpm_limit.map(|value| value as i32))
.bind(key.concurrent_limit)
.bind(key.learned_rpm_limit.map(|value| value as i32))
.bind(&key.allowed_models)
.bind(&key.capabilities)
@@ -1726,8 +1732,8 @@ UPDATE provider_api_keys
SET
provider_id = $2,
api_formats = $3,
auth_type_by_format = $39,
allow_auth_channel_mismatch_formats = $40,
auth_type_by_format = $40,
allow_auth_channel_mismatch_formats = $41,
auth_type = $4,
api_key = $5,
auth_config = $6,
@@ -1737,46 +1743,47 @@ SET
internal_priority = $10,
global_priority_by_format = $11,
rpm_limit = $12,
learned_rpm_limit = $13,
allowed_models = $14,
capabilities = $15,
cache_ttl_minutes = $16,
max_probe_interval_minutes = $17,
auto_fetch_models = $18,
locked_models = $19,
model_include_patterns = $20,
model_exclude_patterns = $21,
proxy = $22,
fingerprint = $23,
upstream_metadata = $24,
concurrent_limit = $13,
learned_rpm_limit = $14,
allowed_models = $15,
capabilities = $16,
cache_ttl_minutes = $17,
max_probe_interval_minutes = $18,
auto_fetch_models = $19,
locked_models = $20,
model_include_patterns = $21,
model_exclude_patterns = $22,
proxy = $23,
fingerprint = $24,
upstream_metadata = $25,
expires_at = CASE
WHEN $38::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($38::double precision)
WHEN $39::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($39::double precision)
END,
oauth_invalid_at = CASE
WHEN $25::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($25::double precision)
WHEN $26::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($26::double precision)
END,
oauth_invalid_reason = $26,
status_snapshot = $27,
concurrent_429_count = COALESCE($28, 0),
rpm_429_count = COALESCE($29, 0),
oauth_invalid_reason = $27,
status_snapshot = $28,
concurrent_429_count = COALESCE($29, 0),
rpm_429_count = COALESCE($30, 0),
last_429_at = CASE
WHEN $30::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($30::double precision)
WHEN $31::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($31::double precision)
END,
last_429_type = $31,
adjustment_history = $32,
utilization_samples = $33,
last_429_type = $32,
adjustment_history = $33,
utilization_samples = $34,
last_probe_increase_at = CASE
WHEN $34::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($34::double precision)
WHEN $35::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($35::double precision)
END,
last_rpm_peak = $35,
is_active = $36,
last_rpm_peak = $36,
is_active = $37,
updated_at = CASE
WHEN $37::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($37::double precision)
WHEN $38::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($38::double precision)
END
WHERE id = $1
"#,
@@ -1793,6 +1800,7 @@ WHERE id = $1
.bind(key.internal_priority)
.bind(&key.global_priority_by_format)
.bind(key.rpm_limit.map(|value| value as i32))
.bind(key.concurrent_limit)
.bind(key.learned_rpm_limit.map(|value| value as i32))
.bind(&key.allowed_models)
.bind(&key.capabilities)
@@ -2301,6 +2309,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
})
})
.transpose()?;
let concurrent_limit = row_get::<Option<i32>>(row, "concurrent_limit")?;
let learned_rpm_limit = row_get::<Option<i32>>(row, "learned_rpm_limit")?
.map(|value| {
u32::try_from(value).map_err(|_| {
@@ -2466,6 +2475,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
let mut key = key
.with_rate_limit_fields(
rpm_limit,
concurrent_limit,
learned_rpm_limit,
concurrent_429_count,
rpm_429_count,
@@ -2545,4 +2555,41 @@ mod tests {
assert!(sql.contains("total_cost_usd"));
}
}
#[test]
fn provider_api_keys_concurrent_limit_queries_include_field() {
for sql in [
super::LIST_KEYS_BY_IDS_PREFIX,
super::LIST_KEYS_BY_PROVIDER_IDS_PREFIX,
super::LIST_KEY_SUMMARIES_BY_PROVIDER_IDS_PREFIX,
] {
assert!(sql.contains("concurrent_limit"));
}
let source = include_str!("sql.rs");
assert!(source.contains("concurrent_limit,"));
assert!(source.contains("concurrent_limit = $13"));
assert!(source.contains(".bind(key.concurrent_limit)"));
assert!(source.contains("row_get::<Option<i32>>(row, \"concurrent_limit\")"));
}
#[test]
fn provider_api_keys_concurrent_limit_schema_is_nullable_without_default() {
let migration = include_str!(
"../../../migrations/20260502000000_add_provider_key_auth_channel_mismatch_formats.sql"
);
let concurrent_limit_line = migration
.lines()
.find(|line| line.contains("ADD COLUMN IF NOT EXISTS concurrent_limit"))
.expect("concurrent_limit migration line should exist")
.to_ascii_lowercase();
assert_eq!(
concurrent_limit_line.trim(),
"add column if not exists concurrent_limit integer;"
);
let baseline = include_str!("../../../bootstrap/20260413020000_baseline_v2.sql");
assert!(baseline.contains("CREATE TABLE IF NOT EXISTS public.provider_api_keys"));
assert!(baseline.contains("concurrent_limit integer,"));
}
}

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

220
crates/aether-scheduler-core/src/health.rs
View File

@@ -100,6 +100,18 @@ pub fn count_recent_active_requests_for_provider(
.count()
}
pub fn count_recent_active_requests_for_provider_key(
recent_candidates: &[StoredRequestCandidate],
key_id: &str,
now_unix_secs: u64,
) -> usize {
recent_candidates
.iter()
.filter(|candidate| candidate.key_id.as_deref() == Some(key_id))
.filter(|candidate| is_recently_active(candidate, now_unix_secs))
.count()
}
pub fn count_recent_active_requests_for_api_key(
recent_candidates: &[StoredRequestCandidate],
api_key_id: &str,
@@ -598,7 +610,8 @@ mod tests {
use super::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
count_recent_active_requests_for_provider, count_recent_rpm_requests_for_provider_key,
count_recent_active_requests_for_provider, count_recent_active_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key_since, effective_provider_key_health_score,
effective_provider_key_rpm_limit, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, provider_key_health_bucket, provider_key_health_score,
@@ -782,10 +795,211 @@ mod tests {
);
}
#[test]
fn provider_key_concurrency_counts_only_recent_active_requests() {
let recent_candidates = vec![
StoredRequestCandidate::new(
"pending".to_string(),
"req-pending".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
900_000,
Some(900_000),
None,
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"streaming".to_string(),
"req-streaming".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Streaming,
None,
false,
None,
None,
None,
None,
None,
None,
None,
950_000,
Some(950_000),
None,
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"finished".to_string(),
"req-finished".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
None,
None,
None,
None,
None,
None,
None,
975_000,
Some(975_000),
Some(976_000),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"failed".to_string(),
"req-failed".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(429),
None,
Some("upstream failure".to_string()),
Some(20),
None,
None,
None,
980_000,
Some(980_000),
Some(981_000),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"cancelled".to_string(),
"req-cancelled".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Cancelled,
None,
false,
Some(499),
None,
Some("client cancelled".to_string()),
Some(10),
None,
None,
None,
982_000,
Some(982_000),
Some(983_000),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"stale".to_string(),
"req-stale".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
699_000,
Some(699_000),
None,
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"other-key".to_string(),
"req-other-key".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-b".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
990_000,
Some(990_000),
None,
)
.expect("candidate should build"),
];
assert_eq!(
count_recent_active_requests_for_provider_key(&recent_candidates, "key-a", 1_000),
2
);
}
#[test]
fn fixed_provider_key_rpm_limit_takes_precedence() {
let key = provider_catalog_key("key-a").with_rate_limit_fields(
Some(120),
None,
Some(80),
None,
None,
@@ -801,6 +1015,7 @@ mod tests {
#[test]
fn learned_provider_key_rpm_limit_requires_confidence() {
let low_confidence = provider_catalog_key("key-a").with_rate_limit_fields(
None,
None,
Some(80),
Some(0),
@@ -813,6 +1028,7 @@ mod tests {
assert_eq!(effective_provider_key_rpm_limit(&low_confidence, 100), None);
let mut high_confidence = provider_catalog_key("key-a").with_rate_limit_fields(
None,
None,
Some(80),
Some(0),
@@ -840,6 +1056,7 @@ mod tests {
#[test]
fn learned_provider_key_rpm_limit_uses_confirmed_observations_as_fallback_confidence() {
let key = provider_catalog_key("key-a").with_rate_limit_fields(
None,
None,
Some(80),
Some(0),
@@ -1010,6 +1227,7 @@ mod tests {
None,
None,
None,
None,
Some(5),
Some(5),
);

12
crates/aether-scheduler-core/src/lib.rs
View File

@@ -26,12 +26,12 @@ pub use candidate::{
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
count_recent_active_requests_for_provider, count_recent_rpm_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key_since, effective_provider_key_health_score,
effective_provider_key_rpm_limit, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
count_recent_active_requests_for_provider, count_recent_active_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key, count_recent_rpm_requests_for_provider_key_since,
effective_provider_key_health_score, effective_provider_key_rpm_limit,
is_candidate_in_recent_failure_cooldown, is_provider_key_circuit_open,
provider_key_health_bucket, provider_key_health_score, provider_key_rpm_allows_request,
provider_key_rpm_allows_request_since, ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
};
pub use model::{
candidate_model_names, extract_global_priority_for_format, matches_model_mapping,