mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-10 03:09:50 +08:00
feat: add routing profile scheduling policies
This commit is contained in:
@@ -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));
|
||||
}
|
||||
|
||||
+29
-7
@@ -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;
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user