feat(refit): add immutable revision catalog service for S3 delta weight sync - #610
feat(refit): add immutable revision catalog service for S3 delta weight sync#610nv-hwoo wants to merge 3 commits into
Conversation
Signed-off-by: Hyunjae Woo <hwoo@nvidia.com>
WalkthroughThis change adds a revision protobuf API, manifest models and validation, Redis and in-memory catalog backends, gRPC server integration, generated bindings, and a Python client with lifecycle tests. ChangesRevision catalog
Estimated code review effort: 5 (Critical) | ~120 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (5)
modelexpress_server/src/revision/service.rs (1)
62-76: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe
createdflag is discarded.
publishreturnsPublicationResult { record, created }, and this handler returns onlyrecord. A publisher cannot distinguish a first publication from an idempotent replay. If clients need that signal, expose it in the response message; otherwise this is fine as designed.🤖 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/revision/service.rs` around lines 62 - 76, Review the publish_revision handler and expose PublicationResult.created in the PublishRevision response contract so clients can distinguish new publications from idempotent replays. Update the protobuf response message and generated/server usage as needed, while preserving the existing RevisionRecord payload and publication behavior.modelexpress_server/src/revision/state.rs (2)
173-207: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a test for the invalid lifecycle path.
The tests cover publish idempotence, conflict, commit, idempotent commit, and not-found.
CatalogError::InvalidLifecyclehas no coverage. A test that stores a record with an unrecognized state value would close that gap.🤖 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/revision/state.rs` around lines 173 - 207, Add a test alongside immutable_publication_is_idempotent_and_rejects_conflicts and commit_is_an_idempotent_ready_to_committed_transition that inserts or stores a revision record with an unrecognized state value, invokes the relevant lifecycle operation, and asserts it returns CatalogError::InvalidLifecycle.
62-67: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove
RevisionCatalogState::connectif no external API requires it. The factory connects the backend before state construction, and no repository code calls the state-level method.🤖 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/revision/state.rs` around lines 62 - 67, Remove the unused RevisionCatalogState::connect method, since the factory already connects the backend before constructing the state and no repository code invokes this state-level API.modelexpress_server/src/revision/backend/redis.rs (1)
42-47: 🗄️ Data Integrity & Integration | 🔵 TrivialConsider a retention policy for revision keys.
revision_keycreates one Redis hash per model and target version. Nothing in this file removes or expires those keys, so the catalog grows without bound as versions accumulate. Decide whether a TTL, an explicit delete path, or an external compaction job owns cleanup.🤖 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/revision/backend/redis.rs` around lines 42 - 47, Define and implement a cleanup policy for keys generated by revision_key, such as applying a TTL when they are created, deleting obsolete revisions through an explicit path, or documenting and wiring an external compaction owner. Ensure the selected policy prevents unbounded Redis growth while preserving access to active revision data.modelexpress_server/src/revision/backend.rs (1)
53-79: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReport the actual reason when the memory backend is not available.
If
memory-backendis enabled withoutintegration-tests, the factory returns "revision catalog currently supports only the Redis backend". That message hides the real cause. State that the in-memory catalog needs theintegration-testsfeature.Proposed message change
#[cfg(not(feature = "integration-tests"))] { - Err("revision catalog currently supports only the Redis backend".into()) + Err("in-memory revision catalog requires the 'integration-tests' feature".into()) }🤖 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/revision/backend.rs` around lines 53 - 79, Update the cfg(not(feature = "integration-tests")) branch of create_revision_catalog_backend for BackendConfig::Memory to return an error stating that the in-memory revision catalog requires the integration-tests feature, rather than claiming only Redis is supported.
🤖 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_client/python/modelexpress/refit/catalog.py`:
- Around line 59-80: Add a configurable finite timeout to the catalog client and
pass it as the timeout keyword argument in publish_revision, get_revision, and
commit_revision when invoking the corresponding stub RPCs. Update FakeStub to
record RPC keyword arguments, and extend the tests to assert that each call
receives the configured timeout.
- Around line 42-45: Update the channel initialization around
RevisionCatalogServiceStub to preserve TLS for https:// endpoints: use
grpc.secure_channel with appropriate TLS credentials, or explicitly reject HTTPS
when only plaintext host:port targets are supported. Keep grpc.insecure_channel
for http:// endpoints and add coverage verifying the transport selection.
In `@modelexpress_client/python/modelexpress/revision_pb2_grpc.py`:
- Line 24: Update generate_proto.sh so the generated revision_pb2_grpc.py does
not contain an f-string without replacement fields, or configure Ruff to exclude
generated bindings from F541; then regenerate the binding and ensure the fix is
reproducible from the generation script rather than editing revision_pb2_grpc.py
directly.
In `@modelexpress_server/src/revision/backend/redis.rs`:
- Around line 166-200: Update commit_revision and COMMIT_LUA so the
compare-and-swap does not depend on current.encode_to_vec() matching stored
record bytes exactly: perform the state transition atomically in Lua using the
stored hash field and compare an independent lifecycle-state field or the bytes
observed within the same script invocation. Preserve NotFound, AlreadyCommitted,
InvalidState, successful commit, and conflict retry outcomes, and add backoff
between retries to avoid immediate repeated round trips under contention.
In `@modelexpress_server/src/revision/backend/testing.rs`:
- Around line 4-5: Update the module-level documentation in testing.rs to state
that TestRevisionCatalogBackend is selected by create_revision_catalog_backend
for BackendConfig::Memory when the integration-tests feature is enabled,
removing the inaccurate claim that it is not selectable.
In `@modelexpress_server/src/server.rs`:
- Around line 122-127: Wrap the await of create_revision_catalog_backend in
tokio::time::timeout with the same 10-second duration used by the registry and
P2P connection paths. Handle timeout and connection errors through the existing
error logging and propagation flow so startup fails fast while preserving the
successful revision_backend initialization.
In `@modelexpress_server/tests/in_process_server.rs`:
- Around line 104-109: Bound the connection-retry loop that calls
RevisionCatalogServiceClient::connect with tokio::time::timeout, or reuse the
file’s existing wait helper if available. Preserve the retry and delay behavior
within the timeout, then fail the test with a clear timeout error instead of
waiting indefinitely when run_server does not become reachable.
---
Nitpick comments:
In `@modelexpress_server/src/revision/backend.rs`:
- Around line 53-79: Update the cfg(not(feature = "integration-tests")) branch
of create_revision_catalog_backend for BackendConfig::Memory to return an error
stating that the in-memory revision catalog requires the integration-tests
feature, rather than claiming only Redis is supported.
In `@modelexpress_server/src/revision/backend/redis.rs`:
- Around line 42-47: Define and implement a cleanup policy for keys generated by
revision_key, such as applying a TTL when they are created, deleting obsolete
revisions through an explicit path, or documenting and wiring an external
compaction owner. Ensure the selected policy prevents unbounded Redis growth
while preserving access to active revision data.
In `@modelexpress_server/src/revision/service.rs`:
- Around line 62-76: Review the publish_revision handler and expose
PublicationResult.created in the PublishRevision response contract so clients
can distinguish new publications from idempotent replays. Update the protobuf
response message and generated/server usage as needed, while preserving the
existing RevisionRecord payload and publication behavior.
In `@modelexpress_server/src/revision/state.rs`:
- Around line 173-207: Add a test alongside
immutable_publication_is_idempotent_and_rejects_conflicts and
commit_is_an_idempotent_ready_to_committed_transition that inserts or stores a
revision record with an unrecognized state value, invokes the relevant lifecycle
operation, and asserts it returns CatalogError::InvalidLifecycle.
- Around line 62-67: Remove the unused RevisionCatalogState::connect method,
since the factory already connects the backend before constructing the state and
no repository code invokes this state-level API.
🪄 Autofix
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: 6b8e2e20-8b47-4326-8926-a59458da1fa6
📒 Files selected for processing (21)
modelexpress_client/python/generate_proto.shmodelexpress_client/python/modelexpress/refit/catalog.pymodelexpress_client/python/modelexpress/refit/manifest.pymodelexpress_client/python/modelexpress/revision_pb2.pymodelexpress_client/python/modelexpress/revision_pb2_grpc.pymodelexpress_client/python/tests/test_revision_catalog.pymodelexpress_client/python/tests/test_revision_manifest.pymodelexpress_client/python/tests/test_revision_proto.pymodelexpress_common/build.rsmodelexpress_common/proto/revision.protomodelexpress_common/src/lib.rsmodelexpress_common/src/revision.rsmodelexpress_server/src/lib.rsmodelexpress_server/src/revision.rsmodelexpress_server/src/revision/backend.rsmodelexpress_server/src/revision/backend/redis.rsmodelexpress_server/src/revision/backend/testing.rsmodelexpress_server/src/revision/service.rsmodelexpress_server/src/revision/state.rsmodelexpress_server/src/server.rsmodelexpress_server/tests/in_process_server.rs
Signed-off-by: Hyunjae Woo <hwoo@nvidia.com>
|
Addressed the actionable CodeRabbit findings in Additional nitpick triage:
CI startup root cause and fix: Kubernetes P2P jobs use the Kubernetes metadata backend, but the new Redis-only revision factory was called unconditionally. The revision service is now registered only for supported backends; existing Kubernetes registry/P2P server behavior remains available. Validation includes workspace Rust tests, integration-feature server tests, two real-Redis lifecycle regressions, affected Python revision tests, Clippy/fmt/cargo-check, generated-binding reproduction with grpcio-tools 1.66.2, changed-file Ruff/pre-commit, and secret-pattern scan. |
Signed-off-by: Hyunjae Woo <hwoo@nvidia.com>
|
Follow-up review found and fixed one Redis lifecycle blocker in The first fix made a separate Redis The corrected storage contract is now:
A real-Redis regression now asserts exact stored bytes are unchanged after commit (including an appended unknown protobuf field), a subsequent read reports COMMITTED, and publish replay after commit also reports COMMITTED. Both live-Redis lifecycle tests pass. Additional validation after this correction: full Rust workspace tests, integration-feature server suite, affected Python revision tests (20 passed), Clippy with all targets/features and warnings denied, cargo fmt/check, changed-file pre-commit, diff check, and secret-pattern scan. |
Summary
Add an immutable revision catalog service for coordinating live model-weight updates. The catalog records revision identity and lifecycle state, while weight payloads remain in external object storage. It provides the control-plane contract used by the publisher and receiver in the following PRs.
This is a parallel effort along the current refit API work (#604). A separate follow-up PRs will merge this into the standard refit API.
Changes
RevisionCatalogServicegRPC API:PublishRevisionGetRevisionCommitRevisionREADY -> COMMITTEDrevision lifecycle.Payload bytes do not pass through the catalog. The service stores only immutable revision metadata and lifecycle state.
Correctness
Validation
pre-commit run --from-ref origin/main --to-ref hwoo/mx-delta-revision-prcargo fmtcargo clippycargo checkSummary by CodeRabbit
New Features
Bug Fixes