Skip to content

[TENT] Make UB batch submission all-or-nothing - #56

Open
Connor-Matthew wants to merge 2 commits into
mainfrom
ub/batch-submit-atomic
Open

Connor-Matthew wants to merge 2 commits into
mainfrom
ub/batch-submit-atomic

Conversation

@Connor-Matthew

@Connor-Matthew Connor-Matthew commented Nov 5, 2026 •

Copy link
Copy Markdown

Description

Internal CI run for the UB batch-submission fix. Not for merging into the fork's main.

UbTransport::submitTransferTasks() handed tasks to UbWorkers one at a time. When a later task was rejected, the earlier tasks of the same call were already queued and could only be asked to cancel, so a failed submit could still run part of the batch. UbWorkers::submit() also skipped any slice that could not be marked queued and still reported success for the rest of the task.

Changes:

  • Add UbWorkers::submitBatch(). It checks every task and slice before changing any state, then marks and queues all slices under the queue lock. If a slice changes state, the workers stop, or the queue cannot grow, nothing is published, and any task whose slices were already marked is cancelled so no slice is left stuck in the queued state. submit() now goes through submitBatch().
  • submitTransferTasks() submits the whole request list in one call, then adds the tasks to the sub-batch.
  • The progress callback is armed only after the tasks are in the sub-batch. A very small task could complete before it was queryable, and its notification was lost. The transport notifies once after arming to cover that window.

This fix was split out of the TENT_4 branch (ub/tent4-integration) and contains no receiver-credit code.

Suggested Merge Priority

  • P0 — Urgent: Critical security vulnerability, data loss/corruption, or widespread production outage requiring immediate action
  • P1 — High: Significant correctness or stability issue in existing functionality, or a regression blocking users without a practical workaround
  • P2 — Normal: Routine bug fix, validated performance/compatibility improvement, or new feature for a confirmed use case
  • P3 — Low: Experimental feature, optional enhancement, refactoring, or cleanup with no urgent user impact

Module

  • Transfer Engine (mooncake-transfer-engine)
  • Mooncake Store (mooncake-store)
  • Mooncake Conductor (mooncake-conductor)
  • Reshard (mooncake-reshard)
  • Mooncake EP (mooncake-ep)
  • Mooncake PG (mooncake-pg)
  • Integration (mooncake-integration)
  • P2P Store (mooncake-p2p-store)
  • Python Wheel (mooncake-wheel)
  • Common (mooncake-common)
  • Mooncake RL (mooncake-rl)
  • CI/CD
  • Docs
  • Other

Type of Change

  • Bug fix
  • New feature
  • Refactor
  • Breaking change
  • Documentation update
  • Performance improvement
  • Other

How Has This Been Tested?

Local Linux arm64 container with the CI ub-mock configuration (Debug, -DUSE_UB=ON), full build plus the TENT test directory.

Test commands:

cmake -G Ninja .. -DUSE_TENT=ON -DUSE_HTTP=ON -DBUILD_UNIT_TESTS=ON -DUSE_CUDA=OFF \
  -DWITH_STORE=OFF -DWITH_STORE_RUST=OFF -DUSE_UB=ON -DCMAKE_BUILD_TYPE=Debug
ninja
ctest --test-dir mooncake-transfer-engine/tent/tests -j8 --output-on-failure
./mooncake-transfer-engine/tent/tests/tent_ub_native_data_path_test --gtest_repeat=50

Test results:

  • Unit tests pass: 68 of 69 TENT targets pass. transfer_engine_config_override_test fails only in the local container: its DNS resolves an intentionally invalid hostname, and the container has no NUMA sysfs. The same four cases failed there before this change, and this PR does not touch that code.
  • tent_ub_native_data_path_test passes 50 repeated runs.
  • Mutation check: with submitBatch() changed back to submitting tasks one at a time, the new workers-level test fails (one task queued, slice left in the queued state).

Checklist

  • Every commit I contribute includes my Signed-off-by trailer certifying the DCO
  • I have performed a self-review of my own code
  • I have formatted my code using ./scripts/code_format.sh
  • I have run pre-commit on the files changed in this PR and all hooks pass
  • I have updated the documentation (if applicable)
  • I have added tests to prove my changes are effective
  • For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

  • No AI tools were used
  • AI tools were used (specify below)

An AI coding assistant helped split this fix out of the TENT_4 branch, port the tests, run the local container build and mutation check, and draft this description. The human submitter is responsible for reviewing every changed line.

UbTransport::submitTransferTasks() handed tasks to UbWorkers one at a
time. When a later task was rejected, the earlier tasks of the same call
were already queued and could only be asked to cancel, so a failed submit
could still run part of the batch. UbWorkers::submit() also skipped any
slice that could not be marked queued and reported success for the rest
of the task.

Add UbWorkers::submitBatch(). It checks every task and slice before
changing any state, then marks and queues all slices under the queue
lock. If a slice changes state, the workers stop, or the queue cannot
grow, nothing is published, and any task whose slices were already
marked is cancelled so no slice is left stuck in the queued state.
submit() now goes through submitBatch(), and submitTransferTasks()
submits the whole request list in one call before adding the tasks to
the sub-batch.

Arm the progress callback only after the tasks are in the sub-batch. A
very small task could complete before it was queryable and its
notification was lost; notify once after arming to cover that window.

Co-authored-by: Le1zyCatt <2490162471@qq.com>
Signed-off-by: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
@github-actions github-actions Bot added run-ci Transfer Engine size/M 100-1000 changed lines (additions + deletions) labels Nov 5, 2026
Roll back partial enqueues for allocation and length errors. Cancel marked tasks before constructing the returned error status.

Inject a second-enqueue failure to verify worker rollback and preserve existing transport sub-batch task IDs after a failed submit.

Signed-off-by: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci size/M 100-1000 changed lines (additions + deletions) Transfer Engine

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant