api模型检测
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

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