| 50 | GdrServer::~GdrServer() {} |
| 51 | |
| 52 | Status GdrServer::Init() { |
| 53 | RendezvousMgrCreationFunction rendezvous_mgr_func = |
| 54 | [this](const WorkerEnv* env) { |
| 55 | return new GdrRendezvousMgr(env, remote_memory_manager_.get()); |
| 56 | }; |
| 57 | WorkerCreationFunction worker_func = [this](WorkerEnv* env, |
| 58 | const ConfigProto& config) { |
| 59 | return std::unique_ptr<GdrWorker>( |
| 60 | new GdrWorker(env, config, remote_memory_manager_.get())); |
| 61 | }; |
| 62 | CollectiveMgrCreationFunction collective_mgr_func = |
| 63 | [this](const ConfigProto& config, const WorkerEnv* env, |
| 64 | WorkerCacheInterface* worker_cache) { |
| 65 | string unused; |
| 66 | string default_worker_name; |
| 67 | DeviceNameUtils::SplitDeviceName( |
| 68 | env->device_mgr->ListDevices()[0]->name(), &default_worker_name, |
| 69 | &unused); |
| 70 | |
| 71 | std::unique_ptr<DeviceResolverDistributed> dev_resolver( |
| 72 | new DeviceResolverDistributed(env->device_mgr, worker_cache, |
| 73 | default_worker_name)); |
| 74 | std::unique_ptr<CollectiveParamResolverDistributed> param_resolver( |
| 75 | new CollectiveParamResolverDistributed( |
| 76 | config, env->device_mgr, dev_resolver.get(), worker_cache, |
| 77 | default_worker_name)); |
| 78 | return new GdrCollectiveExecutorMgr( |
| 79 | config, env->device_mgr, std::move(dev_resolver), |
| 80 | std::move(param_resolver), worker_cache, default_worker_name, |
| 81 | remote_memory_manager_.get()); |
| 82 | }; |
| 83 | TF_RETURN_IF_ERROR(remote_memory_manager_->Init()); |
| 84 | |
| 85 | GrpcServerOptions opts; |
| 86 | opts.rendezvous_mgr_func = rendezvous_mgr_func; |
| 87 | opts.collective_mgr_func = collective_mgr_func; |
| 88 | opts.worker_func = worker_func; |
| 89 | return GrpcServer::Init(opts); |
| 90 | } |
| 91 | |
| 92 | Status GdrServer::Start() { |
| 93 | { |
no test coverage detected