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

feat: thread session affinity through scheduler selection

Browse Source
This commit is contained in:
RWDai
2026-05-05 11:23:54 +08:00
parent c596d82dec
commit eb8878695a
3 changed files with 34 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -24,7 +24,7 @@ use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_scheduler_core::{
candidate_model_names, candidate_supports_required_capability, matches_model_mapping,
normalize_api_format, resolve_provider_model_name, select_provider_model_name,
SchedulerMinimalCandidateSelectionCandidate,
ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate,
};
use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot};
use regex::Regex;
@@ -57,6 +57,7 @@ pub(crate) async fn list_selectable_candidates(
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -68,6 +69,7 @@ pub(crate) async fn list_selectable_candidates(
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
)
@@ -89,6 +91,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<
@@ -106,6 +109,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
)
@@ -120,6 +124,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
) -> Result<
(
@@ -138,6 +143,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
candidates,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
ordering_config,
priority_affinity_key,
@@ -218,6 +224,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
require_streaming,
required_capabilities.as_ref(),
auth_snapshot,
None,
now_unix_secs,
false,
)

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

@@ -3,6 +3,7 @@ use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::GatewayError;
use aether_scheduler_core::ClientSessionAffinity;
use super::affinity::{build_scheduler_affinity_cache_key, remember_scheduler_affinity};
use super::enumeration::enumerate_scheduler_candidates;
@@ -41,11 +42,16 @@ pub(super) async fn select_minimal_candidate(
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let affinity_cache_key = build_scheduler_affinity_cache_key(
auth_snapshot,
api_format,
global_model_name,
client_session_affinity,
);
let selected = collect_selectable_candidates(
selection_row_source,
runtime_state,
@@ -54,6 +60,7 @@ pub(super) async fn select_minimal_candidate(
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
)
@@ -74,6 +81,7 @@ pub(super) async fn collect_selectable_candidates(
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
@@ -85,6 +93,7 @@ pub(super) async fn collect_selectable_candidates(
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
)
@@ -100,6 +109,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<
@@ -129,6 +139,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
candidates,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
ordering_config,
priority_affinity_key,
@@ -144,6 +155,7 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
mut candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: crate::scheduler::config::SchedulerOrderingConfig,
priority_affinity_key: Option<&str>,
@@ -157,8 +169,12 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
let runtime_snapshot =
read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
.await?;
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let affinity_cache_key = build_scheduler_affinity_cache_key(
auth_snapshot,
api_format,
global_model_name,
client_session_affinity,
);
let cached_affinity_target = if ordering_config.scheduling_mode
== SchedulerSchedulingMode::CacheAffinity
{
@@ -191,12 +207,8 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
));
}
let (mut selected, skipped) = resolve_scheduler_candidate_selectability(
candidates,
&runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
);
let (mut selected, skipped) =
resolve_scheduler_candidate_selectability(candidates, &runtime_snapshot, now_unix_secs);
rank_scheduler_candidates(
&mut selected,
&runtime_snapshot,

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

@@ -47,6 +47,7 @@ async fn select_candidate(
require_streaming,
None,
auth_snapshot,
None,
now_unix_secs,
false,
)
@@ -70,6 +71,7 @@ async fn collect_selectable_candidates(
require_streaming,
None,
auth_snapshot,
None,
now_unix_secs,
false,
)
@@ -99,6 +101,7 @@ async fn collect_selectable_candidates_with_skip_reasons(
require_streaming,
None,
auth_snapshot,
None,
now_unix_secs,
false,
)
@@ -406,6 +409,7 @@ async fn scheduler_selection_prefers_required_capability_matches_before_priority
false,
Some(&required_capabilities),
None,
None,
100,
false,
)