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

Restrict scheduler affinity to cache affinity mode

Browse Source
This commit is contained in:
fawney19
2026-05-11 14:06:49 +08:00
parent e91c874863
commit 247ea9d1bd
12 changed files with 515 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
View File

@@ -31,6 +31,7 @@ use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_ms;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
use crate::{AppState, GatewayError};
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
@@ -182,6 +183,7 @@ struct GatewayLocalCandidateMaterializationPort<'a, F, G> {
request_auth_channel: Option<&'a str>,
persistence_policy: LocalCandidatePersistencePolicy<'a>,
resolution_mode: LocalCandidateResolutionMode,
scheduler_cache_affinity_enabled: bool,
build_available_extra_data: F,
decorate_skipped_candidate: G,
}
@@ -244,6 +246,9 @@ where
}
fn remember_first_candidate_affinity(&self, candidates: &[Self::Eligible]) {
if !self.scheduler_cache_affinity_enabled {
return;
}
remember_first_local_candidate_affinity(
self.state,
self.auth_snapshot,
@@ -430,6 +435,7 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
{
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
let port = GatewayLocalCandidateMaterializationPort {
state,
trace_id,
@@ -442,6 +448,7 @@ where
request_auth_channel,
persistence_policy,
resolution_mode,
scheduler_cache_affinity_enabled,
build_available_extra_data,
decorate_skipped_candidate,
};
@@ -474,6 +481,7 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
{
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
let _ = build_available_extra_data;
let (candidates, resolved_skipped) = resolve_and_rank_logical_local_execution_candidates(
state,
@@ -505,14 +513,16 @@ where
}
}
remember_first_local_candidate_affinity(
state,
auth_snapshot,
client_session_affinity,
client_api_format,
requested_model,
&candidates,
);
if scheduler_cache_affinity_enabled {
remember_first_local_candidate_affinity(
state,
auth_snapshot,
client_session_affinity,
client_api_format,
requested_model,
&candidates,
);
}
let (items, _) = build_logical_candidate_items(
state,
@@ -597,6 +607,7 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync + 'a,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync + 'a,
{
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
let _ = build_available_extra_data;
let decorate_skipped_candidate = Arc::new(decorate_skipped_candidate);
let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic;
@@ -630,6 +641,7 @@ where
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
scheduler_cache_affinity_enabled,
};
cursor.load_next_page().await;
let candidate_count = cursor.candidate_count;
@@ -665,6 +677,7 @@ struct RequestedModelAttemptPageCursor<'a> {
candidate_count: usize,
next_candidate_index: u32,
remembered_affinity: bool,
scheduler_cache_affinity_enabled: bool,
}
impl<'a> RequestedModelAttemptPageCursor<'a> {
@@ -726,7 +739,10 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
);
}
}
if !self.remembered_affinity && !candidates.is_empty() {
if self.scheduler_cache_affinity_enabled
&& !self.remembered_affinity
&& !candidates.is_empty()
{
remember_first_local_candidate_affinity(
self.state,
Some(&self.auth_snapshot),
@@ -795,6 +811,21 @@ async fn pop_attempt_from_items(
}
}
async fn scheduler_cache_affinity_enabled(state: PlannerAppState<'_>) -> bool {
match read_scheduler_ordering_config(state.app()).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
Err(error) => {
warn!(
event_name = "planner_scheduler_affinity_config_load_failed",
log_type = "event",
error = ?error,
"failed to load scheduler config while checking cache affinity mode"
);
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
}
}
}
pub(crate) fn remember_first_local_candidate_affinity(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
@@ -1211,6 +1242,7 @@ mod tests {
use std::sync::Arc;
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use aether_data::repository::auth::StoredAuthApiKeySnapshot;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
@@ -1219,6 +1251,7 @@ mod tests {
GatewayProviderTransportProvider,
};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankingMode,
SchedulerRankingOutcome,
};
@@ -1227,6 +1260,7 @@ mod tests {
use super::*;
use crate::data::GatewayDataState;
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
fn sample_candidate(key_id: &str) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
@@ -1331,6 +1365,46 @@ mod tests {
}
}
fn sample_auth_snapshot() -> GatewayAuthApiKeySnapshot {
GatewayAuthApiKeySnapshot::from_stored(
StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"api-key-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
None,
None,
None,
)
.expect("stored auth snapshot should build"),
0,
)
}
fn no_extra_data(_: &EligibleLocalExecutionCandidate) -> Option<Value> {
None
}
fn identity_skipped_candidate(
candidate: SkippedLocalExecutionCandidate,
) -> SkippedLocalExecutionCandidate {
candidate
}
#[tokio::test]
async fn pool_group_keys_are_not_persisted_as_available_before_attempt() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
@@ -1368,6 +1442,55 @@ mod tests {
assert_eq!(stored[0].candidate_index, 2);
}
#[test]
fn materialization_port_ignores_scheduler_affinity_when_cache_affinity_disabled() {
let app = AppState::new().expect("state should build");
let auth_snapshot = sample_auth_snapshot();
let port = GatewayLocalCandidateMaterializationPort {
state: PlannerAppState::new(&app),
trace_id: "trace-affinity-disabled",
client_api_format: "openai:chat",
requested_model: Some("gpt-5"),
auth_snapshot: Some(&auth_snapshot),
client_session_affinity: None,
required_capabilities: None,
sticky_session_token: None,
request_auth_channel: None,
persistence_policy: LocalCandidatePersistencePolicy {
available: LocalAvailableCandidatePersistenceContext {
user_id: "user-1",
api_key_id: "api-key-1",
required_capabilities: None,
error_context: "test available",
},
skipped: LocalSkippedCandidatePersistenceContext {
user_id: "user-1",
api_key_id: "api-key-1",
required_capabilities: None,
error_context: "test skipped",
record_runtime_miss_diagnostic: false,
},
},
resolution_mode: LocalCandidateResolutionMode::Standard,
scheduler_cache_affinity_enabled: false,
build_available_extra_data: no_extra_data,
decorate_skipped_candidate: identity_skipped_candidate,
};
let candidate = sample_eligible("key-a", None);
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");
aether_ai_serving::AiCandidateMaterializationPort::remember_first_candidate_affinity(
&port,
&[candidate],
);
assert!(app
.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
.is_none());
}
#[tokio::test]
async fn logical_materialization_does_not_persist_pool_group_representative() {
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());

11
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
View File

@@ -23,8 +23,10 @@ use crate::ai_serving::{
candidate_common_transport_skip_reason, CandidateTransportPolicyFacts, PlannerAppState,
};
use crate::clock::current_unix_ms;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::handlers::shared::provider_pool::read_admin_provider_pool_runtime_state;
use crate::handlers::shared::provider_pool::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
};
use crate::handlers::shared::provider_pool::{
try_claim_admin_provider_pool_key, AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
};
@@ -306,6 +308,9 @@ impl<'a> PoolKeyCursor<'a> {
async fn sticky_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
let pool_config = pool_config_for_candidate(&self.group)?;
if !admin_provider_pool_cache_affinity_enabled(&pool_config) {
return None;
}
let runtime = read_admin_provider_pool_runtime_state(
self.state.app().runtime_state.as_ref(),
self.group.candidate.provider_id.as_str(),
@@ -1206,7 +1211,7 @@ mod tests {
}
#[test]
fn pool_scheduler_promotes_sticky_hit_regardless_distribution_mode() {
fn pool_scheduler_ignores_sticky_hit_without_cache_affinity() {
let key_a = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
@@ -1250,7 +1255,7 @@ mod tests {
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-a", "key-b"]
vec!["key-b", "key-a"]
);
}

52
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -215,6 +215,28 @@ fn admin_provider_pool_lru_enabled(
.any(|item| item.enabled && item.preset.eq_ignore_ascii_case("lru"))
}
pub(crate) fn admin_provider_pool_cache_affinity_enabled(
pool_config: &AdminProviderPoolConfig,
) -> bool {
let mut seen = std::collections::BTreeSet::new();
for item in &pool_config.scheduling_presets {
let preset = item.preset.trim().to_ascii_lowercase();
if preset.is_empty() || !seen.insert(preset.clone()) {
continue;
}
if !item.enabled {
continue;
}
if matches!(
preset.as_str(),
"lru" | "cache_affinity" | "load_balance" | "single_account"
) {
return preset == "cache_affinity";
}
}
false
}
pub(crate) fn admin_provider_pool_config(
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
) -> Option<AdminProviderPoolConfig> {
@@ -331,7 +353,10 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
#[cfg(test)]
mod tests {
use super::{admin_provider_pool_config, admin_provider_pool_config_from_config_value};
use super::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config,
admin_provider_pool_config_from_config_value,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use serde_json::json;
@@ -551,4 +576,29 @@ mod tests {
assert_eq!(config.unschedulable_rules[1].keyword, "review_required");
assert_eq!(config.unschedulable_rules[1].duration_minutes, 5);
}
#[test]
fn cache_affinity_enabled_only_when_it_is_distribution_mode() {
let cache_affinity = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "cache_affinity", "enabled": true},
{"preset": "priority_first", "enabled": true}
]
}
})))
.expect("pool config should parse");
assert!(admin_provider_pool_cache_affinity_enabled(&cache_affinity));
let load_balance = admin_provider_pool_config_from_config_value(Some(&json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "load_balance", "enabled": true},
{"preset": "cache_affinity", "enabled": true}
]
}
})))
.expect("pool config should parse");
assert!(!admin_provider_pool_cache_affinity_enabled(&load_balance));
}
}

39
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -3,6 +3,7 @@ use super::keys::{
pool_cooldown_keys, pool_cost_keys, pool_latency_keys, pool_lru_key, pool_sticky_key,
pool_sticky_pattern,
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
};
@@ -44,11 +45,13 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
let cost_keys = pool_cost_keys(provider_id, key_ids);
let latency_keys = pool_latency_keys(provider_id, key_ids);
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
&& admin_provider_pool_cache_affinity_enabled(pool_config);
if let Some(sticky_session_token) = sticky_session_token
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
.filter(|_| sticky_sessions_enabled)
{
let sticky_key = pool_sticky_key(provider_id, sticky_session_token);
if let Ok(Some(bound_key_id)) = runtime.kv_get(&sticky_key).await {
@@ -77,22 +80,24 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
}
let sticky_keys = runtime
.scan_keys(&pool_sticky_pattern(provider_id), 200)
.await
.unwrap_or_default();
state.total_sticky_sessions = sticky_keys.len();
if !sticky_keys.is_empty() {
let raw_keys = sticky_keys
.iter()
.map(|key| runtime.strip_namespace(key).to_string())
.collect::<Vec<_>>();
if let Ok(values) = runtime.kv_get_many(&raw_keys).await {
for bound_key_id in values.into_iter().flatten() {
*state
.sticky_sessions_by_key
.entry(bound_key_id)
.or_insert(0) += 1;
if sticky_sessions_enabled {
let sticky_keys = runtime
.scan_keys(&pool_sticky_pattern(provider_id), 200)
.await
.unwrap_or_default();
state.total_sticky_sessions = sticky_keys.len();
if !sticky_keys.is_empty() {
let raw_keys = sticky_keys
.iter()
.map(|key| runtime.strip_namespace(key).to_string())
.collect::<Vec<_>>();
if let Ok(values) = runtime.kv_get_many(&raw_keys).await {
for bound_key_id in values.into_iter().flatten() {
*state
.sticky_sessions_by_key
.entry(bound_key_id)
.or_insert(0) += 1;
}
}
}
}

46
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
View File

@@ -2,6 +2,7 @@ use super::keys::{
pool_cooldown_index_key, pool_cooldown_key, pool_cost_key, pool_latency_key, pool_lru_key,
pool_sticky_key, pool_stream_timeout_key,
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolUnschedulableRule,
};
@@ -346,6 +347,7 @@ pub(crate) async fn record_admin_provider_pool_success(
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|_| pool_config.sticky_session_ttl_seconds > 0)
.filter(|_| admin_provider_pool_cache_affinity_enabled(pool_config))
{
let _ = runtime
.kv_set(
@@ -810,6 +812,50 @@ mod tests {
assert!(runtime.lru_score_by_key.contains_key("key-1"));
}
#[tokio::test]
async fn success_feedback_does_not_write_sticky_without_cache_affinity() {
let Some(redis) = start_managed_redis_or_skip().await else {
return;
};
let app = build_runner_app(redis.redis_url(), "pool_runtime_no_sticky_load_balance").await;
let runtime = app.runtime_state.as_ref();
let mut pool_config = sample_pool_config();
pool_config.scheduling_presets = vec![AdminProviderPoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
}];
pool_config.lru_enabled = false;
let key_ids = vec!["key-1".to_string()];
record_admin_provider_pool_success(
runtime,
"provider-1",
"key-1",
&pool_config,
Some("session-1"),
120,
Some(80),
)
.await;
let runtime = read_admin_provider_pool_runtime_state(
runtime,
"provider-1",
&key_ids,
&pool_config,
Some("session-1"),
)
.await;
assert_eq!(runtime.total_sticky_sessions, 0);
assert_eq!(runtime.sticky_bound_key_id, None);
assert_eq!(runtime.sticky_sessions_by_key.get("key-1"), None);
assert_eq!(runtime.cost_window_usage_by_key.get("key-1"), Some(&120));
assert_eq!(runtime.latency_avg_ms_by_key.get("key-1"), Some(&80.0));
assert!(runtime.lru_score_by_key.contains_key("key-1"));
}
#[tokio::test]
async fn error_feedback_respects_retry_after_for_rate_limits() {
let Some(redis) = start_managed_redis_or_skip().await else {

15
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -54,6 +54,7 @@ use crate::headers::{
should_skip_request_header,
};
use crate::router::RequestAdmissionError;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
use crate::{
AppState, FrontdoorUserRpmOutcome, GatewayError, GatewayFallbackMetricKind,
GatewayFallbackReason, LocalExecutionRuntimeMissDiagnostic,
@@ -297,6 +298,20 @@ async fn maybe_forward_public_request_to_tunnel_owner(
}) else {
return Ok(None);
};
let cache_affinity_enabled = match read_scheduler_ordering_config(state).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
Err(err) => {
warn!(
trace_id = %request_context.trace_id,
error = ?err,
"gateway failed to load scheduler config while checking tunnel affinity forwarding mode"
);
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
}
};
if !cache_affinity_enabled {
return Ok(None);
}
let Some(api_format) = decision
.auth_endpoint_signature
.as_deref()

4
apps/aether-gateway/src/handlers/shared/provider_pool.rs
View File

@@ -1,4 +1,6 @@
pub(crate) use super::super::admin::provider::pool::config::admin_provider_pool_config_from_config_value;
pub(crate) use super::super::admin::provider::pool::config::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
};
pub(crate) use super::super::admin::provider::pool::runtime::{
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_runtime_state,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,

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

@@ -31,6 +31,7 @@ use crate::handlers::shared::provider_pool::{
};
use crate::orchestration::local_execution_candidate_metadata_from_report_context;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
use crate::AppState;
#[derive(Debug, Clone, Copy)]
@@ -238,10 +239,28 @@ fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffi
})
}
fn remember_successful_local_scheduler_affinity(
async fn scheduler_cache_affinity_enabled(state: &AppState) -> bool {
match read_scheduler_ordering_config(state).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
Err(error) => {
warn!(
event_name = "orchestration_scheduler_affinity_config_load_failed",
log_type = "event",
error = ?error,
"failed to load scheduler config while checking cache affinity mode"
);
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
}
}
}
async fn remember_successful_local_scheduler_affinity(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
) {
if !scheduler_cache_affinity_enabled(state).await {
return;
}
let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) else {
return;
};
@@ -516,7 +535,7 @@ async fn record_health_success_effect(
context: LocalExecutionEffectContext<'_>,
_effect: LocalHealthSuccessEffect,
) {
remember_successful_local_scheduler_affinity(state, context);
remember_successful_local_scheduler_affinity(state, context).await;
let api_format = context.plan.provider_api_format.trim();
if api_format.is_empty() {
@@ -1362,6 +1381,41 @@ mod tests {
);
}
#[tokio::test]
async fn load_balance_success_does_not_remember_scheduler_affinity_cache() {
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
json!("load_balance"),
)]),
);
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!(state
.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL)
.is_none());
}
#[tokio::test]
async fn success_remembers_session_scoped_scheduler_affinity_cache() {
let state = AppState::new().expect("gateway state should build");

38
apps/aether-gateway/src/scheduler/candidate/selection.rs
View File

@@ -47,34 +47,50 @@ pub(super) async fn select_minimal_candidate(
enable_model_directives: bool,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let affinity_epoch = runtime_state.scheduler_affinity_epoch();
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let affinity_cache_key = build_scheduler_affinity_cache_key(
auth_snapshot,
api_format,
global_model_name,
client_session_affinity,
);
let selected = collect_selectable_candidates(
let priority_affinity_key =
scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
let candidates = enumerate_scheduler_candidates(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
)
.await?;
let selected = collect_selectable_enumerated_candidates_with_skip_reasons(
runtime_state,
api_format,
global_model_name,
candidates,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
ordering_config,
priority_affinity_key,
)
.await?
.0
.into_iter()
.next();
if let Some(candidate) = selected.as_ref() {
remember_scheduler_affinity(
affinity_cache_key.as_deref(),
runtime_state,
candidate,
Some(affinity_epoch),
);
if ordering_config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity {
if let Some(candidate) = selected.as_ref() {
remember_scheduler_affinity(
affinity_cache_key.as_deref(),
runtime_state,
candidate,
Some(affinity_epoch),
);
}
}
Ok(selected)
}

41
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -22,6 +22,7 @@ use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::data::GatewayDataState;
use crate::{AppState, GatewayError};
use super::super::affinity::build_scheduler_affinity_cache_key;
use super::super::runtime::should_skip_provider_quota;
use super::super::selection::{
collect_selectable_candidates as collect_selectable_candidates_impl,
@@ -605,6 +606,46 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
assert_eq!(selected.key_id, "key-b");
}
#[tokio::test]
async fn load_balance_selection_does_not_remember_scheduler_affinity() {
let row = sample_row();
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
row,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
.with_system_config_values_for_tests(vec![(
"scheduling_mode".to_string(),
json!("load_balance"),
)]),
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
let cache_key =
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None)
.expect("scheduler affinity cache key should build");
let selected = select_candidate(
state.data.as_ref(),
&state,
"openai:chat",
"gpt-4.1",
false,
Some(&auth_snapshot),
100,
)
.await
.expect("selection should succeed")
.expect("candidate should exist");
assert_eq!(selected.key_id, "key-1");
assert!(state
.read_scheduler_affinity_target(cache_key.as_str(), Duration::from_secs(300))
.is_none());
}
#[tokio::test]
async fn load_balance_ignores_provider_priority_and_cached_affinity() {
let mut first = sample_row();