mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 01:47:47 +08:00
Merge remote-tracking branch 'origin/main' into codex/fix-antigravity-quota
This commit is contained in:
@@ -47,13 +47,12 @@ use crate::cache::{
|
||||
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::orchestration::{ExecutionAttemptIdentity, POOL_KEY_RETRY_INDEX_STRIDE};
|
||||
use crate::scheduler::candidate::is_auth_api_key_concurrency_limit_skip_reason;
|
||||
use crate::scheduler::config::SchedulerSchedulingMode;
|
||||
use crate::stage_metrics::observe_gateway_stage_ms;
|
||||
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);
|
||||
|
||||
@@ -481,10 +480,6 @@ where
|
||||
type ExtraData = Value;
|
||||
type Error = Infallible;
|
||||
|
||||
fn attempt_slot_count(&self, candidate: &Self::Candidate) -> u32 {
|
||||
local_attempt_slot_count(&candidate.transport)
|
||||
}
|
||||
|
||||
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
|
||||
available_candidate_extra_data_with_dispatch_ref(candidate, &self.build_extra_data)
|
||||
}
|
||||
@@ -1610,7 +1605,8 @@ async fn persist_available_local_execution_candidate_at_index<F>(
|
||||
where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
{
|
||||
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
|
||||
// Exactly one attempt is materialized per candidate; same-key retries are
|
||||
// derived lazily by the attempt loop after a failure.
|
||||
let extra_data = ai_candidate_extra_data_with_ranking(
|
||||
available_candidate_base_extra_data_with_dispatch_ref(&candidate, build_extra_data),
|
||||
candidate.ranking.as_ref(),
|
||||
@@ -1625,53 +1621,34 @@ where
|
||||
Some(candidate_index),
|
||||
extra_data,
|
||||
);
|
||||
let should_persist = should_persist_available_local_candidate(&candidate);
|
||||
let mut attempts = Vec::with_capacity(attempt_slots as usize);
|
||||
let mut owned_candidate = Some(candidate);
|
||||
let retry_index = effective_retry_index(0, candidate.orchestration.pool_key_index);
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let candidate_id = if should_persist_available_local_candidate(&candidate) {
|
||||
state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
context.user_id,
|
||||
context.api_key_id,
|
||||
&candidate.candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
generated_candidate_id.as_str(),
|
||||
context.required_capabilities,
|
||||
extra_data,
|
||||
current_unix_ms(),
|
||||
context.error_context,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
generated_candidate_id
|
||||
};
|
||||
|
||||
for retry_index in 0..attempt_slots {
|
||||
let candidate_ref = owned_candidate
|
||||
.as_ref()
|
||||
.expect("candidate should remain available until final retry");
|
||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||
let candidate_id = if should_persist {
|
||||
state
|
||||
.persist_available_local_candidate(
|
||||
trace_id,
|
||||
context.user_id,
|
||||
context.api_key_id,
|
||||
&candidate_ref.candidate,
|
||||
candidate_index,
|
||||
effective_retry_index(retry_index, candidate_ref.orchestration.pool_key_index),
|
||||
generated_candidate_id.as_str(),
|
||||
context.required_capabilities,
|
||||
extra_data.clone(),
|
||||
current_unix_ms(),
|
||||
context.error_context,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
generated_candidate_id
|
||||
};
|
||||
|
||||
let candidate = if retry_index + 1 == attempt_slots {
|
||||
owned_candidate
|
||||
.take()
|
||||
.expect("final retry should consume owned candidate")
|
||||
} else {
|
||||
candidate_ref.clone()
|
||||
};
|
||||
let retry_index =
|
||||
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
|
||||
attempts.push(LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id,
|
||||
});
|
||||
}
|
||||
|
||||
attempts
|
||||
vec![LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id,
|
||||
}]
|
||||
}
|
||||
|
||||
fn available_candidate_extra_data_with_dispatch_ref<F>(
|
||||
@@ -1840,6 +1817,7 @@ fn routing_trace_for_candidate(
|
||||
CandidateKind::Provider => Some(candidate.key_id.clone()),
|
||||
CandidateKind::PoolGroup => None,
|
||||
},
|
||||
api_format: Some(candidate.endpoint_api_format.clone()),
|
||||
provider_priority: candidate.provider_priority,
|
||||
key_priority: candidate
|
||||
.key_global_priority_for_format
|
||||
@@ -1923,32 +1901,15 @@ fn build_unpersisted_local_execution_candidate_attempts(
|
||||
candidate: EligibleLocalExecutionCandidate,
|
||||
candidate_index: u32,
|
||||
) -> VecDeque<LocalExecutionCandidateAttempt> {
|
||||
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
|
||||
let mut attempts = VecDeque::with_capacity(attempt_slots as usize);
|
||||
let mut owned_candidate = Some(candidate);
|
||||
|
||||
for retry_index in 0..attempt_slots {
|
||||
let candidate = if retry_index + 1 == attempt_slots {
|
||||
owned_candidate
|
||||
.take()
|
||||
.expect("final retry should consume owned candidate")
|
||||
} else {
|
||||
owned_candidate
|
||||
.as_ref()
|
||||
.expect("candidate should remain available until final retry")
|
||||
.clone()
|
||||
};
|
||||
let retry_index =
|
||||
effective_retry_index(retry_index, candidate.orchestration.pool_key_index);
|
||||
attempts.push_back(LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id: Uuid::new_v4().to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
attempts
|
||||
// One attempt per candidate; same-key retries are derived lazily by the
|
||||
// attempt loop after a failure.
|
||||
let retry_index = effective_retry_index(0, candidate.orchestration.pool_key_index);
|
||||
VecDeque::from([LocalExecutionCandidateAttempt {
|
||||
eligible: candidate,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
candidate_id: Uuid::new_v4().to_string(),
|
||||
}])
|
||||
}
|
||||
|
||||
async fn persist_pool_group_exhaustion_skipped_candidate(
|
||||
@@ -2276,6 +2237,8 @@ mod tests {
|
||||
pool_key_index,
|
||||
pool_key_lease: None,
|
||||
scheduler_affinity_epoch: None,
|
||||
// These tests cover persistence shape, not same-key retries.
|
||||
sticky_key_attempts: Some(1),
|
||||
},
|
||||
ranking: None,
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ use aether_ai_serving::{
|
||||
AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig,
|
||||
AiRankingSchedulingMode,
|
||||
};
|
||||
use aether_routing_core::{ResolvedRoutingPolicy, RoutingSchedulingMode, RoutingSetPriorityMode};
|
||||
use aether_routing_core::ResolvedRoutingPolicy;
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::warn;
|
||||
@@ -184,35 +184,16 @@ fn ai_ranking_scheduling_mode(mode: SchedulerSchedulingMode) -> AiRankingSchedul
|
||||
}
|
||||
}
|
||||
|
||||
/// Ordering config for a request. A resolved routing policy is authoritative
|
||||
/// and is never merged with legacy system-config values; without a policy the
|
||||
/// effective default (system-default routing group, then legacy keys) applies.
|
||||
pub(crate) async fn scheduler_ordering_config_for_routing_policy(
|
||||
state: PlannerAppState<'_>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> SchedulerOrderingConfig {
|
||||
let system_config = read_scheduler_ordering_config_or_default(state).await;
|
||||
match routing_policy {
|
||||
Some(policy) => {
|
||||
let mut config = scheduler_ordering_config_from_routing_policy(policy);
|
||||
config.keep_priority_on_conversion |= system_config.keep_priority_on_conversion;
|
||||
config
|
||||
}
|
||||
None => system_config,
|
||||
}
|
||||
}
|
||||
|
||||
fn scheduler_ordering_config_from_routing_policy(
|
||||
policy: &ResolvedRoutingPolicy,
|
||||
) -> SchedulerOrderingConfig {
|
||||
SchedulerOrderingConfig {
|
||||
priority_mode: match policy.priority_mode {
|
||||
RoutingSetPriorityMode::Provider => SchedulerPriorityMode::Provider,
|
||||
RoutingSetPriorityMode::GlobalKey => SchedulerPriorityMode::GlobalKey,
|
||||
},
|
||||
scheduling_mode: match policy.scheduling_mode {
|
||||
RoutingSchedulingMode::FixedOrder => SchedulerSchedulingMode::FixedOrder,
|
||||
RoutingSchedulingMode::CacheAffinity => SchedulerSchedulingMode::CacheAffinity,
|
||||
RoutingSchedulingMode::LoadBalance => SchedulerSchedulingMode::LoadBalance,
|
||||
},
|
||||
keep_priority_on_conversion: policy.keep_priority_on_conversion,
|
||||
Some(policy) => SchedulerOrderingConfig::from_routing_policy(policy),
|
||||
None => read_scheduler_ordering_config_or_default(state).await,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -231,14 +212,26 @@ fn routing_overlaid_candidate(
|
||||
let overlaid_key_priority = match kind {
|
||||
LocalExecutionCandidateKind::SingleKey => policy
|
||||
.ranking_overlay
|
||||
.key_priority_overrides
|
||||
.get(candidate.key_id.as_str()),
|
||||
.key_priority_override_matching_format(candidate.key_id.as_str(), |format| {
|
||||
crate::ai_serving::api_format_alias_matches(
|
||||
format,
|
||||
candidate.endpoint_api_format.as_str(),
|
||||
)
|
||||
})
|
||||
.or_else(|| {
|
||||
policy
|
||||
.ranking_overlay
|
||||
.key_priority_overrides
|
||||
.get(candidate.key_id.as_str())
|
||||
.copied()
|
||||
}),
|
||||
LocalExecutionCandidateKind::PoolGroup => policy
|
||||
.ranking_overlay
|
||||
.pool_priority_overrides
|
||||
.get(candidate.provider_id.as_str()),
|
||||
.get(candidate.provider_id.as_str())
|
||||
.copied(),
|
||||
};
|
||||
if let Some(overlaid_key_priority) = overlaid_key_priority.copied() {
|
||||
if let Some(overlaid_key_priority) = overlaid_key_priority {
|
||||
overlaid.key_internal_priority = overlaid_key_priority;
|
||||
overlaid.key_global_priority_for_format = Some(overlaid_key_priority);
|
||||
}
|
||||
@@ -378,6 +371,7 @@ mod tests {
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
ranking_overlay: aether_routing_core::RankingOverlay::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: BTreeMap::new(),
|
||||
@@ -396,7 +390,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn routing_policy_inherits_global_conversion_priority_override() {
|
||||
async fn routing_policy_ignores_legacy_global_conversion_priority_override() {
|
||||
let data_state = GatewayDataState::default().with_system_config_values_for_tests([(
|
||||
"keep_priority_on_conversion".to_string(),
|
||||
json!(true),
|
||||
@@ -413,6 +407,7 @@ mod tests {
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
@@ -429,7 +424,10 @@ mod tests {
|
||||
ordering.scheduling_mode,
|
||||
crate::scheduler::config::SchedulerSchedulingMode::FixedOrder
|
||||
);
|
||||
assert!(ordering.keep_priority_on_conversion);
|
||||
assert!(
|
||||
!ordering.keep_priority_on_conversion,
|
||||
"a resolved routing policy must not inherit the legacy system-config flag"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -447,6 +445,7 @@ mod tests {
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::GlobalKey,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
ranking_overlay: aether_routing_core::RankingOverlay {
|
||||
pool_priority_overrides: BTreeMap::from([("provider-1".to_string(), 4)]),
|
||||
key_priority_overrides: BTreeMap::from([("representative-key".to_string(), 1)]),
|
||||
|
||||
@@ -23,7 +23,9 @@ use crate::ai_serving::{
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
use crate::stage_metrics::observe_gateway_stage_ms;
|
||||
|
||||
use super::candidate_ranking::rank_eligible_local_execution_candidates;
|
||||
use super::candidate_ranking::{
|
||||
rank_eligible_local_execution_candidates, scheduler_ordering_config_for_routing_policy,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub(crate) struct EligibleLocalExecutionCandidate {
|
||||
@@ -378,8 +380,18 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
|
||||
"candidate_resolution_core",
|
||||
started_at.elapsed().as_millis() as u64,
|
||||
);
|
||||
let sticky_key_attempts = if outcome.eligible_candidates.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(
|
||||
scheduler_ordering_config_for_routing_policy(state, routing_policy)
|
||||
.await
|
||||
.sticky_key_attempts,
|
||||
)
|
||||
};
|
||||
for candidate in &mut outcome.eligible_candidates {
|
||||
candidate.orchestration.scheduler_affinity_epoch = Some(scheduler_affinity_epoch);
|
||||
candidate.orchestration.sticky_key_attempts = sticky_key_attempts;
|
||||
}
|
||||
(outcome.eligible_candidates, outcome.skipped_candidates)
|
||||
}
|
||||
|
||||
@@ -174,6 +174,8 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
|
||||
self.ranking_seed,
|
||||
false,
|
||||
self.request_operation,
|
||||
self.routing_policy
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -1291,6 +1293,9 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
|
||||
.then_some(self.client_session_affinity.as_ref())
|
||||
.flatten(),
|
||||
self.ranking_seed,
|
||||
self.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await?;
|
||||
let skipped_candidates = skipped_candidates
|
||||
@@ -1884,6 +1889,7 @@ mod tests {
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
@@ -1947,6 +1953,7 @@ mod tests {
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
@@ -2657,14 +2664,16 @@ mod tests {
|
||||
candidate_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests("development-key")
|
||||
// Legacy keys deliberately disagree with the routing policy: the
|
||||
// resolved policy must be the only source of scheduler ordering.
|
||||
.with_system_config_values_for_tests([
|
||||
(
|
||||
"scheduling_mode".to_string(),
|
||||
serde_json::json!("fixed_order"),
|
||||
serde_json::json!("cache_affinity"),
|
||||
),
|
||||
(
|
||||
"keep_priority_on_conversion".to_string(),
|
||||
serde_json::json!(true),
|
||||
serde_json::json!(false),
|
||||
),
|
||||
]);
|
||||
let app = AppState::new()
|
||||
@@ -2681,7 +2690,8 @@ mod tests {
|
||||
resolved_model: "gpt-5.4-mini".to_string(),
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: false,
|
||||
keep_priority_on_conversion: true,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
|
||||
@@ -1073,6 +1073,7 @@ fn ensure_report_context_routing_trace(
|
||||
endpoint_id: decision.endpoint_id.clone().unwrap_or_default(),
|
||||
model_id,
|
||||
key_id,
|
||||
api_format: decision.provider_api_format.clone(),
|
||||
provider_priority,
|
||||
key_priority,
|
||||
},
|
||||
|
||||
@@ -26,6 +26,7 @@ use crate::ai_serving::{
|
||||
};
|
||||
use crate::client_session_affinity::client_session_affinity_from_api_request;
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::{
|
||||
@@ -140,6 +141,10 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
current_unix_secs(),
|
||||
false,
|
||||
spec.operation.map(|operation| operation.as_str()),
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
@@ -246,6 +251,10 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
|
||||
current_unix_secs(),
|
||||
false,
|
||||
spec.operation.map(|operation| operation.as_str()),
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await?;
|
||||
|
||||
|
||||
@@ -203,6 +203,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
client_session_affinity: input.client_session_affinity.as_ref(),
|
||||
routing_policy: input.routing_policy.as_ref(),
|
||||
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
|
||||
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
|
||||
@@ -4,7 +4,7 @@ use aether_ai_serving::{
|
||||
build_ai_execution_report_context,
|
||||
insert_provider_stream_event_api_format as insert_ai_provider_stream_event_api_format,
|
||||
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
|
||||
AiExecutionReportContextParts, AiRequestOrigin,
|
||||
AiExecutionReportContextParts, AiRequestOrigin, STICKY_KEY_ATTEMPTS_REPORT_FIELD,
|
||||
};
|
||||
use aether_routing_core::ResolvedRoutingPolicy;
|
||||
use aether_runtime_state::RuntimeLockLease;
|
||||
@@ -59,6 +59,9 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
pub(crate) client_session_affinity: Option<&'a ClientSessionAffinity>,
|
||||
pub(crate) routing_policy: Option<&'a ResolvedRoutingPolicy>,
|
||||
pub(crate) scheduler_affinity_epoch: Option<u64>,
|
||||
/// Routing policy sticky-key attempt budget; read back by the attempt
|
||||
/// loop to derive same-key retries lazily.
|
||||
pub(crate) sticky_key_attempts: Option<u32>,
|
||||
pub(crate) client_requested_stream: bool,
|
||||
pub(crate) upstream_is_stream: bool,
|
||||
pub(crate) has_envelope: bool,
|
||||
@@ -124,6 +127,12 @@ pub(crate) fn build_local_execution_report_context(
|
||||
Value::Number(epoch.into()),
|
||||
);
|
||||
}
|
||||
if let Some(sticky_key_attempts) = parts.sticky_key_attempts {
|
||||
extra_fields.insert(
|
||||
STICKY_KEY_ATTEMPTS_REPORT_FIELD.to_string(),
|
||||
Value::Number(sticky_key_attempts.into()),
|
||||
);
|
||||
}
|
||||
insert_request_path_fields(
|
||||
&mut extra_fields,
|
||||
parts.request_path,
|
||||
@@ -330,6 +339,7 @@ mod tests {
|
||||
client_session_affinity: Some(&client_session_affinity),
|
||||
routing_policy: None,
|
||||
scheduler_affinity_epoch: None,
|
||||
sticky_key_attempts: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
@@ -413,6 +423,7 @@ mod tests {
|
||||
client_session_affinity: None,
|
||||
routing_policy: None,
|
||||
scheduler_affinity_epoch: None,
|
||||
sticky_key_attempts: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: true,
|
||||
has_envelope: false,
|
||||
@@ -480,6 +491,7 @@ mod tests {
|
||||
client_session_affinity: None,
|
||||
routing_policy: None,
|
||||
scheduler_affinity_epoch: None,
|
||||
sticky_key_attempts: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
|
||||
@@ -109,6 +109,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
client_session_affinity: input.client_session_affinity.as_ref(),
|
||||
routing_policy: input.routing_policy.as_ref(),
|
||||
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
|
||||
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
has_envelope: false,
|
||||
|
||||
@@ -26,6 +26,7 @@ use crate::ai_serving::{
|
||||
};
|
||||
use crate::client_session_affinity::client_session_affinity_from_parts;
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalGeminiFilesCandidateAttempt;
|
||||
@@ -108,6 +109,10 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
@@ -181,6 +186,10 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await?;
|
||||
Ok(build_local_execution_candidate_attempt_source_with_serving(
|
||||
|
||||
@@ -122,6 +122,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
client_session_affinity: input.client_session_affinity.as_ref(),
|
||||
routing_policy: input.routing_policy.as_ref(),
|
||||
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
|
||||
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
upstream_is_stream,
|
||||
has_envelope: false,
|
||||
|
||||
@@ -27,6 +27,7 @@ use crate::ai_serving::{
|
||||
};
|
||||
use crate::client_session_affinity::client_session_affinity_from_parts;
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
use crate::{AppState, GatewayError};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
@@ -127,6 +128,10 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
false,
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -201,6 +206,10 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
false,
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -90,6 +90,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
client_session_affinity: input.client_session_affinity.as_ref(),
|
||||
routing_policy: input.routing_policy.as_ref(),
|
||||
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
|
||||
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
|
||||
@@ -29,6 +29,7 @@ use crate::ai_serving::{
|
||||
};
|
||||
use crate::client_session_affinity::client_session_affinity_from_parts;
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalVideoCreateCandidateAttempt;
|
||||
@@ -133,6 +134,10 @@ pub(super) async fn list_local_video_create_candidate_attempts(
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
false,
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -190,6 +195,10 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>(
|
||||
input.client_session_affinity.as_ref(),
|
||||
current_unix_secs(),
|
||||
false,
|
||||
input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -142,6 +142,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
client_session_affinity: input.client_session_affinity.as_ref(),
|
||||
routing_policy: input.routing_policy.as_ref(),
|
||||
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
|
||||
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
|
||||
@@ -195,6 +195,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
client_session_affinity: input.client_session_affinity.as_ref(),
|
||||
routing_policy: input.routing_policy.as_ref(),
|
||||
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
|
||||
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
|
||||
+1
@@ -184,6 +184,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
client_session_affinity: input.client_session_affinity.as_ref(),
|
||||
routing_policy: input.routing_policy.as_ref(),
|
||||
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
|
||||
sticky_key_attempts: eligible.orchestration.sticky_key_attempts,
|
||||
client_requested_stream: body_json
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
|
||||
@@ -8,9 +8,14 @@ use crate::constants::{
|
||||
API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS, API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS,
|
||||
};
|
||||
use crate::scheduler::candidate::SchedulerSkippedCandidate;
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
use crate::GatewayError;
|
||||
|
||||
impl<'a> PlannerAppState<'a> {
|
||||
/// `ordering_config` is the request's routing-policy derived scheduler
|
||||
/// config (see `SchedulerOrderingConfig::from_routing_policy`). `None`
|
||||
/// falls back to the runtime default.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates(
|
||||
self,
|
||||
api_format: &str,
|
||||
@@ -21,6 +26,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
crate::scheduler::candidate::list_selectable_candidates(
|
||||
self.app().data.as_ref(),
|
||||
@@ -33,10 +39,12 @@ impl<'a> PlannerAppState<'a> {
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
self,
|
||||
api_format: &str,
|
||||
@@ -47,6 +55,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -64,10 +73,12 @@ impl<'a> PlannerAppState<'a> {
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
None,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_operation(
|
||||
self,
|
||||
api_format: &str,
|
||||
@@ -79,6 +90,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
request_operation: Option<&str>,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -103,6 +115,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
attempt_now_unix_secs,
|
||||
enable_model_directives,
|
||||
request_operation,
|
||||
ordering_config,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -123,6 +136,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
self,
|
||||
api_format: &str,
|
||||
@@ -132,6 +146,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -148,10 +163,12 @@ impl<'a> PlannerAppState<'a> {
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
self,
|
||||
candidate_api_format: &str,
|
||||
@@ -160,6 +177,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
|
||||
let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
|
||||
@@ -176,6 +194,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
attempt_now_unix_secs,
|
||||
ordering_config,
|
||||
)
|
||||
.await?;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user