Merge remote-tracking branch 'origin/main' into codex/fix-antigravity-quota

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