mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-16 16:07:45 +08:00
Merge pull request #824 from hkxiaoyao/fix/fixed-order-target-select
fix(routing): preserve fixed order for streaming chat
This commit is contained in:
@@ -93,6 +93,16 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
|
|||||||
&self,
|
&self,
|
||||||
candidate: Self::Candidate,
|
candidate: Self::Candidate,
|
||||||
) -> Self::Skipped {
|
) -> Self::Skipped {
|
||||||
|
warn!(
|
||||||
|
event_name = "local_candidate_skipped",
|
||||||
|
log_type = "event",
|
||||||
|
provider_id = %candidate.provider_id,
|
||||||
|
endpoint_id = %candidate.endpoint_id,
|
||||||
|
key_id = %candidate.key_id,
|
||||||
|
api_format = %candidate.endpoint_api_format,
|
||||||
|
skip_reason = "transport_snapshot_missing",
|
||||||
|
"local execution candidate skipped during planning"
|
||||||
|
);
|
||||||
SkippedLocalExecutionCandidate {
|
SkippedLocalExecutionCandidate {
|
||||||
candidate,
|
candidate,
|
||||||
skip_reason: "transport_snapshot_missing",
|
skip_reason: "transport_snapshot_missing",
|
||||||
@@ -145,6 +155,16 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
|
|||||||
transport: Self::Transport,
|
transport: Self::Transport,
|
||||||
skip_reason: &'static str,
|
skip_reason: &'static str,
|
||||||
) -> Self::Skipped {
|
) -> Self::Skipped {
|
||||||
|
warn!(
|
||||||
|
event_name = "local_candidate_skipped",
|
||||||
|
log_type = "event",
|
||||||
|
provider_id = %candidate.provider_id,
|
||||||
|
endpoint_id = %candidate.endpoint_id,
|
||||||
|
key_id = %candidate.key_id,
|
||||||
|
api_format = %candidate.endpoint_api_format,
|
||||||
|
skip_reason,
|
||||||
|
"local execution candidate skipped during planning"
|
||||||
|
);
|
||||||
SkippedLocalExecutionCandidate {
|
SkippedLocalExecutionCandidate {
|
||||||
candidate,
|
candidate,
|
||||||
skip_reason,
|
skip_reason,
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
use aether_routing_core::RoutingExecutionPolicy;
|
use aether_routing_core::{RoutingExecutionPolicy, RoutingSchedulingMode};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use std::collections::VecDeque;
|
use std::collections::VecDeque;
|
||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
@@ -207,7 +207,12 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
|
|||||||
async fn next_raw_attempt_with_target_select(
|
async fn next_raw_attempt_with_target_select(
|
||||||
&mut self,
|
&mut self,
|
||||||
) -> Result<Option<LocalOpenAiChatCandidateAttempt>, GatewayError> {
|
) -> Result<Option<LocalOpenAiChatCandidateAttempt>, GatewayError> {
|
||||||
let select_window = openai_chat_stream_target_select_window();
|
let select_window = openai_chat_stream_target_select_window_for_mode(
|
||||||
|
self.input
|
||||||
|
.routing_policy
|
||||||
|
.as_ref()
|
||||||
|
.map(|policy| policy.scheduling_mode),
|
||||||
|
);
|
||||||
if select_window <= 1 {
|
if select_window <= 1 {
|
||||||
return self.next_raw_attempt_linear().await;
|
return self.next_raw_attempt_linear().await;
|
||||||
}
|
}
|
||||||
@@ -365,6 +370,15 @@ fn openai_chat_stream_target_select_window() -> usize {
|
|||||||
.clamp(1, MAX_OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW)
|
.clamp(1, MAX_OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn openai_chat_stream_target_select_window_for_mode(
|
||||||
|
scheduling_mode: Option<RoutingSchedulingMode>,
|
||||||
|
) -> usize {
|
||||||
|
if scheduling_mode == Some(RoutingSchedulingMode::FixedOrder) {
|
||||||
|
return 1;
|
||||||
|
}
|
||||||
|
openai_chat_stream_target_select_window()
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Clone, Copy)]
|
#[derive(Clone, Copy)]
|
||||||
struct TargetSelectCandidateIdentity<'a> {
|
struct TargetSelectCandidateIdentity<'a> {
|
||||||
provider_id: &'a str,
|
provider_id: &'a str,
|
||||||
@@ -574,4 +588,14 @@ mod tests {
|
|||||||
|
|
||||||
assert_eq!(select_target_index(19, &choices), 1);
|
assert_eq!(select_target_index(19, &choices), 1);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn fixed_order_disables_stream_target_selection() {
|
||||||
|
assert_eq!(
|
||||||
|
openai_chat_stream_target_select_window_for_mode(Some(
|
||||||
|
RoutingSchedulingMode::FixedOrder,
|
||||||
|
)),
|
||||||
|
1
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -634,7 +634,9 @@ impl<'a> PoolKeyCursor<'a> {
|
|||||||
|
|
||||||
if !self.score_phase_exhausted {
|
if !self.score_phase_exhausted {
|
||||||
if let Some(score_candidates) = self.next_score_candidates().await {
|
if let Some(score_candidates) = self.next_score_candidates().await {
|
||||||
return Some(score_candidates);
|
if !score_candidates.is_empty() {
|
||||||
|
return Some(score_candidates);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -4168,6 +4170,115 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn inactive_pool_key_with_stale_score_does_not_exhaust_pool() {
|
||||||
|
let provider_config = Some(json!({
|
||||||
|
"pool_advanced": {
|
||||||
|
"score_top_n": 128,
|
||||||
|
"scheduling_presets": [
|
||||||
|
{"preset": "single_account", "enabled": true},
|
||||||
|
{"preset": "priority_first", "enabled": true}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
let (provider, endpoint, mut keys, mut rows) =
|
||||||
|
large_pool_fixture(2, provider_config.clone());
|
||||||
|
keys[1].is_active = false;
|
||||||
|
rows.retain(|row| row.key_id != "key-00001");
|
||||||
|
let scores = vec![
|
||||||
|
sample_provider_key_pool_score("provider-pool", "key-00000", 5.0),
|
||||||
|
sample_provider_key_pool_score("provider-pool", "key-00001", 20.0),
|
||||||
|
];
|
||||||
|
let data_state =
|
||||||
|
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
|
||||||
|
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![provider],
|
||||||
|
vec![endpoint],
|
||||||
|
keys,
|
||||||
|
)),
|
||||||
|
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
|
||||||
|
)
|
||||||
|
.with_pool_score_repository_for_tests(Arc::new(
|
||||||
|
InMemoryPoolMemberScoreRepository::seed(scores),
|
||||||
|
))
|
||||||
|
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
|
||||||
|
let app = AppState::new()
|
||||||
|
.expect("state should build")
|
||||||
|
.with_data_state_for_tests(data_state);
|
||||||
|
let group = sample_eligible_candidate(
|
||||||
|
"provider-pool",
|
||||||
|
"endpoint-1",
|
||||||
|
"pool-group",
|
||||||
|
10,
|
||||||
|
provider_config,
|
||||||
|
);
|
||||||
|
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
|
||||||
|
|
||||||
|
let candidate = cursor
|
||||||
|
.next_key()
|
||||||
|
.await
|
||||||
|
.expect("active key must stay schedulable beside a stale inactive score");
|
||||||
|
|
||||||
|
assert_eq!(candidate.candidate.key_id, "key-00000");
|
||||||
|
assert_eq!(
|
||||||
|
cursor.skip_reason_counts.get("pool_score_member_missing"),
|
||||||
|
Some(&1)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn stale_inactive_score_only_does_not_exhaust_pool() {
|
||||||
|
let provider_config = Some(json!({
|
||||||
|
"pool_advanced": {
|
||||||
|
"score_top_n": 128,
|
||||||
|
"scheduling_presets": [
|
||||||
|
{"preset": "single_account", "enabled": true},
|
||||||
|
{"preset": "priority_first", "enabled": true}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
let (provider, endpoint, mut keys, mut rows) =
|
||||||
|
large_pool_fixture(2, provider_config.clone());
|
||||||
|
keys[1].is_active = false;
|
||||||
|
rows.retain(|row| row.key_id != "key-00001");
|
||||||
|
let scores = vec![sample_provider_key_pool_score(
|
||||||
|
"provider-pool",
|
||||||
|
"key-00001",
|
||||||
|
20.0,
|
||||||
|
)];
|
||||||
|
let data_state =
|
||||||
|
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
|
||||||
|
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![provider],
|
||||||
|
vec![endpoint],
|
||||||
|
keys,
|
||||||
|
)),
|
||||||
|
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
|
||||||
|
)
|
||||||
|
.with_pool_score_repository_for_tests(Arc::new(
|
||||||
|
InMemoryPoolMemberScoreRepository::seed(scores),
|
||||||
|
))
|
||||||
|
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
|
||||||
|
let app = AppState::new()
|
||||||
|
.expect("state should build")
|
||||||
|
.with_data_state_for_tests(data_state);
|
||||||
|
let group = sample_eligible_candidate(
|
||||||
|
"provider-pool",
|
||||||
|
"endpoint-1",
|
||||||
|
"pool-group",
|
||||||
|
10,
|
||||||
|
provider_config,
|
||||||
|
);
|
||||||
|
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
|
||||||
|
|
||||||
|
let candidate = cursor
|
||||||
|
.next_key()
|
||||||
|
.await
|
||||||
|
.expect("catalog rows must remain schedulable when the only score is stale");
|
||||||
|
|
||||||
|
assert_eq!(candidate.candidate.key_id, "key-00000");
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn score_candidates_continue_across_pool_windows() {
|
async fn score_candidates_continue_across_pool_windows() {
|
||||||
let provider_config = Some(json!({
|
let provider_config = Some(json!({
|
||||||
|
|||||||
@@ -23,6 +23,16 @@ pub(super) fn resolve_scheduler_candidate_selectability(
|
|||||||
if let Some(skip_reason) =
|
if let Some(skip_reason) =
|
||||||
current_candidate_runtime_skip_reason(&candidate, runtime_snapshot, now_unix_secs)
|
current_candidate_runtime_skip_reason(&candidate, runtime_snapshot, now_unix_secs)
|
||||||
{
|
{
|
||||||
|
tracing::debug!(
|
||||||
|
event_name = "scheduler_candidate_skipped",
|
||||||
|
log_type = "event",
|
||||||
|
provider_id = %candidate.provider_id,
|
||||||
|
endpoint_id = %candidate.endpoint_id,
|
||||||
|
key_id = %candidate.key_id,
|
||||||
|
api_format = %candidate.endpoint_api_format,
|
||||||
|
skip_reason,
|
||||||
|
"scheduler candidate skipped during runtime selectability resolution"
|
||||||
|
);
|
||||||
if emitted_skipped_keys.insert(key) {
|
if emitted_skipped_keys.insert(key) {
|
||||||
skipped.push(SchedulerSkippedCandidate {
|
skipped.push(SchedulerSkippedCandidate {
|
||||||
candidate,
|
candidate,
|
||||||
|
|||||||
@@ -804,12 +804,43 @@ impl AppState {
|
|||||||
if updated.is_some() {
|
if updated.is_some() {
|
||||||
self.invalidate_provider_routing_caches();
|
self.invalidate_provider_routing_caches();
|
||||||
}
|
}
|
||||||
|
if let Some(key) = updated.as_ref().filter(|key| !key.is_active) {
|
||||||
|
self.delete_inactive_provider_catalog_key_pool_scores(
|
||||||
|
key.provider_id.as_str(),
|
||||||
|
key.id.as_str(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
match updated {
|
match updated {
|
||||||
Some(key) => self.open_provider_catalog_key(key).await.map(Some),
|
Some(key) => self.open_provider_catalog_key(key).await.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn delete_inactive_provider_catalog_key_pool_scores(
|
||||||
|
&self,
|
||||||
|
provider_id: &str,
|
||||||
|
key_id: &str,
|
||||||
|
) {
|
||||||
|
if let Err(err) = self
|
||||||
|
.data
|
||||||
|
.delete_pool_member_scores_for_member(
|
||||||
|
&pool_scores::PoolMemberIdentity::provider_api_key(
|
||||||
|
provider_id.to_string(),
|
||||||
|
key_id.to_string(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
warn!(
|
||||||
|
provider_id,
|
||||||
|
key_id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway provider catalog key deactivate: failed to delete pool member scores"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn compare_and_update_provider_catalog_key_admin_state(
|
pub(crate) async fn compare_and_update_provider_catalog_key_admin_state(
|
||||||
&self,
|
&self,
|
||||||
update: &provider_catalog::ProviderCatalogKeyAdminCasUpdate,
|
update: &provider_catalog::ProviderCatalogKeyAdminCasUpdate,
|
||||||
@@ -843,6 +874,15 @@ impl AppState {
|
|||||||
if updated.as_ref().is_some_and(|keys| !keys.is_empty()) {
|
if updated.as_ref().is_some_and(|keys| !keys.is_empty()) {
|
||||||
self.invalidate_provider_routing_caches();
|
self.invalidate_provider_routing_caches();
|
||||||
}
|
}
|
||||||
|
if let Some(keys) = updated.as_ref() {
|
||||||
|
for key in keys.iter().filter(|key| !key.is_active) {
|
||||||
|
self.delete_inactive_provider_catalog_key_pool_scores(
|
||||||
|
key.provider_id.as_str(),
|
||||||
|
key.id.as_str(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
match updated {
|
match updated {
|
||||||
Some(keys) => self.open_provider_catalog_keys(keys).await.map(Some),
|
Some(keys) => self.open_provider_catalog_keys(keys).await.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
|||||||
Reference in New Issue
Block a user