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

Hold a reference to ensure the eager client outlives the function data op.

Browse Source 
PiperOrigin-RevId: 287567727
Change-Id: I86f060e22ab3ef7421c34a53f40008f418b99fa0
This commit is contained in:
Haoyu Zhang 2019-12-30 10:27:19 -08:00 committed by Jacques Pienaar
parent 3ae602667e
commit 8f23523ff8
2 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/distributed_runtime/eager/cluster_function_library_runtime.cc
View File

@ -131,7 +131,7 @@ void EagerClusterFunctionLibraryRuntime::Run(
function_data = &function_data_[handle]; function_data = &function_data_[handle];
} }
EagerClient* eager_client = function_data->eager_client; EagerClient* eager_client = function_data->eager_client.get();
if (eager_client == nullptr) { if (eager_client == nullptr) {
done(errors::Internal("Could not find eager client")); done(errors::Internal("Could not find eager client"));
return; return;
@ -195,7 +195,7 @@ void EagerClusterFunctionLibraryRuntime::CleanUp(
function_data = &function_data_[handle]; function_data = &function_data_[handle];
} }
EagerClient* eager_client = function_data->eager_client; EagerClient* eager_client = function_data->eager_client.get();
if (eager_client == nullptr) { if (eager_client == nullptr) {
done(errors::Internal("Could not find eager client")); done(errors::Internal("Could not find eager client"));
return; return;

8
tensorflow/core/distributed_runtime/eager/cluster_function_library_runtime.h
View File

@ -70,12 +70,16 @@ class EagerClusterFunctionLibraryRuntime
struct FunctionData { struct FunctionData {
const string target; const string target;
EagerClient* eager_client = nullptr; core::RefCountPtr<EagerClient> eager_client;
std::unique_ptr<EagerOperation> op; std::unique_ptr<EagerOperation> op;
FunctionData(const string& target, EagerClient* eager_client, FunctionData(const string& target, EagerClient* eager_client,
std::unique_ptr<EagerOperation> op) std::unique_ptr<EagerOperation> op)
: target(target), eager_client(eager_client), op(std::move(op)) {} : target(target),
eager_client(core::RefCountPtr<EagerClient>(eager_client)),
op(std::move(op)) {
eager_client->Ref();
}
}; };
mutable mutex mu_; mutable mutex mu_;