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

Track scheduler affinity epochs and key sorting

Browse Source
This commit is contained in:
fawney19
2026-05-11 01:45:49 +08:00
parent 7b81c77424
commit e3574e1918
38 changed files with 656 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
apps/aether-gateway/src/ai_serving/planner/candidate_affinity_cache.rs
View File

@@ -40,6 +40,26 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
client_api_format: &str,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) {
remember_scheduler_affinity_for_candidate_at_epoch(
state,
auth_snapshot,
client_session_affinity,
client_api_format,
requested_model,
candidate,
None,
);
}
pub(crate) fn remember_scheduler_affinity_for_candidate_at_epoch(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
client_api_format: &str,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
expected_epoch: Option<u64>,
) {
let Some(api_key_id) = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
@@ -56,7 +76,7 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
return;
};
state.app().remember_scheduler_affinity_target(
let _ = state.app().remember_scheduler_affinity_target_for_epoch(
&cache_key,
SchedulerAffinityTarget {
provider_id: candidate.provider_id.clone(),
@@ -65,5 +85,6 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
},
SCHEDULER_AFFINITY_TTL,
PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES,
expected_epoch,
);
}

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

@@ -15,7 +15,7 @@ use std::sync::Arc;
use tracing::warn;
use uuid::Uuid;
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate_at_epoch;
use crate::ai_serving::planner::candidate_resolution::{
resolve_and_rank_logical_local_execution_candidates, EligibleLocalExecutionCandidate,
LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
@@ -810,13 +810,14 @@ pub(crate) fn remember_first_local_candidate_affinity(
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(first_candidate.candidate.global_model_name.as_str());
remember_scheduler_affinity_for_candidate(
remember_scheduler_affinity_for_candidate_at_epoch(
state,
auth_snapshot,
client_session_affinity,
client_api_format,
affinity_requested_model,
&first_candidate.candidate,
first_candidate.orchestration.scheduler_affinity_epoch,
);
}
@@ -1324,6 +1325,7 @@ mod tests {
candidate_group_id: pool_key_index.map(|_| "pool-group".to_string()),
pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch: None,
},
ranking: None,
}

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

@@ -332,6 +332,7 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let scheduler_affinity_epoch = state.app().scheduler_affinity_epoch();
let port = GatewayLocalCandidateResolutionPort {
state,
requested_model,
@@ -350,7 +351,12 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
};
match run_ai_candidate_resolution(&port, candidates, request).await {
Ok(outcome) => (outcome.eligible_candidates, outcome.skipped_candidates),
Ok(mut outcome) => {
for candidate in &mut outcome.eligible_candidates {
candidate.orchestration.scheduler_affinity_epoch = Some(scheduler_affinity_epoch);
}
(outcome.eligible_candidates, outcome.skipped_candidates)
}
Err(error) => match error {},
}
}

1
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs
View File

@@ -119,6 +119,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)

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

@@ -938,10 +938,12 @@ fn apply_ai_pool_orchestration(
mut candidate: EligibleLocalExecutionCandidate,
orchestration: AiPoolCandidateOrchestration,
) -> EligibleLocalExecutionCandidate {
let scheduler_affinity_epoch = candidate.orchestration.scheduler_affinity_epoch;
candidate.orchestration = LocalExecutionCandidateMetadata {
candidate_group_id: orchestration.candidate_group_id,
pool_key_index: orchestration.pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch,
};
candidate
}
@@ -1127,6 +1129,7 @@ mod tests {
),
pool_key_index: Some(0),
pool_key_lease: None,
scheduler_affinity_epoch: None,
}
);
assert_eq!(reordered[1].orchestration.pool_key_index, Some(1));
@@ -1143,6 +1146,7 @@ mod tests {
),
pool_key_index: None,
pool_key_lease: None,
scheduler_affinity_epoch: None,
}
);
}

15
apps/aether-gateway/src/ai_serving/planner/report_context.rs
View File

@@ -18,7 +18,10 @@ use crate::ai_serving::{
use crate::client_session_affinity::{
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
use crate::orchestration::{insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity};
use crate::orchestration::{
insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity,
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
};
pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
@@ -52,6 +55,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) original_request_body_json: Option<&'a Value>,
pub(crate) original_request_body_base64: Option<&'a str>,
pub(crate) client_session_affinity: Option<&'a ClientSessionAffinity>,
pub(crate) scheduler_affinity_epoch: Option<u64>,
pub(crate) client_requested_stream: bool,
pub(crate) upstream_is_stream: bool,
pub(crate) has_envelope: bool,
@@ -89,6 +93,12 @@ pub(crate) fn build_local_execution_report_context(
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
insert_pool_key_lease_report_context_fields(&mut extra_fields, parts.pool_key_lease);
if let Some(epoch) = parts.scheduler_affinity_epoch {
extra_fields.insert(
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD.to_string(),
Value::Number(epoch.into()),
);
}
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
@@ -278,6 +288,7 @@ mod tests {
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
client_session_affinity: Some(&client_session_affinity),
scheduler_affinity_epoch: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
@@ -359,6 +370,7 @@ mod tests {
})),
original_request_body_base64: None,
client_session_affinity: None,
scheduler_affinity_epoch: None,
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
@@ -424,6 +436,7 @@ mod tests {
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
client_session_affinity: None,
scheduler_affinity_epoch: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,

1
apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs
View File

@@ -101,6 +101,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
original_request_body_json: Some(body_json),
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream: spec_metadata.require_streaming,
has_envelope: false,

1
apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs
View File

@@ -121,6 +121,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,

1
apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs
View File

@@ -82,6 +82,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,

1
apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs
View File

@@ -127,6 +127,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)

1
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
View File

@@ -116,6 +116,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)

1
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs
View File

@@ -112,6 +112,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)