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 : #include "urma_direct_transport.h"
11 : #include "serializable.h"
12 : #include "exchange_ub_buffer_dto.h"
13 : #include "exchange_ub_conn_dto.h"
14 : #include "local_ub_rma_buffer.h"
15 : #include "orion_adapter_hccp.h"
16 :
17 : namespace Hccl {
18 : constexpr u32 FINISH_MSG_SIZE = 128;
19 : constexpr char_t FINISH_MSG[FINISH_MSG_SIZE] = "Ub Comm Pipe ready!";
20 :
21 : constexpr size_t RMT_BUFFER_VEC_SIZE = 3;
22 : constexpr size_t RMT_BUFFER_INDEX = 2;
23 : constexpr size_t CONN_NUM= 1;
24 : constexpr u32 WQE_SIZE = 64;
25 :
26 0 : UrmaDirectTransport::UrmaDirectTransport(CommonLocRes &commonLocRes, Attribution &attr, const LinkData &linkData,
27 0 : const Socket &socket, RdmaHandle rdmaHandle1)
28 0 : : BaseMemTransport(commonLocRes, attr, linkData, socket, TransportType::UB), rdmaHandle(rdmaHandle1)
29 0 : {}
30 :
31 0 : UrmaDirectTransport::UrmaDirectTransport(CommonLocRes &commonLocRes, Attribution &attr, const LinkData &linkData,
32 : const Socket &socket, RdmaHandle rdmaHandle1,
33 0 : std::function<void(u32 streamId, u32 taskId, const TaskParam &taskParam)> callback)
34 0 : : BaseMemTransport(commonLocRes, attr, linkData, socket, TransportType::UB, callback), rdmaHandle(rdmaHandle1)
35 0 : {}
36 :
37 0 : RemoteUbRmaBuffer* UrmaDirectTransport::GetRmtBuffer() const
38 : {
39 0 : HCCL_INFO("[%s] start", __func__);
40 0 : if (rmtBufferVec.size() != RMT_BUFFER_VEC_SIZE) {
41 0 : THROW<InternalException>(
42 0 : StringFormat("[%s] rmtBufferVec is not [%u], size[%u]", __func__, rmtBufferVec.size(), RMT_BUFFER_VEC_SIZE));
43 : }
44 0 : auto rmtBuf = rmtBufferVec[RMT_BUFFER_INDEX].get();
45 0 : CHECK_NULLPTR(rmtBuf, "[UrmaDirectTransport::GetRmtBuffer] rmtBuf is nullptr!");
46 0 : return rmtBuf;
47 : }
48 :
49 0 : std::string UrmaDirectTransport::Describe() const
50 : {
51 : string msg = StringFormat("UbMemTransport=[commonLocRes=%s, urmaStatus=%s, ",
52 0 : commonLocRes.Describe().c_str(), urmaStatus.Describe().c_str());
53 0 : msg += StringFormat("exchangeDataSize=%u, ", exchangeDataSize);
54 0 : return msg;
55 0 : }
56 :
57 0 : HcclAiRMAWQ UrmaDirectTransport::GetAiRMAWQ()
58 : {
59 0 : if (baseStatus != TransportStatus::READY) {
60 0 : MACRO_THROW(InternalException, StringFormat(
61 : "[UrmaDirectTransport::%s]transport status is not ready, please check", __func__));
62 : }
63 :
64 0 : HcclAiRMAWQ wq = {0};
65 0 : wq.wqeSize = WQE_SIZE;
66 :
67 0 : size_t connNum = commonLocRes.connVec.size();
68 0 : if (connNum != CONN_NUM) {
69 0 : THROW<InternalException>("[UrmaDirectTransport::%s] connNum[%llu] is not [%llu]",
70 : __func__, connNum, CONN_NUM);
71 : }
72 0 : auto conn = reinterpret_cast<DevUbCtpConnection *>(commonLocRes.connVec[0]);
73 0 : CHECK_NULLPTR(conn, StringFormat("[UrmaDirectTransport::%s] failed, connection pointer is nullptr", __func__));
74 0 : conn->SetWqInfo(wq);
75 :
76 0 : for (auto &it : commonLocRes.bufferVec) {
77 0 : if (it != nullptr) {
78 0 : LocalUbRmaBuffer* localBuffer = dynamic_cast<LocalUbRmaBuffer*>(it);
79 0 : CHECK_NULLPTR(localBuffer,
80 0 : StringFormat("[UrmaDirectTransport::%s] failed, localBuffer pointer is nullptr", __func__));
81 0 : HCCL_INFO("get local buffer, %s", localBuffer->Describe().c_str());
82 0 : wq.localTokenId = localBuffer->GetTokenId();
83 : }
84 : }
85 :
86 0 : for (auto &it : rmtBufferVec) {
87 0 : if (it != nullptr) {
88 0 : HCCL_INFO("get remote buffer, %s", it->Describe().c_str());
89 0 : wq.rmtObjId = it->GetTokenId();
90 0 : wq.rmtTokenValue = it->GetTokenValue();
91 : }
92 : }
93 :
94 0 : return wq;
95 : }
96 :
97 0 : HcclAiRMACQ UrmaDirectTransport::GetAiRMACQ()
98 : {
99 0 : if (baseStatus != TransportStatus::READY) {
100 0 : MACRO_THROW(InternalException, StringFormat(
101 : "[UrmaDirectTransport::%s]transport status is not ready, please check", __func__));
102 : }
103 0 : size_t connNum = commonLocRes.connVec.size();
104 0 : if (connNum != CONN_NUM) {
105 0 : THROW<InternalException>("[UrmaDirectTransport::%s] connNum[%llu] is not [%llu]",
106 : __func__, connNum, CONN_NUM);
107 : }
108 0 : auto conn = reinterpret_cast<DevUbCtpConnection *>(commonLocRes.connVec[0]);
109 0 : CHECK_NULLPTR(conn, StringFormat("[UrmaDirectTransport::%s] failed, connection pointer is nullptr", __func__));
110 :
111 0 : HcclAiRMACQ cq = {0};
112 0 : conn->SetCqInfo(cq);
113 0 : return cq;
114 : }
115 :
116 0 : void UrmaDirectTransport::SendExchangeData()
117 : {
118 0 : bufferNum = commonLocRes.bufferVec.size(); // 需要交换的buffer数量
119 0 : connNum = commonLocRes.connVec.size();
120 :
121 0 : HCCL_INFO("bufferNum=%u, connNum=%u", bufferNum, connNum);
122 :
123 0 : BinaryStream binaryStream;
124 0 : HandshakeMsgPack(binaryStream);
125 0 : BufferVecPack(binaryStream);
126 0 : ConnVecPack(binaryStream);
127 :
128 0 : binaryStream.Dump(sendData);
129 0 : socket->SendAsync(sendData.data(), sendData.size());
130 0 : exchangeDataSize = sendData.size();
131 :
132 0 : HCCL_INFO("send data %s, size=%llu", GetLinkDescInfo().c_str(), exchangeDataSize);
133 0 : }
134 :
135 0 : void UrmaDirectTransport::BufferVecPack(BinaryStream &binaryStream)
136 : {
137 0 : binaryStream << bufferNum;
138 0 : HCCL_INFO("start pack %s bufferVec", transportType.Describe().c_str());
139 0 : u32 pos = 0;
140 0 : for (auto &it : commonLocRes.bufferVec) {
141 0 : binaryStream << pos;
142 0 : if (it != nullptr) { // 非空的buffer，从buffer中获取 dto
143 0 : std::unique_ptr<Serializable> dto = it->GetExchangeDto();
144 0 : dto->Serialize(binaryStream);
145 0 : HCCL_INFO("pack buffer pos=%u dto %s", pos, dto->Describe().c_str());
146 0 : } else { // 空的buffer，dto所有字段为0(size=0)
147 0 : ExchangeUbBufferDto exchangeDto;
148 0 : exchangeDto.Serialize(binaryStream);
149 0 : HCCL_INFO("pack buffer pos=%u, dto is null %s", pos, exchangeDto.Describe().c_str());
150 0 : }
151 0 : pos++;
152 : }
153 0 : }
154 :
155 0 : bool UrmaDirectTransport::IsResReady()
156 : {
157 0 : for (auto &it : commonLocRes.connVec) {
158 0 : CHECK_NULLPTR(it,
159 0 : StringFormat("[UrmaDirectTransport::%s] failed, connection pointer is nullptr", __func__));
160 :
161 0 : RmaConnType connType = it->GetRmaConnType();
162 0 : if (connType != RmaConnType::UB) {
163 0 : THROW<InternalException>("[UrmaDirectTransport::%s] connection type[%s] is not ub",
164 0 : __func__, connType.Describe().c_str());
165 : }
166 :
167 0 : auto status = it->GetStatus();
168 0 : if (status != RmaConnStatus::EXCHANGEABLE &&
169 0 : status != RmaConnStatus::READY) {
170 0 : return false;
171 : }
172 : }
173 :
174 0 : HCCL_INFO("[UrmaDirectTransport::IsResReady] all resources ready.");
175 0 : return true;
176 : }
177 :
178 0 : void UrmaDirectTransport::RecvExchangeData()
179 : {
180 0 : recvData.resize(exchangeDataSize);
181 0 : socket->RecvAsync(reinterpret_cast<u8 *>(recvData.data()), recvData.size());
182 :
183 0 : HCCL_INFO("recv data %s, size=%llu", GetLinkDescInfo().c_str(), recvData.size());
184 0 : }
185 :
186 0 : bool UrmaDirectTransport::ConnVecUnpackProc(BinaryStream &binaryStream)
187 : {
188 : u32 rmtConnNum;
189 0 : binaryStream >> rmtConnNum;
190 0 : HCCL_INFO("start unpack conn %s connNum=%u, rmtConnNum=%u", GetLinkDescInfo().c_str(), connNum, rmtConnNum);
191 0 : if (connNum != rmtConnNum) {
192 0 : MACRO_THROW(InvalidParamsException,
193 : StringFormat("connNum=%u is not equal to rmtConnNum=%u", connNum, rmtConnNum));
194 : }
195 :
196 0 : bool result = false; // 不需要发送 finish
197 0 : for (u32 i = 0; i < rmtConnNum; i++) {
198 : u32 pos;
199 0 : binaryStream >> pos;
200 0 : ExchangeUbConnDto rmtDto;
201 0 : rmtDto.Deserialize(binaryStream);
202 0 : HCCL_INFO("unpack connection pos=%u dto %s", pos, rmtDto.Describe().c_str());
203 0 : if (commonLocRes.connVec[i]->GetStatus() != RmaConnStatus::READY) {
204 0 : HCCL_INFO("parse and import pos=%u, rmt dto to connection[%s]", pos,
205 : commonLocRes.connVec[i]->Describe().c_str());
206 0 : commonLocRes.connVec[i]->ParseRmtExchangeDto(rmtDto);
207 0 : commonLocRes.connVec[i]->ImportRmtDto();
208 0 : result = true; // connection 建链，需要发送finish
209 : }
210 0 : }
211 0 : return result;
212 : }
213 :
214 0 : void UrmaDirectTransport::RmtBufferVecUnpackProc(u32 locNum, BinaryStream &binaryStream, RemoteBufferVec &bufferVec)
215 : {
216 : u32 rmtNum;
217 0 : binaryStream >> rmtNum;
218 :
219 0 : HCCL_INFO("unpack BUFFER %s, locNum=%u, rmtNum=%u", GetLinkDescInfo().c_str(), locNum, rmtNum);
220 0 : if (rmtNum != locNum) {
221 0 : MACRO_THROW(InvalidParamsException,
222 : StringFormat("BUFFER, locNum=%u is not equal to rmtNum=%u", locNum, rmtNum));
223 : }
224 :
225 0 : for (u32 i = 0; i < rmtNum; i++) {
226 : u32 pos;
227 0 : binaryStream >> pos;
228 0 : ExchangeUbBufferDto dto;
229 0 : dto.Deserialize(binaryStream);
230 0 : if (bufferVec.size() > pos) {
231 : // 对于之前已经加过的资源，无需追加
232 0 : continue;
233 : }
234 :
235 0 : HCCL_INFO("unpack BUFFER pos=%u, dto %s", pos, dto.Describe().c_str());
236 0 : if (dto.size == 0) { // size为0，则为 remote 空buffer
237 0 : HCCL_INFO("unpack nullptr, pos=%u", pos);
238 0 : bufferVec.push_back(nullptr);
239 0 : rmtRmaBufferVec.push_back(nullptr);
240 : } else { // size非0，则构造一个remote buffer
241 0 : bufferVec.push_back(make_unique<RemoteUbRmaBuffer>(rdmaHandle, dto));
242 0 : rmtRmaBufferVec.push_back(bufferVec.back().get());
243 0 : HCCL_INFO("unpack buffer pos=%u, rmtRmaBuffer=%s", pos, bufferVec.back()->Describe().c_str());
244 : }
245 0 : }
246 0 : }
247 :
248 0 : bool UrmaDirectTransport::RecvDataProcess()
249 : {
250 0 : HCCL_INFO("RecvDataProcess: link=%s, size=%llu, exchangeDataSize=%u", GetLinkDescInfo().c_str(), recvData.size(),
251 : exchangeDataSize);
252 0 : BinaryStream binaryStream(recvData);
253 0 : HandshakeMsgUnpack(binaryStream);
254 0 : RmtBufferVecUnpackProc(bufferNum, binaryStream, rmtBufferVec);
255 0 : return ConnVecUnpackProc(binaryStream);
256 0 : }
257 :
258 0 : bool UrmaDirectTransport::IsConnsReady()
259 : {
260 0 : for (u32 i = 0; i < connNum; i++) {
261 0 : if (commonLocRes.connVec[i]->GetStatus() != RmaConnStatus::READY) {
262 0 : return false;
263 : }
264 : }
265 0 : HCCL_INFO("conns are ready.");
266 0 : return true;
267 : }
268 :
269 0 : void UrmaDirectTransport::SendFinish()
270 : {
271 0 : HCCL_INFO("start send Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
272 0 : sendFinishMsg = std::vector<char>(FINISH_MSG, FINISH_MSG + FINISH_MSG_SIZE);
273 0 : socket->SendAsync(sendFinishMsg.data(), FINISH_MSG_SIZE);
274 0 : HCCL_INFO("end send Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
275 0 : }
276 :
277 0 : void UrmaDirectTransport::RecvFinish()
278 : {
279 0 : recvFinishMsg.resize(FINISH_MSG_SIZE);
280 0 : HCCL_INFO("start recv Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
281 0 : socket->RecvAsync(reinterpret_cast<u8 *>(recvFinishMsg.data()), FINISH_MSG_SIZE);
282 0 : HCCL_INFO("end recv Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
283 0 : }
284 :
285 0 : TransportStatus UrmaDirectTransport::GetStatus()
286 : {
287 0 : if (baseStatus == TransportStatus::READY) {
288 0 : return baseStatus;
289 0 : } else if (baseStatus == TransportStatus::INIT) {
290 0 : urmaStatus = UrmaStatus::INIT;
291 : }
292 0 : if (!IsSocketReady()) {
293 0 : return baseStatus;
294 : }
295 0 : switch (urmaStatus) {
296 0 : case UrmaStatus::INIT:
297 0 : urmaStatus = UrmaStatus::SOCKET_OK;
298 0 : baseStatus = TransportStatus::SOCKET_OK;
299 0 : break;
300 0 : case UrmaStatus::SOCKET_OK:
301 0 : if (IsResReady()) {
302 0 : urmaStatus = UrmaStatus::SEND_DATA;
303 0 : SendExchangeData();
304 : }
305 0 : break;
306 0 : case UrmaStatus::SEND_DATA:
307 0 : RecvExchangeData();
308 0 : urmaStatus = UrmaStatus::RECV_DATA;
309 0 : break;
310 0 : case UrmaStatus::RECV_DATA:
311 0 : if (RecvDataProcess()) { // 收消息中，如果设置到connection的建链，则需要发送 finish
312 0 : urmaStatus = UrmaStatus::PROCESS_DATA;
313 : } else { // 不需要发送finish，则将transport状态调整为 ready
314 0 : urmaStatus = UrmaStatus::RECV_FIN;
315 0 : SetBaseStatusReady();
316 : }
317 0 : break;
318 0 : case UrmaStatus::PROCESS_DATA:
319 0 : if (IsConnsReady()) {
320 0 : urmaStatus = UrmaStatus::CONN_OK;
321 0 : SendFinish();
322 : }
323 0 : break;
324 0 : case UrmaStatus::CONN_OK:
325 0 : RecvFinish();
326 0 : urmaStatus = UrmaStatus::SEND_FIN;
327 0 : break;
328 0 : case UrmaStatus::SEND_FIN:
329 0 : urmaStatus = UrmaStatus::RECV_FIN;
330 0 : SetBaseStatusReady();
331 0 : break;
332 0 : default:
333 0 : break;
334 : }
335 0 : return baseStatus;
336 : }
337 :
338 : } // namespace Hccl
|
