experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix verbs compile error (#9791)

Browse Source
This commit is contained in:
Liangliang He 2017-05-12 13:11:33 +08:00 committed by Benoit Steiner
parent eb0a59f003
commit a33022c147
4 changed files with 15 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/contrib/verbs/rdma.cc
View File

@ -778,11 +778,8 @@ void RdmaTensorBuffer::SendNextItem() {
EnqueueItem(key_with_step_id);
}
};
// Use default session (legacy_session_)
// TODO use WorkerSessionForSession
// need to pass in session handle
channel_->adapter_->worker_env_->session_mgr->LegacySession()
->rendezvous_mgr->RecvLocalAsync(step_id, parsed, cb);
channel_->adapter_->worker_env_->rendezvous_mgr
->RecvLocalAsync(step_id, parsed, cb);
}
}

15
tensorflow/contrib/verbs/rdma_rendezvous_mgr.cc
View File

@ -29,9 +29,9 @@ namespace tensorflow {
class RdmaRemoteRendezvous : public BaseRemoteRendezvous {
public:
RdmaRemoteRendezvous(const WorkerEnv* env, const string& worker_name,
RdmaRemoteRendezvous(const WorkerEnv* env,
int64 step_id, RdmaMgr* rdma_mgr)
: BaseRemoteRendezvous(env, worker_name, step_id, true),
: BaseRemoteRendezvous(env, step_id, true),
rdma_mgr_(rdma_mgr) {}
protected:
@ -133,15 +133,12 @@ void RdmaRemoteRendezvous::RecvFromRemoteAsync(
rb->SendNextItem();
}
RdmaRendezvousMgr::RdmaRendezvousMgr(const WorkerEnv* env,
const string& worker_name,
WorkerCacheInterface* worker_cache)
: BaseRendezvousMgr(env, worker_name) {}
RdmaRendezvousMgr::RdmaRendezvousMgr(const WorkerEnv* env)
: BaseRendezvousMgr(env) {}
BaseRemoteRendezvous* RdmaRendezvousMgr::Create(int64 step_id,
const WorkerEnv* worker_env,
const string& worker_name) {
return new RdmaRemoteRendezvous(worker_env, worker_name, step_id, rdma_mgr_);
const WorkerEnv* worker_env) {
return new RdmaRemoteRendezvous(worker_env, step_id, rdma_mgr_);
}
} // end namespace tensorflow

7
tensorflow/contrib/verbs/rdma_rendezvous_mgr.h
View File

@ -45,13 +45,12 @@ namespace tensorflow {
// RendezvousMgr must have keys generated by Rendezvous::CreateKey.
class RdmaRendezvousMgr : public BaseRendezvousMgr {
public:
explicit RdmaRendezvousMgr(const WorkerEnv* env, const string& worker_name,
WorkerCacheInterface* worker_cache);
explicit RdmaRendezvousMgr(const WorkerEnv* env);
void SetRdmaMgr(RdmaMgr* rdma_mgr) { rdma_mgr_ = rdma_mgr; }
protected:
BaseRemoteRendezvous* Create(int64 step_id, const WorkerEnv* worker_env,
const string& worker_name) override;
BaseRemoteRendezvous* Create(int64 step_id,
const WorkerEnv* worker_env) override;
private:
RdmaMgr* rdma_mgr_;

14
tensorflow/contrib/verbs/verbs_server_lib.cc
View File

@ -27,10 +27,8 @@ namespace tensorflow {
namespace {
// static utility function
RendezvousMgrInterface* NewRdmaRendezvousMgr(
const WorkerEnv* env, const string& worker_name,
WorkerCacheInterface* worker_cache) {
return new RdmaRendezvousMgr(env, worker_name, worker_cache);
RendezvousMgrInterface* NewRdmaRendezvousMgr(const WorkerEnv* env) {
return new RdmaRendezvousMgr(env);
}
} // namespace
@ -56,7 +54,7 @@ Status VerbsServer::ChannelCacheFactory(const ServerDef& server_def,
TF_RETURN_IF_ERROR(ParseChannelSpec(server_def, &channel_spec));
*channel_cache =
NewGrpcChannelCache(channel_spec, GetChannelCreationFunction(server_def));
NewGrpcChannelCache(channel_spec, GetChannelCreationFunction());
const string host_port = (*channel_cache)->TranslateTask(name_prefix);
int requested_port;
@ -86,11 +84,7 @@ Status VerbsServer::Init(ServiceInitFunction service_func,
rdma_mgr_ = new RdmaMgr(worker_env(), channel_cache_);
// set rdma_mgr for verbs_service and rdma_rendezvous_mgr
verbs_service_->SetRdmaMgr(rdma_mgr_);
// hardcoded to default session (legacy_session_)
// TODO: use WorkerSessionForSession
// need to pass in session handle
dynamic_cast<RdmaRendezvousMgr*>(
worker_env()->session_mgr->LegacySession()->rendezvous_mgr.get())
dynamic_cast<RdmaRendezvousMgr*>(worker_env()->rendezvous_mgr)
->SetRdmaMgr(rdma_mgr_);
}
return s;