@@ -15,36 +15,51 @@ #include "../../../../../../legacy/common/utils/string_util.h" #include "../../../../../../legacy/unified_platform/resource/mem/user_remote_mem_getter.h" #include "comm_mems.h" +#include "env_config/env_config.h" namespace hcomm { AivUbMemTransport::AivUbMemTransport(Hccl::Socket *socket, HcommChannelDesc &channelDesc) : socket_(socket), channelDesc_(channelDesc) {} +HcclResult AivUbMemTransport::FillTagVec(HcommMemHandle *memHandles, uint32_t bufferNum, + std::vector &bufferVec, std::vector> &tagVec) +{ + uint32_t totalBufferNum = localRmaBufferVec_.size() + bufferNum; + localUserMemTag_.reserve(totalBufferNum); + if (UNLIKELY(totalBufferNum > MAX_BUFFER_NUM)) { + HCCL_ERROR("[AivUbMemTransport][FillTagVec] totalBufferNum[%u] exceeds limit[%u]", totalBufferNum, MAX_BUFFER_NUM); + return HCCL_E_PARA; + } + for (uint32_t i = 0; i < bufferNum; ++i) { + auto locMemInfo = reinterpret_cast(memHandles[i]); + CHK_PTR_NULL(locMemInfo); + auto localIpcRmaBuffer = reinterpret_cast(locMemInfo->bufferHandle); + CHK_PTR_NULL(localIpcRmaBuffer); + bufferVec.push_back(localIpcRmaBuffer); + std::array memTag{}; + std::string tag = locMemInfo->memTag; + if (UNLIKELY(tag.size() >= HCCL_RES_TAG_MAX_LEN)) { + HCCL_ERROR("[AivUbMemTransport][FillTagVec] tagSize exceeds limit[%u]", HCCL_RES_TAG_MAX_LEN); + return HCCL_E_PARA; + } + CHK_SAFETY_FUNC_RET(memcpy_s(memTag.data(), memTag.size(), tag.c_str(), tag.size())); + HCCL_INFO("[AivUbMemTransport][FillTagVec] memHandleNum[%u] memTag[%s]", i, memTag.data()); + tagVec.push_back(memTag); + } + return HCCL_SUCCESS; +} + HcclResult AivUbMemTransport::Init() { // 拷贝memTag信息 uint32_t bufferNum = channelDesc_.memHandleNum; - localUserMemTag_.reserve(bufferNum); if (bufferNum == 0) { HCCL_ERROR("[AivUbMemTransport][Init] bufferNum is 0."); return HCCL_E_PARA; } - HCCL_INFO("channelDesc_.memHandleNum: %u", bufferNum); - for (uint32_t i = 0; i < bufferNum; ++i) { - auto locMemInfo = reinterpret_cast(channelDesc_.memHandles[i]); - auto localIpcRmaBuffer = reinterpret_cast(locMemInfo->bufferHandle); - localRmaBufferVec_.push_back(localIpcRmaBuffer); - std::array memTag{}; - std::string tag = locMemInfo->memTag; - if (UNLIKELY(tag.size() >= HCCL_RES_TAG_MAX_LEN)) { - HCCL_ERROR("[AivUbMemTransport][Init] tagSize exceeds limit[%u]", HCCL_RES_TAG_MAX_LEN); - return HCCL_E_PARA; - } - CHK_SAFETY_FUNC_RET(memcpy_s(memTag.data(), memTag.size(), tag.c_str(), tag.size())); - HCCL_INFO("[AivUbMemTransport][Init] memHandleNum[%u] memTag[%s]", i, memTag.data()); - localUserMemTag_.push_back(memTag); - } + HCCL_INFO("[AivUbMemTransport][Init] channelDesc_.memHandleNum: %u", bufferNum); + CHK_RET(FillTagVec(channelDesc_.memHandles, bufferNum, localRmaBufferVec_, localUserMemTag_)); baseStatus_ = Hccl::TransportStatus::INIT; return HCCL_SUCCESS; @@ -53,7 +68,7 @@ HcclResult AivUbMemTransport::Init() HcclResult AivUbMemTransport::IsSocketReady(bool &isReady) { CHK_PTR_NULL(socket_); - + EXCEPTION_HANDLE_BEGIN Hccl::SocketStatus socketStatus = socket_->GetAsyncStatus(); if (socketStatus == Hccl::SocketStatus::OK) { baseStatus_ = Hccl::TransportStatus::SOCKET_OK; @@ -62,7 +77,7 @@ HcclResult AivUbMemTransport::IsSocketReady(bool &isReady) baseStatus_ = Hccl::TransportStatus::SOCKET_TIMEOUT; isReady = false; } - + EXCEPTION_HANDLE_END return HCCL_SUCCESS; } @@ -87,6 +102,14 @@ Hccl::TransportStatus AivUbMemTransport::GetStatus() baseStatus_ = Hccl::TransportStatus::SOCKET_OK; break; case AivUbMemTransportStatus::SOCKET_OK: + SendDataSize(); + aivUbStatus_ = AivUbMemTransportStatus::SEND_DATA_SIZE; + break; + case AivUbMemTransportStatus::SEND_DATA_SIZE: + RecvDataSize(); + aivUbStatus_ = AivUbMemTransportStatus::RECV_DATA_SIZE; + break; + case AivUbMemTransportStatus::RECV_DATA_SIZE: SendMemInfo(); aivUbStatus_ = AivUbMemTransportStatus::SEND_MEM_INFO; break; @@ -109,40 +132,63 @@ Hccl::TransportStatus AivUbMemTransport::GetStatus() return baseStatus_; } -HcclResult AivUbMemTransport::SendMemInfo() +HcclResult AivUbMemTransport::SendDataSize() { HCCL_INFO("[%s] start", __func__); Hccl::BinaryStream binaryStream; - BufferPack(binaryStream); // 名字？ + BufferPack(binaryStream, localRmaBufferVec_, localUserMemTag_); binaryStream.Dump(sendData_); + u32 sendSize = sendData_.size(); EXCEPTION_HANDLE_BEGIN - socket_->SendAsync(reinterpret_cast(&sendData_[0]), sendData_.size()); + socket_->SendAsync(reinterpret_cast(&sendSize), sizeof(sendSize)); EXCEPTION_HANDLE_END - exchangeDataSize_ = sendData_.size(); HCCL_INFO("[%s] finished", __func__); return HCCL_SUCCESS; } -void AivUbMemTransport::BufferPack(Hccl::BinaryStream &binaryStream) +HcclResult AivUbMemTransport::RecvDataSize() { - u32 vecSize = localRmaBufferVec_.size(); + HCCL_INFO("[%s] start", __func__); + + EXCEPTION_HANDLE_BEGIN + socket_->RecvAsync(reinterpret_cast(&exchangeDataSize_), sizeof(exchangeDataSize_)); + EXCEPTION_HANDLE_END + HCCL_INFO("[%s] finished", __func__); + return HCCL_SUCCESS; +} + +HcclResult AivUbMemTransport::SendMemInfo() +{ + HCCL_INFO("[%s] start", __func__); + + EXCEPTION_HANDLE_BEGIN + socket_->SendAsync(reinterpret_cast(&sendData_[0]), sendData_.size()); + EXCEPTION_HANDLE_END + HCCL_INFO("[%s] finished", __func__); + return HCCL_SUCCESS; +} + +void AivUbMemTransport::BufferPack(Hccl::BinaryStream &binaryStream, std::vector &bufferVec, + std::vector> &tagVec) +{ + u32 vecSize = bufferVec.size(); binaryStream << vecSize; HCCL_RUN_INFO("BufferPack vecSize=%u", vecSize); - for (uint32_t i = 0; i < vecSize; ++i) { - std::unique_ptr dto = localRmaBufferVec_[i]->GetExchangeDto(); - HCCL_INFO("[%s] dto[%s]", __func__, dto->Describe().c_str()); - dto->Serialize(binaryStream); - } - - for (const auto& tag : localUserMemTag_) { + for (const auto& tag : tagVec) { // 逐个字节传输 for (uint32_t i = 0; i < HCCL_RES_TAG_MAX_LEN; i++) { binaryStream << static_cast(tag[i]); } } + + for (uint32_t i = 0; i < vecSize; ++i) { + std::unique_ptr dto = bufferVec[i]->GetExchangeDto(); + HCCL_INFO("[%s] dto[%s]", __func__, dto->Describe().c_str()); + dto->Serialize(binaryStream); + } } HcclResult AivUbMemTransport::RecvMemInfo() @@ -158,6 +204,9 @@ HcclResult AivUbMemTransport::RecvMemInfo() HcclResult AivUbMemTransport::RecvDataProcess() { Hccl::BinaryStream binaryStream(recvData_); + rmtBufferVec_.clear(); + rmtRmaBufferVec_.clear(); + remoteUserMemTag_.clear(); EXCEPTION_HANDLE_BEGIN RmtBufferUnpackProc(binaryStream); EXCEPTION_HANDLE_END @@ -166,36 +215,40 @@ HcclResult AivUbMemTransport::RecvDataProcess() void AivUbMemTransport::RmtBufferUnpackProc(Hccl::BinaryStream &binaryStream) { - rmtBufferVec_.clear(); // 这两个可以归一吗？ - rmtRmaBufferVec_.clear(); u32 vecSize{0}; binaryStream >> vecSize; HCCL_RUN_INFO("vecSize=%u", vecSize); - if (UNLIKELY(vecSize > MAX_BUFFER_NUM)) { + uint32_t totalBufferNum = remoteUserMemTag_.size() + vecSize; + if (UNLIKELY(totalBufferNum > MAX_BUFFER_NUM)) { EXCEPTION_THROW_IF_ERR(HCCL_E_PARA, "[AivUbMemTransport][RmtBufferUnpackProc] vecSize exceeds limit."); } + + rmtTagTemp_.resize(vecSize); + for (auto& tag : rmtTagTemp_) { + for (uint32_t i = 0; i < HCCL_RES_TAG_MAX_LEN; i++) { + u8 byte; + binaryStream >> byte; + tag[i] = static_cast(byte); + } + } + remoteUserMemTag_.insert(remoteUserMemTag_.end(), rmtTagTemp_.begin(), rmtTagTemp_.end()); + for (u32 pos = 0; pos < vecSize; ++pos) { Hccl::ExchangeIpcBufferDto dto; dto.Deserialize(binaryStream); HCCL_INFO("[%s] dto[%s]", __func__, dto.Describe().c_str()); + const char* src = rmtTagTemp_[pos].data(); + std::string bufTag(src, strnlen(src, HCCL_RES_TAG_MAX_LEN)); if (dto.size == 0) { // size为0，则为 remote 空buffer HCCL_INFO("unpack nullptr, pos=%u", pos); rmtBufferVec_.push_back(nullptr); rmtRmaBufferVec_.push_back(nullptr); } else { // size非0，则构造一个remote buffer - rmtBufferVec_.push_back(std::make_unique(dto, "UbMemory")); + HCCL_INFO("[AivUbMemTransport][RmtBufferUnpackProc] unpack buffer tag[%s]", bufTag.c_str()); + rmtBufferVec_.push_back(std::make_unique(dto, bufTag)); rmtRmaBufferVec_.push_back(rmtBufferVec_.back().get()); } } - - remoteUserMemTag_.resize(vecSize); - for (auto& tag : remoteUserMemTag_) { - for (uint32_t i = 0; i < HCCL_RES_TAG_MAX_LEN; i++) { - u8 byte; - binaryStream >> byte; - tag[i] = static_cast(byte); - } - } } HcclResult AivUbMemTransport::GetRemoteMem(HcclMem **remoteMem, uint32_t *memNum, char **memTags) @@ -271,4 +324,66 @@ HcclResult AivUbMemTransport::GetUserRemoteMem(CommMem **remoteMem, char ***memT CHK_RET(Hccl::GetRemoteUserMem(remoteMemCtx)); return HCCL_SUCCESS; } + +HcclResult AivUbMemTransport::CheckSocketStatus() +{ + CHK_PTR_NULL(socket_); + auto timeout = std::chrono::seconds(Hccl::EnvConfig::GetInstance().GetSocketConfig().GetLinkTimeOut()); + auto startTime = std::chrono::steady_clock::now(); + uint32_t retryCount = 0; + while(true) { + EXCEPTION_HANDLE_BEGIN + Hccl::SocketStatus socketStatus = socket_->GetAsyncStatus(); + if (socketStatus == Hccl::SocketStatus::OK) { + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - startTime).count(); + HCCL_INFO("[AivUbMemTransport][%s] success, elapsed[%lld]ms, retryCount[%u]", + __func__, elapsed, retryCount); + break; + } + if ((std::chrono::steady_clock::now() - startTime) >= timeout || + socketStatus == Hccl::SocketStatus::TIMEOUT) { + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - startTime).count(); + HCCL_ERROR("[AivUbMemTransport][%s] channel connect timeout after %lld sec, elapsed[%lld]ms, retryCount[%u]", + __func__, timeout, elapsed, retryCount); + return HCCL_E_TIMEOUT; + } + EXCEPTION_HANDLE_END + retryCount++; + } + return HCCL_SUCCESS; +} + +HcclResult AivUbMemTransport::UpdateMemInfo(HcommMemHandle *memHandles, uint32_t memHandleNum) +{ + if (memHandleNum == 0) { + HCCL_WARNING("[AivUbMemTransport][UpdateMemInfo] bufferNum is 0."); + return HCCL_SUCCESS; + } + locMemTemp_.clear(); + locTagTemp_.clear(); + CHK_RET(FillTagVec(memHandles, memHandleNum, locMemTemp_, locTagTemp_)); + HCCL_INFO("[AivUbMemTransport][UpdateMemInfo] bufferNum[%zu]", locMemTemp_.size()); + sendData_.clear(); + Hccl::BinaryStream sendStream; + BufferPack(sendStream, locMemTemp_, locTagTemp_); + sendStream.Dump(sendData_); + u32 sendSize = sendData_.size(); + EXCEPTION_HANDLE_BEGIN + socket_->SendAsync(reinterpret_cast(&sendSize), sizeof(sendSize)); + EXCEPTION_HANDLE_END + CHK_RET(CheckSocketStatus()); + CHK_RET(RecvDataSize()); + CHK_RET(CheckSocketStatus()); + CHK_RET(SendMemInfo()); + CHK_RET(CheckSocketStatus()); + CHK_RET(RecvMemInfo()); + CHK_RET(CheckSocketStatus()); + Hccl::BinaryStream recvStream(recvData_); + RmtBufferUnpackProc(recvStream); + localRmaBufferVec_.insert(localRmaBufferVec_.end(), locMemTemp_.begin(), locMemTemp_.end()); + localUserMemTag_.insert(localUserMemTag_.end(), locTagTemp_.begin(), locTagTemp_.end()); + return HCCL_SUCCESS; +} }