mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 17:37:47 +08:00
fix(gateway): add provider key concurrency and cache affinity modes
This commit is contained in:
@@ -126,7 +126,7 @@ use crate::orchestration::{
|
||||
LocalOAuthSuccessEffect, LocalPoolErrorEffect,
|
||||
};
|
||||
use crate::provider_pool_demand::{
|
||||
acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard,
|
||||
acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard,
|
||||
};
|
||||
use crate::request_candidate_runtime::{
|
||||
ensure_execution_request_candidate_slot, persist_local_request_candidate_status_record,
|
||||
@@ -3772,6 +3772,41 @@ async fn execute_execution_runtime_stream_inner(
|
||||
plan_kind,
|
||||
report_context.as_ref(),
|
||||
);
|
||||
let candidate_started_unix_secs = current_request_candidate_unix_ms();
|
||||
let provider_in_flight_started_at = Instant::now();
|
||||
let mut provider_pool_in_flight_guard =
|
||||
match acquire_provider_pool_execution_guard(state, &plan).await? {
|
||||
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
|
||||
ProviderPoolInFlightAdmission::Saturated { limit } => {
|
||||
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
||||
*retry_scope = AiAttemptRetryScope::Candidate;
|
||||
}
|
||||
if let Some(snapshot) = request_candidate_status_snapshot.as_ref() {
|
||||
record_local_request_candidate_status_snapshot(
|
||||
state,
|
||||
snapshot,
|
||||
SchedulerRequestCandidateStatusUpdate {
|
||||
status: RequestCandidateStatus::Skipped,
|
||||
status_code: Some(http::StatusCode::TOO_MANY_REQUESTS.as_u16()),
|
||||
error_type: Some("provider_key_concurrency_limit_reached".to_string()),
|
||||
error_message: Some(format!(
|
||||
"provider key concurrency limit reached: {limit}"
|
||||
)),
|
||||
latency_ms: Some(0),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
observe_gateway_stage_trace_ms(
|
||||
&mut stage_trace,
|
||||
"stream_provider_in_flight",
|
||||
provider_in_flight_started_at.elapsed().as_millis() as u64,
|
||||
);
|
||||
// Inline passthrough records its lifecycle seed after upstream headers are
|
||||
// available. Avoid constructing a throwaway seed on the common path.
|
||||
let mut lifecycle_seed = (!defer_stream_pending_for_direct_inline)
|
||||
@@ -3781,7 +3816,6 @@ async fn execute_execution_runtime_stream_inner(
|
||||
record_stream_pending_lifecycle(state, seed, &mut stage_trace).await;
|
||||
lifecycle_pending_recorded = true;
|
||||
}
|
||||
let candidate_started_unix_secs = current_request_candidate_unix_ms();
|
||||
if let Some(snapshot) = request_candidate_status_snapshot.clone() {
|
||||
record_local_request_candidate_status_snapshot(
|
||||
state,
|
||||
@@ -3810,20 +3844,6 @@ async fn execute_execution_runtime_stream_inner(
|
||||
.and_then(|context| context.candidate_index)
|
||||
.map(|value| value.to_string())
|
||||
.unwrap_or_else(|| "-".to_string());
|
||||
let provider_in_flight_started_at = Instant::now();
|
||||
let mut provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
|
||||
state.runtime_state.clone(),
|
||||
&plan.provider_id,
|
||||
plan.request_id.as_str(),
|
||||
plan.candidate_id.as_deref(),
|
||||
key_id.as_str(),
|
||||
)
|
||||
.await;
|
||||
observe_gateway_stage_trace_ms(
|
||||
&mut stage_trace,
|
||||
"stream_provider_in_flight",
|
||||
provider_in_flight_started_at.elapsed().as_millis() as u64,
|
||||
);
|
||||
match maybe_execute_grok_stream(&plan, report_context.as_ref()).await {
|
||||
Ok(Some(grok_stream)) => {
|
||||
return execute_stream_from_frame_stream_with_retry_scope(
|
||||
|
||||
@@ -78,7 +78,9 @@ use crate::orchestration::{
|
||||
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
|
||||
LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect,
|
||||
};
|
||||
use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard;
|
||||
use crate::provider_pool_demand::{
|
||||
acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission,
|
||||
};
|
||||
use crate::request_candidate_runtime::{
|
||||
ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data,
|
||||
record_local_request_candidate_status, record_local_request_candidate_status_snapshot,
|
||||
@@ -1995,6 +1997,32 @@ async fn execute_execution_runtime_sync_impl(
|
||||
.unwrap_or_else(|| "-".to_string());
|
||||
let candidate_started_at = Instant::now();
|
||||
let candidate_started_unix_secs = current_request_candidate_unix_ms();
|
||||
let _provider_pool_in_flight_guard = match acquire_provider_pool_execution_guard(state, &plan)
|
||||
.await?
|
||||
{
|
||||
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
|
||||
ProviderPoolInFlightAdmission::Saturated { limit } => {
|
||||
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
||||
*retry_scope = AiAttemptRetryScope::Candidate;
|
||||
}
|
||||
record_local_request_candidate_status(
|
||||
state,
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
SchedulerRequestCandidateStatusUpdate {
|
||||
status: RequestCandidateStatus::Skipped,
|
||||
status_code: Some(StatusCode::TOO_MANY_REQUESTS.as_u16()),
|
||||
error_type: Some("provider_key_concurrency_limit_reached".to_string()),
|
||||
error_message: Some(format!("provider key concurrency limit reached: {limit}")),
|
||||
latency_ms: Some(0),
|
||||
started_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
finished_at_unix_ms: Some(candidate_started_unix_secs),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
|
||||
let usage_data = state.usage_lifecycle_data_state().as_ref().clone();
|
||||
state
|
||||
@@ -2024,14 +2052,6 @@ async fn execute_execution_runtime_sync_impl(
|
||||
candidate_started_at,
|
||||
);
|
||||
let result = (async {
|
||||
let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
|
||||
state.runtime_state.clone(),
|
||||
&plan.provider_id,
|
||||
plan_request_id.as_str(),
|
||||
plan_candidate_id.as_deref(),
|
||||
key_id.as_str(),
|
||||
)
|
||||
.await;
|
||||
record_sync_execution_active(
|
||||
state,
|
||||
&plan,
|
||||
|
||||
Reference in New Issue
Block a user