mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix: keep scheduler affinity out of runtime checks
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use aether_scheduler_core::{candidate_key, SchedulerAffinityTarget};
|
||||
use aether_scheduler_core::candidate_key;
|
||||
|
||||
use super::runtime::{current_candidate_runtime_skip_reason, CandidateRuntimeSelectionSnapshot};
|
||||
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerSkippedCandidate};
|
||||
@@ -9,7 +9,6 @@ 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>,
|
||||
@@ -21,12 +20,9 @@ pub(super) fn resolve_scheduler_candidate_selectability(
|
||||
|
||||
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 let Some(skip_reason) =
|
||||
current_candidate_runtime_skip_reason(&candidate, runtime_snapshot, now_unix_secs)
|
||||
{
|
||||
if emitted_skipped_keys.insert(key) {
|
||||
skipped.push(SchedulerSkippedCandidate {
|
||||
candidate,
|
||||
|
||||
@@ -8,7 +8,7 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
|
||||
use aether_scheduler_core::{
|
||||
auth_api_key_concurrency_limit_reached, build_provider_concurrent_limit_map,
|
||||
candidate_is_selectable_with_runtime_state, candidate_runtime_skip_reason_with_state,
|
||||
CandidateRuntimeSelectabilityInput, SchedulerAffinityTarget,
|
||||
CandidateRuntimeSelectabilityInput,
|
||||
};
|
||||
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
@@ -100,7 +100,6 @@ pub(super) fn is_candidate_selectable(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
snapshot: &CandidateRuntimeSelectionSnapshot,
|
||||
now_unix_secs: u64,
|
||||
cached_affinity_target: Option<&SchedulerAffinityTarget>,
|
||||
) -> bool {
|
||||
let pool_group = snapshot
|
||||
.pool_provider_ids
|
||||
@@ -111,7 +110,6 @@ pub(super) fn is_candidate_selectable(
|
||||
provider_concurrent_limits: &snapshot.provider_concurrent_limits,
|
||||
provider_key_rpm_states: &snapshot.provider_key_rpm_states,
|
||||
now_unix_secs,
|
||||
cached_affinity_target,
|
||||
provider_quota_blocks_requests: snapshot
|
||||
.provider_quota_blocks_requests
|
||||
.get(candidate.provider_id.as_str())
|
||||
@@ -145,7 +143,6 @@ pub(super) fn current_candidate_runtime_skip_reason(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
snapshot: &CandidateRuntimeSelectionSnapshot,
|
||||
now_unix_secs: u64,
|
||||
cached_affinity_target: Option<&SchedulerAffinityTarget>,
|
||||
) -> Option<&'static str> {
|
||||
let pool_group = snapshot
|
||||
.pool_provider_ids
|
||||
@@ -171,7 +168,6 @@ pub(super) fn current_candidate_runtime_skip_reason(
|
||||
provider_concurrent_limits: &snapshot.provider_concurrent_limits,
|
||||
provider_key_rpm_states: &snapshot.provider_key_rpm_states,
|
||||
now_unix_secs,
|
||||
cached_affinity_target,
|
||||
provider_quota_blocks_requests,
|
||||
account_quota_exhausted: !pool_group
|
||||
&& snapshot
|
||||
|
||||
@@ -11,9 +11,10 @@ use aether_data_contracts::repository::candidates::{
|
||||
};
|
||||
use aether_scheduler_core::{
|
||||
apply_scheduler_candidate_ranking, candidate_affinity_hash,
|
||||
enumerate_minimal_candidate_selection, EnumerateMinimalCandidateSelectionInput,
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankableCandidate,
|
||||
SchedulerRankingContext, SchedulerRankingMode,
|
||||
enumerate_minimal_candidate_selection, ClientSessionAffinity,
|
||||
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
|
||||
SchedulerPriorityMode, SchedulerRankableCandidate, SchedulerRankingContext,
|
||||
SchedulerRankingMode,
|
||||
};
|
||||
|
||||
use crate::cache::SchedulerAffinityTarget;
|
||||
@@ -45,12 +46,49 @@ async fn select_candidate(
|
||||
require_streaming,
|
||||
None,
|
||||
auth_snapshot,
|
||||
None,
|
||||
now_unix_secs,
|
||||
false,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scheduler_affinity_cache_key_uses_session_scope_when_available() {
|
||||
let auth_snapshot = sample_auth_snapshot("api-key-1");
|
||||
let first_session = ClientSessionAffinity::new(
|
||||
Some("generic".to_string()),
|
||||
Some("session=conversation-a".to_string()),
|
||||
);
|
||||
let second_session = ClientSessionAffinity::new(
|
||||
Some("generic".to_string()),
|
||||
Some("session=conversation-b".to_string()),
|
||||
);
|
||||
|
||||
let first_key = build_scheduler_affinity_cache_key(
|
||||
Some(&auth_snapshot),
|
||||
"openai:chat",
|
||||
"gpt-4.1",
|
||||
Some(&first_session),
|
||||
)
|
||||
.expect("first key should build");
|
||||
let second_key = build_scheduler_affinity_cache_key(
|
||||
Some(&auth_snapshot),
|
||||
"openai:chat",
|
||||
"gpt-4.1",
|
||||
Some(&second_session),
|
||||
)
|
||||
.expect("second key should build");
|
||||
let legacy_key =
|
||||
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None)
|
||||
.expect("legacy key should build");
|
||||
|
||||
assert_ne!(first_key, second_key);
|
||||
assert_ne!(first_key, legacy_key);
|
||||
assert!(first_key.starts_with("scheduler_affinity:v2:"));
|
||||
assert!(!first_key.contains("conversation-a"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn same_priority_candidates_are_distributed_by_affinity_key() {
|
||||
let mut first = sample_row();
|
||||
@@ -168,7 +206,7 @@ async fn reuses_cached_scheduler_affinity_candidate_before_sorted_fallback() {
|
||||
|
||||
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
|
||||
let cache_key =
|
||||
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1")
|
||||
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None)
|
||||
.expect("cache key should build");
|
||||
state.scheduler_affinity_cache.insert(
|
||||
cache_key,
|
||||
@@ -199,7 +237,7 @@ async fn reuses_cached_scheduler_affinity_candidate_before_sorted_fallback() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cached_affinity_candidate_can_use_reserved_provider_key_rpm_capacity() {
|
||||
async fn cached_affinity_candidate_cannot_use_reserved_provider_key_rpm_capacity() {
|
||||
let mut first = sample_row();
|
||||
first.provider_id = "provider-a".to_string();
|
||||
first.provider_name = "openai-a".to_string();
|
||||
@@ -273,7 +311,7 @@ async fn cached_affinity_candidate_can_use_reserved_provider_key_rpm_capacity()
|
||||
|
||||
let auth_snapshot = sample_auth_snapshot("api-key-cached-user");
|
||||
let cache_key =
|
||||
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1")
|
||||
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None)
|
||||
.expect("cache key should build");
|
||||
state.scheduler_affinity_cache.insert(
|
||||
cache_key,
|
||||
@@ -299,6 +337,6 @@ async fn cached_affinity_candidate_can_use_reserved_provider_key_rpm_capacity()
|
||||
.expect("selection should succeed")
|
||||
.expect("candidate should exist");
|
||||
|
||||
assert_eq!(selected.provider_id, "provider-a");
|
||||
assert_eq!(selected.key_id, "key-a");
|
||||
assert_eq!(selected.provider_id, "provider-b");
|
||||
assert_eq!(selected.key_id, "key-b");
|
||||
}
|
||||
|
||||
@@ -311,8 +311,8 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
|
||||
"candidate/runtime.rs should depend on core selectable predicate helper"
|
||||
);
|
||||
assert!(
|
||||
runtime.contains("SchedulerAffinityTarget"),
|
||||
"candidate/runtime.rs should depend on core SchedulerAffinityTarget"
|
||||
!runtime.contains("SchedulerAffinityTarget"),
|
||||
"candidate/runtime.rs should keep affinity out of runtime eligibility checks"
|
||||
);
|
||||
assert!(
|
||||
runtime.contains("auth_api_key_concurrency_limit_reached"),
|
||||
|
||||
Reference in New Issue
Block a user