From fe1ff3c01a78801d7b1de4eb88f1835cd46cfd04 Mon Sep 17 00:00:00 2001 From: Derek Murray Date: Fri, 1 May 2020 15:04:34 -0700 Subject: [PATCH] [tf.data] Apply single-device multi component function optimization in `RunSync()`. Previously, `RunSync()` was missing the optimization from e5c6881c778aad9664da016075a0c6b43d523432 that enables a single-component multi-device function from being executed as a local function. When 624f1a0f82e79fc92c27b39c5db95a0a45b7a9d7 switched all tf.data functions to be considered as multi-device functions, the optimization became load-bearing, because occasionally it will dispatch a function to a ProcessFunctionLibraryRuntime that does not actually support multi-device functions, because it has no rendezvous factory. The fix is simple: update the call to `GetHandleOnDevice()` in `PrepareRunSync()` to pass the optional `include_multi_device = true` argument. PiperOrigin-RevId: 309481098 Change-Id: Id12d01cffc50399bab5711c7de4756c542873ec2 --- tensorflow/core/common_runtime/function.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/common_runtime/function.cc b/tensorflow/core/common_runtime/function.cc index 08619371833..72a08b4dc9d 100644 --- a/tensorflow/core/common_runtime/function.cc +++ b/tensorflow/core/common_runtime/function.cc @@ -1235,7 +1235,8 @@ Status FunctionLibraryRuntimeImpl::PrepareRunSync( run_opts->create_rendezvous = false; } - LocalHandle local_handle = parent_->GetHandleOnDevice(device_name_, handle); + LocalHandle local_handle = parent_->GetHandleOnDevice( + device_name_, handle, /*include_multi_device=*/true); if (local_handle == kInvalidLocalHandle) { *out_item = nullptr; return Status::OK();