pytorch/test/cpp/rpc/test_e2e_process_group.cpp
Wanchao Liang 553ccccc54 [c10d] switch ProcessGroup to be managed by intrusive_ptr (#47343)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/47343

Test Plan: Imported from OSS

Reviewed By: gmagogsfm

Differential Revision: D24723418

Pulled By: wanchaol

fbshipit-source-id: 0463819b96c53b12bdbb3905431110d7b21beb77
2020-11-12 07:36:23 -08:00

46 lines
1.3 KiB
C++

#include <gtest/gtest.h>
#include "e2e_test_base.h"
#include <c10d/ProcessGroupGloo.hpp>
#include <torch/csrc/distributed/rpc/process_group_agent.h>
#include <torch/csrc/distributed/rpc/request_callback_no_python.h>
#include <torch/torch.h>
namespace torch {
namespace distributed {
namespace rpc {
class TestE2EProcessGroup : public TestE2EBase {
protected:
void buildRpcAgent() override {
c10d::ProcessGroupGloo::Options options;
options.devices.push_back(
::c10d::ProcessGroupGloo::createDeviceForHostname(serverAddress));
std::chrono::milliseconds rpcTimeout(30000);
options.timeout = rpcTimeout;
// Initialize server rpc agent.
auto pg = c10::make_intrusive<c10d::ProcessGroupGloo>(
store, 0, numWorkers, options);
rpcAgent = std::make_shared<ProcessGroupAgent>(
"worker",
pg,
std::max(16U, std::thread::hardware_concurrency()),
rpcTimeout,
std::make_unique<RequestCallbackNoPython>());
}
};
// End to end training loop test in C++ so that we can run LSAN on this test to
// catch memory leaks. Enabling LSAN with python multiprocessing has been
// challenging and we don't have a good solution yet.
TEST_F(TestE2EProcessGroup, TestTrainingLoop) {
runTrainingLoop();
}
} // namespace rpc
} // namespace distributed
} // namespace torch