feat: add NGC model provider to support NGC Downloads (#222)

Cargo.toml

@@ -55,6 +55,7 @@ futures = "0.3"
 uuid = { version = "1.17", features = ["v4", "serde"] }
 thiserror = "2.0"
 redis = { version = "0.27", features = ["tokio-comp", "connection-manager"] }
+reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream"] }
 
 # Kubernetes client (kube-rs) for CRD-based metadata backend
 kube = { version = "0.98", features = [

docs/ARCHITECTURE.md

@@ -11,7 +11,7 @@ Detailed reference document for the ModelExpress codebase. For deployment and co
 
 ModelExpress is a Rust-based model cache management service and GPU-to-GPU model weight transfer system. It serves two roles:
 
-- **Model Cache Service** - A sidecar alongside inference solutions (vLLM, SGLang, NVIDIA Dynamo) that accelerates model downloads from HuggingFace with SQLite-backed tracking and LRU cache eviction.
+- **Model Cache Service** - A sidecar alongside inference solutions (vLLM, SGLang, NVIDIA Dynamo) that accelerates model downloads from HuggingFace and NGC with SQLite-backed tracking and LRU cache eviction.
 - **P2P Weight Transfer** - GPU-to-GPU model weight transfers between vLLM instances using NVIDIA NIXL over RDMA/InfiniBand, enabling ~15-second transfers for 681GB models.
 
 ### Current Status
@@ -29,6 +29,7 @@ graph TD
         C1[Client CLI / Library] -->|gRPC| S1[ModelExpress Server]
         S1 --> DB[(SQLite)]
         S1 --> HF[HuggingFace Hub]
+        S1 --> NGC[NVIDIA NGC]
         S1 --> Cache[Model Cache Dir]
     end
 
@@ -150,9 +151,10 @@ ModelExpress/
 │       ├── config.rs                   # Config trait utilities
 │       ├── download.rs                 # Download orchestration (smart-fallback, direct, server-only)
 │       ├── models.rs                   # Status, ModelProvider, ModelStatus, ModelStatusResponse
+│       ├── providers.rs               # ModelProviderTrait definition, re-exports
 │       └── providers/
-│           ├── mod.rs                  # ModelProviderTrait definition
-│           └── huggingface.rs          # HuggingFaceProvider implementation
+│           ├── huggingface.rs          # HuggingFaceProvider implementation
+│           └── ngc.rs                  # NgcProvider implementation
 │
 ├── workspace-tests/
 │   ├── Cargo.toml
@@ -251,7 +253,7 @@ Four proto files define four services, all compiled via `tonic-build` in `modele
 | `StreamModelFiles` | `ModelFilesRequest` | stream `FileChunk` | Stream model file contents (1MB chunks) |
 | `ListModelFiles` | `ModelFilesRequest` | `ModelFileList` | List files with sizes |
 
-Key message types: `ModelProvider` (HuggingFace), `ModelStatus` (Downloading, Downloaded, Error), `ModelStatusUpdate`, `FileChunk`.
+Key message types: `ModelProvider` (HuggingFace, NGC), `ModelStatus` (Downloading, Downloaded, Error), `ModelStatusUpdate`, `FileChunk`.
 
 ### p2p.proto - P2pService
 
@@ -409,7 +411,7 @@ Output formats: `--format human` (default), `--format json`, `--format json-pret
 | `config` | Config trait utilities |
 | `download` | Download orchestration with strategy pattern |
 | `models` | `Status`, `ModelProvider`, `ModelStatus`, `ModelStatusResponse` |
-| `providers` | `ModelProviderTrait` + `HuggingFaceProvider` |
+| `providers` | `ModelProviderTrait` + `HuggingFaceProvider` + `NgcProvider` |
 | `grpc` | Generated tonic stubs for all 4 services |
 | `constants` | `DEFAULT_GRPC_PORT` (8001), `DEFAULT_TIMEOUT_SECS` (30), `DEFAULT_TRANSFER_CHUNK_SIZE` (32KB) |
 
@@ -428,7 +430,9 @@ pub trait ModelProviderTrait: Send + Sync {
 }
 ```
 
-Currently one implementation: `HuggingFaceProvider` (uses `hf-hub` crate with high-CPU download mode).
+Two implementations:
+- `HuggingFaceProvider` — uses the `hf-hub` crate with high-CPU download mode.
+- `NgcProvider` — downloads from NVIDIA NGC via the V2 artifact API (Bearer-authenticated `/files/{path}` for team artifacts; presigned S3 URLs for org-level artifacts). Falls back to `checksums.blake3` manifest enumeration when bulk file listing returns 400. Resolves the NGC API key from `NGC_API_KEY`, `NGC_CLI_API_KEY`, or `~/.ngc/config`.
 
 ### ClientConfig / ClientArgs
 

docs/DEPLOYMENT.md

@@ -131,7 +131,9 @@ The CLI client also uses layered configuration: CLI args > env vars > config fil
 | `MODEL_EXPRESS_NO_SHARED_STORAGE` | `false` | Use gRPC streaming instead of shared storage |
 | `MODEL_EXPRESS_TRANSFER_CHUNK_SIZE` | `32768` | Transfer chunk size (bytes) |
 
-Cache directory resolution: `MODEL_EXPRESS_CACHE_DIRECTORY` -> `HF_HUB_CACHE` -> `~/.cache/huggingface/hub`.
+Cache directory resolution for HuggingFace: `MODEL_EXPRESS_CACHE_DIRECTORY` -> `HF_HUB_CACHE` -> `~/.cache/huggingface/hub`.
+
+Cache directory resolution for NGC: `MODEL_EXPRESS_CACHE_DIRECTORY` -> `~/.cache/ngc`.
 
 See [`CLI.md`](CLI.md) for full CLI usage documentation.
 
@@ -185,6 +187,23 @@ kubectl create secret generic hf-token-secret \
   -n ${NAMESPACE}
 ```
 
+### NGC API Key
+
+To download models from NVIDIA NGC, set an NGC API key. The server resolves it in this order:
+
+1. `NGC_API_KEY` environment variable
+2. `NGC_CLI_API_KEY` environment variable
+3. `~/.ngc/config` (written by `ngc config set`)
+
+```bash
+export NGC_API_KEY=your_ngc_api_key
+kubectl create secret generic ngc-api-key-secret \
+  --from-literal=NGC_API_KEY=${NGC_API_KEY} \
+  -n ${NAMESPACE}
+```
+
+Pass it to the server pod via `envFrom` or individual `env` entries in your deployment manifest.
+
 ### Helm Chart
 
 The `helm/` directory provides a full Helm chart with configurable replicas, PVC, ingress, and resource limits.

helm/values-local-storage.yaml

@@ -78,6 +78,20 @@ extraEnv:
     valueFrom:
       fieldRef:
         fieldPath: metadata.namespace
+  # Uncomment to enable HuggingFace downloads (private/gated models):
+  #   kubectl create secret generic hf-token-secret --from-literal=HF_TOKEN=<your-token>
+  # - name: HF_TOKEN
+  #   valueFrom:
+  #     secretKeyRef:
+  #       name: hf-token-secret
+  #       key: HF_TOKEN
+  # Uncomment to enable NGC model downloads:
+  #   kubectl create secret generic ngc-api-key-secret --from-literal=NGC_API_KEY=<your-key>
+  # - name: NGC_API_KEY
+  #   valueFrom:
+  #     secretKeyRef:
+  #       name: ngc-api-key-secret
+  #       key: NGC_API_KEY
 
 livenessProbe:
   enabled: true

helm/values.yaml

@@ -107,6 +107,21 @@ extraVolumeMounts: []
 # Additional volumes
 extraVolumes: []
 # Additional environment variables
+# To enable HuggingFace model downloads (private/gated models), create a secret and uncomment:
+#   kubectl create secret generic hf-token-secret --from-literal=HF_TOKEN=<your-token>
+# To enable NGC model downloads, create a secret and uncomment:
+#   kubectl create secret generic ngc-api-key-secret --from-literal=NGC_API_KEY=<your-key>
+# extraEnv:
+#   - name: HF_TOKEN
+#     valueFrom:
+#       secretKeyRef:
+#         name: hf-token-secret
+#         key: HF_TOKEN
+#   - name: NGC_API_KEY
+#     valueFrom:
+#       secretKeyRef:
+#         name: ngc-api-key-secret
+#         key: NGC_API_KEY
 extraEnv: []
 # Additional container arguments
 args: []

modelexpress_common/Cargo.toml

@@ -30,7 +30,9 @@ async-trait = { workspace = true }
 chrono = { workspace = true }
 clap = { workspace = true }
 config = { workspace = true }
+futures = { workspace = true }
 jiff = { workspace = true }
+reqwest = { workspace = true }
 
 [dev-dependencies]
 mockall = "0.13"

modelexpress_common/proto/model.proto

@@ -43,6 +43,7 @@ enum ModelStatus {
 // Enum representing model providers
 enum ModelProvider {
   HUGGING_FACE = 0;
+  NGC = 1;
 }
 
 // Request for streaming model files

modelexpress_common/src/cache.rs

@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    Utils, constants, models::ModelProvider, providers::huggingface::HuggingFaceProviderCache,
+    Utils, constants,
+    models::ModelProvider,
+    providers::{huggingface::HuggingFaceProviderCache, ngc::NgcProviderCache},
 };
 use anyhow::{Context, Result};
 use serde::{Deserialize, Serialize};
@@ -228,8 +230,9 @@ impl CacheConfig {
             });
         }
 
-        let provider = ModelProvider::HuggingFace;
-        models.extend(cache_for_provider(provider).list_models(&self.local_path)?);
+        for provider in [ModelProvider::HuggingFace, ModelProvider::Ngc] {
+            models.extend(cache_for_provider(provider).list_models(&self.local_path)?);
+        }
 
         models.sort_by(|left, right| {
             provider_sort_key(left.provider)
@@ -334,6 +337,7 @@ pub(crate) trait ProviderCache: Send + Sync {
 pub(crate) fn cache_for_provider(provider: ModelProvider) -> &'static dyn ProviderCache {
     match provider {
         ModelProvider::HuggingFace => &HuggingFaceProviderCache,
+        ModelProvider::Ngc => &NgcProviderCache,
     }
 }
 
@@ -366,6 +370,7 @@ pub(crate) fn directory_size(path: &Path) -> Result<u64> {
 fn provider_sort_key(provider: ModelProvider) -> u8 {
     match provider {
         ModelProvider::HuggingFace => 0,
+        ModelProvider::Ngc => 1,
     }
 }
 

modelexpress_common/src/download.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::models::ModelProvider;
-use crate::providers::{HuggingFaceProvider, ModelProviderTrait};
+use crate::providers::{HuggingFaceProvider, ModelProviderTrait, NgcProvider};
 use anyhow::Result;
 use std::path::PathBuf;
 use tracing::{info, warn};
@@ -12,6 +12,7 @@ use tracing::{info, warn};
 pub fn get_provider(provider: ModelProvider) -> Box<dyn ModelProviderTrait> {
     match provider {
         ModelProvider::HuggingFace => Box::new(HuggingFaceProvider),
+        ModelProvider::Ngc => Box::new(NgcProvider),
     }
 }
 

modelexpress_common/src/lib.rs

@@ -141,6 +141,7 @@ impl From<models::ModelProvider> for grpc::model::ModelProvider {
     fn from(provider: models::ModelProvider) -> Self {
         match provider {
             models::ModelProvider::HuggingFace => grpc::model::ModelProvider::HuggingFace,
+            models::ModelProvider::Ngc => grpc::model::ModelProvider::Ngc,
         }
     }
 }
@@ -149,6 +150,7 @@ impl From<grpc::model::ModelProvider> for models::ModelProvider {
     fn from(provider: grpc::model::ModelProvider) -> Self {
         match provider {
             grpc::model::ModelProvider::HuggingFace => models::ModelProvider::HuggingFace,
+            grpc::model::ModelProvider::Ngc => models::ModelProvider::Ngc,
         }
     }
 }

modelexpress_common/src/models.rs

@@ -30,13 +30,16 @@ pub enum ModelProvider {
     /// Hugging Face model hub
     #[default]
     HuggingFace,
+    /// NVIDIA NGC catalog
+    Ngc,
 }
 
 impl ModelProvider {
     #[must_use]
     pub const fn as_str(self) -> &'static str {
         match self {
             Self::HuggingFace => "hugging-face",
+            Self::Ngc => "ngc",
         }
     }
 }
@@ -49,7 +52,7 @@ impl Display for ModelProvider {
 
 impl ValueEnum for ModelProvider {
     fn value_variants<'a>() -> &'a [Self] {
-        &[Self::HuggingFace]
+        &[Self::HuggingFace, Self::Ngc]
     }
 
     fn to_possible_value(&self) -> Option<PossibleValue> {
@@ -98,15 +101,26 @@ mod tests {
     #[test]
     fn test_model_provider_display() {
         assert_eq!(ModelProvider::HuggingFace.to_string(), "hugging-face");
+        assert_eq!(ModelProvider::Ngc.to_string(), "ngc");
     }
 
     #[test]
     fn test_model_provider_value_enum_matches_display() {
-        let provider = ModelProvider::HuggingFace;
-        let parsed = ModelProvider::from_str(provider.as_str(), false)
-            .expect("Failed to parse ModelProvider from clap value");
+        for provider in [ModelProvider::HuggingFace, ModelProvider::Ngc] {
+            let parsed = ModelProvider::from_str(provider.as_str(), false)
+                .expect("Failed to parse ModelProvider from clap value");
+            assert_eq!(parsed, provider);
+        }
+    }
 
-        assert_eq!(parsed, provider);
+    #[test]
+    fn test_model_provider_ngc_serialization() {
+        let provider = ModelProvider::Ngc;
+        let serialized =
+            serde_json::to_string(&provider).expect("Failed to serialize ModelProvider");
+        let deserialized: ModelProvider =
+            serde_json::from_str(&serialized).expect("Failed to deserialize ModelProvider");
+        assert_eq!(provider, deserialized);
     }
 
     #[test]

modelexpress_common/src/providers.rs

@@ -77,8 +77,10 @@ pub trait ModelProviderTrait: Send + Sync {
 }
 
 pub mod huggingface;
+pub mod ngc;
 
 pub use huggingface::HuggingFaceProvider;
+pub use ngc::NgcProvider;
 
 #[cfg(test)]
 mod tests {