diff options
author | Jiri Simsa <jsimsa@google.com> | 2018-05-29 18:10:27 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-05-29 18:12:55 -0700 |
commit | 2c75dbfd2d37a3c06d34cc4b12682a63a75503f7 (patch) | |
tree | 710ffd2e3ec5376b650e38935fef44dde1f75e78 /tensorflow/core/util/rpc/call_container.h | |
parent | 7a4d278a3dbb71c0d707e2c5e99423489099f441 (diff) |
Making RPC op handle the case where cancellation manager is not initialized in OpKernelContext.
Fixes #19496
PiperOrigin-RevId: 198488860
Diffstat (limited to 'tensorflow/core/util/rpc/call_container.h')
-rw-r--r-- | tensorflow/core/util/rpc/call_container.h | 21 |
1 files changed, 14 insertions, 7 deletions
diff --git a/tensorflow/core/util/rpc/call_container.h b/tensorflow/core/util/rpc/call_container.h index e1226a7f16..39ead10815 100644 --- a/tensorflow/core/util/rpc/call_container.h +++ b/tensorflow/core/util/rpc/call_container.h @@ -102,7 +102,9 @@ CallContainer<Call>::CallContainer( typename CallContainer<Call>::StartCallFn start_call_fn) : ctx_(ctx), done_(std::move(done)), - token_(ctx->cancellation_manager()->get_cancellation_token()), + token_(ctx->cancellation_manager() != nullptr + ? ctx->cancellation_manager()->get_cancellation_token() + : CancellationManager::kInvalidToken), fail_fast_(fail_fast), try_rpc_(try_rpc), callback_destroyed_(new Notification) { @@ -110,7 +112,9 @@ CallContainer<Call>::CallContainer( // This will run when all RPCs are finished. reffed_status_callback_ = new ReffedStatusCallback([this](const Status& s) { - ctx_->cancellation_manager()->DeregisterCallback(token_); + if (token_ != CancellationManager::kInvalidToken) { + ctx_->cancellation_manager()->DeregisterCallback(token_); + } ctx_->SetStatus(s); done_(); callback_destroyed_->WaitForNotification(); @@ -125,11 +129,14 @@ CallContainer<Call>::CallContainer( std::shared_ptr<internal::NotifyWhenDestroyed> notify_when_destroyed( new internal::NotifyWhenDestroyed(callback_destroyed_)); std::shared_ptr<Notification> calls_started(new Notification); - bool is_cancelled = !ctx_->cancellation_manager()->RegisterCallback( - token_, [this, calls_started, notify_when_destroyed]() { - calls_started->WaitForNotification(); - StartCancel(); - }); + bool is_cancelled = false; + if (token_ != CancellationManager::kInvalidToken) { + is_cancelled = !ctx_->cancellation_manager()->RegisterCallback( + token_, [this, calls_started, notify_when_destroyed]() { + calls_started->WaitForNotification(); + StartCancel(); + }); + } for (int i = 0; i < num_calls; ++i) { create_call_fn(this, i); |