mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
Unify candidate ranking pipeline
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -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
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -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(),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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(),
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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(),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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(),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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(),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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(),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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::{
|
||||||
|
|||||||
@@ -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(),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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",
|
||||||
}))));
|
}))));
|
||||||
|
|||||||
28
apps/aether-gateway/src/scheduler/candidate/enumeration.rs
Normal file
28
apps/aether-gateway/src/scheduler/candidate/enumeration.rs
Normal 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()))
|
||||||
|
}
|
||||||
@@ -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;
|
||||||
|
|
||||||
|
|||||||
76
apps/aether-gateway/src/scheduler/candidate/ranking.rs
Normal file
76
apps/aether-gateway/src/scheduler/candidate/ranking.rs
Normal 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,
|
||||||
|
}
|
||||||
|
}
|
||||||
45
apps/aether-gateway/src/scheduler/candidate/resolution.rs
Normal file
45
apps/aether-gateway/src/scheduler/candidate/resolution.rs
Normal 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)
|
||||||
|
}
|
||||||
@@ -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>(
|
||||||
|
|||||||
@@ -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();
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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!(
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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
124
crates/aether-scheduler-core/src/candidate/capability.rs
Normal file
124
crates/aether-scheduler-core/src/candidate/capability.rs
Normal 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"
|
||||||
|
)
|
||||||
|
}
|
||||||
163
crates/aether-scheduler-core/src/candidate/enumeration.rs
Normal file
163
crates/aether-scheduler-core/src/candidate/enumeration.rs
Normal 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()
|
||||||
|
}
|
||||||
39
crates/aether-scheduler-core/src/candidate/identity.rs
Normal file
39
crates/aether-scheduler-core/src/candidate/identity.rs
Normal 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),
|
||||||
|
)
|
||||||
|
}
|
||||||
482
crates/aether-scheduler-core/src/candidate/mod.rs
Normal file
482
crates/aether-scheduler-core/src/candidate/mod.rs
Normal 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,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
204
crates/aether-scheduler-core/src/candidate/selectability.rs
Normal file
204
crates/aether-scheduler-core/src/candidate/selectability.rs
Normal 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)
|
||||||
|
}
|
||||||
41
crates/aether-scheduler-core/src/candidate/types.rs
Normal file
41
crates/aether-scheduler-core/src/candidate/types.rs
Normal 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,
|
||||||
|
}
|
||||||
@@ -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,
|
||||||
|
|||||||
@@ -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>(
|
||||||
|
|||||||
12
crates/aether-scheduler-core/src/ranking/format.rs
Normal file
12
crates/aether-scheduler-core/src/ranking/format.rs
Normal 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))
|
||||||
|
}
|
||||||
294
crates/aether-scheduler-core/src/ranking/mod.rs
Normal file
294
crates/aether-scheduler-core/src/ranking/mod.rs
Normal 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"]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
111
crates/aether-scheduler-core/src/ranking/modes.rs
Normal file
111
crates/aether-scheduler-core/src/ranking/modes.rs
Normal 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;
|
||||||
|
}
|
||||||
|
}
|
||||||
52
crates/aether-scheduler-core/src/ranking/priority.rs
Normal file
52
crates/aether-scheduler-core/src/ranking/priority.rs
Normal 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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
28
crates/aether-scheduler-core/src/ranking/reasons.rs
Normal file
28
crates/aether-scheduler-core/src/ranking/reasons.rs
Normal 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)
|
||||||
|
}
|
||||||
131
crates/aether-scheduler-core/src/ranking/types.rs
Normal file
131
crates/aether-scheduler-core/src/ranking/types.rs
Normal 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>,
|
||||||
|
}
|
||||||
@@ -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 ''
|
||||||
|
|||||||
Reference in New Issue
Block a user