From 0a96e11e98a9306010773364f0b91e6007291f7e Mon Sep 17 00:00:00 2001 From: orangels Date: Wed, 6 May 2026 21:22:19 +0800 Subject: [PATCH] feat: run rpm limiter modes --- src/cli.rs | 330 +++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 285 insertions(+), 45 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index a493b61..93dbeb0 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -6,6 +6,10 @@ use crate::report::{ BenchmarkReport, BenchmarkSummaryReport, DatasetReport, LatencyReport, RpmReport, RpmRunReport, RpmSummaryReport, RunReport, WrongCaseReport, write_benchmark_report, write_rpm_report, }; +use crate::rpm_modes::{ + ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule, + sustained_schedule, token_bucket_schedule, window_boundary_plan, +}; use crate::runner::{ModelRequest, run_model_request}; use anyhow::{Context, Result, bail}; use chrono::Utc; @@ -50,15 +54,25 @@ pub enum Command { Rpm { #[arg(long, default_value = "config.yaml")] config: PathBuf, + #[arg(long, value_enum, default_value_t = RpmMode::Sustained)] + mode: RpmMode, #[arg(long)] provider: Option, #[arg(long)] model: Option, #[arg(long)] - rpm: u32, - #[arg(long)] + rpm: Option, + #[arg(long, default_value = "60s")] duration: String, #[arg(long)] + burst: Option, + #[arg(long)] + probe_seconds: Option, + #[arg(long, default_value_t = 500)] + window_offset_ms: u64, + #[arg(long)] + concurrency: Option, + #[arg(long)] prompt: String, }, } @@ -133,12 +147,34 @@ pub async fn dispatch(cli: Cli) -> Result<()> { Command::Bench { command } => dispatch_bench(command).await, Command::Rpm { config, + mode, provider, model, rpm, duration, + burst, + probe_seconds, + window_offset_ms, + concurrency, prompt, - } => run_rpm(config, provider, model, rpm, duration, prompt).await, + } => { + run_rpm( + config, + RpmCommandOptions { + mode, + provider, + model, + rpm, + duration, + burst, + probe_seconds, + window_offset_ms, + concurrency, + prompt, + }, + ) + .await + } } } @@ -325,46 +361,54 @@ async fn run_gpqa_benchmark( Ok(()) } -async fn run_rpm( - config_path: PathBuf, +struct RpmCommandOptions { + mode: RpmMode, provider: Option, model: Option, - rpm: u32, + rpm: Option, duration: String, + burst: Option, + probe_seconds: Option, + window_offset_ms: u64, + concurrency: Option, prompt: String, -) -> Result<()> { - if rpm == 0 { - bail!("rpm must be greater than 0"); - } - let duration = parse_duration(&duration)?; +} + +async fn run_rpm(config_path: PathBuf, options: RpmCommandOptions) -> Result<()> { + let mode_plan = build_rpm_mode_plan( + options.mode, + options.rpm, + &options.duration, + options.burst, + options.probe_seconds, + options.window_offset_ms, + )?; let config = AppConfig::load(&config_path)?; - let provider_name = provider_name(&config, provider.as_deref())?; + let provider_name = provider_name(&config, options.provider.as_deref())?; let provider_config = config.resolved_provider(Some(&provider_name))?; - let model = model.unwrap_or_else(|| provider_config.default_model.clone()); + let model = options + .model + .unwrap_or_else(|| provider_config.default_model.clone()); let request = ModelRequest { - prompt, + prompt: options.prompt, ..request_template(&provider_config, &model, 0.0, 1024) }; - let schedule = rpm_start_schedule(duration, rpm); let started_at = Utc::now(); let started = Instant::now(); - let tokio_started = TokioInstant::now(); - let in_flight_limit = schedule.len().max(1); let mut metrics = Metrics::new(); - let results = stream::iter(schedule.into_iter().map(|offset| { - let request = request.clone(); - async move { - sleep_until(tokio_started + offset).await; - run_model_request(provider_config.protocol, request).await - } - })) - .buffer_unordered(in_flight_limit) - .collect::>() + let results = run_scheduled_requests( + provider_config.protocol, + request, + mode_plan.probes, + options.concurrency.unwrap_or(mode_plan.default_concurrency), + ) .await; for result in results { - match result { + let _phase = result.phase; + let _second = result.second; + match result.result { Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64), Err(error) => metrics.record_failure(error_code(&error)), } @@ -378,7 +422,7 @@ async fn run_rpm( run: RpmRunReport { started_at, duration_ms: started.elapsed().as_millis(), - target_rpm: rpm, + target_rpm: mode_plan.target_rpm, temperature: 0.0, max_tokens: 1024, }, @@ -395,6 +439,119 @@ async fn run_rpm( Ok(()) } +#[derive(Debug)] +struct RpmModePlan { + target_rpm: u32, + probes: Vec, + default_concurrency: usize, +} + +#[derive(Debug)] +struct ScheduledResult { + phase: ProbePhase, + second: Option, + result: Result, +} + +fn build_rpm_mode_plan( + mode: RpmMode, + rpm: Option, + duration: &str, + burst: Option, + probe_seconds: Option, + window_offset_ms: u64, +) -> Result { + let target_rpm = rpm.unwrap_or(0); + let burst = burst.unwrap_or(target_rpm); + let probes = match mode { + RpmMode::Sustained => { + let rpm = require_positive("rpm", rpm)?; + let duration = parse_duration(duration)?; + sustained_schedule(duration, rpm) + .into_iter() + .map(|offset| ScheduledProbe { + offset, + phase: ProbePhase::RefillProbe, + second: Some(offset.as_secs()), + }) + .collect() + } + RpmMode::Burst => { + let burst = require_positive_value("burst", burst)?; + burst_schedule(burst) + .into_iter() + .map(|offset| ScheduledProbe { + offset, + phase: ProbePhase::Burst, + second: Some(0), + }) + .collect() + } + RpmMode::TokenBucket => { + let rpm = require_positive("rpm", rpm)?; + let burst = require_positive_value("burst", burst)?; + token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(30)) + } + RpmMode::SlidingWindow => { + let burst = require_positive_value("burst", burst)?; + sliding_window_schedule(burst, probe_seconds.unwrap_or(90)) + } + RpmMode::WindowBoundary => { + let burst = require_positive_value("burst", burst)?; + window_boundary_plan(Utc::now(), burst, window_offset_ms).probes + } + RpmMode::Diagnose => { + let rpm = require_positive("rpm", rpm)?; + let burst = require_positive_value("burst", burst)?; + let mut probes = token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(90)); + probes.extend(window_boundary_plan(Utc::now(), burst, window_offset_ms).probes); + probes + } + }; + + let default_concurrency = probes.len().max(1); + Ok(RpmModePlan { + target_rpm, + default_concurrency, + probes, + }) +} + +fn require_positive(name: &str, value: Option) -> Result { + let value = value.with_context(|| format!("{name} is required for this rpm mode"))?; + require_positive_value(name, value) +} + +fn require_positive_value(name: &str, value: u32) -> Result { + if value == 0 { + bail!("{name} must be greater than 0"); + } + Ok(value) +} + +async fn run_scheduled_requests( + protocol: crate::config::ProtocolKind, + request: ModelRequest, + starts: Vec, + max_in_flight: usize, +) -> Vec { + let tokio_started = TokioInstant::now(); + stream::iter(starts.into_iter().map(|start| { + let request = request.clone(); + async move { + sleep_until(tokio_started + start.offset).await; + ScheduledResult { + phase: start.phase, + second: start.second, + result: run_model_request(protocol, request).await, + } + } + })) + .buffer_unordered(nonzero_concurrency(max_in_flight)) + .collect::>() + .await +} + fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result { match provider { Some(provider) => Ok(provider.to_string()), @@ -590,22 +747,6 @@ fn parse_duration(value: &str) -> Result { Ok(Duration::from_secs(seconds)) } -fn rpm_start_schedule(duration: Duration, rpm: u32) -> Vec { - let interval_nanos = 60_000_000_000u128 / u128::from(rpm); - if interval_nanos == 0 { - return Vec::new(); - } - - let duration_nanos = duration.as_nanos(); - let mut offsets = Vec::new(); - let mut offset = 0u128; - while offset < duration_nanos { - offsets.push(Duration::from_nanos(offset.min(u128::from(u64::MAX)) as u64)); - offset += interval_nanos; - } - offsets -} - fn dataset_data_dir(config_path: &Path) -> Result { if !config_path.exists() { return Ok(PathBuf::from("data/benchmarks")); @@ -648,7 +789,7 @@ mod tests { #[test] fn computes_rpm_start_schedule_from_run_start() { - let schedule = rpm_start_schedule(Duration::from_secs(60), 120); + let schedule = sustained_schedule(Duration::from_secs(60), 120); assert_eq!(schedule.len(), 120); assert_eq!(schedule[0], Duration::ZERO); @@ -662,4 +803,103 @@ mod tests { assert!(error.to_string().contains("duration")); } + + #[test] + fn rpm_command_defaults_to_sustained_mode() { + let cli = Cli::try_parse_from([ + "lq_token_test", + "rpm", + "--provider", + "anthropic", + "--rpm", + "120", + "--prompt", + "hello", + ]) + .expect("parse rpm"); + + let Command::Rpm { mode, rpm, .. } = cli.command else { + panic!("expected rpm command"); + }; + + assert_eq!(mode, RpmMode::Sustained); + assert_eq!(rpm, Some(120)); + } + + #[test] + fn rpm_command_parses_token_bucket_mode() { + let cli = Cli::try_parse_from([ + "lq_token_test", + "rpm", + "--mode", + "token-bucket", + "--provider", + "anthropic", + "--rpm", + "120", + "--burst", + "120", + "--probe-seconds", + "30", + "--prompt", + "hello", + ]) + .expect("parse token bucket rpm"); + + let Command::Rpm { + mode, + burst, + probe_seconds, + .. + } = cli.command + else { + panic!("expected rpm command"); + }; + + assert_eq!(mode, RpmMode::TokenBucket); + assert_eq!(burst, Some(120)); + assert_eq!(probe_seconds, Some(30)); + } + + #[test] + fn rpm_command_parses_window_boundary_offset() { + let cli = Cli::try_parse_from([ + "lq_token_test", + "rpm", + "--mode", + "window-boundary", + "--provider", + "anthropic", + "--burst", + "10", + "--window-offset-ms", + "250", + "--prompt", + "hello", + ]) + .expect("parse window boundary rpm"); + + let Command::Rpm { + mode, + window_offset_ms, + .. + } = cli.command + else { + panic!("expected rpm command"); + }; + + assert_eq!(mode, RpmMode::WindowBoundary); + assert_eq!(window_offset_ms, 250); + } + + #[test] + fn rpm_mode_validation_rejects_zero_values() { + let zero_rpm = build_rpm_mode_plan(RpmMode::Sustained, Some(0), "60s", None, None, 500) + .expect_err("zero rpm should fail"); + let zero_burst = build_rpm_mode_plan(RpmMode::Burst, None, "60s", Some(0), None, 500) + .expect_err("zero burst should fail"); + + assert!(zero_rpm.to_string().contains("rpm")); + assert!(zero_burst.to_string().contains("burst")); + } }