mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 18:37:46 +08:00
Merge remote-tracking branch 'origin/main' into codex/fix-antigravity-quota
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user