Line data Source code
1 : /**
2 : * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3 : * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 : * CANN Open Software License Agreement Version 2.0 (the "License").
5 : * Please refer to the License for details. You may not use this file except in compliance with the License.
6 : * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7 : * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8 : * See LICENSE in the root of the software repository for the full text of the License.
9 : */
10 :
11 : #include "aiv_ub_mem_transport.h"
12 : #include "exception_handler.h"
13 : #include "../../../../../../legacy/ascend950/unified_platform/resource/socket/socket.h"
14 : #include "../../../../../../legacy/ascend950/unified_platform/resource/buffer/exchange_ipc_buffer_dto.h"
15 : #include "../../../../../../legacy/ascend950/unified_platform/resource/mem/user_remote_mem_getter.h"
16 : #include "env_config/env_config_v2.h"
17 :
18 : namespace hcomm {
19 :
20 10 : AivUbMemTransport::AivUbMemTransport(Hccl::Socket* socket, HcommChannelDesc& channelDesc)
21 10 : : socket_(socket),
22 10 : channelDesc_(channelDesc)
23 10 : {}
24 :
25 2 : HcclResult AivUbMemTransport::FillBufferVec(
26 : HcommMemHandle* memHandles, uint32_t bufferNum, std::vector<Hccl::LocalIpcRmaBuffer*>& bufferVec)
27 : {
28 2 : uint32_t totalBufferNum = localRmaBufferVec_.size() + bufferNum;
29 2 : if (UNLIKELY(totalBufferNum > MAX_BUFFER_NUM)) {
30 0 : HCCL_ERROR(
31 : "[AivUbMemTransport][FillBufferVec] totalBufferNum[%u] exceeds limit[%u]", totalBufferNum, MAX_BUFFER_NUM);
32 0 : return HCCL_E_PARA;
33 : }
34 5 : for (uint32_t i = 0; i < bufferNum; ++i) {
35 3 : auto localIpcRmaBuffer = reinterpret_cast<Hccl::LocalIpcRmaBuffer*>(memHandles[i]);
36 3 : CHK_PTR_NULL(localIpcRmaBuffer);
37 3 : auto buf = localIpcRmaBuffer->GetBuf();
38 3 : CHK_PTR_NULL(buf);
39 3 : bufferVec.push_back(localIpcRmaBuffer);
40 3 : HCCL_INFO(
41 : "[AivUbMemTransport][FillBufferVec] memHandleNum[%u] buffer[%s]", i, localIpcRmaBuffer->Describe().data());
42 : }
43 2 : return HCCL_SUCCESS;
44 : }
45 :
46 1 : HcclResult AivUbMemTransport::Init()
47 : {
48 1 : uint32_t bufferNum = channelDesc_.memHandleNum;
49 1 : if (bufferNum == 0) {
50 1 : HCCL_ERROR("[AivUbMemTransport][Init] bufferNum is 0.");
51 1 : return HCCL_E_PARA;
52 : }
53 0 : HCCL_INFO("[AivUbMemTransport][Init] channelDesc_.memHandleNum: %u", bufferNum);
54 0 : CHK_RET(FillBufferVec(channelDesc_.memHandles, bufferNum, localRmaBufferVec_));
55 :
56 0 : baseStatus_ = Hccl::TransportStatus::INIT;
57 0 : return HCCL_SUCCESS;
58 : }
59 :
60 8 : HcclResult AivUbMemTransport::IsSocketReady(bool& isReady)
61 : {
62 8 : CHK_PTR_NULL(socket_);
63 : EXCEPTION_HANDLE_BEGIN
64 8 : Hccl::SocketStatus socketStatus = socket_->GetAsyncStatus();
65 7 : if (socketStatus == Hccl::SocketStatus::OK) {
66 7 : baseStatus_ = Hccl::TransportStatus::SOCKET_OK;
67 7 : isReady = true;
68 0 : } else if (socketStatus == Hccl::SocketStatus::TIMEOUT) {
69 0 : baseStatus_ = Hccl::TransportStatus::SOCKET_TIMEOUT;
70 0 : isReady = false;
71 : }
72 1 : EXCEPTION_HANDLE_END
73 7 : return HCCL_SUCCESS;
74 : }
75 :
76 6 : void AivUbMemTransport::CheckStatusFuncResult(std::string funcName, HcclResult ret)
77 : {
78 6 : if (UNLIKELY(ret != HCCL_SUCCESS)) {
79 1 : HCCL_ERROR(
80 : "[%s] fail ret[%d], aivUbStatus_[%d], baseStatus_[%d]", funcName.c_str(), ret, aivUbStatus_, baseStatus_);
81 1 : baseStatus_ = Hccl::TransportStatus::INVALID;
82 : }
83 6 : }
84 :
85 8 : Hccl::TransportStatus AivUbMemTransport::GetStatus()
86 : {
87 8 : if (baseStatus_ == Hccl::TransportStatus::READY || baseStatus_ == Hccl::TransportStatus::INVALID) {
88 0 : return baseStatus_;
89 8 : } else if (baseStatus_ == Hccl::TransportStatus::INIT) {
90 0 : aivUbStatus_ = AivUbMemTransportStatus::INIT;
91 : }
92 :
93 8 : bool isReady = false;
94 8 : if (UNLIKELY(IsSocketReady(isReady) != HCCL_SUCCESS)) {
95 1 : HCCL_ERROR("[%s] IsSocketReady fail, aivUbStatus_[%d], baseStatus_[%d]", __func__, aivUbStatus_, baseStatus_);
96 1 : baseStatus_ = Hccl::TransportStatus::INVALID;
97 1 : return baseStatus_;
98 : }
99 7 : if (!isReady) {
100 0 : return baseStatus_;
101 : }
102 7 : return UpdateStatus();
103 : }
104 :
105 7 : Hccl::TransportStatus AivUbMemTransport::UpdateStatus()
106 : {
107 7 : HCCL_INFO(
108 : "%s aivUbStatus_[%d], baseStatus_[%d] start, aivUbStatus_::SOCKET_OK[%d]", __func__, aivUbStatus_, baseStatus_,
109 : AivUbMemTransportStatus::SOCKET_OK);
110 : HcclResult ret;
111 7 : switch (aivUbStatus_) {
112 0 : case AivUbMemTransportStatus::INIT:
113 0 : aivUbStatus_ = AivUbMemTransportStatus::SOCKET_OK;
114 0 : baseStatus_ = Hccl::TransportStatus::SOCKET_OK;
115 0 : break;
116 1 : case AivUbMemTransportStatus::SOCKET_OK:
117 1 : ret = SendDataSize();
118 1 : CheckStatusFuncResult("SendDataSize", ret);
119 1 : aivUbStatus_ = AivUbMemTransportStatus::SEND_DATA_SIZE;
120 1 : break;
121 1 : case AivUbMemTransportStatus::SEND_DATA_SIZE:
122 1 : ret = RecvDataSize();
123 1 : CheckStatusFuncResult("RecvDataSize", ret);
124 1 : aivUbStatus_ = AivUbMemTransportStatus::RECV_DATA_SIZE;
125 1 : break;
126 1 : case AivUbMemTransportStatus::RECV_DATA_SIZE:
127 1 : ret = SendMemInfo();
128 1 : CheckStatusFuncResult("SendMemInfo", ret);
129 1 : aivUbStatus_ = AivUbMemTransportStatus::SEND_MEM_INFO;
130 1 : break;
131 1 : case AivUbMemTransportStatus::SEND_MEM_INFO:
132 1 : ret = RecvMemInfo();
133 1 : CheckStatusFuncResult("RecvMemInfo", ret);
134 1 : aivUbStatus_ = AivUbMemTransportStatus::RECV_MEM_INFO;
135 1 : break;
136 2 : case AivUbMemTransportStatus::RECV_MEM_INFO:
137 2 : ret = RecvDataProcess();
138 2 : CheckStatusFuncResult("RecvDataProcess", ret);
139 2 : aivUbStatus_ = AivUbMemTransportStatus::RECV_MEM_FIN;
140 2 : break;
141 1 : case AivUbMemTransportStatus::RECV_MEM_FIN:
142 1 : aivUbStatus_ = AivUbMemTransportStatus::READY;
143 1 : baseStatus_ = Hccl::TransportStatus::READY;
144 1 : break;
145 0 : default:
146 0 : break;
147 : }
148 7 : HCCL_INFO("%s aivUbStatus_[%d], baseStatus_[%d]", __func__, aivUbStatus_, baseStatus_);
149 7 : return baseStatus_;
150 : }
151 :
152 1 : HcclResult AivUbMemTransport::SendDataSize()
153 : {
154 1 : HCCL_INFO("[%s] start", __func__);
155 :
156 1 : Hccl::BinaryStream binaryStream;
157 1 : CHK_RET(BufferPack(binaryStream, localRmaBufferVec_));
158 :
159 1 : binaryStream.Dump(sendData_);
160 1 : u32 sendSize = sendData_.size();
161 : EXCEPTION_HANDLE_BEGIN
162 1 : socket_->SendAsync(&sendSize, sizeof(sendSize));
163 0 : EXCEPTION_HANDLE_END
164 1 : HCCL_INFO("[%s] finished", __func__);
165 1 : return HCCL_SUCCESS;
166 1 : }
167 :
168 2 : HcclResult AivUbMemTransport::RecvDataSize()
169 : {
170 2 : HCCL_INFO("[%s] start", __func__);
171 :
172 : EXCEPTION_HANDLE_BEGIN
173 2 : socket_->RecvAsync(reinterpret_cast<u8*>(&exchangeDataSize_), sizeof(exchangeDataSize_));
174 0 : EXCEPTION_HANDLE_END
175 2 : HCCL_INFO("[%s] finished", __func__);
176 2 : return HCCL_SUCCESS;
177 : }
178 :
179 2 : HcclResult AivUbMemTransport::SendMemInfo()
180 : {
181 2 : HCCL_INFO("[%s] start", __func__);
182 :
183 : EXCEPTION_HANDLE_BEGIN
184 2 : socket_->SendAsync(&sendData_[0], sendData_.size());
185 0 : EXCEPTION_HANDLE_END
186 2 : HCCL_INFO("[%s] finished", __func__);
187 2 : return HCCL_SUCCESS;
188 : }
189 :
190 : HcclResult
191 3 : AivUbMemTransport::BufferPack(Hccl::BinaryStream& binaryStream, std::vector<Hccl::LocalIpcRmaBuffer*>& bufferVec)
192 : {
193 3 : u32 vecSize = bufferVec.size();
194 3 : binaryStream << vecSize;
195 3 : HCCL_INFO("BufferPack vecSize=%u", vecSize);
196 :
197 6 : for (uint32_t i = 0; i < vecSize; ++i) {
198 3 : std::unique_ptr<Hccl::Serializable> dto = bufferVec[i]->GetExchangeDto();
199 3 : CHK_PTR_NULL(dto);
200 3 : dto->Serialize(binaryStream);
201 3 : HCCL_INFO("[%s] dto[%s]", __func__, dto->Describe().c_str());
202 3 : }
203 3 : return HCCL_SUCCESS;
204 : }
205 :
206 2 : HcclResult AivUbMemTransport::RecvMemInfo()
207 : {
208 2 : recvData_.resize(exchangeDataSize_);
209 : EXCEPTION_HANDLE_BEGIN
210 2 : socket_->RecvAsync(reinterpret_cast<u8*>(&recvData_[0]), recvData_.size());
211 0 : EXCEPTION_HANDLE_END
212 : // HCCL_INFO("recv data, size=%llu, data=%s", data.size(), Hccl::Bytes2hex(data.data(), data.size()).c_str());
213 2 : return HCCL_SUCCESS;
214 : }
215 :
216 2 : HcclResult AivUbMemTransport::RecvDataProcess()
217 : {
218 2 : Hccl::BinaryStream binaryStream(recvData_);
219 2 : rmtBufferVec_.clear();
220 2 : rmtRmaBufferVec_.clear();
221 : EXCEPTION_HANDLE_BEGIN
222 2 : RmtBufferUnpackProc(binaryStream);
223 1 : EXCEPTION_HANDLE_END
224 1 : return HCCL_SUCCESS;
225 2 : }
226 :
227 2 : void AivUbMemTransport::RmtBufferUnpackProc(Hccl::BinaryStream& binaryStream)
228 : {
229 2 : u32 vecSize{0};
230 2 : binaryStream >> vecSize;
231 2 : HCCL_INFO("vecSize=%u", vecSize);
232 2 : uint32_t totalBufferNum = rmtBufferVec_.size() + vecSize;
233 2 : if (UNLIKELY(totalBufferNum > MAX_BUFFER_NUM)) {
234 0 : EXCEPTION_THROW_IF_ERR(HCCL_E_PARA, "[AivUbMemTransport][RmtBufferUnpackProc] vecSize exceeds limit.");
235 : }
236 :
237 3 : for (u32 pos = 0; pos < vecSize; ++pos) {
238 1 : Hccl::ExchangeIpcBufferDto dto;
239 1 : dto.Deserialize(binaryStream);
240 1 : HCCL_INFO("[%s] dto[%s]", __func__, dto.Describe().c_str());
241 1 : if (dto.size == 0) { // size为0，则为 remote 空buffer
242 0 : HCCL_INFO("unpack nullptr, pos=%u", pos);
243 0 : rmtBufferVec_.push_back(nullptr);
244 0 : rmtRmaBufferVec_.push_back(nullptr);
245 : } else { // size非0，则构造一个remote buffer
246 1 : HCCL_INFO("[AivUbMemTransport][RmtBufferUnpackProc] unpack buffer memInfo[%s]", dto.memInfo.c_str());
247 1 : rmtBufferVec_.push_back(std::make_unique<Hccl::RemoteIpcRmaBuffer>(dto));
248 1 : rmtBufferVec_.back()->SetPathMode(channelDesc_.ubMemAttr.pathMode);
249 1 : rmtRmaBufferVec_.push_back(rmtBufferVec_.back().get());
250 : }
251 1 : }
252 2 : }
253 :
254 2 : HcclResult AivUbMemTransport::GetRemoteMems(uint32_t* memNum, CommMem** remoteMem, char*** memInfos)
255 : {
256 2 : std::lock_guard<std::mutex> lock(remoteMemsMutex_);
257 : Hccl::RemoteMemCtx<std::unique_ptr<Hccl::RemoteIpcRmaBuffer>> remoteMemCtx{
258 2 : cacheValid_, rmtBufferVec_, remoteUserMems_, memInfoCopies_, memInfoPointers_, remoteMem, memInfos, memNum};
259 2 : CHK_RET(GetRemoteUserMems(remoteMemCtx));
260 2 : return HCCL_SUCCESS;
261 2 : }
262 :
263 5 : HcclResult AivUbMemTransport::CheckSocketStatus(std::string socketOperator)
264 : {
265 5 : CHK_PTR_NULL(socket_);
266 5 : auto timeout = std::chrono::seconds(Hccl::EnvConfig::GetInstance().GetSocketConfig().GetLinkTimeOut());
267 5 : auto startTime = std::chrono::steady_clock::now();
268 5 : uint32_t retryCount = 0;
269 : while (true) {
270 : EXCEPTION_HANDLE_BEGIN
271 5 : Hccl::SocketStatus socketStatus = socket_->GetAsyncStatus();
272 5 : if (socketStatus == Hccl::SocketStatus::OK) {
273 : auto elapsed
274 4 : = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - startTime)
275 4 : .count();
276 4 : HCCL_INFO(
277 : "[AivUbMemTransport][%s] socket transport operation[%s] success, elapsed[%lld]ms, retryCount[%u]",
278 : __func__, socketOperator.c_str(), elapsed, retryCount);
279 4 : break;
280 : }
281 1 : if ((std::chrono::steady_clock::now() - startTime) >= timeout || socketStatus == Hccl::SocketStatus::TIMEOUT) {
282 : auto elapsed
283 1 : = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - startTime)
284 1 : .count();
285 1 : HCCL_ERROR(
286 : "[AivUbMemTransport][%s] socket transport operation[%s] timeout after %lld sec, elapsed[%lld]ms, "
287 : "retryCount[%u]",
288 : __func__, socketOperator.c_str(), timeout, elapsed, retryCount);
289 1 : return HCCL_E_TIMEOUT;
290 : }
291 0 : EXCEPTION_HANDLE_END
292 0 : retryCount++;
293 0 : }
294 4 : return HCCL_SUCCESS;
295 : }
296 :
297 3 : HcclResult AivUbMemTransport::UpdateMemInfo(HcommMemHandle* memHandles, uint32_t memHandleNum)
298 : {
299 3 : if (memHandleNum == 0) {
300 1 : HCCL_WARNING("[AivUbMemTransport][UpdateMemInfo] bufferNum is 0.");
301 1 : return HCCL_SUCCESS;
302 : }
303 2 : locMemTemp_.clear();
304 2 : CHK_RET(FillBufferVec(memHandles, memHandleNum, locMemTemp_));
305 2 : HCCL_INFO("[AivUbMemTransport][UpdateMemInfo] bufferNum[%zu]", locMemTemp_.size());
306 2 : sendData_.clear();
307 2 : Hccl::BinaryStream sendStream;
308 2 : CHK_RET(BufferPack(sendStream, locMemTemp_));
309 2 : sendStream.Dump(sendData_);
310 2 : u32 sendSize = sendData_.size();
311 : EXCEPTION_HANDLE_BEGIN
312 2 : socket_->SendAsync(&sendSize, sizeof(sendSize));
313 0 : EXCEPTION_HANDLE_END
314 4 : CHK_RET(CheckSocketStatus("SendDataSize"));
315 1 : CHK_RET(RecvDataSize());
316 2 : CHK_RET(CheckSocketStatus("RecvDataSize"));
317 1 : CHK_RET(SendMemInfo());
318 2 : CHK_RET(CheckSocketStatus("SendMemInfo"));
319 1 : CHK_RET(RecvMemInfo());
320 2 : CHK_RET(CheckSocketStatus("RecvMemInfo"));
321 1 : Hccl::BinaryStream recvStream(recvData_);
322 : EXCEPTION_HANDLE_BEGIN
323 1 : RmtBufferUnpackProc(recvStream);
324 0 : EXCEPTION_HANDLE_END
325 1 : localRmaBufferVec_.insert(localRmaBufferVec_.end(), locMemTemp_.begin(), locMemTemp_.end());
326 : // 流程中已有新增内存数量判断，故执行到此位置一定存在新增内存，需要将标识置位false，使得再次调用GetRemoteMems时重新构造缓存
327 1 : cacheValid_ = false;
328 1 : return HCCL_SUCCESS;
329 2 : }
330 : } // namespace hcomm
|
