i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Lazy load requested model candidates

Browse Source
This commit is contained in:
fawney19
2026-05-03 20:56:31 +08:00
parent a24e4a793d
commit 6f00cabe96
13 changed files with 1047 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files

263
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
View File

@@ -11,6 +11,8 @@ use async_trait::async_trait;
use serde_json::Value;
use std::collections::VecDeque;
use std::convert::Infallible;
use std::sync::Arc;
use tracing::warn;
use uuid::Uuid;
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
@@ -20,6 +22,9 @@ use crate::ai_serving::planner::candidate_resolution::{
resolve_and_rank_logical_local_execution_candidates, EligibleLocalExecutionCandidate,
LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
};
use crate::ai_serving::planner::candidate_source::{
LocalCandidatePreselectionKeyMode, LocalCandidatePreselectionPageCursor,
};
use crate::ai_serving::planner::materialization_policy::LocalCandidatePersistencePolicy;
use crate::ai_serving::planner::pool_scheduler::PoolKeyCursor;
use crate::ai_serving::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
@@ -42,6 +47,10 @@ pub(crate) struct LocalExecutionCandidateAttemptSource<'a> {
items: VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>,
}
type DecorateSkippedCandidateFn<'a> = Arc<
dyn Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync + 'a,
>;
#[async_trait]
pub(crate) trait LocalExecutionAttemptSource<T>: Send {
async fn next_execution_attempt(&mut self) -> Result<Option<T>, GatewayError>;
@@ -58,6 +67,9 @@ enum LocalExecutionCandidateAttemptSourceItem<'a> {
candidate_index: u32,
pending_attempts: VecDeque<LocalExecutionCandidateAttempt>,
},
RequestedModelPage {
cursor: Box<RequestedModelAttemptPageCursor<'a>>,
},
}
impl<'a> LocalExecutionCandidateAttemptSource<'a> {
@@ -93,6 +105,13 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
*candidate_index,
);
}
LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } => {
let Some(attempt) = cursor.next_attempt().await else {
self.items.pop_front();
continue;
};
return Some(attempt);
}
}
}
}
@@ -476,9 +495,34 @@ where
&candidates,
);
let (items, _) = build_logical_candidate_items(
state,
candidates,
0,
sticky_session_token,
requested_model,
request_auth_channel,
);
(
LocalExecutionCandidateAttemptSource { items },
candidate_count,
)
}
fn build_logical_candidate_items<'a>(
state: PlannerAppState<'a>,
candidates: Vec<EligibleLocalExecutionCandidate>,
starting_candidate_index: u32,
sticky_session_token: Option<&str>,
requested_model: Option<&str>,
request_auth_channel: Option<&str>,
) -> (VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>, u32) {
let mut items = VecDeque::new();
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_index = candidate_index as u32;
let mut next_candidate_index = starting_candidate_index;
for candidate in candidates {
let candidate_index = next_candidate_index;
next_candidate_index = next_candidate_index.saturating_add(1);
match candidate.kind {
LocalExecutionCandidateKind::SingleKey => {
let attempts = build_unpersisted_local_execution_candidate_attempts(
@@ -504,13 +548,228 @@ where
}
}
}
(items, next_candidate_index)
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn build_lazy_requested_model_execution_candidate_attempt_source_with_serving<
'a,
F,
G,
>(
state: PlannerAppState<'a>,
trace_id: &str,
client_api_format: &str,
requested_model: &str,
require_streaming: bool,
auth_snapshot: &GatewayAuthApiKeySnapshot,
required_capabilities: Option<&Value>,
sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
persistence_policy: LocalCandidatePersistencePolicy<'_>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
resolution_mode: LocalCandidateResolutionMode,
build_available_extra_data: F,
decorate_skipped_candidate: G,
) -> (LocalExecutionCandidateAttemptSource<'a>, usize)
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync + 'a,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync + 'a,
{
let _ = build_available_extra_data;
let decorate_skipped_candidate = Arc::new(decorate_skipped_candidate);
let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic;
let page_cursor = LocalCandidatePreselectionPageCursor::new(
state,
client_api_format,
requested_model,
require_streaming,
required_capabilities,
auth_snapshot,
use_api_format_alias_match,
key_mode,
)
.await;
let mut cursor = RequestedModelAttemptPageCursor {
state,
trace_id: trace_id.to_string(),
client_api_format: client_api_format.to_string(),
requested_model: requested_model.to_string(),
auth_snapshot: auth_snapshot.clone(),
required_capabilities: required_capabilities.cloned(),
sticky_session_token: sticky_session_token.map(str::to_string),
request_auth_channel: request_auth_channel.map(str::to_string),
record_runtime_miss_diagnostic,
resolution_mode,
decorate_skipped_candidate,
page_cursor,
pending_items: VecDeque::new(),
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
};
cursor.load_next_page().await;
let candidate_count = cursor.candidate_count;
let mut items = VecDeque::new();
if !cursor.pending_items.is_empty() {
items.push_back(
LocalExecutionCandidateAttemptSourceItem::RequestedModelPage {
cursor: Box::new(cursor),
},
);
}
(
LocalExecutionCandidateAttemptSource { items },
candidate_count,
)
}
struct RequestedModelAttemptPageCursor<'a> {
state: PlannerAppState<'a>,
trace_id: String,
client_api_format: String,
requested_model: String,
auth_snapshot: GatewayAuthApiKeySnapshot,
required_capabilities: Option<Value>,
sticky_session_token: Option<String>,
request_auth_channel: Option<String>,
record_runtime_miss_diagnostic: bool,
resolution_mode: LocalCandidateResolutionMode,
decorate_skipped_candidate: DecorateSkippedCandidateFn<'a>,
page_cursor: LocalCandidatePreselectionPageCursor<'a>,
pending_items: VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>,
candidate_count: usize,
next_candidate_index: u32,
remembered_affinity: bool,
}
impl<'a> RequestedModelAttemptPageCursor<'a> {
async fn next_attempt(&mut self) -> Option<LocalExecutionCandidateAttempt> {
loop {
if let Some(attempt) = pop_attempt_from_items(&mut self.pending_items).await {
return Some(attempt);
}
if !self.load_next_page().await {
return None;
}
}
}
async fn load_next_page(&mut self) -> bool {
loop {
let page = match self.page_cursor.next_page().await {
Ok(Some(page)) => page,
Ok(None) => return false,
Err(error) => {
warn!(
trace_id = %self.trace_id,
error = ?error,
"gateway lazy requested-model candidate page read failed"
);
return false;
}
};
let (candidates, resolved_skipped) =
resolve_and_rank_logical_local_execution_candidates(
self.state,
page.candidates,
&self.client_api_format,
Some(&self.requested_model),
Some(&self.auth_snapshot),
self.required_capabilities.as_ref(),
self.sticky_session_token.as_deref(),
self.request_auth_channel.as_deref(),
self.resolution_mode,
)
.await;
let skipped_candidates = page
.skipped_candidates
.into_iter()
.chain(resolved_skipped)
.map(|skipped| (self.decorate_skipped_candidate)(skipped))
.collect::<Vec<_>>();
self.candidate_count = self
.candidate_count
.saturating_add(candidates.len() + skipped_candidates.len());
if self.record_runtime_miss_diagnostic {
for skipped_candidate in &skipped_candidates {
record_local_runtime_candidate_skip_reason(
self.state.app(),
&self.trace_id,
skipped_candidate.skip_reason,
);
}
}
if !self.remembered_affinity && !candidates.is_empty() {
remember_first_local_candidate_affinity(
self.state,
Some(&self.auth_snapshot),
&self.client_api_format,
Some(&self.requested_model),
&candidates,
);
self.remembered_affinity = true;
}
let (items, next_candidate_index) = build_logical_candidate_items(
self.state,
candidates,
self.next_candidate_index,
self.sticky_session_token.as_deref(),
Some(&self.requested_model),
self.request_auth_channel.as_deref(),
);
self.next_candidate_index = next_candidate_index;
if !items.is_empty() {
self.pending_items = items;
return true;
}
}
}
}
async fn pop_attempt_from_items(
items: &mut VecDeque<LocalExecutionCandidateAttemptSourceItem<'_>>,
) -> Option<LocalExecutionCandidateAttempt> {
loop {
let front = items.front_mut()?;
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if let Some(attempt) = attempts.pop_front() {
if attempts.is_empty() {
items.pop_front();
}
return Some(attempt);
}
items.pop_front();
}
LocalExecutionCandidateAttemptSourceItem::Pool {
cursor,
candidate_index,
pending_attempts,
} => {
if let Some(attempt) = pending_attempts.pop_front() {
return Some(attempt);
}
let Some(candidate) = cursor.next_key().await else {
cursor.log_exhausted();
let _ = cursor.take_skipped_candidates();
items.pop_front();
continue;
};
*pending_attempts = build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
);
}
LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { .. } => {
items.pop_front();
}
}
}
}
pub(crate) fn remember_first_local_candidate_affinity(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,

350
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
View File

@@ -1,13 +1,21 @@
use aether_ai_serving::{
run_ai_candidate_preselection, AiCandidatePreselectionOutcome, AiCandidatePreselectionPort,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use aether_scheduler_core::{
enumerate_minimal_candidate_selection_with_model_directives, normalize_api_format,
resolve_requested_global_model_name_with_model_directives,
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
};
use async_trait::async_trait;
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_secs;
use crate::data::candidate_selection::{
read_requested_model_rows_fast_path_page, requested_model_candidate_names,
REQUESTED_MODEL_CANDIDATE_PAGE_SIZE, REQUESTED_MODEL_MAX_SCANNED_ROWS,
};
use crate::scheduler::candidate::SchedulerSkippedCandidate;
use crate::GatewayError;
@@ -174,6 +182,332 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
run_ai_candidate_preselection(&port).await
}
pub(crate) struct LocalCandidatePreselectionPageCursor<'a> {
state: PlannerAppState<'a>,
client_api_format: String,
requested_model: String,
require_streaming: bool,
required_capabilities: Option<serde_json::Value>,
auth_snapshot: GatewayAuthApiKeySnapshot,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
candidate_api_formats: Vec<String>,
model_directive_enabled_api_formats: BTreeSet<String>,
format_index: usize,
requested_name_indexes: BTreeMap<String, usize>,
requested_name_offsets: BTreeMap<String, u32>,
scanned_rows_by_format: BTreeMap<String, u32>,
resolved_global_model_names: BTreeMap<String, String>,
seen_candidate_keys: BTreeSet<String>,
}
impl<'a> LocalCandidatePreselectionPageCursor<'a> {
#[allow(clippy::too_many_arguments)]
pub(crate) async fn new(
state: PlannerAppState<'a>,
client_api_format: &str,
requested_model: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: &GatewayAuthApiKeySnapshot,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
) -> Self {
let candidate_api_formats =
crate::ai_serving::request_candidate_api_formats(client_api_format, require_streaming)
.into_iter()
.map(str::to_string)
.collect::<Vec<_>>();
let mut model_directive_enabled_api_formats = BTreeSet::new();
for api_format in &candidate_api_formats {
if crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state.app(),
api_format,
Some(requested_model),
)
.await
{
model_directive_enabled_api_formats
.insert(crate::ai_serving::normalize_api_format_alias(api_format));
}
}
Self {
state,
client_api_format: client_api_format.to_string(),
requested_model: requested_model.to_string(),
require_streaming,
required_capabilities: required_capabilities.cloned(),
auth_snapshot: auth_snapshot.clone(),
use_api_format_alias_match,
key_mode,
candidate_api_formats,
model_directive_enabled_api_formats,
format_index: 0,
requested_name_indexes: BTreeMap::new(),
requested_name_offsets: BTreeMap::new(),
scanned_rows_by_format: BTreeMap::new(),
resolved_global_model_names: BTreeMap::new(),
seen_candidate_keys: BTreeSet::new(),
}
}
pub(crate) async fn next_page(
&mut self,
) -> Result<
Option<
AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
>,
GatewayError,
> {
while self.format_index < self.candidate_api_formats.len() {
let candidate_api_format = self.candidate_api_formats[self.format_index].clone();
let Some(outcome) = self.next_page_for_api_format(&candidate_api_format).await? else {
self.format_index += 1;
continue;
};
if outcome.candidates.is_empty() && outcome.skipped_candidates.is_empty() {
continue;
}
return Ok(Some(outcome));
}
Ok(None)
}
async fn next_page_for_api_format(
&mut self,
candidate_api_format: &str,
) -> Result<
Option<
AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
>,
GatewayError,
> {
let normalized_api_format = normalize_api_format(candidate_api_format);
if normalized_api_format.is_empty() {
return Ok(None);
}
let enable_model_directives = self.model_directive_enabled_api_formats.contains(
&crate::ai_serving::normalize_api_format_alias(candidate_api_format),
);
let requested_names =
requested_model_candidate_names(&self.requested_model, enable_model_directives);
let scanned = *self
.scanned_rows_by_format
.get(&normalized_api_format)
.unwrap_or(&0);
if scanned >= REQUESTED_MODEL_MAX_SCANNED_ROWS {
return Ok(None);
}
loop {
let requested_name_index = *self
.requested_name_indexes
.entry(normalized_api_format.clone())
.or_insert(0);
let Some(requested_name) = requested_names.get(requested_name_index) else {
return Ok(None);
};
if requested_name.trim().is_empty() {
self.requested_name_indexes
.insert(normalized_api_format.clone(), requested_name_index + 1);
continue;
}
let offset_key = format!("{normalized_api_format}:{requested_name_index}");
let offset = *self
.requested_name_offsets
.entry(offset_key.clone())
.or_insert(0);
let scanned = *self
.scanned_rows_by_format
.get(&normalized_api_format)
.unwrap_or(&0);
let remaining = REQUESTED_MODEL_MAX_SCANNED_ROWS.saturating_sub(scanned);
if remaining == 0 {
return Ok(None);
}
let limit = REQUESTED_MODEL_CANDIDATE_PAGE_SIZE.min(remaining);
let page = read_requested_model_rows_fast_path_page(
self.state.app().data.as_ref(),
&normalized_api_format,
&self.requested_model,
requested_name,
offset,
limit,
enable_model_directives,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.scanned_rows_by_format.insert(
normalized_api_format.clone(),
scanned.saturating_add(page.scanned_rows),
);
self.requested_name_offsets
.insert(offset_key, offset.saturating_add(limit));
if page.end_of_requested_name {
self.requested_name_indexes
.insert(normalized_api_format.clone(), requested_name_index + 1);
}
if page.scanned_rows == 0 {
if requested_name_index + 1 >= requested_names.len() {
return Ok(None);
}
continue;
}
let mut rows = page
.rows
.into_iter()
.filter(|row| {
self.seen_candidate_keys.insert(format!(
"{}:{}:{}:{}",
row.endpoint_id, row.key_id, row.model_id, row.endpoint_api_format
))
})
.collect::<Vec<_>>();
if rows.is_empty() {
continue;
}
let resolved_global_model_name =
if let Some(value) = self.resolved_global_model_names.get(&normalized_api_format) {
value.clone()
} else {
let Some(value) = resolve_requested_global_model_name_with_model_directives(
&rows,
&self.requested_model,
&normalized_api_format,
enable_model_directives,
) else {
continue;
};
self.resolved_global_model_names
.insert(normalized_api_format.clone(), value.clone());
value
};
rows.retain(|row| row.global_model_name == resolved_global_model_name);
if rows.is_empty() {
continue;
}
let auth_constraints = matches_client_api_format(
self.use_api_format_alias_match,
candidate_api_format,
&self.client_api_format,
)
.then_some(&self.auth_snapshot)
.map(crate::data::candidate_selection::auth_snapshot_constraints);
let enumerated_candidates =
enumerate_minimal_candidate_selection_with_model_directives(
EnumerateMinimalCandidateSelectionInput {
rows,
normalized_api_format: &normalized_api_format,
requested_model_name: &self.requested_model,
resolved_global_model_name: resolved_global_model_name.as_str(),
require_streaming: self.require_streaming,
required_capabilities: self.required_capabilities.as_ref(),
auth_constraints: auth_constraints.as_ref(),
},
enable_model_directives,
)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let mut candidates = Vec::new();
for candidate in enumerated_candidates {
if !self.candidate_allowed_for_page(
&candidate,
candidate_api_format,
enable_model_directives,
) {
continue;
}
if !self
.seen_candidate_keys
.insert(local_candidate_preselection_key(&candidate, self.key_mode))
{
continue;
}
candidates.push(candidate);
}
let matches_client_format = matches_client_api_format(
self.use_api_format_alias_match,
candidate_api_format,
&self.client_api_format,
);
let auth_snapshot = matches_client_format.then_some(&self.auth_snapshot);
let (candidates, skipped_candidates) = self
.state
.list_selectable_enumerated_candidates_with_skip_reasons(
candidate_api_format,
&resolved_global_model_name,
candidates,
self.required_capabilities.as_ref(),
auth_snapshot,
current_unix_secs(),
)
.await?;
let skipped_candidates = skipped_candidates
.into_iter()
.map(skipped_local_execution_candidate_from_scheduler_skip)
.filter(|skipped_candidate| {
self.skipped_candidate_allowed_for_page(
skipped_candidate,
candidate_api_format,
enable_model_directives,
)
})
.collect::<Vec<_>>();
return Ok(Some(AiCandidatePreselectionOutcome {
candidates,
skipped_candidates,
}));
}
}
fn candidate_allowed_for_page(
&self,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_api_format: &str,
enable_model_directives: bool,
) -> bool {
matches_client_api_format(
self.use_api_format_alias_match,
candidate_api_format,
&self.client_api_format,
) || auth_snapshot_allows_cross_format_candidate(
&self.auth_snapshot,
&self.requested_model,
candidate,
enable_model_directives,
)
}
fn skipped_candidate_allowed_for_page(
&self,
skipped_candidate: &SkippedLocalExecutionCandidate,
candidate_api_format: &str,
enable_model_directives: bool,
) -> bool {
matches_client_api_format(
self.use_api_format_alias_match,
candidate_api_format,
&self.client_api_format,
) || auth_snapshot_allows_cross_format_candidate(
&self.auth_snapshot,
&self.requested_model,
&skipped_candidate.candidate,
enable_model_directives,
)
}
}
fn skipped_local_execution_candidate_from_scheduler_skip(
skipped_candidate: SchedulerSkippedCandidate,
) -> SkippedLocalExecutionCandidate {
@@ -211,6 +545,18 @@ fn local_candidate_preselection_key(
}
}
fn matches_client_api_format(
use_api_format_alias_match: bool,
candidate_api_format: &str,
client_api_format: &str,
) -> bool {
if use_api_format_alias_match {
crate::ai_serving::api_format_alias_matches(candidate_api_format, client_api_format)
} else {
candidate_api_format == client_api_format
}
}
pub(crate) fn auth_snapshot_allows_cross_format_candidate(
auth_snapshot: &GatewayAuthApiKeySnapshot,
requested_model: &str,

135
apps/aether-gateway/src/ai_serving/planner/standard/family/candidates.rs
View File

@@ -1,7 +1,7 @@
use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::{
build_local_execution_candidate_attempt_source_with_serving,
build_lazy_requested_model_execution_candidate_attempt_source_with_serving,
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
LocalExecutionCandidateAttemptSource,
};
@@ -184,79 +184,70 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::StandardDecision,
);
let preselection = preselect_local_execution_candidates_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::Standard,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.as_str(),
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
Ok(
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
&input.auth_snapshot,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
LocalCandidateResolutionMode::Standard,
move |eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
serde_json::Map::new(),
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
eligible.candidate.endpoint_api_format.as_str(),
))
},
move |mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
spec_metadata.api_format,
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await,
)
.await)
}

1
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision.rs
View File

@@ -7,6 +7,7 @@ mod support;
pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for_candidate;
pub(super) use self::support::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,

79
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs
View File

@@ -1,6 +1,7 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_serving::planner::candidate_materialization::{
build_lazy_requested_model_execution_candidate_attempt_source_with_serving,
build_local_execution_candidate_attempt_source_with_serving,
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -13,6 +14,7 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::candidate_source::LocalCandidatePreselectionKeyMode;
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
@@ -269,3 +271,80 @@ pub(crate) async fn build_local_openai_chat_candidate_attempt_source<'a>(
)
.await
}
pub(crate) async fn build_lazy_local_openai_chat_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiChatDecisionInput,
body_json: &serde_json::Value,
require_streaming: bool,
) -> (LocalOpenAiChatCandidateAttemptSource<'a>, usize) {
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
);
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
"openai:chat",
&input.requested_model,
require_streaming,
&input.auth_snapshot,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
LocalCandidateResolutionMode::Standard,
move |eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: "openai:chat",
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.trim(),
))
},
move |mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
"openai:chat",
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await
}

1
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
View File

@@ -12,6 +12,7 @@ mod decision;
mod plans;
use self::decision::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatCandidateAttempt,

36
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
View File

@@ -2,6 +2,7 @@ use async_trait::async_trait;
use tracing::warn;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
@@ -49,27 +50,10 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
return Ok(None);
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, true).await {
Ok(value) => value,
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat stream decision scheduler selection failed"
);
set_local_openai_chat_miss_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
"scheduler_selection_failed",
);
return Ok(None);
}
};
let candidate_count = candidates.len() + skipped_candidates.len();
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, true,
)
.await;
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
@@ -90,16 +74,6 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
candidate_count,
);
let (candidates, candidate_count) = build_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
Ok(Some((
LocalOpenAiChatStreamAttemptSource {
state,

36
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
View File

@@ -2,6 +2,7 @@ use async_trait::async_trait;
use tracing::warn;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
@@ -58,27 +59,10 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
return Ok(None);
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, false).await {
Ok(value) => value,
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat sync decision scheduler selection failed"
);
set_local_openai_chat_miss_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
"scheduler_selection_failed",
);
return Ok(None);
}
};
let candidate_count = candidates.len() + skipped_candidates.len();
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, false,
)
.await;
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
@@ -99,16 +83,6 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
candidate_count,
);
let (candidates, candidate_count) = build_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
Ok(Some((
LocalOpenAiChatSyncAttemptSource {
state,

133
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/support.rs
View File

@@ -2,6 +2,7 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::{
build_lazy_requested_model_execution_candidate_attempt_source_with_serving,
build_local_execution_candidate_attempt_source_with_serving,
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -240,80 +241,72 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
);
let preselection = preselect_local_execution_candidates_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::Standard,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.as_str(),
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
Ok(
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
&input.auth_snapshot,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
LocalCandidateResolutionMode::Standard,
move |eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
serde_json::Map::new(),
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
eligible.candidate.endpoint_api_format.as_str(),
))
},
move |mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
spec_metadata.api_format,
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await,
)
.await)
}
pub(crate) async fn mark_skipped_local_openai_responses_candidate(

27
apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs
View File

@@ -99,6 +99,33 @@ impl<'a> PlannerAppState<'a> {
}
}
pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
self,
api_format: &str,
global_model_name: &str,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
),
GatewayError,
> {
crate::scheduler::candidate::list_selectable_enumerated_candidates_with_skip_reasons(
self.app(),
api_format,
global_model_name,
candidates,
required_capabilities,
auth_snapshot,
now_unix_secs,
)
.await
}
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
self,
candidate_api_format: &str,

141
apps/aether-gateway/src/data/candidate_selection.rs
View File

@@ -45,8 +45,15 @@ pub(crate) trait MinimalCandidateSelectionRowSource {
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>;
}
const REQUESTED_MODEL_CANDIDATE_PAGE_SIZE: u32 = 256;
const REQUESTED_MODEL_MAX_SCANNED_ROWS: u32 = 2048;
pub(crate) const REQUESTED_MODEL_CANDIDATE_PAGE_SIZE: u32 = 256;
pub(crate) const REQUESTED_MODEL_MAX_SCANNED_ROWS: u32 = 2048;
#[derive(Debug, Clone)]
pub(crate) struct RequestedModelCandidateRowsPage {
pub(crate) rows: Vec<StoredMinimalCandidateSelectionRow>,
pub(crate) scanned_rows: u32,
pub(crate) end_of_requested_name: bool,
}
pub(crate) async fn read_requested_model_rows(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
@@ -125,16 +132,8 @@ async fn read_requested_model_rows_fast_path(
requested_model_name: &str,
enable_model_directives: bool,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let mut requested_names = vec![requested_model_name.trim().to_string()];
if enable_model_directives {
if let Some(base_model) =
crate::ai_serving::model_directive_base_model(requested_model_name)
{
if !requested_names.iter().any(|value| value == &base_model) {
requested_names.push(base_model);
}
}
}
let requested_names =
requested_model_candidate_names(requested_model_name, enable_model_directives);
let mut rows = Vec::new();
let mut seen = BTreeSet::new();
@@ -180,6 +179,58 @@ async fn read_requested_model_rows_fast_path(
Ok(rows)
}
pub(crate) fn requested_model_candidate_names(
requested_model_name: &str,
enable_model_directives: bool,
) -> Vec<String> {
let mut requested_names = vec![requested_model_name.trim().to_string()];
if enable_model_directives {
if let Some(base_model) =
crate::ai_serving::model_directive_base_model(requested_model_name)
{
if !requested_names.iter().any(|value| value == &base_model) {
requested_names.push(base_model);
}
}
}
requested_names
}
pub(crate) async fn read_requested_model_rows_fast_path_page(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
requested_name: &str,
offset: u32,
limit: u32,
enable_model_directives: bool,
) -> Result<RequestedModelCandidateRowsPage, DataLayerError> {
let limit = limit.max(1);
let page = state
.read_minimal_candidate_selection_rows_for_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: api_format.to_string(),
requested_model_name: requested_name.to_string(),
offset,
limit,
},
)
.await?;
let scanned_rows = page.len() as u32;
let end_of_requested_name = scanned_rows < limit;
let rows = filter_rows_for_requested_model(
page,
requested_model_name,
api_format,
enable_model_directives,
);
Ok(RequestedModelCandidateRowsPage {
rows,
scanned_rows,
end_of_requested_name,
})
}
pub(crate) async fn enumerate_minimal_candidate_selection_with_required_capabilities(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
@@ -308,7 +359,9 @@ pub(crate) async fn read_global_model_names_for_api_format(
Ok(model_names.into_iter().collect())
}
fn auth_snapshot_constraints(snapshot: &GatewayAuthApiKeySnapshot) -> SchedulerAuthConstraints {
pub(crate) fn auth_snapshot_constraints(
snapshot: &GatewayAuthApiKeySnapshot,
) -> SchedulerAuthConstraints {
SchedulerAuthConstraints {
allowed_providers: snapshot
.effective_allowed_providers()
@@ -325,8 +378,8 @@ fn auth_snapshot_constraints(snapshot: &GatewayAuthApiKeySnapshot) -> SchedulerA
#[cfg(test)]
mod tests {
use super::{
read_requested_model_rows, MinimalCandidateSelectionRowSource,
StoredMinimalCandidateSelectionRow,
read_requested_model_rows, read_requested_model_rows_fast_path_page,
MinimalCandidateSelectionRowSource, StoredMinimalCandidateSelectionRow,
};
use aether_data::DataLayerError;
use aether_data_contracts::repository::candidate_selection::{
@@ -497,4 +550,62 @@ mod tests {
);
assert_eq!(source.fallback_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn requested_model_rows_page_reads_only_requested_slice() {
let mut rows = Vec::new();
for index in 0..10 {
let mut row = sample_row("gpt-5");
row.key_id = format!("key-{index}");
rows.push(row);
}
let source = CountingSelectionSource::new(rows, Vec::new());
let page = read_requested_model_rows_fast_path_page(
&source,
"openai:chat",
"gpt-5",
"gpt-5",
4,
3,
false,
)
.await
.expect("page read should succeed");
assert_eq!(page.scanned_rows, 3);
assert!(!page.end_of_requested_name);
assert_eq!(
page.rows
.iter()
.map(|row| row.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-4", "key-5", "key-6"]
);
assert_eq!(source.fast_calls.load(Ordering::SeqCst), 1);
assert_eq!(source.fallback_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn requested_model_rows_page_reports_end_of_requested_name() {
let source = CountingSelectionSource::new(vec![sample_row("gpt-5")], Vec::new());
let page = read_requested_model_rows_fast_path_page(
&source,
"openai:chat",
"gpt-5",
"gpt-5",
0,
3,
false,
)
.await
.expect("page read should succeed");
assert_eq!(page.scanned_rows, 1);
assert!(page.end_of_requested_name);
assert_eq!(page.rows.len(), 1);
assert_eq!(source.fast_calls.load(Ordering::SeqCst), 1);
assert_eq!(source.fallback_calls.load(Ordering::SeqCst), 0);
}
}

34
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -1,5 +1,6 @@
use self::selection::{
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
collect_selectable_enumerated_candidates_with_skip_reasons,
};
use super::state::SchedulerRuntimeState;
@@ -111,6 +112,39 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
),
GatewayError,
> {
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let priority_affinity_key =
selection::scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
collect_selectable_enumerated_candidates_with_skip_reasons(
runtime_state,
api_format,
global_model_name,
candidates,
required_capabilities,
auth_snapshot,
now_unix_secs,
ordering_config,
priority_affinity_key,
)
.await
}
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,

36
apps/aether-gateway/src/scheduler/candidate/selection.rs
View File

@@ -112,7 +112,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let priority_affinity_key =
scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
let mut candidates = enumerate_scheduler_candidates(
let candidates = enumerate_scheduler_candidates(
selection_row_source,
api_format,
global_model_name,
@@ -122,6 +122,38 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
enable_model_directives,
)
.await?;
collect_selectable_enumerated_candidates_with_skip_reasons(
runtime_state,
api_format,
global_model_name,
candidates,
required_capabilities,
auth_snapshot,
now_unix_secs,
ordering_config,
priority_affinity_key,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
mut candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
ordering_config: crate::scheduler::config::SchedulerOrderingConfig,
priority_affinity_key: Option<&str>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
),
GatewayError,
> {
let runtime_snapshot =
read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
.await?;
@@ -178,7 +210,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
Ok((selected, skipped))
}
fn scheduling_priority_affinity_key<'a>(
pub(super) fn scheduling_priority_affinity_key<'a>(
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
scheduling_mode: SchedulerSchedulingMode,
) -> Option<&'a str> {