feat: pass session scope into scheduler cache

This commit is contained in:
RWDai
2026-05-05 11:23:54 +08:00
parent 96dd3fa4ca
commit c596d82dec
2 changed files with 14 additions and 5 deletions

View File

@@ -1,7 +1,8 @@
use std::time::Duration;
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, SchedulerAffinityTarget,
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session, ClientSessionAffinity,
SchedulerAffinityTarget,
};
use super::state::SchedulerRuntimeState;
@@ -11,13 +12,15 @@ pub(crate) const SCHEDULER_AFFINITY_TTL: Duration = Duration::from_secs(300);
pub(crate) fn read_cached_scheduler_affinity_target(
state: &(impl SchedulerRuntimeState + ?Sized),
api_key_id: &str,
client_session_affinity: Option<&ClientSessionAffinity>,
api_format: &str,
global_model_name: &str,
) -> Option<SchedulerAffinityTarget> {
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,
api_format,
global_model_name,
client_session_affinity,
)?;
state.read_cached_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
}

View File

@@ -1,6 +1,6 @@
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
matches_affinity_target, SchedulerAffinityTarget,
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session, candidate_affinity_hash,
candidate_key, matches_affinity_target, ClientSessionAffinity, SchedulerAffinityTarget,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
@@ -15,11 +15,17 @@ pub(super) fn build_scheduler_affinity_cache_key(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
api_format: &str,
global_model_name: &str,
client_session_affinity: Option<&ClientSessionAffinity>,
) -> Option<String> {
let api_key_id = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty())?;
build_scheduler_affinity_cache_key_for_api_key_id(api_key_id, api_format, global_model_name)
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
api_key_id,
api_format,
global_model_name,
client_session_affinity,
)
}
pub(super) fn scheduler_candidate_affinity_hash(