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: 抽离 AI pipeline 与调度共享能力逻辑

Browse Source
This commit is contained in:
fawney19
2026-04-10 01:46:14 +08:00
parent b901a6ffc7
commit 5014e2f5fd
255 changed files with 15057 additions and 3115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

202
crates/aether-scheduler-core/src/candidate.rs
View File

@@ -5,6 +5,18 @@ use aether_data_contracts::repository::candidates::StoredRequestCandidate;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data_contracts::DataLayerError;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum SchedulerPriorityMode {
Provider,
GlobalKey,
}
impl Default for SchedulerPriorityMode {
fn default() -> Self {
Self::Provider
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct SchedulerMinimalCandidateSelectionCandidate {
pub provider_id: String,
@@ -63,6 +75,34 @@ pub fn candidate_supports_required_capability(
false
}
pub fn requested_capability_priority_for_candidate(
required_capabilities: Option<&serde_json::Value>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (u32, u32) {
let Some(required_capabilities) = required_capabilities.and_then(serde_json::Value::as_object)
else {
return (0, 0);
};
let mut exclusive_misses = 0u32;
let mut compatible_misses = 0u32;
for (capability, value) in required_capabilities {
if !requested_capability_is_enabled(value) {
continue;
}
if candidate_supports_required_capability(candidate, capability) {
continue;
}
if requested_capability_is_compatible(capability) {
compatible_misses += 1;
} else {
exclusive_misses += 1;
}
}
(exclusive_misses, compatible_misses)
}
pub fn auth_api_key_concurrency_limit_reached(
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
@@ -83,8 +123,10 @@ pub fn build_minimal_candidate_selection(
requested_model_name: &str,
resolved_global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_constraints: Option<&crate::SchedulerAuthConstraints>,
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
if normalized_api_format.is_empty() {
return Ok(Vec::new());
@@ -143,22 +185,57 @@ pub fn build_minimal_candidate_selection(
}
candidates.sort_by(|left, right| {
left.key_global_priority_for_format
requested_capability_priority_for_candidate(required_capabilities, left)
.cmp(&requested_capability_priority_for_candidate(
required_capabilities,
right,
))
.then_with(|| {
compare_candidates_by_priority_mode(left, right, priority_mode, affinity_key)
})
});
Ok(candidates)
}
fn requested_capability_is_enabled(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => value.as_i64().is_some_and(|value| value > 0),
_ => false,
}
}
fn requested_capability_is_compatible(capability: &str) -> bool {
matches!(
capability.trim().to_ascii_lowercase().as_str(),
"cache_1h" | "context_1m"
)
}
pub fn compare_candidates_by_priority_mode(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
affinity_key: Option<&str>,
) -> std::cmp::Ordering {
match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then_with(|| compare_candidate_identity(left, right)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then(left.provider_id.cmp(&right.provider_id))
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
});
Ok(candidates)
.then_with(|| compare_candidate_identity(left, right)),
}
}
pub fn collect_global_model_names_for_required_capability(
@@ -238,23 +315,38 @@ pub fn collect_selectable_candidates_from_keys(
pub fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
required_capabilities: Option<&serde_json::Value>,
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) {
candidates.sort_by(|left, right| {
left.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| compare_provider_key_health_order(left, right, provider_key_rpm_states))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then(left.provider_id.cmp(&right.provider_id))
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
requested_capability_priority_for_candidate(required_capabilities, left)
.cmp(&requested_capability_priority_for_candidate(
required_capabilities,
right,
))
.then_with(|| match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| {
compare_provider_key_health_order(left, right, provider_key_rpm_states)
})
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then_with(|| compare_candidate_identity(left, right)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| {
compare_provider_key_health_order(left, right, provider_key_rpm_states)
})
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| compare_candidate_identity(left, right)),
})
});
}
@@ -362,6 +454,20 @@ fn candidate_provider_key_health_bucket(
})
}
fn compare_candidate_identity(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
) -> std::cmp::Ordering {
left.provider_id
.cmp(&right.provider_id)
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
}
fn candidate_provider_key_health_score(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
@@ -392,6 +498,7 @@ mod tests {
collect_global_model_names_for_required_capability,
collect_selectable_candidates_from_keys, reorder_candidates_by_scheduler_health,
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
};
use crate::SchedulerAuthConstraints;
@@ -478,11 +585,11 @@ mod tests {
fn stored_candidate(
id: &str,
status: RequestCandidateStatus,
created_at_unix_secs: i64,
created_at_unix_ms: i64,
) -> StoredRequestCandidate {
let finished_at_unix_secs = match status {
let finished_at_unix_ms = match status {
RequestCandidateStatus::Pending | RequestCandidateStatus::Streaming => None,
_ => Some(created_at_unix_secs),
_ => Some(created_at_unix_ms),
};
StoredRequestCandidate::new(
id.to_string(),
@@ -506,9 +613,9 @@ mod tests {
None,
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
finished_at_unix_secs,
created_at_unix_ms,
Some(created_at_unix_ms),
finished_at_unix_ms,
)
.expect("candidate should build")
}
@@ -546,8 +653,10 @@ mod tests {
"gpt-5",
"gpt-5",
false,
None,
Some(&constraints),
None,
SchedulerPriorityMode::Provider,
)
.expect("candidate selection should build");
@@ -579,6 +688,35 @@ mod tests {
assert_eq!(model_names, vec!["gpt-5".to_string()]);
}
#[test]
fn minimal_candidate_selection_prefers_matching_requested_capabilities_before_priority() {
let mut missing_capability = sample_row("1");
missing_capability.key_capabilities = Some(serde_json::json!({"cache_1h": false}));
missing_capability.provider_priority = 0;
let mut matching_capability = sample_row("2");
matching_capability.key_capabilities = Some(serde_json::json!({"cache_1h": true}));
matching_capability.provider_priority = 10;
let required_capabilities = serde_json::json!({"cache_1h": true});
let candidates = build_minimal_candidate_selection(
vec![missing_capability, matching_capability],
"openai:chat",
"gpt-5",
"gpt-5",
false,
Some(&required_capabilities),
None,
None,
SchedulerPriorityMode::Provider,
)
.expect("candidate selection should build");
assert_eq!(candidates.len(), 2);
assert_eq!(candidates[0].key_id, "key-2");
assert_eq!(candidates[1].key_id, "key-1");
}
#[test]
fn reorders_candidates_by_health_before_affinity_tiebreak() {
let mut candidates = vec![
@@ -595,7 +733,9 @@ mod tests {
reorder_candidates_by_scheduler_health(
&mut candidates,
&provider_key_rpm_states,
None,
Some("api-key-1"),
SchedulerPriorityMode::GlobalKey,
);
assert_ne!(candidates[0].key_id, "key-2");

66
crates/aether-scheduler-core/src/health.rs
View File

@@ -57,9 +57,10 @@ pub fn is_candidate_in_recent_failure_cooldown(
}
let observed_at_unix_secs = candidate
.finished_at_unix_secs
.or(candidate.started_at_unix_secs)
.unwrap_or(candidate.created_at_unix_secs);
.finished_at_unix_ms
.or(candidate.started_at_unix_ms)
.map(|ms| ms / 1000)
.unwrap_or(candidate.created_at_unix_ms / 1000);
if now_unix_secs.saturating_sub(observed_at_unix_secs) > FAILURE_COOLDOWN_WINDOW_SECS {
continue;
}
@@ -151,8 +152,9 @@ pub fn count_recent_rpm_requests_for_provider_key_since(
continue;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
.started_at_unix_ms
.map(|ms| ms / 1000)
.unwrap_or(candidate.created_at_unix_ms / 1000);
if reset_after_unix_secs.is_some_and(|reset_after| observed_at_unix_secs <= reset_after) {
continue;
}
@@ -308,7 +310,7 @@ fn provider_key_dynamic_reservation_ratio(
}
fn is_recently_active(candidate: &StoredRequestCandidate, now_unix_secs: u64) -> bool {
if candidate.finished_at_unix_secs.is_some() {
if candidate.finished_at_unix_ms.is_some() {
return false;
}
@@ -320,22 +322,21 @@ fn is_recently_active(candidate: &StoredRequestCandidate, now_unix_secs: u64) ->
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
.started_at_unix_ms
.map(|ms| ms / 1000)
.unwrap_or(candidate.created_at_unix_ms / 1000);
now_unix_secs.saturating_sub(observed_at_unix_secs) <= ACTIVE_REQUEST_WINDOW_SECS
}
fn is_recent_rpm_observation(candidate: &StoredRequestCandidate, now_unix_secs: u64) -> bool {
if !candidate
.status
.is_attempted(candidate.started_at_unix_secs)
{
if !candidate.status.is_attempted(candidate.started_at_unix_ms) {
return false;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
.started_at_unix_ms
.map(|ms| ms / 1000)
.unwrap_or(candidate.created_at_unix_ms / 1000);
now_unix_secs.saturating_sub(observed_at_unix_secs) <= PROVIDER_KEY_RPM_WINDOW_SECS
}
@@ -465,6 +466,7 @@ mod tests {
status: RequestCandidateStatus,
created_at_unix_secs: i64,
) -> StoredRequestCandidate {
let created_at_unix_ms = created_at_unix_secs * 1000;
StoredRequestCandidate::new(
id.to_string(),
format!("req-{id}"),
@@ -487,9 +489,9 @@ mod tests {
None,
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
Some(created_at_unix_secs),
created_at_unix_ms,
Some(created_at_unix_ms),
Some(created_at_unix_ms),
)
.expect("candidate should build")
}
@@ -710,9 +712,9 @@ mod tests {
Some(7),
None,
None,
95,
Some(95),
Some(96),
95_000,
Some(95_000),
Some(96_000),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
@@ -737,9 +739,9 @@ mod tests {
None,
None,
None,
98,
Some(98),
Some(99),
98_000,
Some(98_000),
Some(99_000),
)
.expect("candidate should build"),
];
@@ -775,9 +777,9 @@ mod tests {
Some(7),
None,
None,
95,
Some(95),
Some(96),
95_000,
Some(95_000),
Some(96_000),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
@@ -802,9 +804,9 @@ mod tests {
Some(2),
None,
None,
99,
Some(99),
Some(100),
99_000,
Some(99_000),
Some(100_000),
)
.expect("candidate should build"),
];
@@ -854,9 +856,9 @@ mod tests {
Some(9),
None,
None,
95,
Some(95),
Some(96),
95_000,
Some(95_000),
Some(96_000),
)
.expect("candidate should build")];

5
crates/aether-scheduler-core/src/lib.rs
View File

@@ -18,8 +18,9 @@ pub use candidate::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability, collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate,
compare_candidates_by_priority_mode, reorder_candidates_by_scheduler_health,
requested_capability_priority_for_candidate, CandidateRuntimeSelectabilityInput,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,

98
crates/aether-scheduler-core/src/request_candidate.rs
View File

@@ -31,9 +31,9 @@ pub struct SchedulerResolvedReportRequestCandidateSlot {
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub extra_data: Option<Value>,
pub created_at_unix_secs: u64,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
pub created_at_unix_ms: u64,
pub started_at_unix_ms: Option<u64>,
pub finished_at_unix_ms: Option<u64>,
}
pub struct SchedulerExecutionRequestCandidateSeed {
@@ -48,8 +48,8 @@ pub struct SchedulerRequestCandidateStatusUpdate {
pub error_type: Option<String>,
pub error_message: Option<String>,
pub latency_ms: Option<u64>,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
pub started_at_unix_ms: Option<u64>,
pub finished_at_unix_ms: Option<u64>,
}
#[derive(Debug, Clone)]
@@ -63,7 +63,7 @@ pub struct LocalRequestCandidateStatusRecordInput<'a> {
pub struct ReportRequestCandidateStatusRecordInput {
pub slot: SchedulerResolvedReportRequestCandidateSlot,
pub status_update: SchedulerRequestCandidateStatusUpdate,
pub now_unix_secs: u64,
pub now_unix_ms: u64,
}
pub fn execution_error_details(
@@ -115,16 +115,16 @@ pub fn parse_request_candidate_report_context(
pub fn resolve_report_request_candidate_slot(
existing_candidates: &[StoredRequestCandidate],
metadata: SchedulerRequestCandidateReportContext,
now_unix_secs: u64,
now_unix_ms: u64,
generated_candidate_id: String,
) -> Option<SchedulerResolvedReportRequestCandidateSlot> {
let request_id = metadata.request_id.clone()?;
let matched_candidate = match_existing_report_candidate(existing_candidates, &metadata);
let synthesized_extra_data = build_report_candidate_extra_data(&metadata);
let created_at_unix_secs = matched_candidate
let created_at_unix_ms = matched_candidate
.as_ref()
.map(|candidate| candidate.created_at_unix_secs)
.unwrap_or(now_unix_secs);
.map(|candidate| candidate.created_at_unix_ms)
.unwrap_or(now_unix_ms);
let candidate_index = matched_candidate
.as_ref()
.map(|candidate| candidate.candidate_index)
@@ -168,20 +168,20 @@ pub fn resolve_report_request_candidate_slot(
.as_ref()
.and_then(|candidate| candidate.extra_data.clone())
.or(synthesized_extra_data),
created_at_unix_secs,
started_at_unix_secs: matched_candidate
created_at_unix_ms,
started_at_unix_ms: matched_candidate
.as_ref()
.and_then(|candidate| candidate.started_at_unix_secs),
finished_at_unix_secs: matched_candidate
.and_then(|candidate| candidate.started_at_unix_ms),
finished_at_unix_ms: matched_candidate
.as_ref()
.and_then(|candidate| candidate.finished_at_unix_secs),
.and_then(|candidate| candidate.finished_at_unix_ms),
})
}
pub fn build_execution_request_candidate_seed(
plan: &ExecutionPlan,
report_context: Option<&Value>,
started_at_unix_secs: u64,
started_at_unix_ms: u64,
generated_candidate_id: String,
) -> SchedulerExecutionRequestCandidateSeed {
let mut context = report_context
@@ -247,9 +247,9 @@ pub fn build_execution_request_candidate_seed(
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: Some(started_at_unix_secs),
started_at_unix_secs: Some(started_at_unix_secs),
finished_at_unix_secs: None,
created_at_unix_ms: Some(started_at_unix_ms),
started_at_unix_ms: Some(started_at_unix_ms),
finished_at_unix_ms: None,
},
report_context: Value::Object(context),
}
@@ -269,8 +269,8 @@ pub fn build_local_request_candidate_status_record(
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
started_at_unix_ms,
finished_at_unix_ms,
} = status_update;
let candidate_id = plan
@@ -303,9 +303,9 @@ pub fn build_local_request_candidate_status_record(
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: None,
started_at_unix_secs,
finished_at_unix_secs,
created_at_unix_ms: None,
started_at_unix_ms,
finished_at_unix_ms,
})
}
@@ -315,7 +315,7 @@ pub fn build_report_request_candidate_status_record(
let ReportRequestCandidateStatusRecordInput {
slot,
status_update,
now_unix_secs,
now_unix_ms,
} = input;
let SchedulerRequestCandidateStatusUpdate {
status,
@@ -323,16 +323,16 @@ pub fn build_report_request_candidate_status_record(
error_type,
error_message,
latency_ms,
started_at_unix_secs,
finished_at_unix_secs,
started_at_unix_ms,
finished_at_unix_ms,
} = status_update;
let terminal_unix_secs = finished_at_unix_secs.unwrap_or(now_unix_secs);
let started_at_unix_secs = started_at_unix_secs
.or(slot.started_at_unix_secs)
let terminal_unix_secs = finished_at_unix_ms.unwrap_or(now_unix_ms);
let started_at_unix_ms = started_at_unix_ms
.or(slot.started_at_unix_ms)
.or_else(|| status.is_attempted(None).then_some(terminal_unix_secs));
let finished_at_unix_secs = finished_at_unix_secs
.or(slot.finished_at_unix_secs)
let finished_at_unix_ms = finished_at_unix_ms
.or(slot.finished_at_unix_ms)
.or_else(|| is_terminal_candidate_status(status).then_some(terminal_unix_secs));
UpsertRequestCandidateRecord {
@@ -357,9 +357,9 @@ pub fn build_report_request_candidate_status_record(
concurrent_requests: None,
extra_data: slot.extra_data,
required_capabilities: None,
created_at_unix_secs: Some(slot.created_at_unix_secs),
started_at_unix_secs,
finished_at_unix_secs,
created_at_unix_ms: Some(slot.created_at_unix_ms),
started_at_unix_ms,
finished_at_unix_ms,
}
}
@@ -443,7 +443,7 @@ fn match_existing_report_candidate<'a>(
(
candidate.retry_index,
candidate.candidate_index,
candidate.created_at_unix_secs,
candidate.created_at_unix_ms,
)
})
}
@@ -524,8 +524,8 @@ mod tests {
None,
None,
None,
100,
Some(110),
100_000,
Some(110_000),
None,
)
.expect("candidate should build")
@@ -660,8 +660,8 @@ mod tests {
error_type: Some("Upstream5xx".to_string()),
error_message: Some("boom".to_string()),
latency_ms: Some(42),
started_at_unix_secs: Some(100),
finished_at_unix_secs: Some(101),
started_at_unix_ms: Some(100),
finished_at_unix_ms: Some(101),
},
})
.expect("record should build");
@@ -688,9 +688,9 @@ mod tests {
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
extra_data: None,
created_at_unix_secs: 10,
started_at_unix_secs: None,
finished_at_unix_secs: None,
created_at_unix_ms: 10,
started_at_unix_ms: None,
finished_at_unix_ms: None,
},
status_update: SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
@@ -698,15 +698,15 @@ mod tests {
error_type: None,
error_message: None,
latency_ms: Some(12),
started_at_unix_secs: None,
finished_at_unix_secs: None,
started_at_unix_ms: None,
finished_at_unix_ms: None,
},
now_unix_secs: 123,
now_unix_ms: 123,
});
assert_eq!(record.started_at_unix_secs, Some(123));
assert_eq!(record.finished_at_unix_secs, Some(123));
assert_eq!(record.created_at_unix_secs, Some(10));
assert_eq!(record.started_at_unix_ms, Some(123));
assert_eq!(record.finished_at_unix_ms, Some(123));
assert_eq!(record.created_at_unix_ms, Some(10));
assert_eq!(record.status, RequestCandidateStatus::Success);
}
}