Unify candidate ranking pipeline

This commit is contained in:
fawney19
2026-04-27 12:34:03 +08:00
parent 9b866a6d17
commit 3b542434a2
46 changed files with 3635 additions and 1914 deletions

View File

@@ -1,393 +1 @@
use std::sync::Arc; pub(crate) use super::candidate_resolution::*;
use aether_provider_transport::provider_types::provider_type_is_fixed;
use tracing::warn;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
use crate::orchestration::LocalExecutionCandidateMetadata;
use super::candidate_affinity::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
pub(crate) provider_api_format: String,
pub(crate) orchestration: LocalExecutionCandidateMetadata,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct SkippedLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) skip_reason: &'static str,
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
pub(crate) extra_data: Option<serde_json::Value>,
}
impl SkippedLocalExecutionCandidate {
pub(crate) fn transport_ref(&self) -> Option<&GatewayProviderTransportSnapshot> {
self.transport.as_deref()
}
}
pub(crate) async fn filter_and_rank_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: &str,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.trim();
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
required_capabilities,
sticky_session_token,
|candidate, transport, normalized_client_api_format| {
current_local_execution_candidate_skip_reason_with_transport(
candidate,
transport,
normalized_client_api_format,
requested_model,
)
},
)
.await
}
pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: Option<&str>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.map(str::trim);
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
required_capabilities,
sticky_session_token,
|candidate, transport, _normalized_client_api_format| {
current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
requested_model,
)
},
)
.await
}
async fn filter_and_rank_local_execution_candidates_with_gate<F>(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
runtime_skip_reason: F,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
)
where
F: Fn(
&SchedulerMinimalCandidateSelectionCandidate,
&GatewayProviderTransportSnapshot,
&str,
) -> Option<&'static str>,
{
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
let mut selectable = Vec::with_capacity(candidates.len());
let mut skipped = Vec::with_capacity(candidates.len());
for candidate in candidates {
let Some(transport) = read_candidate_transport_snapshot(state, &candidate).await else {
skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason: "transport_snapshot_missing",
transport: None,
extra_data: None,
});
continue;
};
let transport = Arc::new(transport);
if candidate_is_ineligible_due_to_disabled_format_conversion(
transport.as_ref(),
normalized_client_api_format.as_str(),
) {
continue;
}
match runtime_skip_reason(
&candidate,
transport.as_ref(),
normalized_client_api_format.as_str(),
) {
Some(skip_reason) => skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason,
transport: Some(transport),
extra_data: None,
}),
None => selectable.push(EligibleLocalExecutionCandidate {
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
candidate,
transport,
orchestration: LocalExecutionCandidateMetadata::default(),
}),
}
}
let ranked = rank_eligible_local_execution_candidates(
state,
selectable,
normalized_client_api_format.as_str(),
required_capabilities,
)
.await;
let (ranked, pool_skipped) =
apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
skipped.extend(pool_skipped);
(ranked, skipped)
}
pub(crate) fn extract_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
fn non_empty_str(value: Option<&serde_json::Value>) -> Option<&str> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
let object = body_json.as_object()?;
non_empty_str(object.get("prompt_cache_key"))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.or_else(|| non_empty_str(object.get("session_id")))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| {
object
.get("metadata")
.and_then(serde_json::Value::as_object)
.and_then(|metadata| {
non_empty_str(metadata.get("session_id"))
.or_else(|| non_empty_str(metadata.get("conversation_id")))
})
})
.or_else(|| {
object
.get("conversationState")
.and_then(serde_json::Value::as_object)
.and_then(|state| {
non_empty_str(state.get("conversationId"))
.or_else(|| non_empty_str(state.get("sessionId")))
})
})
.map(ToOwned::to_owned)
}
fn current_local_execution_candidate_common_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
requested_model: Option<&str>,
) -> Option<&'static str> {
let requested_model = requested_model.unwrap_or_default();
if !transport.provider.is_active {
return Some("provider_inactive");
}
if !transport.endpoint.is_active {
return Some("endpoint_inactive");
}
if !transport.key.is_active {
return Some("key_inactive");
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if !candidate
.endpoint_api_format
.trim()
.eq_ignore_ascii_case(endpoint_api_format)
&& !api_format_matches(&candidate.endpoint_api_format, endpoint_api_format)
{
return Some("endpoint_api_format_changed");
}
if !transport_key_supports_api_format(transport, endpoint_api_format) {
return Some("key_api_format_disabled");
}
if !transport_key_allows_candidate_model(transport, requested_model, candidate) {
return Some("key_model_disabled");
}
None
}
fn candidate_is_ineligible_due_to_disabled_format_conversion(
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
) -> bool {
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return false;
}
crate::ai_pipeline::conversion::request_conversion_kind(
normalized_client_api_format,
endpoint_api_format,
)
.is_some()
&& crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
normalized_client_api_format,
endpoint_api_format,
)
&& !crate::ai_pipeline::conversion::request_conversion_enabled_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
)
}
fn current_local_execution_candidate_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
requested_model: &str,
) -> Option<&'static str> {
if let Some(skip_reason) = current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
Some(requested_model),
) {
return Some(skip_reason);
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return None;
}
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
) {
return Some("transport_unsupported");
}
None
}
fn transport_key_supports_api_format(
transport: &GatewayProviderTransportSnapshot,
endpoint_api_format: &str,
) -> bool {
let provider_type = transport.provider.provider_type.trim();
let auth_type = transport.key.auth_type.trim();
let inherits_provider_api_formats = provider_type_is_fixed(provider_type)
&& (auth_type.eq_ignore_ascii_case("oauth")
|| (provider_type.eq_ignore_ascii_case("kiro")
&& auth_type.eq_ignore_ascii_case("bearer")
&& transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())));
if inherits_provider_api_formats {
return true;
}
match transport.key.api_formats.as_deref() {
None => true,
Some(formats) => formats
.iter()
.any(|value| api_format_matches(value, endpoint_api_format)),
}
}
fn normalize_api_format_alias(value: &str) -> String {
crate::ai_pipeline::normalize_legacy_openai_format_alias(value)
}
fn api_format_matches(left: &str, right: &str) -> bool {
normalize_api_format_alias(left) == normalize_api_format_alias(right)
}
fn transport_key_allows_candidate_model(
transport: &GatewayProviderTransportSnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
let Some(allowed_models) = transport.key.allowed_models.as_deref() else {
return true;
};
let requested_model = requested_model.trim();
let global_model_name = candidate.global_model_name.trim();
let selected_provider_model_name = candidate.selected_provider_model_name.trim();
let mapping_matched_model = candidate
.mapping_matched_model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
for allowed_model in allowed_models.iter().map(String::as_str).map(str::trim) {
if allowed_model.is_empty() {
continue;
}
if allowed_model == requested_model
|| allowed_model == global_model_name
|| allowed_model == selected_provider_model_name
|| mapping_matched_model.is_some_and(|value| value == allowed_model)
{
return true;
}
}
false
}
pub(crate) async fn read_candidate_transport_snapshot(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Option<GatewayProviderTransportSnapshot> {
match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(transport)) => Some(transport),
Ok(None) => None,
Err(error) => {
warn!(
event_name = "candidate_eligibility_transport_load_failed",
log_type = "event",
provider_id = %candidate.provider_id,
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
error = ?error,
"failed to load provider transport while evaluating local candidate eligibility"
);
None
}
}
}

View File

@@ -1,9 +1,10 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use aether_scheduler_core::SchedulerRankingOutcome;
use serde_json::Value; use serde_json::Value;
use uuid::Uuid; use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::remember_scheduler_affinity_for_candidate; use crate::ai_pipeline::planner::candidate_affinity::remember_scheduler_affinity_for_candidate;
use crate::ai_pipeline::planner::candidate_eligibility::{ use crate::ai_pipeline::planner::candidate_resolution::{
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate, EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
}; };
use crate::ai_pipeline::planner::failure_diagnostic::CandidateFailureDiagnostic; use crate::ai_pipeline::planner::failure_diagnostic::CandidateFailureDiagnostic;
@@ -70,7 +71,10 @@ pub(crate) fn remember_first_local_candidate_affinity(
} }
fn should_persist_available_local_candidate(eligible: &EligibleLocalExecutionCandidate) -> bool { fn should_persist_available_local_candidate(eligible: &EligibleLocalExecutionCandidate) -> bool {
eligible.orchestration.pool_key_index.is_none() eligible
.orchestration
.pool_key_index
.is_none_or(|index| index == 0)
} }
fn should_persist_skipped_local_candidate(candidate: &SkippedLocalExecutionCandidate) -> bool { fn should_persist_skipped_local_candidate(candidate: &SkippedLocalExecutionCandidate) -> bool {
@@ -104,7 +108,10 @@ where
let candidate_index = candidate_index as u32; let candidate_index = candidate_index as u32;
let attempt_slots = local_attempt_slot_count(&eligible.transport); let attempt_slots = local_attempt_slot_count(&eligible.transport);
let pool_key_index = eligible.orchestration.pool_key_index; let pool_key_index = eligible.orchestration.pool_key_index;
let extra_data = build_extra_data(&eligible); let extra_data = local_candidate_extra_data_with_ranking(
build_extra_data(&eligible),
eligible.ranking.as_ref(),
);
let mut owned_eligible = Some(eligible); let mut owned_eligible = Some(eligible);
for retry_index in 0..attempt_slots { for retry_index in 0..attempt_slots {
@@ -153,6 +160,54 @@ where
materialized materialized
} }
fn local_candidate_extra_data_with_ranking(
extra_data: Option<Value>,
ranking: Option<&SchedulerRankingOutcome>,
) -> Option<Value> {
let Some(ranking) = ranking else {
return extra_data;
};
let mut object = match extra_data {
Some(Value::Object(object)) => object,
Some(value) => {
let mut object = serde_json::Map::new();
object.insert("extra".to_string(), value);
object
}
None => serde_json::Map::new(),
};
object.insert(
"ranking_mode".to_string(),
Value::String(format!("{:?}", ranking.ranking_mode)),
);
object.insert(
"priority_mode".to_string(),
Value::String(format!("{:?}", ranking.priority_mode)),
);
object.insert(
"ranking_index".to_string(),
Value::Number(serde_json::Number::from(ranking.ranking_index as u64)),
);
object.insert(
"priority_slot".to_string(),
Value::Number(serde_json::Number::from(i64::from(ranking.priority_slot))),
);
if let Some(promoted_by) = ranking.promoted_by {
object.insert(
"promoted_by".to_string(),
Value::String(promoted_by.to_string()),
);
}
if let Some(demoted_by) = ranking.demoted_by {
object.insert(
"demoted_by".to_string(),
Value::String(demoted_by.to_string()),
);
}
Some(Value::Object(object))
}
pub(crate) async fn persist_available_local_execution_candidates_with_context<F>( pub(crate) async fn persist_available_local_execution_candidates_with_context<F>(
state: PlannerAppState<'_>, state: PlannerAppState<'_>,
trace_id: &str, trace_id: &str,
@@ -356,7 +411,10 @@ mod tests {
GatewayProviderTransportEndpoint, GatewayProviderTransportKey, GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportProvider,
}; };
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; use aether_scheduler_core::{
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankingMode,
SchedulerRankingOutcome,
};
use serde_json::json; use serde_json::json;
use super::*; use super::*;
@@ -456,11 +514,12 @@ mod tests {
candidate_group_id: pool_key_index.map(|_| "pool-group".to_string()), candidate_group_id: pool_key_index.map(|_| "pool-group".to_string()),
pool_key_index, pool_key_index,
}, },
ranking: None,
} }
} }
#[tokio::test] #[tokio::test]
async fn pool_candidates_are_not_persisted_as_available_before_attempt() { async fn pool_group_representatives_are_persisted_as_available_before_attempt() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default()); let repository = Arc::new(InMemoryRequestCandidateRepository::default());
let app = AppState::new() let app = AppState::new()
.expect("state should build") .expect("state should build")
@@ -478,6 +537,7 @@ mod tests {
None, None,
vec![ vec![
sample_eligible("pool-key", Some(0)), sample_eligible("pool-key", Some(0)),
sample_eligible("pool-key-internal", Some(1)),
sample_eligible("normal-key", None), sample_eligible("normal-key", None),
], ],
"persist should not fail", "persist should not fail",
@@ -485,13 +545,74 @@ mod tests {
) )
.await; .await;
assert_eq!(attempts.len(), 2); assert_eq!(attempts.len(), 3);
let stored = app let stored = app
.read_request_candidates_by_request_id("trace-pool-lazy") .read_request_candidates_by_request_id("trace-pool-lazy")
.await .await
.expect("request candidates should read"); .expect("request candidates should read");
assert_eq!(stored.len(), 2);
assert_eq!(stored[0].key_id.as_deref(), Some("pool-key"));
assert_eq!(stored[0].candidate_index, 0);
assert_eq!(stored[1].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[1].candidate_index, 2);
}
#[tokio::test]
async fn available_candidates_persist_ranking_metadata_in_extra_data() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone(
&repository,
)),
);
let mut eligible = sample_eligible("ranked-key", None);
eligible.ranking = Some(SchedulerRankingOutcome {
original_index: 1,
ranking_index: 0,
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
priority_slot: 7,
promoted_by: Some("cached_affinity"),
demoted_by: Some("cross_format"),
});
persist_available_local_execution_candidates(
PlannerAppState::new(&app),
"trace-ranking-extra-data",
"user-1",
"api-key-1",
None,
vec![eligible],
"persist should not fail",
|_| Some(json!({ "existing": "value" })),
)
.await;
let stored = app
.read_request_candidates_by_request_id("trace-ranking-extra-data")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 1); assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key")); let extra_data = stored[0]
.extra_data
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("ranking metadata should persist as object extra data");
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
assert_eq!(
extra_data.get("ranking_mode"),
Some(&json!("CacheAffinity"))
);
assert_eq!(extra_data.get("priority_mode"), Some(&json!("Provider")));
assert_eq!(extra_data.get("ranking_index"), Some(&json!(0)));
assert_eq!(extra_data.get("priority_slot"), Some(&json!(7)));
assert_eq!(
extra_data.get("promoted_by"),
Some(&json!("cached_affinity"))
);
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
} }
#[tokio::test] #[tokio::test]

View File

@@ -2,7 +2,7 @@ use aether_contracts::ProxySnapshot;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::{json, Map, Value}; use serde_json::{json, Map, Value};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate; use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::passthrough::resolve_same_format_provider_transport_unsupported_reason_for_trace; use crate::ai_pipeline::planner::passthrough::resolve_same_format_provider_transport_unsupported_reason_for_trace;
use crate::ai_pipeline::transport::{ use crate::ai_pipeline::transport::{
body_rules_are_locally_supported, header_rules_are_locally_supported, body_rules_are_locally_supported, header_rules_are_locally_supported,

View File

@@ -0,0 +1,414 @@
use std::sync::Arc;
use aether_provider_transport::provider_types::provider_type_is_fixed;
use tracing::warn;
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
use crate::ai_pipeline::{
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, PlannerAppState,
};
use crate::orchestration::LocalExecutionCandidateMetadata;
use super::candidate_affinity::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
pub(crate) provider_api_format: String,
pub(crate) orchestration: LocalExecutionCandidateMetadata,
pub(crate) ranking: Option<SchedulerRankingOutcome>,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct SkippedLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) skip_reason: &'static str,
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
pub(crate) extra_data: Option<serde_json::Value>,
}
impl SkippedLocalExecutionCandidate {
pub(crate) fn transport_ref(&self) -> Option<&GatewayProviderTransportSnapshot> {
self.transport.as_deref()
}
}
pub(crate) async fn filter_and_rank_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: &str,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.trim();
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
Some(requested_model),
auth_snapshot,
required_capabilities,
sticky_session_token,
|candidate, transport, normalized_client_api_format| {
current_local_execution_candidate_skip_reason_with_transport(
candidate,
transport,
normalized_client_api_format,
requested_model,
)
},
)
.await
}
pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: Option<&str>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.map(str::trim);
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
requested_model,
auth_snapshot,
required_capabilities,
sticky_session_token,
|candidate, transport, _normalized_client_api_format| {
current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
requested_model,
)
},
)
.await
}
async fn filter_and_rank_local_execution_candidates_with_gate<F>(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: Option<&str>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
runtime_skip_reason: F,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
)
where
F: Fn(
&SchedulerMinimalCandidateSelectionCandidate,
&GatewayProviderTransportSnapshot,
&str,
) -> Option<&'static str>,
{
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
let mut selectable = Vec::with_capacity(candidates.len());
let mut skipped = Vec::with_capacity(candidates.len());
for candidate in candidates {
let Some(transport) = read_candidate_transport_snapshot(state, &candidate).await else {
skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason: "transport_snapshot_missing",
transport: None,
extra_data: None,
});
continue;
};
let transport = Arc::new(transport);
match runtime_skip_reason(
&candidate,
transport.as_ref(),
normalized_client_api_format.as_str(),
) {
Some(skip_reason) => skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason,
transport: Some(transport),
extra_data: None,
}),
None => selectable.push(EligibleLocalExecutionCandidate {
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
candidate,
transport,
orchestration: LocalExecutionCandidateMetadata::default(),
ranking: None,
}),
}
}
let ranked = rank_eligible_local_execution_candidates(
state,
selectable,
normalized_client_api_format.as_str(),
requested_model,
auth_snapshot,
required_capabilities,
)
.await;
let (ranked, pool_skipped) =
apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
skipped.extend(pool_skipped);
(ranked, skipped)
}
pub(crate) fn extract_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
fn non_empty_str(value: Option<&serde_json::Value>) -> Option<&str> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
let object = body_json.as_object()?;
non_empty_str(object.get("prompt_cache_key"))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.or_else(|| non_empty_str(object.get("session_id")))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| {
object
.get("metadata")
.and_then(serde_json::Value::as_object)
.and_then(|metadata| {
non_empty_str(metadata.get("session_id"))
.or_else(|| non_empty_str(metadata.get("conversation_id")))
})
})
.or_else(|| {
object
.get("conversationState")
.and_then(serde_json::Value::as_object)
.and_then(|state| {
non_empty_str(state.get("conversationId"))
.or_else(|| non_empty_str(state.get("sessionId")))
})
})
.map(ToOwned::to_owned)
}
fn current_local_execution_candidate_common_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
requested_model: Option<&str>,
) -> Option<&'static str> {
let requested_model = requested_model.unwrap_or_default();
if !transport.provider.is_active {
return Some("provider_inactive");
}
if !transport.endpoint.is_active {
return Some("endpoint_inactive");
}
if !transport.key.is_active {
return Some("key_inactive");
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if !candidate
.endpoint_api_format
.trim()
.eq_ignore_ascii_case(endpoint_api_format)
&& !api_format_matches(&candidate.endpoint_api_format, endpoint_api_format)
{
return Some("endpoint_api_format_changed");
}
if !transport_key_supports_api_format(transport, endpoint_api_format) {
return Some("key_api_format_disabled");
}
if !transport_key_allows_candidate_model(transport, requested_model, candidate) {
return Some("key_model_disabled");
}
None
}
fn disabled_format_conversion_skip_reason(
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
) -> Option<&'static str> {
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return None;
}
if crate::ai_pipeline::conversion::request_conversion_kind(
normalized_client_api_format,
endpoint_api_format,
)
.is_none()
{
return None;
}
if crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
normalized_client_api_format,
endpoint_api_format,
) && !crate::ai_pipeline::conversion::request_conversion_enabled_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
) {
return Some("format_conversion_disabled");
}
None
}
fn current_local_execution_candidate_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
requested_model: &str,
) -> Option<&'static str> {
if let Some(skip_reason) = current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
Some(requested_model),
) {
return Some(skip_reason);
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return None;
}
if let Some(skip_reason) =
disabled_format_conversion_skip_reason(transport, normalized_client_api_format)
{
return Some(skip_reason);
}
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
) {
return Some("transport_unsupported");
}
None
}
fn transport_key_supports_api_format(
transport: &GatewayProviderTransportSnapshot,
endpoint_api_format: &str,
) -> bool {
let provider_type = transport.provider.provider_type.trim();
let auth_type = transport.key.auth_type.trim();
let inherits_provider_api_formats = provider_type_is_fixed(provider_type)
&& (auth_type.eq_ignore_ascii_case("oauth")
|| (provider_type.eq_ignore_ascii_case("kiro")
&& auth_type.eq_ignore_ascii_case("bearer")
&& transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())));
if inherits_provider_api_formats {
return true;
}
match transport.key.api_formats.as_deref() {
None => true,
Some(formats) => formats
.iter()
.any(|value| api_format_matches(value, endpoint_api_format)),
}
}
fn normalize_api_format_alias(value: &str) -> String {
crate::ai_pipeline::normalize_legacy_openai_format_alias(value)
}
fn api_format_matches(left: &str, right: &str) -> bool {
normalize_api_format_alias(left) == normalize_api_format_alias(right)
}
fn transport_key_allows_candidate_model(
transport: &GatewayProviderTransportSnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
let Some(allowed_models) = transport.key.allowed_models.as_deref() else {
return true;
};
let requested_model = requested_model.trim();
let global_model_name = candidate.global_model_name.trim();
let selected_provider_model_name = candidate.selected_provider_model_name.trim();
let mapping_matched_model = candidate
.mapping_matched_model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
for allowed_model in allowed_models.iter().map(String::as_str).map(str::trim) {
if allowed_model.is_empty() {
continue;
}
if allowed_model == requested_model
|| allowed_model == global_model_name
|| allowed_model == selected_provider_model_name
|| mapping_matched_model.is_some_and(|value| value == allowed_model)
{
return true;
}
}
false
}
pub(crate) async fn read_candidate_transport_snapshot(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Option<GatewayProviderTransportSnapshot> {
match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(transport)) => Some(transport),
Ok(None) => None,
Err(error) => {
warn!(
event_name = "candidate_resolution_transport_load_failed",
log_type = "event",
provider_id = %candidate.provider_id,
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
error = ?error,
"failed to load provider transport while evaluating local candidate eligibility"
);
None
}
}
}

View File

@@ -9,6 +9,7 @@ mod candidate_eligibility;
mod candidate_materialization; mod candidate_materialization;
mod candidate_metadata; mod candidate_metadata;
mod candidate_preparation; mod candidate_preparation;
mod candidate_resolution;
mod candidate_source; mod candidate_source;
mod common; mod common;
mod decision; mod decision;
@@ -27,7 +28,7 @@ mod specialized;
mod standard; mod standard;
mod state; mod state;
pub(crate) use self::candidate_eligibility::extract_pool_sticky_session_token; pub(crate) use self::candidate_resolution::extract_pool_sticky_session_token;
pub(crate) use self::failure_diagnostic::{ pub(crate) use self::failure_diagnostic::{
CandidateFailureDiagnostic, CandidateFailureDiagnosticKind, CandidateFailureDiagnostic, CandidateFailureDiagnosticKind,
}; };

View File

@@ -1,9 +1,5 @@
use tracing::warn; use tracing::warn;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{ use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context, persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context, persist_skipped_local_execution_candidates_with_context,
@@ -14,6 +10,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate, build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts, LocalExecutionCandidateMetadataParts,
}; };
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::common::extract_requested_model_from_request; use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{ use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input, build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
@@ -112,6 +112,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
candidates, candidates,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
sticky_session_token.as_deref(), sticky_session_token.as_deref(),
) )

View File

@@ -1,9 +1,9 @@
use std::sync::Arc; use std::sync::Arc;
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{ use crate::ai_pipeline::planner::candidate_preparation::{
resolve_candidate_mapped_model, resolve_candidate_oauth_auth, OauthPreparationContext, resolve_candidate_mapped_model, resolve_candidate_oauth_auth, OauthPreparationContext,
}; };
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata; use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_pipeline::transport::kiro::KiroRequestAuth; use crate::ai_pipeline::transport::kiro::KiroRequestAuth;
use crate::ai_pipeline::transport::vertex::resolve_local_vertex_api_key_query_auth; use crate::ai_pipeline::transport::vertex::resolve_local_vertex_api_key_query_auth;

View File

@@ -7,7 +7,7 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
use serde_json::{Map, Value}; use serde_json::{Map, Value};
use tracing::warn; use tracing::warn;
use crate::ai_pipeline::planner::candidate_eligibility::{ use crate::ai_pipeline::planner::candidate_resolution::{
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate, EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
}; };
use crate::ai_pipeline::PlannerAppState; use crate::ai_pipeline::PlannerAppState;
@@ -434,6 +434,7 @@ fn schedule_pool_group(
transport, transport,
provider_api_format, provider_api_format,
orchestration, orchestration,
ranking,
} = eligible; } = eligible;
let key_id = candidate.key_id.clone(); let key_id = candidate.key_id.clone();
let mut key_context = key_context_by_id.get(&key_id).cloned().unwrap_or_default(); let mut key_context = key_context_by_id.get(&key_id).cloned().unwrap_or_default();
@@ -495,6 +496,7 @@ fn schedule_pool_group(
transport, transport,
provider_api_format, provider_api_format,
orchestration, orchestration,
ranking,
}, },
key_context, key_context,
original_index, original_index,
@@ -1040,7 +1042,7 @@ mod tests {
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context, apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
normalize_enabled_pool_presets, PoolCatalogKeyContext, normalize_enabled_pool_presets, PoolCatalogKeyContext,
}; };
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate; use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::PlannerAppState; use crate::ai_pipeline::PlannerAppState;
use crate::data::GatewayDataState; use crate::data::GatewayDataState;
use crate::handlers::shared::provider_pool::{ use crate::handlers::shared::provider_pool::{
@@ -1836,6 +1838,7 @@ mod tests {
}, },
provider_api_format: "openai:chat".to_string(), provider_api_format: "openai:chat".to_string(),
orchestration: LocalExecutionCandidateMetadata::default(), orchestration: LocalExecutionCandidateMetadata::default(),
ranking: None,
transport: Arc::new(crate::ai_pipeline::GatewayProviderTransportSnapshot { transport: Arc::new(crate::ai_pipeline::GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider { provider: GatewayProviderTransportProvider {
id: provider_id.to_string(), id: provider_id.to_string(),

View File

@@ -3,7 +3,6 @@ use serde_json::json;
use tracing::warn; use tracing::warn;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates_without_transport_pair_gate;
use crate::ai_pipeline::planner::candidate_materialization::{ use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate,
mark_skipped_local_execution_candidate_with_failure_diagnostic, mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -15,6 +14,7 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata, build_local_execution_candidate_metadata,
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts, build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
}; };
use crate::ai_pipeline::planner::candidate_resolution::filter_and_rank_local_execution_candidates_without_transport_pair_gate;
use crate::ai_pipeline::planner::decision_input::{ use crate::ai_pipeline::planner::decision_input::{
build_local_authenticated_decision_input, resolve_local_authenticated_decision_input, build_local_authenticated_decision_input, resolve_local_authenticated_decision_input,
}; };
@@ -95,6 +95,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
candidates, candidates,
GEMINI_FILES_CLIENT_API_FORMAT, GEMINI_FILES_CLIENT_API_FORMAT,
None, None,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
None, None,
) )

View File

@@ -1,10 +1,6 @@
use tracing::warn; use tracing::warn;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{ use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate,
mark_skipped_local_execution_candidate_with_failure_diagnostic, mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -16,6 +12,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata, build_local_execution_candidate_metadata,
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts, build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
}; };
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::decision_input::{ use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input, build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
}; };
@@ -157,6 +157,7 @@ async fn materialize_local_openai_image_candidate_attempts(
candidates, candidates,
api_format, api_format,
&input.requested_model, &input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
sticky_session_token.as_deref(), sticky_session_token.as_deref(),
) )

View File

@@ -3,10 +3,6 @@ use tracing::warn;
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec}; use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{ use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate,
mark_skipped_local_execution_candidate_with_failure_diagnostic, mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -18,6 +14,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata, build_local_execution_candidate_metadata,
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts, build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
}; };
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::common::extract_requested_model_from_request; use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{ use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input, build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
@@ -169,6 +169,7 @@ async fn materialize_local_video_create_candidate_attempts(
candidates, candidates,
api_format, api_format,
&input.requested_model, &input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
sticky_session_token.as_deref(), sticky_session_token.as_deref(),
) )

View File

@@ -3,10 +3,6 @@ use std::collections::BTreeSet;
use tracing::warn; use tracing::warn;
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind}; use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{ use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context, persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context, persist_skipped_local_execution_candidates_with_context,
@@ -17,6 +13,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate, build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts, LocalExecutionCandidateMetadataParts,
}; };
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate; use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::planner::common::extract_requested_model_from_request; use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{ use crate::ai_pipeline::planner::decision_input::{
@@ -178,6 +178,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
candidates, candidates,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
sticky_session_token.as_deref(), sticky_session_token.as_deref(),
) )

View File

@@ -4,10 +4,10 @@ use std::sync::Arc;
use serde_json::Value; use serde_json::Value;
use crate::ai_pipeline::conversion::{request_conversion_direct_auth, request_conversion_kind}; use crate::ai_pipeline::conversion::{request_conversion_direct_auth, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{ use crate::ai_pipeline::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext, prepare_header_authenticated_candidate, OauthPreparationContext,
}; };
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND; use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_pipeline::planner::standard::{ use crate::ai_pipeline::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body, apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,

View File

@@ -1,10 +1,6 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{ use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data, mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic, mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -17,6 +13,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate, build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts, LocalExecutionCandidateMetadataParts,
}; };
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::materialization_policy::{ use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind, build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
}; };
@@ -135,6 +135,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
candidates, candidates,
"openai:chat", "openai:chat",
&input.requested_model, &input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
sticky_session_token.as_deref(), sticky_session_token.as_deref(),
) )

View File

@@ -4,7 +4,7 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::super::{GatewayError, LocalOpenAiChatDecisionInput}; use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
use crate::ai_pipeline::conversion::request_candidate_api_formats; use crate::ai_pipeline::conversion::request_candidate_api_formats;
use crate::ai_pipeline::planner::candidate_eligibility::SkippedLocalExecutionCandidate; use crate::ai_pipeline::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate; use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::PlannerAppState; use crate::ai_pipeline::PlannerAppState;
use crate::clock::current_unix_secs; use crate::clock::current_unix_secs;

View File

@@ -5,10 +5,10 @@ use serde_json::Value;
use tracing::debug; use tracing::debug;
use crate::ai_pipeline::conversion::{request_conversion_direct_auth, request_conversion_kind}; use crate::ai_pipeline::conversion::{request_conversion_direct_auth, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{ use crate::ai_pipeline::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext, prepare_header_authenticated_candidate, OauthPreparationContext,
}; };
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider; use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
use crate::ai_pipeline::planner::spec_metadata::local_openai_responses_spec_metadata; use crate::ai_pipeline::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_pipeline::planner::standard::{ use crate::ai_pipeline::planner::standard::{

View File

@@ -5,10 +5,6 @@ use tracing::warn;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext; use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind}; use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{ use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data, mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic, mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -21,6 +17,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate, build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts, LocalExecutionCandidateMetadataParts,
}; };
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate; use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::planner::common::extract_standard_requested_model; use crate::ai_pipeline::planner::common::extract_standard_requested_model;
use crate::ai_pipeline::planner::decision_input::{ use crate::ai_pipeline::planner::decision_input::{
@@ -234,6 +234,7 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
candidates, candidates,
spec_metadata.api_format, spec_metadata.api_format,
&input.requested_model, &input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(), input.required_capabilities.as_ref(),
sticky_session_token.as_deref(), sticky_session_token.as_deref(),
) )

View File

@@ -2,9 +2,10 @@ use aether_data::DataLayerError;
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow; use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_scheduler_core::{ use aether_scheduler_core::{
auth_constraints_allow_api_format, build_minimal_candidate_selection, auth_constraints_allow_api_format, build_minimal_candidate_selection,
collect_global_model_names_for_required_capability, normalize_api_format, collect_global_model_names_for_required_capability, enumerate_minimal_candidate_selection,
resolve_requested_global_model_name, BuildMinimalCandidateSelectionInput, normalize_api_format, resolve_requested_global_model_name, row_supports_requested_model,
SchedulerAuthConstraints, SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, BuildMinimalCandidateSelectionInput, SchedulerAuthConstraints,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
}; };
use async_trait::async_trait; use async_trait::async_trait;
use std::collections::BTreeSet; use std::collections::BTreeSet;
@@ -30,31 +31,24 @@ pub(crate) async fn read_requested_model_rows(
api_format: &str, api_format: &str,
requested_model_name: &str, requested_model_name: &str,
) -> Result<Option<(String, Vec<StoredMinimalCandidateSelectionRow>)>, DataLayerError> { ) -> Result<Option<(String, Vec<StoredMinimalCandidateSelectionRow>)>, DataLayerError> {
let exact_rows = state
.read_minimal_candidate_selection_rows_for_api_format_and_global_model(
api_format,
requested_model_name,
)
.await?;
if !exact_rows.is_empty() {
return Ok(Some((requested_model_name.to_string(), exact_rows)));
}
let rows = state let rows = state
.read_minimal_candidate_selection_rows_for_api_format(api_format) .read_minimal_candidate_selection_rows_for_api_format(api_format)
.await?; .await?;
let rows = rows
.into_iter()
.filter(|row| row_supports_requested_model(row, requested_model_name, api_format))
.collect::<Vec<_>>();
if rows.is_empty() {
return Ok(None);
}
let Some(resolved_global_model_name) = let Some(resolved_global_model_name) =
resolve_requested_global_model_name(&rows, requested_model_name, api_format) resolve_requested_global_model_name(&rows, requested_model_name, api_format)
else { else {
return Ok(None); return Ok(None);
}; };
Ok(Some(( Ok(Some((resolved_global_model_name, rows)))
resolved_global_model_name.clone(),
rows.into_iter()
.filter(|row| row.global_model_name == resolved_global_model_name)
.collect(),
)))
} }
pub(crate) async fn read_minimal_candidate_selection( pub(crate) async fn read_minimal_candidate_selection(
@@ -180,6 +174,45 @@ pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_affi
}) })
} }
pub(crate) async fn enumerate_minimal_candidate_selection_with_required_capabilities(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !auth_constraints_allow_api_format(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
&normalized_api_format,
) {
return Ok(Vec::new());
}
let Some((resolved_global_model_name, rows)) =
read_requested_model_rows(state, &normalized_api_format, requested_model_name).await?
else {
return Ok(Vec::new());
};
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
enumerate_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows,
normalized_api_format: &normalized_api_format,
requested_model_name,
resolved_global_model_name: resolved_global_model_name.as_str(),
require_streaming,
required_capabilities,
auth_constraints: auth_constraints.as_ref(),
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
}
pub(crate) async fn read_global_model_names_for_required_capability( pub(crate) async fn read_global_model_names_for_required_capability(
state: &(impl MinimalCandidateSelectionRowSource + Sync), state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str, api_format: &str,

View File

@@ -271,7 +271,10 @@ where
fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool { fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool {
let metadata = local_execution_candidate_metadata_from_report_context(report_context); let metadata = local_execution_candidate_metadata_from_report_context(report_context);
metadata.candidate_group_id.is_some() && metadata.pool_key_index.is_some() metadata.candidate_group_id.is_some()
&& metadata
.pool_key_index
.is_some_and(|pool_key_index| pool_key_index > 0)
} }
fn resolve_stream_candidate_watchdog_timeout(plan: &aether_contracts::ExecutionPlan) -> Duration { fn resolve_stream_candidate_watchdog_timeout(plan: &aether_contracts::ExecutionPlan) -> Duration {
@@ -484,6 +487,10 @@ mod tests {
"candidate_group_id": "pool-group", "candidate_group_id": "pool-group",
"pool_key_index": 1, "pool_key_index": 1,
})))); }))));
assert!(!should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group",
"pool_key_index": 0,
}))));
assert!(!should_skip_unused_persistence(Some(&json!({ assert!(!should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group", "candidate_group_id": "pool-group",
})))); }))));

View File

@@ -0,0 +1,28 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
enumerate_minimal_candidate_selection_with_required_capabilities,
MinimalCandidateSelectionRowSource,
};
use crate::GatewayError;
pub(super) async fn enumerate_scheduler_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
enumerate_minimal_candidate_selection_with_required_capabilities(
selection_row_source,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
required_capabilities,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}

View File

@@ -1,10 +1,12 @@
use self::affinity::candidate_affinity_hash;
use self::selection::{ use self::selection::{
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons, collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
}; };
use super::state::SchedulerRuntimeState; use super::state::SchedulerRuntimeState;
mod affinity; mod affinity;
mod enumeration;
mod ranking;
mod resolution;
mod runtime; mod runtime;
mod selection; mod selection;

View File

@@ -0,0 +1,76 @@
use aether_scheduler_core::{
apply_scheduler_candidate_ranking, effective_provider_key_health_score,
matches_affinity_target, provider_key_health_bucket,
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
};
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use super::affinity::candidate_affinity_hash;
use super::runtime::CandidateRuntimeSelectionSnapshot;
use super::SchedulerMinimalCandidateSelectionCandidate;
pub(super) fn rank_scheduler_candidates(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
ordering_config: SchedulerOrderingConfig,
required_capabilities: Option<&serde_json::Value>,
priority_affinity_key: Option<&str>,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
now_unix_secs: u64,
) {
let rankables = candidates
.iter()
.enumerate()
.map(|(index, candidate)| {
let provider_key = runtime_snapshot
.provider_key_rpm_states
.get(&candidate.key_id);
SchedulerRankableCandidate::from_candidate(candidate, index)
.with_capability_priority(requested_capability_priority_for_candidate(
required_capabilities,
candidate,
))
.with_cached_affinity_match(
cached_affinity_target
.is_some_and(|target| matches_affinity_target(candidate, target)),
)
.with_affinity_hash(
priority_affinity_key.map(|key| candidate_affinity_hash(key, candidate)),
)
.with_health(
provider_key.and_then(|key| {
provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
}),
provider_key
.and_then(|key| {
effective_provider_key_health_score(
key,
candidate.endpoint_api_format.as_str(),
)
})
.unwrap_or(1.0),
)
})
.collect::<Vec<_>>();
apply_scheduler_candidate_ranking(
candidates,
&rankables,
SchedulerRankingContext {
priority_mode: ordering_config.priority_mode,
ranking_mode: scheduler_ranking_mode(ordering_config.scheduling_mode),
include_health: true,
load_balance_seed: now_unix_secs,
},
);
}
fn scheduler_ranking_mode(mode: SchedulerSchedulingMode) -> SchedulerRankingMode {
match mode {
SchedulerSchedulingMode::FixedOrder => SchedulerRankingMode::FixedOrder,
SchedulerSchedulingMode::CacheAffinity => SchedulerRankingMode::CacheAffinity,
SchedulerSchedulingMode::LoadBalance => SchedulerRankingMode::LoadBalance,
}
}

View File

@@ -0,0 +1,45 @@
use std::collections::BTreeSet;
use aether_scheduler_core::SchedulerAffinityTarget;
use super::affinity::candidate_key;
use super::runtime::{current_candidate_runtime_skip_reason, CandidateRuntimeSelectionSnapshot};
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerSkippedCandidate};
pub(super) fn resolve_scheduler_candidate_selectability(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
) -> (
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
) {
let mut selected = Vec::with_capacity(candidates.len());
let mut skipped = Vec::new();
let mut emitted_selected_keys = BTreeSet::new();
let mut emitted_skipped_keys = BTreeSet::new();
for candidate in candidates {
let key = candidate_key(&candidate);
if let Some(skip_reason) = current_candidate_runtime_skip_reason(
&candidate,
runtime_snapshot,
now_unix_secs,
cached_affinity_target,
) {
if emitted_skipped_keys.insert(key) {
skipped.push(SchedulerSkippedCandidate {
candidate,
skip_reason,
});
}
continue;
}
if emitted_selected_keys.insert(key) {
selected.push(candidate);
}
}
(selected, skipped)
}

View File

@@ -1,27 +1,15 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::{
collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health as reorder_candidates_by_scheduler_health_in_core,
SchedulerPriorityMode,
};
use crate::data::auth::GatewayAuthApiKeySnapshot; use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{ use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities,
MinimalCandidateSelectionRowSource,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL; use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::SchedulerSchedulingMode; use crate::scheduler::config::SchedulerSchedulingMode;
use crate::GatewayError; use crate::GatewayError;
use super::affinity::{ use super::affinity::{build_scheduler_affinity_cache_key, remember_scheduler_affinity};
build_scheduler_affinity_cache_key, candidate_key, remember_scheduler_affinity, use super::enumeration::enumerate_scheduler_candidates;
}; use super::ranking::rank_scheduler_candidates;
use super::resolution::resolve_scheduler_candidate_selectability;
use super::runtime::{ use super::runtime::{
auth_snapshot_concurrency_limit_reached, current_candidate_runtime_skip_reason, auth_snapshot_concurrency_limit_reached, read_candidate_runtime_selection_snapshot,
read_candidate_runtime_selection_snapshot,
}; };
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState}; use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
@@ -44,69 +32,6 @@ pub(super) fn is_exact_all_skipped_by_auth_limit(
.all(|candidate| candidate.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON) .all(|candidate| candidate.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
} }
pub(super) fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
required_capabilities: Option<&serde_json::Value>,
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) {
reorder_candidates_by_scheduler_health_in_core(
candidates,
provider_key_rpm_states,
required_capabilities,
affinity_key,
priority_mode,
);
}
fn apply_load_balance_rotation(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
priority_mode: SchedulerPriorityMode,
now_unix_secs: u64,
) {
if candidates.len() < 2 {
return;
}
let mut start = 0usize;
while start < candidates.len() {
let mut end = start + 1;
while end < candidates.len()
&& candidates_share_load_balance_group(
&candidates[start],
&candidates[end],
priority_mode,
)
{
end += 1;
}
let group_len = end - start;
if group_len > 1 {
let offset = usize::try_from(now_unix_secs).unwrap_or(0) % group_len;
candidates[start..end].rotate_left(offset);
}
start = end;
}
}
fn candidates_share_load_balance_group(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
) -> bool {
match priority_mode {
SchedulerPriorityMode::Provider => {
left.provider_priority == right.provider_priority
&& left.key_internal_priority == right.key_internal_priority
}
SchedulerPriorityMode::GlobalKey => {
left.key_global_priority_for_format == right.key_global_priority_for_format
}
}
}
#[cfg_attr(not(test), allow(dead_code))] #[cfg_attr(not(test), allow(dead_code))]
pub(super) async fn select_minimal_candidate( pub(super) async fn select_minimal_candidate(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync), selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
@@ -182,36 +107,18 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
let ordering_config = runtime_state.read_scheduler_ordering_config().await?; let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let priority_affinity_key = let priority_affinity_key =
scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode); scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
let mut candidates = let mut candidates = enumerate_scheduler_candidates(
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
selection_row_source, selection_row_source,
api_format, api_format,
global_model_name, global_model_name,
require_streaming, require_streaming,
auth_snapshot,
ordering_config.priority_mode,
priority_affinity_key,
required_capabilities, required_capabilities,
auth_snapshot,
) )
.await .await?;
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let runtime_snapshot = let runtime_snapshot =
read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs) read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
.await?; .await?;
reorder_candidates_by_scheduler_health(
&mut candidates,
&runtime_snapshot.provider_key_rpm_states,
required_capabilities,
priority_affinity_key,
ordering_config.priority_mode,
);
if ordering_config.scheduling_mode == SchedulerSchedulingMode::LoadBalance {
apply_load_balance_rotation(
&mut candidates,
ordering_config.priority_mode,
now_unix_secs,
);
}
let affinity_cache_key = let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name); build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let cached_affinity_target = if ordering_config.scheduling_mode let cached_affinity_target = if ordering_config.scheduling_mode
@@ -225,6 +132,15 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
}; };
if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) { if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) {
rank_scheduler_candidates(
&mut candidates,
&runtime_snapshot,
ordering_config,
required_capabilities,
priority_affinity_key,
cached_affinity_target.as_ref(),
now_unix_secs,
);
return Ok(( return Ok((
Vec::new(), Vec::new(),
candidates candidates
@@ -237,37 +153,23 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
)); ));
} }
let mut selected_keys = BTreeSet::new(); let (mut selected, skipped) = resolve_scheduler_candidate_selectability(
let mut skipped = Vec::new(); candidates,
let mut emitted_skipped_keys = BTreeSet::new(); &runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
);
rank_scheduler_candidates(
&mut selected,
&runtime_snapshot,
ordering_config,
required_capabilities,
priority_affinity_key,
cached_affinity_target.as_ref(),
now_unix_secs,
);
for candidate in &candidates { Ok((selected, skipped))
let key = candidate_key(candidate);
if let Some(skip_reason) = current_candidate_runtime_skip_reason(
candidate,
&runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
) {
if emitted_skipped_keys.insert(key) {
skipped.push(SchedulerSkippedCandidate {
candidate: candidate.clone(),
skip_reason,
});
}
continue;
}
selected_keys.insert(key);
}
Ok((
collect_selectable_candidates_from_keys(
candidates,
&selected_keys,
cached_affinity_target.as_ref(),
),
skipped,
))
} }
fn scheduling_priority_affinity_key<'a>( fn scheduling_priority_affinity_key<'a>(

View File

@@ -144,6 +144,54 @@ async fn read_minimal_candidate_selection_resolves_provider_model_alias() {
assert_eq!(selection[0].selected_provider_model_name, "gpt-5.2"); assert_eq!(selection[0].selected_provider_model_name, "gpt-5.2");
} }
#[tokio::test]
async fn read_minimal_candidate_selection_keeps_all_rows_supporting_requested_model() {
let mut exact = sample_row();
exact.provider_id = "provider-exact".to_string();
exact.endpoint_id = "endpoint-exact".to_string();
exact.key_id = "key-exact".to_string();
exact.model_id = "model-exact".to_string();
exact.global_model_id = "global-exact".to_string();
exact.global_model_name = "gpt-5".to_string();
exact.model_provider_model_name = "gpt-5".to_string();
exact.model_provider_model_mappings = None;
let mut mapped = sample_row();
mapped.provider_id = "provider-mapped".to_string();
mapped.endpoint_id = "endpoint-mapped".to_string();
mapped.key_id = "key-mapped".to_string();
mapped.model_id = "model-mapped".to_string();
mapped.global_model_id = "global-mapped".to_string();
mapped.global_model_name = "claude-sonnet".to_string();
mapped.global_model_mappings = Some(vec!["gpt-5".to_string()]);
mapped.model_provider_model_name = "claude-sonnet-upstream".to_string();
mapped.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "claude-sonnet-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
exact, mapped,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let selection = read_minimal_candidate_selection(&state, "openai:chat", "gpt-5", false, None)
.await
.expect("selection should succeed");
let provider_ids = selection
.iter()
.map(|candidate| candidate.provider_id.as_str())
.collect::<Vec<_>>();
assert_eq!(provider_ids, vec!["provider-exact", "provider-mapped"]);
assert_eq!(
selection[1].selected_provider_model_name,
"claude-sonnet-upstream"
);
}
#[tokio::test] #[tokio::test]
async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_snapshot() { async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_snapshot() {
let mut row = sample_row(); let mut row = sample_row();

View File

@@ -384,9 +384,9 @@ async fn fixed_order_disables_same_priority_affinity_hash_tiebreaker() {
second.endpoint_id = "endpoint-b".to_string(); second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string(); second.key_id = "key-b".to_string();
second.key_name = "beta".to_string(); second.key_name = "beta".to_string();
second.provider_priority = 0; second.provider_priority = 1;
second.key_internal_priority = 0; second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 0})); second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second, first, second,
@@ -438,9 +438,9 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
second.endpoint_id = "endpoint-b".to_string(); second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string(); second.key_id = "key-b".to_string();
second.key_name = "beta".to_string(); second.key_name = "beta".to_string();
second.provider_priority = 1; second.provider_priority = 0;
second.key_internal_priority = 0; second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 1})); second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second, first, second,

View File

@@ -476,6 +476,23 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
} }
} }
fn sample_disabled_conversion_candidate_row() -> StoredMinimalCandidateSelectionRow {
let mut row = sample_candidate_row();
row.provider_id = "provider-openai-chat-gemini-local-disabled".to_string();
row.provider_priority = 20;
row.endpoint_id = "endpoint-openai-chat-gemini-local-disabled".to_string();
row.key_id = "key-openai-chat-gemini-local-disabled".to_string();
row.model_id = "model-openai-chat-gemini-local-disabled".to_string();
row.global_model_id = "global-model-openai-chat-gemini-local-disabled".to_string();
row.model_provider_model_name = "gemini-2.5-flash-upstream".to_string();
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gemini-2.5-flash-upstream".to_string(),
priority: 2,
api_formats: Some(vec!["gemini:chat".to_string()]),
}]);
row
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider { fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new( StoredProviderCatalogProvider::new(
"provider-openai-chat-gemini-local-1".to_string(), "provider-openai-chat-gemini-local-1".to_string(),
@@ -497,6 +514,27 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
) )
} }
fn sample_disabled_conversion_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-gemini-local-disabled".to_string(),
"gemini".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint { fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new( StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-gemini-local-1".to_string(), "endpoint-openai-chat-gemini-local-1".to_string(),
@@ -522,6 +560,29 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
.expect("endpoint transport should build") .expect("endpoint transport should build")
} }
fn sample_disabled_conversion_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-gemini-local-disabled".to_string(),
"provider-openai-chat-gemini-local-disabled".to_string(),
"gemini:chat".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
None,
None,
Some(2),
Some("/custom/v1beta/models/gemini-2.5-flash-upstream:generateContent".to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey { fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new( StoredProviderCatalogKey::new(
"key-openai-chat-gemini-local-1".to_string(), "key-openai-chat-gemini-local-1".to_string(),
@@ -550,6 +611,34 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
.expect("key transport should build") .expect("key transport should build")
} }
fn sample_disabled_conversion_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-chat-gemini-local-disabled".to_string(),
"provider-openai-chat-gemini-local-disabled".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:chat"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-chat-gemini-disabled",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"gemini:chat": 2})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>)); let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime); let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let seen_report = Arc::new(Mutex::new(false)); let seen_report = Arc::new(Mutex::new(false));
@@ -730,11 +819,21 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
let candidate_selection_repository = let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(), sample_candidate_row(),
sample_disabled_conversion_candidate_row(),
])); ]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()], vec![
vec![sample_provider_catalog_endpoint()], sample_provider_catalog_provider(),
vec![sample_provider_catalog_key()], sample_disabled_conversion_provider_catalog_provider(),
],
vec![
sample_provider_catalog_endpoint(),
sample_disabled_conversion_provider_catalog_endpoint(),
],
vec![
sample_provider_catalog_key(),
sample_disabled_conversion_provider_catalog_key(),
],
)); ));
let (upstream_url, upstream_handle) = start_server(upstream).await; let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -833,8 +932,31 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
.list_by_request_id("trace-openai-chat-gemini-local-123") .list_by_request_id("trace-openai-chat-gemini-local-123")
.await .await
.expect("request candidate trace should read"); .expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1); assert_eq!(stored_candidates.len(), 2);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success); assert_eq!(
stored_candidates
.iter()
.filter(|candidate| candidate.status == RequestCandidateStatus::Success)
.count(),
1
);
let skipped_candidate = stored_candidates
.iter()
.find(|candidate| candidate.status == RequestCandidateStatus::Skipped)
.expect("disabled conversion candidate should be persisted as skipped");
assert_eq!(
skipped_candidate.skip_reason.as_deref(),
Some("format_conversion_disabled")
);
let extra_data = skipped_candidate
.extra_data
.as_ref()
.expect("skipped cross-format candidate extra_data should exist");
assert_eq!(extra_data["execution_strategy"], "local_cross_format");
assert_eq!(
extra_data["transport_diagnostics"]["request_pair"]["conversion_enabled"],
false
);
tokio::time::sleep(std::time::Duration::from_millis(100)).await; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!( assert!(

View File

@@ -423,6 +423,7 @@ fn ai_pipeline_planner_separates_local_candidate_eligibility_from_affinity_ranki
for pattern in [ for pattern in [
"mod candidate_affinity;", "mod candidate_affinity;",
"mod candidate_eligibility;", "mod candidate_eligibility;",
"mod candidate_resolution;",
"mod candidate_preparation;", "mod candidate_preparation;",
] { ] {
assert!( assert!(
@@ -431,19 +432,30 @@ fn ai_pipeline_planner_separates_local_candidate_eligibility_from_affinity_ranki
); );
} }
let candidate_eligibility = let candidate_resolution =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_eligibility.rs"); read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_resolution.rs");
for pattern in [ for pattern in [
"pub(crate) async fn filter_and_rank_local_execution_candidates(", "pub(crate) async fn filter_and_rank_local_execution_candidates(",
"pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(", "pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(",
"pub(crate) async fn read_candidate_transport_snapshot(", "pub(crate) async fn read_candidate_transport_snapshot(",
] { ] {
assert!( assert!(
candidate_eligibility.contains(pattern), candidate_resolution.contains(pattern),
"planner/candidate_eligibility.rs should own {pattern}" "planner/candidate_resolution.rs should own {pattern}"
); );
} }
let candidate_eligibility =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_eligibility.rs");
assert!(
candidate_eligibility.contains("pub(crate) use super::candidate_resolution::*;"),
"planner/candidate_eligibility.rs should remain a compatibility shim"
);
assert!(
!candidate_eligibility.contains("async fn filter_and_rank_local_execution_candidates("),
"planner/candidate_eligibility.rs should not keep resolution implementation"
);
let candidate_affinity = let candidate_affinity =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs"); read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs");
assert!( assert!(
@@ -1233,11 +1245,11 @@ fn ai_pipeline_specialized_files_attempts_consume_eligible_local_candidates_with
assert!( assert!(
specialized_files_support specialized_files_support
.contains("filter_and_rank_local_execution_candidates_without_transport_pair_gate("), .contains("filter_and_rank_local_execution_candidates_without_transport_pair_gate("),
"specialized files support should source runtime gating from candidate_eligibility" "specialized files support should source runtime gating from candidate_resolution"
); );
assert!( assert!(
!specialized_files_support.contains("rank_local_execution_candidates("), !specialized_files_support.contains("rank_local_execution_candidates("),
"specialized files support should not bypass candidate_eligibility with raw affinity ranking" "specialized files support should not bypass candidate_resolution with raw affinity ranking"
); );
let specialized_files_decision = read_workspace_file( let specialized_files_decision = read_workspace_file(

View File

@@ -254,12 +254,20 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"selection.rs should not depend on gateway-local SchedulerAffinityTarget" "selection.rs should not depend on gateway-local SchedulerAffinityTarget"
); );
assert!( assert!(
selection.contains("reorder_candidates_by_scheduler_health_in_core"), selection.contains("enumerate_scheduler_candidates("),
"selection.rs should depend on core candidate reorder helper" "selection.rs should delegate candidate enumeration"
); );
assert!( assert!(
selection.contains("collect_selectable_candidates_from_keys"), selection.contains("read_candidate_runtime_selection_snapshot("),
"selection.rs should depend on core selectable-candidate collector" "selection.rs should delegate runtime snapshot loading"
);
assert!(
selection.contains("resolve_scheduler_candidate_selectability("),
"selection.rs should delegate selectability resolution"
);
assert!(
selection.contains("rank_scheduler_candidates("),
"selection.rs should delegate final ranking"
); );
for pattern in [ for pattern in [
"async fn collect_selectable_candidates(", "async fn collect_selectable_candidates(",
@@ -282,8 +290,10 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"read_provider_concurrent_limits(", "read_provider_concurrent_limits(",
"read_provider_key_rpm_states(", "read_provider_key_rpm_states(",
"candidate_is_selectable_with_runtime_state", "candidate_is_selectable_with_runtime_state",
"collect_selectable_candidates_from_keys",
"auth_api_key_concurrency_limit_reached", "auth_api_key_concurrency_limit_reached",
"build_provider_concurrent_limit_map(", "build_provider_concurrent_limit_map(",
"reorder_candidates_by_scheduler_health",
] { ] {
assert!( assert!(
!selection.contains(pattern), !selection.contains(pattern),
@@ -558,8 +568,8 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
let planner_candidate_affinity = let planner_candidate_affinity =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs"); read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs");
assert!( assert!(
planner_candidate_affinity planner_candidate_affinity.contains("use aether_scheduler_core::{")
.contains("aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate"), && planner_candidate_affinity.contains("SchedulerMinimalCandidateSelectionCandidate"),
"planner/candidate_affinity.rs should depend directly on core minimal candidate DTO" "planner/candidate_affinity.rs should depend directly on core minimal candidate DTO"
); );
assert!( assert!(

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,124 @@
use super::types::SchedulerMinimalCandidateSelectionCandidate;
#[derive(Debug, Clone, Copy)]
pub(crate) struct RequiredCapabilityDescriptor<'a> {
pub(crate) name: &'a str,
pub(crate) compatible: bool,
}
pub fn candidate_supports_required_capability(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
required_capability: &str,
) -> bool {
let required_capability = required_capability.trim();
if required_capability.is_empty() {
return true;
}
let Some(capabilities) = candidate.key_capabilities.as_ref() else {
return false;
};
if let Some(object) = capabilities.as_object() {
return object.iter().any(|(key, value)| {
key.eq_ignore_ascii_case(required_capability)
&& match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => {
value.as_i64().is_some_and(|value| value > 0)
}
_ => false,
}
});
}
if let Some(items) = capabilities.as_array() {
return items.iter().any(|value| {
value
.as_str()
.is_some_and(|value| value.eq_ignore_ascii_case(required_capability))
});
}
false
}
pub fn requested_capability_priority_for_candidate(
required_capabilities: Option<&serde_json::Value>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (u32, u32) {
let Some(required_capabilities) = required_capabilities.and_then(serde_json::Value::as_object)
else {
return (0, 0);
};
requested_capability_priority_for_candidate_descriptors(
required_capabilities
.iter()
.filter_map(|(capability, value)| {
requested_capability_is_enabled(value).then_some(RequiredCapabilityDescriptor {
name: capability.as_str(),
compatible: requested_capability_is_compatible(capability),
})
}),
candidate,
)
}
pub(crate) fn enabled_required_capabilities(
required_capabilities: Option<&serde_json::Value>,
) -> Vec<RequiredCapabilityDescriptor<'_>> {
let Some(required_capabilities) = required_capabilities.and_then(serde_json::Value::as_object)
else {
return Vec::new();
};
required_capabilities
.iter()
.filter_map(|(capability, value)| {
requested_capability_is_enabled(value).then_some(RequiredCapabilityDescriptor {
name: capability.as_str(),
compatible: requested_capability_is_compatible(capability),
})
})
.collect()
}
pub(crate) fn requested_capability_priority_for_candidate_descriptors<'a, I>(
required_capabilities: I,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (u32, u32)
where
I: IntoIterator<Item = RequiredCapabilityDescriptor<'a>>,
{
let mut exclusive_misses = 0u32;
let mut compatible_misses = 0u32;
for capability in required_capabilities {
if candidate_supports_required_capability(candidate, capability.name) {
continue;
}
if capability.compatible {
compatible_misses += 1;
} else {
exclusive_misses += 1;
}
}
(exclusive_misses, compatible_misses)
}
fn requested_capability_is_enabled(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => value.as_i64().is_some_and(|value| value > 0),
_ => false,
}
}
fn requested_capability_is_compatible(capability: &str) -> bool {
matches!(
capability.trim().to_ascii_lowercase().as_str(),
"cache_1h" | "context_1m"
)
}

View File

@@ -0,0 +1,163 @@
use std::collections::BTreeSet;
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data_contracts::DataLayerError;
use super::capability::{
enabled_required_capabilities, requested_capability_priority_for_candidate_descriptors,
};
use super::types::{
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
};
pub fn build_minimal_candidate_selection(
input: BuildMinimalCandidateSelectionInput<'_>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let priority_mode = input.priority_mode;
let affinity_key = input.affinity_key.map(str::to_string);
let required_capabilities = enabled_required_capabilities(input.required_capabilities);
let mut candidates = enumerate_minimal_candidate_selection(input)?;
let rankables = candidates
.iter()
.enumerate()
.map(|(index, candidate)| {
crate::SchedulerRankableCandidate::from_candidate(candidate, index)
.with_capability_priority(requested_capability_priority_for_candidate_descriptors(
required_capabilities.iter().copied(),
candidate,
))
.with_affinity_hash(
affinity_key
.as_deref()
.map(|key| crate::candidate_affinity_hash(key, candidate)),
)
})
.collect::<Vec<_>>();
crate::apply_scheduler_candidate_ranking(
&mut candidates,
&rankables,
crate::SchedulerRankingContext {
priority_mode,
ranking_mode: crate::SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
},
);
Ok(candidates)
}
pub fn enumerate_minimal_candidate_selection(
input: BuildMinimalCandidateSelectionInput<'_>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let BuildMinimalCandidateSelectionInput {
rows,
normalized_api_format,
requested_model_name,
resolved_global_model_name,
require_streaming,
auth_constraints,
..
} = input;
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !crate::auth_constraints_allow_api_format(auth_constraints, normalized_api_format) {
return Ok(Vec::new());
}
if !crate::auth_constraints_allow_model(
auth_constraints,
requested_model_name,
resolved_global_model_name,
) {
return Ok(Vec::new());
}
let mut candidates = Vec::with_capacity(rows.len());
for row in rows {
if !crate::auth_constraints_allow_provider(
auth_constraints,
&row.provider_id,
&row.provider_name,
&row.provider_type,
) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
let Some((selected_provider_model_name, mapping_matched_model)) =
crate::resolve_provider_model_name(&row, requested_model_name, normalized_api_format)
else {
continue;
};
candidates.push(SchedulerMinimalCandidateSelectionCandidate {
provider_id: row.provider_id,
provider_name: row.provider_name,
provider_type: row.provider_type,
provider_priority: row.provider_priority,
endpoint_id: row.endpoint_id,
endpoint_api_format: row.endpoint_api_format,
key_id: row.key_id,
key_name: row.key_name,
key_auth_type: row.key_auth_type,
key_internal_priority: row.key_internal_priority,
key_global_priority_for_format: crate::extract_global_priority_for_format(
row.key_global_priority_by_format.as_ref(),
normalized_api_format,
)?,
key_capabilities: row.key_capabilities,
model_id: row.model_id,
global_model_id: row.global_model_id,
global_model_name: row.global_model_name,
selected_provider_model_name,
mapping_matched_model,
});
}
Ok(candidates)
}
pub fn collect_global_model_names_for_required_capability(
rows: Vec<StoredMinimalCandidateSelectionRow>,
normalized_api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_constraints: Option<&crate::SchedulerAuthConstraints>,
) -> Vec<String> {
if normalized_api_format.is_empty() || required_capability.trim().is_empty() {
return Vec::new();
}
if !crate::auth_constraints_allow_api_format(auth_constraints, normalized_api_format) {
return Vec::new();
}
let mut model_names = BTreeSet::new();
for row in rows {
if !crate::auth_constraints_allow_provider(
auth_constraints,
&row.provider_id,
&row.provider_name,
&row.provider_type,
) {
continue;
}
if !crate::row_supports_required_capability(&row, required_capability) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
if !crate::auth_constraints_allow_model(
auth_constraints,
&row.global_model_name,
&row.global_model_name,
) {
continue;
}
model_names.insert(row.global_model_name);
}
model_names.into_iter().collect()
}

View File

@@ -0,0 +1,39 @@
use super::types::{SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode};
pub fn compare_candidates_by_priority_mode(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
affinity_key: Option<&str>,
) -> std::cmp::Ordering {
match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then_with(|| compare_candidate_identity(left, right)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| compare_candidate_identity(left, right)),
}
}
pub(crate) fn compare_candidate_identity(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
) -> std::cmp::Ordering {
left.provider_id
.cmp(&right.provider_id)
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
}

View File

@@ -0,0 +1,482 @@
pub mod capability;
pub mod enumeration;
pub mod identity;
pub mod selectability;
pub mod types;
pub use capability::{
candidate_supports_required_capability, requested_capability_priority_for_candidate,
};
pub use enumeration::{
build_minimal_candidate_selection, collect_global_model_names_for_required_capability,
enumerate_minimal_candidate_selection,
};
pub use identity::compare_candidates_by_priority_mode;
pub use selectability::{
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
candidate_runtime_skip_reason_with_state, collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health, CandidateRuntimeSelectabilityInput,
};
pub use types::{
BuildMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use super::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability,
collect_selectable_candidates_from_keys, reorder_candidates_by_scheduler_health,
BuildMinimalCandidateSelectionInput, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
};
use crate::SchedulerAuthConstraints;
fn sample_row(id: &str) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{id}"),
provider_name: format!("Provider {id}"),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{id}"),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: format!("key-{id}"),
key_name: format!("prod-{id}"),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: Some(serde_json::json!({"cache_1h": true})),
key_internal_priority: 50,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 2})),
model_id: format!("model-{id}"),
global_model_id: format!("global-model-{id}"),
global_model_name: "gpt-5".to_string(),
global_model_mappings: Some(vec!["gpt-5(?:\\.\\d+)?".to_string()]),
global_model_supports_streaming: Some(true),
model_provider_model_name: format!("gpt-5-upstream-{id}"),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: format!("gpt-5-canary-{id}"),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]),
model_supports_streaming: None,
model_is_active: true,
model_is_available: true,
}
}
fn sample_candidate(
id: &str,
capabilities: Option<serde_json::Value>,
) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
provider_id: format!("provider-{id}"),
provider_name: format!("Provider {id}"),
provider_type: "openai".to_string(),
provider_priority: 0,
endpoint_id: format!("endpoint-{id}"),
endpoint_api_format: "openai:chat".to_string(),
key_id: format!("key-{id}"),
key_name: format!("key-{id}"),
key_auth_type: "bearer".to_string(),
key_internal_priority: 0,
key_global_priority_for_format: None,
key_capabilities: capabilities,
model_id: format!("model-{id}"),
global_model_id: format!("global-model-{id}"),
global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
mapping_matched_model: None,
}
}
fn sample_key(id: &str, health_score: f64) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
format!("key-{id}"),
format!("provider-{id}"),
format!("key-{id}"),
"api_key".to_string(),
None,
true,
)
.expect("provider key should build");
key.health_by_format = Some(serde_json::json!({
"openai:chat": {
"health_score": health_score
}
}));
key
}
fn stored_candidate(
id: &str,
status: RequestCandidateStatus,
created_at_unix_ms: i64,
) -> StoredRequestCandidate {
let finished_at_unix_ms = match status {
RequestCandidateStatus::Pending | RequestCandidateStatus::Streaming => None,
_ => Some(created_at_unix_ms),
};
StoredRequestCandidate::new(
id.to_string(),
format!("req-{id}"),
None,
None,
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
status,
None,
false,
None,
None,
None,
None,
None,
None,
None,
created_at_unix_ms,
Some(created_at_unix_ms),
finished_at_unix_ms,
)
.expect("candidate should build")
}
#[test]
fn reads_required_capability_from_object_and_array_forms() {
assert!(candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!({"vision": true}))),
"vision"
));
assert!(candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!(["vision", "tools"]))),
"tools"
));
assert!(!candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!({"vision": false}))),
"vision"
));
}
#[test]
fn builds_minimal_candidate_selection_with_auth_constraints() {
let mut disallowed = sample_row("2");
disallowed.provider_id = "provider-blocked".to_string();
disallowed.provider_name = "Blocked".to_string();
let constraints = SchedulerAuthConstraints {
allowed_providers: Some(vec!["provider-1".to_string()]),
allowed_api_formats: Some(vec!["OPENAI:CHAT".to_string()]),
allowed_models: Some(vec!["gpt-5".to_string()]),
};
let candidates = build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows: vec![sample_row("1"), disallowed],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
resolved_global_model_name: "gpt-5",
require_streaming: false,
required_capabilities: None,
auth_constraints: Some(&constraints),
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate selection should build");
assert_eq!(candidates.len(), 1);
assert_eq!(candidates[0].provider_id, "provider-1");
assert_eq!(candidates[0].selected_provider_model_name, "gpt-5-canary-1");
}
#[test]
fn enumeration_preserves_theoretical_candidate_order_without_final_sorting() {
let mut later_priority = sample_row("1");
later_priority.provider_priority = 10;
let mut earlier_priority = sample_row("2");
earlier_priority.provider_priority = 0;
let candidates =
super::enumerate_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows: vec![later_priority, earlier_priority],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
resolved_global_model_name: "gpt-5",
require_streaming: false,
required_capabilities: None,
auth_constraints: None,
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate enumeration should build");
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].provider_id, "provider-1");
assert_eq!(candidates[1].provider_id, "provider-2");
}
#[test]
fn collects_global_model_names_for_required_capability_with_auth_constraints() {
let mut disallowed = sample_row("2");
disallowed.global_model_name = "gpt-4.1".to_string();
disallowed.provider_id = "provider-blocked".to_string();
disallowed.provider_name = "Blocked".to_string();
let constraints = SchedulerAuthConstraints {
allowed_providers: Some(vec!["provider-1".to_string()]),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_models: Some(vec!["gpt-5".to_string()]),
};
let model_names = collect_global_model_names_for_required_capability(
vec![sample_row("1"), disallowed],
"openai:chat",
"cache_1h",
false,
Some(&constraints),
);
assert_eq!(model_names, vec!["gpt-5".to_string()]);
}
#[test]
fn minimal_candidate_selection_prefers_matching_requested_capabilities_before_priority() {
let mut missing_capability = sample_row("1");
missing_capability.key_capabilities = Some(serde_json::json!({"cache_1h": false}));
missing_capability.provider_priority = 0;
let mut matching_capability = sample_row("2");
matching_capability.key_capabilities = Some(serde_json::json!({"cache_1h": true}));
matching_capability.provider_priority = 10;
let required_capabilities = serde_json::json!({"cache_1h": true});
let candidates = build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows: vec![missing_capability, matching_capability],
normalized_api_format: "openai:chat",
requested_model_name: "gpt-5",
resolved_global_model_name: "gpt-5",
require_streaming: false,
required_capabilities: Some(&required_capabilities),
auth_constraints: None,
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
.expect("candidate selection should build");
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].key_id, "key-2");
assert_eq!(candidates[1].key_id, "key-1");
}
#[test]
fn reorders_candidates_by_health_before_affinity_tiebreak() {
let mut candidates = vec![
sample_candidate("1", None),
sample_candidate("2", None),
sample_candidate("3", None),
];
let provider_key_rpm_states = BTreeMap::from([
("key-1".to_string(), sample_key("1", 0.95)),
("key-2".to_string(), sample_key("2", 0.40)),
("key-3".to_string(), sample_key("3", 0.95)),
]);
reorder_candidates_by_scheduler_health(
&mut candidates,
&provider_key_rpm_states,
None,
Some("api-key-1"),
SchedulerPriorityMode::GlobalKey,
);
assert_ne!(candidates[0].key_id, "key-2");
assert_ne!(candidates[1].key_id, "key-2");
assert_eq!(candidates[2].key_id, "key-2");
}
#[test]
fn collects_selectable_candidates_with_affinity_priority_and_dedup() {
let candidates = vec![
sample_candidate("1", None),
sample_candidate("2", None),
sample_candidate("1", None),
];
let selectable_keys = BTreeSet::from([
(
"provider-1".to_string(),
"endpoint-1".to_string(),
"key-1".to_string(),
),
(
"provider-2".to_string(),
"endpoint-2".to_string(),
"key-2".to_string(),
),
]);
let selected = collect_selectable_candidates_from_keys(
candidates,
&selectable_keys,
Some(&crate::SchedulerAffinityTarget {
provider_id: "provider-2".to_string(),
endpoint_id: "endpoint-2".to_string(),
key_id: "key-2".to_string(),
}),
);
assert_eq!(selected.len(), 2);
assert_eq!(selected[0].key_id, "key-2");
assert_eq!(selected[1].key_id, "key-1");
}
#[test]
fn candidate_selectability_respects_provider_concurrency_limit() {
let recent_candidates = vec![stored_candidate("one", RequestCandidateStatus::Pending, 95)];
let provider_concurrent_limits = BTreeMap::from([("provider-1".to_string(), 1)]);
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &recent_candidates,
provider_concurrent_limits: &provider_concurrent_limits,
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
},
));
}
#[test]
fn candidate_selectability_rejects_quota_or_zero_health() {
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), sample_key("1", 0.0))]);
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &provider_key_rpm_states,
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
},
));
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: true,
account_quota_exhausted: false,
oauth_invalid: false,
rpm_reset_at: None,
},
));
}
#[test]
fn candidate_selectability_rejects_exhausted_account_quota() {
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: true,
oauth_invalid: false,
rpm_reset_at: None,
},
));
}
#[test]
fn candidate_selectability_rejects_oauth_invalid_keys() {
assert!(!candidate_is_selectable_with_runtime_state(
CandidateRuntimeSelectabilityInput {
candidate: &sample_candidate("1", None),
recent_candidates: &[],
provider_concurrent_limits: &BTreeMap::new(),
provider_key_rpm_states: &BTreeMap::new(),
now_unix_secs: 100,
cached_affinity_target: None,
provider_quota_blocks_requests: false,
account_quota_exhausted: false,
oauth_invalid: true,
rpm_reset_at: None,
},
));
}
#[test]
fn detects_auth_api_key_concurrency_limit_from_recent_active_requests() {
let recent_candidates = vec![StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95,
Some(95),
None,
)
.expect("candidate should build")];
assert!(auth_api_key_concurrency_limit_reached(
&recent_candidates,
100,
"api-key-1",
1,
));
assert!(!auth_api_key_concurrency_limit_reached(
&recent_candidates,
100,
"api-key-1",
2,
));
}
}

View File

@@ -0,0 +1,204 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::candidates::StoredRequestCandidate;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use super::capability::{
enabled_required_capabilities, requested_capability_priority_for_candidate_descriptors,
};
use super::types::{SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode};
pub fn auth_api_key_concurrency_limit_reached(
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
api_key_id: &str,
concurrent_limit: usize,
) -> bool {
if api_key_id.trim().is_empty() || concurrent_limit == 0 {
return false;
}
crate::count_recent_active_requests_for_api_key(recent_candidates, api_key_id, now_unix_secs)
>= concurrent_limit
}
pub fn collect_selectable_candidates_from_keys(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
selectable_keys: &BTreeSet<(String, String, String)>,
cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
let mut promoted = None;
let mut selected = Vec::with_capacity(candidates.len());
let mut emitted_keys = BTreeSet::new();
for candidate in candidates {
let key = crate::candidate_key(&candidate);
if !selectable_keys.contains(&key) || !emitted_keys.insert(key) {
continue;
}
if promoted.is_none()
&& cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(&candidate, target))
{
promoted = Some(candidate);
} else {
selected.push(candidate);
}
}
if let Some(candidate) = promoted {
selected.insert(0, candidate);
}
selected
}
pub fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
required_capabilities: Option<&serde_json::Value>,
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) {
let required_capabilities = enabled_required_capabilities(required_capabilities);
let rankables = candidates
.iter()
.enumerate()
.map(|(index, candidate)| {
crate::SchedulerRankableCandidate::from_candidate(candidate, index)
.with_capability_priority(requested_capability_priority_for_candidate_descriptors(
required_capabilities.iter().copied(),
candidate,
))
.with_affinity_hash(
affinity_key.map(|key| crate::candidate_affinity_hash(key, candidate)),
)
.with_health(
provider_key_rpm_states
.get(&candidate.key_id)
.and_then(|key| {
crate::provider_key_health_bucket(
key,
candidate.endpoint_api_format.as_str(),
)
}),
candidate_provider_key_health_score(candidate, Some(provider_key_rpm_states)),
)
})
.collect::<Vec<_>>();
crate::apply_scheduler_candidate_ranking(
candidates,
&rankables,
crate::SchedulerRankingContext {
priority_mode,
ranking_mode: crate::SchedulerRankingMode::CacheAffinity,
include_health: true,
load_balance_seed: 0,
},
);
}
#[derive(Clone, Copy, Debug)]
pub struct CandidateRuntimeSelectabilityInput<'a> {
pub candidate: &'a SchedulerMinimalCandidateSelectionCandidate,
pub recent_candidates: &'a [StoredRequestCandidate],
pub provider_concurrent_limits: &'a BTreeMap<String, usize>,
pub provider_key_rpm_states: &'a BTreeMap<String, StoredProviderCatalogKey>,
pub now_unix_secs: u64,
pub cached_affinity_target: Option<&'a crate::SchedulerAffinityTarget>,
pub provider_quota_blocks_requests: bool,
pub account_quota_exhausted: bool,
pub oauth_invalid: bool,
pub rpm_reset_at: Option<u64>,
}
pub fn candidate_is_selectable_with_runtime_state(
input: CandidateRuntimeSelectabilityInput<'_>,
) -> bool {
candidate_runtime_skip_reason_with_state(input).is_none()
}
pub fn candidate_runtime_skip_reason_with_state(
input: CandidateRuntimeSelectabilityInput<'_>,
) -> Option<&'static str> {
let CandidateRuntimeSelectabilityInput {
candidate,
recent_candidates,
provider_concurrent_limits,
provider_key_rpm_states,
now_unix_secs,
cached_affinity_target,
provider_quota_blocks_requests,
account_quota_exhausted,
oauth_invalid,
rpm_reset_at,
} = input;
if provider_quota_blocks_requests {
return Some("provider_quota_blocked");
}
if account_quota_exhausted {
return Some("account_quota_exhausted");
}
if oauth_invalid {
return Some("oauth_invalid");
}
if crate::is_candidate_in_recent_failure_cooldown(
recent_candidates,
candidate.provider_id.as_str(),
candidate.endpoint_id.as_str(),
candidate.key_id.as_str(),
now_unix_secs,
) {
return Some("recent_failure_cooldown");
}
if provider_concurrent_limits
.get(&candidate.provider_id)
.is_some_and(|limit| {
crate::count_recent_active_requests_for_provider(
recent_candidates,
candidate.provider_id.as_str(),
now_unix_secs,
) >= *limit
})
{
return Some("provider_concurrency_limit_reached");
}
let is_cached_user = cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(candidate, target));
if let Some(provider_key) = provider_key_rpm_states.get(&candidate.key_id) {
if crate::is_provider_key_circuit_open(provider_key, candidate.endpoint_api_format.as_str())
{
return Some("key_circuit_open");
}
if crate::provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())
.is_some_and(|score| score <= 0.0)
{
return Some("key_health_score_zero");
}
if !crate::provider_key_rpm_allows_request_since(
provider_key,
recent_candidates,
now_unix_secs,
is_cached_user,
rpm_reset_at,
) {
return Some("key_rpm_exhausted");
}
}
None
}
fn candidate_provider_key_health_score(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: Option<&BTreeMap<String, StoredProviderCatalogKey>>,
) -> f64 {
provider_key_rpm_states
.and_then(|states| states.get(&candidate.key_id))
.and_then(|key| {
crate::effective_provider_key_health_score(key, candidate.endpoint_api_format.as_str())
})
.unwrap_or(1.0)
}

View File

@@ -0,0 +1,41 @@
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub enum SchedulerPriorityMode {
#[default]
Provider,
GlobalKey,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct SchedulerMinimalCandidateSelectionCandidate {
pub provider_id: String,
pub provider_name: String,
pub provider_type: String,
pub provider_priority: i32,
pub endpoint_id: String,
pub endpoint_api_format: String,
pub key_id: String,
pub key_name: String,
pub key_auth_type: String,
pub key_internal_priority: i32,
pub key_global_priority_for_format: Option<i32>,
pub key_capabilities: Option<serde_json::Value>,
pub model_id: String,
pub global_model_id: String,
pub global_model_name: String,
pub selected_provider_model_name: String,
pub mapping_matched_model: Option<String>,
}
pub struct BuildMinimalCandidateSelectionInput<'a> {
pub rows: Vec<StoredMinimalCandidateSelectionRow>,
pub normalized_api_format: &'a str,
pub requested_model_name: &'a str,
pub resolved_global_model_name: &'a str,
pub require_streaming: bool,
pub required_capabilities: Option<&'a serde_json::Value>,
pub auth_constraints: Option<&'a crate::SchedulerAuthConstraints>,
pub affinity_key: Option<&'a str>,
pub priority_mode: SchedulerPriorityMode,
}

View File

@@ -4,6 +4,7 @@ mod candidate;
mod health; mod health;
mod model; mod model;
mod provider; mod provider;
mod ranking;
mod request_candidate; mod request_candidate;
pub use affinity::{ pub use affinity::{
@@ -19,9 +20,10 @@ pub use candidate::{
candidate_is_selectable_with_runtime_state, candidate_runtime_skip_reason_with_state, candidate_is_selectable_with_runtime_state, candidate_runtime_skip_reason_with_state,
candidate_supports_required_capability, collect_global_model_names_for_required_capability, candidate_supports_required_capability, collect_global_model_names_for_required_capability,
collect_selectable_candidates_from_keys, compare_candidates_by_priority_mode, collect_selectable_candidates_from_keys, compare_candidates_by_priority_mode,
reorder_candidates_by_scheduler_health, requested_capability_priority_for_candidate, enumerate_minimal_candidate_selection, reorder_candidates_by_scheduler_health,
BuildMinimalCandidateSelectionInput, CandidateRuntimeSelectabilityInput, requested_capability_priority_for_candidate, BuildMinimalCandidateSelectionInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
}; };
pub use health::{ pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key, aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
@@ -35,9 +37,17 @@ pub use health::{
pub use model::{ pub use model::{
candidate_model_names, extract_global_priority_for_format, matches_model_mapping, candidate_model_names, extract_global_priority_for_format, matches_model_mapping,
normalize_api_format, resolve_provider_model_name, resolve_requested_global_model_name, normalize_api_format, resolve_provider_model_name, resolve_requested_global_model_name,
row_supports_required_capability, select_provider_model_name, row_supports_requested_model, row_supports_required_capability, select_provider_model_name,
}; };
pub use provider::{build_provider_concurrent_limit_map, should_skip_provider_quota}; pub use provider::{build_provider_concurrent_limit_map, should_skip_provider_quota};
pub use ranking::{
apply_scheduler_candidate_ranking, candidate_priority_slot, candidates_share_priority_group,
compare_candidate_identity_for_ranking, compare_candidate_priority_slot,
scheduler_candidate_ranking_order, scheduler_ranking_outcomes, SchedulerRankableCandidate,
SchedulerRankingContext, SchedulerRankingMode, SchedulerRankingOutcome,
SchedulerTunnelAffinityBucket, RANKING_REASON_CACHED_AFFINITY, RANKING_REASON_CROSS_FORMAT,
RANKING_REASON_LOCAL_TUNNEL,
};
pub use request_candidate::{ pub use request_candidate::{
build_execution_request_candidate_seed, build_local_request_candidate_status_record, build_execution_request_candidate_seed, build_local_request_candidate_status_record,
build_report_request_candidate_status_record, execution_error_details, build_report_request_candidate_status_record, execution_error_details,

View File

@@ -11,30 +11,56 @@ pub fn resolve_requested_global_model_name(
requested_model_name: &str, requested_model_name: &str,
api_format: &str, api_format: &str,
) -> Option<String> { ) -> Option<String> {
resolve_global_model_name_by(rows, |row| { resolve_global_model_name_by(rows, |row| row.global_model_name == requested_model_name)
row.model_provider_model_name == requested_model_name .or_else(|| {
}) resolve_global_model_name_by(rows, |row| {
.or_else(|| { row.model_provider_model_name == requested_model_name
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
}) })
}) })
}) .or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
})
})
}
pub fn row_supports_requested_model(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
) -> bool {
row.global_model_name == requested_model_name
|| row.model_provider_model_name == requested_model_name
|| row
.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
|| row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
} }
fn resolve_global_model_name_by<F>( fn resolve_global_model_name_by<F>(

View File

@@ -0,0 +1,12 @@
use std::cmp::Ordering;
use super::types::SchedulerRankableCandidate;
pub fn compare_format_state(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
) -> Ordering {
left.demote_cross_format
.cmp(&right.demote_cross_format)
.then(left.format_preference.cmp(&right.format_preference))
}

View File

@@ -0,0 +1,294 @@
mod format;
mod modes;
mod priority;
mod reasons;
mod types;
pub use format::compare_format_state;
pub use modes::{apply_load_balance_rotation, compare_rankable_candidates};
pub use priority::{
candidate_priority_slot, candidates_share_priority_group, compare_candidate_priority_slot,
};
pub use reasons::{
demoted_by as ranking_demoted_by, promoted_by as ranking_promoted_by,
RANKING_REASON_CACHED_AFFINITY, RANKING_REASON_CROSS_FORMAT, RANKING_REASON_LOCAL_TUNNEL,
};
pub use types::{
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
SchedulerRankingOutcome, SchedulerTunnelAffinityBucket,
};
pub fn compare_candidate_identity_for_ranking(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
) -> std::cmp::Ordering {
left.provider_id
.cmp(&right.provider_id)
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
}
pub fn scheduler_candidate_ranking_order(
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<usize> {
let mut order = (0..candidates.len()).collect::<Vec<_>>();
order.sort_by(|left, right| {
compare_rankable_candidates(&candidates[*left], &candidates[*right], context)
});
apply_load_balance_rotation(&mut order, candidates, context);
order
}
pub fn scheduler_ranking_outcomes(
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<SchedulerRankingOutcome> {
scheduler_candidate_ranking_order(candidates, context)
.into_iter()
.enumerate()
.map(|(ranking_index, original_index)| {
let candidate = &candidates[original_index];
SchedulerRankingOutcome {
original_index,
ranking_index,
priority_mode: context.priority_mode,
ranking_mode: context.ranking_mode,
priority_slot: candidate_priority_slot(candidate, context.priority_mode),
promoted_by: ranking_promoted_by(candidate, context.ranking_mode),
demoted_by: ranking_demoted_by(candidate),
}
})
.collect()
}
pub fn apply_scheduler_candidate_ranking<T>(
items: &mut [T],
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<SchedulerRankingOutcome> {
let outcomes = scheduler_ranking_outcomes(candidates, context);
apply_order(
items,
outcomes
.iter()
.map(|outcome| outcome.original_index)
.collect(),
);
outcomes
}
fn apply_order<T>(items: &mut [T], sorted_old_indices: Vec<usize>) {
if items.len() < 2 {
return;
}
let mut target_positions = vec![0usize; sorted_old_indices.len()];
for (new_position, old_position) in sorted_old_indices.into_iter().enumerate() {
target_positions[old_position] = new_position;
}
for index in 0..items.len() {
while target_positions[index] != index {
let target = target_positions[index];
items.swap(index, target);
target_positions.swap(index, target);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{SchedulerPriorityMode, SchedulerTunnelAffinityBucket};
fn candidate(
id: &str,
provider_priority: i32,
key_priority: i32,
global_key_priority: Option<i32>,
) -> SchedulerRankableCandidate {
SchedulerRankableCandidate {
provider_id: format!("provider-{id}"),
endpoint_id: format!("endpoint-{id}"),
key_id: format!("key-{id}"),
selected_provider_model_name: "gpt-5".to_string(),
provider_priority,
key_internal_priority: key_priority,
key_global_priority_for_format: global_key_priority,
capability_priority: (0, 0),
cached_affinity_match: false,
affinity_hash: None,
tunnel_bucket: SchedulerTunnelAffinityBucket::Neutral,
demote_cross_format: false,
format_preference: (0, 0),
health_bucket: None,
health_score: 1.0,
original_index: 0,
}
}
fn ranked_ids(
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) -> Vec<String> {
scheduler_candidate_ranking_order(candidates, context)
.into_iter()
.map(|index| candidates[index].provider_id.clone())
.collect()
}
#[test]
fn provider_priority_mode_prefers_provider_priority_slot() {
let candidates = vec![
candidate("global", 10, 0, Some(0)),
candidate("provider", 0, 10, Some(10)),
];
assert_eq!(
ranked_ids(
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-provider", "provider-global"]
);
}
#[test]
fn global_key_priority_mode_prefers_global_key_priority_slot() {
let candidates = vec![
candidate("provider", 0, 10, Some(10)),
candidate("global", 10, 0, Some(0)),
];
assert_eq!(
ranked_ids(
&candidates,
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::GlobalKey,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-global", "provider-provider"]
);
}
#[test]
fn fixed_order_keeps_priority_before_affinity_tunnel_and_format_preference() {
let mut lower_priority = candidate("lower", 10, 0, Some(10));
lower_priority.cached_affinity_match = true;
lower_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::LocalTunnel;
lower_priority.format_preference = (0, 0);
let mut higher_priority = candidate("higher", 0, 0, Some(0));
higher_priority.demote_cross_format = true;
higher_priority.format_preference = (9, 9);
assert_eq!(
ranked_ids(
&[lower_priority, higher_priority],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::FixedOrder,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-higher", "provider-lower"]
);
}
#[test]
fn cache_affinity_can_promote_cached_candidate_and_reports_reason() {
let high_priority = candidate("high", 0, 0, Some(0));
let mut cached = candidate("cached", 10, 0, Some(10));
cached.cached_affinity_match = true;
let candidates = vec![high_priority, cached];
let context = SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
};
let outcomes = scheduler_ranking_outcomes(&candidates, context);
assert_eq!(outcomes[0].original_index, 1);
assert_eq!(
outcomes[0].promoted_by,
Some(RANKING_REASON_CACHED_AFFINITY)
);
}
#[test]
fn cache_affinity_without_cache_hit_keeps_priority_before_tunnel() {
let mut higher_priority = candidate("higher", 0, 0, Some(0));
higher_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::RemoteTunnel;
let mut lower_priority = candidate("lower", 10, 0, Some(10));
lower_priority.tunnel_bucket = SchedulerTunnelAffinityBucket::LocalTunnel;
assert_eq!(
ranked_ids(
&[lower_priority, higher_priority],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-higher", "provider-lower"]
);
}
#[test]
fn cache_affinity_keeps_cross_format_demotion_before_priority() {
let same_format_low_priority = candidate("same", 10, 0, Some(10));
let mut cross_format_high_priority = candidate("cross", 0, 0, Some(0));
cross_format_high_priority.demote_cross_format = true;
assert_eq!(
ranked_ids(
&[cross_format_high_priority, same_format_low_priority],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
},
),
vec!["provider-same", "provider-cross"]
);
}
#[test]
fn load_balance_rotates_only_within_same_priority_group() {
let first = candidate("first", 0, 0, Some(0));
let second = candidate("second", 0, 0, Some(0));
let third = candidate("third", 10, 0, Some(10));
assert_eq!(
ranked_ids(
&[first, second, third],
SchedulerRankingContext {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::LoadBalance,
include_health: false,
load_balance_seed: 1,
},
),
vec!["provider-second", "provider-first", "provider-third"]
);
}
}

View File

@@ -0,0 +1,111 @@
use std::cmp::Ordering;
use super::types::{SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode};
use super::{
candidates_share_priority_group, compare_candidate_identity_for_ranking,
compare_candidate_priority_slot, compare_format_state,
};
pub fn compare_rankable_candidates(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
match context.ranking_mode {
SchedulerRankingMode::FixedOrder => compare_fixed_order(left, right, context),
SchedulerRankingMode::CacheAffinity => compare_cache_affinity(left, right, context),
SchedulerRankingMode::LoadBalance => compare_load_balance_base(left, right, context),
}
}
fn compare_fixed_order(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then_with(|| compare_format_state(left, right))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
fn compare_cache_affinity(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then_with(|| right.cached_affinity_match.cmp(&left.cached_affinity_match))
.then(left.demote_cross_format.cmp(&right.demote_cross_format))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then(left.tunnel_bucket.cmp(&right.tunnel_bucket))
.then(left.format_preference.cmp(&right.format_preference))
.then_with(|| compare_health(left, right, context.include_health))
.then(left.affinity_hash.cmp(&right.affinity_hash))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
fn compare_load_balance_base(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
context: SchedulerRankingContext,
) -> Ordering {
left.capability_priority
.cmp(&right.capability_priority)
.then(left.demote_cross_format.cmp(&right.demote_cross_format))
.then_with(|| compare_candidate_priority_slot(left, right, context.priority_mode))
.then(left.format_preference.cmp(&right.format_preference))
.then_with(|| compare_health(left, right, context.include_health))
.then(left.affinity_hash.cmp(&right.affinity_hash))
.then_with(|| compare_candidate_identity_for_ranking(left, right))
.then(left.original_index.cmp(&right.original_index))
}
fn compare_health(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
include_health: bool,
) -> Ordering {
if !include_health {
return Ordering::Equal;
}
right
.health_bucket
.cmp(&left.health_bucket)
.then_with(|| right.health_score.total_cmp(&left.health_score))
}
pub fn apply_load_balance_rotation(
sorted_indices: &mut [usize],
candidates: &[SchedulerRankableCandidate],
context: SchedulerRankingContext,
) {
if context.ranking_mode != SchedulerRankingMode::LoadBalance || sorted_indices.len() < 2 {
return;
}
let mut start = 0usize;
while start < sorted_indices.len() {
let mut end = start + 1;
while end < sorted_indices.len()
&& candidates_share_priority_group(
&candidates[sorted_indices[start]],
&candidates[sorted_indices[end]],
context.priority_mode,
)
{
end += 1;
}
let group_len = end - start;
if group_len > 1 {
let offset = usize::try_from(context.load_balance_seed).unwrap_or(0) % group_len;
sorted_indices[start..end].rotate_left(offset);
}
start = end;
}
}

View File

@@ -0,0 +1,52 @@
use std::cmp::Ordering;
use crate::SchedulerPriorityMode;
use super::types::SchedulerRankableCandidate;
pub fn candidate_priority_slot(
candidate: &SchedulerRankableCandidate,
priority_mode: SchedulerPriorityMode,
) -> i32 {
match priority_mode {
SchedulerPriorityMode::Provider => candidate.provider_priority,
SchedulerPriorityMode::GlobalKey => {
candidate.key_global_priority_for_format.unwrap_or(i32::MAX)
}
}
}
pub fn compare_candidate_priority_slot(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
priority_mode: SchedulerPriorityMode,
) -> Ordering {
match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority)),
}
}
pub fn candidates_share_priority_group(
left: &SchedulerRankableCandidate,
right: &SchedulerRankableCandidate,
priority_mode: SchedulerPriorityMode,
) -> bool {
match priority_mode {
SchedulerPriorityMode::Provider => {
left.provider_priority == right.provider_priority
&& left.key_internal_priority == right.key_internal_priority
}
SchedulerPriorityMode::GlobalKey => {
left.key_global_priority_for_format == right.key_global_priority_for_format
}
}
}

View File

@@ -0,0 +1,28 @@
use super::types::{
SchedulerRankableCandidate, SchedulerRankingMode, SchedulerTunnelAffinityBucket,
};
pub const RANKING_REASON_CACHED_AFFINITY: &str = "cached_affinity";
pub const RANKING_REASON_LOCAL_TUNNEL: &str = "local_tunnel";
pub const RANKING_REASON_CROSS_FORMAT: &str = "cross_format";
pub fn promoted_by(
candidate: &SchedulerRankableCandidate,
ranking_mode: SchedulerRankingMode,
) -> Option<&'static str> {
if ranking_mode == SchedulerRankingMode::CacheAffinity && candidate.cached_affinity_match {
return Some(RANKING_REASON_CACHED_AFFINITY);
}
if ranking_mode == SchedulerRankingMode::CacheAffinity
&& candidate.tunnel_bucket == SchedulerTunnelAffinityBucket::LocalTunnel
{
return Some(RANKING_REASON_LOCAL_TUNNEL);
}
None
}
pub fn demoted_by(candidate: &SchedulerRankableCandidate) -> Option<&'static str> {
candidate
.demote_cross_format
.then_some(RANKING_REASON_CROSS_FORMAT)
}

View File

@@ -0,0 +1,131 @@
use crate::{
ProviderKeyHealthBucket, SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub enum SchedulerRankingMode {
FixedOrder,
#[default]
CacheAffinity,
LoadBalance,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
pub enum SchedulerTunnelAffinityBucket {
LocalTunnel = 0,
#[default]
Neutral = 1,
RemoteTunnel = 2,
}
#[derive(Debug, Clone, PartialEq)]
pub struct SchedulerRankableCandidate {
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
pub selected_provider_model_name: String,
pub provider_priority: i32,
pub key_internal_priority: i32,
pub key_global_priority_for_format: Option<i32>,
pub capability_priority: (u32, u32),
pub cached_affinity_match: bool,
pub affinity_hash: Option<u64>,
pub tunnel_bucket: SchedulerTunnelAffinityBucket,
pub demote_cross_format: bool,
pub format_preference: (u8, u8),
pub health_bucket: Option<ProviderKeyHealthBucket>,
pub health_score: f64,
pub original_index: usize,
}
impl SchedulerRankableCandidate {
pub fn from_candidate(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
original_index: usize,
) -> Self {
Self {
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
selected_provider_model_name: candidate.selected_provider_model_name.clone(),
provider_priority: candidate.provider_priority,
key_internal_priority: candidate.key_internal_priority,
key_global_priority_for_format: candidate.key_global_priority_for_format,
capability_priority: (0, 0),
cached_affinity_match: false,
affinity_hash: None,
tunnel_bucket: SchedulerTunnelAffinityBucket::Neutral,
demote_cross_format: false,
format_preference: (0, 0),
health_bucket: None,
health_score: 1.0,
original_index,
}
}
pub fn with_capability_priority(mut self, value: (u32, u32)) -> Self {
self.capability_priority = value;
self
}
pub fn with_cached_affinity_match(mut self, value: bool) -> Self {
self.cached_affinity_match = value;
self
}
pub fn with_affinity_hash(mut self, value: Option<u64>) -> Self {
self.affinity_hash = value;
self
}
pub fn with_tunnel_bucket(mut self, value: SchedulerTunnelAffinityBucket) -> Self {
self.tunnel_bucket = value;
self
}
pub fn with_format_state(
mut self,
demote_cross_format: bool,
format_preference: (u8, u8),
) -> Self {
self.demote_cross_format = demote_cross_format;
self.format_preference = format_preference;
self
}
pub fn with_health(mut self, bucket: Option<ProviderKeyHealthBucket>, score: f64) -> Self {
self.health_bucket = bucket;
self.health_score = score;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SchedulerRankingContext {
pub priority_mode: SchedulerPriorityMode,
pub ranking_mode: SchedulerRankingMode,
pub include_health: bool,
pub load_balance_seed: u64,
}
impl Default for SchedulerRankingContext {
fn default() -> Self {
Self {
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
include_health: false,
load_balance_seed: 0,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct SchedulerRankingOutcome {
pub original_index: usize,
pub ranking_index: usize,
pub priority_mode: SchedulerPriorityMode,
pub ranking_mode: SchedulerRankingMode,
pub priority_slot: i32,
pub promoted_by: Option<&'static str>,
pub demoted_by: Option<&'static str>,
}

View File

@@ -228,6 +228,18 @@
</span> </span>
</span> </span>
</div> </div>
<div
v-if="currentAttemptRankingInfo"
class="info-item"
>
<span class="info-label">排序原因</span>
<span class="info-value info-value-stacked">
<code class="format-code">{{ currentAttemptRankingInfo.summary }}</code>
<span class="text-xs text-muted-foreground">
{{ currentAttemptRankingInfo.hint }}
</span>
</span>
</div>
<div <div
v-if="currentAttempt.key_name || currentAttempt.key_id" v-if="currentAttempt.key_name || currentAttempt.key_id"
class="info-item" class="info-item"
@@ -1292,6 +1304,73 @@ const currentAttemptSchedulerInfo = computed<{
} }
}) })
const normalizeMetadataText = (value: unknown): string => {
return typeof value === 'string' ? value.trim() : ''
}
const formatRankingModeLabel = (value: string): string => {
const normalized = value.replace(/([a-z])([A-Z])/g, '$1_$2').toLowerCase()
const labels: Record<string, string> = {
fixed_order: '固定顺序',
cache_affinity: '亲和性优先',
load_balance: '负载均衡',
}
return labels[normalized] || value
}
const formatPriorityModeLabel = (value: string): string => {
const normalized = value.replace(/([a-z])([A-Z])/g, '$1_$2').toLowerCase()
const labels: Record<string, string> = {
provider: 'Provider 优先级',
global_key: '全局 Key 优先级',
}
return labels[normalized] || value
}
const formatRankingReasonLabel = (value: string): string => {
const labels: Record<string, string> = {
cached_affinity: '缓存亲和性命中',
local_tunnel: '本地隧道优先',
cross_format: '跨格式降级',
}
return labels[value] || value
}
const currentAttemptRankingInfo = computed<{
summary: string
hint: string
} | null>(() => {
const attempt = currentAttempt.value
if (!attempt) return null
const extra = extractObject(attempt.extra_data)
if (!extra) return null
const rankingMode = normalizeMetadataText(extra.ranking_mode)
const priorityMode = normalizeMetadataText(extra.priority_mode)
const promotedBy = normalizeMetadataText(extra.promoted_by)
const demotedBy = normalizeMetadataText(extra.demoted_by)
const rankingIndex = normalizePriorityNumber(extra.ranking_index)
const prioritySlot = normalizePriorityNumber(extra.priority_slot)
if (!rankingMode && !priorityMode && !promotedBy && !demotedBy && rankingIndex === null && prioritySlot === null) {
return null
}
const summaryParts: string[] = []
if (rankingMode) summaryParts.push(formatRankingModeLabel(rankingMode))
if (promotedBy) summaryParts.push(formatRankingReasonLabel(promotedBy))
if (demotedBy) summaryParts.push(formatRankingReasonLabel(demotedBy))
const hintParts: string[] = []
if (rankingIndex !== null) hintParts.push(`排序 #${rankingIndex + 1}`)
if (priorityMode) hintParts.push(formatPriorityModeLabel(priorityMode))
if (prioritySlot !== null) hintParts.push(`槽位 ${prioritySlot}`)
return {
summary: summaryParts.length > 0 ? summaryParts.join(' / ') : '排序元数据',
hint: hintParts.length > 0 ? hintParts.join(' · ') : '候选排序由 scheduler ranking engine 生成',
}
})
const currentAttemptKeyFormatsDisplay = computed(() => { const currentAttemptKeyFormatsDisplay = computed(() => {
const attempt = currentAttempt.value const attempt = currentAttempt.value
if (!attempt) return '' if (!attempt) return ''