未验证 提交 f25abbae 编写于 作者: X Xin Pan 提交者: GitHub

Merge pull request #1 from Yancey1989/panxin_polish_sparse

polish sparse update code
...@@ -63,6 +63,9 @@ bool RequestSendHandler::Handle(const std::string& varname, ...@@ -63,6 +63,9 @@ bool RequestSendHandler::Handle(const std::string& varname,
PADDLE_THROW("sync: Can not find server side var"); PADDLE_THROW("sync: Can not find server side var");
return false; return false;
} }
if (invar->IsType<framework::SelectedRows>()) {
rpc_server_->RecordSparseVar(invar);
}
} }
return true; return true;
......
...@@ -73,6 +73,19 @@ void RPCServer::ResetBarrierCounter() { ...@@ -73,6 +73,19 @@ void RPCServer::ResetBarrierCounter() {
t.second = 0; t.second = 0;
} }
} }
void RPCServer::RecordSparseVar(framework::Variable* sparse_var) {
std::unique_lock<std::mutex> lock(mutex_sparse_var_recorder_);
sparse_vars_.push_back(sparse_var);
}
void RPCServer::ResetSparseVarsRecorder() {
VLOG(3) << "RPCServer reset sparse vars recorder.";
std::unique_lock<std::mutex> lock(mutex_sparse_var_recorder_);
for (auto* var : sparse_vars_) {
var->GetMutable<framework::SelectedRows>()->mutable_rows()->clear();
}
sparse_vars_.clear();
}
void RPCServer::RegisterRPC(const std::string& rpc_name, void RPCServer::RegisterRPC(const std::string& rpc_name,
RequestHandler* handler, int thread_num) { RequestHandler* handler, int thread_num) {
......
...@@ -60,7 +60,10 @@ class RPCServer { ...@@ -60,7 +60,10 @@ class RPCServer {
void SetCond(const std::string& rpc_name); void SetCond(const std::string& rpc_name);
void WaitCond(const std::string& rpc_name); void WaitCond(const std::string& rpc_name);
void IncreaseBatchBarrier(const std::string rpc_name); void IncreaseBatchBarrier(const std::string rpc_name);
void ResetBarrierCounter(); void ResetBarrierCounter();
void RecordSparseVar(framework::Variable* sparse_var);
void ResetSparseVarsRecorder();
protected: protected:
virtual void ShutDownImpl() = 0; virtual void ShutDownImpl() = 0;
...@@ -74,6 +77,9 @@ class RPCServer { ...@@ -74,6 +77,9 @@ class RPCServer {
std::atomic<int> cur_cond_; std::atomic<int> cur_cond_;
std::condition_variable rpc_cond_; std::condition_variable rpc_cond_;
std::vector<framework::Variable*> sparse_vars_;
std::mutex mutex_sparse_var_recorder_;
protected: protected:
std::string bind_address_; std::string bind_address_;
std::atomic<int> exit_flag_; std::atomic<int> exit_flag_;
......
...@@ -146,6 +146,7 @@ void ListenAndServOp::RunSyncLoop(framework::Executor *executor, ...@@ -146,6 +146,7 @@ void ListenAndServOp::RunSyncLoop(framework::Executor *executor,
rpc_service_->SetCond(detail::kRequestGet); rpc_service_->SetCond(detail::kRequestGet);
rpc_service_->WaitBarrier(detail::kRequestGet); rpc_service_->WaitBarrier(detail::kRequestGet);
rpc_service_->ResetBarrierCounter(); rpc_service_->ResetBarrierCounter();
rpc_service_->ResetSparseVarsRecorder();
} // while(true) } // while(true)
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册