Skip to content

feat: add GCS model provider#187

Closed
andrewpaprotsky wants to merge 7 commits intomainfrom
apaprotskyi/gcs
Closed

feat: add GCS model provider#187
andrewpaprotsky wants to merge 7 commits intomainfrom
apaprotskyi/gcs

Conversation

@andrewpaprotsky
Copy link
Copy Markdown
Contributor

@andrewpaprotsky andrewpaprotsky commented Mar 27, 2026

  • Add first-class GCS support across common, client, and server layers, including proto/provider enums, DB mapping, cache stats/listing/clear, and provider-aware request routing.
  • Introduce modelexpress_common::providers::gcs with strict gs://bucket/path canonicalization, safe cache layout validation, ancestor/descendant overlap protection, and resumable downloads with CRC32C verification.
  • Store GCS cache metadata in .mx/manifest.json and use that manifest for cache discovery, listing, and download planning.
  • Canonicalize model names at request boundaries and keep the server as the source of truth for Client::request_model, with direct download used only when the initial client connection cannot be established.
  • Update dependency wiring for Google Cloud Storage support and align tests with the final provider-aware client and cache behavior.

BREAKING CHANGE:

  • Client::request_model now requires an explicit provider argument.
  • Removed public APIs: preload_model_to_cache, request_model_with_provider_and_fallback, request_model_server_only, and download_model_directly.
  • Renamed request_model_with_provider to request_model_on_server.

Summary by CodeRabbit

Release Notes

  • New Features

    • Added Google Cloud Storage (GCS) support as a model provider alongside Hugging Face
    • Extended caching system to discover, manage, and cache models from GCS
  • Improvements

    • Enhanced error reporting with clearer transport error messages
    • Improved model download workflows with canonical model-name handling

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented Apr 1, 2026

Walkthrough

This pull request extends the ModelExpress ecosystem with Google Cloud Storage (GCS) support and refactors the client API to enforce provider-aware model requests. It removes cached-preload and provider-fallback methods from the public API, replaces direct download entrypoints with a unified GCS provider implementation, and updates client/server layers to canonicalize and validate model names per provider specifications.

Changes

Cohort / File(s) Summary
CLI & Test Updates
modelexpress_client/src/bin/cli.rs, modelexpress_client/src/bin/fallback_test.rs, modelexpress_client/src/bin/test_client.rs
Added ValueEnum import and test for provider parsing; updated test data to use GCS provider and paths; refactored smart-fallback logic to use ClientConfig::for_testing() and simplified result handling.
Handler Refactoring
modelexpress_client/src/bin/modules/handlers.rs
Updated download_model to apply cache config directly for smart-fallback, switch server-only to new request_model API, and replace direct download with modelexpress_common::download::download_model with explicit cache paths and error wrapping.
Client API Core
modelexpress_client/src/lib.rs
Removed public methods (preload_model_to_cache, request_model_with_provider_and_fallback, request_model_with_provider, request_model_server_only, download_model_directly); added request_model_on_server; changed request_model signature to require explicit provider parameter; updated smart-fallback and streaming logic with new provider-aware flow.
Model & Provider Definitions
modelexpress_common/src/models.rs, modelexpress_common/proto/model.proto
Added ModelProvider::Gcs enum variant; updated as_str() and clap::ValueEnum support to include GCS; extended serialization tests for both HuggingFace and GCS variants.
Provider Infrastructure
modelexpress_common/src/providers.rs
Added canonical_model_name trait method with default implementation; extended is_weight_file to recognize .iop and .gas extensions; added GCS provider module and re-export.
GCS Provider Implementation
modelexpress_common/src/providers/gcs.rs
New 2100\+ LOC module implementing Google Cloud Storage provider with model name parsing/canonicalization, cached directory layout management, concurrent parallel downloads with CRC32C verification, temp-file promotion, cache discovery, and deletion logic; implements ModelProviderTrait and ProviderCache.
Common Download & Caching
modelexpress_common/src/download.rs, modelexpress_common/src/cache.rs
Extended provider factory to construct GcsProvider; added public canonical_model_name API; updated cache stats and clearing logic to handle both HuggingFace and GCS models with provider-specific layout resolution.
Error Handling & Conversions
modelexpress_common/src/lib.rs
Changed Transport error variant from storing tonic::transport::Error to storing formatted String; added format_error_chain utility; implemented From<tonic::transport::Error> conversion; extended gRPC↔model provider conversions to include Gcs.
Server Provider Support
modelexpress_server/src/database.rs, modelexpress_server/src/services.rs
Extended database provider mappings to recognize ModelProvider::Gcs in record decoding/serialization; updated ensure_model_downloaded, stream_model_files, and list_model_files to canonicalize model names; added zero-byte file handling in streaming; enforced cache-path correctness checks; added GCS canonicalization tests.
Dependencies & Build
modelexpress_common/Cargo.toml
Added runtime dependencies: google-cloud-storage, rustls (with ring and std features), crc32c, and futures.
Integration Tests
workspace-tests/tests/integration_tests.rs
Replaced Client::download_model_directly calls with local helper invoking modelexpress_common::download::download_model directly with explicit cache paths and error wrapping.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Poem

🐰 Hops with glee through cloud and cache,
New GCS paths we boldly hash!
Providers now sing in harmony—
Face and Cloud dance wild and free. 🌥️✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 65.22% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Title check ✅ Passed The title 'feat: add GCS model provider' directly and clearly describes the primary change - adding Google Cloud Storage as a new model provider across the codebase.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ 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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
modelexpress_client/src/bin/modules/handlers.rs (1)

188-198: ⚠️ Potential issue | 🟠 Major

Keep DownloadStrategy::ServerOnly on the server-only API.

Client::request_model() now streams files into the local cache when shared_storage is false, so this branch no longer matches the selected strategy. A ServerOnly download can unexpectedly pull the full model to the client.

Proposed fix
         DownloadStrategy::ServerOnly => {
             debug!("Using server-only strategy");
             let mut client = if let Some(cache_config) = cache_config {
                 Client::new_with_cache(config.clone(), cache_config).await?
             } else {
                 Client::new(config.clone()).await?
             };
             client
-                .request_model(&model_name, provider, false)
-                .await
-                .map(|_| ())
+                .request_model_on_server(&model_name, provider, false)
+                .await
         }
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@modelexpress_client/src/bin/modules/handlers.rs` around lines 188 - 198, The
ServerOnly branch is calling Client::request_model(...) which can stream files
to local cache when shared_storage is false and thus violates
DownloadStrategy::ServerOnly; replace that call with the server-only API (or a
dedicated method) that requests the model to remain on the server instead of
downloading it locally — locate the DownloadStrategy::ServerOnly branch and
replace the Client::request_model(&model_name, provider, false).await call with
the appropriate server-only request (e.g., a Client::request_model_server_only
or other RPC that only triggers server-side provisioning) so no files are
streamed to the client cache.
🧹 Nitpick comments (2)
modelexpress_client/src/bin/cli.rs (1)

212-244: Test name doesn't fully reflect test behavior.

The test test_cli_model_clear_defaults_to_hugging_face_provider now also tests explicit --provider gcs parsing (lines 219-221) before testing the default behavior. While the test logic is correct, the name only describes part of what's being tested.

Consider renaming to test_cli_model_clear_provider_parsing_and_default or splitting into two tests for clarity.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@modelexpress_client/src/bin/cli.rs` around lines 212 - 244, Rename or split
the test currently named test_cli_model_clear_defaults_to_hugging_face_provider
because it also verifies explicit provider parsing; either (A) rename the
function to something like test_cli_model_clear_provider_parsing_and_default to
reflect both checks, or (B) split into two tests: one that exercises
Cli::try_parse_from(["modelexpress-cli","model","clear","--provider","gcs",...])
to assert explicit provider parsing, and a second that uses
Cli::try_parse_from(["modelexpress-cli","model","clear","gs://..."]) to assert
default ModelProvider::HuggingFace for
super::modules::args::ModelCommands::Clear by inspecting Commands::Model {
command } and the provider/model_name fields.
modelexpress_common/src/download.rs (1)

171-175: Redundant assertions duplicated from test_get_provider.

Lines 172-174 repeat the same get_provider + provider_name() assertions already covered in test_get_provider (lines 106-110). Consider removing the duplicate assertions to keep test_download_model_routing focused on its original purpose.

Suggested cleanup
     #[tokio::test]
     async fn test_download_model_routing() {
         // Test that download_model function properly routes to the provider
         // Note: This test doesn't actually download from HF to avoid network dependency
         // In a real scenario, you might want to mock the hf-hub dependency
 
         let provider = ModelProvider::HuggingFace;
         let provider_impl = get_provider(provider);
         assert_eq!(provider_impl.provider_name(), "Hugging Face");
-
-        let provider = ModelProvider::Gcs;
-        let provider_impl = get_provider(provider);
-        assert_eq!(provider_impl.provider_name(), "GCS");
     }
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@modelexpress_common/src/download.rs` around lines 171 - 175, Remove the
redundant assertions in test_download_model_routing that duplicate
test_get_provider: delete the get_provider call and assert_eq! on
provider_name() for ModelProvider::Gcs (the lines invoking get_provider,
provider_impl, and provider_impl.provider_name()) so the test focuses only on
download routing logic; the existing test_get_provider already covers
get_provider and provider_name behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@modelexpress_client/src/lib.rs`:
- Around line 581-593: The streaming and direct-fallback branches use
cache_config.local_path (or config.cache.local_path) directly instead of the
provider-aware helper used by Client::get_cache_dir(), causing HF models
resolved via HF_HUB_CACHE to be missed; update the code paths that call
stream_model_files_from_server(&model_name, provider) and the direct fallback
(also at the other occurrence) to first resolve the target cache directory
through the same helper function (get_cache_dir or the provider-aware variant
used by get_model_path) and pass that resolved path into the streaming/download
logic so both branches consult the identical provider-aware cache directory
resolution (ensure you update calls around needs_streaming,
stream_model_files_from_server, and the direct fallback near the mentioned
block).

In `@modelexpress_common/Cargo.toml`:
- Around line 34-36: The dependencies google-cloud-storage, rustls, and crc32c
are pinned in this sub-crate; add them to the root Cargo.toml under
[workspace.dependencies] (use cargo add to ensure exact versions and feature
flags are recorded) and then update this file's entries for
google-cloud-storage, rustls, and crc32c to reference the workspace (use {
workspace = true } and keep any needed features configuration here if allowed).
Ensure rustls keeps the features ["ring","std"] at the workspace level if
required and remove direct version pins from modelexpress_common's Cargo.toml so
all three use the workspace definitions.

In `@modelexpress_common/src/providers/gcs.rs`:
- Around line 1238-1273: Wrap the whole GCS-per-model prepare/download sequence
in an OS-backed exclusive per-model lock acquired before calling
current_model_dir.has_model_marker() and held until after the marker is written
/ cleanup completes (i.e., spanning ensure_available(), DownloadContext::new,
the with_prepared() closure that invokes Downloader::new and
download_matching_objects(), and the subsequent marker write); implement this as
a per-model claim file (with advisory flock or native file lock) named from the
model/context.model_dir, use an RAII guard to automatically release the lock on
success or error, and surface any lock acquisition errors (with optional
timeout) so concurrent callers serialize access and cannot race during marker
checks or directory cleanup.

---

Outside diff comments:
In `@modelexpress_client/src/bin/modules/handlers.rs`:
- Around line 188-198: The ServerOnly branch is calling
Client::request_model(...) which can stream files to local cache when
shared_storage is false and thus violates DownloadStrategy::ServerOnly; replace
that call with the server-only API (or a dedicated method) that requests the
model to remain on the server instead of downloading it locally — locate the
DownloadStrategy::ServerOnly branch and replace the
Client::request_model(&model_name, provider, false).await call with the
appropriate server-only request (e.g., a Client::request_model_server_only or
other RPC that only triggers server-side provisioning) so no files are streamed
to the client cache.

---

Nitpick comments:
In `@modelexpress_client/src/bin/cli.rs`:
- Around line 212-244: Rename or split the test currently named
test_cli_model_clear_defaults_to_hugging_face_provider because it also verifies
explicit provider parsing; either (A) rename the function to something like
test_cli_model_clear_provider_parsing_and_default to reflect both checks, or (B)
split into two tests: one that exercises
Cli::try_parse_from(["modelexpress-cli","model","clear","--provider","gcs",...])
to assert explicit provider parsing, and a second that uses
Cli::try_parse_from(["modelexpress-cli","model","clear","gs://..."]) to assert
default ModelProvider::HuggingFace for
super::modules::args::ModelCommands::Clear by inspecting Commands::Model {
command } and the provider/model_name fields.

In `@modelexpress_common/src/download.rs`:
- Around line 171-175: Remove the redundant assertions in
test_download_model_routing that duplicate test_get_provider: delete the
get_provider call and assert_eq! on provider_name() for ModelProvider::Gcs (the
lines invoking get_provider, provider_impl, and provider_impl.provider_name())
so the test focuses only on download routing logic; the existing
test_get_provider already covers get_provider and provider_name behavior.
🪄 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: 14302338-a141-4e2c-ae03-606fde320329

📥 Commits

Reviewing files that changed from the base of the PR and between 125d322 and 798f519.

⛔ Files ignored due to path filters (1)
  • Cargo.lock is excluded by !**/*.lock
📒 Files selected for processing (16)
  • modelexpress_client/src/bin/cli.rs
  • modelexpress_client/src/bin/fallback_test.rs
  • modelexpress_client/src/bin/modules/handlers.rs
  • modelexpress_client/src/bin/test_client.rs
  • modelexpress_client/src/lib.rs
  • modelexpress_common/Cargo.toml
  • modelexpress_common/proto/model.proto
  • modelexpress_common/src/cache.rs
  • modelexpress_common/src/download.rs
  • modelexpress_common/src/lib.rs
  • modelexpress_common/src/models.rs
  • modelexpress_common/src/providers.rs
  • modelexpress_common/src/providers/gcs.rs
  • modelexpress_server/src/database.rs
  • modelexpress_server/src/services.rs
  • workspace-tests/tests/integration_tests.rs

Comment thread modelexpress_client/src/lib.rs
Comment thread modelexpress_common/Cargo.toml Outdated
Comment thread modelexpress_common/src/providers/gcs.rs Outdated
@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 3, 2026

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds first-class Google Cloud Storage (GCS) support as a model provider across the common, client, and server layers, including provider-aware canonicalization, cache layout/validation, and download/streaming behavior.

Changes:

  • Introduce a new GCS provider implementation with gs:// canonicalization, safe cache layout validation, manifest-based cache discovery, and parallel/resumable downloads with CRC32C verification.
  • Extend shared model/provider enums (Rust + proto) and propagate provider-aware handling through server services, client APIs, and cache stats/listing/clear flows.
  • Refactor client request/download APIs to make provider explicit and tighten “server as source of truth” behavior with a limited direct-download fallback.

Reviewed changes

Copilot reviewed 12 out of 18 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
workspace-tests/tests/integration_tests.rs Updates integration tests to use common direct-download path and provider-aware APIs.
modelexpress_server/src/services.rs Canonicalizes model names at request boundaries; validates cache layout before streaming; handles zero-byte files in streaming; adds targeted tests.
modelexpress_server/src/database.rs Adds DB mapping support for the new Gcs provider with tests.
modelexpress_common/src/providers/gcs.rs New GCS provider implementation: URL parsing/canonicalization, cache layout rules, manifest handling, overlap protection, parallel/resumable download + CRC32C checks, and extensive tests.
modelexpress_common/src/providers.rs Adds provider-level canonicalization hook; exports GCS provider; extends weight-file detection.
modelexpress_common/src/models.rs Adds ModelProvider::Gcs and updates clap/serde tests.
modelexpress_common/src/lib.rs Improves transport error formatting by capturing error chains; adds provider enum conversions for GCS.
modelexpress_common/src/download.rs Routes provider selection to include GCS; adds provider-based canonicalization helper + tests.
modelexpress_common/src/cache.rs Extends cache stats/listing/clear + path resolution to include GCS layout; adds tests.
modelexpress_common/proto/model.proto Adds GCS provider enum value to the API.
modelexpress_common/Cargo.toml Wires in GCS-related dependencies for the common crate.
modelexpress_client/src/lib.rs Makes request_model provider-aware; refactors fallback behavior; updates streaming logic and adds tests for GCS behavior.
modelexpress_client/src/bin/test_client.rs Updates CLI test client to new provider-explicit APIs and direct download path.
modelexpress_client/src/bin/modules/handlers.rs Updates CLI handlers to use provider-aware client APIs and common direct download.
modelexpress_client/src/bin/fallback_test.rs Updates fallback test binary to new smart-fallback semantics.
modelexpress_client/src/bin/cli.rs Updates CLI parsing tests for provider enum values.
Cargo.toml Adds workspace deps needed for GCS integration.
Cargo.lock Updates lockfile to include new dependencies and transitive changes.

Comment thread modelexpress_common/Cargo.toml
Comment thread modelexpress_client/src/lib.rs
Comment thread modelexpress_server/src/services.rs Outdated
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 12 out of 18 changed files in this pull request and generated 2 comments.

Comment thread modelexpress_common/src/providers/gcs.rs
Comment thread workspace-tests/tests/integration_tests.rs
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 12 out of 18 changed files in this pull request and generated 4 comments.

Comment thread modelexpress_server/src/services.rs Outdated
Comment thread modelexpress_common/src/providers/gcs.rs Outdated
Comment thread modelexpress_common/src/providers/gcs.rs Outdated
Comment thread modelexpress_server/src/services.rs Outdated
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 12 out of 18 changed files in this pull request and generated 3 comments.

Comment thread modelexpress_common/src/providers/gcs.rs
Comment thread modelexpress_common/src/providers/gcs.rs
Comment thread modelexpress_server/src/services.rs
- Add first-class GCS support across common, client, and server layers,
  including proto/provider enums, DB mapping, cache stats/listing/clear, and
  provider-aware request routing.
- Introduce `modelexpress_common::providers::gcs` with strict
  `gs://bucket/path` canonicalization, safe cache layout validation,
  ancestor/descendant overlap protection, and resumable downloads with CRC32C
  verification.
- Store GCS cache metadata in `.mx/manifest.json` and use that manifest for
  cache discovery, listing, and download planning instead of marker-file
  metadata.
- Canonicalize model names at request boundaries and keep the server as the
  source of truth for `Client::request_model`, with direct download used only
  when the initial client connection cannot be established.
- Update dependency wiring for Google Cloud Storage support and align tests
  with the final provider-aware client and cache behavior.

BREAKING CHANGE:

- `Client::request_model` now requires an explicit `provider` argument.
- Removed public APIs: `preload_model_to_cache`,
  `request_model_with_provider_and_fallback`, `request_model_server_only`, and
  `download_model_directly`.
- Renamed `request_model_with_provider` to `request_model_on_server`.

Signed-off-by: Andrew Paprotsky <apaprotskyi@nvidia.com>
Signed-off-by: Andrew Paprotsky <apaprotskyi@nvidia.com>
Signed-off-by: Andrew Paprotsky <apaprotskyi@nvidia.com>
Signed-off-by: Andrew Paprotsky <apaprotskyi@nvidia.com>
Signed-off-by: Andrew Paprotsky <apaprotskyi@nvidia.com>
Signed-off-by: Andrew Paprotsky <apaprotskyi@nvidia.com>
Signed-off-by: Andrew Paprotsky <apaprotskyi@nvidia.com>
@andrewpaprotsky
Copy link
Copy Markdown
Contributor Author

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants