mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-09 12:40:20 +08:00
Merge codex/routing-strategy-consolidation into main
This commit is contained in:
@@ -9,7 +9,7 @@ use aether_ai_serving::{
|
||||
use aether_dispatch_core::{DispatchSequence, DispatchSequenceItem};
|
||||
use aether_routing_core::{
|
||||
rank_vector_for_candidate, CandidateKind, ResolvedRoutingPolicy, RoutingCandidateFacts,
|
||||
RoutingCandidateTrace, RoutingDecisionTrace,
|
||||
RoutingCandidateTrace, RoutingDecisionTrace, RoutingExecutionPolicy,
|
||||
};
|
||||
use aether_scheduler_core::{
|
||||
ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome,
|
||||
@@ -79,6 +79,13 @@ type DecorateSkippedCandidateFn<'a> = Arc<
|
||||
pub(crate) trait LocalExecutionAttemptSource<T>: Send {
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<T>, GatewayError>;
|
||||
|
||||
/// Returns the request-scoped execution behaviour selected by routing.
|
||||
/// Execution wrappers use this snapshot before consuming the first
|
||||
/// attempt, avoiding a second lookup against mutable system settings.
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn drain_execution_attempts(&mut self) -> Result<Vec<T>, GatewayError>;
|
||||
|
||||
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError>;
|
||||
@@ -1237,9 +1244,7 @@ async fn scheduler_cache_affinity_enabled(
|
||||
state: PlannerAppState<'_>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> bool {
|
||||
scheduler_ordering_config_for_routing_policy(state, routing_policy)
|
||||
.await
|
||||
.scheduling_mode
|
||||
scheduler_ordering_config_for_routing_policy(routing_policy).scheduling_mode
|
||||
== SchedulerSchedulingMode::CacheAffinity
|
||||
}
|
||||
|
||||
|
||||
@@ -6,14 +6,11 @@ use aether_ai_serving::{
|
||||
use aether_routing_core::ResolvedRoutingPolicy;
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::clock::current_unix_ms;
|
||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
||||
use crate::scheduler::config::{
|
||||
read_scheduler_ordering_config, SchedulerOrderingConfig, SchedulerSchedulingMode,
|
||||
};
|
||||
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
|
||||
use aether_scheduler_core::{
|
||||
matches_affinity_target, ClientSessionAffinity, SchedulerAffinityTarget,
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankableCandidate,
|
||||
@@ -133,7 +130,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> Vec<EligibleLocalExecutionCandidate> {
|
||||
let ordering_config = scheduler_ordering_config_for_routing_policy(state, routing_policy).await;
|
||||
let ordering_config = scheduler_ordering_config_for_routing_policy(routing_policy);
|
||||
let port = GatewayLocalCandidateRankingPort {
|
||||
state,
|
||||
requested_model,
|
||||
@@ -184,16 +181,24 @@ 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<'_>,
|
||||
/// Return the immutable scheduler snapshot carried by a resolved routing
|
||||
/// policy. A missing policy is a programming error in production request
|
||||
/// paths; unit tests may use the scheduler default for isolated ranking tests.
|
||||
pub(crate) fn scheduler_ordering_config_for_routing_policy(
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> SchedulerOrderingConfig {
|
||||
match routing_policy {
|
||||
Some(policy) => SchedulerOrderingConfig::from_routing_policy(policy),
|
||||
None => read_scheduler_ordering_config_or_default(state).await,
|
||||
None => {
|
||||
#[cfg(test)]
|
||||
{
|
||||
SchedulerOrderingConfig::default()
|
||||
}
|
||||
#[cfg(not(test))]
|
||||
{
|
||||
panic!("resolved routing policy is required before candidate scheduling")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -238,23 +243,6 @@ fn routing_overlaid_candidate(
|
||||
overlaid
|
||||
}
|
||||
|
||||
async fn read_scheduler_ordering_config_or_default(
|
||||
state: PlannerAppState<'_>,
|
||||
) -> SchedulerOrderingConfig {
|
||||
match read_scheduler_ordering_config(state.app()).await {
|
||||
Ok(config) => config,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "planner_scheduler_ordering_config_load_failed",
|
||||
log_type = "event",
|
||||
error = ?error,
|
||||
"failed to load scheduler ordering config while ranking local execution candidates"
|
||||
);
|
||||
SchedulerOrderingConfig::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
@@ -263,10 +251,16 @@ mod tests {
|
||||
use aether_ai_serving::{
|
||||
ai_ranking_context, build_ai_rankable_candidate, AiRankableCandidateParts,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::{
|
||||
provider_catalog::InMemoryProviderCatalogReadRepository,
|
||||
routing_profiles::InMemoryRoutingGroupRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupRecord, RoutingGroupWriteRepository,
|
||||
};
|
||||
use aether_scheduler_core::{
|
||||
apply_scheduler_candidate_ranking,
|
||||
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
|
||||
@@ -296,7 +290,11 @@ mod tests {
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
|
||||
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
|
||||
let ordering_config = super::read_scheduler_ordering_config_or_default(state).await;
|
||||
let ordering_config =
|
||||
crate::scheduler::config::read_system_default_routing_ordering_config(state.app())
|
||||
.await
|
||||
.expect("routing strategy should load")
|
||||
.unwrap_or_default();
|
||||
let mut candidates = candidates;
|
||||
let mut rankables = Vec::with_capacity(candidates.len());
|
||||
let mut ordering_cache = CandidateTransportRankingFactsCache::default();
|
||||
@@ -372,6 +370,7 @@ mod tests {
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
execution_policy: Default::default(),
|
||||
ranking_overlay: aether_routing_core::RankingOverlay::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: BTreeMap::new(),
|
||||
@@ -408,17 +407,14 @@ mod tests {
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
execution_policy: Default::default(),
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
matched_rules: Vec::new(),
|
||||
};
|
||||
|
||||
let ordering = super::scheduler_ordering_config_for_routing_policy(
|
||||
PlannerAppState::new(&state),
|
||||
Some(&policy),
|
||||
)
|
||||
.await;
|
||||
let ordering = super::scheduler_ordering_config_for_routing_policy(Some(&policy));
|
||||
|
||||
assert_eq!(
|
||||
ordering.scheduling_mode,
|
||||
@@ -446,6 +442,7 @@ mod tests {
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
execution_policy: Default::default(),
|
||||
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)]),
|
||||
@@ -917,7 +914,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_execution_ranking_keeps_cross_format_priority_when_global_override_is_enabled() {
|
||||
async fn local_execution_ranking_keeps_cross_format_priority_when_strategy_override_is_enabled()
|
||||
{
|
||||
let provider_catalog = InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider_with_options("provider-same", false, 10),
|
||||
@@ -932,14 +930,32 @@ mod tests {
|
||||
sample_key_for_provider("provider-cross", "key-cross", ""),
|
||||
],
|
||||
);
|
||||
let routing_repository = std::sync::Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
routing_repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "strategy-default".to_string(),
|
||||
name: "strategy-default".to_string(),
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
sort_order: 0,
|
||||
config_json: json!({
|
||||
"default_policy": {
|
||||
"keep_priority_on_conversion": true
|
||||
}
|
||||
}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.expect("routing strategy should be created");
|
||||
let data_state = GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
std::sync::Arc::new(provider_catalog),
|
||||
"development-key",
|
||||
)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"keep_priority_on_conversion".to_string(),
|
||||
json!(true),
|
||||
)]);
|
||||
.with_routing_group_repository_for_tests(routing_repository);
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
|
||||
@@ -384,8 +384,7 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
|
||||
None
|
||||
} else {
|
||||
Some(
|
||||
scheduler_ordering_config_for_routing_policy(state, routing_policy)
|
||||
.await
|
||||
scheduler_ordering_config_for_routing_policy(routing_policy)
|
||||
.sticky_key_attempts,
|
||||
)
|
||||
};
|
||||
|
||||
@@ -174,8 +174,9 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
|
||||
self.ranking_seed,
|
||||
false,
|
||||
self.request_operation,
|
||||
self.routing_policy
|
||||
.map(SchedulerOrderingConfig::from_routing_policy),
|
||||
super::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
self.routing_policy,
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -427,11 +428,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
|
||||
);
|
||||
|
||||
let ordering_config =
|
||||
super::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
state,
|
||||
routing_policy,
|
||||
)
|
||||
.await;
|
||||
super::candidate_ranking::scheduler_ordering_config_for_routing_policy(routing_policy);
|
||||
|
||||
Self {
|
||||
state,
|
||||
@@ -1293,9 +1290,7 @@ 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),
|
||||
self.ordering_config,
|
||||
)
|
||||
.await?;
|
||||
let skipped_candidates = skipped_candidates
|
||||
@@ -1890,6 +1885,7 @@ mod tests {
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
execution_policy: Default::default(),
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
@@ -1954,6 +1950,7 @@ mod tests {
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
execution_policy: Default::default(),
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
@@ -2692,6 +2689,7 @@ mod tests {
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::FixedOrder,
|
||||
keep_priority_on_conversion: true,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
execution_policy: Default::default(),
|
||||
ranking_overlay: Default::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: Default::default(),
|
||||
|
||||
@@ -625,21 +625,17 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
|
||||
GatewayRoutingSelectionError::NotFound(explicit_group.unwrap_or_default()),
|
||||
));
|
||||
}
|
||||
None
|
||||
return Err(routing_selection_error(
|
||||
GatewayRoutingSelectionError::NoDefault,
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let Some((group_id, group_version, group_config_json, selection_source)) = selected_group
|
||||
else {
|
||||
input.client_session_affinity = client_session_affinity_from_api_request(
|
||||
client_api_format,
|
||||
&parts.headers,
|
||||
Some(body_json),
|
||||
);
|
||||
input.routing_policy = None;
|
||||
input.routing_trace_seed = None;
|
||||
input.routing_context = None;
|
||||
return Ok(());
|
||||
return Err(routing_selection_error(
|
||||
GatewayRoutingSelectionError::NoDefault,
|
||||
));
|
||||
};
|
||||
|
||||
if try_attach_static_default_routing_policy_to_input(
|
||||
@@ -863,6 +859,10 @@ fn routing_selection_error(error: GatewayRoutingSelectionError) -> GatewayError
|
||||
GatewayRoutingSelectionError::Repository(message) => {
|
||||
GatewayError::Internal(format!("routing group repository lookup failed: {message}"))
|
||||
}
|
||||
GatewayRoutingSelectionError::NoDefault => GatewayError::Client {
|
||||
status: StatusCode::SERVICE_UNAVAILABLE,
|
||||
message: "no enabled routing strategy is configured for this request".to_string(),
|
||||
},
|
||||
error => GatewayError::Client {
|
||||
status: StatusCode::FORBIDDEN,
|
||||
message: error.to_string(),
|
||||
@@ -1207,6 +1207,7 @@ mod tests {
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: false,
|
||||
sort_order: 0,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
|
||||
+6
-9
@@ -26,7 +26,6 @@ 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::{
|
||||
@@ -141,10 +140,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
@@ -251,10 +249,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
|
||||
|
||||
@@ -26,6 +26,7 @@ use super::{
|
||||
LocalSameFormatProviderCandidateAttemptSource, LocalSameFormatProviderDecisionInput,
|
||||
LocalSameFormatProviderSpec,
|
||||
};
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
|
||||
pub(crate) struct LocalSameFormatProviderSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
@@ -189,6 +190,13 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalSameFormatProviderSyncAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_sync_attempt(attempt).await? {
|
||||
@@ -234,6 +242,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalSameFormatProviderSyncA
|
||||
impl LocalExecutionAttemptSource<AiStreamAttempt>
|
||||
for LocalSameFormatProviderStreamAttemptSource<'_>
|
||||
{
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_stream_attempt(attempt).await? {
|
||||
|
||||
@@ -21,7 +21,8 @@ use crate::client_session_affinity::{
|
||||
};
|
||||
use crate::orchestration::{
|
||||
insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity,
|
||||
ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
|
||||
ROUTING_EXECUTION_POLICY_REPORT_FIELD, ROUTING_POOL_POLICY_OVERRIDE_REPORT_FIELD,
|
||||
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
|
||||
};
|
||||
use crate::scheduler::affinity::insert_scheduler_affinity_policy_report_context_field;
|
||||
|
||||
@@ -112,6 +113,11 @@ pub(crate) fn build_local_execution_report_context(
|
||||
}
|
||||
insert_pool_key_lease_report_context_fields(&mut extra_fields, parts.pool_key_lease);
|
||||
insert_scheduler_affinity_policy_report_context_field(&mut extra_fields, parts.routing_policy);
|
||||
if let Some(policy) = parts.routing_policy {
|
||||
if let Ok(value) = serde_json::to_value(policy.execution_policy) {
|
||||
extra_fields.insert(ROUTING_EXECUTION_POLICY_REPORT_FIELD.to_string(), value);
|
||||
}
|
||||
}
|
||||
if let Some(override_policy) = parts
|
||||
.routing_policy
|
||||
.and_then(|policy| policy.pool_policy_overrides.get(parts.provider_id))
|
||||
|
||||
@@ -17,6 +17,7 @@ use crate::ai_serving::{
|
||||
resolve_gemini_files_sync_spec as resolve_sync_spec, LocalGeminiFilesSpec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
|
||||
use self::decision::maybe_build_local_gemini_files_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
@@ -174,6 +175,13 @@ pub(crate) async fn build_local_gemini_files_stream_attempt_source_for_kind<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalGeminiFilesSyncAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_sync_attempt(attempt).await? {
|
||||
@@ -212,6 +220,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalGeminiFilesSyncAttemptS
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalGeminiFilesStreamAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_stream_attempt(attempt).await? {
|
||||
|
||||
@@ -26,7 +26,6 @@ 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;
|
||||
@@ -109,10 +108,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
let outcome = materialize_local_execution_candidates_with_serving(
|
||||
@@ -186,10 +184,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
Ok(build_local_execution_candidate_attempt_source_with_serving(
|
||||
|
||||
@@ -19,6 +19,7 @@ use crate::ai_serving::{
|
||||
resolve_local_image_sync_spec as resolve_sync_spec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
|
||||
use self::decision::maybe_build_local_openai_image_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
@@ -252,6 +253,13 @@ pub(crate) async fn build_local_image_stream_attempt_source_for_kind<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiImageSyncAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_sync_attempt(attempt).await? {
|
||||
@@ -290,6 +298,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiImageSyncAttemptS
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiImageStreamAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_stream_attempt(attempt).await? {
|
||||
|
||||
@@ -27,7 +27,6 @@ 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;
|
||||
|
||||
@@ -128,10 +127,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -206,10 +204,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -16,6 +16,7 @@ use crate::ai_serving::{
|
||||
LocalVideoCreateSpec,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
|
||||
use self::decision::maybe_build_local_video_create_decision_payload_for_candidate;
|
||||
use self::support::{
|
||||
@@ -104,6 +105,13 @@ pub(crate) async fn build_local_video_sync_attempt_source_for_kind<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalVideoCreateSyncAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_sync_attempt(attempt).await? {
|
||||
|
||||
@@ -29,7 +29,6 @@ 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;
|
||||
@@ -134,10 +133,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -195,10 +193,9 @@ 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),
|
||||
crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy(
|
||||
input.routing_policy.as_ref(),
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -18,6 +18,7 @@ use crate::ai_serving::planner::spec_metadata::{
|
||||
};
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
|
||||
use super::candidates::{
|
||||
build_local_standard_candidate_attempt_source, resolve_local_standard_decision_input,
|
||||
@@ -177,6 +178,13 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalStandardSyncAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_sync_attempt(attempt).await? {
|
||||
@@ -220,6 +228,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalStandardSyncAttemptSour
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalStandardStreamAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_stream_attempt(attempt).await? {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
use async_trait::async_trait;
|
||||
use std::collections::VecDeque;
|
||||
use tracing::warn;
|
||||
@@ -119,6 +120,13 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiChatStreamAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
let select_started_at = std::time::Instant::now();
|
||||
let selected = self.next_execution_attempt_with_target_select().await?;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
use async_trait::async_trait;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -92,6 +93,13 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiChatSyncAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_sync_attempt(attempt).await? {
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
use async_trait::async_trait;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -161,6 +162,13 @@ pub(super) async fn build_local_stream_attempt_source<'a>(
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_sync_attempt(attempt).await? {
|
||||
@@ -204,6 +212,13 @@ impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalOpenAiResponsesSyncAtte
|
||||
|
||||
#[async_trait]
|
||||
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiResponsesStreamAttemptSource<'_> {
|
||||
fn routing_execution_policy(&self) -> Option<RoutingExecutionPolicy> {
|
||||
self.input
|
||||
.routing_policy
|
||||
.as_ref()
|
||||
.map(|policy| policy.execution_policy)
|
||||
}
|
||||
|
||||
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
|
||||
while let Some(attempt) = self.candidates.next_attempt().await? {
|
||||
match self.build_stream_attempt(attempt).await? {
|
||||
|
||||
@@ -12,9 +12,8 @@ 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.
|
||||
/// `ordering_config` is the immutable scheduler snapshot derived from the
|
||||
/// request's resolved routing policy.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates(
|
||||
self,
|
||||
@@ -26,7 +25,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
crate::scheduler::candidate::list_selectable_candidates(
|
||||
self.app().data.as_ref(),
|
||||
@@ -55,7 +54,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -90,7 +89,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
request_operation: Option<&str>,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -146,7 +145,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -177,7 +176,7 @@ impl<'a> PlannerAppState<'a> {
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
ordering_config: 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));
|
||||
|
||||
@@ -5077,6 +5077,7 @@ mod tests {
|
||||
scheduling_mode: RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
sticky_key_attempts: aether_routing_core::DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
execution_policy: Default::default(),
|
||||
ranking_overlay: RankingOverlay {
|
||||
allowed_keys: key_ids.into_iter().map(str::to_string).collect(),
|
||||
..RankingOverlay::default()
|
||||
|
||||
@@ -118,8 +118,7 @@ use crate::execution_runtime::{
|
||||
use crate::log_ids::short_request_id;
|
||||
use crate::orchestration::{
|
||||
apply_local_execution_effect, build_local_error_flow_metadata, classify_failure_disposition,
|
||||
cyber_continue_failover_enabled, spawn_local_oauth_success_effect,
|
||||
trace_upstream_response_body, with_error_flow_report_context,
|
||||
spawn_local_oauth_success_effect, trace_upstream_response_body, with_error_flow_report_context,
|
||||
with_upstream_response_report_context, FailureDisposition, FailureTokenAction,
|
||||
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
|
||||
LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis,
|
||||
@@ -6173,7 +6172,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
|
||||
}
|
||||
let prefetch_for_cyber_failover =
|
||||
is_openai_responses_family_format(plan.provider_api_format.as_str())
|
||||
&& cyber_continue_failover_enabled(state).await;
|
||||
&& crate::orchestration::routing_execution_policy_from_report_context(
|
||||
report_context.as_ref(),
|
||||
)
|
||||
.is_some_and(|policy| policy.cyber_continue_failover);
|
||||
let stream_commit_policy = StreamCommitPolicy::for_response(
|
||||
direct_stream_finalize_kind.is_some(),
|
||||
upstream_content_type,
|
||||
@@ -8501,14 +8503,6 @@ mod tests {
|
||||
Arc::new(provider_catalog),
|
||||
"development-key",
|
||||
);
|
||||
let data_state = if continue_failover {
|
||||
data_state.with_system_config_values_for_tests([(
|
||||
crate::orchestration::CYBER_CONTINUE_FAILOVER_CONFIG_KEY.to_string(),
|
||||
json!(true),
|
||||
)])
|
||||
} else {
|
||||
data_state
|
||||
};
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
@@ -8557,7 +8551,10 @@ mod tests {
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"provider_api_format": "openai:responses",
|
||||
"client_api_format": "openai:responses"
|
||||
"client_api_format": "openai:responses",
|
||||
"routing_execution_policy": {
|
||||
"cyber_continue_failover": continue_failover
|
||||
}
|
||||
})),
|
||||
crate::clock::current_unix_ms(),
|
||||
Instant::now(),
|
||||
@@ -10515,7 +10512,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prefetched_codex_cyber_policy_violation_retries_when_system_setting_is_enabled() {
|
||||
async fn prefetched_codex_cyber_policy_violation_retries_when_routing_strategy_is_enabled() {
|
||||
assert!(
|
||||
execute_prefetched_codex_cyber_policy_failure(true)
|
||||
.await
|
||||
|
||||
@@ -54,13 +54,10 @@ use crate::executor::{
|
||||
record_failed_usage_for_exhausted_request, LocalExecutionExhaustion,
|
||||
LocalExecutionRequestOutcome,
|
||||
};
|
||||
use crate::handlers::shared::system_config_bool;
|
||||
use crate::request_diagnostics::{current_request_diagnostics, scope_request_diagnostics_with};
|
||||
use crate::stage_metrics::observe_gateway_stage_ms;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
const ENABLE_OPENAI_IMAGE_SYNC_HEARTBEAT_CONFIG_KEY: &str = "enable_openai_image_sync_heartbeat";
|
||||
const ENABLE_STANDARD_TEXT_SYNC_HEARTBEAT_CONFIG_KEY: &str = "enable_standard_text_sync_heartbeat";
|
||||
const OPENAI_IMAGE_SYNC_HEARTBEAT_INTERNAL_ERROR_STATUS: u16 = 502;
|
||||
const OPENAI_IMAGE_SYNC_HEARTBEAT_EXHAUSTED_STATUS: u16 = 503;
|
||||
const OPENAI_IMAGE_SYNC_HEARTBEAT_ERROR_MESSAGE_LIMIT: usize = 4096;
|
||||
@@ -107,7 +104,10 @@ pub(crate) async fn maybe_execute_sync_via_local_decision(
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await {
|
||||
if standard_text_sync_heartbeat_should_wrap(
|
||||
plan_kind,
|
||||
attempt_source.routing_execution_policy(),
|
||||
) {
|
||||
let parts_for_task = parts.clone();
|
||||
let body_json_for_task = body_json.clone();
|
||||
let transfer_tracker_for_task = transfer_tracker.clone();
|
||||
@@ -264,7 +264,10 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision(
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await {
|
||||
if standard_text_sync_heartbeat_should_wrap(
|
||||
plan_kind,
|
||||
attempt_source.routing_execution_policy(),
|
||||
) {
|
||||
let parts_for_task = parts.clone();
|
||||
let body_json_for_task = body_json.clone();
|
||||
let transfer_tracker_for_task = transfer_tracker.clone();
|
||||
@@ -381,7 +384,10 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision(
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await {
|
||||
if standard_text_sync_heartbeat_should_wrap(
|
||||
plan_kind,
|
||||
attempt_source.routing_execution_policy(),
|
||||
) {
|
||||
let parts_for_task = parts.clone();
|
||||
let body_json_for_task = body_json.clone();
|
||||
let transfer_tracker_for_task = transfer_tracker.clone();
|
||||
@@ -609,7 +615,10 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision(
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if standard_text_sync_heartbeat_should_wrap(state, plan_kind).await {
|
||||
if standard_text_sync_heartbeat_should_wrap(
|
||||
plan_kind,
|
||||
attempt_source.routing_execution_policy(),
|
||||
) {
|
||||
let parts_for_task = parts.clone();
|
||||
let body_json_for_task = body_json.clone();
|
||||
let transfer_tracker_for_task = transfer_tracker.clone();
|
||||
@@ -746,42 +755,6 @@ pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
|
||||
.await
|
||||
}
|
||||
|
||||
async fn openai_image_sync_heartbeat_enabled(state: &AppState) -> bool {
|
||||
match state
|
||||
.read_system_config_json_value(ENABLE_OPENAI_IMAGE_SYNC_HEARTBEAT_CONFIG_KEY)
|
||||
.await
|
||||
{
|
||||
Ok(value) => system_config_bool(value.as_ref(), false),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
event_name = "openai_image_sync_heartbeat_config_read_failed",
|
||||
log_type = "ops",
|
||||
error = ?err,
|
||||
"gateway failed to read sync image heartbeat config; defaulting disabled"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn standard_text_sync_heartbeat_enabled(state: &AppState) -> bool {
|
||||
match state
|
||||
.read_system_config_json_value(ENABLE_STANDARD_TEXT_SYNC_HEARTBEAT_CONFIG_KEY)
|
||||
.await
|
||||
{
|
||||
Ok(value) => system_config_bool(value.as_ref(), false),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
event_name = "standard_text_sync_heartbeat_config_read_failed",
|
||||
log_type = "ops",
|
||||
error = ?err,
|
||||
"gateway failed to read standard text sync heartbeat config; defaulting disabled"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn standard_text_sync_heartbeat_applies_to_plan_kind(plan_kind: &str) -> bool {
|
||||
matches!(
|
||||
plan_kind,
|
||||
@@ -795,9 +768,12 @@ fn standard_text_sync_heartbeat_applies_to_plan_kind(plan_kind: &str) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
async fn standard_text_sync_heartbeat_should_wrap(state: &AppState, plan_kind: &str) -> bool {
|
||||
fn standard_text_sync_heartbeat_should_wrap(
|
||||
plan_kind: &str,
|
||||
execution_policy: Option<aether_routing_core::RoutingExecutionPolicy>,
|
||||
) -> bool {
|
||||
standard_text_sync_heartbeat_applies_to_plan_kind(plan_kind)
|
||||
&& standard_text_sync_heartbeat_enabled(state).await
|
||||
&& execution_policy.is_some_and(|policy| policy.enable_cf_heartbeat)
|
||||
}
|
||||
|
||||
fn standard_text_sync_heartbeat_client_api_format_for_plan_kind(plan_kind: &str) -> &'static str {
|
||||
@@ -1329,7 +1305,10 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
|
||||
return Ok(LocalExecutionRequestOutcome::NoPath);
|
||||
};
|
||||
|
||||
if openai_image_sync_heartbeat_enabled(state).await {
|
||||
if attempt_source
|
||||
.routing_execution_policy()
|
||||
.is_some_and(|policy| policy.enable_cf_heartbeat)
|
||||
{
|
||||
let mut attempts = Vec::new();
|
||||
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
|
||||
attempts.push(attempt);
|
||||
@@ -1867,11 +1846,10 @@ mod tests {
|
||||
assert_eq!(body, json!({"data": [{"b64_json": "x"}]}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_image_sync_heartbeat_missing_config_defaults_disabled() {
|
||||
let state = AppState::new().expect("state should build");
|
||||
|
||||
assert!(!openai_image_sync_heartbeat_enabled(&state).await);
|
||||
#[test]
|
||||
fn openai_image_sync_heartbeat_missing_routing_policy_defaults_disabled() {
|
||||
assert!(!Option::<aether_routing_core::RoutingExecutionPolicy>::None
|
||||
.is_some_and(|policy| policy.enable_cf_heartbeat));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -2149,23 +2127,9 @@ mod tests {
|
||||
assert_eq!(body, json!({"data": [{"b64_json": "fallback-provider"}]}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standard_text_sync_heartbeat_missing_config_defaults_disabled() {
|
||||
let state = AppState::new().expect("state should build");
|
||||
|
||||
assert!(!standard_text_sync_heartbeat_enabled(&state).await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standard_text_sync_heartbeat_no_local_candidates_preserves_no_path() {
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::disabled().with_system_config_values_for_tests([(
|
||||
ENABLE_STANDARD_TEXT_SYNC_HEARTBEAT_CONFIG_KEY.to_string(),
|
||||
json!(true),
|
||||
)]),
|
||||
);
|
||||
let state = AppState::new().expect("state should build");
|
||||
let (parts, _) = http::Request::builder()
|
||||
.method(http::Method::POST)
|
||||
.uri("/v1/responses")
|
||||
|
||||
@@ -67,11 +67,14 @@ pub(crate) async fn build_admin_global_model_routing_payload(
|
||||
.push(key);
|
||||
}
|
||||
|
||||
// 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
|
||||
.unwrap_or_default();
|
||||
// The admin view reports the system-default routing strategy.
|
||||
let ordering_config =
|
||||
match crate::scheduler::config::read_system_default_routing_ordering_config(state.app())
|
||||
.await
|
||||
{
|
||||
Ok(Some(config)) => config,
|
||||
Ok(None) | Err(_) => crate::scheduler::config::SchedulerOrderingConfig::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;
|
||||
|
||||
@@ -265,7 +265,9 @@ pub(super) async fn build_admin_monitoring_cache_snapshot(
|
||||
state: &AdminAppState<'_>,
|
||||
) -> Result<AdminMonitoringCacheSnapshot, GatewayError> {
|
||||
let ordering_config =
|
||||
crate::scheduler::config::read_scheduler_ordering_config(state.app()).await?;
|
||||
crate::scheduler::config::read_system_default_routing_ordering_config(state.app())
|
||||
.await?
|
||||
.unwrap_or_default();
|
||||
let scheduling_mode = ordering_config.scheduling_mode_str().to_string();
|
||||
let provider_priority_mode = ordering_config.priority_mode_str().to_string();
|
||||
|
||||
|
||||
@@ -44,6 +44,8 @@ struct AdminRoutingGroupCreateRequest {
|
||||
#[serde(default)]
|
||||
is_system_default: bool,
|
||||
#[serde(default)]
|
||||
sort_order: i64,
|
||||
#[serde(default)]
|
||||
config_json: Option<Value>,
|
||||
}
|
||||
|
||||
@@ -137,6 +139,7 @@ async fn maybe_build_routing_groups_response(
|
||||
description: payload.description,
|
||||
enabled: payload.enabled,
|
||||
is_system_default: payload.is_system_default,
|
||||
sort_order: payload.sort_order,
|
||||
config_json,
|
||||
version: 1,
|
||||
created_at: now,
|
||||
@@ -464,6 +467,9 @@ fn build_routing_group_update_patch(
|
||||
if let Some(value) = object.get("is_system_default") {
|
||||
patch.is_system_default = Some(required_bool(value, "is_system_default")?);
|
||||
}
|
||||
if let Some(value) = object.get("sort_order") {
|
||||
patch.sort_order = Some(required_i64(value, "sort_order")?.max(0));
|
||||
}
|
||||
if let Some(value) = object.get("config_json") {
|
||||
validate_config_json(value)?;
|
||||
patch.config_json = Some(value.clone());
|
||||
@@ -604,6 +610,7 @@ fn routing_group_payload(group: &StoredRoutingGroup) -> Value {
|
||||
"description": group.description,
|
||||
"enabled": group.enabled,
|
||||
"is_system_default": group.is_system_default,
|
||||
"sort_order": group.sort_order,
|
||||
"config_json": group.config_json,
|
||||
"version": group.version,
|
||||
"created_at": group.created_at,
|
||||
|
||||
@@ -68,7 +68,6 @@ use crate::scheduler::candidate::{
|
||||
is_auth_api_key_concurrency_limit_skip_reason, AUTH_API_KEY_CONCURRENCY_LIMIT_SKIP_REASON,
|
||||
LEGACY_API_KEY_CONCURRENCY_LIMIT_SKIP_REASON,
|
||||
};
|
||||
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
|
||||
use crate::stage_metrics::observe_gateway_stage_ms;
|
||||
use crate::{
|
||||
AppState, FrontdoorUserRpmOutcome, GatewayError, GatewayFallbackMetricKind,
|
||||
@@ -411,27 +410,7 @@ async fn maybe_forward_public_request_to_tunnel_owner(
|
||||
policy_context,
|
||||
)
|
||||
} else {
|
||||
let cache_affinity_enabled = match read_scheduler_ordering_config(state).await {
|
||||
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %request_context.trace_id,
|
||||
error = ?err,
|
||||
"gateway failed to load scheduler config while checking tunnel affinity forwarding mode"
|
||||
);
|
||||
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
|
||||
}
|
||||
};
|
||||
if !cache_affinity_enabled {
|
||||
return Ok(None);
|
||||
}
|
||||
crate::scheduler::affinity::read_cached_scheduler_affinity_target(
|
||||
state,
|
||||
&auth_context.api_key_id,
|
||||
affinity_context.client_session_affinity.as_ref(),
|
||||
api_format,
|
||||
&affinity_context.requested_model,
|
||||
)
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(target) = target else {
|
||||
return Ok(None);
|
||||
|
||||
+362
-71
@@ -76,6 +76,29 @@ enum DatabaseDriverArg {
|
||||
Postgres,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
|
||||
enum DatabaseModeArg {
|
||||
Auto,
|
||||
VerifyOnly,
|
||||
}
|
||||
|
||||
fn resolve_database_mode(
|
||||
configured: Option<DatabaseModeArg>,
|
||||
legacy_auto_prepare: Option<bool>,
|
||||
) -> DatabaseModeArg {
|
||||
if let Some(configured) = configured {
|
||||
return configured;
|
||||
}
|
||||
if let Some(legacy_auto_prepare) = legacy_auto_prepare {
|
||||
return if legacy_auto_prepare {
|
||||
DatabaseModeArg::Auto
|
||||
} else {
|
||||
DatabaseModeArg::VerifyOnly
|
||||
};
|
||||
}
|
||||
DatabaseModeArg::Auto
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
|
||||
enum ExportDomainArg {
|
||||
Users,
|
||||
@@ -538,46 +561,71 @@ fn automatic_sql_pool_config_for_parallelism(
|
||||
|
||||
#[derive(ClapArgs, Debug, Clone)]
|
||||
struct GatewayDataArgs {
|
||||
#[arg(long, env = "AETHER_DATABASE_DRIVER")]
|
||||
#[arg(long, env = "AETHER_DATABASE_DRIVER", global = true)]
|
||||
database_driver: Option<DatabaseDriverArg>,
|
||||
|
||||
#[arg(long, env = "AETHER_DATABASE_URL")]
|
||||
#[arg(long, env = "AETHER_DATABASE_URL", global = true)]
|
||||
database_url: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_URL")]
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_URL", global = true)]
|
||||
postgres_url: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_ENCRYPTION_KEY")]
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_ENCRYPTION_KEY", global = true)]
|
||||
encryption_key: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_REDIS_URL")]
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_REDIS_URL", global = true)]
|
||||
redis_url: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_REDIS_KEY_PREFIX")]
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_REDIS_KEY_PREFIX", global = true)]
|
||||
redis_key_prefix: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS")]
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS",
|
||||
global = true
|
||||
)]
|
||||
postgres_min_connections: Option<u32>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS")]
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS",
|
||||
global = true
|
||||
)]
|
||||
postgres_max_connections: Option<u32>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_ACQUIRE_TIMEOUT_MS")]
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_ACQUIRE_TIMEOUT_MS",
|
||||
global = true
|
||||
)]
|
||||
postgres_acquire_timeout_ms: Option<u64>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_IDLE_TIMEOUT_MS")]
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_IDLE_TIMEOUT_MS",
|
||||
global = true
|
||||
)]
|
||||
postgres_idle_timeout_ms: Option<u64>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_MAX_LIFETIME_MS")]
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_MAX_LIFETIME_MS",
|
||||
global = true
|
||||
)]
|
||||
postgres_max_lifetime_ms: Option<u64>,
|
||||
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATA_POSTGRES_STATEMENT_CACHE_CAPACITY")]
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_STATEMENT_CACHE_CAPACITY",
|
||||
global = true
|
||||
)]
|
||||
postgres_statement_cache_capacity: Option<usize>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_DATA_POSTGRES_REQUIRE_SSL",
|
||||
default_value_t = false
|
||||
default_value_t = false,
|
||||
global = true
|
||||
)]
|
||||
postgres_require_ssl: bool,
|
||||
}
|
||||
@@ -1144,13 +1192,26 @@ enum DataCommand {
|
||||
Import(DataImportArgs),
|
||||
/// Copy persistent SQL data directly between two databases without a JSONL file.
|
||||
Copy(DataCopyArgs),
|
||||
/// Inspect or prepare the configured database.
|
||||
Db(DatabaseCommandArgs),
|
||||
}
|
||||
|
||||
#[derive(ClapArgs, Debug, Clone)]
|
||||
struct DatabaseCommandArgs {
|
||||
#[command(subcommand)]
|
||||
command: DatabaseCommand,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
enum DatabaseCommand {
|
||||
/// Show whether schema migrations and data backfills are current.
|
||||
Status,
|
||||
/// Apply pending schema migrations and data backfills.
|
||||
Prepare,
|
||||
}
|
||||
|
||||
#[derive(ClapArgs, Debug, Clone)]
|
||||
struct DataExportArgs {
|
||||
#[command(flatten)]
|
||||
data: GatewayDataArgs,
|
||||
|
||||
#[arg(long)]
|
||||
output: PathBuf,
|
||||
|
||||
@@ -1160,9 +1221,6 @@ struct DataExportArgs {
|
||||
|
||||
#[derive(ClapArgs, Debug, Clone)]
|
||||
struct DataImportArgs {
|
||||
#[command(flatten)]
|
||||
data: GatewayDataArgs,
|
||||
|
||||
#[arg(long)]
|
||||
input: PathBuf,
|
||||
}
|
||||
@@ -1284,18 +1342,25 @@ struct Args {
|
||||
)]
|
||||
node_role: NodeRoleArg,
|
||||
|
||||
#[arg(long, default_value_t = false)]
|
||||
#[arg(long, hide = true, default_value_t = false)]
|
||||
migrate: bool,
|
||||
|
||||
#[arg(long, default_value_t = false)]
|
||||
#[arg(long, hide = true, default_value_t = false)]
|
||||
apply_backfills: bool,
|
||||
|
||||
/// Database startup policy. Defaults to auto when neither this nor the legacy setting is set.
|
||||
#[arg(long, env = "AETHER_GATEWAY_DATABASE_MODE", value_enum)]
|
||||
database_mode: Option<DatabaseModeArg>,
|
||||
|
||||
/// Legacy compatibility switch. Prefer --database-mode.
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_GATEWAY_AUTO_PREPARE_DATABASE",
|
||||
default_value_t = false
|
||||
hide = true,
|
||||
num_args = 0..=1,
|
||||
default_missing_value = "true"
|
||||
)]
|
||||
auto_prepare_database: bool,
|
||||
auto_prepare_database: Option<bool>,
|
||||
|
||||
/// Path to frontend static files directory (SPA). When set, the gateway
|
||||
/// serves the frontend directly without nginx.
|
||||
@@ -1405,6 +1470,10 @@ struct Args {
|
||||
}
|
||||
|
||||
impl Args {
|
||||
fn effective_database_mode(&self) -> DatabaseModeArg {
|
||||
resolve_database_mode(self.database_mode, self.auto_prepare_database)
|
||||
}
|
||||
|
||||
fn effective_runtime_backend(
|
||||
&self,
|
||||
database: Option<&SqlDatabaseConfig>,
|
||||
@@ -1475,15 +1544,7 @@ impl Args {
|
||||
}
|
||||
|
||||
fn runtime_config(&self) -> Result<ServiceRuntimeConfig, std::io::Error> {
|
||||
let default_log_filter = if self.command.is_some()
|
||||
|| self.migrate
|
||||
|| self.apply_backfills
|
||||
|| self.auto_prepare_database
|
||||
{
|
||||
"aether_gateway=info,aether_data=info"
|
||||
} else {
|
||||
"aether_gateway=info"
|
||||
};
|
||||
let default_log_filter = "aether_gateway=info,aether_data=info";
|
||||
let config = self
|
||||
.logging
|
||||
.apply_to_runtime_config(ServiceRuntimeConfig::new(
|
||||
@@ -1786,7 +1847,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let args = Args::parse();
|
||||
if let Some(command) = args.command.as_ref() {
|
||||
init_service_runtime(args.runtime_config()?)?;
|
||||
return run_data_command(command).await;
|
||||
return run_data_command(command, &args.data).await;
|
||||
}
|
||||
if args.migrate {
|
||||
init_service_runtime(args.runtime_config()?)?;
|
||||
@@ -2086,7 +2147,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
execution_runtime_configured = state.execution_runtime_configured(),
|
||||
"aether-gateway data layer configured"
|
||||
);
|
||||
prepare_database_startup_requirements(&state, args.auto_prepare_database).await?;
|
||||
prepare_database_startup_requirements(&state, args.effective_database_mode()).await?;
|
||||
state.warm_database_pools().await?;
|
||||
let reset_stale_proxy_nodes = state.reset_stale_proxy_node_tunnel_statuses().await?;
|
||||
if reset_stale_proxy_nodes > 0 {
|
||||
@@ -2101,16 +2162,11 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
info!(
|
||||
group_id = %group.id,
|
||||
group_name = %group.name,
|
||||
"created system default routing group from legacy scheduler config"
|
||||
"created system default routing group from routing strategy defaults"
|
||||
);
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
error = %err,
|
||||
"failed to bootstrap system default routing group; scheduler falls back to legacy system config"
|
||||
);
|
||||
}
|
||||
Err(err) => return Err(err.into()),
|
||||
}
|
||||
match state.prewarm_chat_pii_redaction_runtime_config().await {
|
||||
Ok(enabled) => {
|
||||
@@ -2208,14 +2264,77 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_data_command(command: &DataCommand) -> Result<(), Box<dyn std::error::Error>> {
|
||||
async fn run_data_command(
|
||||
command: &DataCommand,
|
||||
data: &GatewayDataArgs,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
match command {
|
||||
DataCommand::Export(args) => run_data_export(args).await,
|
||||
DataCommand::Import(args) => run_data_import(args).await,
|
||||
DataCommand::Export(args) => run_data_export(args, data).await,
|
||||
DataCommand::Import(args) => run_data_import(args, data).await,
|
||||
DataCommand::Copy(args) => run_data_copy(args).await,
|
||||
DataCommand::Db(args) => run_database_command(args, data).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_database_command(
|
||||
args: &DatabaseCommandArgs,
|
||||
data: &GatewayDataArgs,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
match args.command {
|
||||
DatabaseCommand::Status => run_database_status(data).await,
|
||||
DatabaseCommand::Prepare => run_database_prepare(data).await,
|
||||
}
|
||||
}
|
||||
|
||||
fn database_maintenance_state(
|
||||
data: &GatewayDataArgs,
|
||||
) -> Result<(DatabaseDriver, AppState), Box<dyn std::error::Error>> {
|
||||
let database = required_sql_database_config(data)?;
|
||||
let driver = database.driver;
|
||||
let state = AppState::new()?.with_data_config(data.to_config())?;
|
||||
Ok((driver, state))
|
||||
}
|
||||
|
||||
async fn run_database_status(data: &GatewayDataArgs) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let (driver, state) = database_maintenance_state(data)?;
|
||||
let pending_migrations = state
|
||||
.pending_database_migrations()
|
||||
.await?
|
||||
.unwrap_or_default();
|
||||
|
||||
if let Some(next) = pending_migrations.first() {
|
||||
println!("database {driver}: preparation required");
|
||||
println!("pending migrations: {}", pending_migrations.len());
|
||||
println!("next migration: {} ({})", next.version, next.description);
|
||||
println!("pending backfills: not checked until migrations are current");
|
||||
println!("run `aether-gateway db prepare`");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let pending_backfills = state
|
||||
.pending_database_backfills()
|
||||
.await?
|
||||
.unwrap_or_default();
|
||||
if let Some(next) = pending_backfills.first() {
|
||||
println!("database {driver}: preparation required");
|
||||
println!("pending migrations: 0");
|
||||
println!("pending backfills: {}", pending_backfills.len());
|
||||
println!("next backfill: {} ({})", next.version, next.description);
|
||||
println!("run `aether-gateway db prepare`");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!("database {driver}: ready (schema and backfills are current)");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_database_prepare(data: &GatewayDataArgs) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let (driver, state) = database_maintenance_state(data)?;
|
||||
prepare_database_startup_requirements(&state, DatabaseModeArg::Auto).await?;
|
||||
println!("database {driver}: ready (schema and backfills are current)");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn required_sql_database_config(
|
||||
data: &GatewayDataArgs,
|
||||
) -> Result<SqlDatabaseConfig, Box<dyn std::error::Error>> {
|
||||
@@ -2242,8 +2361,11 @@ fn current_unix_secs() -> Result<u64, std::time::SystemTimeError> {
|
||||
.as_secs())
|
||||
}
|
||||
|
||||
async fn run_data_export(args: &DataExportArgs) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = required_sql_database_config(&args.data)?;
|
||||
async fn run_data_export(
|
||||
args: &DataExportArgs,
|
||||
data: &GatewayDataArgs,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = required_sql_database_config(data)?;
|
||||
let driver = database.driver;
|
||||
let domains = requested_export_domains(args);
|
||||
let created_at_unix_secs = current_unix_secs()?;
|
||||
@@ -2265,8 +2387,11 @@ async fn run_data_export(args: &DataExportArgs) -> Result<(), Box<dyn std::error
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_data_import(args: &DataImportArgs) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = required_sql_database_config(&args.data)?;
|
||||
async fn run_data_import(
|
||||
args: &DataImportArgs,
|
||||
data: &GatewayDataArgs,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = required_sql_database_config(data)?;
|
||||
let driver = database.driver;
|
||||
let input = tokio::fs::read_to_string(&args.input).await?;
|
||||
let imported = import_database_jsonl(database, &input).await?;
|
||||
@@ -2426,17 +2551,15 @@ async fn run_explicit_backfills(args: &Args) -> Result<(), Box<dyn std::error::E
|
||||
|
||||
async fn prepare_database_startup_requirements(
|
||||
state: &AppState,
|
||||
auto_prepare_database: bool,
|
||||
database_mode: DatabaseModeArg,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
if !auto_prepare_database {
|
||||
if matches!(database_mode, DatabaseModeArg::VerifyOnly) {
|
||||
ensure_database_schema_is_current(state).await?;
|
||||
ensure_database_backfills_are_current(state).await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
info!(
|
||||
"auto database preparation enabled; applying pending migrations and backfills before serving traffic"
|
||||
);
|
||||
info!("database preparation enabled; applying pending migrations and backfills");
|
||||
|
||||
let Some(pending_migrations) = state.prepare_database_for_startup().await? else {
|
||||
return Ok(());
|
||||
@@ -2450,10 +2573,10 @@ async fn prepare_database_startup_requirements(
|
||||
next_version = next.version,
|
||||
next_description = %next.description,
|
||||
pending_versions = %format_pending_migrations(&pending_migrations),
|
||||
"running database migrations during service startup..."
|
||||
"running database migrations during database preparation..."
|
||||
);
|
||||
if state.run_database_migrations().await? {
|
||||
info!("database migrations complete during service startup");
|
||||
info!("database migrations complete");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2472,10 +2595,10 @@ async fn prepare_database_startup_requirements(
|
||||
next_version = next.version,
|
||||
next_description = %next.description,
|
||||
pending_versions = %format_pending_backfills(&pending_backfills),
|
||||
"running database backfills during service startup..."
|
||||
"running database backfills during database preparation..."
|
||||
);
|
||||
if state.run_database_backfills().await? {
|
||||
info!("database backfills complete during service startup");
|
||||
info!("database backfills complete");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -2520,7 +2643,7 @@ async fn ensure_database_backfills_are_current(
|
||||
async fn ensure_database_schema_is_current(
|
||||
state: &AppState,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let Some(pending) = state.prepare_database_for_startup().await? else {
|
||||
let Some(pending) = state.pending_database_migrations().await? else {
|
||||
return Ok(());
|
||||
};
|
||||
if pending.is_empty() {
|
||||
@@ -2539,7 +2662,7 @@ fn pending_schema_error(
|
||||
next_description: &str,
|
||||
) -> std::io::Error {
|
||||
std::io::Error::other(format!(
|
||||
"database schema is behind by {} migration(s); next pending migration is {} ({})\nrun `aether-gateway --migrate` before starting the service",
|
||||
"database schema is behind by {} migration(s); next pending migration is {} ({})\nrun `aether-gateway db prepare` before starting the service",
|
||||
pending_count, next_version, next_description
|
||||
))
|
||||
}
|
||||
@@ -2550,7 +2673,7 @@ fn pending_backfills_error(
|
||||
next_description: &str,
|
||||
) -> std::io::Error {
|
||||
std::io::Error::other(format!(
|
||||
"database backfills are behind by {} backfill(s); next pending backfill is {} ({})\nrun `aether-gateway --apply-backfills` before starting the service",
|
||||
"database backfills are behind by {} backfill(s); next pending backfill is {} ({})\nrun `aether-gateway db prepare` before starting the service",
|
||||
pending_count, next_version, next_description
|
||||
))
|
||||
}
|
||||
@@ -2562,8 +2685,9 @@ mod tests {
|
||||
automatic_gateway_request_concurrency_for_parallelism, automatic_sql_pool_config,
|
||||
automatic_sql_pool_config_for_parallelism, automatic_usage_queue_workers_for_parallelism,
|
||||
ensure_database_backfills_are_current, ensure_database_schema_is_current,
|
||||
pending_backfills_error, pending_schema_error, resolve_healthcheck_url,
|
||||
usage_database_config_for_role, Args, DatabaseDriverArg, DeploymentTopologyArg,
|
||||
pending_backfills_error, pending_schema_error, resolve_database_mode,
|
||||
resolve_healthcheck_url, usage_database_config_for_role, Args, DataCommand,
|
||||
DatabaseCommand, DatabaseDriverArg, DatabaseModeArg, DeploymentTopologyArg,
|
||||
GatewayDataArgs, GatewayFrontdoorArgs, GatewayLogDestinationArg, GatewayLogFormatArg,
|
||||
GatewayLogRotationArg, GatewayLoggingArgs, GatewayRateLimitArgs, GatewayUsageArgs,
|
||||
NodeRoleArg, RuntimeBackendArg, VideoTaskTruthSourceArg,
|
||||
@@ -2574,6 +2698,7 @@ mod tests {
|
||||
};
|
||||
use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
|
||||
use aether_gateway::AppState;
|
||||
use clap::Parser;
|
||||
|
||||
fn test_args() -> Args {
|
||||
Args {
|
||||
@@ -2588,7 +2713,8 @@ mod tests {
|
||||
node_role: NodeRoleArg::All,
|
||||
migrate: false,
|
||||
apply_backfills: false,
|
||||
auto_prepare_database: false,
|
||||
database_mode: None,
|
||||
auto_prepare_database: None,
|
||||
static_dir: None,
|
||||
video_task_truth_source_mode: VideoTaskTruthSourceArg::PythonSyncReport,
|
||||
video_task_poller_interval_ms: 5_000,
|
||||
@@ -2690,6 +2816,19 @@ mod tests {
|
||||
.expect("test database config should build")
|
||||
}
|
||||
|
||||
fn temporary_sqlite_args(label: &str) -> (Args, std::path::PathBuf) {
|
||||
let mut args = test_args();
|
||||
let nonce = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("clock should be available")
|
||||
.as_nanos();
|
||||
let database_path =
|
||||
std::env::temp_dir().join(format!("aether-{label}-{}-{nonce}.db", std::process::id()));
|
||||
args.data.database_driver = Some(DatabaseDriverArg::Sqlite);
|
||||
args.data.database_url = Some(format!("sqlite://{}", database_path.display()));
|
||||
(args, database_path)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_healthcheck_url_from_app_port() {
|
||||
assert_eq!(
|
||||
@@ -2801,11 +2940,14 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_runtime_config_keeps_gateway_only_logs() {
|
||||
fn normal_runtime_config_includes_database_lifecycle_logs() {
|
||||
let config = test_args()
|
||||
.runtime_config()
|
||||
.expect("runtime config should build");
|
||||
assert_eq!(config.default_log_filter, "aether_gateway=info");
|
||||
assert_eq!(
|
||||
config.default_log_filter,
|
||||
"aether_gateway=info,aether_data=info"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2822,7 +2964,7 @@ mod tests {
|
||||
#[test]
|
||||
fn auto_prepare_database_runtime_config_enables_data_logs() {
|
||||
let mut args = test_args();
|
||||
args.auto_prepare_database = true;
|
||||
args.auto_prepare_database = Some(true);
|
||||
let config = args.runtime_config().expect("runtime config should build");
|
||||
assert_eq!(
|
||||
config.default_log_filter,
|
||||
@@ -2830,6 +2972,87 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn database_mode_defaults_to_auto_and_preserves_legacy_false() {
|
||||
assert_eq!(resolve_database_mode(None, None), DatabaseModeArg::Auto);
|
||||
assert_eq!(
|
||||
resolve_database_mode(None, Some(false)),
|
||||
DatabaseModeArg::VerifyOnly
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_database_mode(Some(DatabaseModeArg::Auto), Some(false)),
|
||||
DatabaseModeArg::Auto
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_database_commands_and_verify_only_mode() {
|
||||
let status = Args::try_parse_from(["aether-gateway", "db", "status"])
|
||||
.expect("db status should parse");
|
||||
assert!(matches!(
|
||||
status.command,
|
||||
Some(DataCommand::Db(args))
|
||||
if matches!(args.command, DatabaseCommand::Status)
|
||||
));
|
||||
|
||||
let verify_only =
|
||||
Args::try_parse_from(["aether-gateway", "--database-mode", "verify-only"])
|
||||
.expect("verify-only mode should parse");
|
||||
assert_eq!(
|
||||
verify_only.effective_database_mode(),
|
||||
DatabaseModeArg::VerifyOnly
|
||||
);
|
||||
|
||||
let legacy_false =
|
||||
Args::try_parse_from(["aether-gateway", "--auto-prepare-database=false"])
|
||||
.expect("legacy false setting should parse");
|
||||
assert_eq!(
|
||||
legacy_false.effective_database_mode(),
|
||||
DatabaseModeArg::VerifyOnly
|
||||
);
|
||||
|
||||
let prepare = Args::try_parse_from(["aether-gateway", "db", "prepare"])
|
||||
.expect("db prepare should parse");
|
||||
assert!(matches!(
|
||||
prepare.command,
|
||||
Some(DataCommand::Db(args))
|
||||
if matches!(args.command, DatabaseCommand::Prepare)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn database_arguments_are_global_for_database_commands() {
|
||||
let before = Args::try_parse_from([
|
||||
"aether-gateway",
|
||||
"--database-driver",
|
||||
"sqlite",
|
||||
"--database-url",
|
||||
"sqlite:///tmp/before.db",
|
||||
"db",
|
||||
"status",
|
||||
])
|
||||
.expect("database arguments before db should parse");
|
||||
assert_eq!(
|
||||
before.data.database_url.as_deref(),
|
||||
Some("sqlite:///tmp/before.db")
|
||||
);
|
||||
|
||||
let after = Args::try_parse_from([
|
||||
"aether-gateway",
|
||||
"db",
|
||||
"prepare",
|
||||
"--database-driver",
|
||||
"sqlite",
|
||||
"--database-url",
|
||||
"sqlite:///tmp/after.db",
|
||||
])
|
||||
.expect("database arguments after db prepare should parse");
|
||||
assert_eq!(
|
||||
after.data.database_url.as_deref(),
|
||||
Some("sqlite:///tmp/after.db")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_data_pool_auto_sizes_sqlite_to_single_connection() {
|
||||
let mut args = test_args();
|
||||
@@ -3500,17 +3723,17 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_schema_error_mentions_explicit_migrate_command() {
|
||||
fn pending_schema_error_mentions_database_prepare_command() {
|
||||
let error = pending_schema_error(2, 20260413020000, "squash usage schema split");
|
||||
let message = error.to_string();
|
||||
assert!(message.contains("database schema is behind by 2 migration(s)"));
|
||||
assert!(message.contains("20260413020000"));
|
||||
assert!(message.contains("squash usage schema split"));
|
||||
assert!(message.contains("aether-gateway --migrate"));
|
||||
assert!(message.contains("aether-gateway db prepare"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_backfills_error_mentions_explicit_apply_backfills_command() {
|
||||
fn pending_backfills_error_mentions_database_prepare_command() {
|
||||
let message = pending_backfills_error(
|
||||
1,
|
||||
20260422110000,
|
||||
@@ -3520,7 +3743,7 @@ mod tests {
|
||||
assert!(message.contains("database backfills are behind by 1 backfill(s)"));
|
||||
assert!(message.contains("20260422110000"));
|
||||
assert!(message.contains("backfill stats aggregate read path support"));
|
||||
assert!(message.contains("aether-gateway --apply-backfills"));
|
||||
assert!(message.contains("aether-gateway db prepare"));
|
||||
assert!(message.contains("before starting the service"));
|
||||
}
|
||||
|
||||
@@ -3543,11 +3766,79 @@ mod tests {
|
||||
#[tokio::test]
|
||||
async fn auto_prepare_database_is_noop_without_database_pool() {
|
||||
let state = AppState::new().expect("state should build");
|
||||
super::prepare_database_startup_requirements(&state, true)
|
||||
super::prepare_database_startup_requirements(&state, DatabaseModeArg::Auto)
|
||||
.await
|
||||
.expect("disabled data backend should not block startup");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn verify_only_does_not_prepare_fresh_sqlite_database() {
|
||||
let (args, database_path) = temporary_sqlite_args("verify-only");
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_config(args.data.to_config())
|
||||
.expect("sqlite state should build");
|
||||
let pending_before = state
|
||||
.pending_database_migrations()
|
||||
.await
|
||||
.expect("pending migrations should load")
|
||||
.expect("sqlite should expose migration state");
|
||||
assert!(!pending_before.is_empty());
|
||||
|
||||
let error =
|
||||
super::prepare_database_startup_requirements(&state, DatabaseModeArg::VerifyOnly)
|
||||
.await
|
||||
.expect_err("verify-only should reject a fresh database");
|
||||
assert!(error.to_string().contains("aether-gateway db prepare"));
|
||||
|
||||
let pending_after = state
|
||||
.pending_database_migrations()
|
||||
.await
|
||||
.expect("pending migrations should reload")
|
||||
.expect("sqlite should expose migration state");
|
||||
assert_eq!(pending_after, pending_before);
|
||||
drop(state);
|
||||
let _ = std::fs::remove_file(database_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn auto_mode_prepares_fresh_sqlite_database() {
|
||||
let (args, database_path) = temporary_sqlite_args("auto-prepare");
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_config(args.data.to_config())
|
||||
.expect("sqlite state should build");
|
||||
|
||||
super::prepare_database_startup_requirements(&state, DatabaseModeArg::Auto)
|
||||
.await
|
||||
.expect("auto mode should prepare a fresh database");
|
||||
assert!(state
|
||||
.pending_database_migrations()
|
||||
.await
|
||||
.expect("pending migrations should load")
|
||||
.expect("sqlite should expose migration state")
|
||||
.is_empty());
|
||||
assert!(state
|
||||
.pending_database_backfills()
|
||||
.await
|
||||
.expect("pending backfills should load")
|
||||
.expect("sqlite should expose backfill state")
|
||||
.is_empty());
|
||||
drop(state);
|
||||
let _ = std::fs::remove_file(database_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn database_prepare_requires_database_url() {
|
||||
let data = test_args().data;
|
||||
let error = super::run_database_prepare(&data)
|
||||
.await
|
||||
.expect_err("missing database URL should fail");
|
||||
assert!(error
|
||||
.to_string()
|
||||
.contains("AETHER_DATABASE_DRIVER/AETHER_DATABASE_URL"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_migrate_requires_database_url() {
|
||||
let args = test_args();
|
||||
|
||||
@@ -53,7 +53,6 @@ use crate::scheduler::affinity::{
|
||||
scheduler_affinity_policy_context_from_report_context, SCHEDULER_AFFINITY_POLICY_REPORT_FIELD,
|
||||
SCHEDULER_AFFINITY_TTL,
|
||||
};
|
||||
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
|
||||
use crate::AppState;
|
||||
|
||||
const POOL_SCORE_FEEDBACK_GATE_MAX_ENTRIES: usize = 50_000;
|
||||
@@ -763,36 +762,19 @@ async fn local_scheduler_affinity_matches_failed_target(
|
||||
local_execution_plan_uses_pool(state, plan).await
|
||||
}
|
||||
|
||||
async fn scheduler_cache_affinity_enabled(
|
||||
state: &AppState,
|
||||
report_context: Option<&Value>,
|
||||
) -> bool {
|
||||
if report_context
|
||||
fn scheduler_cache_affinity_enabled(report_context: Option<&Value>) -> bool {
|
||||
report_context
|
||||
.and_then(|context| context.get(SCHEDULER_AFFINITY_POLICY_REPORT_FIELD))
|
||||
.is_some()
|
||||
{
|
||||
return scheduler_affinity_policy_context_from_report_context(report_context)
|
||||
.is_some_and(|context| context.cache_affinity_enabled());
|
||||
}
|
||||
match read_scheduler_ordering_config(state).await {
|
||||
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "orchestration_scheduler_affinity_config_load_failed",
|
||||
log_type = "event",
|
||||
error = ?error,
|
||||
"failed to load scheduler config while checking cache affinity mode"
|
||||
);
|
||||
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
|
||||
}
|
||||
}
|
||||
&& scheduler_affinity_policy_context_from_report_context(report_context)
|
||||
.is_some_and(|context| context.cache_affinity_enabled())
|
||||
}
|
||||
|
||||
async fn remember_successful_local_scheduler_affinity(
|
||||
state: &AppState,
|
||||
context: LocalExecutionEffectContext<'_>,
|
||||
) {
|
||||
if !scheduler_cache_affinity_enabled(state, context.report_context).await {
|
||||
if !scheduler_cache_affinity_enabled(context.report_context) {
|
||||
return;
|
||||
}
|
||||
let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) else {
|
||||
|
||||
@@ -56,10 +56,11 @@ pub(crate) use self::oauth_error::{
|
||||
};
|
||||
pub(crate) use self::policy::{
|
||||
append_local_failover_policy_to_value, codex_cyber_flag_passthrough_enabled,
|
||||
cyber_continue_failover_enabled, local_failover_policy_from_report_context,
|
||||
local_failover_policy_from_transport, resolve_local_failover_policy,
|
||||
responses_websocket_adapter, LocalFailoverPolicy, LocalFailoverRegexRule,
|
||||
ResponsesWebSocketAdapter, CYBER_CONTINUE_FAILOVER_CONFIG_KEY, RESPONSES_WEBSOCKET_CONFIG_KEY,
|
||||
local_failover_policy_from_report_context, local_failover_policy_from_transport,
|
||||
resolve_local_failover_policy, responses_websocket_adapter,
|
||||
routing_execution_policy_from_report_context, LocalFailoverPolicy, LocalFailoverRegexRule,
|
||||
ResponsesWebSocketAdapter, RESPONSES_WEBSOCKET_CONFIG_KEY,
|
||||
ROUTING_EXECUTION_POLICY_REPORT_FIELD,
|
||||
};
|
||||
pub(crate) use self::recovery::{
|
||||
analyze_local_failover, analyze_local_transport_error, apply_provider_failure_disposition,
|
||||
|
||||
@@ -4,11 +4,13 @@ use aether_contracts::ExecutionPlan;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::debug;
|
||||
|
||||
use aether_routing_core::RoutingExecutionPolicy;
|
||||
|
||||
use crate::provider_transport::GatewayProviderTransportSnapshot;
|
||||
use crate::AppState;
|
||||
|
||||
pub(crate) const CYBER_CONTINUE_FAILOVER_CONFIG_KEY: &str = "cyber_continue_failover";
|
||||
pub(crate) const RESPONSES_WEBSOCKET_CONFIG_KEY: &str = "responses_websocket";
|
||||
pub(crate) const ROUTING_EXECUTION_POLICY_REPORT_FIELD: &str = "routing_execution_policy";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct LocalFailoverPolicy {
|
||||
@@ -50,7 +52,7 @@ pub(crate) struct LocalFailoverRegexRule {
|
||||
pub(crate) async fn resolve_local_failover_policy(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
_report_context: Option<&serde_json::Value>,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
) -> LocalFailoverPolicy {
|
||||
let mut policy = match state
|
||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
||||
@@ -59,7 +61,8 @@ pub(crate) async fn resolve_local_failover_policy(
|
||||
Ok(Some(transport)) => local_failover_policy_from_transport(&transport),
|
||||
Ok(None) | Err(_) => LocalFailoverPolicy::default(),
|
||||
};
|
||||
let cyber_continue_failover = cyber_continue_failover_enabled(state).await;
|
||||
let cyber_continue_failover = routing_execution_policy_from_report_context(report_context)
|
||||
.is_some_and(|policy| policy.cyber_continue_failover);
|
||||
policy.stop_cyber_policy_errors = !cyber_continue_failover;
|
||||
debug!(
|
||||
event_name = "local_failover_policy_loaded",
|
||||
@@ -83,15 +86,13 @@ pub(crate) async fn resolve_local_failover_policy(
|
||||
policy
|
||||
}
|
||||
|
||||
pub(crate) async fn cyber_continue_failover_enabled(state: &AppState) -> bool {
|
||||
state
|
||||
.read_system_config_json_value(CYBER_CONTINUE_FAILOVER_CONFIG_KEY)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.as_ref()
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
pub(crate) fn routing_execution_policy_from_report_context(
|
||||
report_context: Option<&Value>,
|
||||
) -> Option<RoutingExecutionPolicy> {
|
||||
report_context
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|object| object.get(ROUTING_EXECUTION_POLICY_REPORT_FIELD))
|
||||
.and_then(|value| serde_json::from_value(value.clone()).ok())
|
||||
}
|
||||
|
||||
pub(crate) fn local_failover_policy_from_transport(
|
||||
|
||||
@@ -95,6 +95,7 @@ pub(crate) fn resolve_gateway_static_default_routing_policy(
|
||||
scheduling_mode: default_policy.scheduling_mode,
|
||||
keep_priority_on_conversion: default_policy.keep_priority_on_conversion,
|
||||
sticky_key_attempts: default_policy.sticky_key_attempts,
|
||||
execution_policy: default_policy.execution_policy,
|
||||
ranking_overlay: RankingOverlay::default(),
|
||||
mutation_plan: MutationPlan::default(),
|
||||
pool_policy_overrides: BTreeMap::new(),
|
||||
@@ -108,8 +109,10 @@ fn static_default_policy_fields(
|
||||
let Some(object) = config_json.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !routing_array_field_is_missing_or_empty(object, "allowed_models")
|
||||
|| !routing_array_field_is_missing_or_empty(object, "model_policies")
|
||||
// A strategy's default policy applies to every model. Only model policies
|
||||
// and rules require the request-context-aware resolver; unknown legacy
|
||||
// fields (including the removed group allowlist) are intentionally ignored.
|
||||
if !routing_array_field_is_missing_or_empty(object, "model_policies")
|
||||
|| !routing_array_field_is_missing_or_empty(object, "rules")
|
||||
{
|
||||
return Ok(None);
|
||||
@@ -145,15 +148,47 @@ fn static_default_policy_fields(
|
||||
})?,
|
||||
None => DEFAULT_STICKY_KEY_ATTEMPTS,
|
||||
};
|
||||
let enable_cf_heartbeat = routing_bool_field(
|
||||
default_policy.get("enable_cf_heartbeat"),
|
||||
"enable_cf_heartbeat",
|
||||
)?;
|
||||
// Older strategies stored separate image/text heartbeat flags. Treat
|
||||
// either legacy flag as enabling the unified CF heartbeat setting while
|
||||
// allowing newly saved strategies to use only the canonical key.
|
||||
let legacy_image_heartbeat = routing_bool_field(
|
||||
default_policy.get("enable_openai_image_sync_heartbeat"),
|
||||
"enable_openai_image_sync_heartbeat",
|
||||
)?;
|
||||
let legacy_text_heartbeat = routing_bool_field(
|
||||
default_policy.get("enable_standard_text_sync_heartbeat"),
|
||||
"enable_standard_text_sync_heartbeat",
|
||||
)?;
|
||||
let execution_policy = aether_routing_core::RoutingExecutionPolicy {
|
||||
enable_cf_heartbeat: enable_cf_heartbeat || legacy_image_heartbeat || legacy_text_heartbeat,
|
||||
cyber_continue_failover: routing_bool_field(
|
||||
default_policy.get("cyber_continue_failover"),
|
||||
"cyber_continue_failover",
|
||||
)?,
|
||||
};
|
||||
|
||||
Ok(Some(RoutingDefaultPolicy {
|
||||
priority_mode,
|
||||
scheduling_mode,
|
||||
keep_priority_on_conversion,
|
||||
sticky_key_attempts,
|
||||
execution_policy,
|
||||
}))
|
||||
}
|
||||
|
||||
fn routing_bool_field(value: Option<&Value>, field: &str) -> Result<bool, GatewayError> {
|
||||
match value {
|
||||
Some(value) => value
|
||||
.as_bool()
|
||||
.ok_or_else(|| invalid_routing_group_config(format!("{field} must be a boolean"))),
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn routing_array_field_is_missing_or_empty(
|
||||
object: &serde_json::Map<String, Value>,
|
||||
key: &str,
|
||||
@@ -201,7 +236,7 @@ mod tests {
|
||||
"scheduling_mode": "load_balance",
|
||||
"keep_priority_on_conversion": true
|
||||
},
|
||||
"allowed_models": [],
|
||||
"allowed_models": ["legacy-model"],
|
||||
"model_policies": [],
|
||||
"rules": []
|
||||
});
|
||||
|
||||
@@ -8,6 +8,8 @@ pub(crate) const ROUTING_GROUP_HEADER: &str = "x-aether-scheduler-group";
|
||||
|
||||
#[derive(Debug, Error, Clone, PartialEq, Eq)]
|
||||
pub(crate) enum GatewayRoutingSelectionError {
|
||||
#[error("no enabled routing strategy is configured for this request")]
|
||||
NoDefault,
|
||||
#[error("routing group was explicitly requested but was not found: {0}")]
|
||||
NotFound(String),
|
||||
#[error("routing group was explicitly requested but is not enabled: {0}")]
|
||||
@@ -239,6 +241,7 @@ mod tests {
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: false,
|
||||
sort_order: 0,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
@@ -287,6 +290,7 @@ mod tests {
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
sort_order: 0,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
@@ -322,6 +326,7 @@ mod tests {
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: false,
|
||||
sort_order: 0,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
@@ -444,6 +449,7 @@ mod tests {
|
||||
description: None,
|
||||
enabled: false,
|
||||
is_system_default: false,
|
||||
sort_order: 0,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
@@ -481,6 +487,7 @@ mod tests {
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: false,
|
||||
sort_order: 0,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
use self::selection::{
|
||||
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;
|
||||
@@ -56,8 +55,7 @@ enum RequiredCapabilityMatchMode {
|
||||
}
|
||||
|
||||
/// `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).
|
||||
/// config. Every production scheduling pass must provide this snapshot.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn list_selectable_candidates(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
@@ -70,7 +68,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>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
collect_selectable_candidates(
|
||||
selection_row_source,
|
||||
@@ -107,7 +105,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>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -145,7 +143,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>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -180,7 +178,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>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -188,8 +186,6 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
|
||||
),
|
||||
GatewayError,
|
||||
> {
|
||||
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,
|
||||
@@ -220,7 +216,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>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
Ok(
|
||||
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
|
||||
@@ -249,7 +245,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>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
|
||||
let normalized_api_format = normalize_api_format(candidate_api_format);
|
||||
if normalized_api_format.is_empty() {
|
||||
|
||||
@@ -47,7 +47,7 @@ pub(super) fn is_exact_all_skipped_by_auth_limit(
|
||||
.all(|candidate| is_auth_api_key_concurrency_limit_skip_reason(candidate.skip_reason))
|
||||
}
|
||||
|
||||
#[cfg_attr(not(test), allow(dead_code))]
|
||||
#[cfg(test)]
|
||||
pub(super) async fn select_minimal_candidate(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -59,20 +59,8 @@ pub(super) async fn select_minimal_candidate(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
let affinity_epoch = runtime_state.scheduler_affinity_epoch();
|
||||
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
|
||||
let affinity_cache_key = build_scheduler_affinity_cache_key(
|
||||
auth_snapshot,
|
||||
api_format,
|
||||
global_model_name,
|
||||
client_session_affinity,
|
||||
);
|
||||
let priority_affinity_key = scheduling_priority_affinity_key(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
ordering_config.scheduling_mode,
|
||||
);
|
||||
let candidates = enumerate_scheduler_candidates(
|
||||
selection_row_source,
|
||||
api_format,
|
||||
@@ -84,7 +72,7 @@ pub(super) async fn select_minimal_candidate(
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
let selected = collect_selectable_enumerated_candidates_with_skip_reasons(
|
||||
Ok(collect_selectable_enumerated_candidates_with_skip_reasons(
|
||||
runtime_state,
|
||||
api_format,
|
||||
global_model_name,
|
||||
@@ -94,25 +82,16 @@ pub(super) async fn select_minimal_candidate(
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
ordering_config,
|
||||
priority_affinity_key,
|
||||
scheduling_priority_affinity_key(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
ordering_config.scheduling_mode,
|
||||
),
|
||||
)
|
||||
.await?
|
||||
.0
|
||||
.into_iter()
|
||||
.next();
|
||||
if ordering_config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity
|
||||
&& has_explicit_session_affinity(client_session_affinity)
|
||||
{
|
||||
if let Some(candidate) = selected.as_ref() {
|
||||
remember_scheduler_affinity(
|
||||
affinity_cache_key.as_deref(),
|
||||
runtime_state,
|
||||
candidate,
|
||||
Some(affinity_epoch),
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(selected)
|
||||
.next())
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
@@ -127,7 +106,7 @@ pub(super) async fn collect_selectable_candidates(
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
Ok(
|
||||
collect_selectable_candidates_with_skip_reasons_and_ordering(
|
||||
@@ -149,10 +128,8 @@ pub(super) async fn collect_selectable_candidates(
|
||||
)
|
||||
}
|
||||
|
||||
/// 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)]
|
||||
#[cfg(test)]
|
||||
pub(super) async fn collect_selectable_candidates_with_skip_reasons(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &impl SchedulerRuntimeState,
|
||||
@@ -184,24 +161,11 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
|
||||
now_unix_secs,
|
||||
enable_model_directives,
|
||||
request_operation,
|
||||
None,
|
||||
SchedulerOrderingConfig::default(),
|
||||
)
|
||||
.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),
|
||||
@@ -215,7 +179,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons_and_ordering
|
||||
now_unix_secs: u64,
|
||||
enable_model_directives: bool,
|
||||
request_operation: Option<&str>,
|
||||
ordering_config: Option<SchedulerOrderingConfig>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
) -> Result<
|
||||
(
|
||||
Vec<SchedulerMinimalCandidateSelectionCandidate>,
|
||||
@@ -223,8 +187,6 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons_and_ordering
|
||||
),
|
||||
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,
|
||||
|
||||
@@ -23,6 +23,7 @@ use crate::data::candidate_selection::{
|
||||
read_requested_model_rows, MinimalCandidateSelectionRowSource,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::super::affinity::build_scheduler_affinity_cache_key;
|
||||
@@ -50,6 +51,7 @@ async fn select_candidate(
|
||||
client_session_affinity,
|
||||
now_unix_secs,
|
||||
false,
|
||||
SchedulerOrderingConfig::default(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -65,7 +65,7 @@ async fn compatible_required_capability_prefers_matching_keys_without_hard_filte
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
crate::scheduler::config::SchedulerOrderingConfig::default(),
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -121,7 +121,7 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
|
||||
None,
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
crate::scheduler::config::SchedulerOrderingConfig::default(),
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -198,7 +198,7 @@ async fn required_capability_without_model_uses_session_scoped_affinity() {
|
||||
Some(&auth_snapshot),
|
||||
Some(&client_session_affinity),
|
||||
100,
|
||||
None,
|
||||
crate::scheduler::config::SchedulerOrderingConfig::default(),
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
@@ -276,7 +276,7 @@ async fn required_capability_reports_auth_limit_signal_when_every_model_is_block
|
||||
Some(&auth_snapshot),
|
||||
None,
|
||||
100,
|
||||
None,
|
||||
crate::scheduler::config::SchedulerOrderingConfig::default(),
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed");
|
||||
|
||||
@@ -5,6 +5,7 @@ use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelect
|
||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
|
||||
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
|
||||
};
|
||||
@@ -13,6 +14,9 @@ use aether_data_contracts::repository::candidates::{
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupRecord, RoutingGroupWriteRepository,
|
||||
};
|
||||
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -20,6 +24,7 @@ use crate::cache::SchedulerAffinityTarget;
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::scheduler::config::SchedulerOrderingConfig;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
use super::super::affinity::build_scheduler_affinity_cache_key;
|
||||
@@ -31,6 +36,39 @@ use super::super::selection::{
|
||||
};
|
||||
use super::support::{sample_auth_snapshot, sample_key, sample_provider, sample_row};
|
||||
|
||||
async fn state_with_routing_default_policy(
|
||||
data_state: GatewayDataState,
|
||||
default_policy: serde_json::Value,
|
||||
) -> AppState {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "selection-test-default".to_string(),
|
||||
name: "selection-test-default".to_string(),
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
sort_order: 0,
|
||||
config_json: json!({"default_policy": default_policy}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.expect("routing strategy should be created");
|
||||
AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state.with_routing_group_repository_for_tests(repository))
|
||||
}
|
||||
|
||||
async fn ordering_config(state: &AppState) -> SchedulerOrderingConfig {
|
||||
crate::scheduler::config::read_system_default_routing_ordering_config(state)
|
||||
.await
|
||||
.expect("routing strategy should load")
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
async fn select_candidate(
|
||||
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
|
||||
runtime_state: &AppState,
|
||||
@@ -40,6 +78,7 @@ async fn select_candidate(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
let ordering_config = ordering_config(runtime_state).await;
|
||||
select_candidate_impl(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
@@ -51,6 +90,7 @@ async fn select_candidate(
|
||||
None,
|
||||
now_unix_secs,
|
||||
false,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -64,6 +104,7 @@ async fn collect_selectable_candidates(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
|
||||
let ordering_config = ordering_config(runtime_state).await;
|
||||
collect_selectable_candidates_impl(
|
||||
selection_row_source,
|
||||
runtime_state,
|
||||
@@ -75,7 +116,7 @@ async fn collect_selectable_candidates(
|
||||
None,
|
||||
now_unix_secs,
|
||||
false,
|
||||
None,
|
||||
ordering_config,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -289,15 +330,11 @@ async fn selects_by_provider_priority_when_priority_mode_is_provider() {
|
||||
global_key_first,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"provider_priority_mode".to_string(),
|
||||
json!("provider"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"priority_mode": "provider"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let selected = select_candidate(
|
||||
state.data.as_ref(),
|
||||
@@ -343,15 +380,11 @@ async fn selects_by_global_key_priority_when_priority_mode_is_global_key() {
|
||||
global_key_first,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"provider_priority_mode".to_string(),
|
||||
json!("global_key"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"priority_mode": "global_key"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let selected = select_candidate(
|
||||
state.data.as_ref(),
|
||||
@@ -415,6 +448,7 @@ async fn scheduler_selection_prefers_required_capability_matches_before_priority
|
||||
None,
|
||||
100,
|
||||
false,
|
||||
SchedulerOrderingConfig::default(),
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed")
|
||||
@@ -450,15 +484,11 @@ async fn fixed_order_ignores_cached_scheduler_affinity_promotion() {
|
||||
first, second,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"scheduling_mode".to_string(),
|
||||
json!("fixed_order"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"scheduling_mode": "fixed_order"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
|
||||
state.remember_scheduler_affinity_target(
|
||||
@@ -515,15 +545,11 @@ async fn fixed_order_disables_same_priority_affinity_hash_tiebreaker() {
|
||||
first, second,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"scheduling_mode".to_string(),
|
||||
json!("fixed_order"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"scheduling_mode": "fixed_order"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
|
||||
let selection = collect_selectable_candidates(
|
||||
@@ -569,15 +595,11 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
|
||||
first, second,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"scheduling_mode".to_string(),
|
||||
json!("cache_affinity"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"scheduling_mode": "cache_affinity"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
|
||||
let client_session_affinity = ClientSessionAffinity::from_session_key("session-1");
|
||||
@@ -610,6 +632,7 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
|
||||
Some(&client_session_affinity),
|
||||
100,
|
||||
false,
|
||||
ordering_config(&state).await,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed")
|
||||
@@ -645,15 +668,11 @@ async fn cache_affinity_ignores_cached_scheduler_affinity_without_client_session
|
||||
first, second,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"scheduling_mode".to_string(),
|
||||
json!("cache_affinity"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"scheduling_mode": "cache_affinity"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
|
||||
state.remember_scheduler_affinity_target(
|
||||
@@ -691,15 +710,11 @@ async fn load_balance_selection_does_not_remember_scheduler_affinity() {
|
||||
row,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"scheduling_mode".to_string(),
|
||||
json!("load_balance"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"scheduling_mode": "load_balance"}),
|
||||
)
|
||||
.await;
|
||||
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
|
||||
let client_session_affinity = ClientSessionAffinity::from_session_key("session-1");
|
||||
let cache_key = build_scheduler_affinity_cache_key(
|
||||
@@ -721,6 +736,7 @@ async fn load_balance_selection_does_not_remember_scheduler_affinity() {
|
||||
Some(&client_session_affinity),
|
||||
100,
|
||||
false,
|
||||
ordering_config(&state).await,
|
||||
)
|
||||
.await
|
||||
.expect("selection should succeed")
|
||||
@@ -758,15 +774,11 @@ async fn load_balance_ignores_provider_priority_and_cached_affinity() {
|
||||
first, second,
|
||||
]));
|
||||
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas)
|
||||
.with_system_config_values_for_tests(vec![(
|
||||
"scheduling_mode".to_string(),
|
||||
json!("load_balance"),
|
||||
)]),
|
||||
);
|
||||
let state = state_with_routing_default_policy(
|
||||
GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas),
|
||||
json!({"scheduling_mode": "load_balance"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
|
||||
state.remember_scheduler_affinity_target(
|
||||
|
||||
@@ -55,7 +55,7 @@ impl Default for SchedulerOrderingConfig {
|
||||
|
||||
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.
|
||||
/// the single source of truth for request scheduling.
|
||||
pub(crate) fn from_routing_policy(policy: &ResolvedRoutingPolicy) -> Self {
|
||||
Self {
|
||||
priority_mode: scheduler_priority_mode_from_routing(policy.priority_mode),
|
||||
@@ -87,6 +87,7 @@ impl SchedulerOrderingConfig {
|
||||
},
|
||||
keep_priority_on_conversion: self.keep_priority_on_conversion,
|
||||
sticky_key_attempts: self.sticky_key_attempts,
|
||||
execution_policy: aether_routing_core::RoutingExecutionPolicy::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -114,61 +115,6 @@ fn scheduler_scheduling_mode_from_routing(mode: RoutingSchedulingMode) -> Schedu
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_scheduler_priority_mode(
|
||||
value: Option<&serde_json::Value>,
|
||||
) -> SchedulerPriorityMode {
|
||||
match value
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.as_deref()
|
||||
{
|
||||
Some("global_key") => SchedulerPriorityMode::GlobalKey,
|
||||
_ => SchedulerPriorityMode::Provider,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_keep_priority_on_conversion(value: Option<&serde_json::Value>) -> bool {
|
||||
value.and_then(serde_json::Value::as_bool).unwrap_or(false)
|
||||
}
|
||||
|
||||
pub(crate) fn parse_scheduler_scheduling_mode(
|
||||
value: Option<&serde_json::Value>,
|
||||
) -> SchedulerSchedulingMode {
|
||||
match value
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.as_deref()
|
||||
{
|
||||
Some("fixed_order") => SchedulerSchedulingMode::FixedOrder,
|
||||
Some("load_balance") => SchedulerSchedulingMode::LoadBalance,
|
||||
_ => SchedulerSchedulingMode::CacheAffinity,
|
||||
}
|
||||
}
|
||||
|
||||
/// 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,
|
||||
@@ -201,38 +147,6 @@ pub(crate) async fn read_system_default_routing_ordering_config(
|
||||
)))
|
||||
}
|
||||
|
||||
/// 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
|
||||
.read_system_config_json_value("provider_priority_mode")
|
||||
.await?
|
||||
.as_ref(),
|
||||
);
|
||||
let scheduling_mode = parse_scheduler_scheduling_mode(
|
||||
state
|
||||
.read_system_config_json_value("scheduling_mode")
|
||||
.await?
|
||||
.as_ref(),
|
||||
);
|
||||
let keep_priority_on_conversion = parse_keep_priority_on_conversion(
|
||||
state
|
||||
.read_system_config_json_value("keep_priority_on_conversion")
|
||||
.await?
|
||||
.as_ref(),
|
||||
);
|
||||
Ok(SchedulerOrderingConfig {
|
||||
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;
|
||||
@@ -247,14 +161,6 @@ mod tests {
|
||||
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,
|
||||
@@ -267,6 +173,7 @@ mod tests {
|
||||
description: None,
|
||||
enabled,
|
||||
is_system_default: true,
|
||||
sort_order: 0,
|
||||
config_json,
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
@@ -278,7 +185,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn system_default_routing_group_overrides_legacy_keys() {
|
||||
async fn system_default_routing_group_exposes_strategy_ordering() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(
|
||||
&repository,
|
||||
@@ -293,12 +200,13 @@ mod tests {
|
||||
)
|
||||
.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),
|
||||
GatewayDataState::disabled().with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
let config = read_system_default_routing_ordering_config(&state)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(config.priority_mode, SchedulerPriorityMode::Provider);
|
||||
assert_eq!(config.scheduling_mode, SchedulerSchedulingMode::FixedOrder);
|
||||
@@ -310,18 +218,19 @@ mod tests {
|
||||
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),
|
||||
GatewayDataState::disabled().with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
|
||||
let config = read_scheduler_ordering_config(&state).await.unwrap();
|
||||
let config = read_system_default_routing_ordering_config(&state)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(config, SchedulerOrderingConfig::default());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_or_missing_system_default_group_falls_back_to_legacy_keys() {
|
||||
async fn disabled_or_missing_system_default_group_uses_routing_defaults() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
create_system_default(
|
||||
&repository,
|
||||
@@ -330,28 +239,25 @@ mod tests {
|
||||
)
|
||||
.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()),
|
||||
GatewayDataState::disabled().with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
let without_repository = AppState::new()
|
||||
.unwrap()
|
||||
.with_data_state_for_tests(GatewayDataState::disabled());
|
||||
|
||||
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);
|
||||
let config = read_system_default_routing_ordering_config(&state)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(config.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bootstrap_creates_system_default_group_from_legacy_keys_once() {
|
||||
async fn bootstrap_creates_system_default_group_from_routing_defaults_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()),
|
||||
);
|
||||
|
||||
@@ -365,9 +271,9 @@ mod tests {
|
||||
assert_eq!(
|
||||
created.config_json["default_policy"],
|
||||
json!({
|
||||
"priority_mode": "global_key",
|
||||
"scheduling_mode": "load_balance",
|
||||
"keep_priority_on_conversion": true,
|
||||
"priority_mode": "provider",
|
||||
"scheduling_mode": "cache_affinity",
|
||||
"keep_priority_on_conversion": false,
|
||||
"sticky_key_attempts": DEFAULT_STICKY_KEY_ATTEMPTS
|
||||
})
|
||||
);
|
||||
@@ -386,9 +292,39 @@ mod tests {
|
||||
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);
|
||||
let config = read_system_default_routing_ordering_config(&state)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(config, SchedulerOrderingConfig::default());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bootstrap_does_not_migrate_legacy_scheduler_keys() {
|
||||
let repository = Arc::new(InMemoryRoutingGroupRepository::default());
|
||||
let state = AppState::new().unwrap().with_data_state_for_tests(
|
||||
GatewayDataState::disabled()
|
||||
.with_system_config_values_for_tests([
|
||||
("provider_priority_mode".to_string(), json!("global_key")),
|
||||
("scheduling_mode".to_string(), json!("load_balance")),
|
||||
("keep_priority_on_conversion".to_string(), json!(true)),
|
||||
])
|
||||
.with_routing_group_repository_for_tests(repository),
|
||||
);
|
||||
|
||||
let created = state
|
||||
.ensure_system_default_routing_group_inner()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("bootstrap should create the strategy");
|
||||
assert_eq!(
|
||||
created.config_json["default_policy"],
|
||||
json!({
|
||||
"priority_mode": "provider",
|
||||
"scheduling_mode": "cache_affinity",
|
||||
"keep_priority_on_conversion": false,
|
||||
"sticky_key_attempts": DEFAULT_STICKY_KEY_ATTEMPTS
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,8 +10,6 @@ use async_trait::async_trait;
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
use super::config::SchedulerOrderingConfig;
|
||||
|
||||
#[async_trait]
|
||||
pub(crate) trait SchedulerRuntimeState {
|
||||
async fn read_provider_quota_snapshot(
|
||||
@@ -60,7 +58,4 @@ pub(crate) trait SchedulerRuntimeState {
|
||||
max_entries: usize,
|
||||
expected_epoch: Option<u64>,
|
||||
) -> bool;
|
||||
|
||||
async fn read_scheduler_ordering_config(&self)
|
||||
-> Result<SchedulerOrderingConfig, GatewayError>;
|
||||
}
|
||||
|
||||
@@ -79,12 +79,7 @@ const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(30);
|
||||
// five minutes of total age. Direct database edits that bypass AppState
|
||||
// invalidation can therefore take at most this bounded interval to appear.
|
||||
const SYSTEM_CONFIG_CACHE_MAX_STALENESS: Duration = Duration::from_secs(5 * 60);
|
||||
const SCHEDULER_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
|
||||
"enable_format_conversion",
|
||||
"keep_priority_on_conversion",
|
||||
"provider_priority_mode",
|
||||
"scheduling_mode",
|
||||
];
|
||||
const SCHEDULER_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &["enable_format_conversion"];
|
||||
const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
|
||||
crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY,
|
||||
crate::constants::ANTIGRAVITY_BEARER_BRIDGE_CONFIG_KEY,
|
||||
@@ -4350,12 +4345,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn system_config_entry_write_refreshes_cache_and_scheduler_affinity_for_routing_keys() {
|
||||
async fn system_config_entry_write_refreshes_cache_and_scheduler_affinity_for_format_conversion(
|
||||
) {
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::disabled().with_system_config_values_for_tests([(
|
||||
"keep_priority_on_conversion".to_string(),
|
||||
"enable_format_conversion".to_string(),
|
||||
json!(false),
|
||||
)]),
|
||||
);
|
||||
@@ -4364,7 +4360,7 @@ mod tests {
|
||||
|
||||
assert_eq!(
|
||||
state
|
||||
.read_system_config_json_value("keep_priority_on_conversion")
|
||||
.read_system_config_json_value("enable_format_conversion")
|
||||
.await
|
||||
.expect("system config read should succeed"),
|
||||
Some(json!(false))
|
||||
@@ -4385,13 +4381,13 @@ mod tests {
|
||||
|
||||
let initial_epoch = state.scheduler_affinity_epoch();
|
||||
state
|
||||
.upsert_system_config_entry("keep_priority_on_conversion", &json!(true), None)
|
||||
.upsert_system_config_entry("enable_format_conversion", &json!(true), None)
|
||||
.await
|
||||
.expect("admin config write should succeed");
|
||||
|
||||
assert_eq!(
|
||||
state
|
||||
.read_system_config_json_value("keep_priority_on_conversion")
|
||||
.read_system_config_json_value("enable_format_conversion")
|
||||
.await
|
||||
.expect("system config read should use refreshed cache"),
|
||||
Some(json!(true))
|
||||
|
||||
@@ -680,10 +680,4 @@ impl SchedulerRuntimeState for AppState {
|
||||
expected_epoch,
|
||||
)
|
||||
}
|
||||
|
||||
async fn read_scheduler_ordering_config(
|
||||
&self,
|
||||
) -> Result<crate::scheduler::config::SchedulerOrderingConfig, GatewayError> {
|
||||
crate::scheduler::config::read_scheduler_ordering_config(self).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,10 +15,9 @@ 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).
|
||||
/// When none exists, one is created from the routing defaults.
|
||||
/// 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> {
|
||||
@@ -44,16 +43,12 @@ impl AppState {
|
||||
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"
|
||||
"no system default routing group exists and routing storage is read-only; scheduler uses routing defaults"
|
||||
);
|
||||
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 = RoutingGroupConfig::default();
|
||||
let config_json = serde_json::to_value(config)
|
||||
.map_err(|err| GatewayError::Internal(format!("serialize routing config: {err}")))?;
|
||||
|
||||
@@ -75,9 +70,10 @@ impl AppState {
|
||||
self.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
name,
|
||||
description: Some("自动从旧版调度配置迁移生成的系统默认策略".to_string()),
|
||||
description: Some("系统默认调度策略".to_string()),
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
sort_order: 0,
|
||||
config_json,
|
||||
version: 1,
|
||||
created_at: now,
|
||||
|
||||
Reference in New Issue
Block a user