Add Allgather to collective communicator (#8765)

* Add Allgather to collective communicator
This commit is contained in:
Rong Ou
2023-02-08 19:31:22 -08:00
committed by GitHub
parent 48cefa012e
commit cbf98cb9c6
14 changed files with 187 additions and 4 deletions

View File

@@ -46,6 +46,25 @@ class FederatedClient {
}()},
rank_{rank} {}
std::string Allgather(std::string const &send_buffer) {
AllgatherRequest request;
request.set_sequence_number(sequence_number_++);
request.set_rank(rank_);
request.set_send_buffer(send_buffer);
AllgatherReply reply;
grpc::ClientContext context;
context.set_wait_for_ready(true);
grpc::Status status = stub_->Allgather(&context, request, &reply);
if (status.ok()) {
return reply.receive_buffer();
} else {
std::cout << status.error_code() << ": " << status.error_message() << '\n';
throw std::runtime_error("Allgather RPC failed");
}
}
std::string Allreduce(std::string const &send_buffer, DataType data_type,
ReduceOperation reduce_operation) {
AllreduceRequest request;