i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 09:50:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor pool candidate scheduling

Browse Source
This commit is contained in:
fawney19
2026-05-03 20:14:29 +08:00
parent 8ebee9922c
commit a24e4a793d
55 changed files with 4825 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/ai_serving/planner/passthrough/mod.rs
View File

@@ -3,7 +3,9 @@
mod provider;
pub(crate) use self::provider::{
build_local_stream_attempt_source as build_local_same_format_stream_attempt_source,
build_local_stream_plan_and_reports as build_local_same_format_stream_plan_and_reports,
build_local_sync_attempt_source as build_local_same_format_sync_attempt_source,
build_local_sync_plan_and_reports as build_local_same_format_sync_plan_and_reports,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
maybe_build_stream_local_same_format_provider_decision_payload,

9
apps/aether-gateway/src/ai_serving/planner/passthrough/provider.rs
View File

@@ -62,17 +62,20 @@ mod plans;
mod request;
pub(crate) use self::family::{
build_local_same_format_provider_candidate_attempt_source,
materialize_local_same_format_provider_candidate_attempts,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
resolve_local_same_format_provider_decision_input, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec,
resolve_local_same_format_provider_decision_input, LocalSameFormatProviderCandidateAttempt,
LocalSameFormatProviderCandidateAttemptSource, LocalSameFormatProviderDecisionInput,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
};
pub(crate) use self::family::{
maybe_build_stream_local_same_format_provider_decision_payload,
maybe_build_sync_local_same_format_provider_decision_payload,
};
pub(crate) use self::plans::{
build_local_stream_plan_and_reports, build_local_sync_plan_and_reports,
build_local_stream_attempt_source, build_local_stream_plan_and_reports,
build_local_sync_attempt_source, build_local_sync_plan_and_reports,
};
const ANTIGRAVITY_ENVELOPE_NAME: &str = "antigravity:v1internal";

95
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/candidates.rs
View File

@@ -1,7 +1,9 @@
use tracing::warn;
use crate::ai_serving::planner::candidate_materialization::{
build_local_execution_candidate_attempt_source_with_serving,
materialize_local_execution_candidates_with_serving, LocalCandidateResolutionMode,
LocalExecutionCandidateAttemptSource,
};
use crate::ai_serving::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata,
@@ -169,3 +171,96 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
Ok((outcome.attempts, outcome.candidate_count))
}
pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a>(
state: &'a AppState,
trace_id: &str,
input: &LocalSameFormatProviderDecisionInput,
body_json: &serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Result<(LocalExecutionCandidateAttemptSource<'a>, usize), GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let sticky_session_token = extract_pool_sticky_session_token(body_json);
let persistence_policy = build_local_candidate_persistence_policy(
&input.auth_context,
input.required_capabilities.as_ref(),
LocalCandidatePersistencePolicyKind::SameFormatProviderDecision,
);
let (candidates, preselection_skipped) = planner_state
.list_selectable_candidates_with_skip_reasons(
spec_metadata.api_format,
&input.requested_model,
spec_metadata.require_streaming,
input.required_capabilities.as_ref(),
Some(&input.auth_snapshot),
current_unix_secs(),
)
.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,
candidates,
preselection_skipped
.into_iter()
.map(|item| SkippedLocalExecutionCandidate {
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),
LocalCandidateResolutionMode::Standard,
|eligible| {
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
spec_metadata.api_format,
);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
extra_fields: serde_json::Map::new(),
},
execution_strategy,
conversion_mode,
spec_metadata.api_format,
))
},
|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(|| spec_metadata.api_format.to_string());
let (execution_strategy, conversion_mode) = ai_local_execution_contract_for_formats(
spec_metadata.api_format,
provider_api_format.as_str(),
);
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)
}

2
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/mod.rs
View File

@@ -8,10 +8,12 @@ pub(crate) use self::build::{
maybe_build_sync_local_same_format_provider_decision_payload,
};
pub(crate) use self::candidates::{
build_local_same_format_provider_candidate_attempt_source,
materialize_local_same_format_provider_candidate_attempts,
resolve_local_same_format_provider_decision_input,
};
pub(crate) use self::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalSameFormatProviderCandidateAttempt;
pub(crate) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalSameFormatProviderCandidateAttemptSource;
pub(crate) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalSameFormatProviderDecisionInput;
pub(crate) use crate::ai_serving::{LocalSameFormatProviderFamily, LocalSameFormatProviderSpec};

295
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs
View File

@@ -1,6 +1,10 @@
use async_trait::async_trait;
use tracing::warn;
use crate::ai_serving::planner::common::extract_requested_model_from_request;
use crate::ai_serving::planner::candidate_materialization::LocalExecutionAttemptSource;
use crate::ai_serving::planner::common::{
extract_requested_model_from_request, RequestedModelFamily,
};
use crate::ai_serving::planner::runtime_miss::{
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal,
apply_local_runtime_candidate_terminal_reason, set_local_runtime_miss_diagnostic_reason,
@@ -15,12 +19,299 @@ pub(crate) use crate::ai_serving::{
};
use super::{
build_local_same_format_provider_candidate_attempt_source,
materialize_local_same_format_provider_candidate_attempts,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
resolve_local_same_format_provider_decision_input, AiStreamAttempt, AiSyncAttempt, AppState,
GatewayControlDecision, GatewayError, LocalSameFormatProviderSpec,
GatewayControlDecision, GatewayError, LocalSameFormatProviderCandidateAttempt,
LocalSameFormatProviderCandidateAttemptSource, LocalSameFormatProviderDecisionInput,
LocalSameFormatProviderSpec,
};
pub(crate) struct LocalSameFormatProviderSyncAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
input: LocalSameFormatProviderDecisionInput,
spec: LocalSameFormatProviderSpec,
requested_model_family: RequestedModelFamily,
candidates: LocalSameFormatProviderCandidateAttemptSource<'a>,
}
pub(crate) struct LocalSameFormatProviderStreamAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
input: LocalSameFormatProviderDecisionInput,
spec: LocalSameFormatProviderSpec,
requested_model_family: RequestedModelFamily,
candidates: LocalSameFormatProviderCandidateAttemptSource<'a>,
}
pub(crate) async fn build_local_sync_attempt_source<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Result<Option<(LocalSameFormatProviderSyncAttemptSource<'a>, usize)>, GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("same-format provider spec metadata should include requested-model family");
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
else {
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(None);
};
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (candidates, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
candidate_count,
);
if candidate_count == 0 {
return Ok(None);
}
Ok(Some((
LocalSameFormatProviderSyncAttemptSource {
state,
parts,
trace_id,
body_json,
input,
spec,
requested_model_family,
candidates,
},
candidate_count,
)))
}
pub(crate) async fn build_local_stream_attempt_source<'a>(
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
decision: &'a GatewayControlDecision,
body_json: &'a serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Result<Option<(LocalSameFormatProviderStreamAttemptSource<'a>, usize)>, GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let requested_model_family = spec_metadata
.requested_model_family
.expect("same-format provider spec metadata should include requested-model family");
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
else {
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
decision,
spec_metadata.decision_kind,
extract_requested_model_from_request(parts, body_json, requested_model_family)
.as_deref(),
"decision_input_unavailable",
);
return Ok(None);
};
set_local_runtime_miss_diagnostic_reason(
state,
trace_id,
decision,
spec_metadata.decision_kind,
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (candidates, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
state,
trace_id,
candidate_count,
);
if candidate_count == 0 {
return Ok(None);
}
Ok(Some((
LocalSameFormatProviderStreamAttemptSource {
state,
parts,
trace_id,
body_json,
input,
spec,
requested_model_family,
candidates,
},
candidate_count,
)))
}
#[async_trait]
impl LocalExecutionAttemptSource<AiSyncAttempt> for LocalSameFormatProviderSyncAttemptSource<'_> {
async fn next_execution_attempt(&mut self) -> Result<Option<AiSyncAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await {
match self.build_sync_attempt(attempt).await? {
Some(attempt) => return Ok(Some(attempt)),
None => continue,
}
}
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_sync_plans",
);
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiSyncAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_sync_attempt(attempt).await? {
drained.push(attempt);
}
}
Ok(drained)
}
}
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt>
for LocalSameFormatProviderStreamAttemptSource<'_>
{
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
while let Some(attempt) = self.candidates.next_attempt().await {
match self.build_stream_attempt(attempt).await? {
Some(attempt) => return Ok(Some(attempt)),
None => continue,
}
}
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_stream_plans",
);
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiStreamAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_stream_attempt(attempt).await? {
drained.push(attempt);
}
}
Ok(drained)
}
}
impl LocalSameFormatProviderSyncAttemptSource<'_> {
async fn build_sync_attempt(
&self,
attempt: LocalSameFormatProviderCandidateAttempt,
) -> Result<Option<AiSyncAttempt>, GatewayError> {
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.input,
attempt,
self.spec,
)
.await
else {
return Ok(None);
};
match build_sync_plan_from_requested_model_family(
self.requested_model_family,
self.parts,
self.body_json,
payload,
) {
Ok(value) => Ok(value),
Err(err) => {
warn!(
trace_id = %self.trace_id,
error = ?err,
"gateway local same-format sync decision plan build failed"
);
Ok(None)
}
}
}
}
impl LocalSameFormatProviderStreamAttemptSource<'_> {
async fn build_stream_attempt(
&self,
attempt: LocalSameFormatProviderCandidateAttempt,
) -> Result<Option<AiStreamAttempt>, GatewayError> {
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.input,
attempt,
self.spec,
)
.await
else {
return Ok(None);
};
match build_stream_plan_from_requested_model_family(
self.requested_model_family,
self.parts,
self.body_json,
payload,
) {
Ok(value) => Ok(value),
Err(err) => {
warn!(
trace_id = %self.trace_id,
error = ?err,
"gateway local same-format stream decision plan build failed"
);
Ok(None)
}
}
}
}
pub(crate) async fn build_local_sync_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,