bootstrap.h

#pragma once

#include <string.h>
#include "base.h"
#include "socket.h"
#include "bootstrap_utils.h"
#include "bootstrap_net.h"
#include "thread_pool.h"
#include "ipc_socket.h"

namespace sccl {
namespace hardware {
namespace topology {
namespace bootstrap {

typedef sccl::hardware::net::ipc_socket::scclIpcSocket_t scclIpcSocket_t;

///////////////////////////////////// 用于初始化时的功能函数 //////////////////////////////////////////
scclResult_t bootstrapGetUniqueId(BootstrapHandle_t* handle);
scclResult_t bootstrapCreateRoot(BootstrapHandle_t* handle);

////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// 定义结构体 scclRankInfo，用于存储每个rank的通信节点的信息
// 每个rank必定选定一个GPU、一个RDMA设备、一个CPU，作为通信节点
typedef struct scclRankInfo {
    struct {
        scclSocket_t listen_sock; // 监听套接字
    } cpu;                        // CPU节点
    struct {
        int dev;           // NVML设备编号
        char name[8];      // 设备名称
        char gcn[8];       // GCN架构名称
        int compCap;       // CUDA计算能力
        int64_t pciBusId;  // PCI总线ID以int64_t格式表示
        char pciPath[128]; // PCI设备在/sys中的路径。
    } gpu;                 // GPU节点
    struct {
        int count;          // 网卡数量
        char name[8];       // 主要用于日志记录。
        char pciPath[128];  // PCI设备在/sys中的路径。
        uint64_t guid;      // NIC芯片的唯一标识符。对于具有多个PCI功能（物理或虚拟）的卡非常重要。
        uint8_t ptrSupport; // [SCCL_PTR_HOST|SCCL_PTR_CUDA|SCCL_PTR_DMABUF]
        int speed;          // 端口速度，单位为Mbps。
        int port;           // 端口号。
        float latency;      // 网络延迟
        int maxComms;       // 可以创建的最大通信数量
        int maxRecvs;       // 最大分组接收数量。
    } net;                  // 网络节点

    int rank      = -1; // 当前节点的全局排名
    int localRank = -1; // 当前节点在本地计算节点中的排名

    uint64_t hostHash = 0; // 主机哈希值
    uint64_t pidHash  = 0; // 进程 ID 哈希值
} scclRankInfo_t;

// 所有节点的信息
typedef struct scclRankPhysSet {
    // 构造函数声明
    scclRankPhysSet(int nRanks);

public:
    std::vector<scclRankInfo_t> rank_info_vec;
} scclRankPhysSet_t;

// BootstrapComm 结构体定义，用于存储引导通信信息
typedef struct BootstrapComm {
    void init(int rank, int nRanks, int localRank, int nLocalRanks);
    void destroy();

public:
    scclNet_t* scclNet               = nullptr;
    scclRankPhysSet_t* rank_phys_set = nullptr;

    cpu_set_t cpuAffinity; // CPU亲和性
    int rank        = -1;  // 当前节点的全局排名
    int nRanks      = 0;   // 总的节点数量
    int localRank   = -1;  // 当前节点在本地计算节点中的排名
    int nLocalRanks = 0;   // 本地计算节点中的节点总数
    int interRank   = -1;  // 整个节点在全部节点中的位置
    int nInterRanks = 0;   // 全局拥有节点的个数

    int hipDev    = -1; // CUDA 设备 ID
    int deviceCnt = 0;  // 设备数量

    uint64_t magic; // 魔术数，用于验证结构体
} BootstrapComm_t;

///////////////////////////////////// 用于初始化时的类 //////////////////////////////////////////
class Bootstrap {
public:
    Bootstrap(const BootstrapHandle_t*, int rank, int nRanks);
    virtual ~Bootstrap();

    // 初始化bootstrap通信环境
    scclResult_t init(BootstrapComm_t* bootstrap_comm);

    // 实现跨节点的AllGather通信操作
    scclResult_t bootstrapAllGather(const void* src_data, void* dst_data, int data_size) const;

private:
    // 执行根节点的聚集和广播操作
    scclResult_t bootstrapRootGatherAndBroadcast(BootstrapNodeBasic_t* send_data_basic);

    // 初始化节点通信信息
    scclResult_t bootstrapCommInitNodeInfo(scclNet_t* scclNet, scclRankInfo_t* rank_info);

public:
    int rank, nRanks;           // 初始化阶段获取MPI的值
    int localRank, nLocalRanks; // 通过bootstrapRootGatherAndBroadcast函数确定值
    int interRank, nInterRanks; // 整个节点在全部节点中的位置

private:
    // TODO: 用于控制套接字终端的变量，目前不知道在哪里使用
    volatile uint32_t* abortFlag; // 中止标志，非阻塞套接字设置

    // 外部传入的0号节点的基础信息
    const BootstrapHandle_t* root_handle;
    // 节点内所有进程的基础ip信息
    BootstrapNodeBasic_t* all_node_basic = nullptr;

    // 初始化标志
    bool socketInitDone;
    // 互斥锁，用于保护初始化过程的线程安全
    pthread_mutex_t bootstrapMutex = PTHREAD_MUTEX_INITIALIZER;
    pthread_cond_t bootstrapCond   = PTHREAD_COND_INITIALIZER;

    // 节点内通信的类
    scclIpcSocket_t* ipcsocket = nullptr; // 指向scclIpcSocket类实例的指针，初始值为nullptr
};

// 打印唯一的拓扑信息
scclResult_t printRankInfo(const std::string& prefix, scclRankInfo_t* info);

} // namespace bootstrap
} // namespace topology
} // namespace hardware
} // namespace sccl