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

refactor: isolate dispatch scheduling core

Browse Source
This commit is contained in:
fawney19
2026-05-12 13:15:41 +08:00
parent 81ee27cdea
commit fa73655134
50 changed files with 4467 additions and 3202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,14 +4,17 @@ use aether_ai_serving::{
run_ai_available_candidate_persistence, run_ai_candidate_materialization,
run_ai_skipped_candidate_persistence, AiAvailableCandidatePersistencePort,
AiCandidateMaterializationOutcome, AiCandidateMaterializationPort,
AiSkippedCandidatePersistencePort,
AiCandidatePreselectionOutcome, AiSkippedCandidatePersistencePort,
};
use aether_dispatch_core::{DispatchSequence, DispatchSequenceItem};
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
use async_trait::async_trait;
use serde_json::Value;
use std::collections::VecDeque;
use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Instant;
use tracing::warn;
use uuid::Uuid;
@@ -29,12 +32,16 @@ use crate::ai_serving::planner::runtime_miss::record_local_runtime_candidate_ski
use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_ms;
use crate::dispatch::refs::dispatch_ref_for_local_candidate;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
use crate::scheduler::candidate::API_KEY_CONCURRENCY_LIMIT_SKIP_REASON;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
use crate::{AppState, GatewayError};
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
const AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET: Duration = Duration::from_millis(100);
const AUTH_API_KEY_CONCURRENCY_RETRY_DELAY: Duration = Duration::from_millis(10);
#[derive(Debug, Clone)]
pub(crate) struct LocalExecutionCandidateAttempt {
@@ -61,12 +68,12 @@ pub(crate) trait LocalExecutionAttemptSource<T>: Send {
enum LocalExecutionCandidateAttemptSourceItem<'a> {
Static {
attempts: VecDeque<LocalExecutionCandidateAttempt>,
attempts: DispatchSequence<LocalExecutionCandidateAttempt>,
},
Pool {
cursor: PoolKeyCursor<'a>,
candidate_index: u32,
pending_attempts: VecDeque<LocalExecutionCandidateAttempt>,
pending_attempts: DispatchSequence<LocalExecutionCandidateAttempt>,
},
RequestedModelPage {
cursor: Box<RequestedModelAttemptPageCursor<'a>>,
@@ -80,7 +87,7 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
let mut items = VecDeque::new();
if !attempts.is_empty() {
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
attempts: VecDeque::from(attempts),
attempts: dispatch_sequence_from_attempts(attempts),
});
}
Self { items }
@@ -91,8 +98,8 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
let front = self.items.front_mut()?;
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if let Some(attempt) = attempts.pop_front() {
if attempts.is_empty() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(attempts) {
if dispatch_sequence_exhausted(attempts) {
self.items.pop_front();
}
return Some(attempt);
@@ -104,7 +111,7 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
candidate_index,
pending_attempts,
} => {
if let Some(attempt) = pending_attempts.pop_front() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
let Some(candidate) = cursor.next_key().await else {
@@ -113,9 +120,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
self.items.pop_front();
continue;
};
*pending_attempts = build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
)
.into(),
);
}
LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } => {
@@ -311,10 +321,7 @@ where
}
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
ai_candidate_extra_data_with_ranking(
(self.build_extra_data)(candidate),
candidate.ranking.as_ref(),
)
available_candidate_extra_data_with_dispatch_ref(candidate, &self.build_extra_data)
}
fn generate_candidate_id(&self) -> String {
@@ -506,15 +513,6 @@ where
.map(decorate_skipped_candidate)
.collect::<Vec<_>>();
let candidate_count = candidates.len() + skipped_candidate_count;
if persistence_policy.skipped.record_runtime_miss_diagnostic {
for skipped_candidate in &skipped_candidates {
record_local_runtime_candidate_skip_reason(
state.app(),
trace_id,
skipped_candidate.skip_reason,
);
}
}
if scheduler_cache_affinity_enabled {
remember_first_local_candidate_affinity(
@@ -526,6 +524,14 @@ where
&candidates,
);
}
persist_skipped_local_execution_candidates_with_context(
state.app(),
trace_id,
persistence_policy.skipped,
u32::try_from(candidates.len()).unwrap_or(u32::MAX),
skipped_candidates,
)
.await;
let (items, _) = build_logical_candidate_items(
state,
@@ -566,7 +572,9 @@ fn build_logical_candidate_items<'a>(
candidate_index,
);
if !attempts.is_empty() {
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static { attempts });
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
attempts: dispatch_sequence_from_attempts(attempts.into()),
});
}
}
LocalExecutionCandidateKind::PoolGroup => {
@@ -585,7 +593,7 @@ fn build_logical_candidate_items<'a>(
items.push_back(LocalExecutionCandidateAttemptSourceItem::Pool {
cursor,
candidate_index,
pending_attempts: VecDeque::new(),
pending_attempts: DispatchSequence::new(Vec::new()),
});
}
}
@@ -646,6 +654,10 @@ where
required_capabilities: required_capabilities.cloned(),
sticky_session_token: sticky_session_token.map(str::to_string),
request_auth_channel: request_auth_channel.map(str::to_string),
skipped_user_id: persistence_policy.skipped.user_id.to_string(),
skipped_api_key_id: persistence_policy.skipped.api_key_id.to_string(),
skipped_required_capabilities: persistence_policy.skipped.required_capabilities.cloned(),
skipped_error_context: persistence_policy.skipped.error_context,
record_runtime_miss_diagnostic,
resolution_mode,
decorate_skipped_candidate,
@@ -655,6 +667,7 @@ where
next_candidate_index: 0,
remembered_affinity: false,
scheduler_cache_affinity_enabled,
auth_api_key_concurrency_wait_deadline: None,
};
cursor.load_next_page().await;
let candidate_count = cursor.candidate_count;
@@ -682,6 +695,10 @@ struct RequestedModelAttemptPageCursor<'a> {
required_capabilities: Option<Value>,
sticky_session_token: Option<String>,
request_auth_channel: Option<String>,
skipped_user_id: String,
skipped_api_key_id: String,
skipped_required_capabilities: Option<Value>,
skipped_error_context: &'static str,
record_runtime_miss_diagnostic: bool,
resolution_mode: LocalCandidateResolutionMode,
decorate_skipped_candidate: DecorateSkippedCandidateFn<'a>,
@@ -691,6 +708,7 @@ struct RequestedModelAttemptPageCursor<'a> {
next_candidate_index: u32,
remembered_affinity: bool,
scheduler_cache_affinity_enabled: bool,
auth_api_key_concurrency_wait_deadline: Option<Instant>,
}
impl<'a> RequestedModelAttemptPageCursor<'a> {
@@ -720,6 +738,15 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
}
};
if page_is_exact_auth_api_key_concurrency_limited(&page) {
if self.wait_for_auth_api_key_concurrency_retry().await {
continue;
}
self.persist_final_auth_api_key_concurrency_skips(page.skipped_candidates)
.await;
return false;
}
let (candidates, resolved_skipped) =
resolve_and_rank_logical_local_execution_candidates(
self.state,
@@ -740,18 +767,10 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
.chain(resolved_skipped)
.map(|skipped| (self.decorate_skipped_candidate)(skipped))
.collect::<Vec<_>>();
let skipped_candidate_count = skipped_candidates.len();
self.candidate_count = self
.candidate_count
.saturating_add(candidates.len() + skipped_candidates.len());
if self.record_runtime_miss_diagnostic {
for skipped_candidate in &skipped_candidates {
record_local_runtime_candidate_skip_reason(
self.state.app(),
&self.trace_id,
skipped_candidate.skip_reason,
);
}
}
.saturating_add(candidates.len() + skipped_candidate_count);
if self.scheduler_cache_affinity_enabled
&& !self.remembered_affinity
&& !candidates.is_empty()
@@ -776,13 +795,90 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
Some(&self.requested_model),
self.request_auth_channel.as_deref(),
);
self.next_candidate_index = next_candidate_index;
self.next_candidate_index = next_candidate_index
.saturating_add(u32::try_from(skipped_candidate_count).unwrap_or(u32::MAX));
if !items.is_empty() {
self.pending_items = items;
return true;
}
let skipped_starting_candidate_index = next_candidate_index;
let skipped_persistence = LocalSkippedCandidatePersistenceContext {
user_id: self.skipped_user_id.as_str(),
api_key_id: self.skipped_api_key_id.as_str(),
required_capabilities: self.skipped_required_capabilities.as_ref(),
error_context: self.skipped_error_context,
record_runtime_miss_diagnostic: self.record_runtime_miss_diagnostic,
};
persist_skipped_local_execution_candidates_with_context(
self.state.app(),
&self.trace_id,
skipped_persistence,
skipped_starting_candidate_index,
skipped_candidates,
)
.await;
}
}
async fn wait_for_auth_api_key_concurrency_retry(&mut self) -> bool {
let now = Instant::now();
let deadline = *self
.auth_api_key_concurrency_wait_deadline
.get_or_insert(now + AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET);
if now >= deadline {
return false;
}
let sleep_duration =
AUTH_API_KEY_CONCURRENCY_RETRY_DELAY.min(deadline.saturating_duration_since(now));
tokio::time::sleep(sleep_duration).await;
self.page_cursor.restart_scan();
true
}
async fn persist_final_auth_api_key_concurrency_skips(
&mut self,
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
) {
let skipped_candidates = skipped_candidates
.into_iter()
.map(|skipped| (self.decorate_skipped_candidate)(skipped))
.collect::<Vec<_>>();
let skipped_candidate_count = skipped_candidates.len();
self.candidate_count = self.candidate_count.saturating_add(skipped_candidate_count);
let skipped_persistence = LocalSkippedCandidatePersistenceContext {
user_id: self.skipped_user_id.as_str(),
api_key_id: self.skipped_api_key_id.as_str(),
required_capabilities: self.skipped_required_capabilities.as_ref(),
error_context: self.skipped_error_context,
record_runtime_miss_diagnostic: self.record_runtime_miss_diagnostic,
};
persist_skipped_local_execution_candidates_with_context(
self.state.app(),
&self.trace_id,
skipped_persistence,
self.next_candidate_index,
skipped_candidates,
)
.await;
self.next_candidate_index = self
.next_candidate_index
.saturating_add(u32::try_from(skipped_candidate_count).unwrap_or(u32::MAX));
}
}
fn page_is_exact_auth_api_key_concurrency_limited(
page: &AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
) -> bool {
page.candidates.is_empty()
&& !page.skipped_candidates.is_empty()
&& page
.skipped_candidates
.iter()
.all(|skipped| skipped.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
}
async fn pop_attempt_from_items(
@@ -792,8 +888,8 @@ async fn pop_attempt_from_items(
let front = items.front_mut()?;
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if let Some(attempt) = attempts.pop_front() {
if attempts.is_empty() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(attempts) {
if dispatch_sequence_exhausted(attempts) {
items.pop_front();
}
return Some(attempt);
@@ -805,7 +901,7 @@ async fn pop_attempt_from_items(
candidate_index,
pending_attempts,
} => {
if let Some(attempt) = pending_attempts.pop_front() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
let Some(candidate) = cursor.next_key().await else {
@@ -814,9 +910,12 @@ async fn pop_attempt_from_items(
items.pop_front();
continue;
};
*pending_attempts = build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
)
.into(),
);
}
LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { .. } => {
@@ -1005,7 +1104,7 @@ where
{
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
let extra_data = ai_candidate_extra_data_with_ranking(
build_extra_data(&candidate),
available_candidate_base_extra_data_with_dispatch_ref(&candidate, build_extra_data),
candidate.ranking.as_ref(),
);
let should_persist = should_persist_available_local_candidate(&candidate);
@@ -1057,6 +1156,70 @@ where
attempts
}
fn available_candidate_extra_data_with_dispatch_ref<F>(
candidate: &EligibleLocalExecutionCandidate,
build_extra_data: &F,
) -> Option<Value>
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
{
ai_candidate_extra_data_with_ranking(
available_candidate_base_extra_data_with_dispatch_ref(candidate, build_extra_data),
candidate.ranking.as_ref(),
)
}
fn available_candidate_base_extra_data_with_dispatch_ref<F>(
candidate: &EligibleLocalExecutionCandidate,
build_extra_data: &F,
) -> Option<Value>
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
{
let dispatch_ref = serde_json::to_value(dispatch_ref_for_local_candidate(candidate)).ok()?;
let mut object = match build_extra_data(candidate) {
Some(Value::Object(object)) => object,
Some(value) => {
let mut object = serde_json::Map::new();
object.insert("extra".to_string(), value);
object
}
None => serde_json::Map::new(),
};
object.insert("dispatch_ref".to_string(), dispatch_ref);
Some(Value::Object(object))
}
fn dispatch_sequence_from_attempts(
attempts: Vec<LocalExecutionCandidateAttempt>,
) -> DispatchSequence<LocalExecutionCandidateAttempt> {
DispatchSequence::new(
attempts
.into_iter()
.map(|attempt| DispatchSequenceItem {
candidate_index: attempt.candidate_index,
retry_index: attempt.retry_index,
candidate: attempt,
mark: aether_dispatch_core::DispatchSequenceMark::Pending,
})
.collect(),
)
}
fn next_attempt_from_dispatch_sequence(
sequence: &mut DispatchSequence<LocalExecutionCandidateAttempt>,
) -> Option<LocalExecutionCandidateAttempt> {
let attempt = sequence.next()?.candidate.clone();
let _ = sequence.mark_succeeded();
Some(attempt)
}
fn dispatch_sequence_exhausted(
sequence: &mut DispatchSequence<LocalExecutionCandidateAttempt>,
) -> bool {
sequence.next().is_none()
}
fn build_unpersisted_local_execution_candidate_attempts(
candidate: EligibleLocalExecutionCandidate,
candidate_index: u32,
@@ -1457,6 +1620,16 @@ mod tests {
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 2);
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
}
#[test]
@@ -1562,6 +1735,16 @@ mod tests {
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 1);
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
}
#[test]
@@ -1643,15 +1826,26 @@ mod tests {
Some(&json!("cached_affinity"))
);
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
assert_eq!(
extra_data
.get("dispatch_ref")
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("ranked-key"))
);
}
#[tokio::test]
async fn dynamic_attempt_source_does_not_drain_unexecuted_single_keys() {
let mut source = LocalExecutionCandidateAttemptSource {
items: VecDeque::from([LocalExecutionCandidateAttemptSourceItem::Static {
attempts: build_unpersisted_local_execution_candidate_attempts(
sample_eligible("normal-key", None),
0,
attempts: dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
sample_eligible("normal-key", None),
0,
)
.into(),
),
}]),
};

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

@@ -21,7 +21,6 @@ use crate::ai_serving::{
use crate::orchestration::LocalExecutionCandidateMetadata;
use super::candidate_ranking::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
@@ -61,7 +60,6 @@ struct GatewayLocalCandidateResolutionPort<'a> {
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&'a ClientSessionAffinity>,
required_capabilities: Option<&'a serde_json::Value>,
sticky_session_token: Option<&'a str>,
request_auth_channel: Option<&'a str>,
}
@@ -182,14 +180,7 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
&self,
candidates: Vec<Self::Eligible>,
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
Ok(apply_local_execution_pool_scheduler(
self.state,
candidates,
self.sticky_session_token,
self.requested_model,
self.request_auth_channel,
)
.await)
Ok((candidates, Vec::new()))
}
}
@@ -201,7 +192,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
@@ -216,7 +207,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
AiCandidateResolutionMode::Standard,
)
@@ -231,7 +222,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transpor
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
@@ -246,7 +237,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transpor
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
AiCandidateResolutionMode::WithoutTransportPairGate,
)
@@ -261,7 +252,7 @@ pub(crate) async fn resolve_and_rank_logical_local_execution_candidates(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
mode: AiCandidateResolutionMode,
) -> (
@@ -276,7 +267,7 @@ pub(crate) async fn resolve_and_rank_logical_local_execution_candidates(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
mode,
false,
@@ -292,7 +283,7 @@ async fn resolve_and_rank_local_execution_candidates_with_mode(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
mode: AiCandidateResolutionMode,
) -> (
@@ -307,10 +298,10 @@ async fn resolve_and_rank_local_execution_candidates_with_mode(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
mode,
true,
false,
)
.await
}
@@ -324,7 +315,7 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
mode: AiCandidateResolutionMode,
expand_pool_groups: bool,
@@ -339,7 +330,6 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
request_auth_channel,
};

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

@@ -323,6 +323,16 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
Ok(None)
}
pub(crate) fn restart_scan(&mut self) {
self.format_index = 0;
self.requested_name_indexes.clear();
self.requested_name_offsets.clear();
self.scanned_rows_by_format.clear();
self.resolved_global_model_names.clear();
self.fallback_scanned_api_formats.clear();
self.seen_candidate_keys.clear();
}
async fn next_page_for_api_format(
&mut self,
candidate_api_format: &str,

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

@@ -26,6 +26,10 @@ mod standard;
mod state;
pub(crate) use self::candidate_materialization::LocalExecutionAttemptSource;
pub(crate) use self::candidate_resolution::{
candidate_auth_channel_skip_reason, read_candidate_transport_snapshot,
EligibleLocalExecutionCandidate, LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
};
pub(crate) use self::passthrough::{
build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports,
build_local_same_format_sync_attempt_source, build_local_same_format_sync_plan_and_reports,
@@ -38,7 +42,11 @@ pub(crate) use self::plan_builders::{
AiStreamAttempt, AiSyncAttempt,
};
pub(crate) use self::pool_scores::build_provider_key_pool_score_upsert;
pub(crate) use self::pool_scores::provider_key_pool_score_scope;
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::runtime_miss::{
apply_local_runtime_candidate_terminal_reason, record_local_runtime_candidate_skip_reason,
};
pub(crate) use self::specialized::{
build_local_gemini_files_stream_attempt_source_for_kind,
build_local_gemini_files_stream_plan_and_reports_for_kind,

10
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs
View File

@@ -9,7 +9,7 @@ use crate::{AiExecutionDecision, AppState, GatewayError};
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
use super::candidates::{
materialize_local_same_format_provider_candidate_attempts,
build_local_same_format_provider_candidate_attempt_source,
resolve_local_same_format_provider_decision_input,
};
use super::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
@@ -55,7 +55,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -65,7 +65,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
candidate_count,
);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) =
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
@@ -122,7 +122,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -132,7 +132,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
candidate_count,
);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) =
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,

9
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs
View File

@@ -20,7 +20,6 @@ pub(crate) use crate::ai_serving::{
use super::{
build_local_same_format_provider_candidate_attempt_source,
materialize_local_same_format_provider_candidate_attempts,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
resolve_local_same_format_provider_decision_input, AiStreamAttempt, AiSyncAttempt, AppState,
GatewayControlDecision, GatewayError, LocalSameFormatProviderCandidateAttempt,
@@ -348,7 +347,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -362,7 +361,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -433,7 +432,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -447,7 +446,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)

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

File diff suppressed because it is too large Load Diff

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

@@ -318,10 +318,10 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
return Ok(None);
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,
@@ -359,11 +359,11 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
return Ok(None);
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
let empty_body_json = serde_json::Value::Null;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,
@@ -407,11 +407,11 @@ async fn build_local_sync_plan_and_reports(
return Ok(Vec::new());
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,
@@ -459,12 +459,12 @@ async fn build_local_stream_plan_and_reports(
return Ok(Vec::new());
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
let mut plans = Vec::new();
let empty_body_json = serde_json::Value::Null;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,

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

@@ -405,7 +405,7 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
return Ok(None);
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -413,12 +413,12 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(None);
};
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,
@@ -465,7 +465,7 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
return Ok(None);
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -473,12 +473,12 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(None);
};
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,
@@ -521,7 +521,7 @@ async fn build_local_sync_plan_and_reports(
return Ok(Vec::new());
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -529,13 +529,13 @@ async fn build_local_sync_plan_and_reports(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(Vec::new());
};
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,
@@ -597,7 +597,7 @@ async fn build_local_stream_plan_and_reports(
return Ok(Vec::new());
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -605,13 +605,13 @@ async fn build_local_stream_plan_and_reports(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(Vec::new());
};
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,

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

@@ -180,7 +180,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
return Ok(None);
};
let Some(attempts) = list_local_video_create_candidate_attempts(
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
state,
trace_id,
&input,
@@ -188,12 +188,12 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(None);
};
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
state, parts, body_json, trace_id, &input, attempt, spec,
)
@@ -223,7 +223,7 @@ async fn build_local_sync_plan_and_reports(
return Ok(Vec::new());
};
let Some(attempts) = list_local_video_create_candidate_attempts(
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
state,
trace_id,
&input,
@@ -231,13 +231,13 @@ async fn build_local_sync_plan_and_reports(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(Vec::new());
};
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
state, parts, body_json, trace_id, &input, attempt, spec,
)

27
apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs
View File

@@ -20,8 +20,7 @@ use crate::ai_serving::GatewayControlDecision;
use crate::{AiExecutionDecision, AppState, GatewayError};
use super::candidates::{
build_local_standard_candidate_attempt_source, materialize_local_standard_candidate_attempts,
resolve_local_standard_decision_input,
build_local_standard_candidate_attempt_source, resolve_local_standard_decision_input,
};
use super::payload::maybe_build_local_standard_decision_payload_for_candidate;
use super::{LocalStandardDecisionInput, LocalStandardSpec};
@@ -323,12 +322,12 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -372,12 +371,12 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -427,15 +426,15 @@ pub(crate) async fn build_local_sync_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -501,15 +500,15 @@ pub(crate) async fn build_local_stream_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)

75
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
View File

@@ -1,28 +1,23 @@
use serde_json::Value;
use tracing::warn;
use crate::ai_serving::planner::common::{
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::runtime_miss::set_local_runtime_execution_exhausted_diagnostic;
use crate::ai_serving::GatewayControlDecision;
use crate::{AiExecutionDecision, AppState, GatewayError};
use tracing::warn;
mod decision;
mod plans;
use self::decision::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};
use self::plans::{
build_local_openai_chat_stream_attempt_source, build_local_openai_chat_stream_plan_and_reports,
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
list_local_openai_chat_candidates, resolve_local_openai_chat_decision_input,
set_local_openai_chat_miss_diagnostic,
resolve_local_openai_chat_decision_input,
};
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports_for_kind(
@@ -146,32 +141,17 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
return Ok(None);
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, false).await {
Ok(value) => value,
Err(err) => {
warn!(
event_name = "local_openai_chat_scheduler_selection_failed",
log_type = "event",
trace_id = %trace_id,
error = ?err,
"gateway local openai chat sync decision scheduler selection failed"
);
return Ok(None);
}
};
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, false,
)
.await;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
@@ -181,7 +161,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
attempt,
OPENAI_CHAT_SYNC_PLAN_KIND,
"openai_chat_sync_success",
false,
upstream_is_stream,
)
.await
{
@@ -212,32 +192,17 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
return Ok(None);
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, true).await {
Ok(value) => value,
Err(err) => {
warn!(
event_name = "local_openai_chat_scheduler_selection_failed",
log_type = "event",
trace_id = %trace_id,
error = ?err,
"gateway local openai chat stream decision scheduler selection failed"
);
return Ok(None);
}
};
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, true,
)
.await;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
@@ -247,7 +212,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
upstream_is_stream,
)
.await
{

2
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans.rs
View File

@@ -22,7 +22,7 @@ pub(super) use self::sync::{
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
};
fn openai_chat_upstream_is_stream_for_candidate(
pub(super) fn openai_chat_upstream_is_stream_for_candidate(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
client_is_stream: bool,

86
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
View File

@@ -3,13 +3,10 @@ use tracing::warn;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
};
use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
@@ -176,27 +173,12 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
return Ok(Vec::new());
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, true).await {
Ok(value) => value,
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat stream decision scheduler selection failed"
);
set_local_openai_chat_miss_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
"scheduler_selection_failed",
);
return Ok(Vec::new());
}
};
if candidates.is_empty() && skipped_candidates.is_empty() {
let Some((mut attempt_source, candidate_count)) =
build_local_openai_chat_stream_attempt_source(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
else {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
@@ -206,59 +188,13 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
0,
);
return Ok(Vec::new());
}
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidates.len() + skipped_candidates.len(),
);
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
};
let mut plans = Vec::new();
for attempt in attempts {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
trace_id,
body_json,
&input,
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
upstream_is_stream,
)
.await
else {
continue;
};
match build_openai_chat_stream_plan_from_decision(parts, body_json, payload) {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat stream decision plan build failed"
);
}
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
plans.push(attempt);
if plans.len() >= candidate_count {
break;
}
}

85
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
View File

@@ -3,13 +3,10 @@ use tracing::warn;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
};
use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
@@ -176,27 +173,11 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
return Ok(Vec::new());
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, false).await {
Ok(value) => value,
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat sync decision scheduler selection failed"
);
set_local_openai_chat_miss_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
"scheduler_selection_failed",
);
return Ok(Vec::new());
}
};
if candidates.is_empty() && skipped_candidates.is_empty() {
let Some((mut attempt_source, candidate_count)) = build_local_openai_chat_sync_attempt_source(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
else {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
@@ -206,59 +187,13 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
0,
);
return Ok(Vec::new());
}
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidates.len() + skipped_candidates.len(),
);
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
};
let mut plans = Vec::new();
for attempt in attempts {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
trace_id,
body_json,
&input,
attempt,
OPENAI_CHAT_SYNC_PLAN_KIND,
"openai_chat_sync_success",
upstream_is_stream,
)
.await
else {
continue;
};
match build_openai_chat_sync_plan_from_decision(parts, body_json, payload) {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat sync decision plan build failed"
);
}
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
plans.push(attempt);
if plans.len() >= candidate_count {
break;
}
}

10
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs
View File

@@ -6,7 +6,7 @@ mod decision;
mod plans;
use self::decision::{
materialize_local_openai_responses_candidate_attempts,
build_local_openai_responses_candidate_attempt_source,
maybe_build_local_openai_responses_decision_payload_for_candidate,
resolve_local_openai_responses_decision_input,
};
@@ -108,12 +108,12 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
return Ok(None);
};
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -146,12 +146,12 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
return Ok(None);
};
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)

9
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs
View File

@@ -3,7 +3,6 @@ use tracing::warn;
use super::decision::{
build_local_openai_responses_candidate_attempt_source,
materialize_local_openai_responses_candidate_attempts,
maybe_build_local_openai_responses_decision_payload_for_candidate,
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
@@ -312,7 +311,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
let (mut source, candidate_count) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -322,7 +321,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -384,7 +383,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
let (mut source, candidate_count) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -394,7 +393,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)