feat: add s3 model provider#440
Open
Broduker wants to merge 3 commits into
Open
Conversation
Signed-off-by: shenls <shenlinshan@kanzhun.com>
Contributor
WalkthroughAdds S3 as a new model provider across the full stack: extends the ChangesS3 Model Provider End-to-End
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
modelexpress_common/src/providers/s3.rs (1)
208-212: 💤 Low valueBlocking I/O in async context.
fs::create_dir_allis a blocking call within an async function. While acceptable for quick operations, consider usingtokio::fs::create_dir_allfor consistency with the async operations below (e.g.,tokio::fs::writeat line 240).♻️ Suggested refactor
- fs::create_dir_all(&cache_dir).with_context(|| { + tokio::fs::create_dir_all(&cache_dir).await.with_context(|| { format!("Failed to create cache directory: {}", cache_dir.display()) })?;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_common/src/providers/s3.rs` around lines 208 - 212, Replace the blocking fs::create_dir_all call with its async equivalent tokio::fs::create_dir_all to avoid blocking I/O in the async context. Since tokio::fs::create_dir_all is an async function, you will need to await it. This ensures consistency with other async file operations used later in the function, such as tokio::fs::write.modelexpress_client/src/bin/modules/handlers.rs (2)
62-63: ⚡ Quick winAdd a dedicated S3 validation-path regression test.
The new S3 inference/canonicalization branches are untested in this module, while GCS has explicit coverage. Add a sibling test to lock in behavior for mixed-case
S3://...inputs.✅ Suggested test
+ #[test] + fn test_validate_resolves_s3_model_name_to_s3_cache_path() { + let temp_dir = TempDir::new().expect("Failed to create temp dir"); + let model_name = "S3://bucket/foo/bar"; + let expected_path = temp_dir + .path() + .join("s3") + .join("bucket") + .join("foo") + .join("bar"); + fs::create_dir_all(&expected_path).expect("Failed to create S3 cache path"); + + let (provider, model_path) = resolve_validation_model_path(temp_dir.path(), model_name); + + assert_eq!(provider, ModelProvider::S3); + assert_eq!(model_path, expected_path); + assert!(model_path.exists()); + assert!(!temp_dir.path().join("S3://bucket/foo/bar").exists()); + }Also applies to: 84-85
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/src/bin/modules/handlers.rs` around lines 62 - 63, The new S3 inference/canonicalization branches (ModelProvider::S3 handling) in the handlers module lack test coverage while GCS has explicit regression tests. Add a sibling test to the existing GCS test cases that validates the behavior for mixed-case S3:// prefixes (such as S3://, s3://, S3://, etc.) to ensure the strip_ascii_prefix_ignore_case function correctly identifies and canonicalizes S3 model provider inputs, mirroring the test coverage pattern already established for GCS.
58-69: ⚡ Quick winUse the shared provider resolver to prevent contract drift.
Line 58-Line 69 duplicates scheme-to-provider logic that already lives in
modelexpress_common/src/models.rs(ModelProvider::resolve_for_model_name). Reusing it here keeps client validation aligned with the core provider contract as new schemes are added.♻️ Suggested refactor
fn infer_provider_from_model_name(model_name: &str) -> ModelProvider { - let model_name = model_name.trim_start(); - if strip_ascii_prefix_ignore_case(model_name, "gs://").is_some() { - ModelProvider::Gcs - } else if strip_ascii_prefix_ignore_case(model_name, "s3://").is_some() { - ModelProvider::S3 - } else if strip_ascii_prefix_ignore_case(model_name, "ngc://").is_some() { - ModelProvider::Ngc - } else { - ModelProvider::HuggingFace - } + ModelProvider::resolve_for_model_name(model_name, ModelProvider::HuggingFace) }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/src/bin/modules/handlers.rs` around lines 58 - 69, The function infer_provider_from_model_name duplicates scheme-to-provider resolution logic that is already implemented in the shared ModelProvider::resolve_for_model_name function. Remove the duplicated if-else chain in infer_provider_from_model_name and instead directly call ModelProvider::resolve_for_model_name with the trimmed model_name, returning its result. This prevents contract drift and ensures the client validation stays synchronized with the core provider contract as new schemes are added.modelexpress_client/src/lib.rs (1)
281-281: ⚡ Quick winDeduplicate effective-provider resolution across client entry points.
The same
ModelProvider::resolve_for_model_namecall is repeated in four methods. A small helper keeps this contract in one place and reduces future drift.♻️ Suggested refactor
impl Client { + #[inline] + fn resolve_provider_for_model_name(model_name: &str, provider: ModelProvider) -> ModelProvider { + ModelProvider::resolve_for_model_name(model_name, provider) + } + pub async fn get_model_path( &self, model_name: &str, provider: ModelProvider, ) -> anyhow::Result<PathBuf> { - let provider = ModelProvider::resolve_for_model_name(model_name, provider); + let provider = Self::resolve_provider_for_model_name(model_name, provider); @@ pub async fn request_model_on_server( @@ - let provider = ModelProvider::resolve_for_model_name(&model_name, provider); + let provider = Self::resolve_provider_for_model_name(&model_name, provider); @@ pub async fn request_model( @@ - let provider = ModelProvider::resolve_for_model_name(&model_name, provider); + let provider = Self::resolve_provider_for_model_name(&model_name, provider); @@ pub async fn request_model_with_smart_fallback( @@ - let provider = ModelProvider::resolve_for_model_name(&model_name, provider); + let provider = Self::resolve_provider_for_model_name(&model_name, provider);Also applies to: 635-635, 706-706, 741-741
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/src/lib.rs` at line 281, The ModelProvider::resolve_for_model_name call is duplicated across four client entry point methods (appearing at lines 281, 635, 706, and 741). Create a private helper method in the client that wraps this ModelProvider::resolve_for_model_name call with the model_name and provider parameters, then replace all four duplicate calls with invocations of this new helper method. This centralizes the contract in one place and prevents future inconsistencies across these entry points.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelexpress_server/src/registry/backend/kubernetes.rs`:
- Around line 31-36: The ALL_PROVIDERS constant now includes ModelProvider::S3,
expanding the array to 4 providers which means candidate_cr_names() will now
return 5 entries (4 providers plus a legacy entry). Update the test that
validates candidate_cr_names() output (currently expecting 4 entries) to expect
5 entries instead, and add membership assertions to verify that
ModelProvider::S3 is included in the candidate names returned by the function.
This ensures the test expectations align with the expanded provider list.
In `@modelexpress_server/src/registry/backend/redis.rs`:
- Around line 42-47: The test for the `candidate_keys()` function needs to be
updated to reflect the addition of S3 to the `ALL_PROVIDERS` constant. Update
the test assertions around line 594 to expect 5 keys instead of 4 (accounting
for the 4 providers in `ALL_PROVIDERS` plus the legacy key), and add an
assertion to verify that the S3 provider key is included in the returned
candidate keys list.
---
Nitpick comments:
In `@modelexpress_client/src/bin/modules/handlers.rs`:
- Around line 62-63: The new S3 inference/canonicalization branches
(ModelProvider::S3 handling) in the handlers module lack test coverage while GCS
has explicit regression tests. Add a sibling test to the existing GCS test cases
that validates the behavior for mixed-case S3:// prefixes (such as S3://, s3://,
S3://, etc.) to ensure the strip_ascii_prefix_ignore_case function correctly
identifies and canonicalizes S3 model provider inputs, mirroring the test
coverage pattern already established for GCS.
- Around line 58-69: The function infer_provider_from_model_name duplicates
scheme-to-provider resolution logic that is already implemented in the shared
ModelProvider::resolve_for_model_name function. Remove the duplicated if-else
chain in infer_provider_from_model_name and instead directly call
ModelProvider::resolve_for_model_name with the trimmed model_name, returning its
result. This prevents contract drift and ensures the client validation stays
synchronized with the core provider contract as new schemes are added.
In `@modelexpress_client/src/lib.rs`:
- Line 281: The ModelProvider::resolve_for_model_name call is duplicated across
four client entry point methods (appearing at lines 281, 635, 706, and 741).
Create a private helper method in the client that wraps this
ModelProvider::resolve_for_model_name call with the model_name and provider
parameters, then replace all four duplicate calls with invocations of this new
helper method. This centralizes the contract in one place and prevents future
inconsistencies across these entry points.
In `@modelexpress_common/src/providers/s3.rs`:
- Around line 208-212: Replace the blocking fs::create_dir_all call with its
async equivalent tokio::fs::create_dir_all to avoid blocking I/O in the async
context. Since tokio::fs::create_dir_all is an async function, you will need to
await it. This ensures consistency with other async file operations used later
in the function, such as tokio::fs::write.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b44da810-a72d-4969-b574-15ea710dad44
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (13)
Cargo.tomlmodelexpress_client/src/bin/modules/handlers.rsmodelexpress_client/src/lib.rsmodelexpress_common/Cargo.tomlmodelexpress_common/proto/model.protomodelexpress_common/src/cache.rsmodelexpress_common/src/download.rsmodelexpress_common/src/lib.rsmodelexpress_common/src/models.rsmodelexpress_common/src/providers.rsmodelexpress_common/src/providers/s3.rsmodelexpress_server/src/registry/backend/kubernetes.rsmodelexpress_server/src/registry/backend/redis.rs
Comment on lines
+31
to
36
| const ALL_PROVIDERS: [ModelProvider; 4] = [ | ||
| ModelProvider::HuggingFace, | ||
| ModelProvider::Ngc, | ||
| ModelProvider::Gcs, | ||
| ModelProvider::S3, | ||
| ]; |
Contributor
There was a problem hiding this comment.
Update candidate-name test expectations after expanding ALL_PROVIDERS.
Adding ModelProvider::S3 here increases candidate_cr_names() output to 5 entries (4 providers + legacy), but the test at Line 705 still expects 4 and omits S3 membership checks. This will fail tests consistently.
Suggested test fix
- assert_eq!(candidates.len(), 4);
+ assert_eq!(candidates.len(), 5);
for p in [
ModelProvider::HuggingFace,
ModelProvider::Ngc,
ModelProvider::Gcs,
+ ModelProvider::S3,
] {
assert!(candidates.contains(&KubernetesRegistryBackend::cr_name_for(p, n)));
}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_server/src/registry/backend/kubernetes.rs` around lines 31 - 36,
The ALL_PROVIDERS constant now includes ModelProvider::S3, expanding the array
to 4 providers which means candidate_cr_names() will now return 5 entries (4
providers plus a legacy entry). Update the test that validates
candidate_cr_names() output (currently expecting 4 entries) to expect 5 entries
instead, and add membership assertions to verify that ModelProvider::S3 is
included in the candidate names returned by the function. This ensures the test
expectations align with the expanded provider list.
Comment on lines
+42
to
47
| const ALL_PROVIDERS: [ModelProvider; 4] = [ | ||
| ModelProvider::HuggingFace, | ||
| ModelProvider::Ngc, | ||
| ModelProvider::Gcs, | ||
| ModelProvider::S3, | ||
| ]; |
Contributor
There was a problem hiding this comment.
Fix stale candidate-key test after adding S3 to provider fan-out.
With S3 added to ALL_PROVIDERS, candidate_keys() now returns 5 keys (4 providers + legacy), but the test at Line 594 still asserts 4 and misses the S3 key assertion.
Suggested test fix
- assert_eq!(keys.len(), 4);
+ assert_eq!(keys.len(), 5);
assert!(keys.contains(&"mx:model:HuggingFace:org/model".to_string()));
assert!(keys.contains(&"mx:model:Ngc:org/model".to_string()));
assert!(keys.contains(&"mx:model:Gcs:org/model".to_string()));
+ assert!(keys.contains(&"mx:model:S3:org/model".to_string()));
assert!(keys.contains(&"mx:model:org/model".to_string())); // legacy🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_server/src/registry/backend/redis.rs` around lines 42 - 47, The
test for the `candidate_keys()` function needs to be updated to reflect the
addition of S3 to the `ALL_PROVIDERS` constant. Update the test assertions
around line 594 to expect 5 keys instead of 4 (accounting for the 4 providers in
`ALL_PROVIDERS` plus the legacy key), and add an assertion to verify that the S3
provider key is included in the returned candidate keys list.
Signed-off-by: shenls <shenlinshan@kanzhun.com>
Signed-off-by: shenls <shenlinshan@kanzhun.com>
Author
|
coud anyone review it? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview:
Adds S3-compatible object storage as a ModelExpress model provider for metadata downloads and provider-aware request routing.
Details:
ModelProvider::S3to the protobuf and Rust provider model.object_storefor metadata-only downloads froms3://bucket/prefixinto the ModelExpress cache.ModelProvider::resolve_for_model_nameso clients resolves3://,gs://, andngc://URIs before sending requests.Where should the reviewer start?
Start with
modelexpress_common/src/providers/s3.rsfor the new provider implementation, thenmodelexpress_common/src/models.rsandmodelexpress_client/src/lib.rsfor provider resolution and request routing.Related Issues
Summary by CodeRabbit
s3://URLs