mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-10 19:29:50 +08:00
refactor: isolate dispatch scheduling core
This commit is contained in:
@@ -20,8 +20,7 @@ use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::candidates::{
|
||||
build_local_standard_candidate_attempt_source, materialize_local_standard_candidate_attempts,
|
||||
resolve_local_standard_decision_input,
|
||||
build_local_standard_candidate_attempt_source, resolve_local_standard_decision_input,
|
||||
};
|
||||
use super::payload::maybe_build_local_standard_decision_payload_for_candidate;
|
||||
use super::{LocalStandardDecisionInput, LocalStandardSpec};
|
||||
@@ -323,12 +322,12 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
@@ -372,12 +371,12 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
@@ -427,15 +426,15 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
@@ -501,15 +500,15 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (attempts, candidate_count) =
|
||||
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
|
||||
@@ -1,28 +1,23 @@
|
||||
use serde_json::Value;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::common::{
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
};
|
||||
use crate::ai_serving::planner::runtime_miss::set_local_runtime_execution_exhausted_diagnostic;
|
||||
use crate::ai_serving::GatewayControlDecision;
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
use tracing::warn;
|
||||
|
||||
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,
|
||||
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
|
||||
};
|
||||
use self::plans::{
|
||||
build_local_openai_chat_stream_attempt_source, build_local_openai_chat_stream_plan_and_reports,
|
||||
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
|
||||
list_local_openai_chat_candidates, resolve_local_openai_chat_decision_input,
|
||||
set_local_openai_chat_miss_diagnostic,
|
||||
resolve_local_openai_chat_decision_input,
|
||||
};
|
||||
|
||||
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports_for_kind(
|
||||
@@ -146,32 +141,17 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (candidates, skipped_candidates) =
|
||||
match list_local_openai_chat_candidates(state, &input, false).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "local_openai_chat_scheduler_selection_failed",
|
||||
log_type = "event",
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision scheduler selection failed"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, false,
|
||||
)
|
||||
.await;
|
||||
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
false,
|
||||
);
|
||||
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
@@ -181,7 +161,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
|
||||
attempt,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
"openai_chat_sync_success",
|
||||
false,
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -212,32 +192,17 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (candidates, skipped_candidates) =
|
||||
match list_local_openai_chat_candidates(state, &input, true).await {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "local_openai_chat_scheduler_selection_failed",
|
||||
log_type = "event",
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision scheduler selection failed"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, true,
|
||||
)
|
||||
.await;
|
||||
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
true,
|
||||
);
|
||||
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
@@ -247,7 +212,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
true,
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -22,7 +22,7 @@ pub(super) use self::sync::{
|
||||
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
|
||||
};
|
||||
|
||||
fn openai_chat_upstream_is_stream_for_candidate(
|
||||
pub(super) fn openai_chat_upstream_is_stream_for_candidate(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
client_is_stream: bool,
|
||||
|
||||
@@ -3,13 +3,10 @@ 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,
|
||||
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
|
||||
LocalOpenAiChatDecisionInput,
|
||||
};
|
||||
use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
@@ -176,27 +173,12 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
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(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() && skipped_candidates.is_empty() {
|
||||
let Some((mut attempt_source, candidate_count)) =
|
||||
build_local_openai_chat_stream_attempt_source(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
@@ -206,59 +188,13 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
0,
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
candidates.len() + skipped_candidates.len(),
|
||||
);
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
true,
|
||||
);
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
"openai_chat_stream_success",
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_stream_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat stream decision plan build failed"
|
||||
);
|
||||
}
|
||||
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
|
||||
plans.push(attempt);
|
||||
if plans.len() >= candidate_count {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -3,13 +3,10 @@ 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,
|
||||
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
|
||||
LocalOpenAiChatDecisionInput,
|
||||
};
|
||||
use super::candidates::list_local_openai_chat_candidates;
|
||||
use super::diagnostic::{
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
|
||||
};
|
||||
@@ -176,27 +173,11 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
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(Vec::new());
|
||||
}
|
||||
};
|
||||
if candidates.is_empty() && skipped_candidates.is_empty() {
|
||||
let Some((mut attempt_source, candidate_count)) = build_local_openai_chat_sync_attempt_source(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
@@ -206,59 +187,13 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
0,
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
set_local_openai_chat_candidate_evaluation_diagnostic(
|
||||
state,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
Some(input.requested_model.as_str()),
|
||||
candidates.len() + skipped_candidates.len(),
|
||||
);
|
||||
|
||||
let attempts = materialize_local_openai_chat_candidate_attempts(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
candidates,
|
||||
skipped_candidates,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
|
||||
&attempt.eligible.transport,
|
||||
attempt.eligible.provider_api_format.as_str(),
|
||||
false,
|
||||
);
|
||||
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
&input,
|
||||
attempt,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
"openai_chat_sync_success",
|
||||
upstream_is_stream,
|
||||
)
|
||||
.await
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match build_openai_chat_sync_plan_from_decision(parts, body_json, payload) {
|
||||
Ok(Some(value)) => plans.push(value),
|
||||
Ok(None) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat sync decision plan build failed"
|
||||
);
|
||||
}
|
||||
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
|
||||
plans.push(attempt);
|
||||
if plans.len() >= candidate_count {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ mod decision;
|
||||
mod plans;
|
||||
|
||||
use self::decision::{
|
||||
materialize_local_openai_responses_candidate_attempts,
|
||||
build_local_openai_responses_candidate_attempt_source,
|
||||
maybe_build_local_openai_responses_decision_payload_for_candidate,
|
||||
resolve_local_openai_responses_decision_input,
|
||||
};
|
||||
@@ -108,12 +108,12 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
|
||||
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
@@ -146,12 +146,12 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
|
||||
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
|
||||
@@ -3,7 +3,6 @@ use tracing::warn;
|
||||
|
||||
use super::decision::{
|
||||
build_local_openai_responses_candidate_attempt_source,
|
||||
materialize_local_openai_responses_candidate_attempts,
|
||||
maybe_build_local_openai_responses_decision_payload_for_candidate,
|
||||
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
|
||||
LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
|
||||
@@ -312,7 +311,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
|
||||
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
|
||||
let (mut source, candidate_count) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
@@ -322,7 +321,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
@@ -384,7 +383,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
|
||||
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
|
||||
let (mut source, candidate_count) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
.await?;
|
||||
@@ -394,7 +393,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
|
||||
}
|
||||
|
||||
let mut plans = Vec::new();
|
||||
for attempt in attempts {
|
||||
while let Some(attempt) = source.next_attempt().await {
|
||||
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user