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

This commit is contained in:
ZheFox
2026-09-03 10:46:49 +08:00
57 changed files with 1850 additions and 696 deletions
@@ -1,7 +1,9 @@
use self::selection::{
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons_and_ordering,
collect_selectable_enumerated_candidates_with_skip_reasons,
resolve_preselection_ordering_config,
};
use super::config::SchedulerOrderingConfig;
use super::state::SchedulerRuntimeState;
mod affinity;
@@ -53,6 +55,10 @@ enum RequiredCapabilityMatchMode {
Exclusive,
}
/// `ordering_config` carries the request's routing-policy derived scheduler
/// config. `None` falls back to the runtime default (system-default routing
/// group, then legacy system-config keys).
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -64,6 +70,7 @@ pub(crate) async fn list_selectable_candidates(
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
collect_selectable_candidates(
selection_row_source,
@@ -76,6 +83,7 @@ pub(crate) async fn list_selectable_candidates(
client_session_affinity,
now_unix_secs,
enable_model_directives,
ordering_config,
)
.await
}
@@ -87,6 +95,7 @@ pub(crate) fn is_exact_all_skipped_by_auth_limit(
selection::is_exact_all_skipped_by_auth_limit(selected, skipped)
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_with_skip_reasons(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -98,6 +107,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -105,7 +115,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
),
GatewayError,
> {
collect_selectable_candidates_with_skip_reasons(
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
@@ -117,10 +127,12 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
now_unix_secs,
enable_model_directives,
None,
ordering_config,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_operation(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -133,6 +145,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
now_unix_secs: u64,
enable_model_directives: bool,
request_operation: Option<&str>,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -140,7 +153,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
),
GatewayError,
> {
collect_selectable_candidates_with_skip_reasons(
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
@@ -152,6 +165,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons_for_request_ope
now_unix_secs,
enable_model_directives,
request_operation,
ordering_config,
)
.await
}
@@ -166,6 +180,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -173,7 +188,8 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
),
GatewayError,
> {
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let ordering_config =
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
let priority_affinity_key = selection::scheduling_priority_affinity_key(
auth_snapshot,
client_session_affinity,
@@ -194,6 +210,7 @@ pub(crate) async fn list_selectable_enumerated_candidates_with_skip_reasons(
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -203,6 +220,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
Ok(
list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
@@ -214,12 +232,14 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
auth_snapshot,
client_session_affinity,
now_unix_secs,
ordering_config,
)
.await?
.0,
)
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -229,6 +249,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<(Vec<SchedulerMinimalCandidateSelectionCandidate>, bool), GatewayError> {
let normalized_api_format = normalize_api_format(candidate_api_format);
if normalized_api_format.is_empty() {
@@ -262,20 +283,22 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
let mut all_attempts_blocked_by_auth_limit = !model_names.is_empty();
for global_model_name in model_names {
let (candidates, skipped_candidates) = collect_selectable_candidates_with_skip_reasons(
selection_row_source,
runtime_state,
&normalized_api_format,
&global_model_name,
require_streaming,
required_capabilities.as_ref(),
auth_snapshot,
client_session_affinity,
now_unix_secs,
false,
None,
)
.await?;
let (candidates, skipped_candidates) =
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
&normalized_api_format,
&global_model_name,
require_streaming,
required_capabilities.as_ref(),
auth_snapshot,
client_session_affinity,
now_unix_secs,
false,
None,
ordering_config,
)
.await?;
all_attempts_blocked_by_auth_limit &=
is_exact_all_skipped_by_auth_limit(&candidates, &skipped_candidates);
match capability_mode {
@@ -1,7 +1,7 @@
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use crate::GatewayError;
use aether_scheduler_core::ClientSessionAffinity;
@@ -115,6 +115,7 @@ pub(super) async fn select_minimal_candidate(
Ok(selected)
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn collect_selectable_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -126,24 +127,32 @@ pub(super) async fn collect_selectable_candidates(
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
Ok(collect_selectable_candidates_with_skip_reasons(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
None,
Ok(
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
None,
ordering_config,
)
.await?
.0,
)
.await?
.0)
}
/// Legacy-shaped entrypoint that resolves the ordering config from the
/// runtime state. Prefer `collect_selectable_candidates_with_skip_reasons_and_ordering`
/// and pass the request's routing-policy config explicitly.
#[allow(clippy::too_many_arguments)]
pub(super) async fn collect_selectable_candidates_with_skip_reasons(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
@@ -163,7 +172,59 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
),
GatewayError,
> {
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source,
runtime_state,
api_format,
global_model_name,
require_streaming,
required_capabilities,
auth_snapshot,
client_session_affinity,
now_unix_secs,
enable_model_directives,
request_operation,
None,
)
.await
}
/// Resolve the ordering config for a preselection pass: the routing-policy
/// derived config wins when the caller has one; otherwise fall back to the
/// runtime default (system-default routing group, then legacy keys).
pub(super) async fn resolve_preselection_ordering_config(
runtime_state: &impl SchedulerRuntimeState,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<SchedulerOrderingConfig, GatewayError> {
match ordering_config {
Some(config) => Ok(config),
None => runtime_state.read_scheduler_ordering_config().await,
}
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn collect_selectable_candidates_with_skip_reasons_and_ordering(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
runtime_state: &impl SchedulerRuntimeState,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
enable_model_directives: bool,
request_operation: Option<&str>,
ordering_config: Option<SchedulerOrderingConfig>,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
),
GatewayError,
> {
let ordering_config =
resolve_preselection_ordering_config(runtime_state, ordering_config).await?;
let priority_affinity_key = scheduling_priority_affinity_key(
auth_snapshot,
client_session_affinity,
@@ -205,7 +266,7 @@ pub(super) async fn collect_selectable_enumerated_candidates_with_skip_reasons(
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
now_unix_secs: u64,
ordering_config: crate::scheduler::config::SchedulerOrderingConfig,
ordering_config: SchedulerOrderingConfig,
priority_affinity_key: Option<&str>,
) -> Result<
(
@@ -65,6 +65,7 @@ async fn compatible_required_capability_prefers_matching_keys_without_hard_filte
None,
None,
100,
None,
)
.await
.expect("selection should succeed");
@@ -120,6 +121,7 @@ async fn exclusive_required_capability_keeps_hard_filtering_only_matching_keys()
None,
None,
100,
None,
)
.await
.expect("selection should succeed");
@@ -196,6 +198,7 @@ async fn required_capability_without_model_uses_session_scoped_affinity() {
Some(&auth_snapshot),
Some(&client_session_affinity),
100,
None,
)
.await
.expect("selection should succeed");
@@ -273,6 +276,7 @@ async fn required_capability_reports_auth_limit_signal_when_every_model_is_block
Some(&auth_snapshot),
None,
100,
None,
)
.await
.expect("selection should succeed");
@@ -75,6 +75,7 @@ async fn collect_selectable_candidates(
None,
now_unix_secs,
false,
None,
)
.await
}