1use std::collections::{BTreeMap, HashMap, VecDeque};
22use std::future::Future;
23use std::time::{Duration, Instant};
24
25use tokio::task::JoinSet;
26
27use crate::protocol::{RpcError, RunResult, TranscriptSummary};
28use crate::{Params, Trial};
29
30const GROW_THRESHOLD: usize = 3;
32const MAX_BACKOFF_STEPS: u32 = 6;
34
35#[derive(Clone, Debug)]
38pub struct CaseSpec {
39 pub eval: String,
40 pub sample: String,
41 pub target: String,
42 pub provider: String,
45 pub params: Params,
46 pub trial: Trial,
49 pub timeout: Option<Duration>,
55}
56
57impl CaseSpec {
58 pub fn key(&self) -> String {
60 format!("{}{}", self.logical_key(), self.trial.key_suffix())
61 }
62
63 pub fn logical_key(&self) -> String {
65 crate::case_key(&self.eval, &self.sample, &self.target, &self.params)
66 }
67}
68
69#[derive(Clone, Debug)]
71pub struct Concurrency {
72 pub global: usize,
74 pub per_provider: BTreeMap<String, usize>,
76 pub default_per_provider: usize,
78 pub adaptive: bool,
80 pub max_retries: u32,
82 pub base_backoff: Duration,
84}
85
86impl Concurrency {
87 pub fn new(global: usize) -> Self {
91 let global = global.max(1);
92 Self {
93 global,
94 per_provider: BTreeMap::new(),
95 default_per_provider: global,
96 adaptive: true,
97 max_retries: 4,
98 base_backoff: Duration::from_millis(500),
99 }
100 }
101
102 pub fn provider(mut self, provider: impl Into<String>, limit: usize) -> Self {
104 self.per_provider.insert(provider.into(), limit.max(1));
105 self
106 }
107}
108
109impl Default for Concurrency {
110 fn default() -> Self {
111 Self::new(8)
112 }
113}
114
115#[derive(Debug)]
117struct ProviderState {
118 in_flight: usize,
119 limit: usize,
121 ceiling: usize,
123 ok_streak: usize,
124 backoff_steps: u32,
126 backoff_until: Option<Instant>,
128}
129
130struct Limiter {
133 providers: HashMap<String, ProviderState>,
134 per_provider: BTreeMap<String, usize>,
135 default_per_provider: usize,
136 global_in_flight: usize,
137 global_max: usize,
138 adaptive: bool,
139 base_backoff: Duration,
140}
141
142impl Limiter {
143 fn new(cfg: &Concurrency) -> Self {
144 Self {
145 providers: HashMap::new(),
146 per_provider: cfg.per_provider.clone(),
147 default_per_provider: cfg.default_per_provider.max(1),
148 global_in_flight: 0,
149 global_max: cfg.global.max(1),
150 adaptive: cfg.adaptive,
151 base_backoff: cfg.base_backoff,
152 }
153 }
154
155 fn ceiling_for(&self, provider: &str) -> usize {
156 self.per_provider
157 .get(provider)
158 .copied()
159 .unwrap_or(self.default_per_provider)
160 .clamp(1, self.global_max)
161 }
162
163 fn state(&mut self, provider: &str) -> &mut ProviderState {
164 let ceiling = self.ceiling_for(provider);
165 self.providers
166 .entry(provider.to_string())
167 .or_insert_with(|| ProviderState {
168 in_flight: 0,
169 limit: ceiling,
170 ceiling,
171 ok_streak: 0,
172 backoff_steps: 0,
173 backoff_until: None,
174 })
175 }
176
177 fn can_start(&mut self, provider: &str, now: Instant) -> bool {
180 if self.global_in_flight >= self.global_max {
181 return false;
182 }
183 let st = self.state(provider);
184 st.in_flight < st.limit && st.backoff_until.is_none_or(|t| now >= t)
185 }
186
187 fn start(&mut self, provider: &str) {
188 self.global_in_flight += 1;
189 self.state(provider).in_flight += 1;
190 }
191
192 fn finish(&mut self, provider: &str, rate_limited: bool, now: Instant) {
194 let adaptive = self.adaptive;
195 let base = self.base_backoff;
196 self.global_in_flight = self.global_in_flight.saturating_sub(1);
197 let st = self.state(provider);
198 st.in_flight = st.in_flight.saturating_sub(1);
199 if !adaptive {
200 return;
201 }
202 if rate_limited {
203 st.limit = (st.limit / 2).max(1);
205 st.backoff_steps = (st.backoff_steps + 1).min(MAX_BACKOFF_STEPS);
206 let mult = 1u32 << (st.backoff_steps - 1);
207 st.backoff_until = Some(now + base * mult);
208 st.ok_streak = 0;
209 } else {
210 st.ok_streak += 1;
212 if st.ok_streak >= GROW_THRESHOLD {
213 st.ok_streak = 0;
214 st.backoff_steps = st.backoff_steps.saturating_sub(1);
215 if st.limit < st.ceiling {
216 st.limit += 1;
217 }
218 }
219 }
220 }
221
222 fn earliest_ready(&self, pending: &VecDeque<(CaseSpec, u32)>) -> Option<Instant> {
225 pending
226 .iter()
227 .filter_map(|(c, _)| self.providers.get(&c.provider))
228 .filter_map(|s| s.backoff_until)
229 .min()
230 }
231}
232
233fn outcome_rate_limited(res: &Result<RunResult, RpcError>) -> bool {
236 match res {
237 Err(e) => crate::is_rate_limited(&e.message),
238 Ok(r) => r
239 .transcript
240 .error
241 .as_deref()
242 .is_some_and(crate::is_rate_limited),
243 }
244}
245
246fn outcome_retryable(res: &Result<RunResult, RpcError>) -> bool {
254 match res {
255 Err(e) => e.retryable || crate::is_rate_limited(&e.message),
256 Ok(r) => {
257 r.transcript.error_kind == crate::ErrorKind::Infra
258 || r.transcript
259 .error
260 .as_deref()
261 .is_some_and(crate::is_rate_limited)
262 }
263 }
264}
265
266fn failed_result(case: &CaseSpec, error: RpcError) -> RunResult {
270 let infra = error.retryable || crate::is_rate_limited(&error.message);
271 RunResult {
272 eval: case.eval.clone(),
273 sample: case.sample.clone(),
274 target: case.target.clone(),
275 params: case.params.clone(),
276 trial: case.trial.index,
277 trials: case.trial.count,
278 seed: case.trial.seed,
279 input: Vec::new(),
280 expected: None,
281 passed: false,
282 aggregate: 0.0,
283 scores: Vec::new(),
284 transcript: TranscriptSummary {
285 error: Some(error.message),
286 error_kind: if infra {
287 crate::ErrorKind::Infra
288 } else {
289 crate::ErrorKind::Subject
290 },
291 ..Default::default()
292 },
293 skipped: false,
294 }
295}
296
297pub async fn run_cases<F, Fut>(
305 cases: Vec<CaseSpec>,
306 cfg: &Concurrency,
307 run: F,
308 mut on_done: impl FnMut(&CaseSpec, RunResult),
309) where
310 F: Fn(CaseSpec) -> Fut,
311 Fut: Future<Output = Result<RunResult, RpcError>> + Send + 'static,
312{
313 let mut limiter = Limiter::new(cfg);
314 let mut pending: VecDeque<(CaseSpec, u32)> = cases.into_iter().map(|c| (c, 0)).collect();
315 let mut tasks: JoinSet<Result<RunResult, RpcError>> = JoinSet::new();
316 let mut inflight: HashMap<tokio::task::Id, (CaseSpec, u32)> = HashMap::new();
319
320 loop {
321 loop {
323 let now = Instant::now();
324 let idx = pending
325 .iter()
326 .position(|(c, _)| limiter.can_start(&c.provider, now));
327 let Some(idx) = idx else { break };
328 let (case, attempts) = pending.remove(idx).expect("index in bounds");
329 limiter.start(&case.provider);
330 let task_case = case.clone();
331 let timeout = case.timeout;
332 let fut = run(case);
333 let task = async move {
339 match timeout {
340 Some(dur) => match tokio::time::timeout(dur, fut).await {
341 Ok(res) => res,
342 Err(_) => Err(RpcError::new(format!(
343 "timed out after {}s (target timeout)",
344 dur.as_secs()
345 ))),
346 },
347 None => fut.await,
348 }
349 };
350 let id = tasks.spawn(task).id();
351 inflight.insert(id, (task_case, attempts));
352 }
353
354 if tasks.is_empty() {
355 if pending.is_empty() {
356 break;
357 }
358 match limiter.earliest_ready(&pending) {
360 Some(t) => {
361 tokio::time::sleep_until(t.into()).await;
362 continue;
363 }
364 None => break,
366 }
367 }
368
369 let Some(joined) = tasks.join_next_with_id().await else {
370 continue;
371 };
372 let (case, attempts, res) = match joined {
375 Ok((id, res)) => {
376 let (case, attempts) = inflight.remove(&id).expect("task id tracked");
377 (case, attempts, res)
378 }
379 Err(join_err) => {
380 let (case, attempts) = inflight.remove(&join_err.id()).expect("task id tracked");
381 (
382 case,
383 attempts,
384 Err(RpcError::new(format!("task panicked: {join_err}"))),
385 )
386 }
387 };
388
389 let rate_limited = outcome_rate_limited(&res);
390 limiter.finish(&case.provider, rate_limited, Instant::now());
391
392 if attempts < cfg.max_retries && outcome_retryable(&res) {
397 pending.push_back((case, attempts + 1));
398 continue;
399 }
400
401 let result = match res {
402 Ok(result) => result,
403 Err(error) => failed_result(&case, error),
404 };
405 on_done(&case, result);
406 }
407}
408
409#[cfg(test)]
410mod tests {
411 use super::*;
412 use std::sync::Arc;
413 use std::sync::atomic::{AtomicUsize, Ordering};
414 use tokio::sync::Mutex;
415
416 fn case(provider: &str, id: &str) -> CaseSpec {
417 CaseSpec {
418 eval: "e".into(),
419 sample: id.into(),
420 target: format!("{provider}/m"),
421 provider: provider.into(),
422 params: Params::new(),
423 trial: Trial::single(),
424 timeout: None,
425 }
426 }
427
428 fn ok_result(case: &CaseSpec) -> RunResult {
429 RunResult {
430 eval: case.eval.clone(),
431 sample: case.sample.clone(),
432 target: case.target.clone(),
433 params: case.params.clone(),
434 trial: case.trial.index,
435 trials: case.trial.count,
436 seed: case.trial.seed,
437 input: Vec::new(),
438 expected: None,
439 passed: true,
440 aggregate: 1.0,
441 scores: Vec::new(),
442 transcript: TranscriptSummary::default(),
443 skipped: false,
444 }
445 }
446
447 #[test]
448 fn limiter_buckets_per_provider() {
449 let cfg = Concurrency::new(10).provider("anthropic", 2);
450 let mut lim = Limiter::new(&cfg);
451 let now = Instant::now();
452 assert!(lim.can_start("anthropic", now));
453 lim.start("anthropic");
454 lim.start("anthropic");
455 assert!(!lim.can_start("anthropic", now));
457 assert!(lim.can_start("openai", now));
459 }
460
461 #[test]
462 fn rate_limit_halves_and_backs_off() {
463 let cfg = Concurrency::new(8); let mut lim = Limiter::new(&cfg);
465 let now = Instant::now();
466 lim.start("anthropic");
467 lim.finish("anthropic", true, now);
468 let st = &lim.providers["anthropic"];
469 assert_eq!(st.limit, 4); assert!(st.backoff_until.is_some());
471 assert!(!lim.can_start("anthropic", now));
473 }
474
475 #[test]
476 fn sustained_success_grows_back() {
477 let cfg = Concurrency::new(8);
478 let mut lim = Limiter::new(&cfg);
479 let now = Instant::now();
480 lim.start("anthropic");
482 lim.finish("anthropic", true, now); assert_eq!(lim.providers["anthropic"].limit, 4);
484 for _ in 0..GROW_THRESHOLD {
486 lim.start("anthropic");
487 lim.finish("anthropic", false, now);
488 }
489 assert_eq!(lim.providers["anthropic"].limit, 5);
490 }
491
492 #[tokio::test]
493 async fn runs_every_case_once() {
494 let cases: Vec<CaseSpec> = (0..20).map(|i| case("sim", &i.to_string())).collect();
495 let cfg = Concurrency::new(4);
496 let seen = Arc::new(AtomicUsize::new(0));
497 let seen2 = seen.clone();
498 let mut done = Vec::new();
499 run_cases(
500 cases,
501 &cfg,
502 move |c| {
503 let seen = seen2.clone();
504 async move {
505 seen.fetch_add(1, Ordering::SeqCst);
506 Ok(ok_result(&c))
507 }
508 },
509 |_, r| done.push(r),
510 )
511 .await;
512 assert_eq!(seen.load(Ordering::SeqCst), 20);
513 assert_eq!(done.len(), 20);
514 assert!(done.iter().all(|r| r.passed));
515 }
516
517 #[tokio::test]
518 async fn respects_global_concurrency_cap() {
519 let cases: Vec<CaseSpec> = (0..30).map(|i| case("sim", &i.to_string())).collect();
521 let cfg = Concurrency::new(3);
522 let active = Arc::new(AtomicUsize::new(0));
523 let peak = Arc::new(AtomicUsize::new(0));
524 let (a, p) = (active.clone(), peak.clone());
525 let mut done = 0usize;
526 run_cases(
527 cases,
528 &cfg,
529 move |c| {
530 let (a, p) = (a.clone(), p.clone());
531 async move {
532 let n = a.fetch_add(1, Ordering::SeqCst) + 1;
533 p.fetch_max(n, Ordering::SeqCst);
534 tokio::time::sleep(Duration::from_millis(5)).await;
535 a.fetch_sub(1, Ordering::SeqCst);
536 Ok(ok_result(&c))
537 }
538 },
539 |_, _| done += 1,
540 )
541 .await;
542 assert_eq!(done, 30);
543 assert!(peak.load(Ordering::SeqCst) <= 3, "peak exceeded global cap");
544 }
545
546 #[tokio::test]
547 async fn retries_rate_limited_case_then_succeeds() {
548 let cfg = Concurrency {
550 base_backoff: Duration::from_millis(1),
551 ..Concurrency::new(2)
552 };
553 let attempts = Arc::new(Mutex::new(0usize));
554 let a = attempts.clone();
555 let mut results = Vec::new();
556 run_cases(
557 vec![case("anthropic", "x")],
558 &cfg,
559 move |c| {
560 let a = a.clone();
561 async move {
562 let mut n = a.lock().await;
563 *n += 1;
564 if *n == 1 {
565 Err(RpcError::new("HTTP 429 rate limit"))
566 } else {
567 Ok(ok_result(&c))
568 }
569 }
570 },
571 |_, r| results.push(r),
572 )
573 .await;
574 assert_eq!(*attempts.lock().await, 2);
575 assert_eq!(results.len(), 1);
576 assert!(results[0].passed);
577 }
578
579 #[tokio::test]
582 async fn retries_infra_errored_case_then_succeeds() {
583 let cfg = Concurrency::new(2);
584 let attempts = Arc::new(Mutex::new(0usize));
585 let a = attempts.clone();
586 let mut results = Vec::new();
587 run_cases(
588 vec![case("sim", "x")],
589 &cfg,
590 move |c| {
591 let a = a.clone();
592 async move {
593 let mut n = a.lock().await;
594 *n += 1;
595 if *n == 1 {
596 let mut r = ok_result(&c);
597 r.passed = false;
598 r.transcript.error = Some("provider 503 unavailable".into());
599 r.transcript.error_kind = crate::ErrorKind::Infra;
600 Ok(r)
601 } else {
602 Ok(ok_result(&c))
603 }
604 }
605 },
606 |_, r| results.push(r),
607 )
608 .await;
609 assert_eq!(*attempts.lock().await, 2); assert_eq!(results.len(), 1);
611 assert!(results[0].passed);
612 }
613
614 #[tokio::test]
619 async fn retries_retryable_rpc_error_then_succeeds() {
620 let cfg = Concurrency::new(2);
621 let attempts = Arc::new(Mutex::new(0usize));
622 let a = attempts.clone();
623 let mut results = Vec::new();
624 run_cases(
625 vec![case("sim", "x")],
626 &cfg,
627 move |c| {
628 let a = a.clone();
629 async move {
630 let mut n = a.lock().await;
631 *n += 1;
632 if *n == 1 {
633 Err(RpcError::new("provider outage").retryable())
634 } else {
635 Ok(ok_result(&c))
636 }
637 }
638 },
639 |_, r| results.push(r),
640 )
641 .await;
642 assert_eq!(*attempts.lock().await, 2); assert_eq!(results.len(), 1);
644 assert!(results[0].passed);
645 }
646
647 #[tokio::test]
650 async fn non_retryable_rpc_error_is_not_requeued() {
651 let cfg = Concurrency::new(2);
652 let count = Arc::new(AtomicUsize::new(0));
653 let c2 = count.clone();
654 let mut results = Vec::new();
655 run_cases(
656 vec![case("sim", "x")],
657 &cfg,
658 move |_| {
659 let c2 = c2.clone();
660 async move {
661 c2.fetch_add(1, Ordering::SeqCst);
662 Err(RpcError::new("bad run params"))
663 }
664 },
665 |_, r| results.push(r),
666 )
667 .await;
668 assert_eq!(count.load(Ordering::SeqCst), 1); assert_eq!(results.len(), 1);
670 assert!(!results[0].passed);
671 assert_eq!(results[0].transcript.error_kind, crate::ErrorKind::Subject);
672 }
673
674 #[tokio::test]
675 async fn panicking_case_is_recorded_and_frees_its_slot() {
676 let cfg = Concurrency::new(1);
679 let cases = vec![case("sim", "boom"), case("sim", "ok")];
680 let mut results = Vec::new();
681 run_cases(
682 cases,
683 &cfg,
684 move |c| async move {
685 if c.sample == "boom" {
686 panic!("subject blew up");
687 }
688 Ok(ok_result(&c))
689 },
690 |c, r| results.push((c.sample.clone(), r)),
691 )
692 .await;
693 assert_eq!(results.len(), 2);
694 let boom = results.iter().find(|(s, _)| s == "boom").unwrap();
695 assert!(!boom.1.passed);
696 assert!(
697 boom.1
698 .transcript
699 .error
700 .as_deref()
701 .unwrap()
702 .contains("panic")
703 );
704 let ok = results.iter().find(|(s, _)| s == "ok").unwrap();
705 assert!(ok.1.passed);
706 }
707
708 #[tokio::test]
709 async fn gives_up_after_max_retries() {
710 let cfg = Concurrency {
711 base_backoff: Duration::from_millis(1),
712 max_retries: 2,
713 ..Concurrency::new(1)
714 };
715 let count = Arc::new(AtomicUsize::new(0));
716 let c2 = count.clone();
717 let mut results = Vec::new();
718 run_cases(
719 vec![case("anthropic", "x")],
720 &cfg,
721 move |_| {
722 let c2 = c2.clone();
723 async move {
724 c2.fetch_add(1, Ordering::SeqCst);
725 Err(RpcError::new("429 too many requests"))
726 }
727 },
728 |_, r| results.push(r),
729 )
730 .await;
731 assert_eq!(count.load(Ordering::SeqCst), 3);
733 assert_eq!(results.len(), 1);
734 assert!(!results[0].passed);
735 assert!(results[0].transcript.error.is_some());
736 }
737
738 #[tokio::test]
742 async fn times_out_slow_case_without_retrying() {
743 let cfg = Concurrency {
744 base_backoff: Duration::from_millis(1),
745 max_retries: 4,
746 ..Concurrency::new(2)
747 };
748 let count = Arc::new(AtomicUsize::new(0));
749 let c2 = count.clone();
750 let slow = CaseSpec {
751 timeout: Some(Duration::from_millis(10)),
752 ..case("sim", "slow")
753 };
754 let mut results = Vec::new();
755 run_cases(
756 vec![slow],
757 &cfg,
758 move |c| {
759 let c2 = c2.clone();
760 async move {
761 c2.fetch_add(1, Ordering::SeqCst);
762 tokio::time::sleep(Duration::from_secs(60)).await;
764 Ok(ok_result(&c))
765 }
766 },
767 |_, r| results.push(r),
768 )
769 .await;
770 assert_eq!(count.load(Ordering::SeqCst), 1, "no retry after timeout");
771 assert_eq!(results.len(), 1);
772 assert!(!results[0].passed);
773 let err = results[0].transcript.error.as_deref().unwrap();
774 assert!(err.contains("timed out"), "{err}");
775 assert_eq!(results[0].transcript.error_kind, crate::ErrorKind::Subject);
777 }
778
779 #[tokio::test]
781 async fn fast_case_under_timeout_passes() {
782 let cfg = Concurrency::new(2);
783 let fast = CaseSpec {
784 timeout: Some(Duration::from_secs(30)),
785 ..case("sim", "fast")
786 };
787 let mut results = Vec::new();
788 run_cases(
789 vec![fast],
790 &cfg,
791 move |c| async move { Ok(ok_result(&c)) },
792 |_, r| results.push(r),
793 )
794 .await;
795 assert_eq!(results.len(), 1);
796 assert!(results[0].passed);
797 }
798}