[ENH]: Change retry rate limits in s3 client (#6917)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
- Increases the token bucket capacity that s3 client uses for retries by
10x
    - Adds a token bucket fill rate of 1000 tokens / sec
- New functionality
    - ...

## Test plan

- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Migration plan

_None_

## Observability plan

_None_

## Documentation Changes

_None_

Co-authored-by: Sanket Kedia <sanketkedia@Sankets-MacBook-Pro.local>

Author: sanketkedia

rust/storage/src/config.rs

@@ -104,6 +104,10 @@ pub struct S3StorageConfig {
     pub upload_part_size_bytes: usize,
     #[serde(default = "S3StorageConfig::default_download_part_size_bytes")]
     pub download_part_size_bytes: usize,
+    #[serde(default = "S3StorageConfig::default_retry_token_bucket_capacity")]
+    pub retry_token_bucket_capacity: usize,
+    #[serde(default = "S3StorageConfig::default_retry_token_refill_rate")]
+    pub retry_token_refill_rate: f32,
 }
 
 impl S3StorageConfig {
@@ -147,6 +151,14 @@ impl S3StorageConfig {
     fn default_download_part_size_bytes() -> usize {
         8 * 1024 * 1024
     }
+
+    fn default_retry_token_bucket_capacity() -> usize {
+        5000
+    }
+
+    fn default_retry_token_refill_rate() -> f32 {
+        1000.0
+    }
 }
 
 impl Default for S3StorageConfig {
@@ -163,6 +175,8 @@ impl Default for S3StorageConfig {
             stall_upload_enabled: S3StorageConfig::default_stall_upload_enabled(),
             upload_part_size_bytes: S3StorageConfig::default_upload_part_size_bytes(),
             download_part_size_bytes: S3StorageConfig::default_download_part_size_bytes(),
+            retry_token_bucket_capacity: S3StorageConfig::default_retry_token_bucket_capacity(),
+            retry_token_refill_rate: S3StorageConfig::default_retry_token_refill_rate(),
         }
     }
 }

rust/storage/src/s3.rs

@@ -17,6 +17,7 @@ use async_trait::async_trait;
 use aws_config::retry::RetryConfig;
 use aws_config::timeout::TimeoutConfigBuilder;
 use aws_sdk_s3;
+use aws_sdk_s3::config::retry::{RetryPartition, TokenBucket};
 use aws_sdk_s3::config::StalledStreamProtectionConfig;
 use aws_sdk_s3::error::ProvideErrorMetadata;
 use aws_sdk_s3::error::SdkError;
@@ -1194,6 +1195,15 @@ impl Configurable<StorageConfig> for S3Storage {
                 let retry_config =
                     RetryConfig::standard().with_max_attempts(s3_config.request_retry_count);
 
+                let retry_partition = RetryPartition::custom("chroma-s3")
+                    .token_bucket(
+                        TokenBucket::builder()
+                            .capacity(s3_config.retry_token_bucket_capacity)
+                            .refill_rate(s3_config.retry_token_refill_rate)
+                            .build(),
+                    )
+                    .build();
+
                 let client = match &s3_config.credentials {
                     super::config::S3CredentialsConfig::Minio
                     | super::config::S3CredentialsConfig::Localhost => {
@@ -1228,14 +1238,14 @@ impl Configurable<StorageConfig> for S3Storage {
                         aws_sdk_s3::Client::from_conf(config)
                     }
                     super::config::S3CredentialsConfig::AWS => {
-                        let config = aws_config::load_from_env().await;
-                        let config = config
-                            .to_builder()
+                        let sdk_config = aws_config::load_from_env().await;
+                        let config = aws_sdk_s3::config::Builder::from(&sdk_config)
                             .timeout_config(timeout_config)
                             .stalled_stream_protection(stalled_config)
                             .retry_config(retry_config)
+                            .retry_partition(retry_partition)
                             .build();
-                        aws_sdk_s3::Client::new(&config)
+                        aws_sdk_s3::Client::from_conf(config)
                     }
                     super::config::S3CredentialsConfig::Explicit {
                         access_key_id,
@@ -1258,7 +1268,8 @@ impl Configurable<StorageConfig> for S3Storage {
                             .region(aws_sdk_s3::config::Region::new(region.clone()))
                             .timeout_config(timeout_config)
                             .stalled_stream_protection(stalled_config)
-                            .retry_config(retry_config);
+                            .retry_config(retry_config)
+                            .retry_partition(retry_partition);
 
                         if let Some(url) = custom_endpoint {
                             config_builder =