api模型检测
Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

1265 строки
39 KiB

  1. use crate::benchmarks;
  2. use crate::benchmarks::judge;
  3. use crate::config::AppConfig;
  4. use crate::metrics::{LatencySummary, Metrics, MetricsSummary};
  5. use crate::report::{
  6. BenchmarkParamsReport, BenchmarkReport, BenchmarkSummaryReport, CorrectCaseReport,
  7. DatasetReport, LatencyReport, LimiterInferenceKind, LimiterInferenceReport,
  8. PhaseSummaryReport, ProbeSecondReport, RpmModeDetailReport, RpmParamsReport, RpmReport,
  9. RpmRunReport, RpmSummaryReport, RunReport, WindowBoundaryReport, WrongCaseReport,
  10. write_benchmark_report, write_rpm_report,
  11. };
  12. use crate::rpm_modes::{
  13. ProbePhase, RpmMode, ScheduledProbe, burst_schedule, sliding_window_schedule,
  14. sustained_schedule, token_bucket_schedule, window_boundary_plan,
  15. };
  16. use crate::runner::{ModelRequest, run_model_request};
  17. use anyhow::{Context, Result, bail};
  18. use chrono::Utc;
  19. use clap::{Parser, Subcommand};
  20. use futures::{StreamExt, stream};
  21. use indicatif::{ProgressBar, ProgressStyle};
  22. use regex::Regex;
  23. use std::collections::BTreeMap;
  24. use std::path::{Path, PathBuf};
  25. use std::time::{Duration, Instant};
  26. use tokio::time::{Instant as TokioInstant, sleep_until};
  27. #[derive(Debug, Parser)]
  28. #[command(
  29. name = "lq_token_test",
  30. version,
  31. about = "Test LLM relay protocols, RPM, and benchmark accuracy"
  32. )]
  33. pub struct Cli {
  34. #[command(subcommand)]
  35. pub command: Command,
  36. }
  37. #[derive(Debug, Subcommand)]
  38. pub enum Command {
  39. Check {
  40. #[arg(long, default_value = "config.yaml")]
  41. config: PathBuf,
  42. #[arg(long)]
  43. provider: Option<String>,
  44. #[arg(long)]
  45. model: Option<String>,
  46. #[arg(long)]
  47. prompt: String,
  48. #[arg(long)]
  49. stream: Option<bool>,
  50. },
  51. Dataset {
  52. #[command(subcommand)]
  53. command: DatasetCommand,
  54. },
  55. Bench {
  56. #[command(subcommand)]
  57. command: BenchCommand,
  58. },
  59. Rpm {
  60. #[arg(long, default_value = "config.yaml")]
  61. config: PathBuf,
  62. #[arg(long, value_enum, default_value_t = RpmMode::Sustained)]
  63. mode: RpmMode,
  64. #[arg(long)]
  65. provider: Option<String>,
  66. #[arg(long)]
  67. model: Option<String>,
  68. #[arg(long)]
  69. rpm: Option<u32>,
  70. #[arg(long, default_value = "60s")]
  71. duration: String,
  72. #[arg(long)]
  73. burst: Option<u32>,
  74. #[arg(long)]
  75. probe_seconds: Option<u64>,
  76. #[arg(long, default_value_t = 500)]
  77. window_offset_ms: u64,
  78. #[arg(long)]
  79. concurrency: Option<usize>,
  80. #[arg(long)]
  81. prompt: String,
  82. #[arg(long)]
  83. stream: Option<bool>,
  84. },
  85. }
  86. #[derive(Debug, Subcommand)]
  87. pub enum DatasetCommand {
  88. Fetch { dataset: String },
  89. }
  90. #[derive(Debug, Subcommand)]
  91. pub enum BenchCommand {
  92. Aime2026 {
  93. #[arg(long, default_value = "config.yaml")]
  94. config: PathBuf,
  95. #[arg(long)]
  96. provider: Option<String>,
  97. #[arg(long)]
  98. model: Option<String>,
  99. #[arg(long, default_value_t = 4)]
  100. concurrency: usize,
  101. #[arg(long)]
  102. limit: Option<usize>,
  103. #[arg(long)]
  104. stream: Option<bool>,
  105. #[arg(long, default_value_t = 32768)]
  106. max_tokens: u32,
  107. },
  108. GpqaDiamond {
  109. #[arg(long, default_value = "config.yaml")]
  110. config: PathBuf,
  111. #[arg(long)]
  112. provider: Option<String>,
  113. #[arg(long)]
  114. model: Option<String>,
  115. #[arg(long, default_value_t = 4)]
  116. concurrency: usize,
  117. #[arg(long)]
  118. limit: Option<usize>,
  119. #[arg(long)]
  120. stream: Option<bool>,
  121. #[arg(long, default_value_t = 32768)]
  122. max_tokens: u32,
  123. },
  124. }
  125. pub async fn dispatch(cli: Cli) -> Result<()> {
  126. match cli.command {
  127. Command::Check {
  128. config,
  129. provider,
  130. model,
  131. prompt,
  132. stream,
  133. } => {
  134. let config = AppConfig::load(&config)?;
  135. let provider = config.resolved_provider(provider.as_deref())?;
  136. let request = ModelRequest {
  137. base_url: provider.base_url.clone(),
  138. api_token: provider.api_token.clone(),
  139. model: model.unwrap_or_else(|| provider.default_model.clone()),
  140. prompt,
  141. temperature: 0.0,
  142. max_tokens: 1024,
  143. stream: stream.unwrap_or(provider.stream),
  144. };
  145. let response = run_model_request(provider.protocol, request).await?;
  146. println!("status: {}", response.status);
  147. println!("elapsed_ms: {}", response.elapsed_ms);
  148. if let Some(ttft) = response.first_token_ms {
  149. println!("first_token_ms: {}", ttft);
  150. }
  151. println!("{}", response.text);
  152. Ok(())
  153. }
  154. Command::Dataset {
  155. command: DatasetCommand::Fetch { dataset },
  156. } => {
  157. let data_dir = dataset_data_dir(Path::new("config.yaml"))?;
  158. let path = benchmarks::fetch_dataset(&dataset, &data_dir).await?;
  159. println!("{}", path.display());
  160. Ok(())
  161. }
  162. Command::Bench { command } => dispatch_bench(command).await,
  163. Command::Rpm {
  164. config,
  165. mode,
  166. provider,
  167. model,
  168. rpm,
  169. duration,
  170. burst,
  171. probe_seconds,
  172. window_offset_ms,
  173. concurrency,
  174. prompt,
  175. stream,
  176. } => {
  177. run_rpm(
  178. config,
  179. RpmCommandOptions {
  180. mode,
  181. provider,
  182. model,
  183. rpm,
  184. duration,
  185. burst,
  186. probe_seconds,
  187. window_offset_ms,
  188. concurrency,
  189. prompt,
  190. stream,
  191. },
  192. )
  193. .await
  194. }
  195. }
  196. }
  197. async fn dispatch_bench(command: BenchCommand) -> Result<()> {
  198. match command {
  199. BenchCommand::Aime2026 {
  200. config,
  201. provider,
  202. model,
  203. concurrency,
  204. limit,
  205. stream,
  206. max_tokens,
  207. } => run_aime_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await,
  208. BenchCommand::GpqaDiamond {
  209. config,
  210. provider,
  211. model,
  212. concurrency,
  213. limit,
  214. stream,
  215. max_tokens,
  216. } => run_gpqa_benchmark(config, provider, model, concurrency, limit, stream, max_tokens).await,
  217. }
  218. }
  219. async fn run_aime_benchmark(
  220. config_path: PathBuf,
  221. provider: Option<String>,
  222. model: Option<String>,
  223. concurrency: usize,
  224. limit: Option<usize>,
  225. stream: Option<bool>,
  226. max_tokens: u32,
  227. ) -> Result<()> {
  228. let config = AppConfig::load(&config_path)?;
  229. let provider_name = provider_name(&config, provider.as_deref())?;
  230. let provider_config = config.resolved_provider(Some(&provider_name))?;
  231. let model = model.unwrap_or_else(|| provider_config.default_model.clone());
  232. let loaded = benchmarks::aime::load_cases(Path::new(&config.benchmarks.data_dir))?;
  233. let dataset = dataset_report(
  234. config
  235. .benchmarks
  236. .aime2026
  237. .as_ref()
  238. .map(|dataset| (dataset.source.as_str(), dataset.split.as_str())),
  239. &loaded.local_path,
  240. );
  241. let cases = apply_limit(loaded.cases, limit);
  242. let total = cases.len() as u64;
  243. let started_at = Utc::now();
  244. let started = Instant::now();
  245. let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens);
  246. base_request.stream = stream.unwrap_or(provider_config.stream);
  247. let protocol = provider_config.protocol;
  248. let pb = ProgressBar::new(total);
  249. pb.set_style(
  250. ProgressStyle::default_bar()
  251. .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
  252. .expect("valid template"),
  253. );
  254. let mut results = stream::iter(cases)
  255. .map(|case| {
  256. let mut request = base_request.clone();
  257. request.prompt = case.prompt();
  258. async move {
  259. let result = run_model_request(protocol, request).await;
  260. (case, result)
  261. }
  262. })
  263. .buffer_unordered(nonzero_concurrency(concurrency));
  264. let mut metrics = Metrics::new();
  265. let mut wrong_cases = Vec::new();
  266. let mut correct_samples = Vec::new();
  267. while let Some((case, result)) = results.next().await {
  268. pb.inc(1);
  269. match result {
  270. Ok(response) => {
  271. metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64));
  272. let actual = judge::extract_final_integer(&response.text)
  273. .unwrap_or_else(|| "no_answer".to_string());
  274. let correct = judge::judge_integer(&response.text, &case.answer);
  275. metrics.record_judgement(correct);
  276. if correct {
  277. if correct_samples.len() < 5 {
  278. correct_samples.push(CorrectCaseReport {
  279. id: case.id,
  280. question: case.problem,
  281. expected: case.answer,
  282. raw_output: response.text,
  283. });
  284. }
  285. } else {
  286. wrong_cases.push(WrongCaseReport {
  287. id: case.id,
  288. question: case.problem,
  289. expected: case.answer,
  290. actual,
  291. raw_output: response.text,
  292. });
  293. }
  294. }
  295. Err(error) => metrics.record_failure(error_code(&error)),
  296. }
  297. }
  298. pb.finish_and_clear();
  299. let summary = metrics.summary();
  300. let report = benchmark_report(BenchmarkReportInput {
  301. benchmark: "aime2026",
  302. provider: provider_name,
  303. model,
  304. stream: base_request.stream,
  305. dataset,
  306. started_at,
  307. duration_ms: started.elapsed().as_millis(),
  308. concurrency,
  309. limit,
  310. max_tokens,
  311. summary,
  312. correct_samples,
  313. wrong_cases,
  314. });
  315. let report_path = write_benchmark_report(Path::new("."), &report)?;
  316. print_benchmark_report(&report, &report_path);
  317. Ok(())
  318. }
  319. async fn run_gpqa_benchmark(
  320. config_path: PathBuf,
  321. provider: Option<String>,
  322. model: Option<String>,
  323. concurrency: usize,
  324. limit: Option<usize>,
  325. stream: Option<bool>,
  326. max_tokens: u32,
  327. ) -> Result<()> {
  328. let config = AppConfig::load(&config_path)?;
  329. let provider_name = provider_name(&config, provider.as_deref())?;
  330. let provider_config = config.resolved_provider(Some(&provider_name))?;
  331. let model = model.unwrap_or_else(|| provider_config.default_model.clone());
  332. let loaded = benchmarks::gpqa::load_cases(Path::new(&config.benchmarks.data_dir))?;
  333. let dataset = dataset_report(
  334. config
  335. .benchmarks
  336. .gpqa_diamond
  337. .as_ref()
  338. .map(|dataset| (dataset.source.as_str(), dataset.split.as_str())),
  339. &loaded.local_path,
  340. );
  341. let cases = apply_limit(loaded.cases, limit);
  342. let total = cases.len() as u64;
  343. let started_at = Utc::now();
  344. let started = Instant::now();
  345. let mut base_request = request_template(&provider_config, &model, 0.0, max_tokens);
  346. base_request.stream = stream.unwrap_or(provider_config.stream);
  347. let protocol = provider_config.protocol;
  348. let pb = ProgressBar::new(total);
  349. pb.set_style(
  350. ProgressStyle::default_bar()
  351. .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
  352. .expect("valid template"),
  353. );
  354. let mut results = stream::iter(cases)
  355. .map(|case| {
  356. let mut request = base_request.clone();
  357. request.prompt = case.prompt();
  358. async move {
  359. let result = run_model_request(protocol, request).await;
  360. (case, result)
  361. }
  362. })
  363. .buffer_unordered(nonzero_concurrency(concurrency));
  364. let mut metrics = Metrics::new();
  365. let mut wrong_cases = Vec::new();
  366. let mut correct_samples = Vec::new();
  367. while let Some((case, result)) = results.next().await {
  368. pb.inc(1);
  369. match result {
  370. Ok(response) => {
  371. metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64));
  372. let actual = judge::extract_choice(&response.text)
  373. .map(|choice| choice.to_string())
  374. .unwrap_or_else(|| "no_answer".to_string());
  375. let expected = case.answer.to_string();
  376. let correct = judge::judge_choice(&response.text, case.answer);
  377. metrics.record_judgement(correct);
  378. if correct {
  379. if correct_samples.len() < 5 {
  380. correct_samples.push(CorrectCaseReport {
  381. id: case.id,
  382. question: case.question,
  383. expected,
  384. raw_output: response.text,
  385. });
  386. }
  387. } else {
  388. wrong_cases.push(WrongCaseReport {
  389. id: case.id,
  390. question: case.question,
  391. expected,
  392. actual,
  393. raw_output: response.text,
  394. });
  395. }
  396. }
  397. Err(error) => metrics.record_failure(error_code(&error)),
  398. }
  399. }
  400. pb.finish_and_clear();
  401. let summary = metrics.summary();
  402. let report = benchmark_report(BenchmarkReportInput {
  403. benchmark: "gpqa-diamond",
  404. provider: provider_name,
  405. model,
  406. stream: base_request.stream,
  407. dataset,
  408. started_at,
  409. duration_ms: started.elapsed().as_millis(),
  410. concurrency,
  411. limit,
  412. max_tokens,
  413. summary,
  414. correct_samples,
  415. wrong_cases,
  416. });
  417. let report_path = write_benchmark_report(Path::new("."), &report)?;
  418. print_benchmark_report(&report, &report_path);
  419. Ok(())
  420. }
  421. struct RpmCommandOptions {
  422. mode: RpmMode,
  423. provider: Option<String>,
  424. model: Option<String>,
  425. rpm: Option<u32>,
  426. duration: String,
  427. burst: Option<u32>,
  428. probe_seconds: Option<u64>,
  429. window_offset_ms: u64,
  430. concurrency: Option<usize>,
  431. prompt: String,
  432. stream: Option<bool>,
  433. }
  434. async fn run_rpm(config_path: PathBuf, options: RpmCommandOptions) -> Result<()> {
  435. let config = AppConfig::load(&config_path)?;
  436. let mode_plan = build_rpm_mode_plan(
  437. options.mode,
  438. options.rpm,
  439. &options.duration,
  440. options.burst,
  441. options.probe_seconds,
  442. options.window_offset_ms,
  443. )?;
  444. let provider_name = provider_name(&config, options.provider.as_deref())?;
  445. let provider_config = config.resolved_provider(Some(&provider_name))?;
  446. let model = options
  447. .model
  448. .unwrap_or_else(|| provider_config.default_model.clone());
  449. let stream_enabled = options.stream.unwrap_or(provider_config.stream);
  450. let concurrency = options.concurrency.unwrap_or(mode_plan.default_concurrency);
  451. let request = ModelRequest {
  452. prompt: options.prompt.clone(),
  453. stream: stream_enabled,
  454. ..request_template(&provider_config, &model, 0.0, 1024)
  455. };
  456. let started_at = Utc::now();
  457. let started = Instant::now();
  458. let mut metrics = Metrics::new();
  459. let mut mode_summary = RpmModeSummaryBuilder::default();
  460. let results = run_scheduled_requests(
  461. provider_config.protocol,
  462. request,
  463. mode_plan.probes,
  464. concurrency,
  465. )
  466. .await;
  467. for result in results {
  468. let success = result.result.is_ok();
  469. mode_summary.record(result.phase, result.second, success);
  470. match result.result {
  471. Ok(response) => metrics.record_success(response.status, response.elapsed_ms as u64, response.first_token_ms.map(|ms| ms as u64)),
  472. Err(error) => metrics.record_failure(error_code(&error)),
  473. }
  474. }
  475. let summary = metrics.summary();
  476. let report = RpmReport {
  477. benchmark: "rpm".to_string(),
  478. provider: provider_name,
  479. model,
  480. params: RpmParamsReport {
  481. prompt: options.prompt,
  482. stream: stream_enabled,
  483. duration: options.duration,
  484. burst: options.burst,
  485. concurrency,
  486. probe_seconds: options.probe_seconds,
  487. window_offset_ms: options.window_offset_ms,
  488. },
  489. run: RpmRunReport {
  490. started_at,
  491. duration_ms: started.elapsed().as_millis(),
  492. target_rpm: mode_plan.target_rpm,
  493. actual_rpm: actual_rpm(summary.total, started.elapsed()),
  494. temperature: 0.0,
  495. max_tokens: 1024,
  496. },
  497. summary: RpmSummaryReport {
  498. actual_requests: summary.total,
  499. success: summary.success,
  500. failure: summary.failed,
  501. latency_ms: latency_report(&summary.latency_ms),
  502. ttft_ms: latency_report(&summary.ttft_ms),
  503. },
  504. mode: mode_plan.mode_name.to_string(),
  505. mode_detail: mode_summary.into_report(options.mode),
  506. errors: summary.errors,
  507. };
  508. let report_path = write_rpm_report(Path::new("."), &report)?;
  509. print_rpm_report(&report, &report_path);
  510. Ok(())
  511. }
  512. #[derive(Debug)]
  513. struct RpmModePlan {
  514. mode_name: &'static str,
  515. target_rpm: u32,
  516. probes: Vec<ScheduledProbe>,
  517. default_concurrency: usize,
  518. }
  519. #[derive(Debug)]
  520. struct ScheduledResult {
  521. phase: ProbePhase,
  522. second: Option<u64>,
  523. result: Result<crate::runner::ModelResponse>,
  524. }
  525. fn build_rpm_mode_plan(
  526. mode: RpmMode,
  527. rpm: Option<u32>,
  528. duration: &str,
  529. burst: Option<u32>,
  530. probe_seconds: Option<u64>,
  531. window_offset_ms: u64,
  532. ) -> Result<RpmModePlan> {
  533. let target_rpm = rpm.unwrap_or(0);
  534. let burst = burst.unwrap_or(target_rpm);
  535. let probes = match mode {
  536. RpmMode::Sustained => {
  537. let rpm = require_positive("rpm", rpm)?;
  538. let duration = parse_duration(duration)?;
  539. sustained_schedule(duration, rpm)
  540. .into_iter()
  541. .map(|offset| ScheduledProbe {
  542. offset,
  543. phase: ProbePhase::RefillProbe,
  544. second: Some(offset.as_secs()),
  545. })
  546. .collect()
  547. }
  548. RpmMode::Burst => {
  549. let burst = require_positive_value("burst", burst)?;
  550. burst_schedule(burst)
  551. .into_iter()
  552. .map(|offset| ScheduledProbe {
  553. offset,
  554. phase: ProbePhase::Burst,
  555. second: Some(0),
  556. })
  557. .collect()
  558. }
  559. RpmMode::TokenBucket => {
  560. let rpm = require_positive("rpm", rpm)?;
  561. let burst = require_positive_value("burst", burst)?;
  562. token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(30))
  563. }
  564. RpmMode::SlidingWindow => {
  565. let burst = require_positive_value("burst", burst)?;
  566. sliding_window_schedule(burst, probe_seconds.unwrap_or(90))
  567. }
  568. RpmMode::WindowBoundary => {
  569. let burst = require_positive_value("burst", burst)?;
  570. window_boundary_plan(Utc::now(), burst, window_offset_ms).probes
  571. }
  572. RpmMode::Diagnose => {
  573. let rpm = require_positive("rpm", rpm)?;
  574. let burst = require_positive_value("burst", burst)?;
  575. let mut probes = token_bucket_schedule(rpm, burst, probe_seconds.unwrap_or(90));
  576. probes.extend(sliding_window_schedule(burst, probe_seconds.unwrap_or(90)));
  577. probes.extend(window_boundary_plan(Utc::now(), burst, window_offset_ms).probes);
  578. probes
  579. }
  580. };
  581. let default_concurrency = probes.len().max(1);
  582. Ok(RpmModePlan {
  583. mode_name: mode_name(mode),
  584. target_rpm,
  585. default_concurrency,
  586. probes,
  587. })
  588. }
  589. fn mode_name(mode: RpmMode) -> &'static str {
  590. match mode {
  591. RpmMode::Sustained => "sustained",
  592. RpmMode::Burst => "burst",
  593. RpmMode::TokenBucket => "token-bucket",
  594. RpmMode::SlidingWindow => "sliding-window",
  595. RpmMode::WindowBoundary => "window-boundary",
  596. RpmMode::Diagnose => "diagnose",
  597. }
  598. }
  599. #[derive(Default)]
  600. struct RpmModeSummaryBuilder {
  601. phases: BTreeMap<&'static str, PhaseAccumulator>,
  602. refill_seconds: BTreeMap<u64, PhaseAccumulator>,
  603. sliding_seconds: BTreeMap<u64, PhaseAccumulator>,
  604. }
  605. impl RpmModeSummaryBuilder {
  606. fn record(&mut self, phase: ProbePhase, second: Option<u64>, success: bool) {
  607. match phase {
  608. ProbePhase::Burst => self.phase("burst").record(success),
  609. ProbePhase::RefillProbe => {
  610. self.phase("refill_probe").record(success);
  611. self.refill_seconds
  612. .entry(second.unwrap_or(0))
  613. .or_default()
  614. .record(success);
  615. }
  616. ProbePhase::SlidingProbe => {
  617. self.phase("sliding_probe").record(success);
  618. self.sliding_seconds
  619. .entry(second.unwrap_or(0))
  620. .or_default()
  621. .record(success);
  622. }
  623. ProbePhase::BeforeBoundary => self.phase("before_boundary").record(success),
  624. ProbePhase::AfterBoundary => self.phase("after_boundary").record(success),
  625. }
  626. }
  627. fn into_report(self, mode: RpmMode) -> Option<RpmModeDetailReport> {
  628. if mode == RpmMode::Sustained {
  629. return None;
  630. }
  631. let burst = self.phases.get("burst").map(PhaseAccumulator::to_report);
  632. let refill_probe = probe_seconds_report(self.refill_seconds);
  633. let sliding_probe = probe_seconds_report(self.sliding_seconds);
  634. let window_boundary = match (
  635. self.phases.get("before_boundary"),
  636. self.phases.get("after_boundary"),
  637. ) {
  638. (Some(before), Some(after)) => Some(WindowBoundaryReport {
  639. before: before.to_report(),
  640. after: after.to_report(),
  641. }),
  642. _ => None,
  643. };
  644. let inference = if mode == RpmMode::Diagnose {
  645. Some(infer_limiter(
  646. burst.as_ref(),
  647. &refill_probe,
  648. &sliding_probe,
  649. window_boundary.as_ref(),
  650. ))
  651. } else {
  652. None
  653. };
  654. Some(RpmModeDetailReport {
  655. burst,
  656. refill_probe,
  657. sliding_probe,
  658. window_boundary,
  659. inference,
  660. })
  661. }
  662. fn phase(&mut self, name: &'static str) -> &mut PhaseAccumulator {
  663. self.phases.entry(name).or_default()
  664. }
  665. }
  666. #[derive(Default)]
  667. struct PhaseAccumulator {
  668. sent: u64,
  669. success: u64,
  670. failure: u64,
  671. }
  672. impl PhaseAccumulator {
  673. fn record(&mut self, success: bool) {
  674. self.sent += 1;
  675. if success {
  676. self.success += 1;
  677. } else {
  678. self.failure += 1;
  679. }
  680. }
  681. fn to_report(&self) -> PhaseSummaryReport {
  682. PhaseSummaryReport {
  683. sent: self.sent,
  684. success: self.success,
  685. failure: self.failure,
  686. }
  687. }
  688. fn success_rate(&self) -> f64 {
  689. if self.sent == 0 {
  690. 0.0
  691. } else {
  692. self.success as f64 / self.sent as f64
  693. }
  694. }
  695. }
  696. fn probe_seconds_report(seconds: BTreeMap<u64, PhaseAccumulator>) -> Vec<ProbeSecondReport> {
  697. seconds
  698. .into_iter()
  699. .map(|(second, accumulator)| ProbeSecondReport {
  700. second,
  701. sent: accumulator.sent,
  702. success: accumulator.success,
  703. failure: accumulator.failure,
  704. })
  705. .collect()
  706. }
  707. fn infer_limiter(
  708. _burst: Option<&PhaseSummaryReport>,
  709. refill_probe: &[ProbeSecondReport],
  710. sliding_probe: &[ProbeSecondReport],
  711. window_boundary: Option<&WindowBoundaryReport>,
  712. ) -> LimiterInferenceReport {
  713. let mut signals = Vec::new();
  714. if let Some(boundary) = window_boundary {
  715. let before_rate = phase_success_rate(&boundary.before);
  716. let after_rate = phase_success_rate(&boundary.after);
  717. if after_rate > before_rate + 0.3 {
  718. signals.push("after-boundary success rate was much higher than before-boundary".into());
  719. return LimiterInferenceReport {
  720. likely_limiter: LimiterInferenceKind::FixedWindow,
  721. confidence: "medium".to_string(),
  722. signals,
  723. };
  724. }
  725. }
  726. let refill_sent: u64 = refill_probe.iter().map(|probe| probe.sent).sum();
  727. let refill_success: u64 = refill_probe.iter().map(|probe| probe.success).sum();
  728. if refill_sent > 0 && refill_success as f64 / refill_sent as f64 >= 0.5 {
  729. signals.push("refill probes recovered at a steady rate".into());
  730. return LimiterInferenceReport {
  731. likely_limiter: LimiterInferenceKind::TokenBucket,
  732. confidence: "medium".to_string(),
  733. signals,
  734. };
  735. }
  736. let early = sliding_probe
  737. .iter()
  738. .filter(|probe| probe.second <= 30)
  739. .fold(PhaseAccumulator::default(), |mut acc, probe| {
  740. acc.sent += probe.sent;
  741. acc.success += probe.success;
  742. acc.failure += probe.failure;
  743. acc
  744. });
  745. let late = sliding_probe
  746. .iter()
  747. .filter(|probe| probe.second >= 60)
  748. .fold(PhaseAccumulator::default(), |mut acc, probe| {
  749. acc.sent += probe.sent;
  750. acc.success += probe.success;
  751. acc.failure += probe.failure;
  752. acc
  753. });
  754. if late.sent > 0 && late.success_rate() > early.success_rate() + 0.3 {
  755. signals.push("probe recovery improved near the 60 second rolling window".into());
  756. return LimiterInferenceReport {
  757. likely_limiter: LimiterInferenceKind::SlidingWindow,
  758. confidence: "medium".to_string(),
  759. signals,
  760. };
  761. }
  762. signals.push("signals did not clearly match a limiter model".into());
  763. LimiterInferenceReport {
  764. likely_limiter: LimiterInferenceKind::Unknown,
  765. confidence: "low".to_string(),
  766. signals,
  767. }
  768. }
  769. fn phase_success_rate(phase: &PhaseSummaryReport) -> f64 {
  770. if phase.sent == 0 {
  771. 0.0
  772. } else {
  773. phase.success as f64 / phase.sent as f64
  774. }
  775. }
  776. fn actual_rpm(total_requests: u64, elapsed: Duration) -> Option<f64> {
  777. let elapsed_seconds = elapsed.as_secs_f64();
  778. if elapsed_seconds == 0.0 {
  779. None
  780. } else {
  781. Some(total_requests as f64 / elapsed_seconds * 60.0)
  782. }
  783. }
  784. fn require_positive(name: &str, value: Option<u32>) -> Result<u32> {
  785. let value = value.with_context(|| format!("{name} is required for this rpm mode"))?;
  786. require_positive_value(name, value)
  787. }
  788. fn require_positive_value(name: &str, value: u32) -> Result<u32> {
  789. if value == 0 {
  790. bail!("{name} must be greater than 0");
  791. }
  792. Ok(value)
  793. }
  794. async fn run_scheduled_requests(
  795. protocol: crate::config::ProtocolKind,
  796. request: ModelRequest,
  797. starts: Vec<ScheduledProbe>,
  798. max_in_flight: usize,
  799. ) -> Vec<ScheduledResult> {
  800. let total = starts.len() as u64;
  801. let pb = ProgressBar::new(total);
  802. pb.set_style(
  803. ProgressStyle::default_bar()
  804. .template("[{elapsed_precise}] {bar:40} {pos}/{len} ({eta})")
  805. .expect("valid template"),
  806. );
  807. let tokio_started = TokioInstant::now();
  808. let mut results = Vec::with_capacity(starts.len());
  809. let mut s = stream::iter(starts.into_iter().map(|start| {
  810. let request = request.clone();
  811. async move {
  812. sleep_until(tokio_started + start.offset).await;
  813. ScheduledResult {
  814. phase: start.phase,
  815. second: start.second,
  816. result: run_model_request(protocol, request).await,
  817. }
  818. }
  819. }))
  820. .buffer_unordered(nonzero_concurrency(max_in_flight));
  821. while let Some(result) = s.next().await {
  822. pb.inc(1);
  823. results.push(result);
  824. }
  825. pb.finish_and_clear();
  826. results
  827. }
  828. fn provider_name(config: &AppConfig, provider: Option<&str>) -> Result<String> {
  829. match provider {
  830. Some(provider) => Ok(provider.to_string()),
  831. None => config
  832. .default_provider
  833. .clone()
  834. .context("no provider specified and config has no default_provider"),
  835. }
  836. }
  837. fn request_template(
  838. provider: &crate::config::ProviderConfig,
  839. model: &str,
  840. temperature: f32,
  841. max_tokens: u32,
  842. ) -> ModelRequest {
  843. ModelRequest {
  844. base_url: provider.base_url.clone(),
  845. api_token: provider.api_token.clone(),
  846. model: model.to_string(),
  847. prompt: String::new(),
  848. temperature,
  849. max_tokens,
  850. stream: provider.stream,
  851. }
  852. }
  853. fn apply_limit<T>(cases: Vec<T>, limit: Option<usize>) -> Vec<T> {
  854. match limit {
  855. Some(limit) => cases.into_iter().take(limit).collect(),
  856. None => cases,
  857. }
  858. }
  859. fn dataset_report(config: Option<(&str, &str)>, local_path: &Path) -> DatasetReport {
  860. let (source, split) = config.unwrap_or(("local", "train"));
  861. DatasetReport {
  862. source: source.to_string(),
  863. split: split.to_string(),
  864. revision: None,
  865. local_path: local_path.display().to_string(),
  866. }
  867. }
  868. struct BenchmarkReportInput {
  869. benchmark: &'static str,
  870. provider: String,
  871. model: String,
  872. stream: bool,
  873. dataset: DatasetReport,
  874. started_at: chrono::DateTime<Utc>,
  875. duration_ms: u128,
  876. concurrency: usize,
  877. limit: Option<usize>,
  878. max_tokens: u32,
  879. summary: MetricsSummary,
  880. correct_samples: Vec<CorrectCaseReport>,
  881. wrong_cases: Vec<WrongCaseReport>,
  882. }
  883. fn benchmark_report(input: BenchmarkReportInput) -> BenchmarkReport {
  884. BenchmarkReport {
  885. benchmark: input.benchmark.to_string(),
  886. provider: input.provider,
  887. model: input.model,
  888. params: BenchmarkParamsReport {
  889. stream: input.stream,
  890. },
  891. dataset: input.dataset,
  892. run: RunReport {
  893. started_at: input.started_at,
  894. duration_ms: input.duration_ms,
  895. concurrency: input.concurrency,
  896. limit: input.limit,
  897. temperature: 0.0,
  898. max_tokens: input.max_tokens,
  899. },
  900. summary: BenchmarkSummaryReport {
  901. accuracy: input.summary.accuracy,
  902. success: input.summary.success,
  903. total: input.summary.total,
  904. correct: input.summary.correct,
  905. wrong: input.summary.wrong,
  906. failed: input.summary.failed,
  907. latency_ms: latency_report(&input.summary.latency_ms),
  908. ttft_ms: latency_report(&input.summary.ttft_ms),
  909. },
  910. errors: input.summary.errors,
  911. correct_samples: input.correct_samples,
  912. wrong_cases: input.wrong_cases,
  913. }
  914. }
  915. fn latency_report(summary: &LatencySummary) -> LatencyReport {
  916. LatencyReport {
  917. p50: summary.p50,
  918. p95: summary.p95,
  919. p99: summary.p99,
  920. }
  921. }
  922. fn print_benchmark_report(report: &BenchmarkReport, report_path: &Path) {
  923. println!("benchmark: {}", report.benchmark);
  924. println!(
  925. "accuracy: {}",
  926. report
  927. .summary
  928. .accuracy
  929. .map(|accuracy| format!("{:.2}%", accuracy * 100.0))
  930. .unwrap_or_else(|| "n/a".to_string())
  931. );
  932. println!(
  933. "success: {}/{} (failed: {})",
  934. report.summary.success, report.summary.total, report.summary.failed
  935. );
  936. println!(
  937. "latency_ms: p50={} p95={} p99={}",
  938. format_optional_latency(report.summary.latency_ms.p50),
  939. format_optional_latency(report.summary.latency_ms.p95),
  940. format_optional_latency(report.summary.latency_ms.p99)
  941. );
  942. if report.summary.ttft_ms.p50.is_some() {
  943. println!(
  944. "ttft_ms: p50={} p95={} p99={}",
  945. format_optional_latency(report.summary.ttft_ms.p50),
  946. format_optional_latency(report.summary.ttft_ms.p95),
  947. format_optional_latency(report.summary.ttft_ms.p99)
  948. );
  949. }
  950. println!("errors:");
  951. if report.errors.is_empty() {
  952. println!(" none");
  953. } else {
  954. for error in &report.errors {
  955. println!(" {}: {}", error.code, error.count);
  956. }
  957. }
  958. println!("wrong_cases:");
  959. if report.wrong_cases.is_empty() {
  960. println!(" none");
  961. } else {
  962. for case in &report.wrong_cases {
  963. println!(
  964. " {} expected={} actual={}",
  965. case.id, case.expected, case.actual
  966. );
  967. }
  968. }
  969. println!("report: {}", report_path.display());
  970. }
  971. fn print_rpm_report(report: &RpmReport, report_path: &Path) {
  972. println!("mode: {}", report.mode);
  973. println!("target_rpm: {}", report.run.target_rpm);
  974. println!(
  975. "actual_rpm: {}",
  976. report
  977. .run
  978. .actual_rpm
  979. .map(|rpm| format!("{rpm:.2}"))
  980. .unwrap_or_else(|| "n/a".to_string())
  981. );
  982. println!("actual_requests: {}", report.summary.actual_requests);
  983. println!(
  984. "success: {} failed: {}",
  985. report.summary.success, report.summary.failure
  986. );
  987. println!(
  988. "latency_ms: p50={} p95={} p99={}",
  989. format_optional_latency(report.summary.latency_ms.p50),
  990. format_optional_latency(report.summary.latency_ms.p95),
  991. format_optional_latency(report.summary.latency_ms.p99)
  992. );
  993. if report.summary.ttft_ms.p50.is_some() {
  994. println!(
  995. "ttft_ms: p50={} p95={} p99={}",
  996. format_optional_latency(report.summary.ttft_ms.p50),
  997. format_optional_latency(report.summary.ttft_ms.p95),
  998. format_optional_latency(report.summary.ttft_ms.p99)
  999. );
  1000. }
  1001. println!("errors:");
  1002. if report.errors.is_empty() {
  1003. println!(" none");
  1004. } else {
  1005. for error in &report.errors {
  1006. println!(" {}: {}", error.code, error.count);
  1007. }
  1008. }
  1009. println!("report: {}", report_path.display());
  1010. }
  1011. fn format_optional_latency(value: Option<u64>) -> String {
  1012. value
  1013. .map(|value| value.to_string())
  1014. .unwrap_or_else(|| "n/a".to_string())
  1015. }
  1016. fn error_code(error: &anyhow::Error) -> String {
  1017. let message = error.to_string();
  1018. let status_regex = Regex::new(r"status\s+(\d{3})").expect("valid status regex");
  1019. status_regex
  1020. .captures(&message)
  1021. .and_then(|captures| captures.get(1))
  1022. .map(|code| code.as_str().to_string())
  1023. .unwrap_or_else(|| "request_error".to_string())
  1024. }
  1025. fn nonzero_concurrency(concurrency: usize) -> usize {
  1026. concurrency.max(1)
  1027. }
  1028. fn parse_duration(value: &str) -> Result<Duration> {
  1029. let value = value.trim();
  1030. let Some(number) = value.strip_suffix('s') else {
  1031. if let Some(number) = value.strip_suffix('m') {
  1032. let minutes = number
  1033. .parse::<u64>()
  1034. .with_context(|| format!("invalid duration: {value}"))?;
  1035. return Ok(Duration::from_secs(minutes * 60));
  1036. }
  1037. bail!("invalid duration: expected values like 60s or 5m");
  1038. };
  1039. let seconds = number
  1040. .parse::<u64>()
  1041. .with_context(|| format!("invalid duration: {value}"))?;
  1042. Ok(Duration::from_secs(seconds))
  1043. }
  1044. fn dataset_data_dir(config_path: &Path) -> Result<PathBuf> {
  1045. if !config_path.exists() {
  1046. return Ok(PathBuf::from("data/benchmarks"));
  1047. }
  1048. let config = AppConfig::load(config_path)?;
  1049. Ok(PathBuf::from(config.benchmarks.data_dir))
  1050. }
  1051. #[cfg(test)]
  1052. mod tests {
  1053. use super::*;
  1054. #[test]
  1055. fn dataset_data_dir_defaults_when_config_is_missing() {
  1056. let temp_dir = tempfile::tempdir().expect("create temp dir");
  1057. let missing_config = temp_dir.path().join("missing-config.yaml");
  1058. let data_dir = dataset_data_dir(&missing_config).expect("default data dir");
  1059. assert_eq!(data_dir, PathBuf::from("data/benchmarks"));
  1060. }
  1061. #[test]
  1062. fn dataset_data_dir_propagates_invalid_existing_config() {
  1063. let temp_dir = tempfile::tempdir().expect("create temp dir");
  1064. let config_path = temp_dir.path().join("config.yaml");
  1065. std::fs::write(&config_path, "providers: [").expect("write invalid config");
  1066. let error = dataset_data_dir(&config_path).expect_err("invalid config should fail");
  1067. assert!(error.to_string().contains("failed to parse config"));
  1068. }
  1069. #[test]
  1070. fn parses_duration_seconds_and_minutes() {
  1071. assert_eq!(parse_duration("60s").expect("seconds").as_secs(), 60);
  1072. assert_eq!(parse_duration("5m").expect("minutes").as_secs(), 300);
  1073. }
  1074. #[test]
  1075. fn computes_rpm_start_schedule_from_run_start() {
  1076. let schedule = sustained_schedule(Duration::from_secs(60), 120);
  1077. assert_eq!(schedule.len(), 120);
  1078. assert_eq!(schedule[0], Duration::ZERO);
  1079. assert_eq!(schedule[1], Duration::from_millis(500));
  1080. assert_eq!(schedule[119], Duration::from_millis(59_500));
  1081. }
  1082. #[test]
  1083. fn rejects_invalid_duration() {
  1084. let error = parse_duration("one hour").expect_err("invalid duration");
  1085. assert!(error.to_string().contains("duration"));
  1086. }
  1087. #[test]
  1088. fn rpm_command_defaults_to_sustained_mode() {
  1089. let cli = Cli::try_parse_from([
  1090. "lq_token_test",
  1091. "rpm",
  1092. "--provider",
  1093. "anthropic",
  1094. "--rpm",
  1095. "120",
  1096. "--prompt",
  1097. "hello",
  1098. ])
  1099. .expect("parse rpm");
  1100. let Command::Rpm { mode, rpm, .. } = cli.command else {
  1101. panic!("expected rpm command");
  1102. };
  1103. assert_eq!(mode, RpmMode::Sustained);
  1104. assert_eq!(rpm, Some(120));
  1105. }
  1106. #[test]
  1107. fn rpm_command_parses_token_bucket_mode() {
  1108. let cli = Cli::try_parse_from([
  1109. "lq_token_test",
  1110. "rpm",
  1111. "--mode",
  1112. "token-bucket",
  1113. "--provider",
  1114. "anthropic",
  1115. "--rpm",
  1116. "120",
  1117. "--burst",
  1118. "120",
  1119. "--probe-seconds",
  1120. "30",
  1121. "--prompt",
  1122. "hello",
  1123. ])
  1124. .expect("parse token bucket rpm");
  1125. let Command::Rpm {
  1126. mode,
  1127. burst,
  1128. probe_seconds,
  1129. ..
  1130. } = cli.command
  1131. else {
  1132. panic!("expected rpm command");
  1133. };
  1134. assert_eq!(mode, RpmMode::TokenBucket);
  1135. assert_eq!(burst, Some(120));
  1136. assert_eq!(probe_seconds, Some(30));
  1137. }
  1138. #[test]
  1139. fn rpm_command_parses_window_boundary_offset() {
  1140. let cli = Cli::try_parse_from([
  1141. "lq_token_test",
  1142. "rpm",
  1143. "--mode",
  1144. "window-boundary",
  1145. "--provider",
  1146. "anthropic",
  1147. "--burst",
  1148. "10",
  1149. "--window-offset-ms",
  1150. "250",
  1151. "--prompt",
  1152. "hello",
  1153. ])
  1154. .expect("parse window boundary rpm");
  1155. let Command::Rpm {
  1156. mode,
  1157. window_offset_ms,
  1158. ..
  1159. } = cli.command
  1160. else {
  1161. panic!("expected rpm command");
  1162. };
  1163. assert_eq!(mode, RpmMode::WindowBoundary);
  1164. assert_eq!(window_offset_ms, 250);
  1165. }
  1166. #[test]
  1167. fn rpm_mode_validation_rejects_zero_values() {
  1168. let zero_rpm = build_rpm_mode_plan(RpmMode::Sustained, Some(0), "60s", None, None, 500)
  1169. .expect_err("zero rpm should fail");
  1170. let zero_burst = build_rpm_mode_plan(RpmMode::Burst, None, "60s", Some(0), None, 500)
  1171. .expect_err("zero burst should fail");
  1172. assert!(zero_rpm.to_string().contains("rpm"));
  1173. assert!(zero_burst.to_string().contains("burst"));
  1174. }
  1175. }