Skip to content

feat: Add TransferEngine backend to P2P metadata#157

Merged
ishandhanani merged 4 commits intomainfrom
ishan/transfer-engine-backend
Mar 17, 2026
Merged

feat: Add TransferEngine backend to P2P metadata#157
ishandhanani merged 4 commits intomainfrom
ishan/transfer-engine-backend

Conversation

@ishandhanani
Copy link
Copy Markdown
Contributor

@ishandhanani ishandhanani commented Mar 5, 2026

Summary

  • Extend WorkerMetadata in p2p.proto with oneof backend_metadata supporting both NIXL (bytes) and TransferEngine (session_id string)
  • Update all metadata backends (memory, redis, kubernetes, layered) to handle the new oneof variant
  • Regenerate Python protobuf stubs for protobuf 5.x compatibility

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 oneof vs opaque bytes

We use a typed oneof rather than opaque bytes 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 the oneof gives us compile-time type safety for that future.

Changes

  • p2p.proto: WorkerMetadata.backend_metadata oneof with nixl_metadata (field 2) and transfer_engine_session_id (field 10)
  • metadata_backend.rs: New BackendMetadataRecord enum (Nixl/TransferEngine/None) with from_flat() helper for deserialization
  • k8s_types.rs: Added transfer_engine_session_id: Option<String> to WorkerStatus
  • redis.rs: Added transfer_engine_session_id to WorkerRecordJson
  • All backend From conversions updated
  • Python stubs regenerated with grpcio-tools<=1.66.2

Test plan

  • 170 Rust tests pass, 0 clippy warnings
  • E2E roundtrip: publish metadata with transfer_engine_session_id, get it back with correct oneof
  • Backward compat: existing NIXL metadata path unaffected
  • SGLang E2E: Qwen3-235B-A22B-FP8, TP4 EP4, 22x speedup (0.74s RDMA vs 16.43s disk)

Summary by CodeRabbit

  • New Features

    • Added support for TransferEngine session IDs as an alternative metadata option for worker configuration.
  • Refactor

    • Restructured worker metadata to support multiple backend types as mutually exclusive options.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented Mar 6, 2026

Walkthrough

The 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

Cohort / File(s) Summary
Protocol Buffer Definition
modelexpress_common/proto/p2p.proto
Added transfer_engine_session_id field as part of a new backend_metadata oneof alongside nixl_metadata, making them mutually exclusive options for backend-specific metadata.
Generated Protobuf Code
modelexpress_client/python/modelexpress/p2p_pb2.py
Updated serialized descriptor offsets and blob data across multiple message types (_WORKERMETADATA, _PUBLISHMETADATAREQUEST, _PUBLISHMETADATARESPONSE, _GETMETADATA*, _GETREADY*, _P2PSERVICE) to reflect new field additions.
Python Types
modelexpress_client/python/modelexpress/types.py
Added transfer_engine_session_id: str field to WorkerMetadata and applied default value b"" to existing nixl_metadata field.
Rust K8s Types
modelexpress_server/src/k8s_types.rs
Added transfer_engine_session_id: Option<String> field to WorkerStatus with serde rename to transferEngineSessionId.
Rust Metadata Backend Core
modelexpress_server/src/metadata_backend.rs
Introduced BackendMetadataRecord enum with variants Nixl, TransferEngine, and None; replaced WorkerRecord.nixl_metadata with backend_metadata: BackendMetadataRecord; added conversion logic in From implementations.
Rust Metadata Backend Implementations
modelexpress_server/src/metadata_backend/kubernetes.rs, modelexpress_server/src/metadata_backend/redis.rs
Updated worker metadata handling to extract and store both nixl_metadata and transfer_engine_session_id separately, then reconstruct via BackendMetadataRecord::from_flat() when deserializing.
Rust Test Metadata Backends
modelexpress_server/src/metadata_backend/layered.rs, modelexpress_server/src/metadata_backend/memory.rs
Updated test constructors to use backend_metadata: Some(BackendMetadata::NixlMetadata(...)) instead of direct nixl_metadata field; added necessary imports.
Rust State/API
modelexpress_server/src/state.rs
Publicly re-exported BackendMetadataRecord and updated WorkerRecord public API to expose backend_metadata: BackendMetadataRecord instead of nixl_metadata; updated test expectations accordingly.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Poem

🐰 A session of hop and transfer,
Where backends no longer confer,
One metadata blooms, not two—
NIXL or Engine shines through,
New oneof hops structure anew! 🌟

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 53.85% 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
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately summarizes the main change: adding TransferEngine as a new backend option to P2P metadata handling, which is reflected throughout all modified files.

✏️ 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: 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 a WorkerMetadata -> WorkerRecord -> WorkerMetadata case 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![]) and BackendMetadataRecord::None both serialize here as nixl_metadata = "" and transfer_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

📥 Commits

Reviewing files that changed from the base of the PR and between 3376ee5 and 5ca91fd.

⛔ Files ignored due to path filters (1)
  • Cargo.lock is excluded by !**/*.lock
📒 Files selected for processing (10)
  • modelexpress_client/python/modelexpress/p2p_pb2.py
  • modelexpress_client/python/modelexpress/types.py
  • modelexpress_common/proto/p2p.proto
  • modelexpress_server/src/k8s_types.rs
  • modelexpress_server/src/metadata_backend.rs
  • modelexpress_server/src/metadata_backend/kubernetes.rs
  • modelexpress_server/src/metadata_backend/layered.rs
  • modelexpress_server/src/metadata_backend/memory.rs
  • modelexpress_server/src/metadata_backend/redis.rs
  • modelexpress_server/src/state.rs

Comment thread modelexpress_server/src/k8s_types.rs Outdated
Comment thread modelexpress_server/src/metadata_backend/kubernetes.rs
@AndyDai-nv AndyDai-nv changed the title Add TransferEngine backend to P2P metadata feat: Add TransferEngine backend to P2P metadata Mar 6, 2026
@github-actions github-actions Bot added the feat label Mar 6, 2026
Copy link
Copy Markdown
Contributor

@AndyDai-nv AndyDai-nv left a comment

Choose a reason for hiding this comment

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

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

Comment thread modelexpress_server/src/metadata_backend/redis.rs
Comment thread modelexpress_server/src/metadata_backend/kubernetes.rs
@KavinKrishnan
Copy link
Copy Markdown
Contributor

Look pretty good overall

Some questions (maybe can be disregarded if this PR is just to setup subsequent followup for SGL support):

  1. How are you handling configuration of which backend to use (NIXL or TE)? IIUC we are always using NIXL currently
  2. How are we ensuring target node and source are using the same backend (and respective routing logic)?

Some recommendations:

  1. Env variables for backend configuration
  2. Maybe have a dedicated TransferEngineManager like NixlTransferManager

ishandhanani and others added 3 commits March 15, 2026 21:07
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.
@ishandhanani ishandhanani merged commit 63643c7 into main Mar 17, 2026
16 checks passed
@ishandhanani ishandhanani deleted the ishan/transfer-engine-backend branch March 17, 2026 20:25
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.

3 participants