Skip to content

[Feat] Asu: add transport sub-batch foundation#997

Open
yuanzhg078 wants to merge 1 commit into
ModelEngine-Group:feature_26h1from
yuanzhg078:dev_asu_trans_pr_1
Open

[Feat] Asu: add transport sub-batch foundation#997
yuanzhg078 wants to merge 1 commit into
ModelEngine-Group:feature_26h1from
yuanzhg078:dev_asu_trans_pr_1

Conversation

@yuanzhg078
Copy link
Copy Markdown
Contributor

@yuanzhg078 yuanzhg078 commented Jun 5, 2026

Purpose

Add the ASU transport sub-batch foundation for splitting large requests, tracking sub-batch state, mapping ASU response/result codes, and finalizing transport task completion.

Modifications

Add TransportSubBatchContext and TransportSubBatchState to track per-sub-batch CID, operation type, connection, buffers, status, and entry results.
Add ASU batch limits by operation type
Add IoScheduler to split entry and key batches into SSU-compliant sub-batches.
Add ASU response status mapping from CQE status and result-buffer entries to transport Status and QueryResult.
Add task completion utilities to release send/flag buffers, release connection inflight counts, count terminal sub-batches, and finalize task state.
Update BufferManager SGE lkey assignment and related test expectation.
Wire ASU transport unit tests into the CMake test target.

Test

Added AsuResponseStatusTest, IoSchedulerTest, TransportTaskCompletionTest, and updated BufferManagerTest coverage

Introduce ASU sub-batch task state, IO scheduling limits, response status mapping, and completion finalization primitives.

Add focused tests for scheduler splitting, CQE/result status conversion, buffer release, and terminal task state handling.
@yuanzhg078 yuanzhg078 requested a review from Infinite666 June 5, 2026 06:43
@pyxyzc pyxyzc self-requested a review June 5, 2026 07:27
set(Ascend_ROOT "/usr/local/Ascend/ascend-toolkit/latest" CACHE PATH "Path to Ascend root directory")
endif()

find_path(ASU_ASCEND_INCLUDE_DIR
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.

Require acl/acl.h before checking the runtime. Move the Ascend header discovery inside the Ascend runtime block, or make the option explicitly reject non-Ascend runtimes.

Comment thread ucm/transport/kv/asu/trans/src/asu_response_status.cpp
Comment thread ucm/transport/kv/asu/trans/src/asu_response_status.cpp
Comment thread ucm/transport/kv/asu/trans/src/transport_task_completion.cpp
Comment thread ucm/transport/kv/asu/trans/src/transport_task_completion.cpp
Comment thread ucm/transport/kv/asu/trans/src/transport_task_completion.cpp
Comment thread ucm/transport/kv/asu/trans/src/transport_task_completion.cpp
Comment thread ucm/transport/kv/asu/trans/src/transport_task_completion.cpp
} // namespace

std::vector<IoScheduler::ScheduledIoBatch> IoScheduler::SplitForAsu(
const BatchView<KVBuffer>& entries, std::size_t maxIoNum) const
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.

Can we set maxIoNum when IoScheduler init like config instead of passing it every time?

Comment thread ucm/transport/kv/asu/trans/src/io_scheduler.cpp
Comment thread ucm/transport/kv/asu/trans/src/transport_task_manager.h
Comment thread ucm/transport/kv/asu/trans/src/transport_task_manager.h
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants