Lazy load requested model candidates

This commit is contained in:
fawney19
2026-05-03 20:56:31 +08:00
parent a24e4a793d
commit 6f00cabe96
13 changed files with 1047 additions and 225 deletions
@@ -7,6 +7,7 @@ mod support;
pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for_candidate;
pub(super) use self::support::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
@@ -1,6 +1,7 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_serving::planner::candidate_materialization::{
build_lazy_requested_model_execution_candidate_attempt_source_with_serving,
build_local_execution_candidate_attempt_source_with_serving,
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -13,6 +14,7 @@ use crate::ai_serving::planner::candidate_metadata::{
LocalExecutionCandidateMetadataParts,
};
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::candidate_source::LocalCandidatePreselectionKeyMode;
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
@@ -269,3 +271,80 @@ pub(crate) async fn build_local_openai_chat_candidate_attempt_source<'a>(
)
.await
}
pub(crate) async fn build_lazy_local_openai_chat_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalOpenAiChatDecisionInput,
body_json: &serde_json::Value,
require_streaming: bool,
) -> (LocalOpenAiChatCandidateAttemptSource<'a>, usize) {
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let auth_context: &ExecutionRuntimeAuthContext = &input.auth_context;
let persistence_policy = build_local_candidate_persistence_policy(
auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiChatDecision,
);
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
"openai:chat",
&input.requested_model,
require_streaming,
&input.auth_snapshot,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
LocalCandidateResolutionMode::Standard,
move |eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: "openai:chat",
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.trim(),
))
},
move |mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
"openai:chat",
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await
}
@@ -12,6 +12,7 @@ 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,
@@ -2,6 +2,7 @@ use async_trait::async_trait;
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,
@@ -49,27 +50,10 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
return Ok(None);
};
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(None);
}
};
let candidate_count = candidates.len() + skipped_candidates.len();
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, true,
)
.await;
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
@@ -90,16 +74,6 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
candidate_count,
);
let (candidates, candidate_count) = build_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
Ok(Some((
LocalOpenAiChatStreamAttemptSource {
state,
@@ -2,6 +2,7 @@ use async_trait::async_trait;
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,
@@ -58,27 +59,10 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
return Ok(None);
};
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(None);
}
};
let candidate_count = candidates.len() + skipped_candidates.len();
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, false,
)
.await;
if candidate_count == 0 {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
@@ -99,16 +83,6 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
candidate_count,
);
let (candidates, candidate_count) = build_local_openai_chat_candidate_attempt_source(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
Ok(Some((
LocalOpenAiChatSyncAttemptSource {
state,
@@ -2,6 +2,7 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::{
build_lazy_requested_model_execution_candidate_attempt_source_with_serving,
build_local_execution_candidate_attempt_source_with_serving,
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -240,80 +241,72 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::OpenAiResponsesDecision,
);
let preselection = preselect_local_execution_candidates_with_serving(
planner_state,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
&input.auth_snapshot,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
)
.await?;
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
preselection.candidates,
preselection.skipped_candidates,
LocalCandidateResolutionMode::Standard,
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
eligible.candidate.endpoint_api_format.as_str(),
))
},
|mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
Ok(
build_lazy_requested_model_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
&input.auth_snapshot,
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
LocalCandidateResolutionMode::Standard,
move |eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
serde_json::Map::new(),
&provider_api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
eligible.candidate.endpoint_api_format.as_str(),
))
},
move |mut skipped_candidate| {
let provider_api_format = skipped_candidate
.transport
.as_ref()
.map(|transport| transport.endpoint.api_format.trim().to_ascii_lowercase())
.unwrap_or_else(|| {
skipped_candidate
.candidate
.endpoint_api_format
.trim()
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
&provider_api_format,
);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
spec_metadata.api_format,
serde_json::Map::new(),
execution_strategy,
conversion_mode,
provider_api_format.as_str(),
),
);
skipped_candidate
},
)
.await,
)
.await)
}
pub(crate) async fn mark_skipped_local_openai_responses_candidate(