diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/hccl_transport/device_listener.h b/mooncake-transfer-engine/include/transport/ascend_transport/hccl_transport/device_listener.h new file mode 100644 index 000000000..636d531d3 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/ascend_transport/hccl_transport/device_listener.h @@ -0,0 +1,74 @@ +// Copyright 2025 Huawei Technologies Co., Ltd +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +#include "hccl/hccl_types.h" + +namespace mooncake::detail { + +// HCCL supports Listen on port 0 and GetLocalPort after a successful Listen. +// Bind the NIC first, then use its assigned port for the VNIC. Only publish +// the pair once both listeners are ready. The factory is also the SDK seam +// used by hardware-free regression tests. +template +HcclResult listenOnAvailableDevicePort(Factory createSocket, + std::shared_ptr& nicListener, + std::shared_ptr& vnicListener, + uint64_t& devicePort) { + constexpr int kMaxAttempts = 32; + for (int attempt = 0; attempt < kMaxAttempts; ++attempt) { + auto nic = createSocket(/*vnic=*/false, /*port=*/0); + if (!nic) return HCCL_E_MEMORY; + auto result = nic->Init(); + if (result != HCCL_SUCCESS) return result; + result = nic->Listen(); + // In HCCL's RA listen adapter, an already-bound port is mapped to + // HCCL_E_UNAVAIL. Do not retry parameter, allocation or device errors. + if (result == HCCL_E_UNAVAIL) continue; + if (result != HCCL_SUCCESS) return result; + + const uint32_t port = nic->GetLocalPort(); + if (port == 0 || port > UINT16_MAX) return HCCL_E_PARA; + auto vnic = createSocket(/*vnic=*/true, port); + if (!vnic) return HCCL_E_MEMORY; + result = vnic->Init(); + if (result != HCCL_SUCCESS) return result; + result = vnic->Listen(); + if (result == HCCL_E_UNAVAIL) { + // DeInit stops a listener; Close alone only closes connections. + // Refuse another attempt if rollback itself fails. + const auto vnicCleanup = vnic->DeInit(); + const auto nicCleanup = nic->DeInit(); + if (vnicCleanup != HCCL_SUCCESS) return vnicCleanup; + if (nicCleanup != HCCL_SUCCESS) return nicCleanup; + continue; + } + if (result != HCCL_SUCCESS) return result; + if (vnic->GetLocalPort() != port) return HCCL_E_PARA; + + nicListener = std::move(nic); + vnicListener = std::move(vnic); + devicePort = port; + return HCCL_SUCCESS; + } + return HCCL_E_UNAVAIL; +} + +} // namespace mooncake::detail diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/hccl_transport/hccl_transport.h b/mooncake-transfer-engine/include/transport/ascend_transport/hccl_transport/hccl_transport.h index 94a09787f..63658d842 100644 --- a/mooncake-transfer-engine/include/transport/ascend_transport/hccl_transport/hccl_transport.h +++ b/mooncake-transfer-engine/include/transport/ascend_transport/hccl_transport/hccl_transport.h @@ -33,7 +33,8 @@ #define THREAD_NUM 1 #define ASCEND_DEFAULT_HOST_PORT 10000 -#define ASCEND_DEFAULT_DEVICE_PORT 16666 +// Port 0 asks HCCL to atomically allocate an available listener port. +#define ASCEND_DEFAULT_DEVICE_PORT 0 namespace mooncake { class TransferMetadata; @@ -103,4 +104,4 @@ class HcclTransport : public Transport { RankInfo remote_rank_info_; }; } // namespace mooncake -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/hccl_transport_mem_c.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/hccl_transport_mem_c.cpp index 35688b369..48297cfe8 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/hccl_transport_mem_c.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/ascend_transport_c/hccl_transport_mem_c.cpp @@ -27,8 +27,8 @@ #include #include #include -#include "mpi.h" #include "transport/ascend_transport/hccl_transport/hccl_transport_mem_c.h" +#include "transport/ascend_transport/hccl_transport/device_listener.h" #ifdef __cplusplus extern "C" { @@ -137,16 +137,6 @@ static int initServerNetSocket(RankInfo *local_rank_info) { local_rank_info->deviceLogicId, localIp), "HcclNetOpenDev DEVICE_NIC_TYPE failed"); - nicServerSocket_ = std::make_shared( - nicNetDevCtx_, local_rank_info->devicePort); - if (nicServerSocket_ == NULL) { - LOG(ERROR) << "make nicNetDevCtx_ failed"; - return -1; - } - - RETRY_CALL(nicServerSocket_->Init(), "nicServerSocket_ Init failed"); - RETRY_CALL(nicServerSocket_->Listen(), "nicServerSocket_ Listen failed"); - // Use virtual network card within HCCS hccl::HcclIpAddress localVnicIp(local_rank_info->devicePhyId); RETRY_CALL( @@ -160,17 +150,19 @@ static int initServerNetSocket(RankInfo *local_rank_info) { local_rank_info->deviceLogicId, localVnicIp), "HcclNetOpenDev vnicNetDevCtx_ failed"); - // control plane connection, creat serversocket, listening client - vnicServerSocket_ = std::make_shared( - vnicNetDevCtx_, local_rank_info->devicePort); - if (vnicServerSocket_ == NULL) { - LOG(ERROR) << "vnicServerSocket_ make failed"; + const auto listenResult = + mooncake::detail::listenOnAvailableDevicePort( + [](bool vnic, uint32_t port) { + return std::make_shared( + vnic ? vnicNetDevCtx_ : nicNetDevCtx_, port); + }, + nicServerSocket_, vnicServerSocket_, local_rank_info->devicePort); + if (listenResult != HCCL_SUCCESS) { + LOG(ERROR) << "NIC/VNIC listener initialization failed, HcclResult: " + << listenResult; return -1; } - RETRY_CALL(vnicServerSocket_->Init(), "vnicServerSocket_ Init failed"); - RETRY_CALL(vnicServerSocket_->Listen(), "vnicServerSocket_ Listen failed"); - RETRY_CALL(HcclDispatcherInit(DispatcherType::DISPATCHER_NORMAL, local_rank_info->devicePhyId, &dispatcher_), "client HcclDispatcherInit failed"); diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 3686051af..7986553ac 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -670,6 +670,17 @@ add_executable(common_test ${WORKSPACE}/common_test.cpp) target_link_libraries(common_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME common_test COMMAND common_test) +# Uses the real listener selection helper with a controlled socket boundary; +# requires HCCL headers but no accelerator or transport engine initialization. +if(USE_ASCEND) + add_executable(hccl_device_listener_test + ${WORKSPACE}/hccl_device_listener_test.cpp) + target_include_directories(hccl_device_listener_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + target_link_libraries(hccl_device_listener_test PRIVATE gtest gtest_main) + add_test(NAME hccl_device_listener_test COMMAND hccl_device_listener_test) +endif() + if(USE_ASCEND_DIRECT) # AscendDirectTransport unit test with mock implementations Mock # implementations are included in the test file via anonymous namespace diff --git a/mooncake-transfer-engine/tests/hccl_device_listener_test.cpp b/mooncake-transfer-engine/tests/hccl_device_listener_test.cpp new file mode 100644 index 000000000..23f7cec54 --- /dev/null +++ b/mooncake-transfer-engine/tests/hccl_device_listener_test.cpp @@ -0,0 +1,182 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "transport/ascend_transport/hccl_transport/device_listener.h" + +namespace mooncake { +namespace { + +struct SocketState { + HcclResult initResult = HCCL_SUCCESS; + HcclResult listenResult = HCCL_SUCCESS; + int vnicConflicts = 0; + int creations = 0; + int listeners = 0; +}; + +// Two loopback IPs stand in for the separate NIC and VNIC endpoints. Binding +// and listening are real OS operations; only SDK errors are controlled. +class TestSocket { + public: + TestSocket(SocketState& state, bool vnic, uint32_t port) + : state_(state), vnic_(vnic), port_(port) { + ++state_.creations; + } + ~TestSocket() { DeInit(); } + + HcclResult Init() { + if (state_.initResult != HCCL_SUCCESS) return state_.initResult; + fd_ = socket(AF_INET, SOCK_STREAM, 0); + return fd_ < 0 ? HCCL_E_SYSCALL : HCCL_SUCCESS; + } + + HcclResult Listen() { + if (state_.listenResult != HCCL_SUCCESS) return state_.listenResult; + if (vnic_ && state_.vnicConflicts > 0) { + --state_.vnicConflicts; + return HCCL_E_UNAVAIL; + } + sockaddr_in address{}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(vnic_ ? 0x7f000002 : 0x7f000001); + address.sin_port = htons(port_); + if (bind(fd_, reinterpret_cast(&address), sizeof(address)) != + 0) { + return errno == EADDRINUSE ? HCCL_E_UNAVAIL : HCCL_E_SYSCALL; + } + if (listen(fd_, 1) != 0) return HCCL_E_SYSCALL; + socklen_t size = sizeof(address); + if (getsockname(fd_, reinterpret_cast(&address), &size) != 0) + return HCCL_E_SYSCALL; + port_ = ntohs(address.sin_port); + listening_ = true; + ++state_.listeners; + return HCCL_SUCCESS; + } + + HcclResult DeInit() { + if (fd_ >= 0) close(fd_); + fd_ = -1; + if (listening_) --state_.listeners; + listening_ = false; + return HCCL_SUCCESS; + } + + uint32_t GetLocalPort() const { return port_; } + + private: + SocketState& state_; + bool vnic_; + uint32_t port_; + int fd_ = -1; + bool listening_ = false; +}; + +HcclResult createListeners(SocketState& state, std::shared_ptr& nic, + std::shared_ptr& vnic, uint64_t& port) { + return detail::listenOnAvailableDevicePort( + [&state](bool isVnic, uint32_t candidate) { + return std::make_shared(state, isVnic, candidate); + }, + nic, vnic, port); +} + +TEST(HcclDeviceListenerTest, OccupiedDefaultPortAllowsConcurrentListenerPairs) { + SocketState blockerState; + TestSocket nicBlocker(blockerState, false, 16666); + TestSocket vnicBlocker(blockerState, true, 16666); + ASSERT_EQ(nicBlocker.Init(), HCCL_SUCCESS); + ASSERT_EQ(vnicBlocker.Init(), HCCL_SUCCESS); + const auto nicBlocked = nicBlocker.Listen(); + const auto vnicBlocked = vnicBlocker.Listen(); + ASSERT_TRUE(nicBlocked == HCCL_SUCCESS || nicBlocked == HCCL_E_UNAVAIL); + ASSERT_TRUE(vnicBlocked == HCCL_SUCCESS || vnicBlocked == HCCL_E_UNAVAIL); + + constexpr int kWorkers = 8; + std::array states; + std::array, kWorkers> nics; + std::array, kWorkers> vnics; + std::array ports{}; + std::array results; + std::array workers; + for (int i = 0; i < kWorkers; ++i) { + workers[i] = std::thread([&, i] { + results[i] = + createListeners(states[i], nics[i], vnics[i], ports[i]); + }); + } + for (auto& worker : workers) worker.join(); + for (int i = 0; i < kWorkers; ++i) { + ASSERT_EQ(results[i], HCCL_SUCCESS); + EXPECT_NE(ports[i], 16666); + EXPECT_GT(ports[i], 0); + EXPECT_EQ(nics[i]->GetLocalPort(), ports[i]); + EXPECT_EQ(vnics[i]->GetLocalPort(), ports[i]); + for (int j = 0; j < i; ++j) EXPECT_NE(ports[i], ports[j]); + } +} + +TEST(HcclDeviceListenerTest, RollsBackNicWhenVnicPortIsBusy) { + SocketState state; + state.vnicConflicts = 1; + std::shared_ptr nic, vnic; + uint64_t port = 0; + ASSERT_EQ(createListeners(state, nic, vnic, port), HCCL_SUCCESS); + EXPECT_EQ(state.creations, 4); + EXPECT_EQ(state.listeners, 2); + nic.reset(); + vnic.reset(); + EXPECT_EQ(state.listeners, 0); +} + +TEST(HcclDeviceListenerTest, BoundsPortRetriesAndPreservesOtherErrors) { + SocketState state; + state.vnicConflicts = 32; + std::shared_ptr nic, vnic; + uint64_t port = 0; + EXPECT_EQ(createListeners(state, nic, vnic, port), HCCL_E_UNAVAIL); + EXPECT_EQ(state.creations, 64); + EXPECT_EQ(state.listeners, 0); + EXPECT_EQ(port, 0); + EXPECT_EQ(nic, nullptr); + EXPECT_EQ(vnic, nullptr); + + for (auto error : + {HCCL_E_MEMORY, HCCL_E_PARA, HCCL_E_SYSCALL, HCCL_E_RUNTIME}) { + SocketState failedState; + failedState.initResult = error; + EXPECT_EQ(createListeners(failedState, nic, vnic, port), error); + EXPECT_EQ(failedState.creations, 1); + EXPECT_EQ(failedState.listeners, 0); + failedState.initResult = HCCL_SUCCESS; + failedState.listenResult = error; + EXPECT_EQ(createListeners(failedState, nic, vnic, port), error); + EXPECT_EQ(failedState.creations, 2); + EXPECT_EQ(failedState.listeners, 0); + EXPECT_EQ(port, 0); + } +} + +} // namespace +} // namespace mooncake