Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2026 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_RANK_INFO_DETECT_SERVICE_H
      12              : #define HCCLV2_RANK_INFO_DETECT_SERVICE_H
      13              : 
      14              : #include "socket.h"
      15              : #include <unordered_map>
      16              : #include "new_rank_info.h"
      17              : #include "rank_table_info.h"
      18              : #include "json_parser.h"
      19              : #include "internal_exception.h"
      20              : #include "timeout_exception.h"
      21              : #include "socket_exception.h"
      22              : #include "socket_agent.h"
      23              : 
      24              : namespace Hccl {
      25              : 
      26              : class RankInfoDetectService {
      27              : public:
      28           29 :     RankInfoDetectService(
      29              :         u32 devPhyId, std::shared_ptr<Socket> serverSocket, std::string identifier, vector<RaSocketWhitelist> wlistInfo)
      30           29 :         : devPhyId_(devPhyId),
      31           29 :           serverSocket_(serverSocket),
      32           29 :           hostIp_(serverSocket_->GetLocalIp()),
      33           29 :           identifier_(identifier),
      34           58 :           wlistInfo_(wlistInfo)
      35           29 :     {}
      36              :     ~RankInfoDetectService();
      37              : 
      38              :     void Setup();
      39              : 
      40              : private:
      41              :     u32 devPhyId_{0};
      42              :     shared_ptr<Socket> serverSocket_{nullptr};
      43              :     std::unordered_map<std::string, std::shared_ptr<Socket>> connSockets_{};
      44              :     IpAddress hostIp_{};
      45              :     RankTableInfo rankTable_{};
      46              :     std::string failedAgentIdList_{};
      47              :     u32 currentStep_{0};
      48              :     std::string identifier_{};
      49              :     vector<RaSocketWhitelist> wlistInfo_{};
      50              : 
      51              :     void GetConnections();
      52              :     void GetRankTable();
      53              :     void BroadcastRankTable();
      54              :     void Disconnect();
      55              :     void TearDown();
      56              : 
      57              :     bool RecvRemoteAgentId(SocketAgent& connSocketAgent, std::string& agentId);
      58              :     bool RecvRemoteRankSize(SocketAgent& connSocketAgent, u32& rankSize);
      59              :     void RecvRankInfoMsg(SocketAgent& connSocketAgent, vector<char>& rankInfoMsg);
      60              :     void SortRankTable();
      61              :     void ParseRankTable(vector<char>& rankInfoMsg);
      62              : 
      63              :     // 异常流程处理方法
      64              :     void FailedConnectionAgentIdString(u32 rankSize);
      65              : 
      66              :     // 校验相关方法
      67              :     bool RecvAndVerifyRemoteAgentIdAndRankSize(
      68              :         std::shared_ptr<Socket> connSocket, u32& expectedSocketNum, u32& previousRankSize);
      69              :     bool VerifyRemoteRankSize(u32& previousRankSize, u32 remoteRankSize) const;
      70              : 
      71              :     // DFX相关方法
      72              :     void DisplayConnectingStatus(u32 totalSockets, u32 waitSockets);
      73              :     void DisplayConnectedRanks();
      74              : };
      75              : 
      76              : } // namespace Hccl
      77              : #endif // HCCLV2_RANK_INFO_DETECT_SERVICE_H