mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Make pool probing request-driven
This commit is contained in:
@@ -407,7 +407,6 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
||||
overload_cooldown_seconds: 30,
|
||||
health_policy_enabled: true,
|
||||
probing_enabled: false,
|
||||
probing_interval_minutes: 10,
|
||||
probing_target_percent: None,
|
||||
probing_target_count: None,
|
||||
probe_concurrency: 4,
|
||||
@@ -475,12 +474,6 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
||||
.get("probing_enabled")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
probing_interval_minutes: pool_advanced
|
||||
.get("probing_interval_minutes")
|
||||
.and_then(json_u64)
|
||||
.filter(|value| *value > 0)
|
||||
.map(|value| value.min(1440))
|
||||
.unwrap_or(10),
|
||||
probing_target_percent: parse_pool_probe_target_percent(pool_advanced),
|
||||
probing_target_count: parse_pool_probe_target_count(pool_advanced),
|
||||
probe_concurrency: pool_advanced
|
||||
@@ -592,7 +585,6 @@ mod tests {
|
||||
"overload_cooldown_seconds": 45,
|
||||
"health_policy_enabled": false,
|
||||
"probing_enabled": true,
|
||||
"probing_interval_minutes": 20,
|
||||
"probing_target_percent": 25,
|
||||
"probing_target_count": 3,
|
||||
"probe_concurrency": 6,
|
||||
@@ -634,7 +626,6 @@ mod tests {
|
||||
assert_eq!(config.overload_cooldown_seconds, 45);
|
||||
assert!(!config.health_policy_enabled);
|
||||
assert!(config.probing_enabled);
|
||||
assert_eq!(config.probing_interval_minutes, 20);
|
||||
assert_eq!(config.probing_target_percent, Some(25.0));
|
||||
assert_eq!(config.probing_target_count, Some(3));
|
||||
assert_eq!(config.probe_concurrency, 6);
|
||||
@@ -656,7 +647,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clamps_pool_quota_probe_interval_to_python_range() {
|
||||
fn ignores_legacy_pool_quota_probe_interval() {
|
||||
let provider = sample_provider(json!({
|
||||
"pool_advanced": {
|
||||
"probing_enabled": true,
|
||||
@@ -664,7 +655,7 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
|
||||
assert_eq!(config.probing_interval_minutes, 1440);
|
||||
assert!(config.probing_enabled);
|
||||
|
||||
let provider = sample_provider(json!({
|
||||
"pool_advanced": {
|
||||
@@ -673,7 +664,7 @@ mod tests {
|
||||
}
|
||||
}));
|
||||
let config = admin_provider_pool_config(&provider).expect("pool config should exist");
|
||||
assert_eq!(config.probing_interval_minutes, 10);
|
||||
assert!(config.probing_enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -614,7 +614,6 @@ mod tests {
|
||||
overload_cooldown_seconds: 30,
|
||||
health_policy_enabled: true,
|
||||
probing_enabled: false,
|
||||
probing_interval_minutes: 10,
|
||||
probing_target_percent: None,
|
||||
probing_target_count: None,
|
||||
probe_concurrency: 4,
|
||||
|
||||
@@ -45,7 +45,6 @@ pub(crate) struct AdminProviderPoolConfig {
|
||||
pub(crate) overload_cooldown_seconds: u64,
|
||||
pub(crate) health_policy_enabled: bool,
|
||||
pub(crate) probing_enabled: bool,
|
||||
pub(crate) probing_interval_minutes: u64,
|
||||
pub(crate) probing_target_percent: Option<f64>,
|
||||
pub(crate) probing_target_count: Option<u64>,
|
||||
pub(crate) probe_concurrency: u64,
|
||||
|
||||
@@ -29,7 +29,7 @@ use crate::handlers::shared::provider_pool::{
|
||||
};
|
||||
use crate::provider_pool_demand::{
|
||||
provider_pool_burst_pending_key, read_provider_pool_demand_snapshot,
|
||||
sample_provider_pool_demand,
|
||||
sample_provider_pool_demand, ProviderPoolDemandSnapshot,
|
||||
};
|
||||
|
||||
use super::pool_score_rebuild::ensure_provider_key_pool_scores_for_keys;
|
||||
@@ -44,6 +44,9 @@ const POOL_QUOTA_PROBE_BURST_TRIGGER_LOCK_TTL_MS: u64 = 30_000;
|
||||
const POOL_QUOTA_PROBE_BURST_PENDING_PREFIX: &str = "ap:quota_probe:burst_pending";
|
||||
const POOL_QUOTA_PROBE_BURST_PENDING_TTL_SECONDS: u64 = 30;
|
||||
const POOL_QUOTA_PROBE_BURST_RETRY_GUARD_SECONDS: u64 = 15;
|
||||
const POOL_QUOTA_PROBE_AUTO_MIN_INTERVAL_SECONDS: u64 = 30;
|
||||
const POOL_QUOTA_PROBE_AUTO_MAX_INTERVAL_SECONDS: u64 = 10 * 60;
|
||||
const POOL_QUOTA_PROBE_AUTO_MAX_PRESSURE: u64 = 64;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum PoolQuotaProbeMode {
|
||||
@@ -214,6 +217,34 @@ fn pool_quota_probe_selection_limit_for_mode(
|
||||
}
|
||||
}
|
||||
|
||||
fn pool_quota_probe_auto_interval_seconds(demand_snapshot: &ProviderPoolDemandSnapshot) -> u64 {
|
||||
let live_pressure = u64::try_from(demand_snapshot.in_flight).unwrap_or(u64::MAX);
|
||||
let ema_pressure =
|
||||
if demand_snapshot.ema_in_flight.is_finite() && demand_snapshot.ema_in_flight > 0.0 {
|
||||
demand_snapshot
|
||||
.ema_in_flight
|
||||
.ceil()
|
||||
.clamp(0.0, POOL_QUOTA_PROBE_AUTO_MAX_PRESSURE as f64) as u64
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let request_pressure = live_pressure.max(ema_pressure);
|
||||
if request_pressure == 0 {
|
||||
return POOL_QUOTA_PROBE_AUTO_MAX_INTERVAL_SECONDS;
|
||||
}
|
||||
|
||||
let hot_pressure = u64::try_from(demand_snapshot.desired_hot).unwrap_or(u64::MAX);
|
||||
let pressure = request_pressure
|
||||
.max(hot_pressure)
|
||||
.clamp(1, POOL_QUOTA_PROBE_AUTO_MAX_PRESSURE);
|
||||
POOL_QUOTA_PROBE_AUTO_MAX_INTERVAL_SECONDS
|
||||
.saturating_div(pressure)
|
||||
.clamp(
|
||||
POOL_QUOTA_PROBE_AUTO_MIN_INTERVAL_SECONDS,
|
||||
POOL_QUOTA_PROBE_AUTO_MAX_INTERVAL_SECONDS,
|
||||
)
|
||||
}
|
||||
|
||||
fn active_probe_member_remains_valid(score: Option<&StoredPoolMemberScore>) -> bool {
|
||||
match score.map(|score| score.hard_state) {
|
||||
Some(PoolMemberHardState::Available | PoolMemberHardState::Unknown) | None => true,
|
||||
@@ -894,10 +925,7 @@ async fn select_keys_for_provider(
|
||||
.get(key_id.as_str())
|
||||
.is_none_or(|last_probe_ts| {
|
||||
now_ts.saturating_sub(*last_probe_ts)
|
||||
>= pool_config
|
||||
.probing_interval_minutes
|
||||
.clamp(1, 1440)
|
||||
.saturating_mul(60)
|
||||
>= pool_quota_probe_auto_interval_seconds(&demand_snapshot)
|
||||
})
|
||||
}
|
||||
PoolQuotaProbeMode::Burst => {
|
||||
@@ -925,10 +953,7 @@ async fn select_keys_for_provider(
|
||||
}
|
||||
|
||||
let stamp_interval_seconds = match mode {
|
||||
PoolQuotaProbeMode::Base => pool_config
|
||||
.probing_interval_minutes
|
||||
.clamp(1, 1440)
|
||||
.saturating_mul(60),
|
||||
PoolQuotaProbeMode::Base => pool_quota_probe_auto_interval_seconds(&demand_snapshot),
|
||||
PoolQuotaProbeMode::Burst => POOL_QUOTA_PROBE_BURST_PENDING_TTL_SECONDS,
|
||||
};
|
||||
mark_probe_timestamps(
|
||||
@@ -1756,6 +1781,42 @@ mod tests {
|
||||
assert_eq!(pool_quota_probe_target_count(3, Some(80.0), Some(10)), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pool_quota_probe_auto_interval_tracks_request_pressure() {
|
||||
let idle = ProviderPoolDemandSnapshot {
|
||||
in_flight: 0,
|
||||
ema_in_flight: 0.0,
|
||||
desired_hot: 2,
|
||||
sampled_at_unix_ms: 0,
|
||||
};
|
||||
assert_eq!(
|
||||
pool_quota_probe_auto_interval_seconds(&idle),
|
||||
POOL_QUOTA_PROBE_AUTO_MAX_INTERVAL_SECONDS
|
||||
);
|
||||
|
||||
let active = ProviderPoolDemandSnapshot {
|
||||
in_flight: 1,
|
||||
ema_in_flight: 1.0,
|
||||
desired_hot: 2,
|
||||
sampled_at_unix_ms: 0,
|
||||
};
|
||||
assert!(
|
||||
pool_quota_probe_auto_interval_seconds(&active)
|
||||
< POOL_QUOTA_PROBE_AUTO_MAX_INTERVAL_SECONDS
|
||||
);
|
||||
|
||||
let saturated = ProviderPoolDemandSnapshot {
|
||||
in_flight: 128,
|
||||
ema_in_flight: 128.0,
|
||||
desired_hot: 128,
|
||||
sampled_at_unix_ms: 0,
|
||||
};
|
||||
assert_eq!(
|
||||
pool_quota_probe_auto_interval_seconds(&saturated),
|
||||
POOL_QUOTA_PROBE_AUTO_MIN_INTERVAL_SECONDS
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_only_pool_out_keys_to_fill_active_probe_target() {
|
||||
let key_ids = vec![
|
||||
|
||||
@@ -47,7 +47,6 @@ use crate::maintenance::spawn_gemini_file_mapping_cleanup_worker;
|
||||
use crate::maintenance::spawn_oauth_token_refresh_worker;
|
||||
use crate::maintenance::spawn_pending_cleanup_worker;
|
||||
use crate::maintenance::spawn_pool_monitor_worker;
|
||||
use crate::maintenance::spawn_pool_quota_probe_worker;
|
||||
use crate::maintenance::spawn_pool_score_rebuild_worker;
|
||||
use crate::maintenance::spawn_provider_checkin_worker;
|
||||
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
|
||||
@@ -1173,10 +1172,6 @@ impl AppState {
|
||||
crate::task_runtime::TASK_KEY_POOL_MONITOR,
|
||||
spawn_pool_monitor_worker(self.data.clone()),
|
||||
);
|
||||
supervise_worker(
|
||||
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
|
||||
spawn_pool_quota_probe_worker(self.clone()),
|
||||
);
|
||||
supervise_worker(
|
||||
crate::task_runtime::TASK_KEY_ACCOUNT_SELF_CHECK,
|
||||
spawn_account_self_check_worker(self.clone()),
|
||||
|
||||
@@ -21,7 +21,6 @@ pub(crate) const TASK_KEY_USAGE_QUEUE_WORKER: &str = "usage.queue.worker";
|
||||
pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
|
||||
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
|
||||
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
|
||||
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
|
||||
pub(crate) const TASK_KEY_ACCOUNT_SELF_CHECK: &str = "account.self_check.worker";
|
||||
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
|
||||
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
|
||||
@@ -95,14 +94,6 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_POOL_QUOTA_PROBE,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_ACCOUNT_SELF_CHECK,
|
||||
TaskKind::Scheduled,
|
||||
|
||||
Reference in New Issue
Block a user