| 25 | namespace oneflow { |
| 26 | |
| 27 | Maybe<void> CtrlBootstrap::InitProcessCtx(int64_t port, ProcessCtx* ret_process_ctx) { |
| 28 | std::vector<WorkerProcessInfo> worker_process_info_list; |
| 29 | worker_process_info_list.reserve(world_size()); |
| 30 | if (rank() == 0) { |
| 31 | WorkerProcessInfo worker_process_info; |
| 32 | { |
| 33 | worker_process_info.set_rank(rank()); |
| 34 | worker_process_info.set_port(port); |
| 35 | JUST(SetCurrentHostByMaster(&worker_process_info)); |
| 36 | } |
| 37 | worker_process_info_list.emplace_back(worker_process_info); |
| 38 | for (int64_t world_rank = 1; world_rank < world_size(); ++world_rank) { |
| 39 | std::string key = std::string("GetWorkerProcessInfo") + std::to_string(world_rank); |
| 40 | WorkerProcessInfo cur_work_process_info; |
| 41 | mut_bootstrap_client()->PullMasterKV(key, &cur_work_process_info); |
| 42 | CHECK_EQ_OR_RETURN(world_rank, worker_process_info_list.size()); |
| 43 | CHECK_EQ_OR_RETURN(world_rank, cur_work_process_info.rank()); |
| 44 | worker_process_info_list.emplace_back(cur_work_process_info); |
| 45 | } |
| 46 | } else { |
| 47 | std::string key = std::string("GetWorkerProcessInfo") + std::to_string(rank()); |
| 48 | WorkerProcessInfo cur_work_process_info; |
| 49 | { |
| 50 | cur_work_process_info.set_rank(rank()); |
| 51 | cur_work_process_info.set_port(port); |
| 52 | JUST(SetCurrentHostByWorker(&cur_work_process_info)); |
| 53 | } |
| 54 | mut_bootstrap_client()->PushMasterKV(key, cur_work_process_info); |
| 55 | } |
| 56 | |
| 57 | mut_bootstrap_client()->Barrier(__FILE__ ":" OF_PP_STRINGIZE(__LINE__)); |
| 58 | |
| 59 | if (rank() == 0) { |
| 60 | ret_process_ctx->set_rank(rank()); |
| 61 | ret_process_ctx->mutable_ctrl_addr()->Clear(); |
| 62 | for (const auto& worker_process_info : worker_process_info_list) { |
| 63 | Address* addr = ret_process_ctx->mutable_ctrl_addr()->Add(); |
| 64 | if (worker_process_info.has_host()) { addr->set_host(worker_process_info.host()); } |
| 65 | addr->set_port(worker_process_info.port()); |
| 66 | JUST(SetHostByMaster(addr, worker_process_info.rank())); |
| 67 | } |
| 68 | JUST(SetNodeSize(ret_process_ctx)); |
| 69 | mut_bootstrap_client()->PushMasterKV("BroadcastProcessCtx", *ret_process_ctx); |
| 70 | } else { |
| 71 | mut_bootstrap_client()->PullMasterKV("BroadcastProcessCtx", ret_process_ctx); |
| 72 | ret_process_ctx->set_rank(rank()); |
| 73 | } |
| 74 | |
| 75 | mut_bootstrap_client()->Barrier(__FILE__ ":" OF_PP_STRINGIZE(__LINE__)); |
| 76 | |
| 77 | VLOG(2) << "\n" << ret_process_ctx->DebugString(); |
| 78 | return Maybe<void>::Ok(); |
| 79 | } |
| 80 | |
| 81 | HostListCtrlBootstrap::HostListCtrlBootstrap(const EnvDesc& env_desc) : CtrlBootstrap() { |
| 82 | bootstrap_server_.reset(new HostListBootstrapServer(env_desc)); |