feat: Add TransferEngine backend to P2P metadata#157
Conversation
WalkthroughThe changes introduce support for TransferEngine session IDs as an alternative to NIXL metadata in WorkerMetadata. A new oneof structure is added to protobuf definitions, a BackendMetadataRecord enum is introduced to represent mutually exclusive metadata options, and all conversion layers are updated to handle the new field. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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 |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
modelexpress_server/src/state.rs (1)
239-266: Add a focused round-trip test for the TransferEngine branch.This only pins the existing NIXL path. Since this PR adds
transfer_engine_session_id, add aWorkerMetadata -> WorkerRecord -> WorkerMetadatacase for that variant here so the backwards-compat wrapper catches regressions without relying on E2E coverage.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelexpress_server/src/state.rs` around lines 239 - 266, Update the test_worker_record_conversion unit test to include and verify the new TransferEngine branch: construct a WorkerMetadata that sets backend_metadata to the TransferEngine variant (with a transfer_engine_session_id value), convert it to WorkerRecord via WorkerRecord::from(...) and back into WorkerMetadata via Into, and assert that worker_rank, backend_metadata (specifically the TransferEngine session id), and tensors round-trip correctly; reference the existing test_worker_record_conversion, WorkerMetadata, WorkerRecord, and BackendMetadataRecord symbols when adding the new case so the backwards-compat wrapper will catch regressions.modelexpress_server/src/metadata_backend/kubernetes.rs (1)
263-269: The flattened CRD payload loses the oneof discriminator for empty NIXL blobs.
BackendMetadataRecord::Nixl(vec![])andBackendMetadataRecord::Noneboth serialize here asnixl_metadata = ""andtransfer_engine_session_id = None, so the read side cannot reconstruct which variant was originally set. If empty NIXL payloads are valid, persist the backend kind explicitly instead of inferring it from the flattened fields.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@modelexpress_server/src/metadata_backend/kubernetes.rs` around lines 263 - 269, The serialization flattener in kubernetes.rs (the match assigning nixl_metadata and transfer_engine_session_id) loses the discriminator because Nixl(vec![]) and None both become nixl_metadata = "" and transfer_engine_session_id = None; modify this match to also emit an explicit backend kind/discriminator (e.g., backend_kind = "Nixl" | "TransferEngine" | "None") when constructing the flattened CRD payload so the reader can reconstruct the original BackendMetadataRecord; set backend_kind alongside nixl_metadata and transfer_engine_session_id in the same match arm and ensure the reader/deserialize path consumes that discriminator to restore the correct enum variant.
🤖 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_server/src/k8s_types.rs`:
- Around line 94-96: The WorkerStatus field transfer_engine_session_id
(serialized as transferEngineSessionId) was added in Rust but is missing from
the CRD OpenAPI schema; update the workers array item schema in
crd-modelmetadata.yaml to add a transferEngineSessionId property (string,
optional/nullable) under the WorkerStatus properties so the CRD matches the
WorkerStatus struct, include an appropriate description, and do not mark it as
required.
In `@modelexpress_server/src/metadata_backend/kubernetes.rs`:
- Around line 263-274: WorkerStatus.nixl_metadata is currently serialized as
base64 but deserialization (where ModelMetadata.status.workers[].nixlMetadata is
read) assumes base64 and will fail for previously persisted non-base64 entries;
add a compatibility deserializer: annotate the nixl_metadata field with
serde(deserialize_with = "deserialize_nixl_metadata"), implement a
deserialize_nixl_metadata function that tries base64::decode and deserializes
the decoded bytes, and if base64 decoding fails or the decoded bytes are
invalid, fall back to interpreting the input as a raw UTF-8/JSON string (or
already-encoded form) and produce the same internal string representation;
ensure the deserializer returns the same type used by WorkerStatus.nixl_metadata
and include unit tests covering both legacy and base64 formats.
---
Nitpick comments:
In `@modelexpress_server/src/metadata_backend/kubernetes.rs`:
- Around line 263-269: The serialization flattener in kubernetes.rs (the match
assigning nixl_metadata and transfer_engine_session_id) loses the discriminator
because Nixl(vec![]) and None both become nixl_metadata = "" and
transfer_engine_session_id = None; modify this match to also emit an explicit
backend kind/discriminator (e.g., backend_kind = "Nixl" | "TransferEngine" |
"None") when constructing the flattened CRD payload so the reader can
reconstruct the original BackendMetadataRecord; set backend_kind alongside
nixl_metadata and transfer_engine_session_id in the same match arm and ensure
the reader/deserialize path consumes that discriminator to restore the correct
enum variant.
In `@modelexpress_server/src/state.rs`:
- Around line 239-266: Update the test_worker_record_conversion unit test to
include and verify the new TransferEngine branch: construct a WorkerMetadata
that sets backend_metadata to the TransferEngine variant (with a
transfer_engine_session_id value), convert it to WorkerRecord via
WorkerRecord::from(...) and back into WorkerMetadata via Into, and assert that
worker_rank, backend_metadata (specifically the TransferEngine session id), and
tensors round-trip correctly; reference the existing
test_worker_record_conversion, WorkerMetadata, WorkerRecord, and
BackendMetadataRecord symbols when adding the new case so the backwards-compat
wrapper will catch regressions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 6e4f7ec9-4a14-415b-a2c4-d791459d0882
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (10)
modelexpress_client/python/modelexpress/p2p_pb2.pymodelexpress_client/python/modelexpress/types.pymodelexpress_common/proto/p2p.protomodelexpress_server/src/k8s_types.rsmodelexpress_server/src/metadata_backend.rsmodelexpress_server/src/metadata_backend/kubernetes.rsmodelexpress_server/src/metadata_backend/layered.rsmodelexpress_server/src/metadata_backend/memory.rsmodelexpress_server/src/metadata_backend/redis.rsmodelexpress_server/src/state.rs
AndyDai-nv
left a comment
There was a problem hiding this comment.
Could you do a sign off to pass DCO check? overall LGTM. Let's wait for @KavinKrishnan's review as well, just for a double check
|
Look pretty good overall Some questions (maybe can be disregarded if this PR is just to setup subsequent followup for SGL support):
Some recommendations:
|
Extend WorkerMetadata with oneof backend_metadata supporting both NIXL (bytes) and TransferEngine (session_id string). Update all metadata backends (memory, redis, kubernetes, layered) and regenerate Python protobuf stubs for protobuf 5.x compatibility.
- Extract BackendMetadataRecord::from_flat() to deduplicate deserialization logic shared between Redis and Kubernetes backends - Restore SPDX license headers on generated Python stubs
Add explicit backend_type field ("nixl", "transfer_engine", "none") to
WorkerRecordJson (Redis), WorkerStatus (K8s), and CRD YAML instead of
inferring the backend variant from which fields are populated.
from_flat() uses backend_type as authoritative discriminator when present,
falling back to field-inference for backwards compatibility with records
written before this change.
Also adds transferEngineSessionId to CRD schema and unit tests for
TransferEngine round-trip and from_flat discriminator logic.
a9f1142 to
cfbe428
Compare
Summary
WorkerMetadatainp2p.protowithoneof backend_metadatasupporting both NIXL (bytes) and TransferEngine (session_id string)Context
This enables SGLang's Mooncake TransferEngine to use ModelExpress for metadata coordination, alongside the existing NIXL support. The seed instance publishes its TransferEngine session ID + tensor descriptors to MX, and target instances query MX to discover the seed and perform RDMA reads.
Design: typed
oneofvs opaque bytesWe use a typed
oneofrather than opaquebytes backend_metadata+string backend_type. Today MX acts as a metadata store, but the plan is for the MX client to own the transfer pipeline (compute transfer plan, execute RDMA, apply fixups). When MX owns transfers, it needs to interpret the backend metadata to select the right transfer backend -- so theoneofgives us compile-time type safety for that future.Changes
p2p.proto:WorkerMetadata.backend_metadataoneof withnixl_metadata(field 2) andtransfer_engine_session_id(field 10)metadata_backend.rs: NewBackendMetadataRecordenum (Nixl/TransferEngine/None) withfrom_flat()helper for deserializationk8s_types.rs: Addedtransfer_engine_session_id: Option<String>toWorkerStatusredis.rs: Addedtransfer_engine_session_idtoWorkerRecordJsonFromconversions updatedgrpcio-tools<=1.66.2Test plan
Summary by CodeRabbit
New Features
Refactor