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() { return bodyLen != 0 && headerSended == headerLen && bodySended == bodyLen; }
      52              :     };
      53              : 
      54              :     struct FdContext {
      55              :         std::shared_ptr<HcclSocket> socket;
      56              :         SendState txState;
      57              :     };
      58              : 
      59              :     using WorkerTask = std::function<HcclResult(void)>;
      60              : 
      61              : public:
      62              :     static constexpr u32 DEFAULT_THREAD_NUM = 1;
      63              :     static constexpr u32 MAX_THREAD_NUM = 4;
      64              :     static constexpr s32 INVALID_EPOLL_EVENT_FD = -1;
      65              :     static constexpr s32 EPOLL_TIMEOUT_MS = 100;    // 100ms
      66              :     static constexpr s32 LAST_EPOLL_TIMEOUT_MS = 5; // 5ms
      67              :     static constexpr s32 RANK_CAPACITY_PER_THREAD = 512;
      68              : 
      69           16 :     explicit TopoInfoExchangeDispather(
      70              :         TopoInfoExchangeServer* topoInfoExchangeServer, u32 threadNum = DEFAULT_THREAD_NUM)
      71           16 :         : topoInfoExchangeServer_(topoInfoExchangeServer),
      72           16 :           threadNum_(threadNum)
      73           16 :     {}
      74              :     ~TopoInfoExchangeDispather();
      75              : 
      76              :     HcclResult BroadcastRankTable(
      77              :         const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets, const RankTable_t& clusterInfo,
      78              :         const std::string& failedAgentIdList);
      79              :     HcclResult BroadcastGroupLeaderInfo(
      80              :         const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets, const GroupLeader_t& leaderInfo);
      81              : 
      82              : private:
      83              :     void InitWorkerThread();
      84              :     void WorkerWait(int workId);
      85              :     void WakeWoker();
      86              :     void RunWorkerThread(int workId);
      87              :     bool GetTask(WorkerTask& workTask);
      88              :     HcclResult PrepareResource(
      89              :         const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets, const RankTable_t& clusterInfo,
      90              :         const std::string& failedAgentIdList);
      91              :     HcclResult PrepareLeaderResource(
      92              :         const std::map<std::string, std::shared_ptr<HcclSocket>> connectSockets, const GroupLeader_t& leaderInfo);
      93              :     HcclResult SendOnce();
      94              :     HcclResult ProcessOneSendEvent(s32 epollFd, FdHandle& fdHandle);
      95              :     HcclResult ProcessSend();
      96              :     void CleanResource();
      97              :     HcclResult CloseEpollFd();
      98              : 
      99              :     TopoInfoExchangeServer* topoInfoExchangeServer_;
     100              :     u32 threadNum_ = 1;
     101              :     u32 rankNum_ = 0;
     102              :     std::vector<std::thread> workerThreads_;
     103              :     std::queue<WorkerTask> taskQueue_;
     104              :     std::mutex taskQueueMutex_;
     105              : 
     106              :     std::unordered_map<FdHandle, FdContext> fdHandleToFdContextMap_;
     107              :     std::mutex fdHandleMapMutex_;
     108              :     s32 epollFds_ = INVALID_EPOLL_EVENT_FD;
     109              :     std::atomic<u32> sendDoneCount_{0};
     110              : 
     111              :     std::string rankTableJson_;
     112              : 
     113              :     std::mutex wakeMutex_;
     114              :     std::atomic<bool> ready_{false};
     115              :     std::atomic<bool> stop_{false};
     116              :     std::condition_variable wakeManager_;
     117              : };
     118              : } // namespace hccl
     119              : 
     120              : #endif /* TOPOINFO_EXCHANGE_DISPATCHER_H */