feat: add routing profile scheduling policies

This commit is contained in:
fawney19
2026-05-18 11:03:49 +08:00
parent a2f91b4108
commit 92813e6122
124 changed files with 11681 additions and 578 deletions
@@ -33,7 +33,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
set_local_runtime_miss_diagnostic_reason(
state,
@@ -55,6 +55,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let body_json = input.effective_body_json(body_json);
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
@@ -70,7 +71,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -100,7 +101,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
set_local_runtime_miss_diagnostic_reason(
state,
@@ -122,6 +123,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let body_json = input.effective_body_json(body_json);
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
@@ -137,7 +139,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
{
return Ok(Some(payload));
}
@@ -13,6 +13,7 @@ use crate::ai_serving::planner::candidate_metadata::{
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::planner::common::extract_requested_model_from_request;
use crate::ai_serving::planner::decision_input::{
attach_routing_policy_to_local_requested_model_input,
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
use crate::ai_serving::planner::materialization_policy::{
@@ -39,19 +40,21 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
spec: LocalSameFormatProviderSpec,
) -> Option<LocalSameFormatProviderDecisionInput> {
) -> Result<Option<LocalSameFormatProviderDecisionInput>, GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
return None;
return Ok(None);
};
let requested_model = extract_requested_model_from_request(
let Some(requested_model) = extract_requested_model_from_request(
parts,
body_json,
spec_metadata
.requested_model_family
.expect("same-format provider specs should declare requested-model family"),
)?;
) else {
return Ok(None);
};
let resolved_input = match resolve_local_authenticated_decision_input(
state,
@@ -62,7 +65,7 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
.await
{
Ok(Some(resolved_input)) => resolved_input,
Ok(None) => return None,
Ok(None) => return Ok(None),
Err(err) => {
warn!(
trace_id = %trace_id,
@@ -70,14 +73,31 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
error = ?err,
"gateway local same-format decision auth snapshot read failed"
);
return None;
return Err(err);
}
};
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
input.request_auth_channel = decision.request_auth_channel.clone();
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
Some(input)
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
state,
parts,
&mut input,
body_json,
spec_metadata.api_format,
)
.await
{
warn!(
trace_id = %trace_id,
api_format = spec_metadata.api_format,
error = ?err,
"gateway local same-format decision routing profile resolution failed"
);
return Err(err);
}
Ok(Some(input))
}
pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
@@ -114,6 +134,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
@@ -211,6 +232,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
input.required_capabilities.as_ref(),
input.routing_policy.as_ref(),
sticky_session_token.as_deref(),
input.request_auth_channel.as_deref(),
persistence_policy,
@@ -6,6 +6,7 @@ use crate::ai_serving::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
};
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
use crate::ai_serving::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
@@ -22,7 +23,7 @@ use crate::ai_serving::transport::{
};
use crate::{
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
AiExecutionDecision, AppState,
AiExecutionDecision, AppState, GatewayError,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
@@ -40,7 +41,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
input: &LocalSameFormatProviderDecisionInput,
attempt: LocalSameFormatProviderCandidateAttempt,
spec: LocalSameFormatProviderSpec,
) -> Option<AiExecutionDecision> {
) -> Result<Option<AiExecutionDecision>, GatewayError> {
let spec_metadata = local_same_format_provider_spec_metadata(spec);
let LocalSameFormatProviderCandidateAttempt {
eligible,
@@ -51,10 +52,13 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
let candidate = &eligible.candidate;
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
let Some(resolved) = resolve_local_same_format_provider_candidate_payload_parts(
state, parts, trace_id, body_json, input, &attempt, spec,
)
.await?;
.await
else {
return Ok(None);
};
let prompt_cache_key = resolved
.provider_request_body
@@ -85,6 +89,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
);
}
let provider_api_format = resolved.provider_api_format.clone();
let effective_headers = input.effective_headers(&parts.headers);
let report_context = append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
@@ -112,7 +117,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
original_headers: effective_headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
@@ -151,41 +156,41 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_request_body,
} = resolved;
Some(build_ai_execution_decision_response(
AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.to_string(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url,
provider_request_method: None,
auth_header,
auth_value,
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model,
prompt_cache_key,
provider_request_headers,
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream,
report_kind: Some(report_kind.to_string()),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
},
))
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
decision_kind: spec_metadata.decision_kind.to_string(),
execution_strategy,
conversion_mode,
request_id: trace_id.to_string(),
candidate_id: candidate_id.to_string(),
provider_name: transport.provider.name.clone(),
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
upstream_base_url: transport.endpoint.base_url.clone(),
upstream_url,
provider_request_method: None,
auth_header,
auth_value,
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model,
prompt_cache_key,
provider_request_headers,
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
upstream_is_stream,
report_kind: Some(report_kind.to_string()),
report_context: Some(report_context),
auth_context: input.auth_context.clone(),
});
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
Ok(Some(decision))
}
pub(super) async fn mark_skipped_local_same_format_provider_candidate(
@@ -125,6 +125,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
Some(&input.requested_model),
)
.await;
let effective_headers = input.effective_headers(&parts.headers);
let Some(mut base_provider_request_body) =
super::super::request::build_same_format_provider_request_body(
@@ -133,7 +134,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
&prepared.mapped_model,
spec,
prepared.transport.endpoint.body_rules.as_ref(),
Some(&parts.headers),
Some(effective_headers),
prepared.upstream_is_stream,
prepared.force_body_stream_field,
prepared.kiro_auth.as_ref(),
@@ -279,7 +280,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
.unwrap_or_default();
let Some(provider_request_headers) =
build_same_format_provider_headers(SameFormatProviderHeadersInput {
headers: &parts.headers,
headers: effective_headers,
provider_request_body: &provider_request_body,
original_request_body: body_json,
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
@@ -31,7 +31,7 @@ pub(crate) struct LocalSameFormatProviderSyncAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
body_json: serde_json::Value,
input: LocalSameFormatProviderDecisionInput,
spec: LocalSameFormatProviderSpec,
requested_model_family: RequestedModelFamily,
@@ -42,7 +42,7 @@ pub(crate) struct LocalSameFormatProviderStreamAttemptSource<'a> {
state: &'a AppState,
parts: &'a http::request::Parts,
trace_id: &'a str,
body_json: &'a serde_json::Value,
body_json: serde_json::Value,
input: LocalSameFormatProviderDecisionInput,
spec: LocalSameFormatProviderSpec,
requested_model_family: RequestedModelFamily,
@@ -64,7 +64,7 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
set_local_runtime_miss_diagnostic_reason(
state,
@@ -85,8 +85,13 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let effective_body_json = input.effective_body_json(body_json).clone();
let (candidates, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
state,
trace_id,
&input,
&effective_body_json,
spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
@@ -103,7 +108,7 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
state,
parts,
trace_id,
body_json,
body_json: effective_body_json,
input,
spec,
requested_model_family,
@@ -128,7 +133,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
set_local_runtime_miss_diagnostic_reason(
state,
@@ -149,8 +154,13 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let effective_body_json = input.effective_body_json(body_json).clone();
let (candidates, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
state,
trace_id,
&input,
&effective_body_json,
spec,
)
.await?;
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
@@ -167,7 +177,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
state,
parts,
trace_id,
body_json,
body_json: effective_body_json,
input,
spec,
requested_model_family,
@@ -244,12 +254,12 @@ impl LocalSameFormatProviderSyncAttemptSource<'_> {
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.body_json,
&self.input,
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -257,7 +267,7 @@ impl LocalSameFormatProviderSyncAttemptSource<'_> {
match build_sync_plan_from_requested_model_family(
self.requested_model_family,
self.parts,
self.body_json,
&self.body_json,
payload,
) {
Ok(value) => Ok(value),
@@ -282,12 +292,12 @@ impl LocalSameFormatProviderStreamAttemptSource<'_> {
self.state,
self.parts,
self.trace_id,
self.body_json,
&self.body_json,
&self.input,
attempt,
self.spec,
)
.await
.await?
else {
return Ok(None);
};
@@ -295,7 +305,7 @@ impl LocalSameFormatProviderStreamAttemptSource<'_> {
match build_stream_plan_from_requested_model_family(
self.requested_model_family,
self.parts,
self.body_json,
&self.body_json,
payload,
) {
Ok(value) => Ok(value),
@@ -326,7 +336,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
set_local_runtime_miss_diagnostic_reason(
state,
@@ -347,6 +357,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let body_json = input.effective_body_json(body_json);
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
@@ -365,7 +376,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};
@@ -411,7 +422,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
let Some(input) = resolve_local_same_format_provider_decision_input(
state, parts, trace_id, decision, body_json, spec,
)
.await
.await?
else {
set_local_runtime_miss_diagnostic_reason(
state,
@@ -432,6 +443,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let body_json = input.effective_body_json(body_json);
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
@@ -450,7 +462,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
.await
.await?
else {
continue;
};