Skip to content

Commit

Permalink
Rename part_size_bytes -> upload_part_size_bytes
Browse files Browse the repository at this point in the history
  • Loading branch information
@codetheweb
codetheweb committed Jul 31, 2024
1 parent a38b817 commit adb5028
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 30 deletions.
4 changes: 2 additions & 2 deletions rust/worker/chroma_config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ query_service:
credentials: "Minio"
connect_timeout_ms: 5000
request_timeout_ms: 30000 # 1 minute
part_size_bytes: 8388608 # 8MB
upload_part_size_bytes: 8388608 # 8MB
log:
Grpc:
host: "logservice.chroma"
Expand Down Expand Up @@ -81,7 +81,7 @@ compaction_service:
credentials: "Minio"
connect_timeout_ms: 5000
request_timeout_ms: 60000 # 1 minute
part_size_bytes: 8388608 # 8MB
upload_part_size_bytes: 8388608 # 8MB
log:
Grpc:
host: "logservice.chroma"
Expand Down
18 changes: 9 additions & 9 deletions rust/worker/src/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ mod tests {
credentials: Minio
connect_timeout_ms: 5000
request_timeout_ms: 1000
part_size_bytes: 8388608
upload_part_size_bytes: 8388608
log:
Grpc:
host: "localhost"
Expand Down Expand Up @@ -237,7 +237,7 @@ mod tests {
credentials: Minio
connect_timeout_ms: 5000
request_timeout_ms: 1000
part_size_bytes: 8388608
upload_part_size_bytes: 8388608
log:
Grpc:
host: "localhost"
Expand Down Expand Up @@ -316,7 +316,7 @@ mod tests {
credentials: Minio
connect_timeout_ms: 5000
request_timeout_ms: 1000
part_size_bytes: 8388608
upload_part_size_bytes: 8388608
log:
Grpc:
host: "localhost"
Expand Down Expand Up @@ -369,7 +369,7 @@ mod tests {
credentials: Minio
connect_timeout_ms: 5000
request_timeout_ms: 1000
part_size_bytes: 8388608
upload_part_size_bytes: 8388608
log:
Grpc:
host: "localhost"
Expand Down Expand Up @@ -466,7 +466,7 @@ mod tests {
credentials: Minio
connect_timeout_ms: 5000
request_timeout_ms: 1000
part_size_bytes: 8388608
upload_part_size_bytes: 8388608
log:
Grpc:
host: "localhost"
Expand Down Expand Up @@ -519,7 +519,7 @@ mod tests {
credentials: Minio
connect_timeout_ms: 5000
request_timeout_ms: 1000
part_size_bytes: 8388608
upload_part_size_bytes: 8388608
log:
Grpc:
host: "localhost"
Expand Down Expand Up @@ -577,7 +577,7 @@ mod tests {
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__PART_SIZE_BYTES",
"CHROMA_COMPACTION_SERVICE__STORAGE__S3__upload_part_size_bytes",
format!("{}", 1024 * 1024 * 8),
);
let _ = jail.set_env(
Expand Down Expand Up @@ -614,7 +614,7 @@ mod tests {
credentials: Minio
connect_timeout_ms: 5000
request_timeout_ms: 1000
part_size_bytes: 8388608
upload_part_size_bytes: 8388608
log:
Grpc:
host: "localhost"
Expand Down Expand Up @@ -709,7 +709,7 @@ mod tests {
);
assert_eq!(s.connect_timeout_ms, 5000);
assert_eq!(s.request_timeout_ms, 1000);
assert_eq!(s.part_size_bytes, 1024 * 1024 * 8);
assert_eq!(s.upload_part_size_bytes, 1024 * 1024 * 8);
}
_ => panic!("Invalid storage config"),
}
Expand Down
2 changes: 1 addition & 1 deletion rust/worker/src/storage/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ pub(crate) struct S3StorageConfig {
pub(crate) credentials: S3CredentialsConfig,
pub(crate) connect_timeout_ms: u64,
pub(crate) request_timeout_ms: u64,
pub(crate) part_size_bytes: u64,
pub(crate) upload_part_size_bytes: u64,
}

#[derive(Deserialize, Debug)]
Expand Down
41 changes: 23 additions & 18 deletions rust/worker/src/storage/s3.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use thiserror::Error;
pub(crate) struct S3Storage {
bucket: String,
client: aws_sdk_s3::Client,
part_size_bytes: u64,
upload_part_size_bytes: u64,
}

#[derive(Error, Debug)]
Expand Down Expand Up @@ -67,11 +67,11 @@ impl ChromaError for S3GetError {
}

impl S3Storage {
fn new(bucket: &str, client: aws_sdk_s3::Client, part_size_bytes: u64) -> S3Storage {
fn new(bucket: &str, client: aws_sdk_s3::Client, upload_part_size_bytes: u64) -> S3Storage {
return S3Storage {
bucket: bucket.to_string(),
client,
part_size_bytes,
upload_part_size_bytes,
};
}

Expand Down Expand Up @@ -298,24 +298,24 @@ impl S3Storage {
&self,
total_size_bytes: u64,
) -> impl Iterator<Item = (i32, u64, u64)> {
let part_size_bytes = self.part_size_bytes.clone();
let mut part_count = (total_size_bytes / part_size_bytes) + 1;
let mut size_of_last_part = total_size_bytes % part_size_bytes;
let upload_part_size_bytes = self.upload_part_size_bytes.clone();
let mut part_count = (total_size_bytes / upload_part_size_bytes) + 1;
let mut size_of_last_part = total_size_bytes % upload_part_size_bytes;
if size_of_last_part == 0 {
size_of_last_part = part_size_bytes;
size_of_last_part = upload_part_size_bytes;
part_count -= 1;
}

(0..part_count).map(move |part_index| {
let this_part = if part_count - 1 == part_index {
size_of_last_part
} else {
part_size_bytes
upload_part_size_bytes
};
(
// Part numbers start at 1
(part_index + 1) as i32,
part_index * part_size_bytes,
part_index * upload_part_size_bytes,
this_part,
)
})
Expand Down Expand Up @@ -377,7 +377,8 @@ impl Configurable<StorageConfig> for S3Storage {
aws_sdk_s3::Client::new(&config)
}
};
let storage = S3Storage::new(&s3_config.bucket, client, s3_config.part_size_bytes);
let storage =
S3Storage::new(&s3_config.bucket, client, s3_config.upload_part_size_bytes);
// for minio we create the bucket since it is only used for testing
match &s3_config.credentials {
super::config::S3CredentialsConfig::Minio => {
Expand Down Expand Up @@ -440,7 +441,7 @@ mod tests {
let storage = S3Storage {
bucket: "test".to_string(),
client,
part_size_bytes: 1024 * 1024 * 8,
upload_part_size_bytes: 1024 * 1024 * 8,
};
storage.create_bucket().await.unwrap();

Expand Down Expand Up @@ -468,13 +469,13 @@ mod tests {
assert_eq!(buf, test_data);
}

async fn test_put_file(file_size: usize, part_size_bytes: u64) {
async fn test_put_file(file_size: usize, upload_part_size_bytes: u64) {
let client = get_s3_client();

let storage = S3Storage {
bucket: "test".to_string(),
client,
part_size_bytes,
upload_part_size_bytes,
};
storage.create_bucket().await.unwrap();

Expand Down Expand Up @@ -517,16 +518,20 @@ mod tests {
#[tokio::test]
#[cfg(CHROMA_KUBERNETES_INTEGRATION)]
async fn test_put_file_scenarios() {
let test_part_size_bytes = 1024 * 1024 * 8; // 8MB
let test_upload_part_size_bytes = 1024 * 1024 * 8; // 8MB

// Under part size
test_put_file(1024, test_part_size_bytes).await;
test_put_file(1024, test_upload_part_size_bytes).await;
// At part size
test_put_file(test_part_size_bytes as usize, test_part_size_bytes).await;
test_put_file(
test_upload_part_size_bytes as usize,
test_upload_part_size_bytes,
)
.await;
// Over part size
test_put_file(
(test_part_size_bytes as f64 * 2.5) as usize,
test_part_size_bytes,
(test_upload_part_size_bytes as f64 * 2.5) as usize,
test_upload_part_size_bytes,
)
.await;
}
Expand Down

0 comments on commit adb5028

Please sign in to comment.