refactor: isolate dispatch scheduling core

This commit is contained in:
fawney19
2026-05-12 13:15:41 +08:00
parent 81ee27cdea
commit fa73655134
50 changed files with 4467 additions and 3202 deletions
@@ -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,
)