mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 01:17:46 +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?;
|
||||
|
||||
|
||||
@@ -756,19 +756,28 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let api_format = self.group.candidate.endpoint_api_format.as_str();
|
||||
rows.sort_by(|left, right| {
|
||||
let left_priority = self
|
||||
.routing_overlay
|
||||
.as_ref()
|
||||
.map_or(left.key_internal_priority, |overlay| {
|
||||
overlay.key_priority(&left.key_id, left.key_internal_priority)
|
||||
});
|
||||
let right_priority = self
|
||||
.routing_overlay
|
||||
.as_ref()
|
||||
.map_or(right.key_internal_priority, |overlay| {
|
||||
overlay.key_priority(&right.key_id, right.key_internal_priority)
|
||||
});
|
||||
let left_priority = self.routing_overlay.as_ref().map_or(
|
||||
left.key_internal_priority,
|
||||
|overlay| {
|
||||
overlay.key_priority_for_format(
|
||||
&left.key_id,
|
||||
api_format,
|
||||
left.key_internal_priority,
|
||||
)
|
||||
},
|
||||
);
|
||||
let right_priority = self.routing_overlay.as_ref().map_or(
|
||||
right.key_internal_priority,
|
||||
|overlay| {
|
||||
overlay.key_priority_for_format(
|
||||
&right.key_id,
|
||||
api_format,
|
||||
right.key_internal_priority,
|
||||
)
|
||||
},
|
||||
);
|
||||
left_priority
|
||||
.cmp(&right_priority)
|
||||
.then(left.key_id.cmp(&right.key_id))
|
||||
@@ -1987,11 +1996,13 @@ fn apply_pool_orchestration(
|
||||
orchestration: PoolCandidateOrchestration,
|
||||
) -> EligibleLocalExecutionCandidate {
|
||||
let scheduler_affinity_epoch = candidate.orchestration.scheduler_affinity_epoch;
|
||||
let sticky_key_attempts = candidate.orchestration.sticky_key_attempts;
|
||||
candidate.orchestration = LocalExecutionCandidateMetadata {
|
||||
candidate_group_id: orchestration.candidate_group_id,
|
||||
pool_key_index: orchestration.pool_key_index,
|
||||
pool_key_lease: None,
|
||||
scheduler_affinity_epoch,
|
||||
sticky_key_attempts,
|
||||
};
|
||||
candidate
|
||||
}
|
||||
@@ -2244,6 +2255,7 @@ mod tests {
|
||||
pool_key_index: Some(0),
|
||||
pool_key_lease: None,
|
||||
scheduler_affinity_epoch: None,
|
||||
sticky_key_attempts: None,
|
||||
}
|
||||
);
|
||||
assert_eq!(reordered[1].orchestration.pool_key_index, Some(1));
|
||||
@@ -2261,6 +2273,7 @@ mod tests {
|
||||
pool_key_index: None,
|
||||
pool_key_lease: None,
|
||||
scheduler_affinity_epoch: None,
|
||||
sticky_key_attempts: None,
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -5157,6 +5170,7 @@ mod tests {
|
||||
priority_mode: RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
ranking_overlay: RankingOverlay {
|
||||
allowed_keys: key_ids.into_iter().map(str::to_string).collect(),
|
||||
..RankingOverlay::default()
|
||||
|
||||
@@ -208,6 +208,7 @@ mod tests {
|
||||
pool_key_index: None,
|
||||
pool_key_lease: None,
|
||||
scheduler_affinity_epoch: None,
|
||||
sticky_key_attempts: None,
|
||||
},
|
||||
ranking: None,
|
||||
}
|
||||
|
||||
@@ -252,6 +252,10 @@ where
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn next_same_key_retry(&self, attempt: &T) -> Result<Option<T>, Self::Error> {
|
||||
Ok(crate::orchestration::next_same_key_retry_attempt(attempt))
|
||||
}
|
||||
|
||||
async fn record_attempt_failed(&self, attempt: &T) -> Result<(), Self::Error> {
|
||||
record_provider_transfer_attempt_failed(
|
||||
self.state,
|
||||
@@ -788,18 +792,31 @@ where
|
||||
{
|
||||
let mut last_attempted = None;
|
||||
let mut fallback_response = None;
|
||||
// A same-key retry derived after a candidate-scoped failure runs before
|
||||
// the source is asked for the next candidate.
|
||||
let mut pending_same_key_retry: Option<Attempt> = None;
|
||||
|
||||
loop {
|
||||
let next_started_at = std::time::Instant::now();
|
||||
let next_attempt =
|
||||
next_execution_attempt_with_timeout(source, trace_id, plan_kind, planning_timeout)
|
||||
let attempt = match pending_same_key_retry.take() {
|
||||
Some(attempt) => attempt,
|
||||
None => {
|
||||
let next_started_at = std::time::Instant::now();
|
||||
let next_attempt = next_execution_attempt_with_timeout(
|
||||
source,
|
||||
trace_id,
|
||||
plan_kind,
|
||||
planning_timeout,
|
||||
)
|
||||
.await?;
|
||||
observe_gateway_stage_ms(
|
||||
"stream_candidate_next",
|
||||
next_started_at.elapsed().as_millis() as u64,
|
||||
);
|
||||
let Some(attempt) = next_attempt else {
|
||||
break;
|
||||
observe_gateway_stage_ms(
|
||||
"stream_candidate_next",
|
||||
next_started_at.elapsed().as_millis() as u64,
|
||||
);
|
||||
let Some(attempt) = next_attempt else {
|
||||
break;
|
||||
};
|
||||
attempt
|
||||
}
|
||||
};
|
||||
if port.should_skip_attempt(&attempt).await? {
|
||||
let provider_id = attempt.execution_plan().provider_id.clone();
|
||||
@@ -839,6 +856,9 @@ where
|
||||
if attempt_fallback_response.is_some() {
|
||||
fallback_response = attempt_fallback_response;
|
||||
}
|
||||
if scope == AiAttemptRetryScope::Candidate {
|
||||
pending_same_key_retry = port.next_same_key_retry(&attempt).await?;
|
||||
}
|
||||
apply_attempt_retry_scope(source, &attempt, scope).await?;
|
||||
}
|
||||
}
|
||||
@@ -952,6 +972,10 @@ where
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn next_same_key_retry(&self, attempt: &T) -> Result<Option<T>, Self::Error> {
|
||||
Ok(crate::orchestration::next_same_key_retry_attempt(attempt))
|
||||
}
|
||||
|
||||
async fn record_attempt_failed(&self, attempt: &T) -> Result<(), Self::Error> {
|
||||
record_provider_transfer_attempt_failed(
|
||||
self.state,
|
||||
|
||||
@@ -1663,6 +1663,22 @@ mod tests {
|
||||
candidate_index: u32,
|
||||
endpoint_id: &str,
|
||||
candidate_id: &str,
|
||||
) -> AiSyncAttempt {
|
||||
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
|
||||
candidate_index,
|
||||
endpoint_id,
|
||||
candidate_id,
|
||||
1,
|
||||
)
|
||||
}
|
||||
|
||||
/// `sticky_key_attempts` is pinned so these tests exercise candidate
|
||||
/// failover; the default same-key retry is covered separately.
|
||||
fn test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
|
||||
candidate_index: u32,
|
||||
endpoint_id: &str,
|
||||
candidate_id: &str,
|
||||
sticky_key_attempts: u32,
|
||||
) -> AiSyncAttempt {
|
||||
AiSyncAttempt {
|
||||
plan: test_openai_image_heartbeat_plan(endpoint_id, candidate_id),
|
||||
@@ -1670,6 +1686,7 @@ mod tests {
|
||||
report_context: Some(json!({
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
"sticky_key_attempts": sticky_key_attempts,
|
||||
})),
|
||||
}
|
||||
}
|
||||
@@ -1828,6 +1845,9 @@ mod tests {
|
||||
report_context: Some(json!({
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": 0,
|
||||
// Pin to a single attempt so this helper exercises candidate
|
||||
// failover rather than the default same-key retry.
|
||||
"sticky_key_attempts": 1,
|
||||
"client_api_format": client_api_format,
|
||||
"provider_api_format": client_api_format,
|
||||
})),
|
||||
@@ -1980,6 +2000,90 @@ mod tests {
|
||||
assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_retries_sticky_key_lazily_before_failover() {
|
||||
let seen_plans = Arc::new(std::sync::Mutex::new(Vec::<(String, Option<String>)>::new()));
|
||||
let seen_plans_for_override = Arc::clone(&seen_plans);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_execution_runtime_sync_override_for_tests(move |plan| {
|
||||
seen_plans_for_override
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.push((plan.endpoint_id.clone(), plan.candidate_id.clone()));
|
||||
if plan.endpoint_id == "endpoint-retry" {
|
||||
Ok(test_openai_image_execution_result(
|
||||
plan,
|
||||
StatusCode::TOO_MANY_REQUESTS.as_u16(),
|
||||
json!({"error": {"message": "retry this candidate"}}),
|
||||
))
|
||||
} else {
|
||||
Ok(test_openai_image_execution_result(
|
||||
plan,
|
||||
StatusCode::OK.as_u16(),
|
||||
json!({"data": [{"b64_json": "second-candidate"}]}),
|
||||
))
|
||||
}
|
||||
});
|
||||
// Three total attempts on the sticky key; only one attempt is
|
||||
// materialized up front, the other two are derived after each failure.
|
||||
let attempts = vec![
|
||||
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
|
||||
0,
|
||||
"endpoint-retry",
|
||||
"candidate-retry",
|
||||
3,
|
||||
),
|
||||
test_openai_image_heartbeat_attempt_with_sticky_key_attempts(
|
||||
1,
|
||||
"endpoint-success",
|
||||
"candidate-success",
|
||||
3,
|
||||
),
|
||||
];
|
||||
let outcome = execute_openai_image_sync_heartbeat_attempts(
|
||||
state,
|
||||
"/v1/images/generations".to_string(),
|
||||
"trace-image-heartbeat-sticky-retry".to_string(),
|
||||
test_openai_image_heartbeat_decision(),
|
||||
TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(),
|
||||
attempts,
|
||||
ProviderTransferTracker::default(),
|
||||
Instant::now(),
|
||||
)
|
||||
.await
|
||||
.expect("heartbeat attempts should execute");
|
||||
let LocalExecutionRequestOutcome::Responded(response) = outcome else {
|
||||
panic!("second candidate should return a response");
|
||||
};
|
||||
let bytes = openai_image_sync_heartbeat_response_body_bytes(response).await;
|
||||
let body: Value = serde_json::from_slice(&bytes).expect("body should decode");
|
||||
|
||||
let seen_plans = seen_plans.lock().expect("mutex should lock").clone();
|
||||
assert_eq!(
|
||||
seen_plans
|
||||
.iter()
|
||||
.map(|(endpoint_id, _)| endpoint_id.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
[
|
||||
"endpoint-retry",
|
||||
"endpoint-retry",
|
||||
"endpoint-retry",
|
||||
"endpoint-success"
|
||||
]
|
||||
);
|
||||
let sticky_candidate_ids = seen_plans[..3]
|
||||
.iter()
|
||||
.map(|(_, candidate_id)| candidate_id.clone())
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
assert_eq!(
|
||||
sticky_candidate_ids.len(),
|
||||
3,
|
||||
"each derived same-key retry must carry a fresh candidate id"
|
||||
);
|
||||
assert_eq!(body, json!({"data": [{"b64_json": "second-candidate"}]}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_honors_provider_transfer_limit() {
|
||||
let call_count = Arc::new(AtomicUsize::new(0));
|
||||
@@ -2013,6 +2117,7 @@ mod tests {
|
||||
attempt.report_context = Some(json!({
|
||||
"candidate_index": index,
|
||||
"retry_index": 0,
|
||||
"sticky_key_attempts": 1,
|
||||
"local_failover_policy": {
|
||||
"max_transfer_count": 1,
|
||||
"max_transfer_timeout_seconds": 0
|
||||
|
||||
@@ -67,27 +67,14 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
.push(key);
|
||||
}
|
||||
|
||||
let scheduling_mode = state
|
||||
.read_system_config_json_value("scheduling_mode")
|
||||
// Effective default scheduling: system-default routing group first, then
|
||||
// legacy system-config keys.
|
||||
let ordering_config = crate::scheduler::config::read_scheduler_ordering_config(state.app())
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|value| value.as_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| "cache_affinity".to_string());
|
||||
let priority_mode = state
|
||||
.read_system_config_json_value("provider_priority_mode")
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|value| value.as_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| "provider".to_string());
|
||||
let keep_priority_on_conversion = state
|
||||
.read_system_config_json_value("keep_priority_on_conversion")
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false);
|
||||
.unwrap_or_default();
|
||||
let scheduling_mode = ordering_config.scheduling_mode_str().to_string();
|
||||
let priority_mode = ordering_config.priority_mode_str().to_string();
|
||||
let keep_priority_on_conversion = ordering_config.keep_priority_on_conversion;
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs())
|
||||
|
||||
@@ -264,16 +264,10 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
|
||||
pub(super) async fn build_admin_monitoring_cache_snapshot(
|
||||
state: &AdminAppState<'_>,
|
||||
) -> Result<AdminMonitoringCacheSnapshot, GatewayError> {
|
||||
let scheduling_mode = state
|
||||
.read_system_config_json_value("scheduling_mode")
|
||||
.await?
|
||||
.and_then(|value| value.as_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| "cache_affinity".to_string());
|
||||
let provider_priority_mode = state
|
||||
.read_system_config_json_value("provider_priority_mode")
|
||||
.await?
|
||||
.and_then(|value| value.as_str().map(ToOwned::to_owned))
|
||||
.unwrap_or_else(|| "provider".to_string());
|
||||
let ordering_config =
|
||||
crate::scheduler::config::read_scheduler_ordering_config(state.app()).await?;
|
||||
let scheduling_mode = ordering_config.scheduling_mode_str().to_string();
|
||||
let provider_priority_mode = ordering_config.priority_mode_str().to_string();
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
let usage_summary = if state.has_usage_data_reader() {
|
||||
|
||||
@@ -2096,6 +2096,22 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
);
|
||||
}
|
||||
state.bootstrap_admin_from_env().await?;
|
||||
match state.ensure_system_default_routing_group().await {
|
||||
Ok(Some(group)) => {
|
||||
info!(
|
||||
group_id = %group.id,
|
||||
group_name = %group.name,
|
||||
"created system default routing group from legacy scheduler config"
|
||||
);
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
error = %err,
|
||||
"failed to bootstrap system default routing group; scheduler falls back to legacy system config"
|
||||
);
|
||||
}
|
||||
}
|
||||
match state.prewarm_chat_pii_redaction_runtime_config().await {
|
||||
Ok(enabled) => {
|
||||
info!(
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
use aether_ai_serving::{AiExecutionAttempt, STICKY_KEY_ATTEMPTS_REPORT_FIELD};
|
||||
use aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS;
|
||||
use aether_runtime_state::RuntimeLockLease;
|
||||
use aether_scheduler_core::parse_request_candidate_report_context;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::provider_transport::GatewayProviderTransportSnapshot;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) struct ExecutionAttemptIdentity {
|
||||
@@ -32,6 +33,9 @@ pub(crate) struct LocalExecutionCandidateMetadata {
|
||||
pub(crate) pool_key_index: Option<u32>,
|
||||
pub(crate) pool_key_lease: Option<RuntimeLockLease>,
|
||||
pub(crate) scheduler_affinity_epoch: Option<u64>,
|
||||
/// Routing-policy `sticky_key_attempts` in effect for this request. `None`
|
||||
/// means the policy default applies.
|
||||
pub(crate) sticky_key_attempts: Option<u32>,
|
||||
}
|
||||
|
||||
pub(crate) const SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD: &str = "scheduler_affinity_epoch";
|
||||
@@ -42,6 +46,11 @@ pub(crate) const POOL_KEY_LEASE_TOKEN_REPORT_FIELD: &str = "pool_key_lease_token
|
||||
pub(crate) const POOL_KEY_LEASE_FENCING_REPORT_FIELD: &str = "pool_key_lease_fencing_token";
|
||||
pub(crate) const POOL_KEY_LEASE_TTL_MS_REPORT_FIELD: &str = "pool_key_lease_ttl_ms";
|
||||
|
||||
/// Pool-expanded keys encode `pool_key_index * STRIDE + retry_index` into the
|
||||
/// persisted `retry_index` so a pool group's keys stay ordered in one candidate
|
||||
/// slot. Same-key retries on a pool key are therefore bounded by the stride.
|
||||
pub(crate) const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
|
||||
|
||||
pub(crate) fn attempt_identity_from_report_context(
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<ExecutionAttemptIdentity> {
|
||||
@@ -72,6 +81,10 @@ pub(crate) fn local_execution_candidate_metadata_from_report_context(
|
||||
scheduler_affinity_epoch: report_context
|
||||
.and_then(|value| value.get(SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD))
|
||||
.and_then(Value::as_u64),
|
||||
sticky_key_attempts: report_context
|
||||
.and_then(|value| value.get(STICKY_KEY_ATTEMPTS_REPORT_FIELD))
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|value| u32::try_from(value).ok()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -140,58 +153,55 @@ fn pool_key_lease_from_report_context(report_context: Option<&Value>) -> Option<
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_attempt_identities(
|
||||
candidate_index: u32,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Vec<ExecutionAttemptIdentity> {
|
||||
let attempt_slots = local_attempt_slot_count(transport);
|
||||
(0..attempt_slots)
|
||||
.map(|retry_index| ExecutionAttemptIdentity::new(candidate_index, retry_index))
|
||||
.collect()
|
||||
/// Retry index of the next same-key attempt, or `None` when the sticky-key
|
||||
/// budget for this candidate is used up.
|
||||
///
|
||||
/// Only the first-ranked candidate (index `0`, the cache-affinity sticky key)
|
||||
/// is retried on the same key; every later candidate gets exactly one attempt
|
||||
/// so that once failover has started it keeps advancing. `sticky_key_attempts`
|
||||
/// is the *total* attempt count on that key: `2` means one retry, `0` and `1`
|
||||
/// mean none. There is no upper bound: attempts are derived one at a time
|
||||
/// after each failure, never materialized ahead of time.
|
||||
///
|
||||
/// Inside a pool group only the first key (`pool_key_index == 0`) is treated
|
||||
/// as sticky, and its retries stay below `POOL_KEY_RETRY_INDEX_STRIDE` so the
|
||||
/// encoded retry index never collides with the next pool key.
|
||||
pub(crate) fn next_same_key_retry_index(
|
||||
identity: ExecutionAttemptIdentity,
|
||||
sticky_key_attempts: Option<u32>,
|
||||
) -> Option<u32> {
|
||||
if identity.candidate_index != 0 {
|
||||
return None;
|
||||
}
|
||||
let pool_limit = match identity.pool_key_index {
|
||||
None => u32::MAX,
|
||||
Some(0) => POOL_KEY_RETRY_INDEX_STRIDE,
|
||||
Some(_) => return None,
|
||||
};
|
||||
let budget = sticky_key_attempts.unwrap_or(DEFAULT_STICKY_KEY_ATTEMPTS);
|
||||
let attempts_so_far = identity.retry_index.checked_add(1)?;
|
||||
if attempts_so_far >= budget || attempts_so_far >= pool_limit {
|
||||
return None;
|
||||
}
|
||||
Some(attempts_so_far)
|
||||
}
|
||||
|
||||
pub(crate) fn local_attempt_slot_count(transport: &GatewayProviderTransportSnapshot) -> u32 {
|
||||
local_attempt_slots_from_transport(transport).unwrap_or(1)
|
||||
}
|
||||
|
||||
/// For endpoint/provider table fields, `2` is the legacy admin default and is
|
||||
/// treated as "not explicitly configured" so existing local-execution behaviour
|
||||
/// (one attempt slot per candidate) stays unchanged. Values `0`, `1`, and `>2`
|
||||
/// are treated as explicit.
|
||||
const LEGACY_DEFAULT_MAX_RETRIES: u32 = 2;
|
||||
|
||||
/// Upper bound on local attempt slots. This is intentionally stricter than
|
||||
/// admin max_retries validation to prevent unbounded pre-materialization from
|
||||
/// arbitrarily large JSON config values.
|
||||
const MAX_LOCAL_ATTEMPT_SLOTS: u32 = 99;
|
||||
|
||||
fn local_attempt_slots_from_transport(transport: &GatewayProviderTransportSnapshot) -> Option<u32> {
|
||||
let rules = transport
|
||||
.provider
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|config| config.get("failover_rules"))
|
||||
.and_then(Value::as_object);
|
||||
|
||||
rules
|
||||
.and_then(|value| value.get("max_retries"))
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|value| u32::try_from(value).ok())
|
||||
.or_else(|| {
|
||||
transport
|
||||
.endpoint
|
||||
.max_retries
|
||||
.and_then(|value| u32::try_from(value).ok())
|
||||
.filter(|&value| value != LEGACY_DEFAULT_MAX_RETRIES)
|
||||
})
|
||||
.or_else(|| {
|
||||
transport
|
||||
.provider
|
||||
.max_retries
|
||||
.and_then(|value| u32::try_from(value).ok())
|
||||
.filter(|&value| value != LEGACY_DEFAULT_MAX_RETRIES)
|
||||
})
|
||||
.map(|value| value.clamp(1, MAX_LOCAL_ATTEMPT_SLOTS))
|
||||
/// Derive the next same-key attempt for `attempt` after a candidate-scoped
|
||||
/// failure, reading the attempt identity and sticky budget from its report
|
||||
/// context. Returns `None` when no further same-key retry is allowed.
|
||||
pub(crate) fn next_same_key_retry_attempt<A: AiExecutionAttempt>(attempt: &A) -> Option<A> {
|
||||
let owned_report_context = attempt
|
||||
.report_context_ref()
|
||||
.is_none()
|
||||
.then(|| attempt.report_context())
|
||||
.flatten();
|
||||
let report_context = attempt
|
||||
.report_context_ref()
|
||||
.or(owned_report_context.as_ref());
|
||||
let identity = attempt_identity_from_report_context(report_context)?;
|
||||
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
|
||||
let retry_index = next_same_key_retry_index(identity, metadata.sticky_key_attempts)?;
|
||||
attempt.with_same_key_retry(retry_index, Uuid::new_v4().to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -199,277 +209,143 @@ mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
attempt_identity_from_report_context, build_local_attempt_identities,
|
||||
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
|
||||
LocalExecutionCandidateMetadata,
|
||||
};
|
||||
use crate::provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||
attempt_identity_from_report_context,
|
||||
local_execution_candidate_metadata_from_report_context, next_same_key_retry_attempt,
|
||||
next_same_key_retry_index, ExecutionAttemptIdentity, LocalExecutionCandidateMetadata,
|
||||
POOL_KEY_RETRY_INDEX_STRIDE,
|
||||
};
|
||||
use aether_ai_serving::{AiExecutionAttempt, AiSyncAttempt};
|
||||
use aether_runtime_state::RuntimeLockLease;
|
||||
|
||||
fn sample_transport(
|
||||
provider_max_retries: Option<i32>,
|
||||
endpoint_max_retries: Option<i32>,
|
||||
provider_config: Option<serde_json::Value>,
|
||||
) -> GatewayProviderTransportSnapshot {
|
||||
GatewayProviderTransportSnapshot {
|
||||
provider: GatewayProviderTransportProvider {
|
||||
id: "provider-1".to_string(),
|
||||
name: "OpenAI".to_string(),
|
||||
provider_type: "llm".to_string(),
|
||||
website: None,
|
||||
is_active: true,
|
||||
keep_priority_on_conversion: false,
|
||||
enable_format_conversion: true,
|
||||
concurrent_limit: None,
|
||||
max_retries: provider_max_retries,
|
||||
proxy: None,
|
||||
request_timeout_secs: None,
|
||||
stream_first_byte_timeout_secs: None,
|
||||
config: provider_config,
|
||||
},
|
||||
endpoint: GatewayProviderTransportEndpoint {
|
||||
id: "endpoint-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
api_format: "openai:chat".to_string(),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
is_active: true,
|
||||
base_url: "https://example.com".to_string(),
|
||||
header_rules: None,
|
||||
body_rules: None,
|
||||
max_retries: endpoint_max_retries,
|
||||
custom_path: None,
|
||||
config: None,
|
||||
format_acceptance_config: None,
|
||||
proxy: None,
|
||||
},
|
||||
key: GatewayProviderTransportKey {
|
||||
id: "key-1".to_string(),
|
||||
provider_id: "provider-1".to_string(),
|
||||
name: "primary".to_string(),
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
allow_auth_channel_mismatch_formats: None,
|
||||
#[test]
|
||||
fn first_candidate_defaults_to_one_same_key_retry() {
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 0), None),
|
||||
Some(1)
|
||||
);
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 1), None),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
global_priority_by_format: None,
|
||||
expires_at_unix_secs: None,
|
||||
proxy: None,
|
||||
fingerprint: None,
|
||||
upstream_metadata: None,
|
||||
decrypted_api_key: "secret".to_string(),
|
||||
decrypted_auth_config: None,
|
||||
},
|
||||
#[test]
|
||||
fn first_candidate_uses_policy_sticky_key_attempts_without_upper_bound() {
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 2), Some(3)),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 1), Some(3)),
|
||||
Some(2)
|
||||
);
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 4_999), Some(10_000)),
|
||||
Some(5_000)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_and_one_mean_single_attempt() {
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 0), Some(0)),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(0, 0), Some(1)),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failover_candidates_never_retry_on_the_same_key() {
|
||||
for candidate_index in 1..5 {
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(ExecutionAttemptIdentity::new(candidate_index, 0), None),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
next_same_key_retry_index(
|
||||
ExecutionAttemptIdentity::new(candidate_index, 0),
|
||||
Some(50)
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_defaults_to_single_attempt() {
|
||||
let identities = build_local_attempt_identities(3, &sample_transport(None, None, None));
|
||||
|
||||
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(3, 0)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_prefer_failover_rules_over_endpoint_and_provider() {
|
||||
let identities = build_local_attempt_identities(
|
||||
1,
|
||||
&sample_transport(
|
||||
Some(5),
|
||||
Some(4),
|
||||
Some(json!({
|
||||
"failover_rules": {
|
||||
"max_retries": 2
|
||||
}
|
||||
})),
|
||||
),
|
||||
);
|
||||
fn pool_groups_only_retry_their_first_key_within_the_stride() {
|
||||
let first_pool_key = ExecutionAttemptIdentity::new(0, 0).with_pool_key_index(Some(0));
|
||||
assert_eq!(next_same_key_retry_index(first_pool_key, Some(3)), Some(1));
|
||||
|
||||
let at_stride_limit = ExecutionAttemptIdentity::new(0, POOL_KEY_RETRY_INDEX_STRIDE - 1)
|
||||
.with_pool_key_index(Some(0));
|
||||
assert_eq!(
|
||||
identities,
|
||||
vec![
|
||||
ExecutionAttemptIdentity::new(1, 0),
|
||||
ExecutionAttemptIdentity::new(1, 1),
|
||||
]
|
||||
next_same_key_retry_index(at_stride_limit, Some(10_000)),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_falls_back_to_endpoint_max_retries() {
|
||||
let identities =
|
||||
build_local_attempt_identities(2, &sample_transport(Some(5), Some(3), None));
|
||||
|
||||
let second_pool_key = ExecutionAttemptIdentity::new(0, POOL_KEY_RETRY_INDEX_STRIDE)
|
||||
.with_pool_key_index(Some(1));
|
||||
assert_eq!(
|
||||
identities,
|
||||
vec![
|
||||
ExecutionAttemptIdentity::new(2, 0),
|
||||
ExecutionAttemptIdentity::new(2, 1),
|
||||
ExecutionAttemptIdentity::new(2, 2),
|
||||
]
|
||||
next_same_key_retry_index(second_pool_key, Some(10_000)),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_falls_back_to_provider_max_retries() {
|
||||
let identities = build_local_attempt_identities(0, &sample_transport(Some(4), None, None));
|
||||
fn next_same_key_retry_attempt_rewrites_candidate_id_and_retry_index() {
|
||||
let attempt = AiSyncAttempt {
|
||||
plan: aether_contracts::ExecutionPlan {
|
||||
request_id: "trace-1".to_string(),
|
||||
candidate_id: Some("candidate-a".to_string()),
|
||||
provider_name: None,
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://example.com".to_string(),
|
||||
headers: Default::default(),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: aether_contracts::RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: false,
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
provider_api_format: "openai:chat".to_string(),
|
||||
model_name: None,
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
},
|
||||
report_kind: None,
|
||||
report_context: Some(json!({
|
||||
"candidate_id": "candidate-a",
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"sticky_key_attempts": 2,
|
||||
})),
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
identities,
|
||||
vec![
|
||||
ExecutionAttemptIdentity::new(0, 0),
|
||||
ExecutionAttemptIdentity::new(0, 1),
|
||||
ExecutionAttemptIdentity::new(0, 2),
|
||||
ExecutionAttemptIdentity::new(0, 3),
|
||||
]
|
||||
let retry = next_same_key_retry_attempt(&attempt).expect("one same-key retry remains");
|
||||
let retry_candidate_id = retry.plan.candidate_id.clone().expect("fresh candidate id");
|
||||
assert_ne!(retry_candidate_id, "candidate-a");
|
||||
assert_eq!(retry.plan.key_id, "key-1");
|
||||
let context = retry.report_context_ref().expect("context retained");
|
||||
assert_eq!(context["candidate_id"], json!(retry_candidate_id));
|
||||
assert_eq!(context["retry_index"], json!(1));
|
||||
assert_eq!(context["candidate_index"], json!(0));
|
||||
|
||||
assert!(
|
||||
next_same_key_retry_attempt(&retry).is_none(),
|
||||
"budget of 2 attempts is exhausted after one retry"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_endpoint_overrides_provider() {
|
||||
let identities =
|
||||
build_local_attempt_identities(7, &sample_transport(Some(10), Some(3), None));
|
||||
|
||||
assert_eq!(
|
||||
identities,
|
||||
vec![
|
||||
ExecutionAttemptIdentity::new(7, 0),
|
||||
ExecutionAttemptIdentity::new(7, 1),
|
||||
ExecutionAttemptIdentity::new(7, 2),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_default_two_treated_as_unset() {
|
||||
let identities =
|
||||
build_local_attempt_identities(5, &sample_transport(Some(2), Some(2), None));
|
||||
|
||||
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(5, 0)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_endpoint_two_falls_back_to_provider_ten() {
|
||||
let identities =
|
||||
build_local_attempt_identities(1, &sample_transport(Some(10), Some(2), None));
|
||||
|
||||
assert_eq!(
|
||||
identities,
|
||||
vec![
|
||||
ExecutionAttemptIdentity::new(1, 0),
|
||||
ExecutionAttemptIdentity::new(1, 1),
|
||||
ExecutionAttemptIdentity::new(1, 2),
|
||||
ExecutionAttemptIdentity::new(1, 3),
|
||||
ExecutionAttemptIdentity::new(1, 4),
|
||||
ExecutionAttemptIdentity::new(1, 5),
|
||||
ExecutionAttemptIdentity::new(1, 6),
|
||||
ExecutionAttemptIdentity::new(1, 7),
|
||||
ExecutionAttemptIdentity::new(1, 8),
|
||||
ExecutionAttemptIdentity::new(1, 9),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_failover_rules_zero_produces_one_slot() {
|
||||
let identities = build_local_attempt_identities(
|
||||
1,
|
||||
&sample_transport(
|
||||
Some(5),
|
||||
Some(4),
|
||||
Some(json!({
|
||||
"failover_rules": {
|
||||
"max_retries": 0
|
||||
}
|
||||
})),
|
||||
),
|
||||
);
|
||||
|
||||
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(1, 0)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_endpoint_zero_produces_one_slot() {
|
||||
let identities =
|
||||
build_local_attempt_identities(3, &sample_transport(Some(5), Some(0), None));
|
||||
|
||||
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(3, 0)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_provider_zero_produces_one_slot() {
|
||||
let identities = build_local_attempt_identities(3, &sample_transport(Some(0), None, None));
|
||||
|
||||
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(3, 0)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_provider_ten_creates_ten_slots() {
|
||||
let identities = build_local_attempt_identities(2, &sample_transport(Some(10), None, None));
|
||||
|
||||
assert_eq!(identities.len(), 10);
|
||||
assert_eq!(identities[0], ExecutionAttemptIdentity::new(2, 0));
|
||||
assert_eq!(identities[9], ExecutionAttemptIdentity::new(2, 9));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_failover_rules_over_limit_clamped_to_max() {
|
||||
let identities = build_local_attempt_identities(
|
||||
0,
|
||||
&sample_transport(
|
||||
Some(3),
|
||||
Some(5),
|
||||
Some(json!({
|
||||
"failover_rules": {
|
||||
"max_retries": 1000
|
||||
}
|
||||
})),
|
||||
),
|
||||
);
|
||||
|
||||
assert_eq!(identities.len(), 99);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_failover_rules_u32_max_clamped_to_max() {
|
||||
let identities = build_local_attempt_identities(
|
||||
0,
|
||||
&sample_transport(
|
||||
None,
|
||||
None,
|
||||
Some(json!({
|
||||
"failover_rules": {
|
||||
"max_retries": u32::MAX
|
||||
}
|
||||
})),
|
||||
),
|
||||
);
|
||||
|
||||
assert_eq!(identities.len(), 99);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_endpoint_over_limit_clamped_to_max() {
|
||||
let identities =
|
||||
build_local_attempt_identities(0, &sample_transport(None, Some(2000), None));
|
||||
|
||||
assert_eq!(identities.len(), 99);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_local_attempt_identities_provider_over_limit_clamped_to_max() {
|
||||
let identities =
|
||||
build_local_attempt_identities(0, &sample_transport(Some(5000), None, None));
|
||||
|
||||
assert_eq!(identities.len(), 99);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_attempt_identity_from_report_context_reads_candidate_and_retry_indices() {
|
||||
let identity = attempt_identity_from_report_context(Some(&json!({
|
||||
@@ -499,6 +375,7 @@ mod tests {
|
||||
"pool_key_lease_token": "gateway-1:token-1",
|
||||
"pool_key_lease_fencing_token": 7,
|
||||
"pool_key_lease_ttl_ms": 900000,
|
||||
"sticky_key_attempts": 3,
|
||||
})));
|
||||
|
||||
assert_eq!(
|
||||
@@ -514,6 +391,7 @@ mod tests {
|
||||
ttl_ms: 900000,
|
||||
}),
|
||||
scheduler_affinity_epoch: None,
|
||||
sticky_key_attempts: Some(3),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -20,11 +20,10 @@ pub(crate) use self::adaptive::{
|
||||
LocalAdaptiveRateLimitProjection, LocalAdaptiveSuccessProjection,
|
||||
};
|
||||
pub(crate) use self::attempt::{
|
||||
attempt_identity_from_report_context, build_local_attempt_identities,
|
||||
insert_pool_key_lease_report_context_fields, local_attempt_slot_count,
|
||||
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
|
||||
LocalExecutionCandidateMetadata, ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD,
|
||||
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
|
||||
attempt_identity_from_report_context, insert_pool_key_lease_report_context_fields,
|
||||
local_execution_candidate_metadata_from_report_context, next_same_key_retry_attempt,
|
||||
ExecutionAttemptIdentity, LocalExecutionCandidateMetadata, POOL_KEY_RETRY_INDEX_STRIDE,
|
||||
ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
|
||||
};
|
||||
pub(crate) use self::classifier::{
|
||||
classify_anthropic_failure_disposition, classify_failure_disposition, classify_local_failover,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use aether_routing_core::{
|
||||
resolve_routing_policy, MutationPlan, RankingOverlay, ResolvedRoutingPolicy,
|
||||
RoutingGroupConfig, RoutingPolicyInput, RoutingRulePhase, RoutingSchedulingMode,
|
||||
RoutingSetPriorityMode,
|
||||
RoutingDefaultPolicy, RoutingGroupConfig, RoutingPolicyInput, RoutingRulePhase,
|
||||
RoutingSchedulingMode, RoutingSetPriorityMode, DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
};
|
||||
use http::StatusCode;
|
||||
use serde_json::Value;
|
||||
@@ -81,9 +81,7 @@ pub(crate) fn resolve_gateway_routing_policy(
|
||||
pub(crate) fn resolve_gateway_static_default_routing_policy(
|
||||
input: GatewayStaticRoutingPolicyInput<'_>,
|
||||
) -> Result<Option<ResolvedRoutingPolicy>, GatewayError> {
|
||||
let Some((priority_mode, scheduling_mode, keep_priority_on_conversion)) =
|
||||
static_default_policy_fields(input.group_config_json)?
|
||||
else {
|
||||
let Some(default_policy) = static_default_policy_fields(input.group_config_json)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
@@ -93,9 +91,10 @@ pub(crate) fn resolve_gateway_static_default_routing_policy(
|
||||
selection_source: input.selection_source.to_string(),
|
||||
requested_model: input.requested_model.to_string(),
|
||||
resolved_model: input.resolved_model.to_string(),
|
||||
priority_mode,
|
||||
scheduling_mode,
|
||||
keep_priority_on_conversion,
|
||||
priority_mode: default_policy.priority_mode,
|
||||
scheduling_mode: default_policy.scheduling_mode,
|
||||
keep_priority_on_conversion: default_policy.keep_priority_on_conversion,
|
||||
sticky_key_attempts: default_policy.sticky_key_attempts,
|
||||
ranking_overlay: RankingOverlay::default(),
|
||||
mutation_plan: MutationPlan::default(),
|
||||
pool_policy_overrides: BTreeMap::new(),
|
||||
@@ -105,7 +104,7 @@ pub(crate) fn resolve_gateway_static_default_routing_policy(
|
||||
|
||||
fn static_default_policy_fields(
|
||||
config_json: &Value,
|
||||
) -> Result<Option<(RoutingSetPriorityMode, RoutingSchedulingMode, bool)>, GatewayError> {
|
||||
) -> Result<Option<RoutingDefaultPolicy>, GatewayError> {
|
||||
let Some(object) = config_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -117,11 +116,7 @@ fn static_default_policy_fields(
|
||||
}
|
||||
|
||||
let Some(default_policy) = object.get("default_policy") else {
|
||||
return Ok(Some((
|
||||
RoutingSetPriorityMode::default(),
|
||||
RoutingSchedulingMode::default(),
|
||||
false,
|
||||
)));
|
||||
return Ok(Some(RoutingDefaultPolicy::default()));
|
||||
};
|
||||
let Some(default_policy) = default_policy.as_object() else {
|
||||
return Ok(None);
|
||||
@@ -141,12 +136,22 @@ fn static_default_policy_fields(
|
||||
})?,
|
||||
None => false,
|
||||
};
|
||||
let sticky_key_attempts = match default_policy.get("sticky_key_attempts") {
|
||||
Some(value) => value
|
||||
.as_u64()
|
||||
.and_then(|value| u32::try_from(value).ok())
|
||||
.ok_or_else(|| {
|
||||
invalid_routing_group_config("sticky_key_attempts must be a non-negative integer")
|
||||
})?,
|
||||
None => DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
};
|
||||
|
||||
Ok(Some((
|
||||
Ok(Some(RoutingDefaultPolicy {
|
||||
priority_mode,
|
||||
scheduling_mode,
|
||||
keep_priority_on_conversion,
|
||||
)))
|
||||
sticky_key_attempts,
|
||||
}))
|
||||
}
|
||||
|
||||
fn routing_array_field_is_missing_or_empty(
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use self::selection::{
|
||||
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
|
||||
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons_and_ordering,
|
||||
collect_selectable_enumerated_candidates_with_skip_reasons,
|
||||
resolve_preselection_ordering_config,
|
||||
};
|
||||
use super::config::SchedulerOrderingConfig;
|
||||
use super::state::SchedulerRuntimeState;
|
||||
|
||||
mod affinity;
|
||||
@@ -53,6 +55,10 @@ enum RequiredCapabilityMatchMode {
|
||||
Exclusive,
|
||||
}
|
||||
|
||||
/// `ordering_config` carries the request's routing-policy derived scheduler
|
||||
/// config. `None` falls back to the runtime default (system-default routing
|
||||
/// group, then legacy system-config keys).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -64,6 +70,7 @@ pub(crate) async fn list_selectable_candidates(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
collect_selectable_candidates(
|
||||
selection_row_source,
|
||||
@@ -76,6 +83,7 @@ pub(crate) async fn list_selectable_candidates(
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -87,6 +95,7 @@ pub(crate) fn is_exact_all_skipped_by_auth_limit(
|
||||
selection::is_exact_all_skipped_by_auth_limit(selected, skipped)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -98,6 +107,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -105,7 +115,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
collect_selectable_candidates_with_skip_reasons(
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
@@ -117,10 +127,12 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
|
||||
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(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -133,6 +145,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
request_operation: Option<&str>,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -140,7 +153,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
collect_selectable_candidates_with_skip_reasons(
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
@@ -152,6 +165,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
request_operation,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -166,6 +180,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -173,7 +188,8 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
|
||||
let ordering_config =
|
||||
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
|
||||
let priority_affinity_key = selection::scheduling_priority_affinity_key(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
@@ -194,6 +210,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -203,6 +220,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
Ok(
|
||||
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
@@ -214,12 +232,14 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
ordering_config,
|
||||
)
|
||||
.await?
|
||||
.0,
|
||||
)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -229,6 +249,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
|
||||
let normalized_api_format = normalize_api_format(candidate_api_format);
|
||||
if normalized_api_format.is_empty() {
|
||||
@@ -262,20 +283,22 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
|
||||
let mut all_attempts_blocked_by_auth_limit = !model_names.is_empty();
|
||||
|
||||
for global_model_name in model_names {
|
||||
let (candidates, skipped_candidates) = collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
&normalized_api_format,
|
||||
&global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
false,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
let (candidates, skipped_candidates) =
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
&normalized_api_format,
|
||||
&global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities.as_ref(),
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
false,
|
||||
None,
|
||||
ordering_config,
|
||||
)
|
||||
.await?;
|
||||
all_attempts_blocked_by_auth_limit &=
|
||||
is_exact_all_skipped_by_auth_limit(&candidates, &skipped_candidates);
|
||||
match capability_mode {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
|
||||
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
|
||||
use crate::scheduler::config::SchedulerSchedulingMode;
|
||||
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
|
||||
use crate::GatewayError;
|
||||
use aether_scheduler_core::ClientSessionAffinity;
|
||||
|
||||
@@ -115,6 +115,7 @@ pub(super) async fn select_minimal_candidate(
|
||||
Ok(selected)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn collect_selectable_candidates(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -126,24 +127,32 @@ pub(super) async fn collect_selectable_candidates(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
Ok(collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
None,
|
||||
Ok(
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
None,
|
||||
ordering_config,
|
||||
)
|
||||
.await?
|
||||
.0,
|
||||
)
|
||||
.await?
|
||||
.0)
|
||||
}
|
||||
|
||||
/// Legacy-shaped entrypoint that resolves the ordering config from the
|
||||
/// runtime state. Prefer `collect_selectable_candidates_with_skip_reasons_and_ordering`
|
||||
/// and pass the request's routing-policy config explicitly.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -163,7 +172,59 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
request_operation,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Resolve the ordering config for a preselection pass: the routing-policy
|
||||
/// derived config wins when the caller has one; otherwise fall back to the
|
||||
/// runtime default (system-default routing group, then legacy keys).
|
||||
pub(super) async fn resolve_preselection_ordering_config(
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<SchedulerOrderingConfig, GatewayError> {
|
||||
match ordering_config {
|
||||
Some(config) => Ok(config),
|
||||
None => runtime_state.read_scheduler_ordering_config().await,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
api_format: &str,
|
||||
global_model_name: &str,
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
request_operation: Option<&str>,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
Vec<SchedulerSkippedCandidate>,
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
let ordering_config =
|
||||
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
|
||||
let priority_affinity_key = scheduling_priority_affinity_key(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
@@ -205,7 +266,7 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: crate::scheduler::config::SchedulerOrderingConfig,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
priority_affinity_key: Option<&str>,
|
||||
) -> Result<
|
||||
(
|
||||
|
||||
@@ -65,6 +65,7 @@ async fn compatible_required_capability_prefers_matching_keys_without_hard_filte
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -120,6 +121,7 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -196,6 +198,7 @@ async fn required_capability_without_model_uses_session_scoped_affinity() {
|
||||
Some(&auth_snapshot),
|
||||
Some(&client_session_affinity),
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -273,6 +276,7 @@ async fn required_capability_reports_auth_limit_signal_when_every_model_is_block
|
||||
Some(&auth_snapshot),
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
@@ -75,6 +75,7 @@ async fn collect_selectable_candidates(
|
||||
None,
|
||||
now_unix_secs,
|
||||
false,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -1,4 +1,10 @@
|
||||
use aether_data_contracts::repository::routing_profiles::RoutingGroupLookupKey;
|
||||
use aether_routing_core::{
|
||||
ResolvedRoutingPolicy, RoutingDefaultPolicy, RoutingSchedulingMode, RoutingSetPriorityMode,
|
||||
DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerPriorityMode;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
@@ -10,11 +16,30 @@ pub(crate) enum SchedulerSchedulingMode {
|
||||
LoadBalance,
|
||||
}
|
||||
|
||||
impl SchedulerSchedulingMode {
|
||||
pub(crate) fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::FixedOrder => "fixed_order",
|
||||
Self::CacheAffinity => "cache_affinity",
|
||||
Self::LoadBalance => "load_balance",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn scheduler_priority_mode_as_str(mode: SchedulerPriorityMode) -> &'static str {
|
||||
match mode {
|
||||
SchedulerPriorityMode::Provider => "provider",
|
||||
SchedulerPriorityMode::GlobalKey => "global_key",
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) struct SchedulerOrderingConfig {
|
||||
pub(crate) priority_mode: SchedulerPriorityMode,
|
||||
pub(crate) scheduling_mode: SchedulerSchedulingMode,
|
||||
pub(crate) keep_priority_on_conversion: bool,
|
||||
/// Total attempts on the first-ranked (sticky) candidate before failover.
|
||||
pub(crate) sticky_key_attempts: u32,
|
||||
}
|
||||
|
||||
impl Default for SchedulerOrderingConfig {
|
||||
@@ -23,10 +48,72 @@ impl Default for SchedulerOrderingConfig {
|
||||
priority_mode: SchedulerPriorityMode::Provider,
|
||||
scheduling_mode: SchedulerSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SchedulerOrderingConfig {
|
||||
/// Ordering config derived from a resolved routing policy. The policy is
|
||||
/// the single source of truth: no legacy system-config value is merged in.
|
||||
pub(crate) fn from_routing_policy(policy: &ResolvedRoutingPolicy) -> Self {
|
||||
Self {
|
||||
priority_mode: scheduler_priority_mode_from_routing(policy.priority_mode),
|
||||
scheduling_mode: scheduler_scheduling_mode_from_routing(policy.scheduling_mode),
|
||||
keep_priority_on_conversion: policy.keep_priority_on_conversion,
|
||||
sticky_key_attempts: policy.sticky_key_attempts,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn from_routing_default_policy(policy: &RoutingDefaultPolicy) -> Self {
|
||||
Self {
|
||||
priority_mode: scheduler_priority_mode_from_routing(policy.priority_mode),
|
||||
scheduling_mode: scheduler_scheduling_mode_from_routing(policy.scheduling_mode),
|
||||
keep_priority_on_conversion: policy.keep_priority_on_conversion,
|
||||
sticky_key_attempts: policy.sticky_key_attempts,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn to_routing_default_policy(self) -> RoutingDefaultPolicy {
|
||||
RoutingDefaultPolicy {
|
||||
priority_mode: match self.priority_mode {
|
||||
SchedulerPriorityMode::Provider => RoutingSetPriorityMode::Provider,
|
||||
SchedulerPriorityMode::GlobalKey => RoutingSetPriorityMode::GlobalKey,
|
||||
},
|
||||
scheduling_mode: match self.scheduling_mode {
|
||||
SchedulerSchedulingMode::FixedOrder => RoutingSchedulingMode::FixedOrder,
|
||||
SchedulerSchedulingMode::CacheAffinity => RoutingSchedulingMode::CacheAffinity,
|
||||
SchedulerSchedulingMode::LoadBalance => RoutingSchedulingMode::LoadBalance,
|
||||
},
|
||||
keep_priority_on_conversion: self.keep_priority_on_conversion,
|
||||
sticky_key_attempts: self.sticky_key_attempts,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn priority_mode_str(self) -> &'static str {
|
||||
scheduler_priority_mode_as_str(self.priority_mode)
|
||||
}
|
||||
|
||||
pub(crate) fn scheduling_mode_str(self) -> &'static str {
|
||||
self.scheduling_mode.as_str()
|
||||
}
|
||||
}
|
||||
|
||||
fn scheduler_priority_mode_from_routing(mode: RoutingSetPriorityMode) -> SchedulerPriorityMode {
|
||||
match mode {
|
||||
RoutingSetPriorityMode::Provider => SchedulerPriorityMode::Provider,
|
||||
RoutingSetPriorityMode::GlobalKey => SchedulerPriorityMode::GlobalKey,
|
||||
}
|
||||
}
|
||||
|
||||
fn scheduler_scheduling_mode_from_routing(mode: RoutingSchedulingMode) -> SchedulerSchedulingMode {
|
||||
match mode {
|
||||
RoutingSchedulingMode::FixedOrder => SchedulerSchedulingMode::FixedOrder,
|
||||
RoutingSchedulingMode::CacheAffinity => SchedulerSchedulingMode::CacheAffinity,
|
||||
RoutingSchedulingMode::LoadBalance => SchedulerSchedulingMode::LoadBalance,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_scheduler_priority_mode(
|
||||
value: Option<&serde_json::Value>,
|
||||
) -> SchedulerPriorityMode {
|
||||
@@ -62,8 +149,62 @@ pub(crate) fn parse_scheduler_scheduling_mode(
|
||||
}
|
||||
}
|
||||
|
||||
/// Effective scheduler ordering config for requests that carry no resolved
|
||||
/// routing policy.
|
||||
///
|
||||
/// Resolution order:
|
||||
/// 1. the enabled system-default routing group's `default_policy`;
|
||||
/// 2. the legacy system-config keys (`provider_priority_mode`,
|
||||
/// `scheduling_mode`, `keep_priority_on_conversion`).
|
||||
///
|
||||
/// Step 2 only exists so deployments that never created a routing group keep
|
||||
/// their behaviour; once the legacy keys are removed this function collapses
|
||||
/// to step 1 plus `SchedulerOrderingConfig::default()`.
|
||||
pub(crate) async fn read_scheduler_ordering_config(
|
||||
state: &AppState,
|
||||
) -> Result<SchedulerOrderingConfig, GatewayError> {
|
||||
if let Some(config) = read_system_default_routing_ordering_config(state).await? {
|
||||
return Ok(config);
|
||||
}
|
||||
read_legacy_scheduler_ordering_config(state).await
|
||||
}
|
||||
|
||||
/// Ordering config from the enabled system-default routing group, if any.
|
||||
pub(crate) async fn read_system_default_routing_ordering_config(
|
||||
state: &AppState,
|
||||
) -> Result<Option<SchedulerOrderingConfig>, GatewayError> {
|
||||
let Some(group) = state
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await?
|
||||
.filter(|group| group.enabled)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let default_policy = match group.config_json.get("default_policy") {
|
||||
None | Some(serde_json::Value::Null) => RoutingDefaultPolicy::default(),
|
||||
Some(value) => match serde_json::from_value::<RoutingDefaultPolicy>(value.clone()) {
|
||||
Ok(policy) => policy,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "scheduler_system_default_routing_policy_invalid",
|
||||
log_type = "event",
|
||||
group_id = %group.id,
|
||||
error = %error,
|
||||
"system default routing group has an invalid default_policy; ignoring it"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
},
|
||||
};
|
||||
Ok(Some(SchedulerOrderingConfig::from_routing_default_policy(
|
||||
&default_policy,
|
||||
)))
|
||||
}
|
||||
|
||||
/// Legacy system-config based ordering config. Kept only as a migration
|
||||
/// fallback; see `read_scheduler_ordering_config`.
|
||||
pub(crate) async fn read_legacy_scheduler_ordering_config(
|
||||
state: &AppState,
|
||||
) -> Result<SchedulerOrderingConfig, GatewayError> {
|
||||
let priority_mode = parse_scheduler_priority_mode(
|
||||
state
|
||||
@@ -87,5 +228,167 @@ pub(crate) async fn read_scheduler_ordering_config(
|
||||
priority_mode,
|
||||
scheduling_mode,
|
||||
keep_priority_on_conversion,
|
||||
// Legacy config never carried a sticky-key setting; use the routing default.
|
||||
sticky_key_attempts: DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupRecord, RoutingGroupLookupKey, RoutingGroupReadRepository,
|
||||
RoutingGroupWriteRepository,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
fn legacy_values() -> [(String, serde_json::Value); 3] {
|
||||
[
|
||||
("provider_priority_mode".to_string(), json!("global_key")),
|
||||
("scheduling_mode".to_string(), json!("load_balance")),
|
||||
("keep_priority_on_conversion".to_string(), json!(true)),
|
||||
]
|
||||
}
|
||||
|
||||
async fn create_system_default(
|
||||
repository: &InMemoryRoutingGroupRepository,
|
||||
enabled: bool,
|
||||
config_json: serde_json::Value,
|
||||
) {
|
||||
repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "system-default".to_string(),
|
||||
name: "system-default".to_string(),
|
||||
description: None,
|
||||
enabled,
|
||||
is_system_default: true,
|
||||
config_json,
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn system_default_routing_group_overrides_legacy_keys() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(
|
||||
&repository,
|
||||
true,
|
||||
json!({
|
||||
"default_policy": {
|
||||
"priority_mode": "provider",
|
||||
"scheduling_mode": "fixed_order",
|
||||
"keep_priority_on_conversion": false
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
|
||||
assert_eq!(config.priority_mode, SchedulerPriorityMode::Provider);
|
||||
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::FixedOrder);
|
||||
assert!(!config.keep_priority_on_conversion);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_default_policy_in_system_default_group_uses_routing_defaults() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(&repository, true, json!({})).await;
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
|
||||
assert_eq!(config, SchedulerOrderingConfig::default());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_or_missing_system_default_group_falls_back_to_legacy_keys() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(
|
||||
&repository,
|
||||
false,
|
||||
json!({"default_policy": {"scheduling_mode": "fixed_order"}}),
|
||||
)
|
||||
.await;
|
||||
let with_disabled_group = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
let without_repository = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled().with_system_config_values_for_tests(legacy_values()),
|
||||
);
|
||||
|
||||
for state in [with_disabled_group, without_repository] {
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
assert_eq!(config.priority_mode, SchedulerPriorityMode::GlobalKey);
|
||||
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::LoadBalance);
|
||||
assert!(config.keep_priority_on_conversion);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bootstrap_creates_system_default_group_from_legacy_keys_once() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests(legacy_values())
|
||||
.with_routing_group_repository_for_tests(repository.clone()),
|
||||
);
|
||||
|
||||
let created = state
|
||||
.ensure_system_default_routing_group_inner()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("first bootstrap should create the system default group");
|
||||
assert!(created.enabled);
|
||||
assert!(created.is_system_default);
|
||||
assert_eq!(
|
||||
created.config_json["default_policy"],
|
||||
json!({
|
||||
"priority_mode": "global_key",
|
||||
"scheduling_mode": "load_balance",
|
||||
"keep_priority_on_conversion": true,
|
||||
"sticky_key_attempts": DEFAULT_STICKY_KEY_ATTEMPTS
|
||||
})
|
||||
);
|
||||
|
||||
let second = state
|
||||
.ensure_system_default_routing_group_inner()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(second.is_none(), "bootstrap must be idempotent");
|
||||
assert_eq!(
|
||||
repository
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await
|
||||
.unwrap()
|
||||
.map(|group| group.id),
|
||||
Some(created.id)
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
assert_eq!(config.priority_mode, SchedulerPriorityMode::GlobalKey);
|
||||
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::LoadBalance);
|
||||
assert!(config.keep_priority_on_conversion);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,11 +4,89 @@ use aether_data_contracts::repository::routing_profiles::{
|
||||
StoredRoutingGroup, StoredRoutingGroupBinding, StoredRoutingGroupVersion,
|
||||
UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
|
||||
};
|
||||
use aether_routing_core::RoutingGroupConfig;
|
||||
use std::sync::Arc;
|
||||
use tracing::warn;
|
||||
|
||||
use super::{AppState, GatewayError};
|
||||
|
||||
const BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME: &str = "system-default";
|
||||
|
||||
impl AppState {
|
||||
/// Make sure an enabled system-default routing group exists.
|
||||
///
|
||||
/// When none exists, one is created from the legacy scheduler system-config
|
||||
/// keys so that removing those keys later does not change behaviour. Returns
|
||||
/// the created group, or `None` when nothing had to be created (no routing
|
||||
/// storage, no writer, or a system default already exists).
|
||||
pub async fn ensure_system_default_routing_group(
|
||||
&self,
|
||||
) -> Result<Option<StoredRoutingGroup>, std::io::Error> {
|
||||
self.ensure_system_default_routing_group_inner()
|
||||
.await
|
||||
.map_err(|err| std::io::Error::other(format!("{err:?}")))
|
||||
}
|
||||
|
||||
pub(crate) async fn ensure_system_default_routing_group_inner(
|
||||
&self,
|
||||
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
|
||||
if !self.has_routing_group_data_reader() {
|
||||
return Ok(None);
|
||||
}
|
||||
if self
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await?
|
||||
.is_some()
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
if !self.has_routing_group_data_writer() {
|
||||
warn!(
|
||||
event_name = "routing_system_default_bootstrap_skipped",
|
||||
log_type = "event",
|
||||
"no system default routing group exists and routing storage is read-only; scheduler falls back to legacy system config"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let legacy = crate::scheduler::config::read_legacy_scheduler_ordering_config(self).await?;
|
||||
let config = RoutingGroupConfig {
|
||||
default_policy: legacy.to_routing_default_policy(),
|
||||
..RoutingGroupConfig::default()
|
||||
};
|
||||
let config_json = serde_json::to_value(config)
|
||||
.map_err(|err| GatewayError::Internal(format!("serialize routing config: {err}")))?;
|
||||
|
||||
let name = if self
|
||||
.find_routing_group(RoutingGroupLookupKey::Name(
|
||||
BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME,
|
||||
))
|
||||
.await?
|
||||
.is_some()
|
||||
{
|
||||
format!(
|
||||
"{BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME}-{}",
|
||||
&uuid::Uuid::new_v4().simple().to_string()[..8]
|
||||
)
|
||||
} else {
|
||||
BOOTSTRAP_SYSTEM_DEFAULT_ROUTING_GROUP_NAME.to_string()
|
||||
};
|
||||
let now = crate::clock::current_unix_secs() as i64;
|
||||
self.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
name,
|
||||
description: Some("自动从旧版调度配置迁移生成的系统默认策略".to_string()),
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
config_json,
|
||||
version: 1,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
published_at: Some(now),
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn has_routing_group_data_reader(&self) -> bool {
|
||||
self.data.has_routing_group_reader()
|
||||
}
|
||||
|
||||
@@ -2220,7 +2220,9 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
|
||||
.to_string(),
|
||||
});
|
||||
|
||||
let frames = if attempt == 1 {
|
||||
// The primary key gets two attempts under the default
|
||||
// sticky_key_attempts; both must fail to reach the backup.
|
||||
let frames = if attempt <= 2 {
|
||||
concat!(
|
||||
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":429,\"headers\":{\"content-type\":\"application/json\"}}}\n",
|
||||
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"error\\\":{\\\"message\\\":\\\"rate limited\\\",\\\"type\\\":\\\"rate_limit_error\\\"}}\"}}\n",
|
||||
@@ -2362,14 +2364,25 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.len()
|
||||
>= 2
|
||||
>= 3
|
||||
})
|
||||
.await;
|
||||
let seen_execution_runtime_requests = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone();
|
||||
assert_eq!(seen_execution_runtime_requests.len(), 2);
|
||||
// Default sticky_key_attempts is 2: the primary key is retried once on
|
||||
// the same key, then failover moves to the backup with a single attempt.
|
||||
assert_eq!(seen_execution_runtime_requests.len(), 3);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_requests
|
||||
.iter()
|
||||
.filter(|request| {
|
||||
request.url == "https://api.openai.primary.example/chat/completions"
|
||||
})
|
||||
.count(),
|
||||
2
|
||||
);
|
||||
let primary_request = seen_execution_runtime_requests
|
||||
.iter()
|
||||
.find(|request| request.url == "https://api.openai.primary.example/chat/completions")
|
||||
@@ -2404,7 +2417,15 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
|
||||
.list_by_request_id("trace-openai-chat-local-stream-failover-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 2);
|
||||
assert_eq!(stored_candidates.len(), 3);
|
||||
assert_eq!(
|
||||
stored_candidates
|
||||
.iter()
|
||||
.filter(|candidate| candidate.status == RequestCandidateStatus::Failed)
|
||||
.count(),
|
||||
2,
|
||||
"both sticky-key attempts on the primary should be recorded as failed"
|
||||
);
|
||||
let failed_candidate = stored_candidates
|
||||
.iter()
|
||||
.find(|candidate| {
|
||||
@@ -2465,7 +2486,7 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
|
||||
|
||||
assert_eq!(
|
||||
execution_runtime_hits.load(std::sync::atomic::Ordering::SeqCst),
|
||||
2
|
||||
3
|
||||
);
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -112,7 +112,7 @@ async fn gateway_skips_unsupported_local_openai_chat_sync_candidate_before_tryin
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
@@ -134,7 +134,7 @@ async fn gateway_skips_unsupported_local_openai_chat_sync_candidate_before_tryin
|
||||
"https://api.openai.skip.example".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
@@ -520,7 +520,7 @@ async fn gateway_surfaces_local_execution_runtime_miss_reason_when_all_openai_ch
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
@@ -542,7 +542,7 @@ async fn gateway_surfaces_local_execution_runtime_miss_reason_when_all_openai_ch
|
||||
"https://chatgpt.com/backend-api/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
@@ -802,7 +802,7 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
@@ -828,7 +828,7 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
|
||||
base_url.to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
@@ -970,7 +970,9 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
|
||||
.to_string(),
|
||||
});
|
||||
|
||||
if attempt == 1 {
|
||||
// The primary key gets two attempts under the default
|
||||
// sticky_key_attempts; both must fail to reach the backup.
|
||||
if attempt <= 2 {
|
||||
return Json(json!({
|
||||
"request_id": "trace-openai-chat-local-failover-123",
|
||||
"status_code": 401,
|
||||
@@ -1125,56 +1127,63 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone();
|
||||
assert_eq!(seen_execution_runtime_requests.len(), 2);
|
||||
// Default sticky_key_attempts is 2: the primary key is retried once on
|
||||
// the same key, then failover moves to the backup with a single attempt.
|
||||
assert_eq!(seen_execution_runtime_requests.len(), 3);
|
||||
for primary_request in &seen_execution_runtime_requests[..2] {
|
||||
assert_eq!(
|
||||
primary_request.trace_id,
|
||||
"trace-openai-chat-local-failover-123"
|
||||
);
|
||||
assert_eq!(
|
||||
primary_request.url,
|
||||
"https://api.openai.primary.example/chat/completions"
|
||||
);
|
||||
assert_eq!(
|
||||
primary_request.authorization,
|
||||
"Bearer sk-upstream-openai-primary"
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
seen_execution_runtime_requests[0].trace_id,
|
||||
"trace-openai-chat-local-failover-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_requests[0].url,
|
||||
"https://api.openai.primary.example/chat/completions"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_requests[0].authorization,
|
||||
"Bearer sk-upstream-openai-primary"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_requests[1].url,
|
||||
seen_execution_runtime_requests[2].url,
|
||||
"https://api.openai.backup.example/chat/completions"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_requests[1].model,
|
||||
seen_execution_runtime_requests[2].model,
|
||||
"gpt-5-upstream-backup"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_requests[1].authorization,
|
||||
seen_execution_runtime_requests[2].authorization,
|
||||
"Bearer sk-upstream-openai-backup"
|
||||
);
|
||||
let stored_candidates = request_candidate_repository
|
||||
.list_by_request_id("trace-openai-chat-local-failover-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 2);
|
||||
assert_eq!(stored_candidates[0].candidate_index, 0);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(stored_candidates[0].status_code, Some(401));
|
||||
assert_eq!(
|
||||
stored_candidates[0].error_message.as_deref(),
|
||||
Some("invalid auth token")
|
||||
);
|
||||
let failed_upstream_response = stored_candidates[0]
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("upstream_response"))
|
||||
.expect("failed candidate should keep its upstream response");
|
||||
assert_eq!(failed_upstream_response["status_code"], json!(401));
|
||||
assert_eq!(
|
||||
failed_upstream_response["body"]["error"]["message"],
|
||||
json!("invalid auth token")
|
||||
);
|
||||
assert_eq!(stored_candidates[1].candidate_index, 1);
|
||||
assert_eq!(stored_candidates[1].status, RequestCandidateStatus::Success);
|
||||
assert_eq!(stored_candidates[1].status_code, Some(200));
|
||||
assert_eq!(stored_candidates.len(), 3);
|
||||
for (retry_index, failed_candidate) in stored_candidates[..2].iter().enumerate() {
|
||||
assert_eq!(failed_candidate.candidate_index, 0);
|
||||
assert_eq!(failed_candidate.retry_index, retry_index as u32);
|
||||
assert_eq!(failed_candidate.status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(failed_candidate.status_code, Some(401));
|
||||
assert_eq!(
|
||||
failed_candidate.error_message.as_deref(),
|
||||
Some("invalid auth token")
|
||||
);
|
||||
let failed_upstream_response = failed_candidate
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("upstream_response"))
|
||||
.expect("failed candidate should keep its upstream response");
|
||||
assert_eq!(failed_upstream_response["status_code"], json!(401));
|
||||
assert_eq!(
|
||||
failed_upstream_response["body"]["error"]["message"],
|
||||
json!("invalid auth token")
|
||||
);
|
||||
}
|
||||
assert_eq!(stored_candidates[2].candidate_index, 1);
|
||||
assert_eq!(stored_candidates[2].status, RequestCandidateStatus::Success);
|
||||
assert_eq!(stored_candidates[2].status_code, Some(200));
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
@@ -1184,7 +1193,7 @@ async fn gateway_retries_next_local_openai_chat_sync_candidate_after_auth_failur
|
||||
|
||||
assert_eq!(
|
||||
*execution_runtime_hits.lock().expect("mutex should lock"),
|
||||
2
|
||||
3
|
||||
);
|
||||
assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
@@ -140,7 +140,7 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
Some(900.0),
|
||||
None,
|
||||
@@ -162,7 +162,7 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
|
||||
"https://chatgpt.com/backend-api/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
@@ -570,9 +570,12 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
|
||||
.filter(|plan| plan["request_id"] == "trace-search-failover-1")
|
||||
.map(|plan| plan["provider_id"].clone())
|
||||
.collect::<Vec<_>>();
|
||||
// Default sticky_key_attempts is 2: the first provider is retried once on
|
||||
// the same key before failover advances to the second provider.
|
||||
assert_eq!(
|
||||
failover_plans,
|
||||
vec![
|
||||
json!("provider-codex-search-1"),
|
||||
json!("provider-codex-search-1"),
|
||||
json!("provider-codex-search-2")
|
||||
]
|
||||
@@ -581,17 +584,16 @@ async fn gateway_executes_codex_search_with_responses_permission_and_search_cont
|
||||
.list_by_request_id("trace-search-failover-1")
|
||||
.await
|
||||
.expect("failover request candidates should read");
|
||||
assert_eq!(failover_candidates.len(), 2);
|
||||
assert_eq!(failover_candidates.len(), 3);
|
||||
for failed_candidate in &failover_candidates[..2] {
|
||||
assert_eq!(failed_candidate.status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(failed_candidate.status_code, Some(500));
|
||||
}
|
||||
assert_eq!(
|
||||
failover_candidates[0].status,
|
||||
RequestCandidateStatus::Failed
|
||||
);
|
||||
assert_eq!(failover_candidates[0].status_code, Some(500));
|
||||
assert_eq!(
|
||||
failover_candidates[1].status,
|
||||
failover_candidates[2].status,
|
||||
RequestCandidateStatus::Success
|
||||
);
|
||||
assert_eq!(failover_candidates[1].status_code, Some(200));
|
||||
assert_eq!(failover_candidates[2].status_code, Some(200));
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
|
||||
@@ -1714,7 +1714,6 @@ fn ai_serving_candidate_materialization_owns_affinity_and_candidate_runtime_pers
|
||||
"pub fn ai_should_persist_available_candidate_for_pool_key",
|
||||
"pub fn ai_should_persist_skipped_candidate_for_pool_membership",
|
||||
"pub fn ai_candidate_extra_data_with_ranking",
|
||||
"attempt_slot_count",
|
||||
"should_persist_available_candidate",
|
||||
"persist_available_candidate",
|
||||
"build_attempt",
|
||||
|
||||
@@ -122,7 +122,7 @@ pub(super) fn sample_local_openai_provider() -> StoredProviderCatalogProvider {
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
@@ -144,7 +144,7 @@ pub(super) fn sample_local_openai_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
"https://api.openai.example/v1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
|
||||
@@ -879,9 +879,13 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha
|
||||
.list_by_request_id("trace-openai-chat-local-report-sync-failure-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(stored_candidates[0].status_code, Some(503));
|
||||
// The only candidate is the sticky first key: the default policy retries
|
||||
// it once on the same key before the request is exhausted.
|
||||
assert_eq!(stored_candidates.len(), 2);
|
||||
for candidate in &stored_candidates {
|
||||
assert_eq!(candidate.status, RequestCandidateStatus::Failed);
|
||||
assert_eq!(candidate.status_code, Some(503));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -957,7 +961,8 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable
|
||||
let response = send_request(gateway, request).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(*execution_hits.lock().expect("mutex should lock"), 1);
|
||||
// The sticky first key is retried once on the same key before exhaustion.
|
||||
assert_eq!(*execution_hits.lock().expect("mutex should lock"), 2);
|
||||
|
||||
let stored_usage = wait_for_usage_status(
|
||||
usage_repository.as_ref(),
|
||||
@@ -982,15 +987,15 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable
|
||||
.list_by_request_id("trace-openai-chat-local-transport-unavailable-123")
|
||||
.await
|
||||
.expect("request candidate trace should read");
|
||||
assert_eq!(stored_candidates.len(), 1);
|
||||
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
|
||||
assert!(stored_candidates[0]
|
||||
.latency_ms
|
||||
.is_some_and(|value| value >= 5));
|
||||
assert_eq!(
|
||||
stored_candidates[0].error_type.as_deref(),
|
||||
Some("execution_runtime_unavailable")
|
||||
);
|
||||
assert_eq!(stored_candidates.len(), 2);
|
||||
for candidate in &stored_candidates {
|
||||
assert_eq!(candidate.status, RequestCandidateStatus::Failed);
|
||||
assert!(candidate.latency_ms.is_some_and(|value| value >= 5));
|
||||
assert_eq!(
|
||||
candidate.error_type.as_deref(),
|
||||
Some("execution_runtime_unavailable")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1786,7 +1791,7 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
@@ -1808,7 +1813,7 @@ async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_s
|
||||
"https://right.codes/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
Some("/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
@@ -2111,7 +2116,7 @@ fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
@@ -2133,7 +2138,7 @@ fn gateway_keeps_failed_usage_request_capture_lightweight_for_large_local_claude
|
||||
"https://right.codes/codex".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some(1),
|
||||
Some("/v1/messages".to_string()),
|
||||
None,
|
||||
None,
|
||||
|
||||
Reference in New Issue
Block a user