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

Cap Codex pool cooldowns and add key circuit breaker

Browse Source
This commit is contained in:
fawney19
2026-05-04 02:15:32 +08:00
parent 1d62722d47
commit 099653f732
10 changed files with 536 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -506,6 +506,49 @@ mod tests {
);
}
#[test]
fn key_circuit_allows_probe_after_next_probe_time() {
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,
"next_probe_at_unix_secs": 100
}
}));
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: 99,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
}),
Some("key_circuit_open")
);
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,
}),
None
);
}
#[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))]);

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

@@ -107,8 +107,11 @@ pub fn candidate_runtime_skip_reason_with_state(
let is_cached_user = cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(candidate, target));
if let Some(provider_key) = provider_key {
if crate::is_provider_key_circuit_open(provider_key, candidate.endpoint_api_format.as_str())
{
if crate::is_provider_key_circuit_open_at(
provider_key,
candidate.endpoint_api_format.as_str(),
now_unix_secs,
) {
return Some("key_circuit_open");
}
if crate::provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())

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

@@ -284,6 +284,33 @@ pub fn is_provider_key_circuit_open(key: &StoredProviderCatalogKey, api_format:
.unwrap_or(false)
}
pub fn is_provider_key_circuit_open_at(
key: &StoredProviderCatalogKey,
api_format: &str,
now_unix_secs: u64,
) -> bool {
let Some(payload) = key
.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|values| values.get(api_format))
.and_then(serde_json::Value::as_object)
else {
return false;
};
if !payload
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return false;
}
payload
.get("next_probe_at_unix_secs")
.and_then(serde_json::Value::as_u64)
.is_none_or(|next_probe_at| now_unix_secs < next_probe_at)
}
fn available_provider_key_rpm_slots_for_new_user(
key: &StoredProviderCatalogKey,
current_usage: usize,

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

@@ -31,8 +31,9 @@ pub use health::{
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,
is_provider_key_circuit_open_at, 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,