| @@ -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<String>, | |||
| #[arg(long)] | |||
| model: Option<String>, | |||
| #[arg(long)] | |||
| rpm: u32, | |||
| #[arg(long)] | |||
| rpm: Option<u32>, | |||
| #[arg(long, default_value = "60s")] | |||
| duration: String, | |||
| #[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, | |||
| }, | |||
| } | |||
| @@ -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<String>, | |||
| model: Option<String>, | |||
| rpm: u32, | |||
| rpm: Option<u32>, | |||
| duration: String, | |||
| burst: Option<u32>, | |||
| probe_seconds: Option<u64>, | |||
| window_offset_ms: u64, | |||
| concurrency: Option<usize>, | |||
| 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::<Vec<_>>() | |||
| 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<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> { | |||
| match provider { | |||
| Some(provider) => Ok(provider.to_string()), | |||
| @@ -590,22 +747,6 @@ fn parse_duration(value: &str) -> Result<Duration> { | |||
| 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> { | |||
| 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")); | |||
| } | |||
| } | |||