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

feat: scope planner affinity cache by session

Browse Source
This commit is contained in:
RWDai
2026-05-05 11:23:54 +08:00
parent bcd542f7ee
commit 58d6b2add6
2 changed files with 75 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
apps/aether-gateway/src/ai_serving/planner/candidate_affinity_cache.rs
View File

@@ -1,6 +1,6 @@
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
SchedulerMinimalCandidateSelectionCandidate,
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session, ClientSessionAffinity,
SchedulerAffinityTarget, SchedulerMinimalCandidateSelectionCandidate,
};
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
@@ -11,6 +11,7 @@ const PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES: usize = 10_000;
pub(crate) fn read_cached_scheduler_affinity_target(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
client_api_format: &str,
requested_model: Option<&str>,
) -> Option<SchedulerAffinityTarget> {
@@ -20,10 +21,11 @@ pub(crate) fn read_cached_scheduler_affinity_target(
let api_key_id = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty())?;
let cache_key = build_scheduler_affinity_cache_key_for_api_key_id(
let cache_key = build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
api_key_id,
client_api_format,
requested_model,
client_session_affinity,
)?;
state
@@ -34,6 +36,7 @@ pub(crate) fn read_cached_scheduler_affinity_target(
pub(crate) fn remember_scheduler_affinity_for_candidate(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
client_api_format: &str,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
@@ -44,10 +47,11 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
else {
return;
};
let Some(cache_key) = build_scheduler_affinity_cache_key_for_api_key_id(
let Some(cache_key) = build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
api_key_id,
client_api_format,
requested_model,
client_session_affinity,
) else {
return;
};

70
apps/aether-gateway/src/ai_serving/planner/candidate_ranking.rs
View File

@@ -14,8 +14,9 @@ use crate::scheduler::config::{
read_scheduler_ordering_config, SchedulerOrderingConfig, SchedulerSchedulingMode,
};
use aether_scheduler_core::{
matches_affinity_target, SchedulerAffinityTarget, SchedulerMinimalCandidateSelectionCandidate,
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingOutcome,
matches_affinity_target, ClientSessionAffinity, SchedulerAffinityTarget,
SchedulerMinimalCandidateSelectionCandidate, SchedulerRankableCandidate,
SchedulerRankingContext, SchedulerRankingOutcome,
};
use super::candidate_affinity_cache::read_cached_scheduler_affinity_target;
@@ -28,6 +29,7 @@ struct GatewayLocalCandidateRankingPort<'a> {
state: PlannerAppState<'a>,
requested_model: Option<&'a str>,
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&'a ClientSessionAffinity>,
required_capabilities: Option<&'a serde_json::Value>,
ordering_config: SchedulerOrderingConfig,
}
@@ -58,6 +60,7 @@ impl AiCandidateRankingPort for GatewayLocalCandidateRankingPort<'_> {
Ok(read_cached_scheduler_affinity_target(
self.state,
self.auth_snapshot,
self.client_session_affinity,
normalized_client_api_format,
affinity_requested_model,
))
@@ -116,6 +119,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
normalized_client_api_format: &str,
requested_model: Option<&str>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&serde_json::Value>,
) -> Vec<EligibleLocalExecutionCandidate> {
let ordering_config = read_scheduler_ordering_config_or_default(state).await;
@@ -123,6 +127,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
state,
requested_model,
auth_snapshot,
client_session_affinity,
required_capabilities,
ordering_config,
};
@@ -212,7 +217,9 @@ mod tests {
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_scheduler_core::{
apply_scheduler_candidate_ranking, RANKING_REASON_CACHED_AFFINITY,
apply_scheduler_candidate_ranking,
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
ClientSessionAffinity, RANKING_REASON_CACHED_AFFINITY,
};
use serde_json::json;
@@ -1052,6 +1059,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1130,6 +1138,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1204,6 +1213,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1269,6 +1279,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1350,6 +1361,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1398,6 +1410,7 @@ mod tests {
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
None,
"openai:chat",
"gpt-4.1",
&cached_candidate,
@@ -1422,6 +1435,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1498,6 +1512,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1567,6 +1582,7 @@ mod tests {
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
None,
"claude:messages",
"gpt-4.1",
&cached_cross_format,
@@ -1591,6 +1607,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1660,6 +1677,7 @@ mod tests {
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
None,
"openai:chat",
"gpt-4.1",
&cached_candidate,
@@ -1692,6 +1710,7 @@ mod tests {
None,
None,
None,
None,
)
.await;
@@ -1762,6 +1781,7 @@ mod tests {
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
None,
"openai:chat",
"gpt-4.1",
&cached_candidate,
@@ -1786,6 +1806,7 @@ mod tests {
None,
None,
None,
None,
aether_ai_serving::AiCandidateResolutionMode::Standard,
)
.await;
@@ -1852,6 +1873,7 @@ mod tests {
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
None,
"openai:chat",
"gpt-4.1",
&cached_candidate,
@@ -1876,6 +1898,7 @@ mod tests {
None,
None,
None,
None,
aether_ai_serving::AiCandidateResolutionMode::Standard,
)
.await;
@@ -1902,6 +1925,7 @@ mod tests {
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
None,
"openai:chat",
"gpt-5",
&candidate,
@@ -1917,4 +1941,44 @@ mod tests {
assert_eq!(remembered.endpoint_id, "endpoint-1");
assert_eq!(remembered.key_id, "key-1");
}
#[tokio::test]
async fn remembers_scheduler_affinity_for_client_session_scope() {
let state = AppState::new().expect("state should build");
let auth_snapshot = sample_auth_snapshot();
let client_session_affinity = ClientSessionAffinity::new(
Some("generic".to_string()),
Some("session=conversation-1;agent=coder".to_string()),
);
let candidate = sample_candidate("endpoint-session", "key-session");
remember_scheduler_affinity_for_candidate(
PlannerAppState::new(&state),
Some(&auth_snapshot),
Some(&client_session_affinity),
"openai:chat",
"gpt-5",
&candidate,
);
let session_key = build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
"api-key-1",
"openai:chat",
"gpt-5",
Some(&client_session_affinity),
)
.expect("session key should build");
let remembered = state
.read_scheduler_affinity_target(&session_key, SCHEDULER_AFFINITY_TTL)
.expect("session affinity target should be cached");
assert_eq!(remembered.provider_id, "provider-1");
assert_eq!(remembered.endpoint_id, "endpoint-session");
assert_eq!(remembered.key_id, "key-session");
assert!(state
.read_scheduler_affinity_target(
"scheduler_affinity:api-key-1:openai:chat:gpt-5",
SCHEDULER_AFFINITY_TTL,
)
.is_none());
}
}