#include <xgboost/tree_updater.h>
#include <string>
#include <vector>
#include "../collective/broadcast.h"
#include "../collective/communicator-inl.h"
#include "../common/io.h"
#include "xgboost/json.h"
namespace xgboost::tree {
DMLC_REGISTRY_FILE_TAG(updater_sync);
class TreeSyncher : public TreeUpdater {
public:
explicit TreeSyncher(Context const* tparam) : TreeUpdater(tparam) {}
void Configure(const Args&) override {}
void LoadConfig(Json const&) override {}
void SaveConfig(Json*) const override {}
[[nodiscard]] char const* Name() const override { return "prune"; }
void Update(TrainParam const*, linalg::Matrix<GradientPair>*, DMatrix*,
common::Span<HostDeviceVector<bst_node_t>> ,
const std::vector<RegTree*>& trees) override {
if (collective::GetWorldSize() == 1) return;
std::string s_model;
common::MemoryBufferStream fs(&s_model);
int rank = collective::GetRank();
if (rank == 0) {
for (auto tree : trees) {
tree->Save(&fs);
}
}
fs.Seek(0);
auto rc = collective::Broadcast(ctx_, linalg::MakeVec(s_model.data(), s_model.size()), 0);
SafeColl(rc);
for (auto tree : trees) {
tree->Load(&fs);
}
}
};
XGBOOST_REGISTER_TREE_UPDATER(TreeSyncher, "sync")
.describe("Syncher that synchronize the tree in all distributed nodes.")
.set_body([](Context const* ctx, auto) { return new TreeSyncher(ctx); });
}