mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 17:37:47 +08:00
feat(routing): consolidate scheduling strategy configuration
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));
|
||||
|
||||
Reference in New Issue
Block a user