mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-11 19:59:50 +08:00
Lazy load requested model candidates
This commit is contained in:
@@ -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,
|
||||
|
||||
+63
-70
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user