Repository navigation
feat: add NGC model provider to support NGC Downloads (#222) - #232
Merged
Merged
Conversation
Add NVIDIA NGC as a model provider alongside HuggingFace, enabling ModelExpress to download, cache, and manage models from the NGC catalog. NGC provider: - Implement NgcProvider (ModelProviderTrait) and NgcProviderCache (ProviderCache) - NGC API integration: token exchange via authn.nvidia.com, paginated file listing, streaming downloads with concurrent file transfers - Support V1 (presigned URL) and V2 (Bearer auth) NGC storage backends - Support all NGC API key types: personal (nvapi-), legacy, and service keys - UAM fallback: when bulk file listing is gated (HTTP 400/401), enumerate files via checksums.blake3 manifest - Model name parsing: org/name/version, org/team/name/version, org/team/type/name/version, and ngc:// URI scheme - API key resolution: NGC_API_KEY, NGC_CLI_API_KEY env vars, ~/.ngc/config (INI and JSON formats) - Parallel file downloads using buffer_unordered with min(num_cpus, 10) concurrency, matching nim-sdk patterns - Thread-safe Bearer token refresh via RwLock with double-checked locking to handle tokens expiring during large model downloads - Atomic writes (temp file + rename) to prevent partial files on interruption - Path traversal protection in model name parsing - Input whitespace trimming for robustness Integration: - Extend ModelProvider enum, proto (NGC = 1), gRPC conversions, database serialization - Wire into cache.rs, download.rs, lib.rs, providers.rs Helm: - Document HF_TOKEN and NGC_API_KEY secret injection in values.yaml and values-local-storage.yaml via extraEnv Documentation: - Update ARCHITECTURE.md and DEPLOYMENT.md for NGC provider - Fix provider directory tree to reflect providers.rs + providers/ layout Tests: - 31 unit tests covering parsing, cache operations, API response handling, V1/V2 storage, round-trip model names, path traversal rejection Signed-off-by: Sunil Upare <sunilupare45@gmail.com>
Contributor
WalkthroughThis PR extends model download infrastructure to support NVIDIA NGC models alongside existing HuggingFace support. Changes include new provider implementation, domain model extensions, provider routing, deployment configuration documentation, and server-side database integration. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~70 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 1
🧹 Nitpick comments (4)
docs/DEPLOYMENT.md (1)
192-197: DocumentNGC_CLI_HOMEoverride in key resolution order.The precedence list is almost complete; the implementation also supports overriding the config directory via
NGC_CLI_HOMEbefore reading.ngc/config. Adding that note will keep docs fully aligned with runtime behavior.Suggested doc tweak
-3. `~/.ngc/config` (written by `ngc config set`) +3. `${NGC_CLI_HOME}/.ngc/config` if `NGC_CLI_HOME` is set +4. `~/.ngc/config` (written by `ngc config set`)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/DEPLOYMENT.md` around lines 192 - 197, The docs list NGC API key resolution order but omit the NGC_CLI_HOME override; update the precedence description to mention that before reading ~/.ngc/config (created by ngc config set) the server will respect an alternate config directory set via the NGC_CLI_HOME environment variable, so the full resolution order becomes: NGC_API_KEY env var, NGC_CLI_API_KEY env var, then ~/.ngc/config (or the directory pointed to by NGC_CLI_HOME) — mention NGC_CLI_HOME explicitly and that it overrides the default .ngc config path.modelexpress_common/src/lib.rs (1)
140-155: Add an explicitNgcround-trip conversion test.The new match arms are correct, but the current provider conversion test only asserts
HuggingFace. AddingNgcintest_model_provider_conversion_both_wayswill lock in this mapping.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelexpress_common/src/lib.rs` around lines 140 - 155, Add an explicit assertion for the Ngc variant in the existing round-trip conversion test: update test_model_provider_conversion_both_ways to also assert that models::ModelProvider::Ngc converts to grpc::model::ModelProvider::Ngc and back to models::ModelProvider::Ngc, verifying the impl From conversions between models::ModelProvider and grpc::model::ModelProvider both ways.modelexpress_common/src/providers/ngc.rs (2)
648-668: Consider adding temp file cleanup on error.If an error occurs after creating the temp file (e.g., during write or flush), the
.tmpfile will be left behind. While this isn't critical (subsequent downloads will overwrite it), explicit cleanup on error would be cleaner.Optional: Add cleanup on error
let temp_path = dest.with_extension("tmp"); let mut file = tokio::fs::File::create(&temp_path) .await .with_context(|| format!("Failed to create {temp_path:?}"))?; let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { - let bytes = chunk.with_context(|| format!("Error reading response body for {rel_path}"))?; - file.write_all(&bytes) - .await - .with_context(|| format!("Failed to write to {temp_path:?}"))?; + let write_result = async { + while let Some(chunk) = stream.next().await { + let bytes = chunk.with_context(|| format!("Error reading response body for {rel_path}"))?; + file.write_all(&bytes) + .await + .with_context(|| format!("Failed to write to {temp_path:?}"))?; + } + file.flush() + .await + .with_context(|| format!("Failed to flush {temp_path:?}"))?; + Ok::<_, anyhow::Error>(()) + }.await; + + if let Err(e) = write_result { + let _ = tokio::fs::remove_file(&temp_path).await; + return Err(e); } - - file.flush() - .await - .with_context(|| format!("Failed to flush {temp_path:?}"))?; drop(file);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelexpress_common/src/providers/ngc.rs` around lines 648 - 668, The temp file at temp_path can be left behind if an error occurs after creation; wrap the download/write/flush/rename sequence in a block that captures the result, and on Err call tokio::fs::remove_file(&temp_path).await.ok() before returning the error. Concretely, after creating temp_path and file, perform the stream loop, flush, and rename inside a let result = (async { ... }).await or a match so you can do if let Err(e) = result { tokio::fs::remove_file(&temp_path).await.ok(); return Err(e); } — reference symbols: temp_path, dest, file, stream, response.bytes_stream(), and the final tokio::fs::rename call.
965-971: Edge case: unreadable directory returns as cached.If
read_dirfails (e.g., permission denied), the code treats the directory as non-empty and returns success. This might mask permission issues, though such cases are rare in practice.Optional: Propagate read_dir errors
- let is_empty = match tokio::fs::read_dir(&path).await { - Ok(mut rd) => rd.next_entry().await.ok().flatten().is_none(), - Err(_) => false, - }; + let mut rd = tokio::fs::read_dir(&path) + .await + .with_context(|| format!("Failed to read cache directory {path:?}"))?; + let is_empty = rd.next_entry().await.ok().flatten().is_none();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelexpress_common/src/providers/ngc.rs` around lines 965 - 971, The current is_empty check swallows tokio::fs::read_dir errors and treats unreadable directories as non-empty; instead surface failures from tokio::fs::read_dir so permission/IO errors aren’t masked. Change the match around tokio::fs::read_dir(&path).await in ngc.rs (the is_empty calculation) to propagate the Err (e.g., return/bail with the error using anyhow::Context or map the error into anyhow::Error) rather than converting Err(_) to false, so read_dir failures cause the caller to see the error and bail.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@helm/values-local-storage.yaml`:
- Around line 81-94: Update the commented secret names to match the repo
convention by replacing usages of hf-token and ngc-api-key with hf-token-secret
and ngc-api-key-secret: change the example kubectl create secret commands and
the corresponding commented env var blocks (the secretKeyRef.name values in the
HF_TOKEN and NGC_API_KEY sections) so they reference hf-token-secret and
ngc-api-key-secret instead of hf-token / ngc-api-key, ensuring consistent naming
across examples.
---
Nitpick comments:
In `@docs/DEPLOYMENT.md`:
- Around line 192-197: The docs list NGC API key resolution order but omit the
NGC_CLI_HOME override; update the precedence description to mention that before
reading ~/.ngc/config (created by ngc config set) the server will respect an
alternate config directory set via the NGC_CLI_HOME environment variable, so the
full resolution order becomes: NGC_API_KEY env var, NGC_CLI_API_KEY env var,
then ~/.ngc/config (or the directory pointed to by NGC_CLI_HOME) — mention
NGC_CLI_HOME explicitly and that it overrides the default .ngc config path.
In `@modelexpress_common/src/lib.rs`:
- Around line 140-155: Add an explicit assertion for the Ngc variant in the
existing round-trip conversion test: update
test_model_provider_conversion_both_ways to also assert that
models::ModelProvider::Ngc converts to grpc::model::ModelProvider::Ngc and back
to models::ModelProvider::Ngc, verifying the impl From conversions between
models::ModelProvider and grpc::model::ModelProvider both ways.
In `@modelexpress_common/src/providers/ngc.rs`:
- Around line 648-668: The temp file at temp_path can be left behind if an error
occurs after creation; wrap the download/write/flush/rename sequence in a block
that captures the result, and on Err call
tokio::fs::remove_file(&temp_path).await.ok() before returning the error.
Concretely, after creating temp_path and file, perform the stream loop, flush,
and rename inside a let result = (async { ... }).await or a match so you can do
if let Err(e) = result { tokio::fs::remove_file(&temp_path).await.ok(); return
Err(e); } — reference symbols: temp_path, dest, file, stream,
response.bytes_stream(), and the final tokio::fs::rename call.
- Around line 965-971: The current is_empty check swallows tokio::fs::read_dir
errors and treats unreadable directories as non-empty; instead surface failures
from tokio::fs::read_dir so permission/IO errors aren’t masked. Change the match
around tokio::fs::read_dir(&path).await in ngc.rs (the is_empty calculation) to
propagate the Err (e.g., return/bail with the error using anyhow::Context or map
the error into anyhow::Error) rather than converting Err(_) to false, so
read_dir failures cause the caller to see the error and bail.
🪄 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: Pro
Run ID: 80d8e3b9-2777-4819-91a8-c700a5c5224b
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (14)
Cargo.tomldocs/ARCHITECTURE.mddocs/DEPLOYMENT.mdhelm/values-local-storage.yamlhelm/values.yamlmodelexpress_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/ngc.rsmodelexpress_server/src/database.rs
… fix helm secret names Signed-off-by: Sunil Upare <sunilupare45@gmail.com>
KavinKrishnan
approved these changes
Apr 15, 2026
This was referenced Jun 1, 2026
This was referenced Sep 30, 2026
This branch was previously deployed
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
2 participants
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.
Implement NgcProvider and NgcProviderCache to download models from NVIDIA NGC alongside the existing HuggingFace provider.
Few Caveats:
Summary by CodeRabbit
New Features
Documentation