aggregator.h

 
#pragma once
#include <xgboost/data.h>
 
#include <limits>
#include <string>
#include <utility>
#include <vector>
 
#include "communicator-inl.h"
 
namespace xgboost {
namespace collective {
 
template <typename Function>
void ApplyWithLabels(MetaInfo const& info, void* buffer, size_t size, Function&& function) {
  if (info.IsVerticalFederated()) {
    // We assume labels are only available on worker 0, so the calculation is done there and result
    // broadcast to other workers.
    std::string message;
    if (collective::GetRank() == 0) {
      try {
        std::forward<Function>(function)();
      } catch (dmlc::Error& e) {
        message = e.what();
      }
    }
 
    collective::Broadcast(&message, 0);
    if (message.empty()) {
      collective::Broadcast(buffer, size, 0);
    } else {
      LOG(FATAL) << &message[0];
    }
  } else {
    std::forward<Function>(function)();
  }
}
 
template <typename T>
T GlobalMax(MetaInfo const& info, T value) {
  if (info.IsRowSplit()) {
    collective::Allreduce<collective::Operation::kMax>(&value, 1);
  }
  return value;
}
 
template <typename T>
void GlobalSum(MetaInfo const& info, T* values, size_t size) {
  if (info.IsRowSplit()) {
    collective::Allreduce<collective::Operation::kSum>(values, size);
  }
}
 
template <typename Container>
void GlobalSum(MetaInfo const& info, Container* values) {
  GlobalSum(info, values->data(), values->size());
}
 
template <typename T>
T GlobalRatio(MetaInfo const& info, T dividend, T divisor) {
  std::array<T, 2> results{dividend, divisor};
  GlobalSum(info, &results);
  std::tie(dividend, divisor) = std::tuple_cat(results);
  if (divisor <= 0) {
    return std::numeric_limits<T>::quiet_NaN();
  } else {
    return dividend / divisor;
  }
}
 
}  // namespace collective
}  // namespace xgboost