i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 17:30:23 +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

119
apps/aether-gateway/src/executor/candidate_loop.rs
View File

@@ -11,6 +11,7 @@ use axum::http::Response;
use tokio::time::{timeout, Duration};
use tracing::{debug, warn, Instrument};
use crate::ai_serving::LocalExecutionAttemptSource;
use crate::clock::current_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{execute_execution_runtime_stream, execute_execution_runtime_sync};
@@ -80,6 +81,42 @@ where
.await
}
pub(crate) async fn execute_sync_attempt_source<T, S>(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
mut source: S,
) -> Result<LocalExecutionRequestOutcome, GatewayError>
where
T: AiExecutionAttempt + Send + Sync + 'static,
S: LocalExecutionAttemptSource<T>,
{
let span = tracing::debug_span!("candidates", trace_id = %trace_id, plan_kind);
async move {
tracing::debug!(
event_name = "candidate_loop_started",
log_type = "event",
trace_id = %trace_id,
plan_kind,
"dynamic candidate loop started"
);
let port = SyncAttemptLoopPort {
state,
parts,
trace_id,
decision,
plan_kind,
};
run_dynamic_attempt_loop(&port, &mut source).await
}
.instrument(span)
.await
}
struct SyncAttemptLoopPort<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
@@ -182,6 +219,75 @@ where
.await
}
pub(crate) async fn execute_stream_attempt_source<T, S>(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
mut source: S,
) -> Result<LocalExecutionRequestOutcome, GatewayError>
where
T: AiExecutionAttempt + Send + Sync + 'static,
S: LocalExecutionAttemptSource<T>,
{
let span = tracing::debug_span!("candidates", trace_id = %trace_id, plan_kind);
async move {
tracing::debug!(
event_name = "candidate_loop_started",
log_type = "event",
trace_id = %trace_id,
plan_kind,
"dynamic candidate loop started"
);
let port = StreamAttemptLoopPort {
state,
trace_id,
decision,
plan_kind,
};
run_dynamic_attempt_loop(&port, &mut source).await
}
.instrument(span)
.await
}
async fn run_dynamic_attempt_loop<Port, Source, Attempt>(
port: &Port,
source: &mut Source,
) -> Result<LocalExecutionRequestOutcome, GatewayError>
where
Port: AiAttemptLoopPort<
Attempt,
Response = Response<Body>,
Exhaustion = crate::executor::LocalExecutionExhaustion,
Error = GatewayError,
>,
Source: LocalExecutionAttemptSource<Attempt>,
Attempt: AiExecutionAttempt + Send + Sync + 'static,
{
let mut last_attempted = None;
while let Some(attempt) = source.next_execution_attempt().await? {
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
if let Some(response) = port.execute_attempt(&attempt).await? {
let remaining = source.drain_execution_attempts().await?;
port.mark_unused_attempts(remaining).await?;
return Ok(LocalExecutionRequestOutcome::responded(response));
}
}
let Some((last_plan, last_report_context)) = last_attempted else {
return Ok(LocalExecutionRequestOutcome::NoPath);
};
Ok(LocalExecutionRequestOutcome::Exhausted(
port.build_exhaustion(last_plan, last_report_context)
.await?,
))
}
struct StreamAttemptLoopPort<'a> {
state: &'a AppState,
trace_id: &'a str,
@@ -307,10 +413,7 @@ where
fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool {
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
metadata.candidate_group_id.is_some()
&& metadata
.pool_key_index
.is_some_and(|pool_key_index| pool_key_index > 0)
metadata.candidate_group_id.is_some() && metadata.pool_key_index.is_some()
}
fn resolve_stream_candidate_watchdog_timeout(plan: &aether_contracts::ExecutionPlan) -> Duration {
@@ -520,13 +623,13 @@ mod tests {
#[test]
fn unused_persistence_skips_pool_internal_candidates() {
assert!(should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group",
"pool_key_index": 1,
}))));
assert!(!should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group",
"pool_key_index": 0,
}))));
assert!(should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group",
"pool_key_index": 1,
}))));
assert!(!should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group",
}))));

249
apps/aether-gateway/src/executor/orchestration.rs
View File

@@ -1,24 +1,29 @@
use crate::ai_serving::api::{
build_local_gemini_files_stream_plan_and_reports_for_kind,
build_local_gemini_files_sync_plan_and_reports_for_kind,
build_local_image_stream_plan_and_reports_for_kind,
build_local_image_sync_plan_and_reports_for_kind,
build_local_gemini_files_stream_attempt_source_for_kind,
build_local_gemini_files_sync_attempt_source_for_kind,
build_local_image_stream_attempt_source_for_kind,
build_local_image_sync_attempt_source_for_kind,
build_local_openai_chat_stream_attempt_source_for_kind,
build_local_openai_chat_stream_plan_and_reports_for_kind,
build_local_openai_chat_sync_attempt_source_for_kind,
build_local_openai_chat_sync_plan_and_reports_for_kind,
build_local_openai_responses_stream_attempt_source_for_kind,
build_local_openai_responses_stream_plan_and_reports_for_kind,
build_local_openai_responses_sync_attempt_source_for_kind,
build_local_openai_responses_sync_plan_and_reports_for_kind,
build_local_same_format_stream_plan_and_reports, build_local_same_format_sync_plan_and_reports,
build_local_video_sync_plan_and_reports_for_kind,
build_standard_family_stream_plan_and_reports, build_standard_family_sync_plan_and_reports,
parse_direct_request_body, resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_gemini_stream_spec, resolve_gemini_sync_spec, resolve_local_same_format_stream_spec,
build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports,
build_local_same_format_sync_attempt_source, build_local_same_format_sync_plan_and_reports,
build_local_video_sync_attempt_source_for_kind, build_standard_family_stream_attempt_source,
build_standard_family_sync_attempt_source, parse_direct_request_body,
resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_gemini_stream_spec,
resolve_gemini_sync_spec, resolve_local_same_format_stream_spec,
resolve_local_same_format_sync_spec, set_local_openai_chat_execution_exhausted_diagnostic,
AiStreamAttempt, AiSyncAttempt, LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
};
use crate::control::GatewayControlDecision;
use crate::executor::candidate_loop::{
execute_stream_plan_and_reports, execute_sync_plan_and_reports,
execute_stream_attempt_source, execute_sync_attempt_source, execute_sync_plan_and_reports,
};
use crate::executor::LocalExecutionRequestOutcome;
use crate::{AiExecutionDecision, AppState, GatewayError};
@@ -52,28 +57,33 @@ pub(crate) async fn maybe_execute_sync_via_local_decision(
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports = build_local_openai_chat_sync_plan_and_reports_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
let Some((attempt_source, candidate_count)) =
build_local_openai_chat_sync_attempt_source_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
let plan_count = plan_and_reports.len();
let outcome = execute_sync_plan_and_reports(
let outcome = execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
plan_and_reports,
attempt_source,
)
.await?;
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
set_local_openai_chat_execution_exhausted_diagnostic(
state, trace_id, decision, plan_kind, body_json, plan_count,
state,
trace_id,
decision,
plan_kind,
body_json,
candidate_count,
);
}
@@ -88,22 +98,32 @@ pub(crate) async fn maybe_execute_stream_via_local_decision(
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports = build_local_openai_chat_stream_plan_and_reports_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
let Some((attempt_source, candidate_count)) =
build_local_openai_chat_stream_attempt_source_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let outcome = execute_stream_attempt_source::<AiStreamAttempt, _>(
state,
trace_id,
decision,
plan_kind,
attempt_source,
)
.await?;
if plan_and_reports.is_empty() {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
let plan_count = plan_and_reports.len();
let outcome =
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports)
.await?;
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
set_local_openai_chat_execution_exhausted_diagnostic(
state, trace_id, decision, plan_kind, body_json, plan_count,
state,
trace_id,
decision,
plan_kind,
body_json,
candidate_count,
);
}
@@ -118,22 +138,22 @@ pub(crate) async fn maybe_execute_sync_via_local_openai_responses_decision(
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<AiSyncAttempt> =
build_local_openai_responses_sync_plan_and_reports_for_kind(
let Some((attempt_source, _candidate_count)) =
build_local_openai_responses_sync_attempt_source_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_sync_plan_and_reports(
execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
plan_and_reports,
attempt_source,
)
.await
}
@@ -146,16 +166,23 @@ pub(crate) async fn maybe_execute_stream_via_local_openai_responses_decision(
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<AiStreamAttempt> =
build_local_openai_responses_stream_plan_and_reports_for_kind(
let Some((attempt_source, _candidate_count)) =
build_local_openai_responses_stream_attempt_source_for_kind(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
execute_stream_attempt_source::<AiStreamAttempt, _>(
state,
trace_id,
decision,
plan_kind,
attempt_source,
)
.await
}
pub(crate) async fn maybe_execute_sync_via_standard_family_decision(
@@ -171,21 +198,21 @@ pub(crate) async fn maybe_execute_sync_via_standard_family_decision(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<AiSyncAttempt> = build_standard_family_sync_plan_and_reports(
let Some((attempt_source, _candidate_count)) = build_standard_family_sync_attempt_source(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_sync_plan_and_reports(
execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
plan_and_reports,
attempt_source,
)
.await
}
@@ -203,15 +230,22 @@ pub(crate) async fn maybe_execute_stream_via_standard_family_decision(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<AiStreamAttempt> = build_standard_family_stream_plan_and_reports(
let Some((attempt_source, _candidate_count)) = build_standard_family_stream_attempt_source(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
execute_stream_attempt_source::<AiStreamAttempt, _>(
state,
trace_id,
decision,
plan_kind,
attempt_source,
)
.await
}
pub(crate) async fn maybe_execute_sync_via_local_standard_decision(
@@ -328,21 +362,21 @@ pub(crate) async fn maybe_execute_sync_via_local_same_format_provider_decision(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<AiSyncAttempt> = build_local_same_format_sync_plan_and_reports(
let Some((attempt_source, _candidate_count)) = build_local_same_format_sync_attempt_source(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_sync_plan_and_reports(
execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
plan_and_reports,
attempt_source,
)
.await
}
@@ -359,15 +393,22 @@ pub(crate) async fn maybe_execute_stream_via_local_same_format_provider_decision
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let plan_and_reports: Vec<AiStreamAttempt> = build_local_same_format_stream_plan_and_reports(
let Some((attempt_source, _candidate_count)) = build_local_same_format_stream_attempt_source(
state, parts, trace_id, decision, body_json, spec,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
execute_stream_attempt_source::<AiStreamAttempt, _>(
state,
trace_id,
decision,
plan_kind,
attempt_source,
)
.await
}
pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
@@ -380,8 +421,8 @@ pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<AiSyncAttempt> =
build_local_gemini_files_sync_plan_and_reports_for_kind(
let Some((attempt_source, _candidate_count)) =
build_local_gemini_files_sync_attempt_source_for_kind(
state,
parts,
body_json,
@@ -391,18 +432,18 @@ pub(crate) async fn maybe_execute_sync_via_local_gemini_files_decision(
decision,
plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_sync_plan_and_reports(
execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
plan_and_reports,
attempt_source,
)
.await
}
@@ -416,7 +457,7 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<AiSyncAttempt> = build_local_image_sync_plan_and_reports_for_kind(
let Some((attempt_source, _candidate_count)) = build_local_image_sync_attempt_source_for_kind(
state,
parts,
body_json,
@@ -425,18 +466,18 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
decision,
plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_sync_plan_and_reports(
execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
plan_and_reports,
attempt_source,
)
.await
}
@@ -448,16 +489,23 @@ pub(crate) async fn maybe_execute_stream_via_local_gemini_files_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<AiStreamAttempt> =
build_local_gemini_files_stream_plan_and_reports_for_kind(
let Some((attempt_source, _candidate_count)) =
build_local_gemini_files_stream_attempt_source_for_kind(
state, parts, trace_id, decision, plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
execute_stream_attempt_source::<AiStreamAttempt, _>(
state,
trace_id,
decision,
plan_kind,
attempt_source,
)
.await
}
pub(crate) async fn maybe_execute_stream_via_local_image_decision(
@@ -469,8 +517,8 @@ pub(crate) async fn maybe_execute_stream_via_local_image_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<AiStreamAttempt> =
build_local_image_stream_plan_and_reports_for_kind(
let Some((attempt_source, _candidate_count)) =
build_local_image_stream_attempt_source_for_kind(
state,
parts,
body_json,
@@ -479,12 +527,19 @@ pub(crate) async fn maybe_execute_stream_via_local_image_decision(
decision,
plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_stream_plan_and_reports(state, trace_id, decision, plan_kind, plan_and_reports).await
execute_stream_attempt_source::<AiStreamAttempt, _>(
state,
trace_id,
decision,
plan_kind,
attempt_source,
)
.await
}
pub(crate) async fn maybe_execute_sync_via_local_video_decision(
@@ -495,21 +550,21 @@ pub(crate) async fn maybe_execute_sync_via_local_video_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let plan_and_reports: Vec<AiSyncAttempt> = build_local_video_sync_plan_and_reports_for_kind(
let Some((attempt_source, _candidate_count)) = build_local_video_sync_attempt_source_for_kind(
state, parts, body_json, trace_id, decision, plan_kind,
)
.await?;
if plan_and_reports.is_empty() {
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
}
};
execute_sync_plan_and_reports(
execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
decision,
plan_kind,
plan_and_reports,
attempt_source,
)
.await
}