Browse Source

feat: run rpm limiter modes

main
orangels 1 week ago
parent
commit
0a96e11e98
1 changed files with 285 additions and 45 deletions
  1. +285
    -45
      src/cli.rs

+ 285
- 45
src/cli.rs View File

@@ -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"));
}
}

Loading…
Cancel
Save