[TENT] Add rule-based transport and device selection by alogfans · Pull Request #2079 · kvcache-ai/Mooncake

alogfans · 2026-05-11T07:17:32Z

Description

Add TransportSelector for flexible, configuration-driven transport selection policy while maintaining full backward compatibility.

Key features:

Configuration-driven transport selection via JSON policy rules
Support for segment_type filtering, device allocation, and transport priority
Legacy mode option (use_legacy_transport_selection) for exact original behavior
Default policies match original hardcoded behavior

Changes:

Add TransportSelector class with SelectionContext, SelectionPolicy, SelectionResult
Integrate TransportSelector into TransferEngineImpl
Add legacy mode support to preserve original code path

Module

Type of Change

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have formatted my own code using ./scripts/code_format.sh before submitting.
I have updated the documentation.
I have added tests to prove my changes are effective.

Add TransportSelector for flexible, configuration-driven transport selection policy while maintaining full backward compatibility. Key features: - Configuration-driven transport selection via JSON policy rules - Support for segment_type filtering, device allocation, and transport priority - Legacy mode option (use_legacy_transport_selection) for exact original behavior - Default policies match original hardcoded behavior Changes: - Add TransportSelector class with SelectionContext, SelectionPolicy, SelectionResult - Integrate TransportSelector into TransferEngineImpl - Add legacy mode support to preserve original code path - Restore TaskInfo fields (xport_priority, failover_count) for backward compatibility - Add max_failover_attempts configuration option Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

gemini-code-assist

Code Review

This pull request introduces a configuration-driven TransportSelector to replace hardcoded transport selection logic, allowing for more flexible policies based on segment type, memory patterns, and transfer size. Key changes include the integration of the selector into TransferEngineImpl, updates to the TaskInfo structure for device masking, and the addition of a legacy selection mode. Feedback focuses on correcting transport capability checks that incorrectly included file-specific flags for memory transfers, improving the robustness of the same_machine verification, handling NPU memory patterns, and ensuring safer JSON parsing when loading policies.

gemini-code-assist · 2026-05-11T07:19:02Z

+        bool same_machine =
+            (desc->machine_id ==
+             metadata_->segmentManager().getLocal()->machine_id);


The simplified same_machine check is less robust than the legacy implementation. It lacks checks for empty machine_id strings and does not verify if the local segment descriptor is valid. If machine_id is empty for both segments, they would incorrectly be considered as being on the same machine.

auto local_desc = metadata_->segmentManager().getLocal(); bool same_machine = local_desc && !desc->machine_id.empty() && !local_desc->machine_id.empty() && desc->machine_id == local_desc->machine_id;

gemini-code-assist · 2026-05-11T07:19:02Z

+bool TransportSelector::matchesMemoryPattern(const std::string& pattern,
+                                             MemoryType type) const {
+    if (pattern == kMemoryTypeWildcard) {
+        return true;
+    }
+
+    // Convert MemoryType to string for comparison
+    std::string type_str;
+    switch (type) {
+        case MTYPE_CPU:
+            type_str = kMemoryTypeCpu;
+            break;
+        case MTYPE_CUDA:
+            type_str = kMemoryTypeCuda;
+            break;
+        case MTYPE_ROCM:
+            type_str = "rocm";
+            break;
+        default:
+            type_str = "unknown";
+            break;
+    }
+
+    return pattern == type_str;
+}


The pattern matching logic for memory types does not handle the npu pattern, even though kMemoryTypeNpu is defined. Since MTYPE_CUDA is used to represent both CUDA and NPU devices in this context, the npu pattern should be explicitly matched against MTYPE_CUDA.

bool TransportSelector::matchesMemoryPattern(const std::string& pattern, MemoryType type) const { if (pattern == kMemoryTypeWildcard) { return true; } if (pattern == kMemoryTypeNpu && type == MTYPE_CUDA) { return true; } // Convert MemoryType to string for comparison std::string type_str; switch (type) { case MTYPE_CPU: type_str = kMemoryTypeCpu; break; case MTYPE_CUDA: type_str = kMemoryTypeCuda; break; case MTYPE_ROCM: type_str = "rocm"; break; default: type_str = "unknown"; break; } return pattern == type_str; }

gemini-code-assist · 2026-05-11T07:19:02Z

+        if (policy_json.contains("same_machine")) {
+            policy.same_machine = policy_json["same_machine"].get<bool>();
+        } else {
+            policy.same_machine = std::nullopt;
+        }


Using .get<bool>() (and similar .get<T>() calls for min_size, max_size, and priority) without verifying the JSON element's type can lead to a nlohmann::json::type_error exception if the configuration file contains unexpected types. It is safer to use .is_boolean() check or the .value() method with a default value.

Add transport_selector_test.cpp with test coverage for: - Default policies matching original behavior (File/Memory segments) - Transport type name parsing - Legacy mode enable/disable - Transport availability based on capabilities - Priority offset for fallback scenarios - Device mask handling - NVLINK same-machine constraint - ROCm memory type support - GPU-to-GPU, CPU-to-CPU, CPU-to-GPU, GPU-to-CPU transfers Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Fix compilation error by accessing Transport::caps (protected) through helper methods instead of a separate public member. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

codecov-commenter · 2026-05-11T09:15:42Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 50.00000% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...r-engine/tests/endpoint_store_integration_test.cpp	50.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

…elector

staryxchen · 2026-05-23T07:20:57Z

+            if (checkAvailability(transport_list_[GDS], local_mtype)) {
+                if (transport_index-- == 0) result.transport = GDS;
+            } else if (checkAvailability(transport_list_[IOURING],
+                                         local_mtype)) {


This changes the transport fallback behavior in legacy mode. Is this by design?

staryxchen · 2026-05-23T07:27:11Z

+    const std::vector<TransportType>*
+        buffer_transports;  // Pointer to transports in buffer
+    size_t transfer_size;   // Transfer size in bytes
+    int priority_level;     // Request priority level (higher = more urgent)


The QoS system uses PRIO_HIGH = 0 as the most urgent. The comment "higher = more urgent" is the opposite of the actual encoding.

staryxchen · 2026-05-23T07:28:45Z

    auto prev_type = task.type;

-    if (++task.failover_count > max_failover_attempts_) {
+    if (task.failover_count >= max_failover_attempts_) {


Why change this condition?

alogfans · 2026-05-25T02:46:10Z

@staryxchen We have addressed your issues by reverting some of the modifcations.

staryxchen

LGTM. But CI failed, I'm not sure if it is related with this PR.

When ibv_get_device_list returns a non-NULL list but num_devices == 0, we need to call ibv_free_device_list before returning to avoid leaking the allocated memory. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* [TENT] Add configuration-driven transport selector Add TransportSelector for flexible, configuration-driven transport selection policy while maintaining full backward compatibility. Key features: - Configuration-driven transport selection via JSON policy rules - Support for segment_type filtering, device allocation, and transport priority - Legacy mode option (use_legacy_transport_selection) for exact original behavior - Default policies match original hardcoded behavior Changes: - Add TransportSelector class with SelectionContext, SelectionPolicy, SelectionResult - Integrate TransportSelector into TransferEngineImpl - Add legacy mode support to preserve original code path - Restore TaskInfo fields (xport_priority, failover_count) for backward compatibility - Add max_failover_attempts configuration option Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * [TENT] Add comprehensive unit tests for TransportSelector Add transport_selector_test.cpp with test coverage for: - Default policies matching original behavior (File/Memory segments) - Transport type name parsing - Legacy mode enable/disable - Transport availability based on capabilities - Priority offset for fallback scenarios - Device mask handling - NVLINK same-machine constraint - ROCm memory type support - GPU-to-GPU, CPU-to-CPU, CPU-to-GPU, GPU-to-CPU transfers Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * [TENT] Fix FakeTransport to use protected caps member Fix compilation error by accessing Transport::caps (protected) through helper methods instead of a separate public member. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix * fix format issues * Add priority-based rule * Reformat * Fix review comments * Fix memory leak in endpoint_store_integration_test When ibv_get_device_list returns a non-NULL list but num_devices == 0, we need to call ibv_free_device_list before returning to avoid leaking the allocated memory. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Trigger CI --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

alogfans requested review from chestnut-Q and doujiang24 as code owners May 11, 2026 07:17

github-actions Bot added run-ci Transfer Engine labels May 11, 2026

gemini-code-assist Bot reviewed May 11, 2026

View reviewed changes

alogfans and others added 4 commits May 11, 2026 15:21

[TENT] Fix FakeTransport to use protected caps member

5112154

Fix compilation error by accessing Transport::caps (protected) through helper methods instead of a separate public member. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

fix

951f5ab

fix format issues

fe49b07

alogfans and others added 4 commits May 22, 2026 14:58

Merge branch 'main' into feat/tent-transport-selector

55705eb

Merge remote-tracking branch 'origin/main' into feat/tent-transport-s…

259baab

…elector

Add priority-based rule

0c9bece

Reformat

f2c6c6a

alogfans requested review from 00fish0, ShangmingCai, dtcccc, staryxchen, stmatengss and ykwd as code owners May 22, 2026 08:12

alogfans changed the title ~~[TENT] Add configuration-driven transport selector~~ [TENT] Add rule-based transport and device selection May 22, 2026

staryxchen reviewed May 23, 2026

View reviewed changes

Fix review comments

227a78e

staryxchen approved these changes May 25, 2026

View reviewed changes

stmatengss mentioned this pull request May 25, 2026

[RoadMap][Call For Contribution] Mooncake Project Overall Roadmap #1883

Open

69 tasks

alogfans and others added 2 commits May 26, 2026 14:39

Fix memory leak in endpoint_store_integration_test

d088b05

When ibv_get_device_list returns a non-NULL list but num_devices == 0, we need to call ibv_free_device_list before returning to avoid leaking the allocated memory. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Trigger CI

40e551f

alogfans merged commit 54411af into kvcache-ai:main May 27, 2026
17 of 19 checks passed

catyans mentioned this pull request Jun 23, 2026

[RFC]: Traffic class hint API for TransferEngine — bridging upper-layer semantics (KV / EP / CTRL) to lower-layer SL/TC #2568

Open

catyans mentioned this pull request Jul 12, 2026

[RFC]: TENT Backend Capability Model and First-Class Execution Plan #2863

Open

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[TENT] Add rule-based transport and device selection #2079

[TENT] Add rule-based transport and device selection #2079
alogfans merged 12 commits into
kvcache-ai:mainfrom
alogfans:feat/tent-transport-selector

alogfans commented May 11, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot May 11, 2026

Uh oh!

gemini-code-assist Bot May 11, 2026

Uh oh!

gemini-code-assist Bot May 11, 2026

Uh oh!

codecov-commenter commented May 11, 2026 •

edited

Loading

Uh oh!

staryxchen May 23, 2026

Uh oh!

staryxchen May 23, 2026

Uh oh!

staryxchen May 23, 2026

Uh oh!

alogfans commented May 25, 2026

Uh oh!

staryxchen left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

alogfans commented May 11, 2026

Description

Module

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot May 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot May 11, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

staryxchen May 23, 2026

Choose a reason for hiding this comment

Uh oh!

staryxchen May 23, 2026

Choose a reason for hiding this comment

Uh oh!

staryxchen May 23, 2026

Choose a reason for hiding this comment

Uh oh!

alogfans commented May 25, 2026

Uh oh!

staryxchen left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov-commenter commented May 11, 2026 •

edited

Loading