| @@ -6,6 +6,10 @@ use crate::report::{ | |||||
| BenchmarkReport, BenchmarkSummaryReport, DatasetReport, LatencyReport, RpmReport, RpmRunReport, | BenchmarkReport, BenchmarkSummaryReport, DatasetReport, LatencyReport, RpmReport, RpmRunReport, | ||||
| RpmSummaryReport, RunReport, WrongCaseReport, write_benchmark_report, write_rpm_report, | 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 crate::runner::{ModelRequest, run_model_request}; | ||||
| use anyhow::{Context, Result, bail}; | use anyhow::{Context, Result, bail}; | ||||
| use chrono::Utc; | use chrono::Utc; | ||||
| @@ -50,15 +54,25 @@ pub enum Command { | |||||
| Rpm { | Rpm { | ||||
| #[arg(long, default_value = "config.yaml")] | #[arg(long, default_value = "config.yaml")] | ||||
| config: PathBuf, | config: PathBuf, | ||||
| #[arg(long, value_enum, default_value_t = RpmMode::Sustained)] | |||||
| mode: RpmMode, | |||||
| #[arg(long)] | #[arg(long)] | ||||
| provider: Option<String>, | provider: Option<String>, | ||||
| #[arg(long)] | #[arg(long)] | ||||
| model: Option<String>, | model: Option<String>, | ||||
| #[arg(long)] | #[arg(long)] | ||||
| rpm: u32, | |||||
| #[arg(long)] | |||||
| rpm: Option<u32>, | |||||
| #[arg(long, default_value = "60s")] | |||||
| duration: String, | duration: String, | ||||
| #[arg(long)] | #[arg(long)] | ||||
| burst: Option<u32>, | |||||
| #[arg(long)] | |||||
| probe_seconds: Option<u64>, | |||||
| #[arg(long, default_value_t = 500)] | |||||
| window_offset_ms: u64, | |||||
| #[arg(long)] | |||||
| concurrency: Option<usize>, | |||||
| #[arg(long)] | |||||
| prompt: String, | prompt: String, | ||||
| }, | }, | ||||
| } | } | ||||
| @@ -133,12 +147,34 @@ pub async fn dispatch(cli: Cli) -> Result<()> { | |||||
| Command::Bench { command } => dispatch_bench(command).await, | Command::Bench { command } => dispatch_bench(command).await, | ||||
| Command::Rpm { | Command::Rpm { | ||||
| config, | config, | ||||
| mode, | |||||
| provider, | provider, | ||||
| model, | model, | ||||
| rpm, | rpm, | ||||
| duration, | duration, | ||||
| burst, | |||||
| probe_seconds, | |||||
| window_offset_ms, | |||||
| concurrency, | |||||
| prompt, | 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(()) | Ok(()) | ||||
| } | } | ||||
| async fn run_rpm( | |||||
| config_path: PathBuf, | |||||
| struct RpmCommandOptions { | |||||
| mode: RpmMode, | |||||
| provider: Option<String>, | provider: Option<String>, | ||||
| model: Option<String>, | model: Option<String>, | ||||
| rpm: u32, | |||||
| rpm: Option<u32>, | |||||
| duration: String, | duration: String, | ||||
| burst: Option<u32>, | |||||
| probe_seconds: Option<u64>, | |||||
| window_offset_ms: u64, | |||||
| concurrency: Option<usize>, | |||||
| prompt: String, | 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 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 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 { | let request = ModelRequest { | ||||
| prompt, | |||||
| prompt: options.prompt, | |||||
| ..request_template(&provider_config, &model, 0.0, 1024) | ..request_template(&provider_config, &model, 0.0, 1024) | ||||
| }; | }; | ||||
| let schedule = rpm_start_schedule(duration, rpm); | |||||
| let started_at = Utc::now(); | let started_at = Utc::now(); | ||||
| let started = Instant::now(); | let started = Instant::now(); | ||||
| let tokio_started = TokioInstant::now(); | |||||
| let in_flight_limit = schedule.len().max(1); | |||||
| let mut metrics = Metrics::new(); | 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::<Vec<_>>() | |||||
| let results = run_scheduled_requests( | |||||
| provider_config.protocol, | |||||
| request, | |||||
| mode_plan.probes, | |||||
| options.concurrency.unwrap_or(mode_plan.default_concurrency), | |||||
| ) | |||||
| .await; | .await; | ||||
| for result in results { | 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), | Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64), | ||||
| Err(error) => metrics.record_failure(error_code(&error)), | Err(error) => metrics.record_failure(error_code(&error)), | ||||
| } | } | ||||
| @@ -378,7 +422,7 @@ async fn run_rpm( | |||||
| run: RpmRunReport { | run: RpmRunReport { | ||||
| started_at, | started_at, | ||||
| duration_ms: started.elapsed().as_millis(), | duration_ms: started.elapsed().as_millis(), | ||||
| target_rpm: rpm, | |||||
| target_rpm: mode_plan.target_rpm, | |||||
| temperature: 0.0, | temperature: 0.0, | ||||
| max_tokens: 1024, | max_tokens: 1024, | ||||
| }, | }, | ||||
| @@ -395,6 +439,119 @@ async fn run_rpm( | |||||
| Ok(()) | Ok(()) | ||||
| } | } | ||||
| #[derive(Debug)] | |||||
| struct RpmModePlan { | |||||
| target_rpm: u32, | |||||
| probes: Vec<ScheduledProbe>, | |||||
| default_concurrency: usize, | |||||
| } | |||||
| #[derive(Debug)] | |||||
| struct ScheduledResult { | |||||
| phase: ProbePhase, | |||||
| second: Option<u64>, | |||||
| result: Result<crate::runner::ModelResponse>, | |||||
| } | |||||
| fn build_rpm_mode_plan( | |||||
| mode: RpmMode, | |||||
| rpm: Option<u32>, | |||||
| duration: &str, | |||||
| burst: Option<u32>, | |||||
| probe_seconds: Option<u64>, | |||||
| window_offset_ms: u64, | |||||
| ) -> Result<RpmModePlan> { | |||||
| 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<u32>) -> Result<u32> { | |||||
| 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<u32> { | |||||
| 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<ScheduledProbe>, | |||||
| max_in_flight: usize, | |||||
| ) -> Vec<ScheduledResult> { | |||||
| 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::<Vec<_>>() | |||||
| .await | |||||
| } | |||||
| fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result<String> { | fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result<String> { | ||||
| match provider { | match provider { | ||||
| Some(provider) => Ok(provider.to_string()), | Some(provider) => Ok(provider.to_string()), | ||||
| @@ -590,22 +747,6 @@ fn parse_duration(value: &str) -> Result<Duration> { | |||||
| Ok(Duration::from_secs(seconds)) | Ok(Duration::from_secs(seconds)) | ||||
| } | } | ||||
| fn rpm_start_schedule(duration: Duration, rpm: u32) -> Vec<Duration> { | |||||
| 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<PathBuf> { | fn dataset_data_dir(config_path: &Path) -> Result<PathBuf> { | ||||
| if !config_path.exists() { | if !config_path.exists() { | ||||
| return Ok(PathBuf::from("data/benchmarks")); | return Ok(PathBuf::from("data/benchmarks")); | ||||
| @@ -648,7 +789,7 @@ mod tests { | |||||
| #[test] | #[test] | ||||
| fn computes_rpm_start_schedule_from_run_start() { | 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.len(), 120); | ||||
| assert_eq!(schedule[0], Duration::ZERO); | assert_eq!(schedule[0], Duration::ZERO); | ||||
| @@ -662,4 +803,103 @@ mod tests { | |||||
| assert!(error.to_string().contains("duration")); | 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")); | |||||
| } | |||||
| } | } | ||||