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

refactor: lazy pool key scheduling

Browse Source
This commit is contained in:
fawney19
2026-05-11 00:12:05 +08:00
parent 1a0f1a7b72
commit bacb14e5f0
26 changed files with 1856 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

171
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
View File

@@ -3,7 +3,7 @@ use aether_ai_serving::{
ai_should_persist_skipped_candidate_for_pool_membership,
run_ai_available_candidate_persistence, run_ai_candidate_materialization,
run_ai_skipped_candidate_persistence, AiAvailableCandidatePersistencePort,
AiCandidateMaterializationOutcome, AiCandidateMaterializationPort, AiCandidateResolutionMode,
AiCandidateMaterializationOutcome, AiCandidateMaterializationPort,
AiSkippedCandidatePersistencePort,
};
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
@@ -17,8 +17,6 @@ use uuid::Uuid;
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
use crate::ai_serving::planner::candidate_resolution::{
resolve_and_rank_local_execution_candidates,
resolve_and_rank_local_execution_candidates_without_transport_pair_gate,
resolve_and_rank_logical_local_execution_candidates, EligibleLocalExecutionCandidate,
LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
};
@@ -225,37 +223,19 @@ where
&self,
candidates: Vec<Self::Candidate>,
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
let requested_model = self.requested_model.map(str::to_string);
let resolved = match self.resolution_mode {
AiCandidateResolutionMode::Standard => {
resolve_and_rank_local_execution_candidates(
self.state,
candidates,
self.client_api_format,
requested_model.as_deref().unwrap_or_default(),
self.auth_snapshot,
self.client_session_affinity,
self.required_capabilities,
self.sticky_session_token,
self.request_auth_channel,
)
.await
}
AiCandidateResolutionMode::WithoutTransportPairGate => {
resolve_and_rank_local_execution_candidates_without_transport_pair_gate(
self.state,
candidates,
self.client_api_format,
requested_model.as_deref(),
self.auth_snapshot,
self.client_session_affinity,
self.required_capabilities,
self.sticky_session_token,
self.request_auth_channel,
)
.await
}
};
let resolved = resolve_and_rank_logical_local_execution_candidates(
self.state,
candidates,
self.client_api_format,
self.requested_model,
self.auth_snapshot,
self.client_session_affinity,
self.required_capabilities,
self.sticky_session_token,
self.request_auth_channel,
self.resolution_mode,
)
.await;
Ok(resolved)
}
@@ -278,11 +258,14 @@ where
&self,
candidates: Vec<Self::Eligible>,
) -> Result<Vec<Self::Attempt>, Self::Error> {
Ok(persist_available_local_execution_candidates_with_context(
Ok(materialize_logical_local_execution_candidate_attempts(
self.state,
self.trace_id,
self.persistence_policy.available,
candidates,
self.sticky_session_token,
self.requested_model,
self.request_auth_channel,
&self.build_available_extra_data,
)
.await)
@@ -491,6 +474,7 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
{
let _ = build_available_extra_data;
let (candidates, resolved_skipped) = resolve_and_rank_logical_local_execution_candidates(
state,
candidates,
@@ -901,6 +885,64 @@ where
.await
}
#[allow(clippy::too_many_arguments)]
async fn materialize_logical_local_execution_candidate_attempts<F>(
state: PlannerAppState<'_>,
trace_id: &str,
context: LocalAvailableCandidatePersistenceContext<'_>,
candidates: Vec<EligibleLocalExecutionCandidate>,
sticky_session_token: Option<&str>,
requested_model: Option<&str>,
request_auth_channel: Option<&str>,
build_extra_data: &F,
) -> Vec<LocalExecutionCandidateAttempt>
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
{
let mut attempts = Vec::new();
for (candidate_index, candidate) in candidates.into_iter().enumerate() {
let candidate_index = u32::try_from(candidate_index).unwrap_or(u32::MAX);
match candidate.kind {
LocalExecutionCandidateKind::SingleKey => {
attempts.extend(
persist_available_local_execution_candidate_at_index(
state,
trace_id,
context,
candidate,
candidate_index,
build_extra_data,
)
.await,
);
}
LocalExecutionCandidateKind::PoolGroup => {
let mut cursor = PoolKeyCursor::new(
state,
candidate,
sticky_session_token,
requested_model,
request_auth_channel,
);
let attempt_count_before_pool = attempts.len();
while let Some(candidate) = cursor.next_key().await {
attempts.extend(build_unpersisted_local_execution_candidate_attempts(
candidate,
candidate_index,
));
}
let _ = cursor.take_skipped_candidates();
if attempts.len() == attempt_count_before_pool {
cursor.log_exhausted();
}
}
}
}
attempts
}
async fn persist_available_local_execution_candidate_at_index<F>(
state: PlannerAppState<'_>,
trace_id: &str,
@@ -1167,7 +1209,10 @@ mod tests {
use std::collections::VecDeque;
use std::sync::Arc;
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
@@ -1278,6 +1323,7 @@ mod tests {
orchestration: LocalExecutionCandidateMetadata {
candidate_group_id: pool_key_index.map(|_| "pool-group".to_string()),
pool_key_index,
pool_key_lease: None,
},
ranking: None,
}
@@ -1320,6 +1366,61 @@ mod tests {
assert_eq!(stored[0].candidate_index, 2);
}
#[tokio::test]
async fn logical_materialization_does_not_persist_pool_group_representative() {
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
Arc::new(InMemoryAuthApiKeySnapshotRepository::default()),
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::default()),
Arc::new(InMemoryProviderCatalogReadRepository::seed(
Vec::new(),
Vec::new(),
Vec::new(),
)),
Arc::clone(&request_candidate_repository),
"test-encryption-key",
),
);
let mut pool_group = sample_eligible("pool-group", None);
pool_group.kind = LocalExecutionCandidateKind::PoolGroup;
pool_group.transport = sample_transport(
"pool-group",
Some(json!({ "pool_advanced": { "scheduling_presets": [] } })),
);
let attempts = materialize_logical_local_execution_candidate_attempts(
PlannerAppState::new(&app),
"trace-logical-pool",
LocalAvailableCandidatePersistenceContext {
user_id: "user-1",
api_key_id: "api-key-1",
required_capabilities: None,
error_context: "persist should not fail",
},
vec![pool_group, sample_eligible("normal-key", None)],
None,
Some("gpt-5"),
None,
&|_| None,
)
.await;
assert_eq!(attempts.len(), 1);
assert_eq!(attempts[0].candidate_index, 1);
assert_eq!(attempts[0].eligible.candidate.key_id, "normal-key");
let stored = app
.read_request_candidates_by_request_id("trace-logical-pool")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 1);
}
#[test]
fn pool_key_attempts_use_distinct_effective_retry_indices() {
let first = build_unpersisted_local_execution_candidate_attempts(

1
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs
View File

@@ -105,6 +105,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),

1052
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
View File

File diff suppressed because it is too large Load Diff

8
apps/aether-gateway/src/ai_serving/planner/report_context.rs
View File

@@ -6,6 +6,7 @@ use aether_ai_serving::{
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
AiExecutionReportContextParts, AiRequestOrigin,
};
use aether_runtime_state::RuntimeLockLease;
use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
use serde_json::{Map, Value};
@@ -17,7 +18,7 @@ use crate::ai_serving::{
use crate::client_session_affinity::{
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
use crate::orchestration::ExecutionAttemptIdentity;
use crate::orchestration::{insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity};
pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
@@ -37,6 +38,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) client_api_format: &'a str,
pub(crate) mapped_model: Option<&'a str>,
pub(crate) candidate_group_id: Option<&'a str>,
pub(crate) pool_key_lease: Option<&'a RuntimeLockLease>,
pub(crate) ranking: Option<&'a SchedulerRankingOutcome>,
pub(crate) upstream_url: Option<&'a str>,
pub(crate) header_rules: Option<&'a Value>,
@@ -86,6 +88,7 @@ pub(crate) fn build_local_execution_report_context(
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
insert_pool_key_lease_report_context_fields(&mut extra_fields, parts.pool_key_lease);
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
@@ -258,6 +261,7 @@ mod tests {
client_api_format: "openai:chat",
mapped_model: None,
candidate_group_id: None,
pool_key_lease: None,
ranking: None,
upstream_url: None,
header_rules: None,
@@ -337,6 +341,7 @@ mod tests {
client_api_format: "gemini:generate_content",
mapped_model: None,
candidate_group_id: None,
pool_key_lease: None,
ranking: None,
upstream_url: None,
header_rules: None,
@@ -405,6 +410,7 @@ mod tests {
client_api_format: "openai:chat",
mapped_model: None,
candidate_group_id: None,
pool_key_lease: None,
ranking: None,
upstream_url: None,
header_rules: None,

1
apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs
View File

@@ -87,6 +87,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
mapped_model: None,
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: None,
header_rules: transport.endpoint.header_rules.as_ref(),

1
apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs
View File

@@ -107,6 +107,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),

1
apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs
View File

@@ -68,6 +68,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: None,
header_rules: transport.endpoint.header_rules.as_ref(),

1
apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs
View File

@@ -113,6 +113,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),

1
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
View File

@@ -102,6 +102,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
client_api_format: "openai:chat",
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),

1
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs
View File

@@ -98,6 +98,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
pool_key_lease: eligible.orchestration.pool_key_lease.as_ref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),

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

@@ -16,6 +16,7 @@ use crate::clock::current_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{execute_execution_runtime_stream, execute_execution_runtime_sync};
use crate::executor::{build_local_execution_exhaustion, LocalExecutionRequestOutcome};
use crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease;
use crate::log_ids::short_request_id;
use crate::orchestration::local_execution_candidate_metadata_from_report_context;
use crate::request_candidate_runtime::{
@@ -403,7 +404,19 @@ where
{
for plan_and_report in remaining {
let report_context = plan_and_report.report_context();
if should_skip_unused_persistence(report_context.as_ref()) {
let metadata =
local_execution_candidate_metadata_from_report_context(report_context.as_ref());
if let Some(lease) = metadata.pool_key_lease.as_ref() {
if let Err(err) =
release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), lease).await
{
warn!(
error = ?err,
"gateway candidate loop: failed to release unused pool key lease"
);
}
}
if should_skip_unused_persistence_from_metadata(&metadata) {
continue;
}
record_local_request_candidate_status(
@@ -426,6 +439,12 @@ where
fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool {
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
should_skip_unused_persistence_from_metadata(&metadata)
}
fn should_skip_unused_persistence_from_metadata(
metadata: &crate::orchestration::LocalExecutionCandidateMetadata,
) -> bool {
metadata.candidate_group_id.is_some() && metadata.pool_key_index.is_some()
}

4
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/keys.rs
View File

@@ -14,6 +14,10 @@ pub(super) fn pool_cooldown_key(provider_id: &str, key_id: &str) -> String {
format!("ap:{provider_id}:cooldown:{key_id}")
}
pub(super) fn pool_lease_key(provider_id: &str, key_id: &str) -> String {
format!("ap:{provider_id}:lease:{key_id}")
}
pub(super) fn pool_cooldown_index_key(provider_id: &str) -> String {
format!("ap:{provider_id}:cooldown_idx")
}

27
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/leases.rs Normal file
View File

@@ -0,0 +1,27 @@
use super::keys::pool_lease_key;
use aether_runtime_state::{DataLayerError, RuntimeLockLease, RuntimeState};
use std::time::Duration;
pub(crate) const ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS: u64 = 15 * 60 * 1000;
pub(crate) async fn try_claim_admin_provider_pool_key(
runtime: &RuntimeState,
provider_id: &str,
key_id: &str,
owner: &str,
) -> Result<Option<RuntimeLockLease>, DataLayerError> {
runtime
.lock_try_acquire(
&pool_lease_key(provider_id, key_id),
owner,
Duration::from_millis(ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS),
)
.await
}
pub(crate) async fn release_admin_provider_pool_key_lease(
runtime: &RuntimeState,
lease: &RuntimeLockLease,
) -> Result<bool, DataLayerError> {
runtime.lock_release(lease).await
}

5
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/mod.rs
View File

@@ -1,9 +1,14 @@
mod keys;
mod leases;
mod mutations;
mod reads;
mod status;
mod writes;
pub(crate) use self::leases::{
release_admin_provider_pool_key_lease, try_claim_admin_provider_pool_key,
ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS,
};
pub(crate) use self::mutations::{
clear_admin_provider_pool_cooldown, reset_admin_provider_pool_cost,
};

3
apps/aether-gateway/src/handlers/shared/provider_pool.rs
View File

@@ -2,7 +2,8 @@ pub(crate) use super::super::admin::provider::pool::config::admin_provider_pool_
pub(crate) use super::super::admin::provider::pool::runtime::{
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_runtime_state,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success,
record_admin_provider_pool_success, release_admin_provider_pool_key_lease,
try_claim_admin_provider_pool_key, ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS,
};
pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,

74
apps/aether-gateway/src/orchestration/attempt.rs
View File

@@ -1,3 +1,4 @@
use aether_runtime_state::RuntimeLockLease;
use aether_scheduler_core::parse_request_candidate_report_context;
use serde_json::Value;
@@ -29,8 +30,14 @@ impl ExecutionAttemptIdentity {
pub(crate) struct LocalExecutionCandidateMetadata {
pub(crate) candidate_group_id: Option<String>,
pub(crate) pool_key_index: Option<u32>,
pub(crate) pool_key_lease: Option<RuntimeLockLease>,
}
pub(crate) const POOL_KEY_LEASE_KEY_REPORT_FIELD: &str = "pool_key_lease_key";
pub(crate) const POOL_KEY_LEASE_OWNER_REPORT_FIELD: &str = "pool_key_lease_owner";
pub(crate) const POOL_KEY_LEASE_TOKEN_REPORT_FIELD: &str = "pool_key_lease_token";
pub(crate) const POOL_KEY_LEASE_TTL_MS_REPORT_FIELD: &str = "pool_key_lease_ttl_ms";
pub(crate) fn attempt_identity_from_report_context(
report_context: Option<&Value>,
) -> Option<ExecutionAttemptIdentity> {
@@ -57,9 +64,65 @@ pub(crate) fn local_execution_candidate_metadata_from_report_context(
.and_then(|value| value.get("pool_key_index"))
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok()),
pool_key_lease: pool_key_lease_from_report_context(report_context),
}
}
pub(crate) fn insert_pool_key_lease_report_context_fields(
extra_fields: &mut serde_json::Map<String, Value>,
lease: Option<&RuntimeLockLease>,
) {
let Some(lease) = lease else {
return;
};
extra_fields.insert(
POOL_KEY_LEASE_KEY_REPORT_FIELD.to_string(),
Value::String(lease.key.clone()),
);
extra_fields.insert(
POOL_KEY_LEASE_OWNER_REPORT_FIELD.to_string(),
Value::String(lease.owner.clone()),
);
extra_fields.insert(
POOL_KEY_LEASE_TOKEN_REPORT_FIELD.to_string(),
Value::String(lease.token.clone()),
);
extra_fields.insert(
POOL_KEY_LEASE_TTL_MS_REPORT_FIELD.to_string(),
Value::Number(lease.ttl_ms.into()),
);
}
fn pool_key_lease_from_report_context(report_context: Option<&Value>) -> Option<RuntimeLockLease> {
let report_context = report_context?;
let key = report_context
.get(POOL_KEY_LEASE_KEY_REPORT_FIELD)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let owner = report_context
.get(POOL_KEY_LEASE_OWNER_REPORT_FIELD)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let token = report_context
.get(POOL_KEY_LEASE_TOKEN_REPORT_FIELD)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let ttl_ms = report_context
.get(POOL_KEY_LEASE_TTL_MS_REPORT_FIELD)
.and_then(Value::as_u64)
.filter(|value| *value > 0)?;
Some(RuntimeLockLease {
key: key.to_string(),
owner: owner.to_string(),
token: token.to_string(),
ttl_ms,
})
}
pub(crate) fn build_local_attempt_identities(
candidate_index: u32,
transport: &GatewayProviderTransportSnapshot,
@@ -127,6 +190,7 @@ mod tests {
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use aether_runtime_state::RuntimeLockLease;
fn sample_transport(
provider_max_retries: Option<i32>,
@@ -412,6 +476,10 @@ mod tests {
let metadata = local_execution_candidate_metadata_from_report_context(Some(&json!({
"candidate_group_id": "group-1",
"pool_key_index": 3,
"pool_key_lease_key": "ap:provider-1:lease:key-1",
"pool_key_lease_owner": "gateway-1",
"pool_key_lease_token": "gateway-1:token-1",
"pool_key_lease_ttl_ms": 900000,
})));
assert_eq!(
@@ -419,6 +487,12 @@ mod tests {
LocalExecutionCandidateMetadata {
candidate_group_id: Some("group-1".to_string()),
pool_key_index: Some(3),
pool_key_lease: Some(RuntimeLockLease {
key: "ap:provider-1:lease:key-1".to_string(),
owner: "gateway-1".to_string(),
token: "gateway-1:token-1".to_string(),
ttl_ms: 900000,
}),
}
);
}

24
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -27,8 +27,9 @@ use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_conf
use crate::handlers::shared::provider_pool::{
admin_provider_pool_key_circuit_breaker_reason, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
AdminProviderPoolConfig,
release_admin_provider_pool_key_lease, AdminProviderPoolConfig,
};
use crate::orchestration::local_execution_candidate_metadata_from_report_context;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
@@ -129,19 +130,40 @@ pub(crate) async fn apply_local_execution_effect(
}
LocalExecutionEffect::PoolSuccessSync { payload } => {
record_sync_pool_success_effect(state, context, payload).await;
release_pool_key_lease_effect(state, context).await;
}
LocalExecutionEffect::PoolSuccessStream { payload } => {
record_stream_pool_success_effect(state, context, payload).await;
release_pool_key_lease_effect(state, context).await;
}
LocalExecutionEffect::PoolError(effect) => {
record_pool_error_effect(state, context, effect).await;
release_pool_key_lease_effect(state, context).await;
}
LocalExecutionEffect::PoolStreamTimeout => {
record_pool_stream_timeout_effect(state, context).await;
release_pool_key_lease_effect(state, context).await;
}
}
}
async fn release_pool_key_lease_effect(state: &AppState, context: LocalExecutionEffectContext<'_>) {
let metadata = local_execution_candidate_metadata_from_report_context(context.report_context);
let Some(lease) = metadata.pool_key_lease else {
return;
};
if let Err(err) =
release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), &lease).await
{
warn!(
error = ?err,
provider_id = %context.plan.provider_id,
key_id = %context.plan.key_id,
"gateway orchestration effects: failed to release pool key lease"
);
}
}
fn report_context_string_field<'a>(
report_context: Option<&'a Value>,
field: &str,

3
apps/aether-gateway/src/orchestration/mod.rs
View File

@@ -17,7 +17,8 @@ pub(crate) use self::adaptive::{
LocalAdaptiveRateLimitProjection, LocalAdaptiveSuccessProjection,
};
pub(crate) use self::attempt::{
attempt_identity_from_report_context, build_local_attempt_identities, local_attempt_slot_count,
attempt_identity_from_report_context, build_local_attempt_identities,
insert_pool_key_lease_report_context_fields, local_attempt_slot_count,
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
LocalExecutionCandidateMetadata,
};