Skip to content

[TENT] Add best-effort RDMA task cancellation#2851

Open
catyans wants to merge 3 commits into
kvcache-ai:mainfrom
catyans:feature/tent-rdma-cancel
Open

[TENT] Add best-effort RDMA task cancellation#2851
catyans wants to merge 3 commits into
kvcache-ai:mainfrom
catyans:feature/tent-rdma-cancel

Conversation

@catyans

@catyans catyans commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

Add task-level, best-effort cancellation to TENT, with the first transport implementation for RDMA.

  • expose cancellation through the C++, C, and Python APIs
  • cancel owners that are still waiting in the local admission queue without dispatching them
  • let RDMA workers suppress queued, grouped, and retrying slices before ibv_post_send
  • drain already-posted WRs through the normal CQ path instead of reporting a false cancellation
  • preserve partial transferred_bytes, prevent failover after cancellation, and make repeated cancellation idempotent
  • release RDMA device quota exactly once when an unposted slice is canceled or abandoned

This is a lifecycle follow-up to #2132. It does not change staging cancellation or add cancellation support to non-RDMA transports.

Semantics

cancelTransfer(batch_id, task_id) returning OK means the request was accepted, not that the task is already terminal. Callers continue polling before freeing the batch.

  • admission-queued task: terminal CANCELED, never dispatched
  • worker-owned but unposted RDMA slice: terminal CANCELED
  • mixed posted/unposted RDMA task: posted WRs drain; final status is CANCELED, with completed bytes retained
  • fully posted RDMA task: may still finish COMPLETED
  • merged public tasks: canceling any alias cancels their shared physical task
  • staging or direct non-RDMA task: NotImplemented

Validation

  • admission_queue_test, tent_runtime_queue_dispatch_test, and rdma_cancel_test: pass
  • the three suites above repeated 100 times: pass
  • targeted ASan + UBSan cancellation tests: pass
  • real mlx5 RDMA integration on a server with eight active bonded RNICs:
    • existing 4 MiB WRITE then READ data-integrity path: pass
    • 1 worker / 1 QP / 1 outstanding WR, 16 tasks: tasks 0-14 completed and task 15 reached CANCELED
    • repeated in 10 independent containers: 10/10 pass
  • pre-commit: whitespace, EOF, YAML, merge markers, large files, ruff, codespell, clang-format, and cmake-format pass
  • clang-format 20 was run in the Linux validation environment and produced no diff

Compatibility

Default behavior is unchanged unless a caller invokes the new cancellation API. The base Transport implementation reports NotImplemented, so transports opt in explicitly.

AI assistance

This change was developed with AI assistance. The implementation, edge-case semantics, tests, sanitizer results, and real-RDMA behavior were reviewed and validated by the submitter.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Code Review

This pull request introduces best-effort task cancellation for queued and RDMA tasks across the Transfer Engine, implementing cancellation APIs in the C++, C, and Python interfaces, and updating the admission queue, RDMA transport, and worker threads. The review highlights several important issues: a data race on non-atomic volatile int fields (success_slices and resolved_slices) in RdmaTask that should be atomic; a potential null-pointer dereference of worker_context_ in Workers::cancel if called after shutdown; the incorrect use of CHECK_POINTER on the non-pointer batch_id in the C API; and a missing null check for slice->task in cancelUnpostedSlice.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp Outdated
Comment thread mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp
Comment thread mooncake-transfer-engine/tent/src/transfer_engine_c.cpp Outdated
Comment thread mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp Outdated
@github-actions github-actions Bot added documentation Improvements or additions to documentation run-ci Transfer Engine labels Jul 10, 2026
@codecov-commenter

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@catyans

catyans commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator Author

Addressed the Gemini cancellation review issues in a3755091:

  • changed RdmaTask::success_slices / resolved_slices from volatile ints to atomics and updated status accounting to use atomic fetch/load
  • made Workers::cancel() best-effort safe after worker shutdown by checking running_, worker_context_, and num_workers_
  • added a slice->task guard in cancelUnpostedSlice()
  • replaced C API CHECK_POINTER(batch_id) checks with scalar if (!batch_id) return -1 checks for cancel/status entry points

Validation:

  • clang-format and git diff --check passed
  • remote Linux container build with USE_TENT=ON, BUILD_UNIT_TESTS=ON, USE_CUDA=OFF
  • built rdma_cancel_test and tent_runtime_queue_dispatch_test
  • ctest -R "rdma_cancel_test|tent_runtime_queue_dispatch_test" --output-on-failure: 2/2 passed

@catyans catyans marked this pull request as ready for review July 11, 2026 04:23
@catyans

catyans commented Jul 12, 2026

Copy link
Copy Markdown
Collaborator Author

Merged the latest upstream/main in 159712b7 and resolved the only conflict in tent/tests/CMakeLists.txt by retaining both rdma_cancel_test and the newly added thread_local_storage_test.

Post-merge validation on the H20/RDMA cluster:

  • clean configure/build with USE_TENT=ON, BUILD_UNIT_TESTS=ON, USE_CUDA=OFF
  • focused ctest suite: 5/5 passed (rdma_cancel_test, tent_rdma_transport_test, admission_queue_test, tent_runtime_queue_dispatch_test, thread_local_storage_test)
  • rdma_cancel_test: 100 shuffled repeats passed
  • tent_rdma_transport_test: 20 shuffled repeats passed
  • git diff --check passed

GitHub CI has been retriggered for the merge commit.

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

Labels

documentation Improvements or additions to documentation run-ci Transfer Engine

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants