eval: Add jitter to retry attempts (#32542)

Ben Brandt created 6 months ago

Adds some jitter to avoid the issue that all requests will retry at
roughly the same time in eval where we have a lot of concurrent
requests.

Release Notes:

- N/A

Change summary

crates/assistant_tools/src/edit_agent/evals.rs | 11 ++++++++---
1 file changed, 8 insertions(+), 3 deletions(-)

Detailed changes

crates/assistant_tools/src/edit_agent/evals.rs

@@ -1634,15 +1634,20 @@ impl EditAgentTest {
 }
 
 async fn retry_on_rate_limit<R>(mut request: impl AsyncFnMut() -> Result<R>) -> Result<R> {
+    let mut attempt = 0;
     loop {
+        attempt += 1;
         match request().await {
             Ok(result) => return Ok(result),
             Err(err) => match err.downcast::<LanguageModelCompletionError>() {
                 Ok(err) => match err {
                     LanguageModelCompletionError::RateLimit(duration) => {
-                        // Wait until after we are allowed to try again
-                        eprintln!("Rate limit exceeded. Waiting for {duration:?}...",);
-                        Timer::after(duration).await;
+                        // Wait for the duration supplied, with some jitter to avoid all requests being made at the same time.
+                        let jitter = duration.mul_f64(rand::thread_rng().gen_range(0.0..0.5));
+                        eprintln!(
+                            "Attempt #{attempt}: Rate limit exceeded. Retry after {duration:?} + jitter of {jitter:?}"
+                        );
+                        Timer::after(duration + jitter).await;
                         continue;
                     }
                     _ => return Err(err.into()),