Make pool probing request-driven

This commit is contained in:
fawney19
2026-05-16 00:18:47 +08:00
parent 43d891bee1
commit d53546d56f
10 changed files with 78 additions and 69 deletions

View File

@@ -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]

View File

@@ -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,

View File

@@ -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,

View File

@@ -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![

View File

@@ -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()),

View File

@@ -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,