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: pass session affinity through planner inputs

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

10
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
View File

@@ -3,7 +3,7 @@ use aether_ai_serving::{
};
use aether_scheduler_core::{
enumerate_minimal_candidate_selection_with_model_directives, normalize_api_format,
resolve_requested_global_model_name_with_model_directives,
resolve_requested_global_model_name_with_model_directives, ClientSessionAffinity,
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
};
use async_trait::async_trait;
@@ -32,6 +32,7 @@ struct GatewayLocalCandidatePreselectionPort<'a> {
require_streaming: bool,
required_capabilities: Option<&'a serde_json::Value>,
auth_snapshot: &'a GatewayAuthApiKeySnapshot,
client_session_affinity: Option<&'a ClientSessionAffinity>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
candidate_api_formats: Vec<String>,
@@ -73,6 +74,7 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
self.require_streaming,
self.required_capabilities,
auth_snapshot,
self.client_session_affinity,
current_unix_secs(),
)
.await?;
@@ -139,6 +141,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: &GatewayAuthApiKeySnapshot,
client_session_affinity: Option<&ClientSessionAffinity>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
) -> Result<
@@ -173,6 +176,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
use_api_format_alias_match,
key_mode,
candidate_api_formats,
@@ -189,6 +193,7 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> {
require_streaming: bool,
required_capabilities: Option<serde_json::Value>,
auth_snapshot: GatewayAuthApiKeySnapshot,
client_session_affinity: Option<ClientSessionAffinity>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
candidate_api_formats: Vec<String>,
@@ -210,6 +215,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: &GatewayAuthApiKeySnapshot,
client_session_affinity: Option<&ClientSessionAffinity>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
) -> Self {
@@ -239,6 +245,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
require_streaming,
required_capabilities: required_capabilities.cloned(),
auth_snapshot: auth_snapshot.clone(),
client_session_affinity: client_session_affinity.cloned(),
use_api_format_alias_match,
key_mode,
candidate_api_formats,
@@ -449,6 +456,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
candidates,
self.required_capabilities.as_ref(),
auth_snapshot,
self.client_session_affinity.as_ref(),
current_unix_secs(),
)
.await?;

3
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
View File

@@ -1,4 +1,5 @@
use aether_ai_serving::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
use aether_scheduler_core::ClientSessionAffinity;
use async_trait::async_trait;
use crate::ai_serving::{ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, PlannerAppState};
@@ -19,6 +20,7 @@ pub(crate) struct LocalRequestedModelDecisionInput {
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
pub(crate) required_capabilities: Option<serde_json::Value>,
pub(crate) request_auth_channel: Option<String>,
pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
}
#[derive(Debug, Clone)]
@@ -95,6 +97,7 @@ pub(crate) fn build_local_requested_model_decision_input(
auth_snapshot: resolved_input.auth_snapshot,
required_capabilities: resolved_input.required_capabilities,
request_auth_channel: None,
client_session_affinity: None,
}
}

8
apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs
View File

@@ -1,4 +1,4 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
use std::time::Duration;
use tokio::time::Instant;
@@ -18,6 +18,7 @@ impl<'a> PlannerAppState<'a> {
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let enable_model_directives =
@@ -35,6 +36,7 @@ impl<'a> PlannerAppState<'a> {
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
)
@@ -48,6 +50,7 @@ impl<'a> PlannerAppState<'a> {
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
) -> Result<
(
@@ -77,6 +80,7 @@ impl<'a> PlannerAppState<'a> {
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
attempt_now_unix_secs,
enable_model_directives,
)
@@ -106,6 +110,7 @@ impl<'a> PlannerAppState<'a> {
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
) -> Result<
(
@@ -121,6 +126,7 @@ impl<'a> PlannerAppState<'a> {
candidates,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
)
.await