[TE]feat: ascend direct transport add async tranfer task limit

[ascend-direct-dev]

Description

Asynchronous transmission uses stream resources which are limited, so a async task limit mechanism has been added.

Type of Change

• Types
  • Bug fix
  • New feature
    • Transfer Engine
    • Mooncake Store
    • Mooncake EP
    • Integration
    • P2P Store
    • Python Wheel
  • Breaking change
  • CI/CD
  • Documentation update
  • Other

How Has This Been Tested?

Checklist

• I have performed a self-review of my own code.
• I have updated the documentation.
• I have added tests to prove my changes are effective.

[gemini-code-assist]

Summary of Changes

Hello @ascend-direct-dev, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a critical asynchronous task limiting mechanism within the AscendDirectTransport to prevent resource over-utilization during high-volume data transfers. It also refines the configuration loading process by centralizing environment variable parsing and making the worker thread pool size dynamically configurable, leading to more robust and adaptable transport behavior.

Highlights

• Asynchronous Task Limiting: Implemented a mechanism to limit the number of concurrent asynchronous transfer tasks in AscendDirectTransport to kAsyncTaskLimit (defaulting to 120), preventing resource exhaustion and ensuring stable operation.
• Environment Variable Configuration Refinement: Centralized and improved the parsing of several environment variables (ASCEND_CONNECT_TIMEOUT, ASCEND_TRANSFER_TIMEOUT, ASCEND_USE_SHORT_CONNECTION, ASCEND_THREAD_POOL_SIZE, ASCEND_BASE_PORT, ASCEND_RT_VISIBLE_DEVICES) using a new parseFromString utility, enhancing robustness and providing better error logging.
• Dynamic Thread Pool Sizing: The worker thread pool size is now configurable via the ASCEND_THREAD_POOL_SIZE environment variable, with a default of 8 and a maximum allowed value of 16, offering more flexibility in resource allocation.
• Removed Connection Retry Logic: The explicit retry mechanism for adxl::NOT_CONNECTED status in connectAndTransfer has been removed, simplifying the connection handling logic.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to limit the number of concurrent asynchronous transfer tasks in the AscendDirectTransport, which is a valuable addition for managing limited stream resources. The implementation of the task limiting logic is sound. My review includes suggestions to improve maintainability by refactoring duplicated code, a question about the removal of connection retry logic which could impact resilience, and a suggestion to restore a warning log for better debuggability.

I am having trouble creating individual review comments. Click here to see my feedback.

mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp (808-815)

The retry logic for adxl::NOT_CONNECTED has been removed. This could decrease the system's resilience to transient connection issues. Was this removal intentional? If so, it would be good to document this change in behavior. If not, please consider reintroducing the retry mechanism to handle temporary connection drops gracefully.

mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp (114-169)

There is significant code duplication in parsing environment variables (ASCEND_CONNECT_TIMEOUT, ASCEND_TRANSFER_TIMEOUT, ASCEND_USE_SHORT_CONNECTION, ASCEND_THREAD_POOL_SIZE). This can be refactored to improve code clarity and maintainability. Consider creating a helper function to handle parsing values from environment variables. For example, for the timeout variables:

void parse_timeout_from_env(const char* env_var_name, int32_t& timeout_variable, const std::string& log_message) {
    char* env_value_str = std::getenv(env_var_name);
    if (env_value_str) {
        std::optional<int32_t> timeout_value = parseFromString<int32_t>(env_value_str);
        if (timeout_value.has_value()) {
            timeout_variable = timeout_value.value();
            LOG(INFO) << log_message << timeout_variable;
        }
    }
}

// Usage:
parse_timeout_from_env("ASCEND_CONNECT_TIMEOUT", connect_timeout_, "Set connection timeout to:");
parse_timeout_from_env("ASCEND_TRANSFER_TIMEOUT", transfer_timeout_, "Set transfer timeout to:");

A similar approach could be used for the other variables, possibly with a template to handle different types and validation logic.

mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp (576-587)

The else block that handles an out-of-bounds dev_id has been removed. This removes a useful warning log which could aid in debugging misconfigurations of ASCEND_RT_VISIBLE_DEVICES. Please consider adding it back to warn users about potential out-of-bounds access.

        if (dev_id < static_cast<int32_t>(device_list.size())) {
            std::optional<int32_t> dev_id_opt =
                parseFromString<int32_t>(device_list[dev_id]);
            if (dev_id_opt.has_value()) {
                dev_id = dev_id_opt.value();
                LOG(INFO) << "Set device id to:" << dev_id;
            } else {
                LOG(WARNING) << "Device id is " << dev_id
                             << ", ASCEND_RT_VISIBLE_DEVICES is "
                             << rt_visible_devices << ", which is unexpected.";
            }
        } else {
            LOG(WARNING) << "Device id is " << dev_id
                         << ", ASCEND_RT_VISIBLE_DEVICES is "
                         << rt_visible_devices << ", which is unexpected.";
        }

[codecov-commenter]

⚠️ Please install the 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!