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 TOPOINFO_EXCHANGE_DISPATCHER_H
      12              : #define TOPOINFO_EXCHANGE_DISPATCHER_H
      13              :  
      14              : #include <map>
      15              : #include <atomic>
      16              : #include <vector>
      17              : #include <mutex>
      18              : #include <thread>
      19              : #include <climits>
      20              : #include <condition_variable>
      21              : 
      22              : #include "adapter_hccp_common.h"
      23              : #include "externalinput_pub.h"
      24              : #include "hccl_socket.h"
      25              : #include "hccl_common.h"
      26              : #include "topoinfo_exchange_server.h"
      27              : #include "../json_utils.h"
      28              : 
      29              : namespace hccl {
      30              : class TopoInfoExchangeDispather {
      31              :     // avoid the struct name pollution hccl namespace, so use the struct in class
      32              : public:
      33              :     struct SendState {
      34              :         u32 rankId;
      35              :         u32 header;
      36              :         u32 identify          = UINT_MAX;    // 默认UINT_MAX时，不发送identify
      37              :         size_t headerLen       = sizeof(u32); // the header need to send
      38              :         size_t headerSended    = 0;           // the header have sended length
      39              :         size_t bodyLen         = 0;           // the whole data length
      40              :         size_t bodySended      = 0;           // the data have sended
      41              :         size_t identifyLen    = sizeof(u32); // the identify need to send (MasterInfo mode)
      42              :         size_t identifySended = 0;           // the identify have sended
      43              :         void *data;                           // data pointer
      44              :         bool firstSendFlag_ =  true;
      45              :  
      46              :         HcclResult Send(std::shared_ptr<HcclSocket> socket);
      47              :         HcclResult SendHeader(std::shared_ptr<HcclSocket> socket);
      48              :         HcclResult SendBody(std::shared_ptr<HcclSocket> socket);
      49              :         HcclResult SendIdentify(std::shared_ptr<HcclSocket> socket);
      50              :         HcclResult SendHelper(std::shared_ptr<HcclSocket> socket, char *buf, size_t dataLen, size_t &sendedLen);
      51            0 :         bool IsOk()
      52              :         {
      53            0 :             return bodyLen != 0 && headerSended == headerLen && bodySended == bodyLen;
      54              :         }
      55              :     };
      56              : 
      57              :     struct FdContext {
      58              :         std::shared_ptr<HcclSocket> socket;
      59              :         SendState txState;
      60              :     };
      61              : 
      62              :     using WorkerTask = std::function<HcclResult(void)>;
      63              :  
      64              : public:
      65              :     static constexpr u32 DEFAULT_THREAD_NUM = 1;
      66              :     static constexpr u32 MAX_THREAD_NUM = 4;
      67              :     static constexpr s32 INVALID_EPOLL_EVENT_FD = -1;
      68              :     static constexpr s32 EPOLL_TIMEOUT_MS = 100; // 100ms
      69              :     static constexpr s32 LAST_EPOLL_TIMEOUT_MS = 5; // 5ms
      70              :     static constexpr s32 RANK_CAPACITY_PER_THREAD = 512;
      71              :  
      72           16 :     explicit TopoInfoExchangeDispather(TopoInfoExchangeServer *topoInfoExchangeServer,
      73              :         u32 threadNum = DEFAULT_THREAD_NUM)
      74           16 :         : topoInfoExchangeServer_(topoInfoExchangeServer), threadNum_(threadNum)
      75              :     {
      76           16 :     }
      77              :     ~TopoInfoExchangeDispather();
      78              :  
      79              :     HcclResult BroadcastRankTable(const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets,
      80              :         const RankTable_t &clusterInfo, const std::string &failedAgentIdList);
      81              :     HcclResult BroadcastGroupLeaderInfo(const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets,
      82              :         const GroupLeader_t &leaderInfo);
      83              :  
      84              : private:
      85              :     void InitWorkerThread();
      86              :     void WorkerWait(int workId);
      87              :     void WakeWoker();
      88              :     void RunWorkerThread(int workId);
      89              :     bool GetTask(WorkerTask &workTask);
      90              :     HcclResult PrepareResource(const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets,
      91              :         const RankTable_t &clusterInfo, const std::string &failedAgentIdList);
      92              :     HcclResult PrepareLeaderResource(const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets,
      93              :         const GroupLeader_t &leaderInfo);
      94              :     HcclResult SendOnce();
      95              :     HcclResult ProcessOneSendEvent(s32 epollFd, FdHandle &fdHandle);
      96              :     HcclResult ProcessSend();
      97              :     void CleanResource();
      98              :     HcclResult CloseEpollFd();
      99              : 
     100              :     TopoInfoExchangeServer *topoInfoExchangeServer_;
     101              :     u32 threadNum_ = 1;
     102              :     u32 rankNum_   = 0;
     103              :     std::vector<std::thread> workerThreads_;
     104              :     std::queue<WorkerTask> taskQueue_;
     105              :     std::mutex taskQueueMutex_;
     106              : 
     107              :     std::unordered_map<FdHandle, FdContext> fdHandleToFdContextMap_;
     108              :     std::mutex fdHandleMapMutex_;
     109              :     s32 epollFds_ = INVALID_EPOLL_EVENT_FD;
     110              :     std::atomic<u32> sendDoneCount_{0};
     111              : 
     112              :     std::string rankTableJson_;
     113              : 
     114              :     std::mutex wakeMutex_;
     115              :     std::atomic<bool> ready_{false};
     116              :     std::atomic<bool> stop_{false};
     117              :     std::condition_variable wakeManager_;
     118              : };
     119              : } // namespace hccl
     120              : 
     121              : #endif /* TOPOINFO_EXCHANGE_DISPATCHER_H */