in_memory_handler.h

 
#pragma once
#include <condition_variable>
#include <string>
 
#include "communicator.h"
 
namespace xgboost {
namespace collective {
 
class InMemoryHandler {
 public:
  InMemoryHandler() = default;
 
  explicit InMemoryHandler(int worldSize) : world_size_{worldSize} {}
 
  void Init(int world_size, int rank);
 
  void Shutdown(uint64_t sequence_number, int rank);
 
  void Allgather(char const* input, std::size_t bytes, std::string* output,
                 std::size_t sequence_number, int rank);
 
  void Allreduce(char const* input, std::size_t bytes, std::string* output,
                 std::size_t sequence_number, int rank, DataType data_type, Operation op);
 
  void Broadcast(char const* input, std::size_t bytes, std::string* output,
                 std::size_t sequence_number, int rank, int root);
 
 private:
  template <class HandlerFunctor>
  void Handle(char const* input, std::size_t size, std::string* output, std::size_t sequence_number,
              int rank, HandlerFunctor const& functor);
 
  int world_size_{};                    
  int received_{};                      
  int sent_{};                          
  std::string buffer_{};                
  uint64_t sequence_number_{};          
  mutable std::mutex mutex_;            
  mutable std::condition_variable cv_;  
};
 
}  // namespace collective
}  // namespace xgboost