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

fix(routing): harden routed pool scheduling

Browse Source
This commit is contained in:
elky
2026-07-27 22:06:28 +08:00
parent 550cc36760
commit 4148ab1931
62 changed files with 5510 additions and 850 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/proxy/mod.rs
+105 -27
View File
@@ -346,20 +346,6 @@ async fn maybe_forward_public_request_to_tunnel_owner(
}) else {
return Ok(None);
};
let cache_affinity_enabled = match read_scheduler_ordering_config(state).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
Err(err) => {
warn!(
trace_id = %request_context.trace_id,
error = ?err,
"gateway failed to load scheduler config while checking tunnel affinity forwarding mode"
);
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
}
};
if !cache_affinity_enabled {
return Ok(None);
}
let Some(api_format) = decision
.auth_endpoint_signature
.as_deref()
@@ -382,21 +368,66 @@ async fn maybe_forward_public_request_to_tunnel_owner(
crate::headers::decoded_request_body_bytes(&parts.headers, body.as_ref()).ok()?;
serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok()
});
let client_session_affinity =
crate::client_session_affinity::client_session_affinity_from_api_request(
api_format,
&parts.headers,
body_json.as_ref(),
);
let Some(target) = crate::scheduler::affinity::read_cached_scheduler_affinity_target(
let empty_body_json = serde_json::Value::Null;
let affinity_context = match crate::ai_serving::resolve_tunnel_scheduler_affinity_context(
state,
&auth_context.api_key_id,
client_session_affinity.as_ref(),
parts,
decision,
requested_model,
body_json.as_ref().unwrap_or(&empty_body_json),
api_format,
&requested_model,
) else {
)
.await
{
Ok(Some(context)) => context,
Ok(None) => return Ok(None),
Err(err) => {
warn!(
trace_id = %request_context.trace_id,
error = ?err,
"gateway failed to resolve routing policy while checking tunnel affinity forwarding"
);
return Ok(None);
}
};
let target = if let Some(policy_context) = affinity_context.policy_context.as_ref() {
crate::scheduler::affinity::read_cached_scheduler_affinity_target_with_policy_context(
state,
&auth_context.api_key_id,
affinity_context.client_session_affinity.as_ref(),
api_format,
&affinity_context.requested_model,
policy_context,
)
} else {
let cache_affinity_enabled = match read_scheduler_ordering_config(state).await {
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
Err(err) => {
warn!(
trace_id = %request_context.trace_id,
error = ?err,
"gateway failed to load scheduler config while checking tunnel affinity forwarding mode"
);
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
}
};
if !cache_affinity_enabled {
return Ok(None);
}
crate::scheduler::affinity::read_cached_scheduler_affinity_target(
state,
&auth_context.api_key_id,
affinity_context.client_session_affinity.as_ref(),
api_format,
&affinity_context.requested_model,
)
};
let Some(target) = target else {
return Ok(None);
};
if !routing_overlay_allows_affinity_target(affinity_context.routing_overlay.as_ref(), &target) {
return Ok(None);
}
let transport = match state
.read_provider_transport_snapshot(&target.provider_id, &target.endpoint_id, &target.key_id)
@@ -547,6 +578,16 @@ async fn maybe_forward_public_request_to_tunnel_owner(
Ok(Some(response))
}
fn routing_overlay_allows_affinity_target(
routing_overlay: Option<&aether_routing_core::RankingOverlay>,
target: &aether_scheduler_core::SchedulerAffinityTarget,
) -> bool {
routing_overlay.is_none_or(|overlay| {
overlay.provider_allowed(target.provider_id.as_str())
&& overlay.key_allowed(target.key_id.as_str())
})
}
fn owner_forward_request_is_stream(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
@@ -2327,14 +2368,51 @@ mod tests {
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
owner_forward_request_is_stream, restore_redacted_stream_execution_response,
restore_redacted_sync_execution_response, GatewayControlDecision,
LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, RequestBodyBufferPolicy,
restore_redacted_sync_execution_response, routing_overlay_allows_affinity_target,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
RequestBodyBufferPolicy,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::{header, HeaderMap, HeaderValue, Method, Response};
use serde_json::json;
use tokio::sync::Semaphore;
#[test]
fn routing_overlay_blocks_disallowed_tunnel_affinity_target() {
let target = aether_scheduler_core::SchedulerAffinityTarget {
provider_id: "provider-allowed".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-allowed".to_string(),
};
let matching = aether_routing_core::RankingOverlay {
allowed_providers: vec!["provider-allowed".to_string()],
allowed_keys: vec!["key-allowed".to_string()],
..aether_routing_core::RankingOverlay::default()
};
let wrong_provider = aether_routing_core::RankingOverlay {
allowed_providers: vec!["provider-other".to_string()],
..matching.clone()
};
let wrong_key = aether_routing_core::RankingOverlay {
allowed_keys: vec!["key-other".to_string()],
..matching.clone()
};
assert!(routing_overlay_allows_affinity_target(None, &target));
assert!(routing_overlay_allows_affinity_target(
Some(&matching),
&target
));
assert!(!routing_overlay_allows_affinity_target(
Some(&wrong_provider),
&target
));
assert!(!routing_overlay_allows_affinity_target(
Some(&wrong_key),
&target
));
}
#[test]
fn owner_forward_uses_search_protocol_timeout_semantics() {
let request = http::Request::builder()