refactor: 抽离 AI pipeline 与调度共享能力逻辑

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

@@ -4,13 +4,16 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
 use aether_scheduler_core::{
     collect_selectable_candidates_from_keys,
     reorder_candidates_by_scheduler_health as reorder_candidates_by_scheduler_health_in_core,
+    SchedulerPriorityMode,
 };
 
 use crate::data::auth::GatewayAuthApiKeySnapshot;
 use crate::data::candidate_selection::{
-    read_minimal_candidate_selection, MinimalCandidateSelectionRowSource,
+    read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities,
+    MinimalCandidateSelectionRowSource,
 };
 use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
+use crate::scheduler::config::SchedulerSchedulingMode;
 use crate::GatewayError;
 
 use super::affinity::{
@@ -25,18 +28,66 @@ use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
 pub(super) fn reorder_candidates_by_scheduler_health(
     candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
     provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
-    auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+    required_capabilities: Option<&serde_json::Value>,
+    affinity_key: Option<&str>,
+    priority_mode: SchedulerPriorityMode,
 ) {
-    let affinity_key = auth_snapshot
-        .map(|snapshot| snapshot.api_key_id.trim())
-        .filter(|value| !value.is_empty());
     reorder_candidates_by_scheduler_health_in_core(
         candidates,
         provider_key_rpm_states,
+        required_capabilities,
         affinity_key,
+        priority_mode,
     );
 }
 
+fn apply_load_balance_rotation(
+    candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
+    priority_mode: SchedulerPriorityMode,
+    now_unix_secs: u64,
+) {
+    if candidates.len() < 2 {
+        return;
+    }
+
+    let mut start = 0usize;
+    while start < candidates.len() {
+        let mut end = start + 1;
+        while end < candidates.len()
+            && candidates_share_load_balance_group(
+                &candidates[start],
+                &candidates[end],
+                priority_mode,
+            )
+        {
+            end += 1;
+        }
+
+        let group_len = end - start;
+        if group_len > 1 {
+            let offset = usize::try_from(now_unix_secs).unwrap_or(0) % group_len;
+            candidates[start..end].rotate_left(offset);
+        }
+        start = end;
+    }
+}
+
+fn candidates_share_load_balance_group(
+    left: &SchedulerMinimalCandidateSelectionCandidate,
+    right: &SchedulerMinimalCandidateSelectionCandidate,
+    priority_mode: SchedulerPriorityMode,
+) -> bool {
+    match priority_mode {
+        SchedulerPriorityMode::Provider => {
+            left.provider_priority == right.provider_priority
+                && left.key_internal_priority == right.key_internal_priority
+        }
+        SchedulerPriorityMode::GlobalKey => {
+            left.key_global_priority_for_format == right.key_global_priority_for_format
+        }
+    }
+}
+
 #[cfg_attr(not(test), allow(dead_code))]
 pub(super) async fn select_minimal_candidate(
     selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
@@ -44,6 +95,7 @@ pub(super) async fn select_minimal_candidate(
     api_format: &str,
     global_model_name: &str,
     require_streaming: bool,
+    required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
     now_unix_secs: u64,
 ) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -55,6 +107,7 @@ pub(super) async fn select_minimal_candidate(
         api_format,
         global_model_name,
         require_streaming,
+        required_capabilities,
         auth_snapshot,
         now_unix_secs,
     )
@@ -73,31 +126,54 @@ pub(super) async fn collect_selectable_candidates(
     api_format: &str,
     global_model_name: &str,
     require_streaming: bool,
+    required_capabilities: Option<&serde_json::Value>,
     auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
     now_unix_secs: u64,
 ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
-    let mut candidates = read_minimal_candidate_selection(
+    let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
+    let priority_affinity_key =
+        scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
+    let mut candidates =
+        read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
         selection_row_source,
         api_format,
         global_model_name,
         require_streaming,
         auth_snapshot,
+        ordering_config.priority_mode,
+        priority_affinity_key,
+        required_capabilities,
     )
-    .await
-    .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        .await
+        .map_err(|err| GatewayError::Internal(err.to_string()))?;
     let runtime_snapshot =
         read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
             .await?;
     reorder_candidates_by_scheduler_health(
         &mut candidates,
         &runtime_snapshot.provider_key_rpm_states,
-        auth_snapshot,
+        required_capabilities,
+        priority_affinity_key,
+        ordering_config.priority_mode,
     );
+    if ordering_config.scheduling_mode == SchedulerSchedulingMode::LoadBalance {
+        apply_load_balance_rotation(
+            &mut candidates,
+            ordering_config.priority_mode,
+            now_unix_secs,
+        );
+    }
     let affinity_cache_key =
         build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
-    let cached_affinity_target = affinity_cache_key.as_deref().and_then(|cache_key| {
-        runtime_state.read_cached_scheduler_affinity_target(cache_key, SCHEDULER_AFFINITY_TTL)
-    });
+    let cached_affinity_target = if ordering_config.scheduling_mode
+        == SchedulerSchedulingMode::CacheAffinity
+    {
+        affinity_cache_key.as_deref().and_then(|cache_key| {
+            runtime_state.read_cached_scheduler_affinity_target(cache_key, SCHEDULER_AFFINITY_TTL)
+        })
+    } else {
+        None
+    };
 
     if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) {
         return Ok(Vec::new());
@@ -123,3 +199,16 @@ pub(super) async fn collect_selectable_candidates(
         cached_affinity_target.as_ref(),
     ))
 }
+
+fn scheduling_priority_affinity_key<'a>(
+    auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
+    scheduling_mode: SchedulerSchedulingMode,
+) -> Option<&'a str> {
+    if scheduling_mode == SchedulerSchedulingMode::FixedOrder {
+        return None;
+    }
+
+    auth_snapshot
+        .map(|snapshot| snapshot.api_key_id.trim())
+        .filter(|value| !value.is_empty())
+}