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