Track scheduler affinity epochs and key sorting

apps/aether-gateway/src/scheduler/candidate/affinity.rs

@@ -47,13 +47,14 @@ pub(super) fn remember_scheduler_affinity(
     affinity_cache_key: Option<&str>,
     state: &(impl SchedulerRuntimeState + ?Sized),
     candidate: &SchedulerMinimalCandidateSelectionCandidate,
+    expected_epoch: Option<u64>,
 ) {
     let Some(cache_key) = affinity_cache_key else {
         return;
     };
     let (provider_id, endpoint_id, key_id) = candidate_key(candidate);
 
-    state.remember_scheduler_affinity_target(
+    let _ = state.remember_scheduler_affinity_target_for_epoch(
         cache_key,
         SchedulerAffinityTarget {
             provider_id,
@@ -62,5 +63,6 @@ pub(super) fn remember_scheduler_affinity(
         },
         SCHEDULER_AFFINITY_TTL,
         SCHEDULER_AFFINITY_MAX_ENTRIES,
+        expected_epoch,
     );
 }

apps/aether-gateway/src/scheduler/candidate/selection.rs

@@ -46,6 +46,7 @@ pub(super) async fn select_minimal_candidate(
     now_unix_secs: u64,
     enable_model_directives: bool,
 ) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
+    let affinity_epoch = runtime_state.scheduler_affinity_epoch();
     let affinity_cache_key = build_scheduler_affinity_cache_key(
         auth_snapshot,
         api_format,
@@ -68,7 +69,12 @@ pub(super) async fn select_minimal_candidate(
     .into_iter()
     .next();
     if let Some(candidate) = selected.as_ref() {
-        remember_scheduler_affinity(affinity_cache_key.as_deref(), runtime_state, candidate);
+        remember_scheduler_affinity(
+            affinity_cache_key.as_deref(),
+            runtime_state,
+            candidate,
+            Some(affinity_epoch),
+        );
     }
     Ok(selected)
 }

apps/aether-gateway/src/scheduler/state.rs

@@ -42,6 +42,8 @@ pub(crate) trait SchedulerRuntimeState {
         ttl: Duration,
     ) -> Option<SchedulerAffinityTarget>;
 
+    fn scheduler_affinity_epoch(&self) -> u64;
+
     fn remember_scheduler_affinity_target(
         &self,
         cache_key: &str,
@@ -50,6 +52,15 @@ pub(crate) trait SchedulerRuntimeState {
         max_entries: usize,
     );
 
+    fn remember_scheduler_affinity_target_for_epoch(
+        &self,
+        cache_key: &str,
+        target: SchedulerAffinityTarget,
+        ttl: Duration,
+        max_entries: usize,
+        expected_epoch: Option<u64>,
+    ) -> bool;
+
     async fn read_scheduler_ordering_config(&self)
         -> Result<SchedulerOrderingConfig, GatewayError>;
 }