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 : #ifndef HCCL_SOCKET_MANAGER_H
12 : #define HCCL_SOCKET_MANAGER_H
13 :
14 : #include <map>
15 : #include <vector>
16 : #include <unordered_map>
17 : #include <string>
18 : #include <memory>
19 : #include <mutex>
20 : #include <atomic>
21 : #include <functional>
22 : #include <hccl/hccl_types.h>
23 : #include "hccl_common.h"
24 : #include "hccl_ip_address.h"
25 : #include "hccl_socket.h"
26 : #include "common.h"
27 :
28 : namespace hccl {
29 : class PortInfo {
30 : public:
31 206 : PortInfo(const HcclIpAddress &ip, u32 listenPort)
32 206 : : ip(ip), listenPort(listenPort)
33 205 : {}
34 296 : ~PortInfo()
35 296 : {}
36 :
37 : bool operator==(const PortInfo &portInfo) const
38 : {
39 : return listenPort == portInfo.listenPort && ip == portInfo.ip;
40 : }
41 :
42 : bool operator!=(const PortInfo &portInfo) const
43 : {
44 : return !(portInfo == *this);
45 : }
46 :
47 272 : bool operator<(const PortInfo &portInfo) const
48 : {
49 272 : if (ip < portInfo.ip) {
50 36 : return true;
51 : }
52 236 : if (portInfo.ip < ip) {
53 12 : return false;
54 : }
55 224 : return listenPort < portInfo.listenPort;
56 : }
57 :
58 : bool operator>(const PortInfo &portInfo) const
59 : {
60 : return portInfo < *this;
61 : }
62 :
63 : bool operator<=(const PortInfo &portInfo) const
64 : {
65 : return !(portInfo < *this);
66 : }
67 :
68 : bool operator>=(const PortInfo &portInfo) const
69 : {
70 : return !(*this < portInfo);
71 : }
72 :
73 : HcclIpAddress ip;
74 : u32 listenPort;
75 : };
76 : using NicHandleInfo = struct NicHandleInfoDef {
77 : HcclIpAddress ip;
78 : SocketHandle nicSocketHandle;
79 : NicType socketType;
80 :
81 : NicHandleInfoDef() : ip(), nicSocketHandle(nullptr), socketType(NicType::DEVICE_NIC_TYPE)
82 : {}
83 : };
84 :
85 : class HcclSocketManager {
86 : public:
87 : explicit HcclSocketManager(NICDeployment nicDeployment, s32 deviceLogicId, u32 devicePhyId, u32 userRank);
88 : virtual ~HcclSocketManager();
89 :
90 : HcclResult AddWhiteList(const std::string &commTag,
91 : const HcclNetDevCtx netDevCtx,
92 : HcclRankLinkInfo remoteRankInfo);
93 : void DestroySockets(const std::string &commTag);
94 : void DestroySockets(const std::string &commTag, u32 rank);
95 : HcclResult CreateSockets(const std::string &commTag, bool isInterLink,
96 : const HcclNetDevCtx netDevCtx,
97 : const std::map<u32, HcclRankLinkInfo> &dstServerMap,
98 : const std::map<u32, HcclRankLinkInfo> &dstClientMap,
99 : std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > &serverSocketsMap,
100 : std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > &clientSocketsMap,
101 : bool isSupportReuse = false, bool isWaitEstablished = true);
102 : HcclResult GetListenPortByIp(
103 : const NICDeployment nicDeployment, const HcclIpAddress &ipAddr, std::set<u32> &listenedPort);
104 :
105 : void GetSocketsByRankIP(const std::string &commTag, u32 remoteRank, const HcclIpAddress &remoteIp,
106 : u32 socketsPerLink, std::vector<std::shared_ptr<HcclSocket> > &ipSockets, u32 &gotLinkNum);
107 : void GetSocketsByRankIP(const HcclIpAddress &remoteIp, u32 socketsPerLink,
108 : std::vector<std::shared_ptr<HcclSocket>> &rankSockets, std::vector<std::shared_ptr<HcclSocket>> &ipSockets,
109 : u32 &gotLinkNum);
110 :
111 : HcclResult ServerInit(const HcclNetDevCtx netDevCtx, u32 port);
112 : HcclResult ServerDeInit(const HcclNetDevCtx netDevCtx, u32 port);
113 :
114 : HcclResult CreateSingleLinkSocket(const std::string &commTag,
115 : const HcclNetDevCtx netDevCtx,
116 : HcclRankLinkInfo rmtRank,
117 : std::vector<std::shared_ptr<HcclSocket> > &connectSockets,
118 : bool isWaitEstablished = true,
119 : bool isSupportReuse = false,
120 : s32 timeout = 0, uint32_t connectMode = 0);
121 :
122 : HcclResult WaitLinksEstablishCompleted(HcclSocketRole localRole,
123 : std::map <u32, std::vector<std::shared_ptr<HcclSocket> > > &socketsMap, std::map<u32, u32> &dstRankToUserRank,
124 : const RankInfo &loaclRankInfo, const RankInfo &remoteRankInfo, const HcclNetDevCtx &netDevCtx);
125 : void DestroySockets();
126 :
127 : void AbortAndDeleteSocket(const std::string &commTag, HcclSocketRole role,
128 : const std::map <u32, std::vector<std::shared_ptr<HcclSocket> > > &socketsMap);
129 :
130 : HcclResult SetStopFlag(bool value);
131 : bool GetStopFlag();
132 0 : HcclResult WaitLinkEstablish(std::shared_ptr<HcclSocket> socket, std::function<bool()> needStop = []() { return false; },
133 : s32 timeout = 0);
134 : HcclResult ServerDeInit(const HcclIpAddress& localIp, u32 port);
135 : private:
136 : HcclResult AddWhiteList(const std::string &commTag, bool isInterLink, NicType socketType,
137 : const HcclIpAddress &localIp, const std::map<u32, HcclRankLinkInfo> &whiteListMap, uint32_t connectMode = 0);
138 : HcclResult DelWhiteList(const std::string &commTag);
139 : HcclResult CreateSockets(const std::string &commTag, bool isInterLink, const HcclNetDevCtx netDevCtx,
140 : NicType socketType, HcclSocketRole localRole, const HcclIpAddress &localIp,
141 : const HcclRankLinkInfo &remoteLinkInfo, std::vector<std::shared_ptr<HcclSocket> > &ipSockets,
142 : bool isSupportReuse, uint32_t connectMode = 0);
143 : HcclResult CreateSockets(const std::string &commTag, bool isInterLink, const HcclNetDevCtx netDevCtx,
144 : NicType socketType, HcclSocketRole localRole, const HcclIpAddress &localIp,
145 : const std::map<u32, HcclRankLinkInfo> &remoteInfos,
146 : std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > &socketsMap,
147 : std::map<u32, u32> &dstRankToUserRank, bool isSupportReuse, uint32_t connectMode = 0);
148 : void DestroySockets(std::vector<std::shared_ptr<HcclSocket> > rankSockets);
149 : void TransformSocketStatus(HcclSocketStatus status, std::string &stringStatus) const;
150 : void PrintSocketsInfo(const std::string &localRole,
151 : u32 rank, std::vector<std::shared_ptr<HcclSocket> > ipSockets, std::string &sTlsStatus) const;
152 : void PrintErrorConnectionInfo(HcclSocketRole localRole,
153 : std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > &rankSocketsMap,
154 : std::map<u32, u32> &dstRankToUserRank, TlsStatus &tlsStatus) const;
155 : void PrintErrorConnection(HcclSocketRole localRole,
156 : std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > &rankSocketsMap,
157 : std::map<u32, u32> &dstRankToUserRank, TlsStatus &tlsStatus) const;
158 : u32 GetConnLimit(NicType socketType);
159 : std::string MakeUniqueConnTag(const std::string &commTag, bool isInterLink, u32 rank, u32 indexForLink);
160 : HcclResult ConstructWhiteList(const std::string &commTag,
161 : bool isInterLink, NicType socketType,
162 : const HcclRankLinkInfo &dstRankLinkInfo, std::vector<SocketWlistInfo> &wlistInfosVec, uint32_t connectMode = 0);
163 : void SaveWhiteListInfo(const std::string &commTag, std::shared_ptr<HcclSocket> &socket,
164 : const std::vector<SocketWlistInfo> wlistInfos);
165 : HcclResult ConstructSockets(const std::string &commTag, bool isInterLink, const HcclNetDevCtx netDevCtx,
166 : u32 socketsPerLink, NicType socketType, u32 dstRank, const HcclIpAddress &remoteIp, u32 remotePort,
167 : const HcclIpAddress &localIp, HcclSocketRole localRole, std::vector<std::shared_ptr<HcclSocket>> &socketList,
168 : uint32_t connectMode = 0);
169 : void SaveSockets(const std::string &commTag, u32 remoteRank, const HcclIpAddress &remoteIp,
170 : std::vector<std::shared_ptr<HcclSocket> > &ipSockets);
171 :
172 : HcclResult WaitLinksEstablishCompleted(HcclSocketRole localRole,
173 : std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > &rankSocketsMap, s32 timeout = 0);
174 :
175 : void AddIpQueue(RankInfo &localRankInfo, RankInfo &remoteRankInfo, NicType nicType, s32 deviceLogicId);
176 : NICDeployment nicDeployment_;
177 : s32 deviceLogicId_;
178 : u32 devicePhyId_;
179 : u32 userRank_;
180 :
181 : // 后继这个放在HcclSocket中管理
182 : std::map<std::string, std::map<std::shared_ptr<HcclSocket>, std::vector<SocketWlistInfo>>>
183 : wlistInfosMap_;
184 : std::map<std::string, std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > > commSocketsMap_;
185 : std::mutex wlistMapMutex_;
186 : std::mutex socketsMapMutex_;
187 :
188 : static std::mutex serverMapMutex_;
189 : static std::map<PortInfo, std::shared_ptr<HcclSocket>> serverSocketMap_;
190 : static std::map<PortInfo, Referenced> serverSocketRefMap_;
191 :
192 : std::atomic<bool> stopFlag_{false};
193 : };
194 :
195 : using IntraExchanger = struct IntraExchangerDef {
196 : std::map<u32, std::vector<std::shared_ptr<HcclSocket> > > socketsMap;
197 : std::shared_ptr<HcclSocketManager> socketManager;
198 27 : IntraExchangerDef() : socketsMap(), socketManager()
199 27 : {}
200 : };
201 :
202 : using RegisterDetectCallBack =
203 : void (*)(RankInfo &localRankInfo, RankInfo &remoteRankInfo, NicType nicType,
204 : s32 deviceLogicId);
205 : #ifdef __cplusplus
206 : extern "C" {
207 : #endif // __cplusplus
208 : void DetectCallBack(RegisterDetectCallBack p1);
209 : #ifdef __cplusplus
210 : }
211 : #endif // __cplusplus
212 :
213 : } // namespace hccl
214 : #endif /* * HCCL_SOCKET_MANAGER_H */
|