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 HCCLV2_SOCKET_MANAGER_H
      12              : #define HCCLV2_SOCKET_MANAGER_H
      13              : 
      14              : #include <string>
      15              : #include <vector>
      16              : #include <set>
      17              : #include <unordered_map>
      18              : #include <memory>
      19              : 
      20              : #include "../../unified_platform/resource/socket/socket.h"
      21              : #include "virtual_topo.h"
      22              : #include "socket_config.h"
      23              : #include "env_func.h"
      24              : #include "orion_adapter_hccp.h"
      25              : 
      26              : namespace Hccl {
      27              : 
      28              : using RankIpPortMap = std::unordered_map<u32, std::unordered_map<IpAddress, u32>>;
      29              : using RankIpPortMapPtr = std::shared_ptr<RankIpPortMap>;
      30              : 
      31              : class CommunicatorImpl;
      32              : class SocketManager {
      33              : public:
      34            1 :     SocketManager() = default;
      35              :     SocketManager(
      36              :         const CommunicatorImpl& communicator, u32 localRank, u32 devicePhyId, u32 deviceLogicId,
      37              :         std::function<shared_ptr<Socket>(
      38              :             IpAddress& localIpAddress, IpAddress& remoteIpAddress, u32 listenPort, SocketHandle socketHandle,
      39              :             const std::string& tag, SocketRole socketRole, NicType nicType)>
      40              :             socketProducer
      41              :         = nullptr);
      42              : 
      43              :     SocketManager(u32 localRank, u32 devicePhyId, u32 deviceLogicId, const std::string& socketTag);
      44              : 
      45              :     void
      46              :     SetDeviceServerListenPortMap(const std::unordered_map<u32, std::unordered_map<IpAddress, u32>>& rankListenPortMap);
      47              : 
      48              :     std::unordered_map<u32, std::unordered_map<IpAddress, u32>>
      49              :     GetSubCommDeviceServerListenPortMap(const std::vector<u32>& rankIds) const;
      50              : 
      51              :     u32 GetDeviceListenPort(const u32& rankId, const IpAddress& ipAddress);
      52              : 
      53              :     void BatchCreateSockets(const vector<LinkData>& links);
      54              :     void ServerListen(const SocketConfig& socketConfig);
      55              :     void ConnectSockets(const SocketConfig& socketConfig);
      56              : 
      57              :     void BatchCreateSockets(const SocketConfig& socketConfig);
      58              : 
      59              :     void ServerInit(PortData& localPort);
      60              : 
      61              :     static void ServerInitAll(NewRankInfo& rankInfo);
      62              : 
      63              :     void ServerDeInit(PortData& localPort) const;
      64              : 
      65              :     Socket* CreateConnectedSocket(const SocketConfig& socketConfig);
      66              : 
      67              :     Socket* GetConnectedSocket(const SocketConfig& socketConfig) const;
      68              : 
      69              :     bool CheckServerPortListening(const PortData& portData, const uint32_t port) const;
      70              : 
      71              :     // 供通信域初始化阶段登记已抢占的 host 监听 socket，供算子下发阶段复用，避免跨阶段端口竞争
      72              :     bool RegisterHostListenSocket(const PortData& portData, std::shared_ptr<Socket> listenSocket);
      73              : 
      74              :     void DestroyAll();
      75              : 
      76              :     void AddWhiteList(PortData& localPort, vector<RaSocketWhitelist>& wlistInfoVec) const;
      77              : 
      78              :     bool DelWhiteList(PortData& localPort, vector<RaSocketWhitelist>& wlistInfoVec) const;
      79              : 
      80              :     ~SocketManager();
      81              : 
      82              :     SocketManager(const SocketManager& socketManager) = delete;
      83              : 
      84              :     SocketManager& operator=(const SocketManager& socketManager) = delete;
      85              : 
      86              : private:
      87              :     void PrepareLinkAndServerInit(const SocketConfig& socketConfig);
      88              :     void BatchServerInit(const vector<LinkData>& links);
      89              :     void BatchAddWhiteList(const vector<LinkData>& links);
      90              :     void BatchCreateConnectedSockets(const vector<LinkData>& links);
      91              :     void AddWhiteList(const SocketConfig& socketConfig);
      92              :     const CommunicatorImpl* comm;
      93              :     static std::unordered_map<PortData, shared_ptr<Socket>>& GetServerSocketMap();
      94              :     u32 localRank;
      95              :     u32 devicePhyId;
      96              :     u32 deviceLogicId_;
      97              :     std::unordered_map<u32, std::unordered_map<IpAddress, u32>> rankListenPortMap_{};
      98              :     std::function<shared_ptr<Socket>(
      99              :         IpAddress& localIpAddress, IpAddress& remoteIpAddress, u32 listenPort, SocketHandle socketHandle,
     100              :         const std::string& tag, SocketRole socketRole, NicType nicType)>
     101              :         socketProducer
     102           12 :         = [](IpAddress& localIpAddress, IpAddress& remoteIpAddress, u32 listenPort, SocketHandle socketHandle,
     103              :              const std::string& tag, SocketRole socketRole, NicType nicType) -> shared_ptr<Socket> {
     104              :         auto tmpSocket = std::make_shared<Socket>(
     105           12 :             socketHandle, localIpAddress, listenPort, remoteIpAddress, tag, socketRole, nicType);
     106           36 :         HCCL_INFO("create socket with role %u", static_cast<u32>(socketRole));
     107           12 :         return tmpSocket;
     108            0 :     };
     109              : 
     110              :     std::unordered_map<SocketConfig, shared_ptr<Socket>> connectedSocketMap;
     111              :     std::unordered_map<PortData, vector<RaSocketWhitelist>> socketWlistMap{};
     112              : 
     113              :     Socket* GetServerListenSocket(const PortData& localPort) const;
     114              :     std::set<LinkData> availableLinks;
     115              : 
     116              :     std::string socketTag_{};
     117              :     static std::mutex socketLock;
     118              : };
     119              : 
     120              : } // namespace Hccl
     121              : 
     122              : #endif // HCCLV2_SOCKET_MANAGER_H