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

fix scheduler affinity candidate selection

Browse Source
This commit is contained in:
fawney19
2026-04-30 16:27:24 +08:00
parent 558abfcfa3
commit 33aa70c22b
24 changed files with 1715 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,6 +4,7 @@ use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, count_recent_rpm_requests_for_provider_key,
SchedulerAffinityTarget,
};
use aether_usage_runtime::{
build_stream_terminal_usage_outcome, build_sync_terminal_usage_outcome,
@@ -24,6 +25,7 @@ use crate::handlers::shared::provider_pool::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success, AdminProviderPoolConfig,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
#[derive(Debug, Clone, Copy)]
@@ -96,6 +98,7 @@ struct PoolFeedbackContext {
}
const ADAPTIVE_RPM_RECENT_CANDIDATE_LIMIT: usize = 512;
const LOCAL_EXECUTION_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
pub(crate) async fn apply_local_execution_effect(
state: &AppState,
@@ -155,6 +158,40 @@ fn local_scheduler_affinity_cache_key(report_context: Option<&Value>) -> Option<
)
}
fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffinityTarget> {
let provider_id = plan.provider_id.trim();
let endpoint_id = plan.endpoint_id.trim();
let key_id = plan.key_id.trim();
if provider_id.is_empty() || endpoint_id.is_empty() || key_id.is_empty() {
return None;
}
Some(SchedulerAffinityTarget {
provider_id: provider_id.to_string(),
endpoint_id: endpoint_id.to_string(),
key_id: key_id.to_string(),
})
}
fn remember_successful_local_scheduler_affinity(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
) {
let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) else {
return;
};
let Some(target) = local_scheduler_affinity_target(context.plan) else {
return;
};
state.remember_scheduler_affinity_target(
&cache_key,
target,
SCHEDULER_AFFINITY_TTL,
LOCAL_EXECUTION_SCHEDULER_AFFINITY_MAX_ENTRIES,
);
}
fn pool_feedback_request_body<'a>(
plan: &'a ExecutionPlan,
report_context: Option<&'a Value>,
@@ -415,6 +452,8 @@ async fn record_health_success_effect(
context: LocalExecutionEffectContext<'_>,
_effect: LocalHealthSuccessEffect,
) {
remember_successful_local_scheduler_affinity(state, context);
let api_format = context.plan.provider_api_format.trim();
if api_format.is_empty() {
return;
@@ -1039,6 +1078,102 @@ mod tests {
.is_some());
}
#[tokio::test]
async fn success_remembers_scheduler_affinity_cache_for_final_candidate() {
let state = AppState::new().expect("gateway state should build");
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
})
);
}
#[tokio::test]
async fn fallback_success_rewarms_scheduler_affinity_after_failed_candidate_invalidates() {
let state = AppState::new().expect("gateway state should build");
let failed_plan = sample_plan();
let mut success_plan = sample_plan();
success_plan.provider_id = "prov-2".to_string();
success_plan.endpoint_id = "ep-2".to_string();
success_plan.key_id = "key-2".to_string();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key.clone(),
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
},
SCHEDULER_AFFINITY_TTL,
16,
);
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &failed_plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: 429,
classification: LocalFailoverClassification::RetryUpstreamFailure,
}),
)
.await;
assert!(state
.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
.is_none());
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &success_plan,
report_context: Some(&report_context),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(SchedulerAffinityTarget {
provider_id: "prov-2".to_string(),
endpoint_id: "ep-2".to_string(),
key_id: "key-2".to_string(),
})
);
}
#[test]
fn semantic_client_error_does_not_penalize_pool_feedback() {
assert!(!local_candidate_failure_should_record_pool_error(