[ENH] multipart S3 file uploads (#2590)

chroma-core · Aug 1, 2024 · 9472830 · 9472830
1 parent c0a1104
commit 9472830
Show file tree

Hide file tree

Showing 6 changed files with 236 additions and 58 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/rust/worker/Cargo.toml b/rust/worker/Cargo.toml
@@ -66,6 +66,7 @@ rayon = "1.8.0"
 criterion = "0.3"
 random-port = "0.1.1"
 serial_test = "3.1.1"
+rand_xorshift = "0.3.0"
 
 [build-dependencies]
 tonic-build = "0.10"

diff --git a/rust/worker/chroma_config.yaml b/rust/worker/chroma_config.yaml
@@ -28,6 +28,7 @@ query_service:
             credentials: "Minio"
             connect_timeout_ms: 5000
             request_timeout_ms: 30000 # 1 minute
+            upload_part_size_bytes: 8388608 # 8MB
     log:
         Grpc:
             host: "logservice.chroma"
@@ -80,6 +81,7 @@ compaction_service:
             credentials: "Minio"
             connect_timeout_ms: 5000
             request_timeout_ms: 60000 # 1 minute
+            upload_part_size_bytes: 8388608 # 8MB
     log:
         Grpc:
             host: "logservice.chroma"

diff --git a/rust/worker/src/config.rs b/rust/worker/src/config.rs
@@ -184,6 +184,7 @@ mod tests {
                             credentials: Minio
                             connect_timeout_ms: 5000
                             request_timeout_ms: 1000
+                            upload_part_size_bytes: 8388608
                     log:
                         Grpc:
                             host: "localhost"
@@ -236,6 +237,7 @@ mod tests {
                             credentials: Minio
                             connect_timeout_ms: 5000
                             request_timeout_ms: 1000
+                            upload_part_size_bytes: 8388608
                     log:
                         Grpc:
                             host: "localhost"
@@ -314,6 +316,7 @@ mod tests {
                             credentials: Minio
                             connect_timeout_ms: 5000
                             request_timeout_ms: 1000
+                            upload_part_size_bytes: 8388608
                     log:
                         Grpc:
                             host: "localhost"
@@ -366,6 +369,7 @@ mod tests {
                             credentials: Minio
                             connect_timeout_ms: 5000
                             request_timeout_ms: 1000
+                            upload_part_size_bytes: 8388608
                     log:
                         Grpc:
                             host: "localhost"
@@ -462,6 +466,7 @@ mod tests {
                             credentials: Minio
                             connect_timeout_ms: 5000
                             request_timeout_ms: 1000
+                            upload_part_size_bytes: 8388608
                     log:
                         Grpc:
                             host: "localhost"
@@ -514,6 +519,7 @@ mod tests {
                             credentials: Minio
                             connect_timeout_ms: 5000
                             request_timeout_ms: 1000
+                            upload_part_size_bytes: 8388608
                     log:
                         Grpc:
                             host: "localhost"
@@ -570,6 +576,10 @@ mod tests {
             let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__MY_PORT", 50051);
             let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__BUCKET", "buckets!");
             let _ = jail.set_env("CHROMA_COMPACTION_SERVICE__STORAGE__S3__CREDENTIALS", "AWS");
+            let _ = jail.set_env(
+                "CHROMA_COMPACTION_SERVICE__STORAGE__S3__upload_part_size_bytes",
+                format!("{}", 1024 * 1024 * 8),
+            );
             let _ = jail.set_env(
                 "CHROMA_COMPACTION_SERVICE__STORAGE__S3__CONNECT_TIMEOUT_MS",
                 5000,
@@ -604,6 +614,7 @@ mod tests {
                             credentials: Minio
                             connect_timeout_ms: 5000
                             request_timeout_ms: 1000
+                            upload_part_size_bytes: 8388608
                     log:
                         Grpc:
                             host: "localhost"
@@ -698,6 +709,7 @@ mod tests {
                     );
                     assert_eq!(s.connect_timeout_ms, 5000);
                     assert_eq!(s.request_timeout_ms, 1000);
+                    assert_eq!(s.upload_part_size_bytes, 1024 * 1024 * 8);
                 }
                 _ => panic!("Invalid storage config"),
             }

diff --git a/rust/worker/src/storage/config.rs b/rust/worker/src/storage/config.rs
@@ -30,6 +30,7 @@ pub(crate) struct S3StorageConfig {
     pub(crate) credentials: S3CredentialsConfig,
     pub(crate) connect_timeout_ms: u64,
     pub(crate) request_timeout_ms: u64,
+    pub(crate) upload_part_size_bytes: usize,
 }
 
 #[derive(Deserialize, Debug)]